diff --git a/.cspell.json b/.cspell.json index 4354043..2e53f73 100644 --- a/.cspell.json +++ b/.cspell.json @@ -52,6 +52,14 @@ "stdenv", "demisto", "toomanyrequests", - "envrc" + "envrc", + "MNIST", + "torchvision", + "makedirs", + "optim", + "argmax", + "randperm", + "shutil", + "pytest" ] } diff --git a/README.md b/README.md index 3022521..3bd3733 100644 --- a/README.md +++ b/README.md @@ -14,12 +14,14 @@ You are encouraged to fork the repo and experiment with adjusting the workflows ## Examples -| Example Name | Description | -|------------------------------------------------------------|--------------------------------------------------------------------------------------------------------------------------------------------| -| [cronworkflow-example]( examples/cronworkflow-example/ ) | Runs a basic CronWorkflow. | -| [dag-diamond]( examples/dag-diamond/ ) | Runs a basic DAG Workflow. | -| [external-logs]( examples/external-logs/ ) | Workflows deploys a kubernetes job. You can see the logs from the job within Pipekit. | -| [fan-out-fan-in]( examples/fan-out-fan-in/ ) | Shows how S3 artifact processing can be parallelized with Argo Workflows using a fan-out approach. | -| [get-versions](examples/get-versions/) | A workflow that outputs the versions of software installed in a Pipekit free trial cluster. Available as both a Hera and a native Workflow | -| [hera-coinflip](examples/hera-coinflip/) | Runs a basic coinflip example using the Hera Python framework. | -| [hera-notebook-forecast](examples/hera-notebook-forecast/) | Press-play data job from a Jupyter notebook: aggregate energy demand and forecast the next day. Shows a platform abstraction over Hera. | +| Example Name | Description | +| ---------------------------------------------------------- | ----------------------------------------------------------------------------------------------------------------------------------------------------- | +| [ci-pr-check](examples/ci-pr-check/) | A pull-request check pipeline: checkout, then lint and unit tests in parallel, with a combined status on exit. Runs in the free trial cluster. | +| [cronworkflow-example](examples/cronworkflow-example/) | Runs a basic CronWorkflow. | +| [dag-diamond](examples/dag-diamond/) | Runs a basic DAG Workflow. | +| [external-logs](examples/external-logs/) | Workflows deploys a kubernetes job. You can see the logs from the job within Pipekit. | +| [fan-out-fan-in](examples/fan-out-fan-in/) | Shows how S3 artifact processing can be parallelized with Argo Workflows using a fan-out approach. | +| [get-versions](examples/get-versions/) | A workflow that outputs the versions of software installed in a Pipekit free trial cluster. Available as both a Hera and a native Workflow | +| [hera-coinflip](examples/hera-coinflip/) | Runs a basic coinflip example using the Hera Python framework. | +| [hera-notebook-forecast](examples/hera-notebook-forecast/) | Press-play data job from a Jupyter notebook: aggregate energy demand and forecast the next day. Shows a platform abstraction over Hera. | +| [pytorch-fashion-mnist](examples/pytorch-fashion-mnist/) | A PyTorch training pipeline: prep, parallel train, and register the best model. Authored with the Argo Workflows Python SDK and as a native Workflow. | diff --git a/examples/ci-pr-check/README.md b/examples/ci-pr-check/README.md new file mode 100644 index 0000000..e50f525 --- /dev/null +++ b/examples/ci-pr-check/README.md @@ -0,0 +1,47 @@ +[![Pipekit Logo](../../assets/images/pipekit-logo.png)](https://pipekit.io) + +# CI pull-request check + +A pull-request check pipeline that runs on Argo Workflows. It checks out a repository, runs lint and unit tests in parallel, and reports a combined pass or fail status. + +```text +checkout -> lint + unit-test (in parallel) +onExit: report +``` + +- `checkout` clones the repository and passes the working tree to the next steps as an artifact. +- `lint` and `unit-test` run in parallel, each against the checked-out code. +- `report` is an exit handler, so it runs whether the checks pass or fail. It prints the overall result and, if given a GitHub token, posts a commit status. + +This is a simplified, self-contained version of the CI that builds this examples repository (see the `ci/` directory). It runs inside the free trial cluster with no secrets: by default it checks the public `pipekit/examples` repository. The full build, test, and deploy pipeline, including multi-cluster promotion, is described in the Pipekit docs under [CI/CD](https://docs.pipekit.io/use-cases/ci-cd). + +## Log into Pipekit via the CLI + +With the [CLI installed](https://docs.pipekit.io/reference/cli), log in once: + +```bash +pipekit login +``` + +## Run the Workflow + +```bash +pipekit submit -w --cluster-name=free-trial-cluster --pipe-name=ci-pr-check-example examples/ci-pr-check/workflow.yaml +``` + +## Check your own repository + +Override the parameters to point at your project and its commands: + +```bash +pipekit submit -w --cluster-name=free-trial-cluster --pipe-name=ci-pr-check-example \ + -p repo_url=https://github.com/your-org/your-repo.git \ + -p ref=main \ + -p lint_cmd="pip install --quiet ruff && ruff check ." \ + -p test_cmd="pytest" \ + examples/ci-pr-check/workflow.yaml +``` + +## Post a real commit status + +The `report` step posts a GitHub commit status when three values are present in the pod: `GITHUB_TOKEN`, `GITHUB_REPOSITORY` (as `owner/repo`), and `GIT_SHA`. Add the token as a [secret on the pipe](https://docs.pipekit.io/using-pipekit/pipes/edit/secrets) and pass the other two as parameters or environment values. Without them the step prints the status it would post, so the example still runs with no setup. diff --git a/examples/ci-pr-check/workflow.yaml b/examples/ci-pr-check/workflow.yaml new file mode 100644 index 0000000..6ad3955 --- /dev/null +++ b/examples/ci-pr-check/workflow.yaml @@ -0,0 +1,131 @@ +--- +# A pull-request check pipeline as a single Argo Workflow: +# +# checkout -> lint + unit-test (in parallel) +# onExit: report (posts a combined pass/fail status) +# +# checkout clones a repository and passes the code to the next steps as an artifact. +# lint and unit-test run in parallel. An exit handler reports the overall result, so the +# status reflects the whole run whether it passed or failed. +# +# This is a simplified, self-contained version of the CI that builds this examples +# repository (see ci/ci.yaml). It runs inside the Pipekit free trial cluster with no +# secrets: by default it checks the public pipekit/examples repository. The full +# build, test, and deploy pipeline is described in the Pipekit docs under CI/CD. +apiVersion: argoproj.io/v1alpha1 +kind: Workflow +metadata: + generateName: ci-pr-check- + namespace: argo +spec: + entrypoint: main + onExit: report + serviceAccountName: argo-workflow + arguments: + parameters: + - name: repo_url + value: "https://github.com/pipekit/examples.git" + - name: ref + value: "main" + # The command each check runs, from the repository root. Override these to point + # at your own project's linter and tests. + - name: lint_cmd + value: "pip install --quiet ruff && ruff check ." + - name: test_cmd + value: "python examples/hera-notebook-forecast/test_demand_forecast.py" + templates: + - name: main + dag: + tasks: + - name: checkout + template: checkout + - name: lint + template: check + depends: checkout + arguments: + parameters: + - name: title + value: "lint" + - name: command + value: "{{workflow.parameters.lint_cmd}}" + - name: unit-test + template: check + depends: checkout + arguments: + parameters: + - name: title + value: "unit-test" + - name: command + value: "{{workflow.parameters.test_cmd}}" + + # checkout: shallow-clone the repository and pass the working tree on as an artifact. + - name: checkout + container: + image: python:3.11-slim + command: [bash, -c] + args: + - | + set -e + apt-get update -qq && apt-get install -y -qq git >/dev/null + git clone --depth 1 --branch {{workflow.parameters.ref}} {{workflow.parameters.repo_url}} /tmp/src + echo "checked out {{workflow.parameters.repo_url}} at {{workflow.parameters.ref}}" + resources: + requests: { cpu: "250m", memory: 256Mi, ephemeral-storage: 256Mi } + limits: { cpu: "500m", memory: 512Mi, ephemeral-storage: 1Gi } + outputs: + artifacts: + - name: source + path: /tmp/src + s3: + key: "{{workflow.uid}}/source" + + # check: run one command against the checked-out code. Reused by lint and unit-test. + - name: check + inputs: + parameters: + - name: title + - name: command + artifacts: + - name: source + path: /tmp/src + s3: + key: "{{workflow.uid}}/source" + container: + image: python:3.11-slim + workingDir: /tmp/src + command: [bash, -c] + args: + - | + set -e + echo "=== {{inputs.parameters.title}} ===" + {{inputs.parameters.command}} + resources: + requests: { cpu: "250m", memory: 256Mi, ephemeral-storage: 256Mi } + limits: { cpu: "1", memory: 768Mi, ephemeral-storage: 1Gi } + + # report: exit handler. Prints the final status of the whole run. If a GitHub token + # is provided as the GITHUB_TOKEN secret on the pipe, it also posts a commit status; + # without one it prints the status it would post, so the example runs with no setup. + - name: report + container: + image: python:3.11-slim + command: [bash, -c] + args: + - | + set -e + STATUS="{{workflow.status}}" + if [ "$STATUS" = "Succeeded" ]; then STATE="success"; else STATE="failure"; fi + echo "pull-request check finished: $STATUS -> commit status '$STATE'" + if [ -n "${GITHUB_TOKEN:-}" ] && [ -n "${GITHUB_REPOSITORY:-}" ] && [ -n "${GIT_SHA:-}" ]; then + apt-get update -qq && apt-get install -y -qq curl >/dev/null + curl -sS -X POST \ + -H "Authorization: Bearer ${GITHUB_TOKEN}" \ + -H "Accept: application/vnd.github+json" \ + "https://api.github.com/repos/${GITHUB_REPOSITORY}/statuses/${GIT_SHA}" \ + -d "{\"state\":\"${STATE}\",\"context\":\"pipekit/ci\"}" + else + echo "GITHUB_TOKEN not set; skipping the GitHub commit status post" + fi + resources: + requests: { cpu: "100m", memory: 128Mi, ephemeral-storage: 128Mi } + limits: { cpu: "250m", memory: 256Mi, ephemeral-storage: 256Mi } diff --git a/examples/pytorch-fashion-mnist/README.md b/examples/pytorch-fashion-mnist/README.md new file mode 100644 index 0000000..63ad24f --- /dev/null +++ b/examples/pytorch-fashion-mnist/README.md @@ -0,0 +1,88 @@ +[![Pipekit Logo](../../assets/images/pipekit-logo.png)](https://pipekit.io) + +# PyTorch Fashion-MNIST training pipeline + +A small machine-learning pipeline that trains an image classifier on Argo Workflows and picks the best model. It runs inside the Pipekit free trial cluster. + +The pipeline has four stages as one DAG: + +```text +prep -> train (x3, in parallel) -> register +``` + +- `prep` downloads the Fashion-MNIST dataset once and saves it as an artifact. +- `train` fans out over three hyperparameter configs. Each pod trains a small PyTorch model on the CPU and writes its model and metrics as a candidate artifact. +- `register` reads every candidate, picks the most accurate, and writes the winning model and a model card to the artifact store. + +This is the standard "data prep, train, evaluate, register" shape, scaled down to fit the free trial cluster's resource limits. The training logic is generic, so you can swap in your own model and dataset. + +## Files + +- `workflow.py`: the pipeline authored with the Argo Workflows Python SDK. Data and ML teams usually author in Python, so this is the primary path. +- `workflow.yaml`: the same DAG as a plain manifest, for the GitOps path or the CLI. +- `model.py`: the model and training logic as plain PyTorch functions, the tested source of truth. The cluster steps mirror this logic. +- `test_model.py`: a local test of that logic. Runs in milliseconds, no cluster. + +## PyTorch on the free trial cluster + +PyTorch is installed at run time from `download.pytorch.org` (the CPU-only wheel), so the only Docker Hub image is `python:3.11-slim`. The free trial cluster shares one egress IP, and pulling large images from Docker Hub there can hit the anonymous rate limit. Fetching the wheel from pytorch.org avoids that. + +The `torch_index` workflow parameter sets the pip index. Point it at an internal mirror to run offline or on a locked-down cluster. + +## Log into Pipekit via the CLI + +With the [CLI installed](https://docs.pipekit.io/reference/cli), log in once: + +```bash +pipekit login +``` + +## Run it with the Python SDK + +From the repo root: + +```bash +pip install 'hera' 'pipekit-sdk>=7.1.0' +python examples/pytorch-fashion-mnist/workflow.py +``` + +The SDK package installs as `hera` today. The script submits the pipeline, prints a link to watch it live in the Pipekit UI, and exits. Set `PIPEKIT_CLUSTER` to target a cluster other than `free-trial-cluster`. Print the manifest without submitting with `RENDER_ONLY=1 python workflow.py`. + +## Run the native Workflow + +```bash +pipekit submit -w --cluster-name=free-trial-cluster --pipe-name=fashion-mnist-example examples/pytorch-fashion-mnist/workflow.yaml +``` + +## What you get + +The run finishes in a few minutes. Each `train` pod logs its test accuracy, for example: + +```text +train config 0: lr=0.05 hidden=64 accuracy=0.6280 +train config 1: lr=0.1 hidden=128 accuracy=0.6955 +train config 2: lr=0.2 hidden=256 accuracy=0.7290 +``` + +The `register` step selects the most accurate model and exposes it as an output parameter, `selected`, so the choice is visible in the run graph even after the pods are cleaned up: + +```json +{ + "framework": "pytorch", + "dataset": "fashion-mnist", + "selected": { "idx": "2", "lr": 0.2, "hidden": 256, "accuracy": 0.723 } +} +``` + +## Test the logic locally + +```bash +pip install torch +python examples/pytorch-fashion-mnist/test_model.py +``` + +The model and training functions are plain PyTorch, so you can test them in milliseconds before the job ever touches the cluster. + +## Scaling up + +The free trial version keeps the model and dataset small so it fits the trial cluster's limits. For the production shape of this pipeline, including a model registry and GPU training, see [ML Pipelines](https://docs.pipekit.io/use-cases/ml-pipelines) in the Pipekit docs. diff --git a/examples/pytorch-fashion-mnist/model.py b/examples/pytorch-fashion-mnist/model.py new file mode 100644 index 0000000..abbf723 --- /dev/null +++ b/examples/pytorch-fashion-mnist/model.py @@ -0,0 +1,57 @@ +"""Pure PyTorch logic for the Fashion-MNIST classifier. No Argo, no cluster. + +These functions are the tested source of truth (see ``test_model.py``). The workflow +steps in ``workflow.yaml`` and ``workflow.py`` mirror this logic so the same maths runs +on the cluster. This is the same split the ``hera-notebook-forecast`` example uses: +plain, tested Python here, and the cluster step inlines an equivalent. +""" + +import torch +from torch import nn + +CLASSES = 10 +PIXELS = 28 * 28 + + +def build_model(hidden): + """A small multi-layer perceptron: 784 inputs, one hidden layer, 10 outputs. + + Big enough to learn Fashion-MNIST on a CPU, small enough to train in seconds. + """ + return nn.Sequential( + nn.Flatten(), + nn.Linear(PIXELS, hidden), + nn.ReLU(), + nn.Linear(hidden, CLASSES), + ) + + +def train(model, images, labels, lr, epochs=1, batch_size=128, seed=0): + """Train the model in place with plain SGD and return it. + + ``images`` is a float tensor shaped [N, 1, 28, 28] scaled to 0-1. ``labels`` is a + long tensor shaped [N]. The seed fixes the batch order so a given config is + reproducible. + """ + torch.manual_seed(seed) + optimizer = torch.optim.SGD(model.parameters(), lr=lr) + loss_fn = nn.CrossEntropyLoss() + sample_count = images.shape[0] + model.train() + for _ in range(epochs): + order = torch.randperm(sample_count) + for start in range(0, sample_count, batch_size): + batch = order[start : start + batch_size] + optimizer.zero_grad() + loss = loss_fn(model(images[batch]), labels[batch]) + loss.backward() + optimizer.step() + return model + + +def accuracy(model, images, labels): + """Return the fraction of correct predictions on the given set.""" + model.eval() + with torch.no_grad(): + predictions = model(images).argmax(dim=1) + return float((predictions == labels).float().mean()) diff --git a/examples/pytorch-fashion-mnist/test_model.py b/examples/pytorch-fashion-mnist/test_model.py new file mode 100644 index 0000000..88de888 --- /dev/null +++ b/examples/pytorch-fashion-mnist/test_model.py @@ -0,0 +1,38 @@ +"""Local test of the model logic. Random tensors, no download, runs in milliseconds. + +Run it before the job ever touches the cluster: + + python test_model.py +""" + +import torch + +from model import CLASSES, accuracy, build_model, train + + +def test_build_model_shapes(): + model = build_model(hidden=32) + output = model(torch.zeros(4, 1, 28, 28)) + assert output.shape == (4, CLASSES) + + +def test_model_learns_a_separable_pattern(): + # Two classes that are trivially separable: all-black images are class 0, + # all-white images are class 1. A working training loop must reach high accuracy. + half = 128 + images = torch.cat([torch.zeros(half, 1, 28, 28), torch.ones(half, 1, 28, 28)]) + labels = torch.cat([torch.zeros(half), torch.ones(half)]).long() + + model = build_model(hidden=16) + before = accuracy(model, images, labels) + train(model, images, labels, lr=0.1, epochs=5, seed=0) + after = accuracy(model, images, labels) + + assert after > before + assert after > 0.9 + + +if __name__ == "__main__": + test_build_model_shapes() + test_model_learns_a_separable_pattern() + print("ok") diff --git a/examples/pytorch-fashion-mnist/workflow.py b/examples/pytorch-fashion-mnist/workflow.py new file mode 100644 index 0000000..cb1ed18 --- /dev/null +++ b/examples/pytorch-fashion-mnist/workflow.py @@ -0,0 +1,237 @@ +"""The same Fashion-MNIST pipeline authored with the Argo Workflows Python SDK. + +This builds the identical prep -> train (fan-out) -> register DAG as ``workflow.yaml`` +and submits it to Pipekit. Data and ML teams usually author in Python, so this is the +canonical way to write the pipeline; ``workflow.yaml`` is the same graph as a plain +manifest for the GitOps path. + +Run it from the repo dev shell: + + pip install 'hera' 'pipekit-sdk>=7.1.0' && pipekit login + python examples/pytorch-fashion-mnist/workflow.py + +The SDK package installs as ``hera`` today. Set PIPEKIT_CLUSTER to target a cluster +other than the free trial one. Print the manifest instead of submitting with +``RENDER_ONLY=1 python workflow.py``. +""" + +import os + +from hera.workflows import DAG, Container, Parameter, S3Artifact, Workflow +from hera.workflows.models import ArchiveStrategy, NoneStrategy, ResourceRequirements, ValueFrom + +TORCH_INDEX = "https://download.pytorch.org/whl/cpu" + +# One pod per config. Each trains a small MLP and scores it on the test set. +CONFIGS = [ + {"idx": "0", "lr": "0.05", "hidden": "64"}, + {"idx": "1", "lr": "0.10", "hidden": "128"}, + {"idx": "2", "lr": "0.20", "hidden": "256"}, +] + + +def _resources(memory, ephemeral, cpu_limit="1"): + return ResourceRequirements( + requests={"cpu": "500m", "memory": "512Mi", "ephemeral-storage": "512Mi"}, + limits={"cpu": cpu_limit, "memory": memory, "ephemeral-storage": ephemeral}, + ) + + +def _none_archive(): + return ArchiveStrategy(none=NoneStrategy()) + + +PREP_SCRIPT = r""" +set -e +pip install --quiet --no-cache-dir torch torchvision --index-url {{workflow.parameters.torch_index}} +python - <<'PY' +import os +import torch +from torchvision import datasets, transforms + +to_tensor = transforms.ToTensor() +train_ds = datasets.FashionMNIST("/tmp/d", train=True, download=True, transform=to_tensor) +test_ds = datasets.FashionMNIST("/tmp/d", train=False, download=True, transform=to_tensor) + +def take(dataset, count): + images = torch.stack([dataset[i][0] for i in range(count)]) + labels = torch.tensor([dataset[i][1] for i in range(count)]) + return images, labels + +train_x, train_y = take(train_ds, 8000) +test_x, test_y = take(test_ds, 2000) +os.makedirs("/tmp/out", exist_ok=True) +torch.save({"train_x": train_x, "train_y": train_y, "test_x": test_x, "test_y": test_y}, "/tmp/out/data.pt") +print("prep: train", tuple(train_x.shape), "test", tuple(test_x.shape)) +PY +""" + +TRAIN_SCRIPT = r""" +set -e +pip install --quiet --no-cache-dir torch --index-url {{workflow.parameters.torch_index}} +python - <<'PY' +import json +import os +import torch +from torch import nn + +idx = "{{inputs.parameters.idx}}" +lr = float("{{inputs.parameters.lr}}") +hidden = int("{{inputs.parameters.hidden}}") +pixels, classes = 28 * 28, 10 + +data = torch.load("/tmp/in/data.pt") +model = nn.Sequential(nn.Flatten(), nn.Linear(pixels, hidden), nn.ReLU(), nn.Linear(hidden, classes)) +torch.manual_seed(0) +optimizer = torch.optim.SGD(model.parameters(), lr=lr) +loss_fn = nn.CrossEntropyLoss() +train_x, train_y = data["train_x"], data["train_y"] +sample_count = train_x.shape[0] +model.train() +order = torch.randperm(sample_count) +for start in range(0, sample_count, 128): + batch = order[start : start + 128] + optimizer.zero_grad() + loss_fn(model(train_x[batch]), train_y[batch]).backward() + optimizer.step() + +model.eval() +with torch.no_grad(): + predictions = model(data["test_x"]).argmax(dim=1) + accuracy = float((predictions == data["test_y"]).float().mean()) + +os.makedirs("/tmp/out", exist_ok=True) +torch.save(model.state_dict(), "/tmp/out/model.pt") +with open("/tmp/out/metrics.json", "w") as handle: + json.dump({"idx": idx, "lr": lr, "hidden": hidden, "accuracy": round(accuracy, 4)}, handle) +print(f"train config {idx}: lr={lr} hidden={hidden} accuracy={accuracy:.4f}") +PY +""" + +REGISTER_SCRIPT = r""" +set -e +python - <<'PY' +import glob +import json +import os +import shutil + +best = None +for metrics_path in sorted(glob.glob("/tmp/candidates/*/metrics.json")): + with open(metrics_path) as handle: + metrics = json.load(handle) + print("candidate:", metrics) + if best is None or metrics["accuracy"] > best["accuracy"]: + best = dict(metrics, _dir=os.path.dirname(metrics_path)) + +os.makedirs("/tmp/out", exist_ok=True) +shutil.copy(os.path.join(best["_dir"], "model.pt"), "/tmp/out/model.pt") +card = { + "framework": "pytorch", + "dataset": "fashion-mnist", + "selected": {key: best[key] for key in ("idx", "lr", "hidden", "accuracy")}, +} +with open("/tmp/out/model-card.json", "w") as handle: + json.dump(card, handle, indent=2) +print("registered best model:", card["selected"]) +PY +""" + + +def build() -> Workflow: + with Workflow( + generate_name="fashion-mnist-", + entrypoint="main", + namespace="argo", + service_account_name="argo-workflow", + arguments={"torch_index": TORCH_INDEX}, + ) as workflow: + prep = Container( + name="prep", + image="python:3.11-slim", + command=["bash", "-c"], + args=[PREP_SCRIPT], + resources=_resources(memory="1500Mi", ephemeral="3Gi"), + outputs=[ + S3Artifact( + name="dataset", + path="/tmp/out", + key="{{workflow.uid}}/dataset", + archive=_none_archive(), + ) + ], + ) + train = Container( + name="train", + image="python:3.11-slim", + command=["bash", "-c"], + args=[TRAIN_SCRIPT], + resources=_resources(memory="2Gi", ephemeral="3Gi"), + inputs=[ + Parameter(name="idx"), + Parameter(name="lr"), + Parameter(name="hidden"), + S3Artifact(name="dataset", path="/tmp/in", key="{{workflow.uid}}/dataset"), + ], + outputs=[ + S3Artifact( + name="candidate", + path="/tmp/out", + key="{{workflow.uid}}/candidates/{{inputs.parameters.idx}}", + archive=_none_archive(), + ) + ], + ) + register = Container( + name="register", + image="python:3.11-slim", + command=["bash", "-c"], + args=[REGISTER_SCRIPT], + resources=ResourceRequirements( + requests={"cpu": "250m", "memory": "256Mi", "ephemeral-storage": "256Mi"}, + limits={"cpu": "500m", "memory": "512Mi", "ephemeral-storage": "512Mi"}, + ), + inputs=[ + S3Artifact( + name="candidates", path="/tmp/candidates", key="{{workflow.uid}}/candidates" + ) + ], + outputs=[ + Parameter(name="selected", value_from=ValueFrom(path="/tmp/out/model-card.json")), + S3Artifact( + name="registered", + path="/tmp/out", + key="{{workflow.uid}}/registered", + archive=_none_archive(), + ), + ], + ) + + with DAG(name="main"): + prep_task = prep() + train_task = train( + arguments={ + "idx": "{{item.idx}}", + "lr": "{{item.lr}}", + "hidden": "{{item.hidden}}", + }, + with_items=CONFIGS, + ) + register_task = register() + prep_task >> train_task >> register_task + + return workflow + + +if __name__ == "__main__": + workflow = build() + if os.environ.get("RENDER_ONLY"): + print(workflow.to_yaml()) + else: + from pipekit_sdk.service import PipekitService + + cluster = os.environ.get("PIPEKIT_CLUSTER", "free-trial-cluster") + pipekit = PipekitService() + pipe_run = pipekit.submit(workflow, cluster) + print(f"submitted run {pipe_run.uuid}") + print(f"watch live: https://pipekit.io/pipes/{pipe_run.pipe_uuid}/runs/{pipe_run.uuid}") diff --git a/examples/pytorch-fashion-mnist/workflow.yaml b/examples/pytorch-fashion-mnist/workflow.yaml new file mode 100644 index 0000000..4639e46 --- /dev/null +++ b/examples/pytorch-fashion-mnist/workflow.yaml @@ -0,0 +1,231 @@ +--- +# A four-stage machine-learning pipeline as a single Argo Workflow: +# +# prep -> train (x3, in parallel) -> register +# +# prep downloads Fashion-MNIST once and saves it as an artifact. train fans out over +# three hyperparameter configs; each pod trains a small PyTorch model on the CPU and +# writes its model plus metrics as an artifact. register reads every candidate, picks +# the most accurate, and writes the winning model and a model card to the artifact +# store. This is the "data prep, train, evaluate, register" shape, scaled to run inside +# the Pipekit free trial cluster. +# +# PyTorch is installed at run time from download.pytorch.org (the CPU wheel), so the +# only Docker Hub image is python:3.11-slim. The free trial cluster shares one egress +# IP, and pulling large images from Docker Hub there can hit the anonymous rate limit; +# fetching the wheel from pytorch.org avoids that. +apiVersion: argoproj.io/v1alpha1 +kind: Workflow +metadata: + generateName: fashion-mnist- + namespace: argo +spec: + entrypoint: main + serviceAccountName: argo-workflow + arguments: + parameters: + # The pip index for the CPU-only PyTorch build. Kept as a parameter so you can + # point it at an internal mirror without editing every step. + - name: torch_index + value: "https://download.pytorch.org/whl/cpu" + templates: + - name: main + dag: + tasks: + - name: prep + template: prep + - name: train + template: train + depends: prep + arguments: + parameters: + - name: idx + value: "{{item.idx}}" + - name: lr + value: "{{item.lr}}" + - name: hidden + value: "{{item.hidden}}" + withItems: + - { idx: "0", lr: "0.05", hidden: "64" } + - { idx: "1", lr: "0.10", hidden: "128" } + - { idx: "2", lr: "0.20", hidden: "256" } + - name: register + template: register + depends: train + + # prep: download Fashion-MNIST, take a subset to keep the job light, and save the + # tensors as one artifact the train pods read. + - name: prep + container: + image: python:3.11-slim + command: [bash, -c] + args: + - | + set -e + pip install --quiet --no-cache-dir torch torchvision --index-url {{workflow.parameters.torch_index}} + python - <<'PY' + import os + import torch + from torchvision import datasets, transforms + + to_tensor = transforms.ToTensor() + train_ds = datasets.FashionMNIST("/tmp/d", train=True, download=True, transform=to_tensor) + test_ds = datasets.FashionMNIST("/tmp/d", train=False, download=True, transform=to_tensor) + + def take(dataset, count): + images = torch.stack([dataset[i][0] for i in range(count)]) + labels = torch.tensor([dataset[i][1] for i in range(count)]) + return images, labels + + train_x, train_y = take(train_ds, 8000) + test_x, test_y = take(test_ds, 2000) + os.makedirs("/tmp/out", exist_ok=True) + torch.save( + {"train_x": train_x, "train_y": train_y, "test_x": test_x, "test_y": test_y}, + "/tmp/out/data.pt", + ) + print("prep: train", tuple(train_x.shape), "test", tuple(test_x.shape)) + PY + resources: + requests: { cpu: "500m", memory: 512Mi, ephemeral-storage: 512Mi } + limits: { cpu: "1", memory: 1500Mi, ephemeral-storage: 3Gi } + outputs: + artifacts: + - name: dataset + path: /tmp/out + archive: + none: {} + s3: + key: "{{workflow.uid}}/dataset" + + # train: one pod per hyperparameter config. Trains a small MLP, scores it on the + # test set, and writes the model and its metrics as a candidate artifact. + - name: train + inputs: + parameters: + - name: idx + - name: lr + - name: hidden + artifacts: + - name: dataset + path: /tmp/in + s3: + key: "{{workflow.uid}}/dataset" + container: + image: python:3.11-slim + command: [bash, -c] + args: + - | + set -e + pip install --quiet --no-cache-dir torch --index-url {{workflow.parameters.torch_index}} + python - <<'PY' + import json + import os + import torch + from torch import nn + + idx = "{{inputs.parameters.idx}}" + lr = float("{{inputs.parameters.lr}}") + hidden = int("{{inputs.parameters.hidden}}") + pixels, classes = 28 * 28, 10 + + data = torch.load("/tmp/in/data.pt") + model = nn.Sequential( + nn.Flatten(), + nn.Linear(pixels, hidden), + nn.ReLU(), + nn.Linear(hidden, classes), + ) + torch.manual_seed(0) + optimizer = torch.optim.SGD(model.parameters(), lr=lr) + loss_fn = nn.CrossEntropyLoss() + train_x, train_y = data["train_x"], data["train_y"] + sample_count = train_x.shape[0] + model.train() + order = torch.randperm(sample_count) + for start in range(0, sample_count, 128): + batch = order[start : start + 128] + optimizer.zero_grad() + loss_fn(model(train_x[batch]), train_y[batch]).backward() + optimizer.step() + + model.eval() + with torch.no_grad(): + predictions = model(data["test_x"]).argmax(dim=1) + accuracy = float((predictions == data["test_y"]).float().mean()) + + os.makedirs("/tmp/out", exist_ok=True) + torch.save(model.state_dict(), "/tmp/out/model.pt") + with open("/tmp/out/metrics.json", "w") as handle: + json.dump({"idx": idx, "lr": lr, "hidden": hidden, "accuracy": round(accuracy, 4)}, handle) + print(f"train config {idx}: lr={lr} hidden={hidden} accuracy={accuracy:.4f}") + PY + resources: + requests: { cpu: "500m", memory: 512Mi, ephemeral-storage: 512Mi } + limits: { cpu: "1", memory: 2Gi, ephemeral-storage: 3Gi } + outputs: + artifacts: + - name: candidate + path: /tmp/out + archive: + none: {} + s3: + key: "{{workflow.uid}}/candidates/{{inputs.parameters.idx}}" + + # register: read every candidate's metrics, pick the most accurate, and write the + # winning model and a model card. No PyTorch needed here, so this step is quick. + - name: register + inputs: + artifacts: + - name: candidates + path: /tmp/candidates + s3: + key: "{{workflow.uid}}/candidates" + container: + image: python:3.11-slim + command: [bash, -c] + args: + - | + set -e + python - <<'PY' + import glob + import json + import os + import shutil + + best = None + for metrics_path in sorted(glob.glob("/tmp/candidates/*/metrics.json")): + with open(metrics_path) as handle: + metrics = json.load(handle) + print("candidate:", metrics) + if best is None or metrics["accuracy"] > best["accuracy"]: + best = dict(metrics, _dir=os.path.dirname(metrics_path)) + + os.makedirs("/tmp/out", exist_ok=True) + shutil.copy(os.path.join(best["_dir"], "model.pt"), "/tmp/out/model.pt") + card = { + "framework": "pytorch", + "dataset": "fashion-mnist", + "selected": {key: best[key] for key in ("idx", "lr", "hidden", "accuracy")}, + } + with open("/tmp/out/model-card.json", "w") as handle: + json.dump(card, handle, indent=2) + print("registered best model:", card["selected"]) + PY + resources: + requests: { cpu: "250m", memory: 256Mi, ephemeral-storage: 256Mi } + limits: { cpu: "500m", memory: 512Mi, ephemeral-storage: 512Mi } + outputs: + parameters: + # Exposed as a node output so the winning config is visible in the run graph + # even after the pod is garbage-collected on success. + - name: selected + valueFrom: + path: /tmp/out/model-card.json + artifacts: + - name: registered + path: /tmp/out + archive: + none: {} + s3: + key: "{{workflow.uid}}/registered"