@aws/ml-container-creator 0.2.0

package/templates/code/serve

@@ -0,0 +1,300 @@
+#!/bin/bash
+# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+# SPDX-License-Identifier: Apache-2.0
+
+<% if (modelServer === 'vllm') { %>
+echo "Starting vLLM server"
+<% } else if (modelServer === 'sglang') { %>
+echo "Starting SGLang server"
+<% } else if (modelServer === 'tensorrt-llm') { %>
+echo "Starting TensorRT-LLM server"
+<% } else if (modelServer === 'lmi') { %>
+echo "Starting LMI (Large Model Inference) server"
+<% } else if (modelServer === 'djl') { %>
+echo "Starting DJL Serving server"
+<% } %>
+
+<% if (modelServer === 'lmi' || modelServer === 'djl') { %>
+# LMI/DJL containers use serving.properties for configuration
+# The configuration file should be at /opt/ml/model/serving.properties
+# DJL Serving will automatically start with this configuration
+
+if [ ! -f /opt/ml/model/serving.properties ]; then
+    echo "Error: serving.properties not found at /opt/ml/model/serving.properties"
+    exit 1
+fi
+
+echo "Using configuration from /opt/ml/model/serving.properties"
+cat /opt/ml/model/serving.properties
+
+# DJL Serving is already configured in the base image
+# This script is not typically needed for LMI/DJL as they have their own entrypoint
+# But we provide it for consistency with other model servers
+exit 0
+<% } else { %>
+
+<% if (typeof modelSource !== 'undefined' && modelSource !== 'huggingface') { %>
+# ---------------------------------------------------------------------------
+# download_model_from_s3 — Download model artifacts from S3 to a local path
+# ---------------------------------------------------------------------------
+download_model_from_s3() {
+    local s3_uri="$1"
+    local dest_path="$2"
+    local start_time
+    start_time=$(date +%s)
+
+    if [ -z "$s3_uri" ] || [ -z "$dest_path" ]; then
+        echo "Error: download_model_from_s3 requires S3 URI and destination path" >&2
+        return 1
+    fi
+
+    echo "Downloading model from ${s3_uri} to ${dest_path}..." >&2
+    mkdir -p "${dest_path}"
+
+    if [[ "$s3_uri" == *.tar.gz ]] || [[ "$s3_uri" == *.tgz ]]; then
+        # Tarball: download and extract
+        if ! aws s3 cp "$s3_uri" /tmp/model_archive.tar.gz; then
+            echo "Error: Failed to download tarball from ${s3_uri}" >&2
+            return 1
+        fi
+        if ! tar -xzf /tmp/model_archive.tar.gz -C "$dest_path"; then
+            echo "Error: Failed to extract tarball from ${s3_uri}" >&2
+            rm -f /tmp/model_archive.tar.gz
+            return 1
+        fi
+        rm -f /tmp/model_archive.tar.gz
+    elif [[ "$s3_uri" == */ ]] || ! aws s3 ls "$s3_uri" 2>/dev/null | grep -q "^[0-9]"; then
+        # Directory prefix: sync
+        if ! aws s3 sync "$s3_uri" "$dest_path"; then
+            echo "Error: Failed to sync from ${s3_uri}" >&2
+            return 1
+        fi
+    else
+        # Single file: copy
+        if ! aws s3 cp "$s3_uri" "$dest_path/"; then
+            echo "Error: Failed to copy ${s3_uri}" >&2
+            return 1
+        fi
+    fi
+
+    local duration
+    duration=$(( $(date +%s) - start_time ))
+    echo "Download complete: ${s3_uri} → ${dest_path} (${duration}s)" >&2
+}
+<% } %>
+
+# ---------------------------------------------------------------------------
+# Model Loading Adapter — resolve model based on MODEL_SOURCE env var
+# ---------------------------------------------------------------------------
+MODEL_SOURCE="${MODEL_SOURCE:-huggingface}"
+MODEL_ARTIFACT_URI="${MODEL_ARTIFACT_URI:-}"
+LOCAL_MODEL_PATH="/opt/ml/model"
+
+<% if (modelServer === 'vllm') { %>
+_MODEL_VAR="VLLM_MODEL"
+<% } else if (modelServer === 'sglang') { %>
+_MODEL_VAR="SGLANG_MODEL_PATH"
+<% } else if (modelServer === 'tensorrt-llm') { %>
+_MODEL_VAR="TRTLLM_MODEL"
+<% } %>
+
+resolve_model() {
+    case "$MODEL_SOURCE" in
+        huggingface)
+            # Pass model name directly — server fetches from HF Hub
+            echo "${!_MODEL_VAR}"
+            return
+            ;;
+        s3|jumpstart|jumpstart-hub|registry)
+            # Check for pre-mounted artifacts first
+            if [ -d "$LOCAL_MODEL_PATH" ] && [ "$(ls -A $LOCAL_MODEL_PATH 2>/dev/null)" ]; then
+                echo "Using pre-mounted model artifacts at $LOCAL_MODEL_PATH" >&2
+                echo "$LOCAL_MODEL_PATH"
+                return
+            fi
+
+            # For registry:// models, resolve artifact URI at runtime via SageMaker API
+            if [ "$MODEL_SOURCE" = "registry" ] && [ -z "$MODEL_ARTIFACT_URI" ]; then
+                local model_uri="${!_MODEL_VAR}"
+                local registry_prefix="registry://"
+                if [[ "$model_uri" == "${registry_prefix}"* ]]; then
+                    local registry_path="${model_uri#${registry_prefix}}"
+                    local group_name="${registry_path%%/*}"
+                    local version="${registry_path#*/}"
+                    local region="${AWS_REGION:-${AWS_DEFAULT_REGION:-us-east-1}}"
+
+                    # Get account ID for ARN construction
+                    local account_id
+                    account_id=$(aws sts get-caller-identity --query Account --output text 2>/dev/null) || {
+                        echo "Error: Failed to get AWS account ID for model package ARN" >&2
+                        exit 1
+                    }
+
+                    local package_arn="arn:aws:sagemaker:${region}:${account_id}:model-package/${group_name}/${version}"
+                    echo "Resolving ${model_uri} via SageMaker DescribeModelPackage..." >&2
+                    echo "   ARN: ${package_arn}" >&2
+
+                    local describe_output
+                    describe_output=$(aws sagemaker describe-model-package \
+                        --model-package-name "$package_arn" \
+                        --region "$region" \
+                        --output json 2>/dev/null) || {
+                        echo "Error: Failed to describe model package: ${package_arn}" >&2
+                        exit 1
+                    }
+
+                    # Try ModelDataUrl first, then S3DataSource.S3Uri, then description
+                    MODEL_ARTIFACT_URI=$(echo "$describe_output" | python3 -c "
+import sys, json, re
+try:
+    pkg = json.load(sys.stdin)
+    uri = ''
+    # Check InferenceSpecification.Containers[0]
+    containers = pkg.get('InferenceSpecification', {}).get('Containers', [])
+    if containers:
+        c = containers[0]
+        uri = c.get('ModelDataUrl', '')
+        if not uri:
+            uri = c.get('ModelDataSource', {}).get('S3DataSource', {}).get('S3Uri', '')
+    # Fallback: extract S3 URI from ModelPackageDescription
+    if not uri:
+        desc = pkg.get('ModelPackageDescription', '')
+        m = re.search(r's3://[^\s]+', desc)
+        if m:
+            uri = m.group(0)
+    # Fallback: check ModelCard hyperparameters for model_artifacts_s3
+    if not uri:
+        try:
+            card = pkg.get('ModelCard', {})
+            content = card.get('ModelCardContent', '{}')
+            card_data = json.loads(content) if isinstance(content, str) else content
+            params = card_data.get('training_details', {}).get('training_job_details', {}).get('hyper_parameters', [])
+            for p in params:
+                if p.get('name') == 'model_artifacts_s3':
+                    uri = p.get('value', '')
+                    break
+        except:
+            pass
+    print(uri)
+except:
+    print('')
+" 2>/dev/null)
+
+                    if [ -n "$MODEL_ARTIFACT_URI" ] && [ "$MODEL_ARTIFACT_URI" != "None" ]; then
+                        echo "Resolved artifact URI: ${MODEL_ARTIFACT_URI}" >&2
+                    else
+                        echo "Error: No model artifact URI found in model package: ${package_arn}" >&2
+                        echo "   Checked: InferenceSpecification.Containers[0].ModelDataUrl" >&2
+                        echo "   Checked: InferenceSpecification.Containers[0].ModelDataSource.S3DataSource.S3Uri" >&2
+                        exit 1
+                    fi
+                fi
+            fi
+
+            # Need artifact URI for download
+            if [ -z "$MODEL_ARTIFACT_URI" ]; then
+                echo "Error: ${MODEL_SOURCE} model requires artifact URI or pre-mounted artifacts at $LOCAL_MODEL_PATH" >&2
+                exit 1
+            fi
+            # Download from S3
+            if ! download_model_from_s3 "$MODEL_ARTIFACT_URI" "$LOCAL_MODEL_PATH"; then
+                echo "Error: Failed to download model from ${MODEL_ARTIFACT_URI}" >&2
+                exit 1
+            fi
+            echo "$LOCAL_MODEL_PATH"
+            ;;
+        *)
+            # Unrecognized source — treat as huggingface
+            echo "${!_MODEL_VAR}"
+            return
+            ;;
+    esac
+}
+
+_RESOLVED_MODEL=$(resolve_model) || exit 1
+export "${_MODEL_VAR}=${_RESOLVED_MODEL}"
+echo "Resolved ${_MODEL_VAR}=${_RESOLVED_MODEL} (source: ${MODEL_SOURCE})"
+unset _MODEL_VAR _RESOLVED_MODEL
+
+# Initialize server arguments
+<% if (modelServer === 'tensorrt-llm') { %>
+# port 8081 for internal TensorRT-LLM server (nginx proxies on 8080)
+SERVER_ARGS=(--host 0.0.0.0 --port 8081)
+<% } else { %>
+# port 8080 required by SageMaker: https://docs.aws.amazon.com/sagemaker/latest/dg/your-algorithms-inference-code.html#your-algorithms-inference-code-container-response
+SERVER_ARGS=(--host 0.0.0.0 --port 8080)
+<% } %>
+
+# Define the prefix for environment variables to look for
+<% if (modelServer === 'vllm') { %>
+PREFIX="VLLM_"
+<% } else if (modelServer === 'sglang') { %>
+PREFIX="SGLANG_"
+<% } else if (modelServer === 'tensorrt-llm') { %>
+PREFIX="TRTLLM_"
+<% } %>
+ARG_PREFIX="--"
+
+# Define environment variables to exclude (internal variables set by base images)
+<% if (modelServer === 'vllm') { %>
+EXCLUDE_VARS=("VLLM_USAGE_SOURCE")
+<% } else if (modelServer === 'sglang') { %>
+EXCLUDE_VARS=()
+<% } else if (modelServer === 'tensorrt-llm') { %>
+# Exclude TRTLLM_MODEL as it's used as the positional MODEL argument
+EXCLUDE_VARS=("TRTLLM_MODEL")
+<% } %>
+
+# Declare and populate array of matching environment variables
+mapfile -t env_vars < <(env | grep "^${PREFIX}")
+
+# Loop through the array and convert to command-line arguments
+for var in "${env_vars[@]}"; do
+    IFS='=' read -r key value <<< "$var"
+    
+    # Skip excluded variables
+    skip=false
+    for exclude in "${EXCLUDE_VARS[@]}"; do
+        if [ "$key" = "$exclude" ]; then
+            skip=true
+            break
+        fi
+    done
+    
+    if [ "$skip" = true ]; then
+        continue
+    fi
+    
+    # Remove prefix, convert to lowercase, and replace underscores with dashes
+    arg_name=$(echo "${key#"${PREFIX}"}" | tr '[:upper:]' '[:lower:]' | tr '_' '-')
+    SERVER_ARGS+=("${ARG_PREFIX}${arg_name}")
+    if [ -n "$value" ]; then
+        SERVER_ARGS+=("$value")
+    fi
+done
+
+echo "-------------------------------------------------------------------"
+<% if (modelServer === 'vllm') { %>
+echo "vLLM engine args: [${SERVER_ARGS[@]}]"
+<% } else if (modelServer === 'sglang') { %>
+echo "SGLang engine args: [${SERVER_ARGS[@]}]"
+<% } else if (modelServer === 'tensorrt-llm') { %>
+echo "TensorRT-LLM engine args: [${SERVER_ARGS[@]}]"
+<% } %>
+echo "-------------------------------------------------------------------"
+
+# Pass the collected arguments to the main entrypoint
+<% if (modelServer === 'vllm') { %>
+exec python3 -m vllm.entrypoints.openai.api_server "${SERVER_ARGS[@]}"
+<% } else if (modelServer === 'sglang') { %>
+exec python3 -m sglang.launch_server "${SERVER_ARGS[@]}"
+<% } else if (modelServer === 'tensorrt-llm') { %>
+# TensorRT-LLM requires the model as a positional argument
+# Syntax: trtllm-serve serve MODEL [OPTIONS]
+if [ -z "$TRTLLM_MODEL" ]; then
+    echo "Error: TRTLLM_MODEL environment variable is not set"
+    exit 1
+fi
+exec trtllm-serve serve "$TRTLLM_MODEL" "${SERVER_ARGS[@]}"
+<% } %>
+<% } %>

