Add stage sync barriers and unbuffered host-ranked diagnostic logging

Author: amepas

src/maxdiffusion/pipelines/flux/flux2klein_pipeline.py

@@ -242,6 +242,14 @@ def __call__(
       trace["prompt_encoding"] = time.perf_counter() - t0
       max_logging.log(f" -> [TIMING] Prompt Encoding (Qwen3): {trace['prompt_encoding']:.4f} seconds ⏱️")
 
+      proc_id = jax.process_index()
+      proc_cnt = jax.process_count()
+      host_prefix = f"[HOST {proc_id}/{proc_cnt}] "
+
+      # Stage Sync 1: Phase A Complete
+      multihost_utils.sync_global_devices("phase_a_complete")
+      max_logging.log(f"{host_prefix} Passed Phase A Sync Barrier (phase_a_complete) successfully! ✅", flush=True)
+
       # Shard pipeline batch inputs across data axis ("data") for SPMD multi-host execution
       data_sharding = jax.sharding.NamedSharding(self.mesh, P("data"))
 
@@ -257,11 +265,25 @@ def put_data_on_devices(x, sharding):
       txt_ids_val = put_data_on_devices(txt_ids_val, data_sharding)
       img_ids_val = put_data_on_devices(img_ids_val, data_sharding)
 
+      max_logging.log(
+          f"{host_prefix} DIAGNOSTIC TENSORS BEFORE PHASE B:\n"
+          f"  latents_jax: shape={latents_jax.shape}, dtype={latents_jax.dtype}, sharding={getattr(latents_jax, 'sharding', None)}\n"
+          f"  prompt_embeds_jax: shape={prompt_embeds_jax.shape}, dtype={prompt_embeds_jax.dtype}, sharding={getattr(prompt_embeds_jax, 'sharding', None)}\n"
+          f"  txt_ids_val: shape={txt_ids_val.shape}, dtype={txt_ids_val.dtype}, sharding={getattr(txt_ids_val, 'sharding', None)}\n"
+          f"  img_ids_val: shape={img_ids_val.shape}, dtype={img_ids_val.dtype}, sharding={getattr(img_ids_val, 'sharding', None)}",
+          flush=True,
+      )
+
+      # Stage Sync 2: Pre-Phase B Start
+      multihost_utils.sync_global_devices("pre_phase_b_start")
+      max_logging.log(f"{host_prefix} Passed Pre-Phase B Sync Barrier (pre_phase_b_start) successfully! ✅", flush=True)
+
       # ---------------------------------------------------------------------
       # PHASE B: Denoising Loop (Flux Transformer - Fused JAX fori_loop)
       # ---------------------------------------------------------------------
       max_logging.log(
-          f"[PHASE B] Running {num_inference_steps}-step E2E Denoising Loop on a batch of {batch_size} images..."
+          f"{host_prefix} [PHASE B] Running {num_inference_steps}-step E2E Denoising Loop on a batch of {batch_size} images...",
+          flush=True,
       )
       t0 = time.perf_counter()
 
@@ -291,15 +313,27 @@ def step_fn(step_idx, current_latents):
 
         return jax.lax.fori_loop(0, num_inference_steps, step_fn, latents)
 
-      latents_jax = denoise_loop(
-          params,
-          latents_jax,
-          prompt_embeds_jax,
-          txt_ids_val,
-          img_ids_val,
-          scheduler_state,
-      )
-      latents_jax.block_until_ready()
+      try:
+        latents_jax = denoise_loop(
+            params,
+            latents_jax,
+            prompt_embeds_jax,
+            txt_ids_val,
+            img_ids_val,
+            scheduler_state,
+        )
+        latents_jax.block_until_ready()
+      except Exception as e:
+        max_logging.log(f"❌ {host_prefix} EXCEPTION IN DENOISE LOOP: {e}", flush=True)
+        import traceback
+
+        traceback.print_exc()
+        sys.stdout.flush()
+        raise e
+
+      # Stage Sync 3: Phase B Complete
+      multihost_utils.sync_global_devices("phase_b_complete")
+      max_logging.log(f"{host_prefix} Passed Phase B Sync Barrier (phase_b_complete) successfully! ✅", flush=True)
 
       trace["denoise_loop"] = time.perf_counter() - t0
       max_logging.log(f" -> [TIMING] Denoising Loop (Flux): {trace['denoise_loop']:.4f} seconds ⏱️")