feat(py + cli): download datasets as files by default, add --zip flag

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: <ORG>/<TEAMSPACE>/<DATASET_NAME> or
     <ORG>/<TEAMSPACE>/<DATASET_NAME>/<VERSION>. 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}")

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
             ``<ORGANIZATION>/<TEAMSPACE>/<DATASET-NAME>`` or
             ``<ORGANIZATION>/<TEAMSPACE>/<DATASET-NAME>/<VERSION>``.
             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:

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,24 +180,29 @@ 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")
             if not file_url:
                 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,

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.",
     )
 
 

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(