Merge pull request #43 from koxu1996/feature/job-agglomeration

[k8s] Agglomerating jobs

Author: balis

functions/kubernetes/buffer.js

@@ -0,0 +1,64 @@
+/**
+ * Class used for buffering jobs - kind of agglomeration.
+ */
+class BufferCountWithTimeout {
+  constructor(count, idleTimeoutMs, cb) {
+    this.elements = []
+    this.triggerCount = count;
+    this.idleTimeoutMs = idleTimeoutMs;
+    this.cb = cb;
+    this._rescheduleTimeout();
+  }
+
+  addItem(item) {
+    this.elements.push(item);
+    if (this.elements.length >= this.triggerCount) {
+      console.log("Running callback [reached count]");
+      let elementsCopy = this.elements;
+      this.elements = [];
+      this.cb(elementsCopy);
+    }
+    this._rescheduleTimeout();
+  }
+
+  _rescheduleTimeout() {
+    if (this.timeoutId) {
+      clearInterval(this.timeoutId);
+    }
+    this.timeoutId = setTimeout(() => {
+      console.log("Running callback [reached timeout]");
+      let elementsCopy = this.elements;
+      this.timeoutId = null;
+      this.elements = [];
+      this.cb(elementsCopy);
+    }, this.idleTimeoutMs);
+  }
+}
+
+/**
+ * Testing function.
+ * TODO: REMOVE
+ */
+async function testBuffer() {
+  let fn = (items) => {
+    console.log("Got from buffer:", items);
+  }
+  let test = new BufferCountWithTimeout(3, 1000, fn);
+  test.addItem(1);
+  test.addItem(2);
+  test.addItem(3);
+  test.addItem(4);
+  test.addItem(5);
+
+  await new Promise(resolve => setTimeout(resolve, 2000));
+
+  test.addItem(6);
+  test.addItem(7);
+  test.addItem(8);
+  test.addItem(9);
+  test.addItem(10);
+  test.addItem(12);
+  test.addItem(13);
+}
+
+exports.BufferCountWithTimeout = BufferCountWithTimeout;

functions/kubernetes/buffer_manager.js

@@ -0,0 +1,108 @@
+var Buffer = require('./buffer.js').BufferCountWithTimeout;
+
+/**
+ * Class used to store multiple buffers and add items to them
+ * according to specified configuration.
+ */
+class BufferManager {
+  constructor(cb) {
+    this.buffers = []
+    this.taskBufferMap = {}
+    this.configured = false;
+    this.cb = cb;
+  }
+
+  setCallback(cb) {
+    if (this.cb != undefined) {
+      throw Error("Callback is already set");
+    }
+    this.cb = cb;
+    return;
+  }
+
+  isConfigured() {
+    return this.configured;
+  }
+
+  configure(buffersConf) {
+    /** Configuration cannot be executed more than once. */
+    if (this.configured == true) {
+      throw Error("BufferManager can be configured only once");
+    }
+    this.configured = true;
+
+    /** Parse configuration. */
+    if (Array.isArray(buffersConf) == false) {
+      throw Error("Buffers configuration should be an array");
+    }
+    for (let i = 0; i < buffersConf.length; i++) {
+      let matchTask = buffersConf[i]['matchTask'];
+      let size = buffersConf[i]['size'];
+      let timeoutMs = buffersConf[i]['timeoutMs'];
+      if (matchTask == undefined || size === undefined || timeoutMs == undefined) {
+        throw Error("Following keys are required: matchTask, size, timeoutMs");
+      }
+      let res = this.buffers.push(new Buffer(size, timeoutMs, this.cb));
+      let buffIndex = res - 1;
+
+      /** Build map of taskName -> bufferId. */
+      for (let j = 0; j < matchTask.length; j++) {
+        let taskName = matchTask[j];
+        if (this.taskBufferMap[taskName] != undefined) {
+          console.log("WARNING: task", taskName, "is already matched in another buffer, ignoring");
+          continue;
+        }
+        this.taskBufferMap[taskName] = buffIndex;
+      }
+    }
+
+    return;
+  }
+
+  addItem(taskName, item) {
+    let bufferId = this.taskBufferMap[taskName];
+    /** If task is not buffered, then execute callback immediately. */
+    if (bufferId == undefined) {
+      this.cb([item]);
+      return;
+    }
+
+    /** Buffering item. */
+    this.buffers[bufferId].addItem(item);
+
+    return;
+  }
+}
+
+async function testBufferManager() {
+  let cb = (items) => {
+    console.log("Got from buffer:", items);
+  }
+  buffersConf = [
+    {
+      matchTask: ['job_a', 'job_b'],
+      size: 2,
+      timeoutMs: 3000,
+    },
+    {
+      matchTask: ['job_b', 'job_c'],
+      size: 3,
+      timeoutMs: 2500,
+    },
+  ];
+  let test = new BufferManager(cb)
+  test.configure(buffersConf);
+  test.addItem('job_a', 1);
+  test.addItem('job_a', 2);
+  test.addItem('job_b', 3);
+  test.addItem('job_b', 4);
+  test.addItem('job_b', 5);
+
+  await new Promise(resolve => setTimeout(resolve, 6000));
+
+  test.addItem('job_a', 6);
+  test.addItem('job_c', 7);
+  test.addItem('job_c', 8);
+}
+
+exports.BufferManager = BufferManager;

