Skip to content

Commit c9e64a3

Browse files
aws-asoliduaws-asolidu
authored
Merge pull request #106 from arunbhati/feature/smus-sagemaker-extension
feat(session): add SageMaker Unified Studio portal redirect
2 parents 310e702 + 76c3d33 commit c9e64a3

File tree

4 files changed

+266
-23
lines changed

4 files changed

+266
-23
lines changed

‎patched-vscode/extensions/sagemaker-extension/src/constant.ts‎

Lines changed: 41 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -27,6 +27,10 @@ export const FIVE_MINUTES_INTERVAL_MILLIS = 5 * 60 * 1000;
2727

2828
export const SAGEMAKER_METADATA_PATH = '/opt/ml/metadata/resource-metadata.json';
2929

30+
// Service name identifier for SageMaker Unified Studio
31+
export const SMUS_SERVICE_NAME = 'SageMakerUnifiedStudio';
32+
export const SERVICE_NAME_ENV_VAR = 'SERVICE_NAME';
33+
3034
export class SagemakerCookie {
3135
authMode: string
3236
expiryTime: number
@@ -56,6 +60,11 @@ export class SagemakerResourceMetadata {
5660
ResourceArn?: string
5761
ResourceName?: string
5862
AppImageVersion?: string
63+
AdditionalMetadata?: {
64+
DataZoneDomainId?: string
65+
DataZoneProjectId?: string
66+
DataZoneDomainRegion?: string
67+
}
5968
};
6069
export function isSSOMode(cookie: SagemakerCookie) {
6170
return (cookie.authMode === AUTH_MODE.SSO)
@@ -69,4 +78,35 @@ export function getExpiryTime(cookie: SagemakerCookie): number {
6978
} else {
7079
return -1;
7180
}
72-
}
81+
}
82+
83+
/**
84+
* Constructs the SMUS portal URL using domain, region, and project information
85+
* Returns null if not in SMUS environment or if required fields are missing
86+
*/
87+
export const getSmusVscodePortalUrl = (metadata: SagemakerResourceMetadata | null): string | null => {
88+
if (process.env[SERVICE_NAME_ENV_VAR] !== SMUS_SERVICE_NAME) {
89+
return null;
90+
}
91+
92+
if (!metadata || !metadata.AdditionalMetadata) {
93+
// fail silently not to block users
94+
console.error('[SMUS] Metadata is undefined or null');
95+
return null;
96+
}
97+
98+
const { DataZoneDomainId, DataZoneDomainRegion, DataZoneProjectId } = metadata.AdditionalMetadata;
99+
100+
if (!DataZoneDomainId || !DataZoneDomainRegion || !DataZoneProjectId) {
101+
// fail silently not to block users
102+
// TODO: add monitoring to detect such cases
103+
console.error('[SMUS] Required fields missing in metadata:', {
104+
DataZoneDomainId: !!DataZoneDomainId,
105+
DataZoneDomainRegion: !!DataZoneDomainRegion,
106+
DataZoneProjectId: !!DataZoneProjectId
107+
});
108+
return null;
109+
}
110+
111+
return `https://${DataZoneDomainId}.sagemaker.${DataZoneDomainRegion}.on.aws/projects/${DataZoneProjectId}/overview`;
112+
}

‎patched-vscode/extensions/sagemaker-extension/src/extension.ts‎

Lines changed: 42 additions & 22 deletions
Original file line numberDiff line numberDiff line change
@@ -11,14 +11,33 @@ import {
1111
WARNING_BUTTON_SAVE_AND_RENEW_SESSION,
1212
SagemakerCookie,
1313
SagemakerResourceMetadata,
14-
getExpiryTime
14+
getExpiryTime,
15+
getSmusVscodePortalUrl
1516
} from "./constant";
1617
import * as console from "console";
1718

1819

1920
const PARSE_SAGEMAKER_COOKIE_COMMAND = 'sagemaker.parseCookies';
2021
const ENABLE_AUTO_UPDATE_COMMAND = 'workbench.extensions.action.enableAutoUpdate';
2122

