truefoundry 0.11.11rc1__py3-none-any.whl → 0.11.12__py3-none-any.whl

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

@@ -1,15 +1,15 @@
-from typing import List, Optional
-
 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_secret_mounts,
     generate_uv_pip_install_command,
-    get_available_secrets,
 )
 from truefoundry.deploy.v2.lib.patched_models import (
     _resolve_requirements_path,
@@ -57,13 +57,8 @@ def get_additional_pip_packages(build_configuration: TaskPySparkBuild):
 def generate_dockerfile_content(
     build_configuration: TaskPySparkBuild,
     package_manager: str = ENV_VARS.TFY_PYTHON_BUILD_PACKAGE_MANAGER,
-    docker_build_extra_args: Optional[List[str]] = None,
+    mount_python_package_manager_conf_secret: bool = False,
 ) -> str:
-    # Get available secrets from docker build extra args
-    available_secrets = set()
-    if docker_build_extra_args:
-        available_secrets = get_available_secrets(docker_build_extra_args)
-
     # TODO (chiragjn): Handle recursive references to other requirements files e.g. `-r requirements-gpu.txt`
     requirements_path = _resolve_requirements_path(
         build_context_path="",
@@ -83,13 +78,13 @@ def generate_dockerfile_content(
         python_packages_install_command = generate_pip_install_command(
             requirements_path=requirements_destination_path,
             pip_packages=pip_packages,
-            available_secrets=available_secrets,
+            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,
-            available_secrets=available_secrets,
+            mount_uv_conf_secret=mount_python_package_manager_conf_secret,
         )
     else:
         raise ValueError(f"Unsupported package manager: {package_manager}")
@@ -106,12 +101,17 @@ def generate_dockerfile_content(
         "python_packages_install_command": python_packages_install_command,
     }
 
-    if available_secrets:
-        template_args["package_manager_config_secret_mount"] = generate_secret_mounts(
-            available_secrets=available_secrets,
-            python_dependencies_type="pip",
-            package_manager=package_manager,
-        )
+    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"] = ""
 

truefoundry/deploy/builder/constants.py

@@ -13,15 +13,3 @@ UV_CONF_BUILDKIT_SECRET_MOUNT = (
 UV_CONF_SECRET_MOUNT_AS_ENV = (
     f"UV_CONFIG_FILE=/run/secrets/{BUILDKIT_SECRET_MOUNT_UV_CONF_ID}"
 )
-
-BUILDKIT_SECRET_MOUNT_UV_ENV_ID = "uv.env"
-UV_ENV_BUILDKIT_SECRET_MOUNT = (
-    f"--mount=type=secret,id={BUILDKIT_SECRET_MOUNT_UV_ENV_ID}"
-)
-UV_ENV_SECRET_MOUNT_AS_ENV = f". /run/secrets/{BUILDKIT_SECRET_MOUNT_UV_ENV_ID}"
-
-BUILDKIT_SECRET_MOUNT_POETRY_ENV_ID = "poetry.env"
-POETRY_ENV_BUILDKIT_SECRET_MOUNT = (
-    f"--mount=type=secret,id={BUILDKIT_SECRET_MOUNT_POETRY_ENV_ID}"
-)
-POETRY_ENV_SECRET_MOUNT_AS_ENV = f". /run/secrets/{BUILDKIT_SECRET_MOUNT_POETRY_ENV_ID}"

truefoundry/deploy/builder/utils.py

@@ -1,20 +1,12 @@
 import shlex
-from typing import List, Optional, Set
+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_POETRY_ENV_ID,
     BUILDKIT_SECRET_MOUNT_UV_CONF_ID,
-    BUILDKIT_SECRET_MOUNT_UV_ENV_ID,
-    PIP_CONF_BUILDKIT_SECRET_MOUNT,
     PIP_CONF_SECRET_MOUNT_AS_ENV,
-    POETRY_ENV_BUILDKIT_SECRET_MOUNT,
-    POETRY_ENV_SECRET_MOUNT_AS_ENV,
-    UV_CONF_BUILDKIT_SECRET_MOUNT,
     UV_CONF_SECRET_MOUNT_AS_ENV,
-    UV_ENV_BUILDKIT_SECRET_MOUNT,
-    UV_ENV_SECRET_MOUNT_AS_ENV,
 )
 
 
@@ -36,129 +28,29 @@ def _get_id_from_buildkit_secret_value(value: str) -> Optional[str]:
 
 def has_python_package_manager_conf_secret(docker_build_extra_args: List[str]) -> bool:
     args = [arg.strip() for arg in docker_build_extra_args]
-    for i, arg in enumerate(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])
-                in (
-                    BUILDKIT_SECRET_MOUNT_PIP_CONF_ID,
-                    BUILDKIT_SECRET_MOUNT_UV_CONF_ID,
-                    BUILDKIT_SECRET_MOUNT_UV_ENV_ID,
-                    BUILDKIT_SECRET_MOUNT_POETRY_ENV_ID,
-                )
+                in (BUILDKIT_SECRET_MOUNT_PIP_CONF_ID, BUILDKIT_SECRET_MOUNT_UV_CONF_ID)
             )
         ):
             return True
     return False
 
 
