freesolo-flash-dev 0.2.56__tar.gz → 0.2.58__tar.gz

freesolo_flash_dev-0.2.58/.claude/skills/verify/SKILL.md

@@ -0,0 +1,11 @@
+# flash verification
+
+Use isolated loopback services to exercise the public FastAPI and CLI surfaces.
+
+1. set `HOME` to a temporary directory so run JSON and `server.db` stay isolated.
+2. launch a synthetic serving backend on an unused localhost port that implements `/adapters` and `/v1/chat/completions` and records request paths.
+3. launch `flash.server.app.create_app()` with uvicorn through a temporary launcher. patch startup credential and external-provider preflights only, and set `FREESOLO_INTERNAL_KEY`, `FLASH_DEPLOY_SYNC=1`, and `FREESOLO_SERVING_URL` to the synthetic backend.
+4. drive the real HTTP routes with curl using the internal bearer key. create dry-run records through `POST /v1/runs`; when a completed artifact is needed, edit only the isolated temporary run JSON before driving deploy, chat, list, cancel, and undeploy routes.
+5. capture API response bodies and the synthetic backend request log. stop both loopback processes after verification.
+
+Use `uv run --project <repo> python ...` for all Python launchers. Never use production credentials or network services.

{freesolo_flash_dev-0.2.56 → freesolo_flash_dev-0.2.58}/.github/workflows/auto-rebake.yml

@@ -105,9 +105,9 @@ jobs:
           FP_BASE: ${{ steps.base.outputs.fp_base }}
           REV: ${{ steps.base.outputs.revision }}
         run: |
-          # arch list is the single source of truth in deps.py (mirrored by the bake matrix; a test
-          # pins that mirror). textual parse -> no flash import / uv sync needed.
-          arches=$(python3 -c "import re,pathlib; t=pathlib.Path('flash/providers/runpod/train/deps.py').read_text(); m=re.search(r'BAKED_PER_SM_ARCHES\s*=\s*frozenset\(\{([^}]+)\}\)', t); print(' '.join(re.findall(r'sm\d+', m.group(1))))")
+          # flash/providers/_worker.py is the single source of truth for the arch list; the bake
+          # matrix mirrors it. read the canonical assignment statically without importing flash.
+          arches=$(python3 docker/kernel_fingerprint.py --print-baked-arches)
           relayer=(); gpu=()
           summary="| arch | published fp_cache | published fp_base | action |\n|---|---|---|---|\n"
           for sm in $arches; do

{freesolo_flash_dev-0.2.56 → freesolo_flash_dev-0.2.58}/.github/workflows/bake-kernel-cache.yml

@@ -5,7 +5,7 @@ name: flash worker kernel-cache (per-arch)
 #   1. offloads the warmup to a RunPod GPU of that arch (GH runners have none) -> build/kernel_cache,
 #   2. docker build --build-arg BUILD_KERNEL_CACHE=true -> ghcr.io/freesolo-co/flash-worker:cu128-<sm>,
 #   3. pushes it. The allocator then selects the cu128-<sm> tag per GPU class for baked arches
-#      (see flash/providers/runpod/train/deps.py).
+#      (see flash/providers/_worker.py).
 #
 # Depends on the base ghcr.io/freesolo-co/flash-worker:cu128 image existing (worker-image.yml) — the
 # warmup runs INSIDE it so the cache matches the image's pinned torch/triton/fla/liger toolchain.

{freesolo_flash_dev-0.2.56 → freesolo_flash_dev-0.2.58}/.github/workflows/worker-image.yml

@@ -1,11 +1,10 @@
 name: flash worker image
 
-# Builds the prebuilt flash worker image (full training stack + flash-attn baked in) and
-# pushes it to GHCR, so Vast/RunPod cold-start skips the per-host dep install (the dominant
-# cold-start cost). Training and serving both default to the pinned WORKER_IMAGE constant
-# (flash/providers/runpod/train.py) — bump that to the new tag. FLASH_WORKER_IMAGE is an
-# operator env override honored by BOTH training (providers/runpod/{jobs,train}.py) and serving
-# (flash/serve/deploy.py).
+# builds the prebuilt flash worker image (full training stack + flash-attn baked in) and
+# pushes it to ghcr, so runpod/vast/lambda training cold-start skips the per-host dependency
+# install (the dominant cold-start cost). training providers select the pinned image and operator
+# override through flash/providers/_worker.py. serving has no consumer for this training-worker
+# image or override.
 #
 # Triggered automatically on flash/** changes (worker training code is baked in; skipping a
 # rebuild leaves GPU workers running stale code) and manually for Dockerfile/dep changes.
@@ -95,9 +94,9 @@ jobs:
           context: .
           file: Dockerfile.worker
           push: true
-          # Pin the freesolo-co namespace: both RunPod and Vast pull the hardcoded
-          # WORKER_IMAGE = ghcr.io/freesolo-co/flash-worker:cu128 (flash/providers/runpod/train.py),
-          # so a fork run must not push to its own owner namespace and leave the real image stale.
+          # pin the freesolo-co namespace: runpod, vast, and lambda training paths pull the image
+          # declared in flash/providers/_worker.py, so a fork run must not push to its own owner
+          # namespace and leave the real image stale.
           tags: ghcr.io/freesolo-co/flash-worker:${{ github.event.inputs.tag || 'cu128' }}
           provenance: false
           # Stamp the kernel-cache fingerprints + the source commit so auto-rebake.yml can read what

{freesolo_flash_dev-0.2.56 → freesolo_flash_dev-0.2.58}/Dockerfile.worker

@@ -1,7 +1,7 @@
 # Flash prebuilt WORKER image — the full training stack baked in for RunPod workers.
 # The worker still fetches the flash `code/**` from HF at runtime, so only DEPS live here.
 #
-# This image is referenced by name in flash/providers/runpod/train/deps.py:WORKER_IMAGE. It MUST be
+# this image name is defined in flash/providers/_worker.py. it must be
 # published to GHCR and the package set PUBLIC (no registry login is configured) so RunPod can
 # pull it. Normally built by
 # .github/workflows/worker-image.yml (manual dispatch); to build/push by hand:
@@ -12,7 +12,7 @@
 #
 # Base is torch 2.10 / cu12.8 / cudnn9 (-devel = has nvcc for any source builds); it ships the
 # CUDA 12.8 runtime libs the PyPI wheels link against. The pinned stack mirrors
-# flash/providers/runpod/train/deps.py:WORKER_DEPS.
+# flash/providers/_worker.py:WORKER_DEPS.
 FROM pytorch/pytorch:2.10.0-cuda12.8-cudnn9-devel
 
 ENV DEBIAN_FRONTEND=noninteractive \
@@ -78,7 +78,7 @@ RUN pip install --no-cache-dir \
 # pure-PyTorch fallback (H100 SXM, Qwen3.5 hidden-2560 LoRA): seq4096 4.2x faster / 1.6x less mem,
 # seq16384 12.6x / 1.9x, loss matches to 1.8e-4. fla is from git: the PyPI wheel is a broken stub
 # (no fla.modules). The git URL is PINNED to a commit (not the moving default branch) for
-# reproducible image builds — keep this SHA in lockstep with flash/providers/runpod/train/deps.py:WORKER_DEPS.
+# reproducible image builds; keep this sha in lockstep with the worker dependencies in flash/providers/_worker.py.
 
 # FlashAttention: PREFER a prebuilt wheel. There is no OFFICIAL Dao-AILab wheel for torch 2.10 yet,
 # but community CI publishes matching wheels (cu128/torch2.10/cp312); CI passes one as FLASH_ATTN_SPEC
@@ -144,13 +144,13 @@ RUN TORCH_CUDA_ARCH_LIST="8.0 8.6 8.9 9.0 10.0 12.0" CAUSAL_CONV1D_FORCE_BUILD=T
     || { echo "causal_conv1d: build/import failed -> removing + GDN packing stays off (Qwen3.5/3.6 train unpacked)"; \
          pip uninstall -y causal-conv1d 2>/dev/null || true; }
 
