freesolo-flash-dev 1.0.3__tar.gz → 1.0.4__tar.gz

{freesolo_flash_dev-1.0.3 → freesolo_flash_dev-1.0.4}/Dockerfile.worker

@@ -55,7 +55,7 @@ RUN pip install --no-cache-dir \
     "flashinfer-python==0.6.6" \
     "bitsandbytes>=0.49" \
     "datasets>=4.7,<6" \
-    "huggingface_hub>=0.25" \
+    "huggingface_hub>=1.2.0" \
     "accelerate>=1.4" \
     "wandb>=0.17" \
     "flash-linear-attention @ git+https://github.com/fla-org/flash-linear-attention.git@f0e213dbd8b5fb90c3c7eca869ac1706d5377139" \

{freesolo_flash_dev-1.0.3 → freesolo_flash_dev-1.0.4}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: freesolo-flash-dev
-Version: 1.0.3
+Version: 1.0.4
 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
@@ -20,7 +20,7 @@ Requires-Dist: datasets>=2.19; extra == 'dev'
 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: huggingface-hub>=1.2.0; 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'
@@ -35,7 +35,7 @@ Requires-Dist: accelerate>=1.4; extra == 'gpu'
 Requires-Dist: bitsandbytes>=0.49; extra == 'gpu'
 Requires-Dist: datasets>=2.19; extra == 'gpu'
 Requires-Dist: freesolo>=0.2.54; extra == 'gpu'
-Requires-Dist: huggingface-hub>=0.34; extra == 'gpu'
+Requires-Dist: huggingface-hub>=1.2.0; extra == 'gpu'
 Requires-Dist: peft>=0.19; extra == 'gpu'
 Requires-Dist: torch==2.10.0; extra == 'gpu'
 Requires-Dist: transformers<5.11,>=5.6; extra == 'gpu'
@@ -46,7 +46,7 @@ Requires-Dist: datasets>=2.19; extra == 'server'
 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: huggingface-hub>=1.2.0; 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'

{freesolo_flash_dev-1.0.3 → freesolo_flash_dev-1.0.4}/flash/cli/__init__.py

@@ -424,7 +424,11 @@ def _build_parser() -> argparse.ArgumentParser:
     deployments.set_defaults(func=cmd_deployments)
 
     chat = sub.add_parser("chat", help="chat with a deployed adapter")
-    chat.add_argument("run_id")
+    chat.add_argument(
+        "run_id",
+        help="a bare RUN_ID (its current deployment), RUN_ID/step-N (a specific deployed "
+        "checkpoint), or a full immutable adapter revision",
+    )
     chat.add_argument("-m", "--message", required=True)
     chat.add_argument(
         "--system",

{freesolo_flash_dev-1.0.3 → freesolo_flash_dev-1.0.4}/flash/cli/commands.py

@@ -160,8 +160,9 @@ def cmd_gpus(args) -> int:
         (info for info in GPU_INFO.values() if info.enum_member), key=lambda g: g.hourly_usd
     )
     tip = (
-        "Tip: GPU class selection is fully automatic — the submit-time allocator always picks the\n"
-        "cheapest validated managed class that fits the model, so you don't pin a GPU type."
+        "Tip: GPU allocation is automatic by default.\n"
+        "The allocator picks the cheapest validated class that fits. Pin a specific class by "
+        'adding exact_type = "<CLASS>" to the [gpu] section.'
     )
     if render.styled():
         rows = [(info.name, info.vram_gb, runpod_rates.get(info.name)) for info in infos]