-def get_available_secrets(docker_build_extra_args: List[str]) -> Set[str]:
-    available_secrets = set()
-    args = [arg.strip() for arg in docker_build_extra_args]
-    for i, arg in enumerate(args):
-        if arg == "--secret" and i + 1 < len(args):
-            secret_id = _get_id_from_buildkit_secret_value(args[i + 1])
-            if secret_id:
-                available_secrets.add(secret_id)
-    return available_secrets
-
-
-def generate_secret_mounts(
-    available_secrets: Set[str], python_dependencies_type: str, package_manager: str
-) -> str:
-    mounts = []
-
-    if python_dependencies_type == "pip":
-        if (
-            package_manager == "pip"
-            and BUILDKIT_SECRET_MOUNT_PIP_CONF_ID in available_secrets
-        ):
-            mounts.append(PIP_CONF_BUILDKIT_SECRET_MOUNT)
-        elif (
-            package_manager == "uv"
-            and BUILDKIT_SECRET_MOUNT_UV_CONF_ID in available_secrets
-        ):
-            mounts.append(UV_CONF_BUILDKIT_SECRET_MOUNT)
-    elif python_dependencies_type == "uv":
-        if BUILDKIT_SECRET_MOUNT_UV_CONF_ID in available_secrets:
-            mounts.append(UV_CONF_BUILDKIT_SECRET_MOUNT)
-        if BUILDKIT_SECRET_MOUNT_UV_ENV_ID in available_secrets:
-            mounts.append(UV_ENV_BUILDKIT_SECRET_MOUNT)
-    elif python_dependencies_type == "poetry":
-        if BUILDKIT_SECRET_MOUNT_POETRY_ENV_ID in available_secrets:
-            mounts.append(POETRY_ENV_BUILDKIT_SECRET_MOUNT)
-
-    return " ".join(mounts)
-
-
-def generate_secret_env_commands(
-    available_secrets: Set[str], python_dependencies_type: str, package_manager: str
-) -> List[str]:
-    env_commands = []
-
-    if python_dependencies_type == "pip":
-        if (
-            package_manager == "pip"
-            and BUILDKIT_SECRET_MOUNT_PIP_CONF_ID in available_secrets
-        ):
-            env_commands.append(PIP_CONF_SECRET_MOUNT_AS_ENV)
-        elif (
-            package_manager == "uv"
-            and BUILDKIT_SECRET_MOUNT_UV_CONF_ID in available_secrets
-        ):
-            env_commands.append(UV_CONF_SECRET_MOUNT_AS_ENV)
-    elif python_dependencies_type == "uv":
-        if BUILDKIT_SECRET_MOUNT_UV_CONF_ID in available_secrets:
-            env_commands.append(UV_CONF_SECRET_MOUNT_AS_ENV)
-        if BUILDKIT_SECRET_MOUNT_UV_ENV_ID in available_secrets:
-            env_commands.append(UV_ENV_SECRET_MOUNT_AS_ENV)
-    elif python_dependencies_type == "poetry":
-        if BUILDKIT_SECRET_MOUNT_POETRY_ENV_ID in available_secrets:
-            env_commands.append(POETRY_ENV_SECRET_MOUNT_AS_ENV)
-
-    return env_commands
-
-
-def generate_shell_command_with_secrets(
-    env_commands: List[str], command: List[str]
-) -> str:
-    """Generate a shell command that properly handles source commands and environment variables."""
-    if not env_commands:
-        return shlex.join(command)
-
-    # Separate source commands from env var assignments
-    source_commands = [
-        cmd for cmd in env_commands if cmd.startswith(". ") or cmd.startswith("source ")
-    ]
-    env_assignments = [
-        cmd
-        for cmd in env_commands
-        if not cmd.startswith(". ") and not cmd.startswith("source ")
-    ]
-
-    if source_commands:
-        # If we have source commands, join them with the command
-        shell_command_parts = source_commands.copy()
-
-        # Add environment variable assignments and command together
-        if env_assignments:
-            final_command = shlex.join(env_assignments + command)
-        else:
-            final_command = shlex.join(command)
-
-        shell_command_parts.append(final_command)
-        return " && ".join(shell_command_parts)
-
-    return shlex.join(env_assignments + command)
-
-
 def generate_pip_install_command(
     requirements_path: Optional[str],
     pip_packages: Optional[List[str]],
-    available_secrets: Optional[Set[str]] = None,
+    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:
@@ -171,37 +63,27 @@ def generate_pip_install_command(
     if not args:
         return None
 
-    secret_env_commands = []
-    if available_secrets:
-        secret_env_commands = generate_secret_env_commands(
-            available_secrets, python_dependencies_type="pip", package_manager="pip"
-        )
-
-    final_pip_install_command = generate_shell_command_with_secrets(
-        secret_env_commands, command + args
+    final_pip_install_command = shlex.join(envs + command + args)
+    final_docker_run_command = " && ".join(
+        [upgrade_pip_command, final_pip_install_command]
     )
-
-    return final_pip_install_command
+    return final_docker_run_command
 
 
 def generate_uv_pip_install_command(
     requirements_path: Optional[str],
     pip_packages: Optional[List[str]],
-    available_secrets: Optional[Set[str]] = None,
+    mount_uv_conf_secret: bool = False,
 ) -> Optional[str]:
-    uv_mount = f"--mount=from={ENV_VARS.TFY_PYTHON_BUILD_UV_IMAGE_REPO}:{ENV_VARS.TFY_PYTHON_BUILD_UV_IMAGE_TAG},source=/uv,target=/usr/local/bin/uv"
-
+    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",
     ]
-
-    secret_env_commands = []
-    if available_secrets:
-        secret_env_commands = generate_secret_env_commands(
-            available_secrets, python_dependencies_type="pip", package_manager="uv"
-        )
+    if mount_uv_conf_secret:
+        envs.append(UV_CONF_SECRET_MOUNT_AS_ENV)
 
     command = ["uv", "pip", "install", "--no-cache-dir"]
 
@@ -217,10 +99,9 @@ def generate_uv_pip_install_command(
     if not args:
         return None
 
-    uv_pip_install_command = generate_shell_command_with_secrets(
-        secret_env_commands, envs + command + args
-    )
-    final_docker_run_command = " ".join([uv_mount, uv_pip_install_command])
+    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
 
@@ -237,73 +118,3 @@ def generate_apt_install_command(apt_packages: Optional[List[str]]) -> Optional[
     return " && ".join(
         [apt_update_command, apt_install_command, clear_apt_lists_command]
     )
-
-
-def generate_command_to_install_from_uv_lock(
-    sync_options: Optional[str],
-    uv_version: Optional[str],
-    install_project: bool = False,
-    available_secrets: Optional[Set[str]] = None,
-):
-    uv_image_uri = f"{ENV_VARS.TFY_PYTHON_BUILD_UV_IMAGE_REPO}:{uv_version if uv_version is not None else ENV_VARS.TFY_PYTHON_BUILD_UV_IMAGE_TAG}"
-    uv_mount = f"--mount=from={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",
-    ]
-
-    secret_env_commands = []
-    if available_secrets:
-        secret_env_commands = generate_secret_env_commands(
-            available_secrets, python_dependencies_type="uv", package_manager="uv"
-        )
-
-    command = ["uv", "sync"]
-    sync_options_list = shlex.split(sync_options or "")
-    if "--active" not in sync_options_list:
-        sync_options_list.append("--active")
-
-    if not install_project and "--no-install-project" not in sync_options_list:
-        sync_options_list.append("--no-install-project")
-
-    command.extend(sync_options_list)
-
-    uv_sync_install_command = generate_shell_command_with_secrets(
-        secret_env_commands, envs + command
-    )
-    final_docker_run_command = " ".join([uv_mount, uv_sync_install_command])
-
-    return final_docker_run_command
-
-
-def generate_poetry_install_command(
-    install_options: Optional[str],
-    install_project: bool = False,
-    available_secrets: Optional[Set[str]] = None,
-) -> Optional[str]:
-    command = ["poetry", "install"]
-    install_options_list = shlex.split(install_options or "")
-
-    if "--no-interaction" not in install_options_list:
-        command.append("--no-interaction")
-
-    if not install_project and "--no-root" not in install_options_list:
-        command.append("--no-root")
-
-    command.extend(install_options_list)
-
-    secret_env_commands = []
-    if available_secrets:
-        secret_env_commands = generate_secret_env_commands(
-            available_secrets=available_secrets,
-            python_dependencies_type="poetry",
-            package_manager="poetry",
-        )
-
-    poetry_install_cmd = generate_shell_command_with_secrets(
-        secret_env_commands, command
-    )
-
-    return poetry_install_cmd

truefoundry/deploy/lib/util.py

@@ -1,11 +1,6 @@
-import os
 import re
 from typing import Union
 
-from truefoundry.deploy._autogen.models import (
-    PythonBuild,
-)
-
 
 def get_application_fqn_from_deployment_fqn(deployment_fqn: str) -> str:
     if not re.search(r":\d+$", deployment_fqn):
@@ -34,33 +29,3 @@ def find_list_paths(data, parent_key="", sep="."):
             new_key = f"{parent_key}[{i}]"
             list_paths.extend(find_list_paths(value, new_key, sep))
     return list_paths
-
-
-def _validate_file_path(project_root_path: str, file_path: str):
-    file_absolute_path = os.path.join(project_root_path, file_path)
-    if not os.path.exists(file_absolute_path):
-        raise FileNotFoundError(
-            f"file {file_path} not found. file path should be relative to your project root path: {project_root_path}."
-        )
-
-
-def validate_paths(python_build: PythonBuild, source_dir: str):
-    if not python_build.python_dependencies:
-        return
-
-    if not os.path.exists(source_dir):
-        raise ValueError(f"project root path {source_dir!r} of does not exist")
-    if (
-        python_build.python_dependencies.type == "pip"
-        and python_build.python_dependencies.requirements_path
-    ):
-        _validate_file_path(
-            source_dir,
-            python_build.python_dependencies.requirements_path,
-        )
-    if python_build.python_dependencies.type == "uv":
-        _validate_file_path(source_dir, "uv.lock")
-        _validate_file_path(source_dir, "pyproject.toml")
-    if python_build.python_dependencies.type == "poetry":
-        _validate_file_path(source_dir, "pyproject.toml")
-        _validate_file_path(source_dir, "poetry.lock")

truefoundry/deploy/v2/lib/deploy.py

@@ -13,9 +13,7 @@ from truefoundry.deploy.lib.clients.servicefoundry_client import (
 )
 from truefoundry.deploy.lib.dao.workspace import get_workspace_by_fqn
 from truefoundry.deploy.lib.model.entity import Deployment, DeploymentTransitionStatus
-from truefoundry.deploy.lib.util import (
-    get_application_fqn_from_deployment_fqn,
-)
+from truefoundry.deploy.lib.util import get_application_fqn_from_deployment_fqn
 from truefoundry.deploy.v2.lib.models import BuildResponse
 from truefoundry.deploy.v2.lib.source import (
     local_source_to_image,

truefoundry/deploy/v2/lib/patched_models.py

@@ -24,7 +24,7 @@ build=Build(
     ...
     build_spec=PythonBuild(
         ...
-        python_dependencies=Pip(requirements_path={requirements_txt_path!r})
+        requirements_path={requirements_txt_path!r}
     )
 )
 ```
@@ -36,44 +36,12 @@ build:
   type: build
   build_spec:
     type: tfy-python-buildpack
-    python_dependencies:
-      type: pip
-      requirements_path: {requirements_txt_path!r}
+    requirements_path: {requirements_txt_path!r}
     ...
 ...
 ```
 
-or set it to None if you don't want to use any requirements file.
-"""
-
-SPECS_UPGRADE_WARNING_MESSAGE_TEMPLATE = """\
-The `requirements_path` and `pip_packages` fields are deprecated and will be removed in a future release.
-Please use the `python_dependencies` field instead. Please use the following format:
-
-```python
-build=Build(
-    ...
-    build_spec=PythonBuild(
-        ...
-        python_dependencies=Pip(requirements_path={requirements_txt_path!r}, pip_packages={pip_packages!r})
-    )
-)
-```
-
-OR
-
-```yaml
-build:
-  type: build
-  build_spec:
-    type: tfy-python-buildpack
-    python_dependencies:
-      type: pip
-      requirements_path: {requirements_txt_path!r}
-      pip_packages: {pip_packages!r}
-    ...
-...
-```
+or set it to None if you don't want use any requirements file.
 """
 
 
@@ -106,18 +74,6 @@ def _resolve_requirements_path(
     return None
 
 
-def check_whether_poetry_toml_exists(
-    build_context_path: str,
-) -> bool:
-    required_filename = "poetry.toml"
-    possible_path = os.path.join(build_context_path, required_filename)
-
-    if os.path.isfile(possible_path):
-        return True
-
-    return False
-
-
 class CUDAVersion(str, enum.Enum):
     CUDA_11_0_CUDNN8 = "11.0-cudnn8"
     CUDA_11_1_CUDNN8 = "11.1-cudnn8"
@@ -209,23 +165,6 @@ class PythonBuild(models.PythonBuild, PatchedModelBase):
             build_context_path=values.get("build_context_path") or "./",
             requirements_path=values.get("requirements_path"),
         )
-
-        if not values.get("python_dependencies"):
-            warnings.warn(
-                SPECS_UPGRADE_WARNING_MESSAGE_TEMPLATE.format(
-                    requirements_txt_path=values.get("requirements_path"),
-                    pip_packages=values.get("pip_packages"),
-                ),
-                category=TrueFoundryDeprecationWarning,
-                stacklevel=2,
-            )
-            values["python_dependencies"] = Pip(
-                type="pip",
-                requirements_path=values.get("requirements_path"),
-                pip_packages=values.get("pip_packages"),
-            )
-            values.pop("pip_packages", None)
-            values.pop("requirements_path", None)
         return values
 
 
@@ -635,15 +574,3 @@ class SparkExecutorDynamicScaling(models.SparkExecutorDynamicScaling, PatchedMod
 
 class TaskPySparkBuild(models.TaskPySparkBuild, PatchedModelBase):
     type: Literal["task-pyspark-build"] = "task-pyspark-build"
-
-
-class Pip(models.Pip, PatchedModelBase):
-    type: Literal["pip"] = "pip"
-
-
-class UV(models.UV, PatchedModelBase):
-    type: Literal["uv"] = "uv"
-
-
-class Poetry(models.Poetry, PatchedModelBase):
-    type: Literal["poetry"] = "poetry"

truefoundry/deploy/v2/lib/source.py

@@ -21,7 +21,6 @@ from truefoundry.deploy.lib.clients.servicefoundry_client import (
     ServiceFoundryServiceClient,
 )
 from truefoundry.deploy.lib.dao import workspace as workspace_lib
-from truefoundry.deploy.lib.util import validate_paths
 from truefoundry.deploy.v2.lib.patched_models import Image, RemoteSource
 from truefoundry.logger import logger
 
@@ -210,9 +209,6 @@ def local_source_to_image(
                 f"of component {component_name!r} does not exist"
             )
 
-    if isinstance(build.build_spec, models.PythonBuild):
-        validate_paths(build.build_spec, source_dir)
-
     client = ServiceFoundryServiceClient()
 
     workspace = workspace_lib.get_workspace_by_fqn(workspace_fqn=workspace_fqn)

truefoundry/ml/artifact/truefoundry_artifact_repo.py

@@ -40,6 +40,7 @@ from truefoundry.common.storage_provider_utils import (
     azure_multi_part_upload,
     decide_file_parts,
     s3_compatible_multipart_upload,
+    truncate_path_for_progress,
 )
 from truefoundry.ml._autogen.client import (  # type: ignore[attr-defined]
     ApiClient,
@@ -125,14 +126,14 @@ def _signed_url_upload_file(
             augmented_raise_for_status(response, exception_class=MlFoundryException)  # type: ignore
         return
 
-    task_progress_bar = progress_bar.add_task(
-        f"[green]Uploading {local_file}:", start=True
+    task_id = progress_bar.add_task(
+        f"[green]⬆ {truncate_path_for_progress(local_file, 64, relpath=True)}",
+        start=True,
+        visible=True,
     )
 
     def callback(length):
-        progress_bar.update(
-            task_progress_bar, advance=length, total=os.stat(local_file).st_size
-        )
+        progress_bar.update(task_id, advance=length, total=os.stat(local_file).st_size)
         if abort_event and abort_event.is_set():
             raise Exception("aborting upload")
 
@@ -147,6 +148,9 @@ def _signed_url_upload_file(
         ) as response:
             augmented_raise_for_status(response, exception_class=MlFoundryException)  # type: ignore
 
+    if progress_bar is not None:
+        progress_bar.refresh()
+
 
 def _download_file_using_http_uri(
     http_uri,
@@ -167,7 +171,11 @@ def _download_file_using_http_uri(
         exception_class=MlFoundryException,  # type: ignore
     ) as response:
         augmented_raise_for_status(response, exception_class=MlFoundryException)  # type: ignore
-        file_size = int(response.headers.get("Content-Length", 0))
+        file_size = int(response.headers.get("Content-Length", -1))
+        if file_size == 0 and callback:
+            # special case for empty files
+            callback(1, 1)
+        file_size = file_size if file_size > 0 else 0
         with open(download_path, "wb") as output_file:
             for chunk in response.iter_content(chunk_size=chunk_size):
                 if callback:
@@ -658,15 +666,18 @@ class MlFoundryArtifactsRepository:
             logger.info("Downloading %s to %s", remote_file_path, local_path)
 
         if progress_bar is not None:
-            download_progress_bar = progress_bar.add_task(
-                f"[green]Downloading to {remote_file_path}:", start=True
+            task_id = progress_bar.add_task(
+                f"[green]⬇ {truncate_path_for_progress(remote_file_path, 64)}",
+                start=True,
+                visible=True,
             )
 
-        def callback(chunk, total_file_size):
+        def callback(chunk_size: int, total_file_size: int):
+            nonlocal task_id
             if progress_bar is not None:
                 progress_bar.update(
-                    download_progress_bar,
-                    advance=chunk,
+                    task_id,
+                    advance=chunk_size,
                     total=total_file_size,
                 )
             if abort_event and abort_event.is_set():
@@ -677,6 +688,9 @@ class MlFoundryArtifactsRepository:
             download_path=local_path,
             callback=callback,
         )
+
+        if progress_bar is not None:
+            progress_bar.refresh()
         logger.debug("Downloaded %s to %s", remote_file_path, local_path)
 
     def _download_artifact(

truefoundry/ml/log_types/artifacts/artifact.py

@@ -347,10 +347,11 @@ class ArtifactVersion:
             artifact_version.download(path="<your-desired-download-path>")
             ```
         """
-        download_path = self._download(
+        download_info = self._download(
             path=path, overwrite=overwrite, progress=progress
         )
-        return download_path.download_dir
+        logger.info("Downloaded artifact contents to %s", download_info.download_dir)
+        return download_info.download_dir
 
     def delete(self) -> bool:
         """

truefoundry/ml/log_types/artifacts/model.py

@@ -440,6 +440,7 @@ class ModelVersion:
         download_info = self._download(
             path=path, overwrite=overwrite, progress=progress
         )
+        logger.info("Downloaded model contents to %s", download_info.model_dir)
         return download_info
 
     def delete(self) -> bool: