truefoundry 0.4.1__py3-none-any.whl → 0.4.2__py3-none-any.whl

truefoundry/common/auth_service_client.py

@@ -8,7 +8,7 @@ from truefoundry.common.constants import VERSION_PREFIX
 from truefoundry.common.entities import DeviceCode, Token
 from truefoundry.common.exceptions import BadRequestException
 from truefoundry.common.request_utils import request_handling, requests_retry_session
-from truefoundry.common.utils import poll_for_function
+from truefoundry.common.utils import poll_for_function, relogin_error_message
 from truefoundry.logger import logger
 
 
@@ -16,6 +16,7 @@ class AuthServiceClient(ABC):
     def __init__(self, tenant_name: str):
         self._tenant_name = tenant_name
 
+    # TODO (chiragjn): Rename base_url to tfy_host
     @classmethod
     def from_base_url(cls, base_url: str) -> "AuthServiceClient":
         from truefoundry.common.servicefoundry_client import (
@@ -54,7 +55,9 @@ class ServiceFoundryServerAuthServiceClient(AuthServiceClient):
         if not token.refresh_token:
             # TODO: Add a way to propagate error messages without traceback to the output interface side
             raise Exception(
-                f"Unable to resume login session. Please log in again using `tfy login {host_arg_str} --relogin`"
+                relogin_error_message(
+                    "Unable to resume login session.", host=host_arg_str
+                )
             )
         url = f"{self._api_server_url}/{VERSION_PREFIX}/oauth2/token"
         data = {
@@ -70,7 +73,9 @@ class ServiceFoundryServerAuthServiceClient(AuthServiceClient):
             return Token.parse_obj(response_data)
         except BadRequestException as ex:
             raise Exception(
-                f"Unable to resume login session. Please log in again using `tfy login {host_arg_str} --relogin`"
+                relogin_error_message(
+                    "Unable to resume login session.", host=host_arg_str
+                )
             ) from ex
 
     def get_device_code(self) -> DeviceCode:
@@ -127,7 +132,9 @@ class AuthServerServiceClient(AuthServiceClient):
         if not token.refresh_token:
             # TODO: Add a way to propagate error messages without traceback to the output interface side
             raise Exception(
-                f"Unable to resume login session. Please log in again using `tfy login {host_arg_str} --relogin`"
+                relogin_error_message(
+                    "Unable to resume login session.", host=host_arg_str
+                )
             )
         url = f"{self._auth_server_url}/api/{VERSION_PREFIX}/oauth/token/refresh"
         data = {
@@ -141,7 +148,9 @@ class AuthServerServiceClient(AuthServiceClient):
             return Token.parse_obj(response_data)
         except BadRequestException as ex:
             raise Exception(
-                f"Unable to resume login session. Please log in again using `tfy login {host_arg_str} --relogin`"
+                relogin_error_message(
+                    "Unable to resume login session.", host=host_arg_str
+                )
             ) from ex
 
     def get_device_code(self) -> DeviceCode:

truefoundry/common/constants.py

@@ -5,8 +5,9 @@ CREDENTIAL_FILEPATH = TFY_CONFIG_DIR / "credentials.json"
 
 TFY_HOST_ENV_KEY = "TFY_HOST"
 TFY_API_KEY_ENV_KEY = "TFY_API_KEY"
-SERVICEFOUNDRY_SERVER_URL_ENV_KEY = "SERVICEFOUNDRY_SERVER_URL"
+TFY_CLI_LOCAL_DEV_MODE_ENV_KEY = "TFY_CLI_LOCAL_DEV_MODE"
 
 API_SERVER_RELATIVE_PATH = "api/svc"
+MLFOUNDRY_SERVER_RELATIVE_PATH = "api/ml"
 VERSION_PREFIX = "v1"
 SERVICEFOUNDRY_CLIENT_MAX_RETRIES = 2

truefoundry/common/credential_file_manager.py

@@ -10,6 +10,7 @@ from filelock import FileLock, Timeout
 
 from truefoundry.common.constants import CREDENTIAL_FILEPATH
 from truefoundry.common.entities import CredentialsFileContent
+from truefoundry.common.utils import relogin_error_message
 from truefoundry.logger import logger
 
 
@@ -87,9 +88,9 @@ class CredentialsFileManager:
             return CredentialsFileContent.parse_file(self._credentials_file_path)
         except Exception as ex:
             raise Exception(
-                "Error while reading the credentials file "
-                f"{self._credentials_file_path}. Please login again "
-                "using `tfy login --relogin` or `tfy.login(relogin=True)` function"
+                relogin_error_message(
+                    f"Error while reading the credentials file {self._credentials_file_path}"
+                )
             ) from ex
 
     @_ensure_lock_taken

truefoundry/common/credential_provider.py

@@ -6,7 +6,7 @@ from truefoundry.common.auth_service_client import AuthServiceClient
 from truefoundry.common.constants import TFY_API_KEY_ENV_KEY
 from truefoundry.common.credential_file_manager import CredentialsFileManager
 from truefoundry.common.entities import CredentialsFileContent, Token
-from truefoundry.common.utils import resolve_base_url
+from truefoundry.common.utils import resolve_tfy_host
 from truefoundry.logger import logger
 
 TOKEN_REFRESH_LOCK = threading.RLock()
@@ -17,6 +17,7 @@ class CredentialProvider(ABC):
     @abstractmethod
     def token(self) -> Token: ...
 
+    # TODO (chiragjn): Rename base_url to tfy_host
     @property
     @abstractmethod
     def base_url(self) -> str: ...
@@ -34,10 +35,8 @@ class EnvCredentialProvider(CredentialProvider):
             raise Exception(
                 f"Value of {TFY_API_KEY_ENV_KEY} env var should be non-empty string"
             )
-        # TODO: Read host from cred file as well.
-        base_url = resolve_base_url().strip("/")
-        self._host = base_url
-        self._auth_service = AuthServiceClient.from_base_url(base_url=base_url)
+        self._host = resolve_tfy_host()
+        self._auth_service = AuthServiceClient.from_base_url(base_url=self._host)
         self._token: Token = Token(access_token=api_key, refresh_token=None)  # type: ignore[call-arg]
 
     @staticmethod

truefoundry/common/servicefoundry_client.py

@@ -1,7 +1,6 @@
 from __future__ import annotations
 
 import functools
-from urllib.parse import urlparse
 
 import requests
 from packaging import version
@@ -16,7 +15,7 @@ from truefoundry.common.entities import (
 )
 from truefoundry.common.request_utils import request_handling, requests_retry_session
 from truefoundry.common.utils import (
-    append_servicefoundry_path_to_base_url,
+    get_tfy_servers_config,
     timed_lru_cache,
 )
 from truefoundry.logger import logger
@@ -49,10 +48,10 @@ def session_with_retries() -> requests.Session:
 
 
 @timed_lru_cache(seconds=30 * 60)
-def _cached_get_tenant_info(api_server_url: str) -> TenantInfo:
+def _cached_get_tenant_info(api_server_url: str, tenant_host: str) -> TenantInfo:
     res = session_with_retries().get(
         url=f"{api_server_url}/{VERSION_PREFIX}/tenant-id",
-        params={"hostName": urlparse(api_server_url).netloc},
+        params={"hostName": tenant_host},
     )
     res = request_handling(res)
     return TenantInfo.parse_obj(res)
@@ -68,9 +67,12 @@ def _cached_get_python_sdk_config(api_server_url: str) -> PythonSDKConfig:
 
 
 class ServiceFoundryServiceClient:
+    # TODO (chiragjn): Rename base_url to tfy_host
     def __init__(self, base_url: str):
         self._base_url = base_url.strip("/")
-        self._api_server_url = append_servicefoundry_path_to_base_url(self._base_url)
+        tfy_servers_config = get_tfy_servers_config(self._base_url)
+        self._tenant_host = tfy_servers_config.tenant_host
+        self._api_server_url = tfy_servers_config.servicefoundry_server_url
 
     @property
     def base_url(self) -> str:
@@ -78,11 +80,16 @@ class ServiceFoundryServiceClient:
 
     @property
     def tenant_info(self) -> TenantInfo:
-        return _cached_get_tenant_info(self._api_server_url)
+        return _cached_get_tenant_info(
+            self._api_server_url,
+            tenant_host=self._tenant_host,
+        )
 
     @property
     def python_sdk_config(self) -> PythonSDKConfig:
-        return _cached_get_python_sdk_config(self._api_server_url)
+        return _cached_get_python_sdk_config(
+            self._api_server_url,
+        )
 
     @functools.cached_property
     def _min_cli_version_required(self) -> str:

truefoundry/common/utils.py

@@ -3,17 +3,62 @@ import time
 from functools import lru_cache, wraps
 from time import monotonic_ns
 from typing import Callable, Generator, Optional, TypeVar
-from urllib.parse import urljoin, urlsplit
+from urllib.parse import urljoin, urlparse
 
 from truefoundry.common.constants import (
     API_SERVER_RELATIVE_PATH,
-    SERVICEFOUNDRY_SERVER_URL_ENV_KEY,
+    MLFOUNDRY_SERVER_RELATIVE_PATH,
+    TFY_CLI_LOCAL_DEV_MODE_ENV_KEY,
     TFY_HOST_ENV_KEY,
 )
+from truefoundry.pydantic_v1 import BaseSettings
 
 T = TypeVar("T")
 
 
+class _TFYServersConfig(BaseSettings):
+    class Config:
+        env_prefix = "TFY_CLI_LOCAL_"
+        env_file = ".tfy-cli-local.env"
+
+    tenant_host: str
+    servicefoundry_server_url: str
+    mlfoundry_server_url: str
+
+    @classmethod
+    def from_base_url(cls, base_url: str) -> "_TFYServersConfig":
+        base_url = base_url.strip("/")
+        return cls(
+            tenant_host=urlparse(base_url).netloc,
+            servicefoundry_server_url=urljoin(base_url, API_SERVER_RELATIVE_PATH),
+            mlfoundry_server_url=urljoin(base_url, MLFOUNDRY_SERVER_RELATIVE_PATH),
+        )
+
+
+_tfy_servers_config = None
+
+
+def get_tfy_servers_config(base_url: str) -> _TFYServersConfig:
+    global _tfy_servers_config
+    if _tfy_servers_config is None:
+        if os.getenv(TFY_CLI_LOCAL_DEV_MODE_ENV_KEY):
+            _tfy_servers_config = _TFYServersConfig()  # type: ignore[call-arg]
+        else:
+            _tfy_servers_config = _TFYServersConfig.from_base_url(base_url)
+    return _tfy_servers_config
+
+
+def relogin_error_message(message: str, host: str = "HOST") -> str:
+    suffix = ""
+    if host == "HOST":
+        suffix = " where HOST is TrueFoundry platform URL"
+    return (
+        f"{message}\n"
+        f"Please login again using `tfy login --host {host} --relogin` "
+        f"or `truefoundry.login(host={host!r}, relogin=True)` function" + suffix
+    )
+
+
 def timed_lru_cache(
     seconds: int = 300, maxsize: Optional[int] = None
 ) -> Callable[[Callable[..., T]], Callable[..., T]]:
@@ -42,15 +87,19 @@ def poll_for_function(
         time.sleep(poll_after_secs)
 
 
-def resolve_base_url(host: Optional[str] = None) -> str:
-    if not host and not os.getenv(TFY_HOST_ENV_KEY):
+def validate_tfy_host(tfy_host: str) -> None:
+    if not (tfy_host.startswith("https://") or tfy_host.startswith("http://")):
         raise ValueError(
-            f"Either `host` should be provided by --host <value>, or `{TFY_HOST_ENV_KEY}` env must be set"
+            f"Invalid host {tfy_host!r}. It should start with https:// or http://"
         )
-    return host or os.getenv(TFY_HOST_ENV_KEY)
 
 
-def append_servicefoundry_path_to_base_url(base_url: str):
-    if urlsplit(base_url).netloc.startswith("localhost"):
-        return os.getenv(SERVICEFOUNDRY_SERVER_URL_ENV_KEY)
-    return urljoin(base_url, API_SERVER_RELATIVE_PATH)
+def resolve_tfy_host(tfy_host: Optional[str] = None) -> str:
+    if not tfy_host and not os.getenv(TFY_HOST_ENV_KEY):
+        raise ValueError(
+            f"Either `host` should be provided using `--host <value>`, or `{TFY_HOST_ENV_KEY}` env must be set"
+        )
+    tfy_host = tfy_host or os.getenv(TFY_HOST_ENV_KEY)
+    tfy_host = tfy_host.strip("/")
+    validate_tfy_host(tfy_host)
+    return tfy_host

truefoundry/deploy/auto_gen/models.py

@@ -1,6 +1,6 @@
 # generated by datamodel-codegen:
 #   filename:  application.json
-#   timestamp: 2024-09-29T10:28:22+00:00
+#   timestamp: 2024-10-09T13:00:44+00:00
 
 from __future__ import annotations
 
@@ -363,10 +363,12 @@ class HelmRepo(BaseModel):
         description="+label=Helm repository URL\n+sort=1\n+message=Needs to be a valid URL.",
     )
     chart: str = Field(
-        ..., description="+label=Chart name\n+sort=2\n+usage=The helm chart name"
+        ...,
+        description='+label=Chart name\n+sort=2\n+usage=The helm chart name\n+uiType=InputSelect\n+uiProps={"creatable":true, "searchable":true}',
     )
     version: str = Field(
-        ..., description="+label=Version\n+sort=3\n+usage=Helm chart version"
+        ...,
+        description='+label=Version\n+sort=3\n+usage=Helm chart version\n+uiType=InputSelect\n+uiProps={"creatable":true, "searchable":true}',
     )
 
 
@@ -448,10 +450,14 @@ class JobAlert(BaseModel):
     +label=Alert
     """
 
-    notification_channel: str = Field(
+    notification_channel: constr(min_length=1) = Field(
         ...,
         description='+label=Notification Channel\n+usage=Specify the notification channel to send alerts to\n+uiType=IntegrationSelect\n+uiProps={"integrationType":"notification-channel"}\n+sort=660',
     )
+    to_emails: Optional[List[constr(min_length=1)]] = Field(
+        None,
+        description="+label=To Emails\n+usage=List of recipients' email addresses if the notification channel is Email.\n+docs=Specify the emails to send alerts to\n+sort=665",
+    )
     on_start: bool = Field(
         False,
         description="+label=On Start\n+usage=Send an alert when the job starts\n+sort=670",
@@ -1546,7 +1552,7 @@ class SSHServer(BaseWorkbenchInput):
     image: WorkbenchImage
     ssh_public_key: str = Field(
         ...,
-        description="+label: SSH Public Key\n+usage=Add Your SSH Public Key, this will be used to authenticate you to the SSH Server. You can find it using `cat ~/.ssh/id_rsa.pub`\n+sort=4",
+        description="+label: SSH Public Key\n+usage=Add Your SSH Public Key, this will be used to authenticate you to the SSH Server.  \\\nYou can find it using `cat ~/.ssh/id_rsa.pub` in Mac/Linux or `type $home\\.ssh\\id_rsa.pub` in Windows Powershell.  \\\nYou can also generate a new SSH key pair using `ssh-keygen -t rsa` in your local terminal. (same for both Mac/Linux and Windows Powershell)\n+uiType=TextArea\n+sort=4",
     )
 
 

truefoundry/deploy/builder/__init__.py

@@ -1,4 +1,4 @@
-from typing import Dict, List, Optional, Union
+from typing import Any, Dict, List, Optional, Union
 
 from truefoundry.deploy.auto_gen.models import (
     DockerFileBuild,
@@ -24,7 +24,7 @@ class _BuildConfig(BaseModel):
 
 
 def build(
-    build_configuration: Union[BaseModel, Dict],
+    build_configuration: Union[BaseModel, Dict[str, Any]],
     tag: str,
     extra_opts: Optional[List[str]] = None,
 ):

truefoundry/deploy/builder/builders/tfy_notebook_buildpack/__init__.py

@@ -8,16 +8,20 @@ from truefoundry.deploy.builder.builders.tfy_notebook_buildpack.dockerfile_templ
     NotebookImageBuild,
     generate_dockerfile_content,
 )
+from truefoundry.deploy.builder.utils import has_pip_conf_secret
 
 __all__ = ["generate_dockerfile_content", "build"]
 
 
 def _convert_to_dockerfile_build_config(
-    build_configuration: NotebookImageBuild, local_dir: str
+    build_configuration: NotebookImageBuild,
+    local_dir: str,
+    mount_pip_conf_secret: bool = False,
 ) -> DockerFileBuild:
     dockerfile_content = generate_dockerfile_content(
         build_configuration=build_configuration,
         local_dir=local_dir,
+        mount_pip_conf_secret=mount_pip_conf_secret,
     )
     dockerfile_path = os.path.join(local_dir, "Dockerfile")
     with open(dockerfile_path, "w", encoding="utf8") as fp:
@@ -34,10 +38,12 @@ def build(
     build_configuration: NotebookImageBuild,
     extra_opts: Optional[List[str]] = None,
 ):
+    mount_pip_conf_secret = has_pip_conf_secret(extra_opts) if extra_opts else False
     with TemporaryDirectory() as local_dir:
         docker_build_configuration = _convert_to_dockerfile_build_config(
             build_configuration,
             local_dir=local_dir,
+            mount_pip_conf_secret=mount_pip_conf_secret,
         )
         dockerfile.build(
             tag=tag,

truefoundry/deploy/builder/builders/tfy_notebook_buildpack/dockerfile_template.py

@@ -4,6 +4,10 @@ from typing import Literal, Optional
 
 from mako.template import Template
 
+from truefoundry.deploy.builder.constants import (
+    PIP_CONF_BUILDKIT_SECRET_MOUNT,
+    PIP_CONF_SECRET_MOUNT_AS_ENV,
+)
 from truefoundry.pydantic_v1 import BaseModel
 
 
@@ -28,31 +32,40 @@ USER $NB_UID
 
 
 def generate_build_script_docker_commands(
-    build_script: Optional[str], local_dir: str
+    build_script: Optional[str],
+    local_dir: str,
+    mount_pip_conf_secret: bool = False,
 ) -> Optional[str]:
     if not build_script:
         return None
-    build_script_path = None
-    if build_script:
-        # we add build script's hash to the file name to ensure docker cache invalidation
-        script_hash = hashlib.sha256(build_script.encode("utf-8")).hexdigest()
-        build_script_path = os.path.join(local_dir, f"build-script-{script_hash}.sh")
-        with open(build_script_path, "w") as fp:
-            fp.write(build_script)
-        build_script_path = os.path.relpath(build_script_path, local_dir)
+    # we add build script's hash to the file name to ensure docker cache invalidation
+    script_hash = hashlib.sha256(build_script.encode("utf-8")).hexdigest()
+    build_script_path = os.path.join(local_dir, f"build-script-{script_hash}.sh")
+    with open(build_script_path, "w") as fp:
+        fp.write(build_script)
+    build_script_path = os.path.relpath(build_script_path, local_dir)
+    buildkit_secret_mounts = ""
+    environment_variables = ["DEBIAN_FRONTEND=noninteractive"]
+    if mount_pip_conf_secret:
+        buildkit_secret_mounts = PIP_CONF_BUILDKIT_SECRET_MOUNT
+        environment_variables.append(PIP_CONF_SECRET_MOUNT_AS_ENV)
+    all_environment_variables = " ".join(environment_variables)
     run_build_script_command = f"""\
 COPY {build_script_path} /tmp/user-build-script.sh
-RUN mkdir -p /var/log/ && DEBIAN_FRONTEND=noninteractive bash -ex /tmp/user-build-script.sh 2>&1 | tee /var/log/user-build-script-output.log
+RUN {buildkit_secret_mounts} mkdir -p /var/log/ && {all_environment_variables} bash -ex /tmp/user-build-script.sh 2>&1 | tee /var/log/user-build-script-output.log
 """
     return run_build_script_command
 
 
 def generate_dockerfile_content(
-    build_configuration: NotebookImageBuild, local_dir: str
+    build_configuration: NotebookImageBuild,
+    local_dir: str,
+    mount_pip_conf_secret: bool = False,
 ) -> str:
     build_script_docker_commands = generate_build_script_docker_commands(
         build_script=build_configuration.build_script,
         local_dir=local_dir,
+        mount_pip_conf_secret=mount_pip_conf_secret,
     )
 
     template_args = {
@@ -62,5 +75,5 @@ def generate_dockerfile_content(
 
     template = DOCKERFILE_TEMPLATE
 
-    dockerfile_content = template.render(**template_args)
+    dockerfile_content: str = template.render(**template_args)
     return dockerfile_content

truefoundry/deploy/builder/builders/tfy_python_buildpack/__init__.py

@@ -7,6 +7,7 @@ from truefoundry.deploy.builder.builders import dockerfile
 from truefoundry.deploy.builder.builders.tfy_python_buildpack.dockerfile_template import (
     generate_dockerfile_content,
 )
+from truefoundry.deploy.builder.utils import has_pip_conf_secret
 
 __all__ = ["generate_dockerfile_content", "build"]
 
@@ -14,9 +15,11 @@ __all__ = ["generate_dockerfile_content", "build"]
 def _convert_to_dockerfile_build_config(
     build_configuration: PythonBuild,
     dockerfile_path: str,
+    mount_pip_conf_secret: bool = False,
 ) -> DockerFileBuild:
     dockerfile_content = generate_dockerfile_content(
-        build_configuration=build_configuration
+        build_configuration=build_configuration,
+        mount_pip_conf_secret=mount_pip_conf_secret,
     )
     with open(dockerfile_path, "w", encoding="utf8") as fp:
         fp.write(dockerfile_content)
@@ -33,9 +36,12 @@ def build(
     build_configuration: PythonBuild,
     extra_opts: Optional[List[str]] = None,
 ):
+    mount_pip_conf_secret = has_pip_conf_secret(extra_opts) if extra_opts else False
     with TemporaryDirectory() as local_dir:
         docker_build_configuration = _convert_to_dockerfile_build_config(
-            build_configuration, dockerfile_path=os.path.join(local_dir, "Dockerfile")
+            build_configuration,
+            dockerfile_path=os.path.join(local_dir, "Dockerfile"),
+            mount_pip_conf_secret=mount_pip_conf_secret,
         )
         dockerfile.build(
             tag=tag,

truefoundry/deploy/builder/builders/tfy_python_buildpack/dockerfile_template.py

@@ -4,6 +4,10 @@ from typing import Dict, List, Optional
 from mako.template import Template
 
 from truefoundry.deploy.auto_gen.models import PythonBuild
+from truefoundry.deploy.builder.constants import (
+    PIP_CONF_BUILDKIT_SECRET_MOUNT,
+    PIP_CONF_SECRET_MOUNT_AS_ENV,
+)
 from truefoundry.deploy.v2.lib.patched_models import CUDAVersion
 
 # TODO (chiragjn): Switch to a non-root user inside the container
@@ -18,7 +22,7 @@ COPY ${requirements_path} ${requirements_destination_path}
 % endif
 
 % if pip_install_command is not None:
-RUN ${pip_install_command}
+RUN ${pip_config_secret_mount} ${pip_install_command}
 % endif
 
 COPY . /app
@@ -27,7 +31,7 @@ WORKDIR /app
 
 DOCKERFILE_TEMPLATE = Template(
     """
-FROM --platform=linux/amd64 python:${python_version}
+FROM --platform=linux/amd64 public.ecr.aws/docker/library/python:${python_version}
 ENV PATH=/virtualenvs/venv/bin:$PATH
 RUN apt update && \
     DEBIAN_FRONTEND=noninteractive apt install -y --no-install-recommends git && \
@@ -65,6 +69,11 @@ CUDA_VERSION_TO_IMAGE_TAG: Dict[str, str] = {
     CUDAVersion.CUDA_12_0_CUDNN8.value: "12.0.1-cudnn8-runtime-ubuntu22.04",
     CUDAVersion.CUDA_12_1_CUDNN8.value: "12.1.1-cudnn8-runtime-ubuntu22.04",
     CUDAVersion.CUDA_12_2_CUDNN8.value: "12.2.2-cudnn8-runtime-ubuntu22.04",
+    CUDAVersion.CUDA_12_3_CUDNN9.value: "12.3.2-cudnn9-runtime-ubuntu22.04",
+    # From 12.4+ onwards, the image tags drop the cudnn version
+    CUDAVersion.CUDA_12_4_CUDNN9.value: "12.4.1-cudnn-runtime-ubuntu22.04",
+    CUDAVersion.CUDA_12_5_CUDNN9.value: "12.5.1-cudnn-runtime-ubuntu22.04",
+    CUDAVersion.CUDA_12_6_CUDNN9.value: "12.6.1-cudnn-runtime-ubuntu22.04",
 }
 
 
@@ -100,7 +109,9 @@ def generate_apt_install_command(apt_packages: Optional[List[str]]) -> Optional[
 
 
 def generate_pip_install_command(
-    requirements_path: Optional[str], pip_packages: Optional[List[str]]
+    requirements_path: Optional[str],
+    pip_packages: Optional[List[str]],
+    mount_pip_conf_secret: bool = False,
 ) -> Optional[str]:
     upgrade_pip_command = "python -m pip install -U pip setuptools wheel"
     final_pip_install_command = None
@@ -119,11 +130,17 @@ def generate_pip_install_command(
     if not final_pip_install_command:
         return None
 
+    if mount_pip_conf_secret:
+        final_pip_install_command = (
+            f"{PIP_CONF_SECRET_MOUNT_AS_ENV} {final_pip_install_command}"
+        )
+
     return " && ".join([upgrade_pip_command, final_pip_install_command])
 
 
 def generate_dockerfile_content(
     build_configuration: PythonBuild,
+    mount_pip_conf_secret: bool = False,
 ) -> str:
     # TODO (chiragjn): Handle recursive references to other requirements files e.g. `-r requirements-gpu.txt`
     requirements_path = resolve_requirements_txt_path(build_configuration)
@@ -133,6 +150,7 @@ def generate_dockerfile_content(
     pip_install_command = generate_pip_install_command(
         requirements_path=requirements_destination_path,
         pip_packages=build_configuration.pip_packages,
+        mount_pip_conf_secret=mount_pip_conf_secret,
     )
     apt_install_command = generate_apt_install_command(
         apt_packages=build_configuration.apt_packages
@@ -146,6 +164,11 @@ def generate_dockerfile_content(
         "pip_install_command": pip_install_command,
     }
 
+    if mount_pip_conf_secret:
+        template_args["pip_config_secret_mount"] = PIP_CONF_BUILDKIT_SECRET_MOUNT
+    else:
+        template_args["pip_config_secret_mount"] = ""
+
     if build_configuration.cuda_version:
         template = CUDA_DOCKERFILE_TEMPLATE
         template_args["nvidia_cuda_image_tag"] = CUDA_VERSION_TO_IMAGE_TAG.get(

truefoundry/deploy/builder/constants.py

@@ -0,0 +1,7 @@
+BUILDKIT_SECRET_MOUNT_PIP_CONF_ID = "pip.conf"
+PIP_CONF_BUILDKIT_SECRET_MOUNT = (
+    f"--mount=type=secret,id={BUILDKIT_SECRET_MOUNT_PIP_CONF_ID}"
+)
+PIP_CONF_SECRET_MOUNT_AS_ENV = (
+    f"PIP_CONFIG_FILE=/run/secrets/{BUILDKIT_SECRET_MOUNT_PIP_CONF_ID}"
+)

truefoundry/deploy/builder/utils.py

@@ -0,0 +1,32 @@
+from typing import List, Optional
+
+from truefoundry.deploy.builder.constants import BUILDKIT_SECRET_MOUNT_PIP_CONF_ID
+
+
+def _get_id_from_buildkit_secret_value(value: str) -> Optional[str]:
+    parts = value.split(",")
+    secret_config = {}
+    for part in parts:
+        kv = part.split("=", 1)
+        if len(kv) != 2:
+            continue
+        key, value = kv
+        secret_config[key] = value
+
+    if "id" in secret_config and "src" in secret_config:
+        return secret_config["id"]
+
+    return None
+
+
+def has_pip_conf_secret(docker_build_extra_args: List[str]) -> bool:
+    args = [arg.strip() for arg in docker_build_extra_args]
+    for i, arg in enumerate(docker_build_extra_args):
+        if (
+            arg == "--secret"
+            and i + 1 < len(args)
+            and _get_id_from_buildkit_secret_value(args[i + 1])
+            == BUILDKIT_SECRET_MOUNT_PIP_CONF_ID
+        ):
+            return True
+    return False