From 151d62de11797fc6e577c43328262379e2b7f37d Mon Sep 17 00:00:00 2001 From: DWarez Date: Thu, 9 Jul 2026 17:18:52 +0200 Subject: [PATCH] add: tutorial for async batched inference using sagemaker Signed-off-by: DWarez --- .../sagemaker-notebook.ipynb | 772 ++++++++++++++++++ 1 file changed, 772 insertions(+) create mode 100644 docs/sagemaker/notebooks/sagemaker-sdk/async-inference-embedding-tei/sagemaker-notebook.ipynb diff --git a/docs/sagemaker/notebooks/sagemaker-sdk/async-inference-embedding-tei/sagemaker-notebook.ipynb b/docs/sagemaker/notebooks/sagemaker-sdk/async-inference-embedding-tei/sagemaker-notebook.ipynb new file mode 100644 index 000000000..831b0d67c --- /dev/null +++ b/docs/sagemaker/notebooks/sagemaker-sdk/async-inference-embedding-tei/sagemaker-notebook.ipynb @@ -0,0 +1,772 @@ +{ + "cells": [ + { + "cell_type": "markdown", + "id": "cell-01", + "metadata": {}, + "source": "# Batch-embed a corpus with SageMaker asynchronous inference\n\n## What asynchronous inference is\n\nMost inference is synchronous: a client opens a connection, sends a request,\nwaits while the model runs, and reads the response on the same connection.\nThat works well for interactive traffic, but it struggles when the work is\nslow, the payloads are large, or requests arrive in unpredictable bursts. The\ncaller has to hold a connection open for the whole run, and the endpoint has\nto be sized for peak load even while it sits idle the rest of the time.\n\n**Asynchronous inference** decouples the request from the result. The client\nuploads the payload to object storage (S3) and sends the endpoint a *pointer*\nto it. The endpoint puts the request on an internal queue, processes it when\ncapacity is available, and writes the result back to S3. The client picks the\nresult up later by polling S3, or by reacting to a success/error\nnotification.\n\n![SageMaker asynchronous inference architecture](https://docs.aws.amazon.com/images/sagemaker/latest/dg/images/async-architecture.png)\n\nThat indirection is what makes the pattern useful:\n\n- **Large payloads, long runtimes.** Inputs are read from S3, not from a\n size-capped HTTP body, and there is no client connection to time out.\n- **Bursty, queue-shaped traffic.** Requests accumulate in the queue and\n drain at the endpoint's own pace, instead of every spike forcing an\n immediate scale-out.\n- **Scale to zero.** When the queue is empty the endpoint can run zero\n instances and cost nothing, then wake up when new requests land.\n\nThe cost is latency: you trade an immediate answer for throughput and\nelasticity. So asynchronous inference is for offline and background work, not\nthe interactive request path.\n\n## The use case we picked\n\nTo keep things concrete, this tutorial uses asynchronous inference to\n**embed a text corpus for semantic search**. Embedding a corpus is a textbook\nfit: it is a large, one-off batch that runs whenever the content changes, and\nnobody is waiting on any single vector. (Embedding the user's *live* query is\nthe opposite shape — small and latency-sensitive — and belongs on a small\nreal-time endpoint. This notebook is only the offline half.)\n\nThe corpus is\n[`sentence-transformers/natural-questions`](https://huggingface.co/datasets/sentence-transformers/natural-questions):\nreal Google search queries paired with the Wikipedia passages that answer\nthem. We embed the passages through an asynchronous endpoint, then embed a\nheld-out query the same way and check that its matching passage ranks near\nthe top. The embedding model is served with\n[Text Embeddings Inference (TEI)](https://huggingface.co/docs/text-embeddings-inference),\nHugging Face's container for embedding models.\n\nReferences:\n\n- [SageMaker asynchronous inference](https://docs.aws.amazon.com/sagemaker/latest/dg/async-inference.html)\n- [SageMaker async autoscaling](https://docs.aws.amazon.com/sagemaker/latest/dg/async-inference-autoscale.html)\n- [Natural Questions dataset](https://huggingface.co/datasets/sentence-transformers/natural-questions)" + }, + { + "cell_type": "markdown", + "id": "cell-02", + "metadata": {}, + "source": "## Prerequisites\n\nRun the next cell before importing the SDK. It installs the SageMaker Python\nSDK and `datasets` into the active kernel.\n\nYou also need an existing SageMaker execution role with access to SageMaker,\nS3, CloudWatch, Application Auto Scaling, and the ECR repository that hosts\nthe selected serving DLC." + }, + { + "cell_type": "code", + "execution_count": null, + "id": "cell-03", + "metadata": {}, + "outputs": [], + "source": [ + "%pip install \"sagemaker>=3.0.0\" datasets --upgrade --quiet" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "cell-04", + "metadata": {}, + "outputs": [], + "source": [ + "import datetime as dt\n", + "import json\n", + "import math\n", + "import os\n", + "import time\n", + "import uuid\n", + "from urllib.parse import urlparse\n", + "\n", + "import boto3\n", + "from botocore.exceptions import ClientError\n", + "from datasets import load_dataset\n", + "\n", + "from sagemaker.core import image_uris\n", + "from sagemaker.core.helper.session_helper import Session, get_execution_role\n", + "from sagemaker.core.inference_config import AsyncInferenceConfig\n", + "from sagemaker.serve import ModelBuilder, ModelServer\n", + "from sagemaker.serve.builder.schema_builder import SchemaBuilder" + ] + }, + { + "cell_type": "markdown", + "id": "cell-05", + "metadata": {}, + "source": "## The corpus and the embedding model\n\nEach record in\n[`sentence-transformers/natural-questions`](https://huggingface.co/datasets/sentence-transformers/natural-questions)\nis a search query paired with a Wikipedia passage that answers it. We embed\nthe passages to build the search index and keep the queries to test retrieval\nafterwards. The next cell loads a small slice; raise `DATASET_SIZE` to index\nmore.\n\nThe embedding model is `BAAI/bge-small-en-v1.5`, which produces\n384-dimensional vectors and runs on a CPU instance. To use a different model,\nset `HF_MODEL_ID` and set `EMBEDDING_DIM` to its output dimension." + }, + { + "cell_type": "code", + "execution_count": null, + "id": "cell-06", + "metadata": {}, + "outputs": [], + "source": [ + "PROJECT = \"hf-async-rag\"\n", + "RUN_ID = dt.datetime.now(dt.timezone.utc).strftime(\"%Y%m%d%H%M%S\")\n", + "\n", + "DATASET_ID = \"sentence-transformers/natural-questions\"\n", + "DATASET_SIZE = int(os.getenv(\"DATASET_SIZE\", \"96\"))\n", + "TEXTS_PER_INVOCATION = int(os.getenv(\"TEXTS_PER_INVOCATION\", \"8\"))\n", + "MAX_CONTEXT_CHARS = int(os.getenv(\"MAX_CONTEXT_CHARS\", \"1200\"))\n", + "\n", + "MODEL_ID = os.getenv(\"HF_MODEL_ID\", \"BAAI/bge-small-en-v1.5\")\n", + "EMBEDDING_DIM = int(os.getenv(\"EMBEDDING_DIM\", \"384\"))\n", + "TEI_VERSION = os.getenv(\"TEI_VERSION\", \"1.8.2\")\n", + "INSTANCE_TYPE = os.getenv(\"SAGEMAKER_INSTANCE_TYPE\", \"ml.c6i.xlarge\")\n", + "ENDPOINT_NAME = os.getenv(\"SAGEMAKER_ENDPOINT_NAME\", f\"{PROJECT}-{RUN_ID}\")\n", + "\n", + "MAX_INSTANCE_COUNT = int(os.getenv(\"MAX_INSTANCE_COUNT\", \"4\"))\n", + "BACKLOG_PER_INSTANCE_TARGET = float(os.getenv(\"BACKLOG_PER_INSTANCE_TARGET\", \"5\"))\n", + "MAX_CONCURRENT_INVOCATIONS_PER_INSTANCE = int(\n", + " os.getenv(\"MAX_CONCURRENT_INVOCATIONS_PER_INSTANCE\", \"4\")\n", + ")\n", + "\n", + "SUCCESS_SNS_TOPIC_ARN = os.getenv(\"SUCCESS_SNS_TOPIC_ARN\")\n", + "ERROR_SNS_TOPIC_ARN = os.getenv(\"ERROR_SNS_TOPIC_ARN\")\n", + "ALARM_SNS_TOPIC_ARN = os.getenv(\"ALARM_SNS_TOPIC_ARN\")\n", + "\n", + "# Keep cleanup on when running this file as a script. Set CLEANUP=false if you\n", + "# want to inspect the endpoint after the tutorial finishes.\n", + "CLEANUP = os.getenv(\"CLEANUP\", \"true\").lower() not in {\"0\", \"false\", \"no\"}" + ] + }, + { + "cell_type": "markdown", + "id": "cell-07", + "metadata": {}, + "source": "## Set up the SageMaker session\n\nThe endpoint runs under a SageMaker execution role: an IAM role that grants\naccess to S3, ECR, and CloudWatch. Set `SAGEMAKER_EXECUTION_ROLE_ARN` to the\nrole you want to use, or `SAGEMAKER_EXECUTION_ROLE_NAME` if you only have its\nname. Inside SageMaker Studio or a notebook instance you can leave both unset\nand the role is detected automatically." + }, + { + "cell_type": "code", + "execution_count": null, + "id": "cell-08", + "metadata": {}, + "outputs": [], + "source": [ + "requested_region = os.getenv(\"AWS_REGION\") or os.getenv(\"AWS_DEFAULT_REGION\")\n", + "boto_session = boto3.Session(region_name=requested_region) if requested_region else boto3.Session()\n", + "sess = Session(boto_session=boto_session)\n", + "region = sess.boto_region_name\n", + "\n", + "s3 = boto_session.client(\"s3\")\n", + "sm = boto_session.client(\"sagemaker\")\n", + "logs = boto_session.client(\"logs\")\n", + "cloudwatch = boto_session.client(\"cloudwatch\")\n", + "autoscaling = boto_session.client(\"application-autoscaling\")\n", + "\n", + "\n", + "def resolve_role(session, sagemaker_session):\n", + " role_arn = os.getenv(\"SAGEMAKER_EXECUTION_ROLE_ARN\")\n", + " if role_arn:\n", + " return role_arn\n", + "\n", + " role_name = os.getenv(\"SAGEMAKER_EXECUTION_ROLE_NAME\")\n", + " if role_name:\n", + " iam = session.client(\"iam\")\n", + " return iam.get_role(RoleName=role_name)[\"Role\"][\"Arn\"]\n", + "\n", + " return get_execution_role(sagemaker_session=sagemaker_session)\n", + "\n", + "\n", + "role = resolve_role(boto_session, sess)\n", + "bucket = sess.default_bucket()\n", + "\n", + "base_s3_uri = os.getenv(\"SAGEMAKER_ASYNC_BASE_S3_URI\", f\"s3://{bucket}/{PROJECT}/{RUN_ID}\")\n", + "input_s3_prefix = f\"{base_s3_uri.rstrip('/')}/input\"\n", + "output_s3_prefix = f\"{base_s3_uri.rstrip('/')}/output\"\n", + "failure_s3_prefix = f\"{base_s3_uri.rstrip('/')}/failure\"\n", + "index_s3_uri = f\"{base_s3_uri.rstrip('/')}/index/documents-with-embeddings.jsonl\"\n", + "\n", + "print(f\"region: {region}\")\n", + "print(f\"role: {role}\")\n", + "print(f\"endpoint: {ENDPOINT_NAME}\")\n", + "print(f\"async input: {input_s3_prefix}\")\n", + "print(f\"async output: {output_s3_prefix}\")" + ] + }, + { + "cell_type": "markdown", + "id": "cell-09", + "metadata": {}, + "source": "## Load the corpus\n\nLoad the slice and reshape each row into a record with a stable `id`, its\nquery, and the passage (truncated to `MAX_CONTEXT_CHARS`). The passages are\nwhat we embed; the first record's query is set aside for the retrieval test at\nthe end." + }, + { + "cell_type": "code", + "execution_count": null, + "id": "cell-10", + "metadata": {}, + "outputs": [], + "source": [ + "raw_dataset = load_dataset(DATASET_ID, \"pair\", split=f\"train[:{DATASET_SIZE}]\")\n", + "\n", + "records = []\n", + "for row_index, row in enumerate(raw_dataset):\n", + " passage = row[\"answer\"].strip()\n", + " query = row[\"query\"].strip()\n", + " if not passage or not query:\n", + " continue\n", + "\n", + " records.append(\n", + " {\n", + " \"id\": f\"nq-{row_index:05d}\",\n", + " \"question\": query,\n", + " \"context\": passage[:MAX_CONTEXT_CHARS],\n", + " }\n", + " )\n", + "\n", + "if len(records) < 2:\n", + " raise ValueError(\"Need at least two usable records for the retrieval smoke test.\")\n", + "\n", + "query_record = records[0]\n", + "document_records = records\n", + "\n", + "print(f\"loaded records: {len(document_records)}\")\n", + "print(f\"held-out query: {query_record['question']}\")" + ] + }, + { + "cell_type": "markdown", + "id": "cell-11", + "metadata": {}, + "source": [ + "## Select the TEI serving container\n", + "\n", + "Asynchronous inference does not prescribe a model server. For an embedding\n", + "workload the Hugging Face choice is Text Embeddings Inference (TEI): CPU\n", + "instances use `huggingface-tei-cpu`, GPU instances use `huggingface-tei`. The\n", + "helper below picks the right one for the configured instance type.\n" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "cell-12", + "metadata": {}, + "outputs": [], + "source": [ + "def is_gpu_instance(instance_type):\n", + " return instance_type.startswith((\"ml.g\", \"ml.p\"))\n", + "\n", + "\n", + "def get_tei_image_uri(instance_type):\n", + " framework = \"huggingface-tei\" if is_gpu_instance(instance_type) else \"huggingface-tei-cpu\"\n", + " return image_uris.retrieve(\n", + " framework=framework,\n", + " region=region,\n", + " version=TEI_VERSION,\n", + " image_scope=\"inference\",\n", + " instance_type=instance_type,\n", + " )\n", + "\n", + "\n", + "image_uri = get_tei_image_uri(INSTANCE_TYPE)\n", + "print(image_uri)" + ] + }, + { + "cell_type": "markdown", + "id": "cell-13", + "metadata": {}, + "source": "## Build the model\n\n`ModelBuilder` describes the SageMaker model: the Hub model ID, the serving\ncontainer, the model server, and a small input/output example. The container\ndownloads the model from the Hub when the endpoint starts. For gated or\nprivate models, set `HF_TOKEN` (or `HUGGING_FACE_HUB_TOKEN`) before running\nthe notebook." + }, + { + "cell_type": "code", + "execution_count": null, + "id": "cell-14", + "metadata": {}, + "outputs": [], + "source": [ + "hf_token = os.getenv(\"HF_TOKEN\") or os.getenv(\"HUGGING_FACE_HUB_TOKEN\")\n", + "\n", + "env_vars = {\n", + " \"HF_MODEL_ID\": MODEL_ID,\n", + " \"MAX_BATCH_TOKENS\": os.getenv(\"MAX_BATCH_TOKENS\", \"16384\"),\n", + " \"MAX_CLIENT_BATCH_SIZE\": os.getenv(\"MAX_CLIENT_BATCH_SIZE\", \"32\"),\n", + "}\n", + "\n", + "if hf_token:\n", + " env_vars[\"HF_TOKEN\"] = hf_token\n", + " env_vars[\"HUGGING_FACE_HUB_TOKEN\"] = hf_token\n", + "\n", + "resource_tags = [\n", + " {\"Key\": \"Project\", \"Value\": PROJECT},\n", + " {\"Key\": \"ModelId\", \"Value\": MODEL_ID},\n", + " {\"Key\": \"CreatedBy\", \"Value\": \"hf-sagemaker-docs\"},\n", + "]\n", + "\n", + "model_builder = ModelBuilder(\n", + " model=MODEL_ID,\n", + " role_arn=role,\n", + " sagemaker_session=sess,\n", + " instance_type=INSTANCE_TYPE,\n", + " image_uri=image_uri,\n", + " model_server=ModelServer.TEI,\n", + " env_vars=env_vars,\n", + " schema_builder=SchemaBuilder(\n", + " sample_input={\"inputs\": [\"who wrote the origin of species\"]},\n", + " sample_output=[[0.0] * EMBEDDING_DIM],\n", + " ),\n", + ")\n", + "\n", + "tei_model = model_builder.build(model_name=f\"{PROJECT}-model-{RUN_ID}\")" + ] + }, + { + "cell_type": "markdown", + "id": "cell-15", + "metadata": {}, + "source": [ + "## Deploy the asynchronous endpoint\n", + "\n", + "The request body lives in S3, `invoke_async` sends SageMaker a pointer to that\n", + "object, and SageMaker later writes either the response or the failure payload\n", + "back to S3. Once `AsyncInferenceConfig` is attached to the endpoint\n", + "configuration, the endpoint accepts async invocations only.\n", + "\n", + "We deploy with one instance so the container can start and the first\n", + "retrieval test does not wait for scale-out. The autoscaling policy in the next\n", + "section lets the same endpoint scale to zero after the queue is empty.\n" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "cell-16", + "metadata": {}, + "outputs": [], + "source": [ + "notification_config = {}\n", + "if SUCCESS_SNS_TOPIC_ARN:\n", + " notification_config[\"SuccessTopic\"] = SUCCESS_SNS_TOPIC_ARN\n", + "if ERROR_SNS_TOPIC_ARN:\n", + " notification_config[\"ErrorTopic\"] = ERROR_SNS_TOPIC_ARN\n", + "\n", + "async_config = AsyncInferenceConfig(\n", + " output_path=output_s3_prefix,\n", + " failure_path=failure_s3_prefix,\n", + " max_concurrent_invocations_per_instance=MAX_CONCURRENT_INVOCATIONS_PER_INSTANCE,\n", + " notification_config=notification_config or None,\n", + ")\n", + "\n", + "endpoint = model_builder.deploy(\n", + " endpoint_name=ENDPOINT_NAME,\n", + " initial_instance_count=1,\n", + " instance_type=INSTANCE_TYPE,\n", + " inference_config=async_config,\n", + " container_timeout_in_seconds=900,\n", + " tags=resource_tags,\n", + " wait=True,\n", + ")\n", + "\n", + "endpoint_description = sm.describe_endpoint(EndpointName=ENDPOINT_NAME)\n", + "endpoint_config_name = endpoint_description[\"EndpointConfigName\"]\n", + "model_name = tei_model.model_name\n", + "\n", + "print(f\"endpoint status: {endpoint_description['EndpointStatus']}\")\n", + "print(f\"endpoint config: {endpoint_config_name}\")\n", + "print(f\"model: {model_name}\")" + ] + }, + { + "cell_type": "markdown", + "id": "cell-17", + "metadata": {}, + "source": [ + "## Autoscale and scale to zero\n", + "\n", + "Scale-to-zero is one of the reasons asynchronous inference is attractive for\n", + "batch workloads. The target-tracking policy scales with queue depth, while the\n", + "step-scaling policy wakes the endpoint from zero as soon as a backlog appears.\n", + "Without that wake-up alarm, the endpoint might wait until the queue exceeds\n", + "the target-tracking threshold before it adds the first instance.\n" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "cell-18", + "metadata": {}, + "outputs": [], + "source": [ + "variant_name = \"AllTraffic\"\n", + "resource_id = f\"endpoint/{ENDPOINT_NAME}/variant/{variant_name}\"\n", + "\n", + "autoscaling.register_scalable_target(\n", + " ServiceNamespace=\"sagemaker\",\n", + " ResourceId=resource_id,\n", + " ScalableDimension=\"sagemaker:variant:DesiredInstanceCount\",\n", + " MinCapacity=0,\n", + " MaxCapacity=MAX_INSTANCE_COUNT,\n", + ")\n", + "\n", + "autoscaling.put_scaling_policy(\n", + " PolicyName=f\"{ENDPOINT_NAME}-backlog-target-tracking\",\n", + " ServiceNamespace=\"sagemaker\",\n", + " ResourceId=resource_id,\n", + " ScalableDimension=\"sagemaker:variant:DesiredInstanceCount\",\n", + " PolicyType=\"TargetTrackingScaling\",\n", + " TargetTrackingScalingPolicyConfiguration={\n", + " \"TargetValue\": BACKLOG_PER_INSTANCE_TARGET,\n", + " \"CustomizedMetricSpecification\": {\n", + " \"MetricName\": \"ApproximateBacklogSizePerInstance\",\n", + " \"Namespace\": \"AWS/SageMaker\",\n", + " \"Dimensions\": [{\"Name\": \"EndpointName\", \"Value\": ENDPOINT_NAME}],\n", + " \"Statistic\": \"Average\",\n", + " },\n", + " \"ScaleInCooldown\": 300,\n", + " \"ScaleOutCooldown\": 60,\n", + " },\n", + ")\n", + "\n", + "step_policy = autoscaling.put_scaling_policy(\n", + " PolicyName=f\"{ENDPOINT_NAME}-wake-from-zero\",\n", + " ServiceNamespace=\"sagemaker\",\n", + " ResourceId=resource_id,\n", + " ScalableDimension=\"sagemaker:variant:DesiredInstanceCount\",\n", + " PolicyType=\"StepScaling\",\n", + " StepScalingPolicyConfiguration={\n", + " \"AdjustmentType\": \"ChangeInCapacity\",\n", + " \"MetricAggregationType\": \"Average\",\n", + " \"Cooldown\": 300,\n", + " \"StepAdjustments\": [{\"MetricIntervalLowerBound\": 0, \"ScalingAdjustment\": 1}],\n", + " },\n", + ")\n", + "\n", + "wake_alarm_name = f\"{ENDPOINT_NAME}-has-backlog-without-capacity\"\n", + "backlog_alarm_name = f\"{ENDPOINT_NAME}-async-backlog-high\"\n", + "failure_alarm_name = f\"{ENDPOINT_NAME}-async-failures\"\n", + "alarm_names = [wake_alarm_name, backlog_alarm_name, failure_alarm_name]\n", + "\n", + "cloudwatch.put_metric_alarm(\n", + " AlarmName=wake_alarm_name,\n", + " AlarmDescription=\"Wake async endpoint from zero when requests are queued.\",\n", + " Namespace=\"AWS/SageMaker\",\n", + " MetricName=\"HasBacklogWithoutCapacity\",\n", + " Dimensions=[{\"Name\": \"EndpointName\", \"Value\": ENDPOINT_NAME}],\n", + " Statistic=\"Average\",\n", + " Period=60,\n", + " EvaluationPeriods=2,\n", + " DatapointsToAlarm=2,\n", + " Threshold=1,\n", + " ComparisonOperator=\"GreaterThanOrEqualToThreshold\",\n", + " TreatMissingData=\"missing\",\n", + " AlarmActions=[step_policy[\"PolicyARN\"]],\n", + ")\n", + "\n", + "backlog_alarm = {\n", + " \"AlarmName\": backlog_alarm_name,\n", + " \"AlarmDescription\": \"Async queue is growing faster than the endpoint can drain it.\",\n", + " \"Namespace\": \"AWS/SageMaker\",\n", + " \"MetricName\": \"ApproximateBacklogSize\",\n", + " \"Dimensions\": [{\"Name\": \"EndpointName\", \"Value\": ENDPOINT_NAME}],\n", + " \"Statistic\": \"Average\",\n", + " \"Period\": 60,\n", + " \"EvaluationPeriods\": 3,\n", + " \"DatapointsToAlarm\": 2,\n", + " \"Threshold\": 50,\n", + " \"ComparisonOperator\": \"GreaterThanThreshold\",\n", + " \"TreatMissingData\": \"notBreaching\",\n", + "}\n", + "\n", + "failure_alarm = {\n", + " \"AlarmName\": failure_alarm_name,\n", + " \"AlarmDescription\": \"Async inference requests are failing.\",\n", + " \"Namespace\": \"AWS/SageMaker\",\n", + " \"MetricName\": \"InvocationsFailed\",\n", + " \"Dimensions\": [{\"Name\": \"EndpointName\", \"Value\": ENDPOINT_NAME}],\n", + " \"Statistic\": \"Sum\",\n", + " \"Period\": 60,\n", + " \"EvaluationPeriods\": 1,\n", + " \"DatapointsToAlarm\": 1,\n", + " \"Threshold\": 5,\n", + " \"ComparisonOperator\": \"GreaterThanThreshold\",\n", + " \"TreatMissingData\": \"notBreaching\",\n", + "}\n", + "\n", + "if ALARM_SNS_TOPIC_ARN:\n", + " backlog_alarm[\"AlarmActions\"] = [ALARM_SNS_TOPIC_ARN]\n", + " failure_alarm[\"AlarmActions\"] = [ALARM_SNS_TOPIC_ARN]\n", + "\n", + "cloudwatch.put_metric_alarm(**backlog_alarm)\n", + "cloudwatch.put_metric_alarm(**failure_alarm)\n", + "\n", + "print(f\"registered scalable target: {resource_id}\")\n", + "print(f\"alarms: {', '.join(alarm_names)}\")" + ] + }, + { + "cell_type": "markdown", + "id": "cell-19", + "metadata": {}, + "source": "## Submit the embedding requests\n\nAn asynchronous endpoint does not take the request body directly. Each batch\nis uploaded to S3 first, and `invoke_async` sends the endpoint that S3\nlocation instead of the payload. The container still receives an ordinary\nembedding request: `{\"inputs\": [...]}`.\n\n`invoke_async` returns right away with a pointer to where the output will be\nwritten; it does not wait for the vectors." + }, + { + "cell_type": "code", + "execution_count": null, + "id": "cell-20", + "metadata": {}, + "outputs": [], + "source": [ + "def parse_s3_uri(uri):\n", + " parsed = urlparse(uri)\n", + " if parsed.scheme != \"s3\" or not parsed.netloc or not parsed.path.strip(\"/\"):\n", + " raise ValueError(f\"expected S3 URI, got {uri!r}\")\n", + " return parsed.netloc, parsed.path.lstrip(\"/\")\n", + "\n", + "\n", + "def s3_join(prefix, *parts):\n", + " return \"/\".join([prefix.rstrip(\"/\"), *(part.strip(\"/\") for part in parts)])\n", + "\n", + "\n", + "def put_json(uri, payload):\n", + " target_bucket, key = parse_s3_uri(uri)\n", + " s3.put_object(\n", + " Bucket=target_bucket,\n", + " Key=key,\n", + " Body=json.dumps(payload).encode(\"utf-8\"),\n", + " ContentType=\"application/json\",\n", + " )\n", + "\n", + "\n", + "def put_text(uri, text):\n", + " target_bucket, key = parse_s3_uri(uri)\n", + " s3.put_object(Bucket=target_bucket, Key=key, Body=text.encode(\"utf-8\"))\n", + "\n", + "\n", + "def batched(items, size):\n", + " for start in range(0, len(items), size):\n", + " yield items[start : start + size]\n", + "\n", + "\n", + "assert parse_s3_uri(\"s3://example-bucket/path/file.json\") == (\n", + " \"example-bucket\",\n", + " \"path/file.json\",\n", + ")\n", + "\n", + "document_jobs = []\n", + "for batch_index, batch in enumerate(batched(document_records, TEXTS_PER_INVOCATION), start=1):\n", + " payload = {\"inputs\": [record[\"context\"] for record in batch]}\n", + " input_uri = s3_join(input_s3_prefix, \"documents\", f\"batch-{batch_index:04d}.json\")\n", + " put_json(input_uri, payload)\n", + "\n", + " response = endpoint.invoke_async(\n", + " input_location=input_uri,\n", + " content_type=\"application/json\",\n", + " accept=\"application/json\",\n", + " inference_id=f\"documents-{batch_index:04d}-{uuid.uuid4()}\",\n", + " invocation_timeout_seconds=900,\n", + " session=boto_session,\n", + " region=region,\n", + " )\n", + " document_jobs.append((batch, response))\n", + " print(f\"submitted {input_uri} -> {response.output_location}\")" + ] + }, + { + "cell_type": "markdown", + "id": "cell-21", + "metadata": {}, + "source": [ + "## Collect the embeddings\n", + "\n", + "The async response objects give us both the output path and the failure path.\n", + "Polling S3 keeps the notebook simple, while production pipelines often use\n", + "SNS, EventBridge, or a workflow engine to react to completed outputs.\n", + "\n", + "Once every batch returns, the vectors are joined back to the metadata that\n", + "created them and written as JSON Lines. A downstream indexer could read that\n", + "file and push the embeddings into OpenSearch, PostgreSQL with pgvector, a\n", + "vector database, or another retrieval store.\n" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "cell-22", + "metadata": {}, + "outputs": [], + "source": [ + "def is_missing_key_error(error):\n", + " return error.response.get(\"Error\", {}).get(\"Code\") in {\"NoSuchKey\", \"404\", \"NotFound\"}\n", + "\n", + "\n", + "def read_s3_text(uri):\n", + " source_bucket, key = parse_s3_uri(uri)\n", + " response = s3.get_object(Bucket=source_bucket, Key=key)\n", + " return response[\"Body\"].read().decode(\"utf-8\")\n", + "\n", + "\n", + "def wait_for_async_json(response, timeout=1800, poll=10):\n", + " deadline = time.time() + timeout\n", + "\n", + " while time.time() < deadline:\n", + " for uri, failed in (\n", + " (response.output_location, False),\n", + " (response.failure_location, True),\n", + " ):\n", + " if not uri:\n", + " continue\n", + "\n", + " try:\n", + " body = read_s3_text(uri)\n", + " except ClientError as error:\n", + " if is_missing_key_error(error):\n", + " continue\n", + " raise\n", + "\n", + " if failed:\n", + " raise RuntimeError(f\"async inference failed: {body}\")\n", + " return json.loads(body)\n", + "\n", + " print(\"waiting for async output...\")\n", + " time.sleep(poll)\n", + "\n", + " raise TimeoutError(f\"no async result after {timeout} seconds: {response.output_location}\")\n", + "\n", + "\n", + "index_records = []\n", + "for batch, response in document_jobs:\n", + " vectors = wait_for_async_json(response)\n", + " if len(vectors) != len(batch):\n", + " raise RuntimeError(f\"expected {len(batch)} vectors, got {len(vectors)}\")\n", + "\n", + " for record, vector in zip(batch, vectors):\n", + " index_records.append(\n", + " {\n", + " \"id\": record[\"id\"],\n", + " \"question\": record[\"question\"],\n", + " \"context\": record[\"context\"],\n", + " \"embedding\": vector,\n", + " }\n", + " )\n", + "\n", + "assert len(index_records) == len(document_records)\n", + "assert all(len(record[\"embedding\"]) == EMBEDDING_DIM for record in index_records)\n", + "\n", + "put_text(index_s3_uri, \"\\n\".join(json.dumps(record) for record in index_records))\n", + "\n", + "print(f\"embedded documents: {len(index_records)}\")\n", + "print(f\"embedding dimensions: {len(index_records[0]['embedding'])}\")\n", + "print(f\"index written to: {index_s3_uri}\")" + ] + }, + { + "cell_type": "markdown", + "id": "cell-23", + "metadata": {}, + "source": "## Validate with a retrieval test\n\nTo check that the embeddings support search, we embed the held-out query\nthrough the same asynchronous endpoint, score it against every passage\nembedding with cosine similarity, and rank the passages. The passage that\noriginally answered the query should come out on top." + }, + { + "cell_type": "code", + "execution_count": null, + "id": "cell-24", + "metadata": {}, + "outputs": [], + "source": [ + "def cosine(a, b):\n", + " dot = sum(x * y for x, y in zip(a, b))\n", + " norm_a = math.sqrt(sum(x * x for x in a))\n", + " norm_b = math.sqrt(sum(y * y for y in b))\n", + " if norm_a == 0 or norm_b == 0:\n", + " raise ValueError(\"cosine similarity is undefined for a zero vector\")\n", + " return dot / (norm_a * norm_b)\n", + "\n", + "\n", + "assert cosine([1.0, 0.0], [0.0, 1.0]) == 0.0\n", + "\n", + "query_uri = s3_join(input_s3_prefix, \"queries\", \"held-out-question.json\")\n", + "put_json(query_uri, {\"inputs\": [query_record[\"question\"]]})\n", + "\n", + "query_response = endpoint.invoke_async(\n", + " input_location=query_uri,\n", + " content_type=\"application/json\",\n", + " accept=\"application/json\",\n", + " inference_id=f\"query-{uuid.uuid4()}\",\n", + " invocation_timeout_seconds=900,\n", + " session=boto_session,\n", + " region=region,\n", + ")\n", + "query_embedding = wait_for_async_json(query_response)[0]\n", + "\n", + "ranked = sorted(\n", + " (\n", + " {\n", + " \"id\": record[\"id\"],\n", + " \"score\": cosine(query_embedding, record[\"embedding\"]),\n", + " \"question\": record[\"question\"],\n", + " \"context\": record[\"context\"],\n", + " }\n", + " for record in index_records\n", + " ),\n", + " key=lambda item: item[\"score\"],\n", + " reverse=True,\n", + ")\n", + "\n", + "for hit in ranked[:5]:\n", + " print(f\"{hit['score']:.3f} {hit['id']}: {hit['question']}\")\n", + "\n", + "top_ids = [hit[\"id\"] for hit in ranked[:5]]\n", + "assert query_record[\"id\"] in top_ids" + ] + }, + { + "cell_type": "markdown", + "id": "cell-27", + "metadata": {}, + "source": [ + "## Clean up\n", + "\n", + "Delete the endpoint, endpoint configuration, model, autoscaling target, and\n", + "tutorial alarms when you are done. The S3 inputs and outputs are left in\n", + "place because they are useful for inspection and because many teams hand\n", + "those objects to the next indexing stage.\n", + "\n" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "cell-28", + "metadata": {}, + "outputs": [], + "source": [ + "def ignore_not_found(error):\n", + " code = error.response.get(\"Error\", {}).get(\"Code\", \"\")\n", + " message = error.response.get(\"Error\", {}).get(\"Message\", \"\")\n", + " return code in {\"ResourceNotFound\", \"ResourceNotFoundException\"} or \"not exist\" in message\n", + "\n", + "\n", + "def cleanup_resources():\n", + " print(\"deleting CloudWatch alarms\")\n", + " try:\n", + " cloudwatch.delete_alarms(AlarmNames=alarm_names)\n", + " except ClientError as error:\n", + " if not ignore_not_found(error):\n", + " raise\n", + "\n", + " print(\"deregistering scalable target\")\n", + " try:\n", + " autoscaling.deregister_scalable_target(\n", + " ServiceNamespace=\"sagemaker\",\n", + " ResourceId=resource_id,\n", + " ScalableDimension=\"sagemaker:variant:DesiredInstanceCount\",\n", + " )\n", + " except ClientError as error:\n", + " if not ignore_not_found(error):\n", + " raise\n", + "\n", + " print(\"deleting endpoint\")\n", + " try:\n", + " sm.delete_endpoint(EndpointName=ENDPOINT_NAME)\n", + " sm.get_waiter(\"endpoint_deleted\").wait(\n", + " EndpointName=ENDPOINT_NAME,\n", + " WaiterConfig={\"Delay\": 30, \"MaxAttempts\": 60},\n", + " )\n", + " except ClientError as error:\n", + " if not ignore_not_found(error):\n", + " raise\n", + "\n", + " print(\"deleting endpoint config\")\n", + " try:\n", + " sm.delete_endpoint_config(EndpointConfigName=endpoint_config_name)\n", + " except ClientError as error:\n", + " if not ignore_not_found(error):\n", + " raise\n", + "\n", + " print(\"deleting model\")\n", + " try:\n", + " sm.delete_model(ModelName=model_name)\n", + " except ClientError as error:\n", + " if not ignore_not_found(error):\n", + " raise\n", + "\n", + "\n", + "if CLEANUP:\n", + " cleanup_resources()\n", + "else:\n", + " print(f\"left endpoint running: {ENDPOINT_NAME}\")\n" + ] + } + ], + "metadata": { + "kernelspec": { + "display_name": "SageMaker v3 (.venv-sm-v3)", + "language": "python", + "name": "sagemaker-v3" + }, + "language_info": { + "codemirror_mode": { + "name": "ipython", + "version": 3 + }, + "file_extension": ".py", + "mimetype": "text/x-python", + "name": "python", + "nbconvert_exporter": "python", + "pygments_lexer": "ipython3", + "version": "3.12.13" + } + }, + "nbformat": 4, + "nbformat_minor": 5 +} \ No newline at end of file