Skip to content

Commit 85aaf8f

Browse files
claytonparnellclaytonparnell
authored
Merge SMUS changes into template (#605)
1 parent 45959ee commit 85aaf8f

File tree

6 files changed

+234
-9
lines changed

6 files changed

+234
-9
lines changed

‎template/v2/dirs/etc/sagemaker-ui/sagemaker_ui_post_startup.sh

Lines changed: 59 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,48 @@
11
#!/bin/bash
22
set -eux
33

4+
# Writes script status to file. This file is read by an IDE extension responsible for dispatching UI post-startup-status to the user.
5+
write_status_to_file() {
6+
local status="$1"
7+
local message="$2"
8+
local file="/tmp/.post-startup-status.json"
9+
10+
# Check if the file exists, if not, create it
11+
if [ ! -f "$file" ]; then
12+
touch "$file" || {
13+
echo "Failed to create $file" >&2
14+
return 0
15+
}
16+
fi
17+
18+
# Ensure the file is writable
19+
if [ ! -w "$file" ]; then
20+
echo "Error: $file is not writable" >&2
21+
return 0
22+
fi
23+
24+
# Create the JSON object and write to file
25+
jq -n --arg status "$status" --arg message "$message" '{"status":$status,"message":$message}' > "$file"
26+
27+
}
28+
29+
# checks if the script status is "in-progress". If so, no errors were detected and it can be marked successful.
30+
write_status_to_file_on_script_complete() {
31+
local file="/tmp/.post-startup-status.json"
32+
local check_key="status"
33+
local check_value="in-progress"
34+
35+
36+
if jq -e --arg key "$check_key" --arg value "$check_value" '.[$key] == $value' "$file" > /dev/null; then
37+
write_status_to_file "success" "IDE configured successfully."
38+
echo "Post-startup script completed successfully. Success status written to $file"
39+
else
40+
echo "Skipping writing post-startup script "success" status. An error was detected during execution and written to $file".
41+
fi
42+
}
43+
44+
write_status_to_file "in-progress" "IDE configuration in progress."
45+
446
sourceMetaData=/opt/ml/metadata/resource-metadata.json
547

648
# Extract the required fields from meta data stored in opt/ml/metadata.
@@ -56,6 +98,7 @@ set -x
5698

5799
if [ $responseStatus -ne 0 ]; then
58100
echo "Failed to fetch domain execution role credentials. Will skip adding new credentials profile: DomainExecutionRoleCreds."
101+
write_status_to_file "error" "Network issue detected. Your domain may be using a public subnet, which affects IDE functionality. Please contact your administrator."
59102
else
60103
aws configure set credential_process "sagemaker-studio credentials get-domain-execution-role-credential-in-space --domain-id $dataZoneDomainId --profile default" --profile DomainExecutionRoleCreds
61104
echo "Successfully configured DomainExecutionRoleCreds profile"
@@ -71,19 +114,19 @@ else
71114
response=$( aws datazone get-user-profile --domain-identifier "$dataZoneDomainId" --user-identifier "$dataZoneUserId" --region "$dataZoneDomainRegion" )
72115
fi
73116

74-
# Extract the Auth Mode from the response. Unified Studio currently supports IAM, SSO and SAML.
117+
# Extract the Auth Mode from the response. Unified Studio currently supports IAM, SSO and SAML.
75118
auth_mode=$(echo "$response" | jq -r '.type')
76119

77120
case "$auth_mode" in
78121
"IAM")
79-
# For IAM users - extract IAM ARN from response. Response does not contain username or email.
122+
# For IAM users - extract IAM ARN from response. Response does not contain username or email.
80123
arn=$(echo "$response" | jq -r '.details.iam.arn')
81124
# Split ARN by / and return the last field
82125
username=$(echo "$arn" | awk -F'/' '{print $NF}')
83126
email="$arn"
84127
;;
85128
"SSO"|"SAML")
86-
# For SSO and SAML user, extract username and email if present in response.
129+
# For SSO and SAML user, extract username and email if present in response.
87130
username=$(echo "$response" | jq -r '.details.sso.username')
88131
email=$(echo "$response" | jq -r '.details.sso.email')
89132
# Setting up the email as username if email is not present
@@ -112,14 +155,21 @@ fi
112155