@@ -383,18 +384,26 @@ def _log_follow_progress(status: dict | None, fallback_state: str) -> tuple[str,
     parts = [state]
     heartbeat = status.get("last_heartbeat") if isinstance(status, dict) else None
     if isinstance(heartbeat, dict):
+        heartbeat_age_seconds = render._heartbeat_age_seconds(heartbeat.get("ts"))
         stage = heartbeat.get("stage")
         if stage:
             parts.append(f"stage={stage}")
+            if state == "running":
+                warmup = render.warmup_message(
+                    stage,
+                    heartbeat_age_seconds,
+                    render.heartbeat_is_current_attempt(status, heartbeat),
+                )
+                if warmup:
+                    parts.append(warmup)
         step = heartbeat.get("step")
         if step is not None:
             parts.append(f"step={step}")
         # live heartbeat age so a long quiet phase reads as "alive, throttled" not "frozen".
         # minute granularity: the non-TTY follow path prints a line whenever this string changes,
         # so a seconds-precision age would emit one line per poll.
-        ts = heartbeat.get("ts")
-        if isinstance(ts, (int, float)) and ts > 0:
-            mins = int(max(0.0, time.time() - ts) // 60)
+        if heartbeat_age_seconds is not None:
+            mins = int(heartbeat_age_seconds // 60)
             parts.append(f"hb={mins}m" if mins else "hb=<1m")
     realized = status.get("realized_cost_usd")
     if realized is not None:
@@ -737,18 +746,11 @@ def cmd_chat(args) -> int:
     if revision is None and parsed is None:
         print(
             f"invalid chat target {args.run_id!r} "
-            "(expected a bare <run_id> or full immutable adapter revision)",
-            file=sys.stderr,
-        )
-        return 1
-    if revision is None and parsed[1] is not None:
-        print(
-            "RUN_ID/step-N is not a valid chat target because it would route through the mutable "
-            "run alias; use the full immutable adapter revision returned by `flash deployments`",
+            "(expected a bare <run_id>, <run_id>/step-N, or full immutable adapter revision)",
             file=sys.stderr,
         )
         return 1
-    chat_target = args.run_id if revision is not None else parsed[0]
+    chat_target = args.run_id
     client = client_from_config()
     messages = [{"role": "user", "content": args.message}]
     system = getattr(args, "system", None)

{freesolo_flash_dev-1.0.3 → freesolo_flash_dev-1.0.4}/flash/cli/envpush.py

@@ -11,6 +11,7 @@ import os
 import re
 import sys
 import tempfile
+from collections.abc import Iterator
 from pathlib import Path
 from typing import TYPE_CHECKING, Literal, TypeAlias
 
@@ -174,6 +175,10 @@ def cmd_env_delete(args) -> int:
 
 
 _ENV_ENTRYPOINT = "environment.py"
+_ENV_SYSPATH_BOOTSTRAP = (
+    "import os as _flash_os, sys as _flash_sys\n"
+    "_flash_sys.path.insert(0, _flash_os.path.dirname(__file__))\n"
+)
 _ENV_PUSH_IGNORED_NAMES = frozenset(
     {
         ".prime",
@@ -183,26 +188,47 @@ _ENV_PUSH_IGNORED_NAMES = frozenset(
         ".venv",
         ".mypy_cache",
         ".pytest_cache",
-        "pyproject.toml",
-        "source",
+        ".ruff_cache",
+        ".idea",
+        ".vscode",
+        ".ds_store",
+        ".git-worktrees",
+        ".env",
+        ".netrc",
+        ".aws",
+        ".ssh",
+        "venv",
+        "node_modules",
+        "credentials",
+        "secrets",
     }
 )
-# ``.md`` is included so the ``TRAINING.md`` playbook `flash env setup` scaffolds (and any
-# user-authored README/NOTES) travels with the env into the hub and back out through
-# ``flash env pull`` — a published env should carry its own training guidance, not just code+data.
-_ENV_PUSH_SIDECAR_SUFFIXES = frozenset(
-    {
-        ".csv",
-        ".json",
-        ".jsonl",
-        ".md",
-        ".parquet",
-        ".tsv",
-        ".txt",
-        ".yaml",
-        ".yml",
-    }
+_ENV_PUSH_ROOT_IGNORED_NAMES = frozenset({"pyproject.toml", "source"})
+_ENV_PUSH_SECRET_PATTERNS = (
+    ".env.*",
+    "*.env",
+    "*.pem",
+    "*.key",
+    "*.pfx",
+    "*.p12",
+    "credentials*",
+    "id_rsa*",
+    "id_dsa*",
+    "id_ecdsa*",
+    "id_ed25519*",
 )
+# secret patterns match secret data FILES only (see _ignore_env_push_path): they never prune a
+# directory, so a package tree like credentials_store/ still ships, and python source is exempt,
+# so helper modules like credentials.py or id_ed25519_loader.py travel instead of being silently
+# dropped and breaking the worker with ModuleNotFoundError. that keeps the prefixes broad enough
+# to catch backup variants (credentials_backup, id_rsa_backup) without eating code or packages.
+_ENV_PUSH_CODE_SUFFIXES = frozenset({".py", ".pyi"})
+_ENV_PUSH_MAX_TOTAL_BYTES = 256 * 1024 * 1024
+
+try:
+    from flash.envs.archive_policy import ARCHIVE_MEMBER_LIMIT as _ENV_PUSH_MAX_FILES
+except ImportError:
+    _ENV_PUSH_MAX_FILES = 5000
 
 
 def _normalize_env_name(raw: str) -> str | None:
@@ -227,16 +253,12 @@ def _normalize_env_name(raw: str) -> str | None:
 
 def _with_syspath_bootstrap(env_source: str) -> str:
     """Prepend a sys.path bootstrap so a published env can resolve shipped sibling helpers."""
-    bootstrap = (
-        "import os as _flash_os, sys as _flash_sys\n"
-        "_flash_sys.path.insert(0, _flash_os.path.dirname(__file__))\n"
-    )
     import ast
 
     try:
         tree = ast.parse(env_source)
     except SyntaxError:
-        return bootstrap + env_source
+        return _ENV_SYSPATH_BOOTSTRAP + env_source
     insert_after = 0
     body = tree.body
     i = 0
@@ -252,7 +274,16 @@ def _with_syspath_bootstrap(env_source: str) -> str:
         insert_after = body[i].end_lineno or insert_after
         i += 1
     lines = env_source.splitlines(keepends=True)
-    return "".join(lines[:insert_after]) + bootstrap + "".join(lines[insert_after:])
+    return (
+        "".join(lines[:insert_after])
+        + _ENV_SYSPATH_BOOTSTRAP
+        + "".join(lines[insert_after:])
+    )
+
+
+def _raise_walk_error(error: OSError) -> None:
+    """fail env packaging loudly instead of silently skipping an unreadable directory."""
+    raise error
 
 
 def _tar_b64(directory: Path) -> str:
@@ -266,44 +297,130 @@ def _tar_b64(directory: Path) -> str:
     with tarfile.open(fileobj=buf, mode="w:gz") as tar:
         for root, dirs, files in os.walk(directory):
             root_path = Path(root)
-            dirs[:] = sorted(d for d in dirs if d not in _ENV_PUSH_IGNORED_NAMES)
-            for name in dirs:
-                child = root_path / name
-                tar.add(child, arcname=str(child.relative_to(directory)), recursive=False)
+            dirs[:] = sorted(d for d in dirs if d.lower() not in _ENV_PUSH_IGNORED_NAMES)
             for name in sorted(files):
                 child = root_path / name
                 tar.add(child, arcname=str(child.relative_to(directory)), recursive=False)
     return base64.b64encode(buf.getvalue()).decode()
 
 
-def _copy_env_sidecars(env_root: Path, dest: Path, *, entrypoint: Path) -> None:
-    """Copy helper code and data sidecars beside environment.py."""
-    import shutil
+def _ignore_env_push_path(path: Path, *, env_root: Path, entrypoint: Path) -> bool:
+    """Return whether a source path must not be included in an env package."""
+    import fnmatch
+
+    name = path.name
+    lowered_name = name.lower()
+    if path.is_symlink():
+        return True
+    if path.is_dir() and (path / "pyvenv.cfg").is_file():
+        return True
+    if path == entrypoint or (
+        path.parent == env_root and lowered_name == _ENV_ENTRYPOINT.lower()
+    ):
+        return True
+    if lowered_name.startswith(".") or lowered_name in _ENV_PUSH_IGNORED_NAMES:
+        return True
+    if path.parent == env_root and lowered_name in _ENV_PUSH_ROOT_IGNORED_NAMES:
+        return True
+    if (
+        path.is_file()
+        and path.suffix.lower() not in _ENV_PUSH_CODE_SUFFIXES
+        and any(fnmatch.fnmatchcase(lowered_name, pattern) for pattern in _ENV_PUSH_SECRET_PATTERNS)
+    ):
+        return True
+    return not (path.is_dir() or path.is_file())
+
+
+def _iter_env_sidecar_files(
+    env_root: Path, *, entrypoint: Path, include_full_tree: bool
+) -> Iterator[tuple[Path, Path]]:
+    """Yield publishable sidecar files and their package-relative paths."""
+    import os
 
-    for child in sorted(env_root.iterdir()):
-        if (
-            child == entrypoint
-            or child.name == _ENV_ENTRYPOINT
-            or child.name in _ENV_PUSH_IGNORED_NAMES
+    roots = [env_root]
+    if not include_full_tree:
+        # a single-file push carries only the entrypoint, its dataset sidecars, and its own
+        # docs. ship user-authored docs so the synthesized stub readme does not stand in for
+        # real training guidance the user shipped next to the module.
+        for doc_name in ("README.md", "TRAINING.md"):
+            doc = env_root / doc_name
+            if doc.is_file() and not _ignore_env_push_path(
+                doc, env_root=env_root, entrypoint=entrypoint
+            ):
+                yield doc, doc.relative_to(env_root)
+        roots = [
+            env_root / name
+            for name in ("dataset", "datasets")
+            if (env_root / name).is_dir() and not (env_root / name).is_symlink()
+        ]
+
+    for walk_root in roots:
+        for root, dirs, files in os.walk(
+            walk_root, topdown=True, followlinks=False, onerror=_raise_walk_error
         ):
-            continue
-        if child.name.startswith("."):
-            continue
-        target = dest / child.name
-        if child.is_dir():
-            if child.name == "dataset":
-                shutil.copytree(
-                    child,
-                    target,
-                    ignore=shutil.ignore_patterns(*_ENV_PUSH_IGNORED_NAMES),
+            root_path = Path(root)
+            dirs[:] = sorted(
+                name
+                for name in dirs
+                if not _ignore_env_push_path(
+                    root_path / name, env_root=env_root, entrypoint=entrypoint
                 )
-            continue
-        if not child.is_file():
-            continue
-        if (
-            child.suffix == ".py" and not child.name.startswith("__")
-        ) or child.suffix.lower() in _ENV_PUSH_SIDECAR_SUFFIXES:
-            shutil.copy2(child, target)
+            )
+            for name in sorted(files):
+                child = root_path / name
+                if _ignore_env_push_path(child, env_root=env_root, entrypoint=entrypoint):
+                    continue
+                yield child, child.relative_to(env_root)
+
+
+def _check_env_push_limits(
+    env_root: Path, *, entrypoint: Path, include_full_tree: bool, env_name: str
+) -> None:
+    """Reject oversized source trees before copying or building an in-memory archive."""
+    files = 1
+    total_bytes = entrypoint.stat().st_size + len(_ENV_SYSPATH_BOOTSTRAP.encode())
+    directories: set[Path] = set()
+    has_readme = False
+    for child, relative in _iter_env_sidecar_files(
+        env_root, entrypoint=entrypoint, include_full_tree=include_full_tree
+    ):
+        files += 1
+        total_bytes += child.stat().st_size
+        has_readme = has_readme or relative == Path("README.md")
+        # the server counts every file AND directory against the same member cap when it
+        # repackages the checkout for download, so count unique parent dirs here too.
+        directories.update(parent for parent in relative.parents if parent != Path("."))
+    if not has_readme:
+        files += 1
+        total_bytes += len(f"# {env_name}\n\nFlash Freesolo environment.\n".encode())
+
+    members = files + len(directories)
+    if total_bytes > _ENV_PUSH_MAX_TOTAL_BYTES:
+        raise ValueError(
+            f"environment package totals {_human_bytes(total_bytes)} "
+            f"(limit {_human_bytes(_ENV_PUSH_MAX_TOTAL_BYTES)}); "
+            "remove large artifacts or use a smaller dataset"
+        )
+    if members > _ENV_PUSH_MAX_FILES:
+        raise ValueError(
+            f"selected environment contains {members:,} files and directories "
+            f"(limit {_ENV_PUSH_MAX_FILES:,}); "
+            "remove large artifacts or use a smaller dataset"
+        )
+
+
+def _copy_env_sidecars(
+    env_root: Path, dest: Path, *, entrypoint: Path, include_full_tree: bool
+) -> None:
+    """Copy publishable sidecar files beside environment.py, creating directories lazily."""
+    import shutil
+
+    for child, relative in _iter_env_sidecar_files(
+        env_root, entrypoint=entrypoint, include_full_tree=include_full_tree
+    ):
+        target = dest / relative
+        target.parent.mkdir(parents=True, exist_ok=True)
+        shutil.copy2(child, target)
 
 
 def _human_bytes(n: int) -> str:
@@ -385,8 +502,11 @@ def cmd_env_push(args) -> int:
     src = Path(args.path)
     if not src.exists():
         return _err(f"no such path: {src}")
+    if src.is_symlink():
+        return _err(f"cannot publish {src}: symlinks are not allowed")
 
-    if src.is_dir():
+    include_full_tree = src.is_dir()
+    if include_full_tree:
         canonical_entrypoint = src / _ENV_ENTRYPOINT
         if canonical_entrypoint.is_file():
             entrypoint = canonical_entrypoint
@@ -410,11 +530,29 @@ def cmd_env_push(args) -> int:
     else:
         return _err(f"cannot publish {src}: expected a Freesolo .py module or an env directory.")
 
+    if entrypoint.is_symlink():
+        return _err(f"cannot publish {entrypoint}: symlinks are not allowed")
+
+    try:
+        _check_env_push_limits(
+            env_root,
+            entrypoint=entrypoint,
+            include_full_tree=include_full_tree,
+            env_name=env_name,
+        )
+    except (OSError, ValueError) as exc:
+        return _err(f"cannot publish {src}: {exc}")
+
     with tempfile.TemporaryDirectory(prefix="flash-env-push-") as tmp:
         pkg = Path(tmp)
         module_source = entrypoint.read_text()
         (pkg / _ENV_ENTRYPOINT).write_text(_with_syspath_bootstrap(module_source))
-        _copy_env_sidecars(env_root, pkg, entrypoint=entrypoint)
+        _copy_env_sidecars(
+            env_root,
+            pkg,
+            entrypoint=entrypoint,
+            include_full_tree=include_full_tree,
+        )
         # Only synthesize a stub README when the env didn't ship its own (now carried as a
         # ``.md`` sidecar) — don't clobber a user-authored README with boilerplate.
         readme = pkg / "README.md"

{freesolo_flash_dev-1.0.3 → freesolo_flash_dev-1.0.4}/flash/cli/render.py

@@ -533,21 +533,76 @@ def _humanize_ts(value) -> str | None:
     return datetime.datetime.fromtimestamp(value, datetime.UTC).strftime("%Y-%m-%d %H:%M UTC")
 
 
-def _humanize_age(value) -> str | None:
-    """Format an epoch seconds value as a compact age ("3m ago"), None for non-numbers."""
+def _heartbeat_age_seconds(value: object) -> float | None:
+    """Return heartbeat age in seconds, or None for an unusable timestamp."""
     if not isinstance(value, (int, float)) or value <= 0:
         return None
-    secs = max(0.0, time.time() - value)
-    if secs < 90:
-        return f"{int(secs)}s ago"
-    if secs < 5400:
-        return f"{int(secs // 60)}m ago"
-    return f"{secs / 3600:.1f}h ago"
+    return max(0.0, time.time() - value)
+
+
+def _humanize_age_seconds(seconds: float | None) -> str | None:
+    if seconds is None:
+        return None
+    if seconds < 90:
+        return f"{int(seconds)}s ago"
+    if seconds < 5400:
+        return f"{int(seconds // 60)}m ago"
+    return f"{seconds / 3600:.1f}h ago"
+
+
+def _humanize_age(value) -> str | None:
+    """Format an epoch seconds value as a compact age ("3m ago"), None for non-numbers."""
+    return _humanize_age_seconds(_heartbeat_age_seconds(value))
 
 
 # heartbeat age past which the panel reminds that quiet is normal: worker uploads are throttled
-# (240s quiet phases, up to 900s mid-training), so a frozen ts is usually NOT a dead worker.
+# (240s quiet phases, up to 900s mid-training), so a frozen ts is usually not a dead worker.
 _HB_QUIET_HINT_AFTER_S = 300.0
+_WARMUP_HEARTBEAT_FRESH_FOR_S = 1200.0
+_WARMUP_STAGES = frozenset({"rl_train_start", "rl_initializing"})
+
+
+def heartbeat_is_current_attempt(obj: dict, heartbeat: dict) -> bool:
+    """False only when the heartbeat provably belongs to a superseded retry attempt.
+
+    Retries reuse the seed's heartbeat path, so a recovered run flips back to ``running`` for the
+    replacement worker while ``last_heartbeat`` can still be the previous attempt's setup ping until
+    the new worker publishes one. ``remote.attempt`` is the live attempt; ``last_heartbeat.attempt``
+    is the one that produced the ping. When the live attempt is known, keep the reassurance only for a
+    heartbeat whose attempt matches it. When it is unknown (e.g. a managed status payload that omits
+    ``remote``), fall back to ``warmup_message``'s age gating rather than suppress, so warmup still
+    reassures on planes that do not surface a live attempt.
+    """
+    # reuse the poller's attempt-identity contract (a bounded nonnegative int, never a bool, string,
+    # or float) so the status display and stall detection agree on what a valid attempt is.
+    from flash.providers._poll import _attempt_int
+
+    remote = obj.get("remote")
+    current_attempt = _attempt_int(remote.get("attempt")) if isinstance(remote, dict) else None
+    if current_attempt is None:
+        return True
+    return _attempt_int(heartbeat.get("attempt")) == current_attempt
+
+
+def warmup_message(
+    stage: object,
+    heartbeat_age_seconds: float | None,
+    from_current_attempt: bool = True,
+) -> str | None:
+    """Explain healthy RL setup stages only while the heartbeat is fresh and from the live attempt."""
+    stage_name = str(stage)
+    if stage_name not in _WARMUP_STAGES:
+        return None
+    if not from_current_attempt:
+        return None
+    if heartbeat_age_seconds is None:
+        return None
+    if heartbeat_age_seconds > _WARMUP_HEARTBEAT_FRESH_FOR_S:
+        return None
+    return (
+        f"warming up (stage={stage_name}): initializing model, vLLM, and training kernels - "
+        "typically several minutes, sometimes 15-20 min; setup is not billed; do not cancel"
+    )
 
 
 def _heartbeat_pairs(obj: dict) -> list[tuple[str, str]]:
@@ -562,11 +617,19 @@ def _heartbeat_pairs(obj: dict) -> list[tuple[str, str]]:
     if hb.get("liveness"):
         worker += " · alive ping"
     pairs = [("worker", worker)]
-    ts = hb.get("ts")
-    age = _humanize_age(ts)
+    heartbeat_age_seconds = _heartbeat_age_seconds(hb.get("ts"))
+    running = str(obj.get("state") or "") == "running"
+    if running:
+        warmup = warmup_message(
+            hb.get("stage"),
+            heartbeat_age_seconds,
+            heartbeat_is_current_attempt(obj, hb),
+        )
+        if warmup:
+            pairs.append(("warmup", warmup))
+    age = _humanize_age_seconds(heartbeat_age_seconds)
     if age:
-        running = str(obj.get("state") or "") == "running"
-        if running and isinstance(ts, (int, float)) and (time.time() - ts) > _HB_QUIET_HINT_AFTER_S:
+        if running and heartbeat_age_seconds > _HB_QUIET_HINT_AFTER_S:
             age += _dim(
                 "  (heartbeat uploads are throttled; quiet is not dead - check flash log -f <run-id>)"
             )

{freesolo_flash_dev-1.0.3 → freesolo_flash_dev-1.0.4}/flash/client/http.py

@@ -121,6 +121,9 @@ class _ProgressReader:
         return chunk
 
 
+_CHAT_STEP_SELECTOR_CAPABILITY = "chat_step_selector_v1"
+
+
 def _validate_chat_messages(messages: list[dict]) -> None:
     if not isinstance(messages, list):
         raise ClientError("chat messages must be a list")
@@ -129,24 +132,20 @@ def _validate_chat_messages(messages: list[dict]) -> None:
             raise ClientError(f"chat messages[{index}] must be an object")
 
 
-def _parse_chat_target(target: str) -> tuple[str, str | None]:
+def _parse_chat_target(target: str) -> tuple[str, str | None, int | None]:
     from flash.schema import parse_adapter_revision, parse_checkpoint_ref
 
     revision = parse_adapter_revision(target)
     if revision is not None:
-        return revision[0], target.strip()
+        return revision[0], target.strip(), None
     parsed = parse_checkpoint_ref(target)
     if parsed is None:
         raise ClientError(
-            "invalid run id: expected a bare RUN_ID or full immutable adapter revision"
+            "invalid run id: expected a bare RUN_ID, RUN_ID/step-N, or a full immutable adapter "
+            "revision"
         )
     run_id, step = parsed
-    if step is not None:
-        raise ClientError(
-            "RUN_ID/step-N is not a valid chat target because it would route through the mutable "
-            "run alias; use the full immutable adapter revision returned by flash deployments"
-        )
-    return run_id, None
+    return run_id, None, step
 
 
 def _prepare_chat_request(
@@ -157,7 +156,7 @@ def _prepare_chat_request(
     *,
     stream: bool = False,
 ) -> tuple[str, dict[str, Any]]:
-    base_run_id, adapter_revision = _parse_chat_target(target)
+    base_run_id, adapter_revision, step = _parse_chat_target(target)
     _validate_chat_messages(messages)
     body: dict[str, Any] = {
         "messages": messages,
@@ -168,6 +167,8 @@ def _prepare_chat_request(
         body["stream"] = True
     if adapter_revision is not None:
         body["adapter_revision"] = adapter_revision
+    elif step is not None:
+        body["step"] = step
     return base_run_id, body
 
 
@@ -289,6 +290,15 @@ class ApiClient:
     def health(self) -> dict:
         return self._request("GET", "/v1/health", timeout=10.0)
 
+    def _require_chat_step_selector(self) -> None:
+        capabilities = self.health().get("capabilities")
+        if not isinstance(capabilities, list) or _CHAT_STEP_SELECTOR_CAPABILITY not in capabilities:
+            raise ClientError(
+                "chat checkpoint selectors require a control plane that advertises "
+                f"{_CHAT_STEP_SELECTOR_CAPABILITY}; use a full immutable adapter revision or "
+                "upgrade the control plane"
+            )
+
     def publish_env(
         self,
         *,
@@ -458,6 +468,8 @@ class ApiClient:
             temperature,
             max_tokens,
         )
+        if "step" in body:
+            self._require_chat_step_selector()
         return self._request(
             "POST",
             f"/v1/runs/{base_run_id}/chat",
@@ -479,6 +491,8 @@ class ApiClient:
             max_tokens,
             stream=True,
         )
+        if "step" in body:
+            self._require_chat_step_selector()
         headers = {"Content-Type": "application/json", **self._auth_headers()}
         req = urllib.request.Request(
             f"{self.api_url}/v1/runs/{base_run_id}/chat",