diff --git a/routes/cookbook_helpers.py b/routes/cookbook_helpers.py index b2401d5..7847e35 100644 --- a/routes/cookbook_helpers.py +++ b/routes/cookbook_helpers.py @@ -124,6 +124,79 @@ def _local_tooling_path_export(executable: str) -> str: return f'export PATH="{esc}:$PATH"' +def _cached_model_scan_script(model_dirs: list[str] | None = None) -> str: + """Build the standalone Python scanner used by /api/model/cached.""" + lines = [ + "import json, os", + "models = []", + "seen = set()", + "BLOCKED_ROOTS = ('/sys', '/proc', '/dev', '/run', '/var/run')", + "def safe_path(p):", + " try:", + " rp = os.path.realpath(os.path.expanduser(p))", + " return not any(rp == b or rp.startswith(b + os.sep) for b in BLOCKED_ROOTS)", + " except Exception:", + " return False", + "def safe_walk(top):", + " if not safe_path(top): return", + " for root, dirs, fns in os.walk(top, followlinks=False):", + " dirs[:] = [d for d in dirs if not os.path.islink(os.path.join(root, d)) and safe_path(os.path.join(root, d))]", + " yield root, dirs, fns", + "def scan_hf(cache):", + " if not os.path.isdir(cache): return", + " for d in sorted(os.listdir(cache)):", + " if not d.startswith('models--'): continue", + " rid = d.replace('models--','').replace('--','/')", + " if rid in seen: continue", + " seen.add(rid)", + " blobs = os.path.join(cache, d, 'blobs')", + " sz, nf, ic = 0, 0, False", + " if os.path.isdir(blobs):", + " for f in os.scandir(blobs):", + " if f.is_file(): nf += 1; sz += f.stat().st_size", + " if f.name.endswith('.incomplete'): ic = True", + " snap = os.path.join(cache, d, 'snapshots')", + " is_diffusion = False; is_gguf = False", + " if os.path.isdir(snap):", + " for sd in os.listdir(snap):", + " sf = os.path.join(snap, sd)", + " if not os.path.isdir(sf): continue", + " if os.path.exists(os.path.join(sf, 'model_index.json')): is_diffusion = True", + " try:", + " if any(x.endswith('.gguf') for x in os.listdir(sf)): is_gguf = True", + " except Exception: pass", + " models.append({'repo_id':rid,'size_bytes':sz,'nb_files':nf,'has_incomplete':ic,'path':cache,'is_diffusion':is_diffusion,'is_gguf':is_gguf})", + "def scan_dir(p):", + " if not os.path.isdir(p) or not safe_path(p): return", + " for d in sorted(os.listdir(p)):", + " if d.startswith('.'): continue", + " if d.startswith('models--'): continue", + " fp = os.path.join(p, d)", + " if not os.path.isdir(fp) or os.path.islink(fp) or not safe_path(fp): continue", + " if d in seen: continue", + " is_model = False; is_gguf = False", + " for root, dirs, fns in safe_walk(fp):", + " for fn in fns:", + " if fn.endswith('.gguf'): is_gguf = True; is_model = True", + " elif fn == 'config.json' or fn.endswith('.safetensors') or fn.endswith('.bin'): is_model = True", + " if is_model: break", + " if not is_model: continue", + " seen.add(d)", + " sz, nf = 0, 0", + " for dp, _, fns in safe_walk(fp):", + " for fn in fns:", + " try: nf += 1; sz += os.path.getsize(os.path.join(dp, fn))", + " except Exception: pass", + " is_diff = os.path.exists(os.path.join(fp, 'model_index.json'))", + " models.append({'repo_id':d,'size_bytes':sz,'nb_files':nf,'has_incomplete':False,'path':p,'is_local_dir':True,'is_diffusion':is_diff,'is_gguf':is_gguf})", + "scan_hf(os.path.expanduser('~/.cache/huggingface/hub'))", + ] + for model_dir in model_dirs or []: + lines.append(f"scan_dir(os.path.expanduser({model_dir!r}))") + lines.append("print(json.dumps(models))") + return "\n".join(lines) + "\n" + + def _ps_squote(v: str) -> str: """Escape a value for PowerShell single-quoted string interpolation. Belt-and-suspenders on top of _validate_token's regex — if the regex diff --git a/routes/cookbook_routes.py b/routes/cookbook_routes.py index 3c6bf5b..cc10763 100644 --- a/routes/cookbook_routes.py +++ b/routes/cookbook_routes.py @@ -37,7 +37,7 @@ from routes.cookbook_helpers import ( _validate_local_dir, _validate_ssh_port, _validate_gpus, _shell_path, _ps_squote, _bash_squote, _validate_serve_cmd, _parse_serve_phase, _safe_env_prefix, _local_tooling_path_export, _append_serve_preflight_exit_lines, - _append_serve_exit_code_lines, + _append_serve_exit_code_lines, _cached_model_scan_script, ModelDownloadRequest, ServeRequest, ) @@ -647,84 +647,13 @@ def setup_cookbook_routes() -> APIRouter: raise HTTPException(400, "Invalid ssh_port") TMUX_LOG_DIR.mkdir(parents=True, exist_ok=True) - paths_code = "import json, os\n" - paths_code += "models = []\n" - paths_code += "seen = set()\n" - paths_code += "BLOCKED_ROOTS = ('/sys', '/proc', '/dev', '/run', '/var/run')\n" - paths_code += "def safe_path(p):\n" - paths_code += " try:\n" - paths_code += " rp = os.path.realpath(os.path.expanduser(p))\n" - paths_code += " return not any(rp == b or rp.startswith(b + os.sep) for b in BLOCKED_ROOTS)\n" - paths_code += " except Exception:\n" - paths_code += " return False\n" - paths_code += "def safe_walk(top):\n" - paths_code += " if not safe_path(top): return\n" - paths_code += " for root, dirs, fns in os.walk(top, followlinks=False):\n" - paths_code += " dirs[:] = [d for d in dirs if not os.path.islink(os.path.join(root, d)) and safe_path(os.path.join(root, d))]\n" - paths_code += " yield root, dirs, fns\n" - # Scan HF cache format (models-- directories with blobs/) - paths_code += "def scan_hf(cache):\n" - paths_code += " if not os.path.isdir(cache): return\n" - paths_code += " for d in sorted(os.listdir(cache)):\n" - paths_code += " if not d.startswith('models--'): continue\n" - paths_code += " rid = d.replace('models--','').replace('--','/')\n" - paths_code += " if rid in seen: continue\n" - paths_code += " seen.add(rid)\n" - paths_code += " blobs = os.path.join(cache, d, 'blobs')\n" - paths_code += " sz, nf, ic = 0, 0, False\n" - paths_code += " if os.path.isdir(blobs):\n" - paths_code += " for f in os.scandir(blobs):\n" - paths_code += " if f.is_file(): nf += 1; sz += f.stat().st_size\n" - paths_code += " if f.name.endswith('.incomplete'): ic = True\n" - paths_code += " # Check if it's an LLM (has config.json with model_type) vs diffusion (has model_index.json)\n" - paths_code += " snap = os.path.join(cache, d, 'snapshots')\n" - paths_code += " is_diffusion = False; is_gguf = False\n" - paths_code += " if os.path.isdir(snap):\n" - paths_code += " for sd in os.listdir(snap):\n" - paths_code += " sf = os.path.join(snap, sd)\n" - paths_code += " if not os.path.isdir(sf): continue\n" - paths_code += " if os.path.exists(os.path.join(sf, 'model_index.json')): is_diffusion = True\n" - paths_code += " try:\n" - paths_code += " if any(x.endswith('.gguf') for x in os.listdir(sf)): is_gguf = True\n" - paths_code += " except Exception: pass\n" - paths_code += " models.append({'repo_id':rid,'size_bytes':sz,'nb_files':nf,'has_incomplete':ic,'path':cache,'is_diffusion':is_diffusion,'is_gguf':is_gguf})\n" - # Scan plain directory (each subdirectory = a model if it has model files) - paths_code += "def scan_dir(p):\n" - paths_code += " if not os.path.isdir(p) or not safe_path(p): return\n" - paths_code += " for d in sorted(os.listdir(p)):\n" - paths_code += " if d.startswith('.'): continue\n" - paths_code += " fp = os.path.join(p, d)\n" - paths_code += " if not os.path.isdir(fp) or os.path.islink(fp) or not safe_path(fp): continue\n" - paths_code += " if d in seen: continue\n" - paths_code += " # Check if it looks like a model (has config.json, safetensors, bin, or gguf)\n" - paths_code += " is_model = False; is_gguf = False\n" - paths_code += " for root, dirs, fns in safe_walk(fp):\n" - paths_code += " for fn in fns:\n" - paths_code += " if fn.endswith('.gguf'): is_gguf = True; is_model = True\n" - paths_code += " elif fn == 'config.json' or fn.endswith('.safetensors') or fn.endswith('.bin'): is_model = True\n" - paths_code += " if is_model: break\n" - paths_code += " if not is_model: continue\n" - paths_code += " seen.add(d)\n" - paths_code += " sz, nf = 0, 0\n" - paths_code += " for dp, _, fns in safe_walk(fp):\n" - paths_code += " for fn in fns:\n" - paths_code += " try: nf += 1; sz += os.path.getsize(os.path.join(dp, fn))\n" - paths_code += " except Exception: pass\n" - paths_code += " is_diff = os.path.exists(os.path.join(fp, 'model_index.json'))\n" - paths_code += " models.append({'repo_id':d,'size_bytes':sz,'nb_files':nf,'has_incomplete':False,'path':p,'is_local_dir':True,'is_diffusion':is_diff,'is_gguf':is_gguf})\n" - # Always scan HF cache - paths_code += "scan_hf(os.path.expanduser('~/.cache/huggingface/hub'))\n" - # Also scan custom model dirs (comma-separated) if specified + model_dirs = [] if model_dir: for d in model_dir.split(','): d = d.strip() - if d and d != '~/.cache/huggingface/hub': - # repr() encodes the dir as a properly-escaped Python string - # literal. The old f"...'{d}'..." broke out of the quotes on - # any `'` in the value, injecting arbitrary Python that then - # ran locally or over ssh. - paths_code += f"scan_dir(os.path.expanduser({d!r}))\n" - paths_code += "print(json.dumps(models))\n" + if d: + model_dirs.append(d) + paths_code = _cached_model_scan_script(model_dirs) scan_py = TMUX_LOG_DIR / "scan_cache.py" scan_py.write_text(paths_code, encoding="utf-8") @@ -779,6 +708,8 @@ def setup_cookbook_routes() -> APIRouter: } if m.get("is_local_dir"): entry["is_local_dir"] = True + if m.get("is_gguf"): + entry["is_gguf"] = True models.append(entry) except Exception as e: logger.warning(f"Failed to parse cached models: {e}") diff --git a/tests/test_cookbook_helpers.py b/tests/test_cookbook_helpers.py index 566b99f..5935115 100644 --- a/tests/test_cookbook_helpers.py +++ b/tests/test_cookbook_helpers.py @@ -1,7 +1,12 @@ +import json +import subprocess +import sys + import pytest from fastapi import HTTPException from routes.cookbook_helpers import ( + _cached_model_scan_script, _append_serve_exit_code_lines, _append_serve_preflight_exit_lines, _local_tooling_path_export, @@ -92,3 +97,29 @@ def test_serve_runner_preserves_command_exit_code(): assert "ODYSSEUS_CMD_EXIT=$?" in script assert 'echo "=== Process exited with code $ODYSSEUS_CMD_EXIT ==="' in script assert 'echo "=== Process exited with code $? ==="' not in script + + +def test_cached_model_scan_reports_plain_dir_gguf(tmp_path): + """Custom download dirs may sit inside the HF hub cache and contain plain + per-model folders. They must show up in Serve and keep the GGUF signal.""" + plain = tmp_path / "Qwen3.6-27B" + plain.mkdir() + (plain / "Qwen3.6-27B-Q4_K_M.gguf").write_bytes(b"gguf") + + hf_internal = tmp_path / "models--Qwen--Qwen3.6-27B" + (hf_internal / "snapshots" / "abc").mkdir(parents=True) + (hf_internal / "snapshots" / "abc" / "model.safetensors").write_bytes(b"safe") + + scan_py = tmp_path / "scan_cache.py" + scan_py.write_text(_cached_model_scan_script([str(tmp_path)]), encoding="utf-8") + proc = subprocess.run( + [sys.executable, str(scan_py)], + check=True, + capture_output=True, + text=True, + ) + + by_repo = {m["repo_id"]: m for m in json.loads(proc.stdout)} + assert "models--Qwen--Qwen3.6-27B" not in by_repo + assert by_repo["Qwen3.6-27B"]["is_local_dir"] is True + assert by_repo["Qwen3.6-27B"]["is_gguf"] is True