Examples

"""Example of using litlogger in a standalone script.

This example demonstrates the dict-like API:
- experiment["key"].append(value) to log time-series metrics
- experiment["key"] = "value" to set metadata
- experiment["key"] = File(path) to log static files
"""

import random
import time

import litlogger
from litlogger import File

# Initialize experiment
experiment = litlogger.init(
    name="standalone-demo",
    # teamspace="my-team",  # Optional: specify teamspace
)

# Set metadata
experiment["learning_rate"] = "0.001"
experiment["batch_size"] = "32"
experiment["model"] = "resnet50"

print("Training started...")

# Simulate a training loop
for epoch in range(10):
    for step in range(20):
        # Simulate training metrics
        loss = 2.0 * (1.0 / (epoch * 20 + step + 1))
        accuracy = min(0.95, (epoch * 20 + step) / 200.0)

        global_step = epoch * 20 + step
        experiment["train/loss"].append(loss + random.uniform(-0.1, 0.1), step=global_step)
        experiment["train/accuracy"].append(accuracy + random.uniform(-0.05, 0.05), step=global_step)

        time.sleep(0.01)  # Simulate training time

    # Log validation metrics at the end of each epoch
    val_loss = loss * 0.9
    val_accuracy = accuracy + 0.05

    val_step = (epoch + 1) * 20
    experiment["val/loss"].append(val_loss, step=val_step)
    experiment["val/accuracy"].append(val_accuracy, step=val_step)
    experiment["epoch"].append(epoch, step=val_step)

    print(f"Epoch {epoch + 1}/10 - Loss: {loss:.4f}, Accuracy: {accuracy:.4f}")

# Log a file (e.g., final config)
config_file = "/tmp/config.txt"
with open(config_file, "w") as f:
    f.write("learning_rate: 0.001\nbatch_size: 32\nmodel: resnet50\n")

experiment["config"] = File(config_file)

# Read back data
print(f"Model: {experiment['model']}")
print(f"Train losses logged: {len(experiment['train/loss'])}")
print(f"All metrics: {list(experiment.metrics.keys())}")

# Finalize the experiment
experiment.finalize()

"""Example of logging files, media, and models with the dict API."""

import base64
import tempfile
import uuid
from pathlib import Path

import litlogger
from litlogger import File, Image, Model, Text

experiment_name = f"file-media-model-demo-{uuid.uuid4().hex[:8]}"
workdir = Path(tempfile.mkdtemp(prefix="litlogger-file-media-model-"))


experiment = litlogger.init(name=experiment_name)
(workdir / "staging").mkdir(parents=True, exist_ok=True)

experiment["workflow"] = "new-api"
experiment["framework"] = "litlogger"
experiment["notes"] = "files, media, and models"
experiment["summary"] = Text("Static files, media series, metrics, and the new Model API.")

readme_path = workdir / "readme.txt"
readme_path.write_text("dict api example\n", encoding="utf-8")
experiment["artifacts/readme"] = File(str(readme_path))

png_path = workdir / "preview.png"
png_path.write_bytes(
    base64.b64decode("iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAQAAAC1HAwCAAAAC0lEQVR42mP8/x8AAusB9Wn4nWQAAAAASUVORK5CYII=")
)
experiment["media/preview"] = Image(str(png_path))

for step in range(3):
    loss = 1.0 / (step + 1)
    experiment["train/loss"].append(loss, step=step)

    report_path = workdir / f"report-{step}.txt"
    report_path.write_text(f"step={step}\nloss={loss:.4f}\n", encoding="utf-8")
    experiment["reports"].append(File(str(report_path)), step=step)

    experiment["captions"].append(Text(f"checkpoint at step {step}"), step=step)

(workdir / "model_artifact").mkdir(parents=True, exist_ok=True)
(workdir / "model_artifact" / "weights.bin").write_bytes(b"fake-weights")
(workdir / "model_artifact" / "config.json").write_text('{"layers": 2}\n', encoding="utf-8")
experiment[f"{experiment_name}/models/latest"] = Model(
    str(workdir / "model_artifact"),
    name=f"{experiment_name}-models-latest",
    version="artifact-v1",
)

experiment[f"{experiment_name}/models/object"] = Model(
    {"kind": "object", "layers": [4, 8, 2], "activation": "relu"},
    name=f"{experiment_name}-models-object",
    staging_dir=str(workdir / "staging" / "object"),
    version="object-v1",
)

checkpoint0 = workdir / "checkpoint-0.ckpt"
checkpoint0.write_text("checkpoint step 0\n", encoding="utf-8")
experiment[f"{experiment_name}/checkpoints"].append(
    Model(str(checkpoint0), name=f"{experiment_name}-checkpoints", version="v1"),
    step=0,
)

checkpoint1 = workdir / "checkpoint-1.ckpt"
checkpoint1.write_text("checkpoint step 1\n", encoding="utf-8")
experiment[f"{experiment_name}/checkpoints"].append(
    Model(str(checkpoint1), name=f"{experiment_name}-checkpoints", version="v2"),
    step=1,
)

print("metadata:", experiment.metadata)
print("artifacts:", sorted(experiment.artifacts))
print("metrics:", sorted(experiment.metrics))

experiment.finalize()

