Skip to content

feat(evaluator): Add dedicated Gym task image for agent evaluation - #1366

Closed
JashG wants to merge 11 commits into
jgulabrai/AALGO-494-gym-integration-testsfrom
AALGO-494-gym-task-image/jgulabrai
Closed

feat(evaluator): Add dedicated Gym task image for agent evaluation#1366
JashG wants to merge 11 commits into
jgulabrai/AALGO-494-gym-integration-testsfrom
AALGO-494-gym-task-image/jgulabrai

Conversation

@JashG

@JashG JashG commented Aug 18, 2026

Copy link
Copy Markdown
Contributor

Summary

NOTE: Target branch is set to #1315 as this is a follow-up to that PR.

Gym agent evaluations currently depend on a CI-specific variant of the shared CPU task image. This PR replaces it with a dedicated, production-shaped nmp-gym-tasks image that isolates Gym and Ray dependencies while remaining CI-only (i.e. only built and used for e2e tests in CI).

Changes

Gym task image

  • Add a dedicated nmp-gym-tasks image based on nmp-cpu-tasks that installs Gym and Ray into an isolated /opt/gym-venv.
  • Add a committed image-specific dependency lock for reproducible amd64 and arm64 builds.
  • Keep test and model configuration outside the image and supply it per job.

Evaluator routing

  • Updated agent routing: GymRunnerTarget evaluations rout to nmp-gym-tasks, all others route to nmp-cpu-tasks.
  • Add NEMO_EVALUATOR_GYM_TASKS_IMAGE for overriding the fully qualified Gym image.

CI and validation

  • Build and publish the SHA-tagged Gym image through the CPU CI workflow.
  • Pre-pull the Gym image into the dedicated Kind E2E cluster.
  • Add a cache-only smoke target that verifies the Gym CLI and imports Gym, Ray, and tiktoken.
  • Include the CI-only GHCR package in stale-image cleanup.
  • Remove the previous nmp-cpu-tasks-gym-e2e image and tag-alias workflow.

Dependency maintenance

  • Document how to update Gym dependencies, regenerate the lock, validate both architectures, and run image smoke and E2E tests.
  • Keep the image excluded from release targets until the Ray security blocker is resolved or approved.

Common Patterns with Customizer Images

Customizer also generates training images with Gym as a dependency. I identified a few common patterns that made sense to follow in this PR:

  • Use a dedicated task image rather than adding specialized dependencies to a shared platform image.
  • Separate the publishable runtime stage from a cache-only smoke-test stage.
  • Build both runtime and smoke targets through docker-bake.hcl.
  • Cache image layers independently and validate the actual runtime environment during the build.
  • Support a fully qualified image override for internal testing and deployment.
  • Keep dependency ownership explicit. Evaluator owns a small image-specific lock because the root workspace excludes Ray, while Customizer images consume their applicable workspace or upstream dependency authorities.

Type of Change

  • Code change (feature, bug fix, or refactor)
  • Code change with documentation updates
  • Documentation only
  • Contributor tooling or automation
  • CI, build, or test infrastructure

Quality Gates

  • Tests added or updated for changed behavior
  • Existing tests cover changed behavior — justification:
  • Tests not applicable — justification:
  • Documentation updated for user-visible behavior
  • Documentation not applicable — justification:

Verification

  • Pull request title follows the repository's Conventional Commit format
  • Every commit includes an appropriate Signed-off-by: trailer
  • uv run pre-commit run -a passes, or any blocked checks are identified below
  • Targeted tests pass, or tests are marked not applicable above
  • No secrets, API keys, or credentials are included

Targeted validation:

  • All tests pass
  • flox -q activate -- actionlint .github/workflows/ci.yaml — passed.
  • docker buildx bake nmp-gym-tasks-smoke-test — passed; imported nemo-gym 0.5.0, ray 2.57.0, and tiktoken 0.13.0.
  • Gym lock checks and dry-run synchronization passed for Linux amd64 and arm64.
  • Full pre-commit was not run; targeted lint, type, workflow, copyright, and whitespace checks passed.

JashG added 11 commits August 17, 2026 13:01
Signed-off-by: Jash Gulabrai <jgulabrai@nvidia.com>
The rebase onto main's promptless-row handling (which independently
touched discover_gym_tasks) surfaced tests still asserting gym_row
lives under inputs. It moved to metadata earlier in this branch, since
the wire-facing TaskInputs DTO only recognizes `instruction` and has
nowhere else for an arbitrary dataset row to travel through a
submitted job spec.

Signed-off-by: Jash Gulabrai <jgulabrai@nvidia.com>
Signed-off-by: Jash Gulabrai <jgulabrai@nvidia.com>
Signed-off-by: Jash Gulabrai <jgulabrai@nvidia.com>
Signed-off-by: Jash Gulabrai <jgulabrai@nvidia.com>
Signed-off-by: Jash Gulabrai <jgulabrai@nvidia.com>
Signed-off-by: Jash Gulabrai <jgulabrai@nvidia.com>
Signed-off-by: Jash Gulabrai <jgulabrai@nvidia.com>
Signed-off-by: Jash Gulabrai <jgulabrai@nvidia.com>
Signed-off-by: Jash Gulabrai <jgulabrai@nvidia.com>
Signed-off-by: Jash Gulabrai <jgulabrai@nvidia.com>
@github-actions github-actions Bot added the feat label Aug 18, 2026
@JashG
JashG force-pushed the jgulabrai/AALGO-494-gym-integration-tests branch from 4a722ca to 7d26389 Compare August 18, 2026 17:51
@JashG JashG closed this Aug 19, 2026
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant