Skip to content

Auto enable s3ag plugin for SMAI/SMUS JL #885

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
wants to merge 7 commits into
base: main
Choose a base branch
from
Open

Auto enable s3ag plugin for SMAI/SMUS JL #885

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
7 commits
Select commit Hold shift + click to select a range
19b57d5
Update entry scripts to autoenable s3ag
claytonparnell Aug 12, 2025
7b247e4
Update pysdk defaults
claytonparnell Aug 13, 2025
91aff2d
Merge remote-tracking branch 'origin/main' into auto-enable-s3ag
claytonparnell Aug 27, 2025
9faf358
Formatting
claytonparnell Aug 27, 2025
37ba87d
Add to start scripts as well as entrypoint
claytonparnell Sep 3, 2025
6866428
Just set the variable in the start- scripts
claytonparnell Sep 4, 2025
1f06c3c
same changes for v3
claytonparnell Sep 4, 2025
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
2 changes: 1 addition & 1 deletion src/main.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -263,7 +263,7 @@ def _build_local_images(
# Minimal patch build, use .patch Dockerfiles
dockerfile = f"./Dockerfile-{image_type}.patch"
else:
dockerfile="./Dockerfile"
dockerfile = "./Dockerfile"
try:
image, log_gen = _docker_client.images.build(
path=target_version_dir, dockerfile=dockerfile, rm=True, pull=True, buildargs=config["build_args"]
Expand Down
32 changes: 16 additions & 16 deletions template/v2/dirs/etc/sagemaker-ui/sagemaker_ui_post_startup.sh
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -212,7 +212,7 @@ mkdir -p "$HOME/.config" # Create config directory if it doesn't exist
jq -n \
--arg smusProjectDirectory "$SMUS_PROJECT_DIR" \
--arg isGitProject "$IS_GIT_PROJECT" \
'{
'{
smusProjectDirectory: $smusProjectDirectory,
isGitProject: ($isGitProject == "true")
}' > "$HOME/.config/smus-storage-metadata.json"
Expand Down Expand Up @@ -247,7 +247,7 @@ if [ -f "$q_settings_file" ]; then
q_auth_mode=$(jq -r '.auth_mode' < $q_settings_file)
if [ "$q_auth_mode" == "IAM" ]; then
export AMAZON_Q_SIGV4=true
else
else
export AMAZON_Q_SIGV4=false
fi
else
Expand All @@ -260,7 +260,7 @@ if $AMAZON_Q_SIGV4; then
else
echo export AMAZON_Q_SIGV4=$AMAZON_Q_SIGV4 >> ~/.bashrc
fi
else
else
# Remove from .bashrc if it exists
sed -i '/^export AMAZON_Q_SIGV4=/d' ~/.bashrc
fi
Expand Down Expand Up @@ -293,7 +293,7 @@ if [ -f "$source_file" ]; then
if [ -f "$target_file" ]; then
# Target file exists - merge configurations
echo "Existing MCP configuration found, merging configurations..."

# Check if it's valid JSON first
if jq empty "$target_file" 2>/dev/null; then
# Initialize mcpServers object if it doesn't exist
Expand All @@ -302,9 +302,9 @@ if [ -f "$source_file" ]; then
jq '. + {"mcpServers":{}}' "$target_file" > "$target_file.tmp"
mv "$target_file.tmp" "$target_file"
fi

servers=$(jq '.mcpServers | keys[]' "$source_file" | tr -d '"')

# Add each server from source to target if it doesn't exist
for server in $servers; do
if ! jq -e ".mcpServers.\"$server\"" "$target_file" >/dev/null 2>&1; then
Expand All @@ -326,7 +326,7 @@ if [ -f "$source_file" ]; then
cp "$source_file" "$target_file"
echo "Created new MCP configuration with default servers"
fi

echo "Successfully configured MCP for SageMaker"
else
echo "Warning: MCP configuration file not found at $source_file"
Expand All @@ -339,10 +339,10 @@ agents_source_file="/etc/sagemaker-ui/sagemaker-mcp/default.json"

if [ -f "$agents_source_file" ]; then
mkdir -p "$HOME/.aws/amazonq/agents/"

if [ -f "$agents_target_file" ]; then
echo "Existing Amazon Q agents configuration found, merging mcpServers..."

# Check if target file is valid JSON
if jq empty "$agents_target_file" 2>/dev/null; then
# Initialize mcpServers object if it doesn't exist in target
Expand All @@ -351,11 +351,11 @@ if [ -f "$agents_source_file" ]; then
jq '. + {"mcpServers":{}}' "$agents_target_file" > "$agents_target_file.tmp"
mv "$agents_target_file.tmp" "$agents_target_file"
fi

# Add servers from source that don't exist in target and update tools
if jq -e '.mcpServers' "$agents_source_file" >/dev/null 2>&1; then
source_server_names=$(jq -r '.mcpServers | keys[]' "$agents_source_file")

for server_name in $source_server_names; do
if ! jq -e ".mcpServers.\"$server_name\"" "$agents_target_file" >/dev/null 2>&1; then
# Server doesn't exist in target - add it
Expand All @@ -364,7 +364,7 @@ if [ -f "$agents_source_file" ]; then
'.mcpServers[$name] = $config' "$agents_target_file" > "$agents_target_file.tmp"
mv "$agents_target_file.tmp" "$agents_target_file"
echo "Added server '$server_name' to agents configuration"

# Check if source has tools that reference this server and add them
server_tool_ref="@$server_name"
if jq -e --arg tool "$server_tool_ref" '.tools | index($tool)' "$agents_source_file" >/dev/null 2>&1; then
Expand All @@ -373,7 +373,7 @@ if [ -f "$agents_source_file" ]; then
jq '. + {"tools":[]}' "$agents_target_file" > "$agents_target_file.tmp"
mv "$agents_target_file.tmp" "$agents_target_file"
fi

# Add tool reference if it doesn't exist
if ! jq -e --arg tool "$server_tool_ref" '.tools | index($tool)' "$agents_target_file" >/dev/null 2>&1; then
jq --arg tool "$server_tool_ref" '.tools += [$tool]' "$agents_target_file" > "$agents_target_file.tmp"
Expand All @@ -385,7 +385,7 @@ if [ -f "$agents_source_file" ]; then
echo "Server '$server_name' already exists in configuration, skipping"
fi
done

echo "Successfully added missing mcpServers and tools from default.json to agents configuration"
else
echo "No mcpServers found in source configuration"
Expand All @@ -398,7 +398,7 @@ if [ -f "$agents_source_file" ]; then
cp "$agents_source_file" "$agents_target_file"
echo "Created new Amazon Q agents configuration file"
fi

echo "Successfully migrated MCP configuration to Amazon Q agents"
else
echo "Warning: Source configuration file not found at $agents_source_file"
Expand All @@ -420,7 +420,7 @@ if [ "${SAGEMAKER_APP_TYPE_LOWERCASE}" = "jupyterlab" ] && [ "$is_express_mode"

# write unexpected error to file if any of the remaining scripts fail.
trap 'write_status_to_file "error" "An unexpected error occurred. Please stop and restart your space to retry."' ERR

# Install conda and pip dependencies if lib mgmt config existing
bash /etc/sagemaker-ui/libmgmt/install-lib.sh

Expand Down
7 changes: 6 additions & 1 deletion template/v2/dirs/etc/sagemaker/sm_pysdk_default_config.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -36,7 +36,12 @@ def generate_intelligent_default_config(metadata: str) -> dict:
"Model": {"ExecutionRoleArn": metadata["UserRoleArn"]},
"ModelPackage": {"ValidationSpecification": {"ValidationRole": metadata["UserRoleArn"]}},
"ProcessingJob": {"RoleArn": metadata["UserRoleArn"]},
"TrainingJob": {"RoleArn": metadata["UserRoleArn"]},
"TrainingJob": {
"RoleArn": metadata["UserRoleArn"],
"Environment": {
"BOTOCORE_EXPERIMENTAL__PLUGINS": "S3AccessGrantsPlugin=aws_s3_access_grants_boto3_plugin.s3_access_grants_plugin"
},
},
},
}

