fix(docker): support non-root runtime execution

Signed-off-by: oliver könig <okoenig@nvidia.com>

Author: ko3n1g

.dockerignore

@@ -4,6 +4,7 @@ __pycache__
 *.pyd
 .Python
 env
+.venv
 pip-log.txt
 pip-delete-this-directory.txt
 .tox

docker/Dockerfile

@@ -11,6 +11,12 @@
 #     docker buildx build --build-context nemo-gym=. -f docker/Dockerfile \
 #       --tag <registry>/nemo-gym:latest --push .
 #
+#   Non-root runtime override (the image still defaults to root):
+#     docker run --rm --user 65532:65532 \
+#       --env HOME=/home/nemo-runtime \
+#       --env XDG_CACHE_HOME=/opt/nemo-gym/cache/xdg \
+#       <registry>/nemo-gym:latest --help
+#
 
 ARG BASE_IMAGE=nvcr.io/nvidia/cuda-dl-base:26.03-cuda13.2-devel-ubuntu24.04
 
@@ -68,9 +74,28 @@ EOF
 
 ARG UV_VERSION=0.11.29
 ARG PYTHON_VERSION=3.13.14
-ENV PATH="/root/.local/bin:$PATH"
-RUN curl -LsSf https://astral.sh/uv/${UV_VERSION}/install.sh | sh && \
-    uv python install ${PYTHON_VERSION}
+ENV UV_PYTHON_INSTALL_DIR=/opt/uv/python
+ENV UV_CACHE_DIR=/opt/uv/cache
+RUN curl -LsSf https://astral.sh/uv/${UV_VERSION}/install.sh | \
+        env UV_INSTALL_DIR=/usr/local/bin UV_NO_MODIFY_PATH=1 sh && \
+    /usr/local/bin/uv python install ${PYTHON_VERSION}
+ENV PATH="/usr/local/bin:$PATH"
+
+ARG RUNTIME_UID=65532
+ARG RUNTIME_GID=65532
+RUN if ! getent group "${RUNTIME_GID}" >/dev/null; then \
+        groupadd --gid "${RUNTIME_GID}" nemo-runtime; \
+    fi && \
+    if ! getent passwd "${RUNTIME_UID}" >/dev/null; then \
+        useradd --no-log-init --uid "${RUNTIME_UID}" --gid "${RUNTIME_GID}" \
+            --create-home --home-dir /home/nemo-runtime --shell /bin/bash nemo-runtime; \
+    fi && \
+    install -d -o "${RUNTIME_UID}" -g "${RUNTIME_GID}" \
+        /home/nemo-runtime \
+        /opt/nemo-gym \
+        /opt/nemo-gym/cache \
+        /opt/nemo-gym/results \
+        /opt/uv/cache
 
 ENV RAY_USAGE_STATS_ENABLED=0
 
@@ -81,13 +106,16 @@ WORKDIR /opt/nemo-gym
 
 ARG BASE_IMAGE
 ARG UV_VERSION
+ARG RUNTIME_UID
+ARG RUNTIME_GID
 
 ENV UV_PROJECT_ENVIRONMENT=/opt/nemo_gym_venv
 ENV UV_LINK_MODE=copy
 
 # Copy only dependency metadata first for layer caching.
-COPY --from=nemo-gym pyproject.toml uv.lock ./
-COPY --from=nemo-gym nemo_gym/__init__.py nemo_gym/package_info.py ./nemo_gym/
+COPY --from=nemo-gym --chown=${RUNTIME_UID}:${RUNTIME_GID} pyproject.toml uv.lock ./
+COPY --from=nemo-gym --chown=${RUNTIME_UID}:${RUNTIME_GID} \
+    nemo_gym/__init__.py nemo_gym/package_info.py ./nemo_gym/
 
 RUN uv venv --seed && \
     uv sync --link-mode symlink --locked --extra vllm --no-install-project
@@ -111,11 +139,19 @@ LABEL com.nvidia.build.ref="${NVIDIA_BUILD_REF}"
 
 # Copy full source. Exclude pyproject.toml and uv.lock since they are
 # already present from the hermetic stage.
