add priority to Bach, Job, and run_async

[QingengWei]

Greptile Summary

Updated On: 2025-09-03 11:18:45 UTC

This PR adds priority parameter support to the Tidy3D web API's batch processing system, specifically for vGPU queue management. The implementation enables users to specify task priority levels from 1 (lowest) to 10 (highest) when submitting simulations to the vGPU queue, with priority only affecting vGPU licenses and not FlexCredits simulations.

The changes follow a clear architectural pattern, threading the priority parameter through the entire web API call chain: run_async() → Batch.run() → Batch.start() → individual Job.start() → webapi.start() → SimulationTask.submit(). The parameter is consistently implemented as optional with None as the default value, maintaining backward compatibility throughout the system.

Key modifications include:

• Core web API functions (webapi.py): Added priority validation (1-10 range) and proper parameter forwarding to the underlying task submission layer
• Job and Batch classes (container.py): Both run() and start() methods now accept priority parameters with proper ThreadPoolExecutor implementation ensuring all jobs in a batch inherit the same priority
• High-level interfaces (asynchronous.py): The run_async() function supports priority specification and passes it through to batch operations
• Task submission layer (task_core.py): SimulationTask.submit() includes priority in the HTTP request payload sent to the server
• Autograd module (autograd.py): Priority parameter added to wrapper functions, though implementation has gaps in internal execution paths
• Test infrastructure (test_webapi.py): Updated mock functions to handle the new parameter signature

The implementation integrates well with existing codebase architecture and follows established patterns for parameter passing. Documentation clearly specifies the vGPU-specific nature of this feature, helping users understand when priority applies.

Important Files Changed

Changed Files

Filename	Score	Overview
tidy3d/web/core/task_core.py	5/5	Added priority parameter to SimulationTask.submit() with proper HTTP payload integration
tidy3d/web/api/webapi.py	4/5	Added priority support to core webapi functions with validation and clear documentation
tidy3d/web/api/container.py	4/5	Added priority parameter to Job and Batch classes with threading implementation
tidy3d/web/api/asynchronous.py	4/5	Added priority parameter to run_async function with proper batch integration
tidy3d/web/api/autograd/autograd.py	2/5	Added priority parameter but with incomplete implementation in autograd-specific paths
tests/test_web/test_webapi.py	5/5	Updated mock function to accept variable arguments for test compatibility

Confidence score: 3/5

• This PR requires careful review due to incomplete implementation in the autograd module that could cause inconsistent behavior
• Score lowered due to critical issue where priority parameter is not passed through autograd-specific execution paths in _run_tidy3d() and _run_async_tidy3d() calls
• Pay close attention to tidy3d/web/api/autograd/autograd.py where priority is missing from several internal function calls that could result in priority being ignored in autograd workflows

Sequence Diagram

sequenceDiagram
    participant User
    participant run_async
    participant Batch
    participant Job
    participant webapi
    participant SimulationTask
    
    User->>run_async: "run_async(simulations, priority=5)"
    run_async->>Batch: "Batch(simulations, folder_name, ...)"
    run_async->>Batch: "batch.run(path_dir, priority)"
    Batch->>Batch: "upload()"
    Batch->>Batch: "start(priority)"
    loop for each job in batch
        Batch->>Job: "job.start(priority)"
        Job->>webapi: "web.start(task_id, solver_version, pay_type, priority)"
        webapi->>SimulationTask: "task.submit(solver_version, worker_group, pay_type, priority)"
        SimulationTask->>SimulationTask: "HTTP POST with priority in payload"
    end
    Batch->>Batch: "monitor()"
    Batch->>User: "return BatchData"

Loading

[QingengWei]

@e-g-melo @momchil-flex @tylerflex The pr is for https://flow360.atlassian.net/browse/SCEM-10928

[greptile-apps]

_{2 files reviewed, 6 comments}

_{Edit Code Review Bot Settings | Greptile}

[github-actions]

Diff Coverage

Diff: origin/develop...HEAD, staged and unstaged changes

• tidy3d/web/api/asynchronous.py (100%)
• tidy3d/web/api/autograd/autograd.py (33.3%): Missing lines 310,1298,1318,1320,1337-1338
• tidy3d/web/api/container.py (78.6%): Missing lines 251,642,770

Summary

• Total: 24 lines
• Missing: 9 lines
• Coverage: 62%

tidy3d/web/api/autograd/autograd.py

Lines 306-314

  306         Interface for submitting several :class:`Simulation` objects to sever.
  307     """
  308     # validate priority if specified
  309     if priority is not None and (priority < 1 or priority > 10):
! 310         raise ValueError("Priority must be between '1' and '10' if specified.")
  311 
  312     if is_valid_for_autograd_async(simulations):
  313         return _run_async(
  314             simulations=simulations,

---

Lines 1294-1302

  1294     """Run a batch of simulations using regular web.run()."""
  1295 
  1296     batch_init_kwargs = parse_run_kwargs(**run_kwargs)
  1297     path_dir = run_kwargs.pop("path_dir", None)
! 1298     priority = run_kwargs.get("priority")
  1299     batch = Batch(simulations=simulations, **batch_init_kwargs)
  1300     td.log.info(f"running {batch.simulation_type} batch with '_run_async_tidy3d()'")
  1301 
  1302     if batch.simulation_type == "autograd_fwd":

---

Lines 1314-1324

  1314             task_id = task_ids[task_name]
  1315             upload_sim_fields_keys(sim_fields_keys, task_id=task_id, verbose=verbose)
  1316 
  1317     if path_dir:
! 1318         batch_data = batch.run(path_dir, priority=priority)
  1319     else:
! 1320         batch_data = batch.run(priority=priority)
  1321 
  1322     task_ids = {key: job.task_id for key, job in batch.jobs.items()}
  1323     return batch_data, task_ids

---

Lines 1333-1342

  1333     _ = run_kwargs.pop("path_dir", None)
  1334     batch = Batch(simulations=simulations, **batch_init_kwargs)
  1335     td.log.info(f"running {batch.simulation_type} batch with '_run_async_tidy3d_bwd()'")
  1336 
! 1337     priority = run_kwargs.get("priority")
! 1338     batch.start(priority=priority)
  1339     batch.monitor()
  1340 
  1341     vjp_traced_fields_dict = {}
  1342     for task_name, job in batch.jobs.items():

---

tidy3d/web/api/container.py

Lines 247-255

  247         self.upload()
  248         if priority is None:
  249             self.start()
  250         else:
! 251             self.start(priority=priority)
  252         self.monitor()
  253         return self.load(path=path)
  254 
  255     @cached_property

---

Lines 638-646

  638         self.to_file(self._batch_path(path_dir=path_dir))
  639         if priority is None:
  640             self.start()
  641         else:
! 642             self.start(priority=priority)
  643         self.monitor()
  644         return self.load(path_dir=path_dir)
  645 
  646     @cached_property

---

Lines 766-774

  766             for _, job in self.jobs.items():
  767                 if priority is None:
  768                     executor.submit(job.start)
  769                 else:
! 770                     executor.submit(job.start, priority=priority)
  771 
  772     def get_run_info(self) -> dict[TaskName, RunInfo]:
  773         """get information about a each of the tasks in the :class:`Batch`.

---

[greptile-apps]

_{6 files reviewed, no comments}

_{Edit Code Review Bot Settings | Greptile}

[greptile-apps]

_{6 files reviewed, 1 comment}

_{Edit Code Review Bot Settings | Greptile}