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 a built-in session that can build, run, monitor, and repair pipelines.

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

Installation

ATHENA_VERSION="$(athena --version | awk '{print $2}')"
pip install "buildathena-sdk==${ATHENA_VERSION}"

Requires Python 3.11+.

Quick Start

Define a block, read resolved config from ctx.config, emit metrics and progress, and register an explicit artifact when you want a durable named asset:

from athena import BlockContext, ConfigRef, block

@block(
    name="TrainModel",
    outputs=["checkpoint"],
    config=ConfigRef(search_path="conf", config_name="train"),
)
async def train_model(ctx: BlockContext) -> dict:
    epochs = int(ctx.config.get("epochs", 100))
    learning_rate = float(ctx.config.get("learning_rate", 1e-3))

    for epoch in range(epochs):
        loss = train_epoch(lr=learning_rate)
        await ctx.emit_metric("loss", loss, step=epoch)
        await ctx.emit_progress(epoch + 1, epochs)
        await ctx.check_pause()  # cooperative pause point

    checkpoint = await ctx.artifacts.register(
        "model.pt",
        name="checkpoint",
        format="pickle",
        mime_type="application/octet-stream",
        tags=["training", "final"],
    )
    return {"checkpoint": checkpoint.as_ref()}

Portable resource requests may include CPU, system memory, /dev/shm shared memory, ephemeral local storage, and accelerator requirements. memory is the total container RAM reservation/limit. shared_memory is only the /dev/shm mount-size ceiling: it may be used without memory or be larger than memory, and it does not create an additional physical RAM pool. When memory is omitted, RAM remains unreserved.

from athena import BlockEnvironment, ResourceSpec

environment = BlockEnvironment(
    backends=["docker", "kubernetes"],
    image="us-docker.pkg.dev/acme/athena/train:2026-08-24",
    resources=ResourceSpec(
        cpu="2",
        memory="4Gi",
        shared_memory="8Gi",
        ephemeral_storage="16Gi",
    ),
)

Docker supports shared_memory and intentionally does not match requests that contain ephemeral_storage. Kubernetes supports both. Process accepts resource-bearing work but does not enforce or reserve those fields. Omit process from BlockEnvironment.backends when enforcement is required.

Docker and Kubernetes mount the exact run code bundle at /athena/workspace, set ATHENA_WORKSPACE=/athena/workspace, and execute with that directory as the working directory. The Dockerfile owns dependency installation and environment layout; the runtime mount hides anything the image placed under the reserved path. Ensure the declared executable resolves through PATH or use an absolute path. BlockEnvironment.executable overrides the image's ENTRYPOINT and CMD on both backends, and Athena appends -m athena.harness. Include any required runtime wrapper explicitly in that prefix. Process workers set ATHENA_WORKSPACE to a fresh detached, attempt-owned worktree at the exact code_ref commit. An ignored authoring .venv is not copied there; reuse requires a prepared interpreter available on the execution endpoint.

Consuming Inputs

Block inputs come from the Python signature after ctx. Athena hydrates those arguments from upstream outputs before invoking the block:

@block(name="Evaluate", outputs=["report"])
async def evaluate(ctx: BlockContext, checkpoint) -> dict:
    checkpoint_ref = 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", outputs=["dataset"], secrets=["API_KEY"])
async def fetch_data(ctx: BlockContext) -> dict:
    key = ctx.secrets["API_KEY"]
    data = await download(api_key=key)
    return {"dataset": data}

Declare outputs on the decorator and return a mapping with exactly those keys.

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(epochs):
    train_step()
    await ctx.check_pause()  # yields control if a pause was requested

BlockContext API

as_data(), resolve(), and load() apply to managed artifacts. URI-backed external artifacts use as_ref() and expose their location through ref.uri for code that can access the bound store or filesystem.

Block-Scoped Athena Client

Athena injects the canonical Repo, Session, and private Chat client into an executing block. Chat collections paginate transparently, sending returns after durable acceptance, and receipt waits target only the submitted turn:

repo = await ctx.athena.repos.get("https://github.com/acme/training.git")
session = await ctx.athena.sessions.create(
    title="Candidate search",
    repos=[repo.at("candidate", new_branch=True)],
)

proposal_branch = session.repos[0].branch_name
chat = await session.chats.create(model_key="claude-opus-5")
receipt = await chat.send(
    f"""Athena prepared `{proposal_branch}` as a fresh proposal branch seeded from the
exact current head of `candidate`. Work directly in the prepared branch. After the
edit, use athena_attachment_git first with action `commit` and a concise message,
then with action `publish` and no message.""",
)
turn = await receipt.wait(timeout=None)

async for existing_chat in session.chats.list():
    print(existing_chat.id)

async for item in chat.items():
    consume(item)

async for item in turn.items():
    consume_turn_item(item)

