Retry failed Spot jobs on standard capacity without restarting cancelled jobs

src/gcprunner/job.py

@@ -83,7 +83,7 @@ def resolve_image(
     "FAILED": "FAILED",
     "DELETION_IN_PROGRESS": "RUNNING",
     "CANCELLATION_IN_PROGRESS": "RUNNING",
-    "CANCELLED": "FAILED",
+    "CANCELLED": "CANCELLED",
 }
 
 

src/gcprunner/runner.py

@@ -25,6 +25,8 @@ def multi_run(
 
     Mirrors ``slurmrunner.multi_run``: jobs run independently, no dependency
     graph. Batch already reaped each VM, so there's nothing to clean up here.
+    A Spot job that still fails after its Batch retries is resubmitted once as
+    standard.
 
     Raises ``RuntimeError``, naming the jobs, if any fails to launch or reaches
     a terminal state other than ``COMPLETED`` (Batch already applied its in-task
@@ -75,6 +77,27 @@ def multi_run(
                 if job.status in {"PENDING", "RUNNING"}:
                     still.append(job)
                     continue
+                if job.status == "FAILED" and job.spec.provisioning == "spot":
+                    # Spot retries exhausted (typically a preemption storm):
+                    # resubmit once on-demand, which Batch never reclaims.
+                    logger.warning(
+                        "Job %s failed on Spot; resubmitting as standard", job.name
+                    )
+                    job.spec.provisioning = "standard"
+                    try:
+                        job.launch()
+                    except Exception as e:
+                        logger.error("Failed to relaunch %s: %s", job.name, e)
+                        job.status = "FAILED"
+                        failed.append(job.name)
+                        submission_errors.append(f"{job.name}: {e}")
+                        continue
+                    if verbose:
+                        print(
+                            f"Relaunched as standard: {job.name} (id: {job.short_id})"
+                        )
+                    still.append(job)
+                    continue
                 if verbose:
                     print(f"Completed: {job.name} - {job.status}")
                 logger.info("Job %s -> %s", job.name, job.status)

src/tests/test_gcprunner.py

@@ -13,7 +13,7 @@
     build_batch_job_dict,
 )
 from gcprunner.gpus import GPU_PROFILES, resolve_gpu
-from gcprunner.job import Job, _parse_slurm_time, _job_id_safe
+from gcprunner.job import Job, _STATE_MAP, _parse_slurm_time, _job_id_safe
 from gcprunner.runner import multi_run
 
 
@@ -693,15 +693,18 @@ def test_gcp_retry_unset_omits_kwarg():
 class _FakeJob:
     """Minimal Job stand-in for multi_run: terminal status set on update."""
 
-    def __init__(self, name, final_status):
+    def __init__(self, name, final_status, provisioning="standard"):
         self.name = name
         self._final = final_status
         self.status = "PENDING"
         self.short_id = name
         self.cmd = "echo hi"
+        self.spec = SimpleNamespace(provisioning=provisioning)
+        self.launched = []
 
     def launch(self):
-        pass
+        self.launched.append(self.spec.provisioning)
+        self.status = "PENDING"
 
     def update_status(self, *args, **kwargs):
         self.status = self._final
@@ -722,6 +725,36 @@ def test_multi_run_all_completed_ok():
     multi_run(jobs, max_proc=3, launch_sleep=0, update_sleep=0)  # no raise
 
 
+def test_multi_run_spot_failure_falls_back_to_standard():
+    class _PreemptedOnceJob(_FakeJob):
+        def update_status(self, *args, **kwargs):
+            self.status = "FAILED" if self.spec.provisioning == "spot" else "COMPLETED"
+
+    job = _PreemptedOnceJob("j0", None, provisioning="spot")
+    multi_run([job], launch_sleep=0, update_sleep=0)  # no raise
+    assert job.launched == ["spot", "standard"]
+
+
+def test_multi_run_spot_fallback_fails_once():
+    job = _FakeJob("j0", "FAILED", provisioning="spot")
+    with pytest.raises(RuntimeError, match="j0"):
+        multi_run([job], launch_sleep=0, update_sleep=0)
+    assert job.launched == ["spot", "standard"]
+
+
+@pytest.mark.parametrize("raise_on_failure", [True, False])
+def test_multi_run_cancelled_spot_job_is_not_relaunched(raise_on_failure):
+    job = _FakeJob("j0", _STATE_MAP["CANCELLED"], provisioning="spot")
+    kwargs = dict(launch_sleep=0, update_sleep=0, raise_on_failure=raise_on_failure)
+    if raise_on_failure:
+        with pytest.raises(RuntimeError, match="j0"):
+            multi_run([job], **kwargs)
+    else:
+        multi_run([job], **kwargs)
+    assert job.launched == ["spot"]
+    assert job.status == "CANCELLED"
+
+
 def test_multi_run_raises_on_launch_failure():
     class _UnlaunchableJob(_FakeJob):
         def launch(self):

src/tests/test_snp_folds.py

@@ -34,7 +34,7 @@ def __init__(self, **kwargs):
         self.snp_stats = "logSUM"
         self.targets_file = None
         self.crosses = 1
-        self.conda_env = "torch126"
+        self.conda_env = None
         self.embed = False
         self.num_folds = None
         self.fold_subset_list = None