experiment = litlogger.init(name=experiment_name)
(workdir / "downloads").mkdir(parents=True, exist_ok=True)
(workdir / "downloads" / "model-cache").mkdir(parents=True, exist_ok=True)
(workdir / "downloads" / "checkpoints").mkdir(parents=True, exist_ok=True)

print(experiment["artifacts/readme"].save(str(workdir / "downloads" / "readme.txt")))
print(experiment["media/preview"].save(str(workdir / "downloads" / "preview.png")))
print(experiment[f"{experiment_name}/models/latest"].save(str(workdir / "downloads" / "artifact")))
print(experiment[f"{experiment_name}/models/object"].load(str(workdir / "downloads" / "model-cache" / "object")))

checkpoint_series = experiment[f"{experiment_name}/checkpoints"]
print("checkpoint series length:", len(checkpoint_series))
print(checkpoint_series[0].save(str(workdir / "downloads" / "checkpoints" / "checkpoint-0")))
print(checkpoint_series[1].save(str(workdir / "downloads" / "checkpoints" / "checkpoint-1")))
print(checkpoint_series[-1].save(str(workdir / "downloads" / "checkpoints" / "latest")))

experiment.finalize()

import os

import lightning as L
from lightning.pytorch.loggers import LitLogger
from torch import nn, optim
from torch.utils.data import DataLoader
from torchvision.datasets import MNIST
from torchvision.transforms import ToTensor


class LitAutoEncoder(L.LightningModule):
    def __init__(self, lr=1e-3, inp_size=28):
        super().__init__()

        self.encoder = nn.Sequential(nn.Linear(inp_size * inp_size, 64), nn.ReLU(), nn.Linear(64, 3))
        self.decoder = nn.Sequential(nn.Linear(3, 64), nn.ReLU(), nn.Linear(64, inp_size * inp_size))
        self.lr = lr
        self.save_hyperparameters()

    def training_step(self, batch, batch_idx):
        x, y = batch
        x = x.view(x.size(0), -1)
        z = self.encoder(x)
        x_hat = self.decoder(z)
        loss = nn.functional.mse_loss(x_hat, x)
        # log metrics
        self.log("train_loss", loss)
        return loss

    def configure_optimizers(self):
        return optim.Adam(self.parameters(), lr=self.lr)


if __name__ == "__main__":
    # init the autoencoder
    autoencoder = LitAutoEncoder(lr=1e-3, inp_size=28)

    # setup data
    train_loader = DataLoader(
        dataset=MNIST(os.getcwd(), download=True, transform=ToTensor()),
        batch_size=32,
        shuffle=True,
        num_workers=os.cpu_count(),
        persistent_workers=True,
    )

    # configure the logger
    lit_logger = LitLogger()

    # pass logger to the Trainer
    trainer = L.Trainer(
        limit_train_batches=100,
        max_epochs=3,
        logger=lit_logger,
    )

    # train the model
    trainer.fit(model=autoencoder, train_dataloaders=train_loader)

import time

import litgpt
import litlogger
import litserve as ls


class SimpleLitAPI(ls.LitAPI):
    def setup(self, device):
        self.llm = litgpt.LLM.load("EleutherAI/pythia-14m")
        self.experiment = litlogger.init(
            name="litserve-inference",
            metadata={"model": "pythia-14m"},
        )

    def decode_request(self, request):
        return request["prompt"]

    def predict(self, prompt):
        start = time.perf_counter()
        output = self.llm.generate(prompt, max_new_tokens=200)
        self.experiment["generation_time"].append(time.perf_counter() - start)
        self.experiment["input_tokens"].append(self.llm.preprocessor.encode(prompt).numel())
        self.experiment["output_tokens"].append(self.llm.preprocessor.encode(output).numel())
        return {"output": output}

    def encode_response(self, output):
        return {"output": output}

    def teardown(self, devices):
        self.experiment.finalize()


if __name__ == "__main__":
    server = ls.LitServer(SimpleLitAPI())
    server.run(port=8000)

curl -X POST http://127.0.0.1:8000/predict \
  -H "Content-Type: application/json" \
  -d '{"prompt": "What do llamas eat"}'

import json

import litlogger
from litlogger import File

config = {"model": "ResNet50", "lr": 0.001}
with open("config.json", "w") as f:
    json.dump(config, f)

experiment = litlogger.init(name="image-classifier")

# Set metadata
experiment["model"] = "ResNet50"
experiment["dataset"] = "CIFAR10"

# Log a static file
experiment["config"] = File("config.json")

# Log training metrics
for epoch in range(10):
    experiment["loss"].append(1.0 / (epoch + 1), step=epoch)

experiment.finalize()

"""Fetch data from a previously logged experiment.

This script resumes the 'image-classifier' experiment and reads back
metadata and artifacts using the dict-like API.
"""

import json

import litlogger

experiment = litlogger.init(name="image-classifier")

# Read metadata via property
metadata = experiment.metadata
print(f"Metadata: {metadata}")

# Read artifacts and download
artifacts = experiment.artifacts
if "config" in artifacts:
    artifacts["config"].save("/tmp/config.json")
    with open("/tmp/config.json") as f:
        print(f"Config: {json.load(f)}")

experiment.finalize()