Skip to content

Commit d2c8ec2

Browse files
mndeshmu97mndeshmu97
committed
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
1 parent 3cd6872 commit d2c8ec2

File tree

3 files changed

+149
-2
lines changed

3 files changed

+149
-2
lines changed

‎patched-vscode/src/vs/server/node/webClientServer.ts‎

Lines changed: 56 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -3,9 +3,11 @@
33
* Licensed under the MIT License. See License.txt in the project root for license information.
44
*--------------------------------------------------------------------------------------------*/
55

6-
import { createReadStream } from 'fs';
6+
import { createReadStream, existsSync, writeFileSync } from 'fs';
77
import {readFile } from 'fs/promises';
88
import { Promises } from 'vs/base/node/pfs';
9+
import { spawn } from 'child_process';
10+
import * as fs from 'fs';
911
import * as path from 'path';
1012
import * as http from 'http';
1113
import * as url from 'url';
@@ -39,6 +41,10 @@ const textMimeType: { [ext: string]: string | undefined } = {
3941
'.svg': 'image/svg+xml',
4042
};
4143

44+
const enum ServiceName {
45+
SAGEMAKER_UNIFIED_STUDIO = 'SageMakerUnifiedStudio',
46+
}
47+
4248
/**
4349
* Return an error to the client.
4450
*/
@@ -102,6 +108,7 @@ export class WebClientServer {
102108
private readonly _callbackRoute: string;
103109
private readonly _webExtensionRoute: string;
104110
private readonly _idleRoute: string;
111+
private readonly _postStartupScriptRoute: string;
105112

106113
constructor(
107114
private readonly _connectionToken: ServerConnectionToken,
@@ -118,6 +125,7 @@ export class WebClientServer {
118125
this._callbackRoute = `${serverRootPath}/callback`;
119126
this._webExtensionRoute = `${serverRootPath}/web-extension-resource`;
120127
this._idleRoute = '/api/idle';
128+
this._postStartupScriptRoute = '/api/poststartup';
121129
}
122130

123131
/**
@@ -146,6 +154,9 @@ export class WebClientServer {
146154
// extension resource support
147155
return this._handleWebExtensionResource(req, res, parsedUrl);
148156
}
157+
if (pathname === this._postStartupScriptRoute) {
158+
return this._handlePostStartupScriptInvocation(req, res);
159+
}
149160

150161
return serveError(req, res, 404, 'Not found.');
151162
} catch (error) {
@@ -459,12 +470,20 @@ export class WebClientServer {
459470
}
460471

461472
/**
462-
* Handles API requests to retrieve the last activity timestamp.
473+
* Handles API requests to retrieve the last activity timestamp.
463474
*/
464475
private async _handleIdle(req: http.IncomingMessage, res: http.ServerResponse): Promise<void> {
465476
try {
466477
const tmpDirectory = '/tmp/'
467478
const idleFilePath = path.join(tmpDirectory, '.sagemaker-last-active-timestamp');
479+
480+
// If idle shutdown file does not exist, this indicates the app UI may never been opened
481+
// Create the initial metadata file
482+
if (!existsSync(idleFilePath)) {
483+
const timestamp = new Date().toISOString();
484+
writeFileSync(idleFilePath, timestamp);
485+
}
486+
468487
const data = await readFile(idleFilePath, 'utf8');
469488

470489
res.statusCode = 200;
@@ -474,6 +493,41 @@ export class WebClientServer {
474493
serveError(req, res, 500, error.message)
475494
}
476495
}
496+
497+
/**
498+
* Handles API requests to run the post-startup script in SMD.
499+
*/
500+
private async _handlePostStartupScriptInvocation(req: http.IncomingMessage, res: http.ServerResponse): Promise<void> {
501+
const postStartupScripPath = '/etc/sagemaker-ui/sagemaker_ui_post_startup.sh'
502+
const logPath = '/var/log/apps/post_startup_default.log';
503+
const logStream = fs.createWriteStream(logPath, { flags: 'a' });
504+
505+
// Only trigger post-startup script invocation for SageMakerUnifiedStudio app.
506+
if (process.env['SERVICE_NAME'] != ServiceName.SAGEMAKER_UNIFIED_STUDIO) {
507+
return serveError(req, res, 403, 'Forbidden');
508+
} else {
509+
//If postStartupScriptFile doesn't exist, it will throw FileNotFoundError (404)
510+
//If exists, it will start the execution and add the execution logs in logFile.
511+
try {
512+
if (fs.existsSync(postStartupScripPath)) {
513+
// Adding 0o755 to make script file executable
514+
fs.chmodSync(postStartupScripPath, 0o755);
515+
516+
const subprocess = spawn('bash', [`${postStartupScripPath}`], { cwd: '/' });
517+
subprocess.stdout.pipe(logStream);
518+
subprocess.stderr.pipe(logStream);
519+
520+
res.statusCode = 200;
521+
res.setHeader('Content-Type', 'application/json');
522+
res.end(JSON.stringify({ 'success': 'true' }));
523+
} else {
524+
serveError(req, res, 500, 'Poststartup script file not found at ' + postStartupScripPath);
525+
}
526+
} catch (error) {
527+
serveError(req, res, 500, error.message);
528+
}
529+
}
530+
}
477531
}
478532

479533
/**

‎patches/sagemaker-ui-post-startup.patch‎

Lines changed: 92 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,92 @@
1+
Index: sagemaker-code-editor/vscode/src/vs/server/node/webClientServer.ts
2+
===================================================================
3+
--- sagemaker-code-editor.orig/vscode/src/vs/server/node/webClientServer.ts
4+
+++ sagemaker-code-editor/vscode/src/vs/server/node/webClientServer.ts
5+
@@ -6,6 +6,8 @@
6+
import { createReadStream, existsSync, writeFileSync } from 'fs';
7+
import {readFile } from 'fs/promises';
8+
import { Promises } from 'vs/base/node/pfs';
9+
+import { spawn } from 'child_process';
10+
+import * as fs from 'fs';
11+
import * as path from 'path';
12+
import * as http from 'http';
13+
import * as url from 'url';
14+
@@ -39,6 +41,10 @@ const textMimeType: { [ext: string]: str
15+
'.svg': 'image/svg+xml',
16+
};
17+
18+
+const enum ServiceName {
19+
+ SAGEMAKER_UNIFIED_STUDIO = 'SageMakerUnifiedStudio',
20+
+}
21+
+
22+
/**
23+
* Return an error to the client.
24+
*/
25+
@@ -102,6 +108,7 @@ export class WebClientServer {
26+
private readonly _callbackRoute: string;
27+
private readonly _webExtensionRoute: string;
28+
private readonly _idleRoute: string;
29+
+ private readonly _postStartupScriptRoute: string;
30+
31+
constructor(
32+
private readonly _connectionToken: ServerConnectionToken,
33+
@@ -118,6 +125,7 @@ export class WebClientServer {
34+
this._callbackRoute = `${serverRootPath}/callback`;
35+
this._webExtensionRoute = `${serverRootPath}/web-extension-resource`;
36+
this._idleRoute = '/api/idle';
37+
+ this._postStartupScriptRoute = '/api/poststartup';
38+
}
39+
40+
/**
41+
@@ -146,6 +154,9 @@ export class WebClientServer {
42+
// extension resource support
43+
return this._handleWebExtensionResource(req, res, parsedUrl);
44+
}
45+
+ if (pathname === this._postStartupScriptRoute) {
46+
+ return this._handlePostStartupScriptInvocation(req, res);
47+
+ }
48+
49+
return serveError(req, res, 404, 'Not found.');
50+
} catch (error) {
51+
@@ -482,6 +493,41 @@ export class WebClientServer {
52+
serveError(req, res, 500, error.message)
53+
}
54+
}
55+
+
56+
+ /**
57+
+ * Handles API requests to run the post-startup script in SMD.
58+
+ */
59+
+ private async _handlePostStartupScriptInvocation(req: http.IncomingMessage, res: http.ServerResponse): Promise<void> {
60+
+ const postStartupScripPath = '/etc/sagemaker-ui/sagemaker_ui_post_startup.sh'
61+
+ const logPath = '/var/log/apps/post_startup_default.log';
62+
+ const logStream = fs.createWriteStream(logPath, { flags: 'a' });
63+
+
64+
+ // Only trigger post-startup script invocation for SageMakerUnifiedStudio app.
65+
+ if (process.env['SERVICE_NAME'] != ServiceName.SAGEMAKER_UNIFIED_STUDIO) {
66+
+ return serveError(req, res, 403, 'Forbidden');
67+
+ } else {
68+
+ //If postStartupScriptFile doesn't exist, it will throw FileNotFoundError (404)
69+
+ //If exists, it will start the execution and add the execution logs in logFile.
70+
+ try {
71+
+ if (fs.existsSync(postStartupScripPath)) {
72+
+ // Adding 0o755 to make script file executable
73+
+ fs.chmodSync(postStartupScripPath, 0o755);
74+
+
75+
+ const subprocess = spawn('bash', [`${postStartupScripPath}`], { cwd: '/' });
76+
+ subprocess.stdout.pipe(logStream);
77+
+ subprocess.stderr.pipe(logStream);
78+
+
79+
+ res.statusCode = 200;
80+
+ res.setHeader('Content-Type', 'application/json');
81+
+ res.end(JSON.stringify({ 'success': 'true' }));
82+
+ } else {
83+
+ serveError(req, res, 500, 'Poststartup script file not found at ' + postStartupScripPath);
84+
+ }
85+
+ } catch (error) {
86+
+ serveError(req, res, 500, error.message);
87+
+ }
88+
+ }
89+
+ }
90+
}
91+
92+
/**

‎patches/series‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -12,3 +12,4 @@ terminal-crash-mitigation.patch
1212
sagemaker-open-notebook-extension.patch
1313
security.diff
1414
sagemaker-ui-dark-theme.patch
15+
sagemaker-ui-post-startup.patch

0 commit comments

Comments
 (0)