Tao Run Automl
NVIDIAOfficialCreated by NVIDIA Corporation
- Skill
- Workflow
- Software development
- Apache-2.0
- Medium risk
- Has scripts
AutoGineer summary
Written by AutoGineer from the skill’s own files, not by the author
The tao-run-automl skill provides a structured, platform-agnostic workflow for running hyperparameter optimization (HPO) on NVIDIA TAO models. It integrates with various platforms like Brev, SLURM, Kubernetes, and Docker to manage the lifecycle of training, evaluation, and compression tasks. The skill handles algorithm selection, experiment tracking via WandB, and provides a rigorous pre-launch review process to ensure compute resources are used efficiently. It supports advanced search strategies including Bayesian, Hyperband, ASHA, and LLM-guided optimization, ensuring that model tuning is both systematic and data-driven.
What it does
- Automates hyperparameter search using algorithms like Bayesian, ASHA, and BOHB.
- Executes platform-agnostic jobs across Docker, Kubernetes, SLURM, and Brev.
- Performs mandatory baseline evaluations before launching optimization trials.
- Integrates LLM-based analysis for periodic experiment review and search space refinement.
- Manages experiment tracking and result interpretation for model training, distillation, pruning, and quantization.
- Enforces strict pre-launch reviews and resource budget checks.
When to use it
- When you need to optimize hyperparameters for a NVIDIA TAO network.
- When you want to automate the tuning of training, evaluation, inference, distillation, pruning, or quantization tasks.
- When you require a systematic approach to model performance improvement using advanced search algorithms.
- When you need to track and compare multiple experiment configurations efficiently.
Advantages
- Reduces manual effort in finding optimal model configurations.
- Provides consistent, reproducible experiment tracking.
- Supports diverse compute environments without changing the workflow logic.
- Prevents resource waste through pre-launch validation and baseline evaluation.
- Offers flexible search strategies tailored to different budget and performance requirements.
Usage examples
- “Run AutoML for my object detection model on Kubernetes using Bayesian search.”
- The agent will initiate the preflight checks for the Kubernetes platform, validate the model schema, perform a baseline evaluation, and begin the Bayesian optimization process.
- “Optimize the pruning parameters for my classification model using ASHA.”
- The agent will configure the ASHA algorithm, set up the pruning action, and manage the trial lifecycle to find the best-performing pruned model.
Install
Download the skill as a ZIP:
https://mp-staging.autogineer.com/agent-skills/nvidia/tao-run-automl/downloadUnzip it into the cross-client skills folder of your project:
mkdir -p .agents/skills && unzip tao-run-automl.zip -d .agents/skills
Documentation by NVIDIA Corporation
TAO AutoML
Standalone install? If this session was not initialized by the TAO skill bank plugin, run
tao-setupfirst (host preflight, credentials, cross-skill discovery).
Run automated hyperparameter optimization for a TAO model by combining:
The selected model skill under
skills/models/<model_skill>/.The selected platform skill under
skills/platform/<platform>/.AutoMLRunner, which generates recommendations, launches selected action jobs, extracts metrics, and feeds results back to the optimizer.
Do not launch until model metadata, platform preflight, data visibility, credentials, image choice, and compute shape are all proven.
Execution Runtime — Hard Gate
Every recommendation, baseline evaluation, per-recommendation evaluation, and
final evaluation runs in the selected model action's resolved
container_image by default. Resolve it from the model skill before any
training-environment setup. A local checkpoint or Hugging Face model ID does not
change this rule.
Use venv-based model execution only when explicitly requested. Never infer
venv mode from local-docker, local GPUs, Python, or pyproject.toml. If
absent, execution is container-backed. A host/controller venv for tao_automl,
TAO SDK, or a platform adapter is
control-plane-only; keep child model actions in the resolved container image.
Reference Map
references/skill_info.yaml: this workflow's structured metadata.Split detailed references:
automl-preflight-concepts.mdfor prerequisites and support checks;automl-intent-algorithms.mdfor search policy;automl-compression-literature.mdfor distill/prune/quantize algorithm sufficiency and future compression-search roadmap;automl-runner-configuration.mdfor runner/API/WandB details;automl-advanced-monitoring.mdfor hooks, resume, and pitfalls; andautoml-examples.mdfor conversation examples; andautoml-common-pitfalls.mdfor recurring safety checks.detailed-guide.mdis only the map.skills/models/<network>/SKILL.md: model-specific dataset requirements, metrics, HPO notes, checkpoint handoff, and known failures.skills/models/<network>/references/skill_info.yaml: action contract, container image, inputs, outputs, upload exclusions, andmode.skills/platform/<platform>/SKILL.md: selected platform preflight, credentials, resource shape, monitoring, and cancellation.skills/core/tao-launch-workflow/SKILL.md: shared intake pattern for platform, credentials, dataset visibility, image confirmation, and user confirmation.
Preflight
Run the shared launch intake. If the user has not chosen a platform, ask; Brev, SLURM, Kubernetes, and Docker are equal peers.
Run the selected platform skill's preflight before generating runner files.
Verify
nvidia-tao-automlimports:
python -c "import tao_automl; from tao_automl.runner import AutoMLRunner; print('OK')"Then verify the selected platform's SDK constructs — importing tao_automl does
not prove the platform backend is installed (e.g. DockerSDK() raises
CredentialError without the docker package). See
automl-preflight-concepts.md.
If missing, show the exact install command from versions.yaml and ask before
installing:
SB="${TAO_SKILL_BANK_PATH:-~/tao-skill-bank}"
pip install "$($SB/scripts/resolve_versions_key.py wheels.tao_automl_<platform>)"Valid platform wheel keys are tao_automl_brev, tao_automl_slurm,
tao_automl_kubernetes, tao_automl_docker, and tao_automl_all. Use
all only for development machines that need every backend. Add ,llm only
when the user requests LLM-guided algorithms.
Model Support Gate
Before every run:
Read the model
SKILL.mdandreferences/skill_info.yaml.Confirm
automl_enabled: truefor the model or that the model skill explicitly routes the selected action to AutoML.Confirm
<skill_dir>/schemas/<action>.schema.jsonexists and parses. This is the AutoML search-space gate.For non-TAO-Core models such as Cosmos-RL and CLIP, also require
references/spec_template_<action>.yaml; otherwise the runner has no complete action defaults.If any gate fails, do not improvise a search space. Report the missing package artifact.
Inputs
Collect these before runner construction:
| Input | Requirement |
|---|---|
model_skill | Resolved model skill directory under skills/models/. Resolve user aliases such as network_arch to the packaged skill directory first. |
network_arch | Read from the resolved model skill metadata. |
action | Action to optimize — train, evaluate, inference, distill, prune, or quantize with a packaged schema/template. |
platform | One of the supported TAO platform skills. |
train_dataset / eval_dataset / action inputs | Use model-specific spec keys and layout. Non-train actions may also need parent/teacher checkpoints, calibration data, or pruned artifacts. |
results_root | Local, Lustre, or S3 path appropriate for the platform. |
gpu_count, num_nodes | Respect model and platform limits. |
container_image | Resolve through model metadata and versions.yaml; show it to the user. |
automl_algorithm | Default bayesian unless user asks for another algorithm or the model skill recommends one. |
metric, direction | Prefer the model skill's validation/task metric. |
automl_budget | Recommendation count, max epochs/rungs, concurrency, or population size as required by the algorithm. |
Never ask for secret values. Verify required env vars with
[ -n "$VAR_NAME" ] && echo SET || echo UNSET.
Pre-Launch Review Gate
Before launching any recommendation jobs, show a concrete launch review and get user confirmation. This gate applies to every AutoML run for every AutoML-supported model/network; it is not Cosmos-specific and must not be scoped to a single model skill. This applies even when platform and image preflight already passed. The review must include:
model/network, platform, image, GPU/node shape, and result/workspace root
dataset mode and concrete spec keys, including train/eval sample counts when they can be read cheaply
algorithm, budget, max concurrent jobs, metric, and direction
searchable parameters and ranges, including default values when the user did not provide an explicit search space
exact generated recommendation configs for the initial launch batch, produced in a review-only step before any recommendation job is submitted
estimated runtime per recommendation and total expected wall time, with the assumptions used
the automatic baseline eval job id, metric value, and result path from the post-preflight eval job, or an explicit blocker if the model has no runnable evaluate action or validation data
the post-AutoML final evaluation plan for the selected best checkpoint/model, including metric, dataset, and record path
If the estimate is longer than the user's stated limit or materially longer than a normal interactive run, ask whether to reduce recommendations, epochs, dataset size, validation frequency, or search space before launch. Do not hide multi-day estimates in logs.
Automatic Baseline Eval Job
After platform, image, credential, data, and model preflight pass, run the model's evaluate action once on the selected validation/eval data before submitting any AutoML recommendation jobs. This is required AutoML setup, not an optional "pretrained eval" question for the user. Use the same base model or checkpoint that the AutoML training run starts from, the model skill's evaluate spec/template, and the selected platform's normal job submission path. If the model skill recommends a smaller shape for evaluation than training, use that shape and call it out in the launch review.
Share the eval metric number in the launch review before asking for confirmation. If a starting checkpoint exists but the baseline cannot be produced — no packaged evaluate action, missing eval dataset, failed eval job — stop and report the blocker instead of silently falling back to a training-loss-only run.
For training from scratch, record the baseline as unavailable and proceed; do not evaluate an empty checkpoint.
The runner owns final evaluation. When eval is runnable, pass
final_eval_fn(best_rec, train_job_id) to AutoMLRunner.run; the result then
carries result["final_evaluation"]. See automl-preflight-concepts.md for
the callback, checkpoint, baseline, and from-scratch rules.
Dependency And Data Preflight
If the selected workflow needs object storage or a platform CLI and the tool is
missing, report the missing dependency and offer the exact install command
before continuing. After user approval, rerun
scripts/check_tao_launch_preflight.py with --install-missing-tools so it
installs the smallest needed package and immediately retries path verification.
For S3 paths, verify both credentials and path readability from the launch
platform before creating runner artifacts.
For models that read large media archives or directories during every training
trial, stage or extract the dataset once to storage visible from the execution
platform, then point all recommendation specs at that staged path. Record the
source URI, staged path, byte/file-count evidence when available, and timestamp
in <workspace>/evaluations/data_staging.json. If staging is not possible,
include the repeated S3 I/O risk in the pre-launch review and ask before
spending a long AutoML budget on it.
When the model skill defines sample-count-sensitive constraints, enforce them
before launch. Reject or cap every batch-size recommendation that would create
zero training steps for the selected dataset and GPU shard count. Use
scripts/check_tao_launch_preflight.py --effective-batch-limit train_annotation=<batch_size>,<shard_count> for each generated recommendation
before submitting it. If a recommendation later fails because the data is too
small for the effective batch size, classify it as an invalid configuration,
replace or adjust it only when remaining budget exists, and report the
correction in the final summary.
When train sample count is known from an annotation file or cheap manifest read,
pass it as automl_settings["train_sample_count"] to AutoMLRunner.run so the
runner can cap impossible recommendations before submitting a job and record the
adjustment in result["history"][i]["adjustments"].
Algorithm Policy
| Algorithm | Good fit | Required knobs |
|---|---|---|
bayesian | Default for small/medium budgets and few parameters. | num_recommendations, metric, direction |
hyperband, asha | Many configs, cheap early rungs; ASHA is parallel-friendly. | max_epochs, reduction_factor, optional max_concurrent |
bohb, dehb | Mixed Bayesian/evolutionary search with multi-fidelity budgets. | same rung budget fields as Hyperband |
pbt | Long training where schedules should mutate during training. | population and generation budget |
llm, hybrid, autoresearch | User explicitly wants LLM-guided search with a configured endpoint. | LLM endpoint config plus budget |
For evaluate or inference, default to Bayesian/BFBO-style search over the
selected action's prompt, decoding, preprocessing, or runtime config knobs.
Use a task metric from the action outputs/logs and set direction explicitly
when the metric name is ambiguous. Do not use training-loss assumptions for
actions that do not update weights.
For distill, use the same train-like policy when the distill action performs
epoch-based optimization and writes checkpoints. For single-shot prune and
quantize, default to bayesian or bfbo unless the action schema/model skill
declares an epoch-like or calibration-budget field that makes
hyperband/asha/bohb/dehb meaningful. Use eval_fn when the selected
metric must be computed by a follow-up evaluate/inference action after the
compression action completes.
Prefer the model skill's recommendation over generic defaults. Avoid ASHA or Hyperband when the model skill says startup, validation, or checkpoint cost dominates short trials.
Spec And Search Space
Build specs as nested dictionaries. If a model skill lists paths in dotted notation for readability, walk the path and assign the nested leaf; do not store flat dotted strings as spec keys.
Use the packaged selected-action schema for:
automl_default_parametersautoml_disabled_parametersvalid min/max ranges
enums, option weights, conditions, dependencies, and popular parameters
User-provided search spaces must stay inside schema constraints. For integer knobs with discrete choices, include the schema's required integer option shape instead of a loose list if the model skill calls that out.
Data source overrides are mandatory unless the model skill says the launcher can derive them. Preserve exact user-provided spec keys when the dataset uses direct annotation/media paths.
Metric Policy
Training loss is cheap but can be misleading. Prefer the model skill's task metric. Use one of these:
Log metric:
metric=<name>,direction=maximize|minimize.metric_extractor(logs, metric_name): parse the model's logs when the default resolver is ambiguous.eval_fn(rec, train_job_id): run the model's evaluate action after each recommendation when the user wants a downstream task metric.
Do not map kpi to a metric unless the model skill explicitly defines that
mapping.
The final report must compare the baseline metric, each recommendation's
metric, and the selected best metric so users can see the impact of tuning. For
model skills that require an eval_fn to compute the real task metric, use
that evaluator instead of optimizing a convenient training loss unless the user
explicitly accepts the proxy metric.
Runner Construction
Use the selected platform SDK only after its preflight passes. Construct SDKs without embedding credentials in code.
sdk is the platform SDK object; containerless venv models use
VirtualEnvSDK(venv_path=..., work_dir=...). Always pass work_dir -- the
default ~/.tao_sdk/virtualenv fills the home directory with trial
checkpoints. See automl-runner-configuration.md.
import sys
from pathlib import Path
from tao_automl.runner import AutoMLRunner
skill_bank = Path("<absolute-tao-skill-bank>")
model_skill = "<resolved-model-skill-directory>"
skill_dir = skill_bank / "skills" / "models" / model_skill
sys.path.insert(
0, str(skill_bank / "skills/applications/tao-run-automl/scripts")
)
from resolve_automl_session import validate_session_settings
runner = AutoMLRunner(
sdk=sdk,
skill_dir=str(skill_dir),
action=action, # train, distill, prune, quantize, ...
)
workspace_path = Path("<automl_workspace>")
resume = False
# Mandatory fail-closed gate from this skill's bundled scripts directory.
validate_session_settings(
automl_settings,
resume=resume,
workspace=workspace_path if resume else None,
)
result = runner.run(
workspace_path=str(workspace_path), # timestamp it to avoid collisions
automl_settings=automl_settings, # must contain an explicit session_id
spec_overrides=spec_overrides,
automl_hyperparameters=automl_hyperparameters,
custom_param_ranges=custom_param_ranges,
metric_extractor=metric_extractor, # optional
eval_fn=eval_fn, # optional
final_eval_fn=final_eval_fn, # optional but required when final eval is runnable
resume=resume,
)Set automl_settings["session_id"] explicitly and call
validate_session_settings before every run. Generate a fresh ID once with
scripts/resolve_automl_session.py new. Resume only when explicitly requested;
resolve its controller with
scripts/resolve_automl_session.py resolve --workspace <full-run-path>.
Missing or ambiguous state is a blocker. See the
resume section of references/automl-advanced-monitoring.md for the complete
fresh/resume pattern.
Monitoring
Use runner status output and the platform SDK's get_job_status,
get_job_logs, and get_failure_analysis. For active jobs, report:
recommendation id / trial id
platform job id
status
current metric
best metric so far
selected hyperparameters for the current/best recommendation
elapsed time and updated ETA when enough timing data exists
On failure, classify whether it is infrastructure, data visibility, image, credential, spec/schema, or model-code failure. Fix only the minimal cause and do not silently spend additional budget on repeated invalid recommendations. If a blocker is fixed during run setup, continue from the original task after showing the updated preflight/launch review instead of leaving the user to restate the request.
For LLM-based algorithms, inspect the brain logs before calling the run valid. Verify that LLM calls succeeded, proposals were generated, prior metrics were used to choose later parameter changes, and logs show keep/discard or equivalent algorithm decisions. If the brain falls back to random sampling, classify the LLM workflow as failed or blocked instead of treating it as a valid LLM-guided run.
Result Handoff
At completion:
Identify the best recommendation by the selected metric and direction.
Return the best child job id and its result path.
Resolve the model checkpoint or action artifact using the model skill's checkpoint/artifact metadata and SDK helpers; do not guess filenames such as
latest.Report the exact search space, algorithm, budget, metric, and platform.
Report the automatic baseline eval job id/result path/metric, all recommendation metrics, final evaluation status/result path/metric, failed recommendations and root causes, elapsed time, and final runtime notes.
If this feeds a workflow such as AutoML + DEFT, pass the winning spec overrides and checkpoint through the workflow's declared handoff fields.
With the default retention policy, verify that cleanup-supported, safely prunable terminal trial artifacts were deleted and that the winning training artifacts remain. Report protected promotion/resume parents or conservative Hybrid results explicitly. A remote bind, named volume, or other output route the SDK cannot reclaim must fail retention preflight before the first trial rather than be silently retained.
Common Pitfalls
See references/automl-common-pitfalls.md before launching or recovering an
AutoML run.
Files
- ATTRIBUTION.mdAdded by AutoGineer912 B
- BENCHMARK.md7.6 KiB
- LICENSELicence file added by AutoGineer11.2 KiB
- SKILL.md19.4 KiB
- cards/00-preflight.md1.9 KiB
- cards/10-baseline-eval.md3.7 KiB
- cards/20-launch-recs.md7.0 KiB
- cards/30-interpret.md3.2 KiB
- cards/README.md3.3 KiB
- cards/driver.shScript8.0 KiB
- config/skillspector-baseline.yaml1.6 KiB
- eval.config9.8 KiB
- evals/evals.json1.7 KiB
- references/automl-advanced-monitoring.md22.6 KiB
- references/automl-common-pitfalls.md908 B
- references/automl-compression-literature.md5.6 KiB
- references/automl-examples.md6.0 KiB
- references/automl-intent-algorithms.md26.1 KiB
- references/automl-preflight-concepts.md13.1 KiB
- references/automl-runner-configuration.md17.3 KiB
- references/best_rec_adapter.pyScript4.1 KiB
- references/detailed-guide.md989 B
- references/skill_info.yaml1.8 KiB
- references/tests/test_best_rec_adapter.pyScript3.7 KiB
- scripts/resolve_automl_session.pyScript5.6 KiB
- skill-card.md5.2 KiB
- skill.oms.sig9.5 KiB
- tests/test_resolve_automl_session.pyScript5.3 KiB
references/automl-runner-configuration.md
TAO AutoML Runner Configuration
Runner construction, SDK handoff examples, settings keys, metrics, custom parameter ranges, spec overrides, and WandB setup.
Load this file only when the compact SKILL.md points here for the current task. If this reference conflicts with SKILL.md, skill_info.yaml, schemas, or platform/model skills, the compact/current source wins.
Contents
Minimal Example
Pick whichever SDK matches where you want trials to run. AutoMLRunner is
platform-agnostic — none of the SDKs is a default; the user picks.
from tao_sdk.platforms.slurm import SlurmSDK # SLURM cluster
from tao_sdk.platforms.kubernetes import KubernetesSDK # K8s (EKS / GKE / on-prem)
from tao_sdk.platforms.docker import DockerSDK # local Docker daemon
VirtualEnvSDK (containerless venv runs)
Full Example (all options)
LLM-Powered Algorithm Example
Programmatic API (without runner)
automl_settingskeyskpimetric resolutioncustom_param_rangesformatModel-specific search-space rules
LLM Analyzer (server-side range narrowing)
spec_overridesWandB Experiment Tracking
Setup
or (when reinstalling tao-run-automl with the wandb extra):
pip install "$(${TAO_SKILL_BANK_PATH:?}/scripts/resolve_versions_key.py wheels.tao_automl_brev | sed 's/]/,wandb]/')"How it works
Minimal WandB setup
Option 1: via config dict
Option 2: environment variable (simpler)
export WANDB_API_KEY=your-key
Dashboard features
Step 3: Configure and Run
Minimal Example
from datetime import datetime
from pathlib import Path
# Pick whichever SDK matches where you want trials to run. AutoMLRunner is
# platform-agnostic — none of the SDKs is a default; the user picks.
from tao_sdk.platforms.brev import BrevSDK # Brev GPU instances
# from tao_sdk.platforms.slurm import SlurmSDK # SLURM cluster
# from tao_sdk.platforms.kubernetes import KubernetesSDK # K8s (EKS / GKE / on-prem)
# from tao_sdk.platforms.docker import DockerSDK # local Docker daemon
from tao_automl.runner import AutoMLRunner
TIMESTAMP = datetime.now().strftime("%Y%m%d_%H%M%S")
sdk = BrevSDK() # reads platform credentials from env
runner = AutoMLRunner(
sdk=sdk,
skill_dir=SKILL_BANK / "skills" / "models" / model_skill, # resolved skill dir
action=action, # train, evaluate, inference, distill, prune, quantize, ...
)
result = runner.run(
train_dataset_uri=train_dataset_uri,
automl_settings={
"algorithm": algorithm,
"metric": metric,
"automl_max_recommendations": max_recommendations,
"session_id": session_id, # explicit, generated once
},
workspace_path=f"./automl_workspace/{TIMESTAMP}", # timestamped to avoid collisions
# Platform-specific create_job kwargs go here as **platform_kwargs.
# See each platform's SKILL.md for the kwargs each accepts.
gpu_count=8,
num_nodes=1,
gpu_type="H100", # Brev-specific
)VirtualEnvSDK (containerless venv runs)
Model skills with container_image: null and execution.type: python_script
(the Kumo-Forecast and Kumo-Anomaly skills) run trials as local
Python processes instead of containers. Their SDK is VirtualEnvSDK, and
unlike the credential-reading platform SDKs it takes constructor arguments:
from tao_sdk.platforms.virtualenv import VirtualEnvSDK
from tao_automl.runner import AutoMLRunner
sdk = VirtualEnvSDK(
venv_path="/abs/path/to/model-venv", # must contain pyvenv.cfg + bin/python
work_dir="/scratch/automl-jobs", # ALWAYS set this - see caution below
)
runner = AutoMLRunner(sdk=sdk, skill_dir=str(skill_dir), action="train")Full signature: VirtualEnvSDK(venv_path, work_dir=None, state_file=None).
Caution -
work_dirdefaults into your home directory. Whenwork_diris omitted the SDK resolves it to$TAO_SDK_STATE_DIR/virtualenv, or~/.tao_sdk/virtualenvwhen that env var is unset, and creates<work_dir>/jobs/on construction. Every trial's checkpoints then land under$HOME- roughly 1.4 GB per Kumo-Forecast trial. Always pass an explicitwork_diron a filesystem with room forautoml_max_recommendationscheckpoints. SettingTAO_SDK_STATE_DIRrelocates the default, but an explicitwork_diris the reviewable form and is what the launch review must show.
Job layout (this is where trial artifacts actually land):
<work_dir>/jobs/<job-id>/spec/spec.<yaml|json|toml>
<work_dir>/jobs/<job-id>/results/<declared_output_key_with_dots_as_underscores>/
<work_dir>/jobs/<job-id>/logs/job.logDeclared outputs from references/skill_info.yaml are rewritten into that
per-job results directory, so train.output_dir becomes
<work_dir>/jobs/<job-id>/results/train_output_dir. Resolve it at runtime with
sdk.get_job_results_dir(job_id); do not precompute it.
Do not pass image= for python_script execution - the venv is the runtime.
Verify construction during preflight. Importing the class is not enough: the
constructor is what validates venv_path, so an import-only check reports
success against a venv that cannot run a single trial.
VENV_PATH=/abs/path/to/model-venv \
WORK_DIR=/scratch/automl-jobs \
python - <<'PY'
import os
from tao_sdk.platforms.virtualenv import VirtualEnvSDK
VirtualEnvSDK(
venv_path=os.environ["VENV_PATH"],
work_dir=os.environ["WORK_DIR"],
)
print("OK")
PYUse the same venv_path and work_dir the runner will use. A missing or
malformed venv exits non-zero with
ValueError: Virtual environment does not exist: <venv_path>, so the check
fails before any trial launches. Construction also creates <work_dir>/jobs/,
which is the second reason to pass work_dir explicitly here.
The virtualenv extra is not exposed as its own versions.yaml wheel key;
wheels.tao_automl_all is currently the only key that includes it.
Full Example (all options)
def my_eval(rec, train_job_id):
"""Optional post-training evaluator. Return a float (the real metric)
or None to fall back to the log-based extractor."""
# e.g. read a results file uploaded by the container and compute the requested metric
...
return 0.71
result = runner.run(
# --- Required ---
train_dataset_uri=train_dataset_uri,
# --- Dataset + resources ---
eval_dataset_uri=eval_dataset_uri,
base_checkpoint="",
image=image, # only set to override skill_info's container_image
# --- AutoML config ---
automl_settings={
"algorithm": algorithm,
"metric": metric,
"direction": direction, # explicit when needed
"automl_max_recommendations": max_recommendations,
"session_id": session_id, # explicit, generated/resolved once
},
automl_hyperparameters=automl_hyperparameters, # from model skill / schema
custom_param_ranges=custom_param_ranges, # from model skill / user constraints
# --- Per-rec spec overrides ---
spec_overrides=spec_overrides, # mandatory model-specific overrides from model skill
# --- State + durability ---
workspace_path=f"./my_experiment/{TIMESTAMP}", # ALWAYS timestamp to avoid collisions
resume=False, # True → recovers in-flight jobs
# --- Hooks (all optional, opt-in) ---
metric_extractor=None, # custom log→metric parser
eval_fn=my_eval, # post-training real-metric eval
on_recommendation=lambda r: print(f"launching rec {r.id}: {r.specs}"),
on_result=lambda r, metric, status: print(f"rec {r.id} {status} → {metric}"),
# --- Platform create_job kwargs (forwarded as **platform_kwargs) ---
# SLURM: partition, account, num_nodes, gpu_count
# Kubernetes: namespace, node_selector, tolerations, num_nodes, gpu_count
# Docker: mounts, gpu_count
# Brev: instance_id, gpu_type, gpu_count
gpu_count=8,
num_nodes=1,
gpu_type="H100",
)LLM-Powered Algorithm Example
For llm, hybrid, or autoresearch, use the same generic runner shape as above, plus the required LLM endpoint, model, and key in automl_settings. All model-specific hyperparameters, metric extractors, and spec_overrides must still come from the model skill.
LLM endpoint configuration (in order of precedence):
automl_settingskeys:llm_endpoint,llm_model,llm_api_key; aliasesbase_url,model, andapi_keyare also accepted.Environment variables:
AUTOML_LLM_ENDPOINT,AUTOML_LLM_MODEL,AUTOML_LLM_API_KEYFallback env var for API key:
NVIDIA_API_KEYDefaults: NVIDIA inference endpoint (
https://inference-api.nvidia.com) withgcp/google/gemini-3.1-pro-preview. Always pass the endpoint, model, and key explicitly for reproducible LLM-based testing.
Programmatic API (without runner)
For tighter control, use the AutoML class directly:
from tao_automl import AutoML
automl = AutoML(
workspace="/tmp/my_experiment",
network=network_arch,
train_specs=my_train_spec_dict,
action=action,
settings={
"algorithm": "bayesian",
"metric": "loss",
"automl_max_recommendations": 10,
},
wandb_config={"enabled": True, "project": "my-project"},
)
while not automl.is_complete():
recs = automl.next_recommendation()
for rec in recs:
metric_value = train_model(rec.specs) # your training function
automl.report_result(rec.id, metric_value)
automl.finish() # close WandB run
print("Best:", automl.get_best().specs)automl_settings keys
| Key | Type | Default | Description |
|---|---|---|---|
algorithm | str | required | bayesian, hyperband, bohb, asha, bfbo, dehb, pbt, hyperband_es, llm, hybrid, autoresearch |
session_id | str | required by this skill | Stable controller identity. Generate it once for a fresh run with scripts/resolve_automl_session.py new; on resume recover it from the full run workspace. Never rely on the wheel's random default. |
metric | str | "loss" | Metric name. The implicit rule for direction is "contains 'loss' → minimize, else maximize". Override with direction. |
direction | "minimize" | "maximize" | inferred | Explicit direction. Required only when it disagrees with the implicit rule. The runner transparently inverts reported values so callers always see their metric in its original scale. |
automl_max_recommendations | int | 20 | Max trials (bayesian, bfbo, llm) |
automl_max_epochs | int | 27 | Epoch budget (hyperband, bohb, asha, dehb) |
automl_reduction_factor | int | 3 | Halving factor (hyperband variants) |
automl_max_concurrent | int | 4 | Max parallel configs (asha only) |
automl_population_size | int | 10 | Population size (pbt only) |
automl_max_experiments | int | 50 | Max experiments (autoresearch only) |
llm_endpoint | str | https://inference-api.nvidia.com | OpenAI-compatible API endpoint (llm, hybrid, autoresearch) |
llm_model | str | gcp/google/gemini-3.1-pro-preview | LLM model name (llm, hybrid, autoresearch) |
llm_api_key | str | from env | API key for the LLM endpoint |
research_program | str | None | Free-text research directives for the autoresearch agent |
evolvable_text_parameters | str | list[str] | [] | Searchable text parameters that autoresearch may rewrite beyond schema seed enums. |
automl_delete_intermediate_ckpt | bool | True | Automatically delete cleanup-supported terminal artifacts once they are safe to prune. The active trial/current best, every non-dominated Pareto-front member in a multi-objective search, and Hyperband-family promotion/resume inputs are protected; Hybrid conservatively retains successful trials when full-fidelity provenance is ambiguous. Cleanup-aware SDKs reject remote Docker binds, named volumes, unsafe output ownership, and unbound S3 identities before the first trial because deletion cannot be verified. Set False only when all trial artifacts and their external cleanup lifecycle are intentionally owned. |
override_automl_disabled_params | bool | False | Include params whose schema automl_enabled is False. For advanced users who want to search over params the network author didn't flag for AutoML. |
AutoMLRunner.run(..., feedback_fn=...) accepts a callback with signature
(recommendation, job_id) -> JSON-safe value. For reflective searches, return
compact training-split records containing the input/query, generated output,
and a failure-mode comment that does not reveal the correct answer. Do not return
gold/expected labels, item or media identifiers, or paths. The runner removes
common ground-truth and identifier fields, persists the sanitized feedback, and
supplies it to the next autoresearch proposal; it does not replace the numeric
validation selection metric.
kpi metric resolution
When metric="kpi", the controller resolves the actual metric key from the network config's metrics.monitoring_metric field. Whether kpi is appropriate, and whether a custom metric_extractor is needed, is model-specific. Follow the model skill's AutoML / HPO Notes.
custom_param_ranges format
Each entry can include:
| Field | Type | Description |
|---|---|---|
valid_min | float/int/list | Min value. For list-valued parameters, pass the list shape required by the schema. |
valid_max | float/int/list | Max value. Same list rules as min. |
valid_options | list[str] | For categorical/ordered params: restrict to these values |
option_weights | list[float] | Sampling weights for valid_options. Must match length. Higher weight = more likely to be sampled. |
disable_list | bool | For params that can be float OR list: True keeps it as a single float for optimization, bypassing network list helpers. Use only when supported by the schema/model skill. |
Example with all features:
custom_param_ranges={
"<float_param>": {"valid_min": min_value, "valid_max": max_value, "disable_list": True},
"<categorical_param>": {
"valid_options": ["option_a", "option_b"],
"option_weights": [0.7, 0.3],
},
"<list_param>": {"valid_min": [min_a, min_b], "valid_max": [max_a, max_b]},
}Model-specific search-space rules
Some networks have built-in search-space exclusions or algorithm restrictions. Do not document them here; read the model skill's AutoML / HPO Notes and let schema validation report unsupported combinations.
LLM Analyzer (server-side range narrowing)
The controller supports automatic range narrowing via the LLM analyzer. Enable via environment variables before launching:
os.environ["AUTOML_LLM_ANALYZER_ENABLED"] = "true"
os.environ["AUTOML_LLM_ANALYZER_INTERVAL"] = "5" # analyze every 5 completed recs
os.environ["AUTOML_LLM_ANALYZER_NARROW_RANGES"] = "true" # auto-tighten custom_param_rangesWhen enabled, after every N completed experiments the analyzer reviews patterns, assesses convergence, and optionally narrows search ranges to focus on promising regions. This happens server-side and persists the narrowed ranges.
spec_overrides
spec_overrides keys are model-specific. Read the model skill's Training Requirements, Per-Action Dataset Requirements, and Typical Spec Overrides sections, then pass only the keys required or recommended there. Do not infer override keys from examples in this AutoML skill.
Every key you pass is validated against the skill's spec schema. Typos that look like existing keys raise ValueError with a suggestion; genuinely-new keys are accepted with a warning.
WandB Experiment Tracking
AutoML optionally integrates with Weights & Biases to track all experiments in a single dashboard.
Setup
pip install wandb
# or (when reinstalling tao-run-automl with the wandb extra):
# pip install "$("${TAO_SKILL_BANK_PATH:?}/scripts/resolve_versions_key.py" wheels.tao_automl_brev | sed 's/]/,wandb]/')"How it works
When wandb_config={"enabled": True} is passed:
The controller creates a WandB run named
automl_brainin the specified project.All recommendations are grouped under a WandB group (e.g.
automl_abc123) so parent + child training runs appear together in the dashboard.After every result, a WandB table (
automl_experiments) is logged containing:experiment_id,job_id,status, metric value,best_epoch_numberAll varying hyperparameter values
Call
automl.finish()(or letrunner.run()complete) to finalize the WandB run.
Minimal WandB setup
# Option 1: via config dict
result = runner.run(
...,
wandb_config={
"enabled": True,
"project": "tao-hpo",
"api_key": "your-key", # or set WANDB_API_KEY env var
},
)
# Option 2: environment variable (simpler)
# export WANDB_API_KEY=your-key
result = runner.run(
...,
wandb_config={"enabled": True, "project": "tao-hpo"},
)Dashboard features
Once tracking is active, you can:
Compare all trials side-by-side in the WandB table view
Sort by metric to find the best config instantly
Group by hyperparameter to see which values correlate with good results
Link to child training runs if the compute backend also logs to WandB (group name is available via
automl.wandb_group)
Requirements and compatibility
Works with: Claude Code, Claude apps, OpenAI Codex, GitHub Copilot, Cursor, Gemini CLI, Kiro, OpenCode, Goose, Amp, Roo Code, JetBrains Junie, Factory Droid, Devin (ex-Windsurf), Mistral Vibe, OpenHands, Letta
“Works with” means the tool reads the open Agent Skills format; a skill written for one tool may use features others lack.
Needs: Declares allowed tools, Needs Python, Needs Node.js
The author says: Requires docker + nvidia-container-toolkit. Workflows declare additional requirements.
Security
What our scanner noticed· Medium risk
- Phrasing that asks for actions without confirmation
- Large base64-encoded blobs
- References to credential files
- Instructions that use the network
- Scripts that your agent may run
This skill includes scripts that your agent may run. Read them before you use it.
A static scan is not a guarantee. Skills are third-party content: review a skill before you let an agent use it.
Attribution and licence
- Created by
- NVIDIA Corporation
- Publisher
- NVIDIA (Official)
- Source
- NVIDIA/skills · skills/tao-run-automl · commit fce3cca
- Licence
- Apache License 2.0 (redistribution allowed)
- LICENSE (added by AutoGineer from LICENSE-APACHE)
- ATTRIBUTION.md (written by AutoGineer)
- Synced
AutoGineer is not affiliated with or endorsed by NVIDIA Corporation or NVIDIA.
Download ZIP
The ZIP holds the skill’s files, its licence files and an ATTRIBUTION.md. Download tao-run-automl.zip