truefoundry 0.11.3rc1__py3-none-any.whl → 0.11.4__py3-none-any.whl

truefoundry/common/constants.py

@@ -68,6 +68,8 @@ class TrueFoundrySdkEnv(BaseSettings):
     # TODO(gw): Use another image with more linient rate limits
     TFY_SPARK_BUILD_SPARK_IMAGE_REPO: str = "public.ecr.aws/bitnami/spark"
 
+    TFY_TASK_PYSPARK_BUILD_SPARK_IMAGE_REPO: str = "public.ecr.aws/bitnami/spark"
+
     # For local development, this enables futher configuration via _TFYServersConfig
     TFY_CLI_LOCAL_DEV_MODE: bool = False
 

truefoundry/deploy/__init__.py

@@ -120,6 +120,7 @@ from truefoundry.deploy.v2.lib.patched_models import (
     SparkJobPythonEntrypoint,
     SparkJobPythonNotebookEntrypoint,
     SparkJobScalaEntrypoint,
+    SparkJobScalaNotebookEntrypoint,
     SQSInputConfig,
     SQSOutputConfig,
     SQSQueueMetricConfig,

truefoundry/deploy/_autogen/models.py

@@ -1,6 +1,6 @@
 # generated by datamodel-codegen:
 #   filename:  application.json
-#   timestamp: 2025-07-16T09:53:23+00:00
+#   timestamp: 2025-07-28T14:11:36+00:00
 
 from __future__ import annotations
 
@@ -103,7 +103,7 @@ class AsyncProcessorSidecar(BaseModel):
 
 
 class Autoshutdown(BaseModel):
-    wait_time: conint(ge=0) = Field(
+    wait_time: conint(ge=0, le=604800) = Field(
         900,
         description="The period to wait after the last received request before scaling the replicas to 0. This value should be high enough to allow for the replicas of the service to come up to avoid premature scaling down.",
     )
@@ -957,6 +957,14 @@ class SparkJobScalaEntrypoint(BaseModel):
     )
 
 
+class SparkJobScalaNotebookEntrypoint(BaseModel):
+    type: Literal["scala-notebook"] = Field(..., description="")
+    main_application_file: str = Field(
+        ...,
+        description="The main application file to be executed by the spark job. Relative path in case of git repository.",
+    )
+
+
 class StaticVolumeConfig(BaseModel):
     type: Literal["static"] = Field(..., description="Volume Type for the volume.")
     persistent_volume_name: str = Field(
@@ -1489,6 +1497,7 @@ class SparkJob(BaseModel):
         SparkJobScalaEntrypoint,
         SparkJobJavaEntrypoint,
         SparkJobPythonNotebookEntrypoint,
+        SparkJobScalaNotebookEntrypoint,
     ] = Field(..., description="")
     driver_config: SparkDriverConfig
     executor_config: SparkExecutorConfig

truefoundry/deploy/builder/__init__.py

@@ -5,6 +5,7 @@ from truefoundry.deploy._autogen.models import (
     PythonBuild,
     SparkBuild,
     TaskDockerFileBuild,
+    TaskPySparkBuild,
     TaskPythonBuild,
 )
 from truefoundry.deploy.builder.builders import get_builder
@@ -22,6 +23,7 @@ class _BuildConfig(BaseModel):
         TaskPythonBuild,
         TaskDockerFileBuild,
         SparkBuild,
+        TaskPySparkBuild,
     ] = Field(discriminator="type")
 
 

truefoundry/deploy/builder/builders/__init__.py

