KAITO - Model Management & Chat Features

docs/book/src/features/kaito-integration.md

@@ -11,7 +11,7 @@ Right click on your AKS cluster and select **Deploy an LLM with KAITO** and then
 
 ![Installation Page](../resources/kaito-install-page.png)
 
-Once on the page, click **Install** and the Kaito installation progress will begin. You will be notified when Kaito has been successfully installed. 
+Once on the page, click **Install** and the KAITO installation progress will begin. You will be notified when KAITO has been successfully installed. 
 
 ### Deploy a model
 

package.json

@@ -287,15 +287,15 @@
             {
                 "command": "aks.aksCreateClusterFromCopilot",
                 "title": "Command to create an AKS cluster"
-            },  
+            },
             {
                 "command": "aks.aksDeployManifest",
                 "title": "Deploy application manifest from Github Copilot Chat"
-            },               
-            {            
+            },
+            {
                 "command": "aks.aksOpenKubectlPanel",
                 "title": "Run Kubectl Commands from Github Copilot Chat"
-            },                   
+            },
             {
                 "command": "aks.aksTCPDump",
                 "title": "Collect TCP Dumps"
@@ -316,9 +316,17 @@
                 "command": "aks.aksKaitoCreateCRD",
                 "title": "Create KAITO Workspace"
             },
+            {
+                "command": "aks.aksKaitoManage",
+                "title": "Manage KAITO Models"
+            },
             {
                 "command": "aks.aksKaito",
                 "title": "Install KAITO"
+            },
+            {
+                "command": "aks.aksKaitoTest",
+                "title": "Test KAITO models"
             }
         ],
         "menus": {
@@ -442,6 +450,10 @@
                 {
                     "command": "aks.aksKaitoCreateCRD",
                     "group": "navigation"
+                },
+                {
+                    "command": "aks.aksKaitoManage",
+                    "group": "navigation"
                 }
             ],
             "aks.detectorsSubMenu": [
@@ -546,7 +558,7 @@
             },
             {
                 "id": "aks.kaitoInstallSubMenu",
-                "label": "Deploy a LLM with KAITO (Beta)"
+                "label": "Deploy a LLM with KAITO"
             }
         ]
     },

src/commands/aksKaito/aksKaito.ts

@@ -50,11 +50,7 @@ export default async function aksKaito(_context: IActionContext, target: unknown
         return;
     }
 
