From 4f9272a5ce61e454a9524b5dd108cad8e88e727b Mon Sep 17 00:00:00 2001
From: Nikola Jokic <nikola-jokic@github.com>
Date: Fri, 3 Jun 2022 11:38:59 +0200
Subject: [PATCH] Added resource requirement based on --cpus and --memory,-m

---
 packages/k8s/src/hooks/prepare-job.ts | 64 +++++++++++++++++++++++++++
 1 file changed, 64 insertions(+)

diff --git a/packages/k8s/src/hooks/prepare-job.ts b/packages/k8s/src/hooks/prepare-job.ts
index df92e44..1e2fab6 100644
--- a/packages/k8s/src/hooks/prepare-job.ts
+++ b/packages/k8s/src/hooks/prepare-job.ts
@@ -175,10 +175,15 @@ function createPodSpec(
     args: container.entryPointArgs,
     ports: containerPorts(container)
   } as k8s.V1Container
+
   if (container.workingDirectory) {
     podContainer.workingDir = container.workingDirectory
   }
 
+  if (container.createOptions) {
+    podContainer.resources = getResourceRequirements(container.createOptions)
+  }
+
   podContainer.env = []
   for (const [key, value] of Object.entries(
     container['environmentVariables']
@@ -195,3 +200,62 @@ function createPodSpec(
 
   return podContainer
 }
+
+function getResourceRequirements(
+  createOptions: string
+): k8s.V1ResourceRequirements {
+  const rr = new k8s.V1ResourceRequirements()
+  rr.limits = {}
+  rr.requests = {}
+
+  const options = parseOptions(createOptions)
+  for (const [key, value] of Object.entries(options)) {
+    switch (key) {
+      case '--cpus':
+        rr.requests.cpu = value
+        break
+      case '--memory':
+      case '-m':
+        rr.limits.memory = value
+        break
+      default:
+        core.warning(
+          `Container option ${key} is not supported. Supported options are ['--cpus', '--memory', '-m']`
+        )
+    }
+  }
+
+  return rr
+}
+
+function parseOptions(options: string): { [option: string]: string } {
+  const rv: { [option: string]: string } = {}
+
+  const spaceSplit = options.split(' ')
+  for (let i = 0; i < spaceSplit.length; i++) {
+    if (!spaceSplit[i].startsWith('-')) {
+      throw new Error(`Options specified in wrong format: ${options}`)
+    }
+
+    const optSplit = spaceSplit[i].split('=')
+    const optName = optSplit[0]
+    let optValue = ''
+    switch (optSplit.length) {
+      case 1:
+        if (spaceSplit.length <= i + 1) {
+          throw new Error(`Option ${optName} must have a value`)
+        }
+        optValue = spaceSplit[++i]
+        break
+      case 2:
+        optValue = optSplit[1]
+        break
+      default:
+        throw new Error(`failed to parse option ${spaceSplit[i]}`)
+    }
+
+    rv[optName] = optValue
+  }
+
+  return rv
+}