feature: install FrameworkProcessor requirements with uv for faster job startup

sagemaker-core/src/sagemaker/core/utils/install_requirements.py

@@ -22,7 +22,7 @@
 
     - ``configure_pip()`` — returns an authenticated pip index URL (or ``None``).
       Use when you need to build your own pip command with custom flags.
-    - ``install_requirements(path)`` — configures pip and runs ``pip install -r``.
+    - ``install_requirements(path)`` — configures pip and installs with ``uv``.
       Use when you just want requirements installed.
 
     ::
@@ -38,6 +38,7 @@
 import logging
 import os
 import re
+import shutil
 import subprocess
 import sys
 
@@ -111,23 +112,44 @@ def _login_awscli(region, account, domain, repo):
     )
 
 
-def configure_pip(auth_method=CodeArtifactAuthMethod.AUTO):
+def _set_pip_index(python_executable, index_url):
+    """Persist an authenticated index URL into pip config.
+
+    Mirrors what ``aws codeartifact login`` does for the CLI path: writes
+    ``global.index-url`` so the container's pip — including the ``uv`` bootstrap in
+    :func:`_ensure_uv` — pulls from CodeArtifact. This matters in isolated
+    environments with no public PyPI access. (``uv`` ignores ``pip.conf``, so the
+    index is still surfaced to uv separately via ``UV_INDEX_URL``.)
+    """
+    subprocess.check_call(
+        [python_executable, "-m", "pip", "config", "set", "global.index-url", index_url]
+    )
+
+
+def configure_pip(auth_method=CodeArtifactAuthMethod.AUTO, python_executable=None):
     """Configure pip for CodeArtifact if ``CA_REPOSITORY_ARN`` is set.
 
+    Both auth paths persist the authenticated index into pip config, so any later
+    pip invocation — including the ``uv`` bootstrap in :func:`_ensure_uv` — uses
+    CodeArtifact even in isolated environments without public PyPI access.
+
     Args:
         auth_method: Authentication mechanism to use. Defaults to ``CodeArtifactAuthMethod.AUTO``
             (try boto3 first, fall back to AWS CLI).
+        python_executable: Python executable whose pip config is written on the
+            boto3 path. Defaults to ``sys.executable``.
 
     Returns:
-        An authenticated pip index URL (str) when boto3 succeeds,
-        ``None`` when AWS CLI was used (pip config modified globally),
+        An authenticated pip index URL (str) when boto3 succeeds (also written to
+        pip config), ``None`` when AWS CLI was used (pip config modified globally),
         or ``None`` when ``CA_REPOSITORY_ARN`` is not set.
 
     Raises:
         SystemExit: When ``CA_REPOSITORY_ARN`` is set but the requested
             auth method is not available.
         ValueError: When the ARN format is invalid.
     """
+    python_executable = python_executable or sys.executable
     arn = os.environ.get(CA_REPOSITORY_ARN_ENV)
     if not arn:
         return None
@@ -144,12 +166,15 @@ def configure_pip(auth_method=CodeArtifactAuthMethod.AUTO):
 
     if auth_method in (CodeArtifactAuthMethod.BOTO3, CodeArtifactAuthMethod.AUTO):
         try:
-            return _get_index_boto3(region, account, domain, repo)
+            index = _get_index_boto3(region, account, domain, repo)
         except ImportError:
             if auth_method == CodeArtifactAuthMethod.BOTO3:
                 logger.error("boto3 is not available")
                 sys.exit(1)
             logger.info("boto3 not available, trying AWS CLI fallback")
+        else:
+            _set_pip_index(python_executable, index)
+            return index
 
     if auth_method in (CodeArtifactAuthMethod.AWS_CLI, CodeArtifactAuthMethod.AUTO):
         try:
@@ -170,23 +195,111 @@ def configure_pip(auth_method=CodeArtifactAuthMethod.AUTO):
     sys.exit(1)
 
 
