[Autoscaler] Pass custom resources to "ray start" multi instance autoscaling (#9986)

python/ray/autoscaler/command_runner.py

@@ -1,6 +1,6 @@
 from getpass import getuser
 from shlex import quote
-from typing import List, Tuple
+from typing import List, Tuple, Dict
 import click
 import hashlib
 import logging
@@ -66,6 +66,35 @@ def set_using_login_shells(val):
     _config["use_login_shells"] = val
 
 
+def _with_environment_variables(cmd: str,
+                                environment_variables: Dict[str, object]):
+    """Prepend environment variables to a shell command.
+
+    Args:
+        cmd (str): The base command.
+        environment_variables (Dict[str, object]): The set of environment
+            variables. If an environment variable value is a dict, it will
+            automatically be converted to a one line yaml string.
+    """
+
+    def dict_as_one_line_yaml(d):
+        items = []
+        for key, val in d.items():
+            item_str = "{}: {}".format(quote(str(key)), quote(str(val)))
+            items.append(item_str)
+
+        return "{" + ",".join(items) + "}"
+
+    as_strings = []
+    for key, val in environment_variables.items():
+        if isinstance(val, dict):
+            val = dict_as_one_line_yaml(val)
+        s = "export {}={};".format(key, quote(val))
+        as_strings.append(s)
+    all_vars = "".join(as_strings)
+    return all_vars + cmd
+
+
 def _with_interactive(cmd):
     force_interactive = ("true && source ~/.bashrc && "
                          "export OMP_NUM_THREADS=1 PYTHONWARNINGS=ignore && ")
@@ -85,6 +114,7 @@ class CommandRunnerInterface:
             exit_on_fail: bool = False,
             port_forward: List[Tuple[int, int]] = None,
             with_output: bool = False,
+            environment_variables: Dict[str, object] = None,
             run_env: str = "auto",
             ssh_options_override_ssh_key: str = "",
     ) -> str:
@@ -100,6 +130,8 @@ class CommandRunnerInterface:
             port_forward (list): List of (local, remote) ports to forward, or
                 a single tuple.
             with_output (bool): Whether to return output.
+            environment_variables (Dict[str, str | int | Dict[str, str]):
+                Environment variables that `cmd` should be run with.
             run_env (str): Options: docker/host/auto. Used in
                 DockerCommandRunner to determine the run environment.
             ssh_options_override_ssh_key (str): if provided, overwrites
@@ -147,6 +179,7 @@ class KubernetesCommandRunner(CommandRunnerInterface):
             exit_on_fail=False,
             port_forward=None,
             with_output=False,
+            environment_variables: Dict[str, object] = None,
             run_env="auto",  # Unused argument.
             ssh_options_override_ssh_key="",  # Unused argument.
     ):
@@ -180,6 +213,9 @@ class KubernetesCommandRunner(CommandRunnerInterface):
                 self.node_id,
                 "--",
             ]
+            cmd = _with_interactive(cmd)
+            if environment_variables:
+                cmd = _with_environment_variables(cmd, environment_variables)
             final_cmd += _with_interactive(cmd)
             logger.info(self.log_prefix + "Running {}".format(final_cmd))
             try:
@@ -434,6 +470,7 @@ class SSHCommandRunner(CommandRunnerInterface):
             exit_on_fail=False,
             port_forward=None,
             with_output=False,
+            environment_variables: Dict[str, object] = None,
             run_env="auto",  # Unused argument.
             ssh_options_override_ssh_key="",
     ):
@@ -472,6 +509,8 @@ class SSHCommandRunner(CommandRunnerInterface):
             "{}@{}".format(self.ssh_user, self.ssh_ip)
         ]
         if cmd:
+            if environment_variables:
+                cmd = _with_environment_variables(cmd, environment_variables)
             if is_using_login_shells():
                 final_cmd += _with_interactive(cmd)
             else:
@@ -544,12 +583,16 @@ class DockerCommandRunner(SSHCommandRunner):
             exit_on_fail=False,
             port_forward=None,
             with_output=False,
+            environment_variables: Dict[str, object] = None,
             run_env="auto",
             ssh_options_override_ssh_key="",
     ):
         if run_env == "auto":
             run_env = "host" if cmd.find("docker") == 0 else "docker"
 
+        if environment_variables:
+            cmd = _with_environment_variables(cmd, environment_variables)
+
         if run_env == "docker":
             cmd = self._docker_expand_user(cmd, any_char=True)
             cmd = " ".join(_with_interactive(cmd))