Allow users to pass task payload via deep storage instead of environment variable

distribution/docker/peon.sh

@@ -150,6 +150,6 @@ then
 fi
 
 # take the ${TASK_JSON} environment variable and base64 decode, unzip and throw it in ${TASK_DIR}/task.json
-mkdir -p ${TASK_DIR}; echo ${TASK_JSON} | base64 -d | gzip -d > ${TASK_DIR}/task.json;
+mkdir -p ${TASK_DIR}; [ -n "$TASK_JSON" ] && echo ${TASK_JSON} | base64 -d | gzip -d > ${TASK_DIR}/task.json;
 
 exec bin/run-java ${JAVA_OPTS} -cp $COMMON_CONF_DIR:$SERVICE_CONF_DIR:lib/*: org.apache.druid.cli.Main internal peon $@

extensions-contrib/kubernetes-overlord-extensions/src/main/java/org/apache/druid/k8s/overlord/KubernetesPeonLifecycle.java

@@ -30,6 +30,7 @@
 import org.apache.commons.io.IOUtils;
 import org.apache.druid.indexer.TaskLocation;
 import org.apache.druid.indexer.TaskStatus;
+import org.apache.druid.indexing.common.config.TaskConfig;
 import org.apache.druid.indexing.common.task.Task;
 import org.apache.druid.java.util.common.StringUtils;
 import org.apache.druid.java.util.emitter.EmittingLogger;
@@ -42,6 +43,7 @@
 
 import java.io.IOException;
 import java.io.InputStream;
+import java.nio.charset.Charset;
 import java.nio.charset.StandardCharsets;
 import java.nio.file.Files;
 import java.nio.file.Path;
@@ -88,7 +90,7 @@ protected enum State
   private final KubernetesPeonClient kubernetesClient;
   private final ObjectMapper mapper;
   private final TaskStateListener stateListener;
-
+  private final TaskConfig taskConfig;
   @MonotonicNonNull
   private LogWatch logWatch;
 
@@ -99,7 +101,8 @@ protected KubernetesPeonLifecycle(
       KubernetesPeonClient kubernetesClient,
       TaskLogs taskLogs,
       ObjectMapper mapper,
-      TaskStateListener stateListener
+      TaskStateListener stateListener,
+      TaskConfig taskConfig
   )
   {
     this.taskId = new K8sTaskId(task);
@@ -108,6 +111,7 @@ protected KubernetesPeonLifecycle(
     this.taskLogs = taskLogs;
     this.mapper = mapper;
     this.stateListener = stateListener;
+    this.taskConfig = taskConfig;
   }
 
   /**
@@ -124,6 +128,10 @@ protected synchronized TaskStatus run(Job job, long launchTimeout, long timeout)
     try {
       updateState(new State[]{State.NOT_STARTED}, State.PENDING);
 
+      if (taskConfig.isEnableTaskPayloadManagerPerTask()) {
+        writeTaskPayload(task);
+      }
+
       // In case something bad happens and run is called twice on this KubernetesPeonLifecycle, reset taskLocation.
       taskLocation = null;
       kubernetesClient.launchPeonJobAndWaitForStart(
@@ -145,6 +153,28 @@ protected synchronized TaskStatus run(Job job, long launchTimeout, long timeout)
     }
   }
 
+  private void writeTaskPayload(Task task)
+  {
+    try {
+      Path file = Files.createTempFile(taskId.getOriginalTaskId(), "task.json");
+      try {
+        FileUtils.writeStringToFile(file.toFile(), mapper.writeValueAsString(task), Charset.defaultCharset());
+        taskLogs.pushTaskPayload(task.getId(), file.toFile());
+      }
+      catch (Exception e) {
+        log.error("Failed to write task payload for task: %s", taskId.getOriginalTaskId());
+        throw new RuntimeException(e);
+      }
+      finally {
+        Files.deleteIfExists(file);
+      }
+    }
+    catch (IOException e) {
+      log.error("Failed to write task payload for task: %s", taskId.getOriginalTaskId());
+      throw new RuntimeException(e);
+    }
+  }
+
   /**
    * Join existing Kubernetes Job
    *

extensions-contrib/kubernetes-overlord-extensions/src/main/java/org/apache/druid/k8s/overlord/KubernetesPeonLifecycleFactory.java

@@ -20,6 +20,7 @@
 package org.apache.druid.k8s.overlord;
 
 import com.fasterxml.jackson.databind.ObjectMapper;
+import org.apache.druid.indexing.common.config.TaskConfig;
 import org.apache.druid.indexing.common.task.Task;
 import org.apache.druid.k8s.overlord.common.KubernetesPeonClient;
 import org.apache.druid.tasklogs.TaskLogs;
@@ -29,16 +30,19 @@ public class KubernetesPeonLifecycleFactory implements PeonLifecycleFactory
   private final KubernetesPeonClient client;
   private final TaskLogs taskLogs;
   private final ObjectMapper mapper;
+  private final TaskConfig taskConfig;
 
   public KubernetesPeonLifecycleFactory(
       KubernetesPeonClient client,
       TaskLogs taskLogs,
-      ObjectMapper mapper
+      ObjectMapper mapper,
+      TaskConfig taskConfig
   )
   {
     this.client = client;
     this.taskLogs = taskLogs;
     this.mapper = mapper;
+    this.taskConfig = taskConfig;
   }
 
   @Override
@@ -49,7 +53,8 @@ public KubernetesPeonLifecycle build(Task task, KubernetesPeonLifecycle.TaskStat
         client,
         taskLogs,
         mapper,
-        stateListener
+        stateListener,
+        taskConfig
     );
   }
 }

extensions-contrib/kubernetes-overlord-extensions/src/main/java/org/apache/druid/k8s/overlord/KubernetesTaskRunnerConfig.java

@@ -90,6 +90,11 @@ public class KubernetesTaskRunnerConfig
   // how long to wait for the peon k8s job to launch
   private Period k8sjobLaunchTimeout = new Period("PT1H");
 
+  // Whether to pass the task.json payload to the peon K8s Jobs as a environment variable.
+  @JsonProperty
+  @NotNull
+  private Boolean taskPayloadAsEnvVariable = Boolean.TRUE;
+
   @JsonProperty
   // ForkingTaskRunner inherits the monitors from the MM, in k8s mode
   // the peon inherits the monitors from the overlord, so if someone specifies
@@ -135,7 +140,8 @@ private KubernetesTaskRunnerConfig(
       List<String> javaOptsArray,
       Map<String, String> labels,
       Map<String, String> annotations,
-      Integer capacity
+      Integer capacity,
+      Boolean taskPayloadAsEnvVariable
   )
   {
     this.namespace = namespace;
@@ -196,6 +202,10 @@ private KubernetesTaskRunnerConfig(
         capacity,
         this.capacity
     );
+    this.taskPayloadAsEnvVariable = ObjectUtils.defaultIfNull(
+        taskPayloadAsEnvVariable,
+        Boolean.TRUE
+    );
   }
 
   public String getNamespace()
@@ -279,6 +289,11 @@ public Integer getCapacity()
     return capacity;
   }
 
+  public Boolean isTaskPayloadAsEnvVariable()
+  {
+    return taskPayloadAsEnvVariable;
+  }
+
   public static Builder builder()
   {
     return new Builder();
@@ -302,6 +317,7 @@ public static class Builder
     private Map<String, String> labels;
     private Map<String, String> annotations;
     private Integer capacity;
+    private Boolean taskPayloadAsEnvVariable;
 
     public Builder()
     {
@@ -403,6 +419,12 @@ public Builder withCapacity(@Min(0) @Max(Integer.MAX_VALUE) Integer capacity)
       return this;
     }
 
+    public Builder withTaskPayloadAsEnvVariable(Boolean taskPayloadAsEnvVariable)
+    {
+      this.taskPayloadAsEnvVariable = taskPayloadAsEnvVariable;
+      return this;
+    }
+
     public KubernetesTaskRunnerConfig build()
     {
       return new KubernetesTaskRunnerConfig(
@@ -421,7 +443,8 @@ public KubernetesTaskRunnerConfig build()
           this.javaOptsArray,
           this.labels,
           this.annotations,
-          this.capacity
+          this.capacity,
+          this.taskPayloadAsEnvVariable
       );
     }
   }

extensions-contrib/kubernetes-overlord-extensions/src/main/java/org/apache/druid/k8s/overlord/KubernetesTaskRunnerFactory.java

@@ -58,7 +58,6 @@ public class KubernetesTaskRunnerFactory implements TaskRunnerFactory<Kubernetes
   private final ServiceEmitter emitter;
   private KubernetesTaskRunner runner;
 
-
   @Inject
   public KubernetesTaskRunnerFactory(
       @Smile ObjectMapper smileMapper,
@@ -101,7 +100,7 @@ public KubernetesTaskRunner build()
         kubernetesTaskRunnerConfig,
         peonClient,
         httpClient,
-        new KubernetesPeonLifecycleFactory(peonClient, taskLogs, smileMapper),
+        new KubernetesPeonLifecycleFactory(peonClient, taskLogs, smileMapper, taskConfig),
         emitter
     );
     return runner;
@@ -137,15 +136,17 @@ private TaskAdapter buildTaskAdapter(DruidKubernetesClient client)
           taskConfig,
           startupLoggingConfig,
           druidNode,
-          smileMapper
+          smileMapper,
+          taskLogs
       );
     } else if (PodTemplateTaskAdapter.TYPE.equals(adapter)) {
       return new PodTemplateTaskAdapter(
           kubernetesTaskRunnerConfig,
           taskConfig,
           druidNode,
           smileMapper,
-          properties
+          properties,
+          taskLogs
       );
     } else {
       return new SingleContainerTaskAdapter(
@@ -154,7 +155,8 @@ private TaskAdapter buildTaskAdapter(DruidKubernetesClient client)
           taskConfig,
           startupLoggingConfig,
           druidNode,
-          smileMapper
+          smileMapper,
+          taskLogs
       );
     }
   }

extensions-contrib/kubernetes-overlord-extensions/src/main/java/org/apache/druid/k8s/overlord/taskadapter/K8sTaskAdapter.java

@@ -41,13 +41,15 @@
 import io.fabric8.kubernetes.api.model.ResourceRequirementsBuilder;
 import io.fabric8.kubernetes.api.model.batch.v1.Job;
 import io.fabric8.kubernetes.api.model.batch.v1.JobBuilder;
+import org.apache.commons.io.IOUtils;
 import org.apache.commons.lang3.StringUtils;
 import org.apache.druid.indexing.common.config.TaskConfig;
 import org.apache.druid.indexing.common.task.Task;
 import org.apache.druid.indexing.overlord.ForkingTaskRunner;
 import org.apache.druid.indexing.overlord.QuotableWhiteSpaceSplitter;
 import org.apache.druid.indexing.overlord.config.ForkingTaskRunnerConfig;
 import org.apache.druid.java.util.common.HumanReadableBytes;
+import org.apache.druid.java.util.common.IOE;
 import org.apache.druid.java.util.emitter.EmittingLogger;
 import org.apache.druid.k8s.overlord.KubernetesTaskRunnerConfig;
 import org.apache.druid.k8s.overlord.common.Base64Compression;
@@ -57,8 +59,11 @@
 import org.apache.druid.k8s.overlord.common.PeonCommandContext;
 import org.apache.druid.server.DruidNode;
 import org.apache.druid.server.log.StartupLoggingConfig;
+import org.apache.druid.tasklogs.TaskLogs;
 
 import java.io.IOException;
+import java.io.InputStream;
+import java.nio.charset.Charset;
 import java.util.ArrayList;
 import java.util.Collections;
 import java.util.HashMap;
@@ -89,14 +94,16 @@ public abstract class K8sTaskAdapter implements TaskAdapter
   protected final StartupLoggingConfig startupLoggingConfig;
   protected final DruidNode node;
   protected final ObjectMapper mapper;
+  protected final TaskLogs taskLogs;
 
   public K8sTaskAdapter(
       KubernetesClientApi client,
       KubernetesTaskRunnerConfig taskRunnerConfig,
       TaskConfig taskConfig,
       StartupLoggingConfig startupLoggingConfig,
       DruidNode node,
-      ObjectMapper mapper
+      ObjectMapper mapper,
+      TaskLogs taskLogs
   )
   {
     this.client = client;
@@ -105,6 +112,7 @@ public K8sTaskAdapter(
     this.startupLoggingConfig = startupLoggingConfig;
     this.node = node;
     this.mapper = mapper;
+    this.taskLogs = taskLogs;
   }
 
   @Override
@@ -126,6 +134,14 @@ public Job fromTask(Task task) throws IOException
   @Override
   public Task toTask(Job from) throws IOException
   {
+    if (taskConfig.isEnableTaskPayloadManagerPerTask()) {
+      com.google.common.base.Optional<InputStream> taskBody = taskLogs.streamTaskPayload(getTaskId(from).getOriginalTaskId());
+      if (!taskBody.isPresent()) {
+        throw new IOE("Could not find task payload in task logs for job [%s]", from.getMetadata().getName());
+      }
+      String task = IOUtils.toString(taskBody.get(), Charset.defaultCharset());
+      return mapper.readValue(task, Task.class);
+    }
     PodSpec podSpec = from.getSpec().getTemplate().getSpec();
     massageSpec(podSpec, "main");
     List<EnvVar> envVars = podSpec.getContainers().get(0).getEnv();
@@ -137,6 +153,20 @@ public Task toTask(Job from) throws IOException
     return mapper.readValue(Base64Compression.decompressBase64(contents), Task.class);
   }
 
+  @Override
+  public K8sTaskId getTaskId(Job from) throws IOException
+  {
+    Map<String, String> annotations = from.getSpec().getTemplate().getMetadata().getAnnotations();
+    if (annotations == null) {
+      throw new IOE("No annotations found on pod spec for job [%s]", from.getMetadata().getName());
+    }
+    String taskId = annotations.get(DruidK8sConstants.TASK_ID);
+    if (taskId == null) {
+      throw new IOE("No task_id annotation found on pod spec for job [%s]", from.getMetadata().getName());
+    }
+    return new K8sTaskId(taskId);
+  }
+
   @VisibleForTesting
   abstract Job createJobFromPodSpec(PodSpec podSpec, Task task, PeonCommandContext context) throws IOException;
 
@@ -219,15 +249,11 @@ void addEnvironmentVariables(Container mainContainer, PeonCommandContext context
                                      .build());
     }
 
-    mainContainer.getEnv().addAll(Lists.newArrayList(
+    List<EnvVar> envVars = Lists.newArrayList(
         new EnvVarBuilder()
             .withName(DruidK8sConstants.TASK_DIR_ENV)
             .withValue(context.getTaskDir().getAbsolutePath())
             .build(),
-        new EnvVarBuilder()
-            .withName(DruidK8sConstants.TASK_JSON_ENV)
-            .withValue(taskContents)
-            .build(),
         new EnvVarBuilder()
             .withName(DruidK8sConstants.JAVA_OPTS)
             .withValue(Joiner.on(" ").join(context.getJavaOpts()))
@@ -244,7 +270,16 @@ void addEnvironmentVariables(Container mainContainer, PeonCommandContext context
                 null,
                 "metadata.name"
             )).build()).build()
-    ));
+    );
+    if (taskRunnerConfig.isTaskPayloadAsEnvVariable()) {
+      envVars.add(
+          new EnvVarBuilder()
+              .withName(DruidK8sConstants.TASK_JSON_ENV)
+              .withValue(taskRunnerConfig.isTaskPayloadAsEnvVariable() ? taskContents : "")
+              .build()
+      );
+    }
+    mainContainer.getEnv().addAll(envVars);
   }
 
   protected Container setupMainContainer(
@@ -403,6 +438,9 @@ private List<String> generateCommand(Task task)
       command.add("--loadBroadcastSegments");
       command.add("true");
     }
+
+    command.add("--taskId");
+    command.add(task.getId());
     log.info(
         "Peon Command for K8s job: %s",
         ForkingTaskRunner.getMaskedCommand(startupLoggingConfig.getMaskProperties(), command)