Add opt-in same-node scheduling preference for job and step pods

packages/k8s/src/k8s/index.ts

@@ -16,6 +16,8 @@ import {
   PodPhase,
   mergePodSpecWithOptions,
   mergeObjectMeta,
+  injectSameNodePreference,
+  useSameNodePreference,
   fixArgs,
   listDirAllCommand,
   sleep,
@@ -35,6 +37,30 @@ const k8sAuthorizationV1Api = kc.makeApiClient(k8s.AuthorizationV1Api)
 
 const DEFAULT_WAIT_FOR_POD_TIME_SECONDS = 10 * 60 // 10 min
 
+async function applySameNodePreference(spec: k8s.V1PodSpec): Promise<void> {
+  if (!useSameNodePreference()) {
+    return
+  }
+  try {
+    const runnerPodName = process.env.ACTIONS_RUNNER_POD_NAME
+    if (!runnerPodName) {
+      return
+    }
+    const runnerPod = await k8sApi.readNamespacedPod({
+      name: runnerPodName,
+      namespace: namespace()
+    })
+    const runnerNodeName = runnerPod.spec?.nodeName
+    if (runnerNodeName) {
+      injectSameNodePreference(spec, runnerNodeName)
+    }
+  } catch (err) {
+    core.warning(
+      `Could not look up runner pod node for same-node scheduling: ${err}`
+    )
+  }
+}
+
 export const requiredPermissions = [
   {
     group: '',
@@ -176,6 +202,8 @@ export async function createJobPod(
     mergePodSpecWithOptions(appPod.spec, extension.spec)
   }
 
+  await applySameNodePreference(appPod.spec)
+
   return await k8sApi.createNamespacedPod({
     namespace: namespace(),
     body: appPod
@@ -229,6 +257,8 @@ export async function createContainerStepPod(
     mergePodSpecWithOptions(appPod.spec, extension.spec)
   }
 
+  await applySameNodePreference(appPod.spec)
+
   return await k8sApi.createNamespacedPod({
     namespace: namespace(),
     body: appPod

packages/k8s/src/k8s/utils.ts

@@ -12,6 +12,7 @@ export const DEFAULT_CONTAINER_ENTRY_POINT = 'tail'
 
 export const ENV_HOOK_TEMPLATE_PATH = 'ACTIONS_RUNNER_CONTAINER_HOOK_TEMPLATE'
 export const ENV_USE_KUBE_SCHEDULER = 'ACTIONS_RUNNER_USE_KUBE_SCHEDULER'
+export const ENV_SAME_NODE_PREFERENCE = 'ACTIONS_RUNNER_SAME_NODE_PREFERENCE'
 
 export const EXTERNALS_VOLUME_NAME = 'externals'
 export const GITHUB_VOLUME_NAME = 'github'
@@ -269,6 +270,42 @@ export function useKubeScheduler(): boolean {
   return process.env[ENV_USE_KUBE_SCHEDULER] === 'true'
 }
 
+export function useSameNodePreference(): boolean {
+  return process.env[ENV_SAME_NODE_PREFERENCE] === 'true'
+}
+
+export function injectSameNodePreference(
+  spec: k8s.V1PodSpec,
+  nodeName: string
+): void {
+  if (!spec.affinity) {
+    spec.affinity = {}
+  }
+  if (!spec.affinity.nodeAffinity) {
+    spec.affinity.nodeAffinity = {}
+  }
+  if (
+    !spec.affinity.nodeAffinity.preferredDuringSchedulingIgnoredDuringExecution
+  ) {
+    spec.affinity.nodeAffinity.preferredDuringSchedulingIgnoredDuringExecution =
+      []
+  }
+  spec.affinity.nodeAffinity.preferredDuringSchedulingIgnoredDuringExecution.push(
+    {
+      weight: 100,
+      preference: {
+        matchExpressions: [
+          {
+            key: 'kubernetes.io/hostname',
+            operator: 'In',
+            values: [nodeName]
+          }
+        ]
+      }
+    }
+  )
+}
+
 export enum PodPhase {
   PENDING = 'Pending',
   RUNNING = 'Running',

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

@@ -6,6 +6,9 @@ import {
   mergePodSpecWithOptions,
   mergeContainerWithOptions,
   readExtensionFromFile,
+  injectSameNodePreference,
+  useSameNodePreference,
+  ENV_SAME_NODE_PREFERENCE,
   ENV_HOOK_TEMPLATE_PATH
 } from '../src/k8s/utils'
 import * as k8s from '@kubernetes/client-node'
@@ -339,6 +342,123 @@ spec:
     expect(base).toStrictEqual(expectContainer)
   })
 
+  describe('useSameNodePreference', () => {
+    const originalEnv = process.env
+
+    afterEach(() => {
+      process.env = originalEnv
+    })
+
+    it('should return false when env var is not set', () => {
+      delete process.env[ENV_SAME_NODE_PREFERENCE]
+      expect(useSameNodePreference()).toBe(false)
+    })
+
+    it('should return false when env var is empty', () => {
+      process.env[ENV_SAME_NODE_PREFERENCE] = ''
+      expect(useSameNodePreference()).toBe(false)
+    })
+
+    it('should return false when env var is not "true"', () => {
+      process.env[ENV_SAME_NODE_PREFERENCE] = 'false'
+      expect(useSameNodePreference()).toBe(false)
+
+      process.env[ENV_SAME_NODE_PREFERENCE] = 'TRUE'
+      expect(useSameNodePreference()).toBe(false)
+
+      process.env[ENV_SAME_NODE_PREFERENCE] = '1'
+      expect(useSameNodePreference()).toBe(false)
+    })
+
+    it('should return true when env var is "true"', () => {
+      process.env[ENV_SAME_NODE_PREFERENCE] = 'true'
+      expect(useSameNodePreference()).toBe(true)
+    })
+  })
+
+  describe('injectSameNodePreference', () => {
+    it('should add nodeAffinity to empty spec', () => {
+      const spec = { containers: [] } as k8s.V1PodSpec
+      injectSameNodePreference(spec, 'node-1')
+
+      const preferred =
+        spec.affinity?.nodeAffinity
+          ?.preferredDuringSchedulingIgnoredDuringExecution
+      expect(preferred).toHaveLength(1)
+      expect(preferred![0].weight).toBe(100)
+      expect(
+        preferred![0].preference.matchExpressions![0]
+      ).toStrictEqual({
+        key: 'kubernetes.io/hostname',
+        operator: 'In',
+        values: ['node-1']
+      })
+    })
+
+    it('should append to existing preferred scheduling entries', () => {
+      const spec = {
+        containers: [],
+        affinity: {
+          nodeAffinity: {
+            preferredDuringSchedulingIgnoredDuringExecution: [
+              {
+                weight: 50,
+                preference: {
+                  matchExpressions: [
+                    {
+                      key: 'node.kubernetes.io/instance-type',
+                      operator: 'In',
+                      values: ['m5.xlarge']
+                    }
+                  ]
+                }
+              }
+            ]
+          }
+        }
+      } as k8s.V1PodSpec
+
+      injectSameNodePreference(spec, 'node-2')
+
+      const preferred =
+        spec.affinity!.nodeAffinity!
+          .preferredDuringSchedulingIgnoredDuringExecution!
+      expect(preferred).toHaveLength(2)
+      expect(preferred[0].weight).toBe(50)
+      expect(preferred[1].weight).toBe(100)
+      expect(
+        preferred[1].preference.matchExpressions![0].values
+      ).toStrictEqual(['node-2'])
+    })
+
+    it('should not touch requiredDuringScheduling', () => {
+      const required = {
+        nodeSelectorTerms: [
+          {
+            matchExpressions: [
+              { key: 'zone', operator: 'In', values: ['us-east-1a'] }
+            ]
+          }
+        ]
+      }
+      const spec = {
+        containers: [],
+        affinity: {
+          nodeAffinity: {
+            requiredDuringSchedulingIgnoredDuringExecution: required
+          }
+        }
+      } as k8s.V1PodSpec
+
+      injectSameNodePreference(spec, 'node-3')
+
+      expect(
+        spec.affinity!.nodeAffinity!
+          .requiredDuringSchedulingIgnoredDuringExecution
+      ).toStrictEqual(required)
+    })
+  })
+
   it('should merge pod spec', () => {
     const base = {
       containers: [