@aws/ml-container-creator 0.9.0 → 0.10.0

package/src/lib/template-variable-resolver.js

@@ -0,0 +1,398 @@
+// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+// SPDX-License-Identifier: Apache-2.0
+
+import fs from 'fs';
+import path from 'path';
+import { fileURLToPath } from 'url';
+import { isTuneSupported } from './tune-catalog-validator.js';
+
+const __filename = fileURLToPath(import.meta.url);
+const __dirname = path.dirname(__filename);
+
+/**
+ * Finds model configuration by exact match or glob-pattern match.
+ *
+ * @param {string} modelName - Model ID to look up
+ * @param {object} registryConfigManager - Registry configuration manager
+ * @returns {object|null} Model configuration or null
+ */
+function _findModelConfig(modelName, registryConfigManager) {
+    if (!registryConfigManager?.modelRegistry) return null;
+
+    // Exact match first
+    const exact = registryConfigManager.modelRegistry[modelName];
+    if (exact) return exact;
+
+    // Pattern matching with glob-style wildcards
+    for (const [pattern, config] of Object.entries(registryConfigManager.modelRegistry)) {
+        if (pattern.includes('*')) {
+            const regex = new RegExp(`^${pattern.replace(/\*/g, '.*')}$`);
+            if (regex.test(modelName)) {
+                return config;
+            }
+        }
+    }
+
+    return null;
+}
+
+/**
+ * Merges environment variables from all catalog sources with correct precedence.
+ * Precedence (lowest → highest):
+ *   1. catalog defaults (Image_Entry defaults.envVars)
+ *   2. framework profile (Image_Entry profiles[selectedProfile].envVars)
+ *   3. model entry (model catalog entry envVars)
+ *   4. model profile (model catalog entry profiles[selectedProfile].envVars)
+ *   5. CLI overrides (existing answers.envVars from user CLI input)
+ *
+ * @param {object} answers - Configuration answers
+ * @param {object|null} registryConfigManager - Registry configuration manager
+ */
+export async function _mergeEnvVarsWithPrecedence(answers, registryConfigManager) {
+    if (!registryConfigManager) return;
+
+    // Capture CLI-provided env vars before merging (highest precedence)
+    const cliEnvVars = { ...answers.envVars };
+
+    // Resolve the framework config for the selected framework + version
+    const frameworkName = answers.framework || answers.deploymentConfig;
+    const frameworkVersion = answers.frameworkVersion;
+    let frameworkConfig = null;
+
+    if (frameworkName && registryConfigManager.frameworkRegistry) {
+        const frameworkVersions = registryConfigManager.frameworkRegistry[frameworkName];
+        if (frameworkVersions) {
+            if (frameworkVersion && frameworkVersions[frameworkVersion]) {
+                frameworkConfig = frameworkVersions[frameworkVersion];
+            } else {
+                // Fall back to latest version for Triton and other non-versioned lookups
+                const versions = Object.keys(frameworkVersions).sort((a, b) =>
+                    b.localeCompare(a, undefined, { numeric: true })
+                );
+                if (versions.length > 0) {
+                    frameworkConfig = frameworkVersions[versions[0]];
+                }
+            }
+        }
+    }
+
+    // Resolve the model config (exact match or pattern match)
+    let modelConfig = null;
+    if (answers.modelName && registryConfigManager.modelRegistry) {
+        modelConfig = _findModelConfig(answers.modelName, registryConfigManager);
+    }
+
+    // Layer 1: catalog defaults (Image_Entry defaults.envVars)
+    const catalogDefaults = frameworkConfig?.envVars || {};
+
+    // Layer 2: framework profile envVars
+    let frameworkProfileEnvVars = {};
+    if (answers.frameworkProfile && frameworkConfig?.profiles) {
+        const profile = frameworkConfig.profiles[answers.frameworkProfile];
+        if (profile?.envVars) {
+            frameworkProfileEnvVars = profile.envVars;
+        }
+    }
+
+    // Layer 3: model entry envVars
+    const modelEntryEnvVars = modelConfig?.envVars || {};
+
+    // Layer 4: model profile envVars
+    let modelProfileEnvVars = {};
+    if (answers.modelProfile && modelConfig?.profiles) {
+        const profile = modelConfig.profiles[answers.modelProfile];
+        if (profile?.envVars) {
+            modelProfileEnvVars = profile.envVars;
+        }
+    }
+
+    // Layer 5: CLI overrides (captured above)
+
+    // Merge in precedence order: each layer overrides the previous
+    answers.envVars = {
+        ...catalogDefaults,
+        ...frameworkProfileEnvVars,
+        ...modelEntryEnvVars,
+        ...modelProfileEnvVars,
+        ...cliEnvVars
+    };
+}
+
+/**
+ * Validates environment variables using the registry system.
+ * Displays errors and warnings to the user.
+ *
+ * @param {object} answers - Configuration answers
+ * @param {object} registryConfigManager - Registry configuration manager
+ */
+export async function _validateEnvironmentVariables(answers, registryConfigManager) {
+    // Get framework configuration
+    // For Triton configs, look up using deploymentConfig key (e.g. 'triton-fil')
+    let frameworkConfig;
+    if (answers.architecture === 'triton' && answers.deploymentConfig) {
+        const tritonEntry = registryConfigManager.frameworkRegistry?.[answers.deploymentConfig];
+        if (tritonEntry) {
+            const versions = Object.keys(tritonEntry);
+            if (versions.length > 0) {
+                frameworkConfig = tritonEntry[versions[0]];
+            }
+        }
+    }
+    if (!frameworkConfig) {
+        frameworkConfig = registryConfigManager.frameworkRegistry?.[answers.framework]?.[answers.frameworkVersion];
+    }
+
+    if (!frameworkConfig || !frameworkConfig.envVars) {
+        return; // No env vars to validate
+    }
+
+    console.log('\n🔍 Validating environment variables...');
+
+    // Validate environment variables
+    const validationResult = registryConfigManager.validateEnvironmentVariables(
+        frameworkConfig.envVars,
+        frameworkConfig
+    );
+
+    // Display validation results
+    if (validationResult.errors && validationResult.errors.length > 0) {
+        console.log('\n❌ Environment Variable Validation Errors:');
+        validationResult.errors.forEach(error => {
+            console.log(`   • ${error.key}: ${error.message}`);
+        });
+    }
+
+    if (validationResult.warnings && validationResult.warnings.length > 0) {
+        console.log('\n⚠️  Environment Variable Validation Warnings:');
+        validationResult.warnings.forEach(warning => {
+            console.log(`   • ${warning.key ? `${warning.key}: ` : ''}${warning.message}`);
+        });
+    }
+
+    if (validationResult.strategiesUsed && validationResult.strategiesUsed.length > 0) {
+        console.log(`\n✅ Validation methods used: ${validationResult.strategiesUsed.join(', ')}`);
+    }
+
+    if (!validationResult.errors || validationResult.errors.length === 0) {
+        if (!validationResult.warnings || validationResult.warnings.length === 0) {
+            console.log('   ✅ All environment variables validated successfully');
+        }
+    }
+
+    // In non-interactive mode (skip-prompts), throw on errors
+    if (validationResult.errors && validationResult.errors.length > 0) {
+        throw new Error('Environment variable validation failed. Please fix the errors and try again.');
+    }
+}
+
+/**
+ * Ensures all template variables have proper defaults to prevent
+ * "undefined" errors in EJS templates. Also enriches answers with
+ * registry data (env var merging, HuggingFace data, Triton base image).
+ *
+ * @param {object} answers - Answers object to fill defaults into
+ * @param {object|null} registryConfigManager - Registry configuration manager (or null)
+ */
+export async function _ensureTemplateVariables(answers, registryConfigManager = null) {
+    const defaults = {
+        chatTemplate: null,
+        chatTemplateSource: null,
+        hfToken: null,
+        hfTokenArn: null,
+        ngcApiKey: null,
+        ngcTokenArn: null,
+        envVars: {},
+        inferenceAmiVersion: null,
+        accelerator: null,
+        frameworkVersion: null,
+        validationLevel: 'unknown',
+        configSources: [],
+        recommendedInstanceTypes: [],
+        roleArn: null,
+        deploymentConfig: '',
+        architecture: null,
+        backend: null,
+        engine: null,
+        codebuildComputeType: null,
+        codebuildProjectName: null,
+        modelName: null,
+        modelFormat: null,
+        includeSampleModel: true,
+        includeTesting: true,
+        testTypes: [],
+        buildTimestamp: new Date().toISOString(),
+        buildTarget: 'codebuild',
+        deploymentTarget: 'realtime-inference',
+        hyperPodCluster: null,
+        hyperPodNamespace: 'default',
+        hyperPodReplicas: 1,
+        fsxVolumeHandle: null,
+        baseImage: null,
+        modelSource: 'huggingface',
+        artifactUri: '',
+        modelLoadStrategy: 'runtime',
+        existingEndpointName: null,
+        enableLora: false,
+        maxLoras: 30,
+        maxLoraRank: 64
+    };
+
+    Object.entries(defaults).forEach(([key, value]) => {
+        if (answers[key] === undefined) {
+            answers[key] = value;
+        }
+    });
+
+    // Backward compatibility: populate framework and modelServer from architecture/backend
+    if (!answers.framework && answers.architecture) {
+        answers.framework = answers.architecture;
+    }
+    if (!answers.modelServer && answers.backend) {
+        answers.modelServer = answers.backend;
+    }
+
+    // Always include testing with all available test types
+    answers.includeTesting = true;
+    if (!answers.testTypes || answers.testTypes.length === 0) {
+        if (answers.architecture === 'transformers' || answers.framework === 'transformers') {
+            answers.testTypes = ['hosted-model-endpoint'];
+        } else {
+            answers.testTypes = ['local-model-cli', 'local-model-server', 'hosted-model-endpoint'];
+        }
+    }
+
+    // Merge catalog env vars into answers.envVars with correct precedence
+    await _mergeEnvVarsWithPrecedence(answers, registryConfigManager);
+
+    // For Triton architecture, set default base image fallback
+    if (answers.architecture === 'triton' && !answers.baseImage) {
+        // Try to look up base image from framework registry using deployment-config key
+        const tritonRegistryKey = answers.deploymentConfig;
+        if (tritonRegistryKey && registryConfigManager?.frameworkRegistry) {
+            const tritonFrameworkConfig = registryConfigManager.frameworkRegistry[tritonRegistryKey];
+            if (tritonFrameworkConfig) {
+                const versions = Object.keys(tritonFrameworkConfig).sort((a, b) =>
+                    b.localeCompare(a, undefined, { numeric: true })
+                );
+                if (versions.length > 0) {
+                    const latestConfig = tritonFrameworkConfig[versions[0]];
+                    if (latestConfig.baseImage) {
+                        answers.baseImage = latestConfig.baseImage;
+                    }
+                    if (latestConfig.inferenceAmiVersion && !answers.inferenceAmiVersion) {
+                        answers.inferenceAmiVersion = latestConfig.inferenceAmiVersion;
+                    }
+                    if (latestConfig.accelerator) {
+                        answers.accelerator = latestConfig.accelerator;
+                    }
+                }
+            }
+        }
+        // Final fallback: hardcoded default Triton base image
+        if (!answers.baseImage) {
+            answers.baseImage = 'nvcr.io/nvidia/tritonserver:24.08-py3';
+        }
+    }
+
+    // For transformer models, enrich with HuggingFace data and non-envVar metadata
+    if (answers.framework === 'transformers' && answers.modelName && registryConfigManager) {
+        try {
+            // Fetch HuggingFace data for model-specific info
+            const hfData = await registryConfigManager._fetchHuggingFaceData(answers.modelName);
+
+            // Merge chatTemplate if available and not already set
+            if (hfData && hfData.chatTemplate && !answers.chatTemplate) {
+                answers.chatTemplate = hfData.chatTemplate;
+                answers.chatTemplateSource = 'HuggingFace_Hub_API';
+            }
+
+            // Check Model Registry for chatTemplate overrides
+            if (registryConfigManager.modelRegistry) {
+                const modelConfig = _findModelConfig(answers.modelName, registryConfigManager);
+
+                if (modelConfig && modelConfig.chatTemplate) {
+                    answers.chatTemplate = modelConfig.chatTemplate;
+                    answers.chatTemplateSource = 'Model_Registry';
+                }
+            }
+
+            // Set framework-level metadata (non-envVar fields)
+            if (answers.frameworkVersion && registryConfigManager.frameworkRegistry) {
+                const frameworkConfig = registryConfigManager.frameworkRegistry[answers.framework]?.[answers.frameworkVersion];
+
+                if (frameworkConfig) {
+                    if (frameworkConfig.inferenceAmiVersion && !answers.inferenceAmiVersion) {
+                        answers.inferenceAmiVersion = frameworkConfig.inferenceAmiVersion;
+                    }
+                    if (frameworkConfig.accelerator) {
+                        answers.accelerator = frameworkConfig.accelerator;
+                    }
+                }
+            }
+        } catch (error) {
+            // Silently continue - defaults are already set
+        }
+    }
+
+    // Populate baseImage from the catalog when still falsy (covers --skip-prompts and
+    // cases where MCP/CLI/config did not provide a base image).
+    // Precedence: MCP > CLI > config > catalog default (this block).
+    if (!answers.baseImage && registryConfigManager?.frameworkRegistry) {
+        const backendKey = answers.backend || answers.modelServer;
+        if (backendKey) {
+            const frameworkVersions = registryConfigManager.frameworkRegistry[backendKey];
+            if (frameworkVersions) {
+                let resolvedConfig = null;
+                if (answers.frameworkVersion && frameworkVersions[answers.frameworkVersion]) {
+                    resolvedConfig = frameworkVersions[answers.frameworkVersion];
+                } else {
+                    // Fall back to latest version
+                    const versions = Object.keys(frameworkVersions).sort((a, b) =>
+                        b.localeCompare(a, undefined, { numeric: true })
+                    );
+                    if (versions.length > 0) {
+                        resolvedConfig = frameworkVersions[versions[0]];
+                    }
+                }
+                if (resolvedConfig?.baseImage) {
+                    answers.baseImage = resolvedConfig.baseImage;
+                }
+            }
+        }
+    }
+
+    // Populate icGpuCount from instance catalog when not explicitly set.
+    // The deploy template uses IC_GPU_COUNT unconditionally for NumberOfAcceleratorDevicesRequired,
+    // so it must always have a value for GPU deployments.
+    if ((answers.icGpuCount === null || answers.icGpuCount === undefined) && answers.instanceType) {
+        // Use gpuCount from instance-sizer recommendation if available
+        if (answers.gpuCount) {
+            answers.icGpuCount = answers.gpuCount;
+        } else {
+            // Look up from instances catalog
+            try {
+                const catalogPath = path.resolve(__dirname, '..', '..', 'servers', 'lib', 'catalogs', 'instances.json');
+                const catalogData = JSON.parse(fs.readFileSync(catalogPath, 'utf-8'));
+                const instanceInfo = catalogData?.catalog?.[answers.instanceType];
+                if (instanceInfo?.gpus && instanceInfo.gpus > 0) {
+                    answers.icGpuCount = instanceInfo.gpus;
+                }
+            } catch {
+                // Silently continue — template fallback handles missing value
+            }
+        }
+    }
+
+    // Determine tune support based on model presence in the tune catalog.
+    // Used by the do/config template to write TUNE_SUPPORTED=true|false.
+    if (answers.tuneSupported === undefined) {
+        try {
+            const tuneCatalogPath = path.resolve(__dirname, '..', '..', 'config', 'tune-catalog.json');
+            const tuneCatalog = JSON.parse(fs.readFileSync(tuneCatalogPath, 'utf-8'));
+            const modelId = answers.modelName || '';
+            answers.tuneSupported = isTuneSupported(modelId, tuneCatalog);
+        } catch {
+            answers.tuneSupported = false;
+        }
+    }
+}

