From 7f0f2fd8d3c167bc1b9ea2403cf50b6323dd250d Mon Sep 17 00:00:00 2001 From: ZealSV Date: Fri, 30 Oct 2026 09:44:33 -0700 Subject: [PATCH] feat(serve): describe an endpoint's traffic and size workloads from it Add describe_endpoint_traffic(), which describes the traffic one model on an endpoint served over a window, relative (lookback back from end_time, seven days by default) or absolute (start_time to end_time), up to seven days. Times must be timezone-aware and are reported in UTC. - Concurrency for the whole deployment: each hour's maximum ConcurrentRequestsPerCopy or ConcurrentRequestsPerModel times the current copies or instances. - With detailed observability on, the request count and the prompt and generated tokens of the requests the model completed, from the serving container's token histograms. They're read through CloudWatch's Prometheus-compatible query API with the session's credentials, over the CloudWatch client's HTTP session (its proxy, CA bundle and timeouts), with retries on throttling, server errors and dropped connections. The read is best effort: a failure is reported in token_metrics_error and logged. The histograms are listed per serving framework (vLLM today), and the result names the framework read. They aren't read when the endpoint configuration turns detailed observability off. - Tokens and concurrency are TrafficStats: min, p25, p50, p70, p80, p90, p95, p99, max (p100), mean and standard deviation. Token percentiles interpolate within the histograms' buckets. Concurrency's are over the window's hourly maximums, so its max is the busiest hour's. - The variant's Data Capture prefix, when the endpoint captures requests. The histograms are cumulative, so a series' traffic in the window is the sum of its increments between daily snapshots, from the start's to the window's last. A snapshot below the one before it follows a restart of its container and counts from zero. An adapter's requests run in its base component's container, which publishes the concurrency and the token histograms for them, so an adapter is described by its base component, and finding the component skips adapters. Add Workload.from_endpoint(), which builds a synthetic workload from the token metrics at the busiest hour's concurrency, or with use_data_capture=True replays the variant's captured requests, so a benchmark or recommendation job can be sized from what an endpoint serves today. Other workload parameters, such as tokenizer and hf_token, pass through to the workload. Integration tests: a no-GPU check that the signed token query is accepted and parsed by CloudWatch's query API, and an e2e step that describes the benchmarked endpoint's traffic and runs a benchmark sized from it. --- .../ai_inference_recommender/__init__.py | 8 + .../serve/ai_inference_recommender/traffic.py | 688 +++++++++++++ .../ai_inference_recommender/workload.py | 115 +++ ...ce_recommender_enhancements_integration.py | 73 +- .../test_traffic.py | 915 ++++++++++++++++++ 5 files changed, 1797 insertions(+), 2 deletions(-) create mode 100644 sagemaker-serve/src/sagemaker/serve/ai_inference_recommender/traffic.py create mode 100644 sagemaker-serve/tests/unit/test_ai_inference_recommender/test_traffic.py diff --git a/sagemaker-serve/src/sagemaker/serve/ai_inference_recommender/__init__.py b/sagemaker-serve/src/sagemaker/serve/ai_inference_recommender/__init__.py index 12132b5542..d8082e72ab 100644 --- a/sagemaker-serve/src/sagemaker/serve/ai_inference_recommender/__init__.py +++ b/sagemaker-serve/src/sagemaker/serve/ai_inference_recommender/__init__.py @@ -39,6 +39,11 @@ compare_benchmarks, ) from sagemaker.serve.ai_inference_recommender.secrets import Secret +from sagemaker.serve.ai_inference_recommender.traffic import ( + EndpointTraffic, + TrafficStats, + describe_endpoint_traffic, +) from sagemaker.serve.ai_inference_recommender.workload import Workload from sagemaker.serve.ai_inference_recommender._model_builder_methods import ( start_benchmark, @@ -51,14 +56,17 @@ "BenchmarkMetrics", "BenchmarkResult", "BenchmarkSearchResult", + "EndpointTraffic", "FeatureGatedError", "InferenceFramework", "PerformanceTarget", "RecommendationJob", "Secret", + "TrafficStats", "Workload", "WorkloadValidationError", "compare_benchmarks", + "describe_endpoint_traffic", "list_benchmarks", "list_recommendations", "start_benchmark", diff --git a/sagemaker-serve/src/sagemaker/serve/ai_inference_recommender/traffic.py b/sagemaker-serve/src/sagemaker/serve/ai_inference_recommender/traffic.py new file mode 100644 index 0000000000..aba9feaa71 --- /dev/null +++ b/sagemaker-serve/src/sagemaker/serve/ai_inference_recommender/traffic.py @@ -0,0 +1,688 @@ +# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"). You +# may not use this file except in compliance with the License. A copy of +# the License is located at +# +# http://aws.amazon.com/apache2.0/ +# +# or in the "license" file accompanying this file. This file is +# distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF +# ANY KIND, either express or implied. See the License for the specific +# language governing permissions and limitations under the License. +"""Describe the traffic one model on a SageMaker endpoint served, to size a workload from it. + +Concurrency comes from the endpoint's invocation metrics in CloudWatch: each hour's maximum +``ConcurrentRequestsPerCopy`` (an inference component) or ``ConcurrentRequestsPerModel`` (a model +on a production variant), times the current copies or instances. + +Token sizes come from detailed observability. With it on (the endpoint configuration's +``MetricsConfig.EnableDetailedObservability``), the serving container's histograms of each +completed request's prompt and generated tokens are published to CloudWatch, which serves them +through its Prometheus-compatible query API. Those calls are signed with the session's +credentials and need ``cloudwatch:GetMetricData`` and ``cloudwatch:ListMetrics``. +``_TOKEN_METRICS`` lists the histograms read for each serving framework. +""" + +from __future__ import absolute_import + +import itertools +import json +import logging +import math +import time +import urllib.parse +from dataclasses import dataclass, field +from datetime import datetime, timedelta, timezone +from typing import Any, Callable, Dict, List, Optional, Tuple, Union + +import boto3 +from botocore.auth import SigV4Auth +from botocore.awsrequest import AWSRequest +from botocore.exceptions import BotoCoreError, ClientError + +from sagemaker.core.helper.session_helper import Session +from sagemaker.core.telemetry.constants import Feature +from sagemaker.core.telemetry.telemetry_logging import _telemetry_emitter +from sagemaker.serve.ai_inference_recommender.listing import _boto_session + +logger = logging.getLogger(__name__) + +# The longest window that can be read: a CloudWatch PromQL query spans at most seven days. +MAX_WINDOW = timedelta(days=7) +_DAY_SECONDS = 86_400 +_HOUR_SECONDS = 3_600 +_ENDPOINT_LABEL = "aws.sagemaker.endpoint.name" +_COMPONENT_LABEL = "aws.sagemaker.inference_component.name" +_VARIANT_LABEL = "aws.sagemaker.variant.name" +_STAT_LABEL = "stat" +# Attempts per query, with exponential backoff, on throttling, server errors and dropped connections. +_QUERY_ATTEMPTS = 3 +_RETRY_BASE_SECONDS = 0.5 +_RETRYABLE_STATUS = frozenset({429, 500, 502, 503, 504}) + + +@dataclass(frozen=True) +class _TokenMetrics: + """The histograms a serving framework publishes of each request's prompt and generated tokens.""" + + framework: str + prompt_tokens: str + generation_tokens: str + + +# Read in this order; the first framework with requests in the window is used. To support +# another framework, add its histograms here. +_TOKEN_METRICS = ( + _TokenMetrics("vllm", "vllm:request_prompt_tokens", "vllm:request_generation_tokens"), +) + + +@dataclass(frozen=True) +class TrafficStats: + """The distribution of one quantity across a window: percentiles, mean and standard deviation. + + ``max`` is the 100th percentile. + """ + + min: float + p25: float + p50: float + p70: float + p80: float + p90: float + p95: float + p99: float + max: float + mean: float + stddev: float + + +_PERCENTILES = ("p25", "p50", "p70", "p80", "p90", "p95", "p99") + + +@dataclass(frozen=True) +class EndpointTraffic: + """The traffic one model on an endpoint served from ``start_time`` to ``end_time``. + + Attributes: + endpoint_name: The endpoint described. + inference_component_name: The inference component described (an adapter's base + component when an adapter is passed), or None for a model on a production variant. + variant_name: The production variant the model runs on. + start_time: The window's start, in UTC. + end_time: The window's end, in UTC. + concurrency: The concurrent requests on the whole deployment across the window's hours: + each hour's per-copy or per-instance maximum times the current copies or instances (at + least one). ``max`` is the busiest hour's. None when the window had no traffic. If the + copies or instances autoscale, the count in an hour may differ from the current one. + detailed_observability: The endpoint configuration's + ``MetricsConfig.EnableDetailedObservability``, or None when the configuration doesn't + say or can't be read. Token sizes aren't read when it's False. + framework: The serving framework whose token histograms were read, or None without them. + requests: Requests the model completed in the window, from its token histograms. 0 + without them. + input_tokens: Prompt tokens per request, or None without token histograms. + output_tokens: Generated tokens per request, or None without token histograms. + token_metrics_error: Why the token histograms couldn't be read, for example a missing + permission or an unreachable query API, or None when the read didn't fail. + data_capture_s3_uri: Where SageMaker Data Capture writes the variant's requests, or None + when the endpoint doesn't capture them. The variant's captures include the requests of + every inference component on it. + """ + + endpoint_name: str + inference_component_name: Optional[str] + variant_name: Optional[str] + start_time: datetime + end_time: datetime + concurrency: Optional[TrafficStats] + detailed_observability: Optional[bool] + framework: Optional[str] + requests: int + input_tokens: Optional[TrafficStats] + output_tokens: Optional[TrafficStats] + token_metrics_error: Optional[str] + data_capture_s3_uri: Optional[str] + + +@dataclass(frozen=True) +class _Target: + """The model on an endpoint whose traffic is read, and how many copies or instances it has.""" + + endpoint_name: str + inference_component_name: Optional[str] + variant_name: Optional[str] + units: int + + +@dataclass +class _Histogram: + """One side's requests in the window: their count, token sum and weighted buckets. + + ``buckets`` maps a bucket's ``(lower, upper)`` bounds to its requests in the window. + """ + + count: float = 0.0 + total: float = 0.0 + buckets: Dict[Tuple[float, float], float] = field(default_factory=dict) + + +@_telemetry_emitter( + feature=Feature.INFERENCE_RECOMMENDER, + func_name="ai_inference_recommender.describe_endpoint_traffic", +) +def describe_endpoint_traffic( + endpoint_name: str, + *, + inference_component_name: Optional[str] = None, + variant_name: Optional[str] = None, + lookback: Optional[timedelta] = None, + start_time: Optional[datetime] = None, + end_time: Optional[datetime] = None, + sagemaker_session: Optional[Union[boto3.session.Session, Session]] = None, +) -> EndpointTraffic: + """Describe the concurrency and token sizes of the traffic one model on an endpoint served. + + Pass the inference component or the production variant to describe. When the endpoint has a + single base inference component, or no inference components and a single variant, it is + found for you. An adapter's requests run in its base component's container, so passing an + adapter describes its base component, whose traffic includes every adapter on it. + + The window is relative or absolute: ``lookback`` back from ``end_time``, or ``start_time`` to + ``end_time``. ``end_time`` defaults to now and ``lookback`` to seven days, the longest window. + Times must be timezone-aware; they are compared and reported in UTC. + + The request count and the token means are exact, unless a serving container restarted in the + window: the requests it completed between the previous daily snapshot and the restart aren't + counted. The token percentiles and standard deviations are estimated from the histograms' + buckets. Metrics can take a few minutes to arrive, so a window that ends in the last few + minutes can miss its newest requests. Token sizes need detailed observability on and a serving + container whose token histograms are read (vLLM); without them, ``input_tokens`` and + ``output_tokens`` are None. Reading them is best effort: when the read fails, + ``token_metrics_error`` says why and a warning is logged. + + Args: + endpoint_name: The SageMaker endpoint name. + inference_component_name: The inference component to describe, for an endpoint that has + them. + variant_name: The production variant to describe, for a model deployed on a variant. + lookback: How far back from ``end_time`` to read, up to seven days. Can't be combined + with ``start_time``. + start_time: The window's start, a timezone-aware ``datetime``. + end_time: The window's end, a timezone-aware ``datetime`` no later than now. + sagemaker_session: Optional session — a ``boto3.session.Session`` or a sagemaker + ``Session`` wrapping one. A default is created if omitted. Use one in the endpoint's + region. + + Returns: + An ``EndpointTraffic``. + + Raises: + ValueError: If the arguments are invalid, or the model to describe can't be determined. + botocore.exceptions.ClientError: If a SageMaker or CloudWatch call fails, for example + when the endpoint doesn't exist. + """ + if not endpoint_name: + raise ValueError("endpoint_name is required.") + if inference_component_name and variant_name: + raise ValueError("Pass at most one of inference_component_name and variant_name.") + start, end = _window( + lookback, start_time, end_time, datetime.fromtimestamp(time.time(), tz=timezone.utc) + ) + + boto_session = _boto_session(sagemaker_session) or boto3.session.Session() + sagemaker = boto_session.client("sagemaker") + endpoint = sagemaker.describe_endpoint(EndpointName=endpoint_name) + target = _resolve_target(sagemaker, endpoint, inference_component_name, variant_name) + observability = _detailed_observability(sagemaker, endpoint) + cloudwatch = boto_session.client("cloudwatch") + concurrency = _concurrency(cloudwatch, target, start, end) + + framework, requests, input_tokens, output_tokens, token_error = None, 0, None, None, None + if observability is not False: + try: + snapshots = _snapshots(_PromQLClient(boto_session, cloudwatch), target, start, end) + except RuntimeError as error: + token_error = str(error) + logger.warning( + "Couldn't read the token metrics of endpoint %s: %s", endpoint_name, error + ) + else: + for metrics in _TOKEN_METRICS: + prompt = _window_histogram(snapshots, f"{metrics.framework}/input") + generation = _window_histogram(snapshots, f"{metrics.framework}/output") + if prompt is not None and generation is not None: + framework, requests = metrics.framework, round(prompt.count) + input_tokens = _histogram_stats(prompt) + output_tokens = _histogram_stats(generation) + break + return EndpointTraffic( + endpoint_name=endpoint_name, + inference_component_name=target.inference_component_name, + variant_name=target.variant_name, + start_time=start, + end_time=end, + concurrency=concurrency, + detailed_observability=observability, + framework=framework, + requests=requests, + input_tokens=input_tokens, + output_tokens=output_tokens, + token_metrics_error=token_error, + data_capture_s3_uri=_data_capture_uri(endpoint, target.variant_name), + ) + + +def _window( + lookback: Optional[timedelta], + start_time: Optional[datetime], + end_time: Optional[datetime], + now: datetime, +) -> Tuple[datetime, datetime]: + """The window's start and end in UTC, to the second.""" + for name, value in (("start_time", start_time), ("end_time", end_time)): + if value is not None and (not isinstance(value, datetime) or value.utcoffset() is None): + raise ValueError( + f"{name} must be a timezone-aware datetime, for example " + f"datetime(2026, 1, 31, tzinfo=timezone.utc); got {value!r}." + ) + if lookback is not None: + if not isinstance(lookback, timedelta) or lookback <= timedelta(0): + raise ValueError(f"lookback must be a positive timedelta, got {lookback!r}.") + if start_time is not None: + raise ValueError("Pass lookback or start_time, not both.") + end = (end_time or now).astimezone(timezone.utc).replace(microsecond=0) + if end > now: + raise ValueError(f"end_time can't be later than now; got {end_time!r}.") + if start_time is not None: + start = start_time.astimezone(timezone.utc).replace(microsecond=0) + else: + start = end - (lookback or MAX_WINDOW) + if start >= end: + raise ValueError(f"start_time must be before end_time; got {start} and {end}.") + if end - start > MAX_WINDOW: + raise ValueError(f"The window can be at most {MAX_WINDOW.days} days; got {end - start}.") + return start, end + + +def _resolve_target( + client, endpoint: Dict[str, Any], component: Optional[str], variant: Optional[str] +) -> _Target: + """Find the inference component or production variant to read, and its copies or instances.""" + endpoint_name = endpoint["EndpointName"] + variants = {v["VariantName"]: v for v in endpoint.get("ProductionVariants", [])} + if component is None and variant is None: + components = _base_components(client, endpoint_name) + if len(components) > 1: + raise ValueError( + f"Endpoint {endpoint_name!r} has inference components {components}; pass " + "inference_component_name to choose one." + ) + if components: + component = components[0] + elif len(variants) > 1: + raise ValueError( + f"Endpoint {endpoint_name!r} has production variants {sorted(variants)}; pass " + "variant_name to choose one." + ) + elif variants: + variant = next(iter(variants)) + else: + raise ValueError(f"Endpoint {endpoint_name!r} has no production variant.") + if component is not None: + described = client.describe_inference_component(InferenceComponentName=component) + if described.get("EndpointName") != endpoint_name: + raise ValueError( + f"Inference component {component!r} is not on endpoint {endpoint_name!r}." + ) + base = (described.get("Specification") or {}).get("BaseInferenceComponentName") + if base: + # An adapter's requests run in its base component's container, which publishes the + # concurrency and the token histograms for them. + component = base + described = client.describe_inference_component(InferenceComponentName=base) + copies = (described.get("RuntimeConfig") or {}).get("CurrentCopyCount") or 0 + return _Target(endpoint_name, component, described.get("VariantName"), copies) + if variant not in variants: + raise ValueError( + f"Endpoint {endpoint_name!r} has no production variant {variant!r}; it has " + f"{sorted(variants)}." + ) + instances = variants[variant].get("CurrentInstanceCount") or 0 + return _Target(endpoint_name, None, variant, instances) + + +def _base_components(client, endpoint_name: str) -> List[str]: + """The endpoint's inference components that aren't adapters. + + An adapter's traffic is read from its base component, so only base components are choices. + """ + components = _inference_components(client, endpoint_name) + if len(components) <= 1: + return components + return [ + name + for name in components + if not ( + client.describe_inference_component(InferenceComponentName=name).get("Specification") + or {} + ).get("BaseInferenceComponentName") + ] + + +def _inference_components(client, endpoint_name: str) -> List[str]: + """The names of the endpoint's inference components.""" + names: List[str] = [] + kwargs: Dict[str, Any] = {"EndpointNameEquals": endpoint_name} + while True: + page = client.list_inference_components(**kwargs) + names += [ic["InferenceComponentName"] for ic in page.get("InferenceComponents", [])] + if not page.get("NextToken"): + return names + kwargs["NextToken"] = page["NextToken"] + + +def _detailed_observability(client, endpoint: Dict[str, Any]) -> Optional[bool]: + """The endpoint configuration's ``EnableDetailedObservability``, or None if it doesn't say. + + An endpoint configuration can be deleted while an endpoint runs on it; then it can't say. + """ + try: + config = client.describe_endpoint_config(EndpointConfigName=endpoint["EndpointConfigName"]) + except ClientError: + return None + enabled = (config.get("MetricsConfig") or {}).get("EnableDetailedObservability") + return enabled if isinstance(enabled, bool) else None + + +def _data_capture_uri(endpoint: Dict[str, Any], variant: Optional[str]) -> Optional[str]: + """Where Data Capture writes the variant's requests, or None when capture is off.""" + capture = endpoint.get("DataCaptureConfig") or {} + destination = capture.get("DestinationS3Uri") + if not capture.get("EnableCapture") or not destination or not variant: + return None + return f"{destination.rstrip('/')}/{endpoint['EndpointName']}/{variant}/" + + +def _concurrency(client, target: _Target, start: datetime, end: datetime) -> Optional[TrafficStats]: + """The deployment's concurrent requests across the window's hours, or None without traffic. + + SageMaker publishes concurrency per copy or per instance, so each hour's maximum is + multiplied by the current copies or instances. + """ + if target.inference_component_name is not None: + metric = "ConcurrentRequestsPerCopy" + dimensions = [{"Name": "InferenceComponentName", "Value": target.inference_component_name}] + else: + metric = "ConcurrentRequestsPerModel" + dimensions = [ + {"Name": "EndpointName", "Value": target.endpoint_name}, + {"Name": "VariantName", "Value": target.variant_name}, + ] + response = client.get_metric_statistics( + Namespace="AWS/SageMaker", + MetricName=metric, + Dimensions=dimensions, + StartTime=start, + EndTime=end, + Period=_HOUR_SECONDS, + Statistics=["Maximum"], + ) + units = max(target.units, 1) + hourly = sorted(point["Maximum"] * units for point in response.get("Datapoints", [])) + if not hourly or hourly[-1] <= 0: + return None + return _sample_stats(hourly) + + +def _sample_stats(values: List[float]) -> TrafficStats: + """The distribution of sorted samples; percentiles interpolate linearly between ranks.""" + + def percentile(fraction: float) -> float: + rank = fraction * (len(values) - 1) + low = math.floor(rank) + high = min(low + 1, len(values) - 1) + return values[low] + (values[high] - values[low]) * (rank - low) + + mean = sum(values) / len(values) + stddev = math.sqrt(sum((value - mean) ** 2 for value in values) / len(values)) + return _stats(percentile, values[0], values[-1], mean, stddev) + + +def _stats( + percentile: Callable[[float], float], low: float, high: float, mean: float, stddev: float +) -> TrafficStats: + """A ``TrafficStats`` from a percentile function, rounded to one decimal place.""" + values = {name: percentile(int(name[1:]) / 100) for name in _PERCENTILES} + return TrafficStats( + min=round(low, 1), + **{name: round(value, 1) for name, value in values.items()}, + max=round(high, 1), + mean=round(mean, 1), + stddev=round(stddev, 1), + ) + + +def _snapshots( + client: "_PromQLClient", target: _Target, start: datetime, end: datetime +) -> List[Dict[Tuple, Dict[str, Any]]]: + """Each token histogram series' last snapshot a day at a time, newest first, then at the start. + + All but the last entry are read at the window's end and a day at a time back while after its + start; together they hold each series' newest snapshot in the window. The last is read at the + start. The histograms are cumulative, so a series' newest snapshot in the window less its + snapshot at the start is the window's requests. + """ + labels = {_ENDPOINT_LABEL: target.endpoint_name} + if target.inference_component_name is not None: + labels[_COMPONENT_LABEL] = target.inference_component_name + else: + labels[_VARIANT_LABEL] = target.variant_name + parts = [ + f'label_replace(last_over_time({_selector(metric, labels)}[1d]), "{_STAT_LABEL}", ' + f'"{metrics.framework}/{side}", "", "")' + for metrics in _TOKEN_METRICS + for side, metric in ( + ("input", metrics.prompt_tokens), + ("output", metrics.generation_tokens), + ) + ] + query = " or ".join(parts) + first, last = int(end.timestamp()), int(start.timestamp()) + times = list(range(first, last, -_DAY_SECONDS)) + [last] + return [client.histograms(query, at) for at in times] + + +def _window_histogram( + snapshots: List[Dict[Tuple, Dict[str, Any]]], stat: str +) -> Optional[_Histogram]: + """One side's requests in the window, or None without requests. + + The histograms are cumulative, so a series' requests between two of its snapshots are the later + one less the earlier, summed over its snapshots from the start's through the window's days. A + snapshot whose count, sum or any bucket is below the one before it follows a restart of its + container, and counts from zero, as does a series that started in the window. + """ + *window, before = snapshots + previous = {key: sample for key, sample in before.items() if dict(key).get(_STAT_LABEL) == stat} + histogram = _Histogram() + for snapshot in reversed(window): + for key, sample in snapshot.items(): + if dict(key).get(_STAT_LABEL) != stat: + continue + earlier = previous.get(key) + _add(histogram, sample, 1.0) + if earlier is not None and not _restarted(earlier, sample): + _add(histogram, earlier, -1.0) + previous[key] = sample + if histogram.count <= 0 or histogram.total < 0 or not any(histogram.buckets.values()): + return None + return histogram + + +def _restarted(earlier: Dict[str, Any], later: Dict[str, Any]) -> bool: + """Whether a cumulative snapshot went down from an earlier one: its container restarted.""" + if float(later.get("count", 0)) < float(earlier.get("count", 0)): + return True + if float(later.get("sum", 0)) < float(earlier.get("sum", 0)): + return True + counts = {(float(low), float(high)): float(n) for _, low, high, n in earlier.get("buckets", [])} + return any( + float(n) < counts.get((float(low), float(high)), 0.0) + for _, low, high, n in later.get("buckets", []) + ) + + +def _add(histogram: _Histogram, sample: Dict[str, Any], sign: float) -> None: + """Add a snapshot's count, sum and buckets to a histogram, or subtract them.""" + histogram.count += sign * float(sample.get("count", 0)) + histogram.total += sign * float(sample.get("sum", 0)) + for _, lower, upper, observations in sample.get("buckets", []): + bounds = (float(lower), float(upper)) + histogram.buckets[bounds] = histogram.buckets.get(bounds, 0.0) + sign * float(observations) + + +def _histogram_stats(histogram: _Histogram) -> TrafficStats: + """A histogram's distribution: the exact mean, and the rest estimated from its buckets. + + Requests are spread evenly within a bucket; an open-ended bucket counts at its finite bound. + Percentiles interpolate linearly within buckets, and the standard deviation uses the buckets' + midpoints. ``min`` and ``max`` are the bounds of the lowest and highest buckets with requests. + """ + count = histogram.count + tolerance = count * 1e-9 + buckets = [ + (*_finite(lower, upper), weight) + for (lower, upper), weight in histogram.buckets.items() + if abs(weight) > tolerance + ] + + def below(point: float, inclusive: bool) -> float: + total = 0.0 + for low, high, weight in buckets: + if low == high: + total += weight if point > high or (inclusive and point == high) else 0.0 + elif point >= high: + total += weight + elif point > low: + total += weight * (point - low) / (high - low) + return total + + # The cumulative count just below and at each bound, so a bucket collapsed to a point is a + # step. A running maximum keeps it from dipping where buckets were resized between snapshots. + bounds = sorted({bound for low, high, _ in buckets for bound in (low, high)}) + knots = [(bound, below(bound, inclusive)) for bound in bounds for inclusive in (False, True)] + seen = list(itertools.accumulate((cumulative for _, cumulative in knots), max)) + + def percentile(fraction: float) -> float: + target = fraction * count - tolerance + index = next((i for i, value in enumerate(seen) if value >= target), len(knots) - 1) + if index == 0 or seen[index] <= seen[index - 1]: + return knots[index][0] + low, high = knots[index - 1][0], knots[index][0] + before, after = seen[index - 1], seen[index] + return low + (high - low) * (target - before) / (after - before) + + first = next((i for i, value in enumerate(seen) if value > tolerance), 0) + lowest = knots[first - 1][0] if first > 0 else knots[0][0] + midpoint_sum = sum(weight * (low + high) / 2 for low, high, weight in buckets) + square_sum = sum(weight * ((low + high) / 2) ** 2 for low, high, weight in buckets) + spread = math.sqrt(max(square_sum / count - (midpoint_sum / count) ** 2, 0.0)) + return _stats(percentile, lowest, percentile(1.0), histogram.total / count, spread) + + +def _finite(lower: float, upper: float) -> Tuple[float, float]: + """A bucket's bounds, with an open-ended bucket collapsed to its finite bound.""" + if math.isinf(upper): + return lower, lower + if math.isinf(lower): + return upper, upper + return lower, upper + + +def _selector(metric: str, labels: Dict[str, str]) -> str: + """A PromQL series selector: the quoted metric name and one exact-match matcher per label.""" + matchers = [_quote(metric)] + matchers += [f"{_quote(name)}={_quote(value)}" for name, value in labels.items()] + return "{" + ", ".join(matchers) + "}" + + +def _quote(text: str) -> str: + """A PromQL string literal: backslashes and double quotes escaped.""" + return '"' + text.replace("\\", "\\\\").replace('"', '\\"') + '"' + + +class _PromQLClient: + """Signed GET requests to CloudWatch's Prometheus-compatible query API in one region. + + Requests go through the CloudWatch client's HTTP session, so they use its proxy, CA bundle and + timeout settings, and a query is retried with backoff on throttling, server errors and dropped + connections. + """ + + def __init__(self, boto_session: boto3.session.Session, cloudwatch) -> None: + credentials = boto_session.get_credentials() + if credentials is None: + raise RuntimeError("No AWS credentials found. Configure credentials and try again.") + if not boto_session.region_name: + raise RuntimeError( + "The session has no region. Set the endpoint's region and try again." + ) + self._credentials = credentials + self._region = boto_session.region_name + # The CloudWatch client resolves the region's endpoint in its partition. + self._base_url = cloudwatch.meta.endpoint_url + self._http = cloudwatch._endpoint.http_session # pylint: disable=protected-access + + def histograms(self, expression: str, at: int) -> Dict[Tuple, Dict[str, Any]]: + """Evaluate an expression at one time; return {series labels: histogram sample}.""" + data = self._get("/api/v1/query", {"query": expression, "time": str(at)}) + return { + tuple(sorted(series.get("metric", {}).items())): series["histogram"][1] + for series in data.get("result", []) + if "histogram" in series + } + + def _get(self, path: str, params: Dict[str, str]) -> Dict[str, Any]: + query = urllib.parse.urlencode(params, quote_via=urllib.parse.quote) + url = f"{self._base_url}{path}?{query}" + attempt = 1 + while True: + request = AWSRequest(method="GET", url=url) + SigV4Auth( + self._credentials.get_frozen_credentials(), "monitoring", self._region + ).add_auth(request) + try: + response = self._http.send(request.prepare()) + except BotoCoreError as error: + if attempt == _QUERY_ATTEMPTS: + raise RuntimeError( + f"Couldn't reach CloudWatch's query API in {self._region}: {error}" + ) from None + else: + if response.status_code not in _RETRYABLE_STATUS or attempt == _QUERY_ATTEMPTS: + return _query_data(response) + time.sleep(_RETRY_BASE_SECONDS * 2 ** (attempt - 1)) + attempt += 1 + + +def _query_data(response) -> Dict[str, Any]: + """The ``data`` of a query API response, or a RuntimeError saying why there's none.""" + detail = response.content.decode("utf-8", "replace")[:300] + if response.status_code == 403: + raise RuntimeError( + "CloudWatch denied the query (HTTP 403). Reading token metrics needs " + f"cloudwatch:GetMetricData and cloudwatch:ListMetrics. {detail}" + ) + if response.status_code != 200: + raise RuntimeError(f"CloudWatch's query API returned HTTP {response.status_code}: {detail}") + try: + body = json.loads(response.content) + except ValueError: + raise RuntimeError( + f"CloudWatch's query API returned a body that isn't JSON: {detail}" + ) from None + if body.get("status") != "success": + raise RuntimeError(f"The CloudWatch query failed: {body.get('error', detail)}") + return body.get("data", {}) diff --git a/sagemaker-serve/src/sagemaker/serve/ai_inference_recommender/workload.py b/sagemaker-serve/src/sagemaker/serve/ai_inference_recommender/workload.py index c3ecfe3ce8..174cf6bd9d 100644 --- a/sagemaker-serve/src/sagemaker/serve/ai_inference_recommender/workload.py +++ b/sagemaker-serve/src/sagemaker/serve/ai_inference_recommender/workload.py @@ -15,15 +15,21 @@ from __future__ import absolute_import import json +import logging +import math import os import re import uuid +from datetime import datetime, timedelta from typing import Any, Dict, List, Optional, Union from urllib.parse import urlparse from pydantic import BaseModel, ConfigDict, Field, field_validator from sagemaker.serve.ai_inference_recommender.secrets import Secret +from sagemaker.serve.ai_inference_recommender.traffic import describe_endpoint_traffic + +logger = logging.getLogger(__name__) # Default input-data channel names; the channel is mounted at # {_CONTAINER_INPUT_DATA_DIR}/{channel_name}/ inside the benchmark container. @@ -207,6 +213,115 @@ def from_dataset( ], ) + @classmethod + def from_endpoint( + cls, + endpoint_name: str, + *, + inference_component_name: Optional[str] = None, + variant_name: Optional[str] = None, + lookback: Optional[timedelta] = None, + start_time: Optional[datetime] = None, + end_time: Optional[datetime] = None, + concurrency: Optional[int] = None, + use_data_capture: bool = False, + sagemaker_session: Optional[Any] = None, + **params: Any, + ) -> "Workload": + """Build a workload from the traffic one model on an endpoint served. + + Describes the endpoint with :func:`describe_endpoint_traffic`. The workload uses synthetic + prompts, as :meth:`Workload.synthetic` does, with the mean and standard deviation of the + prompt and response lengths of the requests the model completed in the window; those need + detailed observability on the endpoint and a vLLM serving container. With + ``use_data_capture``, it replays the requests SageMaker Data Capture recorded instead, as + :meth:`Workload.from_dataset` does. Either way, the concurrency is the window's busiest + hour for the whole deployment. + + Args: + endpoint_name: The SageMaker endpoint name. + inference_component_name: The inference component to read, for an endpoint that has + them. Found for you when the endpoint has only one. + variant_name: The production variant to read, for a model deployed on a variant. + Found for you when the endpoint has only one. + lookback: How far back from ``end_time`` to read, up to seven days (the default). + start_time: The window's start, a timezone-aware ``datetime``, instead of + ``lookback``. + end_time: The window's end, a timezone-aware ``datetime``. Defaults to now. + concurrency: Number of in-flight requests. Defaults to the busiest hour's. + use_data_capture: Replay the requests SageMaker Data Capture recorded for the variant, + when the endpoint captures them, instead of using synthetic prompts. Off by + default: the replay is every request under the variant's capture prefix, from all + of its inference components and all the history Data Capture kept, not just the + window, and the captures can hold your endpoint users' data. If they're encrypted + with a KMS key, the job's role needs ``kms:Decrypt`` on it. + sagemaker_session: Optional session used to read the endpoint — a + ``boto3.session.Session`` or a sagemaker ``Session`` wrapping one. + **params: Passed on to :meth:`Workload.from_dataset` or :meth:`Workload.synthetic`, + for example ``tokenizer`` (required for synthetic prompts), ``hf_token``, + ``request_count`` and ``streaming``. + + Returns: + A ``Workload`` that replays captured requests or uses synthetic prompts. + + Raises: + ValueError: If the endpoint served no traffic in the window and ``concurrency`` + isn't given, or the workload needs synthetic prompts and the endpoint published + no token metrics in the window. + """ + traffic = describe_endpoint_traffic( + endpoint_name, + inference_component_name=inference_component_name, + variant_name=variant_name, + lookback=lookback, + start_time=start_time, + end_time=end_time, + sagemaker_session=sagemaker_session, + ) + window = f"from {traffic.start_time} to {traffic.end_time}" + if concurrency is None: + if traffic.concurrency is None: + raise ValueError( + f"Endpoint {endpoint_name!r} served no requests {window}; pass concurrency." + ) + concurrency = math.ceil(traffic.concurrency.max) + if use_data_capture and traffic.data_capture_s3_uri: + logger.warning( + "The workload replays every request SageMaker Data Capture recorded under %s. They " + "can hold your endpoint users' data.", + traffic.data_capture_s3_uri, + ) + return cls.from_dataset( + traffic.data_capture_s3_uri, + custom_dataset_type="sagemaker-datacapture", + concurrency=concurrency, + **params, + ) + if use_data_capture: + logger.warning( + "Endpoint %s doesn't capture its requests, so the workload uses synthetic prompts.", + endpoint_name, + ) + if traffic.input_tokens is None or traffic.output_tokens is None: + if traffic.token_metrics_error: + reason = f"they couldn't be read ({traffic.token_metrics_error})" + elif traffic.detailed_observability is False: + reason = "its endpoint configuration turns detailed observability off" + else: + reason = "they need detailed observability on and a vLLM serving container" + raise ValueError( + f"Endpoint {endpoint_name!r} published no token metrics {window}: {reason}. " + "Use Workload.synthetic(...) with your own token sizes instead." + ) + return cls.synthetic( + concurrency=concurrency, + prompt_input_tokens_mean=round(traffic.input_tokens.mean), + prompt_input_tokens_stddev=traffic.input_tokens.stddev, + output_tokens_mean=round(traffic.output_tokens.mean), + output_tokens_stddev=traffic.output_tokens.stddev, + **params, + ) + @classmethod def template( cls, diff --git a/sagemaker-serve/tests/integ/test_ai_inference_recommender_enhancements_integration.py b/sagemaker-serve/tests/integ/test_ai_inference_recommender_enhancements_integration.py index 2385581983..38922279cb 100644 --- a/sagemaker-serve/tests/integ/test_ai_inference_recommender_enhancements_integration.py +++ b/sagemaker-serve/tests/integ/test_ai_inference_recommender_enhancements_integration.py @@ -20,8 +20,10 @@ import logging import time import uuid +from datetime import datetime, timedelta, timezone import pytest +from botocore.exceptions import ClientError from sagemaker.core.helper.session_helper import Session, get_execution_role from sagemaker.core.jumpstart.configs import JumpStartConfig @@ -36,10 +38,12 @@ from sagemaker.serve.ai_inference_recommender import ( Workload, compare_benchmarks, + describe_endpoint_traffic, list_benchmarks, list_recommendations, start_benchmark, ) +from sagemaker.serve.ai_inference_recommender import traffic as traffic_module from sagemaker.serve.model_builder import ModelBuilder from sagemaker.train.configs import Compute @@ -95,6 +99,39 @@ def test_list_benchmarks_and_recommendations_plumbing(): logger.info("Non-matching filters correctly returned empty lists.") +def test_describe_endpoint_traffic_plumbing(): + """The token metrics query is signed and parsed by CloudWatch's query API, and a missing + endpoint raises SageMaker's error. + + No GPU and no resources: the query names an endpoint that doesn't exist, so the API answers + with no series, which still exercises the selector quoting, the label_replace/or expression, + SigV4 signing for the ``monitoring`` service and the response parsing. + """ + boto_session = Session().boto_session + missing = f"no-such-endpoint-{uuid.uuid4().hex}" + end = datetime.now(timezone.utc).replace(microsecond=0) + client = traffic_module._PromQLClient(boto_session, boto_session.client("cloudwatch")) + target = traffic_module._Target(missing, None, "AllTraffic", 1) + + snapshots = traffic_module._snapshots(client, target, end - timedelta(days=1), end) + + assert snapshots == [{}, {}] + + with pytest.raises(ClientError) as raised: + describe_endpoint_traffic(missing, sagemaker_session=boto_session) + assert raised.value.response["Error"]["Code"] == "ValidationException" + + +def _wait_for_traffic(endpoint_name, timeout_seconds=600): + """describe_endpoint_traffic over the last 3 hours, once its concurrency has arrived.""" + deadline = time.time() + timeout_seconds + while True: + traffic = describe_endpoint_traffic(endpoint_name, lookback=timedelta(hours=3)) + if traffic.concurrency is not None or time.time() > deadline: + return traffic + time.sleep(30) + + @pytest.mark.gpu_intensive @pytest.mark.xfail_on_insufficient_capacity def test_recommendation_deploy_best_and_compare_e2e(): @@ -119,6 +156,8 @@ def test_recommendation_deploy_best_and_compare_e2e(): bench_b_name = f"air-enh-bench-b-{unique_id}" bench_a_wl = f"air-enh-bench-a-wl-{unique_id}" bench_b_wl = f"air-enh-bench-b-wl-{unique_id}" + bench_c_name = f"air-enh-bench-c-{unique_id}" + bench_c_wl = f"air-enh-bench-c-wl-{unique_id}" source_model = None endpoint = None @@ -222,6 +261,36 @@ def test_recommendation_deploy_best_and_compare_e2e(): f"Benchmark {job.get_name()} did not complete: " f"{job.ai_benchmark_job_status}" ) + # (4b) describe_endpoint_traffic sees the two benchmarks: their concurrency always, and + # their requests when the container publishes vLLM's token histograms. A workload built + # from that traffic is accepted by a benchmark. + traffic = _wait_for_traffic(endpoint_name) + logger.info("describe_endpoint_traffic: %s", traffic) + assert traffic.concurrency is not None and traffic.concurrency.max >= 1 + assert traffic.token_metrics_error is None + if traffic.framework == "vllm": + assert traffic.requests > 0 + assert traffic.input_tokens.p50 > 0 and traffic.output_tokens.p50 > 0 + sized = Workload.from_endpoint( + endpoint_name, + lookback=timedelta(hours=3), + tokenizer=WORKLOAD_TOKENIZER, + request_count=10, + ) + assert sized.parameters["concurrency"] >= 1 + bench_c = start_benchmark( + endpoint=endpoint_name, + workload=sized, + role=role, + name=bench_c_name, + workload_config_name=bench_c_wl, + wait=True, + ) + assert bench_c.ai_benchmark_job_status == "Completed", ( + f"Benchmark {bench_c.get_name()} did not complete: " + f"{bench_c.ai_benchmark_job_status}" + ) + # (5) list_benchmarks(endpoint=...) filters by this run's (unique) # endpoint. Hard-assert the filter contract: every returned job targets # this endpoint. Finding both specific jobs is best-effort (List is @@ -269,7 +338,7 @@ def test_recommendation_deploy_best_and_compare_e2e(): ) if source_model is not None: _delete_quietly(lambda: source_model, f"Model {src_model_name}") - for job_name in (bench_a_name, bench_b_name): + for job_name in (bench_a_name, bench_b_name, bench_c_name): _delete_quietly( lambda n=job_name: AIBenchmarkJob.get(ai_benchmark_job_name=n), f"AIBenchmarkJob {job_name}", @@ -278,7 +347,7 @@ def test_recommendation_deploy_best_and_compare_e2e(): lambda: AIRecommendationJob.get(ai_recommendation_job_name=rec_job_name), f"AIRecommendationJob {rec_job_name}", ) - for wl in (rec_wl_name, bench_a_wl, bench_b_wl): + for wl in (rec_wl_name, bench_a_wl, bench_b_wl, bench_c_wl): _delete_quietly( lambda n=wl: AIWorkloadConfig.get(ai_workload_config_name=n), f"AIWorkloadConfig {wl}", diff --git a/sagemaker-serve/tests/unit/test_ai_inference_recommender/test_traffic.py b/sagemaker-serve/tests/unit/test_ai_inference_recommender/test_traffic.py new file mode 100644 index 0000000000..73a69ed3b7 --- /dev/null +++ b/sagemaker-serve/tests/unit/test_ai_inference_recommender/test_traffic.py @@ -0,0 +1,915 @@ +# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"). You +# may not use this file except in compliance with the License. A copy of +# the License is located at +# +# http://aws.amazon.com/apache2.0/ +# +# or in the "license" file accompanying this file. This file is +# distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF +# ANY KIND, either express or implied. See the License for the specific +# language governing permissions and limitations under the License. +"""Unit tests for describe_endpoint_traffic and Workload.from_endpoint.""" + +from __future__ import absolute_import + +import json +import math +import urllib.parse +from datetime import datetime, timedelta, timezone +from unittest.mock import MagicMock, patch + +import boto3 +import pytest +from botocore.exceptions import ClientError, EndpointConnectionError + +from sagemaker.serve.ai_inference_recommender import ( + EndpointTraffic, + TrafficStats, + Workload, + describe_endpoint_traffic, +) +from sagemaker.serve.ai_inference_recommender import traffic as traffic_module + +DAY = 86_400 +NOW = 1_800_000_000 +NOW_UTC = datetime.fromtimestamp(NOW, tz=timezone.utc) +ENDPOINT = "my-endpoint" +COMPONENT = "my-ic" + + +@pytest.fixture(autouse=True) +def _offline_telemetry(): + """Keep the telemetry decorator from sending anything.""" + with patch("sagemaker.core.telemetry.telemetry_logging._send_telemetry_request"), patch( + "sagemaker.core.telemetry.telemetry_logging._get_default_sagemaker_session", + return_value=None, + ): + yield + + +def _histogram(values): + """A cumulative snapshot of these observations, one bucket of width 10 per distinct value.""" + buckets = {} + for value in values: + buckets[value] = buckets.get(value, 0) + 1 + return { + "count": str(len(values)), + "sum": str(sum(values)), + "buckets": [[0, str(v - 5), str(v + 5), str(n)] for v, n in sorted(buckets.items())], + } + + +def _both(container, values, framework="vllm"): + """The same snapshot for both sides, outputs a tenth of the inputs.""" + return { + (f"{framework}/input", container): _histogram(values), + (f"{framework}/output", container): _histogram([v / 10 for v in values]), + } + + +def _response(status_code=200, body=None): + response = MagicMock() + response.status_code = status_code + response.content = json.dumps(body if body is not None else {}).encode() + return response + + +def _params(prepared): + return dict(urllib.parse.parse_qsl(urllib.parse.urlsplit(prepared.url).query)) + + +class _FakeQueryAPI: + """Answers instant queries from {days_ago: {(stat, container): snapshot}}; records requests.""" + + def __init__(self, days=None): + self.days = days or {} + self.requests = [] + + def send(self, prepared): + self.requests.append(prepared) + at = int(_params(prepared)["time"]) + result = [ + { + "metric": {"stat": stat, "aws.sagemaker.container.id": container}, + "histogram": [at, snapshot], + } + for (stat, container), snapshot in self.days.get((NOW - at) // DAY, {}).items() + ] + return _response( + body={"status": "success", "data": {"resultType": "vector", "result": result}} + ) + + +def _sagemaker_client( + components=(COMPONENT,), + copies=2, + variants=(("AllTraffic", 3),), + observability=True, + capture=None, +): + client = MagicMock() + client.describe_endpoint.return_value = { + "EndpointName": ENDPOINT, + "EndpointConfigName": "my-config", + "ProductionVariants": [ + {"VariantName": name, "CurrentInstanceCount": count} for name, count in variants + ], + **({"DataCaptureConfig": capture} if capture is not None else {}), + } + client.describe_endpoint_config.return_value = ( + {"MetricsConfig": {"EnableDetailedObservability": observability}} + if observability is not None + else {} + ) + client.list_inference_components.return_value = { + "InferenceComponents": [{"InferenceComponentName": name} for name in components] + } + client.describe_inference_component.return_value = { + "EndpointName": ENDPOINT, + "VariantName": "AllTraffic", + "RuntimeConfig": {"CurrentCopyCount": copies}, + } + return client + + +def _cloudwatch_client(peaks=(1.0, 2.5), api=None): + client = MagicMock() + client.meta.endpoint_url = "https://monitoring.us-east-1.amazonaws.com" + client._endpoint.http_session = api if api is not None else _FakeQueryAPI() + client.get_metric_statistics.return_value = { + "Datapoints": [{"Maximum": peak} for peak in peaks] + } + return client + + +def _session(sagemaker_client, cloudwatch_client): + session = boto3.session.Session( + aws_access_key_id="AKIDEXAMPLE", + aws_secret_access_key="SECRET", + aws_session_token="TOKEN", + region_name="us-east-1", + ) + clients = {"sagemaker": sagemaker_client, "cloudwatch": cloudwatch_client} + session.client = MagicMock(side_effect=lambda name, **_: clients[name]) + return session + + +def _describe(api, sagemaker_client=None, cloudwatch_client=None, **kwargs): + cloudwatch = cloudwatch_client or _cloudwatch_client() + cloudwatch._endpoint.http_session = api + session = _session(sagemaker_client or _sagemaker_client(), cloudwatch) + with patch.object(traffic_module.time, "time", return_value=NOW), patch.object( + traffic_module.time, "sleep" + ): + return describe_endpoint_traffic(ENDPOINT, sagemaker_session=session, **kwargs) + + +def _times(api): + return [int(_params(prepared)["time"]) for prepared in api.requests] + + +class TestTokenSizes: + def test_the_window_is_each_series_newest_snapshot_less_its_snapshot_at_the_start(self): + # Container a served 100 and 300 before the window and 200, 200 and 800 in it; container + # b started in the window and served 500. + api = _FakeQueryAPI( + { + 0: {**_both("a", [100, 300, 200, 200, 800]), **_both("b", [500])}, + 3: _both("a", [100, 300, 200]), + 7: _both("a", [100, 300]), + } + ) + + traffic = _describe(api) + + window = [200, 200, 800, 500] + mean = sum(window) / len(window) + stddev = round(math.sqrt(sum((x - mean) ** 2 for x in window) / len(window)), 1) + assert traffic.framework == "vllm" + assert traffic.requests == 4 + tokens = traffic.input_tokens + assert (tokens.min, tokens.p50, tokens.max) == (195.0, 205.0, 805.0) + assert (tokens.mean, tokens.stddev) == (mean, stddev) + assert traffic.output_tokens.mean == mean / 10 + assert traffic.output_tokens.stddev == round(stddev / 10, 1) + + def test_percentiles_interpolate_within_buckets(self): + # 100 requests, one in each bucket from [0, 10] to [990, 1000]. + api = _FakeQueryAPI({0: _both("a", list(range(5, 1000, 10)))}) + + tokens = _describe(api).input_tokens + + assert tokens == TrafficStats( + min=0.0, + p25=250.0, + p50=500.0, + p70=700.0, + p80=800.0, + p90=900.0, + p95=950.0, + p99=990.0, + max=1000.0, + mean=500.0, + stddev=tokens.stddev, + ) + assert tokens.stddev == pytest.approx(288.7, abs=0.1) + + def test_an_open_ended_bucket_counts_at_its_finite_bound(self): + snapshot = { + "count": "4", + "sum": "10", + "buckets": [[0, "-Inf", "0", "1"], [0, "0", "4", "2"], [0, "4", "+Inf", "1"]], + } + api = _FakeQueryAPI({0: {("vllm/input", "a"): snapshot, ("vllm/output", "a"): snapshot}}) + + tokens = _describe(api).input_tokens + + assert (tokens.min, tokens.p25, tokens.p50, tokens.p70, tokens.p80, tokens.max) == ( + 0.0, + 0.0, + 2.0, + 3.6, + 4.0, + 4.0, + ) + assert (tokens.mean, tokens.stddev) == (2.5, round(math.sqrt(2), 1)) + + def test_buckets_resized_between_snapshots_keep_the_percentiles_in_order(self): + # The start snapshot's buckets are finer than the end's: [0, 5] and [5, 10] against [0, 10]. + start = {"count": "2", "sum": "8", "buckets": [[0, "0", "5", "1"], [0, "5", "10", "1"]]} + end = {"count": "6", "sum": "48", "buckets": [[0, "0", "10", "4"], [0, "10", "20", "2"]]} + api = _FakeQueryAPI( + { + 0: {("vllm/input", "a"): end, ("vllm/output", "a"): end}, + 7: {("vllm/input", "a"): start, ("vllm/output", "a"): start}, + } + ) + + traffic = _describe(api) + + tokens = traffic.input_tokens + assert traffic.requests == 4 + assert tokens.mean == 10.0 + percentiles = [tokens.min, tokens.p25, tokens.p50, tokens.p90, tokens.p99, tokens.max] + assert percentiles == sorted(percentiles) + assert tokens.max == 20.0 + + def test_buckets_short_of_the_count_put_the_top_percentiles_at_the_highest_bound(self): + # The snapshot counts 4 requests but its buckets hold 2. + snapshot = { + "count": "4", + "sum": "40", + "buckets": [[0, "0", "10", "1"], [0, "10", "20", "1"]], + } + api = _FakeQueryAPI({0: {("vllm/input", "a"): snapshot, ("vllm/output", "a"): snapshot}}) + + tokens = _describe(api).input_tokens + + assert (tokens.p25, tokens.p50, tokens.p99, tokens.max) == (10.0, 20.0, 20.0, 20.0) + + def test_reads_each_day_back_from_the_end_and_the_start_with_signed_instant_queries(self): + api = _FakeQueryAPI({0: _both("a", [10])}) + + _describe(api, lookback=timedelta(days=3)) + + assert _times(api) == [NOW, NOW - DAY, NOW - 2 * DAY, NOW - 3 * DAY] + for prepared in api.requests: + url = urllib.parse.urlsplit(prepared.url) + query = _params(prepared)["query"] + assert url.netloc == "monitoring.us-east-1.amazonaws.com" + assert url.path == "/api/v1/query" + assert '"aws.sagemaker.endpoint.name"="my-endpoint"' in query + assert '"aws.sagemaker.inference_component.name"="my-ic"' in query + assert 'last_over_time({"vllm:request_prompt_tokens"' in query + assert 'last_over_time({"vllm:request_generation_tokens"' in query + assert '"stat", "vllm/input"' in query and '"stat", "vllm/output"' in query + assert "increase(" not in query + assert "/us-east-1/monitoring/aws4_request" in prepared.headers["Authorization"] + + def test_a_model_on_a_variant_is_selected_by_its_variant_label(self): + api = _FakeQueryAPI({0: _both("a", [10])}) + + traffic = _describe( + api, sagemaker_client=_sagemaker_client(components=()), variant_name="AllTraffic" + ) + + query = _params(api.requests[0])["query"] + assert traffic.inference_component_name is None + assert traffic.variant_name == "AllTraffic" + assert '"aws.sagemaker.variant.name"="AllTraffic"' in query + assert "inference_component" not in query + + def test_a_restarted_series_counts_from_zero(self): + # The container restarted after the window began: its count went from 5 to 2. + api = _FakeQueryAPI({0: _both("a", [40, 60]), 7: _both("a", [1, 1, 1, 1, 1])}) + + traffic = _describe(api) + + assert traffic.requests == 2 + assert traffic.input_tokens.mean == 50.0 + + def test_a_restart_between_days_adds_the_requests_from_before_it(self): + # The series had served 2 requests before the window. In the window it served 3 more, + # restarted, and served 2 after the restart. + api = _FakeQueryAPI( + { + 7: _both("a", [10, 10]), + 3: _both("a", [10, 10, 30, 30, 30]), + 0: _both("a", [50, 50]), + } + ) + + traffic = _describe(api) + + assert traffic.requests == 5 + assert traffic.input_tokens.mean == (3 * 30 + 2 * 50) / 5 + assert traffic.input_tokens.min == 25.0 + + def test_a_bucket_that_goes_down_is_a_restart(self): + # The count and sum grew, but the bucket around 10 went from 3 to 1: the container + # restarted, so the later snapshot counts from zero. + api = _FakeQueryAPI({7: _both("a", [10, 10, 10]), 0: _both("a", [10, 90, 90, 90])}) + + traffic = _describe(api) + + assert traffic.requests == 4 + assert traffic.input_tokens.mean == 70.0 + + def test_a_sum_that_goes_down_is_a_restart(self): + api = _FakeQueryAPI({7: _both("a", [50, 50]), 0: _both("a", [10, 10, 10])}) + + traffic = _describe(api) + + assert traffic.requests == 3 + assert traffic.input_tokens.mean == 10.0 + + def test_a_series_seen_only_before_the_window_is_left_out(self): + api = _FakeQueryAPI({0: _both("a", [10]), 7: {**_both("a", []), **_both("old", [999])}}) + + traffic = _describe(api) + + assert traffic.requests == 1 + assert traffic.input_tokens.mean == 10.0 + + def test_the_newest_snapshot_in_the_window_is_the_one_read(self): + # The series stopped two days ago; the snapshot from then is its newest. + api = _FakeQueryAPI({2: _both("a", [10, 30]), 5: _both("a", [10])}) + + traffic = _describe(api) + + assert traffic.requests == 2 + assert traffic.input_tokens.mean == 20.0 + + def test_no_histogram_in_the_window_leaves_the_token_sizes_out(self): + traffic = _describe(_FakeQueryAPI()) + + assert (traffic.framework, traffic.requests) == (None, 0) + assert traffic.input_tokens is None and traffic.output_tokens is None + + def test_detailed_observability_off_skips_the_token_metrics(self): + api = _FakeQueryAPI({0: _both("a", [10])}) + + traffic = _describe(api, sagemaker_client=_sagemaker_client(observability=False)) + + assert api.requests == [] + assert traffic.detailed_observability is False + assert traffic.input_tokens is None and traffic.framework is None + assert traffic.concurrency is not None + + def test_an_unset_detailed_observability_still_reads_the_token_metrics(self): + api = _FakeQueryAPI({0: _both("a", [10])}) + + traffic = _describe(api, sagemaker_client=_sagemaker_client(observability=None)) + + assert traffic.detailed_observability is None + assert traffic.input_tokens.mean == 10.0 + + def test_a_deleted_endpoint_configuration_still_reads_the_token_metrics(self): + sagemaker = _sagemaker_client() + sagemaker.describe_endpoint_config.side_effect = ClientError( + {"Error": {"Code": "ValidationException", "Message": "Could not find"}}, + "DescribeEndpointConfig", + ) + api = _FakeQueryAPI({0: _both("a", [10])}) + + traffic = _describe(api, sagemaker_client=sagemaker) + + assert traffic.detailed_observability is None + assert traffic.input_tokens.mean == 10.0 + + def test_another_frameworks_histograms_are_read_when_listed(self): + frameworks = ( + traffic_module._TokenMetrics("vllm", "vllm:prompt", "vllm:generation"), + traffic_module._TokenMetrics("other", "other:prompt", "other:generation"), + ) + api = _FakeQueryAPI({0: _both("a", [10], framework="other")}) + + with patch.object(traffic_module, "_TOKEN_METRICS", frameworks): + traffic = _describe(api) + + query = _params(api.requests[0])["query"] + assert '{"other:prompt"' in query and '"stat", "other/output"' in query + assert traffic.framework == "other" + assert traffic.input_tokens.mean == 10.0 + + def test_label_values_are_escaped(self): + assert traffic_module._selector("m", {"l": 'a"b\\c'}) == '{"m", "l"="a\\"b\\\\c"}' + + +class TestWindow: + def test_the_default_is_the_seven_days_up_to_now_in_utc(self): + cloudwatch = _cloudwatch_client() + + traffic = _describe(_FakeQueryAPI(), cloudwatch_client=cloudwatch) + + assert (traffic.start_time, traffic.end_time) == (NOW_UTC - timedelta(days=7), NOW_UTC) + assert traffic.end_time.tzinfo == timezone.utc + kwargs = cloudwatch.get_metric_statistics.call_args.kwargs + assert (kwargs["StartTime"], kwargs["EndTime"]) == (traffic.start_time, traffic.end_time) + + def test_a_lookback_ends_at_the_end_time(self): + end = NOW_UTC - timedelta(days=10) + api = _FakeQueryAPI() + + traffic = _describe(api, lookback=timedelta(days=2), end_time=end) + + assert (traffic.start_time, traffic.end_time) == (end - timedelta(days=2), end) + assert _times(api) == [NOW - 10 * DAY, NOW - 11 * DAY, NOW - 12 * DAY] + + def test_an_absolute_window_in_another_time_zone_is_read_in_utc(self): + pacific = timezone(timedelta(hours=-8)) + start = datetime(2027, 1, 10, 0, 0, tzinfo=pacific) + end = datetime(2027, 1, 12, 12, 0, tzinfo=pacific) + api = _FakeQueryAPI() + + traffic = _describe(api, start_time=start, end_time=end) + + assert traffic.start_time == datetime(2027, 1, 10, 8, 0, tzinfo=timezone.utc) + assert traffic.end_time == datetime(2027, 1, 12, 20, 0, tzinfo=timezone.utc) + first, last = int(end.timestamp()), int(start.timestamp()) + assert _times(api) == [first, first - DAY, first - 2 * DAY, last] + + def test_a_window_of_part_of_a_day_reads_the_end_and_the_start(self): + api = _FakeQueryAPI() + + _describe(api, lookback=timedelta(hours=6)) + + assert _times(api) == [NOW, NOW - 6 * 3_600] + + @pytest.mark.parametrize( + "kwargs, message", + [ + ({"start_time": datetime(2027, 1, 10)}, "timezone-aware"), + ({"end_time": datetime(2027, 1, 10)}, "timezone-aware"), + ({"end_time": "2027-01-10"}, "timezone-aware"), + ({"lookback": timedelta(0)}, "positive timedelta"), + ({"lookback": 3}, "positive timedelta"), + ( + {"lookback": timedelta(days=1), "start_time": NOW_UTC - timedelta(days=1)}, + "not both", + ), + ({"end_time": NOW_UTC + timedelta(minutes=1)}, "later than now"), + ({"start_time": NOW_UTC, "end_time": NOW_UTC}, "before end_time"), + ({"lookback": timedelta(days=7, seconds=1)}, "at most 7 days"), + ({"start_time": NOW_UTC - timedelta(days=8)}, "at most 7 days"), + ], + ) + def test_bad_windows_are_refused(self, kwargs, message): + with pytest.raises(ValueError, match=message): + _describe(_FakeQueryAPI(), **kwargs) + + +class TestConcurrency: + def test_a_components_hourly_peaks_per_copy_are_multiplied_by_its_copies(self): + cloudwatch = _cloudwatch_client(peaks=(2.5, 1.0, 0.5, 2.0)) + + traffic = _describe(_FakeQueryAPI(), cloudwatch_client=cloudwatch) + + # Each hour's maximum per copy, times 2 copies: 1, 2, 4 and 5. + concurrency = traffic.concurrency + assert (concurrency.min, concurrency.p50, concurrency.max) == (1.0, 3.0, 5.0) + assert (concurrency.p25, concurrency.p90, concurrency.mean) == (1.8, 4.7, 3.0) + assert concurrency.stddev == round(math.sqrt(2.5), 1) + call = cloudwatch.get_metric_statistics.call_args.kwargs + assert call["MetricName"] == "ConcurrentRequestsPerCopy" + assert call["Dimensions"] == [{"Name": "InferenceComponentName", "Value": COMPONENT}] + assert (call["Namespace"], call["Period"], call["Statistics"]) == ( + "AWS/SageMaker", + 3600, + ["Maximum"], + ) + + def test_a_variants_hourly_peaks_per_model_are_multiplied_by_its_instances(self): + cloudwatch = _cloudwatch_client(peaks=(1.5,)) + + traffic = _describe( + _FakeQueryAPI(), + sagemaker_client=_sagemaker_client(components=()), + cloudwatch_client=cloudwatch, + ) + + assert traffic.concurrency.max == 4.5 + assert traffic.concurrency.p25 == 4.5 + call = cloudwatch.get_metric_statistics.call_args.kwargs + assert call["MetricName"] == "ConcurrentRequestsPerModel" + assert call["Dimensions"] == [ + {"Name": "EndpointName", "Value": ENDPOINT}, + {"Name": "VariantName", "Value": "AllTraffic"}, + ] + + def test_no_traffic_is_no_concurrency(self): + for peaks in ((), (0.0, 0.0)): + traffic = _describe(_FakeQueryAPI(), cloudwatch_client=_cloudwatch_client(peaks=peaks)) + + assert traffic.concurrency is None + + def test_a_component_with_no_copies_now_counts_one(self): + traffic = _describe( + _FakeQueryAPI(), + sagemaker_client=_sagemaker_client(copies=0), + cloudwatch_client=_cloudwatch_client(peaks=(3.0,)), + ) + + assert traffic.concurrency.max == 3.0 + + +class TestTarget: + def test_the_only_component_is_found(self): + sagemaker = _sagemaker_client() + + traffic = _describe(_FakeQueryAPI(), sagemaker_client=sagemaker) + + assert (traffic.inference_component_name, traffic.variant_name) == ( + COMPONENT, + "AllTraffic", + ) + sagemaker.list_inference_components.assert_any_call(EndpointNameEquals=ENDPOINT) + sagemaker.describe_endpoint_config.assert_called_once_with(EndpointConfigName="my-config") + + @staticmethod + def _with_adapter(sagemaker): + """The endpoint's components: base my-ic with 2 copies, and adapter my-adapter on it.""" + described = { + COMPONENT: { + "EndpointName": ENDPOINT, + "VariantName": "AllTraffic", + "RuntimeConfig": {"CurrentCopyCount": 2}, + }, + "my-adapter": { + "EndpointName": ENDPOINT, + "VariantName": "AllTraffic", + "Specification": {"BaseInferenceComponentName": COMPONENT}, + }, + } + sagemaker.list_inference_components.return_value = { + "InferenceComponents": [ + {"InferenceComponentName": COMPONENT}, + {"InferenceComponentName": "my-adapter"}, + ] + } + sagemaker.describe_inference_component.side_effect = lambda InferenceComponentName: ( + described[InferenceComponentName] + ) + return sagemaker + + def test_an_adapter_reads_its_base_component(self): + api = _FakeQueryAPI({0: _both("a", [10])}) + cloudwatch = _cloudwatch_client(peaks=(3.0,)) + + traffic = _describe( + api, + sagemaker_client=self._with_adapter(_sagemaker_client()), + cloudwatch_client=cloudwatch, + inference_component_name="my-adapter", + ) + + assert traffic.inference_component_name == COMPONENT + assert traffic.concurrency.max == 6.0 + assert cloudwatch.get_metric_statistics.call_args.kwargs["Dimensions"] == [ + {"Name": "InferenceComponentName", "Value": COMPONENT} + ] + assert ( + '"aws.sagemaker.inference_component.name"="my-ic"' in _params(api.requests[0])["query"] + ) + + def test_finding_the_component_skips_adapters(self): + traffic = _describe( + _FakeQueryAPI(), sagemaker_client=self._with_adapter(_sagemaker_client()) + ) + + assert traffic.inference_component_name == COMPONENT + + def test_several_components_need_a_choice(self): + with pytest.raises(ValueError, match="inference_component_name"): + _describe(_FakeQueryAPI(), sagemaker_client=_sagemaker_client(components=("a", "b"))) + + def test_several_variants_need_a_choice(self): + sagemaker = _sagemaker_client(components=(), variants=(("a", 1), ("b", 1))) + + with pytest.raises(ValueError, match="variant_name"): + _describe(_FakeQueryAPI(), sagemaker_client=sagemaker) + + def test_a_component_on_another_endpoint_is_refused(self): + sagemaker = _sagemaker_client() + sagemaker.describe_inference_component.return_value["EndpointName"] = "other" + + with pytest.raises(ValueError, match="is not on endpoint"): + _describe( + _FakeQueryAPI(), sagemaker_client=sagemaker, inference_component_name=COMPONENT + ) + + def test_an_unknown_variant_is_refused(self): + with pytest.raises(ValueError, match="no production variant 'nope'"): + _describe(_FakeQueryAPI(), variant_name="nope") + + def test_an_endpoint_without_variants_is_refused(self): + sagemaker = _sagemaker_client(components=(), variants=()) + + with pytest.raises(ValueError, match="no production variant"): + _describe(_FakeQueryAPI(), sagemaker_client=sagemaker) + + def test_components_are_listed_across_pages(self): + sagemaker = _sagemaker_client() + sagemaker.list_inference_components.side_effect = [ + {"InferenceComponents": [{"InferenceComponentName": "a"}], "NextToken": "t"}, + {"InferenceComponents": [{"InferenceComponentName": "b"}]}, + ] + + with pytest.raises(ValueError, match=r"\['a', 'b'\]"): + _describe(_FakeQueryAPI(), sagemaker_client=sagemaker) + assert sagemaker.list_inference_components.call_args.kwargs["NextToken"] == "t" + + def test_a_component_and_a_variant_together_are_refused(self): + with pytest.raises(ValueError, match="at most one"): + _describe(_FakeQueryAPI(), inference_component_name="a", variant_name="b") + + def test_an_endpoint_name_is_required(self): + with pytest.raises(ValueError, match="endpoint_name"): + describe_endpoint_traffic("") + + +class TestDataCapture: + def test_the_variants_capture_prefix_is_reported(self): + capture = {"EnableCapture": True, "DestinationS3Uri": "s3://bucket/capture/"} + + traffic = _describe(_FakeQueryAPI(), sagemaker_client=_sagemaker_client(capture=capture)) + + assert traffic.data_capture_s3_uri == "s3://bucket/capture/my-endpoint/AllTraffic/" + + @pytest.mark.parametrize( + "capture", + [ + None, + {"EnableCapture": False, "DestinationS3Uri": "s3://bucket/capture"}, + {"EnableCapture": True}, + ], + ) + def test_no_capture_is_no_prefix(self, capture): + traffic = _describe(_FakeQueryAPI(), sagemaker_client=_sagemaker_client(capture=capture)) + + assert traffic.data_capture_s3_uri is None + + +class TestQueryAPIErrors: + """Token metrics are read best effort: a failed read is recorded, and the rest is returned.""" + + def test_access_denied_names_the_permissions(self, caplog): + api = MagicMock() + api.send.return_value = _response(403, {"message": "not authorized"}) + + traffic = _describe(api) + + assert "cloudwatch:GetMetricData and cloudwatch:ListMetrics" in traffic.token_metrics_error + assert traffic.input_tokens is None and traffic.concurrency is not None + assert "Couldn't read the token metrics" in caplog.text + assert api.send.call_count == 1 + + def test_other_http_errors_and_failed_queries_are_recorded(self): + api = MagicMock() + api.send.return_value = _response(422, {"error": "range too long"}) + assert "HTTP 422" in _describe(api).token_metrics_error + + api.send.return_value = _response(200, {"status": "error", "error": "parse error"}) + assert "parse error" in _describe(api).token_metrics_error + + def test_a_body_that_isnt_json_is_recorded(self): + api = MagicMock() + response = _response(200) + response.content = b"proxy error" + api.send.return_value = response + + assert "isn't JSON" in _describe(api).token_metrics_error + + def test_throttling_and_server_errors_are_retried(self): + fake = _FakeQueryAPI({0: _both("a", [10])}) + failures = iter([_response(429), _response(503)]) + api = MagicMock() + api.send.side_effect = lambda prepared: next(failures, None) or fake.send(prepared) + + traffic = _describe(api) + + assert traffic.token_metrics_error is None + assert traffic.input_tokens.mean == 10.0 + + def test_a_query_gives_up_after_its_attempts(self): + api = MagicMock() + api.send.return_value = _response(503) + + traffic = _describe(api) + + assert "HTTP 503" in traffic.token_metrics_error + assert api.send.call_count == traffic_module._QUERY_ATTEMPTS + + def test_an_unreachable_api_is_retried_then_recorded(self): + api = MagicMock() + api.send.side_effect = EndpointConnectionError(endpoint_url="https://monitoring.example") + + traffic = _describe(api) + + assert "Couldn't reach" in traffic.token_metrics_error + assert api.send.call_count == traffic_module._QUERY_ATTEMPTS + + def test_queries_go_through_the_cloudwatch_clients_http_session(self): + api = _FakeQueryAPI({0: _both("a", [10])}) + + _describe(api) + + assert api.requests and all( + "/us-east-1/monitoring/aws4_request" in p.headers["Authorization"] for p in api.requests + ) + + def test_no_credentials_raise(self): + session = MagicMock(spec=boto3.session.Session) + session.get_credentials.return_value = None + session.region_name = "us-east-1" + + with pytest.raises(RuntimeError, match="No AWS credentials"): + traffic_module._PromQLClient(session, _cloudwatch_client()) + + def test_no_region_raises(self): + session = MagicMock(spec=boto3.session.Session) + session.region_name = None + + with pytest.raises(RuntimeError, match="no region"): + traffic_module._PromQLClient(session, _cloudwatch_client()) + + +class TestSessionHandling: + def test_a_sagemaker_session_is_unwrapped(self): + boto_session = _session(_sagemaker_client(), _cloudwatch_client()) + sm_session = MagicMock() + sm_session.boto_session = boto_session + sm_session.sagemaker_config = {} + with patch.object(traffic_module.time, "time", return_value=NOW): + traffic = describe_endpoint_traffic(ENDPOINT, sagemaker_session=sm_session) + + assert traffic.inference_component_name == COMPONENT + boto_session.client.assert_any_call("sagemaker") + + +def _stats(mean, stddev, peak=None): + """A TrafficStats with this mean, standard deviation and maximum.""" + peak = mean if peak is None else peak + return TrafficStats( + min=0.0, + p25=mean, + p50=mean, + p70=mean, + p80=mean, + p90=mean, + p95=mean, + p99=peak, + max=peak, + mean=mean, + stddev=stddev, + ) + + +class TestWorkloadFromEndpoint: + _DESCRIBE = "sagemaker.serve.ai_inference_recommender.workload.describe_endpoint_traffic" + + def _traffic(self, **overrides): + values = { + "endpoint_name": ENDPOINT, + "inference_component_name": COMPONENT, + "variant_name": "AllTraffic", + "start_time": NOW_UTC - timedelta(days=7), + "end_time": NOW_UTC, + "concurrency": _stats(3.0, 1.2, peak=7.5), + "detailed_observability": True, + "framework": "vllm", + "requests": 1234, + "input_tokens": _stats(812.4, 230.1), + "output_tokens": _stats(196.6, 75.3), + "token_metrics_error": None, + "data_capture_s3_uri": None, + } + values.update(overrides) + return EndpointTraffic(**values) + + def test_sizes_a_synthetic_workload_from_the_token_metrics(self): + session = MagicMock() + start, end = NOW_UTC - timedelta(days=3), NOW_UTC - timedelta(days=1) + with patch(self._DESCRIBE, return_value=self._traffic()) as describe: + workload = Workload.from_endpoint( + ENDPOINT, + inference_component_name=COMPONENT, + start_time=start, + end_time=end, + sagemaker_session=session, + tokenizer="org/model", + hf_token="arn:aws:secretsmanager:us-east-1:123456789012:secret:hf-AbCdEf", + request_count=200, + ) + + describe.assert_called_once_with( + ENDPOINT, + inference_component_name=COMPONENT, + variant_name=None, + lookback=None, + start_time=start, + end_time=end, + sagemaker_session=session, + ) + assert workload.parameters == { + "tokenizer": "org/model", + "concurrency": 8, + "request_count": 200, + "prompt_input_tokens_mean": 812, + "prompt_input_tokens_stddev": 230.1, + "output_tokens_mean": 197, + "output_tokens_stddev": 75.3, + "streaming": True, + } + assert workload.secrets == { + "hf_token": "arn:aws:secretsmanager:us-east-1:123456789012:secret:hf-AbCdEf" + } + assert workload.dataset_channels == [] + + _CAPTURES = "s3://bucket/capture/my-endpoint/AllTraffic/" + + def test_captured_requests_are_replayed_when_asked(self, caplog): + with patch(self._DESCRIBE, return_value=self._traffic(data_capture_s3_uri=self._CAPTURES)): + workload = Workload.from_endpoint(ENDPOINT, use_data_capture=True, request_count=50) + + assert [channel.s3_uri for channel in workload.dataset_channels] == [self._CAPTURES] + assert workload.parameters["custom_dataset_type"] == "sagemaker-datacapture" + assert (workload.parameters["concurrency"], workload.parameters["request_count"]) == (8, 50) + assert self._CAPTURES in caplog.text + + def test_captured_requests_arent_replayed_by_default(self): + with patch(self._DESCRIBE, return_value=self._traffic(data_capture_s3_uri=self._CAPTURES)): + workload = Workload.from_endpoint(ENDPOINT, tokenizer="org/model") + + assert workload.dataset_channels == [] + assert workload.parameters["prompt_input_tokens_mean"] == 812 + + def test_asking_for_a_replay_without_captures_uses_synthetic_prompts(self, caplog): + with patch(self._DESCRIBE, return_value=self._traffic()): + workload = Workload.from_endpoint( + ENDPOINT, use_data_capture=True, tokenizer="org/model" + ) + + assert workload.dataset_channels == [] + assert "doesn't capture its requests" in caplog.text + + def test_a_given_concurrency_wins(self): + with patch(self._DESCRIBE, return_value=self._traffic(concurrency=None)): + workload = Workload.from_endpoint(ENDPOINT, tokenizer="org/model", concurrency=4) + + assert workload.parameters["concurrency"] == 4 + + @pytest.mark.parametrize( + "observability, reason", + [(False, "turns detailed observability off"), (None, "need detailed observability on")], + ) + def test_no_token_metrics_is_refused(self, observability, reason): + traffic = self._traffic( + input_tokens=None, + output_tokens=None, + requests=0, + framework=None, + detailed_observability=observability, + ) + with patch(self._DESCRIBE, return_value=traffic): + with pytest.raises(ValueError, match=f"published no token metrics .*{reason}"): + Workload.from_endpoint(ENDPOINT, tokenizer="org/model") + + def test_a_failed_token_read_says_why(self): + traffic = self._traffic( + input_tokens=None, + output_tokens=None, + requests=0, + framework=None, + token_metrics_error="CloudWatch denied the query (HTTP 403).", + ) + with patch(self._DESCRIBE, return_value=traffic): + with pytest.raises(ValueError, match=r"couldn't be read \(CloudWatch denied"): + Workload.from_endpoint(ENDPOINT, tokenizer="org/model") + + def test_no_traffic_without_a_concurrency_is_refused(self): + with patch(self._DESCRIBE, return_value=self._traffic(concurrency=None)): + with pytest.raises(ValueError, match="pass concurrency"): + Workload.from_endpoint(ENDPOINT, tokenizer="org/model")