-    const clusterName = clusterNode.result.name;
-    const armId = clusterNode.result.armId;
-    const subscriptionId = clusterNode.result.subscriptionId;
-    const resourceGroupName = clusterNode.result.resourceGroupName;
-
+    const { name: clusterName, armId, subscriptionId, resourceGroupName } = clusterNode.result;
     const panel = new KaitoPanel(extension.result.extensionUri);
 
     const filterKaitoPodNames = await filterPodName(

src/commands/aksKaito/aksKaitoCreateCRD.ts

@@ -47,10 +47,9 @@ export default async function aksKaitoCreateCRD(_context: IActionContext, target
         return;
     }
 
-    const clusterName = clusterNode.result.name;
-    const armId = clusterNode.result.armId;
-    const subscriptionId = clusterNode.result.subscriptionId;
-    const resourceGroupName = clusterNode.result.resourceGroupName;
+    const { name: clusterName, armId, subscriptionId, resourceGroupName } = clusterNode.result;
+
+    // "kaito-" is assuming kaito pods are still named kaito-workspace & kaito-gpu-provisioner
     const filterKaitoPodNames = await longRunning(`Checking if KAITO is installed.`, () => {
         return filterPodName(sessionProvider.result, kubectl, subscriptionId, resourceGroupName, clusterName, "kaito-");
     });
@@ -62,17 +61,17 @@ export default async function aksKaitoCreateCRD(_context: IActionContext, target
 
     if (filterKaitoPodNames.result.length === 0) {
         vscode.window.showWarningMessage(
-            `Please install Kaito for cluster ${clusterName}. \n \n Kaito Workspace generation is only enabled when kaito is installed. Skipping generation.`,
+            `Please install Kaito for cluster ${clusterName}. \n \n KAITO Workspace generation is only enabled when KAITO is installed.`,
         );
         return;
     }
+
     const clusterInfo = await getKubernetesClusterInfo(sessionProvider.result, target, cloudExplorer, clusterExplorer);
     if (failed(clusterInfo)) {
         vscode.window.showErrorMessage(clusterInfo.error);
         return;
     }
     const kubeConfigFile = await tmpfile.createTempFile(clusterInfo.result.kubeconfigYaml, "yaml");
-
     const panel = new KaitoModelsPanel(extension.result.extensionUri);
     const dataProvider = new KaitoModelsPanelDataProvider(
         clusterName,
@@ -82,6 +81,7 @@ export default async function aksKaitoCreateCRD(_context: IActionContext, target
         kubectl,
         kubeConfigFile.filePath,
         sessionProvider.result,
+        target,
     );
     panel.show(dataProvider);
 }

src/commands/aksKaito/aksKaitoManage.ts

@@ -0,0 +1,113 @@
+import { IActionContext } from "@microsoft/vscode-azext-utils";
+import * as vscode from "vscode";
+import * as tmpfile from "../utils/tempfile";
+import * as k8s from "vscode-kubernetes-tools-api";
+import { getKubernetesClusterInfo } from "../utils/clusters";
+import { getReadySessionProvider } from "../../auth/azureAuth";
+import { KaitoManagePanelDataProvider } from "../../panels/KaitoManagePanel";
+import { KaitoManagePanel } from "../../panels/KaitoManagePanel";
+import { getAksClusterTreeNode } from "../utils/clusters";
+import { failed } from "../utils/errorable";
+import { getExtension } from "../utils/host";
+import { getConditions, convertAgeToMinutes } from "../../panels/utilities/KaitoHelpers";
+import { invokeKubectlCommand } from "../utils/kubectl";
+import { getKaitoInstallationStatus } from "../utils/kaitoValidationHelpers";
+
+export default async function aksKaitoManage(_context: IActionContext, target: unknown): Promise<void> {
+    const cloudExplorer = await k8s.extension.cloudExplorer.v1;
+    const clusterExplorer = await k8s.extension.clusterExplorer.v1;
+    const kubectl = await k8s.extension.kubectl.v1;
+
+    const sessionProvider = await getReadySessionProvider();
+    if (failed(sessionProvider)) {
+        vscode.window.showErrorMessage(sessionProvider.error);
+        return;
+    }
+
+    const clusterNode = getAksClusterTreeNode(target, cloudExplorer);
+    if (failed(clusterNode)) {
+        vscode.window.showErrorMessage(clusterNode.error);
+        return;
+    }
+
+    const extension = getExtension();
+    if (failed(extension)) {
+        vscode.window.showErrorMessage(extension.error);
+        return;
+    }
+
+    if (!kubectl.available) {
+        vscode.window.showWarningMessage(`Kubectl is unavailable.`);
+        return;
+    }
+
+    if (!cloudExplorer.available) {
+        vscode.window.showWarningMessage(`Cloud explorer is unavailable.`);
+        return;
+    }
+
+    if (!clusterExplorer.available) {
+        vscode.window.showWarningMessage(`Cluster explorer is unavailable.`);
+        return;
+    }
+
+    const clusterInfo = await getKubernetesClusterInfo(sessionProvider.result, target, cloudExplorer, clusterExplorer);
+    if (failed(clusterInfo)) {
+        vscode.window.showErrorMessage(clusterInfo.error);
+        return;
+    }
+
+    const kubeConfigFile = await tmpfile.createTempFile(clusterInfo.result.kubeconfigYaml, "yaml");
+    const { name: clusterName, armId, subscriptionId, resourceGroupName } = clusterNode.result;
+
+    // Returns an object with the status of the kaito pods
+    const kaitoStatus = await getKaitoInstallationStatus(
+        sessionProvider,
+        kubectl,
+        subscriptionId,
+        resourceGroupName,
+        clusterName,
+        clusterInfo,
+    );
+
+    // Only proceed if kaito is installed and both the workspace & gpu-provisioner pods are running
+    if (!kaitoStatus.kaitoInstalled || !kaitoStatus.kaitoWorkspaceReady || !kaitoStatus.kaitoGPUProvisionerReady) {
+        return;
+    }
+
+    // The logic below is to acquire the initial deployment data.
+    const command = `get workspace -o json`;
+    const kubectlresult = await invokeKubectlCommand(kubectl, kubeConfigFile.filePath, command);
+    if (failed(kubectlresult)) {
+        vscode.window.showErrorMessage(`Error retrieving workspaces: ${kubectlresult.error}`);
+        return;
+    }
+    const models = [];
+    const data = JSON.parse(kubectlresult.result.stdout);
+    for (const item of data.items) {
+        const conditions: Array<{ type: string; status: string }> = item.status?.conditions || [];
+        const { resourceReady, inferenceReady, workspaceReady } = getConditions(conditions);
+        // The data below is used to indicate the current progress of the active model deployments
+        models.push({
+            name: item.inference?.preset?.name,
+            instance: item.resource?.instanceType,
+            resourceReady: resourceReady,
+            inferenceReady: inferenceReady,
+            workspaceReady: workspaceReady,
+            age: convertAgeToMinutes(item.metadata?.creationTimestamp),
+        });
+    }
+
+    const panel = new KaitoManagePanel(extension.result.extensionUri);
+    const dataProvider = new KaitoManagePanelDataProvider(
+        clusterName,
+        subscriptionId,
+        resourceGroupName,
+        armId,
+        kubectl,
+        kubeConfigFile.filePath,
+        models,
+        target,
+    );
+    panel.show(dataProvider);
+}

src/commands/aksKaito/aksKaitoTest.ts

@@ -0,0 +1,72 @@
+import { IActionContext } from "@microsoft/vscode-azext-utils";
+import * as vscode from "vscode";
+import * as tmpfile from "../utils/tempfile";
+import * as k8s from "vscode-kubernetes-tools-api";
+import { getKubernetesClusterInfo } from "../utils/clusters";
+import { getReadySessionProvider } from "../../auth/azureAuth";
+import { getAksClusterTreeNode } from "../utils/clusters";
+import { failed } from "../utils/errorable";
+import { getExtension } from "../utils/host";
+import { KaitoTestPanel, KaitoTestPanelDataProvider } from "../../panels/KaitoTestPanel";
+
+export default async function aksKaitoTest(
+    _context: IActionContext,
+    { target, modelName }: { target: unknown; modelName: string },
+): Promise<void> {
+    const cloudExplorer = await k8s.extension.cloudExplorer.v1;
+    if (!cloudExplorer.available) {
+        vscode.window.showWarningMessage(`Cloud explorer is unavailable.`);
+        return;
+    }
+
+    const clusterExplorer = await k8s.extension.clusterExplorer.v1;
+    if (!clusterExplorer.available) {
+        vscode.window.showWarningMessage(`Cluster explorer is unavailable.`);
+        return;
+    }
+
+    const kubectl = await k8s.extension.kubectl.v1;
+    if (!kubectl.available) {
+        vscode.window.showWarningMessage(`Kubectl is unavailable.`);
+        return;
+    }
+
+    const sessionProvider = await getReadySessionProvider();
+    if (failed(sessionProvider)) {
+        vscode.window.showErrorMessage(sessionProvider.error);
+        return;
+    }
+
+    const clusterNode = getAksClusterTreeNode(target, cloudExplorer);
+    if (failed(clusterNode)) {
+        vscode.window.showErrorMessage(clusterNode.error);
+        return;
+    }
+
+    const extension = getExtension();
+    if (failed(extension)) {
+        vscode.window.showErrorMessage(extension.error);
+        return;
+    }
+
+    const { name: clusterName, armId, subscriptionId, resourceGroupName } = clusterNode.result;
+    const clusterInfo = await getKubernetesClusterInfo(sessionProvider.result, target, cloudExplorer, clusterExplorer);
+    if (failed(clusterInfo)) {
+        vscode.window.showErrorMessage(clusterInfo.error);
+        return;
+    }
+    const kubeConfigFile = await tmpfile.createTempFile(clusterInfo.result.kubeconfigYaml, "yaml");
+
+    const panel = new KaitoTestPanel(extension.result.extensionUri);
+    const dataProvider = new KaitoTestPanelDataProvider(
+        clusterName,
+        subscriptionId,
+        resourceGroupName,
+        armId,
+        kubectl,
+        kubeConfigFile.filePath,
+        sessionProvider.result,
+        modelName,
+    );
+    panel.show(dataProvider);
+}

src/commands/aksKaito/akskaitoGenerateYaml.ts

@@ -39,6 +39,7 @@ export default async function aksKaitoGenerateYaml(_context: IActionContext, tar
     const subscriptionId = clusterNode.result.subscriptionId;
     const resourceGroupName = clusterNode.result.resourceGroupName;
 
+    // "kaito-" is assuming kaito pods are still named kaito-workspace & kaito-gpu-provisioner
     const filterKaitoPodNames = await longRunning(`Checking if KAITO is installed.`, () => {
         return filterPodName(sessionProvider.result, kubectl, subscriptionId, resourceGroupName, clusterName, "kaito-");
     });