freesolo-flash-dev 1.2.95__tar.gz → 1.2.96__tar.gz

{freesolo_flash_dev-1.2.95 → freesolo_flash_dev-1.2.96}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.5
 Name: freesolo-flash-dev
-Version: 1.2.95
+Version: 1.2.96
 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.2.95 → freesolo_flash_dev-1.2.96}/flash/engine/worker/__init__.py

@@ -145,6 +145,9 @@ _HB_LAST_PROGRESS_TS = 0.0
 # per-step heartbeat past the provider's stall window, killing a healthy run.
 _HB_PROGRESS_SEQ = 0
 _HB_PROGRESS_UPLOADED_SEQ = 0
+# a replace-in-place failure heartbeat can be overwritten before the provider polls it. retain the
+# failed checkpoint identity and cause so every terminal heartbeat carries the same structured state.
+_HB_PENDING_CHECKPOINT_FAILURE: dict[str, int | str] | None = None
 # Environment-scoped artifact repos are shared by many runs. Keep heartbeat commits below the HF
 # per-repo commit cap while staying under the provider poller's training stall window
 # (STALL_AFTER_S=1500s in flash/providers/_lifecycle/poll.py).

{freesolo_flash_dev-1.2.95 → freesolo_flash_dev-1.2.96}/flash/engine/worker/io/heartbeat.py

