Adding New Features¶

Guide for contributors who want to add new frameworks, model servers, or other features to ML Container Creator.

Table of Contents¶

Adding a New ML Framework (predictor and serve-engine plugin paths)
Adding a New Test Type
Adding a New Deployment Target
Adding a New Secret Type
Testing Your Changes

Adding a New ML Framework¶

Two paths — pick the right one. MLCC serves models two ways, and each has its own plugin system and its own authoring guide. Do NOT add a framework by editing a branched model_handler.py or hardcoding format maps in config-validator.js/model-prompts.js — that duplication was deliberately removed. Both paths are descriptor-driven:

You want to add… Path Authoring guide

A classical-ML framework (sklearn/xgboost-like) served over HTTP via Flask/FastAPI Predictor-framework plugin (templates/code/predictors.d/<framework>/) Predictor-Framework Plugins

An LLM/transformer serving engine (vLLM/SGLang/llama.cpp-like) Serve-engine plugin (templates/code/serve.d/<engine>/) Serve-Engine Plugin Authoring

Classical ML (HTTP predictor) — the short version¶

A predictor framework is a descriptor + handler plugin. Adding PyTorch is a two-file drop plus the deployment-config registration:

Descriptor — templates/code/predictors.d/pytorch/manifest.json (schema: predictors.d/manifest.schema.json):

{
  "framework": "pytorch",
  "display_name": "PyTorch",
  "model_formats": ["pt", "pth", "torchscript"],
  "default_model_format": "pt",
  "pip_dependencies": ["torch==2.4.0"],
  "base_image_catalog": "python-slim",
  "handler": "handler.py",
  "test_payload": "{\"instances\": [[1.0, 2.0, 3.0, 4.0]]}"
}

Handler — templates/code/predictors.d/pytorch/handler.py: the ModelHandler class (load_model/is_loaded/preprocess/predict/ postprocess). This is the imperative code the manifest references; it is EJS-rendered into the generated project as code/model_handler.py (so you may use <%= modelFormat %> in its file-glob, like the shipped handlers).
Everything else is derived. src/lib/predictor-manifest-reader.js discovers the framework; config-validator.js (accepted formats, engines, default), model-prompts.js (engine choice, format choices), config-manager.js (format→engine inference, default), and requirements.txt (pip deps) all read the descriptor. You do not edit per-framework literals in those files.
Verify — npx mocha test/unit/predictor-framework-conformance.test.js (the drift guard), npm run lint, and an end-to-end generate (--deployment-config=http-flask --model-format=pt …). See the Predictor-Framework Plugins guide for the full checklist and the two-axis (framework × web-server) model.

LLM / transformer serving engine — the short version¶

An LLM engine is a serve-engine plugin (templates/code/serve.d/<engine>/ manifest.json + <engine>.ejs). The base image already contains the server, so the contract is pure manifest data. Adding one touches five registration surfaces (enum, resolver, prefix map, template-manager allow-list, plus the manifest). See the Serve-Engine Plugin Authoring Guide for the full end-to-end checklist (it is maintained there, not duplicated here).

Adding a web server (Flask/FastAPI alternative)¶

The web-server axis (flask vs fastapi) is orthogonal to the framework axis and is NOT part of the predictor plugin. It lives in templates/code/serve.py and templates/code/start_server.py as a modelServer EJS branch (serve.py is framework-agnostic — it only references ModelHandler). Adding a third HTTP web server means adding a branch there and to the http-<server> deployment-config enum, not a new predictor plugin.

Adding a New Test Type¶

Let's add integration tests.

Step 1: Update SUPPORTED_OPTIONS¶

SUPPORTED_OPTIONS = {
    testTypes: [
        'local-model-cli',
        'local-model-server',
        'hosted-model-endpoint',
        'integration-tests'
    ],
    // ...
}

Step 2: Add to Test Type Prompt¶

{
    type: 'checkbox',
    name: 'testTypes',
    message: 'Test type?',
    choices: (answers) => {
        const baseTests = [
            'local-model-cli',
            'local-model-server',
            'hosted-model-endpoint',
            'integration-tests'
        ];
        // Filter based on framework
        return baseTests;
    }
}

Step 3: Create Test Template¶

# templates/test/test_integration.sh

#!/bin/bash
set -e

echo "Running integration tests for <%= projectName %>"

# Test 1: Build container
echo "Test 1: Building container..."
docker build -t <%= projectName %>-test .

# Test 2: Start container
echo "Test 2: Starting container..."
CONTAINER_ID=$(docker run -d -p 8080:8080 <%= projectName %>-test)
sleep 10

# Test 3: Health check
echo "Test 3: Testing health endpoint..."
curl -f http://localhost:8080/ping || exit 1

# Test 4: Inference
echo "Test 4: Testing inference..."
curl -X POST http://localhost:8080/invocations \
  -H "Content-Type: application/json" \
  -d '{"instances": [[1.0, 2.0, 3.0]]}' || exit 1

# Test 5: Load test
echo "Test 5: Running load test..."
for i in {1..100}; do
  curl -X POST http://localhost:8080/invocations \
    -H "Content-Type: application/json" \
    -d '{"instances": [[1.0, 2.0, 3.0]]}' &
done
wait

# Cleanup
echo "Cleaning up..."
docker stop $CONTAINER_ID
docker rm $CONTAINER_ID

echo "All integration tests passed!"

Step 4: Update Conditional Logic¶

// In writing() method
if (!this.answers.testTypes.includes('integration-tests')) {
    ignorePatterns.push('**/test/test_integration.sh');
}

Adding a New Deployment Target¶

Let's add support for AWS Lambda.

Step 1: Update SUPPORTED_OPTIONS¶

SUPPORTED_OPTIONS = {
    deployment: ['sagemaker', 'lambda'],
    // ...
}

Step 2: Add to Deployment Prompt¶

{
    type: 'list',
    name: 'deployTarget',
    message: 'Deployment target?',
    choices: ['sagemaker', 'lambda'],
    default: 'sagemaker'
}

Step 3: Create Lambda-Specific Templates¶

# templates/code/lambda_handler.py

import json
import base64
from model_handler import ModelHandler

# Initialize model once (outside handler)
model_handler = ModelHandler('/opt/ml/model')

def lambda_handler(event, context):
    """AWS Lambda handler function."""
    try:
        # Parse input
        body = json.loads(event['body'])
        instances = body.get('instances', [])

        # Run inference
        predictions = model_handler.predict(instances)

        # Return response
        return {
            'statusCode': 200,
            'body': json.dumps({
                'predictions': predictions
            })
        }
    except Exception as e:
        return {
            'statusCode': 500,
            'body': json.dumps({
                'error': str(e)
            })
        }

Step 4: Create Lambda Deployment Script¶

# templates/deploy/deploy_lambda.sh

#!/bin/bash
set -e

PROJECT_NAME="<%= projectName %>"
REGION="<%= awsRegion %>"
ROLE_ARN=$1

if [ -z "$ROLE_ARN" ]; then
    echo "Usage: ./deploy_lambda.sh <lambda-execution-role-arn>"
    exit 1
fi

echo "Deploying ${PROJECT_NAME} to AWS Lambda..."

# Create deployment package
echo "Creating deployment package..."
cd code
zip -r ../deployment.zip . -x "*.pyc" "__pycache__/*"
cd ..

# Create Lambda function
echo "Creating Lambda function..."
aws lambda create-function \
    --function-name ${PROJECT_NAME} \
    --runtime python3.9 \
    --role ${ROLE_ARN} \
    --handler lambda_handler.lambda_handler \
    --zip-file fileb://deployment.zip \
    --timeout 30 \
    --memory-size 512 \
    --region ${REGION}

echo "Lambda function deployed successfully!"
echo "Test with: aws lambda invoke --function-name ${PROJECT_NAME} --payload '{...}' output.json"

Step 5: Update Conditional Logic¶

// In writing() method
if (this.answers.deployTarget === 'lambda') {
    ignorePatterns.push('**/deploy/build_and_push.sh');
    ignorePatterns.push('**/deploy/deploy.sh');
    ignorePatterns.push('**/nginx-predictors.conf');
    ignorePatterns.push('**/nginx-tensorrt.conf');
} else {
    ignorePatterns.push('**/deploy/deploy_lambda.sh');
    ignorePatterns.push('**/code/lambda_handler.py');
}

Adding a New Secret Type¶

The secrets system uses a registry-driven architecture. Adding a new secret type requires only adding an entry to the Secret_Classification registry — the CLI, prompt flow, and do-script templates derive behavior from this registry automatically.

Let's walk through adding support for a private PyPI token used during build-time pip install.

Understanding the Registry¶

The Secret_Classification registry lives at src/lib/secret-classification.js. Each entry defines all metadata needed for end-to-end integration:

// src/lib/secret-classification.js
export const SECRET_CLASSIFICATIONS = Object.freeze([
    {
        identifier: 'hf-token',           // Unique key — used in naming convention and CLI
        displayName: 'HuggingFace Token',  // Human-readable label for prompts
        stages: ['build-time', 'runtime'], // When the secret is needed
        purpose: 'Gated model download from HuggingFace Hub',
        cliFlag: 'hf-token-arn',           // CLI flag for ARN input (--hf-token-arn)
        cliFlagPlaintext: 'hf-token',      // Existing CLI flag for plaintext (--hf-token)
        envVar: 'HF_TOKEN',               // Environment variable for plaintext value
        envVarArn: 'HF_TOKEN_ARN',        // Environment variable for ARN reference
        promptLabel: 'HuggingFace token'   // Label shown in interactive prompts
    },
    // ... other entries
]);

Field	Type	Description
`identifier`	`string`	Unique key (e.g., `pypi-token`) — used in naming convention and type selection
`displayName`	`string`	Human-readable label shown in prompts and output
`stages`	`string[]`	When the secret is needed: `['build-time']`, `['runtime']`, or `['build-time', 'runtime']`
`purpose`	`string`	Description of what the secret is used for
`cliFlag`	`string`	CLI flag name for ARN input (e.g., `pypi-token-arn`)
`cliFlagPlaintext`	`string`	CLI flag name for plaintext input (e.g., `pypi-token`)
`envVar`	`string`	Environment variable name for plaintext (e.g., `PIP_INDEX_TOKEN`)
`envVarArn`	`string`	Environment variable name for ARN (e.g., `PIP_INDEX_TOKEN_ARN`)
`promptLabel`	`string`	Label shown in interactive prompts

Step 1: Add Registry Entry¶

Edit src/lib/secret-classification.js and add a new entry to the SECRET_CLASSIFICATIONS array:

export const SECRET_CLASSIFICATIONS = Object.freeze([
    // ... existing entries ...
    {
        identifier: 'pypi-token',
        displayName: 'Private PyPI Token',
        stages: ['build-time'],
        purpose: 'Authenticating with private PyPI registry during pip install',
        cliFlag: 'pypi-token-arn',
        cliFlagPlaintext: 'pypi-token',
        envVar: 'PIP_INDEX_TOKEN',
        envVarArn: 'PIP_INDEX_TOKEN_ARN',
        promptLabel: 'Private PyPI token'
    }
]);

This single addition automatically enables: - The type appears in secrets create interactive prompts - The prompt flow queries for managed secrets of this type - The secrets list command includes secrets of this type

Step 2: Register CLI Flags¶

Edit bin/cli.js to add the new ARN and plaintext flags on the root command:

.addOption(new Option('--pypi-token-arn <arn>', 'Private PyPI token ARN from Secrets Manager'))
.addOption(new Option('--pypi-token <value>', 'Private PyPI token (plaintext)'))

Add mutual exclusion validation alongside the existing checks:

if (options.pypiToken && options.pypiTokenArn) {
    console.error('❌ Cannot specify both --pypi-token and --pypi-token-arn');
    process.exit(1);
}

Step 3: Update Do-Script Templates¶

`templates/do/config`¶

Add the ARN vs plaintext export logic:

<% if (pypiTokenArn) { %>
# Private PyPI token — resolved from Secrets Manager at build-time
export PIP_INDEX_TOKEN_ARN="<%= pypiTokenArn %>"
<% } else if (pypiToken) { %>
export PIP_INDEX_TOKEN="<%= pypiToken %>"
<% } %>

`templates/do/build`¶

Add a resolution block before docker build (alongside existing secret resolution blocks):

if [ -n "${PIP_INDEX_TOKEN_ARN:-}" ]; then
    echo "🔐 Resolving Private PyPI token from Secrets Manager..."
    PIP_INDEX_TOKEN=$(aws secretsmanager get-secret-value \
        --secret-id "${PIP_INDEX_TOKEN_ARN}" \
        --query SecretString --output text) || {
        echo "❌ Failed to resolve Private PyPI token from Secrets Manager"
        exit 3
    }
    export PIP_INDEX_TOKEN
fi

Since pypi-token is build-time only, no changes are needed in templates/do/serve.

Step 4: IAM Considerations¶

