Skip container

Signed-off-by: Hongxin Liang <[email protected]>
flyteorg · Dec 12, 2023 · b0d9b93 · b0d9b93
1 parent 70ad715
commit b0d9b93
Showing 1 changed file with 10 additions and 12 deletions.
diff --git a/jflyte-utils/src/main/java/org/flyte/jflyte/utils/ProtoUtil.java b/jflyte-utils/src/main/java/org/flyte/jflyte/utils/ProtoUtil.java
@@ -19,7 +19,6 @@
 import static com.google.common.base.Strings.emptyToNull;
 import static com.google.common.base.Strings.nullToEmpty;
 import static java.time.format.DateTimeFormatter.ISO_DATE_TIME;
-import static java.util.Objects.requireNonNull;
 import static org.flyte.jflyte.utils.MoreCollectors.mapValues;
 import static org.flyte.jflyte.utils.MoreCollectors.toUnmodifiableList;
 import static org.flyte.jflyte.utils.MoreCollectors.toUnmodifiableMap;
@@ -312,19 +311,18 @@ static TaskOuterClass.TaskSpec serialize(TaskSpec spec) {
   }
 
   static Tasks.TaskTemplate serialize(TaskTemplate taskTemplate) {
-    Container container =
-        requireNonNull(
-            taskTemplate.container(), "Only container based task templates are supported");
-
     TaskMetadata metadata = serializeTaskMetadata(taskTemplate);
 
-    return Tasks.TaskTemplate.newBuilder()
-        .setContainer(serialize(container))
-        .setMetadata(metadata)
-        .setInterface(serialize(taskTemplate.interface_()))
-        .setType(taskTemplate.type())
-        .setCustom(serializeStruct(taskTemplate.custom()))
-        .build();
+    final Tasks.TaskTemplate.Builder builder =
+        Tasks.TaskTemplate.newBuilder()
+            .setMetadata(metadata)
+            .setInterface(serialize(taskTemplate.interface_()))
+            .setType(taskTemplate.type())
+            .setCustom(serializeStruct(taskTemplate.custom()));
+    if (taskTemplate.container() == null) {
+      return builder.setContainer(Tasks.Container.getDefaultInstance()).build();
+    }
+    return builder.setContainer(serialize(taskTemplate.container())).build();
   }
 
   private static TaskMetadata serializeTaskMetadata(TaskTemplate taskTemplate) {