Skip to content

Commit 09d4ac5

Browse files
toluojo-worktoluojo-work
authored
feat: Add recovery mode support for Unified Studio (#735)
1 parent 566b8ea commit 09d4ac5

File tree

4 files changed

+66
-22
lines changed

4 files changed

+66
-22
lines changed

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

Lines changed: 9 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -9,8 +9,15 @@ eval "$(micromamba shell hook --shell=bash)"
99
# apply patches for SMUS
1010
/etc/patches/apply_patches.sh smus && sudo rm -rf /etc/patches
1111

12-
# Activate conda environment 'base', where supervisord is installed
13-
micromamba activate base
12+
# Activate conda environment depending on if we are in Recovery or Standard mode.
13+
if [ -n "$SAGEMAKER_RECOVERY_MODE" ]; then
14+
export HOME=$SAGEMAKER_RECOVERY_MODE_HOME
15+
# Activate conda environment `sagemaker-recovery-mode`
16+
micromamba activate sagemaker-recovery-mode
17+
else
18+
# Activate conda environment 'base', where supervisord is installed
19+
micromamba activate base
20+
fi
1421

1522
export SAGEMAKER_APP_TYPE_LOWERCASE=$(echo $SAGEMAKER_APP_TYPE | tr '[:upper:]' '[:lower:]')
1623
export SERVICE_NAME='SageMakerUnifiedStudio'

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

Lines changed: 24 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -3,8 +3,14 @@ set -e
33

44
eval "$(micromamba shell hook --shell=bash)"
55

6-
# Activate conda environment 'base', which is the default environment for Cosmos
7-
micromamba activate base
6+
# Activate conda environment depending on if we are in Recovery or Standard mode.
7+
if [ -n "$SAGEMAKER_RECOVERY_MODE" ]; then
8+
# Activate conda environment `sagemaker-recovery-mode`
9+
micromamba activate sagemaker-recovery-mode
10+
else
11+
# Activate conda environment 'base' which is the default for Cosmos
12+
micromamba activate base
13+
fi
814

915
sudo cp -r /etc/sagemaker-ui/kernels/. /opt/conda/share/jupyter/kernels/
1016
sudo cp /etc/sagemaker-ui/jupyter/server/jupyter_server_config.py /opt/conda/etc/jupyter/
@@ -28,10 +34,19 @@ if [[ $(jupyter kernelspec list | grep glue_pyspark) ]]; then
2834
jupyter-kernelspec remove -f -y glue_pyspark
2935
fi
3036

31-
jupyter lab --ip 0.0.0.0 --port 8888 \
32-
--ServerApp.base_url="/$SAGEMAKER_APP_TYPE_LOWERCASE/default" \
33-
--ServerApp.token='' \
34-
--ServerApp.allow_origin='*' \
35-
--SchedulerApp.db_url='sqlite:////tmp/.jupyter_scheduler_do_not_delete.sqlite' \
36-
--collaborative \
37-
--ServerApp.identity_provider_class='sagemaker_jupyter_server_extension.identity.SageMakerIdentityProvider'
37+
if [ -n "$SAGEMAKER_RECOVERY_MODE" ]; then
38+
# Disabling collaboration and identity_provider_class flags for recovery mode
39+
jupyter lab --ip 0.0.0.0 --port 8888 \
40+
--ServerApp.base_url="/$SAGEMAKER_APP_TYPE_LOWERCASE/default" \
41+
--ServerApp.token='' \
42+
--ServerApp.allow_origin='*' \
43+
--SchedulerApp.db_url='sqlite:////tmp/.jupyter_scheduler_do_not_delete.sqlite'
44+
else
45+
jupyter lab --ip 0.0.0.0 --port 8888 \
46+
--ServerApp.base_url="/$SAGEMAKER_APP_TYPE_LOWERCASE/default" \
47+
--ServerApp.token='' \
48+
--ServerApp.allow_origin='*' \
49+
--SchedulerApp.db_url='sqlite:////tmp/.jupyter_scheduler_do_not_delete.sqlite' \
50+
--collaborative \
51+
--ServerApp.identity_provider_class='sagemaker_jupyter_server_extension.identity.SageMakerIdentityProvider'
52+
fi

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

Lines changed: 9 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -9,8 +9,15 @@ eval "$(micromamba shell hook --shell=bash)"
99
# apply patches for SMUS
1010
/etc/patches/apply_patches.sh smus && sudo rm -rf /etc/patches
1111

12-
# Activate conda environment 'base', where supervisord is installed
13-
micromamba activate base
12+
# Activate conda environment depending on if we are in Recovery or Standard mode.
13+
if [ -n "$SAGEMAKER_RECOVERY_MODE" ]; then
14+
export HOME=$SAGEMAKER_RECOVERY_MODE_HOME
15+
# Activate conda environment `sagemaker-recovery-mode`
16+
micromamba activate sagemaker-recovery-mode
17+
else
18+
# Activate conda environment 'base', where supervisord is installed
19+
micromamba activate base
20+
fi
1421

1522
export SAGEMAKER_APP_TYPE_LOWERCASE=$(echo $SAGEMAKER_APP_TYPE | tr '[:upper:]' '[:lower:]')
1623
export SERVICE_NAME='SageMakerUnifiedStudio'

‎template/v3/dirs/usr/local/bin/start-sagemaker-ui-jupyter-server

Lines changed: 24 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -3,8 +3,14 @@ set -e
33

44
eval "$(micromamba shell hook --shell=bash)"
55

6-
# Activate conda environment 'base', which is the default environment for Cosmos
7-
micromamba activate base
6+
# Activate conda environment depending on if we are in Recovery or Standard mode.
7+
if [ -n "$SAGEMAKER_RECOVERY_MODE" ]; then
8+
# Activate conda environment `sagemaker-recovery-mode`
9+
micromamba activate sagemaker-recovery-mode
10+
else
11+
# Activate conda environment 'base' which is the default for Cosmos
12+
micromamba activate base
13+
fi
814

915
sudo cp -r /etc/sagemaker-ui/kernels/. /opt/conda/share/jupyter/kernels/
1016
sudo cp /etc/sagemaker-ui/jupyter/server/jupyter_server_config.py /opt/conda/etc/jupyter/
@@ -28,10 +34,19 @@ if [[ $(jupyter kernelspec list | grep glue_pyspark) ]]; then
2834
jupyter-kernelspec remove -f -y glue_pyspark
2935
fi
3036

31-
jupyter lab --ip 0.0.0.0 --port 8888 \
32-
--ServerApp.base_url="/$SAGEMAKER_APP_TYPE_LOWERCASE/default" \
33-
--ServerApp.token='' \
34-
--ServerApp.allow_origin='*' \
35-
--SchedulerApp.db_url='sqlite:////tmp/.jupyter_scheduler_do_not_delete.sqlite' \
36-
--collaborative \
37-
--ServerApp.identity_provider_class='sagemaker_jupyter_server_extension.identity.SageMakerIdentityProvider'
37+
if [ -n "$SAGEMAKER_RECOVERY_MODE" ]; then
38+
# Disabling collaboration and identity_provider_class flags for recovery mode
39+
jupyter lab --ip 0.0.0.0 --port 8888 \
40+
--ServerApp.base_url="/$SAGEMAKER_APP_TYPE_LOWERCASE/default" \
41+
--ServerApp.token='' \
42+
--ServerApp.allow_origin='*' \
43+
--SchedulerApp.db_url='sqlite:////tmp/.jupyter_scheduler_do_not_delete.sqlite'
44+
else
45+
jupyter lab --ip 0.0.0.0 --port 8888 \
46+
--ServerApp.base_url="/$SAGEMAKER_APP_TYPE_LOWERCASE/default" \
47+
--ServerApp.token='' \
48+
--ServerApp.allow_origin='*' \
49+
--SchedulerApp.db_url='sqlite:////tmp/.jupyter_scheduler_do_not_delete.sqlite' \
50+
--collaborative \
51+
--ServerApp.identity_provider_class='sagemaker_jupyter_server_extension.identity.SageMakerIdentityProvider'
52+
fi

0 commit comments

Comments
 (0)