Enable GPU for local runner (#377)

PR that adds GPU resource for the local runner when specified in the componentOP based on: https://docs.docker.com/compose/gpu-support/
ml6team · Aug 23, 2023 · 2cbecac · 2cbecac
1 parent 50f3a97
commit 2cbecac
Show file tree

Hide file tree

Showing 4 changed files with 29 additions and 1 deletion.
diff --git a/src/fondant/compiler.py b/src/fondant/compiler.py
@@ -73,7 +73,7 @@ def ignore_aliases(self, data):
                 return True
 
         with open(output_path, "w") as outfile:
-            yaml.dump(spec, outfile, Dumper=NoAliasDumper)
+            yaml.dump(spec, outfile, Dumper=NoAliasDumper, default_flow_style=False)
 
         logger.info(f"Successfully compiled to {output_path}")
 
@@ -176,6 +176,21 @@ def _generate_spec(
                 "volumes": volumes,
             }
 
+            if component_op.number_of_gpus is not None:
+                services[component_name]["deploy"] = {
+                    "resources": {
+                        "reservations": {
+                            "devices": [
+                                {
+                                    "driver": "nvidia",
+                                    "count": component_op.number_of_gpus,
+                                    "capabilities": ["gpu"],
+                                },
+                            ],
+                        },
+                    },
+                }
+
             if component_op.dockerfile_path is not None:
                 logger.info(
                     f"Found Dockerfile for {component_name}, adding build step.",
@@ -186,6 +201,7 @@ def _generate_spec(
                 }
             else:
                 services[component_name]["image"] = component_op.component_spec.image
+
         return {
             "name": pipeline.name,
             "version": "3.8",

diff --git a/tests/example_pipelines/compiled_pipeline/example_1/docker-compose.yml b/tests/example_pipelines/compiled_pipeline/example_1/docker-compose.yml
@@ -20,6 +20,14 @@ services:
       {"type": "binary"}}}, "captions": {"fields": {"data": {"type": "string"}}}},
       "args": {"storage_args": {"description": "Storage arguments", "type": "str"}}}'
     depends_on: {}
+    deploy:
+      resources:
+        reservations:
+          devices:
+          - capabilities:
+            - gpu
+            count: 1
+            driver: nvidia
     volumes: []
   second_component:
     build:

diff --git a/tests/example_pipelines/compiled_pipeline/example_1/kubeflow_pipeline.yml b/tests/example_pipelines/compiled_pipeline/example_1/kubeflow_pipeline.yml
@@ -38,6 +38,9 @@ spec:
       - --output_manifest_path
       - /tmp/outputs/output_manifest_path/data
       image: example_component:latest
+      resources:
+        limits:
+          nvidia.com/gpu: 1
     inputs:
       artifacts:
       - name: input_manifest_path

diff --git a/tests/test_compiler.py b/tests/test_compiler.py
@@ -20,6 +20,7 @@
                 Path(COMPONENTS_PATH / "example_1" / "first_component"),
                 arguments={"storage_args": "a dummy string arg"},
                 input_partition_rows="disable",
+                number_of_gpus=1,
             ),
             ComponentOp(
                 Path(COMPONENTS_PATH / "example_1" / "second_component"),