---
name: paidf-orchestration-write-dag
description: Use when a user describes a custom PAIDF 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 PAIDF `launchable.md`).
version: "1.0.0"
license: CC-BY-4.0 AND Apache-2.0
metadata:
  owner: NVIDIA
  service: physical-ai-data-factory
  version: 1.0.0
  reviewed: '2026-09-02'
  author: NVIDIA
  tags:
    - physical-ai
    - paidf-orchestration
    - airflow
    - dag
---

# Write DAG

## What This Skill Does

Composes a new Kubernetes-only Airflow DAG from the **existing, shared task groups** in
`airflow/dags/shared/task_groups/` — the same building blocks `image_attribute_augmentation_dag.py`
and `event_video_generation_dag.py` are built from. 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 only.** This repository's `main`/current branches check in K8s manifests only (see
`image-attribute-augmentation-workflow` and `event-video-generation-workflow` skills) — no NVCF, no
OSMO. Do not offer those backends unless the user explicitly says their environment supports them,
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 IAA and EVG today; a custom DAG can
mix stages from either lineage (e.g. `CosmosTaskGroup` + `AutoLabelingTaskGroup` with no
`DetectionAndTrackingTaskGroup`, or `ImageAttributeAugmentationTaskGroup` with no `CosmosTaskGroup`
for a labeling-only pass over already-augmented crops).

**These DAGs are not checked in as permanent artifacts.** Unlike `image_attribute_augmentation_dag`
or `event_video_generation_dag`, a DAG this skill generates is a one-off, user-requested pipeline —
it is not added to `image-attribute-augmentation-workflow`, `event-video-generation-workflow`, or
any other run-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 task groups here are thin Airflow wrappers around two component pipelines that have their own
skills, external to this repo — **`augmentation`** (image-edit / Cosmos Transfer / Cosmos Predict /
image2video container) and **`auto-labeling`** (detection, captioning, visual QA, person attribute
search container). The DAG's `container_args` and generated config files are literally the CLI
contract those skills document — this skill does not reimplement that contract, it points at it.

**Those two skills are only reachable when this session's workspace happens to include their
repos** — there's no submodule, plugin registration, or other durable link from `sdg-workflow` to
them. Try `Skill(skill="augmentation")` / `Skill(skill="auto-labeling")` first; if the skill isn't
found, fall back to
[`references/component-skills-excerpt.md`](references/component-skills-excerpt.md) — a small,
frozen excerpt of the schema/prompt/troubleshooting content this skill has actually needed so far.
Treat it as a stale snapshot, not a substitute for the real skill: if what you need isn't in it, say
so and ask rather than extrapolating.

**Always invoke 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:

| Writing... | Invoke skill | Read |
|---|---|---|
| `CosmosTaskGroup` config — image-edit (IAA-style) | `augmentation` | `references/configuration-schema.md` (endpoints list, adapters), `references/image-attribute-augmentation.md` (image-edit prompts, MCQ `exclude_variables`, distribution configs) |
| `CosmosTaskGroup` config — image2video / Cosmos Transfer/Predict (EVG-style) | `augmentation` | `references/event-video-gen.md`, `references/config-decision-tree.md` (which model/adapter for the input→output shape you actually have) |
| `ImageAttributeAugmentationTaskGroup` / `EventAndPersonAttributeSearchTaskGroup` prompts, question banks | `auto-labeling` | `references/prompt-authoring.md`, `references/event-and-person-attribute-search.md`, `references/stages/person-attribute-search.md` |
| `DetectionAndTrackingTaskGroup` / `CaptioningTaskGroup` / `VisualQATaskGroup` / `AutoLabelingTaskGroup` config | `auto-labeling` | the matching file under `references/stages/` (`detection-and-tracking.md`, `captioning.md`, `visual-qa.md`) |
| Config validation or runtime errors from either container | matching skill | its `references/troubleshooting.md` |

