cdf-orchestration-write-dag

Files (8)

• BENCHMARK.md
• SKILL.md
• evals.json · evals/
• component-skills-excerpt.md · references/
• gpu-footprint.md · references/
• trigger-examples.md · references/
• skill-card.md
• skill.oms.sig

SKILL.md (first lines)

# Write DAG

## What This Skill Does

Composes a new Kubernetes-only Airflow DAG from the **existing task groups and stages** that the
shipped DAGs are built from: the shared task groups in `airflow/dags/shared/task_groups/`, plus
the workflow-local stages the shipped DAGs define for themselves (for example DIG's simulation and
AnomalyGen tasks, AV's curation and retrieval tasks). It does not invent new task-group *logic*.
Output is:

1. `airflow/dags/workflows/<dag_name>_dag/configs/<dag_name>_k8s_manifest.yaml` — compute + component manifest
2. `airflow/dags/workflows/<dag_name>_dag/<dag_name>_dag.py` — the DAG builder
3. `airflow/dags/workflows/<dag_name>_dag/models/payload.py` — the payload pydantic model
4. `airflow/dags/workflows/<dag_name>_dag/callables/` — **required workflow-local glue**, not
   optional. Every `*TaskGroup` used here takes one or more `Callable` constructor args with **no
   default** (e.g. `InputPreparationTaskGroup(prepare_input_callable=...)`,
   `CosmosTaskGroup(config_generation_callable=..., output_validation_callable=...)`,
   `ValidatedOutputTaskGroup(validation_callable=...)`). These callables are genuinely new,
   per-DAG code — they are what turns generic container-arg building into *this* pipeline's actual
   input listing / config generation / output validation. Do not skip writing them, and do not
   treat "don't invent task logic" as license to invent them freely either: for each task group
   you use, find the closest existing DAG that uses the *same* task group, read its equivalent
   callable, and adapt it — trim the parts that don't apply (e.g. dropping an
   `output_group_id`/downstream-XCom concern when there is no next stage) rather than writing from
   a blank page.

**Kubernetes, and OSMO for workflows that support it.** This repository checks in K8s manifests for
every workflow, and additionally checks in OSMO manifests for Event Video Generation and Video
Data Augmentation (`event_video_generation_dag_osmo`, `video_data_augmentation_dag_osmo`) — see the
`cdf-event-video-generation` skill for the pattern. OSMO DAGs are **external-services-only**:
their manifests define OSMO batch tasks but no in-cluster VLM/LLM/image-edit/evaluator service
components, so an OSMO variant must reject `external_services: false` payloads and cannot use
`service_lifecycle` to deploy anything. A user asking for an OSMO variant of a new custom DAG is
in scope — replicate the external-services-only contract from an existing OSMO manifest/payload
rather than refusing. NVCF is still unsupported — do not offer it unless the user explicitly says
their environment supports it, and if so, treat it as new scope requiring its own investigation
rather than a flag on this skill.

**Not scoped to one workflow.** Task groups are shared across the shipped DAGs; a custom DAG can
mix stages from any of them (e.g. `CosmosTaskGroup` + `CosmosEvaluatorTaskGroup` with no
auto-labeling stages, or `ImageAttributeAugmentationTaskGroup` with no `CosmosTaskGroup` for a
labeling-only pass over already-augmented crops).

