feat(tidy3d): FXC-3004 Container-aware web.run returning lazy results

Co-authored-by: daquinteroflex <dario@flexcompute.com>

Author: 2 people

CHANGELOG.md

@@ -34,6 +34,8 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
   - `CustomCurrentIntegral2D` → `Custom2DCurrentIntegral`
   - Path integral and impedance calculator classes have been refactored and moved from `tidy3d.plugins.microwave` to `tidy3d.components.microwave`. They are now publicly exported via the top-level package `__init__.py`, so you can import them directly, e.g. `from tidy3d import ImpedanceCalculator, AxisAlignedVoltageIntegral, AxisAlignedCurrentIntegral, Custom2DVoltageIntegral, Custom2DCurrentIntegral, Custom2DPathIntegral`.
 - `DirectivityMonitor` now forces `far_field_approx` to `True`, which was previously configurable.
+- Unified run submission API: `web.run(...)` is now a container-aware wrapper that accepts a single simulation or arbitrarily nested containers (`list`, `tuple`, `dict` values) and returns results in the same shape.
+- `web.Batch(ComponentModeler)` and `web.Job(ComponentModeler)` native support
 
 ### Fixed
 - More robust `Sellmeier` and `Debye` material model, and prevent very large pole parameters in `PoleResidue` material model.

tests/test_web/test_webapi.py

@@ -2,6 +2,7 @@
 from __future__ import annotations
 
 import os
+from types import SimpleNamespace
 
 import numpy as np
 import pytest
@@ -22,7 +23,9 @@
 from tidy3d.exceptions import SetupError
 from tidy3d.web import common
 from tidy3d.web.api.asynchronous import run_async
-from tidy3d.web.api.container import Batch, Job
+from tidy3d.web.api.container import Batch, Job, WebContainer
+from tidy3d.web.api.run import _collect_by_hash, run
+from tidy3d.web.api.tidy3d_stub import Tidy3dStubData
 from tidy3d.web.api.webapi import (
     abort,
     delete,
@@ -38,7 +41,6 @@
     load_simulation,
     monitor,
     real_cost,
-    run,
     start,
     upload,
 )
@@ -55,6 +57,7 @@
 EST_FLEX_UNIT = 11.11
 FILE_SIZE_GB = 4.0
 common.CONNECTION_RETRY_TIME = 0.1
+INVALID_TASK_ID = "INVALID_TASK_ID"
 
 task_core_path = "tidy3d.web.core.task_core"
 api_path = "tidy3d.web.api.webapi"
@@ -768,13 +771,95 @@ def save_sim_to_path(path: str) -> None:
 @responses.activate
 def test_load_invalid_task_raises(mock_webapi):
     """Ensure that load() raises TaskNotFoundError for a non-existent task ID."""
-    fake_id = "INVALID_TASK_ID"
 
     responses.add(
         responses.GET,
-        f"{Env.current.web_api_endpoint}/tidy3d/tasks/{fake_id}/detail",
+        f"{Env.current.web_api_endpoint}/tidy3d/tasks/{INVALID_TASK_ID}/detail",
         json={"error": "Task not found"},
         status=404,
     )
     with pytest.raises(WebNotFoundError, match="Resource not found"):
