@aws/ml-container-creator 0.10.0 → 0.12.1

package/servers/base-image-picker/index.js

@@ -18,18 +18,18 @@
  *   Returns: { values, choices, metadata }
  */
 
-import { McpServer } from '@modelcontextprotocol/sdk/server/mcp.js'
-import { StdioServerTransport } from '@modelcontextprotocol/sdk/server/stdio.js'
-import { z } from 'zod'
-import { readFileSync } from 'node:fs'
-import { fileURLToPath } from 'node:url'
-import { resolve, dirname } from 'node:path'
-import { DynamicResolver as DynamicResolverBase } from '../lib/dynamic-resolver.js'
+import { McpServer } from '@modelcontextprotocol/sdk/server/mcp.js';
+import { StdioServerTransport } from '@modelcontextprotocol/sdk/server/stdio.js';
+import { z } from 'zod';
+import { readFileSync } from 'node:fs';
+import { fileURLToPath } from 'node:url';
+import { resolve, dirname } from 'node:path';
+import { DynamicResolver as DynamicResolverBase } from '../lib/dynamic-resolver.js';
 
 // ── Catalog loader ───────────────────────────────────────────────────────────
 
-const __filename = fileURLToPath(import.meta.url)
-const __dirname = dirname(__filename)
+const __filename = fileURLToPath(import.meta.url);
+const __dirname = dirname(__filename);
 
 /**
  * Load and parse a JSON catalog file relative to the server directory.
@@ -39,17 +39,17 @@ const __dirname = dirname(__filename)
  * @returns {any} Parsed JSON content
  */
 function loadCatalog(relativePath) {
-    const fullPath = resolve(__dirname, relativePath)
-    let raw
+    const fullPath = resolve(__dirname, relativePath);
+    let raw;
     try {
-        raw = readFileSync(fullPath, 'utf8')
+        raw = readFileSync(fullPath, 'utf8');
     } catch (err) {
-        throw new Error(`Catalog file not found: ${fullPath}`)
+        throw new Error(`Catalog file not found: ${fullPath}`);
     }
     try {
-        return JSON.parse(raw)
+        return JSON.parse(raw);
     } catch (err) {
-        throw new Error(`Failed to parse catalog ${fullPath}: ${err.message}`)
+        throw new Error(`Failed to parse catalog ${fullPath}: ${err.message}`);
     }
 }
 
