@aws/ml-container-creator 0.2.0

package/src/lib/cli-handler.js

@@ -0,0 +1,529 @@
+// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+// SPDX-License-Identifier: Apache-2.0
+
+/**
+ * CLI Handler - Handles special CLI arguments and commands
+ * 
+ * Supports behavioral CLI arguments like:
+ * - ml-container-creator configure
+ * - ml-container-creator generate-empty-config
+ * - ml-container-creator --help
+ */
+
+import fs from 'fs';
+// import path from 'path'; // Unused import removed
+import McpCommandHandler from './mcp-command-handler.js';
+
+export default class CliHandler {
+    constructor(generator) {
+        this.generator = generator;
+    }
+
+    /**
+     * Handles special CLI arguments
+     * @returns {boolean} True if special command was handled, false to continue normal flow
+     */
+    async handleCliArguments() {
+        const args = this.generator.args;
+        const options = this.generator.options;
+        
+        // Check for help flags in options first
+        if (options.help || options.h) {
+            this._showHelp();
+            return true;
+        }
+        
+        if (args.length === 0) {
+            return false; // Normal generation flow
+        }
+
+        const command = args[0].toLowerCase();
+
+        switch (command) {
+        case 'configure':
+            await this._handleConfigure();
+            return true;
+                
+        case 'generate-empty-config':
+            await this._handleGenerateEmptyConfig();
+            return true;
+                
+        case 'mcp': {
+            const mcpHandler = new McpCommandHandler(this.generator);
+            await mcpHandler.handle(args.slice(1), options);
+            return true;
+        }
+
+        case 'registry': {
+            const { default: RegistryCommandHandler } = await import('./registry-command-handler.js');
+            const registryHandler = new RegistryCommandHandler();
+            await registryHandler.handle(args.slice(1), options);
+            return true;
+        }
+
+        case 'bootstrap': {
+            const { default: BootstrapCommandHandler } = await import('./bootstrap-command-handler.js');
+            const bootstrapHandler = new BootstrapCommandHandler();
+            await bootstrapHandler.handle(args.slice(1), options);
+            return true;
+        }
+
+        case 'help':
+        case '--help':
+        case '-h':
+            this._showHelp();
+            return true;
+                
+        default:
+            // If it's not a special command, treat as project name
+            return false;
+        }
+    }
+
+    /**
+     * Interactive configuration setup
+     * @private
+     */
+    async _handleConfigure() {
+        console.log('\n🔧 ML Container Creator Configuration');
+        console.log('\nThis will help you set up configuration files for your project.\n');
+
+        const answers = await this.generator.prompt([
+            {
+                type: 'list',
+                name: 'configType',
+                message: 'What type of configuration would you like to create?',
+                choices: [
+                    { name: 'Package.json section (recommended for projects)', value: 'package' },
+                    { name: 'Custom config file (config/mcp.json)', value: 'custom' },
+                    { name: 'Show environment variable examples', value: 'env' },
+                    { name: 'Show CLI option examples', value: 'cli' }
+                ]
+            }
+        ]);
+
+        switch (answers.configType) {
+        case 'package':
+            await this._createPackageJsonConfig();
+            break;
+        case 'custom':
+            await this._createCustomConfig();
+            break;
+        case 'env':
+            this._showEnvironmentExamples();
+            break;
+        case 'cli':
+            this._showCliExamples();
+            break;
+        }
+    }
+
+    /**
+     * Generate empty configuration file
+     * @private
+     */
+    async _handleGenerateEmptyConfig() {
+        const answers = await this.generator.prompt([
+            {
+                type: 'list',
+                name: 'configType',
+                message: 'Which configuration file format?',
+                choices: [
+                    { name: 'Custom config (config/mcp.json)', value: 'custom' }
+                ]
+            }
+        ]);
+
+        const emptyConfig = this._getEmptyConfig();
+
+        if (answers.configType === 'custom') {
+            const configPath = 'config/mcp.json';
+            if (!fs.existsSync('config')) {
+                fs.mkdirSync('config', { recursive: true });
+            }
+            fs.writeFileSync(configPath, JSON.stringify(emptyConfig, null, 2));
+            console.log(`\n✅ Created empty configuration file: ${configPath}`);
+        }
+
+        console.log('\n📝 Edit the configuration file and run the generator again.');
+    }
+
+    /**
+     * Show comprehensive help
+     * @private
+     */
+    _showHelp() {
+        console.log(`
+🚀 ML Container Creator Generator
+
+USAGE:
+  ml-container-creator [command] [project-name] [options]
+
+COMMANDS:
+  configure              Interactive configuration setup
+  generate-empty-config  Generate empty configuration file
+  bootstrap              Set up shared AWS infrastructure (IAM role, ECR repo, S3 buckets)
+  help                   Show this help message
+
+EXAMPLES:
+  ml-container-creator                           # Interactive mode
+  ml-container-creator my-project               # With project name
+  ml-container-creator --skip-prompts           # Skip prompts, use config
+  ml-container-creator --framework=sklearn      # With CLI options
+  ml-container-creator --config=prod.json       # With config file
+
+CLI OPTIONS:
+  --skip-prompts              Skip interactive prompts
+  --config=<file>             Load configuration from file
+  --project-name=<name>       Project name
+  --project-dir=<dir>         Output directory path
+  --framework=<framework>     ML framework (sklearn|xgboost|tensorflow|transformers)
+  --model-name=<name>         HuggingFace model name (for transformers framework)
+  --model-server=<server>     Model server (flask|fastapi|vllm|sglang|tensorrt-llm|lmi|djl)
+  --model-format=<format>     Model format (depends on framework)
+  --include-sample            Include sample model code
+  --include-testing           Include test suite
+  --test-types=<types>        Comma-separated test types (local-model-cli,local-model-server,hosted-model-endpoint)
+  --build-target=<target>     Build target (codebuild)
+  --codebuild-compute-type=<type> CodeBuild compute type (BUILD_GENERAL1_SMALL|BUILD_GENERAL1_MEDIUM|BUILD_GENERAL1_LARGE)
+  --codebuild-project-name=<name> CodeBuild project name
+  --instance-type=<type>      SageMaker instance type (e.g., ml.m5.large, ml.g5.xlarge)
+  --region=<region>           AWS region
+  --role-arn=<arn>            AWS IAM role ARN for SageMaker execution
+  --deployment-target=<target> Deployment target (managed-inference|hyperpod-eks)
+  --hyperpod-cluster=<name> HyperPod EKS cluster name
+  --hyperpod-namespace=<ns> Kubernetes namespace for HyperPod (default: default)
+  --hyperpod-replicas=<n>   Number of replicas for HyperPod (default: 1)
+  --fsx-volume-handle=<id>  FSx for Lustre volume handle for HyperPod storage
+  --hf-token=<token>          HuggingFace token (or "$HF_TOKEN" for env var)
+
+VALIDATION OPTIONS:
+  --validate-env-vars         Enable environment variable validation (default: true)
+  --validate-with-docker      Enable Docker introspection validation (default: false, opt-in)
+  --offline                   Disable HuggingFace API lookups (default: false)
+
+MCP OPTIONS:
+  --smart                     Enable Bedrock-powered smart mode on all MCP servers
+  --discover                  Enable live registry lookups (e.g. Docker Hub) on MCP servers that support it
+
+REGISTRY SYSTEM:
+  The generator includes built-in registries for frameworks, models, and instance types:
+  
+  Framework Registry:
+    - Provides base images, environment variables, and accelerator requirements
+    - Supports version-specific configurations (e.g., vllm 0.3.0, 0.4.0)
+    - Includes optimization profiles (low-latency, high-throughput)
+    - Validates accelerator compatibility with instance types
+  
+  Model Registry:
+    - Stores model-specific configurations and chat templates
+    - Supports pattern matching (e.g., "meta-llama/Llama-2-*")
+    - Provides recommended instance types per model
+    - Includes validation levels (tested, community-validated, experimental)
+    - Offers model profiles for different deployment scenarios
+  
+  Instance Accelerator Mapping:
+    - Maps AWS instance types to accelerator capabilities
+    - Validates framework-instance compatibility
+    - Provides hardware specifications (GPU type, memory, vCPUs)
+  
+  Registry Features:
+    - Automatic environment variable injection from registries
+    - Chat template detection from HuggingFace Hub or Model Registry
+    - Validation of framework-accelerator-instance compatibility
+    - Inline documentation comments in generated Dockerfiles
+
+ENVIRONMENT VARIABLES:
+  ML_INSTANCE_TYPE           Instance type
+  ML_BUILD_TARGET            Build target
+  ML_CODEBUILD_COMPUTE_TYPE  CodeBuild compute type
+  AWS_REGION                 AWS region
+  AWS_ROLE                   AWS IAM role ARN
+  ML_CONTAINER_CREATOR_CONFIG Config file path
+  VALIDATE_ENV_VARS          Enable env var validation (true|false)
+  VALIDATE_WITH_DOCKER       Enable Docker validation (true|false)
+  OFFLINE_MODE               Disable HuggingFace API (true|false)
+  HF_TOKEN                   HuggingFace authentication token
+
+CONFIGURATION FILES (in precedence order):
+  1. CLI options (highest precedence)
+  2. Environment variables
+  3. --config file
+  4. config/mcp.json
+  5. package.json "ml-container-creator" section
+  6. Registry defaults (framework/model registries)
+  7. Generator defaults
+  8. Interactive prompts (lowest precedence)
+
+TRANSFORMER MODEL EXAMPLES:
+  # vLLM with Llama-2 7B
+  ml-container-creator --framework=transformers \\
+    --model-name=meta-llama/Llama-2-7b-chat-hf \\
+    --model-server=vllm \\
+    --instance-type=ml.g5.xlarge \\
+    --hf-token=$HF_TOKEN \\
+    --skip-prompts
+
+  # TensorRT-LLM with custom model
+  ml-container-creator --framework=transformers \\
+    --model-name=openai/gpt-oss-20b \\
+    --model-server=tensorrt-llm \\
+    --instance-type=ml.g6.12xlarge \\
+    --skip-prompts
+
+  # SGLang with Mistral
+  ml-container-creator --framework=transformers \\
+    --model-name=mistralai/Mistral-7B-Instruct-v0.2 \\
+    --model-server=sglang \\
+    --instance-type=ml.g5.xlarge \\
+    --skip-prompts
+
+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
+  
+  See docs/REGISTRY_CONTRIBUTION_GUIDE.md for detailed contribution guidelines.
+
+For more information, visit: https://github.com/awslabs/ml-container-creator
+`);
+    }
+
+    /**
+     * Create package.json configuration section
+     * @private
+     */
+    async _createPackageJsonConfig() {
+        const packageJsonPath = 'package.json';
+        let packageJson = {};
+
+        if (fs.existsSync(packageJsonPath)) {
+            packageJson = JSON.parse(fs.readFileSync(packageJsonPath, 'utf8'));
+        }
+
+        const config = await this._promptForConfig();
+        packageJson['ml-container-creator'] = config;
+
+        fs.writeFileSync(packageJsonPath, JSON.stringify(packageJson, null, 2));
+        console.log(`\n✅ Added configuration to ${packageJsonPath}`);
+        console.log('\n📝 You can now run: ml-container-creator --skip-prompts');
+    }
+
+    /**
+     * Create custom configuration file
+     * @private
+     */
+    async _createCustomConfig() {
+        const config = await this._promptForConfig();
+        const configPath = 'config/mcp.json';
+        
+        if (!fs.existsSync('config')) {
+            fs.mkdirSync('config', { recursive: true });
+        }
+        fs.writeFileSync(configPath, JSON.stringify(config, null, 2));
+        console.log(`\n✅ Created configuration file: ${configPath}`);
+        console.log('\n📝 You can now run: ml-container-creator --skip-prompts');
+    }
+
+
+
+    /**
+     * Prompt for configuration values
+     * @private
+     */
+    async _promptForConfig() {
+        return await this.generator.prompt([
+            {
+                type: 'input',
+                name: 'projectName',
+                message: 'Project name:',
+                default: 'my-ml-project'
+            },
+            {
+                type: 'list',
+                name: 'framework',
+                message: 'ML framework:',
+                choices: ['sklearn', 'xgboost', 'tensorflow', 'transformers']
+            },
+            {
+                type: 'list',
+                name: 'modelServer',
+                message: 'Model server:',
+                choices: (answers) => {
+                    if (answers.framework === 'transformers') {
+                        return ['vllm', 'sglang'];
+                    }
+                    return ['flask', 'fastapi'];
+                }
+            },
+            {
+                type: 'list',
+                name: 'modelFormat',
+                message: 'Model format:',
+                choices: (answers) => {
+                    const formatMap = {
+                        'xgboost': ['json', 'model', 'ubj'],
+                        'sklearn': ['pkl', 'joblib'],
+                        'tensorflow': ['keras', 'h5', 'SavedModel']
+                    };
+                    return formatMap[answers.framework] || [];
+                },
+                when: answers => answers.framework !== 'transformers'
+            },
+            {
+                type: 'confirm',
+                name: 'includeSampleModel',
+                message: 'Include sample model:',
+                default: false,
+                when: answers => answers.framework !== 'transformers'
+            },
+            {
+                type: 'confirm',
+                name: 'includeTesting',
+                message: 'Include testing:',
+                default: true
+            },
+            {
+                type: 'list',
+                name: 'instanceType',
+                message: 'Instance type:',
+                choices: (answers) => {
+                    if (answers.framework === 'transformers') {
+                        return ['ml.g5.xlarge'];
+                    }
+                    return ['ml.m5.large', 'ml.g5.xlarge'];
+                },
+                default: 'ml.m5.large'
+            },
+            {
+                type: 'list',
+                name: 'awsRegion',
+                message: 'AWS region:',
+                choices: ['us-east-1'],
+                default: 'us-east-1'
+            }
+        ]);
+    }
+
+    /**
+     * Show environment variable examples
+     * @private
+     */
+    _showEnvironmentExamples() {
+        console.log(`
+🌍 Environment Variable Examples:
+
+# Basic configuration
+# Note: Core parameters (framework, model-server, etc.) are only supported via CLI options and config files
+
+# Optional settings
+export ML_INSTANCE_TYPE="ml.m5.large"
+export AWS_REGION="us-east-1"
+export AWS_ROLE="arn:aws:iam::123456789012:role/SageMakerRole"
+export ML_CONTAINER_CREATOR_CONFIG="./my-config.json"
+
+# Then run with CLI options for core parameters:
+ml-container-creator --framework=sklearn --model-server=flask --model-format=pkl --skip-prompts
+`);
+    }
+
+    /**
+     * Show CLI option examples
+     * @private
+     */
+    _showCliExamples() {
+        console.log(`
+💻 CLI Option Examples:
+
+# Basic sklearn project
+ml-container-creator my-sklearn-project \\
+  --framework=sklearn \\
+  --model-server=flask \\
+  --model-format=pkl \\
+  --skip-prompts
+
+# Transformers project with vLLM
+ml-container-creator my-llm-project \\
+  --framework=transformers \\
+  --model-name=meta-llama/Llama-2-7b-chat-hf \\
+  --model-server=vllm \\
+  --instance-type=ml.g5.xlarge \\
+  --hf-token=$HF_TOKEN \\
+  --region=us-east-1 \\
+  --skip-prompts
+
+# TensorRT-LLM with custom instance type
+ml-container-creator my-tensorrt-project \\
+  --framework=transformers \\
+  --model-name=openai/gpt-oss-20b \\
+  --model-server=tensorrt-llm \\
+  --instance-type=ml.g6.12xlarge \\
+  --skip-prompts
+
+# SGLang with Mistral model
+ml-container-creator my-sglang-project \\
+  --framework=transformers \\
+  --model-name=mistralai/Mistral-7B-Instruct-v0.2 \\
+  --model-server=sglang \\
+  --instance-type=ml.g5.xlarge \\
+  --skip-prompts
+
+# XGBoost with FastAPI and custom role
+ml-container-creator my-xgb-project \\
+  --framework=xgboost \\
+  --model-server=fastapi \\
+  --model-format=json \\
+  --include-sample \\
+  --include-testing \\
+  --role-arn=arn:aws:iam::123456789012:role/SageMakerRole \\
+  --skip-prompts
+
+# CodeBuild deployment target
+ml-container-creator my-codebuild-project \\
+  --framework=sklearn \\
+  --model-server=flask \\
+  --model-format=pkl \\
+  --build-target=codebuild \\
+  --codebuild-compute-type=BUILD_GENERAL1_MEDIUM \\
+  --codebuild-project-name=my-build-project \\
+  --skip-prompts
+
+# Using configuration file
+ml-container-creator --config=production.json --skip-prompts
+
+# With validation options
+ml-container-creator \\
+  --framework=transformers \\
+  --model-name=meta-llama/Llama-2-7b-chat-hf \\
+  --model-server=vllm \\
+  --instance-type=ml.g5.xlarge \\
+  --validate-env-vars=true \\
+  --validate-with-docker=false \\
+  --offline=false \\
+  --skip-prompts
+`);
+    }
+
+    /**
+     * Get empty configuration template
+     * @private
+     */
+    _getEmptyConfig() {
+        return {
+            'projectName': 'my-ml-project',
+            'framework': 'sklearn',
+            'modelServer': 'flask',
+            'modelFormat': 'pkl',
+            'includeSampleModel': false,
+            'includeTesting': true,
+            'testTypes': ['local-model-cli', 'local-model-server', 'hosted-model-endpoint'],
+            'buildTarget': 'codebuild',
+            'codebuildComputeType': 'BUILD_GENERAL1_MEDIUM',
+            'codebuildProjectName': 'my-build-project',
+            'instanceType': 'ml.m5.large',
+            'awsRegion': 'us-east-1'
+        };
+    }
+}
+