From 463ce2e551968e8082ba0711b7d2f2672bc11f03 Mon Sep 17 00:00:00 2001 From: Eason Ma Date: Thu, 24 Jul 2025 16:17:42 -0700 Subject: [PATCH 1/7] feat: add Amazon Q agent configuration and MCP migration - Add default.json configuration files for Amazon Q agents in v2 and v3 templates - Add migration logic in post-startup scripts to merge MCP configurations - Support intelligent merging of existing agent configurations --- .../sagemaker-ui/sagemaker-mcp/default.json | 48 ++++++++++++++++ .../sagemaker-ui/sagemaker_ui_post_startup.sh | 55 +++++++++++++++++++ .../sagemaker-ui/sagemaker-mcp/default.json | 48 ++++++++++++++++ .../sagemaker-ui/sagemaker_ui_post_startup.sh | 55 +++++++++++++++++++ 4 files changed, 206 insertions(+) create mode 100644 template/v2/dirs/etc/sagemaker-ui/sagemaker-mcp/default.json mode change 100755 => 100644 template/v2/dirs/etc/sagemaker-ui/sagemaker_ui_post_startup.sh create mode 100644 template/v3/dirs/etc/sagemaker-ui/sagemaker-mcp/default.json diff --git a/template/v2/dirs/etc/sagemaker-ui/sagemaker-mcp/default.json b/template/v2/dirs/etc/sagemaker-ui/sagemaker-mcp/default.json new file mode 100644 index 000000000..c5d5cb4e9 --- /dev/null +++ b/template/v2/dirs/etc/sagemaker-ui/sagemaker-mcp/default.json @@ -0,0 +1,48 @@ +{ + "name": "default-agent", + "version": "1.0.0", + "description": "Default agent configuration", + "mcpServers": { + "smus-local-mcp": { + "command": "python", + "args": ["/etc/sagemaker-ui/sagemaker-mcp/smus-mcp.py"] + } + }, + "tools": [ + "fsRead", + "fsWrite", + "fsReplace", + "listDirectory", + "fileSearch", + "executeBash" + ], + "allowedTools": [ + "fsRead", + "fsWrite", + "fsReplace", + "listDirectory", + "fileSearch" + ], + "toolsSettings": { + "execute_bash": { + "alwaysAllow": [ + { + "preset": "readOnly" + } + ] + }, + "use_aws": { + "alwaysAllow": [ + { + "preset": "readOnly" + } + ] + } + }, + "includedFiles": [ + "AmazonQ.md", + "README.md", + ".amazonq/rules/**/*.md" + ], + "resources": [] + } \ No newline at end of file diff --git a/template/v2/dirs/etc/sagemaker-ui/sagemaker_ui_post_startup.sh b/template/v2/dirs/etc/sagemaker-ui/sagemaker_ui_post_startup.sh old mode 100755 new mode 100644 index 22100ef4c..410c094bb --- a/template/v2/dirs/etc/sagemaker-ui/sagemaker_ui_post_startup.sh +++ b/template/v2/dirs/etc/sagemaker-ui/sagemaker_ui_post_startup.sh @@ -303,6 +303,61 @@ else echo "Warning: MCP configuration file not found at $source_file" fi +# Migrating MCP configuration to new config file +echo "Migrating MCP configuration to Amazon Q agents..." +agents_target_file="$HOME/.aws/amazonq/agents/default.json" +agents_source_file="/etc/sagemaker-ui/sagemaker-mcp/default.json" + +if [ -f "$agents_source_file" ]; then + mkdir -p "$HOME/.aws/amazonq/agents/" + + if [ -f "$agents_target_file" ]; then + echo "Existing Amazon Q agents configuration found, merging mcpServers..." + + # Check if target file is valid JSON + if jq empty "$agents_target_file" 2>/dev/null; then + # Initialize mcpServers object if it doesn't exist in target + if ! jq -e '.mcpServers' "$agents_target_file" >/dev/null 2>&1; then + echo "Creating mcpServers object in existing agents configuration" + jq '. + {"mcpServers":{}}' "$agents_target_file" > "$agents_target_file.tmp" + mv "$agents_target_file.tmp" "$agents_target_file" + fi + + # Add servers from source that don't exist in target + if jq -e '.mcpServers' "$agents_source_file" >/dev/null 2>&1; then + source_server_names=$(jq -r '.mcpServers | keys[]' "$agents_source_file") + + for server_name in $source_server_names; do + if ! jq -e ".mcpServers.\"$server_name\"" "$agents_target_file" >/dev/null 2>&1; then + # Server doesn't exist in target - add it + server_config=$(jq ".mcpServers.\"$server_name\"" "$agents_source_file") + jq --arg name "$server_name" --argjson config "$server_config" \ + '.mcpServers[$name] = $config' "$agents_target_file" > "$agents_target_file.tmp" + mv "$agents_target_file.tmp" "$agents_target_file" + echo "Added server '$server_name' to agents configuration" + else + echo "Server '$server_name' already exists in configuration, skipping" + fi + done + + echo "Successfully added missing mcpServers from default.json to agents configuration" + else + echo "No mcpServers found in source configuration" + fi + else + echo "Warning: Existing agents configuration is not valid JSON, replacing with default configuration" + cp "$agents_source_file" "$agents_target_file" + fi + else + cp "$agents_source_file" "$agents_target_file" + echo "Created new Amazon Q agents configuration file" + fi + + echo "Successfully migrated MCP configuration to Amazon Q agents" +else + echo "Warning: Source configuration file not found at $agents_source_file" +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 and domain is not in express mode diff --git a/template/v3/dirs/etc/sagemaker-ui/sagemaker-mcp/default.json b/template/v3/dirs/etc/sagemaker-ui/sagemaker-mcp/default.json new file mode 100644 index 000000000..c5d5cb4e9 --- /dev/null +++ b/template/v3/dirs/etc/sagemaker-ui/sagemaker-mcp/default.json @@ -0,0 +1,48 @@ +{ + "name": "default-agent", + "version": "1.0.0", + "description": "Default agent configuration", + "mcpServers": { + "smus-local-mcp": { + "command": "python", + "args": ["/etc/sagemaker-ui/sagemaker-mcp/smus-mcp.py"] + } + }, + "tools": [ + "fsRead", + "fsWrite", + "fsReplace", + "listDirectory", + "fileSearch", + "executeBash" + ], + "allowedTools": [ + "fsRead", + "fsWrite", + "fsReplace", + "listDirectory", + "fileSearch" + ], + "toolsSettings": { + "execute_bash": { + "alwaysAllow": [ + { + "preset": "readOnly" + } + ] + }, + "use_aws": { + "alwaysAllow": [ + { + "preset": "readOnly" + } + ] + } + }, + "includedFiles": [ + "AmazonQ.md", + "README.md", + ".amazonq/rules/**/*.md" + ], + "resources": [] + } \ No newline at end of file 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 index 2913f7d91..a0aa5f694 100755 --- a/template/v3/dirs/etc/sagemaker-ui/sagemaker_ui_post_startup.sh +++ b/template/v3/dirs/etc/sagemaker-ui/sagemaker_ui_post_startup.sh @@ -304,6 +304,61 @@ else echo "Warning: MCP configuration file not found at $source_file" fi +# Migrating MCP configuration to new config file +echo "Migrating MCP configuration to Amazon Q agents..." +agents_target_file="$HOME/.aws/amazonq/agents/default.json" +agents_source_file="/etc/sagemaker-ui/sagemaker-mcp/default.json" + +if [ -f "$agents_source_file" ]; then + mkdir -p "$HOME/.aws/amazonq/agents/" + + if [ -f "$agents_target_file" ]; then + echo "Existing Amazon Q agents configuration found, merging mcpServers..." + + # Check if target file is valid JSON + if jq empty "$agents_target_file" 2>/dev/null; then + # Initialize mcpServers object if it doesn't exist in target + if ! jq -e '.mcpServers' "$agents_target_file" >/dev/null 2>&1; then + echo "Creating mcpServers object in existing agents configuration" + jq '. + {"mcpServers":{}}' "$agents_target_file" > "$agents_target_file.tmp" + mv "$agents_target_file.tmp" "$agents_target_file" + fi + + # Add servers from source that don't exist in target + if jq -e '.mcpServers' "$agents_source_file" >/dev/null 2>&1; then + source_server_names=$(jq -r '.mcpServers | keys[]' "$agents_source_file") + + for server_name in $source_server_names; do + if ! jq -e ".mcpServers.\"$server_name\"" "$agents_target_file" >/dev/null 2>&1; then + # Server doesn't exist in target - add it + server_config=$(jq ".mcpServers.\"$server_name\"" "$agents_source_file") + jq --arg name "$server_name" --argjson config "$server_config" \ + '.mcpServers[$name] = $config' "$agents_target_file" > "$agents_target_file.tmp" + mv "$agents_target_file.tmp" "$agents_target_file" + echo "Added server '$server_name' to agents configuration" + else + echo "Server '$server_name' already exists in configuration, skipping" + fi + done + + echo "Successfully added missing mcpServers from default.json to agents configuration" + else + echo "No mcpServers found in source configuration" + fi + else + echo "Warning: Existing agents configuration is not valid JSON, replacing with default configuration" + cp "$agents_source_file" "$agents_target_file" + fi + else + cp "$agents_source_file" "$agents_target_file" + echo "Created new Amazon Q agents configuration file" + fi + + echo "Successfully migrated MCP configuration to Amazon Q agents" +else + echo "Warning: Source configuration file not found at $agents_source_file" +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 and domain is not in express mode From db84309d5a5cf44f37c81fa41ca2def29b9976aa Mon Sep 17 00:00:00 2001 From: Eason Ma Date: Thu, 24 Jul 2025 17:11:53 -0700 Subject: [PATCH 2/7] fix: add smus local mcp into tools list --- .../sagemaker-ui/sagemaker-mcp/default.json | 22 ++--- .../sagemaker-ui/sagemaker_ui_post_startup.sh | 21 ++++- .../sagemaker-ui/sagemaker-mcp/default.json | 90 +++++++++---------- .../sagemaker-ui/sagemaker_ui_post_startup.sh | 21 ++++- 4 files changed, 94 insertions(+), 60 deletions(-) diff --git a/template/v2/dirs/etc/sagemaker-ui/sagemaker-mcp/default.json b/template/v2/dirs/etc/sagemaker-ui/sagemaker-mcp/default.json index c5d5cb4e9..07516d45f 100644 --- a/template/v2/dirs/etc/sagemaker-ui/sagemaker-mcp/default.json +++ b/template/v2/dirs/etc/sagemaker-ui/sagemaker-mcp/default.json @@ -3,18 +3,22 @@ "version": "1.0.0", "description": "Default agent configuration", "mcpServers": { - "smus-local-mcp": { - "command": "python", - "args": ["/etc/sagemaker-ui/sagemaker-mcp/smus-mcp.py"] - } - }, + "smus-local-mcp": { + "command": "python", + "args": [ + "/etc/sagemaker-ui/sagemaker-mcp/smus-mcp.py" + ], + "env": {} + } + }, "tools": [ "fsRead", "fsWrite", "fsReplace", "listDirectory", "fileSearch", - "executeBash" + "executeBash", + "@smus-local-mcp" ], "allowedTools": [ "fsRead", @@ -39,10 +43,6 @@ ] } }, - "includedFiles": [ - "AmazonQ.md", - "README.md", - ".amazonq/rules/**/*.md" - ], + "includedFiles": [], "resources": [] } \ No newline at end of file diff --git a/template/v2/dirs/etc/sagemaker-ui/sagemaker_ui_post_startup.sh b/template/v2/dirs/etc/sagemaker-ui/sagemaker_ui_post_startup.sh index 410c094bb..2c2c146e5 100644 --- a/template/v2/dirs/etc/sagemaker-ui/sagemaker_ui_post_startup.sh +++ b/template/v2/dirs/etc/sagemaker-ui/sagemaker_ui_post_startup.sh @@ -323,7 +323,7 @@ if [ -f "$agents_source_file" ]; then mv "$agents_target_file.tmp" "$agents_target_file" fi - # Add servers from source that don't exist in target + # Add servers from source that don't exist in target and update tools if jq -e '.mcpServers' "$agents_source_file" >/dev/null 2>&1; then source_server_names=$(jq -r '.mcpServers | keys[]' "$agents_source_file") @@ -335,12 +335,29 @@ if [ -f "$agents_source_file" ]; then '.mcpServers[$name] = $config' "$agents_target_file" > "$agents_target_file.tmp" mv "$agents_target_file.tmp" "$agents_target_file" echo "Added server '$server_name' to agents configuration" + + # Check if source has tools that reference this server and add them + server_tool_ref="@$server_name" + if jq -e --arg tool "$server_tool_ref" '.tools | index($tool)' "$agents_source_file" >/dev/null 2>&1; then + # Initialize tools array if it doesn't exist + if ! jq -e '.tools' "$agents_target_file" >/dev/null 2>&1; then + jq '. + {"tools":[]}' "$agents_target_file" > "$agents_target_file.tmp" + mv "$agents_target_file.tmp" "$agents_target_file" + fi + + # Add tool reference if it doesn't exist + if ! jq -e --arg tool "$server_tool_ref" '.tools | index($tool)' "$agents_target_file" >/dev/null 2>&1; then + jq --arg tool "$server_tool_ref" '.tools += [$tool]' "$agents_target_file" > "$agents_target_file.tmp" + mv "$agents_target_file.tmp" "$agents_target_file" + echo "Added tool reference '$server_tool_ref' to agents configuration" + fi + fi else echo "Server '$server_name' already exists in configuration, skipping" fi done - echo "Successfully added missing mcpServers from default.json to agents configuration" + echo "Successfully added missing mcpServers and tools from default.json to agents configuration" else echo "No mcpServers found in source configuration" fi diff --git a/template/v3/dirs/etc/sagemaker-ui/sagemaker-mcp/default.json b/template/v3/dirs/etc/sagemaker-ui/sagemaker-mcp/default.json index c5d5cb4e9..0b0bebfee 100644 --- a/template/v3/dirs/etc/sagemaker-ui/sagemaker-mcp/default.json +++ b/template/v3/dirs/etc/sagemaker-ui/sagemaker-mcp/default.json @@ -1,48 +1,48 @@ { - "name": "default-agent", - "version": "1.0.0", - "description": "Default agent configuration", - "mcpServers": { - "smus-local-mcp": { - "command": "python", - "args": ["/etc/sagemaker-ui/sagemaker-mcp/smus-mcp.py"] + "name": "default-agent", + "version": "1.0.0", + "description": "Default agent configuration", + "mcpServers": { + "smus-local-mcp": { + "command": "python", + "args": [ + "/etc/sagemaker-ui/sagemaker-mcp/smus-mcp.py" + ], + "env": {} + } + }, + "tools": [ + "fsRead", + "fsWrite", + "fsReplace", + "listDirectory", + "fileSearch", + "executeBash", + "@smus-local-mcp" + ], + "allowedTools": [ + "fsRead", + "fsWrite", + "fsReplace", + "listDirectory", + "fileSearch" + ], + "toolsSettings": { + "execute_bash": { + "alwaysAllow": [ + { + "preset": "readOnly" } - }, - "tools": [ - "fsRead", - "fsWrite", - "fsReplace", - "listDirectory", - "fileSearch", - "executeBash" - ], - "allowedTools": [ - "fsRead", - "fsWrite", - "fsReplace", - "listDirectory", - "fileSearch" - ], - "toolsSettings": { - "execute_bash": { - "alwaysAllow": [ - { - "preset": "readOnly" - } - ] - }, - "use_aws": { - "alwaysAllow": [ - { - "preset": "readOnly" - } - ] - } + ] }, - "includedFiles": [ - "AmazonQ.md", - "README.md", - ".amazonq/rules/**/*.md" - ], - "resources": [] - } \ No newline at end of file + "use_aws": { + "alwaysAllow": [ + { + "preset": "readOnly" + } + ] + } + }, + "includedFiles": [], + "resources": [] +} \ No newline at end of file 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 index a0aa5f694..b310fd360 100755 --- a/template/v3/dirs/etc/sagemaker-ui/sagemaker_ui_post_startup.sh +++ b/template/v3/dirs/etc/sagemaker-ui/sagemaker_ui_post_startup.sh @@ -324,7 +324,7 @@ if [ -f "$agents_source_file" ]; then mv "$agents_target_file.tmp" "$agents_target_file" fi - # Add servers from source that don't exist in target + # Add servers from source that don't exist in target and update tools if jq -e '.mcpServers' "$agents_source_file" >/dev/null 2>&1; then source_server_names=$(jq -r '.mcpServers | keys[]' "$agents_source_file") @@ -336,12 +336,29 @@ if [ -f "$agents_source_file" ]; then '.mcpServers[$name] = $config' "$agents_target_file" > "$agents_target_file.tmp" mv "$agents_target_file.tmp" "$agents_target_file" echo "Added server '$server_name' to agents configuration" + + # Check if source has tools that reference this server and add them + server_tool_ref="@$server_name" + if jq -e --arg tool "$server_tool_ref" '.tools | index($tool)' "$agents_source_file" >/dev/null 2>&1; then + # Initialize tools array if it doesn't exist + if ! jq -e '.tools' "$agents_target_file" >/dev/null 2>&1; then + jq '. + {"tools":[]}' "$agents_target_file" > "$agents_target_file.tmp" + mv "$agents_target_file.tmp" "$agents_target_file" + fi + + # Add tool reference if it doesn't exist + if ! jq -e --arg tool "$server_tool_ref" '.tools | index($tool)' "$agents_target_file" >/dev/null 2>&1; then + jq --arg tool "$server_tool_ref" '.tools += [$tool]' "$agents_target_file" > "$agents_target_file.tmp" + mv "$agents_target_file.tmp" "$agents_target_file" + echo "Added tool reference '$server_tool_ref' to agents configuration" + fi + fi else echo "Server '$server_name' already exists in configuration, skipping" fi done - echo "Successfully added missing mcpServers from default.json to agents configuration" + echo "Successfully added missing mcpServers and tools from default.json to agents configuration" else echo "No mcpServers found in source configuration" fi From 9f15e2b24733f9114a7ee9ee7fd89a77b43c8f60 Mon Sep 17 00:00:00 2001 From: Eason Ma Date: Fri, 8 Aug 2025 14:03:47 -0700 Subject: [PATCH 3/7] Updated QCLI download version to . It was pinned to avoid known issue. --- template/v2/Dockerfile | 2 +- template/v3/Dockerfile | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/template/v2/Dockerfile b/template/v2/Dockerfile index 929c89e3d..2afa47548 100644 --- a/template/v2/Dockerfile +++ b/template/v2/Dockerfile @@ -63,7 +63,7 @@ RUN apt-get update && apt-get upgrade -y && \ rm -rf aws awscliv2.zip && \ : && \ # Install Q CLI - curl --proto '=https' --tlsv1.2 -sSf "https://desktop-release.q.us-east-1.amazonaws.com/1.12.7/q-x86_64-linux.zip" -o "q.zip" && \ + curl --proto '=https' --tlsv1.2 -sSf "https://desktop-release.q.us-east-1.amazonaws.com/latest/q-x86_64-linux.zip" -o "q.zip" && \ unzip q.zip && \ Q_INSTALL_GLOBAL=true ./q/install.sh --no-confirm && \ rm -rf q q.zip && \ diff --git a/template/v3/Dockerfile b/template/v3/Dockerfile index 644caf8aa..65d4da7ea 100644 --- a/template/v3/Dockerfile +++ b/template/v3/Dockerfile @@ -63,7 +63,7 @@ RUN apt-get update && apt-get upgrade -y && \ rm -rf aws awscliv2.zip && \ : && \ # Install Q CLI - curl --proto '=https' --tlsv1.2 -sSf "https://desktop-release.q.us-east-1.amazonaws.com/1.12.7/q-x86_64-linux.zip" -o "q.zip" && \ + curl --proto '=https' --tlsv1.2 -sSf "https://desktop-release.q.us-east-1.amazonaws.com/latest/q-x86_64-linux.zip" -o "q.zip" && \ unzip q.zip && \ Q_INSTALL_GLOBAL=true ./q/install.sh --no-confirm && \ rm -rf q q.zip && \ From 04eacca64202331855f02be5b34c535e056cd900 Mon Sep 17 00:00:00 2001 From: Eason Ma Date: Mon, 11 Aug 2025 15:56:18 -0700 Subject: [PATCH 4/7] Emitting metrics for SMUS local mcp via Q chat extension log file entry --- .../sagemaker-ui/sagemaker-mcp/smus-mcp.py | 59 +++++++++++++++++++ .../sagemaker-ui/sagemaker-mcp/smus-mcp.py | 59 +++++++++++++++++++ 2 files changed, 118 insertions(+) diff --git a/template/v2/dirs/etc/sagemaker-ui/sagemaker-mcp/smus-mcp.py b/template/v2/dirs/etc/sagemaker-ui/sagemaker-mcp/smus-mcp.py index 1fb237e3f..ad1a98db5 100644 --- a/template/v2/dirs/etc/sagemaker-ui/sagemaker-mcp/smus-mcp.py +++ b/template/v2/dirs/etc/sagemaker-ui/sagemaker-mcp/smus-mcp.py @@ -7,6 +7,7 @@ import logging import os import re +import time from typing import Any, Dict from mcp.server.fastmcp import FastMCP @@ -83,6 +84,54 @@ def create_smus_context_identifiers_response(domain_id: str, project_id: str, re Again, include only required parameters. Any extra parameters may cause the API to fail. Stick strictly to the schema.""" +def write_log_entry(operation: str, tool_name: str, error_count: int = 0, domain_id: str = "", latency: float = 0) -> None: + """ + Write a log entry to the SageMaker GenAI Jupyter Lab extension log file. + + Args: + operation: The operation being performed + tool_name: The name of the tool being used + error_count: Number of errors (0 for success, 1 for failure) + domain_id: The domain identifier (optional) + latency: Execution latency in milliseconds + """ + try: + log_entry = { + "Operation": operation, + "ToolName": tool_name, + "ErrorCount": error_count, + "DomainId": domain_id, + "Latency": latency, + "_aws": { + "Timestamp": int(time.time() * 1000), + "CloudWatchMetrics": [{ + "Dimensions": [["Operation"], ["ToolName"]], + "Metrics": [ + { + "Name": "ErrorCount", + "Unit": "Count" + }, + { + "Name": "Latency", + "Unit": "Milliseconds" + } + ], + "Namespace": "SagemakerGenAIJupyterLab" + }] + }, + } + + # Write to log file + log_file_path = "/var/log/studio/sagemaker_genai_jl_ext/sagemaker_genai_jl_extension.log" + os.makedirs(os.path.dirname(log_file_path), exist_ok=True) + with open(log_file_path, "a") as log_file: + log_file.write(json.dumps(log_entry) + "\n") + + except Exception: + # Fail quietly - logging failures shouldn't affect main functionality + pass + + async def aws_context_provider() -> Dict[str, Any]: """ AWS Context Provider - MUST BE CALLED BEFORE ANY use_aws OPERATIONS @@ -101,16 +150,26 @@ async def aws_context_provider() -> Dict[str, Any]: Returns: dict: Parameter context to be used with subsequent use_aws operations """ + start_time = time.time() identifiers_response = "" + domain_id = "" try: ctx = ProjectContext() domain_id = safe_get_attr(ctx, "domain_id", "") project_id = safe_get_attr(ctx, "project_id", "") region = safe_get_attr(ctx, "region", "") identifiers_response = create_smus_context_identifiers_response(domain_id, project_id, region) + + latency = (time.time() - start_time) * 1000 + write_log_entry("SMUS-Local-MCP", "aws_context_provider", 0, domain_id, latency) + return {"response": identifiers_response} except Exception as e: logger.error(f"Error providing SMUS context identifiers: {e}") + # Calculate latency and log error + latency = (time.time() - start_time) * 1000 + error_domain_id = domain_id if domain_id else "unable to retrieve domain id" + write_log_entry("SMUS-Local-MCP", "aws_context_provider", 1, error_domain_id, latency) return {"response": identifiers_response, "error": "Error providing SMUS context identifiers"} diff --git a/template/v3/dirs/etc/sagemaker-ui/sagemaker-mcp/smus-mcp.py b/template/v3/dirs/etc/sagemaker-ui/sagemaker-mcp/smus-mcp.py index 1fb237e3f..ad1a98db5 100644 --- a/template/v3/dirs/etc/sagemaker-ui/sagemaker-mcp/smus-mcp.py +++ b/template/v3/dirs/etc/sagemaker-ui/sagemaker-mcp/smus-mcp.py @@ -7,6 +7,7 @@ import logging import os import re +import time from typing import Any, Dict from mcp.server.fastmcp import FastMCP @@ -83,6 +84,54 @@ def create_smus_context_identifiers_response(domain_id: str, project_id: str, re Again, include only required parameters. Any extra parameters may cause the API to fail. Stick strictly to the schema.""" +def write_log_entry(operation: str, tool_name: str, error_count: int = 0, domain_id: str = "", latency: float = 0) -> None: + """ + Write a log entry to the SageMaker GenAI Jupyter Lab extension log file. + + Args: + operation: The operation being performed + tool_name: The name of the tool being used + error_count: Number of errors (0 for success, 1 for failure) + domain_id: The domain identifier (optional) + latency: Execution latency in milliseconds + """ + try: + log_entry = { + "Operation": operation, + "ToolName": tool_name, + "ErrorCount": error_count, + "DomainId": domain_id, + "Latency": latency, + "_aws": { + "Timestamp": int(time.time() * 1000), + "CloudWatchMetrics": [{ + "Dimensions": [["Operation"], ["ToolName"]], + "Metrics": [ + { + "Name": "ErrorCount", + "Unit": "Count" + }, + { + "Name": "Latency", + "Unit": "Milliseconds" + } + ], + "Namespace": "SagemakerGenAIJupyterLab" + }] + }, + } + + # Write to log file + log_file_path = "/var/log/studio/sagemaker_genai_jl_ext/sagemaker_genai_jl_extension.log" + os.makedirs(os.path.dirname(log_file_path), exist_ok=True) + with open(log_file_path, "a") as log_file: + log_file.write(json.dumps(log_entry) + "\n") + + except Exception: + # Fail quietly - logging failures shouldn't affect main functionality + pass + + async def aws_context_provider() -> Dict[str, Any]: """ AWS Context Provider - MUST BE CALLED BEFORE ANY use_aws OPERATIONS @@ -101,16 +150,26 @@ async def aws_context_provider() -> Dict[str, Any]: Returns: dict: Parameter context to be used with subsequent use_aws operations """ + start_time = time.time() identifiers_response = "" + domain_id = "" try: ctx = ProjectContext() domain_id = safe_get_attr(ctx, "domain_id", "") project_id = safe_get_attr(ctx, "project_id", "") region = safe_get_attr(ctx, "region", "") identifiers_response = create_smus_context_identifiers_response(domain_id, project_id, region) + + latency = (time.time() - start_time) * 1000 + write_log_entry("SMUS-Local-MCP", "aws_context_provider", 0, domain_id, latency) + return {"response": identifiers_response} except Exception as e: logger.error(f"Error providing SMUS context identifiers: {e}") + # Calculate latency and log error + latency = (time.time() - start_time) * 1000 + error_domain_id = domain_id if domain_id else "unable to retrieve domain id" + write_log_entry("SMUS-Local-MCP", "aws_context_provider", 1, error_domain_id, latency) return {"response": identifiers_response, "error": "Error providing SMUS context identifiers"} From 31e12972e0a52789332b20e78c4bff19107b3907 Mon Sep 17 00:00:00 2001 From: Eason Ma Date: Mon, 11 Aug 2025 16:09:47 -0700 Subject: [PATCH 5/7] Revert "Updated QCLI download version to . It was pinned to avoid known issue." This reverts commit 9f15e2b24733f9114a7ee9ee7fd89a77b43c8f60. --- template/v2/Dockerfile | 2 +- template/v3/Dockerfile | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/template/v2/Dockerfile b/template/v2/Dockerfile index 2afa47548..929c89e3d 100644 --- a/template/v2/Dockerfile +++ b/template/v2/Dockerfile @@ -63,7 +63,7 @@ RUN apt-get update && apt-get upgrade -y && \ rm -rf aws awscliv2.zip && \ : && \ # Install Q CLI - curl --proto '=https' --tlsv1.2 -sSf "https://desktop-release.q.us-east-1.amazonaws.com/latest/q-x86_64-linux.zip" -o "q.zip" && \ + curl --proto '=https' --tlsv1.2 -sSf "https://desktop-release.q.us-east-1.amazonaws.com/1.12.7/q-x86_64-linux.zip" -o "q.zip" && \ unzip q.zip && \ Q_INSTALL_GLOBAL=true ./q/install.sh --no-confirm && \ rm -rf q q.zip && \ diff --git a/template/v3/Dockerfile b/template/v3/Dockerfile index 65d4da7ea..644caf8aa 100644 --- a/template/v3/Dockerfile +++ b/template/v3/Dockerfile @@ -63,7 +63,7 @@ RUN apt-get update && apt-get upgrade -y && \ rm -rf aws awscliv2.zip && \ : && \ # Install Q CLI - curl --proto '=https' --tlsv1.2 -sSf "https://desktop-release.q.us-east-1.amazonaws.com/latest/q-x86_64-linux.zip" -o "q.zip" && \ + curl --proto '=https' --tlsv1.2 -sSf "https://desktop-release.q.us-east-1.amazonaws.com/1.12.7/q-x86_64-linux.zip" -o "q.zip" && \ unzip q.zip && \ Q_INSTALL_GLOBAL=true ./q/install.sh --no-confirm && \ rm -rf q q.zip && \ From a087289f8f6250063904192e4eabf319dc3fdeb3 Mon Sep 17 00:00:00 2001 From: Eason Ma Date: Mon, 11 Aug 2025 16:16:50 -0700 Subject: [PATCH 6/7] black reformat for local mcp script --- .../sagemaker-ui/sagemaker-mcp/smus-mcp.py | 38 +++++++++---------- .../sagemaker-ui/sagemaker-mcp/smus-mcp.py | 38 +++++++++---------- 2 files changed, 36 insertions(+), 40 deletions(-) diff --git a/template/v2/dirs/etc/sagemaker-ui/sagemaker-mcp/smus-mcp.py b/template/v2/dirs/etc/sagemaker-ui/sagemaker-mcp/smus-mcp.py index ad1a98db5..e37dbae19 100644 --- a/template/v2/dirs/etc/sagemaker-ui/sagemaker-mcp/smus-mcp.py +++ b/template/v2/dirs/etc/sagemaker-ui/sagemaker-mcp/smus-mcp.py @@ -84,10 +84,12 @@ def create_smus_context_identifiers_response(domain_id: str, project_id: str, re Again, include only required parameters. Any extra parameters may cause the API to fail. Stick strictly to the schema.""" -def write_log_entry(operation: str, tool_name: str, error_count: int = 0, domain_id: str = "", latency: float = 0) -> None: +def write_log_entry( + operation: str, tool_name: str, error_count: int = 0, domain_id: str = "", latency: float = 0 +) -> None: """ Write a log entry to the SageMaker GenAI Jupyter Lab extension log file. - + Args: operation: The operation being performed tool_name: The name of the tool being used @@ -104,29 +106,25 @@ def write_log_entry(operation: str, tool_name: str, error_count: int = 0, domain "Latency": latency, "_aws": { "Timestamp": int(time.time() * 1000), - "CloudWatchMetrics": [{ - "Dimensions": [["Operation"], ["ToolName"]], - "Metrics": [ - { - "Name": "ErrorCount", - "Unit": "Count" - }, - { - "Name": "Latency", - "Unit": "Milliseconds" - } - ], - "Namespace": "SagemakerGenAIJupyterLab" - }] + "CloudWatchMetrics": [ + { + "Dimensions": [["Operation"], ["ToolName"]], + "Metrics": [ + {"Name": "ErrorCount", "Unit": "Count"}, + {"Name": "Latency", "Unit": "Milliseconds"}, + ], + "Namespace": "SagemakerGenAIJupyterLab", + } + ], }, } - + # Write to log file log_file_path = "/var/log/studio/sagemaker_genai_jl_ext/sagemaker_genai_jl_extension.log" os.makedirs(os.path.dirname(log_file_path), exist_ok=True) with open(log_file_path, "a") as log_file: log_file.write(json.dumps(log_entry) + "\n") - + except Exception: # Fail quietly - logging failures shouldn't affect main functionality pass @@ -159,10 +157,10 @@ async def aws_context_provider() -> Dict[str, Any]: project_id = safe_get_attr(ctx, "project_id", "") region = safe_get_attr(ctx, "region", "") identifiers_response = create_smus_context_identifiers_response(domain_id, project_id, region) - + latency = (time.time() - start_time) * 1000 write_log_entry("SMUS-Local-MCP", "aws_context_provider", 0, domain_id, latency) - + return {"response": identifiers_response} except Exception as e: logger.error(f"Error providing SMUS context identifiers: {e}") diff --git a/template/v3/dirs/etc/sagemaker-ui/sagemaker-mcp/smus-mcp.py b/template/v3/dirs/etc/sagemaker-ui/sagemaker-mcp/smus-mcp.py index ad1a98db5..e37dbae19 100644 --- a/template/v3/dirs/etc/sagemaker-ui/sagemaker-mcp/smus-mcp.py +++ b/template/v3/dirs/etc/sagemaker-ui/sagemaker-mcp/smus-mcp.py @@ -84,10 +84,12 @@ def create_smus_context_identifiers_response(domain_id: str, project_id: str, re Again, include only required parameters. Any extra parameters may cause the API to fail. Stick strictly to the schema.""" -def write_log_entry(operation: str, tool_name: str, error_count: int = 0, domain_id: str = "", latency: float = 0) -> None: +def write_log_entry( + operation: str, tool_name: str, error_count: int = 0, domain_id: str = "", latency: float = 0 +) -> None: """ Write a log entry to the SageMaker GenAI Jupyter Lab extension log file. - + Args: operation: The operation being performed tool_name: The name of the tool being used @@ -104,29 +106,25 @@ def write_log_entry(operation: str, tool_name: str, error_count: int = 0, domain "Latency": latency, "_aws": { "Timestamp": int(time.time() * 1000), - "CloudWatchMetrics": [{ - "Dimensions": [["Operation"], ["ToolName"]], - "Metrics": [ - { - "Name": "ErrorCount", - "Unit": "Count" - }, - { - "Name": "Latency", - "Unit": "Milliseconds" - } - ], - "Namespace": "SagemakerGenAIJupyterLab" - }] + "CloudWatchMetrics": [ + { + "Dimensions": [["Operation"], ["ToolName"]], + "Metrics": [ + {"Name": "ErrorCount", "Unit": "Count"}, + {"Name": "Latency", "Unit": "Milliseconds"}, + ], + "Namespace": "SagemakerGenAIJupyterLab", + } + ], }, } - + # Write to log file log_file_path = "/var/log/studio/sagemaker_genai_jl_ext/sagemaker_genai_jl_extension.log" os.makedirs(os.path.dirname(log_file_path), exist_ok=True) with open(log_file_path, "a") as log_file: log_file.write(json.dumps(log_entry) + "\n") - + except Exception: # Fail quietly - logging failures shouldn't affect main functionality pass @@ -159,10 +157,10 @@ async def aws_context_provider() -> Dict[str, Any]: project_id = safe_get_attr(ctx, "project_id", "") region = safe_get_attr(ctx, "region", "") identifiers_response = create_smus_context_identifiers_response(domain_id, project_id, region) - + latency = (time.time() - start_time) * 1000 write_log_entry("SMUS-Local-MCP", "aws_context_provider", 0, domain_id, latency) - + return {"response": identifiers_response} except Exception as e: logger.error(f"Error providing SMUS context identifiers: {e}") From 413ff574135a84515b2cc4176bc0fcec070c139d Mon Sep 17 00:00:00 2001 From: Eason Ma Date: Mon, 11 Aug 2025 16:20:14 -0700 Subject: [PATCH 7/7] Helping to reformat pysdk files --- .../etc/sagemaker/sm_pysdk_default_config.py | 35 +++++++++---------- .../etc/sagemaker/sm_pysdk_default_config.py | 35 +++++++++---------- 2 files changed, 34 insertions(+), 36 deletions(-) diff --git a/template/v2/dirs/etc/sagemaker/sm_pysdk_default_config.py b/template/v2/dirs/etc/sagemaker/sm_pysdk_default_config.py index 0489d4028..7038730e7 100644 --- a/template/v2/dirs/etc/sagemaker/sm_pysdk_default_config.py +++ b/template/v2/dirs/etc/sagemaker/sm_pysdk_default_config.py @@ -6,8 +6,13 @@ def generate_intelligent_default_config(metadata: str) -> dict: - has_vpc = metadata["SecurityGroupIds"] and metadata["Subnets"] and metadata["SecurityGroupIds"] != [''] and metadata["Subnets"] != [''] - + has_vpc = ( + metadata["SecurityGroupIds"] + and metadata["Subnets"] + and metadata["SecurityGroupIds"] != [""] + and metadata["Subnets"] != [""] + ) + config = { "SchemaVersion": "1.0", "SageMaker": { @@ -34,31 +39,21 @@ def generate_intelligent_default_config(metadata: str) -> dict: "TrainingJob": {"RoleArn": metadata["UserRoleArn"]}, }, } - + if has_vpc: vpc_config = {"SecurityGroupIds": metadata["SecurityGroupIds"], "Subnets": metadata["Subnets"]} config["SageMaker"]["PythonSDK"]["Modules"]["RemoteFunction"]["VpcConfig"] = vpc_config config["SageMaker"]["PythonSDK"]["Modules"]["NotebookJob"]["VpcConfig"] = vpc_config config["SageMaker"]["MonitoringSchedule"] = { - "MonitoringScheduleConfig": { - "MonitoringJobDefinition": { - "NetworkConfig": {"VpcConfig": vpc_config} - } - } - } - config["SageMaker"]["AutoMLJob"] = { - "AutoMLJobConfig": { - "SecurityConfig": {"VpcConfig": vpc_config} - } - } - config["SageMaker"]["AutoMLJobV2"] = { - "SecurityConfig": {"VpcConfig": vpc_config} + "MonitoringScheduleConfig": {"MonitoringJobDefinition": {"NetworkConfig": {"VpcConfig": vpc_config}}} } + config["SageMaker"]["AutoMLJob"] = {"AutoMLJobConfig": {"SecurityConfig": {"VpcConfig": vpc_config}}} + config["SageMaker"]["AutoMLJobV2"] = {"SecurityConfig": {"VpcConfig": vpc_config}} config["SageMaker"]["CompilationJob"] = {"VpcConfig": vpc_config} config["SageMaker"]["Model"]["VpcConfig"] = vpc_config config["SageMaker"]["ProcessingJob"]["NetworkConfig"] = {"VpcConfig": vpc_config} config["SageMaker"]["TrainingJob"]["VpcConfig"] = vpc_config - + return config @@ -93,7 +88,11 @@ def generate_intelligent_default_config(metadata: str) -> dict: } # Not create config file when invalid value exists in metadata - empty_values = [key for key, value in metadata.items() if key not in ["SecurityGroupIds", "Subnets"] and (value == "" or value == [""])] + empty_values = [ + key + for key, value in metadata.items() + if key not in ["SecurityGroupIds", "Subnets"] and (value == "" or value == [""]) + ] if empty_values: raise AttributeError(f"There are empty values in the metadata: {empty_values}") diff --git a/template/v3/dirs/etc/sagemaker/sm_pysdk_default_config.py b/template/v3/dirs/etc/sagemaker/sm_pysdk_default_config.py index 328dddb99..35b559e96 100644 --- a/template/v3/dirs/etc/sagemaker/sm_pysdk_default_config.py +++ b/template/v3/dirs/etc/sagemaker/sm_pysdk_default_config.py @@ -6,9 +6,14 @@ def generate_intelligent_default_config(metadata: str) -> dict: - has_vpc = metadata["SecurityGroupIds"] and metadata["Subnets"] and metadata["SecurityGroupIds"] != [''] and metadata["Subnets"] != [''] + has_vpc = ( + metadata["SecurityGroupIds"] + and metadata["Subnets"] + and metadata["SecurityGroupIds"] != [""] + and metadata["Subnets"] != [""] + ) vpc_config = {"SecurityGroupIds": metadata["SecurityGroupIds"], "Subnets": metadata["Subnets"]} if has_vpc else None - + config = { "SchemaVersion": "1.0", "SageMaker": { @@ -35,30 +40,20 @@ def generate_intelligent_default_config(metadata: str) -> dict: "TrainingJob": {"RoleArn": metadata["UserRoleArn"]}, }, } - + if has_vpc: config["SageMaker"]["PythonSDK"]["Modules"]["RemoteFunction"]["VpcConfig"] = vpc_config config["SageMaker"]["PythonSDK"]["Modules"]["NotebookJob"]["VpcConfig"] = vpc_config config["SageMaker"]["MonitoringSchedule"] = { - "MonitoringScheduleConfig": { - "MonitoringJobDefinition": { - "NetworkConfig": {"VpcConfig": vpc_config} - } - } - } - config["SageMaker"]["AutoMLJob"] = { - "AutoMLJobConfig": { - "SecurityConfig": {"VpcConfig": vpc_config} - } - } - config["SageMaker"]["AutoMLJobV2"] = { - "SecurityConfig": {"VpcConfig": vpc_config} + "MonitoringScheduleConfig": {"MonitoringJobDefinition": {"NetworkConfig": {"VpcConfig": vpc_config}}} } + config["SageMaker"]["AutoMLJob"] = {"AutoMLJobConfig": {"SecurityConfig": {"VpcConfig": vpc_config}}} + config["SageMaker"]["AutoMLJobV2"] = {"SecurityConfig": {"VpcConfig": vpc_config}} config["SageMaker"]["CompilationJob"] = {"VpcConfig": vpc_config} config["SageMaker"]["Model"]["VpcConfig"] = vpc_config config["SageMaker"]["ProcessingJob"]["NetworkConfig"] = {"VpcConfig": vpc_config} config["SageMaker"]["TrainingJob"]["VpcConfig"] = vpc_config - + return config @@ -93,7 +88,11 @@ def generate_intelligent_default_config(metadata: str) -> dict: } # Not create config file when invalid value exists in metadata - empty_values = [key for key, value in metadata.items() if key not in ["SecurityGroupIds", "Subnets"] and (value == "" or value == [""])] + empty_values = [ + key + for key, value in metadata.items() + if key not in ["SecurityGroupIds", "Subnets"] and (value == "" or value == [""]) + ] if empty_values: raise AttributeError(f"There are empty values in the metadata: {empty_values}")