freesolo-flash-dev 1.0.100__tar.gz → 1.1.2__tar.gz

{freesolo_flash_dev-1.0.100 → freesolo_flash_dev-1.1.2}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: freesolo-flash-dev
-Version: 1.0.100
+Version: 1.1.2
 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

{freesolo_flash_dev-1.0.100 → freesolo_flash_dev-1.1.2}/flash/engine/worker/sft_train.py

@@ -549,6 +549,9 @@ class _VerlCheckpointWatcher:
     ) -> None:
         self.local_dir = local_dir
         self.export_root = export_root
+        # sibling of export_root so it shares its filesystem: _staged_source hardlinks into it, and
+        # os.link cannot cross a mount. holds one step at a time and is cleaned up after each publish.
+        self.staging_root = os.path.join(export_root, "_staging")
         self.python_bin = python_bin
         self.model_id = model_id
         self.model_revision = model_revision
@@ -605,10 +608,58 @@ class _VerlCheckpointWatcher:
     def _should_publish(self, step: int) -> bool:
         return not self.required_steps or step in self.required_steps
 
+    def _staged_source(self, step: int, checkpoint_dir: str) -> str:
+        """Pin a completed checkpoint against verl's retention before the slow work reads it.
+
+        This watcher publishes on its own thread, so the merge and the multi-GB upload below both
+        overlap training. verl owns `local_dir` and prunes it from the CHILD process: under
+        `trainer.max_ckpt_to_keep=1`, `register_checkpoint` calls `shutil.rmtree` on global_step_N as
+        soon as N+1 finishes saving. Nothing the parent sets can pin a directory across that, so a
+        required save still exporting or uploading when the next save lands reads a tree being
+        deleted underneath it, and an otherwise successful paid run fails its required save. Close
+        `save_at_steps` (say `[2, 3]`) make that likely rather than theoretical.
+
+        HARDLINKS, not a copy and not a rename:
+        - a copy of tens of GB is itself slow enough to lose the same race it is meant to close.
+        - a rename would win the race but break in-pod resume: verl resolves
+          `latest_checkpointed_iteration.txt` to `global_step_N` (`find_latest_ckpt_path`), so moving
+          the directory leaves the tracker pointing at nothing if the pod is preempted afterwards.
+
+        Linking is O(1) per file and leaves the original in place, so verl's retention and resume see
+        exactly what they saw before while the link keeps the DATA alive for as long as this watcher
+        needs it. A pruned original just drops its own reference; the inode survives.
+
+        Best-effort by design: any failure returns the original path, which is the previous
+        (race-exposed) behaviour. Degrading beats failing a run that would otherwise complete.
+        """
+        staged = os.path.join(self.staging_root, f"global_step_{step}")
+        try:
+            shutil.rmtree(staged, ignore_errors=True)
+            shutil.copytree(checkpoint_dir, staged, copy_function=os.link)
+        except OSError as error:
+            print(
+                f"[ckpt] step {step} could not be staged out of verl's retention tree "
+                f"({error}); publishing from verl's copy",
+                flush=True,
+            )
+            shutil.rmtree(staged, ignore_errors=True)
+            return checkpoint_dir
+        return staged
+
     def _publish(self, step: int, checkpoint_dir: str) -> None:
         if not self._should_publish(step):
             self.processed_steps.add(step)
             return
