@aws/ml-container-creator 0.2.5 → 0.3.0

package/src/app.js

@@ -119,6 +119,23 @@ export async function run(projectName, options) {
     let answers;
     if (configManager.shouldSkipPrompts()) {
         console.log('\n🚀 Skipping prompts - using configuration from other sources');
+
+        // Fail-fast if required parameters are missing
+        const missing = configManager.getMissingRequiredParameters();
+        if (missing.length > 0) {
+            console.error('\n❌ Cannot skip prompts — required parameters are missing:\n');
+            for (const param of missing) {
+                const matrix = configManager._getParameterMatrix()[param];
+                const cliFlag = matrix?.cliOption ? `--${matrix.cliOption}` : '';
+                const envVar = matrix?.envVar || '';
+                const hints = [cliFlag, envVar].filter(Boolean).join(' or ');
+                console.error(`   • ${param}${hints ? ` (${hints})` : ''}`);
+            }
+            console.error('\n   Provide these via CLI flags, environment variables, or a config file.');
+            console.error('   Run "ml-container-creator --help" for available options.\n');
+            process.exit(1);
+        }
+
         answers = configManager.getFinalConfiguration();
 
         // Infer modelSource from model name prefix if not set
@@ -188,6 +205,23 @@ export async function run(projectName, options) {
 
     // --- Phase: Writing ---
     const destDir = path.resolve(answers.destinationDir);
+
+    // Safety guard: refuse to generate into the generator's own directory
+    const destPkgPath = path.join(destDir, 'package.json');
+    if (fs.existsSync(destPkgPath)) {
+        try {
+            const destPkg = JSON.parse(fs.readFileSync(destPkgPath, 'utf8'));
+            if (destPkg.name === '@aws/ml-container-creator') {
+                console.log('\n❌ Refusing to generate into the generator\'s own directory.');
+                console.log('   This would overwrite the generator source files.');
+                console.log('   Use --project-dir or provide a project name instead.\n');
+                return;
+            }
+        } catch {
+            // If we can't read/parse package.json, it's not the generator dir — proceed
+        }
+    }
+
     fs.mkdirSync(destDir, { recursive: true });
 
     await writeProject(TEMPLATE_DIR, destDir, answers, registryConfigManager, tritonBackends, configManager);
@@ -476,7 +510,9 @@ async function _ensureTemplateVariables(answers, registryConfigManager = null) {
         chatTemplate: null,
         chatTemplateSource: null,
         hfToken: null,
+        hfTokenArn: null,
         ngcApiKey: null,
+        ngcTokenArn: null,
         envVars: {},
         inferenceAmiVersion: null,
         accelerator: null,

package/src/lib/architecture-sync.js

@@ -0,0 +1,171 @@
+// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+// SPDX-License-Identifier: Apache-2.0
+
+/**
+ * Architecture Sync
+ *
+ * Fetches model registry source files from server GitHub repositories
+ * and extracts supported model_type values into the model-servers catalog.
+ */
+
+import { readFileSync, writeFileSync } from 'node:fs';
+
+/**
+ * Parse vLLM's model registry Python source to extract model_type keys.
+ *
+ * vLLM's registry maps architecture class names to (module, impl_class) tuples:
+ *   "LlamaForCausalLM": ("llama", "LlamaForCausalLM"),
+ *   "Qwen2ForCausalLM": ("qwen2", "Qwen2ForCausalLM"),
+ *
+ * The module name (first tuple element) corresponds to the model_type.
+ * Also matches older formats where model_type is used directly as dict key.
+ *
+ * @param {string} source - Python source code content
+ * @returns {string[]} Sorted array of model_type strings
+ */
+export const parseVllmRegistry = (source) => {
+    const modelTypes = new Set();
+    const patterns = [
+        // Tuple value format: ("module_name", "ClassName") — extract module_name
+        /\("([a-z][a-z0-9_]*)"\s*,\s*"[A-Z]/g,
+        // Direct lowercase key format (older registries): "model_type": (
+        /"([a-z][a-z0-9_]*)":\s*\(/g,
+        // Direct lowercase key format: "model_type": ClassName
+        /"([a-z][a-z0-9_]*)":\s*[A-Z]/g,
+        // Direct lowercase key format: "model_type": [
+        /"([a-z][a-z0-9_]*)":\s*\[/g
+    ];
+    for (const pattern of patterns) {
+        let match;
+        while ((match = pattern.exec(source)) !== null) {
+            modelTypes.add(match[1]);
+        }
+    }
+    return [...modelTypes].sort();
+};
+
+/**
+ * Parse SGLang's model_registry.py to extract model_type keys.
+ *
+ * Matches patterns like:
+ *   "llama": ModelClass,
+ *   "qwen2": (ModulePath, ClassName),
+ *
+ * @param {string} source - Python source code content
+ * @returns {string[]} Sorted array of model_type strings
+ */
+export const parseSglangRegistry = (source) => {
+    const modelTypes = new Set();
+    const patterns = [
+        /"([a-z][a-z0-9_]*)":\s*\(/g,
+        /"([a-z][a-z0-9_]*)":\s*[A-Z]/g,
+        /"([a-z][a-z0-9_]*)":\s*\[/g
+    ];
+    for (const pattern of patterns) {
+        let match;
+        while ((match = pattern.exec(source)) !== null) {
+            modelTypes.add(match[1]);
+        }
+    }
+    return [...modelTypes].sort();
+};
+
+/**
+ * Parse TensorRT-LLM's models __init__.py to extract model_type keys.
+ *
+ * Matches patterns from MODEL_MAP or similar dict structures:
+ *   "llama": LlamaForCausalLM,
+ *   "gpt2": GPT2LMHeadModel,
+ *
+ * @param {string} source - Python source code content
+ * @returns {string[]} Sorted array of model_type strings
+ */
+export const parseTensorRTRegistry = (source) => {
+    const modelTypes = new Set();
+    const patterns = [
+        /"([a-z][a-z0-9_]*)":\s*[A-Z]/g,
+        /"([a-z][a-z0-9_]*)":\s*\(/g,
+        /'([a-z][a-z0-9_]*)':\s*[A-Z]/g,
+        /'([a-z][a-z0-9_]*)':\s*\(/g
+    ];
+    for (const pattern of patterns) {
+        let match;
+        while ((match = pattern.exec(source)) !== null) {
+            modelTypes.add(match[1]);
+        }
+    }
+    return [...modelTypes].sort();
+};
+
+/**
+ * Configuration mapping each server to its GitHub repository,
+ * registry file path, tag prefix, and parser function.
+ */
+export const SERVER_REGISTRY_SOURCES = {
+    vllm: {
+        repo: 'vllm-project/vllm',
+        file: 'vllm/model_executor/models/registry.py',
+        tagPrefix: 'v',
+        parser: parseVllmRegistry
+    },
+    sglang: {
+        repo: 'sgl-project/sglang',
+        file: 'python/sglang/srt/models/model_registry.py',
+        tagPrefix: 'v',
+        parser: parseSglangRegistry
+    },
+    'tensorrt-llm': {
+        repo: 'NVIDIA/TensorRT-LLM',
+        file: 'tensorrt_llm/models/__init__.py',
+        tagPrefix: 'v',
+        parser: parseTensorRTRegistry
+    }
+};
+
+/**
+ * Sync supported model architectures from server GitHub repositories
+ * into the model-servers catalog.
+ *
+ * For each server entry in the catalog that has a matching source config,
+ * fetches the model registry file from GitHub at the version tag and
+ * parses it to extract supported model_type values.
+ *
+ * @param {string} catalogPath - Path to model-servers.json
+ * @returns {object} Summary with counts and failures
+ */
+export const syncArchitectures = async (catalogPath) => {
+    const catalog = JSON.parse(readFileSync(catalogPath, 'utf8'));
+    const summary = { servers: [], failures: [] };
+
+    for (const [server, entries] of Object.entries(catalog)) {
+        const source = SERVER_REGISTRY_SOURCES[server];
+        if (!source) continue;
+
+        for (const entry of entries) {
+            const version = entry.labels?.framework_version;
+            if (!version) continue;
+
+            const tag = `${source.tagPrefix}${version}`;
+            const url = `https://raw.githubusercontent.com/${source.repo}/${tag}/${source.file}`;
+
+            try {
+                const response = await fetch(url);
+                if (!response.ok) {
+                    summary.failures.push({ server, version, reason: `HTTP ${response.status}` });
+                    console.log(`   ⚠️  ${server} ${version}: fetch failed (HTTP ${response.status})`);
+                    continue;
+                }
+                const content = await response.text();
+                entry.supportedModelTypes = source.parser(content);
+                summary.servers.push({ server, version, count: entry.supportedModelTypes.length });
+                console.log(`   ✓ ${server} ${version}: ${entry.supportedModelTypes.length} architectures`);
+            } catch (err) {
+                summary.failures.push({ server, version, reason: err.message });
+                console.log(`   ⚠️  ${server} ${version}: fetch failed (${err.message})`);
+            }
+        }
+    }
+
+    writeFileSync(catalogPath, JSON.stringify(catalog, null, 4));
+    return summary;
+};

package/src/lib/arn-detection.js

@@ -0,0 +1,22 @@
+// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+// SPDX-License-Identifier: Apache-2.0
+
+/**
+ * ARN Detection Utility
+ *
+ * Provides a pure function for distinguishing AWS Secrets Manager ARNs
+ * from plaintext values. Used by the prompt flow and CLI to determine
+ * whether user input should be treated as a secret reference or a
+ * literal token value.
+ */
+
+const SECRETS_MANAGER_ARN_PREFIX = 'arn:aws:secretsmanager:';
+
+/**
+ * Determines if a value is a Secrets Manager ARN.
+ * @param {*} value - The input value to check
+ * @returns {boolean} True if the value is a Secrets Manager ARN
+ */
+export function isSecretsManagerArn(value) {
+    return typeof value === 'string' && value.startsWith(SECRETS_MANAGER_ARN_PREFIX);
+}

package/src/lib/bootstrap-command-handler.js

@@ -24,6 +24,8 @@ import { fileURLToPath } from 'node:url';
 import BootstrapConfig from './bootstrap-config.js';
 import AwsProfileParser from './aws-profile-parser.js';
 import AssetManager from './asset-manager.js';
+import McpCommandHandler from './mcp-command-handler.js';
+import RegistryCommandHandler from './registry-command-handler.js';
 import { runPrompts } from '../prompt-adapter.js';
 
 const __filename = fileURLToPath(import.meta.url);
@@ -45,6 +47,12 @@ export default class BootstrapCommandHandler {
      * @param {object} options - Parsed CLI options
      */
     async handle(args, options) {
+        // Handle legacy --sync-schemas flag for backward compatibility
+        if (options['sync-schemas']) {
+            await this._handleSyncSchemas();
+            if (args.length === 0) return;
+        }
+
         if (args.length === 0) {
             await this._handleInteractiveSetup(options);
             return;
@@ -74,6 +82,9 @@ export default class BootstrapCommandHandler {
         case 'update':
             await this._handleUpdate(options);
             break;
+        case 'sync-schemas':
+            await this._handleSyncSchemas();
+            break;
         default:
             console.log(`Unknown bootstrap subcommand: ${subcommand}`);
             this._showHelp();
@@ -302,6 +313,9 @@ export default class BootstrapCommandHandler {
 
         // Display summary
         this._displaySummary(profileName, profileData);
+
+        // Step 6: Post-setup chain (mcp init → sync-architectures → sync-schemas)
+        await this._runPostSetupChain(options);
     }
 
     /**
@@ -927,6 +941,35 @@ export default class BootstrapCommandHandler {
         console.log(`\n  Done. ${toRemove.length} removed, ${after.length} remaining.`);
     }
 
+    /**
+     * Handle sync-schemas subcommand: download service models and verify AWS CLI.
+     */
+    async _handleSyncSchemas() {
+        console.log('\n📦 Schema Sync — Downloading AWS service models...\n');
+
+        // Verify AWS CLI is installed
+        try {
+            const version = execSync('aws --version', { encoding: 'utf8', stdio: ['pipe', 'pipe', 'pipe'] }).trim();
+            console.log(`  AWS CLI: ${version}`);
+        } catch {
+            console.log('  ⚠️  AWS CLI not found.');
+            console.log('  Install: https://docs.aws.amazon.com/cli/latest/userguide/getting-started-install.html');
+            console.log('  Continuing without AWS CLI verification...\n');
+        }
+
+        // Dynamic import to avoid circular dependencies
+        const { syncSchemas } = await import('./schema-sync.js');
+        const result = await syncSchemas();
+
+        if (result.success) {
+            console.log('\n  ✅ Schema sync complete.');
+        } else {
+            console.log('\n  ⚠️  Schema sync completed with errors (some services may be unavailable).');
+        }
+
+        console.log(`  Manifest written: lastSynced = ${result.manifest.lastSynced}\n`);
+    }
+
     /**
      * Re-deploy bootstrap infrastructure using the active profile.
      * No prompts — reads all config from the existing profile and re-applies
@@ -1016,6 +1059,74 @@ export default class BootstrapCommandHandler {
         // Save updated profile
         this.config.setProfile(name, profileConfig);
         console.log(`\n✅ Update complete for profile "${name}"`);
+
+        // Re-run post-setup chain after updating AWS resources
+        await this._runPostSetupChain(options);
+    }
+
+    /**
+     * Run the post-setup chain: mcp init → registry sync-architectures → sync-schemas.
+     * Each step is independent — failures are collected and reported at the end.
+     *
+     * @param {object} options - Parsed CLI options (checks skipPostSetup)
+     */
+    async _runPostSetupChain(options = {}) {
+        if (options['skip-post-setup']) {
+            console.log('\n⏭️  Skipping post-setup chain (--skip-post-setup)');
+            return;
+        }
+
+        console.log('\n🔗 Running post-setup configuration...\n');
+
+        const failures = [];
+
+        // 1. MCP init — register bundled MCP servers
+        console.log('📡 Registering MCP servers...');
+        try {
+            const generatorAdapter = {
+                destinationPath(...segments) {
+                    return path.resolve(process.cwd(), ...segments);
+                }
+            };
+            const mcpHandler = new McpCommandHandler(generatorAdapter);
+            await mcpHandler.handle(['init'], {});
+        } catch (error) {
+            failures.push({ step: 'mcp init', error: error.message });
+            console.log(`  ⚠️  mcp init failed: ${error.message}`);
+        }
+
+        // 2. Registry sync-architectures — populate supportedModelTypes
+        console.log('\n📋 Syncing model architecture registry...');
+        try {
+            const registryHandler = new RegistryCommandHandler();
+            await registryHandler.handle(['sync-architectures'], {});
+        } catch (error) {
+            failures.push({ step: 'registry sync-architectures', error: error.message });
+            console.log(`  ⚠️  registry sync-architectures failed: ${error.message}`);
+        }
+
+        // 3. Schema sync — download AWS service models
+        console.log('\n📐 Syncing service schemas...');
+        try {
+            await this._handleSyncSchemas();
+        } catch (error) {
+            failures.push({ step: 'sync-schemas', error: error.message });
+            console.log(`  ⚠️  sync-schemas failed: ${error.message}`);
+        }
+
+        // Report results
+        if (failures.length === 0) {
+            console.log('\n✅ Bootstrap complete — all systems operational');
+        } else {
+            console.log(`\n⚠️  Bootstrap complete with ${failures.length} warning${failures.length === 1 ? '' : 's'}:`);
+            for (const { step, error } of failures) {
+                console.log(`  • ${step}: ${error}`);
+            }
+            console.log('\n  These steps can be re-run individually:');
+            console.log('    ml-container-creator mcp init');
+            console.log('    ml-container-creator registry sync-architectures');
+            console.log('    ml-container-creator bootstrap sync-schemas');
+        }
     }
 
     /**
@@ -1204,12 +1315,20 @@ export default class BootstrapCommandHandler {
                     Effect: 'Allow',
                     Action: [
                         's3:GetObject',
+                        's3:PutObject',
+                        's3:AbortMultipartUpload',
                         's3:ListBucket'
                     ],
                     Resource: [
                         'arn:aws:s3:::ml-container-creator-*',
                         'arn:aws:s3:::ml-container-creator-*/*'
                     ]
+                },
+                {
+                    Sid: 'SNSPublish',
+                    Effect: 'Allow',
+                    Action: 'sns:Publish',
+                    Resource: 'arn:aws:sns:*:*:ml-container-creator-*'
                 }
             ]
         };
@@ -1611,6 +1730,7 @@ SETUP OPTIONS:
   --skip-s3                           Skip S3 bucket creation
   --ci                                Provision CI testing infrastructure
   --skip-ci                           Skip CI infrastructure provisioning
+  --skip-post-setup                   Skip post-setup chain (mcp init, sync-architectures, sync-schemas)
 
 STATUS OPTIONS:
   --verify                            Check each active resource against AWS APIs for drift detection

package/src/lib/cli-handler.js

@@ -280,9 +280,9 @@ TRANSFORMER MODEL EXAMPLES:
 
 REGISTRY CONTRIBUTION:
   To contribute to the catalogs:
-  - Framework Catalog: servers/base-image-picker/catalogs/model-servers.json
-  - Model Catalog: servers/model-picker/catalogs/popular-transformers.json
-  - Instance Catalog: servers/instance-recommender/catalogs/instances.json
+  - Framework Catalog: servers/lib/catalogs/model-servers.json
+  - Model Catalog: servers/lib/catalogs/popular-transformers.json
+  - Instance Catalog: servers/lib/catalogs/instances.json
   
   See docs/REGISTRY_CONTRIBUTION_GUIDE.md for detailed contribution guidelines.
 

package/src/lib/config-manager.js

@@ -29,7 +29,7 @@ import ParameterSchemaValidator from './parameter-schema-validator.js';
 
 const __configMgrFilename = fileURLToPath(import.meta.url);
 const __configMgrDir = dirname(__configMgrFilename);
-const tritonBackendsCatalogPath = resolve(__configMgrDir, '../../servers/base-image-picker/catalogs/triton-backends.json');
+const tritonBackendsCatalogPath = resolve(__configMgrDir, '../../servers/lib/catalogs/triton-backends.json');
 
 function loadTritonBackendsFromCatalog() {
     try {
@@ -300,6 +300,15 @@ export default class ConfigManager {
             finalConfig.hfToken = this._resolveHfToken(finalConfig.hfToken);
         }
 
+        // Mutual exclusion: ARN takes precedence over plaintext when both are set
+        // (CLI validation should prevent this, but enforce at config level too)
+        if (finalConfig.hfTokenArn) {
+            finalConfig.hfToken = null;
+        }
+        if (finalConfig.ngcTokenArn) {
+            finalConfig.ngcApiKey = null;
+        }
+
         // Map awsRoleArn to roleArn for templates
         if (finalConfig.awsRoleArn) {
             finalConfig.roleArn = finalConfig.awsRoleArn;
@@ -643,6 +652,28 @@ export default class ConfigManager {
                 default: null,
                 valueSpace: 'bounded'
             },
+            hfTokenArn: {
+                cliOption: 'hf-token-arn',
+                envVar: null,
+                configFile: true,
+                packageJson: false,
+                mcp: false,
+                promptable: false,
+                required: false,
+                default: null,
+                valueSpace: 'bounded'
+            },
+            ngcTokenArn: {
+                cliOption: 'ngc-token-arn',
+                envVar: null,
+                configFile: true,
+                packageJson: false,
+                mcp: false,
+                promptable: false,
+                required: false,
+                default: null,
+                valueSpace: 'bounded'
+            },
             deploymentTarget: {
                 cliOption: 'deployment-target',
                 envVar: 'ML_DEPLOYMENT_TARGET',
@@ -1675,6 +1706,18 @@ export default class ConfigManager {
             }
         }
 
+        // Validate mutual exclusion: plaintext token and ARN cannot both be set
+        if (this.config.hfToken && this.config.hfTokenArn) {
+            errors.push('Cannot specify both --hf-token and --hf-token-arn. Use one or the other.');
+        }
+        if (this.config.ngcTokenArn) {
+            // Check ngcToken from CLI options (Commander converts --ngc-token to ngcToken)
+            const ngcTokenFromCli = this.options['ngc-token'];
+            if (ngcTokenFromCli) {
+                errors.push('Cannot specify both --ngc-token and --ngc-token-arn. Use one or the other.');
+            }
+        }
+
         // Validate AWS Role ARN format if provided
         if (this.config.awsRoleArn) {
             try {
@@ -1927,6 +1970,9 @@ export default class ConfigManager {
             // For required parameters: fill auto-generatable values
             if (this.config[param] === undefined || this.config[param] === null) {
                 if (param === 'instanceType') {
+                    // If instance-sizer is configured and model is known, defer to sizer
+                    // The sizer query happens in PromptRunner after model is selected
+                    // For now, set a heuristic default that may be overridden by the sizer
                     const arch = architecture || 'http';
                     this.config[param] = arch === 'http' ? 'ml.m5.large' : 'ml.g5.xlarge';
                 } else if (param === 'modelFormat') {

package/src/lib/configuration-manager.js

@@ -619,8 +619,8 @@ export default class ConfigurationManager {
      */
     _generateSubmissionInstructions(registryType, config, configEntry) {
         const registryFile = registryType === 'framework' 
-            ? 'servers/base-image-picker/catalogs/model-servers.json'
-            : 'servers/model-picker/catalogs/popular-transformers.json';
+            ? 'servers/lib/catalogs/model-servers.json'
+            : 'servers/lib/catalogs/popular-transformers.json';
         
         const registryName = registryType === 'framework' 
             ? 'Framework_Catalog'