Expand Down
6 changes: 6 additions & 0 deletions template/v2/dirs/usr/local/bin/start-jupyter-server
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -14,6 +14,12 @@ else
jupyter labextension disable @amzn/sagemaker-data-explorer-jl-plugin sagemaker-data-explorer:plugin amzn/sagemaker-ui-theme-jlplugin @amzn/sagemaker-ui-doc-manager-jl-plugin @amzn/sagemaker-connection-magics-jlextension @amzn/sagemaker_gen_ai_jupyterlab_extension @amzn/sagemaker-post-startup-notification-plugin:plugin @amzn/sagemaker-post-startup-notification-plugin:custom-terminal-launcher-plugin @amzn/sagemaker-studio-scheduler:scheduler @amzn/sagemaker-studio-scheduler:schedulerTelemetry
fi

# Enable S3AG plugin if TIP is enabled
if [ -n "$TRUSTED_IDENTITY_PROPOGATION_ENABLED" ]; then
export BOTOCORE_EXPERIMENTAL__PLUGINS=S3AccessGrantsPlugin=aws_s3_access_grants_boto3_plugin.s3_access_grants_plugin
echo BOTOCORE_EXPERIMENTAL__PLUGINS=S3AccessGrantsPlugin=aws_s3_access_grants_boto3_plugin.s3_access_grants_plugin >> ~/.bashrc
fi

