@aws/ml-container-creator 0.8.0 → 0.9.0

package/src/lib/train-config-persistence.js

@@ -0,0 +1,143 @@
+// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+// SPDX-License-Identifier: Apache-2.0
+
+/**
+ * Train Config Persistence
+ *
+ * JavaScript module that models the config persistence logic from the bash
+ * `_update_config_var` function in `templates/do/train`. This module provides
+ * a pure JavaScript implementation for property-based testing of the config
+ * persistence behavior after job submission and completion.
+ *
+ * The config file uses the format:
+ *   export VAR_NAME="value"
+ *
+ * Behavior:
+ * - If the variable already exists: update its value in-place
+ * - If the variable doesn't exist: append it to the end
+ * - Existing variables in the config are preserved
+ *
+ * Requirements: 3.4, 5.1
+ */
+
+import { readFileSync, writeFileSync } from 'node:fs';
+
+/**
+ * Update or add a config variable in a do/config-style file.
+ * Mimics the bash _update_config_var() function from templates/do/train:
+ *
+ *   if grep -q "^export ${var_name}=" "${config_file}"; then
+ *       sed -i.bak "s|^export ${var_name}=.*|export ${var_name}=\"${var_value}\"|" "${config_file}"
+ *       rm -f "${config_file}.bak"
+ *   else
+ *       echo "export ${var_name}=\"${var_value}\"" >> "${config_file}"
+ *   fi
+ *
+ * @param {string} configContent - Current content of the config file
+ * @param {string} varName - Variable name (e.g., TRAIN_JOB_NAME)
+ * @param {string} varValue - Variable value
+ * @returns {string} Updated config content
+ */
+export function updateConfigVar(configContent, varName, varValue) {
+    const pattern = new RegExp(`^export ${varName}=.*$`, 'm');
+
+    if (pattern.test(configContent)) {
+        // Variable exists — update in-place
+        return configContent.replace(pattern, `export ${varName}="${varValue}"`);
+    } else {
+        // Variable doesn't exist — append
+        let result = configContent;
+        if (result.length > 0 && !result.endsWith('\n')) {
+            result += '\n';
+        }
+        result += `export ${varName}="${varValue}"\n`;
+        return result;
+    }
+}
+
+/**
+ * Read a config variable value from a do/config-style file content.
+ *
+ * @param {string} configContent - Content of the config file
+ * @param {string} varName - Variable name to read
+ * @returns {string|null} The variable value, or null if not found
+ */
+export function readConfigVar(configContent, varName) {
+    const pattern = new RegExp(`^export ${varName}="([^"]*)"`, 'm');
+    const match = configContent.match(pattern);
+    return match ? match[1] : null;
+}
+
+/**
+ * Simulate the config writes that happen after a successful training job submission.
+ * This mirrors the behavior in do/train's _submit_job() function which calls:
+ *   _update_config_var "TRAIN_JOB_NAME" "${JOB_NAME}"
+ *
+ * @param {string} configContent - Current content of the config file
+ * @param {object} params - Submission parameters
+ * @param {string} params.jobName - Generated job name (pattern: ${PROJECT_NAME}-train-${TIMESTAMP})
+ * @returns {string} Updated config content
+ */
+export function persistTrainSubmission(configContent, { jobName }) {
+    return updateConfigVar(configContent, 'TRAIN_JOB_NAME', jobName);
+}
+
+/**
+ * Simulate the config writes that happen after a training job completes.
+ * This mirrors the behavior in do/train's _handle_completion() function which calls:
+ *   _update_config_var "TRAIN_OUTPUT_PATH" "${output_path}"
+ *
+ * @param {string} configContent - Current content of the config file
+ * @param {object} params - Completion parameters
+ * @param {string} params.outputPath - S3 path to the output artifacts
+ * @returns {string} Updated config content
+ */
+export function persistTrainCompletion(configContent, { outputPath }) {
+    return updateConfigVar(configContent, 'TRAIN_OUTPUT_PATH', outputPath);
+}
+
+/**
+ * Generate a training job name following the pattern used by do/train.
+ * Pattern: ${projectName}-train-YYYYMMDD-HHMMSS
+ *
+ * @param {string} projectName - Project name
+ * @param {Date} [timestamp] - Optional timestamp (defaults to now)
+ * @returns {string} Generated job name
+ */
+export function generateTrainJobName(projectName, timestamp = new Date()) {
+    const year = timestamp.getFullYear().toString();
+    const month = (timestamp.getMonth() + 1).toString().padStart(2, '0');
+    const day = timestamp.getDate().toString().padStart(2, '0');
+    const hours = timestamp.getHours().toString().padStart(2, '0');
+    const minutes = timestamp.getMinutes().toString().padStart(2, '0');
+    const seconds = timestamp.getSeconds().toString().padStart(2, '0');
+    const dateStr = `${year}${month}${day}`;
+    const timeStr = `${hours}${minutes}${seconds}`;
+    return `${projectName}-train-${dateStr}-${timeStr}`;
+}
+
+/**
+ * File-based version of updateConfigVar that reads/writes to disk.
+ * Used for integration-style tests that need actual file I/O.
+ *
+ * @param {string} configPath - Path to the config file
+ * @param {string} varName - Variable name
+ * @param {string} varValue - Variable value
+ */
+export function updateConfigVarFile(configPath, varName, varValue) {
+    const content = readFileSync(configPath, 'utf8');
+    const updated = updateConfigVar(content, varName, varValue);
+    writeFileSync(configPath, updated, 'utf8');
+}
+
+/**
+ * File-based version of readConfigVar that reads from disk.
+ *
+ * @param {string} configPath - Path to the config file
+ * @param {string} varName - Variable name to read
+ * @returns {string|null} The variable value, or null if not found
+ */
+export function readConfigVarFile(configPath, varName) {
+    const content = readFileSync(configPath, 'utf8');
+    return readConfigVar(content, varName);
+}

