Select between virtual GPU and regular queue on Python client

tidy3d/plugins/adjoint/web.py

@@ -85,6 +85,7 @@ def _run(
     path: str = "simulation_data.hdf5",
     callback_url: str = None,
     verbose: bool = True,
+    use_credits: bool = None,
 ) -> JaxSimulationData:
     """Split the provided ``JaxSimulation`` into a regular ``Simulation`` and a ``JaxInfo`` part,
     run using ``tidy3d_run_fn``, which runs on the server by default but can be monkeypatched,
@@ -111,6 +112,7 @@ def run(
     path: str = "simulation_data.hdf5",
     callback_url: str = None,
     verbose: bool = True,
+    use_credits: bool = None,
 ) -> JaxSimulationData:
     """Submits a :class:`.JaxSimulation` to server, starts running, monitors progress, downloads,
     and loads results as a :class:`.JaxSimulationData` object.
@@ -131,7 +133,10 @@ def run(
         fields ``{'id', 'status', 'name', 'workUnit', 'solverVersion'}``.
     verbose : bool = True
         If `True`, will print progressbars and status, otherwise, will run silently.
-
+    use_credits: bool = None
+        None:  Run with reserved GPU if user has the license, otherwise run with credits pay.
+        True: Pay for the task with credits;
+        False: Run with reserved GPU if user has the license.
     Returns
     -------
     :class:`.JaxSimulationData`
@@ -148,6 +153,7 @@ def run(
         path=path,
         callback_url=callback_url,
         verbose=verbose,
+        use_credits=use_credits,
     )
 
 
@@ -271,16 +277,23 @@ class AdjointJob(Job):
         description="Container of information needed to reconstruct jax simulation.",
     )
 
-    def start(self) -> None:
+    def start(self, use_credits: bool = None) -> None:
         """Start running a :class:`AdjointJob`. after uploading jax info.
 
+        Parameters
+        ----------
+
+        use_credits: bool = None
+            None:  Run with reserved GPU if user has the license, otherwise run with credits pay.
+            True: Pay for the task with credits;
+            False: Run with reserved GPU if user has the license.
         Note
         ----
         To monitor progress of the :class:`Job`, call :meth:`Job.monitor` after started.
         """
         if self.jax_info is not None:
             upload_jax_info(task_id=self.task_id, jax_info=self.jax_info, verbose=self.verbose)
-        super().start()
+        super().start(use_credits=use_credits)
 
 
 class AdjointBatch(Batch):
@@ -309,17 +322,25 @@ class AdjointBatch(Batch):
 
     _job_type = AdjointJob
 
-    def start(self) -> None:
+    def start(self,use_credits: bool = None,) -> None:
         """Start running a :class:`AdjointBatch`. after uploading jax info for each job.
 
+        Parameters
+        ----------
+
+        use_credits: bool = None
+            None:  Run with reserved GPU if user has the license, otherwise run with credits pay.
+            True: Pay for the task with credits;
+            False: Run with reserved GPU if user has the license.
+
         Note
         ----
         To monitor progress of the :class:`Batch`, call :meth:`Batch.monitor` after started.
         """
         for task_name, job in self.jobs.items():
             jax_info = self.jax_infos.get(task_name)
             upload_jax_info(task_id=job.task_id, jax_info=jax_info, verbose=job.verbose)
