vLLM Deployment And Infrastructure For MangaAssist

# Base CUDA environment (pre-installed in our Docker base image)
# NVIDIA Driver: 535.104.12
# CUDA Toolkit: 12.1

# vLLM and dependencies
pip install vllm==0.4.3 \
    torch==2.3.0 \
    transformers==4.40.0 \
    safetensors==0.4.3 \
    uvicorn==0.29.0

# For AWQ quantization (offline, not needed on serving nodes)
pip install autoawq==0.2.4

# For LoRA adapter training (offline, not needed on serving nodes)
pip install peft==0.10.0

┌────────────────────────────────────────────────────────┐
│                    VRAM Budget (24 GB)                  │
├────────────────────────────────────────────────────────┤
│ Model weights (AWQ INT4 Llama-3-8B)          4.5 GB   │
│ LoRA adapters (3 × ~40 MB)                   0.12 GB  │
│ KV cache blocks (128 seqs × block_size=16)  ~14.0 GB  │
│ CUDA workspace + activations                  1.5 GB   │
│ CUDA graph memory                             1.0 GB   │
│ Safety headroom (gpu_memory_utilization=0.92) 1.92 GB  │
│ Unallocated (OS/driver overhead)              0.96 GB  │
├────────────────────────────────────────────────────────┤
│ Total                                        24.0 GB   │
└────────────────────────────────────────────────────────┘

# ======================================
# Stage 1: Build vLLM with CUDA support
# ======================================
FROM nvidia/cuda:12.1.1-devel-ubuntu22.04 AS builder

