truefoundry 0.4.1rc1__tar.gz → 0.4.2__tar.gz

{truefoundry-0.4.1rc1 → truefoundry-0.4.2}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: truefoundry
-Version: 0.4.1rc1
+Version: 0.4.2
 Summary: Truefoundry CLI
 Author: Abhishek Choudhary
 Author-email: abhishek@truefoundry.com
@@ -31,7 +31,6 @@ Requires-Dist: openai (>=1.16.2,<2.0.0)
 Requires-Dist: packaging (>=20.0,<25.0)
 Requires-Dist: pandas (>=1.0.0,<3.0.0) ; python_version < "3.10"
 Requires-Dist: pandas (>=1.4.0,<3.0.0) ; python_version >= "3.10"
-Requires-Dist: psutil (>=5.9.0,<7.0.0)
 Requires-Dist: pydantic (>=1.8.2,<3.0.0)
 Requires-Dist: pygments (>=2.12.0,<3.0.0)
 Requires-Dist: python-dateutil (>=2.8.2,<3.0.0)

{truefoundry-0.4.1rc1 → truefoundry-0.4.2}/pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "truefoundry"
-version = "0.4.1rc1"
+version = "0.4.2"
 description = "Truefoundry CLI"
 authors = ["Abhishek Choudhary <abhishek@truefoundry.com>"]
 readme = "README.md"
@@ -21,7 +21,6 @@ importlib-resources = ">=5.2.0,<7.0.0"
 Mako = ">=1.1.6,<2.0.0"
 openai = ">=1.16.2,<2.0.0"
 packaging = ">=20.0,<25.0"
-psutil = ">=5.9.0,<7.0.0"
 pydantic = ">=1.8.2,<3.0.0"
 pygments = ">=2.12.0,<3.0.0"
 PyJWT = ">=2.0.0,<3.0.0"

{truefoundry-0.4.1rc1 → truefoundry-0.4.2}/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-0.4.1rc1 → truefoundry-0.4.2}/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-0.4.1rc1 → truefoundry-0.4.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-0.4.1rc1 → truefoundry-0.4.2}/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-0.4.1rc1 → truefoundry-0.4.2}/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-0.4.2/truefoundry/common/utils.py

@@ -0,0 +1,105 @@
+import os
+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, urlparse
+
+from truefoundry.common.constants import (
+    API_SERVER_RELATIVE_PATH,
+    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]]:
+    def wrapper_cache(func: Callable[..., T]) -> Callable[..., T]:
+        func = lru_cache(maxsize=maxsize)(func)
+        func.delta = seconds * 10**9
+        func.expiration = monotonic_ns() + func.delta
+
+        @wraps(func)
+        def wrapped_func(*args, **kwargs):
+            if monotonic_ns() >= func.expiration:
+                func.cache_clear()
+                func.expiration = monotonic_ns() + func.delta
+            return func(*args, **kwargs)
+
+        return wrapped_func
+
+    return wrapper_cache
+
+
+def poll_for_function(
+    func: Callable[..., T], poll_after_secs: int = 5, *args, **kwargs
+) -> Generator[T, None, None]:
+    while True:
+        yield func(*args, **kwargs)
+        time.sleep(poll_after_secs)
+
+
+def validate_tfy_host(tfy_host: str) -> None:
+    if not (tfy_host.startswith("https://") or tfy_host.startswith("http://")):
+        raise ValueError(
+            f"Invalid host {tfy_host!r}. It should start with https:// or http://"
+        )
+
+
+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-0.4.1rc1 → truefoundry-0.4.2}/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-0.4.1rc1 → truefoundry-0.4.2}/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-0.4.1rc1 → truefoundry-0.4.2}/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-0.4.2/truefoundry/deploy/builder/builders/tfy_notebook_buildpack/dockerfile_template.py

@@ -0,0 +1,79 @@
+import hashlib
+import os
+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
+
+
+class NotebookImageBuild(BaseModel):
+    type: Literal["tfy-notebook-buildpack"] = "tfy-notebook-buildpack"
+    base_image_uri: str
+    build_script: Optional[str] = None
+
+
+DOCKERFILE_TEMPLATE = Template(
+    """
+FROM ${base_image_uri}
+USER root
+
+% if build_script_docker_commands is not None:
+${build_script_docker_commands}
+% endif
+
+USER $NB_UID
+"""
+)
+
+
+def generate_build_script_docker_commands(
+    build_script: Optional[str],
+    local_dir: str,
+    mount_pip_conf_secret: bool = False,
+) -> Optional[str]:
+    if not build_script:
+        return None
+    # 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 {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,
+    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 = {
+        "base_image_uri": build_configuration.base_image_uri,
+        "build_script_docker_commands": build_script_docker_commands,
+    }
+
+    template = DOCKERFILE_TEMPLATE
+
+    dockerfile_content: str = template.render(**template_args)
+    return dockerfile_content

{truefoundry-0.4.1rc1 → truefoundry-0.4.2}/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-0.4.1rc1 → truefoundry-0.4.2}/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-0.4.2/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-0.4.2/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