fix(math_with_judge): restore trajectory collection config cited by README

README.md

@@ -326,6 +326,7 @@ The Dataset column links to publicly available datasets (e.g., on HuggingFace).
 | Math With Autograder                          | math                  | Math QA verified by a Skills-style autograder LLM judge with math-verify symbolic fallback                                                                                                                                                                                                                                                                                                                                      | Score hard-math benchmarks (e.g. IMO AnswerBench) where the judge is a unidirectional Correct/Incorrect grader                        | -     | -          | -                                                         | <a href='resources_servers/math_with_autograder/configs/math_with_autograder.yaml'>math_with_autograder.yaml</a>                                                                                                                    | -                                                                                                                                                              |
 | Math With Code                                | math                  | Model solves competitive math problems using simple calculator tools                                                                                                                                                                                                                                                                                                                                                            | Improve math and simple tool use capabilities                                                                                         | ✓     | -          | Apache 2.0                                                | <a href='resources_servers/math_with_code/configs/math_with_code.yaml'>math_with_code.yaml</a>                                                                                                                                      | -                                                                                                                                                              |
 | Math With Judge                               | math                  | DAPO17k math dataset with math-verify                                                                                                                                                                                                                                                                                                                                                                                           | Improve math capabilities including AIME 24 / 25                                                                                      | ✓     | ✓          | Apache 2.0                                                | <a href='resources_servers/math_with_judge/configs/dapo17k.yaml'>dapo17k.yaml</a>                                                                                                                                                   | -                                                                                                                                                              |
+| Math With Judge                               | math                  | DAPO17k train problems repeated for trajectory collection with math-verify                                                                                                                                                                                                                                                                                                                                                      | Collect rollouts on DAPO17k without training (NeMo RL is_trajectory_collection)                                                       | -     | ✓          | -                                                         | <a href='resources_servers/math_with_judge/configs/dapo17k_trajectory_collection.yaml'>dapo17k_trajectory_collection.yaml</a>                                                                                                       | -                                                                                                                                                              |
 | Math With Judge                               | math                  | MathStackOverflow math dataset with math-verify                                                                                                                                                                                                                                                                                                                                                                                 | Improve math capabilities including AIME 24 / 25                                                                                      | ✓     | ✓          | Creative Commons Attribution-ShareAlike 4.0 International | <a href='resources_servers/math_with_judge/configs/math_stack_overflow.yaml'>math_stack_overflow.yaml</a>                                                                                                                           | <a href='https://huggingface.co/datasets/nvidia/Nemotron-RL-math-stack_overflow'>Nemotron-RL-math-stack_overflow</a>                                           |
 | Math With Judge                               | math                  | OpenMathReasoning math dataset with math-verify and LLM-as-a-judge                                                                                                                                                                                                                                                                                                                                                              | Improve math capabilities including AIME 24 / 25                                                                                      | ✓     | ✓          | Creative Commons Attribution 4.0 International            | <a href='resources_servers/math_with_judge/configs/math_with_judge.yaml'>math_with_judge.yaml</a>                                                                                                                                   | <a href='https://huggingface.co/datasets/nvidia/Nemotron-RL-math-OpenMathReasoning'>Nemotron-RL-math-OpenMathReasoning</a>                                     |
 | Mcqa                                          | knowledge             | Multi-choice question answering problems                                                                                                                                                                                                                                                                                                                                                                                        | Improve benchmarks like MMLU / GPQA / HLE                                                                                             | ✓     | ✓          | Apache 2.0                                                | <a href='resources_servers/mcqa/configs/mcqa.yaml'>mcqa.yaml</a>                                                                                                                                                                    | <a href='https://huggingface.co/datasets/nvidia/Nemotron-RL-knowledge-mcqa'>Nemotron-RL-knowledge-mcqa</a>                                                     |

nemo_gym/gitlab_utils.py

@@ -108,7 +108,10 @@ def download_jsonl_dataset(
         download_link,
         headers={"Authorization": f"Bearer {MLFLOW_TRACKING_TOKEN.get()}"},
     )