23+
// Global redirect URL for SMUS environment
24+
let smusRedirectUrl: string | null = null;
25+
26+
function fetchMetadata(): SagemakerResourceMetadata | null {
27+
try {
28+
const data = fs.readFileSync(SAGEMAKER_METADATA_PATH, 'utf-8');
29+
return JSON.parse(data) as SagemakerResourceMetadata;
30+
} catch (error) {
31+
// fail silently not to block users
32+
console.error('Error reading metadata file:', error);
33+
return null;
34+
}
35+
}
36+
37+
function initializeSmusRedirectUrl() {
38+
smusRedirectUrl = getSmusVscodePortalUrl(fetchMetadata());
39+
}
40+
2241
function showWarningDialog() {
2342
vscode.commands.executeCommand(PARSE_SAGEMAKER_COOKIE_COMMAND).then(response => {
2443

@@ -59,11 +78,12 @@ function showWarningDialog() {
5978
}
6079

6180
function signInError(sagemakerCookie: SagemakerCookie) {
81+
const redirectUrl = getRedirectUrl(sagemakerCookie);
6282
// The session has expired
6383
SessionWarning.signInWarning(sagemakerCookie)
6484
.then((selection) => {
6585
if (selection === SIGN_IN_BUTTON) {
66-
vscode.env.openExternal(vscode.Uri.parse(<string>sagemakerCookie.redirectURL));
86+
vscode.env.openExternal(vscode.Uri.parse(redirectUrl));
6787
}
6888
});
6989
}
@@ -94,32 +114,21 @@ function saveWorkspace() {
94114
});
95115
}
96116
function renewSession(sagemakerCookie: SagemakerCookie) {
117+
const redirectUrl = getRedirectUrl(sagemakerCookie);
97118
// TODO: Log and trigger a Signin
98-
vscode.env.openExternal(vscode.Uri.parse(<string>sagemakerCookie.redirectURL));
119+
vscode.env.openExternal(vscode.Uri.parse(redirectUrl));
99120
// Trigger the function to show the warning again after 5 minutes again to validate.
100121
setTimeout(showWarningDialog, FIVE_MINUTES_INTERVAL_MILLIS);
101122
}
102123

103124
function updateStatusItemWithMetadata(context: vscode.ExtensionContext) {
104-
fs.readFile(SAGEMAKER_METADATA_PATH, 'utf-8', (err, data) => {
105-
if (err) {
106-
// fail silently not to block users
107-
} else {
108-
try {
109-
const jsonData = JSON.parse(data) as SagemakerResourceMetadata;
110-
const spaceName = jsonData.SpaceName;
111-
112-
if (spaceName != null) {
113-
let spaceNameStatusBarItem = vscode.window.createStatusBarItem(vscode.StatusBarAlignment.Left, 100);
114-
spaceNameStatusBarItem.text = `Space: ${spaceName}`;
115-
spaceNameStatusBarItem.show();
116-
context.subscriptions.push(spaceNameStatusBarItem);
117-
}
118-
} catch (jsonError) {
119-
// fail silently not to block users
120-
}
121-
}
122-
});
125+
const metadata = fetchMetadata();
126+
if (metadata?.SpaceName) {
127+
let spaceNameStatusBarItem = vscode.window.createStatusBarItem(vscode.StatusBarAlignment.Left, 100);
128+
spaceNameStatusBarItem.text = `Space: ${metadata.SpaceName}`;
129+
spaceNameStatusBarItem.show();
130+
context.subscriptions.push(spaceNameStatusBarItem);
131+
}
123132
}
124133

125134
// Render warning message regarding auto upgrade disabled
@@ -158,6 +167,9 @@ export function activate(context: vscode.ExtensionContext) {
158167
// TODO: log activation of extension
159168
console.log('Activating Sagemaker Extension...');
160169

170+
// First set smusRedirectUrl if we are in SMUS environment
171+
initializeSmusRedirectUrl();
172+
161173
// execute the get cookie command and save the data to cookies
162174
vscode.commands.executeCommand(PARSE_SAGEMAKER_COOKIE_COMMAND).then(r => {
163175

@@ -170,3 +182,11 @@ export function activate(context: vscode.ExtensionContext) {
170182
// render warning message regarding auto upgrade disabled
171183
renderExtensionAutoUpgradeDisabledNotification();
172184
}
185+
186+
/**
187+
* Returns the appropriate redirect URL based on the environment
188+
* Uses SMUS URL if available, falls back to original redirect URL
189+
*/
190+
function getRedirectUrl(sagemakerCookie: SagemakerCookie): string {
191+
return smusRedirectUrl || sagemakerCookie.redirectURL;
192+
}

‎patches/sagemaker-extension-smus-support.patch‎

Lines changed: 182 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,182 @@
1+
Index: sagemaker-code-editor/vscode/extensions/sagemaker-extension/src/constant.ts
2+
===================================================================
3+
--- sagemaker-code-editor.orig/vscode/extensions/sagemaker-extension/src/constant.ts
4+
+++ sagemaker-code-editor/vscode/extensions/sagemaker-extension/src/constant.ts
5+
@@ -27,6 +27,10 @@ export const FIVE_MINUTES_INTERVAL_MILLI
6+
7+
export const SAGEMAKER_METADATA_PATH = '/opt/ml/metadata/resource-metadata.json';
8+
9+
+// Service name identifier for SageMaker Unified Studio
10+
+export const SMUS_SERVICE_NAME = 'SageMakerUnifiedStudio';
11+
+export const SERVICE_NAME_ENV_VAR = 'SERVICE_NAME';
12+
+
13+
export class SagemakerCookie {
14+
authMode: string
15+
expiryTime: number
16+
@@ -56,6 +60,11 @@ export class SagemakerResourceMetadata {
17+
ResourceArn?: string
18+
ResourceName?: string
19+
AppImageVersion?: string
20+
+ AdditionalMetadata?: {
21+
+ DataZoneDomainId?: string
22+
+ DataZoneProjectId?: string
23+
+ DataZoneDomainRegion?: string
24+
+ }
25+
};
26+
export function isSSOMode(cookie: SagemakerCookie) {
27+
return (cookie.authMode === AUTH_MODE.SSO)
28+
@@ -69,4 +78,35 @@ export function getExpiryTime(cookie: Sa
29+
} else {
30+
return -1;
31+
}
32+
-}
33+
\ No newline at end of file
34+
+}
35+
+
36+
+/**
37+
+ * Constructs the SMUS portal URL using domain, region, and project information
38+
+ * Returns null if not in SMUS environment or if required fields are missing
39+
+ */
40+
+export const getSmusVscodePortalUrl = (metadata: SagemakerResourceMetadata | null): string | null => {
41+
+ if (process.env[SERVICE_NAME_ENV_VAR] !== SMUS_SERVICE_NAME) {
42+
+ return null;
43+
+ }
44+
+
45+
+ if (!metadata || !metadata.AdditionalMetadata) {
46+
+ // fail silently not to block users
47+
+ console.error('[SMUS] Metadata is undefined or null');
48+
+ return null;
49+
+ }
50+
+
51+
+ const { DataZoneDomainId, DataZoneDomainRegion, DataZoneProjectId } = metadata.AdditionalMetadata;
52+
+
53+
+ if (!DataZoneDomainId || !DataZoneDomainRegion || !DataZoneProjectId) {
54+
+ // fail silently not to block users
55+
+ // TODO: add monitoring to detect such cases
56+
+ console.error('[SMUS] Required fields missing in metadata:', {
57+
+ DataZoneDomainId: !!DataZoneDomainId,
58+
+ DataZoneDomainRegion: !!DataZoneDomainRegion,
59+
+ DataZoneProjectId: !!DataZoneProjectId
60+
+ });
61+
+ return null;
62+
+ }
63+
+
64+
+ return `https://${DataZoneDomainId}.sagemaker.${DataZoneDomainRegion}.on.aws/projects/${DataZoneProjectId}/overview`;
65+
+}
66+
Index: sagemaker-code-editor/vscode/extensions/sagemaker-extension/src/extension.ts
67+
===================================================================
68+
--- sagemaker-code-editor.orig/vscode/extensions/sagemaker-extension/src/extension.ts
69+
+++ sagemaker-code-editor/vscode/extensions/sagemaker-extension/src/extension.ts
70+
@@ -11,7 +11,8 @@ import {
71+
WARNING_BUTTON_SAVE_AND_RENEW_SESSION,
72+
SagemakerCookie,
73+
SagemakerResourceMetadata,
74+
- getExpiryTime
75+
+ getExpiryTime,
76+
+ getSmusVscodePortalUrl
77+
} from "./constant";
78+
import * as console from "console";
79+
80+
@@ -19,6 +20,24 @@ import * as console from "console";
81+
const PARSE_SAGEMAKER_COOKIE_COMMAND = 'sagemaker.parseCookies';
82+
const ENABLE_AUTO_UPDATE_COMMAND = 'workbench.extensions.action.enableAutoUpdate';
83+
84+
+// Global redirect URL for SMUS environment
85+
+let smusRedirectUrl: string | null = null;
86+
+
87+
+function fetchMetadata(): SagemakerResourceMetadata | null {
88+
+ try {
89+
+ const data = fs.readFileSync(SAGEMAKER_METADATA_PATH, 'utf-8');
90+
+ return JSON.parse(data) as SagemakerResourceMetadata;
91+
+ } catch (error) {
92+
+ // fail silently not to block users
93+
+ console.error('Error reading metadata file:', error);
94+
+ return null;
95+
+ }
96+
+}
97+
+
98+
+function initializeSmusRedirectUrl() {
99+
+ smusRedirectUrl = getSmusVscodePortalUrl(fetchMetadata());
100+
+}
101+
+
102+
function showWarningDialog() {
103+
vscode.commands.executeCommand(PARSE_SAGEMAKER_COOKIE_COMMAND).then(response => {
104+
105+
@@ -59,11 +78,12 @@ function showWarningDialog() {
106+
}
107+
108+
function signInError(sagemakerCookie: SagemakerCookie) {
109+
+ const redirectUrl = getRedirectUrl(sagemakerCookie);
110+
// The session has expired
111+
SessionWarning.signInWarning(sagemakerCookie)
112+
.then((selection) => {
113+
if (selection === SIGN_IN_BUTTON) {
114+
- vscode.env.openExternal(vscode.Uri.parse(<string>sagemakerCookie.redirectURL));
115+
+ vscode.env.openExternal(vscode.Uri.parse(redirectUrl));
116+
}
117+
});
118+
}
119+
@@ -94,32 +114,21 @@ function saveWorkspace() {
120+
});
121+
}
122+
function renewSession(sagemakerCookie: SagemakerCookie) {
123+
+ const redirectUrl = getRedirectUrl(sagemakerCookie);
124+
// TODO: Log and trigger a Signin
125+
- vscode.env.openExternal(vscode.Uri.parse(<string>sagemakerCookie.redirectURL));
126+
+ vscode.env.openExternal(vscode.Uri.parse(redirectUrl));
127+
// Trigger the function to show the warning again after 5 minutes again to validate.
128+
setTimeout(showWarningDialog, FIVE_MINUTES_INTERVAL_MILLIS);
129+
}
130+
131+
function updateStatusItemWithMetadata(context: vscode.ExtensionContext) {
132+
- fs.readFile(SAGEMAKER_METADATA_PATH, 'utf-8', (err, data) => {
133+
- if (err) {
134+
- // fail silently not to block users
135+
- } else {
136+
- try {
137+
- const jsonData = JSON.parse(data) as SagemakerResourceMetadata;
138+
- const spaceName = jsonData.SpaceName;
139+
-
140+
- if (spaceName != null) {
141+
- let spaceNameStatusBarItem = vscode.window.createStatusBarItem(vscode.StatusBarAlignment.Left, 100);
142+
- spaceNameStatusBarItem.text = `Space: ${spaceName}`;
143+
- spaceNameStatusBarItem.show();
144+
- context.subscriptions.push(spaceNameStatusBarItem);
145+
- }
146+
- } catch (jsonError) {
147+
- // fail silently not to block users
148+
- }
149+
- }
150+
- });
151+
+ const metadata = fetchMetadata();
152+
+ if (metadata?.SpaceName) {
153+
+ let spaceNameStatusBarItem = vscode.window.createStatusBarItem(vscode.StatusBarAlignment.Left, 100);
154+
+ spaceNameStatusBarItem.text = `Space: ${metadata.SpaceName}`;
155+
+ spaceNameStatusBarItem.show();
156+
+ context.subscriptions.push(spaceNameStatusBarItem);
157+
+ }
158+
}
159+
160+
// Render warning message regarding auto upgrade disabled
161+
@@ -158,6 +167,9 @@ export function activate(context: vscode
162+
// TODO: log activation of extension
163+
console.log('Activating Sagemaker Extension...');
164+
165+
+ // First set smusRedirectUrl if we are in SMUS environment
166+
+ initializeSmusRedirectUrl();
167+
+
168+
// execute the get cookie command and save the data to cookies
169+
vscode.commands.executeCommand(PARSE_SAGEMAKER_COOKIE_COMMAND).then(r => {
170+
171+
@@ -170,3 +182,11 @@ export function activate(context: vscode
172+
// render warning message regarding auto upgrade disabled
173+
renderExtensionAutoUpgradeDisabledNotification();
174+
}
175+
+
176+
+/**
177+
+ * Returns the appropriate redirect URL based on the environment
178+
+ * Uses SMUS URL if available, falls back to original redirect URL
179+
+ */
180+
+function getRedirectUrl(sagemakerCookie: SagemakerCookie): string {
181+
+ return smusRedirectUrl || sagemakerCookie.redirectURL;
182+
+}

‎patches/series‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -13,3 +13,4 @@ sagemaker-open-notebook-extension.patch
1313
security.diff
1414
sagemaker-ui-dark-theme.patch
1515
sagemaker-ui-post-startup.patch
16+
sagemaker-extension-smus-support.patch

0 commit comments

Comments
 (0)