feat(data): check that taskset task IDs are stable and unique

fern/versions/latest/pages/environment-tutorials/resources-server-apis.mdx

@@ -483,10 +483,17 @@ overrides, `num_repeats_add_seed`, and run-level `skills`. Set response paramete
 in the source row before collation; benchmark prompt configs still apply to input messages.
 
 The task's qualified ID is `<taskset>/<id>`. The ID is taken from `task_id`,
-`problem_id`, or `instance_id`, in that order. Without an explicit ID, an existing
+`problem_id`, or `instance_id`, in that order, and the field stays in `task_input`.
+Without an explicit ID, an existing
 `_ng_task_index` takes precedence; otherwise collation uses a zero-based source-row
 index across all datasets of the same type, before repeats. Adding or reordering
-other datasets in that collation can shift these positional IDs.
+other datasets in that collation can shift these positional IDs, so collation warns
+once for each dataset that uses them. Write a stable `task_id` in the preparation script instead.
+
+A task ID must name one task in its taskset. Collation reports IDs that appear on two
+source rows, including the same file declared twice under one taskset. Repeats of one
+row are not duplicates. The report fails collation when `task_data_validation` is
+`error`, the default for `gym dataset collate +mode=example_validation`, and is printed otherwise.
 
 These tasksets route to an Environment Server, not to an agent, so their datasets
 do not set `agent`. Benchmark discovery, `gym eval prepare`, and collation check the

nemo_gym/task_materialization.py

@@ -16,22 +16,30 @@
 )
 
 
+# Source fields that carry a task's own ID, in order of precedence.
+TASK_ID_FIELDS = ("task_id", "problem_id", "instance_id")
+
+
+def source_task_id(row: Mapping[str, JsonValue]) -> str | None:
+    """Return the ID a flat row carries in its own fields, or None when only its position identifies it."""
+    return next((str(row[key]) for key in TASK_ID_FIELDS if row.get(key) is not None), None)
+
+
 def materialize_task(
     row: Mapping[str, JsonValue], *, taskset: str, task_index: int | None = None
 ) -> dict[str, JsonValue]:
     """Preserve legacy task identity and task fields without embedding runtime routing.
 
-    Rows without an explicit ID use the collector's ``_ng_task_index``, or a source-row
-    index across all datasets of the same type in the collation, before repeats.
-    Adding or reordering other datasets shifts these positional IDs; editing prompt text does not.
+    The ID comes from the row's own fields (``TASK_ID_FIELDS``). Rows without one use the
+    collector's ``_ng_task_index``, or a source-row index across all datasets of the same type in
+    the collation, before repeats. These positional IDs are legacy: adding or reordering other
+    datasets shifts them, so collation warns when it uses them.
+    The source ID field stays in ``task_input``, because some verifiers read it as a task field.
     Prepared source files are never modified by this conversion.
     """
     if "task_input" in row or isinstance(row.get("task_id"), Mapping):
         raise ValueError("Expected a flat dataset row, not an already materialized task")
-    task_id = next(
-        (str(row[key]) for key in ("task_id", "problem_id", "instance_id") if row.get(key) is not None),
-        None,
-    )
+    task_id = source_task_id(row)
     if task_id is None:
         index = row.get(TASK_INDEX_KEY_NAME, task_index)
         if not isinstance(index, int) or isinstance(index, bool) or index < 0:

nemo_gym/train_data_utils.py

@@ -65,11 +65,12 @@
 from nemo_gym.prompt import apply_prompt_to_row, load_prompt_config, validate_prompt_compatibility
 from nemo_gym.task_data import (
     TaskDataSchemaError,
+    TaskDataValidationReport,
     TaskDataValidator,
     find_server_dir,
     load_task_data_schema,
 )
-from nemo_gym.task_materialization import materialize_task
+from nemo_gym.task_materialization import TASK_ID_FIELDS, materialize_task, source_task_id
 
 
 class TrainDataProcessorConfig(BaseNeMoGymCLIConfig):
