ML extension - Improving predict parameter mapping experience (#10264)

extensions/machine-learning/src/views/controllerBase.ts

@@ -22,14 +22,14 @@ export abstract class ControllerBase {
 	/**
 	 * Executes an action and sends back callback event to the view
 	 */
-	public async executeAction<T extends ViewBase>(dialog: T, eventName: string, func: (...args: any[]) => Promise<any>, ...args: any[]): Promise<void> {
+	public async executeAction<T extends ViewBase>(dialog: T, eventName: string, inputArgs: any, func: (...args: any[]) => Promise<any>, ...args: any[]): Promise<void> {
 		const callbackEvent = ViewBase.getCallbackEventName(eventName);
 		try {
 			let result = await func(...args);
-			dialog.sendCallbackRequest(callbackEvent, { data: result });
+			dialog.sendCallbackRequest(callbackEvent, { inputArgs: inputArgs, data: result });
 
 		} catch (error) {
-			dialog.sendCallbackRequest(callbackEvent, { error: error });
+			dialog.sendCallbackRequest(callbackEvent, { inputArgs: inputArgs, error: error });
 		}
 	}
 
@@ -39,7 +39,7 @@ export abstract class ControllerBase {
 	 */
 	public registerEvents(view: ViewBase): void {
 		view.on(LocalPathsEventName, async (args) => {
-			await this.executeAction(view, LocalPathsEventName, this.getLocalPaths, this._apiWrapper, args);
+			await this.executeAction(view, LocalPathsEventName, args, this.getLocalPaths, this._apiWrapper, args);
 		});
 	}
 

extensions/machine-learning/src/views/models/azureModelsTable.ts

@@ -65,9 +65,22 @@ export class AzureModelsTable extends ModelViewBase implements IDataComponent<Wo
 								...constants.cssStyles.tableRow
 							},
 						},
+						{ // Framework
+							displayName: constants.modelFramework,
+							ariaLabel: constants.modelFramework,
+							valueType: azdata.DeclarativeDataType.string,
+							isReadOnly: true,
+							width: 100,
+							headerCssStyles: {
+								...constants.cssStyles.tableHeader
+							},
+							rowCssStyles: {
+								...constants.cssStyles.tableRow
+							},
+						},
 						{ // Version
-							displayName: constants.modelVersion,
-							ariaLabel: constants.modelVersion,
+							displayName: constants.modelFrameworkVersion,
+							ariaLabel: constants.modelFrameworkVersion,
 							valueType: azdata.DeclarativeDataType.string,
 							isReadOnly: true,
 							width: 100,
@@ -159,7 +172,7 @@ export class AzureModelsTable extends ModelViewBase implements IDataComponent<Wo
 				selectModelButton = radioButton;
 			}
 
-			return [model.name, model.createdTime, model.frameworkVersion, selectModelButton];
+			return [model.name, model.createdTime, model.framework, model.frameworkVersion, selectModelButton];
 		}
 
 		return [];

extensions/machine-learning/src/views/models/azureResourceFilterComponent.ts

@@ -144,6 +144,9 @@ export class AzureResourceFilterComponent extends ModelViewBase implements IData
 			let values = this._azureAccounts.map(a => { return { displayName: a.displayInfo.displayName, name: a.key.accountId }; });
 			this._accounts.values = values;
 			this._accounts.value = values[0];
+		} else {
+			this._accounts.values = [];
+			this._accounts.value = undefined;
 		}
 		await this.onAccountSelected();
 	}
@@ -161,6 +164,9 @@ export class AzureResourceFilterComponent extends ModelViewBase implements IData
 			let values = this._azureSubscriptions.map(s => { return { displayName: s.name, name: s.id }; });
 			this._subscriptions.values = values;
 			this._subscriptions.value = values[0];
+		} else {
+			this._subscriptions.values = [];
+			this._subscriptions.value = undefined;
 		}
 		await this.onSubscriptionSelected();
 	}
@@ -171,6 +177,9 @@ export class AzureResourceFilterComponent extends ModelViewBase implements IData
 			let values = this._azureGroups.map(s => { return { displayName: s.name, name: s.id }; });
 			this._groups.values = values;
 			this._groups.value = values[0];
+		} else {
+			this._groups.values = [];
+			this._groups.value = undefined;
 		}
 		await this.onGroupSelected();
 	}
@@ -181,6 +190,9 @@ export class AzureResourceFilterComponent extends ModelViewBase implements IData
 			let values = this._azureWorkspaces.map(s => { return { displayName: s.name || '', name: s.id || '' }; });
 			this._workspaces.values = values;
 			this._workspaces.value = values[0];
+		} else {
+			this._workspaces.values = [];
+			this._workspaces.value = undefined;
 		}
 		this.onWorkspaceSelectedChanged();
 	}

extensions/machine-learning/src/views/models/localModelsComponent.ts

@@ -39,10 +39,7 @@ export class LocalModelsComponent extends ModelViewBase implements IDataComponen
 		}).component();
 		this._localBrowse = modelBuilder.button().withProperties({
 			label: constants.browseModels,
-			width: this.browseButtonMaxLength,
-			CSSStyles: {
-				'text-align': 'end'
-			}
+			width: this.browseButtonMaxLength
 		}).component();
 		this._localBrowse.onDidClick(async () => {
 
@@ -65,7 +62,7 @@ export class LocalModelsComponent extends ModelViewBase implements IDataComponen
 			.withLayout({
 				flexFlow: 'row',
 				justifyContent: 'space-between',
-				width: this.componentMaxLength
+				width: this.componentMaxLength + 200
 			}).withItems([
 				this._localPath, this._localBrowse]
 			).component();
@@ -80,9 +77,9 @@ export class LocalModelsComponent extends ModelViewBase implements IDataComponen
 	public addComponents(formBuilder: azdata.FormBuilder) {
 		if (this._flex) {
 			formBuilder.addFormItem({
-				title: '',
+				title: constants.modelLocalSourceTitle,
 				component: this._flex
-			});
+			}, { info: constants.modelLocalSourceTooltip });
 		}
 	}
 

extensions/machine-learning/src/views/models/manageModels/currentModelsComponent.ts

