/* Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. * * Licensed under the Apache License, Version 2.0 (the "License"). You * may not use this file except in compliance with the License. A copy of * the License is located at * * http://aws.amazon.com/apache2.0/ * * or in the "license" file accompanying this file. This file is * distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF * ANY KIND, either express or implied. See the License for the specific * language governing permissions and limitations under the License. */ import * as React from 'react'; import { RunsModel } from '../models/RunsModel'; import { NotebookPanel } from '@jupyterlab/notebook'; import { ServerConnection } from '@jupyterlab/services'; import { URLExt } from '@jupyterlab/coreutils'; import { IStateDB } from '@jupyterlab/statedb'; import { JSONValue, ReadonlyJSONObject, JSONArray } from '@lumino/coreutils'; import { Signal } from '@lumino/signaling'; import { Widget } from '@lumino/widgets'; import { UploadNotebookResponse, InvokeRequest, InvokeResponse, CreateRuleRequest, CreateRuleResponse, ErrorResponse, } from '../server'; import { InputColumn, LabeledTextInput } from './InputColumn'; import { ParameterEditor, ParameterKV } from './ParameterEditor'; import { Alert, AlertProps } from './Alert'; import { runSidebarSectionClass, runSidebarNotebookNameClass, runSidebarNoHeaderClass, sidebarButtonClass, alertAreaClass, runSidebarNoNotebookClass, flexButtonsClass, } from '../style/SchedulePanel'; import { JupyterFrontEnd, ILabShell } from '@jupyterlab/application'; import { RulesModel } from '../models/RulesModel'; const KEY = 'sagemaker-run-notebook:schedule-sidebar:data'; /** Interface for SchedulePanel component props */ export interface ISchedulePanelProps { app: JupyterFrontEnd; shell: ILabShell; runsModel: RunsModel; rulesModel: RulesModel; stateDB: IStateDB; } interface PersistentState { image: string; role: string; instanceType: string; ruleName: string; schedule: string; eventPattern: string; } interface ISchedulePanelState extends PersistentState { notebook: string; notebookPanel: NotebookPanel; parameters: ParameterKV[]; alerts: (AlertProps & { key: string })[]; } // convertParameters turns the parameters we use here into a map that the server and containter expect. // TODO: fix the container to take lists and delete this function so that parameters are always in the order specified // eslint-disable-next-line @typescript-eslint/no-explicit-any function convertParameters(params: ParameterKV[]): Record { // eslint-disable-next-line @typescript-eslint/no-explicit-any const result: Record = {}; params.forEach((param) => { result[param.name] = param.value; }); return result; } /** A React component for the schedule extension's main display */ export class SchedulePanel extends React.Component { constructor(props: ISchedulePanelProps) { super(props); this.app = props.app; this.shell = props.shell; this.currentNotebookChanged = new Signal(this); this.setCurrentWidget(this.shell.currentWidget); this.state = { notebook: this.notebook, notebookPanel: this.currentNotebookPanel, image: '', parameters: null, role: '', instanceType: '', ruleName: '', schedule: '', eventPattern: '', alerts: [], }; this.loadState(); this.shell.currentChanged.connect(this.onCurrentWidgetChanged, this); } //TODO: track notebook renames private onCurrentWidgetChanged(sender: ILabShell, args: ILabShell.IChangedArgs) { const newWidget = args.newValue; const label = newWidget && newWidget.title.label; console.log(`current widget changed to ${label}`); this.setCurrentWidget(newWidget); this.setState({ notebook: this.notebook, notebookPanel: this.currentNotebookPanel }); } private setCurrentWidget(newWidget: Widget): void { const context = newWidget && (newWidget as NotebookPanel).context; const session = context && context.sessionContext && context.sessionContext.session; const isNotebook = session && session.type === 'notebook'; if (isNotebook) { this.currentNotebookPanel = newWidget as NotebookPanel; this.notebook = session.name; } else { this.currentNotebookPanel = null; this.notebook = null; } this.currentNotebookChanged.emit(this.currentNotebookPanel); } /** * Renders the component. * * @returns React element */ render = (): React.ReactElement => { const notebookIndependent = (
{this.renderViewButtons()} {this.renderCurrentNotebook()}
); const notebookDependent = this.currentNotebookPanel ? (
{this.renderRunParameters()} {this.renderScheduleParameters()} {this.renderExecuteButtons()} {this.renderAlerts()}
) : (

Select or create a notebook to enable execution and scheduling.

); return (
{notebookIndependent} {notebookDependent}
); }; private renderViewButtons() { return (
View
); } private renderCurrentNotebook() { const notebook = this.state.notebook; let notebookDisplay: React.ReactElement; if (notebook != null) { notebookDisplay = {notebook}; } else { notebookDisplay = No notebook selected; } return (
Current Notebook

{notebookDisplay}

); } private renderRunParameters() { return (
Notebook Execution
); } private renderScheduleParameters() { return (
Schedule Rule
); } private renderExecuteButtons() { return (
); } private renderAlerts() { return (
{this.state.alerts.map((alert) => ( ))}
); } private onParametersChange = (editor: ParameterEditor): void => { this.setState({ parameters: editor.value }); }; private onImageChange = (event: React.ChangeEvent): void => { this.setState({ image: event.target.value }, () => this.saveState()); }; private onRoleChange = (event: React.ChangeEvent): void => { this.setState({ role: event.target.value }, () => this.saveState()); }; private onInstanceTypeChange = (event: React.ChangeEvent): void => { this.setState({ instanceType: event.target.value }, () => this.saveState()); }; private onRuleNameChange = (event: React.ChangeEvent): void => { this.setState({ ruleName: event.target.value }, () => this.saveState()); }; private onScheduleChange = (event: React.ChangeEvent): void => { this.setState({ schedule: event.target.value }, () => this.saveState()); }; private onEventPatternChange = (event: React.ChangeEvent): void => { this.setState({ eventPattern: event.target.value }, () => this.saveState()); }; private onRunListClick = async (): Promise => { this.app.commands.execute('sagemaker_run_notebook:open_list_runs'); }; private onScheduleListClick = async (): Promise => { this.app.commands.execute('sagemaker_run_notebook:open_list_schedules'); }; private onRunClick = async (): Promise => { console.log('Run!!!'); this.clearAlerts(); try { this.addAlert({ type: 'notice', message: `Starting notebook run for "${this.state.notebook}"`, }); const content = this.currentNotebookPanel.model.toJSON(); const s3Object = await this.uploadNotebook(content); console.log(`notebook uploaded to ${s3Object}`); // TODO: clean up non-camel-case entries in the server requests /* eslint-disable @typescript-eslint/camelcase */ const request: InvokeRequest = { image: this.state.image, input_path: s3Object, notebook: this.state.notebook, parameters: convertParameters(this.state.parameters), role: this.state.role, instance_type: this.state.instanceType, }; /* eslint-enable @typescript-eslint/camelcase */ const jobName = await this.invokeNotebook(request); this.addAlert({ message: `Started notebook run "${jobName}"` }); console.log(`started job ${jobName}`); this.props.runsModel.refresh(); } catch (e) { this.addAlert({ type: 'error', message: `Error starting run for "${this.state.notebook}": ${e.message}`, }); } }; private onScheduleClick = async (): Promise => { console.log('Create schedule!!!'); this.clearAlerts(); try { this.addAlert({ type: 'notice', message: `Creating rule "${this.state.ruleName}"`, }); const content = this.currentNotebookPanel.model.toJSON(); const s3Object = await this.uploadNotebook(content); console.log(`notebook uploaded to ${s3Object}`); /* eslint-disable @typescript-eslint/camelcase */ const request: CreateRuleRequest = { image: this.state.image, input_path: s3Object, notebook: this.state.notebook, parameters: convertParameters(this.state.parameters), role: this.state.role, instance_type: this.state.instanceType, }; const schedule = this.state.schedule; if (schedule !== '') { request.schedule = schedule; } const eventPattern = this.state.eventPattern; if (eventPattern !== '') { request.event_pattern = eventPattern; } /* eslint-enable @typescript-eslint/camelcase */ const ruleName = await this.createRule(this.state.ruleName, request); this.addAlert({ message: `Created rule "${ruleName}"` }); console.log(`created rule ${ruleName}`); this.props.rulesModel.refresh(); } catch (e) { this.addAlert({ type: 'error', message: `Error creating rule "${this.state.ruleName}": ${e.message}`, }); } }; private async uploadNotebook(notebook: JSONValue): Promise { const settings = ServerConnection.makeSettings(); const response = await ServerConnection.makeRequest( URLExt.join(settings.baseUrl, 'sagemaker-scheduler', 'upload'), { method: 'PUT', body: JSON.stringify(notebook) }, settings, ); if (!response.ok) { const error = (await response.json()) as ErrorResponse; let errorMessage: string; if (error.error) { errorMessage = error.error.message; } else { errorMessage = JSON.stringify(error); } throw Error('Uploading notebook to S3 failed: ' + errorMessage); } const data = (await response.json()) as UploadNotebookResponse; return data.s3Object; } // Figure out if the current notebook has a parameter cell marked private hasParameterCell(): boolean { if (!this.currentNotebookPanel) { return false; } const cells = this.currentNotebookPanel.model.cells; for (let i = 0; i < cells.length; i++) { const tags = cells.get(i).metadata.get('tags') as JSONArray; if (tags && tags.includes('parameters')) { return true; } } return false; } private runReady(): string[] { const result: string[] = []; if (!this.state.image) { result.push('missing container image'); } if (!this.state.instanceType) { result.push('missing instance type'); } if (this.state.parameters.length > 0 && !this.hasParameterCell()) { result.push(`no parameter cell defined in ${this.notebook}`); } return result; } private async invokeNotebook(request: InvokeRequest): Promise { const errors = this.runReady(); if (errors.length > 0) { throw new Error(errors.join(', ')); } const settings = ServerConnection.makeSettings(); const response = await ServerConnection.makeRequest( URLExt.join(settings.baseUrl, 'sagemaker-scheduler', 'run'), { method: 'POST', body: JSON.stringify(request) }, settings, ); if (!response.ok) { const error = (await response.json()) as ErrorResponse; let errorMessage: string; if (error.error) { errorMessage = error.error.message; } else { errorMessage = JSON.stringify(error); } throw Error(errorMessage); } const data = (await response.json()) as InvokeResponse; return data.job_name; } // Return an array of reasons that we can't create a schedule. private scheduleReady(): string[] { const result = this.runReady(); if (!this.state.ruleName) { result.push('missing schedule name'); } if (!(this.state.schedule || this.state.eventPattern)) { result.push('must have either a schedule or an event pattern'); } return result; } private async createRule(ruleName: string, request: CreateRuleRequest): Promise { const errors = this.scheduleReady(); if (errors.length > 0) { throw new Error(errors.join(', ')); } const settings = ServerConnection.makeSettings(); const response = await ServerConnection.makeRequest( URLExt.join(settings.baseUrl, 'sagemaker-scheduler', 'schedule', ruleName), { method: 'POST', body: JSON.stringify(request) }, settings, ); if (!response.ok) { const error = (await response.json()) as ErrorResponse; let errorMessage: string; if (error.error) { errorMessage = error.error.message; } else { errorMessage = JSON.stringify(error); } throw Error(errorMessage); } const data = (await response.json()) as CreateRuleResponse; return data.rule_name; } private alertKey = 0; private addAlert(alert: AlertProps) { const key = this.alertKey++; const keyedAlert: AlertProps & { key: string } = { ...alert, key: `alert-${key}` }; this.setState({ alerts: [keyedAlert] }); } private clearAlerts() { this.setState({ alerts: [] }); } private saveState() { const state = { image: this.state.image, role: this.state.role, instanceType: this.state.instanceType, ruleName: this.state.ruleName, schedule: this.state.schedule, eventPattern: this.state.eventPattern, }; this.props.stateDB.save(KEY, state); } private loadState() { this.props.stateDB.fetch(KEY).then((s) => { const state = s as ReadonlyJSONObject; if (state) { this.setState({ image: state['image'] as string, role: state['role'] as string, instanceType: state['instanceType'] as string, ruleName: state['ruleName'] as string, schedule: state['schedule'] as string, eventPattern: state['eventPattern'] as string, }); } }); } private app: JupyterFrontEnd; private shell: ILabShell; private currentNotebookPanel: NotebookPanel; private notebook: string; private currentNotebookChanged: Signal; }