# Start Jupyter server in rtc mode for shared spaces
if [ -n "$SAGEMAKER_APP_TYPE_LOWERCASE" ] && [ "$SAGEMAKER_SPACE_TYPE_LOWERCASE" == "shared" ]; then
jupyter labextension enable @jupyter/collaboration-extension
Expand Down
8 changes: 6 additions & 2 deletions template/v2/dirs/usr/local/bin/start-sagemaker-ui-jupyter-server
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -3,7 +3,7 @@ set -e

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

# Activate conda environment depending on if we are in Recovery or Standard mode.
# Activate conda environment depending on if we are in Recovery or Standard mode.
if [ -n "$SAGEMAKER_RECOVERY_MODE" ]; then
# Activate conda environment `sagemaker-recovery-mode`
micromamba activate sagemaker-recovery-mode
Expand Down Expand Up @@ -37,6 +37,10 @@ if [[ $(jupyter kernelspec list | grep glue_pyspark) ]]; then
jupyter-kernelspec remove -f -y glue_pyspark
fi

# Enable S3 Access Grant plugin by default
export BOTOCORE_EXPERIMENTAL__PLUGINS=S3AccessGrantsPlugin=aws_s3_access_grants_boto3_plugin.s3_access_grants_plugin
echo BOTOCORE_EXPERIMENTAL__PLUGINS=S3AccessGrantsPlugin=aws_s3_access_grants_boto3_plugin.s3_access_grants_plugin >> ~/.bashrc

if [ -n "$SAGEMAKER_RECOVERY_MODE" ]; then
# Disabling collaboration and identity_provider_class flags for recovery mode
jupyter lab --ip 0.0.0.0 --port 8888 \
Expand All @@ -46,7 +50,7 @@ if [ -n "$SAGEMAKER_RECOVERY_MODE" ]; then
--SchedulerApp.db_url='sqlite:////tmp/.jupyter_scheduler_do_not_delete.sqlite' \
--SQLiteYStore.db_path='/tmp/.ydoc_db_do_not_delete.sqlite' \
--ArbitraryFileIdManager.db_path='/tmp/.fileid_do_not_delete.sqlite'
else
else
jupyter lab --ip 0.0.0.0 --port 8888 \
--ServerApp.base_url="/$SAGEMAKER_APP_TYPE_LOWERCASE/default" \
--ServerApp.token='' \
Expand Down
32 changes: 16 additions & 16 deletions template/v3/dirs/etc/sagemaker-ui/sagemaker_ui_post_startup.sh
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -213,7 +213,7 @@ mkdir -p "$HOME/.config" # Create config directory if it doesn't exist
jq -n \
--arg smusProjectDirectory "$SMUS_PROJECT_DIR" \
--arg isGitProject "$IS_GIT_PROJECT" \
'{
'{
smusProjectDirectory: $smusProjectDirectory,
isGitProject: ($isGitProject == "true")
}' > "$HOME/.config/smus-storage-metadata.json"
Expand Down Expand Up @@ -248,7 +248,7 @@ if [ -f "$q_settings_file" ]; then
q_auth_mode=$(jq -r '.auth_mode' < $q_settings_file)
if [ "$q_auth_mode" == "IAM" ]; then
export AMAZON_Q_SIGV4=true
else
else
export AMAZON_Q_SIGV4=false
fi
else
Expand All @@ -261,7 +261,7 @@ if $AMAZON_Q_SIGV4; then
else
echo export AMAZON_Q_SIGV4=$AMAZON_Q_SIGV4 >> ~/.bashrc
fi
else
else
# Remove from .bashrc if it exists
sed -i '/^export AMAZON_Q_SIGV4=/d' ~/.bashrc
fi
Expand Down Expand Up @@ -294,7 +294,7 @@ if [ -f "$source_file" ]; then
if [ -f "$target_file" ]; then
# Target file exists - merge configurations
echo "Existing MCP configuration found, merging configurations..."

# Check if it's valid JSON first
if jq empty "$target_file" 2>/dev/null; then
# Initialize mcpServers object if it doesn't exist
Expand All @@ -303,9 +303,9 @@ if [ -f "$source_file" ]; then
jq '. + {"mcpServers":{}}' "$target_file" > "$target_file.tmp"
mv "$target_file.tmp" "$target_file"
fi

