@aws/ml-container-creator 0.8.0 → 0.9.1

package/src/lib/train-request-builder.js

@@ -0,0 +1,120 @@
+// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+// SPDX-License-Identifier: Apache-2.0
+
+/**
+ * Train Request Builder
+ *
+ * JavaScript module that replicates the Python helper's (.train_build_request.py)
+ * logic for constructing a CreateTrainingJob JSON request from a parsed config.
+ *
+ * This module mirrors the behavior of the Python build_request() function,
+ * providing a testable implementation of the config-to-API mapping logic.
+ */
+
+/**
+ * Build a CreateTrainingJob request from a parsed training config.
+ *
+ * Maps config fields to the SageMaker CreateTrainingJob API structure:
+ * - image → AlgorithmSpecification.TrainingImage
+ * - instance_type → ResourceConfig.InstanceType
+ * - instance_count → ResourceConfig.InstanceCount
+ * - dataset → InputDataConfig[0].DataSource.S3DataSource.S3Uri
+ * - output_path → OutputDataConfig.S3OutputPath
+ * - hyperparameters → HyperParameters (string key-value pairs)
+ * - max_runtime_seconds → StoppingCondition.MaxRuntimeInSeconds
+ * - enable_spot=true → EnableManagedSpotTraining = true
+ * - enable_spot=true → StoppingCondition.MaxWaitTimeInSeconds
+ * - checkpoint_path → CheckpointConfig.S3Uri
+ * - metric_definitions → AlgorithmSpecification.MetricDefinitions
+ * - environment → Environment
+ * - tags → Tags (converted from {k:v} to [{Key:k, Value:v}])
+ *
+ * @param {object} options - Build options
+ * @param {string} options.jobName - Training job name
+ * @param {string} options.roleArn - SageMaker execution role ARN
+ * @param {object} options.config - Parsed training config (from parseTrainingConfig)
+ * @returns {object} CreateTrainingJob request body
+ */
+export function buildTrainingJobRequest({ jobName, roleArn, config }) {
+    const request = {
+        TrainingJobName: jobName,
+        RoleArn: roleArn,
+        AlgorithmSpecification: {
+            TrainingImage: config.image,
+            TrainingInputMode: 'File'
+        },
+        InputDataConfig: [
+            {
+                ChannelName: 'training',
+                DataSource: {
+                    S3DataSource: {
+                        S3DataType: 'S3Prefix',
+                        S3Uri: config.dataset,
+                        S3DataDistributionType: 'FullyReplicated'
+                    }
+                }
+            }
+        ],
+        OutputDataConfig: {
+            S3OutputPath: config.output_path
+        },
+        ResourceConfig: {
+            InstanceType: config.instance_type,
+            InstanceCount: parseInt(config.instance_count, 10),
+            VolumeSizeInGB: parseInt(config.volume_size_gb, 10)
+        },
+        StoppingCondition: {
+            MaxRuntimeInSeconds: parseInt(config.max_runtime_seconds, 10)
+        }
+    };
+
+    // Hyperparameters — ensure all values are strings (SageMaker requirement)
+    const hyperparams = config.hyperparameters || {};
+    if (Object.keys(hyperparams).length > 0) {
+        request.HyperParameters = {};
+        for (const [k, v] of Object.entries(hyperparams)) {
+            request.HyperParameters[String(k)] = String(v);
+        }
+    }
+
+    // Managed spot training
+    const enableSpot = config.enable_spot === 'true' || config.enable_spot === true;
+    if (enableSpot) {
+        request.EnableManagedSpotTraining = true;
+        request.StoppingCondition.MaxWaitTimeInSeconds = parseInt(config.max_wait_seconds, 10);
+    }
+
+    // Checkpoint configuration (for spot training resumption)
+    const checkpointPath = config.checkpoint_path || '';
+    if (checkpointPath) {
+        request.CheckpointConfig = {
+            S3Uri: checkpointPath
+        };
+    }
+
+    // Metric definitions (custom CloudWatch metrics)
+    const metricDefs = config.metric_definitions || [];
+    if (Array.isArray(metricDefs) && metricDefs.length > 0) {
+        request.AlgorithmSpecification.MetricDefinitions = metricDefs.map(m => ({
+            Name: m.name,
+            Regex: m.regex
+        }));
+    }
+
+    // Environment variables for the container
+    const environment = config.environment || {};
+    if (Object.keys(environment).length > 0) {
+        request.Environment = environment;
+    }
+
+    // Tags — convert from {key: value} map to [{Key: k, Value: v}] array
+    const tags = config.tags || {};
+    if (Object.keys(tags).length > 0) {
+        request.Tags = Object.entries(tags).map(([k, v]) => ({
+            Key: String(k),
+            Value: String(v)
+        }));
+    }
+
+    return request;
+}

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

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

@@ -0,0 +1,47 @@
+# ---------------------------------------------------------------------------
+# SGLang Server Configuration
+# ---------------------------------------------------------------------------
+# Env prefix: SGLANG_
+# Entrypoint: python3 -m sglang.launch_server
+# Port: 8080 (SageMaker requirement)
+# ---------------------------------------------------------------------------
+
+PREFIX="SGLANG_"
+ARG_PREFIX="--"
+
+EXCLUDE_VARS=()
+
+# Declare and populate array of matching environment variables
+mapfile -t env_vars < <(env | grep "^${PREFIX}")
+
+# Convert SGLANG_ env vars to CLI 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, 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 "-------------------------------------------------------------------"
+echo "SGLang engine args: [${SERVER_ARGS[@]}]"
+echo "-------------------------------------------------------------------"
+
+exec python3 -m sglang.launch_server "${SERVER_ARGS[@]}"

package/templates/code/serve.d/tensorrt-llm.ejs

@@ -0,0 +1,53 @@
+# ---------------------------------------------------------------------------
+# TensorRT-LLM Server Configuration
+# ---------------------------------------------------------------------------
+# Env prefix: TRTLLM_
+# Entrypoint: trtllm-serve serve MODEL [OPTIONS]
+# Port: 8081 (nginx proxies to 8080 for SageMaker)
+# ---------------------------------------------------------------------------
+
+PREFIX="TRTLLM_"
+ARG_PREFIX="--"
+
+# TRTLLM_MODEL is used as the positional argument, not a --flag
+EXCLUDE_VARS=("TRTLLM_MODEL")
+
+# Declare and populate array of matching environment variables
+mapfile -t env_vars < <(env | grep "^${PREFIX}")
+
+# Convert TRTLLM_ env vars to CLI 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, 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 "-------------------------------------------------------------------"
+echo "TensorRT-LLM engine args: [${SERVER_ARGS[@]}]"
+echo "-------------------------------------------------------------------"
+
+# TensorRT-LLM requires the model as a positional argument
+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.d/vllm.ejs

@@ -0,0 +1,48 @@
+# ---------------------------------------------------------------------------
+# vLLM Server Configuration
+# ---------------------------------------------------------------------------
+# Env prefix: VLLM_
+# Entrypoint: python3 -m vllm.entrypoints.openai.api_server
+# Port: 8080 (SageMaker requirement)
+# ---------------------------------------------------------------------------
+
+PREFIX="VLLM_"
+ARG_PREFIX="--"
+
+# Internal variables set by the base image — not CLI args
+EXCLUDE_VARS=("VLLM_USAGE_SOURCE" "VLLM_ENABLE_CUDA_COMPATIBILITY")
+
+# Declare and populate array of matching environment variables
+mapfile -t env_vars < <(env | grep "^${PREFIX}")
+
+# Convert VLLM_ env vars to CLI 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, 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 "-------------------------------------------------------------------"
+echo "vLLM engine args: [${SERVER_ARGS[@]}]"
+echo "-------------------------------------------------------------------"
+
+exec python3 -m vllm.entrypoints.openai.api_server "${SERVER_ARGS[@]}"

package/templates/do/.train_build_request.py

