Fetch task template in dynamic workflow task

flytekit-examples-scala/src/main/resources/META-INF/services/org.flyte.flytekit.SdkRunnableTask

@@ -2,3 +2,4 @@ org.flyte.examples.flytekitscala.HelloWorldTask
 org.flyte.examples.flytekitscala.SumTask
 org.flyte.examples.flytekitscala.GreetTask
 org.flyte.examples.flytekitscala.AddQuestionTask
+org.flyte.examples.flytekitscala.NoInputsTask

flytekit-examples/src/main/java/org/flyte/examples/DynamicFibonacciWorkflowTask.java

@@ -16,8 +16,9 @@
  */
 package org.flyte.examples;
 
-import static org.flyte.examples.FlyteEnvironment.DOMAIN;
+import static org.flyte.examples.FlyteEnvironment.DEVELOPMENT_DOMAIN;
 import static org.flyte.examples.FlyteEnvironment.PROJECT;
+import static org.flyte.examples.FlyteEnvironment.STAGING_DOMAIN;
 
 import com.google.auto.service.AutoService;
 import com.google.auto.value.AutoValue;
@@ -65,22 +66,34 @@ public Output run(SdkWorkflowBuilder builder, Input input) {
     } else if (input.n().get() == 0) {
       return Output.create(SdkBindingDataFactory.of(0));
     } else {
+      // remote task that is discoverable in current classpath
       SdkNode<Void> hello =
           builder.apply(
               "hello",
               SdkRemoteTask.create(
-                  DOMAIN,
+                  DEVELOPMENT_DOMAIN,
                   PROJECT,
                   HelloWorldTask.class.getName(),
                   SdkTypes.nulls(),
                   SdkTypes.nulls()));
+      // a fully remote task
+      SdkNode<Void> world =
+          builder.apply(
+              "world",
+              SdkRemoteTask.create(
+                      STAGING_DOMAIN,
+                      PROJECT,
+                      "org.flyte.examples.flytekitscala.NoInputsTask",
+                      SdkTypes.nulls(),
+                      SdkTypes.nulls())
+                  .withUpstreamNode(hello));
       @Var SdkBindingData<Long> prev = SdkBindingDataFactory.of(0);
       @Var SdkBindingData<Long> value = SdkBindingDataFactory.of(1);
       for (int i = 2; i <= input.n().get(); i++) {
         SdkBindingData<Long> next =
             builder
                 .apply(
-                    "fib-" + i, new SumTask().withUpstreamNode(hello), SumInput.create(value, prev))
+                    "fib-" + i, new SumTask().withUpstreamNode(world), SumInput.create(value, prev))
                 .getOutputs();
         prev = value;
         value = next;

flytekit-examples/src/main/java/org/flyte/examples/FlyteEnvironment.java

@@ -18,7 +18,8 @@
 
 public final class FlyteEnvironment {
 
-  public static final String DOMAIN = "development";
+  public static final String DEVELOPMENT_DOMAIN = "development";
+  public static final String STAGING_DOMAIN = "staging";
   public static final String PROJECT = "flytesnacks";
 
   private FlyteEnvironment() {

integration-tests/pom.xml

@@ -58,6 +58,11 @@
       <artifactId>flytekit-examples</artifactId>
       <scope>test</scope>
     </dependency>
+    <dependency>
+      <groupId>org.flyte</groupId>
+      <artifactId>flytekit-examples-scala_2.13</artifactId>
+      <scope>test</scope>
+    </dependency>
     <dependency>
       <groupId>org.flyte</groupId>
       <artifactId>jflyte</artifactId>

integration-tests/src/test/java/org/flyte/JavaExamplesIT.java

@@ -17,6 +17,7 @@
 package org.flyte;
 
 import static org.flyte.FlyteContainer.CLIENT;
+import static org.flyte.examples.FlyteEnvironment.STAGING_DOMAIN;
 import static org.flyte.utils.Literal.ofIntegerMap;
 import static org.hamcrest.MatcherAssert.assertThat;
 import static org.hamcrest.Matchers.equalTo;
@@ -29,11 +30,13 @@
 
 @TestInstance(TestInstance.Lifecycle.PER_CLASS)
 public class JavaExamplesIT {
-  private static final String CLASSPATH = "flytekit-examples/target/lib";
+  private static final String CLASSPATH_EXAMPLES = "flytekit-examples/target/lib";
+  private static final String CLASSPATH_EXAMPLES_SCALA = "flytekit-examples-scala/target/lib";
 
   @BeforeAll
   public static void beforeAll() {
-    CLIENT.registerWorkflows(CLASSPATH);
+    CLIENT.registerWorkflows(CLASSPATH_EXAMPLES);
+    CLIENT.registerWorkflows(CLASSPATH_EXAMPLES_SCALA, STAGING_DOMAIN);
   }
 
   @Test

integration-tests/src/test/java/org/flyte/utils/FlyteSandboxClient.java

@@ -16,7 +16,7 @@
  */
 package org.flyte.utils;
 
-import static org.flyte.examples.FlyteEnvironment.DOMAIN;
+import static org.flyte.examples.FlyteEnvironment.DEVELOPMENT_DOMAIN;
 import static org.flyte.examples.FlyteEnvironment.PROJECT;
 
 import flyteidl.admin.ExecutionOuterClass;
@@ -59,7 +59,7 @@ public Literals.LiteralMap createTaskExecution(String name, Literals.LiteralMap
     return createExecution(
         IdentifierOuterClass.Identifier.newBuilder()
             .setResourceType(IdentifierOuterClass.ResourceType.TASK)
-            .setDomain(DOMAIN)
+            .setDomain(DEVELOPMENT_DOMAIN)
             .setProject(PROJECT)
             .setName(name)
             .setVersion(version)
@@ -71,7 +71,7 @@ public Literals.LiteralMap createExecution(String name, Literals.LiteralMap inpu
     return createExecution(
         IdentifierOuterClass.Identifier.newBuilder()
             .setResourceType(IdentifierOuterClass.ResourceType.LAUNCH_PLAN)
-            .setDomain(DOMAIN)
+            .setDomain(DEVELOPMENT_DOMAIN)
             .setProject(PROJECT)
             .setName(name)
             .setVersion(version)
@@ -84,7 +84,7 @@ private Literals.LiteralMap createExecution(
     ExecutionOuterClass.ExecutionCreateResponse response =
         stub.createExecution(
             ExecutionOuterClass.ExecutionCreateRequest.newBuilder()
-                .setDomain(DOMAIN)
+                .setDomain(DEVELOPMENT_DOMAIN)
                 .setProject(PROJECT)
                 .setInputs(inputs)
                 .setSpec(ExecutionOuterClass.ExecutionSpec.newBuilder().setLaunchPlan(id).build())
@@ -148,21 +148,25 @@ private boolean isRunning(Execution.WorkflowExecution.Phase phase) {
     return false;
   }
 
-  public void registerWorkflows(String classpath) {
+  public void registerWorkflows(String classpath, String domain) {
     try {
       jflyte(
           "jflyte",
           "register",
           "workflows",
           "-p=" + PROJECT,
-          "-d=" + DOMAIN,
+          "-d=" + domain,
           "-v=" + version,
           "-cp=" + classpath);
     } catch (Exception e) {
       throw new RuntimeException("Could not register workflows from: " + classpath, e);
     }
   }
 
+  public void registerWorkflows(String classpath) {
+    registerWorkflows(classpath, DEVELOPMENT_DOMAIN);
+  }
+
   public void serializeWorkflows(String classpath, String folder) {
     jflyte("jflyte", "serialize", "workflows", "-cp=" + classpath, "-f=" + folder);
   }

jflyte-utils/src/main/java/org/flyte/jflyte/utils/FlyteAdminClient.java

@@ -25,7 +25,6 @@
 import flyteidl.admin.LaunchPlanOuterClass;
 import flyteidl.admin.TaskOuterClass;
 import flyteidl.admin.WorkflowOuterClass;
-import flyteidl.core.IdentifierOuterClass;
 import flyteidl.service.AdminServiceGrpc;
 import io.grpc.Channel;
 import io.grpc.ClientInterceptor;
@@ -185,34 +184,38 @@ public TaskIdentifier fetchLatestTaskId(NamedEntityIdentifier taskId) {
     return fetchLatestResource(
         taskId,
         request -> stub.listTasks(request).getTasksList(),
-        TaskOuterClass.Task::getId,
-        ProtoUtil::deserializeTaskId);
+        task -> ProtoUtil.deserializeTaskId(task.getId()));
+  }
+
+  @Nullable
+  public TaskTemplate fetchLatestTaskTemplate(NamedEntityIdentifier taskId) {
+    return fetchLatestResource(
+        taskId,
+        request -> stub.listTasks(request).getTasksList(),
+        task -> ProtoUtil.deserialize(task.getClosure().getCompiledTask().getTemplate()));
   }
 
   @Nullable
   public WorkflowIdentifier fetchLatestWorkflowId(NamedEntityIdentifier workflowId) {
     return fetchLatestResource(
         workflowId,
         request -> stub.listWorkflows(request).getWorkflowsList(),
-        WorkflowOuterClass.Workflow::getId,
-        ProtoUtil::deserializeWorkflowId);
+        workflow -> ProtoUtil.deserializeWorkflowId(workflow.getId()));
   }
 
   @Nullable
   public LaunchPlanIdentifier fetchLatestLaunchPlanId(NamedEntityIdentifier launchPlanId) {
     return fetchLatestResource(
         launchPlanId,
         request -> stub.listLaunchPlans(request).getLaunchPlansList(),
-        LaunchPlanOuterClass.LaunchPlan::getId,
-        ProtoUtil::deserializeLaunchPlanId);
+        launchPlan -> ProtoUtil.deserializeLaunchPlanId(launchPlan.getId()));
   }
 
   @Nullable
   private <T, RespT> T fetchLatestResource(
       NamedEntityIdentifier nameId,
       Function<ResourceListRequest, List<RespT>> performRequestFn,
-      Function<RespT, IdentifierOuterClass.Identifier> extractIdFn,
-      Function<IdentifierOuterClass.Identifier, T> deserializeFn) {
+      Function<RespT, T> deserializeFn) {
     ResourceListRequest request =
         ResourceListRequest.newBuilder()
             .setLimit(1)
@@ -230,8 +233,7 @@ private <T, RespT> T fetchLatestResource(
       return null;
     }
 
-    IdentifierOuterClass.Identifier id = extractIdFn.apply(list.get(0));
-    return deserializeFn.apply(id);
+    return deserializeFn.apply(list.get(0));
   }
 
   private <T> void idempotentCreate(String label, Object id, GrpcRetries.Retryable<T> retryable) {

jflyte-utils/src/main/java/org/flyte/jflyte/utils/ProjectClosure.java

@@ -47,6 +47,7 @@
 import java.util.Set;
 import java.util.concurrent.atomic.AtomicInteger;
 import java.util.function.BiConsumer;
+import java.util.function.Function;
 import java.util.function.Supplier;
 import java.util.stream.Stream;
 import org.flyte.api.v1.Container;
@@ -295,8 +296,8 @@ static void checkCycles(Map<WorkflowIdentifier, WorkflowTemplate> allWorkflows)
                     checkCycles(
                         workflowId,
                         allWorkflows,
-                        /*beingVisited=*/ new HashSet<>(),
-                        /*visited=*/ new HashSet<>()))
+                        /* beingVisited= */ new HashSet<>(),
+                        /* visited= */ new HashSet<>()))
             .findFirst();
     if (cycle.isPresent()) {
       throw new IllegalArgumentException(
@@ -374,8 +375,10 @@ public static Map<WorkflowIdentifier, WorkflowTemplate> collectSubWorkflows(
         .collect(toUnmodifiableMap());
   }
 
-  public static Map<TaskIdentifier, TaskTemplate> collectTasks(
-      List<Node> rewrittenNodes, Map<TaskIdentifier, TaskTemplate> allTasks) {
+  public static Map<TaskIdentifier, TaskTemplate> collectDynamicWorkflowTasks(
+      List<Node> rewrittenNodes,
+      Map<TaskIdentifier, TaskTemplate> allTasks,
+      Function<TaskIdentifier, TaskTemplate> remoteTaskTemplateFetcher) {
     return collectTaskIds(rewrittenNodes).stream()
         // all identifiers should be rewritten at this point
         .map(
@@ -389,7 +392,9 @@ public static Map<TaskIdentifier, TaskTemplate> collectTasks(
         .distinct()
         .map(
             taskId -> {
-              TaskTemplate taskTemplate = allTasks.get(taskId);
+              TaskTemplate taskTemplate =
+                  Optional.ofNullable(allTasks.get(taskId))
+                      .orElseGet(() -> remoteTaskTemplateFetcher.apply(taskId));
 
               if (taskTemplate == null) {
                 throw new NoSuchElementException("Can't find referenced task " + taskId);

jflyte-utils/src/test/java/org/flyte/jflyte/utils/Fixtures.java

@@ -0,0 +1,62 @@
+/*
+ * Copyright 2023 Flyte Authors.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied.  See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+package org.flyte.jflyte.utils;
+
+import static java.util.Collections.emptyMap;
+
+import com.google.common.collect.ImmutableList;
+import com.google.common.collect.ImmutableMap;
+import org.flyte.api.v1.Container;
+import org.flyte.api.v1.KeyValuePair;
+import org.flyte.api.v1.RetryStrategy;
+import org.flyte.api.v1.SimpleType;
+import org.flyte.api.v1.Struct;
+import org.flyte.api.v1.TaskTemplate;
+import org.flyte.api.v1.TypedInterface;
+
+final class Fixtures {
+  static final String IMAGE_NAME = "alpine:latest";
+  static final String COMMAND = "date";
+
+  static final Container CONTAINER =
+      Container.builder()
+          .command(ImmutableList.of(COMMAND))
+          .args(ImmutableList.of())
+          .image(IMAGE_NAME)
+          .env(ImmutableList.of(KeyValuePair.of("key", "value")))
+          .build();
+  static final TypedInterface INTERFACE_ =
+      TypedInterface.builder()
+          .inputs(ImmutableMap.of("x", ApiUtils.createVar(SimpleType.STRING)))
+          .outputs(ImmutableMap.of("y", ApiUtils.createVar(SimpleType.INTEGER)))
+          .build();
+  static final RetryStrategy RETRIES = RetryStrategy.builder().retries(4).build();
+  static final TaskTemplate TASK_TEMPLATE =
+      TaskTemplate.builder()
+          .container(CONTAINER)
+          .type("custom-task")
+          .interface_(INTERFACE_)
+          .custom(Struct.of(emptyMap()))
+          .retries(RETRIES)
+          .discoverable(false)
+          .cacheSerializable(false)
+          .build();
+
+  private Fixtures() {
+    throw new UnsupportedOperationException();
+  }
+}