`config-decision-tree.md` is the one to open first whenever the task group is doing something other than
a straight image-edit (IAA's case) — it's what tells you Cosmos Transfer vs. Predict vs. image2video,
not the task-group name.

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`. Always required. | 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.auto_labeling` (or the manifest's detection task key — check the base manifest) |
| `CaptioningTaskGroup` | `captioning` | VLM captioning on detection/tracking output. | `components.tasks.auto_labeling` |
| `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` |
| `AutoLabelingTaskGroup` | `auto_labeling` | Generic mapped auto-labeling container execution (detection+tracking as one component) — used when a workflow doesn't need the split `DetectionAndTrackingTaskGroup`/`CaptioningTaskGroup` stages. | `components.tasks.auto_labeling` |
| `ValidatedOutputTaskGroup` | `validated_output` | Runs a workflow-specific output validation callable. Always last processing step before reporting. | 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 whichever of `image_attribute_augmentation_dag.py` or
`event_video_generation_dag.py` already uses that group — 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/` directory, not just `shared/task_groups/`, for stage-specific groups
like this, and decide inclusion the same way: gate on manifest task presence, don't hardcode it in.

**Constructor kwarg defaults are not guaranteed to match your manifest.** e.g.
`CosmosTaskGroup`'s `internal_augmentation_pool` default is `internal_image_edit_service_pool`,
but IAA's manifest profile is actually named `iaa_internal_image_edit_service_pool` — the reference
DAG overrides the kwarg explicitly. Always diff every pool/profile-name-shaped default against the
manifest you wrote in Step 3, not just against the task group's source.

### 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`.

---

## 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: `image_attribute_augmentation_dag_k8s`
already runs `input_preparation → cosmos_augmentation → cosmos_post_processing →
event_and_person_attribute_search → generate_augmented_dataset → validated_output`, and
`event_video_generation_dag_k8s` already runs its own full chain (see that DAG's file for the exact
task groups). If the requested pipeline's stage sequence — after mapping it to task groups per the
table below — **matches one of those chains exactly**, this is not a `write-dag` case, no matter
what triggered the invocation:

- Different `input_path`/`output_directory`/service URLs/model overrides/`num_augmentation`/
  `variable_distribution`/`max_imgs` are **payload differences**. Both checked-in DAGs already
  parameterize all of that through their existing payload schema — none of it justifies a new DAG.
- Only a **structural** difference — a task group added, removed, or reordered relative to the
  existing chain — is a genuine `write-dag` case.

If the check matches an existing DAG: stop, do not write any files, and tell the user to use
`image-attribute-augmentation-workflow` or `event-video-generation-workflow` (whichever chain
matched) instead — with the payload differences they described, since that skill's payload already
covers them. If it's a genuine structural difference, proceed to Step 1.

---

## Step 1: Gather Requirements

Ask all of the following together, not one at a time:

- **Pipeline** — the ordered sequence of stages, described in the user's own words. Map each
  stage to a task group from the table above. If ambiguous (e.g. "augmentation and labeling" —
  does that include a verification/QA pass?), ask before writing anything.
- **DAG name** — snake_case identifier.
- **Input path, output directory, max items to process.**
- **Service mode** — `external_services: true` (user supplies endpoint URLs) or `false` (DAG
  deploys VLM/LLM/image-edit pods in-cluster). If external, gather every service URL the selected
  task groups need.
- **Container images / models** — for each processing task group, confirm the container image and
  served model against the **component skill's current default** (`augmentation` or
  `auto-labeling`, per the table above), not a hardcoded value. Ask the user to confirm or override.
- **Prompts / question banks / MCQ variables** — if the pipeline needs domain-specific prompts
  (edit prompt template, verification questions, person-attribute question bank), author them with
  the matching component skill's `prompt-authoring.md`/`image-attribute-augmentation.md` guidance,
  not from scratch.

## Step 2: Read the Base Manifest and Component Skill References

Read the closest existing K8s manifest in full before writing anything —
`airflow/dags/workflows/image_attribute_augmentation_dag/configs/image_attribute_augmentation_k8s_manifest.yaml`
if the pipeline includes `CosmosTaskGroup`/`ImageAttributeAugmentationTaskGroup`, or
`airflow/dags/workflows/event_video_generation_dag/configs/event_video_generation_k8s_manifest.yaml`
if it includes detection/tracking/captioning/visual-QA. These are the source of truth for
`deployment.profiles`, pool names, `container_args`, secrets, and GPU configuration — copy fields
verbatim, do not invent new profile shapes.

In parallel, invoke the component skill(s) selected in Step 1 and read the referenced files. This
is where model defaults, endpoint `adapter` values (`nim`, `openai.chat.completions`,
`openai.images.edits`), and prompt/config schema come from — the manifest's `container_args` and
the DAG's generated config file must agree with what that skill documents.

## Step 3: Write the K8s Manifest

**Path:** `airflow/dags/workflows/<dag_name>_dag/configs/<dag_name>_k8s_manifest.yaml`

Copy the base manifest chosen in Step 2 verbatim, then:

1. **Trim `components.tasks`** to only the entries the selected task groups need (see the table's
   "Manifest component(s) needed" column).
2. **Trim `components.endpoints`** to only what the remaining tasks need.
3. **Apply container image / model overrides** confirmed in Step 1 — cross-check the resulting
   `container_args` (served model name, required flags such as `--omni` for `vllm-omni serve
   Qwen/Qwen-Image-Edit-2511`, endpoint adapter matching the exposed API path) against the
   component skill reference and against any external spec the user cited (see **Parity Check**
   below).
4. Do not change profiles, pools, secrets, or connection settings.

## Step 4: Write the DAG Python File

**Path:** `airflow/dags/workflows/<dag_name>_dag/<dag_name>_dag.py`

Follow the structure of whichever reference DAG matches the pipeline's task groups —
`image_attribute_augmentation_dag.py` for an image-edit/IAA-style chain,
`event_video_generation_dag.py` for a detection/captioning/visual-QA/EVG-style chain, or the closer
match by shared task groups when the pipeline mixes both lineages. Either way, it's the current,
working pattern — do not resurrect an older webserver/NVCF-era template:

- A `<DagName>DAGBuilder` class wrapping a `ComponentBuilder(manifest_path=...)`.
- `build_dag()` instantiates `ValidatePayloadTaskGroup`, `InputPreparationTaskGroup`,
  `ServiceLifecycleTaskGroup` (with a `ServiceLifecycleSpec` tuple matching only the endpoints this
  DAG uses), then the chosen processing task groups in order, then `ValidatedOutputTaskGroup` and
  optionally `PerformanceReportingTaskGroup`.
- Wire `fail_pipeline`/`pipeline_success`/shutdown exactly as in the reference DAG — reporting is a
  terminal branch off `validated_output`, never upstream of the outcome tasks.
- Module-level registration uses the manifest-existence guard:

```python
_DAG_NAME_K8S_MANIFEST = os.environ.get(
    "<DAG_NAME_UPPER>_K8S_MANIFEST_PATH",
    <DAG_NAME_CONFIG_DIR> / "<dag_name>_k8s_manifest.yaml",
)
<dag_name>_dag_k8s = _create_<dag_name>_dag_if_manifest_exists(
    manifest_path=_DAG_NAME_K8S_MANIFEST,
    platform="k8s",
    description="...",
    tags=["sdg", "kubernetes", "<dag_name>"],
)
```

A missing manifest means the DAG is silently absent from Airflow, not broken — this is the only
"registration" mechanism in the current codebase. There is no webserver database, no DAG-seeding
step, and no NVCF/OSMO branch to add.

**Path:** `airflow/dags/workflows/<dag_name>_dag/callables/`

For every task group in the chosen chain, write the required callable(s) by adapting the closest
existing DAG's equivalent (see **Workflow-local task groups** above for the sibling `tasks/`
directory used for non-shared groups like `CosmosPostProcessingTaskGroup`). This is real, new code
— it is expected to differ per DAG — but its *shape* (signature, XCom pull/push pattern, error
handling via `AirflowFailException`) should mirror the reference implementation, not be designed
from scratch.

**Watch for payload-schema coupling when reusing a `tasks/` callable across DAGs.** IAA's
`generate_image_attribute_augmentation_performance_report` re-validates the templated payload
against `ImageAttributeAugmentationDagPayloadConfig` by name inside its own `_parse_payload` — it is
not schema-agnostic despite living next to schema-agnostic reporting primitives. Reusing it
unmodified from a different DAG's payload model re-validates against the *wrong* schema; it may not
even fail loudly, since a before-validator on the foreign model can silently backfill fields your
payload never had. Any `tasks/` callable that imports a specific `*DagPayloadConfig` by name needs
the same treatment as th