feat(core): allow setting device reuqest on docker execution

kestra-io · Oct 13, 2022 · ff1e91e · ff1e91e
1 parent 744b86e
commit ff1e91e
Show file tree

Hide file tree

Showing 2 changed files with 33 additions and 0 deletions.
diff --git a/core/src/main/java/io/kestra/core/tasks/scripts/AbstractBash.java b/core/src/main/java/io/kestra/core/tasks/scripts/AbstractBash.java
@@ -425,5 +425,23 @@ public static class DockerOptions {
         @PluginProperty(dynamic = false)
         @Builder.Default
         protected Boolean pullImage = true;
+
+        @Schema(
+            title = "A list of request for devices to be sent to device drivers"
+        )
+        @PluginProperty(dynamic = false)
+        protected List<DeviceRequest> deviceRequests;
+
+        @SuperBuilder
+        @NoArgsConstructor
+        @Getter
+        @Introspected
+        public static class DeviceRequest {
+            private String driver;
+            private Integer count;
+            private List<String> deviceIds;
+            private List<List<String>> capabilities;
+            private Map<String, String> options;
+        }
     }
 }
diff --git a/core/src/main/java/io/kestra/core/tasks/scripts/runners/DockerScriptRunner.java b/core/src/main/java/io/kestra/core/tasks/scripts/runners/DockerScriptRunner.java
@@ -194,6 +194,21 @@ public RunResult run(
                 );
             }
 
+            if (abstractBash.getDockerOptions().getDeviceRequests() != null) {
+                hostConfig.withDeviceRequests(abstractBash.getDockerOptions()
+                    .getDeviceRequests()
+                    .stream()
+                    .map(throwFunction(deviceRequest -> new DeviceRequest()
+                        .withDriver(runContext.render(deviceRequest.getDriver()))
+                        .withCount(deviceRequest.getCount())
+                        .withDeviceIds(deviceRequest.getDeviceIds())
+                        .withCapabilities(deviceRequest.getCapabilities())
+                        .withOptions(deviceRequest.getOptions())
+                    ))
+                    .collect(Collectors.toList())
+                );
+            }
+
             if (abstractBash.getDockerOptions().getNetworkMode() != null) {
                 hostConfig.withNetworkMode(runContext.render(abstractBash.getDockerOptions().getNetworkMode(), additionalVars));
             }