Move yaml code from container_service -> vm_util

PiperOrigin-RevId: 846416972

Author: hubatish

perfkitbenchmarker/container_service/kubernetes.py

@@ -14,7 +14,6 @@
 """Contains classes related to managed kubernetes container services."""
 
 import calendar
-import collections
 import dataclasses
 import datetime
 import functools
@@ -487,18 +486,6 @@ def _WriteAndApplyManifest(
       )
       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
@@ -515,87 +502,24 @@ def ConvertManifestToYamlDicts(
     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
-
-    def _ConvertPossiblyEmptyValue(value: Any) -> Any | None:
-      """Converts a given value to the dictionary type, or None if empty."""
-      if not bool(value) and value != 0:
-        return None
-      converted_value = KubernetesClusterCommands._ConvertToDictType(
-          value, dict_lambda
-      )
-      if not bool(converted_value) and converted_value != 0:
-        return None
-      return converted_value
-
-    yaml_list = []
-    if isinstance(yaml_doc, list):
-      for item in yaml_doc:
-        converted_value = _ConvertPossiblyEmptyValue(item)
-        if converted_value is None:
-          continue
-        yaml_list.append(converted_value)
-      return yaml_list
-    yaml_dict = dict_lambda()
-    for key, value in yaml_doc.items():
-      converted_value = _ConvertPossiblyEmptyValue(value)
-      if converted_value is None:
-        continue
-      yaml_dict[key] = converted_value
-    return yaml_dict
+    return vm_util.ReadYamlAsDicts(manifest)
 
   @staticmethod
   def ApplyYaml(
-      yaml_dicts: list[dict[str, Any]], should_log_file: bool = True
+      yaml_dicts: list[dict[str, Any]], **logging_kwargs
   ) -> 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.
+      **logging_kwargs: Keyword arguments passed to file writing.
 
     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
-    )
+    vm_util.IncrementStackLevel(**logging_kwargs)
+    yaml_file = vm_util.WriteYaml(yaml_dicts, **logging_kwargs)
+    return KubernetesClusterCommands._ApplyRenderedManifest(yaml_file)
 
   @staticmethod
   def WaitForResource(

perfkitbenchmarker/os_mixin.py

@@ -570,7 +570,7 @@ def PushDataFile(self, data_file, remote_path='', should_double_copy=None):
       self.PushFile(file_path, remote_path)
 
   def RenderTemplate(
-      self, template_path, remote_path, context, should_log_file: bool = False
+      self, template_path, remote_path, context, **logging_kwargs
   ):
     """Renders a local Jinja2 template and copies it to the remote host.
 
@@ -581,16 +581,19 @@ def RenderTemplate(
       template_path: string. Local path to jinja2 template.
       remote_path: string. Remote path for rendered file on the remote vm.
       context: dict. Variables to pass to the Jinja2 template during rendering.
-      should_log_file: bool. Whether to log the file after rendering.
+      **logging_kwargs: dict. Keyword arguments passed to WriteTemporaryFile.
 
     Raises:
       jinja2.UndefinedError: if template contains variables not present in
         'context'.
       RemoteCommandError: If there was a problem copying the file.
     """
     context = {'vm': self} | context
+    vm_util.IncrementStackLevel(**logging_kwargs)
+    if 'should_log_file' not in logging_kwargs:
+      logging_kwargs['should_log_file'] = False
     rendered_template = vm_util.RenderTemplate(
-        template_path, context, should_log_file
+        template_path, context, **logging_kwargs
     )
     self.RemoteCopy(rendered_template, remote_path)
 

perfkitbenchmarker/vm_util.py

@@ -15,6 +15,7 @@
 """Set of utility functions for working with virtual machines."""
 
 
+import collections
 import contextlib
 import enum
 import logging
@@ -36,6 +37,7 @@
 from perfkitbenchmarker import errors
 from perfkitbenchmarker import log_util
 from perfkitbenchmarker import temp_dir
+import yaml
 
 FLAGS = flags.FLAGS
 # Using logger rather than logging.info to avoid stack_level problems.
@@ -925,11 +927,48 @@ def DictionaryToEnvString(dictionary, joiner=' '):
   )
 
 
