diff --git a/python/lightning_sdk/cli/dataset/download.py b/python/lightning_sdk/cli/dataset/download.py index 78937ade..3bdc03cb 100644 --- a/python/lightning_sdk/cli/dataset/download.py +++ b/python/lightning_sdk/cli/dataset/download.py @@ -23,27 +23,39 @@ "--target_path", default=".", type=click.Path(file_okay=False, dir_okay=True), - help="Local directory to download the dataset zip into.", + help="Local directory to download the dataset into.", +) +@click.option( + "--zip/--no-zip", + "as_zip", + default=False, + help="Package the dataset into a single .zip archive instead of a directory of files.", ) def download_dataset_cmd( name: str, cluster_id: Optional[str] = None, target_path: str = ".", + as_zip: bool = False, ) -> None: - """Download a dataset version as a zip file. + """Download a dataset version. NAME must be a Lightning path: // or ///. If no version specified, defaults to most recent version. + By default the files are downloaded into a directory. Pass --zip to + package them into a single .zip archive instead. + Usage: lightning dataset download my-org/my-teamspace/my-dataset lightning dataset download my-org/my-teamspace/my-dataset/v3 lightning dataset download my-org/my-teamspace/my-dataset/v3 --target-path ./data + lightning dataset download my-org/my-teamspace/my-dataset --zip """ info = download_dataset( name=name, target_path=target_path, cluster_id=cluster_id, + as_zip=as_zip, ) click.echo(f"Downloaded dataset '{info.name}' version '{info.version}' to {info.path}") diff --git a/python/lightning_sdk/datasets.py b/python/lightning_sdk/datasets.py index 950c21d0..b6f2351f 100644 --- a/python/lightning_sdk/datasets.py +++ b/python/lightning_sdk/datasets.py @@ -36,7 +36,7 @@ class DownloadedDatasetInfo: """Metadata returned after a successful dataset download. Attributes: - path: The local file path to the downloaded zip file. + path: The local path the dataset was downloaded to (a directory, or a zip file when downloaded with as_zip). name: The dataset name. version: The version tag that was downloaded (e.g. ``"v1"``). """ @@ -111,17 +111,22 @@ def download_dataset( name: str, target_path: str = ".", cluster_id: Optional[str] = None, + as_zip: bool = False, ) -> DownloadedDatasetInfo: - """Download a dataset version as a zip file from Lightning Datasets. + """Download a dataset version from Lightning Datasets. + + By default the files are downloaded into a directory. Set ``as_zip`` to + package them into a single zip archive instead. Args: name: Lightning path to dataset in the format ``//`` or ``///``. If no version specified, defaults to the most recent version. - target_path: Local directory to download the dataset zip into. + target_path: Local directory to download the dataset into. Defaults to the current directory (``"."``). cluster_id: Optional cluster ID to download from. + as_zip: If True, package the files into a single zip archive. Returns: DownloadedDatasetInfo: Metadata about the downloaded dataset, including @@ -136,15 +141,17 @@ def download_dataset( import os - zip_path = os.path.join(target_path, f"{dataset_name}_{version}.zip") + base_name = f"{dataset_name}_{version}" + out_path = os.path.join(target_path, f"{base_name}.zip" if as_zip else base_name) _download_dataset_version( project_id=project_id, dataset_name=dataset_name, version=version, - target_path=zip_path, + target_path=out_path, cluster_id=cluster_id, + as_zip=as_zip, ) - return DownloadedDatasetInfo(path=zip_path, name=dataset_name, version=version) + return DownloadedDatasetInfo(path=out_path, name=dataset_name, version=version) def list_datasets(name: str) -> list: diff --git a/python/lightning_sdk/lightning_cloud/utils/dataset.py b/python/lightning_sdk/lightning_cloud/utils/dataset.py index 06e79fb9..973a50b1 100644 --- a/python/lightning_sdk/lightning_cloud/utils/dataset.py +++ b/python/lightning_sdk/lightning_cloud/utils/dataset.py @@ -115,19 +115,23 @@ def _download_dataset_version( version: str, target_path: str, cluster_id: Optional[str] = None, + as_zip: bool = False, ): """ - Download a dataset version as a zip file from the API. + Download a dataset version from the API. - Fetches presigned file URLs from the files endpoint, downloads each file, - and packages them into a proper zip archive at the target path. + Fetches presigned file URLs from the files endpoint and downloads each file. + By default the files are written into ``target_path`` as a directory. When + ``as_zip`` is set, they are packaged into a zip archive at ``target_path`` + instead. Args: project_id: The project ID. dataset_name: The dataset name. version: The dataset version to download. - target_path: Local file path where the downloaded zip will be saved. + target_path: Local directory to download into, or the zip path when as_zip is set. cluster_id: Optional cluster ID. + as_zip: If True, package the files into a single zip archive. """ import shutil import tempfile @@ -176,8 +180,11 @@ def _download_dataset_version( if not files_list: raise ValueError(f"No files found for dataset '{dataset_name}' version '{version}' in project '{project_id}'.") - # create a zip archive from the downloaded files - tmp_dir = tempfile.mkdtemp() + # Download into a directory. When zipping, use a temp dir and archive it + # afterwards; otherwise download straight into target_path. + dest_dir = tempfile.mkdtemp() if as_zip else target_path + if not as_zip: + os.makedirs(dest_dir, exist_ok=True) try: for i, file_info in enumerate(files_list): file_url = file_info.get("url") @@ -185,15 +192,17 @@ def _download_dataset_version( continue filepath = file_info.get("filepath", f"file_{i}") filepath = filepath.lstrip("/") - local_tmp = os.path.join(tmp_dir, filepath) - os.makedirs(os.path.dirname(local_tmp), exist_ok=True) - urllib.request.urlretrieve(file_url, local_tmp) - base = target_path - if base.endswith(".zip"): - base = base[:-4] - shutil.make_archive(base, "zip", tmp_dir) + local_path = os.path.join(dest_dir, filepath) + os.makedirs(os.path.dirname(local_path), exist_ok=True) + urllib.request.urlretrieve(file_url, local_path) + if as_zip: + base = target_path + if base.endswith(".zip"): + base = base[:-4] + shutil.make_archive(base, "zip", dest_dir) finally: - shutil.rmtree(tmp_dir, ignore_errors=True) + if as_zip: + shutil.rmtree(dest_dir, ignore_errors=True) def _create_dataset( project_id: str, diff --git a/python/tests/cli/dataset/test_download.py b/python/tests/cli/dataset/test_download.py index 4563eb3a..25737a69 100644 --- a/python/tests/cli/dataset/test_download.py +++ b/python/tests/cli/dataset/test_download.py @@ -6,8 +6,9 @@ def test_dataset_download_help() -> None: assert_help_contains( "lightning dataset download --help", "Usage: lightning dataset download", - "Download a dataset version as a zip file.", + "Download a dataset version.", "NAME must be a Lightning path:", + "--zip", ) @@ -16,7 +17,7 @@ def test_datasets_download_help() -> None: assert_help_contains( "lightning datasets download --help", "Usage: lightning datasets download", - "Download a dataset version as a zip file.", + "Download a dataset version.", ) diff --git a/python/tests/core/test_teamspace.py b/python/tests/core/test_teamspace.py index c8509cf3..ffbb3ae8 100644 --- a/python/tests/core/test_teamspace.py +++ b/python/tests/core/test_teamspace.py @@ -1235,6 +1235,7 @@ def test_download_dataset_version( version="3", target_path=target_path, cluster_id="aws-us-east", + as_zip=True, ) mock_api_client.request.assert_any_call( @@ -1257,6 +1258,60 @@ def test_download_dataset_version( os.unlink(target_path) +@mock.patch.dict(os.environ, {"LIGHTNING_CLOUD_URL": "https://lightning.ai", "LIGHTNING_AUTH_TOKEN": "test-token"}) +@mock.patch("lightning_sdk.lightning_cloud.utils.dataset.LightningClient") +@mock.patch("lightning_sdk.lightning_cloud.rest_client.Auth", new=mock.MagicMock()) +def test_download_dataset_version_default_downloads_files_no_zip( + mock_lightning_client, + internal_teamspace_api_list_mocker, + internal_user_api_mocker, +): + import json + import shutil + import tempfile + import urllib + + from lightning_sdk.lightning_cloud.utils.dataset import _download_dataset_version + + mock_list_response = mock.MagicMock() + mock_list_response.data = json.dumps({"datasets": []}) + + mock_files_response = mock.MagicMock() + mock_files_response.data = json.dumps( + {"files": [{"filepath": "data.csv", "url": "https://presigned.example.com/data.csv"}]} + ) + + mock_api_client = mock.MagicMock() + mock_api_client.request.side_effect = [mock_list_response, mock_files_response] + mock_api_client.default_headers = {"Authorization": "Bearer test"} + + mock_client_instance = mock.MagicMock() + mock_client_instance.api_client = mock_api_client + mock_lightning_client.return_value = mock_client_instance + + target_path = tempfile.mkdtemp() + try: + with mock.patch.object(urllib.request, "urlretrieve", return_value=None) as mock_urlretrieve, mock.patch( + "shutil.make_archive" + ) as mock_make_archive: + # as_zip defaults to False: files land in target_path, no archive is created + _download_dataset_version( + project_id="proj-1", + dataset_name="ds-1", + version="3", + target_path=target_path, + cluster_id="aws-us-east", + ) + + mock_make_archive.assert_not_called() + mock_urlretrieve.assert_called_once() + args = mock_urlretrieve.call_args.args + assert args[0] == "https://presigned.example.com/data.csv" + assert args[1] == os.path.join(target_path, "data.csv") + finally: + shutil.rmtree(target_path, ignore_errors=True) + + @mock.patch.dict(os.environ, {"LIGHTNING_CLOUD_URL": "https://lightning.ai"}) @mock.patch("lightning_sdk.lightning_cloud.utils.dataset.LightningClient") @mock.patch("lightning_sdk.lightning_cloud.rest_client.Auth", new=mock.MagicMock()) @@ -1296,6 +1351,7 @@ def test_download_dataset_version_no_token_no_cluster( dataset_name="ds-2", version="1", target_path=target_path, + as_zip=True, ) mock_api_client.request.assert_any_call(