@@ -0,0 +1,141 @@
+#!/usr/bin/env python3
+# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+# SPDX-License-Identifier: Apache-2.0
+
+"""
+Build the CreateTrainingJob JSON request for SageMaker.
+
+This helper is called by do/train to construct the full API request body.
+It handles conditional fields (spot training, metric definitions, environment,
+tags) and writes the result to a JSON file for use with:
+    aws sagemaker create-training-job --cli-input-json file://path.json
+"""
+
+import argparse
+import json
+import sys
+
+
+def parse_args():
+    """Parse command-line arguments."""
+    parser = argparse.ArgumentParser(description='Build CreateTrainingJob request JSON')
+    parser.add_argument('--job-name', required=True, help='Training job name')
+    parser.add_argument('--role-arn', required=True, help='SageMaker execution role ARN')
+    parser.add_argument('--image', required=True, help='Training container image URI')
+    parser.add_argument('--instance-type', required=True, help='Instance type')
+    parser.add_argument('--instance-count', required=True, help='Instance count')
+    parser.add_argument('--volume-size', required=True, help='Volume size in GB')
+    parser.add_argument('--dataset', required=True, help='S3 URI for training dataset')
+    parser.add_argument('--output-path', required=True, help='S3 URI for output')
+    parser.add_argument('--max-runtime', required=True, help='Max runtime in seconds')
+    parser.add_argument('--hyperparams', required=True, help='Hyperparameters as JSON string')
+    parser.add_argument('--enable-spot', required=True, help='Enable spot training (true/false)')
+    parser.add_argument('--max-wait', required=True, help='Max wait time for spot in seconds')
+    parser.add_argument('--checkpoint-path', required=True, help='S3 checkpoint path')
+    parser.add_argument('--metric-definitions', required=True, help='Metric definitions as JSON array')
+    parser.add_argument('--environment', required=True, help='Environment variables as JSON object')
+    parser.add_argument('--tags', required=True, help='Tags as JSON object (key-value map)')
+    parser.add_argument('--output-file', required=True, help='Output file path for the JSON')
+    return parser.parse_args()
+
+
+def build_request(args):
+    """Construct the CreateTrainingJob request dictionary."""
+    # Parse JSON inputs
+    hyperparams = json.loads(args.hyperparams) if args.hyperparams else {}
+    metric_definitions = json.loads(args.metric_definitions) if args.metric_definitions else []
+    environment = json.loads(args.environment) if args.environment else {}
+    tags = json.loads(args.tags) if args.tags else {}
+
+    # Base request structure
+    request = {
+        'TrainingJobName': args.job_name,
+        'RoleArn': args.role_arn,
+        'AlgorithmSpecification': {
+            'TrainingImage': args.image,
+            'TrainingInputMode': 'File'
+        },
+        'InputDataConfig': [
+            {
+                'ChannelName': 'training',
+                'DataSource': {
+                    'S3DataSource': {
+                        'S3DataType': 'S3Prefix',
+                        'S3Uri': args.dataset,
+                        'S3DataDistributionType': 'FullyReplicated'
+                    }
+                }
+            }
+        ],
+        'OutputDataConfig': {
+            'S3OutputPath': args.output_path
+        },
+        'ResourceConfig': {
+            'InstanceType': args.instance_type,
+            'InstanceCount': int(args.instance_count),
+            'VolumeSizeInGB': int(args.volume_size)
+        },
+        'StoppingCondition': {
+            'MaxRuntimeInSeconds': int(args.max_runtime)
+        }
+    }
+
+    # Hyperparameters — ensure all values are strings (SageMaker requirement)
+    if hyperparams:
+        request['HyperParameters'] = {
+            str(k): str(v) for k, v in hyperparams.items()
+        }
+
+    # Managed spot training
+    if args.enable_spot == 'true':
+        request['EnableManagedSpotTraining'] = True
+        request['StoppingCondition']['MaxWaitTimeInSeconds'] = int(args.max_wait)
+
+    # Checkpoint configuration (for spot training resumption)
+    if args.checkpoint_path:
+        request['CheckpointConfig'] = {
+            'S3Uri': args.checkpoint_path
+        }
+
+    # Metric definitions (custom CloudWatch metrics)
+    if metric_definitions and metric_definitions != []:
+        request['AlgorithmSpecification']['MetricDefinitions'] = [
+            {'Name': m['name'], 'Regex': m['regex']}
+            for m in metric_definitions
+        ]
+
+    # Environment variables for the container
+    if environment and environment != {}:
+        request['Environment'] = environment
+
+    # Tags — convert from {key: value} map to [{Key: k, Value: v}] array
+    if tags and tags != {}:
+        request['Tags'] = [
+            {'Key': str(k), 'Value': str(v)}
+            for k, v in tags.items()
+        ]
+
+    return request
+
+
+def main():
+    """Main entry point."""
+    args = parse_args()
+
+    try:
+        request = build_request(args)
+    except (json.JSONDecodeError, ValueError) as e:
+        print(f'❌ Failed to build request: {e}', file=sys.stderr)
+        sys.exit(1)
+
+    # Write the JSON request to the output file
+    try:
+        with open(args.output_file, 'w') as f:
+            json.dump(request, f, indent=2)
+    except IOError as e:
+        print(f'❌ Failed to write request file: {e}', file=sys.stderr)
+        sys.exit(1)
+
+
+if __name__ == '__main__':
+    main()