Skip to content

install sm-spark-cli only if sagemaker workflows is enabled #915

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

Merged
claytonparnell merged 1 commit into from
Sep 29, 2025
Merged

install sm-spark-cli only if sagemaker workflows is enabled #915

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
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 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 @@ -425,7 +425,7 @@ if [ "${SAGEMAKER_APP_TYPE_LOWERCASE}" = "jupyterlab" ] && [ "$is_express_mode"
bash /etc/sagemaker-ui/libmgmt/install-lib.sh

# Install sm-spark-cli
bash /etc/sagemaker-ui/workflows/sm-spark-cli-install.sh
bash /etc/sagemaker-ui/workflows/sm-spark-cli-install.sh || echo "Warning: sm-spark-cli installation failed, continuing..."
Copy link
Contributor

@reganbaum reganbaum Sep 12, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

nit: can we put this print statement in the install script to keep the post startup script clean?

Copy link
Contributor Author

@danobi06 danobi06 Sep 12, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Extending this one line is as clean as we can get to safe installing the script.

fi

# Execute network validation script, to check if any required AWS Services are unreachable
Expand Down
23 changes: 16 additions & 7 deletions template/v2/dirs/etc/sagemaker-ui/workflows/sm-spark-cli-install.sh
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,9 +1,18 @@
#!/bin/bash
RESOURCE_METADATA_FILE=/opt/ml/metadata/resource-metadata.json
DZ_DOMAIN_ID=$(jq -r '.AdditionalMetadata.DataZoneDomainId' < $RESOURCE_METADATA_FILE)
DZ_PROJECT_ID=$(jq -r '.AdditionalMetadata.DataZoneProjectId' < $RESOURCE_METADATA_FILE)
DZ_DOMAIN_REGION=$(jq -r '.AdditionalMetadata.DataZoneDomainRegion' < $RESOURCE_METADATA_FILE)
DZ_ENDPOINT=$(jq -r '.AdditionalMetadata.DataZoneEndpoint' < $RESOURCE_METADATA_FILE)

# install sm-spark-cli
sudo curl -LO https://github.com/aws-samples/amazon-sagemaker-spark-ui/releases/download/v0.9.1/amazon-sagemaker-spark-ui.tar.gz && \
sudo tar -xvzf amazon-sagemaker-spark-ui.tar.gz && \
sudo chmod +x amazon-sagemaker-spark-ui/install-scripts/studio/install-history-server.sh && \
sudo amazon-sagemaker-spark-ui/install-scripts/studio/install-history-server.sh && \
rm -rf ~/.m2 && \
sudo rm -rf amazon-sagemaker-spark-ui*
# install sm-spark-cli if workflows blueprint is enabled
if [ "$(python /etc/sagemaker-ui/workflows/workflow_client.py check-blueprint --region "$DZ_DOMAIN_REGION" --domain-id "$DZ_DOMAIN_ID" --endpoint "$DZ_ENDPOINT" --project-id "$DZ_PROJECT_ID")" = "True" ]; then
echo "Workflows blueprint is enabled. Installing sm-spark-cli."
# install sm-spark-cli
sudo curl -LO https://github.com/aws-samples/amazon-sagemaker-spark-ui/releases/download/v0.9.1/amazon-sagemaker-spark-ui.tar.gz && \
sudo tar -xvzf amazon-sagemaker-spark-ui.tar.gz && \
sudo chmod +x amazon-sagemaker-spark-ui/install-scripts/studio/install-history-server.sh && \
sudo amazon-sagemaker-spark-ui/install-scripts/studio/install-history-server.sh && \
rm -rf ~/.m2 && \
sudo rm -rf amazon-sagemaker-spark-ui*
fi
8 changes: 6 additions & 2 deletions template/v2/dirs/etc/sagemaker-ui/workflows/start-workflows-container.sh
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -3,7 +3,11 @@ set -eu

# Get project directory based on storage type
PROJECT_DIR=${SMUS_PROJECT_DIR:-"$HOME/src"}
MOUNT_DIR=$(readlink -f "$PROJECT_DIR") # get the symlink source if it's symlink
if [ -z "${SMUS_PROJECT_DIR:-}" ]; then
MOUNT_DIR=$PROJECT_DIR
else
MOUNT_DIR=$(readlink -f "$PROJECT_DIR") # get the symlink source
fi

# Datazone project metadata
RESOURCE_METADATA_FILE=/opt/ml/metadata/resource-metadata.json
Expand Down Expand Up @@ -75,7 +79,7 @@ if [ ! -f "${WORKFLOW_HEALTH_PATH}/status.json" ]; then
fi

# Only start local runner if Workflows blueprint is enabled
if [ "$(python /etc/sagemaker-ui/workflows/workflow_client.py check-blueprint --region "$DZ_DOMAIN_REGION" --domain-id "$DZ_DOMAIN_ID" --endpoint "$DZ_ENDPOINT")" = "False" ]; then
if [ "$(python /etc/sagemaker-ui/workflows/workflow_client.py check-blueprint --region "$DZ_DOMAIN_REGION" --domain-id "$DZ_DOMAIN_ID" --endpoint "$DZ_ENDPOINT" --project-id "$DZ_PROJECT_ID")" = "False" ]; then
echo "Workflows blueprint is not enabled. Workflows will not start."
handle_workflows_startup_error 0
fi
Expand Down
41 changes: 33 additions & 8 deletions template/v2/dirs/etc/sagemaker-ui/workflows/workflow_client.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -47,23 +47,45 @@ def stop_local_runner(session: requests.Session, **kwargs):
)
return _validate_response("StopLocalRunner", response)


