[OPENVINO CODE] add device options on Frontend and Backend for choosing a preferred device while running a model

modules/openvino_code/package.json

@@ -1,7 +1,7 @@
 {
   "publisher": "OpenVINO",
   "name": "openvino-code-completion",
-  "version": "0.0.11",
+  "version": "0.0.12",
   "displayName": "OpenVINO Code Completion",
   "description": "VSCode extension for AI code completion with OpenVINO",
   "icon": "media/logo.png",
@@ -199,6 +199,17 @@
             ],
             "description": "Which model to use for code generation."
           },
+          "openvinoCode.device": {
+            "order": 0,
+            "type": "string",
+            "default": "GPU",
+            "enum": [
+              "CPU",
+              "GPU",
+              "NPU"
+            ],
+            "description": "Which device to use for code generation."
+          },
           "openvinoCode.serverUrl": {
             "order": 1,
             "type": "string",

modules/openvino_code/server/main.py

@@ -1,8 +1,12 @@
 from src.utils import get_parser, setup_logger
+import logging
+
 
 
 # Logger should be set up before other imports to propagate logging config to other packages
 setup_logger()
+logger = logging.getLogger("")
+
 
 import uvicorn  # noqa: E402
 
@@ -14,6 +18,9 @@ def main():
     args = get_parser().parse_args()
 
     # temporary solution for cli args passing
+
+    logger.error(args.model)
+    logger.error(args.device)
     generator_dependency = get_generator_dependency(args.model, args.device, args.tokenizer_checkpoint, args.assistant)
     app.dependency_overrides[get_generator_dummy] = generator_dependency
 

modules/openvino_code/server/src/generators.py

@@ -86,7 +86,7 @@ class OVGenerator(GeneratorFunctor):
     def __init__(
         self,
         checkpoint: str,
-        device: str = "CPU",
+        device: str = "GPU",
         tokenizer_checkpoint: Optional[str] = None,
         assistant_checkpoint: Optional[str] = None,
         summarize_stop_tokens: Optional[Container[str]] = SUMMARIZE_STOP_TOKENS,

modules/openvino_code/shared/side-panel-message.ts

@@ -10,6 +10,7 @@ export enum SidePanelMessageTypes {
   GENERATE_COMPLETION_CLICK = `${sidePanelMessagePrefix}.generateCompletionClick`,
   SETTINGS_CLICK = `${sidePanelMessagePrefix}.settingsClick`,
   MODEL_CHANGE = `${sidePanelMessagePrefix}.modelChange`,
+  DEVICE_CHANGE = `${sidePanelMessagePrefix}.deviceChange`,
 }
 
 export interface ISidePanelMessage<P = unknown> {

modules/openvino_code/side-panel-ui/src/components/sections/ServerSection/ServerSection.tsx

@@ -7,6 +7,9 @@
 import './ServerSection.css';
 import { ModelSelect } from './ModelSelect/ModelSelect';
 import { ModelName } from '@shared/model';
+import { DeviceSelect } from './DeviceSelect/DeviceSelect';
+import { DeviceName } from '@shared/device';
+
 
 interface ServerSectionProps {
   state: IExtensionState | null;
@@ -46,6 +49,15 @@
     });
   };
 
+  const handleDeviceChange = (deviceName: DeviceName) => {
+    vscode.postMessage({
+      type: SidePanelMessageTypes.DEVICE_CHANGE,
+      payload: {
+        deviceName,
+      },
+    });
+  };
+
   if (!state) {
     return <>Extension state is not available</>;
   }
@@ -64,6 +76,13 @@
         supportedFeatures={state.features.supportedList}
         serverStatus={state.server.status}
       ></ModelSelect>
+      <DeviceSelect
+        disabled={!isServerStopped}
+        onChange={handleDeviceChange}
+        selectedDeviceName={state.config.device}
+        supportedFeatures={state.features.supportedList}
+        serverStatus={state.server.status}
+      ></DeviceSelect>
       {isServerStarting && <StartingStages currentStage={state.server.stage}></StartingStages>}
       <div className="button-group">
         {isServerStopped && <button onClick={handleStartServerClick}>Start Server</button>}

modules/openvino_code/src/configuration.ts

@@ -1,4 +1,5 @@
 import { ModelName } from '@shared/model';
+import { DeviceName } from '@shared/device';
 import { WorkspaceConfiguration, workspace } from 'vscode';
 import { CONFIG_KEY } from './constants';
 
@@ -7,6 +8,7 @@
  */
 export type CustomConfiguration = {
   model: ModelName;
+  device: DeviceName;
   serverUrl: string;
   serverRequestTimeout: number;
   streamInlineCompletion: boolean;

modules/openvino_code/src/python-server/python-server-runner.ts

@@ -13,14 +13,15 @@
 import { MODEL_NAME_TO_ID_MAP, ModelName } from '@shared/model';
 import { extensionState } from '../state';
 import { clearLruCache } from '../lru-cache.decorator';
+import { DEVICE_NAME_TO_ID_MAP, DeviceName } from '@shared/device';
 
 const SERVER_STARTED_STDOUT_ANCHOR = 'OpenVINO Code Server started';
 
 interface ServerHooks {
   onStarted: () => void;
 }
 
-async function runServer(modelName: ModelName, config: PythonServerConfiguration, hooks?: ServerHooks) {
+async function runServer(modelName: ModelName, deviceName: DeviceName, config: PythonServerConfiguration, hooks?: ServerHooks) {
   const { serverDir, proxyEnv, abortSignal, logger } = config;
   logger.info('Starting server...');
 
@@ -40,8 +41,9 @@
   }
 
   const model = MODEL_NAME_TO_ID_MAP[modelName];
+  const device = DEVICE_NAME_TO_ID_MAP[deviceName];
 
-  await spawnCommand(venvPython, ['main.py', '--model', model], {
+  await spawnCommand(venvPython, ['main.py', '--model', model, '--device', device], {
     logger,
     cwd: serverDir,
     abortSignal,
@@ -149,8 +151,9 @@
     this._stateController.setStage(ServerStartingStage.START_SERVER);
 
     const modelName = extensionState.config.model;
+    const deviceName = extensionState.config.device;
 
-    await runServer(modelName, config, {
+    await runServer(modelName, deviceName, config, {
       onStarted: () => {
         this._stateController.setStatus(ServerStatus.STARTED);
         this._stateController.setStage(null);

modules/openvino_code/src/side-panel/side-panel-message-handler.ts

@@ -4,6 +4,8 @@ import { Webview, commands } from 'vscode';
 import { settingsService } from '../settings/settings.service';
 import { COMMANDS } from '../constants';
 import { ModelName } from '@shared/model';
+import { DeviceName } from '@shared/device';
+
 
 type SidePanelMessageHandlerType = (webview: Webview, payload?: ISidePanelMessage['payload']) => void;
 
@@ -12,6 +14,8 @@ const sidePanelMessageHandlers: Record<SidePanelMessageTypes, SidePanelMessageHa
   [SidePanelMessageTypes.SETTINGS_CLICK]: () => settingsService.openSettings(),
   [SidePanelMessageTypes.MODEL_CHANGE]: (_, payload) =>
     settingsService.updateSetting('model', (payload as { modelName: ModelName }).modelName),
+  [SidePanelMessageTypes.DEVICE_CHANGE]: (_, payload) =>
+    settingsService.updateSetting('device', (payload as { deviceName: DeviceName }).deviceName),  
   [SidePanelMessageTypes.START_SERVER_CLICK]: () => void commands.executeCommand(COMMANDS.START_SERVER_NATIVE),
   [SidePanelMessageTypes.STOP_SERVER_CLICK]: () => void commands.executeCommand(COMMANDS.STOP_SERVER_NATIVE),
   [SidePanelMessageTypes.SHOW_SERVER_LOG_CLICK]: () => void commands.executeCommand(COMMANDS.SHOW_SERVER_LOG),