feat: select optimization policy based on extension

tenstorrent · Nov 1, 2024 · 5c52a3e · 5c52a3e
1 parent 63d489e
commit 5c52a3e
Show file tree

Hide file tree

Showing 3 changed files with 9 additions and 7 deletions.
diff --git a/src/ui/src/common/model_loader_service_interface.ts b/src/ui/src/common/model_loader_service_interface.ts
@@ -41,7 +41,7 @@ export interface ModelLoaderServiceInterface {
   get loadedGraphCollections(): WritableSignal<GraphCollection[] | undefined>;
   get models(): WritableSignal<ModelItem[]>;
   get changesToUpload(): WritableSignal<ChangesPerGraphAndNode>;
-  get optimizationPolicies(): string[];
+  getOptimizationPolicies(extensionId: string): string[];
   get selectedOptimizationPolicy(): WritableSignal<string>;
   get graphErrors(): WritableSignal<string[] | undefined>;
   get hasChangesToUpload(): boolean;

diff --git a/src/ui/src/components/visualizer/graph_selector.ts b/src/ui/src/components/visualizer/graph_selector.ts
@@ -446,11 +446,13 @@ export class GraphSelector {
   }
 
   get selectedOptimizationPolicy(): string {
-    return this.modelLoaderService.selectedOptimizationPolicy() || (this.modelLoaderService.optimizationPolicies[0] ?? 'Default');
+    const curExtensionId = this.getCurrentGraphInformation().models[0].selectedAdapter?.id ?? '';
+    return this.modelLoaderService.selectedOptimizationPolicy() || (this.modelLoaderService.getOptimizationPolicies(curExtensionId)[0] ?? 'Default');
   }
 
   get optimizationPolicies(): string[] {
-    return this.modelLoaderService.optimizationPolicies;
+    const curExtensionId = this.getCurrentGraphInformation().models[0].selectedAdapter?.id ?? '';
+    return this.modelLoaderService.getOptimizationPolicies(curExtensionId);
   }
 
   private getLabelWidth(label: string, fontSize = 12, bold = false): number {

diff --git a/src/ui/src/services/model_loader_service.ts b/src/ui/src/services/model_loader_service.ts
@@ -82,14 +82,14 @@ export class ModelLoaderService implements ModelLoaderServiceInterface {
     readonly extensionService: ExtensionService,
   ) {}
 
-  get optimizationPolicies(): string[] {
-    return [...this.extensionService.extensionSettings.values()].flatMap(({ optimizationPolicies }) => optimizationPolicies ?? []);
-  }
-
   get hasChangesToUpload() {
     return Object.keys(this.changesToUpload()).length > 0;
   }
 
+  getOptimizationPolicies(extensionId: string): string[] {
+    return this.extensionService.extensionSettings.get(extensionId)?.optimizationPolicies ?? [];
+  }
+
   async executeModel(modelItem: ModelItem) {
     modelItem.status.set(ModelItemStatus.PROCESSING);
     let updatedPath = modelItem.path;