resolve comments

Author: YixuanWang-99

src/maxtext/common/goodput.py

@@ -135,21 +135,22 @@ def create_goodput_recorder(config):
     if config.enable_goodput_recording and jax.process_index() == 0:
       max_logging.log("[GOODPUT NO-OP] recorder skipped (decoupled stub).")
     return None
-  if config.enable_goodput_recording:
-    logger_name = f"goodput_{config.run_name}"
 
-    # Detect if we should use the elastic-aware recorder
-    recorder = None
-    if config.elastic_enabled:
-      try:
-        from ml_goodput_measurement import goodput_elastic
-        from maxtext.utils import elastic_utils
-        recorder = goodput_elastic.ElasticGoodputRecorder(config.run_name, logger_name, jax.process_index() == 0)      
-        elastic_utils.record_slice_state(recorder)
-      except ImportError:
-        pass
-    
-    if recorder is None:
-      recorder = goodput.GoodputRecorder(config.run_name, logger_name, jax.process_index() == 0)
-    return recorder
-  return None
+  if not config.enable_goodput_recording:
+    return None
+
+  logger_name = f"goodput_{config.run_name}"
+
+  # Detect if we should use the elastic-aware recorder
+  if config.elastic_enabled:
+    try:
+      from ml_goodput_measurement import goodput_elastic
+      from maxtext.utils import elastic_utils
+      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
+
+  return goodput.GoodputRecorder(config.run_name, logger_name, jax.process_index() == 0)

tests/unit/elastic_utils_test.py

@@ -72,6 +72,7 @@ def setUp(self):
     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 = {0: [FakeDevice()], 1: [FakeDevice()]}
     self.fake_jax.process_index.return_value = 0
 
     # Inject fakes into elastic_utils namespace