@@ -74,8 +74,8 @@ class ImageResolver extends DynamicResolverBase {
      * @param {string} [options.searchCriteria] - Optional filter string
      * @returns {Promise<{images: object[], defaultImage: string|null}>}
      */
-    async fetchImages(framework, options = {}) {
-        throw new Error('fetchImages() must be implemented by subclass')
+    async fetchImages(framework, _options = {}) {
+        throw new Error('fetchImages() must be implemented by subclass');
     }
 
     /**
@@ -83,33 +83,33 @@ class ImageResolver extends DynamicResolverBase {
      * @returns {string[]}
      */
     supportedFrameworks() {
-        throw new Error('supportedFrameworks() must be implemented by subclass')
+        throw new Error('supportedFrameworks() must be implemented by subclass');
     }
 
     // ── DynamicResolver interface bridge ─────────────────────────────────
 
     async fetch(key, options = {}) {
-        return this.fetchImages(key, options)
+        return this.fetchImages(key, options);
     }
 
     supportedKeys() {
-        return this.supportedFrameworks()
+        return this.supportedFrameworks();
     }
 }
 
 // ── Load catalogs from JSON files ─────────────────────────────────────────────
 
-let TRANSFORMER_IMAGE_CATALOG
-let PYTHON_SLIM_CATALOG
-let TRITON_IMAGE_CATALOG
+let TRANSFORMER_IMAGE_CATALOG;
+let PYTHON_SLIM_CATALOG;
+let TRITON_IMAGE_CATALOG;
 
 try {
-    TRANSFORMER_IMAGE_CATALOG = loadCatalog('../lib/catalogs/model-servers.json')
-    PYTHON_SLIM_CATALOG = loadCatalog('../lib/catalogs/python-slim.json')
-    TRITON_IMAGE_CATALOG = loadCatalog('../lib/catalogs/triton.json')
+    TRANSFORMER_IMAGE_CATALOG = loadCatalog('../lib/catalogs/model-servers.json');
+    PYTHON_SLIM_CATALOG = loadCatalog('../lib/catalogs/python-slim.json');
+    TRITON_IMAGE_CATALOG = loadCatalog('../lib/catalogs/triton.json');
 } catch (err) {
-    process.stderr.write(`[base-image-picker] Fatal: ${err.message}\n`)
-    process.exit(1)
+    process.stderr.write(`[base-image-picker] Fatal: ${err.message}\n`);
+    process.exit(1);
 }
 
 // ── DynamicResolver ──────────────────────────────────────────────────────────
@@ -123,39 +123,39 @@ try {
  */
 class DynamicResolver extends ImageResolver {
     constructor(options = {}) {
-        super()
-        this._timeout = options.timeout || 5000
+        super();
+        this._timeout = options.timeout || 5000;
         // Registry API endpoints per framework
         this._registryEndpoints = {
             'vllm': 'https://hub.docker.com/v2/repositories/vllm/vllm-openai/tags',
             'sglang': 'https://hub.docker.com/v2/repositories/lmsysorg/sglang/tags',
             'djl': 'https://hub.docker.com/v2/repositories/deepjavalibrary/djl-serving/tags'
             // tensorrt-llm and lmi require auth — not supported in V1 discover
-        }
+        };
     }
 
     async fetchImages(framework, options = {}) {
-        const { limit = 5 } = options
-        const endpoint = this._registryEndpoints[framework]
+        const { limit = 5 } = options;
+        const endpoint = this._registryEndpoints[framework];
         if (!endpoint) {
-            return { images: [], defaultImage: null }
+            return { images: [], defaultImage: null };
         }
 
-        const controller = new AbortController()
-        const timer = setTimeout(() => controller.abort(), this._timeout)
+        const controller = new AbortController();
+        const timer = setTimeout(() => controller.abort(), this._timeout);
 
         try {
             const response = await fetch(
                 `${endpoint}?page_size=${limit}&ordering=-last_updated`,
                 { signal: controller.signal }
-            )
-            clearTimeout(timer)
+            );
+            clearTimeout(timer);
 
             if (!response.ok) {
-                throw new Error(`Registry API returned ${response.status}`)
+                throw new Error(`Registry API returned ${response.status}`);
             }
 
-            const data = await response.json()
+            const data = await response.json();
             const images = (data.results || []).map(tag => ({
                 image: `${this._repoForFramework(framework)}:${tag.name}`,
                 tag: tag.name,
@@ -164,22 +164,22 @@ class DynamicResolver extends ImageResolver {
                 labels: {},
                 registry: 'dockerhub',
                 repository: this._repoForFramework(framework)
-            }))
+            }));
 
             return {
                 images: images.slice(0, limit),
                 defaultImage: images[0]?.image || null
-            }
+            };
         } catch (err) {
-            log(`[discover] Registry API failed for ${framework}: ${err.message}`)
-            return { images: [], defaultImage: null }
+            log(`[discover] Registry API failed for ${framework}: ${err.message}`);
+            return { images: [], defaultImage: null };
         } finally {
-            clearTimeout(timer)
+            clearTimeout(timer);
         }
     }
 
     supportedFrameworks() {
-        return Object.keys(this._registryEndpoints)
+        return Object.keys(this._registryEndpoints);
     }
 
     _repoForFramework(framework) {
@@ -187,8 +187,8 @@ class DynamicResolver extends ImageResolver {
             'vllm': 'vllm/vllm-openai',
             'sglang': 'lmsysorg/sglang',
             'djl': 'deepjavalibrary/djl-serving'
-        }
-        return map[framework] || framework
+        };
+        return map[framework] || framework;
     }
 }
 
