autograd support for Job.run() and Batch.run()
#2722
Conversation
![greptile-apps[bot]](https://avatars.githubusercontent.com/in/867647?s=60&v=4){kind=small}
| ```python | ||
| # tidy3d/web/api/batch_data.py | ||
| from __future__ import annotations | ||
| from typing import Mapping | ||
| from tidy3d.components.base import Tidy3dBaseModel | ||
| from tidy3d.web.core.constants import TaskName | ||
| from .tidy3d_stub import SimulationDataType | ||
| import tidy3d.web.api.webapi as web # Only import webapi, not container | ||
|
|
||
| class BatchData(Tidy3dBaseModel, Mapping): | ||
| # Move entire BatchData class here with minimal dependencies | ||
| ``` |
There was a problem hiding this comment.
logic: The import structure in the BatchData extraction example may create new dependencies. Consider whether importing webapi directly might recreate circular dependencies.
| # Test that autograd works through emulated Job.run() (demonstrates end goal) | ||
| val, grad = ag.value_and_grad(objective)(params0) | ||
| print(f"Job.run() with emulation - value: {val}, grad: {grad}") | ||
| # After refactor, this should work without emulation by calling autograd.run() internally |
There was a problem hiding this comment.
style: Remove this comment about 'after refactor' since the refactor appears to be implemented based on the test assertions below
Context Used: Rule - Remove commented-out or obsolete code; rely on version control for history. (link)
| 2. **Simplify Batch class**: | ||
| ```python | ||
| # container.py | ||
| class Batch(WebContainer): | ||
| def upload(self) -> None: | ||
| """Thin wrapper around asynchronous.upload_batch()""" | ||
| from .asynchronous import upload_batch | ||
| upload_batch(self.simulations, **self._get_batch_kwargs()) | ||
|
|
||
| def run(self, path_dir: str = DEFAULT_DATA_DIR) -> BatchData: | ||
| """Thin wrapper that calls asynchronous functions""" | ||
| from .asynchronous import run_async | ||
| return run_async(self.simulations, path_dir=path_dir, **self._get_batch_kwargs()) | ||
| ``` |
There was a problem hiding this comment.
logic: The Batch.run() implementation imports run_async from asynchronous module, but Step 3 is about moving batch logic TO asynchronous.py. This creates a circular reference where Batch calls asynchronous.run_async which may need Batch internals.
| # Test that autograd works through emulated Batch.run() (demonstrates end goal) | ||
| val, grad = ag.value_and_grad(objective)(params0) | ||
| print(f"Batch.run() with emulation - value: {val}, grad: {grad}") | ||
| # After refactor, this should work without emulation by calling autograd.run_async() internally |
There was a problem hiding this comment.
style: Remove this comment about 'after refactor' since the refactor appears to be implemented based on the test assertions below
Context Used: Rule - Remove commented-out or obsolete code; rely on version control for history. (link)
| 1. **Lazy imports**: Use function-level imports in container.py to avoid circular dependencies | ||
| 2. **Backwards compatibility**: Keep existing methods working for users who don't need autograd | ||
| 3. **Clean separation**: BatchData becomes independent, asynchronous.py contains core logic | ||
| 4. **Minimal changes**: Autograd functions remain unchanged, containers adapt to them | ||
|
|
||
| ### Potential Issues and Solutions | ||
|
|
||
| 1. **Import timing**: If circular imports persist, use `TYPE_CHECKING` and runtime imports | ||
| 2. **Test mocking**: May need to update mock strategy to handle the new call paths | ||
| 3. **Plugin compatibility**: Verify that S-matrix and other plugins still work correctly | ||
| 4. **Performance**: The function-level imports might add small overhead |
There was a problem hiding this comment.
style: Function-level imports will indeed add runtime overhead on every method call. Consider caching the imported functions as class/module attributes after first import to minimize performance impact.
|
|
||
| # We need to get the task_id from the run somehow | ||
| # For now, let's extract it from run_kwargs if it's there, or generate a placeholder | ||
| task_id = run_kwargs.get("task_id", f"autograd_{task_name}") |
There was a problem hiding this comment.
logic: Using placeholder task_id generation instead of getting actual task_id from webapi.run() could cause issues with autograd_fwd simulation types that need real task IDs for file uploads
| batch_data = run_async_webapi(simulations=sims, **web_kwargs) | ||
|
|
||
| # Generate task_ids - we'll need to get these from somewhere or create placeholders | ||
| task_ids = {task_name: f"autograd_batch_{task_name}" for task_name in simulations.keys()} |
There was a problem hiding this comment.
logic: Placeholder task_id generation may not work correctly for autograd_fwd cases that require real task IDs for sim_fields_keys uploads
tidy3d/web/api/webapi.py
Outdated
| # For now, use a placeholder task_id since the individual run() calls handle everything | ||
| task_ids[task_name] = f"run_async_{task_name}" |
There was a problem hiding this comment.
logic: Using placeholder task IDs like f"run_async_{task_name}" may cause issues if other code expects real task IDs. Consider using the actual task IDs returned by the run() function calls.
| "aborted", | ||
| ) | ||
|
|
||
| max_task_name = max(len(task_name) for task_name in task_ids.keys()) |
There was a problem hiding this comment.
logic: Potential infinite loop if task_ids is empty - max() will raise ValueError
| simulations=simulations, | ||
| folder_name=folder_name, | ||
| callback_url=callback_url, | ||
| num_workers=num_workers or len(simulations), |
There was a problem hiding this comment.
style: Using num_workers or len(simulations) can be problematic if num_workers=0 - should use num_workers if num_workers is not None else len(simulations)
| num_workers=num_workers or len(simulations), | |
| num_workers=num_workers if num_workers is not None else len(simulations), |
Greptile Summary
This PR implements autograd (automatic differentiation) support for
Job.run()andBatch.run()methods in the Tidy3D web API. The core change enables gradient computation through these high-level object-oriented interfaces, which is essential for inverse design and optimization workflows.The implementation follows a comprehensive refactoring plan documented in
AUTOGRAD_REFACTOR_PLAN.md. The main architectural changes include:Module Reorganization:
BatchDataclass is extracted fromcontainer.pyinto its ownbatch_data.pymodule to break circular import dependencies that prevented container classes from importing autograd functions.Function Delegation:
Job.run()andBatch.run()methods are refactored to delegate to autograd-compatible functions (run_autograd()andrun_async_autograd()) that automatically detect if simulations contain traced parameters and route them appropriately.Async Function Extraction: Batch processing logic is extracted from the
Batchclass into standalone async functions (upload_async,start_async,monitor_async, etc.) to enable autograd compatibility through function-based primitives.Sequential Processing: A new
run_async()function inwebapi.pyprovides autograd-compatible batch processing by sequentially calling the proven single-simulationrun()function for each simulation in a batch.The changes maintain full backward compatibility - existing code using
Job.run()andBatch.run()will continue to work exactly as before, but now these methods can be used within JAX-transformable functions for gradient-based optimization. The refactoring also reduces code duplication by consolidating async operations into reusable functions.Important Files Changed
Files and Ratings
tidy3d/web/api/container.pytidy3d/web/api/autograd/autograd.pytidy3d/web/api/asynchronous.pytidy3d/web/api/webapi.pytidy3d/web/api/batch_data.pytidy3d/web/__init__.pytidy3d/plugins/adjoint/web.pytests/test_components/test_autograd.pyAUTOGRAD_REFACTOR_PLAN.mdtests/utils.pytidy3d/plugins/design/design.pytests/test_plugins/test_adjoint.pytidy3d/plugins/smatrix/component_modelers/modal.pytidy3d/plugins/smatrix/component_modelers/base.pytidy3d/plugins/smatrix/component_modelers/terminal.pydocs/notebooksConfidence score: 3/5
asynchronous.pymonitoring loop andautograd.pytask ID generation for potential infinite loops or missing task referencesSequence Diagram
sequenceDiagram participant User participant Job participant Batch participant AutogradRun as "autograd.run" participant AutogradRunAsync as "autograd.run_async" participant WebAPI as "webapi functions" participant Server as "Tidy3D Server" Note over User, Server: Job.run() - Single Simulation Flow User->>+Job: job.run(path="data.hdf5") Note over Job: Job.run() now calls autograd.run() internally Job->>+AutogradRun: run(simulation, task_name, path, ...) Note over AutogradRun: Checks if simulation is autograd-compatible alt Simulation has autograd tracers AutogradRun->>AutogradRun: _run_primitive() - forward pass AutogradRun->>+WebAPI: _run_tidy3d() - actual simulation WebAPI->>+Server: upload/start/monitor/download Server-->>-WebAPI: simulation_data WebAPI-->>-AutogradRun: simulation_data AutogradRun->>AutogradRun: postprocess_run() - insert tracers else No autograd tracers AutogradRun->>+WebAPI: run_webapi() - regular run WebAPI->>+Server: upload/start/monitor/download Server-->>-WebAPI: simulation_data WebAPI-->>-AutogradRun: simulation_data end AutogradRun-->>-Job: simulation_data (autograd-compatible) Job-->>-User: simulation_data Note over User, Server: Batch.run() - Multiple Simulations Flow User->>+Batch: batch.run(path_dir="./data/") Note over Batch: Batch.run() now calls autograd.run_async() internally Batch->>+AutogradRunAsync: run_async(simulations, path_dir, ...) Note over AutogradRunAsync: Checks if all simulations are autograd-compatible alt All simulations have autograd tracers AutogradRunAsync->>AutogradRunAsync: _run_async_primitive() - forward pass AutogradRunAsync->>+WebAPI: _run_async_tidy3d() - actual batch WebAPI->>+Server: batch upload/start/monitor/download Server-->>-WebAPI: batch_data WebAPI-->>-AutogradRunAsync: batch_data AutogradRunAsync->>AutogradRunAsync: postprocess batch - insert tracers else No autograd tracers in simulations AutogradRunAsync->>+WebAPI: run_async_webapi() - regular batch WebAPI->>+Server: batch upload/start/monitor/download Server-->>-WebAPI: batch_data WebAPI-->>-AutogradRunAsync: batch_data end AutogradRunAsync-->>-Batch: batch_data (autograd-compatible) Batch-->>-User: batch_data Note over User, Server: Gradient Computation (when tracers present) User->>+AutogradRun: autograd.grad(objective)(params) Note over AutogradRun: VJP backward pass triggered AutogradRun->>AutogradRun: _run_bwd() - construct adjoint sims AutogradRun->>+WebAPI: run adjoint simulations WebAPI->>+Server: adjoint upload/start/monitor/download Server-->>-WebAPI: adjoint_data WebAPI-->>-AutogradRun: adjoint_data AutogradRun->>AutogradRun: postprocess_adj() - compute gradients AutogradRun-->>-User: gradients