Add logic to extract TPU requests and add to container spec

Author: quoctruong

packages/k8s/src/hooks/prepare-job.ts

@@ -36,7 +36,8 @@ import {
   getNumberOfHost,
   sleep,
   generateServicesName,
-  getEntryPointAndArgs
+  getEntryPointAndArgs,
+  getTpuRequest
 } from '../k8s/utils'
 import {
   CONTAINER_EXTENSION_PREFIX,
@@ -308,6 +309,7 @@ export function createContainerSpec(
     }
   }
 
+  let tpuRequest = 0
   if (!jobContainer && createOptions && createOptions?.length > 0) {
     core.debug(
       `overriding service container ${JSON.stringify(
@@ -320,6 +322,8 @@ export function createContainerSpec(
       container.entryPoint = entryPointAndArgs[0]
       container.entryPointArgs = entryPointAndArgs.slice(1)
     }
+    tpuRequest = getTpuRequest(createOptions)
+    core.debug(`TPU request from service container is ${tpuRequest}`)
   }
 
   const podContainer = {
@@ -339,6 +343,18 @@ export function createContainerSpec(
     podContainer.args = fixArgs(container.entryPointArgs)
   }
 
+  if (tpuRequest > 0) {
+    core.debug(`assigning ${tpuRequest} to podContainer`)
+    podContainer.resources = {
+      limits: {
+        'google.com/tpu': String(tpuRequest)
+      },
+      requests: {
+        'google.com/tpu': String(tpuRequest)
+      }
+    }
+  }
+
   podContainer.env = []
   for (const [key, value] of Object.entries(
     container['environmentVariables']

packages/k8s/src/k8s/utils.ts

@@ -467,14 +467,25 @@ export function createScriptExecutorContainer(
 }
 
 const entrypointRegex = /--entrypoint=\[(.*?)\]/
+const tpuRegex = /--tpu=([0-9]+)/
 
 export function getEntryPointAndArgs(createOptions: string): string[] {
   const match = createOptions.match(entrypointRegex)
 
   if (!match || match[1] === undefined) {
-    core.debug(`no match for createoptions ${createOptions}`)
+    core.debug(`no entrypoint found for createOptions ${createOptions}`)
     return []
   }
 
   return match[1].split(',').map(item => item.trim())
 }
+
+export function getTpuRequest(createOptions: string): number {
+  const match = createOptions.match(tpuRegex)
+  if (!match || match[1] === undefined) {
+    core.debug(`no tpu found for ${createOptions}`)
+    return 0
+  }
+
+  return Number(match[1])
+}

packages/k8s/tests/k8s-utils-test.ts

@@ -13,7 +13,8 @@ import {
   getNumberOfHost,
   ENV_NUMBER_OF_HOSTS,
   generateServicesName,
-  getEntryPointAndArgs
+  getEntryPointAndArgs,
+  getTpuRequest
 } from '../src/k8s/utils'
 import * as k8s from '@kubernetes/client-node'
 import { TestHelper } from './test-setup'
@@ -389,6 +390,13 @@ describe('k8s utils', () => {
     })
   })
 
+  describe('getTpuRequest', () => {
+    it('should return correct number of TPU requested', () => {
+      expect(getTpuRequest('--tpu=4')).toEqual(4)
+      expect(getTpuRequest('--nothing')).toEqual(0)
+    })
+  })
+
   describe('create script executor container', () => {
     it('should install script executor at the volume mount location', () => {
       const executorVolumeMount = new k8s.V1VolumeMount()