-    with open(config.output_fpath, "w") as f:
+    # Relative outputs such as `data/x.jsonl` only exist in the Gym root; create the directory like the HF download.
+    output_fpath = Path(config.output_fpath)
+    output_fpath.parent.mkdir(parents=True, exist_ok=True)
+    with open(output_fpath, "w") as f:
         f.write(response.content.decode())
 
 

resources_servers/math_with_judge/README.md

@@ -17,6 +17,10 @@ gym env start \
     +math_with_judge.resources_servers.math_with_judge.judge_model_server.name=policy_model
 ```
 
+The OpenMathReasoning download below and `--download` under "Prepare for trajectory collection" fetch from the
+GitLab model registry. They need the `mlflow` extra (`uv sync --extra mlflow`) and `mlflow_tracking_uri` and
+`mlflow_tracking_token` in `env.yaml`.
+
 To download the OpenMathReasoning dataset, the following command can be run:
 ```bash
 gym dataset download --storage gitlab \
@@ -36,6 +40,9 @@ gym eval run --no-serve \
 ```
 
 ## Prepare for trajectory collection
+This writes `validation.jsonl` with each DAPO17k train problem repeated 16 times. NeMo RL reads only the
+validation split when `is_trajectory_collection: true`. `--download` fetches `dapo17k` 0.0.1 from the GitLab
+model registry.
 ```bash
 gym dataset collate \
     --config resources_servers/math_with_judge/configs/dapo17k_trajectory_collection.yaml \

resources_servers/math_with_judge/configs/dapo17k_trajectory_collection.yaml

@@ -0,0 +1,48 @@
+math_with_judge:
+  resources_servers:
+    math_with_judge:
+      entrypoint: app.py
+      # Grading uses math-verify only (should_use_judge: false), so the judge reference just needs to resolve.
+      judge_model_server:
+        type: responses_api_models
+        name: policy_model
+      judge_responses_create_params: {
+        input: []
+      }
+      should_use_judge: false
+      domain: math
+      verified: false
+      description: DAPO17k train problems repeated for trajectory collection with math-verify
+      value: Collect rollouts on DAPO17k without training (NeMo RL is_trajectory_collection)
+math_with_judge_simple_agent:
+  responses_api_agents:
+    simple_agent:
+      entrypoint: app.py
+      resources_server:
+        type: resources_servers
+        name: math_with_judge
+      model_server:
+        type: responses_api_models
+        name: policy_model
+      datasets:
+      # NeMo RL trajectory collection reads only the validation split, so the train problems go there.
+      # Own local file: collate checks each dataset against `<jsonl stem>_metrics.json`, and
+      # dapo17k_train_metrics.json describes dapo17k.yaml's train split.
+      - name: validation
+        type: validation
+        jsonl_fpath: resources_servers/math_with_judge/data/dapo17k_trajectory_collection_train.jsonl
+        # From YouJiacheng/DAPO-Math-17k-dedup
+        gitlab_identifier:
+          dataset_name: dapo17k
+          version: 0.0.1
+          artifact_fpath: dapo17k_train.jsonl
+        num_repeats: 16
+        license: Apache 2.0
+
+math_with_judge_environment_server:
+  environment_servers:
+    legacy_agent:
+      entrypoint: app.py
+      agent_server:
+        type: responses_api_agents
+        name: math_with_judge_simple_agent

resources_servers/math_with_judge/data/dapo17k_trajectory_collection_train_metrics.json