The existing IAM policy in config/bootstrap-stack.json already covers new secret types because it scopes to the mlcc/* naming prefix:

{
    "Sid": "SecretsManagerRead",
    "Effect": "Allow",
    "Action": [
        "secretsmanager:GetSecretValue",
        "secretsmanager:DescribeSecret"
    ],
    "Resource": "arn:aws:secretsmanager:*:*:secret:mlcc/*",
    "Condition": {
        "StringEquals": {
            "aws:ResourceTag/mlcc:managed-by": "ml-container-creator"
        }
    }
}

No IAM changes are needed as long as your new secret follows the mlcc/<type>/<label> naming convention. The policy grants access to any secret under the mlcc/ prefix that carries the mlcc:managed-by tag.

Tagging Schema¶

Every secret created by the CLI automatically receives these tags:

Tag Key	Value	Purpose
`mlcc:managed-by`	`ml-container-creator`	Identifies mlcc-managed secrets for IAM scoping
`mlcc:created-by`	`secrets`	Identifies the creation source
`mlcc:secret-type`	`<identifier>` (e.g., `pypi-token`)	Links to the Secret_Classification entry

These tags are applied automatically by the SecretsCommandHandler and cannot be overridden by the user. If a user provides conflicting mlcc: tags via --json, the system values take precedence and a warning is displayed.

Naming Convention¶

All managed secrets follow the pattern:

mlcc/<secret-type>/<user-provided-label>

Examples: - mlcc/hf-token/production - mlcc/ngc-token/team-shared - mlcc/pypi-token/ci-pipeline

This naming convention maps directly to the IAM policy resource pattern arn:aws:secretsmanager:*:*:secret:mlcc/*, ensuring that any new secret type is automatically covered without IAM policy changes.

Step 5: Add Tests¶

Add a test verifying the new entry has all required fields in test/property/secret-classification-completeness.property.test.js (the existing property test validates all entries automatically).

Add unit tests for the new do-script resolution logic in test/unit/secrets-template-output.test.js.

Validation Checklist¶

After adding a new secret type, verify:

The new type appears in secrets create interactive type selection prompt
secrets create --type <new-type> --name test --secret-value xxx creates a secret with the correct name (mlcc/<type>/test)
The prompt flow lists managed secrets of the new type during project generation
Do-scripts resolve the secret at the correct stage (build-time, runtime, or both)
The IAM policy covers the new naming path (automatic if using mlcc/ prefix)
Mutual exclusion validation rejects both --<type> and --<type>-arn flags together
The do/config template exports the correct _ARN variable when an ARN is configured
Existing plaintext flows remain unchanged when no ARN is provided
All property tests pass (npm test)

Files to Update Summary¶

File	Change
`src/lib/secret-classification.js`	Add new entry to `SECRET_CLASSIFICATIONS` array
`bin/cli.js`	Add `--<type>-arn` and `--<type>` flags, add mutual exclusion check
`templates/do/config`	Add ARN/plaintext export conditional
`templates/do/build`	Add resolution block (if stages include `build-time`)
`templates/do/serve`	Add resolution block (if stages include `runtime`)
`test/unit/secrets-template-output.test.js`	Add tests for new template output

Testing Your Changes¶

Unit Tests¶

# Run all tests
npm test

# Run specific test
npm test -- test/pytorch-generator.js

# Run with coverage
npm test -- --coverage

Manual Testing¶

# Link CLI globally
npm link

# Test generation
ml-container-creator

# Test all combinations
# - New framework + flask
# - New framework + fastapi
# - With/without sample model
# - With/without tests

Integration Testing¶

# Generate project
ml-container-creator

# Build container
cd generated-project
docker build -t test-model .

# Test locally
docker run -p 8080:8080 test-model

# Test inference
curl -X POST http://localhost:8080/invocations \
  -H "Content-Type: application/json" \
  -d '{"instances": [[1.0, 2.0, 3.0]]}'

Validation Checklist¶

Submitting Your Changes¶

1. Create Feature Branch¶

git checkout -b feature/add-pytorch-support

2. Make Changes¶

Follow the steps above for your feature.

3. Test Thoroughly¶

npm test
npm run lint

4. Commit Changes¶

git add .
git commit -m "feat: add PyTorch framework support

- Add pytorch to supported frameworks
- Add pt, pth, torchscript model formats
- Create PyTorch-specific model handler
- Add tests for PyTorch generation
- Update documentation"

5. Push and Create PR¶

git push origin feature/add-pytorch-support

Then create a Pull Request on GitHub with: - Clear description of changes - Screenshots/examples if applicable - Link to related issues - Test results

Best Practices¶

Code Organization¶

Keep generator logic in src/app.js and src/lib/
Keep templates in templates/ (plugins under templates/code/serve.d/ and templates/code/predictors.d/)
Keep tests in test/
Keep documentation in docs/

Template Design¶

Use EJS for all dynamic content
Keep templates simple and readable
Add comments explaining conditional logic
Test all template paths

Documentation¶

Update all relevant documentation
Add examples for new features
Include troubleshooting tips
Update steering files for AI context

Testing¶

Write unit tests for generator logic
Test all configuration combinations
Test generated projects end-to-end
Include edge cases

Backward Compatibility¶

Don't break existing features
Provide migration guide if needed
Deprecate features gracefully
Version breaking changes appropriately

Getting Help¶

Review existing code for patterns
Check Architecture
Ask in GitHub Discussions
Open a draft PR for early feedback

Resources¶

Adding a New `do/` Script¶

Every do/ script in MLCC has a machine-readable contract that controls runtime behavior and advisory agent suggestions. See do-script-contract.md for the full reference.

Step 1: Define the contract¶

Answer four questions: 1. What does my script operate on? → type (model-centric, deployment-centric, hybrid) 2. What must exist before it runs? → guard (none, artifact-ready, model-staged, deployment-active, training-infra) 3. When in the lifecycle does it fit? → lifecycle 4. Which targets does it apply to? → targets (all, or comma-separated list)

Step 2: Create the script from template¶

#!/usr/bin/env bash
# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
# SPDX-License-Identifier: Apache-2.0
#
# do/myscript — One-line description.
#
# @mlcc-script
# type: <model-centric | deployment-centric | hybrid>
# guard: <none | artifact-ready | model-staged | deployment-active | training-infra>
# lifecycle: <value from schema>
# targets: <all | target1, target2>

SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
source "${SCRIPT_DIR}/lib/script-contract.sh"   # ← auto-enforces declared guard
source "${SCRIPT_DIR}/config"
source "${SCRIPT_DIR}/lib/profile.sh"

# Your script logic here

Step 3: Add the script to `templates/do/README.md`¶

Step 4: Add contract tests¶

Add tests in test/unit/do-script-contracts.test.js: - Guard enforcement test (exits code 3 when guard not met) - Happy path test (succeeds when guard is satisfied)

You want to add…	Path	Authoring guide
A classical-ML framework (sklearn/xgboost-like) served over HTTP via Flask/FastAPI	Predictor-framework plugin (`templates/code/predictors.d/<framework>/`)	Predictor-Framework Plugins
An LLM/transformer serving engine (vLLM/SGLang/llama.cpp-like)	Serve-engine plugin (`templates/code/serve.d/<engine>/`)	Serve-Engine Plugin Authoring

Adding New Features¶

Table of Contents¶

Adding a New ML Framework¶

Classical ML (HTTP predictor) — the short version¶

LLM / transformer serving engine — the short version¶

Adding a web server (Flask/FastAPI alternative)¶

Adding a New Test Type¶

Step 1: Update SUPPORTED_OPTIONS¶

Step 2: Add to Test Type Prompt¶

Step 3: Create Test Template¶

Step 4: Update Conditional Logic¶

Adding a New Deployment Target¶

Step 1: Update SUPPORTED_OPTIONS¶

Step 2: Add to Deployment Prompt¶

Step 3: Create Lambda-Specific Templates¶

Step 4: Create Lambda Deployment Script¶

Step 5: Update Conditional Logic¶

Adding a New Secret Type¶

Understanding the Registry¶

Step 1: Add Registry Entry¶

Step 2: Register CLI Flags¶

Step 3: Update Do-Script Templates¶

templates/do/config¶

templates/do/build¶

Step 4: IAM Considerations¶

Tagging Schema¶

Naming Convention¶

Step 5: Add Tests¶

Validation Checklist¶

Files to Update Summary¶

Testing Your Changes¶

Unit Tests¶

Manual Testing¶

Integration Testing¶

Validation Checklist¶

Submitting Your Changes¶

1. Create Feature Branch¶

2. Make Changes¶

3. Test Thoroughly¶

4. Commit Changes¶

5. Push and Create PR¶

Best Practices¶

Code Organization¶

Template Design¶

Documentation¶

Testing¶

Backward Compatibility¶

Getting Help¶

Resources¶

Adding a New do/ Script¶

Step 1: Define the contract¶

Step 2: Create the script from template¶

Step 3: Add the script to templates/do/README.md¶

Step 4: Add contract tests¶

`templates/do/config`¶

`templates/do/build`¶

Adding a New `do/` Script¶

Step 3: Add the script to `templates/do/README.md`¶