Add select / de-select all option and add a submit button to the hosts sidenav for trace_viewer

PiperOrigin-RevId: 820313088

Author: muditgokhale2

frontend/app/common/interfaces/navigation_event.ts

@@ -5,6 +5,8 @@ export declare interface NavigationEvent {
   run?: string;
   tag?: string;
   host?: string;
+  // Added to support multi-host functionality for trace_viewer.
+  hosts?: string[];
   // Graph Viewer crosslink params
   opName?: string;
   moduleName?: string;

frontend/app/components/sidenav/sidenav.ng.html

@@ -30,7 +30,7 @@
   </mat-form-field>
 </div>
 
-<div class="item-container">
+<div class="item-container" *ngIf="!isMultiHostsEnabled">
   <div [ngClass]="{'mat-subheading-2': true, 'disabled': !hosts.length}">
     Hosts ({{hosts.length}})
   </div>
@@ -43,6 +43,35 @@
   </mat-form-field>
 </div>
 
+<div class="item-container" *ngIf="isMultiHostsEnabled">
+  <div [ngClass]="{'mat-subheading-2': true, 'disabled': !hosts.length}">
+    Hosts ({{selectedHostsInternal.length}})
+  </div>
+  <mat-form-field class="full-width" appearance="outline">
+    <mat-select #multiSelectHost="matSelect" panelClass="multi-host-select-panel" [value]="selectedHostsPending" [disabled]="!hosts.length" (selectionChange)="onHostsSelectionChange($event.value)" multiple>
+
+      <div class="select-panel-content">
+        <div class="select-all-button-container" (click)="$event.stopPropagation()">
+          <button mat-button class="full-width" (click)="toggleAllHosts()">
+            <span *ngIf="hosts.length === selectedHostsPending.length">Deselect All</span>
+            <span *ngIf="hosts.length !== selectedHostsPending.length">Select All</span>
+          </button>
+          <mat-divider></mat-divider>
+        </div>
+        <mat-option *ngFor="let host of hosts" [value]="host">
+          {{host}}
+        </mat-option>
+
+        <div class="submit-button-container" (click)="$event.stopPropagation()">
+          <button mat-flat-button color="primary" class="full-width" (click)="onSubmitHosts(); multiSelectHost.close()">
+            Apply Host Selection
+          </button>
+        </div>
+      </div>
+    </mat-select>
+  </mat-form-field>
+</div>
+
 <!-- TODO(xprof): Remove module selector from sidenav once it's been moved to the memory viewer control. -->
 <div [hidden]="!is_hlo_tool" class="item-container">
   <div [ngClass]="{'mat-subheading-2': true, 'disabled': !moduleList.length}">
@@ -66,4 +95,3 @@
 </div>
 
 <br>
-

frontend/app/components/sidenav/sidenav.scss

@@ -13,3 +13,50 @@
 .mat-subheading-2 {
   margin-bottom: 0;
 }
+
+// Target the custom panel class defined in sidenav.ng.html
+.multi-host-select-panel {
+  .select-panel-content {
+    display: flex;
+    flex-direction: column;
+  }
+
+  // Container for the Select All/Deselect All button at the top
+  .select-all-button-container {
+    position: sticky;
+    top: 0;
+    z-index: 10;
+    background: #fff;
+    padding: 0 8px;
+    box-shadow: 0 2px 5px rgba(0, 0, 0, 0.05);
+
+    .full-width {
+      width: 100%;
+      text-align: left;
+      padding-left: 0;
+    }
+
+    mat-divider {
+      margin: 0 -8px; // Extend divider to panel edge
+    }
+  }
+
+  // Container for the submit button at the bottom of the dropdown
+  .submit-button-container {
+    position: sticky;
+    bottom: 0;
+    padding: 8px;
+
+    background: #fff;
+    box-shadow: 0 -2px 5px rgba(0, 0, 0, 0.1);
+
+    .full-width {
+      width: 100%;
+    }
+  }
+
+  // Ensure options remain clickable and are not covered by the sticky button's padding.
+  mat-option {
+    flex-shrink: 0;
+  }
+}

frontend/app/components/sidenav/sidenav.ts

@@ -32,8 +32,11 @@ export class SideNav implements OnInit, OnDestroy {
   selectedRunInternal = '';
   selectedTagInternal = '';
   selectedHostInternal = '';
+  selectedHostsInternal: string[] = [];
+  selectedHostsPending: string[] = [];
   selectedModuleInternal = '';
   navigationParams: {[key: string]: string|boolean} = {};
+  multiHostEnabledTools: string[] = ['trace_viewer', 'trace_viewer@'];
 
   hideCaptureProfileButton = false;
 
@@ -65,6 +68,11 @@ export class SideNav implements OnInit, OnDestroy {
     return HLO_TOOLS.includes(this.selectedTag);
   }
 
+  get isMultiHostsEnabled() {
+    const tag = this.selectedTag || '';
+    return this.multiHostEnabledTools.includes(tag);
+  }
+
   // Getter for valid run given url router or user selection.
   get selectedRun() {
     return this.runs.find(validRun => validRun === this.selectedRunInternal) ||
@@ -90,6 +98,10 @@ export class SideNav implements OnInit, OnDestroy {
         this.moduleList[0] || '';
   }
 
+  get selectedHosts() {
+    return this.selectedHostsInternal;
+  }
+
   // https://github.com/angular/angular/issues/11023#issuecomment-752228784
   mergeRouteParams(): Map<string, string> {
     const params = new Map<string, string>();
@@ -119,20 +131,25 @@ export class SideNav implements OnInit, OnDestroy {
     const run = params.get('run') || '';
     const tag = params.get('tool') || params.get('tag') || '';
     const host = params.get('host') || '';
+    const hostsParam = params.get('hosts');
     const opName = params.get('node_name') || params.get('opName') || '';
     const moduleName = params.get('module_name') || '';
     this.navigationParams['firstLoad'] = true;
     if (opName) {
       this.navigationParams['opName'] = opName;
     }
-    if (this.selectedRunInternal === run && this.selectedTagInternal === tag &&
-        this.selectedHostInternal === host) {
-      return;
-    }
     this.selectedRunInternal = run;
     this.selectedTagInternal = tag;
-    this.selectedHostInternal = host;
     this.selectedModuleInternal = moduleName;
+
+    if (this.isMultiHostsEnabled) {
+      if (hostsParam) {
+        this.selectedHostsInternal = hostsParam.split(',');
+      }
+      this.selectedHostsPending = [...this.selectedHostsInternal];
+    } else {
+      this.selectedHostInternal = host;
+    }
     this.update();
   }
 
@@ -153,9 +170,13 @@ export class SideNav implements OnInit, OnDestroy {
     const navigationEvent: NavigationEvent = {
       run: this.selectedRun,
       tag: this.selectedTag,
-      host: this.selectedHost,
       ...this.navigationParams,
     };
+    if (this.isMultiHostsEnabled) {
+      navigationEvent.hosts = this.selectedHosts;
+    } else {
+      navigationEvent.host = this.selectedHost;
+    }
     if (this.is_hlo_tool) {
       navigationEvent.moduleName = this.selectedModule;
     }
@@ -242,8 +263,21 @@ export class SideNav implements OnInit, OnDestroy {
     this.afterUpdateTag();
   }
 
-  onTagSelectionChange(tag: string) {
+  async onTagSelectionChange(tag: string) {
     this.selectedTagInternal = tag;
+    this.selectedHostsInternal = [];
+    this.selectedHostsPending = [];
+    this.selectedHostInternal = '';
+
+    if (this.isMultiHostsEnabled) {
+      this.hosts = await this.getHostsForSelectedTag();
+      if (this.hosts.length > 0) {
+        this.selectedHostsInternal = [this.hosts[0]];
+      } else {
+        this.selectedHostsInternal = [];
+      }
+      this.selectedHostsPending = [...this.selectedHostsInternal];
+    }
     this.afterUpdateTag();
   }
 
@@ -255,18 +289,51 @@ export class SideNav implements OnInit, OnDestroy {
   // Keep them under the same update function as initial step of the separation.
   async updateHosts() {
     this.hosts = await this.getHostsForSelectedTag();
+    if (this.isMultiHostsEnabled) {
+      if (this.selectedHostsInternal.length === 0 && this.hosts.length > 0) {
+        this.selectedHostsInternal = [this.hosts[0]];
+      }
+      this.selectedHostsPending = [...this.selectedHostsInternal];
+    } else {
+      if (!this.selectedHostInternal && this.hosts.length > 0) {
+        this.selectedHostInternal = this.hosts[0];
+      }
+    }
     if (this.is_hlo_tool) {
       this.moduleList = await this.getModuleListForSelectedTag();
     }
 
     this.afterUpdateHost();
   }
 
-  onHostSelectionChange(host: string) {
-    this.selectedHostInternal = host;
+  onHostSelectionChange(selection: string) {
+    this.selectedHostInternal = selection;
+    this.navigateTools();
+  }
+
+  onHostsSelectionChange(selection: string[]) {
+    this.selectedHostsPending =
+        Array.isArray(selection) ? selection : [selection];
+  }
+
+  onSubmitHosts() {
+    this.selectedHostsInternal = [...this.selectedHostsPending];
     this.navigateTools();
   }
 
+  toggleAllHosts() {
+    const allAvailableHosts = this.hosts;
+
+    const areAllSelected = allAvailableHosts.length > 0 &&
+        allAvailableHosts.length === this.selectedHostsPending.length;
+
+    if (areAllSelected) {
+      this.selectedHostsPending = [];
+    } else {
+      this.selectedHostsPending = [...allAvailableHosts];
+    }
+  }
+
   onModuleSelectionChange(module: string) {
     this.selectedModuleInternal = module;
     this.navigateTools();
@@ -276,26 +343,65 @@ export class SideNav implements OnInit, OnDestroy {
     this.navigateTools();
   }
 
+  // Helper function to serialize query parameters
+  private serializeQueryParams(
+      params: {[key: string]: string|string[]|boolean|undefined}): string {
+    const searchParams = new URLSearchParams();
+    for (const key in params) {
+      if (params.hasOwnProperty(key)) {
+        const value = params[key];
+        // Only include non-null/non-undefined values
+        if (value !== undefined && value !== null) {
+          if (Array.isArray(value)) {
+            // Arrays are handled as comma-separated strings (like 'hosts')
+            searchParams.set(key, value.join(','));
+          } else if (typeof value === 'boolean') {
+            // Only set boolean flags if they are explicitly true
+            if (value === true) {
+              searchParams.set(key, 'true');
+            }
+          } else {
+            searchParams.set(key, String(value));
+          }
+        }
+      }
+    }
+    const queryString = searchParams.toString();
+    return queryString ? `?${queryString}` : '';
+  }
+
   updateUrlHistory() {
-    // TODO(xprof): change to camel case when constructing url
-    const toolQueryParams = Object.keys(this.navigationParams)
-                                .map(key => {
-                                  return `${key}=${this.navigationParams[key]}`;
-                                })
-                                .join('&');
-    const toolQueryParamsString =
-        toolQueryParams.length ? `&${toolQueryParams}` : '';
-    const moduleNameQuery =
-        this.is_hlo_tool ? `&module_name=${this.selectedModule}` : '';
-    const url = `${window.parent.location.origin}?tool=${
-        this.selectedTag}&host=${this.selectedHost}&run=${this.selectedRun}${
-        toolQueryParamsString}${moduleNameQuery}#profile`;
+    const navigationEvent = this.getNavigationEvent();
+    const queryParams: {[key: string]: string|string[]|boolean|
+                        undefined} = {...navigationEvent};
+
+    if (this.isMultiHostsEnabled) {
+      // For multi-host enabled tools, ensure 'hosts' is a comma-separated string in the URL
+      if (queryParams['hosts'] && Array.isArray(queryParams['hosts'])) {
+        queryParams['hosts'] = (queryParams['hosts'] as string[]).join(',');
+      }
+      delete queryParams['host'];  // Remove single host param
+    } else {
+      // For other tools, ensure 'host' is used
+      delete queryParams['hosts'];  // Remove multi-host param
+    }
+
+    // Get current path to avoid changing the base URL
+    const pathname = window.parent.location.pathname;
+
+    // Use the custom serialization helper
+    const queryString = this.serializeQueryParams(queryParams);
+    const url = pathname + queryString;
+
     window.parent.history.pushState({}, '', url);
   }
 
   navigateTools() {
     const navigationEvent = this.getNavigationEvent();
     this.communicationService.onNavigateReady(navigationEvent);
+
+    // This router.navigate call remains, as it's responsible for Angular
+    // routing
     this.router.navigate(
         [
           this.selectedTag || 'empty',

frontend/app/components/trace_viewer/trace_viewer.ts

@@ -1,5 +1,4 @@
 import {PlatformLocation} from '@angular/common';
-import {HttpParams} from '@angular/common/http';
 import {Component, inject, Injector, OnDestroy} from '@angular/core';
 import {ActivatedRoute} from '@angular/router';
 import {API_PREFIX, DATA_API, PLUGIN_NAME} from 'org_xprof/frontend/app/common/constants/constants';
@@ -38,11 +37,19 @@ export class TraceViewer implements OnDestroy {
 
   update(event: NavigationEvent) {
     const isStreaming = (event.tag === 'trace_viewer@');
-    const params = new HttpParams()
-                       .set('run', event.run!)
-                       .set('tag', event.tag!)
-                       .set('host', event.host!);
-    const traceDataUrl = this.pathPrefix + DATA_API + '?' + params.toString();
+    const run = event.run || '';
+    const tag = event.tag || '';
+
+    let queryString = `run=${run}&tag=${tag}`;
+
+    if (event.hosts && typeof event.hosts === 'string') {
+      // Since event.hosts is a comma-separated string, we can use it directly.
+      queryString += `&hosts=${event.hosts}`;
+    } else if (event.host) {
+      queryString += `&host=${event.host}`;
+    }
+
+    const traceDataUrl = `${this.pathPrefix}${DATA_API}?${queryString}`;
     this.url = this.pathPrefix + API_PREFIX + PLUGIN_NAME +
         '/trace_viewer_index.html' +
         '?is_streaming=' + isStreaming.toString() + '&is_oss=true' +

plugin/xprof/convert/raw_to_tool_data.py

@@ -116,10 +116,9 @@ def xspace_to_tool_data(
     if success:
       data = process_raw_trace(raw_data)
   elif tool == 'trace_viewer@':
-    # Streaming trace viewer handles one host at a time.
-    assert len(xspace_paths) == 1
     options = params.get('trace_viewer_options', {})
     options['use_saved_result'] = params.get('use_saved_result', True)
+    options['hosts'] = params.get('hosts', [])
     raw_data, success = xspace_wrapper_func(xspace_paths, tool, options)
     if success:
       data = raw_data