Skip to content

Commit

Permalink
changes for installing kaito via arm (#46)
Browse files Browse the repository at this point in the history
  • Loading branch information
@hsubramanianaks
hsubramanianaks authored Aug 27, 2024
1 parent 5c09a50 commit fbbea78
Show file tree
Hide file tree
Showing 10 changed files with 253 additions and 72 deletions.
15 changes: 7 additions & 8 deletions src/panels/CreateClusterPanel.ts
View file
Original file line number Diff line number Diff line change
Expand Up @@ -2,24 +2,23 @@ import { ContainerServiceClient, KubernetesVersion } from "@azure/arm-containers
import { ResourceGroup as ARMResourceGroup, ResourceManagementClient } from "@azure/arm-resources";
import { RestError } from "@azure/storage-blob";
import { Uri, window } from "vscode";
import { getEnvironment } from "../auth/azureAuth";
import { ReadyAzureSessionProvider } from "../auth/types";
import { getAksClient, getResourceManagementClient } from "../commands/utils/arm";
import { getPortalResourceUrl } from "../commands/utils/env";
import { failed, getErrorMessage } from "../commands/utils/errorable";
import { getResourceGroups } from "../commands/utils/resourceGroups";
import { MessageHandler, MessageSink } from "../webview-contract/messaging";
import {
InitialState,
Preset,
ProgressEventType,
ToVsCodeMsgDef,
ToWebViewMsgDef,
ResourceGroup as WebviewResourceGroup,
} from "../webview-contract/webviewDefinitions/createCluster";
import { BasePanel, PanelDataProvider } from "./BasePanel";
import { ClusterDeploymentBuilder, ClusterSpec } from "./utilities/ClusterSpecCreationBuilder";
import { getPortalResourceUrl } from "../commands/utils/env";
import { TelemetryDefinition } from "../webview-contract/webviewTypes";
import { getResourceGroups } from "../commands/utils/resourceGroups";
import { getAksClient, getResourceManagementClient } from "../commands/utils/arm";
import { getEnvironment } from "../auth/azureAuth";
import { ReadyAzureSessionProvider } from "../auth/types";
import { BasePanel, PanelDataProvider } from "./BasePanel";
import { ClusterDeploymentBuilder, ClusterSpec, Preset } from "./utilities/ClusterSpecCreationBuilder";

export class CreateClusterPanel extends BasePanel<"createCluster"> {
constructor(extensionUri: Uri) {
Expand Down
192 changes: 141 additions & 51 deletions src/panels/KaitoPanel.ts
View file
Original file line number Diff line number Diff line change
@@ -1,13 +1,16 @@
import { FeatureClient } from "@azure/arm-features";
import { ResourceManagementClient } from "@azure/arm-resources";
import { RestError } from "@azure/storage-blob";
import * as vscode from "vscode";
import { ReadyAzureSessionProvider } from "../auth/types";
import { getFeatureClient, getResourceManagementClient } from "../commands/utils/arm";
import { getErrorMessage } from "../commands/utils/errorable";
import { longRunning } from "../commands/utils/host";
import { MessageHandler, MessageSink } from "../webview-contract/messaging";
import { InitialState, ToVsCodeMsgDef, ToWebViewMsgDef } from "../webview-contract/webviewDefinitions/kaito";
import { TelemetryDefinition } from "../webview-contract/webviewTypes";
import { BasePanel, PanelDataProvider } from "./BasePanel";
import { getFeatureClient, getResourceManagementClient } from "../commands/utils/arm";
import { FeatureClient } from "@azure/arm-features";
import { longRunning } from "../commands/utils/host";
import { ResourceManagementClient } from "@azure/arm-resources";
import { ClusterDeploymentBuilder, ClusterSpec, Preset } from "./utilities/ClusterSpecCreationBuilder";

export class KaitoPanel extends BasePanel<"kaito"> {
constructor(extensionUri: vscode.Uri) {
Expand Down Expand Up @@ -107,13 +110,55 @@ export class KaitoPanelDataProvider implements PanelDataProvider<"kaito"> {
});
}
private async handleKaitoInstallation(webview: MessageSink<ToWebViewMsgDef>) {

// register feature
const featureRegister = await longRunning(`Register KAITO Feature.`, () =>
this.featureClient.features.register("Microsoft.ContainerService", "AIToolchainOperatorPreview"),
// const featureRegister = await longRunning(`Register KAITO Feature.`, () =>
// this.featureClient.features.register("Microsoft.ContainerService", "AIToolchainOperatorPreview"),
// );

// if (featureRegister.properties?.state !== "Registered") {
// webview.postKaitoInstallProgressUpdate({
// operationDescription: "Installing Kaito",
// event: 3,
// errorMessage: "Failed to register feature",
// models: [],
// });
// return;
// }

// // Install kaito enablement
// // Get current json
// const currentJson = await longRunning(`Get current json.`, () => {
// return this.resourceManagementClient.resources.getById(this.armId, "2023-08-01");
// });
// console.log(currentJson);

// // Update json
// if (currentJson.properties) {
// currentJson.properties.aiToolchainOperatorProfile = { enabled: true };
// }

// const updateJson = await longRunning(`Update json.`, () => {
// return this.resourceManagementClient.resources.beginCreateOrUpdateByIdAndWait(
// this.armId,
// "2023-08-01",
// currentJson,
// );
// });
// console.log(updateJson);
const subscriptionFeatureRegistrationType = {
properties: {},
};
const options = {
subscriptionFeatureRegistrationType,
};

const featureRegistrationPoller = await this.featureClient.subscriptionFeatureRegistrations.createOrUpdate(
"Microsoft.ContainerService",
"AIToolchainOperatorPreview",
options,
);

if (featureRegister.properties?.state !== "Registered") {
if (featureRegistrationPoller.properties?.state !== "Registered") {
webview.postKaitoInstallProgressUpdate({
operationDescription: "Installing Kaito",
event: 3,
Expand All @@ -123,58 +168,103 @@ export class KaitoPanelDataProvider implements PanelDataProvider<"kaito"> {
return;
}

// Install kaito enablement
// Get current json
const currentJson = await longRunning(`Get current json.`, () => {
return this.resourceManagementClient.resources.getById(this.armId, "2023-08-01");
});
console.log(currentJson);

// Update json
if (currentJson.properties) {
currentJson.properties.aiToolchainOperatorProfile = { enabled: true };
}

const updateJson = await longRunning(`Update json.`, () => {
return this.resourceManagementClient.resources.beginCreateOrUpdateByIdAndWait(this.armId, "2023-08-01", currentJson);
});
console.log(updateJson);

// const kaitoEnablement = await longRunning(`Enable KAITO Feature.`, () =>
// this.resourceManagementClient.deployments.beginCreateOrUpdate(
// this.resourceGroupName,
// "Microsoft.ContainerService",
// "",
// "providers/Microsoft.ContainerService/enableKaito",
// "2021-11-01-preview",
// {},
// ),
// );
const clusterSpec: ClusterSpec = {
location: "eastus2euap", //TODO get location from cluster
name: this.clusterName,
resourceGroupName: this.resourceGroupName,
subscriptionId: this.subscriptionId,
kubernetesVersion: "1.28", // TODO k8s version from cluster
};

// install kaito
webview.postKaitoInstallProgressUpdate({
operationDescription: "Installing Kaito",
event: 1,
errorMessage: null,
models: [],
});
const deploymentName = `${this.clusterName}-${Math.random().toString(36).substring(5)}`;

// simulate kaito installation and success
setTimeout(() => {
const deploymentSpec = new ClusterDeploymentBuilder()
.buildCommonParametersForKaito(clusterSpec)
.buildTemplate(Preset.KaitoAddon)
.getDeployment();
try {
const poller = await this.resourceManagementClient.deployments.beginCreateOrUpdate(
this.resourceGroupName,
deploymentName,
deploymentSpec,
);
// kaito installation in progress
webview.postKaitoInstallProgressUpdate({
operationDescription: "Kaito installed",
event: 4,
errorMessage: null,
models: [
{
family: "family",
modelName: "modelName",
minimumGpu: 1,
kaitoVersion: "v1.0",
modelSource: "modelSource",
},
],
operationDescription: "Installing Kaito",
event: 1,
errorMessage: undefined,
models: [],
});
poller.onProgress((state) => {
if (state.status === "succeeded") {
webview.postKaitoInstallProgressUpdate({
operationDescription: "Installing Kaito succeeded",
event: 4,
errorMessage: undefined,
models: [],
});
} else if (state.status === "failed") {
webview.postKaitoInstallProgressUpdate({
operationDescription: "Installing Kaito failed",
event: 3,
errorMessage: state.error?.message,
models: [],
});
}
});
}, 5000);
} catch (ex) {
const errorMessage = isInvalidTemplateDeploymentError(ex)
? getInvalidTemplateErrorMessage(ex)
: getErrorMessage(ex);
vscode.window.showErrorMessage(`Error installing Kaito addon for ${this.clusterName}: ${errorMessage}`);
webview.postKaitoInstallProgressUpdate({
operationDescription: "Installing Kaito failed",
event: 3,
errorMessage: ex instanceof Error ? ex.message : String(ex),
models: [],
});
}
}
}

function getInvalidTemplateErrorMessage(ex: InvalidTemplateDeploymentRestError): string {
const innerDetails = ex.details.error?.details || [];
if (innerDetails.length > 0) {
const details = innerDetails.map((d) => `${d.code}: ${d.message}`).join("\n");
return `Invalid template:\n${details}`;
}

const innerError = ex.details.error?.message || "";
if (innerError) {
return `Invalid template:\n${innerError}`;
}

return `Invalid template: ${getErrorMessage(ex)}`;
}

type InvalidTemplateDeploymentRestError = RestError & {
details: {
error?: {
code: "InvalidTemplateDeployment";
message?: string;
details?: {
code?: string;
message?: string;
}[];
};
};
};

function isInvalidTemplateDeploymentError(ex: unknown): ex is InvalidTemplateDeploymentRestError {
return isRestError(ex) && ex.code === "InvalidTemplateDeployment";
}

function isRestError(ex: unknown): ex is RestError {
return typeof ex === "object" && ex !== null && ex.constructor.name === "RestError";
}
3 changes: 3 additions & 0 deletions src/panels/templates/DevTestCreateCluster.json
View file
Original file line number Diff line number Diff line change
Expand Up @@ -776,6 +776,9 @@
"networkPolicy": "[parameters('networkPolicy')]",
"serviceCidr": "[parameters('serviceCidr')]",
"dnsServiceIP": "[parameters('dnsServiceIP')]"
},
"AiToolchainOperatorProfile": {
"enabled": true
}
},
"tags": "[parameters('clusterTags')]"
Expand Down
59 changes: 59 additions & 0 deletions src/panels/templates/UpdateClusterKaitoAddon.json
View file
Original file line number Diff line number Diff line change
@@ -0,0 +1,59 @@
{
"$schema": "https://schema.management.azure.com/schemas/2019-04-01/deploymentTemplate.json#",
"contentVersion": "1.0.0.0",
"parameters": {
"subscriptionId": {
"type": "string",
"metadata": {
"description": "Subscription ID"
}
},
"location": {
"type": "string",
"metadata": {
"description": "Location for the resources"
}
},
"resourceName":{
"type": "string",
"metadata": {
"description": "Name of the resource"
}
},

"resourceGroupName":{
"type": "string",
"metadata": {
"description": "Name of the resource group"
}
},

"apiVersion":{
"type": "string",
"metadata": {
"description": "API version"
}
},
"kubernetesVersion": {
"type": "string",
"metadata": {
"description": "Kubernetes version"
}
}
},
"resources": [
{
"type": "Microsoft.ContainerService/managedClusters",
"apiVersion": "[parameters('apiVersion')]",
"name": "[parameters('resourceName')]",
"properties": {
"aiToolchainOperatorProfile": {
"enabled": true
},
"oidcIssuerProfile": {
"enabled": true
}
}
}
]
}
36 changes: 34 additions & 2 deletions src/panels/utilities/ClusterSpecCreationBuilder.ts
View file
Original file line number Diff line number Diff line change
@@ -1,23 +1,29 @@
import { Deployment } from "@azure/arm-resources";
import { Preset } from "../../webview-contract/webviewDefinitions/createCluster";
import devTestTemplate from "../templates/DevTestCreateCluster.json";
import kaitoTemplate from "../templates/UpdateClusterKaitoAddon.json";

export type ClusterSpec = {
location: string;
name: string;
resourceGroupName: string;
subscriptionId: string;
kubernetesVersion: string;
username: string;
username?: string;
};

type TemplateContent = Record<string, unknown>;

const deploymentApiVersion = "2023-08-01";
const presetTemplates: Record<Preset, TemplateContent> = {
dev: devTestTemplate,
kaitoAddon: kaitoTemplate,
};

export enum Preset {
Dev = "dev",
KaitoAddon = "kaitoAddon",
}

export class ClusterDeploymentBuilder {
private deployment: Deployment = {
properties: {
Expand All @@ -27,6 +33,32 @@ export class ClusterDeploymentBuilder {
},
};

public buildCommonParametersForKaito(clusterSpec: ClusterSpec): ClusterDeploymentBuilder {
this.deployment.properties.parameters = {
...this.deployment.properties.parameters,
location: {
value: clusterSpec.location,
},
resourceName: {
value: clusterSpec.name,
},
apiVersion: {
value: deploymentApiVersion,
},
subscriptionId: {
value: clusterSpec.subscriptionId,
},
resourceGroupName: {
value: clusterSpec.resourceGroupName,
},
kubernetesVersion: {
value: clusterSpec.kubernetesVersion,
},
};

return this;
}

public buildCommonParameters(clusterSpec: ClusterSpec): ClusterDeploymentBuilder {
this.deployment.properties.parameters = {
...this.deployment.properties.parameters,
Expand Down
Loading

0 comments on commit fbbea78

Please sign in to comment.