tribunal-kit 9.2.3 → 9.2.5

package/skills/hf-cloud-sagemaker-production-defaults/scripts/deploy_async.py

@@ -0,0 +1,434 @@
+#!/usr/bin/env python
+"""Create a SageMaker async-inference endpoint with production defaults.
+
+Minimal usage:
+    python deploy_async.py --model-name <name> --image-uri <uri> \
+        --role-arn <arn> --instance-type ml.g5.xlarge --region <region> \
+        --output-s3-uri s3://bucket/path/
+
+Async differs from real-time in three structural ways:
+  1. AsyncInferenceConfig on the endpoint config (S3 output path, optional
+     SNS notification topics, optional failure path)
+  2. Autoscaling MinCapacity=0 is supported — endpoints can genuinely scale
+     to zero between batches. Target-tracking on ApproximateBacklogSizePerInstance.
+  3. Scale-from-zero (0→1) requires a SEPARATE step-scaling policy bound to
+     a HasBacklogWithoutCapacity alarm. Target-tracking alone can't transition
+     from zero, so without the step policy the endpoint won't wake up.
+
+See SKILL.md for the recommended chained-with-resolver pattern.
+"""
+
+import argparse
+import json
+import sys
+from typing import Any
+
+import boto3
+from botocore.exceptions import ClientError
+
+from _common import (
+    build_tags,
+    create_model,
+    log as _log,
+    make_endpoint_name,
+    parse_env,
+    wait_for_endpoint,
+)
+
+
+# Defaults — change here, not at call sites.
+DEFAULTS = {
+    "initial_instance_count": 1,
+    "min_capacity": 0,                          # async genuinely supports scale-to-zero
+    "max_capacity": 4,
+    "backlog_per_instance_target": 5,           # target queue depth per instance
+    "scale_in_cooldown_seconds": 600,           # slower scale-in for async — batches are bursty
+    "scale_out_cooldown_seconds": 60,
+    "wake_from_zero_step_size": 1,              # 0→1 instance when backlog appears
+    "max_concurrent_invocations_per_instance": 4,
+    # Alarms
+    "alarm_backlog_size_threshold": 50,         # queue too deep
+    "alarm_failed_invocations_threshold": 5,    # repeated failures
+    "alarm_evaluation_periods": 1,
+    "alarm_period_seconds": 60,                 # async needs faster eval for wake-from-zero
+    "environment_tag": "dev",
+}
+
+
+def log(msg: str) -> None:
+    _log("deploy_async", msg)
+
+
+def create_async_endpoint_config(
+    sm: Any, *, config_name: str, model_name: str, instance_type: str,
+    initial_instance_count: int, inference_ami_version: str | None,
+    output_s3_uri: str, failure_s3_uri: str | None,
+    success_topic_arn: str | None, error_topic_arn: str | None,
+    max_concurrent_invocations: int,
+    tags: list[dict],
+) -> str:
+    log(f"Creating async endpoint config: {config_name}")
+
+    production_variant: dict[str, Any] = {
+        "VariantName": "AllTraffic",
+        "ModelName": model_name,
+        "InstanceType": instance_type,
+        "InitialInstanceCount": initial_instance_count,
+        "InitialVariantWeight": 1.0,
+    }
+    if inference_ami_version:
+        production_variant["InferenceAmiVersion"] = inference_ami_version
+        log(f"  InferenceAmiVersion set to: {inference_ami_version}")
+
+    output_config: dict[str, Any] = {"S3OutputPath": output_s3_uri}
+    if failure_s3_uri:
+        output_config["S3FailurePath"] = failure_s3_uri
+    if success_topic_arn or error_topic_arn:
+        notification: dict[str, str] = {}
+        if success_topic_arn:
+            notification["SuccessTopic"] = success_topic_arn
+        if error_topic_arn:
+            notification["ErrorTopic"] = error_topic_arn
+        output_config["NotificationConfig"] = notification
+
+    async_config: dict[str, Any] = {
+        "OutputConfig": output_config,
+        "ClientConfig": {
+            "MaxConcurrentInvocationsPerInstance": max_concurrent_invocations,
+        },
+    }
+
+    kwargs: dict[str, Any] = {
+        "EndpointConfigName": config_name,
+        "ProductionVariants": [production_variant],
+        "AsyncInferenceConfig": async_config,
+        "Tags": tags,
+    }
+
+    try:
+        sm.create_endpoint_config(**kwargs)
+    except ClientError as e:
+        if "Cannot create already existing endpoint configuration" in str(e):
+            log(f"Endpoint config {config_name} already exists — reusing")
+        else:
+            raise
+    return config_name
+
+
+def create_endpoint(sm: Any, *, endpoint_name: str, config_name: str, tags: list[dict]) -> None:
+    log(f"Creating endpoint: {endpoint_name}")
+    sm.create_endpoint(
+        EndpointName=endpoint_name,
+        EndpointConfigName=config_name,
+        Tags=tags,
+    )
+
+
+def register_async_autoscaling(
+    *, endpoint_name: str, variant_name: str, min_capacity: int, max_capacity: int,
+    backlog_per_instance_target: int, scale_in_cooldown: int, scale_out_cooldown: int,
+    wake_step_size: int, region: str,
+) -> None:
+    """Register two autoscaling policies on the variant:
+
+    1. Target-tracking on ApproximateBacklogSizePerInstance — handles
+       ongoing scaling between MinCapacity and MaxCapacity.
+    2. Step-scaling that increments capacity by `wake_step_size` when
+       triggered — needed for 0→1 wake-from-zero, because target-tracking
+       can't transition from zero capacity.
+
+    The step-scaling policy is invoked by a HasBacklogWithoutCapacity alarm
+    (created separately in create_async_alarms).
+    """
+    log(f"Registering async autoscaling: min={min_capacity} max={max_capacity} backlog-target={backlog_per_instance_target}")
+    appscaling = boto3.client("application-autoscaling", region_name=region)
+    resource_id = f"endpoint/{endpoint_name}/variant/{variant_name}"
+
+    appscaling.register_scalable_target(
+        ServiceNamespace="sagemaker",
+        ResourceId=resource_id,
+        ScalableDimension="sagemaker:variant:DesiredInstanceCount",
+        MinCapacity=min_capacity,
+        MaxCapacity=max_capacity,
+    )
+
+    # (1) Target-tracking for ongoing scaling between min and max
+    appscaling.put_scaling_policy(
+        PolicyName=f"{endpoint_name}-backlog-target-tracking",
+        ServiceNamespace="sagemaker",
+        ResourceId=resource_id,
+        ScalableDimension="sagemaker:variant:DesiredInstanceCount",
+        PolicyType="TargetTrackingScaling",
+        TargetTrackingScalingPolicyConfiguration={
+            "TargetValue": float(backlog_per_instance_target),
+            "CustomizedMetricSpecification": {
+                "MetricName": "ApproximateBacklogSizePerInstance",
+                "Namespace": "AWS/SageMaker",
+                "Dimensions": [{"Name": "EndpointName", "Value": endpoint_name}],
+                "Statistic": "Average",
+            },
+            "ScaleInCooldown": scale_in_cooldown,
+            "ScaleOutCooldown": scale_out_cooldown,
+        },
+    )
+
+    # (2) Step-scaling for wake-from-zero. The CloudWatch alarm bound to
+    # this policy is created in create_async_alarms.
+    appscaling.put_scaling_policy(
+        PolicyName=f"{endpoint_name}-step-wake-from-zero",
+        ServiceNamespace="sagemaker",
+        ResourceId=resource_id,
+        ScalableDimension="sagemaker:variant:DesiredInstanceCount",
+        PolicyType="StepScaling",
+        StepScalingPolicyConfiguration={
+            "AdjustmentType": "ChangeInCapacity",
+            "MetricAggregationType": "Maximum",
+            "Cooldown": scale_out_cooldown,
+            "StepAdjustments": [
+                {
+                    "MetricIntervalLowerBound": 0,
+                    "ScalingAdjustment": wake_step_size,
+                },
+            ],
+        },
+    )
+
+
+def create_async_alarms(
+    *, endpoint_name: str, variant_name: str, sns_topic_arn: str | None,
+    region: str, wake_alarm_arns_for_step_policy: list[str],
+) -> None:
+    """Create CloudWatch alarms for the async endpoint, including the
+    HasBacklogWithoutCapacity alarm that drives the wake-from-zero policy.
+    """
+    log(f"Creating CloudWatch alarms for {endpoint_name}")
+    cw = boto3.client("cloudwatch", region_name=region)
+    actions = [sns_topic_arn] if sns_topic_arn else []
+    endpoint_dim = [{"Name": "EndpointName", "Value": endpoint_name}]
+
+    # 1. Backlog too deep — capacity is keeping up poorly with incoming requests
+    cw.put_metric_alarm(
+        AlarmName=f"{endpoint_name}-ApproximateBacklogSize",
+        AlarmDescription="Async queue backlog too deep",
+        MetricName="ApproximateBacklogSize",
+        Namespace="AWS/SageMaker",
+        Dimensions=endpoint_dim,
+        Statistic="Average",
+        Period=DEFAULTS["alarm_period_seconds"],
+        EvaluationPeriods=DEFAULTS["alarm_evaluation_periods"],
+        Threshold=DEFAULTS["alarm_backlog_size_threshold"],
+        ComparisonOperator="GreaterThanThreshold",
+        TreatMissingData="notBreaching",
+        AlarmActions=actions,
+    )
+
+    # 2. InvocationsFailed — repeated client/server errors in async processing
+    cw.put_metric_alarm(
+        AlarmName=f"{endpoint_name}-InvocationsFailed",
+        AlarmDescription="Async invocations failing repeatedly",
+        MetricName="InvocationsFailed",
+        Namespace="AWS/SageMaker",
+        Dimensions=endpoint_dim,
+        Statistic="Sum",
+        Period=DEFAULTS["alarm_period_seconds"],
+        EvaluationPeriods=DEFAULTS["alarm_evaluation_periods"],
+        Threshold=DEFAULTS["alarm_failed_invocations_threshold"],
+        ComparisonOperator="GreaterThanThreshold",
+        TreatMissingData="notBreaching",
+        AlarmActions=actions,
+    )
+
+    # 3. HasBacklogWithoutCapacity — drives the step-scaling wake-from-zero
+    # policy. The Alarm action is the step policy ARN, not the SNS topic.
+    cw.put_metric_alarm(
+        AlarmName=f"{endpoint_name}-HasBacklogWithoutCapacity",
+        AlarmDescription="Async backlog exists but no instances running — triggers wake-from-zero",
+        MetricName="HasBacklogWithoutCapacity",
+        Namespace="AWS/SageMaker",
+        Dimensions=endpoint_dim,
+        Statistic="Average",
+        Period=60,                              # tight evaluation: we want fast wake-up
+        EvaluationPeriods=1,
+        Threshold=1,
+        ComparisonOperator="GreaterThanOrEqualToThreshold",
+        TreatMissingData="notBreaching",
+        AlarmActions=wake_alarm_arns_for_step_policy,   # step policy ARNs
+    )
+
+    if not sns_topic_arn:
+        log("WARNING: no --sns-alarm-topic — backlog/failure alarms exist but won't notify anyone.")
+
+
+def get_step_policy_arn(*, endpoint_name: str, variant_name: str, region: str) -> str:
+    """Look up the ARN of the step-scaling policy we just created.
+
+    `put_scaling_policy` returns the ARN on creation but we don't capture it
+    above to keep the function signature small. Re-fetch via describe.
+    """
+    appscaling = boto3.client("application-autoscaling", region_name=region)
+    resp = appscaling.describe_scaling_policies(
+        ServiceNamespace="sagemaker",
+        ResourceId=f"endpoint/{endpoint_name}/variant/{variant_name}",
+        PolicyNames=[f"{endpoint_name}-step-wake-from-zero"],
+    )
+    policies = resp.get("ScalingPolicies", [])
+    if not policies:
+        raise RuntimeError(f"Step-scaling policy not found for {endpoint_name} — was register_async_autoscaling called?")
+    return policies[0]["PolicyARN"]
+
+
+def main() -> int:
+    p = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
+
+    # Required
+    p.add_argument("--model-name", required=True)
+    p.add_argument("--image-uri", required=True, help="From hf-cloud-serving-image-selection")
+    p.add_argument("--role-arn", required=True, help="From hf-cloud-sagemaker-iam-preflight")
+    p.add_argument("--instance-type", required=True, help="e.g. ml.g5.xlarge")
+    p.add_argument("--region", required=True, help="From hf-cloud-aws-context-discovery")
+    p.add_argument("--output-s3-uri", required=True,
+                   help="S3 path where async results are written. e.g. s3://amzn-s3-demo-async-output/async-output/")
+
+    # Conditional
+    p.add_argument("--model-s3-uri", default=None,
+                   help="S3 model artifact extracted to /opt/ml/model (preferred); omit only for Hub-at-runtime")
+    p.add_argument("--env", action="append", default=[], help="KEY=VALUE; repeatable")
+    p.add_argument("--inference-ami-version", default=None,
+                   help="REQUIRED for vLLM DLC with CUDA 13+ (e.g. al2-ami-sagemaker-inference-gpu-3-1)")
+    p.add_argument("--failure-s3-uri", default=None,
+                   help="S3 path for failed async invocations (optional; default: errors written to OutputConfig.S3OutputPath)")
+    p.add_argument("--success-sns-topic", default=None,
+                   help="SNS topic ARN notified when async invocation succeeds")
+    p.add_argument("--error-sns-topic", default=None,
+                   help="SNS topic ARN notified when async invocation fails")
+
+    # Naming
+    p.add_argument("--endpoint-name", default=None, help="Default: <model-name>-<timestamp>")
+    p.add_argument("--project", default=None, help="Tag value (default: model name)")
+    p.add_argument("--environment", default=DEFAULTS["environment_tag"])
+
+    # Capacity / scaling
+    p.add_argument("--initial-instance-count", type=int, default=DEFAULTS["initial_instance_count"])
+    p.add_argument("--min-capacity", type=int, default=DEFAULTS["min_capacity"],
+                   help="Async supports 0 (scale-to-zero between batches). Default 0.")
+    p.add_argument("--max-capacity", type=int, default=DEFAULTS["max_capacity"])
+    p.add_argument("--backlog-per-instance-target", type=int,
+                   default=DEFAULTS["backlog_per_instance_target"],
+                   help="Target queue depth per instance for autoscaling")
+    p.add_argument("--max-concurrent-invocations-per-instance", type=int,
+                   default=DEFAULTS["max_concurrent_invocations_per_instance"])
+    p.add_argument("--no-autoscaling", action="store_true",
+                   help="NOT RECOMMENDED for async — endpoint won't scale to zero")
+
+    # Alarms
+    p.add_argument("--sns-alarm-topic", default=None, help="SNS topic ARN for backlog/failure alarms")
+    p.add_argument("--no-alarms", action="store_true")
+
+    args = p.parse_args()
+
+    env_dict = parse_env(args.env)
+    endpoint_name = make_endpoint_name(args.model_name, args.endpoint_name)
+    config_name = f"{endpoint_name}-config"
+
+    sts = boto3.client("sts", region_name=args.region)
+    sm = boto3.client("sagemaker", region_name=args.region)
+    caller_arn = sts.get_caller_identity()["Arn"]
+
+    if args.min_capacity == 0 and args.no_autoscaling:
+        raise SystemExit(
+            "--min-capacity 0 requires autoscaling to be enabled (the wake-from-zero "
+            "policy is what brings instances up). Either remove --no-autoscaling or "
+            "set --min-capacity 1."
+        )
+
+    tags = build_tags(
+        project=args.project or args.model_name,
+        caller_arn=caller_arn,
+        environment=args.environment,
+        model_s3_uri=args.model_s3_uri,
+        extra={"InferenceMode": "async"},
+    )
+
+    create_model(
+        sm, model_name=args.model_name, image_uri=args.image_uri,
+        role_arn=args.role_arn, model_s3_uri=args.model_s3_uri,
+        env=env_dict, tags=tags, log_prefix="deploy_async",
+    )
+    create_async_endpoint_config(
+        sm, config_name=config_name, model_name=args.model_name,
+        instance_type=args.instance_type,
+        initial_instance_count=args.initial_instance_count,
+        inference_ami_version=args.inference_ami_version,
+        output_s3_uri=args.output_s3_uri,
+        failure_s3_uri=args.failure_s3_uri,
+        success_topic_arn=args.success_sns_topic,
+        error_topic_arn=args.error_sns_topic,
+        max_concurrent_invocations=args.max_concurrent_invocations_per_instance,
+        tags=tags,
+    )
+    create_endpoint(sm, endpoint_name=endpoint_name, config_name=config_name, tags=tags)
+    wait_for_endpoint(sm, endpoint_name, log_prefix="deploy_async")
+
+    if not args.no_autoscaling:
+        register_async_autoscaling(
+            endpoint_name=endpoint_name, variant_name="AllTraffic",
+            min_capacity=args.min_capacity, max_capacity=args.max_capacity,
+            backlog_per_instance_target=args.backlog_per_instance_target,
+            scale_in_cooldown=DEFAULTS["scale_in_cooldown_seconds"],
+            scale_out_cooldown=DEFAULTS["scale_out_cooldown_seconds"],
+            wake_step_size=DEFAULTS["wake_from_zero_step_size"],
+            region=args.region,
+        )
+        # The wake-from-zero alarm needs the step policy's ARN as its action.
+        step_policy_arn = get_step_policy_arn(
+            endpoint_name=endpoint_name, variant_name="AllTraffic", region=args.region,
+        )
+    else:
+        log("WARNING: autoscaling skipped. Endpoint will NOT scale (in either direction).")
+        step_policy_arn = None
+
+    if not args.no_alarms:
+        create_async_alarms(
+            endpoint_name=endpoint_name, variant_name="AllTraffic",
+            sns_topic_arn=args.sns_alarm_topic, region=args.region,
+            wake_alarm_arns_for_step_policy=[step_policy_arn] if step_policy_arn else [],
+        )
+
+    # Summary
+    log("")
+    log(f"Async deployment complete: {endpoint_name}")
+    log(f"  Instance:           {args.instance_type}")
+    if args.no_autoscaling:
+        autoscaling_summary = "OFF"
+    else:
+        zero_label = "enabled" if args.min_capacity == 0 else "disabled"
+        autoscaling_summary = f"{args.min_capacity}-{args.max_capacity} instances (scale-to-zero {zero_label})"
+    log(f"  Autoscaling:        {autoscaling_summary}")
+    log(f"  Output S3 path:     {args.output_s3_uri}")
+    log(f"  Notifications:      success={args.success_sns_topic or 'none'} error={args.error_sns_topic or 'none'}")
+    log("")
+    log("Invoke (async, via S3 input location):")
+    log(f"  aws sagemaker-runtime invoke-endpoint-async \\")
+    log(f"    --endpoint-name {endpoint_name} \\")
+    log(f"    --input-location s3://YOUR-INPUT-BUCKET/path/to/input.json \\")
+    log(f"    --content-type application/json --region {args.region}")
+    log(f"  # Result will land at: {args.output_s3_uri}")
+    log("  # Write input.json as BOM-free UTF-8 (on Windows, NOT 'Set-Content -Encoding UTF8')")
+    log("")
+    log(f"Teardown: python3 teardown.py {endpoint_name} {args.region}")
+
+    print(json.dumps({
+        "endpoint_name": endpoint_name,
+        "endpoint_config_name": config_name,
+        "model_name": args.model_name,
+        "region": args.region,
+        "instance_type": args.instance_type,
+        "inference_mode": "async",
+        "output_s3_uri": args.output_s3_uri,
+    }))
+
+    return 0
+
+
+if __name__ == "__main__":
+    sys.exit(main())