def check_blueprint(region: str, domain_id: str, endpoint: str, **kwargs):
def check_blueprint(region: str, domain_id: str, endpoint: str, project_id: str, **kwargs):
DZ_CLIENT = boto3.client("datazone")
# add correct endpoint for gamma env
if endpoint != "":
DZ_CLIENT = boto3.client("datazone", endpoint_url=endpoint)
try:
# check if workflows blueprint is enabled in project profile
project_profile_id = DZ_CLIENT.get_project(
domainIdentifier=domain_id, identifier=project_id
)["projectProfileId"]
project_blueprints = DZ_CLIENT.get_project_profile(
domainIdentifier=domain_id, identifier=project_profile_id
)['environmentConfigurations']
proj_blueprint_ids = [proj_env_config["environmentBlueprintId"] for proj_env_config in project_blueprints]
blueprint_id = DZ_CLIENT.list_environment_blueprints(
managed=True, domainIdentifier=domain_id, name="Workflows"
)["items"][0]["id"]
blueprint_config = DZ_CLIENT.get_environment_blueprint_configuration(
domainIdentifier=domain_id, environmentBlueprintIdentifier=blueprint_id
)
enabled_regions = blueprint_config["enabledRegions"]
print(str(region in enabled_regions))

if blueprint_id in proj_blueprint_ids:
blueprint_config = DZ_CLIENT.get_environment_blueprint_configuration(
domainIdentifier=domain_id, environmentBlueprintIdentifier=blueprint_id
)
enabled_regions = blueprint_config["enabledRegions"]
print(str(region in enabled_regions))
else:
print("False")
except:
print("False")
# fallback to checking if only workflows blueprint exists
try:
blueprint_id = DZ_CLIENT.list_environment_blueprints(
managed=True, domainIdentifier=domain_id, name="Workflows"
Copy link
Contributor

@reganbaum reganbaum Sep 12, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

IIRC users can modify blueprint names, and we've had issues in the past relying on the blueprint name. Can we get the environment blueprint by looking at the type or another field instead of the name?

Copy link
Contributor Author

@danobi06 danobi06 Sep 12, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

The workflows blueprint is a managed blueprint provided by sagemaker, customers won't be able to edit the blueprint. Unfortunately list environment blueprint only accepts searching by managed, domainIdentifier and name. ref
In the case users adds their own custom blueprint based off of workflows, then I agree this check won't suffice.

Copy link
Contributor

@reganbaum reganbaum Sep 15, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Got it, thanks for clarifying. But the benefit is at least we don't hang anymore, so while not 100% fullproof, does still fix this edge case.

)["items"][0]["id"]
blueprint_config = DZ_CLIENT.get_environment_blueprint_configuration(
domainIdentifier=domain_id, environmentBlueprintIdentifier=blueprint_id
)
enabled_regions = blueprint_config["enabledRegions"]
print(str(region in enabled_regions))
except:
print("False")


