[perf] use uv for venv creation and pip install

sky/setup_files/setup.py

@@ -216,9 +216,15 @@ def parse_readme(readme: str) -> str:
     # We need azure-identity>=1.13.0 to enable the customization of the
     # timeout of AzureCliCredential.
     'azure': [
-        'azure-cli>=2.65.0', 'azure-core>=1.31.0', 'azure-identity>=1.19.0',
-        'azure-mgmt-network>=27.0.0', 'azure-mgmt-compute>=33.0.0',
-        'azure-storage-blob>=12.23.1', 'msgraph-sdk'
+        # If you update the azure-cli dependency, update in sky/skylet/constants.py
+        # SKYPILOT_WHEEL_INSTALLATION_COMMANDS as well.
+        'azure-cli>=2.65.0',
+        'azure-core>=1.31.0',
+        'azure-identity>=1.19.0',
+        'azure-mgmt-network>=27.0.0',
+        'azure-mgmt-compute>=33.0.0',
+        'azure-storage-blob>=12.23.1',
+        'msgraph-sdk',
     ] + local_ray,
     # We need google-api-python-client>=2.69.0 to enable 'discardLocalSsd'
     # parameter for stopping instances.

sky/skylet/constants.py

@@ -39,6 +39,7 @@
                            'which python3')
 # Python executable, e.g., /opt/conda/bin/python3
 SKY_PYTHON_CMD = f'$({SKY_GET_PYTHON_PATH_CMD})'
+# Prefer SKY_UV_PIP_CMD, which is faster. TODO(cooper): remove all usages.
 SKY_PIP_CMD = f'{SKY_PYTHON_CMD} -m pip'
 # Ray executable, e.g., /opt/conda/bin/ray
 # We need to add SKY_PYTHON_CMD before ray executable because:
@@ -50,6 +51,14 @@
 SKY_REMOTE_PYTHON_ENV_NAME = 'skypilot-runtime'
 SKY_REMOTE_PYTHON_ENV = f'~/{SKY_REMOTE_PYTHON_ENV_NAME}'
 ACTIVATE_SKY_REMOTE_PYTHON_ENV = f'source {SKY_REMOTE_PYTHON_ENV}/bin/activate'
