@aws/ml-container-creator 0.2.0

package/templates/triton/Dockerfile

@@ -0,0 +1,128 @@
+# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+# SPDX-License-Identifier: Apache-2.0
+
+# NVIDIA Triton Inference Server Dockerfile
+# Backend: <%= backend %>
+
+<% if (comments && comments.acceleratorInfo) { %>
+<%= comments.acceleratorInfo %>
+<% } %>
+
+<% if (comments && comments.validationInfo) { %>
+<%= comments.validationInfo %>
+<% } %>
+
+# Triton Inference Server base image from NVIDIA NGC
+# Public image - no NGC authentication required
+ARG BASE_IMAGE=<%= baseImage || 'nvcr.io/nvidia/tritonserver:24.08-py3' %>
+FROM ${BASE_IMAGE}
+
+# Set a docker label to name this project, postpended with the build time
+LABEL project.name="<%= projectName %>-<%= buildTimestamp %>" \
+      project.base-name="<%= projectName %>" \
+      project.build-time="<%= buildTimestamp %>"
+
+# Set a docker label to advertise multi-model support on the container
+LABEL com.amazonaws.sagemaker.capabilities.multi-models=true
+# Set a docker label to enable container to use SAGEMAKER_BIND_TO_PORT environment variable if present
+LABEL com.amazonaws.sagemaker.capabilities.accept-bind-to-port=true
+
+# Set working directory
+WORKDIR /opt/ml
+
+<% if (backend === 'vllm' || backend === 'tensorrtllm') { %>
+# HuggingFace model configuration for LLM backends
+ENV HF_MODEL_ID="<%= modelName %>"
+<% if (hfToken) { %>
+# Set HuggingFace authentication token for gated models
+ENV HF_TOKEN="<%= hfToken %>"
+<% } %>
+<% } %>
+
+<% if (backend === 'python') { %>
+# Install Python backend dependencies
+COPY triton/requirements.txt /tmp/triton_requirements.txt
+RUN pip install --no-cache-dir -r /tmp/triton_requirements.txt && \
+    rm /tmp/triton_requirements.txt
+<% } %>
+
+# Set up model repository directory structure
+# Triton expects models at: /opt/ml/model/model_repository/<model-name>/<version>/
+RUN mkdir -p /opt/ml/model/model_repository/<%= modelName || 'model' %>/1
+
+# Set permissions for model repository
+RUN chmod -R 755 /opt/ml/model/model_repository
+
+# Copy Triton model configuration
+COPY triton/config.pbtxt /opt/ml/model/model_repository/<%= modelName || 'model' %>/config.pbtxt
+
+<% if (backend === 'python') { %>
+# Copy Python backend model implementation
+COPY triton/model.py /opt/ml/model/model_repository/<%= modelName || 'model' %>/1/model.py
+<% } %>
+
+<% if (includeSampleModel) { %>
+# Copy sample model artifact
+<% if (backend === 'fil') { %>
+<% if (modelFormat === 'xgboost_json') { %>
+COPY sample_model/abalone_model.json /opt/ml/model/model_repository/<%= modelName || 'model' %>/1/xgboost.json
+<% } else if (modelFormat === 'xgboost_ubj') { %>
+COPY sample_model/abalone_model.ubj /opt/ml/model/model_repository/<%= modelName || 'model' %>/1/xgboost.ubj
+<% } else if (modelFormat === 'lightgbm_txt') { %>
+COPY sample_model/abalone_model.txt /opt/ml/model/model_repository/<%= modelName || 'model' %>/1/model.txt
+<% } %>
+<% } else if (backend === 'onnxruntime') { %>
+COPY sample_model/abalone_model.onnx /opt/ml/model/model_repository/<%= modelName || 'model' %>/1/model.onnx
+<% } else if (backend === 'tensorflow') { %>
+COPY sample_model/abalone_model.savedmodel /opt/ml/model/model_repository/<%= modelName || 'model' %>/1/model.savedmodel/
+<% } else if (backend === 'pytorch') { %>
+COPY sample_model/abalone_model.pt /opt/ml/model/model_repository/<%= modelName || 'model' %>/1/model.pt
+<% } else if (backend === 'python') { %>
+<% if (modelFormat === 'pkl') { %>
+COPY sample_model/abalone_model.pkl /opt/ml/model/model_repository/<%= modelName || 'model' %>/1/model.pkl
+<% } else if (modelFormat === 'joblib') { %>
+COPY sample_model/abalone_model.joblib /opt/ml/model/model_repository/<%= modelName || 'model' %>/1/model.joblib
+<% } %>
+<% } %>
+# Also copy training script for reference
+COPY sample_model/ /opt/ml/sample_model/
+<% } else { %>
+# Model artifacts should be placed in:
+# /opt/ml/model/model_repository/<%= modelName || 'model' %>/1/
+# COPY your_model_files /opt/ml/model/model_repository/<%= modelName || 'model' %>/1/
+<% } %>
+
+<% if (comments && comments.envVarExplanations && Object.keys(comments.envVarExplanations).length > 0) { %>
+# Environment Variables Configuration
+<% for (const [category, comment] of Object.entries(comments.envVarExplanations)) { %>
+<%= comment %>
+<% } %>
+<% } %>
+
+# Triton environment variables
+ENV TRITON_MODEL_REPOSITORY=/opt/ml/model/model_repository
+
+<% if (orderedEnvVars && orderedEnvVars.length > 0) { %>
+# Additional environment variables from configuration
+<% orderedEnvVars.forEach(({ key, value }) => { %>
+ENV <%= key %>=<%= value %>
+<% }); %>
+<% } %>
+
+# Expose port 8080 for SageMaker compatibility
+# Triton default ports: 8000 (HTTP), 8001 (gRPC), 8002 (metrics)
+# SageMaker requires port 8080
+EXPOSE 8080
+
+<% if (comments && comments.troubleshooting) { %>
+<%= comments.troubleshooting %>
+<% } %>
+
+# Start Triton Inference Server
+# --http-port=8080: SageMaker requires port 8080
+# --model-repository: Path to model repository
+# --strict-model-config=false: Allow Triton to auto-complete config for some backends
+ENTRYPOINT ["tritonserver", \
+            "--http-port=8080", \
+            "--model-repository=/opt/ml/model/model_repository", \
+            "--strict-model-config=false"]

