update gcs bucket upload logic to ensure there's a local copy if it fails

Author: eltsai

src/maxdiffusion/max_utils.py

@@ -87,13 +87,36 @@ def l2norm_pytree(x):
 
 def activate_profiler(config):
   if jax.process_index() == 0 and config.enable_profiler:
-    jax.profiler.start_trace(config.tensorboard_dir)
+    trace_dir = config.tensorboard_dir
+    if trace_dir.startswith("gs://"):
+      trace_dir = "/tmp/profiler_traces"
+    os.makedirs(trace_dir, exist_ok=True)
+    max_logging.log(f"Starting profiler trace in: {trace_dir}")
+    jax.profiler.start_trace(trace_dir)
 
 
 def deactivate_profiler(config):
   if jax.process_index() == 0 and config.enable_profiler:
     jax.profiler.stop_trace()
 
+    trace_dir = config.tensorboard_dir
+    if trace_dir.startswith("gs://"):
+      local_dir = "/tmp/profiler_traces"
+      if os.path.exists(local_dir):
+        max_logging.log(f"Uploading profiler traces from {local_dir} to {trace_dir}...")
+        client = storage.Client()
+        bucket_name, prefix = parse_gcs_bucket_and_prefix(trace_dir)
+        bucket = client.bucket(bucket_name)
+
+        for root, _, files in os.walk(local_dir):
+          for file in files:
+            local_file = os.path.join(root, file)
+            rel_path = os.path.relpath(local_file, local_dir)
+            blob_name = os.path.join(prefix, rel_path)
+            blob = bucket.blob(blob_name)
+            blob.upload_from_filename(local_file)
+            max_logging.log(f"Uploaded {local_file} to gs://{bucket_name}/{blob_name}")
+
 
 def initialize_summary_writer(config):
   return writer.SummaryWriter(config.tensorboard_dir) if jax.process_index() == 0 else None