Skip to content

feat(py): download datasets as files with optional unzip #32

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
ethanwharris merged 2 commits into from
Jul 23, 2026
Merged

feat(py): download datasets as files with optional unzip #32

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
cd3b6d6
feat(py): support directory and extracted dataset downloads
ravivooda Jul 21, 2026
1ac5adf
refactor(py): remove post-download zip output
ravivooda Jul 22, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
16 changes: 14 additions & 2 deletions python/lightning_sdk/cli/dataset/download.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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(
"--unzip",
is_flag=True,
default=False,
help="Extract a dataset stored as exactly one .zip artifact.",
)
def download_dataset_cmd(
name: str,
cluster_id: Optional[str] = None,
target_path: str = ".",
unzip: 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, files are downloaded into a directory. Pass --unzip to
explicitly extract a version stored as exactly one ZIP artifact.

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 --unzip
"""
info = download_dataset(
name=name,
target_path=target_path,
cluster_id=cluster_id,
unzip=unzip,
)
click.echo(f"Downloaded dataset '{info.name}' version '{info.version}' to {info.path}")
24 changes: 17 additions & 7 deletions python/lightning_sdk/datasets.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -39,7 +39,7 @@ class DownloadedDatasetInfo:
"""Metadata returned after a successful dataset download.

Attributes:
path: The local file path to the downloaded zip file.
path: The local directory containing the dataset.
name: The dataset name.
version: The version tag that was downloaded (e.g. ``"v1"``).
"""
Expand Down Expand Up @@ -123,25 +123,34 @@ def download_dataset(
cluster_id: Optional[str] = None,
num_workers: int = _DEFAULT_DOWNLOAD_WORKERS,
part_size: int = _DEFAULT_DOWNLOAD_PART_SIZE,
unzip: bool = False,
) -> DownloadedDatasetInfo:
"""Download a dataset version as a zip file from Lightning Datasets.
"""Download a dataset version from Lightning Datasets.

Files download concurrently, each fetched via chunked HTTP-Range requests.
By default, files are downloaded into a directory. Set ``unzip=True`` to
safely extract a dataset version stored as exactly one ZIP artifact. Files
download concurrently, each fetched via chunked HTTP-Range requests.

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.
num_workers: number of concurrent download threads (default 16).
part_size: byte-range part size for splitting large files (default 64 MB).
unzip: Extract a version stored as exactly one ZIP artifact into a
directory. This is never enabled automatically.

Returns:
DownloadedDatasetInfo: Metadata about the downloaded dataset, including
the local path, dataset name, and version.

Raises:
ValueError: If ``unzip=True`` is used with a version that is not exactly
one ZIP artifact.
"""
org_name, ts_name, dataset_name, version = _parse_dataset_path(name)
teamspace = _get_teamspace(name=ts_name, organization=org_name)
Expand All @@ -154,18 +163,19 @@ def download_dataset(

import os

zip_path = os.path.join(target_path, f"{dataset_name}_{version}.zip")
output_path = os.path.join(target_path, f"{dataset_name}_{version}")
_download_dataset_version(
project_id=project_id,
dataset_name=dataset_name,
version=version,
target_path=zip_path,
target_path=output_path,
cluster_id=cluster_id,
dataset_id=dataset_id,
num_workers=num_workers,
part_size=part_size,
unzip=unzip,
)
return DownloadedDatasetInfo(path=zip_path, name=dataset_name, version=version)
return DownloadedDatasetInfo(path=output_path, name=dataset_name, version=version)


def list_datasets(name: str) -> list:
Expand Down
158 changes: 120 additions & 38 deletions python/lightning_sdk/lightning_cloud/utils/dataset.py
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,27 +1,28 @@
from typing import Callable, List, Optional, Union
import concurrent.futures
import json
import math
import os
import shutil
import tempfile
import zipfile
from concurrent.futures import ThreadPoolExecutor
from functools import partial
from pathlib import Path
from pathlib import Path, PurePosixPath, PureWindowsPath
from typing import Callable, List, Optional, Union

import backoff
import requests
from tqdm.auto import tqdm

from lightning_sdk.lightning_cloud.rest_client import LightningClient
from lightning_sdk.lightning_cloud.openapi.api_client import ApiClient
from lightning_sdk.lightning_cloud import env

from lightning_sdk.api.utils import (
_BYTES_PER_MB,
_MAX_BATCH_SIZE,
_MAX_SIZE_MULTI_PART_CHUNK,
_MAX_WORKERS,
_SIZE_LIMIT_SINGLE_PART,
)
from lightning_sdk.lightning_cloud import env
from lightning_sdk.lightning_cloud.openapi.api_client import ApiClient
from lightning_sdk.lightning_cloud.rest_client import LightningClient

# Total upload concurrency budget, split across files x within-file parts.
_DEFAULT_UPLOAD_WORKERS = 16
Expand All @@ -32,6 +33,65 @@
_DEFAULT_DOWNLOAD_PART_SIZE = 64 * 1024 * 1024 # 64 MiB byte-range parts


def _normalize_relative_path(path: str, source: str, strip_leading_slashes: bool = False) -> str:
"""Normalize a relative path and reject paths that could escape a destination."""
if not isinstance(path, str) or not path or "\x00" in path:
raise ValueError(f"Unsafe {source} path {path!r}: expected a non-empty relative path.")

if strip_leading_slashes:
path = path.lstrip("/")

# Dataset and ZIP paths use forward slashes, but validate backslashes as
# separators too so the same input cannot become unsafe on Windows.
normalized = path.replace("\\", "/")
posix_path = PurePosixPath(normalized)
windows_path = PureWindowsPath(path)
if (
posix_path.is_absolute()
or windows_path.is_absolute()
or windows_path.drive
or any(part == ".." for part in posix_path.parts)
):
raise ValueError(
f"Unsafe {source} path {path!r}: absolute paths and parent-directory traversal are not allowed."
)

relative_path = str(posix_path)
if relative_path in {"", "."}:
raise ValueError(f"Unsafe {source} path {path!r}: expected a non-empty relative path.")
return relative_path


def _safe_destination_path(destination: str, relative_path: str, source: str) -> Path:
"""Return a path contained by destination, accounting for existing symlinks."""
root = Path(destination).resolve()
target = root.joinpath(*PurePosixPath(relative_path).parts)
try:
target.resolve(strict=False).relative_to(root)
except ValueError as ex:
raise ValueError(f"Unsafe {source} path {relative_path!r}: path escapes the destination.") from ex
return target


def _extract_zip_safely(archive_path: str, destination: str) -> None:
"""Extract a ZIP after validating every member path."""
with zipfile.ZipFile(archive_path) as archive:
members = []
for member in archive.infolist():
relative_path = _normalize_relative_path(member.filename, "ZIP member")
target = _safe_destination_path(destination, relative_path, "ZIP member")
members.append((member, target))

Path(destination).mkdir(parents=True, exist_ok=True)
for member, target in members:
if member.is_dir():
target.mkdir(parents=True, exist_ok=True)
continue
target.parent.mkdir(parents=True, exist_ok=True)
with archive.open(member) as source_file, target.open("wb") as target_file:
shutil.copyfileobj(source_file, target_file)


def _parse_dataset_path(name: str) -> tuple:
"""Parse a dataset path like 'org/teamspace/dataset/version' into components.

Expand Down Expand Up @@ -61,9 +121,7 @@ def _parse_dataset_path(name: str) -> tuple:
return org_name, ts_name, dataset_name, version


def _resolve_dataset_id_and_version(
project_id: str, dataset_name: str, version: Optional[str] = None
) -> tuple:
def _resolve_dataset_id_and_version(project_id: str, dataset_name: str, version: Optional[str] = None) -> tuple:
"""List datasets once and return ``(dataset_id, resolved_version)``.

Combines name->id and default-version resolution into a single API round-trip
Expand Down Expand Up @@ -125,8 +183,13 @@ def _download_dataset_files(

def _to_entry(indexed_file: tuple) -> tuple:
index, file_info = indexed_file
raw_path = file_info.get("filepath")
if raw_path is None:
raw_path = f"file_{index}"
relative_path = _normalize_relative_path(raw_path, "dataset filepath", strip_leading_slashes=True)
return (
file_info.get("filepath", f"file_{index}").lstrip("/"),
relative_path,
_safe_destination_path(dest_dir, relative_path, "dataset filepath"),
file_info.get("url"),
int(file_info.get("size") or 0),
)
Expand All @@ -135,21 +198,19 @@ def _to_entry(indexed_file: tuple) -> tuple:
total_bytes = 0
urls = {}
tasks = []
# Lazily normalize each entry while pre-allocating its file and building the
# flat byte-range task list, so files_list is traversed only once.
for rel, url, size in entries:
# Pre-allocate each downloadable file and build the flat byte-range task list.
for rel, local_path, url, size in entries:
if not url:
continue
total_bytes += size
urls[rel] = url
local_path = os.path.join(dest_dir, rel)
os.makedirs(os.path.dirname(local_path) or dest_dir, exist_ok=True)
with open(local_path, "wb") as f:
local_path.parent.mkdir(parents=True, exist_ok=True)
with local_path.open("wb") as f:
if size:
f.truncate(size)
for start in range(0, size, part_size):
end = min(start + part_size, size) - 1
tasks.append((rel, local_path, start, end))
tasks.append((rel, str(local_path), start, end))

if not urls:
return
Expand Down Expand Up @@ -193,6 +254,7 @@ def _download_part(task: tuple) -> None:
finally:
pbar.close()


def _download_dataset_version(
project_id: str,
dataset_name: str,
Expand All @@ -202,29 +264,28 @@ def _download_dataset_version(
dataset_id: Optional[str] = None,
num_workers: int = _DEFAULT_DOWNLOAD_WORKERS,
part_size: int = _DEFAULT_DOWNLOAD_PART_SIZE,
):
"""
Download a dataset version as a zip file from the API.
unzip: bool = False,
) -> None:
"""Download a dataset version from the API.

Fetches presigned file URLs from the files endpoint and downloads the files
concurrently, each fetched with chunked HTTP-Range requests (see
``_download_dataset_files``), then packages them into a zip archive at
``target_path``.
``_download_dataset_files``). By default files are written into
``target_path``. With ``unzip``, the version must contain one ZIP artifact,
which is safely extracted into ``target_path``.

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: Destination directory.
cluster_id: Optional cluster ID.
dataset_id: The resolved dataset ID; when provided, skips an extra
datasets-list API round-trip to resolve the name.
num_workers: number of concurrent download threads (default 16).
part_size: byte-range part size for splitting large files (default 64 MB).
unzip: Extract a dataset stored as exactly one ZIP artifact.
"""
import shutil
import tempfile

cloud_url = env.LIGHTNING_CLOUD_URL
client = LightningClient(retry=False)
api_client: ApiClient = client.api_client
Expand Down Expand Up @@ -262,13 +323,26 @@ def _download_dataset_version(

try:
files_data = json.loads(resp.data) if resp.data else {}
except json.JSONDecodeError:
raise ValueError(f"Failed to parse response from {files_url}: {resp.data[:200]}")
except json.JSONDecodeError as ex:
raise ValueError(f"Failed to parse response from {files_url}: {resp.data[:200]}") from ex

files_list = files_data.get("files", [])
if not files_list:
raise ValueError(f"No files found for dataset '{dataset_name}' version '{version}' in project '{project_id}'.")

zip_relative_path = None
if unzip:
if len(files_list) != 1:
raise ValueError("`unzip=True` requires the dataset version to contain exactly one .zip artifact.")
raw_path = files_list[0].get("filepath")
if raw_path is None:
raw_path = "file_0"
zip_relative_path = _normalize_relative_path(raw_path, "dataset filepath", strip_leading_slashes=True)
if not zip_relative_path.lower().endswith(".zip"):
raise ValueError("`unzip=True` requires the dataset version to contain exactly one .zip artifact.")
if not files_list[0].get("url"):
raise ValueError("The dataset ZIP artifact does not have a download URL.")

# Re-fetch a fresh presigned URL for one file (used when a URL expires mid-download).
def _refresh_file(rel_path: str) -> dict:
r = api_client.request(
Expand All @@ -279,22 +353,30 @@ def _refresh_file(rel_path: str) -> dict:
_preload_content=True,
)
for f in (json.loads(r.data) if r.data else {}).get("files", []):
if f.get("filepath", "").lstrip("/") == rel_path:
filepath = f.get("filepath")
if (
filepath is not None
and _normalize_relative_path(filepath, "dataset filepath", strip_leading_slashes=True) == rel_path
):
return {"url": f.get("url"), "size": int(f.get("size") or 0)}
return {"url": None, "size": 0}

# Download in parallel into a temp dir, then package into a zip archive.
tmp_dir = tempfile.mkdtemp()
# ZIP extraction stages the archive in a temporary directory.
# Raw directory downloads write directly to the requested output directory.
staging = unzip
dest_dir = tempfile.mkdtemp(prefix="lightning-dataset-") if staging else target_path
if not staging:
os.makedirs(dest_dir, exist_ok=True)
try:
_download_dataset_files(
files_list, tmp_dir, project_id, _refresh_file, num_workers=num_workers, part_size=part_size
files_list, dest_dir, project_id, _refresh_file, num_workers=num_workers, part_size=part_size
)
base = target_path
if base.endswith(".zip"):
base = base[:-4]
shutil.make_archive(base, "zip", tmp_dir)
if unzip:
archive_path = _safe_destination_path(dest_dir, zip_relative_path, "dataset filepath")
_extract_zip_safely(str(archive_path), target_path)
finally:
shutil.rmtree(tmp_dir, ignore_errors=True)
if staging:
shutil.rmtree(dest_dir, ignore_errors=True)


def _create_dataset(
Expand Down
Loading
Loading