---
name: cdf-orchestration-write-dag
description: Use when a user describes a custom CDF Orchestration pipeline — a specific ordered combination of stages such as augmentation only, auto-labeling only, detection+captioning only, or image attribute augmentation without full auto-labeling — that no existing DAG in airflow/dags/workflows/ covers, and asks for a new Kubernetes DAG. Also use to check that a generated or existing DAG's model/container/prompt choices match an external spec document (e.g. a CDF `launchable.md`). Not for requests that an existing DAG already covers, including payload-only changes such as different inputs, endpoints, models, or variant counts — route those to the matching workflow skill, unless the user still wants a separate DAG after being pointed to it.
version: "1.0.1"
license: CC-BY-4.0 AND Apache-2.0
metadata:
  owner: NVIDIA
  service: cosmos-data-factory
  version: 1.0.1
  reviewed: '2026-10-08'
  author: NVIDIA
  tags:
    - physical-ai
    - cdf-orchestration
    - airflow
    - dag
---

# 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:

- DIG `pipeline_mode`: `simulation` (simulation only), `sim2real` (simulation, then image edit),
  `full` (adds AnomalyGen, the default), `defect` (AnomalyGen only).
- AV `cosmos.enabled`: `false` runs curation and retrieval only; `true` adds Cosmos3-Nano
  generation and physical-plausibility evaluation.

If the user's request doesn't map onto a documented config/prompt pattern in the component skill,
say so and ask rather than inventing prompt text or config keys.

---

## Available Task Groups

Import each directly from its submodule (matches every existing DAG file — do not import from the
`task_groups` package `__init__`, which is an incomplete re-export and omits
`EventAndPersonAttributeSearchTaskGroup`).

