@aws/ml-container-creator 1.0.0 → 1.0.2

package/package.json

@@ -1,6 +1,6 @@
 {
   "name": "@aws/ml-container-creator",
-  "version": "1.0.0",
+  "version": "1.0.2",
   "description": "Build and deploy custom ML containers on AWS SageMaker with minimal configuration.",
   "main": "src/index.js",
   "bin": {

package/servers/endpoint-picker/index.js

@@ -200,22 +200,37 @@ async function fetchEndpoints(client, { limit = 10, showFull = false } = {}) {
 
             const variantName = primaryVariant.VariantName || 'AllTraffic';
             let instanceType = primaryVariant.InstanceType || null;
+            let instancePools = primaryVariant.InstancePools || null;
 
             // For IC-based endpoints, InstanceType may not be in the variant runtime response.
-            // Fall back to DescribeEndpointConfig which always has it.
-            if (!instanceType && detail.EndpointConfigName) {
+            // Fall back to DescribeEndpointConfig which has either InstanceType or InstancePools.
+            if (!instanceType && !instancePools && detail.EndpointConfigName) {
                 try {
                     const ecCmd = new _DescribeEndpointConfigCommand({ EndpointConfigName: detail.EndpointConfigName });
                     const ecDetail = await client.send(ecCmd);
                     const ecVariant = (ecDetail.ProductionVariants || [])[0];
                     if (ecVariant?.InstanceType) {
                         instanceType = ecVariant.InstanceType;
+                    } else if (ecVariant?.InstancePools && ecVariant.InstancePools.length > 0) {
+                        instancePools = ecVariant.InstancePools;
                     }
                 } catch (ecErr) {
                     log(`Warning: could not describe endpoint config for "${endpointName}": ${ecErr.message}`);
                 }
             }
-            instanceType = instanceType || 'unknown';
+
+            // Resolve instanceType display string from pools if needed
+            if (!instanceType && instancePools && instancePools.length > 0) {
+                // Sort by priority, use highest-priority (lowest number) for GPU lookup
+                const sorted = [...instancePools].sort((a, b) => (a.Priority || 99) - (b.Priority || 99));
+                instanceType = sorted[0].InstanceType || 'unknown';
+                // Build display string showing the pool: "ml.g5.12xl (pool: 3 types)"
+                if (sorted.length > 1) {
+                    instanceType = `${instanceType} (pool: ${sorted.length} types)`;
+                }
+            } else {
+                instanceType = instanceType || 'unknown';
+            }
 
             const instanceCount = primaryVariant.CurrentInstanceCount ?? primaryVariant.DesiredInstanceCount ?? 1;
             const hasInstancePools = !!(primaryVariant.InstancePools && primaryVariant.InstancePools.length > 0);
@@ -244,7 +259,12 @@ async function fetchEndpoints(client, { limit = 10, showFull = false } = {}) {
             } while (icNextToken);
 
             // Capacity estimation
-            const gpusPerInstance = getGpusForInstance(instanceType);
+            // For pool endpoints, instanceType may be "ml.g5.12xlarge (pool: 3 types)"
+            // Extract the raw type for catalog lookup
+            const instanceTypeForLookup = instanceType.includes(' (pool:')
+                ? instanceType.split(' (pool:')[0]
+                : instanceType;
+            const gpusPerInstance = getGpusForInstance(instanceTypeForLookup);
             let availableGpus;
             if (gpusPerInstance === null) {
                 availableGpus = '?';

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

@@ -52,6 +52,7 @@ export default class BootstrapCommandHandler {
     _setupS3Buckets() { return this.provisioners._setupS3Buckets(); }
     _createS3Bucket(name, tags) { return this.provisioners._createS3Bucket(name, tags); }
     _verifyCliV2() { return this.provisioners._verifyCliV2(); }
+    _provisionAiRegistryHub(profileData) { return this.provisioners.provisionAiRegistryHub(profileData); }
 
     // ── ProfileManager delegations (backward compat for tests) ──────
 
@@ -357,6 +358,9 @@ export default class BootstrapCommandHandler {
             console.log('     Tune jobs will still work but experiment tracking may not be available.');
         }
 
+        // Step 4c: AI Registry Hub
+        await this._provisionAiRegistryHub(profileData);
+
         // Step 5: CI Infrastructure setup (separate CDK stack — unchanged)
         this._displayProgress('🧪', 'CI Testing Infrastructure...');
         try {
@@ -714,6 +718,10 @@ export default class BootstrapCommandHandler {
             console.log(`  ⚠️  MLflow App setup skipped: ${error.message}`);
         }
 
+        // Ensure AI Registry hub exists
+        this._currentProfile = profileConfig.awsProfile;
+        await this._provisionAiRegistryHub(profileConfig);
+
         // Save updated profile
         this.config.setProfile(name, profileConfig);
         console.log(`\n✅ Update complete for profile "${name}"`);

package/src/lib/bootstrap-profile-manager.js

@@ -172,6 +172,23 @@ export default class BootstrapProfileManager {
             }
         }
 
+        // Check AI Registry hub status
+        if (profile.config.aiRegistryHubName) {
+            try {
+                const hubExists = this.handler._resourceExists(
+                    `sagemaker describe-hub --hub-name ${profile.config.aiRegistryHubName} --region ${profile.config.awsRegion}`,
+                    profile.config.awsProfile
+                );
+                console.log(hubExists
+                    ? `  ✅ AI Registry hub: ${profile.config.aiRegistryHubName}`
+                    : `  ⚠️  AI Registry hub: ${profile.config.aiRegistryHubName} — missing`);
+            } catch {
+                console.log(`  ⚠️  AI Registry hub: ${profile.config.aiRegistryHubName} — could not validate`);
+            }
+        } else {
+            console.log('  ℹ️  AI Registry hub: not provisioned (run bootstrap to create)');
+        }
+
         // Display deployed resources from manifest
         console.log('\n📦 Deployed Resources:');
 

package/src/lib/bootstrap-provisioners.js

@@ -405,6 +405,54 @@ export default class BootstrapProvisioners {
         }
     }
 
+    /**
+     * Provision a deterministic SageMaker AI Registry Hub.
+     * Idempotent: checks if `mlcc-registry-{accountId}` already exists before creating.
+     * Non-fatal: catches all errors and prints a warning — bootstrap continues regardless.
+     *
+     * @param {object} profileData - Profile data object (mutated in place with hub info)
+     */
+    async provisionAiRegistryHub(profileData) {
+        const hubName = `mlcc-registry-${profileData.accountId}`;
+        const region = profileData.awsRegion;
+
+        console.log('\n📦 Provisioning AI Registry hub...');
+
+        try {
+            // Check if hub already exists (idempotent)
+            const hubExists = this.handler._resourceExists(
+                `sagemaker describe-hub --hub-name ${hubName} --region ${region}`,
+                this.handler._currentProfile
+            );
+
+            if (hubExists) {
+                const hubInfo = this.handler._execAws(
+                    `sagemaker describe-hub --hub-name ${hubName} --region ${region}`,
+                    this.handler._currentProfile
+                );
+                console.log(`  ✅ AI Registry hub already provisioned: ${hubName}`);
+                profileData.aiRegistryHubName = hubName;
+                profileData.aiRegistryHubArn = hubInfo.HubArn;
+                return;
+            }
+
+            // Create new hub (always — no adopt-existing logic)
+            const tags = this._buildResourceTags();
+            const tagsFile = this.handler._formatTagsForCli(tags);
+            const createResult = this.handler._execAws(
+                `sagemaker create-hub --hub-name ${hubName} --hub-display-name "MCC AI Registry" --hub-description "Dataset, evaluator, and model versioning for ml-container-creator" --tags ${tagsFile} --region ${region}`,
+                this.handler._currentProfile
+            );
+            console.log(`  ✅ AI Registry hub "${hubName}" — created`);
+            profileData.aiRegistryHubName = hubName;
+            profileData.aiRegistryHubArn = createResult.HubArn;
+        } catch (err) {
+            const message = err.message || String(err);
+            console.log(`  ⚠️  Could not provision AI Registry hub (non-fatal): ${message}`);
+            console.log('     Dataset registration will use local JSON registry.');
+        }
+    }
+
     /**
      * Build the standard resource tag set.
      * @returns {Array<{Key: string, Value: string}>} Tag array

package/src/lib/path-prover-brain.js

@@ -16,6 +16,24 @@ const __dirname = dirname(__filename);
  * classifies failures, gates tune/adapter stages, and builds
  * Athena-compatible records with run_type='path_prove'.
  *
+ * ## Module Status (AC-1.4)
+ *
+ * ALL exported functions are FULLY FUNCTIONAL:
+ * - `identifyGaps()` — Cartesian product gap finder, prioritized by neighbor count
+ * - `findNearestSubstitution()` — Hamming distance nearest-neighbor, same-family constraint
+ * - `classifyFailure()` — regex pattern matching to 6 categories (capacity, timeout, oom, code_bug, model_incompatibility, service_limitation)
+ * - `shouldExecuteTuneStages()` — gating logic for tune/adapter stages
+ * - `hammingDistance()` — config vector comparison across CONFIG_DIMENSIONS
+ * - `buildPathProverRecord()` — Athena record construction with run_type='path_prove'
+ * - `findUnfeasibleRecord()` — checks if a config is known-unfeasible to prevent repeated attempts
+ * - `getNextPriorityConfig()` — priority queue management for v1 validation mode
+ * - `updatePriorityStatus()` — updates target status after prove attempts
+ * - `getPriorityQueueStatus()` — summary counts for priority queue
+ * - `loadPriorityTargets()` — file-based priority target loading
+ * - `resolveProveTpDegree()` — TP degree auto-resolution from instance catalog
+ *
+ * This is stabilization (tests + docs), not implementation. No new logic needed.
+ *
  * Feature: ci-benchmark-pipeline
  * Requirements: 8.1–8.12
  */
@@ -611,6 +629,45 @@ export function loadPriorityTargets(configPath) {
     }
 }
 
+// ── Optimization Space Schema (Task 3 — AC-3.5) ─────────────────────────────
+
+/**
+ * Load the optimization search space schema from config/optimization-space.json.
+ *
+ * Returns the parsed schema with dimensions, version, and description.
+ * Used by gap identification to enumerate sweepable dimensions and their
+ * allowed values for the optimization/prove sweep.
+ *
+ * @returns {object|null} Parsed schema object, or null if file not found/invalid
+ */
+export function loadOptimizationSpace() {
+    try {
+        const schemaPath = resolve(__dirname, '..', '..', 'config', 'optimization-space.json');
+        const raw = readFileSync(schemaPath, 'utf8');
+        return JSON.parse(raw);
+    } catch {
+        return null;
+    }
+}
+
+/**
+ * Get the list of sweepable dimension names from the optimization space schema.
+ *
+ * Filters dimensions by status === 'sweepable' and returns their keys.
+ * Useful for verifying sync between CONFIG_DIMENSIONS and the schema.
+ *
+ * @param {object} [schema] - Pre-loaded schema (loads from file if omitted)
+ * @returns {string[]} Array of sweepable dimension names
+ */
+export function getSweepableDimensions(schema = null) {
+    const data = schema || loadOptimizationSpace();
+    if (!data || !data.dimensions) return [];
+
+    return Object.keys(data.dimensions).filter(
+        key => data.dimensions[key].status === 'sweepable'
+    );
+}
+
 // ── TP Degree Auto-Resolution at Prove-Time (Task 6.5) ──────────────────────
 
 /**

package/src/lib/prove-pipeline-executor.js

@@ -8,6 +8,25 @@
  * Handles stage-specific logic including idempotency checks, status tracking,
  * and fail-fast behavior.
  *
+ * ## Module Status (AC-1.4)
+ *
+ * FUNCTIONAL stages:
+ * - `executeStageStep()` — fully wired with idempotency via `.mlcc/staged-assets.json`
+ * - `isAlreadyStaged()` — checks staged assets existence and validity
+ * - `getStagingState()` — resolves current staging state from filesystem + step results
+ * - `isValidLifecycleStage()` — validates individual stage names
+ * - `validateStagesArray()` — validates arrays of stage names
+ * - `formatStagingStatus()` — formats staging state for display
+ * - `buildTargetStatus()` — builds status summary for a prove target
+ *
+ * INTENTIONALLY INCOMPLETE (post-v1 scope):
+ * - Other lifecycle stage executors (build, push, deploy, test, tune, adapter,
+ *   test-adapter, benchmark, register, clean) are NOT implemented.
+ * - Only the `stage` step has execution logic. Other stages are recognized in
+ *   validation but have no executor function.
+ * - This is not "broken" — these were never finished before the laptop was bricked.
+ *   They are explicitly post-v1 scope.
+ *
  * Feature: s3-model-loading
  * Requirements: 5.1, 5.2, 5.3, 5.4, 5.5
  */
@@ -40,6 +59,22 @@ export const VALID_LIFECYCLE_STAGES = [
     'clean'
 ];
 
+// TODO(post-v1): Implement executor functions for lifecycle stages beyond 'stage'.
+// The following stages are recognized for validation purposes but have no execution logic:
+//   - generate: Should invoke `mcc generate` to produce project scaffolding
+//   - build: Should run `do/build` to build the Docker container
+//   - push: Should run `do/push` to push container to ECR
+//   - deploy: Should run `do/deploy` to create SageMaker endpoint
+//   - test: Should run `do/test` to invoke endpoint and verify correctness
+//   - tune: Should run `do/tune` for fine-tuning jobs (gated by shouldExecuteTuneStages)
+//   - adapter: Should run `do/adapter` for LoRA adapter serving
+//   - test-adapter: Should test adapter endpoints after deployment
+//   - benchmark: Should run `do/benchmark` for performance measurement
+//   - register: Should register proven config in Athena/DynamoDB
+//   - clean: Should tear down deployed resources
+// These were never finished before the original developer's laptop was bricked.
+// They are explicitly post-v1 scope, not "broken" code.
+
 /**
  * Possible staging states for status output.
  */

package/templates/do/.benchmark_writer.py

@@ -1,5 +1,3 @@
-#!/usr/bin/env python3
-# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
 # SPDX-License-Identifier: Apache-2.0
 
 """Benchmark Writer — Converts do/benchmark output to enriched Parquet for Athena.
@@ -340,7 +338,7 @@ def _extract_base_image_version(base_image):
     return ''
 
 
-def enrich_records(config, results, run_timestamp=None):
+def enrich_records(config, results, run_timestamp=None, instance_catalog=None):
     """Build enriched records from config context and benchmark results.
 
     Each metrics entry becomes one enriched record with all Athena columns populated.
@@ -349,6 +347,7 @@ def enrich_records(config, results, run_timestamp=None):
         config: dict with config context fields (project_name, model_name, etc.)
         results: dict with benchmark results (job_name, metrics array)
         run_timestamp: Optional datetime for run_timestamp. Defaults to now UTC.
+        instance_catalog: Optional pre-loaded instance catalog dict. If None, loaded from disk.
 
     Returns:
         list of enriched record dicts (one per concurrency level).
@@ -364,10 +363,21 @@ def enrich_records(config, results, run_timestamp=None):
 
     # Derived fields
     model_family = derive_model_family(model_name)
+    instance_family = derive_instance_family(instance_type)
+
+    # Resolve instance metadata from catalog (AC-2.8)
+    hw_meta = resolve_instance_metadata(instance_type, instance_catalog)
+    gpu_count = hw_meta['gpu_count']
+    gpu_type = hw_meta['gpu_type']
+    gpu_memory_gb = hw_meta['gpu_memory_gb']
 
     # Optional context fields
     deployment_target = config.get('deployment_target', 'realtime-inference')
-    tensor_parallel_degree = config.get('tensor_parallel_degree', 1)
+    try:
+        tensor_parallel_degree = int(config.get('tensor_parallel_degree', 1))
+    except (ValueError, TypeError):
+        tensor_parallel_degree = 1
+
     quantization = config.get('quantization', 'none')
     enable_lora = config.get('enable_lora', False)
     base_image = config.get('base_image', '')
@@ -377,6 +387,11 @@ def enrich_records(config, results, run_timestamp=None):
     ci_run_id = config.get('ci_run_id', '')
     account_id = config.get('account_id', '')
 
+    # Configuration dimensions (nullable)
+    max_model_len_raw = config.get('max_model_len')
+    max_model_len = int(max_model_len_raw) if max_model_len_raw not in (None, '', 0) else None
+    kv_cache_dtype = config.get('kv_cache_dtype') or None
+
 
     # Get metrics from results
     metrics = results.get('metrics', []) if isinstance(results, dict) else []
@@ -447,6 +462,13 @@ def enrich_records(config, results, run_timestamp=None):
             'deployment_target': deployment_target,
             'quantization': quantization,
             'tensor_parallel_degree': tensor_parallel_degree,
+            'instance_family': instance_family,
+            'gpu_count': gpu_count,
+            'gpu_type': gpu_type,
+            'gpu_memory_gb': gpu_memory_gb,
+            'max_model_len': max_model_len,
+            'enable_lora': enable_lora,
+            'kv_cache_dtype': kv_cache_dtype,
             'serving_config': json.dumps(serving_config_dict),
             'workload': config.get('workload', 'manual'),
             'concurrency': concurrency,
@@ -481,6 +503,7 @@ def enrich_records(config, results, run_timestamp=None):
             'output_sequence_length_avg': scalar(metric.get('output_sequence_length', metric.get('output_sequence_length_avg', 0.0))),
             'input_sequence_length_avg': scalar(metric.get('input_sequence_length', metric.get('input_sequence_length_avg', 0.0))),
             'error_rate': error_rate,
+            'cost_per_1m_tokens': cost,
             'benchmark_duration_sec': metric.get('benchmark_duration_sec', duration_seconds),
             'run_type': run_type,
             'benchmark_job_name': results.get('job_name', '') if isinstance(results, dict) else '',
@@ -792,6 +815,54 @@ def register_partition(bucket, model, instance, target,
 # ── Parquet Serialization ─────────────────────────────────────────────────────
 
 
+def load_instance_catalog():
+    """Load the instance catalog from servers/lib/catalogs/instances.json.
+
+    Resolves the path relative to the project root (two levels up from templates/do/).
+    Returns the 'catalog' dict mapping instance_type → metadata, or empty dict on failure.
+
+    Returns:
+        dict mapping instance type strings to their metadata dicts.
+    """
+    # Resolve relative to this file: templates/do/.benchmark_writer.py → project root
+    this_dir = os.path.dirname(os.path.abspath(__file__))
+    # Navigate up from templates/do/ to project root
+    project_root = os.path.normpath(os.path.join(this_dir, '..', '..'))
+    catalog_path = os.path.join(project_root, 'servers', 'lib', 'catalogs', 'instances.json')
+
+    try:
+        with open(catalog_path, 'r') as f:
+            data = json.load(f)
+        return data.get('catalog', {})
+    except (FileNotFoundError, json.JSONDecodeError, IOError):
+        return {}
+
+
+def resolve_instance_metadata(instance_type, instance_catalog=None):
+    """Resolve GPU metadata from the instance catalog for a given instance_type.
+
+    Args:
+        instance_type: SageMaker instance type (e.g., 'ml.g5.xlarge').
+        instance_catalog: Optional pre-loaded catalog dict. If None, loads from disk.
+
+    Returns:
+        dict with keys: gpu_count (int|None), gpu_type (str|None), gpu_memory_gb (float|None).
+        All values are None if instance_type is not found in catalog.
+    """
+    if instance_catalog is None:
+        instance_catalog = load_instance_catalog()
+
+    entry = instance_catalog.get(instance_type)
+    if entry is None:
+        return {'gpu_count': None, 'gpu_type': None, 'gpu_memory_gb': None}
+
+    return {
+        'gpu_count': entry.get('gpus'),
+        'gpu_type': entry.get('gpuType'),
+        'gpu_memory_gb': entry.get('gpuMemoryGb'),
+    }
+
+
 def get_parquet_schema():
     """Return the pyarrow schema matching the Athena DDL for benchmark_results.
 
@@ -814,6 +885,17 @@ def get_parquet_schema():
         pa.field("quantization", pa.string()),
         pa.field("tensor_parallel_degree", pa.int32()),
 
+        # Hardware metadata (resolved from instance catalog at write time)
+        pa.field("instance_family", pa.string()),
+        pa.field("gpu_count", pa.int32()),
+        pa.field("gpu_type", pa.string()),
+        pa.field("gpu_memory_gb", pa.float64()),
+
+        # Configuration dimensions (top-level for Athena queryability)
+        pa.field("max_model_len", pa.int32()),
+        pa.field("enable_lora", pa.bool_()),
+        pa.field("kv_cache_dtype", pa.string()),
+
         # Full serving config (extensible JSON blob)
         pa.field("serving_config", pa.string()),
 
@@ -852,6 +934,7 @@ def get_parquet_schema():
         pa.field("output_sequence_length_avg", pa.float64()),
         pa.field("input_sequence_length_avg", pa.float64()),
         pa.field("error_rate", pa.float64()),
+        pa.field("cost_per_1m_tokens", pa.float64()),
         pa.field("benchmark_duration_sec", pa.float64()),
 
         # Run Metadata
@@ -1182,6 +1265,9 @@ def cmd_write(args):
     if args.adapter_name:
         input_data['adapter_name'] = args.adapter_name
 
+    if getattr(args, 'instance_type', None):
+        input_data['instance_type'] = args.instance_type
+
     # ── Validate before any S3 interaction ────────────────────────────────
     errors = validate_benchmark_input(input_data)
     if errors:
@@ -1391,6 +1477,8 @@ def _load_config_file(config_path):
                     'MODEL_NAME': 'model_name',
                     'HF_MODEL_ID': 'hf_model_id',
                     'INSTANCE_TYPE': 'instance_type',
+                    'INSTANCE_POOLS': 'instance_pools',
+                    'BENCHMARK_INSTANCE_TYPE': 'benchmark_instance_type',
                     'DEPLOYMENT_CONFIG': 'deployment_config',
                     'DEPLOYMENT_TARGET': 'deployment_target',
                     'AWS_REGION': 'region',
@@ -1429,6 +1517,24 @@ def _load_config_file(config_path):
         parts = context['model_name'].rstrip('/').split('/')
         context['model_name'] = parts[-1] if parts else context['model_name']
 
+    # Resolve instance_type precedence:
+    #   BENCHMARK_INSTANCE_TYPE (live-resolved, persisted by do/benchmark) > INSTANCE_TYPE > INSTANCE_POOLS fallback
+    if context.get('benchmark_instance_type'):
+        context['instance_type'] = context.pop('benchmark_instance_type')
+    # Fall back to INSTANCE_POOLS when neither is set.
+    # Heterogeneous pool configs may not have a standalone INSTANCE_TYPE value
+    # but always define INSTANCE_POOLS as a JSON array with Priority fields.
+    if not context.get('instance_type') and context.get('instance_pools'):
+        try:
+            pools = json.loads(context['instance_pools'])
+            if pools:
+                # Pick the highest-priority (lowest number) instance
+                best = min(pools, key=lambda p: p.get('Priority', 999))
+                context['instance_type'] = best.get('InstanceType', '')
+        except (json.JSONDecodeError, TypeError, KeyError):
+            pass
+    context.pop('instance_pools', None)  # Don't leak raw JSON into record
+
     # Also scan IC config files (do/ic/*.conf) for IC_ENV_* serving params
     # These override do/config values for serving-specific settings
     try:
@@ -1505,6 +1611,10 @@ def main():
         help='LoRA adapter name (differentiates adapter benchmarks from base model in Athena)'
     )
 
+    write_parser.add_argument(
+        '--instance-type', dest='instance_type', default=None,
+        help='Override instance type (use when actual provisioned instance differs from config, e.g. heterogeneous pools)'
+    )
     write_parser.add_argument(
         '--dry-run', dest='dry_run', action='store_true',
         help='Output enriched records as JSON without writing to S3'