Skip to main content

Athena Labs Python SDK for agentic ML workflow orchestration

Project description

buildathena-sdk

PyPI version Python 3.11+ Status: Alpha

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 an explicit artifact when you want a durable named asset:

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

    checkpoint = await ctx.artifacts.register("model.pt", schema_type="checkpoint")
    return {"checkpoint": checkpoint.as_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:
    checkpoint_ref = ctx.inputs["checkpoint"]
    checkpoint_path = await ctx.artifacts.resolve(checkpoint_ref)
    model = load_model(checkpoint_path)
    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(source, schema_type=, format=, name=, tags=, metadata=) Create a durable artifact
artifact.as_ref() / artifact.as_data() Choose pointer or hydrated downstream delivery
await ctx.artifacts.resolve(ref) Resolve any Athena-readable artifact to a local path
await ctx.artifacts.load(ref) Load and deserialize a formatted 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

License

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

Project details


Download files

Download the file for your platform. If you're not sure which to choose, learn more about installing packages.

Source Distribution

buildathena_sdk-0.2.30.tar.gz (75.3 kB view details)

Uploaded Source

Built Distribution

If you're not sure about the file name format, learn more about wheel file names.

buildathena_sdk-0.2.30-py3-none-any.whl (48.0 kB view details)

Uploaded Python 3

File details

Details for the file buildathena_sdk-0.2.30.tar.gz.

File metadata

  • Download URL: buildathena_sdk-0.2.30.tar.gz
  • Upload date:
  • Size: 75.3 kB
  • Tags: Source
  • Uploaded using Trusted Publishing? No
  • Uploaded via: uv/0.8.0

File hashes

Hashes for buildathena_sdk-0.2.30.tar.gz
Algorithm Hash digest
SHA256 98cd58ed9df06e5d34878aada4f9e218333e71f3461d9b148cce029724fc2354
MD5 0eaa01c7d65fc9e12e7f473540bbfecb
BLAKE2b-256 960b6daab68ef52d81743a28c85b2dbe6f7e31a5b5f583cfcd16681448567452

See more details on using hashes here.

File details

Details for the file buildathena_sdk-0.2.30-py3-none-any.whl.

File metadata

File hashes

Hashes for buildathena_sdk-0.2.30-py3-none-any.whl
Algorithm Hash digest
SHA256 d6550ff1e151f5fd3249e9cf14aea511d1001de57b32527783a03f60f2273a2d
MD5 4fa7ab1e746d3a9d8e44774e68a78b0b
BLAKE2b-256 fe6bfc522e33cccfb2ff73d475810cfe599d9174420e3e60caa4b738c8c44694

See more details on using hashes here.

Supported by

AWS Cloud computing and Security Sponsor Datadog Monitoring Depot Continuous Integration Fastly CDN Google Download Analytics Pingdom Monitoring Sentry Error logging StatusPage Status page