Part III — Searching for the best architecture

Part III — Searching for the best architecture#

Tired of the julia-fractal example? Good news, we are going to use a new one: training a convolutional neural network (CNN) on an image dataset (CIFAR-10). Such a CNN can be configured with various numbers of layers or kernel sizes, and can be trained with various hyperparameters. But how do we know which architecture is the best? This matters: if the goal of an experiment is to compare a new approach to image classification against a state-of-the-art CNN, we better be sure to configure the latter as best as we can.

This part of the tutorial will introduce a number of features, including uv-based setups (as opposed to the previously used spack-based ones), domain descriptions, parameter sweeps, and parameter-space exploration using Bayesian optimization.

We build exp.sh up as we go: Steps 1–2 start with a train job that searches three parameters, and Step 8 grows it to eight. So the command output shown along the way reflects the file at that step, not the finished script — the complete version appears at the end, in The complete experiment (Part III).

Note

This part runs on a laptop. Unlike Part I it needs no compiler or MPI, but it does need the uv plugin’s prerequisites (a network connection and the ability to run uv, which bootstrap provisions for you) and enough time and memory to train a small network on CPU. A GPU is not required. As in Part I, every step’s code is shown in full and was validated end to end during development; the console output is illustrative.

The program#

Like julia-fractal, the trainer lives in its own public repository, image-classifier-example (the cifar-cnn program). We do not reproduce its source here; we treat it as a black box with a simple contract. It is a real PyTorch program that:

reads a CNN architecture from command-line arguments (number of convolution blocks, channels per block, kernel size, and — later — dropout, activation, learning rate, and more);
trains that architecture on CIFAR-10 for a few epochs;
prints its final validation accuracy on a stable, machine-readable line, RESULT val_accuracy=<percent>, which the job reads back and records.

It is installable as a Python package, so the setup can point uv straight at the repository and get a pinned, reproducible environment and the program itself in one step. That contract — these arguments in, an accuracy out — is all we need; the PyTorch internals stay out of view.

Step 1 — One training run#

The first step provisions the environment and runs a single architecture. Three registrations do it: a setup that installs the program, a resource that downloads the dataset, and a job that trains one network and records its accuracy.

Part I provisioned a C++ toolchain with the Spack plugin. An ML stack is a better fit for uv, so Part III imports that plugin instead and declares a uv-backed setup. The setup names one requirement — the cifar-cnn package, installed directly from its Git repository — and has no body of its own; uv building and activating the environment is all it does:

Shorthand

# A uv-backed setup installs the classifier. @with_uv_specs declares a uv Python
# environment whose single requirement is the cifar-cnn package, installed
# straight from its git repository; knit builds and activates it as the setup's
# first step, so the `cifar-cnn` console script is on PATH before any job that
# declares this setup runs. It captures the resolved pyproject.toml / uv.lock as
# provenance. The setup needs no body of its own --- building the environment is
# all it does --- so @empty stands in for the body function.
@setup "cnnenv" "A uv environment with the cifar-cnn classifier installed."
@with_uv_specs "cifar-cnn @ git+https://github.com/knit-sh/image-classifier-example.git"
@empty
@done

Long form

# A uv-backed setup installs the classifier. @with_uv_specs declares a uv Python
# environment whose single requirement is the cifar-cnn package, installed
# straight from its git repository; knit builds and activates it as the setup's
# first step, so the `cifar-cnn` console script is on PATH before any job that
# declares this setup runs. It captures the resolved pyproject.toml / uv.lock as
# provenance. The setup needs no body of its own --- building the environment is
# all it does --- so @empty stands in for the body function.
knit_register_setup "cnnenv" knit_empty "A uv environment with the cifar-cnn classifier installed."
@with_uv_specs "cifar-cnn @ git+https://github.com/knit-sh/image-classifier-example.git"
knit_done

Because a uv-backed setup is declared, bootstrap provisions a knit-private uv automatically. uv installs the program and pins the exact dependency set (pyproject.toml / uv.lock) as provenance, so a reproducer gets the same environment.

The CIFAR-10 dataset is an input, so it is a resource. @with_url makes cifar a downloadable resource with a default URL; knit fetch downloads and (with --uncompress) unpacks it once into resources/, recording the URL:

Shorthand

# A resource declares HOW to acquire an input. @with_url makes `cifar` a
# downloadable resource whose default URL is the CIFAR-10 python archive;
# `knit fetch --name cifar -- cifar --uncompress` downloads and unpacks it into
# resources/cifar/ once, then records the URL for provenance. The train job
# reads the staged directory back through knit_resource_path.
@resource "cifar" "The CIFAR-10 dataset, downloaded and unpacked."
@with_url "https://www.cs.toronto.edu/~kriz/cifar-10-python.tar.gz"
@done

Long form

# A resource declares HOW to acquire an input. @with_url makes `cifar` a
# downloadable resource whose default URL is the CIFAR-10 python archive;
# `knit fetch --name cifar -- cifar --uncompress` downloads and unpacks it into
# resources/cifar/ once, then records the URL for provenance. The train job
# reads the staged directory back through knit_resource_path.
knit_register_resource "cifar" "The CIFAR-10 dataset, downloaded and unpacked."
knit_with_url "https://www.cs.toronto.edu/~kriz/cifar-10-python.tar.gz"
knit_done

Note

Probably due to CIFAR-10’s popularity, download is usually slow from its source website.

Finally the train job. It depends on the setup (so cifar-cnn is on PATH) and the resource (so the dataset is staged), runs cifar-cnn train, parses the RESULT val_accuracy= line, and records the number as a declared output — the column every later step reads as the objective:

Shorthand

# The job that trains one network and records its validation accuracy. It
# declares the cnnenv setup (so cifar-cnn is on PATH) and the cifar resource (so
# the dataset is staged). Its body runs `cifar-cnn train`, reads the
# machine-readable "RESULT val_accuracy=..." line the program prints, and records
# it with knit_output --- that recorded column is the objective later steps rank
# and optimize.
#
# Each tunable parameter also carries a --domain: its search space, a range
# [lo..hi] or a set {a,b,c}. The job runs fine without mentioning domains (every
# parameter has a default); the domains come into play in Step 2, when we search.
# There is deliberately no @with_table here: @job already declares the job's own
# table, which is where the objective is read; a bare @with_table would split
# recording from that table.
@job "train" "Train the CNN and record its validation accuracy."
@with_setup    "cnnenv"
@with_resource "data:cifar" "Name of the fetched CIFAR instance."
@with_optional "dataset:string"        "cifar10" "Dataset loader (not searched)."
@with_optional "num-blocks:integer"    "3"       "Convolution blocks."          --domain '[1..5]'
@with_optional "base-channels:integer" "32"      "Channels in the first block." --domain '{16,32,64}'
@with_optional "kernel-size:integer"   "3"       "Convolution kernel size."     --domain '{3,5}'
@with_optional "epochs:integer"        "3"       "Training epochs (not searched)."
@with_output   "val_accuracy:real"     "0"       "Final-epoch validation accuracy (%)."
_train_job() {
    local datadir dataset out acc
    datadir="$(knit_resource_path "$(knit_get_parameter data "$@")")"
    dataset="$(knit_get_parameter dataset "$@")"

    out="$(cifar-cnn train \
        --model         "${KNIT_JOB_PREFIX%/}/model.pt" \
        --dataset       "${dataset}" \
        --data          "${datadir}" \
        --num-blocks    "$(knit_get_parameter num-blocks "$@")" \
        --base-channels "$(knit_get_parameter base-channels "$@")" \
        --kernel-size   "$(knit_get_parameter kernel-size "$@")" \
        --epochs        "$(knit_get_parameter epochs "$@")" \
        2>&1 | tee /dev/stderr)"

    acc="$(grep -oE 'RESULT val_accuracy=[0-9.]+' <<<"${out}" | tail -1 | cut -d= -f2)"
    knit_output "val_accuracy" "${acc:-0}"
}
@done

Long form

# The job that trains one network and records its validation accuracy. It
# declares the cnnenv setup (so cifar-cnn is on PATH) and the cifar resource (so
# the dataset is staged). Its body runs `cifar-cnn train`, reads the
# machine-readable "RESULT val_accuracy=..." line the program prints, and records
# it with knit_output --- that recorded column is the objective later steps rank
# and optimize.
#
# Each tunable parameter also carries a --domain: its search space, a range
# [lo..hi] or a set {a,b,c}. The job runs fine without mentioning domains (every
# parameter has a default); the domains come into play in Step 2, when we search.
# There is deliberately no @with_table here: @job already declares the job's own
# table, which is where the objective is read; a bare @with_table would split
# recording from that table.
knit_register_job "train" _train_job "Train the CNN and record its validation accuracy."
knit_with_setup    "cnnenv"
knit_with_resource "data:cifar" "Name of the fetched CIFAR instance."
knit_with_optional "dataset:string"        "cifar10" "Dataset loader (not searched)."
knit_with_optional "num-blocks:integer"    "3"       "Convolution blocks."          --domain '[1..5]'
knit_with_optional "base-channels:integer" "32"      "Channels in the first block." --domain '{16,32,64}'
knit_with_optional "kernel-size:integer"   "3"       "Convolution kernel size."     --domain '{3,5}'
knit_with_optional "epochs:integer"        "3"       "Training epochs (not searched)."
knit_with_output   "val_accuracy:real"     "0"       "Final-epoch validation accuracy (%)."
_train_job() {
    local datadir dataset out acc
    datadir="$(knit_resource_path "$(knit_get_parameter data "$@")")"
    dataset="$(knit_get_parameter dataset "$@")"

    out="$(cifar-cnn train \
        --model         "${KNIT_JOB_PREFIX%/}/model.pt" \
        --dataset       "${dataset}" \
        --data          "${datadir}" \
        --num-blocks    "$(knit_get_parameter num-blocks "$@")" \
        --base-channels "$(knit_get_parameter base-channels "$@")" \
        --kernel-size   "$(knit_get_parameter kernel-size "$@")" \
        --epochs        "$(knit_get_parameter epochs "$@")" \
        2>&1 | tee /dev/stderr)"

    acc="$(grep -oE 'RESULT val_accuracy=[0-9.]+' <<<"${out}" | tail -1 | cut -d= -f2)"
    knit_output "val_accuracy" "${acc:-0}"
}
knit_done

Bootstrap, fetch the data, build the environment, and train one network:

$ ./exp.sh bootstrap
$ ./exp.sh fetch --name cifar -- cifar --uncompress
$ ./exp.sh setup --name cnnenv -- cnnenv
$ ./exp.sh submit --setup cnnenv --wait -- train --data cifar
...
RESULT val_accuracy=63.7

That run used every parameter’s default. The recorded val_accuracy is the single number the rest of this tutorial works to maximize.

Note

You may have noticed each parameter already carries a --domain annotation. The job runs fine without it — every parameter has a default — so we glossed over it just now. It is the subject of the next step.

Step 2 — Declare the search space#

To search for the best architecture we first have to say what “all architectures” means: the set of values each parameter may take. That is the parameter’s domain, which you may already have spotted on the train job above. It is declared with --domain right where the parameter is declared — a range [lo..hi] for a count (num-blocks), or a set {a,b,c} for a discrete choice (base-channels, kernel-size). Quote the descriptor so the shell does not brace-expand {...}. Parameters that are not searched — dataset (fixed per run) and epochs (kept small so a run finishes fast) — simply carry no domain. A domain changes nothing about how the job runs on its own; it is a declaration the search tools read.

Those declarations are the single source of truth for the rest of the part, and you can read them back at any time with ./exp.sh describe --space:

