diff --git a/README.md b/README.md index 7410b18c..ec8cddcd 100644 --- a/README.md +++ b/README.md @@ -615,13 +615,33 @@ grid --remote allocator join --dedicated # Controller/relay host: register exact artifacts, inspect, then enable actuation grid --local allocator model set --grid allocator-control \ --memory-mb --artifact-sha256 \ - --artifact-source hf://owner/repo/model.gguf --artifact-size-mb \ + --artifact-source hf://owner/repo@/model.gguf --artifact-size-mb \ --min-replicas 0 --max-replicas 3 grid --local allocator mode recommend --grid allocator-control grid --local allocator status --grid allocator-control grid --local allocator mode automatic --grid allocator-control ``` +The optional model scout watches trusted Hugging Face publishers for immutable GGUF and vLLM +releases, rejects unknown licenses and mutable revisions, checks the current fleet's runtime, +memory, and disk fit, and writes ranked proposals. Hub popularity is discovery priority—not model +quality. Qualification always sends real requests through Grid to a bounded canary and records the +result as allocator evaluation evidence: + +```bash +grid --local allocator scout run --grid allocator-control --search coder +grid --local allocator scout status --grid allocator-control +grid --local allocator scout benchmark --grid allocator-control \ + --inference-grid --deploy-canary + +# Or refresh proposals every six hours. This discovers only; it never silently replaces a model. +grid --local allocator scout watch --grid allocator-control --interval 21600 +``` + +Canary profiles are exact-revision and exact-digest pinned, have zero required replicas and at most +one replica, and do not retire incumbents. A candidate enters the live portfolio only after it +passes the real quality floor; replacement remains a separate, reviewable allocation decision. + Allocator enrollment verifies the managed llama.cpp runtime and installs Grid's version- and SHA-256-pinned build when it is absent. It also waits briefly for a just-started provider identity to become visible at the relay, so the two fresh-node commands above are safe to run back to back. diff --git a/cli/allocator_scout.py b/cli/allocator_scout.py new file mode 100644 index 00000000..52c89a68 --- /dev/null +++ b/cli/allocator_scout.py @@ -0,0 +1,325 @@ +"""Operator-facing autonomous model discovery and real-canary qualification.""" + +from __future__ import annotations + +import argparse +import json +import subprocess +import time +from pathlib import Path +from typing import Any +from urllib.parse import quote + +from local import config, runtime +from shared import paths +from shared.allocator.models import ModelProfile, stable_digest +from shared.allocator.scout import ( + DEFAULT_ALLOWED_LICENSES, + DEFAULT_TRUSTED_AUTHORS, + HuggingFaceDiscovery, + ScoutPolicy, + ScoutProposal, + benchmark_candidate, + build_proposals, + load_scout_state, + proposals_from_state, + save_scout_state, +) + + +def cmd_allocator_scout_run(args: argparse.Namespace) -> int: + cfg = config.select_grid(getattr(args, "grid", None)) + policy = _policy(args) + from .allocator import _request + + status = _request(cfg, "GET", "/allocator/status") + discovery = HuggingFaceDiscovery(base_url=args.hub_url) + try: + candidates = discovery.discover(policy, search=args.search) + finally: + discovery.close() + proposals = build_proposals(candidates, status) + state_path = _state_path(cfg, getattr(args, "state_file", None)) + save_scout_state(state_path, proposals, policy=policy) + payload = { + "state_file": str(state_path), + "discovered": len(candidates), + "benchmark_ready": sum(item.state == "benchmark-ready" for item in proposals), + "discovery_issues": list(discovery.issues), + "proposals": [item.to_dict() for item in proposals], + } + if args.json: + print(json.dumps(payload, indent=2)) + else: + print( + f"Allocator scout: {len(candidates)} immutable candidates · " + f"{payload['benchmark_ready']} fit this fleet" + ) + _print_proposals(proposals) + if discovery.issues: + print(f" warning {len(discovery.issues)} Hub source(s) skipped") + for issue in discovery.issues[:5]: + print(f" {issue}") + print(f" state {state_path}") + return 0 + + +def cmd_allocator_scout_status(args: argparse.Namespace) -> int: + cfg = config.select_grid(getattr(args, "grid", None)) + state_path = _state_path(cfg, getattr(args, "state_file", None)) + state = load_scout_state(state_path) + proposals = proposals_from_state(state) + if args.json: + print(json.dumps({**state, "state_file": str(state_path)}, indent=2)) + else: + print( + f"Allocator scout: {len(proposals)} proposals · " + f"{sum(item.state == 'qualified' for item in proposals)} qualified" + ) + _print_proposals(proposals) + print(f" state {state_path}") + return 0 + + +def cmd_allocator_scout_benchmark(args: argparse.Namespace) -> int: + cfg = config.select_grid(getattr(args, "grid", None)) + state_path = _state_path(cfg, getattr(args, "state_file", None)) + state = load_scout_state(state_path) + proposals = list(proposals_from_state(state)) + matches = [item for item in proposals if item.proposal_id == args.proposal] + if not matches: + raise SystemExit(f"Scout proposal not found: {args.proposal}") + proposal = matches[0] + if proposal.state not in ("benchmark-ready", "benchmark-failed", "qualified"): + raise SystemExit( + f"Scout proposal {proposal.proposal_id} is {proposal.state}; it cannot be benchmarked." + ) + if not args.deploy_canary: + raise SystemExit( + "Real benchmarking requires --deploy-canary. It creates a min=0/max=1 immutable " + "profile and never retires an incumbent." + ) + + _put_canary_profile(cfg, proposal, args) + runner = _GridBenchmarkRunner( + grid=args.inference_grid, + startup_timeout=args.startup_timeout, + request_timeout=args.request_timeout, + ) + updated, samples = benchmark_candidate( + proposal, + runner, + workloads=args.workloads, + ) + _record_evaluations(cfg, updated, samples, args) + proposals[proposals.index(proposal)] = updated + sample_state = dict(state.get("benchmark_samples") or {}) + sample_state[updated.proposal_id] = samples + policy = _policy_from_state(state) + save_scout_state( + state_path, + proposals, + policy=policy, + benchmark_samples=sample_state, + ) + payload = { + "proposal": updated.to_dict(), + "samples": [ + { + "workload": sample.workload, + "quality": sample.quality, + "latency_ms": sample.latency_ms, + "output_units": sample.output_units, + "error": sample.error, + } + for sample in samples + ], + } + if args.json: + print(json.dumps(payload, indent=2)) + else: + print( + f"Scout benchmark {updated.state}: {updated.candidate.model_id} · " + f"quality {updated.benchmark_quality:.2f} · " + f"median {updated.benchmark_latency_ms or 0:.0f} ms" + ) + if updated.compare_models: + print(" compare with " + ", ".join(updated.compare_models)) + return 0 if updated.state == "qualified" else 1 + + +def cmd_allocator_scout_watch(args: argparse.Namespace) -> int: + """Run repeated discovery cycles; canaries remain explicit and operator-observable.""" + + cycles = 0 + try: + while args.max_cycles == 0 or cycles < args.max_cycles: + cmd_allocator_scout_run(args) + cycles += 1 + if args.max_cycles and cycles >= args.max_cycles: + break + time.sleep(args.interval) + except KeyboardInterrupt: + print("Allocator scout stopped.") + return 0 + + +def _policy(args: argparse.Namespace) -> ScoutPolicy: + return ScoutPolicy( + workloads=tuple(args.workloads or ("coding", "general", "research")), + runtimes=tuple(args.runtimes or ("llama.cpp", "vllm")), + trusted_authors=tuple(args.authors or DEFAULT_TRUSTED_AUTHORS), + allowed_licenses=tuple(args.licenses or DEFAULT_ALLOWED_LICENSES), + quantizations=tuple(args.quantizations or ("Q4_K_M", "Q5_K_M", "Q4_K_S", "Q5_K_S", "Q6_K")), + max_results=args.limit, + max_repositories=args.inspect, + max_artifact_size_mb=args.max_artifact_size_mb, + min_downloads=args.min_downloads, + ) + + +def _policy_from_state(state: dict[str, Any]) -> ScoutPolicy | None: + raw = state.get("policy") + if not isinstance(raw, dict) or not raw: + return None + try: + return ScoutPolicy(**raw) + except (TypeError, ValueError): + return None + + +def _state_path(cfg: dict[str, Any], explicit: str | None) -> Path: + if explicit: + return Path(explicit).expanduser() + identity = str(cfg.get("grid_id") or runtime.grid_url(cfg)) + scope = stable_digest({"grid": identity})[:20] + return paths.grid_home() / "allocator-scout" / f"{scope}.json" + + +def _put_canary_profile( + cfg: dict[str, Any], proposal: ScoutProposal, args: argparse.Namespace +) -> None: + from .allocator import _control_token, _request + + candidate = proposal.candidate + backends = ("cuda",) if candidate.runtime == "vllm" else ("cpu", "cuda", "metal") + profile = ModelProfile( + model_id=candidate.model_id, + memory_mb=candidate.estimated_memory_mb, + runtimes=(candidate.runtime,), + backends=backends, + min_replicas=0, + max_replicas=1, + target_utilization=0.70, + replica_concurrency=1, + expected_service_seconds=15.0, + latency_slo_ms=float(args.request_timeout) * 1_000.0, + priority=50, + load_seconds=float(args.startup_timeout), + warm_seconds=min(120.0, float(args.startup_timeout)), + min_residency_seconds=300.0, + scale_down_cooldown_seconds=300.0, + min_failure_domains=1, + min_gpu_count=1 if candidate.runtime == "vllm" else 0, + artifact_sha256=candidate.artifact_sha256, + artifact_source=candidate.artifact_source, + artifact_size_mb=candidate.artifact_size_mb, + max_colocated_models=1 if candidate.estimated_memory_mb >= 16_000 else 0, + workload_scores=candidate.workload_scores, + ) + _request( + cfg, + "PUT", + f"/allocator/models/{quote(candidate.model_id, safe='')}", + body=profile.to_dict(), + token=_control_token(cfg, getattr(args, "token_file", None)), + allow_insecure_http=getattr(args, "allow_insecure_http", False), + ) + + +def _record_evaluations(cfg, proposal, samples, args) -> None: + from .allocator import _control_token, _request + + token = _control_token(cfg, getattr(args, "token_file", None)) + for sample in samples: + _request( + cfg, + "POST", + "/allocator/evaluations", + body={ + "model_id": proposal.candidate.model_id, + "workload": sample.workload, + "artifact_sha256": proposal.candidate.artifact_sha256, + "quality": sample.quality, + "error": sample.error, + "latency_ms": sample.latency_ms, + "output_units": sample.output_units, + }, + token=token, + allow_insecure_http=getattr(args, "allow_insecure_http", False), + ) + + +class _GridBenchmarkRunner: + def __init__(self, *, grid: str, startup_timeout: float, request_timeout: float) -> None: + self.grid = grid + self.startup_timeout = startup_timeout + self.request_timeout = request_timeout + + def __call__(self, model: str, prompt: str) -> tuple[str, float]: + deadline = time.monotonic() + self.startup_timeout + alias = model.removesuffix(".gguf") + last_error = "" + while True: + started = time.monotonic() + command = runtime.cli_command() + [ + "--remote", + "chat", + "--grid", + self.grid, + "-m", + alias, + "--allow-self-provider", + "--timeout", + str(self.request_timeout), + "--json", + prompt, + ] + completed = subprocess.run( + command, + capture_output=True, + text=True, + timeout=self.request_timeout + 10.0, + check=False, + ) + elapsed_ms = (time.monotonic() - started) * 1_000.0 + if completed.returncode == 0: + try: + payload = json.loads(completed.stdout) + text = str(payload["choices"][0]["message"]["content"]) + timings = payload.get("timings") or {} + latency = float(timings.get("predicted_ms") or elapsed_ms) + return text, latency + except (KeyError, IndexError, TypeError, ValueError, json.JSONDecodeError) as exc: + last_error = f"malformed Grid response: {exc}" + else: + last_error = completed.stderr.strip() or completed.stdout.strip() + remaining = deadline - time.monotonic() + if remaining <= 0: + raise RuntimeError(last_error or "canary did not become routable") + time.sleep(min(5.0, remaining)) + + +def _print_proposals(proposals: list[ScoutProposal] | tuple[ScoutProposal, ...]) -> None: + for item in proposals[:10]: + ready_hosts = sum(fit.fits for fit in item.fits) + quality = ( + f" · quality {item.benchmark_quality:.2f}" + if item.benchmark_quality is not None + else "" + ) + print( + f" {item.state:<16} {item.candidate.model_id} · {item.candidate.runtime} · " + f"{item.candidate.artifact_size_mb} MB · {ready_hosts} hosts{quality}" + ) + print(f" id {item.proposal_id} · {item.candidate.repo_id}@{item.candidate.revision[:12]}") diff --git a/cli/parser.py b/cli/parser.py index 00628bbe..91db6a2d 100644 --- a/cli/parser.py +++ b/cli/parser.py @@ -30,6 +30,12 @@ cmd_allocator_tick, cmd_allocator_token_write, ) +from .allocator_scout import ( + cmd_allocator_scout_benchmark, + cmd_allocator_scout_run, + cmd_allocator_scout_status, + cmd_allocator_scout_watch, +) from shared.allocator.scenario import SCENARIO_STRATEGIES from .allocator_scenario import ( @@ -514,7 +520,7 @@ def _add_allocator(sub) -> None: metavar="URI", help=( "Authenticated immutable source for autonomous loading; managed llama.cpp accepts " - "an exact hf://owner/repo/path.gguf URI." + "an exact hf://owner/repo[@commit]/path.gguf URI." ), ) set_model.add_argument( @@ -623,6 +629,66 @@ def _add_allocator(sub) -> None: tick.add_argument("--json", action="store_true", help="Emit machine-readable JSON.") tick.set_defaults(handler=cmd_allocator_tick) + scout = allocator_sub.add_parser( + "scout", + help="Discover immutable open-weight releases and qualify real canaries", + ) + scout_sub = scout.add_subparsers(dest="allocator_scout_command", required=True) + + def add_scout_discovery_options(command) -> None: + command.add_argument("--search", default="", help="Optional Hub model search text.") + command.add_argument("--workload", action="append", dest="workloads", default=[]) + command.add_argument( + "--runtime", action="append", dest="runtimes", choices=("llama.cpp", "vllm"), default=[] + ) + command.add_argument("--author", action="append", dest="authors", default=[]) + command.add_argument("--license", action="append", dest="licenses", default=[]) + command.add_argument("--quantization", action="append", dest="quantizations", default=[]) + command.add_argument("--limit", type=int, default=30) + command.add_argument("--inspect", type=int, default=12) + command.add_argument("--max-artifact-size-mb", type=int, default=100_000) + command.add_argument("--min-downloads", type=int, default=0) + command.add_argument("--hub-url", default="https://huggingface.co") + command.add_argument("--state-file", default=None) + _add_allocator_grid(command) + command.add_argument("--json", action="store_true") + + scout_run = scout_sub.add_parser( + "run", help="Run one discovery and fleet-fit proposal cycle" + ) + add_scout_discovery_options(scout_run) + scout_run.set_defaults(handler=cmd_allocator_scout_run) + + scout_watch = scout_sub.add_parser( + "watch", help="Continuously refresh discovery and proposals" + ) + add_scout_discovery_options(scout_watch) + scout_watch.add_argument("--interval", type=float, default=21_600.0) + scout_watch.add_argument( + "--max-cycles", type=int, default=0, help="Stop after N cycles; 0 runs until interrupted." + ) + scout_watch.set_defaults(handler=cmd_allocator_scout_watch) + + scout_status = scout_sub.add_parser("status", help="Show persisted scout proposals") + scout_status.add_argument("--state-file", default=None) + _add_allocator_grid(scout_status) + scout_status.add_argument("--json", action="store_true") + scout_status.set_defaults(handler=cmd_allocator_scout_status) + + scout_benchmark = scout_sub.add_parser( + "benchmark", help="Deploy one bounded canary and record real evaluation evidence" + ) + scout_benchmark.add_argument("proposal") + scout_benchmark.add_argument("--inference-grid", required=True) + scout_benchmark.add_argument("--workload", action="append", dest="workloads", default=[]) + scout_benchmark.add_argument("--deploy-canary", action="store_true") + scout_benchmark.add_argument("--startup-timeout", type=float, default=900.0) + scout_benchmark.add_argument("--request-timeout", type=float, default=120.0) + scout_benchmark.add_argument("--state-file", default=None) + _add_allocator_grid(scout_benchmark, token=True) + scout_benchmark.add_argument("--json", action="store_true") + scout_benchmark.set_defaults(handler=cmd_allocator_scout_benchmark) + token = allocator_sub.add_parser("token", help="Provision the node control capability") token_sub = token.add_subparsers(dest="allocator_token_command", required=True) token_write = token_sub.add_parser( diff --git a/docs/allocator.md b/docs/allocator.md index 5c15daa4..21b37cb4 100644 --- a/docs/allocator.md +++ b/docs/allocator.md @@ -634,8 +634,8 @@ failed destructive actions remain backoff-protected. `load` never infers a mutable download source from a display name. It verifies an existing GGUF, or the authenticated profile may provide all three autonomous-transfer fields: an exact -`hf://owner/repo/path.gguf` source, an immutable SHA-256, and a maximum artifact size. The llama.cpp -adapter downloads under an artifact-addressed staging name, resumes bounded partial transfers, +`hf://owner/repo[@commit]/path.gguf` source, an immutable SHA-256, and a maximum artifact size. The +llama.cpp adapter downloads under an artifact-addressed staging name, resumes bounded partial transfers, rejects streams above the size ceiling, hashes the complete file, and atomically publishes it only after verification. A wrong digest never replaces the prior cache. When a profile declares `artifact_sha256`, both `load` and `warm` hash the exact cached file before process launch. A @@ -867,7 +867,7 @@ grid --local allocator model set \ --grid allocator-control \ --memory-mb \ --artifact-sha256 <64-hex-digest> \ - --artifact-source hf://owner/repo/path/to/model.gguf \ + --artifact-source hf://owner/repo@/path/to/model.gguf \ --artifact-size-mb \ --runtime llama.cpp \ --min-replicas 0 \ @@ -935,7 +935,7 @@ budget for one replica, not the file's compressed size: grid allocator model set \ --memory-mb 12000 \ --artifact-sha256 <64-hex-digest> \ - --artifact-source hf://owner/repo/path/to/model.gguf \ + --artifact-source hf://owner/repo@/path/to/model.gguf \ --artifact-size-mb 9000 \ --workload-score coding=1 \ --workload-score research=.8 \ @@ -979,6 +979,52 @@ latency, or error pressure still requests at least one replica beyond the curren If `--runtime` is omitted, it defaults to `llama.cpp`. Once the flag is present, only the listed runtimes are eligible. +### Open-weight model scout + +The allocator can continuously look for newer open-weight releases without treating a trending +name as proof that it is better. The scout queries public Hugging Face metadata, permits only +configured publishers and licenses, resolves a full repository commit, obtains the LFS SHA-256 for +an exact non-sharded GGUF (or a deterministic immutable vLLM snapshot identity), and rejects an +artifact whose size is unknown. It downloads no model bytes during discovery. + +Run one discovery cycle against the controller's current node inventory: + +```bash +grid --local allocator scout run --grid allocator-control --search coder +grid --local allocator scout status --grid allocator-control +``` + +Each proposal explains whether any accepting node has the runtime, free memory, and free disk to +host it, and identifies configured models serving an overlapping workload for later comparison. +Popularity and release recency affect only which candidates are inspected first. They never count +as quality evidence. + +Qualification is deliberately a separate real-inference step: + +```bash +grid --local allocator scout benchmark \ + --grid allocator-control \ + --inference-grid \ + --deploy-canary +``` + +The command creates an immutable `min_replicas=0`, `max_replicas=1` profile, waits for the allocator +to place and warm it, sends the bounded coding/research/general cases through the normal Grid route, +and records the outputs' quality, error, latency, and size through the authenticated evaluation API. +It never drains or deletes an incumbent. The scout state is written owner-only below +`~/.grid/allocator-scout/`, so discovery can be reviewed or resumed without trusting terminal text. + +For periodic discovery, run: + +```bash +grid --local allocator scout watch --grid allocator-control --interval 21600 +``` + +`watch` refreshes proposals every six hours and remains discovery-only. Canary deployment is +explicit because it consumes real fleet capacity. A production replacement should require the +candidate's recorded benchmark evidence plus a warm-before-drain allocator plan; mutable `main` +revisions are never emitted by the scout. + Use the three modes as a rollout sequence: ```bash diff --git a/shared/allocator/runtime.py b/shared/allocator/runtime.py index 59218c94..efc31040 100644 --- a/shared/allocator/runtime.py +++ b/shared/allocator/runtime.py @@ -499,7 +499,7 @@ def fetch_artifact( raise RuntimeError( "managed llama.cpp artifact ids must be plain .gguf filenames" ) - repo, filename = _parse_hugging_face_artifact_source(source) + repo, revision, filename = _parse_hugging_face_artifact_source(source) from shared.models import download target = download.local_path(model_id) @@ -532,6 +532,7 @@ def fetch_artifact( filename, out=staging, max_bytes=maximum_bytes, + revision=revision, ) except SystemExit as exc: raise RuntimeError(str(exc)) from None @@ -2937,7 +2938,7 @@ def _cached_model_path(model_id: str) -> Path | None: ) -def _parse_hugging_face_artifact_source(source: str) -> tuple[str, str]: +def _parse_hugging_face_artifact_source(source: str) -> tuple[str, str, str]: parsed = urlsplit(str(source or "")) if ( parsed.scheme != "hf" @@ -2948,19 +2949,32 @@ def _parse_hugging_face_artifact_source(source: str) -> tuple[str, str]: or parsed.fragment ): raise RuntimeError( - "managed llama.cpp artifact_source must be hf://owner/repo/path.gguf" + "managed llama.cpp artifact_source must be " + "hf://owner/repo[@commit]/path.gguf" ) parts = [unquote(item) for item in parsed.path.split("/") if item] if len(parts) < 2: raise RuntimeError( - "managed llama.cpp artifact_source must be hf://owner/repo/path.gguf" + "managed llama.cpp artifact_source must be " + "hf://owner/repo[@commit]/path.gguf" ) if any(item in (".", "..") or "\0" in item for item in parts): raise RuntimeError("artifact_source contains an unsafe path component") + repo_name, separator, revision = parts[0].partition("@") + if separator and ( + len(revision) != 40 + or any(character not in "0123456789abcdefABCDEF" for character in revision) + ): + raise RuntimeError("Hugging Face artifact revision must be a full 40-hex commit") + if not repo_name: + raise RuntimeError( + "managed llama.cpp artifact_source must be " + "hf://owner/repo[@commit]/path.gguf" + ) filename = "/".join(parts[1:]) if not filename.lower().endswith(".gguf"): raise RuntimeError("managed llama.cpp artifact_source must name an exact .gguf") - return f"{parsed.netloc}/{parts[0]}", filename + return f"{parsed.netloc}/{repo_name}", (revision.lower() if separator else "main"), filename def _sha256_file(path: Path) -> str: diff --git a/shared/allocator/scout.py b/shared/allocator/scout.py new file mode 100644 index 00000000..99fa408c --- /dev/null +++ b/shared/allocator/scout.py @@ -0,0 +1,672 @@ +"""Immutable model discovery, fleet-fit analysis, and benchmark proposal state. + +The scout deliberately stops short of silently replacing a production model. Hub metadata is +useful for finding candidates, not evidence of answer quality. A candidate becomes proposal-ready +only after its repository revision and artifact digest are immutable, its license passes policy, +the current fleet can host it, and a real canary benchmark records fresh evaluation evidence. +""" + +from __future__ import annotations + +import hashlib +import math +import re +import statistics +import time +from dataclasses import asdict, dataclass +from datetime import datetime, timezone +from pathlib import Path +from typing import Any, Callable, Iterable, Mapping + +import httpx + +from shared import jsonio + + +SCOUT_SCHEMA_VERSION = 1 +DEFAULT_HUB_URL = "https://huggingface.co" +DEFAULT_ALLOWED_LICENSES = ( + "apache-2.0", + "bsd-2-clause", + "bsd-3-clause", + "cc-by-4.0", + "mit", +) +DEFAULT_TRUSTED_AUTHORS = ( + "deepseek-ai", + "google", + "ggml-org", + "meta-llama", + "microsoft", + "mistralai", + "nvidia", + "qwen", +) +DEFAULT_QUANTIZATIONS = ("Q4_K_M", "Q5_K_M", "Q4_K_S", "Q5_K_S", "Q6_K") +MAX_DISCOVERY_RESULTS = 100 +MAX_REPOSITORIES_INSPECTED = 40 +MAX_ARTIFACT_SIZE_MB = 500_000 +_FULL_SHA = re.compile(r"^[0-9a-fA-F]{40}$") +_FULL_SHA256 = re.compile(r"^[0-9a-fA-F]{64}$") +_SPLIT_GGUF = re.compile(r"-\d{5}-of-\d{5}\.gguf$", re.IGNORECASE) +_PARAMETERS = re.compile(r"(? None: + for name in ("workloads", "runtimes", "trusted_authors", "allowed_licenses"): + values = tuple(sorted({str(item).strip().lower() for item in getattr(self, name) if str(item).strip()})) + if not values: + raise ValueError(f"{name} must not be empty") + object.__setattr__(self, name, values) + quantizations = tuple(dict.fromkeys(str(item).strip().upper() for item in self.quantizations if str(item).strip())) + if not quantizations: + raise ValueError("quantizations must not be empty") + object.__setattr__(self, "quantizations", quantizations) + if not 1 <= self.max_results <= MAX_DISCOVERY_RESULTS: + raise ValueError(f"max_results must be in [1, {MAX_DISCOVERY_RESULTS}]") + if not 1 <= self.max_repositories <= MAX_REPOSITORIES_INSPECTED: + raise ValueError(f"max_repositories must be in [1, {MAX_REPOSITORIES_INSPECTED}]") + if not 1 <= self.max_artifact_size_mb <= MAX_ARTIFACT_SIZE_MB: + raise ValueError(f"max_artifact_size_mb must be in [1, {MAX_ARTIFACT_SIZE_MB}]") + if self.min_downloads < 0: + raise ValueError("min_downloads must be non-negative") + + +@dataclass(frozen=True, slots=True) +class ModelCandidate: + candidate_id: str + repo_id: str + revision: str + runtime: str + model_id: str + artifact_path: str + artifact_source: str + artifact_sha256: str + artifact_size_mb: int + estimated_memory_mb: int + quantization: str = "" + parameter_billions: float = 0.0 + license: str = "" + downloads: int = 0 + likes: int = 0 + last_modified: str = "" + pipeline_tag: str = "" + workload_scores: tuple[tuple[str, float], ...] = () + + def to_dict(self) -> dict[str, Any]: + return asdict(self) + + @classmethod + def from_dict(cls, value: Mapping[str, Any]) -> "ModelCandidate": + fields = dict(value) + fields["workload_scores"] = tuple(tuple(item) for item in fields.get("workload_scores") or ()) + return cls(**fields) + + +@dataclass(frozen=True, slots=True) +class FleetFit: + node_id: str + fits: bool + runtime: str + headroom_mb: int + disk_headroom_mb: int | None + score: float + reason: str + + +@dataclass(frozen=True, slots=True) +class ScoutProposal: + proposal_id: str + candidate: ModelCandidate + state: str + score: float + fits: tuple[FleetFit, ...] + compare_models: tuple[str, ...] = () + reasons: tuple[str, ...] = () + benchmark_quality: float | None = None + benchmark_latency_ms: float | None = None + benchmark_samples: int = 0 + + def to_dict(self) -> dict[str, Any]: + return { + **asdict(self), + "candidate": self.candidate.to_dict(), + "fits": [asdict(item) for item in self.fits], + } + + @classmethod + def from_dict(cls, value: Mapping[str, Any]) -> "ScoutProposal": + fields = dict(value) + fields["candidate"] = ModelCandidate.from_dict(fields["candidate"]) + fields["fits"] = tuple(FleetFit(**item) for item in fields.get("fits") or ()) + fields["compare_models"] = tuple(fields.get("compare_models") or ()) + fields["reasons"] = tuple(fields.get("reasons") or ()) + return cls(**fields) + + +@dataclass(frozen=True, slots=True) +class BenchmarkCase: + workload: str + prompt: str + required_fragments: tuple[str, ...] = () + + +@dataclass(frozen=True, slots=True) +class BenchmarkSample: + workload: str + quality: float + latency_ms: float + output_units: int + error: bool + + +BENCHMARK_CASES: tuple[BenchmarkCase, ...] = ( + BenchmarkCase( + "coding", + "Write a Python function named clamp(value, low, high). Return only the function.", + ("def clamp", "return"), + ), + BenchmarkCase( + "research", + "State the difference between correlation and causation in two concise sentences.", + ("correlation", "caus"), + ), + BenchmarkCase( + "general", + "Reply with exactly the four-character string GRID.", + ("grid",), + ), +) + + +class HuggingFaceDiscovery: + """Read public Hub metadata through its documented API; no artifact bytes are downloaded.""" + + def __init__( + self, + *, + base_url: str = DEFAULT_HUB_URL, + client: httpx.Client | None = None, + timeout: float = 20.0, + ) -> None: + self.base_url = base_url.rstrip("/") + self._owns_client = client is None + self.client = client or httpx.Client(timeout=timeout, follow_redirects=True, trust_env=True) + self.issues: list[str] = [] + + def close(self) -> None: + if self._owns_client: + self.client.close() + + def discover(self, policy: ScoutPolicy, *, search: str = "") -> tuple[ModelCandidate, ...]: + self.issues = [] + queries: list[dict[str, str]] = [] + if "llama.cpp" in policy.runtimes: + queries.append({"filter": "gguf"}) + if "vllm" in policy.runtimes: + queries.append({"apps": "vllm"}) + rows_by_id: dict[str, Mapping[str, Any]] = {} + # Query publishers independently. A single global "newest GGUF" page is dominated by + # community conversions and could starve every trusted publisher out of a bounded scan. + for query in queries: + for author in policy.trusted_authors: + try: + response = self.client.get( + f"{self.base_url}/api/models", + params={ + **query, + # Hub owner matching is case-sensitive for some organizations even though + # policy comparison is intentionally canonical and case-insensitive. + "author": _HUB_AUTHOR_CASE.get(author, author), + "sort": "lastModified", + "direction": "-1", + "limit": str(min(policy.max_results, policy.max_repositories)), + **({"search": search} if search else {}), + }, + ) + response.raise_for_status() + rows = response.json() + except (httpx.HTTPError, ValueError) as exc: + self.issues.append(_discovery_issue("listing", author, exc)) + continue + if not isinstance(rows, list): + self.issues.append(f"listing {author}: non-list response") + continue + for row in rows: + if isinstance(row, Mapping): + repo_id = str(row.get("id") or row.get("modelId") or "") + if repo_id: + rows_by_id.setdefault(repo_id, row) + candidates: list[ModelCandidate] = [] + inspected = 0 + for summary in rows_by_id.values(): + if inspected >= policy.max_repositories: + break + if not isinstance(summary, Mapping): + continue + repo_id = str(summary.get("id") or summary.get("modelId") or "") + author = repo_id.partition("/")[0].lower() + if not repo_id or author not in policy.trusted_authors: + continue + if bool(summary.get("gated")) or int(summary.get("downloads") or 0) < policy.min_downloads: + continue + inspected += 1 + try: + detail_response = self.client.get( + f"{self.base_url}/api/models/{repo_id}", params={"blobs": "true"} + ) + detail_response.raise_for_status() + detail = detail_response.json() + except (httpx.HTTPError, ValueError) as exc: + self.issues.append(_discovery_issue("detail", repo_id, exc)) + continue + if not isinstance(detail, Mapping): + self.issues.append(f"detail {repo_id}: non-object response") + continue + candidates.extend(_candidates_from_hub_detail(detail, policy)) + return tuple(sorted(candidates, key=_candidate_discovery_sort_key)) + + +def _discovery_issue(kind: str, source: str, exc: Exception) -> str: + if isinstance(exc, httpx.HTTPStatusError): + return f"{kind} {source}: HTTP {exc.response.status_code}" + return f"{kind} {source}: {type(exc).__name__}" + + +def _candidates_from_hub_detail( + detail: Mapping[str, Any], policy: ScoutPolicy +) -> tuple[ModelCandidate, ...]: + repo_id = str(detail.get("id") or detail.get("modelId") or "") + revision = str(detail.get("sha") or "").lower() + if not repo_id or not _FULL_SHA.fullmatch(revision) or bool(detail.get("gated")): + return () + license_name = _license_of(detail) + if license_name.lower() not in policy.allowed_licenses: + return () + siblings = detail.get("siblings") + if not isinstance(siblings, list): + return () + out: list[ModelCandidate] = [] + if "llama.cpp" in policy.runtimes: + for item in siblings: + candidate = _gguf_candidate(detail, item, policy, repo_id, revision, license_name) + if candidate is not None: + out.append(candidate) + if "vllm" in policy.runtimes: + candidate = _vllm_candidate(detail, siblings, policy, repo_id, revision, license_name) + if candidate is not None: + out.append(candidate) + return tuple(out) + + +def _gguf_candidate( + detail: Mapping[str, Any], + item: object, + policy: ScoutPolicy, + repo_id: str, + revision: str, + license_name: str, +) -> ModelCandidate | None: + if not isinstance(item, Mapping): + return None + path = str(item.get("rfilename") or item.get("path") or "") + if not path.lower().endswith(".gguf") or _SPLIT_GGUF.search(path): + return None + quant = next((value for value in policy.quantizations if value in path.upper()), "") + if not quant: + return None + size_bytes, digest = _artifact_identity(item) + size_mb = math.ceil(size_bytes / (1024 * 1024)) if size_bytes > 0 else 0 + if not size_mb or size_mb > policy.max_artifact_size_mb or not _FULL_SHA256.fullmatch(digest): + return None + model_id = Path(path).name + scores = _workload_scores(detail, repo_id, policy.workloads) + identity = _candidate_id(repo_id, revision, path, digest) + return ModelCandidate( + candidate_id=identity, + repo_id=repo_id, + revision=revision, + runtime="llama.cpp", + model_id=model_id, + artifact_path=path, + artifact_source=f"hf://{repo_id}@{revision}/{path}", + artifact_sha256=digest.lower(), + artifact_size_mb=size_mb, + estimated_memory_mb=min(MAX_ARTIFACT_SIZE_MB, math.ceil(size_mb * 1.20) + 768), + quantization=quant, + parameter_billions=_parameter_billions(detail, repo_id), + license=license_name, + downloads=int(detail.get("downloads") or 0), + likes=int(detail.get("likes") or 0), + last_modified=str(detail.get("lastModified") or detail.get("last_modified") or ""), + pipeline_tag=str(detail.get("pipeline_tag") or ""), + workload_scores=scores, + ) + + +def _vllm_candidate( + detail: Mapping[str, Any], + siblings: list[Any], + policy: ScoutPolicy, + repo_id: str, + revision: str, + license_name: str, +) -> ModelCandidate | None: + tags = {str(item).lower() for item in detail.get("tags") or ()} + has_config = any( + isinstance(item, Mapping) + and str(item.get("rfilename") or item.get("path") or "") == "config.json" + for item in siblings + ) + safetensors = [ + item + for item in siblings + if isinstance(item, Mapping) + and str(item.get("rfilename") or item.get("path") or "").endswith(".safetensors") + ] + if not has_config or not safetensors or not ({"transformers", "vllm"} & tags): + return None + total_bytes = sum(_artifact_identity(item)[0] for item in safetensors) + size_mb = math.ceil(total_bytes / (1024 * 1024)) if total_bytes > 0 else 0 + if not size_mb or size_mb > policy.max_artifact_size_mb: + return None + digest = hashlib.sha256(f"hf://{repo_id}@{revision}".encode()).hexdigest() + return ModelCandidate( + candidate_id=_candidate_id(repo_id, revision, "snapshot", digest), + repo_id=repo_id, + revision=revision, + runtime="vllm", + model_id=repo_id, + artifact_path="", + artifact_source=f"hf://{repo_id}@{revision}", + artifact_sha256=digest, + artifact_size_mb=size_mb, + estimated_memory_mb=min(MAX_ARTIFACT_SIZE_MB, math.ceil(size_mb * 1.12) + 2048), + parameter_billions=_parameter_billions(detail, repo_id), + license=license_name, + downloads=int(detail.get("downloads") or 0), + likes=int(detail.get("likes") or 0), + last_modified=str(detail.get("lastModified") or detail.get("last_modified") or ""), + pipeline_tag=str(detail.get("pipeline_tag") or ""), + workload_scores=_workload_scores(detail, repo_id, policy.workloads), + ) + + +def analyze_fleet_fit( + candidate: ModelCandidate, nodes: Iterable[Mapping[str, Any]] +) -> tuple[FleetFit, ...]: + fits: list[FleetFit] = [] + for node in nodes: + node_id = str(node.get("node_id") or "") + runtimes = {str(item).lower() for item in node.get("runtimes") or ()} + capacity = _safe_nonnegative_int(node.get("capacity_mb")) + reserved = _safe_nonnegative_int(node.get("reserved_mb")) + headroom = max(0, capacity - reserved) + raw_disk = node.get("disk_available_mb") + disk = _safe_nonnegative_int(raw_disk) if raw_disk is not None else None + reason = "" + if str(node.get("state") or "accepting") != "accepting": + reason = "node is not accepting placements" + elif candidate.runtime.lower() not in runtimes: + reason = f"runtime {candidate.runtime} is unavailable" + elif headroom < candidate.estimated_memory_mb: + reason = f"needs {candidate.estimated_memory_mb} MB memory; {headroom} MB available" + elif disk is None: + reason = "disk availability is unknown" + elif disk < candidate.artifact_size_mb: + reason = f"needs {candidate.artifact_size_mb} MB disk; {disk} MB available" + score = 0.0 + if not reason: + memory_margin = min(1.0, (headroom - candidate.estimated_memory_mb) / max(1, candidate.estimated_memory_mb)) + disk_margin = min(1.0, (disk - candidate.artifact_size_mb) / max(1, candidate.artifact_size_mb)) + bandwidth = max(0.0, float(node.get("memory_bandwidth_gbps") or 0.0)) + score = round(0.50 + 0.20 * memory_margin + 0.10 * disk_margin + 0.20 * min(1.0, bandwidth / 2_000.0), 6) + fits.append(FleetFit(node_id, not reason, candidate.runtime, headroom, disk, score, reason)) + return tuple(sorted(fits, key=lambda item: (-int(item.fits), -item.score, item.node_id))) + + +def build_proposals( + candidates: Iterable[ModelCandidate], status: Mapping[str, Any] +) -> tuple[ScoutProposal, ...]: + nodes = [item for item in status.get("nodes") or () if isinstance(item, Mapping)] + profiles = [item for item in status.get("models") or () if isinstance(item, Mapping)] + proposals: list[ScoutProposal] = [] + for candidate in candidates: + fits = analyze_fleet_fit(candidate, nodes) + compatible = tuple(item for item in fits if item.fits) + compare = _comparison_models(candidate, profiles) + reasons: list[str] = [] + if not compatible: + reasons.append("no current allocator host can safely fit the immutable artifact") + if not candidate.workload_scores: + reasons.append("no configured workload suitability could be inferred") + state = "benchmark-ready" if not reasons else "blocked" + score = _proposal_score(candidate, compatible) + proposals.append( + ScoutProposal( + proposal_id=f"proposal-{candidate.candidate_id}", + candidate=candidate, + state=state, + score=score, + fits=fits, + compare_models=compare, + reasons=tuple(reasons), + ) + ) + return tuple(sorted(proposals, key=lambda item: (-int(item.state == "benchmark-ready"), -item.score, item.proposal_id))) + + +def benchmark_candidate( + proposal: ScoutProposal, + runner: Callable[[str, str], tuple[str, float]], + *, + workloads: Iterable[str] = (), +) -> tuple[ScoutProposal, tuple[BenchmarkSample, ...]]: + selected = {str(item).lower() for item in workloads if str(item)} + cases = [case for case in BENCHMARK_CASES if not selected or case.workload in selected] + if not cases: + raise ValueError("no benchmark cases match the requested workloads") + samples: list[BenchmarkSample] = [] + for case in cases: + started = time.monotonic() + try: + text, reported_latency_ms = runner(proposal.candidate.model_id, case.prompt) + elapsed_ms = max(0.0, (time.monotonic() - started) * 1_000.0) + latency_ms = max(elapsed_ms, float(reported_latency_ms or 0.0)) + lowered = str(text).lower() + hits = sum(fragment.lower() in lowered for fragment in case.required_fragments) + quality = hits / max(1, len(case.required_fragments)) + samples.append(BenchmarkSample(case.workload, quality, latency_ms, len(str(text)), False)) + except Exception: + samples.append(BenchmarkSample(case.workload, 0.0, 0.0, 0, True)) + quality = statistics.fmean(item.quality for item in samples) + successful_latency = [item.latency_ms for item in samples if not item.error] + latency = statistics.median(successful_latency) if successful_latency else None + qualified = quality >= 0.80 and not any(item.error for item in samples) + updated = ScoutProposal( + proposal_id=proposal.proposal_id, + candidate=proposal.candidate, + state="qualified" if qualified else "benchmark-failed", + score=proposal.score, + fits=proposal.fits, + compare_models=proposal.compare_models, + reasons=proposal.reasons + (() if qualified else ("real canary benchmark did not meet the quality floor",)), + benchmark_quality=round(quality, 6), + benchmark_latency_ms=round(latency, 3) if latency is not None else None, + benchmark_samples=len(samples), + ) + return updated, tuple(samples) + + +def load_scout_state(path: Path) -> dict[str, Any]: + state = jsonio.load_json(path) + if not state: + return {"schema_version": SCOUT_SCHEMA_VERSION, "updated_at": 0.0, "proposals": []} + if int(state.get("schema_version", 0)) != SCOUT_SCHEMA_VERSION: + raise ValueError("unsupported allocator scout state schema") + proposals = state.get("proposals") + if not isinstance(proposals, list): + raise ValueError("allocator scout proposals must be a list") + return state + + +def save_scout_state( + path: Path, + proposals: Iterable[ScoutProposal], + *, + policy: ScoutPolicy | None = None, + benchmark_samples: Mapping[str, Iterable[BenchmarkSample]] | None = None, +) -> None: + path.parent.mkdir(parents=True, exist_ok=True) + payload = { + "schema_version": SCOUT_SCHEMA_VERSION, + "updated_at": time.time(), + "policy": asdict(policy) if policy is not None else {}, + "proposals": [item.to_dict() for item in proposals], + "benchmark_samples": { + key: [ + dict(sample) if isinstance(sample, Mapping) else asdict(sample) + for sample in values + ] + for key, values in (benchmark_samples or {}).items() + }, + } + jsonio.atomic_write_json(path, payload, mode=0o600) + + +def proposals_from_state(state: Mapping[str, Any]) -> tuple[ScoutProposal, ...]: + return tuple(ScoutProposal.from_dict(item) for item in state.get("proposals") or ()) + + +def _artifact_identity(item: Mapping[str, Any]) -> tuple[int, str]: + lfs = item.get("lfs") + if isinstance(lfs, Mapping): + size = _safe_nonnegative_int(lfs.get("size") or item.get("size")) + digest = str(lfs.get("sha256") or "") + else: + size = _safe_nonnegative_int(item.get("size")) + digest = str(item.get("blobId") or item.get("blob_id") or "") + return size, digest + + +def _license_of(detail: Mapping[str, Any]) -> str: + card = detail.get("cardData") or detail.get("card_data") + if isinstance(card, Mapping): + value = card.get("license") + if isinstance(value, list): + return str(value[0] if value else "") + if value: + return str(value) + for tag in detail.get("tags") or (): + if str(tag).startswith("license:"): + return str(tag).partition(":")[2] + return "" + + +def _parameter_billions(detail: Mapping[str, Any], repo_id: str) -> float: + safetensors = detail.get("safetensors") + if isinstance(safetensors, Mapping): + total = safetensors.get("total") + try: + if total is not None: + return round(float(total) / 1_000_000_000.0, 3) + except (TypeError, ValueError, OverflowError): + pass + matches = _PARAMETERS.findall(repo_id.replace("-", " ")) + return float(matches[-1]) if matches else 0.0 + + +def _workload_scores( + detail: Mapping[str, Any], repo_id: str, workloads: Iterable[str] +) -> tuple[tuple[str, float], ...]: + corpus = " ".join( + [repo_id, str(detail.get("pipeline_tag") or ""), *(str(item) for item in detail.get("tags") or ())] + ).lower() + scores: dict[str, float] = {} + for workload in workloads: + if workload == "coding" and any(word in corpus for word in ("code", "coder", "program")): + scores[workload] = 1.0 + elif workload == "research" and any(word in corpus for word in ("reason", "research", "math", "science")): + scores[workload] = 0.9 + elif workload in ("general", "research", "coding") and any( + word in corpus for word in ("text-generation", "causal-lm", "instruct", "chat") + ): + scores[workload] = 0.7 if workload != "general" else 0.9 + return tuple(sorted(scores.items())) + + +def _comparison_models( + candidate: ModelCandidate, profiles: Iterable[Mapping[str, Any]] +) -> tuple[str, ...]: + candidate_workloads = {name for name, score in candidate.workload_scores if score > 0} + matches: list[tuple[float, str]] = [] + for profile in profiles: + model_id = str(profile.get("model_id") or "") + if not model_id or model_id == candidate.model_id: + continue + scores = { + str(item[0]): float(item[1]) + for item in profile.get("workload_scores") or () + if isinstance(item, (list, tuple)) and len(item) == 2 + } + overlap = sum(scores.get(workload, 0.0) for workload in candidate_workloads) + if overlap: + matches.append((overlap, model_id)) + return tuple(model_id for _score, model_id in sorted(matches, key=lambda item: (-item[0], item[1]))[:3]) + + +def _proposal_score(candidate: ModelCandidate, fits: tuple[FleetFit, ...]) -> float: + popularity = min(1.0, math.log10(candidate.downloads + 1) / 7.0) + approval = min(1.0, math.log10(candidate.likes + 1) / 4.0) + recency = _recency_score(candidate.last_modified) + suitability = max((score for _name, score in candidate.workload_scores), default=0.0) + fleet = max((item.score for item in fits), default=0.0) + # This is only a discovery priority. Quality remains absent until benchmark_candidate records it. + return round(0.15 * popularity + 0.10 * approval + 0.20 * recency + 0.25 * suitability + 0.30 * fleet, 6) + + +def _recency_score(value: str) -> float: + try: + parsed = datetime.fromisoformat(value.replace("Z", "+00:00")) + if parsed.tzinfo is None: + parsed = parsed.replace(tzinfo=timezone.utc) + age_days = max(0.0, (datetime.now(timezone.utc) - parsed).total_seconds() / 86_400.0) + return math.exp(-age_days / 90.0) + except (TypeError, ValueError, OverflowError): + return 0.0 + + +def _candidate_discovery_sort_key(candidate: ModelCandidate) -> tuple[Any, ...]: + quant_rank = ( + DEFAULT_QUANTIZATIONS.index(candidate.quantization) + if candidate.quantization in DEFAULT_QUANTIZATIONS + else len(DEFAULT_QUANTIZATIONS) + ) + return (-_recency_score(candidate.last_modified), quant_rank, -candidate.downloads, candidate.candidate_id) + + +def _candidate_id(repo_id: str, revision: str, path: str, digest: str) -> str: + value = f"{repo_id}\0{revision}\0{path}\0{digest}".encode() + return hashlib.sha256(value).hexdigest()[:20] + + +def _safe_nonnegative_int(value: Any) -> int: + if isinstance(value, bool): + return 0 + try: + return max(0, int(value or 0)) + except (TypeError, ValueError, OverflowError): + return 0 diff --git a/shared/models/download.py b/shared/models/download.py index 58cb0085..7ab78f4f 100644 --- a/shared/models/download.py +++ b/shared/models/download.py @@ -4,6 +4,7 @@ import sys from pathlib import Path +from urllib.parse import quote import httpx @@ -16,8 +17,10 @@ DOWNLOAD_READ_TIMEOUT_SECONDS = 30.0 -def hf_url(repo: str, quantized_file: str) -> str: - return f"{HF_BASE}/{repo}/resolve/main/{quantized_file}" +def hf_url(repo: str, quantized_file: str, revision: str = "main") -> str: + """An exact Hub resolve URL; callers may pin a full commit instead of following ``main``.""" + + return f"{HF_BASE}/{repo}/resolve/{quote(revision, safe='')}/{quantized_file}" DEFAULT_QUANT = "Q4_K_M" @@ -144,6 +147,7 @@ def download( out: Path | None = None, on_progress=None, max_bytes: int | None = None, + revision: str = "main", ) -> Path: if max_bytes is not None and ( isinstance(max_bytes, bool) or not isinstance(max_bytes, int) or max_bytes <= 0 @@ -162,7 +166,7 @@ def download( return target part = target.with_suffix(target.suffix + ".part") target.parent.mkdir(parents=True, exist_ok=True) - url = hf_url(repo, quantized_file) + url = hf_url(repo, quantized_file, revision) have = part.stat().st_size if part.exists() else 0 if max_bytes is not None and have > max_bytes: diff --git a/tests/test_allocator_scout.py b/tests/test_allocator_scout.py new file mode 100644 index 00000000..a955181d --- /dev/null +++ b/tests/test_allocator_scout.py @@ -0,0 +1,317 @@ +from __future__ import annotations + +import json +import os +from dataclasses import replace +from datetime import datetime, timezone + +import httpx +import pytest + +from cli import parser +from shared.allocator.scout import ( + HuggingFaceDiscovery, + ModelCandidate, + ScoutPolicy, + analyze_fleet_fit, + benchmark_candidate, + build_proposals, + load_scout_state, + proposals_from_state, + save_scout_state, +) +from shared.allocator.runtime import _parse_hugging_face_artifact_source +from shared.models.download import hf_url + + +REVISION = "a" * 40 +DIGEST = "b" * 64 + + +def _detail(*, repo: str = "Qwen/Qwen-Coder-GGUF", license_name: str = "apache-2.0") -> dict: + return { + "id": repo, + "sha": REVISION, + "downloads": 10_000, + "likes": 500, + "lastModified": datetime.now(timezone.utc).isoformat(), + "pipeline_tag": "text-generation", + "tags": ["gguf", "transformers", "code", f"license:{license_name}"], + "cardData": {"license": license_name}, + "siblings": [ + {"rfilename": "config.json", "size": 100}, + { + "rfilename": "Qwen-Coder-Q4_K_M.gguf", + "size": 4_000_000_000, + "lfs": {"size": 4_000_000_000, "sha256": DIGEST}, + }, + { + "rfilename": "model-00001-of-00002-Q4_K_M.gguf", + "lfs": {"size": 2_000_000_000, "sha256": "c" * 64}, + }, + { + "rfilename": "model.safetensors", + "size": 5_000_000_000, + "lfs": {"size": 5_000_000_000, "sha256": "d" * 64}, + }, + ], + } + + +def _candidate(**changes) -> ModelCandidate: + value = ModelCandidate( + candidate_id="candidate-a", + repo_id="Qwen/Qwen-Coder-GGUF", + revision=REVISION, + runtime="llama.cpp", + model_id="Qwen-Coder-Q4_K_M.gguf", + artifact_path="Qwen-Coder-Q4_K_M.gguf", + artifact_source=f"hf://Qwen/Qwen-Coder-GGUF@{REVISION}/Qwen-Coder-Q4_K_M.gguf", + artifact_sha256=DIGEST, + artifact_size_mb=4_000, + estimated_memory_mb=5_568, + quantization="Q4_K_M", + license="apache-2.0", + downloads=10_000, + likes=500, + last_modified=datetime.now(timezone.utc).isoformat(), + pipeline_tag="text-generation", + workload_scores=(("coding", 1.0), ("general", 0.7)), + ) + return replace(value, **changes) + + +def _node(**changes) -> dict: + value = { + "node_id": "gpu-a", + "state": "accepting", + "capacity_mb": 24_000, + "reserved_mb": 4_000, + "disk_available_mb": 80_000, + "runtimes": ["llama.cpp", "vllm"], + "memory_bandwidth_gbps": 1_000, + } + value.update(changes) + return value + + +def test_policy_normalizes_and_bounds_external_input(): + policy = ScoutPolicy( + workloads=("Coding", "coding"), + runtimes=("VLLM",), + trusted_authors=("Qwen",), + allowed_licenses=("MIT",), + ) + assert policy.workloads == ("coding",) + assert policy.runtimes == ("vllm",) + assert policy.trusted_authors == ("qwen",) + with pytest.raises(ValueError, match="max_results"): + ScoutPolicy(max_results=101) + + +def test_discovery_resolves_exact_gguf_and_vllm_snapshot_without_downloading_bytes(): + requests: list[httpx.Request] = [] + + def handler(request: httpx.Request) -> httpx.Response: + requests.append(request) + if request.url.path == "/api/models": + return httpx.Response(200, json=[{"id": "Qwen/Qwen-Coder-GGUF", "downloads": 10_000}]) + assert request.url.path == "/api/models/Qwen/Qwen-Coder-GGUF" + return httpx.Response(200, json=_detail()) + + client = httpx.Client(transport=httpx.MockTransport(handler)) + discovery = HuggingFaceDiscovery(base_url="https://hub.test", client=client) + candidates = discovery.discover( + ScoutPolicy(trusted_authors=("qwen",), runtimes=("llama.cpp", "vllm")) + ) + + assert {item.runtime for item in candidates} == {"llama.cpp", "vllm"} + gguf = next(item for item in candidates if item.runtime == "llama.cpp") + assert gguf.artifact_sha256 == DIGEST + assert gguf.artifact_source == ( + f"hf://Qwen/Qwen-Coder-GGUF@{REVISION}/Qwen-Coder-Q4_K_M.gguf" + ) + vllm = next(item for item in candidates if item.runtime == "vllm") + assert vllm.artifact_source == f"hf://Qwen/Qwen-Coder-GGUF@{REVISION}" + assert len(vllm.artifact_sha256) == 64 + # One GGUF listing + one vLLM listing, but the duplicate repo is inspected only once. + assert sum(request.url.path == "/api/models" for request in requests) == 2 + assert sum(request.url.path.endswith("Qwen/Qwen-Coder-GGUF") for request in requests) == 1 + + +def test_discovery_rejects_untrusted_gated_mutable_and_unknown_license_repositories(): + details = _detail(license_name="other") + details["sha"] = "main" + + def handler(request: httpx.Request) -> httpx.Response: + if request.url.path == "/api/models": + return httpx.Response( + 200, + json=[ + {"id": "random/model", "downloads": 99_000}, + {"id": "Qwen/gated", "gated": True, "downloads": 99_000}, + {"id": "Qwen/mutable", "downloads": 99_000}, + ], + ) + return httpx.Response(200, json=details) + + discovery = HuggingFaceDiscovery( + base_url="https://hub.test", + client=httpx.Client(transport=httpx.MockTransport(handler)), + ) + assert discovery.discover(ScoutPolicy(trusted_authors=("qwen",), runtimes=("llama.cpp",))) == () + + +def test_discovery_keeps_good_candidates_when_one_repository_is_rate_limited(): + def handler(request: httpx.Request) -> httpx.Response: + if request.url.path == "/api/models": + return httpx.Response( + 200, + json=[ + {"id": "Qwen/rate-limited", "downloads": 10_000}, + {"id": "Qwen/Qwen-Coder-GGUF", "downloads": 10_000}, + ], + ) + if request.url.path.endswith("/rate-limited"): + return httpx.Response(429) + return httpx.Response(200, json=_detail()) + + discovery = HuggingFaceDiscovery( + base_url="https://hub.test", + client=httpx.Client(transport=httpx.MockTransport(handler)), + ) + candidates = discovery.discover( + ScoutPolicy(trusted_authors=("qwen",), runtimes=("llama.cpp",)) + ) + + assert candidates + assert {candidate.repo_id for candidate in candidates} == {"Qwen/Qwen-Coder-GGUF"} + assert discovery.issues == ["detail Qwen/rate-limited: HTTP 429"] + + +def test_fleet_fit_fails_closed_on_runtime_memory_disk_and_state(): + candidate = _candidate() + fits = analyze_fleet_fit( + candidate, + [ + _node(), + _node(node_id="no-runtime", runtimes=["ollama"]), + _node(node_id="no-memory", capacity_mb=5_000, reserved_mb=1_000), + _node(node_id="no-disk", disk_available_mb=1_000), + _node(node_id="unknown-disk", disk_available_mb=None), + _node(node_id="paused", state="paused"), + ], + ) + assert fits[0].node_id == "gpu-a" and fits[0].fits + reasons = {item.node_id: item.reason for item in fits} + assert "runtime" in reasons["no-runtime"] + assert "memory" in reasons["no-memory"] + assert "disk" in reasons["no-disk"] + assert "unknown" in reasons["unknown-disk"] + assert "not accepting" in reasons["paused"] + + +def test_proposals_compare_only_models_serving_the_same_workload(): + status = { + "nodes": [_node()], + "models": [ + {"model_id": "old-coder", "workload_scores": [["coding", 0.9]]}, + {"model_id": "image", "workload_scores": [["image", 1.0]]}, + ], + } + proposal = build_proposals([_candidate()], status)[0] + assert proposal.state == "benchmark-ready" + assert proposal.compare_models == ("old-coder",) + assert proposal.score > 0 + + +def test_benchmark_requires_real_outputs_to_meet_every_case_floor(): + proposal = build_proposals([_candidate()], {"nodes": [_node()], "models": []})[0] + + def passing(_model: str, prompt: str) -> tuple[str, float]: + if "Python" in prompt: + return "def clamp(value, low, high):\n return max(low, min(high, value))", 12 + if "correlation" in prompt: + return "Correlation is association. Causation means one factor causes another.", 18 + return "GRID", 4 + + qualified, samples = benchmark_candidate(proposal, passing) + assert qualified.state == "qualified" + assert qualified.benchmark_quality == 1.0 + assert qualified.benchmark_samples == 3 + assert len(samples) == 3 + + failed, _ = benchmark_candidate(proposal, lambda _model, _prompt: ("wrong", 1)) + assert failed.state == "benchmark-failed" + assert failed.benchmark_quality == 0.0 + + +def test_benchmark_transport_failure_is_evidence_not_an_exception(): + proposal = build_proposals([_candidate()], {"nodes": [_node()], "models": []})[0] + + def broken(_model: str, _prompt: str): + raise RuntimeError("relay offline") + + failed, samples = benchmark_candidate(proposal, broken, workloads=("coding",)) + assert failed.state == "benchmark-failed" + assert samples[0].error is True + + +def test_scout_state_round_trips_and_is_owner_only(tmp_path): + proposal = build_proposals([_candidate()], {"nodes": [_node()], "models": []})[0] + path = tmp_path / "scout.json" + save_scout_state(path, [proposal], policy=ScoutPolicy()) + loaded = load_scout_state(path) + restored = proposals_from_state(loaded) + assert restored == (proposal,) + if os.name == "posix": + assert path.stat().st_mode & 0o077 == 0 + + +def test_scout_state_rejects_unknown_schema(tmp_path): + path = tmp_path / "scout.json" + path.write_text(json.dumps({"schema_version": 99, "proposals": []})) + with pytest.raises(ValueError, match="unsupported"): + load_scout_state(path) + + +def test_commit_pinned_gguf_source_reaches_exact_download_revision(): + source = f"hf://Qwen/Qwen-Coder-GGUF@{REVISION}/nested/model.gguf" + assert _parse_hugging_face_artifact_source(source) == ( + "Qwen/Qwen-Coder-GGUF", + REVISION, + "nested/model.gguf", + ) + assert f"/resolve/{REVISION}/" in hf_url("Qwen/repo", "model.gguf", REVISION) + with pytest.raises(RuntimeError, match="full 40-hex"): + _parse_hugging_face_artifact_source( + "hf://Qwen/Qwen-Coder-GGUF@main/model.gguf" + ) + + +def test_existing_unpinned_gguf_sources_remain_compatible(): + assert _parse_hugging_face_artifact_source("hf://owner/repo/model.gguf") == ( + "owner/repo", + "main", + "model.gguf", + ) + + +def test_parser_exposes_scout_run_watch_status_and_real_benchmark(): + cli = parser.build_parser() + assert cli.parse_args(["allocator", "scout", "run"]).handler.__name__ == "cmd_allocator_scout_run" + assert cli.parse_args(["allocator", "scout", "watch"]).handler.__name__ == "cmd_allocator_scout_watch" + assert cli.parse_args(["allocator", "scout", "status"]).handler.__name__ == "cmd_allocator_scout_status" + args = cli.parse_args( + [ + "allocator", + "scout", + "benchmark", + "proposal-1", + "--inference-grid", + "forge", + "--deploy-canary", + ] + ) + assert args.handler.__name__ == "cmd_allocator_scout_benchmark" + assert args.deploy_canary is True