feat(design): add runtime priority control

Author: yaugenst-flex

CHANGELOG.md

@@ -33,6 +33,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 - Added a Gaussian inverse design filter option with autograd gradients and complete padding mode coverage.
 - Added support for argument passing to DRC file when running checks with `DRCRunner.run(..., drc_args={key: value})` in klayout plugin.
 - Added support for `nonlinear_spec` in `CustomMedium` and `CustomDispersiveMedium`.
+- `tidy3d.plugins.design.DesignSpace.run(..., fn_post=...)` now accepts a `priority` keyword to propagate vGPU queue priority to all automatically batched simulations.
 
 ### Breaking Changes
 - Edge singularity correction at PEC and lossy metal edges defaults to `True`.

tests/test_plugins/test_design.py

@@ -463,6 +463,22 @@ def test_sweep(sweep_method, monkeypatch):
     assert float_label in sweep_results_df, "didn't assign column header properly for float"
 
 
+def test_priority_forwarded_to_batch(monkeypatch):
+    captured_priority = {}
+
+    def run_with_priority(batch, path_dir: Optional[str] = None, priority: Optional[int] = None):
+        captured_priority["value"] = priority
+        return emulated_batch_run(batch, path_dir=path_dir)
+
+    monkeypatch.setattr(web.Batch, "run", run_with_priority)
+
+    design_space = init_design_space(sweep_method=tdd.MethodGrid())
+
+    design_space.run(scs_pre, scs_post, verbose=False, priority=7)
+
+    assert captured_priority["value"] == 7
+
+
 def emulated_estimate_cost_return(self, verbose=True):
     return 0.5
 

tidy3d/plugins/design/README.md

@@ -95,6 +95,10 @@ With the design parameters and our method defined, we can combine everything int
 design_space = tdd.DesignSpace(parameters=[param_n, param_r], method=method)
 ```
 
+Execution-time options such as vGPU queue priority are supplied when calling `run`. For example,
+`design_space.run(fn_pre, fn_post, priority=7)` forwards that priority to every automatically batched
+simulation (only available when both `fn` and `fn_post` are provided).
+
 ## Results
 
 The `DesignSpace.run()` function returns a `Result` object, which is basically a dataset containing the function inputs, outputs, source code, and any task ID information corresponding to each data point.

tidy3d/plugins/design/design.py

@@ -144,7 +144,13 @@ def get_fn_source(function: Callable) -> str:
         except (TypeError, OSError):
             return None
 
-    def run(self, fn: Callable, fn_post: Optional[Callable] = None, verbose: bool = True) -> Result:
+    def run(
+        self,
+        fn: Callable,
+        fn_post: Optional[Callable] = None,
+        verbose: bool = True,
+        priority: Optional[int] = None,
+    ) -> Result:
         """Explore a parameter space with a supplied method using the user supplied function.
         Supplied functions are used to evaluate the design space and are called within the method.
         For optimization methods these functions act as the fitness function. A single function can be
@@ -209,26 +215,43 @@ def run(self, fn: Callable, fn_post: Optional[Callable] = None, verbose: bool =
         verbose : bool = True
             Toggle the output of statements stored in the logging console.
 
+        priority : int = None
+            Optional vGPU queue priority (1 = lowest, 10 = highest) applied to simulations run via
+            ``fn`` / ``fn_post``. Ignored when ``fn_post`` is not provided because no automatic
+            batching occurs in that mode.
+
         Returns
         -------
         :class:`Result`
             Object containing the results of the design space exploration.
             Can be converted to ``pandas.DataFrame`` with ``.to_dataframe()``.
         """
 
+        if priority is not None and (priority < 1 or priority > 10):
+            raise ValueError("'priority' must be between 1 and 10 (inclusive).")
+
         # Get the console
         # Method.run checks for console is None instead of being passed console and verbose
         console = get_logging_console() if verbose else None
 
         # Run based on how many functions the user provides
         if fn_post is None:
+            if priority is not None:
+                log.warning(
+                    "'priority' has no effect unless both fn and fn_post are provided;"
+                    " please split your function to allow automatic batching.",
+                    log_once=True,
+                )
             fn_args, fn_values, aux_values, opt_output = self.run_single(fn, console)
             sim_names = None
             sim_paths = None
 
         else:
             fn_args, fn_values, aux_values, opt_output, sim_names, sim_paths = self.run_pre_post(
-                fn_pre=fn, fn_post=fn_post, console=console
+                fn_pre=fn,
+                fn_post=fn_post,
+                console=console,
+                priority=priority,
             )
 
             if len(sim_names) == 0:
@@ -252,12 +275,20 @@ def run_single(self, fn: Callable, console: Console) -> tuple(list[dict], list,
         evaluate_fn = self._get_evaluate_fn_single(fn=fn)
         return self.method._run(run_fn=evaluate_fn, parameters=self.parameters, console=console)
 
-    def run_pre_post(self, fn_pre: Callable, fn_post: Callable, console: Console) -> tuple(
-        list[dict], list[dict], list[Any]
-    ):
+    def run_pre_post(
+        self,
+        fn_pre: Callable,
+        fn_post: Callable,
+        console: Console,
+        priority: Optional[int] = None,
+    ) -> tuple(list[dict], list[dict], list[Any]):
         """Run a function with Tidy3D implicitly called in between."""
         handler = self._get_evaluate_fn_pre_post(
-            fn_pre=fn_pre, fn_post=fn_post, fn_mid=self._fn_mid, console=console
+            fn_pre=fn_pre,
+            fn_post=fn_post,
+            fn_mid=self._fn_mid,
+            console=console,
+            priority=priority,
         )
         fn_args, fn_values, aux_values, opt_output = self.method._run(
             run_fn=handler.fn_combined, parameters=self.parameters, console=console
@@ -276,16 +307,22 @@ def evaluate(args_list: list) -> list[Any]:
         return evaluate
 
     def _get_evaluate_fn_pre_post(
-        self, fn_pre: Callable, fn_post: Callable, fn_mid: Callable, console: Console
+        self,
+        fn_pre: Callable,
+        fn_post: Callable,
+        fn_mid: Callable,
+        console: Console,
+        priority: Optional[int],
     ):
         """Get function that tries to use batch processing on a set of arguments."""
 
         class Pre_Post_Handler:
-            def __init__(self, console) -> None:
+            def __init__(self, console, priority) -> None:
                 self.sim_counter = 0
                 self.sim_names = []
                 self.sim_paths = []
                 self.console = console
+                self.priority = priority
 
             def fn_combined(self, args_list: list[dict[str, Any]]) -> list[Any]:
                 """Compute fn_pre and fn_post functions and capture other outputs."""
@@ -300,26 +337,37 @@ def fn_combined(self, args_list: list[dict[str, Any]]) -> list[Any]:
                     )
 
                 data, task_names, task_paths, sim_counter = fn_mid(
-                    sim_dict, self.sim_counter, self.console
+                    sim_dict,
+                    self.sim_counter,
+                    self.console,
+                    self.priority,
                 )
                 self.sim_names.extend(task_names)
                 self.sim_paths.extend(task_paths)
                 self.sim_counter = sim_counter
                 post_out = [fn_post(val) for val in data.values()]
                 return post_out
 
-        handler = Pre_Post_Handler(console)
+        handler = Pre_Post_Handler(console, priority)
 
         return handler
 
     @staticmethod
-    def _run_batch(batch: Batch, path_dir: str) -> BatchData:
+    def _run_batch(
+        batch: Batch,
+        path_dir: str,
+        priority: Optional[int] = None,
+    ) -> BatchData:
         """Run a batch and return the BatchData."""
-        batch_out = batch.run(path_dir=path_dir)
+        batch_out = batch.run(path_dir=path_dir, priority=priority)
         return batch_out
 
     def _fn_mid(
-        self, pre_out: dict[int, Any], sim_counter: int, console: Console
+        self,
+        pre_out: dict[int, Any],
+        sim_counter: int,
+        console: Console,
+        priority: Optional[int],
     ) -> Union[dict[int, Any], BatchData]:
         """A function of the output of ``fn_pre`` that gives the input to ``fn_post``."""
 
@@ -399,11 +447,11 @@ def _find_and_map(
             simulation_type="tidy3d_design",
             verbose=False,  # Using a custom output instead of Batch.monitor updates
         )
-        sims_out = self._run_batch(batch, path_dir=self.path_dir)
+        sims_out = self._run_batch(batch, path_dir=self.path_dir, priority=priority)
 
         batch_results = {}
         for batch_key, batch in batches.items():
-            batch_out = self._run_batch(batch, path_dir=self.path_dir)
+            batch_out = self._run_batch(batch, path_dir=self.path_dir, priority=priority)
             batch_results[batch_key] = batch_out
 
         def _return_to_dict(return_dict: dict, key: str, return_obj: Any) -> None:
@@ -468,6 +516,7 @@ def run_batch(
             Union[SimulationData, list[SimulationData], dict[str, SimulationData]], Any
         ],
         path_dir: str = ".",
+        priority: Optional[int] = None,
         **batch_kwargs: Any,
     ) -> Result:
         """
@@ -485,7 +534,7 @@ def run_batch(
             )
 
         new_self = self.updated_copy(path_dir=path_dir)
-        result = new_self.run(fn=fn_pre, fn_post=fn_post)
+        result = new_self.run(fn=fn_pre, fn_post=fn_post, priority=priority)
 
         return result