feat(web): add priority parameter for vGPU task queue management

Author: yaugenst-flex

CHANGELOG.md

@@ -10,6 +10,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 ### Added
 - Add support for `np.unwrap` in `tidy3d.plugins.autograd`.
 - Add Nunley variant to germanium material library based on Nunley et al. 2016 data.
+- Added `priority` parameter to `web.run()` and related functions to allow vGPU users to set task priority (1-10) in the queue.
 
 ### Changed
 - Switched to an analytical gradient calculation for spatially-varying pole-residue models (`CustomPoleResidue`).

tests/test_web/test_tidy3d_task.py

@@ -218,6 +218,7 @@ def test_submit(set_api_key):
                     "workerGroup": None,
                     "enableCaching": Env.current.enable_caching,
                     "payType": PayType.AUTO,
+                    "priority": None,
                 }
             )
         ],

tests/test_web/test_webapi.py

@@ -182,29 +182,36 @@ def mock_get_info(monkeypatch, set_api_key):
 def mock_start(monkeypatch, set_api_key, mock_get_info):
     """Mocks webapi.start."""
 
-    responses.add(
-        responses.POST,
-        f"{Env.current.web_api_endpoint}/tidy3d/tasks/{TASK_ID}/submit",
-        match=[
-            matchers.json_params_matcher(
-                {
-                    "solverVersion": None,
-                    "workerGroup": None,
-                    "protocolVersion": td.version.__version__,
-                    "enableCaching": Env.current.enable_caching,
-                    "payType": PayType.AUTO,
+    def add_mock_response(priority=None):
+        expected_body = {
+            "solverVersion": None,
+            "workerGroup": None,
+            "protocolVersion": td.version.__version__,
+            "enableCaching": Env.current.enable_caching,
+            "payType": PayType.AUTO,
+            "priority": priority,
+        }
+
+        responses.add(
+            responses.POST,
+            f"{Env.current.web_api_endpoint}/tidy3d/tasks/{TASK_ID}/submit",
+            match=[matchers.json_params_matcher(expected_body)],
+            json={
+                "data": {
+                    "taskId": TASK_ID,
+                    "taskName": TASK_NAME,
+                    "createdAt": CREATED_AT,
                 }
-            )
-        ],
-        json={
-            "data": {
-                "taskId": TASK_ID,
-                "taskName": TASK_NAME,
-                "createdAt": CREATED_AT,
-            }
-        },
-        status=200,
-    )
+            },
+            status=200,
+        )
+
+    # Add response for calls without priority
+    add_mock_response(None)
+
+    # Add responses for calls with specific priority values
+    for priority in [1, 5, 10]:
+        add_mock_response(priority)
 
 
 @pytest.fixture
@@ -322,6 +329,39 @@ def test_start(mock_start):
     start(TASK_ID)
 
 
+@responses.activate
+@pytest.mark.parametrize("priority", [1, 5, 10, None])
+def test_start_with_valid_priority(mock_start, priority):
+    """Test start with valid priority values."""
+    start(TASK_ID, priority=priority)
+
+
+@responses.activate
+@pytest.mark.parametrize("priority", [0, -1, 11, 15])
+def test_start_with_invalid_priority(mock_start, priority):
+    """Test start with invalid priority values."""
+    with pytest.raises(ValueError, match="Priority must be between '1' and '10' if specified."):
+        start(TASK_ID, priority=priority)
+
+
+@responses.activate
+@pytest.mark.parametrize("priority", [5, None])
+def test_run_with_valid_priority(mock_webapi, monkeypatch, priority):
+    """Test run with valid priority parameter."""
+    monkeypatch.setattr(f"{api_path}.load", lambda *args, **kwargs: True)
+    sim = make_sim()
+    run(sim, TASK_NAME, folder_name=PROJECT_NAME, priority=priority)
+
+
+@responses.activate
+@pytest.mark.parametrize("priority", [0, -1, 11, 15])
+def test_run_with_invalid_priority(mock_webapi, priority):
+    """Test run with invalid priority values."""
+    sim = make_sim()
+    with pytest.raises(ValueError, match="Priority must be between '1' and '10' if specified."):
+        run(sim, TASK_NAME, folder_name=PROJECT_NAME, priority=priority)
+
+
 @responses.activate
 def test_get_run_info(mock_get_run_info):
     assert get_run_info(TASK_ID) == (100, 0)

tests/test_web/test_webapi_eme.py

@@ -131,6 +131,7 @@ def mock_start(monkeypatch, set_api_key, mock_get_info):
                     "protocolVersion": td.version.__version__,
                     "enableCaching": Env.current.enable_caching,
                     "payType": PayType.AUTO,
+                    "priority": None,
                 }
             )
         ],

tests/test_web/test_webapi_heat.py

@@ -128,6 +128,7 @@ def mock_start(monkeypatch, set_api_key, mock_get_info):
                     "protocolVersion": td.version.__version__,
                     "enableCaching": Env.current.enable_caching,
                     "payType": PayType.AUTO,