+        staged_dir = self._staged_source(step, checkpoint_dir)
+        try:
+            self._publish_from(step, staged_dir)
+        finally:
+            if staged_dir != checkpoint_dir:
+                # the links have served their purpose once the upload returns; the originals and the
+                # exported adapter are untouched by this.
+                shutil.rmtree(staged_dir, ignore_errors=True)
+
+    def _publish_from(self, step: int, checkpoint_dir: str) -> None:
         actor_dir = resolve_checkpoint_actor_dir(checkpoint_dir)
         adapter_dir = os.path.join(self.export_root, f"step-{step}")
         _export_checkpoint_adapter(

{freesolo_flash_dev-1.0.100 → freesolo_flash_dev-1.1.2}/flash/runner/__init__.py

@@ -2080,18 +2080,24 @@ def _persist_metrics(spec: JobSpec, metrics: dict) -> float:
     # case _gpu_rate prices off whichever configured provider offers the class.
     provider = str(metrics.get("allocated_provider") or "")
     rate = _gpu_rate(gpu_type, provider)
+    # `hourly_rate` is per CARD, so a sharded run costs the wall times the rate times the number of
+    # cards it actually occupied. `allocated_gpu_count` is worker/lifecycle-stamped for the same
+    # reason `allocated_gpu` is: the spec's gpu.count is only a ceiling and allocation may pick
+    # fewer. Absent on records predating the stamp, where one card is the correct reading.
+    gpu_count = max(1, int(metrics.get("allocated_gpu_count") or 1))
     cost = metrics.get("cost_usd")
     if cost:
         cost = float(cost or 0.0)
     else:
         wall = float(metrics.get("wall_seconds") or 0.0)
-        cost = wall / 3600.0 * rate
+        cost = wall / 3600.0 * rate * gpu_count
         metrics = {**metrics, "cost_usd": cost}
         metrics.setdefault("notes", {})
         if isinstance(metrics["notes"], dict):
             metrics["notes"]["provider"] = provider or "unknown"
             metrics["notes"]["gpu_rate_usd_hr"] = rate
             metrics["notes"]["gpu"] = gpu_type
+            metrics["notes"]["gpu_count"] = gpu_count
     with open(os.path.join(dest, "metrics.json"), "w") as f:
         json.dump(metrics, f, indent=2)
     with contextlib.suppress(Exception):

{freesolo_flash_dev-1.0.100 → freesolo_flash_dev-1.1.2}/flash/runner/deploy.py

@@ -1110,6 +1110,7 @@ def attach_run(run_id: str, log_stream=None) -> RunStatus:
             raise ValueError("persisted attempt identity is missing or invalid")
         next_attempt = recovered_attempt + 1
         allocated_gpu = remote.pop("allocated_gpu", None)
+        allocated_gpu_count = remote.pop("allocated_gpu_count", None)
         handle = JobHandle.from_dict(remote)
         try:
             poll_spec = _spec_with_remaining_wall(worker_spec, require_provider_minimum=False)
@@ -1256,6 +1257,8 @@ def attach_run(run_id: str, log_stream=None) -> RunStatus:
             return status_for_return()
         if allocated_gpu and isinstance(res.metrics, dict):
             res.metrics.setdefault("allocated_gpu", allocated_gpu)
+        if allocated_gpu_count and isinstance(res.metrics, dict):
+            res.metrics.setdefault("allocated_gpu_count", int(allocated_gpu_count))
         if not _adopt_completed_attempt(
             run_id,
             worker_spec,

{freesolo_flash_dev-1.0.100 → freesolo_flash_dev-1.1.2}/flash/runner/lifecycle.py

@@ -316,6 +316,9 @@ def _runpod_completed_metrics(handle, *, deadline_at: float | None = None) -> di
         allocated_gpu = original.get("allocated_gpu")
         if allocated_gpu:
             metrics.setdefault("allocated_gpu", allocated_gpu)
+        allocated_count = original.get("allocated_gpu_count")
+        if allocated_count:
+            metrics.setdefault("allocated_gpu_count", int(allocated_count))
         return metrics
     except _CompletedAttemptPending:
         raise
@@ -684,6 +687,11 @@ def _submit_seed_supervised(
                 **canonical_handle,
                 "seed": int(seed),
                 "allocated_gpu": current_gpu.get("name"),
+                # carried beside the gpu name for the same reason and by the same route: the
+                # canonical provider handle drops unknown keys, so a recovering process can only
+                # learn the shape from what was persisted here. without the count a run adopted
+                # after a control-plane restart prices its wall as a single card.
+                "allocated_gpu_count": current_gpu.get("count"),
                 "on_last_gpu": bool(current_on_last_gpu["value"]),
                 "code_prefix": code_prefix,
             }
@@ -742,6 +750,11 @@ def _submit_seed_supervised(
             metrics.setdefault("allocated_gpu", current_gpu["name"])
         if current_gpu.get("provider"):
             metrics.setdefault("allocated_provider", current_gpu["provider"])
+        # the runpod serverless route returns here rather than through the `res.ok` stamp below, so
+        # the card count has to be recorded on both or a sharded serverless run is still priced as
+        # one card. same source either way: the candidate allocation actually chose.
+        if current_gpu.get("count"):
+            metrics.setdefault("allocated_gpu_count", int(current_gpu["count"]))
         return metrics
 
     max_retries = int(spec.gpu.max_retries)
@@ -984,6 +997,7 @@ def _submit_seed_supervised(
                 raise
             current_gpu["name"] = chosen.gpu
             current_gpu["provider"] = chosen.provider
+            current_gpu["count"] = int(getattr(chosen, "gpu_count", 1) or 1)
             current_attempt["value"] = attempt
             retry_delay = 0
             submission_lock = _deploy_lock(spec.run_id)
@@ -1113,6 +1127,11 @@ def _submit_seed_supervised(
                 # the provider that actually billed this run, so cost attribution prices the
                 # class on ITS substrate rather than assuming RunPod's table.
                 res.metrics.setdefault("allocated_provider", chosen.provider)
+                # and how many cards of it. `hourly_rate(gpu_type)` is per CARD, so a fallback that
+                # prices the wall once records a 2x/4x run at half or a quarter of its real spend.
+                # the spec's own gpu.count cannot stand in: it is a ceiling, and allocation
+                # routinely picks fewer (see TRAINING.md, "a ceiling, not an exact count").
+                res.metrics.setdefault("allocated_gpu_count", int(getattr(chosen, "gpu_count", 1)))
             return res.metrics
         # cancel wins over any retry-shaped failure.
         try:

{freesolo_flash_dev-1.0.100 → freesolo_flash_dev-1.1.2}/pyproject.toml

@@ -1,7 +1,7 @@
 [project]
 name = "freesolo-flash-dev"
-version = "1.0.100"
-description = "Flash — managed LoRA post-training (SFT/GRPO/OPD) for Freesolo environments, driven by the `flash` CLI"
+version = "1.1.2"
+description ="Flash — managed LoRA post-training (SFT/GRPO/OPD) for Freesolo environments, driven by the `flash` CLI"
 readme = "README.md"
 # RunPod Flash supports Python 3.11-3.12 (not 3.13 yet).
 requires-python = ">=3.11,<3.13"
@@ -121,7 +121,7 @@ artifacts = ["flash/py.typed"]
 # auto-published by .github/workflows/publish-dev.yml: a push to `dev` whose version isn't yet on
 # PyPI publishes it; pushes that don't change it find it already published and no-op. So cutting a
 # dev-channel release = bump this (and [project].version) and merge to `dev`.
-version = "1.0.100"
+version = "1.1.2"
 
 [tool.pytest.ini_options]
 pythonpath = ["."]

{freesolo_flash_dev-1.0.100 → freesolo_flash_dev-1.1.2}/tests/test_opd_resume_safety.py

@@ -656,6 +656,8 @@ def test_opd_automatic_retry_after_teardown_requires_all_markers_absent(monkeypa
         # stamped alongside the gpu so cost attribution prices the class on the substrate
         # that actually billed it.
         "allocated_provider": "runpod",
+        # and the card count, since the hourly rate is per card.
+        "allocated_gpu_count": 1,
     }
     assert provider.attempts == [0, 1]
     retry_submit = provider.events.index(("submit", 1))
@@ -744,6 +746,8 @@ def test_opd_retry_passes_gate_revision_and_overwrites_spoofed_value(monkeypatch
         # stamped alongside the gpu so cost attribution prices the class on the substrate
         # that actually billed it.
         "allocated_provider": "runpod",
+        # and the card count, since the hourly rate is per card.
+        "allocated_gpu_count": 1,
     }
     broker_transport = {
         "FLASH_CONTROL_PANEL_URL": "https://broker.example",

{freesolo_flash_dev-1.0.100 → freesolo_flash_dev-1.1.2}/tests/test_provider_routing.py

@@ -796,6 +796,57 @@ def test_runpod_cost_projection_flows_into_run_status(orch, monkeypatch):
     assert cost == pytest.approx(0.345)  # 0.5 hr x $0.69/hr (RTX 4090)
 
 
+def test_a_multi_card_run_is_costed_for_every_card_it_occupied(orch):
+    """`hourly_rate` is per CARD, so the measured cost must multiply by the allocated count.
+
+    without it a 4-card run records a quarter of its real spend, and the cost analytics that compare
+    estimates against actuals are corrupted by exactly that factor. the count has to come from the
+    metrics stamp rather than `spec.gpu.count`, which is only a ceiling -- allocation routinely picks
+    fewer cards than authored.
+    """
+    spec = _spec()
+    _seed_status(orch, spec)
+    base = {"train_tokens": 4096, "wall_seconds": 1800, "allocated_gpu": "RTX 4090"}
+
+    four_card = orch._persist_metrics(spec, {**base, "allocated_gpu_count": 4})
+
+    assert four_card == pytest.approx(0.345 * 4), "a 4-card run was priced as one card"
+    # a record predating the stamp still reads as one card rather than zero or a crash
+    assert orch._persist_metrics(spec, dict(base)) == pytest.approx(0.345)
+
+
+def test_the_allocated_card_count_reaches_the_metrics_the_cost_is_read_from(orch, monkeypatch):
+    """The multiply above is inert unless allocation actually stamps the count it chose.
+
+    this drives the real submit path rather than handing `_persist_metrics` a literal, because the
+    two halves fail independently: pricing can multiply correctly by a count that is never recorded.
+    """
+    from flash.providers import allocator
+    from flash.providers.base import Candidate, PollResult
+    from flash.providers.runpod import jobs as rp_jobs
+
+    monkeypatch.setattr(
+        allocator,
+        "allocate",
+        lambda *a, **k: _alloc(candidates=(Candidate("runpod", "RTX 4090", 0.69, 24, 4),)),
+    )
+
+    def fake_runpod_submit(run_spec, seed, log=None, on_handle=None, attempt=0, **_):
+        if on_handle:
+            on_handle(_runpod_handle(attempt=attempt))
+        return PollResult(True, metrics={"train_tokens": 4096, "wall_seconds": 1800})
+
+    monkeypatch.setattr(rp_jobs, "submit_run", fake_runpod_submit)
+    spec = _spec(provider="runpod", type="RTX 4090", count=4)
+    _seed_status(orch, spec)
+
+    metrics = orch._submit_seed_supervised(spec, spec.seed, io.StringIO())
+
+    assert metrics["allocated_gpu_count"] == 4
+    # and the stamp is what pricing then reads, so the two halves compose
+    assert orch._persist_metrics(spec, metrics) == pytest.approx(0.345 * 4)
+
+
 def test_infra_retry_walks_to_next_runpod_class_and_deletes_endpoint(orch, monkeypatch):
     from flash.providers import allocator
     from flash.providers.base import Candidate, PollResult

{freesolo_flash_dev-1.0.100 → freesolo_flash_dev-1.1.2}/tests/test_sft_train.py

@@ -5,7 +5,9 @@ from __future__ import annotations
 import contextlib
 import importlib.util
 import os
+import pathlib
 import re
+import shutil
 import sys
 import types
 from dataclasses import replace
@@ -959,10 +961,13 @@ def test_checkpoint_watcher_exports_and_uploads_required_step(monkeypatch, tmp_p
 
     watcher._publish(5, str(checkpoint_dir))
 
-    assert exported[0][0] == str(actor_dir)
+    # both read the STAGED copy rather than verl's directory (see _staged_source); what matters is
+    # that they agree on one source and that it carries the `actor/` level this layout puts there.
+    assert os.path.basename(exported[0][0]) == "actor"
+    assert exported[0][0] == os.path.join(uploaded[0][1], "actor")
     assert published[0][1] == 5
     assert published[0][2]["required"] is True
-    assert uploaded == [(5, str(checkpoint_dir))]
+    assert [step for step, _ in uploaded] == [5]
     assert watcher.processed_steps == {5}
 
 
@@ -1002,7 +1007,72 @@ def test_checkpoint_watcher_exports_the_sft_layout(monkeypatch, tmp_path):
 
     watcher._publish(5, str(checkpoint_dir))
 
-    assert exported == [str(checkpoint_dir)]
+    # the export reads a STAGED copy, not verl's own directory (see _staged_source), so assert the
+    # layout this test is about -- global_step_N itself, no `actor/` level -- rather than the path.
+    assert [os.path.basename(path) for path in exported] == ["global_step_5"]
+
+
+def test_a_required_save_survives_verl_pruning_it_mid_publish(monkeypatch, tmp_path):
+    """verl deleting global_step_N while it is being published must not fail the run.
+
+    the watcher publishes on its own thread, so the merge and the multi-GB upload overlap training.
+    under `trainer.max_ckpt_to_keep=1` verl's `register_checkpoint` rmtree's global_step_N from the
+    CHILD process the moment N+1 finishes saving, and the parent has no way to pin it. with close
+    `save_at_steps` (here [2, 3]) that lands squarely inside the publish window.
+
+    this reproduces the interleave rather than asserting on the staging mechanism: the fake upload
+    deletes verl's directory the way verl would, then reads the source it was handed. before staging,
+    that read fails and takes the whole paid run down with a required-save error.
+    """
+    import flash.engine.worker as worker
+    from flash.engine.worker import sft_train
+
+    local_dir = tmp_path / "checkpoints"
+    checkpoint_dir = local_dir / "global_step_2"
+    (checkpoint_dir / "actor" / "huggingface").mkdir(parents=True)
+    (checkpoint_dir / "actor" / "model.safetensors").write_bytes(b"weights")
+
+    monkeypatch.setattr(
+        sft_train,
+        "_export_checkpoint_adapter",
+        lambda actor, adapter, **kwargs: os.makedirs(adapter, exist_ok=True),
+    )
+    monkeypatch.setattr(
+        worker, "publish_deployable_checkpoint", lambda adapter, step, **kwargs: None
+    )
+    read_back = {}
+
+    def fake_upload(step, checkpoint, **kwargs):
+        kwargs["before_upload"]()
+        # verl prunes global_step_2 because global_step_3 just landed. this is the whole race.
+        shutil.rmtree(checkpoint_dir, ignore_errors=True)
+        # the upload streams the folder AFTER that, so it must still have something to read.
+        read_back["payload"] = (
+            pathlib.Path(checkpoint, "actor", "model.safetensors").read_bytes()
+            if os.path.isdir(checkpoint)
+            else None
+        )
+        return True
+
+    monkeypatch.setattr(worker, "upload_resume_checkpoint", fake_upload)
+    watcher = sft_train._VerlCheckpointWatcher(
+        local_dir=str(local_dir),
+        export_root=str(tmp_path / "exports"),
+        python_bin="/verl/python",
+        model_id="org/model",
+        model_revision="commit",
+        required_steps=(2, 3),
+    )
+
+    watcher._publish(2, str(checkpoint_dir))
+
+    assert read_back["payload"] == b"weights", (
+        "the upload lost its source when verl pruned the checkpoint, so a required save would fail "
+        "an otherwise successful paid run"
+    )
+    assert watcher.processed_steps == {2}
+    # the staging links are transient: they must not outlive the publish and accumulate on the pod.
+    assert not os.path.exists(os.path.join(str(tmp_path / "exports"), "_staging", "global_step_2"))
 
 
 def test_resume_credits_only_required_saves_that_are_durable(monkeypatch):

{freesolo_flash_dev-1.0.100 → freesolo_flash_dev-1.1.2}/uv.lock

@@ -1025,7 +1025,7 @@ wheels = [
 
 [[package]]
 name = "freesolo-flash"
-version = "1.0.100"
+version = "1.1.2"
 source = { editable = "." }
 
 [package.optional-dependencies]