fix: disable usage stats more forcefully since container env took precedence (#25)

Signed-off-by: Terry Kong <terryk@nvidia.com>

Author: terrykong

docker/Dockerfile

@@ -9,6 +9,9 @@ RUN apt-get update && sudo apt-get install -y jq
 
 RUN pip install uv
 RUN echo "unset RAY_RUNTIME_ENV_HOOK" >> /home/ray/.bashrc
+# Disable usage stats by default for users who are sensitive to sharing usage.
+# Users are encouraged to enable if the wish.
+ENV RAY_USAGE_STATS_ENABLED=0
 
 FROM base AS hermetic
 # hermetic creates a virtual environment with the default dependencies pre-installed for convenience
@@ -38,4 +41,4 @@ EOF
 COPY --chown=ray --chmod=755 . /opt/reinforcer
 RUN uv pip install --no-deps --editable /opt/reinforcer
 
-FROM base
+FROM base

nemo_reinforcer/__init__.py

@@ -1,3 +1,4 @@
+import os
 from nemo_reinforcer.package_info import (
     __contact_emails__,
     __contact_names__,
@@ -11,3 +12,5 @@
     __shortversion__,
     __version__,
 )
+
+os.environ["RAY_USAGE_STATS_ENABLED"] = "0"

ray.sub

@@ -19,13 +19,6 @@ MOUNTS=$MOUNTS
 COMMAND=${COMMAND:-}  # This is a script relative to the SLURM_SUBMIT_DIR. If left empty, it will leave the cluster idle after it's brought up.
 ########################################################
 
-########################################################
-# Global settings
-########################################################
-# Disable usage stats by default so that it's opt-in
-export RAY_USAGE_STATS_ENABLED=${RAY_USAGE_STATS_ENABLED:-0}
-########################################################
-
 COMMON_SRUN_ARGS=""
 COMMON_SRUN_ARGS+=" --export=ALL"
 COMMON_SRUN_ARGS+=" --no-container-mount-home"
@@ -70,6 +63,7 @@ head_cmd=$(cat <<EOF
 env
 cat <<EOFINNER | tee /launch-head.sh
 ray start --head \
+--disable-usage-stats \
 --num-cpus=0 \
 --num-gpus=0 \
 --node-ip-address="$head_node_ip" \
@@ -104,9 +98,11 @@ for ((i = 0; i < SLURM_JOB_NUM_NODES; i++)); do
 env
 cat <<EOFINNER | tee /launch-worker.sh
 ray start --address "$ip_head" \
+          --disable-usage-stats \
           --resources="{\"worker_units\": $gpus_per_node}" \
           --min-worker-port=$min_worker_port \
-          --max-worker-port=$max_worker_port --block
+          --max-worker-port=$max_worker_port \
+          --block
 EOFINNER
 
 count=0

tests/unit/test_meta.py

@@ -0,0 +1,34 @@
+# Copyright (c) 2025, NVIDIA CORPORATION.  All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+# This module tests things outside of any package (e.g., things in the root __init__.py)
+
+import pytest
+import os
+
+
+def test_usage_stats_disabled_by_default():
+    import nemo_reinforcer
+
+    assert os.environ["RAY_USAGE_STATS_ENABLED"] == "0", (
+        f"Our dockerfile, slurm submission script and default environment setting when importing reinforcer should all disable usage stats collection. This failing is not expected."
+    )
+
+
+def test_usage_stats_disabled_in_tests():
+    import tests
+
+    assert os.environ["RAY_USAGE_STATS_ENABLED"] == "0", (
+        f"Our dockerfile, slurm submission script and default environment setting when importing reinforcer should all disable usage stats collection. This failing is not expected."
+    )