Skip to content

feat: add decoupled workflows config #1019

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
jwemmons wants to merge 2 commits into
base: main
Choose a base branch
from
Open

feat: add decoupled workflows config #1019

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
3183cc0
feat: add decoupled workflows config
jwemmons Jan 6, 2026
d73f437
Merge branch 'aws:main' into main
jwemmons Jan 13, 2026
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
43 changes: 33 additions & 10 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 @@ -37,6 +37,12 @@ USER_REQUIREMENTS_FILE="${WORKFLOW_CONFIG_PATH}/requirements.txt"
USER_PLUGINS_FOLDER="${WORKFLOW_CONFIG_PATH}/plugins"
USER_STARTUP_FILE="${WORKFLOW_CONFIG_PATH}/startup.sh"

DZ_PROJECT_SYS_S3PATH=$(sed 's|/dev$|/sys|' <<< "$DZ_PROJECT_S3PATH")
WORKFLOWS_ARTIFACTS_SOURCE_S3PATH="${DZ_PROJECT_SYS_S3PATH}/workflows"
WORKFLOW_PLUGINS_SOURCE_S3PATH="${WORKFLOWS_ARTIFACTS_SOURCE_S3PATH}/plugins"
WORKFLOW_REQUIREMENTS_SOURCE_S3PATH="${WORKFLOWS_ARTIFACTS_SOURCE_S3PATH}/requirements.txt"
WORKFLOW_STARTUP_SOURCE_S3PATH="${WORKFLOWS_ARTIFACTS_SOURCE_S3PATH}/startup.sh"


handle_workflows_startup_error() {
local step=$1
Expand Down Expand Up @@ -134,17 +140,30 @@ cat >>"$WORKFLOW_DAG_PATH/.airflowignore" <<'END'
.ipynb_checkpoints
END

#copy plugins from conda
cp $WORKFLOW_PLUGINS_SOURCE_PATH $WORKFLOW_PLUGINS_PATH
#copy requirements from conda
cp $WORKFLOW_REQUIREMENTS_SOURCE_PATH $WORKFLOW_REQUIREMENTS_PATH
V2_CONFIG_ENABLED=$(python /etc/sagemaker-ui/workflows/workflow_client.py check-config-v2-enabled --domain-id "$DZ_DOMAIN_ID" --endpoint "$DZ_ENDPOINT" --project-id "$DZ_PROJECT_ID")

# Copy system startup
cp /etc/sagemaker-ui/workflows/startup/startup.sh $WORKFLOW_STARTUP_PATH
if [ $V2_CONFIG_ENABLED = "True" ]; then
# copy config files from project bucket
aws s3 sync $WORKFLOW_PLUGINS_SOURCE_S3PATH $WORKFLOW_PLUGINS_PATH --delete
aws s3 cp $WORKFLOW_REQUIREMENTS_SOURCE_S3PATH $WORKFLOW_REQUIREMENTS_PATH
aws s3 cp $WORKFLOW_STARTUP_SOURCE_S3PATH $WORKFLOW_STARTUP_PATH
Comment on lines +147 to +149
Copy link
Contributor

@reganbaum reganbaum Jan 8, 2026

Choose a reason for hiding this comment

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

we should probably wrap these calls in a try/catch, and if it fails, then assume the FF isn't enabled as a fallback.

It's also probably useful for debugging purposes later on to include some echo statements throughout this logic so if we ever need to get the logs from the customer, we can see what happened more easily. Could we add in the try/catch and echo statements (e.g., print if the v2config is enabled, when it starts/finishes something, after receiving the blueprint id, etc)

Copy link
Author

@jwemmons jwemmons Jan 12, 2026

Choose a reason for hiding this comment

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

There is overall exception handling for this block here: https://github.com/aws/sagemaker-distribution/blob/main/template/v2/dirs/etc/sagemaker-ui/workflows/start-workflows-container.sh#L173

I don't know that falling back to the pre-FF behavior is the correct approach. Say this new logic has already successfully ran previously and now the s3 calls fail. Now you'll end up copying/using the old config files instead of the new ones and your local and shared environments will differ. I think the current logic of catching the failure and exiting immediately is more appropriate.

Will add additional logging though.

else
#copy plugins from conda
cp $WORKFLOW_PLUGINS_SOURCE_PATH $WORKFLOW_PLUGINS_PATH
#copy requirements from conda
cp $WORKFLOW_REQUIREMENTS_SOURCE_PATH $WORKFLOW_REQUIREMENTS_PATH

# Copy system startup
cp /etc/sagemaker-ui/workflows/startup/startup.sh $WORKFLOW_STARTUP_PATH
fi

# Append user's custom startup script, if exists
# Add user's custom startup script, if exists
if [ -f $USER_STARTUP_FILE ]; then
tail -n +2 $USER_STARTUP_FILE >> "${WORKFLOW_STARTUP_PATH}startup.sh"
if [ $V2_CONFIG_ENABLED = "True" ]; then
cp $USER_STARTUP_FILE "${WORKFLOW_PLUGINS_PATH}/user-startup.sh"
else
tail -n +2 $USER_STARTUP_FILE >> "${WORKFLOW_STARTUP_PATH}startup.sh"
fi
else
# Give the user a template startup script
echo "#!/bin/bash" > "${USER_STARTUP_FILE}"
Expand All @@ -154,9 +173,13 @@ else
echo "# pip install dbt-core" >> "${USER_STARTUP_FILE}"
fi

# Append user's custom requirements, if exists
# Add user's custom requirements, if exists
if [ -f $USER_REQUIREMENTS_FILE ]; then
cat $USER_REQUIREMENTS_FILE >> "${WORKFLOW_REQUIREMENTS_PATH}requirements.txt"
if [ $V2_CONFIG_ENABLED = "True" ]; then
cp $USER_REQUIREMENTS_FILE "${WORKFLOW_PLUGINS_PATH}/user-requirements.txt"
else
cat $USER_REQUIREMENTS_FILE >> "${WORKFLOW_REQUIREMENTS_PATH}requirements.txt"
fi
else
# Give the user a template requirements.txt file
echo "# Place any requirements you'd like included in your workflows environment here" > "${USER_REQUIREMENTS_FILE}"
Expand Down
48 changes: 41 additions & 7 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,24 +47,25 @@ def stop_local_runner(session: requests.Session, **kwargs):
)
return _validate_response("StopLocalRunner", response)


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']
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"]

if blueprint_id in proj_blueprint_ids:
blueprint_config = DZ_CLIENT.get_environment_blueprint_configuration(
domainIdentifier=domain_id, environmentBlueprintIdentifier=blueprint_id
Expand All @@ -88,11 +89,39 @@ def check_blueprint(region: str, domain_id: str, endpoint: str, project_id: str,
print("False")


def check_config_v2_enabled(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:
# get workflows blueprint environment resources
blueprint_id = DZ_CLIENT.list_environment_blueprints(managed=True, domainIdentifier=domain_id, name="Tooling")[
"items"
][0]["id"]
environment_id = DZ_CLIENT.list_environments(
domainIdentifier=domain_id, projectIdentifier=project_id, environmentBlueprintIdentifier=blueprint_id
)["items"][0]["id"]
environment_resources = DZ_CLIENT.get_environment(domainIdentifier=domain_id, identifier=environment_id)[
"provisionedResources"
]

for resource in environment_resources:
if resource["name"] == "enableWorkflowsConfigV2":
print(resource["value"] == "true")
return

print("False")
except Exception:
print("False")


COMMAND_REGISTRY = {
"update-local-runner-status": update_local_runner_status,
"start-local-runner": start_local_runner,
"stop-local-runner": stop_local_runner,
"check-blueprint": check_blueprint,
"check-config-v2-enabled": check_config_v2_enabled,
}


Expand Down Expand Up @@ -120,6 +149,11 @@ def main():
"--project-id", type=str, required=True, help="Datazone Project ID for blueprint check"
)

check_config_v2_enabled_parser = subparsers.add_parser("check-config-v2-enabled", help="Check if config v2 enabled")
check_config_v2_enabled_parser.add_argument("--domain-id", type=str, required=True, help="Datazone Domain ID")
check_config_v2_enabled_parser.add_argument("--endpoint", type=str, required=True, help="Datazone endpoint")
check_config_v2_enabled_parser.add_argument("--project-id", type=str, required=True, help="Datazone Project ID")

args = parser.parse_args()

# create the request session
Expand Down
43 changes: 33 additions & 10 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 @@ -37,6 +37,12 @@ USER_REQUIREMENTS_FILE="${WORKFLOW_CONFIG_PATH}/requirements.txt"
USER_PLUGINS_FOLDER="${WORKFLOW_CONFIG_PATH}/plugins"
USER_STARTUP_FILE="${WORKFLOW_CONFIG_PATH}/startup.sh"

DZ_PROJECT_SYS_S3PATH=$(sed 's|/dev$|/sys|' <<< "$DZ_PROJECT_S3PATH")
WORKFLOWS_ARTIFACTS_SOURCE_S3PATH="${DZ_PROJECT_SYS_S3PATH}/workflows"
WORKFLOW_PLUGINS_SOURCE_S3PATH="${WORKFLOWS_ARTIFACTS_SOURCE_S3PATH}/plugins"
WORKFLOW_REQUIREMENTS_SOURCE_S3PATH="${WORKFLOWS_ARTIFACTS_SOURCE_S3PATH}/requirements.txt"
WORKFLOW_STARTUP_SOURCE_S3PATH="${WORKFLOWS_ARTIFACTS_SOURCE_S3PATH}/startup.sh"


handle_workflows_startup_error() {
local step=$1
Expand Down Expand Up @@ -134,17 +140,30 @@ cat >>"$WORKFLOW_DAG_PATH/.airflowignore" <<'END'
.ipynb_checkpoints
END

#copy plugins from conda
cp $WORKFLOW_PLUGINS_SOURCE_PATH $WORKFLOW_PLUGINS_PATH
#copy requirements from conda
cp $WORKFLOW_REQUIREMENTS_SOURCE_PATH $WORKFLOW_REQUIREMENTS_PATH
V2_CONFIG_ENABLED=$(python /etc/sagemaker-ui/workflows/workflow_client.py check-config-v2-enabled --domain-id "$DZ_DOMAIN_ID" --endpoint "$DZ_ENDPOINT" --project-id "$DZ_PROJECT_ID")

# Copy system startup
cp /etc/sagemaker-ui/workflows/startup/startup.sh $WORKFLOW_STARTUP_PATH
if [ $V2_CONFIG_ENABLED = "True" ]; then
# copy config files from project bucket
aws s3 sync $WORKFLOW_PLUGINS_SOURCE_S3PATH $WORKFLOW_PLUGINS_PATH --delete
aws s3 cp $WORKFLOW_REQUIREMENTS_SOURCE_S3PATH $WORKFLOW_REQUIREMENTS_PATH
aws s3 cp $WORKFLOW_STARTUP_SOURCE_S3PATH $WORKFLOW_STARTUP_PATH
else
#copy plugins from conda
cp $WORKFLOW_PLUGINS_SOURCE_PATH $WORKFLOW_PLUGINS_PATH
#copy requirements from conda
cp $WORKFLOW_REQUIREMENTS_SOURCE_PATH $WORKFLOW_REQUIREMENTS_PATH

# Copy system startup
cp /etc/sagemaker-ui/workflows/startup/startup.sh $WORKFLOW_STARTUP_PATH
fi

# Append user's custom startup script, if exists
# Add user's custom startup script, if exists
if [ -f $USER_STARTUP_FILE ]; then
tail -n +2 $USER_STARTUP_FILE >> "${WORKFLOW_STARTUP_PATH}startup.sh"
if [ $V2_CONFIG_ENABLED = "True" ]; then
cp $USER_STARTUP_FILE "${WORKFLOW_PLUGINS_PATH}/user-startup.sh"
else
tail -n +2 $USER_STARTUP_FILE >> "${WORKFLOW_STARTUP_PATH}startup.sh"
fi
else
# Give the user a template startup script
echo "#!/bin/bash" > "${USER_STARTUP_FILE}"
Expand All @@ -154,9 +173,13 @@ else
echo "# pip install dbt-core" >> "${USER_STARTUP_FILE}"
fi

# Append user's custom requirements, if exists
# Add user's custom requirements, if exists
if [ -f $USER_REQUIREMENTS_FILE ]; then
cat $USER_REQUIREMENTS_FILE >> "${WORKFLOW_REQUIREMENTS_PATH}requirements.txt"
if [ $V2_CONFIG_ENABLED = "True" ]; then
cp $USER_REQUIREMENTS_FILE "${WORKFLOW_PLUGINS_PATH}/user-requirements.txt"
else
cat $USER_REQUIREMENTS_FILE >> "${WORKFLOW_REQUIREMENTS_PATH}requirements.txt"
fi
else
# Give the user a template requirements.txt file
echo "# Place any requirements you'd like included in your workflows environment here" > "${USER_REQUIREMENTS_FILE}"
Expand Down
48 changes: 41 additions & 7 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,24 +47,25 @@ def stop_local_runner(session: requests.Session, **kwargs):
)
return _validate_response("StopLocalRunner", response)


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']
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"]

if blueprint_id in proj_blueprint_ids:
blueprint_config = DZ_CLIENT.get_environment_blueprint_configuration(
domainIdentifier=domain_id, environmentBlueprintIdentifier=blueprint_id
Expand All @@ -88,11 +89,39 @@ def check_blueprint(region: str, domain_id: str, endpoint: str, project_id: str,
print("False")


def check_config_v2_enabled(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:
# get workflows blueprint environment resources
blueprint_id = DZ_CLIENT.list_environment_blueprints(managed=True, domainIdentifier=domain_id, name="Tooling")[
"items"
][0]["id"]
environment_id = DZ_CLIENT.list_environments(
domainIdentifier=domain_id, projectIdentifier=project_id, environmentBlueprintIdentifier=blueprint_id
)["items"][0]["id"]
environment_resources = DZ_CLIENT.get_environment(domainIdentifier=domain_id, identifier=environment_id)[
"provisionedResources"
]

for resource in environment_resources:
if resource["name"] == "enableWorkflowsConfigV2":
print(resource["value"] == "true")
return

print("False")
except Exception:
print("False")


COMMAND_REGISTRY = {
"update-local-runner-status": update_local_runner_status,
"start-local-runner": start_local_runner,
"stop-local-runner": stop_local_runner,
"check-blueprint": check_blueprint,
"check-config-v2-enabled": check_config_v2_enabled,
}


Expand Down Expand Up @@ -120,6 +149,11 @@ def main():
"--project-id", type=str, required=True, help="Datazone Project ID for blueprint check"
)

check_config_v2_enabled_parser = subparsers.add_parser("check-config-v2-enabled", help="Check if config v2 enabled")
check_config_v2_enabled_parser.add_argument("--domain-id", type=str, required=True, help="Datazone Domain ID")
check_config_v2_enabled_parser.add_argument("--endpoint", type=str, required=True, help="Datazone endpoint")
check_config_v2_enabled_parser.add_argument("--project-id", type=str, required=True, help="Datazone Project ID")

args = parser.parse_args()

# create the request session
Expand Down
Loading