[PyTorch] Add user_metadata display to memory visualizer (#165939)

Summary: Enhanced the PyTorch CUDA memory visualizer to display user_metadata alongside stack frames when inspecting allocations. The user_metadata field is now shown in all views (Allocator State History, Active Memory Timeline, etc.) with consistent formatting. The implementation handles both string and object metadata types, displaying strings directly and objects as key-value pairs.

Test Plan:
1. Generate a memory snapshot with user_metadata
2. Open the memory visualizer in a browser
3. Load the snapshot file
4. Verify user_metadata appears
5. Test with both string metadata ("testing") and object metadata ({"key": "value"})
6. Verify formatting shows "User Metadata:\n  <value>" for strings

 {F1982860439}

Differential Revision: D85095152

Pull Request resolved: https://github.com/pytorch/pytorch/pull/165939
Approved by: https://github.com/yushangdi

torch/utils/viz/MemoryViz.js

@@ -33,12 +33,12 @@ function version_space() {
   };
 }
 
-function Segment(addr, size, stream, frames, version) {
+function Segment(addr, size, stream, frames, version, user_metadata) {
-  return {addr, size, stream, version, frames};
+  return {addr, size, stream, version, frames, user_metadata};
 }
 
-function Block(addr, size, requested_size, frames, free_requested, version) {
+function Block(addr, size, requested_size, frames, free_requested, version, user_metadata) {
-  return {addr, size, requested_size, frames, free_requested, version};
+  return {addr, size, requested_size, frames, free_requested, version, user_metadata};
 }
 
 function EventSelector(outer, events, stack_info, memory_view) {
@@ -140,7 +140,9 @@ function eventStack(e, allocated, reserved) {
       reserved,
     )} reserved)\n${event}`;
   }
-  return event + '\n' + format_frames(e.frames);
+  const user_metadata_str = format_user_metadata(e.user_metadata);
+  const frames_str = format_frames(e.frames);
+  return event + '\n' + (user_metadata_str ? user_metadata_str + '\n' : '') + frames_str;
 }
 
 function hashCode(num) {
@@ -216,6 +218,7 @@ function MemoryView(outer, stack_info, snapshot, device) {
         seg.stream,
         seg.frames || [],
         seg.version,
+        seg.user_metadata,
       ),
     );
     for (const b of seg.blocks) {
@@ -229,6 +232,7 @@ function MemoryView(outer, stack_info, snapshot, device) {
         b.frames,
         b.state === 'active_pending_free',
         b.version,
+        b.user_metadata,
       );
     }
   }
@@ -307,6 +311,7 @@ function MemoryView(outer, stack_info, snapshot, device) {
             event.frames,
             false,
             event.version,
+            event.user_metadata,
           );
           break;
         case 'free_requested':
@@ -320,6 +325,7 @@ function MemoryView(outer, stack_info, snapshot, device) {
             event.frames,
             true,
             event.version,
+            event.user_metadata,
           );
           break;
         case 'alloc':
@@ -335,6 +341,7 @@ function MemoryView(outer, stack_info, snapshot, device) {
               event.stream,
               event.frames,
               event.version,
+              event.user_metadata,
             ),
           );
           break;
@@ -348,6 +355,7 @@ function MemoryView(outer, stack_info, snapshot, device) {
               event.stream,
               event.frames,
               event.version,
+              event.user_metadata,
             ),
           );
           break;
@@ -426,13 +434,17 @@ function MemoryView(outer, stack_info, snapshot, device) {
           if (t.internal_free > 0) {
             internal = ` (${(t.internal_free / free) * 100}% internal)`;
           }
+          const user_metadata_str = format_user_metadata(t.user_metadata);
+          const frames_str = format_frames(t.frames);
           return (
             `s${t.addr.toString(16)}_${t.version}: segment ${formatSize(
               t.size,
             )} allocated, ` +
             `${formatSize(free)} free${internal} (stream ${
               t.stream
-            })\n${format_frames(t.frames)}`
+            })\n` +
+            (user_metadata_str ? user_metadata_str + '\n' : '') +
+            frames_str
           );
         },
         d => {
@@ -493,12 +505,15 @@ function MemoryView(outer, stack_info, snapshot, device) {
           if (t.free_requested) {
             requested = ' (block freed but waiting due to record_stream)';
           }
+          const user_metadata_str = format_user_metadata(t.user_metadata);
+          const frames_str = format_frames(t.frames);
           return (
             `b${t.addr.toString(16)}_${t.version} ` +
             `${formatSize(t.requested_size)} allocation${requested} (stream ${
               t.segment.stream
             })\n` +
-            format_frames(t.frames)
+            (user_metadata_str ? user_metadata_str + '\n' : '') +
+            frames_str
           );
         },
         removeStroke,
@@ -524,12 +539,15 @@ function MemoryView(outer, stack_info, snapshot, device) {
         d => {
           addStroke(d);
           const t = d.datum();
+          const user_metadata_str = format_user_metadata(t.user_metadata);
+          const frames_str = format_frames(t.frames);
           return (
             `Free space lost due to rounding ${formatSize(
               t.size - t.requested_size,
             )}` +
             ` (stream ${t.segment.stream})\n` +
-            format_frames(t.frames)
+            (user_metadata_str ? user_metadata_str + '\n' : '') +
+            frames_str
           );
         },
         removeStroke,
@@ -760,6 +778,23 @@ function frameFilter({name, filename}) {
   return true;
 }
 
+function format_user_metadata(user_metadata) {
+  if (!user_metadata) {
+    return '';
+  }
+  // Handle string metadata
+  if (typeof user_metadata === 'string') {
+    return `User Metadata:\n  ${user_metadata}`;
+  }
+  // Handle object metadata
+  if (typeof user_metadata === 'object' && Object.keys(user_metadata).length === 0) {
+    return '';
+  }
+  const metadata_lines = Object.entries(user_metadata)
+    .map(([key, value]) => `  ${key}: ${value}`);
+  return 'User Metadata:\n' + metadata_lines.join('\n');
+}
+
 function format_frames(frames) {
   if (frames.length === 0) {
     return (
@@ -992,6 +1027,10 @@ function process_alloc_data(snapshot, device, plot_segments, max_entries) {
       if (!elem.action.includes('alloc')) {
         text = `${text}\nalloc not recorded, stack trace for free:`;
       }
+      const user_metadata_str = format_user_metadata(elem.user_metadata);
+      if (user_metadata_str) {
+        text = `${text}\n${user_metadata_str}`;
+      }
       text = `${text}\n${format_frames(elem.frames)}`;
       return text;
     },