servers=$(jq '.mcpServers | keys[]' "$source_file" | tr -d '"')

# Add each server from source to target if it doesn't exist
for server in $servers; do
if ! jq -e ".mcpServers.\"$server\"" "$target_file" >/dev/null 2>&1; then
Expand All @@ -327,7 +327,7 @@ if [ -f "$source_file" ]; then
cp "$source_file" "$target_file"
echo "Created new MCP configuration with default servers"
fi

echo "Successfully configured MCP for SageMaker"
else
echo "Warning: MCP configuration file not found at $source_file"
Expand All @@ -340,10 +340,10 @@ agents_source_file="/etc/sagemaker-ui/sagemaker-mcp/default.json"

if [ -f "$agents_source_file" ]; then
mkdir -p "$HOME/.aws/amazonq/agents/"

if [ -f "$agents_target_file" ]; then
echo "Existing Amazon Q agents configuration found, merging mcpServers..."

# Check if target file is valid JSON
if jq empty "$agents_target_file" 2>/dev/null; then
# Initialize mcpServers object if it doesn't exist in target
Expand All @@ -352,11 +352,11 @@ if [ -f "$agents_source_file" ]; then
jq '. + {"mcpServers":{}}' "$agents_target_file" > "$agents_target_file.tmp"
mv "$agents_target_file.tmp" "$agents_target_file"
fi

# Add servers from source that don't exist in target and update tools
if jq -e '.mcpServers' "$agents_source_file" >/dev/null 2>&1; then
source_server_names=$(jq -r '.mcpServers | keys[]' "$agents_source_file")

for server_name in $source_server_names; do
if ! jq -e ".mcpServers.\"$server_name\"" "$agents_target_file" >/dev/null 2>&1; then
# Server doesn't exist in target - add it
Expand All @@ -365,7 +365,7 @@ if [ -f "$agents_source_file" ]; then
'.mcpServers[$name] = $config' "$agents_target_file" > "$agents_target_file.tmp"
mv "$agents_target_file.tmp" "$agents_target_file"
echo "Added server '$server_name' to agents configuration"

# Check if source has tools that reference this server and add them
server_tool_ref="@$server_name"
if jq -e --arg tool "$server_tool_ref" '.tools | index($tool)' "$agents_source_file" >/dev/null 2>&1; then
Expand All @@ -374,7 +374,7 @@ if [ -f "$agents_source_file" ]; then
jq '. + {"tools":[]}' "$agents_target_file" > "$agents_target_file.tmp"
mv "$agents_target_file.tmp" "$agents_target_file"
fi

# Add tool reference if it doesn't exist
if ! jq -e --arg tool "$server_tool_ref" '.tools | index($tool)' "$agents_target_file" >/dev/null 2>&1; then
jq --arg tool "$server_tool_ref" '.tools += [$tool]' "$agents_target_file" > "$agents_target_file.tmp"
Expand All @@ -386,7 +386,7 @@ if [ -f "$agents_source_file" ]; then
echo "Server '$server_name' already exists in configuration, skipping"
fi
done

echo "Successfully added missing mcpServers and tools from default.json to agents configuration"
else
echo "No mcpServers found in source configuration"
Expand All @@ -399,7 +399,7 @@ if [ -f "$agents_source_file" ]; then
cp "$agents_source_file" "$agents_target_file"
echo "Created new Amazon Q agents configuration file"
fi

echo "Successfully migrated MCP configuration to Amazon Q agents"
else
echo "Warning: Source configuration file not found at $agents_source_file"
Expand All @@ -421,7 +421,7 @@ if [ "${SAGEMAKER_APP_TYPE_LOWERCASE}" = "jupyterlab" ] && [ "$is_express_mode"

# write unexpected error to file if any of the remaining scripts fail.
trap 'write_status_to_file "error" "An unexpected error occurred. Please stop and restart your space to retry."' ERR

# Install conda and pip dependencies if lib mgmt config existing
bash /etc/sagemaker-ui/libmgmt/install-lib.sh

Expand Down
7 changes: 6 additions & 1 deletion template/v3/dirs/etc/sagemaker/sm_pysdk_default_config.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -37,7 +37,12 @@ def generate_intelligent_default_config(metadata: str) -> dict:
"Model": {"ExecutionRoleArn": metadata["UserRoleArn"]},
"ModelPackage": {"ValidationSpecification": {"ValidationRole": metadata["UserRoleArn"]}},
"ProcessingJob": {"RoleArn": metadata["UserRoleArn"]},
"TrainingJob": {"RoleArn": metadata["UserRoleArn"]},
"TrainingJob": {
"RoleArn": metadata["UserRoleArn"],
"Environment": {
"BOTOCORE_EXPERIMENTAL__PLUGINS": "S3AccessGrantsPlugin=aws_s3_access_grants_boto3_plugin.s3_access_grants_plugin"
},
},
},
}