+def WriteTemporaryFile(
+    file_contents: str,
+    origin: str = '',
+    should_log_file: bool = True,
+    stack_level: int = 1,
+) -> str:
+  """Writes the string file_contents to a temporary file.
+
+  Args:
+    file_contents: string. The contents of the file to write.
+    origin: str. The origin of the file, used for logging.
+    should_log_file: bool. Whether to log the file before writing.
+    stack_level: int. The stack level to use for logging.
+
+  Returns:
+    The name of the temporary file.
+  """
+  stack_level += 1
+  prefix = 'pkb-'
+  with NamedTemporaryFile(
+      prefix=prefix, dir=GetTempDir(), delete=False, mode='w'
+  ) as tf:
+    if origin:
+      origin = f'from {origin}'
+    if should_log_file:
+      logging.info(
+          'Writing temporary file %s to name %s with contents:\n%s',
+          origin,
+          tf.name,
+          file_contents,
+          stacklevel=stack_level,
+      )
+    tf.write(file_contents)
+    tf.close()
+    return tf.name
+
+
 def RenderTemplate(
     template_path,
     context,
-    should_log_file: bool = False,
     trim_spaces: bool = False,
+    **logging_kwargs: dict[str, Any],
 ) -> str:
   """Renders a local Jinja2 template and returns its file name.
 
@@ -938,8 +977,8 @@ def RenderTemplate(
   Args:
     template_path: string. Local path to jinja2 template.
     context: dict. Variables to pass to the Jinja2 template during rendering.
-    should_log_file: bool. Whether to log the file after rendering.
     trim_spaces: bool. Value for both trim_blocks and lstrip_blocks.
+    **logging_kwargs: dict. Keyword arguments passed to WriteTemporaryFile.
 
   Raises:
     jinja2.UndefinedError: if template contains variables not present in
@@ -951,24 +990,14 @@ def RenderTemplate(
   rendered_template = ReadAndRenderJinja2Template(
       template_path, trim_spaces, **context
   )
-  prefix = 'pkb-' + os.path.basename(template_path)
-  with NamedTemporaryFile(
-      prefix=prefix, dir=GetTempDir(), delete=False, mode='w'
-  ) as tf:
-    if should_log_file:
-      logging.info(
-          'Rendered template from %s to %s with full text:\n%s',
-          template_path,
-          tf.name,
-          rendered_template,
-          stacklevel=2,
-      )
-    tf.write(rendered_template)
-    tf.close()
-    return tf.name
+  IncrementStackLevel(**logging_kwargs)
+  return WriteTemporaryFile(
+      rendered_template,
+      origin=f'jinja2 template {template_path}',
+      **logging_kwargs,
+  )
 
 
-@staticmethod
 def ReadAndRenderJinja2Template(
     file_path: str, trim_spaces: bool = False, **kwargs
 ) -> str:
@@ -1003,3 +1032,100 @@ def ReadLocalFile(filename: str) -> str:
   file_path = posixpath.join(GetTempDir(), filename)
   stdout, _, _ = IssueCommand(['cat', file_path])
   return stdout
+
+
+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(RecursiveDict)
+
+
+def ReadYamlAsDicts(
+    file_contents: str
+) -> list[dict[str, Any]]:
+  """Reads file contents & converts it to a list of recursive YAML doc dicts.
+
+  Args:
+    file_contents: The yaml string.
+
+  Returns:
+    The various YAML documents as a list of recursive dictionaries.
+  """
+  yaml_docs = yaml.safe_load_all(file_contents)
+  return_yamls = []
+  for yaml_doc in yaml_docs:
+    return_yamls.append(ConvertToDictType(yaml_doc, RecursiveDict))
+  return return_yamls
+
+
+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
+
+  def _ConvertPossiblyEmptyValue(value: Any) -> Any | None:
+    """Converts a given value to the dictionary type, or None if empty."""
+    if not bool(value) and value != 0:
+      return None
+    converted_value = ConvertToDictType(value, dict_lambda)
+    if not bool(converted_value) and converted_value != 0:
+      return None
+    return converted_value
+
+  yaml_list = []
+  if isinstance(yaml_doc, list):
+    for item in yaml_doc:
+      converted_value = _ConvertPossiblyEmptyValue(item)
+      if converted_value is None:
+        continue
+      yaml_list.append(converted_value)
+    return yaml_list
+  yaml_dict = dict_lambda()
+  for key, value in yaml_doc.items():
+    converted_value = _ConvertPossiblyEmptyValue(value)
+    if converted_value is None:
+      continue
+    yaml_dict[key] = converted_value
+  return yaml_dict
+
+
+def WriteYaml(
+    yaml_dicts: list[dict[str, Any]], **kwargs
+) -> str:
+  """Writes yaml to a file & returns the name of that file.
+
+  Args:
+    yaml_dicts: A list of YAML documents.
+    **kwargs: Keyword arguments passed to file writing.
+
+  Returns:
+    Names of the written file.
+  """
+  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(ConvertToDictType(yaml_dict, dict))
+  manifest = yaml.dump_all(normal_dicts)
+  IncrementStackLevel(**kwargs)
+  return WriteTemporaryFile(
+      manifest, origin='yaml', **kwargs
+  )

tests/container_service_test.py

@@ -1,4 +1,5 @@
 import os
+import tempfile
 import time
 from typing import Callable, Iterable, Protocol, Tuple
 import unittest
@@ -163,6 +164,13 @@ def test_apply_manifest_yaml_logs(self):
             ),
         )
     )
+    self.enter_context(
+        mock.patch.object(
+            vm_util,
+            'GetTempDir',
+            return_value=tempfile.gettempdir(),
+        )
+    )
     with self.assertLogs(level='INFO') as logs:
       yamls = self.kubernetes_cluster.ConvertManifestToYamlDicts(
           'tests/data/kube_apply.yaml.j2',