Add memory allocation timeline in TensorBoard

PiperOrigin-RevId: 628155729
tensorflow · Apr 25, 2024 · d67c004 · d67c004
1 parent dfca241
commit d67c004
Show file tree

Hide file tree

Showing 7 changed files with 30 additions and 6 deletions.
diff --git a/frontend/app/components/main_page/main_page_module.ts b/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},

diff --git a/frontend/app/components/memory_viewer/memory_usage/memory_usage.ts b/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|null, currentHost: string|null) {
     this.nColor = 0;
 
     this.peakHeapSizeBytes = 0;
@@ -76,20 +77,24 @@ 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|null,
+      currentHost: string|null) {
     // Copy the fields from preprocessed result.
     this.moduleName = preprocess.moduleName || '';
     this.timelineUrl = preprocess.allocationTimeline || '';
     if (!this.timelineUrl.startsWith('/memory_viewer.json')) {
-      this.timelineUrl = '';
+      // redirecting memory allocation timeline to this url on TensorBoard
+      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;

diff --git a/frontend/app/components/memory_viewer/memory_viewer.ng.html b/frontend/app/components/memory_viewer/memory_viewer.ng.html
@@ -1,3 +1,5 @@
 <memory-viewer-main
     [memoryViewerPreprocessResult]="memoryViewerPreprocessResult"
+    [currentRun]="currentRun"
+    [currentHost]="currentHost"
 ></memory-viewer-main>
diff --git a/frontend/app/components/memory_viewer/memory_viewer.ts b/frontend/app/components/memory_viewer/memory_viewer.ts
@@ -16,6 +16,8 @@ import {takeUntil} from 'rxjs/operators';
 })
 export class MemoryViewer implements OnDestroy {
   memoryViewerPreprocessResult: MemoryViewerPreprocessResult|null = null;
+  currentRun = '';
+  currentHost = '';
 
   /** Handles on-destroy Subject, used to unsubscribe. */
   private readonly destroyed = new ReplaySubject<void>(1);
@@ -51,6 +53,8 @@ export class MemoryViewer implements OnDestroy {
           if (!data) return;
           this.memoryViewerPreprocessResult =
               data as MemoryViewerPreprocessResult | null;
+          this.currentRun = event.run || '';
+          this.currentHost = event.host || '';
         });
   }
 

diff --git a/frontend/app/components/memory_viewer/memory_viewer_main/memory_viewer_main.ts b/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;
     }

diff --git a/plugin/tensorboard_plugin_profile/convert/raw_to_tool_data.py b/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)

diff --git a/plugin/tensorboard_plugin_profile/profile_plugin.py b/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',