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")