package/src/lib/train-config-validator.js

@@ -0,0 +1,112 @@
+// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+// SPDX-License-Identifier: Apache-2.0
+
+/**
+ * Train Config Validator
+ *
+ * Validates training configuration objects parsed from do/training/config.yaml.
+ * Checks that all required fields are present and provides descriptive error
+ * messages naming the specific missing field.
+ *
+ * This module mirrors the validation logic in the bash `_validate_config`
+ * function in templates/do/train, enabling property-based testing of the
+ * validation rules in isolation.
+ *
+ * Requirements: 2.12, 10.1
+ */
+
+/**
+ * Required fields for a valid training configuration.
+ * Each entry maps the field name to a human-readable description and expected format.
+ */
+export const REQUIRED_FIELDS = {
+    image: {
+        description: 'The container image URI',
+        format: 'image: "123456789012.dkr.ecr.us-east-1.amazonaws.com/my-training:latest"'
+    },
+    script: {
+        description: 'The training script S3 path',
+        format: 'script: "s3://my-bucket/scripts/train.py"'
+    },
+    instance_type: {
+        description: 'The SageMaker instance type',
+        format: 'instance_type: "ml.g5.xlarge"'
+    },
+    dataset: {
+        description: 'The S3 dataset path',
+        format: 'dataset: "s3://my-bucket/data/train/"'
+    },
+    output_path: {
+        description: 'The S3 output path',
+        format: 'output_path: "s3://my-bucket/output/"'
+    }
+};
+
+/**
+ * Validate that all required fields are present in a training config.
+ *
+ * @param {Object} config - The parsed training configuration object
+ * @returns {{ valid: boolean, errors: Array<{ field: string, message: string }> }}
+ *   - valid: true if all required fields are present and non-empty
+ *   - errors: array of error objects, each naming the missing field
+ */
+export function validateRequiredFields(config) {
+    const errors = [];
+
+    for (const [field, meta] of Object.entries(REQUIRED_FIELDS)) {
+        const value = config ? config[field] : undefined;
+
+        if (value === undefined || value === null || value === '') {
+            errors.push({
+                field,
+                message: `Missing required field: ${field}\n   ${meta.description} is required in do/training/config.yaml\n\n   Expected format: ${meta.format}`
+            });
+        }
+    }
+
+    return {
+        valid: errors.length === 0,
+        errors
+    };
+}
+
+/**
+ * Validate spot training checkpoint requirement.
+ * When enable_spot is true, checkpoint_path must be specified.
+ *
+ * @param {Object} config - The parsed training configuration object
+ * @returns {{ valid: boolean, errors: Array<{ field: string, message: string }> }}
+ */
+export function validateSpotConfig(config) {
+    const errors = [];
+
+    if (config && config.enable_spot === true && (!config.checkpoint_path || config.checkpoint_path === '')) {
+        errors.push({
+            field: 'checkpoint_path',
+            message: 'Checkpoint path required for spot training\n   When enable_spot is true, a checkpoint S3 path must be specified\n   so training can resume after spot interruptions.'
+        });
+    }
+
+    return {
+        valid: errors.length === 0,
+        errors
+    };
+}
+
+/**
+ * Full validation of a training config — checks required fields and spot config.
+ *
+ * @param {Object} config - The parsed training configuration object
+ * @returns {{ valid: boolean, errors: Array<{ field: string, message: string }> }}
+ */
+export function validateTrainingConfig(config) {
+    const requiredResult = validateRequiredFields(config);
+    const spotResult = validateSpotConfig(config);
+
+    const allErrors = [...requiredResult.errors, ...spotResult.errors];
+
+    return {
+        valid: allErrors.length === 0,
+        errors: allErrors
+    };
+}

package/src/lib/train-feedback.js

@@ -0,0 +1,46 @@
+// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+// SPDX-License-Identifier: Apache-2.0
+
+/**
+ * Train Feedback Loop — JavaScript equivalent of do/lib/feedback.sh
+ *
+ * Generates post-completion feedback output with artifact locations
+ * and deployment suggestions based on artifact type.
+ */
+
+/**
+ * Generate completion feedback output for a training/tuning job.
+ *
+ * Replicates the logic of print_completion_feedback() in do/lib/feedback.sh.
+ *
+ * @param {object} params
+ * @param {string} params.outputPath - S3 URI to the output artifacts
+ * @param {string} params.outputType - "adapter" or "full-model"
+ * @param {string} params.jobName - Job name for reference
+ * @param {string} [params.modelPackageArn] - Optional model package ARN
+ * @returns {string} The formatted feedback output
+ */
+export function generateCompletionFeedback({ outputPath, outputType, jobName, modelPackageArn = '' }) {
+    const lines = [];
+
+    lines.push('');
+    lines.push(`✅ Training complete: ${jobName}`);
+    lines.push('');
+    lines.push(`   Artifacts: ${outputPath}`);
+    if (modelPackageArn) {
+        lines.push(`   Model Package: ${modelPackageArn}`);
+    }
+    lines.push('');
+    lines.push('   Next steps:');
+
+    if (outputType === 'adapter') {
+        lines.push(`     • Deploy as LoRA adapter:  ./do/adapter add my-adapter --weights ${outputPath}`);
+        lines.push('     • (Requires running endpoint with LoRA enabled)');
+    } else if (outputType === 'full-model') {
+        lines.push(`     • Deploy as new IC:        ./do/add-ic my-model --model-data ${outputPath}`);
+        lines.push(`     • Replace current base:    ./do/deploy --force-ic --model-data ${outputPath}`);
+    }
+    lines.push('');
+
+    return lines.join('\n');
+}

package/src/lib/train-idempotency.js

@@ -0,0 +1,97 @@
+// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+// SPDX-License-Identifier: Apache-2.0
+
+/**
+ * Train Idempotency Decision Logic
+ *
+ * Models the idempotency check logic from the bash `_check_idempotency` function
+ * in `templates/do/train` as a pure JavaScript function for property-based testing.
+ *
+ * The idempotency pattern:
+ * - If --force is set, always create a new job regardless of existing status
+ * - If no existing job, create a new job
+ * - If existing job is InProgress, poll it
+ * - If existing job is Completed, display results
+ * - If existing job is Failed or Stopped, display failure and suggest --force
+ *
+ * Requirements: 5.1–5.5
+ */
+
+/**
+ * Valid existing job statuses that SageMaker can report.
+ */
+export const JOB_STATUSES = ['InProgress', 'Completed', 'Failed', 'Stopped'];
+
+/**
+ * Possible actions the train script can take after the idempotency check.
+ */
+export const ACTIONS = {
+    CREATE_NEW_JOB: 'create_new_job',
+    POLL_EXISTING: 'poll_existing',
+    DISPLAY_RESULTS: 'display_results',
+    DISPLAY_FAILURE: 'display_failure'
+};
+
+/**
+ * Determine the action to take based on existing job status and force flag.
+ *
+ * This mirrors the bash `_check_idempotency` logic in a testable form:
+ * - force=true → always create_new_job
+ * - no existing status (null/empty) → create_new_job
+ * - InProgress → poll_existing
+ * - Completed → display_results
+ * - Failed → display_failure
+ * - Stopped → display_failure
+ *
+ * @param {string|null|undefined} existingStatus - The current job status from DescribeTrainingJob
+ * @param {boolean} forceFlag - Whether --force was specified
+ * @returns {{ action: string, reason: string }}
+ *   - action: one of ACTIONS values
+ *   - reason: human-readable explanation of why this action was chosen
+ */
+export function determineAction(existingStatus, forceFlag) {
+    // Force flag always overrides — create a new job regardless of existing status
+    if (forceFlag === true) {
+        return {
+            action: ACTIONS.CREATE_NEW_JOB,
+            reason: '--force specified, creating new job regardless of existing status'
+        };
+    }
+
+    // No existing job — create a new one
+    if (!existingStatus || existingStatus === '') {
+        return {
+            action: ACTIONS.CREATE_NEW_JOB,
+            reason: 'No existing job found, creating new job'
+        };
+    }
+
+    // Existing job found — action depends on status
+    switch (existingStatus) {
+    case 'InProgress':
+        return {
+            action: ACTIONS.POLL_EXISTING,
+            reason: `Existing job is ${existingStatus}, resuming polling`
+        };
+
+    case 'Completed':
+        return {
+            action: ACTIONS.DISPLAY_RESULTS,
+            reason: `Existing job is ${existingStatus}, displaying results`
+        };
+
+    case 'Failed':
+    case 'Stopped':
+        return {
+            action: ACTIONS.DISPLAY_FAILURE,
+            reason: `Existing job is ${existingStatus}, suggest --force to create new job`
+        };
+
+    default:
+        // Unknown status — treat as failure, suggest --force
+        return {
+            action: ACTIONS.DISPLAY_FAILURE,
+            reason: `Unexpected job status: ${existingStatus}, suggest --force`
+        };
+    }
+}

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/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()