Add missing SMUS code changes from v2.6.0 to v3.1.0

aws-tianquaw · aws-tianquaw · commit 4d11c394d86d · 2025-05-08T21:45:17.000Z
diff --git a/build_artifacts/v3/v3.1/v3.1.0/dirs/etc/sagemaker-ui/sagemaker_ui_post_startup.sh b/build_artifacts/v3/v3.1/v3.1.0/dirs/etc/sagemaker-ui/sagemaker_ui_post_startup.sh
@@ -1,6 +1,30 @@
 #!/bin/bash
 set -eux
 
+# Only run the following commands if SAGEMAKER_APP_TYPE_LOWERCASE is jupyterlab
+if [ "${SAGEMAKER_APP_TYPE_LOWERCASE}" = "jupyterlab" ]; then
+    # Override Jupyter AI config file with specific content only for 2.6
+    # this is a potential workaround for Q chat issue where the chat does not
+    # load since it does not pick up the latest config file
+    NB_USER=sagemaker-user
+    # Check if Jupyter AI config file exists, override with specific content if so only for 2.6
+    # this is a potential workaround for Q chat issue
+    JUPYTER_AI_CONFIG_PATH=/home/${NB_USER}/.local/share/jupyter/jupyter_ai/config.json
+    JUPYTER_AI_CONFIG_CONTENT='{
+        "model_provider_id": "amazon-q:Q-Developer",
+        "embeddings_provider_id": null,
+        "send_with_shift_enter": false,
+        "fields": {},
+        "api_keys": {},
+        "completions_model_provider_id": null,
+        "completions_fields": {},
+        "embeddings_fields": {}
+    }'
+
+    # Overwrite the file if it exists (or create it if it doesn't)
+    echo "$JUPYTER_AI_CONFIG_CONTENT" > "$JUPYTER_AI_CONFIG_PATH"
+fi
+
 # Writes script status to file. This file is read by an IDE extension responsible for dispatching UI post-startup-status to the user.
 write_status_to_file() {
     local status="$1"
@@ -91,14 +115,14 @@ set +x
 
 # Note: The $? check immediately follows the sagemaker-studio command to ensure we're checking its exit status.
 # Adding commands between these lines could lead to incorrect error handling.
-response=$( sagemaker-studio credentials get-domain-execution-role-credential-in-space --domain-id "$dataZoneDomainId" --profile default)
+response=$(timeout 30 sagemaker-studio credentials get-domain-execution-role-credential-in-space --domain-id "$dataZoneDomainId" --profile default)
 responseStatus=$?
 
 set -x
 
 if [ $responseStatus -ne 0 ]; then
         echo "Failed to fetch domain execution role credentials. Will skip adding new credentials profile: DomainExecutionRoleCreds."
-        write_status_to_file "error" "Network issue detected. Your domain may be using a public subnet, which affects IDE functionality. Please contact your administrator."
+        write_status_to_file "error" "Network issue detected. Your domain may be using a public subnet, which affects IDE functionality. Please contact your admin."
 else
         aws configure set credential_process "sagemaker-studio credentials get-domain-execution-role-credential-in-space --domain-id $dataZoneDomainId --profile default" --profile DomainExecutionRoleCreds
         echo "Successfully configured DomainExecutionRoleCreds profile"
@@ -153,23 +177,28 @@ else
   echo readonly LOGNAME >> ~/.bashrc
 fi
 
-set -e
-
-# 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
-
 # Generate sagemaker pysdk intelligent default config
 nohup python /etc/sagemaker/sm_pysdk_default_config.py &
 # Only run the following commands if SAGEMAKER_APP_TYPE_LOWERCASE is jupyterlab
 if [ "${SAGEMAKER_APP_TYPE_LOWERCASE}" = "jupyterlab" ]; then
+    # do not fail immediately for non-zero exit code returned
+    # by start-workflows-container. An expected non-zero exit
+    # code will be returned if there is not a minimum of 2
+    # CPU cores available.
     # Start workflows local runner
     bash /etc/sagemaker-ui/workflows/start-workflows-container.sh
 
+    # ensure functions inherit traps and fail immediately
+    set -eE
+
+    # 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 $HOME/src
 
     # Install sm-spark-cli
     bash /etc/sagemaker-ui/workflows/sm-spark-cli-install.sh
 fi
 
-write_status_to_file_on_script_complete
+write_status_to_file_on_script_complete
diff --git a/build_artifacts/v3/v3.1/v3.1.0/dirs/etc/sagemaker-ui/workflows/start-workflows-container.sh b/build_artifacts/v3/v3.1/v3.1.0/dirs/etc/sagemaker-ui/workflows/start-workflows-container.sh
@@ -42,6 +42,7 @@ DZ_DOMAIN_ID=$(jq -r '.AdditionalMetadata.DataZoneDomainId' < $RESOURCE_METADATA
 DZ_PROJECT_ID=$(jq -r '.AdditionalMetadata.DataZoneProjectId' < $RESOURCE_METADATA_FILE)
 DZ_ENV_ID=$(jq -r '.AdditionalMetadata.DataZoneEnvironmentId' < $RESOURCE_METADATA_FILE)
 DZ_DOMAIN_REGION=$(jq -r '.AdditionalMetadata.DataZoneDomainRegion' < $RESOURCE_METADATA_FILE)
+DZ_ENDPOINT=$(jq -r '.AdditionalMetadata.DataZoneEndpoint' < $RESOURCE_METADATA_FILE)
 DZ_PROJECT_S3PATH=$(jq -r '.AdditionalMetadata.ProjectS3Path' < $RESOURCE_METADATA_FILE)
 WORKFLOW_DAG_PATH="/home/sagemaker-user/${HOME_FOLDER_NAME}/workflows/dags"
 WORKFLOW_CONFIG_PATH="/home/sagemaker-user/${HOME_FOLDER_NAME}/workflows/config"
@@ -67,7 +68,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 --domain-id "$DZ_DOMAIN_ID")" = "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")" = "False" ]; then
     echo "Workflows blueprint is not enabled. Workflows will not start."
     handle_workflows_startup_error 0
 fi
diff --git a/build_artifacts/v3/v3.1/v3.1.0/dirs/etc/sagemaker-ui/workflows/workflow_client.py b/build_artifacts/v3/v3.1/v3.1.0/dirs/etc/sagemaker-ui/workflows/workflow_client.py
@@ -8,7 +8,6 @@
 JUPYTERLAB_URL = "http://default:8888/jupyterlab/default/"
 WORKFLOWS_API_ENDPOINT = "api/sagemaker/workflows"
 TIMESTAMP_FORMAT = "%Y-%m-%d %H:%M:%S.%f%z"
-DZ_CLIENT = boto3.client("datazone")
 
 
 def _validate_response(function_name: str, response: requests.Response):
@@ -49,12 +48,20 @@ def stop_local_runner(session: requests.Session, **kwargs):
     return _validate_response("StopLocalRunner", response)
 
 
-def check_blueprint(domain_id: str, **kwargs):
+def check_blueprint(region: str, domain_id: str, endpoint: str, **kwargs):
+    DZ_CLIENT = boto3.client("datazone")
+    # add correct endpoint for gamma env
+    if endpoint != "":
+        DZ_CLIENT = boto3.client("datazone", endpoint_url=endpoint)
     try:
-        workflow_blueprint = DZ_CLIENT.list_environment_blueprints(domainIdentifier=domain_id, name="Workflows")[
-            "items"
-        ]
-        print(str(bool(workflow_blueprint)))
+        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")
 
@@ -83,6 +90,10 @@ def main():
     check_blueprint_parser.add_argument(
         "--domain-id", type=str, required=True, help="Datazone Domain ID for blueprint check"
     )
+    check_blueprint_parser.add_argument("--region", type=str, required=True, help="Datazone Domain region")
+    check_blueprint_parser.add_argument(
+        "--endpoint", type=str, required=True, help="Datazone endpoint for blueprint check"
+    )
 
     args = parser.parse_args()
 
diff --git a/template/v3/dirs/etc/sagemaker-ui/sagemaker_ui_post_startup.sh b/template/v3/dirs/etc/sagemaker-ui/sagemaker_ui_post_startup.sh
@@ -1,6 +1,30 @@
 #!/bin/bash
 set -eux
 
+# Only run the following commands if SAGEMAKER_APP_TYPE_LOWERCASE is jupyterlab
+if [ "${SAGEMAKER_APP_TYPE_LOWERCASE}" = "jupyterlab" ]; then
+    # Override Jupyter AI config file with specific content only for 2.6
+    # this is a potential workaround for Q chat issue where the chat does not
+    # load since it does not pick up the latest config file
+    NB_USER=sagemaker-user
+    # Check if Jupyter AI config file exists, override with specific content if so only for 2.6
+    # this is a potential workaround for Q chat issue
+    JUPYTER_AI_CONFIG_PATH=/home/${NB_USER}/.local/share/jupyter/jupyter_ai/config.json
+    JUPYTER_AI_CONFIG_CONTENT='{
+        "model_provider_id": "amazon-q:Q-Developer",
+        "embeddings_provider_id": null,
+        "send_with_shift_enter": false,
+        "fields": {},
+        "api_keys": {},
+        "completions_model_provider_id": null,
+        "completions_fields": {},
+        "embeddings_fields": {}
+    }'
+
+    # Overwrite the file if it exists (or create it if it doesn't)
+    echo "$JUPYTER_AI_CONFIG_CONTENT" > "$JUPYTER_AI_CONFIG_PATH"
+fi
+
 # Writes script status to file. This file is read by an IDE extension responsible for dispatching UI post-startup-status to the user.
 write_status_to_file() {
     local status="$1"
@@ -91,14 +115,14 @@ set +x
 
 # Note: The $? check immediately follows the sagemaker-studio command to ensure we're checking its exit status.
 # Adding commands between these lines could lead to incorrect error handling.
-response=$( sagemaker-studio credentials get-domain-execution-role-credential-in-space --domain-id "$dataZoneDomainId" --profile default)
+response=$(timeout 30 sagemaker-studio credentials get-domain-execution-role-credential-in-space --domain-id "$dataZoneDomainId" --profile default)
 responseStatus=$?
 
 set -x
 
 if [ $responseStatus -ne 0 ]; then
         echo "Failed to fetch domain execution role credentials. Will skip adding new credentials profile: DomainExecutionRoleCreds."
-        write_status_to_file "error" "Network issue detected. Your domain may be using a public subnet, which affects IDE functionality. Please contact your administrator."
+        write_status_to_file "error" "Network issue detected. Your domain may be using a public subnet, which affects IDE functionality. Please contact your admin."
 else
         aws configure set credential_process "sagemaker-studio credentials get-domain-execution-role-credential-in-space --domain-id $dataZoneDomainId --profile default" --profile DomainExecutionRoleCreds
         echo "Successfully configured DomainExecutionRoleCreds profile"
@@ -153,23 +177,28 @@ else
   echo readonly LOGNAME >> ~/.bashrc
 fi
 
-set -e
-
-# 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
-
 # Generate sagemaker pysdk intelligent default config
 nohup python /etc/sagemaker/sm_pysdk_default_config.py &
 # Only run the following commands if SAGEMAKER_APP_TYPE_LOWERCASE is jupyterlab
 if [ "${SAGEMAKER_APP_TYPE_LOWERCASE}" = "jupyterlab" ]; then
+    # do not fail immediately for non-zero exit code returned
+    # by start-workflows-container. An expected non-zero exit
+    # code will be returned if there is not a minimum of 2
+    # CPU cores available.
     # Start workflows local runner
     bash /etc/sagemaker-ui/workflows/start-workflows-container.sh
 
+    # ensure functions inherit traps and fail immediately
+    set -eE
+
+    # 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 $HOME/src
 
     # Install sm-spark-cli
     bash /etc/sagemaker-ui/workflows/sm-spark-cli-install.sh
 fi
 
-write_status_to_file_on_script_complete
+write_status_to_file_on_script_complete
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
@@ -42,6 +42,7 @@ DZ_DOMAIN_ID=$(jq -r '.AdditionalMetadata.DataZoneDomainId' < $RESOURCE_METADATA
 DZ_PROJECT_ID=$(jq -r '.AdditionalMetadata.DataZoneProjectId' < $RESOURCE_METADATA_FILE)
 DZ_ENV_ID=$(jq -r '.AdditionalMetadata.DataZoneEnvironmentId' < $RESOURCE_METADATA_FILE)
 DZ_DOMAIN_REGION=$(jq -r '.AdditionalMetadata.DataZoneDomainRegion' < $RESOURCE_METADATA_FILE)
+DZ_ENDPOINT=$(jq -r '.AdditionalMetadata.DataZoneEndpoint' < $RESOURCE_METADATA_FILE)
 DZ_PROJECT_S3PATH=$(jq -r '.AdditionalMetadata.ProjectS3Path' < $RESOURCE_METADATA_FILE)
 WORKFLOW_DAG_PATH="/home/sagemaker-user/${HOME_FOLDER_NAME}/workflows/dags"
 WORKFLOW_CONFIG_PATH="/home/sagemaker-user/${HOME_FOLDER_NAME}/workflows/config"
@@ -67,7 +68,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 --domain-id "$DZ_DOMAIN_ID")" = "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")" = "False" ]; then
     echo "Workflows blueprint is not enabled. Workflows will not start."
     handle_workflows_startup_error 0
 fi
diff --git a/template/v3/dirs/etc/sagemaker-ui/workflows/workflow_client.py b/template/v3/dirs/etc/sagemaker-ui/workflows/workflow_client.py
@@ -8,7 +8,6 @@
 JUPYTERLAB_URL = "http://default:8888/jupyterlab/default/"
 WORKFLOWS_API_ENDPOINT = "api/sagemaker/workflows"
 TIMESTAMP_FORMAT = "%Y-%m-%d %H:%M:%S.%f%z"
-DZ_CLIENT = boto3.client("datazone")
 
 
 def _validate_response(function_name: str, response: requests.Response):
@@ -49,12 +48,20 @@ def stop_local_runner(session: requests.Session, **kwargs):
     return _validate_response("StopLocalRunner", response)
 
 
-def check_blueprint(domain_id: str, **kwargs):
+def check_blueprint(region: str, domain_id: str, endpoint: str, **kwargs):
+    DZ_CLIENT = boto3.client("datazone")
+    # add correct endpoint for gamma env
+    if endpoint != "":
+        DZ_CLIENT = boto3.client("datazone", endpoint_url=endpoint)
     try:
-        workflow_blueprint = DZ_CLIENT.list_environment_blueprints(domainIdentifier=domain_id, name="Workflows")[
-            "items"
-        ]
-        print(str(bool(workflow_blueprint)))
+        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")
 
@@ -83,6 +90,10 @@ def main():
     check_blueprint_parser.add_argument(
         "--domain-id", type=str, required=True, help="Datazone Domain ID for blueprint check"
     )
+    check_blueprint_parser.add_argument("--region", type=str, required=True, help="Datazone Domain region")
+    check_blueprint_parser.add_argument(
+        "--endpoint", type=str, required=True, help="Datazone endpoint for blueprint check"
+    )
 
     args = parser.parse_args()
 