@@ -0,0 +1,56 @@
+{
+    "name": "validation",
+    "type": "validation",
+    "jsonl_fpath": "resources_servers/math_with_judge/data/dapo17k_trajectory_collection_train.jsonl",
+    "num_repeats": 16,
+    "source": {
+        "type": "gitlab",
+        "dataset_name": "dapo17k",
+        "version": "0.0.1",
+        "artifact_fpath": "dapo17k_train.jsonl"
+    },
+    "gitlab_identifier": {
+        "dataset_name": "dapo17k",
+        "version": "0.0.1",
+        "artifact_fpath": "dapo17k_train.jsonl"
+    },
+    "huggingface_identifier": null,
+    "license": "Apache 2.0",
+    "Number of examples": 278368,
+    "Number of tools": {
+        "Total # non-null values": 0,
+        "Average": 0.0,
+        "Min": 0.0,
+        "Max": 0.0,
+        "Standard deviation": 0.0
+    },
+    "Json-dumped number of words (proxy for token count)": {
+        "Total # non-null values": 278368,
+        "Average": 85.63,
+        "Min": 45.0,
+        "Max": 322.0,
+        "Standard deviation": 26.94
+    },
+    "Number of turns": {
+        "Total # non-null values": 278368,
+        "Average": 1.0,
+        "Min": 1.0,
+        "Max": 1.0,
+        "Standard deviation": 0.0
+    },
+    "Temperature": {
+        "Total # non-null values": 0,
+        "Average": 0.0,
+        "Min": 0.0,
+        "Max": 0.0,
+        "Standard deviation": 0.0
+    },
+    "question": {
+        "unique_count": 17398,
+        "total_count": 278368
+    },
+    "expected_answer": {
+        "unique_count": 2251,
+        "total_count": 278368
+    }
+}

tests/unit_tests/test_gitlab_utils.py

@@ -13,12 +13,14 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 import sys
+from pathlib import Path
+from types import ModuleType, SimpleNamespace
 
 import pytest
 from pytest import MonkeyPatch
 
-from nemo_gym.config_types import ConfigError, MLFlowConfig
-from nemo_gym.gitlab_utils import create_mlflow_client
+from nemo_gym.config_types import ConfigError, DownloadJsonlDatasetGitlabConfig, MLFlowConfig
+from nemo_gym.gitlab_utils import create_mlflow_client, download_jsonl_dataset
 
 
 # TODO: Eventually we want to add more tests to ensure that the Gitlab flow does not break
@@ -40,3 +42,39 @@ def test_create_mlflow_client_gives_an_actionable_error_when_mlflow_is_missing(
 
         with pytest.raises(ConfigError, match=r"pip install nemo-gym\[mlflow\]"):
             create_mlflow_client()
+
+    def test_download_creates_the_output_directory(self, monkeypatch: MonkeyPatch, tmp_path: Path) -> None:
+        # README commands download to a relative `data/...` path, which exists only in the Gym root.
+        # Fake the registry lookup and the HTTP fetch; mlflow itself is an optional extra.
+        client = SimpleNamespace(
+            tracking_uri="https://gitlab.invalid/api/v4/projects/1/ml/mlflow",
+            get_model_version=lambda name, version: SimpleNamespace(run_id="run"),
+        )
+        repository = SimpleNamespace(repo=SimpleNamespace(artifact_uri="https://gitlab.invalid/artifacts/"))
+        monkeypatch.setattr("nemo_gym.gitlab_utils.create_mlflow_client", lambda: client)
+        monkeypatch.setitem(sys.modules, "mlflow", ModuleType("mlflow"))
+        monkeypatch.setitem(
+            sys.modules, "mlflow.artifacts", SimpleNamespace(get_artifact_repository=lambda **_: repository)
+        )
+        monkeypatch.setitem(
+            sys.modules,
+            "mlflow.environment_variables",
+            SimpleNamespace(MLFLOW_TRACKING_TOKEN=SimpleNamespace(get=lambda: "token")),
+        )
+        requested = []
+
+        def fake_get(url: str, headers: dict) -> SimpleNamespace:
+            requested.append(url)
+            return SimpleNamespace(content=b'{"question": "q"}\n')
+
+        monkeypatch.setattr("nemo_gym.gitlab_utils.requests.get", fake_get)
+        output = tmp_path / "data" / "nested" / "rows.jsonl"
+
+        download_jsonl_dataset(
+            DownloadJsonlDatasetGitlabConfig(
+                dataset_name="dataset", version="0.0.1", artifact_fpath="rows.jsonl", output_fpath=str(output)
+            )
+        )
+
+        assert requested == ["https://gitlab.invalid/artifacts/rows.jsonl"]
+        assert output.read_text() == '{"question": "q"}\n'