$ ./exp.sh describe --only 'submit train' --space
submit train
------------
  [job, user]  Train the CNN and record its validation accuracy.

  Options
  -------
    --help                   Print this help message and exit.
    --data <value>           [required, resource: cifar] Name of the fetched CIFAR instance.
    --dataset <value>        [default: 'cifar10'] Dataset loader (not searched).
    --num-blocks <value>     [default: '3', range [1,5]] Convolution blocks.
    --base-channels <value>  [default: '32', values {16,32,64}] Channels in the first block.
    --kernel-size <value>    [default: '3', values {3,5}] Convolution kernel size.
    --epochs <value>         [default: '3'] Training epochs (not searched).

  Outputs
  -------
    val-accuracy  [real, default: '0'] Final-epoch validation accuracy (%).

With the space declared, the obvious next move is to try more than one point in it.

Step 3 — A sweep by hand#

One training run tells us nothing about which architecture is best; we need to compare several. The first thing anyone reaches for is a shell loop over the two discrete parameters:

$ for nb in 2 3 4; do
    for ks in 3 5; do
      ./exp.sh submit --setup cnnenv --wait -- train --data cifar \
          --num-blocks "$nb" --kernel-size "$ks"
    done
  done

This works — it trains all six architectures and records six val_accuracy rows — but it is unsatisfying in three ways. The grid is hard-coded in shell, not derived from the domains we just declared, so it can silently drift from them. The six runs share no common label, so nothing ties them together as one experiment. And it is not resumable: if the machine reboots after the fourth run, the loop starts over from the first. Each of these is a reason to let Knit build the grid for us.

Step 4 — `knit sweep`#

knit sweep replaces the loop. It runs the existing train job once per point of a cartesian product, where each swept parameter takes sweep-syntax in place of a value: an explicit set '{2,3,4}', or @domain to draw from the parameter’s declared --domain. The --group grid on the submit line tags every job the sweep produces with one campaign label, so you can later drain, inspect, or remove them as a unit. Here is the six-point grid, run two at a time:

$ ./exp.sh sweep --max-inflight 2 -- submit --group grid --setup cnnenv --wait \
      -- train --data cifar --num-blocks '{2,3,4}' --kernel-size '{3,5}'

Quote the sets so the shell does not brace-expand them. Preview the expansion without running anything by inserting --dry-run right after sweep; it prints the reconstructed command lines and exits. Notice that the entire submit line — --group, --setup, and all — is replicated verbatim for each point, with only the swept parameters changing:

$ ./exp.sh sweep --dry-run -- submit --group grid --setup cnnenv -- train \
      --data cifar --num-blocks '{2,3,4}' --kernel-size '{3,5}'
submit --group grid --setup cnnenv -- train --data cifar --num-blocks 2 --kernel-size 3
submit --group grid --setup cnnenv -- train --data cifar --num-blocks 2 --kernel-size 5
submit --group grid --setup cnnenv -- train --data cifar --num-blocks 3 --kernel-size 3
submit --group grid --setup cnnenv -- train --data cifar --num-blocks 3 --kernel-size 5
submit --group grid --setup cnnenv -- train --data cifar --num-blocks 4 --kernel-size 3
submit --group grid --setup cnnenv -- train --data cifar --num-blocks 4 --kernel-size 5

This is where the domains of Step 2 pay off. Give a parameter @domain instead of an explicit set and the sweep draws its values straight from the declared domain, so the grid can never drift from it. --num-blocks @domain expands to the whole [1..5] range:

$ ./exp.sh sweep --dry-run -- submit -- train \
      --data cifar --num-blocks @domain --kernel-size '{3,5}'
submit -- train --data cifar --num-blocks 1 --kernel-size 3
submit -- train --data cifar --num-blocks 1 --kernel-size 5
... six more ...
submit -- train --data cifar --num-blocks 5 --kernel-size 3
submit -- train --data cifar --num-blocks 5 --kernel-size 5

Beyond running the jobs, the sweep records itself: each child job carries a provenance edge back to the sweep, so the whole grid can be queried as one unit later (we use that in the inspection step). One command, one label, and the grid is derived from the declared space — two of the three complaints from Step 3 are gone. But this ephemeral form still runs the whole sweep inside one foreground process. Lose that process — a closed laptop, a dropped connection — and the sweep stops where it was. The third complaint, resumability, needs one more move.

Step 5 — Durable prepare and drain#

The fix is to separate planning the sweep from running it. Point the sweep at prepare instead of submit and it records one job per grid point in the prepared state, contacting no scheduler:

$ ./exp.sh sweep -- prepare --group grid --setup cnnenv -- train \
      --data cifar --num-blocks '{2,3,4}' --kernel-size '{3,5}'

The whole submission spec — the --group label, the setup, and on a cluster the nodes, walltime, and queue — is frozen at prepare time, which is why those options go on the prepare line. The six jobs now sit durably in the database as a to-do list. Release them with submit drain, which keeps up to --max-inflight jobs alive at once:

$ ./exp.sh submit drain --group grid --max-inflight 2

Resumability is now a property of the queue, not of any one process. If a drain is interrupted, the finished jobs stay recorded and the unreleased ones stay prepared; running submit drain --group grid again picks up exactly where it left off. This prepare-then-drain shape is the backbone of the rest of the part.

Step 6 — Run the sweep inside a pilot job#

Everything so far runs fine on a laptop. On a cluster, though, submit drain releases each train job to the batch scheduler, so every job waits in the queue on its own — six queue waits for a six-point grid, and far more once the search grows. We would rather wait in the queue once, get an allocation, and run the whole sweep inside it.

That is what a pilot job does. A job marked @with_pilot takes one allocation and starts the in-allocation knit-pilotd daemon over its nodes; from inside its body, every knit submit — including the ones submit drain issues — targets that daemon instead of the batch scheduler, so the drained jobs pack onto the pilot’s own nodes with no second queue wait. To fill those nodes rather than use them one at a time, the drain runs several jobs at once: the body sizes --max-inflight to the allocation’s node count (one single-node train per node), read with knit_job_hostnames from inside the pilot:

Shorthand

