Add default resources support to DynamicWorkflowTask and related classes

Signed-off-by: Rafael Ribeiro Raposo <rafaelraposo@spotify.com>
flyteorg · Feb 10, 2025 · 156b00f · 156b00f
1 parent f034f75
commit 156b00f
Show file tree

Hide file tree

Showing 5 changed files with 47 additions and 0 deletions.
diff --git a/flytekit-api/src/main/java/org/flyte/api/v1/DynamicWorkflowTask.java b/flytekit-api/src/main/java/org/flyte/api/v1/DynamicWorkflowTask.java
@@ -26,4 +26,8 @@ public interface DynamicWorkflowTask {
   DynamicJobSpec run(Map<String, Literal> inputs);
 
   RetryStrategy getRetries();
+
+  default Resources getResources() {
+    return Resources.builder().build();
+  }
 }
diff --git a/flytekit-java/src/main/java/org/flyte/flytekit/SdkDynamicWorkflowTask.java b/flytekit-java/src/main/java/org/flyte/flytekit/SdkDynamicWorkflowTask.java
@@ -93,4 +93,8 @@ public SdkNode<OutputT> apply(
   public int getRetries() {
     return 0;
   }
+
+  public SdkResources getResources() {
+    return SdkResources.empty();
+  }
 }
diff --git a/flytekit-java/src/main/java/org/flyte/flytekit/SdkDynamicWorkflowTaskRegistrar.java b/flytekit-java/src/main/java/org/flyte/flytekit/SdkDynamicWorkflowTaskRegistrar.java
@@ -33,6 +33,7 @@
 import org.flyte.api.v1.DynamicWorkflowTaskRegistrar;
 import org.flyte.api.v1.Literal;
 import org.flyte.api.v1.Node;
+import org.flyte.api.v1.Resources;
 import org.flyte.api.v1.RetryStrategy;
 import org.flyte.api.v1.TaskIdentifier;
 import org.flyte.api.v1.TypedInterface;
@@ -112,6 +113,11 @@ public DynamicJobSpec run(Map<String, Literal> inputs) {
     public RetryStrategy getRetries() {
       return RetryStrategy.builder().retries(sdkDynamicWorkflow.getRetries()).build();
     }
+
+    @Override
+    public Resources getResources() {
+      return sdkDynamicWorkflow.getResources().toIdl();
+    }
   }
 
   /**

diff --git a/flytekit-java/src/test/java/org/flyte/flytekit/SdkDynamicWorkflowTaskRegistrarTest.java b/flytekit-java/src/test/java/org/flyte/flytekit/SdkDynamicWorkflowTaskRegistrarTest.java
@@ -22,6 +22,7 @@
 import static org.hamcrest.Matchers.hasSize;
 
 import com.google.errorprone.annotations.Var;
+import java.util.HashMap;
 import java.util.List;
 import java.util.Map;
 import org.flyte.api.v1.Binding;
@@ -30,6 +31,7 @@
 import org.flyte.api.v1.Literal;
 import org.flyte.api.v1.OutputReference;
 import org.flyte.api.v1.Primitive;
+import org.flyte.api.v1.Resources;
 import org.flyte.api.v1.Scalar;
 import org.flyte.api.v1.TaskIdentifier;
 import org.flyte.api.v1.TypedInterface;
@@ -70,6 +72,14 @@ void shouldLoad() {
                 .inputs(SdkLiteralTypes.integers().asSdkType("n").getVariableMap())
                 .outputs(SdkLiteralTypes.integers().asSdkType("2n").getVariableMap())
                 .build()));
+    assertThat(
+        dynWf.getResources(),
+        equalTo(
+            Resources.builder()
+                .requests(resources("0.5", "2Gi"))
+                .limits(resources("2", "5Gi"))
+                .build()));
+    ;
     var spec =
         dynWf.run(Map.of("n", Literal.ofScalar(Scalar.ofPrimitive(Primitive.ofIntegerValue(3)))));
     assertThat(spec.nodes(), hasSize(3));
@@ -103,6 +113,14 @@ public SdkBindingData<Long> run(SdkWorkflowBuilder builder, SdkBindingData<Long>
       }
       return x;
     }
+
+    @Override
+    public SdkResources getResources() {
+      return SdkResources.builder()
+          .requests(sdkResources("0.5", "2Gi"))
+          .limits(sdkResources("2", "5Gi"))
+          .build();
+    }
   }
 
   static class Mult2 extends SdkRunnableTask<SdkBindingData<Long>, SdkBindingData<Long>> {
@@ -117,4 +135,18 @@ public SdkBindingData<Long> run(SdkBindingData<Long> input) {
       return SdkBindingDataFactory.of(input.get() * 2);
     }
   }
+
+  private static Map<Resources.ResourceName, String> resources(String cpu, String memory) {
+    Map<Resources.ResourceName, String> limits = new HashMap<>();
+    limits.put(Resources.ResourceName.CPU, cpu);
+    limits.put(Resources.ResourceName.MEMORY, memory);
+    return limits;
+  }
+
+  private static Map<SdkResources.ResourceName, String> sdkResources(String cpu, String memory) {
+    Map<SdkResources.ResourceName, String> limits = new HashMap<>();
+    limits.put(SdkResources.ResourceName.CPU, cpu);
+    limits.put(SdkResources.ResourceName.MEMORY, memory);
+    return limits;
+  }
 }
diff --git a/jflyte-utils/src/main/java/org/flyte/jflyte/utils/ProjectClosure.java b/jflyte-utils/src/main/java/org/flyte/jflyte/utils/ProjectClosure.java
@@ -577,6 +577,7 @@ private static TaskTemplate createTaskTemplateForDynamicWorkflow(
                     "{{.taskTemplatePath}}"))
             .image(image)
             .env(emptyList())
+            .resources(task.getResources())
             .build();
 
     return TaskTemplate.builder()