**These DAGs are not checked in as permanent artifacts.** Unlike the five workflow DAGs that ship
with this repository (listed in [Step 0](#step-0-confirm-no-existing-dag-already-covers-this)), a
DAG this skill generates is a one-off, user-requested pipeline — it is not added to any workflow
skill's DAG-ID table. Because nothing else will ever guide the user through deploying
or running it, **this skill owns the full lifecycle, not just authoring** — Step 8 below is not
optional follow-through, it's the only place readiness-check/deploy/trigger/monitor for this DAG
will ever happen.

---

## Making Use of Component Skills

The shared task groups are thin Airflow wrappers around CDF components. Each component owns its
config schema, model defaults, prompts, and errors in its own skill; the DAG's `container_args`
and generated config files are the contract that skill documents. This skill does not reimplement
that contract.

**Open the matching component skill before writing manifest task entries, config templates,
prompts, or question banks** — never invent config schema, model defaults, or prompt wording from
memory. [`references/component-skills-excerpt.md`](references/component-skills-excerpt.md) links
every component skill and keeps an offline summary of each. If the component repo is unreachable,
use that summary; if what you need isn't in it, say so and ask rather than extrapolating.

Use this table to pick the component and to see which shipped DAG already runs that stage — the
closest shipped DAG is also the one to copy callables and manifest entries from:

| Task group(s) or stage | Component | Component skill | Shipped DAGs that already use it |
|---|---|---|---|
| `CosmosTaskGroup` | CDF Augmentation | `cdf-augmentation` | IAA (image edit), DIG (image edit), EVG (image-to-video), VDA (Cosmos Transfer); AV subclasses it as workflow-local `AVCosmosTaskGroup` (Cosmos3-Nano with RGB or WSM controls) |
| `CosmosEvaluatorTaskGroup` | CDF Augmentation (Cosmos Evaluator) | `cdf-augmentation` | IAA, EVG, VDA; AV subclasses it as workflow-local `AVEvaluatorTaskGroup` (physical plausibility) |
| `DetectionAndTrackingTaskGroup`, `CaptioningTaskGroup`, `VisualQATaskGroup` | CDF Auto-Labeling | `cdf-auto-labeling` | EVG, VDA |
| `ImageAttributeAugmentationTaskGroup`, `EventAndPersonAttributeSearchTaskGroup` | CDF Auto-Labeling | `cdf-auto-labeling` | IAA, EVG |
| `SuperResolutionTaskGroup`, `ReasoningTaskGroup`, `TrainingExportTaskGroup` | CDF Auto-Labeling | `cdf-auto-labeling` | VDA |
| Simulation render (manifest task `image_generation`) | CDF Simulation | `cdf-simulation` | DIG |
| AnomalyGen defect generation (manifest task `anomaly_generation`) | CDF AnomalyGen | `cdf-anomalygen` | DIG |
| Curation and retrieval (manifest tasks `curator_model_download`, `curator`, `media_probe`, `cds_ingest`; endpoints `milvus_service`, `cds_service`) | CDF Curation and Retrieval | `cdf-curation-and-retrieval` | AV |

For `CosmosTaskGroup`, the input → output shape decides the model and adapter, not the task-group
name: Cosmos Transfer for video scene changes, Cosmos Predict for new or continued video, image
edit for still images, and image-to-video to animate a first frame.

Simulation, AnomalyGen, and curation and retrieval have no class in `shared/task_groups/`. Each is
a manifest task (or a few) wired directly in the DIG or AV DAG with `self.builder.build_task` /
`partial_task`, plus workflow-local callables that prepare its inputs and validate its outputs.
Reuse them the same way as any workflow-local stage (see **Workflow-local task groups** below):
copy the manifest task entries, the wiring, and the callables from that DAG, and adapt the
callables to your payload model.

- **Simulation and AnomalyGen** — `defect_image_generation_smt_dag.py` with its manifest
  `configs/smt/defect_image_generation_k8s_manifest.yaml`. The `generate_configs` task
  (`generate_smt_configs`) builds one simulation config per `golden_generation` variant from
  `configs/smt/image_generation_config.yaml`, and one AnomalyGen config per `defect_generation`
  variant from `configs/smt/anomaly_gen_config.yaml`. Simulation (`image_generation`) is mapped over the
  `image_generation_environments` list, selected by `select_generated_config_list`. AnomalyGen
  (`anomaly_generation`) is mapped over the `anomaly_package` tasks (`package_anomalygen_inputs`),
  and the `validate_anomalygen_outputs` task checks its results. Callables live in
  `callables/smt/`, and the payload model is `SmtDefectImageGenerationPayload` in
  `callables/smt/models.py`. DIG's image-edit stage between them is a `CosmosTaskGroup` that, in
  internal mode, also needs the `image_edit_checkpoint_download` task.
- **Curation and retrieval** — `autonomous_vehicle_dataset_generation_dag.py` with
  `configs/autonomous_vehicle_dataset_generation_k8s_manifest.yaml` and
  `configs/cosmos_curator_config.yaml`. The chain is `generate_configs` (`create_avdg_configs`) and
  `input_preparation` (`prepare_avdg_input`) → `build_curator_args` → `curator_model_download`
  (task `download_curator_model`) → `curator` (task `run_curator`) → `prepare_curator_probe` →
  `media_probe` (task `probe_curator_outputs`) → `reconcile_curator_outputs` → `cds_ingest` (task
  `ingest_cds`, after the `milvus_service` and `cds_service` endpoints are ready) →
  `collect_results` → `validate_retrieval`. AV reuses the same `media_probe` task to check WSM
  controls and generated videos. The pods run scripts from the DAG's `callables/` that
  `config_generation.py` base64-encodes into the generated configs (`cds_ingest.py`,
  `msc_credentials.py`, `media_probe.py`, `control_video_probe.py`), and the DAG file injects
  `cds_vector_pipeline.py` and `cds_vector_components.py` into the `cds_service` deployment. Copy
  all of them. These scripts use boto3, so inputs and outputs must be on S3 or an S3-compatible
  store.

DIG and AV also run subsets of their chains through payload switches, and those subsets count as
payload differences in Step 0:

From NVIDIA/skills, CC-BY-4.0 AND Apache-2.0.

More from NVIDIA/skills