fix: add utility capability for local_docker (#906)

Differential Revision: D57034038 Pull Request resolved: #907
pytorch · May 7, 2024 · ec8d0e8 · ec8d0e8
1 parent c13595b
commit ec8d0e8
Show file tree

Hide file tree

Showing 2 changed files with 2 additions and 2 deletions.
diff --git a/torchx/schedulers/docker_scheduler.py b/torchx/schedulers/docker_scheduler.py
@@ -314,7 +314,7 @@ def _submit_dryrun(self, app: AppDef, cfg: DockerOpts) -> AppDryRunInfo[DockerJo
                     c.kwargs["device_requests"] = [
                         DeviceRequest(
                             count=resource.gpu,
-                            capabilities=[["compute"]],
+                            capabilities=[["compute", "utility"]],
                         )
                     ]
                 req.containers.append(c)

diff --git a/torchx/schedulers/test/docker_scheduler_test.py b/torchx/schedulers/test/docker_scheduler_test.py
@@ -91,7 +91,7 @@ def test_submit_dryrun(self) -> None:
                         "device_requests": [
                             DeviceRequest(
                                 count=4,
-                                capabilities=[["compute"]],
+                                capabilities=[["compute", "utility"]],
                             )
                         ],
                         "devices": [