freesolo-flash-dev 0.2.32__tar.gz → 0.2.34__tar.gz

freesolo_flash_dev-0.2.34/.github/workflows/notify-tests-repo.yml

@@ -0,0 +1,35 @@
+name: Notify tests repo
+
+# Flash's cross-repo integration seams (flash<->backend auth, allocator, CLI flow)
+# live in the freesolo-co/tests repo, which carries flash as a pinned submodule.
+# On every push to dev/main, dispatch to that repo so its CI bumps the flash
+# submodule pin and runs the affected suites against the new sha. Mirrors the
+# freesolo-side notify-tests-repo.yml; the tests repo listens for `flash-merged`.
+
+on:
+  push:
+    branches: [dev, main]
+  workflow_dispatch:
+
+permissions:
+  contents: read
+
+jobs:
+  notify:
+    runs-on: ubuntu-latest
+    steps:
+      - name: Dispatch to freesolo-co/tests
+        env:
+          PAT: ${{ secrets.FREESOLO_GITHUB_PAT }}
+          SHA: ${{ github.sha }}
+          REF: ${{ github.ref_name }}
+        run: |
+          if [ -z "$PAT" ]; then
+            echo "No FREESOLO_GITHUB_PAT (likely a fork) — skipping dispatch."
+            exit 0
+          fi
+          curl -fsS -X POST \
+            -H "Authorization: Bearer ${PAT}" \
+            -H "Accept: application/vnd.github+json" \
+            https://api.github.com/repos/freesolo-co/tests/dispatches \
+            -d "{\"event_type\":\"flash-merged\",\"client_payload\":{\"sha\":\"${SHA}\",\"ref\":\"${REF}\"}}"

{freesolo_flash_dev-0.2.32 → freesolo_flash_dev-0.2.34}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: freesolo-flash-dev
-Version: 0.2.32
+Version: 0.2.34
 Summary: Flash — managed LoRA post-training (SFT/GRPO) 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

{freesolo_flash_dev-0.2.32 → freesolo_flash_dev-0.2.34}/flash/catalog.py

@@ -38,6 +38,7 @@ class ServingCapacity:
     serve_model_id: str = ""
     max_num_seqs: int = 0
     max_num_batched_tokens: int = 0
+    tensor_parallel_size: int = 0
     gpu_memory_utilization: float = 0.0
 
 
@@ -104,6 +105,7 @@ class ModelInfo:
                 "serve_model_id",
                 "max_num_seqs",
                 "max_num_batched_tokens",
+                "tensor_parallel_size",
                 "gpu_memory_utilization",
             ):
                 if serving.get(key) in ("", 0, 0.0, None):
@@ -113,6 +115,15 @@ class ModelInfo:
 
 DEFAULT_MODEL = "Qwen/Qwen3.5-4B"
 