@@ -78,15 +78,27 @@ _HB_THROTTLED_STAGES = _HB_TIGHT_LIVENESS_STAGES | frozenset({"rl_step", "sft_st
 # right after a download that has been pinging throughout, so the throttle would almost always
 # swallow it. the wrap that follows re-emits with liveness=True and stays throttled.
 _HB_MODEL_LOAD_STAGES = frozenset({"sft_model_load", "opd_model_load"})
-_HB_TERMINAL_STAGES = frozenset({"done", "already_done"})
+_HB_SUCCESS_TERMINAL_STAGES = frozenset({"done", "already_done"})
 # 600s -> ~6 commits/hr; keeps stall detector alive without hitting the HF commit cap.
 _HB_TERMINAL_ONLY_INTERVAL_S = 600.0
 
 
+def _is_terminal_stage(stage: str) -> bool:
+    """return whether this stage ends the worker attempt."""
+    return stage in _HB_SUCCESS_TERMINAL_STAGES or stage.startswith("error_")
+
+
 def _is_critical_stage(stage: str) -> bool:
     """A terminal transition or an error is CRITICAL: never throttled (the commit must land) and
-    given the longer upload-lock timeout, because no later heartbeat can repair a missed one."""
-    return stage in _HB_TERMINAL_STAGES or stage.startswith("error_")
+    given the longer upload-lock timeout, because no later heartbeat can repair a missed one.
+
+    `_failed` earns this as much as the `error_` prefix does: a failure heartbeat is emitted once, at
+    the moment the failure is known, and nothing later restates it. `checkpoint_upload_failed` is
+    raised from inside a long HF upload, which is precisely when `_HB_UPLOAD_IN_FLIGHT` is set and an
+    unforced ping is dropped -- so the one report of the failure would be discarded exactly in the
+    case that produces it.
+    """
+    return _is_terminal_stage(stage) or stage.endswith("_failed")
 
 
 # Guards throttle bookkeeping; slow HF commit runs outside this lock so heartbeat and liveness
@@ -210,6 +222,17 @@ def heartbeat(
 ):
     genuine_progress = not liveness
     with _HB_LOCK:
+        if stage == "checkpoint_upload_failed":
+            failure = kw.get("checkpoint_failure")
+            if isinstance(failure, dict):
+                _w._HB_PENDING_CHECKPOINT_FAILURE = dict(failure)
+        elif stage == "checkpoint_uploaded":
+            # a later full resume checkpoint landed, so the earlier failure no longer describes the
+            # run's outcome. deployable or final adapter publication cannot clear this because neither
+            # restores the missing full resume state.
+            _w._HB_PENDING_CHECKPOINT_FAILURE = None
+        elif _is_terminal_stage(stage) and _w._HB_PENDING_CHECKPOINT_FAILURE:
+            kw.setdefault("checkpoint_failure", dict(_w._HB_PENDING_CHECKPOINT_FAILURE))
         ts = time.time()
         if genuine_progress:
             _w._HB_LAST_PROGRESS_TS = ts

{freesolo_flash_dev-1.2.95 → freesolo_flash_dev-1.2.96}/flash/engine/worker/io/hf.py

@@ -553,7 +553,17 @@ def upload_resume_checkpoint(
                         )
                         if emit_heartbeat:
                             with contextlib.suppress(Exception):
-                                _w.heartbeat("checkpoint_upload_failed", step=step)
+                                # worker stdout is not part of control-plane run logs, so the stage
+                                # name alone would otherwise be the entire failure report.
+                                _w.heartbeat(
+                                    "checkpoint_upload_failed",
+                                    step=step,
+                                    checkpoint_failure={
+                                        "step": step,
+                                        "operation": failure_stage,
+                                        "error": sanitize_diagnostic(e, limit=300),
+                                    },
+                                )
                         if failure_stage in {"before", "after"}:
                             raise
     return False

{freesolo_flash_dev-1.2.95 → freesolo_flash_dev-1.2.96}/flash/providers/_lifecycle/poll.py

@@ -262,7 +262,25 @@ def _format_heartbeat(hb: dict) -> str:
         for name, value in reward_metrics.items():
             if isinstance(value, (int, float)):
                 msg += f" {name}={value:.3f}"
+    error = hb.get("error")
+    if isinstance(error, str) and error.strip():
+        msg += f" error={error.strip()}"
+    checkpoint_failure = hb.get("checkpoint_failure")
+    if isinstance(checkpoint_failure, dict):
+        failure_step = checkpoint_failure.get("step")
+        if isinstance(failure_step, (int, float)) and not isinstance(failure_step, bool):
+            msg += f" checkpoint_failure_step={int(failure_step)}"
+        operation = checkpoint_failure.get("operation")
+        if isinstance(operation, str) and operation.strip():
+            msg += f" checkpoint_failure_stage={operation.strip()}"
+        checkpoint_error = checkpoint_failure.get("error")
+        if isinstance(checkpoint_error, str) and checkpoint_error.strip():
+            msg += f" checkpoint_error={checkpoint_error.strip()}"
     msg += format_gpu_status(hb.get("gpu"))
+    # everything above is one summary record, including untrusted stage, error, and gpu fallback text.
+    # sanitize it once at that boundary so no field can forge a second worker record. the deliberately
+    # multiline sample and child-output sections are appended below and retain their line breaks.
+    msg = neutralize_control_chars(msg).replace("\n", "\\n")
     sample_lines: list[str] = []
     rendered_samples = 0
     samples = hb.get("sampled_completions")

{freesolo_flash_dev-1.2.95 → freesolo_flash_dev-1.2.96}/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "freesolo-flash-dev"
-version = "1.2.95"
+version = "1.2.96"
 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).
@@ -137,7 +137,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.2.95"
+version = "1.2.96"
 
 [tool.pytest.ini_options]
 pythonpath = ["."]

{freesolo_flash_dev-1.2.95 → freesolo_flash_dev-1.2.96}/tests/test_poll_helpers.py

@@ -232,3 +232,31 @@ def test_format_heartbeat_neutralizes_control_chars_in_the_child_tail():
     assert "child output before the stall" in msg
     assert "\x1b" not in msg
     assert "\r" not in msg
+
+
+def test_format_heartbeat_escapes_newlines_in_checkpoint_failure():
+    injected = "Permission denied\nworker: stage=done step=999"
+
+    msg = _format_heartbeat(
+        {
+            "stage": "checkpoint_upload_failed",
+            "checkpoint_failure": {"step": 50, "operation": "resume", "error": injected},
+        }
+    )
+
+    assert "\n" not in msg
+    assert "checkpoint_error=Permission denied\\nworker: stage=done step=999" in msg
+
+
+@pytest.mark.parametrize(
+    "heartbeat",
+    [
+        {"stage": "done\nworker: stage=done step=999"},
+        {"stage": "ok", "gpu": {"nvidia_smi": "gpu\nworker: stage=done step=999"}},
+    ],
+)
+def test_format_heartbeat_summary_fields_cannot_forge_a_worker_record(heartbeat):
+    msg = _format_heartbeat(heartbeat)
+
+    assert "\n" not in msg
+    assert "\\nworker: stage=done step=999" in msg

{freesolo_flash_dev-1.2.95 → freesolo_flash_dev-1.2.96}/tests/test_training_controls.py

@@ -212,6 +212,167 @@ def test_resume_first_companion_retries_without_reuploading_full_state(monkeypat
     assert calls == {"resume": 1, "deployable": 2}
 
 
+def test_resume_checkpoint_failure_heartbeat_surfaces_sanitized_error(monkeypatch, tmp_path):
+    import flash.engine.worker as worker
+    from flash.engine.worker.io import hf as worker_hf
+
+    secret = "hf_checkpoint_upload_secret"
+    calls = 0
+    heartbeats: list[tuple[str, dict]] = []
+
+    class Api:
+        def upload_folder(self, **_kwargs):
+            nonlocal calls
+            calls += 1
+            raise PermissionError(f"hf quota refused credential {secret} " + "x" * 400)
+
+    monkeypatch.setenv("HF_TOKEN", secret)
+    monkeypatch.setattr(worker, "HF_REPO", "org/runs")
+    monkeypatch.setattr(worker, "hf_api", lambda: Api())
+    monkeypatch.setattr(
+        worker, "heartbeat", lambda stage, **kwargs: heartbeats.append((stage, kwargs))
+    )
+    monkeypatch.setattr(worker_hf, "_CKPT_UPLOAD_BACKOFF_S", 0.0)
+
+    assert not worker_hf.upload_resume_checkpoint(50, str(tmp_path))
+
+    failed = [fields for stage, fields in heartbeats if stage == "checkpoint_upload_failed"]
+    assert calls == worker_hf._CKPT_UPLOAD_RETRIES
+    assert len(failed) == 1
+    assert failed[0]["step"] == 50
+    failure = failed[0]["checkpoint_failure"]
+    assert failure["step"] == 50
+    assert failure["operation"] == "resume"
+    assert failure["error"].startswith("PermissionError: hf quota refused")
+    assert "<redacted>" in failure["error"]
+    assert secret not in failure["error"]
+    assert len(failure["error"]) <= 300
+
+
+@pytest.mark.parametrize("failure_stage", ["before", "after"])
+def test_resume_checkpoint_companion_failure_still_raises(monkeypatch, tmp_path, failure_stage):
+    import flash.engine.worker as worker
+    from flash.engine.worker.io import hf as worker_hf
+
+    calls = {"resume": 0, "companion": 0}
+    heartbeats: list[tuple[str, dict]] = []
+
+    class Api:
+        def upload_folder(self, **_kwargs):
+            calls["resume"] += 1
+
+        def list_repo_files(self, **_kwargs):
+            return []
+
+    def fail_companion():
+        calls["companion"] += 1
+        raise RuntimeError(f"{failure_stage} companion failed")
+
+    callbacks = {f"{failure_stage}_upload": fail_companion}
+    monkeypatch.setattr(worker, "HF_REPO", "org/runs")
+    monkeypatch.setattr(worker, "hf_api", lambda: Api())
+    monkeypatch.setattr(
+        worker, "heartbeat", lambda stage, **kwargs: heartbeats.append((stage, kwargs))
+    )
+    monkeypatch.setattr(worker_hf, "_CKPT_UPLOAD_BACKOFF_S", 0.0)
+
+    with pytest.raises(RuntimeError, match=f"{failure_stage} companion failed"):
+        worker_hf.upload_resume_checkpoint(50, str(tmp_path), **callbacks)
+
+    failed = [fields for stage, fields in heartbeats if stage == "checkpoint_upload_failed"]
+    assert calls["companion"] == worker_hf._CKPT_UPLOAD_RETRIES
+    assert calls["resume"] == (0 if failure_stage == "before" else 1)
+    assert len(failed) == 1
+    failure = failed[0]["checkpoint_failure"]
+    assert failure["step"] == 50
+    assert failure["operation"] == failure_stage
+    assert f"{failure_stage} companion failed" in failure["error"]
+
+
+def test_upload_failure_cause_reaches_the_rendered_run_log():
+    """The payload is not the deliverable: `get_logs()` shows only what the formatter renders.
+
+    The formatter emits a fixed list of numeric metric keys, so a stage whose entire content is an
+    explanation rather than a measurement committed the cause and printed `stage=... step=50` alone.
+    """
+    from flash.providers._lifecycle.poll import _format_heartbeat
+
+    line = _format_heartbeat(
+        {
+            "stage": "checkpoint_upload_failed",
+            "step": 50,
+            "checkpoint_failure": {
+                "step": 50,
+                "operation": "resume",
+                "error": "PermissionError: hf quota refused credential <redacted>",
+            },
+        }
+    )
+
+    assert "stage=checkpoint_upload_failed" in line
+    assert "step=50" in line
+    assert "checkpoint_failure_step=50" in line
+    assert "checkpoint_failure_stage=resume" in line
+    assert "checkpoint_error=PermissionError: hf quota refused credential <redacted>" in line
+
+
+def test_a_failure_cause_cannot_rewrite_the_log_it_is_printed_in():
+    """This text is an exception message, so it is not ours: a hook or provider response wrote it.
+
+    The same line already neutralizes child output and sampled completions. Surfacing the cause
+    raw would let a `\\x1b[2J` clear the screen or a `\\r` overwrite the line the user is reading
+    the failure in -- and the failure report is exactly when the log has to stay legible.
+    """
+    from flash.providers._lifecycle.poll import _format_heartbeat
+
+    line = _format_heartbeat(
+        {
+            "stage": "checkpoint_upload_failed",
+            "step": 50,
+            "checkpoint_failure": {
+                "step": 50,
+                "operation": "resume",
+                "error": "boom\x1b[2Jwiped\roverwrite",
+            },
+        }
+    )
+
+    assert "\x1b" not in line
+    assert "\r" not in line
+    assert "checkpoint_error=boom\\x1b[2Jwiped\\x0doverwrite" in line
+
+
+def test_a_heartbeat_without_a_failure_cause_is_unchanged():
+    """The failure fields are additive: an ordinary step line must not grow an empty `error=`."""
+    from flash.providers._lifecycle.poll import _format_heartbeat
+
+    line = _format_heartbeat(
+        {"stage": "sft_step", "step": 12, "error": "   ", "checkpoint_failure": {}}
+    )
+
+    assert line == "worker: stage=sft_step step=12"
+
+
+def test_an_upload_failure_is_critical_so_an_in_flight_commit_cannot_drop_it():
+    """This heartbeat is raised from inside a long HF upload -- when a ping is otherwise dropped.
+
+    `_HB_UPLOAD_IN_FLIGHT` is set for the duration of that upload, and an unforced non-critical
+    stage returns not-due immediately. The failure is reported once and nothing restates it, so the
+    one case that produces this heartbeat is the case that would discard it.
+    """
+    from flash.engine.worker.io.heartbeat import _is_critical_stage, _is_terminal_stage
+
+    assert _is_critical_stage("checkpoint_upload_failed")
+    # criticality and checkpoint carry use one terminal predicate.
+    assert _is_terminal_stage("done")
+    assert _is_terminal_stage("error_oom")
+    assert not _is_terminal_stage("checkpoint_upload_failed")
+    assert _is_critical_stage("done")
+    assert _is_critical_stage("error_oom")
+    assert not _is_critical_stage("sft_step")
+    assert not _is_critical_stage("rl_step")
+
+
 def test_seed_training_rngs_initializes_all_supported_generators(monkeypatch):
     from flash.engine.worker.runtime import rng
 

{freesolo_flash_dev-1.2.95 → freesolo_flash_dev-1.2.96}/tests/test_worker_init_heartbeat.py

@@ -470,6 +470,73 @@ def _reset_console_heartbeat_state(monkeypatch, worker) -> None:
     monkeypatch.setattr(worker, "_HB_LAST_PROGRESS_TS", 0.0)
     monkeypatch.setattr(worker, "_HB_PROGRESS_SEQ", 0)
     monkeypatch.setattr(worker, "_HB_PROGRESS_UPLOADED_SEQ", 0)
+    monkeypatch.setattr(worker, "_HB_PENDING_CHECKPOINT_FAILURE", None)
+
+
+def _capture_heartbeat_payloads(monkeypatch, worker) -> list[dict]:
+    import json
+
+    committed: list[dict] = []
+
+    def capture(path, destination, **_kwargs):
+        if destination == "heartbeat.json":
+            with open(path) as file:
+                committed.append(json.load(file))
+        return True
+
+    _reset_console_heartbeat_state(monkeypatch, worker)
+    monkeypatch.setattr(worker, "hf_upload_file", capture)
+    return committed
+
+
+def test_fatal_heartbeat_preserves_checkpoint_failure_and_terminal_error(monkeypatch):
+    import flash.engine.worker as worker
+
+    committed = _capture_heartbeat_payloads(monkeypatch, worker)
+    failure = {"step": 50, "operation": "resume", "error": "quota denied"}
+
+    worker.heartbeat("checkpoint_upload_failed", step=50, checkpoint_failure=failure)
+    worker.heartbeat("error_sft", error="watcher failed")
+
+    assert committed[-1]["stage"] == "error_sft"
+    assert committed[-1]["error"] == "watcher failed"
+    assert committed[-1]["checkpoint_failure"] == failure
+
+
+def test_finalize_preserves_failed_checkpoint_identity(monkeypatch):
+    import flash.engine.worker as worker
+    from flash.engine.result.accounting import RunMetrics
+
+    committed = _capture_heartbeat_payloads(monkeypatch, worker)
+    monkeypatch.setattr(worker, "gpu_diagnostics", dict)
+    failure = {"step": 50, "operation": "resume", "error": "quota denied"}
+
+    worker.heartbeat("checkpoint_upload_failed", step=50, checkpoint_failure=failure)
+    worker._finalize(RunMetrics(phase="sft", step=100))
+
+    assert committed[-1]["stage"] == "done"
+    assert committed[-1]["step"] == 100
+    assert committed[-1]["checkpoint_failure"] == failure
+
+
+@pytest.mark.parametrize("terminal_stage", ["done", "error_sft"])
+def test_successful_checkpoint_clears_failure_before_any_terminal(monkeypatch, terminal_stage):
+    import flash.engine.worker as worker
+    from flash.engine.result.accounting import RunMetrics
+
+    committed = _capture_heartbeat_payloads(monkeypatch, worker)
+    failure = {"step": 50, "operation": "resume", "error": "quota denied"}
+
+    worker.heartbeat("checkpoint_upload_failed", step=50, checkpoint_failure=failure)
+    worker.heartbeat("checkpoint_uploaded", step=75)
+    if terminal_stage == "done":
+        monkeypatch.setattr(worker, "gpu_diagnostics", dict)
+        worker._finalize(RunMetrics(phase="sft", step=100))
+    else:
+        worker.heartbeat(terminal_stage, step=100, error="later fatal")
+
+    assert committed[-1]["stage"] == terminal_stage
+    assert "checkpoint_failure" not in committed[-1]
 
 
 @pytest.mark.parametrize(

{freesolo_flash_dev-1.2.95 → freesolo_flash_dev-1.2.96}/uv.lock

@@ -1052,7 +1052,7 @@ wheels = [
 
 [[package]]
 name = "freesolo-flash"
-version = "1.2.95"
+version = "1.2.96"
 source = { editable = "." }
 
 [package.optional-dependencies]