package/templates/code/serve.py

@@ -0,0 +1,175 @@
+#!/usr/bin/env python3
+# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+# SPDX-License-Identifier: Apache-2.0
+"""
+SageMaker inference server script
+"""
+import os
+import logging
+<% if (modelServer === 'flask') { %>
+from flask import Flask, request, jsonify
+<% } else if (modelServer === 'fastapi') { %>
+from fastapi import FastAPI, HTTPException, Request
+from fastapi.responses import JSONResponse
+<% } else if (modelServer === 'sglang') { %>
+from fastapi import FastAPI, HTTPException, Request
+from fastapi.responses import JSONResponse
+from fastapi.responses import JSONResponse
+import asyncio
+from sglang import Runtime
+<% } %>
+from model_handler import ModelHandler
+
+# Configure logging
+logging.basicConfig(level=logging.INFO,
+                   format='%(asctime)s - %(name)s - %(levelname)s - %(message)s')
+logger = logging.getLogger(__name__)
+model_handler = None
+
+<% if (modelServer === 'flask') { %>
+app = Flask(__name__)
+
+def create_app():
+    app = Flask(__name__)
+
+    @app.route('/ping', methods=['GET'])
+    def ping():
+        """Health check endpoint"""
+        if model_handler and model_handler.is_loaded():
+            return jsonify({'status': 'healthy'})
+        return jsonify({'status': 'model not loaded'}), 503
+
+    @app.route('/invocations', methods=['POST'])
+    def invocations():
+        """Main inference endpoint"""
+        if not model_handler or not model_handler.is_loaded():
+            return jsonify({'error': 'Model not loaded'}), 503
+
+        try:
+            data = request.get_json() if request.is_json else request.data
+            result = model_handler.predict(data)
+            return jsonify(result)
+        except ValueError as e:
+            return jsonify({'error': f'Invalid input: {str(e)}'}), 400
+        except Exception as e:
+            logger.exception("Error during inference")
+            return jsonify({'error': str(e)}), 500
+
+    return app
+
+def load_model_for_worker():
+    """Load the model when the server starts"""
+    global model_handler
+    model_path = "/opt/ml/model"
+    logger.info(f"Loading model from {model_path}")
+    model_handler = ModelHandler(model_path)
+    model_handler.load_model()
+    logger.info("Model loaded successfully")
+<% } else if (modelServer === 'fastapi') { %>
+app = FastAPI()
+
+@app.on_event("startup")
+async def startup_event():
+    """Load the model when the server starts"""
+    global model_handler
+    model_path = "/opt/ml/model"
+    logger.info(f"Loading model from {model_path}")
+    model_handler = ModelHandler(model_path)
+    model_handler.load_model()
+    logger.info("Model loaded successfully")
+
+@app.get('/ping')
+async def ping():
+    """Health check endpoint"""
+    if model_handler and model_handler.is_loaded():
+        return {'status': 'healthy'}
+    raise HTTPException(status_code=503, detail={'status': 'model not loaded'})
+
+@app.post('/invocations')
+async def invocations(request: Request):
+    """Main inference endpoint"""
+    if not model_handler or not model_handler.is_loaded():
+        raise HTTPException(status_code=503, detail={'error': 'Model not loaded'})
+
+    try:
+        content_type = request.headers.get('content-type', '')
+        if 'application/json' in content_type:
+            data = await request.json()
+        else:
+            data = await request.body()
+        
+        result = model_handler.predict(data)
+        return result
+    except ValueError as e:
+        raise HTTPException(status_code=400, detail={'error': f'Invalid input: {str(e)}'})
+    except Exception as e:
+        logger.exception("Error during inference")
+        raise HTTPException(status_code=500, detail={'error': str(e)})
+
+<% } else if (modelServer === 'sglang') { %>
+app = FastAPI()
+sglang_runtime = None
+
+@app.on_event("startup")
+async def startup_event():
+    """Initialize SGLang runtime when the server starts"""
+    global sglang_runtime
+    model_id = "<%= modelName %>"
+    logger.info(f"Initializing SGLang runtime with model: {model_id}")
+
+    sglang_runtime = Runtime(
+        model_path=model_id,
+        tokenizer_path=model_id,
+        device="cuda",
+        mem_fraction_static=0.8
+    )
+    logger.info("SGLang runtime initialized successfully")
+
+@app.get('/ping')
+async def ping():
+    """Health check endpoint"""
+    if sglang_runtime:
+        return {'status': 'healthy'}
+    raise HTTPException(status_code=503, detail={'status': 'runtime not loaded'})
+
+@app.post('/invocations')
+async def invocations(request: Request):
+    """Main inference endpoint"""
+    if not sglang_runtime:
+        raise HTTPException(status_code=503, detail={'error': 'Runtime not loaded'})
+
+    try:
+        content_type = request.headers.get('content-type', '')
+        if 'application/json' in content_type:
+            data = await request.json()
+        else:
+            data = await request.body()
+
+        # Extract prompts from SageMaker format
+        if isinstance(data, dict):
+            prompts = data.get('instances', data.get('inputs', [data]))
+        else:
+            prompts = [data]
+
+        # Generate responses
+        outputs = sglang_runtime.generate(prompts)
+        return {'predictions': outputs}
+
+    except ValueError as e:
+        raise HTTPException(status_code=400, detail={'error': f'Invalid input: {str(e)}'})
+    except Exception as e:
+        logger.exception("Error during inference")
+        raise HTTPException(status_code=500, detail={'error': str(e)})
+<% } %>
+
+if __name__ == '__main__':
+<% if (modelServer === 'flask') { %>
+    app = create_app()
+    load_model_for_worker()  # Load model for development server
+    port = int(os.environ.get("SAGEMAKER_BIND_TO_PORT", 8080))
+    app.run(host='0.0.0.0', port=port, debug=False)
+<% } else if (modelServer === 'fastapi' || modelServer === 'sglang') { %>
+    import uvicorn
+    port = int(os.environ.get("SAGEMAKER_BIND_TO_PORT", 8080))
+    uvicorn.run(app, host='0.0.0.0', port=port)
+<% } %>

