Enable GPU for local runner

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",

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:

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

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"),