axl-workflows 0.6.0

AXL Workflows (axl) is the foundation layer of the AXL AI Platform — a full AI platform from development to deployment. Define data and ML workflows as plain Python classes; axl compiles them to a backend-agnostic IR, packages them as a portable .axlp bundle, and ships them to any runtime provider.

Build once, ship everywhere.

• Local runtime → fast iteration on your machine.
• Argo Workflows → production Kubernetes pipelines.
• Kubeflow Pipelines (v2) → KFP- and Vertex-compatible execution.

Write once → compile to a portable IR → run anywhere. No YAML, no vendor lock-in. Portability is the point: the IR and the .axlp pack are the product. See the vision and roadmap.

axl-workflows is the foundation layer of the AXL AI Platform, and the only layer that exists today.

🚀 Quick Start

# Install
pip install axl-workflows

# Or with uv
uv pip install axl-workflows

# Create your first workflow
axl --help

✨ Key Features

• Class-based DSL: Define workflows as Python classes, with steps as methods and a dag() to wire them.

• Simple params: Declare defaults once on the class (params = {"lr": 0.1, "epochs": 3} — the defaults fix the types) and override them per run without a rebuild: a Run Config's params:, or axl submit --param epochs=10. Steps read self.workflow_params; the params-as-a-step pattern (a step returning a Pydantic model or dict built from them) keeps working. No special Param/Artifact classes; unknown names and wrong types are refused before anything runs.

