diff --git a/README.md b/README.md index 60588437..e88687f6 100644 --- a/README.md +++ b/README.md @@ -67,6 +67,7 @@ rapidpipe run show "$run_id" `rapidpipe.launch.batch` reads its deployment configuration from the environment only, never from a committed value: - `RAPIDPIPE_BATCH_JOB_QUEUE` -- the Batch job queue name or ARN. +- `RAPIDPIPE_BATCH_RECLAIM_QUEUE` -- optional: the queue for every further attempt of a unit once one of its attempts, in the run or a run it was seeded from (`--only-failed`, `loop --retry-failed`), was lost to a Spot reclaim (a FAILED job whose `statusReason` starts `Host EC2`, recorded `transient` with `scheduler_metadata.batch.reclaim`). Set it to an on-demand queue while `RAPIDPIPE_BATCH_JOB_QUEUE` names a Spot queue, and a reclaimed unit moves to on-demand capacity and never retries on Spot; any other transient keeps the job queue. Unset, or equal to `RAPIDPIPE_BATCH_JOB_QUEUE`, every attempt goes to the job queue, as before. `run reconcile --resolve-jobless` searches both queues when they differ. (Ruling, supervisor step 3, 2026-10-04: the resubmit's queue is the gap between a classified reclaim and on-demand placement; an EventBridge resubmit would bypass the attempt record, and a lane-to-queue map is a bigger change than the gap.) - `RAPIDPIPE_BATCH_JOB_DEFINITION_SCRATCH` / `RAPIDPIPE_BATCH_JOB_DEFINITION_PRODUCTION` -- the Batch job definition name or ARN for a scratch or production run; `RAPIDPIPE_BATCH_JOB_DEFINITION` is the scratch fallback only. - `RAPIDPIPE_OUTPUTS_ROOT_SCRATCH` / `RAPIDPIPE_OUTPUTS_ROOT_PRODUCTION` -- the `s3://bucket/prefix` under which `runs////` lives for a scratch or production run; `RAPIDPIPE_OUTPUTS_ROOT` is the scratch fallback only. A production run never falls back to an unsuffixed variable. - `RAPIDPIPE_SCRATCH_BUCKET` -- optional: the only bucket `run delete` may remove objects from (default: the bucket of the scratch outputs root). diff --git a/rapidpipe/launch/batch.py b/rapidpipe/launch/batch.py index 9835b38a..5d9103b9 100644 --- a/rapidpipe/launch/batch.py +++ b/rapidpipe/launch/batch.py @@ -30,6 +30,13 @@ hostnames are injected at deploy time, never committed to `rapid`."): - ``RAPIDPIPE_BATCH_JOB_QUEUE`` -- the Batch job queue name or ARN. +- ``RAPIDPIPE_BATCH_RECLAIM_QUEUE`` -- optional: the queue for every + further attempt of a unit once one of its attempts, in the run or a run + it was seeded from, was lost to a Spot reclaim (:func:`is_host_reclaim`), + so a reclaimed unit moves to on-demand capacity and never retries on + Spot. Unset, or equal to + ``RAPIDPIPE_BATCH_JOB_QUEUE``, every attempt goes to the job queue and + no reclaim lookup is made (:func:`queue_for_unit`). - ``RAPIDPIPE_BATCH_JOB_DEFINITION_SCRATCH`` / ``RAPIDPIPE_BATCH_JOB_DEFINITION_PRODUCTION`` -- the Batch job definition name or ARN for a run of that kind (:func:`job_definition_for`). @@ -393,10 +400,10 @@ def submit_unit( if outputs_root is None or job_definition is None: kind = _run_kind(conn, run_id) outputs_root = outputs_root or outputs_root_for(kind) - job_queue = _require_env("RAPIDPIPE_BATCH_JOB_QUEUE") + job_queue = queue_for_unit(conn, run_id, stage, unit_id) job_definition = job_definition or job_definition_for(kind) else: - job_queue = _require_env("RAPIDPIPE_BATCH_JOB_QUEUE") + job_queue = queue_for_unit(conn, run_id, stage, unit_id) job_name_prefix = os.environ.get("RAPIDPIPE_BATCH_JOB_NAME_PREFIX", "rapid") batch = client if client is not None else batch_client() if released is not None: @@ -629,19 +636,93 @@ def _last_container_exit_code(job: dict[str, Any]) -> int | None: return container.get("exitCode") +def _status_reasons(job: dict[str, Any]) -> list[Any]: + """The job's own ``statusReason`` and its last attempt's.""" + attempts = job.get("attempts") or [] + last = attempts[-1] if attempts else {} + return [job.get("statusReason"), last.get("statusReason")] + + +def is_host_reclaim(job: dict[str, Any]) -> bool: + """Whether ``job`` lost its EC2 host (a Spot reclaim): a ``statusReason``, + the job's or its last attempt's, starts with one of + :data:`RETRYABLE_INFRASTRUCTURE_REASONS`' ``statusReason`` prefixes.""" + prefixes = RETRYABLE_INFRASTRUCTURE_REASONS["statusReason"] + return any(isinstance(r, str) and r.startswith(prefixes) for r in _status_reasons(job)) + + def is_retryable_infrastructure_failure(job: dict[str, Any]) -> bool: """Whether a FAILED ``job`` with no container exit code failed for one of :data:`RETRYABLE_INFRASTRUCTURE_REASONS`.""" + if is_host_reclaim(job): + return True attempts = job.get("attempts") or [] last = attempts[-1] if attempts else {} - status_reasons = [job.get("statusReason"), last.get("statusReason")] container_reasons = [(last.get("container") or {}).get("reason"), (job.get("container") or {}).get("reason")] - for reasons, prefixes in ((status_reasons, RETRYABLE_INFRASTRUCTURE_REASONS["statusReason"]), - (container_reasons, RETRYABLE_INFRASTRUCTURE_REASONS["reason"])): - if any(isinstance(r, str) and r.startswith(prefixes) for r in reasons): - return True - return False + prefixes = RETRYABLE_INFRASTRUCTURE_REASONS["reason"] + return any(isinstance(r, str) and r.startswith(prefixes) for r in container_reasons) + + +def _unit_was_reclaimed(conn, run_id: str, stage: str, unit_id: str) -> bool: + """Whether any attempt of ``unit_id`` in ``run_id``'s ``stage``, or of + the unit it was seeded from (``units.seeded_from_unit``, followed back + through every seed, so a ``--retry-failed`` re-run keeps it), was + recorded ``transient`` for a Spot reclaim (:func:`reconcile` marks it + ``scheduler_metadata.batch.reclaim``). + + A module-level function so the database-free unit tests can + monkeypatch it alongside the repository calls. + """ + with conn.cursor() as cur: + cur.execute( + """ + WITH RECURSIVE lineage (id, seed) AS ( + SELECT id, seeded_from_unit FROM units + WHERE run = %s AND stage = %s AND unit_id = %s + UNION + SELECT u.id, u.seeded_from_unit + FROM units u JOIN lineage l ON u.id = l.seed + ) + SELECT EXISTS ( + SELECT 1 + FROM lineage l + JOIN attempts a ON a.unit = l.id + JOIN execution_records e ON e.attempt = a.id + WHERE a.disposition = 'transient' + AND (e.scheduler_metadata -> 'batch' ->> 'reclaim') = 'true' + ) + """, + (run_id, stage, unit_id)) + (reclaimed,) = cur.fetchone() + return bool(reclaimed) + + +def _search_queues() -> list[str]: + """The queues a job of this deployment can be on: the job queue, then + the reclaim queue when it is set and differs.""" + job_queue = _require_env("RAPIDPIPE_BATCH_JOB_QUEUE") + reclaim_queue = os.environ.get("RAPIDPIPE_BATCH_RECLAIM_QUEUE") or job_queue + return [job_queue] if reclaim_queue == job_queue else [job_queue, reclaim_queue] + + +def queue_for_unit(conn, run_id: str, stage: str, unit_id: str) -> str: + """The Batch queue the next attempt of ``unit_id`` is submitted to. + + ``RAPIDPIPE_BATCH_JOB_QUEUE``, unless ``RAPIDPIPE_BATCH_RECLAIM_QUEUE`` + is set, differs from it, and an earlier attempt of this unit, in this + run or a run it was seeded from, was lost to a Spot reclaim + (:func:`_unit_was_reclaimed`): then the + reclaim queue, for every further attempt, so a reclaimed unit runs on + on-demand capacity and never retries on Spot. Any other transient + (exit 75, a container-start failure, ``lost``) on a unit never + reclaimed keeps the job queue. Unset or equal, no lookup is made. + """ + queues = _search_queues() + if len(queues) == 1: + return queues[0] + job_queue, reclaim_queue = queues + return reclaim_queue if _unit_was_reclaimed(conn, run_id, stage, unit_id) else job_queue def _epoch_ms_to_iso(value: Any) -> str | None: @@ -889,11 +970,15 @@ def reconcile( # records this as a termination without a stage exit # code (stage contract, "Exit codes"). disposition = "killed" + execution_record = _with_batch_scheduler_metadata( + _execution_record_with_defaults(conn, run_id), job) + if disposition == "transient" and is_host_reclaim(job): + # What queue_for_unit reads to send the unit's further + # attempts to the reclaim queue. + execution_record["scheduler_metadata"]["batch"]["reclaim"] = True record_attempt_result( conn, attempt_id, exit_code, disposition, output_location, - _with_batch_scheduler_metadata( - _execution_record_with_defaults(conn, run_id), job), - scheduler_job_id=job_id) + execution_record, scheduler_job_id=job_id) conn.commit() results.append(Reconciled( attempt_id=attempt_id, job_id=job_id, batch_status=status, @@ -956,8 +1041,10 @@ def resolve_jobless( ``scheduler_job_id``: its allocation committed but the Batch submission (or recording its job id) failed or has not happened yet, so :func:`reconcile` never looks at it and ``run cancel`` cannot terminate - it. For each one, the Batch queue (``RAPIDPIPE_BATCH_JOB_QUEUE``) is - searched, every status, for the name :func:`submit_unit` gives its job + it. For each one, the Batch queue (``RAPIDPIPE_BATCH_JOB_QUEUE``, and + ``RAPIDPIPE_BATCH_RECLAIM_QUEUE`` too when it is set and differs, since + a reclaimed unit's attempt is submitted there) is searched, every + status, for the name :func:`submit_unit` gives its job (:func:`_job_name` with ``RAPIDPIPE_BATCH_JOB_NAME_PREFIX``): - exactly one job: its id is recorded on the attempt @@ -994,13 +1081,17 @@ def resolve_jobless( if not jobless: return [] - job_queue = _require_env("RAPIDPIPE_BATCH_JOB_QUEUE") + queues = _search_queues() prefix = os.environ.get("RAPIDPIPE_BATCH_JOB_NAME_PREFIX", "rapid") batch = client if client is not None else batch_client() note = f"no scheduler job after {older_than_seconds:g} s" results: list[Reconciled] = [] for attempt_id, stage in jobless: - found = _jobs_named(batch, job_queue, _job_name(prefix, stage, attempt_id)) + job_name = _job_name(prefix, stage, attempt_id) + # Distinct ids: a queue's name and its ARN are two strings for one + # queue, and the same job found twice is not an ambiguity. + found = list(dict.fromkeys( + job_id for queue in queues for job_id in _jobs_named(batch, queue, job_name))) if len(found) > 1: results.append(Reconciled( attempt_id=attempt_id, job_id=",".join(found), batch_status="AMBIGUOUS", diff --git a/tests/db/test_launch.py b/tests/db/test_launch.py index bd5a90b4..01211b17 100644 --- a/tests/db/test_launch.py +++ b/tests/db/test_launch.py @@ -446,6 +446,108 @@ def test_reconcile_missing_job_id_is_lost(conn, batch_env): assert disposition == "lost" +def test_reclaimed_unit_resubmits_to_the_reclaim_queue_for_good(conn, batch_env, monkeypatch): + monkeypatch.setenv("RAPIDPIPE_BATCH_RECLAIM_QUEUE", "test-reclaim-queue") + run_id = _make_run(conn, kind="scratch", max_attempts=4) + conn.commit() + + fake_batch = FakeBatch() + first = _submit_one(conn, fake_batch, run_id=run_id, unit_id="reclaim-001/SCA07") + conn.commit() + assert fake_batch.submitted[-1]["jobQueue"] == "test-queue" + + fake_batch.set_status(first.job_id, "FAILED", + status_reason="Host EC2 (instance i-0123) terminated.") + launch_batch.reconcile(conn, run_id=run_id, client=fake_batch, s3_client=FakeS3()) + conn.commit() + + with conn.cursor() as cur: + cur.execute( + "SELECT a.disposition, e.scheduler_metadata -> 'batch' -> 'reclaim' " + "FROM attempts a JOIN execution_records e ON e.attempt = a.id WHERE a.id = %s", + (first.attempt_id,)) + assert cur.fetchone() == ("transient", True) + + second = _submit_one(conn, fake_batch, run_id=run_id, unit_id="reclaim-001/SCA07") + conn.commit() + assert fake_batch.submitted[-1]["jobQueue"] == "test-reclaim-queue" + + # The reclaim is the unit's own: another unit of the run, and the same + # unit id in another run, stay on the job queue. + _submit_one(conn, fake_batch, run_id=run_id, unit_id="reclaim-001/SCA08") + conn.commit() + assert fake_batch.submitted[-1]["jobQueue"] == "test-queue" + other_run = _make_run(conn, kind="scratch", max_attempts=2) + conn.commit() + _submit_one(conn, fake_batch, run_id=other_run, unit_id="reclaim-001/SCA07") + conn.commit() + assert fake_batch.submitted[-1]["jobQueue"] == "test-queue" + + # A later transient that is not a reclaim does not send it back. + fake_batch.set_status(second.job_id, "FAILED", container_exit_code=75) + launch_batch.reconcile(conn, run_id=run_id, client=fake_batch, s3_client=FakeS3()) + conn.commit() + _submit_one(conn, fake_batch, run_id=run_id, unit_id="reclaim-001/SCA07") + conn.commit() + assert fake_batch.submitted[-1]["jobQueue"] == "test-reclaim-queue" + + +def test_reclaim_carries_into_a_run_seeded_from_it(conn, batch_env, monkeypatch): + monkeypatch.setenv("RAPIDPIPE_BATCH_RECLAIM_QUEUE", "test-reclaim-queue") + seed_run = _make_run(conn, kind="scratch", max_attempts=1) + conn.commit() + + fake_batch = FakeBatch() + first = _submit_one(conn, fake_batch, run_id=seed_run, unit_id="reclaim-003/SCA07") + conn.commit() + fake_batch.set_status(first.job_id, "FAILED", + status_reason="Host EC2 (instance i-0123) terminated.") + launch_batch.reconcile(conn, run_id=seed_run, client=fake_batch, s3_client=FakeS3()) + conn.commit() + + with conn.cursor() as cur: + cur.execute("SELECT unit FROM attempts WHERE id = %s", (first.attempt_id,)) + (seed_unit,) = cur.fetchone() + + # Two generations of re-run: a re-run of a re-run still finds the reclaim. + middle_run = _make_run(conn, kind="scratch", max_attempts=1) + repo.add_unit(conn, middle_run, "admit", "detector-image", "reclaim-003/SCA07", + seeded_from_unit=seed_unit) + conn.commit() + with conn.cursor() as cur: + cur.execute("SELECT id FROM units WHERE run = %s", (middle_run,)) + (middle_unit,) = cur.fetchone() + rerun = _make_run(conn, kind="scratch", max_attempts=1) + repo.add_unit(conn, rerun, "admit", "detector-image", "reclaim-003/SCA07", + seeded_from_unit=middle_unit) + conn.commit() + + _submit_one(conn, fake_batch, run_id=rerun, unit_id="reclaim-003/SCA07") + conn.commit() + assert fake_batch.submitted[-1]["jobQueue"] == "test-reclaim-queue" + + +def test_non_reclaim_transient_resubmits_to_the_job_queue(conn, batch_env, monkeypatch): + monkeypatch.setenv("RAPIDPIPE_BATCH_RECLAIM_QUEUE", "test-reclaim-queue") + run_id = _make_run(conn, kind="scratch", max_attempts=2) + conn.commit() + + fake_batch = FakeBatch() + first = _submit_one(conn, fake_batch, run_id=run_id, unit_id="reclaim-002/SCA07") + conn.commit() + + fake_batch.set_status(first.job_id, "FAILED", + container_reason="CannotPullContainerError: pull rate limit") + results = launch_batch.reconcile( + conn, run_id=run_id, client=fake_batch, s3_client=FakeS3()) + conn.commit() + assert results[0].disposition == "transient" + + _submit_one(conn, fake_batch, run_id=run_id, unit_id="reclaim-002/SCA07") + conn.commit() + assert fake_batch.submitted[-1]["jobQueue"] == "test-queue" + + def test_reconcile_running_job_left_untouched(conn, batch_env): run_id = _make_run(conn, kind="scratch", max_attempts=2) conn.commit() diff --git a/tests/unit/test_launch.py b/tests/unit/test_launch.py index e0c3c9d9..d7f77b75 100644 --- a/tests/unit/test_launch.py +++ b/tests/unit/test_launch.py @@ -598,6 +598,160 @@ def test_reconcile_failed_exit_code_wins_over_an_infrastructure_reason(monkeypat monkeypatch, status="FAILED", container_exit_code=70, status_reason="Host EC2 (instance i-0123) terminated.") assert result.disposition == "failed" + assert "reclaim" not in recorded["execution_record"]["scheduler_metadata"]["batch"] + + +# ====================================================================== +# Spot reclaims: the reclaim flag and the reclaim queue +# ====================================================================== + +@pytest.mark.parametrize("reasons", [ + {"status_reason": "Host EC2 (instance i-0123) terminated."}, + {"attempt_status_reason": "Host EC2 (instance i-0123) terminated."}, +]) +def test_reconcile_marks_a_host_reclaim(monkeypatch, reasons): + result, recorded = _reconcile_one( + monkeypatch, status="FAILED", container_exit_code=None, **reasons) + assert result.disposition == "transient" + assert recorded["execution_record"]["scheduler_metadata"]["batch"]["reclaim"] is True + + +@pytest.mark.parametrize("reasons", [ + {"container_reason": "CannotPullContainerError: pull rate limit"}, + {"container_reason": "DockerTimeoutError: Could not transition to started"}, +]) +def test_reconcile_does_not_mark_other_infrastructure_failures(monkeypatch, reasons): + result, recorded = _reconcile_one( + monkeypatch, status="FAILED", container_exit_code=None, **reasons) + assert result.disposition == "transient" + assert "reclaim" not in recorded["execution_record"]["scheduler_metadata"]["batch"] + + +def _submit_for_queue(monkeypatch, *, reclaim_queue, reclaimed): + monkeypatch.setenv("RAPIDPIPE_BATCH_JOB_QUEUE", "queue1") + monkeypatch.setenv("RAPIDPIPE_BATCH_JOB_DEFINITION", "def1") + monkeypatch.setenv("RAPIDPIPE_OUTPUTS_ROOT", "s3://bucket/prefix") + if reclaim_queue is None: + monkeypatch.delenv("RAPIDPIPE_BATCH_RECLAIM_QUEUE", raising=False) + else: + monkeypatch.setenv("RAPIDPIPE_BATCH_RECLAIM_QUEUE", reclaim_queue) + _patch_repository(monkeypatch) + lookups = [] + + def _fake_unit_was_reclaimed(conn, run_id, stage, unit_id): + lookups.append((run_id, stage, unit_id)) + if reclaimed is None: + raise AssertionError("no reclaim lookup expected") + return reclaimed + + monkeypatch.setattr(launch_batch, "_unit_was_reclaimed", _fake_unit_was_reclaimed) + fake = FakeBatch() + launch_batch.submit_unit( + _FakeConn(), run_id="RUN01", stage="admit", unit_kind="detector-image", + unit_id="u1", inputs_location="s3://in/pre", client=fake) + return fake.submitted[0]["jobQueue"], lookups + + +@pytest.mark.parametrize("reclaim_queue", [None, "", "queue1"]) +def test_submit_unit_without_a_distinct_reclaim_queue_is_unchanged(monkeypatch, reclaim_queue): + # Unset, empty or equal: the job queue, and no reclaim lookup at all. + queue, lookups = _submit_for_queue(monkeypatch, reclaim_queue=reclaim_queue, reclaimed=None) + assert queue == "queue1" + assert lookups == [] + + +def test_submit_unit_sends_a_reclaimed_unit_to_the_reclaim_queue(monkeypatch): + queue, lookups = _submit_for_queue(monkeypatch, reclaim_queue="ondemand", reclaimed=True) + assert queue == "ondemand" + assert lookups == [("RUN01", "admit", "u1")] + + +def test_submit_unit_keeps_a_never_reclaimed_unit_on_the_job_queue(monkeypatch): + queue, lookups = _submit_for_queue(monkeypatch, reclaim_queue="ondemand", reclaimed=False) + assert queue == "queue1" + assert lookups == [("RUN01", "admit", "u1")] + + +class _QueueAwareBatch(FakeBatch): + """FakeBatch whose ``list_jobs`` finds a job only on the queue it was + placed on, so a test can tell which queues were searched.""" + + def __init__(self, placed): + super().__init__() + self.placed = placed # job name -> (queue, job id) + self.searched = [] + + def list_jobs(self, *, jobQueue, filters=None, nextToken=None, **_): + name = next(f["values"][0] for f in filters or [] if f["name"] == "JOB_NAME") + self.searched.append(jobQueue) + queue, job_id = self.placed.get(name, (None, None)) + jobs = [{"jobId": job_id, "jobName": name}] if queue == jobQueue else [] + return {"jobSummaryList": jobs} + + +def _resolve_one_jobless(monkeypatch, *, reclaim_queue, placed_on, fake=None): + monkeypatch.setenv("RAPIDPIPE_BATCH_JOB_QUEUE", "queue1") + monkeypatch.delenv("RAPIDPIPE_BATCH_JOB_NAME_PREFIX", raising=False) + if reclaim_queue is None: + monkeypatch.delenv("RAPIDPIPE_BATCH_RECLAIM_QUEUE", raising=False) + else: + monkeypatch.setenv("RAPIDPIPE_BATCH_RECLAIM_QUEUE", reclaim_queue) + + class _Cursor: + def __enter__(self): + return self + + def __exit__(self, *a): + return False + + def execute(self, query, params=None): + pass + + def fetchall(self): + return [("ATTEMPT01", "admit")] + + def fetchone(self): + return ("loc-1", None, None, True) + + class _Conn(_FakeConn): + def cursor(self): + return _Cursor() + + recorded = [] + monkeypatch.setattr(launch_batch, "record_scheduler_job", + lambda conn, attempt_id, job_id: recorded.append((attempt_id, job_id))) + name = launch_batch._job_name("rapid", "admit", "ATTEMPT01") + if fake is None: + fake = _QueueAwareBatch({name: (placed_on, "job-9")}) + else: + fake.add_job(name) + results = launch_batch.resolve_jobless(_Conn(), run_id="r1", client=fake) + return results, recorded, getattr(fake, "searched", None) + + +def test_resolve_jobless_finds_a_job_on_the_reclaim_queue(monkeypatch): + results, recorded, searched = _resolve_one_jobless( + monkeypatch, reclaim_queue="ondemand", placed_on="ondemand") + assert searched == ["queue1", "ondemand"] + assert recorded == [("ATTEMPT01", "job-9")] + assert results[0].batch_status == "REPAIRED" + + +def test_resolve_jobless_same_job_on_two_queue_spellings_is_not_ambiguous(monkeypatch): + # A queue's name and its ARN are different strings for one queue: + # FakeBatch's list_jobs ignores the queue, so both searches find the job. + results, recorded, _ = _resolve_one_jobless( + monkeypatch, reclaim_queue="arn:aws:batch:region:acct:job-queue/queue1", + placed_on=None, fake=FakeBatch()) + assert results[0].batch_status == "REPAIRED" + assert len(recorded) == 1 + + +def test_resolve_jobless_searches_one_queue_when_the_reclaim_queue_is_the_same(monkeypatch): + results, recorded, searched = _resolve_one_jobless( + monkeypatch, reclaim_queue="queue1", placed_on="queue1") + assert searched == ["queue1"] + assert recorded == [("ATTEMPT01", "job-9")] def test_reconcile_missing_job_id_is_lost(monkeypatch):