diff --git a/frontend/app/components/main_page/main_page_module.ts b/frontend/app/components/main_page/main_page_module.ts index 6ff34df3b..798d016dd 100644 --- 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 index b6ba40972..c0cabcd1b 100644 --- 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 index 8a7069840..69f4afebe 100644 --- a/frontend/app/components/memory_viewer/memory_viewer.ng.html +++ b/frontend/app/components/memory_viewer/memory_viewer.ng.html @@ -1,3 +1,5 @@ diff --git a/frontend/app/components/memory_viewer/memory_viewer.ts b/frontend/app/components/memory_viewer/memory_viewer.ts index fe4d9174c..12268dcfc 100644 --- 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(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 index 9b9a3e1f1..a6411bee4 100644 --- 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 index d7d750ee7..0190ce940 100644 --- 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 index 41c1b8d87..68c498039 100644 --- 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',