+def _ensure_uv(python_executable):
+    """Return a path to a ``uv`` executable, bootstrapping it with pip if absent.
+
+    Some containers don't ship ``uv``. When it's missing we install it with the
+    container's pip (which has already been pointed at CodeArtifact, if configured).
+
+    Args:
+        python_executable: Python executable whose pip is used to bootstrap ``uv``.
+
+    Returns:
+        Path to the ``uv`` executable (str).
+    """
+    uv = shutil.which("uv")
+    if uv:
+        return uv
+    logger.info("uv not found; bootstrapping it with pip")
+    subprocess.check_call([python_executable, "-m", "pip", "install", "uv"])
+    return shutil.which("uv") or "uv"
+
+
+def _pip_config_get(python_executable, key):
+    """Read a pip config value (e.g. set by ``aws codeartifact login``).
+
+    ``uv`` ignores ``pip.conf``, so any index configured globally on pip must be
+    read back and propagated to ``uv`` explicitly. Returns ``None`` when the key
+    is unset or pip can't report it.
+
+    Args:
+        python_executable: Python executable whose pip config is queried.
+        key: pip config key, e.g. ``global.index-url``.
+    """
+    try:
+        out = subprocess.check_output(
+            [python_executable, "-m", "pip", "config", "get", key],
+            stderr=subprocess.DEVNULL,
+        )
+    except (subprocess.CalledProcessError, FileNotFoundError):
+        return None
+    value = out.decode().strip()
+    return value or None
+
+
+def _build_uv_index_env(python_executable, index):
+    """Build the ``UV_*`` environment for index configuration.
+
+    Bridges both CodeArtifact auth paths and any pre-existing pip config into the
+    environment variables ``uv`` understands:
+
+    - ``index`` (from boto3 auth) becomes ``UV_INDEX_URL``.
+    - When boto3 didn't yield an index (CLI login wrote ``pip.conf``, or the user
+      pre-configured pip), read ``global.index-url`` back from pip config.
+    - ``global.extra-index-url`` / ``global.trusted-host`` are always propagated
+      when present, to stay general across private-index setups.
+
+    Args:
+        python_executable: Python executable whose pip config is consulted.
+        index: Authenticated index URL from :func:`configure_pip`, or ``None``.
+
+    Returns:
+        A dict of ``UV_*`` overrides to merge into the subprocess environment.
+    """
+    env = {}
+
+    index_url = index or _pip_config_get(python_executable, "global.index-url")
+    if index_url:
+        env["UV_INDEX_URL"] = index_url
+
+    extra_index_url = _pip_config_get(python_executable, "global.extra-index-url")
+    if extra_index_url:
+        env["UV_EXTRA_INDEX_URL"] = extra_index_url
+
+    trusted_host = _pip_config_get(python_executable, "global.trusted-host")
+    if trusted_host:
+        env["UV_INSECURE_HOST"] = trusted_host
+
+    return env
+
+
 def install_requirements(
     requirements_file="requirements.txt", python_executable=None, auth_method=CodeArtifactAuthMethod.AUTO
 ):
-    """Install pip requirements with optional CodeArtifact authentication.
+    """Install requirements with ``uv`` and optional CodeArtifact authentication.
+
+    Configures CodeArtifact (if ``CA_REPOSITORY_ARN`` is set), bootstraps ``uv``
+    when the container doesn't ship it, propagates any private-index configuration
+    into ``uv``'s environment, and installs the requirements into the system
+    interpreter.
 
     Args:
         requirements_file: Path to the requirements file.
