truefoundry 0.4.4rc9__py3-none-any.whl → 0.4.4rc10__py3-none-any.whl

truefoundry/ml/artifact/truefoundry_artifact_repo.py

@@ -16,6 +16,7 @@ from typing import (
     List,
     NamedTuple,
     Optional,
+    Sequence,
     Tuple,
     Union,
 )
@@ -88,13 +89,6 @@ _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
@@ -239,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):
@@ -340,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:
@@ -410,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):
@@ -510,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=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,
+    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,
@@ -685,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,
             )
 
@@ -697,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,
@@ -709,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,
@@ -745,7 +661,7 @@ class MlFoundryArtifactsRepository:
         else:
             raise NotImplementedError()
 
-    def _log_artifact(
+    def _upload_file(
         self,
         local_file: str,
         artifact_path: str,
@@ -784,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),
@@ -797,49 +717,349 @@ 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:
+        ) 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)
+
+            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._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(
+                "Cannot log a directory as an artifact. Use `log_artifacts` instead"
+            )
+        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(
+                "Cannot log a file as a directory. Use `log_artifact` instead"
+            )
+        dest_path = artifact_path
+        dest_path = dest_path.lstrip(posixpath.sep)
+
+        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,
+                )
+
+    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,
+        )
+
+    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,
+                )
+            )
+
+    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
+
+    def _is_directory(self, artifact_path):
+        for _ in self.list_artifacts(artifact_path, page_size=3):
+            return True
+        return False
+
+    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.
+
+        :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(
+        self,
+        remote_file_path: str,
+        local_path: str,
+        progress_bar: Optional[Progress],
+        signed_url: Optional[SignedURLDto],
+        abort_event: Optional[Event] = None,
+    ):
+        if not remote_file_path:
+            raise MlFoundryException(
+                f"remote_file_path cannot be None or empty str {remote_file_path}"
+            )
+        if not signed_url:
+            signed_url = self.get_signed_urls_for_read(
+                artifact_identifier=self.artifact_identifier, paths=[remote_file_path]
+            )[0]
+
+        if progress_bar is None or not progress_bar.disable:
+            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
+            )
+
+        def callback(chunk, total_file_size):
+            if progress_bar is not None:
+                progress_bar.update(
+                    download_progress_bar,
+                    advance=chunk,
+                    total=total_file_size,
+                )
+            if abort_event and abort_event.is_set():
+                raise Exception("aborting download")
+
+        _download_file_using_http_uri(
+            http_uri=signed_url.signed_url,
+            download_path=local_path,
+            callback=callback,
+        )
+        logger.debug("Downloaded %s to %s", remote_file_path, local_path)
+
+    def _download_artifact(
+        self,
+        src_artifact_path: str,
+        dst_local_dir_path: str,
+        signed_url: Optional[SignedURLDto],
+        progress_bar: Optional[Progress] = None,
+        abort_event=None,
+    ) -> str:
+        """
+        Download the file artifact specified by `src_artifact_path` to the local filesystem
+        directory specified by `dst_local_dir_path`.
+        :param src_artifact_path: A relative, POSIX-style path referring to a file 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: Absolute path of the local filesystem destination directory
+                                    to which to download the specified artifact. The downloaded
+                                    artifact may be written to a subdirectory of
+                                    `dst_local_dir_path` if `src_artifact_path` contains
+                                    subdirectories.
+        :param progress_bar: An instance of a Rich progress bar used to visually display the
+                              progress of the file download.
+        :return: A local filesystem path referring to the downloaded file.
+        """
+        local_destination_file_path = self._create_download_destination(
+            src_artifact_path=src_artifact_path, dst_local_dir_path=dst_local_dir_path
+        )
+        self._download_file(
+            remote_file_path=src_artifact_path,
+            local_path=local_destination_file_path,
+            signed_url=signed_url,
+            abort_event=abort_event,
+            progress_bar=progress_bar,
+        )
+        return local_destination_file_path
+
+    def _get_file_paths_recur(self, src_artifact_dir_path, dst_local_dir_path):
+        local_dir = os.path.join(dst_local_dir_path, src_artifact_dir_path)
+        dir_content = [  # prevent infinite loop, sometimes the dir is recursively included
+            file_info
+            for file_info in self.list_artifacts(src_artifact_dir_path)
+            if file_info.path != "." and file_info.path != src_artifact_dir_path
+        ]
+        if not dir_content:  # empty dir
+            if not os.path.exists(local_dir):
+                os.makedirs(local_dir, exist_ok=True)
+        else:
+            for file_info in dir_content:
+                if file_info.is_dir:
+                    yield from self._get_file_paths_recur(
+                        src_artifact_dir_path=file_info.path,
+                        dst_local_dir_path=dst_local_dir_path,
+                    )
+                else:
+                    yield file_info.path, dst_local_dir_path
+
+    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(user_choice=progress)
+
+        is_dir_temp = False
+        if dst_path is None:
             dst_path = tempfile.mkdtemp()
             is_dir_temp = True
 
@@ -963,205 +1183,3 @@ class MlFoundryArtifactsRepository:
 
         finally:
             progress_bar.stop()
-
-    # noinspection PyMethodOverriding
-    def _download_file(
-        self,
-        remote_file_path: str,
-        local_path: str,
-        progress_bar: Optional[Progress],
-        signed_url: Optional[SignedURLDto],
-        abort_event: Optional[Event] = None,
-    ):
-        if not remote_file_path:
-            raise MlFoundryException(
-                f"remote_file_path cannot be None or empty str {remote_file_path}"
-            )
-        if not signed_url:
-            signed_url = self.get_signed_urls_for_read(
-                artifact_identifier=self.artifact_identifier, paths=[remote_file_path]
-            )[0]
-
-        if progress_bar is None or not progress_bar.disable:
-            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
-            )
-
-        def callback(chunk, total_file_size):
-            if progress_bar is not None:
-                progress_bar.update(
-                    download_progress_bar,
-                    advance=chunk,
-                    total=total_file_size,
-                )
-            if abort_event and abort_event.is_set():
-                raise Exception("aborting download")
-
-        _download_file_using_http_uri(
-            http_uri=signed_url.signed_url,
-            download_path=local_path,
-            callback=callback,
-        )
-        logger.debug("Downloaded %s to %s", remote_file_path, local_path)
-
-    def _download_artifact(
-        self,
-        src_artifact_path,
-        dst_local_dir_path,
-        signed_url: Optional[SignedURLDto],
-        progress_bar: Optional[Progress] = None,
-        abort_event=None,
-    ) -> str:
-        """
-        Download the file artifact specified by `src_artifact_path` to the local filesystem
-        directory specified by `dst_local_dir_path`.
-        :param src_artifact_path: A relative, POSIX-style path referring to a file 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: Absolute path of the local filesystem destination directory
-                                    to which to download the specified artifact. The downloaded
-                                    artifact may be written to a subdirectory of
-                                    `dst_local_dir_path` if `src_artifact_path` contains
-                                    subdirectories.
-        :param progress_bar: An instance of a Rich progress bar used to visually display the
-                              progress of the file download.
-        :return: A local filesystem path referring to the downloaded file.
-        """
-        local_destination_file_path = self._create_download_destination(
-            src_artifact_path=src_artifact_path, dst_local_dir_path=dst_local_dir_path
-        )
-        self._download_file(
-            remote_file_path=src_artifact_path,
-            local_path=local_destination_file_path,
-            signed_url=signed_url,
-            abort_event=abort_event,
-            progress_bar=progress_bar,
-        )
-        return local_destination_file_path
-
-    def _get_file_paths_recur(self, src_artifact_dir_path, dst_local_dir_path):
-        local_dir = os.path.join(dst_local_dir_path, src_artifact_dir_path)
-        dir_content = [  # prevent infinite loop, sometimes the dir is recursively included
-            file_info
-            for file_info in self.list_artifacts(src_artifact_dir_path)
-            if file_info.path != "." and file_info.path != src_artifact_dir_path
-        ]
-        if not dir_content:  # empty dir
-            if not os.path.exists(local_dir):
-                os.makedirs(local_dir, exist_ok=True)
-        else:
-            for file_info in dir_content:
-                if file_info.is_dir:
-                    yield from self._get_file_paths_recur(
-                        src_artifact_dir_path=file_info.path,
-                        dst_local_dir_path=dst_local_dir_path,
-                    )
-                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(
-        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
-
-    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
-
-    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 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,
-                )
-            )