# @with_pilot marks this job a pilot. Inside its body every `knit submit` --- here
# the ones `submit drain` issues --- targets the in-allocation knit-pilotd daemon
# over the pilot's own nodes, not the machine scheduler, so the prepared sweep
# packs onto this one allocation instead of queueing job by job. The pilot trains
# nothing itself, so it declares no setup; each drained `train` job still carries
# the --setup frozen into it at prepare time.
#
# To actually pack the allocation, the drain runs several jobs at once:
# --max-inflight is sized to the pilot's node count (one single-node `train` per
# node), read with knit_job_hostnames inside the allocation. A plain
# `submit drain` with no --max-inflight releases jobs one at a time, leaving every
# node but one idle.
#
# The body is a pure drainer, written to be re-runnable. Resubmitting the same
# pilot after a walltime kill resumes the sweep in three moves:
#   1. job reconcile --- a hard kill leaves jobs stuck in `running`/`submitting`;
#      reconcile asks the scheduler their true fate and corrects them to `killed`.
#   2. job reset --- move those `killed` jobs back to `prepared` so they re-run.
#   3. submit drain --- drain the still-`prepared` remainder of the sweep, N jobs
#      at a time, where N is the allocation's node count.
# On the first run reconcile and reset find nothing and are no-ops; on a resume
# they recover the interrupted jobs. The sweep's product is prepared once; this
# pilot just chews through it across as many walltime windows as it takes.
@job "sweep-pilot" "Drain a prepared sweep inside one pilot allocation."
@without_setup
@with_pilot
@with_optional "group:string" "grid" "Sweep group to drain."
_sweep_pilot_job() {
    local group nodes
    group="$(knit_get_parameter group "$@")"
    # One single-node `train` job per node: size the drain's concurrency to the
    # pilot's allocation so the whole allocation works at once.
    nodes="$(knit_job_hostnames | wc -l)"
    knit job reconcile --group "${group}"
    knit job reset     --group "${group}"
    knit submit drain  --group "${group}" --max-inflight "${nodes}"
}
@done

Long form

# @with_pilot marks this job a pilot. Inside its body every `knit submit` --- here
# the ones `submit drain` issues --- targets the in-allocation knit-pilotd daemon
# over the pilot's own nodes, not the machine scheduler, so the prepared sweep
# packs onto this one allocation instead of queueing job by job. The pilot trains
# nothing itself, so it declares no setup; each drained `train` job still carries
# the --setup frozen into it at prepare time.
#
# To actually pack the allocation, the drain runs several jobs at once:
# --max-inflight is sized to the pilot's node count (one single-node `train` per
# node), read with knit_job_hostnames inside the allocation. A plain
# `submit drain` with no --max-inflight releases jobs one at a time, leaving every
# node but one idle.
#
# The body is a pure drainer, written to be re-runnable. Resubmitting the same
# pilot after a walltime kill resumes the sweep in three moves:
#   1. job reconcile --- a hard kill leaves jobs stuck in `running`/`submitting`;
#      reconcile asks the scheduler their true fate and corrects them to `killed`.
#   2. job reset --- move those `killed` jobs back to `prepared` so they re-run.
#   3. submit drain --- drain the still-`prepared` remainder of the sweep, N jobs
#      at a time, where N is the allocation's node count.
# On the first run reconcile and reset find nothing and are no-ops; on a resume
# they recover the interrupted jobs. The sweep's product is prepared once; this
# pilot just chews through it across as many walltime windows as it takes.
knit_register_job "sweep-pilot" _sweep_pilot_job "Drain a prepared sweep inside one pilot allocation."
knit_without_setup
knit_with_pilot
knit_with_optional "group:string" "grid" "Sweep group to drain."
_sweep_pilot_job() {
    local group nodes
    group="$(knit_get_parameter group "$@")"
    # One single-node `train` job per node: size the drain's concurrency to the
    # pilot's allocation so the whole allocation works at once.
    nodes="$(knit_job_hostnames | wc -l)"
    knit job reconcile --group "${group}"
    knit job reset     --group "${group}"
    knit submit drain  --group "${group}" --max-inflight "${nodes}"
}
knit_done

The pilot trains nothing itself, so it declares @without_setup; each train job it drains still carries the --setup cnnenv frozen into it at prepare time. Prepare the sweep as in Step 5, then submit the pilot over as many nodes as you want the jobs to pack onto:

$ ./exp.sh sweep -- prepare --group grid --setup cnnenv -- train \
      --data cifar --num-blocks '{2,3,4}' --kernel-size '{3,5}'
$ ./exp.sh submit --nodes 4 --wait -- sweep-pilot --group grid

The submit drain call is the part that does the draining. The two calls before it — job reconcile and job reset — are there for a different reason: surviving a walltime kill, which is the subject of the next step.

Note

By default the daemon shares the pilot’s first node with the jobs it drains onto the allocation. The daemon is not compute-intensive, so this is usually fine. But if there is any chance it perturbs the sub-jobs in a way that matters — most often when a sub-job’s own run time is part of what the experiment measures — it is better to keep that node for the daemon alone: allocate one extra node and declare the pilot with @with_pilot --exclude-launch-host, which drops the daemon’s host from the pool so the sub-jobs run on the other nodes only.

Step 7 — One pilot, resumed until the sweep is done#

A pilot runs under a wall-clock limit. A large sweep may not finish inside one allocation’s walltime, and when the limit is reached the scheduler kills the pilot with SIGKILL. That hard kill gives the compute-side bookkeeping no chance to run, so the jobs that happened to be mid-flight are left stranded: their rows stay in running or submitting state forever, even though nothing is running them. The still-prepared remainder of the sweep is untouched, but those stranded jobs would never run again on their own.

The pilot body handles this in three moves, which is why it is written the way it is:

job reconcile asks the scheduler the true fate of every running / submitting job in the group and corrects the stranded rows to killed.
job reset moves those killed jobs back to prepared, so they are eligible to run again.
submit drain drains the whole prepared set — the recovered jobs plus whatever never got released the first time.

On the pilot’s first run there is nothing stranded and nothing killed, so the first two calls are no-ops and the third simply drains the sweep. On a resume they recover the interrupted jobs and the drain carries on. So resuming a sweep that ran out of walltime is just resubmitting the very same pilot, with more walltime, or as many times as it takes:

$ ./exp.sh submit --walltime 06:00:00 --nodes 4 --wait -- sweep-pilot --group grid

The sweep’s product is prepared once and lives durably in the database; the pilot is a stateless drainer that re-enters the reconcile → reset → drain cycle every time it starts. That division — a durable plan, a disposable runner — is the operational payoff of the whole prepare/drain design.

Note

The three calls can be collapsed into one command: submit drain takes --reconcile (do the reconcile pass first) and --retry killed,failed (re-queue those states before draining), so the body could read simply knit submit drain --group grid --reconcile --retry killed,failed. The three-step form is spelled out here because it shows exactly which state each call repairs.

Step 8 — When the grid explodes: Bayesian optimization#

The grid in Steps 3–5 swept two parameters. A good architecture depends on more than two, though. Suppose we also let the search vary the dropout rate, the activation function, whether the network uses residual connections, the learning rate, and the batch size — five more parameters on the train job:

Shorthand

@with_optional "dropout:real"          "0.3"     "Dropout before the classifier head." --domain '[0.0..0.5]'
@with_optional "activation:string"     "relu"    "Activation function."                --domain '{relu,gelu,leaky-relu}'
@with_optional "residual:string"       "false"   "Residual skip connections."          --domain '{false,true}'
@with_optional "lr:real"               "0.001"   "Initial learning rate."              --domain '[0.0001..0.1]' --log
@with_optional "batch-size:integer"    "128"     "Mini-batch size."                    --domain '{32,64,128,256}'

Long form

knit_with_optional "dropout:real"          "0.3"     "Dropout before the classifier head." --domain '[0.0..0.5]'
knit_with_optional "activation:string"     "relu"    "Activation function."                --domain '{relu,gelu,leaky-relu}'
knit_with_optional "residual:string"       "false"   "Residual skip connections."          --domain '{false,true}'
knit_with_optional "lr:real"               "0.001"   "Initial learning rate."              --domain '[0.0001..0.1]' --log
knit_with_optional "batch-size:integer"    "128"     "Mini-batch size."                    --domain '{32,64,128,256}'

The search space is now eight-dimensional. Even at a handful of values per parameter the cartesian product runs to thousands of trainings — days or weeks of compute to enumerate exhaustively, most of it spent on architectures that were never going to be good. A grid is the wrong tool at this size.

The right tool is adaptive search: instead of enumerating every point, build a model of how the parameters affect accuracy, and let it propose the next architecture to try based on everything seen so far. That is Bayesian optimization, and Knit provides it through the ax plugin (knit_import ax). It drives the same train job — it only changes how the next set of parameters is chosen.

Like the durable sweep, an ax search is prepared, then run. ax prepare records a campaign: the objective column and its direction (--objective val_accuracy --maximize), the campaign’s lifetime trial budget (--max-trials), a quasi-random warm-up (--sobol-trials), and an RNG --seed for reproducibility. The downstream is an ordinary submit -- train line, with @domain on each parameter ax should search (drawing from the --domain declared in Step 2) and a fixed value on the rest:

$ ./exp.sh ax prepare --campaign tune --objective val_accuracy --maximize \
      --max-trials 40 --sobol-trials 8 --seed 1 \
      -- submit --setup cnnenv -- train --data cifar \
         --num-blocks @domain --base-channels @domain --kernel-size @domain \
         --dropout @domain --activation @domain --residual @domain \
         --lr @domain --batch-size @domain

ax optimize then runs the loop: it asks the model for a set of parameters, submits a train job with them, reads back val_accuracy, tells the model the result, and repeats until the budget is spent — converging on a strong architecture in dozens of trainings rather than thousands:

$ ./exp.sh ax optimize --campaign tune

The optimizer runs in a private uv-backed environment that is built during bootstrap. It is resumable: interrupt it and run the same command again to continue from the campaign’s persisted state. Two per-call options shape a run: --max-inflight N keeps N trials in flight at once, and --trials K stops after K more trials (handy for advancing a long campaign in bounded chunks). When a campaign reaches its --max-trials budget but you want to keep searching, raise the budget with ax extend --campaign tune --max-trials 80 and run ax optimize again.

Running it inside a pilot works exactly as in Step 6 — a pilot will call ax optimize instead of submit drain, so every trial the loop submits packs onto the pilot’s allocation:

Shorthand

# The ax equivalent of the sweep-draining pilot from Step 6. Its body runs
# `ax optimize` instead of `submit drain`: the ask -> submit -> observe loop
# submits each trial job, and inside the pilot allocation those submits pack onto
# the pilot's own nodes. The pilot trains nothing itself (@without_setup); each
# trial job carries the --setup frozen into the campaign at prepare time.
# Resuming is just resubmitting the same pilot at the same campaign: ax reads its
# persisted state and continues from the trials already done.
@job "tune-pilot" "Optimize an ax campaign inside one pilot allocation."
@without_setup
@with_pilot
@with_optional "campaign:string" "tune" "Name of a prepared ax campaign to optimize."
@with_optional "max-inflight:integer" "4" "Trials to keep in flight on the pilot's nodes."
_tune_pilot_job() {
    knit ax optimize \
        --campaign     "$(knit_get_parameter campaign "$@")" \
        --max-inflight "$(knit_get_parameter max-inflight "$@")"
}
@done

Long form

# The ax equivalent of the sweep-draining pilot from Step 6. Its body runs
# `ax optimize` instead of `submit drain`: the ask -> submit -> observe loop
# submits each trial job, and inside the pilot allocation those submits pack onto
# the pilot's own nodes. The pilot trains nothing itself (@without_setup); each
# trial job carries the --setup frozen into the campaign at prepare time.
# Resuming is just resubmitting the same pilot at the same campaign: ax reads its
# persisted state and continues from the trials already done.
knit_register_job "tune-pilot" _tune_pilot_job "Optimize an ax campaign inside one pilot allocation."
knit_without_setup
knit_with_pilot
knit_with_optional "campaign:string" "tune" "Name of a prepared ax campaign to optimize."
knit_with_optional "max-inflight:integer" "4" "Trials to keep in flight on the pilot's nodes."
_tune_pilot_job() {
    knit ax optimize \
        --campaign     "$(knit_get_parameter campaign "$@")" \
        --max-inflight "$(knit_get_parameter max-inflight "$@")"
}
knit_done

