Import datasets

Lightning AI offers various ways to import datasets to the platform.

Connect S3 buckets

Connect S3 buckets into Lightning AI to instantly access the data in a teamspace. This process takes a few minutes.

To connect an S3 bucket, click Add data on the top right, select AWS S3 Bucket and follow the instructions.

An s3 connection has access scoped to a teamspace for security reasons. This allows you to restrict data access to different team members.

Access the connected s3 data from any Studio in that teamspace (/teamspace/s3_connections).

Writing to S3 connections

Write access is on by default when you connect a bucket — it's the Write checkbox in the Add data form. If a connection is read-only, turn it on from the drive:

1. In the Drive, open the connection's Settings.
2. Under Permissions, switch on Write access.
3. Restart any running Studios for the change to take effect.

/teamspace/s3_connections/<name> is then writable from Studios and Jobs like any other folder. Publicly accessible buckets are the exception which can't be made writable.

Writing directly with boto3

To write to the bucket from your code instead of through the drive, use the boto3 library.

Build the client explicitly. Some Studios come with an AWS profile pointing at Lightning Storage rather than Amazon, and boto3 will apply it, rejecting a valid AWS key with InvalidAccessKeyId. Always pass the bucket's region, endpoint and credentials.

import os
import boto3

s3 = boto3.client(
    's3',
    region_name='us-east-1',
    endpoint_url='https://s3.us-east-1.amazonaws.com',
    aws_access_key_id=os.environ['AWS_ACCESS_KEY_ID'],
    aws_secret_access_key=os.environ['AWS_SECRET_ACCESS_KEY'],
)

s3.upload_file(Filename='demo.txt', Bucket='s3-demo-bucket', Key='test')

Upload via web app

Upload datasets up to 200GB using the Lightning Drive on the web app.

For uploads exceeding 200GB, the Drive may work, but we recommend to use the SDK as described below.

Upload via SDK

Upload datasets programmatically with the SDK:

With studio.upload_file(path_to_local_file, path_to_remote_file) you can upload files. This allows for programmatic filtering or parallelizing uploads. A full example is:

#! pip install lightning_sdk
# LIGHTNING_USER_ID=000000-000000-000000-00000-000000000
# LIGHTNING_API_KEY=111111-111111-111111-11111-111111111
import os
from lightning_sdk import Studio

studio = Studio(name="my-studio-name", teamspace="my-teamspace", user="my-user")
total_root = "path/to/demo"
for root, _, files in os.walk(total_root):
    for f in files:
        # upload in remote dir where everything will be relative to home
        remote_path = os.path.join(os.path.relpath(root, total_root), f)
        studio.upload_file(os.path.join(root, f), remote_path)

Here, the code iterates over all files in the demo path (directories will be created on the fly if necessary). When uploading a file, the remote path is always relative to the home directory and must be a subdirectory of it. By specifying the remote_path as os.path.join(os.path.relpath(root, total_root), f), every file will be uploaded to the path relative to the root of the path we want to upload under the Studio's home directory.

Download third-party data

To download data from services like Hugging Face datasets, Kaggle, or others—consider using a Data Prep machine with at least 32 CPUs and 3TB storage. After downloading and processing, put the Studio to sleep.

No need to restart the Studio unless you want to change the underlying data.

Example: Kaggle

This example shows how to download any Kaggle dataset to a Studio.

Access the dataset from any other Studio in that teamspace.

⚡ ~ python
Python 3.10.10 (main, Mar 21 2023, 18:45:11) [GCC 11.2.0] on linux
Type "help", "copyright", "credits" or "license" for more information.
>>> import os
>>> os.listdir("/teamspace/studios/imagenet-1m/data")
['test', 'train', 'val']
>>> os.listdir("/teamspace/studios/imagenet-1m/data/train")
['n01440764', 'n01443537', 'n01484850', 'n01491361', ..., 'n0187240']

Example: Huggingface

Use any Hugging Face dataset in Studios by installing the datasets library.

⚡ ~ pip install datasets

Add the cache_dir argument to the load_dataset for the data to be downloaded and persisted within the Studio.

from datasets import load_dataset

dataset = load_dataset("open_subtitles", lang1="fi", lang2="hi", cache_dir="./data")

When running the code above, the data gets downloaded within the Studio.

⚡ ~ python main.py
You can avoid this message in future by passing the argument `trust_remote_code=True`.
Passing `trust_remote_code=True` will be mandatory to load this dataset from the next major release of `datasets`.
  warnings.warn(
Downloading builder script: 100%|█████████████████████████████████████████████████████████████████████████████████████████████████| 6.22k/6.22k [00:00<00:00, 19.8MB/s]
Downloading readme: 100%|█████████████████████████████████████████████████████████████████████████████████████████████████████████| 7.45k/7.45k [00:00<00:00, 25.5MB/s]
Downloading data: 100%|████████████████████████████████████████████████████████████████████████████████████████████████████████████| 2.00M/2.00M [00:02<00:00, 905kB/s]
Generating train split: 55597 examples [00:03, 14771.18 examples/s]

Some datasets require your access token from Hugging Face.