feat: Persist batch with task IDs before start to allow recovery

Author: FilipeFcp

CHANGELOG.md

@@ -12,6 +12,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 ### Fixed
 - Fixed missing amplitude factor and handling of negative normal direction case when making adjoint sources from `DiffractionMonitor`.
+- Improved the robustness of batch jobs. The batch state, including all `task_ids`, is now saved to `batch.hdf5` immediately after upload. This fixes an issue where an interrupted batch (e.g., due to a kernel crash or network loss) would be unrecoverable.
 
 ## [2.9.0] - 2025-08-04
 

tests/test_web/test_webapi.py

@@ -1,6 +1,8 @@
 # Tests webapi and things that depend on it
 from __future__ import annotations
 
+import os
+
 import numpy as np
 import pytest
 import responses
@@ -656,6 +658,40 @@ def test_create_output_dirs(mock_webapi, tmp_path, monkeypatch):
     assert non_existent_dirs_batch.is_dir()
 
 
+@responses.activate
+def test_batch_run_saves_file_after_upload(mock_webapi, mock_job_status, tmp_path, monkeypatch):
+    """Test that batch.run() saves batch file with task_ids immediately after upload."""
+    sims = {TASK_NAME: make_sim()}
+    batch = Batch(simulations=sims, folder_name=PROJECT_NAME)
+
+    batch_file_saved = {"saved": False, "has_task_ids": False}
+    original_to_file = Batch.to_file
+
+    def track_to_file(self, fname):
+        batch_file_saved["saved"] = True
+        batch_file_saved["has_task_ids"] = self.jobs is not None and TASK_NAME in self.jobs
+        return original_to_file(self, fname)
+
+    # mock start to interrupt run() after upload and to_file
+    def mock_start_interrupt(self):
+        # at this point, upload() and to_file() should have been called
+        assert batch_file_saved["saved"], "Batch file should be saved before start()"
+        assert batch_file_saved["has_task_ids"], "Batch file should have task_ids"
+        # verify file actually exists and can be loaded
+        batch_path = self._batch_path(path_dir=str(tmp_path))
+        assert os.path.exists(batch_path)
+        recovered = Batch.from_file(batch_path)
+        assert recovered.jobs[TASK_NAME].task_id == TASK_ID
+        raise RuntimeError("Simulated interruption after upload")
+
+    monkeypatch.setattr(Batch, "to_file", track_to_file)
+    monkeypatch.setattr(Batch, "start", mock_start_interrupt)
+
+    # run should save the batch file after upload, even if interrupted
+    with pytest.raises(RuntimeError, match="Simulated interruption"):
+        batch.run(path_dir=str(tmp_path))
+
+
 """ Async """
 
 

tidy3d/web/api/container.py

@@ -601,6 +601,7 @@ def run(self, path_dir: str = DEFAULT_DATA_DIR) -> BatchData:
         """
         self._check_path_dir(path_dir)
         self.upload()
+        self.to_file(self._batch_path(path_dir=path_dir))
         self.start()
         self.monitor()
         return self.load(path_dir=path_dir)
@@ -988,7 +989,6 @@ def load(self, path_dir: str = DEFAULT_DATA_DIR, replace_existing: bool = False)
         allowing one to load this :class:`Batch` later using ``batch = Batch.from_file()``.
         """
         self._check_path_dir(path_dir=path_dir)
-        self.download(path_dir=path_dir, replace_existing=replace_existing)
 
         if self.jobs is None:
             raise DataError("Can't load batch results, hasn't been uploaded.")
@@ -1010,6 +1010,8 @@ def load(self, path_dir: str = DEFAULT_DATA_DIR, replace_existing: bool = False)
                 job_data = data[task_name]
                 job.simulation._patch_data(data=job_data)
 
+        self.download(path_dir=path_dir, replace_existing=replace_existing)
+
         return data
 
     def delete(self) -> None: