diff --git a/template/v2/dirs/etc/sagemaker-ui/workflows/start-workflows-container.sh b/template/v2/dirs/etc/sagemaker-ui/workflows/start-workflows-container.sh index 6fe0ca696..f2850c6e4 100755 --- a/template/v2/dirs/etc/sagemaker-ui/workflows/start-workflows-container.sh +++ b/template/v2/dirs/etc/sagemaker-ui/workflows/start-workflows-container.sh @@ -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 @@ -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}" @@ -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}" diff --git a/template/v2/dirs/etc/sagemaker-ui/workflows/workflow_client.py b/template/v2/dirs/etc/sagemaker-ui/workflows/workflow_client.py index 36b2b509f..9558006e3 100644 --- a/template/v2/dirs/etc/sagemaker-ui/workflows/workflow_client.py +++ b/template/v2/dirs/etc/sagemaker-ui/workflows/workflow_client.py @@ -47,6 +47,7 @@ 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 @@ -54,17 +55,17 @@ def check_blueprint(region: str, domain_id: str, endpoint: str, project_id: str, 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 @@ -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, } @@ -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 diff --git a/template/v3/dirs/etc/sagemaker-ui/workflows/start-workflows-container.sh b/template/v3/dirs/etc/sagemaker-ui/workflows/start-workflows-container.sh index 6fe0ca696..f2850c6e4 100755 --- a/template/v3/dirs/etc/sagemaker-ui/workflows/start-workflows-container.sh +++ b/template/v3/dirs/etc/sagemaker-ui/workflows/start-workflows-container.sh @@ -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 @@ -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}" @@ -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}" diff --git a/template/v3/dirs/etc/sagemaker-ui/workflows/workflow_client.py b/template/v3/dirs/etc/sagemaker-ui/workflows/workflow_client.py index 36b2b509f..9558006e3 100644 --- a/template/v3/dirs/etc/sagemaker-ui/workflows/workflow_client.py +++ b/template/v3/dirs/etc/sagemaker-ui/workflows/workflow_client.py @@ -47,6 +47,7 @@ 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 @@ -54,17 +55,17 @@ def check_blueprint(region: str, domain_id: str, endpoint: str, project_id: str, 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 @@ -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, } @@ -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