Expand Down
6 changes: 6 additions & 0 deletions template/v3/dirs/usr/local/bin/start-jupyter-server
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -14,6 +14,12 @@ else
jupyter labextension disable @amzn/sagemaker-data-explorer-jl-plugin sagemaker-data-explorer:plugin amzn/sagemaker-ui-theme-jlplugin @amzn/sagemaker-ui-doc-manager-jl-plugin @amzn/sagemaker-connection-magics-jlextension @amzn/sagemaker_gen_ai_jupyterlab_extension @amzn/sagemaker-post-startup-notification-plugin:plugin @amzn/sagemaker-post-startup-notification-plugin:custom-terminal-launcher-plugin @amzn/sagemaker-studio-scheduler:scheduler @amzn/sagemaker-studio-scheduler:schedulerTelemetry
fi

# Enable S3AG plugin if TIP is enabled
if [ -n "$TRUSTED_IDENTITY_PROPOGATION_ENABLED" ]; then
export BOTOCORE_EXPERIMENTAL__PLUGINS=S3AccessGrantsPlugin=aws_s3_access_grants_boto3_plugin.s3_access_grants_plugin
echo BOTOCORE_EXPERIMENTAL__PLUGINS=S3AccessGrantsPlugin=aws_s3_access_grants_boto3_plugin.s3_access_grants_plugin >> ~/.bashrc
fi

# Start Jupyter server in rtc mode for shared spaces
if [ -n "$SAGEMAKER_APP_TYPE_LOWERCASE" ] && [ "$SAGEMAKER_SPACE_TYPE_LOWERCASE" == "shared" ]; then
jupyter labextension enable @jupyter/collaboration-extension
Expand Down
14 changes: 9 additions & 5 deletions template/v3/dirs/usr/local/bin/start-sagemaker-ui-jupyter-server
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -3,7 +3,7 @@ set -e

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

# Activate conda environment depending on if we are in Recovery or Standard mode.
# Activate conda environment depending on if we are in Recovery or Standard mode.
if [ -n "$SAGEMAKER_RECOVERY_MODE" ]; then
# Activate conda environment `sagemaker-recovery-mode`
micromamba activate sagemaker-recovery-mode
Expand All @@ -13,7 +13,7 @@ else

# Disable jupyter-ai in favor of sagemaker_gen_ai_jupyterlab_extension
jupyter labextension disable @jupyter-ai/core

# Enable RTC to allow async Q file updates
jupyter labextension enable @jupyter/docprovider-extension
fi
Expand All @@ -40,19 +40,23 @@ if [[ $(jupyter kernelspec list | grep glue_pyspark) ]]; then
jupyter-kernelspec remove -f -y glue_pyspark
fi

# Enable S3 Access Grant plugin by default
export BOTOCORE_EXPERIMENTAL__PLUGINS=S3AccessGrantsPlugin=aws_s3_access_grants_boto3_plugin.s3_access_grants_plugin
echo BOTOCORE_EXPERIMENTAL__PLUGINS=S3AccessGrantsPlugin=aws_s3_access_grants_boto3_plugin.s3_access_grants_plugin >> ~/.bashrc

if [ -n "$SAGEMAKER_RECOVERY_MODE" ]; then
# Disabling collaboration and identity_provider_class flags for recovery mode
jupyter lab --ip 0.0.0.0 --port 8888 \
--ServerApp.base_url="/$SAGEMAKER_APP_TYPE_LOWERCASE/default" \
--ServerApp.token='' \
--ServerApp.allow_origin='*' \
--SchedulerApp.db_url='sqlite:////tmp/.jupyter_scheduler_do_not_delete.sqlite'
else
else
jupyter lab --ip 0.0.0.0 --port 8888 \
--ServerApp.base_url="/$SAGEMAKER_APP_TYPE_LOWERCASE/default" \
--ServerApp.token='' \
--ServerApp.allow_origin='*' \
--SchedulerApp.db_url='sqlite:////tmp/.jupyter_scheduler_do_not_delete.sqlite' \
--collaborative \
--ServerApp.identity_provider_class='sagemaker_jupyter_server_extension.identity.SageMakerIdentityProvider'
fi
--ServerApp.identity_provider_class='sagemaker_jupyter_server_extension.identity.SageMakerIdentityProvider'
fi