SageMaker PyTorch MME Ignores Custom `inference.py` and Falls Back to Default Handler

Was the above answer helpful, is the issue resolved? If not I'll add some additional items.

1. Model Format Verification: The original error mentions TorchScript format specifically, so verifying that your models are in the correct format and properly saved using torch.save() is crucial.

2. Logging Implementation: trying a "minimal 'hello world' inference.py with only logging" - implementing proper logging statements in your custom inference.py could help debug where exactly the process is failing.

3. Lower-level API Approach: attempting to use "the lower-level sagemaker.model.Model class with explicit SAGEMAKER_PROGRAM and MMS_DEFAULT_HANDLER environment variables" - this could be an alternative solution path worth exploring if the higher-level API continues to fail.

4. Model Artifact Structure: the importance of having clean model artifacts in S3 (.tar.gz files) without any Python code, which is a critical consideration for proper deployment.

5. Python Path Issues: potential Python path issues, suggesting that the directory structure and Python path configuration could be important factors to consider when troubleshooting. Edit: Apologies as you've tried some of these. Please let us know if the issue persists.

When deploying a PyTorch Multi-Model Endpoint (MME), there are several key issues that could cause your custom inference.py script to be ignored:

1. Code Directory Structure: For MME deployments, your code structure is critical. The inference.py file must be placed in a code/ directory within your model tarball. The SageMaker container expects this specific structure to properly load your custom handler.

2. Model vs. Code Separation: For MME, there's an important distinction between model artifacts and code. In your case, you mentioned your model artifacts in S3 contain no Python code, which is actually part of the problem. For MME deployments, the code needs to be included with each model artifact or provided separately in a way that ensures it's available at /opt/ml/model/code.

3. Container Initialization: When the MME container starts, it looks for the code in specific locations. If it can't find your custom handler at /opt/ml/model/code, it will fall back to the default handler, which expects TorchScript models.

4. Network Isolation: If you have enable_network_isolation=True, this could affect how your code is packaged and deployed. For MME, you might need to set this to False to ensure proper separation of code and model artifacts.

Here's what you should try:

1. Modify your deployment code to ensure the source directory is properly packaged:

pytorch_model_def = PyTorchModel(
entry_point='inference.py',
source_dir='./code',
role=role_arn,
sagemaker_session=sagemaker_session,
model_data=s3_model_repository + "/initial_model.tar.gz",  # Provide an initial model
framework_version='2.0',
py_version='py310',
enable_network_isolation=False  # Important for MME
)

2. Ensure your model_fn function in inference.py is properly implemented to handle your non-TorchScript models. For example:

def model_fn(model_dir):
# Custom code to load your model from state_dict
model = YourModelClass()
weights_path = os.path.join(model_dir, "model.pt")
model.load_state_dict(torch.load(weights_path))
model.eval()
return model

3. Check your CloudWatch logs for any permission issues or file path errors. The container might be looking for your code in a different location than expected.

4. If you're still having issues, try creating a simple test case with a single model deployment (non-MME) first to verify your custom handler works, then migrate to the MME setup.

The key insight is that for MME, the code and model separation works differently than for standard endpoints, and ensuring your code is available at the right location is crucial for the container to use your custom handler instead of falling back to the default one.
Sources
FileNotFoundError: [Errno 2] No such file or directory: '/opt/ml/model/code' when deploying SageMaker Multi-model Endpoint (MME) | AWS re:Post
Create a SageMaker inference endpoint with custom model & extended container | AWS Machine Learning Blog
Build Your Own Container for SageMaker AI Multi-Model Endpoints - Amazon SageMaker AI

This is my actual full deploy.py

# deploy_tgn_mme.py

import sagemaker
from sagemaker.model import Model
from sagemaker.multidatamodel import MultiDataModel
from sagemaker.pytorch.model import PyTorchPredictor
import time
import boto3
import logging
import os
import tarfile


logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(name)s - %(message)s')
logger = logging.getLogger(__name__)

def create_dummy_model_artifact():
    dummy_dir = './dummy_model'
    dummy_tar_path = './dummy_model.tar.gz'
    os.makedirs(dummy_dir, exist_ok=True)
    with open(os.path.join(dummy_dir, 'dummy.txt'), 'w') as f:
        f.write('This is a dummy model artifact for MME setup.')
    with tarfile.open(dummy_tar_path, 'w:gz') as tar:
        tar.add(dummy_dir, arcname='.')
    import shutil
    shutil.rmtree(dummy_dir)
    return dummy_tar_path

def main():

    s3_model_repository = "s3://.../tgn-mme-models/"
    endpoint_name = f"tgn-recommender-mme-cpu-{time.strftime('%Y%m%d-%H%M%S')}"
    role_arn = "..."
    aws_region = "..."

    instance_type = 'ml.c5.xlarge' 
    instance_count = 1



    boto_session = boto3.Session(region_name=aws_region)
    sagemaker_session = sagemaker.Session(boto_session=boto_session)
    


    dummy_tar_path = create_dummy_model_artifact()
    s3_client = boto3.client('s3', region_name=aws_region)
    dummy_s3_key = f"dummy-models/dummy_model_{os.path.basename(dummy_tar_path)}"
    s3_client.upload_file(dummy_tar_path, '...-recommender', dummy_s3_key)
    dummy_model_data = f"s3://...-recommender/{dummy_s3_key}"
    os.remove(dummy_tar_path)


    image_uri = sagemaker.image_uris.retrieve(
        framework="pytorch",
        region=aws_region,
        version="2.0",
        py_version="py310",
        instance_type=instance_type, 
        image_scope="inference"
    )


   
    base_model = Model(
        image_uri=image_uri,
        role=role_arn,
        model_data=dummy_model_data,
        source_dir='./code', 
        entry_point='inference.py',
        sagemaker_session=sagemaker_session,
        env={
            'SAGEMAKER_CONTAINER_LOG_LEVEL': '20',
            'SAGEMAKER_REGION': aws_region,
            'SAGEMAKER_PROGRAM': 'inference.py',
        }
    )
    
    multi_data_model = MultiDataModel(
        name=f"tgn-mme-container-{time.strftime('%Y%m%d-%H%M%S')}",
        model_data_prefix=s3_model_repository,
        model=base_model
    )
    try:
        multi_data_model.deploy(
            initial_instance_count=instance_count,
            instance_type=instance_type,
            endpoint_name=endpoint_name,
            wait=True
        )

        predictor = PyTorchPredictor(
            endpoint_name=endpoint_name,
            sagemaker_session=sagemaker_session
        )

        logger.info("\n==========================================")
        logger.info("✅ Done!")
        logger.info(f"Endpoint Name: {predictor.endpoint_name}")
        logger.info("==========================================")
        
    except Exception as e:
        logger.error("\n==========================================")
        logger.error(f"❌ Deployment failed: {e}", exc_info=True)
        logger.error("==========================================")
        
        client = sagemaker_session.sagemaker_client
        try:
            client.delete_endpoint(EndpointName=endpoint_name)
        except client.exceptions.ClientError: pass
        try:
            client.delete_endpoint_config(EndpointConfigName=endpoint_name)
        except client.exceptions.ClientError: pass
        try:
            client.delete_model(ModelName=multi_data_model.name)
        except client.exceptions.ClientError: pass

if __name__ == "__main__":
    main()