113156
set -e
114157

158+
# write unexpected error to file if any of the remaining scripts fail.
159+
trap 'write_status_to_file "error" "An unexpected error occurred. Please stop and restart your space to retry."' ERR
160+
115161
# Generate sagemaker pysdk intelligent default config
116162
nohup python /etc/sagemaker/sm_pysdk_default_config.py &
163+
# Only run the following commands if SAGEMAKER_APP_TYPE_LOWERCASE is jupyterlab
164+
if [ "${SAGEMAKER_APP_TYPE_LOWERCASE}" = "jupyterlab" ]; then
165+
# Start workflows local runner
166+
bash /etc/sagemaker-ui/workflows/start-workflows-container.sh
117167

118-
# Start workflows local runner
119-
bash /etc/sagemaker-ui/workflows/start-workflows-container.sh
168+
# Install conda and pip dependencies if lib mgmt config existing
169+
bash /etc/sagemaker-ui/libmgmt/install-lib.sh $HOME/src
120170

121-
# Install conda and pip dependencies if lib mgmt config existing
122-
bash /etc/sagemaker-ui/libmgmt/install-lib.sh $HOME/src
171+
# Install sm-spark-cli
172+
bash /etc/sagemaker-ui/workflows/sm-spark-cli-install.sh
173+
fi
123174

124-
# Install sm-spark-cli
125-
bash /etc/sagemaker-ui/workflows/sm-spark-cli-install.sh
175+
write_status_to_file_on_script_complete

‎template/v2/dirs/etc/sagemaker-ui/set_code_editor_theme.sh

Lines changed: 23 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,23 @@
1+
#!/bin/bash
2+
set -e
3+
4+
USER_SETTINGS_FILE="/home/sagemaker-user/sagemaker-code-editor-server-data/data/User/settings.json"
5+
COLOR_THEME_KEY="workbench.colorTheme"
6+
COLOR_THEME_VALUE="Default Dark Modern"
7+
8+
# Check if the settings file exists
9+
if [ ! -f "$USER_SETTINGS_FILE" ]; then
10+
echo "Code Editor user settings file not found. Creating..."
11+
mkdir -p "$(dirname "$USER_SETTINGS_FILE")"
12+
echo "{}" > "$USER_SETTINGS_FILE"
13+
fi
14+
15+
EXISTING_COLOR_THEME_VALUE=$(jq -r --arg key "$COLOR_THEME_KEY" '.[$key] // empty' "$USER_SETTINGS_FILE")
16+
17+
if [[ -n "$EXISTING_COLOR_THEME_VALUE" ]]; then
18+
echo "Theme is already set in user settings as '$EXISTING_COLOR_THEME_VALUE'. No changes made."
19+
else
20+
# Set theme
21+
jq --arg key "$COLOR_THEME_KEY" --arg value "$COLOR_THEME_VALUE" '.[$key] = $value' "$USER_SETTINGS_FILE" > "${USER_SETTINGS_FILE}.tmp" && mv "${USER_SETTINGS_FILE}.tmp" "$USER_SETTINGS_FILE"
22+
echo "Successfully set Code Editor theme to '$COLOR_THEME_VALUE'."
23+
fi

‎template/v2/dirs/etc/supervisor/conf.d/supervisord-sagemaker-ui-code-editor.conf

