chore: import upstream snapshot with attribution

examples/provider-amazon-sagemaker/deploy-test-model.py

@@ -0,0 +1,225 @@
+#!/usr/bin/env python3
+"""
+SageMaker Deployment Helper Script
+
+This script helps deploy a test model on Amazon SageMaker for testing the promptfoo SageMaker provider.
+It uses the Hugging Face integration with SageMaker to deploy models.
+
+Prerequisites:
+- AWS CLI configured
+- Required Python packages: sagemaker, boto3
+- SageMaker execution role with appropriate permissions
+
+Usage:
+  python deploy-test-model.py --model-id meta-llama/Llama-2-7b-chat-hf --task text-generation
+
+"""
+
+import argparse
+import json
+import logging
+import time
+from datetime import datetime
+
+import boto3
+
+# Configure logging
+logging.basicConfig(
+    level=logging.INFO, format="%(asctime)s - %(levelname)s - %(message)s"
+)
+logger = logging.getLogger(__name__)
+
+# Parse command line arguments
+parser = argparse.ArgumentParser(description="Deploy a Hugging Face model to SageMaker")
+parser.add_argument(
+    "--model-id",
+    required=True,
+    help="Hugging Face model ID (e.g., meta-llama/Llama-2-7b-chat-hf)",
+)
+parser.add_argument(
+    "--task",
+    default="text-generation",
+    help="Task of the model (default: text-generation)",
+)
+parser.add_argument(
+    "--instance-type",
+    default="ml.g5.2xlarge",
+    help="SageMaker instance type (default: ml.g5.2xlarge)",
+)
+parser.add_argument(
+    "--endpoint-name", help="Custom endpoint name (default: based on model name)"
+)
+parser.add_argument("--region", help="AWS region (default: from AWS CLI configuration)")
+parser.add_argument(
+    "--role-name",
+    help="SageMaker execution IAM role name (if not specified, will try to find one)",
+)
+args = parser.parse_args()
+
+# Get the AWS region
+session = boto3.session.Session()
+region = args.region or session.region_name
+if not region:
+    logger.error("AWS region not specified and not found in AWS CLI configuration")
+    raise SystemExit(1)
+
+# Generate endpoint name if not provided
+if not args.endpoint_name:
+    # Extract model name from model ID and create a timestamp
+    model_name = args.model_id.split("/")[-1].lower()
+    timestamp = datetime.now().strftime("%m%d%H%M")
+    args.endpoint_name = f"promptfoo-test-{model_name}-{timestamp}"
+
+# Connect to AWS services
+iam = boto3.client("iam", region_name=region)
+sagemaker_client = boto3.client("sagemaker", region_name=region)
+
+
+# Find or get SageMaker execution role
+def get_sagemaker_role() -> str:
+    """Return the requested role ARN or discover a SageMaker execution role."""
+    if args.role_name:
+        try:
+            role = iam.get_role(RoleName=args.role_name)
+            return role["Role"]["Arn"]
+        except Exception as e:
+            logger.error(f"Error getting specified role: {e}")
+            raise SystemExit(1)
+
+    # Try to find a SageMaker execution role
+    try:
+        roles = iam.list_roles()
+        for role in roles["Roles"]:
+            if "AmazonSageMaker-ExecutionRole" in role["RoleName"]:
+                logger.info(f"Found SageMaker role: {role['RoleName']}")
+                return role["Arn"]
+
+        logger.error(
+            "No SageMaker execution role found. Please specify a role with --role-name"
+        )
+        raise SystemExit(1)
+    except Exception as e:
+        logger.error(f"Error finding SageMaker role: {e}")
+        raise SystemExit(1)
+
+
+role_arn = get_sagemaker_role()
+logger.info(f"Using role ARN: {role_arn}")
+
+# Create model
+try:
+    logger.info(f"Creating model: {args.model_id}")
+
+    # HuggingFace Container settings
+    # Find the latest container image for the region
+    transformers_version = "4.28.1"
+    pytorch_version = "2.0.0"
+    python_version = "py310"
+
+    hf_container = f"763104351884.dkr.ecr.{region}.amazonaws.com/huggingface-pytorch-tgi:{transformers_version}-transformers{pytorch_version}-cuda11.8-{python_version}"
+
+    # Create model
+    model_name = f"promptfoo-test-model-{int(time.time())}"
+
+    # Hub config for text generation interface
+    hub_config = {"HF_MODEL_ID": args.model_id, "HF_TASK": args.task}
+
+    # Generation parameters for text generation models
+    if args.task == "text-generation":
+        hub_config["PARAMETERS"] = {
+            "max_new_tokens": 256,
+            "temperature": 0.7,
+            "return_full_text": False,
+        }
+
+    # Create the model
+    create_model_response = sagemaker_client.create_model(
+        ModelName=model_name,
+        PrimaryContainer={
+            "Image": hf_container,
+            "Environment": {
+                "HF_MODEL_ID": args.model_id,
+                "HF_TASK": args.task,
+                "HF_MODEL_QUANTIZE": "bitsandbytes",  # Optional: for quantization
+                "SM_NUM_GPUS": "1",
+                "MAX_INPUT_LENGTH": "2048",
+                "MAX_TOTAL_TOKENS": "4096",
+                "HF_HUB_CONFIG": json.dumps(hub_config),
+            },
+        },
+        ExecutionRoleArn=role_arn,
+    )
+
+    logger.info(f"Model created: {model_name}")
+
+    # Create endpoint configuration
+    endpoint_config_name = f"{args.endpoint_name}-config"
+
+    logger.info(f"Creating endpoint configuration: {endpoint_config_name}")
+    sagemaker_client.create_endpoint_config(
+        EndpointConfigName=endpoint_config_name,
+        ProductionVariants=[
+            {
+                "VariantName": "AllTraffic",
+                "ModelName": model_name,
+                "InstanceType": args.instance_type,
+                "InitialInstanceCount": 1,
+            }
+        ],
+    )
+
+    # Create endpoint
+    logger.info(
+        f"Creating endpoint: {args.endpoint_name} (this will take several minutes)"
+    )
+    sagemaker_client.create_endpoint(
+        EndpointName=args.endpoint_name, EndpointConfigName=endpoint_config_name
+    )
+
+    # Wait for endpoint to be in service
+    logger.info("Waiting for endpoint to be in service...")
+
+    status = None
+    while status != "InService":
+        response = sagemaker_client.describe_endpoint(EndpointName=args.endpoint_name)
+        status = response["EndpointStatus"]
+
+        if status == "Failed":
+            logger.error(
+                f"Endpoint creation failed: {response.get('FailureReason', 'Unknown reason')}"
+            )
+            raise SystemExit(1)
+
+        if status != "InService":
+            logger.info(f"Endpoint status: {status}. Waiting...")
+            time.sleep(60)
+
+    logger.info(f"✅ Endpoint {args.endpoint_name} is now InService!")
+    logger.info("\nTest with promptfoo using:")
+    logger.info(
+        f"""
+providers:
+  - id: sagemaker:{args.task.replace("-", ":")}:{args.endpoint_name}
+    config:
+      region: {region}
+      modelType: {"openai" if "llama" in args.model_id.lower() else "custom"}
+      maxTokens: 256
+      temperature: 0.7
+"""
+    )
+
+    logger.info("\nOr use the test script:")
+    logger.info(
+        f"node test-sagemaker-provider.js --endpoint={args.endpoint_name} --region={region} --model-type={'openai' if 'llama' in args.model_id.lower() else 'custom'}"
+    )
+
+    logger.info(
+        "\nTo delete this endpoint when done testing (to avoid unnecessary charges):"
+    )
+    logger.info(
+        f"aws sagemaker delete-endpoint --endpoint-name {args.endpoint_name} --region {region}"
+    )
+
+except Exception as e:
+    logger.error(f"Error deploying model: {e}")
+    raise SystemExit(1)