Armory SDK

AI-Native MLOps Platform SDK for Python - A drop-in replacement for Weights & Biases.

Installation

# From GitHub (recommended)
pip install git+https://github.com/sionic-ai/armory.git#subdirectory=sdk/python

# With PyTorch Lightning support
pip install "armory-sdk[pytorch-lightning] @ git+https://github.com/sionic-ai/armory.git#subdirectory=sdk/python"

# With HuggingFace Transformers support
pip install "armory-sdk[transformers] @ git+https://github.com/sionic-ai/armory.git#subdirectory=sdk/python"

Configuration

Set the Armory server URL:

export ARMORY_API_URL=http://localhost:8080

Or in Python:

import armory
armory.setup(api_url="http://localhost:8080")

Quick Start

import armory

# Initialize a run
run = armory.init(
    project="my-embeddings",
    name="experiment-v1",
    config={
        "learning_rate": 1e-4,
        "batch_size": 32,
        "epochs": 10,
    }
)

# Log metrics during training
for epoch in range(10):
    train_loss = train_one_epoch()
    val_loss = evaluate()

    run.log({
        "epoch": epoch,
        "train_loss": train_loss,
        "val_loss": val_loss,
    })

# Finish the run
run.finish()

Logging API

Basic Metrics

# Single metric
run.log({"loss": 0.5})

# Multiple metrics
run.log({
    "loss": 0.5,
    "accuracy": 0.92,
    "learning_rate": 1e-4,
})

# With step number
run.log({"loss": 0.5}, step=100)

Rich Media

import armory

# Images
run.log({"sample": armory.Image(numpy_array)})
run.log({"sample": armory.Image("path/to/image.png")})
run.log({"sample": armory.Image(pil_image, caption="Generated sample")})

# Tables
run.log({
    "predictions": armory.Table(
        columns=["id", "label", "prediction", "confidence"],
        data=[
            [1, "cat", "cat", 0.95],
            [2, "dog", "cat", 0.51],
        ]
    )
})

# Histograms
run.log({"weights": armory.Histogram(model.layer.weight.detach().numpy())})
run.log({"activations": armory.Histogram(activations, num_bins=64)})

Training Events

Log training lifecycle events for real-time monitoring:

# Training lifecycle
run.log_event("train_start")
run.log_event("epoch_start", f"Starting epoch {epoch}")
run.log_event("epoch_end", metrics={"loss": 0.5, "accuracy": 0.9})

# Evaluation
run.log_event("eval_start")
run.log_event("eval_end", metrics={"ndcg@10": 0.85, "mrr": 0.72})

# Checkpoints
run.log_event("checkpoint", "Saved to checkpoints/epoch_10.pt")

# Custom events
run.log_event("custom", "Early stopping triggered", metrics={"patience": 5})

run.log_event("train_end")

Summary (Final Metrics)

# Set summary values (persisted as final results)
run.summary["best_accuracy"] = 0.95
run.summary["best_epoch"] = 8
run.summary["total_params"] = 110_000_000

# Access summary
print(run.summary["best_accuracy"])

Config

# Access config
print(run.config["learning_rate"])

# Update config (before training starts)
run.config.update({"warmup_steps": 1000})

Artifacts

# Log model artifacts
run.log_artifact("model.pt", type="model")
run.log_artifact("tokenizer/", type="tokenizer")  # Directory

# Log with metadata
run.log_artifact(
    "best_model.pt",
    type="model",
    metadata={"accuracy": 0.95, "epoch": 10}
)

Context Manager

with armory.init(project="my-project") as run:
    for step in range(100):
        run.log({"loss": 1.0 / (step + 1)})
# Automatically finished when exiting the context

Framework Integrations

PyTorch Lightning

from armory.integrations.lightning import ArmoryLightningCallback

trainer = pl.Trainer(
    max_epochs=10,
    callbacks=[
        ArmoryLightningCallback(
            project="my-project",
            name="lightning-run",
            config={"model": "bert-base"}
        )
    ]
)
trainer.fit(model, datamodule)

HuggingFace Transformers

from armory.integrations.transformers import ArmoryTrainerCallback

trainer = Trainer(
    model=model,
    args=training_args,
    train_dataset=train_dataset,
    callbacks=[
        ArmoryTrainerCallback(
            project="my-project",
            name="hf-training",
            config=dict(training_args)
        )
    ]
)
trainer.train()

Hyperparameter Sweeps

import armory

sweep_config = {
    "method": "bayes",  # "grid", "random", "bayes"
    "metric": {"name": "val_loss", "goal": "minimize"},
    "parameters": {
        "learning_rate": {"min": 1e-5, "max": 1e-2, "distribution": "log_uniform"},
        "batch_size": {"values": [16, 32, 64]},
        "epochs": {"value": 10},
    }
}

def train():
    run = armory.init()
    lr = run.config["learning_rate"]
    batch_size = run.config["batch_size"]

    # ... training code ...

    run.log({"val_loss": val_loss})
    run.finish()

# Run sweep
sweep_id = armory.sweep(sweep_config, project="my-project")
armory.agent(sweep_id, function=train, count=20)

W&B Migration Guide

Armory SDK is designed as a drop-in replacement for W&B:

# Before (W&B)
import wandb
wandb.init(project="my-project")
wandb.log({"loss": 0.5})
wandb.finish()

# After (Armory)
import armory
armory.init(project="my-project")
armory.log({"loss": 0.5})
armory.finish()

License

MIT