freesolo-flash-dev 1.2.58__tar.gz → 1.2.59__tar.gz

{freesolo_flash_dev-1.2.58 → freesolo_flash_dev-1.2.59}/CHANGELOG.md

@@ -21,6 +21,17 @@ This file starts at 1.1.35. Earlier releases are not reconstructed here; use
 
 ### Changed
 
+- A RunPod attempt's stall window is now measured from evidence timestamps rather than from when
+  the control plane happened to look at it. The window was seeded from the moment polling started
+  and was reset again by the first status read, and a heartbeat was credited at the time it was
+  read rather than the time it was written. Every one of those is a wall-clock reading, so each
+  reattach handed a worker that was already wedged a complete fresh setup window, and a supervisor
+  that reattached repeatedly could keep a dead worker on a paid GPU until the run's wall deadline.
+  The window is now anchored to the persisted launch, a heartbeat is credited at its own timestamp
+  clamped to that launch, and only a genuine status transition counts as progress. Instance
+  polling (Vast/Lambda) already worked this way; this brings RunPod onto the same rule. A job that
+  is still queued is unaffected: the capacity wait keeps its own exemption, so a worker granted
+  late still gets its full cold-start budget.
 - An attempt's terminal artifacts (its strict marker plus `metrics.json`) are now interpreted by
   one shared protocol instead of being decoded separately by live instance polling and by each
   recovery path. The two decoders disagreed about identical bytes. A marker that could not be tied

{freesolo_flash_dev-1.2.58 → freesolo_flash_dev-1.2.59}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.5
 Name: freesolo-flash-dev
-Version: 1.2.58
+Version: 1.2.59
 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.58 → freesolo_flash_dev-1.2.59}/flash/providers/runpod/job_execution.py

