update string quote; added process=0 check

Author: Zephyr271828

src/MaxText/pyconfig.py

@@ -120,7 +120,7 @@ def _prepare_for_pydantic(raw_keys: dict[str, Any]) -> dict[str, Any]:
   for key, value in raw_keys.items():
     if key not in valid_fields:
       logger.warning("Ignoring invalid/unsupported field from YAML/CLI: %s", repr(key))
-      raise ValueError(f"{key!r} not in {', '.join(map(repr, valid_fields))}.")
+      raise ValueError(f"{key!r} not in {", ".join(map(repr, valid_fields))}.")
 
     new_value = value
     if isinstance(new_value, str) and new_value.lower() == "none":

src/maxtext/common/metric_logger.py

@@ -104,8 +104,7 @@ def __init__(self, config, learning_rate_schedule):
     if self.config.managed_mldiagnostics:
       ManagedMLDiagnostics(config)  # Initialize the MLRun instance.
 
-    self.enable_wandb = self.config.enable_wandb and jax.process_index() == 0
-    if self.enable_wandb: 
+    if self.config.enable_wandb and jax.process_index() == 0: 
       wandb.init(
         project=config.wandb_project_name,
         name=config.wandb_run_name,
@@ -133,7 +132,7 @@ def write_metrics(self, metrics, step, is_training=True):
       if self.config.managed_mldiagnostics:
         self.write_metrics_to_managed_mldiagnostics(metrics, step)
 
-      if self.enable_wandb:
+      if self.config.enable_wandb and jax.process_index() == 0:
         self.write_metrics_to_wandb(metrics, step)
 
   def log_metrics(self, metrics, step, is_training):