@@ -209,14 +209,14 @@ class DynamicResolver extends ImageResolver {
  * @returns {object[]} Merged, deduplicated image list
  */
 function mergeStaticAndDynamic(staticImages, dynamicImages, limit) {
-    const staticIds = new Set(staticImages.map(e => e.image))
-    const netNew = dynamicImages.filter(e => !staticIds.has(e.image))
+    const staticIds = new Set(staticImages.map(e => e.image));
+    const netNew = dynamicImages.filter(e => !staticIds.has(e.image));
 
     // Sort net-new by created desc
-    netNew.sort((a, b) => new Date(b.created) - new Date(a.created))
+    netNew.sort((a, b) => new Date(b.created) - new Date(a.created));
 
-    const merged = [...staticImages, ...netNew]
-    return limit != null ? merged.slice(0, limit) : merged
+    const merged = [...staticImages, ...netNew];
+    return limit !== null && limit !== undefined ? merged.slice(0, limit) : merged;
 }
 
 // ── StaticCatalogResolver ────────────────────────────────────────────────────
@@ -229,29 +229,29 @@ function mergeStaticAndDynamic(staticImages, dynamicImages, limit) {
  */
 class StaticCatalogResolver extends ImageResolver {
     constructor(transformerCatalog, pythonSlimCatalog, tritonCatalog) {
-        super()
-        this._transformerCatalog = transformerCatalog
-        this._pythonSlimCatalog = pythonSlimCatalog
-        this._tritonCatalog = tritonCatalog || []
+        super();
+        this._transformerCatalog = transformerCatalog;
+        this._pythonSlimCatalog = pythonSlimCatalog;
+        this._tritonCatalog = tritonCatalog || [];
     }
 
     async fetchImages(framework, options = {}) {
-        const { limit = 5, searchCriteria } = options
+        const { limit = 5, searchCriteria } = options;
 
         if (framework === 'python-slim') {
-            return this._resolvePythonSlim(limit, searchCriteria)
+            return this._resolvePythonSlim(limit, searchCriteria);
         }
 
         if (framework === 'triton') {
-            return this._resolveTriton(limit, searchCriteria)
+            return this._resolveTriton(limit, searchCriteria);
         }
 
-        const catalog = this._transformerCatalog[framework] || []
-        const sliced = catalog.slice(0, limit)
+        const catalog = this._transformerCatalog[framework] || [];
+        const sliced = catalog.slice(0, limit);
         return {
             images: sliced,
             defaultImage: sliced[0]?.image || null
-        }
+        };
     }
 
     supportedFrameworks() {
@@ -259,46 +259,46 @@ class StaticCatalogResolver extends ImageResolver {
             ...Object.keys(this._transformerCatalog),
             'python-slim',
             'triton'
-        ]
+        ];
     }
 
     _resolvePythonSlim(limit, searchCriteria) {
-        let catalog = [...this._pythonSlimCatalog]
+        let catalog = [...this._pythonSlimCatalog];
 
         if (searchCriteria && searchCriteria.trim()) {
-            const query = searchCriteria.trim().toLowerCase()
+            const query = searchCriteria.trim().toLowerCase();
             catalog = catalog.filter(entry =>
                 entry.tag.toLowerCase().includes(query) ||
                 entry.image.toLowerCase().includes(query) ||
                 (entry.labels.python_version && entry.labels.python_version.toLowerCase().includes(query))
-            )
+            );
         }
 
-        const sliced = catalog.slice(0, limit)
+        const sliced = catalog.slice(0, limit);
         return {
             images: sliced,
             defaultImage: sliced[0]?.image || null
-        }
+        };
     }
 
     _resolveTriton(limit, searchCriteria) {
-        let catalog = [...this._tritonCatalog]
+        let catalog = [...this._tritonCatalog];
 
         if (searchCriteria && searchCriteria.trim()) {
-            const query = searchCriteria.trim().toLowerCase()
+            const query = searchCriteria.trim().toLowerCase();
             catalog = catalog.filter(entry =>
                 entry.tag.toLowerCase().includes(query) ||
                 entry.image.toLowerCase().includes(query) ||
                 (entry.labels.triton_version && entry.labels.triton_version.toLowerCase().includes(query)) ||
                 (entry.labels.cuda_version && entry.labels.cuda_version.toLowerCase().includes(query))
-            )
+            );
         }
 
