Setting the file. One moment. Collect Diagnostics · AWS AI ML · aws/agent-toolkit-for-aws · Skills Docs10
Setup DevOps Agent
24.7
Code Output Guide · references
RDS Oracle
references/endpoint-diagnostics/scripts/collect_diagnostics.py
Python·348 lines·13 KB
17 datetime(2024, 3, 15, 14, 0, tzinfo=timezone.utc),
18 datetime(2024, 3, 15, 15, 0, tzinfo=timezone.utc),
19 ),
20 )
21
22Note: Log collection retrieves up to 100 events without pagination.
23This provides a representative sample, not exhaustive logs.
24"""
25
26from __future__ import annotations
27
28from dataclasses import dataclass
29from datetime import datetime, timedelta, timezone
30from typing import Any
31
32import boto3
33from botocore.exceptions import ClientError
34
35
36@dataclass(frozen=True)
37class TimeRange:
38 """Time window for data collection. Use factory methods to construct.
39
40 Either specify an absolute window (start + end) or a relative lookback
41 (minutes from now). Do not mix — if start is set, minutes is ignored.
42 """
43
44 start: datetime | None = None
45 end: datetime | None = None
46 minutes: int = 5
47
48 @classmethod
49 def last(cls, minutes: int) -> TimeRange:
50 """Relative window: collect the last N minutes of data."""
51 return cls(minutes=minutes)
52
53 @classmethod
54 def between(cls, start: datetime, end: datetime) -> TimeRange:
55 """Absolute window: collect data between two UTC timestamps."""
56 return cls(start=start, end=end)
57
58 def resolve(self) -> tuple[datetime, datetime]:
59 """Resolve to concrete (start, end) datetimes."""
60 now = datetime.now(timezone.utc)
61 if self.start is not None:
62 return self.start, self.end if self.end is not None else now
63 return now - timedelta(minutes=self.minutes), now
64
65
66def collect_endpoint_diagnostics(
67 endpoint_name: str,
68 region: str,
69 time_range: TimeRange | None = None,
70) -> dict[str, Any]:
71 """Collect endpoint status, CloudWatch metrics, and recent logs.
72
73 Args:
74 endpoint_name: The SageMaker endpoint name.
75 region: AWS region where the endpoint is deployed.
76 time_range: Time window for metrics and logs. Defaults to last 5 min.
77 Use TimeRange.last(N) for relative or TimeRange.between(s, e)
78 for absolute UTC windows.
79
80 Returns:
81 Dict with keys: endpoint_status, metrics, recent_logs, errors,
82 reference_docs.
83 """
84 if not endpoint_name or not region:
85 raise ValueError("endpoint_name and region are required")
86
87 if time_range is None:
88 time_range = TimeRange()
89
90 window_start, window_end = time_range.resolve()
91
92 sm_client = boto3.client("sagemaker", region_name=region)
93 cw_client = boto3.client("cloudwatch", region_name=region)
94 logs_client = boto3.client("logs", region_name=region)
95
96 errors: list[str] = []
97 endpoint_status: dict[str, Any] = {}
98 metrics: dict[str, Any] = {}
99 recent_logs: list[dict[str, Any]] = []
100 first_variant_name: str = ""
101
102 # Step 1: DescribeEndpoint
103 try:
104 resp = sm_client.describe_endpoint(EndpointName=endpoint_name)
105 endpoint_status = {
106 "status": resp.get("EndpointStatus"),
107 "failure_reason": resp.get("FailureReason"),
108 "creation_time": str(resp.get("CreationTime")),
109 "last_modified_time": str(resp.get("LastModifiedTime")),
110 "production_variants": [
111 {
112 "name": v.get("VariantName"),
113 "instance_type": v.get("InstanceType"),
114 "current_instance_count": v.get("CurrentInstanceCount"),
115 "current_weight": v.get("CurrentWeight"),
116 }
117 for v in resp.get("ProductionVariants", [])
118 ],
119 }
120 variants = resp.get("ProductionVariants", [])
121 if variants:
122 first_variant_name = variants[0].get("VariantName", "")
123 except ClientError as e:
124 errors.append(f"DescribeEndpoint failed: {e.response['Error']['Message']}")
125
126 # Step 2: GetMetricData
127 try:
128 # SageMaker publishes invocation metrics with {EndpointName, VariantName} dimensions.
129 # CloudWatch requires an exact dimension-set match, so we must include VariantName.
130 variant_name = first_variant_name or "AllTraffic"
131 endpoint_variant_dim = [
132 {"Name": "EndpointName", "Value": endpoint_name},
133 {"Name": "VariantName", "Value": variant_name},
134 ]
135
136 if not first_variant_name:
137 errors.append(
138 f"Variant name unavailable from DescribeEndpoint; defaulting to 'AllTraffic' for metrics"
139 )
140
141 # Invocation metrics (require EndpointName + VariantName)
142 metric_queries = [
143 {
144 "Id": "invocations",
145 "MetricStat": {
146 "Metric": {
147 "Namespace": "AWS/SageMaker",
148 "MetricName": "Invocations",
149 "Dimensions": endpoint_variant_dim,
150 },
151 "Period": 60,
152 "Stat": "Sum",
153 },
154 },
155 {
156 "Id": "errors_4xx",
157 "MetricStat": {
158 "Metric": {
159 "Namespace": "AWS/SageMaker",
160 "MetricName": "Invocation4XXErrors",
161 "Dimensions": endpoint_variant_dim,
162 },
163 "Period": 60,
164 "Stat": "Sum",
165 },
166 },
167 {
168 "Id": "errors_5xx",
169 "MetricStat": {
170 "Metric": {
171 "Namespace": "AWS/SageMaker",
172 "MetricName": "Invocation5XXErrors",
173 "Dimensions": endpoint_variant_dim,
174 },
175 "Period": 60,
176 "Stat": "Sum",
177 },
178 },
179 {
180 "Id": "invocation_model_errors",
181 "MetricStat": {
182 "Metric": {
183 "Namespace": "AWS/SageMaker",
184 "MetricName": "InvocationModelErrors",
185 "Dimensions": endpoint_variant_dim,
186 },
187 "Period": 60,
188 "Stat": "Sum",
189 },
190 },
191 {
192 "Id": "model_latency_avg",
193 "MetricStat": {
194 "Metric": {
195 "Namespace": "AWS/SageMaker",
196 "MetricName": "ModelLatency",
197 "Dimensions": endpoint_variant_dim,
198 },
199 "Period": 60,
200 "Stat": "Average",
201 },
202 },
203 {
204 "Id": "model_latency_p99",
205 "MetricStat": {
206 "Metric": {
207 "Namespace": "AWS/SageMaker",
208 "MetricName": "ModelLatency",
209 "Dimensions": endpoint_variant_dim,
210 },
211 "Period": 60,
212 "Stat": "p99",
213 },
214 },
215 {
216 "Id": "overhead_latency",
217 "MetricStat": {
218 "Metric": {
219 "Namespace": "AWS/SageMaker",
220 "MetricName": "OverheadLatency",
221 "Dimensions": endpoint_variant_dim,
222 },
223 "Period": 60,
224 "Stat": "Average",
225 },
226 },
227 ]
228
229 # Instance-level metrics (CPU/Memory/GPU) use /aws/sagemaker/Endpoints namespace
230 if first_variant_name:
231 variant_dim = [
232 {"Name": "EndpointName", "Value": endpoint_name},
233 {"Name": "VariantName", "Value": first_variant_name},
234 ]
235 metric_queries.extend(
236 [
237 {
238 "Id": "cpu_util",
239 "MetricStat": {
240 "Metric": {
241 "Namespace": "/aws/sagemaker/Endpoints",
242 "MetricName": "CPUUtilization",
243 "Dimensions": variant_dim,
244 },
245 "Period": 60,
246 "Stat": "Average",
247 },
248 },
249 {
250 "Id": "memory_util",
251 "MetricStat": {
252 "Metric": {
253 "Namespace": "/aws/sagemaker/Endpoints",
254 "MetricName": "MemoryUtilization",
255 "Dimensions": variant_dim,
256 },
257 "Period": 60,
258 "Stat": "Average",
259 },
260 },
261 {
262 "Id": "gpu_util",
263 "MetricStat": {
264 "Metric": {
265 "Namespace": "/aws/sagemaker/Endpoints",
266 "MetricName": "GPUUtilization",
267 "Dimensions": variant_dim,
268 },
269 "Period": 60,
270 "Stat": "Average",
271 },
272 },
273 {
274 "Id": "gpu_memory_util",
275 "MetricStat": {
276 "Metric": {
277 "Namespace": "/aws/sagemaker/Endpoints",
278 "MetricName": "GPUMemoryUtilization",
279 "Dimensions": variant_dim,
280 },
281 "Period": 60,
282 "Stat": "Average",
283 },
284 },
285 ]
286 )
287 else:
288 errors.append(
289 "Variant name unavailable; skipped instance-level metrics (CPU/Memory/GPU)"
290 )
291
292 # Warn if multiple variants exist — metrics are only for the first
293 variants = endpoint_status.get("production_variants", [])
294 if len(variants) > 1:
295 variant_names = [v["name"] for v in variants]
296 errors.append(
297 f"Multiple variants detected: {variant_names}. "
298 f"Metrics shown are for '{variant_name}' only. "
299 f"Ask the user which variant to diagnose if needed."
300 )
301
302 cw_resp = cw_client.get_metric_data(
303 MetricDataQueries=metric_queries,
304 StartTime=window_start,
305 EndTime=window_end,
306 )
307 for result in cw_resp.get("MetricDataResults", []):
308 metrics[result["Id"]] = result.get("Values", [])
309 except ClientError as e:
310 errors.append(f"GetMetricData failed: {e.response['Error']['Message']}")
311
312 # Step 3: FilterLogEvents — use a wider window for logs
313 try:
314 log_group = f"/aws/sagemaker/Endpoints/{endpoint_name}"
315 if time_range.start is not None:
316 # Respect user-specified absolute window boundaries
317 log_window_start = max(window_start, window_end - timedelta(minutes=15))
318 else:
319 # For relative windows, always use 15-minute log window per SKILL.md
320 log_window_start = window_end - timedelta(minutes=15)
321 log_resp = logs_client.filter_log_events(
322 logGroupName=log_group,
323 startTime=int(log_window_start.timestamp() * 1000),
324 endTime=int(window_end.timestamp() * 1000),
325 limit=100,
326 )
327 recent_logs = [
328 {
329 "message": e.get("message"),
330 "timestamp": datetime.fromtimestamp(
331 e.get("timestamp", 0) / 1000, tz=timezone.utc
332 ).isoformat(),
333 }
334 for e in log_resp.get("events", [])
335 ]
336 except ClientError as e:
337 errors.append(f"FilterLogEvents failed: {e.response['Error']['Message']}")
338
339 return {
340 "endpoint_status": endpoint_status,
341 "metrics": metrics,
342 "recent_logs": recent_logs,
343 "errors": errors,
344 "reference_docs": {
345 "metrics_definitions": "https://docs.aws.amazon.com/sagemaker/latest/dg/monitoring-cloudwatch.html",
346 "troubleshooting": "https://docs.aws.amazon.com/sagemaker/latest/dg/deploy-model-troubleshoot.html",
347 },
348 }