-        load(fake_id)
+        load(INVALID_TASK_ID, replace_existing=True)
+
+
+def _fake_load_factory(tmp_root, taskid_to_sim: dict):
+    def _fake_load(task_id, path="simulation_data.hdf5", lazy=False, **kwargs):
+        abs_path = path if os.path.isabs(path) else os.path.join(tmp_root, path)
+        abs_path = os.path.normpath(abs_path)
+        os.makedirs(os.path.dirname(abs_path), exist_ok=True)
+
+        sim_for_this = taskid_to_sim.get(task_id)
+
+        log = "- Time step    827 / time 4.13e-14s (  4 % done), field decay: 0.110e+00"
+        sim_data = SimulationData(simulation=sim_for_this, data=[], diverged=False, log=log)
+
+        sim_data.to_file(abs_path)
+        return Tidy3dStubData.postprocess(abs_path, lazy=lazy)
+
+    return _fake_load
+
+
+def apply_common_patches(
+    monkeypatch, tmp_root, *, api_path="tidy3d.web.api.webapi", path_to_sim=None, taskid_to_sim=None
+):
+    """Patch start/monitor/get_info/estimate_cost/upload/_check_folder/_modesolver_patch/load."""
+    monkeypatch.setattr(f"{api_path}.start", lambda *a, **k: True)
+    monkeypatch.setattr(f"{api_path}.monitor", lambda *a, **k: True)
+    monkeypatch.setattr(f"{api_path}.get_info", lambda *a, **k: SimpleNamespace(status="success"))
+    monkeypatch.setattr(f"{api_path}.estimate_cost", lambda *a, **k: 0.0)
+    monkeypatch.setattr(f"{api_path}.upload", lambda *a, **k: k["task_name"])
+    monkeypatch.setattr(WebContainer, "_check_folder", lambda *a, **k: True)
+    monkeypatch.setattr(f"{api_path}._modesolver_patch", lambda *_, **__: None, raising=False)
+    monkeypatch.setattr(
+        f"{api_path}.load",
+        _fake_load_factory(tmp_root=str(tmp_root), taskid_to_sim=taskid_to_sim),
+        raising=False,
+    )
+
+
+@responses.activate
+def test_run_with_flexible_containers_offline_lazy(monkeypatch, tmp_path):
+    sim1 = make_sim()
+    sim2 = sim1.updated_copy(run_time=sim1.run_time / 2)
+    sim_container = [sim1, {"sim": sim1, "sim2": sim2}, (sim1, [sim2])]
+
+    h2sim = _collect_by_hash(sim_container)
+    task_name = "T"
+    out_dir = tmp_path / "out"
+
+    taskid_to_sim = {f"{task_name}_{h}": s for h, s in h2sim.items()}
+
+    apply_common_patches(monkeypatch, tmp_path, taskid_to_sim=taskid_to_sim)
+
+    data = run(sim_container, task_name=task_name, folder_name="PROJECT", path=str(out_dir))
+
+    assert isinstance(data, list) and len(data) == 3
+
+    assert isinstance(data[0], SimulationData)
+    assert data[0].__class__.__name__ == "SimulationDataProxy"
+
+    assert isinstance(data[1], dict)
+    assert "sim2" in data[1]
+    assert isinstance(data[1]["sim2"], SimulationData)
+    assert data[1]["sim2"].__class__.__name__ == "SimulationDataProxy"
+
+    assert isinstance(data[2], tuple)
+    assert data[2][0].__class__.__name__ == "SimulationDataProxy"
+    assert isinstance(data[2][1], list)
+    assert data[2][1][0].__class__.__name__ == "SimulationDataProxy"
+
+    assert data[0].simulation == sim1
+    assert data[1]["sim2"].simulation == sim2
+
+
+@responses.activate
+def test_run_single_offline_eager(monkeypatch, tmp_path):
+    sim = make_sim()
+    single_file = str(tmp_path / "sim.hdf5")
+    task_name = "single"
+    apply_common_patches(monkeypatch, tmp_path, taskid_to_sim={task_name: sim})
+
+    sim_data = run(sim, task_name=task_name, path=single_file)
+
+    assert isinstance(sim_data, SimulationData)
+    assert sim_data.__class__.__name__ == "SimulationData"  # no proxy

tests/test_web/test_webapi_mode.py

@@ -12,6 +12,7 @@
 from tidy3d.plugins.mode import ModeSolver
 from tidy3d.web.api.asynchronous import run_async
 from tidy3d.web.api.container import Batch, Job
+from tidy3d.web.api.run import run
 from tidy3d.web.api.webapi import (
     abort,
     download_json,
@@ -20,7 +21,6 @@
     get_reduced_simulation,
     get_run_info,
     load_simulation,
-    run,
     upload,
 )
 from tidy3d.web.core.environment import Env

tidy3d/web/__init__.py

@@ -13,8 +13,9 @@
 
 # from .api.asynchronous import run_async # NOTE: we use autograd one now (see below)
 # autograd compatible wrappers for run and run_async
-from .api.autograd.autograd import run, run_async
+from .api.autograd.autograd import run_async
 from .api.container import Batch, BatchData, Job
