Agent Jobs

March 22, 2026 ยท View on GitHub

wxtrain now has a canonical agent-facing job layer for dataset and model planning.

The goal is to let an agent take a high-level request such as:

build a week of HRRR severe-training data for a Swin transformer

and translate it into deterministic manifests instead of hand-managed ad hoc configs.

Commands

Initialize a starter job spec:

cargo run -p wx-cli --bin wxtrain -- train job-init --output examples\agent_job_swin.json --architecture swin-transformer --task forecasting --dataset-name hrrr_swin_demo

Plan the dataset and model recipe:

cargo run -p wx-cli --bin wxtrain -- train job-plan --spec examples\agent_job_swin.json

Build a dataset directly when the job spec uses existing_grib_files:

cargo run -p wx-cli --bin wxtrain -- train job-build --spec examples\agent_job_classical.json --output-dir examples\agent_classical_build --colormap heat

Job Spec Shape

An agent job spec includes:

  • the dataset/job name
  • the data source
  • feature profiles and custom channels
  • labels
  • the target model architecture and learning task
  • dataset export preferences

Supported architectures today:

  • classical_ml
  • diffusion
  • swin_transformer
  • forecast_graph_network
  • custom

Supported source modes today:

  • existing_grib_files: directly executable through train job-build
  • model_window: planner-only collection/fetch expansion outline for future orchestration

Outputs

train job-plan emits:

  • an architecture-aware dataset/export plan
  • a model recipe with trainer-family/loss/input-layout defaults
  • a ready-to-run GribDatasetBuildRequest when the source is executable locally
  • planned feature expansion and execution notes for the agent

For local existing_grib_files jobs, train job-build now materializes the planned training channels instead of only exporting raw decoded GRIB messages. That includes:

  • raw channels from the selected feature profiles
  • derived map/profile diagnostics such as CAPE/CIN, SRH, shear, STP/SCP, PWAT, and pressure-map fields
  • supported custom SRH contracts like custom_srh_500m and custom_srh_1km

Current Shape

The job spec uses these top-level keys:

  • job_name
  • dataset_name
  • description
  • data_source
  • features
  • labels
  • model
  • output

features is an object, not an array. It supports:

  • profiles
  • extra_channels
  • custom_features

labels are currently {name, source} contracts. Label type or loss behavior is carried by the model task, not per-label metadata.

train job-build writes:

  • job_plan.json
  • model_recipe.json
  • dataset_request.json
  • the standard dataset build artifacts and shards

The runnable classical example uses the synthetic severe fixture generated by:

powershell -NoProfile -File examples\make_agent_job_sample.ps1