truefoundry 0.4.3__py3-none-any.whl → 0.4.4__py3-none-any.whl

truefoundry/ml/artifact/truefoundry_artifact_repo.py

@@ -16,6 +16,7 @@ from typing import (
     List,
     NamedTuple,
     Optional,
+    Sequence,
     Tuple,
     Union,
 )
@@ -34,6 +35,7 @@ from rich.progress import (
 )
 from tqdm.utils import CallbackIOWrapper
 
+from truefoundry.common.constants import ENV_VARS
 from truefoundry.ml.autogen.client import (  # type: ignore[attr-defined]
     ApiClient,
     CreateMultiPartUploadForDatasetRequestDto,
@@ -58,14 +60,12 @@ from truefoundry.ml.clients.utils import (
     augmented_raise_for_status,
     cloud_storage_http_request,
 )
-from truefoundry.ml.env_vars import DISABLE_MULTIPART_UPLOAD
 from truefoundry.ml.exceptions import MlFoundryException
 from truefoundry.ml.logger import logger
 from truefoundry.ml.session import _get_api_client
 from truefoundry.pydantic_v1 import BaseModel, root_validator
 
 _MIN_BYTES_REQUIRED_FOR_MULTIPART = 100 * 1024 * 1024
-_MULTIPART_DISABLED = os.getenv(DISABLE_MULTIPART_UPLOAD, "").lower() == "true"
 # GCP/S3 Maximum number of parts per upload	10,000
 # Maximum number of blocks in a block blob 50,000 blocks
 # TODO: This number is artificially limited now. Later
@@ -84,21 +84,11 @@ _MAX_NUM_PARTS_FOR_MULTIPART = 1000
 # Azure Maximum size of a block in a block blob	4000 MiB
 # GCP/S3 Maximum size of an individual part in a multipart upload 5 GiB
 _MAX_PART_SIZE_BYTES_FOR_MULTIPART = 4 * 1024 * 1024 * 1000
-_cpu_count = os.cpu_count() or 2
-_MAX_WORKERS_FOR_UPLOAD = max(min(32, _cpu_count * 2), 4)
-_MAX_WORKERS_FOR_DOWNLOAD = max(min(32, _cpu_count * 2), 4)
 _LIST_FILES_PAGE_SIZE = 500
 _GENERATE_SIGNED_URL_BATCH_SIZE = 50
 DEFAULT_PRESIGNED_URL_EXPIRY_TIME = 3600
 
 
-def _get_relpath_if_in_tempdir(path: str) -> str:
-    tempdir = tempfile.gettempdir()
-    if path.startswith(tempdir):
-        return os.path.relpath(path, tempdir)
-    return path
-
-
 def _can_display_progress(user_choice: Optional[bool] = None) -> bool:
     if user_choice is False:
         return False
@@ -199,7 +189,10 @@ class _FileMultiPartInfo(NamedTuple):
 
 def _decide_file_parts(file_path: str) -> _FileMultiPartInfo:
     file_size = os.path.getsize(file_path)
-    if file_size < _MIN_BYTES_REQUIRED_FOR_MULTIPART or _MULTIPART_DISABLED:
+    if (
+        file_size < _MIN_BYTES_REQUIRED_FOR_MULTIPART
+        or ENV_VARS.TFY_ARTIFACTS_DISABLE_MULTIPART_UPLOAD
+    ):
         return _FileMultiPartInfo(1, part_size=file_size, file_size=file_size)
 
     ideal_num_parts = math.ceil(file_size / _PART_SIZE_BYTES_FOR_MULTIPART)
@@ -240,7 +233,7 @@ def _signed_url_upload_file(
         return
 
     task_progress_bar = progress_bar.add_task(
-        f"[green]Uploading {_get_relpath_if_in_tempdir(local_file)}:", start=True
+        f"[green]Uploading {local_file}:", start=True
     )
 
     def callback(length):
@@ -262,7 +255,7 @@ def _signed_url_upload_file(
 def _download_file_using_http_uri(
     http_uri,
     download_path,
-    chunk_size=100000000,
+    chunk_size=ENV_VARS.TFY_ARTIFACTS_DOWNLOAD_CHUNK_SIZE_BYTES,
     callback: Optional[Callable[[int, int], Any]] = None,
 ):
     """
@@ -283,6 +276,9 @@ def _download_file_using_http_uri(
                 if not chunk:
                     break
                 output_file.write(chunk)
+                if ENV_VARS.TFY_ARTIFACTS_DOWNLOAD_FSYNC_CHUNKS:
+                    output_file.flush()
+                    os.fsync(output_file.fileno())
 
 
 class _CallbackIOWrapperForMultiPartUpload(CallbackIOWrapper):
@@ -338,7 +334,7 @@ def _s3_compatible_multipart_upload(
     parts = []
 
     multi_part_upload_progress = progress_bar.add_task(
-        f"[green]Uploading {_get_relpath_if_in_tempdir(local_file)}:", start=True
+        f"[green]Uploading {local_file}:", start=True
     )
 
     def upload(part_number: int, seek: int) -> None:
@@ -408,7 +404,7 @@ def _azure_multi_part_upload(
     abort_event = abort_event or Event()
 
     multi_part_upload_progress = progress_bar.add_task(
-        f"[green]Uploading {_get_relpath_if_in_tempdir(local_file)}:", start=True
+        f"[green]Uploading {local_file}:", start=True
     )
 
     def upload(part_number: int, seek: int):
@@ -508,164 +504,56 @@ class MlFoundryArtifactsRepository:
             api_client=self._api_client
         )
 
-    def _create_download_destination(
-        self, src_artifact_path, dst_local_dir_path=None
-    ) -> str:
-        """
-        Creates a local filesystem location to be used as a destination for downloading the artifact
-        specified by `src_artifact_path`. The destination location is a subdirectory of the
-        specified `dst_local_dir_path`, which is determined according to the structure of
-        `src_artifact_path`. For example, if `src_artifact_path` is `dir1/file1.txt`, then the
-        resulting destination path is `<dst_local_dir_path>/dir1/file1.txt`. Local directories are
-        created for the resulting destination location if they do not exist.
-
-        :param src_artifact_path: A relative, POSIX-style path referring to an artifact stored
-                                  within the repository's artifact root location.
-                                  `src_artifact_path` should be specified relative to the
-                                  repository's artifact root location.
-        :param dst_local_dir_path: The absolute path to a local filesystem directory in which the
-                                   local destination path will be contained. The local destination
-                                   path may be contained in a subdirectory of `dst_root_dir` if
-                                   `src_artifact_path` contains subdirectories.
-        :return: The absolute path to a local filesystem location to be used as a destination
-                 for downloading the artifact specified by `src_artifact_path`.
-        """
-        src_artifact_path = src_artifact_path.rstrip(
-            "/"
-        )  # Ensure correct dirname for trailing '/'
-        dirpath = posixpath.dirname(src_artifact_path)
-        local_dir_path = os.path.join(dst_local_dir_path, dirpath)
-        local_file_path = os.path.join(dst_local_dir_path, src_artifact_path)
-        if not os.path.exists(local_dir_path):
-            os.makedirs(local_dir_path, exist_ok=True)
-        return local_file_path
-
-    # these methods should be named list_files, log_directory, log_file, etc
-    def list_artifacts(
-        self, path=None, page_size=_LIST_FILES_PAGE_SIZE, **kwargs
-    ) -> Iterator[FileInfoDto]:
-        page_token = None
-        started = False
-        while not started or page_token is not None:
-            started = True
-            page = self.list_files(
-                artifact_identifier=self.artifact_identifier,
-                path=path,
-                page_size=page_size,
-                page_token=page_token,
+    # TODO (chiragjn): Refactor these methods - if else is very inconvenient
+    def get_signed_urls_for_read(
+        self,
+        artifact_identifier: ArtifactIdentifier,
+        paths,
+    ) -> List[SignedURLDto]:
+        if artifact_identifier.artifact_version_id:
+            signed_urls_response = self._mlfoundry_artifacts_api.get_signed_urls_for_read_post(
+                get_signed_urls_for_artifact_version_read_request_dto=GetSignedURLsForArtifactVersionReadRequestDto(
+                    id=str(artifact_identifier.artifact_version_id), paths=paths
+                )
             )
-            for file_info in page.files:
-                yield file_info
-            page_token = page.next_page_token
-
-    def log_artifacts(  # noqa: C901
-        self, local_dir, artifact_path=None, progress=None
-    ):
-        show_progress = _can_display_progress(progress)
-
-        dest_path = artifact_path or ""
-        dest_path = dest_path.lstrip(posixpath.sep)
-
-        files_for_normal_upload: List[Tuple[str, str, _FileMultiPartInfo]] = []
-        files_for_multipart_upload: List[Tuple[str, str, _FileMultiPartInfo]] = []
-
-        for root, _, file_names in os.walk(local_dir):
-            upload_path = dest_path
-            if root != local_dir:
-                rel_path = os.path.relpath(root, local_dir)
-                rel_path = relative_path_to_artifact_path(rel_path)
-                upload_path = posixpath.join(dest_path, rel_path)
-            for file_name in file_names:
-                local_file = os.path.join(root, file_name)
-                multipart_info = _decide_file_parts(local_file)
-
-                final_upload_path = upload_path or ""
-                final_upload_path = final_upload_path.lstrip(posixpath.sep)
-                final_upload_path = posixpath.join(
-                    final_upload_path, os.path.basename(local_file)
+            signed_urls = signed_urls_response.signed_urls
+        elif artifact_identifier.dataset_fqn:
+            signed_urls_dataset_response = self._mlfoundry_artifacts_api.get_signed_urls_dataset_read_post(
+                get_signed_urls_for_dataset_read_request_dto=GetSignedURLsForDatasetReadRequestDto(
+                    dataset_fqn=artifact_identifier.dataset_fqn, paths=paths
                 )
+            )
+            signed_urls = signed_urls_dataset_response.signed_urls
+        else:
+            raise ValueError(
+                "Invalid artifact type - both `artifact_version_id` and `dataset_fqn` both are None"
+            )
+        return signed_urls
 
-                if multipart_info.num_parts == 1:
-                    files_for_normal_upload.append(
-                        (final_upload_path, local_file, multipart_info)
-                    )
-                else:
-                    files_for_multipart_upload.append(
-                        (final_upload_path, local_file, multipart_info)
-                    )
-
-        abort_event = Event()
-
-        with Progress(
-            "[progress.description]{task.description}",
-            BarColumn(bar_width=None),
-            "[progress.percentage]{task.percentage:>3.0f}%",
-            DownloadColumn(),
-            TransferSpeedColumn(),
-            TimeRemainingColumn(),
-            TimeElapsedColumn(),
-            refresh_per_second=1,
-            disable=not show_progress,
-            expand=True,
-        ) as progress_bar, ThreadPoolExecutor(
-            max_workers=_MAX_WORKERS_FOR_UPLOAD
-        ) as executor:
-            futures: List[Future] = []
-            # Note: While this batching is beneficial when there is a large number of files, there is also
-            # a rare case risk of the signed url expiring before a request is made to it
-            _batch_size = _GENERATE_SIGNED_URL_BATCH_SIZE
-            for start_idx in range(0, len(files_for_normal_upload), _batch_size):
-                end_idx = min(start_idx + _batch_size, len(files_for_normal_upload))
-                if _any_future_has_failed(futures):
-                    break
-                logger.debug("Generating write signed urls for a batch ...")
-                remote_file_paths = [
-                    files_for_normal_upload[idx][0] for idx in range(start_idx, end_idx)
-                ]
-                signed_urls = self.get_signed_urls_for_write(
-                    artifact_identifier=self.artifact_identifier,
-                    paths=remote_file_paths,
+    def get_signed_urls_for_write(
+        self,
+        artifact_identifier: ArtifactIdentifier,
+        paths: List[str],
+    ) -> List[SignedURLDto]:
+        if artifact_identifier.artifact_version_id:
+            signed_urls_response = self._mlfoundry_artifacts_api.get_signed_urls_for_write_post(
+                get_signed_urls_for_artifact_version_write_request_dto=GetSignedURLsForArtifactVersionWriteRequestDto(
+                    id=str(artifact_identifier.artifact_version_id), paths=paths
                 )
-                for idx, signed_url in zip(range(start_idx, end_idx), signed_urls):
-                    (
-                        upload_path,
-                        local_file,
-                        multipart_info,
-                    ) = files_for_normal_upload[idx]
-                    future = executor.submit(
-                        self._log_artifact,
-                        local_file=local_file,
-                        artifact_path=upload_path,
-                        multipart_info=multipart_info,
-                        signed_url=signed_url,
-                        abort_event=abort_event,
-                        executor_for_multipart_upload=None,
-                        progress_bar=progress_bar,
-                    )
-                    futures.append(future)
-
-            done, not_done = wait(futures, return_when=FIRST_EXCEPTION)
-            if len(not_done) > 0:
-                abort_event.set()
-            for future in not_done:
-                future.cancel()
-            for future in done:
-                if future.exception() is not None:
-                    raise future.exception()
-
-            for (
-                upload_path,
-                local_file,
-                multipart_info,
-            ) in files_for_multipart_upload:
-                self._log_artifact(
-                    local_file=local_file,
-                    artifact_path=upload_path,
-                    signed_url=None,
-                    multipart_info=multipart_info,
-                    executor_for_multipart_upload=executor,
-                    progress_bar=progress_bar,
+            )
+            signed_urls = signed_urls_response.signed_urls
+        elif artifact_identifier.dataset_fqn:
+            signed_urls_dataset_response = self._mlfoundry_artifacts_api.get_signed_urls_for_dataset_write_post(
+                get_signed_url_for_dataset_write_request_dto=GetSignedURLForDatasetWriteRequestDto(
+                    dataset_fqn=artifact_identifier.dataset_fqn, paths=paths
                 )
+            )
+            signed_urls = signed_urls_dataset_response.signed_urls
+        else:
+            raise ValueError(
+                "Invalid artifact type - both `artifact_version_id` and `dataset_fqn` both are None"
+            )
+        return signed_urls
 
     def _normal_upload(
         self,
@@ -683,7 +571,7 @@ class MlFoundryArtifactsRepository:
         if progress_bar.disable:
             logger.info(
                 "Uploading %s to %s",
-                _get_relpath_if_in_tempdir(local_file),
+                local_file,
                 artifact_path,
             )
 
@@ -695,6 +583,36 @@ class MlFoundryArtifactsRepository:
         )
         logger.debug("Uploaded %s to %s", local_file, artifact_path)
 
+    def _create_multipart_upload_for_identifier(
+        self,
+        artifact_identifier: ArtifactIdentifier,
+        path,
+        num_parts,
+    ) -> MultiPartUploadDto:
+        if artifact_identifier.artifact_version_id:
+            create_multipart_response: MultiPartUploadResponseDto = self._mlfoundry_artifacts_api.create_multi_part_upload_post(
+                create_multi_part_upload_request_dto=CreateMultiPartUploadRequestDto(
+                    artifact_version_id=str(artifact_identifier.artifact_version_id),
+                    path=path,
+                    num_parts=num_parts,
+                )
+            )
+            multipart_upload = create_multipart_response.multipart_upload
+        elif artifact_identifier.dataset_fqn:
+            create_multipart_for_dataset_response = self._mlfoundry_artifacts_api.create_multipart_upload_for_dataset_post(
+                create_multi_part_upload_for_dataset_request_dto=CreateMultiPartUploadForDatasetRequestDto(
+                    dataset_fqn=artifact_identifier.dataset_fqn,
+                    path=path,
+                    num_parts=num_parts,
+                )
+            )
+            multipart_upload = create_multipart_for_dataset_response.multipart_upload
+        else:
+            raise ValueError(
+                "Invalid artifact type - both `artifact_version_id` and `dataset_fqn` both are None"
+            )
+        return multipart_upload
+
     def _multipart_upload(
         self,
         local_file: str,
@@ -707,11 +625,11 @@ class MlFoundryArtifactsRepository:
         if progress_bar.disable:
             logger.info(
                 "Uploading %s to %s using multipart upload",
-                _get_relpath_if_in_tempdir(local_file),
+                local_file,
                 artifact_path,
             )
 
-        multipart_upload = self.create_multipart_upload_for_identifier(
+        multipart_upload = self._create_multipart_upload_for_identifier(
             artifact_identifier=self.artifact_identifier,
             path=artifact_path,
             num_parts=multipart_info.num_parts,
@@ -743,7 +661,7 @@ class MlFoundryArtifactsRepository:
         else:
             raise NotImplementedError()
 
-    def _log_artifact(
+    def _upload_file(
         self,
         local_file: str,
         artifact_path: str,
@@ -763,7 +681,9 @@ class MlFoundryArtifactsRepository:
             )
 
         if not executor_for_multipart_upload:
-            with ThreadPoolExecutor(max_workers=_MAX_WORKERS_FOR_UPLOAD) as executor:
+            with ThreadPoolExecutor(
+                max_workers=ENV_VARS.TFY_ARTIFACTS_UPLOAD_MAX_WORKERS
+            ) as executor:
                 return self._multipart_upload(
                     local_file=local_file,
                     artifact_path=artifact_path,
@@ -780,10 +700,14 @@ class MlFoundryArtifactsRepository:
             progress_bar=progress_bar,
         )
 
-    def log_artifact(self, local_file: str, artifact_path: Optional[str] = None):
-        upload_path = artifact_path or ""
-        upload_path = upload_path.lstrip(posixpath.sep)
-        upload_path = posixpath.join(upload_path, os.path.basename(local_file))
+    def _upload(
+        self,
+        files_for_normal_upload: Sequence[Tuple[str, str, _FileMultiPartInfo]],
+        files_for_multipart_upload: Sequence[Tuple[str, str, _FileMultiPartInfo]],
+        progress: Optional[bool] = None,
+    ):
+        abort_event = Event()
+        show_progress = _can_display_progress(progress)
         with Progress(
             "[progress.description]{task.description}",
             BarColumn(bar_width=None),
@@ -793,170 +717,221 @@ class MlFoundryArtifactsRepository:
             TimeRemainingColumn(),
             TimeElapsedColumn(),
             refresh_per_second=1,
-            disable=True,
+            disable=not show_progress,
             expand=True,
-        ) as progress_bar:
-            self._log_artifact(
-                local_file=local_file,
-                artifact_path=upload_path,
-                multipart_info=_decide_file_parts(local_file),
-                progress_bar=progress_bar,
-            )
-
-    def _is_directory(self, artifact_path):
-        for _ in self.list_artifacts(artifact_path, page_size=3):
-            return True
-        return False
-
-    def download_artifacts(  # noqa: C901
-        self,
-        artifact_path: str,
-        dst_path: Optional[str] = None,
-        overwrite: bool = False,
-        progress: Optional[bool] = None,
-    ) -> str:
-        """
-        Download an artifact file or directory to a local directory if applicable, and return a
-        local path for it. The caller is responsible for managing the lifecycle of the downloaded artifacts.
-
-        Args:
-            artifact_path: Relative source path to the desired artifacts.
-            dst_path: Absolute path of the local filesystem destination directory to which to
-                download the specified artifacts. This directory must already exist.
-                If unspecified, the artifacts will either be downloaded to a new
-                uniquely-named directory.
-            overwrite: if to overwrite the files at/inside `dst_path` if they exist
-            progress: Show or hide progress bar
-
-        Returns:
-            str: Absolute path of the local filesystem location containing the desired artifacts.
-        """
-
-        show_progress = _can_display_progress()
-
-        is_dir_temp = False
-        if dst_path is None:
-            dst_path = tempfile.mkdtemp()
-            is_dir_temp = True
+        ) as progress_bar, ThreadPoolExecutor(
+            max_workers=ENV_VARS.TFY_ARTIFACTS_UPLOAD_MAX_WORKERS
+        ) as executor:
+            futures: List[Future] = []
+            # Note: While this batching is beneficial when there is a large number of files, there is also
+            # a rare case risk of the signed url expiring before a request is made to it
+            _batch_size = _GENERATE_SIGNED_URL_BATCH_SIZE
+            for start_idx in range(0, len(files_for_normal_upload), _batch_size):
+                end_idx = min(start_idx + _batch_size, len(files_for_normal_upload))
+                if _any_future_has_failed(futures):
+                    break
+                logger.debug("Generating write signed urls for a batch ...")
+                remote_file_paths = [
+                    files_for_normal_upload[idx][0] for idx in range(start_idx, end_idx)
+                ]
+                signed_urls = self.get_signed_urls_for_write(
+                    artifact_identifier=self.artifact_identifier,
+                    paths=remote_file_paths,
+                )
+                for idx, signed_url in zip(range(start_idx, end_idx), signed_urls):
+                    (
+                        upload_path,
+                        local_file,
+                        multipart_info,
+                    ) = files_for_normal_upload[idx]
+                    future = executor.submit(
+                        self._upload_file,
+                        local_file=local_file,
+                        artifact_path=upload_path,
+                        multipart_info=multipart_info,
+                        signed_url=signed_url,
+                        abort_event=abort_event,
+                        executor_for_multipart_upload=None,
+                        progress_bar=progress_bar,
+                    )
+                    futures.append(future)
 
-        dst_path = os.path.abspath(dst_path)
-        if is_dir_temp:
-            logger.info(
-                f"Using temporary directory {dst_path} as the download directory"
-            )
+            done, not_done = wait(futures, return_when=FIRST_EXCEPTION)
+            if len(not_done) > 0:
+                abort_event.set()
+            for future in not_done:
+                future.cancel()
+            for future in done:
+                if future.exception() is not None:
+                    raise future.exception()
 
-        if not os.path.exists(dst_path):
+            for (
+                upload_path,
+                local_file,
+                multipart_info,
+            ) in files_for_multipart_upload:
+                self._upload_file(
+                    local_file=local_file,
+                    artifact_path=upload_path,
+                    signed_url=None,
+                    multipart_info=multipart_info,
+                    executor_for_multipart_upload=executor,
+                    progress_bar=progress_bar,
+                )
+
+    def _add_file_for_upload(
+        self,
+        local_file: str,
+        artifact_path: str,
+        files_for_normal_upload: List[Tuple[str, str, _FileMultiPartInfo]],
+        files_for_multipart_upload: List[Tuple[str, str, _FileMultiPartInfo]],
+    ):
+        local_file = os.path.realpath(local_file)
+        if os.path.isdir(local_file):
             raise MlFoundryException(
-                message=(
-                    "The destination path for downloaded artifacts does not"
-                    " exist! Destination path: {dst_path}".format(dst_path=dst_path)
-                ),
+                "Cannot log a directory as an artifact. Use `log_artifacts` instead"
             )
-        elif not os.path.isdir(dst_path):
+        upload_path = artifact_path
+        upload_path = upload_path.lstrip(posixpath.sep)
+        multipart_info = _decide_file_parts(local_file)
+        if multipart_info.num_parts == 1:
+            files_for_normal_upload.append((upload_path, local_file, multipart_info))
+        else:
+            files_for_multipart_upload.append((upload_path, local_file, multipart_info))
+
+    def _add_dir_for_upload(
+        self,
+        local_dir: str,
+        artifact_path: str,
+        files_for_normal_upload: List[Tuple[str, str, _FileMultiPartInfo]],
+        files_for_multipart_upload: List[Tuple[str, str, _FileMultiPartInfo]],
+    ):
+        local_dir = os.path.realpath(local_dir)
+        if not os.path.isdir(local_dir):
             raise MlFoundryException(
-                message=(
-                    "The destination path for downloaded artifacts must be a directory!"
-                    " Destination path: {dst_path}".format(dst_path=dst_path)
-                ),
+                "Cannot log a file as a directory. Use `log_artifact` instead"
             )
+        dest_path = artifact_path
+        dest_path = dest_path.lstrip(posixpath.sep)
 
-        progress_bar = Progress(
-            "[progress.description]{task.description}",
-            BarColumn(),
-            "[progress.percentage]{task.percentage:>3.0f}%",
-            DownloadColumn(),
-            TransferSpeedColumn(),
-            TimeRemainingColumn(),
-            TimeElapsedColumn(),
-            refresh_per_second=1,
-            disable=not show_progress,
-            expand=True,
-        )
-
-        try:
-            progress_bar.start()
-            # Check if the artifacts points to a directory
-            if self._is_directory(artifact_path):
-                futures: List[Future] = []
-                file_paths: List[Tuple[str, str]] = []
-                abort_event = Event()
+        for root, _, file_names in os.walk(local_dir):
+            upload_path = dest_path
+            if root != local_dir:
+                rel_path = os.path.relpath(root, local_dir)
+                rel_path = relative_path_to_artifact_path(rel_path)
+                upload_path = posixpath.join(dest_path, rel_path)
+            for file_name in file_names:
+                local_file = os.path.join(root, file_name)
+                self._add_file_for_upload(
+                    local_file=local_file,
+                    artifact_path=upload_path,
+                    files_for_normal_upload=files_for_normal_upload,
+                    files_for_multipart_upload=files_for_multipart_upload,
+                )
 
-                # Check if any file is being overwritten before downloading them
-                for file_path, download_dest_path in self._get_file_paths_recur(
-                    src_artifact_dir_path=artifact_path, dst_local_dir_path=dst_path
-                ):
-                    final_file_path = os.path.join(download_dest_path, file_path)
+    def log_artifacts(
+        self,
+        src_dest_pairs: Sequence[Tuple[str, Optional[str]]],
+        progress: Optional[bool] = None,
+    ):
+        files_for_normal_upload: List[Tuple[str, str, _FileMultiPartInfo]] = []
+        files_for_multipart_upload: List[Tuple[str, str, _FileMultiPartInfo]] = []
+        for src, dest in src_dest_pairs:
+            if os.path.isdir(src):
+                self._add_dir_for_upload(
+                    local_dir=src,
+                    artifact_path=dest or "",
+                    files_for_normal_upload=files_for_normal_upload,
+                    files_for_multipart_upload=files_for_multipart_upload,
+                )
+            else:
+                self._add_file_for_upload(
+                    local_file=src,
+                    artifact_path=dest or "",
+                    files_for_normal_upload=files_for_normal_upload,
+                    files_for_multipart_upload=files_for_multipart_upload,
+                )
+        self._upload(
+            files_for_normal_upload=files_for_normal_upload,
+            files_for_multipart_upload=files_for_multipart_upload,
+            progress=progress,
+        )
 
-                    # There would be no overwrite if temp directory is being used
-                    if (
-                        not is_dir_temp
-                        and os.path.exists(final_file_path)
-                        and not overwrite
-                    ):
-                        raise MlFoundryException(
-                            f"File already exists at {final_file_path}, aborting download "
-                            f"(set `overwrite` flag to overwrite this and any subsequent files)"
-                        )
-                    file_paths.append((file_path, download_dest_path))
+    def _list_files(
+        self, artifact_identifier: ArtifactIdentifier, path, page_size, page_token
+    ) -> Union[ListFilesForDatasetResponseDto, ListFilesForArtifactVersionsResponseDto]:
+        if artifact_identifier.dataset_fqn:
+            return self._mlfoundry_artifacts_api.list_files_for_dataset_post(
+                list_files_for_dataset_request_dto=ListFilesForDatasetRequestDto(
+                    dataset_fqn=artifact_identifier.dataset_fqn,
+                    path=path,
+                    max_results=page_size,
+                    page_token=page_token,
+                )
+            )
+        else:
+            return self._mlfoundry_artifacts_api.list_files_for_artifact_version_post(
+                list_files_for_artifact_version_request_dto=ListFilesForArtifactVersionRequestDto(
+                    id=str(artifact_identifier.artifact_version_id),
+                    path=path,
+                    max_results=page_size,
+                    page_token=page_token,
+                )
+            )
 
-                with ThreadPoolExecutor(_MAX_WORKERS_FOR_DOWNLOAD) as executor:
-                    # Note: While this batching is beneficial when there is a large number of files, there is also
-                    # a rare case risk of the signed url expiring before a request is made to it
-                    batch_size = _GENERATE_SIGNED_URL_BATCH_SIZE
-                    for start_idx in range(0, len(file_paths), batch_size):
-                        end_idx = min(start_idx + batch_size, len(file_paths))
-                        if _any_future_has_failed(futures):
-                            break
-                        logger.debug("Generating read signed urls for a batch ...")
-                        remote_file_paths = [
-                            file_paths[idx][0] for idx in range(start_idx, end_idx)
-                        ]
-                        signed_urls = self.get_signed_urls_for_read(
-                            artifact_identifier=self.artifact_identifier,
-                            paths=remote_file_paths,
-                        )
-                        for idx, signed_url in zip(
-                            range(start_idx, end_idx), signed_urls
-                        ):
-                            file_path, download_dest_path = file_paths[idx]
-                            future = executor.submit(
-                                self._download_artifact,
-                                src_artifact_path=file_path,
-                                dst_local_dir_path=download_dest_path,
-                                signed_url=signed_url,
-                                abort_event=abort_event,
-                                progress_bar=progress_bar,
-                            )
-                            futures.append(future)
+    def list_artifacts(
+        self, path=None, page_size=_LIST_FILES_PAGE_SIZE, **kwargs
+    ) -> Iterator[FileInfoDto]:
+        page_token = None
+        started = False
+        while not started or page_token is not None:
+            started = True
+            page = self._list_files(
+                artifact_identifier=self.artifact_identifier,
+                path=path,
+                page_size=page_size,
+                page_token=page_token,
+            )
+            for file_info in page.files:
+                yield file_info
+            page_token = page.next_page_token
 
-                    done, not_done = wait(futures, return_when=FIRST_EXCEPTION)
-                    if len(not_done) > 0:
-                        abort_event.set()
-                    for future in not_done:
-                        future.cancel()
-                    for future in done:
-                        if future.exception() is not None:
-                            raise future.exception()
+    def _is_directory(self, artifact_path):
+        for _ in self.list_artifacts(artifact_path, page_size=3):
+            return True
+        return False
 
-                    output_dir = os.path.join(dst_path, artifact_path)
-                    return output_dir
-            else:
-                return self._download_artifact(
-                    src_artifact_path=artifact_path,
-                    dst_local_dir_path=dst_path,
-                    signed_url=None,
-                    progress_bar=progress_bar,
-                )
-        except Exception as err:
-            if is_dir_temp:
-                logger.info(
-                    f"Error encountered, removing temporary download directory at {dst_path}"
-                )
-                rmtree(dst_path)  # remove temp directory alongside it's contents
-            raise err
+    def _create_download_destination(
+        self, src_artifact_path: str, dst_local_dir_path: str
+    ) -> str:
+        """
+        Creates a local filesystem location to be used as a destination for downloading the artifact
+        specified by `src_artifact_path`. The destination location is a subdirectory of the
+        specified `dst_local_dir_path`, which is determined according to the structure of
+        `src_artifact_path`. For example, if `src_artifact_path` is `dir1/file1.txt`, then the
+        resulting destination path is `<dst_local_dir_path>/dir1/file1.txt`. Local directories are
+        created for the resulting destination location if they do not exist.
 
-        finally:
-            progress_bar.stop()
+        :param src_artifact_path: A relative, POSIX-style path referring to an artifact stored
+                                  within the repository's artifact root location.
+                                  `src_artifact_path` should be specified relative to the
+                                  repository's artifact root location.
+        :param dst_local_dir_path: The absolute path to a local filesystem directory in which the
+                                   local destination path will be contained. The local destination
+                                   path may be contained in a subdirectory of `dst_root_dir` if
+                                   `src_artifact_path` contains subdirectories.
+        :return: The absolute path to a local filesystem location to be used as a destination
+                 for downloading the artifact specified by `src_artifact_path`.
+        """
+        src_artifact_path = src_artifact_path.rstrip(
+            "/"
+        )  # Ensure correct dirname for trailing '/'
+        dirpath = posixpath.dirname(src_artifact_path)
+        local_dir_path = os.path.join(dst_local_dir_path, dirpath)
+        local_file_path = os.path.join(dst_local_dir_path, src_artifact_path)
+        if not os.path.exists(local_dir_path):
+            os.makedirs(local_dir_path, exist_ok=True)
+        return local_file_path
 
     # noinspection PyMethodOverriding
     def _download_file(
@@ -1003,8 +978,8 @@ class MlFoundryArtifactsRepository:
 
     def _download_artifact(
         self,
-        src_artifact_path,
-        dst_local_dir_path,
+        src_artifact_path: str,
+        dst_local_dir_path: str,
         signed_url: Optional[SignedURLDto],
         progress_bar: Optional[Progress] = None,
         abort_event=None,
@@ -1057,105 +1032,154 @@ class MlFoundryArtifactsRepository:
                 else:
                     yield file_info.path, dst_local_dir_path
 
-    # TODO (chiragjn): Refactor these methods - if else is very inconvenient
-    def get_signed_urls_for_read(
+    def download_artifacts(  # noqa: C901
         self,
-        artifact_identifier: ArtifactIdentifier,
-        paths,
-    ) -> List[SignedURLDto]:
-        if artifact_identifier.artifact_version_id:
-            signed_urls_response = self._mlfoundry_artifacts_api.get_signed_urls_for_read_post(
-                get_signed_urls_for_artifact_version_read_request_dto=GetSignedURLsForArtifactVersionReadRequestDto(
-                    id=str(artifact_identifier.artifact_version_id), paths=paths
-                )
-            )
-            signed_urls = signed_urls_response.signed_urls
-        elif artifact_identifier.dataset_fqn:
-            signed_urls_dataset_response = self._mlfoundry_artifacts_api.get_signed_urls_dataset_read_post(
-                get_signed_urls_for_dataset_read_request_dto=GetSignedURLsForDatasetReadRequestDto(
-                    dataset_fqn=artifact_identifier.dataset_fqn, paths=paths
-                )
-            )
-            signed_urls = signed_urls_dataset_response.signed_urls
-        else:
-            raise ValueError(
-                "Invalid artifact type - both `artifact_version_id` and `dataset_fqn` both are None"
-            )
-        return signed_urls
+        artifact_path: str,
+        dst_path: Optional[str] = None,
+        overwrite: bool = False,
+        progress: Optional[bool] = None,
+    ) -> str:
+        """
+        Download an artifact file or directory to a local directory if applicable, and return a
+        local path for it. The caller is responsible for managing the lifecycle of the downloaded artifacts.
 
-    def get_signed_urls_for_write(
-        self,
-        artifact_identifier: ArtifactIdentifier,
-        paths: List[str],
-    ) -> List[SignedURLDto]:
-        if artifact_identifier.artifact_version_id:
-            signed_urls_response = self._mlfoundry_artifacts_api.get_signed_urls_for_write_post(
-                get_signed_urls_for_artifact_version_write_request_dto=GetSignedURLsForArtifactVersionWriteRequestDto(
-                    id=str(artifact_identifier.artifact_version_id), paths=paths
-                )
-            )
-            signed_urls = signed_urls_response.signed_urls
-        elif artifact_identifier.dataset_fqn:
-            signed_urls_dataset_response = self._mlfoundry_artifacts_api.get_signed_urls_for_dataset_write_post(
-                get_signed_url_for_dataset_write_request_dto=GetSignedURLForDatasetWriteRequestDto(
-                    dataset_fqn=artifact_identifier.dataset_fqn, paths=paths
-                )
-            )
-            signed_urls = signed_urls_dataset_response.signed_urls
-        else:
-            raise ValueError(
-                "Invalid artifact type - both `artifact_version_id` and `dataset_fqn` both are None"
-            )
-        return signed_urls
+        Args:
+            artifact_path: Relative source path to the desired artifacts.
+            dst_path: Absolute path of the local filesystem destination directory to which to
+                download the specified artifacts. This directory must already exist.
+                If unspecified, the artifacts will either be downloaded to a new
+                uniquely-named directory.
+            overwrite: if to overwrite the files at/inside `dst_path` if they exist
+            progress: Show or hide progress bar
 
-    def create_multipart_upload_for_identifier(
-        self,
-        artifact_identifier: ArtifactIdentifier,
-        path,
-        num_parts,
-    ) -> MultiPartUploadDto:
-        if artifact_identifier.artifact_version_id:
-            create_multipart_response: MultiPartUploadResponseDto = self._mlfoundry_artifacts_api.create_multi_part_upload_post(
-                create_multi_part_upload_request_dto=CreateMultiPartUploadRequestDto(
-                    artifact_version_id=str(artifact_identifier.artifact_version_id),
-                    path=path,
-                    num_parts=num_parts,
-                )
+        Returns:
+            str: Absolute path of the local filesystem location containing the desired artifacts.
+        """
+
+        show_progress = _can_display_progress(user_choice=progress)
+
+        is_dir_temp = False
+        if dst_path is None:
+            dst_path = tempfile.mkdtemp()
+            is_dir_temp = True
+
+        dst_path = os.path.abspath(dst_path)
+        if is_dir_temp:
+            logger.info(
+                f"Using temporary directory {dst_path} as the download directory"
             )
-            multipart_upload = create_multipart_response.multipart_upload
-        elif artifact_identifier.dataset_fqn:
-            create_multipart_for_dataset_response = self._mlfoundry_artifacts_api.create_multipart_upload_for_dataset_post(
-                create_multi_part_upload_for_dataset_request_dto=CreateMultiPartUploadForDatasetRequestDto(
-                    dataset_fqn=artifact_identifier.dataset_fqn,
-                    path=path,
-                    num_parts=num_parts,
-                )
+
+        if not os.path.exists(dst_path):
+            raise MlFoundryException(
+                message=(
+                    "The destination path for downloaded artifacts does not"
+                    " exist! Destination path: {dst_path}".format(dst_path=dst_path)
+                ),
             )
-            multipart_upload = create_multipart_for_dataset_response.multipart_upload
-        else:
-            raise ValueError(
-                "Invalid artifact type - both `artifact_version_id` and `dataset_fqn` both are None"
+        elif not os.path.isdir(dst_path):
+            raise MlFoundryException(
+                message=(
+                    "The destination path for downloaded artifacts must be a directory!"
+                    " Destination path: {dst_path}".format(dst_path=dst_path)
+                ),
             )
-        return multipart_upload
 
-    def list_files(
-        self, artifact_identifier: ArtifactIdentifier, path, page_size, page_token
-    ) -> Union[ListFilesForDatasetResponseDto, ListFilesForArtifactVersionsResponseDto]:
-        if artifact_identifier.dataset_fqn:
-            return self._mlfoundry_artifacts_api.list_files_for_dataset_post(
-                list_files_for_dataset_request_dto=ListFilesForDatasetRequestDto(
-                    dataset_fqn=artifact_identifier.dataset_fqn,
-                    path=path,
-                    max_results=page_size,
-                    page_token=page_token,
+        progress_bar = Progress(
+            "[progress.description]{task.description}",
+            BarColumn(),
+            "[progress.percentage]{task.percentage:>3.0f}%",
+            DownloadColumn(),
+            TransferSpeedColumn(),
+            TimeRemainingColumn(),
+            TimeElapsedColumn(),
+            refresh_per_second=1,
+            disable=not show_progress,
+            expand=True,
+        )
+
+        try:
+            progress_bar.start()
+            # Check if the artifacts points to a directory
+            if self._is_directory(artifact_path):
+                futures: List[Future] = []
+                file_paths: List[Tuple[str, str]] = []
+                abort_event = Event()
+
+                # Check if any file is being overwritten before downloading them
+                for file_path, download_dest_path in self._get_file_paths_recur(
+                    src_artifact_dir_path=artifact_path, dst_local_dir_path=dst_path
+                ):
+                    final_file_path = os.path.join(download_dest_path, file_path)
+
+                    # There would be no overwrite if temp directory is being used
+                    if (
+                        not is_dir_temp
+                        and os.path.exists(final_file_path)
+                        and not overwrite
+                    ):
+                        raise MlFoundryException(
+                            f"File already exists at {final_file_path}, aborting download "
+                            f"(set `overwrite` flag to overwrite this and any subsequent files)"
+                        )
+                    file_paths.append((file_path, download_dest_path))
+
+                with ThreadPoolExecutor(
+                    max_workers=ENV_VARS.TFY_ARTIFACTS_DOWNLOAD_MAX_WORKERS
+                ) as executor:
+                    # Note: While this batching is beneficial when there is a large number of files, there is also
+                    # a rare case risk of the signed url expiring before a request is made to it
+                    batch_size = _GENERATE_SIGNED_URL_BATCH_SIZE
+                    for start_idx in range(0, len(file_paths), batch_size):
+                        end_idx = min(start_idx + batch_size, len(file_paths))
+                        if _any_future_has_failed(futures):
+                            break
+                        logger.debug("Generating read signed urls for a batch ...")
+                        remote_file_paths = [
+                            file_paths[idx][0] for idx in range(start_idx, end_idx)
+                        ]
+                        signed_urls = self.get_signed_urls_for_read(
+                            artifact_identifier=self.artifact_identifier,
+                            paths=remote_file_paths,
+                        )
+                        for idx, signed_url in zip(
+                            range(start_idx, end_idx), signed_urls
+                        ):
+                            file_path, download_dest_path = file_paths[idx]
+                            future = executor.submit(
+                                self._download_artifact,
+                                src_artifact_path=file_path,
+                                dst_local_dir_path=download_dest_path,
+                                signed_url=signed_url,
+                                abort_event=abort_event,
+                                progress_bar=progress_bar,
+                            )
+                            futures.append(future)
+
+                    done, not_done = wait(futures, return_when=FIRST_EXCEPTION)
+                    if len(not_done) > 0:
+                        abort_event.set()
+                    for future in not_done:
+                        future.cancel()
+                    for future in done:
+                        if future.exception() is not None:
+                            raise future.exception()
+
+                    output_dir = os.path.join(dst_path, artifact_path)
+                    return output_dir
+            else:
+                return self._download_artifact(
+                    src_artifact_path=artifact_path,
+                    dst_local_dir_path=dst_path,
+                    signed_url=None,
+                    progress_bar=progress_bar,
                 )
-            )
-        else:
-            return self._mlfoundry_artifacts_api.list_files_for_artifact_version_post(
-                list_files_for_artifact_version_request_dto=ListFilesForArtifactVersionRequestDto(
-                    id=str(artifact_identifier.artifact_version_id),
-                    path=path,
-                    max_results=page_size,
-                    page_token=page_token,
+        except Exception as err:
+            if is_dir_temp:
+                logger.info(
+                    f"Error encountered, removing temporary download directory at {dst_path}"
                 )
-            )
+                rmtree(dst_path)  # remove temp directory alongside it's contents
+            raise err
+
+        finally:
+            progress_bar.stop()