+                    "priority": None,
                 }
             )
         ],

tests/test_web/test_webapi_mode.py

@@ -164,6 +164,7 @@ def mock_start(monkeypatch, set_api_key, mock_get_info):
                     "protocolVersion": td.version.__version__,
                     "enableCaching": Env.current.enable_caching,
                     "payType": PayType.AUTO,
+                    "priority": None,
                 }
             )
         ],

tests/test_web/test_webapi_mode_sim.py

@@ -160,6 +160,7 @@ def mock_start(monkeypatch, set_api_key, mock_get_info):
                     "protocolVersion": td.version.__version__,
                     "enableCaching": Env.current.enable_caching,
                     "payType": PayType.AUTO,
+                    "priority": None,
                 }
             )
         ],

tidy3d/web/api/autograd/autograd.py

@@ -108,6 +108,7 @@ def run(
     max_num_adjoint_per_fwd: int = MAX_NUM_ADJOINT_PER_FWD,
     reduce_simulation: typing.Literal["auto", True, False] = "auto",
     pay_type: typing.Union[PayType, str] = PayType.AUTO,
+    priority: typing.Optional[int] = None,
 ) -> SimulationDataType:
     """
     Submits a :class:`.Simulation` to server, starts running, monitors progress, downloads,
@@ -147,6 +148,8 @@ def run(
         Whether to reduce structures in the simulation to the simulation domain only. Note: currently only implemented for the mode solver.
     pay_type: typing.Union[PayType, str] = PayType.AUTO
         Which method to pay for the simulation.
+    priority: int = None
+        Task priority for vGPU queue (1=lowest, 10=highest).
     Returns
     -------
     Union[:class:`.SimulationData`, :class:`.HeatSimulationData`, :class:`.EMESimulationData`]
@@ -191,6 +194,8 @@ def run(
     :meth:`tidy3d.web.api.container.Batch.monitor`
         Monitor progress of each of the running tasks.
     """
+    if priority is not None and (priority < 1 or priority > 10):
+        raise ValueError("Priority must be between '1' and '10' if specified.")
     if is_valid_for_autograd(simulation):
         return _run(
             simulation=simulation,
@@ -225,6 +230,7 @@ def run(
         parent_tasks=parent_tasks,
         reduce_simulation=reduce_simulation,
         pay_type=pay_type,
+        priority=priority,
     )
 
 

tidy3d/web/api/webapi.py

@@ -83,6 +83,7 @@ def run(
     parent_tasks: Optional[list[str]] = None,
     reduce_simulation: Literal["auto", True, False] = "auto",
     pay_type: Union[PayType, str] = PayType.AUTO,
+    priority: Optional[int] = None,
 ) -> SimulationDataType:
     """
     Submits a :class:`.Simulation` to server, starts running, monitors progress, downloads,
@@ -117,7 +118,8 @@ def run(
         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
        Which method to pay the simulation.
-
+    priority: int = None
+        Task priority for vGPU queue (1=lowest, 10=highest).
     Returns
     -------
     Union[:class:`.SimulationData`, :class:`.HeatSimulationData`, :class:`.EMESimulationData`]
@@ -179,6 +181,7 @@ def run(
         solver_version=solver_version,
         worker_group=worker_group,
         pay_type=pay_type,
+        priority=priority,
     )
     monitor(task_id, verbose=verbose)
     data = load(
@@ -372,6 +375,7 @@ def start(
     solver_version: Optional[str] = None,
     worker_group: Optional[str] = None,
     pay_type: Union[PayType, str] = PayType.AUTO,
+    priority: Optional[int] = None,
 ) -> None:
     """Start running the simulation associated with task.
 
@@ -386,17 +390,22 @@ def start(
         worker group
     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).
     Note
     ----
     To monitor progress, can call :meth:`monitor` after starting simulation.
     """
+    if priority is not None and (priority < 1 or priority > 10):
+        raise ValueError("Priority must be between '1' and '10' if specified.")
     task = SimulationTask.get(task_id)
     if not task:
         raise ValueError("Task not found.")
     task.submit(
         solver_version=solver_version,
         worker_group=worker_group,
         pay_type=pay_type,
+        priority=priority,
     )
 
 

tidy3d/web/core/task_core.py

@@ -429,6 +429,7 @@ def submit(
         solver_version: Optional[str] = None,
         worker_group: Optional[str] = None,
         pay_type: Union[PayType, str] = PayType.AUTO,
+        priority: Optional[int] = None,
     ):
         """Kick off this task.
 
@@ -444,6 +445,8 @@ def submit(
             worker group
         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).
         """
         pay_type = PayType(pay_type) if not isinstance(pay_type, PayType) else pay_type
 
@@ -460,6 +463,7 @@ def submit(
                 "protocolVersion": protocol_version,
                 "enableCaching": Env.current.enable_caching,
                 "payType": pay_type.value,
+                "priority": priority,
             },
         )