## What does this PR do? This PR extends the experiment tracking possibilities by adding MLFlow backend. Till now there was a W&B exporter wired in the code, that tracks configuration, metrics and rollouts. This PR adds an exporter abstraction, moves existing functionality into W&B backend and adds MLFlow backend. Both can be configured at the same time - artifacts are sent to both tracking servers. **Deprecated field:** `upload_rollouts_to_wandb` has been renamed to `upload_rollouts`. If the user provides an old key, it is mapped to the new one and a deprecation warning is emitted. ## Followup 1. Export file artifacts. The original W&B code logs rollouts only, while NEL was logging more. This PR just adds a new backend and doesn't provide full feature parity with NEL. Tracked with #2619 2. Placement of export initialization seems sub-optimal. In my opinion it belongs to CLI, not config parsing. As a side effect we don't have the user-friendly error handling on config error in exporter. Tracked with #2620 3. Clean `RewardProfiler` up. It looks like it uses `wandb.Histrogram` for no good reason and carries some other obsolete code. Tracked with #2621 ## Usage example ``` gym eval run \ --benchmark gpqa \ --model-type openai_model \ --output gpqa-results/rollouts.jsonl \ --split benchmark \ +mlflow_tracking_uri=https://mlflow.my-server.com/ \ +mlflow_experiment_name=my-eval \ +mlflow_run_name=gpqa ``` This will run the evaluation as usual and create an MLFlow experiment for the run: ``` ... INFO: Shutting down INFO: Waiting for application shutdown. INFO: Application shutdown complete. INFO: Finished server process [2090575] NeMo Gym finished! Shutting down Ray cluster spun up by NeMo Gym... 🏃 View run gpqa at: https://mlflow.my-server.com/#/experiments/1740/runs/f088e196471048c9bf20e9beb9b771f5 🧪 View experiment at: https://mlflow.my-server.com/#/experiments/1740 ``` ## Checklist - [x] I have read the [contributing guidelines](https://docs.nvidia.com/nemo/gym/latest/contribute/development-setup). - [x] The change is focused; unrelated "drive-by" edits are tracked as separate issues/PRs. - [x] Tests added or updated and pass locally, or N/A for docs-only / non-code changes (so CI unit/server checks pass when applicable). - [x] Pre-commit checks pass locally (`pre-commit run --all-files`) (so CI lint/format/copyright pass). - [x] All commits have DCO sign-off (`git commit -s`) (so the DCO check passes). --------- Signed-off-by: Marta Stepniewska-Dziubinska <martas@nvidia.com> Co-authored-by: bxyu-nvidia <bxyu@nvidia.com>

Uh oh!

Extend exporters to push more artifacts #2619

Description

Metadata

Metadata

Assignees

Labels

Type

Projects

Milestone

Relationships

Development

Issue actions