buildathena-sdk 0.1.40

Athena Labs Python SDK for agentic ML workflow orchestration

buildathena-sdk

Python SDK for building blocks and workflows on the Athena Labs ML orchestration platform. Athena Labs makes ML workflows reproducible, observable, interruptible, and composable through a DAG execution engine with an AI agent that can build, run, monitor, and fix pipelines.

Alpha software — APIs may change between releases. Pin your version in production.

Installation

pip install buildathena-sdk

Requires Python 3.11+.

Quick Start

Define a block with a Pydantic config class, emit metrics and progress, and register artifacts:

from pydantic import BaseModel
from athena import block, BlockContext

class TrainConfig(BaseModel):
    epochs: int = 100
    learning_rate: float = 1e-3

@block(name="TrainModel", outputs=["checkpoint"], config=TrainConfig)
async def train_model(ctx: BlockContext, config: TrainConfig) -> dict:
    for epoch in range(config.epochs):
        loss = train_epoch(lr=config.learning_rate)
        await ctx.emit_metric("loss", loss, epoch=epoch)
        await ctx.emit_progress(epoch + 1, config.epochs)
        await ctx.check_pause()  # cooperative pause point

    ref = await ctx.artifacts.register("model.pt", schema_type="checkpoint")
    return {"checkpoint": ref}

Consuming Inputs

Blocks declare their inputs and access upstream outputs through ctx.inputs:

@block(name="Evaluate", inputs=["checkpoint"], outputs=["report"])
async def evaluate(ctx: BlockContext, config) -> dict:
    model = load_model(ctx.inputs["checkpoint"])
    score = run_eval(model)
    await ctx.emit_metric("accuracy", score)
    return {"report": {"accuracy": score}}

Credentials

Declare required secrets in the @block decorator and access them at runtime via ctx.secrets. Credentials are encrypted at rest and injected only during execution:

@block(name="FetchData", secrets=["API_KEY"], outputs=["dataset"])
async def fetch_data(ctx: BlockContext, config) -> dict:
    key = ctx.secrets["API_KEY"]
    data = await download(api_key=key)
    return {"dataset": data}

Cooperative Pause

Call check_pause() inside long-running loops to let Athena Labs pause the block between iterations without losing progress:

for epoch in range(config.epochs):
    train_step()
    await ctx.check_pause()  # yields control if a pause was requested

BlockContext API

Method / Accessor	Description
ctx.inputs["name"]	Access upstream block outputs
ctx.secrets["KEY"]	Access declared secrets
await ctx.emit_metric(name, value, epoch=, step=, **labels)	Emit a named metric
await ctx.emit_progress(current, total, message=)	Emit progress (current/total)
await ctx.emit_event(type, name, payload)	Emit a custom event
await ctx.check_pause()	Cooperative pause checkpoint
await ctx.artifacts.register(path, schema_type=, name=, tags=, metadata=)	Register an artifact

Config System

Athena Labs supports YAML configuration with composition, inheritance, and variable substitution:

# base.yaml
training:
  epochs: 100
  optimizer: adam

# experiment.yaml
$extends: base.yaml
training:
  epochs: 200
  learning_rate: ${LR:-1e-3}

See the full config docs for $include, $extends, and ${ref} substitution.

Other Decorators

@handler — Custom Storage Backends

Resolve custom URI schemes (e.g., s3://, gs://) for artifact storage:

from athena import handler

@handler(schemes=["s3"], name="S3Handler")
class S3Handler:
    async def get(self, key: str) -> bytes: ...
    async def put(self, key: str, data: bytes) -> str: ...

@inspector_backend — Interactive UI Plugins

Build interactive inspector panels that attach to running workflows:

from athena import inspector_backend

@inspector_backend
class LossPlotBackend:
    async def initialize(self, ctx) -> None: ...
    async def handle_message(self, msg: dict) -> dict: ...
    async def cleanup(self) -> None: ...

Documentation

• Getting Started — installation and first workflow
• SDK Reference — decorators, config, and discovery
• BlockContext API — full method reference
• Features — gates, inspectors, caching, and more

License

Proprietary - Copyright (c) 2026 Athena Labs Research Inc. All rights reserved.