agent_sandbox: single-source the sandbox scheduling target

The gVisor scheduling selector and taint were duplicated as literal strings
across the benchmark config, the installer DaemonSet, and the sandbox template,
with nothing keeping them in sync. Untangle scheduling from runtime identity:

- Scheduling: select the sandbox nodepool via the pkb_nodepool label PKB
  already injects on every pool, and derive the pod toleration from a single
  taint constant. nodeSelector/tolerations are now injected in Python (like
  _configure_controller_manifest) instead of being hardcoded in the manifests.
- Runtime identity: runtimeClassName stays runsc, used only for the
  RuntimeClass, containerd registration, and the pod runtimeClassName.

PKB does not yet apply nodepool taints to nodes (that lands in a follow-up), so
the canonical taint lives in a _SANDBOX_TAINT constant with a TODO to read it
from the nodepool config once that wiring exists. The SandboxWarmPool is
unchanged: it inherits scheduling from the SandboxTemplate podTemplate.

Author: geojaz

perfkitbenchmarker/data/agent_sandbox/gvisor-installer/daemonset.yaml

@@ -6,8 +6,9 @@
 # data/agent_sandbox/gvisor-installer/install.sh before this DaemonSet is
 # applied). The ConfigMap key is "install.sh", mounted at /scripts.
 #
-# Targets nodes labelled sandbox.gke.io/runtime=runsc (the label the
-# benchmark applies to the sandbox node pool).
+# nodeSelector and tolerations are injected at apply time (see
+# _render_gvisor_daemonset in resources/kubernetes/k8s_agent_sandbox.py) so the
+# DaemonSet targets the sandbox node pool via the pkb_nodepool label.
 apiVersion: apps/v1
 kind: DaemonSet
 metadata:
@@ -25,13 +26,6 @@ spec:
         app.kubernetes.io/name: gvisor-installer
     spec:
       hostPID: true
-      tolerations:
-        - key: sandbox.gke.io/runtime
-          operator: Equal
-          value: runsc
-          effect: NoSchedule
-      nodeSelector:
-        sandbox.gke.io/runtime: runsc
       initContainers:
         - name: install
           image: docker.io/library/ubuntu:24.04

perfkitbenchmarker/data/agent_sandbox/sandbox-template.yaml.j2

@@ -19,13 +19,6 @@ spec:
       automountServiceAccountToken: false
       securityContext:
         runAsNonRoot: true
-      nodeSelector:
-        sandbox.gke.io/runtime: {{ runtime_class }}
-      tolerations:
-        - key: sandbox.gke.io/runtime
-          operator: Equal
-          value: {{ runtime_class }}
-          effect: NoSchedule
       containers:
         - name: python-runtime
           image: {{ image }}

perfkitbenchmarker/linux_benchmarks/agent_sandbox_benchmark.py

@@ -78,8 +78,6 @@
           AWS:
             machine_type: m8i.4xlarge
             zone: us-east-1a
-        node_labels:
-          sandbox.gke.io/runtime: runsc
         node_taints:
           - sandbox.gke.io/runtime=runsc:NoSchedule
     agent_sandbox:

perfkitbenchmarker/resources/kubernetes/k8s_agent_sandbox.py

@@ -79,6 +79,18 @@
 
 _SANDBOX_NAME = 'agent-sandbox'
 