• IO Handlers: Steps return plain Python objects; axl serializes them via an io_handler — pickle (default), cloudpickle, json, or your own.

  • Per-step override (@step(io_handler="json"))
  • Input modes: input_mode="path" is accepted but not wired through the runtimes yet (#14) — steps currently always receive objects.

• Intermediate Representation (IR): Backend-agnostic DAG model (nodes, edges, resources, IO metadata).

• Multiple backends:

  • Local runtime → develop and iterate quickly.
  • Argo Workflows → YAML generation for production Kubernetes pipelines.
  • Kubeflow Pipelines (v2) → KFP- and Vertex-compatible pipeline packages.

• Unified runner image: One container executes steps locally and in Argo pods.

• Resource & retry hints: Declare CPU, memory, caching, retries, and conditions at the step level.

• CLI tools: Compile, validate, pack, bake images, and run locally.

• Run it on a cluster, from axl (v0.5.0): axl submit hands a compiled artifact to a Kubeflow Pipelines cluster and returns a run handle; axl status and axl logs follow it without kubectl or the KFP UI. Retried submissions re-attach to the in-flight run instead of launching a duplicate. (Artifact fetch on Kubeflow is not available yet — axl outputs refuses loudly rather than returning an empty directory.)

• Artifact Stores and the Run Config (v0.6.0): where a run's artifacts live is environment configuration, not workflow code. A repo-checked --config file selects the store (local, S3, MinIO/seaweedfs, GCS) plus workflow-wide resource/retry/env defaults, and the same workflow compiles unchanged for dev and prod. On Kubeflow the store becomes the pipeline root. (ADR-0013, contract)

• axl.sdk, frozen (v0.5.0): the same lifecycle as typed Python — build_ir, validate, compile, pack, submit, status, logs, outputs — with frozen result models and a stability contract, so other tools build on axl without shelling out. See the SDK guide.

📦 Example Workflow (params as a step, with Pydantic)

# examples/churn_workflow.py
from axl import Workflow, step
from pydantic import BaseModel

# Parameters are just a normal step output (typed with Pydantic for convenience).
class TrainParams(BaseModel):
    seed: int = 42
    input_path: str = "data/raw.csv"

class ChurnTrain(Workflow):
    # Workflow configuration via class attributes
    name = "churn-train"
    image = "ghcr.io/axl-platform/axl-workflows/runner:0.6.0"
    io_handler = "pickle"

    @step
    def params(self) -> TrainParams:
        # Use defaults here; optionally read from YAML/env if you prefer.
        return TrainParams()

    @step  # default io_handler = pickle
    def preprocess(self, p: TrainParams):
        import pandas as pd
        df = pd.read_csv(p.input_path)
        # ... feature engineering ...
        return df  # persisted via pickle (default)

    @step
    def train(self, features, p: TrainParams):
        from sklearn.ensemble import RandomForestClassifier
        import numpy as np
        y = (features.sum(axis=1) > features.sum(axis=1).median()).astype(int)
        X = features.select_dtypes(include=[np.number]).fillna(0)
        model = RandomForestClassifier(n_estimators=50, random_state=p.seed).fit(X, y)
        return model  # persisted via pickle

    @step
    def evaluate(self, model) -> float:
        # pretend evaluation
        return 0.9123

    def dag(self):
        p = self.params()
        feats = self.preprocess(p)
        model = self.train(feats, p)
        return self.evaluate(model)

Variations

• Receive a file path instead of an object:

  from pathlib import Path
  
  @step(input_mode={"features": "path"})
  def profile(self, features: Path) -> dict:
      return {"bytes": Path(features).stat().st_size}
  

• Override the io handler per step. A handler is a codec over a binary stream — it never sees a path or a store — and registering one by name makes it usable in @step. Register it in the workflow's module, so the pod runtime sees it too:

  import pandas as pd
  from axl.io import BaseIOHandler, registry
  
  class ParquetIOHandler(BaseIOHandler):
      def dump(self, obj, stream):
          obj.to_parquet(stream)
  
      def load(self, stream):
          return pd.read_parquet(stream)
  
  registry.register("parquet", ParquetIOHandler)
  
  @step(io_handler="parquet")
  def preprocess(self, p: TrainParams):
      return pd.read_csv(p.input_path)  # stored as Parquet; downstream gets a DataFrame
  

• Choose where artifacts live per environment, without touching the workflow:

  # configs/prod.yaml
  config_version: "1.0"
  artifact_store: {type: s3, bucket: ml-artifacts, prefix: churn}
  defaults:
    resources: {cpu: "1", memory: 2Gi}   # @step(resources=...) still wins per step
    retries: 2
  params:
    seed: 7
  

🛠 CLI

# Compile to Argo Workflows YAML
axl compile -m examples/churn_workflow.py:ChurnTrain --target argo --out churn.yaml

# Compile to Kubeflow Pipelines v2 package
axl compile -m examples/churn_workflow.py:ChurnTrain --target kfp --out pipeline.yaml

# Run locally
axl run local -m examples/churn_workflow.py:ChurnTrain

# Same, with a Run Config: artifact store, step defaults, params (v0.6.0)
axl compile -m examples/churn_workflow.py:ChurnTrain --target kfp --config configs/prod.yaml --out pipeline.yaml
axl run local -m examples/churn_workflow.py:ChurnTrain --config configs/dev.yaml

# Change a remote run's parameters at submit time — no rebuild (validated against the declaration)
axl submit pipeline.yaml --target kfp --endpoint http://localhost:8080 --param seed=11 --param model=knn

# Validate workflow definition (with --config, prints the resolved per-step values)
axl validate -m examples/churn_workflow.py:ChurnTrain

# Package a portable .axlp (IR + runtime metadata + optional source)
axl pack -m examples/churn_workflow.py:ChurnTrain --out churn.axlp

Running on a cluster (v0.5.0)

compile seals an artifact; submit hands that artifact — never your source — to a provider, so what you inspected is byte-for-byte what runs (ADR-0007). The submitting machine needs only the file and a flag:

# Bake your code into an image, then compile against it
axl build-image examples/churn_workflow.py:ChurnTrain --tag <registry>/churn:v1 --push
axl compile -m examples/churn_workflow.py:ChurnTrain --target kfp \
    --image-mode baked --image-tag <registry>/churn:v1 \
    --config configs/prod.yaml --out pipeline.yaml   # store → pipeline root

# Hand it to a Kubeflow Pipelines cluster
axl submit pipeline.yaml --target kfp --endpoint http://localhost:8080

axl status                 # recent runs
axl status <run-id>        # one run, in axl's own vocabulary
axl logs <run-id>          # what your steps printed

axl run stays local, blocking, and in-process — it never submits, and it rejects a compiled artifact with a pointer to submit.

With a Run Config's artifact_store, a Kubeflow run writes its artifacts under that store's prefix instead of the deployment default. Argo cannot express a store yet and refuses at compile time (#112). Changing only the store does not re-materialize cached steps — the cache is keyed on the step, not on where its outputs go.

Provider support today: Kubeflow Pipelines is Certified for submit / status / logs, proven by a conformance suite and a local≡cluster parity test that run nightly against a real cluster. axl outputs (artifact fetch) is not available on Kubeflow yet — KFP assigns artifact locations at runtime and exposes them through no API, so the command refuses with an explanation instead of returning nothing (ADR-0010). Argo remains a Compile-only target: compile with axl, submit with argo submit.

Cluster lifecycle and storage setup are out of scope for axl. Today, follow the upstream install steps in the Argo and KFP guides.

📐 Architecture

axl-workflows is Layer 1 of the AXL AI Platform:

┌─────────────────────────────────────────────────┐
│  LAYER 4: MONITOR                 not started   │
├─────────────────────────────────────────────────┤
│  LAYER 3: SERVE                   not started   │
├─────────────────────────────────────────────────┤
│  LAYER 2: MANAGE                  not started   │
├─────────────────────────────────────────────────┤
│  LAYER 1: AUTHOR → COMPILE → RUN  ← (here)      │
├─────────────────────────────────────────────────┤
│  OPS (cross-cutting)  axlctl      not started   │
└─────────────────────────────────────────────────┘

Only Layer 1 exists today; the rest is stated direction.

Within this repo, the layers are:

1. Authoring Layer

   • Python DSL: @step decorator, Workflow base class
   • Params are a normal step (often a Pydantic model)
   • Configuration via class attributes (name, image, io_handler)
   • IO handled by io_handlers (default: pickle) — how an object becomes bytes
   • Run Config (--config) — where artifacts live, step defaults, params; never what is computed
   • Wire dependencies via dag() (auto-inferred in v0.3.0+)

2. IR (Intermediate Representation)

   • Backend-agnostic DAG: nodes, edges, inputs/outputs, resources, retry policies, IO metadata

3. Compilers

   • Argo: IR → Argo Workflow YAML
   • KFP: IR → Kubeflow Pipelines v2 package
   • Plugin architecture — add any target via entry points

4. Runtime

   • Unified runner image (axl-runner) executes steps in pods and locally
   • Handles env (via uv), IO handler dump/load into an Artifact Store, structured logging, retries

5. CLI

   • axl compile, axl run local, axl validate
   • axl pack, axl build-image
   • axl submit, axl status, axl logs, axl outputs (v0.5.0)
   • --config on compile, validate, run local (v0.6.0)

📂 Project Structure

axl/
  core/          # DSL: decorators, base classes, typing
  io/            # io_handlers (pickle, cloudpickle, json) + artifact stores (local, memory, gcs, s3/minio)
  run_config.py  # Run Config loader (--config)
  ir/            # Intermediate Representation (nodes, edges, workflows)
  compiler/      # Backend compilers (Argo, Kubeflow)
  runtime/       # Runner container + IO + env setup (uv)
  cli.py         # CLI entrypoint
examples/
  churn_workflow.py
tests/
  test_core.py   # Tests for DSL components
  test_ir.py     # Tests for IR components
pyproject.toml
README.md

🎯 Why AXL Workflows?

• Local development is fast and simple.

• Argo/KFP is production-grade but YAML is verbose and hard to get started with.

• axl bridges the gap:

  • Simple, class-based DSL — no YAML, no vendor-specific decorators
  • Params as a normal step — no special Param/Artifact classes
  • IO handlers for painless object ↔ file persistence
  • Backend-agnostic IR — one workflow definition, multiple compile targets
  • Compile once, run anywhere

📄 License

This project is licensed under the MIT License - see the LICENSE file for details.