package/templates/triton/config.pbtxt

@@ -0,0 +1,163 @@
+# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+# SPDX-License-Identifier: Apache-2.0
+
+# Triton Model Configuration
+# Backend: <%= backend %>
+# Model: <%= modelName || 'model' %>
+
+name: "<%= modelName || 'model' %>"
+backend: "<%= backend %>"
+
+<% if (backend === 'vllm' || backend === 'tensorrtllm') { %>
+# LLM backends (vllm, tensorrtllm) auto-configure most settings
+# Minimal configuration is sufficient
+max_batch_size: 0
+
+<% if (backend === 'vllm') { %>
+# vLLM backend parameters
+parameters: {
+    key: "model"
+    value: { string_value: "<%= modelName %>" }
+}
+parameters: {
+    key: "gpu_memory_utilization"
+    value: { string_value: "0.9" }
+}
+<% } else if (backend === 'tensorrtllm') { %>
+# TensorRT-LLM backend parameters
+parameters: {
+    key: "model"
+    value: { string_value: "<%= modelName %>" }
+}
+<% } %>
+<% } else { %>
+# Maximum batch size (0 = batching disabled)
+max_batch_size: 8
+
+<% if (backend === 'fil') { %>
+# FIL (Forest Inference Library) backend for tree-based models
+# Supports XGBoost, LightGBM, and scikit-learn random forests
+input [
+    {
+        name: "input__0"
+        data_type: TYPE_FP32
+        dims: [ -1 ]  # Dynamic feature dimension
+    }
+]
+output [
+    {
+        name: "output__0"
+        data_type: TYPE_FP32
+        dims: [ 1 ]  # Single prediction value
+    }
+]
+
+# FIL-specific parameters
+parameters: {
+    key: "model_type"
+<% if (modelFormat === 'xgboost_json' || modelFormat === 'xgboost_ubj') { %>
+    value: { string_value: "xgboost_json" }
+<% } else if (modelFormat === 'lightgbm_txt') { %>
+    value: { string_value: "lightgbm" }
+<% } %>
+}
+parameters: {
+    key: "output_class"
+    value: { string_value: "false" }
+}
+<% } else if (backend === 'onnxruntime') { %>
+# ONNX Runtime backend
+# Input/output shapes depend on your specific model
+input [
+    {
+        name: "input"
+        data_type: TYPE_FP32
+        dims: [ -1 ]  # Dynamic input dimension
+    }
+]
+output [
+    {
+        name: "output"
+        data_type: TYPE_FP32
+        dims: [ -1 ]  # Dynamic output dimension
+    }
+]
+<% } else if (backend === 'tensorflow') { %>
+# TensorFlow SavedModel backend
+# Input/output shapes depend on your specific model
+input [
+    {
+        name: "input"
+        data_type: TYPE_FP32
+        dims: [ -1 ]  # Dynamic input dimension
+    }
+]
+output [
+    {
+        name: "output"
+        data_type: TYPE_FP32
+        dims: [ -1 ]  # Dynamic output dimension
+    }
+]
+
+# TensorFlow-specific parameters
+parameters: {
+    key: "TF_INTER_OP_PARALLELISM"
+    value: { string_value: "0" }
+}
+parameters: {
+    key: "TF_INTRA_OP_PARALLELISM"
+    value: { string_value: "0" }
+}
+<% } else if (backend === 'pytorch') { %>
+# PyTorch TorchScript backend
+# Input/output shapes depend on your specific model
+input [
+    {
+        name: "INPUT__0"
+        data_type: TYPE_FP32
+        dims: [ -1 ]  # Dynamic input dimension
+    }
+]
+output [
+    {
+        name: "OUTPUT__0"
+        data_type: TYPE_FP32
+        dims: [ -1 ]  # Dynamic output dimension
+    }
+]
+<% } else if (backend === 'python') { %>
+# Python backend
+# Custom model implementation in model.py
+input [
+    {
+        name: "INPUT"
+        data_type: TYPE_FP32
+        dims: [ -1 ]  # Dynamic input dimension
+    }
+]
+output [
+    {
+        name: "OUTPUT"
+        data_type: TYPE_FP32
+        dims: [ -1 ]  # Dynamic output dimension
+    }
+]
+<% } %>
+
+# Instance group configuration
+# Specifies how many instances of the model to run and on which device
+instance_group [
+    {
+        count: 1
+        kind: KIND_AUTO  # Automatically select CPU or GPU based on availability
+    }
+]
+
+# Dynamic batching configuration
+# Enables automatic request batching for improved throughput
+dynamic_batching {
+    preferred_batch_size: [ 4, 8 ]
+    max_queue_delay_microseconds: 100
+}
+<% } %>

