Skip to content

feat(evaulator): Add dedicated Gym task image for agent evaluation - #1384

Merged
JashG merged 4 commits into
jgulabrai/AALGO-494-gym-integration-testsfrom
jgulabrai/AALGO-489-gym-base-image
Aug 19, 2026
Merged

feat(evaulator): Add dedicated Gym task image for agent evaluation#1384
JashG merged 4 commits into
jgulabrai/AALGO-494-gym-integration-testsfrom
jgulabrai/AALGO-489-gym-base-image

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, releasable nmp-gym-tasks image that isolates Gym and Ray dependencies from the shared nmp-cpu-tasks image.

Changes

Gym task image

  • Add a dedicated nmp-gym-tasks image based on nmp-cpu-tasks.
  • Install 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.
  • Keep Gym and Ray dependencies out of the shared CPU task image.

Evaluator routing

  • Route GymRunnerTarget evaluations to nmp-gym-tasks.
  • Continue routing all other agent evaluations to nmp-cpu-tasks.
  • Add NEMO_EVALUATOR_GYM_TASKS_IMAGE for overriding the fully qualified Gym image.

Image publication

  • Include nmp-gym-tasks in the standard CPU image build group.
  • Add the image to nightly and stable release publication.
  • Add NGC catalog metadata for the released image.
  • Continue publishing SHA-tagged images through the CPU CI workflow.
  • Include the GHCR package in stale-image cleanup.

CI and validation

  • 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.
  • Keep the smoke target in the CI-specific CPU build group while publishing the runtime image through the standard group.
  • Remove the previous nmp-cpu-tasks-gym-e2e image and tag-alias workflow.

Dependency maintenance

  • Document how to update Gym dependencies and regenerate the image-specific lock.
  • Validate dependency resolution for Linux amd64 and arm64.
  • Document image smoke-test and E2E commands.

Common Patterns with Customizer Images

Customizer also generates training images with Gym as a dependency. This PR follows the same relevant patterns:

  • Use a dedicated task image instead of adding specialized dependencies to a shared platform image.
  • Separate the publishable runtime stage from the cache-only smoke-test stage.
  • Build 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 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:

  • Local Gym agent-evaluation E2E passed against a locally built nmp-gym-tasks image: 1 passed in 177.20s.
  • docker buildx bake nmp-gym-tasks-smoke-test passed and 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.
  • docker buildx bake --print docker-cpu docker-cpu-ci confirmed the runtime and smoke targets are wired into the expected build groups.
  • NGC metadata tests passed: 14 passed.
  • NGC metadata dry-run discovered nmp-gym-tasks.

@github-actions github-actions Bot added the feat label Aug 18, 2026
Signed-off-by: Jash Gulabrai <jgulabrai@nvidia.com>
@JashG
JashG force-pushed the jgulabrai/AALGO-489-gym-base-image branch from 17e8691 to bf0e0bf Compare August 19, 2026 14:19
@JashG
JashG marked this pull request as ready for review August 19, 2026 14:19
@JashG
JashG requested review from a team as code owners August 19, 2026 14:19
JashG added 3 commits August 19, 2026 10:35
Signed-off-by: Jash Gulabrai <jgulabrai@nvidia.com>
Signed-off-by: Jash Gulabrai <jgulabrai@nvidia.com>
…pulls image

Signed-off-by: Jash Gulabrai <jgulabrai@nvidia.com>
@JashG
JashG merged commit eaa30a3 into jgulabrai/AALGO-494-gym-integration-tests Aug 19, 2026
4 checks passed
@JashG
JashG deleted the jgulabrai/AALGO-489-gym-base-image branch August 19, 2026 17:11
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