diff --git a/docs/uk_replay/RECIPE.md b/docs/uk_replay/RECIPE.md index 380aa25..dcfe915 100644 --- a/docs/uk_replay/RECIPE.md +++ b/docs/uk_replay/RECIPE.md @@ -598,10 +598,51 @@ fields and head effects. `--allow-changed-existing` permits replacement only after that report has been reviewed; original bytes remain in the backup. `--allow-partial` permits read-only inspection of a focused download, but partial or noncanonical outputs cannot be adopted. Keep raw receipts for -provenance. Each Modal invocation uses one replay worker. Count every -active invocation and local simulation toward the overall maximum of two; -separate invocations do not share a global concurrency gate. A full event -plus one focused repeat uses both slots. Remote wall time is recorded +provenance. Each Modal invocation uses one replay worker and first takes one of two +simulation slots in the Modal Dict `scorecard-replay-simulation-slots` +(`modal_uk_event.SimulationSlots`), waiting while both are held, so Modal +invocations never run more than two simulations at once. How the slots hold: + +* A slot has one chain of claims per calendar day, + `claim-{slot}-{day}-{generation}`. A slot changes hands only through + `put(..., skip_if_exists=True)`, Modal's atomic create-if-absent: a caller + creates a chain's next generation only when its newest one is released or + lapsed. Nothing is deleted, there is no pointer that a late write could + roll back, and a holder writes only its own claims, to release or shorten + them. A lease that crosses midnight holds a claim in both days' chains. +* A day's chain matters for at most a day plus a lease, far inside the 7 days + after which Modal expires an idle entry, so no claim that still matters + can have expired, and a write that arrives days late lands in a chain that + no longer governs anything. The guarantee needs no assumption about how + long a step stalls. +* Leases run on Modal's clock: each invocation measures its offset from + Modal with a small remote function, checks it again after waiting for a + slot (a drift over 30 seconds stops the run, exit 75), and takes over a + lapsed claim only 2 minutes after its expiry. +* The lease is 4 hours. The remote function refuses to start unless its + whole 3-hour timeout plus a 10-minute margin fits inside it, so a run + ends before its slot can pass on, even if the client dies. +* A finished or never-started run frees its slot at once. When `app.run()` + exits after ctrl-C, Modal's client has already asked the server to stop + the app's tasks, and the slot frees 10 minutes later. Anything else after + the run started (a lost connection, which Modal's runner re-raises + without that request; a remote error; a failed stop request; a client + killed outright) leaves the claim until the lease lapses. A failed + release never costs a returned result: it is downloaded and the claim + lapses. +* The slot is taken with the app already open, after the image build, so a + build never eats into a lease; an invocation waits for a slot (polling + every 30 seconds, in no set order) with its app open and idle. +* `tests/test_modal_uk_event.py` checks this call by call: Hypothesis + interleaves several holders' Dict calls, each on its own clock within the + enforced drift, with clock jumps of any size at any point, midnight + crossings, crashes, interrupts, lost connections and Modal's 7-day expiry, + and asserts that no two runs overlap on a slot and that both slots can be + taken once every lease has lapsed. + +Local simulations are not in that lease: count them yourself toward the +overall maximum of two. A full event plus one focused repeat uses both +slots. Remote wall time is recorded separately in `MODAL_RECEIPT.json`; use the completed event's observed timings when planning the next lane. diff --git a/pipeline/modal_uk_event.py b/pipeline/modal_uk_event.py index d5f4361..66f1d33 100644 --- a/pipeline/modal_uk_event.py +++ b/pipeline/modal_uk_event.py @@ -40,6 +40,7 @@ import sys import tempfile import time +from dataclasses import dataclass from pathlib import Path, PurePosixPath ROOT = Path(__file__).resolve().parent.parent @@ -100,6 +101,220 @@ } +SLOTS_DICT = "scorecard-replay-simulation-slots" +SLOT_COUNT = 2 +# The function timeout (10,800 s) plus the guard margin and time for the +# container to be scheduled. The remote function refuses to start unless its +# whole timeout and the margin fit inside the lease, so a run always ends +# before its slot can pass to anyone else. +LEASE_SECONDS = 4 * 60 * 60 +# The remote guard leaves this much of the lease unused. +GUARD_MARGIN_SECONDS = 10 * 60 +# Claims are grouped by calendar day (UTC, on Modal's clock). A day's claims +# matter for at most EPOCH_SECONDS + LEASE_SECONDS, far inside the 7 days +# after which Modal expires an idle Dict entry, so no claim that still +# matters can have expired. +EPOCH_SECONDS = 24 * 60 * 60 +# Every invocation measures its offset from Modal's clock and uses Modal +# time for leases. A claim is treated as lapsed only this long after its +# expiry, which covers the measurement error and drift of two invocations. +CLOCK_MARGIN_SECONDS = 2 * 60 +MAX_CLOCK_DRIFT_SECONDS = 30 +# When app.run() exits after the runner caught ctrl-C, it has already asked +# the server to stop the ephemeral app's tasks (modal/runner.py, +# _status_based_disconnect); the slot frees this long afterwards. +INTERRUPT_GRACE_SECONDS = 10 * 60 + + +class LeaseLost(RuntimeError): + """This invocation cannot safely hold a simulation slot.""" + + +@dataclass(frozen=True) +class Lease: + slot: int + expires: float + claims: tuple # ((epoch, generation), ...), ascending by epoch + + +class SimulationSlots: + """At most two replay simulations at once across Modal invocations. + + Each invocation runs one worker, and separate invocations share no + Modal-side limit, so a third concurrent run would exceed the two the + shared account allows. Before starting, an invocation takes one of two + slots in a named Modal Dict. + + A slot has one chain of claims per calendar day (epoch): + ``claim-{slot}-{epoch}-{generation}``. The only operation that takes a + slot is ``put(key, value, skip_if_exists=True)``, an atomic + create-if-absent on the server (modal/dict.py, + ``DictUpdateRequest(if_not_exists=...)``): a caller creates a chain's + next generation only when its newest one is released or lapsed, and + exactly one caller creates any generation. Nothing is deleted, there is + no pointer to go stale, and a holder writes only its own claims, to + release or shorten them. A lease that crosses midnight holds a claim in + both days' chains, taken in order and given back together. + + Why this is safe: + + * Within a chain, generations are contiguous and each newer one was + created only after the one before it was free, so at most one claim + is live at a time. A run at time t holds the live claim in t's chain. + * A chain governs only its own day, and its keys are at most + EPOCH_SECONDS + LEASE_SECONDS old while it does, so Modal's expiry of + entries idle for 7 days cannot remove one that still matters. A write + that arrives late lands in a chain that no longer governs anything. + * A step that stalls, however long, either finds its generation taken + (the create fails) or was unopposed all along; a lease is granted + only with enough time left for a run, on the clock read after the + last write. + * Clocks: ``clock`` must be Modal's time (main() measures the offset). + Takeover waits CLOCK_MARGIN_SECONDS past a claim's expiry, and the + remote guard ends every run GUARD_MARGIN_SECONDS before it. + + Every Dict call is one step of a generator (``_try_steps``, + ``_release_steps``), so tests interleave several holders call by call. + """ + + def __init__( + self, + store, + *, + holder, + count=SLOT_COUNT, + lease=LEASE_SECONDS, + clock=time.time, + minimum=None, + ): + self.store, self.holder, self.count = store, holder, count + self.lease, self.clock = lease, clock + # A lease with less than this left is not worth holding: the remote + # guard would refuse the run. + self.minimum = ( + RESOURCE_LIMITS["timeout_seconds"] + GUARD_MARGIN_SECONDS + 60 + if minimum is None + else minimum + ) + + @staticmethod + def claim_key(slot: int, epoch: int, generation: int) -> str: + return f"claim-{slot}-{epoch}-{generation}" + + @staticmethod + def epochs(start: float, expires: float) -> range: + """The days a lease touches, padded by the clock margin.""" + first = int((start - CLOCK_MARGIN_SECONDS) // EPOCH_SECONDS) + last = int((expires + CLOCK_MARGIN_SECONDS) // EPOCH_SECONDS) + return range(first, last + 1) + + def _claim(self, expires: float, *, released: bool = False) -> dict: + return {"holder": self.holder, "expires": expires, "released": released} + + def _give_back(self, slot: int, claims, expires: float): + for epoch, generation in claims: + key = self.claim_key(slot, epoch, generation) + yield ("put", key, self._claim(expires, released=True), False) + + def _try_steps(self, slot: int): + expires = self.clock() + self.lease + made = [] + try: + for epoch in self.epochs(expires - self.lease, expires): + generation, newest = -1, None + while True: + key = self.claim_key(slot, epoch, generation + 1) + following = yield ("get", key) + if following is None: + break + generation, newest = generation + 1, following + free = ( + newest is None + or newest["released"] + or newest["expires"] + CLOCK_MARGIN_SECONDS <= self.clock() + ) + key = self.claim_key(slot, epoch, generation + 1) + if not free or not (yield ("put", key, self._claim(expires), True)): + yield from self._give_back(slot, made, expires) + return None + made.append((epoch, generation + 1)) + except Exception: + # A failed call: hold nothing rather than a slot with no run. + yield from self._give_back(slot, made, expires) + raise + if expires - self.clock() < self.minimum: + # The step itself took too long for a run to fit in the lease. + yield from self._give_back(slot, made, expires) + return None + return Lease(slot, expires, tuple(made)) + + def _release_steps(self, lease: Lease, *, grace: float = 0.0): + for epoch, generation in lease.claims: + key = self.claim_key(lease.slot, epoch, generation) + if grace: + claim = self._claim(min(lease.expires, self.clock() + grace)) + else: + claim = self._claim(lease.expires, released=True) + yield ("put", key, claim, False) + + def _drive(self, steps): + """Run one protocol generator against the store.""" + try: + call = next(steps) + while True: + try: + if call[0] == "get": + result = self.store.get(call[1]) + else: + result = self.store.put( + call[1], call[2], skip_if_exists=call[3] + ) + except Exception as error: + # Let the step see a failed call, so it can give back + # the claims it just made. + call = steps.throw(error) + continue + call = steps.send(result) + except StopIteration as done: + return done.value + + def acquire( + self, + *, + wait_seconds: float = 6 * 60 * 60, + poll_seconds: float = 30, + sleep=time.sleep, + ) -> Lease: + deadline = self.clock() + wait_seconds + while True: + for slot in range(self.count): + lease = self._drive(self._try_steps(slot)) + if lease is not None: + return lease + if self.clock() >= deadline: + raise TimeoutError("no simulation slot became free") + print("[modal] both simulation slots are held; waiting", flush=True) + sleep(poll_seconds) + + def release(self, lease: Lease, *, grace: float = 0.0) -> None: + """Free the slot now, or ``grace`` seconds from now when the remote + run may still be stopping.""" + self._drive(self._release_steps(lease, grace=grace)) + + +def clock_offset(remote_now, *, local=time.time, attempts: int = 4) -> float: + """Modal's clock minus this machine's, from a remote call bracketed by + two local readings; the first call may include a cold start, so a slow + round trip is retried.""" + for _ in range(attempts): + before = local() + remote = remote_now() + after = local() + if after - before <= 10: + return remote - (before + after) / 2 + raise LeaseLost("could not measure the clock offset from Modal") + + def canonical_bytes(value: object) -> bytes: return ( json.dumps(value, indent=1, sort_keys=True, allow_nan=False) + "\n" @@ -472,6 +687,8 @@ def modal_function(mounts: dict[str, Path]): for relative, path in sorted(mounts.items()): image = image.add_local_file(path, f"{REMOTE_ROOT}/{relative}", copy=False) app = modal.App("scorecard-uk-event-replay", include_source=False) + timeout_seconds = RESOURCE_LIMITS["timeout_seconds"] + guard_seconds = timeout_seconds + GUARD_MARGIN_SECONDS @app.function( image=image, @@ -479,19 +696,23 @@ def modal_function(mounts: dict[str, Path]): include_source=False, cpu=(2, 2), memory=(32768, 32768), - timeout=10800, + timeout=timeout_seconds, max_containers=1, retries=0, block_network=True, restrict_modal_access=True, ) @modal.concurrent(max_inputs=1) - def replay(payload: dict) -> dict: + def replay(payload: dict, lease_expires: float) -> dict: import hashlib import json import subprocess + import time from pathlib import Path + # A run that could outlive its simulation slot does not start. + if time.time() + guard_seconds > lease_expires: + raise RuntimeError("the simulation slot lease does not cover the timeout") # Inherit stdout/stderr so Modal streams phase progress as it happens. subprocess.run( [ @@ -525,7 +746,20 @@ def replay(payload: dict) -> dict: } return {"files": files} - return modal, app, replay + @app.function( + image=modal.Image.debian_slim(python_version="3.12"), + serialized=True, + include_source=False, + timeout=60, + block_network=True, + restrict_modal_access=True, + ) + def modal_time() -> float: + import time + + return time.time() + + return modal, app, replay, modal_time def download_prefix(download_root: str, *, root: Path = ROOT) -> Path: @@ -639,10 +873,56 @@ def main(argv: list[str] | None = None) -> int: raise ValueError("download destination escapes its ignored prefix") if target.exists() and any(target.iterdir()): raise ValueError("choose an empty download prefix before a remote run") - modal, app, replay = modal_function(mounts) + modal, app, replay, modal_time = modal_function(mounts) response = None - with modal.enable_output(), app.run(): - response = replay.remote(payload) + slots = lease = None + started = exited = False + try: + with modal.enable_output(), app.run(): + # Leases run on Modal's clock, so two invocations on machines + # whose clocks disagree still agree on when a claim lapses. + offset = clock_offset(modal_time.remote) + slots = SimulationSlots( + modal.Dict.from_name(SLOTS_DICT, create_if_missing=True), + holder=f"{platform.node()}:{os.getpid()}:{digest_bytes(canonical_bytes(request))[:12]}", + clock=lambda: time.time() + offset, + ) + lease = slots.acquire() + # Waiting for a slot can take hours: check the clock again. + if abs(clock_offset(modal_time.remote) - offset) > MAX_CLOCK_DRIFT_SECONDS: + raise LeaseLost("this machine's clock moved against Modal's") + started = True + response = replay.remote(payload, lease.expires) + # app.run() returned: the run finished, or the runner caught ctrl-C + # after asking the server to stop the app's tasks. + exited = True + except LeaseLost as error: + print(f"Modal run not started: {error}", file=sys.stderr) + return 75 + except Exception as error: + if response is None: + raise + # The run returned its result; only leaving the app failed. + print(f"Modal app exit failed after the run returned: {error}", file=sys.stderr) + finally: + # A finished or never-started run frees its slot now; a caught + # ctrl-C frees it after the grace. Otherwise (a lost connection, + # which the runner re-raises without a stop request; a remote error; + # an exception during the stop request) the remote run may still be + # going, and the claim stays until the lease lapses, which the + # remote guard puts after the run's end. A failed release never + # costs a result: the claim simply lapses. + if lease is not None: + try: + if not started or response is not None: + slots.release(lease) + elif exited: + slots.release(lease, grace=INTERRUPT_GRACE_SECONDS) + except Exception as error: + print( + f"simulation slot not released ({error}); it will lapse", + file=sys.stderr, + ) if response is None: print("Modal run interrupted; no output downloaded", file=sys.stderr) return 130 diff --git a/tests/test_modal_uk_event.py b/tests/test_modal_uk_event.py index 4a75e47..ed363e9 100644 --- a/tests/test_modal_uk_event.py +++ b/tests/test_modal_uk_event.py @@ -1,13 +1,18 @@ """Offline Modal transport checks; no Modal installation or country imports.""" import hashlib +import itertools import json +import subprocess # noqa: F401 - patched in the factory test import sys +import time from contextlib import nullcontext from pathlib import Path from types import SimpleNamespace import pytest +from hypothesis import given, settings +from hypothesis import strategies as st from pipeline import modal_uk_event as runner @@ -262,7 +267,7 @@ def __enter__(self): def __exit__(self, kind, value, traceback): return kind is KeyboardInterrupt - def interrupted_remote(payload): + def interrupted_remote(payload, lease_expires): raise KeyboardInterrupt monkeypatch.setattr(runner, "validate_selection", lambda request: None) @@ -270,13 +275,18 @@ def interrupted_remote(payload): runner, "input_manifest", lambda event, hf_cache=None: ({"event": event}, {}) ) monkeypatch.setattr(runner, "download_prefix", lambda value: tmp_path / "download") + slots = FakeDict() monkeypatch.setattr( runner, "modal_function", lambda mounts: ( - SimpleNamespace(enable_output=nullcontext), + SimpleNamespace( + enable_output=nullcontext, + Dict=SimpleNamespace(from_name=lambda name, create_if_missing: slots), + ), SimpleNamespace(run=SuppressInterrupt), SimpleNamespace(remote=interrupted_remote), + SimpleNamespace(remote=time.time), ), ) monkeypatch.setattr( @@ -286,6 +296,13 @@ def interrupted_remote(payload): ) assert runner.main(["--event", "autumn_budget_2024", "--execute"]) == 130 assert "interrupted; no output downloaded" in capsys.readouterr().err + # The interrupted run's slot frees once Modal has had time to stop it. + claims = [v for k, v in slots.data.items() if k.startswith("claim-")] + assert claims and all(claim["released"] is False for claim in claims) + assert all( + claim["expires"] <= time.time() + runner.INTERRUPT_GRACE_SECONDS + 1 + for claim in claims + ) def test_modal_factory_has_one_blocked_network_worker_and_explicit_uv_venv( @@ -321,7 +338,7 @@ def __init__(self, name, *, include_source): assert include_source is False def function(self, **kwargs): - recorded["function"] = kwargs + recorded.setdefault("functions", []).append(kwargs) return lambda fn: fn def concurrent(**kwargs): @@ -333,8 +350,21 @@ def concurrent(**kwargs): monkeypatch.setattr( runner.importlib.metadata, "version", lambda name: runner.MODAL_VERSION ) - runner.modal_function(mounts) - configuration = recorded["function"] + _, _, replay, modal_time = runner.modal_function(mounts) + assert abs(modal_time() - time.time()) < 5 + # The remote refuses to start unless its whole timeout fits in the lease. + monkeypatch.setattr( + sys.modules["subprocess"], + "run", + lambda *args, **kwargs: pytest.fail("a short lease must not start a run"), + ) + timeout = runner.RESOURCE_LIMITS["timeout_seconds"] + for slack in (-60, runner.GUARD_MARGIN_SECONDS - 60): + with pytest.raises(RuntimeError, match="lease does not cover"): + replay({"request": {}}, time.time() + timeout + slack) + configuration, clock_configuration = recorded["functions"] + assert clock_configuration["block_network"] is True + assert clock_configuration["timeout"] == 60 assert configuration["cpu"] == (2, 2) assert configuration["memory"] == (32768, 32768) assert ( @@ -359,3 +389,596 @@ def concurrent(**kwargs): for key in recorded["env"] ) assert len(recorded["mounts"]) == len(mounts) + 1 # build-time frozen requirements + + +class FakeDict: + """A Modal Dict's get/put, with put(skip_if_exists=True) atomic and, given + a clock, entries expiring after 7 days without a read or write (as + modal/dict.py documents).""" + + TTL = 7 * 24 * 60 * 60 + + def __init__(self, clock=None): + self.data, self.touched, self.clock = {}, {}, clock + + def _expire(self): + if self.clock is None: + return + for key in [k for k, t in self.touched.items() if self.clock() - t > self.TTL]: + del self.data[key], self.touched[key] + + def _touch(self, key): + if self.clock is not None and key in self.data: + self.touched[key] = self.clock() + + def put(self, key, value, *, skip_if_exists=False): + self._expire() + if skip_if_exists and key in self.data: + self._touch(key) + return False + self.data[key] = value + self._touch(key) + return True + + def get(self, key, default=None): + self._expire() + self._touch(key) + return self.data.get(key, default) + + +class Clock: + """True (Modal) time. Starts mid-morning, well inside a day.""" + + def __init__(self, now=10 * 3600.0): + self.now = now + + def __call__(self): + return self.now + + +DAY = runner.EPOCH_SECONDS +TIMEOUT = runner.RESOURCE_LIMITS["timeout_seconds"] +MARGIN = runner.CLOCK_MARGIN_SECONDS + + +def slots_for(store, clock, holder, *, offset=0.0, **kwargs): + """A holder whose own clock reads true time plus ``offset``.""" + return runner.SimulationSlots( + store, holder=holder, clock=lambda: clock() + offset, **kwargs + ) + + +def run_steps(store, steps, *, limit=None): + """Drive a protocol generator for at most ``limit`` Dict calls; return + (finished, value).""" + calls = 0 + try: + call = next(steps) + while True: + if limit is not None and calls == limit: + return False, None + calls += 1 + if call[0] == "get": + result = store.get(call[1]) + else: + result = store.put(call[1], call[2], skip_if_exists=call[3]) + call = steps.send(result) + except StopIteration as done: + return True, done.value + + +def holders_of(store, clock): + """Each slot's live claim in today's chain, by true time.""" + held = {} + today = int(clock() // DAY) + for slot in range(runner.SLOT_COUNT): + generation, newest = -1, None + while True: + claim = store.data.get(f"claim-{slot}-{today}-{generation + 1}") + if claim is None: + break + generation, newest = generation + 1, claim + if newest and not newest["released"] and newest["expires"] > clock(): + held[slot] = newest["holder"] + return held + + +def test_the_constants_leave_room_for_a_run_and_for_modal_expiry(): + assert runner.LEASE_SECONDS >= TIMEOUT + runner.GUARD_MARGIN_SECONDS + 30 * 60 + assert runner.INTERRUPT_GRACE_SECONDS < runner.LEASE_SECONDS + # A day's claims matter for a day plus a lease: far inside Modal's expiry. + assert runner.EPOCH_SECONDS + 2 * runner.LEASE_SECONDS < FakeDict.TTL / 2 + # Takeover waits out the clock error of two invocations. + assert runner.CLOCK_MARGIN_SECONDS >= 2 * runner.MAX_CLOCK_DRIFT_SECONDS + 20 + assert runner.GUARD_MARGIN_SECONDS > runner.CLOCK_MARGIN_SECONDS + + +def test_two_slots_then_a_third_run_waits(): + clock = Clock() + store = FakeDict(clock) + a, b, c = (slots_for(store, clock, h) for h in "abc") + leases = {"a": a.acquire(), "b": b.acquire()} + assert {lease.slot for lease in leases.values()} == {0, 1} + sleeps = [] + + def sleep(seconds): + sleeps.append(seconds) + clock.now += seconds + + with pytest.raises(TimeoutError): + c.acquire(wait_seconds=60, poll_seconds=30, sleep=sleep) + assert sleeps == [30, 30] + a.release(leases["a"]) + assert c.acquire(wait_seconds=0, sleep=sleep).slot == leases["a"].slot + assert holders_of(store, clock) == { + leases["a"].slot: "c", + leases["b"].slot: "b", + } + + +def test_a_lapsed_claim_is_taken_only_after_the_clock_margin(): + clock = Clock() + store = FakeDict(clock) + dead = slots_for(store, clock, "dead", count=1).acquire() + nxt = slots_for(store, clock, "next", count=1) + clock.now = dead.expires + MARGIN - 1 + assert run_steps(store, nxt._try_steps(0)) == (True, None) + clock.now = dead.expires + MARGIN + assert run_steps(store, nxt._try_steps(0))[1] is not None + + +def test_a_stale_release_cannot_free_a_newer_holder(): + clock = Clock() + store = FakeDict(clock) + late = slots_for(store, clock, "late", count=1) + old = late.acquire() + clock.now = old.expires + MARGIN + slots_for(store, clock, "successor", count=1).acquire() + late.release(old) + assert holders_of(store, clock) == {0: "successor"} + third = slots_for(store, clock, "third", count=1) + assert run_steps(store, third._try_steps(0)) == (True, None) + + +def test_a_contender_cannot_displace_a_claim_made_after_it_looked(): + clock = Clock() + store = FakeDict(clock) + holder = slots_for(store, clock, "holder", count=1) + lease = holder.acquire() + a = slots_for(store, clock, "a", count=1) + c = slots_for(store, clock, "c", count=1) + clock.now = lease.expires + MARGIN + steps = a._try_steps(0) + call = next(steps) + while call[0] == "get": # A reads up to its create + call = steps.send(store.get(call[1])) + assert run_steps(store, c._try_steps(0))[1] is not None + with pytest.raises(StopIteration) as done: # A's create now fails + steps.send(store.put(call[1], call[2], skip_if_exists=call[3])) + assert done.value.value is None + assert holders_of(store, clock) == {0: "c"} + + +@pytest.mark.parametrize("calls", range(6)) +def test_a_holder_that_dies_at_any_step_frees_its_slot_after_the_lease(calls): + clock = Clock() + store = FakeDict(clock) + dying = slots_for(store, clock, "dying", count=1) + run_steps(store, dying._try_steps(0), limit=calls) + clock.now += runner.LEASE_SECONDS + MARGIN + assert slots_for(store, clock, "next", count=1).acquire(wait_seconds=0).slot == 0 + + +def test_a_lease_across_midnight_holds_both_days(): + clock = Clock(now=22 * 3600.0) # two hours before the day ends + store = FakeDict(clock) + night = slots_for(store, clock, "night", count=1) + lease = night.acquire() + assert [epoch for epoch, _ in lease.claims] == [0, 1] + clock.now = DAY + 600 # the next day, the run still going + assert holders_of(store, clock) == {0: "night"} + morning = slots_for(store, clock, "morning", count=1) + assert run_steps(store, morning._try_steps(0)) == (True, None) + night.release(lease) + assert run_steps(store, morning._try_steps(0))[1] is not None + + +def test_a_refused_second_day_gives_the_first_day_back(): + # Someone already holds tomorrow's chain: taking today's alone would + # hold a slot this run cannot keep past midnight. + clock = Clock(now=22 * 3600.0) + store = FakeDict(clock) + store.put( + "claim-0-1-0", + {"holder": "other", "expires": DAY + 3 * 3600, "released": False}, + ) + a = slots_for(store, clock, "a", count=1) + assert run_steps(store, a._try_steps(0)) == (True, None) + assert store.data["claim-0-0-0"]["released"] is True + + +def test_an_interrupt_frees_the_slot_after_the_grace_period(): + clock = Clock() + store = FakeDict(clock) + slots = slots_for(store, clock, "a", count=1) + lease = slots.acquire() + slots.release(lease, grace=runner.INTERRUPT_GRACE_SECONDS) + assert holders_of(store, clock) == {0: "a"} + other = slots_for(store, clock, "b", count=1) + clock.now += runner.INTERRUPT_GRACE_SECONDS + MARGIN - 1 + assert run_steps(store, other._try_steps(0)) == (True, None) + clock.now += 1 + assert run_steps(store, other._try_steps(0))[1] is not None + + +def test_a_step_that_outlasts_its_lease_grants_nothing(): + # A machine asleep in the middle of taking a slot: whatever it wrote, + # it holds nothing afterwards, and nobody is displaced. + clock = Clock() + store = FakeDict(clock) + a = slots_for(store, clock, "a", count=1) + steps = a._try_steps(0) + call = next(steps) + while call[0] == "get": + call = steps.send(store.get(call[1])) + clock.now += 9 * DAY # asleep with the create in flight + b = slots_for(store, clock, "b", count=1).acquire() + try: + while True: + if call[0] == "get": + call = steps.send(store.get(call[1])) + else: + call = steps.send(store.put(call[1], call[2], skip_if_exists=call[3])) + except StopIteration as done: + assert done.value is None + assert holders_of(store, clock) == {0: "b"} + assert b.slot == 0 + + +def test_a_week_late_write_lands_in_a_day_that_no_longer_matters(): + # Review of bac8974, 1: with no pointer to roll back, a write that + # arrives days late cannot change which claim is newest today. + clock = Clock() + store = FakeDict(clock) + early = slots_for(store, clock, "early", count=1) + lease = early.acquire() + clock.now += 6.5 * DAY + current = slots_for(store, clock, "current", count=1).acquire() + early.release(lease) # the late write + stalled = slots_for(store, clock, "stalled", count=1) + steps = stalled._try_steps(0) + call = next(steps) + clock.now += 0.9 * DAY - runner.LEASE_SECONDS # stalls, then resumes + result = None + try: + while True: + if call[0] == "get": + call = steps.send(store.get(call[1])) + else: + call = steps.send(store.put(call[1], call[2], skip_if_exists=call[3])) + except StopIteration as done: + result = done.value + live = current.expires > clock.now + assert result is None or not live + + +def test_a_fast_clock_cannot_take_an_occupied_slot(): + # Review of bac8974, 2: main() runs every holder on Modal's clock; within + # the drift it enforces, a contender cannot see a live claim as lapsed + # before the run under it has ended. + clock = Clock() + store = FakeDict(clock) + drift = runner.MAX_CLOCK_DRIFT_SECONDS + holder = slots_for(store, clock, "holder", count=1, offset=-drift) + lease = holder.acquire() + run_must_end_by = lease.expires - runner.GUARD_MARGIN_SECONDS # remote guard + fast = slots_for(store, clock, "fast", count=1, offset=drift) + clock.now = run_must_end_by + assert run_steps(store, fast._try_steps(0)) == (True, None) + taken_at = None + while taken_at is None: + clock.now += 1 + if run_steps(store, fast._try_steps(0))[1] is not None: + taken_at = clock.now + assert taken_at >= run_must_end_by + runner.GUARD_MARGIN_SECONDS - 2 * drift + + +def test_a_machine_five_hours_fast_still_sees_modal_time(): + # Review of bac8974, 2, as reported: B's own clock is five hours ahead. + # main() gives the slots Modal's time (local plus the measured offset), + # so B sees A's claim as live. + clock = Clock() + store = FakeDict(clock) + slots_for(store, clock, "a", count=1).acquire() + + def b_local(): + return clock() + 5 * 3600 + + offset = runner.clock_offset(clock, local=b_local) + assert offset == -5 * 3600 + b = runner.SimulationSlots( + store, holder="b", count=1, clock=lambda: b_local() + offset + ) + assert run_steps(store, b._try_steps(0)) == (True, None) + assert holders_of(store, clock) == {0: "a"} + + +def test_the_clock_offset_is_measured_against_modal(): + local = iter([100.0, 104.0]) + assert runner.clock_offset(lambda: 1_102.0, local=lambda: next(local)) == 1_000.0 + # A slow first round trip (a cold start) is retried. + local = iter([0.0, 50.0, 60.0, 61.0]) + assert runner.clock_offset(lambda: 5.5, local=lambda: next(local)) == -55.0 + slow = iter(range(0, 1000, 20)) + with pytest.raises(runner.LeaseLost): + runner.clock_offset(lambda: 0.0, local=lambda: float(next(slow))) + + +class FailingPuts(FakeDict): + """A Dict whose writes start failing after ``allowed`` of them.""" + + def __init__(self, clock, allowed): + super().__init__(clock) + self.allowed = allowed + + def put(self, key, value, *, skip_if_exists=False): + if self.allowed == 0: + raise OSError("connection lost") + self.allowed -= 1 + return super().put(key, value, skip_if_exists=skip_if_exists) + + +def test_a_failed_second_claim_gives_the_first_back_when_it_can(): + clock = Clock(now=22 * 3600.0) # the lease spans two days: two creates + store = FailingPuts(clock, allowed=1) + a = slots_for(store, clock, "a", count=1) + with pytest.raises(OSError): + a.acquire(wait_seconds=0) + store.allowed = 10 + # The first claim stays (its release failed too) but lapses with the lease. + clock.now += runner.LEASE_SECONDS + MARGIN + assert slots_for(store, clock, "b", count=1).acquire(wait_seconds=0).slot == 0 + + +def holder_program(slots): + """main()'s use of the slots, one Dict call or event at a time.""" + lease = None + while lease is None: + for slot in range(slots.count): + lease = yield from slots._try_steps(slot) + if lease is not None: + break + else: + yield ("pause",) + outcome = yield ("start", lease) + if outcome == "running": + outcome = yield ("await",) + # main(): a returned run frees the slot now; after a caught ctrl-C + # (app.run() exited, the stop request sent) it frees after the grace; + # after a lost connection or a remote error nothing is written. + if outcome == "returned": + yield from slots._release_steps(lease) + elif outcome == "interrupted": + yield from slots._release_steps(lease, grace=runner.INTERRUPT_GRACE_SECONDS) + + +# Modal's disconnect stops an ephemeral app's tasks "in a few seconds" +# (modal/runner.py); the model allows a minute, inside the grace period. +STOP_AFTER_DISCONNECT = 60 + + +@settings(max_examples=600, deadline=None) +@given(st.data()) +def test_interleaved_holders_never_overlap_on_a_slot_and_never_lose_one(data): + """Model check: holders interleaved call by call, each on its own clock + (within the drift main() enforces), with clock jumps of any size at any + point (a stalled step, a closed laptop, midnight), crashes (no release; + a started run continues to its end), interrupts, lost connections and + Dict entry expiry. Safety: two runs never overlap on one slot, so at + most SLOT_COUNT run at once. Liveness: once every lease has lapsed, both + slots can be taken again.""" + clock = Clock(now=float(data.draw(st.sampled_from([3600, 12 * 3600, 23 * 3600])))) + store = FakeDict(clock) + drift = runner.MAX_CLOCK_DRIFT_SECONDS + count = data.draw(st.integers(2, 4), label="holders") + actors = [] + for index in range(count): + offset = data.draw(st.sampled_from([-drift, 0, drift]), label="offset") + program = holder_program(slots_for(store, clock, f"h{index}", offset=offset)) + actors.append({"program": program, "call": next(program), "done": False}) + runs = [] + + def send(actor, value): + try: + actor["call"] = actor["program"].send(value) + except StopIteration: + actor["done"] = True + + for _ in range(data.draw(st.integers(0, 150), label="actions")): + action = data.draw( + st.sampled_from( + ["step", "step", "step", "tick", "crash", "interrupt", "disconnect"] + ) + ) + if action == "tick": + clock.now += data.draw( + st.sampled_from( + [1, 60, 600, 3600, TIMEOUT, 4 * 3600, 20 * 3600, 6.5 * DAY, 8 * DAY] + ) + ) + continue + live = [a for a in actors if not a["done"]] + if not live: + break + if action in ("interrupt", "disconnect"): + awaiting = [a for a in live if a["call"][0] == "await"] + if awaiting: + actor = data.draw(st.sampled_from(awaiting)) + if action == "interrupt": # ctrl-C: the stop request is sent + run = actor["run"] + run["end"] = min(run["end"], clock.now + STOP_AFTER_DISCONNECT) + send(actor, "interrupted") + else: # connection lost: the run goes on to its own end + send(actor, "connection_lost") + continue + actor = data.draw(st.sampled_from(live)) + if action == "crash": + actor["done"] = True + continue + call = actor["call"] + if call[0] == "get": + send(actor, store.get(call[1])) + elif call[0] == "put": + send(actor, store.put(call[1], call[2], skip_if_exists=call[3])) + elif call[0] == "pause": + send(actor, None) + elif call[0] == "start": + lease = call[1] + # The remote guard, on Modal's (true) clock. + if clock.now + TIMEOUT + runner.GUARD_MARGIN_SECONDS > lease.expires: + send(actor, "refused") + continue + duration = data.draw(st.integers(1, TIMEOUT), label="run seconds") + run = {"slot": lease.slot, "start": clock.now, "end": clock.now + duration} + for other in runs: + if other["slot"] == run["slot"]: + assert other["end"] <= run["start"], (other, run) + runs.append(run) + actor["run"] = run + send(actor, "running") + elif call[0] == "await" and clock.now >= actor["run"]["end"]: + send(actor, "returned") + running = [r for r in runs if r["start"] <= clock.now < r["end"]] + assert len(running) <= runner.SLOT_COUNT + + clock.now += runner.LEASE_SECONDS + MARGIN + drift + fresh = [slots_for(store, clock, f"fresh{n}") for n in range(runner.SLOT_COUNT)] + taken = {s.acquire(wait_seconds=0).slot for s in fresh} + assert taken == set(range(runner.SLOT_COUNT)), (taken, clock.now, store.data) + + +def modal_fakes(slots, remote, run=nullcontext): + return lambda mounts: ( + SimpleNamespace( + enable_output=nullcontext, + Dict=SimpleNamespace(from_name=lambda name, create_if_missing: slots), + ), + SimpleNamespace(run=run), + SimpleNamespace(remote=remote), + SimpleNamespace(remote=time.time), + ) + + +@pytest.fixture +def stub_main(monkeypatch, tmp_path): + monkeypatch.setattr(runner, "validate_selection", lambda request: None) + monkeypatch.setattr( + runner, "input_manifest", lambda event, hf_cache=None: ({"event": event}, {}) + ) + monkeypatch.setattr(runner, "download_prefix", lambda value: tmp_path / "download") + saved = [] + monkeypatch.setattr(runner, "save_download", lambda *args: saved.append(args) or []) + return saved + + +ARGS = ["--event", "autumn_budget_2024", "--execute"] + + +def claims_in(slots): + return [v for k, v in slots.data.items() if k.startswith("claim-")] + + +def test_a_finished_run_frees_its_slot_at_once(monkeypatch, stub_main): + slots = FakeDict() + seen = {} + + def remote(payload, lease_expires): + seen["expires"] = lease_expires + return {"files": {}} + + monkeypatch.setattr(runner, "modal_function", modal_fakes(slots, remote)) + assert runner.main(ARGS) == 0 + assert len(stub_main) == 1 + assert all(claim["released"] is True for claim in claims_in(slots)) + # The run was told a lease long enough for its timeout and the margin. + assert seen["expires"] - time.time() > TIMEOUT + runner.GUARD_MARGIN_SECONDS + + +def test_a_lost_connection_keeps_the_claim_until_the_lease_lapses( + monkeypatch, stub_main +): + # The runner re-raises a lost connection without asking the server to + # stop the app, so the run may still be going. + class LostConnection(Exception): + pass + + def lost(payload, lease_expires): + raise LostConnection + + slots = FakeDict() + monkeypatch.setattr(runner, "modal_function", modal_fakes(slots, lost)) + with pytest.raises(LostConnection): + runner.main(ARGS) + assert all(claim["released"] is False for claim in claims_in(slots)) + assert all( + claim["expires"] >= time.time() + runner.LEASE_SECONDS - 60 + for claim in claims_in(slots) + ) + + +def test_a_returned_result_survives_a_failed_app_exit(monkeypatch, stub_main): + class ExitFails: + def __enter__(self): + return self + + def __exit__(self, kind, value, traceback): + raise OSError("stop request failed") + + slots = FakeDict() + monkeypatch.setattr( + runner, + "modal_function", + modal_fakes(slots, lambda payload, lease_expires: {"files": {}}, run=ExitFails), + ) + assert runner.main(ARGS) == 0 + assert len(stub_main) == 1 + assert all(claim["released"] is True for claim in claims_in(slots)) + + +def test_a_returned_result_survives_a_failed_release(monkeypatch, stub_main, capsys): + # Review of bac8974, 3: the Dict failing at release must not cost a + # completed result; the claim simply lapses. + class ReleaseFails(FakeDict): + def put(self, key, value, *, skip_if_exists=False): + if not skip_if_exists: + raise OSError("connection lost") + return super().put(key, value, skip_if_exists=skip_if_exists) + + slots = ReleaseFails() + monkeypatch.setattr( + runner, + "modal_function", + modal_fakes(slots, lambda payload, lease_expires: {"files": {}}), + ) + assert runner.main(ARGS) == 0 + assert len(stub_main) == 1 + assert "not released" in capsys.readouterr().err + + +def test_a_clock_that_moved_while_waiting_stops_the_run(monkeypatch, stub_main, capsys): + slots = FakeDict() + readings = iter([0.0, 120.0]) # Modal minus local: 0 at first, then 120 s + + def modal_now(): + return time.time() + next(readings) + + def fakes(mounts): + modal, app, replay, _ = modal_fakes( + slots, lambda payload, lease_expires: pytest.fail("the run must not start") + )(mounts) + return modal, app, replay, SimpleNamespace(remote=modal_now) + + monkeypatch.setattr(runner, "modal_function", fakes) + assert runner.main(ARGS) == 75 + assert "clock moved" in capsys.readouterr().err + assert all(claim["released"] is True for claim in claims_in(slots))