-        super().start()
+        super().start(use_credits=use_credits)
 
 
 def webapi_run_adjoint_fwd(
@@ -588,6 +609,7 @@ def webapi_run_async_adjoint_bwd(
     callback_url: str,
     verbose: bool,
     parent_tasks: List[List[str]],
+    use_credits: bool = None,
 ) -> List[JaxSimulation]:
     """Runs the forward simulations on our servers, stores the gradient data for later."""
 
@@ -635,6 +657,7 @@ def run_local(
     callback_url: str = None,
     verbose: bool = True,
     num_proc: int = NUM_PROC_LOCAL,
+    use_credits: bool = None,
 ) -> JaxSimulationData:
     """Submits a :class:`.JaxSimulation` to server, starts running, monitors progress, downloads,
     and loads results as a :class:`.JaxSimulationData` object.
@@ -658,6 +681,10 @@ def run_local(
     num_proc: int = 1
         Number of processes to use for the gradient computations.
 
+    use_credits: bool = None
+        None:  Run with reserved GPU if user has the license, otherwise run with credits pay.
+        True: Pay for the task with credits;
+        False: Run with reserved GPU if user has the license.
     Returns
     -------
     :class:`.JaxSimulationData`
@@ -675,6 +702,7 @@ def run_local(
         path=path,
         callback_url=callback_url,
         verbose=verbose,
+        use_credits=use_credits,
     )
 
     # convert back to jax type and return

tidy3d/web/api/container.py

@@ -210,21 +210,26 @@ 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, use_credits: bool = 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.
 
+        use_credits: bool = None
+            None:  Run with reserved GPU if user has the license, otherwise run with credits pay.
+            True: Pay for the task with credits;
+            False: Run with reserved GPU if user has the license.
+
         Returns
         -------
         Union[:class:`.SimulationData`, :class:`.HeatSimulationData`, :class:`.EMESimulationData`]
             Object containing simulation results.
         """
         self.upload()
-        self.start()
+        self.start(use_credits=use_credits)
         self.monitor()
         return self.load(path=path)
 
@@ -263,14 +268,25 @@ def status(self):
         """Return current status of :class:`Job`."""
         return self.get_info().status
 
-    def start(self) -> None:
+    def start(
+        self,
+        use_credits: bool = None,
+    ) -> None:
         """Start running a :class:`Job`.
 
+        Parameters
+        ----------
+
+        use_credits: bool = None
+            None:  Run with reserved GPU if user has the license, otherwise run with credits pay.
+            True: Pay for the task with credits;
+            False: Run with reserved GPU if user has the license.
+
         Note
         ----
         To monitor progress of the :class:`Job`, call :meth:`Job.monitor` after started.
         """
-        web.start(self.task_id, solver_version=self.solver_version)
+        web.start(self.task_id, solver_version=self.solver_version, use_credits=use_credits)
 
     def get_run_info(self) -> RunInfo:
         """Return information about the running :class:`Job`.
@@ -561,14 +577,21 @@ 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,
+        use_credits: bool = 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.
-
+        use_credits: bool = None
+            None:  Run with reserved GPU if user has the license, otherwise run with credits pay.
+            True: Pay for the task with credits;
+            False: Run with reserved GPU if user has the license.
         Returns
         ------
         :class:`BatchData`
@@ -591,7 +614,7 @@ def run(self, path_dir: str = DEFAULT_DATA_DIR) -> BatchData:
         """
         self._check_path_dir(path_dir)
         self.upload()
-        self.start()
+        self.start(use_credits=use_credits)
         self.monitor()
         self.download(path_dir=path_dir)
         return self.load(path_dir=path_dir)
@@ -694,13 +717,26 @@ def get_info(self) -> Dict[TaskName, TaskInfo]:
             info_dict[task_name] = task_info
         return info_dict
 
-    def start(self) -> None:
+    def start(
+        self,
+        use_credits: bool = None,
+    ) -> None:
         """Start running all tasks in the :class:`Batch`.
 
+        Parameters
+        ----------
+
+        use_credits: bool = None
+            None:  Run with reserved GPU if user has the license, otherwise run with credits pay.
+            True: Pay for the task with credits;
+            False: Run with reserved GPU if user has the license.
+
         Note
         ----
         To monitor the running simulations, can call :meth:`Batch.monitor`.
         """
+        for _, job in self.jobs.items():
+            job.start(use_credits=use_credits)
         if self.verbose:
             console = get_logging_console()
             console.log(f"Started working on Batch containing {self.num_jobs} tasks.")

tidy3d/web/api/mode.py

@@ -56,6 +56,7 @@ def run(
     progress_callback_upload: Callable[[float], None] = None,
     progress_callback_download: Callable[[float], None] = None,
     reduce_simulation: Literal["auto", True, False] = "auto",
+    use_credits: bool = None,
 ) -> ModeSolverData:
     """Submits a :class:`.ModeSolver` to server, starts running, monitors progress, downloads,
     and loads results as a :class:`.ModeSolverData` object.
@@ -81,6 +82,10 @@ def run(
     reduce_simulation : Literal["auto", True, False] = "auto"
         Restrict simulation to mode solver region. If "auto", then simulation is automatically
         restricted if it contains custom mediums.
+    use_credits: bool = None
+        None:  Run with reserved GPU if user has the license, otherwise run with credits pay.
+        True: Pay for the task with credits;
+        False: Run with reserved GPU if user has the license.
     Returns
     -------
     :class:`.ModeSolverData`
@@ -114,7 +119,7 @@ def run(
             f"Mode solver created with task_id='{task.task_id}', solver_id='{task.solver_id}'."
         )
     task.upload(verbose=verbose, progress_callback=progress_callback_upload)
-    task.submit()
+    task.submit(use_credits=use_credits)
 
     # Wait for task to finish
     prev_status = "draft"
@@ -461,15 +466,29 @@ def upload(
         finally:
             os.unlink(file_name)
 
-    def submit(self):
+    def submit(
+        self,
+        use_credits: bool = None,
+    ):
         """Start the execution of this task.
 
         The mode solver must be uploaded to the server with the :meth:`ModeSolverTask.upload` method
         before this step.
+
+        Parameters
+        ----------
+
+        use_credits: bool = None
+            None:  Run with reserved GPU if user has the license, otherwise run with credits pay.
+            True: Pay for the task with credits;
+            False: Run with reserved GPU if user has the license.
         """
         http.post(
             f"{MODESOLVER_API}/{self.task_id}/{self.solver_id}/run",
-            {"enableCaching": Env.current.enable_caching},
+            {
+                "enableCaching": Env.current.enable_caching,
+                "useCredits": use_credits,
+            },
         )
 
     def delete(self):

tidy3d/web/api/webapi.py

@@ -85,6 +85,7 @@ def run(
     simulation_type: str = "tidy3d",
     parent_tasks: list[str] = None,
     reduce_simulation: Literal["auto", True, False] = "auto",
+    use_credits: bool = None,
 ) -> SimulationDataType:
     """
     Submits a :class:`.Simulation` to server, starts running, monitors progress, downloads,
@@ -117,6 +118,10 @@ def run(
         worker group
     reduce_simulation : Literal["auto", True, False] = "auto"
         Whether to reduce structures in the simulation to the simulation domain only. Note: currently only implemented for the mode solver.
+    use_credits: bool = None
+        None:  Run with reserved GPU if user has the license, otherwise run with credits pay.
+        True: Pay for the task with credits;
+        False: Run with reserved GPU if user has the license.
 
     Returns
     -------
@@ -178,6 +183,7 @@ def run(
         task_id,
         solver_version=solver_version,
         worker_group=worker_group,
+        use_credits=use_credits,
     )
     monitor(task_id, verbose=verbose)
     data = load(
@@ -371,6 +377,7 @@ def start(
     task_id: TaskId,
     solver_version: str = None,
     worker_group: str = None,
+    use_credits: bool = None,
 ) -> None:
     """Start running the simulation associated with task.
 
@@ -383,6 +390,11 @@ def start(
         target solver version.
     worker_group: str = None
         worker group
+    use_credits: bool = None
+        None:  Run with reserved GPU if user has the license, otherwise run with credits pay.
+        True: Pay for the task with credits;
+        False: Run with reserved GPU if user has the license.
+
     Note
     ----
     To monitor progress, can call :meth:`monitor` after starting simulation.
@@ -393,6 +405,7 @@ def start(
     task.submit(
         solver_version=solver_version,
         worker_group=worker_group,
+        use_credits=use_credits,
     )
 
 

tidy3d/web/core/task_core.py

@@ -413,6 +413,7 @@ def submit(
         self,
         solver_version: str = None,
         worker_group: str = None,
+        use_credits: bool = None,
     ):
         """Kick off this task.
 
@@ -426,6 +427,10 @@ def submit(
             target solver version.
         worker_group: str = None
             worker group
+        use_credits: bool = None
+            None:  Run with reserved GPU if user has the license, otherwise run with credits pay.
+            True: Pay for the task with credits;
+            False: Run with reserved GPU if user has the license.
         """
 
         if solver_version:
@@ -440,6 +445,7 @@ def submit(
                 "workerGroup": worker_group,
                 "protocolVersion": protocol_version,
                 "enableCaching": Env.current.enable_caching,
+                "useCredits": use_credits,
             },
         )