From 7a82833c5243d6e597d91da1550525734c135040 Mon Sep 17 00:00:00 2001
From: Terry Kong <terryk@nvidia.com>
Date: Fri, 21 Mar 2025 13:26:30 -0700
Subject: [PATCH] fix: disable usage stats more forcefully since container env
 took precedence

Signed-off-by: Terry Kong <terryk@nvidia.com>
---
 docker/Dockerfile           |  5 ++++-
 nemo_reinforcer/__init__.py |  3 +++
 ray.sub                     | 12 ++++--------
 tests/unit/test_meta.py     | 34 ++++++++++++++++++++++++++++++++++
 4 files changed, 45 insertions(+), 9 deletions(-)
 create mode 100644 tests/unit/test_meta.py

diff --git a/docker/Dockerfile b/docker/Dockerfile
index 0abb6ebd9e..5842eed006 100644
--- a/docker/Dockerfile
+++ b/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
\ No newline at end of file
+FROM base
diff --git a/nemo_reinforcer/__init__.py b/nemo_reinforcer/__init__.py
index 76b1dda065..457028a211 100644
--- a/nemo_reinforcer/__init__.py
+++ b/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"
diff --git a/ray.sub b/ray.sub
index 0e254c6e67..dac0899a35 100644
--- a/ray.sub
+++ b/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
diff --git a/tests/unit/test_meta.py b/tests/unit/test_meta.py
new file mode 100644
index 0000000000..4b9e5bcc23
--- /dev/null
+++ b/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."
+    )