package/templates/code/serving.properties

@@ -0,0 +1,105 @@
+# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+# SPDX-License-Identifier: Apache-2.0
+
+# LMI/DJL Serving Configuration
+# Documentation: https://docs.djl.ai/master/docs/serving/serving/docs/lmi/deployment_guide/index.html
+
+<% if (modelServer === 'lmi') { %>
+# LMI Container Configuration
+# LMI provides optimized inference with multiple backend options
+
+# Model Configuration
+<% if (modelSource === 'huggingface' || !modelSource) { %>
+option.model_id=<%= modelName %>
+<% } else if (artifactUri) { %>
+option.model_id=<%= artifactUri %>
+<% } else { %>
+# Model will be loaded from /opt/ml/model at runtime
+# (JumpStart model without artifact URI — requires SageMaker ModelDataUrl)
+# option.model_id=/opt/ml/model
+<% } %>
+
+# Backend Selection
+# Options: vllm, lmi-dist (DeepSpeed), tensorrt-llm, transformers-neuronx
+# Leave unset to let LMI auto-select the best backend for your model
+# option.rolling_batch=vllm
+
+# Tensor Parallel Degree
+# Set to number of GPUs for multi-GPU inference
+# option.tensor_parallel_degree=1
+
+# Data Type
+# Options: fp16, bf16, fp32, int8
+# option.dtype=fp16
+
+# Max Rolling Batch Size
+# Maximum number of concurrent requests
+# option.max_rolling_batch_size=32
+
+<% if (hfToken && (!modelSource || modelSource === 'huggingface')) { %>
+# HuggingFace Authentication
+option.hf_token=<%= hfToken %>
+<% } %>
+
+<% if (chatTemplate) { %>
+# Chat Template
+# Custom chat template for formatting messages
+option.chat_template=<%= chatTemplate %>
+<% } %>
+
+# Performance Tuning
+# Uncomment and adjust based on your needs:
+# option.max_model_len=4096
+# option.gpu_memory_utilization=0.9
+# option.enable_chunked_prefill=true
+
+<% } else if (modelServer === 'djl') { %>
+# DJL Serving Configuration
+# DJL provides flexible model serving with multiple framework support
+
+# Model Configuration
+<% if (modelSource === 'huggingface' || !modelSource) { %>
+option.model_id=<%= modelName %>
+<% } else if (artifactUri) { %>
+option.model_id=<%= artifactUri %>
+<% } else { %>
+# Model will be loaded from /opt/ml/model at runtime
+# (JumpStart model without artifact URI — requires SageMaker ModelDataUrl)
+# option.model_id=/opt/ml/model
+<% } %>
+
+# Engine Selection
+# Options: Python, PyTorch, TensorFlow, MXNet
+engine=Python
+
+# Batch Size
+# Maximum batch size for inference
+# batch_size=1
+
+# Job Queue Size
+# Maximum number of queued requests
+# job_queue_size=100
+
+<% if (hfToken && (!modelSource || modelSource === 'huggingface')) { %>
+# HuggingFace Authentication
+option.hf_token=<%= hfToken %>
+<% } %>
+
+<% if (chatTemplate) { %>
+# Chat Template
+option.chat_template=<%= chatTemplate %>
+<% } %>
+
+# GPU Configuration
+# option.tensor_parallel_degree=1
+# option.device_map=auto
+
+<% } %>
+
+# Additional Environment-Specific Configuration
+<% if (orderedEnvVars && orderedEnvVars.length > 0) { %>
+# Custom environment variables from configuration
+<% orderedEnvVars.forEach(({ key, value }) => { %>
+# <%= key %>=<%= value %>
+<% }); %>
+<% } %>

package/templates/code/start_server.py

@@ -0,0 +1,39 @@
+#!/usr/bin/env python3
+# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+# SPDX-License-Identifier: Apache-2.0
+"""
+Script to start the SageMaker inference server
+"""
+import signal
+import subprocess
+import sys
+
+
+def signal_handler(signum, frame):
+    """Handle shutdown signals"""
+    print(f"Received signal {signum}, shutting down gracefully...")
+    sys.exit(0)
+
+
+if __name__ == '__main__':
+    # Register signal handlers
+    signal.signal(signal.SIGTERM, signal_handler)
+    signal.signal(signal.SIGINT, signal_handler)
+
+<% if (modelServer === 'flask') { %>
+    print("Starting SageMaker inference server with Gunicorn")
+    subprocess.run([
+        'gunicorn',
+        '--config', '/opt/ml/code/gunicorn_config.py',
+        'wsgi:application'
+    ])
+<% } else if (modelServer === 'fastapi' || modelServer === 'sglang') { %>
+    print("Starting SageMaker inference server with Uvicorn")
+    subprocess.run([
+        'uvicorn',
+        'serve:app',
+        '--host', '0.0.0.0',
+        '--port', '8080',
+        '--workers', '4'
+    ])
+<% } %>

package/templates/code/start_server.sh

