chore: import upstream snapshot with attribution

release/ray_release/command_runner/command_runner.py

@@ -0,0 +1,85 @@
+import abc
+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.util import exponential_backoff_retry
+
+
+class CommandRunner(abc.ABC):
+    """This is run on Buildkite runners."""
+
+    def __init__(
+        self,
+        cluster_manager: ClusterManager,
+        working_dir: str,
+        artifact_path: Optional[str] = None,
+    ):
+        self.cluster_manager = cluster_manager
+        self.working_dir = working_dir
+
+    def prepare_remote_env(self):
+        """Prepare remote environment, e.g. upload files."""
+        raise NotImplementedError
+
+    def wait_for_nodes(self, num_nodes: int, timeout: float = 900.0):
+        """Wait for cluster nodes to be up.
+
+        Args:
+            num_nodes: Number of nodes to wait for.
+            timeout: Timeout in seconds to wait for nodes before
+             raising a ``PrepareCommandTimeoutError``.
+
+        Returns:
+            None
+
+        Raises:
+            PrepareCommandTimeoutError
+        """
+        raise NotImplementedError
+
+    def run_command(
+        self,
+        command: str,
+        env: Optional[Dict] = None,
+        timeout: float = 3600.0,
+        raise_on_timeout: bool = True,
+    ) -> float:
+        """Run command."""
+        raise NotImplementedError
+
+    def run_prepare_command(
+        self, command: str, env: Optional[Dict] = None, timeout: float = 3600.0
+    ):
+        """Run prepare command.
+
+        Command runners may choose to run this differently than the
+        test command.
+        """
+        return exponential_backoff_retry(
+            lambda: self.run_command(command, env, timeout),
+            ClickException,
+            initial_retry_delay_s=5,
+            max_retries=3,
+        )
+
+    def get_last_logs(self) -> Optional[str]:
+        try:
+            return self.get_last_logs_ex()
+        except Exception as e:
+            logger.exception(f"Error fetching logs: {e}")
+            return None
+
+    def get_last_logs_ex(self):
+        raise NotImplementedError
+
+    def fetch_results(self) -> Dict[str, Any]:
+        raise NotImplementedError
+
+    def fetch_metrics(self) -> Dict[str, Any]:
+        raise NotImplementedError
+
+    def fetch_artifact(self) -> None:
+        raise NotImplementedError