diff --git a/src/maxtext/common/goodput.py b/src/maxtext/common/goodput.py index eadbe503e1..95828effec 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 _resolve_goodput_monitor(config): + """Resolves the GoodputMonitor class to use and its class-specific kwargs. + + Returns (monitor_class, monitor_kwargs), where monitor_kwargs holds only + the constructor arguments specific to that class. + """ + 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 + + return monitoring_elastic.ElasticGoodputMonitor, {"include_slice_efficiency": True} + except Exception as e: # pylint: disable=broad-exception-caught + max_logging.log(f"Goodput: could not resolve elastic goodput monitor, falling back to base monitor: {e}") + + return monitoring.GoodputMonitor, {"pathway_enabled": config.enable_pathways_goodput} + + @contextlib.contextmanager def maybe_monitor_goodput(config): """Monitor cumulative goodput if enabled on the lead host. @@ -66,16 +86,9 @@ 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!") + monitor_class, monitor_kwargs = _resolve_goodput_monitor(config) + max_logging.log(f"Goodput: using {monitor_class.__name__} for job: {config.run_name}") kwargs = { "job_name": config.run_name, "logger_name": f"goodput_{config.run_name}", @@ -86,9 +99,8 @@ def maybe_monitor_goodput(config): "include_step_deviation": config.monitor_step_time_deviation, "step_deviation_interval_seconds": config.step_deviation_interval_seconds, "gcp_options": gcp_options, + **monitor_kwargs, } - if monitor_class == monitoring.GoodputMonitor: - kwargs["pathway_enabled"] = config.enable_pathways_goodput goodput_monitor = monitor_class(**kwargs) goodput_monitor.start_goodput_uploader() @@ -143,14 +155,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..0907005b9c 100644 --- a/src/maxtext/utils/elastic_utils.py +++ b/src/maxtext/utils/elastic_utils.py @@ -57,7 +57,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 +69,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 +79,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..50a2e72ff1 100644 --- a/tests/unit/elastic_utils_test.py +++ b/tests/unit/elastic_utils_test.py @@ -465,6 +465,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.""" diff --git a/tests/unit/goodput_utils_test.py b/tests/unit/goodput_utils_test.py index cf998aabb6..4780f36883 100644 --- a/tests/unit/goodput_utils_test.py +++ b/tests/unit/goodput_utils_test.py @@ -24,11 +24,16 @@ GoodputEvent, RECORD_JOB_END_TIME, RECORD_JOB_START_TIME, + _resolve_goodput_monitor, create_goodput_recorder, maybe_monitor_goodput, maybe_record_goodput, record_goodput, ) +from ml_goodput_measurement import goodput +from ml_goodput_measurement import goodput_elastic +from ml_goodput_measurement import monitoring +from ml_goodput_measurement import monitoring_elastic from tests.utils.test_helpers import get_test_config_path, get_test_base_output_directory pytestmark = [pytest.mark.external_training] @@ -93,6 +98,118 @@ 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 test_resolve_goodput_monitor_non_elastic(self): + """A McJAX config must resolve to the base monitor.""" + self.assertFalse(self.config.elastic_enabled) + + monitor_class, monitor_kwargs = _resolve_goodput_monitor(self.config) + + self.assertIs(monitor_class, monitoring.GoodputMonitor) + self.assertIn("pathway_enabled", monitor_kwargs) + self.assertNotIn("include_slice_efficiency", monitor_kwargs) + + @mock.patch("maxtext.utils.elastic_utils.should_use_elastic") + def test_resolve_goodput_monitor_elastic(self, mock_should_use_elastic): + """elastic_enabled=True on an actual Pathways run must resolve to the elastic monitor.""" + mock_should_use_elastic.return_value = True + self.config.elastic_enabled = True + + monitor_class, monitor_kwargs = _resolve_goodput_monitor(self.config) + + self.assertIs(monitor_class, monitoring_elastic.ElasticGoodputMonitor) + self.assertEqual(monitor_kwargs, {"include_slice_efficiency": True}) + self.assertNotIn("pathway_enabled", monitor_kwargs) + + @mock.patch("maxtext.utils.elastic_utils.should_use_elastic") + def test_resolve_goodput_monitor_elastic_enabled_not_pathways_falls_back(self, mock_should_use_elastic): + """elastic_enabled=True but not actually on Pathways must fall back to the base monitor.""" + mock_should_use_elastic.return_value = False + self.config.elastic_enabled = True + + monitor_class, monitor_kwargs = _resolve_goodput_monitor(self.config) + + self.assertIs(monitor_class, monitoring.GoodputMonitor) + self.assertIn("pathway_enabled", monitor_kwargs) + + @mock.patch("ml_goodput_measurement.monitoring_elastic.ElasticGoodputMonitor", side_effect=RuntimeError("boom")) + @mock.patch("maxtext.utils.elastic_utils.should_use_elastic") + def test_resolve_goodput_monitor_elastic_construction_failure_falls_back( + self, mock_should_use_elastic, unused_mock_elastic_monitor_cls + ): + """A failure constructing the elastic monitor must fall back rather than propagate.""" + mock_should_use_elastic.return_value = True + self.config.elastic_enabled = True + + monitor_class, _ = _resolve_goodput_monitor(self.config) + + self.assertIs(monitor_class, 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() + self.config.elastic_enabled = True + + with maybe_monitor_goodput(self.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 + self.config.elastic_enabled = True + + recorder = create_goodput_recorder(self.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 + self.config.elastic_enabled = True + + recorder = create_goodput_recorder(self.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 + self.config.elastic_enabled = True + + recorder = create_goodput_recorder(self.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")