# Stop interrupts current work but leaves the Chat open.
stop = await chat.stop(reason="operator requested")

# Close only after the Chat is quiescent, then read current state explicitly.
await chat.close()
closed = await session.chats.get(chat.id)
assert closed.closed_at is not None

# Exact lookup is the durable pointer; lookup alone does not reopen the Chat.
same_chat = await session.chats.get(closed.id)
await same_chat.resume()
resumed = await session.chats.get(same_chat.id)
assert resumed.closed_at is None

model_key is the stable model identity shown in Athena. With no connection, Athena uses its platform connection. To use a team connection, pass its permanent name, for example connection="Acme inference". Athena resolves that pair to an exact route when the Chat is created and keeps the route pinned on the Chat.

title is optional. If it is omitted, the server assigns a stable friendly title such as quiet-snail-cafe; the generated title is returned on the Session object and in browser projections. Titles are trimmed and limited to 255 characters.

Block code does not author runtime call keys or application idempotency keys. Invoke a Block or nested Workflow directly with child(...); there is no separate .call(...) authoring path. Session, Chat, and independent Workflow.run() operations likewise derive their internal replay identity. Ordinary mutations use a stable logical call slot; a response to a pending Chat request uses that exact request's identity. Keep construction order deterministic so the same source slot continues to represent the same logical operation after a controller retry. Browser-generated user intents and transport-level request keys are separate internal boundaries, not block SDK arguments.

repo.at(parent, new_branch=True) creates a fresh proposal branch from the exact current head of parent; the returned attachment reports the generated branch name. Work directly in that prepared branch. Agent-authored changes are committed with a message and then published without a message; reconciliation is only for a published remote head that changed after preparation.

chat.items() and turn.items() expose only the typed, durable, user-visible transcript projection. They are not live watches or raw event/trace access. Canonical values decode to ordinary Python primitives where lossless; inline bytes, opaque blobs, encoded values, and artifacts remain explicit InlineBytes, BlobRef, EncodedValue, and ArtifactRef wrappers rather than generated Protobuf messages or implicitly fetched bulk data.

chat.stop() durably requests interruption of work active at admission time; it does not close the Chat or prevent later messages and wakes. chat.close() is a synchronous lifecycle mutation that requires the Chat to have no active or queued agent work. While closed, that exact Chat rejects new ingress and cannot receive gate wakes, but it remains readable through session.chats.get(chat_id) and chat.items(). Only explicit chat.resume() reopens it. Closing one Chat does not disable a Session policy that is defined to create a different fresh Chat.

Independent Runs

Workflow.run() returns a durable Run handle. Lifecycle reads return an immutable RunSnapshot; outputs and metrics have focused accessors:

run = await candidate.workflow("workflows/evaluate.py:evaluate").run(profile=None)

snapshot = await run.wait(timeout=3600)
outputs = await run.result()
fitness = await run.metric("fitness")

print(snapshot.status, snapshot.is_terminal)
print(snapshot.source.repo_url)
print(snapshot.source.branch_name, snapshot.source.commit_sha)

This evaluation passes profile=None because the controller polls the Run directly and does not need profile gates or agent wakes.

RunSnapshot carries id, lab_id, workflow_key, status, is_terminal, failure, held_node_ids, created_at, started_at, completed_at, and source. Its source contains repo_url, branch_name, and commit_sha for that observation. run.source is the latest source observed by the handle.

Use status() for a current snapshot, wait() for a terminal snapshot, result() for hydrated named outputs, metric(name) for the latest metric, and kill() or release_held_nodes(...) for lifecycle controls.

Config

Blocks use Hydra-backed ConfigRef values. Discovery passes Hydra's exact composed mapping to ctx.config; it does not unwrap directory-named keys or pass config values as block arguments. Parameters after ctx remain workflow-wired dataflow inputs. Hydra owns composition and interpolation, while per-run block overrides live under exact root node IDs in RunProfile.config.nodes and are deep-merged over the discovered mapping.

from athena import ConfigRef, block

@block(
    config=ConfigRef(
        loader="hydra",
        search_path="configs",
        config_name="train",
    )
)
async def train(ctx):
    epochs = ctx.config["epochs"]

config_name is relative to search_path. For a flat mapping from configs/blocks/train.yaml, use search_path="configs/blocks" and config_name="train". Using search_path="configs" with config_name="blocks/train" intentionally follows Hydra package semantics and can produce a top-level blocks key.

See the config docs for discovery and launch layers.

Documentation

• Getting Started — installation and first workflow
• SDK Reference — decorators, config, and discovery
• Generated Python API Reference — exact public signatures, types, and defaults
• BlockContext API — full method reference
• Evolutionary Controllers — candidate branches, private Chats, evaluation Runs, and lineage
• Features — gates, inspectors, caching, and more
• knowledge-base/current/architecture/sdk.md — internal architecture and implementation reference

License

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