fix(flux2klein): route component XProf traces to output/flux2klein_generation/tensorboard/profile_* subdirectories

Author: amepas

src/maxdiffusion/flux2klein_pipeline.py

@@ -404,7 +404,8 @@ def put_data_on_devices(x, sharding):
         prompt_mask = put_data_on_devices(prompt_mask, data_sharding)
         do_prof_qwen3 = profile_target in ("all", "qwen3")
         if do_prof_qwen3:
-          jax.profiler.start_trace("/tmp/xprof_qwen3")
+          tb_dir = getattr(self._config, "tensorboard_dir", "/tmp")
+          jax.profiler.start_trace(os.path.join(tb_dir, "profile_qwen3"))
         with jax.named_scope("qwen3_text_encoder"):
           prompt_embeds_jax = self._jitted_qwen3_forward(qwen3_params, prompt_ids, prompt_mask)
         prompt_embeds_jax.block_until_ready()
@@ -468,7 +469,8 @@ def put_data_on_devices(x, sharding):
 
         do_prof_denoise = profile_target in ("all", "denoise")
         if do_prof_denoise:
-          jax.profiler.start_trace("/tmp/xprof_denoise")
+          tb_dir = getattr(self._config, "tensorboard_dir", "/tmp")
+          jax.profiler.start_trace(os.path.join(tb_dir, "profile_denoise"))
         with jax.named_scope("fused_flux_denoise_loop"):
           latents_jax = self._jitted_fused_denoise_loop(
               params,
@@ -519,7 +521,8 @@ def put_data_on_devices(x, sharding):
 
     do_prof_vae = profile_target in ("all", "vae")
     if do_prof_vae:
-      jax.profiler.start_trace("/tmp/xprof_vae")
+      tb_dir = getattr(self._config, "tensorboard_dir", "/tmp")
+      jax.profiler.start_trace(os.path.join(tb_dir, "profile_vae"))
     with jax.named_scope("vae_decoder"):
       decoded_out = self._jitted_vae_decode(vae_params, latents_jax, vae_bn_mean_jax, vae_bn_std_jax, height, width)
     images_rgb = decoded_out.sample

src/maxdiffusion/pipelines/flux/flux2klein_pipeline.py

@@ -404,7 +404,8 @@ def put_data_on_devices(x, sharding):
         prompt_mask = put_data_on_devices(prompt_mask, data_sharding)
         do_prof_qwen3 = profile_target in ("all", "qwen3")
         if do_prof_qwen3:
-          jax.profiler.start_trace("/tmp/xprof_qwen3")
+          tb_dir = getattr(self._config, "tensorboard_dir", "/tmp")
+          jax.profiler.start_trace(os.path.join(tb_dir, "profile_qwen3"))
         with jax.named_scope("qwen3_text_encoder"):
           prompt_embeds_jax = self._jitted_qwen3_forward(qwen3_params, prompt_ids, prompt_mask)
         prompt_embeds_jax.block_until_ready()
@@ -468,7 +469,8 @@ def put_data_on_devices(x, sharding):
 
         do_prof_denoise = profile_target in ("all", "denoise")
         if do_prof_denoise:
-          jax.profiler.start_trace("/tmp/xprof_denoise")
+          tb_dir = getattr(self._config, "tensorboard_dir", "/tmp")
+          jax.profiler.start_trace(os.path.join(tb_dir, "profile_denoise"))
         with jax.named_scope("fused_flux_denoise_loop"):
           latents_jax = self._jitted_fused_denoise_loop(
               params,
@@ -519,7 +521,8 @@ def put_data_on_devices(x, sharding):
 
     do_prof_vae = profile_target in ("all", "vae")
     if do_prof_vae:
-      jax.profiler.start_trace("/tmp/xprof_vae")
+      tb_dir = getattr(self._config, "tensorboard_dir", "/tmp")
+      jax.profiler.start_trace(os.path.join(tb_dir, "profile_vae"))
     with jax.named_scope("vae_decoder"):
       decoded_out = self._jitted_vae_decode(vae_params, latents_jax, vae_bn_mean_jax, vae_bn_std_jax, height, width)
     images_rgb = decoded_out.sample