add priority to Bach, Job, and run_async

tests/test_web/test_webapi.py

@@ -673,7 +673,7 @@ def track_to_file(self, fname):
         return original_to_file(self, fname)
 
     # mock start to interrupt run() after upload and to_file
-    def mock_start_interrupt(self):
+    def mock_start_interrupt(self, *args, **kwargs):
         # 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"

tidy3d/web/api/asynchronous.py

@@ -23,6 +23,7 @@ def run_async(
     parent_tasks: Optional[dict[str, list[str]]] = None,
     reduce_simulation: Literal["auto", True, False] = "auto",
     pay_type: Union[PayType, str] = PayType.AUTO,
+    priority: Optional[int] = None,
 ) -> 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.
@@ -52,7 +53,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: Union[PayType, str] = PayType.AUTO
         Specify the payment method.
-
+    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 doesn’t impact simulations using FlexCredits..
     Returns
     ------
     :class:`BatchData`
@@ -90,5 +93,5 @@ def run_async(
         pay_type=pay_type,
     )
 
-    batch_data = batch.run(path_dir=path_dir)
+    batch_data = batch.run(path_dir=path_dir, priority=priority)
     return batch_data

tidy3d/web/api/container.py

@@ -217,21 +217,25 @@ def to_file(self, fname: str) -> None:
         self = self.updated_copy(task_id_cached=task_id_cached)
         super(Job, self).to_file(fname=fname)  # noqa: UP008
 
-    def run(self, path: str = DEFAULT_DATA_PATH) -> SimulationDataType:
+    def run(
+        self, path: str = DEFAULT_DATA_PATH, priority: Optional[int] = None
+    ) -> SimulationDataType:
         """Run :class:`Job` all the way through and return data.
 
         Parameters
         ----------
         path_dir : str = "./simulation_data.hdf5"
             Base directory where data will be downloaded, by default current working directory.
-
+        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 doesn’t impact simulations using FlexCredits.
         Returns
         -------
         Union[:class:`.SimulationData`, :class:`.HeatSimulationData`, :class:`.EMESimulationData`]
             Object containing simulation results.
         """
         self.upload()
-        self.start()
+        self.start(priority=priority)
         self.monitor()
         return self.load(path=path)
 
@@ -270,14 +274,25 @@ def status(self):
         """Return current status of :class:`Job`."""
         return self.get_info().status
 
-    def start(self) -> None:
+    def start(self, priority: Optional[int] = None) -> None:
         """Start running a :class:`Job`.
 
+        Parameters
+        ----------
+
+        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 doesn’t impact simulations using FlexCredits.
         Note
         ----
         To monitor progress of the :class:`Job`, call :meth:`Job.monitor` after started.
         """
-        web.start(self.task_id, solver_version=self.solver_version, pay_type=self.pay_type)
+        web.start(
+            self.task_id,
+            solver_version=self.solver_version,
+            pay_type=self.pay_type,
+            priority=priority,
+        )
 
     def get_run_info(self) -> RunInfo:
         """Return information about the running :class:`Job`.
@@ -571,14 +586,20 @@ class Batch(WebContainer):
 
     _job_type = Job
 
-    def run(self, path_dir: str = DEFAULT_DATA_DIR) -> BatchData:
+    def run(
+        self,
+        path_dir: str = DEFAULT_DATA_DIR,
+        priority: Optional[int] = None,
+    ) -> BatchData:
         """Upload and run each simulation in :class:`Batch`.
 
         Parameters
         ----------
         path_dir : str
             Base directory where data will be downloaded, by default current working directory.
-
+        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 doesn’t impact simulations using FlexCredits.
         Returns
         ------
         :class:`BatchData`
@@ -602,7 +623,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.start(priority=priority)
         self.monitor()
         return self.load(path_dir=path_dir)
 
@@ -705,9 +726,18 @@ def get_info(self) -> dict[TaskName, TaskInfo]:
             info_dict[task_name] = task_info
         return info_dict
 
-    def start(self) -> None:
+    def start(
+        self,
+        priority: Optional[int] = None,
+    ) -> None:
         """Start running all tasks in the :class:`Batch`.
 
+        Parameters
+        ----------
+
+        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 doesn’t impact simulations using FlexCredits.
         Note
         ----
         To monitor the running simulations, can call :meth:`Batch.monitor`.
@@ -718,7 +748,7 @@ def start(self) -> None:
 
         with ThreadPoolExecutor(max_workers=self.num_workers) as executor:
             for _, job in self.jobs.items():
-                executor.submit(job.start)
+                executor.submit(job.start, priority=priority)
 
     def get_run_info(self) -> dict[TaskName, RunInfo]:
         """get information about a each of the tasks in the :class:`Batch`.

tidy3d/web/api/webapi.py

@@ -393,7 +393,8 @@ def start(
     pay_type: Union[PayType, str] = PayType.AUTO
         Which method to pay the simulation
     priority: int = None
-        Task priority for vGPU queue (1=lowest, 10=highest).
+        Priority of the simulation in the Virtual GPU (vGPU) queue (1 = lowest, 10 = highest).
+        It affects only simulations from vGPU licenses and doesn’t impact simulations using FlexCredits.
     Note
     ----
     To monitor progress, can call :meth:`monitor` after starting simulation.

tidy3d/web/core/task_core.py

@@ -447,7 +447,8 @@ def submit(
         pay_type: Union[PayType, str] = PayType.AUTO
             Which method to pay the simulation.
         priority: int = None
-            Task priority for vGPU queue (1=lowest, 10=highest).
+            Priority of the simulation in the Virtual GPU (vGPU) queue (1 = lowest, 10 = highest).
+            It affects only simulations from vGPU licenses and doesn’t impact simulations using FlexCredits.
         """
         pay_type = PayType(pay_type) if not isinstance(pay_type, PayType) else pay_type