Lines changed: 28 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,28 @@
1+
[supervisord]
2+
nodaemon=true
3+
4+
pidfile=/var/run/supervisord/supervisord.pid
5+
logfile=%(ENV_STUDIO_LOGGING_DIR)s/%(ENV_SAGEMAKER_APP_TYPE_LOWERCASE)s/supervisord/supervisord.log
6+
logfile_maxbytes=5MB
7+
logfile_backups=10
8+
redirect_stderr=true
9+
10+
[unix_http_server]
11+
file=/var/run/supervisord/supervisor.sock
12+
chmod=0700
13+
14+
[supervisorctl]
15+
serverurl=unix:///var/run/supervisord/supervisor.sock
16+
17+
[rpcinterface:supervisor]
18+
supervisor.rpcinterface_factory = supervisor.rpcinterface:make_main_rpcinterface
19+
20+
[program:codeeditorserver]
21+
directory=%(ENV_HOME)s
22+
command=start-sagemaker-ui-code-editor
23+
autostart=true
24+
autorestart=true
25+
stdout_logfile=/dev/stdout
26+
stdout_logfile_maxbytes=0
27+
stderr_logfile=/dev/stderr
28+
stderr_logfile_maxbytes=0

‎template/v2/dirs/usr/local/bin/entrypoint-sagemaker-ui-code-editor

Lines changed: 16 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,16 @@
1+
#!/bin/bash
2+
3+
set -e
4+
5+
# Generate and execute the shell code to modifies shell variables to include
6+
# micromamba commands (e.g. using `micromamba activate` to activate environments)
7+
eval "$(micromamba shell hook --shell=bash)"
8+
9+
# Activate conda environment 'base', where supervisord is installed
10+
micromamba activate base
11+
12+
export SAGEMAKER_APP_TYPE_LOWERCASE=$(echo $SAGEMAKER_APP_TYPE | tr '[:upper:]' '[:lower:]')
13+
export SERVICE_NAME='SageMakerUnifiedStudio'
14+
15+
mkdir -p $STUDIO_LOGGING_DIR/$SAGEMAKER_APP_TYPE_LOWERCASE/supervisord
16+
exec supervisord -c /etc/supervisor/conf.d/supervisord-sagemaker-ui-code-editor.conf -n

‎template/v2/dirs/usr/local/bin/entrypoint-sagemaker-ui-jupyter-server

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -10,6 +10,7 @@ eval "$(micromamba shell hook --shell=bash)"
1010
micromamba activate base
1111

1212
export SAGEMAKER_APP_TYPE_LOWERCASE=$(echo $SAGEMAKER_APP_TYPE | tr '[:upper:]' '[:lower:]')
13+
export SERVICE_NAME='SageMakerUnifiedStudio'
1314

1415
mkdir -p $STUDIO_LOGGING_DIR/$SAGEMAKER_APP_TYPE_LOWERCASE/supervisord
1516
exec supervisord -c /etc/supervisor/conf.d/supervisord-sagemaker-ui.conf -n

‎template/v2/dirs/usr/local/bin/start-sagemaker-ui-code-editor

