Metadata-Version: 2.4
Name: research-pipelines
Version: 0.1.0
Summary: Lightweight research project pipeline framework with DAG tracing
Author: Leander Kurscheidt
License: MIT License
        
        Copyright (c) 2026 [fullname]
        
        Permission is hereby granted, free of charge, to any person obtaining a copy
        of this software and associated documentation files (the "Software"), to deal
        in the Software without restriction, including without limitation the rights
        to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
        copies of the Software, and to permit persons to whom the Software is
        furnished to do so, subject to the following conditions:
        
        The above copyright notice and this permission notice shall be included in all
        copies or substantial portions of the Software.
        
        THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
        IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
        FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
        AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
        LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
        OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
        SOFTWARE.
        
Classifier: Programming Language :: Python :: 3.11
Classifier: Programming Language :: Python :: 3.12
Classifier: License :: OSI Approved :: MIT License
Classifier: Operating System :: OS Independent
Requires-Python: >=3.11
Description-Content-Type: text/markdown
License-File: LICENSE
Provides-Extra: example
Requires-Dist: torch>=2.2.0; extra == "example"
Provides-Extra: viz
Requires-Dist: matplotlib; extra == "viz"
Requires-Dist: networkx; extra == "viz"
Provides-Extra: wandb
Requires-Dist: wandb>=0.15.0; extra == "wandb"
Provides-Extra: dev
Requires-Dist: pytest>=9.0; extra == "dev"
Requires-Dist: pytest-cov; extra == "dev"
Dynamic: license-file

# Research Pipelines

A lightweight Python framework for tracing the DAG (directed acyclic graph) of research experiments. Automatically track datasets, models, and evaluations arguments and function-dependencies, then persist everything to wandb or local storage. This is especially useful for plotting of further evaluation of a trained model, as we can recreate the a function call or just the arguments of a traced function. By design, it is a pickle-free solution that relies on recording primitve arguments.

Just decorate function during training like:
```python
@evaluation()
def evaluate(model_obj, metric: str):
    return {"score": 0.95}
```

It turns a huge, messy notebook into something simple like:

```python
# (select a traced run and load its saved configurations)
# rebuild the arguments such that we can call evaluate ourselves
model_obj, metric = query.build_arguments(
    evaluate
)
# load saved weights
model_obj.load_state_dict(state_dict)
# call evaluate
evaluate(model_obj, metric)
# do some plotting
```




## Features

- **Automatic DAG Tracing**: Decorators automatically detect when traced objects are used as dependencies
- **Configuration Persistence**: Basic types (str, int, float, bool, None) are automatically captured and stored
- **Flexible Rebuilding**: The query backend allows for calling the traced functions again, even if they depend on other traced functions
- **Pluggable Backends**: Use PickleBackend for testing or WandBBackend for production wandb integration
- **Zero Boilerplate**: Apply decorators and your functions/classes are automatically traced
- **Recursive Dependency Resolution**: Full transitive closure of all dependencies

## Installation (Dev)

```bash
# Clone or create the project
cd research_pipelines

# Create conda environment
conda create -n research_pipelines python=3.11

# Activate environment
conda activate research_pipelines

# Install package in editable mode
pip install -e .

# Optional: Install the Torch example extra
pip install ".[example]"

# Optional: Install wandb backend
pip install ".[wandb]"
```

## Quick Start

```python
from research_pipelines.decorators import dataset, model, evaluation
from research_pipelines.dag import build_dag

# Decorate your functions
@dataset()
def load_data(path: str, split: str):
    # Load your data...
    return {"data": [...], "metadata": {...}}

@model()
def train_model(train_data, architecture: str, lr: float):
    # Non-basic args (train_data) become dependencies
    # Basic args (architecture, lr) become config
    return trained_model

@evaluation()
def evaluate(model_obj, metric: str):
    return {"score": 0.95}

# Execute your pipeline
data = load_data(path="/data/train.csv", split="train")
model = train_model(train_data=data, architecture="bert", lr=0.001)
results = evaluate(model_obj=model, metric="accuracy")

# Print the DAG
dag = build_dag()
for obj_id, obj in dag.items():
    print(f"{obj['type']}: {obj['config']}, depends on: {obj['dependencies']}")
```

### Rebuild the traced object
The traced objects are not pickled, instead the arguments the functions are called with are saved.

```python
import research_pipelines.query as query

# we can now easily call the functions with the recorded arguments via build()
dataset = query.build(
    load_data
)

# or just get the arguments such that we can call it ourselves
model_obj, metric = query.build_arguments(
    evaluate
)
model_obj.load_state_dict(state_dict)
evaluate(model_obj, metric)
```

## How It Works

### 1. Decoration

Apply `@dataset()`, `@model()`, `@evaluation()`, or generic `@traced(traced_type="...")` to your functions or class constructors:

```python
@dataset()
def load_data(path: str, split: str):
    return load_from_disk(path)

@model()
class MyModel:
    def __init__(self, layers: int, dataset_input):
        self.layers = layers
        self.data = dataset_input
```

### 2. Automatic Tracing

