← Open source

torchestrator

PythonPyTorchDocker

torchestrator is a Dockerfile-like declarative experiment runner for PyTorch. One experiment is a "torchfile": a config that names the model, dataset, optimizer, loss, metrics and training loop in named blocks (MODEL, DATA, TRAIN, EVAL, ENV, LOGGING, METERS), resolved to a typed config and built by a dependency-injection engine over a name-based registry. Components talk through named slots via call / outs mappings, and shapes, dtypes and parameter counts are verified at build time before any training starts.

A torchfile

MODEL:
  type: Sequential
  call: { input: image }
  outs: [logits]
  data:
    layers:
      flatten: { type: Flatten }
      fc1: { type: Linear, data: { in_features: 784, out_features: 128 } }
      relu: { type: ReLU }
      fc2: { type: Linear, data: { in_features: 128, out_features: 10 } }

DATA:
  dataset: { type: MNIST, root: ./data, train: true }
  batch_size: 128
  outs: [image, label]

TRAIN:
  epochs: 10
  optimizer: { type: Adam, data: { lr: 0.001 } }
  loss:
    type: CrossEntropyLoss
    call: { input: logits, target: label }
    outs: [loss]

EVAL:
  metrics:
    - type: MulticlassAccuracy
      call: { input: logits, target: label }
      outs: [accuracy]

Layered FROM chains and a content-addressed hub

Torchfiles inherit with FROM (and override with DELETE), so configs stack as datasets -> baselines -> your experiment. A successful run pickles its live model, dataset, optimizer and scheduler into a hub at ~/.cache/torchestrator/, keyed by the torchfile's content hash; a child that FROMs a parent inherits those trained objects on cascade load, with no state_dict surgery or manual checkpoint paths. Named layers (hub://name) and remote pulls (hub+ssh://host/name) let you, for example, train on an HPC cluster and load the result on a laptop for eval, and a hub_ref marker splices a parent module into a child's model slot (e.g. SimCLR replace-the-head finetune).

Highlights

Auto-registration: drop .py files next to a torchfile and custom models,

losses, hooks and datasets are registered with no boilerplate.

Built-in extensions: optimizers (LARS, LAMB, Lion), losses (Focal, Dice,

Tversky, InfoNCE), schedulers (cosine/linear/poly warmup), hooks (EMA, SWA, early stopping, gradient clipping).

MODEL.meta.freeze freezes module prefixes; the optimizer is built only

over still-trainable parameters. Held-out EVAL.data runs a final no-grad pass.

Pluggable logging sinks (console, JSONL, CSV, TensorBoard, W&B, MLflow,

Logfire, OpenTelemetry) and a METERS section of probes for sub-step timing, GPU memory and gradient norms.

export emits a standalone, dependency-free PyTorch script (custom types

inlined via inspect.getsource).

Pinned Docker and Apptainer images (CUDA 11.8 / Python 3.10 / PyTorch

2.0.1) for CI, cloud GPUs and HPC.

CLI

Command	Purpose
`run <config>`	Run an experiment locally
`validate <config>`	Validate a config without running
`export <config>`	Generate a standalone PyTorch script
`status <config>`	Cache status for a torchfile / FROM chain
`cache info` / `cache purge`	Inspect or clear cached checkpoints
`plot <logfile>`	Plot training curves from JSONL / CSV logs
`diagnose`	Print Python / PyTorch / CUDA diagnostics

It depends on the sister projects torchdatasets and registry-pattern, and is the runner the paper-reproductions monorepo builds on. By design it is not an experiment tracker, hyperparameter sweeper or job scheduler: it wires out to W&B / MLflow for tracking, pairs with Optuna / Ray Tune for search, and delegates execution to srun / apptainer / docker. (v0.7.0, alpha; repository currently private.)

GitHub