Background

Currently, data transfer between tasks lacks proper abstraction:

• Manual data handling in scripts
• No automatic data transfer
• Lack of transparency in data paths
• Incomplete cost calculation

Proposed Approaches

Approach 1: Task Output Registration

This approach introduces automatic data transfer by establishing default paths and enabling customization with set_output and set_input. Environment variables provide easy path access, reducing hardcoded paths.

Custom Path Configuration with set_output and set_input

For cases where users prefer not to modify existing scripts with fixed paths, set_output and set_input allow specifying custom paths directly in the DAG. Suppose preprocess.py saves data to /mnt/volume1/features.parquet and train.py expects the data in /mnt/data/features.parquet:

# preprocess.py - Fixed path
data.save('/mnt/volume1/features.parquet')

# train.py - Fixed path
features = pd.read_parquet('/mnt/data/features.parquet')

Previously, a manual script was required to move data between clusters, but now we can configure this directly and enables automatic data transfer without modifying existing scripts:

# DAG configuration with custom paths
preprocess.set_output('/mnt/volume1', estimated_size_gb=2.0)
train.set_input(preprocess, '/mnt/data')

preprocess >> train

Default Path Mapping with Environment Variables

For newly created scripts, the default path setup minimizes the need for custom configurations. By default, each task binds its output to .sky/{task_name}/output, making this path available to downstream tasks. In the following DAG, preprocess has its output bound to .sky/preprocess/output, and this same path structure is mapped to train as its input:

preprocess = Task(name='preprocess', run='python preprocess.py')
train = Task(name='train', run='python train.py')
preprocess.set_output(estimated_size_gb=2.0)
preprocess >> train

In this setup, users should configure preprocess.py to write output data to .sky/preprocess/output and train.py to read from this location. Additionally, environment variables provide dynamic access to paths, allowing user scripts to remain unaffected by any changes to input or output paths set in the DAG. By referencing TASK_{TASK_NAME}_OUTPUT_PATH, users avoid hardcoded paths, creating a setup where scripts are path-agnostic:

# preprocess.py - Using environment variable for output path
output_dir = os.environ['TASK_PREPROCESS_OUTPUT_PATH']
data.save(f'{output_dir}/features.parquet')

# train.py - Using environment variable for input path
input_dir = os.environ['TASK_PREPROCESS_OUTPUT_PATH']
features = pd.read_parquet(f'{input_dir}/features.parquet')

This configuration allows users to modify paths solely within the DAG while scripts automatically adapt, providing seamless data transfer with minimal setup.

Pros & Cons

✓ Simple, intuitive API
✓ Transparent path handling
✓ Automatic data transfer
✓ Clean separation of mechanism and convenience
✗ Lacks support for partitioned outputs to multiple downstreams

Approach 2: Edge-Based Data Flow

Specify data transfer on edges between tasks, enabling different paths for different downstream tasks.

# DAG definition
preprocess = Task(name='preprocess', run='python preprocess.py')
train_a = Task(name='train_a', run='python train.py')
train_b = Task(name='train_b', run='python train.py')

(preprocess >> train_a).with_data('/data/model_a', size_gb=2.0)
(preprocess >> train_b).with_data('/data/model_b', size_gb=2.0)

# train.py
input_path = os.environ['EDGE_DATA_PATH']
data = pd.read_parquet(input_path)

Pros & Cons

✓ Explicit data routing
✓ Different paths for different downstream tasks
✓ Clear data flow visualization
✗ More verbose for simple cases

Recommendation

Approach 1 provides the best balance of simplicity and functionality.

Possible Extensions & Discussion

Dynamic Resource Requirements: Scale downstream resources based on actual data size. Particularly useful when data size varies significantly between runs.

train.set_resources(lambda size_gb: Resources(
    memory=f'{size_gb * 2}GB'
))

Python-Level Data API: Replace file operations with Python objects for cleaner data access.

data = task.get_data('features')
processed = transform(data)
task.write_data(processed)

Computation Hiding: Built on Python-level API, enable downstream tasks to start as soon as upstream produces partial results.

# Python API enables batch-level control
loader = task.get_data('features').as_stream()
for batch in loader:  # Upstream compute -> transfer -> downstream compute
    train(batch)

Data Contracts: Validate and filter data dependencies between tasks. Also enhances approach 1 with a cleaner abstraction for multi-downstream data routing.

@inputs
class TrainData:
    features: Data = path('features/*')
    metrics: Data = path('metrics.json')

train.use_inputs(TrainData.features)  # Validates and filters
eval.use_inputs(TrainData.metrics)    # Different data for different tasks

Discussion

We welcome feedback on the proposed designs and extensions.