-COPY --from=nemo-gym --exclude=pyproject.toml --exclude=uv.lock . /opt/nemo-gym
+COPY --from=nemo-gym --chown=${RUNTIME_UID}:${RUNTIME_GID} \
+    --exclude=pyproject.toml --exclude=uv.lock . /opt/nemo-gym
 
 # Install the nemo-gym project itself along with the vllm extra.
 RUN UV_LINK_MODE=symlink uv sync --locked --extra vllm
 
+# Keep runtime uv and XDG caches separate from the root-owned build cache that
+# backs the symlinked project environment.
+ENV UV_CACHE_DIR=/opt/nemo-gym/cache/uv
+RUN install -d -o "${RUNTIME_UID}" -g "${RUNTIME_GID}" \
+        "${UV_CACHE_DIR}" \
+        /opt/nemo-gym/cache/xdg
+
 # Optional: pre-warm per-server venvs at build time so they are ready at runtime.
 # Pass --build-arg NEMO_GYM_PREFETCH_CONFIGS="path/to/config1.yaml,path/to/config2.yaml"
 # Uses 'gym env prefetch' which accepts the same config format as 'gym env start',
@@ -126,5 +162,53 @@ if [[ -n "${NEMO_GYM_PREFETCH_CONFIGS:-}" ]]; then
 fi
 EOF
 
+# Prefetched caches and venvs are built as root. Own the project tree and the
+# project environment's directories so the runtime identity can mutate them.
+RUN venv_site_packages="$(/opt/nemo_gym_venv/bin/python -c \
+        'import site; print(site.getsitepackages()[0])')" && \
+    touch "${venv_site_packages}/.nonroot-uninstall-probe" && \
+    chown -R "${RUNTIME_UID}:${RUNTIME_GID}" /opt/nemo-gym && \
+    find /opt/nemo_gym_venv -type d \
+        -exec chown "${RUNTIME_UID}:${RUNTIME_GID}" {} +
+
+USER ${RUNTIME_UID}:${RUNTIME_GID}
+
+# Runtime callers select this UID and set HOME/cache variables explicitly; the
+# published image remains root by default. Prove that the CLI, uv toolchain,
+# project environment, source tree, and runtime paths support that contract.
+RUN export HOME=/home/nemo-runtime \
+        XDG_CACHE_HOME=/opt/nemo-gym/cache/xdg && \
+    test "$(id -u)" = "${RUNTIME_UID}" && \
+    test -w "${HOME}" && \
+    test -w "${XDG_CACHE_HOME}" && \
+    test -w "${UV_CACHE_DIR}" && \
+    test -w /opt/nemo-gym/cache && \
+    test -w /opt/nemo-gym/results && \
+    test -x /usr/local/bin/uv && \
+    test -x /usr/local/bin/uvx && \
+    test -x /opt/nemo_gym_venv/bin/python && \
+    venv_site_packages="$(python -c 'import site; print(site.getsitepackages()[0])')" && \
+    rm "${venv_site_packages}/.nonroot-uninstall-probe" && \
+    touch "${HOME}/.nonroot-write-probe" \
+        "${XDG_CACHE_HOME}/.nonroot-write-probe" \
+        "${UV_CACHE_DIR}/.nonroot-write-probe" \
+        /opt/nemo-gym/.nonroot-write-probe \
+        /opt/nemo-gym/cache/.nonroot-write-probe \
+        /opt/nemo-gym/results/.nonroot-write-probe \
+        "${venv_site_packages}/.nonroot-write-probe" && \
+    gym --help >/dev/null && \
+    uv venv /opt/nemo-gym/cache/container-smoke && \
+    /opt/nemo-gym/cache/container-smoke/bin/python -c \
+        "import sys; assert sys.version_info[:2] == (3, 13)" && \
+    rm -r /opt/nemo-gym/cache/container-smoke \
+        "${HOME}/.nonroot-write-probe" \
+        "${XDG_CACHE_HOME}/.nonroot-write-probe" \
+        "${UV_CACHE_DIR}/.nonroot-write-probe" \
+        /opt/nemo-gym/.nonroot-write-probe \
+        /opt/nemo-gym/cache/.nonroot-write-probe \
+        /opt/nemo-gym/results/.nonroot-write-probe \
+        "${venv_site_packages}/.nonroot-write-probe"
+
+USER root
 ENTRYPOINT ["gym"]
 CMD ["--help"]