RUN apt-get update && apt-get install -y \
    python3.10 python3.10-dev python3-pip git && \
    rm -rf /var/lib/apt/lists/*

COPY requirements.txt /build/requirements.txt
RUN pip install --no-cache-dir -r /build/requirements.txt

# ======================================
# Stage 2: Download and validate model
# ======================================
FROM builder AS model-prep

# Model artifacts are downloaded during image build, NOT at container runtime.
# This is a critical decision: runtime download adds 3-5 minutes to startup.
ARG MODEL_ID=meta-llama/Llama-3-8b-instruct-awq
ARG HF_TOKEN

RUN python3 -c "
from huggingface_hub import snapshot_download
snapshot_download(
    repo_id='${MODEL_ID}',
    local_dir='/models/base',
    token='${HF_TOKEN}',
    ignore_patterns=['*.md', '*.txt', 'original/*'],
)
"

# Download LoRA adapters
COPY scripts/download_adapters.py /build/download_adapters.py
RUN python3 /build/download_adapters.py \
    --output-dir /models/lora \
    --adapters manga_domain_v3 general_support_v2 jp_style_v1

# Validate model integrity
RUN python3 -c "
from transformers import AutoTokenizer
tokenizer = AutoTokenizer.from_pretrained('/models/base')
assert tokenizer.vocab_size > 0, 'Tokenizer validation failed'
print(f'Model validated: vocab_size={tokenizer.vocab_size}')
"

# ======================================
# Stage 3: Slim runtime image
# ======================================
FROM nvidia/cuda:12.1.1-runtime-ubuntu22.04 AS runtime

RUN apt-get update && apt-get install -y \
    python3.10 python3-pip && \
    rm -rf /var/lib/apt/lists/*

# Copy only runtime dependencies (no build tools)
COPY --from=builder /usr/local/lib/python3.10/dist-packages /usr/local/lib/python3.10/dist-packages
COPY --from=builder /usr/local/bin /usr/local/bin

# Copy validated model artifacts
COPY --from=model-prep /models /opt/ml/model

# Copy application code
COPY app/ /app/
COPY scripts/entrypoint.sh /entrypoint.sh

# Health check port
EXPOSE 8080

# Prometheus metrics port
EXPOSE 9090

ENV PYTHONUNBUFFERED=1
ENV MODEL_PATH=/opt/ml/model/base
ENV LORA_PATH=/opt/ml/model/lora
ENV VLLM_PORT=8080

HEALTHCHECK --interval=30s --timeout=10s --retries=3 \
    CMD curl -f http://localhost:8080/ping || exit 1

ENTRYPOINT ["/entrypoint.sh"]

#!/bin/bash
set -euo pipefail

echo "[$(date)] Starting vLLM server for MangaAssist"

# Phase 1: Environment validation
echo "[$(date)] Phase 1: Validating environment"
nvidia-smi --query-gpu=name,memory.total,driver_version --format=csv,noheader
python3 -c "import torch; assert torch.cuda.is_available(), 'CUDA not available'"

# Phase 2: Start vLLM OpenAI-compatible server
echo "[$(date)] Phase 2: Starting vLLM engine"
python3 -m vllm.entrypoints.openai.api_server \
    --model "${MODEL_PATH}" \
    --host 0.0.0.0 \
    --port "${VLLM_PORT:-8080}" \
    --tensor-parallel-size 1 \
    --gpu-memory-utilization 0.92 \
    --max-model-len 4096 \
    --max-num-seqs 128 \
    --max-num-batched-tokens 8192 \
    --block-size 16 \
    --enable-prefix-caching \
    --quantization awq \
    --no-custom-all-reduce \
    --enable-lora \
    --max-loras 4 \
    --max-lora-rank 16 \
    --lora-modules \
        manga_domain_v3="${LORA_PATH}/manga_domain_v3" \
        general_support_v2="${LORA_PATH}/general_support_v2" \
        jp_style_v1="${LORA_PATH}/jp_style_v1" \
    --disable-log-requests \
    2>&1 | tee /var/log/vllm/server.log &

VLLM_PID=$!

# Phase 3: Wait for readiness
echo "[$(date)] Phase 3: Waiting for engine readiness"
MAX_WAIT=180
WAITED=0
while [ $WAITED -lt $MAX_WAIT ]; do
    if curl -sf http://localhost:${VLLM_PORT}/health > /dev/null 2>&1; then
        echo "[$(date)] Engine is ready after ${WAITED}s"
        break
    fi
    sleep 2
    WAITED=$((WAITED + 2))
done

if [ $WAITED -ge $MAX_WAIT ]; then
    echo "[$(date)] ERROR: Engine failed to become ready in ${MAX_WAIT}s"
    kill $VLLM_PID 2>/dev/null
    exit 1
fi

# Phase 4: Warmup requests
echo "[$(date)] Phase 4: Running warmup requests"
python3 /app/scripts/warmup.py \
    --endpoint "http://localhost:${VLLM_PORT}/v1/chat/completions" \
    --num-requests 10 \
    --timeout 30

echo "[$(date)] Phase 5: Server is ready to accept traffic"

# Handle graceful shutdown
trap 'echo "[$(date)] SIGTERM received, draining..."; kill -SIGTERM $VLLM_PID; wait $VLLM_PID' SIGTERM

wait $VLLM_PID

"""
Warmup script that sends synthetic requests to prime the vLLM engine.

Purpose:
- Trigger CUDA graph capture for common sequence lengths
- Prime the KV block allocator
- Warm the prefix cache with the system prompt
- Validate that all LoRA adapters load correctly
"""

import argparse
import time

import httpx


WARMUP_REQUESTS = [
    {
        "model": "manga_domain_v3",
        "messages": [
            {"role": "system", "content": "You are a manga shopping assistant."},
            {"role": "user", "content": "Recommend action manga."},
        ],
        "max_tokens": 50,
        "stream": False,
    },
    {
        "model": "general_support_v2",
        "messages": [
            {"role": "system", "content": "You are a manga shopping assistant."},
            {"role": "user", "content": "What is the return policy?"},
        ],
        "max_tokens": 50,
        "stream": False,
    },
    {
        "model": "jp_style_v1",
        "messages": [
            {"role": "system", "content": "You are a manga shopping assistant."},
            {"role": "user", "content": "Recommend shounen manga in Japanese."},
        ],
        "max_tokens": 50,
        "stream": False,
    },
]


def run_warmup(endpoint: str, num_requests: int, timeout: int) -> None:
    client = httpx.Client(timeout=timeout)
    start = time.time()

    for i in range(num_requests):
        request = WARMUP_REQUESTS[i % len(WARMUP_REQUESTS)]
        try:
            response = client.post(endpoint, json=request)
            response.raise_for_status()
            elapsed = time.time() - start
            print(f"  Warmup {i+1}/{num_requests}: {response.status_code} ({elapsed:.1f}s)")
        except Exception as e:
            print(f"  Warmup {i+1}/{num_requests}: FAILED - {e}")
            raise

    total = time.time() - start
    print(f"Warmup complete: {num_requests} requests in {total:.1f}s")


if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument("--endpoint", required=True)
    parser.add_argument("--num-requests", type=int, default=10)
    parser.add_argument("--timeout", type=int, default=30)
    args = parser.parse_args()
    run_warmup(args.endpoint, args.num_requests, args.timeout)

"""
Health check implementation for vLLM on SageMaker.

Three levels of health:
1. Liveness: Process is running and can respond to HTTP
2. Readiness: Model is loaded, CUDA graphs captured, warmup complete
3. GPU health: CUDA device is accessible and not in error state
"""

import asyncio
import subprocess
from enum import Enum

from fastapi import FastAPI, Response


class HealthState(Enum):
    INITIALIZING = "initializing"
    WARMING_UP = "warming_up"
    READY = "ready"
    DEGRADED = "degraded"
    UNHEALTHY = "unhealthy"


class HealthChecker:
    def __init__(self) -> None:
        self.state = HealthState.INITIALIZING
        self.model_loaded = False
        self.cuda_graphs_ready = False
        self.warmup_complete = False
        self.last_gpu_check_ok = True
        self._gpu_check_interval = 60  # seconds

    def mark_model_loaded(self) -> None:
        self.model_loaded = True
        self._update_state()

    def mark_cuda_graphs_ready(self) -> None:
        self.cuda_graphs_ready = True
        self._update_state()

    def mark_warmup_complete(self) -> None:
        self.warmup_complete = True
        self._update_state()

    def _update_state(self) -> None:
        if self.model_loaded and self.cuda_graphs_ready and self.warmup_complete:
            self.state = HealthState.READY
        elif self.model_loaded:
            self.state = HealthState.WARMING_UP

    async def check_gpu_health(self) -> bool:
        """Check GPU accessibility via nvidia-smi. Run periodically, not per-request."""
        try:
            result = subprocess.run(
                ["nvidia-smi", "--query-gpu=gpu_bus_id", "--format=csv,noheader"],
                capture_output=True,
                text=True,
                timeout=5,
            )
            self.last_gpu_check_ok = result.returncode == 0
        except (subprocess.TimeoutExpired, FileNotFoundError):
            self.last_gpu_check_ok = False
        return self.last_gpu_check_ok

    def is_ready(self) -> bool:
        return self.state == HealthState.READY and self.last_gpu_check_ok

    def is_alive(self) -> bool:
        return self.state != HealthState.UNHEALTHY


health_checker = HealthChecker()


def register_health_routes(app: FastAPI) -> None:
    @app.get("/ping")
    async def ping(response: Response) -> dict:
        """SageMaker liveness check. Returns 200 if process is alive."""
        if not health_checker.is_alive():
            response.status_code = 503
        return {
            "status": health_checker.state.value,
            "model_loaded": health_checker.model_loaded,
            "cuda_graphs_ready": health_checker.cuda_graphs_ready,
            "warmup_complete": health_checker.warmup_complete,
            "gpu_ok": health_checker.last_gpu_check_ok,
        }

    @app.get("/ready")
    async def ready(response: Response) -> dict:
        """Readiness check. Returns 200 only when fully ready to serve."""
        if not health_checker.is_ready():
            response.status_code = 503
        return {"ready": health_checker.is_ready()}

sequenceDiagram
    participant SM as SageMaker
    participant EP as Entrypoint
    participant GW as vLLM Gateway
    participant ENGINE as vLLM Engine

    SM->>EP: SIGTERM
    EP->>GW: Stop accepting new requests
    Note over GW: /ping returns 503\n/ready returns 503
    GW->>ENGINE: Wait for in-flight requests (max 30s)
    Note over ENGINE: Active generations\ncomplete normally
    ENGINE-->>GW: All generations done
    GW->>EP: Clean exit
    EP-->>SM: Exit 0
    Note over SM: Instance removed\nfrom endpoint

import asyncio
import signal
import logging

logger = logging.getLogger(__name__)


class GracefulShutdown:
    def __init__(self, engine, max_drain_seconds: int = 30) -> None:
        self.engine = engine
        self.max_drain_seconds = max_drain_seconds
        self.shutting_down = False

    def install_signal_handlers(self) -> None:
        loop = asyncio.get_event_loop()
        for sig in (signal.SIGTERM, signal.SIGINT):
            loop.add_signal_handler(sig, lambda s=sig: asyncio.create_task(self.shutdown(s)))

    async def shutdown(self, sig: signal.Signals) -> None:
        logger.info(f"Received {sig.name}, starting graceful shutdown")
        self.shutting_down = True

        # Stop accepting new requests (health checks return 503)
        logger.info("Stopped accepting new requests")

        # Wait for in-flight requests to complete
        waited = 0
        while waited < self.max_drain_seconds:
            active = self.engine.get_num_unfinished_requests()
            if active == 0:
                logger.info(f"All requests drained after {waited}s")
                break
            logger.info(f"Draining: {active} requests still in-flight ({waited}s)")
            await asyncio.sleep(1)
            waited += 1

        if waited >= self.max_drain_seconds:
            active = self.engine.get_num_unfinished_requests()
            logger.warning(f"Drain timeout: {active} requests still active, forcing shutdown")

        logger.info("Shutdown complete")

import sagemaker
from sagemaker.model import Model
from sagemaker.predictor import Predictor


def deploy_vllm_endpoint(
    role: str,
    image_uri: str,
    model_data: str = None,  # None when model is baked into image
    instance_type: str = "ml.g5.xlarge",
    initial_instance_count: int = 4,
) -> Predictor:
    model = Model(
        image_uri=image_uri,
        role=role,
        env={
            "MODEL_PATH": "/opt/ml/model/base",
            "LORA_PATH": "/opt/ml/model/lora",
            "VLLM_PORT": "8080",
        },
    )

    predictor = model.deploy(
        initial_instance_count=initial_instance_count,
        instance_type=instance_type,
        endpoint_name="mangaassist-vllm-prod",
        container_startup_health_check_timeout=300,  # 5 min for model load + warmup
        volume_size=50,  # GB, for model artifacts if not baked in
    )

    return predictor

from sagemaker.session import production_variant


def deploy_canary(image_uri_prod: str, image_uri_canary: str, role: str) -> None:
    """Deploy canary variant alongside production for safe rollouts."""
    variants = [
        production_variant(
            model_name="mangaassist-vllm-prod",
            instance_type="ml.g5.xlarge",
            initial_instance_count=4,
            variant_name="primary",
            initial_weight=95,  # 95% of traffic
        ),
        production_variant(
            model_name="mangaassist-vllm-canary",
            instance_type="ml.g5.xlarge",
            initial_instance_count=1,
            variant_name="canary",
            initial_weight=5,  # 5% of traffic
        ),
    ]

    sagemaker.Session().endpoint_from_production_variants(
        name="mangaassist-vllm-prod",
        production_variants=variants,
    )

import boto3

sm_client = boto3.client("sagemaker")

# Enable warm pool for faster scaling
sm_client.update_endpoint(
    EndpointName="mangaassist-vllm-prod",
    RetainAllVariantProperties=True,
    DeploymentConfig={
        "BlueGreenUpdatePolicy": {
            "TrafficRoutingConfiguration": {
                "Type": "ALL_AT_ONCE",
                "WaitIntervalInSeconds": 60,
            },
            "TerminationWaitInSeconds": 120,
            "MaximumExecutionTimeoutInSeconds": 600,
        },
    },
)

import boto3

asg_client = boto3.client("application-autoscaling")

# Register scalable target
asg_client.register_scalable_target(
    ServiceNamespace="sagemaker",
    ResourceId="endpoint/mangaassist-vllm-prod/variant/primary",
    ScalableDimension="sagemaker:variant:DesiredInstanceCount",
    MinCapacity=2,   # Always keep 2 for redundancy
    MaxCapacity=8,   # Cost ceiling
)

# Scale-out policy: react fast to queue buildup
asg_client.put_scaling_policy(
    PolicyName="mangaassist-vllm-scale-out",
    ServiceNamespace="sagemaker",
    ResourceId="endpoint/mangaassist-vllm-prod/variant/primary",
    ScalableDimension="sagemaker:variant:DesiredInstanceCount",
    PolicyType="TargetTrackingScaling",
    TargetTrackingScalingPolicyConfiguration={
        "TargetValue": 50.0,  # Target invocations per instance
        "PredefinedMetricSpecification": {
            "PredefinedMetricType": "SageMakerVariantInvocationsPerInstance",
        },
        "ScaleOutCooldown": 60,   # Fast reaction to spikes
        "ScaleInCooldown": 300,   # Slow scale-in to avoid flapping
    },
)

# Pre-provision for JP evening peak
asg_client.put_scheduled_action(
    ServiceNamespace="sagemaker",
    ScheduledActionName="jp-evening-peak",
    ResourceId="endpoint/mangaassist-vllm-prod/variant/primary",
    ScalableDimension="sagemaker:variant:DesiredInstanceCount",
    Schedule="cron(45 8 * * ? *)",  # 17:45 JST = 08:45 UTC
    ScalableTargetAction={"MinCapacity": 6, "MaxCapacity": 8},
)

# Scale back after JP evening
asg_client.put_scheduled_action(
    ServiceNamespace="sagemaker",
    ScheduledActionName="jp-evening-end",
    ResourceId="endpoint/mangaassist-vllm-prod/variant/primary",
    ScalableDimension="sagemaker:variant:DesiredInstanceCount",
    Schedule="cron(0 14 * * ? *)",  # 23:00 JST = 14:00 UTC
    ScalableTargetAction={"MinCapacity": 2, "MaxCapacity": 8},
)

from fastapi import FastAPI, Request, Response
import httpx
import os

app = FastAPI()
vllm_base = os.getenv("VLLM_BASE_URL", "http://127.0.0.1:8080")


@app.get("/ping")
async def ping() -> Response:
    async with httpx.AsyncClient(timeout=2.0) as client:
        upstream = await client.get(f"{vllm_base}/health")
    return Response(status_code=200 if upstream.status_code == 200 else 503)


@app.post("/invocations")
async def invocations(request: Request) -> Response:
    payload = await request.json()

    vllm_payload = {
        "model": payload.get("model", "manga_domain_v3"),
        "messages": payload["messages"],
        "max_tokens": payload.get("max_tokens", 256),
        "temperature": payload.get("temperature", 0.2),
        "stream": payload.get("stream", False),
    }

    async with httpx.AsyncClient(timeout=70.0) as client:
        upstream = await client.post(
            f"{vllm_base}/v1/chat/completions",
            json=vllm_payload,
        )

    return Response(
        content=upstream.text,
        status_code=upstream.status_code,
        media_type="application/json",
    )

s3://mangaassist-bedrock-import/llama3-domain-v1/
  config.json
  generation_config.json
  tokenizer.json
  tokenizer_config.json
  tokenizer.model
  model-00001-of-00004.safetensors
  model-00002-of-00004.safetensors
  model-00003-of-00004.safetensors
  model-00004-of-00004.safetensors
  model.safetensors.index.json

import boto3

bedrock = boto3.client("bedrock", region_name="us-east-1")

response = bedrock.create_model_import_job(
    jobName="mangaassist-llama3-import-v1",
    importedModelName="mangaassist-llama3-domain-v1",
    roleArn="arn:aws:iam::123456789012:role/bedrock-model-import-role",
    modelDataSource={
        "s3DataSource": {
            "s3Uri": "s3://mangaassist-bedrock-import/llama3-domain-v1/"
        }
    },
)

job_arn = response["jobArn"]
print(job_arn)

import time

while True:
    status = bedrock.get_model_import_job(jobIdentifier=job_arn)
    print(status["status"])

    if status["status"] == "Completed":
        imported_model_arn = status["importedModelArn"]
        break
    if status["status"] == "Failed":
        raise RuntimeError(status["failureMessage"])

    time.sleep(30)

import boto3
import json

runtime = boto3.client("bedrock-runtime", region_name="us-east-1")

response = runtime.converse(
    modelId=imported_model_arn,
    system=[
        {
            "text": "You are a manga shopping assistant."
        }
    ],
    messages=[
        {
            "role": "user",
            "content": [
                {
                    "text": "Recommend action manga under $20."
                }
            ],
        }
    ],
    inferenceConfig={
        "maxTokens": 200,
        "temperature": 0.2,
    },
)

print(json.dumps(response, indent=2, default=str))