From d2c8ec2fb960cc095cfb7f4e91a0512b5d67bf73 Mon Sep 17 00:00:00 2001 From: Mrunal Nitin Deshmukh Date: Thu, 27 Mar 2025 14:34:02 -0700 Subject: [PATCH 1/2] Add SageMaker-UI poststartup endpoint patch **Description** - Post-startup script stored in SMD (SageMaker Distribution) enables the required customizations for SMUS applications. - This functionality is currently operational for JupyterLab apps. - We are extending this capability to CodeEditor apps. **Testing Done** - Tested building the local image for these changes and tested on personal LL stack using BYOI - Verified script execution logs file generation at /var/log/apps after app launch --- .../src/vs/server/node/webClientServer.ts | 58 +++++++++++- patches/sagemaker-ui-post-startup.patch | 92 +++++++++++++++++++ patches/series | 1 + 3 files changed, 149 insertions(+), 2 deletions(-) create mode 100644 patches/sagemaker-ui-post-startup.patch diff --git a/patched-vscode/src/vs/server/node/webClientServer.ts b/patched-vscode/src/vs/server/node/webClientServer.ts index 159574451..09fca3bf5 100644 --- a/patched-vscode/src/vs/server/node/webClientServer.ts +++ b/patched-vscode/src/vs/server/node/webClientServer.ts @@ -3,9 +3,11 @@ * Licensed under the MIT License. See License.txt in the project root for license information. *--------------------------------------------------------------------------------------------*/ -import { createReadStream } from 'fs'; +import { createReadStream, existsSync, writeFileSync } from 'fs'; import {readFile } from 'fs/promises'; import { Promises } from 'vs/base/node/pfs'; +import { spawn } from 'child_process'; +import * as fs from 'fs'; import * as path from 'path'; import * as http from 'http'; import * as url from 'url'; @@ -39,6 +41,10 @@ const textMimeType: { [ext: string]: string | undefined } = { '.svg': 'image/svg+xml', }; +const enum ServiceName { + SAGEMAKER_UNIFIED_STUDIO = 'SageMakerUnifiedStudio', +} + /** * Return an error to the client. */ @@ -102,6 +108,7 @@ export class WebClientServer { private readonly _callbackRoute: string; private readonly _webExtensionRoute: string; private readonly _idleRoute: string; + private readonly _postStartupScriptRoute: string; constructor( private readonly _connectionToken: ServerConnectionToken, @@ -118,6 +125,7 @@ export class WebClientServer { this._callbackRoute = `${serverRootPath}/callback`; this._webExtensionRoute = `${serverRootPath}/web-extension-resource`; this._idleRoute = '/api/idle'; + this._postStartupScriptRoute = '/api/poststartup'; } /** @@ -146,6 +154,9 @@ export class WebClientServer { // extension resource support return this._handleWebExtensionResource(req, res, parsedUrl); } + if (pathname === this._postStartupScriptRoute) { + return this._handlePostStartupScriptInvocation(req, res); + } return serveError(req, res, 404, 'Not found.'); } catch (error) { @@ -459,12 +470,20 @@ export class WebClientServer { } /** - * Handles API requests to retrieve the last activity timestamp. + * Handles API requests to retrieve the last activity timestamp. */ private async _handleIdle(req: http.IncomingMessage, res: http.ServerResponse): Promise { try { const tmpDirectory = '/tmp/' const idleFilePath = path.join(tmpDirectory, '.sagemaker-last-active-timestamp'); + + // If idle shutdown file does not exist, this indicates the app UI may never been opened + // Create the initial metadata file + if (!existsSync(idleFilePath)) { + const timestamp = new Date().toISOString(); + writeFileSync(idleFilePath, timestamp); + } + const data = await readFile(idleFilePath, 'utf8'); res.statusCode = 200; @@ -474,6 +493,41 @@ export class WebClientServer { serveError(req, res, 500, error.message) } } + + /** + * Handles API requests to run the post-startup script in SMD. + */ + private async _handlePostStartupScriptInvocation(req: http.IncomingMessage, res: http.ServerResponse): Promise { + const postStartupScripPath = '/etc/sagemaker-ui/sagemaker_ui_post_startup.sh' + const logPath = '/var/log/apps/post_startup_default.log'; + const logStream = fs.createWriteStream(logPath, { flags: 'a' }); + + // Only trigger post-startup script invocation for SageMakerUnifiedStudio app. + if (process.env['SERVICE_NAME'] != ServiceName.SAGEMAKER_UNIFIED_STUDIO) { + return serveError(req, res, 403, 'Forbidden'); + } else { + //If postStartupScriptFile doesn't exist, it will throw FileNotFoundError (404) + //If exists, it will start the execution and add the execution logs in logFile. + try { + if (fs.existsSync(postStartupScripPath)) { + // Adding 0o755 to make script file executable + fs.chmodSync(postStartupScripPath, 0o755); + + const subprocess = spawn('bash', [`${postStartupScripPath}`], { cwd: '/' }); + subprocess.stdout.pipe(logStream); + subprocess.stderr.pipe(logStream); + + res.statusCode = 200; + res.setHeader('Content-Type', 'application/json'); + res.end(JSON.stringify({ 'success': 'true' })); + } else { + serveError(req, res, 500, 'Poststartup script file not found at ' + postStartupScripPath); + } + } catch (error) { + serveError(req, res, 500, error.message); + } + } + } } /** diff --git a/patches/sagemaker-ui-post-startup.patch b/patches/sagemaker-ui-post-startup.patch new file mode 100644 index 000000000..ba47f6e15 --- /dev/null +++ b/patches/sagemaker-ui-post-startup.patch @@ -0,0 +1,92 @@ +Index: sagemaker-code-editor/vscode/src/vs/server/node/webClientServer.ts +=================================================================== +--- sagemaker-code-editor.orig/vscode/src/vs/server/node/webClientServer.ts ++++ sagemaker-code-editor/vscode/src/vs/server/node/webClientServer.ts +@@ -6,6 +6,8 @@ + import { createReadStream, existsSync, writeFileSync } from 'fs'; + import {readFile } from 'fs/promises'; + import { Promises } from 'vs/base/node/pfs'; ++import { spawn } from 'child_process'; ++import * as fs from 'fs'; + import * as path from 'path'; + import * as http from 'http'; + import * as url from 'url'; +@@ -39,6 +41,10 @@ const textMimeType: { [ext: string]: str + '.svg': 'image/svg+xml', + }; + ++const enum ServiceName { ++ SAGEMAKER_UNIFIED_STUDIO = 'SageMakerUnifiedStudio', ++} ++ + /** + * Return an error to the client. + */ +@@ -102,6 +108,7 @@ export class WebClientServer { + private readonly _callbackRoute: string; + private readonly _webExtensionRoute: string; + private readonly _idleRoute: string; ++ private readonly _postStartupScriptRoute: string; + + constructor( + private readonly _connectionToken: ServerConnectionToken, +@@ -118,6 +125,7 @@ export class WebClientServer { + this._callbackRoute = `${serverRootPath}/callback`; + this._webExtensionRoute = `${serverRootPath}/web-extension-resource`; + this._idleRoute = '/api/idle'; ++ this._postStartupScriptRoute = '/api/poststartup'; + } + + /** +@@ -146,6 +154,9 @@ export class WebClientServer { + // extension resource support + return this._handleWebExtensionResource(req, res, parsedUrl); + } ++ if (pathname === this._postStartupScriptRoute) { ++ return this._handlePostStartupScriptInvocation(req, res); ++ } + + return serveError(req, res, 404, 'Not found.'); + } catch (error) { +@@ -482,6 +493,41 @@ export class WebClientServer { + serveError(req, res, 500, error.message) + } + } ++ ++ /** ++ * Handles API requests to run the post-startup script in SMD. ++ */ ++ private async _handlePostStartupScriptInvocation(req: http.IncomingMessage, res: http.ServerResponse): Promise { ++ const postStartupScripPath = '/etc/sagemaker-ui/sagemaker_ui_post_startup.sh' ++ const logPath = '/var/log/apps/post_startup_default.log'; ++ const logStream = fs.createWriteStream(logPath, { flags: 'a' }); ++ ++ // Only trigger post-startup script invocation for SageMakerUnifiedStudio app. ++ if (process.env['SERVICE_NAME'] != ServiceName.SAGEMAKER_UNIFIED_STUDIO) { ++ return serveError(req, res, 403, 'Forbidden'); ++ } else { ++ //If postStartupScriptFile doesn't exist, it will throw FileNotFoundError (404) ++ //If exists, it will start the execution and add the execution logs in logFile. ++ try { ++ if (fs.existsSync(postStartupScripPath)) { ++ // Adding 0o755 to make script file executable ++ fs.chmodSync(postStartupScripPath, 0o755); ++ ++ const subprocess = spawn('bash', [`${postStartupScripPath}`], { cwd: '/' }); ++ subprocess.stdout.pipe(logStream); ++ subprocess.stderr.pipe(logStream); ++ ++ res.statusCode = 200; ++ res.setHeader('Content-Type', 'application/json'); ++ res.end(JSON.stringify({ 'success': 'true' })); ++ } else { ++ serveError(req, res, 500, 'Poststartup script file not found at ' + postStartupScripPath); ++ } ++ } catch (error) { ++ serveError(req, res, 500, error.message); ++ } ++ } ++ } + } + + /** diff --git a/patches/series b/patches/series index e4cd73f85..fc0b16b2f 100644 --- a/patches/series +++ b/patches/series @@ -12,3 +12,4 @@ terminal-crash-mitigation.patch sagemaker-open-notebook-extension.patch security.diff sagemaker-ui-dark-theme.patch +sagemaker-ui-post-startup.patch From 14911e349ae48f78f802a57fab8d9ab49ec744d1 Mon Sep 17 00:00:00 2001 From: Mrunal Nitin Deshmukh Date: Thu, 27 Mar 2025 14:57:03 -0700 Subject: [PATCH 2/2] Fix typo for postStartupScriptPath --- patched-vscode/src/vs/server/node/webClientServer.ts | 10 +++++----- patches/sagemaker-ui-post-startup.patch | 10 +++++----- 2 files changed, 10 insertions(+), 10 deletions(-) diff --git a/patched-vscode/src/vs/server/node/webClientServer.ts b/patched-vscode/src/vs/server/node/webClientServer.ts index 09fca3bf5..744e112e0 100644 --- a/patched-vscode/src/vs/server/node/webClientServer.ts +++ b/patched-vscode/src/vs/server/node/webClientServer.ts @@ -498,7 +498,7 @@ export class WebClientServer { * Handles API requests to run the post-startup script in SMD. */ private async _handlePostStartupScriptInvocation(req: http.IncomingMessage, res: http.ServerResponse): Promise { - const postStartupScripPath = '/etc/sagemaker-ui/sagemaker_ui_post_startup.sh' + const postStartupScriptPath = '/etc/sagemaker-ui/sagemaker_ui_post_startup.sh' const logPath = '/var/log/apps/post_startup_default.log'; const logStream = fs.createWriteStream(logPath, { flags: 'a' }); @@ -509,11 +509,11 @@ export class WebClientServer { //If postStartupScriptFile doesn't exist, it will throw FileNotFoundError (404) //If exists, it will start the execution and add the execution logs in logFile. try { - if (fs.existsSync(postStartupScripPath)) { + if (fs.existsSync(postStartupScriptPath)) { // Adding 0o755 to make script file executable - fs.chmodSync(postStartupScripPath, 0o755); + fs.chmodSync(postStartupScriptPath, 0o755); - const subprocess = spawn('bash', [`${postStartupScripPath}`], { cwd: '/' }); + const subprocess = spawn('bash', [`${postStartupScriptPath}`], { cwd: '/' }); subprocess.stdout.pipe(logStream); subprocess.stderr.pipe(logStream); @@ -521,7 +521,7 @@ export class WebClientServer { res.setHeader('Content-Type', 'application/json'); res.end(JSON.stringify({ 'success': 'true' })); } else { - serveError(req, res, 500, 'Poststartup script file not found at ' + postStartupScripPath); + serveError(req, res, 500, 'Poststartup script file not found at ' + postStartupScriptPath); } } catch (error) { serveError(req, res, 500, error.message); diff --git a/patches/sagemaker-ui-post-startup.patch b/patches/sagemaker-ui-post-startup.patch index ba47f6e15..56e5ba58e 100644 --- a/patches/sagemaker-ui-post-startup.patch +++ b/patches/sagemaker-ui-post-startup.patch @@ -57,7 +57,7 @@ Index: sagemaker-code-editor/vscode/src/vs/server/node/webClientServer.ts + * Handles API requests to run the post-startup script in SMD. + */ + private async _handlePostStartupScriptInvocation(req: http.IncomingMessage, res: http.ServerResponse): Promise { -+ const postStartupScripPath = '/etc/sagemaker-ui/sagemaker_ui_post_startup.sh' ++ const postStartupScriptPath = '/etc/sagemaker-ui/sagemaker_ui_post_startup.sh' + const logPath = '/var/log/apps/post_startup_default.log'; + const logStream = fs.createWriteStream(logPath, { flags: 'a' }); + @@ -68,11 +68,11 @@ Index: sagemaker-code-editor/vscode/src/vs/server/node/webClientServer.ts + //If postStartupScriptFile doesn't exist, it will throw FileNotFoundError (404) + //If exists, it will start the execution and add the execution logs in logFile. + try { -+ if (fs.existsSync(postStartupScripPath)) { ++ if (fs.existsSync(postStartupScriptPath)) { + // Adding 0o755 to make script file executable -+ fs.chmodSync(postStartupScripPath, 0o755); ++ fs.chmodSync(postStartupScriptPath, 0o755); + -+ const subprocess = spawn('bash', [`${postStartupScripPath}`], { cwd: '/' }); ++ const subprocess = spawn('bash', [`${postStartupScriptPath}`], { cwd: '/' }); + subprocess.stdout.pipe(logStream); + subprocess.stderr.pipe(logStream); + @@ -80,7 +80,7 @@ Index: sagemaker-code-editor/vscode/src/vs/server/node/webClientServer.ts + res.setHeader('Content-Type', 'application/json'); + res.end(JSON.stringify({ 'success': 'true' })); + } else { -+ serveError(req, res, 500, 'Poststartup script file not found at ' + postStartupScripPath); ++ serveError(req, res, 500, 'Poststartup script file not found at ' + postStartupScriptPath); + } + } catch (error) { + serveError(req, res, 500, error.message);