Prepare the campaign, then submit the pilot over enough nodes for the trials to overlap — the pilot’s ax optimize --max-inflight (see the job body above) decides how many run at once. Resuming a campaign the walltime cut short is, again, just resubmitting the same pilot at the same campaign — ax keeps the trials already done and continues the rest:

$ ./exp.sh ax prepare --campaign tune-p --objective val_accuracy --maximize \
      --max-trials 40 --sobol-trials 8 --seed 1 \
      -- submit --setup cnnenv -- train --data cifar \
         --num-blocks @domain --base-channels @domain --kernel-size @domain \
         --dropout @domain --activation @domain --residual @domain \
         --lr @domain --batch-size @domain
$ ./exp.sh submit --nodes 4 --wait -- tune-pilot --campaign tune-p --max-inflight 4

Note

Trial run times can spread widely: a small --batch-size does more steps per epoch, and a large --base-channels or --num-blocks makes a heavier CNN, so one trial can take many times longer than another. Size the --max-trials budget and the walltime for the slow end of the range, not the average.

Step 9 — State the finding#

The search is done; now let’s read the answer back. ax show summarizes a campaign — its status, how many trials ran, and the best objective value seen:

$ ./exp.sh ax show --campaign tune
campaign:   tune
status:     completed
objective:  val_accuracy (maximize)
trials:     40/40
seed:       1
best:       78.4
opt-id:     tune
command:    submit --setup cnnenv -- train --data cifar --num-blocks @domain ...

That gives the best score, but we asked a question about an architecture. Every train run recorded its parameters and accuracy in the train table (as in Part I), so a one-line query returns the winning configuration — the best of everything we tried, grid and campaign alike:

$ ./exp.sh query sql --format column --header --exec "
    SELECT num_blocks, base_channels, kernel_size, activation, residual,
           round(lr, 4) AS lr, batch_size, val_accuracy
    FROM train ORDER BY val_accuracy DESC LIMIT 1"
num_blocks  base_channels  kernel_size  activation  residual  lr      batch_size  val_accuracy
----------  -------------  -----------  ----------  --------  ------  ----------  ------------
4           64             3            gelu        true      0.0123  128         78.4

We can now answer the question this part opened with, in one plain sentence: on CIFAR-10, the best architecture we found is a four-block CNN with 64 base channels, 3×3 kernels, GELU activations and residual connections, trained at a learning rate near 0.012 — it reaches 78.4% validation accuracy. Record the finding alongside the runs so it travels with the experiment:

$ ./exp.sh metadata store --key best-architecture \
      --value "4 blocks, 64 channels, k3, gelu, residual, lr~0.012 -> 78.4% val acc (campaign tune)"

That is where Part III ends: a finding, stated in English and backed by every recorded run that led to it.

Note

Beyond finding the best point. The same campaign machinery can fit a surrogate model of the objective instead of chasing its optimum. Prepare with ax prepare --model (in place of --maximize / --minimize) and the loop explores where the model is least certain. You then ask the fitted model for the objective at any in-space point with ax predict — written as the prepare command with a concrete value in each varying position — and check the model’s fit with ax show --diagnostics. The Optimize a job with Bayesian search recipe in the Stitch Guide covers it.

The complete experiment (Part III)#

Here is the whole experiment in one file — the setup, the resource, the eight-parameter train job, and the two pilots (sweep-pilot for the grid, tune-pilot for the campaign). Save it as exp.sh next to a copy of knit.sh:

Shorthand

#!/bin/bash

# doc-check: source-only
#
# The complete Part III experiment in one file: a uv-backed setup that installs
# the cifar-cnn classifier, a cifar resource, the train job with its full
# eight-parameter search space, and the two pilot jobs --- sweep-pilot (drains a
# prepared grid, resumable) and tune-pilot (runs an ax campaign). This is the
# final state of the experiment built up over Steps 1--8.
#
# Source-only: training needs uv, the dataset, and real compute, pilots need a
# scheduler, and "ax optimize" drives a real Ax environment, none of which run in
# plain CI, so check-docs only syntax-checks this file. Behavior was validated
# end to end during development; the ax loop's convergence is covered live by the
# 40_ax_optimize integration test.

source knit.sh

# uv and ax ship as plugins. Importing uv adds the uv environment manager the
# setup uses; importing ax adds the `ax` optimizer commands (ax imports uv too).
knit_import uv
knit_import ax

knit_set_program_description \
    "Tune a CNN image classifier on CIFAR-10 with a sweep and Bayesian optimization."

# --- Setup ------------------------------------------------------------------
# A uv environment with the cifar-cnn classifier installed from its git repo.
@setup "cnnenv" "A uv environment with the cifar-cnn classifier installed."
@with_uv_specs "cifar-cnn @ git+https://github.com/knit-sh/image-classifier-example.git"
@empty
@done

# --- Resource ---------------------------------------------------------------
# The CIFAR-10 dataset, downloaded and unpacked by `knit fetch`.
@resource "cifar" "The CIFAR-10 dataset, downloaded and unpacked."
@with_url "https://www.cs.toronto.edu/~kriz/cifar-10-python.tar.gz"
@done

