From 8276299875d87479e0362c3fe68dbde345069b1d Mon Sep 17 00:00:00 2001 From: Dipannita Shaw Date: Tue, 11 Aug 2026 17:24:53 -0700 Subject: [PATCH] Add more elastic goodput tests --- src/maxtext/common/goodput.py | 55 ++++++---- src/maxtext/utils/elastic_utils.py | 13 +-- tests/unit/elastic_utils_test.py | 49 ++++++--- tests/unit/goodput_utils_test.py | 155 +++++++++++++++++++++++++++++ 4 files changed, 233 insertions(+), 39 deletions(-) diff --git a/src/maxtext/common/goodput.py b/src/maxtext/common/goodput.py index eadbe503e1..cf4d8c0043 100644 --- a/src/maxtext/common/goodput.py +++ b/src/maxtext/common/goodput.py @@ -42,6 +42,26 @@ class GoodputEvent(Enum): RECORD_JOB_END_TIME = f"record_{GoodputEvent.JOB.value}_end_time" +def _construct_goodput_monitor(config, common_kwargs): + """Constructs a GoodputMonitor, preferring the elastic-aware monitor when applicable.""" + if config.elastic_enabled: + try: + from maxtext.utils import elastic_utils # pylint: disable=import-outside-toplevel + + if elastic_utils.should_use_elastic(config): + from ml_goodput_measurement import monitoring_elastic # pylint: disable=import-outside-toplevel + + monitor = monitoring_elastic.ElasticGoodputMonitor(include_slice_efficiency=True, **common_kwargs) + max_logging.log(f"Goodput: using ElasticGoodputMonitor for job: {config.run_name}") + return monitor + except Exception as e: # pylint: disable=broad-exception-caught + max_logging.log(f"Goodput: could not create elastic goodput monitor, falling back to base monitor: {e}") + + monitor = monitoring.GoodputMonitor(pathway_enabled=config.enable_pathways_goodput, **common_kwargs) + max_logging.log(f"Goodput: using GoodputMonitor for job: {config.run_name}") + return monitor + + @contextlib.contextmanager def maybe_monitor_goodput(config): """Monitor cumulative goodput if enabled on the lead host. @@ -66,17 +86,8 @@ def maybe_monitor_goodput(config): enable_gcp_goodput_metrics=config.enable_gcp_goodput_metrics, enable_gcp_step_deviation_metrics=config.enable_gcp_step_deviation_metrics, ) - monitor_class = monitoring.GoodputMonitor - - if config.elastic_enabled: - try: - from ml_goodput_measurement import monitoring_elastic # pylint: disable=import-outside-toplevel - - monitor_class = monitoring_elastic.ElasticGoodputMonitor - except ImportError: - max_logging.log("Elastic monitor failed!") - kwargs = { + common_kwargs = { "job_name": config.run_name, "logger_name": f"goodput_{config.run_name}", "tensorboard_dir": config.tensorboard_dir, @@ -87,10 +98,8 @@ def maybe_monitor_goodput(config): "step_deviation_interval_seconds": config.step_deviation_interval_seconds, "gcp_options": gcp_options, } - if monitor_class == monitoring.GoodputMonitor: - kwargs["pathway_enabled"] = config.enable_pathways_goodput - goodput_monitor = monitor_class(**kwargs) + goodput_monitor = _construct_goodput_monitor(config, common_kwargs) goodput_monitor.start_goodput_uploader() max_logging.log("Started Goodput upload to Tensorboard & GCM in the background!") yield @@ -143,14 +152,18 @@ def create_goodput_recorder(config): # Detect if we should use the elastic-aware recorder if config.elastic_enabled: try: - from ml_goodput_measurement import goodput_elastic # pylint: disable=import-outside-toplevel from maxtext.utils import elastic_utils # pylint: disable=import-outside-toplevel - recorder = goodput_elastic.ElasticGoodputRecorder(config.run_name, logger_name, jax.process_index() == 0) - elastic_utils.record_slice_state(recorder) - except ImportError as e: - max_logging.log(f"Could not create elastic goodput recorder: {e}") - else: - return recorder + if elastic_utils.should_use_elastic(config): + from ml_goodput_measurement import goodput_elastic # pylint: disable=import-outside-toplevel + + recorder = goodput_elastic.ElasticGoodputRecorder(config.run_name, logger_name, jax.process_index() == 0) + elastic_utils.record_slice_state(recorder) + max_logging.log(f"Goodput: created ElasticGoodputRecorder for job: {config.run_name}") + return recorder + except Exception as e: # pylint: disable=broad-exception-caught + max_logging.log(f"Goodput: could not create elastic goodput recorder, falling back to base recorder: {e}") - return goodput.GoodputRecorder(config.run_name, logger_name, jax.process_index() == 0) + recorder = goodput.GoodputRecorder(config.run_name, logger_name, jax.process_index() == 0) + max_logging.log(f"Goodput: created base GoodputRecorder for job: {config.run_name}") + return recorder diff --git a/src/maxtext/utils/elastic_utils.py b/src/maxtext/utils/elastic_utils.py index fc81135c8a..9d0d62fc28 100644 --- a/src/maxtext/utils/elastic_utils.py +++ b/src/maxtext/utils/elastic_utils.py @@ -22,6 +22,7 @@ from maxtext.utils import gcs_utils from maxtext.utils import max_logging import pathwaysutils +from pathwaysutils.elastic import elastic from pathwaysutils.elastic import manager elastic_manager: manager.Manager | None = None @@ -39,11 +40,11 @@ def record_slice_state(recorder, active_slices_override: int | None = None) -> N ): return - available_slices = len(pathwaysutils.elastic.get_active_slice_indices()) + available_slices = len(elastic.get_active_slice_indices()) active_slices = ( active_slices_override if active_slices_override is not None else len(elastic_manager.active_slice_indices) ) - total_slices = len(pathwaysutils.elastic.get_slice_to_devices(jax.devices())) + total_slices = len(elastic.get_slice_to_devices(jax.devices())) recorder.record_elastic_slice_counts( available_slices=available_slices, @@ -57,7 +58,7 @@ def record_elastic_event_start(recorder, config) -> None: global pending_elastic_event_type event_type = "elastic_scale_up" if is_scale_up_event(config) else "elastic_slice_down" pending_elastic_event_type = event_type - if recorder: + if recorder and hasattr(recorder, "record_elastic_wait_start_time"): recorder.record_elastic_wait_start_time(event_type=event_type) record_slice_state(recorder, active_slices_override=0) @@ -69,7 +70,7 @@ def record_elastic_wait_end_and_reinit_start(recorder) -> None: return event_type = pending_elastic_event_type pending_elastic_event_type = None - if recorder: + if recorder and hasattr(recorder, "record_elastic_wait_end_time"): recorder.record_elastic_wait_end_time(event_type=event_type) recorder.record_elastic_reinit_start_time() record_slice_state(recorder) @@ -79,10 +80,10 @@ def record_elastic_wait_end_and_reinit_start(recorder) -> None: def record_elastic_reinit_end() -> None: """Records end of elastic reinitialization event.""" global pending_reinit_recorder - if pending_reinit_recorder is not None: + if pending_reinit_recorder is not None and hasattr(pending_reinit_recorder, "record_elastic_reinit_end_time"): pending_reinit_recorder.record_elastic_reinit_end_time() record_slice_state(pending_reinit_recorder) - pending_reinit_recorder = None + pending_reinit_recorder = None def elastic_enabled(config) -> bool: diff --git a/tests/unit/elastic_utils_test.py b/tests/unit/elastic_utils_test.py index d9c90d3f2c..b795bf9dc4 100644 --- a/tests/unit/elastic_utils_test.py +++ b/tests/unit/elastic_utils_test.py @@ -61,6 +61,7 @@ def setUp(self): self.original_jax = elastic_utils.jax self.original_gcs_utils = elastic_utils.gcs_utils self.original_max_logging = elastic_utils.max_logging + self.original_elastic = elastic_utils.elastic self.original_manager_class = pathwaysutils.elastic.manager.Manager self.original_scale_up_signal_error = getattr(pathwaysutils.elastic.manager, "ScaleUpSignalError", None) @@ -72,12 +73,12 @@ def setUp(self): self.fake_jax = create_autospec(self.original_jax) self.fake_manager = create_autospec(self.original_manager_class, instance=True) self.fake_manager.available_inactive_slices = set() + self.fake_elastic = create_autospec(self.original_elastic) # Configure default behaviors if needed self.fake_pathwaysutils.is_pathways_backend_used.return_value = True - self.fake_pathwaysutils.elastic = Mock() - self.fake_pathwaysutils.elastic.get_active_slice_indices.return_value = [0, 1] - self.fake_pathwaysutils.elastic.get_slice_to_devices.return_value = { + self.fake_elastic.get_active_slice_indices.return_value = [0, 1] + self.fake_elastic.get_slice_to_devices.return_value = { 0: [FakeDevice()], 1: [FakeDevice()], } @@ -89,6 +90,7 @@ def setUp(self): self.fake_jax.errors.JaxRuntimeError = MockJaxRuntimeError elastic_utils.gcs_utils = self.fake_gcs_utils elastic_utils.max_logging = self.fake_logging + elastic_utils.elastic = self.fake_elastic # Hook up pathwaysutils.elastic.manager.Manager to return our fake_manager pathwaysutils.elastic.manager.Manager = lambda *args, **kwargs: self.fake_manager # pyrefly: ignore[bad-assignment] @@ -102,6 +104,7 @@ def tearDown(self): elastic_utils.jax = self.original_jax elastic_utils.gcs_utils = self.original_gcs_utils elastic_utils.max_logging = self.original_max_logging + elastic_utils.elastic = self.original_elastic pathwaysutils.elastic.manager.Manager = self.original_manager_class pathwaysutils.elastic.manager.ScaleUpSignalError = ( # pyrefly: ignore[bad-assignment] self.original_scale_up_signal_error, @@ -115,7 +118,7 @@ def test_record_slice_state(self): elastic_utils.elastic_manager = self.fake_manager self.fake_manager.active_slice_indices = {0} self.fake_manager.slice_to_devices = {0: [FakeDevice()], 1: [FakeDevice()]} - self.fake_pathwaysutils.elastic.get_active_slice_indices.return_value = {0, 1} + self.fake_elastic.get_active_slice_indices.return_value = {0, 1} fake_recorder = Mock() fake_recorder.record_elastic_slice_counts = Mock() @@ -131,8 +134,6 @@ def test_record_slice_state(self): def test_elastic_enabled(self): config = FakeConfig() self.fake_pathwaysutils.is_pathways_backend_used.return_value = True - self.fake_pathwaysutils.elastic = Mock() - self.fake_pathwaysutils.elastic.get_active_slice_indices.return_value = [0, 1] config.elastic_enabled = True self.assertTrue(elastic_utils.elastic_enabled(config)) @@ -177,8 +178,6 @@ def test_live_devices_no_pathways(self): def test_live_devices_pathways(self): """Tests live_devices when pathways is used.""" self.fake_pathwaysutils.is_pathways_backend_used.return_value = True - self.fake_pathwaysutils.elastic = Mock() - self.fake_pathwaysutils.elastic.get_active_slice_indices.return_value = [0, 1] device0 = FakeDevice(slice_index=0) device1 = FakeDevice(slice_index=1) self.fake_jax.devices.return_value = [device0, device1] @@ -203,8 +202,6 @@ def test_live_devices_disabled(self): def test_elastic_retry_disabled(self): """Tests elastic_retry when disabled but pathways is used.""" self.fake_pathwaysutils.is_pathways_backend_used.return_value = True - self.fake_pathwaysutils.elastic = Mock() - self.fake_pathwaysutils.elastic.get_active_slice_indices.return_value = [0, 1] config = FakeConfig() config.elastic_enabled = False msg = ( @@ -465,6 +462,36 @@ def test_record_elastic_reinit_end_on_cold_start(self): elastic_utils.record_elastic_reinit_end() + def test_record_elastic_event_start_non_elastic_recorder_noop(self): + """A recorder lacking the elastic API (e.g. the ImportError fallback) must not raise.""" + elastic_utils.elastic_manager = self.fake_manager + self.fake_manager.available_inactive_slices = set() + non_elastic_recorder = Mock(spec=[]) # No record_elastic_* attributes. + config = FakeConfig() + + elastic_utils.record_elastic_event_start(non_elastic_recorder, config) # Must not raise. + + self.assertEqual(elastic_utils.pending_elastic_event_type, "elastic_slice_down") + + def test_record_elastic_wait_end_and_reinit_start_non_elastic_recorder_noop(self): + """A recorder lacking the elastic API must not raise, but is still tracked as pending.""" + elastic_utils.pending_elastic_event_type = "elastic_slice_down" # pyrefly: ignore[bad-assignment] + non_elastic_recorder = Mock(spec=[]) + + elastic_utils.record_elastic_wait_end_and_reinit_start(non_elastic_recorder) # Must not raise. + + self.assertIs(elastic_utils.pending_reinit_recorder, non_elastic_recorder) + self.assertIsNone(elastic_utils.pending_elastic_event_type) + + def test_record_elastic_reinit_end_non_elastic_recorder_noop(self): + """A recorder lacking the elastic API must not raise, and pending state is still cleared.""" + non_elastic_recorder = Mock(spec=[]) + elastic_utils.pending_reinit_recorder = non_elastic_recorder + + elastic_utils.record_elastic_reinit_end() # Must not raise. + + self.assertIsNone(elastic_utils.pending_reinit_recorder) + def test_ensure_elastic_manager_initialized_readonly_config(self): """Tests that ensure_elastic_manager_initialized works with read-only config.""" @@ -479,8 +506,6 @@ def __setattr__(self, name, value): config = ReadOnlyConfig() self.fake_pathwaysutils.is_pathways_backend_used.return_value = True - self.fake_pathwaysutils.elastic = Mock() - self.fake_pathwaysutils.elastic.get_active_slice_indices.return_value = [0, 1] # Should not raise ValueError elastic_utils.ensure_elastic_manager_initialized(config) diff --git a/tests/unit/goodput_utils_test.py b/tests/unit/goodput_utils_test.py index cf998aabb6..485b6d29bc 100644 --- a/tests/unit/goodput_utils_test.py +++ b/tests/unit/goodput_utils_test.py @@ -14,16 +14,23 @@ """Tests for goodput_utils.py""" +import tempfile import unittest from unittest import mock import pytest +goodput = pytest.importorskip("ml_goodput_measurement.goodput") +goodput_elastic = pytest.importorskip("ml_goodput_measurement.goodput_elastic") +monitoring = pytest.importorskip("ml_goodput_measurement.monitoring") +monitoring_elastic = pytest.importorskip("ml_goodput_measurement.monitoring_elastic") + from maxtext.configs import pyconfig from maxtext.common.goodput import ( GoodputEvent, RECORD_JOB_END_TIME, RECORD_JOB_START_TIME, + _construct_goodput_monitor, create_goodput_recorder, maybe_monitor_goodput, maybe_record_goodput, @@ -34,6 +41,19 @@ pytestmark = [pytest.mark.external_training] +class _ConfigOverride: + """Wrapper exposing overridden attributes on top of a read-only config.""" + + def __init__(self, base_config, **overrides): + self._base_config = base_config + self._overrides = overrides + + def __getattr__(self, name): + if name in self._overrides: + return self._overrides[name] + return getattr(self._base_config, name) + + class GoodputUtilsTest(unittest.TestCase): """Tests for Goodput monitoring and recording.""" @@ -93,6 +113,141 @@ def test_job_recording_constants(self): self.assertEqual(RECORD_JOB_START_TIME, "record_job_start_time") self.assertEqual(RECORD_JOB_END_TIME, "record_job_end_time") + def _common_monitor_kwargs(self): + """Helper to construct goodput monitor.""" + return { + "job_name": self.config.run_name, + "logger_name": f"goodput_{self.config.run_name}", + "tensorboard_dir": tempfile.mkdtemp(), + "upload_interval": self.config.goodput_upload_interval_seconds, + "monitoring_enabled": True, + "include_badput_breakdown": True, + "include_step_deviation": self.config.monitor_step_time_deviation, + "step_deviation_interval_seconds": self.config.step_deviation_interval_seconds, + "gcp_options": monitoring.GCPOptions(), + } + + @mock.patch("google.cloud.logging.Client") + def test_construct_goodput_monitor_non_elastic(self, mock_cloud_logger): + """A McJAX config must construct the base monitor.""" + mock_cloud_logger.return_value = mock.MagicMock() + self.assertFalse(self.config.elastic_enabled) + + monitor = _construct_goodput_monitor(self.config, self._common_monitor_kwargs()) + + self.assertIsInstance(monitor, monitoring.GoodputMonitor) + self.assertNotIsInstance(monitor, monitoring_elastic.ElasticGoodputMonitor) + + @mock.patch("maxtext.utils.elastic_utils.should_use_elastic") + @mock.patch("google.cloud.logging.Client") + def test_construct_goodput_monitor_elastic(self, mock_cloud_logger, mock_should_use_elastic): + """elastic_enabled=True on an actual Pathways run must construct the elastic monitor.""" + mock_cloud_logger.return_value = mock.MagicMock() + mock_should_use_elastic.return_value = True + config = _ConfigOverride(self.config, elastic_enabled=True) + + monitor = _construct_goodput_monitor(config, self._common_monitor_kwargs()) + + self.assertIsInstance(monitor, monitoring_elastic.ElasticGoodputMonitor) + + @mock.patch("maxtext.utils.elastic_utils.should_use_elastic") + @mock.patch("google.cloud.logging.Client") + def test_construct_goodput_monitor_elastic_enabled_not_pathways_falls_back( + self, mock_cloud_logger, mock_should_use_elastic + ): + """elastic_enabled=True but not actually on Pathways must fall back to the base monitor.""" + mock_cloud_logger.return_value = mock.MagicMock() + mock_should_use_elastic.return_value = False + config = _ConfigOverride(self.config, elastic_enabled=True) + + monitor = _construct_goodput_monitor(config, self._common_monitor_kwargs()) + + self.assertIsInstance(monitor, monitoring.GoodputMonitor) + self.assertNotIsInstance(monitor, monitoring_elastic.ElasticGoodputMonitor) + + @mock.patch("ml_goodput_measurement.monitoring_elastic.ElasticGoodputMonitor", side_effect=RuntimeError("boom")) + @mock.patch("maxtext.utils.elastic_utils.should_use_elastic") + @mock.patch("google.cloud.logging.Client") + def test_construct_goodput_monitor_elastic_construction_failure_falls_back( + self, mock_cloud_logger, mock_should_use_elastic, unused_mock_elastic_monitor_cls + ): + """A failure constructing the elastic monitor must fall back rather than propagate.""" + mock_cloud_logger.return_value = mock.MagicMock() + mock_should_use_elastic.return_value = True + config = _ConfigOverride(self.config, elastic_enabled=True) + + monitor = _construct_goodput_monitor(config, self._common_monitor_kwargs()) # Must not raise. + + # Note: not asserting assertNotIsInstance(monitor, monitoring_elastic.ElasticGoodputMonitor) + # here - that name is itself patched to a Mock for this test, so it isn't a usable type. + self.assertIsInstance(monitor, monitoring.GoodputMonitor) + + @mock.patch("ml_goodput_measurement.monitoring_elastic.ElasticGoodputMonitor.stop_goodput_uploader") + @mock.patch("ml_goodput_measurement.monitoring_elastic.ElasticGoodputMonitor.start_goodput_uploader") + @mock.patch("maxtext.utils.elastic_utils.should_use_elastic") + def test_monitor_goodput_elastic( + self, mock_should_use_elastic, mock_start_goodput_uploader, mock_stop_goodput_uploader + ): + """maybe_monitor_goodput actually starts/stops an ElasticGoodputMonitor when elastic is active.""" + mock_should_use_elastic.return_value = True + mock_start_goodput_uploader.return_value = mock.MagicMock() + config = _ConfigOverride(self.config, elastic_enabled=True) + + with maybe_monitor_goodput(config): + mock_start_goodput_uploader.assert_called() + mock_stop_goodput_uploader.assert_called() + + @mock.patch("google.cloud.logging.Client") + def test_create_goodput_recorder_non_elastic(self, mock_cloud_logger): + """Regular (non-Pathways/McJAX) config must get the base recorder.""" + mock_cloud_logger.return_value = mock.MagicMock() + self.assertFalse(self.config.elastic_enabled) + + recorder = create_goodput_recorder(self.config) + + self.assertNotIsInstance(recorder, goodput_elastic.ElasticGoodputRecorder) + + @mock.patch("maxtext.utils.elastic_utils.should_use_elastic") + @mock.patch("google.cloud.logging.Client") + def test_create_goodput_recorder_elastic(self, mock_cloud_logger, mock_should_use_elastic): + """elastic_enabled=True on an actual Pathways run must get the elastic recorder.""" + mock_cloud_logger.return_value = mock.MagicMock() + mock_should_use_elastic.return_value = True + config = _ConfigOverride(self.config, elastic_enabled=True) + + recorder = create_goodput_recorder(config) + + self.assertIsInstance(recorder, goodput_elastic.ElasticGoodputRecorder) + + @mock.patch("maxtext.utils.elastic_utils.should_use_elastic") + @mock.patch("google.cloud.logging.Client") + def test_create_goodput_recorder_elastic_enabled_not_pathways_falls_back( + self, mock_cloud_logger, mock_should_use_elastic + ): + """elastic_enabled=True but not actually on Pathways (e.g. McJAX) must get the base recorder.""" + mock_cloud_logger.return_value = mock.MagicMock() + mock_should_use_elastic.return_value = False + config = _ConfigOverride(self.config, elastic_enabled=True) + + recorder = create_goodput_recorder(config) + + self.assertNotIsInstance(recorder, goodput_elastic.ElasticGoodputRecorder) + + @mock.patch("ml_goodput_measurement.goodput_elastic.ElasticGoodputRecorder", side_effect=RuntimeError("boom")) + @mock.patch("maxtext.utils.elastic_utils.should_use_elastic") + @mock.patch("google.cloud.logging.Client") + def test_create_goodput_recorder_elastic_construction_failure_falls_back( + self, mock_cloud_logger, mock_should_use_elastic, unused_mock_elastic_recorder_cls + ): + """A failure constructing the elastic recorder must fall back rather than propagate.""" + mock_cloud_logger.return_value = mock.MagicMock() + mock_should_use_elastic.return_value = True + config = _ConfigOverride(self.config, elastic_enabled=True) + + recorder = create_goodput_recorder(config) # Must not raise. + + self.assertIsInstance(recorder, goodput.GoodputRecorder) + @mock.patch("ml_goodput_measurement.goodput.GoodputRecorder.record_job_end_time") @mock.patch("ml_goodput_measurement.goodput.GoodputRecorder.record_job_start_time") @mock.patch("google.cloud.logging.Client")