@@ -18,6 +18,7 @@ from flash.providers._lifecycle.deadline import (
 from flash.providers._lifecycle.poll import (
     PollErrorTracker,
     _attempt_int,
+    heartbeat_progress_ts,
     is_training_heartbeat,
     make_say,
     surface_heartbeat,
@@ -547,6 +548,20 @@ def _classify_queue_state(
     return None
 
 
+def _progress_ts(context: _PollContext, hb_key: tuple) -> float:
+    """The progress this heartbeat proves: its own ts, clamped to [launch, now].
+
+    The companion ``fresh`` flag is deliberately dropped. Its attempt half is already decided by the
+    caller, which returns early on any heartbeat that is not this attempt's, and its ts half is
+    already expressed in the clamp: a heartbeat stamped before launch cannot describe work this
+    attempt did, so it is worth exactly the launch anchor and no more. What remains is the
+    unusable-ts case, where the helper falls back to now; that matches the previous behaviour of
+    treating an untimestamped heartbeat as present-tense progress.
+    """
+    ts, _fresh = heartbeat_progress_ts(hb_key, context.launch_ts, context.current_attempt)
+    return ts
+
+
 def _update_heartbeat(context: _PollContext, state: _PollState) -> None:
     new_key, stage = surface_heartbeat(context.heartbeat_reader, state.last_hb_key, context.say)
     if new_key == state.last_hb_key:
@@ -575,17 +590,29 @@ def _update_heartbeat(context: _PollContext, state: _PollState) -> None:
     hb_ts = new_key[2]
     hb_step = new_key[1]
     is_training_hb = is_training_heartbeat(stage, hb_step)
+    # Credit the heartbeat's OWN timestamp, clamped to [launch, now], never the time we happened to
+    # read it. A heartbeat written long before a reattach and only seen now describes progress that
+    # already happened; paying it out at read time hands a wedged worker a full fresh stall window
+    # on every reattach. `_progress_ts` returns that clamped value, deliberately distinct from the
+    # raw `hb_ts` above, which stays the uploaded-ts bookkeeping the advance checks below compare.
+    # A heartbeat with no usable ts still clamps to now, preserving the existing behaviour that an
+    # untimestamped heartbeat counts as progress. It is computed only where it is used: this loop's
+    # clock is observable through `_jobs.time.time`, so an unused read would be a visible side effect.
+    #
+    # heartbeat ordering is independent of other progress sources. a queue exemption or status
+    # transition may already have advanced `last_progress` beyond this heartbeat's timestamp, so
+    # preserve the newer stall anchor while updating heartbeat-specific bookkeeping.
     if hb_attempt > state.last_hb_attempt:
         # fresh attempt: reset ts baseline and re-derive seen_training_hb so cold-start grace rearms.
         state.last_hb_attempt = hb_attempt
         state.last_hb_ts = hb_ts or 0.0
-        state.last_progress = _jobs.time.time()
+        state.last_progress = max(state.last_progress, _progress_ts(context, new_key))
         state.seen_training_hb = is_training_hb
     elif hb_attempt == state.last_hb_attempt and (hb_ts is None or hb_ts > state.last_hb_ts):
         # gate progress on ts advancing: a stale late upload must not buy a fresh stall window.
         if hb_ts is not None:
             state.last_hb_ts = hb_ts
-        state.last_progress = _jobs.time.time()
+        state.last_progress = max(state.last_progress, _progress_ts(context, new_key))
         if is_training_hb:
             state.seen_training_hb = True
 
@@ -702,7 +729,13 @@ def poll_job(
         last_hb_ts=0.0,
         # -1 sentinel < any real attempt; gates out prior-attempt leftover heartbeats
         last_hb_attempt=-1,
-        last_progress=_jobs.time.time(),
+        # Seed from the persisted LAUNCH, not this poll's start. A reattach has been billing since
+        # launch, so a worker already wedged before we attached must not be handed a fresh setup
+        # window by the mere act of looking at it -- that is what let a restart loop keep a dead
+        # worker alive to the absolute deadline. For a job queued from the start this value is
+        # overwritten anyway by the queued exemption in `_classify_stall`, which re-anchors on every
+        # pre-grant poll; it is the ALREADY-PLACED job at poll start where the difference is real.
+        last_progress=launch_ts,
         seen_training_hb=False,
         last_health_probe=0.0,
         unhealthy_timer=_jobs.GraceTimer(),
@@ -741,8 +774,13 @@ def poll_job(
             state.ever_saw_worker = True
         if status != state.last_status:
             say(f"job {handle.job_id}: {status}")
+            # A genuine TRANSITION is progress, so it moves the clock. The FIRST observation is not a
+            # transition -- `last_status` simply starts unset -- and crediting it would throw away the
+            # launch anchor on poll one, which is exactly how a delayed reattach used to win back a
+            # whole setup window before any evidence of progress had been seen.
+            if state.last_status is not None:
+                state.last_progress = _jobs.time.time()
             state.last_status = status
-            state.last_progress = _jobs.time.time()
         terminal = _classify_terminal_status(context, state, provider_status, status)
         if terminal is not None:
             return terminal

{freesolo_flash_dev-1.2.58 → freesolo_flash_dev-1.2.59}/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "freesolo-flash-dev"
-version = "1.2.58"
+version = "1.2.59"
 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.58"
+version = "1.2.59"
 
 [tool.pytest.ini_options]
 pythonpath = ["."]

{freesolo_flash_dev-1.2.58 → freesolo_flash_dev-1.2.59}/tests/test_jobs.py

@@ -7,6 +7,7 @@ import io
 import re
 import sys
 import tempfile
+import time
 from types import SimpleNamespace
 
 import pytest
@@ -16,6 +17,17 @@ from tests._helpers.runner import provisioned_status
 _RUNPOD_FINGERPRINT = "rpk-0123456789ab"
 
 
+def _live_clock_handle(jobs):
+    """A handle for a test that runs on the REAL clock rather than a faked one.
+
+    The poll loop measures its stall window from the persisted launch, so the fixed 1.0 default
+    below describes a job launched in 1970: fine for a test that also fakes `jobs.time.time` onto a
+    synthetic timeline, instantly "wedged" for one that does not. Tests asserting on the status
+    machine rather than on timing use this so their launch sits on the same clock they run against.
+    """
+    return _runpod_handle(jobs, started_ts=time.time())
+
+
 def _runpod_handle(
     jobs, endpoint_id="ep", endpoint_name="name", job_id="job", attempt=0, started_ts=1.0
 ):
@@ -184,7 +196,7 @@ def _poll(monkeypatch, statuses, heartbeats=None, stall_after_s=10.0):
     hb_iter = iter(heartbeats) if heartbeats is not None else None
 
     reader = (lambda force=False: next(hb_iter, None)) if hb_iter is not None else None
-    h = _runpod_handle(jobs)
+    h = _live_clock_handle(jobs)
     ok_payload = {"acc": 1.0}
     return jobs.poll_job(
         h,
@@ -235,7 +247,7 @@ def test_poll_job_surfaces_heartbeat_before_terminal_return(monkeypatch):
         return heartbeat
 
     res = jobs.poll_job(
-        _runpod_handle(jobs),
+        _live_clock_handle(jobs),
         interval_s=0,
         heartbeat_reader=read_heartbeat,
     )
@@ -314,7 +326,7 @@ def test_poll_job_failure_detail_redacts_secrets_in_provider_error_and_stdout(mo
     monkeypatch.setattr(jobs.time, "sleep", lambda s: None)
 
     res = jobs.poll_job(
-        _runpod_handle(jobs),
+        _live_clock_handle(jobs),
         interval_s=0,
         heartbeat_reader=lambda force=False: None,
     )
@@ -345,7 +357,7 @@ def test_poll_job_failure_surfaces_forced_heartbeat(monkeypatch):
     }
 
     res = jobs.poll_job(
-        _runpod_handle(jobs),
+        _live_clock_handle(jobs),
         interval_s=0,
         heartbeat_reader=lambda force=False: hb,
         log=log,
@@ -383,7 +395,7 @@ def test_poll_job_failure_appends_worker_artifacts(monkeypatch):
         )
 
     res = jobs.poll_job(
-        _runpod_handle(jobs),
+        _live_clock_handle(jobs),
         interval_s=0,
         heartbeat_reader=lambda force=False: None,
         failure_detail_reader=failure_detail_reader,
@@ -423,7 +435,7 @@ def test_poll_job_platform_preempt_does_not_read_worker_artifacts(monkeypatch):
         raise AssertionError("platform terminations should not read worker artifacts")
 
     res = jobs.poll_job(
-        _runpod_handle(jobs),
+        _live_clock_handle(jobs),
         interval_s=0,
         heartbeat_reader=lambda force=False: None,
         failure_detail_reader=failure_detail_reader,
@@ -437,12 +449,17 @@ def _poll_failed_with_heartbeat(monkeypatch, hb):
     from flash.providers.runpod import api as runpod_api
     from flash.providers.runpod import jobs
 
-    hb = {"ts": 2.0, "attempt": 0, **hb}
+    # The heartbeat has to sit on the same timeline as the launch: `worker_flagged_retriable` only
+    # honours a heartbeat stamped at or after the launch it claims to describe, so a fixed ts=2.0
+    # against a real-clock launch is a 1970 artifact that is correctly ignored, and the retriable
+    # flag under test would never be read at all.
+    launch_ts = time.time()
+    hb = {"ts": launch_ts + 1.0, "attempt": 0, **hb}
     seq = iter([{"status": "IN_PROGRESS"}, {"status": "FAILED", "error": "boom"}])
     monkeypatch.setattr(runpod_api, "job_status", lambda eid, jid, **_kw: next(seq))
     monkeypatch.setattr(jobs.time, "sleep", lambda s: None)
     return jobs.poll_job(
-        _runpod_handle(jobs),
+        _runpod_handle(jobs, started_ts=launch_ts),
         interval_s=0,
         heartbeat_reader=lambda force=False: hb,
         stall_after_s=10.0,
@@ -500,10 +517,17 @@ def test_poll_job_completed_decode_error_consults_worker_flags(monkeypatch):
     monkeypatch.setattr(
         runpod_api, "job_status", lambda eid, jid, **_kw: {"status": "COMPLETED", "output": bad}
     )
+    # the heartbeat shares the launch's timeline: a retriable flag stamped before launch is stale
+    # evidence and is correctly ignored, which would make this assert on the wrong reason.
+    launch_ts = time.time()
     res = jobs.poll_job(
-        _runpod_handle(jobs),
+        _runpod_handle(jobs, started_ts=launch_ts),
         interval_s=0,
-        heartbeat_reader=lambda force=False: {"retriable": True, "attempt": 0, "ts": 2.0},
+        heartbeat_reader=lambda force=False: {
+            "retriable": True,
+            "attempt": 0,
+            "ts": launch_ts + 1.0,
+        },
         failure_detail_reader=lambda force=False: "--- error_sft.txt ---\nCUDA out of memory",
     )
     assert res.failure == "job_preempted"
@@ -1861,7 +1885,7 @@ def test_poll_job_in_queue_then_progress_does_not_false_stall(monkeypatch):
         "endpoint_health_for_fingerprint",
         lambda eid, _fingerprint, **_kw: {"workers": {"ready": 1, "running": 1}},
     )
-    h = _runpod_handle(jobs)
+    h = _live_clock_handle(jobs)
     res = jobs.poll_job(h, interval_s=0, heartbeat_reader=lambda: None, queue_grace_s=900.0)
     assert res.ok
 
@@ -2219,6 +2243,176 @@ def test_poll_job_stale_late_heartbeat_does_not_reset_progress(monkeypatch):
     assert fresh_run > none_run, "a genuinely newer heartbeat does reset progress (stalls later)"
 
 
+@pytest.mark.parametrize(
+    ("last_hb_attempt", "last_hb_ts"),
+    [(-1, 0.0), (0, 50.0)],
+    ids=["first-current-attempt-heartbeat", "later-current-attempt-heartbeat"],
+)
+def test_older_heartbeat_cannot_regress_status_progress_anchor(
+    monkeypatch, last_hb_attempt, last_hb_ts
+):
+    # a queue exemption or status transition can advance the shared progress anchor beyond the
+    # worker timestamp. a heartbeat that becomes visible later may still advance heartbeat-specific
+    # bookkeeping, but neither credit branch may move the stall anchor backward.
+    from flash.providers.runpod import jobs
+
+    job_execution = sys.modules["flash.providers.runpod.job_execution"]
+    heartbeat_key = ("boot", None, 100.0, 0)
+    monkeypatch.setattr(
+        job_execution,
+        "surface_heartbeat",
+        lambda _reader, _last_key, _say: (heartbeat_key, "boot"),
+    )
+    monkeypatch.setattr(jobs.time, "time", lambda: 600.0)
+    context = SimpleNamespace(
+        heartbeat_reader=lambda: None,
+        say=lambda _message: None,
+        current_attempt=0,
+        launch_ts=1.0,
+    )
+    state = SimpleNamespace(
+        last_hb_key=None,
+        last_hb_ts=last_hb_ts,
+        last_hb_attempt=last_hb_attempt,
+        last_progress=500.0,
+        seen_training_hb=False,
+        ever_saw_worker=True,
+    )
+
+    job_execution._update_heartbeat(context, state)
+
+    assert state.last_progress == 500.0
+
+
+def _stall_clock_at_giveup(
+    monkeypatch, *, started_ts, heartbeats=(), step_s=100.0, clock_at=5000.0
+):
+    """Poll an ALREADY-PLACED job and return the absolute fake-clock time it gave up at.
+
+    IN_PROGRESS from the first read, so the queued-capacity exemption never applies and the value
+    seeded into ``last_progress`` is the only thing anchoring the stall window. ``clock_at`` is where
+    this poll's wall clock starts; a ``started_ts`` well below it is a launch in the past, i.e. a
+    delayed reattach. The handle rejects a non-positive launch, so the clock starts high rather than
+    the launch going negative.
+    """
+    from flash.providers.runpod import api as runpod_api
+    from flash.providers.runpod import jobs
+
+    monkeypatch.setattr(runpod_api, "job_status", lambda eid, jid, **_kw: {"status": "IN_PROGRESS"})
+    monkeypatch.setattr(jobs.time, "sleep", lambda s: None)
+    state = {"t": clock_at}
+
+    def _time():
+        state["t"] += step_s
+        return state["t"]
+
+    monkeypatch.setattr(jobs.time, "time", _time)
+    hbs = iter(list(heartbeats))
+    res = jobs.poll_job(
+        _runpod_handle(jobs, started_ts=started_ts),
+        interval_s=0,
+        heartbeat_reader=lambda: next(hbs, None),
+        stall_after_s=150.0,
+        setup_grace_s=600.0,
+    )
+    assert res.failure == "stalled", res
+    return state["t"]
+
+
+def test_a_reattached_wedged_worker_is_measured_from_launch_not_from_this_polls_start(monkeypatch):
+    # A reattach has been BILLING since launch, so the stall window has to be measured from the
+    # persisted launch timestamp. Seeding it from `time.time()` at poll start meant every reattach
+    # handed a worker that was already wedged a complete fresh setup window, and a supervisor that
+    # reattaches repeatedly could keep a dead worker on a paid GPU until the run's wall deadline.
+    #
+    # The job is IN_PROGRESS from the first read, so the queued-capacity exemption -- which re-anchors
+    # `last_progress` on every pre-grant poll, and is what made the wrong seed survivable for a job
+    # queued from the start -- never fires here. This is the placed-at-attach case where the seed is
+    # the only anchor.
+    #
+    # Assert on the absolute clock at give-up, not on the "no worker progress for Ns" figure in the
+    # detail: that figure is derived from `last_progress` and reads ~the same either way, so it
+    # cannot tell the fixed and broken code apart.
+    # attached at launch: the setup grace has not been spent, so the full window remains.
+    at_launch = _stall_clock_at_giveup(monkeypatch, started_ts=5000.0)
+    # same wedge, but launched 5000s before this poll attached: that grace is already gone.
+    long_after_launch = _stall_clock_at_giveup(monkeypatch, started_ts=1.0)
+    assert long_after_launch < at_launch, (
+        f"a worker wedged since a launch 5000s ago was torn down at {long_after_launch}s, no sooner "
+        f"than one attached at launch ({at_launch}s); the stall clock anchors on poll start, not launch"
+    )
+
+
+def test_a_heartbeat_written_before_the_reattach_buys_no_fresh_stall_window(monkeypatch):
+    # Crediting progress at READ time means a heartbeat written long ago and surfaced now pays out a
+    # full fresh window, so a worker that stopped producing heartbeats before we attached looks alive
+    # for one more grace period. Credit the heartbeat's OWN ts (clamped to [launch, now]) instead:
+    # it describes when the worker made progress, not when we got around to looking.
+    #
+    # Both runs see exactly ONE staged heartbeat, so the only difference is its timestamp.
+    stale = _stall_clock_at_giveup(
+        monkeypatch,
+        started_ts=5000.0,
+        heartbeats=[{"stage": "train", "step": 1, "ts": 5001.0, "attempt": 0}],
+    )
+    current = _stall_clock_at_giveup(
+        monkeypatch,
+        started_ts=5000.0,
+        heartbeats=[{"stage": "train", "step": 1, "ts": 5400.0, "attempt": 0}],
+    )
+    assert stale < current, (
+        f"a heartbeat stamped at launch (ts=5001.0) held the run to {stale}s, as long as one stamped "
+        f"at ts=5400.0 ({current}s); progress is credited at read time, not at the heartbeat's own ts"
+    )
+
+
+def test_a_heartbeat_from_a_prior_attempt_buys_this_attempt_no_progress(monkeypatch):
+    # A retry reuses the same heartbeat path, so a leftover heartbeat from attempt 0 is still
+    # readable while attempt 1 is running. Crediting it would let the PREVIOUS attempt's work stand
+    # as proof that THIS one is progressing. The attempt identity has to gate the credit, not just
+    # the OOM flags.
+    #
+    # Both runs are attempt 1 and see exactly one heartbeat; only its `attempt` field differs. The
+    # own-attempt run gives up SOONER, which is the tell: its heartbeat is a training heartbeat, so
+    # it both moves the clock and tightens the window from setup grace to the stall limit. The
+    # prior-attempt heartbeat must do neither, leaving the wider setup window running from launch.
+    # Asserting "stale lasts longer" therefore proves the heartbeat was ignored end to end; the
+    # reverse ordering would mean prior-attempt evidence had been credited.
+    from flash.providers.runpod import api as runpod_api
+    from flash.providers.runpod import jobs
+
+    def _giveup_with(hb_attempt):
+        monkeypatch.setattr(
+            runpod_api, "job_status", lambda eid, jid, **_kw: {"status": "IN_PROGRESS"}
+        )
+        monkeypatch.setattr(jobs.time, "sleep", lambda s: None)
+        state = {"t": 5000.0}
+
+        def _time():
+            state["t"] += 100.0
+            return state["t"]
+
+        monkeypatch.setattr(jobs.time, "time", _time)
+        hbs = iter([{"stage": "train", "step": 1, "ts": 5400.0, "attempt": hb_attempt}])
+        res = jobs.poll_job(
+            _runpod_handle(jobs, attempt=1, started_ts=5000.0),
+            interval_s=0,
+            heartbeat_reader=lambda: next(hbs, None),
+            stall_after_s=150.0,
+            setup_grace_s=600.0,
+            current_attempt=1,
+        )
+        assert res.failure == "stalled", res
+        return state["t"]
+
+    stale_attempt = _giveup_with(0)
+    this_attempt = _giveup_with(1)
+    assert stale_attempt > this_attempt, (
+        f"attempt 1 given only a stale attempt=0 heartbeat behaved identically to one given its "
+        f"own ({stale_attempt}s vs {this_attempt}s); prior-attempt evidence is being credited"
+    )
+
+
 def test_poll_job_gapfill_step0_does_not_tighten(monkeypatch):
     # The train-liveness gap-filler emits rl_step/sft_step at step=0 throughout the silent FIRST step
     # (a cold vLLM rollout can run many minutes before global_step ticks to 1). That step=0 ping is a
@@ -6650,7 +6844,7 @@ def test_runpod_completed_metrics_probes_after_expired_recovery_grace(monkeypatc
 
     monkeypatch.setattr(runpod_api, "job_status", completed_status)
     metrics = lifecycle._runpod_completed_metrics(
-        _runpod_handle(jobs),
+        _live_clock_handle(jobs),
         deadline_at=now - lifecycle._RECOVERY_MARKER_GRACE_S - 1.0,
     )
 
@@ -6678,7 +6872,7 @@ def test_runpod_completed_metrics_caps_probe_deadline_when_wall_deadline_far_fut
 
     monkeypatch.setattr(runpod_api, "job_status", completed_status)
     metrics = lifecycle._runpod_completed_metrics(
-        _runpod_handle(jobs),
+        _live_clock_handle(jobs),
         deadline_at=now + 7_200.0,  # wall deadline hours in the future
     )
 

{freesolo_flash_dev-1.2.58 → freesolo_flash_dev-1.2.59}/uv.lock

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