From 8b75205d3b2ffe663cfa2050b9436b076427dc66 Mon Sep 17 00:00:00 2001 From: Dan Saunders Date: Mon, 24 Feb 2025 00:53:58 +0000 Subject: [PATCH] sleep on all ranks in distributed setting --- src/axolotl/telemetry/manager.py | 19 ++++++++++--------- tests/telemetry/test_manager.py | 4 +++- 2 files changed, 13 insertions(+), 10 deletions(-) diff --git a/src/axolotl/telemetry/manager.py b/src/axolotl/telemetry/manager.py index 2b4460d7b..dd2e4b4fe 100644 --- a/src/axolotl/telemetry/manager.py +++ b/src/axolotl/telemetry/manager.py @@ -68,11 +68,6 @@ class TelemetryManager: self.enabled, self.explicit_enable = self._check_telemetry_enabled() if self.enabled: - # Warn about telemetry collection - if not self.explicit_enable: - LOG.warning(ENABLED_WARNING) - time.sleep(ENABLED_WARNING_SLEEP_SECONDS) - self.run_id = str(uuid.uuid4()) self.whitelist = self._load_whitelist() self.system_info = self._get_system_info() @@ -105,10 +100,6 @@ class TelemetryManager: - Boolean denoting whether telemetry is enabled or disabled. - Boolean denoting whether telemetry is explicitly enabled or not. """ - # In the distributed setting, check whether we're running on rank 0 - if not is_main_process(): - return False, False - # Parse relevant env vars and fill opt-out default values axolotl_do_not_track = os.getenv("AXOLOTL_DO_NOT_TRACK") do_not_track = os.getenv("DO_NOT_TRACK") @@ -128,6 +119,16 @@ class TelemetryManager: "true", ) and do_not_track.lower() not in ("1", "true") + # Show warning (and sleep on all ranks) unless explicitly enabled + if enabled and not explicit_enabled: + if is_main_process(): + LOG.warning(ENABLED_WARNING) + time.sleep(ENABLED_WARNING_SLEEP_SECONDS) + + # Only rank 0 will send telemetry + if not is_main_process(): + return False, False + return enabled, explicit_enabled def _load_whitelist(self) -> dict: diff --git a/tests/telemetry/test_manager.py b/tests/telemetry/test_manager.py index 629016510..0085f7b0f 100644 --- a/tests/telemetry/test_manager.py +++ b/tests/telemetry/test_manager.py @@ -76,7 +76,9 @@ def test_telemetry_disabled_with_do_not_track(telemetry_manager_class): def test_telemetry_disabled_for_non_main_process(telemetry_manager_class): """Test that telemetry is disabled for non-main processes""" - with patch("axolotl.telemetry.manager.is_main_process", return_value=False): + with patch.dict(os.environ, {"AXOLOTL_DO_NOT_TRACK": "0"}), patch( + "axolotl.telemetry.manager.is_main_process", return_value=False + ): manager = telemetry_manager_class() assert not manager.enabled