-# Weight cache (this PR, RUNTIME not baked): when the shared weight-cache network volume is attached
-# (flash runner._assign_weight_cache_volume), the control plane points HF_HOME at the persistent mount
-# per-run (providers/.../deps.weight_cache_env), so the MODEL DOWNLOAD is reused across runs landing in
-# the same region. ONLY HF_HOME (inert model blobs) is redirected — the executable kernel-JIT caches
-# below are deliberately NOT shared (the volume is multi-tenant; a poisoned compiled artifact could
-# affect another tenant). The redirect can't be a static image ENV: the volume mounts at a fixed path
-# only at runtime and only when attached, so a baked ENV would break cache-less (no-volume) cold runs.
+# weight cache (runtime, not baked): when the shared weight-cache network volume is attached,
+# flash/providers/_worker.py:weight_cache_env redirects only base-model prefetching through
+# FLASH_WEIGHT_CACHE_DIR to the persistent volume for reuse across runs. HF_HOME remains ephemeral,
+# so environment and reward downloads are not shared. executable kernel-jit caches are deliberately
+# not shared because the volume is multi-tenant; a poisoned compiled artifact could affect another
+# tenant. the redirect cannot be a static image environment variable: the volume mounts at a fixed
+# path only at runtime and only when attached, so a baked redirect would break cache-less cold runs.
 #
 # Compiled-kernel cache: OPT-IN bake, JIT fallback. The fused Triton (fla/chalk),
 # TorchInductor, and Hopper tilelang kernels the trainer builds on first use are content-addressed

{freesolo_flash_dev-0.2.56 → freesolo_flash_dev-0.2.58}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: freesolo-flash-dev
-Version: 0.2.56
+Version: 0.2.58
 Summary: Flash — managed LoRA post-training (SFT/GRPO/OPD) for Freesolo environments, driven by the `flash` CLI
 Project-URL: Homepage, https://github.com/freesolo-co/flash
 Project-URL: Repository, https://github.com/freesolo-co/flash
@@ -21,8 +21,11 @@ Requires-Dist: fastapi; extra == 'dev'
 Requires-Dist: freesolo>=0.2.54; extra == 'dev'
 Requires-Dist: httpx>=0.27; extra == 'dev'
 Requires-Dist: huggingface-hub>=0.34; extra == 'dev'
+Requires-Dist: jsonschema>=4.23; extra == 'dev'
 Requires-Dist: mypy>=1.13.0; extra == 'dev'
 Requires-Dist: pytest>=9.0.3; extra == 'dev'
+Requires-Dist: referencing>=0.28.4; extra == 'dev'
+Requires-Dist: regex>=2024.11.6; extra == 'dev'
 Requires-Dist: ruff>=0.6; extra == 'dev'
 Requires-Dist: runpod-flash; extra == 'dev'
 Requires-Dist: transformers<5.11,>=5.6; extra == 'dev'
@@ -44,6 +47,9 @@ Requires-Dist: fastapi; extra == 'server'
 Requires-Dist: freesolo>=0.2.54; extra == 'server'
 Requires-Dist: httpx>=0.27; extra == 'server'
 Requires-Dist: huggingface-hub>=0.34; extra == 'server'
+Requires-Dist: jsonschema>=4.23; extra == 'server'
+Requires-Dist: referencing>=0.28.4; extra == 'server'
+Requires-Dist: regex>=2024.11.6; extra == 'server'
 Requires-Dist: runpod-flash; extra == 'server'
 Requires-Dist: transformers<5.11,>=5.6; extra == 'server'
 Requires-Dist: uvicorn; extra == 'server'

{freesolo_flash_dev-0.2.56 → freesolo_flash_dev-0.2.58}/docker/bake_kernel_cache.py

@@ -105,7 +105,7 @@ def main() -> int:
     chalk_spec = os.environ.get("FLASH_CHALK_SPEC", "").strip()
     if not chalk_spec:
         try:
-            from flash.providers.runpod.train.deps import DEFAULT_CHALK_SPEC
+            from flash.providers._worker import DEFAULT_CHALK_SPEC
 
             chalk_spec = DEFAULT_CHALK_SPEC
         except Exception as exc:

{freesolo_flash_dev-0.2.56 → freesolo_flash_dev-0.2.58}/docker/kernel_fingerprint.py

@@ -58,6 +58,8 @@ LABEL_REVISION = "org.opencontainers.image.revision"
 # pip specs in Dockerfile.worker's main stack whose kernels land in the mega-cache. everything else
 # in that block is a base-layer dep (fp_base). fla is matched separately (it carries the git sha).
 _CACHE_PKGS = {"tilelang", "apache-tvm-ffi"}
+_BAKED_ARCHES_NAME = "BAKED_PER_SM_ARCHES"
+_BAKED_ARCH_PATTERN = re.compile(r"sm[0-9]+")
 
 
 def _search(pattern: str, text: str, what: str, *, flags: int = 0) -> str:
@@ -110,6 +112,62 @@ def _python_string_constant(text: str, name: str, what: str) -> str:
     return values[name]
 
 
+def parse_baked_per_sm_arches(text: str, *, source: str = "_worker.py") -> list[str]:
+    """Parse the canonical baked-architecture frozenset without importing flash."""
+    tree = ast.parse(text, filename=source)
+    assignments = [
+        node
+        for node in tree.body
+        if (
+            isinstance(node, ast.Assign)
+            and any(isinstance(t, ast.Name) and t.id == _BAKED_ARCHES_NAME for t in node.targets)
+        )
+        or (
+            isinstance(node, ast.AnnAssign)
+            and isinstance(node.target, ast.Name)
+            and node.target.id == _BAKED_ARCHES_NAME
+        )
+    ]
+    context = f"kernel_fingerprint: {_BAKED_ARCHES_NAME} in {source}"
+    if len(assignments) != 1:
+        raise ValueError(
+            f"{context}: expected exactly one top-level assignment, found {len(assignments)}"
+        )
+
+    assignment = assignments[0]
+    if isinstance(assignment, ast.Assign) and (
+        len(assignment.targets) != 1 or not isinstance(assignment.targets[0], ast.Name)
+    ):
+        raise ValueError(f"{context} must use one named target")
+    value = assignment.value
+    if value is None:
+        raise ValueError(f"{context} must have a value")
+    if not (
+        isinstance(value, ast.Call)
+        and isinstance(value.func, ast.Name)
+        and value.func.id == "frozenset"
+        and len(value.args) == 1
+        and not value.keywords
+        and isinstance(value.args[0], ast.Set)
+    ):
+        raise ValueError(f"{context} must be frozenset({{...}}) of string literals")
+
+    arches: list[str] = []
+    for element in value.args[0].elts:
+        if not isinstance(element, ast.Constant) or not isinstance(element.value, str):
+            raise ValueError(f"{context} must contain strings")
+        if _BAKED_ARCH_PATTERN.fullmatch(element.value) is None:
+            raise ValueError(
+                f"{context} has invalid architecture {element.value!r}; expected sm followed by digits"
+            )
+        arches.append(element.value)
+    if not arches:
+        raise ValueError(f"{context} must not be empty")
+    if len(arches) != len(set(arches)):
+        raise ValueError(f"{context} must not contain duplicates")
+    return arches
+
+
 def _pip_stack_specs(dockerfile: str) -> list[str]:
     """The quoted specs from Dockerfile.worker's first `pip install` block (the pinned worker stack)."""
     specs: list[str] = []
@@ -257,6 +315,11 @@ def main(argv: list[str] | None = None) -> int:
     ap.add_argument(
         "--print-label-keys", action="store_true", help="print the OCI label keys + exit"
     )
+    ap.add_argument(
+        "--print-baked-arches",
+        action="store_true",
+        help="print BAKED_PER_SM_ARCHES from _worker.py + exit",
+    )
     args = ap.parse_args(argv)
 
     if args.print_label_keys:
@@ -265,8 +328,20 @@ def main(argv: list[str] | None = None) -> int:
         print(f"revision={LABEL_REVISION}")
         return 0
 
+    root = Path(args.root)
+    if args.print_baked_arches:
+        worker_source = root / "flash" / "providers" / "_worker.py"
+        try:
+            arches = parse_baked_per_sm_arches(
+                worker_source.read_text(encoding="utf-8"), source=str(worker_source)
+            )
+        except (OSError, SyntaxError, ValueError) as exc:
+            ap.error(str(exc))
+        print(" ".join(arches))
+        return 0
+
     fp_cache, fp_base, cache_inputs, base_inputs = fingerprints(
-        Path(args.root), fa2_spec=args.fa2_spec, fa3_spec=args.fa3_spec
+        root, fa2_spec=args.fa2_spec, fa3_spec=args.fa3_spec
     )
     if args.explain:
         print(

{freesolo_flash_dev-0.2.56 → freesolo_flash_dev-0.2.58}/flash/_channel.py

@@ -1,7 +1,5 @@
 """Single switch between prod and dev channels; ``scripts/build_dev_dist.py`` rewrites CHANNEL to "dev"."""
 
-from __future__ import annotations
-
 CHANNEL = "dev"
 
 # Must stay in lockstep with [project.scripts] in pyproject.toml.

{freesolo_flash_dev-0.2.56 → freesolo_flash_dev-0.2.58}/flash/_fileio.py

@@ -1,4 +1,4 @@
-"""Small shared file-IO helpers for credential/manifest JSON under ``~/.flash``."""
+"""Small shared JSON file helpers for data under ``~/.flash``."""
 
 from __future__ import annotations
 

{freesolo_flash_dev-0.2.56 → freesolo_flash_dev-0.2.58}/flash/_update_check.py

@@ -9,7 +9,6 @@ import re
 import sys
 import threading
 import time
-import urllib.error
 import urllib.request
 
 from flash import __version__
@@ -27,7 +26,7 @@ _PYPI_JSON_URL = f"https://pypi.org/pypi/{PACKAGE_NAME}/json"
 CACHE_PATH = CONFIG_DIR / "update_check.json"
 
 _CHECK_INTERVAL_S = 24 * 60 * 60
-# join >= fetch timeout so the worker finishes writing before process exit kills it.
+# keep the join budget above the fetch timeout to allow parsing and persistence before exit.
 _FETCH_TIMEOUT_S = 1.5
 _JOIN_TIMEOUT_S = 2.0
 
@@ -96,9 +95,7 @@ def _clean_version(value: object) -> str | None:
 
 
 def _read_cache() -> dict:
-    # Coerce non-dict results (e.g. []) since _check_due runs before main()'s error handling.
-    cache = read_json_or_empty(CACHE_PATH)
-    return cache if isinstance(cache, dict) else {}
+    return read_json_or_empty(CACHE_PATH)
 
 
 def _check_due(now: float) -> bool:
@@ -122,7 +119,7 @@ def _fetch_latest_version(timeout: float = _FETCH_TIMEOUT_S) -> str | None:
     try:
         with urllib.request.urlopen(req, timeout=timeout) as resp:
             payload = json.loads(resp.read())
-    except (urllib.error.URLError, OSError, ValueError, TimeoutError) as exc:
+    except (OSError, ValueError) as exc:
         logger.debug("update check: PyPI lookup failed: %s", exc)
         return None
     info = payload.get("info") if isinstance(payload, dict) else None
@@ -152,23 +149,16 @@ def _refresh_cache() -> None:
         logger.debug("update check: refresh failed: %s", exc)
 
 
-def _supports_color() -> bool:
-    return not os.environ.get("NO_COLOR")
-
-
-def _red(text: str) -> str:
-    return f"\033[31m{text}\033[0m" if _supports_color() else text
-
-
 def _build_notice() -> str | None:
     """Build the upgrade notice from the cached PyPI version, or ``None`` if up to date."""
     latest = _clean_version(_read_cache().get("pypi_version"))
     if not latest or _is_prerelease(latest) or not _is_newer(latest, __version__):
         return None
-    return _red(
+    notice = (
         f"A new release of {PACKAGE_NAME} is available: {__version__} -> {latest}\n"
         f"Update with `{UPGRADE_COMMAND}`."
     )
+    return f"\033[31m{notice}\033[0m" if not os.environ.get("NO_COLOR") else notice
 
 
 def maybe_start_update_check() -> threading.Thread | None:

{freesolo_flash_dev-0.2.56 → freesolo_flash_dev-0.2.58}/flash/adapter_artifacts.py

@@ -1,7 +1,5 @@
 """Shared adapter artifact filenames used by training and control-plane code."""
 
-from __future__ import annotations
-
 # The PEFT adapter weights file a deployable adapter must carry. Modern saves use safetensors,
 # but older PEFT/default settings may emit adapter_model.bin.
 ADAPTER_WEIGHT_FILES = ("adapter_model.safetensors", "adapter_model.bin")

{freesolo_flash_dev-0.2.56 → freesolo_flash_dev-0.2.58}/flash/catalog.py

@@ -3,7 +3,7 @@
 from __future__ import annotations
 
 import math
-from dataclasses import asdict, dataclass
+from dataclasses import asdict, dataclass, replace
 from typing import Any
 
 ALGORITHMS = ("sft", "grpo", "opd")
@@ -78,7 +78,7 @@ class ModelInfo:
     recommended_gpu: str = DEFAULT_GPU
     # 0 => GRPO uses min_vram_gb like SFT; set when colocated vLLM rollout needs a bigger card.
     grpo_min_vram_gb: int = 0
-    # 0 => SFT sizes from param-based estimate; set only when a model must not down-route to the cheapest card.
+    # 0 => non-grpo sizing uses the param-based estimate; set when a model must not down-route to the cheapest card.
     sft_min_vram_gb: int = 0
     # vLLM sleep mode (offload the colocate rollout engine between GRPO steps) is NON-FUNCTIONAL for
     # this model: the wake/reload HANGS the rollout (a ~70 GB weight reallocation can't be placed in
@@ -124,9 +124,9 @@ class ModelInfo:
 
     def to_dict(self) -> dict[str, Any]:
         data = asdict(self)
-        serving = data.get("serving")
+        serving = data["serving"]
         if serving is None:
-            data.pop("serving", None)
+            del data["serving"]
         else:
             for key in (
                 "serve_model_id",
@@ -162,7 +162,7 @@ MODELS: dict[str, ModelInfo] = {
         params="1.2B dense (Llama arch)",
         params_b=1.2,
         vocab_size=130_560,
-        algos=("sft", "grpo", "opd"),
+        algos=ALGORITHMS,
         min_vram_gb=12,
         recommended_gpu="RTX 4090",
         serving=ServingCapacity(
@@ -181,7 +181,7 @@ MODELS: dict[str, ModelInfo] = {
         params="0.9B (text-only fine-tune)",
         params_b=0.9,
         vocab_size=248_320,
-        algos=("sft", "grpo", "opd"),
+        algos=ALGORITHMS,
         min_vram_gb=12,
         recommended_gpu="RTX 4090",
         serving=ServingCapacity(
@@ -200,7 +200,7 @@ MODELS: dict[str, ModelInfo] = {
         params="2.3B (text-only fine-tune)",
         params_b=2.3,
         vocab_size=248_320,
-        algos=("sft", "grpo", "opd"),
+        algos=ALGORITHMS,
         min_vram_gb=16,
         recommended_gpu="RTX 4090",
         serving=ServingCapacity(
@@ -218,7 +218,7 @@ MODELS: dict[str, ModelInfo] = {
         params="4.7B (text-only fine-tune)",
         params_b=4.7,
         vocab_size=248_320,
-        algos=("sft", "grpo", "opd"),
+        algos=ALGORITHMS,
         min_vram_gb=32,
         recommended_gpu="RTX 5090",
         serving=ServingCapacity(
@@ -240,7 +240,7 @@ MODELS: dict[str, ModelInfo] = {
         params="9.7B (text-only fine-tune)",
         params_b=9.7,
         vocab_size=248_320,
-        algos=("sft", "grpo", "opd"),
+        algos=ALGORITHMS,
         min_vram_gb=48,
         # NOT QLoRA: peft bnb merge during GRPO rollout diverges trainer precision -> TRL ratio collapses to 0.
         grpo_min_vram_gb=80,
@@ -272,7 +272,7 @@ MODELS: dict[str, ModelInfo] = {
         num_layers=40,
         hidden_size=2048,
         vocab_size=248_320,
-        algos=("sft", "grpo", "opd"),
+        algos=ALGORITHMS,
         min_vram_gb=141,
         # Floor to 100 GB so SFT lands on H200, not the thin-margin consumer Blackwell or 80 GB H100.
         sft_min_vram_gb=100,
@@ -328,19 +328,20 @@ def get_model(model_id: str) -> ModelInfo:
         ) from exc
 
 
+def _model_info_for_serving(model: str | ModelInfo | None) -> ModelInfo | None:
+    if isinstance(model, ModelInfo):
+        return model
+    if isinstance(model, str) and model.strip():
+        return MODELS.get(model.strip())
+    return None
+
+
 def serving_lora_rank_cap(model: str | ModelInfo | None) -> int | None:
-    """Return the model's serving LoRA rank cap, or None when Flash has no local cap.
+    """Return the model-specific serving LoRA rank cap.
 
-    Serving capacity is model-specific: small serving models currently allow rank 64, while larger
-    serving paths currently allow rank 32. Unknown/open-policy models intentionally return None instead
-    of inheriting a global fallback.
+    Unknown and open-policy models return None.
     """
-    if isinstance(model, ModelInfo):
-        info = model
-    elif isinstance(model, str) and model.strip():
-        info = MODELS.get(model.strip())
-    else:
-        info = None
+    info = _model_info_for_serving(model)
     if info is None or info.serving is None:
         return None
     return int(info.serving.max_lora_rank)
@@ -355,12 +356,7 @@ def serving_context_cap(model: str | ModelInfo | None) -> int | None:
     ``flash.lora_rank.preflight_train_context_within_serving``). Resolution mirrors
     ``serving_lora_rank_cap``: unknown/open-policy models return None rather than a global fallback.
     """
-    if isinstance(model, ModelInfo):
-        info = model
-    elif isinstance(model, str) and model.strip():
-        info = MODELS.get(model.strip())
-    else:
-        info = None
+    info = _model_info_for_serving(model)
     if info is None or info.serving is None:
         return None
     return int(info.serving.max_model_len)
@@ -372,26 +368,66 @@ def vocab_size_for(model_id: str) -> int:
     return info.vocab_size if info is not None else _DEFAULT_VOCAB_SIZE
 
 
+def resolve_vocab_size(model_id: str, revision: str = "") -> int:
+    """vocab_size for a model, revision-aware when a commit is pinned (mirrors resolve_params_b).
+
+    catalog vocab by default; for a pinned revision, the hf config vocab validated against the
+    catalog (fail-closed, same as the vram path).
+    """
+    info = MODELS.get(model_id)
+    if revision:
+        from flash.engine.vram import _validated_revision_geometry, fetch_hf_model_geometry
+
+        if info is not None:
+            _params_b, vocab = _validated_revision_geometry(model_id, revision, info)
+            return int(vocab or info.vocab_size)
+        _p, vocab, _h, _l = fetch_hf_model_geometry(model_id, revision, strict=True)
+        if vocab:
+            return int(vocab)
+    return vocab_size_for(model_id)
+
+
 def resolve_model(
     model_id: str,
     algorithm: str,
     policy: str = "catalog",
     gpu: str | None = None,
+    model_revision: str = "",
 ) -> ModelInfo:
     """Resolve a model under the configured policy; "allow" accepts any HF model."""
     algo = normalize_algorithm(algorithm)
     if model_id in MODELS:
-        return validate_model_for_algorithm(model_id, algo)
+        info = validate_model_for_algorithm(model_id, algo)
+        if model_revision:
+            from flash.engine.vram import _validated_revision_geometry
+
+            params_b, vocab_size = _validated_revision_geometry(model_id, model_revision, info)
+            info = replace(
+                info,
+                params_b=params_b,
+                params=f"{params_b:.1f}B",
+                vocab_size=vocab_size,
+                min_disk_gb=max(info.min_disk_gb, int(params_b * 2) + 64),
+            )
+        return info
     if policy != "allow":
         return get_model(model_id)
-    return _resolve_open_model(model_id, algo, gpu)
+    return _resolve_open_model(model_id, algo, gpu, model_revision=model_revision)
 
 
-def _resolve_open_model(model_id: str, algo: str, gpu: str | None) -> ModelInfo:
+def _resolve_open_model(
+    model_id: str, algo: str, gpu: str | None, *, model_revision: str = ""
+) -> ModelInfo:
     """Synthesize a ModelInfo for the open-model "allow" policy via a coarse HF VRAM-fit estimate."""
     from flash.engine.vram import check_fit
 
-    est = check_fit(model_id, algo, gpu or DEFAULT_GPU)
+    resolved_gpu = gpu or DEFAULT_GPU
+    est = check_fit(
+        model_id,
+        algo,
+        resolved_gpu,
+        model_revision=model_revision,
+    )
     if est.verdict == "too_big":
         raise ValueError(
             f"{model_id} does not fit the requested GPU: {est.describe()}. "
@@ -399,19 +435,20 @@ def _resolve_open_model(model_id: str, algo: str, gpu: str | None) -> ModelInfo:
         )
     if est.verdict in ("tight", "unknown"):
         print(f"warning: open-model policy: {est.describe()}")
-    params = f"{est.params_b:.1f}B" if est.params_b else "unknown size"
-    min_disk = int(est.params_b * 2) + 64 if est.params_b else 0
+    params_b = est.params_b or 0.0
+    params = f"{params_b:.1f}B" if params_b else "unknown size"
+    min_disk = int(params_b * 2) + 64 if params_b else 0
     return ModelInfo(
         id=model_id,
         display_name=model_id,
         params=params,
         # Carry the estimated/HF param count straight through (0.0 when size is unknown) so downstream
         # sizing reads ``params_b`` directly — no re-parsing the display string.
-        params_b=est.params_b or 0.0,
+        params_b=params_b,
         algos=ALGORITHMS,
         min_vram_gb=math.ceil(est.est_gb) if est.est_gb else 24,
         min_disk_gb=min_disk,
-        recommended_gpu=gpu or DEFAULT_GPU,
+        recommended_gpu=resolved_gpu,
         thinking="unknown",
         notes="unlisted model accepted via the open-model policy (not curated/validated)",
     )

{freesolo_flash_dev-0.2.56 → freesolo_flash_dev-0.2.58}/flash/cli/__init__.py

@@ -11,20 +11,14 @@ from typing import NoReturn
 
 from flash import __version__
 from flash._channel import CLI_NAME
-from flash._logging import configure_logging, get_logger
+from flash._logging import configure_logging
 from flash._update_check import emit_update_notice, maybe_start_update_check
 from flash.cli import render
 
-# Command handlers + the patched client surface live in submodules; re-export them so
-# `flash.cli` stays the single public import surface (and so monkeypatching
-# `flash.cli.commands` reaches the bare globals the handlers read).
+# package-level command imports remain available through flash.cli.
 from flash.cli.commands import (  # noqa: F401
     _CLI_DONE_STATES,
-    _OK_STATES,
-    _STARTER_ENV_PY,
     _USER_ERRORS,
-    _follow_run,
-    _poll_logs,
     client_from_config,
     cmd_cancel,
     cmd_chat,
@@ -32,7 +26,6 @@ from flash.cli.commands import (  # noqa: F401
     cmd_deploy,
     cmd_deployments,
     cmd_env_list,
-    cmd_env_setup,
     cmd_export,
     cmd_gpus,
     cmd_log,
@@ -46,10 +39,9 @@ from flash.cli.commands import (  # noqa: F401
     cmd_whoami,
     verify_freesolo_key,
 )
+from flash.cli.env_setup import cmd_env_setup
 from flash.cli.envpush import cmd_env_delete, cmd_env_pull, cmd_env_push
 
-logger = get_logger("flash.cli")
-
 # Themed `flash --help` catalog. Groups are ordered along the training workflow; each row's
 # summary is the short one-liner the themed grid shows (the verbose per-command text stays on
 # every subparser's own `help=` / `<cmd> --help`). test_cli_help.py asserts these rows stay in
@@ -392,12 +384,6 @@ def _build_parser() -> argparse.ArgumentParser:
         help="run id for the final adapter, or RUN_ID/step-N for a saved checkpoint",
     )
     deploy.add_argument("--dry-run", action="store_true")
-    deploy.add_argument(
-        "--no-verify",
-        action="store_true",
-        help="skip the server-side post-deploy smoke generation (registration alone does NOT "
-        "guarantee the adapter serves; only ready deployments should be scored by evals)",
-    )
     deploy.set_defaults(func=cmd_deploy)
 
     undeploy = sub.add_parser("undeploy", help="tear down a run's serving endpoint")
@@ -430,6 +416,11 @@ def _build_parser() -> argparse.ArgumentParser:
     export.set_defaults(func=cmd_export)
 
     deployments = sub.add_parser("deployments", help="list active serving deployments")
+    deployments.add_argument(
+        "--json",
+        action="store_true",
+        help="print machine-readable deployment records",
+    )
     deployments.set_defaults(func=cmd_deployments)
 
     chat = sub.add_parser("chat", help="chat with a deployed adapter")