Support reading & applying manifests as yaml rather than straight jinja

A number of tricky python pieces in this:
- Convert to the recursive defaultdict to allow for creating new dictionaries when accessed with eg dict1['foo']['bar']['zip'] = 4.
- Convert back to regular dicts so yaml.dump works.
- Swapping from yield to list with ParseApplyOutput as otherwise the generator delayed/avoided execution of the important _RunKubectlCommand statement.

PiperOrigin-RevId: 842884621

Author: hubatish

perfkitbenchmarker/container_service.py

@@ -38,7 +38,6 @@
 from typing import Any, Callable, Dict, Iterable, Iterator, Optional, Sequence
 
 from absl import flags
-import jinja2
 from perfkitbenchmarker import context
 from perfkitbenchmarker import data
 from perfkitbenchmarker import errors
@@ -1001,6 +1000,33 @@ def ApplyManifest(
     Returns:
       Names of the resources, e.g. [deployment.apps/mydeploy, pod/foo]
     """
+    filename = data.ResourcePath(manifest_file)
+    if not filename.endswith('.j2'):
+      if kwargs:
+        raise AssertionError(
+            'kwargs should be empty for non-jinja templates. '
+            f'Found when reading {manifest_file}'
+        )
+      return KubernetesClusterCommands._ApplyRenderedManifest(filename)
+
+    manifest = vm_util.ReadAndRenderJinja2Template(
+        manifest_file, trim_spaces=False, **kwargs
+    )
+    return KubernetesClusterCommands._WriteAndApplyManifest(
+        manifest, should_log_file
+    )
+
+  @staticmethod
+  def _ApplyRenderedManifest(manifest_file: str) -> Iterator[str]:
+    """Applies a rendered Kubernetes manifest & returns resources created.
+
+    Args:
+      manifest_file: The full path of the YAML file.
+
+    Returns:
+      Names of the resources, e.g. [deployment.apps/mydeploy, pod/foo]
+    """
+    out, _, _ = RunKubectlCommand(['apply', '-f', manifest_file])
 
     def _ParseApplyOutput(stdout: str) -> Iterator[str]:
       """Parses the output of kubectl apply to get the name of the resource."""
@@ -1010,27 +1036,131 @@ def _ParseApplyOutput(stdout: str) -> Iterator[str]:
         if match:
           yield match.group(1)
 
-    filename = data.ResourcePath(manifest_file)
-    if not filename.endswith('.j2'):
-      assert not kwargs
-      out, _, _ = RunKubectlCommand(['apply', '-f', filename])
-      return _ParseApplyOutput(out)
+    # Inner function needed to run kubectl command even if output is unused.
+    return _ParseApplyOutput(out)
 
-    environment = jinja2.Environment(undefined=jinja2.StrictUndefined)
-    with open(filename) as template_file, vm_util.NamedTemporaryFile(
+  @staticmethod
+  def _WriteAndApplyManifest(
+      manifest: str, should_log_file: bool = True
+  ) -> Iterator[str]:
+    """Writes the string to file and applies it."""
+    with vm_util.NamedTemporaryFile(
         mode='w', suffix='.yaml'
     ) as rendered_template:
-      manifest = environment.from_string(template_file.read()).render(kwargs)
       rendered_template.write(manifest)
       rendered_template.close()
       if should_log_file:
         logging.info(
             'Rendered manifest file %s with contents:\n%s',
             rendered_template.name,
             manifest,
+            stacklevel=2,
+        )
+      resource_names = KubernetesClusterCommands._ApplyRenderedManifest(
+          rendered_template.name
+      )
+      return resource_names
+
+  @staticmethod
+  def _RecursiveDict() -> dict[str, Any]:
+    """Creates a new dictionary with auto nesting keys.
+
+    See https://stackoverflow.com/a/10218517/2528472.
+
+    Returns:
+      A new dictionary that automatically sets the value for any nested missing
+      keys.
+    """
+    return collections.defaultdict(KubernetesClusterCommands._RecursiveDict)
+
+  @staticmethod
+  def ConvertManifestToYamlDicts(
+      manifest_file: str, **kwargs
+  ) -> list[dict[str, Any]]:
+    """Reads a Kubernetes manifest from a file and converts it to YAML.
+
+    Args:
+      manifest_file: The name of the YAML file or YAML template.
+      **kwargs: Arguments to the jinja template.
+
+    Returns:
+      The various YAML documents as a list of dictionaries.
+    """
+    manifest = vm_util.ReadAndRenderJinja2Template(
+        manifest_file, trim_spaces=False, **kwargs
+    )
+    yaml_docs = yaml.safe_load_all(manifest)
+    return_yamls = []
+    for yaml_doc in yaml_docs:
+      return_yamls.append(
+          KubernetesClusterCommands._ConvertToDictType(
+              yaml_doc, KubernetesClusterCommands._RecursiveDict
+          )
+      )
+    return return_yamls
+
+  @staticmethod
+  def _ConvertToDictType(
+      yaml_doc: Any, dict_lambda: Any
+  ) -> dict[str, Any] | Any:
+    """Converts a YAML document to the given dictionary type.
+
+    In particular a recursive defaultdict can be more easily accessed with e.g.
+    my_dict['a']['b'] = value rather than my_dict.setdefault('a', {})['b'] =
+    value.
+
+    Args:
+      yaml_doc: The YAML document to convert.
+      dict_lambda: A constructor for the dictionary type to convert to.
+
+    Returns:
+      The remade dictionary.
+    """
+    if not isinstance(yaml_doc, dict) and not isinstance(yaml_doc, list):
+      return yaml_doc
+
+    yaml_list = []
+    if isinstance(yaml_doc, list):
+      for item in yaml_doc:
+        if not bool(item) and item != 0:
+          continue
+        yaml_list.append(
+            KubernetesClusterCommands._ConvertToDictType(item, dict_lambda)
         )
-      out, _, _ = RunKubectlCommand(['apply', '-f', rendered_template.name])
-      return _ParseApplyOutput(out)
+      return yaml_list
+    yaml_dict = dict_lambda()
+    for key, value in yaml_doc.items():
+      if not bool(value) and value != 0:
+        continue
+      yaml_dict[key] = KubernetesClusterCommands._ConvertToDictType(
+          value, dict_lambda
+      )
+    return yaml_dict
+
+  @staticmethod
+  def ApplyYaml(
+      yaml_dicts: list[dict[str, Any]], should_log_file: bool = True
+  ) -> Iterator[str]:
+    """Writes yaml to a file and applies it.
+
+    Args:
+      yaml_dicts: A list of YAML documents.
+      should_log_file: Whether to log the rendered manifest to stdout or not.
+
+    Returns:
+      Names of the resources, e.g. [deployment.apps/mydeploy, pod/foo]
+    """
+    normal_dicts = []
+    # Convert back to a normal dict because yaml.dump otherwise adds random
+    # "dictitems:" keys & other python artifacts.
+    for yaml_dict in yaml_dicts:
+      normal_dicts.append(
+          KubernetesClusterCommands._ConvertToDictType(yaml_dict, dict)
+      )
+    manifest = yaml.dump_all(normal_dicts)
+    return KubernetesClusterCommands._WriteAndApplyManifest(
+        manifest, should_log_file
+    )
 
   @staticmethod
   def WaitForResource(

perfkitbenchmarker/linux_benchmarks/kubernetes_scale_benchmark.py

@@ -180,7 +180,7 @@ def ScaleUpPods(
   max_wait_time = _GetScaleTimeout()
   resource_timeout = max_wait_time + 60 * 5  # 5 minutes after waiting to avoid
   # pod delete events from polluting data collection.
-  resource_names = cluster.ApplyManifest(
+  scaleup_yaml = cluster.ConvertManifestToYamlDicts(
       MANIFEST_TEMPLATE,
       Name='kubernetes-scaleup',
       Replicas=num_new_pods,
@@ -197,6 +197,7 @@ def ScaleUpPods(
       ),
       cloud='Azure' if FLAGS.cloud == 'Azure' else None,
   )
+  resource_names = cluster.ApplyYaml(scaleup_yaml)
 
   assert resource_names
   rollout_name = next(resource_names)

perfkitbenchmarker/vm_util.py

@@ -150,9 +150,7 @@
     None,
     'File path to the SSH private key. If None, use the newly generated one.',
 )
-_SSH_KEY_TYPE = flags.DEFINE_string(
-    'ssh_key_type', 'rsa',
-    'SSH key type.')
+_SSH_KEY_TYPE = flags.DEFINE_string('ssh_key_type', 'rsa', 'SSH key type.')
 
 
 class RetryError(Exception):
@@ -950,19 +948,13 @@ def RenderTemplate(
   Returns:
     The name of the temporary file containing the rendered template.
   """
-  with open(template_path) as fp:
-    template_contents = fp.read()
-  environment = jinja2.Environment(
-      undefined=jinja2.StrictUndefined,
-      trim_blocks=trim_spaces,
-      lstrip_blocks=trim_spaces,
+  rendered_template = ReadAndRenderJinja2Template(
+      template_path, trim_spaces, **context
   )
-  template = environment.from_string(template_contents)
   prefix = 'pkb-' + os.path.basename(template_path)
   with NamedTemporaryFile(
       prefix=prefix, dir=GetTempDir(), delete=False, mode='w'
   ) as tf:
-    rendered_template = template.render(**context)
     if should_log_file:
       logging.info(
           'Rendered template from %s to %s with full text:\n%s',
@@ -976,6 +968,24 @@ def RenderTemplate(
     return tf.name
 
 
+@staticmethod
+def ReadAndRenderJinja2Template(
+    file_path: str, trim_spaces: bool = False, **kwargs
+) -> str:
+  """Reads & renders a .j2 file, returning the whole thing as a string."""
+  filename = data.ResourcePath(file_path)
+  with open(filename) as template_file:
+    contents = template_file.read()
+  if file_path.endswith('.j2'):
+    environment = jinja2.Environment(
+        undefined=jinja2.StrictUndefined,
+        trim_blocks=trim_spaces,
+        lstrip_blocks=trim_spaces,
+    )
+    contents = environment.from_string(contents).render(kwargs)
+  return contents
+
+
 def CreateRemoteFile(
     vm, file_contents, file_path, write_mode: Literal['w', 'wb'] = 'w'
 ):

tests/container_service_test.py

@@ -1,3 +1,4 @@
+import os
 import time
 from typing import Callable, Iterable, Protocol, Tuple
 import unittest
@@ -136,7 +137,7 @@ def setUp(self):
         )
     )
 
-  @parameterized.parameters(('created'), ('configured'))
+  @parameterized.parameters('created', 'configured')
   def test_apply_manifest_gets_deployment_name(self, suffix):
     self.MockIssueCommand(
         {'apply -f': [(f'deployment.apps/test-deployment {suffix}', '', 0)]}
@@ -153,6 +154,68 @@ def test_apply_manifest_gets_deployment_name(self, suffix):
     )
     self.assertEqual(next(deploy_ids), 'deployment.apps/test-deployment')
 
+  def test_apply_manifest_logs_jinja(self):
+    self.MockIssueCommand(
+        {'apply -f': [('deployment.apps/test-deployment hello', '', 0)]}
+    )
+    self.enter_context(
+        mock.patch.object(
+            container_service.data,
+            'ResourcePath',
+            return_value=os.path.join(
+                os.path.dirname(__file__), 'data', 'kube_apply.yaml.j2'
+            ),
+        )
+    )
+    with self.assertLogs(level='INFO') as logs:
+      self.kubernetes_cluster.ApplyManifest(
+          'tests/data/kube_apply.yaml.j2',
+          should_log_file=True,
+          name='hello-world',
+          command=['echo', 'hello', 'world'],
+      )
+    # Asserting on logging isn't very important, but is easier than reading the
+    # written file.
+    full_logs = ';'.join(logs.output)
+    self.assertIn('name: hello-world', full_logs)
+    self.assertIn('echo', full_logs)
+
+  def test_apply_manifest_yaml_logs(self):
+    self.MockIssueCommand(
+        {'apply -f': [('deployment.apps/test-deployment hello', '', 0)]}
+    )
+    self.enter_context(
+        mock.patch.object(
+            container_service.data,
+            'ResourcePath',
+            return_value=os.path.join(
+                os.path.dirname(__file__), 'data', 'kube_apply.yaml.j2'
+            ),
+        )
+    )
+    with self.assertLogs(level='INFO') as logs:
+      yamls = self.kubernetes_cluster.ConvertManifestToYamlDicts(
+          'tests/data/kube_apply.yaml.j2',
+          name='hello-world',
+          command=[],
+      )
+      self.assertEqual(yamls[0]['kind'], 'Namespace')
+      yamls[1]['spec']['selector']['app'] = 'hi-world'
+      yamls[1]['spec']['template']['spec']['containers'].append(
+          {'name': 'second-container'}
+      )
+      self.kubernetes_cluster.ApplyYaml(
+          yamls,
+          should_log_file=True,
+      )
+    full_logs = ';'.join(logs.output)
+    self.assertIn('app: hi-world', full_logs)
+    self.assertIn('name: hello-world', full_logs)
+    self.assertIn('name: second-container', full_logs)
+    # Check for no python artifacts.
+    self.assertNotIn('dict', full_logs)
+    self.assertNotIn('null', full_logs)
+
   @mock.patch.object(
       vm_util,
       'IssueCommand',
@@ -450,9 +513,7 @@ def test_KubernetesEventParsing(self):
         },
         'kind': 'Event',
         'lastTimestamp': None,
-        'message': (
-            'Successfully assigned default/deploy-pod to gke-node'
-        ),
+        'message': 'Successfully assigned default/deploy-pod to gke-node',
         'metadata': {
             'creationTimestamp': '2025-10-03T18:05:56Z',
         },
@@ -462,8 +523,7 @@ def test_KubernetesEventParsing(self):
     })
     self.assertIsNotNone(event)
     self.assertEqual(
-        event.message,
-        'Successfully assigned default/deploy-pod to gke-node'
+        event.message, 'Successfully assigned default/deploy-pod to gke-node'
     )
     self.assertEqual(event.reason, 'Scheduled')
     self.assertEqual(event.type, 'Normal')

tests/data/kube_apply.yaml.j2

@@ -0,0 +1,18 @@
+apiVersion: v1
+kind: Namespace
+metadata:
+  name: {{name}}
+---
+apiVersion: v1
+kind: Deployment
+spec:
+  replicas: 1
+  template:
+    spec:
+      containers:
+      - name: {{name}}
+        image: registry.k8s.io/e2e-test-images/agnhost:2.52
+        command:
+          {%- for c in command %}
+          {{ c }}
+          {%- endfor %}