Skip to content

Update: Un-pin QCLI version to use latest #821

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
wants to merge 9 commits into
base: main
Choose a base branch
from
+148 −36
Open

Update: Un-pin QCLI version to use latest #821

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
9 commits
Select commit Hold shift + click to select a range
463ce2e
feat: add Amazon Q agent configuration and MCP migration
YuchengMaUTK Jul 24, 2025
db84309
fix: add smus local mcp into tools list
YuchengMaUTK Jul 25, 2025
7049919
Merge branch 'aws:main' into main
YuchengMaUTK Aug 8, 2025
9f15e2b
Updated QCLI download version to . It was pinned to avoid known issue.
YuchengMaUTK Aug 8, 2025
1a65cc5
Merge branch 'aws:main' into main
YuchengMaUTK Aug 11, 2025
04eacca
Emitting metrics for SMUS local mcp via Q chat extension log file entry
YuchengMaUTK Aug 11, 2025
31e1297
Revert "Updated QCLI download version to . It was pinned to avoid kno…
YuchengMaUTK Aug 11, 2025
a087289
black reformat for local mcp script
YuchengMaUTK Aug 11, 2025
413ff57
Helping to reformat pysdk files
YuchengMaUTK Aug 11, 2025
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
57 changes: 57 additions & 0 deletions template/v2/dirs/etc/sagemaker-ui/sagemaker-mcp/smus-mcp.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -7,6 +7,7 @@
import logging
import os
import re
import time
from typing import Any, Dict

from mcp.server.fastmcp import FastMCP
Expand Down Expand Up @@ -83,6 +84,52 @@ 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
Expand All @@ -101,16 +148,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"}


Expand Down
35 changes: 17 additions & 18 deletions template/v2/dirs/etc/sagemaker/sm_pysdk_default_config.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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": {
Expand All @@ -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


Expand Down Expand Up @@ -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}")

Expand Down
57 changes: 57 additions & 0 deletions template/v3/dirs/etc/sagemaker-ui/sagemaker-mcp/smus-mcp.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -7,6 +7,7 @@
import logging
import os
import re
import time
from typing import Any, Dict

from mcp.server.fastmcp import FastMCP
Expand Down Expand Up @@ -83,6 +84,52 @@ 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
Expand All @@ -101,16 +148,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"}


Expand Down
35 changes: 17 additions & 18 deletions template/v3/dirs/etc/sagemaker/sm_pysdk_default_config.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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": {
Expand All @@ -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


Expand Down Expand Up @@ -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}")

Expand Down