-        python_executable: Python executable to use for pip. Defaults to ``sys.executable``.
+        python_executable: Python executable used to bootstrap ``uv`` and read pip
+            config. Defaults to ``sys.executable``.
         auth_method: Authentication mechanism for CodeArtifact. Defaults to ``CodeArtifactAuthMethod.AUTO``.
     """
     python_executable = python_executable or sys.executable
-    pip_cmd = [python_executable, "-m", "pip", "install", "-r", requirements_file]
-    index = configure_pip(auth_method=auth_method)
-    if index:
-        pip_cmd.extend(["-i", index])
-    logger.info("Running: %s", " ".join(pip_cmd))
-    subprocess.check_call(pip_cmd)
+
+    index = configure_pip(auth_method=auth_method, python_executable=python_executable)
+
+    uv = _ensure_uv(python_executable)
+    env = os.environ.copy()
+    env.update(_build_uv_index_env(python_executable, index))
+
+    install_cmd = [uv, "pip", "install", "--system", "-r", requirements_file]
+    logger.info("Running: %s", " ".join(install_cmd))
+    subprocess.check_call(install_cmd, env=env)
 
 
 def main():

sagemaker-core/tests/unit/test_install_requirements.py

@@ -29,6 +29,8 @@
 
 _MODULE = "sagemaker.core.utils.install_requirements"
 
+UV_PATH = "/usr/bin/uv"
+
 VALID_ARN = "arn:aws:codeartifact:us-west-2:123456789012:repository/my-domain/my-repo"
 PARSED = ("us-west-2", "123456789012", "my-domain", "my-repo")
 FAKE_TOKEN = "fake-auth-token"
@@ -63,8 +65,12 @@ def mock_boto3_ca():
         yield factory, client
 
 
-def _pip_cmd(*extra):
-    return [sys.executable, "-m", "pip", "install", "-r", "reqs.txt", *extra]
+def _uv_cmd(requirements="reqs.txt"):
+    return [UV_PATH, "pip", "install", "--system", "-r", requirements]
+
+
+def _pip_config_set_cmd(index=EXPECTED_INDEX):
+    return [sys.executable, "-m", "pip", "config", "set", "global.index-url", index]
 
 
 # ---------------------------------------------------------------------------
@@ -105,7 +111,8 @@ def test_invalid_arn_raises(self, ca_env):
 
     def test_boto3_success(self, ca_env, mock_boto3_ca):
         factory, client = mock_boto3_ca
-        result = configure_pip()
+        with mock.patch("subprocess.check_call") as mock_call:
+            result = configure_pip()
 
         factory.assert_called_once_with("codeartifact", region_name="us-west-2")
         client.get_authorization_token.assert_called_once_with(
@@ -115,6 +122,9 @@ def test_boto3_success(self, ca_env, mock_boto3_ca):
             domain="my-domain", domainOwner="123456789012", repository="my-repo", format="pypi"
         )
         assert result == EXPECTED_INDEX
+        # boto3 path also persists the index to pip config so the uv bootstrap
+        # (and any other pip call) inherits CodeArtifact in isolated environments.
+        mock_call.assert_called_once_with(_pip_config_set_cmd())
 
     def test_falls_back_to_cli_when_no_boto3(self, ca_env):
         with mock.patch(f"{_MODULE}._get_index_boto3", side_effect=ImportError):
@@ -150,9 +160,16 @@ def test_fails_when_unavailable(self, ca_env):
 
     def test_does_not_try_cli(self, ca_env, mock_boto3_ca):
         with mock.patch(f"{_MODULE}._login_awscli") as mock_cli:
-            configure_pip(auth_method=CodeArtifactAuthMethod.BOTO3)
+            with mock.patch("subprocess.check_call"):
+                configure_pip(auth_method=CodeArtifactAuthMethod.BOTO3)
         mock_cli.assert_not_called()
 
+    def test_writes_pip_config(self, ca_env, mock_boto3_ca):
+        """boto3 path persists the index to pip config for the uv bootstrap."""
+        with mock.patch("subprocess.check_call") as mock_call:
+            configure_pip(auth_method=CodeArtifactAuthMethod.BOTO3)
+        mock_call.assert_called_once_with(_pip_config_set_cmd())
+
 
 # ---------------------------------------------------------------------------
 # configure_pip — AWS_CLI only
@@ -181,44 +198,102 @@ def test_does_not_try_boto3(self, ca_env):
 # ---------------------------------------------------------------------------
 class TestInstallRequirements:
     def test_without_codeartifact(self):
+        """No CA, uv present, no pip index config → bare uv install, clean env."""
         with mock.patch.dict("os.environ", {}, clear=True):
-            with mock.patch("subprocess.check_call") as mock_call:
-                install_requirements("reqs.txt")
-        mock_call.assert_called_once_with(_pip_cmd())
+            with mock.patch(f"{_MODULE}.shutil.which", return_value=UV_PATH):
+                with mock.patch(f"{_MODULE}._pip_config_get", return_value=None):
+                    with mock.patch("subprocess.check_call") as mock_call:
+                        install_requirements("reqs.txt")
+        mock_call.assert_called_once_with(_uv_cmd(), env=mock.ANY)
+        assert "UV_INDEX_URL" not in mock_call.call_args.kwargs["env"]
+
+    def test_with_codeartifact_index_sets_uv_index_url(self):
+        """boto3 index URL is propagated to uv via UV_INDEX_URL."""
+        with mock.patch.dict("os.environ", {}, clear=True):
+            with mock.patch(f"{_MODULE}.configure_pip", return_value=EXPECTED_INDEX):
+                with mock.patch(f"{_MODULE}.shutil.which", return_value=UV_PATH):
+                    with mock.patch(f"{_MODULE}._pip_config_get", return_value=None):
+                        with mock.patch("subprocess.check_call") as mock_call:
+                            install_requirements("reqs.txt")
+        mock_call.assert_called_once_with(_uv_cmd(), env=mock.ANY)
+        assert mock_call.call_args.kwargs["env"]["UV_INDEX_URL"] == EXPECTED_INDEX
 
-    def test_with_codeartifact_index(self):
-        with mock.patch(f"{_MODULE}.configure_pip", return_value=EXPECTED_INDEX):
-            with mock.patch("subprocess.check_call") as mock_call:
-                install_requirements("reqs.txt")
-        mock_call.assert_called_once_with(_pip_cmd("-i", EXPECTED_INDEX))
+    def test_cli_fallback_index_recovered_from_pip_config(self):
+        """CLI login returns no index but writes pip.conf → recovered into UV_INDEX_URL."""
 
-    def test_with_cli_fallback_no_index_flag(self):
-        with mock.patch(f"{_MODULE}.configure_pip", return_value=None):
-            with mock.patch("subprocess.check_call") as mock_call:
-                install_requirements("reqs.txt")
-        mock_call.assert_called_once_with(_pip_cmd())
+        def fake_pip_config(_exe, key):
+            return EXPECTED_INDEX if key == "global.index-url" else None
 
-    def test_custom_python_executable(self):
         with mock.patch.dict("os.environ", {}, clear=True):
-            with mock.patch("subprocess.check_call") as mock_call:
-                install_requirements("reqs.txt", python_executable="/usr/bin/python3")
-        mock_call.assert_called_once_with(
-            ["/usr/bin/python3", "-m", "pip", "install", "-r", "reqs.txt"]
-        )
+            with mock.patch(f"{_MODULE}.configure_pip", return_value=None):
+                with mock.patch(f"{_MODULE}.shutil.which", return_value=UV_PATH):
+                    with mock.patch(f"{_MODULE}._pip_config_get", side_effect=fake_pip_config):
+                        with mock.patch("subprocess.check_call") as mock_call:
+                            install_requirements("reqs.txt")
+        assert mock_call.call_args.kwargs["env"]["UV_INDEX_URL"] == EXPECTED_INDEX
+
+    def test_extra_index_and_trusted_host_propagated(self):
+        """extra-index-url and trusted-host from pip config flow to uv env vars."""
+
+        def fake_pip_config(_exe, key):
+            return {
+                "global.index-url": "https://primary/simple/",
+                "global.extra-index-url": "https://extra/simple/",
+                "global.trusted-host": "extra",
+            }.get(key)
 
-    def test_pip_failure_propagates(self):
+        with mock.patch.dict("os.environ", {}, clear=True):
+            with mock.patch(f"{_MODULE}.configure_pip", return_value=None):
+                with mock.patch(f"{_MODULE}.shutil.which", return_value=UV_PATH):
+                    with mock.patch(f"{_MODULE}._pip_config_get", side_effect=fake_pip_config):
+                        with mock.patch("subprocess.check_call") as mock_call:
+                            install_requirements("reqs.txt")
+        env = mock_call.call_args.kwargs["env"]
+        assert env["UV_INDEX_URL"] == "https://primary/simple/"
+        assert env["UV_EXTRA_INDEX_URL"] == "https://extra/simple/"
+        assert env["UV_INSECURE_HOST"] == "extra"
+
+    def test_bootstraps_uv_when_missing(self):
+        """uv absent → pip install uv, then uv install."""
         with mock.patch.dict("os.environ", {}, clear=True):
             with mock.patch(
-                "subprocess.check_call", side_effect=subprocess.CalledProcessError(1, "pip")
+                f"{_MODULE}.shutil.which", side_effect=[None, UV_PATH]
             ):
-                with pytest.raises(subprocess.CalledProcessError):
-                    install_requirements("reqs.txt")
+                with mock.patch(f"{_MODULE}._pip_config_get", return_value=None):
+                    with mock.patch("subprocess.check_call") as mock_call:
+                        install_requirements("reqs.txt", python_executable="/usr/bin/python3")
+        bootstrap_call = mock.call(["/usr/bin/python3", "-m", "pip", "install", "uv"])
+        install_call = mock.call(_uv_cmd(), env=mock.ANY)
+        mock_call.assert_has_calls([bootstrap_call, install_call])
+
+    def test_custom_python_executable(self):
+        with mock.patch.dict("os.environ", {}, clear=True):
+            with mock.patch(f"{_MODULE}.shutil.which", return_value=UV_PATH):
+                with mock.patch(f"{_MODULE}._pip_config_get", return_value=None):
+                    with mock.patch("subprocess.check_call") as mock_call:
+                        install_requirements("reqs.txt", python_executable="/usr/bin/python3")
+        mock_call.assert_called_once_with(_uv_cmd(), env=mock.ANY)
+
+    def test_install_failure_propagates(self):
+        with mock.patch.dict("os.environ", {}, clear=True):
+            with mock.patch(f"{_MODULE}.shutil.which", return_value=UV_PATH):
+                with mock.patch(f"{_MODULE}._pip_config_get", return_value=None):
+                    with mock.patch(
+                        "subprocess.check_call",
+                        side_effect=subprocess.CalledProcessError(1, "uv"),
+                    ):
+                        with pytest.raises(subprocess.CalledProcessError):
+                            install_requirements("reqs.txt")
 
     def test_auth_method_passed_through(self):
         with mock.patch(f"{_MODULE}.configure_pip", return_value=None) as mock_configure:
-            with mock.patch("subprocess.check_call"):
-                install_requirements("reqs.txt", auth_method=CodeArtifactAuthMethod.BOTO3)
-        mock_configure.assert_called_once_with(auth_method=CodeArtifactAuthMethod.BOTO3)
+            with mock.patch(f"{_MODULE}.shutil.which", return_value=UV_PATH):
+                with mock.patch(f"{_MODULE}._pip_config_get", return_value=None):
+                    with mock.patch("subprocess.check_call"):
+                        install_requirements("reqs.txt", auth_method=CodeArtifactAuthMethod.BOTO3)
+        mock_configure.assert_called_once_with(
+            auth_method=CodeArtifactAuthMethod.BOTO3, python_executable=sys.executable
+        )
 
 
 # ---------------------------------------------------------------------------