@aws/ml-container-creator 0.9.0 → 0.10.0

package/src/lib/cuda-resolver.js

@@ -0,0 +1,140 @@
+// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+// SPDX-License-Identifier: Apache-2.0
+
+/**
+ * CUDA Resolver - Handles CUDA version selection and AMI resolution.
+ * Uses delegation pattern: receives parent PromptRunner reference to access shared state.
+ */
+
+export const CUDA_AMI_MAP = {
+    '11.0': 'al2-ami-sagemaker-inference-gpu-2',
+    '11.4': 'al2-ami-sagemaker-inference-gpu-2-1',
+    '11.8': 'al2-ami-sagemaker-inference-gpu-2-1',
+    '12.1': 'al2-ami-sagemaker-inference-gpu-3-1',
+    '12.2': 'al2-ami-sagemaker-inference-gpu-3-1',
+    '12.4': 'al2-ami-sagemaker-inference-gpu-3-1',
+    '12.6': 'al2-ami-sagemaker-inference-gpu-3-1',
+    '13.0': 'al2023-ami-sagemaker-inference-gpu-4-1'
+};
+
+export default class CudaResolver {
+    constructor(runner) {
+        this.runner = runner;
+    }
+
+    /**
+     * Prompt the user to select a CUDA version when the selected GPU instance
+     * supports multiple versions.
+     *
+     * @param {string} instanceType - Selected instance type (e.g. "ml.g5.2xlarge")
+     * @param {string} framework - Selected framework name
+     * @param {string} frameworkVersion - Selected framework version
+     * @param {string} [baseImageCuda] - CUDA version from selected base image (for auto-resolution)
+     * @returns {Promise<{cudaVersion: string, inferenceAmiVersion: string}|null>}
+     */
+    async _promptCudaVersion(instanceType, framework, frameworkVersion, baseImageCuda) {
+        if (!instanceType) return null;
+
+        const instanceInfo = this.runner._instanceAcceleratorMapping[instanceType];
+        if (!instanceInfo || instanceInfo.accelerator.type !== 'cuda') return null;
+
+        const instanceCudaVersions = instanceInfo.accelerator.versions;
+        if (!instanceCudaVersions || instanceCudaVersions.length === 0) return null;
+
+        // Auto-resolution: when base image specifies a CUDA version, intersect with instance support
+        if (baseImageCuda) {
+            const majorRequired = baseImageCuda.split('.')[0];
+            const intersection = instanceCudaVersions.filter(v => {
+                if (v === baseImageCuda) return true;
+                if (v.startsWith(`${majorRequired}.`)) return true;
+                return false;
+            });
+
+            if (intersection.length > 0) {
+                const exactMatch = intersection.find(v => v === baseImageCuda);
+                const selectedVersion = exactMatch || intersection.sort().pop();
+                const inferenceAmiVersion = CUDA_AMI_MAP[selectedVersion];
+                if (inferenceAmiVersion) {
+                    console.log(`\n🔧 CUDA ${selectedVersion} auto-resolved from base image (requires ${baseImageCuda})`);
+                    console.log(`   AMI: ${inferenceAmiVersion}`);
+                    return { cudaVersion: selectedVersion, inferenceAmiVersion };
+                }
+            } else {
+                console.log(`\n   ⚠️  Base image requires CUDA ${baseImageCuda} but instance ${instanceType} supports: ${instanceCudaVersions.join(', ')}`);
+                console.log('   No compatible CUDA version found. Falling back to manual selection.');
+            }
+        }
+
+        // Get framework CUDA requirements (if available)
+        const registryConfigManager = this.runner.registryConfigManager;
+        const frameworkConfig = registryConfigManager?.frameworkRegistry?.[framework]?.[frameworkVersion];
+        const frameworkAccel = frameworkConfig?.accelerator;
+
+        // Compute compatible CUDA versions
+        let compatibleVersions;
+        if (frameworkAccel?.versionRange) {
+            const { min, max } = frameworkAccel.versionRange;
+            compatibleVersions = instanceCudaVersions.filter(v => {
+                return v >= min && v <= max;
+            });
+        } else {
+            compatibleVersions = [...instanceCudaVersions];
+        }
+
+        if (compatibleVersions.length === 0) {
+            compatibleVersions = [...instanceCudaVersions];
+        }
+
+        // If only one option, auto-select it silently
+        if (compatibleVersions.length === 1) {
+            const cudaVersion = compatibleVersions[0];
+            const inferenceAmiVersion = CUDA_AMI_MAP[cudaVersion];
+            if (inferenceAmiVersion) {
+                console.log(`\n🔧 CUDA ${cudaVersion} auto-selected (only compatible version for ${instanceType})`);
+                console.log(`   AMI: ${inferenceAmiVersion}`);
+            }
+            return inferenceAmiVersion ? { cudaVersion, inferenceAmiVersion } : null;
+        }
+
+        // Multiple options — let the user choose (or auto-select in auto-prompt mode)
+        const defaultVersion = frameworkAccel?.version
+            && compatibleVersions.includes(frameworkAccel.version)
+            ? frameworkAccel.version
+            : instanceInfo.accelerator.default || compatibleVersions[compatibleVersions.length - 1];
+
+        // In auto-prompt mode, auto-select the default without prompting
+        if (this.runner.configManager?.isAutoPrompt()) {
+            const inferenceAmiVersion = CUDA_AMI_MAP[defaultVersion];
+            if (inferenceAmiVersion) {
+                console.log(`\n🔧 CUDA ${defaultVersion} auto-selected (auto-prompt mode)`);
+                console.log(`   AMI: ${inferenceAmiVersion}`);
+            }
+            return inferenceAmiVersion ? { cudaVersion: defaultVersion, inferenceAmiVersion } : null;
+        }
+
+        const choices = compatibleVersions.map(v => {
+            const ami = CUDA_AMI_MAP[v] || 'unknown';
+            const isDefault = v === defaultVersion ? ' (recommended)' : '';
+            return {
+                name: `CUDA ${v}${isDefault}  →  AMI: ${ami}`,
+                value: v,
+                short: `CUDA ${v}`
+            };
+        });
+
+        const { cudaVersion } = await this.runner._runPrompts([{
+            type: 'list',
+            name: 'cudaVersion',
+            message: `Select CUDA version for ${instanceType} (${instanceInfo.accelerator.hardware}):`,
+            choices,
+            default: defaultVersion
+        }]);
+
+        const inferenceAmiVersion = CUDA_AMI_MAP[cudaVersion];
+        if (inferenceAmiVersion) {
+            console.log(`   ✅ CUDA ${cudaVersion} → AMI: ${inferenceAmiVersion}`);
+        }
+
+        return inferenceAmiVersion ? { cudaVersion, inferenceAmiVersion } : null;
+    }
+}