@@ -5,6 +5,7 @@ from truefoundry.deploy.builder.builders import (
     tfy_notebook_buildpack,
     tfy_python_buildpack,
     tfy_spark_buildpack,
+    tfy_task_pyspark_buildpack,
 )
 
 BUILD_REGISTRY: Dict[str, Callable] = {
@@ -12,6 +13,7 @@ BUILD_REGISTRY: Dict[str, Callable] = {
     "tfy-python-buildpack": tfy_python_buildpack.build,
     "tfy-notebook-buildpack": tfy_notebook_buildpack.build,
     "tfy-spark-buildpack": tfy_spark_buildpack.build,
+    "task-pyspark-build": tfy_task_pyspark_buildpack.build,
 }
 
 __all__ = ["get_builder"]

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

@@ -1,5 +1,4 @@
-import shlex
-from typing import Dict, List, Optional
+from typing import Dict
 
 from mako.template import Template
 
@@ -7,9 +6,12 @@ from truefoundry.common.constants import ENV_VARS, PythonPackageManager
 from truefoundry.deploy._autogen.models import PythonBuild
 from truefoundry.deploy.builder.constants import (
     PIP_CONF_BUILDKIT_SECRET_MOUNT,
-    PIP_CONF_SECRET_MOUNT_AS_ENV,
     UV_CONF_BUILDKIT_SECRET_MOUNT,
-    UV_CONF_SECRET_MOUNT_AS_ENV,
+)
+from truefoundry.deploy.builder.utils import (
+    generate_apt_install_command,
+    generate_pip_install_command,
+    generate_uv_pip_install_command,
 )
 from truefoundry.deploy.v2.lib.patched_models import (
     CUDAVersion,
@@ -82,85 +84,6 @@ CUDA_VERSION_TO_IMAGE_TAG: Dict[str, str] = {
 }
 
 
-def generate_apt_install_command(apt_packages: Optional[List[str]]) -> Optional[str]:
-    packages_list = None
-    if apt_packages:
-        packages_list = " ".join(p.strip() for p in apt_packages if p.strip())
-    if not packages_list:
-        return None
-    apt_update_command = "apt update"
-    apt_install_command = f"DEBIAN_FRONTEND=noninteractive apt install -y --no-install-recommends {packages_list}"
-    clear_apt_lists_command = "rm -rf /var/lib/apt/lists/*"
-    return " && ".join(
-        [apt_update_command, apt_install_command, clear_apt_lists_command]
-    )
-
-
-def generate_pip_install_command(
-    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"
-    envs = []
-    if mount_pip_conf_secret:
-        envs.append(PIP_CONF_SECRET_MOUNT_AS_ENV)
-
-    command = ["python", "-m", "pip", "install", "--use-pep517", "--no-cache-dir"]
-    args = []
-    if requirements_path:
-        args.append("-r")
-        args.append(requirements_path)
-
-    if pip_packages:
-        args.extend(pip_packages)
-
-    if not args:
-        return None
-
-    final_pip_install_command = shlex.join(envs + command + args)
-    final_docker_run_command = " && ".join(
-        [upgrade_pip_command, final_pip_install_command]
-    )
-    return final_docker_run_command
-
-
-def generate_uv_pip_install_command(
-    requirements_path: Optional[str],
-    pip_packages: Optional[List[str]],
-    mount_uv_conf_secret: bool = False,
-) -> Optional[str]:
-    upgrade_pip_command = "python -m pip install -U pip setuptools wheel"
-    uv_mount = f"--mount=from={ENV_VARS.TFY_PYTHON_BUILD_UV_IMAGE_URI},source=/uv,target=/usr/local/bin/uv"
-    envs = [
-        "UV_LINK_MODE=copy",
-        "UV_PYTHON_DOWNLOADS=never",
-        "UV_INDEX_STRATEGY=unsafe-best-match",
-    ]
-    if mount_uv_conf_secret:
-        envs.append(UV_CONF_SECRET_MOUNT_AS_ENV)
-
-    command = ["uv", "pip", "install", "--no-cache-dir"]
-
-    args = []
-
-    if requirements_path:
-        args.append("-r")
-        args.append(requirements_path)
-
-    if pip_packages:
-        args.extend(pip_packages)
-
-    if not args:
-        return None
-
-    uv_pip_install_command = shlex.join(envs + command + args)
-    shell_commands = " && ".join([upgrade_pip_command, uv_pip_install_command])
-    final_docker_run_command = " ".join([uv_mount, shell_commands])
-
-    return final_docker_run_command
-
-
 def generate_dockerfile_content(
     build_configuration: PythonBuild,
     package_manager: str = ENV_VARS.TFY_PYTHON_BUILD_PACKAGE_MANAGER,

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

@@ -1,15 +1,14 @@
-import shlex
-from typing import List, Optional
-
 from mako.template import Template
 
 from truefoundry.common.constants import ENV_VARS, PythonPackageManager
 from truefoundry.deploy._autogen.models import SparkBuild
 from truefoundry.deploy.builder.constants import (
     PIP_CONF_BUILDKIT_SECRET_MOUNT,
-    PIP_CONF_SECRET_MOUNT_AS_ENV,
     UV_CONF_BUILDKIT_SECRET_MOUNT,
-    UV_CONF_SECRET_MOUNT_AS_ENV,
+)
+from truefoundry.deploy.builder.utils import (
+    generate_pip_install_command,
+    generate_uv_pip_install_command,
 )
 from truefoundry.deploy.v2.lib.patched_models import (
     _resolve_requirements_path,
@@ -26,8 +25,7 @@ RUN ${package_manager_config_secret_mount} ${python_packages_install_command}
 % endif
 ENV PYTHONDONTWRITEBYTECODE=1
 ENV IPYTHONDIR=/tmp/.ipython
-RUN groupadd --system --gid 1001 spark && useradd --system --uid 1001 --gid spark --no-create-home spark
-USER spark
+USER 1001
 COPY . /app
 """
 
@@ -35,6 +33,20 @@ _POST_USER_TEMPLATE = """
 COPY tfy_execute_notebook.py /app/tfy_execute_notebook.py
 """
 
+_ALMOND_INSTALL_TEMPLATE = """
+ENV COURSIER_CACHE=/opt/coursier-cache
+RUN install_packages curl
+RUN curl -Lo coursier https://git.io/coursier-cli && \
+    chmod +x coursier && \
+    ./coursier launch almond:0.14.1 -- --install --global && \
+    chown -R 1001:0 /usr/local/share/jupyter && \
+    chown -R 1001:0 /opt/coursier-cache && \
+    rm -f coursier
+"""
+
+# Docker image size with almond - 1.26GB
+# Docker image size without almond - 1.1GB
+# Not much harm in packaging almond by default
 DOCKERFILE_TEMPLATE = Template(
     """
 FROM ${spark_image_repo}:${spark_version}
@@ -43,6 +55,7 @@ RUN apt update && \
     DEBIAN_FRONTEND=noninteractive apt install -y --no-install-recommends git && \
     rm -rf /var/lib/apt/lists/*
 """
+    + _ALMOND_INSTALL_TEMPLATE
     + _POST_PYTHON_INSTALL_TEMPLATE
     + _POST_USER_TEMPLATE
 )
@@ -55,71 +68,6 @@ ADDITIONAL_PIP_PACKAGES = [
 ]
 
 
-def generate_pip_install_command(
-    requirements_path: Optional[str],
-    pip_packages: Optional[List[str]],
-    mount_pip_conf_secret: bool = False,
-) -> Optional[str]:
-    upgrade_pip_command = "python3 -m pip install -U pip setuptools wheel"
-    envs = []
-    if mount_pip_conf_secret:
-        envs.append(PIP_CONF_SECRET_MOUNT_AS_ENV)
-
-    command = ["python3", "-m", "pip", "install", "--use-pep517", "--no-cache-dir"]
-    args = []
-    if requirements_path:
-        args.append("-r")
-        args.append(requirements_path)
-
-    if pip_packages:
-        args.extend(pip_packages)
-
-    if not args:
-        return None
-
-    final_pip_install_command = shlex.join(envs + command + args)
-    final_docker_run_command = " && ".join(
-        [upgrade_pip_command, final_pip_install_command]
-    )
-    return final_docker_run_command
-
-
-def generate_uv_pip_install_command(
-    requirements_path: Optional[str],
-    pip_packages: Optional[List[str]],
-    mount_uv_conf_secret: bool = False,
-) -> Optional[str]:
-    upgrade_pip_command = "python3 -m pip install -U pip setuptools wheel"
-    uv_mount = f"--mount=from={ENV_VARS.TFY_PYTHON_BUILD_UV_IMAGE_URI},source=/uv,target=/usr/local/bin/uv"
-    envs = [
-        "UV_LINK_MODE=copy",
-        "UV_PYTHON_DOWNLOADS=never",
-        "UV_INDEX_STRATEGY=unsafe-best-match",
-    ]
-    if mount_uv_conf_secret:
-        envs.append(UV_CONF_SECRET_MOUNT_AS_ENV)
-
-    command = ["uv", "pip", "install", "--no-cache-dir"]
-
-    args = []
-
-    if requirements_path:
-        args.append("-r")
-        args.append(requirements_path)
-
-    if pip_packages:
-        args.extend(pip_packages)
-
-    if not args:
-        return None
-
-    uv_pip_install_command = shlex.join(envs + command + args)
-    shell_commands = " && ".join([upgrade_pip_command, uv_pip_install_command])
-    final_docker_run_command = " ".join([uv_mount, shell_commands])
-
-    return final_docker_run_command
-
-
 def generate_dockerfile_content(
     build_configuration: SparkBuild,
     package_manager: str = ENV_VARS.TFY_PYTHON_BUILD_PACKAGE_MANAGER,

truefoundry/deploy/builder/builders/tfy_spark_buildpack/tfy_execute_notebook.py

@@ -86,12 +86,24 @@ def execute_notebook(notebook_path, output_path="/tmp/output.ipynb", parameters=
         parameters = {}
 
     print(f"Starting execution of notebook: {notebook_path}")
+    notebook_type = os.environ.get("TFY_NOTEBOOK_TYPE", "").lower()
+    kernel_mapping = {"python": "python3", "scala": "scala"}
+
+    if notebook_type not in kernel_mapping:
+        supported_types = ", ".join(kernel_mapping.keys())
+        raise ValueError(
+            f"Unsupported notebook type: '{notebook_type}'. "
+            f"Supported types: [{supported_types}]"
+        )
+
+    kernel_name = kernel_mapping[notebook_type]
+
     pm.execute_notebook(
         input_path=notebook_path,
         output_path=output_path,
         parameters=parameters,
         # TODO(gw): Replace with kernel name for venv
-        kernel_name="python3",
+        kernel_name=kernel_name,
         # Log cell by cell execution output
         # TODO(gw): Output logs to a file instead, so that they aren't merged with the container's logs
         log_output=True,
@@ -107,6 +119,7 @@ def validate_env_vars():
         "TFY_NOTEBOOK_OUTPUT_S3_BUCKET",
         "SPARK_APPLICATION_EVENT_LOG_JWT_TOKEN",
         "TFY_NOTEBOOK_OUTPUT_S3_SECRET_KEY",
+        "TFY_NOTEBOOK_TYPE",
     ]
     unset_keys = [key for key in keys if not os.environ.get(key)]
     if unset_keys:

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

@@ -0,0 +1,52 @@
+import os
+from tempfile import TemporaryDirectory
+from typing import List, Optional
+
+from truefoundry.deploy._autogen.models import DockerFileBuild, TaskPySparkBuild
+from truefoundry.deploy.builder.builders import dockerfile
+from truefoundry.deploy.builder.builders.tfy_task_pyspark_buildpack.dockerfile_template import (
+    generate_dockerfile_content,
+)
+from truefoundry.deploy.builder.utils import has_python_package_manager_conf_secret
+
+__all__ = ["generate_dockerfile_content", "build"]
+
+
+def _convert_to_dockerfile_build_config(
+    build_configuration: TaskPySparkBuild,
+    dockerfile_path: str,
+    mount_python_package_manager_conf_secret: bool = False,
+) -> DockerFileBuild:
+    dockerfile_content = generate_dockerfile_content(
+        build_configuration=build_configuration,
+        mount_python_package_manager_conf_secret=mount_python_package_manager_conf_secret,
+    )
+    with open(dockerfile_path, "w", encoding="utf8") as fp:
+        fp.write(dockerfile_content)
+
+    return DockerFileBuild(
+        type="dockerfile",
+        dockerfile_path=dockerfile_path,
+    )
+
+
+def build(
+    tag: str,
+    build_configuration: TaskPySparkBuild,
+    extra_opts: Optional[List[str]] = None,
+):
+    mount_python_package_manager_conf_secret = (
+        has_python_package_manager_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"),
+            mount_python_package_manager_conf_secret=mount_python_package_manager_conf_secret,
+        )
+        dockerfile.build(
+            tag=tag,
+            build_configuration=docker_build_configuration,
+            extra_opts=extra_opts,
+        )

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

@@ -0,0 +1,121 @@
+from mako.template import Template
+
+from truefoundry.common.constants import ENV_VARS, PythonPackageManager
+from truefoundry.deploy._autogen.models import TaskPySparkBuild
+from truefoundry.deploy.builder.constants import (
+    PIP_CONF_BUILDKIT_SECRET_MOUNT,
+    UV_CONF_BUILDKIT_SECRET_MOUNT,
+)
+from truefoundry.deploy.builder.utils import (
+    generate_apt_install_command,
+    generate_pip_install_command,
+    generate_uv_pip_install_command,
+)
+from truefoundry.deploy.v2.lib.patched_models import (
+    _resolve_requirements_path,
+)
+
+# TODO[GW]: Switch to a non-root user inside the container
+_POST_PYTHON_INSTALL_TEMPLATE = """
+% if apt_install_command is not None:
+RUN ${apt_install_command}
+% endif
+% if requirements_path is not None:
+COPY ${requirements_path} ${requirements_destination_path}
+% endif
+% if python_packages_install_command is not None:
+RUN ${package_manager_config_secret_mount} ${python_packages_install_command}
+% endif
+COPY . /app
+WORKDIR /app
+"""
+
+# TODO[GW]: Check if the entrypoint for the image needs to change
+# Using /opt/venv/ because flyte seems to be using it and this doesn't look configurable
+# TODO[GW]: Double check this^
+DOCKERFILE_TEMPLATE = Template(
+    """
+FROM ${spark_image_repo}:${spark_version}
+ENV PATH=/opt/venv/bin:$PATH
+USER root
+RUN mkdir -p /var/lib/apt/lists/partial && \
+    apt update && \
+    DEBIAN_FRONTEND=noninteractive apt install -y --no-install-recommends git && \
+    python -m venv /opt/venv/ && \
+    rm -rf /var/lib/apt/lists/*
+"""
+    + _POST_PYTHON_INSTALL_TEMPLATE
+)
+
+
+def get_additional_pip_packages(build_configuration: TaskPySparkBuild):
+    return [
+        f"pyspark=={build_configuration.spark_version}",
+    ]
+
+
+def generate_dockerfile_content(
+    build_configuration: TaskPySparkBuild,
+    package_manager: str = ENV_VARS.TFY_PYTHON_BUILD_PACKAGE_MANAGER,
+    mount_python_package_manager_conf_secret: bool = False,
+) -> str:
+    # TODO (chiragjn): Handle recursive references to other requirements files e.g. `-r requirements-gpu.txt`
+    requirements_path = _resolve_requirements_path(
+        build_context_path="",
+        requirements_path=build_configuration.requirements_path,
+    )
+    requirements_destination_path = (
+        "/tmp/requirements.txt" if requirements_path else None
+    )
+    # if not build_configuration.python_version:
+    #     raise ValueError(
+    #         "`python_version` is required for `tfy-python-buildpack` builder"
+    #     )
+    pip_packages = get_additional_pip_packages(build_configuration) + (
+        build_configuration.pip_packages or []
+    )
+    if package_manager == PythonPackageManager.PIP.value:
+        python_packages_install_command = generate_pip_install_command(
+            requirements_path=requirements_destination_path,
+            pip_packages=pip_packages,
+            mount_pip_conf_secret=mount_python_package_manager_conf_secret,
+        )
+    elif package_manager == PythonPackageManager.UV.value:
+        python_packages_install_command = generate_uv_pip_install_command(
+            requirements_path=requirements_destination_path,
+            pip_packages=pip_packages,
+            mount_uv_conf_secret=mount_python_package_manager_conf_secret,
+        )
+    else:
+        raise ValueError(f"Unsupported package manager: {package_manager}")
+
+    apt_install_command = generate_apt_install_command(
+        apt_packages=build_configuration.apt_packages
+    )
+    template_args = {
+        "spark_image_repo": ENV_VARS.TFY_TASK_PYSPARK_BUILD_SPARK_IMAGE_REPO,
+        "spark_version": build_configuration.spark_version,
+        "apt_install_command": apt_install_command,
+        "requirements_path": requirements_path,
+        "requirements_destination_path": requirements_destination_path,
+        "python_packages_install_command": python_packages_install_command,
+    }
+
+    if mount_python_package_manager_conf_secret:
+        if package_manager == PythonPackageManager.PIP.value:
+            template_args["package_manager_config_secret_mount"] = (
+                PIP_CONF_BUILDKIT_SECRET_MOUNT
+            )
+        elif package_manager == PythonPackageManager.UV.value:
+            template_args["package_manager_config_secret_mount"] = (
+                UV_CONF_BUILDKIT_SECRET_MOUNT
+            )
+        else:
+            raise ValueError(f"Unsupported package manager: {package_manager}")
+    else:
+        template_args["package_manager_config_secret_mount"] = ""
+
+    template = DOCKERFILE_TEMPLATE
+
+    dockerfile_content = template.render(**template_args)
+    return dockerfile_content

truefoundry/deploy/builder/utils.py

@@ -1,8 +1,12 @@
+import shlex
 from typing import List, Optional
 
+from truefoundry.common.constants import ENV_VARS
 from truefoundry.deploy.builder.constants import (
     BUILDKIT_SECRET_MOUNT_PIP_CONF_ID,
     BUILDKIT_SECRET_MOUNT_UV_CONF_ID,
+    PIP_CONF_SECRET_MOUNT_AS_ENV,
+    UV_CONF_SECRET_MOUNT_AS_ENV,
 )
 
 
@@ -35,3 +39,82 @@ def has_python_package_manager_conf_secret(docker_build_extra_args: List[str]) -
         ):
             return True
     return False
+
+
+def generate_pip_install_command(
+    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"
+    envs = []
+    if mount_pip_conf_secret:
+        envs.append(PIP_CONF_SECRET_MOUNT_AS_ENV)
+
+    command = ["python", "-m", "pip", "install", "--use-pep517", "--no-cache-dir"]
+    args = []
+    if requirements_path:
+        args.append("-r")
+        args.append(requirements_path)
+
+    if pip_packages:
+        args.extend(pip_packages)
+
+    if not args:
+        return None
+
+    final_pip_install_command = shlex.join(envs + command + args)
+    final_docker_run_command = " && ".join(
+        [upgrade_pip_command, final_pip_install_command]
+    )
+    return final_docker_run_command
+
+
+def generate_uv_pip_install_command(
+    requirements_path: Optional[str],
+    pip_packages: Optional[List[str]],
+    mount_uv_conf_secret: bool = False,
+) -> Optional[str]:
+    upgrade_pip_command = "python -m pip install -U pip setuptools wheel"
+    uv_mount = f"--mount=from={ENV_VARS.TFY_PYTHON_BUILD_UV_IMAGE_URI},source=/uv,target=/usr/local/bin/uv"
+    envs = [
+        "UV_LINK_MODE=copy",
+        "UV_PYTHON_DOWNLOADS=never",
+        "UV_INDEX_STRATEGY=unsafe-best-match",
+    ]
+    if mount_uv_conf_secret:
+        envs.append(UV_CONF_SECRET_MOUNT_AS_ENV)
+
+    command = ["uv", "pip", "install", "--no-cache-dir"]
+
+    args = []
+
+    if requirements_path:
+        args.append("-r")
+        args.append(requirements_path)
+
+    if pip_packages:
+        args.extend(pip_packages)
+
+    if not args:
+        return None
+
+    uv_pip_install_command = shlex.join(envs + command + args)
+    shell_commands = " && ".join([upgrade_pip_command, uv_pip_install_command])
+    final_docker_run_command = " ".join([uv_mount, shell_commands])
+
+    return final_docker_run_command
+
+
+def generate_apt_install_command(apt_packages: Optional[List[str]]) -> Optional[str]:
+    packages_list = None
+    if apt_packages:
+        packages_list = " ".join(p.strip() for p in apt_packages if p.strip())
+    if not packages_list:
+        return None
+    apt_update_command = "apt update"
+    apt_install_command = f"DEBIAN_FRONTEND=noninteractive apt install -y --no-install-recommends {packages_list}"
+    clear_apt_lists_command = "rm -rf /var/lib/apt/lists/*"
+    return " && ".join(
+        [apt_update_command, apt_install_command, clear_apt_lists_command]
+    )

truefoundry/deploy/v2/lib/patched_models.py

@@ -528,6 +528,12 @@ class SparkJobPythonNotebookEntrypoint(
     type: Literal["python-notebook"] = "python-notebook"
 
 
+class SparkJobScalaNotebookEntrypoint(
+    models.SparkJobScalaNotebookEntrypoint, PatchedModelBase
+):
+    type: Literal["scala-notebook"] = "scala-notebook"
+
+
 class PySparkTaskConfig(models.PySparkTaskConfig, PatchedModelBase):
     type: Literal["pyspark-task-config"] = "pyspark-task-config"
 

truefoundry/ml/artifact/truefoundry_artifact_repo.py

@@ -654,7 +654,7 @@ class MlFoundryArtifactsRepository:
                 artifact_identifier=self.artifact_identifier, paths=[remote_file_path]
             )[0]
 
-        if progress_bar is None or not progress_bar.disable:
+        if progress_bar is None or progress_bar.disable:
             logger.info("Downloading %s to %s", remote_file_path, local_path)
 
         if progress_bar is not None:

truefoundry/ml/integrations/huggingface/trainer_callback.py

@@ -0,0 +1,198 @@
+import logging
+import math
+import os
+from typing import TYPE_CHECKING, Any, Dict, Optional
+
+import numpy as np
+
+from truefoundry import ml
+
+try:
+    from transformers.integrations.integration_utils import rewrite_logs
+    from transformers.trainer_callback import TrainerCallback
+except ImportError as e:
+    raise ImportError(
+        "Importing this module requires `transformers` to be installed"
+    ) from e
+
+if TYPE_CHECKING:
+    from transformers.trainer_callback import TrainerControl, TrainerState
+    from transformers.training_args import TrainingArguments
+
+    from truefoundry.ml import MlFoundryRun
+
+logger = logging.getLogger(__name__)
+
+
+class TrueFoundryMLCallback(TrainerCallback):
+    def __init__(
+        self,
+        run: "MlFoundryRun",
+        log_checkpoints: bool = True,
+        checkpoint_artifact_name: Optional[str] = None,
+        auto_end_run_on_train_end: bool = False,
+    ):
+        """
+        Args:
+            run: The run entity to log metrics to.
+            log_checkpoints: Whether to log checkpoints or not, defaults to True.
+            checkpoint_artifact_name: The name of the artifact to log checkpoints to, required if log_checkpoints is True.
+            auto_end_run_on_train_end: Whether to end the run automatically when training ends, defaults to False.
+
+            Usage:
+                from transformers import Trainer
+                from truefoundry.ml.integrations.huggingface.trainer_callback import TrueFoundryMLCallback
+                from truefoundry.ml import get_client
+
+                client = get_client()
+                run = client.create_run(ml_repo="my-ml-repo", run_name="my-run", auto_end=False)
+
+                callback = TrueFoundryMLCallback(
+                    run=run,
+                    log_checkpoints=True,
+                    checkpoint_artifact_name="my-checkpoint",
+                    auto_end_run_on_train_end=True,
+                )
+
+                trainer = Trainer(
+                    ...,
+                    callbacks=[callback]
+                )
+        """
+        self._run = run
+        self._log_checkpoints = log_checkpoints
+        if self._log_checkpoints and not checkpoint_artifact_name:
+            raise ValueError(
+                "`checkpoint_artifact_name` is required when `log_checkpoints` is True"
+            )
+        self._checkpoint_artifact_name = checkpoint_artifact_name
+        self._auto_end_run_on_train_end = auto_end_run_on_train_end
+
+    @classmethod
+    def with_managed_run(
+        cls,
+        ml_repo: str,
+        run_name: Optional[str] = None,
+        log_checkpoints: bool = True,
+        checkpoint_artifact_name: Optional[str] = None,
+        auto_end_run_on_train_end: bool = True,
+    ) -> "TrueFoundryMLCallback":
+        """
+        Args:
+            ml_repo: The name of the ML Repository to log metrics and data to.
+            run_name: The name of the run, if not provided, a random name will be generated.
+            log_checkpoints: Whether to log checkpoints or not, defaults to True.
+            checkpoint_artifact_name: The name of the artifact to log checkpoints to, required if log_checkpoints is True.
+            auto_end_run_on_train_end: Whether to end the run automatically when training ends, defaults to True.
+
+        Usage:
+            from transformers import Trainer
+            from truefoundry.ml.integrations.huggingface.trainer_callback import TrueFoundryMLCallback
+
+            callback = TrueFoundryMLCallback.with_managed_run(
+                ml_repo="my-ml-repo",
+                run_name="my-run",
+                log_checkpoints=True,
+                checkpoint_artifact_name="my-checkpoint",
+                auto_end_run_on_train_end=True,
+            )
+            trainer = Trainer(
+                ...,
+                callbacks=[callback]
+            )
+        """
+        run = ml.get_client().create_run(
+            ml_repo=ml_repo, run_name=run_name, auto_end=False
+        )
+        return cls(
+            run=run,
+            log_checkpoints=log_checkpoints,
+            checkpoint_artifact_name=checkpoint_artifact_name,
+            auto_end_run_on_train_end=auto_end_run_on_train_end,
+        )
+
+    def _drop_non_finite_values(self, dct: Dict[str, Any]) -> Dict[str, Any]:
+        sanitized = {}
+        for k, v in dct.items():
+            if isinstance(v, (int, float, np.integer, np.floating)) and math.isfinite(
+                v
+            ):
+                sanitized[k] = v
+            else:
+                logger.warning(
+                    f'Trainer is attempting to log a value of "{v}" of'
+                    f' type {type(v)} for key "{k}" as a metric.'
+                    " Mlfoundry's log_metric() only accepts finite float and"
+                    " int types so we dropped this attribute."
+                )
+        return sanitized
+
+    @property
+    def run(self) -> "MlFoundryRun":
+        return self._run
+
+    # noinspection PyMethodOverriding
+    def on_log(
+        self,
+        args: "TrainingArguments",
+        state: "TrainerState",
+        control: "TrainerControl",
+        logs: Optional[Dict[str, Any]] = None,
+        **kwargs,
+    ):
+        logs = logs or {}
+        if not state.is_world_process_zero:
+            return
+
+        metrics = self._drop_non_finite_values(logs)
+        self._run.log_metrics(rewrite_logs(metrics), step=state.global_step)
+
+    def on_save(
+        self,
+        args: "TrainingArguments",
+        state: "TrainerState",
+        control: "TrainerControl",
+        **kwargs,
+    ):
+        if not state.is_world_process_zero:
+            return
+
+        if not self._log_checkpoints:
+            return
+
+        if not self._checkpoint_artifact_name:
+            return
+
+        ckpt_dir = f"checkpoint-{state.global_step}"
+        artifact_path = os.path.join(args.output_dir, ckpt_dir)
+        description = None
+        _job_name = os.getenv("TFY_INTERNAL_COMPONENT_NAME")
+        _job_run_name = os.getenv("TFY_INTERNAL_JOB_RUN_NAME")
+        if _job_name:
+            description = f"Checkpoint from job={_job_name} run={_job_run_name}"
+        logger.info(f"Uploading checkpoint {ckpt_dir} ...")
+        metadata = {}
+        for log in state.log_history:
+            if isinstance(log, dict) and log.get("step") == state.global_step:
+                metadata = log.copy()
+        metadata = self._drop_non_finite_values(metadata)
+        self._run.log_artifact(
+            name=self._checkpoint_artifact_name,
+            artifact_paths=[(artifact_path, None)],
+            metadata=metadata,
+            step=state.global_step,
+            description=description,
+        )
+
+    def on_train_end(
+        self,
+        args: "TrainingArguments",
+        state: "TrainerState",
+        control: "TrainerControl",
+        **kwargs,
+    ):
+        """
+        Event called at the end of training.
+        """
+        if self._auto_end_run_on_train_end:
+            self._run.end()

{truefoundry-0.11.3rc1.dist-info → truefoundry-0.11.4.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: truefoundry
-Version: 0.11.3rc1
+Version: 0.11.4
 Summary: TrueFoundry CLI
 Author-email: TrueFoundry Team <abhishek@truefoundry.com>
 Requires-Python: <3.14,>=3.8.1
@@ -30,7 +30,7 @@ Requires-Dist: requirements-parser<0.12.0,>=0.11.0
 Requires-Dist: rich-click<2.0.0,>=1.2.1
 Requires-Dist: rich<14.0.0,>=13.7.1
 Requires-Dist: tqdm<5.0.0,>=4.0.0
-Requires-Dist: truefoundry-sdk<0.2.0,>=0.1.7
+Requires-Dist: truefoundry-sdk<0.2.0,>=0.1.9
 Requires-Dist: typing-extensions>=4.0
 Requires-Dist: urllib3<3,>=1.26.18
 Requires-Dist: yq<4.0.0,>=3.1.0

{truefoundry-0.11.3rc1.dist-info → truefoundry-0.11.4.dist-info}/RECORD

@@ -40,7 +40,7 @@ truefoundry/cli/display_util.py,sha256=9vzN3mbQqU6OhS7qRUiMRana4PTHa4sDTA0Hn7OVj
 truefoundry/cli/util.py,sha256=kEjC20-n_jwxZV9jq-78CxDk4xAySxAoYIXTxZfJzLM,5423
 truefoundry/common/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 truefoundry/common/auth_service_client.py,sha256=N3YxKlx63r6cPZqbgb2lqBOPI69ShB7D7RCIq4FSCjc,7949
-truefoundry/common/constants.py,sha256=nWd3Je71WmHEORRUTCupZy5fWADqEFftjYP6wiYhCIc,4627
+truefoundry/common/constants.py,sha256=3GH-KgEuBp-Gmaq_lreudKCKKPVPdEsGME7xdvIOpBA,4710
 truefoundry/common/credential_file_manager.py,sha256=1yEk1Zm2xS4G0VDFwKSZ4w0VUrcPWQ1nJnoBaz9xyKA,4251
 truefoundry/common/credential_provider.py,sha256=_OhJ2XFlDaVsrUO-FyywxctcGGqDdC2pgcvwEKqQD0Q,4071
 truefoundry/common/entities.py,sha256=b4R6ss06-ygDS3C4Tqa_GOq5LFKDYbt7x4Mghnfz6yo,4007
@@ -52,22 +52,24 @@ truefoundry/common/storage_provider_utils.py,sha256=yURhMw8k0FLFvaviRHDiifhvc6Gn
 truefoundry/common/types.py,sha256=BMJFCsR1lPJAw66IQBSvLyV4I6o_x5oj78gVsUa9si8,188
 truefoundry/common/utils.py,sha256=P0FuAadoJGdpieUORLSN-PiFnkyoGO-K2cS4OPITBWg,6714
 truefoundry/common/warnings.py,sha256=xDMhR_-ZGC40Ycaj6nlFb5MYPexn8WbKCHd4FlflTXQ,705
-truefoundry/deploy/__init__.py,sha256=PVbGPU9S3-dTFn5LvLwaEnfsp2RrGT9iiM7_15kOV84,2837
+truefoundry/deploy/__init__.py,sha256=sP-6Nv-_uV2o3knWcNSGV07j_Hkq0lfUkfZffBg-Hfo,2874
 truefoundry/deploy/python_deploy_codegen.py,sha256=k19_m5DGsUyjOUCSKwIVP8vDna2sq01tHABsUfoVpW4,8019
-truefoundry/deploy/_autogen/models.py,sha256=8sim4XOGFtK68pdZVKscvg6XPQRuGNTisUUSStqcP_A,75542
-truefoundry/deploy/builder/__init__.py,sha256=kgvlkVkiWpMVdim81tIeLrdoACqrFDgwCqHdQVsCsMo,4988
+truefoundry/deploy/_autogen/models.py,sha256=oOBwFb7qumBz1XXllns8wL02-NVvYq5iOOGZlEP3BzU,75893
+truefoundry/deploy/builder/__init__.py,sha256=VR07ZB7ziONEBbVgg1JdRTWY7t4qJjJTMhc2VodXYdA,5036
 truefoundry/deploy/builder/constants.py,sha256=amUkHoHvVKzGv0v_knfiioRuKiJM0V0xW0diERgWiI0,508
 truefoundry/deploy/builder/docker_service.py,sha256=sm7GWeIqyrKaZpxskdLejZlsxcZnM3BTDJr6orvPN4E,3948
-truefoundry/deploy/builder/utils.py,sha256=D68-bqM0NQx-Elg-56mtkENyVyg9faZ9tgTmBuo1Sjs,1076
-truefoundry/deploy/builder/builders/__init__.py,sha256=aomhWdR5L7uSM-GUalw9SnFHD2FQ_n-yFe4NH6nyNxw,715
+truefoundry/deploy/builder/utils.py,sha256=4TO0f3qMFGfFoBK0two1P59jgxlNjUoZYHYRgStcovM,3694
+truefoundry/deploy/builder/builders/__init__.py,sha256=Gp9NODR1E7mUjadhzIe3zzO43bBfHPeNcEDryYF2uo0,807
 truefoundry/deploy/builder/builders/dockerfile.py,sha256=XMbMlPUTMPCyaHl7jJQY1ODtlRkpI61PcvgG6Ck5jNc,1522
 truefoundry/deploy/builder/builders/tfy_notebook_buildpack/__init__.py,sha256=RGWGqY8xOF7vycUPJd10N7ZzahWv24lO0anrOPtLuDU,1796
 truefoundry/deploy/builder/builders/tfy_notebook_buildpack/dockerfile_template.py,sha256=rQgdvKmAT9HArVW4TAG5yd2QTKRs3S5LJ9RQbc_EkHE,2518
 truefoundry/deploy/builder/builders/tfy_python_buildpack/__init__.py,sha256=_fjqHKn80qKi68SAMMALge7_A6e1sTsQWichw8uoGIw,2025
-truefoundry/deploy/builder/builders/tfy_python_buildpack/dockerfile_template.py,sha256=f4l3fH21E2b8W3-JotMKc0AdPcCxV7LRPxxYJa7z_UQ,9134
+truefoundry/deploy/builder/builders/tfy_python_buildpack/dockerfile_template.py,sha256=Kj-ICGFTpDj86v6Juohz7q2TNYpcGIeKBW5HADG7SGE,6704
 truefoundry/deploy/builder/builders/tfy_spark_buildpack/__init__.py,sha256=NEPlM6_vTVxp4ITa18B8DBbgYCn1q5d8be21lbgu5oY,2888
-truefoundry/deploy/builder/builders/tfy_spark_buildpack/dockerfile_template.py,sha256=2zohUaW8Yw_QREHlpRW7Pooomt19HJh44fHjlsiDmwM,6064
-truefoundry/deploy/builder/builders/tfy_spark_buildpack/tfy_execute_notebook.py,sha256=QvawKw30dcHROJ05XQU2KgwH3gtUmEGSkuLxiuPNJ2c,5899
+truefoundry/deploy/builder/builders/tfy_spark_buildpack/dockerfile_template.py,sha256=nMJJfxjy8R7BZK89KicerQQwKLspUSJ3kerWZI3hFxk,4571
+truefoundry/deploy/builder/builders/tfy_spark_buildpack/tfy_execute_notebook.py,sha256=-D37Zjy2SBt3RHxonPEpR1_LR0W7vTSM1kQ1S-fdK-I,6363
+truefoundry/deploy/builder/builders/tfy_task_pyspark_buildpack/__init__.py,sha256=ynnjrFDg1__REd_x-npxxj-5zmFo46z_Ntz7GZ9-DHI,1819
+truefoundry/deploy/builder/builders/tfy_task_pyspark_buildpack/dockerfile_template.py,sha256=cjyPCLXJ8x5blaKbTK4XQ-4tO7DJqe5I9Fw2EJoLmtk,4555
 truefoundry/deploy/cli/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 truefoundry/deploy/cli/commands/__init__.py,sha256=qv818jxqSAygJ3h-6Ul8t-5VOgR_UrSgsVtNCl3e5G0,1408
 truefoundry/deploy/cli/commands/apply_command.py,sha256=DmXmKVokkauyKIiJDtErTwbJ5_LvQeJbTQsG5BjyKpo,2427
@@ -117,7 +119,7 @@ truefoundry/deploy/v2/lib/deploy.py,sha256=Ltm7cpIW14IbmEsR3EAIeWQUch2Z6HLej7heu
 truefoundry/deploy/v2/lib/deploy_workflow.py,sha256=G5BzMIbap8pgDX1eY-TITruUxQdkKhYtBmRwLL6lDeY,14342
 truefoundry/deploy/v2/lib/deployable_patched_models.py,sha256=mUi-OjPf7bc8rzfrPLdFb79LKuDq7F36RxL4V-AXebs,6830
 truefoundry/deploy/v2/lib/models.py,sha256=ogc1UYs1Z2nBdGSKCrde9sk8d0GxFKMkem99uqO5CmM,1148
-truefoundry/deploy/v2/lib/patched_models.py,sha256=TqGzHfHp20xdptmT--4LCWtFVjsWfE-LMT5wNS-fkuU,16751
+truefoundry/deploy/v2/lib/patched_models.py,sha256=bsznDLcUH5GcW8SUEvHETJqoFGlYYJ0j-tyGIqnRraw,16911
 truefoundry/deploy/v2/lib/source.py,sha256=d6-8_6Zn5koBglqrBrY6ZLG_7yyPuLdyEmK4iZTw6xY,9405
 truefoundry/ml/__init__.py,sha256=EEEHV7w58Krpo_W9Chd8Y3TdItfFO3LI6j6Izqc4-P8,2219
 truefoundry/ml/constants.py,sha256=vDq72d4C9FSWqr9MMdjgTF4TuyNFApvo_6RVsSeAjB4,2837
@@ -347,7 +349,7 @@ truefoundry/ml/_autogen/models/schema.py,sha256=a_bp42MMPUbwO3407m0UW2W8EOhnxZXf
 truefoundry/ml/_autogen/models/signature.py,sha256=rBjpxUIsEeWM0sIyYG5uCJB18DKHR4k5yZw8TzuoP48,4987
 truefoundry/ml/_autogen/models/utils.py,sha256=c7RtSLXhOLcP8rjuUtfnMdaKVTZvvbsmw98gPAkAFrs,24371
 truefoundry/ml/artifact/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-truefoundry/ml/artifact/truefoundry_artifact_repo.py,sha256=hbgLxSoihkLVuICzRueuh8iAIc-yruCW5TuMXYQ-aCU,35692
+truefoundry/ml/artifact/truefoundry_artifact_repo.py,sha256=8BFKaXDxutw8bPJLnDI0bO0oNS_xJKo2ijubc2PLFsU,35688
 truefoundry/ml/cli/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 truefoundry/ml/cli/cli.py,sha256=MwpY7z_NEeJE_XIP7XbZELjNeu2vpMmohttHCKDRk54,335
 truefoundry/ml/cli/utils.py,sha256=j6_mZ4Spn114mz3P4QQ8jx0tmorXIuyQnHXVUSDvZi4,1035
@@ -355,6 +357,9 @@ truefoundry/ml/cli/commands/__init__.py,sha256=diDUiRUX4l6TtNLI4iF-ZblczkELM7FRV
 truefoundry/ml/cli/commands/download.py,sha256=N9MhsEQ3U24v_OmnMZT8Q4SoAi38Sm7a21unrACOSDw,2573
 truefoundry/ml/cli/commands/model_init.py,sha256=INyUAU6hiFClI8cZqX5hgnrtNbeKxlZxrjFrjzStU18,2664
 truefoundry/ml/clients/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+truefoundry/ml/integrations/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+truefoundry/ml/integrations/huggingface/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+truefoundry/ml/integrations/huggingface/trainer_callback.py,sha256=Zu5AUbH_ct8I1dHyNYJQZBj9Y__hKo0sc2OxpPXJARE,6952
 truefoundry/ml/log_types/__init__.py,sha256=g4u4D4Jaj0aBK5GtrLV88-qThKZR9pSZ17vFEkN-LmM,125
 truefoundry/ml/log_types/plot.py,sha256=LDh4uy6z2P_a2oPM2lc85c0lt8utVvunohzeMawFjZw,7572
 truefoundry/ml/log_types/pydantic_base.py,sha256=eBlw_AEyAz4iJKDP4zgJOCFWcldwQqpf7FADW1jzIQY,272
@@ -381,7 +386,7 @@ truefoundry/workflow/remote_filesystem/__init__.py,sha256=LQ95ViEjJ7Ts4JcCGOxMPs
 truefoundry/workflow/remote_filesystem/logger.py,sha256=em2l7D6sw7xTLDP0kQSLpgfRRCLpN14Qw85TN7ujQcE,1022
 truefoundry/workflow/remote_filesystem/tfy_signed_url_client.py,sha256=xcT0wQmQlgzcj0nP3tJopyFSVWT1uv3nhiTIuwfXYeg,12342
 truefoundry/workflow/remote_filesystem/tfy_signed_url_fs.py,sha256=nSGPZu0Gyd_jz0KsEE-7w_BmnTD8CVF1S8cUJoxaCbc,13305
-truefoundry-0.11.3rc1.dist-info/METADATA,sha256=ukod0q3x_QWBS0jqVQFGdlYQvCTQxiKRxfvaIhDDzzg,2762
-truefoundry-0.11.3rc1.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
-truefoundry-0.11.3rc1.dist-info/entry_points.txt,sha256=xVjn7RMN-MW2-9f7YU-bBdlZSvvrwzhpX1zmmRmsNPU,98
-truefoundry-0.11.3rc1.dist-info/RECORD,,
+truefoundry-0.11.4.dist-info/METADATA,sha256=RD0XhZ5hvcV7BAguapQ9yYssfoEDIXwNUi11w5riKtc,2759
+truefoundry-0.11.4.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
+truefoundry-0.11.4.dist-info/entry_points.txt,sha256=xVjn7RMN-MW2-9f7YU-bBdlZSvvrwzhpX1zmmRmsNPU,98
+truefoundry-0.11.4.dist-info/RECORD,,