Buckets:

hf-doc-build/doc-dev / sagemaker /pr_2709 /en /examples /sagemaker-sdk-async-inference-embedding-tei.html
download
raw
92.2 kB
<meta charset="utf-8" /><meta name="hf:doc:metadata" content="{&quot;title&quot;:&quot;Batch-embed a corpus with SageMaker asynchronous inference&quot;,&quot;local&quot;:&quot;batch-embed-a-corpus-with-sagemaker-asynchronous-inference&quot;,&quot;sections&quot;:[{&quot;title&quot;:&quot;What asynchronous inference is&quot;,&quot;local&quot;:&quot;what-asynchronous-inference-is&quot;,&quot;sections&quot;:[],&quot;depth&quot;:2},{&quot;title&quot;:&quot;The use case we picked&quot;,&quot;local&quot;:&quot;the-use-case-we-picked&quot;,&quot;sections&quot;:[],&quot;depth&quot;:2},{&quot;title&quot;:&quot;Prerequisites&quot;,&quot;local&quot;:&quot;prerequisites&quot;,&quot;sections&quot;:[],&quot;depth&quot;:2},{&quot;title&quot;:&quot;The corpus and the embedding model&quot;,&quot;local&quot;:&quot;the-corpus-and-the-embedding-model&quot;,&quot;sections&quot;:[],&quot;depth&quot;:2},{&quot;title&quot;:&quot;Set up the SageMaker session&quot;,&quot;local&quot;:&quot;set-up-the-sagemaker-session&quot;,&quot;sections&quot;:[],&quot;depth&quot;:2},{&quot;title&quot;:&quot;Load the corpus&quot;,&quot;local&quot;:&quot;load-the-corpus&quot;,&quot;sections&quot;:[],&quot;depth&quot;:2},{&quot;title&quot;:&quot;Select the TEI serving container&quot;,&quot;local&quot;:&quot;select-the-tei-serving-container&quot;,&quot;sections&quot;:[],&quot;depth&quot;:2},{&quot;title&quot;:&quot;Build the model&quot;,&quot;local&quot;:&quot;build-the-model&quot;,&quot;sections&quot;:[],&quot;depth&quot;:2},{&quot;title&quot;:&quot;Deploy the asynchronous endpoint&quot;,&quot;local&quot;:&quot;deploy-the-asynchronous-endpoint&quot;,&quot;sections&quot;:[],&quot;depth&quot;:2},{&quot;title&quot;:&quot;Autoscale and scale to zero&quot;,&quot;local&quot;:&quot;autoscale-and-scale-to-zero&quot;,&quot;sections&quot;:[],&quot;depth&quot;:2},{&quot;title&quot;:&quot;Submit the embedding requests&quot;,&quot;local&quot;:&quot;submit-the-embedding-requests&quot;,&quot;sections&quot;:[],&quot;depth&quot;:2},{&quot;title&quot;:&quot;Collect the embeddings&quot;,&quot;local&quot;:&quot;collect-the-embeddings&quot;,&quot;sections&quot;:[],&quot;depth&quot;:2},{&quot;title&quot;:&quot;Validate with a retrieval test&quot;,&quot;local&quot;:&quot;validate-with-a-retrieval-test&quot;,&quot;sections&quot;:[],&quot;depth&quot;:2},{&quot;title&quot;:&quot;Clean up&quot;,&quot;local&quot;:&quot;clean-up&quot;,&quot;sections&quot;:[],&quot;depth&quot;:2}],&quot;depth&quot;:1}"/>
<link href="/docs/sagemaker/pr_2709/en/_app/immutable/entry/start.COiK-bUV.js" rel="modulepreload">
<link href="/docs/sagemaker/pr_2709/en/_app/immutable/chunks/UGDyBJsP.js" rel="modulepreload">
<link href="/docs/sagemaker/pr_2709/en/_app/immutable/chunks/Paddn1qZ.js" rel="modulepreload">
<link href="/docs/sagemaker/pr_2709/en/_app/immutable/entry/app.CjYANwwM.js" rel="modulepreload">
<link href="/docs/sagemaker/pr_2709/en/_app/immutable/chunks/Beft82ym.js" rel="modulepreload">
<link href="/docs/sagemaker/pr_2709/en/_app/immutable/chunks/Ds9Xs-bq.js" rel="modulepreload">
<link href="/docs/sagemaker/pr_2709/en/_app/immutable/chunks/DsnmJJEf.js" rel="modulepreload">
<link href="/docs/sagemaker/pr_2709/en/_app/immutable/chunks/DnXYL5Yp.js" rel="modulepreload">
<link href="/docs/sagemaker/pr_2709/en/_app/immutable/nodes/0.CsY3iCn7.js" rel="modulepreload">
<link href="/docs/sagemaker/pr_2709/en/_app/immutable/chunks/CxJVP3KE.js" rel="modulepreload">
<link href="/docs/sagemaker/pr_2709/en/_app/immutable/nodes/4.Bk69HVjz.js" rel="modulepreload">
<link href="/docs/sagemaker/pr_2709/en/_app/immutable/chunks/Cz9gzD1e.js" rel="modulepreload">
<link href="/docs/sagemaker/pr_2709/en/_app/immutable/chunks/CZujaBAM.js" rel="modulepreload">
<!--ga8rw2--><meta name="hf:doc:metadata" content="{&quot;title&quot;:&quot;Batch-embed a corpus with SageMaker asynchronous inference&quot;,&quot;local&quot;:&quot;batch-embed-a-corpus-with-sagemaker-asynchronous-inference&quot;,&quot;sections&quot;:[{&quot;title&quot;:&quot;What asynchronous inference is&quot;,&quot;local&quot;:&quot;what-asynchronous-inference-is&quot;,&quot;sections&quot;:[],&quot;depth&quot;:2},{&quot;title&quot;:&quot;The use case we picked&quot;,&quot;local&quot;:&quot;the-use-case-we-picked&quot;,&quot;sections&quot;:[],&quot;depth&quot;:2},{&quot;title&quot;:&quot;Prerequisites&quot;,&quot;local&quot;:&quot;prerequisites&quot;,&quot;sections&quot;:[],&quot;depth&quot;:2},{&quot;title&quot;:&quot;The corpus and the embedding model&quot;,&quot;local&quot;:&quot;the-corpus-and-the-embedding-model&quot;,&quot;sections&quot;:[],&quot;depth&quot;:2},{&quot;title&quot;:&quot;Set up the SageMaker session&quot;,&quot;local&quot;:&quot;set-up-the-sagemaker-session&quot;,&quot;sections&quot;:[],&quot;depth&quot;:2},{&quot;title&quot;:&quot;Load the corpus&quot;,&quot;local&quot;:&quot;load-the-corpus&quot;,&quot;sections&quot;:[],&quot;depth&quot;:2},{&quot;title&quot;:&quot;Select the TEI serving container&quot;,&quot;local&quot;:&quot;select-the-tei-serving-container&quot;,&quot;sections&quot;:[],&quot;depth&quot;:2},{&quot;title&quot;:&quot;Build the model&quot;,&quot;local&quot;:&quot;build-the-model&quot;,&quot;sections&quot;:[],&quot;depth&quot;:2},{&quot;title&quot;:&quot;Deploy the asynchronous endpoint&quot;,&quot;local&quot;:&quot;deploy-the-asynchronous-endpoint&quot;,&quot;sections&quot;:[],&quot;depth&quot;:2},{&quot;title&quot;:&quot;Autoscale and scale to zero&quot;,&quot;local&quot;:&quot;autoscale-and-scale-to-zero&quot;,&quot;sections&quot;:[],&quot;depth&quot;:2},{&quot;title&quot;:&quot;Submit the embedding requests&quot;,&quot;local&quot;:&quot;submit-the-embedding-requests&quot;,&quot;sections&quot;:[],&quot;depth&quot;:2},{&quot;title&quot;:&quot;Collect the embeddings&quot;,&quot;local&quot;:&quot;collect-the-embeddings&quot;,&quot;sections&quot;:[],&quot;depth&quot;:2},{&quot;title&quot;:&quot;Validate with a retrieval test&quot;,&quot;local&quot;:&quot;validate-with-a-retrieval-test&quot;,&quot;sections&quot;:[],&quot;depth&quot;:2},{&quot;title&quot;:&quot;Clean up&quot;,&quot;local&quot;:&quot;clean-up&quot;,&quot;sections&quot;:[],&quot;depth&quot;:2}],&quot;depth&quot;:1}"/><!---->
<link href="/docs/sagemaker/pr_2709/en/_app/immutable/assets/0.tn0RQdqM.css" rel="modulepreload"> <!--[--><!--[0--><!--[--><!--[0--><!--[--><p></p> <div class="items-center shrink-0 min-w-[100px] max-sm:min-w-[50px] justify-end ml-auto flex" style="float: right; margin-left: 10px; display: inline-flex; position: relative; z-index: 10;"><div class="inline-flex rounded-md max-sm:rounded-sm"><button class="inline-flex items-center gap-1 h-7 max-sm:h-7 px-2 max-sm:px-1.5 text-sm font-medium text-gray-800 border border-r-0 rounded-l-md max-sm:rounded-l-sm border-gray-200 bg-white hover:shadow-inner dark:border-gray-850 dark:bg-gray-950 dark:text-gray-200 dark:hover:bg-gray-800" aria-live="polite"><span class="inline-flex items-center justify-center rounded-md p-0.5 max-sm:p-0 hover:text-gray-800 dark:hover:text-gray-200"><svg class="sm:size-3.5 size-3" xmlns="http://www.w3.org/2000/svg" aria-hidden="true" fill="currentColor" focusable="false" role="img" width="1em" height="1em" preserveAspectRatio="xMidYMid meet" viewBox="0 0 32 32"><path d="M28,10V28H10V10H28m0-2H10a2,2,0,0,0-2,2V28a2,2,0,0,0,2,2H28a2,2,0,0,0,2-2V10a2,2,0,0,0-2-2Z" transform="translate(0)"></path><path d="M4,18H2V4A2,2,0,0,1,4,2H18V4H4Z" transform="translate(0)"></path><rect fill="none" width="32" height="32"></rect></svg><!----></span> <span>Copy page</span></button> <button class="inline-flex items-center justify-center w-6 max-sm:w-5 h-7 max-sm:h-7 disabled:pointer-events-none text-sm text-gray-500 hover:text-gray-700 dark:hover:text-white rounded-r-md max-sm:rounded-r-sm border border-l transition border-gray-200 bg-white hover:shadow-inner dark:border-gray-850 dark:bg-gray-950 dark:text-gray-200 dark:hover:bg-gray-800" aria-haspopup="menu" aria-expanded="false" aria-label="Open copy menu"><svg class="transition-transform text-gray-400 overflow-visible sm:size-3.5 size-3 rotate-0" width="1em" height="1em" viewBox="0 0 12 7" fill="none" xmlns="http://www.w3.org/2000/svg"><path d="M1 1L6 6L11 1" stroke="currentColor"></path></svg><!----></button></div> <!--[-1--><!--]--></div><!----> <!--[0--><h1 class="relative group"><a id="batch-embed-a-corpus-with-sagemaker-asynchronous-inference" class="header-link block pr-1.5 text-lg no-hover:hidden with-hover:absolute with-hover:p-1.5 with-hover:opacity-0 with-hover:group-hover:opacity-100 with-hover:right-full" href="#batch-embed-a-corpus-with-sagemaker-asynchronous-inference"><span><svg xmlns="http://www.w3.org/2000/svg" xmlns:xlink="http://www.w3.org/1999/xlink" aria-hidden="true" role="img" width="1em" height="1em" preserveAspectRatio="xMidYMid meet" viewBox="0 0 256 256"><path d="M167.594 88.393a8.001 8.001 0 0 1 0 11.314l-67.882 67.882a8 8 0 1 1-11.314-11.315l67.882-67.881a8.003 8.003 0 0 1 11.314 0zm-28.287 84.86l-28.284 28.284a40 40 0 0 1-56.567-56.567l28.284-28.284a8 8 0 0 0-11.315-11.315l-28.284 28.284a56 56 0 0 0 79.196 79.197l28.285-28.285a8 8 0 1 0-11.315-11.314zM212.852 43.14a56.002 56.002 0 0 0-79.196 0l-28.284 28.284a8 8 0 1 0 11.314 11.314l28.284-28.284a40 40 0 0 1 56.568 56.567l-28.285 28.285a8 8 0 0 0 11.315 11.314l28.284-28.284a56.065 56.065 0 0 0 0-79.196z" fill="currentColor"></path></svg><!----></span></a> <span>Batch-embed a corpus with SageMaker asynchronous inference</span></h1><!--]--><!----> <p><small>Last updated 2026-08-27</small></p> <!--[1--><h2 class="relative group"><a id="what-asynchronous-inference-is" class="header-link block pr-1.5 text-lg no-hover:hidden with-hover:absolute with-hover:p-1.5 with-hover:opacity-0 with-hover:group-hover:opacity-100 with-hover:right-full" href="#what-asynchronous-inference-is"><span><svg xmlns="http://www.w3.org/2000/svg" xmlns:xlink="http://www.w3.org/1999/xlink" aria-hidden="true" role="img" width="1em" height="1em" preserveAspectRatio="xMidYMid meet" viewBox="0 0 256 256"><path d="M167.594 88.393a8.001 8.001 0 0 1 0 11.314l-67.882 67.882a8 8 0 1 1-11.314-11.315l67.882-67.881a8.003 8.003 0 0 1 11.314 0zm-28.287 84.86l-28.284 28.284a40 40 0 0 1-56.567-56.567l28.284-28.284a8 8 0 0 0-11.315-11.315l-28.284 28.284a56 56 0 0 0 79.196 79.197l28.285-28.285a8 8 0 1 0-11.315-11.314zM212.852 43.14a56.002 56.002 0 0 0-79.196 0l-28.284 28.284a8 8 0 1 0 11.314 11.314l28.284-28.284a40 40 0 0 1 56.568 56.567l-28.285 28.285a8 8 0 0 0 11.315 11.314l28.284-28.284a56.065 56.065 0 0 0 0-79.196z" fill="currentColor"></path></svg><!----></span></a> <span>What asynchronous inference is</span></h2><!--]--><!----> <p>Most inference is synchronous: a client opens a connection, sends a request,
waits while the model runs, and reads the response on the same connection.
That works well for interactive traffic, but it struggles when the work is
slow, the payloads are large, or requests arrive in unpredictable bursts. The
caller has to hold a connection open for the whole run, and the endpoint has
to be sized for peak load even while it sits idle the rest of the time.</p> <p><strong>Asynchronous inference</strong> decouples the request from the result. The client
uploads the payload to object storage (S3) and sends the endpoint a <em>pointer</em> to it. The endpoint puts the request on an internal queue, processes it when
capacity is available, and writes the result back to S3. The client picks the
result up later by polling S3, or by reacting to a success/error
notification.</p> <p><img src="https://docs.aws.amazon.com/images/sagemaker/latest/dg/images/async-architecture.png" alt="SageMaker asynchronous inference architecture"/></p> <p>That indirection is what makes the pattern useful:</p> <ul><li><strong>Large payloads, long runtimes.</strong> Inputs are read from S3, not from a
size-capped HTTP body, and there is no client connection to time out.</li> <li><strong>Bursty, queue-shaped traffic.</strong> Requests accumulate in the queue and
drain at the endpoint’s own pace, instead of every spike forcing an
immediate scale-out.</li> <li><strong>Scale to zero.</strong> When the queue is empty the endpoint can run zero
instances and cost nothing, then wake up when new requests land.</li></ul> <p>The cost is latency: you trade an immediate answer for throughput and
elasticity. So asynchronous inference is for offline and background work, not
the interactive request path.</p> <!--[1--><h2 class="relative group"><a id="the-use-case-we-picked" class="header-link block pr-1.5 text-lg no-hover:hidden with-hover:absolute with-hover:p-1.5 with-hover:opacity-0 with-hover:group-hover:opacity-100 with-hover:right-full" href="#the-use-case-we-picked"><span><svg xmlns="http://www.w3.org/2000/svg" xmlns:xlink="http://www.w3.org/1999/xlink" aria-hidden="true" role="img" width="1em" height="1em" preserveAspectRatio="xMidYMid meet" viewBox="0 0 256 256"><path d="M167.594 88.393a8.001 8.001 0 0 1 0 11.314l-67.882 67.882a8 8 0 1 1-11.314-11.315l67.882-67.881a8.003 8.003 0 0 1 11.314 0zm-28.287 84.86l-28.284 28.284a40 40 0 0 1-56.567-56.567l28.284-28.284a8 8 0 0 0-11.315-11.315l-28.284 28.284a56 56 0 0 0 79.196 79.197l28.285-28.285a8 8 0 1 0-11.315-11.314zM212.852 43.14a56.002 56.002 0 0 0-79.196 0l-28.284 28.284a8 8 0 1 0 11.314 11.314l28.284-28.284a40 40 0 0 1 56.568 56.567l-28.285 28.285a8 8 0 0 0 11.315 11.314l28.284-28.284a56.065 56.065 0 0 0 0-79.196z" fill="currentColor"></path></svg><!----></span></a> <span>The use case we picked</span></h2><!--]--><!----> <p>To keep things concrete, this tutorial uses asynchronous inference to <strong>embed a text corpus for semantic search</strong>. Embedding a corpus is a textbook
fit: it is a large, one-off batch that runs whenever the content changes, and
nobody is waiting on any single vector. (Embedding the user’s <em>live</em> query is
the opposite shape — small and latency-sensitive — and belongs on a small
real-time endpoint. This notebook is only the offline half.)</p> <p>The corpus is <a href="https://huggingface.co/datasets/sentence-transformers/natural-questions" rel="nofollow"><code>sentence-transformers/natural-questions</code></a>:
real Google search queries paired with the Wikipedia passages that answer
them. We embed the passages through an asynchronous endpoint, then embed a
held-out query the same way and check that its matching passage ranks near
the top. The embedding model is served with <a href="https://huggingface.co/docs/text-embeddings-inference" rel="nofollow">Text Embeddings Inference (TEI)</a>,
Hugging Face’s container for embedding models.</p> <p>References:</p> <ul><li><a href="https://docs.aws.amazon.com/sagemaker/latest/dg/async-inference.html" rel="nofollow">SageMaker asynchronous inference</a></li> <li><a href="https://docs.aws.amazon.com/sagemaker/latest/dg/async-inference-autoscale.html" rel="nofollow">SageMaker async autoscaling</a></li> <li><a href="https://huggingface.co/datasets/sentence-transformers/natural-questions" rel="nofollow">Natural Questions dataset</a></li></ul> <!--[1--><h2 class="relative group"><a id="prerequisites" class="header-link block pr-1.5 text-lg no-hover:hidden with-hover:absolute with-hover:p-1.5 with-hover:opacity-0 with-hover:group-hover:opacity-100 with-hover:right-full" href="#prerequisites"><span><svg xmlns="http://www.w3.org/2000/svg" xmlns:xlink="http://www.w3.org/1999/xlink" aria-hidden="true" role="img" width="1em" height="1em" preserveAspectRatio="xMidYMid meet" viewBox="0 0 256 256"><path d="M167.594 88.393a8.001 8.001 0 0 1 0 11.314l-67.882 67.882a8 8 0 1 1-11.314-11.315l67.882-67.881a8.003 8.003 0 0 1 11.314 0zm-28.287 84.86l-28.284 28.284a40 40 0 0 1-56.567-56.567l28.284-28.284a8 8 0 0 0-11.315-11.315l-28.284 28.284a56 56 0 0 0 79.196 79.197l28.285-28.285a8 8 0 1 0-11.315-11.314zM212.852 43.14a56.002 56.002 0 0 0-79.196 0l-28.284 28.284a8 8 0 1 0 11.314 11.314l28.284-28.284a40 40 0 0 1 56.568 56.567l-28.285 28.285a8 8 0 0 0 11.315 11.314l28.284-28.284a56.065 56.065 0 0 0 0-79.196z" fill="currentColor"></path></svg><!----></span></a> <span>Prerequisites</span></h2><!--]--><!----> <p>Run the next cell before importing the SDK. It installs the SageMaker Python
SDK and <code>datasets</code> into the active kernel.</p> <p>You also need an existing SageMaker execution role with access to SageMaker,
S3, CloudWatch, Application Auto Scaling, and the ECR repository that hosts
the selected serving DLC.</p> <div class="code-block relative "><div class="absolute top-2.5 right-4"><button class="inline-flex items-center relative text-sm focus:text-green-500 cursor-pointer focus:outline-none transition duration-200 ease-in-out opacity-0 mx-0.5 text-gray-600 " title="code excerpt" type="button"><svg xmlns="http://www.w3.org/2000/svg" aria-hidden="true" fill="currentColor" focusable="false" role="img" width="1em" height="1em" preserveAspectRatio="xMidYMid meet" viewBox="0 0 32 32"><path d="M28,10V28H10V10H28m0-2H10a2,2,0,0,0-2,2V28a2,2,0,0,0,2,2H28a2,2,0,0,0,2-2V10a2,2,0,0,0-2-2Z" transform="translate(0)"></path><path d="M4,18H2V4A2,2,0,0,1,4,2H18V4H4Z" transform="translate(0)"></path><rect fill="none" width="32" height="32"></rect></svg><!----> <div class=" absolute pointer-events-none transition-opacity bg-black text-white py-1 px-2 leading-tight rounded font-normal shadow left-1/2 top-full transform -translate-x-1/2 translate-y-2 opacity-0 "><div class="absolute bottom-full left-1/2 transform -translate-x-1/2 w-0 h-0 border-black border-4 border-t-0" style="border-left-color: transparent; border-right-color: transparent;"></div> Copied</div><!----></button><!----></div> <pre class="language-python "><!---->---
category: embeddings
---
%pip install <span class="hljs-string">&quot;sagemaker&gt;=3.0.0&quot;</span> datasets --upgrade --quiet<!----></pre></div><!----> <div class="code-block relative "><div class="absolute top-2.5 right-4"><button class="inline-flex items-center relative text-sm focus:text-green-500 cursor-pointer focus:outline-none transition duration-200 ease-in-out opacity-0 mx-0.5 text-gray-600 " title="code excerpt" type="button"><svg xmlns="http://www.w3.org/2000/svg" aria-hidden="true" fill="currentColor" focusable="false" role="img" width="1em" height="1em" preserveAspectRatio="xMidYMid meet" viewBox="0 0 32 32"><path d="M28,10V28H10V10H28m0-2H10a2,2,0,0,0-2,2V28a2,2,0,0,0,2,2H28a2,2,0,0,0,2-2V10a2,2,0,0,0-2-2Z" transform="translate(0)"></path><path d="M4,18H2V4A2,2,0,0,1,4,2H18V4H4Z" transform="translate(0)"></path><rect fill="none" width="32" height="32"></rect></svg><!----> <div class=" absolute pointer-events-none transition-opacity bg-black text-white py-1 px-2 leading-tight rounded font-normal shadow left-1/2 top-full transform -translate-x-1/2 translate-y-2 opacity-0 "><div class="absolute bottom-full left-1/2 transform -translate-x-1/2 w-0 h-0 border-black border-4 border-t-0" style="border-left-color: transparent; border-right-color: transparent;"></div> Copied</div><!----></button><!----></div> <pre class="language-python "><!----><span class="hljs-keyword">import</span> datetime <span class="hljs-keyword">as</span> dt
<span class="hljs-keyword">import</span> json
<span class="hljs-keyword">import</span> math
<span class="hljs-keyword">import</span> os
<span class="hljs-keyword">import</span> time
<span class="hljs-keyword">import</span> uuid
<span class="hljs-keyword">from</span> urllib.parse <span class="hljs-keyword">import</span> urlparse
<span class="hljs-keyword">import</span> boto3
<span class="hljs-keyword">from</span> botocore.exceptions <span class="hljs-keyword">import</span> ClientError
<span class="hljs-keyword">from</span> datasets <span class="hljs-keyword">import</span> load_dataset
<span class="hljs-keyword">from</span> sagemaker.core <span class="hljs-keyword">import</span> image_uris
<span class="hljs-keyword">from</span> sagemaker.core.helper.session_helper <span class="hljs-keyword">import</span> Session, get_execution_role
<span class="hljs-keyword">from</span> sagemaker.core.inference_config <span class="hljs-keyword">import</span> AsyncInferenceConfig
<span class="hljs-keyword">from</span> sagemaker.serve <span class="hljs-keyword">import</span> ModelBuilder, ModelServer
<span class="hljs-keyword">from</span> sagemaker.serve.builder.schema_builder <span class="hljs-keyword">import</span> SchemaBuilder<!----></pre></div><!----> <!--[1--><h2 class="relative group"><a id="the-corpus-and-the-embedding-model" class="header-link block pr-1.5 text-lg no-hover:hidden with-hover:absolute with-hover:p-1.5 with-hover:opacity-0 with-hover:group-hover:opacity-100 with-hover:right-full" href="#the-corpus-and-the-embedding-model"><span><svg xmlns="http://www.w3.org/2000/svg" xmlns:xlink="http://www.w3.org/1999/xlink" aria-hidden="true" role="img" width="1em" height="1em" preserveAspectRatio="xMidYMid meet" viewBox="0 0 256 256"><path d="M167.594 88.393a8.001 8.001 0 0 1 0 11.314l-67.882 67.882a8 8 0 1 1-11.314-11.315l67.882-67.881a8.003 8.003 0 0 1 11.314 0zm-28.287 84.86l-28.284 28.284a40 40 0 0 1-56.567-56.567l28.284-28.284a8 8 0 0 0-11.315-11.315l-28.284 28.284a56 56 0 0 0 79.196 79.197l28.285-28.285a8 8 0 1 0-11.315-11.314zM212.852 43.14a56.002 56.002 0 0 0-79.196 0l-28.284 28.284a8 8 0 1 0 11.314 11.314l28.284-28.284a40 40 0 0 1 56.568 56.567l-28.285 28.285a8 8 0 0 0 11.315 11.314l28.284-28.284a56.065 56.065 0 0 0 0-79.196z" fill="currentColor"></path></svg><!----></span></a> <span>The corpus and the embedding model</span></h2><!--]--><!----> <p>Each record in <a href="https://huggingface.co/datasets/sentence-transformers/natural-questions" rel="nofollow"><code>sentence-transformers/natural-questions</code></a> is a search query paired with a Wikipedia passage that answers it. We embed
the passages to build the search index and keep the queries to test retrieval
afterwards. The next cell loads a small slice; raise <code>DATASET_SIZE</code> to index
more.</p> <p>The embedding model is <code>BAAI/bge-small-en-v1.5</code>, which produces
384-dimensional vectors and runs on a CPU instance. To use a different model,
set <code>HF_MODEL_ID</code> and set <code>EMBEDDING_DIM</code> to its output dimension.</p> <div class="code-block relative "><div class="absolute top-2.5 right-4"><button class="inline-flex items-center relative text-sm focus:text-green-500 cursor-pointer focus:outline-none transition duration-200 ease-in-out opacity-0 mx-0.5 text-gray-600 " title="code excerpt" type="button"><svg xmlns="http://www.w3.org/2000/svg" aria-hidden="true" fill="currentColor" focusable="false" role="img" width="1em" height="1em" preserveAspectRatio="xMidYMid meet" viewBox="0 0 32 32"><path d="M28,10V28H10V10H28m0-2H10a2,2,0,0,0-2,2V28a2,2,0,0,0,2,2H28a2,2,0,0,0,2-2V10a2,2,0,0,0-2-2Z" transform="translate(0)"></path><path d="M4,18H2V4A2,2,0,0,1,4,2H18V4H4Z" transform="translate(0)"></path><rect fill="none" width="32" height="32"></rect></svg><!----> <div class=" absolute pointer-events-none transition-opacity bg-black text-white py-1 px-2 leading-tight rounded font-normal shadow left-1/2 top-full transform -translate-x-1/2 translate-y-2 opacity-0 "><div class="absolute bottom-full left-1/2 transform -translate-x-1/2 w-0 h-0 border-black border-4 border-t-0" style="border-left-color: transparent; border-right-color: transparent;"></div> Copied</div><!----></button><!----></div> <pre class="language-python "><!---->PROJECT = <span class="hljs-string">&quot;hf-async-rag&quot;</span>
RUN_ID = dt.datetime.now(dt.timezone.utc).strftime(<span class="hljs-string">&quot;%Y%m%d%H%M%S&quot;</span>)
DATASET_ID = <span class="hljs-string">&quot;sentence-transformers/natural-questions&quot;</span>
DATASET_SIZE = <span class="hljs-built_in">int</span>(os.getenv(<span class="hljs-string">&quot;DATASET_SIZE&quot;</span>, <span class="hljs-string">&quot;96&quot;</span>))
TEXTS_PER_INVOCATION = <span class="hljs-built_in">int</span>(os.getenv(<span class="hljs-string">&quot;TEXTS_PER_INVOCATION&quot;</span>, <span class="hljs-string">&quot;8&quot;</span>))
MAX_CONTEXT_CHARS = <span class="hljs-built_in">int</span>(os.getenv(<span class="hljs-string">&quot;MAX_CONTEXT_CHARS&quot;</span>, <span class="hljs-string">&quot;1200&quot;</span>))
MODEL_ID = os.getenv(<span class="hljs-string">&quot;HF_MODEL_ID&quot;</span>, <span class="hljs-string">&quot;BAAI/bge-small-en-v1.5&quot;</span>)
EMBEDDING_DIM = <span class="hljs-built_in">int</span>(os.getenv(<span class="hljs-string">&quot;EMBEDDING_DIM&quot;</span>, <span class="hljs-string">&quot;384&quot;</span>))
TEI_VERSION = os.getenv(<span class="hljs-string">&quot;TEI_VERSION&quot;</span>, <span class="hljs-string">&quot;1.8.2&quot;</span>)
INSTANCE_TYPE = os.getenv(<span class="hljs-string">&quot;SAGEMAKER_INSTANCE_TYPE&quot;</span>, <span class="hljs-string">&quot;ml.c6i.xlarge&quot;</span>)
ENDPOINT_NAME = os.getenv(<span class="hljs-string">&quot;SAGEMAKER_ENDPOINT_NAME&quot;</span>, <span class="hljs-string">f&quot;<span class="hljs-subst">{PROJECT}</span>-<span class="hljs-subst">{RUN_ID}</span>&quot;</span>)
MAX_INSTANCE_COUNT = <span class="hljs-built_in">int</span>(os.getenv(<span class="hljs-string">&quot;MAX_INSTANCE_COUNT&quot;</span>, <span class="hljs-string">&quot;4&quot;</span>))
BACKLOG_PER_INSTANCE_TARGET = <span class="hljs-built_in">float</span>(os.getenv(<span class="hljs-string">&quot;BACKLOG_PER_INSTANCE_TARGET&quot;</span>, <span class="hljs-string">&quot;5&quot;</span>))
MAX_CONCURRENT_INVOCATIONS_PER_INSTANCE = <span class="hljs-built_in">int</span>(os.getenv(<span class="hljs-string">&quot;MAX_CONCURRENT_INVOCATIONS_PER_INSTANCE&quot;</span>, <span class="hljs-string">&quot;4&quot;</span>))
SUCCESS_SNS_TOPIC_ARN = os.getenv(<span class="hljs-string">&quot;SUCCESS_SNS_TOPIC_ARN&quot;</span>)
ERROR_SNS_TOPIC_ARN = os.getenv(<span class="hljs-string">&quot;ERROR_SNS_TOPIC_ARN&quot;</span>)
ALARM_SNS_TOPIC_ARN = os.getenv(<span class="hljs-string">&quot;ALARM_SNS_TOPIC_ARN&quot;</span>)
<span class="hljs-comment"># Keep cleanup on when running this file as a script. Set CLEANUP=false if you</span>
<span class="hljs-comment"># want to inspect the endpoint after the tutorial finishes.</span>
CLEANUP = os.getenv(<span class="hljs-string">&quot;CLEANUP&quot;</span>, <span class="hljs-string">&quot;true&quot;</span>).lower() <span class="hljs-keyword">not</span> <span class="hljs-keyword">in</span> {<span class="hljs-string">&quot;0&quot;</span>, <span class="hljs-string">&quot;false&quot;</span>, <span class="hljs-string">&quot;no&quot;</span>}<!----></pre></div><!----> <!--[1--><h2 class="relative group"><a id="set-up-the-sagemaker-session" class="header-link block pr-1.5 text-lg no-hover:hidden with-hover:absolute with-hover:p-1.5 with-hover:opacity-0 with-hover:group-hover:opacity-100 with-hover:right-full" href="#set-up-the-sagemaker-session"><span><svg xmlns="http://www.w3.org/2000/svg" xmlns:xlink="http://www.w3.org/1999/xlink" aria-hidden="true" role="img" width="1em" height="1em" preserveAspectRatio="xMidYMid meet" viewBox="0 0 256 256"><path d="M167.594 88.393a8.001 8.001 0 0 1 0 11.314l-67.882 67.882a8 8 0 1 1-11.314-11.315l67.882-67.881a8.003 8.003 0 0 1 11.314 0zm-28.287 84.86l-28.284 28.284a40 40 0 0 1-56.567-56.567l28.284-28.284a8 8 0 0 0-11.315-11.315l-28.284 28.284a56 56 0 0 0 79.196 79.197l28.285-28.285a8 8 0 1 0-11.315-11.314zM212.852 43.14a56.002 56.002 0 0 0-79.196 0l-28.284 28.284a8 8 0 1 0 11.314 11.314l28.284-28.284a40 40 0 0 1 56.568 56.567l-28.285 28.285a8 8 0 0 0 11.315 11.314l28.284-28.284a56.065 56.065 0 0 0 0-79.196z" fill="currentColor"></path></svg><!----></span></a> <span>Set up the SageMaker session</span></h2><!--]--><!----> <p>The endpoint runs under a SageMaker execution role: an IAM role that grants
access to S3, ECR, and CloudWatch. Set <code>SAGEMAKER_EXECUTION_ROLE_ARN</code> to the
role you want to use, or <code>SAGEMAKER_EXECUTION_ROLE_NAME</code> if you only have its
name. Inside SageMaker Studio or a notebook instance you can leave both unset
and the role is detected automatically.</p> <div class="code-block relative "><div class="absolute top-2.5 right-4"><button class="inline-flex items-center relative text-sm focus:text-green-500 cursor-pointer focus:outline-none transition duration-200 ease-in-out opacity-0 mx-0.5 text-gray-600 " title="code excerpt" type="button"><svg xmlns="http://www.w3.org/2000/svg" aria-hidden="true" fill="currentColor" focusable="false" role="img" width="1em" height="1em" preserveAspectRatio="xMidYMid meet" viewBox="0 0 32 32"><path d="M28,10V28H10V10H28m0-2H10a2,2,0,0,0-2,2V28a2,2,0,0,0,2,2H28a2,2,0,0,0,2-2V10a2,2,0,0,0-2-2Z" transform="translate(0)"></path><path d="M4,18H2V4A2,2,0,0,1,4,2H18V4H4Z" transform="translate(0)"></path><rect fill="none" width="32" height="32"></rect></svg><!----> <div class=" absolute pointer-events-none transition-opacity bg-black text-white py-1 px-2 leading-tight rounded font-normal shadow left-1/2 top-full transform -translate-x-1/2 translate-y-2 opacity-0 "><div class="absolute bottom-full left-1/2 transform -translate-x-1/2 w-0 h-0 border-black border-4 border-t-0" style="border-left-color: transparent; border-right-color: transparent;"></div> Copied</div><!----></button><!----></div> <pre class="language-python "><!---->requested_region = os.getenv(<span class="hljs-string">&quot;AWS_REGION&quot;</span>) <span class="hljs-keyword">or</span> os.getenv(<span class="hljs-string">&quot;AWS_DEFAULT_REGION&quot;</span>)
boto_session = boto3.Session(region_name=requested_region) <span class="hljs-keyword">if</span> requested_region <span class="hljs-keyword">else</span> boto3.Session()
sess = Session(boto_session=boto_session)
region = sess.boto_region_name
s3 = boto_session.client(<span class="hljs-string">&quot;s3&quot;</span>)
sm = boto_session.client(<span class="hljs-string">&quot;sagemaker&quot;</span>)
logs = boto_session.client(<span class="hljs-string">&quot;logs&quot;</span>)
cloudwatch = boto_session.client(<span class="hljs-string">&quot;cloudwatch&quot;</span>)
autoscaling = boto_session.client(<span class="hljs-string">&quot;application-autoscaling&quot;</span>)
<span class="hljs-keyword">def</span> <span class="hljs-title function_">resolve_role</span>(<span class="hljs-params">session, sagemaker_session</span>):
role_arn = os.getenv(<span class="hljs-string">&quot;SAGEMAKER_EXECUTION_ROLE_ARN&quot;</span>)
<span class="hljs-keyword">if</span> role_arn:
<span class="hljs-keyword">return</span> role_arn
role_name = os.getenv(<span class="hljs-string">&quot;SAGEMAKER_EXECUTION_ROLE_NAME&quot;</span>)
<span class="hljs-keyword">if</span> role_name:
iam = session.client(<span class="hljs-string">&quot;iam&quot;</span>)
<span class="hljs-keyword">return</span> iam.get_role(RoleName=role_name)[<span class="hljs-string">&quot;Role&quot;</span>][<span class="hljs-string">&quot;Arn&quot;</span>]
<span class="hljs-keyword">return</span> get_execution_role(sagemaker_session=sagemaker_session)
role = resolve_role(boto_session, sess)
bucket = sess.default_bucket()
base_s3_uri = os.getenv(<span class="hljs-string">&quot;SAGEMAKER_ASYNC_BASE_S3_URI&quot;</span>, <span class="hljs-string">f&quot;s3://<span class="hljs-subst">{bucket}</span>/<span class="hljs-subst">{PROJECT}</span>/<span class="hljs-subst">{RUN_ID}</span>&quot;</span>)
input_s3_prefix = <span class="hljs-string">f&quot;<span class="hljs-subst">{base_s3_uri.rstrip(<span class="hljs-string">&#x27;/&#x27;</span>)}</span>/input&quot;</span>
output_s3_prefix = <span class="hljs-string">f&quot;<span class="hljs-subst">{base_s3_uri.rstrip(<span class="hljs-string">&#x27;/&#x27;</span>)}</span>/output&quot;</span>
failure_s3_prefix = <span class="hljs-string">f&quot;<span class="hljs-subst">{base_s3_uri.rstrip(<span class="hljs-string">&#x27;/&#x27;</span>)}</span>/failure&quot;</span>
index_s3_uri = <span class="hljs-string">f&quot;<span class="hljs-subst">{base_s3_uri.rstrip(<span class="hljs-string">&#x27;/&#x27;</span>)}</span>/index/documents-with-embeddings.jsonl&quot;</span>
<span class="hljs-built_in">print</span>(<span class="hljs-string">f&quot;region: <span class="hljs-subst">{region}</span>&quot;</span>)
<span class="hljs-built_in">print</span>(<span class="hljs-string">f&quot;role: <span class="hljs-subst">{role}</span>&quot;</span>)
<span class="hljs-built_in">print</span>(<span class="hljs-string">f&quot;endpoint: <span class="hljs-subst">{ENDPOINT_NAME}</span>&quot;</span>)
<span class="hljs-built_in">print</span>(<span class="hljs-string">f&quot;async input: <span class="hljs-subst">{input_s3_prefix}</span>&quot;</span>)
<span class="hljs-built_in">print</span>(<span class="hljs-string">f&quot;async output: <span class="hljs-subst">{output_s3_prefix}</span>&quot;</span>)<!----></pre></div><!----> <!--[1--><h2 class="relative group"><a id="load-the-corpus" class="header-link block pr-1.5 text-lg no-hover:hidden with-hover:absolute with-hover:p-1.5 with-hover:opacity-0 with-hover:group-hover:opacity-100 with-hover:right-full" href="#load-the-corpus"><span><svg xmlns="http://www.w3.org/2000/svg" xmlns:xlink="http://www.w3.org/1999/xlink" aria-hidden="true" role="img" width="1em" height="1em" preserveAspectRatio="xMidYMid meet" viewBox="0 0 256 256"><path d="M167.594 88.393a8.001 8.001 0 0 1 0 11.314l-67.882 67.882a8 8 0 1 1-11.314-11.315l67.882-67.881a8.003 8.003 0 0 1 11.314 0zm-28.287 84.86l-28.284 28.284a40 40 0 0 1-56.567-56.567l28.284-28.284a8 8 0 0 0-11.315-11.315l-28.284 28.284a56 56 0 0 0 79.196 79.197l28.285-28.285a8 8 0 1 0-11.315-11.314zM212.852 43.14a56.002 56.002 0 0 0-79.196 0l-28.284 28.284a8 8 0 1 0 11.314 11.314l28.284-28.284a40 40 0 0 1 56.568 56.567l-28.285 28.285a8 8 0 0 0 11.315 11.314l28.284-28.284a56.065 56.065 0 0 0 0-79.196z" fill="currentColor"></path></svg><!----></span></a> <span>Load the corpus</span></h2><!--]--><!----> <p>Load the slice and reshape each row into a record with a stable <code>id</code>, its
query, and the passage (truncated to <code>MAX_CONTEXT_CHARS</code>). The passages are
what we embed; the first record’s query is set aside for the retrieval test at
the end.</p> <div class="code-block relative "><div class="absolute top-2.5 right-4"><button class="inline-flex items-center relative text-sm focus:text-green-500 cursor-pointer focus:outline-none transition duration-200 ease-in-out opacity-0 mx-0.5 text-gray-600 " title="code excerpt" type="button"><svg xmlns="http://www.w3.org/2000/svg" aria-hidden="true" fill="currentColor" focusable="false" role="img" width="1em" height="1em" preserveAspectRatio="xMidYMid meet" viewBox="0 0 32 32"><path d="M28,10V28H10V10H28m0-2H10a2,2,0,0,0-2,2V28a2,2,0,0,0,2,2H28a2,2,0,0,0,2-2V10a2,2,0,0,0-2-2Z" transform="translate(0)"></path><path d="M4,18H2V4A2,2,0,0,1,4,2H18V4H4Z" transform="translate(0)"></path><rect fill="none" width="32" height="32"></rect></svg><!----> <div class=" absolute pointer-events-none transition-opacity bg-black text-white py-1 px-2 leading-tight rounded font-normal shadow left-1/2 top-full transform -translate-x-1/2 translate-y-2 opacity-0 "><div class="absolute bottom-full left-1/2 transform -translate-x-1/2 w-0 h-0 border-black border-4 border-t-0" style="border-left-color: transparent; border-right-color: transparent;"></div> Copied</div><!----></button><!----></div> <pre class="language-python "><!---->raw_dataset = load_dataset(DATASET_ID, <span class="hljs-string">&quot;pair&quot;</span>, split=<span class="hljs-string">f&quot;train[:<span class="hljs-subst">{DATASET_SIZE}</span>]&quot;</span>)
records = []
<span class="hljs-keyword">for</span> row_index, row <span class="hljs-keyword">in</span> <span class="hljs-built_in">enumerate</span>(raw_dataset):
passage = row[<span class="hljs-string">&quot;answer&quot;</span>].strip()
query = row[<span class="hljs-string">&quot;query&quot;</span>].strip()
<span class="hljs-keyword">if</span> <span class="hljs-keyword">not</span> passage <span class="hljs-keyword">or</span> <span class="hljs-keyword">not</span> query:
<span class="hljs-keyword">continue</span>
records.append(
{
<span class="hljs-string">&quot;id&quot;</span>: <span class="hljs-string">f&quot;nq-<span class="hljs-subst">{row_index:05d}</span>&quot;</span>,
<span class="hljs-string">&quot;question&quot;</span>: query,
<span class="hljs-string">&quot;context&quot;</span>: passage[:MAX_CONTEXT_CHARS],
}
)
<span class="hljs-keyword">if</span> <span class="hljs-built_in">len</span>(records) &lt; <span class="hljs-number">2</span>:
<span class="hljs-keyword">raise</span> ValueError(<span class="hljs-string">&quot;Need at least two usable records for the retrieval smoke test.&quot;</span>)
query_record = records[<span class="hljs-number">0</span>]
document_records = records
<span class="hljs-built_in">print</span>(<span class="hljs-string">f&quot;loaded records: <span class="hljs-subst">{<span class="hljs-built_in">len</span>(document_records)}</span>&quot;</span>)
<span class="hljs-built_in">print</span>(<span class="hljs-string">f&quot;held-out query: <span class="hljs-subst">{query_record[<span class="hljs-string">&#x27;question&#x27;</span>]}</span>&quot;</span>)<!----></pre></div><!----> <!--[1--><h2 class="relative group"><a id="select-the-tei-serving-container" class="header-link block pr-1.5 text-lg no-hover:hidden with-hover:absolute with-hover:p-1.5 with-hover:opacity-0 with-hover:group-hover:opacity-100 with-hover:right-full" href="#select-the-tei-serving-container"><span><svg xmlns="http://www.w3.org/2000/svg" xmlns:xlink="http://www.w3.org/1999/xlink" aria-hidden="true" role="img" width="1em" height="1em" preserveAspectRatio="xMidYMid meet" viewBox="0 0 256 256"><path d="M167.594 88.393a8.001 8.001 0 0 1 0 11.314l-67.882 67.882a8 8 0 1 1-11.314-11.315l67.882-67.881a8.003 8.003 0 0 1 11.314 0zm-28.287 84.86l-28.284 28.284a40 40 0 0 1-56.567-56.567l28.284-28.284a8 8 0 0 0-11.315-11.315l-28.284 28.284a56 56 0 0 0 79.196 79.197l28.285-28.285a8 8 0 1 0-11.315-11.314zM212.852 43.14a56.002 56.002 0 0 0-79.196 0l-28.284 28.284a8 8 0 1 0 11.314 11.314l28.284-28.284a40 40 0 0 1 56.568 56.567l-28.285 28.285a8 8 0 0 0 11.315 11.314l28.284-28.284a56.065 56.065 0 0 0 0-79.196z" fill="currentColor"></path></svg><!----></span></a> <span>Select the TEI serving container</span></h2><!--]--><!----> <p>Asynchronous inference does not prescribe a model server. For an embedding
workload the Hugging Face choice is Text Embeddings Inference (TEI): CPU
instances use <code>huggingface-tei-cpu</code>, GPU instances use <code>huggingface-tei</code>. The
helper below picks the right one for the configured instance type.</p> <div class="code-block relative "><div class="absolute top-2.5 right-4"><button class="inline-flex items-center relative text-sm focus:text-green-500 cursor-pointer focus:outline-none transition duration-200 ease-in-out opacity-0 mx-0.5 text-gray-600 " title="code excerpt" type="button"><svg xmlns="http://www.w3.org/2000/svg" aria-hidden="true" fill="currentColor" focusable="false" role="img" width="1em" height="1em" preserveAspectRatio="xMidYMid meet" viewBox="0 0 32 32"><path d="M28,10V28H10V10H28m0-2H10a2,2,0,0,0-2,2V28a2,2,0,0,0,2,2H28a2,2,0,0,0,2-2V10a2,2,0,0,0-2-2Z" transform="translate(0)"></path><path d="M4,18H2V4A2,2,0,0,1,4,2H18V4H4Z" transform="translate(0)"></path><rect fill="none" width="32" height="32"></rect></svg><!----> <div class=" absolute pointer-events-none transition-opacity bg-black text-white py-1 px-2 leading-tight rounded font-normal shadow left-1/2 top-full transform -translate-x-1/2 translate-y-2 opacity-0 "><div class="absolute bottom-full left-1/2 transform -translate-x-1/2 w-0 h-0 border-black border-4 border-t-0" style="border-left-color: transparent; border-right-color: transparent;"></div> Copied</div><!----></button><!----></div> <pre class="language-python "><!----><span class="hljs-keyword">def</span> <span class="hljs-title function_">is_gpu_instance</span>(<span class="hljs-params">instance_type</span>):
<span class="hljs-keyword">return</span> instance_type.startswith((<span class="hljs-string">&quot;ml.g&quot;</span>, <span class="hljs-string">&quot;ml.p&quot;</span>))
<span class="hljs-keyword">def</span> <span class="hljs-title function_">get_tei_image_uri</span>(<span class="hljs-params">instance_type</span>):
framework = <span class="hljs-string">&quot;huggingface-tei&quot;</span> <span class="hljs-keyword">if</span> is_gpu_instance(instance_type) <span class="hljs-keyword">else</span> <span class="hljs-string">&quot;huggingface-tei-cpu&quot;</span>
<span class="hljs-keyword">return</span> image_uris.retrieve(
framework=framework,
region=region,
version=TEI_VERSION,
image_scope=<span class="hljs-string">&quot;inference&quot;</span>,
instance_type=instance_type,
)
image_uri = get_tei_image_uri(INSTANCE_TYPE)
<span class="hljs-built_in">print</span>(image_uri)<!----></pre></div><!----> <!--[1--><h2 class="relative group"><a id="build-the-model" class="header-link block pr-1.5 text-lg no-hover:hidden with-hover:absolute with-hover:p-1.5 with-hover:opacity-0 with-hover:group-hover:opacity-100 with-hover:right-full" href="#build-the-model"><span><svg xmlns="http://www.w3.org/2000/svg" xmlns:xlink="http://www.w3.org/1999/xlink" aria-hidden="true" role="img" width="1em" height="1em" preserveAspectRatio="xMidYMid meet" viewBox="0 0 256 256"><path d="M167.594 88.393a8.001 8.001 0 0 1 0 11.314l-67.882 67.882a8 8 0 1 1-11.314-11.315l67.882-67.881a8.003 8.003 0 0 1 11.314 0zm-28.287 84.86l-28.284 28.284a40 40 0 0 1-56.567-56.567l28.284-28.284a8 8 0 0 0-11.315-11.315l-28.284 28.284a56 56 0 0 0 79.196 79.197l28.285-28.285a8 8 0 1 0-11.315-11.314zM212.852 43.14a56.002 56.002 0 0 0-79.196 0l-28.284 28.284a8 8 0 1 0 11.314 11.314l28.284-28.284a40 40 0 0 1 56.568 56.567l-28.285 28.285a8 8 0 0 0 11.315 11.314l28.284-28.284a56.065 56.065 0 0 0 0-79.196z" fill="currentColor"></path></svg><!----></span></a> <span>Build the model</span></h2><!--]--><!----> <p><code>ModelBuilder</code> describes the SageMaker model: the Hub model ID, the serving
container, the model server, and a small input/output example. The container
downloads the model from the Hub when the endpoint starts. For gated or
private models, set <code>HF_TOKEN</code> (or <code>HUGGING_FACE_HUB_TOKEN</code>) before running
the notebook.</p> <div class="code-block relative "><div class="absolute top-2.5 right-4"><button class="inline-flex items-center relative text-sm focus:text-green-500 cursor-pointer focus:outline-none transition duration-200 ease-in-out opacity-0 mx-0.5 text-gray-600 " title="code excerpt" type="button"><svg xmlns="http://www.w3.org/2000/svg" aria-hidden="true" fill="currentColor" focusable="false" role="img" width="1em" height="1em" preserveAspectRatio="xMidYMid meet" viewBox="0 0 32 32"><path d="M28,10V28H10V10H28m0-2H10a2,2,0,0,0-2,2V28a2,2,0,0,0,2,2H28a2,2,0,0,0,2-2V10a2,2,0,0,0-2-2Z" transform="translate(0)"></path><path d="M4,18H2V4A2,2,0,0,1,4,2H18V4H4Z" transform="translate(0)"></path><rect fill="none" width="32" height="32"></rect></svg><!----> <div class=" absolute pointer-events-none transition-opacity bg-black text-white py-1 px-2 leading-tight rounded font-normal shadow left-1/2 top-full transform -translate-x-1/2 translate-y-2 opacity-0 "><div class="absolute bottom-full left-1/2 transform -translate-x-1/2 w-0 h-0 border-black border-4 border-t-0" style="border-left-color: transparent; border-right-color: transparent;"></div> Copied</div><!----></button><!----></div> <pre class="language-python "><!---->hf_token = os.getenv(<span class="hljs-string">&quot;HF_TOKEN&quot;</span>) <span class="hljs-keyword">or</span> os.getenv(<span class="hljs-string">&quot;HUGGING_FACE_HUB_TOKEN&quot;</span>)
env_vars = {
<span class="hljs-string">&quot;HF_MODEL_ID&quot;</span>: MODEL_ID,
<span class="hljs-string">&quot;MAX_BATCH_TOKENS&quot;</span>: os.getenv(<span class="hljs-string">&quot;MAX_BATCH_TOKENS&quot;</span>, <span class="hljs-string">&quot;16384&quot;</span>),
<span class="hljs-string">&quot;MAX_CLIENT_BATCH_SIZE&quot;</span>: os.getenv(<span class="hljs-string">&quot;MAX_CLIENT_BATCH_SIZE&quot;</span>, <span class="hljs-string">&quot;32&quot;</span>),
}
<span class="hljs-keyword">if</span> hf_token:
env_vars[<span class="hljs-string">&quot;HF_TOKEN&quot;</span>] = hf_token
env_vars[<span class="hljs-string">&quot;HUGGING_FACE_HUB_TOKEN&quot;</span>] = hf_token
resource_tags = [
{<span class="hljs-string">&quot;Key&quot;</span>: <span class="hljs-string">&quot;Project&quot;</span>, <span class="hljs-string">&quot;Value&quot;</span>: PROJECT},
{<span class="hljs-string">&quot;Key&quot;</span>: <span class="hljs-string">&quot;ModelId&quot;</span>, <span class="hljs-string">&quot;Value&quot;</span>: MODEL_ID},
{<span class="hljs-string">&quot;Key&quot;</span>: <span class="hljs-string">&quot;CreatedBy&quot;</span>, <span class="hljs-string">&quot;Value&quot;</span>: <span class="hljs-string">&quot;hf-sagemaker-docs&quot;</span>},
]
model_builder = ModelBuilder(
model=MODEL_ID,
role_arn=role,
sagemaker_session=sess,
instance_type=INSTANCE_TYPE,
image_uri=image_uri,
model_server=ModelServer.TEI,
env_vars=env_vars,
schema_builder=SchemaBuilder(
sample_input={<span class="hljs-string">&quot;inputs&quot;</span>: [<span class="hljs-string">&quot;who wrote the origin of species&quot;</span>]},
sample_output=[[<span class="hljs-number">0.0</span>] * EMBEDDING_DIM],
),
)
tei_model = model_builder.build(model_name=<span class="hljs-string">f&quot;<span class="hljs-subst">{PROJECT}</span>-model-<span class="hljs-subst">{RUN_ID}</span>&quot;</span>)<!----></pre></div><!----> <!--[1--><h2 class="relative group"><a id="deploy-the-asynchronous-endpoint" class="header-link block pr-1.5 text-lg no-hover:hidden with-hover:absolute with-hover:p-1.5 with-hover:opacity-0 with-hover:group-hover:opacity-100 with-hover:right-full" href="#deploy-the-asynchronous-endpoint"><span><svg xmlns="http://www.w3.org/2000/svg" xmlns:xlink="http://www.w3.org/1999/xlink" aria-hidden="true" role="img" width="1em" height="1em" preserveAspectRatio="xMidYMid meet" viewBox="0 0 256 256"><path d="M167.594 88.393a8.001 8.001 0 0 1 0 11.314l-67.882 67.882a8 8 0 1 1-11.314-11.315l67.882-67.881a8.003 8.003 0 0 1 11.314 0zm-28.287 84.86l-28.284 28.284a40 40 0 0 1-56.567-56.567l28.284-28.284a8 8 0 0 0-11.315-11.315l-28.284 28.284a56 56 0 0 0 79.196 79.197l28.285-28.285a8 8 0 1 0-11.315-11.314zM212.852 43.14a56.002 56.002 0 0 0-79.196 0l-28.284 28.284a8 8 0 1 0 11.314 11.314l28.284-28.284a40 40 0 0 1 56.568 56.567l-28.285 28.285a8 8 0 0 0 11.315 11.314l28.284-28.284a56.065 56.065 0 0 0 0-79.196z" fill="currentColor"></path></svg><!----></span></a> <span>Deploy the asynchronous endpoint</span></h2><!--]--><!----> <p>The request body lives in S3, <code>invoke_async</code> sends SageMaker a pointer to that
object, and SageMaker later writes either the response or the failure payload
back to S3. Once <code>AsyncInferenceConfig</code> is attached to the endpoint
configuration, the endpoint accepts async invocations only.</p> <p>We deploy with one instance so the container can start and the first
retrieval test does not wait for scale-out. The autoscaling policy in the next
section lets the same endpoint scale to zero after the queue is empty.</p> <div class="code-block relative "><div class="absolute top-2.5 right-4"><button class="inline-flex items-center relative text-sm focus:text-green-500 cursor-pointer focus:outline-none transition duration-200 ease-in-out opacity-0 mx-0.5 text-gray-600 " title="code excerpt" type="button"><svg xmlns="http://www.w3.org/2000/svg" aria-hidden="true" fill="currentColor" focusable="false" role="img" width="1em" height="1em" preserveAspectRatio="xMidYMid meet" viewBox="0 0 32 32"><path d="M28,10V28H10V10H28m0-2H10a2,2,0,0,0-2,2V28a2,2,0,0,0,2,2H28a2,2,0,0,0,2-2V10a2,2,0,0,0-2-2Z" transform="translate(0)"></path><path d="M4,18H2V4A2,2,0,0,1,4,2H18V4H4Z" transform="translate(0)"></path><rect fill="none" width="32" height="32"></rect></svg><!----> <div class=" absolute pointer-events-none transition-opacity bg-black text-white py-1 px-2 leading-tight rounded font-normal shadow left-1/2 top-full transform -translate-x-1/2 translate-y-2 opacity-0 "><div class="absolute bottom-full left-1/2 transform -translate-x-1/2 w-0 h-0 border-black border-4 border-t-0" style="border-left-color: transparent; border-right-color: transparent;"></div> Copied</div><!----></button><!----></div> <pre class="language-python "><!---->notification_config = {}
<span class="hljs-keyword">if</span> SUCCESS_SNS_TOPIC_ARN:
notification_config[<span class="hljs-string">&quot;SuccessTopic&quot;</span>] = SUCCESS_SNS_TOPIC_ARN
<span class="hljs-keyword">if</span> ERROR_SNS_TOPIC_ARN:
notification_config[<span class="hljs-string">&quot;ErrorTopic&quot;</span>] = ERROR_SNS_TOPIC_ARN
async_config = AsyncInferenceConfig(
output_path=output_s3_prefix,
failure_path=failure_s3_prefix,
max_concurrent_invocations_per_instance=MAX_CONCURRENT_INVOCATIONS_PER_INSTANCE,
notification_config=notification_config <span class="hljs-keyword">or</span> <span class="hljs-literal">None</span>,
)
endpoint = model_builder.deploy(
endpoint_name=ENDPOINT_NAME,
initial_instance_count=<span class="hljs-number">1</span>,
instance_type=INSTANCE_TYPE,
inference_config=async_config,
container_timeout_in_seconds=<span class="hljs-number">900</span>,
tags=resource_tags,
wait=<span class="hljs-literal">True</span>,
)
endpoint_description = sm.describe_endpoint(EndpointName=ENDPOINT_NAME)
endpoint_config_name = endpoint_description[<span class="hljs-string">&quot;EndpointConfigName&quot;</span>]
model_name = tei_model.model_name
<span class="hljs-built_in">print</span>(<span class="hljs-string">f&quot;endpoint status: <span class="hljs-subst">{endpoint_description[<span class="hljs-string">&#x27;EndpointStatus&#x27;</span>]}</span>&quot;</span>)
<span class="hljs-built_in">print</span>(<span class="hljs-string">f&quot;endpoint config: <span class="hljs-subst">{endpoint_config_name}</span>&quot;</span>)
<span class="hljs-built_in">print</span>(<span class="hljs-string">f&quot;model: <span class="hljs-subst">{model_name}</span>&quot;</span>)<!----></pre></div><!----> <!--[1--><h2 class="relative group"><a id="autoscale-and-scale-to-zero" class="header-link block pr-1.5 text-lg no-hover:hidden with-hover:absolute with-hover:p-1.5 with-hover:opacity-0 with-hover:group-hover:opacity-100 with-hover:right-full" href="#autoscale-and-scale-to-zero"><span><svg xmlns="http://www.w3.org/2000/svg" xmlns:xlink="http://www.w3.org/1999/xlink" aria-hidden="true" role="img" width="1em" height="1em" preserveAspectRatio="xMidYMid meet" viewBox="0 0 256 256"><path d="M167.594 88.393a8.001 8.001 0 0 1 0 11.314l-67.882 67.882a8 8 0 1 1-11.314-11.315l67.882-67.881a8.003 8.003 0 0 1 11.314 0zm-28.287 84.86l-28.284 28.284a40 40 0 0 1-56.567-56.567l28.284-28.284a8 8 0 0 0-11.315-11.315l-28.284 28.284a56 56 0 0 0 79.196 79.197l28.285-28.285a8 8 0 1 0-11.315-11.314zM212.852 43.14a56.002 56.002 0 0 0-79.196 0l-28.284 28.284a8 8 0 1 0 11.314 11.314l28.284-28.284a40 40 0 0 1 56.568 56.567l-28.285 28.285a8 8 0 0 0 11.315 11.314l28.284-28.284a56.065 56.065 0 0 0 0-79.196z" fill="currentColor"></path></svg><!----></span></a> <span>Autoscale and scale to zero</span></h2><!--]--><!----> <p>Scale-to-zero is one of the reasons asynchronous inference is attractive for
batch workloads. The target-tracking policy scales with queue depth, while the
step-scaling policy wakes the endpoint from zero as soon as a backlog appears.
Without that wake-up alarm, the endpoint might wait until the queue exceeds
the target-tracking threshold before it adds the first instance.</p> <div class="code-block relative "><div class="absolute top-2.5 right-4"><button class="inline-flex items-center relative text-sm focus:text-green-500 cursor-pointer focus:outline-none transition duration-200 ease-in-out opacity-0 mx-0.5 text-gray-600 " title="code excerpt" type="button"><svg xmlns="http://www.w3.org/2000/svg" aria-hidden="true" fill="currentColor" focusable="false" role="img" width="1em" height="1em" preserveAspectRatio="xMidYMid meet" viewBox="0 0 32 32"><path d="M28,10V28H10V10H28m0-2H10a2,2,0,0,0-2,2V28a2,2,0,0,0,2,2H28a2,2,0,0,0,2-2V10a2,2,0,0,0-2-2Z" transform="translate(0)"></path><path d="M4,18H2V4A2,2,0,0,1,4,2H18V4H4Z" transform="translate(0)"></path><rect fill="none" width="32" height="32"></rect></svg><!----> <div class=" absolute pointer-events-none transition-opacity bg-black text-white py-1 px-2 leading-tight rounded font-normal shadow left-1/2 top-full transform -translate-x-1/2 translate-y-2 opacity-0 "><div class="absolute bottom-full left-1/2 transform -translate-x-1/2 w-0 h-0 border-black border-4 border-t-0" style="border-left-color: transparent; border-right-color: transparent;"></div> Copied</div><!----></button><!----></div> <pre class="language-python "><!---->variant_name = <span class="hljs-string">&quot;AllTraffic&quot;</span>
resource_id = <span class="hljs-string">f&quot;endpoint/<span class="hljs-subst">{ENDPOINT_NAME}</span>/variant/<span class="hljs-subst">{variant_name}</span>&quot;</span>
autoscaling.register_scalable_target(
ServiceNamespace=<span class="hljs-string">&quot;sagemaker&quot;</span>,
ResourceId=resource_id,
ScalableDimension=<span class="hljs-string">&quot;sagemaker:variant:DesiredInstanceCount&quot;</span>,
MinCapacity=<span class="hljs-number">0</span>,
MaxCapacity=MAX_INSTANCE_COUNT,
)
autoscaling.put_scaling_policy(
PolicyName=<span class="hljs-string">f&quot;<span class="hljs-subst">{ENDPOINT_NAME}</span>-backlog-target-tracking&quot;</span>,
ServiceNamespace=<span class="hljs-string">&quot;sagemaker&quot;</span>,
ResourceId=resource_id,
ScalableDimension=<span class="hljs-string">&quot;sagemaker:variant:DesiredInstanceCount&quot;</span>,
PolicyType=<span class="hljs-string">&quot;TargetTrackingScaling&quot;</span>,
TargetTrackingScalingPolicyConfiguration={
<span class="hljs-string">&quot;TargetValue&quot;</span>: BACKLOG_PER_INSTANCE_TARGET,
<span class="hljs-string">&quot;CustomizedMetricSpecification&quot;</span>: {
<span class="hljs-string">&quot;MetricName&quot;</span>: <span class="hljs-string">&quot;ApproximateBacklogSizePerInstance&quot;</span>,
<span class="hljs-string">&quot;Namespace&quot;</span>: <span class="hljs-string">&quot;AWS/SageMaker&quot;</span>,
<span class="hljs-string">&quot;Dimensions&quot;</span>: [{<span class="hljs-string">&quot;Name&quot;</span>: <span class="hljs-string">&quot;EndpointName&quot;</span>, <span class="hljs-string">&quot;Value&quot;</span>: ENDPOINT_NAME}],
<span class="hljs-string">&quot;Statistic&quot;</span>: <span class="hljs-string">&quot;Average&quot;</span>,
},
<span class="hljs-string">&quot;ScaleInCooldown&quot;</span>: <span class="hljs-number">300</span>,
<span class="hljs-string">&quot;ScaleOutCooldown&quot;</span>: <span class="hljs-number">60</span>,
},
)
step_policy = autoscaling.put_scaling_policy(
PolicyName=<span class="hljs-string">f&quot;<span class="hljs-subst">{ENDPOINT_NAME}</span>-wake-from-zero&quot;</span>,
ServiceNamespace=<span class="hljs-string">&quot;sagemaker&quot;</span>,
ResourceId=resource_id,
ScalableDimension=<span class="hljs-string">&quot;sagemaker:variant:DesiredInstanceCount&quot;</span>,
PolicyType=<span class="hljs-string">&quot;StepScaling&quot;</span>,
StepScalingPolicyConfiguration={
<span class="hljs-string">&quot;AdjustmentType&quot;</span>: <span class="hljs-string">&quot;ChangeInCapacity&quot;</span>,
<span class="hljs-string">&quot;MetricAggregationType&quot;</span>: <span class="hljs-string">&quot;Average&quot;</span>,
<span class="hljs-string">&quot;Cooldown&quot;</span>: <span class="hljs-number">300</span>,
<span class="hljs-string">&quot;StepAdjustments&quot;</span>: [{<span class="hljs-string">&quot;MetricIntervalLowerBound&quot;</span>: <span class="hljs-number">0</span>, <span class="hljs-string">&quot;ScalingAdjustment&quot;</span>: <span class="hljs-number">1</span>}],
},
)
wake_alarm_name = <span class="hljs-string">f&quot;<span class="hljs-subst">{ENDPOINT_NAME}</span>-has-backlog-without-capacity&quot;</span>
backlog_alarm_name = <span class="hljs-string">f&quot;<span class="hljs-subst">{ENDPOINT_NAME}</span>-async-backlog-high&quot;</span>
failure_alarm_name = <span class="hljs-string">f&quot;<span class="hljs-subst">{ENDPOINT_NAME}</span>-async-failures&quot;</span>
alarm_names = [wake_alarm_name, backlog_alarm_name, failure_alarm_name]
cloudwatch.put_metric_alarm(
AlarmName=wake_alarm_name,
AlarmDescription=<span class="hljs-string">&quot;Wake async endpoint from zero when requests are queued.&quot;</span>,
Namespace=<span class="hljs-string">&quot;AWS/SageMaker&quot;</span>,
MetricName=<span class="hljs-string">&quot;HasBacklogWithoutCapacity&quot;</span>,
Dimensions=[{<span class="hljs-string">&quot;Name&quot;</span>: <span class="hljs-string">&quot;EndpointName&quot;</span>, <span class="hljs-string">&quot;Value&quot;</span>: ENDPOINT_NAME}],
Statistic=<span class="hljs-string">&quot;Average&quot;</span>,
Period=<span class="hljs-number">60</span>,
EvaluationPeriods=<span class="hljs-number">2</span>,
DatapointsToAlarm=<span class="hljs-number">2</span>,
Threshold=<span class="hljs-number">1</span>,
ComparisonOperator=<span class="hljs-string">&quot;GreaterThanOrEqualToThreshold&quot;</span>,
TreatMissingData=<span class="hljs-string">&quot;missing&quot;</span>,
AlarmActions=[step_policy[<span class="hljs-string">&quot;PolicyARN&quot;</span>]],
)
backlog_alarm = {
<span class="hljs-string">&quot;AlarmName&quot;</span>: backlog_alarm_name,
<span class="hljs-string">&quot;AlarmDescription&quot;</span>: <span class="hljs-string">&quot;Async queue is growing faster than the endpoint can drain it.&quot;</span>,
<span class="hljs-string">&quot;Namespace&quot;</span>: <span class="hljs-string">&quot;AWS/SageMaker&quot;</span>,
<span class="hljs-string">&quot;MetricName&quot;</span>: <span class="hljs-string">&quot;ApproximateBacklogSize&quot;</span>,
<span class="hljs-string">&quot;Dimensions&quot;</span>: [{<span class="hljs-string">&quot;Name&quot;</span>: <span class="hljs-string">&quot;EndpointName&quot;</span>, <span class="hljs-string">&quot;Value&quot;</span>: ENDPOINT_NAME}],
<span class="hljs-string">&quot;Statistic&quot;</span>: <span class="hljs-string">&quot;Average&quot;</span>,
<span class="hljs-string">&quot;Period&quot;</span>: <span class="hljs-number">60</span>,
<span class="hljs-string">&quot;EvaluationPeriods&quot;</span>: <span class="hljs-number">3</span>,
<span class="hljs-string">&quot;DatapointsToAlarm&quot;</span>: <span class="hljs-number">2</span>,
<span class="hljs-string">&quot;Threshold&quot;</span>: <span class="hljs-number">50</span>,
<span class="hljs-string">&quot;ComparisonOperator&quot;</span>: <span class="hljs-string">&quot;GreaterThanThreshold&quot;</span>,
<span class="hljs-string">&quot;TreatMissingData&quot;</span>: <span class="hljs-string">&quot;notBreaching&quot;</span>,
}
failure_alarm = {
<span class="hljs-string">&quot;AlarmName&quot;</span>: failure_alarm_name,
<span class="hljs-string">&quot;AlarmDescription&quot;</span>: <span class="hljs-string">&quot;Async inference requests are failing.&quot;</span>,
<span class="hljs-string">&quot;Namespace&quot;</span>: <span class="hljs-string">&quot;AWS/SageMaker&quot;</span>,
<span class="hljs-string">&quot;MetricName&quot;</span>: <span class="hljs-string">&quot;InvocationsFailed&quot;</span>,
<span class="hljs-string">&quot;Dimensions&quot;</span>: [{<span class="hljs-string">&quot;Name&quot;</span>: <span class="hljs-string">&quot;EndpointName&quot;</span>, <span class="hljs-string">&quot;Value&quot;</span>: ENDPOINT_NAME}],
<span class="hljs-string">&quot;Statistic&quot;</span>: <span class="hljs-string">&quot;Sum&quot;</span>,
<span class="hljs-string">&quot;Period&quot;</span>: <span class="hljs-number">60</span>,
<span class="hljs-string">&quot;EvaluationPeriods&quot;</span>: <span class="hljs-number">1</span>,
<span class="hljs-string">&quot;DatapointsToAlarm&quot;</span>: <span class="hljs-number">1</span>,
<span class="hljs-string">&quot;Threshold&quot;</span>: <span class="hljs-number">5</span>,
<span class="hljs-string">&quot;ComparisonOperator&quot;</span>: <span class="hljs-string">&quot;GreaterThanThreshold&quot;</span>,
<span class="hljs-string">&quot;TreatMissingData&quot;</span>: <span class="hljs-string">&quot;notBreaching&quot;</span>,
}
<span class="hljs-keyword">if</span> ALARM_SNS_TOPIC_ARN:
backlog_alarm[<span class="hljs-string">&quot;AlarmActions&quot;</span>] = [ALARM_SNS_TOPIC_ARN]
failure_alarm[<span class="hljs-string">&quot;AlarmActions&quot;</span>] = [ALARM_SNS_TOPIC_ARN]
cloudwatch.put_metric_alarm(**backlog_alarm)
cloudwatch.put_metric_alarm(**failure_alarm)
<span class="hljs-built_in">print</span>(<span class="hljs-string">f&quot;registered scalable target: <span class="hljs-subst">{resource_id}</span>&quot;</span>)
<span class="hljs-built_in">print</span>(<span class="hljs-string">f&quot;alarms: <span class="hljs-subst">{<span class="hljs-string">&#x27;, &#x27;</span>.join(alarm_names)}</span>&quot;</span>)<!----></pre></div><!----> <!--[1--><h2 class="relative group"><a id="submit-the-embedding-requests" class="header-link block pr-1.5 text-lg no-hover:hidden with-hover:absolute with-hover:p-1.5 with-hover:opacity-0 with-hover:group-hover:opacity-100 with-hover:right-full" href="#submit-the-embedding-requests"><span><svg xmlns="http://www.w3.org/2000/svg" xmlns:xlink="http://www.w3.org/1999/xlink" aria-hidden="true" role="img" width="1em" height="1em" preserveAspectRatio="xMidYMid meet" viewBox="0 0 256 256"><path d="M167.594 88.393a8.001 8.001 0 0 1 0 11.314l-67.882 67.882a8 8 0 1 1-11.314-11.315l67.882-67.881a8.003 8.003 0 0 1 11.314 0zm-28.287 84.86l-28.284 28.284a40 40 0 0 1-56.567-56.567l28.284-28.284a8 8 0 0 0-11.315-11.315l-28.284 28.284a56 56 0 0 0 79.196 79.197l28.285-28.285a8 8 0 1 0-11.315-11.314zM212.852 43.14a56.002 56.002 0 0 0-79.196 0l-28.284 28.284a8 8 0 1 0 11.314 11.314l28.284-28.284a40 40 0 0 1 56.568 56.567l-28.285 28.285a8 8 0 0 0 11.315 11.314l28.284-28.284a56.065 56.065 0 0 0 0-79.196z" fill="currentColor"></path></svg><!----></span></a> <span>Submit the embedding requests</span></h2><!--]--><!----> <p>An asynchronous endpoint does not take the request body directly. Each batch
is uploaded to S3 first, and <code>invoke_async</code> sends the endpoint that S3
location instead of the payload. The container still receives an ordinary
embedding request: <code>{"inputs": [...]}</code>.</p> <p><code>invoke_async</code> returns right away with a pointer to where the output will be
written; it does not wait for the vectors.</p> <div class="code-block relative "><div class="absolute top-2.5 right-4"><button class="inline-flex items-center relative text-sm focus:text-green-500 cursor-pointer focus:outline-none transition duration-200 ease-in-out opacity-0 mx-0.5 text-gray-600 " title="code excerpt" type="button"><svg xmlns="http://www.w3.org/2000/svg" aria-hidden="true" fill="currentColor" focusable="false" role="img" width="1em" height="1em" preserveAspectRatio="xMidYMid meet" viewBox="0 0 32 32"><path d="M28,10V28H10V10H28m0-2H10a2,2,0,0,0-2,2V28a2,2,0,0,0,2,2H28a2,2,0,0,0,2-2V10a2,2,0,0,0-2-2Z" transform="translate(0)"></path><path d="M4,18H2V4A2,2,0,0,1,4,2H18V4H4Z" transform="translate(0)"></path><rect fill="none" width="32" height="32"></rect></svg><!----> <div class=" absolute pointer-events-none transition-opacity bg-black text-white py-1 px-2 leading-tight rounded font-normal shadow left-1/2 top-full transform -translate-x-1/2 translate-y-2 opacity-0 "><div class="absolute bottom-full left-1/2 transform -translate-x-1/2 w-0 h-0 border-black border-4 border-t-0" style="border-left-color: transparent; border-right-color: transparent;"></div> Copied</div><!----></button><!----></div> <pre class="language-python "><!----><span class="hljs-keyword">def</span> <span class="hljs-title function_">parse_s3_uri</span>(<span class="hljs-params">uri</span>):
parsed = urlparse(uri)
<span class="hljs-keyword">if</span> parsed.scheme != <span class="hljs-string">&quot;s3&quot;</span> <span class="hljs-keyword">or</span> <span class="hljs-keyword">not</span> parsed.netloc <span class="hljs-keyword">or</span> <span class="hljs-keyword">not</span> parsed.path.strip(<span class="hljs-string">&quot;/&quot;</span>):
<span class="hljs-keyword">raise</span> ValueError(<span class="hljs-string">f&quot;expected S3 URI, got <span class="hljs-subst">{uri!r}</span>&quot;</span>)
<span class="hljs-keyword">return</span> parsed.netloc, parsed.path.lstrip(<span class="hljs-string">&quot;/&quot;</span>)
<span class="hljs-keyword">def</span> <span class="hljs-title function_">s3_join</span>(<span class="hljs-params">prefix, *parts</span>):
<span class="hljs-keyword">return</span> <span class="hljs-string">&quot;/&quot;</span>.join([prefix.rstrip(<span class="hljs-string">&quot;/&quot;</span>), *(part.strip(<span class="hljs-string">&quot;/&quot;</span>) <span class="hljs-keyword">for</span> part <span class="hljs-keyword">in</span> parts)])
<span class="hljs-keyword">def</span> <span class="hljs-title function_">put_json</span>(<span class="hljs-params">uri, payload</span>):
target_bucket, key = parse_s3_uri(uri)
s3.put_object(
Bucket=target_bucket,
Key=key,
Body=json.dumps(payload).encode(<span class="hljs-string">&quot;utf-8&quot;</span>),
ContentType=<span class="hljs-string">&quot;application/json&quot;</span>,
)
<span class="hljs-keyword">def</span> <span class="hljs-title function_">put_text</span>(<span class="hljs-params">uri, text</span>):
target_bucket, key = parse_s3_uri(uri)
s3.put_object(Bucket=target_bucket, Key=key, Body=text.encode(<span class="hljs-string">&quot;utf-8&quot;</span>))
<span class="hljs-keyword">def</span> <span class="hljs-title function_">batched</span>(<span class="hljs-params">items, size</span>):
<span class="hljs-keyword">for</span> start <span class="hljs-keyword">in</span> <span class="hljs-built_in">range</span>(<span class="hljs-number">0</span>, <span class="hljs-built_in">len</span>(items), size):
<span class="hljs-keyword">yield</span> items[start : start + size]
<span class="hljs-keyword">assert</span> parse_s3_uri(<span class="hljs-string">&quot;s3://example-bucket/path/file.json&quot;</span>) == (
<span class="hljs-string">&quot;example-bucket&quot;</span>,
<span class="hljs-string">&quot;path/file.json&quot;</span>,
)
document_jobs = []
<span class="hljs-keyword">for</span> batch_index, batch <span class="hljs-keyword">in</span> <span class="hljs-built_in">enumerate</span>(batched(document_records, TEXTS_PER_INVOCATION), start=<span class="hljs-number">1</span>):
payload = {<span class="hljs-string">&quot;inputs&quot;</span>: [record[<span class="hljs-string">&quot;context&quot;</span>] <span class="hljs-keyword">for</span> record <span class="hljs-keyword">in</span> batch]}
input_uri = s3_join(input_s3_prefix, <span class="hljs-string">&quot;documents&quot;</span>, <span class="hljs-string">f&quot;batch-<span class="hljs-subst">{batch_index:04d}</span>.json&quot;</span>)
put_json(input_uri, payload)
response = endpoint.invoke_async(
input_location=input_uri,
content_type=<span class="hljs-string">&quot;application/json&quot;</span>,
accept=<span class="hljs-string">&quot;application/json&quot;</span>,
inference_id=<span class="hljs-string">f&quot;documents-<span class="hljs-subst">{batch_index:04d}</span>-<span class="hljs-subst">{uuid.uuid4()}</span>&quot;</span>,
invocation_timeout_seconds=<span class="hljs-number">900</span>,
session=boto_session,
region=region,
)
document_jobs.append((batch, response))
<span class="hljs-built_in">print</span>(<span class="hljs-string">f&quot;submitted <span class="hljs-subst">{input_uri}</span> -&gt; <span class="hljs-subst">{response.output_location}</span>&quot;</span>)<!----></pre></div><!----> <!--[1--><h2 class="relative group"><a id="collect-the-embeddings" class="header-link block pr-1.5 text-lg no-hover:hidden with-hover:absolute with-hover:p-1.5 with-hover:opacity-0 with-hover:group-hover:opacity-100 with-hover:right-full" href="#collect-the-embeddings"><span><svg xmlns="http://www.w3.org/2000/svg" xmlns:xlink="http://www.w3.org/1999/xlink" aria-hidden="true" role="img" width="1em" height="1em" preserveAspectRatio="xMidYMid meet" viewBox="0 0 256 256"><path d="M167.594 88.393a8.001 8.001 0 0 1 0 11.314l-67.882 67.882a8 8 0 1 1-11.314-11.315l67.882-67.881a8.003 8.003 0 0 1 11.314 0zm-28.287 84.86l-28.284 28.284a40 40 0 0 1-56.567-56.567l28.284-28.284a8 8 0 0 0-11.315-11.315l-28.284 28.284a56 56 0 0 0 79.196 79.197l28.285-28.285a8 8 0 1 0-11.315-11.314zM212.852 43.14a56.002 56.002 0 0 0-79.196 0l-28.284 28.284a8 8 0 1 0 11.314 11.314l28.284-28.284a40 40 0 0 1 56.568 56.567l-28.285 28.285a8 8 0 0 0 11.315 11.314l28.284-28.284a56.065 56.065 0 0 0 0-79.196z" fill="currentColor"></path></svg><!----></span></a> <span>Collect the embeddings</span></h2><!--]--><!----> <p>The async response objects give us both the output path and the failure path.
Polling S3 keeps the notebook simple, while production pipelines often use
SNS, EventBridge, or a workflow engine to react to completed outputs.</p> <p>Once every batch returns, the vectors are joined back to the metadata that
created them and written as JSON Lines. A downstream indexer could read that
file and push the embeddings into OpenSearch, PostgreSQL with pgvector, a
vector database, or another retrieval store.</p> <div class="code-block relative "><div class="absolute top-2.5 right-4"><button class="inline-flex items-center relative text-sm focus:text-green-500 cursor-pointer focus:outline-none transition duration-200 ease-in-out opacity-0 mx-0.5 text-gray-600 " title="code excerpt" type="button"><svg xmlns="http://www.w3.org/2000/svg" aria-hidden="true" fill="currentColor" focusable="false" role="img" width="1em" height="1em" preserveAspectRatio="xMidYMid meet" viewBox="0 0 32 32"><path d="M28,10V28H10V10H28m0-2H10a2,2,0,0,0-2,2V28a2,2,0,0,0,2,2H28a2,2,0,0,0,2-2V10a2,2,0,0,0-2-2Z" transform="translate(0)"></path><path d="M4,18H2V4A2,2,0,0,1,4,2H18V4H4Z" transform="translate(0)"></path><rect fill="none" width="32" height="32"></rect></svg><!----> <div class=" absolute pointer-events-none transition-opacity bg-black text-white py-1 px-2 leading-tight rounded font-normal shadow left-1/2 top-full transform -translate-x-1/2 translate-y-2 opacity-0 "><div class="absolute bottom-full left-1/2 transform -translate-x-1/2 w-0 h-0 border-black border-4 border-t-0" style="border-left-color: transparent; border-right-color: transparent;"></div> Copied</div><!----></button><!----></div> <pre class="language-python "><!----><span class="hljs-keyword">def</span> <span class="hljs-title function_">is_missing_key_error</span>(<span class="hljs-params">error</span>):
<span class="hljs-keyword">return</span> error.response.get(<span class="hljs-string">&quot;Error&quot;</span>, {}).get(<span class="hljs-string">&quot;Code&quot;</span>) <span class="hljs-keyword">in</span> {<span class="hljs-string">&quot;NoSuchKey&quot;</span>, <span class="hljs-string">&quot;404&quot;</span>, <span class="hljs-string">&quot;NotFound&quot;</span>}
<span class="hljs-keyword">def</span> <span class="hljs-title function_">read_s3_text</span>(<span class="hljs-params">uri</span>):
source_bucket, key = parse_s3_uri(uri)
response = s3.get_object(Bucket=source_bucket, Key=key)
<span class="hljs-keyword">return</span> response[<span class="hljs-string">&quot;Body&quot;</span>].read().decode(<span class="hljs-string">&quot;utf-8&quot;</span>)
<span class="hljs-keyword">def</span> <span class="hljs-title function_">wait_for_async_json</span>(<span class="hljs-params">response, timeout=<span class="hljs-number">1800</span>, poll=<span class="hljs-number">10</span></span>):
deadline = time.time() + timeout
<span class="hljs-keyword">while</span> time.time() &lt; deadline:
<span class="hljs-keyword">for</span> uri, failed <span class="hljs-keyword">in</span> (
(response.output_location, <span class="hljs-literal">False</span>),
(response.failure_location, <span class="hljs-literal">True</span>),
):
<span class="hljs-keyword">if</span> <span class="hljs-keyword">not</span> uri:
<span class="hljs-keyword">continue</span>
<span class="hljs-keyword">try</span>:
body = read_s3_text(uri)
<span class="hljs-keyword">except</span> ClientError <span class="hljs-keyword">as</span> error:
<span class="hljs-keyword">if</span> is_missing_key_error(error):
<span class="hljs-keyword">continue</span>
<span class="hljs-keyword">raise</span>
<span class="hljs-keyword">if</span> failed:
<span class="hljs-keyword">raise</span> RuntimeError(<span class="hljs-string">f&quot;async inference failed: <span class="hljs-subst">{body}</span>&quot;</span>)
<span class="hljs-keyword">return</span> json.loads(body)
<span class="hljs-built_in">print</span>(<span class="hljs-string">&quot;waiting for async output...&quot;</span>)
time.sleep(poll)
<span class="hljs-keyword">raise</span> TimeoutError(<span class="hljs-string">f&quot;no async result after <span class="hljs-subst">{timeout}</span> seconds: <span class="hljs-subst">{response.output_location}</span>&quot;</span>)
index_records = []
<span class="hljs-keyword">for</span> batch, response <span class="hljs-keyword">in</span> document_jobs:
vectors = wait_for_async_json(response)
<span class="hljs-keyword">if</span> <span class="hljs-built_in">len</span>(vectors) != <span class="hljs-built_in">len</span>(batch):
<span class="hljs-keyword">raise</span> RuntimeError(<span class="hljs-string">f&quot;expected <span class="hljs-subst">{<span class="hljs-built_in">len</span>(batch)}</span> vectors, got <span class="hljs-subst">{<span class="hljs-built_in">len</span>(vectors)}</span>&quot;</span>)
<span class="hljs-keyword">for</span> record, vector <span class="hljs-keyword">in</span> <span class="hljs-built_in">zip</span>(batch, vectors):
index_records.append(
{
<span class="hljs-string">&quot;id&quot;</span>: record[<span class="hljs-string">&quot;id&quot;</span>],
<span class="hljs-string">&quot;question&quot;</span>: record[<span class="hljs-string">&quot;question&quot;</span>],
<span class="hljs-string">&quot;context&quot;</span>: record[<span class="hljs-string">&quot;context&quot;</span>],
<span class="hljs-string">&quot;embedding&quot;</span>: vector,
}
)
<span class="hljs-keyword">assert</span> <span class="hljs-built_in">len</span>(index_records) == <span class="hljs-built_in">len</span>(document_records)
<span class="hljs-keyword">assert</span> <span class="hljs-built_in">all</span>(<span class="hljs-built_in">len</span>(record[<span class="hljs-string">&quot;embedding&quot;</span>]) == EMBEDDING_DIM <span class="hljs-keyword">for</span> record <span class="hljs-keyword">in</span> index_records)
put_text(index_s3_uri, <span class="hljs-string">&quot;\n&quot;</span>.join(json.dumps(record) <span class="hljs-keyword">for</span> record <span class="hljs-keyword">in</span> index_records))
<span class="hljs-built_in">print</span>(<span class="hljs-string">f&quot;embedded documents: <span class="hljs-subst">{<span class="hljs-built_in">len</span>(index_records)}</span>&quot;</span>)
<span class="hljs-built_in">print</span>(<span class="hljs-string">f&quot;embedding dimensions: <span class="hljs-subst">{<span class="hljs-built_in">len</span>(index_records[<span class="hljs-number">0</span>][<span class="hljs-string">&#x27;embedding&#x27;</span>])}</span>&quot;</span>)
<span class="hljs-built_in">print</span>(<span class="hljs-string">f&quot;index written to: <span class="hljs-subst">{index_s3_uri}</span>&quot;</span>)<!----></pre></div><!----> <!--[1--><h2 class="relative group"><a id="validate-with-a-retrieval-test" class="header-link block pr-1.5 text-lg no-hover:hidden with-hover:absolute with-hover:p-1.5 with-hover:opacity-0 with-hover:group-hover:opacity-100 with-hover:right-full" href="#validate-with-a-retrieval-test"><span><svg xmlns="http://www.w3.org/2000/svg" xmlns:xlink="http://www.w3.org/1999/xlink" aria-hidden="true" role="img" width="1em" height="1em" preserveAspectRatio="xMidYMid meet" viewBox="0 0 256 256"><path d="M167.594 88.393a8.001 8.001 0 0 1 0 11.314l-67.882 67.882a8 8 0 1 1-11.314-11.315l67.882-67.881a8.003 8.003 0 0 1 11.314 0zm-28.287 84.86l-28.284 28.284a40 40 0 0 1-56.567-56.567l28.284-28.284a8 8 0 0 0-11.315-11.315l-28.284 28.284a56 56 0 0 0 79.196 79.197l28.285-28.285a8 8 0 1 0-11.315-11.314zM212.852 43.14a56.002 56.002 0 0 0-79.196 0l-28.284 28.284a8 8 0 1 0 11.314 11.314l28.284-28.284a40 40 0 0 1 56.568 56.567l-28.285 28.285a8 8 0 0 0 11.315 11.314l28.284-28.284a56.065 56.065 0 0 0 0-79.196z" fill="currentColor"></path></svg><!----></span></a> <span>Validate with a retrieval test</span></h2><!--]--><!----> <p>To check that the embeddings support search, we embed the held-out query
through the same asynchronous endpoint, score it against every passage
embedding with cosine similarity, and rank the passages. The passage that
originally answered the query should come out on top.</p> <div class="code-block relative "><div class="absolute top-2.5 right-4"><button class="inline-flex items-center relative text-sm focus:text-green-500 cursor-pointer focus:outline-none transition duration-200 ease-in-out opacity-0 mx-0.5 text-gray-600 " title="code excerpt" type="button"><svg xmlns="http://www.w3.org/2000/svg" aria-hidden="true" fill="currentColor" focusable="false" role="img" width="1em" height="1em" preserveAspectRatio="xMidYMid meet" viewBox="0 0 32 32"><path d="M28,10V28H10V10H28m0-2H10a2,2,0,0,0-2,2V28a2,2,0,0,0,2,2H28a2,2,0,0,0,2-2V10a2,2,0,0,0-2-2Z" transform="translate(0)"></path><path d="M4,18H2V4A2,2,0,0,1,4,2H18V4H4Z" transform="translate(0)"></path><rect fill="none" width="32" height="32"></rect></svg><!----> <div class=" absolute pointer-events-none transition-opacity bg-black text-white py-1 px-2 leading-tight rounded font-normal shadow left-1/2 top-full transform -translate-x-1/2 translate-y-2 opacity-0 "><div class="absolute bottom-full left-1/2 transform -translate-x-1/2 w-0 h-0 border-black border-4 border-t-0" style="border-left-color: transparent; border-right-color: transparent;"></div> Copied</div><!----></button><!----></div> <pre class="language-python "><!----><span class="hljs-keyword">def</span> <span class="hljs-title function_">cosine</span>(<span class="hljs-params">a, b</span>):
dot = <span class="hljs-built_in">sum</span>(x * y <span class="hljs-keyword">for</span> x, y <span class="hljs-keyword">in</span> <span class="hljs-built_in">zip</span>(a, b))
norm_a = math.sqrt(<span class="hljs-built_in">sum</span>(x * x <span class="hljs-keyword">for</span> x <span class="hljs-keyword">in</span> a))
norm_b = math.sqrt(<span class="hljs-built_in">sum</span>(y * y <span class="hljs-keyword">for</span> y <span class="hljs-keyword">in</span> b))
<span class="hljs-keyword">if</span> norm_a == <span class="hljs-number">0</span> <span class="hljs-keyword">or</span> norm_b == <span class="hljs-number">0</span>:
<span class="hljs-keyword">raise</span> ValueError(<span class="hljs-string">&quot;cosine similarity is undefined for a zero vector&quot;</span>)
<span class="hljs-keyword">return</span> dot / (norm_a * norm_b)
<span class="hljs-keyword">assert</span> cosine([<span class="hljs-number">1.0</span>, <span class="hljs-number">0.0</span>], [<span class="hljs-number">0.0</span>, <span class="hljs-number">1.0</span>]) == <span class="hljs-number">0.0</span>
query_uri = s3_join(input_s3_prefix, <span class="hljs-string">&quot;queries&quot;</span>, <span class="hljs-string">&quot;held-out-question.json&quot;</span>)
put_json(query_uri, {<span class="hljs-string">&quot;inputs&quot;</span>: [query_record[<span class="hljs-string">&quot;question&quot;</span>]]})
query_response = endpoint.invoke_async(
input_location=query_uri,
content_type=<span class="hljs-string">&quot;application/json&quot;</span>,
accept=<span class="hljs-string">&quot;application/json&quot;</span>,
inference_id=<span class="hljs-string">f&quot;query-<span class="hljs-subst">{uuid.uuid4()}</span>&quot;</span>,
invocation_timeout_seconds=<span class="hljs-number">900</span>,
session=boto_session,
region=region,
)
query_embedding = wait_for_async_json(query_response)[<span class="hljs-number">0</span>]
ranked = <span class="hljs-built_in">sorted</span>(
(
{
<span class="hljs-string">&quot;id&quot;</span>: record[<span class="hljs-string">&quot;id&quot;</span>],
<span class="hljs-string">&quot;score&quot;</span>: cosine(query_embedding, record[<span class="hljs-string">&quot;embedding&quot;</span>]),
<span class="hljs-string">&quot;question&quot;</span>: record[<span class="hljs-string">&quot;question&quot;</span>],
<span class="hljs-string">&quot;context&quot;</span>: record[<span class="hljs-string">&quot;context&quot;</span>],
}
<span class="hljs-keyword">for</span> record <span class="hljs-keyword">in</span> index_records
),
key=<span class="hljs-keyword">lambda</span> item: item[<span class="hljs-string">&quot;score&quot;</span>],
reverse=<span class="hljs-literal">True</span>,
)
<span class="hljs-keyword">for</span> hit <span class="hljs-keyword">in</span> ranked[:<span class="hljs-number">5</span>]:
<span class="hljs-built_in">print</span>(<span class="hljs-string">f&quot;<span class="hljs-subst">{hit[<span class="hljs-string">&#x27;score&#x27;</span>]:<span class="hljs-number">.3</span>f}</span> <span class="hljs-subst">{hit[<span class="hljs-string">&#x27;id&#x27;</span>]}</span>: <span class="hljs-subst">{hit[<span class="hljs-string">&#x27;question&#x27;</span>]}</span>&quot;</span>)
top_ids = [hit[<span class="hljs-string">&quot;id&quot;</span>] <span class="hljs-keyword">for</span> hit <span class="hljs-keyword">in</span> ranked[:<span class="hljs-number">5</span>]]
<span class="hljs-keyword">assert</span> query_record[<span class="hljs-string">&quot;id&quot;</span>] <span class="hljs-keyword">in</span> top_ids<!----></pre></div><!----> <!--[1--><h2 class="relative group"><a id="clean-up" class="header-link block pr-1.5 text-lg no-hover:hidden with-hover:absolute with-hover:p-1.5 with-hover:opacity-0 with-hover:group-hover:opacity-100 with-hover:right-full" href="#clean-up"><span><svg xmlns="http://www.w3.org/2000/svg" xmlns:xlink="http://www.w3.org/1999/xlink" aria-hidden="true" role="img" width="1em" height="1em" preserveAspectRatio="xMidYMid meet" viewBox="0 0 256 256"><path d="M167.594 88.393a8.001 8.001 0 0 1 0 11.314l-67.882 67.882a8 8 0 1 1-11.314-11.315l67.882-67.881a8.003 8.003 0 0 1 11.314 0zm-28.287 84.86l-28.284 28.284a40 40 0 0 1-56.567-56.567l28.284-28.284a8 8 0 0 0-11.315-11.315l-28.284 28.284a56 56 0 0 0 79.196 79.197l28.285-28.285a8 8 0 1 0-11.315-11.314zM212.852 43.14a56.002 56.002 0 0 0-79.196 0l-28.284 28.284a8 8 0 1 0 11.314 11.314l28.284-28.284a40 40 0 0 1 56.568 56.567l-28.285 28.285a8 8 0 0 0 11.315 11.314l28.284-28.284a56.065 56.065 0 0 0 0-79.196z" fill="currentColor"></path></svg><!----></span></a> <span>Clean up</span></h2><!--]--><!----> <p>Delete the endpoint, endpoint configuration, model, autoscaling target, and
tutorial alarms when you are done. The S3 inputs and outputs are left in
place because they are useful for inspection and because many teams hand
those objects to the next indexing stage.</p> <div class="code-block relative "><div class="absolute top-2.5 right-4"><button class="inline-flex items-center relative text-sm focus:text-green-500 cursor-pointer focus:outline-none transition duration-200 ease-in-out opacity-0 mx-0.5 text-gray-600 " title="code excerpt" type="button"><svg xmlns="http://www.w3.org/2000/svg" aria-hidden="true" fill="currentColor" focusable="false" role="img" width="1em" height="1em" preserveAspectRatio="xMidYMid meet" viewBox="0 0 32 32"><path d="M28,10V28H10V10H28m0-2H10a2,2,0,0,0-2,2V28a2,2,0,0,0,2,2H28a2,2,0,0,0,2-2V10a2,2,0,0,0-2-2Z" transform="translate(0)"></path><path d="M4,18H2V4A2,2,0,0,1,4,2H18V4H4Z" transform="translate(0)"></path><rect fill="none" width="32" height="32"></rect></svg><!----> <div class=" absolute pointer-events-none transition-opacity bg-black text-white py-1 px-2 leading-tight rounded font-normal shadow left-1/2 top-full transform -translate-x-1/2 translate-y-2 opacity-0 "><div class="absolute bottom-full left-1/2 transform -translate-x-1/2 w-0 h-0 border-black border-4 border-t-0" style="border-left-color: transparent; border-right-color: transparent;"></div> Copied</div><!----></button><!----></div> <pre class="language-python "><!----><span class="hljs-keyword">def</span> <span class="hljs-title function_">ignore_not_found</span>(<span class="hljs-params">error</span>):
code = error.response.get(<span class="hljs-string">&quot;Error&quot;</span>, {}).get(<span class="hljs-string">&quot;Code&quot;</span>, <span class="hljs-string">&quot;&quot;</span>)
message = error.response.get(<span class="hljs-string">&quot;Error&quot;</span>, {}).get(<span class="hljs-string">&quot;Message&quot;</span>, <span class="hljs-string">&quot;&quot;</span>)
<span class="hljs-keyword">return</span> code <span class="hljs-keyword">in</span> {<span class="hljs-string">&quot;ResourceNotFound&quot;</span>, <span class="hljs-string">&quot;ResourceNotFoundException&quot;</span>} <span class="hljs-keyword">or</span> <span class="hljs-string">&quot;not exist&quot;</span> <span class="hljs-keyword">in</span> message
<span class="hljs-keyword">def</span> <span class="hljs-title function_">cleanup_resources</span>():
<span class="hljs-built_in">print</span>(<span class="hljs-string">&quot;deleting CloudWatch alarms&quot;</span>)
<span class="hljs-keyword">try</span>:
cloudwatch.delete_alarms(AlarmNames=alarm_names)
<span class="hljs-keyword">except</span> ClientError <span class="hljs-keyword">as</span> error:
<span class="hljs-keyword">if</span> <span class="hljs-keyword">not</span> ignore_not_found(error):
<span class="hljs-keyword">raise</span>
<span class="hljs-built_in">print</span>(<span class="hljs-string">&quot;deregistering scalable target&quot;</span>)
<span class="hljs-keyword">try</span>:
autoscaling.deregister_scalable_target(
ServiceNamespace=<span class="hljs-string">&quot;sagemaker&quot;</span>,
ResourceId=resource_id,
ScalableDimension=<span class="hljs-string">&quot;sagemaker:variant:DesiredInstanceCount&quot;</span>,
)
<span class="hljs-keyword">except</span> ClientError <span class="hljs-keyword">as</span> error:
<span class="hljs-keyword">if</span> <span class="hljs-keyword">not</span> ignore_not_found(error):
<span class="hljs-keyword">raise</span>
<span class="hljs-built_in">print</span>(<span class="hljs-string">&quot;deleting endpoint&quot;</span>)
<span class="hljs-keyword">try</span>:
sm.delete_endpoint(EndpointName=ENDPOINT_NAME)
sm.get_waiter(<span class="hljs-string">&quot;endpoint_deleted&quot;</span>).wait(
EndpointName=ENDPOINT_NAME,
WaiterConfig={<span class="hljs-string">&quot;Delay&quot;</span>: <span class="hljs-number">30</span>, <span class="hljs-string">&quot;MaxAttempts&quot;</span>: <span class="hljs-number">60</span>},
)
<span class="hljs-keyword">except</span> ClientError <span class="hljs-keyword">as</span> error:
<span class="hljs-keyword">if</span> <span class="hljs-keyword">not</span> ignore_not_found(error):
<span class="hljs-keyword">raise</span>
<span class="hljs-built_in">print</span>(<span class="hljs-string">&quot;deleting endpoint config&quot;</span>)
<span class="hljs-keyword">try</span>:
sm.delete_endpoint_config(EndpointConfigName=endpoint_config_name)
<span class="hljs-keyword">except</span> ClientError <span class="hljs-keyword">as</span> error:
<span class="hljs-keyword">if</span> <span class="hljs-keyword">not</span> ignore_not_found(error):
<span class="hljs-keyword">raise</span>
<span class="hljs-built_in">print</span>(<span class="hljs-string">&quot;deleting model&quot;</span>)
<span class="hljs-keyword">try</span>:
sm.delete_model(ModelName=model_name)
<span class="hljs-keyword">except</span> ClientError <span class="hljs-keyword">as</span> error:
<span class="hljs-keyword">if</span> <span class="hljs-keyword">not</span> ignore_not_found(error):
<span class="hljs-keyword">raise</span>
<span class="hljs-keyword">if</span> CLEANUP:
cleanup_resources()
<span class="hljs-keyword">else</span>:
<span class="hljs-built_in">print</span>(<span class="hljs-string">f&quot;left endpoint running: <span class="hljs-subst">{ENDPOINT_NAME}</span>&quot;</span>)<!----></pre></div><!----> <hr/> <blockquote class="tip"><p>📍 Find the complete example on GitHub <a href="https://github.com/huggingface/hub-docs/tree/main/docs/sagemaker/notebooks/sagemaker-sdk/async-inference-embedding-tei/sagemaker-notebook.ipynb" rel="nofollow">here</a>!</p><!----></blockquote><!----> <a class="!text-gray-400 !no-underline text-sm flex items-center not-prose mt-4" href="https://github.com/huggingface/hub-docs/blob/main/docs/sagemaker/source/examples/sagemaker-sdk-async-inference-embedding-tei.mdx" target="_blank"><svg class="mr-1" xmlns="http://www.w3.org/2000/svg" aria-hidden="true" fill="currentColor" focusable="false" role="img" width="1em" height="1em" preserveAspectRatio="xMidYMid meet" viewBox="0 0 32 32"><path d="M31,16l-7,7l-1.41-1.41L28.17,16l-5.58-5.59L24,9l7,7z"></path><path d="M1,16l7-7l1.41,1.41L3.83,16l5.58,5.59L8,23l-7-7z"></path><path d="M12.419,25.484L17.639,6.552l1.932,0.518L14.351,26.002z"></path></svg><!----> <span><span class="underline">Update</span> on GitHub</span></a><!----> <p></p><!--]--><!----><!--]--><!--]--><!--]--> <!--[-1--><!--]--><!--]-->
<script>
{
__sveltekit_dwws3f = {
base: "/docs/sagemaker/pr_2709/en",
assets: "/docs/sagemaker/pr_2709/en"
};
const element = document.currentScript.parentElement;
Promise.all([
import("/docs/sagemaker/pr_2709/en/_app/immutable/entry/start.COiK-bUV.js"),
import("/docs/sagemaker/pr_2709/en/_app/immutable/entry/app.CjYANwwM.js")
]).then(([kit, app]) => {
kit.start(app, element, {
node_ids: [0, 4],
data: [null,null],
form: null,
error: null
});
});
}
</script>

Xet Storage Details

Size:
92.2 kB
·
Xet hash:
97b98fed2772e277be559ac1451b082e1a9bba214950e08f0c5b030904c62ce3

Xet efficiently stores files, intelligently splitting them into unique chunks and accelerating uploads and downloads. More info.