package/src/lib/e2e-catalog-validator.js

@@ -12,6 +12,8 @@
  */
 
 import Ajv from 'ajv';
+import { readFileSync } from 'node:fs';
+import path from 'node:path';
 
 const catalogSchema = {
     $schema: 'http://json-schema.org/draft-07/schema#',
@@ -34,20 +36,214 @@ const catalogSchema = {
                         items: { type: 'string', pattern: '^[a-z][a-z0-9-]*$' },
                         minItems: 1
                     },
-                    timeout: { type: 'integer', minimum: 60 }
+                    timeout: { type: 'integer', minimum: 60 },
+                    tuneTimeout: { type: 'integer', minimum: 60 },
+                    tuneConfig: {
+                        type: 'object',
+                        required: ['tuneId', 'technique', 'trainingType', 'dataset'],
+                        properties: {
+                            tuneId: { type: 'string' },
+                            technique: { type: 'string', enum: ['sft', 'dpo', 'rlaif', 'rlvr'] },
+                            trainingType: { type: 'string', enum: ['lora', 'full-rank'] },
+                            dataset: { type: 'string' }
+                        },
+                        additionalProperties: false
+                    }
                 }
             }
         }
     }
 };
 
+/**
+ * Custom validation rules for tune lifecycle entries.
+ *
+ * Enforces:
+ * - Entries with tune lifecycle steps (steps starting with "tune-") must have a `tuneConfig` object
+ * - Entries with tune lifecycle steps must include `--enable-lora` in their `args` field
+ * - If `tuneTimeout` is present, it must be a positive integer >= 60
+ *
+ * @param {Object} catalog - The catalog object with a `configs` array
+ * @param {Array<{ path: string, message: string }>} errors - Errors array to append to
+ */
+export function validateTuneConstraints(catalog, errors) {
+    if (!catalog || !Array.isArray(catalog.configs)) {
+        return;
+    }
+
+    for (let i = 0; i < catalog.configs.length; i++) {
+        const entry = catalog.configs[i];
+        if (!entry || !Array.isArray(entry.lifecycle)) {
+            continue;
+        }
+
+        const hasTuneSteps = entry.lifecycle.some((s) => typeof s === 'string' && s.startsWith('tune-'));
+
+        if (hasTuneSteps) {
+            // Must have tuneConfig
+            if (!entry.tuneConfig) {
+                errors.push({
+                    path: `/configs/${i}`,
+                    message: `entry "${entry.id}" has tune lifecycle steps but no tuneConfig`
+                });
+                continue;
+            }
+
+            // Must have --enable-lora in args
+            if (!entry.args || !entry.args.includes('--enable-lora')) {
+                errors.push({
+                    path: `/configs/${i}/args`,
+                    message: `entry "${entry.id}" has tune steps but args missing --enable-lora`
+                });
+            }
+        }
+
+        // tuneTimeout validation (if present)
+        if (entry.tuneTimeout !== undefined) {
+            if (typeof entry.tuneTimeout !== 'number' || !Number.isInteger(entry.tuneTimeout) || entry.tuneTimeout < 60) {
+                errors.push({
+                    path: `/configs/${i}/tuneTimeout`,
+                    message: `entry "${entry.id}": tuneTimeout must be a positive integer >= 60`
+                });
+            }
+        }
+    }
+}
+
+/**
+ * Validate lifecycle ordering for tune-group steps.
+ *
+ * Tune-group steps are: any step starting with "tune-", "adapter-add", and "test-adapter".
+ * These must appear AFTER the "test" step and BEFORE the "clean" step in the lifecycle array.
+ *
+ * @param {Object} catalog - The catalog object with a `configs` array
+ * @param {Array<{ path: string, message: string }>} errors - Errors array to append to
+ */
+export function validateLifecycleOrdering(catalog, errors) {
+    if (!catalog || !Array.isArray(catalog.configs)) {
+        return;
+    }
+
+    for (let i = 0; i < catalog.configs.length; i++) {
+        const entry = catalog.configs[i];
+        if (!entry || !Array.isArray(entry.lifecycle)) {
+            continue;
+        }
+
+        // Identify tune-group steps in this entry's lifecycle
+        const tuneGroupSteps = entry.lifecycle.filter(
+            (s) => typeof s === 'string' && (s.startsWith('tune-') || s === 'adapter-add' || s === 'test-adapter')
+        );
+
+        if (tuneGroupSteps.length === 0) {
+            continue;
+        }
+
+        const testIdx = entry.lifecycle.indexOf('test');
+        const cleanIdx = entry.lifecycle.indexOf('clean');
+
+        for (const step of tuneGroupSteps) {
+            const stepIdx = entry.lifecycle.indexOf(step);
+
+            if (testIdx >= 0 && stepIdx <= testIdx) {
+                errors.push({
+                    path: `/configs/${i}/lifecycle`,
+                    message: `entry "${entry.id}": "${step}" must come after "test"`
+                });
+            }
+
+            if (cleanIdx >= 0 && stepIdx >= cleanIdx) {
+                errors.push({
+                    path: `/configs/${i}/lifecycle`,
+                    message: `entry "${entry.id}": "${step}" must come before "clean"`
+                });
+            }
+        }
+    }
+}
+
+/**
+ * Validate tuneConfig entries against tune-catalog.json.
+ * This is a "soft" validation — returns an empty errors array if tune-catalog is unavailable.
+ *
+ * For each catalog entry with a tuneConfig:
+ * - Checks that tuneConfig.tuneId exists as a key in the tune-catalog's models object
+ * - Checks that the specified technique is supported for that model
+ * - Checks that the specified trainingType is in the technique's trainingTypes array
+ *
+ * @param {Object} catalog - The catalog object with a configs array
+ * @param {string} tuneCatalogPath - Path to the tune-catalog.json file
+ * @returns {Array<{ path: string, message: string }>} Array of validation errors
+ */
+export function validateTuneCatalogReferences(catalog, tuneCatalogPath) {
+    const errors = [];
+
+    if (!catalog || !Array.isArray(catalog.configs)) {
+        return errors;
+    }
+
+    let tuneCatalog;
+    try {
+        const raw = readFileSync(tuneCatalogPath, 'utf8');
+        tuneCatalog = JSON.parse(raw);
+    } catch {
+        // tune-catalog not available or unparseable — skip cross-reference validation
+        return errors;
+    }
+
+    if (!tuneCatalog || !tuneCatalog.models) {
+        return errors;
+    }
+
+    for (let i = 0; i < catalog.configs.length; i++) {
+        const entry = catalog.configs[i];
+        if (!entry || !entry.tuneConfig) {
+            continue;
+        }
+
+        const { tuneId, technique, trainingType } = entry.tuneConfig;
+
+        // Check tuneId exists in tune-catalog models
+        const tuneModel = tuneCatalog.models[tuneId];
+        if (!tuneModel) {
+            errors.push({
+                path: `/configs/${i}/tuneConfig/tuneId`,
+                message: `entry "${entry.id}": tuneId "${tuneId}" not found in tune-catalog`
+            });
+            continue;
+        }
+
+        // Check technique is supported for this model
+        if (!tuneModel.techniques || !tuneModel.techniques[technique]) {
+            errors.push({
+                path: `/configs/${i}/tuneConfig/technique`,
+                message: `entry "${entry.id}": technique "${technique}" not supported for model "${tuneId}"`
+            });
+            continue;
+        }
+
+        // Check trainingType is in the technique's trainingTypes array
+        const supportedTrainingTypes = tuneModel.techniques[technique].trainingTypes;
+        if (!Array.isArray(supportedTrainingTypes) || !supportedTrainingTypes.includes(trainingType)) {
+            errors.push({
+                path: `/configs/${i}/tuneConfig/trainingType`,
+                message: `entry "${entry.id}": trainingType "${trainingType}" not supported for ${tuneId}/${technique}`
+            });
+        }
+    }
+
+    return errors;
+}
+
 /**
  * Validate an e2e catalog object against the schema and uniqueness constraints.
  *
  * @param {Object} catalog - The catalog object to validate
+ * @param {Object} [options] - Optional configuration
+ * @param {string} [options.tuneCatalogPath] - Custom path to tune-catalog.json for cross-reference validation
  * @returns {{ valid: boolean, errors?: Array<{ path: string, message: string }> }}
  */
