[release test] move command_env getter into AnyscaleJobRunner (#60272)

so that it is not going back and forth between the implementation and
the abstract class, and not implemented as a property.

Signed-off-by: Lonnie Liu <95255098+aslonnie@users.noreply.github.com>

Author: aslonnie

release/ray_release/command_runner/anyscale_job_runner.py

@@ -29,6 +29,7 @@
 from ray_release.file_manager.job_file_manager import JobFileManager
 from ray_release.job_manager.anyscale_job_manager import AnyscaleJobManager
 from ray_release.logger import logger
+from ray_release.reporter.artifacts import DEFAULT_ARTIFACTS_DIR
 from ray_release.util import (
     AZURE_CLOUD_STORAGE,
     AZURE_STORAGE_CONTAINER,
@@ -63,6 +64,24 @@ def _get_env_str(env: Dict[str, str]) -> str:
 
 
 class AnyscaleJobRunner(CommandRunner):
+    # the directory for runners to dump files to (on buildkite runner instances).
+    # Write to this directory. run_release_tests.sh will ensure that the content
+    # shows up under buildkite job's "Artifacts" UI tab.
+    _DEFAULT_ARTIFACTS_DIR = DEFAULT_ARTIFACTS_DIR
+
+    # the artifact file name put under s3 bucket root.
+    # AnyscalejobWrapper will upload user generated artifact to this path
+    # and AnyscaleJobRunner will then download from there.
+    _USER_GENERATED_ARTIFACT = "user_generated_artifact"
+
+    # the path where result json will be written to on both head node
+    # as well as the relative path where result json will be uploaded to on s3.
+    _RESULT_OUTPUT_JSON = "/tmp/release_test_out.json"
+
+    # the path where output json will be written to on both head node
+    # as well as the relative path where metrics json will be uploaded to on s3.
+    _METRICS_OUTPUT_JSON = "/tmp/metrics_test_out.json"
+
     def __init__(
         self,
         cluster_manager: ClusterManager,
@@ -223,17 +242,22 @@ def _handle_command_output(
                 f"with error:\n{error}\n"
             )
 
-    @property
-    def command_env(self):
-        env = super().command_env
-        # Make sure we don't buffer stdout so we don't lose any logs.
-        env["PYTHONUNBUFFERED"] = "1"
-        return env
+    def _get_full_command_env(self, env: Optional[Dict[str, str]] = None):
+        full_env = {
+            "TEST_OUTPUT_JSON": self._RESULT_OUTPUT_JSON,
+            "METRICS_OUTPUT_JSON": self._METRICS_OUTPUT_JSON,
+            "USER_GENERATED_ARTIFACT": self._USER_GENERATED_ARTIFACT,
+            "BUILDKITE_BRANCH": os.environ.get("BUILDKITE_BRANCH", ""),
+            "PYTHONUNBUFFERED": "1",
+        }
+        if env:
+            full_env.update(env)
+        return full_env
 
     def run_command(
         self,
         command: str,
-        env: Optional[Dict] = None,
+        env: Optional[Dict[str, str]] = None,
         timeout: float = 3600.0,
         raise_on_timeout: bool = True,
     ) -> float:
@@ -242,7 +266,7 @@ def run_command(
         # Convert the prepare commands, envs and timeouts into shell-compliant
         # strings that can be passed to the wrapper script
         for prepare_command, prepare_env, prepare_timeout in self.prepare_commands:
-            prepare_env = self.get_full_command_env(prepare_env)
+            prepare_env = self._get_full_command_env(prepare_env)
             env_str = _get_env_str(prepare_env)
             prepare_command_strs.append(f"{env_str} {prepare_command}")
             prepare_command_timeouts.append(prepare_timeout)
@@ -254,7 +278,7 @@ def run_command(
             shlex.quote(str(x)) for x in prepare_command_timeouts
         )
 
-        full_env = self.get_full_command_env(env)
+        full_env = self._get_full_command_env(env)
 
         no_raise_on_timeout_str = (
             " --test-no-raise-on-timeout" if not raise_on_timeout else ""

release/ray_release/command_runner/command_runner.py

@@ -1,36 +1,16 @@
 import abc
-import os
 from typing import Any, Dict, Optional
 
 from click.exceptions import ClickException
 
 from ray_release.cluster_manager.cluster_manager import ClusterManager
 from ray_release.logger import logger
-from ray_release.reporter.artifacts import DEFAULT_ARTIFACTS_DIR
 from ray_release.util import exponential_backoff_retry
 
 
 class CommandRunner(abc.ABC):
     """This is run on Buildkite runners."""
 
-    # the directory for runners to dump files to (on buildkite runner instances).
-    # Write to this directory. run_release_tests.sh will ensure that the content
-    # shows up under buildkite job's "Artifacts" UI tab.
-    _DEFAULT_ARTIFACTS_DIR = DEFAULT_ARTIFACTS_DIR
-
-    # the artifact file name put under s3 bucket root.
-    # AnyscalejobWrapper will upload user generated artifact to this path
-    # and AnyscaleJobRunner will then download from there.
-    _USER_GENERATED_ARTIFACT = "user_generated_artifact"
-
-    # the path where result json will be written to on both head node
-    # as well as the relative path where result json will be uploaded to on s3.
-    _RESULT_OUTPUT_JSON = "/tmp/release_test_out.json"
-
-    # the path where output json will be written to on both head node
-    # as well as the relative path where metrics json will be uploaded to on s3.
-    _METRICS_OUTPUT_JSON = "/tmp/metrics_test_out.json"
-
     def __init__(
         self,
         cluster_manager: ClusterManager,
@@ -40,23 +20,6 @@ def __init__(
         self.cluster_manager = cluster_manager
         self.working_dir = working_dir
 
-    @property
-    def command_env(self):
-        return {
-            "TEST_OUTPUT_JSON": self._RESULT_OUTPUT_JSON,
-            "METRICS_OUTPUT_JSON": self._METRICS_OUTPUT_JSON,
-            "USER_GENERATED_ARTIFACT": self._USER_GENERATED_ARTIFACT,
-            "BUILDKITE_BRANCH": os.environ.get("BUILDKITE_BRANCH", ""),
-        }
-
-    def get_full_command_env(self, env: Optional[Dict] = None):
-        full_env = self.command_env.copy()
-
-        if env:
-            full_env.update(env)
-
-        return full_env
-
     def prepare_remote_env(self):
         """Prepare remote environment, e.g. upload files."""
         raise NotImplementedError