ceos/unsloth
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

Formatting: ruff line-length 100, kwarg-spacing passes, drop blank after short local imports (#6079)

Browse source 
Raise ruff line-length to 100 and extend the local pre-commit format pipeline (def-signature magic-comma normalization, short multi-line assert collapse, kwarg '=' spacing, blank-line-after-short-import removal, adjacent string-literal / f-string+plain merge, redundant-pass pruning). Every transform re-checks the file AST and is dropped if it would differ; the whole-repo reformat is verified AST-identical per file and idempotent.
This commit is contained in:
Daniel Han 2026-06-08 04:24:13 -07:00 • committed by GitHub
commit 3ce187da02
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
377 changed files with 5945 additions and 11859 deletions
Download patch file Download diff file Expand all files Collapse all files

3
.pre-commit-config.yaml
View file

@ -14,5 +14,8 @@ repos:
entry: scripts/run_ruff_format.py
language: python
types: [python]
# Mirror ruff's [tool.ruff] extend-exclude so this hook does not
# half-process files ruff itself skips (which produced churn).
exclude: '(chat_templates|ollama_template_mappers|_auto_install|mapper)\.py$'
additional_dependencies:
- ruff==0.6.9

1
pyproject.toml
View file

@ -1308,6 +1308,7 @@ repository = "https://github.com/unslothai/unsloth"
[tool.ruff]
target-version = "py311"
line-length = 100
force-exclude = true
extend-exclude = [
"*chat_templates.py",

64
scripts/check_frontend_dep_removal.py
View file

@ -52,9 +52,7 @@ EXPECTED_NOISE_FILES = {
}
# Only quoted-string occurrences in these file types can be module specifiers.
JS_LIKE_EXT = re.compile(
r"\.(ts|tsx|js|jsx|mjs|cjs|html|htm|css|scss|sass|json|jsonc)$"
)
JS_LIKE_EXT = re.compile(r"\.(ts|tsx|js|jsx|mjs|cjs|html|htm|css|scss|sass|json|jsonc)$")
# Files where JS-syntactic import patterns (static/dynamic/require/re-export)
# could be a real module reference. Markdown gets a separate gate (.mdx is
# real ESM; .md code fences are not).
@ -273,9 +271,7 @@ def classify(pkg: str, file: str, content: str) -> str | None:
if is_script and re.search(rf"\bimport\(\s*['\"]{esc}{sub}['\"]\s*\)", content):
return "dynamic_import"
# require / require.resolve
if is_script and re.search(
rf"\brequire(?:\.resolve)?\(\s*['\"]{esc}{sub}['\"]\s*\)", content
):
if is_script and re.search(rf"\brequire(?:\.resolve)?\(\s*['\"]{esc}{sub}['\"]\s*\)", content):
return "require"
# Re-exports: `export * from "pkg"`, `export { x } from "pkg"`,
# `export type { Foo } from "pkg"`. Multi-line supported.
@ -289,16 +285,12 @@ def classify(pkg: str, file: str, content: str) -> str | None:
# segment bounded by a quote / `#` / `?` or a subpath `/`, so
# `/node_modules/foo-extra/...` is NOT treated as usage of `foo`.
html_pkg = rf"{esc}(?:/[^'\"#?]*)?(?=['\"#?])"
if is_html and re.search(
rf"<script[^>]*src\s*=\s*['\"][^'\"]*/{html_pkg}", content
):
if is_html and re.search(rf"<script[^>]*src\s*=\s*['\"][^'\"]*/{html_pkg}", content):
return "html_script"
if is_html and re.search(rf"<link[^>]*href\s*=\s*['\"][^'\"]*/{html_pkg}", content):
return "html_link"
# TypeScript triple-slash
if is_ts and re.search(
rf"///\s*<reference\s+types\s*=\s*['\"]{esc}{sub}['\"]", content
):
if is_ts and re.search(rf"///\s*<reference\s+types\s*=\s*['\"]{esc}{sub}['\"]", content):
return "tsc_triple_slash"
# new URL("pkg/...", import.meta.url)
if is_script and re.search(rf"\bnew\s+URL\(\s*['\"]{esc}{sub}['\"]", content):
@ -544,19 +536,13 @@ def _next_real_bin(words: list[str], idx: int) -> str | None:
if first in {"npx", "pnpx", "bunx"} and idx + 1 < len(words):
idx += 1
continue
if (
first in {"pnpm", "yarn"}
and idx + 2 < len(words)
and words[idx + 1] in {"exec", "dlx"}
):
if first in {"pnpm", "yarn"} and idx + 2 < len(words) and words[idx + 1] in {"exec", "dlx"}:
idx += 2
continue
# 3. Wrapper bin (cross-env, dotenv, etc.). Skip the wrapper's
# own flags and any subsequent env-prefix tokens, then re-loop.
bin_token = first.removeprefix("./node_modules/.bin/").removeprefix(
"node_modules/.bin/"
)
bin_token = first.removeprefix("./node_modules/.bin/").removeprefix("node_modules/.bin/")
if bin_token in _SCRIPT_WRAPPERS and bin_token not in seen_wrappers:
seen_wrappers.add(bin_token)
idx += 1
@ -585,9 +571,7 @@ def _next_real_bin(words: list[str], idx: int) -> str | None:
return None
def scripts_bin_refs(
head_pkg: dict, bin_to_pkg: dict[str, str]
) -> dict[str, list[str]]:
def scripts_bin_refs(head_pkg: dict, bin_to_pkg: dict[str, str]) -> dict[str, list[str]]:
"""Return `{package_name: ['scripts.X: cmd', ...]}` listing every
package referenced via its bin name in package.json scripts.
@ -652,11 +636,7 @@ def tsconfig_compiler_types_refs() -> set[str]:
if not isinstance(t, str):
continue
# `vite/client` resolves to `vite` package.
pkg = (
t.split("/", 1)[0]
if not t.startswith("@")
else "/".join(t.split("/", 2)[:2])
)
pkg = t.split("/", 1)[0] if not t.startswith("@") else "/".join(t.split("/", 2)[:2])
out.add(pkg)
return out
@ -820,9 +800,7 @@ _file_lines_cache: dict[str, list[str]] = {}
def _read_file(path: str) -> list[str]:
if path not in _file_lines_cache:
try:
_file_lines_cache[path] = (
Path(path).read_text(errors = "replace").splitlines()
)
_file_lines_cache[path] = Path(path).read_text(errors = "replace").splitlines()
except (OSError, UnicodeDecodeError):
_file_lines_cache[path] = []
return _file_lines_cache[path]
@ -952,18 +930,14 @@ def find_types_runtime_usage(pkg: str, tsc_types: set[str]) -> list[Hit]:
def main() -> int:
p = argparse.ArgumentParser(
description = __doc__, formatter_class = argparse.RawTextHelpFormatter
)
p = argparse.ArgumentParser(description = __doc__, formatter_class = argparse.RawTextHelpFormatter)
p.add_argument(
"--base",
default = "origin/main",
help = "git ref to diff against (default: origin/main). "
"Examples: HEAD~1, main, a-tag, a-sha.",
)
p.add_argument(
"--base-pkg", help = "optional override: read base package.json from this path"
)
p.add_argument("--base-pkg", help = "optional override: read base package.json from this path")
p.add_argument(
"--base-lock",
help = "optional override: read base package-lock.json from this path. "
@ -1057,9 +1031,7 @@ def main() -> int:
print(f" - {w}")
print()
if missing_imports:
print(
f"Imports without a matching package.json dep ({len(missing_imports)}):"
)
print(f"Imports without a matching package.json dep ({len(missing_imports)}):")
for file, ln, spec in missing_imports[:20]:
print(f" - {file}:{ln} imports '{spec}'")
print()
@ -1097,9 +1069,7 @@ def main() -> int:
return 1
return 0
print(
f"Checking {len(removed)} removed package(s) from studio/frontend/package.json"
)
print(f"Checking {len(removed)} removed package(s) from studio/frontend/package.json")
print(f"Base: {args.base} Head: working tree")
print()
@ -1127,9 +1097,7 @@ def main() -> int:
top = f"node_modules/{name}"
top_path = top if top in reachable_paths else None
nested = sorted(
p
for p in reachable_paths
if p != top and p.endswith(f"/node_modules/{name}")
p for p in reachable_paths if p != top and p.endswith(f"/node_modules/{name}")
)
return top_path, nested
@ -1177,9 +1145,7 @@ def main() -> int:
_print_hygiene()
if failures:
print(
f"FAIL: {len(failures)} removed package(s) still referenced and not resolvable"
)
print(f"FAIL: {len(failures)} removed package(s) still referenced and not resolvable")
for name, _ in failures:
print(f" - {name}")
return 1

11
scripts/check_new_install_scripts.py
View file

@ -53,9 +53,7 @@ HIGH = "HIGH"
class Finding:
__slots__ = ("severity", "name", "version", "kind", "detail")
def __init__(
self, severity: str, name: str, version: str, kind: str, detail: str
) -> None:
def __init__(self, severity: str, name: str, version: str, kind: str, detail: str) -> None:
self.severity = severity
self.name = name
self.version = version
@ -206,9 +204,7 @@ def diff_new_install_scripts(base_lock: dict, head_lock: dict) -> list[Finding]:
continue # pre-existing install-script dep; not in scope
name = head[key]
# key is "name@version"; rsplit("@", 1) handles scoped names.
version = (
key[len(name) + 1 :] if key.startswith(name + "@") else "<unversioned>"
)
version = key[len(name) + 1 :] if key.startswith(name + "@") else "<unversioned>"
scripts = _fetch_registry_scripts(name, version)
if scripts:
detail = "; ".join(f"{h}={cmd!r}" for h, cmd in scripts.items())
@ -238,8 +234,7 @@ def diff_new_install_scripts(base_lock: dict, head_lock: dict) -> list[Finding]:
def main(argv: list[str] | None = None) -> int:
parser = argparse.ArgumentParser(
description = (
"Diff two package-lock.json files and refuse any newly-"
"added install-script dep."
"Diff two package-lock.json files and refuse any newly-added install-script dep."
),
)
parser.add_argument(

499
scripts/enforce_kwargs_spacing.py
View file

@ -1,5 +1,10 @@
#!/usr/bin/env python3
"""Ensure keyword arguments use spaces around '=', prune redundant pass statements."""
"""Ensure keyword arguments use spaces around '=', prune redundant pass statements,
drop the blank line after a short indented import block, merge adjacent same-line
string literals, normalize def-signature magic commas (pre-ruff) so a def with
>= 3 params and a default goes one-per-line while everything else stays
collapsible, and collapse a short multi-line assert onto one line (pre-ruff) by
stripping the magic trailing comma that holds it open."""
from __future__ import annotations
@ -123,9 +128,7 @@ def remove_redundant_passes(text: str) -> tuple[str, bool]:
lines = text.splitlines(keepends=True)
changed = False
for node in sorted(
redundant, key=lambda item: (item.lineno, item.col_offset), reverse=True
):
for node in sorted(redundant, key=lambda item: (item.lineno, item.col_offset), reverse=True):
start = node.lineno - 1
end = (node.end_lineno or node.lineno) - 1
if start >= len(lines):
@ -160,7 +163,470 @@ def remove_redundant_passes(text: str) -> tuple[str, bool]:
return "".join(result_lines), changed
def process_file(path: Path) -> bool:
def remove_blank_after_short_import(text: str) -> tuple[str, bool]:
"""Drop blank line(s) after an import block in a *small* nested suite.
Inside an indented suite of <= 3 statements (function/try/if/with/etc., never
module level), when a run of consecutive ``import`` / ``from ... import``
statements is directly followed -- across one or more blank lines and nothing
else -- by another statement in the same suite, remove those blank lines so
the import sits next to the code that uses it. A comment in the gap blocks the
rule. Removing blank lines never changes the AST, so this is always
semantics-preserving.
"""
try:
tree = ast.parse(text)
except SyntaxError:
return text, False
lines = text.splitlines(keepends=True)
import_types = (ast.Import, ast.ImportFrom)
drop: set[int] = set() # 1-based physical line numbers to delete
def suites_of(node: ast.AST) -> list[list[ast.stmt]]:
if isinstance(node, ast.Module):
return [] # module-level import spacing is left alone
out: list[list[ast.stmt]] = []
for attr in ("body", "orelse", "finalbody"):
val = getattr(node, attr, None)
if isinstance(val, list) and val and all(isinstance(s, ast.stmt) for s in val):
out.append(val)
return out
for node in ast.walk(tree):
for suite in suites_of(node):
if len(suite) > 3: # only small blocks
continue
i = 0
while i < len(suite):
if not isinstance(suite[i], import_types):
i += 1
continue
j = i
while j + 1 < len(suite) and isinstance(suite[j + 1], import_types):
j += 1
if j + 1 < len(suite): # an import block followed by another statement
last_imp, nxt = suite[j], suite[j + 1]
gap = range((last_imp.end_lineno or last_imp.lineno) + 1, nxt.lineno)
nums = [n for n in gap if 1 <= n <= len(lines)]
if nums and all(lines[n - 1].strip() == "" for n in nums):
drop.update(nums)
i = j + 1
if not drop:
return text, False
kept = [ln for idx, ln in enumerate(lines, start=1) if idx not in drop]
return "".join(kept), True
_STRING_TRIVIA = (tokenize.NL, tokenize.NEWLINE, tokenize.COMMENT, tokenize.INDENT, tokenize.DEDENT)
_DEF_MIN_PARAMS_FOR_MULTILINE = 3 # signatures with < this many params stay one line
def _def_specs_by_line(tree: ast.AST) -> dict[int, tuple[int, bool]]:
"""Map the line of each def keyword to (param count, has-any-default).
One def per line, so the line is a stable key. ``*`` / ``/`` markers are not
parameters and are not counted. A default exists if any positional default
is present or any keyword-only default is not ``None`` (a ``None`` entry in
``kw_defaults`` means a required keyword-only arg).
"""
out: dict[int, tuple[int, bool]] = {}
for node in ast.walk(tree):
if isinstance(node, (ast.FunctionDef, ast.AsyncFunctionDef)):
a = node.args
count = (
len(a.posonlyargs)
+ len(a.args)
+ len(a.kwonlyargs)
+ (1 if a.vararg else 0)
+ (1 if a.kwarg else 0)
)
has_default = bool(a.defaults) or any(d is not None for d in a.kw_defaults)
out[node.lineno] = (count, has_default)
return out
def normalize_def_trailing_comma(text: str) -> tuple[str, bool]:
"""Force a def / async-def signature one-per-line iff it has >= 3 parameters
AND at least one default value; otherwise keep it collapsible.
Rationale: signatures with defaults read better one parameter per line, but
only once they are non-trivial (< 3 params always stay on one line). A
signature with >= 3 params and a default gets a magic trailing comma added
(ruff then wraps it one-per-line regardless of length); every other
signature has its trailing comma stripped so ruff collapses it onto one line
when it fits (and wraps a genuinely long one by length alone).
Function-definition parameter lists only, never call sites or collection
literals. Parameter counts and defaults come from the AST. Run BEFORE ruff
format. Adding or removing a def trailing comma never changes the AST, which
is re-checked before returning.
"""
try:
tree = ast.parse(text)
toks = list(tokenize.generate_tokens(io.StringIO(text).readline))
except (tokenize.TokenError, IndentationError, SyntaxError):
return text, False
specs = _def_specs_by_line(tree)
n = len(toks)
edits: list[tuple[int, int, str]] = [] # (row, col, "del" | "ins")
i = 0
while i < n:
t = toks[i]
if t.type == tokenize.NAME and t.string == "def" and t.start[0] in specs:
cnt, has_default = specs[t.start[0]]
force_multiline = cnt >= _DEF_MIN_PARAMS_FOR_MULTILINE and has_default
j = i + 1
while j < n and not (toks[j].type == tokenize.OP and toks[j].string == "("):
if toks[j].type == tokenize.NEWLINE:
break
j += 1
if j < n and toks[j].type == tokenize.OP and toks[j].string == "(":
depth = 0
k = j
while k < n:
tk = toks[k]
if tk.type == tokenize.OP and tk.string == "(":
depth += 1
elif tk.type == tokenize.OP and tk.string == ")":
depth -= 1
if depth == 0:
m = k - 1
while m > j and toks[m].type in _STRING_TRIVIA:
m -= 1
last = toks[m]
has_comma = last.type == tokenize.OP and last.string == ","
empty = m == j # nothing between ( and )
if force_multiline and not has_comma and not empty:
edits.append((last.end[0], last.end[1], "ins"))
elif not force_multiline and has_comma:
edits.append((last.start[0], last.start[1], "del"))
break
k += 1
i = k + 1
continue
i += 1
if not edits:
return text, False
lines = text.splitlines(keepends=True)
for row, col, kind in sorted(edits, reverse=True):
ln = lines[row - 1]
if kind == "del":
if col < len(ln) and ln[col] == ",":
lines[row - 1] = ln[:col] + ln[col + 1 :]
else: # ins
lines[row - 1] = ln[:col] + "," + ln[col:]
out = "".join(lines)
try:
if ast.dump(ast.parse(out)) != ast.dump(ast.parse(text)):
return text, False
except SyntaxError:
return text, False
return out, True
def _split_string_token(s: str) -> tuple[str, str, str] | None:
"""Split a string literal's source into (prefix, quote, body).
``prefix`` is the letters before the opening quote (``r``/``f``/``b``/``u``
in any case/order), ``quote`` is the opening delimiter (``'``, ``"``,
``'''`` or ``\"\"\"``) and ``body`` is everything between the delimiters.
Returns ``None`` if ``s`` is not a recognizable string literal.
"""
i = 0
while i < len(s) and s[i] not in ("'", '"'):
i += 1
if i >= len(s):
return None
prefix, rest = s[:i], s[i:]
for q in ('"""', "'''", '"', "'"):
if rest.startswith(q) and rest.endswith(q) and len(rest) >= 2 * len(q):
return prefix, q, rest[len(q) : len(rest) - len(q)]
return None
# A "piece" is one string literal in source: a plain STRING token, or a whole
# f-string spanning FSTRING_START..FSTRING_END. (kind, (row, col0), (row, col1), raw)
def _string_pieces(
toks: list[tokenize.TokenInfo], lines: list[str]
) -> list[tuple[str, tuple[int, int], tuple[int, int], str | None]]:
pieces: list[tuple[str, tuple[int, int], tuple[int, int], str | None]] = []
n = len(toks)
def raw_of(start: tuple[int, int], end: tuple[int, int]) -> str | None:
if start[0] != end[0]: # only single-physical-line pieces are mergeable
return None
return lines[start[0] - 1][start[1] : end[1]]
i = 0
while i < n:
t = toks[i]
if t.type == tokenize.STRING:
pieces.append(("str", t.start, t.end, raw_of(t.start, t.end)))
i += 1
elif t.type == tokenize.FSTRING_START:
depth = 0
j = i
while j < n: # walk to the matching FSTRING_END (f-strings can nest)
if toks[j].type == tokenize.FSTRING_START:
depth += 1
elif toks[j].type == tokenize.FSTRING_END:
depth -= 1
if depth == 0:
break
j += 1
end = toks[j].end
pieces.append(("f", t.start, end, raw_of(t.start, end)))
i = j + 1
else:
pieces.append(("other", t.start, t.end, None))
i += 1
return pieces
def _merge_string_run(pieces: list[tuple[str, str]]) -> str | None:
"""Merge a run of adjacent string pieces into one literal's source text.
``pieces`` is a list of ``(kind, raw_source)`` where kind is ``"str"`` or
``"f"``. Rules: bytes are left side-by-side (return ``None``); a run with no
f-string merges plain/raw/unicode pieces sharing one prefix+quote by simple
body concatenation; a run mixing an f-string with at least one plain string
(and no bytes, no raw) folds into a single f-string -- f pieces keep their
bodies verbatim and plain pieces have their braces escaped (``{`` -> ``{{``).
Runs of only f-strings are left side-by-side. The caller re-checks the file
AST and drops the change if it differs, so any subtle case (e.g. ``\\N{...}``)
that this would mis-handle is caught and skipped.
"""
parsed = []
for kind, raw in pieces:
pqb = _split_string_token(raw)
if pqb is None:
return None
prefix, quote, body = pqb
if "b" in prefix.lower():
return None # bytes: leave side-by-side
parsed.append((kind, prefix, quote, body))
if len({p[2] for p in parsed}) != 1:
return None # mixed quote style: not a safe textual merge
quote = parsed[0][2]
if not any(p[0] == "f" for p in parsed):
# No f-string: merge plain/raw/unicode sharing one prefix by concatenation.
if len({p[1].lower() for p in parsed}) != 1:
return None
return f"{parsed[0][1]}{quote}{''.join(p[3] for p in parsed)}{quote}"
# f-string fold only when a plain string is glued onto an f-string; a run of
# only f-strings is left side-by-side (folding long ones would force ruff to
# re-wrap the surrounding statement).
if all(p[0] == "f" for p in parsed):
return None
# raw mixed with f is too subtle (backslash + brace escaping) -> skip.
if any("r" in p[1].lower() for p in parsed):
return None
body = "".join(
b if kind == "f" else b.replace("{", "{{").replace("}", "}}")
for kind, _pfx, _q, b in parsed
)
return f"f{quote}{body}{quote}"
_LINE_LENGTH = 100 # ruff line-length; an f-fold must not push a statement past it
def _enclosing_stmt(tree: ast.AST, row: int) -> ast.stmt | None:
"""The innermost statement whose physical-line span contains ``row``."""
best: tuple[ast.stmt, int] | None = None
for node in ast.walk(tree):
if isinstance(node, ast.stmt):
lo = node.lineno
hi = node.end_lineno or lo
if lo <= row <= hi and (best is None or hi - lo < best[1]):
best = (node, hi - lo)
return best[0] if best else None
def _fold_collapses(
tree: ast.AST, lines: list[str], row: int, c0: int, c1: int, merged: str
) -> bool:
"""Whether an f-string fold at ``row[c0:c1]`` -> ``merged`` is safe to apply.
Only ``assert`` statements wrap awkwardly when a message is folded: ruff
parenthesizes the *condition* once ``assert cond, msg`` no longer fits on one
line. For every other construct (call argument, ``raise``, assignment, ...) a
folded long message wraps acceptably, so the fold is always allowed. For an
``assert`` the fold is allowed only when the statement is already one physical
line, or its estimated one-line length after folding fits the line length;
otherwise the message is left side-by-side.
"""
stmt = _enclosing_stmt(tree, row)
if not isinstance(stmt, ast.Assert):
return True
lo, hi = stmt.lineno, stmt.end_lineno or stmt.lineno
if lo == hi:
return True
seg = []
for k in range(lo, hi + 1):
ln = lines[k - 1].rstrip("\n")
if k == row:
ln = ln[:c0] + merged + ln[c1:]
seg.append(ln)
indent = len(seg[0]) - len(seg[0].lstrip())
# Conservative over-estimate: join continuation lines with a single space
# (ruff joins bracketed wraps with none), so borderline cases skip the fold.
joined = " ".join(s.strip() for s in seg)
return indent + len(joined) <= _LINE_LENGTH
def merge_adjacent_string_literals(text: str) -> tuple[str, bool]:
"""Merge a run of adjacent string literals on ONE physical line into a single
literal (the ``"a" "b"`` form ruff emits when it collapses an implicit
concatenation). Plain/raw/unicode runs merge by concatenation; a run mixing
an f-string with a plain string folds into one f-string (plain parts' braces
escaped) -- but only when the statement still fits on one line, so a long
message is left side-by-side rather than forcing the statement to re-wrap.
Runs of only f-strings, and bytes, are left side-by-side. The whole file's AST
is re-checked and the change is dropped if it would differ, so the transform
can never change meaning.
"""
try:
toks = list(tokenize.generate_tokens(io.StringIO(text).readline))
tree = ast.parse(text)
except (tokenize.TokenError, IndentationError, SyntaxError):
return text, False
lines = text.splitlines(keepends=True)
pieces = _string_pieces(toks, lines)
# Group consecutive mergeable pieces (str/f, single line, same physical line).
runs: list[list[tuple[str, tuple[int, int], tuple[int, int], str]]] = []
cur: list[tuple[str, tuple[int, int], tuple[int, int], str]] = []
for kind, start, end, raw in pieces:
if kind in ("str", "f") and raw is not None:
if cur and cur[-1][2][0] != start[0]:
if len(cur) >= 2:
runs.append(cur)
cur = []
cur.append((kind, start, end, raw))
else:
if len(cur) >= 2:
runs.append(cur)
cur = []
if len(cur) >= 2:
runs.append(cur)
if not runs:
return text, False
edits = []
for run in runs:
merged = _merge_string_run([(kind, raw) for kind, _s, _e, raw in run])
if merged is None:
continue
row, c0, c1 = run[0][1][0], run[0][1][1], run[-1][2][1]
# An f-string fold must not push its statement onto extra lines; a plain
# concatenation always collapses cleanly so it skips this check.
if any(kind == "f" for kind, _s, _e, _r in run) and not _fold_collapses(
tree, lines, row, c0, c1, merged
):
continue
edits.append((row, c0, c1, merged))
if not edits:
return text, False
for row, c0, c1, repl in sorted(edits, key=lambda e: (e[0], e[1]), reverse=True):
ln = lines[row - 1]
lines[row - 1] = ln[:c0] + repl + ln[c1:]
out = "".join(lines)
try:
if ast.dump(ast.parse(text)) != ast.dump(ast.parse(out)):
return text, False
except SyntaxError:
return text, False
return out, True
def collapse_short_asserts(text: str) -> tuple[str, bool]:
"""Collapse a multi-line ``assert`` onto one line when it would fit.
An ``assert`` is often kept multi-line only by a magic trailing comma inside
a collection / call / tuple-message (``assert x == {\"a\": 1,}`` written
across lines). When the whole statement's estimated one-line length fits the
line length, strip those trailing commas (the comma before a ``)`` / ``]`` /
``}``) so ruff joins it back onto one line on the following format pass.
Run BEFORE ruff format. Skips any assert that contains a comment (a comment
forces ruff to keep it multi-line, which would oscillate). Stripping a
trailing comma is non-semantic except for a one-element tuple ``(x,)``; the
file AST is re-checked and any assert whose strip would change it is left
alone.
"""
try:
tree = ast.parse(text)
toks = list(tokenize.generate_tokens(io.StringIO(text).readline))
except (tokenize.TokenError, IndentationError, SyntaxError):
return text, False
lines = text.splitlines(keepends=True)
multiline = [
(n.lineno, n.end_lineno)
for n in ast.walk(tree)
if isinstance(n, ast.Assert) and (n.end_lineno or n.lineno) > n.lineno
]
if not multiline:
return text, False
comment_rows = {t.start[0] for t in toks if t.type == tokenize.COMMENT}
targets = [] # (lo, hi) spans whose one-line form fits and have no comment
for lo, hi in multiline:
if any(lo <= r <= hi for r in comment_rows):
continue # a comment would keep ruff multi-line -> never collapses
seg = [lines[k].rstrip("\n") for k in range(lo - 1, hi)]
indent = len(seg[0]) - len(seg[0].lstrip())
# Over-estimate (join with a space; keep the comma) so a "fits" verdict
# is always at least as long as ruff's real one-line output -> no fight.
if indent + len(" ".join(s.strip() for s in seg)) <= _LINE_LENGTH:
targets.append((lo, hi))
if not targets:
return text, False
# Trailing commas (a ',' whose next significant token is a closer), grouped
# by the target assert they belong to.
sig = [t for t in toks if t.type not in _STRING_TRIVIA]
by_target: dict[tuple[int, int], list[tuple[int, int]]] = defaultdict(list)
for i, t in enumerate(sig):
if t.type == tokenize.OP and t.string == ",":
nxt = sig[i + 1] if i + 1 < len(sig) else None
if nxt and nxt.type == tokenize.OP and nxt.string in (")", "]", "}"):
for lo, hi in targets:
if lo <= t.start[0] <= hi:
by_target[(lo, hi)].append(t.start)
break
if not by_target:
return text, False
base_dump = ast.dump(tree)
working = lines[:]
changed = False
for positions in by_target.values(): # apply per assert; skip any that break AST
trial = working[:]
for row, col in sorted(positions, reverse=True):
ln = trial[row - 1]
if col < len(ln) and ln[col] == ",":
trial[row - 1] = ln[:col] + ln[col + 1 :]
try:
if ast.dump(ast.parse("".join(trial))) == base_dump:
working, changed = trial, True
except SyntaxError:
pass
return ("".join(working), True) if changed else (text, False)
def process_file(path: Path, pre: bool = False) -> bool:
try:
with tokenize.open(path) as handle:
original = handle.read()
@ -169,9 +635,23 @@ def process_file(path: Path) -> bool:
print(f"Failed to read {path}: {exc}", file=sys.stderr)
return False
if pre:
# Pre-ruff: normalize def-signature magic commas (>=3 params + a default
# add so ruff forces one-per-line; everything else strips so ruff
# collapses), and strip the magic trailing comma from a short multi-line
# assert so ruff joins it onto one line. Everything else runs post-ruff.
updated, normalized = normalize_def_trailing_comma(original)
updated, collapsed = collapse_short_asserts(updated)
if normalized or collapsed:
_atomic_write_text(path, updated, encoding)
return True
return False
updated, changed = enforce_spacing(original)
updated, blanked = remove_blank_after_short_import(updated)
updated, merged = merge_adjacent_string_literals(updated)
updated, removed = remove_redundant_passes(updated)
if changed or removed:
if changed or blanked or merged or removed:
_atomic_write_text(path, updated, encoding)
return True
return False
@ -180,6 +660,11 @@ def process_file(path: Path) -> bool:
def main(argv: list[str]) -> int:
parser = argparse.ArgumentParser(description=__doc__)
parser.add_argument("files", nargs="+", help="Python files to fix")
parser.add_argument(
"--pre",
action="store_true",
help="pre-ruff pass: normalize def-signature commas + collapse short multi-line asserts",
)
args = parser.parse_args(argv)
touched: list[Path] = []
@ -192,7 +677,7 @@ def main(argv: list[str]) -> int:
continue
if not path.exists() or path.is_dir():
continue
if process_file(path):
if process_file(path, pre=args.pre):
touched.append(path)
if touched:

8
scripts/lint_workflow_triggers.py
View file

@ -47,9 +47,7 @@ from pathlib import Path
try:
import yaml
except ImportError:
print(
"ERROR: PyYAML is required. Install with 'pip install pyyaml'", file = sys.stderr
)
print("ERROR: PyYAML is required. Install with 'pip install pyyaml'", file = sys.stderr)
sys.exit(2)
REPO_ROOT = Path(__file__).resolve().parents[1]
@ -153,9 +151,7 @@ def main() -> int:
)
if findings:
print(
"Workflow trigger lint failed with the following issues:", file = sys.stderr
)
print("Workflow trigger lint failed with the following issues:", file = sys.stderr)
for f in findings:
print(f" - {f}", file = sys.stderr)
return 1

9
scripts/lockfile_supply_chain_audit.py
View file

@ -541,9 +541,7 @@ def audit_npm_lockfile(path: Path) -> list[Finding]:
path = str(path),
package = key,
kind = "blocked-known-malicious",
detail = (
f"{pkg_name}@{version} is on the " "BLOCKED_NPM_VERSIONS list"
),
detail = (f"{pkg_name}@{version} is on the BLOCKED_NPM_VERSIONS list"),
)
)
@ -765,10 +763,7 @@ def main(argv: list[str] | None = None) -> int:
"--cargo-lockfile",
action = "append",
default = None,
help = (
"Path to a Cargo.lock (repeatable). "
"Default: studio/src-tauri/Cargo.lock."
),
help = ("Path to a Cargo.lock (repeatable). Default: studio/src-tauri/Cargo.lock."),
)
parser.add_argument(
"--strict",

20
scripts/notebook_to_python.py
View file

@ -186,9 +186,7 @@ def convert_cell_to_python(source: str, *, allow_shell: bool = True) -> str:
cmd_lines.append(lines[i].strip())
full_cmd = "\n".join(cmd_lines)
result.extend(
_emit_shell_command(indent, full_cmd, allow_shell = allow_shell)
)
result.extend(_emit_shell_command(indent, full_cmd, allow_shell = allow_shell))
# %cd path -> os.chdir(path)
elif stripped.startswith("%cd "):
@ -313,9 +311,7 @@ def convert_notebook_to_script(
# Generate output filename
output_filename = filename.replace(".ipynb", ".py")
# Clean up filename
output_filename = (
output_filename.replace("(", "").replace(")", "").replace("-", "_")
)
output_filename = output_filename.replace("(", "").replace(")", "").replace("-", "_")
# Add output directory if specified
if output_dir:
@ -337,9 +333,7 @@ def convert_notebook_to_script(
def main():
import argparse
class Formatter(
argparse.ArgumentDefaultsHelpFormatter, argparse.RawDescriptionHelpFormatter
):
class Formatter(argparse.ArgumentDefaultsHelpFormatter, argparse.RawDescriptionHelpFormatter):
pass
parser = argparse.ArgumentParser(
@ -353,12 +347,8 @@ Examples:
python notebook_to_python.py https://github.com/unslothai/notebooks/blob/main/nb/Oute_TTS_(1B).ipynb
""",
)
parser.add_argument(
"notebooks", nargs = "+", help = "Notebook files or URLs to convert."
)
parser.add_argument(
"-o", "--output", dest = "output_dir", default = ".", help = "Output directory."
)
parser.add_argument("notebooks", nargs = "+", help = "Notebook files or URLs to convert.")
parser.add_argument("-o", "--output", dest = "output_dir", default = ".", help = "Output directory.")
# Default True for backwards compatibility: existing Colab notebooks
# routinely use pipes / redirection / interpolation in `!cmd` lines
# and the converted script needs to keep working. Operators who

61
scripts/notebook_validator.py
View file

@ -92,9 +92,7 @@ COLAB_ORACLE_FILES: dict[str, str] = {
"apt-list-gpu.txt": "colab_apt_list.gpu.txt",
"os-info-gpu.txt": "colab_os_info.gpu.txt",
}
COLAB_ORACLE_BASE_URL = (
"https://raw.githubusercontent.com/googlecolab/backend-info/main/"
)
COLAB_ORACLE_BASE_URL = "https://raw.githubusercontent.com/googlecolab/backend-info/main/"
# ----- Compat tables. PRs add rows as new releases land. ----- #
@ -195,9 +193,7 @@ def install_cells(nb: dict[str, Any]) -> list[tuple[int, str]]:
if first and first[0].strip().startswith("%%capture"):
out.append((i, src))
continue
if re.search(
r"^[ \t]*!\s*(uv\s+)?pip\s+(install|uninstall)\b", src, re.MULTILINE
):
if re.search(r"^[ \t]*!\s*(uv\s+)?pip\s+(install|uninstall)\b", src, re.MULTILINE):
out.append((i, src))
return out
@ -331,9 +327,7 @@ def parse_pip_line(line: str, line_no: int = 0) -> PipInvocation | None:
if t in ("install", "uninstall"):
continue
packages.append(t)
return PipInvocation(
tool = tool, flags = flags, packages = packages, raw = line, line_no = line_no
)
return PipInvocation(tool = tool, flags = flags, packages = packages, raw = line, line_no = line_no)
def _glue_line_continuations(text: str) -> list[tuple[int, str]]:
@ -418,9 +412,7 @@ def pypi_metadata(name: str, version: str) -> dict[str, Any] | None:
return data
def transitive_constraint(
name: str, version: str, target: str
) -> tuple[str | None, list[str]]:
def transitive_constraint(name: str, version: str, target: str) -> tuple[str | None, list[str]]:
"""Return (raw_specifier_string_or_None, list_of_(op,version) tuples)
for the constraint that `name==version` places on `target`.
"""
@ -501,10 +493,7 @@ def resolved_set(install_cell: str, colab: dict[str, str]) -> dict[str, str]:
out[sp.name] = ver
pinned.add(sp.name)
elif op == "<=" and sp.name not in pinned:
if (
sp.name not in upper_bounds
or cmp_versions(ver, upper_bounds[sp.name]) < 0
):
if sp.name not in upper_bounds or cmp_versions(ver, upper_bounds[sp.name]) < 0:
upper_bounds[sp.name] = ver
# Apply upper bounds where Colab's preinstall violates them.
for name, ub in upper_bounds.items():
@ -519,9 +508,7 @@ def resolved_set(install_cell: str, colab: dict[str, str]) -> dict[str, str]:
# ----- Rules ----- #
def rule_inst_001_git_plus(
install_cell: str, file: str, cell_idx: int
) -> list[Finding]:
def rule_inst_001_git_plus(install_cell: str, file: str, cell_idx: int) -> list[Finding]:
findings: list[Finding] = []
for inv in iter_pip_invocations(install_cell):
if any("git+" in p for p in inv.packages) or "git+" in inv.raw:
@ -714,9 +701,7 @@ def rule_inst_005_transformers_tokenizers(
_RE_DOUBLE_BANG = re.compile(r"^[ \t]*!{2,}\s*pip\b", re.MULTILINE)
def rule_inst_006_double_bang(
install_cell: str, file: str, cell_idx: int
) -> list[Finding]:
def rule_inst_006_double_bang(install_cell: str, file: str, cell_idx: int) -> list[Finding]:
findings: list[Finding] = []
for m in _RE_DOUBLE_BANG.finditer(install_cell):
line_no = install_cell.count("\n", 0, m.start()) + 1
@ -813,9 +798,7 @@ POLICY_CLAUSES_DEFAULT = [
]
def extract_policy_clauses(
update_script: pathlib.Path,
) -> list[tuple[str, re.Pattern[str], Any]]:
def extract_policy_clauses(update_script: pathlib.Path) -> list[tuple[str, re.Pattern[str], Any]]:
"""Best-effort: scan update_all_notebooks.py for canonical phrases used by
multiple templates. Falls back to POLICY_CLAUSES_DEFAULT.
@ -879,11 +862,7 @@ def cmd_drift(args: argparse.Namespace) -> int:
print(f"FAIL: {update_script} not found", file = sys.stderr)
return 2
# Stash any pre-existing dirty state, run the updater, diff, restore.
head = (
subprocess.check_output(["git", "rev-parse", "HEAD"], cwd = nbdir)
.decode()
.strip()
)
head = subprocess.check_output(["git", "rev-parse", "HEAD"], cwd = nbdir).decode().strip()
subprocess.run(
["git", "-C", str(nbdir), "stash", "--include-untracked"],
check = False,
@ -990,9 +969,7 @@ def cmd_convert(args: argparse.Namespace) -> int:
hint = proc.stderr[-200:].strip(),
)
)
print(
f"converted {len(notebooks) - len(failed)}/{len(notebooks)} notebooks to {out}"
)
print(f"converted {len(notebooks) - len(failed)}/{len(notebooks)} notebooks to {out}")
_emit(failed)
return 0 if not failed else 1
@ -1002,11 +979,7 @@ def cmd_convert(args: argparse.Namespace) -> int:
def cmd_lint(args: argparse.Namespace) -> int:
nbdir = pathlib.Path(args.notebooks_dir).resolve()
colab_path = (
pathlib.Path(args.colab_pin).resolve()
if args.colab_pin
else COLAB_FALLBACK_FILE
)
colab_path = pathlib.Path(args.colab_pin).resolve() if args.colab_pin else COLAB_FALLBACK_FILE
colab = parse_pip_freeze(colab_path)
if not colab:
print(
@ -1049,13 +1022,9 @@ def cmd_lint(args: argparse.Namespace) -> int:
first_cell = cells[0][0] if cells else None
findings += rule_inst_003_peft_torchao(merged, oracle, rel, first_cell)
findings += rule_inst_004_torchcodec_torch(merged, oracle, rel, first_cell)
findings += rule_inst_005_transformers_tokenizers(
merged, oracle, rel, first_cell
)
findings += rule_inst_005_transformers_tokenizers(merged, oracle, rel, first_cell)
if not args.no_pypi:
findings += rule_inst_002_no_deps_transitive(
merged, oracle, rel, first_cell
)
findings += rule_inst_002_no_deps_transitive(merged, oracle, rel, first_cell)
findings += scan_user_cells(nb, rel)
_emit(findings)
return 0 if not any(f.severity == "error" for f in findings) else 1
@ -1232,9 +1201,7 @@ def cmd_colab_diff(args: argparse.Namespace) -> int:
print(f"::warning::colab-diff: could not fetch {url}: {e}")
continue
if not snap_path.exists():
print(
f"::warning::colab-diff: no committed snapshot at {snap_path}; skipping"
)
print(f"::warning::colab-diff: no committed snapshot at {snap_path}; skipping")
continue
snapshot_text = snap_path.read_text(encoding = "utf-8", errors = "replace")
parser = _COLAB_ORACLE_PARSERS[upstream_name]

16
scripts/run_ruff_format.py
View file

@ -1,5 +1,7 @@
#!/usr/bin/env python3
"""Run `ruff format` followed by kwarg spacing enforcement."""
"""Run a pre-pass (normalize def-signature magic commas + collapse short
multi-line asserts), then `ruff format`, then the kwarg-spacing / import /
string-merge post-pass."""
from __future__ import annotations
@ -15,12 +17,22 @@ def main(argv: list[str]) -> int:
if not files:
return 0
spacing_script = HERE / "enforce_kwargs_spacing.py"
# Pre-ruff: normalize def-signature trailing commas (>=3 params with a
# default -> one-per-line; everything else collapsible) and strip the magic
# comma from a short multi-line assert, so ruff wraps signatures accordingly
# and joins the assert back onto one line.
pre_cmd = [sys.executable, str(spacing_script), "--pre", *files]
pre_proc = subprocess.run(pre_cmd)
if pre_proc.returncode != 0:
return pre_proc.returncode
ruff_cmd = [sys.executable, "-m", "ruff", "format", *files]
ruff_proc = subprocess.run(ruff_cmd)
if ruff_proc.returncode != 0:
return ruff_proc.returncode
spacing_script = HERE / "enforce_kwargs_spacing.py"
spacing_cmd = [sys.executable, str(spacing_script), *files]
spacing_proc = subprocess.run(spacing_cmd)
return spacing_proc.returncode

49
scripts/scan_npm_packages.py
View file

@ -823,8 +823,7 @@ def download_tarball(
written += len(chunk)
if written > max_bytes:
return dest, (
f"download exceeded cap {max_bytes} bytes "
f"after {written} bytes"
f"download exceeded cap {max_bytes} bytes " f"after {written} bytes"
)
h.update(chunk)
out.write(chunk)
@ -926,11 +925,7 @@ def safe_extract(
# get the generous binary cap. We bound BOTH cases.
header = src.read(16)
is_binary = _looks_binary(name, header)
file_cap = (
HARD_MAX_BINARY_FILE_BYTES
if is_binary
else HARD_MAX_TEXT_FILE_BYTES
)
file_cap = HARD_MAX_BINARY_FILE_BYTES if is_binary else HARD_MAX_TEXT_FILE_BYTES
if declared > file_cap:
return (
f"member {name!r} declared size {declared} > "
@ -963,7 +958,11 @@ def safe_extract(
# ─────────────────────────────────────────────────────────────────────
def _evidence(text: str, pat: re.Pattern, max_chars: int = 200) -> str:
def _evidence(
text: str,
pat: re.Pattern,
max_chars: int = 200,
) -> str:
m = pat.search(text)
if not m:
return ""
@ -978,11 +977,7 @@ def _evidence(text: str, pat: re.Pattern, max_chars: int = 200) -> str:
LIFECYCLE_HOOKS = ("preinstall", "install", "postinstall", "prepare")
def scan_package_json(
pkg: PackageEntry,
rel: str,
text: str,
) -> list[Finding]:
def scan_package_json(pkg: PackageEntry, rel: str, text: str) -> list[Finding]:
findings: list[Finding] = []
try:
meta = json.loads(text)
@ -1056,9 +1051,7 @@ def scan_package_json(
if isinstance(opt, dict):
for k, v in opt.items():
if isinstance(v, str) and (
v.startswith("github:")
or v.startswith("git+")
or v.startswith("git://")
v.startswith("github:") or v.startswith("git+") or v.startswith("git://")
):
findings.append(
Finding(
@ -1117,11 +1110,7 @@ def _host_in_outbound_context(text: str, host: str) -> bool:
return False
def scan_text_blob(
pkg: PackageEntry,
rel: str,
text: str,
) -> list[Finding]:
def scan_text_blob(pkg: PackageEntry, rel: str, text: str) -> list[Finding]:
findings: list[Finding] = []
# IOC substrings (literal, case-sensitive).
@ -1190,10 +1179,7 @@ def scan_text_blob(
filename = rel,
pattern = "js-fetch-eval",
evidence = _evidence(text, _JS_FETCH_EVAL),
detail = (
"Function/eval against base64-decoded payload "
"(obfuscated dropper shape)"
),
detail = ("Function/eval against base64-decoded payload (obfuscated dropper shape)"),
)
)
if _JS_ENV_TOKEN.search(text):
@ -1247,10 +1233,7 @@ _TEXT_SUFFIXES = (
)
def scan_extracted_tree(
pkg: PackageEntry,
root: Path,
) -> list[Finding]:
def scan_extracted_tree(pkg: PackageEntry, root: Path) -> list[Finding]:
findings: list[Finding] = []
for path in sorted(root.rglob("*")):
if not path.is_file():
@ -1304,10 +1287,7 @@ def scan_extracted_tree(
# ─────────────────────────────────────────────────────────────────────
def scan_one(
pkg: PackageEntry,
workspace: Path,
) -> tuple[list[Finding], str | None]:
def scan_one(pkg: PackageEntry, workspace: Path) -> tuple[list[Finding], str | None]:
"""Download + extract + scan a single package. Cleans up its dir.
Returns (findings, error). `error` is non-None only on hard
@ -1444,8 +1424,7 @@ def main(argv: list[str] | None = None) -> int:
if hard_errors or blocking:
if blocking:
print(
f"\n[scan-npm] FAIL: {len(blocking)} finding(s) "
f"at or above {threshold}",
f"\n[scan-npm] FAIL: {len(blocking)} finding(s) " f"at or above {threshold}",
file = sys.stderr,
)
return 1

64
scripts/scan_packages.py
View file

@ -86,8 +86,7 @@ RE_SUBPROCESS = re.compile(
# Encoding / obfuscation
RE_BASE64 = re.compile(
r"\bbase64\s*\.\s*(b64decode|decodebytes|b32decode|b16decode)\b"
r"|\bcodecs\s*\.\s*decode\b",
r"\bbase64\s*\.\s*(b64decode|decodebytes|b32decode|b16decode)\b|\bcodecs\s*\.\s*decode\b",
)
# exec / eval
@ -299,9 +298,7 @@ RE_CRYPTO_THEFT = re.compile(
RE_PTH_IMPORT = re.compile(r"^\s*import\s+", re.MULTILINE)
# openssl CLI invocations via subprocess (encrypted exfiltration)
RE_OPENSSL_CLI = re.compile(
r"\bopenssl\s+(enc|rand|rsautl|pkeyutl|genrsa|dgst|s_client)\b"
)
RE_OPENSSL_CLI = re.compile(r"\bopenssl\s+(enc|rand|rsautl|pkeyutl|genrsa|dgst|s_client)\b")
# Write to /tmp then execute (staged dropper)
RE_TEMP_EXEC = re.compile(
@ -962,7 +959,11 @@ def check_py_file(content: str, filename: str, package: str) -> list[Finding]:
return findings
def _extract_evidence(content: str, pattern: re.Pattern, max_matches: int = 3) -> str:
def _extract_evidence(
content: str,
pattern: re.Pattern,
max_matches: int = 3,
) -> str:
"""Pull matching lines as evidence snippets."""
lines = content.splitlines()
matches = []
@ -1266,15 +1267,13 @@ def iter_archive_files(archive_path: str):
# have historically dereferenced them on extract.
if member.issym() or member.islnk():
print(
f" [WARN] {path.name}: refused link member "
f"{member.name!r}",
f" [WARN] {path.name}: refused link member " f"{member.name!r}",
file = sys.stderr,
)
continue
if member.isdev() or member.isfifo():
print(
f" [WARN] {path.name}: refused special member "
f"{member.name!r}",
f" [WARN] {path.name}: refused special member " f"{member.name!r}",
file = sys.stderr,
)
continue
@ -1379,9 +1378,7 @@ def scan_archive(archive_path: str, package: str) -> list[Finding]:
_RE_PYPI_SPEC_VERSION = re.compile(r"==\s*([A-Za-z0-9_.\-+!]+)")
def _check_blocked_pypi_versions(
specs: list[str],
) -> tuple[list[str], list[Finding]]:
def _check_blocked_pypi_versions(specs: list[str]) -> tuple[list[str], list[Finding]]:
"""Filter ``specs`` against ``BLOCKED_PYPI_VERSIONS``.
Returns ``(safe_specs, findings)``. Each blocked spec emits a CRITICAL
@ -1502,9 +1499,7 @@ def download_packages(
env = env,
)
if proc.returncode != 0:
msg = (
f"pip download (with deps) failed: " f"{proc.stderr.strip()[:500]}"
)
msg = f"pip download (with deps) failed: " f"{proc.stderr.strip()[:500]}"
print(f" [ERROR] {msg}", file = sys.stderr)
download_errors.append(msg)
except subprocess.TimeoutExpired:
@ -1547,10 +1542,7 @@ def download_packages(
env = env,
)
if proc.returncode != 0:
msg = (
f"pip download failed for {spec}: "
f"{proc.stderr.strip()[:500]}"
)
msg = f"pip download failed for {spec}: " f"{proc.stderr.strip()[:500]}"
print(f" [ERROR] {msg}", file = sys.stderr)
download_errors.append(msg)
continue
@ -1579,9 +1571,7 @@ def _extract_pkg_name(spec: str) -> str:
"""Extract the package name from a pip spec string."""
m = _RE_NAME.match(spec)
return (
m.group(1)
if m
else spec.split("==")[0].split(">=")[0].split("<=")[0].split("[")[0].strip()
m.group(1) if m else spec.split("==")[0].split(">=")[0].split("<=")[0].split("[")[0].strip()
)
@ -1917,11 +1907,7 @@ def update_req_file(filepath: str, updates: dict[int, str]) -> None:
raise
def _run_fix(
critical_pkgs: set[str],
entries: list[dict],
max_search: int,
) -> None:
def _run_fix(critical_pkgs: set[str], entries: list[dict], max_search: int) -> None:
"""Run the --fix flow: find safe versions, update requirements files."""
# Map package names to their entries for source tracking
pkg_entries: dict[str, list[dict]] = {}
@ -1941,9 +1927,7 @@ def _run_fix(
if git_entries:
for e in git_entries:
src = e["source_file"] or "CLI"
print(
f" [SKIP] {pkg_name} is a git URL dep in {src}, cannot auto-update"
)
print(f" [SKIP] {pkg_name} is a git URL dep in {src}, cannot auto-update")
changes_summary.append(f" SKIP {pkg_name} (git URL)")
continue
@ -1967,9 +1951,7 @@ def _run_fix(
shutil.rmtree(dl_dir, ignore_errors = True)
if not current_ver:
print(
f" [WARN] Cannot determine current version of {pkg_name}, skipping fix"
)
print(f" [WARN] Cannot determine current version of {pkg_name}, skipping fix")
changes_summary.append(f" SKIP {pkg_name} (version unknown)")
continue
@ -1986,9 +1968,7 @@ def _run_fix(
continue
print(f" [OK] {pkg_name}: {current_ver} -> {safe_ver}")
changes_summary.append(
f" FIX {pkg_name}=={current_ver} -> {pkg_name}=={safe_ver}"
)
changes_summary.append(f" FIX {pkg_name}=={current_ver} -> {pkg_name}=={safe_ver}")
# Update all occurrences in requirements files
file_updates: dict[str, dict[int, str]] = {}
@ -2038,9 +2018,7 @@ def _find_requirements_files(root: str) -> list[str]:
dirnames[:] = [
d
for d in dirnames
if not d.startswith(".")
and d not in skip_dirs
and not d.endswith(".egg-info")
if not d.startswith(".") and d not in skip_dirs and not d.endswith(".egg-info")
]
dirname = os.path.basename(dirpath)
for fname in sorted(filenames):
@ -2115,9 +2093,7 @@ def main() -> int:
print(f" {f}")
req_files.extend(found)
else:
print(
f" [WARN] No requirements files found in {scan_dir}/", file = sys.stderr
)
print(f" [WARN] No requirements files found in {scan_dir}/", file = sys.stderr)
# Build unified entry list: list of dicts with source tracking
entries: list[dict] = []
@ -2211,7 +2187,7 @@ def main() -> int:
for err in download_errors:
print(f" [ERROR] {err}", file = sys.stderr)
print(
" Refusing to report 'all clean' on a partial scan; " "exiting 2.",
" Refusing to report 'all clean' on a partial scan; exiting 2.",
file = sys.stderr,
)
return 2

10
scripts/stamp_studio_release.py
View file

@ -17,7 +17,11 @@ import zipfile
from pathlib import Path
def _atomic_write_text(path: Path, data: str, encoding: str = "utf-8") -> None:
def _atomic_write_text(
path: Path,
data: str,
encoding: str = "utf-8",
) -> None:
"""Atomic version of ``Path.write_text``.
A crash or signal mid-write leaves the prior file intact; the
@ -41,9 +45,7 @@ def _atomic_write_text(path: Path, data: str, encoding: str = "utf-8") -> None:
REPO_ROOT = Path(__file__).resolve().parents[1]
BUILD_INFO_PATH = (
REPO_ROOT / "studio" / "backend" / "utils" / "_studio_release_build.py"
)
BUILD_INFO_PATH = REPO_ROOT / "studio" / "backend" / "utils" / "_studio_release_build.py"
BUILD_INFO_SUFFIX = "studio/backend/utils/_studio_release_build.py"
VERSION_RE = re.compile(r"^v\d+\.\d+\.\d+(?:-[0-9A-Za-z.][0-9A-Za-z.-]*)?$")
GIT_DESCRIBE_SUFFIX_RE = re.compile(r"-\d+-g[0-9A-Fa-f]+(?:-dirty)?$")

9
scripts/verify_comment_only_diff.py
View file

@ -128,12 +128,15 @@ def _normalize_yaml_run_strings(obj: Any) -> Any:
return obj
def _walk_yaml_diff(b: Any, a: Any, prefix: str = "") -> None:
def _walk_yaml_diff(
b: Any,
a: Any,
prefix: str = "",
) -> None:
"""Print a path-keyed summary of the first structural / scalar diff."""
if type(b) is not type(a):
print(
f" type-diff at {prefix or '/'}: "
f"{type(b).__name__} -> {type(a).__name__}",
f" type-diff at {prefix or '/'}: " f"{type(b).__name__} -> {type(a).__name__}",
)
return
if isinstance(b, dict):

76
scripts/verify_import_hoist.py
View file

@ -123,9 +123,7 @@ class _Builder(ast.NodeVisitor):
def __init__(self):
self.module = Scope("module", "<module>", None)
self.uses: list[tuple[Scope, str, int]] = [] # (scope, name, lineno) hard loads
self.soft_uses: list[
tuple[Scope, str, int]
] = [] # annotations: count as "used"
self.soft_uses: list[tuple[Scope, str, int]] = [] # annotations: count as "used"
# but never as "unresolved"
# (forward refs / string annos)
@ -166,9 +164,7 @@ class _Builder(ast.NodeVisitor):
def _visit_stmt(self, node: ast.AST, scope: Scope) -> None:
if isinstance(node, (ast.Import, ast.ImportFrom)):
star = isinstance(node, ast.ImportFrom) and any(
a.name == "*" for a in node.names
)
star = isinstance(node, ast.ImportFrom) and any(a.name == "*" for a in node.names)
if star:
scope.star_import = True
for alias in node.names:
@ -356,9 +352,7 @@ class _Builder(ast.NodeVisitor):
self._bind_args(node.args, child)
self._visit_expr(node.body, child)
return
if isinstance(
node, (ast.ListComp, ast.SetComp, ast.GeneratorExp, ast.DictComp)
):
if isinstance(node, (ast.ListComp, ast.SetComp, ast.GeneratorExp, ast.DictComp)):
child = Scope("comp", f"{scope.qualname}.<comp>", scope)
for i, gen in enumerate(node.generators):
# first iterable is evaluated in the enclosing scope
@ -447,9 +441,7 @@ def _legb_chain(scope: Scope) -> list[Scope]:
chain = [scope]
p = scope.parent
while p is not None:
if (
p.kind != "class" or p.parent is None
): # module-level class never happens; keep module
if p.kind != "class" or p.parent is None: # module-level class never happens; keep module
if p.kind != "class":
chain.append(p)
p = p.parent
@ -624,9 +616,7 @@ def compare(before_src: str, after_src: str, path: str) -> list[tuple[str, str]]
for scope, names in b["ambiguous"].items():
new = names - a["ambiguous"].get(scope, set())
for n in sorted(new):
findings.append(
("WARN", f"{path}: AMBIGUOUS-BIND '{n}' import+non-import in {scope}")
)
findings.append(("WARN", f"{path}: AMBIGUOUS-BIND '{n}' import+non-import in {scope}"))
# 6. TARGET-MISSING (informational): a scope stopped resolving to an import
# target. Real bugs are already covered above; remaining cases are code
@ -639,9 +629,7 @@ def compare(before_src: str, after_src: str, path: str) -> list[tuple[str, str]]
if t in added_module_targets
else " [target not re-added here -> likely relocated/deleted]"
)
findings.append(
("INFO", f"{path}: TARGET-MISSING {t} in scope {scope}{relocated}")
)
findings.append(("INFO", f"{path}: TARGET-MISSING {t} in scope {scope}{relocated}"))
return findings
@ -650,45 +638,38 @@ def compare(before_src: str, after_src: str, path: str) -> list[tuple[str, str]]
_SELF_TESTS = {
"dangling_alias": (
# before: inline aliased import, used as _b
"import os\n"
"def f():\n"
" import glob as _b\n"
" return _b.glob('*')\n",
"import os\ndef f():\n import glob as _b\n return _b.glob('*')\n",
# after: hoisted to canonical, but reference NOT normalized -> _b dangles
"import os\n" "import glob\n" "def f():\n" " return _b.glob('*')\n",
"import os\nimport glob\ndef f():\n return _b.glob('*')\n",
"BLOCKER",
),
"rename_clash": (
# before: _b is a deliberate alias; `b` already means something else
"import re as _b\n" "b = 123\n" "def f():\n" " return _b.compile('x'), b\n",
"import re as _b\nb = 123\ndef f():\n return _b.compile('x'), b\n",
# after: someone normalized _b -> b ; now f().b is the int, re is lost
"import re\n" "b = 123\n" "def f():\n" " return b.compile('x'), b\n",
"import re\nb = 123\ndef f():\n return b.compile('x'), b\n",
"BLOCKER", # TARGET-MISSING from:.. or import:re in f
),
"clean_rename": (
"def f():\n" " import glob as _g\n" " return _g.glob('*')\n",
"import glob\n" "def f():\n" " return glob.glob('*')\n",
"def f():\n import glob as _g\n return _g.glob('*')\n",
"import glob\ndef f():\n return glob.glob('*')\n",
None, # expect NO blocker
),
"clean_dedup_redundant": (
"import sys\n" "def f():\n" " import sys\n" " return sys.argv\n",
"import sys\n" "def f():\n" " return sys.argv\n",
"import sys\ndef f():\n import sys\n return sys.argv\n",
"import sys\ndef f():\n return sys.argv\n",
None,
),
"from_import_dangling": (
# from-import alias left un-normalized
"def f():\n"
" from importlib.metadata import version as _v\n"
" return _v('x')\n",
"from importlib.metadata import version\n" "def f():\n" " return _v('x')\n",
"def f():\n from importlib.metadata import version as _v\n return _v('x')\n",
"from importlib.metadata import version\ndef f():\n return _v('x')\n",
"BLOCKER",
),
"local_var_clash": (
# _b renamed to b, but b is a LOCAL variable in f -> import silently unused
"def f(b):\n" " import re as _b\n" " return _b.compile(b)\n",
"import re\n"
"def f(b):\n"
" return b.compile(b)\n", # 'b' is the param, not the module
"def f(b):\n import re as _b\n return _b.compile(b)\n",
"import re\ndef f(b):\n return b.compile(b)\n", # 'b' is the param, not the module
"BLOCKER",
),
"substring_safe": (
@ -705,11 +686,8 @@ _SELF_TESTS = {
),
"attr_access_not_a_use": (
# x._b is attribute access, not a use of name _b; removing import _b is fine
"import os\n"
"def f(x):\n"
" import sys as _b\n"
" return x._b + _b.argv[0]\n",
"import os\n" "import sys\n" "def f(x):\n" " return x._b + sys.argv[0]\n",
"import os\ndef f(x):\n import sys as _b\n return x._b + _b.argv[0]\n",
"import os\nimport sys\ndef f(x):\n return x._b + sys.argv[0]\n",
None,
),
}
@ -797,9 +775,7 @@ def audit_files(paths: list[str]) -> int:
ok = n_err == 0 and n_fp == 0
print(
"\nAUDIT:",
"ROBUST (no crashes, no false positives vs pyflakes)"
if ok
else "NEEDS WORK (see above)",
"ROBUST (no crashes, no false positives vs pyflakes)" if ok else "NEEDS WORK (see above)",
)
return 0 if ok else 1
@ -835,18 +811,12 @@ def main() -> int:
blockers = [f for f in findings if f[0] == "BLOCKER"]
warns = [f for f in findings if f[0] == "WARN"]
infos = [f for f in findings if f[0] == "INFO"]
status = (
"CLEAN"
if not blockers and not warns
else ("BLOCKERS" if blockers else "WARNINGS")
)
status = "CLEAN" if not blockers and not warns else ("BLOCKERS" if blockers else "WARNINGS")
print(f"\n=== {path}: {status} ===")
for sev, m in blockers + warns + infos:
print(f" [{sev}] {m}")
any_blocker = any_blocker or bool(blockers)
print(
"\nOVERALL:", "FAIL (blockers found)" if any_blocker else "PASS (no blockers)"
)
print("\nOVERALL:", "FAIL (blockers found)" if any_blocker else "PASS (no blockers)")
return 1 if any_blocker else 0

12
studio/backend/auth/authentication.py
View file

@ -108,9 +108,7 @@ def create_refresh_token(subject: str, *, desktop: bool = False) -> str:
return token
def refresh_access_token(
refresh_token: str,
) -> Tuple[Optional[str], Optional[str], bool]:
def refresh_access_token(refresh_token: str) -> Tuple[Optional[str], Optional[str], bool]:
"""
Validate a refresh token and issue a new access token.
@ -137,9 +135,7 @@ def reload_secret() -> None:
load_jwt_secret()
async def get_current_subject(
credentials: HTTPAuthorizationCredentials = Depends(security),
) -> str:
async def get_current_subject(credentials: HTTPAuthorizationCredentials = Depends(security)) -> str:
"""Validate JWT and require the password-change flow to be completed."""
return await _get_current_subject(
credentials,
@ -158,9 +154,7 @@ async def get_current_subject_allow_password_change(
async def _get_current_subject(
credentials: HTTPAuthorizationCredentials,
*,
allow_password_change: bool,
credentials: HTTPAuthorizationCredentials, *, allow_password_change: bool
) -> str:
"""
FastAPI dependency to validate the JWT and return the subject.

16
studio/backend/auth/storage.py
View file

@ -151,13 +151,9 @@ def get_connection() -> sqlite3.Connection:
);
"""
)
api_key_columns = {
row["name"] for row in conn.execute("PRAGMA table_info(api_keys)")
}
api_key_columns = {row["name"] for row in conn.execute("PRAGMA table_info(api_keys)")}
if "is_internal" not in api_key_columns:
conn.execute(
"ALTER TABLE api_keys ADD COLUMN is_internal INTEGER NOT NULL DEFAULT 0"
)
conn.execute("ALTER TABLE api_keys ADD COLUMN is_internal INTEGER NOT NULL DEFAULT 0")
conn.execute(
"""
CREATE TABLE IF NOT EXISTS app_secrets (
@ -171,13 +167,9 @@ def get_connection() -> sqlite3.Connection:
conn.execute(
"ALTER TABLE auth_user ADD COLUMN must_change_password INTEGER NOT NULL DEFAULT 0"
)
refresh_columns = {
row["name"] for row in conn.execute("PRAGMA table_info(refresh_tokens)")
}
refresh_columns = {row["name"] for row in conn.execute("PRAGMA table_info(refresh_tokens)")}
if "is_desktop" not in refresh_columns:
conn.execute(
"ALTER TABLE refresh_tokens ADD COLUMN is_desktop INTEGER NOT NULL DEFAULT 0"
)
conn.execute("ALTER TABLE refresh_tokens ADD COLUMN is_desktop INTEGER NOT NULL DEFAULT 0")
conn.commit()
return conn

25
studio/backend/colab.py
View file

@ -44,12 +44,7 @@ def get_colab_url(port: int = 8888) -> str:
try:
url = eval_js(f"google.colab.kernel.proxyPort({port})", timeout_sec = 10)
# A valid Colab proxy URL starts with https:// and embeds the port.
if (
url
and isinstance(url, str)
and url.startswith("https://")
and str(port) in url
):
if url and isinstance(url, str) and url.startswith("https://") and str(port) in url:
return url.rstrip("/")
except Exception as e:
logger.info(f"Note: Could not get Colab URL (attempt {attempt + 1}/3: {e})")
@ -118,11 +113,8 @@ def show_link(port: int = 8888, *, _url: "str | None" = None):
def _is_studio_healthy(port: int, timeout: float = 2.0) -> bool:
"""Return True if a Studio backend is already answering health checks on *port*."""
import urllib.request
try:
with urllib.request.urlopen(
f"http://localhost:{port}/api/health", timeout = timeout
):
with urllib.request.urlopen(f"http://localhost:{port}/api/health", timeout = timeout):
return True
except Exception:
return False
@ -178,7 +170,6 @@ def _show_and_embed(port: int):
# Fallback: Colab's built-in helper (less control, but always works)
try:
from google.colab import output as colab_output
colab_output.serve_kernel_port_as_iframe(port, height = 900, width = "100%")
except ImportError:
pass
@ -200,9 +191,7 @@ def start(port: int = 8888):
# Re-launching would either collide on the port or silently shift to a new
# port and confuse the user. Just re-show the link and iframe instead.
if _is_studio_healthy(port):
logger.info(
f" Studio is already running on port {port} — reusing existing server."
)
logger.info(f" Studio is already running on port {port} — reusing existing server.")
_show_and_embed(port)
try:
for _ in range(10000):
@ -225,9 +214,7 @@ def start(port: int = 8888):
logger.info(" Starting server...")
try:
app = run_server(
host = "0.0.0.0", port = port, frontend_path = frontend_path, silent = True
)
app = run_server(host = "0.0.0.0", port = port, frontend_path = frontend_path, silent = True)
except SystemExit as exc:
logger.error(f"❌ Unsloth Studio failed to start: {exc}")
return
@ -250,9 +237,7 @@ def start(port: int = 8888):
server_ready = False
for _ in range(40):
try:
with urllib.request.urlopen(
f"http://localhost:{actual_port}/api/health", timeout = 1
):
with urllib.request.urlopen(f"http://localhost:{actual_port}/api/health", timeout = 1):
server_ready = True
break
except Exception:

1
studio/backend/core/__init__.py
View file

@ -140,7 +140,6 @@ def __getattr__(name):
# Datasets
if name == "format_and_template_dataset":
from utils.datasets import format_and_template_dataset
globals()["format_and_template_dataset"] = format_and_template_dataset
return format_and_template_dataset

14
studio/backend/core/_torchao_stub.py
View file

@ -64,7 +64,11 @@ def _make_mod_stub(mod_name):
m._unsloth_stub = _STUB_SENTINEL
m.__spec__ = importlib.machinery.ModuleSpec(mod_name, loader = None, is_package = True)
def _ga(attr, _m = m, _n = mod_name):
def _ga(
attr,
_m = m,
_n = mod_name,
):
if attr.startswith("__"):
raise AttributeError(attr)
# Return a stub CLASS (not a module) so that isinstance(x, attr)
@ -89,7 +93,12 @@ class _StubSubpackageLoader(importlib.abc.Loader):
class _StubSubpackageFinder(importlib.abc.MetaPathFinder):
def find_spec(self, fullname, path, target = None):
def find_spec(
self,
fullname,
path,
target = None,
):
if "." not in fullname:
return None
parent = sys.modules.get(fullname.rsplit(".", 1)[0])
@ -118,7 +127,6 @@ def install_torchao_windows_rocm_stub() -> None:
if sys.platform == "win32":
try:
import torch as _torch_probe
_is_win32_rocm = bool(
getattr(getattr(_torch_probe, "version", None), "hip", None)
or "rocm" in getattr(_torch_probe, "__version__", "").lower()

4
studio/backend/core/data_recipe/huggingface.py
View file

@ -36,9 +36,7 @@ def _resolve_recipe_artifact_path(artifact_path: str) -> Path:
if not resolved.exists():
raise RecipeDatasetPublishError("Execution artifacts are no longer available.")
if not resolved.is_dir():
raise RecipeDatasetPublishError(
"Execution artifact path is not a dataset folder."
)
raise RecipeDatasetPublishError("Execution artifact path is not a dataset folder.")
return resolved

47
studio/backend/core/data_recipe/jobs/manager.py
View file

@ -108,9 +108,7 @@ class Subscription:
event_id = self._next_id
body = json.dumps(event, separators = (",", ":"), ensure_ascii = False)
event_type = event.get("type") or "message"
return (
f"id: {event_id}\n" f"event: {event_type}\n" f"data: {body}\n\n"
).encode("utf-8")
return (f"id: {event_id}\n" f"event: {event_type}\n" f"data: {body}\n\n").encode("utf-8")
class JobManager:
@ -158,9 +156,7 @@ class JobManager:
job_id = uuid.uuid4().hex
self._job = Job(job_id = job_id, status = "pending", started_at = time.time())
self._job.progress_columns_total = llm_column_count
self._job.source_progress_estimated_total = _github_source_estimated_total(
recipe
)
self._job.source_progress_estimated_total = _github_source_estimated_total(recipe)
self._job.internal_api_key_id = internal_api_key_id
self._events.clear()
self._seq = 0
@ -187,9 +183,7 @@ class JobManager:
self._pump_thread = threading.Thread(target = self._pump_loop, daemon = True)
self._pump_thread.start()
self._emit(
{"type": EVENT_JOB_ENQUEUED, "ts": time.time(), "job_id": job_id}
)
self._emit({"type": EVENT_JOB_ENQUEUED, "ts": time.time(), "job_id": job_id})
return job_id
def cancel(self, job_id: str) -> bool:
@ -200,9 +194,7 @@ class JobManager:
if self._proc is None or not self._proc.is_alive():
return True
self._job.status = "cancelling"
self._emit(
{"type": EVENT_JOB_CANCELLING, "ts": time.time(), "job_id": job_id}
)
self._emit({"type": EVENT_JOB_CANCELLING, "ts": time.time(), "job_id": job_id})
try:
self._proc.terminate()
except (AttributeError, OSError):
@ -319,19 +311,12 @@ class JobManager:
if not parquet_dir.exists():
return {"error": f"dataset path missing: {parquet_dir}"}
return self._load_dataset_page(
parquet_dir = parquet_dir, limit = limit, offset = offset
)
return self._load_dataset_page(parquet_dir = parquet_dir, limit = limit, offset = offset)
except Exception as exc:
return {"error": f"dataset load failed: {exc}"}
@staticmethod
def _load_dataset_page(
*,
parquet_dir: Path,
limit: int,
offset: int,
) -> dict[str, Any]:
def _load_dataset_page(*, parquet_dir: Path, limit: int, offset: int) -> dict[str, Any]:
dataset_page = JobManager._load_dataset_page_with_duckdb(
parquet_dir = parquet_dir,
limit = limit,
@ -347,10 +332,7 @@ class JobManager:
@staticmethod
def _load_dataset_page_with_duckdb(
*,
parquet_dir: Path,
limit: int,
offset: int,
*, parquet_dir: Path, limit: int, offset: int
) -> dict[str, Any] | None:
parquet_glob = str((parquet_dir / "*.parquet").resolve())
try:
@ -389,10 +371,7 @@ class JobManager:
@staticmethod
def _load_dataset_page_with_data_designer(
*,
parquet_dir: Path,
limit: int,
offset: int,
*, parquet_dir: Path, limit: int, offset: int
) -> dict[str, Any]:
from data_designer.config.utils.io_helpers import read_parquet_dataset
@ -402,7 +381,10 @@ class JobManager:
return {"dataset": to_preview_jsonable(rows), "total": total}
def subscribe(
self, job_id: str, *, after_seq: int | None = None
self,
job_id: str,
*,
after_seq: int | None = None,
) -> Subscription | None:
"""SSE subscribe: get replay buffer + live events stream."""
with self._lock:
@ -497,9 +479,7 @@ class JobManager:
self._job.error = self._job.error or "process exited"
self._job.finished_at = time.time()
event_type = (
EVENT_JOB_CANCELLED
if self._job.status == "cancelled"
else EVENT_JOB_ERROR
EVENT_JOB_CANCELLED if self._job.status == "cancelled" else EVENT_JOB_ERROR
)
self._emit(
{
@ -566,7 +546,6 @@ class JobManager:
return
try:
from auth import storage # deferred: avoids circular import
storage.revoke_internal_api_key(int(key_id))
except Exception:
pass

38
studio/backend/core/data_recipe/jobs/parse.py
View file

@ -119,8 +119,7 @@ def parse_log_message(msg: str) -> ParsedUpdate | None:
page_items = page_items,
rate_remaining = int(m.group("remaining")),
message = (
f"Scraping GitHub source: {repo} "
f"{resource} page {page} (+{page_items})"
f"Scraping GitHub source: {repo} " f"{resource} page {page} (+{page_items})"
),
),
)
@ -134,10 +133,7 @@ def parse_log_message(msg: str) -> ParsedUpdate | None:
source = "github",
status = "rate_limited",
retry_after_sec = seconds,
message = (
"Waiting for GitHub rate limit. "
"Studio will resume automatically."
),
message = ("Waiting for GitHub rate limit. Studio will resume automatically."),
),
)
@ -151,8 +147,7 @@ def parse_log_message(msg: str) -> ParsedUpdate | None:
status = "rate_limited",
retry_after_sec = seconds,
message = (
"Waiting for GitHub secondary rate limit. "
"Studio will resume automatically."
"Waiting for GitHub secondary rate limit. Studio will resume automatically."
),
),
)
@ -166,10 +161,7 @@ def parse_log_message(msg: str) -> ParsedUpdate | None:
source = "github",
status = "rate_limited",
retry_after_sec = seconds,
message = (
"Waiting for GitHub rate limit. "
"Studio will resume automatically."
),
message = ("Waiting for GitHub rate limit. Studio will resume automatically."),
),
)
@ -387,15 +379,13 @@ def _apply_source_progress(job: Job, progress: SourceProgress) -> None:
count_key = f"{progress.repo}:{progress.resource}"
if page_key not in job._source_seen_pages:
job._source_seen_pages.add(page_key)
job._source_counts[count_key] = int(
job._source_counts.get(count_key, 0)
) + int(page_items or 0)
job._source_counts[count_key] = int(job._source_counts.get(count_key, 0)) + int(
page_items or 0
)
fetched_items = sum(job._source_counts.values())
if fetched_items <= 0:
fetched_items = progress.fetched_items or (
previous.fetched_items if previous else None
)
fetched_items = progress.fetched_items or (previous.fetched_items if previous else None)
estimated_total = (
progress.estimated_total
@ -415,14 +405,10 @@ def _apply_source_progress(job: Job, progress: SourceProgress) -> None:
repo = progress.repo or (previous.repo if previous else None),
resource = progress.resource or (previous.resource if previous else None),
page = (
progress.page
if progress.page is not None
else (previous.page if previous else None)
progress.page if progress.page is not None else (previous.page if previous else None)
),
page_items = (
page_items
if page_items is not None
else (previous.page_items if previous else None)
page_items if page_items is not None else (previous.page_items if previous else None)
),
fetched_items = fetched_items,
estimated_total = estimated_total,
@ -453,9 +439,7 @@ def _compute_overall_progress(job: Job, column_progress: Progress) -> Progress:
if len(job._column_done) == 0:
done = current_done
else:
sum_done = sum(
max(0, min(value, total_rows)) for value in job._column_done.values()
)
sum_done = sum(max(0, min(value, total_rows)) for value in job._column_done.values())
done = int(sum_done / total_columns)
prev_done = int(job.progress.done or 0)

23
studio/backend/core/data_recipe/jobs/worker.py
View file

@ -60,9 +60,7 @@ def _slugify_run_name(value: str) -> str:
return slug[:80].strip("-")
def _build_dataset_name(
*, run_name: str | None, job_id: str, artifact_root: Path
) -> str:
def _build_dataset_name(*, run_name: str | None, job_id: str, artifact_root: Path) -> str:
fallback = f"recipe_{job_id}"
slug = _slugify_run_name(run_name or "")
base_name = f"recipe_{slug}" if slug else fallback
@ -74,21 +72,14 @@ def _build_dataset_name(
return candidate
def run_job_process(
*,
event_queue,
recipe: dict[str, Any],
run: dict[str, Any],
) -> None:
def run_job_process(*, event_queue, recipe: dict[str, Any], run: dict[str, Any]) -> None:
"""
Subprocess entrypoint.
Sends events to `event_queue`.
"""
import os
os.environ["PYTHONWARNINGS"] = (
"ignore" # Suppress warnings at C-level before imports
)
os.environ["PYTHONWARNINGS"] = "ignore" # Suppress warnings at C-level before imports
import warnings
from loggers.config import LogConfig
@ -172,14 +163,10 @@ def run_job_process(
}
)
else:
results = designer.create(
builder, num_records = rows, dataset_name = dataset_name
)
results = designer.create(builder, num_records = rows, dataset_name = dataset_name)
analysis = to_jsonable(results.load_analysis().model_dump(mode = "json"))
if merge_batches:
_merge_batches_to_single_parquet(
results.artifact_storage.base_dataset_path
)
_merge_batches_to_single_parquet(results.artifact_storage.base_dataset_path)
artifact_path = str(results.artifact_storage.base_dataset_path)
event_queue.put(
{

2
studio/backend/core/data_recipe/jsonable.py
View file

@ -23,7 +23,6 @@ def _pil_to_preview_payload(image: Any) -> dict[str, Any]:
def _open_pil_image_from_bytes(raw_bytes: bytes):
from PIL import Image # type: ignore
with Image.open(io.BytesIO(raw_bytes)) as image:
return image.copy()
@ -52,7 +51,6 @@ def _to_pil_from_hf_image_dict(value: Any) -> Any | None:
if isinstance(path_value, str) and path_value.strip():
try:
from PIL import Image # type: ignore
with Image.open(Path(path_value)) as image:
return image.copy()
except (OSError, ValueError, TypeError):

50
studio/backend/core/data_recipe/local_callable_validators.py
View file

@ -81,9 +81,7 @@ def split_oxc_local_callable_validators(
def register_oxc_local_callable_validators(
*,
builder,
specs: list[OxcLocalCallableValidatorSpec],
*, builder, specs: list[OxcLocalCallableValidatorSpec]
) -> None:
if not specs:
return
@ -114,10 +112,7 @@ def register_oxc_local_callable_validators(
)
def _parse_oxc_spec(
*,
column: dict[str, Any],
) -> OxcLocalCallableValidatorSpec | None:
def _parse_oxc_spec(*, column: dict[str, Any]) -> OxcLocalCallableValidatorSpec | None:
if str(column.get("column_type") or "").strip() != "validation":
return None
if str(column.get("validator_type") or "").strip() != "local_callable":
@ -138,11 +133,7 @@ def _parse_oxc_spec(
target_columns_raw = column.get("target_columns")
target_columns = (
[
value.strip()
for value in target_columns_raw
if isinstance(value, str) and value.strip()
]
[value.strip() for value in target_columns_raw if isinstance(value, str) and value.strip()]
if isinstance(target_columns_raw, list)
else []
)
@ -182,9 +173,7 @@ def _parse_oxc_validation_marker(fn_name: str) -> tuple[str, str, str]:
return "javascript", "syntax", "auto"
code_lang = parts[0] if parts[0] in _OXC_LANG_TO_NODE_LANG else "javascript"
mode = parts[1] if parts[1] in _OXC_VALIDATION_MODES else "syntax"
code_shape = (
parts[2] if len(parts) >= 3 and parts[2] in _OXC_CODE_SHAPES else "auto"
)
code_shape = parts[2] if len(parts) >= 3 and parts[2] in _OXC_CODE_SHAPES else "auto"
return code_lang, mode, code_shape
@ -205,10 +194,7 @@ def _build_oxc_validation_function(lang: str, validation_mode: str, code_shape:
code_values = (
["" for _ in range(row_count)]
if not code_column
else [
"" if value is None else str(value)
for value in df[code_column].tolist()
]
else ["" if value is None else str(value) for value in df[code_column].tolist()]
)
results = _run_oxc_batch(
@ -224,16 +210,14 @@ def _build_oxc_validation_function(lang: str, validation_mode: str, code_shape:
)
return pd.DataFrame(results)
_validator.__name__ = f"{OXC_VALIDATION_FN_MARKER}_{node_lang}_{mode.replace('+', '_')}_{normalized_code_shape}"
_validator.__name__ = (
f"{OXC_VALIDATION_FN_MARKER}_{node_lang}_{mode.replace('+', '_')}_{normalized_code_shape}"
)
return _validator
def _run_oxc_batch(
*,
node_lang: str,
validation_mode: str,
code_shape: str,
code_values: list[str],
*, node_lang: str, validation_mode: str, code_shape: str, code_values: list[str]
) -> list[dict[str, Any]]:
if not _OXC_RUNNER_PATH.exists():
return _fallback_results(
@ -308,21 +292,13 @@ def _run_oxc_batch(
warning_count_raw = item.get("warning_count")
out.append(
{
"is_valid": bool(is_valid_raw)
if isinstance(is_valid_raw, bool)
else False,
"error_count": int(error_count_raw)
if isinstance(error_count_raw, int)
else 0,
"is_valid": bool(is_valid_raw) if isinstance(is_valid_raw, bool) else False,
"error_count": int(error_count_raw) if isinstance(error_count_raw, int) else 0,
"error_message": str(message_raw or ""),
"severity": str(severity_raw)
if isinstance(severity_raw, str)
else None,
"severity": str(severity_raw) if isinstance(severity_raw, str) else None,
"code": str(code_raw) if isinstance(code_raw, str) else None,
"labels": labels_raw if isinstance(labels_raw, list) else [],
"codeframe": str(codeframe_raw)
if isinstance(codeframe_raw, str)
else None,
"codeframe": str(codeframe_raw) if isinstance(codeframe_raw, str) else None,
"warning_count": int(warning_count_raw)
if isinstance(warning_count_raw, int)
else 0,

39
studio/backend/core/data_recipe/service.py
View file

@ -22,9 +22,7 @@ def _encode_bytes_to_base64(value: bytes | bytearray) -> str:
return base64.b64encode(bytes(value)).decode("utf-8")
def _load_image_file_to_base64(
path_value: str, *, base_path: str | None = None
) -> str | None:
def _load_image_file_to_base64(path_value: str, *, base_path: str | None = None) -> str | None:
try:
path = Path(path_value)
candidates: list[Path] = []
@ -119,9 +117,7 @@ def _apply_data_designer_image_context_patch() -> None:
original_auto_resolve = ImageContext._auto_resolve_context_value
def _patched_auto_resolve(
self: Any, context_value: Any, base_path: str | None
) -> Any:
def _patched_auto_resolve(self: Any, context_value: Any, base_path: str | None) -> Any:
normalized = _normalize_image_context_value(context_value, base_path = base_path)
return original_auto_resolve(self, normalized, base_path)
@ -163,17 +159,12 @@ def _recipe_has_llm_columns(recipe: dict[str, Any]) -> bool:
return False
def _validate_recipe_runtime_support(
recipe: dict[str, Any],
model_providers: list[Any],
) -> None:
def _validate_recipe_runtime_support(recipe: dict[str, Any], model_providers: list[Any]) -> None:
if _recipe_has_llm_columns(recipe) and not model_providers:
raise ValueError("Add a Provider connection block before running this recipe.")
def build_mcp_providers(
recipe: dict[str, Any],
) -> list:
def build_mcp_providers(recipe: dict[str, Any]) -> list:
from data_designer.config.mcp import LocalStdioMCPProvider, MCPProvider # pyright: ignore[reportMissingImports]
# Same gate as the chat MCP path: stdio providers spawn a local subprocess,
@ -259,9 +250,7 @@ def build_config_builder(recipe: dict[str, Any]):
if key not in {"model_providers", "mcp_providers"}
}
recipe_core = _strip_frontend_model_config_metadata(recipe_core)
recipe_core, oxc_local_callable_specs = split_oxc_local_callable_validators(
recipe_core
)
recipe_core, oxc_local_callable_specs = split_oxc_local_callable_validators(recipe_core)
builder = DataDesignerConfigBuilder.from_config({"data_designer": recipe_core})
register_oxc_local_callable_validators(
builder = builder,
@ -285,11 +274,7 @@ def build_config_builder(recipe: dict[str, Any]):
return builder
def create_data_designer(
recipe: dict[str, Any],
*,
artifact_path: str | None = None,
):
def create_data_designer(recipe: dict[str, Any], *, artifact_path: str | None = None):
_apply_data_designer_image_context_patch()
from data_designer.interface.data_designer import DataDesigner # pyright: ignore[reportMissingImports]
@ -302,7 +287,6 @@ def create_data_designer(
# so sampler/expression-only recipes can run without a real provider.
if not model_providers:
from data_designer.config.models import ModelProvider # pyright: ignore[reportMissingImports]
model_providers = [
ModelProvider(
name = "_unused",
@ -326,8 +310,7 @@ def validate_recipe(recipe: dict[str, Any]) -> None:
def preview_recipe(
recipe: dict[str, Any],
num_records: int,
recipe: dict[str, Any], num_records: int
) -> tuple[list[dict[str, Any]], dict[str, Any] | None, dict[str, Any] | None]:
builder = build_config_builder(recipe)
designer = create_data_designer(recipe)
@ -339,14 +322,10 @@ def preview_recipe(
dataset = [to_jsonable(row) for row in raw_rows]
artifacts = (
None
if results.processor_artifacts is None
else to_jsonable(results.processor_artifacts)
None if results.processor_artifacts is None else to_jsonable(results.processor_artifacts)
)
analysis = (
None
if results.analysis is None
else to_jsonable(results.analysis.model_dump(mode = "json"))
None if results.analysis is None else to_jsonable(results.analysis.model_dump(mode = "json"))
)
return dataset, artifacts, analysis

53
studio/backend/core/export/export.py
View file

@ -58,16 +58,11 @@ def _apply_wsl_sudo_patch():
import unsloth_zoo.llama_cpp as llama_cpp_module
def _wsl_do_we_need_sudo(system_type = "debian"):
logger.info(
"WSL detected — skipping sudo check "
"(build deps pre-installed by setup.sh)"
)
logger.info("WSL detected — skipping sudo check (build deps pre-installed by setup.sh)")
return False
llama_cpp_module.do_we_need_sudo = _wsl_do_we_need_sudo
logger.info(
"Applied WSL sudo patch to " "unsloth_zoo.llama_cpp.do_we_need_sudo"
)
logger.info("Applied WSL sudo patch to unsloth_zoo.llama_cpp.do_we_need_sudo")
except Exception as e:
logger.warning(f"Could not apply WSL sudo patch: {e}")
@ -146,7 +141,6 @@ class ExportBackend:
List of tuples: [(model_name, [(display_name, checkpoint_path), ...]), ...]
"""
from utils.models.checkpoints import scan_checkpoints
return scan_checkpoints(outputs_dir = outputs_dir)
def load_checkpoint(
@ -224,7 +218,6 @@ class ExportBackend:
elif self._audio_type == "bicodec":
from unsloth import FastModel
logger.info("Loading as BiCodec (Spark-TTS) audio model...")
model, tokenizer = FastModel.from_pretrained(
model_name = checkpoint_path,
@ -236,7 +229,6 @@ class ExportBackend:
elif self._audio_type == "dac":
from unsloth import FastModel
logger.info("Loading as DAC (OuteTTS) audio model...")
model, tokenizer = FastModel.from_pretrained(
model_name = checkpoint_path,
@ -348,9 +340,7 @@ class ExportBackend:
output_path: Optional[str] = None
try:
if _IS_MLX:
mlx_save_method = (
"merged_4bit" if format_type == "4-bit (FP4)" else "merged_16bit"
)
mlx_save_method = "merged_4bit" if format_type == "4-bit (FP4)" else "merged_16bit"
else:
if format_type == "4-bit (FP4)":
save_method = "merged_4bit_forced"
@ -415,9 +405,7 @@ class ExportBackend:
private = private,
)
else:
hub_save_method = (
save_method if save_method is not None else "merged_16bit"
)
hub_save_method = save_method if save_method is not None else "merged_16bit"
self.current_model.push_to_hub_merged(
repo_id,
self.current_tokenizer,
@ -523,9 +511,7 @@ class ExportBackend:
else:
# Get base model name from request or model config
base_model = (
base_model_id
or self.current_model.config._name_or_path
or "unknown"
base_model_id or self.current_model.config._name_or_path or "unknown"
)
# Create repo
@ -547,9 +533,7 @@ class ExportBackend:
extra = "unsloth",
)
card = ModelCard(content)
card.push_to_hub(
repo_id, token = hf_token, commit_message = "Unsloth Model Card"
)
card.push_to_hub(repo_id, token = hf_token, commit_message = "Unsloth Model Card")
# Upload model files
if save_directory:
@ -614,10 +598,7 @@ class ExportBackend:
LLAMA_CPP_DEFAULT_DIR,
_resolve_local_convert_script, # noqa: F401
)
os.environ.setdefault(
"UNSLOTH_LLAMA_CPP_SCRIPTS_DIR", LLAMA_CPP_DEFAULT_DIR
)
os.environ.setdefault("UNSLOTH_LLAMA_CPP_SCRIPTS_DIR", LLAMA_CPP_DEFAULT_DIR)
except ImportError:
if not _LLAMA_CPP_SCRIPTS_WARNING_EMITTED:
logger.warning(
@ -663,15 +644,11 @@ class ExportBackend:
# Relocate GGUF artifacts into the export directory.
# convert_to_gguf writes .gguf files to cwd (repo root)
# because --outfile is a relative path like "model.Q4_K_M.gguf".
new_ggufs = (
set(glob.glob(os.path.join(cwd, "*.gguf"))) - pre_existing_ggufs
)
new_ggufs = set(glob.glob(os.path.join(cwd, "*.gguf"))) - pre_existing_ggufs
for src in sorted(new_ggufs):
dest = os.path.join(abs_save_dir, os.path.basename(src))
shutil.move(src, dest)
logger.info(
f"Relocated GGUF: {os.path.basename(src)} → {abs_save_dir}/"
)
logger.info(f"Relocated GGUF: {os.path.basename(src)} → {abs_save_dir}/")
# Flatten any .gguf files from subdirectories into abs_save_dir.
# save_pretrained_gguf may create subdirs (e.g. model_gguf/)
@ -701,9 +678,7 @@ class ExportBackend:
# Also relocate Ollama Modelfile if present
modelfile = gguf_dir / "Modelfile"
if modelfile.is_file():
shutil.move(
str(modelfile), os.path.join(abs_save_dir, "Modelfile")
)
shutil.move(str(modelfile), os.path.join(abs_save_dir, "Modelfile"))
logger.info(f"Relocated Modelfile → {abs_save_dir}/")
shutil.rmtree(str(gguf_dir), ignore_errors = True)
logger.info(f"Cleaned up intermediate GGUF dir: {gguf_dir}")
@ -814,12 +789,8 @@ class ExportBackend:
repo_type = "model",
)
else:
self.current_model.push_to_hub(
repo_id, token = hf_token, private = private
)
self.current_tokenizer.push_to_hub(
repo_id, token = hf_token, private = private
)
self.current_model.push_to_hub(repo_id, token = hf_token, private = private)
self.current_tokenizer.push_to_hub(repo_id, token = hf_token, private = private)
logger.info(f"Adapter pushed successfully to {repo_id}")
return True, "LoRA adapter exported successfully", output_path

23
studio/backend/core/export/orchestrator.py
View file

@ -261,7 +261,11 @@ class ExportOrchestrator:
except (EOFError, OSError, ValueError):
return None
def _wait_response(self, expected_type: str, timeout: float = 3600.0) -> dict:
def _wait_response(
self,
expected_type: str,
timeout: float = 3600.0,
) -> dict:
"""Block until a response of the expected type arrives.
Export operations can take a very long time — GGUF conversion for
@ -318,9 +322,7 @@ class ExportOrchestrator:
expected_type,
)
raise RuntimeError(
f"Timeout waiting for '{expected_type}' response after {timeout}s"
)
raise RuntimeError(f"Timeout waiting for '{expected_type}' response after {timeout}s")
def _drain_queue(self) -> list:
"""Drain all pending responses."""
@ -371,9 +373,7 @@ class ExportOrchestrator:
elif self._proc is not None:
self._shutdown_subprocess(timeout = 2)
logger.info(
"Spawning fresh export subprocess for '%s'", checkpoint_path
)
logger.info("Spawning fresh export subprocess for '%s'", checkpoint_path)
self._spawn_subprocess(sub_config)
try:
@ -485,9 +485,7 @@ class ExportOrchestrator:
},
)
def _run_export(
self, export_type: str, params: dict
) -> Tuple[bool, str, Optional[str]]:
def _run_export(self, export_type: str, params: dict) -> Tuple[bool, str, Optional[str]]:
"""Send an export command to the subprocess and wait for result.
Returns ``(success, message, output_path)``. ``output_path`` is the
@ -554,12 +552,9 @@ class ExportOrchestrator:
finally:
self._export_active = False
def scan_checkpoints(
self, outputs_dir: str = str(outputs_root())
) -> List[Tuple[str, list]]:
def scan_checkpoints(self, outputs_dir: str = str(outputs_root())) -> List[Tuple[str, list]]:
"""Scan for checkpoints — no ML imports needed, runs locally."""
from utils.models.checkpoints import scan_checkpoints
return scan_checkpoints(outputs_dir = outputs_dir)

20
studio/backend/core/export/worker.py
View file

@ -362,12 +362,7 @@ def _handle_cleanup(backend, resp_queue: Any) -> None:
)
def run_export_process(
*,
cmd_queue: Any,
resp_queue: Any,
config: dict,
) -> None:
def run_export_process(*, cmd_queue: Any, resp_queue: Any, config: dict) -> None:
"""Subprocess entrypoint. Persistent — runs command loop until shutdown.
Args:
@ -383,9 +378,7 @@ def run_export_process(
_setup_log_capture(resp_queue)
os.environ["TOKENIZERS_PARALLELISM"] = "false"
os.environ["PYTHONWARNINGS"] = (
"ignore" # Suppress warnings at C-level before imports
)
os.environ["PYTHONWARNINGS"] = "ignore" # Suppress warnings at C-level before imports
# Force unbuffered output from any child Python process (e.g. the
# GGUF converter) so their prints surface in the log stream as they
# happen rather than at the end.
@ -430,7 +423,6 @@ def run_export_process(
if sys.platform == "win32":
try:
import triton # noqa: F401
logger.info("Triton available — torch.compile enabled")
except ImportError:
os.environ["TORCHDYNAMO_DISABLE"] = "1"
@ -467,9 +459,7 @@ def run_export_process(
import transformers
logger.info(
"Export subprocess loaded transformers %s", transformers.__version__
)
logger.info("Export subprocess loaded transformers %s", transformers.__version__)
except Exception as exc:
_send_response(
@ -570,9 +560,7 @@ def run_export_process(
)
except Exception as exc:
logger.error(
"Error handling command '%s': %s", cmd_type, exc, exc_info = True
)
logger.error("Error handling command '%s': %s", cmd_type, exc, exc_info = True)
_send_response(
resp_queue,
{

7
studio/backend/core/inference/anthropic_compat.py
View file

@ -42,8 +42,7 @@ def _anthropic_image_block_to_openai_part(block: dict) -> Optional[dict]:
def anthropic_messages_to_openai(
messages: list[dict],
system: Optional[Union[str, list]] = None,
messages: list[dict], system: Optional[Union[str, list]] = None
) -> list[dict]:
"""Convert Anthropic messages + system to OpenAI-format message dicts.
@ -125,9 +124,7 @@ def anthropic_messages_to_openai(
tc = b.get("content", "")
if isinstance(tc, list):
tc = " ".join(
p["text"]
for p in tc
if isinstance(p, dict) and p.get("type") == "text"
p["text"] for p in tc if isinstance(p, dict) and p.get("type") == "text"
)
tool_results.append(
{

33
studio/backend/core/inference/audio_codecs.py
View file

@ -77,12 +77,14 @@ class AudioCodecManager:
return
from snac import SNAC
self._snac_model = (
SNAC.from_pretrained("hubertsiuzdak/snac_24khz").to(device).eval()
)
self._snac_model = SNAC.from_pretrained("hubertsiuzdak/snac_24khz").to(device).eval()
logger.info("Loaded SNAC codec (24kHz)")
def _load_bicodec(self, device: str, model_repo_path: Optional[str] = None) -> None:
def _load_bicodec(
self,
device: str,
model_repo_path: Optional[str] = None,
) -> None:
if self._bicodec_tokenizer is not None:
return
import os
@ -90,9 +92,7 @@ class AudioCodecManager:
# Clone SparkAudio/Spark-TTS GitHub repo for the sparktts Python package
# (same approach as training — the HF model repos don't contain the package)
spark_code_dir = os.path.join(
os.path.dirname(model_repo_path or "."), "Spark-TTS"
)
spark_code_dir = os.path.join(os.path.dirname(model_repo_path or "."), "Spark-TTS")
sparktts_pkg = os.path.join(spark_code_dir, "sparktts")
if not os.path.isdir(sparktts_pkg):
logger.info(f"Cloning SparkAudio/Spark-TTS to {spark_code_dir}...")
@ -177,9 +177,7 @@ class AudioCodecManager:
# ── Decoders ─────────────────────────────────────────────────
def decode_snac(
self, generated_ids: torch.Tensor, device: str
) -> Tuple[bytes, int]:
def decode_snac(self, generated_ids: torch.Tensor, device: str) -> Tuple[bytes, int]:
"""
Decode SNAC tokens (Orpheus) into WAV bytes.
@ -195,9 +193,7 @@ class AudioCodecManager:
cropped = generated_ids[:, token_indices[1][-1] + 1 :]
else:
# Gracefully fall back to using entire output if marker not found
logger.warning(
"No START_OF_SPEECH token (128257) found — using full generated output"
)
logger.warning("No START_OF_SPEECH token (128257) found — using full generated output")
cropped = generated_ids
row = cropped[0]
@ -223,8 +219,7 @@ class AudioCodecManager:
layer_3.append(codes[7 * i + 6] - 24576)
snac_codes = [
torch.tensor(layer).unsqueeze(0).to(device)
for layer in [layer_1, layer_2, layer_3]
torch.tensor(layer).unsqueeze(0).to(device) for layer in [layer_1, layer_2, layer_3]
]
with torch.no_grad():
@ -254,16 +249,12 @@ class AudioCodecManager:
f"BiCodec decode: {len(global_matches)} global tokens, {len(semantic_matches)} semantic tokens"
)
if len(global_matches) < 10:
logger.info(
f"BiCodec generated text (first 500 chars): {generated_text[:500]}"
)
logger.info(f"BiCodec generated text (first 500 chars): {generated_text[:500]}")
if not semantic_matches:
raise ValueError("No bicodec_semantic tokens found in generated output")
semantic_ids = (
torch.tensor([int(t) for t in semantic_matches]).long().unsqueeze(0)
)
semantic_ids = torch.tensor([int(t) for t in semantic_matches]).long().unsqueeze(0)
# Speaker encoder expects exactly 32 global tokens (token_num=32 in BiCodec config).
# Pad with zeros or truncate to 32.

4
studio/backend/core/inference/chat_template_helpers.py
View file

@ -55,6 +55,4 @@ def apply_chat_template_for_generation(
break
if last_exc is not None:
raise last_exc
raise RuntimeError(
"apply_chat_template_for_generation: no attempt produced a result"
)
raise RuntimeError("apply_chat_template_for_generation: no attempt produced a result")

722
studio/backend/core/inference/external_provider.py
View file

File diff suppressed because it is too large Load diff

218
studio/backend/core/inference/inference.py
View file

@ -68,7 +68,13 @@ class HarmonyTextStreamer:
_re.DOTALL,
)
def __init__(self, tokenizer, *, skip_prompt: bool = True, timeout: float = 0.2):
def __init__(
self,
tokenizer,
*,
skip_prompt: bool = True,
timeout: float = 0.2,
):
import queue
self.tokenizer = tokenizer
@ -142,7 +148,6 @@ class HarmonyTextStreamer:
def __next__(self):
from queue import Empty
while True:
try:
val = self._queue.get(timeout = self.timeout)
@ -293,9 +298,7 @@ class InferenceBackend:
if config.is_audio:
audio_type = config.audio_type
adapter_info = " (LoRA adapter)" if config.is_lora else ""
logger.info(
f"Loading audio ({audio_type}) model{adapter_info}: {model_name}"
)
logger.info(f"Loading audio ({audio_type}) model{adapter_info}: {model_name}")
log_gpu_memory(f"Before loading {model_name}")
if audio_type == "csm":
@ -330,9 +333,7 @@ class InferenceBackend:
from huggingface_hub import snapshot_download
local_dir = base_path.split("/")[-1]
repo_path = snapshot_download(
base_path, local_dir = local_dir
)
repo_path = snapshot_download(base_path, local_dir = local_dir)
abs_repo_path = os.path.abspath(repo_path)
logger.info(
@ -443,9 +444,7 @@ class InferenceBackend:
)
# Reject CPU/disk offload for audio models too
raise_if_offloaded(
self.models[model_name]["model"], device_map, "Inference"
)
raise_if_offloaded(self.models[model_name]["model"], device_map, "Inference")
self.active_model_name = model_name
self.loading_models.discard(model_name)
@ -454,9 +453,7 @@ class InferenceBackend:
return True
model_type = "vision" if config.is_vision else "text"
adapter_info = (
" (LoRA adapter)" if self.models[model_name]["is_lora"] else ""
)
adapter_info = " (LoRA adapter)" if self.models[model_name]["is_lora"] else ""
logger.info(f"Loading {model_type} model{adapter_info}: {model_name}")
log_gpu_memory(f"Before loading {model_name}")
@ -482,14 +479,11 @@ class InferenceBackend:
from transformers import ProcessorMixin
if not (
isinstance(processor, ProcessorMixin)
or hasattr(processor, "image_processor")
isinstance(processor, ProcessorMixin) or hasattr(processor, "image_processor")
):
# For LoRA adapters, use the base model. For local merged exports,
# read export_metadata.json to find the original base model.
processor_source = (
config.base_model if config.is_lora else config.identifier
)
processor_source = config.base_model if config.is_lora else config.identifier
if not config.is_lora and config.is_local:
_meta_path = Path(config.path) / "export_metadata.json"
try:
@ -510,9 +504,7 @@ class InferenceBackend:
token = hf_token if hf_token and hf_token.strip() else None,
trust_remote_code = trust_remote_code,
)
logger.info(
f"Loaded {type(processor).__name__} from {processor_source}"
)
logger.info(f"Loaded {type(processor).__name__} from {processor_source}")
self.models[model_name]["model"] = model
self.models[model_name]["tokenizer"] = processor
@ -536,9 +528,7 @@ class InferenceBackend:
self.models[model_name]["model"] = model
self.models[model_name]["tokenizer"] = tokenizer
raise_if_offloaded(
self.models[model_name]["model"], device_map, "Inference"
)
raise_if_offloaded(self.models[model_name]["model"], device_map, "Inference")
# Load chat template info
self._load_chat_template_info(model_name)
@ -588,11 +578,7 @@ class InferenceBackend:
import sys as _sys
from utils.cache_cleanup import clear_unsloth_compiled_cache
_preserve = (
["Unsloth*Trainer.py"]
if _sys.platform in ("win32", "darwin")
else None
)
_preserve = ["Unsloth*Trainer.py"] if _sys.platform in ("win32", "darwin") else None
clear_unsloth_compiled_cache(preserve_patterns = _preserve)
logger.info(f"Model '{model_name}' successfully unloaded.")
@ -665,13 +651,9 @@ class InferenceBackend:
base_model_name = lora_config.base_model
# 1. Load the base model if it's not already in memory
if base_model_name not in self.models or not self.models[
base_model_name
].get("model"):
if base_model_name not in self.models or not self.models[base_model_name].get("model"):
logger.info(f"Base model '{base_model_name}' not loaded, loading now.")
base_config = ModelConfig.from_ui_selection(
base_model_name, None, is_lora = False
)
base_config = ModelConfig.from_ui_selection(base_model_name, None, is_lora = False)
if not self.load_model(
base_config,
max_seq_length,
@ -707,9 +689,7 @@ class InferenceBackend:
logger.error(traceback.format_exc())
return False, None, None
def load_adapter(
self, base_model_name: str, adapter_path: str, adapter_name: str
) -> bool:
def load_adapter(self, base_model_name: str, adapter_path: str, adapter_name: str) -> bool:
"""
Loads an adapter onto the model ONLY if it's not already attached.
"""
@ -790,16 +770,12 @@ class InferenceBackend:
)
model.base_model.disable_adapter_layers()
else:
logger.info(
f"Compare mode: model '{base}' is not a PeftModel, already base"
)
logger.info(f"Compare mode: model '{base}' is not a PeftModel, already base")
elif use_adapter is True:
# Re-enable LoRA layers → adapter output
if isinstance(model, (PeftModel, PeftModelForCausalLM)):
logger.info(
f"Compare mode: enabling adapters on '{base}' for LoRA generation"
)
logger.info(f"Compare mode: enabling adapters on '{base}' for LoRA generation")
model.base_model.enable_adapter_layers()
else:
logger.warning("use_adapter=true but model is not a PeftModel")
@ -807,15 +783,11 @@ class InferenceBackend:
elif isinstance(use_adapter, str):
# Enable adapters and set the specific one active
if isinstance(model, (PeftModel, PeftModelForCausalLM)):
logger.info(
f"Compare mode: enabling adapter '{use_adapter}' on '{base}'"
)
logger.info(f"Compare mode: enabling adapter '{use_adapter}' on '{base}'")
model.base_model.enable_adapter_layers()
self.set_active_adapter(base, use_adapter)
else:
logger.warning(
f"use_adapter='{use_adapter}' but model is not a PeftModel"
)
logger.warning(f"use_adapter='{use_adapter}' but model is not a PeftModel")
def generate_with_adapter_control(
self,
@ -996,8 +968,7 @@ class InferenceBackend:
processor = model_info.get("processor")
has_image_processing = processor is not None and (
isinstance(processor, ProcessorMixin)
or hasattr(processor, "image_processor")
isinstance(processor, ProcessorMixin) or hasattr(processor, "image_processor")
)
if has_image_processing:
yield from self._generate_vision_response(
@ -1029,7 +1000,6 @@ class InferenceBackend:
MODEL_TO_TEMPLATE_MAPPER,
get_tokenizer_chat_template,
)
model_name_lower = self.active_model_name.lower()
# Check if model has a registered template
@ -1053,9 +1023,7 @@ class InferenceBackend:
# Step 2: Format with tokenizer.apply_chat_template()
if system_prompt:
template_messages = [
{"role": "system", "content": system_prompt}
] + messages
template_messages = [{"role": "system", "content": system_prompt}] + messages
else:
template_messages = messages
try:
@ -1119,7 +1087,6 @@ class InferenceBackend:
user_message = ""
if messages and messages[-1]["role"] == "user":
import re
user_message = messages[-1]["content"]
user_message = re.sub(r"<img[^>]*>", "", user_message).strip()
@ -1171,9 +1138,7 @@ class InferenceBackend:
else:
# Text-only for vision model
formatted_prompt = self.format_chat_prompt(messages, system_prompt)
inputs = raw_tokenizer(formatted_prompt, return_tensors = "pt").to(
model.device
)
inputs = raw_tokenizer(formatted_prompt, return_tensors = "pt").to(model.device)
# Stream with TextIteratorStreamer + background thread
try:
@ -1385,7 +1350,9 @@ class InferenceBackend:
yield f"Error: {str(e)}"
def generate_whisper_response(
self, audio_array, cancel_event = None
self,
audio_array,
cancel_event = None,
) -> Generator[str, None, None]:
"""Whisper ASR — takes audio numpy array, yields transcribed text.
@ -1411,7 +1378,6 @@ class InferenceBackend:
def _is_gpt_oss_model(self, model_name: str = None) -> bool:
"""Check if the given (or active) model uses the gpt-oss harmony protocol."""
from utils.datasets import is_gpt_oss_model_name
return is_gpt_oss_model_name(model_name or self.active_model_name or "")
def generate_stream(
@ -1459,9 +1425,7 @@ class InferenceBackend:
timeout = 0.2,
)
except Exception as e:
logger.warning(
f"HarmonyTextStreamer init failed, falling back: {e}"
)
logger.warning(f"HarmonyTextStreamer init failed, falling back: {e}")
streamer = TextIteratorStreamer(
tokenizer,
skip_prompt = True,
@ -1496,7 +1460,6 @@ class InferenceBackend:
StoppingCriteria,
StoppingCriteriaList,
)
class _CancelCriteria(StoppingCriteria):
def __init__(self, ev):
self.ev = ev
@ -1559,9 +1522,7 @@ class InferenceBackend:
cancel_event.set()
thread.join(timeout = 10)
if thread.is_alive():
logger.warning(
"Generation thread did not exit after cancel/join timeout"
)
logger.warning("Generation thread did not exit after cancel/join timeout")
if err.get("msg"):
yield f"Error: {err['msg']}"
@ -1638,21 +1599,12 @@ class InferenceBackend:
raise RuntimeError(f"Unknown audio_type: {audio_type}")
def _generate_snac(
self,
model,
tokenizer,
text,
temperature,
top_p,
max_new_tokens,
repetition_penalty,
self, model, tokenizer, text, temperature, top_p, max_new_tokens, repetition_penalty
):
"""Generate audio using SNAC codec (Orpheus)."""
device = model.device
start_token = torch.tensor([[128259]], device = device) # START_OF_HUMAN
end_tokens = torch.tensor(
[[128009, 128260]], device = device
) # EOT, END_OF_HUMAN
end_tokens = torch.tensor([[128009, 128260]], device = device) # EOT, END_OF_HUMAN
text_ids = tokenizer(text, return_tensors = "pt").input_ids.to(device)
input_ids = torch.cat([start_token, text_ids, end_tokens], dim = 1)
attention_mask = torch.ones_like(input_ids)
@ -1676,20 +1628,12 @@ class InferenceBackend:
inputs = processor(
f"[{speaker_id}]{text}", add_special_tokens = True, return_tensors = "pt"
).to(model.device)
audio_values = model.generate(
**inputs, max_new_tokens = max_new_tokens, output_audio = True
)
audio_values = model.generate(**inputs, max_new_tokens = max_new_tokens, output_audio = True)
return self._audio_codec_manager.decode_csm(audio_values)
def _generate_bicodec(
self, model, tokenizer, text, temperature, top_k, max_new_tokens
):
def _generate_bicodec(self, model, tokenizer, text, temperature, top_k, max_new_tokens):
"""Generate audio using BiCodec (Spark-TTS)."""
prompt = (
"<|task_tts|><|start_content|>"
+ text
+ "<|end_content|><|start_global_token|>"
)
prompt = "<|task_tts|><|start_content|>" + text + "<|end_content|><|start_global_token|>"
inputs = tokenizer([prompt], return_tensors = "pt").to(model.device)
generated = model.generate(
**inputs,
@ -1761,9 +1705,7 @@ class InferenceBackend:
def __init__(self, penalty: float):
self.penalty_last_n = 64
if not isinstance(penalty, float) or penalty <= 0:
raise ValueError(
f"`penalty` has to be a positive float, but is {penalty}"
)
raise ValueError(f"`penalty` has to be a positive float, but is {penalty}")
self.penalty = penalty
@torch.no_grad()
@ -1788,12 +1730,8 @@ class InferenceBackend:
)
return scores
generation_utils.RepetitionPenaltyLogitsProcessor = (
RepetitionPenaltyLogitsProcessorPatch
)
logger.info(
"Patched RepetitionPenaltyLogitsProcessor with 64-token window for OuteTTS"
)
generation_utils.RepetitionPenaltyLogitsProcessor = RepetitionPenaltyLogitsProcessorPatch
logger.info("Patched RepetitionPenaltyLogitsProcessor with 64-token window for OuteTTS")
def _apply_chat_template_for_generation(
self,
@ -1813,7 +1751,6 @@ class InferenceBackend:
from core.inference.chat_template_helpers import (
apply_chat_template_for_generation,
)
return apply_chat_template_for_generation(
tokenizer,
messages,
@ -1823,7 +1760,11 @@ class InferenceBackend:
preserve_thinking = preserve_thinking,
)
def format_chat_prompt(self, messages: list, system_prompt: str = None) -> str:
def format_chat_prompt(
self,
messages: list,
system_prompt: str = None,
) -> str:
if not self.active_model_name or self.active_model_name not in self.models:
logger.error("No active model available")
return ""
@ -1832,9 +1773,7 @@ class InferenceBackend:
logger.error("Tokenizer not loaded for active model")
return ""
chat_template_info = self.models[self.active_model_name].get(
"chat_template_info", {}
)
chat_template_info = self.models[self.active_model_name].get("chat_template_info", {})
tokenizer = self.models[self.active_model_name]["tokenizer"]
tokenizer = getattr(tokenizer, "tokenizer", tokenizer)
@ -1851,14 +1790,11 @@ class InferenceBackend:
if role in ["system", "user", "assistant"] and content.strip():
if role == last_role:
logger.debug(
f"Skipping consecutive {role} message to maintain alternation"
)
logger.debug(f"Skipping consecutive {role} message to maintain alternation")
continue
if role == "user":
import re
clean_content = re.sub(r"<[^>]+>", "", content).strip()
if clean_content:
chat_messages.append({"role": role, "content": clean_content})
@ -1870,9 +1806,7 @@ class InferenceBackend:
continue
if chat_messages and chat_messages[-1]["role"] == "assistant":
logger.debug(
"Removing final assistant message to ensure proper alternation"
)
logger.debug("Removing final assistant message to ensure proper alternation")
chat_messages.pop()
logger.info(f"Sending {len(chat_messages)} messages to tokenizer:")
@ -1887,10 +1821,7 @@ class InferenceBackend:
return formatted_prompt
except Exception as e:
error_msg = str(e).lower()
if (
"chat_template is not set" in error_msg
or "no template argument" in error_msg
):
if "chat_template is not set" in error_msg or "no template argument" in error_msg:
logger.info(
f"Base model detected - no built-in chat template available, using fallback formatting"
)
@ -1901,9 +1832,7 @@ class InferenceBackend:
)
if chat_template_info.get("has_template", False):
logger.info(
"Falling back to manual template formatting based on detected patterns"
)
logger.info("Falling back to manual template formatting based on detected patterns")
template_type = chat_template_info.get("format_type", "generic")
manual_prompt = self._format_chat_manual(
chat_messages,
@ -1916,9 +1845,7 @@ class InferenceBackend:
logger.info("Using generic chat formatting for base model")
return self._format_generic_template(chat_messages, {})
def _format_chat_manual(
self, messages: list, template_type: str, special_tokens: dict
) -> str:
def _format_chat_manual(self, messages: list, template_type: str, special_tokens: dict) -> str:
"""
Manual chat formatting fallback for when tokenizer template fails
@ -1949,9 +1876,7 @@ class InferenceBackend:
for msg in messages:
role = msg["role"]
content = msg["content"]
formatted += (
f"<|start_header_id|>{role}<|end_header_id|>\n\n{content}<|eot_id|>"
)
formatted += f"<|start_header_id|>{role}<|end_header_id|>\n\n{content}<|eot_id|>"
formatted += "<|start_header_id|>assistant<|end_header_id|>\n\n"
return formatted
@ -1980,10 +1905,7 @@ class InferenceBackend:
formatted += f"[INST] {user_content} [/INST]"
if (
i + 1 < len(conversation)
and conversation[i + 1]["role"] == "assistant"
):
if i + 1 < len(conversation) and conversation[i + 1]["role"] == "assistant":
formatted += f" {conversation[i + 1]['content']}</s>"
i += 2
else:
@ -2088,7 +2010,11 @@ class InferenceBackend:
except Exception as e:
logger.warning(f"Could not fully reset generation state: {e}")
def resize_image(self, img, max_size: int = 800):
def resize_image(
self,
img,
max_size: int = 800,
):
"""Resize image while maintaining aspect ratio if either dimension exceeds max_size"""
if img is None:
return None
@ -2107,7 +2033,6 @@ class InferenceBackend:
# Strip harmony protocol tokens and other gpt-oss added tokens
# (e.g. <|return|>) that may leak past the streamer.
import re
text = re.sub(r"<\|[a-z_]+\|>", "", text)
return text.strip()
@ -2119,9 +2044,7 @@ class InferenceBackend:
return text.strip()
def _load_chat_template_info(self, model_name: str):
if model_name not in self.models or not self.models[model_name].get(
"tokenizer"
):
if model_name not in self.models or not self.models[model_name].get("tokenizer"):
return
tokenizer = self.models[model_name]["tokenizer"]
@ -2139,9 +2062,7 @@ class InferenceBackend:
# Try exact match first
model_name_lower = model_name.lower()
if model_name_lower in MODEL_TO_TEMPLATE_MAPPER:
chat_template_info["template_name"] = MODEL_TO_TEMPLATE_MAPPER[
model_name_lower
]
chat_template_info["template_name"] = MODEL_TO_TEMPLATE_MAPPER[model_name_lower]
logger.info(
f"Detected template '{chat_template_info['template_name']}' for {model_name} from mapper"
)
@ -2149,17 +2070,13 @@ class InferenceBackend:
# Try partial match (for variants like model_name-bnb-4bit)
for key in MODEL_TO_TEMPLATE_MAPPER:
if key in model_name_lower or model_name_lower in key:
chat_template_info["template_name"] = MODEL_TO_TEMPLATE_MAPPER[
key
]
chat_template_info["template_name"] = MODEL_TO_TEMPLATE_MAPPER[key]
logger.info(
f"Detected template '{chat_template_info['template_name']}' for {model_name} (partial match)"
)
break
except Exception as e:
logger.warning(
f"Could not detect template from mapper for {model_name}: {e}"
)
logger.warning(f"Could not detect template from mapper for {model_name}: {e}")
try:
if hasattr(tokenizer, "chat_template") and tokenizer.chat_template:
@ -2168,10 +2085,7 @@ class InferenceBackend:
template_str = tokenizer.chat_template.lower()
if (
"start_header_id" in template_str
and "end_header_id" in template_str
):
if "start_header_id" in template_str and "end_header_id" in template_str:
chat_template_info["format_type"] = "llama3"
elif "[inst]" in template_str and "[/inst]" in template_str:
chat_template_info["format_type"] = "mistral"
@ -2198,9 +2112,7 @@ class InferenceBackend:
chat_template_info["special_tokens"] = special_tokens
else:
logger.info(
f"No chat template found for {model_name}, will use generic formatting"
)
logger.info(f"No chat template found for {model_name}, will use generic formatting")
except Exception as e:
logger.error(f"Error loading chat template info for {model_name}: {e}")
@ -2212,9 +2124,7 @@ class InferenceBackend:
f"Chat template loaded for {model_name}: {chat_template_info['format_type']} format"
)
else:
logger.info(
f"No built-in chat template for {model_name}, will use generic formatting"
)
logger.info(f"No built-in chat template for {model_name}, will use generic formatting")
def get_current_model(self) -> Optional[str]:
"""Get currently active model name"""

492
studio/backend/core/inference/llama_cpp.py
View file

File diff suppressed because it is too large Load diff

32
studio/backend/core/inference/llama_server_args.py
View file

@ -126,9 +126,7 @@ def is_managed_flag(flag: str) -> bool:
# stripped from inherited extras so they can't last-wins-override an
# Apply that re-sets the same field.
_CONTEXT_FLAGS: frozenset[str] = frozenset({"-c", "--ctx-size"})
_CACHE_FLAGS: frozenset[str] = frozenset(
{"-ctk", "--cache-type-k", "-ctv", "--cache-type-v"}
)
_CACHE_FLAGS: frozenset[str] = frozenset({"-ctk", "--cache-type-k", "-ctv", "--cache-type-v"})
_SPEC_FLAGS: frozenset[str] = frozenset(
{
"--spec-default",
@ -157,15 +155,11 @@ _TEMPLATE_FLAGS: frozenset[str] = frozenset(
}
)
_SHADOWING_FLAGS: frozenset[str] = (
_CONTEXT_FLAGS | _CACHE_FLAGS | _SPEC_FLAGS | _TEMPLATE_FLAGS
)
_SHADOWING_FLAGS: frozenset[str] = _CONTEXT_FLAGS | _CACHE_FLAGS | _SPEC_FLAGS | _TEMPLATE_FLAGS
# Shadowing flags that take no value -- strip the flag only, never the
# following token.
_BOOLEAN_SHADOWING_FLAGS: frozenset[str] = frozenset(
{"--spec-default", "--jinja", "--no-jinja"}
)
_BOOLEAN_SHADOWING_FLAGS: frozenset[str] = frozenset({"--spec-default", "--jinja", "--no-jinja"})
def parse_ctx_override(args: Optional[Iterable[str]]) -> Optional[int]:
@ -192,31 +186,22 @@ def parse_ctx_override(args: Optional[Iterable[str]]) -> Optional[int]:
i += 1
else:
if i + 1 >= n or _flag_name(tokens[i + 1]) is not None:
raise ValueError(
f"llama-server flag '{flag}' requires an integer value"
)
raise ValueError(f"llama-server flag '{flag}' requires an integer value")
raw_value = tokens[i + 1]
i += 2
try:
value = int(str(raw_value).strip())
except ValueError as exc:
raise ValueError(
f"llama-server flag '{flag}' requires an integer value"
) from exc
raise ValueError(f"llama-server flag '{flag}' requires an integer value") from exc
if value < 0:
raise ValueError(
f"llama-server flag '{flag}' requires a non-negative integer value"
)
raise ValueError(f"llama-server flag '{flag}' requires a non-negative integer value")
override = value
return override
def resolve_requested_ctx(
args: Optional[Iterable[str]],
fallback_n_ctx: int,
) -> int:
def resolve_requested_ctx(args: Optional[Iterable[str]], fallback_n_ctx: int) -> int:
"""Return the context size load_model should treat as requested.
Single source of truth for the two-line ``ctx_override = parse_ctx_override(...);
@ -267,8 +252,7 @@ def parse_cache_override(args: Optional[Iterable[str]]) -> Optional[str]:
def resolve_cache_type_kv(
args: Optional[Iterable[str]],
fallback_cache_type_kv: Optional[str],
args: Optional[Iterable[str]], fallback_cache_type_kv: Optional[str]
) -> Optional[str]:
"""Return the cache type load_model should treat as requested.

17
studio/backend/core/inference/mcp_client.py
View file

@ -34,10 +34,7 @@ def parse_stdio_command(address: str) -> list[str]:
# posix=False keeps backslash paths intact but also keeps the surrounding
# quotes on a token. Strip a matched pair so the argv reaches the
# subprocess clean ('"C:\\Program Files\\node"' -> C:\\Program Files\\node).
parts = [
p[1:-1] if len(p) >= 2 and p[0] == p[-1] and p[0] in "\"'" else p
for p in parts
]
parts = [p[1:-1] if len(p) >= 2 and p[0] == p[-1] and p[0] in "\"'" else p for p in parts]
return parts
@ -104,7 +101,6 @@ async def clear_oauth_tokens_async(url: str) -> None:
failing must not make the delete / update route 500."""
try:
from fastmcp.client.auth import OAuth
auth = OAuth(mcp_url = url, token_storage = _oauth_store())
await auth.token_storage_adapter.clear()
except Exception as exc: # noqa: BLE001
@ -112,7 +108,11 @@ async def clear_oauth_tokens_async(url: str) -> None:
logger.warning("Failed to clear OAuth tokens for %s: %s", url, exc)
def _client(url: str, headers: Optional[dict], use_oauth: bool = False):
def _client(
url: str,
headers: Optional[dict],
use_oauth: bool = False,
):
from fastmcp import Client
if is_stdio(url):
@ -142,13 +142,10 @@ def _client(url: str, headers: Optional[dict], use_oauth: bool = False):
auth = None
if use_oauth:
from fastmcp.client.auth import OAuth
auth = OAuth(mcp_url = url, token_storage = _oauth_store())
transport_cls = (
SSETransport
if infer_transport_type_from_url(url) == "sse"
else StreamableHttpTransport
SSETransport if infer_transport_type_from_url(url) == "sse" else StreamableHttpTransport
)
return Client(transport_cls(url = url, headers = headers or None, auth = auth))

19
studio/backend/core/inference/mlx_inference.py
View file

@ -133,7 +133,6 @@ class MLXInferenceBackend:
if hf_token:
import os
os.environ["HF_TOKEN"] = hf_token
self._configure_memory_limits()
@ -311,9 +310,7 @@ class MLXInferenceBackend:
elif isinstance(content, list):
# Prepend image if not already there
has_image = any(
p.get("type") == "image"
for p in content
if isinstance(p, dict)
p.get("type") == "image" for p in content if isinstance(p, dict)
)
if not has_image:
content.insert(0, {"type": "image"})
@ -383,9 +380,7 @@ class MLXInferenceBackend:
preserve_thinking = preserve_thinking,
)
if prompt is None:
raise RuntimeError(
"apply_chat_template returned None — tokenizer may be incompatible"
)
raise RuntimeError("apply_chat_template returned None — tokenizer may be incompatible")
sampler = make_sampler(
temp = temperature,
@ -441,7 +436,6 @@ class MLXInferenceBackend:
break
except Exception as e:
import traceback
logger.error("stream_generate failed:\n%s", traceback.format_exc())
raise
finally:
@ -538,9 +532,7 @@ class MLXInferenceBackend:
**vlm_kwargs,
):
final_response = response
token_text = (
response.text if hasattr(response, "text") else str(response)
)
token_text = response.text if hasattr(response, "text") else str(response)
cumulative += token_text
yield cumulative
if cancel_event and cancel_event.is_set():
@ -556,7 +548,10 @@ class MLXInferenceBackend:
)
def generate_with_adapter_control(
self, use_adapter = None, cancel_event = None, **gen_kwargs
self,
use_adapter = None,
cancel_event = None,
**gen_kwargs,
) -> Generator[str, None, None]:
# MLX LoRA adapter toggling not yet supported — generate normally
yield from self.generate_chat_response(cancel_event = cancel_event, **gen_kwargs)

60
studio/backend/core/inference/orchestrator.py
View file

@ -63,9 +63,7 @@ class InferenceOrchestrator:
self._resp_queue: Any = None
self._cancel_event: Any = None # mp.Event — set to cancel generation instantly
self._lock = threading.Lock()
self._gen_lock = (
threading.Lock()
) # Serializes generation — one request at a time
self._gen_lock = threading.Lock() # Serializes generation — one request at a time
# Dispatcher state — for compare mode (adapter-controlled requests).
# Instead of serializing via _gen_lock, adapter-controlled requests
@ -95,9 +93,7 @@ class InferenceOrchestrator:
logger.info("InferenceOrchestrator initialized (subprocess mode)")
# Kick off background fetch of top models from HF
threading.Thread(
target = self._fetch_top_models, daemon = True, name = "top-models"
).start()
threading.Thread(target = self._fetch_top_models, daemon = True, name = "top-models").start()
# ------------------------------------------------------------------
# Default models (top GGUFs fetched dynamically from HF)
@ -125,7 +121,6 @@ class InferenceOrchestrator:
"""Fetch top GGUF and non-GGUF repos from unsloth by downloads."""
try:
import httpx
resp = httpx.get(
"https://huggingface.co/api/models",
params = {
@ -140,14 +135,12 @@ class InferenceOrchestrator:
models = resp.json()
# Top 40 GGUFs - frontend pages through them on-demand via
# infinite scroll, so we send a deep pool.
gguf_ids = [
m["id"] for m in models if m.get("id", "").upper().endswith("-GGUF")
][:40]
gguf_ids = [m["id"] for m in models if m.get("id", "").upper().endswith("-GGUF")][
:40
]
# Top 40 non-GGUF hub models
hub_ids = [
m["id"]
for m in models
if not m.get("id", "").upper().endswith("-GGUF")
m["id"] for m in models if not m.get("id", "").upper().endswith("-GGUF")
][:40]
if gguf_ids:
self._top_gguf_cache = gguf_ids
@ -277,7 +270,11 @@ class InferenceOrchestrator:
except (EOFError, OSError, ValueError):
return None
def _wait_response(self, expected_type: str, timeout: float = 300.0) -> dict:
def _wait_response(
self,
expected_type: str,
timeout: float = 300.0,
) -> dict:
"""Block until a response of the expected type arrives.
Also handles 'status' and 'error' events during the wait.
@ -329,8 +326,7 @@ class InferenceOrchestrator:
)
raise RuntimeError(
f"Timeout waiting for '{expected_type}' response "
f"(no activity for {timeout}s)"
f"Timeout waiting for '{expected_type}' response " f"(no activity for {timeout}s)"
)
def _drain_queue(self) -> list:
@ -556,7 +552,11 @@ class InferenceOrchestrator:
with self._mailbox_lock:
self._mailboxes.pop(request_id, None)
def _drain_mailbox(self, mailbox: queue.Queue, timeout: float = 5.0) -> None:
def _drain_mailbox(
self,
mailbox: queue.Queue,
timeout: float = 5.0,
) -> None:
"""Drain a mailbox until gen_done/gen_error, discarding tokens."""
deadline = time.monotonic() + timeout
while time.monotonic() < deadline:
@ -683,8 +683,7 @@ class InferenceOrchestrator:
# First stall and Xet was enabled -> retry with Xet disabled
if attempt == 0 and not disable_xet:
logger.warning(
"Download stalled for '%s' -- retrying with "
"HF_HUB_DISABLE_XET=1",
"Download stalled for '%s' -- retrying with HF_HUB_DISABLE_XET=1",
model_name,
)
self._shutdown_subprocess(timeout = 5)
@ -715,13 +714,9 @@ class InferenceOrchestrator:
# capabilities without re-entering the subprocess.
_tpl_info = model_info.get("chat_template_info")
if isinstance(_tpl_info, dict):
self.models[self.active_model_name]["chat_template_info"] = (
_tpl_info
)
self.models[self.active_model_name]["chat_template_info"] = _tpl_info
self.loading_models.discard(model_name)
logger.info(
"Model '%s' loaded successfully in subprocess", model_name
)
logger.info("Model '%s' loaded successfully in subprocess", model_name)
return True
else:
error = resp.get("error", "Failed to load model")
@ -1157,9 +1152,7 @@ class InferenceOrchestrator:
if resp is None:
if not self._ensure_subprocess_alive():
raise RuntimeError(
"Inference subprocess crashed during audio generation"
)
raise RuntimeError("Inference subprocess crashed during audio generation")
continue
rtype = resp.get("type", "")
@ -1251,9 +1244,7 @@ class InferenceOrchestrator:
# Convert numpy array to list for mp.Queue serialization
audio_data = (
audio_array.tolist()
if hasattr(audio_array, "tolist")
else list(audio_array)
audio_array.tolist() if hasattr(audio_array, "tolist") else list(audio_array)
)
cmd = {
@ -1314,7 +1305,11 @@ class InferenceOrchestrator:
# Local helpers (no subprocess needed)
# ------------------------------------------------------------------
def resize_image(self, img, max_size: int = 800):
def resize_image(
self,
img,
max_size: int = 800,
):
"""Resize image while maintaining aspect ratio.
No ML imports needed — runs locally in parent process.
"""
@ -1357,7 +1352,6 @@ class InferenceOrchestrator:
"""Parent-side gpt-oss detection so the safetensors route can run
the same guard without an IPC round-trip to the subprocess."""
from utils.datasets import is_gpt_oss_model_name
return is_gpt_oss_model_name(model_name or self.active_model_name or "")

21
studio/backend/core/inference/pricing.py
View file

@ -106,11 +106,7 @@ def _lookup(provider: str, model: str) -> Optional[dict[str, float]]:
return None
def calculate_cost(
provider: str,
model: str,
usage: dict[str, Any],
) -> dict[str, float]:
def calculate_cost(provider: str, model: str, usage: dict[str, Any]) -> dict[str, float]:
"""Return a per-turn USD cost breakdown with per-bucket + total
fields so the frontend can render either a single number or a
tooltip without re-doing the math. When the model isn't in the
@ -141,8 +137,7 @@ def calculate_cost(
# Clamp tokens >=0 so corrupted payloads can't produce a negative bill.
cache_creation = max(0, int(usage.get("cache_creation_input_tokens") or 0))
cache_read_native_present = (
"cache_read_input_tokens" in usage
and usage.get("cache_read_input_tokens") is not None
"cache_read_input_tokens" in usage and usage.get("cache_read_input_tokens") is not None
)
cache_read = max(0, int(usage.get("cache_read_input_tokens") or 0))
# Fallback to mirrored prompt_tokens_details only when the native
@ -226,14 +221,10 @@ def calculate_cost(
if cc_5m + cc_1h == 0 and cache_creation > 0:
# No breakdown -- assume default 5m pool.
cc_5m = cache_creation
out["cache_write_usd"] = (
cc_5m / 1_000_000.0
) * base * ANTHROPIC_CACHE_5M_WRITE_MULT + (
out["cache_write_usd"] = (cc_5m / 1_000_000.0) * base * ANTHROPIC_CACHE_5M_WRITE_MULT + (
cc_1h / 1_000_000.0
) * base * ANTHROPIC_CACHE_1H_WRITE_MULT
out["cache_read_usd"] = (
(cache_read / 1_000_000.0) * base * ANTHROPIC_CACHE_READ_MULT
)
out["cache_read_usd"] = (cache_read / 1_000_000.0) * base * ANTHROPIC_CACHE_READ_MULT
# Server-tool surcharges.
srv = usage.get("server_tool_use") or {}
if isinstance(srv, dict):
@ -250,9 +241,7 @@ def calculate_cost(
if cache_read > 0:
non_cached_input = max(0, input_tokens - cache_read)
out["input_usd"] = (non_cached_input / 1_000_000.0) * base
out["cache_read_usd"] = (
(cache_read / 1_000_000.0) * base * OPENAI_CACHE_READ_MULT
)
out["cache_read_usd"] = (cache_read / 1_000_000.0) * base * OPENAI_CACHE_READ_MULT
# OpenAI server-tool surcharges arrive under `openai_tool_use`
# (normalised by the SSE finaliser from output array items).
srv = usage.get("openai_tool_use") or {}

3
studio/backend/core/inference/providers.py
View file

@ -258,8 +258,7 @@ PROVIDER_REGISTRY: dict[str, dict[str, Any]] = {
# sources. The HF /v1/models response is otherwise hundreds of
# ids long (community fine-tunes, mirrors, fp8 variants, etc.).
"model_id_allowlist": re.compile(
r"^(openai|deepseek-ai|google|meta-llama|Qwen|moonshotai|"
r"mistralai|zai-org)/"
r"^(openai|deepseek-ai|google|meta-llama|Qwen|moonshotai|mistralai|zai-org)/"
),
# Cap the post-filter list. /v1/models has no server-side limit
# or popularity sort, so this is just "first N matches" — pair it

23
studio/backend/core/inference/safetensors_agentic.py
View file

@ -86,9 +86,7 @@ def _detect_render_html_tool_start(content: str) -> bool:
if not function_match and tool_call_index < 0:
return False
if function_match and (
tool_call_index < 0 or function_match.start() < tool_call_index
):
if function_match and (tool_call_index < 0 or function_match.start() < tool_call_index):
return function_match.group(1) == "render_html"
if tool_call_index >= 0:
@ -98,7 +96,12 @@ def _detect_render_html_tool_start(content: str) -> bool:
return False
def _coerce_arguments(raw_args, *, heal: bool, tool_name: str = "") -> dict:
def _coerce_arguments(
raw_args,
*,
heal: bool,
tool_name: str = "",
) -> dict:
"""Normalise tool ``arguments`` to a dict.
Some templates emit a JSON string, others a bare query string. With
@ -403,15 +406,11 @@ def run_safetensors_tool_loop(
"final answer."
)
else:
already_ran_ok = any(
k == tc_key and not err for k, err in tool_call_history
)
already_ran_ok = any(k == tc_key and not err for k, err in tool_call_history)
if already_ran_ok:
result = DUPLICATE_CALL_NUDGE
else:
eff_timeout = (
None if tool_call_timeout >= 9999 else tool_call_timeout
)
eff_timeout = None if tool_call_timeout >= 9999 else tool_call_timeout
try:
result = execute_tool(
tool_name,
@ -432,9 +431,7 @@ def run_safetensors_tool_loop(
"result": result,
}
is_error = isinstance(result, str) and result.lstrip().startswith(
TOOL_ERROR_PREFIXES
)
is_error = isinstance(result, str) and result.lstrip().startswith(TOOL_ERROR_PREFIXES)
if tool_name == "render_html" and not is_error:
render_html_succeeded = True
tool_call_history.append((tc_key, is_error))

10
studio/backend/core/inference/tool_call_parser.py
View file

@ -160,9 +160,7 @@ def parse_tool_calls_from_text(content: str, *, id_offset: int = 0) -> list[dict
},
}
if isinstance(tc["function"]["arguments"], dict):
tc["function"]["arguments"] = json.dumps(
tc["function"]["arguments"]
)
tc["function"]["arguments"] = json.dumps(tc["function"]["arguments"])
tool_calls.append(tc)
except (json.JSONDecodeError, ValueError):
pass
@ -179,11 +177,7 @@ def parse_tool_calls_from_text(content: str, *, id_offset: int = 0) -> list[dict
for idx, fm in enumerate(func_starts):
func_name = fm.group(1)
body_start = fm.end()
next_func = (
func_starts[idx + 1].start()
if idx + 1 < len(func_starts)
else len(content)
)
next_func = func_starts[idx + 1].start() if idx + 1 < len(func_starts) else len(content)
end_tag = _TC_END_TAG_RE.search(content[body_start:])
if end_tag:
body_end = body_start + end_tag.start()

144
studio/backend/core/inference/tools.py
View file

@ -121,9 +121,7 @@ _BLOCKED_COMMANDS = (
)
_SHELL_SEPARATORS = frozenset(
{";", "&&", "||", "|", "&", "\n", "(", ")", "`", "{", "}"}
)
_SHELL_SEPARATORS = frozenset({";", "&&", "||", "|", "&", "\n", "(", ")", "`", "{", "}"})
# Bash keywords that introduce a new command position (then $cmd, do $cmd, etc.).
_SHELL_KEYWORDS_AS_SEP = frozenset({"then", "do", "else", "elif"})
# Wrappers whose next non-flag argument is itself the command Bash will exec.
@ -256,9 +254,7 @@ def _find_blocked_commands(command: str) -> set[str]:
tok_lower = token.lower()
# Match -c exactly, or combined flags ending in c (e.g. -lc, -xc)
is_unix_c = tok_lower == "-c" or (
tok_lower.startswith("-")
and tok_lower.endswith("c")
and not tok_lower.startswith("--")
tok_lower.startswith("-") and tok_lower.endswith("c") and not tok_lower.startswith("--")
)
is_win_c = tok_lower == "/c"
if not (is_unix_c or is_win_c) or i < 1 or i + 1 >= len(tokens):
@ -370,18 +366,11 @@ def _sandbox_preexec():
except (ValueError, OSError, AttributeError):
pass
try:
_resource.setrlimit(
_resource.RLIMIT_FSIZE, (100 * 1024 * 1024, 100 * 1024 * 1024)
)
_resource.setrlimit(_resource.RLIMIT_FSIZE, (100 * 1024 * 1024, 100 * 1024 * 1024))
except (ValueError, OSError):
pass
try:
as_bytes = (
int(os.environ.get("UNSLOTH_STUDIO_SANDBOX_AS_GB", "8"))
* 1024
* 1024
* 1024
)
as_bytes = int(os.environ.get("UNSLOTH_STUDIO_SANDBOX_AS_GB", "8")) * 1024 * 1024 * 1024
_resource.setrlimit(_resource.RLIMIT_AS, (as_bytes, as_bytes))
except (ValueError, OSError, AttributeError):
pass
@ -398,9 +387,7 @@ def _sandbox_preexec():
# value (would otherwise leave NOFILE at the parent's default).
nofile = int(os.environ.get("UNSLOTH_STUDIO_SANDBOX_NOFILE", "16384"))
_soft_cur, hard_cur = _resource.getrlimit(_resource.RLIMIT_NOFILE)
target = (
nofile if hard_cur == _resource.RLIM_INFINITY else min(nofile, hard_cur)
)
target = nofile if hard_cur == _resource.RLIM_INFINITY else min(nofile, hard_cur)
_resource.setrlimit(_resource.RLIMIT_NOFILE, (target, target))
except (ValueError, OSError, AttributeError):
pass
@ -432,12 +419,9 @@ def _get_project_workdir(session_id: str) -> str | None:
return None
try:
from storage.studio_db import ensure_chat_project_workspace
project = ensure_chat_project_workspace(project_id)
except Exception:
logger.warning(
"Failed to resolve project sandbox for %s", session_id, exc_info = True
)
logger.warning("Failed to resolve project sandbox for %s", session_id, exc_info = True)
return None
if not project:
return None
@ -468,9 +452,7 @@ def _get_workdir(session_id: str | None = None) -> str:
workdir = project_workdir
elif session_id and _SESSION_ID_RE.match(session_id):
workdir = os.path.join(sandbox_root, session_id)
if not os.path.realpath(workdir).startswith(
os.path.realpath(sandbox_root) + os.sep
):
if not os.path.realpath(workdir).startswith(os.path.realpath(sandbox_root) + os.sep):
workdir = os.path.join(sandbox_root, "_invalid")
elif session_id:
workdir = os.path.join(sandbox_root, "_invalid")
@ -618,9 +600,7 @@ def _mcp_specs_for_server(server: dict, mcp_tools: list[dict]) -> list[dict]:
# Same MCP server returning duplicate tool names would also 400
# OpenAI ("tools[N].function.name duplicates ..."). Drop dupes.
if name in seen_names:
logger.warning(
"Skipping duplicate MCP tool '%s' on '%s'.", raw_name, display
)
logger.warning("Skipping duplicate MCP tool '%s' on '%s'.", raw_name, display)
continue
seen_names.add(name)
specs.append(
@ -629,8 +609,7 @@ def _mcp_specs_for_server(server: dict, mcp_tools: list[dict]) -> list[dict]:
"function": {
"name": name,
"description": f"[{display}] {tool.get('description') or ''}".strip(),
"parameters": tool.get("inputSchema")
or {"type": "object", "properties": {}},
"parameters": tool.get("inputSchema") or {"type": "object", "properties": {}},
},
}
)
@ -708,9 +687,7 @@ def execute_tool(
unset (default) uses ``_EXEC_TIMEOUT`` (300 s).
``session_id``: optional thread/session ID for per-conversation sandbox isolation.
"""
logger.info(
f"execute_tool: name={name}, session_id={session_id}, timeout={timeout}"
)
logger.info(f"execute_tool: name={name}, session_id={session_id}, timeout={timeout}")
effective_timeout = _EXEC_TIMEOUT if timeout is _TIMEOUT_UNSET else timeout
if name == "render_html":
return _render_html_result(arguments)
@ -742,13 +719,9 @@ def execute_tool(
timeout = effective_timeout,
)
if name == "python":
return _python_exec(
arguments.get("code", ""), cancel_event, effective_timeout, session_id
)
return _python_exec(arguments.get("code", ""), cancel_event, effective_timeout, session_id)
if name == "terminal":
return _bash_exec(
arguments.get("command", ""), cancel_event, effective_timeout, session_id
)
return _bash_exec(arguments.get("command", ""), cancel_event, effective_timeout, session_id)
return f"Unknown tool: {name}"
@ -867,7 +840,9 @@ def _validate_and_resolve_host(hostname: str, port: int) -> tuple[bool, str, str
def _fetch_page_text(
url: str, max_chars: int = _MAX_PAGE_CHARS, timeout: int = 30
url: str,
max_chars: int = _MAX_PAGE_CHARS,
timeout: int = 30,
) -> str:
"""Fetch a URL and return plain text content (HTML tags stripped).
@ -921,9 +896,7 @@ def _fetch_page_text(
resp = opener.open(req, timeout = timeout)
except _HTTPError as e:
if e.code not in (301, 302, 303, 307, 308):
return (
f"Failed to fetch URL: HTTP {e.code} {getattr(e, 'reason', '')}"
)
return f"Failed to fetch URL: HTTP {e.code} {getattr(e, 'reason', '')}"
location = e.headers.get("Location")
if not location:
return "Failed to fetch URL: redirect missing Location header."
@ -1188,9 +1161,7 @@ def _check_signal_escape_patterns(code: str):
if func_name:
if func_name in ("signal.signal", "signal"):
if len(node.args) >= 1:
if _ast_name_matches(
node.args[0], ("SIGALRM", "signal.SIGALRM")
):
if _ast_name_matches(node.args[0], ("SIGALRM", "signal.SIGALRM")):
signal_tampering.append(
{
"type": "signal_handler_override",
@ -1200,9 +1171,7 @@ def _check_signal_escape_patterns(code: str):
)
elif func_name in ("signal.setitimer", "setitimer"):
if len(node.args) >= 1:
if _ast_name_matches(
node.args[0], ("ITIMER_REAL", "signal.ITIMER_REAL")
):
if _ast_name_matches(node.args[0], ("ITIMER_REAL", "signal.ITIMER_REAL")):
signal_tampering.append(
{
"type": "timer_manipulation",
@ -1255,9 +1224,7 @@ def _check_signal_escape_patterns(code: str):
else:
has_opaque_kwargs = True
cmd_kw_values = [
v for k, v in expanded_kwargs.items() if k in _CMD_KWARGS
]
cmd_kw_values = [v for k, v in expanded_kwargs.items() if k in _CMD_KWARGS]
all_call_args = list(node.args) + cmd_kw_values
blocked_in_args = _check_args_for_blocked(all_call_args)
@ -1267,9 +1234,7 @@ def _check_signal_escape_patterns(code: str):
{
"type": "shell_escape_dynamic",
"line": node.lineno,
"description": (
f"{shell_func}() called with dynamic **kwargs"
),
"description": (f"{shell_func}() called with dynamic **kwargs"),
}
)
elif blocked_in_args:
@ -1301,8 +1266,7 @@ def _check_signal_escape_patterns(code: str):
)
shell_node = expanded_kwargs.get("shell")
shell_safe = shell_node is None or (
isinstance(shell_node, ast.Constant)
and shell_node.value is False
isinstance(shell_node, ast.Constant) and shell_node.value is False
)
# Dynamic shell-exec args (chr/format/concat bypasses).
if (
@ -1315,15 +1279,10 @@ def _check_signal_escape_patterns(code: str):
if _extract_string_from_node(n) is not None:
return True
if isinstance(n, (ast.List, ast.Tuple)):
return all(
_extract_string_from_node(e) is not None
for e in n.elts
)
return all(_extract_string_from_node(e) is not None for e in n.elts)
return False
has_non_literal = any(
not _is_safe_literal(a) for a in all_call_args
)
has_non_literal = any(not _is_safe_literal(a) for a in all_call_args)
if has_non_literal:
shell_escapes.append(
{
@ -1582,9 +1541,7 @@ def _check_signal_escape_patterns(code: str):
"/etc/sudoers",
"/etc/ssh/",
)
_SENSITIVE_FILE_RE = re.compile(
r"^/proc/(?:self|\d+)/(?:environ|cmdline|task/\d+/environ)$"
)
_SENSITIVE_FILE_RE = re.compile(r"^/proc/(?:self|\d+)/(?:environ|cmdline|task/\d+/environ)$")
def _normalize_host(host: str) -> str:
if not host:
@ -1627,15 +1584,9 @@ def _check_signal_escape_patterns(code: str):
return True
if kw.arg == "data":
v = kw.value
if (
isinstance(v, ast.Call)
and isinstance(v.func, ast.Name)
and v.func.id == "open"
):
if isinstance(v, ast.Call) and isinstance(v.func, ast.Name) and v.func.id == "open":
return True
if isinstance(v, ast.Constant) and isinstance(
v.value, (bytes, bytearray)
):
if isinstance(v, ast.Constant) and isinstance(v.value, (bytes, bytearray)):
return True
return False
@ -1777,9 +1728,7 @@ def _check_signal_escape_patterns(code: str):
"""Whether the path argument resolves to a sandbox-local literal."""
if node is None:
return False
if isinstance(node, ast.Constant) and isinstance(
node.value, (bytes, bytearray)
):
if isinstance(node, ast.Constant) and isinstance(node.value, (bytes, bytearray)):
return True # inline bytes, no file access
if isinstance(node, ast.Constant) and isinstance(node.value, str):
return _is_safe_relative_path(node.value)
@ -1865,11 +1814,7 @@ def _check_signal_escape_patterns(code: str):
)
# Direct sock.connect((host, port)) bypasses the FQ-prefix branch below.
if (
isinstance(node.func, ast.Attribute)
and node.func.attr == "connect"
and node.args
):
if isinstance(node.func, ast.Attribute) and node.func.attr == "connect" and node.args:
a0 = node.args[0]
host_lit = None
if isinstance(a0, ast.Tuple) and a0.elts:
@ -1906,9 +1851,7 @@ def _check_signal_escape_patterns(code: str):
{
"type": "upload_blocked",
"line": getattr(node, "lineno", -1),
"description": (
"Blocked: file upload disallowed in sandbox"
),
"description": ("Blocked: file upload disallowed in sandbox"),
}
)
@ -2010,28 +1953,18 @@ def _check_code_safety(code: str) -> str | None:
if info.get("error"):
return None
reasons = [
item.get("description", "") for item in info.get("signal_tampering", [])
]
shell_reasons = [
item.get("description", "") for item in info.get("shell_escapes", [])
]
reasons = [item.get("description", "") for item in info.get("signal_tampering", [])]
shell_reasons = [item.get("description", "") for item in info.get("shell_escapes", [])]
exception_reasons = [
item.get("description", "") for item in info.get("exception_catching", [])
]
network_reasons = [
item.get("description", "") for item in info.get("network_calls", [])
]
network_reasons = [item.get("description", "") for item in info.get("network_calls", [])]
file_reasons = [
item.get("description", "") for item in info.get("sensitive_file_reads", [])
]
all_reasons = [
r
for r in reasons
+ shell_reasons
+ exception_reasons
+ network_reasons
+ file_reasons
for r in reasons + shell_reasons + exception_reasons + network_reasons + file_reasons
if r
]
if all_reasons:
@ -2063,7 +1996,11 @@ def _kill_process_tree(proc) -> None:
pass
def _cancel_watcher(proc, cancel_event, poll_interval = 0.2):
def _cancel_watcher(
proc,
cancel_event,
poll_interval = 0.2,
):
"""Daemon thread that kills a process when cancel_event is set."""
while proc.poll() is None:
if cancel_event is not None and cancel_event.is_set():
@ -2107,9 +2044,7 @@ def _python_exec(
except OSError:
pass
try:
fd, tmp_path = tempfile.mkstemp(
suffix = ".py", prefix = "studio_exec_", dir = workdir
)
fd, tmp_path = tempfile.mkstemp(suffix = ".py", prefix = "studio_exec_", dir = workdir)
with os.fdopen(fd, "w") as f:
f.write(code)
@ -2170,7 +2105,6 @@ def _python_exec(
new_images.append(_name)
if new_images:
import json as _json
result += f"\n__IMAGES__:{_json.dumps(sorted(new_images))}"
return result

54
studio/backend/core/inference/worker.py
View file

@ -93,9 +93,7 @@ def _build_model_config(config: dict):
return mc
def _get_hf_download_state(
model_names: list[str] | None = None,
) -> tuple[int, bool] | None:
def _get_hf_download_state(model_names: list[str] | None = None) -> tuple[int, bool] | None:
"""Return (total_bytes, has_incomplete) for the HF Hub cache, or None on error.
When *model_names* is provided, only those models' ``blobs/``
@ -124,7 +122,6 @@ def _get_hf_download_state(
if model_names:
from utils.paths import resolve_cached_repo_id_case
for name in model_names:
if not name:
continue
@ -265,14 +262,10 @@ def _handle_load(backend, config: dict, resp_queue: Any) -> None:
adapter_cfg = json.load(f)
training_method = adapter_cfg.get("unsloth_training_method")
if training_method == "lora" and load_in_4bit:
logger.info(
"adapter_config.json says lora — setting load_in_4bit=False"
)
logger.info("adapter_config.json says lora — setting load_in_4bit=False")
load_in_4bit = False
elif training_method == "qlora" and not load_in_4bit:
logger.info(
"adapter_config.json says qlora — setting load_in_4bit=True"
)
logger.info("adapter_config.json says qlora — setting load_in_4bit=True")
load_in_4bit = True
elif not training_method:
if (
@ -401,12 +394,7 @@ def _handle_load(backend, config: dict, resp_queue: Any) -> None:
)
def _handle_generate(
backend,
cmd: dict,
resp_queue: Any,
cancel_event,
) -> None:
def _handle_generate(backend, cmd: dict, resp_queue: Any, cancel_event) -> None:
"""Handle a generate command: stream tokens back via resp_queue.
cancel_event is an mp.Event shared with the parent process.
@ -504,11 +492,7 @@ def _handle_generate(
)
def _handle_generate_audio(
backend,
cmd: dict,
resp_queue: Any,
) -> None:
def _handle_generate_audio(backend, cmd: dict, resp_queue: Any) -> None:
"""Handle TTS audio generation — returns WAV bytes + sample_rate."""
request_id = cmd.get("request_id", "")
try:
@ -551,12 +535,7 @@ def _handle_generate_audio(
)
def _handle_generate_audio_input(
backend,
cmd: dict,
resp_queue: Any,
cancel_event,
) -> None:
def _handle_generate_audio_input(backend, cmd: dict, resp_queue: Any, cancel_event) -> None:
"""Handle audio input generation (ASR/Whisper) — streams text tokens back."""
request_id = cmd.get("request_id", "")
@ -591,9 +570,7 @@ def _handle_generate_audio_input(
for text_chunk in generator:
if cancel_event.is_set():
logger.info(
"Audio input generation cancelled for request %s", request_id
)
logger.info("Audio input generation cancelled for request %s", request_id)
break
_send_response(
@ -660,13 +637,7 @@ def _handle_unload(backend, cmd: dict, resp_queue: Any) -> None:
)
def run_inference_process(
*,
cmd_queue: Any,
resp_queue: Any,
cancel_event,
config: dict,
) -> None:
def run_inference_process(*, cmd_queue: Any, resp_queue: Any, cancel_event, config: dict) -> None:
"""Subprocess entrypoint. Persistent — runs command loop until shutdown.
Args:
@ -676,9 +647,7 @@ def run_inference_process(
config: Initial configuration dict with model info.
"""
os.environ["TOKENIZERS_PARALLELISM"] = "false"
os.environ["PYTHONWARNINGS"] = (
"ignore" # Suppress warnings at C-level before imports
)
os.environ["PYTHONWARNINGS"] = "ignore" # Suppress warnings at C-level before imports
if config.get("disable_xet"):
os.environ["HF_HUB_DISABLE_XET"] = "1"
@ -810,7 +779,6 @@ def run_inference_process(
if sys.platform == "win32":
try:
import triton # noqa: F401
logger.info("Triton available — torch.compile enabled")
except ImportError:
os.environ["TORCHDYNAMO_DISABLE"] = "1"
@ -986,9 +954,7 @@ def run_inference_process(
)
except Exception as exc:
logger.error(
"Error handling command '%s': %s", cmd_type, exc, exc_info = True
)
logger.error("Error handling command '%s': %s", cmd_type, exc, exc_info = True)
_send_response(
resp_queue,
{

10
studio/backend/core/tool_healing.py
View file

@ -87,9 +87,7 @@ def parse_tool_calls_from_text(content: str) -> list[dict]:
},
}
if isinstance(tc["function"]["arguments"], dict):
tc["function"]["arguments"] = json.dumps(
tc["function"]["arguments"]
)
tc["function"]["arguments"] = json.dumps(tc["function"]["arguments"])
tool_calls.append(tc)
except (json.JSONDecodeError, ValueError):
pass
@ -108,11 +106,7 @@ def parse_tool_calls_from_text(content: str) -> list[dict]:
func_name = fm.group(1)
body_start = fm.end()
# Hard boundaries: next <function= tag or </tool_call>
next_func = (
func_starts[idx + 1].start()
if idx + 1 < len(func_starts)
else len(content)
)
next_func = func_starts[idx + 1].start() if idx + 1 < len(func_starts) else len(content)
end_tag = _TC_END_TAG_RE.search(content[body_start:])
if end_tag:
body_end = body_start + end_tag.start()

530
studio/backend/core/training/trainer.py
View file

File diff suppressed because it is too large Load diff

82
studio/backend/core/training/training.py
View file

@ -180,9 +180,7 @@ class TrainingBackend:
if self._pump_thread is not None and self._pump_thread.is_alive():
self._pump_thread.join(timeout = 5.0)
if self._pump_thread.is_alive():
logger.warning(
"Previous pump thread did not exit within 5s — refusing to start"
)
logger.warning("Previous pump thread did not exit within 5s — refusing to start")
return False
self._pump_thread = None
@ -234,9 +232,7 @@ class TrainingBackend:
"train_on_completions": kwargs.get("train_on_completions", False),
"finetune_vision_layers": kwargs.get("finetune_vision_layers", True),
"finetune_language_layers": kwargs.get("finetune_language_layers", True),
"finetune_attention_modules": kwargs.get(
"finetune_attention_modules", True
),
"finetune_attention_modules": kwargs.get("finetune_attention_modules", True),
"finetune_mlp_modules": kwargs.get("finetune_mlp_modules", True),
"enable_wandb": kwargs.get("enable_wandb", False),
"wandb_token": kwargs.get("wandb_token"),
@ -321,9 +317,7 @@ class TrainingBackend:
self._run_finalized = False
self._db_run_created = False
self._db_total_steps_set = False
self._db_config = {
k: v for k, v in config.items() if k not in {"hf_token", "wandb_token"}
}
self._db_config = {k: v for k, v in config.items() if k not in {"hf_token", "wandb_token"}}
self._db_started_at = datetime.now(timezone.utc).isoformat()
# Assign subprocess handles after state reset
@ -353,9 +347,7 @@ class TrainingBackend:
pass
# Update progress immediately for responsive UI
self._progress.status_message = (
"Stopping training and saving checkpoint..."
if save
else "Cancelling training..."
"Stopping training and saving checkpoint..." if save else "Cancelling training..."
)
return True
@ -363,9 +355,7 @@ class TrainingBackend:
"""Force-kill the training subprocess so state can be reset immediately."""
with self._lock:
if self._proc is not None and self._proc.is_alive():
logger.info(
"Force-terminating training subprocess (pid=%s)", self._proc.pid
)
logger.info("Force-terminating training subprocess (pid=%s)", self._proc.pid)
self._proc.terminate()
proc = self._proc
cancelled = self._cancel_requested
@ -525,8 +515,7 @@ class TrainingBackend:
else:
self._progress.is_training = False
self._progress.error = (
self._progress.error
or "Training process exited unexpectedly"
self._progress.error or "Training process exited unexpectedly"
)
self._ensure_db_run_created()
@ -566,9 +555,7 @@ class TrainingBackend:
try:
_safe_lr = float(_raw_lr) if _raw_lr is not None else None
except (TypeError, ValueError):
logger.debug(
"Could not convert learning_rate to float: %s", _raw_lr
)
logger.debug("Could not convert learning_rate to float: %s", _raw_lr)
_safe_lr = None
if _safe_lr is not None and not math.isfinite(_safe_lr):
_safe_lr = None
@ -576,9 +563,7 @@ class TrainingBackend:
self._progress.loss = _safe_loss
if _safe_lr is not None:
self._progress.learning_rate = _safe_lr
self._progress.total_steps = event.get(
"total_steps", self._progress.total_steps
)
self._progress.total_steps = event.get("total_steps", self._progress.total_steps)
self._progress.elapsed_seconds = event.get("elapsed_seconds")
self._progress.eta_seconds = event.get("eta_seconds")
self._progress.grad_norm = event.get("grad_norm")
@ -622,9 +607,7 @@ class TrainingBackend:
try:
eval_loss = float(eval_loss)
except (TypeError, ValueError):
logger.debug(
"Could not convert eval_loss to float: %s", eval_loss
)
logger.debug("Could not convert eval_loss to float: %s", eval_loss)
eval_loss = None
if step > 0 and eval_loss is not None and math.isfinite(eval_loss):
self.eval_loss_history.append(eval_loss)
@ -654,12 +637,9 @@ class TrainingBackend:
"job_id": self.current_job_id,
"model_name": self._db_config["model_name"],
"dataset_name": self._db_config.get("hf_dataset")
or next(
iter(self._db_config.get("local_datasets") or []), "unknown"
),
or next(iter(self._db_config.get("local_datasets") or []), "unknown"),
"config_json": _json.dumps(self._db_config),
"started_at": self._db_started_at
or datetime.now(timezone.utc).isoformat(),
"started_at": self._db_started_at or datetime.now(timezone.utc).isoformat(),
"total_steps": event.get("total_steps"),
}
elif (
@ -737,10 +717,7 @@ class TrainingBackend:
elif db_action == "update_total_steps":
try:
from storage.studio_db import update_run_total_steps
update_run_total_steps(
db_action_kwargs["job_id"], db_action_kwargs["total_steps"]
)
update_run_total_steps(db_action_kwargs["job_id"], db_action_kwargs["total_steps"])
self._db_total_steps_set = True
except Exception:
logger.warning("Failed to update total_steps in DB", exc_info = True)
@ -764,15 +741,12 @@ class TrainingBackend:
model_name = self._db_config["model_name"],
dataset_name = dataset_name,
config_json = _json.dumps(self._db_config),
started_at = self._db_started_at
or datetime.now(timezone.utc).isoformat(),
started_at = self._db_started_at or datetime.now(timezone.utc).isoformat(),
total_steps = self._progress.total_steps or None,
)
self._db_run_created = True
except Exception:
logger.warning(
"Failed to create DB run record for early failure", exc_info = True
)
logger.warning("Failed to create DB run record for early failure", exc_info = True)
def _finalize_run_in_db(
self,
@ -795,10 +769,7 @@ class TrainingBackend:
ended_at = datetime.now(timezone.utc).isoformat(),
final_step = self._progress.step,
final_loss = self._progress.loss
if (
self._progress.loss is not None
and math.isfinite(self._progress.loss)
)
if (self._progress.loss is not None and math.isfinite(self._progress.loss))
else None,
duration_seconds = self._progress.elapsed_seconds,
loss_sparkline = _json.dumps(sparkline),
@ -807,17 +778,11 @@ class TrainingBackend:
)
self._run_finalized = True
except Exception:
logger.warning(
"Failed to finalize run in DB (status=%s)", status, exc_info = True
)
logger.warning("Failed to finalize run in DB (status=%s)", status, exc_info = True)
def _flush_metrics_to_db(self) -> None:
"""Flush buffered metrics to the database and update live progress."""
if (
not self._metric_buffer
or not self.current_job_id
or not self._db_run_created
):
if not self._metric_buffer or not self.current_job_id or not self._db_run_created:
return
# Cap buffer to prevent unbounded memory growth
if len(self._metric_buffer) > 500:
@ -837,10 +802,7 @@ class TrainingBackend:
id = self.current_job_id,
step = self._progress.step,
loss = self._progress.loss
if (
self._progress.loss is not None
and math.isfinite(self._progress.loss)
)
if (self._progress.loss is not None and math.isfinite(self._progress.loss))
else None,
duration_seconds = self._progress.elapsed_seconds,
)
@ -873,7 +835,9 @@ class TrainingBackend:
# ------------------------------------------------------------------
def _create_loss_plot(
self, progress: TrainingProgress, theme: str = "light"
self,
progress: TrainingProgress,
theme: str = "light",
) -> plt.Figure:
"""Create training loss plot with theme-aware styling."""
plt.close("all")
@ -956,9 +920,7 @@ class TrainingBackend:
else:
title = "Training Loss"
ax.set_title(
title, fontsize = 11, fontweight = "bold", pad = 10, color = style["text"]
)
ax.set_title(title, fontsize = 11, fontweight = "bold", pad = 10, color = style["text"])
ax.grid(True, alpha = 0.4, linestyle = "--", color = style["grid_color"])
ax.tick_params(colors = style["text"], which = "both")
ax.spines["top"].set_visible(False)

210
studio/backend/core/training/worker.py
View file

@ -40,9 +40,7 @@ from utils.wheel_utils import (
)
def _output_dir_from_resume_checkpoint(
resume_from_checkpoint: str | None,
) -> str | None:
def _output_dir_from_resume_checkpoint(resume_from_checkpoint: str | None) -> str | None:
if not resume_from_checkpoint:
return None
path = Path(resume_from_checkpoint)
@ -107,9 +105,7 @@ if sys.platform == "win32":
try:
if os.path.isdir(_default_root):
for _ver in sorted(
os.listdir(_default_root), key = _ver_key, reverse = True
):
for _ver in sorted(os.listdir(_default_root), key = _ver_key, reverse = True):
_bin = os.path.join(_default_root, _ver, "bin")
if os.path.isdir(_bin):
_candidates.append(_bin)
@ -259,9 +255,7 @@ def _install_package_wheel_first(
"(this may take several minutes)..."
)
else:
pypi_status_message = (
f"Installing {display_name} from PyPI for faster training..."
)
pypi_status_message = f"Installing {display_name} from PyPI for faster training..."
_send_status(event_queue, pypi_status_message)
@ -353,8 +347,7 @@ def _install_package_wheel_first(
)
_send_status(
event_queue,
f"{display_name} installation timed out after "
f"{_run_kwargs.get('timeout')}s",
f"{display_name} installation timed out after " f"{_run_kwargs.get('timeout')}s",
)
return False
@ -435,7 +428,6 @@ def _flash_linear_attention_importable() -> bool:
try:
import fla.modules # noqa: F401
import fla.ops.gated_delta_rule # noqa: F401
return True
except Exception as exc:
logger.warning(
@ -473,9 +465,7 @@ def _ensure_flash_linear_attention_unconditional(event_queue: Any) -> bool:
if os.getenv(_FLA_SKIP_ENV) == "1":
return False
if sys.platform == "win32":
logger.info(
"Skipping flash-linear-attention install: no prebuilt wheel for Windows"
)
logger.info("Skipping flash-linear-attention install: no prebuilt wheel for Windows")
return False
if sys.version_info < _FLA_MIN_PYTHON:
logger.info(
@ -550,9 +540,7 @@ def _ensure_flash_linear_attention_unconditional(event_queue: Any) -> bool:
)
except _sp.TimeoutExpired:
logger.warning("flash-linear-attention install timed out; continuing")
_send_status(
event_queue, "flash-linear-attention install timed out; continuing"
)
_send_status(event_queue, "flash-linear-attention install timed out; continuing")
return False
if result.returncode != 0:
@ -635,7 +623,6 @@ def _discover_fla_model_types() -> frozenset[str]:
found: set[str] = set()
try:
import transformers
models_root = Path(transformers.__file__).parent / "models"
for modeling in models_root.glob("*/modeling_*.py"):
try:
@ -665,7 +652,6 @@ def _installed_tvm_ffi_version() -> str | None:
"""Installed apache-tvm-ffi version, or None if missing/unimportable."""
try:
from importlib.metadata import version as _pkg_version
return _pkg_version("apache-tvm-ffi")
except Exception:
return None
@ -676,7 +662,6 @@ def _tilelang_importable() -> bool:
try:
import tilelang # noqa: F401
import tvm_ffi # noqa: F401
return True
except Exception as exc:
logger.warning(
@ -694,7 +679,6 @@ def _torch_has_hip() -> bool:
"""
try:
import torch as _torch
return bool(
getattr(_torch.version, "hip", None)
or "rocm" in getattr(_torch, "__version__", "").lower()
@ -734,10 +718,7 @@ def _rocm_classify_unified_memory(props: Any) -> tuple[str, bool]:
# Arch attrs absent — fall back to device-name matching.
dev_lower = (getattr(props, "name", "") or "").lower()
is_unified = (
"890m" in dev_lower
or "880m" in dev_lower
or "8060s" in dev_lower
or "8050s" in dev_lower
"890m" in dev_lower or "880m" in dev_lower or "8060s" in dev_lower or "8050s" in dev_lower
)
return gcn_arch, is_unified
@ -780,9 +761,7 @@ def _run_pip(cmd: list[str], event_queue: Any, label: str) -> bool:
_send_status(event_queue, f"{label} install timed out; continuing")
return False
if result.returncode != 0:
logger.warning(
"%s install failed (continuing without it):\n%s", label, result.stdout
)
logger.warning("%s install failed (continuing without it):\n%s", label, result.stdout)
_send_status(event_queue, f"{label} install failed; continuing")
return False
return True
@ -807,7 +786,6 @@ def _ensure_tilelang_backend_unconditional(event_queue: Any) -> bool:
return False
if not _tilelang_platform_supported():
import platform as _platform
logger.info(
"Skipping tilelang install: no prebuilt wheel for %s/%s",
sys.platform,
@ -883,9 +861,7 @@ def _ensure_tilelang_backend(event_queue: Any, model_name: str) -> None:
# UNSLOTH_STUDIO_SKIP_FAST_PATH_HOOKS=1 falls back to the legacy substring path.
def _rebind_in_already_imported_modules(
*, attr_name: str, old_obj: Any, new_obj: Any
) -> int:
def _rebind_in_already_imported_modules(*, attr_name: str, old_obj: Any, new_obj: Any) -> int:
"""Rebind `attr_name -> new_obj` in every module that already imported `old_obj`.
`from X import Y` creates a local binding that reassigning X.Y won't reach.
@ -958,9 +934,7 @@ def _install_fast_path_hooks(event_queue: Any, model_name: str) -> None:
try:
ok = bool(install_fn(event_queue))
except Exception as exc:
logger.warning(
"%s install raised: %s; falling back to torch", gate_name, exc
)
logger.warning("%s install raised: %s; falling back to torch", gate_name, exc)
ok = False
logger.info("%s hook done; available=%s", gate_name, ok)
# post_available_fn handles "gate already True but ancillary kernel broken" (e.g. tilelang
@ -969,9 +943,7 @@ def _install_fast_path_hooks(event_queue: Any, model_name: str) -> None:
try:
post_available_fn(event_queue)
except Exception as exc:
logger.warning(
"%s post-available step raised: %s; continuing", gate_name, exc
)
logger.warning("%s post-available step raised: %s; continuing", gate_name, exc)
state["installed"] = True
return ok
@ -982,9 +954,7 @@ def _install_fast_path_hooks(event_queue: Any, model_name: str) -> None:
def _fla_install(eq: Any) -> bool:
# FLA alone ~2.35x; +tilelang adds ~26%. tilelang is GDN-only (Qwen3.5 family).
if not _ensure_flash_linear_attention_unconditional(eq):
logger.info(
"FLA install did not produce an importable runtime; skipping TileLang"
)
logger.info("FLA install did not produce an importable runtime; skipping TileLang")
return False
if _model_wants_tilelang(model_name):
_ensure_tilelang_backend_unconditional(eq)
@ -999,10 +969,7 @@ def _install_fast_path_hooks(event_queue: Any, model_name: str) -> None:
# FLA already imports; repair tilelang if missing or on the broken tvm-ffi list.
if not _model_wants_tilelang(model_name):
return
if (
_installed_tvm_ffi_version() not in _TVM_FFI_BROKEN_VERSIONS
and _tilelang_importable()
):
if _installed_tvm_ffi_version() not in _TVM_FFI_BROKEN_VERSIONS and _tilelang_importable():
return
_ensure_tilelang_backend_unconditional(eq)
@ -1018,9 +985,7 @@ def _install_fast_path_hooks(event_queue: Any, model_name: str) -> None:
pypi_version = _CAUSAL_CONV1D_PACKAGE_VERSION,
filename_prefix = "causal_conv1d",
release_tag = _CAUSAL_CONV1D_RELEASE_TAG,
release_base_url = (
"https://github.com/Dao-AILab/causal-conv1d/releases/download"
),
release_base_url = ("https://github.com/Dao-AILab/causal-conv1d/releases/download"),
)
return bool(ok)
@ -1040,9 +1005,7 @@ def _install_fast_path_hooks(event_queue: Any, model_name: str) -> None:
rebound = _rebind_in_already_imported_modules(
attr_name = gate_name, old_obj = original, new_obj = wrapped
)
logger.info(
"Installed fast-path hook on %s (rebound %d modules)", gate_name, rebound
)
logger.info("Installed fast-path hook on %s (rebound %d modules)", gate_name, rebound)
def _should_try_runtime_flash_attn_install(max_seq_length: int) -> bool:
@ -1190,8 +1153,7 @@ def _normalize_mlx_studio_optimizer(value):
except KeyError:
supported = ", ".join(sorted(_MLX_STUDIO_OPTIM_MAP))
raise ValueError(
f"Unsupported optimizer for MLX training: {value!r}. "
f"Supported values: {supported}."
f"Unsupported optimizer for MLX training: {value!r}. " f"Supported values: {supported}."
)
@ -1213,9 +1175,7 @@ def _resolve_mlx_local_dataset_files(file_paths: list) -> list[str]:
all_files: list[str] = []
for dataset_file in file_paths or []:
file_path = (
dataset_file
if os.path.isabs(dataset_file)
else str(resolve_dataset_path(dataset_file))
dataset_file if os.path.isabs(dataset_file) else str(resolve_dataset_path(dataset_file))
)
file_path_obj = Path(file_path)
@ -1313,9 +1273,7 @@ def _run_mlx_training(event_queue, stop_queue, config):
raise NotImplementedError(message)
optim_name = _normalize_mlx_studio_optimizer(config.get("optim", "adamw_8bit"))
lr_scheduler_type = _normalize_mlx_studio_scheduler(
config.get("lr_scheduler_type", "linear")
)
lr_scheduler_type = _normalize_mlx_studio_scheduler(config.get("lr_scheduler_type", "linear"))
# ── 1. Load model ──
# Force text-only if the dataset is not an image dataset, even if the model
@ -1344,8 +1302,7 @@ def _run_mlx_training(event_queue, stop_queue, config):
_send(
"status",
status_message = (
"MLX vision image resize ignored for DeepSeek OCR "
"(uses fixed Gundam preset)."
"MLX vision image resize ignored for DeepSeek OCR (uses fixed Gundam preset)."
),
)
vision_image_size = None
@ -1391,15 +1348,9 @@ def _run_mlx_training(event_queue, stop_queue, config):
finetune_language = config.get("finetune_language_layers", True)
finetune_attention = config.get("finetune_attention_modules", True)
finetune_mlp = config.get("finetune_mlp_modules", True)
finetune_vision = (
config.get("finetune_vision_layers", False) if is_vlm else False
)
finetune_vision = config.get("finetune_vision_layers", False) if is_vlm else False
if (
(finetune_attention or finetune_mlp)
and not finetune_language
and not finetune_vision
):
if (finetune_attention or finetune_mlp) and not finetune_language and not finetune_vision:
finetune_language = True
peft_kwargs["finetune_language_layers"] = finetune_language
@ -1432,9 +1383,7 @@ def _run_mlx_training(event_queue, stop_queue, config):
if len(file_paths) == 1:
p = Path(file_paths[0])
if p.is_dir() and (
(p / "dataset_info.json").exists() or (p / "state.json").exists()
):
if p.is_dir() and ((p / "dataset_info.json").exists() or (p / "state.json").exists()):
return load_from_disk(str(p))
all_files = _resolve_mlx_local_dataset_files(file_paths)
if not all_files:
@ -1474,7 +1423,6 @@ def _run_mlx_training(event_queue, stop_queue, config):
format_type = config.get("format_type", "")
try:
from utils.datasets import format_and_template_dataset
def _fmt_progress(status_message = "", **_kw):
_send("status", status_message = status_message)
@ -1495,9 +1443,7 @@ def _run_mlx_training(event_queue, stop_queue, config):
)
else:
errors = vlm_info.get("errors", [])
raise ValueError(
f"VLM dataset format conversion failed: {'; '.join(errors)}"
)
raise ValueError(f"VLM dataset format conversion failed: {'; '.join(errors)}")
if eval_dataset is not None:
ev_info = format_and_template_dataset(
eval_dataset,
@ -1629,11 +1575,7 @@ def _run_mlx_training(event_queue, stop_queue, config):
)
template_name = MODEL_TO_TEMPLATE_MAPPER.get(model_name.lower())
markers = (
TEMPLATE_TO_RESPONSES_MAPPER.get(template_name)
if template_name
else None
)
markers = TEMPLATE_TO_RESPONSES_MAPPER.get(template_name) if template_name else None
if markers:
trainer = train_on_responses_only(
trainer,
@ -1725,11 +1667,7 @@ def _run_mlx_training(event_queue, stop_queue, config):
"train/tokens_per_sec": tok_s,
"train/peak_gb": peak_gb,
"train/num_tokens": num_tokens,
**(
{"train/grad_norm": grad_norm}
if grad_norm is not None
else {}
),
**({"train/grad_norm": grad_norm} if grad_norm is not None else {}),
},
step = step,
)
@ -1752,9 +1690,7 @@ def _run_mlx_training(event_queue, stop_queue, config):
_send("progress", step = step, eval_loss = eval_loss)
if wandb_run is not None:
try:
wandb_run.log(
{"eval/loss": eval_loss, "eval/perplexity": perplexity}, step = step
)
wandb_run.log({"eval/loss": eval_loss, "eval/perplexity": perplexity}, step = step)
except Exception:
pass
if tb_writer is not None:
@ -1813,12 +1749,7 @@ def _run_mlx_training(event_queue, stop_queue, config):
pass
def run_training_process(
*,
event_queue: Any,
stop_queue: Any,
config: dict,
) -> None:
def run_training_process(*, event_queue: Any, stop_queue: Any, config: dict) -> None:
"""Subprocess entrypoint. Fresh Python — no stale module state.
Args:
@ -1827,9 +1758,7 @@ def run_training_process(
config: Training configuration dict with all parameters.
"""
os.environ["TOKENIZERS_PARALLELISM"] = "false"
os.environ["PYTHONWARNINGS"] = (
"ignore" # Suppress warnings at C-level before imports
)
os.environ["PYTHONWARNINGS"] = "ignore" # Suppress warnings at C-level before imports
# Offline auto-detect: skip ~25s of HF retries per call when DNS is
# dead. Scoped to this subprocess (orchestrator spawns a fresh one).
@ -2002,7 +1931,6 @@ def run_training_process(
# do NOT override on macOS. Windows has no fork at all.
if sys.platform == "linux":
import multiprocessing as _mp
try:
_mp.set_start_method("fork", force = True)
except RuntimeError:
@ -2012,7 +1940,6 @@ def run_training_process(
if sys.platform == "win32":
try:
import triton # noqa: F401
logger.info("Triton available — torch.compile enabled")
except ImportError:
os.environ["TORCHDYNAMO_DISABLE"] = "1"
@ -2044,7 +1971,6 @@ def run_training_process(
try:
import torch.distributed as _td
for _name, _stub in _td_stubs.items():
if not hasattr(_td, _name):
setattr(_td, _name, _stub)
@ -2055,7 +1981,6 @@ def run_training_process(
sys.modules["torch.distributed"] = _td_mock
try:
import torch as _torch
_torch.distributed = _td_mock
except Exception:
pass
@ -2163,9 +2088,7 @@ def run_training_process(
# (e.g. "2.11.0+rocm7.13.0" or "2.9.0+rocmsdk20251116"); fall back
# to that string when version.hip is missing.
def _hip_ver_at_least(major: int, minor: int) -> bool:
_hip_str = getattr(
getattr(_torch_for_rocm, "version", None), "hip", None
)
_hip_str = getattr(getattr(_torch_for_rocm, "version", None), "hip", None)
if not _hip_str:
# Try the standard "+rocmX.Y.Z" embedded version first
# (e.g. "2.11.0+rocm7.13.0").
@ -2227,7 +2150,11 @@ def run_training_process(
_gm_lib = _torch_for_rocm.library.Library("aten", "IMPL")
def _grouped_mm_safe_impl(
self, mat2, offs = None, bias = None, out_dtype = None
self,
mat2,
offs = None,
bias = None,
out_dtype = None,
):
"""Python mm/bmm fallback for _grouped_mm on gfx1200 (null HIP kernel, ROCm ≤ 7.12)."""
_t = _torch_for_rocm
@ -2266,9 +2193,7 @@ def run_training_process(
if prev < self.shape[0]:
a_tail = self[prev:].contiguous()
b_tail = (
mat2[-1].contiguous()
if mat2.dim() == 3
else mat2.contiguous()
mat2[-1].contiguous() if mat2.dim() == 3 else mat2.contiguous()
)
pieces.append(_t.mm(a_tail, b_tail))
result = (
@ -2329,7 +2254,6 @@ def run_training_process(
if _hw.IS_ROCM:
try:
import torch as _torch_mem
if _torch_mem.cuda.is_available():
# Classify unified vs discrete via _rocm_classify_unified_memory.
# See that function's docstring for classification priority.
@ -2539,15 +2463,12 @@ def run_training_process(
if dataset is None or trainer.should_stop:
if trainer.should_stop:
event_queue.put(
{"type": "complete", "output_dir": None, "ts": time.time()}
)
event_queue.put({"type": "complete", "output_dir": None, "ts": time.time()})
else:
event_queue.put(
{
"type": "error",
"error": trainer.training_progress.error
or "Failed to load dataset",
"error": trainer.training_progress.error or "Failed to load dataset",
"stack": "",
"ts": time.time(),
}
@ -2561,7 +2482,6 @@ def run_training_process(
def _monitor_tqdm():
from tqdm.auto import tqdm as _tqdm_cls
while not _tqdm_stop.is_set():
for bar in list(getattr(_tqdm_cls, "_instances", set())):
try:
@ -2569,9 +2489,7 @@ def run_training_process(
desc = getattr(bar, "desc", "") or ""
if total > 0 and n > 0 and desc:
pct = min(int(n * 100 / total), 100)
_send_status(
event_queue, f"{desc.strip()} {pct}% ({n:,}/{total:,})"
)
_send_status(event_queue, f"{desc.strip()} {pct}% ({n:,}/{total:,})")
except (AttributeError, ReferenceError):
pass
_tqdm_stop.wait(3)
@ -2599,9 +2517,7 @@ def run_training_process(
)
if not success or trainer.should_stop:
if trainer.should_stop:
event_queue.put(
{"type": "complete", "output_dir": None, "ts": time.time()}
)
event_queue.put({"type": "complete", "output_dir": None, "ts": time.time()})
else:
error_msg = trainer.training_progress.error or "Failed to load model"
event_queue.put(
@ -2642,9 +2558,7 @@ def run_training_process(
lora_r = config.get("lora_r", 128),
lora_alpha = config.get("lora_alpha", 32),
lora_dropout = config.get("lora_dropout", 0.0),
use_gradient_checkpointing = config.get(
"gradient_checkpointing", "unsloth"
),
use_gradient_checkpointing = config.get("gradient_checkpointing", "unsloth"),
use_rslora = config.get("use_rslora", False),
use_loftq = config.get("use_loftq", False),
)
@ -2654,17 +2568,13 @@ def run_training_process(
use_lora = True,
finetune_vision_layers = config.get("finetune_vision_layers", True),
finetune_language_layers = config.get("finetune_language_layers", True),
finetune_attention_modules = config.get(
"finetune_attention_modules", True
),
finetune_attention_modules = config.get("finetune_attention_modules", True),
finetune_mlp_modules = config.get("finetune_mlp_modules", True),
target_modules = config.get("target_modules"),
lora_r = config.get("lora_r", 16),
lora_alpha = config.get("lora_alpha", 16),
lora_dropout = config.get("lora_dropout", 0.0),
use_gradient_checkpointing = config.get(
"gradient_checkpointing", "unsloth"
),
use_gradient_checkpointing = config.get("gradient_checkpointing", "unsloth"),
use_rslora = config.get("use_rslora", False),
use_loftq = config.get("use_loftq", False),
)
@ -2674,15 +2584,12 @@ def run_training_process(
if not success or trainer.should_stop:
if trainer.should_stop:
event_queue.put(
{"type": "complete", "output_dir": None, "ts": time.time()}
)
event_queue.put({"type": "complete", "output_dir": None, "ts": time.time()})
else:
event_queue.put(
{
"type": "error",
"error": trainer.training_progress.error
or "Failed to prepare model",
"error": trainer.training_progress.error or "Failed to prepare model",
"stack": "",
"ts": time.time(),
}
@ -2738,9 +2645,7 @@ def run_training_process(
ensure_dir(Path(tensorboard_dir))
# Start training (directly — no inner thread, we ARE the subprocess)
dataset_display = (
config.get("hf_dataset", "") or config.get("uploaded_file", "") or ""
)
dataset_display = config.get("hf_dataset", "") or config.get("uploaded_file", "") or ""
_send_status(
event_queue,
f'Training "{model_name}"'
@ -2764,9 +2669,7 @@ def run_training_process(
weight_decay = config.get("weight_decay", 0.001),
random_seed = config.get("random_seed", 3407),
packing = config.get("packing", False),
train_on_completions = False
if is_cpt
else config.get("train_on_completions", False),
train_on_completions = False if is_cpt else config.get("train_on_completions", False),
enable_wandb = config.get("enable_wandb", False),
wandb_project = config.get("wandb_project", "unsloth-training"),
wandb_token = config.get("wandb_token"),
@ -3039,9 +2942,7 @@ def _run_embedding_training(event_queue: Any, stop_queue: Any, config: dict) ->
if candidates:
all_files.extend(str(c) for c in candidates)
continue
raise ValueError(
f"No supported data files in directory: {file_path_obj}"
)
raise ValueError(f"No supported data files in directory: {file_path_obj}")
else:
all_files.append(file_path)
@ -3054,9 +2955,7 @@ def _run_embedding_training(event_queue: Any, stop_queue: Any, config: dict) ->
elif first_ext == ".parquet":
loader = "parquet"
else:
raise ValueError(
f"Unsupported local dataset format: {all_files[0]}"
)
raise ValueError(f"Unsupported local dataset format: {all_files[0]}")
dataset = load_dataset(loader, data_files = all_files, split = "train")
else:
event_queue.put(
@ -3116,9 +3015,7 @@ def _run_embedding_training(event_queue: Any, stop_queue: Any, config: dict) ->
resume_from_checkpoint
)
if not output_dir:
output_dir = str(
resolve_output_dir(f"{model_name.replace('/', '_')}_{int(time.time())}")
)
output_dir = str(resolve_output_dir(f"{model_name.replace('/', '_')}_{int(time.time())}"))
output_dir = str(resolve_output_dir(output_dir))
num_epochs = config.get("num_epochs", 2)
@ -3179,7 +3076,14 @@ def _run_embedding_training(event_queue: Any, stop_queue: Any, config: dict) ->
class _EmbeddingProgressCallback(TrainerCallback):
"""Sends training progress events to the parent process via event_queue."""
def on_log(self, args, state, control, logs = None, **kwargs):
def on_log(
self,
args,
state,
control,
logs = None,
**kwargs,
):
if not logs:
return
loss_value = logs.get("loss", logs.get("train_loss", None))

69
studio/backend/main.py
View file

@ -45,9 +45,7 @@ if sys.platform == "win32":
try:
if os.path.isdir(_default_root):
for _ver in sorted(
os.listdir(_default_root), key = _ver_key, reverse = True
):
for _ver in sorted(os.listdir(_default_root), key = _ver_key, reverse = True):
_bin = os.path.join(_default_root, _ver, "bin")
if os.path.isdir(_bin):
candidates.append(_bin)
@ -83,7 +81,6 @@ if sys.platform == "win32":
_found_rocm_bnb = False
try:
import importlib.util as _ilu
_bnb_spec = _ilu.find_spec("bitsandbytes")
# submodule_search_locations (not spec.origin) handles editable installs.
if _bnb_spec and _bnb_spec.submodule_search_locations:
@ -91,9 +88,7 @@ if sys.platform == "win32":
_all_vers_main: list[str] = []
for _pkg_dir in _bnb_spec.submodule_search_locations:
for _dll in _glob.glob(
os.path.join(_pkg_dir, "libbitsandbytes_rocm*.dll")
):
for _dll in _glob.glob(os.path.join(_pkg_dir, "libbitsandbytes_rocm*.dll")):
_found_rocm_bnb = True
_km = _re_bnb.search(
r"libbitsandbytes_rocm(\d+)\.dll", os.path.basename(_dll)
@ -129,9 +124,7 @@ try:
configure_cpu_threads()
except ValueError as exc:
_raw = os.environ.get("UNSLOTH_CPU_THREADS")
raise SystemExit(
f"Error: Invalid UNSLOTH_CPU_THREADS value {_raw!r}: {exc}"
) from None
raise SystemExit(f"Error: Invalid UNSLOTH_CPU_THREADS value {_raw!r}: {exc}") from None
# Fix for Anaconda/conda-forge Python: seed platform._sys_version_cache before
# any library imports that trigger attrs -> rich -> structlog -> platform crash.
@ -183,9 +176,7 @@ def _read_studio_install_id() -> str:
information for callers reaching /api/health (relevant when Studio
is run with -H 0.0.0.0)."""
try:
token = (
(_STUDIO_ROOT_RESOLVED / "share" / "studio_install_id").read_text().strip()
)
token = (_STUDIO_ROOT_RESOLVED / "share" / "studio_install_id").read_text().strip()
except (OSError, ValueError):
return ""
return token if _STUDIO_INSTALL_ID_RE.fullmatch(token) else ""
@ -269,9 +260,7 @@ def get_unsloth_version() -> str:
except PackageNotFoundError:
pass
version_file = (
_Path(__file__).resolve().parents[2] / "unsloth" / "models" / "_utils.py"
)
version_file = _Path(__file__).resolve().parents[2] / "unsloth" / "models" / "_utils.py"
try:
for line in version_file.read_text(encoding = "utf-8").splitlines():
if line.startswith("__version__ = "):
@ -355,10 +344,7 @@ async def lifespan(app: FastAPI):
print(f"WARNING: {_msg}", flush = True)
except Exception as _probe_exc:
import structlog as _structlog
_structlog.get_logger(__name__).debug(
"llama.cpp startup probes failed: %s", _probe_exc
)
_structlog.get_logger(__name__).debug("llama.cpp startup probes failed: %s", _probe_exc)
from storage.studio_db import cleanup_orphaned_runs
@ -366,10 +352,7 @@ async def lifespan(app: FastAPI):
cleanup_orphaned_runs()
except Exception as exc:
import structlog
structlog.get_logger(__name__).warning(
"cleanup_orphaned_runs failed at startup: %s", exc
)
structlog.get_logger(__name__).warning("cleanup_orphaned_runs failed at startup: %s", exc)
# Pre-cache the helper GGUF model for LLM-assisted dataset detection.
# Runs in a background thread so it doesn't block server startup.
@ -378,7 +361,6 @@ async def lifespan(app: FastAPI):
def _precache():
try:
from utils.datasets.llm_assist import precache_helper_gguf
precache_helper_gguf()
except Exception:
pass # non-critical
@ -477,9 +459,7 @@ def _build_csp(script_nonce: "str | None" = None) -> str:
"https://*.googleusercontent.com wss://*.googleusercontent.com"
)
else:
connect_src = (
"'self' https://huggingface.co https://datasets-server.huggingface.co"
)
connect_src = "'self' https://huggingface.co https://datasets-server.huggingface.co"
return (
"default-src 'self'; "
@ -581,11 +561,7 @@ async def _send_411(send) -> None:
async def _send_413(send, total_bytes: int, max_bytes: int) -> None:
payload = _json_for_413.dumps(
{
"detail": (
f"Request body too large ({total_bytes:,} bytes; max {max_bytes:,})."
)
},
{"detail": (f"Request body too large ({total_bytes:,} bytes; max {max_bytes:,}).")},
).encode("utf-8")
await send(
{
@ -768,9 +744,7 @@ app.include_router(mcp_servers_router, prefix = "/api/mcp/servers", tags = ["mcp
app.include_router(datasets_router, prefix = "/api/datasets", tags = ["datasets"])
app.include_router(data_recipe_router, prefix = "/api/data-recipe", tags = ["data-recipe"])
app.include_router(export_router, prefix = "/api/export", tags = ["export"])
app.include_router(
training_history_router, prefix = "/api/train", tags = ["training-history"]
)
app.include_router(training_history_router, prefix = "/api/train", tags = ["training-history"])
# ============ Health and System Endpoints ============
@ -808,9 +782,7 @@ async def health_check(request: Request):
from auth.authentication import get_current_subject as _gcs
from fastapi.security import HTTPAuthorizationCredentials
creds = HTTPAuthorizationCredentials(
scheme = "Bearer", credentials = auth.split(" ", 1)[1]
)
creds = HTTPAuthorizationCredentials(scheme = "Bearer", credentials = auth.split(" ", 1)[1])
# Must await: a bare coroutine is truthy and would skip the auth check.
subject = await _gcs(creds)
except HTTPException:
@ -843,10 +815,7 @@ def studio_update_status(_current_subject: str = Depends(get_current_subject)):
@app.post("/api/shutdown")
async def shutdown_server(
request: Request,
current_subject: str = Depends(get_current_subject),
):
async def shutdown_server(request: Request, current_subject: str = Depends(get_current_subject)):
"""Gracefully shut down the Unsloth Studio server.
Called by the frontend quit dialog so users can stop the server from the UI
@ -863,7 +832,6 @@ async def shutdown_server(
# Fallback when not launched via run_server() (e.g. direct uvicorn)
import signal
import os
os.kill(os.getpid(), signal.SIGTERM)
request.app.state._shutdown_task = asyncio.create_task(_delayed_shutdown())
@ -871,9 +839,7 @@ async def shutdown_server(
@app.get("/api/system")
async def get_system_info(
current_subject: str = Depends(get_current_subject),
):
async def get_system_info(current_subject: str = Depends(get_current_subject)):
"""Get system information.
Gated behind auth: the response includes platform, Python version,
@ -914,23 +880,18 @@ async def get_system_info(
@app.get("/api/system/gpu-visibility")
async def get_gpu_visibility(
current_subject: str = Depends(get_current_subject),
):
async def get_gpu_visibility(current_subject: str = Depends(get_current_subject)):
return get_backend_visible_gpu_info()
@app.get("/api/system/hardware")
async def get_hardware_info(
current_subject: str = Depends(get_current_subject),
):
async def get_hardware_info(current_subject: str = Depends(get_current_subject)):
"""Return GPU name, total VRAM, and key ML package versions.
Gated behind auth alongside /api/system -- same fingerprinting
concern. /api/system/gpu-visibility is also auth-gated already.
"""
from utils.hardware import get_gpu_summary, get_package_versions
return {
"gpu": get_gpu_summary(),
"versions": get_package_versions(),

12
studio/backend/models/auth.py
View file

@ -26,17 +26,13 @@ class DesktopLoginRequest(BaseModel):
class RefreshTokenRequest(BaseModel):
"""Refresh token payload to obtain new access + refresh tokens."""
refresh_token: str = Field(
..., description = "Refresh token from a previous login or refresh"
)
refresh_token: str = Field(..., description = "Refresh token from a previous login or refresh")
class AuthStatusResponse(BaseModel):
"""Indicate whether the seeded admin auth flow is ready."""
initialized: bool = Field(
..., description = "True if the auth database contains a login user"
)
initialized: bool = Field(..., description = "True if the auth database contains a login user")
default_username: str = Field(
"unsloth",
description = "Default admin username for first-boot UI prefill.",
@ -77,9 +73,7 @@ class ApiKeyResponse(BaseModel):
id: int
name: str
key_prefix: str = Field(
..., description = "First 8 characters after sk-unsloth- for display"
)
key_prefix: str = Field(..., description = "First 8 characters after sk-unsloth- for display")
created_at: str
last_used_at: Optional[str] = None
expires_at: Optional[str] = None

4
studio/backend/models/data_recipe.py
View file

@ -103,9 +103,7 @@ class SeedInspectUploadRequest(BaseModel):
if not self.block_id:
raise ValueError("block_id is required when using file_ids")
if self.file_names is None or len(self.file_ids) != len(self.file_names):
raise ValueError(
"file_names must be provided and same length as file_ids"
)
raise ValueError("file_names must be provided and same length as file_ids")
if has_legacy:
if not self.filename:
raise ValueError("filename is required when using content_base64")

128
studio/backend/models/inference.py
View file

@ -28,9 +28,7 @@ class LoadRequest(BaseModel):
native_path_lease: Optional[str] = Field(
None, description = "Frontend-visible signed native path grant"
)
hf_token: Optional[str] = Field(
None, description = "HuggingFace token for gated models"
)
hf_token: Optional[str] = Field(None, description = "HuggingFace token for gated models")
max_seq_length: int = Field(
0,
ge = 0,
@ -53,9 +51,7 @@ class LoadRequest(BaseModel):
@field_validator("chat_template_override")
@classmethod
def normalize_blank_chat_template_override(
cls, value: Optional[str]
) -> Optional[str]:
def normalize_blank_chat_template_override(cls, value: Optional[str]) -> Optional[str]:
if value is not None and value.strip() == "":
return None
return value
@ -123,9 +119,7 @@ class ValidateModelRequest(BaseModel):
native_path_lease: Optional[str] = Field(
None, description = "Frontend-visible signed native path grant"
)
hf_token: Optional[str] = Field(
None, description = "HuggingFace token for gated models"
)
hf_token: Optional[str] = Field(None, description = "HuggingFace token for gated models")
gguf_variant: Optional[str] = Field(
None, description = "GGUF quantization variant (e.g. 'Q4_K_M')"
)
@ -142,9 +136,7 @@ class ValidateModelResponse(BaseModel):
valid: bool = Field(..., description = "Whether the model identifier looks valid")
message: str = Field(..., description = "Human-readable validation message")
identifier: Optional[str] = Field(None, description = "Resolved model identifier")
display_name: Optional[str] = Field(
None, description = "Display name derived from identifier"
)
display_name: Optional[str] = Field(None, description = "Display name derived from identifier")
is_gguf: bool = Field(False, description = "Whether this is a GGUF model (llama.cpp)")
is_lora: bool = Field(False, description = "Whether this is a LoRA adapter")
is_vision: bool = Field(False, description = "Whether this is a vision-capable model")
@ -162,16 +154,10 @@ class GenerateRequest(BaseModel):
temperature: float = Field(0.6, ge = 0.0, le = 2.0, description = "Sampling temperature")
top_p: float = Field(0.95, ge = 0.0, le = 1.0, description = "Top-p sampling")
top_k: int = Field(20, ge = -1, le = 100, description = "Top-k sampling")
max_new_tokens: int = Field(
2048, ge = 1, le = 4096, description = "Maximum tokens to generate"
)
repetition_penalty: float = Field(
1.0, ge = 1.0, le = 2.0, description = "Repetition penalty"
)
max_new_tokens: int = Field(2048, ge = 1, le = 4096, description = "Maximum tokens to generate")
repetition_penalty: float = Field(1.0, ge = 1.0, le = 2.0, description = "Repetition penalty")
presence_penalty: float = Field(0.0, ge = 0.0, le = 2.0, description = "Presence penalty")
image_base64: Optional[str] = Field(
None, description = "Base64 encoded image for vision models"
)
image_base64: Optional[str] = Field(None, description = "Base64 encoded image for vision models")
class LoadResponse(BaseModel):
@ -182,16 +168,10 @@ class LoadResponse(BaseModel):
display_name: str = Field(..., description = "Display name of the model")
is_vision: bool = Field(False, description = "Whether model is a vision model")
is_lora: bool = Field(False, description = "Whether model is a LoRA adapter")
is_gguf: bool = Field(
False, description = "Whether model is a GGUF model (llama.cpp)"
)
is_gguf: bool = Field(False, description = "Whether model is a GGUF model (llama.cpp)")
is_audio: bool = Field(False, description = "Whether model is a TTS audio model")
audio_type: Optional[str] = Field(
None, description = "Audio codec type: snac, csm, bicodec, dac"
)
has_audio_input: bool = Field(
False, description = "Whether model accepts audio input (ASR)"
)
audio_type: Optional[str] = Field(None, description = "Audio codec type: snac, csm, bicodec, dac")
has_audio_input: bool = Field(False, description = "Whether model accepts audio input (ASR)")
inference: dict = Field(
..., description = "Inference parameters (temperature, top_p, top_k, min_p)"
)
@ -282,17 +262,14 @@ class LoadProgressResponse(BaseModel):
bytes_loaded: int = Field(
0,
description = (
"Bytes of the model already resident in the llama-server "
"process (VmRSS on Linux)."
"Bytes of the model already resident in the llama-server process (VmRSS on Linux)."
),
)
bytes_total: int = Field(
0,
description = "Total bytes across all GGUF shards for the active model.",
)
fraction: float = Field(
0.0, description = "bytes_loaded / bytes_total, clamped to 0..1."
)
fraction: float = Field(0.0, description = "bytes_loaded / bytes_total, clamped to 0..1.")
class InferenceStatusResponse(BaseModel):
@ -305,30 +282,14 @@ class InferenceStatusResponse(BaseModel):
None,
description = "Loadable identifier for the active model.",
)
is_vision: bool = Field(
False, description = "Whether the active model is a vision model"
)
is_gguf: bool = Field(
False, description = "Whether the active model is a GGUF model (llama.cpp)"
)
gguf_variant: Optional[str] = Field(
None, description = "GGUF quantization variant (e.g. Q4_K_M)"
)
is_audio: bool = Field(
False, description = "Whether the active model is a TTS audio model"
)
audio_type: Optional[str] = Field(
None, description = "Audio codec type: snac, csm, bicodec, dac"
)
has_audio_input: bool = Field(
False, description = "Whether model accepts audio input (ASR)"
)
loading: List[str] = Field(
default_factory = list, description = "Models currently being loaded"
)
loaded: List[str] = Field(
default_factory = list, description = "Models currently loaded"
)
is_vision: bool = Field(False, description = "Whether the active model is a vision model")
is_gguf: bool = Field(False, description = "Whether the active model is a GGUF model (llama.cpp)")
gguf_variant: Optional[str] = Field(None, description = "GGUF quantization variant (e.g. Q4_K_M)")
is_audio: bool = Field(False, description = "Whether the active model is a TTS audio model")
audio_type: Optional[str] = Field(None, description = "Audio codec type: snac, csm, bicodec, dac")
has_audio_input: bool = Field(False, description = "Whether model accepts audio input (ASR)")
loading: List[str] = Field(default_factory = list, description = "Models currently being loaded")
loaded: List[str] = Field(default_factory = list, description = "Models currently loaded")
inference: Optional[Dict[str, Any]] = Field(
None, description = "Recommended inference parameters for the active model"
)
@ -353,9 +314,7 @@ class InferenceStatusResponse(BaseModel):
supports_tools: bool = Field(
False, description = "Whether the active model supports tool calling"
)
context_length: Optional[int] = Field(
None, description = "Context length of the active model"
)
context_length: Optional[int] = Field(None, description = "Context length of the active model")
max_context_length: Optional[int] = Field(
None,
description = "Maximum context length currently available for the active model",
@ -563,9 +522,7 @@ class ChatMessage(BaseModel):
``ChatCompletionRequest`` layer by walking back to the preceding assistant.
"""
role: Literal["system", "user", "assistant", "tool"] = Field(
..., description = "Message role"
)
role: Literal["system", "user", "assistant", "tool"] = Field(..., description = "Message role")
content: Optional[Union[str, list[ContentPart]]] = Field(
None, description = "Message content (string or multimodal parts)"
)
@ -666,9 +623,7 @@ class ChatCompletionRequest(BaseModel):
# ── Unsloth extensions (ignored by standard OpenAI clients) ──
top_k: int = Field(20, ge = -1, le = 100, description = "[x-unsloth] Top-k sampling")
min_p: float = Field(
0.01, ge = 0.0, le = 1.0, description = "[x-unsloth] Min-p sampling threshold"
)
min_p: float = Field(0.01, ge = 0.0, le = 1.0, description = "[x-unsloth] Min-p sampling threshold")
repetition_penalty: float = Field(
1.0, ge = 1.0, le = 2.0, description = "[x-unsloth] Repetition penalty"
)
@ -925,9 +880,7 @@ class ChatCompletionRequest(BaseModel):
if not tc_id:
continue
function = tc.get("function")
function_name = (
function.get("name") if isinstance(function, dict) else None
)
function_name = function.get("name") if isinstance(function, dict) else None
if msg.name and function_name == msg.name:
name_match = (tc_id, asst_idx, tc_idx)
break
@ -940,7 +893,6 @@ class ChatCompletionRequest(BaseModel):
break
if picked is None:
import secrets as _secrets
picked = f"call_{_secrets.token_hex(8)}"
msg.tool_call_id = picked
return self
@ -1160,17 +1112,13 @@ class ResponsesFunctionCallInputItem(BaseModel):
"""
type: Literal["function_call"]
id: Optional[str] = Field(
None, description = "Item id assigned by the server (e.g. fc_...)"
)
id: Optional[str] = Field(None, description = "Item id assigned by the server (e.g. fc_...)")
call_id: str = Field(
...,
description = "Correlation id matching a function_call_output on the next turn.",
)
name: str
arguments: str = Field(
..., description = "JSON string of the arguments the model produced."
)
arguments: str = Field(..., description = "JSON string of the arguments the model produced.")
status: Optional[Literal["in_progress", "completed", "incomplete"]] = None
@ -1266,9 +1214,7 @@ class ResponsesRequest(BaseModel):
default = [],
description = "Input text or list of messages / function_call / function_call_output items",
)
instructions: Optional[str] = Field(
None, description = "System / developer instructions"
)
instructions: Optional[str] = Field(None, description = "System / developer instructions")
temperature: Optional[float] = Field(None, ge = 0.0, le = 2.0)
top_p: Optional[float] = Field(None, ge = 0.0, le = 1.0)
max_output_tokens: Optional[int] = Field(None, ge = 1)
@ -1344,9 +1290,7 @@ class ResponsesOutputFunctionCall(BaseModel):
id: str = Field(default_factory = lambda: f"fc_{uuid.uuid4().hex[:12]}")
call_id: str
name: str
arguments: str = Field(
..., description = "JSON string of the arguments the model produced."
)
arguments: str = Field(..., description = "JSON string of the arguments the model produced.")
status: Literal["completed", "in_progress", "incomplete"] = "completed"
@ -1455,16 +1399,12 @@ def _merge_anthropic_system(system: Any, additions: list[str]) -> Any:
if not additions:
return system
addition_blocks = [
{"type": "text", "text": text} for text in additions if text.strip()
]
addition_blocks = [{"type": "text", "text": text} for text in additions if text.strip()]
if not addition_blocks:
return system
if system is None:
return (
addition_blocks[0]["text"] if len(addition_blocks) == 1 else addition_blocks
)
return addition_blocks[0]["text"] if len(addition_blocks) == 1 else addition_blocks
if isinstance(system, str):
return "\n\n".join([system, *[block["text"] for block in addition_blocks]])
if isinstance(system, list):
@ -1543,9 +1483,7 @@ class AnthropicMessagesRequest(BaseModel):
normalized = dict(data)
normalized["messages"] = normalized_messages
normalized["system"] = _merge_anthropic_system(
normalized.get("system"), system_additions
)
normalized["system"] = _merge_anthropic_system(normalized.get("system"), system_additions)
return normalized
@ -1569,9 +1507,7 @@ class AnthropicResponseToolUseBlock(BaseModel):
input: dict
AnthropicResponseBlock = Union[
AnthropicResponseTextBlock, AnthropicResponseToolUseBlock
]
AnthropicResponseBlock = Union[AnthropicResponseTextBlock, AnthropicResponseToolUseBlock]
class AnthropicMessagesResponse(BaseModel):

52
studio/backend/models/models.py
View file

@ -14,9 +14,7 @@ ModelType = Literal["text", "vision", "audio", "embeddings"]
class CheckpointInfo(BaseModel):
"""Information about a discovered checkpoint directory."""
display_name: str = Field(
..., description = "User-friendly checkpoint name (folder name)"
)
display_name: str = Field(..., description = "User-friendly checkpoint name (folder name)")
path: str = Field(..., description = "Full path to the checkpoint directory")
loss: Optional[float] = Field(None, description = "Training loss at this checkpoint")
@ -65,33 +63,23 @@ class ModelDetails(BaseModel):
None, description = "Model identifier (alias for id, for backward compatibility)"
)
name: Optional[str] = Field(None, description = "Display name for the model")
config: Optional[Dict[str, Any]] = Field(
None, description = "Model configuration dictionary"
)
config: Optional[Dict[str, Any]] = Field(None, description = "Model configuration dictionary")
is_vision: bool = Field(False, description = "Whether model is a vision model")
is_embedding: bool = Field(
False, description = "Whether model is an embedding/sentence-transformer model"
)
is_lora: bool = Field(False, description = "Whether model is a LoRA adapter")
is_gguf: bool = Field(
False, description = "Whether model is a GGUF model (llama.cpp format)"
)
is_gguf: bool = Field(False, description = "Whether model is a GGUF model (llama.cpp format)")
is_mlx: bool = Field(
False, description = "Whether model is served via the MLX backend (Apple Silicon)"
)
is_audio: bool = Field(False, description = "Whether model is a TTS audio model")
audio_type: Optional[str] = Field(
None, description = "Audio codec type: snac, csm, bicodec, dac"
)
has_audio_input: bool = Field(
False, description = "Whether model accepts audio input (ASR)"
)
audio_type: Optional[str] = Field(None, description = "Audio codec type: snac, csm, bicodec, dac")
has_audio_input: bool = Field(False, description = "Whether model accepts audio input (ASR)")
model_type: Optional[ModelType] = Field(
None, description = "Collapsed model modality: text, vision, audio, or embeddings"
)
base_model: Optional[str] = Field(
None, description = "Base model if this is a LoRA adapter"
)
base_model: Optional[str] = Field(None, description = "Base model if this is a LoRA adapter")
max_position_embeddings: Optional[int] = Field(
None, description = "Maximum context length supported by the model"
)
@ -104,9 +92,7 @@ class LoRAInfo(BaseModel):
"""LoRA adapter or exported model information"""
display_name: str = Field(..., description = "Display name for the LoRA")
adapter_path: str = Field(
..., description = "Path to the LoRA adapter or exported model"
)
adapter_path: str = Field(..., description = "Path to the LoRA adapter or exported model")
base_model: Optional[str] = Field(None, description = "Base model identifier")
source: Optional[str] = Field(None, description = "'training' or 'exported'")
export_type: Optional[str] = Field(
@ -117,29 +103,21 @@ class LoRAInfo(BaseModel):
class LoRAScanResponse(BaseModel):
"""Response schema for scanning trained LoRA adapters"""
loras: List[LoRAInfo] = Field(
default_factory = list, description = "List of found LoRA adapters"
)
loras: List[LoRAInfo] = Field(default_factory = list, description = "List of found LoRA adapters")
outputs_dir: str = Field(..., description = "Directory that was scanned")
class ModelListResponse(BaseModel):
"""Response schema for listing models"""
models: List[ModelDetails] = Field(
default_factory = list, description = "List of models"
)
default_models: List[str] = Field(
default_factory = list, description = "List of default model IDs"
)
models: List[ModelDetails] = Field(default_factory = list, description = "List of models")
default_models: List[str] = Field(default_factory = list, description = "List of default model IDs")
class GgufVariantDetail(BaseModel):
"""A single GGUF quantization variant in a HuggingFace repo."""
filename: str = Field(
..., description = "GGUF filename (e.g., 'gemma-3-4b-it-Q4_K_M.gguf')"
)
filename: str = Field(..., description = "GGUF filename (e.g., 'gemma-3-4b-it-Q4_K_M.gguf')")
quant: str = Field(..., description = "Quantization label (e.g., 'Q4_K_M')")
size_bytes: int = Field(0, description = "File size in bytes")
downloaded: bool = Field(
@ -185,9 +163,7 @@ class LocalModelInfo(BaseModel):
class LocalModelListResponse(BaseModel):
"""Response schema for listing local/cached models."""
models_dir: str = Field(
..., description = "Directory scanned for custom local models"
)
models_dir: str = Field(..., description = "Directory scanned for custom local models")
hf_cache_dir: Optional[str] = Field(
None,
description = "HF cache root that was scanned",
@ -205,9 +181,7 @@ class LocalModelListResponse(BaseModel):
class AddScanFolderRequest(BaseModel):
"""Request body for adding a custom scan folder."""
path: str = Field(
..., description = "Absolute or relative directory path to scan for models"
)
path: str = Field(..., description = "Absolute or relative directory path to scan for models")
class ScanFolderInfo(BaseModel):

16
studio/backend/models/providers.py
View file

@ -16,9 +16,7 @@ from pydantic import BaseModel, Field
class ProviderRegistryEntry(BaseModel):
"""A supported provider type with its default configuration."""
provider_type: str = Field(
..., description = "Provider identifier (e.g. 'openai', 'mistral')"
)
provider_type: str = Field(..., description = "Provider identifier (e.g. 'openai', 'mistral')")
display_name: str = Field(..., description = "Human-readable provider name")
base_url: str = Field(..., description = "Default API base URL")
default_models: list[str] = Field(
@ -46,9 +44,7 @@ class ProviderCreate(BaseModel):
"""Request to create a saved provider configuration."""
provider_type: str = Field(..., description = "Provider type from the registry")
display_name: str = Field(
..., description = "User-chosen label (e.g. 'My OpenAI Key')"
)
display_name: str = Field(..., description = "User-chosen label (e.g. 'My OpenAI Key')")
base_url: Optional[str] = Field(
None,
description = "Custom base URL (overrides registry default). Omit to use the default.",
@ -60,9 +56,7 @@ class ProviderUpdate(BaseModel):
display_name: Optional[str] = Field(None, description = "New display name")
base_url: Optional[str] = Field(None, description = "New base URL")
is_enabled: Optional[bool] = Field(
None, description = "Enable or disable this provider"
)
is_enabled: Optional[bool] = Field(None, description = "Enable or disable this provider")
class ProviderResponse(BaseModel):
@ -85,9 +79,7 @@ class ProviderModelInfo(BaseModel):
id: str = Field(..., description = "Model ID as expected by the provider API")
display_name: str = Field("", description = "Human-readable model name")
context_length: Optional[int] = Field(
None, description = "Maximum context length in tokens"
)
context_length: Optional[int] = Field(None, description = "Maximum context length in tokens")
owned_by: Optional[str] = Field(None, description = "Model owner/organization")

12
studio/backend/models/responses.py
View file

@ -23,16 +23,10 @@ class TrainingStopResponse(BaseModel):
class TrainingMetricsResponse(BaseModel):
"""Response for training metrics history"""
loss_history: List[float] = Field(
default_factory = list, description = "Loss values per step"
)
lr_history: List[float] = Field(
default_factory = list, description = "Learning rate per step"
)
loss_history: List[float] = Field(default_factory = list, description = "Loss values per step")
lr_history: List[float] = Field(default_factory = list, description = "Learning rate per step")
step_history: List[int] = Field(default_factory = list, description = "Step numbers")
grad_norm_history: List[float] = Field(
default_factory = list, description = "Gradient norm values"
)
grad_norm_history: List[float] = Field(default_factory = list, description = "Gradient norm values")
grad_norm_step_history: List[int] = Field(
default_factory = list, description = "Step numbers for gradient norm values"
)

95
studio/backend/models/training.py
View file

@ -41,9 +41,7 @@ def _parse_lr(v: Any) -> float:
except (TypeError, ValueError):
raise ValueError(f"learning_rate must be parseable as float (got {v!r})")
if not (lr > 0.0):
raise ValueError(
f"learning_rate must be > 0 (got {lr!r}); " "typical range is 1e-6 .. 1e-3"
)
raise ValueError(f"learning_rate must be > 0 (got {lr!r}); typical range is 1e-6 .. 1e-3")
if lr >= _MAX_LR_VALUE:
raise ValueError(
f"learning_rate must be < 1.0 (got {lr!r}); "
@ -59,11 +57,9 @@ class TrainingStartRequest(BaseModel):
model_name: str = Field(
..., description = "Model identifier (e.g., 'unsloth/llama-3-8b-bnb-4bit')"
)
training_type: Literal["LoRA/QLoRA", "Full Finetuning", "Continued Pretraining"] = (
Field(
...,
description = "Training type: 'LoRA/QLoRA', 'Full Finetuning', or 'Continued Pretraining'",
)
training_type: Literal["LoRA/QLoRA", "Full Finetuning", "Continued Pretraining"] = Field(
...,
description = "Training type: 'LoRA/QLoRA', 'Full Finetuning', or 'Continued Pretraining'",
)
hf_token: Optional[str] = Field(None, description = "HuggingFace token")
load_in_4bit: bool = Field(True, description = "Load model in 4-bit quantization")
@ -78,9 +74,7 @@ class TrainingStartRequest(BaseModel):
)
# Dataset parameters
hf_dataset: Optional[str] = Field(
None, description = "HuggingFace dataset identifier"
)
hf_dataset: Optional[str] = Field(None, description = "HuggingFace dataset identifier")
local_datasets: List[str] = Field(
default_factory = list, description = "List of local dataset paths"
)
@ -90,12 +84,8 @@ class TrainingStartRequest(BaseModel):
format_type: str = Field(..., description = "Dataset format type")
subset: Optional[str] = None
train_split: Optional[str] = Field("train", description = "Training split name")
eval_split: Optional[str] = Field(
None, description = "Eval split name. None = auto-detect"
)
eval_steps: float = Field(
0.00, description = "Fraction of total steps between evals (0-1)"
)
eval_split: Optional[str] = Field(None, description = "Eval split name. None = auto-detect")
eval_steps: float = Field(0.00, description = "Fraction of total steps between evals (0-1)")
dataset_slice_start: Optional[int] = Field(
None, description = "Inclusive start row index for dataset slicing"
)
@ -124,9 +114,7 @@ class TrainingStartRequest(BaseModel):
if v is None:
raise ValueError("batch_size is required")
if v < 1 or v > _MAX_BATCH_SIZE:
raise ValueError(
f"batch_size must be in [1, {_MAX_BATCH_SIZE}] (got {v!r})"
)
raise ValueError(f"batch_size must be in [1, {_MAX_BATCH_SIZE}] (got {v!r})")
return v
@field_validator("gradient_accumulation_steps")
@ -136,8 +124,7 @@ class TrainingStartRequest(BaseModel):
return 1
if v < 1 or v > _MAX_GRAD_ACCUM:
raise ValueError(
f"gradient_accumulation_steps must be in [1, {_MAX_GRAD_ACCUM}] "
f"(got {v!r})"
f"gradient_accumulation_steps must be in [1, {_MAX_GRAD_ACCUM}] " f"(got {v!r})"
)
return v
@ -159,18 +146,14 @@ class TrainingStartRequest(BaseModel):
if v is None:
return v
if not isinstance(v, int) or v < 0 or v > _MAX_STEPS:
raise ValueError(
f"max_steps must be a non-negative int <= {_MAX_STEPS} (got {v!r})"
)
raise ValueError(f"max_steps must be a non-negative int <= {_MAX_STEPS} (got {v!r})")
return v
@field_validator("max_seq_length")
@classmethod
def _check_max_seq_length(cls, v: int) -> int:
if v is None or v < 1 or v > _MAX_SEQ_LENGTH:
raise ValueError(
f"max_seq_length must be in [1, {_MAX_SEQ_LENGTH}] (got {v!r})"
)
raise ValueError(f"max_seq_length must be in [1, {_MAX_SEQ_LENGTH}] (got {v!r})")
return v
@field_validator("vision_image_size", mode = "before")
@ -191,7 +174,6 @@ class TrainingStartRequest(BaseModel):
# numpy ints / Integral subclasses, without a hard numpy import.
try:
import numbers
if isinstance(v, numbers.Integral):
coerced = int(v)
elif isinstance(v, numbers.Real) and float(v).is_integer():
@ -214,8 +196,7 @@ class TrainingStartRequest(BaseModel):
return v
if not isinstance(v, int) or v < 0 or v > _MAX_STEPS:
raise ValueError(
f"warmup_steps must be a non-negative int <= {_MAX_STEPS} "
f"(got {v!r})"
f"warmup_steps must be a non-negative int <= {_MAX_STEPS} " f"(got {v!r})"
)
return v
@ -251,9 +232,7 @@ class TrainingStartRequest(BaseModel):
except (TypeError, ValueError):
raise ValueError(f"weight_decay must be a number (got {v!r})")
if wd < 0 or wd > 10.0:
raise ValueError(
f"weight_decay must be in [0, 10] (got {wd!r}); typical 0..0.1"
)
raise ValueError(f"weight_decay must be in [0, 10] (got {wd!r}); typical 0..0.1")
return wd
@field_validator("lora_r")
@ -271,9 +250,7 @@ class TrainingStartRequest(BaseModel):
if v is None:
return 16
if v < 1 or v > _MAX_LORA_ALPHA:
raise ValueError(
f"lora_alpha must be in [1, {_MAX_LORA_ALPHA}] (got {v!r})"
)
raise ValueError(f"lora_alpha must be in [1, {_MAX_LORA_ALPHA}] (got {v!r})")
return v
@field_validator("lora_dropout")
@ -302,9 +279,7 @@ class TrainingStartRequest(BaseModel):
num_epochs: int = Field(1, description = "Number of training epochs")
learning_rate: str = Field("2e-4", description = "Learning rate")
batch_size: int = Field(1, description = "Batch size")
gradient_accumulation_steps: int = Field(
1, description = "Gradient accumulation steps"
)
gradient_accumulation_steps: int = Field(1, description = "Gradient accumulation steps")
warmup_steps: Optional[int] = Field(None, description = "Warmup steps")
warmup_ratio: Optional[float] = Field(None, description = "Warmup ratio")
max_steps: Optional[int] = Field(None, description = "Maximum training steps")
@ -332,31 +307,19 @@ class TrainingStartRequest(BaseModel):
lora_r: int = Field(16, description = "LoRA rank")
lora_alpha: int = Field(16, description = "LoRA alpha")
lora_dropout: float = Field(0.0, description = "LoRA dropout")
target_modules: List[str] = Field(
default_factory = list, description = "Target modules for LoRA"
)
gradient_checkpointing: str = Field(
"", description = "Gradient checkpointing setting"
)
target_modules: List[str] = Field(default_factory = list, description = "Target modules for LoRA")
gradient_checkpointing: str = Field("", description = "Gradient checkpointing setting")
use_rslora: bool = Field(False, description = "Use RSLoRA")
use_loftq: bool = Field(False, description = "Use LoftQ")
train_on_completions: bool = Field(False, description = "Train on completions only")
# Vision-specific LoRA parameters
finetune_vision_layers: bool = Field(False, description = "Finetune vision layers")
finetune_language_layers: bool = Field(
False, description = "Finetune language layers"
)
finetune_attention_modules: bool = Field(
False, description = "Finetune attention modules"
)
finetune_language_layers: bool = Field(False, description = "Finetune language layers")
finetune_attention_modules: bool = Field(False, description = "Finetune attention modules")
finetune_mlp_modules: bool = Field(False, description = "Finetune MLP modules")
is_dataset_image: bool = Field(
False, description = "Whether the dataset contains image data"
)
is_dataset_audio: bool = Field(
False, description = "Whether the dataset contains audio data"
)
is_dataset_image: bool = Field(False, description = "Whether the dataset contains image data")
is_dataset_audio: bool = Field(False, description = "Whether the dataset contains audio data")
is_embedding: bool = Field(
False, description = "Whether model is an embedding/sentence-transformer model"
)
@ -382,9 +345,7 @@ class TrainingStartRequest(BaseModel):
# num_epochs and max_steps each accept 0 as a "use the other one"
# sentinel. If both resolve to 0 there's nothing to train against.
if (self.max_steps is None or self.max_steps == 0) and self.num_epochs == 0:
raise ValueError(
"Either num_epochs or max_steps must be > 0; both cannot be 0."
)
raise ValueError("Either num_epochs or max_steps must be > 0; both cannot be 0.")
return self
@ -411,9 +372,7 @@ class TrainingStatus(BaseModel):
"error",
"stopped",
] = Field(..., description = "Current phase of training pipeline")
is_training_running: bool = Field(
..., description = "True if training loop is actively running"
)
is_training_running: bool = Field(..., description = "True if training loop is actively running")
eval_enabled: bool = Field(
False,
description = "True if evaluation dataset is configured for this training run",
@ -438,9 +397,7 @@ class TrainingProgress(BaseModel):
total_steps: int = Field(..., description = "Total training steps")
loss: Optional[float] = Field(None, description = "Current loss value")
learning_rate: Optional[float] = Field(None, description = "Current learning rate")
progress_percent: float = Field(
..., description = "Progress percentage (0.0 to 100.0)"
)
progress_percent: float = Field(..., description = "Progress percentage (0.0 to 100.0)")
epoch: Optional[float] = Field(None, description = "Current epoch")
elapsed_seconds: Optional[float] = Field(
None, description = "Time elapsed since training started"
@ -449,9 +406,7 @@ class TrainingProgress(BaseModel):
grad_norm: Optional[float] = Field(
None, description = "L2 norm of gradients, computed before gradient clipping"
)
num_tokens: Optional[int] = Field(
None, description = "Total number of tokens processed so far"
)
num_tokens: Optional[int] = Field(None, description = "Total number of tokens processed so far")
eval_loss: Optional[float] = Field(
None, description = "Eval loss from the most recent evaluation step"
)

36
studio/backend/plugins/data-designer-github-repo-seed/src/data_designer_github_repo_seed/scraper.py
View file

@ -91,18 +91,13 @@ def _read_jsonl(path: Path, max_rows: int | None = None):
def _flatten_issue_row(r: dict, repo: str, include_comments: bool, max_c: int) -> dict:
labels = [
l.get("name")
for l in (r.get("labels", {}) or {}).get("nodes", [])
if l.get("name")
]
labels = [l.get("name") for l in (r.get("labels", {}) or {}).get("nodes", []) if l.get("name")]
comments_nodes = (r.get("comments") or {}).get("nodes") or []
comments_text = ""
if include_comments and comments_nodes:
kept = comments_nodes[:max_c]
comments_text = "\n\n".join(
f"[{(c.get('author') or {}).get('login', '?')}]: {c.get('body') or ''}"
for c in kept
f"[{(c.get('author') or {}).get('login', '?')}]: {c.get('body') or ''}" for c in kept
)
return {
"item_type": "issue",
@ -121,18 +116,13 @@ def _flatten_issue_row(r: dict, repo: str, include_comments: bool, max_c: int) -
def _flatten_pr_row(r: dict, repo: str, include_comments: bool, max_c: int) -> dict:
labels = [
l.get("name")
for l in (r.get("labels", {}) or {}).get("nodes", [])
if l.get("name")
]
labels = [l.get("name") for l in (r.get("labels", {}) or {}).get("nodes", []) if l.get("name")]
comments_nodes = (r.get("comments") or {}).get("nodes") or []
comments_text = ""
if include_comments and comments_nodes:
kept = comments_nodes[:max_c]
comments_text = "\n\n".join(
f"[{(c.get('author') or {}).get('login', '?')}]: {c.get('body') or ''}"
for c in kept
f"[{(c.get('author') or {}).get('login', '?')}]: {c.get('body') or ''}" for c in kept
)
return {
"item_type": "pull",
@ -205,14 +195,8 @@ def scrape(cfg: ScrapeConfig, base_dir: Path):
scraper.scrape_prs()
if "commits" in cfg.item_types:
default_ref = repo_meta.get("defaultBranchRef") or {}
default_branch = (
default_ref.get("name") if isinstance(default_ref, dict) else None
)
branch = (
f"refs/heads/{default_branch}"
if default_branch
else "refs/heads/main"
)
default_branch = default_ref.get("name") if isinstance(default_ref, dict) else None
branch = f"refs/heads/{default_branch}" if default_branch else "refs/heads/main"
scraper.scrape_commits(branch = branch)
finally:
scraper.close()
@ -222,16 +206,12 @@ def scrape(cfg: ScrapeConfig, base_dir: Path):
if "issues" in cfg.item_types:
for row in _read_jsonl(repo_dir / "issues.jsonl", read_cap):
all_rows.append(
_flatten_issue_row(
row, repo, cfg.include_comments, cfg.max_comments_per_item
)
_flatten_issue_row(row, repo, cfg.include_comments, cfg.max_comments_per_item)
)
if "pulls" in cfg.item_types:
for row in _read_jsonl(repo_dir / "pull_requests.jsonl", read_cap):
all_rows.append(
_flatten_pr_row(
row, repo, cfg.include_comments, cfg.max_comments_per_item
)
_flatten_pr_row(row, repo, cfg.include_comments, cfg.max_comments_per_item)
)
if "commits" in cfg.item_types:
for row in _read_jsonl(repo_dir / "commits.jsonl", read_cap):

35
studio/backend/plugins/data-designer-github-repo-seed/src/data_designer_github_repo_seed/scraper_impl/gh_client.py
View file

@ -60,9 +60,7 @@ class GitHubClient:
token_source: str | None = None,
):
if token:
self._token_source = (
token_source or "explicit token argument (recipe-level field)"
)
self._token_source = token_source or "explicit token argument (recipe-level field)"
elif os.environ.get("GH_TOKEN"):
self._token_source = "GH_TOKEN environment variable"
token = os.environ["GH_TOKEN"]
@ -72,9 +70,7 @@ class GitHubClient:
else:
raise RuntimeError("GH_TOKEN or GITHUB_TOKEN not set in environment")
self.session = requests.Session()
self.session.headers.update(
{**BASE_HEADERS, "Authorization": f"Bearer {token}"}
)
self.session.headers.update({**BASE_HEADERS, "Authorization": f"Bearer {token}"})
self.min_remaining_graphql = min_remaining_graphql
self.min_remaining_rest = min_remaining_rest
self.graphql_remaining: Optional[int] = None
@ -85,7 +81,11 @@ class GitHubClient:
self.calls_rest = 0
self.retry_count = 0
def _sleep_until(self, reset_ts: int, buffer_s: int = 10) -> None:
def _sleep_until(
self,
reset_ts: int,
buffer_s: int = 10,
) -> None:
now = int(time.time())
wait = max(0, reset_ts - now) + buffer_s
log.warning("Rate limit hit. Sleeping %ds until reset.", wait)
@ -209,9 +209,7 @@ class GitHubClient:
# Retry on RATE_LIMITED
for e in errs:
if e.get("type") == "RATE_LIMITED":
self._sleep_until(
(self.graphql_reset or int(time.time()) + 60)
)
self._sleep_until((self.graphql_reset or int(time.time()) + 60))
break
else:
# No rate-limit error, log and return partial
@ -243,9 +241,7 @@ class GitHubClient:
last_err = None
for attempt in range(max_retries):
try:
r = self.session.request(
method, url, params = params, json = json_body, timeout = 120
)
r = self.session.request(method, url, params = params, json = json_body, timeout = 120)
self.calls_rest += 1
rem = r.headers.get("X-RateLimit-Remaining")
rst = r.headers.get("X-RateLimit-Reset")
@ -269,9 +265,7 @@ class GitHubClient:
if r.status_code in (403, 429):
retry_after = _retry_after_seconds(r.headers.get("Retry-After"))
if retry_after is not None:
log.warning(
"Secondary rate limit on REST. Sleep %ds.", retry_after
)
log.warning("Secondary rate limit on REST. Sleep %ds.", retry_after)
time.sleep(retry_after + 2)
continue
# Check if primary rate
@ -290,7 +284,10 @@ class GitHubClient:
raise RuntimeError(f"REST failed after {max_retries} retries: {last_err}")
def rest_paginate(
self, path: str, params: Optional[Dict[str, Any]] = None, per_page: int = 100
self,
path: str,
params: Optional[Dict[str, Any]] = None,
per_page: int = 100,
) -> Iterator[dict]:
params = dict(params or {})
params.setdefault("per_page", per_page)
@ -298,9 +295,7 @@ class GitHubClient:
while True:
r = self.rest("GET", url, params = params if url == path else None)
if r.status_code != 200:
log.error(
"REST paginate got %s at %s: %s", r.status_code, url, r.text[:200]
)
log.error("REST paginate got %s at %s: %s", r.status_code, url, r.text[:200])
return
items = r.json()
if isinstance(items, dict):

87
studio/backend/plugins/data-designer-github-repo-seed/src/data_designer_github_repo_seed/scraper_impl/scraper.py
View file

@ -86,11 +86,7 @@ class RepoScraper:
return counter >= lim
def _log_rate(self, where: str, data: Dict[str, Any]) -> None:
rl = (
data.get("data", {}).get("rateLimit")
if isinstance(data.get("data"), dict)
else None
)
rl = data.get("data", {}).get("rateLimit") if isinstance(data.get("data"), dict) else None
if rl:
log.debug(
"[%s] rate cost=%s remaining=%s resetAt=%s",
@ -102,9 +98,7 @@ class RepoScraper:
# ----- repo meta -----
def scrape_repo_meta(self) -> Dict[str, Any]:
data = self.client.graphql(
Q.REPO_META_QUERY, {"owner": self.owner, "name": self.name}
)
data = self.client.graphql(Q.REPO_META_QUERY, {"owner": self.owner, "name": self.name})
self._log_rate("repo_meta", data)
repo = data.get("data", {}).get("repository") or {}
repo["_fetchedAt"] = ts()
@ -150,11 +144,7 @@ class RepoScraper:
self._paginate_issue_comments(
it["number"], it["comments"]["pageInfo"]["endCursor"]
)
if (
it.get("timelineItems", {})
.get("pageInfo", {})
.get("hasNextPage")
):
if it.get("timelineItems", {}).get("pageInfo", {}).get("hasNextPage"):
self._paginate_issue_timeline(
it["number"],
it["timelineItems"]["pageInfo"]["endCursor"],
@ -261,30 +251,16 @@ class RepoScraper:
num = pr["number"]
if not self.light:
if pr.get("comments", {}).get("pageInfo", {}).get("hasNextPage"):
self._paginate_pr_comments(
num, pr["comments"]["pageInfo"]["endCursor"]
)
if (
pr.get("timelineItems", {})
.get("pageInfo", {})
.get("hasNextPage")
):
self._paginate_pr_comments(num, pr["comments"]["pageInfo"]["endCursor"])
if pr.get("timelineItems", {}).get("pageInfo", {}).get("hasNextPage"):
self._paginate_pr_timeline(
num, pr["timelineItems"]["pageInfo"]["endCursor"]
)
if pr.get("commits", {}).get("pageInfo", {}).get("hasNextPage"):
self._paginate_pr_commits(
num, pr["commits"]["pageInfo"]["endCursor"]
)
self._paginate_pr_commits(num, pr["commits"]["pageInfo"]["endCursor"])
if pr.get("files", {}).get("pageInfo", {}).get("hasNextPage"):
self._paginate_pr_files(
num, pr["files"]["pageInfo"]["endCursor"]
)
if (
pr.get("reviewThreads", {})
.get("pageInfo", {})
.get("hasNextPage")
):
self._paginate_pr_files(num, pr["files"]["pageInfo"]["endCursor"])
if pr.get("reviewThreads", {}).get("pageInfo", {}).get("hasNextPage"):
self._paginate_pr_review_threads(
num, pr["reviewThreads"]["pageInfo"]["endCursor"]
)
@ -342,9 +318,7 @@ class RepoScraper:
"after": cur,
}
data = self.client.graphql(Q.PR_TIMELINE_QUERY, vars_)
item = ((data.get("data") or {}).get("repository") or {}).get(
"pullRequest"
) or {}
item = ((data.get("data") or {}).get("repository") or {}).get("pullRequest") or {}
tl = item.get("timelineItems") or {}
for ev in tl.get("nodes") or []:
ev["_owner"] = self.owner
@ -367,9 +341,7 @@ class RepoScraper:
"after": cur,
}
data = self.client.graphql(Q.PR_COMMITS_QUERY, vars_)
item = ((data.get("data") or {}).get("repository") or {}).get(
"pullRequest"
) or {}
item = ((data.get("data") or {}).get("repository") or {}).get("pullRequest") or {}
cc = item.get("commits") or {}
for c in cc.get("nodes") or []:
c["_owner"] = self.owner
@ -392,9 +364,7 @@ class RepoScraper:
"after": cur,
}
data = self.client.graphql(Q.PR_FILES_QUERY, vars_)
item = ((data.get("data") or {}).get("repository") or {}).get(
"pullRequest"
) or {}
item = ((data.get("data") or {}).get("repository") or {}).get("pullRequest") or {}
ff = item.get("files") or {}
for f in ff.get("nodes") or []:
f["_owner"] = self.owner
@ -419,9 +389,7 @@ class RepoScraper:
"after": cur,
}
data = self.client.graphql(Q.PR_REVIEW_THREADS_QUERY, vars_)
item = ((data.get("data") or {}).get("repository") or {}).get(
"pullRequest"
) or {}
item = ((data.get("data") or {}).get("repository") or {}).get("pullRequest") or {}
rt = item.get("reviewThreads") or {}
for th in rt.get("nodes") or []:
th["_owner"] = self.owner
@ -461,9 +429,7 @@ class RepoScraper:
d["_fetchedAt"] = ts()
num = d["number"]
if d.get("comments", {}).get("pageInfo", {}).get("hasNextPage"):
self._paginate_discussion_comments(
num, d["comments"]["pageInfo"]["endCursor"]
)
self._paginate_discussion_comments(num, d["comments"]["pageInfo"]["endCursor"])
# paginate replies per comment if needed
for c in d.get("comments", {}).get("nodes", []) or []:
if c.get("replies", {}).get("pageInfo", {}).get("hasNextPage"):
@ -501,9 +467,7 @@ class RepoScraper:
"after": cur,
}
data = self.client.graphql(Q.DISCUSSION_COMMENTS_QUERY, vars_)
disc = ((data.get("data") or {}).get("repository") or {}).get(
"discussion"
) or {}
disc = ((data.get("data") or {}).get("repository") or {}).get("discussion") or {}
cc = disc.get("comments") or {}
for c in cc.get("nodes") or []:
c["_owner"] = self.owner
@ -513,9 +477,7 @@ class RepoScraper:
info = cc.get("pageInfo") or {}
cur = info.get("endCursor") if info.get("hasNextPage") else None
def _paginate_discussion_replies(
self, comment_id: str, after: str, disc_number: int
) -> None:
def _paginate_discussion_replies(self, comment_id: str, after: str, disc_number: int) -> None:
cur = after
while cur:
vars_ = {
@ -650,12 +612,8 @@ def setup_logging(log_file: Path) -> None:
def main():
ap = argparse.ArgumentParser()
ap.add_argument(
"--base-dir", default = "/mnt/disks/unslothai/ubuntu/workspace_34/github_scraper"
)
ap.add_argument(
"--repos", nargs = "+", default = ["unslothai/unsloth", "unslothai/unsloth-zoo"]
)
ap.add_argument("--base-dir", default = "/mnt/disks/unslothai/ubuntu/workspace_34/github_scraper")
ap.add_argument("--repos", nargs = "+", default = ["unslothai/unsloth", "unslothai/unsloth-zoo"])
ap.add_argument("--trial", action = "store_true", help = "Small trial run")
ap.add_argument(
"--only",
@ -688,7 +646,6 @@ def main():
uploader = None
if args.hf_upload_interval > 0:
from hf_uploader import HFUploader
uploader = HFUploader(data_dir, interval_s = args.hf_upload_interval)
uploader.start()
@ -729,15 +686,9 @@ def main():
if not only or "commits" in only:
default_ref = repo_meta.get("defaultBranchRef") or {}
default_branch = (
default_ref.get("name")
if isinstance(default_ref, dict)
else None
)
branch = (
f"refs/heads/{default_branch}"
if default_branch
else "refs/heads/main"
default_ref.get("name") if isinstance(default_ref, dict) else None
)
branch = f"refs/heads/{default_branch}" if default_branch else "refs/heads/main"
scraper.scrape_commits(branch = branch)
finally:
scraper.close()

6
studio/backend/plugins/data-designer-github-repo-seed/src/data_designer_github_repo_seed/scraper_impl/state_store.py
View file

@ -25,7 +25,11 @@ class StateStore:
except Exception:
self._data = {}
def get(self, key: str, default: Any = None) -> Any:
def get(
self,
key: str,
default: Any = None,
) -> Any:
with self._lock:
return self._data.get(key, default)

47
studio/backend/plugins/data-designer-unstructured-seed/src/data_designer_unstructured_seed/chunking.py
View file

@ -19,10 +19,7 @@ _MIN_BREAK_RATIO = 0.6
_CACHE_DIR = unstructured_seed_cache_root()
def resolve_chunking(
chunk_size: Any,
chunk_overlap: Any,
) -> tuple[int, int]:
def resolve_chunking(chunk_size: Any, chunk_overlap: Any) -> tuple[int, int]:
size = _to_int(chunk_size, DEFAULT_CHUNK_SIZE)
size = max(1, min(size, MAX_CHUNK_SIZE))
overlap = _to_int(chunk_overlap, DEFAULT_CHUNK_OVERLAP)
@ -31,11 +28,7 @@ def resolve_chunking(
def build_unstructured_preview_rows(
*,
source_path: Path,
preview_size: int,
chunk_size: Any,
chunk_overlap: Any,
*, source_path: Path, preview_size: int, chunk_size: Any, chunk_overlap: Any
) -> list[dict[str, str]]:
parquet_path, rows = materialize_unstructured_seed_dataset(
source_path = source_path,
@ -49,9 +42,7 @@ def build_unstructured_preview_rows(
try:
import pandas as pd
except ImportError as exc: # pragma: no cover
raise RuntimeError(
f"pandas is required for unstructured seed processing: {exc}"
) from exc
raise RuntimeError(f"pandas is required for unstructured seed processing: {exc}") from exc
dataframe = pd.read_parquet(parquet_path).head(count)
return [
@ -78,10 +69,7 @@ def build_multi_file_preview_rows(
return _round_robin_preview(rows, preview_size)
def _round_robin_preview(
rows: list[dict[str, str]],
preview_size: int,
) -> list[dict[str, str]]:
def _round_robin_preview(rows: list[dict[str, str]], preview_size: int) -> list[dict[str, str]]:
"""Pick preview rows round-robin across source files so every file is represented."""
if not rows or preview_size <= 0:
return []
@ -115,10 +103,7 @@ def _round_robin_preview(
def materialize_unstructured_seed_dataset(
*,
source_path: Path,
chunk_size: Any,
chunk_overlap: Any,
*, source_path: Path, chunk_size: Any, chunk_overlap: Any
) -> tuple[Path, list[dict[str, str]]]:
resolved = source_path.expanduser().resolve()
if not resolved.is_file():
@ -148,9 +133,7 @@ def materialize_unstructured_seed_dataset(
try:
import pandas as pd
except ImportError as exc: # pragma: no cover
raise RuntimeError(
f"pandas is required for unstructured seed processing: {exc}"
) from exc
raise RuntimeError(f"pandas is required for unstructured seed processing: {exc}") from exc
tmp_path = _CACHE_DIR / f"{key}.tmp.parquet"
pd.DataFrame(rows).to_parquet(tmp_path, index = False)
@ -209,12 +192,7 @@ def normalize_unstructured_text(text: str) -> str:
return re.sub(r"\n{3,}", "\n\n", normalized).strip()
def split_text_into_chunks(
*,
text: str,
chunk_size: int,
chunk_overlap: int,
) -> list[str]:
def split_text_into_chunks(*, text: str, chunk_size: int, chunk_overlap: int) -> list[str]:
if not text:
return []
if chunk_size <= 0:
@ -268,12 +246,7 @@ def _to_int(value: Any, fallback: int) -> int:
return parsed
def _compute_cache_key(
*,
source_path: Path,
chunk_size: int,
chunk_overlap: int,
) -> str:
def _compute_cache_key(*, source_path: Path, chunk_size: int, chunk_overlap: int) -> str:
stat = source_path.stat()
payload = "|".join(
[
@ -288,9 +261,7 @@ def _compute_cache_key(
def _compute_multi_file_cache_key(
file_entries: list[tuple[Path, str]],
chunk_size: int,
chunk_overlap: int,
file_entries: list[tuple[Path, str]], chunk_size: int, chunk_overlap: int
) -> str:
parts: list[str] = []
for path, name in sorted(file_entries, key = lambda e: e[1]):

28
studio/backend/routes/auth.py
View file

@ -227,9 +227,7 @@ async def auth_status() -> AuthStatusResponse:
return AuthStatusResponse(
initialized = storage.is_initialized(),
default_username = storage.DEFAULT_ADMIN_USERNAME,
requires_password_change = storage.requires_password_change(
storage.DEFAULT_ADMIN_USERNAME
)
requires_password_change = storage.requires_password_change(storage.DEFAULT_ADMIN_USERNAME)
if storage.is_initialized()
else True,
)
@ -246,10 +244,7 @@ async def login(payload: AuthLoginRequest, request: Request) -> Token:
status_code = status.HTTP_429_TOO_MANY_REQUESTS,
# IP is intentionally not interpolated into the body; behind a
# proxy or NAT it is either misleading or an info leak.
detail = (
f"Too many failed login attempts. "
f"Try again in {blocked_for} seconds."
),
detail = (f"Too many failed login attempts. " f"Try again in {blocked_for} seconds."),
headers = {"Retry-After": str(blocked_for)},
)
@ -285,8 +280,7 @@ async def login(payload: AuthLoginRequest, request: Request) -> Token:
@router.post("/logout", status_code = status.HTTP_204_NO_CONTENT)
async def logout(
request: Request,
current_subject: str = Depends(get_current_subject_allow_password_change),
request: Request, current_subject: str = Depends(get_current_subject_allow_password_change)
) -> Response:
"""Revoke refresh tokens for the subject; the access token is stateless and expires on its own."""
try:
@ -335,9 +329,7 @@ async def refresh(payload: RefreshTokenRequest) -> Token:
access_token = new_access_token,
refresh_token = new_refresh_token,
token_type = "bearer",
must_change_password = False
if is_desktop
else storage.requires_password_change(username),
must_change_password = False if is_desktop else storage.requires_password_change(username),
)
@ -402,8 +394,7 @@ def _row_to_api_key_response(row: dict) -> ApiKeyResponse:
@router.post("/api-keys", response_model = CreateApiKeyResponse)
async def create_api_key(
payload: CreateApiKeyRequest,
current_subject: str = Depends(get_current_subject),
payload: CreateApiKeyRequest, current_subject: str = Depends(get_current_subject)
) -> CreateApiKeyResponse:
"""Create a new API key. The raw key is returned once and cannot be retrieved later."""
expires_at = None
@ -424,9 +415,7 @@ async def create_api_key(
@router.get("/api-keys", response_model = ApiKeyListResponse)
async def list_api_keys(
current_subject: str = Depends(get_current_subject),
) -> ApiKeyListResponse:
async def list_api_keys(current_subject: str = Depends(get_current_subject)) -> ApiKeyListResponse:
"""List all API keys for the authenticated user (raw keys are never exposed)."""
rows = storage.list_api_keys(current_subject)
return ApiKeyListResponse(
@ -435,10 +424,7 @@ async def list_api_keys(
@router.delete("/api-keys/{key_id}")
async def revoke_api_key(
key_id: int,
current_subject: str = Depends(get_current_subject),
) -> dict:
async def revoke_api_key(key_id: int, current_subject: str = Depends(get_current_subject)) -> dict:
"""Revoke (soft-delete) an API key."""
if not storage.revoke_api_key(current_subject, key_id):
raise HTTPException(

52
studio/backend/routes/chat_history.py
View file

@ -163,9 +163,7 @@ class ChatSettingsPayload(BaseModel):
inferenceParams: Optional[ChatInferenceSettings] = None
customPresets: Optional[list[ChatPreset]] = None
activePreset: Optional[str] = None
activePresetSource: Optional[Literal["builtin-default", "custom", "modified"]] = (
None
)
activePresetSource: Optional[Literal["builtin-default", "custom", "modified"]] = None
autoTitle: Optional[bool] = None
reasoningEffort: Optional[
Literal["none", "minimal", "low", "medium", "high", "max", "xhigh"]
@ -228,10 +226,7 @@ async def list_threads(
@router.post("/threads", response_model = ChatThread)
async def save_thread(
payload: ChatThread,
current_subject: str = Depends(get_current_subject),
):
async def save_thread(payload: ChatThread, current_subject: str = Depends(get_current_subject)):
if payload.projectId and get_chat_project(payload.projectId) is None:
raise HTTPException(
status_code = 404,
@ -241,10 +236,7 @@ async def save_thread(
@router.get("/threads/{thread_id}", response_model = ChatThread)
async def get_thread(
thread_id: str,
current_subject: str = Depends(get_current_subject),
):
async def get_thread(thread_id: str, current_subject: str = Depends(get_current_subject)):
thread = get_chat_thread(thread_id)
if thread is None:
raise HTTPException(status_code = 404, detail = f"Thread {thread_id} not found")
@ -277,8 +269,7 @@ async def patch_thread(
@router.delete("/threads")
async def delete_threads(
payload: ChatDeleteRequest,
current_subject: str = Depends(get_current_subject),
payload: ChatDeleteRequest, current_subject: str = Depends(get_current_subject)
):
delete_chat_threads(payload.ids)
return {"status": "deleted"}
@ -286,8 +277,7 @@ async def delete_threads(
@router.get("/projects", response_model = ChatProjectListResponse)
async def list_projects(
include_archived: bool = Query(False),
current_subject: str = Depends(get_current_subject),
include_archived: bool = Query(False), current_subject: str = Depends(get_current_subject)
):
return ChatProjectListResponse(
projects = [
@ -298,18 +288,12 @@ async def list_projects(
@router.post("/projects", response_model = ChatProject)
async def save_project(
payload: ChatProject,
current_subject: str = Depends(get_current_subject),
):
async def save_project(payload: ChatProject, current_subject: str = Depends(get_current_subject)):
return ChatProject(**upsert_chat_project(payload.model_dump()))
@router.get("/projects/{project_id}", response_model = ChatProject)
async def get_project(
project_id: str,
current_subject: str = Depends(get_current_subject),
):
async def get_project(project_id: str, current_subject: str = Depends(get_current_subject)):
project = ensure_chat_project_workspace(project_id)
if project is None:
raise HTTPException(
@ -356,10 +340,7 @@ async def delete_project(
@router.get("/threads/{thread_id}/messages", response_model = ChatMessageListResponse)
async def get_thread_messages(
thread_id: str,
current_subject: str = Depends(get_current_subject),
):
async def get_thread_messages(thread_id: str, current_subject: str = Depends(get_current_subject)):
if get_chat_thread(thread_id) is None:
raise HTTPException(status_code = 404, detail = f"Thread {thread_id} not found")
return ChatMessageListResponse(
@ -369,8 +350,7 @@ async def get_thread_messages(
@router.post("/messages:batch", response_model = ChatMessagesBatchResponse)
async def batch_thread_messages(
payload: ChatMessagesBatchRequest,
current_subject: str = Depends(get_current_subject),
payload: ChatMessagesBatchRequest, current_subject: str = Depends(get_current_subject)
):
"""One round-trip per sidebar/search rebuild instead of N. Unknown thread
ids are returned as empty lists so callers don't need a pre-flight."""
@ -425,14 +405,10 @@ async def replace_thread_messages(
payload: ChatMessageSyncRequest,
current_subject: str = Depends(get_current_subject),
):
mismatched_ids = [
message.id for message in payload.messages if message.threadId != thread_id
]
mismatched_ids = [message.id for message in payload.messages if message.threadId != thread_id]
if mismatched_ids:
preview = ", ".join(mismatched_ids[:5])
suffix = (
"" if len(mismatched_ids) <= 5 else f" (+{len(mismatched_ids) - 5} more)"
)
suffix = "" if len(mismatched_ids) <= 5 else f" (+{len(mismatched_ids) - 5} more)"
raise HTTPException(
status_code = 400,
detail = f"Message threadId mismatch: {preview}{suffix}",
@ -478,8 +454,7 @@ async def get_import_ledger(current_subject: str = Depends(get_current_subject))
@router.post("/import-ledger", response_model = ChatImportLedgerRecordResponse)
async def record_import_ledger(
payload: ChatImportLedgerRecordRequest,
current_subject: str = Depends(get_current_subject),
payload: ChatImportLedgerRecordRequest, current_subject: str = Depends(get_current_subject)
):
"""Mark each legacy thread id as imported. Idempotent."""
accepted, inserted = upsert_chat_legacy_imports(payload.threadIds)
@ -499,8 +474,7 @@ async def get_settings(current_subject: str = Depends(get_current_subject)):
@router.put("/settings", response_model = ChatSettingsResponse)
async def put_settings(
payload: dict[str, Any],
current_subject: str = Depends(get_current_subject),
payload: dict[str, Any], current_subject: str = Depends(get_current_subject)
):
try:
parsed = ChatSettingsPayload.model_validate(payload)

37
studio/backend/routes/data_recipe/jobs.py
View file

@ -159,9 +159,7 @@ def _ensure_selected_local_model_loaded(
) -> None:
model_loaded, active_model, active_variant = _loaded_local_model_identity()
if not model_loaded:
raise ValueError(
"No model loaded in Chat. Load a model first, then run the recipe."
)
raise ValueError("No model loaded in Chat. Load a model first, then run the recipe.")
selection = _single_used_local_model_selection(recipe, local_provider_names)
if selection is None:
@ -171,9 +169,7 @@ def _ensure_selected_local_model_loaded(
variant_matches = not gguf_variant or active_variant == gguf_variant
if active_model.lower() != target.lower() or not variant_matches:
selected = f"{target} ({gguf_variant})" if gguf_variant else target
active = (
f"{active_model} ({active_variant})" if active_variant else active_model
)
active = f"{active_model} ({active_variant})" if active_variant else active_model
raise ValueError(
"Selected local model is not loaded. "
f"Selected {selected}; active {active or 'none'}. "
@ -207,9 +203,7 @@ def _inject_local_structured_response_format(
for mc in model_configs:
if not isinstance(mc, dict):
continue
if mc.get("provider") in local_provider_names and isinstance(
mc.get("alias"), str
):
if mc.get("provider") in local_provider_names and isinstance(mc.get("alias"), str):
alias_to_local_mc[mc["alias"]] = mc
if not alias_to_local_mc:
@ -307,18 +301,12 @@ def _inject_local_providers(recipe: dict[str, Any], request: Request) -> Optiona
# from an LLM column through a model_config. Orphan model_config nodes
# that reference a local provider but that no LLM column uses should
# not block runs; the recipe would never call /v1 for them.
local_names = {
providers[i].get("name") for i in local_indices if providers[i].get("name")
}
local_names = {providers[i].get("name") for i in local_indices if providers[i].get("name")}
used_aliases = _used_llm_model_aliases(recipe)
referenced_providers = {
mc.get("provider")
for mc in recipe.get("model_configs", [])
if (
isinstance(mc, dict)
and mc.get("provider")
and mc.get("alias") in used_aliases
)
if (isinstance(mc, dict) and mc.get("provider") and mc.get("alias") in used_aliases)
}
token = ""
@ -409,9 +397,7 @@ def _normalize_run_name(value: Any) -> str | None:
if value is None:
return None
if not isinstance(value, str):
raise HTTPException(
status_code = 400, detail = "invalid run_name: must be a string"
)
raise HTTPException(status_code = 400, detail = "invalid run_name: must be a string")
trimmed = value.strip()
if not trimmed:
return None
@ -439,7 +425,6 @@ def create_job(payload: RecipePayload, request: Request):
if run_config_raw is not None:
try:
from data_designer.config.run_config import RunConfig
RunConfig.model_validate(run_config_raw)
except (ImportError, ValidationError, TypeError, ValueError) as exc:
raise log_and_http_error(
@ -506,7 +491,6 @@ def _revoke_internal_api_key_safe(key_id: int) -> None:
that revocation failures never mask the caller's own error path."""
try:
from auth import storage # deferred: avoids circular import
storage.revoke_internal_api_key(key_id)
except Exception:
pass
@ -578,9 +562,7 @@ def publish_job_dataset(job_id: str, payload: PublishDatasetRequest):
description = payload.description.strip()
hf_token = payload.hf_token.strip() if isinstance(payload.hf_token, str) else None
artifact_path = (
payload.artifact_path.strip()
if isinstance(payload.artifact_path, str)
else None
payload.artifact_path.strip() if isinstance(payload.artifact_path, str) else None
)
if not repo_id:
@ -591,10 +573,7 @@ def publish_job_dataset(job_id: str, payload: PublishDatasetRequest):
mgr = get_job_manager()
status = mgr.get_status(job_id)
if status is not None:
if (
status.get("status") != "completed"
or status.get("execution_type") != "full"
):
if status.get("status") != "completed" or status.get("execution_type") != "full":
raise HTTPException(
status_code = 409,
detail = "Only completed full runs can be published.",

4
studio/backend/routes/data_recipe/mcp.py
View file

@ -69,9 +69,7 @@ def list_mcp_tools(payload: McpToolsListRequest) -> McpToolsListResponse:
provider = built[0]
try:
tools = mcp_io.list_tools(provider, timeout_sec = payload.timeout_sec)
tool_names = sorted(
{tool.name for tool in tools if getattr(tool, "name", "")}
)
tool_names = sorted({tool.name for tool in tools if getattr(tool, "name", "")})
for tool_name in tool_names:
tool_to_providers[tool_name].append(provider.name)
providers.append(

51
studio/backend/routes/data_recipe/seed.py
View file

@ -63,9 +63,7 @@ _SAFE_ID_RE = re.compile(r"^[a-zA-Z0-9_-]+$")
def _validate_safe_id(value: str, label: str) -> str:
if not value or not _SAFE_ID_RE.match(value):
raise HTTPException(
400, f"Invalid {label}: must be alphanumeric/dash/underscore only"
)
raise HTTPException(400, f"Invalid {label}: must be alphanumeric/dash/underscore only")
return value
@ -75,8 +73,7 @@ def _serialize_preview_value(value: Any) -> Any:
def _serialize_preview_rows(rows: list[dict[str, Any]]) -> list[dict[str, Any]]:
return [
{str(key): _serialize_preview_value(value) for key, value in row.items()}
for row in rows
{str(key): _serialize_preview_value(value) for key, value in row.items()} for row in rows
]
@ -124,7 +121,9 @@ def _select_best_file(data_files: list[str], split: str = DEFAULT_SPLIT) -> str
def _resolve_seed_hf_path(
dataset_name: str, data_files: list[str], split: str = DEFAULT_SPLIT
dataset_name: str,
data_files: list[str],
split: str = DEFAULT_SPLIT,
) -> str | None:
selected = _select_best_file(data_files, split)
if not selected:
@ -164,10 +163,7 @@ def _build_stream_load_kwargs(
def _load_preview_rows(
*,
load_dataset_fn,
load_kwargs: dict[str, Any],
preview_size: int,
*, load_dataset_fn, load_kwargs: dict[str, Any], preview_size: int
) -> list[dict[str, Any]]:
streamed_ds = load_dataset_fn(**load_kwargs)
return [row for row in islice(streamed_ds, preview_size)]
@ -198,9 +194,7 @@ def _decode_base64_payload(content_base64: str) -> bytes:
raise HTTPException(status_code = 400, detail = "invalid base64 payload") from exc
def _read_preview_rows_from_local_file(
path: Path, preview_size: int
) -> list[dict[str, Any]]:
def _read_preview_rows_from_local_file(path: Path, preview_size: int) -> list[dict[str, Any]]:
try:
import pandas as pd
except ImportError as exc:
@ -251,11 +245,7 @@ def _read_preview_rows_from_local_file(
def _read_preview_rows_from_unstructured_file(
*,
path: Path,
preview_size: int,
chunk_size: int | None,
chunk_overlap: int | None,
*, path: Path, preview_size: int, chunk_size: int | None, chunk_overlap: int | None
) -> list[dict[str, Any]]:
if resolve_chunking is None or build_unstructured_preview_rows is None:
raise HTTPException(
@ -302,9 +292,7 @@ def _read_preview_rows_from_multi_files(
for fid, fname in zip(file_ids, file_names):
extracted = block_dir / f"{fid}.extracted.txt"
if not extracted.exists():
raise HTTPException(
404, f"Extracted text not found for file: {fname} (id: {fid})"
)
raise HTTPException(404, f"Extracted text not found for file: {fname} (id: {fid})")
file_entries.append((extracted, fname))
return build_multi_file_preview_rows(
@ -384,9 +372,7 @@ def inspect_seed_dataset(payload: SeedInspectRequest) -> SeedInspectResponse:
) from exc
if not preview_rows:
raise HTTPException(
status_code = 422, detail = "dataset appears empty or unreadable"
)
raise HTTPException(status_code = 422, detail = "dataset appears empty or unreadable")
preview_rows = _serialize_preview_rows(preview_rows)
columns = _extract_columns(preview_rows)
@ -395,9 +381,7 @@ def inspect_seed_dataset(payload: SeedInspectRequest) -> SeedInspectResponse:
else:
resolved_path = _resolve_seed_hf_path(dataset_name, data_files, split)
if not resolved_path:
raise HTTPException(
status_code = 422, detail = "unable to resolve seed dataset path"
)
raise HTTPException(status_code = 422, detail = "unable to resolve seed dataset path")
return SeedInspectResponse(
dataset_name = dataset_name,
@ -415,13 +399,11 @@ def _extract_text_from_file(file_path: Path, ext: str) -> str:
raw = file_path.read_text(encoding = "utf-8", errors = "ignore")
elif ext == ".pdf":
import pymupdf4llm
raw = pymupdf4llm.to_markdown(
str(file_path), write_images = False, show_progress = False, use_ocr = False
)
elif ext == ".docx":
import mammoth
with open(str(file_path), "rb") as f:
result = mammoth.convert_to_markdown(f)
raw = result.value
@ -449,8 +431,7 @@ def _get_block_total_size(block_dir: Path) -> int:
@router.post("/seed/upload-unstructured-file")
async def upload_unstructured_file(
file: UploadFile = FastAPIFile(...),
block_id: str = Form(...),
file: UploadFile = FastAPIFile(...), block_id: str = Form(...)
) -> UnstructuredFileUploadResponse:
_validate_safe_id(block_id, "block_id")
@ -519,9 +500,7 @@ async def upload_unstructured_file(
try:
meta_path = block_dir / f"{file_id}.meta.json"
meta_path.write_text(
json.dumps(
{"original_filename": original_filename, "size_bytes": size_bytes}
),
json.dumps({"original_filename": original_filename, "size_bytes": size_bytes}),
encoding = "utf-8",
)
except OSError:
@ -647,9 +626,7 @@ def inspect_seed_upload(payload: SeedInspectUploadRequest) -> SeedInspectRespons
int(payload.preview_size),
)
if not preview_rows:
raise HTTPException(
status_code = 422, detail = "dataset appears empty or unreadable"
)
raise HTTPException(status_code = 422, detail = "dataset appears empty or unreadable")
columns = _extract_columns(preview_rows)
return SeedInspectResponse(

19
studio/backend/routes/data_recipe/validate.py
View file

@ -21,7 +21,9 @@ from utils.utils import safe_error_detail, safe_curated_detail, log_and_http_err
logger = get_logger(__name__)
router = APIRouter()
_GITHUB_VALIDATE_NOTE = "Recipe shape is valid. GitHub access and rate limits are checked when the run starts."
_GITHUB_VALIDATE_NOTE = (
"Recipe shape is valid. GitHub access and rate limits are checked when the run starts."
)
_GITHUB_ITEM_TYPES = {"issues", "pulls", "commits"}
@ -44,23 +46,17 @@ def _validate_github_seed_static(source: dict[str, Any]) -> list[ValidateError]:
else:
for repo in repos:
if not isinstance(repo, str) or not repo.strip() or "/" not in repo:
errors.append(
ValidateError(message = "GitHub repos must be owner/name strings.")
)
errors.append(ValidateError(message = "GitHub repos must be owner/name strings."))
break
item_types = source.get("item_types")
if not isinstance(item_types, list) or not item_types:
errors.append(
ValidateError(message = "GitHub seed requires at least one item type.")
)
errors.append(ValidateError(message = "GitHub seed requires at least one item type."))
else:
invalid_items = [item for item in item_types if item not in _GITHUB_ITEM_TYPES]
if invalid_items:
errors.append(
ValidateError(
message = "GitHub item types must be issues, pulls, or commits."
)
ValidateError(message = "GitHub item types must be issues, pulls, or commits.")
)
try:
@ -166,8 +162,7 @@ def validate(payload: RecipePayload) -> ValidateResponse:
if not (exc.name or "").startswith("data_designer"):
raise
logger.debug(
"data_designer not installed; deferring full config "
"validation to run start",
"data_designer not installed; deferring full config validation to run start",
missing_module = exc.name,
)
except Exception as exc:

47
studio/backend/routes/datasets.py
View file

@ -99,7 +99,6 @@ def _serialize_preview_value(value):
try:
from PIL.Image import Image as PILImage
if isinstance(value, PILImage):
buffer = io.BytesIO()
value.convert("RGB").save(buffer, format = "JPEG", quality = 85)
@ -261,9 +260,7 @@ def _build_local_dataset_items() -> list[LocalDatasetItem]:
return items
def _load_local_preview_slice(
*, dataset_path: Path, train_split: str, preview_size: int
):
def _load_local_preview_slice(*, dataset_path: Path, train_split: str, preview_size: int):
from datasets import load_dataset
if dataset_path.is_dir():
@ -298,9 +295,7 @@ def _load_local_preview_slice(
elif dataset_path.suffix == ".csv":
dataset = load_dataset("csv", data_files = str(dataset_path), split = train_split)
elif dataset_path.suffix == ".parquet":
dataset = load_dataset(
"parquet", data_files = str(dataset_path), split = train_split
)
dataset = load_dataset("parquet", data_files = str(dataset_path), split = train_split)
else:
raise HTTPException(
status_code = 400, detail = f"Unsupported file format: {dataset_path.suffix}"
@ -320,8 +315,7 @@ def _sanitize_filename(filename: str) -> str:
@router.post("/upload", response_model = UploadDatasetResponse)
async def upload_dataset(
file: UploadFile,
current_subject: str = Depends(get_current_subject),
file: UploadFile, current_subject: str = Depends(get_current_subject)
) -> UploadDatasetResponse:
filename = _sanitize_filename(file.filename or "dataset_upload")
ext = Path(filename).suffix.lower()
@ -378,9 +372,7 @@ def list_local_datasets(
@router.get("/download-progress")
async def get_dataset_download_progress(
repo_id: str = Query(
..., description = "HuggingFace dataset repo ID, e.g. 'unsloth/LaTeX_OCR'"
),
repo_id: str = Query(..., description = "HuggingFace dataset repo ID, e.g. 'unsloth/LaTeX_OCR'"),
current_subject: str = Depends(get_current_subject),
):
"""Return download progress for a HuggingFace dataset repo.
@ -460,10 +452,7 @@ async def get_dataset_download_progress(
@router.post("/check-format", response_model = CheckFormatResponse)
def check_format(
request: CheckFormatRequest,
current_subject: str = Depends(get_current_subject),
):
def check_format(request: CheckFormatRequest, current_subject: str = Depends(get_current_subject)):
"""
Check if a dataset requires manual column mapping.
@ -511,25 +500,19 @@ def check_format(
repo_type = "dataset",
token = request.hf_token or None,
)
data_files = [
f for f in repo_files if any(f.endswith(ext) for ext in DATA_EXTS)
]
data_files = [f for f in repo_files if any(f.endswith(ext) for ext in DATA_EXTS)]
# Prefer tabular formats over archives (e.g. images.zip → ImageFolder
# with synthetic image/label columns that don't match the real schema).
tabular_files = [
f
for f in data_files
if any(f.endswith(ext) for ext in _TABULAR_EXTS)
f for f in data_files if any(f.endswith(ext) for ext in _TABULAR_EXTS)
]
candidates = tabular_files or data_files
# When a subset is specified, narrow to files whose name matches
# (e.g. subset="testmini" → prefer "testmini.parquet").
if request.subset and candidates:
subset_matches = [
f for f in candidates if request.subset in Path(f).stem
]
subset_matches = [f for f in candidates if request.subset in Path(f).stem]
if subset_matches:
candidates = subset_matches
@ -601,9 +584,7 @@ def check_format(
processed = format_result["dataset"]
preview_samples = _serialize_preview_rows(processed)
except Exception as e:
logger.warning(
f"Processed preview generation failed (non-fatal): {e}"
)
logger.warning(f"Processed preview generation failed (non-fatal): {e}")
preview_samples = _serialize_preview_rows(preview_slice)
else:
preview_samples = _serialize_preview_rows(preview_slice)
@ -614,9 +595,7 @@ def check_format(
if image_col and image_col in (result.get("columns") or []):
try:
sample_val = preview_slice[0][image_col]
if isinstance(sample_val, str) and sample_val.startswith(
("http://", "https://")
):
if isinstance(sample_val, str) and sample_val.startswith(("http://", "https://")):
url_warning = (
"This dataset contains image URLs instead of embedded images. "
"Images will be downloaded during training, which may be slow for large datasets."
@ -652,8 +631,7 @@ def check_format(
@router.post("/ai-assist-mapping", response_model = AiAssistMappingResponse)
def ai_assist_mapping(
request: AiAssistMappingRequest,
current_subject: str = Depends(get_current_subject),
request: AiAssistMappingRequest, current_subject: str = Depends(get_current_subject)
):
"""
Run LLM-assisted dataset conversion advisor (user-triggered).
@ -670,8 +648,7 @@ def ai_assist_mapping(
# Truncate sample values for the LLM prompt
truncated = [
{col: str(s.get(col, ""))[:200] for col in request.columns}
for s in request.samples[:5]
{col: str(s.get(col, ""))[:200] for col in request.columns} for s in request.samples[:5]
]
result = llm_conversion_advisor(

36
studio/backend/routes/export.py
View file

@ -54,8 +54,7 @@ logger = get_logger(__name__)
@router.post("/load-checkpoint", response_model = ExportOperationResponse)
async def load_checkpoint(
request: LoadCheckpointRequest,
current_subject: str = Depends(get_current_subject),
request: LoadCheckpointRequest, current_subject: str = Depends(get_current_subject)
):
"""
Load a checkpoint into the export backend.
@ -70,7 +69,6 @@ async def load_checkpoint(
# before loading the export checkpoint (they'd compete for VRAM).
try:
from core.inference import get_inference_backend
inf = get_inference_backend()
if inf.active_model_name:
logger.info(
@ -85,7 +83,6 @@ async def load_checkpoint(
try:
from core.training import get_training_backend
trn = get_training_backend()
if trn.is_training_active():
logger.info("Stopping active training to free GPU memory for export")
@ -96,12 +93,9 @@ async def load_checkpoint(
if not trn.is_training_active():
break
import time
time.sleep(0.5)
else:
logger.warning(
"Training subprocess did not exit within 30s, proceeding anyway"
)
logger.warning("Training subprocess did not exit within 30s, proceeding anyway")
except Exception as e:
logger.warning("Could not stop training: %s", e)
@ -132,9 +126,7 @@ async def load_checkpoint(
@router.post("/cleanup", response_model = ExportOperationResponse)
async def cleanup_export_memory(
current_subject: str = Depends(get_current_subject),
):
async def cleanup_export_memory(current_subject: str = Depends(get_current_subject)):
"""
Cleanup export-related models from memory (GPU/CPU).
@ -165,9 +157,7 @@ async def cleanup_export_memory(
@router.get("/status", response_model = ExportStatusResponse)
async def get_export_status(
current_subject: str = Depends(get_current_subject),
):
async def get_export_status(current_subject: str = Depends(get_current_subject)):
"""
Get current export backend status (loaded checkpoint, model type, PEFT flag).
"""
@ -203,8 +193,7 @@ def _export_details(output_path: Optional[str]) -> Optional[Dict[str, Any]]:
@router.post("/export/merged", response_model = ExportOperationResponse)
async def export_merged_model(
request: ExportMergedModelRequest,
current_subject: str = Depends(get_current_subject),
request: ExportMergedModelRequest, current_subject: str = Depends(get_current_subject)
):
"""
Export a merged PEFT model (e.g., 16-bit or 4-bit) and optionally push to Hub.
@ -243,8 +232,7 @@ async def export_merged_model(
@router.post("/export/base", response_model = ExportOperationResponse)
async def export_base_model(
request: ExportBaseModelRequest,
current_subject: str = Depends(get_current_subject),
request: ExportBaseModelRequest, current_subject: str = Depends(get_current_subject)
):
"""
Export a non-PEFT base model and optionally push to Hub.
@ -283,8 +271,7 @@ async def export_base_model(
@router.post("/export/gguf", response_model = ExportOperationResponse)
async def export_gguf(
request: ExportGGUFRequest,
current_subject: str = Depends(get_current_subject),
request: ExportGGUFRequest, current_subject: str = Depends(get_current_subject)
):
"""
Export the current model to GGUF format and optionally push to Hub.
@ -322,8 +309,7 @@ async def export_gguf(
@router.post("/export/lora", response_model = ExportOperationResponse)
async def export_lora_adapter(
request: ExportLoRAAdapterRequest,
current_subject: str = Depends(get_current_subject),
request: ExportLoRAAdapterRequest, current_subject: str = Depends(get_current_subject)
):
"""
Export only the LoRA adapter (if the loaded model is PEFT).
@ -377,7 +363,11 @@ async def export_lora_adapter(
# directive, and `Last-Event-ID` is honored on reconnect.
def _format_sse(data: str, event: str, event_id: Optional[int] = None) -> str:
def _format_sse(
data: str,
event: str,
event_id: Optional[int] = None,
) -> str:
"""Format a single SSE message with id/event/data fields."""
lines = []
if event_id is not None:

379
studio/backend/routes/inference.py
View file

@ -92,7 +92,9 @@ def _friendly_error(exc: Exception) -> str:
# subprocess is unreachable", which for Studio always means the
# llama-server subprocess crashed or is still coming up.
if isinstance(exc, httpx.RequestError):
return "Lost connection to the model server. It may have crashed -- try reloading the model."
return (
"Lost connection to the model server. It may have crashed -- try reloading the model."
)
msg = str(exc)
m = _re.search(
r"request \((\d+) tokens?\) exceeds the available context size \((\d+) tokens?\)",
@ -105,7 +107,9 @@ def _friendly_error(exc: Exception) -> str:
f"or shorten the conversation."
)
if "Lost connection to llama-server" in msg:
return "Lost connection to the model server. It may have crashed -- try reloading the model."
return (
"Lost connection to the model server. It may have crashed -- try reloading the model."
)
return "An internal error occurred"
@ -353,9 +357,7 @@ async def artifact_preview_frame(
await get_current_subject(creds)
csp = (
_ARTIFACT_PREVIEW_FRAME_NETWORK_CSP
if allow_network
else _ARTIFACT_PREVIEW_FRAME_STRICT_CSP
_ARTIFACT_PREVIEW_FRAME_NETWORK_CSP if allow_network else _ARTIFACT_PREVIEW_FRAME_STRICT_CSP
)
return Response(
content = _ARTIFACT_PREVIEW_FRAME_HTML,
@ -446,9 +448,7 @@ _PENDING_CANCEL_TTL_S = 30.0
def _prune_pending(now: float) -> None:
for k in [
k for k, ts in _PENDING_CANCELS.items() if now - ts > _PENDING_CANCEL_TTL_S
]:
for k in [k for k, ts in _PENDING_CANCELS.items() if now - ts > _PENDING_CANCEL_TTL_S]:
_PENDING_CANCELS.pop(k, None)
@ -590,9 +590,7 @@ _TOOL_XML_RE = _re.compile(
logger = get_logger(__name__)
def _validate_native_mmproj_companion(
mmproj_path: str | None, gguf_path: str | None
) -> None:
def _validate_native_mmproj_companion(mmproj_path: str | None, gguf_path: str | None) -> None:
if not mmproj_path or not gguf_path:
return
import stat as _stat_module
@ -606,9 +604,7 @@ def _validate_native_mmproj_companion(
status_code = 400,
detail = "Native vision companion is no longer accessible.",
) from exc
if _stat_module.S_ISLNK(mm_lstat.st_mode) or not _stat_module.S_ISREG(
mm_lstat.st_mode
):
if _stat_module.S_ISLNK(mm_lstat.st_mode) or not _stat_module.S_ISREG(mm_lstat.st_mode):
raise HTTPException(
status_code = 400,
detail = "Native vision companion must be a regular file.",
@ -636,9 +632,7 @@ def _normalise_settings_str(value: Optional[str]) -> Optional[str]:
return value
def _request_matches_loaded_settings(
request: LoadRequest, llama_backend: LlamaCppBackend
) -> bool:
def _request_matches_loaded_settings(request: LoadRequest, llama_backend: LlamaCppBackend) -> bool:
"""True iff every runtime setting on the request matches the loaded
server. Caller has already checked model+variant+is_loaded. See #5401."""
# Compare requested n_ctx (not effective) so VRAM-cap doesn't mask
@ -664,9 +658,7 @@ def _request_matches_loaded_settings(
if backend_mode in ("mtp", "mtp+ngram") and request.spec_draft_n_max is not None:
if int(request.spec_draft_n_max) != (llama_backend.spec_draft_n_max or 0):
return False
if (request.chat_template_override or None) != (
llama_backend.chat_template_override or None
):
if (request.chat_template_override or None) != (llama_backend.chat_template_override or None):
return False
# llama_extra_args=None means "inherit"; only an explicit list that
# differs forces a reload. On the inherit path, refuse to match if
@ -683,9 +675,7 @@ def _request_matches_loaded_settings(
def _resolve_model_identifier_for_request(
request: LoadRequest | ValidateModelRequest,
*,
operation: str,
request: LoadRequest | ValidateModelRequest, *, operation: str
) -> tuple[str, str, bool]:
if not request.native_path_lease:
return request.model_path, request.model_path, False
@ -705,9 +695,7 @@ def _resolve_model_identifier_for_request(
status_code = 400,
detail = redact_native_paths(str(exc)),
) from exc
display_label = (
grant.display_label or Path(request.model_path).name or "Native model"
)
display_label = grant.display_label or Path(request.model_path).name or "Native model"
return str(grant.canonical_path), display_label, True
@ -788,9 +776,7 @@ async def load_model(
inference_config = load_inference_config(llama_backend.model_identifier)
_gguf_audio = (
llama_backend._audio_type
if hasattr(llama_backend, "_audio_type")
else None
llama_backend._audio_type if hasattr(llama_backend, "_audio_type") else None
)
_gguf_is_audio = getattr(llama_backend, "_is_audio", False)
return LoadResponse(
@ -828,9 +814,7 @@ async def load_model(
backend.active_model_name
and backend.active_model_name.lower() == model_identifier.lower()
):
logger.info(
f"Model already loaded (Unsloth): {model_log_label}, skipping reload"
)
logger.info(f"Model already loaded (Unsloth): {model_log_label}, skipping reload")
inference_config = load_inference_config(backend.active_model_name)
_model_info = backend.models.get(backend.active_model_name, {})
_chat_template = None
@ -847,9 +831,7 @@ async def load_model(
_sf_reasoning_style = _sf_flags["reasoning_style"]
return LoadResponse(
status = "already_loaded",
model = model_log_label
if native_grant_backed
else backend.active_model_name,
model = model_log_label if native_grant_backed else backend.active_model_name,
display_name = model_log_label
if native_grant_backed
else backend.active_model_name,
@ -932,8 +914,7 @@ async def load_model(
)
if not same_source:
logger.info(
"Not inheriting llama_extra_args: stored args came "
"from %s, loading %s",
"Not inheriting llama_extra_args: stored args came from %s, loading %s",
source,
(model_identifier, resolved_variant),
)
@ -951,8 +932,7 @@ async def load_model(
strip_context = "max_seq_length" in fields_set,
strip_cache = "cache_type_kv" in fields_set,
strip_spec = (
"speculative_type" in fields_set
or "spec_draft_n_max" in fields_set
"speculative_type" in fields_set or "spec_draft_n_max" in fields_set
),
strip_template = "chat_template_override" in fields_set,
)
@ -1001,9 +981,7 @@ async def load_model(
else:
# Local mode: llama-server loads via -m <path>
if native_grant_backed and config.gguf_mmproj_file:
_validate_native_mmproj_companion(
config.gguf_mmproj_file, config.gguf_file
)
_validate_native_mmproj_companion(config.gguf_mmproj_file, config.gguf_file)
success = await asyncio.to_thread(
llama_backend.load_model,
gguf_path = config.gguf_file,
@ -1036,9 +1014,7 @@ async def load_model(
# Audio detection moved into load_model under _serial_load_lock (#5642).
_gguf_audio = llama_backend._audio_type
_gguf_is_audio = llama_backend._is_audio
llama_backend._native_display_label = (
model_log_label if native_grant_backed else None
)
llama_backend._native_display_label = model_log_label if native_grant_backed else None
llama_backend._native_grant_backed = bool(native_grant_backed)
if _gguf_is_audio:
logger.info(f"GGUF model detected as audio: audio_type={_gguf_audio}")
@ -1048,9 +1024,7 @@ async def load_model(
return LoadResponse(
status = "loaded",
model = model_log_label if native_grant_backed else config.identifier,
display_name = model_log_label
if native_grant_backed
else config.display_name,
display_name = model_log_label if native_grant_backed else config.display_name,
is_vision = llama_backend.is_vision,
is_lora = False,
is_gguf = True,
@ -1058,9 +1032,7 @@ async def load_model(
audio_type = _gguf_audio,
has_audio_input = llama_backend._has_audio_input,
inference = inference_config,
requires_trust_remote_code = bool(
inference_config.get("trust_remote_code", False)
),
requires_trust_remote_code = bool(inference_config.get("trust_remote_code", False)),
context_length = llama_backend.context_length,
max_context_length = llama_backend.max_context_length,
native_context_length = llama_backend.native_context_length,
@ -1087,12 +1059,9 @@ async def load_model(
# Shut down any export subprocess to free VRAM
try:
from core.export import get_export_backend
exp_backend = get_export_backend()
if exp_backend.current_checkpoint:
logger.info(
"Shutting down export subprocess to free GPU memory for inference"
)
logger.info("Shutting down export subprocess to free GPU memory for inference")
exp_backend._shutdown_subprocess()
exp_backend.current_checkpoint = None
exp_backend.is_vision = False
@ -1162,9 +1131,7 @@ async def load_model(
# Check if YAML says this model needs trust_remote_code
if not request.trust_remote_code:
model_defaults = load_model_defaults(config.identifier)
yaml_trust = model_defaults.get("inference", {}).get(
"trust_remote_code", False
)
yaml_trust = model_defaults.get("inference", {}).get("trust_remote_code", False)
if yaml_trust:
raise HTTPException(
status_code = 400,
@ -1200,9 +1167,7 @@ async def load_model(
return LoadResponse(
status = "loaded",
model = model_log_label if native_grant_backed else config.identifier,
display_name = model_log_label
if native_grant_backed
else config.display_name,
display_name = model_log_label if native_grant_backed else config.display_name,
is_vision = config.is_vision,
is_lora = config.is_lora,
is_gguf = False,
@ -1210,9 +1175,7 @@ async def load_model(
audio_type = config.audio_type,
has_audio_input = config.has_audio_input,
inference = inference_config,
requires_trust_remote_code = bool(
inference_config.get("trust_remote_code", False)
),
requires_trust_remote_code = bool(inference_config.get("trust_remote_code", False)),
supports_reasoning = _sf_flags["supports_reasoning"],
reasoning_style = _sf_flags["reasoning_style"],
reasoning_always_on = _sf_flags["reasoning_always_on"],
@ -1266,8 +1229,7 @@ async def load_model(
@router.post("/validate", response_model = ValidateModelResponse)
async def validate_model(
request: ValidateModelRequest,
current_subject: str = Depends(get_current_subject),
request: ValidateModelRequest, current_subject: str = Depends(get_current_subject)
):
"""
Lightweight validation endpoint for model identifiers.
@ -1342,10 +1304,7 @@ async def validate_model(
@router.post("/unload", response_model = UnloadResponse)
async def unload_model(
request: UnloadRequest,
current_subject: str = Depends(get_current_subject),
):
async def unload_model(request: UnloadRequest, current_subject: str = Depends(get_current_subject)):
"""
Unload a model from memory.
Routes to the correct backend (llama-server for GGUF, Unsloth otherwise).
@ -1355,9 +1314,7 @@ async def unload_model(
llama_backend = get_llama_cpp_backend()
if llama_backend.is_active and (
llama_backend.model_identifier == request.model_path
or is_registered_native_path_label(
llama_backend.model_identifier, request.model_path
)
or is_registered_native_path_label(llama_backend.model_identifier, request.model_path)
or not llama_backend.is_loaded
):
llama_backend.unload_model()
@ -1376,10 +1333,7 @@ async def unload_model(
@studio_router.post("/cancel")
async def cancel_inference(
request: Request,
current_subject: str = Depends(get_current_subject),
):
async def cancel_inference(request: Request, current_subject: str = Depends(get_current_subject)):
"""Cancel in-flight inference requests.
Body (JSON, at least one key required):
@ -1419,8 +1373,7 @@ async def cancel_inference(
@router.post("/generate/stream")
async def generate_stream(
request: GenerateRequest,
current_subject: str = Depends(get_current_subject),
request: GenerateRequest, current_subject: str = Depends(get_current_subject)
):
"""
Generate a chat response with Server-Sent Events (SSE) streaming.
@ -1496,9 +1449,7 @@ async def generate_stream(
@router.get("/status", response_model = InferenceStatusResponse)
async def get_status(
current_subject: str = Depends(get_current_subject),
):
async def get_status(current_subject: str = Depends(get_current_subject)):
"""
Get current inference backend status.
Reports whichever backend (Unsloth or llama-server) is currently active.
@ -1516,7 +1467,6 @@ async def get_status(
_supports_mtp = True # fail open
try:
from utils.llama_cpp_freshness import check_prebuilt_freshness
_freshness = check_prebuilt_freshness(_bin)
except Exception:
_freshness = {}
@ -1590,17 +1540,13 @@ async def get_status(
has_audio_input = model_info.get("has_audio_input", False)
chat_template_info = model_info.get("chat_template_info", {})
chat_template = (
chat_template_info.get("template")
if isinstance(chat_template_info, dict)
else None
chat_template_info.get("template") if isinstance(chat_template_info, dict) else None
)
# Non-GGUF: classify from the loaded template.
_sf_flags = _detect_safetensors_features(backend, chat_template)
inference_config = (
load_inference_config(backend.active_model_name)
if backend.active_model_name
else None
load_inference_config(backend.active_model_name) if backend.active_model_name else None
)
return InferenceStatusResponse(
@ -1635,9 +1581,7 @@ async def get_status(
@router.get("/load-progress", response_model = LoadProgressResponse)
async def get_load_progress(
current_subject: str = Depends(get_current_subject),
):
async def get_load_progress(current_subject: str = Depends(get_current_subject)):
"""
Return the active GGUF load's mmap/upload progress.
@ -1684,9 +1628,7 @@ async def generate_audio(
_, chat_messages, _ = _extract_content_parts(payload.messages)
if not chat_messages:
raise HTTPException(status_code = 400, detail = "No messages provided.")
last_user_msg = next(
(m for m in reversed(chat_messages) if m["role"] == "user"), None
)
last_user_msg = next((m for m in reversed(chat_messages) if m["role"] == "user"), None)
if not last_user_msg:
raise HTTPException(status_code = 400, detail = "No user message found.")
text = last_user_msg["content"]
@ -1711,9 +1653,7 @@ async def generate_audio(
raise HTTPException(status_code = 400, detail = "No model loaded.")
model_info = backend.models.get(backend.active_model_name, {})
if not model_info.get("is_audio"):
raise HTTPException(
status_code = 400, detail = "Active model is not an audio model."
)
raise HTTPException(status_code = 400, detail = "Active model is not an audio model.")
model_name = backend.active_model_name
gen = lambda: backend.generate_audio_response(
text = text,
@ -1727,9 +1667,7 @@ async def generate_audio(
)
try:
wav_bytes, sample_rate = await asyncio.get_event_loop().run_in_executor(
None, gen
)
wav_bytes, sample_rate = await asyncio.get_event_loop().run_in_executor(None, gen)
except Exception as e:
logger.error(f"Audio generation error: {e}", exc_info = True)
raise HTTPException(status_code = 500, detail = safe_error_detail(e))
@ -1795,9 +1733,7 @@ def _decode_audio_base64(b64: str) -> np.ndarray:
return waveform.squeeze(0).numpy()
def _extract_content_parts(
messages: list,
) -> tuple[str, list[dict], "Optional[str]"]:
def _extract_content_parts(messages: list) -> tuple[str, list[dict], "Optional[str]"]:
"""
Parse OpenAI-format messages into components the inference backend expects.
@ -1820,9 +1756,7 @@ def _extract_content_parts(
system_prompt = msg.content
elif isinstance(msg.content, list):
# Unlikely but handle: join text parts
system_prompt = "\n".join(
p.text for p in msg.content if p.type == "text"
)
system_prompt = "\n".join(p.text for p in msg.content if p.type == "text")
continue
# ── User / assistant messages ─────────────────────────
@ -1841,9 +1775,7 @@ def _extract_content_parts(
# data:image/png;base64,<DATA> → extract <DATA>
first_image_b64 = url.split(",", 1)[1] if "," in url else None
else:
logger.warning(
f"Remote image URLs not yet supported: {url[:80]}..."
)
logger.warning(f"Remote image URLs not yet supported: {url[:80]}...")
combined_text = "\n".join(text_parts) if text_parts else ""
chat_messages.append({"role": msg.role, "content": combined_text})
@ -1906,7 +1838,6 @@ def _build_external_messages(
if provider_type == "gemini" and base_url:
try:
from urllib.parse import urlparse as _urlparse
_host = (_urlparse(base_url).hostname or "").lower()
_native_gemini = _host == "generativelanguage.googleapis.com"
except Exception:
@ -2015,11 +1946,7 @@ def _build_external_messages(
# tool_calls (some providers reject empty assistant turns).
# Preserve assistant turns whose only payload is tool_calls
# so multi-turn function-call loops round-trip.
if (
msg.role == "assistant"
and not msg.content.strip()
and not msg.tool_calls
):
if msg.role == "assistant" and not msg.content.strip() and not msg.tool_calls:
continue
out: dict[str, Any] = {"role": msg.role, "content": msg.content}
if msg.role == "assistant" and msg.tool_calls:
@ -2073,9 +2000,7 @@ def _build_external_messages(
"image_url": {"url": part.image_url.url},
}
)
elif (
part.type == "reasoning" and openai and msg.role == "assistant"
):
elif part.type == "reasoning" and openai and msg.role == "assistant":
reasoning: dict[str, Any] = {
"type": "reasoning",
"id": part.id,
@ -2085,9 +2010,7 @@ def _build_external_messages(
reasoning["status"] = part.status
parts.append(reasoning)
elif (
part.type == "image_generation_call"
and openai
and msg.role == "assistant"
part.type == "image_generation_call" and openai and msg.role == "assistant"
):
# ExternalProviderClient maps this onto a top-level
# Responses input item after the current user prompt,
@ -2158,11 +2081,7 @@ def _build_external_messages(
if p.status:
reasoning["status"] = p.status
preserved.append(reasoning)
elif (
p.type == "image_generation_call"
and openai
and msg.role == "assistant"
):
elif p.type == "image_generation_call" and openai and msg.role == "assistant":
image_ref = {"type": "image_generation_call", "id": p.id}
if getattr(p, "response_id", None):
image_ref["response_id"] = p.response_id
@ -2187,9 +2106,7 @@ def _build_external_messages(
_entry_content = entry.get("content")
_has_text = (
isinstance(_entry_content, str) and _entry_content.strip()
) or (
isinstance(_entry_content, list) and len(_entry_content) > 0
)
) or (isinstance(_entry_content, list) and len(_entry_content) > 0)
if not _has_text:
continue
if msg.role == "tool":
@ -2204,8 +2121,7 @@ def _build_external_messages(
async def _proxy_to_external_provider(
payload: ChatCompletionRequest,
request: Request,
payload: ChatCompletionRequest, request: Request
) -> StreamingResponse:
"""
Proxy a chat completion request to an external LLM provider.
@ -2343,9 +2259,7 @@ async def _proxy_to_external_provider(
# ── OpenAI shell-tool container management ───────────────────────
def _resolve_openai_cloud_client(
body: OpenAIContainerRequest,
) -> ExternalProviderClient:
def _resolve_openai_cloud_client(body: OpenAIContainerRequest) -> ExternalProviderClient:
"""
Decrypt the API key + validate the base URL points at OpenAI cloud,
then build an ExternalProviderClient for the three container CRUD
@ -2388,9 +2302,7 @@ def _summarize_container(raw: dict) -> OpenAIContainerSummary:
return OpenAIContainerSummary(
id = str(raw.get("id") or ""),
name = raw.get("name"),
created_at = raw.get("created_at")
if isinstance(raw.get("created_at"), int)
else None,
created_at = raw.get("created_at") if isinstance(raw.get("created_at"), int) else None,
last_active_at = raw.get("last_active_at")
if isinstance(raw.get("last_active_at"), int)
else None,
@ -2404,8 +2316,7 @@ def _summarize_container(raw: dict) -> OpenAIContainerSummary:
response_model = ListOpenAIContainersResponse,
)
async def list_openai_containers(
body: OpenAIContainerRequest,
current_subject: str = Depends(get_current_subject),
body: OpenAIContainerRequest, current_subject: str = Depends(get_current_subject)
) -> ListOpenAIContainersResponse:
"""List the user's OpenAI shell-tool containers."""
client = _resolve_openai_cloud_client(body)
@ -2445,8 +2356,7 @@ async def list_openai_containers(
response_model = OpenAIContainerSummary,
)
async def create_openai_container(
body: CreateOpenAIContainerBody,
current_subject: str = Depends(get_current_subject),
body: CreateOpenAIContainerBody, current_subject: str = Depends(get_current_subject)
) -> OpenAIContainerSummary:
"""Create a named container with the user-chosen idle TTL."""
client = _resolve_openai_cloud_client(body)
@ -2482,8 +2392,7 @@ async def create_openai_container(
@router.post("/external/openai/containers/delete", status_code = 204)
async def delete_openai_container(
body: DeleteOpenAIContainerBody,
current_subject: str = Depends(get_current_subject),
body: DeleteOpenAIContainerBody, current_subject: str = Depends(get_current_subject)
) -> None:
"""Delete a named container by id."""
logger.info(
@ -2671,9 +2580,7 @@ async def openai_chat_completions(
id = completion_id,
created = created,
model = model_name,
choices = [
ChunkChoice(delta = ChoiceDelta(), finish_reason = "stop")
],
choices = [ChunkChoice(delta = ChoiceDelta(), finish_reason = "stop")],
)
yield f"data: {final_chunk.model_dump_json(exclude_none = True)}\n\n"
yield "data: [DONE]\n\n"
@ -2681,9 +2588,7 @@ async def openai_chat_completions(
cancel_event.set()
raise
except Exception as e:
logger.error(
f"Error during audio input streaming: {e}", exc_info = True
)
logger.error(f"Error during audio input streaming: {e}", exc_info = True)
yield f"data: {json.dumps({'error': {'message': _friendly_error(e), 'type': 'server_error'}})}\n\n"
finally:
_tracker.__exit__(None, None, None)
@ -2783,9 +2688,7 @@ async def openai_chat_completions(
)
# ── Parse messages (handles multimodal content parts) ─────
system_prompt, chat_messages, extracted_image_b64 = _extract_content_parts(
payload.messages
)
system_prompt, chat_messages, extracted_image_b64 = _extract_content_parts(payload.messages)
if not chat_messages:
raise HTTPException(
@ -2835,9 +2738,7 @@ async def openai_chat_completions(
tools_to_use = []
elif payload.enabled_tools is not None:
tools_to_use = [
t
for t in ALL_TOOLS
if t["function"]["name"] in payload.enabled_tools
t for t in ALL_TOOLS if t["function"]["name"] in payload.enabled_tools
]
else:
tools_to_use = ALL_TOOLS
@ -2901,8 +2802,7 @@ async def openai_chat_completions(
_nudge = (
_date_line + " "
"You have access to tools. When appropriate, prefer using "
"tools rather than answering from memory. "
+ " ".join(_tool_tip_parts)
"tools rather than answering from memory. " + " ".join(_tool_tip_parts)
)
else:
_nudge = ""
@ -2914,15 +2814,11 @@ async def openai_chat_completions(
system_prompt = system_prompt.rstrip() + "\n\n" + _nudge
else:
system_prompt = _nudge
gguf_messages = _set_or_prepend_system_message(
gguf_messages, system_prompt
)
gguf_messages = _set_or_prepend_system_message(gguf_messages, system_prompt)
# ── Strip stale tool-call XML from conversation history ─
for _msg in gguf_messages:
if _msg.get("role") == "assistant" and isinstance(
_msg.get("content"), str
):
if _msg.get("role") == "assistant" and isinstance(_msg.get("content"), str):
_msg["content"] = _TOOL_XML_RE.sub("", _msg["content"]).strip()
def gguf_generate_with_tools():
@ -3058,9 +2954,7 @@ async def openai_chat_completions(
if _stream_usage or _stream_timings:
usage_obj = CompletionUsage(
prompt_tokens = (_stream_usage or {}).get("prompt_tokens", 0),
completion_tokens = (_stream_usage or {}).get(
"completion_tokens", 0
),
completion_tokens = (_stream_usage or {}).get("completion_tokens", 0),
total_tokens = (_stream_usage or {}).get("total_tokens", 0),
)
usage_chunk = ChatCompletionChunk(
@ -3167,11 +3061,7 @@ async def openai_chat_completions(
else:
logger.warning(
"gguf_stream_chunks: unexpected dict event: %s",
{
k: v
for k, v in cumulative.items()
if k != "timings"
},
{k: v for k, v in cumulative.items() if k != "timings"},
)
continue
new_text = cumulative[len(prev_text) :]
@ -3208,9 +3098,7 @@ async def openai_chat_completions(
if _stream_usage or _stream_timings:
usage_obj = CompletionUsage(
prompt_tokens = (_stream_usage or {}).get("prompt_tokens", 0),
completion_tokens = (_stream_usage or {}).get(
"completion_tokens", 0
),
completion_tokens = (_stream_usage or {}).get("completion_tokens", 0),
total_tokens = (_stream_usage or {}).get("total_tokens", 0),
)
usage_chunk = ChatCompletionChunk(
@ -3270,12 +3158,8 @@ async def openai_chat_completions(
)
],
usage = CompletionUsage(
prompt_tokens = (completion_usage or {}).get("prompt_tokens")
or 0,
completion_tokens = (completion_usage or {}).get(
"completion_tokens"
)
or 0,
prompt_tokens = (completion_usage or {}).get("prompt_tokens") or 0,
completion_tokens = (completion_usage or {}).get("completion_tokens") or 0,
total_tokens = (completion_usage or {}).get("total_tokens") or 0,
),
)
@ -3335,16 +3219,12 @@ async def openai_chat_completions(
# XML -- gpt-oss tools still work via the GGUF path).
_sf_is_gptoss = False
try:
_sf_is_gptoss = bool(
hasattr(backend, "_is_gpt_oss_model") and backend._is_gpt_oss_model()
)
_sf_is_gptoss = bool(hasattr(backend, "_is_gpt_oss_model") and backend._is_gpt_oss_model())
except Exception:
_sf_is_gptoss = False
_sf_tool_budget = (
payload.max_tool_calls_per_message
if payload.max_tool_calls_per_message is not None
else 25
payload.max_tool_calls_per_message if payload.max_tool_calls_per_message is not None else 25
)
# Match the GGUF path: mcp_enabled also opens the tool loop on its own
@ -3426,8 +3306,7 @@ async def openai_chat_completions(
_sf_nudge = (
_sf_date_line + " "
"You have access to tools. When appropriate, prefer using "
"tools rather than answering from memory. "
+ " ".join(_sf_tool_tip_parts)
"tools rather than answering from memory. " + " ".join(_sf_tool_tip_parts)
)
else:
_sf_nudge = ""
@ -3940,9 +3819,7 @@ async def serve_sandbox_file(
@router.get("/models")
async def openai_list_models(
current_subject: str = Depends(get_current_subject),
):
async def openai_list_models(current_subject: str = Depends(get_current_subject)):
"""
OpenAI-compatible model listing endpoint.
@ -3982,10 +3859,7 @@ async def openai_list_models(
@router.post("/completions")
async def openai_completions(
request: Request,
current_subject: str = Depends(get_current_subject),
):
async def openai_completions(request: Request, current_subject: str = Depends(get_current_subject)):
"""
OpenAI-compatible text completions endpoint (non-chat).
@ -4059,10 +3933,7 @@ async def openai_completions(
@router.post("/embeddings")
async def openai_embeddings(
request: Request,
current_subject: str = Depends(get_current_subject),
):
async def openai_embeddings(request: Request, current_subject: str = Depends(get_current_subject)):
"""
OpenAI-compatible embeddings endpoint.
@ -4095,9 +3966,7 @@ async def openai_embeddings(
# =====================================================================
def _translate_responses_tools_to_chat(
tools: Optional[list[dict]],
) -> Optional[list[dict]]:
def _translate_responses_tools_to_chat(tools: Optional[list[dict]]) -> Optional[list[dict]]:
"""Translate Responses-shape function tools to the Chat Completions nested shape.
Responses uses a flat shape per tool entry::
@ -4371,9 +4240,7 @@ def _chat_tool_calls_to_responses_output(tool_calls: list[dict]) -> list[dict]:
async def _responses_non_streaming(
payload: ResponsesRequest,
messages: list[ChatMessage],
request: Request,
payload: ResponsesRequest, messages: list[ChatMessage], request: Request
) -> JSONResponse:
"""Handle a non-streaming Responses API call."""
chat_req = _build_chat_request(payload, messages, stream = False)
@ -4439,9 +4306,7 @@ async def _responses_non_streaming(
async def _responses_stream(
payload: ResponsesRequest,
messages: list[ChatMessage],
request: Request,
payload: ResponsesRequest, messages: list[ChatMessage], request: Request
):
"""Handle a streaming Responses API call, emitting named SSE events.
@ -4491,8 +4356,7 @@ async def _responses_stream(
# Direct pass-through bypasses the openai_chat_completions image gate.
if not llama_backend.is_vision and any(
isinstance(m.content, list)
and any(isinstance(p, ImageContentPart) for p in m.content)
isinstance(m.content, list) and any(isinstance(p, ImageContentPart) for p in m.content)
for m in messages
):
raise HTTPException(
@ -4500,9 +4364,7 @@ async def _responses_stream(
detail = "Image provided but current GGUF model does not support vision.",
)
body = _build_openai_passthrough_body(
chat_req, backend_ctx = llama_backend.context_length
)
body = _build_openai_passthrough_body(chat_req, backend_ctx = llama_backend.context_length)
target_url = f"{llama_backend.base_url}/v1/chat/completions"
async def event_generator():
@ -4861,9 +4723,7 @@ def _anthropic_requested_studio_tools(tools: Optional[list]) -> set[str]:
def _select_anthropic_server_tools(
all_tools: list[dict],
requested_studio_tools: set[str],
enabled_tools: Optional[list[str]],
all_tools: list[dict], requested_studio_tools: set[str], enabled_tools: Optional[list[str]]
) -> list[dict]:
"""Select Studio tools requested through Anthropic tools and extensions."""
if not requested_studio_tools and enabled_tools is None:
@ -4876,9 +4736,7 @@ def _select_anthropic_server_tools(
return [tool for tool in all_tools if tool["function"]["name"] in selected_names]
def _normalize_anthropic_openai_images(
openai_messages: list[dict], is_vision: bool
) -> bool:
def _normalize_anthropic_openai_images(openai_messages: list[dict], is_vision: bool) -> bool:
"""Enforce the vision guard on translated Anthropic messages and
normalize any ``image_url`` parts with base64 data URLs to PNG.
@ -4991,9 +4849,7 @@ async def anthropic_messages(
# Enforce vision guard + re-encode embedded images to PNG so the
# Anthropic endpoint matches the behavior of /v1/chat/completions.
_has_image = _normalize_anthropic_openai_images(
openai_messages, llama_backend.is_vision
)
_has_image = _normalize_anthropic_openai_images(openai_messages, llama_backend.is_vision)
temperature = payload.temperature if payload.temperature is not None else 0.6
top_p = payload.top_p if payload.top_p is not None else 0.95
@ -5002,9 +4858,7 @@ async def anthropic_messages(
repetition_penalty = (
payload.repetition_penalty if payload.repetition_penalty is not None else 1.0
)
presence_penalty = (
payload.presence_penalty if payload.presence_penalty is not None else 0.0
)
presence_penalty = payload.presence_penalty if payload.presence_penalty is not None else 0.0
stop = payload.stop_sequences or None
# Translate Anthropic tool_choice to OpenAI format for forwarding to
@ -5083,9 +4937,7 @@ async def anthropic_messages(
and not _has_image
)
client_tools = (
not server_tools
and len(openai_client_tools) > 0
and llama_backend.supports_tools
not server_tools and len(openai_client_tools) > 0 and llama_backend.supports_tools
)
# ── Client-side pass-through path ─────────────────────────
@ -5265,13 +5117,7 @@ async def anthropic_messages(
)
async def _anthropic_tool_stream(
request,
cancel_event,
run_gen,
message_id,
model_name,
):
async def _anthropic_tool_stream(request, cancel_event, run_gen, message_id, model_name):
"""Streaming response for the tool-calling path."""
_sentinel = object()
@ -5312,13 +5158,7 @@ async def _anthropic_tool_stream(
)
async def _anthropic_plain_stream(
request,
cancel_event,
run_gen,
message_id,
model_name,
):
async def _anthropic_plain_stream(request, cancel_event, run_gen, message_id, model_name):
"""Streaming response for the no-tool path."""
_sentinel = object()
@ -5388,18 +5228,14 @@ async def _anthropic_tool_non_streaming(run_gen, message_id, model_name):
new = clean[len(prev_text) :]
prev_text = clean
if new:
if content_blocks and isinstance(
content_blocks[-1], AnthropicResponseTextBlock
):
if content_blocks and isinstance(content_blocks[-1], AnthropicResponseTextBlock):
content_blocks[-1].text += new
else:
content_blocks.append(AnthropicResponseTextBlock(text = new))
elif etype == "tool_start":
tool_call_id = event["tool_call_id"]
arguments = event.get("arguments", {})
existing_tool_block = (
tool_blocks_by_id.get(tool_call_id) if tool_call_id else None
)
existing_tool_block = tool_blocks_by_id.get(tool_call_id) if tool_call_id else None
if existing_tool_block is not None:
if arguments or not existing_tool_block.input:
existing_tool_block.input = arguments
@ -5500,9 +5336,7 @@ def _build_passthrough_payload(
if stream:
body["stream_options"] = {"include_usage": True}
body["max_tokens"] = (
max_tokens
if max_tokens is not None
else (backend_ctx or _DEFAULT_MAX_TOKENS_FLOOR)
max_tokens if max_tokens is not None else (backend_ctx or _DEFAULT_MAX_TOKENS_FLOOR)
)
body["t_max_predict_ms"] = _DEFAULT_T_MAX_PREDICT_MS
if stop:
@ -5619,9 +5453,7 @@ async def _anthropic_passthrough_stream(
# blocks during llama-server prefill, so the in-loop cancel
# check is unreachable until the first SSE chunk arrives.
# The watcher closes `resp` on cancel, raising in aiter_lines.
cancel_watcher = asyncio.create_task(
_await_cancel_then_close(cancel_event, resp)
)
cancel_watcher = asyncio.create_task(_await_cancel_then_close(cancel_event, resp))
lines_iter = resp.aiter_lines()
async for raw_line in lines_iter:
if cancel_event.is_set():
@ -5855,9 +5687,7 @@ def _strip_provider_synthetic_tool_history(messages: list[dict]) -> list[dict]:
if args_obj.get("_server_tool") is True:
is_synthetic = True
google = args_obj.get("google")
if isinstance(google, dict) and isinstance(
google.get("native_part"), dict
):
if isinstance(google, dict) and isinstance(google.get("native_part"), dict):
is_synthetic = True
if is_synthetic:
tc_id = tc.get("id")
@ -5912,9 +5742,7 @@ def _openai_messages_for_passthrough(payload) -> list[dict]:
transparently.
"""
messages = _strip_provider_synthetic_tool_history(
_drop_empty_assistant_sentinels(
[m.model_dump(exclude_none = True) for m in payload.messages]
)
_drop_empty_assistant_sentinels([m.model_dump(exclude_none = True) for m in payload.messages])
)
if not payload.image_base64:
@ -5964,9 +5792,7 @@ def _openai_messages_for_gguf_chat(payload, is_vision: bool) -> tuple[list[dict]
image attached to its original turn.
"""
messages = _strip_provider_synthetic_tool_history(
_drop_empty_assistant_sentinels(
[m.model_dump(exclude_none = True) for m in payload.messages]
)
_drop_empty_assistant_sentinels([m.model_dump(exclude_none = True) for m in payload.messages])
)
has_message_image = any(
isinstance(msg.get("content"), list)
@ -6048,12 +5874,7 @@ def _build_openai_passthrough_body(payload, backend_ctx = None) -> dict:
async def _openai_passthrough_stream(
request,
cancel_event,
llama_backend,
payload,
model_name,
completion_id,
request, cancel_event, llama_backend, payload, model_name, completion_id
):
"""Streaming client-side pass-through for /v1/chat/completions.
@ -6064,9 +5885,7 @@ async def _openai_passthrough_stream(
observes a standard OpenAI response.
"""
target_url = f"{llama_backend.base_url}/v1/chat/completions"
body = _build_openai_passthrough_body(
payload, backend_ctx = llama_backend.context_length
)
body = _build_openai_passthrough_body(payload, backend_ctx = llama_backend.context_length)
_cancel_keys = (payload.cancel_id, payload.session_id, completion_id)
_tracker = _TrackedCancel(cancel_event, *_cancel_keys)
@ -6139,9 +5958,7 @@ async def _openai_passthrough_stream(
# tiny watcher that closes `resp` as soon as cancel fires,
# unblocking the iterator with a RemoteProtocolError caught
# in the except clause below.
cancel_watcher = asyncio.create_task(
_await_cancel_then_close(cancel_event, resp)
)
cancel_watcher = asyncio.create_task(_await_cancel_then_close(cancel_event, resp))
try:
lines_iter = resp.aiter_lines()
async for raw_line in lines_iter:
@ -6209,11 +6026,7 @@ async def _openai_passthrough_stream(
raise
async def _openai_passthrough_non_streaming(
llama_backend,
payload,
model_name,
):
async def _openai_passthrough_non_streaming(llama_backend, payload, model_name):
"""Non-streaming client-side pass-through for /v1/chat/completions.
Returns llama-server's JSON response verbatim (via JSONResponse) so the
@ -6222,9 +6035,7 @@ async def _openai_passthrough_non_streaming(
token counts.
"""
target_url = f"{llama_backend.base_url}/v1/chat/completions"
body = _build_openai_passthrough_body(
payload, backend_ctx = llama_backend.context_length
)
body = _build_openai_passthrough_body(payload, backend_ctx = llama_backend.context_length)
try:
async with httpx.AsyncClient() as client:

37
studio/backend/routes/mcp_servers.py
View file

@ -112,16 +112,13 @@ def _row_to_response(row: dict) -> McpServerResponse:
@router.get("/", response_model = list[McpServerResponse])
async def list_mcp_servers(
current_subject: str = Depends(get_current_subject),
):
async def list_mcp_servers(current_subject: str = Depends(get_current_subject)):
return [_row_to_response(row) for row in mcp_servers_db.list_servers()]
@router.post("/", response_model = McpServerResponse, status_code = 201)
async def create_mcp_server(
payload: McpServerCreate,
current_subject: str = Depends(get_current_subject),
payload: McpServerCreate, current_subject: str = Depends(get_current_subject)
):
display_name = (payload.display_name or "").strip()
if not display_name:
@ -151,9 +148,7 @@ def _changes_from_payload(payload: McpServerUpdate) -> dict:
if "display_name" in sent:
name = (payload.display_name or "").strip()
if not name:
raise HTTPException(
status_code = 400, detail = "display_name must not be empty"
)
raise HTTPException(status_code = 400, detail = "display_name must not be empty")
changes["display_name"] = name
if "url" in sent:
changes["url"] = _validate_url(payload.url or "")
@ -162,15 +157,11 @@ def _changes_from_payload(payload: McpServerUpdate) -> dict:
changes["headers_json"] = json.dumps(headers) if headers else None
if "is_enabled" in sent:
if payload.is_enabled is None:
raise HTTPException(
status_code = 400, detail = "is_enabled must be true or false"
)
raise HTTPException(status_code = 400, detail = "is_enabled must be true or false")
changes["is_enabled"] = payload.is_enabled
if "use_oauth" in sent:
if payload.use_oauth is None:
raise HTTPException(
status_code = 400, detail = "use_oauth must be true or false"
)
raise HTTPException(status_code = 400, detail = "use_oauth must be true or false")
changes["use_oauth"] = payload.use_oauth
# stdio is OAuth-less: drop a stale OAuth flag when switching to a command.
if "url" in changes and is_stdio(changes["url"]):
@ -203,8 +194,7 @@ async def update_mcp_server(
# disabled; fastmcp keys tokens by URL and would otherwise let a
# re-pointed server silently inherit the old account's credentials.
if bool(old.get("use_oauth")) and (
("url" in changes and changes["url"] != old["url"])
or changes.get("use_oauth") is False
("url" in changes and changes["url"] != old["url"]) or changes.get("use_oauth") is False
):
await clear_oauth_tokens_async(old["url"])
mcp_servers_db.update_server(server_id, changes)
@ -212,10 +202,7 @@ async def update_mcp_server(
@router.delete("/{server_id}", status_code = 204)
async def delete_mcp_server(
server_id: str,
current_subject: str = Depends(get_current_subject),
):
async def delete_mcp_server(server_id: str, current_subject: str = Depends(get_current_subject)):
old = mcp_servers_db.get_server(server_id)
if not old:
raise HTTPException(status_code = 404, detail = "MCP server not found")
@ -226,8 +213,7 @@ async def delete_mcp_server(
@router.post("/{server_id}/refresh", response_model = McpServerProbeResult)
async def refresh_mcp_server_tools(
server_id: str,
current_subject: str = Depends(get_current_subject),
server_id: str, current_subject: str = Depends(get_current_subject)
):
server = mcp_servers_db.get_server(server_id)
if not server:
@ -235,9 +221,7 @@ async def refresh_mcp_server_tools(
# Refresh uses the stored address, so re-check the stdio gate here too: a
# stdio row from a desktop DB must not spawn on a hosted/network host.
if is_stdio(server["url"]) and not stdio_mcp_enabled():
raise HTTPException(
status_code = 400, detail = "stdio MCP servers are disabled on this host"
)
raise HTTPException(status_code = 400, detail = "stdio MCP servers are disabled on this host")
use_oauth = bool(server.get("use_oauth"))
try:
@ -261,8 +245,7 @@ async def refresh_mcp_server_tools(
@router.post("/test", response_model = McpServerProbeResult)
async def test_mcp_server(
payload: McpServerTestRequest,
current_subject: str = Depends(get_current_subject),
payload: McpServerTestRequest, current_subject: str = Depends(get_current_subject)
):
# URL/header validation must surface as 400 like create/update so the
# frontend's create-form pre-flight gets the same error semantics as

141
studio/backend/routes/models.py
View file

@ -141,7 +141,9 @@ logger = get_logger(__name__)
def derive_model_type(
is_vision: bool, audio_type: Optional[str], is_embedding: bool = False
is_vision: bool,
audio_type: Optional[str],
is_embedding: bool = False,
) -> ModelType:
"""Collapse individual capability flags into a single model modality string."""
if is_embedding:
@ -157,7 +159,6 @@ def _resolve_hf_cache_dir() -> Path:
"""Resolve local HF cache root used by hub downloads."""
try:
from huggingface_hub.constants import HF_HUB_CACHE
return Path(HF_HUB_CACHE)
except Exception:
return Path.home() / ".cache" / "huggingface" / "hub"
@ -194,9 +195,7 @@ def _is_model_directory(d: Path) -> bool:
return False
try:
has_config = (d / "config.json").exists() or (
d / "adapter_config.json"
).exists()
has_config = (d / "config.json").exists() or (d / "adapter_config.json").exists()
if not has_config:
return False
return any(_is_weight_file(f) for f in d.iterdir() if f.is_file())
@ -204,11 +203,7 @@ def _is_model_directory(d: Path) -> bool:
return False
def _scan_models_dir(
models_dir: Path,
*,
limit: int | None = None,
) -> List[LocalModelInfo]:
def _scan_models_dir(models_dir: Path, *, limit: int | None = None) -> List[LocalModelInfo]:
if not models_dir.exists() or not models_dir.is_dir():
return []
@ -449,8 +444,7 @@ def _ollama_links_dir(ollama_dir: Path) -> Optional[Path]:
return primary
except OSError as e:
logger.debug(
"Ollama dir %s not writable for .studio_links (%s); "
"falling back to Studio cache",
"Ollama dir %s not writable for .studio_links (%s); falling back to Studio cache",
ollama_dir,
e,
)
@ -475,9 +469,7 @@ def _ollama_links_dir(ollama_dir: Path) -> Optional[Path]:
return None
def _scan_ollama_dir(
ollama_dir: Path, limit: Optional[int] = None
) -> List[LocalModelInfo]:
def _scan_ollama_dir(ollama_dir: Path, limit: Optional[int] = None) -> List[LocalModelInfo]:
"""Scan an Ollama models directory for downloaded models.
Ollama stores models in a content-addressable layout::
@ -571,9 +563,7 @@ def _scan_ollama_dir(
if tmp_path.is_symlink() or tmp_path.exists():
tmp_path.unlink()
except OSError as cleanup_err:
logger.debug(
"Could not clean up tmp path %s: %s", tmp_path, cleanup_err
)
logger.debug("Could not clean up tmp path %s: %s", tmp_path, cleanup_err)
return None
try:
@ -590,11 +580,7 @@ def _scan_ollama_dir(
repo_parts = list(parts[1:-1])
tag = parts[-1]
if (
host == "registry.ollama.ai"
and repo_parts
and repo_parts[0] == "library"
):
if host == "registry.ollama.ai" and repo_parts and repo_parts[0] == "library":
repo_name = "/".join(repo_parts[1:])
elif host == "registry.ollama.ai":
repo_name = "/".join(repo_parts)
@ -651,9 +637,7 @@ def _scan_ollama_dir(
candidate = blobs_dir / digest.replace(":", "-")
if candidate.is_file():
link_name = f"{safe_name}-{tag}{quant}.gguf"
gguf_link_path = _make_link(
model_link_dir, link_name, candidate
)
gguf_link_path = _make_link(model_link_dir, link_name, candidate)
elif media == "application/vnd.ollama.image.projector":
candidate = blobs_dir / digest.replace(":", "-")
@ -726,7 +710,6 @@ async def list_local_models(
allowed_roots.append(hf_default)
try:
from utils.paths import studio_root, outputs_root
allowed_roots.extend([studio_root(), outputs_root()])
except Exception:
pass
@ -785,10 +768,7 @@ async def list_local_models(
+ _scan_hf_cache(folder_path)
+ _scan_lmstudio_dir(folder_path)
)
if not any(
p in (".studio_links", "ollama_links")
for p in Path(m.path).parts
)
if not any(p in (".studio_links", "ollama_links") for p in Path(m.path).parts)
]
custom_models = _generic
if len(custom_models) < _MAX_MODELS_PER_FOLDER:
@ -799,9 +779,7 @@ async def list_local_models(
except OSError as e:
logger.warning("Skipping unreadable scan folder %s: %s", folder_path, e)
continue
local_models += [
m.model_copy(update = {"source": "custom"}) for m in custom_models
]
local_models += [m.model_copy(update = {"source": "custom"}) for m in custom_models]
# Deduplicate models, but always keep custom folder entries so they
# appear in the "Custom Folders" UI section even when the same model
@ -836,19 +814,15 @@ async def list_local_models(
@router.get("/scan-folders")
async def get_scan_folders(
current_subject: str = Depends(get_current_subject),
):
async def get_scan_folders(current_subject: str = Depends(get_current_subject)):
"""List all registered custom model scan folders."""
from storage.studio_db import list_scan_folders
return {"folders": list_scan_folders()}
@router.post("/scan-folders", response_model = ScanFolderInfo, status_code = 201)
async def add_scan_folder_endpoint(
body: AddScanFolderRequest,
current_subject: str = Depends(get_current_subject),
body: AddScanFolderRequest, current_subject: str = Depends(get_current_subject)
):
"""Register a new directory to scan for local models."""
from storage.studio_db import add_scan_folder
@ -867,8 +841,7 @@ async def add_scan_folder_endpoint(
@router.delete("/scan-folders/{folder_id}")
async def remove_scan_folder_endpoint(
folder_id: int,
current_subject: str = Depends(get_current_subject),
folder_id: int, current_subject: str = Depends(get_current_subject)
):
"""Remove a registered custom scan folder."""
from storage.studio_db import remove_scan_folder
@ -879,9 +852,7 @@ async def remove_scan_folder_endpoint(
@router.get("/recommended-folders")
async def get_recommended_folders(
current_subject: str = Depends(get_current_subject),
):
async def get_recommended_folders(current_subject: str = Depends(get_current_subject)):
"""Return well-known model directories that exist on this machine.
Lightweight alternative to ``browse-folders`` for showing quick-pick
@ -1191,9 +1162,7 @@ def _match_browse_child(current: Path, name: str) -> Optional[Path]:
detail = f"Permission denied reading {current.name}",
) from None
except OSError as exc:
logger.warning(
"browse-folders: could not read %s: %s", current, exc, exc_info = True
)
logger.warning("browse-folders: could not read %s: %s", current, exc, exc_info = True)
raise HTTPException(
status_code = 500,
detail = f"Could not read {os.path.basename(str(current))}",
@ -1344,9 +1313,7 @@ async def browse_folders(
detail = f"Permission denied reading {os.path.basename(str(target))}",
)
except OSError as exc:
logger.warning(
"browse-folders: could not read %s: %s", target, exc, exc_info = True
)
logger.warning("browse-folders: could not read %s: %s", target, exc, exc_info = True)
raise HTTPException(
status_code = 500,
detail = f"Could not read {os.path.basename(str(target))}",
@ -1404,9 +1371,7 @@ async def browse_folders(
# would 403 on click. Users can still hop to other allowed roots
# via the suggestion chips below.
parent: Optional[str]
if target.parent == target or not _is_path_inside_allowlist(
target.parent, allowed_roots
):
if target.parent == target or not _is_path_inside_allowlist(target.parent, allowed_roots):
parent = None
else:
parent = str(target.parent)
@ -1475,9 +1440,7 @@ def _looks_like_mlx_repo(model_id: str) -> bool:
@router.get("/list")
async def list_models(
current_subject: str = Depends(get_current_subject),
):
async def list_models(current_subject: str = Depends(get_current_subject)):
"""
List available models (default models and loaded models).
@ -1565,16 +1528,12 @@ def _get_max_position_embeddings(config) -> Optional[int]:
"""Extract max_position_embeddings from a model config, checking text_config fallback."""
if hasattr(config, "max_position_embeddings"):
return config.max_position_embeddings
if hasattr(config, "text_config") and hasattr(
config.text_config, "max_position_embeddings"
):
if hasattr(config, "text_config") and hasattr(config.text_config, "max_position_embeddings"):
return config.text_config.max_position_embeddings
return None
def _get_model_size_bytes(
model_name: str, hf_token: Optional[str] = None
) -> Optional[int]:
def _get_model_size_bytes(model_name: str, hf_token: Optional[str] = None) -> Optional[int]:
"""Get total size of model weight files from HF Hub."""
try:
from huggingface_hub import HfApi
@ -1587,9 +1546,7 @@ def _get_model_size_bytes(
weight_exts = (".safetensors", ".bin", ".pt", ".pth", ".gguf")
total = 0
for sibling in info.siblings:
if sibling.rfilename and any(
sibling.rfilename.endswith(ext) for ext in weight_exts
):
if sibling.rfilename and any(sibling.rfilename.endswith(ext) for ext in weight_exts):
if sibling.size is not None:
total += sibling.size
@ -1777,15 +1734,10 @@ def _loaded_model_matches_deleted_path(active_model: str, deleted_path: Path) ->
)
active_lower = active_model.lower()
target_lower = str(deleted_path).lower()
return active_lower == target_lower or active_lower.startswith(
f"{target_lower}{os.sep}"
)
return active_lower == target_lower or active_lower.startswith(f"{target_lower}{os.sep}")
def _loading_model_matches_deleted_path(
loading_model: object,
deleted_path: Path,
) -> bool:
def _loading_model_matches_deleted_path(loading_model: object, deleted_path: Path) -> bool:
if not loading_model:
return False
return _loaded_model_matches_deleted_path(str(loading_model), deleted_path)
@ -1918,7 +1870,6 @@ async def delete_finetuned_model(
if source == "training":
try:
from core.training import get_training_backend
training_backend = get_training_backend()
if training_backend.is_training_active():
raise HTTPException(
@ -2005,9 +1956,7 @@ async def delete_finetuned_model(
except HTTPException:
raise
except Exception as e:
logger.warning(
"Could not check inference backend loaded model before delete: %s", e
)
logger.warning("Could not check inference backend loaded model before delete: %s", e)
raise HTTPException(
status_code = 503,
detail = "Could not verify model load status before deleting",
@ -2079,10 +2028,7 @@ async def delete_finetuned_model(
@router.get("/loras/{lora_path:path}/base-model", response_model = LoRABaseModelResponse)
async def get_lora_base_model(
lora_path: str,
current_subject: str = Depends(get_current_subject),
):
async def get_lora_base_model(lora_path: str, current_subject: str = Depends(get_current_subject)):
"""
Get the base model for a LoRA adapter.
@ -2115,10 +2061,7 @@ async def get_lora_base_model(
@router.get("/check-vision/{model_name:path}", response_model = VisionCheckResponse)
async def check_vision_model(
model_name: str,
current_subject: str = Depends(get_current_subject),
):
async def check_vision_model(model_name: str, current_subject: str = Depends(get_current_subject)):
"""
Check if a model is a vision model.
@ -2159,9 +2102,7 @@ async def check_embedding_model(
logger.info(f"Checking if embedding model: {model_name}")
is_embedding = is_embedding_model(model_name, hf_token = hf_token)
logger.info(
f"Embedding check result for {model_name}: is_embedding={is_embedding}"
)
logger.info(f"Embedding check result for {model_name}: is_embedding={is_embedding}")
return EmbeddingCheckResponse(
model_name = model_name,
is_embedding = is_embedding,
@ -2182,9 +2123,7 @@ async def get_gguf_variants(
repo_id: str = Query(
..., description = "HuggingFace repo ID (e.g. 'unsloth/gemma-3-4b-it-GGUF')"
),
hf_token: Optional[str] = Query(
None, description = "HuggingFace token for private repos"
),
hf_token: Optional[str] = Query(None, description = "HuggingFace token for private repos"),
current_subject: str = Depends(get_current_subject),
):
"""
@ -2333,11 +2272,7 @@ async def get_gguf_download_progress(
break
total_progress_bytes = downloaded_bytes + in_progress_bytes
progress = (
min(total_progress_bytes / expected_bytes, 0.99)
if expected_bytes > 0
else 0
)
progress = min(total_progress_bytes / expected_bytes, 0.99) if expected_bytes > 0 else 0
# Only report 1.0 when all bytes are in completed files (not in-progress)
if expected_bytes > 0 and downloaded_bytes >= expected_bytes:
progress = 1.0
@ -2480,7 +2415,6 @@ def _all_hf_cache_scans():
try:
# Resolve the active cache dir so we can dedup
from huggingface_hub.constants import HF_HUB_CACHE
seen.add(str(Path(HF_HUB_CACHE).resolve()))
except Exception:
pass
@ -2553,9 +2487,7 @@ def _repo_has_gguf_files(repo_info) -> bool:
@router.get("/cached-gguf")
async def list_cached_gguf(
current_subject: str = Depends(get_current_subject),
):
async def list_cached_gguf(current_subject: str = Depends(get_current_subject)):
"""List GGUF repos downloaded to HF cache, legacy Unsloth cache, and HF default cache."""
try:
cache_scans = _all_hf_cache_scans()
@ -2590,9 +2522,7 @@ async def list_cached_gguf(
@router.get("/cached-models")
async def list_cached_models(
current_subject: str = Depends(get_current_subject),
):
async def list_cached_models(current_subject: str = Depends(get_current_subject)):
"""List non-GGUF model repos downloaded to HF cache, legacy Unsloth cache, and HF default cache."""
_WEIGHT_EXTENSIONS = (".safetensors", ".bin")
@ -2609,9 +2539,7 @@ async def list_cached_models(
if _repo_has_gguf_files(repo_info):
continue
total_size = sum(
(f.size_on_disk or 0)
for rev in repo_info.revisions
for f in rev.files
(f.size_on_disk or 0) for rev in repo_info.revisions for f in rev.files
)
if total_size == 0:
continue
@ -2658,7 +2586,6 @@ async def delete_cached_model(
# Check if model is currently loaded
try:
from routes.inference import get_llama_cpp_backend
llama_backend = get_llama_cpp_backend()
if llama_backend.is_loaded and llama_backend.model_identifier:
loaded_id = llama_backend.model_identifier.lower()

41
studio/backend/routes/providers.py
View file

@ -51,9 +51,7 @@ router = APIRouter()
@router.get("/public-key")
async def get_public_key(
current_subject: str = Depends(get_current_subject),
):
async def get_public_key(current_subject: str = Depends(get_current_subject)):
"""Return the RSA public key PEM for client-side API key encryption.
The ``fingerprint`` field is a short SHA256 of the PEM and is meant
@ -72,9 +70,7 @@ async def get_public_key(
@router.get("/registry", response_model = list[ProviderRegistryEntry])
async def list_registry(
current_subject: str = Depends(get_current_subject),
):
async def list_registry(current_subject: str = Depends(get_current_subject)):
"""List all supported provider types with their default configurations."""
return list_available_providers()
@ -83,9 +79,7 @@ async def list_registry(
@router.get("/pricing")
async def get_pricing_snapshot(
current_subject: str = Depends(get_current_subject),
):
async def get_pricing_snapshot(current_subject: str = Depends(get_current_subject)):
"""Static per-MTok pricing table the frontend uses to convert
upstream usage chunks into a per-turn USD cost. See
``core/inference/pricing.py`` for sourcing notes; values reflect
@ -97,9 +91,7 @@ async def get_pricing_snapshot(
@router.get("/", response_model = list[ProviderResponse])
async def list_provider_configs(
current_subject: str = Depends(get_current_subject),
):
async def list_provider_configs(current_subject: str = Depends(get_current_subject)):
"""List all saved provider configurations."""
rows = providers_db.list_providers()
return [
@ -118,8 +110,7 @@ async def list_provider_configs(
@router.post("/", response_model = ProviderResponse, status_code = 201)
async def create_provider_config(
payload: ProviderCreate,
current_subject: str = Depends(get_current_subject),
payload: ProviderCreate, current_subject: str = Depends(get_current_subject)
):
"""Create a new saved provider configuration (no API key stored)."""
info = get_provider_info(payload.provider_type)
@ -186,8 +177,7 @@ async def update_provider_config(
@router.delete("/{provider_id}", status_code = 204)
async def delete_provider_config(
provider_id: str,
current_subject: str = Depends(get_current_subject),
provider_id: str, current_subject: str = Depends(get_current_subject)
):
"""Delete a saved provider configuration."""
deleted = providers_db.delete_provider(provider_id)
@ -200,8 +190,7 @@ async def delete_provider_config(
@router.post("/test", response_model = ProviderTestResult)
async def test_provider(
payload: ProviderTestRequest,
current_subject: str = Depends(get_current_subject),
payload: ProviderTestRequest, current_subject: str = Depends(get_current_subject)
):
"""
Test connectivity to an external provider.
@ -221,9 +210,7 @@ async def test_provider(
try:
api_key = decrypt_api_key(payload.encrypted_api_key)
except Exception as exc:
logger.warning(
"Failed to decrypt API key (%s): %s", type(exc).__name__, exc
)
logger.warning("Failed to decrypt API key (%s): %s", type(exc).__name__, exc)
raise HTTPException(
status_code = 400,
detail = "Failed to decrypt API key. The public key may have changed — try refreshing the page.",
@ -275,8 +262,7 @@ async def test_provider(
@router.post("/models", response_model = list[ProviderModelInfo])
async def list_provider_models(
payload: ProviderModelsRequest,
current_subject: str = Depends(get_current_subject),
payload: ProviderModelsRequest, current_subject: str = Depends(get_current_subject)
):
"""
List models available from an external provider.
@ -295,9 +281,7 @@ async def list_provider_models(
try:
api_key = decrypt_api_key(payload.encrypted_api_key)
except Exception as exc:
logger.warning(
"Failed to decrypt API key (%s): %s", type(exc).__name__, exc
)
logger.warning("Failed to decrypt API key (%s): %s", type(exc).__name__, exc)
raise HTTPException(
status_code = 400,
detail = "Failed to decrypt API key. The public key may have changed — try refreshing the page.",
@ -338,7 +322,6 @@ async def list_provider_models(
if payload.provider_type == "gemini":
try:
from urllib.parse import urlparse as _urlparse
_host = (_urlparse(base_url).hostname or "").lower()
except Exception:
_host = ""
@ -349,9 +332,7 @@ async def list_provider_models(
if allow_prefixes is not None:
prefix_tuple = tuple(str(p) for p in allow_prefixes if str(p))
if prefix_tuple:
models = [
m for m in models if m.get("id", "").startswith(prefix_tuple)
]
models = [m for m in models if m.get("id", "").startswith(prefix_tuple)]
allowlist = info.get("model_id_allowlist")
if allowlist is not None:
models = [m for m in models if allowlist.match(m.get("id", ""))]

7
studio/backend/routes/settings.py
View file

@ -45,16 +45,13 @@ def _upload_limit_response(limit_mb: int) -> UploadLimitResponse:
@router.get("/upload-limit", response_model = UploadLimitResponse)
def get_upload_limit(
current_subject: str = Depends(get_current_subject),
) -> UploadLimitResponse:
def get_upload_limit(current_subject: str = Depends(get_current_subject)) -> UploadLimitResponse:
return _upload_limit_response(get_upload_limit_mb())
@router.put("/upload-limit", response_model = UploadLimitResponse)
def update_upload_limit(
payload: UploadLimitPayload,
current_subject: str = Depends(get_current_subject),
payload: UploadLimitPayload, current_subject: str = Depends(get_current_subject)
) -> UploadLimitResponse:
try:
limit_mb = set_upload_limit_mb(payload.max_upload_size_mb)

152
studio/backend/routes/training.py
View file

@ -71,9 +71,7 @@ router = APIRouter()
logger = get_logger(__name__)
def _validate_local_dataset_paths(
paths: list[str], label: str = "Local dataset"
) -> list[str]:
def _validate_local_dataset_paths(paths: list[str], label: str = "Local dataset") -> list[str]:
"""Resolve and validate a list of local dataset paths. Returns validated absolute paths."""
validated = []
missing = []
@ -95,9 +93,7 @@ def _validate_local_dataset_paths(
@router.get("/hardware")
async def get_hardware_utilization(
current_subject: str = Depends(get_current_subject),
):
async def get_hardware_utilization(current_subject: str = Depends(get_current_subject)):
"""
Get a live snapshot of GPU hardware utilization.
@ -105,23 +101,18 @@ async def get_hardware_utilization(
Returns live GPU memory usage information for the active backend.
"""
from utils.hardware import get_gpu_utilization
return get_gpu_utilization()
@router.get("/hardware/visible")
async def get_visible_hardware_utilization(
current_subject: str = Depends(get_current_subject),
):
async def get_visible_hardware_utilization(current_subject: str = Depends(get_current_subject)):
from utils.hardware import get_visible_gpu_utilization
return get_visible_gpu_utilization()
@router.post("/start")
async def start_training(
request: TrainingStartRequest,
current_subject: str = Depends(get_current_subject),
request: TrainingStartRequest, current_subject: str = Depends(get_current_subject)
):
"""
Start a training job.
@ -153,9 +144,7 @@ async def start_training(
# Generate job ID — passed into start_training() which sets it on the
# backend only after confirming the old pump thread is dead.
job_id = (
f"job_{datetime.now().strftime('%Y%m%d_%H%M%S')}_{_uuid.uuid4().hex[:8]}"
)
job_id = f"job_{datetime.now().strftime('%Y%m%d_%H%M%S')}_{_uuid.uuid4().hex[:8]}"
# Validate dataset paths if provided
if request.local_datasets:
@ -169,9 +158,7 @@ async def start_training(
resume_output_dir: Optional[str] = None
if request.resume_from_checkpoint:
try:
resume_output_dir = normalize_resume_output_dir(
request.resume_from_checkpoint
)
resume_output_dir = normalize_resume_output_dir(request.resume_from_checkpoint)
except ValueError as e:
# Deliberate user-facing validation message.
validation_message = str(e)
@ -229,9 +216,7 @@ async def start_training(
"lora_r": request.lora_r,
"lora_alpha": request.lora_alpha,
"lora_dropout": request.lora_dropout,
"target_modules": request.target_modules
if request.target_modules
else None,
"target_modules": request.target_modules if request.target_modules else None,
"gradient_checkpointing": request.gradient_checkpointing.strip()
if request.gradient_checkpointing and request.gradient_checkpointing.strip()
else "unsloth",
@ -261,20 +246,15 @@ async def start_training(
# net, consult the YAML directly so models that need it always get it.
if not training_kwargs["trust_remote_code"]:
model_defaults = load_model_defaults(request.model_name)
yaml_trust = model_defaults.get("training", {}).get(
"trust_remote_code", False
)
yaml_trust = model_defaults.get("training", {}).get("trust_remote_code", False)
if yaml_trust:
logger.info(
f"YAML config sets trust_remote_code=True for {request.model_name}"
)
logger.info(f"YAML config sets trust_remote_code=True for {request.model_name}")
training_kwargs["trust_remote_code"] = True
# Free GPU memory: shut down any running inference/export subprocesses
# before training starts (they'd compete for VRAM otherwise)
try:
from core.inference import get_inference_backend
inf_backend = get_inference_backend()
if inf_backend.active_model_name:
logger.info(
@ -289,12 +269,9 @@ async def start_training(
try:
from core.export import get_export_backend
exp_backend = get_export_backend()
if exp_backend.current_checkpoint:
logger.info(
"Shutting down export subprocess to free GPU memory for training"
)
logger.info("Shutting down export subprocess to free GPU memory for training")
exp_backend._shutdown_subprocess()
exp_backend.current_checkpoint = None
exp_backend.is_vision = False
@ -376,9 +353,7 @@ async def stop_training(
@router.post("/reset")
async def reset_training(
current_subject: str = Depends(get_current_subject),
):
async def reset_training(current_subject: str = Depends(get_current_subject)):
"""
Reset training state so the user can return to configuration.
"""
@ -389,14 +364,10 @@ async def reset_training(
if is_active:
if backend._cancel_requested:
# Cancel (save=False) was requested — force-terminate so we can reset immediately
logger.info(
"Force-terminating subprocess for immediate reset (cancel path)"
)
logger.info("Force-terminating subprocess for immediate reset (cancel path)")
backend.force_terminate()
else:
logger.warning(
"Rejected reset while training active: is_active=%s", is_active
)
logger.warning("Rejected reset while training active: is_active=%s", is_active)
raise HTTPException(
status_code = 409,
detail = "Training is still running. Stop training and wait for it to finish before resetting.",
@ -433,9 +404,7 @@ async def reset_training(
@router.get("/status")
async def get_training_status(
current_subject: str = Depends(get_current_subject),
):
async def get_training_status(current_subject: str = Depends(get_current_subject)):
"""
Get the current training status.
"""
@ -467,9 +436,7 @@ async def get_training_status(
msg_lower = status_message.lower()
if "loading" in msg_lower or "importing" in msg_lower:
phase = "loading_model"
elif any(
k in msg_lower for k in ["preparing", "initializing", "configuring"]
):
elif any(k in msg_lower for k in ["preparing", "initializing", "configuring"]):
phase = "configuring"
else:
phase = "training"
@ -528,9 +495,7 @@ async def get_training_status(
@router.get("/metrics", response_model = TrainingMetricsResponse)
async def get_training_metrics(
current_subject: str = Depends(get_current_subject),
):
async def get_training_metrics(current_subject: str = Depends(get_current_subject)):
"""
Get training metrics (loss, learning rate, steps).
"""
@ -572,8 +537,7 @@ async def get_training_metrics(
@router.get("/progress")
async def stream_training_progress(
request: Request,
current_subject: str = Depends(get_current_subject),
request: Request, current_subject: str = Depends(get_current_subject)
):
"""
Stream training progress updates using Server-Sent Events (SSE).
@ -614,14 +578,10 @@ async def stream_training_progress(
if step < 0 or total == 0:
progress_percent = 0.0
else:
progress_percent = (
float(step) / float(total) * 100.0 if total > 0 else 0.0
)
progress_percent = float(step) / float(total) * 100.0 if total > 0 else 0.0
# Get actual values from progress object if available
elapsed_seconds = (
getattr(progress, "elapsed_seconds", None) if progress else None
)
elapsed_seconds = getattr(progress, "elapsed_seconds", None) if progress else None
eta_seconds = getattr(progress, "eta_seconds", None) if progress else None
grad_norm = grad_norm_override
if grad_norm is None and progress:
@ -677,25 +637,15 @@ async def stream_training_progress(
}
for i, step_val in enumerate(backend.step_history):
if step_val > resume_from_step:
loss_val = (
backend.loss_history[i]
if i < len(backend.loss_history)
else None
)
lr_val = (
backend.lr_history[i] if i < len(backend.lr_history) else None
)
loss_val = backend.loss_history[i] if i < len(backend.loss_history) else None
lr_val = backend.lr_history[i] if i < len(backend.lr_history) else None
tp_replay = getattr(
getattr(backend, "trainer", None), "training_progress", None
)
total_replay = (
getattr(tp_replay, "total_steps", step_val)
if tp_replay
else step_val
)
epoch_replay = (
getattr(tp_replay, "epoch", None) if tp_replay else None
getattr(tp_replay, "total_steps", step_val) if tp_replay else step_val
)
epoch_replay = getattr(tp_replay, "epoch", None) if tp_replay else None
payload = build_progress(
step_val,
loss_val,
@ -705,9 +655,7 @@ async def stream_training_progress(
progress = tp_replay,
grad_norm_override = grad_norm_by_step.get(step_val),
)
yield format_sse(
payload.model_dump_json(), event = "progress", event_id = step_val
)
yield format_sse(payload.model_dump_json(), event = "progress", event_id = step_val)
replayed += 1
if replayed:
logger.info(f"SSE reconnect: replayed {replayed} missed steps")
@ -727,21 +675,15 @@ async def stream_training_progress(
epoch = initial_epoch,
progress = tp,
)
yield format_sse(
initial_progress.model_dump_json(), event = "progress", event_id = 0
)
yield format_sse(initial_progress.model_dump_json(), event = "progress", event_id = 0)
# If not active, send final state and exit
if not is_active:
if backend.step_history:
final_step = backend.step_history[-1]
final_loss = (
backend.loss_history[-1] if backend.loss_history else None
)
final_loss = backend.loss_history[-1] if backend.loss_history else None
final_lr = backend.lr_history[-1] if backend.lr_history else None
final_total_steps = (
getattr(tp, "total_steps", final_step) if tp else final_step
)
final_total_steps = getattr(tp, "total_steps", final_step) if tp else final_step
final_epoch = getattr(tp, "epoch", None) if tp else None
payload = build_progress(
final_step,
@ -756,9 +698,7 @@ async def stream_training_progress(
)
else:
yield format_sse(
build_progress(
-1, None, None, 0, progress = tp
).model_dump_json(),
build_progress(-1, None, None, 0, progress = tp).model_dump_json(),
event = "complete",
event_id = 0,
)
@ -767,29 +707,19 @@ async def stream_training_progress(
# ── Live polling loop ────────────────────────────────────
last_step = resume_from_step if resume_from_step is not None else -1
no_update_count = 0
max_no_updates = (
1800 # Timeout after 30 minutes (large models need time for compilation)
)
max_no_updates = 1800 # Timeout after 30 minutes (large models need time for compilation)
while backend.is_training_active():
try:
if backend.step_history:
current_step = backend.step_history[-1]
current_loss = (
backend.loss_history[-1] if backend.loss_history else None
)
current_loss = backend.loss_history[-1] if backend.loss_history else None
current_lr = backend.lr_history[-1] if backend.lr_history else None
tp_inner = getattr(
getattr(backend, "trainer", None), "training_progress", None
)
tp_inner = getattr(getattr(backend, "trainer", None), "training_progress", None)
current_total_steps = (
getattr(tp_inner, "total_steps", current_step)
if tp_inner
else current_step
)
current_epoch = (
getattr(tp_inner, "epoch", None) if tp_inner else None
getattr(tp_inner, "total_steps", current_step) if tp_inner else current_step
)
current_epoch = getattr(tp_inner, "epoch", None) if tp_inner else None
# Only send if step changed
if current_step != last_step:
@ -836,9 +766,7 @@ async def stream_training_progress(
"training_progress",
None,
)
prep_total = (
getattr(tp_prep, "total_steps", 0) if tp_prep else 0
)
prep_total = getattr(tp_prep, "total_steps", 0) if tp_prep else 0
preparing_payload = build_progress(
0,
None,
@ -858,9 +786,7 @@ async def stream_training_progress(
tp_timeout = getattr(
getattr(backend, "trainer", None), "training_progress", None
)
timeout_payload = build_progress(
last_step, None, None, 0, progress = tp_timeout
)
timeout_payload = build_progress(last_step, None, None, 0, progress = tp_timeout)
yield format_sse(
timeout_payload.model_dump_json(),
event = "error",
@ -872,9 +798,7 @@ async def stream_training_progress(
except Exception as e:
logger.error(f"Error in progress stream: {e}", exc_info = True)
tp_error = getattr(
getattr(backend, "trainer", None), "training_progress", None
)
tp_error = getattr(getattr(backend, "trainer", None), "training_progress", None)
error_payload = build_progress(0, None, None, 0, progress = tp_error)
yield format_sse(
error_payload.model_dump_json(),
@ -888,9 +812,7 @@ async def stream_training_progress(
final_loss = backend.loss_history[-1] if backend.loss_history else None
final_lr = backend.lr_history[-1] if backend.lr_history else None
final_tp = getattr(getattr(backend, "trainer", None), "training_progress", None)
final_total_steps = (
getattr(final_tp, "total_steps", final_step) if final_tp else final_step
)
final_total_steps = getattr(final_tp, "total_steps", final_step) if final_tp else final_step
final_epoch = getattr(final_tp, "epoch", None) if final_tp else None
final_payload = build_progress(
final_step,

19
studio/backend/routes/training_history.py
View file

@ -42,19 +42,13 @@ async def list_training_runs(
"""List training runs, newest first."""
result = list_runs(limit = limit, offset = offset)
return TrainingRunListResponse(
runs = [
TrainingRunSummary(**{**r, "can_resume": can_resume_run(r)})
for r in result["runs"]
],
runs = [TrainingRunSummary(**{**r, "can_resume": can_resume_run(r)}) for r in result["runs"]],
total = result["total"],
)
@router.get("/runs/{run_id}", response_model = TrainingRunDetailResponse)
async def get_training_run_detail(
run_id: str,
current_subject: str = Depends(get_current_subject),
):
async def get_training_run_detail(run_id: str, current_subject: str = Depends(get_current_subject)):
"""Get a single training run with full config and metrics."""
run = get_run(run_id)
if run is None:
@ -109,18 +103,13 @@ async def update_training_run(
@router.delete("/runs/{run_id}", response_model = TrainingRunDeleteResponse)
async def delete_training_run(
run_id: str,
current_subject: str = Depends(get_current_subject),
):
async def delete_training_run(run_id: str, current_subject: str = Depends(get_current_subject)):
"""Delete a training run and its metrics (CASCADE)."""
run = get_run(run_id)
if run is None:
raise HTTPException(status_code = 404, detail = f"Run {run_id} not found")
if run["status"] == "running":
raise HTTPException(
status_code = 409, detail = "Cannot delete a running training run"
)
raise HTTPException(status_code = 409, detail = "Cannot delete a running training run")
logger.info("Deleting training run %s", run_id)
delete_run(run_id)
return TrainingRunDeleteResponse(

48
studio/backend/run.py
View file

@ -25,9 +25,7 @@ try:
configure_cpu_threads()
except ValueError as exc:
configured = os.environ.get("UNSLOTH_CPU_THREADS")
raise SystemExit(
f"Error: Invalid UNSLOTH_CPU_THREADS value {configured!r}: {exc}"
) from None
raise SystemExit(f"Error: Invalid UNSLOTH_CPU_THREADS value {configured!r}: {exc}") from None
# Fix for Anaconda/conda-forge Python: seed platform._sys_version_cache before
# any library imports that trigger attrs -> rich -> structlog -> platform crash.
@ -93,9 +91,7 @@ def _install_uvicorn_startup_log_rewrite(bind_host: str, display_host: str) -> N
import re
rewrite_host = (
bind_host in ("0.0.0.0", "::")
and bool(display_host)
and display_host != bind_host
bind_host in ("0.0.0.0", "::") and bool(display_host) and display_host != bind_host
)
new_suffix = "(To stop: press Ctrl+C -- on macOS, Control+C not Command+C)"
old_suffix_re = re.compile(r"\(Press CTRL\+C to quit\)")
@ -136,10 +132,13 @@ def _install_uvicorn_startup_log_rewrite(bind_host: str, display_host: str) -> N
logging.getLogger(name).addFilter(f)
def _local_port_open(host: str, port: int, timeout: float = 1.0) -> bool:
def _local_port_open(
host: str,
port: int,
timeout: float = 1.0,
) -> bool:
"""Return True iff a TCP connection to (host, port) succeeds within timeout."""
import socket
try:
with socket.create_connection((host, port), timeout = timeout):
return True
@ -178,9 +177,7 @@ def _localhost_ipv6_mismatch_url(bind_host: str, port: int) -> "str | None":
return None
try:
addr_info = socket.getaddrinfo(
"localhost", port, socket.AF_UNSPEC, socket.SOCK_STREAM
)
addr_info = socket.getaddrinfo("localhost", port, socket.AF_UNSPEC, socket.SOCK_STREAM)
except Exception:
return None
@ -369,8 +366,7 @@ def _verify_global_reachability(display_host: str, port: int) -> None:
flush = True,
)
print(
f"{dim} ssh -L {port}:localhost:{port} "
f"<user>@{display_host}{reset}",
f"{dim} ssh -L {port}:localhost:{port} " f"<user>@{display_host}{reset}",
flush = True,
)
print(
@ -499,15 +495,17 @@ def _is_port_free(host: str, port: int) -> bool:
return True
def _find_free_port(host: str, start: int, max_attempts: int = 20) -> int:
def _find_free_port(
host: str,
start: int,
max_attempts: int = 20,
) -> int:
"""Find a free port starting from `start`, trying up to max_attempts ports."""
for offset in range(max_attempts):
candidate = start + offset
if _is_port_free(host, candidate):
return candidate
raise RuntimeError(
f"Could not find a free port in range {start}-{start + max_attempts - 1}"
)
raise RuntimeError(f"Could not find a free port in range {start}-{start + max_attempts - 1}")
from utils.paths.storage_roots import studio_root as _studio_root
@ -570,7 +568,6 @@ def _graceful_shutdown(server = None):
# 2. Clean up inference subprocess (if instantiated)
try:
from core.inference.orchestrator import _inference_backend
if _inference_backend is not None:
_inference_backend._shutdown_subprocess(timeout = 5.0)
except Exception as e:
@ -579,7 +576,6 @@ def _graceful_shutdown(server = None):
# 3. Clean up export subprocess (if instantiated)
try:
from core.export.orchestrator import _export_backend
if _export_backend is not None:
_export_backend._shutdown_subprocess(timeout = 5.0)
except Exception as e:
@ -588,7 +584,6 @@ def _graceful_shutdown(server = None):
# 4. Clean up training subprocess (if active)
try:
from core.training.training import _training_backend
if _training_backend is not None:
_training_backend.force_terminate()
except Exception as e:
@ -597,7 +592,6 @@ def _graceful_shutdown(server = None):
# 5. Kill llama-server subprocess (if loaded)
try:
from routes.inference import _llama_cpp_backend
if _llama_cpp_backend is not None:
_llama_cpp_backend._kill_process()
except Exception as e:
@ -651,9 +645,7 @@ def _iter_frontend_fallback_candidates() -> "list[Path]":
# Tolerate single- or multi-line dict literals; [^}]* still
# rejects nested dicts, which the setuptools template never
# emits for editable installs.
m = re.search(
r"^MAPPING\s*(?::[^=]*)?=\s*(\{[^}]*\})", src, re.M | re.S
)
m = re.search(r"^MAPPING\s*(?::[^=]*)?=\s*(\{[^}]*\})", src, re.M | re.S)
if not m:
continue
try:
@ -760,9 +752,7 @@ def run_server(
print("=" * 50)
if blocker:
pid, name = blocker
print(
f"Port {original_port} is already in use by " f"{name} (PID {pid})."
)
print(f"Port {original_port} is already in use by " f"{name} (PID {pid}).")
else:
print(f"Port {original_port} is already in use.")
print(f"Unsloth Studio will use port {port} instead.")
@ -991,9 +981,7 @@ if __name__ == "__main__":
sys.stderr.write("=" * 60 + "\n")
traceback.print_exc(file = sys.stderr)
sys.stderr.write("\n")
sys.stderr.write(
"If a package is missing, try re-running: unsloth studio setup\n"
)
sys.stderr.write("If a package is missing, try re-running: unsloth studio setup\n")
sys.stderr.flush()
sys.exit(1)

4
studio/backend/state/tool_policy.py
View file

@ -21,9 +21,7 @@ def get_tool_policy() -> Optional[bool]:
def set_tool_policy(value: Optional[bool]) -> None:
if value is not None and not isinstance(value, bool):
raise TypeError(
f"tool_policy must be Optional[bool], got {type(value).__name__}"
)
raise TypeError(f"tool_policy must be Optional[bool], got {type(value).__name__}")
global _tool_policy
_tool_policy = value

8
studio/backend/storage/mcp_servers_db.py
View file

@ -29,13 +29,9 @@ def _ensure_schema(conn: sqlite3.Connection) -> None:
"""
)
# use_oauth was added after the first release; backfill for pre-existing DBs.
cols = {
r["name"] for r in conn.execute("PRAGMA table_info(mcp_servers)").fetchall()
}
cols = {r["name"] for r in conn.execute("PRAGMA table_info(mcp_servers)").fetchall()}
if "use_oauth" not in cols:
conn.execute(
"ALTER TABLE mcp_servers ADD COLUMN use_oauth INTEGER NOT NULL DEFAULT 0"
)
conn.execute("ALTER TABLE mcp_servers ADD COLUMN use_oauth INTEGER NOT NULL DEFAULT 0")
def get_connection() -> sqlite3.Connection:

11
studio/backend/storage/providers_db.py
View file

@ -62,12 +62,7 @@ def get_connection() -> sqlite3.Connection:
return conn
def create_provider(
id: str,
provider_type: str,
display_name: str,
base_url: str,
) -> None:
def create_provider(id: str, provider_type: str, display_name: str, base_url: str) -> None:
"""Insert a new provider configuration."""
now = datetime.now(timezone.utc).isoformat()
conn = get_connection()
@ -145,9 +140,7 @@ def list_providers() -> list[dict]:
"""List all provider configurations, ordered by creation time."""
conn = get_connection()
try:
rows = conn.execute(
"SELECT * FROM llm_providers ORDER BY created_at"
).fetchall()
rows = conn.execute("SELECT * FROM llm_providers ORDER BY created_at").fetchall()
return [dict(row) for row in rows]
finally:
conn.close()

72
studio/backend/storage/studio_db.py
View file

@ -89,9 +89,7 @@ def _delete_project_workspace(project: dict) -> None:
try:
root_resolved = root.resolve(strict = False)
except (OSError, RuntimeError, ValueError):
logger.warning(
"Skipping project workspace delete for invalid path %r", root_path
)
logger.warning("Skipping project workspace delete for invalid path %r", root_path)
return
project_id = str(project["id"])
@ -155,9 +153,7 @@ def _ensure_schema(conn: sqlite3.Connection) -> None:
)
"""
)
existing_cols = {
row[1] for row in conn.execute("PRAGMA table_info(training_runs)").fetchall()
}
existing_cols = {row[1] for row in conn.execute("PRAGMA table_info(training_runs)").fetchall()}
if "display_name" not in existing_cols:
conn.execute("ALTER TABLE training_runs ADD COLUMN display_name TEXT")
conn.execute(
@ -177,9 +173,7 @@ def _ensure_schema(conn: sqlite3.Connection) -> None:
)
"""
)
conn.execute(
"CREATE INDEX IF NOT EXISTS idx_metrics_run_id ON training_metrics(run_id)"
)
conn.execute("CREATE INDEX IF NOT EXISTS idx_metrics_run_id ON training_metrics(run_id)")
# Use COLLATE NOCASE on Windows so C:\Models and c:\models dedup via the
# UNIQUE constraint. On Linux/macOS (case-sensitive FS) keep the default
# BINARY collation so /Models and /models remain distinct.
@ -237,13 +231,9 @@ def _ensure_schema(conn: sqlite3.Connection) -> None:
if "project_id" not in chat_thread_cols:
conn.execute("ALTER TABLE chat_threads ADD COLUMN project_id TEXT")
if "openai_code_exec_container_id" not in chat_thread_cols:
conn.execute(
"ALTER TABLE chat_threads ADD COLUMN openai_code_exec_container_id TEXT"
)
conn.execute("ALTER TABLE chat_threads ADD COLUMN openai_code_exec_container_id TEXT")
if "anthropic_code_exec_container_id" not in chat_thread_cols:
conn.execute(
"ALTER TABLE chat_threads ADD COLUMN anthropic_code_exec_container_id TEXT"
)
conn.execute("ALTER TABLE chat_threads ADD COLUMN anthropic_code_exec_container_id TEXT")
conn.execute(
"""
CREATE TABLE IF NOT EXISTS chat_messages (
@ -261,9 +251,7 @@ def _ensure_schema(conn: sqlite3.Connection) -> None:
conn.execute(
"CREATE INDEX IF NOT EXISTS idx_chat_threads_model_type_created_at ON chat_threads(model_type, created_at)"
)
conn.execute(
"CREATE INDEX IF NOT EXISTS idx_chat_threads_pair_id ON chat_threads(pair_id)"
)
conn.execute("CREATE INDEX IF NOT EXISTS idx_chat_threads_pair_id ON chat_threads(pair_id)")
conn.execute(
"CREATE INDEX IF NOT EXISTS idx_chat_threads_project_id ON chat_threads(project_id)"
)
@ -513,9 +501,7 @@ def list_runs(limit: int = 50, offset: int = 0) -> dict:
try:
run["loss_sparkline"] = json.loads(sparkline)
except (json.JSONDecodeError, TypeError):
logger.debug(
"Failed to parse loss_sparkline for run %s", run.get("id")
)
logger.debug("Failed to parse loss_sparkline for run %s", run.get("id"))
run["loss_sparkline"] = None
runs.append(run)
return {"runs": runs, "total": total}
@ -591,9 +577,7 @@ def get_resumable_run_by_output_dir(output_dir: str) -> Optional[dict]:
try:
run["loss_sparkline"] = json.loads(sparkline)
except (json.JSONDecodeError, TypeError):
logger.debug(
"Failed to parse loss_sparkline for output_dir %s", output_dir
)
logger.debug("Failed to parse loss_sparkline for output_dir %s", output_dir)
run["loss_sparkline"] = None
return run
finally:
@ -1130,9 +1114,7 @@ def _parse_chat_setting_json(key: str, value_json: str) -> tuple[bool, Any]:
return False, None
def _load_chat_settings_for_merge(
conn: sqlite3.Connection,
) -> tuple[dict[str, Any], set[str]]:
def _load_chat_settings_for_merge(conn: sqlite3.Connection) -> tuple[dict[str, Any], set[str]]:
rows = conn.execute("SELECT key, value_json FROM chat_settings").fetchall()
current: dict[str, Any] = {}
corrupt: set[str] = set()
@ -1159,9 +1141,7 @@ def _load_chat_settings_for_merge(
def _raise_if_chat_message_thread_conflicts(
conn: sqlite3.Connection,
thread_id: str,
message_ids: list[str],
conn: sqlite3.Connection, thread_id: str, message_ids: list[str]
) -> None:
unique_ids = list(dict.fromkeys(message_ids))
if not unique_ids:
@ -1270,12 +1250,8 @@ def sync_chat_messages(
m.get("parentId"),
m["role"],
json.dumps(m.get("content", [])),
json.dumps(m.get("attachments"))
if m.get("attachments") is not None
else None,
json.dumps(m.get("metadata"))
if m.get("metadata") is not None
else None,
json.dumps(m.get("attachments")) if m.get("attachments") is not None else None,
json.dumps(m.get("metadata")) if m.get("metadata") is not None else None,
int(m["createdAt"]),
)
for m in messages
@ -1355,9 +1331,7 @@ def list_chat_messages_for_threads(thread_ids: list[str]) -> list[dict]:
def get_app_setting(key: str, fallback = None):
conn = get_connection()
try:
row = conn.execute(
"SELECT value_json FROM app_settings WHERE key = ?", (key,)
).fetchone()
row = conn.execute("SELECT value_json FROM app_settings WHERE key = ?", (key,)).fetchone()
if row is None:
return fallback
return _json_loads(row["value_json"], fallback)
@ -1382,9 +1356,7 @@ def upsert_app_settings(settings: dict[str, Any]) -> dict[str, Any]:
[(key, json.dumps(value), now) for key, value in settings.items()],
)
conn.commit()
rows = conn.execute(
"SELECT key, value_json FROM app_settings ORDER BY key"
).fetchall()
rows = conn.execute("SELECT key, value_json FROM app_settings ORDER BY key").fetchall()
return {row["key"]: _json_loads(row["value_json"], None) for row in rows}
finally:
conn.close()
@ -1393,9 +1365,7 @@ def upsert_app_settings(settings: dict[str, Any]) -> dict[str, Any]:
def list_chat_settings() -> dict[str, Any]:
conn = get_connection()
try:
rows = conn.execute(
"SELECT key, value_json FROM chat_settings ORDER BY key"
).fetchall()
rows = conn.execute("SELECT key, value_json FROM chat_settings ORDER BY key").fetchall()
settings: dict[str, Any] = {}
for row in rows:
settings[row["key"]] = _json_loads(row["value_json"], None)
@ -1426,9 +1396,7 @@ def upsert_chat_settings(settings: dict[str, Any]) -> dict[str, Any]:
conn.close()
def _deep_merge_settings(
current: dict[str, Any], updates: dict[str, Any]
) -> dict[str, Any]:
def _deep_merge_settings(current: dict[str, Any], updates: dict[str, Any]) -> dict[str, Any]:
merged = dict(current)
for key, value in updates.items():
current_value = merged.get(key)
@ -1449,9 +1417,7 @@ def upsert_chat_settings_merge(updates: dict[str, Any]) -> dict[str, Any]:
conn.execute("BEGIN IMMEDIATE")
current, corrupt = _load_chat_settings_for_merge(conn)
unsafe_partial_keys = [
key
for key, value in updates.items()
if key in corrupt and isinstance(value, dict)
key for key, value in updates.items() if key in corrupt and isinstance(value, dict)
]
if unsafe_partial_keys:
conn.commit()
@ -1496,9 +1462,7 @@ def list_chat_legacy_imports() -> list[str]:
"""
conn = get_connection()
try:
rows = conn.execute(
"SELECT legacy_thread_id FROM chat_legacy_imports"
).fetchall()
rows = conn.execute("SELECT legacy_thread_id FROM chat_legacy_imports").fetchall()
return [row[0] for row in rows]
finally:
conn.close()

7
studio/backend/tests/test_amd_apu_unified_memory.py
View file

@ -14,7 +14,12 @@ import pytest
from core.inference.llama_cpp import LlamaCppBackend
def _fake_torch(hip, archs, *, cuda_ok = True):
def _fake_torch(
hip,
archs,
*,
cuda_ok = True,
):
t = types.ModuleType("torch")
t.version = types.SimpleNamespace(hip = hip)
t.cuda = types.SimpleNamespace(

8
studio/backend/tests/test_anthropic_citations_edge.py
View file

@ -80,8 +80,7 @@ def _capture(
client = _make_client()
try:
async for line in client.stream_chat_completion(
messages = messages
or [{"role": "user", "content": "what color is grass?"}],
messages = messages or [{"role": "user", "content": "what color is grass?"}],
model = "claude-opus-4-7",
max_tokens = 64,
):
@ -159,10 +158,7 @@ def _citation_payload(body: str) -> dict:
except json.JSONDecodeError:
continue
tool_event = payload.get("_toolEvent") if isinstance(payload, dict) else None
if (
isinstance(tool_event, dict)
and tool_event.get("type") == "document_citations"
):
if isinstance(tool_event, dict) and tool_event.get("type") == "document_citations":
return tool_event
raise AssertionError("document_citations event not parsed out of SSE body")

18
studio/backend/tests/test_anthropic_code_execution.py
View file

@ -118,10 +118,7 @@ def test_code_execution_tool_appended_to_request_body(monkeypatch):
tools = body.get("tools") or []
# Opus 4.7 gets the newer date-pinned variant (`_20260120`) that
# supports REPL state persistence + programmatic tool calling.
assert {
"type": "code_execution_20260120",
"name": "code_execution",
} in tools
assert {"type": "code_execution_20260120", "name": "code_execution"} in tools
# No web_search entry when only code_execution is enabled.
assert all("web_search" not in (t.get("type") or "") for t in tools)
# Beta header still carries the documented flag; both `_20250825`
@ -201,9 +198,7 @@ def test_no_code_execution_tool_when_pill_off(monkeypatch):
assert all("code_execution" not in (t.get("type") or "") for t in tools)
# Beta header must NOT mention code-execution when the tool isn't on
# -- that flag is opt-in only.
assert "code-execution-2025-08-25" not in captured["headers"].get(
"anthropic-beta", ""
)
assert "code-execution-2025-08-25" not in captured["headers"].get("anthropic-beta", "")
def test_bash_code_execution_emits_tool_start_and_end(monkeypatch):
@ -277,11 +272,7 @@ def test_bash_code_execution_emits_tool_start_and_end(monkeypatch):
assert start["tool_call_id"] == "srvtoolu_1"
# `_server_tool: True` marks this as a provider-side synthetic
# tool card for the frontend's history serializer.
assert start["arguments"] == {
"kind": "bash",
"command": "ls -la",
"_server_tool": True,
}
assert start["arguments"] == {"kind": "bash", "command": "ls -la", "_server_tool": True}
assert end["type"] == "tool_end"
assert end["tool_call_id"] == "srvtoolu_1"
@ -308,8 +299,7 @@ def test_text_editor_create_emits_kind_and_status(monkeypatch):
"delta": {
"type": "input_json_delta",
"partial_json": (
'{"command": "create", "path": "new_file.txt", '
'"file_text": "hi"}'
'{"command": "create", "path": "new_file.txt", "file_text": "hi"}'
),
},
},

45
studio/backend/tests/test_anthropic_compaction.py
View file

@ -40,7 +40,12 @@ def _make_client() -> ExternalProviderClient:
)
def _capture(monkeypatch, model: str, threshold, tools = None) -> dict:
def _capture(
monkeypatch,
model: str,
threshold,
tools = None,
) -> dict:
captured: dict = {}
def handler(request: httpx.Request) -> httpx.Response:
@ -120,13 +125,9 @@ def test_supported_model_attaches_compaction_block_and_beta(monkeypatch):
def test_threshold_clamped_to_50k_minimum(monkeypatch):
# Below-min values get clamped UP so we don't 400 upstream.
captured = _capture(monkeypatch, "claude-opus-4-7", 60_000)
assert (
captured["body"]["context_management"]["edits"][0]["trigger"]["value"] == 60_000
)
assert captured["body"]["context_management"]["edits"][0]["trigger"]["value"] == 60_000
captured = _capture(monkeypatch, "claude-opus-4-7", 1)
assert (
captured["body"]["context_management"]["edits"][0]["trigger"]["value"] == 50_000
)
assert captured["body"]["context_management"]["edits"][0]["trigger"]["value"] == 50_000
# ── beta header merge with code execution ────────────────────────────
@ -151,10 +152,7 @@ def test_unsupported_model_silently_drops_compaction(monkeypatch):
captured = _capture(monkeypatch, "claude-haiku-4-5-20251001", 150_000)
assert "context_management" not in captured["body"]
# The beta header must not carry compact-2026-01-12 either.
assert "compact-2026-01-12" not in captured["headers"].get(
"anthropic-beta",
"",
)
assert "compact-2026-01-12" not in captured["headers"].get("anthropic-beta", "")
# ── omitted threshold leaves body untouched ─────────────────────────
@ -163,10 +161,7 @@ def test_unsupported_model_silently_drops_compaction(monkeypatch):
def test_omitted_threshold_no_body_field(monkeypatch):
captured = _capture(monkeypatch, "claude-opus-4-7", None)
assert "context_management" not in captured["body"]
assert "compact-2026-01-12" not in captured["headers"].get(
"anthropic-beta",
"",
)
assert "compact-2026-01-12" not in captured["headers"].get("anthropic-beta", "")
# ── ChatCompletionRequest schema accepts sub-50k threshold ──────────
@ -212,9 +207,7 @@ def test_chat_completion_request_accepts_sub_50k_compaction_threshold():
# ── usage.iterations[] surfaces compaction tokens ──────────────────
def test_message_delta_iterations_array_aggregates_compaction_tokens(
monkeypatch, capsys
):
def test_message_delta_iterations_array_aggregates_compaction_tokens(monkeypatch, capsys):
# When Anthropic compacts mid-stream, the SSE message_delta usage
# payload carries `iterations: [{type:"compaction", ...}, ...]`.
# The top-level input_tokens / output_tokens only account for the
@ -552,9 +545,7 @@ def test_build_external_messages_passes_compaction_for_anthropic_only():
}
)
]
out = _build_external_messages(
msgs, supports_vision = True, provider_type = "anthropic"
)
out = _build_external_messages(msgs, supports_vision = True, provider_type = "anthropic")
assert len(out) == 1
parts = out[0]["content"]
assert parts[0] == {"type": "compaction", "content": "prior summary"}
@ -582,9 +573,7 @@ def test_build_external_messages_strips_compaction_for_non_anthropic_providers()
)
]
for provider in ("openai", "deepseek", "mistral", "gemini", "kimi", "openrouter"):
out = _build_external_messages(
msgs, supports_vision = True, provider_type = provider
)
out = _build_external_messages(msgs, supports_vision = True, provider_type = provider)
assert len(out) == 1, (provider, out)
parts = out[0]["content"]
types = [p.get("type") for p in parts if isinstance(p, dict)]
@ -635,14 +624,10 @@ def test_build_external_messages_non_vision_anthropic_keeps_compaction():
}
)
]
out = _build_external_messages(
msgs, supports_vision = False, provider_type = "anthropic"
)
out = _build_external_messages(msgs, supports_vision = False, provider_type = "anthropic")
parts = out[0]["content"]
assert {"type": "compaction", "content": "prior summary"} in parts
# Non-anthropic + non-vision -> compaction stripped, text collapsed
# back to a string.
out2 = _build_external_messages(
msgs, supports_vision = False, provider_type = "deepseek"
)
out2 = _build_external_messages(msgs, supports_vision = False, provider_type = "deepseek")
assert out2[0]["content"] == "answer", out2

6
studio/backend/tests/test_anthropic_fast_mode_and_refusal.py
View file

@ -58,7 +58,11 @@ def _refusal_sse() -> bytes:
)
def _capture(monkeypatch, sse: bytes = b"", **kwargs) -> tuple[dict, list[str]]:
def _capture(
monkeypatch,
sse: bytes = b"",
**kwargs,
) -> tuple[dict, list[str]]:
"""Install a MockTransport, drive one streamed call, return body+lines."""
captured: dict = {}

19
studio/backend/tests/test_anthropic_fast_mode_edge.py
View file

@ -60,7 +60,11 @@ def _refusal_sse(model: str = "claude-opus-4-7") -> bytes:
)
def _capture(monkeypatch, sse: bytes = b"", **kwargs) -> tuple[dict, list[str]]:
def _capture(
monkeypatch,
sse: bytes = b"",
**kwargs,
) -> tuple[dict, list[str]]:
"""Install a MockTransport, drive one streamed call, return body+lines."""
captured: dict = {}
@ -266,9 +270,7 @@ def test_refusal_notice_appears_before_content_filter_chunk(monkeypatch):
"""The notice content delta must precede the finish_reason chunk."""
_, lines = _capture(monkeypatch, sse = _refusal_sse(), model = "claude-opus-4-7")
notice_idx = next(i for i, l in enumerate(lines) if "stopped by Anthropic" in l)
filter_idx = next(
i for i, l in enumerate(lines) if '"finish_reason": "content_filter"' in l
)
filter_idx = next(i for i, l in enumerate(lines) if '"finish_reason": "content_filter"' in l)
assert notice_idx < filter_idx, (notice_idx, filter_idx, lines)
@ -350,7 +352,6 @@ def test_fast_mode_prefix_tuple_matches_capability_doc(monkeypatch):
"""Tuple must exactly match the two families in the upstream docs:
https://platform.claude.com/docs/en/build-with-claude/fast-mode."""
from core.inference.external_provider import _ANTHROPIC_FAST_MODE_PREFIXES
assert set(_ANTHROPIC_FAST_MODE_PREFIXES) == {
"claude-opus-4-7",
"claude-opus-4-6",
@ -421,9 +422,7 @@ def test_usage_speed_propagates_to_final_usage_chunk_fast(monkeypatch):
def test_usage_speed_propagates_to_final_usage_chunk_standard(monkeypatch):
_, lines = _capture(monkeypatch, sse = _fast_speed_sse(speed = "standard"))
parsed = [
json.loads(l[len("data: ") :])
for l in lines
if l.startswith("data: ") and '"usage"' in l
json.loads(l[len("data: ") :]) for l in lines if l.startswith("data: ") and '"usage"' in l
]
speeds = [p["usage"].get("speed") for p in parsed if "usage" in p]
assert "standard" in speeds, parsed
@ -433,9 +432,7 @@ def test_usage_speed_absent_when_anthropic_does_not_report(monkeypatch):
"""Studio must not invent ``usage.speed`` when upstream omits it."""
_, lines = _capture(monkeypatch)
parsed = [
json.loads(l[len("data: ") :])
for l in lines
if l.startswith("data: ") and '"usage"' in l
json.loads(l[len("data: ") :]) for l in lines if l.startswith("data: ") and '"usage"' in l
]
for p in parsed:
usage = p.get("usage") or {}

61
studio/backend/tests/test_anthropic_messages.py
View file

@ -372,10 +372,7 @@ class TestAnthropicMessagesToOpenAI:
]
result = anthropic_messages_to_openai(msgs)
parts = result[0]["content"]
assert parts[1] == {
"type": "image_url",
"image_url": {"url": "https://x/y.png"},
}
assert parts[1] == {"type": "image_url", "image_url": {"url": "https://x/y.png"}}
def test_image_only_user_message_emits_no_text_part(self):
msgs = [
@ -440,12 +437,7 @@ class TestAnthropicMessagesToOpenAI:
]
result = anthropic_messages_to_openai(msgs)
parts = result[0]["content"]
assert [p["type"] for p in parts] == [
"text",
"image_url",
"text",
"image_url",
]
assert [p["type"] for p in parts] == ["text", "image_url", "text", "image_url"]
assert parts[0]["text"] == "before"
assert parts[2]["text"] == "after"
assert parts[1]["image_url"]["url"] == "data:image/png;base64,AA"
@ -523,15 +515,10 @@ class TestAnthropicToolsToOpenAI:
enabled_tools = ["python"],
)
assert [tool["function"]["name"] for tool in result] == [
"web_search",
"python",
]
assert [tool["function"]["name"] for tool in result] == ["web_search", "python"]
def test_pydantic_model_input(self):
tool = AnthropicTool(
name = "test", description = "desc", input_schema = {"type": "object"}
)
tool = AnthropicTool(name = "test", description = "desc", input_schema = {"type": "object"})
result = anthropic_tools_to_openai([tool])
assert result[0]["function"]["name"] == "test"
@ -631,12 +618,8 @@ class TestAnthropicStreamEmitter:
}
)
first_payloads = [
json.loads(event.split("data: ")[1]) for event in first_events
]
second_payloads = [
json.loads(event.split("data: ")[1]) for event in second_events
]
first_payloads = [json.loads(event.split("data: ")[1]) for event in first_events]
second_payloads = [json.loads(event.split("data: ")[1]) for event in second_events]
tool_starts = [
payload
@ -652,9 +635,7 @@ class TestAnthropicStreamEmitter:
"index": tool_starts[0]["index"],
"delta": {
"type": "input_json_delta",
"partial_json": json.dumps(
{"code": "<!doctype html><html></html>"}
),
"partial_json": json.dumps({"code": "<!doctype html><html></html>"}),
},
}
]
@ -811,9 +792,7 @@ class TestAnthropicToolNonStreaming:
response = asyncio.run(_anthropic_tool_non_streaming(_run_gen, "msg_1", "m"))
body = json.loads(response.body)
tool_blocks = [
block for block in body["content"] if block["type"] == "tool_use"
]
tool_blocks = [block for block in body["content"] if block["type"] == "tool_use"]
assert tool_blocks == [
{
@ -919,26 +898,14 @@ class TestAnthropicPassthroughEmitter:
events1 = e.feed_chunk(
{
"choices": [
{
"delta": {
"tool_calls": [
{"index": 0, "function": {"arguments": '{"cmd'}}
]
}
}
{"delta": {"tool_calls": [{"index": 0, "function": {"arguments": '{"cmd'}}]}}
]
}
)
events2 = e.feed_chunk(
{
"choices": [
{
"delta": {
"tool_calls": [
{"index": 0, "function": {"arguments": '": "ls"}'}}
]
}
}
{"delta": {"tool_calls": [{"index": 0, "function": {"arguments": '": "ls"}'}}]}}
]
}
)
@ -1331,9 +1298,7 @@ class TestAnthropicMessagesToolRouting:
assert exc.value.status_code == 400
assert "Mixing Anthropic server tools" in exc.value.detail
def test_mixed_rejected_when_client_tool_name_collides_with_server_alias(
self, monkeypatch
):
def test_mixed_rejected_when_client_tool_name_collides_with_server_alias(self, monkeypatch):
# Regression: a client tool sharing a name with a mapped server
# tool (e.g. user defines their own "web_search") must still
# trigger the mixed-mode 400 — the post-name filter would
@ -1392,9 +1357,7 @@ class TestAnthropicMessagesToolRouting:
assert exc.value.status_code == 400
assert "name" in exc.value.detail
def test_alias_named_client_tool_without_schema_rejected_with_400(
self, monkeypatch
):
def test_alias_named_client_tool_without_schema_rejected_with_400(self, monkeypatch):
# Regression: a typo'd client tool whose name happens to collide
# with a Studio alias (e.g. user meant a custom "python" tool but
# forgot input_schema) must surface a 400, not silently switch

5
studio/backend/tests/test_anthropic_tool_versions.py
View file

@ -168,10 +168,7 @@ def test_outbound_body_uses_new_versions_on_opus_4_7(monkeypatch):
# Beta header for code execution stays on the existing flag for
# both _20250825 and _20260120; the API uses one header to gate
# the feature, not the date.
assert "code-execution-2025-08-25" in captured["headers"].get(
"anthropic-beta",
"",
)
assert "code-execution-2025-08-25" in captured["headers"].get("anthropic-beta", "")
def test_outbound_body_falls_back_on_haiku_4_5(monkeypatch):

19
studio/backend/tests/test_anthropic_web_fetch.py
View file

@ -104,11 +104,7 @@ def test_web_fetch_tool_appended_to_request_body(monkeypatch):
body = captured["body"]
tools = body.get("tools") or []
# claude-opus-4-7 routes web_fetch to _20260209 (dynamic filtering).
assert {
"type": "web_fetch_20260209",
"name": "web_fetch",
"max_uses": 5,
} in tools
assert {"type": "web_fetch_20260209", "name": "web_fetch", "max_uses": 5} in tools
# web_fetch is GA; no beta header is required.
assert "web-fetch" not in captured["headers"].get("anthropic-beta", "")
@ -183,9 +179,7 @@ def test_no_web_fetch_tool_when_pill_off(monkeypatch):
_drive(run())
tools = captured["body"].get("tools") or []
assert all(
t.get("type") not in ("web_fetch_20250910", "web_fetch_20260209") for t in tools
)
assert all(t.get("type") not in ("web_fetch_20250910", "web_fetch_20260209") for t in tools)
# ── SSE translation ─────────────────────────────────────────────────
@ -253,9 +247,7 @@ def test_web_fetch_success_emits_tool_start_and_end(monkeypatch):
client = _make_client()
return await _collect(
client._stream_anthropic(
messages = [
{"role": "user", "content": "Fetch https://example.com/article"}
],
messages = [{"role": "user", "content": "Fetch https://example.com/article"}],
model = "claude-opus-4-7",
temperature = 0.7,
top_p = 0.95,
@ -273,10 +265,7 @@ def test_web_fetch_success_emits_tool_start_and_end(monkeypatch):
assert start["tool_call_id"] == "srvtoolu_wf1"
# `_server_tool: True` marks this as a provider-side synthetic
# tool card for the frontend's history serializer.
assert start["arguments"] == {
"url": "https://example.com/article",
"_server_tool": True,
}
assert start["arguments"] == {"url": "https://example.com/article", "_server_tool": True}
assert end["type"] == "tool_end"
assert end["tool_call_id"] == "srvtoolu_wf1"
# The source pill uses Title / URL / snippet as parseSourcesFromResult expects.

4
studio/backend/tests/test_audio_token_detection.py
View file

@ -19,9 +19,7 @@ def _classify(tokens: list[str]) -> str | None:
def test_gemma3n_audio_soft_token_is_audio_vlm():
assert (
_classify(["<bos>", "<audio_soft_token>", "<image_soft_token>"]) == "audio_vlm"
)
assert _classify(["<bos>", "<audio_soft_token>", "<image_soft_token>"]) == "audio_vlm"
def test_gemma4_pipe_audio_token_is_audio_vlm():

31
studio/backend/tests/test_cached_gguf_routes.py
View file

@ -66,9 +66,7 @@ def test_iter_gguf_paths_matches_extension_case_insensitively(tmp_path):
assert result == ["Q4_K_M.gguf", "Q8_0.GGUF"]
def test_list_cached_gguf_includes_non_suffix_repo_when_cache_contains_gguf(
monkeypatch, tmp_path
):
def test_list_cached_gguf_includes_non_suffix_repo_when_cache_contains_gguf(monkeypatch, tmp_path):
repo = _repo(
"HauhauCS/Gemma-4-E4B-Uncensored-HauhauCS-Aggressive",
[_file("Q4_K_M.gguf", 5_000), _file("README.md", 10)],
@ -130,9 +128,7 @@ def test_list_cached_gguf_skips_repos_without_positive_gguf_size(monkeypatch, tm
assert result["cached"] == []
def test_list_cached_gguf_keeps_largest_duplicate_repo_across_scans(
monkeypatch, tmp_path
):
def test_list_cached_gguf_keeps_largest_duplicate_repo_across_scans(monkeypatch, tmp_path):
smaller = _repo(
"Org/Dupe",
[_file("Q4_K_M.gguf", 2_000)],
@ -193,9 +189,7 @@ def test_list_cached_gguf_dedupes_shared_blobs_across_revisions(monkeypatch, tmp
]
def test_list_cached_models_skips_non_suffix_repo_when_gguf_files_exist(
monkeypatch, tmp_path
):
def test_list_cached_models_skips_non_suffix_repo_when_gguf_files_exist(monkeypatch, tmp_path):
mixed = _repo(
"Org/MixedRepo",
[
@ -216,9 +210,7 @@ def test_list_cached_models_skips_non_suffix_repo_when_gguf_files_exist(
assert result["cached"] == []
def test_list_cached_gguf_includes_mixed_repo_with_gguf_and_safetensors(
monkeypatch, tmp_path
):
def test_list_cached_gguf_includes_mixed_repo_with_gguf_and_safetensors(monkeypatch, tmp_path):
"""Mirror of the _skips_ test: the mixed repo should still surface in
cached-gguf so the picker can show it as a GGUF download."""
mixed = _repo(
@ -274,9 +266,7 @@ def test_list_cached_gguf_handles_none_size_on_disk(monkeypatch, tmp_path):
]
def test_list_cached_gguf_skips_malformed_repo_without_wiping_response(
monkeypatch, tmp_path
):
def test_list_cached_gguf_skips_malformed_repo_without_wiping_response(monkeypatch, tmp_path):
"""One repo raising during classification must not poison the response
for every other repo in the scan."""
@ -357,17 +347,10 @@ def test_list_cached_models_includes_repo_with_only_mmproj_gguf(monkeypatch, tmp
result = asyncio.run(models_route.list_cached_models(current_subject = "test-user"))
assert result["cached"] == [
{
"repo_id": "Org/MmprojAux",
"size_bytes": 15_000,
}
]
assert result["cached"] == [{"repo_id": "Org/MmprojAux", "size_bytes": 15_000}]
def test_list_cached_gguf_includes_vision_repo_with_main_gguf_and_mmproj(
monkeypatch, tmp_path
):
def test_list_cached_gguf_includes_vision_repo_with_main_gguf_and_mmproj(monkeypatch, tmp_path):
"""A vision-capable GGUF repo (main weight + mmproj adapter) is still
a GGUF repo. The reported size is the main weight size; mmproj is
excluded from the GGUF-size accounting because it is filtered out at

8
studio/backend/tests/test_chat_history_routes.py
View file

@ -109,16 +109,13 @@ def test_chat_inference_settings_covers_frontend_persisted_fields():
pytest.skip("frontend runtime.ts not present")
with open(runtime_ts, encoding = "utf-8") as fh:
block = re.search(
r"interface InferenceParams \{(.*?)\n\}", fh.read(), re.DOTALL
)
block = re.search(r"interface InferenceParams \{(.*?)\n\}", fh.read(), re.DOTALL)
assert block, "InferenceParams interface not found in runtime.ts"
persisted = set(re.findall(r"^\s*(\w+)\??:", block.group(1), re.M)) - {"checkpoint"}
backend = set(chat_history.ChatInferenceSettings.model_fields)
assert persisted == backend, (
f"schema drift: frontend-only {persisted - backend}, "
f"backend-only {backend - persisted}"
f"schema drift: frontend-only {persisted - backend}, " f"backend-only {backend - persisted}"
)
@ -168,7 +165,6 @@ def test_record_import_ledger_returns_accepted_and_inserted(monkeypatch):
def test_record_import_ledger_rejects_oversize_payload():
from pydantic import ValidationError
with pytest.raises(ValidationError):
chat_history.ChatImportLedgerRecordRequest(
threadIds = [f"id-{i}" for i in range(10_001)],

40
studio/backend/tests/test_chat_history_storage.py
View file

@ -13,7 +13,11 @@ import pytest
from storage import studio_db
def _reset_studio_db(tmp_path, monkeypatch, projects_home = None):
def _reset_studio_db(
tmp_path,
monkeypatch,
projects_home = None,
):
monkeypatch.setenv("UNSLOTH_STUDIO_HOME", str(tmp_path))
monkeypatch.setenv(
"UNSLOTH_STUDIO_PROJECTS_HOME",
@ -105,10 +109,7 @@ def test_sync_chat_messages_upserts_without_pruning(tmp_path, monkeypatch):
assert by_id["msg-2"]["content"] == [{"type": "text", "text": "updated text"}]
def test_chat_projects_delete_cascades_threads_and_messages(
tmp_path,
monkeypatch,
):
def test_chat_projects_delete_cascades_threads_and_messages(tmp_path, monkeypatch):
_reset_studio_db(tmp_path, monkeypatch)
project = studio_db.upsert_chat_project(_project())
assert project["rootPath"].startswith(str(tmp_path / "Projects"))
@ -231,23 +232,16 @@ def test_settings_merge_atomic_under_concurrency(tmp_path, monkeypatch):
def test_settings_merge_preserves_nested_keys(tmp_path, monkeypatch):
_reset_studio_db(tmp_path, monkeypatch)
studio_db.upsert_chat_settings_merge(
{"inferenceParams": {"temperature": 0.5, "topP": 0.8}}
)
studio_db.upsert_chat_settings_merge({"inferenceParams": {"temperature": 0.5, "topP": 0.8}})
studio_db.upsert_chat_settings_merge({"inferenceParams": {"temperature": 0.9}})
params = studio_db.list_chat_settings()["inferenceParams"]
assert params == {"temperature": 0.9, "topP": 0.8}
def test_settings_merge_quarantines_corrupt_json_and_rejects_partial_patch(
tmp_path,
monkeypatch,
):
def test_settings_merge_quarantines_corrupt_json_and_rejects_partial_patch(tmp_path, monkeypatch):
_reset_studio_db(tmp_path, monkeypatch)
studio_db.upsert_chat_settings_merge(
{"inferenceParams": {"temperature": 0.5, "topP": 0.8}}
)
studio_db.upsert_chat_settings_merge({"inferenceParams": {"temperature": 0.5, "topP": 0.8}})
conn = studio_db.get_connection()
try:
conn.execute(
@ -295,9 +289,7 @@ def test_settings_merge_replaces_corrupt_scalar_after_quarantine(tmp_path, monke
assert settings["autoTitle"] is True
conn = studio_db.get_connection()
try:
quarantined = conn.execute(
"SELECT key, reason FROM chat_settings_quarantine"
).fetchall()
quarantined = conn.execute("SELECT key, reason FROM chat_settings_quarantine").fetchall()
finally:
conn.close()
assert [(row["key"], row["reason"]) for row in quarantined] == [
@ -353,11 +345,7 @@ def test_legacy_imports_records_and_lists(tmp_path, monkeypatch):
)
assert accepted == 3
assert inserted == 3
assert set(studio_db.list_chat_legacy_imports()) == {
"legacy-a",
"legacy-b",
"legacy-c",
}
assert set(studio_db.list_chat_legacy_imports()) == {"legacy-a", "legacy-b", "legacy-c"}
def test_legacy_imports_is_idempotent(tmp_path, monkeypatch):
@ -371,11 +359,7 @@ def test_legacy_imports_is_idempotent(tmp_path, monkeypatch):
assert (accepted1, inserted1) == (2, 2)
# legacy-b is already in the ledger, only legacy-c is genuinely new.
assert (accepted2, inserted2) == (2, 1)
assert set(studio_db.list_chat_legacy_imports()) == {
"legacy-a",
"legacy-b",
"legacy-c",
}
assert set(studio_db.list_chat_legacy_imports()) == {"legacy-a", "legacy-b", "legacy-c"}
def test_legacy_imports_dedups_input(tmp_path, monkeypatch):

1
studio/backend/tests/test_cleanup_cancelled_checkpoints.py
View file

@ -130,7 +130,6 @@ def test_symlinked_output_dir_skipped(outputs_setup):
def test_missing_output_dir_is_noop(outputs_setup):
from core.training.training import _cleanup_cancelled_checkpoints
_cleanup_cancelled_checkpoints(outputs_setup / "does-not-exist")
# Should not raise; nothing to assert beyond non-failure.

4
studio/backend/tests/test_cpu_threads.py
View file

@ -63,9 +63,7 @@ def test_cpu_thread_cap_is_opt_in(raw):
# Anything that is not a positive integer raises a clear ValueError.
@pytest.mark.parametrize(
"raw", ["zero", "0", "-3", "1.5", "abc", "8a", "0x4", "1e3", "4 0"]
)
@pytest.mark.parametrize("raw", ["zero", "0", "-3", "1.5", "abc", "8a", "0x4", "1e3", "4 0"])
def test_cpu_thread_cap_requires_positive_integer(raw):
with pytest.raises(ValueError, match = "must be a positive integer"):
configure_cpu_threads({"UNSLOTH_CPU_THREADS": raw})

8
studio/backend/tests/test_dataset_upload_limits.py
View file

@ -40,9 +40,7 @@ def isolate_upload_dir(tmp_path, monkeypatch):
def test_dataset_upload_under_configured_cap_succeeds(isolate_upload_dir):
upload = FakeUploadFile("sample.csv", [b"a,b\n1,2\n"])
response = asyncio.run(
datasets_route.upload_dataset(
cast(UploadFile, upload), current_subject = "test-user"
)
datasets_route.upload_dataset(cast(UploadFile, upload), current_subject = "test-user")
)
stored = Path(response.stored_path)
assert response.filename == "sample.csv"
@ -58,9 +56,7 @@ def test_dataset_upload_over_configured_cap_removes_partial_file(isolate_upload_
)
with pytest.raises(HTTPException) as exc:
asyncio.run(
datasets_route.upload_dataset(
cast(UploadFile, upload), current_subject = "test-user"
)
datasets_route.upload_dataset(cast(UploadFile, upload), current_subject = "test-user")
)
assert exc.value.status_code == 413
assert "Maximum is 1MB" in exc.value.detail

44
studio/backend/tests/test_desktop_auth.py
View file

@ -51,12 +51,8 @@ def auth_client():
def data_recipe_jobs_module():
route_path = (
Path(__file__).resolve().parents[1] / "routes" / "data_recipe" / "jobs.py"
)
spec = importlib.util.spec_from_file_location(
"_desktop_data_recipe_jobs", route_path
)
route_path = Path(__file__).resolve().parents[1] / "routes" / "data_recipe" / "jobs.py"
spec = importlib.util.spec_from_file_location("_desktop_data_recipe_jobs", route_path)
jobs_route = importlib.util.module_from_spec(spec)
assert spec.loader is not None
spec.loader.exec_module(jobs_route)
@ -265,9 +261,7 @@ def test_consume_refresh_token_concurrent_only_one_succeeds(tmp_path, monkeypatc
results = list(pool.map(attempt, range(workers)))
successes = [r for r in results if r is not None]
assert (
len(successes) == 1
), f"expected exactly one consumer to win, got {len(successes)}"
assert len(successes) == 1, f"expected exactly one consumer to win, got {len(successes)}"
assert successes[0] == (storage.DEFAULT_ADMIN_USERNAME, False)
@ -285,9 +279,7 @@ def test_desktop_session_uses_real_admin_identity_for_api_keys():
seed_user(must_change_password = True)
raw = storage.create_desktop_secret()
client = auth_client()
token = client.post("/api/auth/desktop-login", json = {"secret": raw}).json()[
"access_token"
]
token = client.post("/api/auth/desktop-login", json = {"secret": raw}).json()["access_token"]
response = client.post(
"/api/auth/api-keys",
@ -322,9 +314,7 @@ def test_local_recipe_token_authenticates_as_admin_for_desktop_user(loaded_local
scheme = "Bearer",
credentials = local_token,
)
assert (
asyncio.run(get_current_subject(credentials)) == storage.DEFAULT_ADMIN_USERNAME
)
assert asyncio.run(get_current_subject(credentials)) == storage.DEFAULT_ADMIN_USERNAME
def test_local_recipe_token_authenticates_as_admin_for_web_user(loaded_local_model):
@ -345,9 +335,7 @@ def test_local_recipe_token_authenticates_as_admin_for_web_user(loaded_local_mod
scheme = "Bearer",
credentials = local_token,
)
assert (
asyncio.run(get_current_subject(credentials)) == storage.DEFAULT_ADMIN_USERNAME
)
assert asyncio.run(get_current_subject(credentials)) == storage.DEFAULT_ADMIN_USERNAME
def test_desktop_login_rejects_invalid_secret():
@ -480,8 +468,7 @@ def test_health_response_reports_desktop_capability_fields(monkeypatch):
def test_provision_desktop_auth_writes_secret_and_creates_db_without_backend_deps(
tmp_path,
monkeypatch,
tmp_path, monkeypatch
):
auth_dir = tmp_path / "auth"
auth_dir.mkdir()
@ -536,12 +523,9 @@ if result.exit_code != 0:
"""
).fetchone()
app_secrets = {
row["key"]: row["value"]
for row in conn.execute("SELECT key, value FROM app_secrets")
}
refresh_columns = {
row["name"] for row in conn.execute("PRAGMA table_info(refresh_tokens)")
row["key"]: row["value"] for row in conn.execute("SELECT key, value FROM app_secrets")
}
refresh_columns = {row["name"] for row in conn.execute("PRAGMA table_info(refresh_tokens)")}
finally:
conn.close()
@ -633,9 +617,7 @@ def test_update_password_clears_desktop_secret():
raw = storage.create_desktop_secret()
assert storage.validate_desktop_secret(raw) == storage.DEFAULT_ADMIN_USERNAME
changed = storage.update_password(
storage.DEFAULT_ADMIN_USERNAME, "new-admin-password"
)
changed = storage.update_password(storage.DEFAULT_ADMIN_USERNAME, "new-admin-password")
assert changed is True
assert storage.validate_desktop_secret(raw) is None
@ -651,11 +633,7 @@ def test_update_password_on_unknown_user_leaves_desktop_secret_intact():
def test_desktop_auth_provision_has_bounded_timeout():
rs_path = (
Path(__file__).resolve().parents[3]
/ "studio"
/ "src-tauri"
/ "src"
/ "desktop_auth.rs"
Path(__file__).resolve().parents[3] / "studio" / "src-tauri" / "src" / "desktop_auth.rs"
)
src = rs_path.read_text()
start = src.index("async fn provision_desktop_auth(")

Some files were not shown because too many files have changed in this diff Show more