@@ -1064,6 +1065,9 @@ def _collate_samples_single_type(
         paths_to_collate = []
         used_prepare_paths: set[Path] = set()
         source_task_index = -1
+        # The first source row of each (taskset, task ID), to report IDs that name two tasks.
+        task_id_locations: Dict[Tuple[str, str], str] = {}
+        duplicate_task_ids: List[str] = []
         for c in server_instance_configs:
             for d in c.datasets:
                 if d.type != type:
@@ -1100,6 +1104,7 @@ def _collate_samples_single_type(
                 # for this dataset comes from the declaration, and passing the stale field
                 # through would leak the old coupling into the clean format.
                 legacy_agent_ref_rows = 0
+                positional_task_id_rows = 0
                 validator = None
                 if task_data_validation != "off":
                     validator = self._task_data_validator_for(c, d, server_instance_configs)
@@ -1120,6 +1125,15 @@ def _collate_samples_single_type(
                         # Count each source row once, using its original JSONL line index.
                         if validator is not None and row_index % d.num_repeats == 0:
                             validator.validate_row(row_index // d.num_repeats, row)
+                        if d.taskset is not None and row_index % d.num_repeats == 0:
+                            task_id = source_task_id(row)
+                            if task_id is None:
+                                positional_task_id_rows += 1
+                            else:
+                                location = f"{c.name}/{d.name} line {row_index // d.num_repeats}"
+                                first = task_id_locations.setdefault((d.taskset, task_id), location)
+                                if first != location:
+                                    duplicate_task_ids.append(f"{d.taskset}/{task_id} ({first}; {location})")
                         if d.taskset is not None:
                             # Keep collector identity outside task_data so shards and retries
                             # retain their capture keys after materialization.
@@ -1143,6 +1157,14 @@ def _collate_samples_single_type(
                         raise ValueError(summary)
                     print(f"[task_data validation]\n{summary}")
 
+                if positional_task_id_rows:
+                    warnings.warn(
+                        f"{d.jsonl_fpath}: {positional_task_id_rows} rows of taskset {d.taskset!r} have no "
+                        f"{', '.join(TASK_ID_FIELDS)}, so their task IDs are positions in this collation and "
+                        "change when datasets are added or reordered. Write a stable task_id in the dataset's "
+                        "preparation script.",
+                        stacklevel=2,
+                    )
                 if legacy_agent_ref_rows:
                     warnings.warn(
                         f"{d.jsonl_fpath}: stripped legacy agent_ref from {legacy_agent_ref_rows} rows "
@@ -1152,6 +1174,17 @@ def _collate_samples_single_type(
                     )
                 paths_to_collate.append(prepare_path)
 
+        if duplicate_task_ids and task_data_validation != "off":
+            # A TaskId names one task in a run; results and reward profiles grouped by it would merge these.
+            summary = (
+                f"{len(duplicate_task_ids)} task IDs name more than one task in their taskset: "
+                + "; ".join(duplicate_task_ids[: TaskDataValidationReport.MAX_RECORDED_ERRORS])
+                + (" ..." if len(duplicate_task_ids) > TaskDataValidationReport.MAX_RECORDED_ERRORS else "")
+            )
+            if task_data_validation == "error":
+                raise ValueError(summary)
+            print(f"[task_data validation]\n{summary}")
+
         return paths_to_collate
 
     def collate_samples(

tests/unit_tests/test_task_materialization.py

@@ -316,6 +316,54 @@ def test_collation_rejects_misplaced_fields_before_materialization(tmp_path, tas
         TrainDataProcessor()._collate_samples_single_type("example", configs, task_data_validation="error")
 
 
+def _collate_taskset(tmp_path, rows, *, mode, datasets=1):
+    source = tmp_path / "source.jsonl"
+    source.write_text("".join(json.dumps(row) + "\n" for row in rows))
+    declarations = [
+        {"name": f"d{i}", "type": "example", "jsonl_fpath": str(source), "taskset": "t", "num_repeats": 2}
+        for i in range(datasets)
+    ]
+    configs = GlobalConfigDictParser().filter_for_server_instance_configs(
+        OmegaConf.create(_resources_server("resources", declarations, implementation="external_custom"))
+    )
+    return TrainDataProcessor()._collate_samples_single_type("example", configs, task_data_validation=mode)
+
+
+def test_collation_warns_once_per_dataset_about_positional_task_ids(tmp_path):
+    with pytest.warns(UserWarning, match=r"2 rows of taskset 't' have no task_id, problem_id, instance_id") as caught:
+        _collate_taskset(tmp_path, [{"q": "a"}, {"q": "b"}, {"task_id": "c", "q": "c"}], mode="error")
+    assert len([w for w in caught if "positions in this collation" in str(w.message)]) == 1
+
+
+def test_collation_with_stable_task_ids_does_not_warn(tmp_path, recwarn):
+    (path,) = _collate_taskset(tmp_path, [{"task_id": "a"}, {"problem_id": "b"}], mode="error")
+    assert not [w for w in recwarn if "positions in this collation" in str(w.message)]
+    # Repeats of one source row share its ID and are not duplicates.
+    assert [json.loads(line)["task_id"]["task_id"] for line in path.read_text().splitlines()] == ["a", "a", "b", "b"]
+
+
+@pytest.mark.parametrize(
+    "rows, datasets",
+    [
+        ([{"task_id": "a", "q": 1}, {"task_id": "a", "q": 2}], 1),
+        # problem_id and task_id resolve to the same TaskId.
+        ([{"task_id": "a"}, {"problem_id": "a"}], 1),
+        # The same file declared twice under one taskset repeats every task.
+        ([{"task_id": "a"}], 2),
+    ],
+)
+def test_collation_rejects_duplicate_task_ids_in_error_mode(tmp_path, rows, datasets):
+    with pytest.raises(ValueError, match=r"1 task IDs name more than one task in their taskset: t/a \(.*line 0;"):
+        _collate_taskset(tmp_path, rows, mode="error", datasets=datasets)
+
+
+def test_collation_reports_duplicate_task_ids_in_warn_mode(tmp_path, capsys):
+    (path,) = _collate_taskset(tmp_path, [{"task_id": "a", "q": 1}, {"task_id": "a", "q": 2}], mode="warn")
+    assert "t/a" in capsys.readouterr().out
+    assert len(path.read_text().splitlines()) == 4
+    _collate_taskset(tmp_path, [{"task_id": "a", "q": 1}, {"task_id": "a", "q": 2}], mode="off")
+
+
 class CustomTaskInput(BaseModel):
     model_config = ConfigDict(extra="forbid")
     scenario: str