Skip to content

feat: add export_mounts and export_image to auto-export config#911

Merged
AdamRajfer merged 1 commit intomainfrom
agronskiy/chore/allow-uv-cache-for-exporters
Apr 14, 2026
Merged

feat: add export_mounts and export_image to auto-export config#911
AdamRajfer merged 1 commit intomainfrom
agronskiy/chore/allow-uv-cache-for-exporters

Conversation

@agronskiy
Copy link
Copy Markdown
Collaborator

@agronskiy agronskiy commented Apr 14, 2026

This tries to address the issue of fast startup of export execution as requested by Gym folks

Approx speedups are here:

Setup apt-get+git install Total
Before (pip, python:3.12-slim) ~1-2 min ~2-3 min 3-5 min
After, cold uv cache 0s ~42s ~42s
After, warm uv cache 0s ~6s ~6s

What was done

  • Add export_mounts config key: mount arbitrary host paths (e.g. a pre-heated uv cache) into the export srun container
  • Add export_image config key: override the default python:3.12.7-slim export image with a custom one (e.g. one with uv+git pre-installed)

Together these eliminate the ~3-5 min apt-get + pip install overhead in the auto-export step, bringing it down to ~6s with a warm uv cache.

To use this, now use this config (with new launcher installed)

execution:
  auto_export:
    destinations: [mlflow]
    export_image: /path/to/uv-git-python3.12.sqsh  # or a registry image
    export_mounts:
      /lustre/.../users/me/.cache/uv: /cache/uv
    launcher_install_cmd: |
      export UV_CACHE_DIR=/cache/uv
      uv pip install --system "nemo-evaluator-launcher[all] @ git+https://github.com/NVIDIA-NeMo/Evaluator.git@main#subdirectory=packages/nemo-evaluator-launcher"

@copy-pr-bot
Copy link
Copy Markdown

copy-pr-bot Bot commented Apr 14, 2026

This pull request requires additional validation before any workflows can run on NVIDIA's runners.

Pull request vetters can view their responsibilities here.

Contributors can view more details about this message here.

Add two new optional keys to `execution.auto_export`:

- `export_mounts`: dict of host_path -> container_path to mount into the
  export srun (e.g. a pre-heated uv cache on Lustre)
- `export_image`: override the default python:3.12.7-slim container with
  a custom image (e.g. one with uv+git pre-installed)

Together these eliminate the ~3-5 min apt-get+pip overhead in the export
step, bringing it down to ~6s with a warm uv cache.

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
Signed-off-by: Alex Gronskiy <agronskiy@nvidia.com>
@agronskiy agronskiy force-pushed the agronskiy/chore/allow-uv-cache-for-exporters branch from fe203d8 to fe004b1 Compare April 14, 2026 15:53
@agronskiy agronskiy marked this pull request as ready for review April 14, 2026 15:53
@agronskiy agronskiy requested review from a team as code owners April 14, 2026 15:53
@AdamRajfer
Copy link
Copy Markdown
Contributor

/ok to test fe004b1

@AdamRajfer AdamRajfer merged commit 076c8de into main Apr 14, 2026
48 checks passed
@AdamRajfer AdamRajfer deleted the agronskiy/chore/allow-uv-cache-for-exporters branch April 14, 2026 19:11
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants