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',