Setting the file. One moment.
Deploy Async · Hf Cloud Sagemaker Production Defaults · huggingface/skills · Skills Docs
Back to the top of the page Script Common
scripts/ deploy_async.py
Python · 434 lines · 18 KB
15 a HasBacklogWithoutCapacity alarm. Target-tracking alone can't transition
16 from zero, so without the step policy the endpoint won't wake up.
17
18 See SKILL.md for the recommended chained-with-resolver pattern.
19 """
20
21 import argparse
22 import json
23 import sys
24 from typing import Any
25
26 import boto3
27 from botocore.exceptions import ClientError
28
29 from _common import (
30 build_tags,
31 create_model,
32 log as _log,
33 make_endpoint_name,
34 parse_env,
35 wait_for_endpoint,
36 )
37
38
39 # Defaults — change here, not at call sites.
40 DEFAULTS = {
41 "initial_instance_count" : 1 ,
42 "min_capacity" : 0 , # async genuinely supports scale-to-zero
43 "max_capacity" : 4 ,
44 "backlog_per_instance_target" : 5 , # target queue depth per instance
45 "scale_in_cooldown_seconds" : 600 , # slower scale-in for async — batches are bursty
46 "scale_out_cooldown_seconds" : 60 ,
47 "wake_from_zero_step_size" : 1 , # 0→1 instance when backlog appears
48 "max_concurrent_invocations_per_instance" : 4 ,
49 # Alarms
50 "alarm_backlog_size_threshold" : 50 , # queue too deep
51 "alarm_failed_invocations_threshold" : 5 , # repeated failures
52 "alarm_evaluation_periods" : 1 ,
53 "alarm_period_seconds" : 60 , # async needs faster eval for wake-from-zero
54 "environment_tag" : "dev" ,
55 }
56
57
58 def log (msg: str ) -> None :
59 _log( "deploy_async" , msg)
60
61
62 def create_async_endpoint_config (
63 sm: Any, * , config_name: str , model_name: str , instance_type: str ,
64 initial_instance_count: int , inference_ami_version: str | None ,
65 output_s3_uri: str , failure_s3_uri: str | None ,
66 success_topic_arn: str | None , error_topic_arn: str | None ,
67 max_concurrent_invocations: int ,
68 tags: list[ dict ],
69 ) -> str :
70 log( f "Creating async endpoint config: { config_name } " )
71
72 production_variant: dict[ str , Any] = {
73 "VariantName" : "AllTraffic" ,
74 "ModelName" : model_name,
75 "InstanceType" : instance_type,
76 "InitialInstanceCount" : initial_instance_count,
77 "InitialVariantWeight" : 1.0 ,
78 }
79 if inference_ami_version:
80 production_variant[ "InferenceAmiVersion" ] = inference_ami_version
81 log( f " InferenceAmiVersion set to: { inference_ami_version } " )
82
83 output_config: dict[ str , Any] = { "S3OutputPath" : output_s3_uri}
84 if failure_s3_uri:
85 output_config[ "S3FailurePath" ] = failure_s3_uri
86 if success_topic_arn or error_topic_arn:
87 notification: dict[ str , str ] = {}
88 if success_topic_arn:
89 notification[ "SuccessTopic" ] = success_topic_arn
90 if error_topic_arn:
91 notification[ "ErrorTopic" ] = error_topic_arn
92 output_config[ "NotificationConfig" ] = notification
93
94 async_config: dict[ str , Any] = {
95 "OutputConfig" : output_config,
96 "ClientConfig" : {
97 "MaxConcurrentInvocationsPerInstance" : max_concurrent_invocations,
98 },
99 }
100
101 kwargs: dict[ str , Any] = {
102 "EndpointConfigName" : config_name,
103 "ProductionVariants" : [production_variant],
104 "AsyncInferenceConfig" : async_config,
105 "Tags" : tags,
106 }
107
108 try :
109 sm.create_endpoint_config( ** kwargs)
110 except ClientError as e:
111 if "Cannot create already existing endpoint configuration" in str (e):
112 log( f "Endpoint config { config_name } already exists — reusing" )
113 else :
114 raise
115 return config_name
116
117
118 def create_endpoint (sm: Any, * , endpoint_name: str , config_name: str , tags: list[ dict ]) -> None :
119 log( f "Creating endpoint: { endpoint_name } " )
120 sm.create_endpoint(
121 EndpointName = endpoint_name,
122 EndpointConfigName = config_name,
123 Tags = tags,
124 )
125
126
127 def register_async_autoscaling (
128 * , endpoint_name: str , variant_name: str , min_capacity: int , max_capacity: int ,
129 backlog_per_instance_target: int , scale_in_cooldown: int , scale_out_cooldown: int ,
130 wake_step_size: int , region: str ,
131 ) -> None :
132 """Register two autoscaling policies on the variant:
133
134 1. Target-tracking on ApproximateBacklogSizePerInstance — handles
135 ongoing scaling between MinCapacity and MaxCapacity.
136 2. Step-scaling that increments capacity by `wake_step_size` when
137 triggered — needed for 0→1 wake-from-zero, because target-tracking
138 can't transition from zero capacity.
139
140 The step-scaling policy is invoked by a HasBacklogWithoutCapacity alarm
141 (created separately in create_async_alarms).
142 """
143 log( f "Registering async autoscaling: min= { min_capacity } max= { max_capacity } backlog-target= { backlog_per_instance_target } " )
144 appscaling = boto3.client( "application-autoscaling" , region_name = region)
145 resource_id = f "endpoint/ { endpoint_name } /variant/ { variant_name } "
146
147 appscaling.register_scalable_target(
148 ServiceNamespace = "sagemaker" ,
149 ResourceId = resource_id,
150 ScalableDimension = "sagemaker:variant:DesiredInstanceCount" ,
151 MinCapacity = min_capacity,
152 MaxCapacity = max_capacity,
153 )
154
155 # (1) Target-tracking for ongoing scaling between min and max
156 appscaling.put_scaling_policy(
157 PolicyName = f " { endpoint_name } -backlog-target-tracking" ,
158 ServiceNamespace = "sagemaker" ,
159 ResourceId = resource_id,
160 ScalableDimension = "sagemaker:variant:DesiredInstanceCount" ,
161 PolicyType = "TargetTrackingScaling" ,
162 TargetTrackingScalingPolicyConfiguration = {
163 "TargetValue" : float (backlog_per_instance_target),
164 "CustomizedMetricSpecification" : {
165 "MetricName" : "ApproximateBacklogSizePerInstance" ,
166 "Namespace" : "AWS/SageMaker" ,
167 "Dimensions" : [{ "Name" : "EndpointName" , "Value" : endpoint_name}],
168 "Statistic" : "Average" ,
169 },
170 "ScaleInCooldown" : scale_in_cooldown,
171 "ScaleOutCooldown" : scale_out_cooldown,
172 },
173 )
174
175 # (2) Step-scaling for wake-from-zero. The CloudWatch alarm bound to
176 # this policy is created in create_async_alarms.
177 appscaling.put_scaling_policy(
178 PolicyName = f " { endpoint_name } -step-wake-from-zero" ,
179 ServiceNamespace = "sagemaker" ,
180 ResourceId = resource_id,
181 ScalableDimension = "sagemaker:variant:DesiredInstanceCount" ,
182 PolicyType = "StepScaling" ,
183 StepScalingPolicyConfiguration = {
184 "AdjustmentType" : "ChangeInCapacity" ,
185 "MetricAggregationType" : "Maximum" ,
186 "Cooldown" : scale_out_cooldown,
187 "StepAdjustments" : [
188 {
189 "MetricIntervalLowerBound" : 0 ,
190 "ScalingAdjustment" : wake_step_size,
191 },
192 ],
193 },
194 )
195
196
197 def create_async_alarms (
198 * , endpoint_name: str , variant_name: str , sns_topic_arn: str | None ,
199 region: str , wake_alarm_arns_for_step_policy: list[ str ],
200 ) -> None :
201 """Create CloudWatch alarms for the async endpoint, including the
202 HasBacklogWithoutCapacity alarm that drives the wake-from-zero policy.
203 """
204 log( f "Creating CloudWatch alarms for { endpoint_name } " )
205 cw = boto3.client( "cloudwatch" , region_name = region)
206 actions = [sns_topic_arn] if sns_topic_arn else []
207 endpoint_dim = [{ "Name" : "EndpointName" , "Value" : endpoint_name}]
208
209 # 1. Backlog too deep — capacity is keeping up poorly with incoming requests
210 cw.put_metric_alarm(
211 AlarmName = f " { endpoint_name } -ApproximateBacklogSize" ,
212 AlarmDescription = "Async queue backlog too deep" ,
213 MetricName = "ApproximateBacklogSize" ,
214 Namespace = "AWS/SageMaker" ,
215 Dimensions = endpoint_dim,
216 Statistic = "Average" ,
217 Period = DEFAULTS [ "alarm_period_seconds" ],
218 EvaluationPeriods = DEFAULTS [ "alarm_evaluation_periods" ],
219 Threshold = DEFAULTS [ "alarm_backlog_size_threshold" ],
220 ComparisonOperator = "GreaterThanThreshold" ,
221 TreatMissingData = "notBreaching" ,
222 AlarmActions = actions,
223 )
224
225 # 2. InvocationsFailed — repeated client/server errors in async processing
226 cw.put_metric_alarm(
227 AlarmName = f " { endpoint_name } -InvocationsFailed" ,
228 AlarmDescription = "Async invocations failing repeatedly" ,
229 MetricName = "InvocationsFailed" ,
230 Namespace = "AWS/SageMaker" ,
231 Dimensions = endpoint_dim,
232 Statistic = "Sum" ,
233 Period = DEFAULTS [ "alarm_period_seconds" ],
234 EvaluationPeriods = DEFAULTS [ "alarm_evaluation_periods" ],
235 Threshold = DEFAULTS [ "alarm_failed_invocations_threshold" ],
236 ComparisonOperator = "GreaterThanThreshold" ,
237 TreatMissingData = "notBreaching" ,
238 AlarmActions = actions,
239 )
240
241 # 3. HasBacklogWithoutCapacity — drives the step-scaling wake-from-zero
242 # policy. The Alarm action is the step policy ARN, not the SNS topic.
243 cw.put_metric_alarm(
244 AlarmName = f " { endpoint_name } -HasBacklogWithoutCapacity" ,
245 AlarmDescription = "Async backlog exists but no instances running — triggers wake-from-zero" ,
246 MetricName = "HasBacklogWithoutCapacity" ,
247 Namespace = "AWS/SageMaker" ,
248 Dimensions = endpoint_dim,
249 Statistic = "Average" ,
250 Period = 60 , # tight evaluation: we want fast wake-up
251 EvaluationPeriods = 1 ,
252 Threshold = 1 ,
253 ComparisonOperator = "GreaterThanOrEqualToThreshold" ,
254 TreatMissingData = "notBreaching" ,
255 AlarmActions = wake_alarm_arns_for_step_policy, # step policy ARNs
256 )
257
258 if not sns_topic_arn:
259 log( "WARNING: no --sns-alarm-topic — backlog/failure alarms exist but won't notify anyone." )
260
261
262 def get_step_policy_arn ( * , endpoint_name: str , variant_name: str , region: str ) -> str :
263 """Look up the ARN of the step-scaling policy we just created.
264
265 `put_scaling_policy` returns the ARN on creation but we don't capture it
266 above to keep the function signature small. Re-fetch via describe.
267 """
268 appscaling = boto3.client( "application-autoscaling" , region_name = region)
269 resp = appscaling.describe_scaling_policies(
270 ServiceNamespace = "sagemaker" ,
271 ResourceId = f "endpoint/ { endpoint_name } /variant/ { variant_name } " ,
272 PolicyNames = [ f " { endpoint_name } -step-wake-from-zero" ],
273 )
274 policies = resp.get( "ScalingPolicies" , [])
275 if not policies:
276 raise RuntimeError ( f "Step-scaling policy not found for { endpoint_name } — was register_async_autoscaling called?" )
277 return policies[ 0 ][ "PolicyARN" ]
278
279
280 def main () -> int :
281 p = argparse.ArgumentParser( description = __doc__ , formatter_class = argparse.RawDescriptionHelpFormatter)
282
283 # Required
284 p.add_argument( "--model-name" , required = True )
285 p.add_argument( "--image-uri" , required = True , help = "From hf-cloud-serving-image-selection" )
286 p.add_argument( "--role-arn" , required = True , help = "From hf-cloud-sagemaker-iam-preflight" )
287 p.add_argument( "--instance-type" , required = True , help = "e.g. ml.g5.xlarge" )
288 p.add_argument( "--region" , required = True , help = "From hf-cloud-aws-context-discovery" )
289 p.add_argument( "--output-s3-uri" , required = True ,
290 help = "S3 path where async results are written. e.g. s3://amzn-s3-demo-async-output/async-output/" )
291
292 # Conditional
293 p.add_argument( "--model-s3-uri" , default = None ,
294 help = "S3 model artifact extracted to /opt/ml/model (preferred); omit only for Hub-at-runtime" )
295 p.add_argument( "--env" , action = "append" , default = [], help = "KEY=VALUE; repeatable" )
296 p.add_argument( "--inference-ami-version" , default = None ,
297 help = "REQUIRED for vLLM DLC with CUDA 13+ (e.g. al2-ami-sagemaker-inference-gpu-3-1)" )
298 p.add_argument( "--failure-s3-uri" , default = None ,
299 help = "S3 path for failed async invocations (optional; default: errors written to OutputConfig.S3OutputPath)" )
300 p.add_argument( "--success-sns-topic" , default = None ,
301 help = "SNS topic ARN notified when async invocation succeeds" )
302 p.add_argument( "--error-sns-topic" , default = None ,
303 help = "SNS topic ARN notified when async invocation fails" )
304
305 # Naming
306 p.add_argument( "--endpoint-name" , default = None , help = "Default: <model-name>-<timestamp>" )
307 p.add_argument( "--project" , default = None , help = "Tag value (default: model name)" )
308 p.add_argument( "--environment" , default = DEFAULTS [ "environment_tag" ])
309
310 # Capacity / scaling
311 p.add_argument( "--initial-instance-count" , type = int , default = DEFAULTS [ "initial_instance_count" ])
312 p.add_argument( "--min-capacity" , type = int , default = DEFAULTS [ "min_capacity" ],
313 help = "Async supports 0 (scale-to-zero between batches). Default 0." )
314 p.add_argument( "--max-capacity" , type = int , default = DEFAULTS [ "max_capacity" ])
315 p.add_argument( "--backlog-per-instance-target" , type = int ,
316 default = DEFAULTS [ "backlog_per_instance_target" ],
317 help = "Target queue depth per instance for autoscaling" )
318 p.add_argument( "--max-concurrent-invocations-per-instance" , type = int ,
319 default = DEFAULTS [ "max_concurrent_invocations_per_instance" ])
320 p.add_argument( "--no-autoscaling" , action = "store_true" ,
321 help = "NOT RECOMMENDED for async — endpoint won't scale to zero" )
322
323 # Alarms
324 p.add_argument( "--sns-alarm-topic" , default = None , help = "SNS topic ARN for backlog/failure alarms" )
325 p.add_argument( "--no-alarms" , action = "store_true" )
326
327 args = p.parse_args()
328
329 env_dict = parse_env(args.env)
330 endpoint_name = make_endpoint_name(args.model_name, args.endpoint_name)
331 config_name = f " { endpoint_name } -config"
332
333 sts = boto3.client( "sts" , region_name = args.region)
334 sm = boto3.client( "sagemaker" , region_name = args.region)
335 caller_arn = sts.get_caller_identity()[ "Arn" ]
336
337 if args.min_capacity == 0 and args.no_autoscaling:
338 raise SystemExit (
339 "--min-capacity 0 requires autoscaling to be enabled (the wake-from-zero "
340 "policy is what brings instances up). Either remove --no-autoscaling or "
341 "set --min-capacity 1."
342 )
343
344 tags = build_tags(
345 project = args.project or args.model_name,
346 caller_arn = caller_arn,
347 environment = args.environment,
348 model_s3_uri = args.model_s3_uri,
349 extra = { "InferenceMode" : "async" },
350 )
351
352 create_model(
353 sm, model_name = args.model_name, image_uri = args.image_uri,
354 role_arn = args.role_arn, model_s3_uri = args.model_s3_uri,
355 env = env_dict, tags = tags, log_prefix = "deploy_async" ,
356 )
357 create_async_endpoint_config(
358 sm, config_name = config_name, model_name = args.model_name,
359 instance_type = args.instance_type,
360 initial_instance_count = args.initial_instance_count,
361 inference_ami_version = args.inference_ami_version,
362 output_s3_uri = args.output_s3_uri,
363 failure_s3_uri = args.failure_s3_uri,
364 success_topic_arn = args.success_sns_topic,
365 error_topic_arn = args.error_sns_topic,
366 max_concurrent_invocations = args.max_concurrent_invocations_per_instance,
367 tags = tags,
368 )
369 create_endpoint(sm, endpoint_name = endpoint_name, config_name = config_name, tags = tags)
370 wait_for_endpoint(sm, endpoint_name, log_prefix = "deploy_async" )
371
372 if not args.no_autoscaling:
373 register_async_autoscaling(
374 endpoint_name = endpoint_name, variant_name = "AllTraffic" ,
375 min_capacity = args.min_capacity, max_capacity = args.max_capacity,
376 backlog_per_instance_target = args.backlog_per_instance_target,
377 scale_in_cooldown = DEFAULTS [ "scale_in_cooldown_seconds" ],
378 scale_out_cooldown = DEFAULTS [ "scale_out_cooldown_seconds" ],
379 wake_step_size = DEFAULTS [ "wake_from_zero_step_size" ],
380 region = args.region,
381 )
382 # The wake-from-zero alarm needs the step policy's ARN as its action.
383 step_policy_arn = get_step_policy_arn(
384 endpoint_name = endpoint_name, variant_name = "AllTraffic" , region = args.region,
385 )
386 else :
387 log( "WARNING: autoscaling skipped. Endpoint will NOT scale (in either direction)." )
388 step_policy_arn = None
389
390 if not args.no_alarms:
391 create_async_alarms(
392 endpoint_name = endpoint_name, variant_name = "AllTraffic" ,
393 sns_topic_arn = args.sns_alarm_topic, region = args.region,
394 wake_alarm_arns_for_step_policy = [step_policy_arn] if step_policy_arn else [],
395 )
396
397 # Summary
398 log( "" )
399 log( f "Async deployment complete: { endpoint_name } " )
400 log( f " Instance: { args.instance_type } " )
401 if args.no_autoscaling:
402 autoscaling_summary = "OFF"
403 else :
404 zero_label = "enabled" if args.min_capacity == 0 else "disabled"
405 autoscaling_summary = f " { args.min_capacity } - { args.max_capacity } instances (scale-to-zero { zero_label } )"
406 log( f " Autoscaling: { autoscaling_summary } " )
407 log( f " Output S3 path: { args.output_s3_uri } " )
408 log( f " Notifications: success= { args.success_sns_topic or 'none' } error= { args.error_sns_topic or 'none' } " )
409 log( "" )
410 log( "Invoke (async, via S3 input location):" )
411 log( f " aws sagemaker-runtime invoke-endpoint-async \\ " )
412 log( f " --endpoint-name { endpoint_name } \\ " )
413 log( f " --input-location s3://YOUR-INPUT-BUCKET/path/to/input.json \\ " )
414 log( f " --content-type application/json --region { args.region } " )
415 log( f " # Result will land at: { args.output_s3_uri } " )
416 log( " # Write input.json as BOM-free UTF-8 (on Windows, NOT 'Set-Content -Encoding UTF8')" )
417 log( "" )
418 log( f "Teardown: python3 teardown.py { endpoint_name } { args.region } " )
419
420 print (json.dumps({
421 "endpoint_name" : endpoint_name,
422 "endpoint_config_name" : config_name,
423 "model_name" : args.model_name,
424 "region" : args.region,
425 "instance_type" : args.instance_type,
426 "inference_mode" : "async" ,
427 "output_s3_uri" : args.output_s3_uri,
428 }))
429
430 return 0
431
432
433 if __name__ == "__main__" :
434 sys.exit(main())