package/templates/code/serve

@@ -10,35 +10,10 @@ echo "$(date -u '+%Y-%m-%dT%H:%M:%SZ') [serve] Container started — PID $$"
 # CUDA compatibility setup (required for newer SageMaker inference AMIs)
 source /usr/bin/cuda_compat.sh 2>/dev/null || true
 
-<% 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"
-<% } %>
+echo "Starting <%= modelServer %> 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
+<%- include('serve.d/lmi') %>
 <% } else { %>
 
 <% if (typeof modelSource !== 'undefined' && modelSource !== 'huggingface') { %>
@@ -60,7 +35,6 @@ download_model_from_s3() {
     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
@@ -72,13 +46,11 @@ download_model_from_s3() {
         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
@@ -109,19 +81,16 @@ _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|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://"
@@ -131,7 +100,6 @@ resolve_model() {
                     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
@@ -151,38 +119,22 @@ resolve_model() {
                         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('')
@@ -192,19 +144,15 @@ except:
                         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
@@ -212,7 +160,6 @@ except:
             echo "$LOCAL_MODEL_PATH"
             ;;
         *)
-            # Unrecognized source — treat as huggingface
             echo "${!_MODEL_VAR}"
             return
             ;;
@@ -226,89 +173,13 @@ 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" "VLLM_ENABLE_CUDA_COMPATIBILITY")
-<% } 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 '_' '-')
-
-    # Boolean handling: true = flag only, false = skip entirely
-    if [ "$value" = "false" ]; then
-        continue
-    fi
-
-    SERVER_ARGS+=("${ARG_PREFIX}${arg_name}")
-    if [ -n "$value" ] && [ "$value" != "true" ]; 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[@]}"
+# --- Server-specific arg conversion and exec ---
+<% if (['vllm', 'sglang', 'tensorrt-llm'].includes(modelServer)) { %>
+<%- include('serve.d/' + modelServer) %>
 <% } %>
 <% } %>

package/templates/code/serve.d/lmi.ejs

@@ -0,0 +1,19 @@
+# ---------------------------------------------------------------------------
+# LMI / DJL Server Configuration
+# ---------------------------------------------------------------------------
+# Config: /opt/ml/model/serving.properties
+# Entrypoint: DJL Serving (built into base image)
+# Port: 8080 (configured in serving.properties)
+# ---------------------------------------------------------------------------
+
+# LMI/DJL containers use serving.properties for 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 entrypoint
+exit 0