+# uv is used for venv and pip, much faster than python implementations.
+SKY_UV_INSTALL_DIR = '"$HOME/.local/bin"'
+SKY_UV_CMD = f'{SKY_UV_INSTALL_DIR}/uv'
+# This won't reinstall uv if it's already installed, so it's safe to re-run.
+SKY_UV_INSTALL_CMD = (f'{SKY_UV_CMD} -V >/dev/null 2>&1 || '
+                      'curl -LsSf https://astral.sh/uv/install.sh '
+                      f'| UV_INSTALL_DIR={SKY_UV_INSTALL_DIR} sh')
+SKY_UV_PIP_CMD = f'VIRTUAL_ENV={SKY_REMOTE_PYTHON_ENV} {SKY_UV_CMD} pip'
 # Deleting the SKY_REMOTE_PYTHON_ENV_NAME from the PATH to deactivate the
 # environment. `deactivate` command does not work when conda is used.
 DEACTIVATE_SKY_REMOTE_PYTHON_ENV = (
@@ -148,28 +157,30 @@
     'echo "Creating conda env with Python 3.10" && '
     f'conda create -y -n {SKY_REMOTE_PYTHON_ENV_NAME} python=3.10 && '
     f'conda activate {SKY_REMOTE_PYTHON_ENV_NAME};'
+    # Install uv for venv management and pip installation.
+    f'{SKY_UV_INSTALL_CMD};'
     # Create a separate conda environment for SkyPilot dependencies.
     f'[ -d {SKY_REMOTE_PYTHON_ENV} ] || '
     # Do NOT use --system-site-packages here, because if users upgrade any
     # packages in the base env, they interfere with skypilot dependencies.
     # Reference: https://github.com/skypilot-org/skypilot/issues/4097
-    f'{SKY_PYTHON_CMD} -m venv {SKY_REMOTE_PYTHON_ENV};'
+    # --seed will include pip and setuptools, which are present in venvs created
+    # with python -m venv.
+    f'{SKY_UV_CMD} venv --seed {SKY_REMOTE_PYTHON_ENV};'
     f'echo "$(echo {SKY_REMOTE_PYTHON_ENV})/bin/python" > {SKY_PYTHON_PATH_FILE};'
 )
 
 _sky_version = str(version.parse(sky.__version__))
 RAY_STATUS = f'RAY_ADDRESS=127.0.0.1:{SKY_REMOTE_RAY_PORT} {SKY_RAY_CMD} status'
 RAY_INSTALLATION_COMMANDS = (
+    f'{SKY_UV_INSTALL_CMD};'
     'mkdir -p ~/sky_workdir && mkdir -p ~/.sky/sky_app;'
-    # Disable the pip version check to avoid the warning message, which makes
-    # the output hard to read.
-    'export PIP_DISABLE_PIP_VERSION_CHECK=1;'
     # Print the PATH in provision.log to help debug PATH issues.
     'echo PATH=$PATH; '
     # Install setuptools<=69.5.1 to avoid the issue with the latest setuptools
     # causing the error:
     #   ImportError: cannot import name 'packaging' from 'pkg_resources'"
-    f'{SKY_PIP_CMD} install "setuptools<70"; '
+    f'{SKY_UV_PIP_CMD} install "setuptools<70"; '
     # Backward compatibility for ray upgrade (#3248): do not upgrade ray if the
     # ray cluster is already running, to avoid the ray cluster being restarted.
     #
@@ -183,10 +194,10 @@
     # 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'{SKY_PIP_CMD} list | grep "ray " | '
+    f'{SKY_UV_PIP_CMD} list | grep "ray " | '
     f'grep {SKY_REMOTE_RAY_VERSION} 2>&1 > /dev/null '
     f'|| {RAY_STATUS} || '
-    f'{SKY_PIP_CMD} install --exists-action w -U ray[default]=={SKY_REMOTE_RAY_VERSION}; '  # pylint: disable=line-too-long
+    f'{SKY_UV_PIP_CMD} install -U ray[default]=={SKY_REMOTE_RAY_VERSION}; '  # pylint: disable=line-too-long
     # In some envs, e.g. pip does not have permission to write under /opt/conda
     # ray package will be installed under ~/.local/bin. If the user's PATH does
     # not include ~/.local/bin (the pip install will have the output: `WARNING:
@@ -202,10 +213,21 @@
     f'which ray > {SKY_RAY_PATH_FILE} || exit 1; }}; ')
 
 SKYPILOT_WHEEL_INSTALLATION_COMMANDS = (
-    f'{{ {SKY_PIP_CMD} list | grep "skypilot " && '
+    f'{SKY_UV_INSTALL_CMD};'
+    f'{{ {SKY_UV_PIP_CMD} list | grep "skypilot " && '
     '[ "$(cat ~/.sky/wheels/current_sky_wheel_hash)" == "{sky_wheel_hash}" ]; } || '  # pylint: disable=line-too-long
-    f'{{ {SKY_PIP_CMD} uninstall skypilot -y; '
-    f'{SKY_PIP_CMD} install "$(echo ~/.sky/wheels/{{sky_wheel_hash}}/'
+    f'{{ {SKY_UV_PIP_CMD} uninstall skypilot; '
+    # uv cannot install azure-cli normally, since it depends on pre-release
+    # packages. Manually install azure-cli with the --prerelease=allow flag
+    # first. This will allow skypilot to successfully install. See
+    # https://docs.astral.sh/uv/pip/compatibility/#pre-release-compatibility.
+    # We don't want to use --prerelease=allow for all packages, because it will
+    # cause uv to use pre-releases for some other packages that have sufficient
+    # stable releases.
+    'if [ "{cloud}" = "azure" ]; then '
+    f'{SKY_UV_PIP_CMD} install --prerelease=allow "azure-cli>=2.65.0"; fi;'
+    # Install skypilot from wheel
+    f'{SKY_UV_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; }; ')
@@ -220,7 +242,7 @@
     # 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'{SKY_PIP_CMD} list | grep "ray " | '
+    f'{SKY_UV_PIP_CMD} list | grep "ray " | '
     f'grep {SKY_REMOTE_RAY_VERSION} 2>&1 > /dev/null && '
     f'{{ {SKY_PYTHON_CMD} -c '
     '"from sky.skylet.ray_patches import patch; patch()" || exit 1; }; ')

sky/templates/kubernetes-ray.yml.j2

@@ -414,7 +414,7 @@ available_node_types:
                 done
                 {{ conda_installation_commands }}
                 {{ ray_installation_commands }}
-                ~/skypilot-runtime/bin/python -m pip install skypilot[kubernetes,remote]
+                VIRTUAL_ENV=~/skypilot-runtime ~/.local/bin/uv pip install skypilot[kubernetes,remote]
                 touch /tmp/ray_skypilot_installation_complete
                 echo "=== Ray and skypilot installation completed ==="