| Task Group | Module | What It Does | Manifest component(s) needed |
|---|---|---|---|
| `ValidatePayloadTaskGroup` | `validate_payload` | Validates the trigger payload against the DAG's pydantic model. Always first. | none |
| `InputPreparationTaskGroup` | `input_preparation` | Lists/normalizes input media from the payload's `input_path`. Required whenever the pipeline reads input media; every shipped DAG uses it except DIG, which generates its inputs by simulation. | none |
| `ServiceLifecycleTaskGroup` | `service_lifecycle` | Deploys/tears down internal VLM/LLM/image-edit/image2video pods and waits for readiness; skipped per-service when `external_services: true`. | `components.endpoints` entries matching the `ServiceLifecycleSpec`s passed in |
| `CosmosTaskGroup` | `cosmos` | Runs the `augmentation` component (image-edit or Cosmos Transfer/Predict/image2video) via a manifest task, dynamically mapped over inputs. | `components.tasks.augmentation` (+ VLM/LLM/image-edit endpoints) |
| `ImageAttributeAugmentationTaskGroup` | `image_attribute_augmentation` | Runs person-attribute search/captioning over cosmos-augmented crops (IAA's auto-labeling stage). Reads a cosmos-output XCom. | `components.tasks.event_and_person_attribute_search` |
| `EventAndPersonAttributeSearchTaskGroup` | `event_and_person_attribute_search` | EVG's equivalent — person/event attribute search over auto-labeled video output. Reads an auto-labeling XCom. | `components.tasks.event_and_person_attribute_search` |
| `DetectionAndTrackingTaskGroup` | `detection_and_tracking` | RFDETR detection + BoxMOT tracking on video/cosmos output. | `components.tasks.detection_and_tracking` |
| `CaptioningTaskGroup` | `captioning` | VLM captioning on detection/tracking output. | `components.tasks.captioning` |
| `VisualQATaskGroup` | `visual_qa` | VLM visual QA; instantiate once per QA purpose with a distinct `group_id`/`component_name` (EVG uses it twice: anomaly QA and person-attribute QA). | `components.tasks.visual_qa` |
| `CosmosEvaluatorTaskGroup` | `cosmos_evaluator` | Deploys and tears down the Cosmos Evaluator (Postgres, database, arbitrator, and the configured checkers: hallucination, attribute verification, physical plausibility) around Cosmos generation. | `components.endpoints.cosmos_evaluator_*` entries for the core services and each selected checker |
| `SuperResolutionTaskGroup` | `super_resolution` | Super-resolves videos with a manifest-defined container task (VDA). | `components.tasks.super_resolution` |
| `ReasoningTaskGroup` | `reasoning` | Produces reasoning labels with a manifest-defined container task (VDA). | `components.tasks.reasoning` |
| `TrainingExportTaskGroup` | `training_export` | Aggregates completed scenes into one training-export container run (VDA). | `components.tasks.training_export` |
| `ValidatedOutputTaskGroup` | `validated_output` | Runs a workflow-specific output validation callable. Last processing step before reporting in every shipped DAG except DIG, which validates in its own `validate_anomalygen_outputs` task. | none |
| `PerformanceReportingTaskGroup` | `reporting` | Optional YAML+HTML performance report, gated by `enable_performance_reporting` in the payload. Hangs off `validated_output` as a terminal branch — never upstream of `fail_pipeline`/`pipeline_success`. | none |

**Before using any task group**, read its `__init__` in the source file — constructor kwargs
(`input_xcom_task_id`, `group_id`, `component_name`, `prepare_args_callable`,
`prepare_args_op_kwargs`) vary per group and determine both manifest wiring and XCom data flow.
Copy the calling pattern from a shipped DAG that already uses that group (see the "Shipped DAGs
that already use it" column above) — do not guess kwargs, and do not assume
every group exposes an `input_xcom_task_id` param: `ValidatedOutputTaskGroup`, for one, doesn't —
it takes a plain `op_kwargs` dict (default `{"payload": ..., "run_id": ...}`), so pulling an
upstream XCom into it means passing a custom `op_kwargs`/callable, not a task-id string.

**Workflow-local task groups exist too, outside `shared/task_groups/`.** IAA's
`CosmosPostProcessingTaskGroup` lives under
`airflow/dags/workflows/image_attribute_augmentation_dag/tasks/` and is wired in *conditionally* —
`image_attribute_augmentation_dag.py` only includes it when
`"cosmos_post_processing" in self.builder.manifest.deployment.components.tasks`. Check each
reference DAG's own `tasks/` and `callables/` directories, not just `shared/task_groups/`, for
stage-specific groups like this — AV's `AVCosmosTaskGroup` (`tasks/cosmos.py`) and
`AVEvaluatorTaskGroup` (`callables/evaluator_lifecycle.py`), and the DIG and AV stages that are
plain manifest tasks wired in the DAG file (see the component table above). Where a stage is
optional, decide inclusion the same way: gate on manifest task presence, don't hardcode it in.

**Constructor kwarg defaults are not guaranteed to match your pipeline.** e.g.
`CosmosTaskGroup`'s `external_augmentation_pool` defaults to `external_image_edit_service_pool`,
which is right for image-edit but wrong for image2video or Cosmos Transfer — EVG passes
`external_image2video_service_pool` and VDA passes `cosmos_transfer_3_service_pool` explicitly.
Internal-mode augmentation always runs on `default_pool`. Always diff every pool/profile-name-shaped
default against `deploy/values.yaml`'s `airflowPools.pools` and the manifest you wrote in Step 3.

### Data flow

`input_preparation` → `service_lifecycle` (parallel with input_preparation, gated by `services_ready`)
→ first processing task group (reads `input_preparation`'s XCom) → ... → last processing task group
→ `validated_output` → `[fail_pipeline, pipeline_success]` and `performance_reporting` (parallel
branches) → `service_lifecycle` shutdown (`ALL_DONE`).

Each processing task group after the first must read its *predecessor's* XCom (`input_xcom_task_id`
pointed at the previous group's group-qualified task id, where that param exists — see above), not
always `input_preparation` — e.g. in IAA, `event_and_person_attribute_search` reads
`cosmos_augmentation.validate_outputs`, not `input_preparation.prepare_input`.

DIG does not follow this shape: it starts from `generate_configs` instead of
`input_preparation`, routes stages with `pipeline_mode` branch and short-circuit tasks, and, in
`full` and `defect` modes, ends at `validate_anomalygen_outputs` instead of `validated_output`. A chain that starts with
simulation should copy DIG's wiring for that part.

---

## Step 0: Confirm No Existing DAG Already Covers This

Run this before Step 1, every time — including when this skill was invoked directly, since nothing
upstream of it is guaranteed to have made this check correctly. **Do not skip it because the
request already says "custom" or "new"; the user's framing is not evidence.**

The test is about the *task-group chain*, not the payload. These DAGs ship with this repository;
read the DAG file for the exact task-group order:

| Shipped DAG | What it runs | Workflow skill |
|---|---|---|
| `image_attribute_augmentation_dag_k8s` | `input_preparation → cosmos_augmentation` (image edit, with attribute-verification evaluation) `→ cosmos_post_processing` (when in the manifest) `→ event_and_person_attribute_search → generate_augmented_dataset → validated_output` | `cdf-image-attribute-augmentation` |
| `event_video_generation_dag_k8s` / `_osmo` | `input_preparation → cosmos_augmentation` (image-to-video, with Cosmos evaluation) `→ auto_labeling` (`detection_and_tracking → captioning → anomaly_visual_qa → person_attribute_visual_qa → person_attribute_search`) `→ generate_anomaly_dataset → validated_output` | `cdf-event-video-generation` |
| `video_data_augmentation_dag_k8s` / `_osmo` | `input_preparation`, then two branches into `validated_output`: auto-labeling of the source videos, and `cosmos_augmentation` (Cosmos Transfer, with Cosmos evaluation) followed by auto-labeling of the generated videos. Each auto-labeling branch runs `super_resolution → detection_and_tracking → captioning → visual_qa → reasoning → training_export` | `cdf-video-data-augmentation` |
| `autonomous_vehicle_dataset_generation_dag_k8s` | Curation and retrieval (`curator → media_probe → cds_ingest → collect_results → validate_retrieval`), then, when `cosmos.enabled`, Cosmos3-Nano generation with physical-plausibility evaluation, then `validated_output` | `cdf-autonomous-vehicle-dataset-generation` |
| `defect_image_generation_smt_dag_k8s` | `generate_configs → image_generation` (simulation) `→ augmentation` (image edit) `→ anomaly_package → anomaly_generation → validate_anomalygen_outputs`, with `pipeline_mode` selecting a subset | `cdf-defect-image-generation` |

If the requested pipeline's stage sequence — after mapping it to task groups and stages per the
Available Task Groups and component tables — **matches one of those chains exactly** (or the chain
of any other DAG under `airflow/dags/workflows/`), do not write a new DAG unless the user still
wants one after the recommendation below:

- Different input or outpu