-        const sliced = catalog.slice(0, limit)
+        const sliced = catalog.slice(0, limit);
         return {
             images: sliced,
             defaultImage: sliced[0]?.image || null
-        }
+        };
     }
 }
 
@@ -312,8 +312,8 @@ class StaticCatalogResolver extends ImageResolver {
  */
 class ResolverRegistry {
     constructor() {
-        this._resolvers = new Map()
-        this._defaultResolver = null
+        this._resolvers = new Map();
+        this._defaultResolver = null;
     }
 
     /**
@@ -322,7 +322,7 @@ class ResolverRegistry {
      */
     register(resolver) {
         for (const framework of resolver.supportedFrameworks()) {
-            this._resolvers.set(framework, resolver)
+            this._resolvers.set(framework, resolver);
         }
     }
 
@@ -331,7 +331,7 @@ class ResolverRegistry {
      * @param {ImageResolver} resolver
      */
     setDefault(resolver) {
-        this._defaultResolver = resolver
+        this._defaultResolver = resolver;
     }
 
     /**
@@ -340,16 +340,16 @@ class ResolverRegistry {
      * @returns {ImageResolver|null}
      */
     getResolver(framework) {
-        return this._resolvers.get(framework) || this._defaultResolver
+        return this._resolvers.get(framework) || this._defaultResolver;
     }
 }
 
 // ── V1 wiring ────────────────────────────────────────────────────────────────
 
-const staticResolver = new StaticCatalogResolver(TRANSFORMER_IMAGE_CATALOG, PYTHON_SLIM_CATALOG, TRITON_IMAGE_CATALOG)
-const registry = new ResolverRegistry()
-registry.register(staticResolver)
-registry.setDefault(staticResolver)
+const staticResolver = new StaticCatalogResolver(TRANSFORMER_IMAGE_CATALOG, PYTHON_SLIM_CATALOG, TRITON_IMAGE_CATALOG);
+const registry = new ResolverRegistry();
+registry.register(staticResolver);
+registry.setDefault(staticResolver);
 
 // ── Discover mode ────────────────────────────────────────────────────────────
 
@@ -358,13 +358,13 @@ registry.setDefault(staticResolver)
  * --discover flag or MCP_DISCOVER=true activates discover mode.
  */
 const discoverMode = process.argv.includes('--discover') ||
-    process.env.MCP_DISCOVER === 'true'
+    process.env.MCP_DISCOVER === 'true';
 
-let dynamicResolver = null
+let dynamicResolver = null;
 
 if (discoverMode) {
-    dynamicResolver = new DynamicResolver()
-    registry.register(dynamicResolver)
+    dynamicResolver = new DynamicResolver();
+    registry.register(dynamicResolver);
 }
 
 // ── Routing logic ────────────────────────────────────────────────────────────
@@ -375,52 +375,52 @@ if (discoverMode) {
  * When discover mode is active, merges static and dynamic results.
  */
 async function resolveBaseImage(context, limit) {
-    const { framework, modelServer, searchCriteria, architecture } = context
+    const { framework, modelServer, searchCriteria, architecture } = context;
 
     // Determine which framework identifier to resolve
-    let resolverKey
+    let resolverKey;
     if (architecture === 'triton') {
-        resolverKey = 'triton'
+        resolverKey = 'triton';
     } else if (architecture === 'diffusors' && modelServer) {
-        resolverKey = modelServer
+        resolverKey = modelServer;
     } else if (framework === 'transformers' && modelServer) {
-        resolverKey = modelServer
+        resolverKey = modelServer;
     } else {
-        resolverKey = 'python-slim'
+        resolverKey = 'python-slim';
     }
 
-    const resolver = registry.getResolver(resolverKey)
+    const resolver = registry.getResolver(resolverKey);
     if (!resolver) {
-        return { values: { baseImage: null }, choices: { baseImage: [] }, metadata: { baseImage: [] } }
+        return { values: { baseImage: null }, choices: { baseImage: [] }, metadata: { baseImage: [] } };
     }
 
-    let resultImages
+    let resultImages;
 
     if (discoverMode && dynamicResolver && dynamicResolver.supportedFrameworks().includes(resolverKey)) {
         // Fetch both static and dynamic results, then merge
-        const staticResult = await staticResolver.fetchImages(resolverKey, { limit, searchCriteria })
-        const dynamicResult = await dynamicResolver.fetchImages(resolverKey, { limit: 5 })
+        const staticResult = await staticResolver.fetchImages(resolverKey, { limit, searchCriteria });
+        const dynamicResult = await dynamicResolver.fetchImages(resolverKey, { limit: 5 });
 
-        resultImages = mergeStaticAndDynamic(staticResult.images, dynamicResult.images, limit)
+        resultImages = mergeStaticAndDynamic(staticResult.images, dynamicResult.images, limit);
     } else {
         // Static-only path (no network calls)
-        const result = await resolver.fetchImages(resolverKey, { limit, searchCriteria })
-        resultImages = result.images
+        const result = await resolver.fetchImages(resolverKey, { limit, searchCriteria });
+        resultImages = result.images;
     }
 
-    const images = resultImages.map(e => e.image)
+    const images = resultImages.map(e => e.image);
     return {
         values: { baseImage: images[0] || null },
         choices: { baseImage: images },
         metadata: { baseImage: resultImages }
-    }
+    };
 }
 
 /**
  * Log to stderr so it doesn't interfere with MCP stdio protocol on stdout.
  */
 function log(message) {
-    process.stderr.write(`[base-image-picker] ${message}\n`)
+    process.stderr.write(`[base-image-picker] ${message}\n`);
 }
 
 // ── MCP Server ───────────────────────────────────────────────────────────────
@@ -428,7 +428,7 @@ function log(message) {
 const server = new McpServer({
     name: 'base-image-picker',
     version: '1.0.0'
-})
+});
 
 server.tool(
     'get_base_images',
@@ -439,15 +439,15 @@ server.tool(
         context: z.record(z.string(), z.any()).optional().describe('Current configuration context (framework, modelServer, searchCriteria)')
     },
     async ({ parameters, limit, context }) => {
-        const values = {}
-        const choices = {}
-        const metadata = {}
+        const values = {};
+        const choices = {};
+        const metadata = {};
 
         if (parameters.includes('baseImage')) {
-            const result = await resolveBaseImage(context || {}, limit)
-            Object.assign(values, result.values)
-            Object.assign(choices, result.choices)
-            Object.assign(metadata, result.metadata)
+            const result = await resolveBaseImage(context || {}, limit);
+            Object.assign(values, result.values);
+            Object.assign(choices, result.choices);
+            Object.assign(metadata, result.metadata);
         }
 
         return {
@@ -455,9 +455,9 @@ server.tool(
                 type: 'text',
                 text: JSON.stringify({ values, choices, metadata })
             }]
-        }
+        };
     }
-)
+);
 
 // ── Exports for testing ──────────────────────────────────────────────────────
 
@@ -476,20 +476,20 @@ export {
     staticResolver,
     dynamicResolver,
     discoverMode
-}
+};
 
-export { DynamicResolverBase as DynamicResolverBase }
+export { DynamicResolverBase as DynamicResolverBase };
 
 // ── Main guard ───────────────────────────────────────────────────────────────
 
-const isMain = process.argv[1] && resolve(process.argv[1]) === __filename
+const isMain = process.argv[1] && resolve(process.argv[1]) === __filename;
 
 if (isMain) {
     if (discoverMode) {
-        log('Discover mode — serving curated catalogs + live registry lookups')
+        log('Discover mode — serving curated catalogs + live registry lookups');
     } else {
-        log('Static mode — serving curated base image catalogs')
+        log('Static mode — serving curated base image catalogs');
     }
-    const transport = new StdioServerTransport()
-    await server.connect(transport)
+    const transport = new StdioServerTransport();
+    await server.connect(transport);
 }