gpu support for docker operator

Tauffer-Consulting · Nov 21, 2023 · ab2d1aa · ab2d1aa
1 parent 7d99b2f
commit ab2d1aa
Show file tree

Hide file tree

Showing 3 changed files with 15 additions and 6 deletions.
diff --git a/frontend/src/features/workflowEditor/components/SidebarForm/ContainerResourceForm.tsx b/frontend/src/features/workflowEditor/components/SidebarForm/ContainerResourceForm.tsx
@@ -73,7 +73,7 @@ const ContainerResourceForm: React.FC = () => {
       </Grid>
       <Grid item xs={6}>
         <NumberInput
-          label="CPU Min"
+          label="CPU Min (m)"
           name="containerResources.cpu.min"
           required
           type="int"
@@ -85,7 +85,7 @@ const ContainerResourceForm: React.FC = () => {
       </Grid>
       <Grid item xs={6}>
         <NumberInput
-          label="CPU Max"
+          label="CPU Max (m)"
           name="containerResources.cpu.max"
           type="int"
           required

diff --git a/src/domino/custom_operators/docker_operator.py b/src/domino/custom_operators/docker_operator.py
@@ -5,8 +5,8 @@
 
 from domino.client.domino_backend_client import DominoBackendRestClient
 from domino.schemas import WorkflowSharedStorage, StorageSource
-
-
+from docker.types import Mount
+import docker
 class DominoDockerOperator(DockerOperator):
 
     def __init__(
@@ -20,6 +20,7 @@ def __init__(
         workspace_id: int,
         piece_input_kwargs: Optional[Dict] = None,
         workflow_shared_storage: WorkflowSharedStorage = None,
+        container_resources: Optional[Dict] = None,
         **docker_operator_kwargs
     ) -> None:
         self.task_id = task_id
@@ -30,6 +31,7 @@ def __init__(
         self.workspace_id = workspace_id
         self.piece_input_kwargs = piece_input_kwargs
         self.workflow_shared_storage = workflow_shared_storage
+        self.container_resources = container_resources or {}
 
         # Environment variables
         self.environment = {
@@ -76,11 +78,17 @@ def __init__(
                 ),
             )
 
+        self.device_requests = []
+        if self.container_resources.get('use_gpu', False):
+            self.device_requests=[
+                docker.types.DeviceRequest(count=-1, capabilities=[['gpu']])
+            ]
         super().__init__(
             **docker_operator_kwargs,
             task_id=task_id,
             docker_url='tcp://docker-proxy:2375',
             mounts=mounts,
+            device_requests=self.device_requests,
             environment=self.environment,
         )
 

diff --git a/src/domino/task.py b/src/domino/task.py
@@ -8,7 +8,6 @@
 from domino.custom_operators.docker_operator import DominoDockerOperator
 from domino.custom_operators.python_operator import PythonOperator
 from domino.custom_operators.worker_operator import DominoWorkerOperator
-from domino.utils import dict_deep_update
 from domino.logger import get_configured_logger
 from domino.schemas import shared_storage_map, StorageSource
 
@@ -140,15 +139,17 @@ def _set_operator(self) -> BaseOperator:
                 workspace_id=self.workspace_id,
                 piece_input_kwargs=self.piece_input_kwargs,
                 workflow_shared_storage=self.workflow_shared_storage,
+                container_resources=self.container_resources,
                 # ----------------- Docker -----------------
+                # TODO uncoment
                 image=self.piece["source_image"],
+                entrypoint=["domino", "run-piece-docker"],
                 do_xcom_push=True,
                 mount_tmp_dir=False,
                 tty=True,
                 xcom_all=False,
                 retrieve_output=True,
                 retrieve_output_path='/airflow/xcom/return.out',
-                entrypoint=["domino", "run-piece-docker"],
             )
 
     def __call__(self) -> Callable: