Add changes from #786, #784, #789, #713

Author: claytonparnell

build_artifacts/v2/v2.8/v2.8.3/dirs/etc/sagemaker-ui/jupyter/lab/settings/page_config.json

@@ -8,7 +8,6 @@
       "@amzn/sagemaker-jupyterlab-extensions:spacemenu": true,
       "@amzn/amazon_sagemaker_sql_editor": true,
       "@sagemaker-studio:EmrCluster": true,
-      "@jupyterlab/scheduler": true,
       "@jupyter/collaboration-extension": true
     }
 }

build_artifacts/v2/v2.8/v2.8.3/dirs/etc/sagemaker-ui/jupyter/server/jupyter_server_config.py

@@ -1,5 +1,11 @@
 import os
 
+from sagemaker_studio_jupyter_scheduler.scheduler import (
+    SagemakerEnvironmentManager,
+    SageMakerJobFilesManager,
+    SageMakerUnifiedStudioScheduler,
+)
+
 c.ServerApp.terminado_settings = {"shell_command": ["/bin/bash"]}
 region = os.environ.get("AWS_REGION")
 csp_rule = os.environ.get("JUPYTERSERVER_CSP_RULE")
@@ -27,4 +33,9 @@
     module_location = os.path.dirname(module.__file__)
     c.LanguageServerManager.extra_node_roots = [f"{module_location}/sql-language-server"]
 except:
+
     pass
+
+c.SchedulerApp.scheduler_class = SageMakerUnifiedStudioScheduler
+c.SchedulerApp.environment_manager_class = SagemakerEnvironmentManager
+c.SchedulerApp.job_files_manager_class = SageMakerJobFilesManager

build_artifacts/v2/v2.8/v2.8.3/dirs/etc/sagemaker-ui/network_validation.sh

@@ -126,14 +126,18 @@ temp_dir=$(mktemp -d)
 # Launch all service API checks in parallel background jobs
 for service in "${!SERVICE_COMMANDS[@]}"; do
   {
-    # Run command with timeout, discard stdout/stderr
-    if timeout "${api_time_out_limit}s" bash -c "${SERVICE_COMMANDS[$service]}" > /dev/null 2>&1; then
+    output_file="$temp_dir/${service}_output"
+
+    # Run command with timeout
+    if timeout "${api_time_out_limit}s" bash -c "${SERVICE_COMMANDS[$service]}" > "$output_file" 2>&1; then
       # Success: write OK to temp file
       echo "OK" > "$temp_dir/$service"
     else
       # Get exit code to differentiate timeout or other errors
       exit_code=$?
-      if [ "$exit_code" -eq 124 ]; then
+      if grep -q "Could not connect to the endpoint URL" "$output_file"; then
+        echo "UNREACHABLE" > "$temp_dir/$service"
+      elif [ "$exit_code" -eq 124 ]; then
         # Timeout exit code
         echo "TIMEOUT" > "$temp_dir/$service"
       else
@@ -155,10 +159,13 @@ for service in "${!SERVICE_COMMANDS[@]}"; do
     if [[ "$result" == "TIMEOUT" ]]; then
       echo "$service API did NOT resolve within ${api_time_out_limit}s. Marking as unreachable."
       unreachable_services+=("$service")
+    elif [[ "$result" == "UNREACHABLE" ]]; then
+      echo "$service API failed to connect to the endpoint. Marking as unreachable."
+      unreachable_services+=("$service")
     elif [[ "$result" == "OK" ]]; then
       echo "$service API is reachable."
     else
-      echo "$service API returned an error (but not a timeout). Ignored for network check."
+      echo "$service API returned an error (but not a timeout or endpoint reachability failure). Ignored for network check."
     fi
   else
     echo "$service check did not produce a result file. Skipping."

build_artifacts/v2/v2.8/v2.8.3/dirs/etc/sagemaker-ui/sagemaker_ui_post_startup.sh

@@ -87,22 +87,43 @@ c.Application.logging_config = {
 EOT
 fi
 
-# Setting this to +x to not log credentials from the response of fetching credentials.
-set +x
+# Add debug call to get domain info
+if [ ! -z "$dataZoneEndPoint" ]; then
+    domain_response=$(aws datazone get-domain --debug --endpoint-url "$dataZoneEndPoint" --identifier "$dataZoneDomainId" --region "$dataZoneDomainRegion" 2>&1)
 
-# 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=$(timeout 30 sagemaker-studio credentials get-domain-execution-role-credential-in-space --domain-id "$dataZoneDomainId" --profile default)
-responseStatus=$?
+else
+    domain_response=$(aws datazone get-domain --debug --identifier "$dataZoneDomainId" --region "$dataZoneDomainRegion" 2>&1)
+fi
+
+# Check if domain is in express mode
+response_body=$(echo "$domain_response" | grep -A1 "Response body:" | tail -n1 | sed 's/^b'\''//;s/'\''$//')
+# Remove leading/trailing whitespace and the 'b' prefix
+cleaned_response=$(echo "$response_body" | sed 's/\\n//g')
+is_express_mode=$(echo "$cleaned_response" | jq -r '.preferences.DOMAIN_MODE == "EXPRESS"')
+
+if [ "$is_express_mode" = "true" ]; then
+    echo "Domain is in express mode. Using default credentials"
+    # Use default credentials - no additional configuration needed
+    aws configure set credential_source EcsContainer --profile DomainExecutionRoleCreds
+    echo "Successfully configured DomainExecutionRoleCreds profile with default credentials"
+else
+    echo "Domain is not in express mode"
+    # Setting this to +x to not log credentials from the response of fetching credentials.
+    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=$(timeout 30 sagemaker-studio credentials get-domain-execution-role-credential-in-space --domain-id "$dataZoneDomainId" --profile default)
+    responseStatus=$?
 
-set -x
+    set -x
 
-if [ $responseStatus -ne 0 ]; then
+    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 admin."
-else
+    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"
+    fi
 fi
 
 # Run AWS CLI command to get the username from DataZone User Profile.
@@ -179,7 +200,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"
@@ -210,8 +231,8 @@ fi
 
 # 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
+# Only run the following commands if SAGEMAKER_APP_TYPE_LOWERCASE is jupyterlab and domain is not in express mode
+if [ "${SAGEMAKER_APP_TYPE_LOWERCASE}" = "jupyterlab" ] && [ "$is_express_mode" != "true" ]; 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
@@ -224,7 +245,7 @@ if [ "${SAGEMAKER_APP_TYPE_LOWERCASE}" = "jupyterlab" ]; then
 
     # 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