# --- Job: train -------------------------------------------------------------
# Trains one architecture and records its validation accuracy. Every searchable
# parameter carries a --domain; dataset and epochs are fixed.
@job "train" "Train the CNN and record its validation accuracy."
@with_setup    "cnnenv"
@with_resource "data:cifar" "Name of the fetched CIFAR instance."
@with_optional "dataset:string"        "cifar10" "Dataset loader (not searched)."
@with_optional "num-blocks:integer"    "3"       "Convolution blocks."          --domain '[1..5]'
@with_optional "base-channels:integer" "32"      "Channels in the first block." --domain '{16,32,64}'
@with_optional "kernel-size:integer"   "3"       "Convolution kernel size."     --domain '{3,5}'
@with_optional "dropout:real"          "0.3"     "Dropout before the classifier head." --domain '[0.0..0.5]'
@with_optional "activation:string"     "relu"    "Activation function."                --domain '{relu,gelu,leaky-relu}'
@with_optional "residual:string"       "false"   "Residual skip connections."          --domain '{false,true}'
@with_optional "lr:real"               "0.001"   "Initial learning rate."              --domain '[0.0001..0.1]' --log
@with_optional "batch-size:integer"    "128"     "Mini-batch size."                    --domain '{32,64,128,256}'
@with_optional "epochs:integer"        "3"       "Training epochs (not searched)."
@with_output   "val_accuracy:real"     "0"       "Final-epoch validation accuracy (%)."
_train_job() {
    local datadir dataset residual out acc
    datadir="$(knit_resource_path "$(knit_get_parameter data "$@")")"
    dataset="$(knit_get_parameter dataset "$@")"

    local -a resflag=()
    residual="$(knit_get_parameter residual "$@")"
    [[ "${residual}" == "true" ]] && resflag=(--residual)

    out="$(cifar-cnn train \
        --model         "${KNIT_JOB_PREFIX%/}/model.pt" \
        --dataset       "${dataset}" \
        --data          "${datadir}" \
        --num-blocks    "$(knit_get_parameter num-blocks "$@")" \
        --base-channels "$(knit_get_parameter base-channels "$@")" \
        --kernel-size   "$(knit_get_parameter kernel-size "$@")" \
        --dropout       "$(knit_get_parameter dropout "$@")" \
        --activation    "$(knit_get_parameter activation "$@")" \
        "${resflag[@]}" \
        --lr            "$(knit_get_parameter lr "$@")" \
        --batch-size    "$(knit_get_parameter batch-size "$@")" \
        --epochs        "$(knit_get_parameter epochs "$@")" \
        2>&1 | tee /dev/stderr)"

    acc="$(grep -oE 'RESULT val_accuracy=[0-9.]+' <<<"${out}" | tail -1 | cut -d= -f2)"
    knit_output "val_accuracy" "${acc:-0}"
}
@done

# --- Job: sweep-pilot -------------------------------------------------------
# Drains a prepared sweep inside one allocation, resuming across walltime
# windows: job reconcile -> job reset -> submit drain. The drain's --max-inflight
# is sized to the pilot's node count (one single-node `train` per node) so the
# whole allocation works at once instead of one job at a time.
@job "sweep-pilot" "Drain a prepared sweep inside one pilot allocation."
@without_setup
@with_pilot
@with_optional "group:string" "grid" "Sweep group to drain."
_sweep_pilot_job() {
    local group nodes
    group="$(knit_get_parameter group "$@")"
    nodes="$(knit_job_hostnames | wc -l)"
    knit job reconcile --group "${group}"
    knit job reset     --group "${group}"
    knit submit drain  --group "${group}" --max-inflight "${nodes}"
}
@done

# --- Job: tune-pilot --------------------------------------------------------
# Runs an ax campaign inside one allocation; resuming is resubmitting at the
# same campaign.
@job "tune-pilot" "Optimize an ax campaign inside one pilot allocation."
@without_setup
@with_pilot
@with_optional "campaign:string" "tune" "Name of a prepared ax campaign to optimize."
@with_optional "max-inflight:integer" "4" "Trials to keep in flight on the pilot's nodes."
_tune_pilot_job() {
    knit ax optimize \
        --campaign     "$(knit_get_parameter campaign "$@")" \
        --max-inflight "$(knit_get_parameter max-inflight "$@")"
}
@done

knit "$@"

Long form

#!/bin/bash

# doc-check: source-only
#
# The complete Part III experiment in one file: a uv-backed setup that installs
# the cifar-cnn classifier, a cifar resource, the train job with its full
# eight-parameter search space, and the two pilot jobs --- sweep-pilot (drains a
# prepared grid, resumable) and tune-pilot (runs an ax campaign). This is the
# final state of the experiment built up over Steps 1--8.
#
# Source-only: training needs uv, the dataset, and real compute, pilots need a
# scheduler, and "ax optimize" drives a real Ax environment, none of which run in
# plain CI, so check-docs only syntax-checks this file. Behavior was validated
# end to end during development; the ax loop's convergence is covered live by the
# 40_ax_optimize integration test.

source knit.sh

# uv and ax ship as plugins. Importing uv adds the uv environment manager the
# setup uses; importing ax adds the `ax` optimizer commands (ax imports uv too).
knit_import uv
knit_import ax

knit_set_program_description \
    "Tune a CNN image classifier on CIFAR-10 with a sweep and Bayesian optimization."

# --- Setup ------------------------------------------------------------------
# A uv environment with the cifar-cnn classifier installed from its git repo.
knit_register_setup "cnnenv" knit_empty "A uv environment with the cifar-cnn classifier installed."
@with_uv_specs "cifar-cnn @ git+https://github.com/knit-sh/image-classifier-example.git"
knit_done

# --- Resource ---------------------------------------------------------------
# The CIFAR-10 dataset, downloaded and unpacked by `knit fetch`.
knit_register_resource "cifar" "The CIFAR-10 dataset, downloaded and unpacked."
knit_with_url "https://www.cs.toronto.edu/~kriz/cifar-10-python.tar.gz"
knit_done

