Add memory allocation timeline in Tensorboard

PiperOrigin-RevId: 628155729

frontend/app/components/main_page/main_page_module.ts

@@ -48,6 +48,7 @@ export const routes: Routes = [
   {path: 'memory_profile^', component: MemoryProfile},
   {path: 'memory_viewer', component: MemoryViewer},
   {path: 'memory_viewer^', component: MemoryViewer},
+  {path: 'allocation_timeline^', component: MemoryViewer},
   {path: 'op_profile', component: OpProfile},
   {path: 'op_profile^', component: OpProfile},
   {path: 'pod_viewer', component: PodViewer},

frontend/app/components/memory_viewer/memory_usage/memory_usage.ts

@@ -41,7 +41,8 @@ export class MemoryUsage {
 
   // Only one of hloProto or preprocess is valid to construct MemoryUsage.
   constructor(
-      preprocess: MemoryViewerPreprocessResult|null, memorySpaceColor: number) {
+      preprocess: MemoryViewerPreprocessResult|null, memorySpaceColor: number,
+      currentRun: string, currentHost: string) {
     this.nColor = 0;
 
     this.peakHeapSizeBytes = 0;
@@ -76,21 +77,21 @@ export class MemoryUsage {
 
     if (preprocess) {
       // Initialize memory viewer from preprocessed data.
-      this.initMemoryUsageFromPrecomputed(preprocess);
+      this.initMemoryUsageFromPrecomputed(preprocess, currentRun, currentHost);
     }
   }
 
   /**
    * Initializes memory usage from precomputed results.
    */
-  private initMemoryUsageFromPrecomputed(preprocess:
-                                             MemoryViewerPreprocessResult) {
+  private initMemoryUsageFromPrecomputed(
+      preprocess: MemoryViewerPreprocessResult, currentRun: string,
+      currentHost: string) {
     // Copy the fields from preprocessed result.
     this.moduleName = preprocess.moduleName || '';
-    this.timelineUrl = preprocess.allocationTimeline || '';
-    if (!this.timelineUrl.startsWith('/memory_viewer.json')) {
-      this.timelineUrl = '';
-    }
+    this.timelineUrl =
+        `${window.parent.location.origin}/data/plugin/profile/data?run=${
+            currentRun}&tag=allocation_timeline%5E&host=${currentHost}`;
     this.peakHeapSizeBytes =
         (preprocess.totalBufferAllocationMib || 0) * 1024 * 1024;
     this.paddingOverhead = (preprocess.peakHeapMib || 0) * 1024 * 1024 -

frontend/app/components/memory_viewer/memory_viewer.ng.html

@@ -1,3 +1,9 @@
+<div class="allocation-timeline">
+    {{allocationTimeline}}
+  </div>
+
 <memory-viewer-main
     [memoryViewerPreprocessResult]="memoryViewerPreprocessResult"
+    [currentRun]="currentRun"
+    [currentHost]="currentHost"
 ></memory-viewer-main>

frontend/app/components/memory_viewer/memory_viewer.ts

@@ -16,6 +16,9 @@ import {takeUntil} from 'rxjs/operators';
 })
 export class MemoryViewer implements OnDestroy {
   memoryViewerPreprocessResult: MemoryViewerPreprocessResult|null = null;
+  currentRun = '';
+  currentHost = '';
+  allocationTimeline: string|null = null;
 
   /** Handles on-destroy Subject, used to unsubscribe. */
   private readonly destroyed = new ReplaySubject<void>(1);
@@ -49,8 +52,14 @@ export class MemoryViewer implements OnDestroy {
             }
           }));
           if (!data) return;
-          this.memoryViewerPreprocessResult =
-              data as MemoryViewerPreprocessResult | null;
+          if (event.tag === 'allocation_timeline') {
+            this.allocationTimeline = data as string;
+          } else {
+            this.memoryViewerPreprocessResult =
+                data as MemoryViewerPreprocessResult | null;
+            this.currentRun = event.run || '';
+            this.currentHost = event.host || '';
+          }
         });
   }
 

frontend/app/components/memory_viewer/memory_viewer_main/memory_viewer_main.ts

@@ -27,6 +27,10 @@ export class MemoryViewerMain implements OnDestroy, OnChanges {
   /** XLA memory space color */
   @Input() memorySpaceColor: number = 0;
 
+  /** Current run and host name */
+  @Input() currentRun = '';
+  @Input() currentHost = '';
+
   moduleName = '';
   peakInfo?: BufferAllocationInfo;
   activeInfo?: BufferAllocationInfo;
@@ -138,7 +142,8 @@ export class MemoryViewerMain implements OnDestroy, OnChanges {
 
   update() {
     this.usage = new MemoryUsage(
-        this.memoryViewerPreprocessResult, this.memorySpaceColor);
+        this.memoryViewerPreprocessResult, this.memorySpaceColor,
+        this.currentRun, this.currentHost);
     if (this.usage.diagnostics.errors.length > 0) {
       return;
     }

plugin/tensorboard_plugin_profile/convert/raw_to_tool_data.py

@@ -186,6 +186,12 @@ def xspace_to_tool_data(
     raw_data, success = xspace_wrapper_func(xspace_paths, tool, options)
     if success:
       data = raw_data
+  elif tool == 'allocation_timeline':
+    options = {'module_name': params.get('host')}
+    raw_data, success = xspace_wrapper_func(xspace_paths, tool, options)
+    if success:
+      data = raw_data
+      content_type = 'text/html'
   elif tool == 'dcn_collective_stats':
     options = {'host_name': params.get('host')}
     raw_data, success = xspace_wrapper_func(xspace_paths, tool, options)

plugin/tensorboard_plugin_profile/profile_plugin.py

@@ -72,6 +72,7 @@
     'overview_page': 'overview_page.json',
     'overview_page@': 'overview_page.pb',
     'memory_viewer': 'memory_viewer.json',
+    'allocation_timeline': 'allocation_timeline.json',
     'pod_viewer': 'pod_viewer.json',
     'framework_op_stats': 'tensorflow_stats.pb',
     'kernel_stats': 'kernel_stats.pb',