-export function validateCatalog(catalog) {
+export function validateCatalog(catalog, options = {}) {
     const ajv = new Ajv({ allErrors: true, strict: false });
     const validate = ajv.compile(catalogSchema);
 
@@ -56,9 +252,19 @@ export function validateCatalog(catalog) {
 
     if (!valid) {
         for (const err of validate.errors) {
+            // Try to include the entry id in the error message for better diagnostics
+            let message = err.message;
+            const pathMatch = (err.instancePath || '').match(/^\/configs\/(\d+)/);
+            if (pathMatch && catalog && Array.isArray(catalog.configs)) {
+                const idx = parseInt(pathMatch[1], 10);
+                const entry = catalog.configs[idx];
+                if (entry && typeof entry.id === 'string') {
+                    message = `entry "${entry.id}": ${err.message}`;
+                }
+            }
             errors.push({
                 path: err.instancePath || '/',
-                message: err.message
+                message
             });
         }
     }
@@ -81,6 +287,17 @@ export function validateCatalog(catalog) {
         }
     }
 
+    // Custom check: tune constraints (runs after Ajv schema validation)
+    validateTuneConstraints(catalog, errors);
+
+    // Custom check: lifecycle ordering for tune-group steps
+    validateLifecycleOrdering(catalog, errors);
+
+    // Cross-reference validation against tune-catalog.json
+    const tuneCatalogPath = options.tuneCatalogPath || path.resolve(process.cwd(), 'config/tune-catalog.json');
+    const crossRefErrors = validateTuneCatalogReferences(catalog, tuneCatalogPath);
+    errors.push(...crossRefErrors);
+
     if (errors.length > 0) {
         return { valid: false, errors };
     }
@@ -101,3 +318,34 @@ export function filterByTier(catalog, tier) {
     }
     return catalog.configs.filter((config) => config.tier === tier);
 }
+
+/**
+ * Filter configs by a specific config id.
+ *
+ * First attempts to find the config within the provided (tier-filtered) configs array.
+ * If not found there, falls back to searching the full catalog — this is a convenience
+ * for re-runs where the user specifies --config without matching the tier.
+ *
+ * @param {Array<Object>} configs - Pre-filtered configs (e.g., from filterByTier)
+ * @param {Object} catalog - The full catalog object with a `configs` array
+ * @param {string} configId - The config id to filter by
+ * @returns {Array<Object>} Matching configs (0 or 1 element)
+ */
+export function filterByConfig(configs, catalog, configId) {
+    if (!configId) {
+        return configs;
+    }
+
+    // First try within the already-filtered set (tier-filtered)
+    const inTier = configs.filter(c => c.id === configId);
+    if (inTier.length > 0) {
+        return inTier;
+    }
+
+    // Fallback: search the full catalog (convenience for --config without matching tier)
+    if (catalog && Array.isArray(catalog.configs)) {
+        return catalog.configs.filter(c => c.id === configId);
+    }
+
+    return [];
+}

package/src/lib/e2e-ci-recorder.js

@@ -0,0 +1,103 @@
+// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+// SPDX-License-Identifier: Apache-2.0
+
+/**
+ * E2E CI Recorder
+ *
+ * Records E2E validation results to the DynamoDB CI table.
+ * Gracefully degrades if the CI table is not provisioned —
+ * all recording calls become no-ops.
+ */
+
+import { computeConfigId } from './ci-register-helpers.js';
+import BootstrapConfig from './bootstrap-config.js';
+
+export class E2ECIRecorder {
+    constructor() {
+        this.config = new BootstrapConfig();
+        this.client = null;
+        this.tableName = null;
+    }
+
+    /**
+     * Initialize the recorder by checking if the CI table is provisioned.
+     * If not provisioned, logs a warning and returns false — all subsequent
+     * recordConfigResult calls become no-ops.
+     *
+     * @returns {Promise<boolean>} true if ready to record, false otherwise
+     */
+    async init() {
+        const profile = this.config.getActiveProfileWithDefaults();
+        if (!profile || !profile.config.ciInfraProvisioned) {
+            console.warn('⚠️  CI table not provisioned — skipping result recording');
+            return false;
+        }
+        this.tableName = profile.config.ciTableName;
+        const { DynamoDBClient } = await import('@aws-sdk/client-dynamodb');
+        this.client = new DynamoDBClient({ region: profile.config.awsRegion });
+        return true;
+    }
+
+    /**
+     * Record a config's E2E result to the DynamoDB CI table.
+     * No-op if init() returned false or was not called.
+     *
+     * @param {object} catalogEntry - The catalog entry that was tested
+     * @param {object} configResult - The result of running the config
+     */
+    async recordConfigResult(catalogEntry, configResult) {
+        if (!this.client) return;
+
+        const configId = this.deriveConfigId(catalogEntry);
+        const item = {
+            configId,
+            schemaVersion: 2,
+            testStatus: configResult.status === 'pass'
+                ? 'pass'
+                : `fail-${configResult.steps.find(s => s.status === 'fail')?.name || 'unknown'}`,
+            lastTestTimestamp: new Date().toISOString(),
+            stageResults: Object.fromEntries(
+                configResult.steps.map(s => [s.name, { status: s.status, duration: s.duration, error: s.error || '' }])
+            ),
+            e2eCatalogId: catalogEntry.id,
+            tier: catalogEntry.tier,
+            duration: configResult.duration
+        };
+
+        try {
+            const { PutItemCommand } = await import('@aws-sdk/client-dynamodb');
+            const { marshall } = await import('@aws-sdk/util-dynamodb');
+            await this.client.send(new PutItemCommand({
+                TableName: this.tableName,
+                Item: marshall(item, { removeUndefinedValues: true })
+            }));
+        } catch (err) {
+            console.warn(`⚠️  Failed to record ${catalogEntry.id} to CI table: ${err.message}`);
+        }
+    }
+
+    /**
+     * Derive a deterministic configId from a catalog entry's args field,
+     * using the same SHA256 algorithm as `do/register --ci`.
+     *
+     * @param {object} catalogEntry - A catalog entry with an `args` string
+     * @returns {string} 16-character hex configId
+     */
+    deriveConfigId(catalogEntry) {
+        const args = Object.fromEntries(
+            catalogEntry.args.split(/\s+/)
+                .filter(a => a.startsWith('--'))
+                .map(a => a.replace(/^--/, '').split('='))
+        );
+        const deploymentTarget = catalogEntry.track === 'realtime'
+            ? 'realtime-inference'
+            : catalogEntry.track;
+        return computeConfigId(
+            args['deployment-config'] || '',
+            args['model-name'] || 'none',
+            args['instance-type'] || '',
+            args['region'] || 'us-west-2',
+            deploymentTarget
+        );
+    }
+}