[Core] Use original Python path for SkyPilot runtime (#3326)

* Use original python path

* Fix python

* cmd

* fix patch

* add test for custom default env

* Fix test

* address comments

* format

* Fix backward compat

* format

* Update sky/backends/backend_utils.py

Co-authored-by: Zongheng Yang <[email protected]>

---------

Co-authored-by: Zongheng Yang <[email protected]>

sky/backends/backend_utils.py

@@ -850,7 +850,7 @@ def write_cluster_config(
 
     # Dump the Ray ports to a file for Ray job submission
     dump_port_command = (
-        f'python -c \'import json, os; json.dump({constants.SKY_REMOTE_RAY_PORT_DICT_STR}, '
+        f'{constants.SKY_PYTHON_CMD} -c \'import json, os; json.dump({constants.SKY_REMOTE_RAY_PORT_DICT_STR}, '
         f'open(os.path.expanduser("{constants.SKY_REMOTE_RAY_PORT_FILE}"), "w", encoding="utf-8"))\''
     )
 
@@ -902,7 +902,9 @@ def write_cluster_config(
                 'ray_dashboard_port': constants.SKY_REMOTE_RAY_DASHBOARD_PORT,
                 'ray_temp_dir': constants.SKY_REMOTE_RAY_TEMPDIR,
                 'dump_port_command': dump_port_command,
-                # Ray version.
+                # Sky-internal constants.
+                'sky_ray_cmd': constants.SKY_RAY_CMD,
+                'sky_pip_cmd': constants.SKY_PIP_CMD,
                 'ray_version': constants.SKY_REMOTE_RAY_VERSION,
                 # Command for waiting ray cluster to be ready on head.
                 'ray_head_wait_initialized_command':

sky/backends/cloud_vm_ray_backend.py

@@ -1933,7 +1933,7 @@ def _ensure_cluster_ray_started(self, handle: 'CloudVmRayResourceHandle',
                 require_outputs=True)
         if returncode == 0:
             return
-        backend.run_on_head(handle, 'ray stop')
+        backend.run_on_head(handle, f'{constants.SKY_RAY_CMD} stop')
 
         # Runs `ray up <kwargs>` with our monkey-patched launch hash
         # calculation. See the monkey patch file for why.
@@ -3095,7 +3095,6 @@ def _exec_code_on_head(
         handle: CloudVmRayResourceHandle,
         codegen: str,
         job_id: int,
-        executable: str,
         detach_run: bool = False,
         spot_dag: Optional['dag.Dag'] = None,
     ) -> None:
@@ -3122,16 +3121,16 @@ def _exec_code_on_head(
         remote_log_dir = self.log_dir
         remote_log_path = os.path.join(remote_log_dir, 'run.log')
 
-        assert executable == 'python3', executable
         cd = f'cd {SKY_REMOTE_WORKDIR}'
 
         job_submit_cmd = (
-            'RAY_DASHBOARD_PORT=$(python -c "from sky.skylet import job_lib; print(job_lib.get_job_submission_port())" 2> /dev/null || echo 8265);'  # pylint: disable=line-too-long
-            f'{cd} && ray job submit '
+            f'RAY_DASHBOARD_PORT=$({constants.SKY_PYTHON_CMD} -c "from sky.skylet import job_lib; print(job_lib.get_job_submission_port())" 2> /dev/null || echo 8265);'  # pylint: disable=line-too-long
+            f'{cd} && {constants.SKY_RAY_CMD} job submit '
             '--address=http://127.0.0.1:$RAY_DASHBOARD_PORT '
             f'--submission-id {job_id}-$(whoami) --no-wait '
             # Redirect stderr to /dev/null to avoid distracting error from ray.
-            f'"{executable} -u {script_path} > {remote_log_path} 2> /dev/null"')
+            f'"{constants.SKY_PYTHON_CMD} -u {script_path} > {remote_log_path} 2> /dev/null"'
+        )
 
         mkdir_code = (f'{cd} && mkdir -p {remote_log_dir} && '
                       f'touch {remote_log_path}')
@@ -3686,7 +3685,8 @@ def teardown_no_lock(self,
                 # We do not check the return code, since Ray returns
                 # non-zero return code when calling Ray stop,
                 # even when the command was executed successfully.
-                self.run_on_head(handle, 'ray stop --force')
+                self.run_on_head(handle,
+                                 f'{constants.SKY_RAY_CMD} stop --force')
             except exceptions.FetchIPError:
                 # This error is expected if the previous cluster IP is
                 # failed to be found,
@@ -4582,7 +4582,6 @@ def _execute_task_one_node(self, handle: CloudVmRayResourceHandle,
         self._exec_code_on_head(handle,
                                 codegen.build(),
                                 job_id,
-                                executable='python3',
                                 detach_run=detach_run,
                                 spot_dag=task.spot_dag)
 
@@ -4646,6 +4645,5 @@ def _execute_task_n_nodes(self, handle: CloudVmRayResourceHandle,
         self._exec_code_on_head(handle,
                                 codegen.build(),
                                 job_id,
-                                executable='python3',
                                 detach_run=detach_run,
                                 spot_dag=task.spot_dag)

sky/clouds/aws.py

@@ -16,6 +16,7 @@
 from sky import skypilot_config
 from sky.adaptors import aws
 from sky.clouds import service_catalog
+from sky.skylet import constants
 from sky.utils import common_utils
 from sky.utils import resources_utils
 from sky.utils import rich_utils
@@ -286,7 +287,7 @@ def get_zone_shell_cmd(cls) -> Optional[str]:
         # https://docs.aws.amazon.com/AWSEC2/latest/UserGuide/instance-identity-documents.html  # pylint: disable=line-too-long
         command_str = (
             'curl -s http://169.254.169.254/latest/dynamic/instance-identity/document'  # pylint: disable=line-too-long
-            ' | python3 -u -c "import sys, json; '
+            f' | {constants.SKY_PYTHON_CMD} -u -c "import sys, json; '
             'print(json.load(sys.stdin)[\'availabilityZone\'])"')
         return command_str
 

sky/execution.py

@@ -678,6 +678,7 @@ def spot_launch(
             'dag_name': dag.name,
             'retry_until_up': retry_until_up,
             'remote_user_config_path': remote_user_config_path,
+            'sky_python_cmd': constants.SKY_PYTHON_CMD,
             **controller_utils.shared_controller_vars_to_fill(
                 'spot',
                 remote_user_config_path=remote_user_config_path,

sky/provision/instance_setup.py

@@ -34,21 +34,22 @@
     'do sudo prlimit --nofile=1048576:1048576 --pid=$id || true; done;')
 
 _DUMP_RAY_PORTS = (
-    'python -c \'import json, os; '
+    f'{constants.SKY_PYTHON_CMD} -c \'import json, os; '
     f'json.dump({constants.SKY_REMOTE_RAY_PORT_DICT_STR}, '
     f'open(os.path.expanduser("{constants.SKY_REMOTE_RAY_PORT_FILE}"), "w", '
     'encoding="utf-8"))\';')
 
 _RAY_PORT_COMMAND = (
-    'RAY_PORT=$(python -c "from sky.skylet import job_lib; '
-    'print(job_lib.get_ray_port())" 2> /dev/null || echo 6379);'
-    'python -c "from sky.utils import common_utils; '
+    f'RAY_PORT=$({constants.SKY_PYTHON_CMD} -c '
+    '"from sky.skylet import job_lib; print(job_lib.get_ray_port())" '
+    '2> /dev/null || echo 6379);'
+    f'{constants.SKY_PYTHON_CMD} -c "from sky.utils import common_utils; '
     'print(common_utils.encode_payload({\'ray_port\': $RAY_PORT}))"')
 
 # Command that calls `ray status` with SkyPilot's Ray port set.
 RAY_STATUS_WITH_SKY_RAY_PORT_COMMAND = (
     f'{_RAY_PORT_COMMAND}; '
-    'RAY_ADDRESS=127.0.0.1:$RAY_PORT ray status')
+    f'RAY_ADDRESS=127.0.0.1:$RAY_PORT {constants.SKY_RAY_CMD} status')
 
 # Command that waits for the ray status to be initialized. Otherwise, a later
 # `sky status -r` may fail due to the ray cluster not being ready.
@@ -59,7 +60,8 @@
     'done;')
 
 # Restart skylet when the version does not match to keep the skylet up-to-date.
-MAYBE_SKYLET_RESTART_CMD = 'python3 -m sky.skylet.attempt_skylet;'
+MAYBE_SKYLET_RESTART_CMD = (f'{constants.SKY_PYTHON_CMD} -m '
+                            'sky.skylet.attempt_skylet;')
 
 
 def _auto_retry(func):
@@ -288,10 +290,11 @@ def start_ray_on_head_node(cluster_name: str, custom_resource: Optional[str],
     # the same credentials. Otherwise, `ray status` will fail to fetch the
     # available nodes.
     # Reference: https://github.com/skypilot-org/skypilot/issues/2441
-    cmd = ('ray stop; unset AWS_ACCESS_KEY_ID AWS_SECRET_ACCESS_KEY; '
+    cmd = (f'{constants.SKY_RAY_CMD} stop; '
+           'unset AWS_ACCESS_KEY_ID AWS_SECRET_ACCESS_KEY; '
            'RAY_SCHEDULER_EVENTS=0 RAY_DEDUP_LOGS=0 '
-           f'ray start --head {ray_options} || exit 1;' + _RAY_PRLIMIT +
-           _DUMP_RAY_PORTS + RAY_HEAD_WAIT_INITIALIZED_COMMAND)
+           f'{constants.SKY_RAY_CMD} start --head {ray_options} || exit 1;' +
+           _RAY_PRLIMIT + _DUMP_RAY_PORTS + RAY_HEAD_WAIT_INITIALIZED_COMMAND)
     logger.info(f'Running command on head node: {cmd}')
     # TODO(zhwu): add the output to log files.
     returncode, stdout, stderr = ssh_runner.run(cmd,
@@ -356,10 +359,11 @@ def start_ray_on_worker_nodes(cluster_name: str, no_restart: bool,
 
     # Unset AWS_ACCESS_KEY_ID AWS_SECRET_ACCESS_KEY, see the comment in
     # `start_ray_on_head_node`.
-    cmd = (f'unset AWS_ACCESS_KEY_ID AWS_SECRET_ACCESS_KEY; '
-           'RAY_SCHEDULER_EVENTS=0 RAY_DEDUP_LOGS=0 '
-           f'ray start --disable-usage-stats {ray_options} || exit 1;' +
-           _RAY_PRLIMIT)
+    cmd = (
+        f'unset AWS_ACCESS_KEY_ID AWS_SECRET_ACCESS_KEY; '
+        'RAY_SCHEDULER_EVENTS=0 RAY_DEDUP_LOGS=0 '
+        f'{constants.SKY_RAY_CMD} start --disable-usage-stats {ray_options} || '
+        'exit 1;' + _RAY_PRLIMIT)
     if no_restart:
         # We do not use ray status to check whether ray is running, because
         # on worker node, if the user started their own ray cluster, ray status

sky/serve/serve_utils.py

@@ -26,6 +26,7 @@
 from sky import status_lib
 from sky.serve import constants
 from sky.serve import serve_state
+from sky.skylet import constants as skylet_constants
 from sky.skylet import job_lib
 from sky.utils import common_utils
 from sky.utils import log_utils
@@ -923,7 +924,8 @@ def stream_serve_process_logs(cls, service_name: str,
     def _build(cls, code: List[str]) -> str:
         code = cls._PREFIX + code
         generated_code = '; '.join(code)
-        return f'python3 -u -c {shlex.quote(generated_code)}'
+        return (f'{skylet_constants.SKY_PYTHON_CMD} '
+                f'-u -c {shlex.quote(generated_code)}')
 
     @classmethod
     def update_service(cls, service_name: str, version: int, mode: str) -> str:

sky/skylet/attempt_skylet.py

@@ -13,12 +13,15 @@ def restart_skylet():
     # TODO(zhwu): make the killing graceful, e.g., use a signal to tell
     # skylet to exit, instead of directly killing it.
     subprocess.run(
-        'ps aux | grep "sky.skylet.skylet" | grep "python3 -m"'
+        # We use -m to grep instead of {constants.SKY_PYTHON_CMD} -m to grep
+        # because need to handle the backward compatibility of the old skylet
+        # started before #3326, which does not use the full path to python.
+        'ps aux | grep "sky.skylet.skylet" | grep " -m "'
         '| awk \'{print $2}\' | xargs kill >> ~/.sky/skylet.log 2>&1',
         shell=True,
         check=False)
     subprocess.run(
-        'nohup python3 -m sky.skylet.skylet'
+        f'nohup {constants.SKY_PYTHON_CMD} -m sky.skylet.skylet'
         ' >> ~/.sky/skylet.log 2>&1 &',
         shell=True,
         check=True)
@@ -27,7 +30,7 @@ def restart_skylet():
 
 
 proc = subprocess.run(
-    'ps aux | grep -v "grep" | grep "sky.skylet.skylet" | grep "python3 -m"',
+    'ps aux | grep -v "grep" | grep "sky.skylet.skylet" | grep " -m"',
     shell=True,
     check=False)
 

sky/skylet/autostop_lib.py

@@ -8,6 +8,7 @@
 
 from sky import sky_logging
 from sky.skylet import configs
+from sky.skylet import constants
 from sky.utils import common_utils
 
 logger = sky_logging.init_logger(__name__)
@@ -121,4 +122,4 @@ def is_autostopping(cls) -> str:
     def _build(cls, code: List[str]) -> str:
         code = cls._PREFIX + code
         code = ';'.join(code)
-        return f'python3 -u -c {shlex.quote(code)}'
+        return f'{constants.SKY_PYTHON_CMD} -u -c {shlex.quote(code)}'

sky/skylet/constants.py

@@ -23,6 +23,21 @@
 SKY_REMOTE_RAY_TEMPDIR = '/tmp/ray_skypilot'
 SKY_REMOTE_RAY_VERSION = '2.9.3'
 
+# We store the absolute path of the python executable (/opt/conda/bin/python3)
+# in this file, so that any future internal commands that need to use python
+# can use this path. This is useful for the case where the user has a custom
+# conda environment as a default environment, which is not the same as the one
+# used for installing SkyPilot runtime (ray and skypilot).
+SKY_PYTHON_PATH_FILE = '~/.sky/python_path'
+SKY_GET_PYTHON_PATH_CMD = (f'cat {SKY_PYTHON_PATH_FILE} 2> /dev/null || '
+                           'which python3')
+# Python executable, e.g., /opt/conda/bin/python3
+SKY_PYTHON_CMD = f'$({SKY_GET_PYTHON_PATH_CMD})'
+_SKY_PYTHON_DIR_CMD = f'$(dirname {SKY_PYTHON_CMD})'
+SKY_PIP_CMD = f'{SKY_PYTHON_CMD} -m pip'
+# Ray executable, e.g., /opt/conda/bin/ray
+SKY_RAY_CMD = f'{_SKY_PYTHON_DIR_CMD}/ray'
+
 # The name for the environment variable that stores the unique ID of the
 # current task. This will stay the same across multiple recoveries of the
 # same spot task.
@@ -78,19 +93,20 @@
     'bash Miniconda3-Linux-x86_64.sh -b && '
     'eval "$(~/miniconda3/bin/conda shell.bash hook)" && conda init && '
     'conda config --set auto_activate_base true); '
-    'grep "# >>> conda initialize >>>" ~/.bashrc || conda init;')
+    'grep "# >>> conda initialize >>>" ~/.bashrc || conda init;'
+    '(type -a python | grep -q python3) || '
+    'echo \'alias python=python3\' >> ~/.bashrc;'
+    '(type -a pip | grep -q pip3) || echo \'alias pip=pip3\' >> ~/.bashrc;'
+    'source ~/.bashrc;'
+    f'[ -f {SKY_PYTHON_PATH_FILE} ] || which python3 > {SKY_PYTHON_PATH_FILE};')
 
 _sky_version = str(version.parse(sky.__version__))
-RAY_STATUS = f'RAY_ADDRESS=127.0.0.1:{SKY_REMOTE_RAY_PORT} ray status'
+RAY_STATUS = f'RAY_ADDRESS=127.0.0.1:{SKY_REMOTE_RAY_PORT} {SKY_RAY_CMD} status'
 # Install ray and skypilot on the remote cluster if they are not already
 # installed. {var} will be replaced with the actual value in
 # backend_utils.write_cluster_config.
 RAY_SKYPILOT_INSTALLATION_COMMANDS = (
-    '(type -a python | grep -q python3) || '
-    'echo \'alias python=python3\' >> ~/.bashrc;'
-    '(type -a pip | grep -q pip3) || echo \'alias pip=pip3\' >> ~/.bashrc;'
     'mkdir -p ~/sky_workdir && mkdir -p ~/.sky/sky_app;'
-    'source ~/.bashrc;'
     # Backward compatibility for ray upgrade (#3248): do not upgrade ray if the
     # ray cluster is already running, to avoid the ray cluster being restarted.
     #
@@ -104,14 +120,15 @@
     # latest ray port 6380, but those existing cluster launched before #1790
     # that has ray cluster on the default port 6379 will be upgraded and
     # restarted.
-    f'pip3 list | grep "ray " | grep {SKY_REMOTE_RAY_VERSION} 2>&1 > /dev/null '
+    f'{SKY_PIP_CMD} list | grep "ray " | '
+    f'grep {SKY_REMOTE_RAY_VERSION} 2>&1 > /dev/null '
     f'|| {RAY_STATUS} || '
-    f'pip3 install --exists-action w -U ray[default]=={SKY_REMOTE_RAY_VERSION}; '  # pylint: disable=line-too-long
+    f'{SKY_PIP_CMD} install --exists-action w -U ray[default]=={SKY_REMOTE_RAY_VERSION}; '  # pylint: disable=line-too-long
     # END ray package check and installation
-    '{ pip3 list | grep "skypilot " && '
+    f'{{ {SKY_PIP_CMD} list | grep "skypilot " && '
     '[ "$(cat ~/.sky/wheels/current_sky_wheel_hash)" == "{sky_wheel_hash}" ]; } || '  # pylint: disable=line-too-long
-    '{ pip3 uninstall skypilot -y; '
-    'pip3 install "$(echo ~/.sky/wheels/{sky_wheel_hash}/'
+    f'{{ {SKY_PIP_CMD} uninstall skypilot -y; '
+    f'{SKY_PIP_CMD} install "$(echo ~/.sky/wheels/{{sky_wheel_hash}}/'
     f'skypilot-{_sky_version}*.whl)[{{cloud}}, remote]" && '
     'echo "{sky_wheel_hash}" > ~/.sky/wheels/current_sky_wheel_hash || '
     'exit 1; }; '
@@ -121,8 +138,8 @@
     # The ray installation above can be skipped due to the existing ray cluster
     # for backward compatibility. In this case, we should not patch the ray
     # files.
-    f'pip3 list | grep "ray " | grep {SKY_REMOTE_RAY_VERSION} 2>&1 > /dev/null '
-    '&& { python3 -c "from sky.skylet.ray_patches import patch; patch()" '
+    f'{SKY_PIP_CMD} list | grep "ray " | grep {SKY_REMOTE_RAY_VERSION} 2>&1 > /dev/null '
+    f'&& {{ {SKY_PYTHON_CMD} -c "from sky.skylet.ray_patches import patch; patch()" '
     '|| exit 1; };')
 
 # The name for the environment variable that stores SkyPilot user hash, which

sky/skylet/job_lib.py

@@ -933,4 +933,4 @@ def get_run_timestamp_with_globbing(cls,
     def _build(cls, code: List[str]) -> str:
         code = cls._PREFIX + code
         code = ';'.join(code)
-        return f'python3 -u -c {shlex.quote(code)}'
+        return f'{constants.SKY_PYTHON_CMD} -u -c {shlex.quote(code)}'

sky/skylet/log_lib.py

@@ -184,8 +184,22 @@ def run_with_log(
             daemon_script = os.path.join(
                 os.path.dirname(os.path.abspath(job_lib.__file__)),
                 'subprocess_daemon.py')
+            if not hasattr(constants, 'SKY_GET_PYTHON_PATH_CMD'):
+                # Backward compatibility: for cluster started before #3326, this
+                # constant does not exist. Since we generate the job script
+                # in backends.cloud_vm_ray_backend with inspect, so the
+                # the lates `run_with_log` will be used, but the `constants` is
+                # not updated. We fallback to `python3` in this case.
+                # TODO(zhwu): remove this after 0.7.0.
+                python_path = 'python3'
+            else:
+                python_path = subprocess.check_output(
+                    constants.SKY_GET_PYTHON_PATH_CMD,
+                    shell=True,
+                    stderr=subprocess.DEVNULL,
+                    encoding='utf-8').strip()
             daemon_cmd = [
-                'python3',
+                python_path,
                 daemon_script,
                 '--parent-pid',
                 str(parent_pid),

sky/spot/spot_utils.py

@@ -741,7 +741,7 @@ def set_pending(cls, job_id: int, spot_dag: 'dag_lib.Dag') -> str:
     def _build(cls, code: List[str]) -> str:
         code = cls._PREFIX + code
         generated_code = '; '.join(code)
-        return f'python3 -u -c {shlex.quote(generated_code)}'
+        return f'{constants.SKY_PYTHON_CMD} -u -c {shlex.quote(generated_code)}'
 
 
 def dump_job_table_cache(job_table: str):

sky/templates/azure-ray.yml.j2

@@ -164,14 +164,14 @@ setup_commands:
 #   current num items (num SSH connections): 2
 head_start_ray_commands:
   # NOTE: --disable-usage-stats in `ray start` saves 10 seconds of idle wait.
-  - ray stop; RAY_SCHEDULER_EVENTS=0 RAY_DEDUP_LOGS=0 ray start --disable-usage-stats --head --port={{ray_port}} --dashboard-port={{ray_dashboard_port}} --object-manager-port=8076 --autoscaling-config=~/ray_bootstrap_config.yaml {{"--num-gpus=%s" % num_gpus if num_gpus}} {{"--resources='%s'" % custom_resources if custom_resources}} --temp-dir {{ray_temp_dir}} || exit 1; 
+  - {{ sky_ray_cmd }} stop; RAY_SCHEDULER_EVENTS=0 RAY_DEDUP_LOGS=0 {{ sky_ray_cmd }} start --disable-usage-stats --head --port={{ray_port}} --dashboard-port={{ray_dashboard_port}} --object-manager-port=8076 --autoscaling-config=~/ray_bootstrap_config.yaml {{"--num-gpus=%s" % num_gpus if num_gpus}} {{"--resources='%s'" % custom_resources if custom_resources}} --temp-dir {{ray_temp_dir}} || exit 1; 
     which prlimit && for id in $(pgrep -f raylet/raylet); do sudo prlimit --nofile=1048576:1048576 --pid=$id || true; done; 
     {{dump_port_command}}; 
     {{ray_head_wait_initialized_command}}
 
 {%- if num_nodes > 1 %}
 worker_start_ray_commands:
-  - ray stop; RAY_SCHEDULER_EVENTS=0 RAY_DEDUP_LOGS=0 ray start --disable-usage-stats --address=$RAY_HEAD_IP:{{ray_port}} --object-manager-port=8076 {{"--num-gpus=%s" % num_gpus if num_gpus}} {{"--resources='%s'" % custom_resources if custom_resources}} --temp-dir {{ray_temp_dir}} || exit 1;
+  - {{ sky_ray_cmd }} stop; RAY_SCHEDULER_EVENTS=0 RAY_DEDUP_LOGS=0 {{ sky_ray_cmd }} start --disable-usage-stats --address=$RAY_HEAD_IP:{{ray_port}} --object-manager-port=8076 {{"--num-gpus=%s" % num_gpus if num_gpus}} {{"--resources='%s'" % custom_resources if custom_resources}} --temp-dir {{ray_temp_dir}} || exit 1;
     which prlimit && for id in $(pgrep -f raylet/raylet); do sudo prlimit --nofile=1048576:1048576 --pid=$id || true; done;
 {%- else %}
 worker_start_ray_commands: []

sky/templates/gcp-ray.yml.j2

@@ -191,7 +191,7 @@ setup_commands:
     source ~/.bashrc;
   {%- if tpu_vm %}
     test -f ~/miniconda3/etc/profile.d/conda.sh && source ~/miniconda3/etc/profile.d/conda.sh && conda activate base || true;
-    pip3 install --upgrade google-api-python-client;
+    {{ sky_pip_cmd }} install --upgrade google-api-python-client;
   {%- endif %}
   {%- if tpu_node_name %}
     grep "export TPU_NAME=" ~/.bashrc && echo "TPU_NAME already set" || echo "export TPU_NAME={{tpu_node_name}}" >> ~/.bashrc;