@aws/ml-container-creator 0.2.0

package/templates/buildspec.yml

@@ -0,0 +1,64 @@
+version: 0.2
+
+env:
+  variables:
+    AWS_DEFAULT_REGION: <%= awsRegion %>
+    AWS_ACCOUNT_ID: ""
+    ECR_REPOSITORY_NAME: "ml-container-creator"
+    PROJECT_NAME: "<%= projectName %>"
+    IMAGE_TAG: "latest"
+
+phases:
+  pre_build:
+    commands:
+      - echo Logging in to Amazon ECR...
+      - AWS_ACCOUNT_ID=$(aws sts get-caller-identity --query Account --output text)
+      - aws ecr get-login-password --region $AWS_DEFAULT_REGION | docker login --username AWS --password-stdin $AWS_ACCOUNT_ID.dkr.ecr.$AWS_DEFAULT_REGION.amazonaws.com
+      - REPOSITORY_URI=$AWS_ACCOUNT_ID.dkr.ecr.$AWS_DEFAULT_REGION.amazonaws.com/$ECR_REPOSITORY_NAME
+      - IMAGE_TAG=${CODEBUILD_RESOLVED_SOURCE_VERSION:-latest}
+      - PROJECT_TAG="$PROJECT_NAME-$(date +%Y%m%d-%H%M%S)"
+      - echo Repository URI is $REPOSITORY_URI
+      - echo Project tag is $PROJECT_TAG
+      - echo Image tag is $IMAGE_TAG
+    on-failure: ABORT
+  build:
+    commands:
+      - echo Build started on `date`
+      - echo Building the Docker image for project $PROJECT_NAME...
+<% if (typeof modelLoadStrategy !== 'undefined' && modelLoadStrategy === 'build-time' && typeof modelSource !== 'undefined' && modelSource && modelSource !== 'huggingface') { %>
+      - |
+        # Export IAM role credentials for docker build
+        eval $(aws configure export-credentials --format env 2>/dev/null || true)
+        docker build \
+          --build-arg AWS_ACCESS_KEY_ID="${AWS_ACCESS_KEY_ID}" \
+          --build-arg AWS_SECRET_ACCESS_KEY="${AWS_SECRET_ACCESS_KEY}" \
+          --build-arg AWS_SESSION_TOKEN="${AWS_SESSION_TOKEN}" \
+          --build-arg AWS_DEFAULT_REGION="${AWS_DEFAULT_REGION}" \
+          -t $REPOSITORY_URI:$PROJECT_TAG .
+<% } else { %>
+      - docker build -t $REPOSITORY_URI:$PROJECT_TAG .
+<% } %>
+      - docker tag $REPOSITORY_URI:$PROJECT_TAG $REPOSITORY_URI:$PROJECT_NAME-latest
+      - docker tag $REPOSITORY_URI:$PROJECT_TAG $REPOSITORY_URI:latest
+      - echo Build completed on `date`
+    on-failure: ABORT
+  post_build:
+    commands:
+      - echo Post-build started on `date`
+      - echo Pushing the Docker images for project $PROJECT_NAME...
+      - docker push $REPOSITORY_URI:$PROJECT_TAG || (echo "Failed to push project tag $PROJECT_TAG" && exit 1)
+      - docker push $REPOSITORY_URI:$PROJECT_NAME-latest || (echo "Failed to push project latest tag" && exit 1)
+      - docker push $REPOSITORY_URI:latest || (echo "Failed to push latest tag" && exit 1)
+      - echo Successfully pushed images to ECR repository $ECR_REPOSITORY_NAME
+      - echo "Available tags:"
+      - echo "  - $PROJECT_TAG (timestamped build)"
+      - echo "  - $PROJECT_NAME-latest (project latest)"
+      - echo "  - latest (global latest)"
+      - echo Writing image definitions file...
+      - printf '[{"name":"%s","imageUri":"%s"}]' $PROJECT_NAME $REPOSITORY_URI:$PROJECT_TAG > imagedefinitions.json
+      - echo Post-build completed on `date`
+
+artifacts:
+  files:
+    - imagedefinitions.json
+  name: <%= projectName %>-artifacts

package/templates/code/chat_template.jinja

@@ -0,0 +1 @@
+<%= chatTemplate %>

package/templates/code/flask/gunicorn_config.py

