freesolo-flash-dev 0.2.43__tar.gz → 0.2.44__tar.gz

{freesolo_flash_dev-0.2.43 → freesolo_flash_dev-0.2.44}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: freesolo-flash-dev
-Version: 0.2.43
+Version: 0.2.44
 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.43 → freesolo_flash_dev-0.2.44}/flash/cli/commands.py

@@ -783,12 +783,32 @@ def cmd_runs(args) -> int:
 
 
 def cmd_cancel(args) -> int:
-    status = client_from_config().cancel_run(args.run_id)
+    client = client_from_config()
+    status = client.cancel_run(args.run_id)
     payload = {"run_id": args.run_id, "state": status["state"]}
+    # A cancelled run is not necessarily worthless: every completed save interval already streamed
+    # a deployable checkpoint, even though the run shows adapter_ref=null / cost=0. Surface the
+    # surviving steps here so the run isn't discarded unseen. Best-effort: cancel never fails on it.
+    checkpoints: list[dict] = []
+    if payload["state"] == "cancelled":
+        try:
+            checkpoints = client.checkpoints(args.run_id)
+        except Exception:
+            checkpoints = []
     if render.styled():
         print(render.cancelled(payload))
     else:
         print(json.dumps(payload, indent=2))
+    if checkpoints:
+        top = max(c["step"] for c in checkpoints)
+        # stderr in the plain path so the machine-readable stdout JSON stays untouched.
+        out = sys.stdout if render.styled() else sys.stderr
+        msg = (
+            f"{len(checkpoints)} deployable checkpoint(s) survive this cancel — list with "
+            f"`flash checkpoints {args.run_id}`, deploy one with "
+            f"`flash deploy {args.run_id}/step-{top}`."
+        )
+        print(render.note(msg) if render.styled() else msg, file=out)
     return 0
 
 

{freesolo_flash_dev-0.2.43 → freesolo_flash_dev-0.2.44}/flash/engine/worker/hf.py

@@ -310,11 +310,13 @@ def publish_deployable_checkpoint(
     return None
 
 
-# Daemon upload threads are killed on worker exit; on_train_end blocks up to this long to flush.
-_CKPT_FLUSH_TIMEOUT_S = 300.0
-# Retry/backoff for the lock-timeout fallback publish (rides out a transient concurrent-commit 409).
+# Retry/backoff for the deployable-adapter publish (rides out a transient concurrent-commit 409).
 _CKPT_FLUSH_RETRIES = 3
 _CKPT_FLUSH_BACKOFF_S = 1.0
+# Retry/backoff for each synchronous checkpoint upload. `on_save` BLOCKS the training loop on the
+# upload, so a transient HF error is retried until the step lands rather than costing the step.
+_CKPT_UPLOAD_RETRIES = 3
+_CKPT_UPLOAD_BACKOFF_S = 5.0
 
 
 def _latest_checkpoint_dir(output_dir: str) -> tuple[int, str] | None:
@@ -368,72 +370,14 @@ def _prune_stale_resume_checkpoints(keep_step: int) -> None:
             api.delete_folder(path_in_repo=folder, repo_id=_w.HF_REPO, repo_type="dataset")
 
 
-class _UploadPump:
-    """Coalescing single-slot upload queue: at most one upload in flight, at most one queued.
-
-    ONE lock guards two fields — ``uploading`` and ``queued`` (newest wins). Every producer path
-    is enqueue-then-pump and the upload thread's finally is done()-then-pump, so a save that
-    lands while the in-flight upload is draining can never be lost (no check-then-set window
-    between separate locks).
-    """
-
-    def __init__(self, start_upload) -> None:
-        self._cond = threading.Condition()
-        self._start_upload = start_upload  # (step, ckpt_dir) -> spawns thread; must call done()
-        self.uploading = False
-        self.queued: tuple[int, str] | None = None
-
-    def enqueue(self, step: int, ckpt_dir: str) -> None:
-        """Queue `step` (replacing any older queued step) and start it if idle."""
-        with self._cond:
-            busy = self.uploading
-            self.queued = (step, ckpt_dir)
-        if busy:
-            print(
-                f"[ckpt] upload busy; queued step {step} (uploads when the in-flight one finishes)"
-            )
-        self.pump()
-
-    def pump(self) -> None:
-        """Atomically start the queued upload iff nothing is in flight."""
-        with self._cond:
-            if self.uploading or self.queued is None:
-                return
-            step, ckpt_dir = self.queued
-            self.queued = None
-            gone = not os.path.isdir(ckpt_dir)
-            if not gone:
-                self.uploading = True
-        if gone:  # pruned by the trainer's save_total_limit meanwhile
-            print(f"[ckpt] queued step {step} checkpoint dir is gone; dropping")
-            return
-        self._start_upload(step, ckpt_dir)
-
-    def done(self) -> None:
-        """Mark the in-flight upload finished, then pump any queued step."""
-        with self._cond:
-            self.uploading = False
-            self._cond.notify_all()
-        self.pump()
-
-    def take_queued(self) -> tuple[int, str] | None:
-        """Remove and return the queued step, if any (for a synchronous train-end flush)."""
-        with self._cond:
-            queued, self.queued = self.queued, None
-            return queued
-
-    def wait_idle(self, timeout: float) -> bool:
-        """Wait up to `timeout` seconds for the in-flight upload (if any) to finish."""
-        with self._cond:
-            return self._cond.wait_for(lambda: not self.uploading, timeout)
-
-
 def make_checkpoint_upload_callback():
     """Return a TrainerCallback that streams each save to HF and publishes deployable per-step adapters.
 
-    Uploads run in a background daemon thread; a save that fires while one is in flight is queued
-    (newest wins) and uploaded next. on_train_end flushes the queue and the final deployable
-    synchronously before exit.
+    Uploads are SYNCHRONOUS: on_save blocks the training loop until the checkpoint is durably on
+    HF. A save therefore never returns while its upload is still running, so there is never a
+    second upload in flight when the next save fires — the "upload busy" contention (and the
+    coalescing/dropping it used to cause) cannot happen by construction. Every `save_every` step is
+    guaranteed to upload, retrying transient HF errors instead of skipping.
     """
     from transformers import TrainerCallback
 
@@ -491,72 +435,57 @@ def make_checkpoint_upload_callback():
         _prune_stale_resume_checkpoints(step)
         _w.heartbeat("checkpoint_uploaded", step=step)
 
-    def _start_upload(step: int, ckpt_dir: str) -> None:
-        """Upload `step` in a background thread; the pump is marked done (and re-pumped) after."""
+    uploaded_steps: set[int] = set()
 
-        def _upload():
+    def _upload_with_retries(step: int, ckpt_dir: str) -> bool:
+        """Run the full upload for `step`, retrying transient failures; True once it lands."""
+        for attempt in range(_CKPT_UPLOAD_RETRIES):
             try:
                 _upload_once(step, ckpt_dir)
+                uploaded_steps.add(step)
+                return True
             except Exception as e:
-                print("ckpt upload warn:", e)
-            finally:
-                pump.done()
-
-        threading.Thread(target=_upload, daemon=True).start()
-
-    # `save_every` must not be advisory: a save that fires while a prior upload is still in
-    # flight is QUEUED (newest wins) and uploaded as soon as the in-flight one finishes, instead
-    # of being dropped with "upload busy; skipping step N". Coalescing to the newest pending step
-    # keeps at most one waiter, so a slow uplink can't pile up threads, while the registered step
-    # list stays as dense as the uplink allows.
-    pump = _UploadPump(_start_upload)
+                if attempt + 1 < _CKPT_UPLOAD_RETRIES:
+                    print(
+                        f"[ckpt] step {step} upload retry {attempt + 1}/{_CKPT_UPLOAD_RETRIES}: {e}"
+                    )
+                    time.sleep(_CKPT_UPLOAD_BACKOFF_S * (attempt + 1))
+                else:
+                    print(
+                        f"[ckpt] step {step} upload FAILED after {_CKPT_UPLOAD_RETRIES} attempts: {e}"
+                    )
+                    with contextlib.suppress(Exception):
+                        _w.heartbeat("checkpoint_upload_failed", step=step)
+        return False
 
     class _CheckpointUpload(TrainerCallback):
         def on_save(self, args, state, control, **kwargs):
+            # SYNCHRONOUS on the training thread: the trainer blocks here until this checkpoint is
+            # uploaded. `save_total_limit=1` rotation (which runs earlier, inside this same save's
+            # `_save_checkpoint`) deletes the PREVIOUS checkpoint, never the current one — and the
+            # next save can't begin until we return — so the dir is stable for the whole upload and
+            # no snapshot/queue is needed. Because the upload finishes before the save returns,
+            # there is never a concurrent upload for a later save to be "busy" against.
             if not _w.HF_REPO:
                 return
             step = int(state.global_step)
             ckpt_dir = os.path.join(args.output_dir, f"checkpoint-{step}")
             if not os.path.isdir(ckpt_dir):
                 return
-            pump.enqueue(step, ckpt_dir)
+            _upload_with_retries(step, ckpt_dir)
 
         def on_train_end(self, args, state, control, **kwargs):
-            # Daemon threads are killed on exit; flush synchronously here.
+            # Safety net for a final checkpoint the trainer wrote without an on_save (e.g. a
+            # load_best_model_at_end / end-of-training save). Synchronous on_save already uploaded
+            # every step it saw, so this only publishes an as-yet-unuploaded latest step.
             if not _w.HF_REPO:
                 return
             latest = _latest_checkpoint_dir(args.output_dir)
             if latest is None:
                 return
             step, ckpt_dir = latest
-            # A queued save the pump never started must not be silently dropped: take it out
-            # of the pump (so a racing done()->pump() can't double-start it) and run it here.
-            queued = pump.take_queued()
-            idle = pump.wait_idle(_CKPT_FLUSH_TIMEOUT_S)
-            if queued is not None and os.path.isdir(queued[1]):
-                q_step, q_dir = queued
-                if idle:
-                    try:
-                        _upload_once(q_step, q_dir)
-                    except Exception as e:
-                        print("ckpt upload warn:", e)
-                else:
-                    # Another upload is still pruning latest-only resume checkpoints. Publish only
-                    # the durable per-step adapter here so concurrent prune passes cannot delete a
-                    # newer resume checkpoint while it is being flushed.
-                    try:
-                        _publish_deployable_recombined(q_dir, q_step, with_retry=True)
-                    except Exception as e:
-                        print("ckpt deployable flush warn:", e)
-                if q_step == step:
-                    return  # the queued flush already published the final deployable
-            if idle:
-                _publish_deployable_recombined(ckpt_dir, step)
-            else:
-                print(
-                    f"[ckpt] upload still in flight after {_CKPT_FLUSH_TIMEOUT_S:.0f}s; publishing "
-                    f"final deployable (step {step}) alongside it"
-                )
-                _publish_deployable_recombined(ckpt_dir, step, with_retry=True)
+            if step in uploaded_steps:
+                return
+            _upload_with_retries(step, ckpt_dir)
 
     return _CheckpointUpload()

{freesolo_flash_dev-0.2.43 → freesolo_flash_dev-0.2.44}/flash/providers/_worker.py

@@ -112,13 +112,14 @@ def _effective_worker_env(spec=None) -> dict[str, str]:
 # Chalk is published to PUBLIC PyPI on every version bump (chalk .github/workflows/publish.yml).
 # Pin the exact PUBLIC version so worker + kernel-cache-bake pods install with NO GitHub auth — the
 # chalk repo is INTERNAL, so a git+https default made installs fail wherever GITHUB_TOKEN is absent
-# (bake pod, tokenless workers). freesolo-chalk 0.5.2 = chalk dev->main sync #44 (standalone, Liger
-# fully replaced; adds arch-TUNED gdn_gated_rmsnorm + lora on sm89/sm90/sm100, MoE grouped-GEMM,
-# serving surface, MiniCPM/Llama kernels). Bump DEFAULT_CHALK_VERSION when the merged kernel surface moves.
-DEFAULT_CHALK_VERSION = "0.5.2"
+# (bake pod, tokenless workers). freesolo-chalk 0.5.3 = chalk dev->main sync #67 (adds the decode
+# serving kernels rope + GDN gated-rmsnorm; GPU-A/B'd launch tuning: flce chunk_mult 16, rope
+# num_warps=2, gdn_conv warps formula, lora fwd autotune candidate; OP_MEAS_OVERRIDE; flce
+# chunk_mult tradeoff docs). Bump DEFAULT_CHALK_VERSION when the merged kernel surface moves.
+DEFAULT_CHALK_VERSION = "0.5.3"
 DEFAULT_CHALK_SPEC = f"freesolo-chalk=={DEFAULT_CHALK_VERSION}"
 # Provenance only (kernel-cache fingerprint / traceability): the chalk commit this version was cut from.
-LATEST_CHALK_MAIN_SHA = "37b2fb9e1a123ceb902ba6ffc4ed182182c6e7d7"
+LATEST_CHALK_MAIN_SHA = "c63f0367e01795006bc0e91e2540ec85e740f19b"
 
 
 def chalk_extra_pip(spec=None) -> list[str]:

{freesolo_flash_dev-0.2.43 → freesolo_flash_dev-0.2.44}/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "freesolo-flash-dev"
-version = "0.2.43"
+version = "0.2.44"
 description = "Flash — managed LoRA post-training (SFT/GRPO) for Freesolo environments, driven by the `flash` CLI"
 readme = "README.md"
 # RunPod Flash supports Python 3.11-3.12 (not 3.13 yet).
@@ -106,7 +106,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 = "0.2.43"
+version = "0.2.44"
 
 [tool.pytest.ini_options]
 pythonpath = ["."]

{freesolo_flash_dev-0.2.43 → freesolo_flash_dev-0.2.44}/tests/test_checkpoints.py

@@ -10,7 +10,6 @@ All HF/network boundaries are stubbed; nothing here touches a GPU or the network
 
 from __future__ import annotations
 
-import time
 from types import SimpleNamespace
 
 import pytest
@@ -197,19 +196,20 @@ def test_latest_checkpoint_dir_picks_highest_step(tmp_path):
 def test_on_train_end_flushes_final_deployable_checkpoint(
     tmp_path, monkeypatch, fake_trainer_callback
 ):
-    """A fast RL run can exit before its last save's async (daemon) upload finishes, so
-    on_train_end must SYNCHRONOUSLY publish the latest on-disk checkpoint as a deployable
-    snapshot — otherwise `flash checkpoints` is empty even though the run trained fine."""
+    """on_train_end is the safety net for a final checkpoint the trainer wrote WITHOUT an on_save
+    (e.g. an end-of-training / load_best_model_at_end save). It must publish the latest on-disk
+    checkpoint as a deployable — otherwise `flash checkpoints` is empty even though the run trained
+    fine. (on_save uploads synchronously, so every step it sees is already up by this point.)"""
     import flash.engine.worker as worker
 
     rec = _RecordingHfApi()
     _prime_worker(monkeypatch, rec)
     out = tmp_path / "out"
     out.mkdir()
-    _make_ckpt_dir(out, 8)  # last save the trainer wrote locally; its async upload was "lost"
+    _make_ckpt_dir(out, 8)  # a final checkpoint on disk that no on_save uploaded
     cb = worker.make_checkpoint_upload_callback()
 
-    # No on_save uploads recorded (simulating their daemon threads being killed at exit).
+    # No on_save fired for this step, so on_train_end must publish it.
     cb.on_train_end(SimpleNamespace(output_dir=str(out)), None, None)
 
     deployable = [
@@ -219,72 +219,38 @@ def test_on_train_end_flushes_final_deployable_checkpoint(
     assert "optimizer.pt" in deployable[0]["ignore_patterns"]
 
 
-def test_on_train_end_publishes_final_deployable_when_lock_is_held(
-    tmp_path, monkeypatch, fake_trainer_callback
-):
-    """Regression: a slow on_save upload can still hold the upload lock at run end (and the final
-    step's own on_save may have been skipped on that busy lock). on_train_end must then publish the
-    final deployable WITHOUT the lock instead of timing out and skipping it — otherwise the worker
-    exits, kills the daemon mid-upload, and `flash checkpoints` is empty despite a successful run."""
-    import threading
-
-    from flash.engine.worker import hf as worker_hf
+def test_on_train_end_no_checkpoints_is_noop(tmp_path, monkeypatch, fake_trainer_callback):
+    import flash.engine.worker as worker
 
     rec = _RecordingHfApi()
-    worker = _prime_worker(monkeypatch, rec)
-    monkeypatch.setattr(worker_hf, "_CKPT_FLUSH_TIMEOUT_S", 0.05)  # give up on the held lock fast
-
+    _prime_worker(monkeypatch, rec)
     out = tmp_path / "out"
     out.mkdir()
-    _make_ckpt_dir(out, 4)  # earlier step; its on_save upload will block, holding the lock
-    _make_ckpt_dir(out, 8)  # the FINAL step (its on_save was skipped on the busy lock)
-
-    holding = threading.Event()
-    release = threading.Event()
-    base_upload = rec.upload_folder
-
-    def upload(**kwargs):
-        base_upload(**kwargs)
-        if kwargs["path_in_repo"].endswith("checkpoint/checkpoint-4"):
-            holding.set()  # the resume upload now holds the lock...
-            release.wait(5)  # ...keep holding it across the on_train_end flush
-
-    rec.upload_folder = upload
-
-    started: list = []
-    real_thread = threading.Thread
-    monkeypatch.setattr(
-        worker.threading,
-        "Thread",
-        lambda *a, **k: started.append(real_thread(*a, **k)) or started[-1],
+    worker.make_checkpoint_upload_callback().on_train_end(
+        SimpleNamespace(output_dir=str(out)), None, None
     )
-
-    cb = worker.make_checkpoint_upload_callback()
-    try:
-        cb.on_save(SimpleNamespace(output_dir=str(out)), SimpleNamespace(global_step=4), None)
-        assert holding.wait(5), "the step-4 upload should be holding the lock"
-        cb.on_train_end(SimpleNamespace(output_dir=str(out)), None, None)
-        deployable = [
-            u for u in rec.uploads if u["path_in_repo"].endswith("checkpoints/step-8/adapter")
-        ]
-        assert len(deployable) == 1, "final deployable must publish even when the lock is held"
-    finally:
-        release.set()
-        for t in started:
-            t.join(5)
+    assert rec.uploads == []
 
 
-def test_on_train_end_no_checkpoints_is_noop(tmp_path, monkeypatch, fake_trainer_callback):
+def test_on_save_uploads_synchronously_before_returning(
+    tmp_path, monkeypatch, fake_trainer_callback
+):
+    """Uploads are synchronous: by the time on_save returns, the checkpoint is fully on HF — no
+    background thread is left running. This is what makes "upload busy" impossible by construction:
+    a later save can never find this one still in flight."""
     import flash.engine.worker as worker
 
     rec = _RecordingHfApi()
     _prime_worker(monkeypatch, rec)
     out = tmp_path / "out"
     out.mkdir()
-    worker.make_checkpoint_upload_callback().on_train_end(
-        SimpleNamespace(output_dir=str(out)), None, None
-    )
-    assert rec.uploads == []
+    _make_ckpt_dir(out, 4)
+    cb = worker.make_checkpoint_upload_callback()
+    cb.on_save(SimpleNamespace(output_dir=str(out)), SimpleNamespace(global_step=4), None)
+    # Fully uploaded the moment on_save returns — no polling, no thread join needed.
+    paths = [u["path_in_repo"] for u in rec.uploads]
+    assert "rl/flash-ckpt-1/checkpoints/step-4/adapter" in paths
+    assert "rl/flash-ckpt-1/checkpoint/checkpoint-4" in paths
 
 
 def test_on_save_publishes_deployable_before_resume(tmp_path, monkeypatch, fake_trainer_callback):
@@ -294,25 +260,11 @@ def test_on_save_publishes_deployable_before_resume(tmp_path, monkeypatch, fake_
 
     rec = _RecordingHfApi()
     _prime_worker(monkeypatch, rec)
-
-    class _SyncThread:  # run the upload body inline so the unit test is deterministic
-        def __init__(self, target=None, daemon=None, **kw):
-            self._target = target
-
-        def start(self):
-            if self._target:
-                self._target()
-
-    monkeypatch.setattr(worker.threading, "Thread", _SyncThread)
     out = tmp_path / "out"
     out.mkdir()
     _make_ckpt_dir(out, 4)
     cb = worker.make_checkpoint_upload_callback()
-    cb.on_save(
-        SimpleNamespace(output_dir=str(out)),
-        SimpleNamespace(global_step=4),
-        None,
-    )
+    cb.on_save(SimpleNamespace(output_dir=str(out)), SimpleNamespace(global_step=4), None)
     paths = [u["path_in_repo"] for u in rec.uploads]
     assert paths == [
         "rl/flash-ckpt-1/checkpoints/step-4/adapter",  # deployable first
@@ -320,115 +272,61 @@ def test_on_save_publishes_deployable_before_resume(tmp_path, monkeypatch, fake_
     ]
 
 
-def test_on_save_queues_busy_step_instead_of_skipping(tmp_path, monkeypatch, fake_trainer_callback):
-    """`save_every` must not be advisory: a save that fires while a prior upload is still in
-    flight is queued and uploaded once the in-flight one finishes — previously it was dropped
-    ("upload busy; skipping step N"), leaving a sparse registered step list."""
-    import threading
+def test_consecutive_saves_each_upload_no_coalescing(tmp_path, monkeypatch, fake_trainer_callback):
+    """Every save uploads its own step. Synchronous uploads make "upload busy" coalescing
+    impossible: back-to-back saves each publish their deployable, so the registered step list is
+    never thinned under contention (the old newest-wins slot dropped intermediate steps)."""
+    import flash.engine.worker as worker
 
     rec = _RecordingHfApi()
-    worker = _prime_worker(monkeypatch, rec)
-
+    _prime_worker(monkeypatch, rec)
     out = tmp_path / "out"
     out.mkdir()
-    _make_ckpt_dir(out, 4)
-    _make_ckpt_dir(out, 8)
-
-    holding = threading.Event()
-    release = threading.Event()
-    base_upload = rec.upload_folder
-
-    def upload(**kwargs):
-        base_upload(**kwargs)
-        if kwargs["path_in_repo"].endswith("checkpoint/checkpoint-4"):
-            holding.set()  # the step-4 resume upload now holds the lock...
-            release.wait(5)  # ...until we let it finish
-
-    rec.upload_folder = upload
-
     cb = worker.make_checkpoint_upload_callback()
-    cb.on_save(SimpleNamespace(output_dir=str(out)), SimpleNamespace(global_step=4), None)
-    assert holding.wait(5), "the step-4 upload should be holding the lock"
-    # step-8 save fires while step-4 is uploading: it must be queued, not skipped.
-    cb.on_save(SimpleNamespace(output_dir=str(out)), SimpleNamespace(global_step=8), None)
-    release.set()
-
-    deadline = time.time() + 10
-    while time.time() < deadline:
-        if any(u["path_in_repo"].endswith("checkpoints/step-8/adapter") for u in rec.uploads):
-            break
-        time.sleep(0.02)
-    deployable = [
-        u for u in rec.uploads if u["path_in_repo"].endswith("checkpoints/step-8/adapter")
-    ]
-    assert len(deployable) == 1, "a busy-lock save must be queued and uploaded, not skipped"
-
-
-def test_upload_pump_no_lost_wakeup(tmp_path):
-    """Regression: a save enqueued in the window where the in-flight upload has passed its own
-    drain check but not yet finished must still be started by done() — with the old two-lock
-    design it could be silently dropped."""
-    from flash.engine.worker.hf import _UploadPump
+    for step in (4, 8, 12):
+        _make_ckpt_dir(out, step)
+        cb.on_save(SimpleNamespace(output_dir=str(out)), SimpleNamespace(global_step=step), None)
+    deployables = sorted(
+        p["path_in_repo"] for p in rec.uploads if p["path_in_repo"].endswith("/adapter")
+    )
+    assert deployables == [
+        "rl/flash-ckpt-1/checkpoints/step-12/adapter",
+        "rl/flash-ckpt-1/checkpoints/step-4/adapter",
+        "rl/flash-ckpt-1/checkpoints/step-8/adapter",
+    ], "each save must publish its own deployable — no step coalesced away"
 
-    ckpt = tmp_path / "checkpoint-8"
-    ckpt.mkdir()
-    started: list[int] = []
-    pump = _UploadPump(lambda step, d: started.append(step))
-    # Simulate an in-flight upload that has already drained the queue (the racy window).
-    pump.uploading = True
-    pump.enqueue(8, str(ckpt))
-    assert started == [], "enqueue while busy must queue, not start"
-    pump.done()  # in-flight finishes: the queued step MUST be pumped, never lost
-    assert started == [8]
-    pump.done()
-    assert started == [8], "queue is single-slot; nothing left to start"
-
-
-def test_on_train_end_flushes_queued_deployable_without_resume_race(
-    tmp_path, monkeypatch, fake_trainer_callback
-):
-    """A save queued behind a still-in-flight upload at run end must publish its deployable adapter
-    without starting a second resume-checkpoint upload/prune race."""
-    import threading
 
+def test_on_save_retries_transient_upload_error(tmp_path, monkeypatch, fake_trainer_callback):
+    """A transient HF error must not cost the step: the synchronous upload retries until it lands,
+    so `save_every` is a guarantee, not best-effort."""
     from flash.engine.worker import hf as worker_hf
 
     rec = _RecordingHfApi()
     worker = _prime_worker(monkeypatch, rec)
-    monkeypatch.setattr(worker_hf, "_CKPT_FLUSH_TIMEOUT_S", 0.05)
+    monkeypatch.setattr(worker_hf, "_CKPT_UPLOAD_BACKOFF_S", 0.0)  # no real sleep between retries
 
     out = tmp_path / "out"
     out.mkdir()
     _make_ckpt_dir(out, 4)
-    _make_ckpt_dir(out, 8)
 
-    holding = threading.Event()
-    release = threading.Event()
+    resume_calls = {"n": 0}
     base_upload = rec.upload_folder
 
-    def upload(**kwargs):
-        base_upload(**kwargs)
+    def flaky(**kwargs):
         if kwargs["path_in_repo"].endswith("checkpoint/checkpoint-4"):
-            holding.set()  # the step-4 resume upload is now in flight...
-            release.wait(5)  # ...and stays in flight across on_train_end
+            resume_calls["n"] += 1
+            if resume_calls["n"] == 1:
+                raise RuntimeError("transient 500 from HF")
+        return base_upload(**kwargs)
 
-    rec.upload_folder = upload
+    rec.upload_folder = flaky
 
     cb = worker.make_checkpoint_upload_callback()
-    try:
-        cb.on_save(SimpleNamespace(output_dir=str(out)), SimpleNamespace(global_step=4), None)
-        assert holding.wait(5), "the step-4 upload should be in flight"
-        cb.on_save(SimpleNamespace(output_dir=str(out)), SimpleNamespace(global_step=8), None)
-        cb.on_train_end(SimpleNamespace(output_dir=str(out)), None, None)
-        paths = [u["path_in_repo"] for u in rec.uploads]
-        assert "rl/flash-ckpt-1/checkpoints/step-8/adapter" in paths, (
-            "queued step's deployable must be flushed at train end"
-        )
-        assert "rl/flash-ckpt-1/checkpoint/checkpoint-8" not in paths, (
-            "queued step's resume checkpoint must not race the in-flight resume prune"
-        )
-    finally:
-        release.set()
+    cb.on_save(SimpleNamespace(output_dir=str(out)), SimpleNamespace(global_step=4), None)
+    assert resume_calls["n"] == 2, "the resume upload must be retried after a transient failure"
+    assert any(u["path_in_repo"].endswith("checkpoint/checkpoint-4") for u in rec.uploads), (
+        "the checkpoint must land on retry, not be skipped"
+    )
 
 
 def test_prune_stale_resume_checkpoints_keeps_only_latest(monkeypatch):
@@ -485,15 +383,6 @@ def test_on_save_skips_deployable_when_vl_recombine_fails(
 
     monkeypatch.setattr(worker, "recombined_warmstart_adapter_dir", _raise)
 
-    class _SyncThread:
-        def __init__(self, target=None, daemon=None, **kw):
-            self._target = target
-
-        def start(self):
-            if self._target:
-                self._target()
-
-    monkeypatch.setattr(worker.threading, "Thread", _SyncThread)
     out = tmp_path / "out"
     out.mkdir()
     _make_ckpt_dir(out, 4)

{freesolo_flash_dev-0.2.43 → freesolo_flash_dev-0.2.44}/tests/test_cli_commands.py

@@ -77,6 +77,10 @@ class _FakeClient:
         self.calls.append(("cancel", run_id))
         return {"run_id": run_id, "state": "cancelled"}
 
+    def checkpoints(self, run_id: str) -> list[dict]:
+        self.calls.append(("checkpoints", run_id))
+        return [{"step": 20}, {"step": 40}]
+
     def deploy(self, run_id: str, **kwargs) -> dict:
         self.calls.append(("deploy", run_id, kwargs))
         return {
@@ -372,6 +376,36 @@ def test_follow_logs_uses_status_progress_when_log_tail_lags(monkeypatch, capsys
     assert "realized_cost=$1.2346" in err
 
 
+def test_cancel_surfaces_surviving_checkpoints(fake_client, capsys) -> None:
+    """`state=cancelled` + adapter_ref=null + cost=0 reads as discardable, yet the per-step
+    deployable checkpoints streamed before the cancel survive it — the cancel output must say
+    so (on stderr in the plain path, keeping the stdout JSON machine-readable)."""
+    import json as _json
+
+    assert _run(["cancel", "flash-1"]) == 0
+    assert ("checkpoints", "flash-1") in fake_client.calls
+    out, err = capsys.readouterr()
+    assert _json.loads(out)["state"] == "cancelled"  # stdout stays pure JSON in the plain path
+    assert "2 deployable checkpoint(s) survive this cancel" in err
+    assert "flash checkpoints flash-1" in err
+    assert "flash deploy flash-1/step-40" in err  # points at the newest surviving step
+
+
+def test_cancel_hint_is_best_effort_when_checkpoint_listing_fails(
+    fake_client, capsys, monkeypatch
+) -> None:
+    """The surviving-checkpoints lookup must never break `flash cancel` itself."""
+
+    def boom(run_id):
+        raise RuntimeError("backend hiccup")
+
+    monkeypatch.setattr(fake_client, "checkpoints", boom)
+    assert _run(["cancel", "flash-1"]) == 0
+    out, err = capsys.readouterr()
+    assert '"state": "cancelled"' in out
+    assert "deployable checkpoint" not in err
+
+
 def test_cancel_deploy_undeploy_deployments(fake_client, capsys) -> None:
     assert _run(["cancel", "flash-1"]) == 0
     assert ("cancel", "flash-1") in fake_client.calls