@@ -0,0 +1,39 @@
+#!/bin/bash
+# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+# SPDX-License-Identifier: Apache-2.0
+
+set -e
+
+echo "Starting TensorRT-LLM server on port 8081..."
+/usr/bin/serve_trtllm &
+TRTLLM_PID=$!
+
+# Wait for TensorRT-LLM to be ready
+echo "Waiting for TensorRT-LLM server to start..."
+for i in {1..300}; do
+    if curl -s http://localhost:8081/health > /dev/null 2>&1; then
+        echo "TensorRT-LLM server is ready!"
+        break
+    fi
+    if [ $i -eq 300 ]; then
+        echo "ERROR: TensorRT-LLM server failed to start within 300 seconds"
+        exit 1
+    fi
+    sleep 1
+done
+
+echo "Starting nginx reverse proxy on port 8080..."
+nginx -c /etc/nginx/nginx.conf &
+NGINX_PID=$!
+
+# Wait for either process to exit (this keeps the container running)
+wait -n $TRTLLM_PID $NGINX_PID
+
+# If we get here, one process exited - this is an error condition
+EXIT_CODE=$?
+echo "ERROR: A critical process exited unexpectedly (exit code: $EXIT_CODE)"
+
+# Kill any remaining processes
+kill $TRTLLM_PID $NGINX_PID 2>/dev/null || true
+
+exit $EXIT_CODE