@@ -36,8 +36,8 @@ export class CurrentModelsComponent extends ModelViewBase implements IPageView {
 	 * @param modelBuilder register the components
 	 */
 	public registerComponent(modelBuilder: azdata.ModelBuilder): azdata.Component {
-		this._tableSelectionComponent = new TableSelectionComponent(this._apiWrapper, this, false);
-		this._tableSelectionComponent.registerComponent(modelBuilder);
+		this._tableSelectionComponent = new TableSelectionComponent(this._apiWrapper, this, { editable: false, preSelected: true });
+		this._tableSelectionComponent.registerComponent(modelBuilder, constants.databaseName, constants.tableName);
 		this._tableSelectionComponent.onSelectedChanged(async () => {
 			await this.onTableSelected();
 		});
@@ -110,7 +110,9 @@ export class CurrentModelsComponent extends ModelViewBase implements IPageView {
 		if (this._tableSelectionComponent?.data) {
 			this.importTable = this._tableSelectionComponent?.data;
 			await this.storeImportConfigTable();
-			await this._dataTable?.refresh();
+			if (this._dataTable) {
+				await this._dataTable.refresh();
+			}
 		}
 	}
 

extensions/machine-learning/src/views/models/manageModels/currentModelsTable.ts

@@ -66,6 +66,32 @@ export class CurrentModelsTable extends ModelViewBase implements IDataComponent<
 					...constants.cssStyles.tableRow
 				},
 			},
+			{ // Version
+				displayName: constants.modelVersion,
+				ariaLabel: constants.modelVersion,
+				valueType: azdata.DeclarativeDataType.string,
+				isReadOnly: true,
+				width: 150,
+				headerCssStyles: {
+					...constants.cssStyles.tableHeader
+				},
+				rowCssStyles: {
+					...constants.cssStyles.tableRow
+				},
+			},
+			{ // Format
+				displayName: constants.modelFramework,
+				ariaLabel: constants.modelFramework,
+				valueType: azdata.DeclarativeDataType.string,
+				isReadOnly: true,
+				width: 150,
+				headerCssStyles: {
+					...constants.cssStyles.tableHeader
+				},
+				rowCssStyles: {
+					...constants.cssStyles.tableRow
+				},
+			},
 			{ // Action
 				displayName: '',
 				valueType: azdata.DeclarativeDataType.component,
@@ -113,13 +139,13 @@ export class CurrentModelsTable extends ModelViewBase implements IDataComponent<
 
 	public addComponents(formBuilder: azdata.FormBuilder) {
 		if (this.component) {
-			formBuilder.addFormItem({ title: constants.modelSourcesTitle, component: this.component });
+			formBuilder.addFormItem({ title: '', component: this.component });
 		}
 	}
 
 	public removeComponents(formBuilder: azdata.FormBuilder) {
 		if (this.component) {
-			formBuilder.removeFormItem({ title: constants.modelSourcesTitle, component: this.component });
+			formBuilder.removeFormItem({ title: '', component: this.component });
 		}
 	}
 
@@ -169,7 +195,7 @@ export class CurrentModelsTable extends ModelViewBase implements IDataComponent<
 	}
 
 	private createTableRow(model: ImportedModel): any[] {
-		let row: any[] = [model.modelName, model.created];
+		let row: any[] = [model.modelName, model.created, model.version, model.framework];
 		if (this._modelBuilder) {
 			const selectButton = this.createSelectButton(model);
 			if (selectButton) {

extensions/machine-learning/src/views/models/manageModels/importModelWizard.ts

@@ -58,11 +58,9 @@ export class ImportModelWizard extends ModelViewBase {
 				validated = this.wizardView ? await this.wizardView.validate(pageInfo) : false;
 			}
 			if (validated && pageInfo.newPage === undefined) {
-				wizard.cancelButton.enabled = false;
-				wizard.backButton.enabled = false;
+				this.onLoading();
 				let result = await this.registerModel();
-				wizard.cancelButton.enabled = true;
-				wizard.backButton.enabled = true;
+				this.onLoaded();
 				if (this._parentView) {
 					this._parentView.importTable = this.importTable;
 					await this._parentView.refresh();
@@ -76,6 +74,21 @@ export class ImportModelWizard extends ModelViewBase {
 		await wizard.open();
 	}
 
+	private onLoading(): void {
+		this.refreshButtons(true);
+	}
+
+	private onLoaded(): void {
+		this.refreshButtons(false);
+	}
+
+	private refreshButtons(loading: boolean): void {
+		if (this.wizardView && this.wizardView.wizard) {
+			this.wizardView.wizard.cancelButton.enabled = !loading;
+			this.wizardView.wizard.backButton.enabled = !loading;
+		}
+	}
+
 	public get modelResources(): ModelSourcesComponent | undefined {
 		return this.modelSourcePage?.modelResources;
 	}

extensions/machine-learning/src/views/models/manageModels/manageModelsDialog.ts

@@ -40,6 +40,7 @@ export class ManageModelsDialog extends ModelViewBase {
 		});
 
 		let dialog = this.dialogView.createDialog(constants.registerModelTitle, [this.currentLanguagesTab]);
+		dialog.isWide = true;
 		dialog.customButtons = [registerModelButton];
 		this.mainViewPanel = dialog;
 		dialog.okButton.hidden = true;

extensions/machine-learning/src/views/models/manageModels/modelDetailsComponent.ts

@@ -78,7 +78,7 @@ export class ModelDetailsComponent extends ModelViewBase implements IDataCompone
 				component: this._createdComponent
 			},
 			{
-				title: constants.modelDeployed,
+				title: constants.modelImported,
 				component: this._deployedComponent
 			}, {
 				title: constants.modelFramework,

extensions/machine-learning/src/views/models/manageModels/modelImportLocationPage.ts

@@ -19,6 +19,7 @@ export class ModelImportLocationPage extends ModelViewBase implements IPageView,
 	private _form: azdata.FormContainer | undefined;
 	private _formBuilder: azdata.FormBuilder | undefined;
 	public tableSelectionComponent: TableSelectionComponent | undefined;
+	private _labelComponent: azdata.TextComponent | undefined;
 
 	constructor(apiWrapper: ApiWrapper, parent: ModelViewBase) {
 		super(apiWrapper, parent.root, parent);
@@ -31,12 +32,35 @@ export class ModelImportLocationPage extends ModelViewBase implements IPageView,
 	public registerComponent(modelBuilder: azdata.ModelBuilder): azdata.Component {
 
 		this._formBuilder = modelBuilder.formContainer();
-		this.tableSelectionComponent = new TableSelectionComponent(this._apiWrapper, this, true);
+		this.tableSelectionComponent = new TableSelectionComponent(this._apiWrapper, this, { editable: true, preSelected: true });
+		this._labelComponent = modelBuilder.text().withProperties({
+			width: 200
+		}).component();
+		const container = modelBuilder.flexContainer().withLayout({
+			width: 800,
+			height: '400px',
+			justifyContent: 'center'
+		}).withItems([
+			this._labelComponent
+		], {
+			CSSStyles: {
+				'align-items': 'center',
+				'padding-top': '30px',
+				'font-size': '16px'
+			}
+		}).component();
+
+
 		this.tableSelectionComponent.onSelectedChanged(async () => {
 			await this.onTableSelected();
 		});
-		this.tableSelectionComponent.registerComponent(modelBuilder);
+		this.tableSelectionComponent.registerComponent(modelBuilder, constants.databaseName, constants.tableName);
 		this.tableSelectionComponent.addComponents(this._formBuilder);
+
+		this._formBuilder.addFormItem({
+			title: '',
+			component: container
+		});
 		this._form = this._formBuilder.component();
 		return this._form;
 	}
@@ -45,6 +69,15 @@ export class ModelImportLocationPage extends ModelViewBase implements IPageView,
 		if (this.tableSelectionComponent?.data) {
 			this.importTable = this.tableSelectionComponent?.data;
 		}
+
+		if (this.importTable && this._labelComponent) {
+			const validated = await this.verifyImportConfigTable(this.importTable);
+			if (validated) {
+				this._labelComponent.value = constants.modelSchemaIsAcceptedMessage;
+			} else {
+				this._labelComponent.value = constants.modelSchemaIsNotAcceptedMessage;
+			}
+		}
 	}
 
 	/**

extensions/machine-learning/src/views/models/modelBrowsePage.ts

@@ -46,7 +46,6 @@ export class ModelBrowsePage extends ModelViewBase implements IPageView, IDataCo
 			editable: false
 		});
 		this.registeredModelsComponent.registerComponent(modelBuilder);
-		this.refresh();
 		this._form = this._formBuilder.component();
 		return this._form;
 	}
@@ -173,6 +172,7 @@ export class ModelBrowsePage extends ModelViewBase implements IPageView, IDataCo
 							fileName: x.model?.name,
 							framework: x.model?.framework,
 							frameworkVersion: x.model?.frameworkVersion,
+							description: x.model?.description,
 							created: x.model?.createdTime
 						},
 						targetImportTable: this.importTable

extensions/machine-learning/src/views/models/modelManagementController.ts

@@ -40,19 +40,19 @@ export class ModelManagementController extends ControllerBase {
 		apiWrapper: ApiWrapper,
 		private _root: string,
 		private _amlService: AzureModelRegistryService,
-		private _registeredModelService: DeployedModelService,
+		private _deployedModelService: DeployedModelService,
 		private _predictService: PredictService) {
 		super(apiWrapper);
 	}
 
 	/**
-	 * Opens the dialog for model registration
+	 * Opens the dialog for model import
 	 * @param parent parent if the view is opened from another view
 	 * @param controller controller
 	 * @param apiWrapper apiWrapper
 	 * @param root root folder path
 	 */
-	public async registerModel(importTable: DatabaseTable | undefined, parent?: ModelViewBase, controller?: ModelManagementController, apiWrapper?: ApiWrapper, root?: string): Promise<ModelViewBase> {
+	public async importModel(importTable: DatabaseTable | undefined, parent?: ModelViewBase, controller?: ModelManagementController, apiWrapper?: ApiWrapper, root?: string): Promise<ModelViewBase> {
 		controller = controller || this;
 		apiWrapper = apiWrapper || this._apiWrapper;
 		root = root || this._root;
@@ -60,7 +60,7 @@ export class ModelManagementController extends ControllerBase {
 		if (importTable) {
 			view.importTable = importTable;
 		} else {
-			view.importTable = await controller._registeredModelService.getRecentImportTable();
+			view.importTable = await controller._deployedModelService.getRecentImportTable();
 		}
 
 		controller.registerEvents(view);
@@ -93,23 +93,31 @@ export class ModelManagementController extends ControllerBase {
 	/**
 	 * Opens the wizard for prediction
 	 */
-	public async predictModel(): Promise<ModelViewBase> {
+	public async predictModel(): Promise<ModelViewBase | undefined> {
 
-		let view = new PredictWizard(this._apiWrapper, this._root);
-		view.importTable = await this._registeredModelService.getRecentImportTable();
+		const onnxSupported = await this._predictService.serverSupportOnnxModel();
+		if (onnxSupported) {
+			await this._deployedModelService.installDependencies();
+			let view = new PredictWizard(this._apiWrapper, this._root);
+			view.importTable = await this._deployedModelService.getRecentImportTable();
 
-		this.registerEvents(view);
-		view.on(LoadModelParametersEventName, async () => {
-			const modelArtifact = await view.getModelFileName();
-			await this.executeAction(view, LoadModelParametersEventName, this.loadModelParameters, this._registeredModelService,
-				modelArtifact?.filePath);
-		});
+			this.registerEvents(view);
 
-		// Open view
-		//
-		await view.open();
-		await view.refresh();
-		return view;
+			view.on(LoadModelParametersEventName, async (args) => {
+				const modelArtifact = await view.getModelFileName();
+				await this.executeAction(view, LoadModelParametersEventName, args, this.loadModelParameters, this._deployedModelService,
+					modelArtifact?.filePath);
+			});
+
+			// Open view
+			//
+			await view.open();
+			await view.refresh();
+			return view;
+		} else {
+			this._apiWrapper.showErrorMessage(constants.onnxNotSupportedError);
+			return undefined;
+		}
 	}
 
 
@@ -122,99 +130,99 @@ export class ModelManagementController extends ControllerBase {
 		// Register events
 		//
 		super.registerEvents(view);
-		view.on(ListAccountsEventName, async () => {
-			await this.executeAction(view, ListAccountsEventName, this.getAzureAccounts, this._amlService);
+		view.on(ListAccountsEventName, async (args) => {
+			await this.executeAction(view, ListAccountsEventName, args, this.getAzureAccounts, this._amlService);
 		});
-		view.on(ListSubscriptionsEventName, async (arg) => {
-			let azureArgs = <AzureResourceEventArgs>arg;
-			await this.executeAction(view, ListSubscriptionsEventName, this.getAzureSubscriptions, this._amlService, azureArgs.account);
+		view.on(ListSubscriptionsEventName, async (args) => {
+			let azureArgs = <AzureResourceEventArgs>args;
+			await this.executeAction(view, ListSubscriptionsEventName, args, this.getAzureSubscriptions, this._amlService, azureArgs.account);
 		});
-		view.on(ListWorkspacesEventName, async (arg) => {
-			let azureArgs = <AzureResourceEventArgs>arg;
-			await this.executeAction(view, ListWorkspacesEventName, this.getWorkspaces, this._amlService, azureArgs.account, azureArgs.subscription, azureArgs.group);
+		view.on(ListWorkspacesEventName, async (args) => {
+			let azureArgs = <AzureResourceEventArgs>args;
+			await this.executeAction(view, ListWorkspacesEventName, args, this.getWorkspaces, this._amlService, azureArgs.account, azureArgs.subscription, azureArgs.group);
 		});
-		view.on(ListGroupsEventName, async (arg) => {
-			let azureArgs = <AzureResourceEventArgs>arg;
-			await this.executeAction(view, ListGroupsEventName, this.getAzureGroups, this._amlService, azureArgs.account, azureArgs.subscription);
+		view.on(ListGroupsEventName, async (args) => {
+			let azureArgs = <AzureResourceEventArgs>args;
+			await this.executeAction(view, ListGroupsEventName, args, this.getAzureGroups, this._amlService, azureArgs.account, azureArgs.subscription);
 		});
-		view.on(ListAzureModelsEventName, async (arg) => {
-			let azureArgs = <AzureResourceEventArgs>arg;
-			await this.executeAction(view, ListAzureModelsEventName, this.getAzureModels, this._amlService
+		view.on(ListAzureModelsEventName, async (args) => {
+			let azureArgs = <AzureResourceEventArgs>args;
+			await this.executeAction(view, ListAzureModelsEventName, args, this.getAzureModels, this._amlService
 				, azureArgs.account, azureArgs.subscription, azureArgs.group, azureArgs.workspace);
 		});
 		view.on(ListModelsEventName, async (args) => {
 			const table = <DatabaseTable>args;
-			await this.executeAction(view, ListModelsEventName, this.getRegisteredModels, this._registeredModelService, table);
+			await this.executeAction(view, ListModelsEventName, args, this.getRegisteredModels, this._deployedModelService, table);
 		});
-		view.on(RegisterLocalModelEventName, async (arg) => {
-			let models = <ModelViewData[]>arg;
-			await this.executeAction(view, RegisterLocalModelEventName, this.registerLocalModel, this._registeredModelService, models);
+		view.on(RegisterLocalModelEventName, async (args) => {
+			let models = <ModelViewData[]>args;
+			await this.executeAction(view, RegisterLocalModelEventName, args, this.registerLocalModel, this._deployedModelService, models);
 			view.refresh();
 		});
 		view.on(RegisterModelEventName, async (args) => {
 			const importTable = <DatabaseTable>args;
-			await this.executeAction(view, RegisterModelEventName, this.registerModel, importTable, view, this, this._apiWrapper, this._root);
+			await this.executeAction(view, RegisterModelEventName, args, this.importModel, importTable, view, this, this._apiWrapper, this._root);
 		});
 		view.on(EditModelEventName, async (args) => {
 			const model = <ImportedModel>args;
-			await this.executeAction(view, EditModelEventName, this.editModel, model, view, this, this._apiWrapper, this._root);
+			await this.executeAction(view, EditModelEventName, args, this.editModel, model, view, this, this._apiWrapper, this._root);
 		});
 		view.on(UpdateModelEventName, async (args) => {
 			const model = <ImportedModel>args;
-			await this.executeAction(view, UpdateModelEventName, this.updateModel, this._registeredModelService, model);
+			await this.executeAction(view, UpdateModelEventName, args, this.updateModel, this._deployedModelService, model);
 		});
 		view.on(DeleteModelEventName, async (args) => {
 			const model = <ImportedModel>args;
-			await this.executeAction(view, DeleteModelEventName, this.deleteModel, this._registeredModelService, model);
+			await this.executeAction(view, DeleteModelEventName, args, this.deleteModel, this._deployedModelService, model);
 		});
-		view.on(RegisterAzureModelEventName, async (arg) => {
-			let models = <ModelViewData[]>arg;
-			await this.executeAction(view, RegisterAzureModelEventName, this.registerAzureModel, this._amlService, this._registeredModelService,
+		view.on(RegisterAzureModelEventName, async (args) => {
+			let models = <ModelViewData[]>args;
+			await this.executeAction(view, RegisterAzureModelEventName, args, this.registerAzureModel, this._amlService, this._deployedModelService,
 				models);
 		});
-		view.on(DownloadAzureModelEventName, async (arg) => {
-			let registerArgs = <AzureModelResource>arg;
-			await this.executeAction(view, DownloadAzureModelEventName, this.downloadAzureModel, this._amlService,
+		view.on(DownloadAzureModelEventName, async (args) => {
+			let registerArgs = <AzureModelResource>args;
+			await this.executeAction(view, DownloadAzureModelEventName, args, this.downloadAzureModel, this._amlService,
 				registerArgs.account, registerArgs.subscription, registerArgs.group, registerArgs.workspace, registerArgs.model);
 		});
-		view.on(ListDatabaseNamesEventName, async () => {
-			await this.executeAction(view, ListDatabaseNamesEventName, this.getDatabaseList, this._predictService);
+		view.on(ListDatabaseNamesEventName, async (args) => {
+			await this.executeAction(view, ListDatabaseNamesEventName, args, this.getDatabaseList, this._predictService);
 		});
-		view.on(ListTableNamesEventName, async (arg) => {
-			let dbName = <string>arg;
-			await this.executeAction(view, ListTableNamesEventName, this.getTableList, this._predictService, dbName);
+		view.on(ListTableNamesEventName, async (args) => {
+			let dbName = <string>args;
+			await this.executeAction(view, ListTableNamesEventName, args, this.getTableList, this._predictService, dbName);
 		});
-		view.on(ListColumnNamesEventName, async (arg) => {
-			let tableColumnsArgs = <DatabaseTable>arg;
-			await this.executeAction(view, ListColumnNamesEventName, this.getTableColumnsList, this._predictService,
+		view.on(ListColumnNamesEventName, async (args) => {
+			let tableColumnsArgs = <DatabaseTable>args;
+			await this.executeAction(view, ListColumnNamesEventName, args, this.getTableColumnsList, this._predictService,
 				tableColumnsArgs);
 		});
-		view.on(PredictModelEventName, async (arg) => {
-			let predictArgs = <PredictModelEventArgs>arg;
-			await this.executeAction(view, PredictModelEventName, this.generatePredictScript, this._predictService,
+		view.on(PredictModelEventName, async (args) => {
+			let predictArgs = <PredictModelEventArgs>args;
+			await this.executeAction(view, PredictModelEventName, args, this.generatePredictScript, this._predictService,
 				predictArgs, predictArgs.model, predictArgs.filePath);
 		});
-		view.on(DownloadRegisteredModelEventName, async (arg) => {
-			let model = <ImportedModel>arg;
-			await this.executeAction(view, DownloadRegisteredModelEventName, this.downloadRegisteredModel, this._registeredModelService,
+		view.on(DownloadRegisteredModelEventName, async (args) => {
+			let model = <ImportedModel>args;
+			await this.executeAction(view, DownloadRegisteredModelEventName, args, this.downloadRegisteredModel, this._deployedModelService,
 				model);
 		});
-		view.on(StoreImportTableEventName, async (arg) => {
-			let importTable = <DatabaseTable>arg;
-			await this.executeAction(view, StoreImportTableEventName, this.storeImportTable, this._registeredModelService,
+		view.on(StoreImportTableEventName, async (args) => {
+			let importTable = <DatabaseTable>args;
+			await this.executeAction(view, StoreImportTableEventName, args, this.storeImportTable, this._deployedModelService,
 				importTable);
 		});
-		view.on(VerifyImportTableEventName, async (arg) => {
-			let importTable = <DatabaseTable>arg;
-			await this.executeAction(view, VerifyImportTableEventName, this.verifyImportTable, this._registeredModelService,
+		view.on(VerifyImportTableEventName, async (args) => {
+			let importTable = <DatabaseTable>args;
+			await this.executeAction(view, VerifyImportTableEventName, args, this.verifyImportTable, this._deployedModelService,
 				importTable);
 		});
-		view.on(SourceModelSelectedEventName, async (arg) => {
-			view.modelSourceType = <ModelSourceType>arg;
+		view.on(SourceModelSelectedEventName, async (args) => {
+			view.modelSourceType = <ModelSourceType>args;
 			await view.refresh();
 		});
-		view.on(SignInToAzureEventName, async () => {
-			await this.executeAction(view, SignInToAzureEventName, this.signInToAzure, this._amlService);
+		view.on(SignInToAzureEventName, async (args) => {
+			await this.executeAction(view, SignInToAzureEventName, args, this.signInToAzure, this._amlService);
 			await view.refresh();
 		});
 	}
@@ -228,7 +236,7 @@ export class ModelManagementController extends ControllerBase {
 		if (importTable) {
 			view.importTable = importTable;
 		} else {
-			view.importTable = await this._registeredModelService.getRecentImportTable();
+			view.importTable = await this._deployedModelService.getRecentImportTable();
 		}
 
 		// Register events

extensions/machine-learning/src/views/models/modelSourcesComponent.ts

@@ -31,12 +31,14 @@ export class ModelSourcesComponent extends ModelViewBase implements IDataCompone
 	 */
 	public registerComponent(modelBuilder: azdata.ModelBuilder): azdata.Component {
 
+		this._sourceType = this._options && this._options.length > 0 ? this._options[0] : ModelSourceType.Local;
+		this.modelSourceType = this._sourceType;
 		this._localModel = modelBuilder.card()
 			.withProperties({
 				value: 'local',
 				name: 'modelLocation',
 				label: constants.localModelSource,
-				selected: this._options[0] === ModelSourceType.Local,
+				selected: this._sourceType === ModelSourceType.Local,
 				cardType: azdata.CardType.VerticalButton,
 				width: 50
 			}).component();
@@ -45,7 +47,7 @@ export class ModelSourcesComponent extends ModelViewBase implements IDataCompone
 				value: 'aml',
 				name: 'modelLocation',
 				label: constants.azureModelSource,
-				selected: this._options[0] === ModelSourceType.Azure,
+				selected: this._sourceType === ModelSourceType.Azure,
 				cardType: azdata.CardType.VerticalButton,
 				width: 50
 			}).component();
@@ -55,7 +57,7 @@ export class ModelSourcesComponent extends ModelViewBase implements IDataCompone
 				value: 'registered',
 				name: 'modelLocation',
 				label: constants.registeredModelsSource,
-				selected: this._options[0] === ModelSourceType.RegisteredModels,
+				selected: this._sourceType === ModelSourceType.RegisteredModels,
 				cardType: azdata.CardType.VerticalButton,
 				width: 50
 			}).component();
@@ -105,9 +107,6 @@ export class ModelSourcesComponent extends ModelViewBase implements IDataCompone
 					break;
 			}
 		});
-		this._sourceType = this._options[0];
-		this.sendRequest(SourceModelSelectedEventName, this._sourceType);
-
 		this._flexContainer = modelBuilder.flexContainer()
 			.withLayout({
 				flexFlow: 'row',

extensions/machine-learning/src/views/models/prediction/columnsSelectionPage.ts

@@ -74,6 +74,18 @@ export class ColumnsSelectionPage extends ModelViewBase implements IPageView, ID
 		}
 	}
 
+	public async validate(): Promise<boolean> {
+		const data = this.data;
+		const validated = data !== undefined && data.databaseName !== undefined && data.inputColumns !== undefined && data.outputColumns !== undefined
+			&& data.tableName !== undefined && data.databaseName !== constants.selectDatabaseTitle && data.tableName !== constants.selectTableTitle
+			&& !data.inputColumns.find(x => x.columnName === constants.selectColumnTitle);
+		if (!validated) {
+			this.showErrorMessage(constants.invalidModelParametersError);
+		}
+
+		return Promise.resolve(validated);
+	}
+
 	public async onEnter(): Promise<void> {
 		await this.inputColumnsComponent?.onLoading();
 		await this.outputColumnsComponent?.onLoading();

extensions/machine-learning/src/views/models/prediction/columnsTable.ts

@@ -173,7 +173,12 @@ export class ColumnsTable extends ModelViewBase implements IDataComponent<Predic
 
 		if (this._table) {
 			if (this._forInput) {
-				const columns = await this.listColumnNames(table);
+				let columns: TableColumn[];
+				try {
+					columns = await this.listColumnNames(table);
+				} catch {
+					columns = [];
+				}
 				if (modelParameters?.inputs && columns) {
 					tableData = tableData.concat(modelParameters.inputs.map(input => this.createInputTableRow(input, columns)));
 				}
@@ -212,6 +217,10 @@ export class ColumnsTable extends ModelViewBase implements IDataComponent<Predic
 			const dataType = dataTypes.find(x => x === modelParameter.type);
 			if (dataType) {
 				nameInput.value = dataType;
+			} else {
+				// Output type not supported
+				//
+				modelParameter.type = dataTypes[0];
 			}
 			this._parameters.push({ columnName: name, paramName: name, dataType: modelParameter.type });
 
@@ -234,7 +243,7 @@ export class ColumnsTable extends ModelViewBase implements IDataComponent<Predic
 					selectedRow.columnName = displayNameInput.value || name;
 				}
 			});
-			return [`${name}(${modelParameter.type ? modelParameter.type : constants.unsupportedModelParameterType})`, displayNameInput, nameInput];
+			return [`${name}(${modelParameter.originalType ? modelParameter.originalType : constants.unsupportedModelParameterType})`, displayNameInput, nameInput];
 		}
 
 		return [];
@@ -242,7 +251,11 @@ export class ColumnsTable extends ModelViewBase implements IDataComponent<Predic
 
 	private createInputTableRow(modelParameter: ModelParameter, columns: TableColumn[] | undefined): any[] {
 		if (this._modelBuilder && columns) {
-			const values = columns.map(c => { return { name: c.columnName, displayName: `${c.columnName}(${c.dataType})` }; });
+
+			let values = columns.map(c => { return { name: c.columnName, displayName: `${c.columnName}(${c.dataType})` }; });
+			if (columns.length > 0 && columns[0].columnName !== constants.selectColumnTitle) {
+				values = [{ displayName: constants.selectColumnTitle, name: '' }].concat(values);
+			}
 			let nameInput = this._modelBuilder.dropDown().withProperties({
 				values: values,
 				width: this.componentMaxLength
@@ -250,11 +263,28 @@ export class ColumnsTable extends ModelViewBase implements IDataComponent<Predic
 			const name = modelParameter.name;
 			let column = values.find(x => x.name === modelParameter.name);
 			if (!column) {
-				column = values[0];
+				column = values.length > 0 ? values[0] : undefined;
 			}
+			const currentColumn = columns.find(x => x.columnName === column?.name);
 			nameInput.value = column;
 
-			this._parameters.push({ columnName: column.name, paramName: name });
+			if (column) {
+				this._parameters.push({ columnName: column.name, paramName: name, paramType: modelParameter.type });
+			}
+			const inputContainer = this._modelBuilder.flexContainer().withLayout({
+				flexFlow: 'row',
+				width: this.componentMaxLength + 20,
+				justifyContent: 'flex-start'
+			}).component();
+			const warningButton = this.createWarningButton();
+			warningButton.onDidClick(() => {
+			});
+
+			const css = {
+				'padding-top': '5px',
+				'padding-right': '5px',
+				'margin': '0px'
+			};
 
 			nameInput.onValueChanged(() => {
 				const selectedColumn = nameInput.value;
@@ -264,12 +294,36 @@ export class ColumnsTable extends ModelViewBase implements IDataComponent<Predic
 				if (selectedRow) {
 					selectedRow.columnName = value || '';
 				}
+
+				const currentColumn = columns.find(x => x.columnName === value);
+				if (currentColumn && modelParameter.type === currentColumn?.dataType) {
+					inputContainer.removeItem(warningButton);
+				} else {
+					inputContainer.addItem(warningButton, {
+						CSSStyles: css
+					});
+				}
 			});
+
 			const label = this._modelBuilder.inputBox().withProperties({
-				value: `${name}(${modelParameter.type ? modelParameter.type : constants.unsupportedModelParameterType})`,
+				value: `${name}(${modelParameter.originalType ? modelParameter.originalType : constants.unsupportedModelParameterType})`,
 				enabled: false,
 				width: this.componentMaxLength
 			}).component();
+
+
+			inputContainer.addItem(label, {
+				CSSStyles: {
+					'padding': '0px',
+					'padding-right': '5px',
+					'margin': '0px'
+				}
+			});
+			if (currentColumn && modelParameter.type !== currentColumn?.dataType) {
+				inputContainer.addItem(warningButton, {
+					CSSStyles: css
+				});
+			}
 			const image = this._modelBuilder.image().withProperties({
 				width: 50,
 				height: 50,
@@ -281,12 +335,28 @@ export class ColumnsTable extends ModelViewBase implements IDataComponent<Predic
 				iconHeight: 20,
 				title: 'maps'
 			}).component();
-			return [nameInput, image, label];
+			return [nameInput, image, inputContainer];
 		}
 
 		return [];
 	}
 
+	private createWarningButton(): azdata.ButtonComponent {
+		const warningButton = this._modelBuilder.button().withProperties({
+			width: '10px',
+			height: '10px',
+			title: constants.columnDataTypeMismatchWarning,
+			iconPath: {
+				dark: this.asAbsolutePath('images/dark/warning_notification_inverse.svg'),
+				light: this.asAbsolutePath('images/light/warning_notification.svg'),
+			},
+			iconHeight: '10px',
+			iconWidth: '10px'
+		}).component();
+
+		return warningButton;
+	}
+
 	/**
 	 * Returns selected data
 	 */

extensions/machine-learning/src/views/models/prediction/inputColumnsComponent.ts

@@ -35,8 +35,8 @@ export class InputColumnsComponent extends ModelViewBase implements IDataCompone
 	 * @param modelBuilder model builder
 	 */
 	public registerComponent(modelBuilder: azdata.ModelBuilder): azdata.Component {
-		this._tableSelectionComponent = new TableSelectionComponent(this._apiWrapper, this, false);
-		this._tableSelectionComponent.registerComponent(modelBuilder);
+		this._tableSelectionComponent = new TableSelectionComponent(this._apiWrapper, this, { editable: false, preSelected: false });
+		this._tableSelectionComponent.registerComponent(modelBuilder, constants.columnDatabase, constants.columnTable);
 		this._tableSelectionComponent.onSelectedChanged(async () => {
 			await this.onTableSelected();
 		});

extensions/machine-learning/src/views/models/prediction/predictWizard.ts

@@ -62,7 +62,7 @@ export class PredictWizard extends ModelViewBase {
 		});
 		wizard.registerNavigationValidator(async (pageInfo: azdata.window.WizardPageChangeInfo) => {
 			let validated: boolean = true;
-			if (pageInfo.newPage > pageInfo.lastPage) {
+			if (pageInfo.newPage === undefined || pageInfo.newPage > pageInfo.lastPage) {
 				validated = this.wizardView ? await this.wizardView.validate(pageInfo) : false;
 			}
 			if (validated) {
@@ -94,7 +94,7 @@ export class PredictWizard extends ModelViewBase {
 	private refreshButtons(loading: boolean): void {
 		if (this.wizardView && this.wizardView.wizard) {
 			this.wizardView.wizard.cancelButton.enabled = !loading;
-			this.wizardView.wizard.cancelButton.enabled = !loading;
+			this.wizardView.wizard.backButton.enabled = !loading;
 		}
 	}
 

extensions/machine-learning/src/views/models/tableSelectionComponent.ts

@@ -7,10 +7,14 @@ import * as azdata from 'azdata';
 import * as vscode from 'vscode';
 import { ModelViewBase } from './modelViewBase';
 import { ApiWrapper } from '../../common/apiWrapper';
-import * as constants from '../../common/constants';
 import { IDataComponent } from '../interfaces';
 import { DatabaseTable } from '../../prediction/interfaces';
+import * as constants from '../../common/constants';
 
+export interface ITableSelectionSettings {
+	editable: boolean,
+	preSelected: boolean
+}
 /**
  * View to render filters to pick an azure resource
  */
@@ -30,7 +34,7 @@ export class TableSelectionComponent extends ModelViewBase implements IDataCompo
 	/**
 	 * Creates a new view
 	 */
-	constructor(apiWrapper: ApiWrapper, parent: ModelViewBase, private _editable: boolean) {
+	constructor(apiWrapper: ApiWrapper, parent: ModelViewBase, private _settings: ITableSelectionSettings) {
 		super(apiWrapper, parent.root, parent);
 	}
 
@@ -38,16 +42,16 @@ export class TableSelectionComponent extends ModelViewBase implements IDataCompo
 	 * Register components
 	 * @param modelBuilder model builder
 	 */
-	public registerComponent(modelBuilder: azdata.ModelBuilder): azdata.Component {
+	public registerComponent(modelBuilder: azdata.ModelBuilder, databaseTitle: string, tableTitle: string): azdata.Component {
 		this._databases = modelBuilder.dropDown().withProperties({
 			width: this.componentMaxLength,
-			editable: this._editable,
-			fireOnTextChange: this._editable
+			editable: this._settings.editable,
+			fireOnTextChange: this._settings.editable
 		}).component();
 		this._tables = modelBuilder.dropDown().withProperties({
 			width: this.componentMaxLength,
-			editable: this._editable,
-			fireOnTextChange: this._editable
+			editable: this._settings.editable,
+			fireOnTextChange: this._settings.editable
 		}).component();
 
 		this._databases.onValueChanged(async () => {
@@ -58,21 +62,21 @@ export class TableSelectionComponent extends ModelViewBase implements IDataCompo
 			// There's an issue with dropdown doesn't set the value in editable mode. this is the workaround
 
 			if (this._tables && value) {
-				this._selectedTableName = this._editable ? value : value.selected;
+				this._selectedTableName = this._settings.editable ? value : value.selected;
 			}
 			await this.onTableSelected();
 		});
 
 		const databaseForm = modelBuilder.formContainer().withFormItems([{
-			title: constants.columnDatabase,
+			title: databaseTitle,
 			component: this._databases,
-		}]).withLayout({
+		}], { info: databaseTitle }).withLayout({
 			padding: '0px'
 		}).component();
 		const tableForm = modelBuilder.formContainer().withFormItems([{
-			title: constants.columnTable,
+			title: tableTitle,
 			component: this._tables
-		}]).withLayout({
+		}], { info: tableTitle }).withLayout({
 			padding: '0px'
 		}).component();
 		this._dbTableComponent = modelBuilder.flexContainer().withItems([
@@ -97,27 +101,21 @@ export class TableSelectionComponent extends ModelViewBase implements IDataCompo
 	}
 
 	public addComponents(formBuilder: azdata.FormBuilder) {
-		if (this._databases && this._tables) {
+		if (this._dbTableComponent) {
 			formBuilder.addFormItems([{
-				title: constants.databaseName,
-				component: this._databases
-			}, {
-				title: constants.tableName,
-				component: this._tables
+				title: '',
+				component: this._dbTableComponent
 			}]);
 		}
 	}
 
 	public removeComponents(formBuilder: azdata.FormBuilder) {
-		if (this._databases && this._tables) {
+		if (this._dbTableComponent) {
 			formBuilder.removeFormItem({
-				title: constants.databaseName,
-				component: this._databases
-			});
-			formBuilder.removeFormItem({
-				title: constants.tableName,
-				component: this._tables
+				title: '',
+				component: this._dbTableComponent
 			});
+
 		}
 	}
 
@@ -140,13 +138,19 @@ export class TableSelectionComponent extends ModelViewBase implements IDataCompo
 	 */
 	public async loadData(): Promise<void> {
 		this._dbNames = await this.listDatabaseNames();
-		if (this._databases && this._dbNames && this._dbNames.length > 0) {
-			this._databases.values = this._dbNames;
-			if (this.importTable) {
+		let dbNames = this._dbNames;
+		if (!this._settings.preSelected && !this._dbNames.find(x => x === constants.selectDatabaseTitle)) {
+			dbNames = [constants.selectDatabaseTitle].concat(this._dbNames);
+		}
+		if (this._databases && dbNames && dbNames.length > 0) {
+			this._databases.values = dbNames;
+
+			if (this.importTable && this._settings.preSelected) {
 				this._databases.value = this.importTable.databaseName;
 			} else {
-				this._databases.value = this._dbNames[0];
+				this._databases.value = dbNames[0];
 			}
+
 		}
 		await this.onDatabaseSelected();
 	}
@@ -160,18 +164,25 @@ export class TableSelectionComponent extends ModelViewBase implements IDataCompo
 
 	private async onDatabaseSelected(): Promise<void> {
 		this._tableNames = await this.listTableNames(this.databaseName || '');
-		if (this._tables && this._tableNames && this._tableNames.length > 0) {
-			this._tables.values = this._tableNames.map(t => this.getTableFullName(t));
+		let tableNames = this._tableNames;
+
+		if (this._tableNames && !this._settings.preSelected && !this._tableNames.find(x => x.tableName === constants.selectTableTitle)) {
+			const firstRow: DatabaseTable = { tableName: constants.selectTableTitle, databaseName: '', schema: '' };
+			tableNames = [firstRow].concat(this._tableNames);
+		}
+
+		if (this._tables && tableNames && tableNames.length > 0) {
+			this._tables.values = tableNames.map(t => this.getTableFullName(t));
 			if (this.importTable) {
-				const selectedTable = this._tableNames.find(t => t.tableName === this.importTable?.tableName && t.schema === this.importTable?.schema);
+				const selectedTable = tableNames.find(t => t.tableName === this.importTable?.tableName && t.schema === this.importTable?.schema);
 				if (selectedTable) {
 					this._selectedTableName = this.getTableFullName(selectedTable);
 					this._tables.value = this.getTableFullName(selectedTable);
 				} else {
-					this._selectedTableName = this._editable ? this.getTableFullName(this.importTable) : this.getTableFullName(this._tableNames[0]);
+					this._selectedTableName = this._settings.editable ? this.getTableFullName(this.importTable) : this.getTableFullName(tableNames[0]);
 				}
 			} else {
-				this._selectedTableName = this.getTableFullName(this._tableNames[0]);
+				this._selectedTableName = this.getTableFullName(tableNames[0]);
 			}
 			this._tables.value = this._selectedTableName;
 		} else if (this._tables) {
@@ -182,7 +193,7 @@ export class TableSelectionComponent extends ModelViewBase implements IDataCompo
 	}
 
 	private getTableFullName(table: DatabaseTable): string {
-		return `${table.schema}.${table.tableName}`;
+		return table.tableName === constants.selectTableTitle ? table.tableName : `${table.schema}.${table.tableName}`;
 	}
 
 	private async onTableSelected(): Promise<void> {

extensions/machine-learning/src/views/viewBase.ts

@@ -12,12 +12,7 @@ import * as path from 'path';
 import { EventEmitterCollection } from '../common/eventEmitter';
 
 export interface CallbackEventArgs {
-	data?: any;
-	error?: (reason?: any) => void;
-}
-
-
-export interface CallbackEventArgs {
+	inputArgs?: any;
 	data?: any;
 	error?: (reason?: any) => void;
 }
@@ -95,28 +90,38 @@ export abstract class ViewBase extends EventEmitterCollection {
 		this.fire(requestType, arg);
 	}
 
-	public sendDataRequest<T>(
+	public async sendDataRequest<T>(
 		eventName: string,
 		arg?: any,
 		callbackEventName?: string): Promise<T> {
-		return new Promise<T>((resolve, reject) => {
+		let emitter: vscode.EventEmitter<any> | undefined;
+		let promise = new Promise<T>((resolve, reject) => {
 			if (!callbackEventName) {
 				callbackEventName = ViewBase.getCallbackEventName(eventName);
 			}
-			this.on(callbackEventName, result => {
+			emitter = this.on(callbackEventName, result => {
 				let callbackArgs = <CallbackEventArgs>result;
 				if (callbackArgs) {
-					if (callbackArgs.error) {
-						reject(callbackArgs.error);
-					} else {
-						resolve(<T>callbackArgs.data);
+					if (callbackArgs.inputArgs === arg) {
+						if (callbackArgs.error) {
+							reject(callbackArgs.error);
+						} else {
+							resolve(<T>callbackArgs.data);
+						}
 					}
 				} else {
 					reject(constants.notSupportedEventArg);
 				}
 			});
+
 			this.fire(eventName, arg);
 		});
+		const result = await promise;
+		if (emitter && callbackEventName) {
+			this.disposeEvent(callbackEventName, emitter);
+		}
+
+		return result;
 	}
 
 	public async getLocalPaths(options: vscode.OpenDialogOptions): Promise<string[]> {

extensions/machine-learning/src/views/widgets/dashboardWidget.ts

@@ -9,6 +9,7 @@ import { ApiWrapper } from '../../common/apiWrapper';
 import * as path from 'path';
 import * as constants from '../../common/constants';
 import * as utils from '../../common/utils';
+import { PredictService } from '../../prediction/predictService';
 
 interface IActionMetadata {
 	title?: string,
@@ -25,53 +26,56 @@ export class DashboardWidget {
 	/**
 	 * Creates new instance of dashboard
 	 */
-	constructor(private _apiWrapper: ApiWrapper, private _root: string) {
+	constructor(private _apiWrapper: ApiWrapper, private _root: string, private _predictService: PredictService) {
 	}
 
-	public register(): void {
-		this._apiWrapper.registerWidget('mls.dashboard', async (view) => {
-			const container = view.modelBuilder.flexContainer().withLayout({
-				flexFlow: 'column',
-				width: '100%',
-				height: '100%'
-			}).component();
-			const header = this.createHeader(view);
-			const tasksContainer = this.createTasks(view);
-			const footerContainer = this.createFooter(view);
-			container.addItem(header, {
-				CSSStyles: {
-					'background-image': `url(${vscode.Uri.file(this.asAbsolutePath('images/background.svg'))})`,
-					'background-repeat': 'no-repeat',
-					'background-position': 'bottom',
-					'width': `${maxWidth}px`,
-					'height': '330px',
-					'background-size': `${maxWidth}px ${headerMaxHeight}px`,
-					'margin-bottom': '-60px'
-				}
-			});
-			container.addItem(tasksContainer, {
-				CSSStyles: {
-					'width': `${maxWidth}px`,
-					'height': '150px',
-				}
-			});
-			container.addItem(footerContainer, {
-				CSSStyles: {
-					'width': `${maxWidth}px`,
-					'height': '500px',
-				}
-			});
-			const mainContainer = view.modelBuilder.flexContainer()
-				.withLayout({
+	public register(): Promise<void> {
+		return new Promise<void>(resolve => {
+			this._apiWrapper.registerWidget('mls.dashboard', async (view) => {
+				const container = view.modelBuilder.flexContainer().withLayout({
 					flexFlow: 'column',
 					width: '100%',
-					height: '100%',
-					position: 'absolute'
+					height: '100%'
 				}).component();
-			mainContainer.addItem(container, {
-				CSSStyles: { 'padding-top': '25px', 'padding-left': '5px' }
+				const header = this.createHeader(view);
+				const tasksContainer = await this.createTasks(view);
+				const footerContainer = this.createFooter(view);
+				container.addItem(header, {
+					CSSStyles: {
+						'background-image': `url(${vscode.Uri.file(this.asAbsolutePath('images/background.svg'))})`,
+						'background-repeat': 'no-repeat',
+						'background-position': 'bottom',
+						'width': `${maxWidth}px`,
+						'height': '330px',
+						'background-size': `${maxWidth}px ${headerMaxHeight}px`,
+						'margin-bottom': '-60px'
+					}
+				});
+				container.addItem(tasksContainer, {
+					CSSStyles: {
+						'width': `${maxWidth}px`,
+						'height': '150px',
+					}
+				});
+				container.addItem(footerContainer, {
+					CSSStyles: {
+						'width': `${maxWidth}px`,
+						'height': '500px',
+					}
+				});
+				const mainContainer = view.modelBuilder.flexContainer()
+					.withLayout({
+						flexFlow: 'column',
+						width: '100%',
+						height: '100%',
+						position: 'absolute'
+					}).component();
+				mainContainer.addItem(container, {
+					CSSStyles: { 'padding-top': '25px', 'padding-left': '5px' }
+				});
+				await view.initializeModel(mainContainer);
+				resolve();
 			});
-			await view.initializeModel(mainContainer);
 		});
 	}
 
@@ -445,7 +449,7 @@ export class DashboardWidget {
 		return path.join(this._root || '', filePath);
 	}
 
-	private createTasks(view: azdata.ModelView): azdata.Component {
+	private async createTasks(view: azdata.ModelView): Promise<azdata.Component> {
 		const tasksContainer = view.modelBuilder.flexContainer().withLayout({
 			flexFlow: 'row',
 			width: '100%',
@@ -489,6 +493,7 @@ export class DashboardWidget {
 				'padding': '10px'
 			}
 		});
+		predictionButton.enabled = await this._predictService.serverSupportOnnxModel();
 
 		return tasksContainer;
 	}
@@ -506,7 +511,7 @@ export class DashboardWidget {
 		const iconContainer = view.modelBuilder.flexContainer().withLayout({
 			flexFlow: 'row',
 			width: maxWidth,
-			height: maxHeight - 20,
+			height: maxHeight - 23,
 			alignItems: 'flex-start'
 		}).component();
 		const labelsContainer = view.modelBuilder.flexContainer().withLayout({
@@ -571,7 +576,7 @@ export class DashboardWidget {
 			}
 		});
 		mainContainer.onDidClick(async () => {
-			if (taskMetaData.command) {
+			if (mainContainer.enabled && taskMetaData.command) {
 				await this._apiWrapper.executeCommand(taskMetaData.command);
 			}
 		});