Lines changed: 107 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,107 @@
1+
#!/bin/bash
2+
set -e
3+
4+
EFS_MOUNT_POINT="/opt/amazon/sagemaker"
5+
EBS_MOUNT_POINT="/home/sagemaker-user"
6+
7+
persistent_settings_folder="${EBS_MOUNT_POINT}/sagemaker-code-editor-server-data"
8+
default_settings_folder="${EFS_MOUNT_POINT}/sagemaker-code-editor-server-data"
9+
10+
override_machine_settings() {
11+
# create a new settings file with preset defaults or merge the defaults into the existing settings file
12+
local settings_relative_path="data/Machine"
13+
local settings_file_path_suffix="${settings_relative_path}/settings.json"
14+
local persistent_machine_settings_file="${persistent_settings_folder}/${settings_file_path_suffix}"
15+
local default_machine_settings_file="${default_settings_folder}/${settings_file_path_suffix}"
16+
17+
if [ ! -f "$persistent_machine_settings_file" ]; then
18+
# copy settings file to EBS if it doesn't exist in EBS
19+
mkdir -p "${persistent_settings_folder}/${settings_relative_path}"
20+
cp "$default_machine_settings_file" "$persistent_machine_settings_file"
21+
echo "Created persistent settings file with default settings at $persistent_machine_settings_file"
22+
else
23+
# if it does exist then merge settings
24+
echo "File already exists: ${persistent_machine_settings_file}. Merging default settings with existing settings."
25+
python3 /usr/local/bin/merge-settings-util.py "$persistent_machine_settings_file" "$default_machine_settings_file"
26+
fi
27+
}
28+
29+
copy_user_settings() {
30+
local settings_relative_path="data/User"
31+
local settings_file_path_suffix="${settings_relative_path}/settings.json"
32+
local persistent_user_settings_file="${persistent_settings_folder}/${settings_file_path_suffix}"
33+
local default_user_settings_file="${default_settings_folder}/${settings_file_path_suffix}"
34+
if [ ! -f "$persistent_user_settings_file" ]; then
35+
# copy user settings file to EBS if it doesn't exist in EBS
36+
mkdir -p "${persistent_settings_folder}/${settings_relative_path}"
37+
cp "$default_user_settings_file" "$persistent_user_settings_file"
38+
echo "Created persistent settings file with default settings at $persistent_user_settings_file"
39+
fi
40+
}
41+
42+
install_prepackaged_extensions() {
43+
local prepackaged_extensions_dir="${default_settings_folder}/extensions"
44+
local persistent_extensions_dir="${persistent_settings_folder}/extensions"
45+
46+
# if extensions directory doesn't exist then this is the first time opening the app
47+
if [ ! -d "${persistent_extensions_dir}" ]; then
48+
mkdir -p $persistent_extensions_dir
49+
50+
# find all pre-packaged extensions folders, symlink them to the persistent volume
51+
find "$prepackaged_extensions_dir" -mindepth 1 -maxdepth 1 -type d | while read -r dir; do
52+
dir_name=$(basename "$dir")
53+
54+
ln -s "$dir" "${persistent_extensions_dir}/${dir_name}"
55+
echo "Added extension: $dir"
56+
done
57+
fi
58+
}
59+
60+
set_code_editor_theme() {
61+
echo "Setting code editor theme"
62+
bash /etc/sagemaker-ui/set_code_editor_theme.sh
63+
}
64+
65+
disableJupyterKernels() {
66+
local kernel_list
67+
kernel_list=$(jupyter kernelspec list)
68+
69+
for kernel in "$@"; do
70+
if echo "$kernel_list" | grep -q "$kernel"; then
71+
echo "Removing kernel: $kernel"
72+
jupyter kernelspec remove -f -y "$kernel"
73+
else
74+
echo "Kernel not found: $kernel"
75+
fi
76+
done
77+
}
78+
79+
eval "$(micromamba shell hook --shell=bash)"
80+
81+
# Activate conda environment 'base', which is the default environment for sagemaker-distribution
82+
micromamba activate base
83+
84+
# Disable Jupyter Kernels
85+
disableJupyterKernels pysparkkernel sparkkernel glue_spark glue_pyspark
86+
87+
# Start code-editor server
88+
if [ -n "$SAGEMAKER_APP_TYPE_LOWERCASE" ]; then
89+
# SAGEMAKER_APP_TYPE is set, indicating the server is running within a SageMaker app.
90+
override_machine_settings
91+
copy_user_settings
92+
install_prepackaged_extensions
93+
set_code_editor_theme
94+
# Configure the base url to be `/<app-type-in-lower-case>/default`.
95+
sagemaker-code-editor --host 0.0.0.0 --port 8888 \
96+
--without-connection-token \
97+
--base-path "/$SAGEMAKER_APP_TYPE_LOWERCASE/default" \
98+
--server-data-dir $persistent_settings_folder \
99+
--extensions-dir ${persistent_settings_folder}/extensions \
100+
--user-data-dir /opt/amazon/sagemaker/sagemaker-code-editor-user-data
101+
else
102+
sagemaker-code-editor --host 0.0.0.0 --port 8888 \
103+
--without-connection-token \
104+
--server-data-dir /opt/amazon/sagemaker/sagemaker-code-editor-server-data \
105+
--extension-dir /opt/amazon/sagemaker/sagemaker-code-editor-server-data/extensions \
106+
--user-data-dir /opt/amazon/sagemaker/sagemaker-code-editor-user-data
107+
fi

0 commit comments

Comments
 (0)