package/templates/triton/model.py

@@ -0,0 +1,130 @@
+# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+# SPDX-License-Identifier: Apache-2.0
+
+"""
+Triton Python Backend Model Implementation
+
+This module implements the TritonPythonModel interface for serving
+custom Python models via NVIDIA Triton Inference Server.
+
+Backend: python
+Model: <%= modelName || 'model' %>
+"""
+
+import json
+import os
+
+import numpy as np
+import triton_python_backend_utils as pb_utils
+
+<% if (modelFormat === 'pkl') { %>
+import pickle
+<% } else if (modelFormat === 'joblib') { %>
+import joblib
+<% } %>
+
+
+class TritonPythonModel:
+    """Triton Python backend model implementation.
+
+    This class implements the required interface for Triton's Python backend:
+    - initialize(): Called once when the model is loaded
+    - execute(): Called for each inference request batch
+    - finalize(): Called once when the model is unloaded
+    """
+
+    def initialize(self, args):
+        """Initialize the model.
+
+        Called once when the model is loaded by Triton. Use this method to
+        load model artifacts and set up any required resources.
+
+        Args:
+            args: Dictionary containing model configuration:
+                - model_config: JSON string of the model configuration
+                - model_instance_kind: Device type (CPU/GPU)
+                - model_instance_device_id: Device ID
+                - model_repository: Path to the model repository
+                - model_version: Model version being loaded
+                - model_name: Name of the model
+        """
+        self.model_config = json.loads(args['model_config'])
+        model_repository = args['model_repository']
+        model_version = args['model_version']
+
+        # Construct path to model artifact
+        model_dir = os.path.join(model_repository, model_version)
+
+<% if (modelFormat === 'pkl') { %>
+        # Load pickle model
+        model_path = os.path.join(model_dir, 'model.pkl')
+        with open(model_path, 'rb') as f:
+            self.model = pickle.load(f)
+<% } else if (modelFormat === 'joblib') { %>
+        # Load joblib model
+        model_path = os.path.join(model_dir, 'model.joblib')
+        self.model = joblib.load(model_path)
+<% } else { %>
+        # Custom model loading
+        # TODO: Implement your model loading logic here
+        # model_path = os.path.join(model_dir, 'your_model_file')
+        self.model = None
+<% } %>
+
+        # Get output configuration from model config
+        output_config = pb_utils.get_output_config_by_name(
+            self.model_config, 'OUTPUT'
+        )
+        self.output_dtype = pb_utils.triton_string_to_numpy(
+            output_config['data_type']
+        )
+
+    def execute(self, requests):
+        """Handle inference requests.
+
+        Called for each batch of inference requests. Processes input tensors
+        and returns output tensors.
+
+        Args:
+            requests: List of pb_utils.InferenceRequest objects
+
+        Returns:
+            List of pb_utils.InferenceResponse objects
+        """
+        responses = []
+
+        for request in requests:
+            # Get input tensor
+            input_tensor = pb_utils.get_input_tensor_by_name(request, 'INPUT')
+            input_data = input_tensor.as_numpy()
+
+<% if (modelFormat === 'pkl' || modelFormat === 'joblib') { %>
+            # Run prediction
+            predictions = self.model.predict(input_data)
+            output_data = np.array(predictions, dtype=self.output_dtype)
+<% } else { %>
+            # Custom inference logic
+            # TODO: Implement your inference logic here
+            output_data = np.zeros(
+                (input_data.shape[0], 1), dtype=self.output_dtype
+            )
+<% } %>
+
+            # Create output tensor
+            output_tensor = pb_utils.Tensor('OUTPUT', output_data)
+
+            # Create inference response
+            inference_response = pb_utils.InferenceResponse(
+                output_tensors=[output_tensor]
+            )
+            responses.append(inference_response)
+
+        return responses
+
+    def finalize(self):
+        """Clean up resources.
+
+        Called once when the model is being unloaded. Use this method to
+        release any resources held by the model.
+        """
+        self.model = None

package/templates/triton/requirements.txt

@@ -0,0 +1,11 @@
+# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+# SPDX-License-Identifier: Apache-2.0
+
+# Python dependencies for Triton Python backend
+numpy>=1.24.0
+<% if (modelFormat === 'pkl') { %>
+scikit-learn>=1.3.0
+<% } else if (modelFormat === 'joblib') { %>
+scikit-learn>=1.3.0
+joblib>=1.3.0
+<% } %>