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

[pre-commit.ci] auto fixes from pre-commit.com hooks

Browse source 
for more information, see https://pre-commit.ci
This commit is contained in:
pre-commit-ci[bot] 2026-07-21 23:18:45 +00:00
commit 18673a562a
682 changed files with 25903 additions and 8390 deletions
Download patch file Download diff file Expand all files Collapse all files

64
scripts/check_frontend_dep_removal.py
View file

@ -52,7 +52,9 @@ EXPECTED_NOISE_FILES = {
}
# File types where a quoted string can be a module specifier.
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 import patterns could be a real module reference (.mdx is
# real ESM; .md code fences are not).
SCRIPT_LIKE_EXT = re.compile(r"\.(ts|tsx|js|jsx|mjs|cjs|mdx)$")
@ -249,7 +251,9 @@ 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`, `export { x } from`, `export type { Foo } from`.
if is_script and re.search(
@ -261,12 +265,16 @@ def classify(pkg: str, file: str, content: str) -> str | None:
# HTML script / link. Match pkg as a complete path segment 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):
@ -479,12 +487,18 @@ 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): skip its flags and env prefixes.
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
@ -510,7 +524,9 @@ 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', ...]}` for every package
referenced via its bin name in package.json scripts.
@ -566,7 +582,11 @@ def tsconfig_compiler_types_refs() -> set[str]:
if not isinstance(t, str):
continue
# `vite/client` resolves to the `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
@ -704,7 +724,9 @@ _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]
@ -819,14 +841,18 @@ 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. "
@ -918,7 +944,9 @@ 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()
@ -956,7 +984,9 @@ 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()
@ -980,7 +1010,9 @@ 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
@ -1026,7 +1058,9 @@ 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

8
scripts/check_new_install_scripts.py
View file

@ -38,7 +38,9 @@ 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
@ -161,7 +163,9 @@ def diff_new_install_scripts(base_lock: dict, head_lock: dict) -> list[Finding]:
if key in base:
continue # pre-existing install-script dep; not in scope
name = head[key]
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())

22
scripts/enforce_kwargs_spacing.py
View file

@ -123,7 +123,9 @@ 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):
@ -181,7 +183,11 @@ def remove_blank_after_short_import(text: str) -> tuple[str, bool]:
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):
if (
isinstance(val, list)
and val
and all(isinstance(s, ast.stmt) for s in val)
):
out.append(val)
return out
@ -199,7 +205,9 @@ def remove_blank_after_short_import(text: str) -> tuple[str, bool]:
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)
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)
@ -211,7 +219,13 @@ def remove_blank_after_short_import(text: str) -> tuple[str, bool]:
return "".join(kept), True
_STRING_TRIVIA = (tokenize.NL, tokenize.NEWLINE, tokenize.COMMENT, tokenize.INDENT, tokenize.DEDENT)
_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

8
scripts/lint_workflow_triggers.py
View file

@ -29,7 +29,9 @@ 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]
@ -133,7 +135,9 @@ 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

8
scripts/lockfile_supply_chain_audit.py
View file

@ -459,7 +459,9 @@ 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"
),
)
)
@ -663,7 +665,9 @@ 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

@ -155,7 +155,9 @@ 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 "):
@ -278,7 +280,9 @@ def convert_notebook_to_script(
source_name = source
output_filename = filename.replace(".ipynb", ".py")
output_filename = output_filename.replace("(", "").replace(")", "").replace("-", "_")
output_filename = (
output_filename.replace("(", "").replace(")", "").replace("-", "_")
)
if output_dir:
output_path = os.path.join(output_dir, output_filename)
@ -297,7 +301,9 @@ 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(
@ -311,8 +317,12 @@ 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 compat; pass --no-allow-shell for untrusted notebooks.
parser.add_argument(
"--allow-shell",

61
scripts/notebook_validator.py
View file

@ -87,7 +87,9 @@ 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. ----- #
@ -187,7 +189,9 @@ 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
@ -318,7 +322,9 @@ 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]]:
@ -403,7 +409,9 @@ 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`.
"""
@ -477,7 +485,10 @@ 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():
@ -492,7 +503,9 @@ 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:
@ -680,7 +693,9 @@ 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
@ -771,7 +786,9 @@ 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 of update_all_notebooks.py for canonical phrases;
falls back to POLICY_CLAUSES_DEFAULT (which we use directly today). The
permissive regexes avoid false positives on template rewords."""
@ -831,7 +848,11 @@ 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,
@ -932,7 +953,9 @@ 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
@ -942,7 +965,11 @@ 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(
@ -982,9 +1009,13 @@ 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
@ -1159,7 +1190,9 @@ 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]

64
scripts/scan_npm_packages.py
View file

@ -770,7 +770,8 @@ 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)
@ -867,7 +868,11 @@ def safe_extract(
# each gets its own cap (both are bounded).
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} > "
@ -1195,7 +1200,11 @@ def _format_match(
def _stream_overflow_digest(
matches, lines: list[str], sl_blanked: list[str], ml_blanked: list[str], nl: list[int]
matches,
lines: list[str],
sl_blanked: list[str],
ml_blanked: list[str],
nl: list[int],
) -> tuple[int, str]:
"""A single digest binding the LOGICAL line (the bound bracket-group context,
not just the regex match text) of every overflow match in the iterable, plus
@ -1211,7 +1220,12 @@ def _stream_overflow_digest(
def _fold_overflow_match(
h, m: re.Match, lines: list[str], sl_blanked: list[str], ml_blanked: list[str], nl: list[int]
h,
m: re.Match,
lines: list[str],
sl_blanked: list[str],
ml_blanked: list[str],
nl: list[int],
) -> None:
"""Fold one overflow match's whitespace-normalized logical-line context into the
running hash ``h``. Shared by _stream_overflow_digest and the inline overflow
@ -1241,11 +1255,14 @@ def _evidence(
return ""
lines, sl_blanked, ml_blanked, nl = _index_text(text)
shown = [
_format_match(text, lines, sl_blanked, ml_blanked, nl, m, max_chars) for m in shown_matches
_format_match(text, lines, sl_blanked, ml_blanked, nl, m, max_chars)
for m in shown_matches
]
# Fold the rest (past the cap) into one digest as they arrive, never building a
# second list. Byte-identical to digesting matches[_MAX_EVIDENCE_MATCHES:].
overflow_count, digest = _stream_overflow_digest(it, lines, sl_blanked, ml_blanked, nl)
overflow_count, digest = _stream_overflow_digest(
it, lines, sl_blanked, ml_blanked, nl
)
if overflow_count:
shown.append(f"(+{overflow_count} more) sha256:{digest}")
return " | ".join(shown)
@ -1291,7 +1308,9 @@ _REGEX_PRECEDING_KEYWORDS = frozenset(
"case",
}
)
_IDENT_CHARS = frozenset("abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ0123456789_$")
_IDENT_CHARS = frozenset(
"abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ0123456789_$"
)
def _slash_is_regex(prev_tok: str) -> bool:
@ -1547,7 +1566,9 @@ def scan_package_json(pkg: PackageEntry, rel: str, text: str) -> list[Finding]:
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(
@ -1612,7 +1633,9 @@ def _outbound_host_evidence(text: str, host: str) -> str:
),
# Host-config form: capture the whole line (path/headers/body), so a
# changed outbound payload on the same hostname line reopens the key.
re.compile(rf"[^\n]*(?:host|hostname)\s*:\s*['\"`]{host_re}['\"`][^\n]*", re.IGNORECASE),
re.compile(
rf"[^\n]*(?:host|hostname)\s*:\s*['\"`]{host_re}['\"`][^\n]*", re.IGNORECASE
),
)
# Record EVERY outbound context for the host, not just the first form that
# matches: a file that already has a baselined URL for the host and later adds
@ -1641,12 +1664,16 @@ def _outbound_host_evidence(text: str, host: str) -> str:
claimed.append((m.start(), m.end()))
chosen.append(m)
else:
_fold_overflow_match(overflow_hash, m, lines, sl_blanked, ml_blanked, nl)
_fold_overflow_match(
overflow_hash, m, lines, sl_blanked, ml_blanked, nl
)
overflow_count += 1
if not chosen:
return host
chosen.sort(key = lambda m: m.start())
shown = [_format_match(text, lines, sl_blanked, ml_blanked, nl, m, 1000) for m in chosen]
shown = [
_format_match(text, lines, sl_blanked, ml_blanked, nl, m, 1000) for m in chosen
]
if overflow_count:
shown.append(f"(+{overflow_count} more) sha256:{overflow_hash.hexdigest()}")
return " | ".join(shown)
@ -1730,7 +1757,9 @@ def scan_text_blob(pkg: PackageEntry, rel: str, text: str) -> list[Finding]:
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):
@ -1871,7 +1900,9 @@ def scan_one(pkg: PackageEntry, workspace: Path) -> tuple[list[Finding], str | N
# Mirrors scan_packages.py. Regenerate with ``--write-baseline``.
# ─────────────────────────────────────────────────────────────────────
_DEFAULT_BASELINE_PATH = str(Path(__file__).resolve().parent / "scan_npm_packages_baseline.json")
_DEFAULT_BASELINE_PATH = str(
Path(__file__).resolve().parent / "scan_npm_packages_baseline.json"
)
# Bumped when the entry-key semantics change. v3 adds an evidence hash so a new
# payload under an already-listed package/path/pattern is not auto-suppressed; v2
@ -1959,7 +1990,9 @@ def _load_baseline(path: str) -> set[tuple[str, str, str, str]]:
if not isinstance(e, dict):
continue
try:
evidence_hash = e.get("evidence_hash") or _evidence_hash(e.get("evidence") or "")
evidence_hash = e.get("evidence_hash") or _evidence_hash(
e.get("evidence") or ""
)
if not e.get("evidence_hash"):
legacy += 1
keys.add(
@ -2206,7 +2239,8 @@ 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

144
scripts/scan_packages.py
View file

@ -160,7 +160,9 @@ RE_EMBEDDED_KEYS = re.compile(
)
# Full PEM block (BEGIN..END), used to pin a multiline key body in evidence.
RE_PEM_BLOCK = re.compile(r"-----BEGIN[^\n]*KEY-----.*?-----END[^\n]*KEY-----", re.DOTALL)
RE_PEM_BLOCK = re.compile(
r"-----BEGIN[^\n]*KEY-----.*?-----END[^\n]*KEY-----", re.DOTALL
)
# Cloud metadata / IMDS endpoints
RE_CLOUD_METADATA = re.compile(
@ -324,7 +326,9 @@ 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(
@ -533,7 +537,9 @@ def check_pth_file(content: str, filename: str, package: str) -> list[Finding]:
# A STRING after one of these tokens (and before a NEWLINE) is a bare
# docstring/doctest/prose statement -- the dominant FP source -- so we blank it.
# A string after `=` or `(` is real code and is never blanked.
_LINE_START_TOKENS = frozenset({tokenize.NEWLINE, tokenize.NL, tokenize.INDENT, tokenize.DEDENT})
_LINE_START_TOKENS = frozenset(
{tokenize.NEWLINE, tokenize.NL, tokenize.INDENT, tokenize.DEDENT}
)
def _is_fstring(tok_string: str) -> bool:
@ -1425,7 +1431,9 @@ def _extract_evidence(
if len(head) > _MAX_LINE_CHARS:
head = head[:_MAX_LINE_CHARS] + "..."
return f"L{start}: {head} sha256:{digest}"
return "\n".join(f"L{start + i}: {_cap_line(ln.rstrip())}" for i, ln in enumerate(span))
return "\n".join(
f"L{start + i}: {_cap_line(ln.rstrip())}" for i, ln in enumerate(span)
)
for i, line in enumerate(lines, 1):
if pattern.search(line):
@ -1484,7 +1492,9 @@ def _embedded_key_evidence(content: str) -> str:
ev = _extract_evidence(content, RE_EMBEDDED_KEYS)
blocks = RE_PEM_BLOCK.findall(content)
if blocks:
digest = hashlib.sha256("\n".join(blocks).encode("utf-8", "replace")).hexdigest()
digest = hashlib.sha256(
"\n".join(blocks).encode("utf-8", "replace")
).hexdigest()
ev = f"{ev} sha256:{digest}" if ev else f"sha256:{digest}"
return ev
@ -1794,13 +1804,15 @@ def iter_archive_files(archive_path: str):
# 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
@ -1975,7 +1987,9 @@ _SDIST_DOWNLOAD_TIMEOUT = 180
# Never fetch an archive larger than we would be willing to scan (iter_archive_files cap).
_MAX_SDIST_BYTES = HARD_MAX_TOTAL_BYTES
# Direct sdist bytes only ever come from PyPI's own CDN; refuse anything else.
_TRUSTED_PYPI_HOSTS = frozenset({"files.pythonhosted.org", "pypi.org", "pypi.python.org"})
_TRUSTED_PYPI_HOSTS = frozenset(
{"files.pythonhosted.org", "pypi.org", "pypi.python.org"}
)
def _spec_pin_version(spec: str) -> str | None:
@ -2014,7 +2028,9 @@ def _release_files(meta: dict, version: str | None) -> list[dict]:
def _release_has_wheel(meta: dict, version: str | None) -> bool:
"""True if the (pinned or latest) release publishes any bdist_wheel."""
return any(f.get("packagetype") == "bdist_wheel" for f in _release_files(meta, version))
return any(
f.get("packagetype") == "bdist_wheel" for f in _release_files(meta, version)
)
def _is_trusted_pypi_url(url: str) -> bool:
@ -2140,10 +2156,14 @@ def _download_sdist_direct(
return None, f"refusing non-PyPI sdist URL for {name}: {url[:80]}"
# basename + sanitize keeps the path inside dest; the char class preserves
# the real `.tar.gz` / `.zip` suffix so the archive reader picks the format.
safe_fname = _RE_PKG_NAME_SANITIZE.sub("_", os.path.basename(fname)) or "sdist.tar.gz"
safe_fname = (
_RE_PKG_NAME_SANITIZE.sub("_", os.path.basename(fname)) or "sdist.tar.gz"
)
out = os.path.join(dest, safe_fname)
try:
req = urllib.request.Request(url, headers = {"Accept": "application/octet-stream"})
req = urllib.request.Request(
url, headers = {"Accept": "application/octet-stream"}
)
with urllib.request.urlopen(req, timeout = _SDIST_DOWNLOAD_TIMEOUT) as resp:
if getattr(resp, "status", 200) != 200:
return None, f"sdist HTTP {getattr(resp, 'status', '?')} for {name}"
@ -2158,7 +2178,10 @@ def _download_sdist_direct(
)
return out, None
except Exception as exc:
return None, f"sdist download failed for {name}: {type(exc).__name__}: {str(exc)[:120]}"
return (
None,
f"sdist download failed for {name}: {type(exc).__name__}: {str(exc)[:120]}",
)
def _pip_download_with_deps(
@ -2179,7 +2202,9 @@ def _pip_download_with_deps(
dest,
] + list(specs)
try:
proc = subprocess.run(cmd, capture_output = True, text = True, timeout = timeout, env = env)
proc = subprocess.run(
cmd, capture_output = True, text = True, timeout = timeout, env = env
)
return proc.returncode, proc.stderr or ""
except subprocess.TimeoutExpired:
return 124, "pip download (with deps) timed out"
@ -2219,7 +2244,9 @@ def _resolve_per_spec_with_deps(
spec,
]
try:
proc = subprocess.run(cmd, capture_output = True, text = True, timeout = 300, env = env)
proc = subprocess.run(
cmd, capture_output = True, text = True, timeout = 300, env = env
)
except subprocess.TimeoutExpired:
download_errors.append(f"per-spec --with-deps timed out for {spec}")
continue
@ -2231,7 +2258,9 @@ def _resolve_per_spec_with_deps(
if fpath is None:
download_errors.append(serr or f"sdist fetch failed for {name}")
continue
sdist_dep_followups.extend(_requires_dist_for(name, version, meta, download_errors))
sdist_dep_followups.extend(
_requires_dist_for(name, version, meta, download_errors)
)
continue
# Has a wheel but the full transitive tree won't co-resolve
# (ResolutionImpossible) -- typically a package the requirement file
@ -2251,7 +2280,9 @@ def _resolve_per_spec_with_deps(
spec,
]
try:
nd = subprocess.run(nd_cmd, capture_output = True, text = True, timeout = 180, env = env)
nd = subprocess.run(
nd_cmd, capture_output = True, text = True, timeout = 180, env = env
)
except subprocess.TimeoutExpired:
download_errors.append(f"per-spec --no-deps timed out for {spec}")
continue
@ -2265,7 +2296,9 @@ def _resolve_per_spec_with_deps(
# which --no-deps skips. Recover the declared deps so that class is
# still scanned (each is fetched as a wheel or direct sdist below).
if meta is not None:
sdist_dep_followups.extend(_requires_dist_for(name, version, meta, download_errors))
sdist_dep_followups.extend(
_requires_dist_for(name, version, meta, download_errors)
)
continue
# --no-deps also failed: last-ditch sdist fetch at the pinned version.
if meta is not None:
@ -2273,7 +2306,8 @@ def _resolve_per_spec_with_deps(
if fpath is not None:
continue
download_errors.append(
f"per-spec failed for {spec} (with-deps and --no-deps): " f"{nd.stderr.strip()[:240]}"
f"per-spec failed for {spec} (with-deps and --no-deps): "
f"{nd.stderr.strip()[:240]}"
)
# Recover the transitive deps of sdist-only packages. A depth-bounded,
@ -2302,7 +2336,9 @@ def _resolve_per_spec_with_deps(
dep,
]
try:
proc = subprocess.run(cmd, capture_output = True, text = True, timeout = 300, env = env)
proc = subprocess.run(
cmd, capture_output = True, text = True, timeout = 300, env = env
)
except subprocess.TimeoutExpired:
print(f" [WARN] dep download timed out for {dep}", file = sys.stderr)
continue
@ -2310,14 +2346,21 @@ def _resolve_per_spec_with_deps(
continue
meta = _pypi_json(dep_name)
if meta is None:
print(f" [WARN] could not resolve indirect dep {dep}; skipping", file = sys.stderr)
print(
f" [WARN] could not resolve indirect dep {dep}; skipping",
file = sys.stderr,
)
continue
if not _release_has_wheel(meta, dep_ver):
fpath, serr = _download_sdist_direct(dep_name, dep_ver, dest, meta = meta)
if fpath is None:
print(f" [WARN] could not fetch sdist dep {dep}: {serr}", file = sys.stderr)
print(
f" [WARN] could not fetch sdist dep {dep}: {serr}", file = sys.stderr
)
elif depth < _MAX_DEP_FOLLOWUP_DEPTH:
worklist.extend((d, depth + 1) for d in _requires_dist_for(dep_name, dep_ver, meta))
worklist.extend(
(d, depth + 1) for d in _requires_dist_for(dep_name, dep_ver, meta)
)
continue
# Wheel published but its tree won't co-resolve (a sdist-only child).
# Fetch the dep alone so it is scanned, then chase its own declared deps.
@ -2333,19 +2376,28 @@ def _resolve_per_spec_with_deps(
dep,
]
try:
nd = subprocess.run(nd_cmd, capture_output = True, text = True, timeout = 180, env = env)
nd = subprocess.run(
nd_cmd, capture_output = True, text = True, timeout = 180, env = env
)
except subprocess.TimeoutExpired:
print(f" [WARN] dep --no-deps timed out for {dep}", file = sys.stderr)
continue
if nd.returncode == 0:
if depth < _MAX_DEP_FOLLOWUP_DEPTH:
worklist.extend((d, depth + 1) for d in _requires_dist_for(dep_name, dep_ver, meta))
worklist.extend(
(d, depth + 1) for d in _requires_dist_for(dep_name, dep_ver, meta)
)
continue
fpath, _serr = _download_sdist_direct(dep_name, dep_ver, dest, meta = meta)
if fpath is None:
print(f" [WARN] could not resolve indirect dep {dep}; skipping", file = sys.stderr)
print(
f" [WARN] could not resolve indirect dep {dep}; skipping",
file = sys.stderr,
)
elif depth < _MAX_DEP_FOLLOWUP_DEPTH:
worklist.extend((d, depth + 1) for d in _requires_dist_for(dep_name, dep_ver, meta))
worklist.extend(
(d, depth + 1) for d in _requires_dist_for(dep_name, dep_ver, meta)
)
def download_packages(
@ -2410,7 +2462,9 @@ def download_packages(
spec,
]
try:
proc = subprocess.run(cmd, capture_output = True, text = True, timeout = 120, env = env)
proc = subprocess.run(
cmd, capture_output = True, text = True, timeout = 120, env = env
)
except subprocess.TimeoutExpired:
download_errors.append(f"pip download timed out for {spec}")
continue
@ -2420,7 +2474,9 @@ def download_packages(
version = _spec_pin_version(spec)
meta = _pypi_json(name)
if meta is not None and not _release_has_wheel(meta, version):
fpath, serr = _download_sdist_direct(name, version, pkg_dir, meta = meta)
fpath, serr = _download_sdist_direct(
name, version, pkg_dir, meta = meta
)
if fpath is not None:
results.append((spec, fpath))
continue
@ -2447,7 +2503,9 @@ 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()
)
@ -2790,7 +2848,9 @@ def _run_fix(critical_pkgs: set[str], entries: list[dict], max_search: int) -> N
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
@ -2815,7 +2875,9 @@ def _run_fix(critical_pkgs: set[str], entries: list[dict], max_search: int) -> N
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
@ -2832,7 +2894,9 @@ def _run_fix(critical_pkgs: set[str], entries: list[dict], max_search: int) -> N
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]] = {}
@ -2879,7 +2943,9 @@ 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):
@ -2953,7 +3019,9 @@ def _canon_evidence(evidence: str) -> str:
def _evidence_hash(evidence: str) -> str:
"""Stable digest of the canonical matched evidence."""
return hashlib.sha256(_canon_evidence(evidence).encode("utf-8", "replace")).hexdigest()
return hashlib.sha256(
_canon_evidence(evidence).encode("utf-8", "replace")
).hexdigest()
def _finding_key(f: Finding) -> tuple[str, str, str, str]:
@ -2995,7 +3063,9 @@ def _load_baseline(path: str) -> set[tuple[str, str, str, str]]:
continue
try:
# Use the reviewed hash; else recompute it from the stored evidence.
evidence_hash = e.get("evidence_hash") or _evidence_hash(e.get("evidence") or "")
evidence_hash = e.get("evidence_hash") or _evidence_hash(
e.get("evidence") or ""
)
if not e.get("evidence_hash"):
legacy += 1
keys.add(
@ -3151,7 +3221,9 @@ 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] = []

12
scripts/stamp_studio_release.py
View file

@ -42,7 +42,9 @@ def _atomic_write_text(
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)?$")
@ -233,7 +235,9 @@ def _read_sdist_member(path: Path) -> str | None:
def verify_dist(expected: str, dist_dir: Path) -> int:
if not is_valid_version(expected):
print(f"Invalid expected Unsloth release version: {expected!r}", file = sys.stderr)
print(
f"Invalid expected Unsloth release version: {expected!r}", file = sys.stderr
)
return 2
artifacts = list(dist_dir.glob("*.whl")) + list(dist_dir.glob("*.tar.gz"))
@ -258,7 +262,9 @@ def verify_dist(expected: str, dist_dir: Path) -> int:
print(failure, file = sys.stderr)
return 2
print(f"Verified Unsloth release version {expected} in {len(artifacts)} artifact(s)")
print(
f"Verified Unsloth release version {expected} in {len(artifacts)} artifact(s)"
)
return 0

24
scripts/sync_allow_scripts_pins.py
View file

@ -74,7 +74,9 @@ def desired_key(name: str, versions: list[str]) -> str:
return f"{name}@{' || '.join(versions)}"
def compute_renames(policy: dict, lock_versions: dict[str, list[str]]) -> dict[str, str]:
def compute_renames(
policy: dict, lock_versions: dict[str, list[str]]
) -> dict[str, str]:
renames: dict[str, str] = {}
for key in policy:
name, rng = split_spec(key)
@ -93,7 +95,9 @@ def main(argv: list[str] | None = None) -> int:
ap = argparse.ArgumentParser(description = __doc__)
mode = ap.add_mutually_exclusive_group(required = True)
mode.add_argument("--check", action = "store_true", help = "exit 1 if pins are stale")
mode.add_argument("--fix", action = "store_true", help = "rewrite package.json in place")
mode.add_argument(
"--fix", action = "store_true", help = "rewrite package.json in place"
)
ap.add_argument(
"--dir",
type = Path,
@ -105,20 +109,26 @@ def main(argv: list[str] | None = None) -> int:
pkg_path = args.dir / "package.json"
lock_path = args.dir / "package-lock.json"
if not pkg_path.exists() or not lock_path.exists():
print(f"sync-allow-scripts: nothing to do ({args.dir} has no package.json + lockfile)")
print(
f"sync-allow-scripts: nothing to do ({args.dir} has no package.json + lockfile)"
)
return 0
pkg = json.loads(pkg_path.read_text(encoding = "utf-8"))
policy = pkg.get("allowScripts")
if not isinstance(policy, dict) or not policy:
print("sync-allow-scripts: no allowScripts policy in package.json, nothing to do")
print(
"sync-allow-scripts: no allowScripts policy in package.json, nothing to do"
)
return 0
lock = json.loads(lock_path.read_text(encoding = "utf-8"))
renames = compute_renames(policy, script_versions_from_lock(lock))
if not renames:
print(f"sync-allow-scripts: {len(policy)} allowScripts entries in sync with the lockfile")
print(
f"sync-allow-scripts: {len(policy)} allowScripts entries in sync with the lockfile"
)
return 0
for old, new in renames.items():
@ -132,7 +142,9 @@ def main(argv: list[str] | None = None) -> int:
return 1
pkg["allowScripts"] = {renames.get(k, k): v for k, v in policy.items()}
pkg_path.write_text(json.dumps(pkg, indent = 2, ensure_ascii = False) + "\n", encoding = "utf-8")
pkg_path.write_text(
json.dumps(pkg, indent = 2, ensure_ascii = False) + "\n", encoding = "utf-8"
)
print(
f"sync-allow-scripts: re-pinned {len(renames)} entr{'y' if len(renames) == 1 else 'ies'} in {pkg_path}"
)

3
scripts/verify_comment_only_diff.py
View file

@ -131,7 +131,8 @@ def _walk_yaml_diff(
"""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):

34
scripts/verify_import_hoist.py
View file

@ -161,7 +161,9 @@ 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:
@ -349,7 +351,9 @@ 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 evaluates in the enclosing scope
@ -608,7 +612,9 @@ def compare(before_src: str, after_src: str, path: str) -> list[tuple[str, str]]
if tbefore and tbefore != tafter and (tbefore - tafter):
lost = tbefore - tafter
gained = tafter - tbefore
relocated = lost <= removed_module_targets and gained <= added_module_targets
relocated = (
lost <= removed_module_targets and gained <= added_module_targets
)
if relocated:
continue
findings.append(
@ -633,7 +639,9 @@ 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 covered above; remaining cases are relocated code.
@ -645,7 +653,9 @@ 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
@ -790,7 +800,9 @@ 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
@ -826,12 +838,18 @@ 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

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

@ -108,7 +108,9 @@ 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.
@ -135,7 +137,9 @@ 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,

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

@ -195,9 +195,13 @@ 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 (
@ -211,9 +215,13 @@ 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

14
studio/backend/auth/terminal_prompt.py
View file

@ -197,7 +197,11 @@ def _read_password(prompt: str, *, out: "TextIO | None" = None) -> str:
def should_prompt_password_change(
*, tunnel_will_start: bool, requires_change: bool, stdin_isatty: bool, stderr_isatty: bool
*,
tunnel_will_start: bool,
requires_change: bool,
stdin_isatty: bool,
stderr_isatty: bool,
) -> bool:
"""Whether to block startup on an interactive terminal password change.
@ -233,7 +237,9 @@ def prompt_for_password_change(
while True:
new_password = _read_password("New password: ", out = out)
if len(new_password) < min_length:
out.write(f"Password must be at least {min_length} characters; try again.\n")
out.write(
f"Password must be at least {min_length} characters; try again.\n"
)
out.flush()
continue
if is_current_password(new_password):
@ -257,7 +263,9 @@ def prompt_for_password_change(
return False
def resolve_supplied_password(cli_value: "str | None", out: "TextIO | None" = None) -> "str | None":
def resolve_supplied_password(
cli_value: "str | None", out: "TextIO | None" = None
) -> "str | None":
"""Resolve a non-interactive initial admin password, or None if unset.
Precedence: an explicit ``--password`` (literal ``-`` reads a line from

8
studio/backend/cloudflare_tunnel.py
View file

@ -87,7 +87,9 @@ def _asset_name() -> Optional[Tuple[str, bool]]:
def _cache_path() -> Optional[Path]:
"""studio_bin_root()/cloudflared(.exe), or None if the studio home is unresolvable."""
try:
from utils.paths.storage_roots import studio_bin_root # lazy: backend-only import
from utils.paths.storage_roots import (
studio_bin_root,
) # lazy: backend-only import
except Exception:
return None
name = "cloudflared.exe" if sys.platform == "win32" else "cloudflared"
@ -273,7 +275,9 @@ class CloudflareTunnel:
if self.url is None:
self.error = "cloudflared exited before emitting a tunnel URL"
elif not self.ready:
self.error = "cloudflared exited before the tunnel connection registered"
self.error = (
"cloudflared exited before the tunnel connection registered"
)
self._url_event.set()
self._ready_event.set()

23
studio/backend/colab.py
View file

@ -41,7 +41,12 @@ def get_colab_url(port: int = 8888) -> str:
try:
url = eval_js(f"google.colab.kernel.proxyPort({port})", timeout_sec = 10)
# Valid proxy URL is 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})")
@ -114,7 +119,9 @@ def _bootstrap_password_pending() -> bool:
from auth.storage import requires_password_change, DEFAULT_ADMIN_USERNAME
return bool(requires_password_change(DEFAULT_ADMIN_USERNAME))
except Exception as e:
logger.info(f"Could not check admin password state ({e}); refusing tunnel to be safe.")
logger.info(
f"Could not check admin password state ({e}); refusing tunnel to be safe."
)
return True
@ -188,7 +195,9 @@ def _is_studio_healthy(port: int, timeout: float = 2.0) -> bool:
"""
import json, urllib.request
try:
with urllib.request.urlopen(f"http://localhost:{port}/api/health", timeout = timeout) as r:
with urllib.request.urlopen(
f"http://localhost:{port}/api/health", timeout = timeout
) as r:
return json.loads(r.read()).get("service") == "Unsloth UI Backend"
except Exception:
return False
@ -295,7 +304,9 @@ def start(port: int = 8888, *, cloudflare: bool = False):
# Fast path: Unsloth already running (cell re-run). Re-launching would collide on
# the port, so just re-show the link and iframe.
if _is_studio_healthy(port):
logger.info(f" Unsloth is already running on port {port} — reusing existing server.")
logger.info(
f" Unsloth is already running on port {port} — reusing existing server."
)
# try/finally: tear the tunnel down even if interrupted mid-start/render.
try:
cf_url = start_cloudflare_tunnel(port) if cloudflare else None
@ -352,7 +363,9 @@ def start(port: int = 8888, *, cloudflare: bool = False):
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:

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

@ -36,7 +36,9 @@ 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

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

@ -111,7 +111,9 @@ 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,7 +160,9 @@ 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
@ -188,7 +192,9 @@ 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:
@ -199,7 +205,9 @@ 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):
@ -316,12 +324,16 @@ 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,
@ -460,8 +472,12 @@ class JobManager:
try:
self._handle_event(job, event)
except Exception:
etype = event.get("type") if isinstance(event, dict) else type(event).__name__
logger.exception("Data-recipe job pump: failed to handle %s event; skipping", etype)
etype = (
event.get("type") if isinstance(event, dict) else type(event).__name__
)
logger.exception(
"Data-recipe job pump: failed to handle %s event; skipping", etype
)
def _pump_loop(self) -> None:
"""Background thread: consume worker events and update the job snapshot.
@ -527,7 +543,9 @@ class JobManager:
if retired_job is not None:
self._retire_workflow_key(retired_job)
except Exception:
logger.exception("Data-recipe job pump: finalization after worker exit failed")
logger.exception(
"Data-recipe job pump: finalization after worker exit failed"
)
return
def _handle_event(self, job: Job, event: dict) -> None:

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

@ -119,7 +119,8 @@ 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})"
),
),
)
@ -133,7 +134,9 @@ def parse_log_message(msg: str) -> ParsedUpdate | None:
source = "github",
status = "rate_limited",
retry_after_sec = seconds,
message = ("Waiting for GitHub rate limit. Unsloth will resume automatically."),
message = (
"Waiting for GitHub rate limit. Unsloth will resume automatically."
),
),
)
@ -161,7 +164,9 @@ def parse_log_message(msg: str) -> ParsedUpdate | None:
source = "github",
status = "rate_limited",
retry_after_sec = seconds,
message = ("Waiting for GitHub rate limit. Unsloth will resume automatically."),
message = (
"Waiting for GitHub rate limit. Unsloth will resume automatically."
),
),
)
@ -379,13 +384,15 @@ 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
@ -405,10 +412,14 @@ 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,
@ -439,7 +450,9 @@ 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)

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

@ -60,7 +60,9 @@ 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
@ -72,7 +74,9 @@ def _build_dataset_name(*, run_name: str | None, job_id: str, artifact_root: Pat
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
@ -160,10 +164,14 @@ def run_job_process(*, event_queue, recipe: dict[str, Any], run: dict[str, Any])
}
)
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(
{

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

@ -134,7 +134,11 @@ def _parse_oxc_spec(*, column: dict[str, Any]) -> OxcLocalCallableValidatorSpec
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 []
)
@ -174,7 +178,9 @@ 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
@ -195,7 +201,10 @@ 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(
@ -211,9 +220,7 @@ 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
@ -306,13 +313,21 @@ 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,

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

@ -24,7 +24,9 @@ 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,7 +121,9 @@ 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)
@ -161,7 +165,9 @@ 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.")
@ -251,7 +257,9 @@ 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,
@ -327,10 +335,14 @@ 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

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

@ -19,7 +19,9 @@ from typing import Optional, Tuple, List
try:
from unsloth import FastLanguageModel, FastVisionModel, _IS_MLX
_UNSLOTH_IMPORT_ERROR = None
except Exception as _unsloth_exc: # ImportError (e.g. missing torch) or a broken native load
except (
Exception
) as _unsloth_exc: # ImportError (e.g. missing torch) or a broken native load
FastLanguageModel = None
FastVisionModel = None
_IS_MLX = False
@ -89,7 +91,9 @@ def _supports_kwarg(fn, name):
params = inspect.signature(fn).parameters
except (TypeError, ValueError):
return False
return name in params or any(p.kind == inspect.Parameter.VAR_KEYWORD for p in params.values())
return name in params or any(
p.kind == inspect.Parameter.VAR_KEYWORD for p in params.values()
)
def _compressed_export_supported():
@ -118,7 +122,10 @@ def _has_nvidia_gpu():
except Exception:
try:
import torch
return bool(torch.cuda.is_available()) and getattr(torch.version, "hip", None) is None
return (
bool(torch.cuda.is_available())
and getattr(torch.version, "hip", None) is None
)
except Exception:
return False
@ -134,14 +141,21 @@ def _hf_offline(timeout = 3):
or os.environ.get("TRANSFORMERS_OFFLINE", "").strip().lower() in _offline
):
return True
if os.environ.get("UNSLOTH_OFFLINE_PROBE", "1").strip().lower() in {"0", "false", "no", "off"}:
if os.environ.get("UNSLOTH_OFFLINE_PROBE", "1").strip().lower() in {
"0",
"false",
"no",
"off",
}:
return False # probe disabled -> assume online; loads still pass local_files_only on env
# Shared bounded, proxy-aware probe (also used by the export worker before version activation).
from utils.transformers_version import hf_endpoint_unreachable
if hf_endpoint_unreachable(timeout):
logger.warning("Hugging Face endpoint unreachable; loading checkpoint in offline mode")
logger.warning(
"Hugging Face endpoint unreachable; loading checkpoint in offline mode"
)
return True
return False
@ -183,7 +197,9 @@ 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
@ -542,7 +558,9 @@ class ExportBackend:
# through it when available; else fall back to the workspace 0.10.x path below.
_shadow_pp = None
try:
from utils.transformers_version import llmcompressor_shadow_pythonpath
from utils.transformers_version import (
llmcompressor_shadow_pythonpath,
)
_shadow_pp = llmcompressor_shadow_pythonpath()
except Exception as e:
logger.warning(f"llm-compressor-main shadow unavailable: {e}")
@ -552,7 +570,9 @@ class ExportBackend:
# No shadow (disabled/offline/failed): the workspace 0.10.x cannot exceed its
# transformers ceiling, so fail fast for sidecar models; default-tier still works.
os.environ.pop(_us._COMPRESSED_QUANTIZE_PYTHONPATH_ENV, None)
_exceeds, _tf_ver = _us._transformers_exceeds_llm_compressor_ceiling()
_exceeds, _tf_ver = (
_us._transformers_exceeds_llm_compressor_ceiling()
)
if _exceeds:
return (
False,
@ -567,7 +587,11 @@ class ExportBackend:
try:
info = _us._normalize_compressed_method(compressed_alias)
except Exception as e:
return False, f"Unsupported compressed export '{compressed_alias}': {e}", None
return (
False,
f"Unsupported compressed export '{compressed_alias}': {e}",
None,
)
if info is None:
return (
False,
@ -577,7 +601,9 @@ class ExportBackend:
compressed_suffix = info[2]
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"
)
elif is_compressed or is_torchao:
save_method = compressed_alias
elif format_type == "4-bit (FP4)":
@ -646,7 +672,11 @@ class ExportBackend:
token = hf_token,
private = private,
)
elif (is_compressed or is_torchao) and output_path and Path(output_path).is_dir():
elif (
(is_compressed or is_torchao)
and output_path
and Path(output_path).is_dir()
):
# Already built in output_path; upload it directly instead of re-running the
# expensive quantization that push_to_hub_merged(save_method=...) would redo.
hf_api = HfApi(token = hf_token)
@ -658,8 +688,12 @@ class ExportBackend:
)
content = MODEL_CARD.format(
username = repo_id.split("/")[0],
base_model = getattr(self.current_model.config, "_name_or_path", "unknown"),
model_type = getattr(self.current_model.config, "model_type", "llm"),
base_model = getattr(
self.current_model.config, "_name_or_path", "unknown"
),
model_type = getattr(
self.current_model.config, "model_type", "llm"
),
method = compressed_alias or format_type,
extra = "unsloth",
)
@ -672,7 +706,9 @@ class ExportBackend:
repo_type = "model",
)
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,
@ -778,7 +814,9 @@ class ExportBackend:
else:
# 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"
)
hf_api = HfApi(token = hf_token)
@ -798,7 +836,9 @@ 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"
)
if save_directory:
hf_api.upload_folder(
@ -870,7 +910,9 @@ class ExportBackend:
try:
# Normalize to a lowercased list so multiple quants come from one model load.
if isinstance(quantization_method, (list, tuple)):
quant_methods = [str(q).lower() for q in quantization_method if str(q).strip()]
quant_methods = [
str(q).lower() for q in quantization_method if str(q).strip()
]
else:
quant_methods = [str(quantization_method).lower()]
if not quant_methods:
@ -885,7 +927,9 @@ 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(
@ -913,12 +957,16 @@ class ExportBackend:
cwd = os.getcwd()
pre_existing_ggufs = set(glob.glob(os.path.join(cwd, "*.gguf")))
pre_existing_subs = {d.name for d in Path(abs_save_dir).iterdir() if d.is_dir()}
pre_existing_subs = {
d.name for d in Path(abs_save_dir).iterdir() if d.is_dir()
}
# Avoid clobbering an existing user-owned model/ directory.
import uuid
_model_tmp = os.path.join(abs_save_dir, f"_tmp_model_{uuid.uuid4().hex[:8]}")
_model_tmp = os.path.join(
abs_save_dir, f"_tmp_model_{uuid.uuid4().hex[:8]}"
)
model_tmp_to_cleanup = _model_tmp
self.current_model.save_pretrained_gguf(
_model_tmp,
@ -928,11 +976,15 @@ class ExportBackend:
)
# Relocate the .gguf that convert_to_gguf wrote to cwd (repo root).
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 GGUF files from subdirs created during this export.
for sub in list(Path(abs_save_dir).iterdir()):
@ -952,7 +1004,10 @@ class ExportBackend:
if self.current_checkpoint:
ckpt = Path(self.current_checkpoint)
gguf_dir = ckpt.parent / f"{ckpt.name}_gguf"
if gguf_dir.is_dir() and gguf_dir.resolve() != Path(abs_save_dir).resolve():
if (
gguf_dir.is_dir()
and gguf_dir.resolve() != Path(abs_save_dir).resolve()
):
for src in gguf_dir.glob("*.gguf"):
dest = os.path.join(abs_save_dir, src.name)
shutil.move(str(src), dest)
@ -960,7 +1015,9 @@ 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}")
@ -1059,7 +1116,9 @@ class ExportBackend:
# getattr so an older build without save_pretrained_gguf returns a clean message
# instead of an AttributeError (a generic 500).
_save_gguf_fn = getattr(self.current_model, "save_pretrained_gguf", None)
if _save_gguf_fn is None or not _supports_kwarg(_save_gguf_fn, "save_method"):
if _save_gguf_fn is None or not _supports_kwarg(
_save_gguf_fn, "save_method"
):
return (
False,
"This Unsloth build does not support GGUF LoRA adapter export. "
@ -1085,11 +1144,14 @@ class ExportBackend:
# Forward the token so convert_lora_to_gguf.py can fetch a gated base's config.
token = hf_token or None,
)
final_ggufs = sorted(glob.glob(os.path.join(save_directory, "*.gguf")))
final_ggufs = sorted(
glob.glob(os.path.join(save_directory, "*.gguf"))
)
logger.info(
"LoRA GGUF export complete. Files in %s:\n %s",
save_directory,
"\n ".join(os.path.basename(f) for f in final_ggufs) or "(none)",
"\n ".join(os.path.basename(f) for f in final_ggufs)
or "(none)",
)
elif _IS_MLX:
# MLX: save adapters.safetensors + tokenizer files
@ -1140,8 +1202,12 @@ 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

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

@ -141,7 +141,9 @@ class ExportOrchestrator:
"""True if the in-flight (or most recent) run was cancelled by the user."""
return self._cancel_requested
def _record_op_finished(self, success: bool, message: str, output_path: Optional[str]) -> None:
def _record_op_finished(
self, success: bool, message: str, output_path: Optional[str]
) -> None:
"""Snapshot the just-finished op so status pollers can recover its outcome.
Called from each op's ``finally`` (with ``_active_op_kind`` still set) BEFORE
@ -150,7 +152,11 @@ class ExportOrchestrator:
"""
with self._op_lock:
self._op_seq += 1
status = "cancelled" if self._cancel_requested else ("success" if success else "error")
status = (
"cancelled"
if self._cancel_requested
else ("success" if success else "error")
)
self._last_op = {
"seq": self._op_seq,
"kind": self._active_op_kind,
@ -220,7 +226,9 @@ class ExportOrchestrator:
# Inside an active op an INSTALL reservation is about to abort on the
# is_export_active check, but a lazy REPAIR has no such check and can be
# rebuilding the sidecar right now, so it must always refuse the spawn.
if _swap_kind == "repair" or (_swap_kind is not None and not self._export_active):
if _swap_kind == "repair" or (
_swap_kind is not None and not self._export_active
):
from utils.transformers_version import SidecarSwapInProgress
raise SidecarSwapInProgress(
"A transformers installation is replacing the latest sidecar; "
@ -397,7 +405,9 @@ 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."""
@ -475,7 +485,9 @@ 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
)
try:
self._spawn_subprocess(sub_config)
except Exception:
@ -502,7 +514,10 @@ class ExportOrchestrator:
self.is_vision = resp.get("is_vision", False)
self.is_peft = resp.get("is_peft", False)
logger.info("Checkpoint '%s' loaded in subprocess", checkpoint_path)
op_success, op_message = True, resp.get("message", "Loaded successfully")
op_success, op_message = (
True,
resp.get("message", "Loaded successfully"),
)
return True, op_message
else:
error = resp.get("message", "Failed to load checkpoint")
@ -609,7 +624,9 @@ 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 and wait for the result.
Returns ``(success, message, output_path)``. ``output_path`` is the on-disk
@ -696,7 +713,9 @@ class ExportOrchestrator:
self._active_op_kind = None
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 — runs locally, no ML imports."""
from utils.models.checkpoints import scan_checkpoints
return scan_checkpoints(outputs_dir = outputs_dir)

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

@ -160,7 +160,9 @@ def _setup_log_capture(resp_queue: Any) -> None:
t_err.start()
def _activate_transformers_version(model_name: str, hf_token: str | None = None) -> None:
def _activate_transformers_version(
model_name: str, hf_token: str | None = None
) -> None:
"""Activate the correct transformers version BEFORE any ML imports."""
# Ensure backend is on sys.path for utils imports.
backend_path = str(Path(__file__).resolve().parent.parent.parent)
@ -187,7 +189,9 @@ def _offline_window_if_unreachable(step = "loading"):
force_ctx = None
try:
from utils.transformers_version import _env_offline, hf_endpoint_unreachable
probe_enabled = os.environ.get("UNSLOTH_OFFLINE_PROBE", "1").strip().lower() not in (
probe_enabled = os.environ.get(
"UNSLOTH_OFFLINE_PROBE", "1"
).strip().lower() not in (
"0",
"false",
"no",
@ -279,7 +283,9 @@ def _handle_load(backend, cmd: dict, resp_queue: Any) -> None:
from utils.models.model_config import get_base_model_from_lora_identifier
# Resolve a LOCAL or REMOTE adapter's base so a remote LoRA base is gated too.
_base = get_base_model_from_lora_identifier(checkpoint_path, cmd.get("hf_token"))
_base = get_base_model_from_lora_identifier(
checkpoint_path, cmd.get("hf_token")
)
if _base:
malware_targets.append(_base)
except Exception as exc:
@ -287,7 +293,9 @@ def _handle_load(backend, cmd: dict, resp_queue: Any) -> None:
_hf_token = cmd.get("hf_token")
for target in dict.fromkeys(malware_targets):
_fs = evaluate_file_security(
target, hf_token = _hf_token, load_subdirs = security_load_subdirs(target, _hf_token)
target,
hf_token = _hf_token,
load_subdirs = security_load_subdirs(target, _hf_token),
)
if _fs.blocked:
_send_response(
@ -313,7 +321,9 @@ def _handle_load(backend, cmd: dict, resp_queue: Any) -> None:
from utils.models.model_config import get_base_model_from_lora_identifier
# Resolve a local or remote adapter's base so its base repo is gated too.
base_model = get_base_model_from_lora_identifier(checkpoint_path, cmd.get("hf_token"))
base_model = get_base_model_from_lora_identifier(
checkpoint_path, cmd.get("hf_token")
)
if base_model:
consent_targets.append(base_model)
except Exception as exc:
@ -541,7 +551,9 @@ def run_export_process(*, cmd_queue: Any, resp_queue: Any, config: dict) -> None
# ── 1. Activate correct transformers version BEFORE any ML imports ──
with _offline_window_if_unreachable(step = "activating transformers"):
try:
_activate_transformers_version(checkpoint_path, config.get("hf_token") or None)
_activate_transformers_version(
checkpoint_path, config.get("hf_token") or None
)
except Exception as exc:
_send_response(
resp_queue,
@ -597,7 +609,9 @@ def run_export_process(*, cmd_queue: Any, resp_queue: Any, config: dict) -> None
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(
@ -703,7 +717,9 @@ def run_export_process(*, cmd_queue: Any, resp_queue: Any, config: dict) -> None
)
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,
{

4
studio/backend/core/import_guards.py
View file

@ -42,7 +42,9 @@ def ensure_real_packages(*names: str) -> None:
saved = list(sys.path)
sys.path[:] = [e for e in sys.path if e not in bad]
for name in shadowed:
for cached in [m for m in list(sys.modules) if m == name or m.startswith(name + ".")]:
for cached in [
m for m in list(sys.modules) if m == name or m.startswith(name + ".")
]:
del sys.modules[cached]
try:
importlib.invalidate_caches()

20
studio/backend/core/inference/_html_to_md.py
View file

@ -366,7 +366,11 @@ class _MarkdownRenderer(HTMLParser):
while self._hidden_marks and self._hidden_marks[-1] >= i:
self._hidden_marks.pop()
break
if self._scope_tags is not None and tag in self._scope_tags and self._scope_depth > 0:
if (
self._scope_tags is not None
and tag in self._scope_tags
and self._scope_depth > 0
):
self._scope_depth -= 1
if self._scope_depth == 0 and self._scope_seg_start is not None:
self.scope_segments.append("".join(self._out[self._scope_seg_start :]))
@ -690,9 +694,13 @@ def _line_is_boilerplate(line: str) -> bool:
normalized = re.sub(r"\s+", " ", line).strip().casefold()
if not normalized:
return False
segments = [segment.strip().rstrip(".!:") for segment in re.split(r"[.!]", normalized)]
segments = [
segment.strip().rstrip(".!:") for segment in re.split(r"[.!]", normalized)
]
segments = [segment for segment in segments if segment]
return bool(segments) and all(segment in _BOILERPLATE_NORMALIZED for segment in segments)
return bool(segments) and all(
segment in _BOILERPLATE_NORMALIZED for segment in segments
)
def _strip_boilerplate_lines(text: str) -> str:
@ -707,7 +715,11 @@ def _strip_boilerplate_lines(text: str) -> str:
in_fence = not in_fence
out.append(line)
continue
if not in_fence and len(line) <= _BOILERPLATE_MAX_LINE_CHARS and _line_is_boilerplate(line):
if (
not in_fence
and len(line) <= _BOILERPLATE_MAX_LINE_CHARS
and _line_is_boilerplate(line)
):
continue
out.append(line)
# Collapse blank runs the dropped lines may have left behind.

8
studio/backend/core/inference/_vulkan_probe.py
View file

@ -50,7 +50,9 @@ def _igpu_flags(base, lib, count: int) -> list[bool]:
for i in range(min(count, dev_count)):
dev = base.ggml_backend_reg_dev_get(reg, i)
if dev:
flags[i] = base.ggml_backend_dev_type(dev) == _GGML_BACKEND_DEVICE_TYPE_IGPU
flags[i] = (
base.ggml_backend_dev_type(dev) == _GGML_BACKEND_DEVICE_TYPE_IGPU
)
except Exception:
# Best-effort: any failure degrades to "discrete" so the memory
# readings still get through instead of crashing the probe.
@ -100,7 +102,9 @@ def main() -> int:
rows = []
for i in range(count):
free, total = ctypes.c_size_t(0), ctypes.c_size_t(0)
lib.ggml_backend_vk_get_device_memory(i, ctypes.byref(free), ctypes.byref(total))
lib.ggml_backend_vk_get_device_memory(
i, ctypes.byref(free), ctypes.byref(total)
)
rows.append("%d\t%d\t%d\t%d" % (i, free.value, int(igpu[i]), total.value))
sys.stdout.write("\n".join(rows))
return 0

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

@ -36,7 +36,11 @@ def anthropic_tool_use_id(upstream_id = None) -> str:
"""Return an Anthropic-style tool_use id (prefix 'toolu_'). Reuses an
upstream id only if it already starts with 'toolu_'; otherwise mints a fresh
'toolu_<24 hex>'."""
if upstream_id and isinstance(upstream_id, str) and upstream_id.startswith("toolu_"):
if (
upstream_id
and isinstance(upstream_id, str)
and upstream_id.startswith("toolu_")
):
return upstream_id
return f"toolu_{uuid.uuid4().hex[:24]}"
@ -149,7 +153,9 @@ 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(
{
@ -459,7 +465,9 @@ class AnthropicStreamEmitter:
events.append(self._close_block())
# Reuse the id published in content_block_start; fall back to mapping
# the raw id only if no tool_start preceded this end.
tool_use_id = self._open_tool_use_id or anthropic_tool_use_id(event.get("tool_call_id", ""))
tool_use_id = self._open_tool_use_id or anthropic_tool_use_id(
event.get("tool_call_id", "")
)
self._open_tool_call_id = None
self._open_tool_use_id = None
self._open_tool_args_sent = False
@ -596,7 +604,11 @@ class AnthropicPassthroughEmitter:
# ── Structured tool calls take precedence over healing ──
# Grammar mode worked: flush anything the healer held (it preceded the
# call in the model's output) and relay verbatim from here on.
if delta.get("tool_calls") and self._healer is not None and not self._healer.dormant:
if (
delta.get("tool_calls")
and self._healer is not None
and not self._healer.dormant
):
for kind, value in self._healer.structured_tool_call_seen():
if kind == "text" and value:
events.extend(self._emit_text_delta(value))

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

@ -183,7 +183,9 @@ class ApiMonitor:
):
# Derive only when no authoritative total has been set;
# a later partial chunk must not clobber a provider total.
entry.total_tokens = (entry.prompt_tokens or 0) + (entry.completion_tokens or 0)
entry.total_tokens = (entry.prompt_tokens or 0) + (
entry.completion_tokens or 0
)
if context_length is not None:
entry.context_length = context_length
entry.updated_at = time.time()
@ -266,7 +268,8 @@ class ApiMonitor:
return sum(
1
for entry in self._entries
if entry.status == "running" and (subject is None or entry.subject == subject)
if entry.status == "running"
and (subject is None or entry.subject == subject)
)
def clear(self) -> None:

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

@ -77,7 +77,9 @@ 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(
@ -92,7 +94,9 @@ class AudioCodecManager:
# Clone SparkAudio/Spark-TTS for the sparktts package (HF model repos
# don't contain it)
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}...")
@ -175,7 +179,9 @@ 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.
Finds the START_OF_SPEECH (128257) marker, extracts codes after it,
@ -188,7 +194,9 @@ class AudioCodecManager:
cropped = generated_ids[:, token_indices[1][-1] + 1 :]
else:
# Fall back to the entire output if the marker is missing
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]
@ -214,7 +222,8 @@ 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():
@ -241,12 +250,16 @@ 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);
# pad with zeros or truncate.

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

@ -75,7 +75,9 @@ def resolve_chat_turn_end_eos_ids_using(template_tokenizer, id_tokenizer) -> lis
original tokenizer, so resolving ids on the mapped tokenizer would store the wrong
(doc-eos) id and let generation run past the real turn marker."""
ids = _eos_id_set(getattr(id_tokenizer, "eos_token_id", None))
template = _collect_template_text(getattr(template_tokenizer, "chat_template", None))
template = _collect_template_text(
getattr(template_tokenizer, "chat_template", None)
)
if not template or any(h in template for h in _HARMONY_MARKERS):
return sorted(ids)
unk = getattr(id_tokenizer, "unk_token_id", None)

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

@ -30,7 +30,9 @@ def _tokenizer_objects(tokenizer) -> tuple:
if tokenizer is None:
return ()
nested = getattr(tokenizer, "tokenizer", None)
return (tokenizer,) if nested is None or nested is tokenizer else (tokenizer, nested)
return (
(tokenizer,) if nested is None or nested is tokenizer else (tokenizer, nested)
)
def _selected_template_strings_from_value(
@ -81,12 +83,18 @@ def _detect_reasoning_channel_markers_from_templates(
templates: tuple[str, ...],
) -> Optional[tuple[str, str]]:
"""Return Gemma native reasoning markers only when a template emits them."""
if any(opener in template for template in templates for opener in _GEMMA_TEMPLATE_OPENERS):
if any(
opener in template
for template in templates
for opener in _GEMMA_TEMPLATE_OPENERS
):
return _GEMMA_THOUGHT_OPEN, _GEMMA_THOUGHT_CLOSE
return None
def detect_reasoning_channel_markers(tokenizer, tools = None) -> Optional[tuple[str, str]]:
def detect_reasoning_channel_markers(
tokenizer, tools = None
) -> Optional[tuple[str, str]]:
"""Return native Gemma thought-channel markers supported by a tokenizer.
Detection uses the active chat template rather than model names or vocabulary
@ -181,7 +189,9 @@ class ReasoningChannelNormalizer:
if not self._buffer:
break
marker = self._closing_marker if self._in_reasoning else self._opening_marker
marker = (
self._closing_marker if self._in_reasoning else self._opening_marker
)
index = self._buffer.find(marker)
if index < 0:
stable, self._buffer = _split_partial_marker(self._buffer, marker)
@ -234,7 +244,9 @@ def normalize_reasoning_snapshots(
normalized_output = ""
for snapshot in stream:
if not snapshot.startswith(raw_output):
raise RuntimeError("Reasoning normalization requires cumulative text snapshots")
raise RuntimeError(
"Reasoning normalization requires cumulative text snapshots"
)
delta = normalizer.feed(snapshot[len(raw_output) :])
raw_output = snapshot
if delta:
@ -373,7 +385,9 @@ 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"
)
try:
return _render(messages)

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

File diff suppressed because it is too large Load diff

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

@ -206,7 +206,9 @@ class ReasoningTextIteratorStreamer(TextIteratorStreamer):
**decode_kwargs,
):
decode_kwargs["skip_special_tokens"] = False
super().__init__(tokenizer, skip_prompt = skip_prompt, timeout = timeout, **decode_kwargs)
super().__init__(
tokenizer, skip_prompt = skip_prompt, timeout = timeout, **decode_kwargs
)
self._normalizer = ReasoningChannelNormalizer(*markers)
self._cancel_event = cancel_event
self._aborted = False
@ -289,11 +291,17 @@ class InferenceBackend:
# Vision models carry the chat_template on the processor, not the inner
# tokenizer. Read markers from whichever has one, but resolve ids on the
# generation tokenizer, else the vision path misses the turn-end token.
template_source = container if getattr(container, "chat_template", None) else tokenizer
template_source = (
container if getattr(container, "chat_template", None) else tokenizer
)
try:
turn_end_ids = resolve_chat_turn_end_eos_ids_using(template_source, tokenizer)
turn_end_ids = resolve_chat_turn_end_eos_ids_using(
template_source, tokenizer
)
except Exception as e: # never block a load on eos resolution
logger.warning("Chat turn-end eos resolution failed for %s: %s", model_name, e)
logger.warning(
"Chat turn-end eos resolution failed for %s: %s", model_name, e
)
return
info["chat_turn_end_eos_ids"] = turn_end_ids
@ -376,7 +384,9 @@ 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":
@ -410,7 +420,9 @@ 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(
@ -521,7 +533,9 @@ 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.models[model_name]["context_length"] = runtime_context_length(
self.models[model_name].get("model"),
max_seq_length,
@ -534,7 +548,9 @@ 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}")
@ -558,10 +574,13 @@ class InferenceBackend:
from transformers import ProcessorMixin
if not (
isinstance(processor, ProcessorMixin) or hasattr(processor, "image_processor")
isinstance(processor, ProcessorMixin)
or hasattr(processor, "image_processor")
):
# LoRA adapters: use base model. Local merged exports: read base from export_metadata.json.
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:
@ -582,7 +601,9 @@ 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
@ -605,7 +626,9 @@ 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"
)
self.models[model_name]["context_length"] = runtime_context_length(
self.models[model_name].get("model"),
max_seq_length,
@ -654,7 +677,11 @@ 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.")
@ -721,9 +748,13 @@ class InferenceBackend:
base_model_name = lora_config.base_model
# 1. Load the base model if 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,
@ -759,7 +790,9 @@ 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:
"""Load an adapter onto the model only if not already attached."""
model = self.models[base_model_name].get("model")
@ -830,12 +863,16 @@ 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")
@ -843,11 +880,15 @@ class InferenceBackend:
elif isinstance(use_adapter, str):
# Enable adapters and set the named 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,
@ -1030,7 +1071,8 @@ 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(
@ -1087,9 +1129,13 @@ class InferenceBackend:
getattr(_gen_tok, "tokenizer", _gen_tok),
)
existing = model_info.get("chat_turn_end_eos_ids") or []
model_info["chat_turn_end_eos_ids"] = sorted(set(existing) | set(refreshed))
model_info["chat_turn_end_eos_ids"] = sorted(
set(existing) | set(refreshed)
)
except Exception as e:
logger.warning(f"Could not refresh chat turn-end eos after template: {e}")
logger.warning(
f"Could not refresh chat turn-end eos after template: {e}"
)
else:
logger.info(
f"No registered Unsloth template for {self.active_model_name}, using tokenizer default"
@ -1099,7 +1145,9 @@ 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
reasoning_channel_markers_resolved = False
@ -1248,7 +1296,9 @@ class InferenceBackend:
else:
# Text-only path for a 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
)
prompt_text = formatted_prompt
# Stream with TextIteratorStreamer + background thread
@ -1288,7 +1338,9 @@ class InferenceBackend:
min_p = min_p,
)
# Presence penalty (GGUF parity) for VLM chat.
_vision_input_ids = inputs.get("input_ids") if hasattr(inputs, "get") else None
_vision_input_ids = (
inputs.get("input_ids") if hasattr(inputs, "get") else None
)
if _vision_input_ids is not None:
_pp = _make_presence_penalty_processor(
presence_penalty, int(_vision_input_ids.shape[1])
@ -1298,7 +1350,9 @@ class InferenceBackend:
stopping_criteria = self._cancel_stopping_criteria(cancel_event)
if stopping_criteria is not None:
generation_kwargs["stopping_criteria"] = stopping_criteria
active_stop_token_ids = self._generation_stop_token_ids(model, generation_kwargs)
active_stop_token_ids = self._generation_stop_token_ids(
model, generation_kwargs
)
err: dict[str, str] = {}
@ -1673,7 +1727,9 @@ class InferenceBackend:
think_prefix = (
""
if self._is_gpt_oss_model()
else detect_think_prefill(prompt, getattr(tokenizer, "all_special_tokens", None))
else detect_think_prefill(
prompt, getattr(tokenizer, "all_special_tokens", None)
)
)
streamer = self._make_text_streamer(
@ -1698,12 +1754,15 @@ class InferenceBackend:
repetition_penalty = repetition_penalty,
do_sample = temperature > 0,
# Resolved once at load (chat_template-derived turn-end tokens).
eos_token_id = model_info.get("chat_turn_end_eos_ids") or tokenizer.eos_token_id,
eos_token_id = model_info.get("chat_turn_end_eos_ids")
or tokenizer.eos_token_id,
pad_token_id = tokenizer.eos_token_id
if tokenizer.pad_token_id is None
else tokenizer.pad_token_id,
)
active_stop_token_ids = self._generation_stop_token_ids(model, generation_kwargs)
active_stop_token_ids = self._generation_stop_token_ids(
model, generation_kwargs
)
# Presence penalty (GGUF parity); prompt_len excludes prompt tokens.
_pp = _make_presence_penalty_processor(
presence_penalty, int(inputs["input_ids"].shape[1])
@ -1794,7 +1853,9 @@ class InferenceBackend:
join_timeout = max(0, cancel_deadline - time.monotonic())
thread.join(timeout = join_timeout)
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"):
raise _GenerationThreadError(err["msg"])
@ -1871,12 +1932,21 @@ 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)
@ -1900,12 +1970,20 @@ 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,
@ -1976,7 +2054,9 @@ 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()
@ -2001,8 +2081,12 @@ 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,
@ -2044,7 +2128,9 @@ 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)
@ -2061,7 +2147,9 @@ 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":
@ -2077,7 +2165,9 @@ 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:")
@ -2092,7 +2182,10 @@ 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"
)
@ -2103,7 +2196,9 @@ 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,
@ -2116,7 +2211,9 @@ 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 when the tokenizer template fails.
Args:
@ -2146,7 +2243,9 @@ class InferenceBackend:
for msg in messages:
role = msg["role"]
content = content_to_text(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
@ -2175,8 +2274,13 @@ class InferenceBackend:
formatted += f"[INST] {user_content} [/INST]"
if i + 1 < len(conversation) and conversation[i + 1]["role"] == "assistant":
formatted += f" {content_to_text(conversation[i + 1]['content'])}</s>"
if (
i + 1 < len(conversation)
and conversation[i + 1]["role"] == "assistant"
):
formatted += (
f" {content_to_text(conversation[i + 1]['content'])}</s>"
)
i += 2
else:
formatted += " "
@ -2362,7 +2466,9 @@ 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"]
@ -2380,7 +2486,9 @@ class InferenceBackend:
# 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"
)
@ -2388,13 +2496,17 @@ class InferenceBackend:
# 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:
@ -2403,7 +2515,10 @@ 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"
@ -2430,7 +2545,9 @@ 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}")
@ -2442,7 +2559,9 @@ 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]:
"""Currently active model name."""

12
studio/backend/core/inference/llama_admission.py
View file

@ -81,7 +81,9 @@ def _bool_env(name: str, default: bool) -> bool:
return default
def _optional_positive_float_env(name: str, default: Optional[float]) -> Optional[float]:
def _optional_positive_float_env(
name: str, default: Optional[float]
) -> Optional[float]:
value = os.environ.get(name)
if value is None or not value.strip():
return default
@ -236,7 +238,9 @@ class LlamaAdmissionQueue:
self._capacity = 1
self._waiters: Deque[_Waiter] = deque()
def reserve(self, *, capacity: int, config: LlamaAdmissionConfig) -> LlamaAdmissionReservation:
def reserve(
self, *, capacity: int, config: LlamaAdmissionConfig
) -> LlamaAdmissionReservation:
capacity = max(1, int(capacity or 1))
if not config.enabled:
return LlamaAdmissionReservation(
@ -332,7 +336,9 @@ class LlamaAdmissionQueue:
def _prune_waiters_locked(self) -> None:
self._waiters = deque(
waiter for waiter in self._waiters if not waiter.cancelled and not waiter.future.done()
waiter
for waiter in self._waiters
if not waiter.cancelled and not waiter.future.done()
)
def _snapshot_locked(self) -> LlamaAdmissionSnapshot:

1088
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_keepwarm.py
View file

@ -114,7 +114,11 @@ def _note_untracked_end() -> None:
def _is_idle(ttl_seconds: float) -> bool:
with _lock:
return _inflight == 0 and _pending == 0 and (time.monotonic() - _last_active) >= ttl_seconds
return (
_inflight == 0
and _pending == 0
and (time.monotonic() - _last_active) >= ttl_seconds
)
def _note_activity() -> None:
@ -240,16 +244,26 @@ def restore_kv_resume(backend, manifest) -> None:
gguf = manifest.get("gguf")
binary = manifest.get("binary")
current = getattr(backend, "_gguf_path", None)
same_gguf = bool(gguf and current) and Path(current).resolve() == Path(gguf).resolve()
same_gguf = (
bool(gguf and current) and Path(current).resolve() == Path(gguf).resolve()
)
if same_gguf:
# Same path is not enough: shards may have been rewritten meanwhile.
identity = getattr(backend, "_gguf_file_identity", None)
same_gguf = callable(identity) and identity(current) == manifest.get("gguf_stat")
same_gguf = callable(identity) and identity(current) == manifest.get(
"gguf_stat"
)
if same_gguf:
# Nor the same file: launch overrides can invalidate KV numerics.
fingerprint = getattr(backend, "_slot_launch_fingerprint", None)
same_gguf = callable(fingerprint) and manifest.get("launch") == fingerprint()
if same_gguf and binary and binary == getattr(backend, "_slot_save_binary", None):
same_gguf = (
callable(fingerprint) and manifest.get("launch") == fingerprint()
)
if (
same_gguf
and binary
and binary == getattr(backend, "_slot_save_binary", None)
):
logger.info("Restoring saved slot KV onto the reloaded model")
backend.restore_slots_for_resume(manifest)
except Exception as exc:
@ -341,7 +355,9 @@ def _loaded_identity(backend):
# Third slot is the advertised id (repo id) an auto-switch load sets on the
# backend; it's the override key, so an idle stash keyed by the concrete load
# path doesn't drop the user's saved launch flags on the alias reload.
advertised = getattr(backend, "_openai_advertised_id", None) or backend.model_identifier
advertised = (
getattr(backend, "_openai_advertised_id", None) or backend.model_identifier
)
return (backend.model_identifier, getattr(backend, "hf_variant", None), advertised)
@ -403,7 +419,9 @@ async def idle_unload_loop(poll_seconds: float = 15.0) -> None:
_set_last_unloaded(freed) # let an alias request reload it
if manifest and freed:
_set_kv_resume({"identity": freed, **manifest})
logger.info("Idle auto-unload: saved slot KV for restore on reload")
logger.info(
"Idle auto-unload: saved slot KV for restore on reload"
)
elif manifest:
_delete_resume_files(manifest)
logger.info("Idle auto-unload: freed GGUF after %ss idle", ttl)

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

@ -196,11 +196,17 @@ _SPLIT_SHADOWING_FLAGS: frozenset[str] = _SPLIT_MODE_FLAGS | _TENSOR_SPLIT_FLAGS
_LAYER_OFFLOAD_FLAGS: frozenset[str] = frozenset(
{"-ngl", "--gpu-layers", "--n-gpu-layers", "-fit", "--fit"}
)
_MOE_OFFLOAD_FLAGS: frozenset[str] = frozenset({"-ncmoe", "--n-cpu-moe", "-cmoe", "--cpu-moe"})
_MOE_OFFLOAD_FLAGS: frozenset[str] = frozenset(
{"-ncmoe", "--n-cpu-moe", "-cmoe", "--cpu-moe"}
)
_OFFLOAD_SHADOWING_FLAGS: frozenset[str] = _LAYER_OFFLOAD_FLAGS | _MOE_OFFLOAD_FLAGS
_SHADOWING_FLAGS: frozenset[str] = (
_CONTEXT_FLAGS | _CACHE_FLAGS | _SPEC_FLAGS | _TEMPLATE_FLAGS | _SPLIT_SHADOWING_FLAGS
_CONTEXT_FLAGS
| _CACHE_FLAGS
| _SPEC_FLAGS
| _TEMPLATE_FLAGS
| _SPLIT_SHADOWING_FLAGS
)
# Shadowing flags that take no value -- strip the flag only, not the next token.
@ -233,16 +239,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
@ -258,7 +270,9 @@ def resolve_requested_ctx(args: Optional[Iterable[str]], fallback_n_ctx: int) ->
return override if override is not None else fallback_n_ctx
def _last_flag_value(args: Optional[Iterable[str]], flags: frozenset[str]) -> Optional[str]:
def _last_flag_value(
args: Optional[Iterable[str]], flags: frozenset[str]
) -> Optional[str]:
"""Return the last-wins string value among ``flags`` in extras, or None.
Handles both ``--flag=value`` and ``--flag value`` forms and raises if a
@ -341,7 +355,9 @@ def parse_split_mode_override(args: Optional[Iterable[str]]) -> Optional[str]:
return _last_flag_value(args, _SPLIT_MODE_FLAGS)
def resolve_tensor_parallel(args: Optional[Iterable[str]], fallback_tensor_parallel: bool) -> bool:
def resolve_tensor_parallel(
args: Optional[Iterable[str]], fallback_tensor_parallel: bool
) -> bool:
"""Return the tensor-parallel state load_model should treat as requested.
A user-supplied ``--split-mode`` in extras last-wins-overrides the

12
studio/backend/core/inference/llama_stats.py
View file

@ -41,7 +41,9 @@ class LlamaServerStatsLogger:
def start(self):
if self._thread is None:
self._thread = threading.Thread(target = self._run, name = "llama-stats", daemon = True)
self._thread = threading.Thread(
target = self._run, name = "llama-stats", daemon = True
)
self._thread.start()
def stop(self):
@ -71,7 +73,9 @@ class LlamaServerStatsLogger:
if not m:
misses += 1
if misses == 3: # transient stall (load/GC); keep polling.
self._log.debug("engine_stats: /metrics scrape failing, still retrying")
self._log.debug(
"engine_stats: /metrics scrape failing, still retrying"
)
continue # real shutdown is driven by stop() from _kill_process
misses = 0
# Generation tokens come from tokens_predicted_total (counter) and
@ -107,7 +111,9 @@ class LlamaServerStatsLogger:
def maybe_start_stats_logger(base_url, logger):
"""Start a stats logger unless UNSLOTH_STUDIO_ENGINE_STATS disables it."""
if (os.environ.get("UNSLOTH_STUDIO_ENGINE_STATS", "1") or "").strip().lower() in _OFF:
if (
os.environ.get("UNSLOTH_STUDIO_ENGINE_STATS", "1") or ""
).strip().lower() in _OFF:
return None
try:
interval = float(os.environ.get("UNSLOTH_STUDIO_ENGINE_STATS_INTERVAL_S", "10"))

26
studio/backend/core/inference/local_model_resolver.py
View file

@ -145,7 +145,11 @@ def _build_index() -> dict[str, _LocalGgufEntry]:
_resolve_hf_cache_dir,
_is_hidden_model,
)
from utils.paths import legacy_hf_cache_dir, hf_default_cache_dir, lmstudio_model_dirs
from utils.paths import (
legacy_hf_cache_dir,
hf_default_cache_dir,
lmstudio_model_dirs,
)
index: dict[str, _LocalGgufEntry] = {}
seen_hf: set[str] = set()
@ -162,7 +166,9 @@ def _build_index() -> dict[str, _LocalGgufEntry]:
return []
seen_hf.add(rp)
return _scan_hf_cache(directory)
except Exception as exc: # a missing/malformed root must skip, never crash the index
except (
Exception
) as exc: # a missing/malformed root must skip, never crash the index
logger.debug("auto-switch: skipping HF cache dir %r: %s", directory, exc)
return []
@ -174,7 +180,11 @@ def _build_index() -> dict[str, _LocalGgufEntry]:
except Exception as exc:
logger.debug("auto-switch: ./models scan failed: %s", exc)
try:
for hf_dir in (_resolve_hf_cache_dir(), legacy_hf_cache_dir(), hf_default_cache_dir()):
for hf_dir in (
_resolve_hf_cache_dir(),
legacy_hf_cache_dir(),
hf_default_cache_dir(),
):
found += _scan_hf_once(hf_dir)
except Exception as exc:
logger.debug("auto-switch: HF cache scan failed: %s", exc)
@ -189,7 +199,9 @@ def _build_index() -> dict[str, _LocalGgufEntry]:
try:
fp = Path(folder["path"])
found += (
_scan_models_dir(fp, limit = 200) + _scan_hf_once(fp) + _scan_lmstudio_dir(fp)
_scan_models_dir(fp, limit = 200)
+ _scan_hf_once(fp)
+ _scan_lmstudio_dir(fp)
)
except Exception as exc:
logger.debug("auto-switch: scan folder %r failed: %s", folder, exc)
@ -214,7 +226,11 @@ def _build_index() -> dict[str, _LocalGgufEntry]:
continue
# Index every alias (including the path) so a client can resolve by any of
# them, even though only the non-path loader_id is advertised.
for key in (raw_id, getattr(info, "model_id", None), getattr(info, "display_name", None)):
for key in (
raw_id,
getattr(info, "model_id", None),
getattr(info, "display_name", None),
):
if key:
index.setdefault(key.strip().lower(), entry)
return index

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

@ -248,7 +248,9 @@ def _client(
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))
@ -266,7 +268,9 @@ _STDIO_WEDGE_MARGIN = 15.0
# the scope includes a caller-supplied thread_id, so an unbounded cache is a
# resource-exhaustion surface. Overridable via env for large deployments.
try:
_STDIO_MAX_SESSIONS = max(1, int(os.environ.get("UNSLOTH_STUDIO_MAX_STDIO_MCP_SESSIONS", "32")))
_STDIO_MAX_SESSIONS = max(
1, int(os.environ.get("UNSLOTH_STUDIO_MAX_STDIO_MCP_SESSIONS", "32"))
)
except ValueError:
_STDIO_MAX_SESSIONS = 32
@ -401,7 +405,11 @@ class _StdioSession:
# wedged loop. No deadline at all when the caller set none -- but poll
# so a session closed under us (server update/delete) can't hang the
# request thread forever on a stopped loop.
deadline = None if timeout is None else time.monotonic() + timeout + _STDIO_WEDGE_MARGIN
deadline = (
None
if timeout is None
else time.monotonic() + timeout + _STDIO_WEDGE_MARGIN
)
try:
while True:
try:
@ -445,7 +453,9 @@ class _StdioSession:
task.cancel()
try:
asyncio.run_coroutine_threadsafe(_shutdown(), loop).result(_STDIO_CLOSE_TIMEOUT)
asyncio.run_coroutine_threadsafe(_shutdown(), loop).result(
_STDIO_CLOSE_TIMEOUT
)
except Exception as exc: # noqa: BLE001
logger.warning(
"MCP stdio session close failed for %s: %s",
@ -547,7 +557,12 @@ def _return_stdio_key_lock(key: tuple, key_lock: _StdioKeyLock) -> None:
def _get_stdio_session(
url: str, headers: Optional[dict], scope: Optional[str], deadline, cancel_event, config_check
url: str,
headers: Optional[dict],
scope: Optional[str],
deadline,
cancel_event,
config_check,
) -> _StdioSession:
"""``deadline`` is the caller's absolute monotonic budget (None = no limit):
the key-lock wait and the connect share it, so a slow startup can't stack
@ -602,10 +617,14 @@ def _get_stdio_session(
current = False
if not current:
session.close()
raise RuntimeError("MCP server was updated or removed while connecting")
raise RuntimeError(
"MCP server was updated or removed while connecting"
)
evicted: list = []
with _stdio_sessions_lock:
closed_while_connecting = _stdio_close_generation(url, headers) != generation
closed_while_connecting = (
_stdio_close_generation(url, headers) != generation
)
if not closed_while_connecting:
session.in_flight = 1
evicted = _evict_stdio_lru_locked() # bound the cache (LRU idle)
@ -613,11 +632,15 @@ def _get_stdio_session(
if not _stdio_reaper_started:
_stdio_reaper_started = True
threading.Thread(
target = _stdio_session_reaper, name = "mcp-stdio-reaper", daemon = True
target = _stdio_session_reaper,
name = "mcp-stdio-reaper",
daemon = True,
).start()
atexit.register(close_stdio_sessions)
for victim in evicted:
logger.info("Evicting LRU idle stdio MCP session: %s", _stdio_log_id(victim.url))
logger.info(
"Evicting LRU idle stdio MCP session: %s", _stdio_log_id(victim.url)
)
victim.close()
if closed_while_connecting:
session.close()
@ -683,7 +706,9 @@ def _evict_stdio_lru_locked() -> list:
cache may transiently overshoot rather than kill an in-flight call."""
victims: list = []
while len(_stdio_sessions) >= _STDIO_MAX_SESSIONS:
idle = [(s.last_used, k) for k, s in _stdio_sessions.items() if s.in_flight == 0]
idle = [
(s.last_used, k) for k, s in _stdio_sessions.items() if s.in_flight == 0
]
if not idle:
break
_, oldest = min(idle, key = lambda item: item[0])
@ -730,7 +755,8 @@ def _reap_idle_stdio_sessions(now: Optional[float] = None) -> None:
expired = [
key
for key, session in _stdio_sessions.items()
if session.in_flight == 0 and now - session.last_used >= _STDIO_SESSION_IDLE_TTL
if session.in_flight == 0
and now - session.last_used >= _STDIO_SESSION_IDLE_TTL
]
sessions = [_stdio_sessions.pop(key) for key in expired]
for key in expired:
@ -778,7 +804,9 @@ _probe_cooloff_until: dict[str, float] = {}
# endpoint/auth used to probe it (url, headers, oauth) or whether it's used at
# all (is_enabled). A rename does not. The update route's eviction and
# get_enabled_mcp_tools' mid-probe guard both key off this so they can't drift.
TOOL_CACHE_INVALIDATING_FIELDS = frozenset({"url", "headers_json", "use_oauth", "is_enabled"})
TOOL_CACHE_INVALIDATING_FIELDS = frozenset(
{"url", "headers_json", "use_oauth", "is_enabled"}
)
def get_cached_tools(server_id: str) -> Optional[list[dict]]:
@ -791,7 +819,11 @@ def cache_tools(server_id: str, tools: list[dict]) -> None:
def record_probe_failure(server_id: str, use_oauth: bool = False) -> None:
cooloff = OAUTH_FAILED_PROBE_COOLOFF_SECONDS if use_oauth else FAILED_PROBE_COOLOFF_SECONDS
cooloff = (
OAUTH_FAILED_PROBE_COOLOFF_SECONDS
if use_oauth
else FAILED_PROBE_COOLOFF_SECONDS
)
_probe_cooloff_until[server_id] = time.monotonic() + cooloff
@ -840,9 +872,13 @@ def _flatten_result(result: Any) -> str:
notes = []
if images:
n = len(images)
notes.append(f"{n} image{'s' if n > 1 else ''} attached; displayed to the user")
notes.append(
f"{n} image{'s' if n > 1 else ''} attached; displayed to the user"
)
if omitted:
notes.append(f"{omitted} image{'s' if omitted > 1 else ''} omitted (too large)")
notes.append(
f"{omitted} image{'s' if omitted > 1 else ''} omitted (too large)"
)
note = f"[{'; '.join(notes)}]"
body = f"{body}\n{note}" if body else note
@ -925,7 +961,9 @@ def _call_stdio_tool(
# attempt 0 may find the cached session stale/dead *before* dispatch and
# reconnect once (safe); attempt 1 is a freshly connected session.
for attempt in (0, 1):
session = _get_stdio_session(url, headers, scope, deadline, cancel_event, config_check)
session = _get_stdio_session(
url, headers, scope, deadline, cancel_event, config_check
)
try:
# Serialize calls per session: overlapping same-scope calls must
# not interleave operations on one stateful server (browser, REPL).
@ -971,7 +1009,9 @@ def _call_stdio_tool(
raise RuntimeError("MCP server connection is not available")
else:
rem = _remaining()
coro = _race_tool_call(session.client.call_tool(name, args), rem, cancel_event)
coro = _race_tool_call(
session.client.call_tool(name, args), rem, cancel_event
)
return session.run(coro, rem)
except (_MCPCancelled, asyncio.TimeoutError):
# _race_tool_call cancels the pending call but cancellation is

38
studio/backend/core/inference/mcp_config_import.py
View file

@ -60,14 +60,19 @@ def _enabled_from_spec(label: str, spec: dict) -> tuple[Optional[bool], Optional
return not disabled, None
def _parse_entry(name: str, spec: object) -> tuple[Optional[ParsedMcpEntry], Optional[str]]:
def _parse_entry(
name: str, spec: object
) -> tuple[Optional[ParsedMcpEntry], Optional[str]]:
label = str(name).strip()
if not label:
return None, "Server entry has an empty name."
if not isinstance(spec, dict):
return None, f"{label}: entry must be an object."
if _has_variable_reference(spec):
return None, f"{label}: VS Code variable references are not supported by import."
return (
None,
f"{label}: VS Code variable references are not supported by import.",
)
is_enabled, error = _enabled_from_spec(label, spec)
if error:
@ -91,8 +96,13 @@ def _parse_entry(name: str, spec: object) -> tuple[Optional[ParsedMcpEntry], Opt
if sandbox_enabled is not None and not isinstance(sandbox_enabled, bool):
return None, f"{label}: 'sandboxEnabled' must be true or false."
if sandbox_enabled:
return None, f"{label}: sandboxed stdio servers cannot be preserved by import."
unsupported = [field for field in _UNSUPPORTED_STDIO_FIELDS if spec.get(field) is not None]
return (
None,
f"{label}: sandboxed stdio servers cannot be preserved by import.",
)
unsupported = [
field for field in _UNSUPPORTED_STDIO_FIELDS if spec.get(field) is not None
]
if unsupported:
return None, f"{label}: import cannot preserve {', '.join(unsupported)}."
if spec.get("oauth") is not None:
@ -104,7 +114,10 @@ def _parse_entry(name: str, spec: object) -> tuple[Optional[ParsedMcpEntry], Opt
if env is not None and not isinstance(env, dict):
return None, f"{label}: 'env' must be an object."
if _has_null_value(env):
return None, f"{label}: null environment values are not supported by import."
return (
None,
f"{label}: null environment values are not supported by import.",
)
url = join_stdio_command([command, *(str(a) for a in args)])
headers = _coerce_str_dict(env) if env else None
return ParsedMcpEntry(label, url, headers, True, is_enabled = is_enabled), None
@ -120,12 +133,21 @@ def _parse_entry(name: str, spec: object) -> tuple[Optional[ParsedMcpEntry], Opt
field for field in _UNSUPPORTED_TIMEOUT_FIELDS if spec.get(field) is not None
]
if unsupported_timeout:
return None, f"{label}: import cannot preserve {', '.join(unsupported_timeout)}."
return (
None,
f"{label}: import cannot preserve {', '.join(unsupported_timeout)}.",
)
url_infers_sse = url.rstrip("/").endswith("/sse")
if entry_type == "sse" and not url_infers_sse:
return None, f"{label}: explicit SSE transport cannot be preserved for this URL."
return (
None,
f"{label}: explicit SSE transport cannot be preserved for this URL.",
)
if entry_type in _HTTP_REMOTE_TYPES and url_infers_sse:
return None, f"{label}: explicit HTTP transport cannot be preserved for this URL."
return (
None,
f"{label}: explicit HTTP transport cannot be preserved for this URL.",
)
oauth_raw = spec.get("oauth")
if oauth_raw is not None and not isinstance(oauth_raw, dict):
return None, f"{label}: 'oauth' must be an object."

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

@ -50,12 +50,16 @@ def _temporary_mlx_adapter_state(model, use_adapter):
"the loaded adapter or False for the base model."
)
if use_adapter is not True and use_adapter is not False:
raise TypeError("Unsloth MLX: use_adapter must be None, True, False, or a string.")
raise TypeError(
"Unsloth MLX: use_adapter must be None, True, False, or a string."
)
adapters, unsupported = _mlx_adapter_modules(model)
if use_adapter is True:
if not adapters and not unsupported:
logger.warning("MLX adapter requested, but the active model has no adapter layers")
logger.warning(
"MLX adapter requested, but the active model has no adapter layers"
)
yield
return
if unsupported:
@ -83,7 +87,11 @@ def _mlx_vlm_model_config(model):
config / _config actually carries a model_type."""
def _model_type(cfg):
return cfg.get("model_type") if isinstance(cfg, dict) else getattr(cfg, "model_type", None)
return (
cfg.get("model_type")
if isinstance(cfg, dict)
else getattr(cfg, "model_type", None)
)
configs = [
cfg
@ -153,10 +161,13 @@ def _prompt_serializes_vlm_media(prompt, messages):
if isinstance(message, dict):
media_reprs.update(_vlm_media_reprs(message.get("content")))
text_content = [
content_to_text(message.get("content")) for message in messages if isinstance(message, dict)
content_to_text(message.get("content"))
for message in messages
if isinstance(message, dict)
]
return any(
prompt.count(media_repr) > sum(content.count(media_repr) for content in text_content)
prompt.count(media_repr)
> sum(content.count(media_repr) for content in text_content)
for media_repr in media_reprs
)
@ -218,7 +229,9 @@ def _mlx_distributed_rank_size(group = None):
if world_size < 1:
raise ValueError(f"Invalid MLX distributed world_size={world_size}.")
if rank < 0 or rank >= world_size:
raise ValueError(f"Invalid MLX distributed rank={rank} for world_size={world_size}.")
raise ValueError(
f"Invalid MLX distributed rank={rank} for world_size={world_size}."
)
return rank, world_size
@ -363,7 +376,9 @@ class MLXInferenceBackend:
self._hf_token = hf_token
model_name = config.identifier if hasattr(config, "identifier") else str(config)
is_vision = getattr(config, "is_vision", False)
distributed_rank, distributed_size = _mlx_distributed_rank_size(distributed_group)
distributed_rank, distributed_size = _mlx_distributed_rank_size(
distributed_group
)
is_distributed = distributed_group is not None and distributed_size > 1
self._distributed_group = distributed_group
self._distributed_rank = distributed_rank
@ -668,7 +683,9 @@ 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"
)
# Parity with the transformers backend: if the template dropped the
# requested tools, fall back to the native template so MLX text models
@ -716,7 +733,9 @@ class MLXInferenceBackend:
)
)
if presence_penalty:
logits_processors.append(_make_mlx_presence_penalty_processor(float(presence_penalty)))
logits_processors.append(
_make_mlx_presence_penalty_processor(float(presence_penalty))
)
if not logits_processors:
logits_processors = None
@ -738,7 +757,10 @@ class MLXInferenceBackend:
type(self._model).__name__,
type(self._tokenizer).__name__,
)
with self._generation_lock, _temporary_mlx_adapter_state(self._model, _adapter_state):
with (
self._generation_lock,
_temporary_mlx_adapter_state(self._model, _adapter_state),
):
final_response = None
try:
# Enter request-scoped model state before yielding any response.
@ -858,7 +880,9 @@ class MLXInferenceBackend:
raise
prompt_error = exc
prompt_issue = (
_vlm_prompt_issue(prompt, messages) if prompt_error is None else "a rendering error"
_vlm_prompt_issue(prompt, messages)
if prompt_error is None
else "a rendering error"
)
if prompt_issue and has_tool_history:
raise RuntimeError(
@ -908,12 +932,16 @@ class MLXInferenceBackend:
)
prompt = recovered_prompt
elif prompt_issue:
raise RuntimeError(f"VLM chat template returned {prompt_issue}.") from prompt_error
raise RuntimeError(
f"VLM chat template returned {prompt_issue}."
) from prompt_error
from core.inference.chat_template_helpers import detect_think_prefill
# Re-emit an open <think> prefill from the prompt (see _generate_text).
cumulative = detect_think_prefill(prompt, getattr(chat_target, "all_special_tokens", None))
cumulative = detect_think_prefill(
prompt, getattr(chat_target, "all_special_tokens", None)
)
logger.info(
"VLM generating: prompt_len=%d, has_image=%s",
len(prompt),
@ -929,7 +957,9 @@ class MLXInferenceBackend:
top_k = int(top_k or 0),
min_p = float(min_p or 0.0),
)
_rep_active = repetition_penalty is not None and float(repetition_penalty) not in (
_rep_active = repetition_penalty is not None and float(
repetition_penalty
) not in (
0.0,
1.0,
)
@ -943,7 +973,9 @@ class MLXInferenceBackend:
_vlm_processors.extend(
make_logits_processors(repetition_penalty = float(repetition_penalty))
)
_vlm_processors.append(_make_mlx_presence_penalty_processor(float(presence_penalty)))
_vlm_processors.append(
_make_mlx_presence_penalty_processor(float(presence_penalty))
)
vlm_kwargs["logits_processors"] = _vlm_processors
elif _rep_active:
vlm_kwargs["repetition_penalty"] = float(repetition_penalty)
@ -953,7 +985,10 @@ class MLXInferenceBackend:
# Hold the generation lock AND the request-scoped adapter state for the
# whole stream so Base-vs-LoRA compare mode honors use_adapter and the
# wrapper tree is restored on completion, cancellation, or close.
with self._generation_lock, _temporary_mlx_adapter_state(self._model, _adapter_state):
with (
self._generation_lock,
_temporary_mlx_adapter_state(self._model, _adapter_state),
):
final_response = None
try:
# Emit any prefilled <think> block before the first token so the
@ -970,7 +1005,11 @@ 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():

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

@ -137,7 +137,9 @@ class InferenceOrchestrator:
atexit.register(self._cleanup)
logger.info("InferenceOrchestrator initialized (subprocess mode)")
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)
@ -175,12 +177,14 @@ class InferenceOrchestrator:
if resp.status_code == 200:
models = resp.json()
# Top 40 GGUFs (deep pool for frontend infinite scroll)
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
@ -361,7 +365,8 @@ class InferenceOrchestrator:
"Try a smaller model, lower context length, or close other GPU-heavy apps."
)
return (
f"{message}{suffix} " f"Details: pid={pid}, signal={sig_name}, exitcode={exitcode}."
f"{message}{suffix} "
f"Details: pid={pid}, signal={sig_name}, exitcode={exitcode}."
)
return f"{message} Details: pid={pid}, exitcode={exitcode}."
@ -444,7 +449,8 @@ 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:
@ -555,7 +561,10 @@ class InferenceOrchestrator:
initial_proc = self._proc
initial_resp_queue = self._resp_queue
while True:
if self._proc is not initial_proc or self._resp_queue is not initial_resp_queue:
if (
self._proc is not initial_proc
or self._resp_queue is not initial_resp_queue
):
yield GenStreamError(
f"Error: {self._subprocess_crash_message(crash_context)}",
public = True,
@ -622,7 +631,10 @@ class InferenceOrchestrator:
# unload_model's _wait_response sees it -- hanging the unload 300s.
if self._unload_pending:
return False
if self._dispatcher_thread is not None and self._dispatcher_thread.is_alive():
if (
self._dispatcher_thread is not None
and self._dispatcher_thread.is_alive()
):
return False
self._dispatcher_stop.clear()
@ -690,7 +702,9 @@ class InferenceOrchestrator:
rtype,
)
except Exception:
logger.exception("Inference dispatcher: failed to route a response; continuing")
logger.exception(
"Inference dispatcher: failed to route a response; continuing"
)
continue
def _generate_dispatched(
@ -720,7 +734,9 @@ class InferenceOrchestrator:
GPU work stays serialized; this only avoids orchestrator lock contention.
"""
if not self._ensure_subprocess_alive():
yield GenStreamError("Error: Inference subprocess is not running", public = True)
yield GenStreamError(
"Error: Inference subprocess is not running", public = True
)
return
if not self.active_model_name:
@ -786,7 +802,8 @@ class InferenceOrchestrator:
# bail when the active model changed or the dispatcher died: a mailbox with no
# dispatcher to route gen_done/gen_error hangs the compare stream.
dispatcher_alive = (
self._dispatcher_thread is not None and self._dispatcher_thread.is_alive()
self._dispatcher_thread is not None
and self._dispatcher_thread.is_alive()
)
unloading = (
self._unload_pending
@ -797,7 +814,9 @@ class InferenceOrchestrator:
self._mailboxes[request_id] = mailbox
# When bailing without a mailbox, note whether any OTHER compare request still
# routes through the dispatcher; if none and this call started it, stop it below.
orphaned_dispatcher = unloading and not dispatcher_preexisting and not self._mailboxes
orphaned_dispatcher = (
unloading and not dispatcher_preexisting and not self._mailboxes
)
if unloading:
# A racing unload can pass its _wait_dispatcher_idle() while the dispatcher was
# stopped, then set _unload_pending. The one we just started would otherwise
@ -919,11 +938,15 @@ class InferenceOrchestrator:
self._send_cmd(cmd)
deadline = None if timeout is None else time.monotonic() + timeout
while deadline is None or time.monotonic() < deadline:
remaining = 1.0 if deadline is None else max(0.1, deadline - time.monotonic())
remaining = (
1.0 if deadline is None else max(0.1, deadline - time.monotonic())
)
resp = self._read_resp(timeout = min(remaining, 1.0))
if resp is None:
if not self._ensure_subprocess_alive():
raise RuntimeError(self._subprocess_crash_message("sharing chat turn"))
raise RuntimeError(
self._subprocess_crash_message("sharing chat turn")
)
continue
rtype = resp.get("type", "")
@ -1148,13 +1171,21 @@ class InferenceOrchestrator:
# 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:
# Worker reports failures (consent gate included) under "message".
error = resp.get("message") or resp.get("error") or "Failed to load model"
error = (
resp.get("message")
or resp.get("error")
or "Failed to load model"
)
self.loading_models.discard(model_name)
self.active_model_name = None
self.models.clear()
@ -1164,7 +1195,10 @@ class InferenceOrchestrator:
self.loading_models.discard(model_name)
from utils.transformers_version import SidecarSwapInProgress
if isinstance(exc, SidecarSwapInProgress) and self._ensure_subprocess_alive():
if (
isinstance(exc, SidecarSwapInProgress)
and self._ensure_subprocess_alive()
):
# Raised before the old worker was torn down: the previous model
# is still live, so keep the mirrors (clearing them would let the
# installer treat the worker as inactive and kill it unreported).
@ -1466,7 +1500,10 @@ class InferenceOrchestrator:
try:
close()
except Exception:
logger.debug("failed to close errored generation stream", exc_info = True)
logger.debug(
"failed to close errored generation stream",
exc_info = True,
)
initial = list(messages)
if system_prompt:
@ -1550,7 +1587,9 @@ class InferenceOrchestrator:
readers don't consume each other's tokens off the shared resp_queue.
"""
if not self._ensure_subprocess_alive():
yield GenStreamError("Error: Inference subprocess is not running", public = True)
yield GenStreamError(
"Error: Inference subprocess is not running", public = True
)
return
if not self.active_model_name:
@ -1677,7 +1716,9 @@ class InferenceOrchestrator:
if resp is None:
if not self._ensure_subprocess_alive():
raise RuntimeError(self._subprocess_crash_message("audio generation"))
raise RuntimeError(
self._subprocess_crash_message("audio generation")
)
continue
rtype = resp.get("type", "")
@ -1756,7 +1797,9 @@ class InferenceOrchestrator:
) -> Generator[str, None, None]:
"""Shared inner logic for audio input generation (Whisper + ASR)."""
if not self._ensure_subprocess_alive():
yield GenStreamError("Error: Inference subprocess is not running", public = True)
yield GenStreamError(
"Error: Inference subprocess is not running", public = True
)
return
if not self.active_model_name:
yield GenStreamError("Error: No active model", public = True)
@ -1774,7 +1817,9 @@ class InferenceOrchestrator:
# numpy array -> 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 = {

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

@ -181,7 +181,9 @@ def _promote(
name = function.get("name") if isinstance(function, dict) else None
if name not in allowed_tools:
continue
arguments = _coerce_promoted_arguments(function.get("arguments"), name, tool_schemas)
arguments = _coerce_promoted_arguments(
function.get("arguments"), name, tool_schemas
)
if arguments is None:
continue
promoted.append(
@ -218,13 +220,17 @@ def heal_openai_message_events(
content = msg.get("content")
if not isinstance(content, str) or not _has_heal_signal(content):
return None
parsed, spans = parse_tool_calls_from_text(content, allow_incomplete = True, with_spans = True)
parsed, spans = parse_tool_calls_from_text(
content, allow_incomplete = True, with_spans = True
)
tool_schemas = _tool_schemas_by_name(tools) if tools is not None else None
events: list = []
pos = 0
call_count = 0
for call, (start, end) in zip(parsed, spans):
promoted = _promote([call], allowed_tools, id_offset = call_count, tool_schemas = tool_schemas)
promoted = _promote(
[call], allowed_tools, id_offset = call_count, tool_schemas = tool_schemas
)
if promoted:
if content[pos:start]:
events.append(("text", content[pos:start]))
@ -545,7 +551,10 @@ def nudge_messages(data: Any, allowed_tools: set) -> list:
is byte-identical and llama-server's slot/prefix cache is reused (same
shape as the enable-tools loop's reprompt).
"""
tool_hint = " or ".join(f"`{name}`" for name in sorted(allowed_tools)) or "an available tool"
tool_hint = (
" or ".join(f"`{name}`" for name in sorted(allowed_tools))
or "an available tool"
)
return [
{"role": "assistant", "content": _last_assistant_text(data)},
{

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

@ -103,7 +103,9 @@ 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 (per-bucket + total).
Unknown model -> ``priced`` False and USD fields 0.0 (token counts still report).
@ -131,7 +133,8 @@ def calculate_cost(provider: str, model: str, usage: dict[str, Any]) -> dict[str
# Clamp >=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))
# Fall back to mirrored prompt_tokens_details only when native
@ -214,10 +217,14 @@ def calculate_cost(provider: str, model: str, usage: dict[str, Any]) -> dict[str
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):
@ -234,7 +241,9 @@ def calculate_cost(provider: str, model: str, usage: dict[str, Any]) -> dict[str
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 items).
srv = usage.get("openai_tool_use") or {}

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

@ -341,7 +341,9 @@ def _reprompt_intent_text(text: str, *, reasoning_prefilled: bool = False) -> st
reasoning_text = "".join(reasoning).strip()
if visible_text:
return visible_text
return "\n".join(part for part in (prefilled_reasoning, reasoning_text) if part).strip()
return "\n".join(
part for part in (prefilled_reasoning, reasoning_text) if part
).strip()
def _looks_like_enabled_bare_json(text: str, enabled_tool_names: Optional[set]) -> bool:
@ -404,7 +406,9 @@ def _detect_render_html_tool_start(content: str) -> bool:
# first call through the parser (it reads top-level names).
arr_calls = parse_tool_calls_from_text(content[mt:])
if arr_calls:
candidates.append((mt, (arr_calls[0].get("function") or {}).get("name") or ""))
candidates.append(
(mt, (arr_calls[0].get("function") or {}).get("name") or "")
)
for rm in _REHEARSAL_RENDER_NAME_RE.finditer(content):
if not _in_think(rm.start(1)):
candidates.append((rm.start(1), rm.group(1)))
@ -532,7 +536,9 @@ def run_safetensors_tool_loop(
# off never prompts, so (like auto) it must not lose first-pass retrieval
# even if a direct caller passes a stale confirm_tool_calls flag.
_skip_autoinject = (
confirm_tool_calls and not bypass_permissions and permission_mode not in ("auto", "off")
confirm_tool_calls
and not bypass_permissions
and permission_mode not in ("auto", "off")
)
_auto = None if _skip_autoinject else build_rag_autoinject(conversation, rag_scope)
if _auto:
@ -570,7 +576,9 @@ def run_safetensors_tool_loop(
def _tool_succeeded(tool_name: str) -> bool:
key_prefix = f"{tool_name}:"
return any(
record.executed and not record.is_error and record.key.startswith(key_prefix)
record.executed
and not record.is_error
and record.key.startswith(key_prefix)
for record in tool_controller.history
)
@ -599,10 +607,14 @@ def run_safetensors_tool_loop(
final_attempt_done = True
active_tools = []
tool_protocol_active = not final_attempt_done and (unrestricted_tools or bool(active_tools))
tool_protocol_active = not final_attempt_done and (
unrestricted_tools or bool(active_tools)
)
tool_xml_signals = TOOL_XML_SIGNALS if tool_protocol_active else ()
# Gate the markerless bare-JSON form on enabled names so an ordinary JSON answer isn't misread as a call.
_enabled_tool_names = None if unrestricted_tools else set(_active_tool_names(active_tools))
_enabled_tool_names = (
None if unrestricted_tools else set(_active_tool_names(active_tools))
)
detect_state = _state_buffering
content_buffer = ""
@ -717,7 +729,10 @@ def run_safetensors_tool_loop(
# Earliest genuine boundary: bare [ARGS] in prose is skipped; a real NAME[ARGS] is
# pulled back to NAME so the name is not flushed.
signal_pos = _earliest_tool_signal(
candidate, tool_xml_signals, _detect_tools, unrestricted = unrestricted_tools
candidate,
tool_xml_signals,
_detect_tools,
unrestricted = unrestricted_tools,
)
if signal_pos >= 0:
before_tool = candidate[:signal_pos]
@ -818,7 +833,9 @@ def run_safetensors_tool_loop(
not is_match
and not is_prefix
and tool_protocol_active
and _is_rehearsal_prefix(stripped, _detect_tools, unrestricted = unrestricted_tools)
and _is_rehearsal_prefix(
stripped, _detect_tools, unrestricted = unrestricted_tools
)
):
is_prefix = True
is_rehearsal_prefix = True
@ -920,7 +937,9 @@ def run_safetensors_tool_loop(
"text": content_accum,
}
_live_args_streamed_upto = len(content_accum)
elif is_prefix and (is_rehearsal_prefix or len(stripped) < _MAX_BUFFER_CHARS):
elif is_prefix and (
is_rehearsal_prefix or len(stripped) < _MAX_BUFFER_CHARS
):
# A rehearsal prefix is self-bounded; the buffer cap must not cut long MCP names short.
continue
else:
@ -976,7 +995,9 @@ def run_safetensors_tool_loop(
if content_buffer:
cumulative_display += content_buffer
cleaned = strip_tool_markup(
cumulative_display, final = True, enabled_tool_names = _enabled_tool_names
cumulative_display,
final = True,
enabled_tool_names = _enabled_tool_names,
)
if len(cleaned) > len(last_emitted):
last_emitted = cleaned
@ -1020,7 +1041,10 @@ def run_safetensors_tool_loop(
len(intent_text),
)
conversation.append({"role": "assistant", "content": intent_text})
tool_hint = " or ".join(_active_tool_names(active_tools)) or "an available tool"
tool_hint = (
" or ".join(_active_tool_names(active_tools))
or "an available tool"
)
conversation.append(
{
"role": "user",
@ -1035,7 +1059,9 @@ def run_safetensors_tool_loop(
# Final answer. If a literal tool marker in prose was buffered but
# never parsed as a call, restore the raw text so the prose surfaces
# in full; route-level cleanup still applies the Auto-Heal policy.
if content_accum and any(sig in content_accum for sig in tool_xml_signals):
if content_accum and any(
sig in content_accum for sig in tool_xml_signals
):
yield {"type": "content", "text": content_accum}
else:
# Turn ended as a plain answer (no [ARGS] followed): the held rehearsal tail is real
@ -1087,7 +1113,9 @@ def run_safetensors_tool_loop(
# Drained bare-JSON call that didn't parse: with Auto-Heal on, drop the fragment
# (plain JSON answers are left untouched); off keeps it visible per the strict contract.
if tool_protocol_active and auto_heal_tool_calls:
_drain_text = strip_leading_bare_json_call(_drain_text, _enabled_tool_names)
_drain_text = strip_leading_bare_json_call(
_drain_text, _enabled_tool_names
)
if _drain_text:
yield {"type": "content", "text": _drain_text}
if provisional_render_html_started and not provisional_resolved:
@ -1112,7 +1140,9 @@ def run_safetensors_tool_loop(
next_call_id += len(tool_calls)
# Strip a leading bare-JSON call from the kept content so it isn't replayed as text or
# next-turn history (``_strip_tool_markup_final`` only knows XML). No-op for plain JSON answers.
content_text = strip_leading_bare_json_call(content_text, _enabled_tool_names)
content_text = strip_leading_bare_json_call(
content_text, _enabled_tool_names
)
if final_attempt_done:
# Final-answer turn re-called a tool -- stop the loop.
@ -1187,14 +1217,18 @@ def run_safetensors_tool_loop(
conversation.append(assistant_msg)
assistant_appended = True
else:
assistant_msg.setdefault("tool_calls", []).append(decision.as_assistant_tool_call())
assistant_msg.setdefault("tool_calls", []).append(
decision.as_assistant_tool_call()
)
# Bypass wins over the confirm gate at the loop level too, so a
# direct internal caller passing both flags never prompts. In
# "auto" mode only calls detected as potentially unsafe pause.
# "off" never prompts (sandbox stays on).
needs_confirm = (
bool(confirm_tool_calls) and not bypass_permissions and permission_mode != "off"
bool(confirm_tool_calls)
and not bypass_permissions
and permission_mode != "off"
)
if needs_confirm and permission_mode == "auto":
from core.inference.tools import is_potentially_unsafe_tool_call
@ -1202,7 +1236,9 @@ def run_safetensors_tool_loop(
decision.tool_name, decision.arguments
)
approval_id = new_approval_id() if needs_confirm else ""
decision_slot = begin_tool_decision(session_id, approval_id) if needs_confirm else None
decision_slot = (
begin_tool_decision(session_id, approval_id) if needs_confirm else None
)
start_event = decision.tool_start_event()
start_event["approval_id"] = approval_id
start_event["awaiting_confirmation"] = needs_confirm
@ -1268,7 +1304,9 @@ def run_safetensors_tool_loop(
)
if _accepts_output_callback(execute_tool):
kwargs["output_callback"] = _output_callback
return execute_tool(_decision.tool_name, _decision.arguments, **kwargs)
return execute_tool(
_decision.tool_name, _decision.arguments, **kwargs
)
try:
result = yield from stream_tool_execution(

4
studio/backend/core/inference/sandbox_site/sitecustomize.py
View file

@ -225,7 +225,9 @@ def _remap(path, notify = True):
for prefix in _PREFIXES + _CONDITIONAL_PREFIXES:
# Heal only while the real prefix directory is absent, so a genuine host
# mount / user directory at that prefix is never shadowed.
if (text == prefix or text.startswith(prefix + "/")) and not os.path.exists(prefix):
if (text == prefix or text.startswith(prefix + "/")) and not os.path.exists(
prefix
):
return _map_onto_cwd(prefix, text, notify = notify)
return path

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

@ -60,9 +60,7 @@ TOOL_XML_SIGNALS = (
# DeepSeek opener variants; shared by parse and strip so a parsed signal is always stripped.
_DEEPSEEK_OPEN_ALT = (
r"tool▁calls▁begin|tool_calls_begin|tool calls begin|tool\\_calls\\_begin|tool▁calls"
)
_DEEPSEEK_OPEN_ALT = r"tool▁calls▁begin|tool_calls_begin|tool calls begin|tool\\_calls\\_begin|tool▁calls"
_DEEPSEEK_OPEN_RE_SRC = r"<｜(?:" + _DEEPSEEK_OPEN_ALT + r")｜>"
# Closed pairs only (mid-stream); _TOOL_ALL_PATS also eats unclosed tails at
@ -84,7 +82,9 @@ _TOOL_CLOSED_PATS = [
# DeepSeek R1 / V3 / V3.1: full envelope (any opener variant) ... end.
re.compile(_DEEPSEEK_OPEN_RE_SRC + r".*?<｜tool▁calls▁end｜>", re.DOTALL),
# Kimi K2: ``<|tool_calls_section_begin|>...<|tool_calls_section_end|>``.
re.compile(r"<\|tool_calls_section_begin\|>.*?<\|tool_calls_section_end\|>", re.DOTALL),
re.compile(
r"<\|tool_calls_section_begin\|>.*?<\|tool_calls_section_end\|>", re.DOTALL
),
# Kimi K2 section-less closed call; else the catch-all below eats trailing prose to EOS.
re.compile(r"<\|tool_call_begin\|>.*?<\|tool_call_end\|>", re.DOTALL),
]
@ -187,7 +187,10 @@ REPROMPT_MAX_CHARS = 2000
def is_short_intent_without_action(text: str) -> bool:
stripped = text.strip()
return 0 < len(stripped) < REPROMPT_MAX_CHARS and INTENT_SIGNAL.search(stripped) is not None
return (
0 < len(stripped) < REPROMPT_MAX_CHARS
and INTENT_SIGNAL.search(stripped) is not None
)
def reprompt_to_act_message(tool_hint: str) -> str:
@ -264,7 +267,9 @@ _DEEPSEEK_R1_CLOSE_RE = re.compile(r"```[\s\r\n]*" + re.escape(_DEEPSEEK_CALL_EN
# direct ``<arg_key>``/``</tool_call>`` (4.7 drops the newline, zero-arg calls close at once).
# Name class ``[\w.\-]+`` keeps prose like ``<tool_call>not a call</tool_call>`` unparsed;
# ``{`` stays with the Qwen JSON parser.
_GLM_TC_OPEN_RE = re.compile(r"<tool_call>\s*([\w.\-]+)\s*(?=\n|<arg_key>|</tool_call>)")
_GLM_TC_OPEN_RE = re.compile(
r"<tool_call>\s*([\w.\-]+)\s*(?=\n|<arg_key>|</tool_call>)"
)
_GLM_TC_CLOSE = "</tool_call>"
_GLM_ARG_KEY_OPEN = "<arg_key>"
_GLM_ARG_KEY_CLOSE = "</arg_key>"
@ -432,7 +437,9 @@ def _strip_mistral_closed_calls(text: str) -> str:
return "".join(out)
def _strip_gemma_wrapperless_calls(text: str, enabled_tool_names: Optional[set] = None) -> str:
def _strip_gemma_wrapperless_calls(
text: str, enabled_tool_names: Optional[set] = None
) -> str:
"""Strip closed wrapper-less Gemma ``call:NAME{...}`` calls with balanced brace
scanning (nested arguments are removed whole). ``enabled_tool_names`` gates the
strip like the parser gate: a disabled/example name stays visible; ``None``
@ -450,7 +457,9 @@ def _strip_gemma_wrapperless_calls(text: str, enabled_tool_names: Optional[set]
if not m:
out.append(text[cursor:])
break
disabled = enabled_tool_names is not None and m.group(1) not in enabled_tool_names
disabled = (
enabled_tool_names is not None and m.group(1) not in enabled_tool_names
)
brace = m.end() - 1 # _GEMMA_BARE_TC_RE consumes through the opening ``{``
# Same boundary scanner as the parser: strip exactly what it consumed.
end = _gemma_body_brace_end(text, brace)
@ -475,7 +484,9 @@ _FUNC_CLOSE_TAG_RE = re.compile(r"</function>")
def _strip_function_xml_calls(text: str, *, final: bool) -> str:
"""Strip ``<function=...>`` calls by mirroring the parser: an opener inside an open ``<parameter>`` is data and each call closes at its first ``</function>`` that is not parameter data; ``final`` drops a trailing unclosed call."""
starts = [
m for m in _TC_FUNC_START_RE.finditer(text) if not _inside_open_parameter(text, m.start())
m
for m in _TC_FUNC_START_RE.finditer(text)
if not _inside_open_parameter(text, m.start())
]
if not starts:
return text
@ -532,7 +543,11 @@ def _glm_value_close(
j = ve + len(_GLM_ARG_VAL_CLOSE)
while j < n and text[j] in " \t\r\n":
j += 1
if j >= n or text.startswith(_GLM_ARG_KEY_OPEN, j) or text.startswith(_GLM_TC_CLOSE, j):
if (
j >= n
or text.startswith(_GLM_ARG_KEY_OPEN, j)
or text.startswith(_GLM_TC_CLOSE, j)
):
while qpos < ve:
ch = text[qpos]
if quote:
@ -541,7 +556,9 @@ def _glm_value_close(
continue
if ch == quote:
quote = ""
elif ch in "\"'" and (prev in ":{[(,=" or (ch == '"' and prev_raw.isspace())):
elif ch in "\"'" and (
prev in ":{[(,=" or (ch == '"' and prev_raw.isspace())
):
quote = ch
if not ch.isspace():
prev = ch
@ -629,7 +646,9 @@ def strip_tool_markup(
# Bare reasoning-rehearsal ``name[ARGS]{json}`` and the Mistral name form promote through
# the shared balanced scan, so strip them the same way (any nesting depth removed whole).
# The rehearsal arm is name-gated: an inactive ``foo[ARGS]{..}`` is prose and is kept.
seg = _tool_healing._strip_bracket_tag_calls(seg, enabled_tool_names = enabled_tool_names)
seg = _tool_healing._strip_bracket_tag_calls(
seg, enabled_tool_names = enabled_tool_names
)
if seg_final:
# Markerless Gemma ``call:NAME{...}`` (name-gated, mirrors the parse gate); end-of-turn only.
seg = _strip_gemma_wrapperless_calls(seg, enabled_tool_names)
@ -667,7 +686,11 @@ def has_tool_signal(text: str) -> bool:
# DeepSeek/Kimi markers must parse as the OUTER call. Detect it opening before the first
# marker so the pre-pass skips it.
_EMBEDDED_MARKER_RE = re.compile(
_DEEPSEEK_OPEN_RE_SRC + "|" + re.escape(_KIMI_SECTION_BEGIN) + "|" + re.escape(_KIMI_CALL_BEGIN)
_DEEPSEEK_OPEN_RE_SRC
+ "|"
+ re.escape(_KIMI_SECTION_BEGIN)
+ "|"
+ re.escape(_KIMI_CALL_BEGIN)
)
# Covers ``<function=NAME>`` and the attribute form. ``<|python_tag|>`` is Llama-3's
# envelope too (built-in ``NAME.call(`` and custom ``{json}``), so a quoted DeepSeek/Kimi
@ -686,7 +709,9 @@ _OUTER_ENVELOPE_CLOSED_PATS = (
)
def _marker_inside_leading_envelope(content: str, enabled_tool_names: Optional[set] = None) -> bool:
def _marker_inside_leading_envelope(
content: str, enabled_tool_names: Optional[set] = None
) -> bool:
first_marker = _EMBEDDED_MARKER_RE.search(content)
if first_marker is None:
return False
@ -700,7 +725,9 @@ def _marker_inside_leading_envelope(content: str, enabled_tool_names: Optional[s
end = _balanced_brace_end(content, i)
if end is not None and i < first_marker.start():
name = _top_level_bare_json_name(content[i : end + 1])
if name is not None and (enabled_tool_names is None or name in enabled_tool_names):
if name is not None and (
enabled_tool_names is None or name in enabled_tool_names
):
# The closed leading call owns the turn: a marker inside it is argument
# data, one after it a trailing example (same rule as the XML envelopes below).
return True
@ -980,7 +1007,9 @@ def parse_tool_calls_from_text(
while i < len(content) and content[i] in " \t\n\r":
i += 1
# The guard guarantees a balanced leading value (object or array).
end = (_balanced_brace_end if content[i] == "{" else _balanced_bracket_end)(content, i)
end = (_balanced_brace_end if content[i] == "{" else _balanced_bracket_end)(
content, i
)
return parse_tool_calls_from_text(
content[end + 1 :],
id_offset = id_offset,
@ -1047,7 +1076,9 @@ def parse_tool_calls_from_text(
]
pre_pass.sort(key = lambda pair: pair[0])
for _pos, parser in pre_pass:
calls = parser(content, id_offset = id_offset, allow_incomplete = allow_incomplete)
calls = parser(
content, id_offset = id_offset, allow_incomplete = allow_incomplete
)
if calls:
return calls
@ -1127,7 +1158,9 @@ def parse_tool_calls_from_text(
_parse_llama3_python_tag, # Llama-3 <|python_tag|>
_parse_mistral_tool_calls, # Mistral [TOOL_CALLS]
):
calls = parser(fallback_content, id_offset = id_offset, allow_incomplete = allow_incomplete)
calls = parser(
fallback_content, id_offset = id_offset, allow_incomplete = allow_incomplete
)
if calls:
return calls
@ -1164,7 +1197,9 @@ def _parse_tool_call_json(
# Strict mode: a balanced JSON body that never closed its ``<tool_call>``
# is a truncated call, not a finished one. Trailing prose after the close
# is still tolerated (matches the GGUF strict path).
if not allow_incomplete and not content[end + 1 :].lstrip().startswith("</tool_call>"):
if not allow_incomplete and not content[end + 1 :].lstrip().startswith(
"</tool_call>"
):
continue
try:
obj = json.loads(content[brace_start : end + 1])
@ -1258,7 +1293,9 @@ def _parse_function_xml(
# group(1) is ``<function=name>``, group(2) is ``<function name="...">``.
func_name = fm.group(1) or fm.group(2)
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)
)
# The call ends at the FIRST </function> / </tool_call> not inside an open
# parameter: a literal close in a code/search argument is skipped as data, and
# prose after the real close isn't folded into the last argument (mirrors
@ -1297,7 +1334,9 @@ def _parse_function_xml(
for pidx, pm in enumerate(param_starts):
val_start = pm.end()
next_param = (
param_starts[pidx + 1].start() if pidx + 1 < len(param_starts) else len(body)
param_starts[pidx + 1].start()
if pidx + 1 < len(param_starts)
else len(body)
)
raw_val = body[val_start:next_param]
if not _TC_PARAM_CLOSE_RE.search(raw_val):
@ -1483,7 +1522,11 @@ def _parse_llama3_python_tag(
cursor = brace + end_offset
continue
name = obj.get("name") or obj.get("function") or ""
args = obj.get("parameters") if "parameters" in obj else obj.get("arguments", {})
args = (
obj.get("parameters")
if "parameters" in obj
else obj.get("arguments", {})
)
# Skip rather than fabricate ``{"value": args}`` for a non-dict/non-string value.
if isinstance(args, dict):
args_str = json.dumps(args)
@ -1634,7 +1677,9 @@ def _parse_mistral_tool_calls(
return out
if content[k] == "[":
return _parse_mistral_array(content, k, id_offset, allow_incomplete = allow_incomplete)
return _parse_mistral_array(
content, k, id_offset, allow_incomplete = allow_incomplete
)
if content[k] == "{":
# Pre-v11 single ``{"name":...}``; fall through without a ``name`` so v11+ still runs.
@ -2024,7 +2069,9 @@ def _top_level_bare_json_name(probe: str) -> Optional[str]:
return function_value
def strip_leading_bare_json_call(text: str, enabled_tool_names: Optional[set] = None) -> str:
def strip_leading_bare_json_call(
text: str, enabled_tool_names: Optional[set] = None
) -> str:
"""Remove leading Llama-3.2 bare-JSON calls (including a ``;``-chained run)
that ``strip_tool_markup`` misses; non-call text is unchanged and
``enabled_tool_names`` gates like the parser. Consuming the whole chain
@ -2124,7 +2171,11 @@ def _gemma_parse_value(
close = text.find(_GEMMA_STR_END, i + len(_GEMMA_STR_BEGIN))
if close < 0:
return text[i + len(_GEMMA_STR_BEGIN) :], len(text), False
return text[i + len(_GEMMA_STR_BEGIN) : close], close + len(_GEMMA_STR_END), True
return (
text[i + len(_GEMMA_STR_BEGIN) : close],
close + len(_GEMMA_STR_END),
True,
)
if text[i] == "{":
return _gemma_parse_mapping(text, i)
if text[i] == "[":
@ -2279,7 +2330,9 @@ def _gemma_parse_stripped_body(body: str) -> dict[str, Any]:
continue
if ch == quote:
quote = ""
elif ch in "\"'" and (prev in ":{[(,=" or (ch == '"' and prev_raw.isspace())):
elif ch in "\"'" and (
prev in ":{[(,=" or (ch == '"' and prev_raw.isspace())
):
quote = ch
elif ch in "{[(":
depth += 1

8
studio/backend/core/inference/tool_loop_controller.py
View file

@ -323,7 +323,9 @@ class ToolLoopController:
self._restrict_to_allowed = tools is not None
self._tools = [copy.deepcopy(dict(tool)) for tool in (tools or [])]
self._allowed_tool_names = {
name for name in (_tool_name_from_schema(tool) for tool in self._tools) if name
name
for name in (_tool_name_from_schema(tool) for tool in self._tools)
if name
}
self._auto_heal_tool_calls = auto_heal_tool_calls
self._one_shot_tools = one_shot_tools
@ -400,7 +402,9 @@ class ToolLoopController:
noop_result = noop,
)
def record_result(self, decision: ToolCallDecision, result: Any) -> ToolCallCompletion:
def record_result(
self, decision: ToolCallDecision, result: Any
) -> ToolCallCompletion:
"""Record a real tool execution and return model/frontend payload helpers."""
result_text = result if isinstance(result, str) else str(result)
failed = is_tool_error(result_text)

8
studio/backend/core/inference/tool_stream_exec.py
View file

@ -70,7 +70,9 @@ TOOL_OUTPUT_STREAM_MAX_CHARS = 400_000
_STREAM_CAPPED_NOTICE = "\n... (further live output not streamed)\n"
def _drain_queue(q: "queue.Queue", sentinel: object, max_chars: int | None) -> tuple[str, bool]:
def _drain_queue(
q: "queue.Queue", sentinel: object, max_chars: int | None
) -> tuple[str, bool]:
"""Pull every currently-queued item, joining chunks in FIFO order.
With ``max_chars`` set, stop concatenating at the budget and discard the
@ -170,7 +172,9 @@ def stream_tool_execution(
# Heartbeats are paced by counting idle queue polls rather than a wall clock
# (tests patch ``time.monotonic`` globally, so the wrapper must not read it).
idle_polls_per_heartbeat = max(1, int(round(heartbeat_interval_s / poll_interval_s)))
idle_polls_per_heartbeat = max(
1, int(round(heartbeat_interval_s / poll_interval_s))
)
idle_polls = 0
streamed_chars = 0
stream_capped = False

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

@ -144,7 +144,9 @@ _BLOCKED_COMMANDS = (
)
_SHELL_SEPARATORS = frozenset({";", "&&", "||", "|", "&", "\n", "(", ")", "`", "{", "}"})
_SHELL_SEPARATORS = frozenset(
{";", "&&", "||", "|", "&", "\n", "(", ")", "`", "{", "}"}
)
# Bash keywords starting 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 the command Bash will exec.
@ -302,7 +304,9 @@ 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):
@ -327,7 +331,9 @@ def _find_blocked_commands(command: str) -> set[str]:
# Directory holding the sandbox ``sitecustomize.py`` shim (code-interpreter
# path remap); placed on the sandboxed child's PYTHONPATH in _build_safe_env.
_SANDBOX_SITE_DIR = os.path.join(os.path.dirname(os.path.abspath(__file__)), "sandbox_site")
_SANDBOX_SITE_DIR = os.path.join(
os.path.dirname(os.path.abspath(__file__)), "sandbox_site"
)
# ── "Approve for me" (permission_mode="auto") safety detection ──────────────
# Auto mode pauses only calls classified here as potentially unsafe. The sandbox
# and hard blocks (blocklist, rlimits) still apply at run time; this gate only
@ -422,7 +428,14 @@ _AUTO_UNSAFE_COMMAND_FLAGS = {
# --files0-from=F makes sort read the NUL-separated list of input files
# named in F, so a crafted list reads arbitrary host files indirectly.
"sort": frozenset(
{"-o", "--output", "--compress-program", "-T", "--temporary-directory", "--files0-from"}
{
"-o",
"--output",
"--compress-program",
"-T",
"--temporary-directory",
"--files0-from",
}
),
"tree": frozenset({"-o"}),
"xxd": frozenset({"-r"}),
@ -471,7 +484,9 @@ _AUTO_UNSAFE_COMMAND_FLAGS = {
),
# fd -x/--exec/-X/--exec-batch run a command per result;
# --base-directory/--search-path move the search root outside the workdir.
"fd": frozenset({"-x", "--exec", "-X", "--exec-batch", "--base-directory", "--search-path"}),
"fd": frozenset(
{"-x", "--exec", "-X", "--exec-batch", "--base-directory", "--search-path"}
),
# date -s/--set writes the clock; display forms (+FORMAT, -d/-u/-R/-r) read.
"date": frozenset({"-s", "--set"}),
# file -C/--compile writes a compiled .mgc magic database; ident forms read.
@ -485,7 +500,9 @@ _AUTO_UNSAFE_COMMAND_FLAGS = {
_AUTO_ARG_SENSITIVE_COMMANDS = frozenset({"hostname", "date"})
# date display flags taking a value token (-d STRING, -r FILE, -f FILE); the
# value is not a clock-setting positional, so it is skipped.
_DATE_DISPLAY_VALUE_FLAGS = frozenset({"-d", "--date", "-r", "--reference", "-f", "--file"})
_DATE_DISPLAY_VALUE_FLAGS = frozenset(
{"-d", "--date", "-r", "--reference", "-f", "--file"}
)
# Commands that write their 2nd positional (uniq [INPUT [OUTPUT]], xxd [infile
# [outfile]]): the 1st file reads to stdout, but a second file positional
# overwrites it, like `sort -o`.
@ -495,14 +512,29 @@ _AUTO_SECOND_POSITIONAL_WRITES = frozenset({"uniq", "xxd"})
# not miscounted as the output-file positional, and, conversely, a file that is
# literally named with digits (uniq 123 out) is still counted.
_SECOND_POSITIONAL_VALUE_FLAGS = {
"uniq": frozenset({"-f", "--skip-fields", "-s", "--skip-chars", "-w", "--check-chars"}),
"uniq": frozenset(
{"-f", "--skip-fields", "-s", "--skip-chars", "-w", "--check-chars"}
),
"xxd": frozenset(
{"-c", "--cols", "-s", "--seek", "-l", "--len", "-g", "--groupsize", "-o", "--offset"}
{
"-c",
"--cols",
"-s",
"--seek",
"-l",
"--len",
"-g",
"--groupsize",
"-o",
"--offset",
}
),
}
# find/fd group with (...) which resets command context, so scan every token for
# these once find/fd appears anywhere.
_AUTO_UNSAFE_FIND_LIKE_FLAGS = _AUTO_UNSAFE_COMMAND_FLAGS["find"] | _AUTO_UNSAFE_COMMAND_FLAGS["fd"]
_AUTO_UNSAFE_FIND_LIKE_FLAGS = (
_AUTO_UNSAFE_COMMAND_FLAGS["find"] | _AUTO_UNSAFE_COMMAND_FLAGS["fd"]
)
# Recursive readers with an absolute-path target escape the workdir onto host
# files (grep -R TOKEN /home, rg TOKEN /), so they ask.
_AUTO_RECURSIVE_SEARCH = frozenset({"grep", "egrep", "fgrep", "rg", "ug", "find", "fd"})
@ -762,7 +794,9 @@ _AUTO_UNSAFE_PY_WRITE_METHODS = frozenset(
# Archive / compressed-file constructors taking the mode as their 2nd arg like
# open: ZipFile(name, "w") / gzip.GzipFile(name, "w") write, so gated only in
# write mode (reading a .gz is fine, so the modules are not blanket-unsafe).
_ARCHIVE_CTOR_NAMES = frozenset({"ZipFile", "TarFile", "GzipFile", "BZ2File", "LZMAFile"})
_ARCHIVE_CTOR_NAMES = frozenset(
{"ZipFile", "TarFile", "GzipFile", "BZ2File", "LZMAFile"}
)
# The stdlib module each archive constructor is imported from.
_ARCHIVE_CTOR_MODULES = {
"zipfile": "ZipFile",
@ -1002,7 +1036,9 @@ def _expand_shell_assignments(command: str) -> str:
var, is_global, pat, rep = m.group(1), m.group(2), m.group(3), m.group(4)
if var not in env or not pat:
return m.group(0)
return env[var].replace(pat, rep) if is_global else env[var].replace(pat, rep, 1)
return (
env[var].replace(pat, rep) if is_global else env[var].replace(pat, rep, 1)
)
def repl_case(m):
var, op = m.group(1), m.group(2)
@ -1025,7 +1061,9 @@ def _expand_shell_assignments(command: str) -> str:
command = _SHELL_PARAM_INDIRECT_RE.sub(repl_indirect, command)
command = _SHELL_PARAM_REPL_RE.sub(repl_pattern, command)
command = _SHELL_PARAM_CASE_RE.sub(repl_case, command)
return _SHELL_VAR_RE.sub(lambda m: env.get(m.group(1) or m.group(2), m.group(0)), command)
return _SHELL_VAR_RE.sub(
lambda m: env.get(m.group(1) or m.group(2), m.group(0)), command
)
def _expand_param_defaults(command: str) -> str:
@ -1164,7 +1202,15 @@ _PATH_CTORS = (
# (os.path.abspath('/etc') -> /etc, Path('/etc').resolve() -> /etc), so folding
# through them keeps a sensitive root visible to the scan.
_PATH_PASSTHROUGH_ATTRS = frozenset(
{"abspath", "normpath", "realpath", "expanduser", "expandvars", "resolve", "absolute"}
{
"abspath",
"normpath",
"realpath",
"expanduser",
"expandvars",
"resolve",
"absolute",
}
)
# pathlib methods that rewrite only the final path component, so the sensitive
# target is never spelled out as a literal (Path('/etc/x').with_name('passwd')
@ -1267,12 +1313,18 @@ def _folded_path(
parts = [base if base is not None else "\x00"]
parts += [(fold(a) or "\x00") for a in node.args]
return "/".join(parts)
if isinstance(func, ast.Attribute) and func.attr in ("glob", "rglob", "iglob"):
if isinstance(func, ast.Attribute) and func.attr in (
"glob",
"rglob",
"iglob",
):
# Path('/etc').glob('passw?') -> the receiver dir joined with the
# glob pattern; _glob_token_sensitive then tests /etc/passw?.
base = fold(func.value)
pattern = fold(node.args[0]) if node.args else "\x00"
return (base if base is not None else "\x00") + "/" + (pattern or "\x00")
return (
(base if base is not None else "\x00") + "/" + (pattern or "\x00")
)
if isinstance(func, ast.Attribute) and func.attr in _PATH_NAME_REWRITES:
# Path('/etc/x').with_name('passwd') -> /etc/passwd; with_stem /
# with_suffix rewrite only the final component. Fold to the
@ -1377,7 +1429,9 @@ def _folded_is_sensitive(folded) -> bool:
# A dynamic segment (NUL) can be the "/" forming a sensitive root:
# open(os.sep + "etc/passwd") folds to "\x00etc/passwd", so re-scan with
# NUL as "/" (a benign "\x00data/file" -> "/data/file" stays safe).
or ("\x00" in folded and _references_sensitive_path(folded.replace("\x00", "/")))
or (
"\x00" in folded and _references_sensitive_path(folded.replace("\x00", "/"))
)
# A dynamic piece can also sit INSIDE a sensitive name: open('/et' +
# chr(99) + '/passwd') folds to "/et\x00/passwd", which none of the above
# catch. Match the literals around each NUL against a credential target,
@ -1409,10 +1463,14 @@ def _terminal_is_potentially_unsafe(command: str) -> bool:
candidates = []
for c in (command, stripped, _decode_ansi_c(command)):
c_param = _expand_param_defaults(c)
candidates.extend((c, c_param, _expand_braces(c_param), _expand_shell_assignments(c_param)))
candidates.extend(
(c, c_param, _expand_braces(c_param), _expand_shell_assignments(c_param))
)
# Run both the literal and glob-sensitive scans over every candidate, so a
# brace-expanded glob (cat /e{t,}c/pass?d -> /etc/pass?d) is caught.
if any(_glob_hits_sensitive(c) or _references_sensitive_path(c) for c in candidates):
if any(
_glob_hits_sensitive(c) or _references_sensitive_path(c) for c in candidates
):
return True
# Newlines (and CR) separate commands in a shell but read as plain
# whitespace to shlex, which would demote "ls\nrm x" to argument position.
@ -1429,7 +1487,9 @@ def _terminal_is_potentially_unsafe(command: str) -> bool:
expanded_command = _expand_shell_assignments(_expand_param_defaults(command))
if expanded_command != command:
try:
elexer = shlex.shlex(expanded_command, posix = True, punctuation_chars = ";&|()")
elexer = shlex.shlex(
expanded_command, posix = True, punctuation_chars = ";&|()"
)
elexer.whitespace_split = True
scan_tokens = list(elexer)
except ValueError:
@ -1438,7 +1498,10 @@ def _terminal_is_potentially_unsafe(command: str) -> bool:
scan_tokens = tokens
# find/fd group with (...) which resets command context, so a trailing
# -delete/-exec could slip past; scan every token when find/fd appears.
if any(os.path.basename(t.strip(";&|()`{}")).lower() in ("find", "fd") for t in scan_tokens):
if any(
os.path.basename(t.strip(";&|()`{}")).lower() in ("find", "fd")
for t in scan_tokens
):
if any(t.split("=", 1)[0] in _AUTO_UNSAFE_FIND_LIKE_FLAGS for t in scan_tokens):
return True
# A recursive reader rooted outside the sandbox reads host files (grep -R
@ -1449,7 +1512,10 @@ def _terminal_is_potentially_unsafe(command: str) -> bool:
# that already asks below.
if any(t.startswith("/") or t.startswith("~") for t in scan_tokens):
token_bases = [os.path.basename(t.strip(";&|()`{}")).lower() for t in tokens]
if any(b in _AUTO_RECURSIVE_SEARCH or b in _AUTO_RECURSIVE_LISTERS for b in token_bases):
if any(
b in _AUTO_RECURSIVE_SEARCH or b in _AUTO_RECURSIVE_LISTERS
for b in token_bases
):
return True
# ls only walks the whole subtree with -R/--recursive (ls -R /home,
# ls -laR /); a non-recursive ls /home lists one level and stays here.
@ -1490,7 +1556,9 @@ def _terminal_is_potentially_unsafe(command: str) -> bool:
# a "--x" prefix of an unsafe long flag fails closed.
is_long_abbrev = flag_head.startswith("--") and len(flag_head) > 2
for uf in _AUTO_UNSAFE_COMMAND_FLAGS.get(current_command, ()):
if flag_head == uf or (len(uf) == 2 and (token.startswith(uf) or uf[1] in cluster)):
if flag_head == uf or (
len(uf) == 2 and (token.startswith(uf) or uf[1] in cluster)
):
return True
if is_long_abbrev and uf.startswith("--") and uf.startswith(flag_head):
return True
@ -1523,7 +1591,9 @@ def _terminal_is_potentially_unsafe(command: str) -> bool:
elif current_command in _AUTO_ARG_SENSITIVE_COMMANDS:
if pending_flag_value:
pending_flag_value = False
elif raw_pos and not (current_command == "date" and raw_pos.startswith("+")):
elif raw_pos and not (
current_command == "date" and raw_pos.startswith("+")
):
return True
continue
if _ASSIGNMENT_RE.match(token):
@ -1654,7 +1724,10 @@ def _python_is_potentially_unsafe(code: str) -> bool:
first = call.args[0]
if not (isinstance(first, ast.Constant) and isinstance(first.value, str)):
return True
return first.value in _AUTO_UNSAFE_PY_ATTRS or first.value in _AUTO_UNSAFE_PY_WRITE_METHODS
return (
first.value in _AUTO_UNSAFE_PY_ATTRS
or first.value in _AUTO_UNSAFE_PY_WRITE_METHODS
)
def _fileinput_inplace(call) -> bool:
# fileinput.input(..., inplace=True) opens each file for in-place rewrite.
@ -1705,7 +1778,9 @@ def _python_is_potentially_unsafe(code: str) -> bool:
# merely passed or printed (print(getattr(o, 'name'))).
if isinstance(arg, ast.Name):
return (
arg.id in open_aliases or arg.id in writer_aliases or arg.id in archive_ctor_aliases
arg.id in open_aliases
or arg.id in writer_aliases
or arg.id in archive_ctor_aliases
)
if isinstance(arg, ast.Attribute):
return (
@ -1800,7 +1875,9 @@ def _python_is_potentially_unsafe(code: str) -> bool:
else:
assign_targets = node.targets
targets = [t.id for t in assign_targets if isinstance(t, ast.Name)]
attr_targets = [t.attr for t in assign_targets if isinstance(t, ast.Attribute)]
attr_targets = [
t.attr for t in assign_targets if isinstance(t, ast.Attribute)
]
if isinstance(value, ast.Name) and value.id in open_aliases:
open_aliases.update(targets)
attr_open_aliases.update(attr_targets) # box.f = open
@ -1836,7 +1913,10 @@ def _python_is_potentially_unsafe(code: str) -> bool:
and value.value.id in builtins_aliases
):
code_exec_aliases.update(targets) # e = builtins.eval
elif isinstance(value, ast.Attribute) and value.attr in _AUTO_UNSAFE_PY_WRITE_METHODS:
elif (
isinstance(value, ast.Attribute)
and value.attr in _AUTO_UNSAFE_PY_WRITE_METHODS
):
writer_aliases.update(targets) # s = np.save
elif isinstance(value, ast.Attribute) and value.attr == "open":
# A captured .open bound method (p = Path('out').open) opens a file
@ -1866,8 +1946,14 @@ def _python_is_potentially_unsafe(code: str) -> bool:
elif (
isinstance(value, ast.Call)
and (
(isinstance(value.func, ast.Name) and value.func.id in partial_aliases)
or (isinstance(value.func, ast.Attribute) and value.func.attr == "partial")
(
isinstance(value.func, ast.Name)
and value.func.id in partial_aliases
)
or (
isinstance(value.func, ast.Attribute)
and value.func.attr == "partial"
)
)
and value.args
and _wraps_write_callable(value.args[0])
@ -1876,7 +1962,10 @@ def _python_is_potentially_unsafe(code: str) -> bool:
elif (
isinstance(value, ast.Call)
and (
(isinstance(value.func, ast.Name) and value.func.id in methodcaller_aliases)
(
isinstance(value.func, ast.Name)
and value.func.id in methodcaller_aliases
)
or (
isinstance(value.func, ast.Attribute)
and value.func.attr == "methodcaller"
@ -1891,14 +1980,20 @@ def _python_is_potentially_unsafe(code: str) -> bool:
# base = '/etc' -> resolve base in a later folded path. A name
# bound more than once is poisoned (\x02) so it fails closed.
for t in targets:
literal_str_vars[t] = "\x02" if t in multi_assigned_names else value.value
literal_str_vars[t] = (
"\x02" if t in multi_assigned_names else value.value
)
elif isinstance(value, (ast.Call, ast.BinOp, ast.Name, ast.JoinedStr)):
# p = Path('/etc'); q = p; r = os.path.join('/etc','x'): record a
# fully-literal folded path so a later reuse (p / 'passwd') folds.
folded = _folded_path(value, literal_str_vars, path_ctor_aliases, pathjoin_aliases)
folded = _folded_path(
value, literal_str_vars, path_ctor_aliases, pathjoin_aliases
)
if folded is not None and "\x00" not in folded and "\x02" not in folded:
for t in targets:
literal_str_vars[t] = "\x02" if t in multi_assigned_names else folded
literal_str_vars[t] = (
"\x02" if t in multi_assigned_names else folded
)
elif isinstance(value, (ast.Tuple, ast.List)):
# Destructuring binds each element like a single assignment, so an
# aliased callable (f, _ = (open, print)) AND a string / path
@ -1906,30 +2001,54 @@ def _python_is_potentially_unsafe(code: str) -> bool:
# the latter a path folded from base/leaf would miss the sensitive
# target and auto-approve.
for target in assign_targets:
if isinstance(target, (ast.Tuple, ast.List)) and len(target.elts) == len(
value.elts
):
if isinstance(target, (ast.Tuple, ast.List)) and len(
target.elts
) == len(value.elts):
for tgt_el, val_el in zip(target.elts, value.elts):
if not isinstance(tgt_el, ast.Name):
continue
tid = tgt_el.id
if isinstance(val_el, ast.Name) and val_el.id in open_aliases:
if (
isinstance(val_el, ast.Name)
and val_el.id in open_aliases
):
open_aliases.add(tid)
elif isinstance(val_el, ast.Name) and val_el.id in getattr_aliases:
elif (
isinstance(val_el, ast.Name)
and val_el.id in getattr_aliases
):
getattr_aliases.add(tid)
elif isinstance(val_el, ast.Name) and val_el.id in partial_aliases:
elif (
isinstance(val_el, ast.Name)
and val_el.id in partial_aliases
):
partial_aliases.add(tid)
elif isinstance(val_el, ast.Name) and val_el.id in writer_aliases:
elif (
isinstance(val_el, ast.Name)
and val_el.id in writer_aliases
):
writer_aliases.add(tid) # s, _ = (save, 1)
elif isinstance(val_el, ast.Name) and val_el.id in archive_ctor_aliases:
elif (
isinstance(val_el, ast.Name)
and val_el.id in archive_ctor_aliases
):
archive_ctor_aliases.add(tid) # z, _ = (ZipFile, 1)
elif isinstance(val_el, ast.Constant) and isinstance(val_el.value, str):
elif isinstance(val_el, ast.Constant) and isinstance(
val_el.value, str
):
literal_str_vars[tid] = (
"\x02" if tid in multi_assigned_names else val_el.value
"\x02"
if tid in multi_assigned_names
else val_el.value
)
elif isinstance(val_el, (ast.Call, ast.BinOp, ast.Name, ast.JoinedStr)):
elif isinstance(
val_el, (ast.Call, ast.BinOp, ast.Name, ast.JoinedStr)
):
folded = _folded_path(
val_el, literal_str_vars, path_ctor_aliases, pathjoin_aliases
val_el,
literal_str_vars,
path_ctor_aliases,
pathjoin_aliases,
)
if (
folded is not None
@ -1937,7 +2056,9 @@ def _python_is_potentially_unsafe(code: str) -> bool:
and "\x02" not in folded
):
literal_str_vars[tid] = (
"\x02" if tid in multi_assigned_names else folded
"\x02"
if tid in multi_assigned_names
else folded
)
elif isinstance(node, (ast.FunctionDef, ast.AsyncFunctionDef, ast.Lambda)):
# A callable captured as a parameter default (def f(o=open): o('x','w'))
@ -2050,13 +2171,17 @@ def _python_is_potentially_unsafe(code: str) -> bool:
# dynamic segment under a sensitive dir (f'/etc/{name}'), or one
# split through a literal variable (base = '/etc'; base+'/passwd').
if _folded_is_sensitive(
_folded_path(node, literal_str_vars, path_ctor_aliases, pathjoin_aliases)
_folded_path(
node, literal_str_vars, path_ctor_aliases, pathjoin_aliases
)
):
return True
elif isinstance(node, ast.Call):
# A sensitive path composed via os.path.join('/etc', name).
if _folded_is_sensitive(
_folded_path(node, literal_str_vars, path_ctor_aliases, pathjoin_aliases)
_folded_path(
node, literal_str_vars, path_ctor_aliases, pathjoin_aliases
)
):
return True
func = node.func
@ -2176,7 +2301,10 @@ def _python_is_potentially_unsafe(code: str) -> bool:
# Path('/home').glob('*') enumerates the receiver dir;
# glob.glob('/home/*') enumerates the pattern's root dir.
_recv = _folded_path(
func.value, literal_str_vars, path_ctor_aliases, pathjoin_aliases
func.value,
literal_str_vars,
path_ctor_aliases,
pathjoin_aliases,
)
if isinstance(_recv, str) and _recv not in ("", "\x00"):
_enum_dir = func.value
@ -2191,7 +2319,10 @@ def _python_is_potentially_unsafe(code: str) -> bool:
_enum_dir = node.args[0]
if _enum_dir is not None:
_folded_dir = _folded_path(
_enum_dir, literal_str_vars, path_ctor_aliases, pathjoin_aliases
_enum_dir,
literal_str_vars,
path_ctor_aliases,
pathjoin_aliases,
)
if isinstance(_folded_dir, str) and (
_folded_dir.startswith("/")
@ -2247,9 +2378,7 @@ _SQL_DDL_OBJECTS = (
)
# Modifiers between the DDL verb and object (CREATE OR REPLACE VIEW, DROP
# MATERIALIZED VIEW, CREATE UNIQUE INDEX).
_SQL_DDL_MODIFIERS = (
r"(?:(?:or\s+replace|unique|temp|temporary|global|local|materialized|recursive)\s+)*"
)
_SQL_DDL_MODIFIERS = r"(?:(?:or\s+replace|unique|temp|temporary|global|local|materialized|recursive)\s+)*"
# A SQL identifier (bare, "quoted", `quoted`, [bracketed]), optionally
# schema-qualified, so UPDATE "users"/public.users/ONLY .../[users] SET all hit.
_SQL_IDENT = r'(?:\w+|"(?:[^"]|"")*"|`(?:[^`]|``)*`|\[[^\]]+\])'
@ -2338,7 +2467,9 @@ _GRAPHQL_COMMENT_RE = re.compile(r"#[^\n]*")
# (mcp__http__get_url {"method": "DELETE"}) mutates an external service even
# though its name looks read-only. GET/HEAD/OPTIONS/TRACE only read.
_MUTATING_HTTP_METHODS = frozenset({"POST", "PUT", "PATCH", "DELETE"})
_HTTP_METHOD_KEYS = frozenset({"method", "http_method", "httpmethod", "verb", "http_verb"})
_HTTP_METHOD_KEYS = frozenset(
{"method", "http_method", "httpmethod", "verb", "http_verb"}
)
def _mcp_arguments_mutate(arguments) -> bool:
@ -2354,7 +2485,9 @@ def _mcp_arguments_mutate(arguments) -> bool:
bool(_MCP_ARG_MUTATION_RE.search(_sql))
or bool(_MCP_ARG_SQLITE_MUTATION_RE.search(_sql))
or bool(_MCP_ARG_SQL_FUNCTION_RE.search(_sql))
or bool(_GRAPHQL_MUTATION_RE.search(_GRAPHQL_COMMENT_RE.sub(" ", value)))
or bool(
_GRAPHQL_MUTATION_RE.search(_GRAPHQL_COMMENT_RE.sub(" ", value))
)
)
if isinstance(value, dict):
for k, v in value.items():
@ -2689,7 +2822,10 @@ def _is_secret_env_value(value: str) -> bool:
"""
if not value:
return False
return _URL_USERINFO_RE.search(value) is not None or _SECRET_VALUE_RE.search(value) is not None
return (
_URL_USERINFO_RE.search(value) is not None
or _SECRET_VALUE_RE.search(value) is not None
)
def _build_bypass_env(workdir: str) -> dict[str, str]:
@ -2763,11 +2899,18 @@ 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
@ -2782,7 +2925,9 @@ def _sandbox_preexec():
# when the parent's hard cap is below the request.
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
@ -2864,7 +3009,9 @@ def _get_project_workdir(session_id: str) -> str | None:
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
@ -2895,7 +3042,9 @@ 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")
@ -2974,7 +3123,8 @@ TERMINAL_TOOL = {
"type": "function",
"function": {
"name": "terminal",
"description": "Execute a terminal command and return stdout/stderr." + _SANDBOX_PATHS_NOTE,
"description": "Execute a terminal command and return stdout/stderr."
+ _SANDBOX_PATHS_NOTE,
"parameters": {
"type": "object",
"properties": {
@ -3082,7 +3232,9 @@ def _mcp_specs_for_server(server: dict, mcp_tools: list[dict]) -> list[dict]:
continue
# Duplicate tool names would also 400 OpenAI; 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(
@ -3091,7 +3243,8 @@ 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": {}},
},
}
)
@ -3110,7 +3263,9 @@ async def get_enabled_mcp_tools() -> list[dict]:
# server gets re-probed -- and blocks the send for the full timeout -- on
# every message.
uncached = [
s for s in servers if get_cached_tools(s["id"]) is None and not in_failure_cooloff(s["id"])
s
for s in servers
if get_cached_tools(s["id"]) is None and not in_failure_cooloff(s["id"])
]
if uncached:
results = await asyncio.gather(
@ -3206,7 +3361,9 @@ def execute_tool(
output). Purely observational: the returned result string is identical
with or without it. Tools without incremental output ignore it.
"""
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 == "search_knowledge_base":
return _search_knowledge_base(arguments, rag_scope)
@ -3407,7 +3564,9 @@ def _message_token_estimate(conversation: list[dict]) -> int:
return total
def _whole_doc_budget(scope: dict | None = None, conversation: list[dict] | None = None) -> int:
def _whole_doc_budget(
scope: dict | None = None, conversation: list[dict] | None = None
) -> int:
try:
from core.rag import config as _rag_config
except Exception: # noqa: BLE001
@ -3447,7 +3606,9 @@ def _last_user_text(conversation: list[dict]) -> str:
return ""
def build_rag_autoinject(conversation: list[dict], rag_scope: dict | None) -> dict | None:
def build_rag_autoinject(
conversation: list[dict], rag_scope: dict | None
) -> dict | None:
"""Pre-retrieve the latest user turn; if a hit clears the cosine floor return
``{"events": [...], "messages": [...]}`` to splice into the loop, else ``None``.
Toggle via ``rag_scope.autoinject`` (else env ``RAG_AUTOINJECT``); floor via
@ -3463,7 +3624,9 @@ def build_rag_autoinject(conversation: list[dict], rag_scope: dict | None) -> di
enabled = _autoinject_enabled()
thread_id = rag_scope.get("thread_id")
whole_doc_requested = (
bool(thread_id) and not rag_scope.get("kb_id") and _thread_whole_doc_enabled(rag_scope)
bool(thread_id)
and not rag_scope.get("kb_id")
and _thread_whole_doc_enabled(rag_scope)
)
if not enabled and not whole_doc_requested:
return None
@ -3474,7 +3637,11 @@ def build_rag_autoinject(conversation: list[dict], rag_scope: dict | None) -> di
from storage import rag_db
if not rag_db.RAG_AVAILABLE:
return None
from core.rag.tool import render_sources, search_for_autoinject, whole_document_context
from core.rag.tool import (
render_sources,
search_for_autoinject,
whole_document_context,
)
except Exception as exc: # noqa: BLE001
logger.warning("RAG auto-inject unavailable: %s", exc)
return None
@ -3517,7 +3684,9 @@ def build_rag_autoinject(conversation: list[dict], rag_scope: dict | None) -> di
**_scope_retrieval_kwargs(rag_scope),
)
except Exception as exc: # noqa: BLE001
logger.warning("RAG project retrieval (whole-doc companion) failed: %s", exc)
logger.warning(
"RAG project retrieval (whole-doc companion) failed: %s", exc
)
proj = None
if proj is not None:
merged = sources + proj[1]
@ -3525,7 +3694,9 @@ def build_rag_autoinject(conversation: list[dict], rag_scope: dict | None) -> di
if max(1, len(merged_text) // 4) <= budget:
sources = merged
text = merged_text
logger.info("RAG auto-inject: whole-document context (%d chunk(s))", len(sources))
logger.info(
"RAG auto-inject: whole-document context (%d chunk(s))", len(sources)
)
if text is None and enabled:
try:
@ -3605,7 +3776,9 @@ _MAX_PDF_FETCH_BYTES = 10 * 1024 * 1024
_MAX_WEB_PDF_PAGES = 50
# Control/undecodable chars, excluding text whitespace and ESC (for ANSI logs).
# Binary when they exceed 12.5%, after allowing 16 minor encoding glitches.
_BINARY_CHAR_RE = re.compile("[\\x00-\\x08\\x0b\\x0c\\x0e-\\x1a\\x1c-\\x1f\\x7f-\\x9f\\ufffd]")
_BINARY_CHAR_RE = re.compile(
"[\\x00-\\x08\\x0b\\x0c\\x0e-\\x1a\\x1c-\\x1f\\x7f-\\x9f\\ufffd]"
)
_MIN_BINARY_CHARS = 16
_BINARY_CHAR_DIVISOR = 8
# Common binary signatures that can otherwise look text-heavy when mislabeled.
@ -4033,7 +4206,11 @@ def _fetch_url_raw(
parsed = urlparse(url)
if parsed.scheme not in ("http", "https"):
return f"Blocked: only http/https URLs are allowed (got {parsed.scheme!r}).", "", ""
return (
f"Blocked: only http/https URLs are allowed (got {parsed.scheme!r}).",
"",
"",
)
if not parsed.hostname:
return "Blocked: URL is missing a hostname.", "", ""
@ -4086,14 +4263,26 @@ def _fetch_url_raw(
resp = opener.open(req, timeout = _fetch_hop_timeout(timeout, deadline))
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.", "", ""
return (
"Failed to fetch URL: redirect missing Location header.",
"",
"",
)
current_url = urljoin(current_url, location)
rp = urlparse(current_url)
if rp.scheme not in ("http", "https") or not rp.hostname:
return "Blocked: redirect target is not a valid http/https URL.", "", ""
return (
"Blocked: redirect target is not a valid http/https URL.",
"",
"",
)
rp_port = rp.port or (443 if rp.scheme == "https" else 80)
ok2, reason2, pinned_ip = _resolve_with_budget(
rp.hostname,
@ -4131,7 +4320,11 @@ def _fetch_url_raw(
# A missing or wrong PDF MIME type is common: once the initial text-sized
# read identifies PDF magic, finish the bounded download to reach the EOF xref.
if not declared_pdf and len(raw_bytes) == max_bytes and _has_pdf_magic(raw_bytes):
if (
not declared_pdf
and len(raw_bytes) == max_bytes
and _has_pdf_magic(raw_bytes)
):
tail_error, tail = _read_capped_body(
resp,
_MAX_PDF_FETCH_BYTES - max_bytes + 1,
@ -4251,7 +4444,9 @@ _HTML_LEADING_TAGS = (
"pre",
"blockquote",
)
_HTML_LEADING_RE = re.compile(r"<(?:!doctype\s+html|/?(?:" + "|".join(_HTML_LEADING_TAGS) + r")\b)")
_HTML_LEADING_RE = re.compile(
r"<(?:!doctype\s+html|/?(?:" + "|".join(_HTML_LEADING_TAGS) + r")\b)"
)
def _looks_like_html(body: str) -> bool:
@ -4335,7 +4530,8 @@ def _fetch_page_text(
readme_body = converted if converted.strip() else body
if readme_body.strip():
return _truncate_page_text(
f"README of {url} (fetched via the GitHub README API):\n\n" + readme_body,
f"README of {url} (fetched via the GitHub README API):\n\n"
+ readme_body,
max_chars,
)
@ -4593,7 +4789,9 @@ 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",
@ -4603,7 +4801,9 @@ 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",
@ -4656,7 +4856,9 @@ 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)
@ -4666,7 +4868,9 @@ 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:
@ -4697,7 +4901,8 @@ 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 (
@ -4710,10 +4915,15 @@ 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(
{
@ -4970,7 +5180,9 @@ 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:
@ -5013,9 +5225,15 @@ 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
@ -5146,7 +5364,9 @@ 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)
@ -5232,7 +5452,11 @@ def _check_signal_escape_patterns(code: str):
)
# Direct sock.connect((host, port)) bypasses the FQ-prefix branch.
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:
@ -5269,7 +5493,9 @@ 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"
),
}
)
@ -5370,18 +5596,28 @@ 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:
@ -5554,7 +5790,9 @@ def _missing_path_hint(output: str, workdir: str | None = None) -> str:
# A convention prefix is an out-of-sandbox signal only when the exact failing
# path could not be isolated; scoped to the failing-path error line(s) so a
# prefix mentioned elsewhere doesn't trigger a misleading hint.
convention = any(prefix in line for line in error_lines for prefix in _MISSING_PATH_PREFIXES)
convention = any(
prefix in line for line in error_lines for prefix in _MISSING_PATH_PREFIXES
)
if abs_path is not None:
# Judge the isolated path against the real workdir even when it matches a
# convention prefix, so a genuine miss inside a project rooted under such
@ -5706,13 +5944,17 @@ 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
)
# utf-8 so non-ASCII in model-written code survives the OS default codec
# (Windows cp1252 would otherwise raise UnicodeEncodeError).
with os.fdopen(fd, "w", encoding = "utf-8") as f:
f.write(code)
safe_env = _build_bypass_env(workdir) if disable_sandbox else _build_safe_env(workdir)
safe_env = (
_build_bypass_env(workdir) if disable_sandbox else _build_safe_env(workdir)
)
if disable_sandbox:
# Match the sandboxed Python path without changing bypass shell I/O.
safe_env = dict(safe_env)
@ -5729,7 +5971,9 @@ def _python_exec(
env = safe_env,
)
if sys.platform != "win32":
popen_kwargs["preexec_fn"] = _bypass_preexec if disable_sandbox else _sandbox_preexec
popen_kwargs["preexec_fn"] = (
_bypass_preexec if disable_sandbox else _sandbox_preexec
)
else:
popen_kwargs["creationflags"] = subprocess.CREATE_NO_WINDOW
@ -5839,7 +6083,9 @@ def _bash_exec(
try:
workdir = _get_workdir(session_id)
safe_env = _build_bypass_env(workdir) if disable_sandbox else _build_safe_env(workdir)
safe_env = (
_build_bypass_env(workdir) if disable_sandbox else _build_safe_env(workdir)
)
popen_kwargs = dict(
stdout = subprocess.PIPE,
stderr = subprocess.STDOUT,
@ -5853,7 +6099,9 @@ def _bash_exec(
env = safe_env,
)
if sys.platform != "win32":
popen_kwargs["preexec_fn"] = _bypass_preexec if disable_sandbox else _sandbox_preexec
popen_kwargs["preexec_fn"] = (
_bypass_preexec if disable_sandbox else _sandbox_preexec
)
else:
popen_kwargs["creationflags"] = subprocess.CREATE_NO_WINDOW

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

@ -40,7 +40,9 @@ def _ensure_backend_on_path() -> None:
sys.path.insert(0, _BACKEND_PATH)
def _activate_transformers_version(model_name: str, hf_token: str | None = None) -> None:
def _activate_transformers_version(
model_name: str, hf_token: str | None = None
) -> None:
"""Activate the correct transformers version BEFORE any ML imports."""
_ensure_backend_on_path()
@ -184,7 +186,9 @@ def _ensure_ssm_kernels(targets: list, resp_queue: Any) -> bool:
try:
from utils.ssm_runtime import ensure_ssm_runtime
except Exception as exc:
logger.debug("ssm_runtime unavailable (%s); skipping SSM kernel pre-install", exc)
logger.debug(
"ssm_runtime unavailable (%s); skipping SSM kernel pre-install", exc
)
return True
_ssm_status = lambda m: _send_response(resp_queue, {"type": "status", "message": m})
@ -304,10 +308,13 @@ def _handle_load(backend, config: dict, resp_queue: Any) -> None:
)
trust_remote_code = config.get("trust_remote_code", False)
if not trust_remote_code and _needs_nemotron_trust(config["model_name"], hf_token = hf_token):
if not trust_remote_code and _needs_nemotron_trust(
config["model_name"], hf_token = hf_token
):
trust_remote_code = True
logger.info(
"Auto-enabled trust_remote_code for Nemotron model: %s", config["model_name"]
"Auto-enabled trust_remote_code for Nemotron model: %s",
config["model_name"],
)
# Authoritative gates over the model + the LoRA base resolved via mc. Must run before
@ -332,7 +339,9 @@ def _handle_load(backend, config: dict, resp_queue: Any) -> None:
from utils.ssm_runtime import ssm_probe_identifier
_ssm_base = (
str(mc.base_model) if (mc.is_lora and getattr(mc, "base_model", None)) else None
str(mc.base_model)
if (mc.is_lora and getattr(mc, "base_model", None))
else None
)
ssm_targets = [ssm_probe_identifier(config["model_name"], _ssm_base)]
if not _ensure_ssm_kernels(ssm_targets, resp_queue):
@ -351,8 +360,12 @@ def _handle_load(backend, config: dict, resp_queue: Any) -> None:
heartbeat_stop = start_watchdog(
repo_ids = watch_repos,
on_stall = lambda msg: _send_response(resp_queue, {"type": "stall", "message": msg}),
on_heartbeat = lambda msg: _send_response(resp_queue, {"type": "status", "message": msg}),
on_stall = lambda msg: _send_response(
resp_queue, {"type": "stall", "message": msg}
),
on_heartbeat = lambda msg: _send_response(
resp_queue, {"type": "status", "message": msg}
),
xet_disabled = os.environ.get("HF_HUB_DISABLE_XET") == "1",
)
try:
@ -386,7 +399,9 @@ def _handle_load(backend, config: dict, resp_queue: Any) -> None:
}
_bm = getattr(backend, "models", {}) or {}
_entry = (
_bm.get(mc.identifier) or _bm.get(getattr(backend, "active_model_name", None)) or {}
_bm.get(mc.identifier)
or _bm.get(getattr(backend, "active_model_name", None))
or {}
)
try:
_context_length = _entry.get("context_length")
@ -659,7 +674,9 @@ def _handle_generate_audio(backend, cmd: dict, resp_queue: Any) -> None:
)
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", "")
@ -694,7 +711,9 @@ def _handle_generate_audio_input(backend, cmd: dict, resp_queue: Any, cancel_eve
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(
@ -777,7 +796,9 @@ def run_inference_process(
than run — the cancel survives the queue handoff.
"""
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"
@ -817,7 +838,10 @@ def run_inference_process(
exc,
)
try:
from core.inference.mlx_inference import MLXInferenceBackend, _init_mlx_distributed
from core.inference.mlx_inference import (
MLXInferenceBackend,
_init_mlx_distributed,
)
backend = MLXInferenceBackend()
if config.get("mlx_distributed"):
@ -961,7 +985,10 @@ def run_inference_process(
if _local_adapter_cfg.is_file():
try:
_lora_base = (
_json.loads(_local_adapter_cfg.read_text()).get("base_model_name_or_path") or None
_json.loads(_local_adapter_cfg.read_text()).get(
"base_model_name_or_path"
)
or None
)
except Exception:
_lora_base = None
@ -994,9 +1021,9 @@ def run_inference_process(
_gate_targets = [model_name]
if _lora_base:
_gate_targets.append(_lora_base)
_trust_remote_code = config.get("trust_remote_code", False) or _needs_nemotron_trust(
model_name, hf_token = _hf_token
)
_trust_remote_code = config.get(
"trust_remote_code", False
) or _needs_nemotron_trust(model_name, hf_token = _hf_token)
if not _run_security_gates(
_gate_targets,
trust_remote_code = _trust_remote_code,
@ -1186,7 +1213,9 @@ 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,
{

12
studio/backend/core/rag/captioner.py
View file

@ -72,7 +72,9 @@ def vision_endpoint() -> tuple[str, str] | None:
try:
from routes.inference import get_llama_cpp_backend
backend = get_llama_cpp_backend()
if getattr(backend, "is_loaded", False) and getattr(backend, "is_vision", False):
if getattr(backend, "is_loaded", False) and getattr(
backend, "is_vision", False
):
return backend.base_url, "local"
except Exception: # noqa: BLE001 - never let discovery break ingestion
return None
@ -139,7 +141,9 @@ def _vision_complete(
return None
def _caption_one(base_url: str, model: str, image_bytes: bytes, timeout: float) -> str | None:
def _caption_one(
base_url: str, model: str, image_bytes: bytes, timeout: float
) -> str | None:
return _vision_complete(
base_url,
model,
@ -150,7 +154,9 @@ def _caption_one(base_url: str, model: str, image_bytes: bytes, timeout: float)
)
def _ocr_one(base_url: str, model: str, image_bytes: bytes, timeout: float) -> str | None:
def _ocr_one(
base_url: str, model: str, image_bytes: bytes, timeout: float
) -> str | None:
return _vision_complete(
base_url,
model,

18
studio/backend/core/rag/chunking.py
View file

@ -27,7 +27,9 @@ class Chunk:
page_char_end: int
def _split(text: str, seps: tuple[str, ...], max_tokens: int, count: TokenCounter) -> list[str]:
def _split(
text: str, seps: tuple[str, ...], max_tokens: int, count: TokenCounter
) -> list[str]:
"""Recursively split into pieces each <= max_tokens (best effort). Pieces
rejoin to ``text`` exactly, so offsets are a running length."""
if count(text) <= max_tokens:
@ -41,7 +43,9 @@ def _split(text: str, seps: tuple[str, ...], max_tokens: int, count: TokenCounte
out: list[str] = []
for p in parts:
out.extend(
[p] if count(p) <= max_tokens else _split(p, seps[i + 1 :], max_tokens, count)
[p]
if count(p) <= max_tokens
else _split(p, seps[i + 1 :], max_tokens, count)
)
return [p for p in out if p]
n = max(1, max_tokens * 4)
@ -49,7 +53,11 @@ def _split(text: str, seps: tuple[str, ...], max_tokens: int, count: TokenCounte
def _merge(
pieces: list[str], starts: list[int], max_tokens: int, overlap: int, count: TokenCounter
pieces: list[str],
starts: list[int],
max_tokens: int,
overlap: int,
count: TokenCounter,
) -> list[tuple[str, int, int]]:
"""Greedy-merge pieces into <= max_tokens chunks with token overlap.
``starts[i]`` is ``pieces[i]``'s page char offset; returns
@ -106,7 +114,9 @@ def chunk_pages(
for piece in pieces:
starts.append(cursor)
cursor += len(piece)
for text, char_start, char_end in _merge(pieces, starts, max_tokens, overlap, count):
for text, char_start, char_end in _merge(
pieces, starts, max_tokens, overlap, count
):
out.append(
Chunk(
text = text,

4
studio/backend/core/rag/config.py
View file

@ -116,7 +116,9 @@ def effective_gguf_repo() -> str:
# llama-server backend only. F16 over Q8_0: faster (no per-block dequant for this
# tiny model) and exact vs fp32, for ~30MB more on disk.
EMBED_GGUF_REPO = os.environ.get("RAG_EMBED_GGUF_REPO", "unsloth/bge-small-en-v1.5-GGUF")
EMBED_GGUF_REPO = os.environ.get(
"RAG_EMBED_GGUF_REPO", "unsloth/bge-small-en-v1.5-GGUF"
)
EMBED_GGUF_VARIANT = os.environ.get("RAG_EMBED_GGUF_VARIANT", "F16")
EMBED_DEVICE = os.environ.get("RAG_EMBED_DEVICE", "auto") # "auto" | "gpu" | "cpu"
EMBED_HOST = os.environ.get("RAG_EMBED_HOST", "127.0.0.1")

24
studio/backend/core/rag/embed_llama_server.py
View file

@ -68,7 +68,9 @@ class LlamaServerBackend:
# Sticky after an auto GPU start fails: later spawns stay on CPU.
self._force_cpu = False
# Pooled client (full URLs per request survive a respawn); trust_env=False skips HTTP(S)_PROXY.
self._client = httpx.Client(timeout = config.EMBED_REQUEST_TIMEOUT_S, trust_env = False)
self._client = httpx.Client(
timeout = config.EMBED_REQUEST_TIMEOUT_S, trust_env = False
)
atexit.register(self._shutdown)
@property
@ -220,7 +222,9 @@ class LlamaServerBackend:
gpus = LlamaCppBackend._get_gpu_free_memory() # [(idx, free_mib)], honors CVD
return any(free >= LlamaServerBackend._MIN_GPU_FREE_MIB for _, free in gpus)
def _build_cmd(self, binary: str, model_path: str, port: int, *, use_gpu: bool) -> list[str]:
def _build_cmd(
self, binary: str, model_path: str, port: int, *, use_gpu: bool
) -> list[str]:
# No --embd-normalize (not in every build; we normalize in Python to match
# the ST path). --fit off: don't auto-resize ctx/offload to device memory.
cmd = [
@ -263,8 +267,12 @@ class LlamaServerBackend:
arch = platform.machine()
lib_dirs = [binary_dir]
for pattern in (
os.path.join(sys.prefix, "lib", "python*", "site-packages", "nvidia", "cu*", "lib"),
os.path.join(sys.prefix, "lib", "python*", "site-packages", "nvidia", "cudnn", "lib"),
os.path.join(
sys.prefix, "lib", "python*", "site-packages", "nvidia", "cu*", "lib"
),
os.path.join(
sys.prefix, "lib", "python*", "site-packages", "nvidia", "cudnn", "lib"
),
):
lib_dirs.extend(d for d in glob.glob(pattern) if os.path.isdir(d))
for cuda_lib in (
@ -378,7 +386,9 @@ class LlamaServerBackend:
def _current(self) -> bool:
"""Alive AND serving the effective repo (a Settings model change makes a
live server stale)."""
return self._process_alive() and self._model_repo == config.effective_gguf_repo()
return (
self._process_alive() and self._model_repo == config.effective_gguf_repo()
)
def _ensure_ready(self) -> None:
"""Guarantee a live server on the effective model, (re)spawning if needed.
@ -449,7 +459,9 @@ class LlamaServerBackend:
raise RuntimeError(
f"llama-server embedder POST {path} -> {e.response.status_code}: {body}"
) from e
raise RuntimeError(f"llama-server embedder POST {path} failed after retry") from last_exc
raise RuntimeError(
f"llama-server embedder POST {path} failed after retry"
) from last_exc
def encode(
self,

12
studio/backend/core/rag/embeddings.py
View file

@ -133,9 +133,13 @@ def _guard_model_security(name: str) -> None:
# directly under a Transformer module dir (0_Transformer/) blocks instead of
# passing as an unreferenced nested shard.
load_subdirs = tuple(
dict.fromkeys((*security_load_subdirs(name, token), *_st_module_subdirs(name, token)))
dict.fromkeys(
(*security_load_subdirs(name, token), *_st_module_subdirs(name, token))
)
)
blocked = evaluate_file_security(name, hf_token = token, load_subdirs = load_subdirs).blocked
blocked = evaluate_file_security(
name, hf_token = token, load_subdirs = load_subdirs
).blocked
except Exception:
return
if blocked:
@ -158,7 +162,9 @@ def _get(model_name: str | None = None):
device = _device()
logger.info("loading embedding model %s on %s", name, device)
_guard_model_security(name)
_model = SentenceTransformer(name, device = device, model_kwargs = dtype_kwargs("float16"))
_model = SentenceTransformer(
name, device = device, model_kwargs = dtype_kwargs("float16")
)
_name = name
return _model

56
studio/backend/core/rag/ingestion.py
View file

@ -117,7 +117,8 @@ def _ocr_scanned_pages(
scanned = [
p.page_number
for p in pages
if p.page_number is not None and len((p.text or "").strip()) < config.OCR_MIN_CHARS
if p.page_number is not None
and len((p.text or "").strip()) < config.OCR_MIN_CHARS
]
if not scanned or captioner.vision_endpoint() is None:
return pages, set()
@ -143,15 +144,21 @@ def _ocr_scanned_pages(
text = texts.get(page.page_number)
if text:
original = (page.text or "").strip()
merged = text if not original or original in text else f"{original}\n\n{text}"
out.append(Page(text = merged, page_number = page.page_number, char_count = len(merged)))
merged = (
text if not original or original in text else f"{original}\n\n{text}"
)
out.append(
Page(text = merged, page_number = page.page_number, char_count = len(merged))
)
ocred.add(page.page_number)
else:
out.append(page)
return out, ocred
def _replace_old_document(conn, replaces: tuple[str, str | None] | None, keep_path: str) -> None:
def _replace_old_document(
conn, replaces: tuple[str, str | None] | None, keep_path: str
) -> None:
"""Drop the document this ingestion replaced (stale embedder / empty prior
ingest), called only after the replacement completed successfully."""
if replaces is None:
@ -214,7 +221,9 @@ def _run(
tiles = []
if tiles:
_progress(conn, job_id, "captioning", 0.28)
captions = captioner.merge_page_captions(captioner.caption_images(tiles))
captions = captioner.merge_page_captions(
captioner.caption_images(tiles)
)
pages = captioner.splice_captions(pages, captions)
_progress(conn, job_id, "chunking", 0.3)
@ -242,12 +251,16 @@ def _run(
from . import locators
regions = locators.pdf_regions_for_chunks(stored_path, pages, chunks)
except Exception:
logger.warning("pdf region location failed for job %s", job_id, exc_info = True)
logger.warning(
"pdf region location failed for job %s", job_id, exc_info = True
)
regions = None
_progress(conn, job_id, "storing", 0.9)
store.add_chunks(conn, scope, document_id, chunks, vectors, regions)
store.set_document_status(conn, document_id, "completed", num_chunks = len(chunks))
store.set_document_status(
conn, document_id, "completed", num_chunks = len(chunks)
)
_replace_old_document(conn, replaces, stored_path)
_set_job(conn, job_id, status = "completed", stage = "done", progress = 1.0)
@ -299,7 +312,9 @@ def start_ingestion(
if existing is not None:
doc = store.get_document(conn, existing)
empty_completed = (
doc is not None and doc.get("status") == "completed" and not doc.get("num_chunks")
doc is not None
and doc.get("status") == "completed"
and not doc.get("num_chunks")
)
# Vectors from a different embedder are stale; re-uploading must
# re-index, not dedupe. NULL (legacy rows) is assumed current. Only
@ -317,13 +332,19 @@ def start_ingestion(
# different model. Re-ingest, don't dedupe.
replaces = (existing, doc.get("stored_path"))
else:
job_id = _new_job(conn, existing, scope, status = "completed", progress = 1.0)
job_id = _new_job(
conn, existing, scope, status = "completed", progress = 1.0
)
_remove_upload(stored_path)
with _jobs_lock:
_jobs[job_id] = queue.Queue()
_emit(
job_id,
{"type": "complete", "num_chunks": doc.get("num_chunks") or 0, "deduped": True},
{
"type": "complete",
"num_chunks": doc.get("num_chunks") or 0,
"deduped": True,
},
)
_emit(job_id, None)
return existing, job_id
@ -354,7 +375,16 @@ def start_ingestion(
# effective_model (not the raw model_name) pins the embedder for the
# whole job: a Settings change mid-ingestion must not switch tokenizer
# or embedder between batches of one document.
args = (job_id, document_id, scope, stored_path, effective_model, ocr, caption, replaces),
args = (
job_id,
document_id,
scope,
stored_path,
effective_model,
ocr,
caption,
replaces,
),
daemon = True,
).start()
return document_id, job_id
@ -437,7 +467,9 @@ def job_events(job_id: str):
# drop a document whose worker is still running. Heartbeat and
# retry on the next poll instead.
logger.warning(
"job_events status read failed for %s; continuing", job_id, exc_info = True
"job_events status read failed for %s; continuing",
job_id,
exc_info = True,
)
yield {"type": "heartbeat"}
continue

8
studio/backend/core/rag/locators.py
View file

@ -120,7 +120,9 @@ def _rects_from_words(page_words: list, indices: list[int], pw: float, ph: float
return out
def _regions_for_match(doc: Any, page_text: str, match: LocatorMatch) -> list[dict[str, Any]]:
def _regions_for_match(
doc: Any, page_text: str, match: LocatorMatch
) -> list[dict[str, Any]]:
try:
if match.page_index < 0 or match.page_index >= len(doc):
return []
@ -147,7 +149,9 @@ def _regions_for_match(doc: Any, page_text: str, match: LocatorMatch) -> list[di
return []
def pdf_regions_for_chunks(pdf_path: Path, pages: list, chunks: list) -> list[list[dict[str, Any]]]:
def pdf_regions_for_chunks(
pdf_path: Path, pages: list, chunks: list
) -> list[list[dict[str, Any]]]:
"""Region rects per chunk (parallel to ``chunks``), keyed off each chunk's
``source_page_index`` / ``page_char_start`` / ``page_char_end``. Non-PDFs and
failures yield [], never an exception."""

22
studio/backend/core/rag/parsers.py
View file

@ -88,7 +88,9 @@ def _markdown_corrupted(text: str) -> bool:
legitimate shaped glyph does not force the fallback)."""
if not text:
return False
threshold = max(_PDF_FALLBACK_MIN_BAD_GLYPHS, _PDF_FALLBACK_BAD_GLYPH_RATIO * len(text))
threshold = max(
_PDF_FALLBACK_MIN_BAD_GLYPHS, _PDF_FALLBACK_BAD_GLYPH_RATIO * len(text)
)
shaped = len(_SHAPED_PRESENTATION_FORMS.findall(text))
return shaped > threshold or text.count("\ufffd") > threshold
@ -135,13 +137,17 @@ def _pdf(
pages: list[Page] = []
images: list[ParsedImage] = []
doc = (
fitz.open(stream = source, filetype = "pdf") if isinstance(source, bytes) else fitz.open(source)
fitz.open(stream = source, filetype = "pdf")
if isinstance(source, bytes)
else fitz.open(source)
)
try:
if doc.needs_pass:
raise ValueError("encrypted PDF requires a password")
total_pages = doc.page_count
page_numbers = range(total_pages if max_pages is None else min(total_pages, max_pages))
page_numbers = range(
total_pages if max_pages is None else min(total_pages, max_pages)
)
if not config.PDF_MARKDOWN:
md = None
elif max_pages is None:
@ -186,7 +192,9 @@ def _pdf(
return pages, images, total_pages
def parse_pdf_bytes(data: bytes, *, max_pages: int | None = None) -> tuple[list[Page], int]:
def parse_pdf_bytes(
data: bytes, *, max_pages: int | None = None
) -> tuple[list[Page], int]:
"""Extract PDF pages from an in-memory download using the ingestion parser.
Returns the (capped) pages plus the document's full page count, so a caller
@ -330,7 +338,11 @@ def render_pdf_figure_tiles(
for clip in clips:
try:
pix = page.get_pixmap(dpi = dpi, clip = clip)
out.append(ParsedImage(image_bytes = pix.tobytes("png"), page_number = num, xref = 0))
out.append(
ParsedImage(
image_bytes = pix.tobytes("png"), page_number = num, xref = 0
)
)
except Exception:
continue
if len(out) >= max_tiles:

21
studio/backend/core/rag/retrieval.py
View file

@ -27,7 +27,10 @@ def retrieve_lexical(
k: int | None = None,
) -> list[Hit]:
k = k or config.TOP_K_LEXICAL
return [Hit(cid, s, lexical_score = s) for cid, s in store.search_lexical(conn, scope, query, k)]
return [
Hit(cid, s, lexical_score = s)
for cid, s in store.search_lexical(conn, scope, query, k)
]
def retrieve_dense(
@ -52,13 +55,19 @@ def _rrf(rankings: list[list[Hit]], rrf_k: int, top_k: int) -> list[Hit]:
best: dict[str, Hit] = {}
for ranking in rankings:
for rank, hit in enumerate(ranking):
fused[hit.chunk_id] = fused.get(hit.chunk_id, 0.0) + 1.0 / (rrf_k + rank + 1)
fused[hit.chunk_id] = fused.get(hit.chunk_id, 0.0) + 1.0 / (
rrf_k + rank + 1
)
cur = best.get(hit.chunk_id)
if cur is None:
best[hit.chunk_id] = Hit(hit.chunk_id, 0.0, hit.lexical_score, hit.dense_score)
best[hit.chunk_id] = Hit(
hit.chunk_id, 0.0, hit.lexical_score, hit.dense_score
)
else:
cur.lexical_score = (
cur.lexical_score if cur.lexical_score is not None else hit.lexical_score
cur.lexical_score
if cur.lexical_score is not None
else hit.lexical_score
)
cur.dense_score = (
cur.dense_score if cur.dense_score is not None else hit.dense_score
@ -89,7 +98,9 @@ def retrieve_hybrid(
if mode == "dense":
return retrieve_dense(conn, scope, query, k, model_name = model_name)
lexical = retrieve_lexical(conn, scope, query, config.TOP_K_LEXICAL)
dense = retrieve_dense(conn, scope, query, config.TOP_K_DENSE, model_name = model_name)
dense = retrieve_dense(
conn, scope, query, config.TOP_K_DENSE, model_name = model_name
)
return _rrf([lexical, dense], config.RRF_K, k)

9
studio/backend/core/rag/store.py
View file

@ -89,7 +89,10 @@ def delete_kb(conn: sqlite3.Connection, kb_id: str) -> None:
"""Delete a knowledge base and every document (+ chunks) under it."""
scope = kb_scope(kb_id)
doc_ids = [
r["id"] for r in conn.execute("SELECT id FROM documents WHERE scope=?", (scope,)).fetchall()
r["id"]
for r in conn.execute(
"SELECT id FROM documents WHERE scope=?", (scope,)
).fetchall()
]
for doc_id in doc_ids:
delete_document(conn, doc_id)
@ -182,7 +185,9 @@ def document_by_hash(conn: sqlite3.Connection, scope: str, sha256: str) -> str |
return row["id"] if row else None
def failed_documents_by_hash(conn: sqlite3.Connection, scope: str, sha256: str) -> list[dict]:
def failed_documents_by_hash(
conn: sqlite3.Connection, scope: str, sha256: str
) -> list[dict]:
rows = conn.execute(
"SELECT id, stored_path FROM documents WHERE scope=? AND sha256=? AND status='failed'",
(scope, sha256),

12
studio/backend/core/rag/tool.py
View file

@ -107,7 +107,9 @@ def render_sources(sources: list[dict]) -> str:
src = quoteattr(s.get("filename") or "unknown")
page = s.get("page")
page_attr = f" page={quoteattr(str(page))}" if page else ""
blocks.append(f'<chunk id="{i}" source={src}{page_attr}>\n{s.get("text") or ""}\n</chunk>')
blocks.append(
f'<chunk id="{i}" source={src}{page_attr}>\n{s.get("text") or ""}\n</chunk>'
)
return "\n\n".join(blocks)
@ -197,10 +199,14 @@ def search_for_autoinject(
mode = mode,
)
strong = [
h for h in hits if h.dense_score is not None and h.dense_score >= min_dense_score
h
for h in hits
if h.dense_score is not None and h.dense_score >= min_dense_score
][:k]
if not strong and hits and mode == "lexical":
probe = retrieval.retrieve_dense(conn, scope, query, 1, model_name = model_name)
probe = retrieval.retrieve_dense(
conn, scope, query, 1, model_name = model_name
)
if (
probe
and probe[0].dense_score is not None

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

@ -37,7 +37,9 @@ _BRACKETED_JSON_ONE_LEVEL = r"\{[^{}]*(?:\{[^{}]*\}[^{}]*)*\}"
_REHEARSAL_CLOSED_STRIP_RE = re.compile(
r"(?<!\[CALL_ID\])\b([\w-]+)\[ARGS\]\s*" + _BRACKETED_JSON_ONE_LEVEL, re.DOTALL
)
_REHEARSAL_TAIL_STRIP_RE = re.compile(r"(?<!\[CALL_ID\])\b([\w-]+)\[ARGS\]\s*(?:\{.*)?$", re.DOTALL)
_REHEARSAL_TAIL_STRIP_RE = re.compile(
r"(?<!\[CALL_ID\])\b([\w-]+)\[ARGS\]\s*(?:\{.*)?$", re.DOTALL
)
# Tool-XML strip patterns; hyphen in the name class covers dashed MCP names.
# Closed-pair patterns are named so _PAT_REQUIRED_TOKEN can skip a doomed lazy rescan when
@ -79,7 +81,9 @@ _TOOL_ALL_PATS = (
)
# Rehearsal strips (name in group 1); name-gated via ``enabled_tool_names``, strip-all when None.
_REHEARSAL_STRIP_PATS = frozenset({_REHEARSAL_CLOSED_STRIP_RE, _REHEARSAL_TAIL_STRIP_RE})
_REHEARSAL_STRIP_PATS = frozenset(
{_REHEARSAL_CLOSED_STRIP_RE, _REHEARSAL_TAIL_STRIP_RE}
)
# Stripped before the quote-aware Gemma helper so a Gemma opener quoted in argument
# data cannot make the helper truncate the block and its tail.
@ -117,7 +121,9 @@ def apply_tool_strip_patterns(
if token is not None and token not in text:
continue
if enabled_tool_names is not None and pat in _REHEARSAL_STRIP_PATS:
text = pat.sub(lambda m: "" if m.group(1) in enabled_tool_names else m.group(0), text)
text = pat.sub(
lambda m: "" if m.group(1) in enabled_tool_names else m.group(0), text
)
else:
text = pat.sub("", text)
return text
@ -148,7 +154,9 @@ _GEMMA_NEXT_KEY_RE = re.compile(r"\s*[A-Za-z_][\w.\-]*\s*:")
# A candidate starting inside a think block is a rehearsal (block kept so literal tags in
# real args survive); ``$`` accepts an unclosed block mid-stream.
_THINK_TAG_RE = re.compile(r"<think>.*?(?:</think>|$)|\[THINK\].*?(?:\[/THINK\]|$)", re.DOTALL)
_THINK_TAG_RE = re.compile(
r"<think>.*?(?:</think>|$)|\[THINK\].*?(?:\[/THINK\]|$)", re.DOTALL
)
# Bare open/close markers for prefilled-reasoning turns (template opens <think> in the prompt).
_THINK_OPEN_RE = re.compile(r"<think>|\[THINK\]")
_THINK_CLOSE_RE = re.compile(r"</think>|\[/THINK\]")
@ -405,7 +413,9 @@ def _quote_gemma_array_elements(body: str) -> str:
# Nested array: normalise its elements too.
inner_end = _balanced_bracket_end(stripped, 0)
if inner_end == len(stripped) - 1:
out.append("[" + _quote_gemma_array_elements(stripped[1:inner_end]) + "]")
out.append(
"[" + _quote_gemma_array_elements(stripped[1:inner_end]) + "]"
)
else:
out.append(element)
continue
@ -493,7 +503,9 @@ def _quote_gemma_object_keys(src: str) -> str:
parts.append(src[i:])
i = len(src)
else:
parts.append("[" + _quote_gemma_array_elements(src[i + 1 : arr_end]) + "]")
parts.append(
"[" + _quote_gemma_array_elements(src[i + 1 : arr_end]) + "]"
)
i = arr_end + 1
elif i < len(src) and src[i] not in '"{':
v_start = i
@ -595,7 +607,9 @@ def _marker_coverage(content: str, markers) -> list[tuple[int, int]]:
if order == 0:
waiting[kind].append(payload) # marker index, now awaiting its close
elif waiting[kind]:
close_end_for[waiting[kind].pop()] = payload # innermost open marker closes here
close_end_for[waiting[kind].pop()] = (
payload # innermost open marker closes here
)
coverage = []
for idx, (start, brace_end, _kind, _m) in enumerate(markers):
if brace_end < 0:
@ -708,7 +722,9 @@ def parse_tool_calls_from_text(
start -= len("<|message_model|>")
else:
name = m.group(1)
arguments = json.dumps(_gemma_arguments_to_json(content[m.end() : brace_end]))
arguments = json.dumps(
_gemma_arguments_to_json(content[m.end() : brace_end])
)
except (json.JSONDecodeError, ValueError):
continue
span_end = brace_end + 1
@ -729,7 +745,9 @@ def parse_tool_calls_from_text(
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()
@ -766,7 +784,9 @@ def parse_tool_calls_from_text(
param_name = pm.group(1)
val_start = pm.end()
next_param = (
param_starts[pidx + 1].start() if pidx + 1 < len(param_starts) else len(body)
param_starts[pidx + 1].start()
if pidx + 1 < len(param_starts)
else len(body)
)
val = body[val_start:next_param]
if not allow_incomplete:
@ -845,11 +865,15 @@ def parse_tool_calls_from_text(
# A bare scalar string stays raw (like the <tool_call> path);
# json.dumps would double-encode it so the arg healer wraps
# "weather" with its literal quotes.
"arguments": args if isinstance(args, str) else json.dumps(args),
"arguments": args
if isinstance(args, str)
else json.dumps(args),
},
}
)
item_end = item_ends[item_idx] if item_idx < len(item_ends) else region_end
item_end = (
item_ends[item_idx] if item_idx < len(item_ends) else region_end
)
last_span_idx = len(call_spans)
call_spans.append((tile_start, item_end))
tile_start = item_end
@ -890,7 +914,9 @@ def _strip_bracket_tag_calls(text: str, enabled_tool_names = None) -> str:
return text
out: list[str] = []
cursor = 0
for start, end, _kind, _m in _iter_bracket_spans(text, enabled_tool_names = enabled_tool_names):
for start, end, _kind, _m in _iter_bracket_spans(
text, enabled_tool_names = enabled_tool_names
):
out.append(text[cursor:start])
cursor = end
out.append(text[cursor:])
@ -941,7 +967,11 @@ def _think_spans_outside_tool_markup(text: str) -> list[tuple[int, int]]:
return think_spans
if not call_spans:
return think_spans
return [(s, e) for (s, e) in think_spans if not any(cs <= s < ce for cs, ce in call_spans)]
return [
(s, e)
for (s, e) in think_spans
if not any(cs <= s < ce for cs, ce in call_spans)
]
def strip_outside_think(text: str, strip_segment) -> str:
@ -1062,7 +1092,9 @@ def _strip_markup_segment(
text = _strip_closed_blocks_outside_gemma(text)
text = _strip_gemma_native_spans(text, final = final)
patterns = _TOOL_ALL_PATS if final else _TOOL_CLOSED_PATS
return apply_tool_strip_patterns(text, patterns, enabled_tool_names = enabled_tool_names)
return apply_tool_strip_patterns(
text, patterns, enabled_tool_names = enabled_tool_names
)
def strip_tool_call_markup(

8
studio/backend/core/training/s3_dataset.py
View file

@ -159,7 +159,9 @@ def prepare_s3_dataset_download(
bucket/prefix contains no supported dataset files.
"""
if not boto3_available():
raise RuntimeError("S3 dataset loading requires boto3. Install it with: pip install boto3")
raise RuntimeError(
"S3 dataset loading requires boto3. Install it with: pip install boto3"
)
bucket = s3_config.get("bucket")
if not bucket:
@ -193,7 +195,9 @@ def prepare_s3_dataset_download(
local_path = _unique_local_path(target_dir, filename, used_paths)
download_kwargs = {}
if cancel_callback is not None:
download_kwargs["Callback"] = lambda _bytes: _raise_if_cancelled(cancel_callback)
download_kwargs["Callback"] = lambda _bytes: _raise_if_cancelled(
cancel_callback
)
client.download_file(bucket, key, local_path, **download_kwargs)
_raise_if_cancelled(cancel_callback)
local_files.append(local_path)

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

File diff suppressed because it is too large Load diff

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

@ -115,9 +115,13 @@ def _coerce_optional_nonneg_float(name: str, value):
try:
coerced = float(value)
except (TypeError, ValueError):
raise ValueError(f"Unsloth: {name}={value!r} must be a non-negative float or None.")
raise ValueError(
f"Unsloth: {name}={value!r} must be a non-negative float or None."
)
if coerced < 0:
raise ValueError(f"Unsloth: {name}={coerced} must be >= 0 (use 0 or None to disable).")
raise ValueError(
f"Unsloth: {name}={coerced} must be >= 0 (use 0 or None to disable)."
)
return coerced
@ -208,7 +212,9 @@ def _build_training_worker_config(values: dict[str, Any]) -> dict[str, Any]:
"tensorboard_dir": values.get("tensorboard_dir", "runs"),
"resume_from_checkpoint": values.get("resume_from_checkpoint"),
"trust_remote_code": values.get("trust_remote_code", False),
"approved_remote_code_fingerprint": values.get("approved_remote_code_fingerprint"),
"approved_remote_code_fingerprint": values.get(
"approved_remote_code_fingerprint"
),
"subject": values.get("subject"),
"gpu_ids": values.get("gpu_ids"),
"s3_config": values.get("s3_config"),
@ -358,12 +364,16 @@ class _MLXTrainerAdapter:
self._pump_thread: Optional[threading.Thread] = None
self._lock = threading.Lock()
def _activate_transformers_for_model(self, model_name: str, hf_token: Optional[str]) -> None:
def _activate_transformers_for_model(
self, model_name: str, hf_token: Optional[str]
) -> None:
try:
from utils.transformers_version import activate_transformers_for_subprocess
activate_transformers_for_subprocess(model_name, hf_token)
except Exception as exc:
logger.warning("MLX trainer adapter Transformers activation failed", error = str(exc))
logger.warning(
"MLX trainer adapter Transformers activation failed", error = str(exc)
)
def add_progress_callback(self, callback: Callable[[TrainingProgress], None]):
self.progress_callbacks.append(callback)
@ -408,10 +418,16 @@ class _MLXTrainerAdapter:
else:
self.is_audio = self._audio_type is not None
self.is_audio_vlm = False
vision = is_vision_model(model_name, hf_token = hf_token) if not self.is_audio else False
vision = (
is_vision_model(model_name, hf_token = hf_token)
if not self.is_audio
else False
)
self.is_vlm = not self.is_audio_vlm and vision and bool(is_dataset_image)
except Exception as exc:
logger.warning("MLX trainer adapter model type detection failed", error = str(exc))
logger.warning(
"MLX trainer adapter model type detection failed", error = str(exc)
)
self.is_vlm = False
self.is_audio = False
self.is_audio_vlm = False
@ -504,7 +520,10 @@ class _MLXTrainerAdapter:
}
self.is_cpt = bool(is_cpt)
self._update_progress(status_message = "Queued MLX dataset load")
return ({"dataset": [], "final_format": "deferred_mlx_cli", "success": True}, None)
return (
{"dataset": [], "final_format": "deferred_mlx_cli", "success": True},
None,
)
def start_training(
self,
@ -512,12 +531,18 @@ class _MLXTrainerAdapter:
eval_dataset = None,
**training_args,
) -> bool:
if self.is_training and self.training_thread and self.training_thread.is_alive():
if (
self.is_training
and self.training_thread
and self.training_thread.is_alive()
):
return False
if self._pump_thread and self._pump_thread.is_alive():
self._pump_thread.join(timeout = 2.0)
if self._pump_thread.is_alive():
self._update_progress(error = "Previous training event pump is still finalizing")
self._update_progress(
error = "Previous training event pump is still finalizing"
)
return False
if not self._model_config:
self._update_progress(error = "Model not loaded")
@ -645,7 +670,9 @@ class _MLXTrainerAdapter:
not self.training_progress.error
and not self.training_progress.is_completed
):
self.training_progress.error = "Training process exited unexpectedly"
self.training_progress.error = (
"Training process exited unexpectedly"
)
self.is_training = False
self._event_queue = None
self._stop_queue = None
@ -673,17 +700,25 @@ class _MLXTrainerAdapter:
step = event.get("step", self.training_progress.step),
epoch = event.get("epoch", self.training_progress.epoch),
loss = event.get("loss", self.training_progress.loss),
learning_rate = event.get("learning_rate", self.training_progress.learning_rate),
total_steps = event.get("total_steps", self.training_progress.total_steps),
learning_rate = event.get(
"learning_rate", self.training_progress.learning_rate
),
total_steps = event.get(
"total_steps", self.training_progress.total_steps
),
elapsed_seconds = event.get(
"elapsed_seconds",
self.training_progress.elapsed_seconds,
),
eta_seconds = event.get("eta_seconds", self.training_progress.eta_seconds),
eta_seconds = event.get(
"eta_seconds", self.training_progress.eta_seconds
),
grad_norm = event.get("grad_norm", self.training_progress.grad_norm),
num_tokens = event.get("num_tokens", self.training_progress.num_tokens),
eval_loss = event.get("eval_loss", self.training_progress.eval_loss),
peak_memory_gb = event.get("peak_memory_gb", self.training_progress.peak_memory_gb),
peak_memory_gb = event.get(
"peak_memory_gb", self.training_progress.peak_memory_gb
),
)
return
if etype == "complete":
@ -718,7 +753,9 @@ class _MLXTrainerAdapter:
if self._stop_queue is not None:
self._stop_queue.put({"type": "stop", "save": save})
status_message = (
"Stopping training and saving checkpoint..." if save else "Cancelling training..."
"Stopping training and saving checkpoint..."
if save
else "Cancelling training..."
)
self._update_progress(status_message = status_message)
return True
@ -842,7 +879,9 @@ 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
# Clear a stale crash flag from a prior died pump so the watchdog can't
@ -879,7 +918,9 @@ class TrainingBackend:
config["resolved_gpu_ids"] = None
config["gpu_selection"] = None
elif gpu_ids:
resolved_gpu_ids, gpu_selection = prepare_gpu_selection(gpu_ids, **gpu_selection_kwargs)
resolved_gpu_ids, gpu_selection = prepare_gpu_selection(
gpu_ids, **gpu_selection_kwargs
)
config["resolved_gpu_ids"] = resolved_gpu_ids
config["gpu_selection"] = gpu_selection
else:
@ -910,7 +951,9 @@ class TrainingBackend:
try:
before_spawn()
except Exception:
logger.warning("before_spawn hook failed; continuing", exc_info = True)
logger.warning(
"before_spawn hook failed; continuing", exc_info = True
)
if defer_auto_selection:
try:
@ -976,7 +1019,9 @@ class TrainingBackend:
self._metric_buffer.clear()
self._run_finalized = False
self._db_run_created = False
self._db_create_in_progress = False # a stale watchdog create can't block this run
self._db_create_in_progress = (
False # a stale watchdog create can't block this run
)
self._db_total_steps_set = False
self._db_config = _sanitize_db_config(config)
self._db_started_at = datetime.now(timezone.utc).isoformat()
@ -1022,7 +1067,9 @@ 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..."
)
# Guarantee the run finalizes even if the worker wedges after saving.
self._start_stop_watchdog(cancel = not save)
@ -1096,7 +1143,9 @@ class TrainingBackend:
reason,
)
else:
logger.warning("Stop watchdog force-terminating stuck training worker: %s", reason)
logger.warning(
"Stop watchdog force-terminating stuck training worker: %s", reason
)
# force_terminate can raise on a wedged child; finalize regardless.
try:
self.force_terminate(target_proc = target_proc)
@ -1161,7 +1210,13 @@ class TrainingBackend:
loss_history = list(self.loss_history)
if claim:
self._finish_stopped_run(
run_id, output_dir, batch, final_step, final_loss, duration, loss_history
run_id,
output_dir,
batch,
final_step,
final_loss,
duration,
loss_history,
)
with self._lock:
if target_proc is None or self._proc is target_proc:
@ -1274,7 +1329,9 @@ class TrainingBackend:
"Model download stalled even over HTTP -- check your network connection"
)
if recover:
logger.warning("Training model-load stalled on Xet; respawning over HTTP: %s", msg)
logger.warning(
"Training model-load stalled on Xet; respawning over HTTP: %s", msg
)
else:
logger.error("Training download stalled with no further fallback: %s", msg)
# Terminate either way so the pump loop proceeds (respawn or finalize).
@ -1302,7 +1359,9 @@ class TrainingBackend:
config = {**config, "disable_xet": True}
self._last_full_config = config
logger.warning("Respawning training worker with HF_HUB_DISABLE_XET=1 after Xet stall")
logger.warning(
"Respawning training worker with HF_HUB_DISABLE_XET=1 after Xet stall"
)
from .worker import run_training_process
@ -1352,7 +1411,9 @@ class TrainingBackend:
new_proc.start()
from utils.process_lifetime import adopt_pid
adopt_pid(new_proc.pid) # bind to parent lifetime (Windows job / sweep)
adopt_pid(
new_proc.pid
) # bind to parent lifetime (Windows job / sweep)
except Exception:
logger.error("Failed to respawn training subprocess", exc_info = True)
self._spawn_in_progress = False
@ -1369,7 +1430,9 @@ class TrainingBackend:
)
return
logger.info("Training subprocess respawned with Xet disabled (pid=%s)", new_proc.pid)
logger.info(
"Training subprocess respawned with Xet disabled (pid=%s)", new_proc.pid
)
new_pump = threading.Thread(target = self._pump_loop, daemon = True)
with self._lock:
self._in_model_load = False
@ -1532,8 +1595,12 @@ class TrainingBackend:
try:
self._handle_event(event)
except Exception:
etype = event.get("type") if isinstance(event, dict) else type(event).__name__
logger.exception("Training event pump: failed to handle %s event; skipping", etype)
etype = (
event.get("type") if isinstance(event, dict) else type(event).__name__
)
logger.exception(
"Training event pump: failed to handle %s event; skipping", etype
)
def _pump_loop(self) -> None:
"""Background thread: consume subprocess events and update state.
@ -1591,7 +1658,8 @@ 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()
@ -1602,7 +1670,9 @@ class TrainingBackend:
else "Training process terminated unexpectedly",
)
except Exception:
logger.exception("Training event pump: finalization after worker exit failed")
logger.exception(
"Training event pump: finalization after worker exit failed"
)
self._pump_running = False
return
@ -1641,7 +1711,9 @@ class TrainingBackend:
except (TypeError, ValueError):
logger.debug("Could not convert loss to float: %s", _raw_loss)
_safe_loss = None
_loss_is_nonfinite = _safe_loss is not None and not math.isfinite(_safe_loss)
_loss_is_nonfinite = _safe_loss is not None and not math.isfinite(
_safe_loss
)
if _loss_is_nonfinite:
# Drop the value rather than laundering it back to the last
# finite loss; clients see loss=None at this step so the NaN
@ -1657,7 +1729,9 @@ 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
@ -1669,7 +1743,9 @@ class TrainingBackend:
self._progress.loss = None
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")
@ -1713,7 +1789,9 @@ 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)
@ -1743,9 +1821,12 @@ 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 (
@ -1830,7 +1911,9 @@ 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)
@ -1856,7 +1939,12 @@ class TrainingBackend:
if step == prev:
return
now = time.monotonic()
if prev >= 0 and step > prev and not is_final and (now - self._last_progress_log_ts) < 30.0:
if (
prev >= 0
and step > prev
and not is_final
and (now - self._last_progress_log_ts) < 30.0
):
return
self._last_progress_log_ts = now
self._last_progress_log_step = step
@ -1908,7 +1996,9 @@ class TrainingBackend:
)
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
)
finally:
with self._lock:
# Publish the flags only if this is still the current run. A killed worker
@ -1917,7 +2007,9 @@ class TrainingBackend:
# (the row was still created by id; the new run owns/creates its own row).
if self.current_job_id == job_id:
if created:
self._db_run_created = True # publish only after the insert commits
self._db_run_created = (
True # publish only after the insert commits
)
self._db_create_in_progress = False
def _finalize_run_in_db(
@ -1936,7 +2028,11 @@ class TrainingBackend:
with self._lock:
if expected_job_id is not None and self.current_job_id != expected_job_id:
return
if not self.current_job_id or not self._db_run_created or self._run_finalized:
if (
not self.current_job_id
or not self._db_run_created
or self._run_finalized
):
return
self._run_finalized = True
run_id = self.current_job_id
@ -1966,7 +2062,9 @@ class TrainingBackend:
except Exception:
with self._lock:
self._run_finalized = False # unclaim so a later flush can retry
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, run_id: Optional[str] = None) -> None:
"""Flush buffered metrics to the DB and update live progress. The target run id,
@ -1995,7 +2093,9 @@ class TrainingBackend:
try:
from storage.studio_db import insert_metrics_batch, update_run_progress
insert_metrics_batch(target, batch)
update_run_progress(id = target, step = step, loss = loss, duration_seconds = duration)
update_run_progress(
id = target, step = step, loss = loss, duration_seconds = duration
)
except Exception:
# Re-queue the claimed batch at the front so it retries on the next flush.
with self._lock:
@ -2125,7 +2225,9 @@ 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

@ -36,7 +36,8 @@ from typing import Any, Callable
if sys.platform.startswith("linux") and "HSA_ENABLE_DXG_DETECTION" not in os.environ:
try:
if os.path.exists("/dev/dxg") and any(
os.path.exists(_p + "/librocdxg.so") for _p in ("/opt/rocm/lib", "/opt/rocm/lib64")
os.path.exists(_p + "/librocdxg.so")
for _p in ("/opt/rocm/lib", "/opt/rocm/lib64")
):
os.environ["HSA_ENABLE_DXG_DETECTION"] = "1"
except Exception:
@ -55,7 +56,9 @@ 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)
@ -118,7 +121,9 @@ 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,7 +264,9 @@ 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)
@ -344,7 +351,8 @@ 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
@ -462,7 +470,9 @@ 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(
@ -537,7 +547,9 @@ 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:
@ -728,7 +740,10 @@ 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
@ -771,7 +786,9 @@ 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
@ -872,7 +889,9 @@ def _ensure_tilelang_backend(event_queue: Any, model_name: str) -> None:
# UNSLOTH_STUDIO_SKIP_FAST_PATH_HOOKS=1 falls back to the 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 imported `old_obj`.
`from X import Y` creates a local binding that reassigning X.Y won't reach.
@ -945,7 +964,9 @@ 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"
@ -954,7 +975,9 @@ 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
@ -965,7 +988,9 @@ 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)
@ -980,7 +1005,10 @@ def _install_fast_path_hooks(event_queue: Any, model_name: str) -> None:
# FLA 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)
@ -996,7 +1024,9 @@ 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)
@ -1016,7 +1046,9 @@ 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:
@ -1050,7 +1082,9 @@ def _ensure_flash_attn_for_long_context(event_queue: Any, max_seq_length: int) -
_send_status(event_queue, "Continuing without flash-attn")
def _activate_transformers_version(model_name: str, hf_token: str | None = None) -> None:
def _activate_transformers_version(
model_name: str, hf_token: str | None = None
) -> None:
"""Activate the correct transformers version BEFORE any ML imports."""
# Ensure backend is on path for utils imports
backend_path = str(Path(__file__).resolve().parent.parent.parent)
@ -1062,7 +1096,9 @@ def _activate_transformers_version(model_name: str, hf_token: str | None = None)
activate_transformers_for_subprocess(model_name, hf_token)
def _activate_transformers_version_or_warn(model_name: str, hf_token: str | None = None) -> None:
def _activate_transformers_version_or_warn(
model_name: str, hf_token: str | None = None
) -> None:
"""Activate the required transformers version for the MLX fast-path.
Unlike the non-MLX path (which treats activation failure as fatal and
@ -1193,7 +1229,10 @@ def _resize_mlx_vlm_images(
image_layout = None,
):
if isinstance(value, list):
return [_resize_mlx_vlm_image(image, resize, image_layout = image_layout) for image in value]
return [
_resize_mlx_vlm_image(image, resize, image_layout = image_layout)
for image in value
]
return _resize_mlx_vlm_image(value, resize, image_layout = image_layout)
@ -1479,7 +1518,9 @@ 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 for non-image datasets even on vision-capable models
@ -1513,7 +1554,9 @@ def _run_mlx_training(event_queue, stop_queue, config):
from utils.models.model_config import get_base_model_from_lora_identifier
# Resolve a LOCAL or REMOTE adapter's base so a remote LoRA base is gated too.
_base = get_base_model_from_lora_identifier(model_name, config.get("hf_token") or None)
_base = get_base_model_from_lora_identifier(
model_name, config.get("hf_token") or None
)
if _base:
malware_targets.append(_base)
except Exception as exc:
@ -1522,7 +1565,9 @@ def _run_mlx_training(event_queue, stop_queue, config):
for target in dict.fromkeys(malware_targets):
_fs = evaluate_file_security(
target, hf_token = hf_token, load_subdirs = security_load_subdirs(target, hf_token)
target,
hf_token = hf_token,
load_subdirs = security_load_subdirs(target, hf_token),
)
if _fs.blocked:
_send(
@ -1637,9 +1682,15 @@ 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
@ -1672,7 +1723,9 @@ 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:
@ -1761,7 +1814,9 @@ 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,
@ -2032,7 +2087,11 @@ 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,
)
@ -2055,7 +2114,9 @@ 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:
@ -2136,7 +2197,9 @@ def run_mlx_training_process(
if not transformers_activated:
# Must precede detect_hardware(): its MLX stack check imports mlx_lm, hence transformers.
_activate_transformers_version_or_warn(model_name, config.get("hf_token") or None)
_activate_transformers_version_or_warn(
model_name, config.get("hf_token") or None
)
from utils.hardware import hardware as _hw
@ -2260,14 +2323,19 @@ def run_training_process(*, event_queue: Any, stop_queue: Any, config: dict) ->
if backend_path not in sys.path:
sys.path.insert(0, backend_path)
from .training import is_apple_silicon_training_platform, should_use_mlx_training_backend
from .training import (
is_apple_silicon_training_platform,
should_use_mlx_training_backend,
)
mlx_backend_requested = is_apple_silicon_training_platform()
mlx_transformers_activated = False
if mlx_backend_requested and _is_current_process_apple_silicon():
# Must precede detect_hardware(): its MLX stack check imports mlx_lm, hence transformers.
_activate_transformers_version_or_warn(model_name, config.get("hf_token") or None)
_activate_transformers_version_or_warn(
model_name, config.get("hf_token") or None
)
mlx_transformers_activated = True
from utils.hardware import hardware as _hw
@ -2328,7 +2396,9 @@ def run_training_process(*, event_queue: Any, stop_queue: Any, config: dict) ->
from utils.models.model_config import get_base_model_from_lora_identifier
# Resolve a LOCAL or REMOTE adapter's base so a remote LoRA base is gated too.
_base = get_base_model_from_lora_identifier(model_name, config.get("hf_token") or None)
_base = get_base_model_from_lora_identifier(
model_name, config.get("hf_token") or None
)
if _base:
malware_targets.append(_base)
except Exception as exc:
@ -2532,7 +2602,9 @@ def run_training_process(*, event_queue: Any, stop_queue: Any, config: dict) ->
if os.path.isfile(os.path.join(_scripts_dir, "hipInfo.exe")):
import shutil as _shutil
if not _shutil.which("hipinfo.exe"):
os.environ["PATH"] = _scripts_dir + os.pathsep + os.environ.get("PATH", "")
os.environ["PATH"] = (
_scripts_dir + os.pathsep + os.environ.get("PATH", "")
)
# BNB picks a rocm DLL from torch.version.hip, but AMD's Windows BNB
# wheel may ship a DLL whose suffix doesn't match. Detect the actual
@ -2573,7 +2645,9 @@ def run_training_process(*, event_queue: Any, stop_queue: Any, config: dict) ->
# so later import fixes can still redetect or opt out. DLL
# with unparsable name -> seeded value or "72".
if _found_rocm_bnb:
_bnb_rocm_ver = _bnb_rocm_ver or os.environ.get("BNB_ROCM_VERSION") or "72"
_bnb_rocm_ver = (
_bnb_rocm_ver or os.environ.get("BNB_ROCM_VERSION") or "72"
)
os.environ["BNB_ROCM_VERSION"] = _bnb_rocm_ver
os.environ["UNSLOTH_BNB_ROCM_VERSION_SOURCE"] = "detected"
logger.info(
@ -2596,7 +2670,9 @@ def run_training_process(*, event_queue: Any, stop_queue: Any, config: dict) ->
# the rocm version embedded in torch.__version__ when version.hip is
# unset (AMD SDK / Radeon wheels).
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.
_ver_match = re.search(r"rocm(\d+)\.(\d+)", _build_version_for_rocm)
@ -2688,7 +2764,9 @@ def run_training_process(*, event_queue: Any, stop_queue: Any, config: dict) ->
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 = (
@ -2860,7 +2938,11 @@ def run_training_process(*, event_queue: Any, stop_queue: Any, config: dict) ->
def _on_progress(progress: TrainingProgress):
has_train_loss = progress.step > 0 and progress.loss is not None
has_eval_loss = progress.eval_loss is not None
if (progress.step == 0 and progress.total_steps > 0) or has_train_loss or has_eval_loss:
if (
(progress.step == 0 and progress.total_steps > 0)
or has_train_loss
or has_eval_loss
):
event_queue.put(
{
"type": "progress",
@ -2971,12 +3053,15 @@ def run_training_process(*, event_queue: Any, stop_queue: Any, config: dict) ->
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(),
}
@ -2997,7 +3082,9 @@ def run_training_process(*, event_queue: Any, stop_queue: Any, config: dict) ->
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)
@ -3053,7 +3140,9 @@ def run_training_process(*, event_queue: Any, stop_queue: Any, config: dict) ->
event_queue.put({"type": "model_load_completed", "ts": time.time()})
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(
@ -3094,7 +3183,9 @@ def run_training_process(*, event_queue: Any, stop_queue: Any, config: dict) ->
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),
)
@ -3104,13 +3195,17 @@ def run_training_process(*, event_queue: Any, stop_queue: Any, config: dict) ->
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),
)
@ -3120,12 +3215,15 @@ def run_training_process(*, event_queue: Any, stop_queue: Any, config: dict) ->
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(),
}
@ -3184,7 +3282,9 @@ def run_training_process(*, event_queue: Any, stop_queue: Any, config: dict) ->
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}"'
@ -3208,7 +3308,9 @@ def run_training_process(*, event_queue: Any, stop_queue: Any, config: dict) ->
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"),
@ -3393,7 +3495,9 @@ def _run_embedding_training(event_queue: Any, stop_queue: Any, config: dict) ->
for target in dict.fromkeys(malware_targets):
_fs = evaluate_file_security(
target, hf_token = hf_token, load_subdirs = security_load_subdirs(target, hf_token)
target,
hf_token = hf_token,
load_subdirs = security_load_subdirs(target, hf_token),
)
if _fs.blocked:
event_queue.put(
@ -3414,7 +3518,9 @@ def _run_embedding_training(event_queue: Any, stop_queue: Any, config: dict) ->
consent_targets = [model_name]
try:
from utils.models.model_config import get_base_model_from_lora_identifier
from utils.models.model_config import (
get_base_model_from_lora_identifier,
)
_cbase = get_base_model_from_lora_identifier(model_name, hf_token)
if _cbase:
consent_targets.append(_cbase)
@ -3541,7 +3647,9 @@ 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)

14
studio/backend/hub/routes/datasets.py
View file

@ -61,14 +61,17 @@ async def list_cached_datasets(current_subject: str = Depends(get_current_subjec
@router.delete("/cached", response_model = DeleteCachedDatasetResponse)
async def delete_cached_dataset(
repo_id: str = Body(..., embed = True), current_subject: str = Depends(get_current_subject)
repo_id: str = Body(..., embed = True),
current_subject: str = Depends(get_current_subject),
):
return await cache_inventory.delete_cached_dataset_response(repo_id)
@router.get("/download-progress", response_model = DownloadProgressResponse)
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'"
),
expected_bytes: int = Query(0, description = "Expected total download size in bytes"),
hf_token: Optional[str] = Depends(get_hf_token),
current_subject: str = Depends(get_current_subject),
@ -89,9 +92,12 @@ async def download_dataset(
return await downloads.download_dataset_response(body, hf_token)
@router.post("/download/cancel", response_model = CancelDatasetDownloadResponse, status_code = 202)
@router.post(
"/download/cancel", response_model = CancelDatasetDownloadResponse, status_code = 202
)
async def cancel_dataset_download(
body: CancelDatasetDownloadRequest, current_subject: str = Depends(get_current_subject)
body: CancelDatasetDownloadRequest,
current_subject: str = Depends(get_current_subject),
):
return await downloads.cancel_dataset_download_response(body)

8
studio/backend/hub/routes/inventory.py
View file

@ -136,7 +136,9 @@ async def cancel_download_model(
@router.get("/download-status", response_model = DownloadJobStatus)
async def get_download_status(
repo_id: str = Query(..., description = "HuggingFace repo ID"),
gguf_variant: str = Query("", description = "Quantization variant (empty for safetensors)"),
gguf_variant: str = Query(
"", description = "Quantization variant (empty for safetensors)"
),
current_subject: str = Depends(get_current_subject),
):
return await downloads.get_download_status_response(repo_id, gguf_variant)
@ -153,7 +155,9 @@ async def get_active_downloads(
@router.get("/transport-status", response_model = TransportStatusResponse)
async def get_model_transport_status(
repo_id: str = Query(..., description = "HuggingFace repo ID"),
gguf_variant: str = Query("", description = "Quantization variant (empty for safetensors)"),
gguf_variant: str = Query(
"", description = "Quantization variant (empty for safetensors)"
),
hf_token: Optional[str] = Depends(get_hf_token),
current_subject: str = Depends(get_current_subject),
):

4
studio/backend/hub/schemas/downloads.py
View file

@ -7,7 +7,9 @@ from pydantic import BaseModel, Field
from typing import List, Literal, Optional
DownloadJobState = Literal["idle", "running", "cancelling", "cancelled", "complete", "error"]
DownloadJobState = Literal[
"idle", "running", "cancelling", "cancelled", "complete", "error"
]
class DownloadModelRequest(BaseModel):

16
studio/backend/hub/schemas/inventory.py
View file

@ -17,13 +17,19 @@ ModelRuntime = Literal["llama_cpp", "transformers", "adapter", "unknown"]
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')")
quant: str = Field(..., description = "Quantization label or internal GGUF variant key")
filename: str = Field(
..., description = "GGUF filename (e.g., 'gemma-3-4b-it-Q4_K_M.gguf')"
)
quant: str = Field(
..., description = "Quantization label or internal GGUF variant key"
)
display_label: Optional[str] = Field(
None, description = "Optional user-facing label when quant is an internal key"
)
size_bytes: int = Field(0, description = "File size in bytes")
download_size_bytes: int = Field(0, description = "Total bytes needed to download this variant")
download_size_bytes: int = Field(
0, description = "Total bytes needed to download this variant"
)
downloaded: bool = Field(
False, description = "Whether this variant is already in the local HF cache"
)
@ -135,7 +141,9 @@ 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",

4
studio/backend/hub/services/__init__.py
View file

@ -12,7 +12,9 @@ from fastapi import HTTPException
from hub.utils.hf_cache_state import resolve_destructive_case_matches
def resolve_destructive_repo_ids(repo_id: str, candidates: Iterable[str], *, noun: str) -> set[str]:
def resolve_destructive_repo_ids(
repo_id: str, candidates: Iterable[str], *, noun: str
) -> set[str]:
"""Cache-dir repo ids a destructive op on *repo_id* may target.
Refuses with 409 on ambiguous case-only matches so a delete never removes

30
studio/backend/hub/services/datasets/cache_inventory.py
View file

@ -194,7 +194,9 @@ def _hf_datasets_cache_roots() -> list[Path]:
if hf_home:
_add(Path(hf_home).expanduser() / "datasets")
xdg_cache = Path(os.environ.get("XDG_CACHE_HOME", Path.home() / ".cache")).expanduser()
xdg_cache = Path(
os.environ.get("XDG_CACHE_HOME", Path.home() / ".cache")
).expanduser()
_add(xdg_cache / "huggingface" / "datasets")
return roots
@ -284,7 +286,11 @@ def _scan_hf_dataset_caches() -> list[dict]:
rev_id = getattr(rev, "commit_hash", None) or str(id(rev))
for f in rev.files:
blob_path = getattr(f, "blob_path", None)
key = str(blob_path) if blob_path else f"{rev_id}:{f.file_name}"
key = (
str(blob_path)
if blob_path
else f"{rev_id}:{f.file_name}"
)
unique_blobs[key] = int(f.size_on_disk or 0)
total_size = sum(unique_blobs.values())
key = repo_info.repo_id.lower()
@ -320,7 +326,9 @@ def _scan_hf_dataset_caches() -> list[dict]:
existing = seen_lower.get(key)
if _prefer_dataset_cache_row(row, existing):
seen_lower[key] = row
elif existing is not None and bool(existing.get("partial")) == bool(row.get("partial")):
elif existing is not None and bool(existing.get("partial")) == bool(
row.get("partial")
):
existing["size_bytes"] = max(existing["size_bytes"], row["size_bytes"])
existing["cache_path"] = existing.get("cache_path") or row.get("cache_path")
if (
@ -332,7 +340,9 @@ def _scan_hf_dataset_caches() -> list[dict]:
for row in _scan_processed_dataset_caches():
key = row["repo_id"].lower()
existing = seen_lower.get(key)
if existing is None or (bool(existing.get("partial")) and not bool(row.get("partial"))):
if existing is None or (
bool(existing.get("partial")) and not bool(row.get("partial"))
):
seen_lower[key] = row
else:
existing["size_bytes"] = max(existing["size_bytes"], row["size_bytes"])
@ -366,7 +376,9 @@ async def delete_cached_dataset_response(repo_id: str) -> dict:
if not _is_valid_repo_id(repo_id):
raise HTTPException(status_code = 400, detail = "Invalid repo_id format")
repo_key = await asyncio.to_thread(resolve_cached_repo_id_case, repo_id, repo_type = "dataset")
repo_key = await asyncio.to_thread(
resolve_cached_repo_id_case, repo_id, repo_type = "dataset"
)
if not downloads.registry.begin_delete(repo_key):
raise HTTPException(
status_code = 400,
@ -401,7 +413,9 @@ def _delete_cached_dataset_blocking(repo_id: str) -> dict:
if str(repo_info.repo_id) not in matched_repo_ids:
continue
try:
strategy = hf_cache.delete_revisions(*(rev.commit_hash for rev in repo_info.revisions))
strategy = hf_cache.delete_revisions(
*(rev.commit_hash for rev in repo_info.revisions)
)
strategy.execute()
deleted = True
except Exception as exc:
@ -430,7 +444,9 @@ def _delete_cached_dataset_blocking(repo_id: str) -> dict:
cache_purged = purge_repo_cache_dirs("dataset", repo_id)
partial_purged = purge_partial_repo("dataset", repo_id)
state_purged = download_manifest.purge_all_state_for_repo("dataset", repo_id) > 0
if not (deleted or processed_deleted or cache_purged or partial_purged or state_purged):
if not (
deleted or processed_deleted or cache_purged or partial_purged or state_purged
):
raise HTTPException(status_code = 404, detail = "Dataset not found in cache")
return {"status": "deleted", "repo_id": repo_id}

40
studio/backend/hub/services/datasets/downloads.py
View file

@ -37,9 +37,7 @@ from hub.utils.snapshot_filters import (
logger = get_logger(__name__)
_dataset_size_cache: "OrderedDict[str, tuple[int, frozenset[str], bool, str, float]]" = (
OrderedDict()
)
_dataset_size_cache: "OrderedDict[str, tuple[int, frozenset[str], bool, str, float]]" = OrderedDict()
_dataset_size_neg_cache: "OrderedDict[tuple[str, str], float]" = OrderedDict()
_DATASET_SIZE_CACHE_MAX = 256
_DATASET_SIZE_POS_TTL = 60.0
@ -88,7 +86,9 @@ def get_dataset_snapshot_metadata_cached(
)
total = total_size_for_siblings(info.siblings)
hashes = blob_hashes_for_siblings(info.siblings)
restricted = bool(getattr(info, "private", False) or getattr(info, "gated", False))
restricted = bool(
getattr(info, "private", False) or getattr(info, "gated", False)
)
except Exception:
with _dataset_size_cache_lock:
_dataset_size_neg_cache[cache_key] = time.monotonic()
@ -132,7 +132,9 @@ async def get_dataset_download_progress_response(
)
def _dataset_status(key: str, *, repo_id: Optional[str] = None) -> DatasetDownloadJobStatus:
def _dataset_status(
key: str, *, repo_id: Optional[str] = None
) -> DatasetDownloadJobStatus:
state, error, generation = download_lifecycle.idle_status(
_registry,
key,
@ -154,7 +156,9 @@ async def download_dataset_response(
detail = f"Invalid repo_id: {repo_id!r}",
)
# Canonicalize so two different-cased paste-ins share one job + cache dir.
repo_id = await asyncio.to_thread(resolve_cached_repo_id_case, repo_id, repo_type = "dataset")
repo_id = await asyncio.to_thread(
resolve_cached_repo_id_case, repo_id, repo_type = "dataset"
)
key = _download_job_key(repo_id)
use_xet = download_lifecycle.resolve_effective_use_xet(body.use_xet)
@ -212,7 +216,9 @@ async def cancel_dataset_download_response(body: CancelDatasetDownloadRequest) -
status_code = 400,
detail = f"Invalid repo_id: {repo_id!r}",
)
repo_id = await asyncio.to_thread(resolve_cached_repo_id_case, repo_id, repo_type = "dataset")
repo_id = await asyncio.to_thread(
resolve_cached_repo_id_case, repo_id, repo_type = "dataset"
)
key = _download_job_key(repo_id)
state = download_lifecycle.cancel_worker(
@ -225,21 +231,29 @@ async def cancel_dataset_download_response(body: CancelDatasetDownloadRequest) -
return {"repo_id": repo_id, "state": state}
async def get_dataset_download_status_response(repo_id: str) -> DatasetDownloadJobStatus:
async def get_dataset_download_status_response(
repo_id: str,
) -> DatasetDownloadJobStatus:
"""Return the latest state of a background dataset download job."""
repo_id = repo_id.strip()
if not _is_valid_repo_id(repo_id):
return DatasetDownloadJobStatus(state = "idle")
repo_id = await asyncio.to_thread(resolve_cached_repo_id_case, repo_id, repo_type = "dataset")
repo_id = await asyncio.to_thread(
resolve_cached_repo_id_case, repo_id, repo_type = "dataset"
)
return _dataset_status(_download_job_key(repo_id), repo_id = repo_id)
async def get_active_dataset_downloads_response(repo_id: str = "") -> ActiveDownloadsResponse:
async def get_active_dataset_downloads_response(
repo_id: str = "",
) -> ActiveDownloadsResponse:
repo_id = repo_id.strip()
if repo_id and not _is_valid_repo_id(repo_id):
return ActiveDownloadsResponse(downloads = [])
canonical_repo_id = (
await asyncio.to_thread(resolve_cached_repo_id_case, repo_id, repo_type = "dataset")
await asyncio.to_thread(
resolve_cached_repo_id_case, repo_id, repo_type = "dataset"
)
if repo_id
else None
)
@ -261,7 +275,9 @@ async def get_dataset_transport_status_response(repo_id: str) -> dict:
return {"has_partial": False, "last_transport": None, "resumable": False}
return {
"has_partial": has_active_incomplete_blobs("dataset", repo_id),
"last_transport": download_registry.read_active_transport_marker("dataset", repo_id),
"last_transport": download_registry.read_active_transport_marker(
"dataset", repo_id
),
"resumable": download_registry.is_resumable_partial("dataset", repo_id),
}

19
studio/backend/hub/services/datasets/formatting.py
View file

@ -177,10 +177,14 @@ def _repo_file_matches_split(path: str, split: str) -> bool:
def _select_tier1_repo_file(
files: list[str], *, subset: Optional[str], train_split: str
) -> Optional[str]:
data_files = sorted(f for f in files if any(f.lower().endswith(ext) for ext in DATA_EXTS))
data_files = sorted(
f for f in files if any(f.lower().endswith(ext) for ext in DATA_EXTS)
)
if not data_files:
return None
tabular_files = [f for f in data_files if any(f.lower().endswith(ext) for ext in _TABULAR_EXTS)]
tabular_files = [
f for f in data_files if any(f.lower().endswith(ext) for ext in _TABULAR_EXTS)
]
candidates = tabular_files or data_files
if subset:
candidates = [f for f in candidates if _repo_file_matches_label(f, subset)]
@ -405,7 +409,9 @@ def check_format_response(
processed = format_dataset_preview(preview_slice)
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)
@ -416,7 +422,9 @@ def check_format_response(
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."
@ -483,7 +491,8 @@ def ai_assist_mapping_response(
from hub.utils.llm_assist import llm_conversion_advisor
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(

17
studio/backend/hub/services/datasets/local.py
View file

@ -223,7 +223,9 @@ def _stream_file_preview_slice(path: Path, preview_size: int):
return Dataset.from_list(rows), None
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
):
# Non-streaming loads take the cached builder lock; use the EACCES-safe wrapper.
from utils.datasets.cache_safe import load_dataset_cache_safe as load_dataset
@ -258,7 +260,9 @@ def _load_local_preview_slice(*, dataset_path: Path, train_split: str, preview_s
# Parquet/Arrow give a cheap exact total_rows via len()+select; JSON/CSV
# carry no such metadata, so stream them and report total_rows=None.
if 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
)
total_rows = len(dataset)
preview_slice = dataset.select(range(min(preview_size, total_rows)))
return preview_slice, total_rows
@ -272,7 +276,9 @@ def _load_local_preview_slice(*, dataset_path: Path, train_split: str, preview_s
)
return preview
raise HTTPException(status_code = 400, detail = f"Unsupported file format: {dataset_path.suffix}")
raise HTTPException(
status_code = 400, detail = f"Unsupported file format: {dataset_path.suffix}"
)
def _sanitize_filename(filename: str) -> str:
@ -285,7 +291,10 @@ def _sanitize_filename(filename: str) -> str:
def _upload_too_large(size_bytes: int) -> HTTPException:
return HTTPException(
status_code = 413,
detail = (f"Upload is too large " f"({size_bytes:,} bytes; max {LOCAL_UPLOAD_MAX_BYTES:,})."),
detail = (
f"Upload is too large "
f"({size_bytes:,} bytes; max {LOCAL_UPLOAD_MAX_BYTES:,})."
),
)

45
studio/backend/hub/services/download_lifecycle.py
View file

@ -44,8 +44,12 @@ def resolve_effective_use_xet(use_xet: bool) -> bool:
def resolve_transport(use_xet: bool) -> str:
transport = download_registry.TRANSPORT_XET if use_xet else download_registry.TRANSPORT_HTTP
unavailable_reason = download_registry.download_transport_unavailable_reason(transport)
transport = (
download_registry.TRANSPORT_XET if use_xet else download_registry.TRANSPORT_HTTP
)
unavailable_reason = download_registry.download_transport_unavailable_reason(
transport
)
if unavailable_reason is not None:
raise HTTPException(status_code = 400, detail = unavailable_reason)
return transport
@ -67,7 +71,9 @@ def spawn_worker(
shared ``.incomplete`` (e.g. bundled mmproj) is never deleted.
"""
cwd = backend_dir()
mode = download_registry.TRANSPORT_XET if use_xet else download_registry.TRANSPORT_HTTP
mode = (
download_registry.TRANSPORT_XET if use_xet else download_registry.TRANSPORT_HTTP
)
env = os.environ.copy()
if protected_blob_hashes:
env["UNSLOTH_PROTECTED_BLOB_HASHES"] = ",".join(sorted(protected_blob_hashes))
@ -95,7 +101,9 @@ def spawn_worker(
if hf_token:
env["HF_TOKEN"] = hf_token
existing_path = env.get("PYTHONPATH", "")
env["PYTHONPATH"] = f"{cwd}{os.pathsep}{existing_path}" if existing_path else str(cwd)
env["PYTHONPATH"] = (
f"{cwd}{os.pathsep}{existing_path}" if existing_path else str(cwd)
)
return subprocess.Popen(
[
sys.executable,
@ -242,7 +250,9 @@ def finalize_worker_exit(
f"{label}: {stderr_text}"
)
else:
logger.info(f"{log_prefix} worker diagnostics for {label}: {stderr_text}")
logger.info(
f"{log_prefix} worker diagnostics for {label}: {stderr_text}"
)
logger.info(f"{log_prefix} complete: {label}")
# Defensive cleanup: the canonical clear is at download-start; this
# catches the rare case where that failed but the download succeeded.
@ -299,7 +309,9 @@ def _set_retry_failure_state(
download_registry.persist_cancel_marker(
repo_type,
repo_id,
metadata.variant if metadata is not None and metadata.variant else fallback_variant,
metadata.variant
if metadata is not None and metadata.variant
else fallback_variant,
metadata.transport
if metadata is not None and metadata.transport
else fallback_transport,
@ -333,7 +345,9 @@ def _try_http_retry(
"""
original_metadata = registry.get_job_metadata(key)
if original_metadata is None:
logger.debug("%s XET retry skipped for %s; metadata unavailable", log_prefix, label)
logger.debug(
"%s XET retry skipped for %s; metadata unavailable", log_prefix, label
)
_set_retry_failure_state(
registry,
key,
@ -580,11 +594,15 @@ def register_worker(
try:
kill_and_reap_process(proc, label = label, logger = logger)
except Exception:
logger.exception("failed to reap worker after watcher crash for %s", key)
logger.exception(
"failed to reap worker after watcher crash for %s", key
)
try:
registry.drop_process(key, proc)
except Exception:
logger.exception("failed to drop worker after watcher crash for %s", key)
logger.exception(
"failed to drop worker after watcher crash for %s", key
)
try:
registry.set_job(key, "error", "download watcher crashed")
except Exception:
@ -718,13 +736,18 @@ def idle_status(
def active_download_refs(
registry: download_registry.DownloadRegistry, repo_id: Optional[str], *, with_variant: bool
registry: download_registry.DownloadRegistry,
repo_id: Optional[str],
*,
with_variant: bool,
) -> list[ActiveDownload]:
downloads: list[ActiveDownload] = []
for ref in registry.active_job_refs(repo_id):
metadata = ref.metadata
if with_variant:
ref_repo_id = metadata.repo_id if metadata is not None else ref.key.split("::", 1)[0]
ref_repo_id = (
metadata.repo_id if metadata is not None else ref.key.split("::", 1)[0]
)
if metadata is not None:
variant = metadata.variant
else:

28
studio/backend/hub/services/models/cache_inventory.py
View file

@ -46,9 +46,7 @@ from utils.hidden_models import is_hidden_model
logger = get_logger(__name__)
_repo_size_cache: "OrderedDict[tuple[str, str, str], tuple[int, frozenset[str], float]]" = (
OrderedDict()
)
_repo_size_cache: "OrderedDict[tuple[str, str, str], tuple[int, frozenset[str], float]]" = OrderedDict()
_repo_size_neg_cache: "OrderedDict[tuple[str, str, str], float]" = OrderedDict()
_REPO_SIZE_CACHE_MAX = 256
_REPO_SIZE_POS_TTL = 60.0
@ -138,7 +136,9 @@ def _cached_repo_file_name(file_obj) -> str:
try:
path = Path(file_path)
parts = path.parts
snapshots_idx = max(i for i, part in enumerate(parts) if part == "snapshots")
snapshots_idx = max(
i for i, part in enumerate(parts) if part == "snapshots"
)
if len(parts) > snapshots_idx + 2:
return Path(*parts[snapshots_idx + 2 :]).as_posix()
except Exception:
@ -155,7 +155,9 @@ def _is_real_cache_blob(blob: Optional[Path], repo_dir: Optional[Path]) -> bool:
if blob is None or repo_dir is None:
return False
try:
return blob.parent.resolve(strict = False) == (repo_dir / "blobs").resolve(strict = False)
return blob.parent.resolve(strict = False) == (repo_dir / "blobs").resolve(
strict = False
)
except OSError:
return False
@ -180,7 +182,9 @@ def local_size_identity(size: int) -> str:
return f"{_LOCAL_SIZE_IDENTITY_PREFIX}{int(size)}"
def _repo_gguf_blob_map(repo_info, *, include_companions: bool = False) -> dict[str, set[str]]:
def _repo_gguf_blob_map(
repo_info, *, include_companions: bool = False
) -> dict[str, set[str]]:
"""Map each cached GGUF file's repo-relative name to the SET of its local
identities across all revisions.
@ -271,7 +275,9 @@ def _scan_cached_gguf() -> list[dict]:
repo_path = Path(repo_info.repo_path)
snapshot_path = _cached_model_snapshot_path(repo_path)
total_size = _repo_gguf_size_bytes(repo_info)
has_variant_state, variant_state_size = _gguf_variant_state_summary(repo_id)
has_variant_state, variant_state_size = _gguf_variant_state_summary(
repo_id
)
is_hidden_infra = _is_hidden_infra_repo(
repo_id,
str(repo_path),
@ -354,7 +360,9 @@ def _repo_non_gguf_model_payload(repo_info) -> _CachedNonGgufPayload:
has_transformers_safetensors = False
has_checkpoint = False
def _record_blob(target: dict[str, int], file_obj, rev_id: str, file_name: str) -> None:
def _record_blob(
target: dict[str, int], file_obj, rev_id: str, file_name: str
) -> None:
blob_path = getattr(file_obj, "blob_path", None)
size = int(file_obj.size_on_disk or 0)
key = str(blob_path) if blob_path else f"{rev_id}:{file_name}"
@ -482,7 +490,9 @@ def _cached_model_local_metadata(repo_path: Path) -> dict:
result["library_name"] = library_name.strip()
tags = card.get("tags")
if isinstance(tags, list):
clean_tags = [tag.strip() for tag in tags if isinstance(tag, str) and tag.strip()]
clean_tags = [
tag.strip() for tag in tags if isinstance(tag, str) and tag.strip()
]
if clean_tags:
result["tags"] = clean_tags
return result

38
studio/backend/hub/services/models/common.py
View file

@ -84,7 +84,9 @@ 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())
@ -193,7 +195,9 @@ def _apply_format_aware_partial(
continue
# GGUF row-level transport is ambiguous (variants may differ); per-variant
# detail lives on GgufVariantDetail.partial_transport via the variants endpoint.
partial_transport = None if row.model_format == "gguf" else snapshot_partial_transport
partial_transport = (
None if row.model_format == "gguf" else snapshot_partial_transport
)
rewritten.append(
row.model_copy(
update = {
@ -217,7 +221,9 @@ def _weight_basename(name: str) -> str:
def _is_adapter_weight_name(name: str) -> bool:
lower = _weight_basename(name)
return lower.startswith("adapter_model") and lower.endswith((".safetensors", ".bin"))
return lower.startswith("adapter_model") and lower.endswith(
(".safetensors", ".bin")
)
def _is_transformers_safetensors_weight_name(name: str) -> bool:
@ -267,7 +273,9 @@ def _classify_non_gguf_model_format(
has_checkpoint_weights: bool,
trusted_hf_cache_repo: bool = False,
) -> Optional[ModelFormat]:
if has_safetensors and (has_config or (trusted_hf_cache_repo and has_transformers_safetensors)):
if has_safetensors and (
has_config or (trusted_hf_cache_repo and has_transformers_safetensors)
):
return "safetensors"
if has_adapter_config and has_adapter_weights:
return "adapter"
@ -278,7 +286,9 @@ def _classify_non_gguf_model_format(
def _is_main_gguf_filename(name: str) -> bool:
return (
_is_gguf_filename(name) and not _is_mmproj_filename(name) and not _is_mtp_drafter_path(name)
_is_gguf_filename(name)
and not _is_mmproj_filename(name)
and not _is_mtp_drafter_path(name)
)
@ -445,7 +455,8 @@ def _local_model_info(
),
load_id = load_id,
model_id = model_id,
display_name = display_name or (scan_path.stem if scan_path.is_file() else scan_path.name),
display_name = display_name
or (scan_path.stem if scan_path.is_file() else scan_path.name),
path = str(load_path),
size_bytes = max(0, int(size_bytes or 0)),
source = source,
@ -520,12 +531,17 @@ def _classify_local_path(
(scan_path / "adapter_config.json").is_file() if scan_path.is_dir() else False
)
adapter_config = _read_adapter_config(scan_path) if has_adapter_config else {}
adapter_base_model = _clean_optional_string(adapter_config.get("base_model_name_or_path"))
adapter_base_model = _clean_optional_string(
adapter_config.get("base_model_name_or_path")
)
adapter_type = _clean_optional_string(adapter_config.get("peft_type"))
training_method = _clean_optional_string(adapter_config.get("unsloth_training_method"))
training_method = _clean_optional_string(
adapter_config.get("unsloth_training_method")
)
has_adapter_weights = any(_is_adapter_weight_file(f) for f in files)
has_safetensors = any(
f.suffix.lower() == ".safetensors" and not _is_adapter_weight_file(f) for f in files
f.suffix.lower() == ".safetensors" and not _is_adapter_weight_file(f)
for f in files
)
has_transformers_safetensors = any(
_is_transformers_safetensors_weight_file(f) and not _is_adapter_weight_file(f)
@ -554,7 +570,9 @@ def _classify_local_path(
if f.suffix.lower() == ".safetensors" and not _is_adapter_weight_file(f)
)
else:
size_bytes = _sum_file_sizes(f for f in files if _is_checkpoint_weight_file(f))
size_bytes = _sum_file_sizes(
f for f in files if _is_checkpoint_weight_file(f)
)
rows.append(
_local_model_info(
scan_path = scan_path,

75
studio/backend/hub/services/models/deletion.py
View file

@ -76,7 +76,9 @@ def _path_exists_or_symlink(path: Path) -> bool:
return False
def _repo_file_matches(target_repo, predicate) -> list[tuple[Path, Optional[Path], str]]:
def _repo_file_matches(
target_repo, predicate
) -> list[tuple[Path, Optional[Path], str]]:
matches: list[tuple[Path, Optional[Path], str]] = []
for rev in getattr(target_repo, "revisions", ()):
for f in getattr(rev, "files", ()):
@ -107,7 +109,9 @@ def _has_remaining_main_gguf(target_repo) -> bool:
)
def _remove_empty_variant_dirs(target_repos: list, variant: str) -> tuple[int, list[str]]:
def _remove_empty_variant_dirs(
target_repos: list, variant: str
) -> tuple[int, list[str]]:
"""Remove now-empty ``snapshots/<rev>/<quant>/`` folders for *variant* (the
quant label names the folder); only empty dirs go, so siblings are safe.
Returns (count removed, removal failures other than a concurrent refill)."""
@ -122,7 +126,9 @@ def _remove_empty_variant_dirs(target_repos: list, variant: str) -> tuple[int, l
if not snapshots.is_dir():
continue
try:
snap_dirs = [s for s in snapshots.iterdir() if s.is_dir() and not s.is_symlink()]
snap_dirs = [
s for s in snapshots.iterdir() if s.is_dir() and not s.is_symlink()
]
except OSError:
continue
for snap in snap_dirs:
@ -164,7 +170,9 @@ def _remove_empty_snapshot_dirs(target_repos: list) -> tuple[int, list[str]]:
if not snapshots.is_dir():
continue
try:
snap_dirs = [s for s in snapshots.iterdir() if s.is_dir() and not s.is_symlink()]
snap_dirs = [
s for s in snapshots.iterdir() if s.is_dir() and not s.is_symlink()
]
except OSError:
continue
for snap in snap_dirs:
@ -192,7 +200,11 @@ def _delete_gguf_variant_from_repos(
completed_hashes: set[str] = set()
for target_repo in target_repos:
repo_dir = Path(target_repo.repo_path) if getattr(target_repo, "repo_path", None) else None
repo_dir = (
Path(target_repo.repo_path)
if getattr(target_repo, "repo_path", None)
else None
)
matched = _repo_file_matches(
target_repo,
lambda name: _is_main_gguf_filename(name)
@ -400,7 +412,11 @@ def reclaim_replaced_gguf_variant(
]
for target_repo in target_repos:
repo_dir = Path(target_repo.repo_path) if getattr(target_repo, "repo_path", None) else None
repo_dir = (
Path(target_repo.repo_path)
if getattr(target_repo, "repo_path", None)
else None
)
stale_matches: list[tuple[Path, Optional[Path], str]] = []
matches = _repo_file_matches(
target_repo,
@ -500,7 +516,10 @@ def _loaded_id_matches_repo(loaded_id: str, repo_id: str) -> bool:
def _loaded_repo_variant_blocks_delete(
loaded_id: str, repo_id: str, delete_variant: Optional[str], loaded_variant: Optional[str]
loaded_id: str,
repo_id: str,
delete_variant: Optional[str],
loaded_variant: Optional[str],
) -> bool:
if not _loaded_id_matches_repo(loaded_id, repo_id):
return False
@ -523,7 +542,9 @@ def _llama_cpp_blocks_delete(repo_id: str, variant: Optional[str]) -> bool:
from routes.inference import get_llama_cpp_backend
backend = get_llama_cpp_backend()
except Exception as e:
logger.debug(f"llama.cpp backend unavailable during delete guard for {repo_id}: {e}")
logger.debug(
f"llama.cpp backend unavailable during delete guard for {repo_id}: {e}"
)
return False
loaded_id = backend.model_identifier
loaded_variant = getattr(backend, "hf_variant", None)
@ -550,7 +571,9 @@ def _inference_backend_blocks_delete(repo_id: str) -> bool:
from core.inference import get_inference_backend
backend = get_inference_backend()
except Exception as e:
logger.debug(f"Inference backend unavailable during delete guard for {repo_id}: {e}")
logger.debug(
f"Inference backend unavailable during delete guard for {repo_id}: {e}"
)
return False
active_name = backend.active_model_name
return bool(active_name) and _loaded_id_matches_repo(active_name, repo_id)
@ -583,7 +606,9 @@ async def delete_cached_model_response(
_inference_backend_blocks_delete(repo_id)
)
except Exception as e:
logger.warning(f"Load-state verification failed for {repo_id}; refusing delete: {e}")
logger.warning(
f"Load-state verification failed for {repo_id}; refusing delete: {e}"
)
raise HTTPException(
status_code = 503,
detail = _LOAD_STATE_UNVERIFIABLE_DETAIL,
@ -594,7 +619,9 @@ async def delete_cached_model_response(
detail = "Unload the model before deleting",
)
repo_key = await asyncio.to_thread(resolve_cached_repo_id_case, repo_id, repo_type = "model")
repo_key = await asyncio.to_thread(
resolve_cached_repo_id_case, repo_id, repo_type = "model"
)
if not downloads.registry.begin_delete(repo_key, variant):
detail = (
f"Cancel the {variant} download before deleting it."
@ -603,7 +630,9 @@ async def delete_cached_model_response(
)
raise HTTPException(status_code = 400, detail = detail)
try:
return await asyncio.to_thread(_delete_cached_model_blocking, repo_id, variant, hf_token)
return await asyncio.to_thread(
_delete_cached_model_blocking, repo_id, variant, hf_token
)
finally:
downloads.registry.end_delete(repo_key, variant)
cache_inventory.invalidate_hf_cache_scans()
@ -642,18 +671,22 @@ def _delete_cached_model_blocking(
if not target_entries:
if variant is None:
cache_purged = purge_repo_cache_dirs("model", repo_id) or purge_partial_repo(
cache_purged = purge_repo_cache_dirs(
"model", repo_id
) or purge_partial_repo("model", repo_id)
state_purged = (
download_manifest.purge_all_state_for_repo("model", repo_id) > 0
)
state_purged = download_manifest.purge_all_state_for_repo("model", repo_id) > 0
if cache_purged or state_purged:
return {"status": "deleted", "repo_id": repo_id}
if variant:
incomplete_result = gguf_variants.delete_variant_incomplete_blobs_result(
repo_id,
variant,
hf_token,
companions = not sibling_active,
incomplete_result = (
gguf_variants.delete_variant_incomplete_blobs_result(
repo_id,
variant,
hf_token,
companions = not sibling_active,
)
)
if incomplete_result.unresolved:
raise HTTPException(
@ -689,7 +722,9 @@ def _delete_cached_model_blocking(
deleted_revisions = False
for hf_cache, repo_info in target_entries:
revision_hashes = [
rev.commit_hash for rev in repo_info.revisions if getattr(rev, "commit_hash", None)
rev.commit_hash
for rev in repo_info.revisions
if getattr(rev, "commit_hash", None)
]
if not revision_hashes:
continue

39
studio/backend/hub/services/models/downloads.py
View file

@ -102,7 +102,9 @@ def _spawn_download_worker(
)
async def download_model_response(body: DownloadModelRequest, hf_token: Optional[str] = None):
async def download_model_response(
body: DownloadModelRequest, hf_token: Optional[str] = None
):
"""Start a background download for a HuggingFace model."""
repo_id = body.repo_id.strip()
if not _is_valid_repo_id(repo_id):
@ -111,7 +113,9 @@ async def download_model_response(body: DownloadModelRequest, hf_token: Optional
detail = f"Invalid repo_id: {repo_id!r}",
)
# Canonicalize so two different-cased paste-ins share one job + cache dir.
repo_id = await asyncio.to_thread(resolve_cached_repo_id_case, repo_id, repo_type = "model")
repo_id = await asyncio.to_thread(
resolve_cached_repo_id_case, repo_id, repo_type = "model"
)
# Avoid concurrent writers to the same HF cache files.
_reject_if_load_in_flight(repo_id)
@ -231,7 +235,9 @@ async def cancel_download_model_response(body: CancelDownloadRequest):
status_code = 400,
detail = f"Invalid repo_id: {repo_id!r}",
)
repo_id = await asyncio.to_thread(resolve_cached_repo_id_case, repo_id, repo_type = "model")
repo_id = await asyncio.to_thread(
resolve_cached_repo_id_case, repo_id, repo_type = "model"
)
variant = (body.gguf_variant or "").strip() or None
if variant is not None and not _is_valid_gguf_variant(variant):
raise HTTPException(
@ -250,12 +256,16 @@ async def cancel_download_model_response(body: CancelDownloadRequest):
return {"job_key": key, "state": state}
async def get_download_status_response(repo_id: str, gguf_variant: str = "") -> DownloadJobStatus:
async def get_download_status_response(
repo_id: str, gguf_variant: str = ""
) -> DownloadJobStatus:
"""Return the latest state of a background download job."""
repo_id = repo_id.strip()
if not _is_valid_repo_id(repo_id):
return DownloadJobStatus(state = "idle")
repo_id = await asyncio.to_thread(resolve_cached_repo_id_case, repo_id, repo_type = "model")
repo_id = await asyncio.to_thread(
resolve_cached_repo_id_case, repo_id, repo_type = "model"
)
variant = (gguf_variant or "").strip() or None
key = _download_job_key(repo_id, variant)
return _job_status(key, repo_id = repo_id, variant = variant)
@ -280,7 +290,9 @@ async def get_active_downloads_response(repo_id: str = "") -> ActiveDownloadsRes
)
def _variant_transport_status(repo_id: str, variant: str, hf_token: Optional[str]) -> dict:
def _variant_transport_status(
repo_id: str, variant: str, hf_token: Optional[str]
) -> dict:
incomplete_hashes = download_registry.incomplete_blob_hashes(
"model",
repo_id,
@ -312,13 +324,16 @@ def _variant_transport_status(repo_id: str, variant: str, hf_token: Optional[str
variant,
)
has_matching_incomplete = bool(
incomplete_hashes and variant_hashes and incomplete_hashes.intersection(variant_hashes)
incomplete_hashes
and variant_hashes
and incomplete_hashes.intersection(variant_hashes)
)
return {
"has_partial": has_partial,
"last_transport": last_transport,
"resumable": (
has_matching_incomplete and last_transport == download_registry.TRANSPORT_HTTP
has_matching_incomplete
and last_transport == download_registry.TRANSPORT_HTTP
),
}
@ -346,7 +361,9 @@ async def get_model_transport_status_response(
return _variant_transport_status(repo_id, variant, hf_token)
return {
"has_partial": has_active_incomplete_blobs("model", repo_id),
"last_transport": download_registry.read_active_transport_marker("model", repo_id),
"last_transport": download_registry.read_active_transport_marker(
"model", repo_id
),
"resumable": download_registry.is_resumable_partial("model", repo_id),
}
@ -390,7 +407,9 @@ async def get_gguf_download_progress_response(
if manifest is not None:
return (
sum(max(0, int(file.size or 0)) for file in manifest.expected_files),
frozenset(file.sha256 for file in manifest.expected_files if file.sha256),
frozenset(
file.sha256 for file in manifest.expected_files if file.sha256
),
)
return (
expected_total,

11
studio/backend/hub/services/models/folder_browser.py
View file

@ -305,7 +305,12 @@ def _browse_relative_parts(requested_path: str, root: Path) -> Optional[list[str
parts = [part for part in rel_text.split(os.sep) if part not in ("", ".")]
altsep = os.altsep
for part in parts:
if part == ".." or "\x00" in part or os.sep in part or (altsep and altsep in part):
if (
part == ".."
or "\x00" in part
or os.sep in part
or (altsep and altsep in part)
):
return None
return parts
@ -528,7 +533,9 @@ def browse_folders_response(
# Parent is None at the FS root and when it would step outside the sandbox,
# so the up-row never 403s on click.
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)

92
studio/backend/hub/services/models/gguf_variants.py
View file

@ -51,9 +51,7 @@ from hub.utils.gguf_plan import (
logger = get_logger(__name__)
_VARIANT_HASH_CACHE: "OrderedDict[tuple[str, str, str, bool], tuple[frozenset[str], float]]" = (
OrderedDict()
)
_VARIANT_HASH_CACHE: "OrderedDict[tuple[str, str, str, bool], tuple[frozenset[str], float]]" = OrderedDict()
_VARIANT_REQUIREMENT_CACHE: "OrderedDict[tuple[str, str, str], tuple[_GgufVariantRequirement, float]]" = OrderedDict()
_VARIANT_REQUIREMENT_NEG_CACHE: "OrderedDict[tuple[str, str], float]" = OrderedDict()
_VARIANT_HASH_MAX = 512
@ -122,7 +120,9 @@ def _variant_requirement_neg_cache_clear(key: tuple[str, str]) -> None:
_VARIANT_REQUIREMENT_NEG_CACHE.pop(key, None)
def _variant_hash_cache_get(key: tuple[str, str, str, bool]) -> Optional[frozenset[str]]:
def _variant_hash_cache_get(
key: tuple[str, str, str, bool],
) -> Optional[frozenset[str]]:
with _VARIANT_HASH_LOCK:
cached = _VARIANT_HASH_CACHE.get(key)
if cached is None:
@ -135,7 +135,9 @@ def _variant_hash_cache_get(key: tuple[str, str, str, bool]) -> Optional[frozens
return hashes
def _variant_hash_cache_set(key: tuple[str, str, str, bool], hashes: frozenset[str]) -> None:
def _variant_hash_cache_set(
key: tuple[str, str, str, bool], hashes: frozenset[str]
) -> None:
with _VARIANT_HASH_LOCK:
_VARIANT_HASH_CACHE[key] = (hashes, time.monotonic())
_VARIANT_HASH_CACHE.move_to_end(key)
@ -143,7 +145,9 @@ def _variant_hash_cache_set(key: tuple[str, str, str, bool], hashes: frozenset[s
_VARIANT_HASH_CACHE.popitem(last = False)
def _variant_requirement_cache_get(key: tuple[str, str, str]) -> Optional[_GgufVariantRequirement]:
def _variant_requirement_cache_get(
key: tuple[str, str, str],
) -> Optional[_GgufVariantRequirement]:
with _VARIANT_HASH_LOCK:
cached = _VARIANT_REQUIREMENT_CACHE.get(key)
if cached is None:
@ -157,7 +161,9 @@ def _variant_requirement_cache_get(key: tuple[str, str, str]) -> Optional[_GgufV
def _variant_requirement_cache_set_many(
repo_id: str, hf_token: Optional[str], requirements: dict[str, _GgufVariantRequirement]
repo_id: str,
hf_token: Optional[str],
requirements: dict[str, _GgufVariantRequirement],
) -> None:
with _VARIANT_HASH_LOCK:
now = time.monotonic()
@ -169,7 +175,9 @@ def _variant_requirement_cache_set_many(
_VARIANT_REQUIREMENT_CACHE.popitem(last = False)
def _build_gguf_variant_requirements(siblings: list) -> dict[str, _GgufVariantRequirement]:
def _build_gguf_variant_requirements(
siblings: list,
) -> dict[str, _GgufVariantRequirement]:
return build_gguf_variant_plans(siblings)
@ -280,7 +288,11 @@ def gguf_variant_blob_hashes(
if requirement is None and allow_remote:
requirement = gguf_variant_requirements(repo_id, variant, hf_token)
if requirement is not None:
hashes = requirement.required_hashes if include_companions else requirement.main_hashes
hashes = (
requirement.required_hashes
if include_companions
else requirement.main_hashes
)
if hashes:
_variant_hash_cache_set(key, hashes)
return hashes
@ -354,11 +366,15 @@ def _size_identity_matches(local_set: set[str], remote_size: int) -> bool:
def _variant_update_available_from_requirement(
local_blobs: dict[str, set[str]], requirement: Optional[_GgufVariantRequirement], variant: str
local_blobs: dict[str, set[str]],
requirement: Optional[_GgufVariantRequirement],
variant: str,
) -> bool:
if requirement is None or not local_blobs:
return False
local_by_posix = {path.replace("\\", "/"): blobs for path, blobs in local_blobs.items()}
local_by_posix = {
path.replace("\\", "/"): blobs for path, blobs in local_blobs.items()
}
for expected in requirement.expected_files:
path = str(expected.path).replace("\\", "/")
if not (
@ -392,7 +408,9 @@ def delete_variant_incomplete_blobs_result(
# With a sibling still downloading, ``companions=False`` keeps a shared mmproj
# from being unlinked out from under it; the repo's last delete reclaims it.
target_hashes = (
gguf_variant_blob_hashes(repo_id, variant, hf_token, include_companions = companions)
gguf_variant_blob_hashes(
repo_id, variant, hf_token, include_companions = companions
)
| extra_hashes
)
if not target_hashes:
@ -402,7 +420,9 @@ def delete_variant_incomplete_blobs_result(
incomplete_blob_hashes = set(),
variant_blob_hashes = frozenset(),
)
has_repo_partials = bool(download_registry.incomplete_blob_hashes("model", repo_id))
has_repo_partials = bool(
download_registry.incomplete_blob_hashes("model", repo_id)
)
return VariantIncompleteDeleteResult(
deleted = 0,
unresolved = has_variant_partial_state and has_repo_partials,
@ -447,7 +467,9 @@ def _mark_empty_dir_cleanables(
variants[i] = v.model_copy(update = {"partial": True})
for key, label in sorted(empty_by_key.items()):
if key not in listed:
variants.append(GgufVariantDetail(filename = f"{label}.gguf", quant = label, partial = True))
variants.append(
GgufVariantDetail(filename = f"{label}.gguf", quant = label, partial = True)
)
return response.model_copy(update = {"variants": variants})
@ -558,7 +580,9 @@ async def get_gguf_variants_response(
)
try:
variants, has_vision, siblings = list_gguf_variants(repo_id, hf_token = hf_token)
variants, has_vision, siblings = list_gguf_variants(
repo_id, hf_token = hf_token
)
except Exception:
cached = list_gguf_variants_from_hf_cache(repo_id)
if cached is not None:
@ -694,17 +718,25 @@ async def get_gguf_variants_response(
partial_quants: set[str] = set()
partial_quant_transports: dict[str, Optional[str]] = {}
try:
incomplete_hashes = download_registry.incomplete_blob_hashes("model", repo_id)
incomplete_hashes = download_registry.incomplete_blob_hashes(
"model", repo_id
)
except Exception as e:
logger.warning(f"Failed to compute partial GGUF variants for {repo_id}: {e}")
logger.warning(
f"Failed to compute partial GGUF variants for {repo_id}: {e}"
)
incomplete_hashes = set()
scan_snapshot_dir = hf_cache_scan.resolve_snapshot_dir_for_scan("model", repo_id)
scan_snapshot_dir = hf_cache_scan.resolve_snapshot_dir_for_scan(
"model", repo_id
)
# Manifest + marker + main incomplete-blob check: catches variants whose
# download was cancelled or whose expected shards are missing/undersized.
for variant in variants:
try:
requirement = requirements_by_quant.get(variant.quant.lower())
variant_hashes = requirement.main_hashes if requirement is not None else None
variant_hashes = (
requirement.main_hashes if requirement is not None else None
)
if variant_hashes is None and incomplete_hashes:
variant_hashes = gguf_variant_blob_hashes(
repo_id,
@ -720,13 +752,16 @@ async def get_gguf_variants_response(
variant_blob_hashes = variant_hashes,
):
partial_quants.add(variant.quant)
partial_quant_transports[variant.quant] = _partial_transport_for_variant(
repo_id,
variant.quant,
partial_quant_transports[variant.quant] = (
_partial_transport_for_variant(
repo_id,
variant.quant,
)
)
except Exception as e:
logger.warning(
f"Manifest-based partial check failed for " f"{repo_id}/{variant.quant}: {e}"
f"Manifest-based partial check failed for "
f"{repo_id}/{variant.quant}: {e}"
)
if incomplete_hashes:
for variant in variants:
@ -736,7 +771,8 @@ async def get_gguf_variants_response(
# companion_hashes adds the MTP drafter (mmproj_hashes covers
# every mmproj precision in the repo, not just the planned one).
if (
(requirement.mmproj_hashes | requirement.companion_hashes) & incomplete_hashes
(requirement.mmproj_hashes | requirement.companion_hashes)
& incomplete_hashes
) and _filenames_cached(
requirement.main_filenames,
requirement.main_size_bytes,
@ -759,7 +795,9 @@ async def get_gguf_variants_response(
display_label = v.display_label,
size_bytes = v.size_bytes,
download_size_bytes = (
requirement.download_size_bytes if requirement is not None else v.size_bytes
requirement.download_size_bytes
if requirement is not None
else v.size_bytes
),
downloaded = downloaded,
update_available = downloaded
@ -769,7 +807,9 @@ async def get_gguf_variants_response(
v.quant,
),
partial = is_partial,
partial_transport = (partial_quant_transports.get(v.quant) if is_partial else None),
partial_transport = (
partial_quant_transports.get(v.quant) if is_partial else None
),
)
return GgufVariantsResponse(

36
studio/backend/hub/services/models/local_inventory.py
View file

@ -99,7 +99,9 @@ def _is_model_directory_for_scan(path: Path, *, entry_limit: int | None) -> bool
if entry_limit is None:
return _is_model_directory(path)
try:
has_config = (path / "config.json").exists() or (path / "adapter_config.json").exists()
has_config = (path / "config.json").exists() or (
path / "adapter_config.json"
).exists()
except OSError:
return False
return has_config and _has_immediate_model_weight(path)
@ -152,7 +154,9 @@ def _scan_models_dir(
break
try:
is_dir = child.is_dir()
is_gguf_file = not is_dir and child.suffix.lower() == ".gguf" and child.is_file()
is_gguf_file = (
not is_dir and child.suffix.lower() == ".gguf" and child.is_file()
)
if not is_dir and not is_gguf_file:
continue
has_model_files = is_gguf_file or _has_immediate_model_signal(child)
@ -202,7 +206,9 @@ def _hf_repo_dir_has_content(repo_dir: Path) -> bool:
return False
def _scan_hf_cache(cache_dir: Path, *, entry_limit: int | None = None) -> List[LocalModelInfo]:
def _scan_hf_cache(
cache_dir: Path, *, entry_limit: int | None = None
) -> List[LocalModelInfo]:
if not _safe_is_dir(cache_dir):
return []
@ -240,7 +246,9 @@ def _scan_hf_cache(cache_dir: Path, *, entry_limit: int | None = None) -> List[L
repo_dir,
)
gguf_partial = hf_cache_scan.is_gguf_repo_partial(model_id, repo_dir)
has_gguf_variant_state, gguf_variant_state_size = _gguf_variant_state_summary(model_id)
has_gguf_variant_state, gguf_variant_state_size = _gguf_variant_state_summary(
model_id
)
snapshot_partial_transport = (
hf_cache_scan.partial_transport_for(
"model",
@ -323,7 +331,9 @@ def _scan_hf_cache(cache_dir: Path, *, entry_limit: int | None = None) -> List[L
return found
def _scan_lmstudio_dir(lm_dir: Path, *, entry_limit: int | None = None) -> List[LocalModelInfo]:
def _scan_lmstudio_dir(
lm_dir: Path, *, entry_limit: int | None = None
) -> List[LocalModelInfo]:
"""Scan an LM Studio models dir (``publisher/model-name`` folders of GGUFs, or top-level standalone GGUFs)."""
if not lm_dir.exists() or not lm_dir.is_dir():
return []
@ -439,7 +449,9 @@ def _resolve_allowed_models_dir(models_dir: str, allowed_roots: list[Path]) -> P
if not models_dir or not models_dir.strip():
raise ValueError("Directory not allowed")
requested = Path(os.path.realpath(os.path.expanduser(normalize_path(models_dir.strip()))))
requested = Path(
os.path.realpath(os.path.expanduser(normalize_path(models_dir.strip())))
)
if any(path_is_same_or_child(requested, root) for root in allowed_roots):
return requested
@ -522,7 +534,9 @@ async def _collect_models_from_default_sources(
and hf_default.resolve() != hf_cache_dir.resolve()
and hf_default.resolve() != legacy_hf.resolve()
):
local_models += await _scan_source("default HF cache", _scan_hf_cache, hf_default)
local_models += await _scan_source(
"default HF cache", _scan_hf_cache, hf_default
)
for lm_dir in lm_dirs:
local_models += await _scan_source("LM Studio", _scan_lmstudio_dir, lm_dir)
@ -633,12 +647,16 @@ def _filter_hidden_models(local_models: List[LocalModelInfo]) -> list[LocalModel
if model.source == "hf_cache"
else None
)
if not is_hidden_model(model.id, model.model_id, model.path, resolved_cache_path):
if not is_hidden_model(
model.id, model.model_id, model.path, resolved_cache_path
):
visible.append(model)
return visible
async def list_local_models_response(models_dir: str = "./models") -> LocalModelListResponse:
async def list_local_models_response(
models_dir: str = "./models",
) -> LocalModelListResponse:
"""List local model candidates from every supported on-device source."""
hf_cache_dir = _resolve_hf_cache_dir()
legacy_hf = legacy_hf_cache_dir()

22
studio/backend/hub/services/models/ollama.py
View file

@ -124,7 +124,9 @@ def _ollama_links_dir(ollama_dir: Path) -> Optional[Path]:
return None
def _make_ollama_blob_link(link_dir: Path, link_name: str, target: Path) -> Optional[str]:
def _make_ollama_blob_link(
link_dir: Path, link_name: str, target: Path
) -> Optional[str]:
"""Create a .gguf-named link to an Ollama blob: tries symlink then hardlink, skips the model if neither works (a full multi-GB copy would block the API). Idempotent."""
try:
link_dir.mkdir(parents = True, exist_ok = True)
@ -137,7 +139,9 @@ def _make_ollama_blob_link(link_dir: Path, link_name: str, target: Path) -> Opti
return None
link_path = _contained_link_path(link_dir, link_name)
if link_path is None:
logger.warning("Refusing unsafe Ollama link name %r under %s", link_name, link_dir)
logger.warning(
"Refusing unsafe Ollama link name %r under %s", link_name, link_dir
)
return None
try:
resolved = target.resolve()
@ -232,7 +236,9 @@ def _ollama_model_info_from_manifest(
model_type = cfg.get("model_type", "")
file_type = cfg.get("file_type", "")
except (json.JSONDecodeError, OSError) as e:
logger.debug("Could not parse Ollama config blob %s: %s", config_blob, e)
logger.debug(
"Could not parse Ollama config blob %s: %s", config_blob, e
)
layers = manifest.get("layers") or []
if not isinstance(layers, list):
@ -260,11 +266,17 @@ def _ollama_model_info_from_manifest(
model_blob = candidate
if materialize_links and model_link_dir is not None:
link_name = f"{safe_name}-{tag}{quant}.gguf"
gguf_link_path = _make_ollama_blob_link(model_link_dir, link_name, candidate)
gguf_link_path = _make_ollama_blob_link(
model_link_dir, link_name, candidate
)
elif materialize_links and media == "application/vnd.ollama.image.projector":
candidate = _ollama_blob_path(blobs_dir, digest)
if candidate is not None and _safe_is_file(candidate) and model_link_dir is not None:
if (
candidate is not None
and _safe_is_file(candidate)
and model_link_dir is not None
):
mmproj_name = f"{safe_name}-{tag}-mmproj.gguf"
_make_ollama_blob_link(model_link_dir, mmproj_name, candidate)

4
studio/backend/hub/services/snapshot_progress.py
View file

@ -41,7 +41,9 @@ _progress_step_lock = threading.Lock()
_last_progress_step: dict[str, int] = {}
def _log_progress_step(job_key: str, repo_id: str, variant: Optional[str], progress: float) -> None:
def _log_progress_step(
job_key: str, repo_id: str, variant: Optional[str], progress: float
) -> None:
step = int(progress * 10)
with _progress_step_lock:
last = _last_progress_step.get(job_key, -1)

12
studio/backend/hub/tests/test_dataset_services.py
View file

@ -198,7 +198,9 @@ def test_delete_cached_dataset_absent_everywhere_raises_404(monkeypatch):
def test_check_format_rejects_invalid_path_as_400():
with pytest.raises(HTTPException) as exc_info:
formatting.check_format_response(CheckFormatRequest(dataset_name = "../../etc/passwd"))
formatting.check_format_response(
CheckFormatRequest(dataset_name = "../../etc/passwd")
)
assert exc_info.value.status_code == 400
@ -286,7 +288,9 @@ def test_dataset_claim_register_cancel_uses_registry_marker_owner(monkeypatch):
)
result = asyncio.run(
downloads.download_dataset_response(SimpleNamespace(repo_id = "Org/Data", use_xet = False))
downloads.download_dataset_response(
SimpleNamespace(repo_id = "Org/Data", use_xet = False)
)
)
assert result["state"] == "cancelled"
@ -328,7 +332,9 @@ def test_upload_dataset_response_writes_non_empty_file(monkeypatch, tmp_path):
payload = b'{"text":"hello"}\n'
monkeypatch.setattr(local, "DATASET_UPLOAD_DIR", tmp_path)
response = asyncio.run(local.upload_dataset_response(_Upload("../train.jsonl", payload)))
response = asyncio.run(
local.upload_dataset_response(_Upload("../train.jsonl", payload))
)
stored_path = Path(response.stored_path)
assert response.filename == "train.jsonl"

11
studio/backend/hub/tests/test_download_lifecycle.py
View file

@ -59,11 +59,18 @@ def test_xet_failure_retries_over_http_for_model_and_dataset(monkeypatch, tmp_pa
register_worker = download_lifecycle.register_worker
for repo_type, repo_id, variant, expected_args in (
("model", "Org/Model", "Q4_K_M", ["--repo-id", "Org/Model", "--variant", "Q4_K_M"]),
(
"model",
"Org/Model",
"Q4_K_M",
["--repo-id", "Org/Model", "--variant", "Q4_K_M"],
),
("dataset", "Org/Data", None, ["--repo-id", "Org/Data", "--dataset"]),
):
registry = download_registry.DownloadRegistry()
key = download_registry.normalize_job_key(f"{repo_id}::{variant}" if variant else repo_id)
key = download_registry.normalize_job_key(
f"{repo_id}::{variant}" if variant else repo_id
)
assert registry.claim(
key,
download_registry.TRANSPORT_XET,

44
studio/backend/hub/tests/test_empty_variant_folder.py
View file

@ -27,13 +27,19 @@ def test_list_empty_gguf_variant_dirs_finds_empty_leftover(tmp_path, monkeypatch
assert gguf.list_empty_gguf_variant_dirs("org/Repo-GGUF") == {"UD-IQ1_S"}
def test_list_empty_excludes_quant_with_files_in_another_snapshot(tmp_path, monkeypatch):
def test_list_empty_excludes_quant_with_files_in_another_snapshot(
tmp_path, monkeypatch
):
snap1 = tmp_path / "s1" / "snapshots" / "rev"
(snap1 / "UD-IQ1_S").mkdir(parents = True) # empty here
snap2 = tmp_path / "s2" / "snapshots" / "rev"
(snap2 / "UD-IQ1_S").mkdir(parents = True)
(snap2 / "UD-IQ1_S" / "m-UD-IQ1_S-00001-of-00001.gguf").write_bytes(b"z") # has shards
monkeypatch.setattr(gguf, "iter_hf_cache_snapshots", lambda repo_id: iter([snap1, snap2]))
(snap2 / "UD-IQ1_S" / "m-UD-IQ1_S-00001-of-00001.gguf").write_bytes(
b"z"
) # has shards
monkeypatch.setattr(
gguf, "iter_hf_cache_snapshots", lambda repo_id: iter([snap1, snap2])
)
assert gguf.list_empty_gguf_variant_dirs("org/Repo-GGUF") == set()
@ -90,10 +96,16 @@ def test_remove_empty_variant_dirs_ignores_concurrent_refill(tmp_path, monkeypat
def test_mark_empty_dir_cleanables_appends_unlisted(monkeypatch):
monkeypatch.setattr(gguf_variants, "list_empty_gguf_variant_dirs", lambda repo_id: {"UD-IQ1_S"})
monkeypatch.setattr(
gguf_variants, "list_empty_gguf_variant_dirs", lambda repo_id: {"UD-IQ1_S"}
)
resp = GgufVariantsResponse(
repo_id = "org/Repo-GGUF",
variants = [GgufVariantDetail(filename = "m-UD-IQ1_M.gguf", quant = "UD-IQ1_M", downloaded = True)],
variants = [
GgufVariantDetail(
filename = "m-UD-IQ1_M.gguf", quant = "UD-IQ1_M", downloaded = True
)
],
)
out = gguf_variants._mark_empty_dir_cleanables("org/Repo-GGUF", resp)
by_q = {v.quant: v for v in out.variants}
@ -102,7 +114,9 @@ def test_mark_empty_dir_cleanables_appends_unlisted(monkeypatch):
def test_mark_empty_dir_cleanables_flips_listed_variant(monkeypatch):
monkeypatch.setattr(gguf_variants, "list_empty_gguf_variant_dirs", lambda repo_id: {"UD-IQ1_S"})
monkeypatch.setattr(
gguf_variants, "list_empty_gguf_variant_dirs", lambda repo_id: {"UD-IQ1_S"}
)
resp = GgufVariantsResponse(
repo_id = "org/Repo-GGUF",
variants = [GgufVariantDetail(filename = "m-UD-IQ1_S.gguf", quant = "UD-IQ1_S")],
@ -120,10 +134,16 @@ def _force_compute_to_raise(monkeypatch):
monkeypatch.setattr(gguf_variants, "list_gguf_variants", _boom, raising = False)
monkeypatch.setattr(
gguf_variants, "list_gguf_variants_from_hf_cache", lambda repo_id: None, raising = False
gguf_variants,
"list_gguf_variants_from_hf_cache",
lambda repo_id: None,
raising = False,
)
monkeypatch.setattr(
gguf_variants, "list_partial_gguf_variants_from_state", lambda repo_id: None, raising = False
gguf_variants,
"list_partial_gguf_variants_from_state",
lambda repo_id: None,
raising = False,
)
@ -133,7 +153,9 @@ def test_get_variants_surfaces_cleanable_when_metadata_fails(monkeypatch):
import asyncio
_force_compute_to_raise(monkeypatch)
monkeypatch.setattr(gguf_variants, "list_empty_gguf_variant_dirs", lambda repo_id: {"UD-IQ1_S"})
monkeypatch.setattr(
gguf_variants, "list_empty_gguf_variant_dirs", lambda repo_id: {"UD-IQ1_S"}
)
resp = asyncio.run(
gguf_variants.get_gguf_variants_response(
@ -152,7 +174,9 @@ def test_get_variants_reraises_when_no_cleanable(monkeypatch):
from fastapi import HTTPException
_force_compute_to_raise(monkeypatch)
monkeypatch.setattr(gguf_variants, "list_empty_gguf_variant_dirs", lambda repo_id: set())
monkeypatch.setattr(
gguf_variants, "list_empty_gguf_variant_dirs", lambda repo_id: set()
)
try:
asyncio.run(

260
studio/backend/hub/tests/test_model_services.py
View file

@ -149,7 +149,9 @@ def test_download_state_preserves_readable_keys_when_safe(monkeypatch, tmp_path)
@pytest.mark.parametrize("variant", ["bad variant with spaces", "q" * 64])
def test_download_state_bounds_long_repo_variant_filenames(monkeypatch, tmp_path, variant):
def test_download_state_bounds_long_repo_variant_filenames(
monkeypatch, tmp_path, variant
):
monkeypatch.setattr(state_dir, "cache_root", lambda: tmp_path)
repo_id = f"{'a' * 96}/{'b' * 96}"
@ -243,7 +245,9 @@ def test_browse_folders_hides_sensitive_dirs(monkeypatch, tmp_path):
(home / ".ssh").mkdir(parents = True)
(home / "models").mkdir()
# Accept and ignore the optional (media_roots, drive_roots) args the caller now passes.
monkeypatch.setattr(folder_browser, "_build_browse_allowlist", lambda *_a, **_k: [home])
monkeypatch.setattr(
folder_browser, "_build_browse_allowlist", lambda *_a, **_k: [home]
)
response = folder_browser.browse_folders_response(str(home), show_hidden = True)
@ -259,15 +263,22 @@ def test_browse_allowlist_includes_linux_run_media_mounts(monkeypatch, tmp_path)
home.mkdir()
model_dir.mkdir(parents = True)
monkeypatch.setattr(folder_browser.Path, "home", lambda: home)
monkeypatch.setattr(folder_browser, "linux_run_media_mount_roots", lambda: [media_root])
monkeypatch.setattr(folder_browser, "_resolve_hf_cache_dir", lambda: tmp_path / "missing-hf")
monkeypatch.setattr(
folder_browser, "linux_run_media_mount_roots", lambda: [media_root]
)
monkeypatch.setattr(
folder_browser, "_resolve_hf_cache_dir", lambda: tmp_path / "missing-hf"
)
monkeypatch.setattr(scan_folders, "list_scan_folders", lambda: [])
monkeypatch.setattr(folder_browser, "well_known_model_dirs", lambda: [])
allowlist = folder_browser._build_browse_allowlist()
assert media_root.resolve() in allowlist
assert folder_browser._resolve_browse_target(str(model_dir), allowlist) == model_dir.resolve()
assert (
folder_browser._resolve_browse_target(str(model_dir), allowlist)
== model_dir.resolve()
)
def test_get_models_folder_response_creates_and_returns_dir(monkeypatch, tmp_path):
@ -344,7 +355,9 @@ def test_make_ollama_blob_link_refuses_escaping_name(tmp_path):
blob.parent.mkdir(parents = True)
blob.write_bytes(b"weights")
escaped = ollama._make_ollama_blob_link(link_dir, "model-tag-../../../pwned.gguf", blob)
escaped = ollama._make_ollama_blob_link(
link_dir, "model-tag-../../../pwned.gguf", blob
)
assert escaped is None
assert not list(tmp_path.rglob("pwned.gguf"))
@ -360,7 +373,9 @@ def test_cached_gguf_scan_dedupes_and_excludes_mmproj_only(monkeypatch, tmp_path
[_file("Q4_K_M.gguf", 300), _file("Q8_0.gguf", 200)],
tmp_path / "large",
)
mmproj_only = _repo("Org/VisionAdapter", [_file("mmproj-F16.gguf", 900)], tmp_path / "mmproj")
mmproj_only = _repo(
"Org/VisionAdapter", [_file("mmproj-F16.gguf", 900)], tmp_path / "mmproj"
)
monkeypatch.setattr(
cache_inventory,
"all_hf_cache_scans",
@ -401,7 +416,9 @@ def test_cached_gguf_scan_preserves_partial_flag(monkeypatch, tmp_path):
assert row["capabilities"]["can_chat"] is False
def test_cached_gguf_scan_includes_variant_state_without_completed_gguf(monkeypatch, tmp_path):
def test_cached_gguf_scan_includes_variant_state_without_completed_gguf(
monkeypatch, tmp_path
):
monkeypatch.setattr(state_dir, "cache_root", lambda: tmp_path / "state")
repo_path = tmp_path / "hub" / "models--Org--PartialGguf"
repo_path.mkdir(parents = True)
@ -417,7 +434,9 @@ def test_cached_gguf_scan_includes_variant_state_without_completed_gguf(monkeypa
[download_manifest.ExpectedFile(path = "model-Q4_K_M.gguf", size = 4096)],
"http",
)
assert download_manifest.write_cancel_marker("model", "Org/PartialGguf", "Q4_K_M", "http")
assert download_manifest.write_cancel_marker(
"model", "Org/PartialGguf", "Q4_K_M", "http"
)
monkeypatch.setattr(
cache_inventory,
"all_hf_cache_scans",
@ -439,7 +458,9 @@ def test_cached_gguf_scan_includes_variant_state_without_completed_gguf(monkeypa
assert row["capabilities"]["requires_variant"] is True
def test_cached_gguf_scan_hides_infra_repos_without_user_downloads(monkeypatch, tmp_path):
def test_cached_gguf_scan_hides_infra_repos_without_user_downloads(
monkeypatch, tmp_path
):
probe = _repo(
"ggml-org/models",
[_file("tinyllamas/stories260K.gguf", 1_200_000)],
@ -467,7 +488,9 @@ def test_cached_gguf_scan_hides_infra_repos_without_user_downloads(monkeypatch,
assert [row["repo_id"] for row in result["cached"]] == ["Org/Chat-GGUF"]
def test_cached_gguf_scan_keeps_infra_repo_with_user_downloaded_variant(monkeypatch, tmp_path):
def test_cached_gguf_scan_keeps_infra_repo_with_user_downloaded_variant(
monkeypatch, tmp_path
):
monkeypatch.setattr(state_dir, "cache_root", lambda: tmp_path / "state")
embedder = _repo(
"unsloth/bge-small-en-v1.5-GGUF",
@ -482,7 +505,11 @@ def test_cached_gguf_scan_keeps_infra_repo_with_user_downloaded_variant(monkeypa
"model",
"unsloth/bge-small-en-v1.5-GGUF",
"Q8_0",
[download_manifest.ExpectedFile(path = "bge-small-en-v1.5-Q8_0.gguf", size = 35_000_000)],
[
download_manifest.ExpectedFile(
path = "bge-small-en-v1.5-Q8_0.gguf", size = 35_000_000
)
],
"http",
)
monkeypatch.setattr(
@ -498,7 +525,9 @@ def test_cached_gguf_scan_keeps_infra_repo_with_user_downloaded_variant(monkeypa
result = {"cached": cache_inventory._scan_cached_gguf()}
assert [row["repo_id"] for row in result["cached"]] == ["unsloth/bge-small-en-v1.5-GGUF"]
assert [row["repo_id"] for row in result["cached"]] == [
"unsloth/bge-small-en-v1.5-GGUF"
]
assert result["cached"][0]["capabilities"]["can_chat"] is False
@ -682,7 +711,9 @@ def test_cached_models_scan_keeps_unrelated_repo_with_custom_generic_embedder(
assert [row["repo_id"] for row in result["cached"]] == ["user/model-chat"]
def test_cached_scans_hide_stale_default_embedder_after_custom_setting(monkeypatch, tmp_path):
def test_cached_scans_hide_stale_default_embedder_after_custom_setting(
monkeypatch, tmp_path
):
from core.rag import config as rag_config
monkeypatch.setattr(rag_config, "effective_embedding_model", lambda: "org/custom")
@ -823,7 +854,9 @@ def test_gguf_variant_blob_hashes_skip_missing_rfilename(monkeypatch):
monkeypatch.setattr(
gguf_variants,
"_fetch_gguf_variant_requirements",
lambda _repo_id, _hf_token = None: gguf_variants._build_gguf_variant_requirements(siblings),
lambda _repo_id, _hf_token = None: gguf_variants._build_gguf_variant_requirements(
siblings
),
)
result = gguf_variants.gguf_variant_blob_hashes("Org/Malformed", "Q4_K_M", None)
@ -867,7 +900,9 @@ def test_download_gguf_variant_purges_only_main_quant_hashes(monkeypatch, tmp_pa
),
)
monkeypatch.setattr(
hf_download, "_verify_completed_download", lambda *args, **kwargs: verified.append(args)
hf_download,
"_verify_completed_download",
lambda *args, **kwargs: verified.append(args),
)
monkeypatch.setattr(
download_registry,
@ -882,7 +917,8 @@ def test_download_gguf_variant_purges_only_main_quant_hashes(monkeypatch, tmp_pa
sys.modules,
"huggingface_hub",
SimpleNamespace(
snapshot_download = lambda **kwargs: snapshot_calls.append(kwargs) or str(tmp_path)
snapshot_download = lambda **kwargs: snapshot_calls.append(kwargs)
or str(tmp_path)
),
)
@ -898,12 +934,20 @@ def test_download_gguf_variant_purges_only_main_quant_hashes(monkeypatch, tmp_pa
},
)
]
assert [file.path for file in written[0][3]] == ["model-Q4_K_M.gguf", "mmproj-F16.gguf"]
assert snapshot_calls[0]["allow_patterns"] == ["model-Q4_K_M.gguf", "mmproj-F16.gguf"]
assert [file.path for file in written[0][3]] == [
"model-Q4_K_M.gguf",
"mmproj-F16.gguf",
]
assert snapshot_calls[0]["allow_patterns"] == [
"model-Q4_K_M.gguf",
"mmproj-F16.gguf",
]
assert verified == [("model", "Org/Vision", "Q4_K_M", str(tmp_path))]
def test_download_gguf_variant_manifest_resume_purges_only_main_quant_hashes(monkeypatch, tmp_path):
def test_download_gguf_variant_manifest_resume_purges_only_main_quant_hashes(
monkeypatch, tmp_path
):
prepare_calls = []
snapshot_calls = []
@ -941,12 +985,15 @@ def test_download_gguf_variant_manifest_resume_purges_only_main_quant_hashes(mon
"prepare_cache_for_transport",
lambda *args, **kwargs: prepare_calls.append((args, kwargs)) or 0,
)
monkeypatch.setattr(hf_download, "_verify_completed_download", lambda *_args, **_kwargs: None)
monkeypatch.setattr(
hf_download, "_verify_completed_download", lambda *_args, **_kwargs: None
)
monkeypatch.setitem(
sys.modules,
"huggingface_hub",
SimpleNamespace(
snapshot_download = lambda **kwargs: snapshot_calls.append(kwargs) or str(tmp_path)
snapshot_download = lambda **kwargs: snapshot_calls.append(kwargs)
or str(tmp_path)
),
)
@ -962,10 +1009,15 @@ def test_download_gguf_variant_manifest_resume_purges_only_main_quant_hashes(mon
},
)
]
assert snapshot_calls[0]["allow_patterns"] == ["model-Q4_K_M.gguf", "mmproj-F16.gguf"]
assert snapshot_calls[0]["allow_patterns"] == [
"model-Q4_K_M.gguf",
"mmproj-F16.gguf",
]
def test_download_snapshot_recovers_manifest_after_metadata_fallback(monkeypatch, tmp_path):
def test_download_snapshot_recovers_manifest_after_metadata_fallback(
monkeypatch, tmp_path
):
metadata_calls = []
written = []
cleared = []
@ -975,11 +1027,15 @@ def test_download_snapshot_recovers_manifest_after_metadata_fallback(monkeypatch
metadata_calls.append(True)
if len(metadata_calls) == 1:
raise RuntimeError("metadata down")
return SimpleNamespace(siblings = [SimpleNamespace(rfilename = "config.json", size = 12)])
return SimpleNamespace(
siblings = [SimpleNamespace(rfilename = "config.json", size = 12)]
)
monkeypatch.setattr(hf_download, "_model_info_with_retry", _metadata)
monkeypatch.setattr(
hf_download, "_verify_completed_download", lambda *args, **kwargs: verified.append(args)
hf_download,
"_verify_completed_download",
lambda *args, **kwargs: verified.append(args),
)
monkeypatch.setattr(
download_registry, "prepare_cache_for_transport", lambda *_args, **_kwargs: 0
@ -1018,7 +1074,9 @@ def test_download_dataset_continues_without_metadata_manifest(monkeypatch, tmp_p
monkeypatch.setattr(hf_download, "_dataset_info_with_retry", _metadata)
monkeypatch.setattr(
hf_download, "_verify_completed_download", lambda *args, **kwargs: verified.append(args)
hf_download,
"_verify_completed_download",
lambda *args, **kwargs: verified.append(args),
)
monkeypatch.setattr(
download_registry, "prepare_cache_for_transport", lambda *_args, **_kwargs: 0
@ -1036,7 +1094,8 @@ def test_download_dataset_continues_without_metadata_manifest(monkeypatch, tmp_p
sys.modules,
"huggingface_hub",
SimpleNamespace(
snapshot_download = lambda **kwargs: snapshot_calls.append(kwargs) or str(tmp_path)
snapshot_download = lambda **kwargs: snapshot_calls.append(kwargs)
or str(tmp_path)
),
)
@ -1070,13 +1129,17 @@ def test_download_snapshot_fails_when_metadata_unavailable_and_partial_remains(
monkeypatch.setattr(hf_download, "_model_info_with_retry", _metadata)
monkeypatch.setattr(
hf_download, "_verify_completed_download", lambda *args, **kwargs: verified.append(args)
hf_download,
"_verify_completed_download",
lambda *args, **kwargs: verified.append(args),
)
monkeypatch.setattr(
download_registry, "prepare_cache_for_transport", lambda *_args, **_kwargs: 0
)
monkeypatch.setattr(download_manifest, "clear_cancel_marker", lambda *_args: None)
monkeypatch.setattr(download_manifest, "read_manifest", lambda *_args, **_kwargs: None)
monkeypatch.setattr(
download_manifest, "read_manifest", lambda *_args, **_kwargs: None
)
monkeypatch.setattr(
download_manifest, "write_manifest", lambda *args: written.append(args) or True
)
@ -1176,7 +1239,9 @@ def test_gguf_download_progress_fallback_logs_warning(monkeypatch):
assert kwargs == {}
def test_gguf_progress_counts_completed_mmproj_with_expected_bytes(monkeypatch, tmp_path):
def test_gguf_progress_counts_completed_mmproj_with_expected_bytes(
monkeypatch, tmp_path
):
"""A finished mmproj companion keeps counting toward progress once the caller
supplies expected bytes; resolving the variant requirement credits it."""
entry = tmp_path / "models--Org--Model-GGUF"
@ -1595,7 +1660,9 @@ def test_gguf_progress_scoped_hashes_exclude_sibling_quant(monkeypatch, tmp_path
assert result["downloaded_bytes"] == 5
def test_gguf_progress_unknown_hashes_does_not_count_foreign_blobs(monkeypatch, tmp_path):
def test_gguf_progress_unknown_hashes_does_not_count_foreign_blobs(
monkeypatch, tmp_path
):
# With a variant's hashes unresolved (metadata flaked, no manifest), the
# shared blobs/ dir's FINALIZED blobs must NOT be counted wholesale: a cached
# sibling quant (``siblinghash``) alongside is the "instant ~900 MB" bug.
@ -1647,7 +1714,9 @@ def test_gguf_progress_unknown_hashes_does_not_count_foreign_blobs(monkeypatch,
assert result["complete_on_disk"] is False
def test_gguf_progress_unknown_hashes_drops_unscoped_incomplete_blob(monkeypatch, tmp_path):
def test_gguf_progress_unknown_hashes_drops_unscoped_incomplete_blob(
monkeypatch, tmp_path
):
# With hashes unresolved, an .incomplete in the shared blobs/ dir can't be
# attributed to this variant (it may be a concurrent sibling's active write),
# so it is dropped, mirroring the finalized-blob guard. In production the
@ -1696,7 +1765,9 @@ def test_gguf_progress_unknown_hashes_drops_unscoped_incomplete_blob(monkeypatch
assert result["completed_bytes"] == 0 # finalized sibling still ignored
def test_gguf_progress_unknown_hashes_no_backward_dip_when_variant_finalizes(monkeypatch, tmp_path):
def test_gguf_progress_unknown_hashes_no_backward_dip_when_variant_finalizes(
monkeypatch, tmp_path
):
# Regression for the two-variant dip: with hashes unresolved, the first quant
# finalizes while the sibling still writes its .incomplete. The sibling's
# bytes used to leak into this numerator, dipping the bar ~99% -> ~78% for
@ -1767,7 +1838,9 @@ def test_hf_cache_model_file_probe_is_bounded(monkeypatch, tmp_path):
model.write_bytes(b"weights")
entries = [first, second, model]
monkeypatch.setattr(model_common.Path, "rglob", lambda _self, _pattern: iter(entries))
monkeypatch.setattr(
model_common.Path, "rglob", lambda _self, _pattern: iter(entries)
)
monkeypatch.setattr(model_common, "_HF_CACHE_MODEL_FILE_PROBE_LIMIT", 2)
bounded = model_common._iter_hf_cache_model_files(snapshot)
@ -1790,7 +1863,9 @@ def test_download_state_lookup_is_repo_case_insensitive(monkeypatch, tmp_path):
None,
[download_manifest.ExpectedFile(path = "config.json", size = 12)],
)
assert download_manifest.write_cancel_marker("model", "Owner/Repo", "Q4_K_M", "http")
assert download_manifest.write_cancel_marker(
"model", "Owner/Repo", "Q4_K_M", "http"
)
manifest = download_manifest.read_manifest("model", "owner/repo", None)
@ -1823,7 +1898,9 @@ def test_hf_cache_scan_fallback_row_uses_local_model_info_alias(monkeypatch, tmp
blobs_dir = repo_dir / "blobs"
blobs_dir.mkdir(parents = True)
(blobs_dir / "blob").write_bytes(b"content")
monkeypatch.setattr(local_inventory, "_classify_local_path", lambda *_args, **_kwargs: [])
monkeypatch.setattr(
local_inventory, "_classify_local_path", lambda *_args, **_kwargs: []
)
monkeypatch.setattr(
local_inventory.hf_cache_scan,
"is_snapshot_partial",
@ -1862,8 +1939,12 @@ def test_hf_cache_scan_uses_gguf_partial_row_for_variant_state(monkeypatch, tmp_
[download_manifest.ExpectedFile(path = "model-Q4_K_M.gguf", size = 8192)],
"http",
)
assert download_manifest.write_cancel_marker("model", "Org/PartialGguf", "Q4_K_M", "http")
monkeypatch.setattr(local_inventory, "_classify_local_path", lambda *_args, **_kwargs: [])
assert download_manifest.write_cancel_marker(
"model", "Org/PartialGguf", "Q4_K_M", "http"
)
monkeypatch.setattr(
local_inventory, "_classify_local_path", lambda *_args, **_kwargs: []
)
monkeypatch.setattr(
local_inventory.hf_cache_scan,
"is_snapshot_partial",
@ -1907,12 +1988,16 @@ def test_local_inventory_filters_custom_embedder_hf_cache_row(monkeypatch, tmp_p
model_id = repo_id,
)
rows = local_inventory._filter_hidden_models([_row("org/embedder"), _row("org/chat-model")])
rows = local_inventory._filter_hidden_models(
[_row("org/embedder"), _row("org/chat-model")]
)
assert [row.model_id for row in rows] == ["org/chat-model"]
def test_local_inventory_filters_embedder_configured_by_snapshot_path(monkeypatch, tmp_path):
def test_local_inventory_filters_embedder_configured_by_snapshot_path(
monkeypatch, tmp_path
):
from core.rag import config as rag_config
embedder_path = tmp_path / "hub" / "models--org--embedder"
@ -1929,7 +2014,9 @@ def test_local_inventory_filters_embedder_configured_by_snapshot_path(monkeypatc
monkeypatch.setattr(
local_inventory.hf_cache_scan,
"resolve_hf_cache_realpath",
lambda path: str(embedder_snapshot) if Path(path) == embedder_path else str(path),
lambda path: str(embedder_snapshot)
if Path(path) == embedder_path
else str(path),
)
def _row(repo_id: str, repo_path: Path):
@ -1957,7 +2044,9 @@ def test_model_download_job_helpers_preserve_idle_shape():
assert status.error is None
def test_gguf_repo_partial_treats_completed_disk_variant_as_clean(monkeypatch, tmp_path):
def test_gguf_repo_partial_treats_completed_disk_variant_as_clean(
monkeypatch, tmp_path
):
monkeypatch.setattr(state_dir, "cache_root", lambda: tmp_path / "state")
snapshot = tmp_path / "cache" / "models--Org--Repo" / "snapshots" / "abc"
snapshot.mkdir(parents = True)
@ -2059,7 +2148,9 @@ def test_variant_partial_accepts_variant_filtered_legacy_hashes(monkeypatch, tmp
)
def test_variant_partial_accepts_completed_variant_in_non_latest_snapshot(monkeypatch, tmp_path):
def test_variant_partial_accepts_completed_variant_in_non_latest_snapshot(
monkeypatch, tmp_path
):
"""A verified GGUF update can prune an older snapshot and make that old
directory the newest by mtime. The variant is still complete when another
snapshot satisfies its manifest."""
@ -2087,13 +2178,17 @@ def test_variant_partial_accepts_completed_variant_in_non_latest_snapshot(monkey
)
def test_gguf_variants_partial_marker_overrides_size_only_downloaded(monkeypatch, tmp_path):
def test_gguf_variants_partial_marker_overrides_size_only_downloaded(
monkeypatch, tmp_path
):
async def _run_inline(fn, *args, **kwargs):
return fn(*args, **kwargs)
monkeypatch.setattr(state_dir, "cache_root", lambda: tmp_path / "state")
monkeypatch.setattr(gguf_variants.asyncio, "to_thread", _run_inline)
assert download_manifest.write_cancel_marker("model", "Org/PartialRepo", "Q4_K_M", "http")
assert download_manifest.write_cancel_marker(
"model", "Org/PartialRepo", "Q4_K_M", "http"
)
snapshot = tmp_path / "cache" / "models--Org--PartialRepo" / "snapshots" / "rev0"
snapshot.mkdir(parents = True)
(snapshot / "model-Q4_K_M.gguf").write_bytes(b"x" * 100)
@ -2423,7 +2518,9 @@ def test_finalize_worker_exit_never_kills_a_healthy_worker(monkeypatch, tmp_path
)
def test_prepare_cache_for_transport_purges_only_requested_hashes(monkeypatch, tmp_path):
def test_prepare_cache_for_transport_purges_only_requested_hashes(
monkeypatch, tmp_path
):
root = tmp_path / "hub"
blobs = root / "models--Org--Repo" / "blobs"
blobs.mkdir(parents = True)
@ -2452,7 +2549,9 @@ def _vision_cache_root(monkeypatch, tmp_path):
return blobs
def test_prepare_cache_for_transport_purges_cross_transport_companion(monkeypatch, tmp_path):
def test_prepare_cache_for_transport_purges_cross_transport_companion(
monkeypatch, tmp_path
):
blobs = _vision_cache_root(monkeypatch, tmp_path)
companion = frozenset({"shared-mmproj"})
@ -2482,7 +2581,9 @@ def test_prepare_cache_for_transport_purges_cross_transport_companion(monkeypatc
assert not (blobs / "shared-mmproj.incomplete").exists()
def test_prepare_cache_for_transport_preserves_same_transport_companion(monkeypatch, tmp_path):
def test_prepare_cache_for_transport_preserves_same_transport_companion(
monkeypatch, tmp_path
):
blobs = _vision_cache_root(monkeypatch, tmp_path)
companion = frozenset({"shared-mmproj"})
@ -2537,7 +2638,9 @@ def test_prepare_cache_for_transport_protects_peer_companion(monkeypatch, tmp_pa
assert (blobs / "shared-mmproj.incomplete").exists()
def test_model_download_records_completed_baseline_for_new_gguf_variant(monkeypatch, tmp_path):
def test_model_download_records_completed_baseline_for_new_gguf_variant(
monkeypatch, tmp_path
):
async def _run_inline(fn, *args, **kwargs):
return fn(*args, **kwargs)
@ -2552,7 +2655,9 @@ def test_model_download_records_completed_baseline_for_new_gguf_variant(monkeypa
downloads.gguf_variants,
"gguf_variant_blob_hashes",
lambda _repo, _variant, _token = None, include_companions = True, **_kwargs: (
frozenset({"mainhash", "mmprojhash"}) if include_companions else frozenset({"mainhash"})
frozenset({"mainhash", "mmprojhash"})
if include_companions
else frozenset({"mainhash"})
),
)
monkeypatch.setattr(
@ -2595,7 +2700,9 @@ def test_model_download_records_completed_baseline_for_new_gguf_variant(monkeypa
registry = _Registry()
monkeypatch.setattr(downloads, "_registry", registry)
monkeypatch.setattr(downloads, "_spawn_download_worker", lambda *_args, **_kwargs: _Proc())
monkeypatch.setattr(
downloads, "_spawn_download_worker", lambda *_args, **_kwargs: _Proc()
)
asyncio.run(
downloads.download_model_response(
@ -2604,7 +2711,9 @@ def test_model_download_records_completed_baseline_for_new_gguf_variant(monkeypa
)
assert registry.claim_kwargs["blob_hashes"] == frozenset({"mainhash"})
assert registry.claim_kwargs["progress_blob_hashes"] == frozenset({"mainhash", "mmprojhash"})
assert registry.claim_kwargs["progress_blob_hashes"] == frozenset(
{"mainhash", "mmprojhash"}
)
assert registry.claim_kwargs["completed_baseline_bytes"] == 30
@ -2638,7 +2747,9 @@ def test_gguf_model_download_skips_completed_baseline_for_variant_resume_state(
downloads.gguf_variants,
"gguf_variant_blob_hashes",
lambda _repo, _variant, _token = None, include_companions = True, **_kwargs: (
frozenset({"mainhash", "mmprojhash"}) if include_companions else frozenset({"mainhash"})
frozenset({"mainhash", "mmprojhash"})
if include_companions
else frozenset({"mainhash"})
),
)
monkeypatch.setattr(
@ -2681,7 +2792,9 @@ def test_gguf_model_download_skips_completed_baseline_for_variant_resume_state(
registry = _Registry()
monkeypatch.setattr(downloads, "_registry", registry)
monkeypatch.setattr(downloads, "_spawn_download_worker", lambda *_args, **_kwargs: _Proc())
monkeypatch.setattr(
downloads, "_spawn_download_worker", lambda *_args, **_kwargs: _Proc()
)
asyncio.run(
downloads.download_model_response(
@ -2695,7 +2808,9 @@ def test_gguf_model_download_skips_completed_baseline_for_variant_resume_state(
def test_model_idle_status_uses_cancel_marker_after_restart(monkeypatch, tmp_path):
monkeypatch.setattr(state_dir, "cache_root", lambda: tmp_path)
monkeypatch.setattr(downloads, "_registry", download_registry.DownloadRegistry())
assert download_manifest.write_cancel_marker("model", "Owner/Repo", "Q4_K_M", "http")
assert download_manifest.write_cancel_marker(
"model", "Owner/Repo", "Q4_K_M", "http"
)
status = asyncio.run(downloads.get_download_status_response("owner/repo", "Q4_K_M"))
@ -2947,7 +3062,9 @@ def test_model_download_watcher_invalidates_hf_cache_scan(monkeypatch):
"_spawn_download_worker",
lambda *_args, **_kwargs: object(),
)
monkeypatch.setattr(downloads.download_lifecycle.threading, "Thread", _ImmediateThread)
monkeypatch.setattr(
downloads.download_lifecycle.threading, "Thread", _ImmediateThread
)
monkeypatch.setattr(
downloads.hf_cache_scan,
"invalidate_hf_cache_scans",
@ -3054,7 +3171,10 @@ def test_two_concurrent_same_repo_variants_both_complete(monkeypatch, tmp_path):
while time.monotonic() < deadline:
s4 = registry.get_job(key_q4).state
s8 = registry.get_job(key_q8).state
if s4 in download_registry.TERMINAL_STATES and s8 in download_registry.TERMINAL_STATES:
if (
s4 in download_registry.TERMINAL_STATES
and s8 in download_registry.TERMINAL_STATES
):
break
time.sleep(0.02)
@ -3170,7 +3290,9 @@ def test_snapshot_progress_filters_stale_blobs(monkeypatch, tmp_path):
assert result["expected_bytes"] == 140
def test_snapshot_progress_confirms_complete_only_with_verified_snapshot(monkeypatch, tmp_path):
def test_snapshot_progress_confirms_complete_only_with_verified_snapshot(
monkeypatch, tmp_path
):
entry = tmp_path / "models--Org--Model"
blobs = entry / "blobs"
snap = entry / "snapshots" / "rev0"
@ -3233,7 +3355,9 @@ def test_expected_files_from_snapshot_dir_records_relative_paths_and_sizes(tmp_p
assert all(f.sha256 is None for f in files)
def test_snapshot_progress_complete_with_manifest_synthesized_from_disk(monkeypatch, tmp_path):
def test_snapshot_progress_complete_with_manifest_synthesized_from_disk(
monkeypatch, tmp_path
):
"""A finished snapshot whose only manifest was synthesized from on-disk files
still verifies as complete, so a refresh finalizes it instead of capping at
99% and evicting it as gone."""
@ -3390,7 +3514,9 @@ def test_download_snapshot_writes_manifest_for_xet(monkeypatch, tmp_path):
),
)
monkeypatch.setattr(
hf_download, "_verify_completed_download", lambda *args, **kwargs: verified.append(args)
hf_download,
"_verify_completed_download",
lambda *args, **kwargs: verified.append(args),
)
monkeypatch.setattr(
download_registry, "prepare_cache_for_transport", lambda *_args, **_kwargs: 0
@ -3425,7 +3551,9 @@ def test_download_gguf_variant_writes_manifest_for_xet(monkeypatch, tmp_path):
),
)
monkeypatch.setattr(
hf_download, "_verify_completed_download", lambda *args, **kwargs: verified.append(args)
hf_download,
"_verify_completed_download",
lambda *args, **kwargs: verified.append(args),
)
monkeypatch.setattr(
download_registry, "prepare_cache_for_transport", lambda *_args, **_kwargs: 0
@ -3460,7 +3588,9 @@ def test_download_dataset_writes_manifest_for_xet(monkeypatch, tmp_path):
),
)
monkeypatch.setattr(
hf_download, "_verify_completed_download", lambda *args, **kwargs: verified.append(args)
hf_download,
"_verify_completed_download",
lambda *args, **kwargs: verified.append(args),
)
monkeypatch.setattr(
download_registry, "prepare_cache_for_transport", lambda *_args, **_kwargs: 0
@ -3498,7 +3628,9 @@ def test_dataset_status_includes_generation(monkeypatch):
lambda repo_id, **_kwargs: repo_id,
)
result = asyncio.run(dataset_downloads.get_dataset_download_status_response("Org/Data"))
result = asyncio.run(
dataset_downloads.get_dataset_download_status_response("Org/Data")
)
assert result.state == "running"
assert result.generation == 4

12
studio/backend/hub/utils/dataset_cache.py
View file

@ -50,7 +50,9 @@ def _matches_label(snapshot: Path, path: Path, label: str) -> bool:
return label in rel
def dataset_snapshot_from_cache_path(local_path: Optional[str], repo_id: str) -> Optional[Path]:
def dataset_snapshot_from_cache_path(
local_path: Optional[str], repo_id: str
) -> Optional[Path]:
if not local_path or not repo_id:
return None
try:
@ -115,7 +117,9 @@ def cached_dataset_candidates(
) -> list[Path]:
try:
files = [
p for p in snapshot.rglob("*") if p.is_file() and p.name.lower().endswith(extensions)
p
for p in snapshot.rglob("*")
if p.is_file() and p.name.lower().endswith(extensions)
]
except OSError:
return []
@ -127,7 +131,9 @@ def cached_dataset_candidates(
def score(path: Path) -> tuple[int, int, str]:
rel = _rel_lower(snapshot, path)
subset_match = bool(subset_lower and _matches_label(snapshot, path, subset_lower))
subset_match = bool(
subset_lower and _matches_label(snapshot, path, subset_lower)
)
split_match = bool(split_lower and split_label_matches(rel, split_lower))
location_rank = 3
if split_match and (not subset_lower or subset_match):

57
studio/backend/hub/utils/dataset_format.py
View file

@ -23,7 +23,10 @@ def _column_names(dataset, sample: Optional[dict] = None) -> list[str]:
def _keyword_in_column(keyword: str, col_name: str) -> bool:
return re.search(r"\b" + re.escape(keyword) + r"\b", col_name, re.IGNORECASE) is not None
return (
re.search(r"\b" + re.escape(keyword) + r"\b", col_name, re.IGNORECASE)
is not None
)
def _unknown_dataset_format(
@ -165,14 +168,19 @@ def detect_custom_format_heuristic(dataset):
def has_keyword(col_name, keywords):
col_lower = col_name.lower()
col_normalized = col_lower.replace("_", "").replace("-", "").replace(" ", "")
return any(keyword in col_lower or keyword in col_normalized for keyword in keywords)
return any(
keyword in col_lower or keyword in col_normalized for keyword in keywords
)
def is_metadata(col_name):
col_lower = col_name.lower()
if col_lower in metadata_exact_match or col_lower in metadata_prefix_patterns:
return True
for pattern in metadata_prefix_patterns:
if col_lower.startswith(pattern.split("_")[0] + "_") and col_lower != pattern:
if (
col_lower.startswith(pattern.split("_")[0] + "_")
and col_lower != pattern
):
if "_" in col_lower:
prefix = col_lower.split("_")[0]
if prefix in ["generation", "pass", "inference"]:
@ -181,7 +189,11 @@ def detect_custom_format_heuristic(dataset):
def get_priority_score(col_name):
col_lower = col_name.lower()
return sum(score for pattern, score in priority_patterns.items() if pattern in col_lower)
return sum(
score
for pattern, score in priority_patterns.items()
if pattern in col_lower
)
def get_content_length(col_name):
try:
@ -195,7 +207,9 @@ def detect_custom_format_heuristic(dataset):
score = 10
if role_type == "user":
col_lower = col_name.lower()
if "task" in col_lower and not any(kw in col_lower for kw in user_words_high_priority):
if "task" in col_lower and not any(
kw in col_lower for kw in user_words_high_priority
):
score -= 15
score += get_priority_score(col_name)
if role_type in ["assistant", "user"]:
@ -219,12 +233,19 @@ def detect_custom_format_heuristic(dataset):
return score
content_columns = [col for col in all_columns if not is_metadata(col)]
assistant_potential = [col for col in content_columns if has_keyword(col, assistant_words)]
assistant_potential = [
col for col in content_columns if has_keyword(col, assistant_words)
]
user_potential = [col for col in content_columns if has_keyword(col, user_words)]
assistant_candidates = [
(col, score)
for col in assistant_potential
if (score := score_column(col, assistant_words, "assistant", len(assistant_potential))) > 0
if (
score := score_column(
col, assistant_words, "assistant", len(assistant_potential)
)
)
> 0
]
if assistant_candidates:
assistant_candidates.sort(key = lambda item: item[1], reverse = True)
@ -399,7 +420,9 @@ def detect_multimodal_dataset(dataset):
audio_columns.append(col_name)
modality_types.add("audio")
if audio_columns:
multimodal_columns = [col for col in multimodal_columns if col not in set(audio_columns)]
multimodal_columns = [
col for col in multimodal_columns if col not in set(audio_columns)
]
detected_text_col = None
if audio_columns:
@ -454,7 +477,9 @@ def detect_vlm_dataset_structure(dataset):
and isinstance(content[0], dict)
and "type" in content[0]
):
has_index = any("index" in item for item in content if isinstance(item, dict))
has_index = any(
"index" in item for item in content if isinstance(item, dict)
)
if has_index and "images" in column_names:
return {
"format": "vlm_messages_llava",
@ -463,7 +488,9 @@ def detect_vlm_dataset_structure(dataset):
"image_column": "images",
"text_column": None,
}
has_image = any("image" in item for item in content if isinstance(item, dict))
has_image = any(
"image" in item for item in content if isinstance(item, dict)
)
if has_image:
return {
"format": "vlm_messages",
@ -555,9 +582,9 @@ def detect_vlm_dataset_structure(dataset):
image_candidates = []
for col in column_names:
value = sample[col]
if any(_keyword_in_column(keyword, col) for keyword in image_keywords) or _is_image_value(
value
):
if any(
_keyword_in_column(keyword, col) for keyword in image_keywords
) or _is_image_value(value):
if hasattr(value, "size") and hasattr(value, "mode"):
score = 100
elif isinstance(value, dict) and ("bytes" in value or "path" in value):
@ -726,7 +753,9 @@ def _standardize_sharegpt_row(row: dict[str, Any], chat_column: str) -> dict[str
if not isinstance(message, dict):
continue
role = message.get("role") or message.get("from")
content = message.get("content") if "content" in message else message.get("value")
content = (
message.get("content") if "content" in message else message.get("value")
)
messages.append(
{
"role": _ROLE_MAP.get(str(role), str(role or "user")),

12
studio/backend/hub/utils/download_manifest.py
View file

@ -59,7 +59,9 @@ _LEGACY_MARKER_VERSION = 1
# Verbatim phrase the worker emits on a degraded completion and the download
# lifecycle escalates to a warning log. Shared so the emit and match stay coupled.
MANIFEST_DEGRADED_MARKER = "completed without a manifest so partial detection is degraded"
MANIFEST_DEGRADED_MARKER = (
"completed without a manifest so partial detection is degraded"
)
@dataclass(frozen = True)
@ -473,14 +475,18 @@ def _iter_variant_state_files(
yield _variant_from_state_file(entry, variant), entry
def iter_variant_manifests(repo_type: RepoType, repo_id: str) -> Iterator[tuple[str, Path]]:
def iter_variant_manifests(
repo_type: RepoType, repo_id: str
) -> Iterator[tuple[str, Path]]:
"""Yield (variant, manifest_path) for every variant-keyed manifest
written for this repo. Used by is_gguf_repo_partial to enumerate all
variants present on disk so the all-variants-broken gate can run."""
yield from _iter_variant_state_files(manifests_dir(), repo_type, repo_id)
def iter_variant_markers(repo_type: RepoType, repo_id: str) -> Iterator[tuple[str, Path]]:
def iter_variant_markers(
repo_type: RepoType, repo_id: str
) -> Iterator[tuple[str, Path]]:
"""Yield (variant, marker_path) for every variant-keyed cancel marker.
Companion to iter_variant_manifests: catches variants cancelled
before download-start ever wrote a manifest (very early failures)."""

59
studio/backend/hub/utils/download_registry.py
View file

@ -113,7 +113,9 @@ def _worker_breadcrumb_path(key: str) -> Optional[Path]:
return parent / f"{safe}.json"
def write_worker_breadcrumb(key: str, pid: int, metadata: Optional["DownloadMetadata"]) -> None:
def write_worker_breadcrumb(
key: str, pid: int, metadata: Optional["DownloadMetadata"]
) -> None:
"""Record a live worker's PID so a restarted backend can reap it. Best
effort: a write failure only forfeits boot-time reaping for this worker,
still covered by the worker's own parent-death watchdog."""
@ -369,7 +371,9 @@ def _iter_active_snapshot_dirs(repo_type: str, repo_id: str) -> Iterator[Path]:
yield snapshot
def _manifest_verifies_against_active_cache(repo_type: str, repo_id: str, manifest) -> bool:
def _manifest_verifies_against_active_cache(
repo_type: str, repo_id: str, manifest
) -> bool:
from hub.utils import download_manifest
for snapshot_dir in _iter_active_snapshot_dirs(repo_type, repo_id):
if download_manifest.verify_against_disk(manifest, snapshot_dir).ok:
@ -377,7 +381,9 @@ def _manifest_verifies_against_active_cache(repo_type: str, repo_id: str, manife
return False
def _manifest_has_active_incomplete_blobs(repo_type: str, repo_id: str, manifest) -> bool:
def _manifest_has_active_incomplete_blobs(
repo_type: str, repo_id: str, manifest
) -> bool:
if not getattr(manifest, "variant", None):
return has_active_incomplete_blobs(repo_type, repo_id)
expected_hashes = frozenset(
@ -386,7 +392,9 @@ def _manifest_has_active_incomplete_blobs(repo_type: str, repo_id: str, manifest
if not expected_hashes:
return has_active_incomplete_blobs(repo_type, repo_id)
return bool(
incomplete_blob_hashes(repo_type, repo_id, active_only = True).intersection(expected_hashes)
incomplete_blob_hashes(repo_type, repo_id, active_only = True).intersection(
expected_hashes
)
)
@ -401,7 +409,9 @@ def _is_transport_marker_file(path: Path) -> bool:
# Matches ".transport", its tmps, and variant-scoped ".transport.gguf-*".
# Real HF cache entries (blobs/refs/snapshots/.no_exist) never start with
# ".transport.".
return path.name == TRANSPORT_MARKER_NAME or path.name.startswith(f"{TRANSPORT_MARKER_NAME}.")
return path.name == TRANSPORT_MARKER_NAME or path.name.startswith(
f"{TRANSPORT_MARKER_NAME}."
)
def _companion_marker_path(entry: Path) -> Path:
@ -519,9 +529,13 @@ def prepare_cache_for_transport(
total_purged += _purge_incomplete_blobs(entry, only_blob_hashes, protected)
else:
if _read_marker(entry, variant) != mode:
total_purged += _purge_incomplete_blobs(entry, only_blob_hashes, protected)
total_purged += _purge_incomplete_blobs(
entry, only_blob_hashes, protected
)
if companion_blob_hashes and _read_companion_marker(entry) != mode:
total_purged += _purge_incomplete_blobs(entry, companion_blob_hashes, protected)
total_purged += _purge_incomplete_blobs(
entry, companion_blob_hashes, protected
)
_write_marker(entry, mode, variant)
if has_companion:
_write_companion_marker(entry, mode)
@ -565,7 +579,9 @@ def purge_empty_marker_dir(
contents = list(entry.iterdir())
except OSError:
continue
if not contents or not all(_is_transport_marker_file(item) for item in contents):
if not contents or not all(
_is_transport_marker_file(item) for item in contents
):
continue
own_name = _marker_path(entry, variant).name
own_markers = [
@ -638,7 +654,9 @@ def incomplete_blob_hashes(
return out
def completed_blob_bytes(repo_type: str, repo_id: str, blob_hashes: frozenset[str]) -> int:
def completed_blob_bytes(
repo_type: str, repo_id: str, blob_hashes: frozenset[str]
) -> int:
"""Sum finalized blob bytes for *blob_hashes* in the active HF cache root.
A worker only writes to the active ``HF_HUB_CACHE`` root, so a baseline must
@ -661,7 +679,9 @@ def completed_blob_bytes(repo_type: str, repo_id: str, blob_hashes: frozenset[st
return total
def existing_blob_bytes(repo_type: str, repo_id: str, blob_hashes: frozenset[str]) -> int:
def existing_blob_bytes(
repo_type: str, repo_id: str, blob_hashes: frozenset[str]
) -> int:
"""Bytes already on disk (finalized + ``.incomplete``) for *blob_hashes* in
the active HF cache root. A blob is in exactly one state, so summing both
candidate names never double-counts. Used to size what a (possibly resumed)
@ -864,7 +884,8 @@ class DownloadRegistry:
pending_generation = self._pending_cancel.get(key)
metadata = self._metadata.get(key)
should_cancel = current == "cancelling" or (
has_pending_cancel and self._generation_matches_locked(key, pending_generation)
has_pending_cancel
and self._generation_matches_locked(key, pending_generation)
)
terminal_state: JobState = "cancelled" if should_cancel else "error"
marker_transport = self._cancel_marker_transports.pop(key, None)
@ -1098,7 +1119,9 @@ class DownloadRegistry:
repo_type = repo_type,
repo_id = repo_id,
variant = variant,
transport = metadata_transport if metadata_transport is not None else transport,
transport = metadata_transport
if metadata_transport is not None
else transport,
cancel_marker_transport = cancel_marker_transport,
blob_hashes = requested_hashes,
progress_blob_hashes = requested_progress_hashes,
@ -1132,7 +1155,9 @@ class DownloadRegistry:
return (metadata.variant or "").strip().lower() or None
return variant_from_key(key)
def _delete_blocked_by_active_locked(self, repo_id: str, variant: Optional[str]) -> bool:
def _delete_blocked_by_active_locked(
self, repo_id: str, variant: Optional[str]
) -> bool:
"""Whether an active download conflicts with deleting *repo_id*/*variant*.
A whole-repo delete (``variant is None``) conflicts with any active
@ -1204,7 +1229,9 @@ class DownloadRegistry:
if repo_key:
candidate_keys = list(self._repo_active.get(repo_key, set()))
else:
candidate_keys = [key for active in self._repo_active.values() for key in active]
candidate_keys = [
key for active in self._repo_active.values() for key in active
]
# An XET->HTTP retry handoff briefly drops its key from _repo_active
# while its job stays active; include those released-but-active jobs
# so the waiting retry still lists and can be adopted or cancelled.
@ -1409,7 +1436,9 @@ class DownloadRegistry:
try:
proc.wait(timeout = max(0.0, deadline - time.monotonic()))
except subprocess.TimeoutExpired:
logger.warning(f"shutdown: {kind} worker for {key} did not exit after kill")
logger.warning(
f"shutdown: {kind} worker for {key} did not exit after kill"
)
except Exception:
pass
# Mark only genuinely interrupted workers (rc != 0, or None on wait

8
studio/backend/hub/utils/gguf.py
View file

@ -303,7 +303,9 @@ def list_empty_gguf_variant_dirs(repo_id: str) -> set[str]:
return {label for key, label in empty.items() if key not in nonempty}
def list_gguf_variants_from_hf_cache(repo_id: str) -> Optional[tuple[list[GgufVariantInfo], bool]]:
def list_gguf_variants_from_hf_cache(
repo_id: str,
) -> Optional[tuple[list[GgufVariantInfo], bool]]:
for snapshot in iter_hf_cache_snapshots(repo_id):
variants, has_vision = list_local_gguf_variants(str(snapshot))
if variants or has_vision:
@ -445,7 +447,9 @@ def list_gguf_variants(
quant = extract_quant_label(filename)
if is_big_endian_gguf_path(filename, quant):
continue
quant_totals[quant] = quant_totals.get(quant, 0) + int(getattr(sibling, "size", 0) or 0)
quant_totals[quant] = quant_totals.get(quant, 0) + int(
getattr(sibling, "size", 0) or 0
)
quant_first_file.setdefault(quant, filename)
for quant, total_size in quant_totals.items():

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