Refactor how worker is initialized

PiperOrigin-RevId: 895415083

Author: Google AI Edge

src/ui/src/components/visualizer/app_service.ts

@@ -128,6 +128,8 @@ export class AppService {
 
   testMode = false;
 
+  onPortal = false;
+
   private groupNodeChildrenCountThresholdFromUrl: string | null = null;
 
   private paneIdToGraph: Record<string, Graph> = {};
@@ -143,7 +145,6 @@ export class AppService {
     private readonly uiStateService: UiStateService,
     private readonly workerService: WorkerService,
   ) {
-    this.listenToWorker();
     this.init();
   }
 
@@ -1106,7 +1107,8 @@ export class AppService {
     this.init();
   }
 
-  private listenToWorker() {
+  setupWorker() {
+    this.workerService.init(this.onPortal);
     this.workerService.worker.addEventListener('message', (event) => {
       const workerEvent = event.data as WorkerEvent;
       switch (workerEvent.eventType) {

src/ui/src/components/visualizer/model_graph_visualizer.ts

@@ -107,6 +107,9 @@ export class ModelGraphVisualizer implements OnInit, OnDestroy, OnChanges {
   /** The sources (file paths) of node data. */
   @Input() nodeDataSources: string[] = [];
 
+  /** Whether the visualizer is running on AI Edge Portal. */
+  @Input() onPortal = false;
+
   /** Triggered when the title is clicked. */
   @Output() readonly titleClicked = new EventEmitter<void>();
 
@@ -253,6 +256,8 @@ export class ModelGraphVisualizer implements OnInit, OnDestroy, OnChanges {
   }
 
   ngOnInit() {
+    this.appService.onPortal = this.onPortal;
+    this.appService.setupWorker();
     this.visualizerThemeService.init(this.el);
 
     this.appService.config.set(this.config || {});

src/ui/src/components/visualizer/title_bar.ts

@@ -108,6 +108,10 @@ export class TitleBar {
       });
   }
 
+  get onPortal(): boolean {
+    return this.appService.onPortal;
+  }
+
   private async runNdpExtension(
     extension: NodeDataProviderExtension,
     runName: string,

src/ui/src/components/visualizer/worker_service.ts

@@ -16,17 +16,24 @@
  * ==============================================================================
  */
 
-import {Injectable} from '@angular/core';
+import {Injectable, OnDestroy} from '@angular/core';
 
 /**
  * Service to manage web worker.
  */
 @Injectable()
-export class WorkerService {
-  readonly worker: Worker;
+export class WorkerService implements OnDestroy {
+  worker!: Worker;
 
-  constructor() {
+  init(onPortal: boolean) {
+    if (this.worker) {
+      this.worker.terminate();
+    }
     this.worker = new Worker(
         new URL('./worker/worker', import.meta.url), {type: 'classic'});
   }
+
+  ngOnDestroy() {
+    this.worker.terminate();
+  }
 }