functions/kubernetes/k8sCommand.js

@@ -1,21 +1,80 @@
 // Runs a job as a Pod (Kubernetes Job) in a Kubernetes cluster
 
 const k8s = require('@kubernetes/client-node');
+var BufferManager = require('./buffer_manager.js').BufferManager;
+var RestartCounter = require('./restart_counter.js').RestartCounter;
 var submitK8sJob = require('./k8sJobSubmit.js').submitK8sJob;
 var fs = require('fs');
 
-async function k8sCommand(ins, outs, context, cb) {
+let bufferManager = new BufferManager();
+
+let backoffLimit = process.env.HF_VAR_BACKOFF_LIMIT || 0;
+let restartCounter = new RestartCounter(backoffLimit);
+
+// Function k8sCommandGroup
+//
+// Inputs:
+// - bufferItems - array containing objects with following properties:
+//   * ins
+//   * outs
+//   * context
+//   * cb
+async function k8sCommandGroup(bufferItems) {
+
+  // No action needed when buffer is empty
+  if (bufferItems.length == 0) {
+    return;
+  }
 
   let startTime = Date.now();
-  console.log("k8sCommand started, time:", startTime);
+  console.log("k8sCommandGroup started, time:", startTime);
+
+  // Function for rebuffering items
+  let restartFn = (bufferIndex) => {
+    let bufferItem = bufferItems[bufferIndex];
+    let taskId = bufferItem.context.taskId;
+    if (restartCounter.isRestartPossible(taskId)) {
+      let restartVal = restartCounter.increase(taskId);
+      console.log("Readding task", taskId, "to buffer (restartCount:", restartVal + ") ...");
+      let itemName = bufferItem.context.name;
+      bufferManager.addItem(itemName, bufferItem);
+    }
+    return;
+  }
+
+  // Extract particular arrays from buffer items
+  let jobArr = [];
+  let taskIdArr = [];
+  let contextArr = [];
+  let cbArr = [];
+  for (let i=0; i<bufferItems.length; i++) {
+    let bufferItem = bufferItems[i];
+    let ins = bufferItem.ins;
+    let outs = bufferItem.outs;
+    let context = bufferItem.context;
+    let cb = bufferItem.cb;
+
+    var job = context.executor; // object containing 'executable', 'args' and others
+    job.name = context.name;
+    job.ins = ins;
+    job.outs = outs;
+
+    jobArr.push(job);
+    taskIdArr.push(context.taskId);
+    contextArr.push(context);
+    cbArr.push(cb);
+  }
+
+  let context = contextArr[0];
+
   // let cluster = await getCluster();
   // const token = await getGCPToken();
 
-  // support for two (or more) clusters (for cloud bursting) 
+  // support for two (or more) clusters (for cloud bursting)
   // if 'partition' is defined, check if there is a custom config file
   // for that partition. This config file may override parameters of the job,
   // possibly even define a path to a different kube_config to be loaded
-  let partition = context.executor.partition; 
+  let partition = context.executor.partition;
   let partitionConfigDir = process.env.HF_VAR_PARTITION_CONFIG_DIR || "/opt/hyperflow/partitions";
   let partitionConfigFile = partitionConfigDir + "/" + "part." + partition + ".config.json";
 
@@ -30,7 +89,7 @@ async function k8sCommand(ins, outs, context, cb) {
   console.log(partitionConfigFile);
   console.log("CUSTOM...", customParams);
 
-  // Set kube_config path if overridden 
+  // Set kube_config path if overridden
   if (customParams.kubeConfigPath) {
       process.env.KUBECONFIG = customParams.kubeConfigPath;
       console.log(process.env.KUBECONFIG);
@@ -39,31 +98,81 @@ async function k8sCommand(ins, outs, context, cb) {
   const kubeconfig = new k8s.KubeConfig();
   kubeconfig.loadFromDefault(); // loadFromString(JSON.stringify(kconfig))
 
-  var job = context.executor; // object containing 'executable', 'args' and others
-  job.name = context.name;
-  job.ins = ins;
-  job.outs = outs;
-
-  let jobExitCode = await submitK8sJob(kubeconfig, job, context.taskId, context, customParams);
+  let jobExitCodes = [];
+  try {
+    jobExitCodes = await submitK8sJob(kubeconfig, jobArr, taskIdArr, contextArr, customParams, restartFn);
+  } catch (err) {
+    console.log("Error when submitting job:", err);
+    throw err;
+  }
 
   let endTime = Date.now();
-  console.log("Ending k8sCommand function, time:", endTime);
+  console.log("Ending k8sCommandGroup function, time:", endTime, "exit codes:", jobExitCodes);
 
   // Stop the entire workflow if a job fails (controlled by an environment variable)
-  if (jobExitCode != 0 && process.env.HF_VAR_STOP_WORKFLOW_WHEN_JOB_FAILED=="1") {
-    console.log('Error: job exited with error code, stopping workflow.');
-    console.log('Error details: job.name: ' + job.name + ', job.args: ' + job.args.join(' '));
-    process.exit(1);
+  for (var i=0; i<jobExitCodes.length; i++) {
+    let jobExitCode = jobExitCodes[i];
+    if (jobExitCode != 0 && process.env.HF_VAR_STOP_WORKFLOW_WHEN_JOB_FAILED=="1") {
+      let taskId = taskIdArr[i];
+      let job = jobArr[i];
+      console.log('Error: job', taskId, 'exited with error code', jobExitCode, ', stopping workflow.');
+      console.log('Error details: job.name: ' + job.name + ', job.args: ' + job.args.join(' '));
+      process.exit(1);
+    }
   }
 
   // if we're here, the job should have succesfully completed -- we write this
   // information to Redis (job executor may make use of it).
+  let markPromises = [];
+  for (var i=0; i<contextArr.length; i++) {
+    // skip failed jobs
+    if (jobExitCodes[i] != 0) {
+      continue;
+    }
+
+    let context = contextArr[i];
+    markPromises.push(context.markTaskCompleted());
+  }
   try {
-    await context.markTaskCompleted();
+    await Promise.all(markPromises);
   } catch {
-    console.error("Marking job", context.taskId, "as completed failed.")
+    console.error("Marking jobs", taskIdArr, "as completed failed.")
+  }
+
+  for (var i=0; i<cbArr.length; i++) {
+    // skip failed jobs
+    if (jobExitCodes[i] != 0) {
+      continue;
+    }
+
+    let cb = cbArr[i];
+    let outs = jobArr[i].outs;
+    cb(null, outs);
+  }
+
+  return;
+}
+
+bufferManager.setCallback((items) => k8sCommandGroup(items));
+
+async function k8sCommand(ins, outs, context, cb) {
+  /** Buffer Manager configuration. */
+  buffersConf = context.appConfig.jobAgglomerations;
+  let alreadyConfigured = bufferManager.isConfigured();
+  if (alreadyConfigured == false && buffersConf != undefined) {
+    bufferManager.configure(buffersConf);
   }
-  cb(null, outs);
+
+  /** Buffer item. */
+  let item = {
+    "ins": ins,
+    "outs": outs,
+    "context": context,
+    "cb": cb
+  };
+  bufferManager.addItem(context.name, item);
+
+  return;
 }
 
 exports.k8sCommand = k8sCommand;