+# Config key of the nodepool the sandbox workload runs on (matches the
+# `nodepools:` block in linux_benchmarks/agent_sandbox_benchmark.py).
+_SANDBOX_NODEPOOL = 'sandbox'
+
+# Taint fencing the sandbox nodepool. PKB does not yet apply nodepool taints to
+# nodes (that wiring lands in PR #6741), so keep the canonical taint string here
+# and derive the pod toleration from it. Keep this in lockstep with the
+# `node_taints` entry in agent_sandbox_benchmark.py BENCHMARK_CONFIG.
+# TODO(#6741): read this from cluster.nodepools[_SANDBOX_NODEPOOL].node_taints
+# once NodepoolConfig carries node_taints, and delete this constant.
+_SANDBOX_TAINT = 'sandbox.gke.io/runtime=runsc:NoSchedule'
+
 
 def _crd_name(filename):
   """Derives a CRD resource name from its manifest filename.
@@ -109,17 +121,88 @@ def _wait_warmpool_ready(warmpool_name, replicas, timeout=600):
   )
 
 
-def install_gvisor():
+def _taint_to_toleration(taint):
+  """Converts a 'key=value:Effect' or 'key:Effect' taint to a toleration dict.
+
+  Parsing mirrors the planned EKS taint parser (PR #6744) but emits a
+  Kubernetes toleration so the node taint and the pod toleration stay symmetric.
+  """
+  spec, sep_effect, effect = taint.rpartition(':')
+  if not sep_effect:
+    raise ValueError(
+        f'Malformed taint, expected key[=value]:Effect: {taint!r}'
+    )
+  key, sep_val, value = spec.partition('=')
+  toleration = {'key': key}
+  if sep_val:
+    toleration['operator'] = 'Equal'
+    toleration['value'] = value
+  else:
+    toleration['operator'] = 'Exists'
+  toleration['effect'] = effect
+  return toleration
+
+
+def _sandbox_scheduling(nodepool_name):
+  """Returns (node_selector, tolerations) targeting the sandbox nodepool.
+
+  The selector rendezvous is the pkb_nodepool label PKB stamps on every node
+  pool (GKE _AddNodeParamsToCmd, EKS _RenderNodeGroupJson), so there is no
+  bespoke label to keep in sync. The toleration is derived from _SANDBOX_TAINT.
+  """
+  node_selector = {'pkb_nodepool': nodepool_name}
+  tolerations = [_taint_to_toleration(_SANDBOX_TAINT)]
+  return node_selector, tolerations
+
+
+def _render_gvisor_daemonset(node_selector, tolerations):
+  """Loads the installer DaemonSet and injects the sandbox scheduling target."""
+  with open(data.ResourcePath(_GVISOR_DAEMONSET)) as manifest_file:
+    manifest = yaml.safe_load(manifest_file)
+  pod_spec = manifest['spec']['template']['spec']
+  pod_spec['nodeSelector'] = node_selector
+  pod_spec['tolerations'] = tolerations
+  return yaml.dump(manifest, default_flow_style=False)
+
+
+def _render_template_manifest(template_spec, node_selector, tolerations):
+  """Renders the SandboxTemplate and injects the sandbox scheduling target.
+
+  nodeSelector/tolerations are set in Python (not the .j2) so the scheduling
+  target has a single source of truth. runtimeClassName stays in the template:
+  it is runtime identity, not scheduling.
+  """
+  labels = template_spec.labels or {'sandbox': 'python-sandbox-bench'}
+  rendered = vm_util.ReadAndRenderJinja2Template(
+      _TEMPLATE_MANIFEST,
+      trim_spaces=False,
+      name=_SANDBOX_NAME,
+      runtime_class=template_spec.runtime_class,
+      image=template_spec.image,
+      cpu_request=template_spec.cpu_request,
+      cpu_limit=template_spec.cpu_limit,
+      memory_request=template_spec.memory_request,
+      memory_limit=template_spec.memory_limit,
+      labels=labels,
+  )
+  manifest = yaml.safe_load(rendered)
+  pod_spec = manifest['spec']['podTemplate']['spec']
+  pod_spec['nodeSelector'] = node_selector
+  pod_spec['tolerations'] = tolerations
+  return yaml.dump(manifest, default_flow_style=False)
+
+
+def install_gvisor(node_selector, tolerations):
   """Installs gVisor onto cluster nodes via the installer DaemonSet.
 
   Creates the gvisor-installer-script ConfigMap (from install.sh) in
   kube-system before applying the DaemonSet. The DaemonSet mounts that
   ConfigMap at /scripts; the init container runs /scripts/install.sh.
-  The ConfigMap must exist before the DaemonSet pods schedule.
+  The ConfigMap must exist before the DaemonSet pods schedule. The DaemonSet's
+  nodeSelector/tolerations are injected so it lands on the sandbox nodepool.
   """
   _create_installer_configmap()
-  # Plain .yaml files: ApplyManifest with NO kwargs.
-  kubernetes_commands.ApplyManifest(_GVISOR_DAEMONSET)
+  _apply_yaml(_render_gvisor_daemonset(node_selector, tolerations))
   kubernetes_commands.ApplyManifest(_GVISOR_RUNTIMECLASS)
   kubernetes_commands.WaitForRollout(
       'daemonset/gvisor-installer', namespace='kube-system'
@@ -302,19 +385,10 @@ def install_controller(
   )
 
 
-def apply_template(template_spec):
+def apply_template(template_spec, node_selector, tolerations):
   """Applies the SandboxTemplate rendered from the template spec."""
-  labels = template_spec.labels or {'sandbox': 'python-sandbox-bench'}
-  kubernetes_commands.ApplyManifest(
-      _TEMPLATE_MANIFEST,
-      name=_SANDBOX_NAME,
-      runtime_class=template_spec.runtime_class,
-      image=template_spec.image,
-      cpu_request=template_spec.cpu_request,
-      cpu_limit=template_spec.cpu_limit,
-      memory_request=template_spec.memory_request,
-      memory_limit=template_spec.memory_limit,
-      labels=labels,
+  _apply_yaml(
+      _render_template_manifest(template_spec, node_selector, tolerations)
   )
 
 
@@ -392,7 +466,8 @@ def _Delete(self):
     pass
 
   def _InstallGvisor(self):
-    install_gvisor()
+    node_selector, tolerations = self._SandboxScheduling()
+    install_gvisor(node_selector, tolerations)
 
   def _InstallCrdsAndRbac(self):
     install_crds_and_rbac(self.spec.manifest_ref)
@@ -405,7 +480,18 @@ def _InstallController(self):
     )
 
   def _ApplyTemplate(self):
-    apply_template(self.spec.sandbox_template)
+    node_selector, tolerations = self._SandboxScheduling()
+    apply_template(self.spec.sandbox_template, node_selector, tolerations)
+
+  def _SandboxScheduling(self):
+    """Resolves (node_selector, tolerations) from the sandbox nodepool."""
+    nodepool = self.cluster.nodepools.get(_SANDBOX_NODEPOOL)
+    if nodepool is None:
+      raise ValueError(
+          f'Agent sandbox requires a nodepool named {_SANDBOX_NODEPOOL!r}; '
+          'add it to the benchmark container_cluster nodepools config.'
+      )
+    return _sandbox_scheduling(nodepool.name)
 
   def _InstallWarmpool(self):
     install_warmpool(self.spec.sandbox_warmpool.replicas)

tests/resources/kubernetes/k8s_agent_sandbox_test.py

@@ -136,7 +136,11 @@ def _Sandbox(self, **template_overrides):
         sandbox_warmpool={'replicas': 3},
         sandbox_template=template_overrides or {'runtime_class': 'runsc'},
     )
-    return k8s_agent_sandbox.K8sAgentSandbox(sandbox_spec, mock.Mock())
+    cluster = mock.Mock()
+    nodepool = mock.Mock()
+    nodepool.name = 'sandbox'
+    cluster.nodepools = {'sandbox': nodepool}
+    return k8s_agent_sandbox.K8sAgentSandbox(sandbox_spec, cluster)
 
   @mock.patch.object(k8s_agent_sandbox, 'install_crds_and_rbac')
   @mock.patch.object(k8s_agent_sandbox, 'install_gvisor')
@@ -231,5 +235,67 @@ def testConfigBuildsK8sAgentSandbox(self):
     self.assertIsInstance(sandbox, k8s_agent_sandbox.K8sAgentSandbox)
 
 
+class SandboxSchedulingTest(pkb_common_test_case.PkbCommonTestCase):
+
+  def testTaintToTolerationWithValue(self):
+    self.assertEqual(
+        k8s_agent_sandbox._taint_to_toleration(
+            'sandbox.gke.io/runtime=runsc:NoSchedule'),
+        {
+            'key': 'sandbox.gke.io/runtime',
+            'operator': 'Equal',
+            'value': 'runsc',
+            'effect': 'NoSchedule',
+        })
+
+  def testTaintToTolerationNoValue(self):
+    self.assertEqual(
+        k8s_agent_sandbox._taint_to_toleration('dedicated:NoSchedule'),
+        {'key': 'dedicated', 'operator': 'Exists', 'effect': 'NoSchedule'})
+
+  def testTaintToTolerationMalformedRaises(self):
+    with self.assertRaises(ValueError):
+      k8s_agent_sandbox._taint_to_toleration('no-effect')
+
+  def testSandboxSchedulingSelectorAndToleration(self):
+    node_selector, tolerations = k8s_agent_sandbox._sandbox_scheduling('sandbox')
+    self.assertEqual(node_selector, {'pkb_nodepool': 'sandbox'})
+    self.assertEqual(tolerations, [{
+        'key': 'sandbox.gke.io/runtime',
+        'operator': 'Equal',
+        'value': 'runsc',
+        'effect': 'NoSchedule',
+    }])
+
+  def testRenderGvisorDaemonsetSchedulesOnPkbNodepool(self):
+    node_selector, tolerations = k8s_agent_sandbox._sandbox_scheduling('sandbox')
+    manifest = yaml.safe_load(
+        k8s_agent_sandbox._render_gvisor_daemonset(node_selector, tolerations))
+    pod_spec = manifest['spec']['template']['spec']
+    self.assertEqual(pod_spec['nodeSelector'], {'pkb_nodepool': 'sandbox'})
+    self.assertEqual(pod_spec['tolerations'], tolerations)
+
+  def testRenderTemplateManifestSchedulingAndRuntimeClass(self):
+    template_spec = mock.Mock()
+    template_spec.runtime_class = 'runsc'
+    template_spec.image = 'img:latest'
+    template_spec.cpu_request = '500m'
+    template_spec.cpu_limit = '2'
+    template_spec.memory_request = '256Mi'
+    template_spec.memory_limit = '1Gi'
+    template_spec.labels = {'sandbox': 'python-sandbox-bench'}
+    node_selector, tolerations = k8s_agent_sandbox._sandbox_scheduling('sandbox')
+    manifest = yaml.safe_load(
+        k8s_agent_sandbox._render_template_manifest(
+            template_spec, node_selector, tolerations))
+    pod_spec = manifest['spec']['podTemplate']['spec']
+    self.assertEqual(pod_spec['nodeSelector'], {'pkb_nodepool': 'sandbox'})
+    self.assertEqual(pod_spec['tolerations'], tolerations)
+    # runtimeClassName stays as runtime identity, not scheduling.
+    self.assertEqual(pod_spec['runtimeClassName'], 'runsc')
+    # The old runtime label is no longer used as a node selector.
+    self.assertNotIn('sandbox.gke.io/runtime', pod_spec['nodeSelector'])
+
+
 if __name__ == '__main__':
   unittest.main()