+SERVING_FP8_MODEL_REPOS: dict[str, str] = {
+    "openbmb/MiniCPM5-1B": "Freesolo-Co/MiniCPM5-1B-FP8",
+    "Qwen/Qwen3.5-0.8B": "Freesolo-Co/Qwen3.5-0.8B-FP8",
+    "Qwen/Qwen3.5-2B": "Freesolo-Co/Qwen3.5-2B-FP8",
+    "Qwen/Qwen3.5-4B": "Freesolo-Co/Qwen3.5-4B-FP8",
+    "Qwen/Qwen3.5-9B": "Freesolo-Co/Qwen3.5-9B-FP8",
+    "Qwen/Qwen3.6-35B-A3B": "Qwen/Qwen3.6-35B-A3B-FP8",
+}
+
 MODELS: dict[str, ModelInfo] = {
     "openbmb/MiniCPM5-1B": ModelInfo(
         id="openbmb/MiniCPM5-1B",
@@ -123,7 +134,13 @@ MODELS: dict[str, ModelInfo] = {
         algos=("sft", "grpo"),
         min_vram_gb=12,
         recommended_gpu="RTX 4090",
-        serving=ServingCapacity(gpu="L4", max_loras=16, max_lora_rank=64, max_model_len=32768),
+        serving=ServingCapacity(
+            gpu="L4",
+            serve_model_id=SERVING_FP8_MODEL_REPOS["openbmb/MiniCPM5-1B"],
+            max_loras=16,
+            max_lora_rank=128,
+            max_model_len=32768,
+        ),
         thinking="hybrid",
         notes="On-device class SLM (131k ctx); standard Llama architecture.",
     ),
@@ -136,7 +153,13 @@ MODELS: dict[str, ModelInfo] = {
         algos=("sft", "grpo"),
         min_vram_gb=12,
         recommended_gpu="RTX 4090",
-        serving=ServingCapacity(gpu="L4", max_loras=16, max_lora_rank=64, max_model_len=32768),
+        serving=ServingCapacity(
+            gpu="L4",
+            serve_model_id=SERVING_FP8_MODEL_REPOS["Qwen/Qwen3.5-0.8B"],
+            max_loras=16,
+            max_lora_rank=128,
+            max_model_len=32768,
+        ),
         thinking="hybrid",
         notes="Smallest Qwen3.5; cheap smoke/dev runs with the modern arch.",
     ),
@@ -149,7 +172,13 @@ MODELS: dict[str, ModelInfo] = {
         algos=("sft", "grpo"),
         min_vram_gb=16,
         recommended_gpu="RTX 4090",
-        serving=ServingCapacity(gpu="L4", max_loras=16, max_lora_rank=64, max_model_len=32768),
+        serving=ServingCapacity(
+            gpu="L4",
+            serve_model_id=SERVING_FP8_MODEL_REPOS["Qwen/Qwen3.5-2B"],
+            max_loras=16,
+            max_lora_rank=128,
+            max_model_len=32768,
+        ),
         thinking="hybrid",
     ),
     "Qwen/Qwen3.5-4B": ModelInfo(
@@ -162,10 +191,10 @@ MODELS: dict[str, ModelInfo] = {
         min_vram_gb=32,
         recommended_gpu="RTX 5090",
         serving=ServingCapacity(
-            gpu="L4",
-            serve_model_id="lovedheart/Qwen3.5-4B-FP8",
+            gpu="L40S",
+            serve_model_id=SERVING_FP8_MODEL_REPOS["Qwen/Qwen3.5-4B"],
             max_loras=64,
-            max_lora_rank=32,
+            max_lora_rank=64,
             max_model_len=8192,
             max_num_seqs=8,
             gpu_memory_utilization=0.98,
@@ -187,10 +216,10 @@ MODELS: dict[str, ModelInfo] = {
         quant="bf16",
         recommended_gpu="A100 PCIe",
         serving=ServingCapacity(
-            gpu="L4",
-            serve_model_id="lovedheart/Qwen3.5-9B-FP8",
+            gpu="L40S",
+            serve_model_id=SERVING_FP8_MODEL_REPOS["Qwen/Qwen3.5-9B"],
             max_loras=44,
-            max_lora_rank=32,
+            max_lora_rank=64,
             max_model_len=8192,
             max_num_seqs=8,
             gpu_memory_utilization=0.98,
@@ -225,13 +254,14 @@ MODELS: dict[str, ModelInfo] = {
         quant="bf16",
         recommended_gpu="H200",
         serving=ServingCapacity(
-            gpu="A100-80GB",
-            serve_model_id="Qwen/Qwen3.6-35B-A3B-FP8",
+            gpu="A100-80GB:2",
+            serve_model_id=SERVING_FP8_MODEL_REPOS["Qwen/Qwen3.6-35B-A3B"],
             max_loras=12,
-            max_lora_rank=32,
+            max_lora_rank=64,
             max_model_len=8192,
             max_num_seqs=8,
             max_num_batched_tokens=4096,
+            tensor_parallel_size=2,
             gpu_memory_utilization=0.98,
         ),
         thinking="hybrid",
@@ -262,8 +292,8 @@ def get_model(model_id: str) -> ModelInfo:
 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.
 
-    Serving capacity is model-specific: small serving models currently allow rank 64, while larger
-    serving paths can cap at rank 32. Unknown/open-policy models intentionally return None instead
+    Serving capacity is model-specific: small serving models currently allow rank 128, while larger
+    serving paths currently allow rank 64. Unknown/open-policy models intentionally return None instead
     of inheriting a global fallback.
     """
     if isinstance(model, ModelInfo):

{freesolo_flash_dev-0.2.32 → freesolo_flash_dev-0.2.34}/flash/cli/__init__.py

@@ -35,6 +35,7 @@ from flash.cli.commands import (  # noqa: F401
     cmd_env_setup,
     cmd_export,
     cmd_gpus,
+    cmd_log,
     cmd_login,
     cmd_models,
     cmd_runs,
@@ -83,7 +84,8 @@ _HELP_GROUPS: list[tuple[str, list[tuple[str, str]]]] = [
         "training",
         [
             ("train", "submit a managed run from a TOML config"),
-            ("status", "show a run's status, logs, or live follow"),
+            ("status", "show a run's current status"),
+            ("log", "print or follow a run's logs"),
             ("runs", "list runs with their state and cost"),
             ("checkpoints", "list a run's deployable RL checkpoints"),
             ("cancel", "cancel a running job"),
@@ -309,21 +311,28 @@ def _build_parser() -> argparse.ArgumentParser:
     )
     train.set_defaults(func=cmd_train)
 
-    status = sub.add_parser("status", help="show a run's status, logs, or follow logs")
+    status = sub.add_parser("status", help="show a run's current status")
     status.add_argument("run_id")
     status.add_argument(
-        "--logs",
+        "-f",
+        "--follow",
         action="store_true",
-        help="print current logs before status — the orchestrator log plus the train-subprocess "
-        "stdout + traceback (console_/error_<phase>.txt) fetched from the run's HF artifact repo",
+        help="poll status until the run ends without replaying logs",
     )
-    status.add_argument(
+    status.set_defaults(func=cmd_status)
+
+    log = sub.add_parser(
+        "log",
+        help="print a run's full logs, including worker console/error artifacts",
+    )
+    log.add_argument("run_id")
+    log.add_argument(
         "-f",
         "--follow",
         action="store_true",
-        help="stream logs until the run ends, then print final status",
+        help="stream logs until the run reaches a terminal state",
     )
-    status.set_defaults(func=cmd_status)
+    log.set_defaults(func=cmd_log)
 
     runs = sub.add_parser("runs", help="list runs and their state/cost")
     runs.set_defaults(func=cmd_runs)

{freesolo_flash_dev-0.2.32 → freesolo_flash_dev-0.2.34}/flash/cli/commands.py

@@ -401,7 +401,7 @@ def cmd_train(args) -> int:
     else:
         print(
             f"run {run_id} submitted; following logs "
-            f"(Ctrl-C detaches, `flash status {run_id} --follow` resumes)",
+            f"(Ctrl-C detaches, `flash log {run_id} --follow` resumes)",
             file=sys.stderr,
         )
     return _follow_run(client, run_id)
@@ -437,28 +437,65 @@ def _follow_run(client: ApiClient, run_id: str) -> int:
     return 0 if state in _OK_STATES else 1
 
 
+def _follow_status(client: ApiClient, run_id: str, interval: float = 2.0) -> int:
+    """Poll run status until terminal, without replaying worker logs."""
+    last_rendered: str | None = None
+    while True:
+        status = client.get_run(run_id)
+        rendered = render.run_status(status) if render.styled() else json.dumps(status, indent=2)
+        if rendered != last_rendered:
+            print(rendered)
+            last_rendered = rendered
+        state = str(status.get("state") or "")
+        if state in _CLI_DONE_STATES:
+            return 0 if state in _OK_STATES else 1
+        time.sleep(interval)
+
+
+def _print_log_text(text: str) -> bool:
+    if not text:
+        return False
+    print(text, end="")
+    if not text.endswith("\n"):
+        print()
+    return True
+
+
+def _print_log_snapshot(client: ApiClient, run_id: str) -> bool:
+    """Print a finite snapshot from the log endpoint without following new bytes."""
+    page = client.get_logs(run_id, offset=0)
+    return _print_log_text(str(page.get("logs") or ""))
+
+
+def _print_worker_output(client: ApiClient, run_id: str, *, printed_any: bool = False) -> bool:
+    for name, text in (client.get_worker_output(run_id) or {}).items():
+        if not text:
+            continue
+        sep = "\n" if printed_any else ""
+        if render.styled():
+            print(f"{sep}{render.log_section(name)}")
+        else:
+            print(f"{sep}----- {name} -----")
+        print(text, end="" if text.endswith("\n") else "\n")
+        printed_any = True
+    return printed_any
+
+
+def cmd_log(args) -> int:
+    client = client_from_config()
+    if getattr(args, "follow", False):
+        state = _poll_logs(client, args.run_id, interval=2.0)
+        _print_worker_output(client, args.run_id, printed_any=True)
+        return 0 if state in _OK_STATES else 1
+    printed_any = _print_log_snapshot(client, args.run_id)
+    _print_worker_output(client, args.run_id, printed_any=printed_any)
+    return 0
+
+
 def cmd_status(args) -> int:
     client = client_from_config()
     if getattr(args, "follow", False):
-        return _follow_run(client, args.run_id)
-    if getattr(args, "logs", False):
-        logs = client.get_logs(args.run_id).get("logs", "")
-        printed_any = False
-        if logs:
-            print(logs, end="")
-            if not logs.endswith("\n"):
-                print()
-            printed_any = True
-        for name, text in (client.get_worker_output(args.run_id) or {}).items():
-            if not text:
-                continue
-            sep = "\n" if printed_any else ""
-            if render.styled():
-                print(f"{sep}{render.log_section(name)}")
-            else:
-                print(f"{sep}----- {name} -----")
-            print(text, end="" if text.endswith("\n") else "\n")
-            printed_any = True
+        return _follow_status(client, args.run_id)
     status = client.get_run(args.run_id)
     if render.styled():
         print(render.run_status(status))

{freesolo_flash_dev-0.2.32 → freesolo_flash_dev-0.2.34}/flash/cli/envpush.py

@@ -190,8 +190,23 @@ _ENV_PUSH_SIDECAR_SUFFIXES = frozenset(
 def _normalize_env_name(raw: str) -> str | None:
     import re
 
-    name = re.sub(r"[^a-z0-9]+", "-", raw.lower()).strip("-")
-    return name or None
+    def normalize_segment(value: str) -> str | None:
+        segment = re.sub(r"[^a-z0-9._-]+", "-", value.lower()).strip("-")
+        if segment in {"", ".", ".."} or not re.search(r"[a-z0-9]", segment):
+            return None
+        return segment
+
+    text = str(raw or "").strip()
+    if "/" not in text:
+        return normalize_segment(text)
+    parts = [part.strip() for part in text.split("/")]
+    if len(parts) != 2 or not all(parts):
+        return None
+    namespace = normalize_segment(parts[0])
+    name = normalize_segment(parts[1])
+    if not namespace or not name:
+        return None
+    return f"{namespace}/{name}"
 
 
 def _with_syspath_bootstrap(env_source: str) -> str:

{freesolo_flash_dev-0.2.32 → freesolo_flash_dev-0.2.34}/flash/cli/render.py

@@ -381,7 +381,7 @@ def version(value: str) -> str:
 def submitted(run_id: str) -> str:
     """The `flash train` hand-off note (printed to stderr before logs start streaming)."""
     head = ok(f"run {_paint(run_id, _ACCENT2)} submitted")
-    hint = _dim(f"following logs — Ctrl-C detaches; resume with `flash status {run_id} --follow`")
+    hint = _dim(f"following logs — Ctrl-C detaches; resume with `flash log {run_id} --follow`")
     return _safe(f"{head}\n{hint}")
 
 
@@ -659,7 +659,7 @@ def chat_label() -> str:
 
 
 def log_section(name: str) -> str:
-    """A themed divider above a passthrough worker-log section in `flash status --logs` — the same
+    """A themed divider above a passthrough worker-log section in `flash log` — the same
     idiom as chat_label sitting above a raw chat reply (the log body stays raw). The machine path
     keeps the plain ``----- name -----`` divider that scripts and tests match on."""
     rule = _paint(_glyph("─", "-") * 3, _FAINT)

{freesolo_flash_dev-0.2.32 → freesolo_flash_dev-0.2.34}/flash/cli/training_doc.py

@@ -127,8 +127,9 @@ flash train configs/rl.toml --background  # submit and return immediately
 
 ```bash
 flash status <run-id>            # state + accrued cost
-flash status <run-id> --logs     # reward/loss trend + worker console/error logs + any traceback
-flash status <run-id> --follow   # stream a live run to completion
+flash log <run-id>               # reward/loss trend + worker console/error logs + any traceback
+flash log <run-id> --follow      # stream a live run to completion
+flash status <run-id>            # current run state, cost, and deployment info
 flash runs                       # all your runs and their state/cost
 flash cancel <run-id>            # stop a run
 ```
@@ -205,8 +206,8 @@ spending another GPU run:
 | All-zero or flat GRPO reward | `reward_mean` stays near 0 and outputs do not improve | Make the reward dense: give partial credit for parse/format/execution/correctness tiers, and log a separate clean `success` metric. Do not keep rerunning an all-zero reward. |
 | Reward rises but behavior is worse | Short, templated, malformed, or reward-hacked outputs score well | Deploy the adapter and probe real examples. Add hard validity gates before judge calls, penalize degenerate shortcuts, and judge the outcome rather than the surface string. |
 | Output is truncated | Correct-looking answers cut off mid-response or JSON is incomplete | Increase `max_tokens` for GRPO rollouts or `max_length` for SFT only after seeing truncation. Oversizing them by default just burns memory/cost. |
-| Infrastructure, CUDA, OOM, vLLM, or kernel failure | Run errors before useful metrics, often during setup/model load | Treat this as infrastructure pressure, not proof the model is too large. Read `flash status <run-id> --logs`, reduce footprint (`max_length`, `max_tokens`, `group_size`) if needed, and let Flash retry/allocate another fitting GPU class. |
-| Run looks stuck after disconnecting | Terminal stopped streaming but the job may still be alive | Ctrl-C detaches. Use `flash status <run-id> --follow` to reattach, `--logs` for the console/error tail, or `flash cancel <run-id>` if you intentionally want to stop it. |
+| Infrastructure, CUDA, OOM, vLLM, or kernel failure | Run errors before useful metrics, often during setup/model load | Treat this as infrastructure pressure, not proof the model is too large. Read `flash log <run-id>`, reduce footprint (`max_length`, `max_tokens`, `group_size`) if needed, and let Flash retry/allocate another fitting GPU class. |
+| Run looks stuck after disconnecting | Terminal stopped streaming but the job may still be alive | Ctrl-C detaches. Use `flash log <run-id> --follow` to reattach, `flash log <run-id>` for the console/error output, or `flash cancel <run-id>` if you intentionally want to stop it. |
 | Final checkpoint regresses | Last step is worse than an earlier checkpoint | Run `flash checkpoints <run-id>`, deploy a specific step with `flash deploy <run-id> --step N`, and compare with held-out probes before exporting or relying on the final adapter. |
 | Export fails before upload | CLI says no HuggingFace token | Pass `flash export --api-key hf_...`, or set `HF_TOKEN` in your shell, `.env`, or `.env.local`. Exports are private unless you pass `--public`. |
 | SFT loss improves but quality does not | Train loss falls while held-out behavior stalls or degrades | Keep a held-out split outside training. Deploy and score that split; if quality drops, reduce epochs or improve data instead of adding more passes. |
@@ -222,15 +223,15 @@ spending another GPU run:
 - **Read the model's outputs, not just the metrics.** A rising reward can come from
   reward-hacking or a degenerate output the reward still credits — metrics alone never
   establish that the model got better. Flash does not expose training-time rollouts
-  through the CLI (`--logs` gives you the metric trend and the worker's console/error
+  through the CLI (`flash log` gives you the metric trend and the worker's console/error
   logs, not the sampled generations), so to read real outputs **deploy the adapter and
   probe it**: `flash deploy <run-id>` then `flash chat <run-id> -m "..."` on at least a
   few real inputs, including ones it should get wrong.
 
   ```bash
   flash status <run-id>            # state + accrued cost
-  flash status <run-id> --logs     # metric trend + worker console/error logs (+ traceback)
-  flash status <run-id> --follow   # stream a live run until completion
+  flash log <run-id>               # metric trend + worker console/error logs (+ traceback)
+  flash log <run-id> --follow      # stream a live run until completion
   flash deploy <run-id>            # serve the adapter, then `flash chat` it to read real outputs
   ```
 
@@ -472,8 +473,8 @@ flash train configs/rl.toml --dry-run # validate the config locally (no GPU, no
 flash train configs/rl.toml --cost    # pre-flight USD estimate, then exit
 flash train configs/rl.toml           # submit and follow logs (Ctrl-C detaches; --background to skip following)
 flash status <run-id>                 # state + accrued cost
-flash status <run-id> --logs          # reward/loss trend + worker console/error logs
-flash status <run-id> --follow        # stream a live run to completion
+flash log <run-id>                    # reward/loss trend + worker console/error logs
+flash log <run-id> --follow           # stream a live run to completion
 flash runs                            # list your runs and their state/cost
 flash cancel <run-id>                 # stop a live run
 flash checkpoints <run-id>            # list deployable RL checkpoints

{freesolo_flash_dev-0.2.32 → freesolo_flash_dev-0.2.34}/flash/client/http.py

@@ -110,6 +110,14 @@ def _validate_checkpoint_step(step: int | float) -> int:
     return int(step)
 
 
+def _validate_chat_messages(messages: list[dict]) -> None:
+    if not isinstance(messages, list):
+        raise ClientError("chat messages must be a list")
+    for index, message in enumerate(messages):
+        if not isinstance(message, dict):
+            raise ClientError(f"chat messages[{index}] must be an object")
+
+
 class ApiClient:
     def __init__(
         self,
@@ -225,14 +233,7 @@ class ApiClient:
         return self._request("GET", f"/v1/runs/{run_id}/logs?offset={int(offset)}")
 
     def get_worker_output(self, run_id: str) -> dict[str, str]:
-        # Tolerate servers that predate /worker: FastAPI returns bare 404 "Not Found" for unknown
-        # paths; real run-not-found 404s carry "unknown run_id: ..." so we can distinguish them.
-        try:
-            return self._request("GET", f"/v1/runs/{run_id}/worker").get("worker", {})
-        except ApiError as exc:
-            if exc.status == 404 and str(exc).strip().lower() == "not found":
-                return {}
-            raise
+        return self._request("GET", f"/v1/runs/{run_id}/worker").get("worker", {})
 
     def cancel_run(self, run_id: str) -> dict:
         return self._request("POST", f"/v1/runs/{run_id}/cancel")
@@ -289,6 +290,7 @@ class ApiClient:
         temperature: float = 0.0,
         max_tokens: int = 512,
     ) -> dict:
+        _validate_chat_messages(messages)
         return self._request(
             "POST",
             f"/v1/runs/{run_id}/chat",
@@ -303,6 +305,7 @@ class ApiClient:
         temperature: float = 0.0,
         max_tokens: int = 512,
     ) -> Iterator[str]:
+        _validate_chat_messages(messages)
         headers = {"Content-Type": "application/json"}
         if self.api_key:
             headers["Authorization"] = f"Bearer {self.api_key}"

{freesolo_flash_dev-0.2.32 → freesolo_flash_dev-0.2.34}/flash/cost/facts.py

@@ -7,6 +7,8 @@ from flash.catalog import MODELS
 from flash.providers.base import GPU_INFO, GpuClass, providers_for
 
 GPU_COMPUTE_TFLOPS: dict[str, float] = {
+    # A10: 125 TFLOPS dense bf16 tensor (NVIDIA spec); Lambda-only 24 GB class, else defaults to 100.
+    "A10": 125.0,
     "RTX 4090": 165.0,
     "RTX 5090": 210.0,
     "A100 PCIe": 312.0,

{freesolo_flash_dev-0.2.32 → freesolo_flash_dev-0.2.34}/flash/envs/adapter.py

@@ -1032,10 +1032,11 @@ class FreesoloEnvironment(BaseEnvironment):
             episode_turns: int | None = int(self._env.max_episode_turns(task))
         except Exception:
             episode_turns = None
+        messages = [dict(message) for message in prompt]
         return {
             "task": task,
-            "prompt": [dict(message) for message in prompt],
-            "messages": [dict(message) for message in prompt],
+            "prompt": prompt,
+            "messages": messages,
             "turns": [],
             "done": False,
             "response_text": "",

{freesolo_flash_dev-0.2.32 → freesolo_flash_dev-0.2.34}/flash/runner/__init__.py

@@ -59,6 +59,66 @@ def _gpu_rate(gpu_type: str) -> float:
         return 0.80
 
 
+def charge_usd_for_spec(spec, *, steps: int | None = None, fallback: float = 0.0) -> float:
+    """The customer charge for a run: the flash.cost estimate (training-only steps x sec/step x $/hr).
+
+    This is the price the run was QUOTED at submit. ``steps=None`` prices the spec's planned steps
+    (a completed run is charged exactly its quote); pass the actual steps that ran to re-price a
+    CANCELLED run at how far it got. Returns ``fallback`` if the spec can't be priced (so a charge is
+    never blocked by a pricing failure)."""
+    try:
+        from flash.cost.analytical import estimate_cost
+        from flash.cost.spec import estimate_for_spec, runconfig_from_spec
+
+        if steps is None:
+            return float(estimate_for_spec(spec).total_usd)
+        n = max(0, int(steps))
+        if n == 0:
+            return 0.0  # cancelled before any training step -> nothing to charge
+        from dataclasses import replace
+
+        cfg = runconfig_from_spec(spec)
+        planned = int(cfg.steps or 0)
+        if planned > 0:
+            n = min(n, planned)
+        # SFT is priced from train_tokens (not steps), so lowering steps ALONE wouldn't prorate a
+        # cancel -- estimate_cost would still charge the full-run token estimate. Scale the token
+        # count to the fraction of steps that ran so a mid-training SFT cancel is charged its share,
+        # mirroring GRPO's steps-based proration.
+        if not cfg.is_grpo and cfg.train_tokens and planned > 0:
+            scaled_tokens = max(1, int(cfg.train_tokens * n / planned))
+            cfg = replace(cfg, steps=n, train_tokens=scaled_tokens)
+        else:
+            cfg = replace(cfg, steps=n)
+        return float(estimate_cost(cfg).total_usd)
+    except Exception:
+        return float(fallback)
+
+
+# Heartbeat stages that mean the worker has entered training (GPU work underway). The per-step
+# `step` field is 1-indexed and only appears once a step COMPLETES, so the expensive first step (a
+# GRPO rollout can be ~17 min) streams one of these stages with NO step yet -- still real GPU time.
+_TRAINING_STAGES = frozenset({"rl_step", "sft_step"})
+
+
+def actual_steps_run(status: RunStatus) -> int:
+    """How many optimizer steps to bill a (cancelled) run for.
+
+    The worker streams a per-step heartbeat whose ``step`` field is the last COMPLETED optimizer step
+    (1-indexed; the last one we persisted is the furthest it reached). Cancelled after N steps -> N.
+    The first step reports no ``step`` until it completes, so a cancel mid-first-step would look like
+    0 steps despite real GPU time -- we floor to 1 whenever a training-stage heartbeat is present.
+    Returns 0 only when no training heartbeat was seen (cancelled during cold-start/setup) -> $0."""
+    hb = status.last_heartbeat if isinstance(status.last_heartbeat, dict) else {}
+    step = hb.get("step")
+    if isinstance(step, (int, float)) and step > 0:
+        return int(step)
+    # Training started (rl_step/sft_step) but no completed step yet -> mid-first-step -> bill 1.
+    if hb.get("stage") in _TRAINING_STAGES:
+        return 1
+    return 0
+
+
 @dataclass
 class RunStatus:
     run_id: str
@@ -74,9 +134,9 @@ class RunStatus:
     remote: dict | None = None
     # Realized provider cost (COGS), pulled from the provider's billing API after the run
     # finishes by the reconciliation job (flash/server/reconcile.py) and reported to the
-    # freesolo backend for estimator accuracy. Distinct from ``cost_usd`` (the wall x $/hr
-    # PROJECTION); ``reconciled_at`` marks that the realized pull has happened so it isn't
-    # re-pulled. Both stay None for un-reconciled / pre-instrumentation runs.
+    # freesolo backend for estimator accuracy. Distinct from ``cost_usd`` (the flash.cost ESTIMATE
+    # we charge the customer); ``reconciled_at`` marks that the realized pull has happened so it
+    # isn't re-pulled. Both stay None for un-reconciled / pre-instrumentation runs.
     realized_cost_usd: float | None = None
     reconciled_at: float | None = None
     # Stamped ONCE on first terminal transition; survives later updated_at bumps from deploy/reconcile.
@@ -234,7 +294,9 @@ _WEIGHT_CACHE_PEAK_FACTOR = 2.0
 def _fits_weight_cache(info: ModelInfo) -> bool:
     """Whether the model's peak download footprint fits the shared weight-cache volume."""
     if not info.params_b:
-        return True  # unknown size -> keep the (attach) default; curated catalog models always set it
+        return (
+            True  # unknown size -> keep the (attach) default; curated catalog models always set it
+        )
     download_gb = info.params_b * 2.0  # bf16: 2 bytes/param (mirrors cost.facts.download_weight_gb)
     return _WEIGHT_CACHE_PEAK_FACTOR * download_gb <= WEIGHT_CACHE_VOLUME_GB