COMMAND_REGISTRY = {
Expand Down Expand Up @@ -94,6 +116,9 @@ def main():
check_blueprint_parser.add_argument(
"--endpoint", type=str, required=True, help="Datazone endpoint for blueprint check"
)
check_blueprint_parser.add_argument(
"--project-id", type=str, required=True, help="Datazone Project ID for blueprint check"
)

args = parser.parse_args()

Expand Down
2 changes: 1 addition & 1 deletion 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 @@ -426,7 +426,7 @@ if [ "${SAGEMAKER_APP_TYPE_LOWERCASE}" = "jupyterlab" ] && [ "$is_express_mode"
bash /etc/sagemaker-ui/libmgmt/install-lib.sh

# Install sm-spark-cli
bash /etc/sagemaker-ui/workflows/sm-spark-cli-install.sh
bash /etc/sagemaker-ui/workflows/sm-spark-cli-install.sh || echo "Warning: sm-spark-cli installation failed, continuing..."
fi

# Execute network validation script, to check if any required AWS Services are unreachable
Expand Down
23 changes: 16 additions & 7 deletions template/v3/dirs/etc/sagemaker-ui/workflows/sm-spark-cli-install.sh
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,9 +1,18 @@
#!/bin/bash
RESOURCE_METADATA_FILE=/opt/ml/metadata/resource-metadata.json
DZ_DOMAIN_ID=$(jq -r '.AdditionalMetadata.DataZoneDomainId' < $RESOURCE_METADATA_FILE)
DZ_PROJECT_ID=$(jq -r '.AdditionalMetadata.DataZoneProjectId' < $RESOURCE_METADATA_FILE)
DZ_DOMAIN_REGION=$(jq -r '.AdditionalMetadata.DataZoneDomainRegion' < $RESOURCE_METADATA_FILE)
DZ_ENDPOINT=$(jq -r '.AdditionalMetadata.DataZoneEndpoint' < $RESOURCE_METADATA_FILE)

# install sm-spark-cli
sudo curl -LO https://github.com/aws-samples/amazon-sagemaker-spark-ui/releases/download/v0.9.1/amazon-sagemaker-spark-ui.tar.gz && \
sudo tar -xvzf amazon-sagemaker-spark-ui.tar.gz && \
sudo chmod +x amazon-sagemaker-spark-ui/install-scripts/studio/install-history-server.sh && \
sudo amazon-sagemaker-spark-ui/install-scripts/studio/install-history-server.sh && \
rm -rf ~/.m2 && \
sudo rm -rf amazon-sagemaker-spark-ui*
# install sm-spark-cli if workflows blueprint is enabled
if [ "$(python /etc/sagemaker-ui/workflows/workflow_client.py check-blueprint --region "$DZ_DOMAIN_REGION" --domain-id "$DZ_DOMAIN_ID" --endpoint "$DZ_ENDPOINT" --project-id "$DZ_PROJECT_ID")" = "True" ]; then
echo "Workflows blueprint is enabled. Installing sm-spark-cli."
# install sm-spark-cli
sudo curl -LO https://github.com/aws-samples/amazon-sagemaker-spark-ui/releases/download/v0.9.1/amazon-sagemaker-spark-ui.tar.gz && \
sudo tar -xvzf amazon-sagemaker-spark-ui.tar.gz && \
sudo chmod +x amazon-sagemaker-spark-ui/install-scripts/studio/install-history-server.sh && \
sudo amazon-sagemaker-spark-ui/install-scripts/studio/install-history-server.sh && \
rm -rf ~/.m2 && \
sudo rm -rf amazon-sagemaker-spark-ui*
fi
8 changes: 6 additions & 2 deletions template/v3/dirs/etc/sagemaker-ui/workflows/start-workflows-container.sh
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -3,7 +3,11 @@ set -eu

# Get project directory based on storage type
PROJECT_DIR=${SMUS_PROJECT_DIR:-"$HOME/src"}
MOUNT_DIR=$(readlink -f "$PROJECT_DIR") # get the symlink source if it's symlink
if [ -z "${SMUS_PROJECT_DIR:-}" ]; then
MOUNT_DIR=$PROJECT_DIR
else
MOUNT_DIR=$(readlink -f "$PROJECT_DIR") # get the symlink source
fi

# Datazone project metadata
RESOURCE_METADATA_FILE=/opt/ml/metadata/resource-metadata.json
Expand Down Expand Up @@ -75,7 +79,7 @@ if [ ! -f "${WORKFLOW_HEALTH_PATH}/status.json" ]; then
fi

# Only start local runner if Workflows blueprint is enabled
if [ "$(python /etc/sagemaker-ui/workflows/workflow_client.py check-blueprint --region "$DZ_DOMAIN_REGION" --domain-id "$DZ_DOMAIN_ID" --endpoint "$DZ_ENDPOINT")" = "False" ]; then
if [ "$(python /etc/sagemaker-ui/workflows/workflow_client.py check-blueprint --region "$DZ_DOMAIN_REGION" --domain-id "$DZ_DOMAIN_ID" --endpoint "$DZ_ENDPOINT" --project-id "$DZ_PROJECT_ID")" = "False" ]; then
echo "Workflows blueprint is not enabled. Workflows will not start."
handle_workflows_startup_error 0
fi
Expand Down
41 changes: 33 additions & 8 deletions template/v3/dirs/etc/sagemaker-ui/workflows/workflow_client.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -47,23 +47,45 @@ def stop_local_runner(session: requests.Session, **kwargs):
)
return _validate_response("StopLocalRunner", response)


def check_blueprint(region: str, domain_id: str, endpoint: str, **kwargs):
def check_blueprint(region: str, domain_id: str, endpoint: str, project_id: str, **kwargs):
DZ_CLIENT = boto3.client("datazone")
# add correct endpoint for gamma env
if endpoint != "":
DZ_CLIENT = boto3.client("datazone", endpoint_url=endpoint)
try:
# check if workflows blueprint is enabled in project profile
project_profile_id = DZ_CLIENT.get_project(
domainIdentifier=domain_id, identifier=project_id
)["projectProfileId"]
project_blueprints = DZ_CLIENT.get_project_profile(
domainIdentifier=domain_id, identifier=project_profile_id
)['environmentConfigurations']
proj_blueprint_ids = [proj_env_config["environmentBlueprintId"] for proj_env_config in project_blueprints]
blueprint_id = DZ_CLIENT.list_environment_blueprints(
managed=True, domainIdentifier=domain_id, name="Workflows"
)["items"][0]["id"]
blueprint_config = DZ_CLIENT.get_environment_blueprint_configuration(
domainIdentifier=domain_id, environmentBlueprintIdentifier=blueprint_id
)
enabled_regions = blueprint_config["enabledRegions"]
print(str(region in enabled_regions))

if blueprint_id in proj_blueprint_ids:
blueprint_config = DZ_CLIENT.get_environment_blueprint_configuration(
domainIdentifier=domain_id, environmentBlueprintIdentifier=blueprint_id
)
enabled_regions = blueprint_config["enabledRegions"]
print(str(region in enabled_regions))
else:
print("False")
except:
print("False")
# fallback to checking if only workflows blueprint exists
try:
blueprint_id = DZ_CLIENT.list_environment_blueprints(
managed=True, domainIdentifier=domain_id, name="Workflows"
)["items"][0]["id"]
blueprint_config = DZ_CLIENT.get_environment_blueprint_configuration(
domainIdentifier=domain_id, environmentBlueprintIdentifier=blueprint_id
)
enabled_regions = blueprint_config["enabledRegions"]
print(str(region in enabled_regions))
except:
print("False")


COMMAND_REGISTRY = {
Expand Down Expand Up @@ -94,6 +116,9 @@ def main():
check_blueprint_parser.add_argument(
"--endpoint", type=str, required=True, help="Datazone endpoint for blueprint check"
)
check_blueprint_parser.add_argument(
"--project-id", type=str, required=True, help="Datazone Project ID for blueprint check"
)

args = parser.parse_args()

Expand Down
Loading