# --- Job: train -------------------------------------------------------------
# Trains one architecture and records its validation accuracy. Every searchable
# parameter carries a --domain; dataset and epochs are fixed.
knit_register_job "train" _train_job "Train the CNN and record its validation accuracy."
knit_with_setup    "cnnenv"
knit_with_resource "data:cifar" "Name of the fetched CIFAR instance."
knit_with_optional "dataset:string"        "cifar10" "Dataset loader (not searched)."
knit_with_optional "num-blocks:integer"    "3"       "Convolution blocks."          --domain '[1..5]'
knit_with_optional "base-channels:integer" "32"      "Channels in the first block." --domain '{16,32,64}'
knit_with_optional "kernel-size:integer"   "3"       "Convolution kernel size."     --domain '{3,5}'
knit_with_optional "dropout:real"          "0.3"     "Dropout before the classifier head." --domain '[0.0..0.5]'
knit_with_optional "activation:string"     "relu"    "Activation function."                --domain '{relu,gelu,leaky-relu}'
knit_with_optional "residual:string"       "false"   "Residual skip connections."          --domain '{false,true}'
knit_with_optional "lr:real"               "0.001"   "Initial learning rate."              --domain '[0.0001..0.1]' --log
knit_with_optional "batch-size:integer"    "128"     "Mini-batch size."                    --domain '{32,64,128,256}'
knit_with_optional "epochs:integer"        "3"       "Training epochs (not searched)."
knit_with_output   "val_accuracy:real"     "0"       "Final-epoch validation accuracy (%)."
_train_job() {
    local datadir dataset residual out acc
    datadir="$(knit_resource_path "$(knit_get_parameter data "$@")")"
    dataset="$(knit_get_parameter dataset "$@")"

    local -a resflag=()
    residual="$(knit_get_parameter residual "$@")"
    [[ "${residual}" == "true" ]] && resflag=(--residual)

    out="$(cifar-cnn train \
        --model         "${KNIT_JOB_PREFIX%/}/model.pt" \
        --dataset       "${dataset}" \
        --data          "${datadir}" \
        --num-blocks    "$(knit_get_parameter num-blocks "$@")" \
        --base-channels "$(knit_get_parameter base-channels "$@")" \
        --kernel-size   "$(knit_get_parameter kernel-size "$@")" \
        --dropout       "$(knit_get_parameter dropout "$@")" \
        --activation    "$(knit_get_parameter activation "$@")" \
        "${resflag[@]}" \
        --lr            "$(knit_get_parameter lr "$@")" \
        --batch-size    "$(knit_get_parameter batch-size "$@")" \
        --epochs        "$(knit_get_parameter epochs "$@")" \
        2>&1 | tee /dev/stderr)"

    acc="$(grep -oE 'RESULT val_accuracy=[0-9.]+' <<<"${out}" | tail -1 | cut -d= -f2)"
    knit_output "val_accuracy" "${acc:-0}"
}
knit_done

# --- Job: sweep-pilot -------------------------------------------------------
# Drains a prepared sweep inside one allocation, resuming across walltime
# windows: job reconcile -> job reset -> submit drain. The drain's --max-inflight
# is sized to the pilot's node count (one single-node `train` per node) so the
# whole allocation works at once instead of one job at a time.
knit_register_job "sweep-pilot" _sweep_pilot_job "Drain a prepared sweep inside one pilot allocation."
knit_without_setup
knit_with_pilot
knit_with_optional "group:string" "grid" "Sweep group to drain."
_sweep_pilot_job() {
    local group nodes
    group="$(knit_get_parameter group "$@")"
    nodes="$(knit_job_hostnames | wc -l)"
    knit job reconcile --group "${group}"
    knit job reset     --group "${group}"
    knit submit drain  --group "${group}" --max-inflight "${nodes}"
}
knit_done

# --- Job: tune-pilot --------------------------------------------------------
# Runs an ax campaign inside one allocation; resuming is resubmitting at the
# same campaign.
knit_register_job "tune-pilot" _tune_pilot_job "Optimize an ax campaign inside one pilot allocation."
knit_without_setup
knit_with_pilot
knit_with_optional "campaign:string" "tune" "Name of a prepared ax campaign to optimize."
knit_with_optional "max-inflight:integer" "4" "Trials to keep in flight on the pilot's nodes."
_tune_pilot_job() {
    knit ax optimize \
        --campaign     "$(knit_get_parameter campaign "$@")" \
        --max-inflight "$(knit_get_parameter max-inflight "$@")"
}
knit_done

knit "$@"

From a cluster login node, the full arc is: provision, fetch the data, build the environment, prepare and drain a grid inside a pilot, then prepare and run a Bayesian campaign inside a pilot, and read the result:

$ ./exp.sh bootstrap  # on a cluster, add --profile <machine>
$ ./exp.sh fetch --name cifar -- cifar --uncompress
$ ./exp.sh setup --name cnnenv -- cnnenv
$ ./exp.sh sweep -- prepare --group grid --setup cnnenv -- train \
      --data cifar --num-blocks '{2,3,4}' --kernel-size '{3,5}'
$ ./exp.sh submit --nodes 4 --wait -- sweep-pilot --group grid
$ ./exp.sh ax prepare --campaign tune --objective val_accuracy --maximize \
      --max-trials 40 --sobol-trials 8 --seed 1 \
      -- submit --setup cnnenv -- train --data cifar \
         --num-blocks @domain --base-channels @domain --kernel-size @domain \
         --dropout @domain --activation @domain --residual @domain \
         --lr @domain --batch-size @domain
$ ./exp.sh submit --nodes 4 --wait -- tune-pilot --campaign tune --max-inflight 4
$ ./exp.sh ax show --campaign tune

On a laptop it is the same script with no pilots: drain the grid directly with submit drain --group grid and run the campaign with ax optimize --campaign tune. The experiment script never changes.

Part III — Searching for the best architecture

Contents

Part III — Searching for the best architecture#

The program#

Step 1 — One training run#

Step 2 — Declare the search space#

Step 3 — A sweep by hand#

Step 4 — knit sweep#

Step 5 — Durable prepare and drain#

Step 6 — Run the sweep inside a pilot job#

Step 7 — One pilot, resumed until the sweep is done#

Step 8 — When the grid explodes: Bayesian optimization#

Step 9 — State the finding#

The complete experiment (Part III)#

Step 4 — `knit sweep`#