When you call a decorated function/constructor:
- **Arguments are classified**:
  - **Basic types** (str, int, float, bool, None): stored as configuration
  - **Traced objects** (returned from other @traced functions): become dependencies
  - **Other types**: ignored (can be supplied manually later)
- **Unique ID** is generated for this object
- **Configuration** (basic args + type) is persisted to backend
- **Dependencies** (other traced object IDs) are recorded

### 3. DAG Structure

The framework automatically builds a DAG:
```
dataset_1 (config: path="/data/train.csv", split="train")
  ↓
model_1 (config: architecture="bert", lr=0.001, depends_on: [dataset_1])
  ↓
eval_1 (config: metric="accuracy", depends_on: [model_1])
```

### 4. Backend Persistence

Choose a backend to persist configurations:

**PickleBackend** (default for testing):
```python
from research_pipelines.backends.pickle_backend import PickleBackend
from research_pipelines.backends.manager import set_backend

backend = PickleBackend(directory=".traced_configs")
set_backend(backend)
```

**WandBBackend** (for wandb integration):
```python
import wandb
from research_pipelines.backends.wandb_backend import WandBBackend
from research_pipelines.backends.manager import set_backend

wandb.init(project="my_project")
backend = WandBBackend()
set_backend(backend)

# Configs are automatically logged to wandb.run.config
```

## API Reference

### Decorators

```python
from research_pipelines.decorators import dataset, model, evaluation, traced

@dataset()
def load_data(...):
    """Traces a dataset creation function/class."""
    pass

@model()
def train(...):
    """Traces a model creation function/class."""
    pass

@evaluation()
def eval(...):
    """Traces an evaluation function/class."""
    pass

@traced(traced_type="custom")
def my_function(...):
    """Generic tracer with custom type."""
    pass
```

### DAG Operations

```python
from research_pipelines.dag import (
    build_dag,
    get_dependencies_recursive,
    detect_circular_dependencies,
    export_dag,
    get_root_objects,
    get_leaf_objects,
    get_objects_by_type,
    get_dependents,
)

# Build full DAG
dag = build_dag()

# Get all transitive dependencies
deps = get_dependencies_recursive(object_id)

# Check for cycles
has_cycles = detect_circular_dependencies()

# Export for serialization
dag_export = export_dag()

# Find roots (datasets with no dependencies)
roots = get_root_objects()

# Find leaves (objects nothing depends on)
leaves = get_leaf_objects()

# Filter by type
datasets = get_objects_by_type("dataset")
models = get_objects_by_type("model")

# Find what depends on an object
dependents = get_dependents(object_id)
```

### Backends

```python
from research_pipelines.backends.manager import get_backend, set_backend

# Get active backend
backend = get_backend()

# Set custom backend
set_backend(my_backend)

# Backend interface
class Backend(ABC):
    def log_config(object_id, config_dict, dependencies):
        """Persist config for an object."""
        pass
    
    def get_config(object_id):
        """Retrieve config for an object."""
        pass
    
    def load_all():
        """Load all configs."""
        pass
    
    def clear():
        """Clear all configs."""
        pass
```

## Configuration Format

Configurations are stored as dictionaries with the following structure:

```python
{
    "object_id_1": {
        "callable": "examples.simple_pipeline:load_dataset"
        "config": {
            "path": "/data/train.csv",
            "split": "train",
            "batch_size": 32,
        },
        "dependencies": [],
    },
    "object_id_2": {
        "callable": "examples.simple_pipeline:create_model"
        "config": {
            "architecture": "bert",
            "learning_rate": 0.001,
        },
        "dependencies": ["object_id_1"],
    },
}
```

When using WandBBackend, this is stored directly in `wandb.run.config`.

## Examples

See [examples/simple_pipeline.py](examples/simple_pipeline.py) for a complete end-to-end example.

Run it:
```bash
conda activate research_pipelines
python examples/simple_pipeline.py
```

## Testing

All tests use PickleBackend and are fully isolated:

```bash
conda activate research_pipelines
pytest tests/ -v
```

## Development

The framework is organized into modules:

- `src/research_pipelines/core.py` - Core tracing logic
- `src/research_pipelines/decorators.py` - @dataset, @model, @evaluation decorators
- `src/research_pipelines/backends/` - Backend implementations
  - `base.py` - Abstract Backend interface
  - `pickle_backend.py` - PickleBackend (testing)
  - `wandb_backend.py` - WandBBackend (wandb integration)
  - `manager.py` - Global backend management
- `src/research_pipelines/dag.py` - DAG utilities
- `tests/` - Test suite (61 tests, all passing)

## Key Design Decisions

1. **Lazy Imports**: wandb is only imported when WandBBackend is used
2. **Automatic Dependency Detection**: Uses Python's `id()` to track object identity
3. **In-Memory Registry**: Separate from backend storage, enables DAG operations
4. **UUID v4 IDs**: Unique, collision-free object identifiers
5. **Type-Based Filtering**: Basic types automatically detected and persisted
6. **Pluggable Backends**: Easy to add custom storage implementations

## Limitations & Future Work

- No support for custom object serialization (by design)
- No execution timing/profiling (configuration-only tracking)
- No automatic versioning/hashing of objects

## License

MIT