+from .api.run import run
 from .api.webapi import (
     abort,
     account,
@@ -29,6 +30,7 @@
     load,
     load_simulation,
     monitor,
+    postprocess_start,
     real_cost,
     start,
     test,
@@ -59,6 +61,7 @@
     "load",
     "load_simulation",
     "monitor",
+    "postprocess_start",
     "real_cost",
     "run",
     "run_async",

tidy3d/web/api/asynchronous.py

@@ -24,6 +24,7 @@ def run_async(
     reduce_simulation: Literal["auto", True, False] = "auto",
     pay_type: Union[PayType, str] = PayType.AUTO,
     priority: Optional[int] = None,
+    lazy: bool = False,
 ) -> BatchData:
     """Submits a set of Union[:class:`.Simulation`, :class:`.HeatSimulation`, :class:`.EMESimulation`] objects to server,
     starts running, monitors progress, downloads, and loads results as a :class:`.BatchData` object.
@@ -56,6 +57,10 @@ def run_async(
     priority: int = None
         Priority of the simulation in the Virtual GPU (vGPU) queue (1 = lowest, 10 = highest).
         It affects only simulations from vGPU licenses and does not impact simulations using FlexCredits.
+    lazy : bool = False
+        Whether to load the actual data (``lazy=False``) or return a proxy that loads
+        the data when accessed (``lazy=True``).
+
     Returns
     ------
     :class:`BatchData`
@@ -91,6 +96,7 @@ def run_async(
         parent_tasks=parent_tasks,
         reduce_simulation=reduce_simulation,
         pay_type=pay_type,
+        lazy=lazy,
     )
 
     batch_data = batch.run(path_dir=path_dir, priority=priority)

tidy3d/web/api/autograd/autograd.py

@@ -106,6 +106,7 @@ def run(
     reduce_simulation: typing.Literal["auto", True, False] = "auto",
     pay_type: typing.Union[PayType, str] = PayType.AUTO,
     priority: typing.Optional[int] = None,
+    lazy: bool = False,
 ) -> WorkflowDataType:
     """
     Submits a :class:`.Simulation` to server, starts running, monitors progress, downloads,
@@ -147,6 +148,10 @@ def run(
         Which method to pay for the simulation.
     priority: int = None
         Task priority for vGPU queue (1=lowest, 10=highest).
+    lazy : bool = False
+        Whether to load the actual data (``lazy=False``) or return a proxy that loads
+        the data when accessed (``lazy=True``).
+
     Returns
     -------
     Union[:class:`.SimulationData`, :class:`.HeatSimulationData`, :class:`.EMESimulationData`, :class:`.ModalComponentModelerData`, :class:`.TerminalComponentModelerData`]
@@ -237,6 +242,7 @@ def run(
             max_num_adjoint_per_fwd=max_num_adjoint_per_fwd,
             pay_type=pay_type,
             priority=priority,
+            lazy=lazy,
         )
 
     return run_webapi(
@@ -255,6 +261,7 @@ def run(
         reduce_simulation=reduce_simulation,
         pay_type=pay_type,
         priority=priority,
+        lazy=lazy,
     )
 
 
@@ -273,6 +280,7 @@ def run_async(
     reduce_simulation: typing.Literal["auto", True, False] = "auto",
     pay_type: typing.Union[PayType, str] = PayType.AUTO,
     priority: typing.Optional[int] = None,
+    lazy: bool = False,
 ) -> BatchData:
     """Submits a set of Union[:class:`.Simulation`, :class:`.HeatSimulation`, :class:`.EMESimulation`] objects to server,
     starts running, monitors progress, downloads, and loads results as a :class:`.BatchData` object.
@@ -307,6 +315,9 @@ def run_async(
         Whether to reduce structures in the simulation to the simulation domain only. Note: currently only implemented for the mode solver.
     pay_type: typing.Union[PayType, str] = PayType.AUTO
         Specify the payment method.
+    lazy : bool = False
+        Whether to load the actual data (``lazy=False``) or return a proxy that loads
+        the data when accessed (``lazy=True``).
 
     Returns
     ------
@@ -349,6 +360,7 @@ def run_async(
             max_num_adjoint_per_fwd=max_num_adjoint_per_fwd,
             pay_type=pay_type,
             priority=priority,
+            lazy=lazy,
         )
 
     return run_async_webapi(
@@ -364,6 +376,7 @@ def run_async(
         reduce_simulation=reduce_simulation,
         pay_type=pay_type,
         priority=priority,
+        lazy=lazy,
     )
 
 

tidy3d/web/api/autograd/engine.py

@@ -10,7 +10,7 @@
 
 def parse_run_kwargs(**run_kwargs):
     """Parse the ``run_kwargs`` to extract what should be passed to the ``Job``/``Batch`` init."""
-    job_fields = [*list(Job._upload_fields), "solver_version", "pay_type"]
+    job_fields = [*list(Job._upload_fields), "solver_version", "pay_type", "lazy"]
     job_init_kwargs = {k: v for k, v in run_kwargs.items() if k in job_fields}
     return job_init_kwargs