@@ -0,0 +1,35 @@
+# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+# SPDX-License-Identifier: Apache-2.0
+
+"""
+Gunicorn configuration file
+"""
+import os
+import multiprocessing
+from serve import load_model_for_worker
+
+# Bind address
+bind = f"0.0.0.0:8080"
+
+# Worker processes - cap at 4 for memory efficiency
+workers = min(multiprocessing.cpu_count(), 4)
+
+# Worker type
+worker_class = 'sync'
+
+# Timeouts
+timeout = 120
+keepalive = 5
+
+# Worker lifecycle management
+max_requests = 500
+max_requests_jitter = 100
+
+# Logging
+accesslog = '-'  # Log to stdout
+errorlog = '-'   # Log to stderr
+loglevel = 'info'
+
+# Worker initialization hook - load model in each worker
+def post_worker_init(worker):
+    load_model_for_worker()

package/templates/code/flask/wsgi.py

@@ -0,0 +1,10 @@
+# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+# SPDX-License-Identifier: Apache-2.0
+
+"""
+WSGI entry point for Gunicorn
+"""
+from serve import create_app
+
+# Create the Flask application
+application = create_app()

package/templates/code/model_handler.py

@@ -0,0 +1,387 @@
+# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+# SPDX-License-Identifier: Apache-2.0
+<% if (framework === 'sklearn') { %>
+"""
+SKLearn model handler for SageMaker inference
+"""
+import os
+import json
+import pickle
+import joblib
+import numpy as np
+from typing import Any, Dict
+import logging
+
+logging.basicConfig(level=logging.INFO)
+logger = logging.getLogger(__name__)
+
+class ModelHandler:
+    """Handle SKLearn model loading and inference"""
+    
+    def __init__(self, model_path: str):
+        self.model_path = model_path
+        self.model = None
+        self._loaded = False
+        
+    def load_model(self):
+        """Load the SKLearn model"""
+        try:
+            model_files = [f for f in os.listdir(self.model_path) if f.endswith('<%= modelFormat %>')]
+            
+            if not model_files:
+                logger.warning("No SKLearn model files found in model directory")
+                logger.warning("Server will start but /invocations will fail until a model is provided")
+                logger.warning("Mount a model directory with: MODEL_DIR=/path/to/model ./do/run")
+                return
+            
+            model_file = os.path.join(self.model_path, model_files[0])
+            logger.info(f"Loading model from {model_file}")
+            
+            # Load with joblib first, fallback to pickle
+            try:
+                self.model = joblib.load(model_file)
+            except:
+                with open(model_file, 'rb') as f:
+                    self.model = pickle.load(f)
+            
+            self._loaded = True
+            logger.info("SKLearn model loaded successfully")
+            
+        except Exception as e:
+            logger.error(f"Error loading model: {str(e)}")
+            raise
+    
+    def is_loaded(self) -> bool:
+        """Check if model is loaded"""
+        return self._loaded and self.model is not None
+    
+    def preprocess(self, raw_data: Any) -> np.ndarray:
+        """Preprocess input data for SKLearn model"""
+        try:
+            if isinstance(raw_data, dict):
+                data = raw_data.get('instances', raw_data.get('data', raw_data))
+            else:
+                data = raw_data
+            
+            if isinstance(data, str):
+                data = json.loads(data)
+            
+            return np.array(data)
+            
+        except Exception as e:
+            logger.error(f"Error in preprocessing: {str(e)}")
+            raise ValueError(f"Invalid input data format: {str(e)}")
+    
+    def postprocess(self, predictions: np.ndarray) -> Dict[str, Any]:
+        """Postprocess SKLearn model predictions"""
+        try:
+            if hasattr(predictions, 'tolist'):
+                predictions = predictions.tolist()
+            
+            return {'predictions': predictions}
+            
+        except Exception as e:
+            logger.error(f"Error in postprocessing: {str(e)}")
+            raise
+    
+    def predict(self, input_data: Any) -> Dict[str, Any]:
+        """Run inference on input data"""
+        if not self.is_loaded():
+            raise RuntimeError("Model is not loaded")
+        
+        try:
+            processed_input = self.preprocess(input_data)
+            predictions = self.model.predict(processed_input)
+            return self.postprocess(predictions)
+            
+        except Exception as e:
+            logger.error(f"Error during inference: {str(e)}")
+            raise
+<% } else if (framework === 'xgboost') { %>
+"""
+XGBoost model handler for SageMaker inference
+"""
+import os
+import json
+import xgboost as xgb
+import numpy as np
+from typing import Any, Dict
+import logging
+
+logging.basicConfig(level=logging.INFO)
+logger = logging.getLogger(__name__)
+
+class ModelHandler:
+    """Handle XGBoost model loading and inference"""
+    
+    def __init__(self, model_path: str):
+        self.model_path = model_path
+        self.model = None
+        self._loaded = False
+        
+    def load_model(self):
+        """Load the XGBoost model"""
+        try:
+            model_files = [f for f in os.listdir(self.model_path) if f.endswith('<%= modelFormat %>')]
+            
+            if not model_files:
+                logger.warning("No XGBoost model files found in model directory")
+                logger.warning("Server will start but /invocations will fail until a model is provided")
+                logger.warning("Mount a model directory with: MODEL_DIR=/path/to/model ./do/run")
+                return
+            
+            model_file = os.path.join(self.model_path, model_files[0])
+            logger.info(f"Loading model from {model_file}")
+            
+            self.model = xgb.Booster()
+            self.model.load_model(model_file)
+            
+            self._loaded = True
+            logger.info("XGBoost model loaded successfully")
+            
+        except Exception as e:
+            logger.error(f"Error loading model: {str(e)}")
+            raise
+    
+    def is_loaded(self) -> bool:
+        """Check if model is loaded"""
+        return self._loaded and self.model is not None
+    
+    def preprocess(self, raw_data: Any) -> xgb.DMatrix:
+        """Preprocess input data for XGBoost model"""
+        try:
+            if isinstance(raw_data, dict):
+                data = raw_data.get('instances', raw_data.get('data', raw_data))
+            else:
+                data = raw_data
+            
+            if isinstance(data, str):
+                data = json.loads(data)
+            
+            return xgb.DMatrix(np.array(data))
+            
+        except Exception as e:
+            logger.error(f"Error in preprocessing: {str(e)}")
+            raise ValueError(f"Invalid input data format: {str(e)}")
+    
+    def postprocess(self, predictions: np.ndarray) -> Dict[str, Any]:
+        """Postprocess XGBoost model predictions"""
+        try:
+            if hasattr(predictions, 'tolist'):
+                predictions = predictions.tolist()
+            
+            return {'predictions': predictions}
+            
+        except Exception as e:
+            logger.error(f"Error in postprocessing: {str(e)}")
+            raise
+    
+    def predict(self, input_data: Any) -> Dict[str, Any]:
+        """Run inference on input data"""
+        if not self.is_loaded():
+            raise RuntimeError("Model is not loaded")
+        
+        try:
+            processed_input = self.preprocess(input_data)
+            predictions = self.model.predict(processed_input)
+            return self.postprocess(predictions)
+            
+        except Exception as e:
+            logger.error(f"Error during inference: {str(e)}")
+            raise
+<% } else if (framework === 'tensorflow') { %>
+"""
+TensorFlow model handler for SageMaker inference
+"""
+import os
+import json
+import tensorflow as tf
+import numpy as np
+from typing import Any, Dict
+import logging
+
+logging.basicConfig(level=logging.INFO)
+logger = logging.getLogger(__name__)
+
+class ModelHandler:
+    """Handle TensorFlow model loading and inference"""
+
+    def __init__(self, model_path: str):
+        self.model_path = model_path
+        self.model = None
+        self._loaded = False
+
+    def load_model(self):
+        """Load the TensorFlow model"""
+        try:
+            logger.info(f"Loading model from {self.model_path}")
+
+            model_files = [f for f in os.listdir(self.model_path) if f.endswith(('.keras', '.h5'))]
+
+            if not model_files:
+                try:
+                    self.model = tf.saved_model.load(self.model_path)
+                    self.model = self.model.signatures['serving_default']
+                except:
+                    logger.warning("No TensorFlow model files found in model directory")
+                    logger.warning("Server will start but /invocations will fail until a model is provided")
+                    logger.warning("Mount a model directory with: MODEL_DIR=/path/to/model ./do/run")
+                    return
+            else:
+                model_file = os.path.join(self.model_path, model_files[0])
+                logger.info(f"Loading model from {model_file}")
+                self.model = tf.keras.models.load_model(model_file, compile=False)
+
+            self._loaded = True
+            logger.info("TensorFlow model loaded successfully")
+
+        except Exception as e:
+            logger.error(f"Error loading model: {str(e)}")
+            raise
+
+    def is_loaded(self) -> bool:
+        """Check if model is loaded"""
+        return self._loaded and self.model is not None
+
+    def preprocess(self, raw_data: Any) -> np.ndarray:
+        """Preprocess input data for TensorFlow model"""
+        try:
+            if isinstance(raw_data, dict):
+                data = raw_data.get('instances', raw_data.get('data', raw_data))
+            else:
+                data = raw_data
+
+            if isinstance(data, str):
+                data = json.loads(data)
+
+            return np.array(data, dtype=np.float32)
+
+        except Exception as e:
+            logger.error(f"Error in preprocessing: {str(e)}")
+            raise ValueError(f"Invalid input data format: {str(e)}")
+
+    def postprocess(self, predictions: np.ndarray) -> Dict[str, Any]:
+        """Postprocess TensorFlow model predictions"""
+        try:
+            if hasattr(predictions, 'numpy'):
+                predictions = predictions.numpy()
+
+            if hasattr(predictions, 'tolist'):
+                predictions = predictions.tolist()
+
+            return {'predictions': predictions}
+
+        except Exception as e:
+            logger.error(f"Error in postprocessing: {str(e)}")
+            raise
+
+    def predict(self, input_data: Any) -> Dict[str, Any]:
+        """Run inference on input data"""
+        if not self.is_loaded():
+            raise RuntimeError("Model is not loaded")
+
+        try:
+            processed_input = self.preprocess(input_data)
+
+            # Handle different model types
+            if hasattr(self.model, 'predict'):
+                # Keras model
+                predictions = self.model.predict(processed_input)
+            else:
+                # SavedModel signature
+                input_tensor = tf.constant(processed_input)
+                result = self.model(input_tensor)
+                predictions = list(result.values())[0]
+
+            return self.postprocess(predictions)
+
+        except Exception as e:
+            logger.error(f"Error during inference: {str(e)}")
+            raise
+<% } else if (framework === 'transformers' && modelServer === 'sglang') { %>
+"""
+SGLang model handler for SageMaker inference
+"""
+import os
+import json
+from typing import Any, Dict, List
+import logging
+from sglang import Runtime
+
+logging.basicConfig(level=logging.INFO)
+logger = logging.getLogger(__name__)
+
+class ModelHandler:
+    """Handle SGLang model loading and inference"""
+
+    def __init__(self, model_path: str):
+        self.model_path = model_path
+        self.runtime = None
+        self._loaded = False
+
+    def load_model(self):
+        """Initialize SGLang runtime"""
+        try:
+            model_id = '<%= modelName %>'
+            logger.info(f"Loading SGLang model: {model_id}")
+
+            self.runtime = Runtime(
+                model_path=model_id,
+                tokenizer_path=model_id,
+                device="cuda",
+                mem_fraction_static=0.8
+            )
+
+            self._loaded = True
+            logger.info("SGLang model loaded successfully")
+
+        except Exception as e:
+            logger.error(f"Error loading model: {str(e)}")
+            raise
+
+    def is_loaded(self) -> bool:
+        """Check if model is loaded"""
+        return self._loaded and self.runtime is not None
+
+    def preprocess(self, raw_data: Any) -> List[str]:
+        """Preprocess input data for SGLang model"""
+        try:
+            if isinstance(raw_data, dict):
+                data = raw_data.get('instances', raw_data.get('inputs', raw_data))
+            else:
+                data = raw_data
+
+            if isinstance(data, str):
+                return [data]
+            elif isinstance(data, list):
+                return data
+            else:
+                raise ValueError("Input must be string or list of strings")
+
+        except Exception as e:
+            logger.error(f"Error in preprocessing: {str(e)}")
+            raise ValueError(f"Invalid input data format: {str(e)}")
+
+    def postprocess(self, outputs: List[str]) -> Dict[str, Any]:
+        """Postprocess SGLang model outputs"""
+        try:
+            return {'predictions': outputs}
+
+        except Exception as e:
+            logger.error(f"Error in postprocessing: {str(e)}")
+            raise
+
+    def predict(self, input_data: Any) -> Dict[str, Any]:
+        """Run inference on input data"""
+        if not self.is_loaded():
+            raise RuntimeError("Model is not loaded")
+
+        try:
+            prompts = self.preprocess(input_data)
+            outputs = self.runtime.generate(prompts)
+            return self.postprocess(outputs)
+
+        except Exception as e:
+            logger.error(f"Error during inference: {str(e)}")
+            raise
+<% } %>