Issue while deploying custom torch model on SageMaker

[tararelan]

I am trying to deploy a custom neural network on SageMaker. However, when I run the deployment code (which I got from the AWS examples), I get the following error:
An error occurred (ModelError) when calling the InvokeEndpoint operation: Received server error (500) from primary and could not load the entire response body. See https://eu-west-2.console.aws.amazon.com/cloudwatch/home?region=eu-west-2#logEventViewer:group=/aws/sagemaker/Endpoints/pytorch-inference-2024-12-30-09-09-34-368 for more information.

This is also the log - there are no logs tagged as errors.
timestamp,message,logStreamName
1735549918460,"2024-12-30 09:11:57,620 [INFO ] main com.amazonaws.ml.mms.ModelServer - ",AllTraffic/i-08ffdaaab186465db
1735549918460,MMS Home: /opt/conda/lib/python3.6/site-packages,AllTraffic/i-08ffdaaab186465db
1735549918460,Current directory: /,AllTraffic/i-08ffdaaab186465db
1735549918460,Temp directory: /home/model-server/tmp,AllTraffic/i-08ffdaaab186465db
1735549918460,Number of GPUs: 0,AllTraffic/i-08ffdaaab186465db
1735549918460,Number of CPUs: 1,AllTraffic/i-08ffdaaab186465db
1735549918460,Max heap size: 3197 M,AllTraffic/i-08ffdaaab186465db
1735549918460,Python executable: /opt/conda/bin/python3.6,AllTraffic/i-08ffdaaab186465db
1735549918460,Config file: /etc/sagemaker-mms.properties,AllTraffic/i-08ffdaaab186465db
1735549918460,Inference address: http://0.0.0.0:8080,AllTraffic/i-08ffdaaab186465db
1735549918460,Management address: http://0.0.0.0:8080,AllTraffic/i-08ffdaaab186465db
1735549918460,Model Store: /.sagemaker/mms/models,AllTraffic/i-08ffdaaab186465db
1735549918460,Initial Models: ALL,AllTraffic/i-08ffdaaab186465db
1735549918460,Log dir: /logs,AllTraffic/i-08ffdaaab186465db
1735549918460,Metrics dir: /logs,AllTraffic/i-08ffdaaab186465db
1735549918460,Netty threads: 0,AllTraffic/i-08ffdaaab186465db
1735549918460,Netty client threads: 0,AllTraffic/i-08ffdaaab186465db
1735549918460,Default workers per model: 1,AllTraffic/i-08ffdaaab186465db
1735549918460,Blacklist Regex: N/A,AllTraffic/i-08ffdaaab186465db
1735549918460,Maximum Response Size: 6553500,AllTraffic/i-08ffdaaab186465db
1735549918460,Maximum Request Size: 6553500,AllTraffic/i-08ffdaaab186465db
1735549918460,Preload model: false,AllTraffic/i-08ffdaaab186465db
1735549918460,Prefer direct buffer: false,AllTraffic/i-08ffdaaab186465db
1735549918460,"2024-12-30 09:11:57,694 [WARN ] W-9000-model com.amazonaws.ml.mms.wlm.WorkerLifeCycle - attachIOStreams() threadName=W-9000-model",AllTraffic/i-08ffdaaab186465db
1735549918460,"2024-12-30 09:11:57,787 [INFO ] W-9000-model-stdout com.amazonaws.ml.mms.wlm.WorkerLifeCycle - model_service_worker started with args: --sock-type unix --sock-name /home/model-server/tmp/.mms.sock.9000 --handler sagemaker_pytorch_serving_container.handler_service --model-path /.sagemaker/mms/models/model --model-name model --preload-model false --tmp-dir /home/model-server/tmp",AllTraffic/i-08ffdaaab186465db
1735549918460,"2024-12-30 09:11:57,788 [INFO ] W-9000-model-stdout com.amazonaws.ml.mms.wlm.WorkerLifeCycle - Listening on port: /home/model-server/tmp/.mms.sock.9000",AllTraffic/i-08ffdaaab186465db
1735549918460,"2024-12-30 09:11:57,788 [INFO ] W-9000-model-stdout com.amazonaws.ml.mms.wlm.WorkerLifeCycle - [PID] 29",AllTraffic/i-08ffdaaab186465db
1735549918460,"2024-12-30 09:11:57,789 [INFO ] W-9000-model-stdout com.amazonaws.ml.mms.wlm.WorkerLifeCycle - MMS worker started.",AllTraffic/i-08ffdaaab186465db
1735549918460,"2024-12-30 09:11:57,789 [INFO ] W-9000-model-stdout com.amazonaws.ml.mms.wlm.WorkerLifeCycle - Python runtime: 3.6.13",AllTraffic/i-08ffdaaab186465db
1735549918460,"2024-12-30 09:11:57,789 [INFO ] main com.amazonaws.ml.mms.wlm.ModelManager - Model model loaded.",AllTraffic/i-08ffdaaab186465db
1735549918460,"2024-12-30 09:11:57,793 [INFO ] main com.amazonaws.ml.mms.ModelServer - Initialize Inference server with: EpollServerSocketChannel.",AllTraffic/i-08ffdaaab186465db
1735549918460,"2024-12-30 09:11:57,802 [INFO ] W-9000-model com.amazonaws.ml.mms.wlm.WorkerThread - Connecting to: /home/model-server/tmp/.mms.sock.9000",AllTraffic/i-08ffdaaab186465db
1735549918460,"2024-12-30 09:11:57,852 [INFO ] main com.amazonaws.ml.mms.ModelServer - Inference API bind to: http://0.0.0.0:8080",AllTraffic/i-08ffdaaab186465db
1735549918460,Model server started.,AllTraffic/i-08ffdaaab186465db
1735549918460,"2024-12-30 09:11:57,859 [INFO ] W-9000-model-stdout com.amazonaws.ml.mms.wlm.WorkerLifeCycle - Connection accepted: /home/model-server/tmp/.mms.sock.9000.",AllTraffic/i-08ffdaaab186465db
1735549918702,"2024-12-30 09:11:57,860 [WARN ] pool-2-thread-1 com.amazonaws.ml.mms.metrics.MetricCollector - worker pid is not available yet.",AllTraffic/i-08ffdaaab186465db
1735549918702,"2024-12-30 09:11:58,598 [INFO ] W-9000-model-stdout com.amazonaws.ml.mms.wlm.WorkerLifeCycle - Model model loaded io_fd=e218b1fffe837cab-00000011-00000000-4bff44452ed61a70-223c9727",AllTraffic/i-08ffdaaab186465db
1735549918702,"2024-12-30 09:11:58,604 [INFO ] W-9000-model com.amazonaws.ml.mms.wlm.WorkerThread - Backend response time: 705",AllTraffic/i-08ffdaaab186465db
1735549920456,"2024-12-30 09:11:58,605 [WARN ] W-9000-model com.amazonaws.ml.mms.wlm.WorkerLifeCycle - attachIOStreams() threadName=W-model-1",AllTraffic/i-08ffdaaab186465db
1735549924460,"2024-12-30 09:12:00,295 [INFO ] pool-1-thread-3 ACCESS_LOG - /169.254.178.2:49624 ""GET /ping HTTP/1.1"" 200 31",AllTraffic/i-08ffdaaab186465db
1735549930461,"2024-12-30 09:12:05,231 [INFO ] pool-1-thread-3 ACCESS_LOG - /169.254.178.2:49624 ""GET /ping HTTP/1.1"" 200 0",AllTraffic/i-08ffdaaab186465db
1735549935241,"2024-12-30 09:12:10,229 [INFO ] pool-1-thread-3 ACCESS_LOG - /169.254.178.2:49624 ""GET /ping HTTP/1.1"" 200 0",AllTraffic/i-08ffdaaab186465db
1735549939461,"2024-12-30 09:12:15,229 [INFO ] pool-1-thread-3 ACCESS_LOG - /169.254.178.2:49624 ""GET /ping HTTP/1.1"" 200 0",AllTraffic/i-08ffdaaab186465db
1735549944505,"2024-12-30 09:12:20,229 [INFO ] pool-1-thread-3 ACCESS_LOG - /169.254.178.2:49624 ""GET /ping HTTP/1.1"" 200 0",AllTraffic/i-08ffdaaab186465db
1735549949461,"2024-12-30 09:12:25,229 [INFO ] pool-1-thread-3 ACCESS_LOG - /169.254.178.2:49624 ""GET /ping HTTP/1.1"" 200 0",AllTraffic/i-08ffdaaab186465db
1735549954461,"2024-12-30 09:12:30,229 [INFO ] pool-1-thread-3 ACCESS_LOG - /169.254.178.2:49624 ""GET /ping HTTP/1.1"" 200 0",AllTraffic/i-08ffdaaab186465db
1735549955791,"2024-12-30 09:12:35,229 [INFO ] pool-1-thread-3 ACCESS_LOG - /169.254.178.2:49624 ""GET /ping HTTP/1.1"" 200 0",AllTraffic/i-08ffdaaab186465db
1735549955791,"2024-12-30 09:12:35,657 [INFO ] W-9000-model com.amazonaws.ml.mms.wlm.WorkerThread - Backend response time: 1",AllTraffic/i-08ffdaaab186465db
1735549960303,"2024-12-30 09:12:35,657 [INFO ] W-9000-model ACCESS_LOG - /169.254.178.2:49624 ""POST /invocations HTTP/1.1"" 500 7",AllTraffic/i-08ffdaaab186465db
1735549964460,"2024-12-30 09:12:40,231 [INFO ] pool-1-thread-3 ACCESS_LOG - /169.254.178.2:44552 ""GET /ping HTTP/1.1"" 200 0",AllTraffic/i-08ffdaaab186465db
1735549969461,"2024-12-30 09:12:45,229 [INFO ] pool-1-thread-3 ACCESS_LOG - /169.254.178.2:44552 ""GET /ping HTTP/1.1"" 200 0",AllTraffic/i-08ffdaaab186465db
1735549974460,"2024-12-30 09:12:50,229 [INFO ] pool-1-thread-3 ACCESS_LOG - /169.254.178.2:44552 ""GET /ping HTTP/1.1"" 200 0",AllTraffic/i-08ffdaaab186465db
1735549979460,"2024-12-30 09:12:55,229 [INFO ] pool-1-thread-3 ACCESS_LOG - /169.254.178.2:44552 ""GET /ping HTTP/1.1"" 200 0",AllTraffic/i-08ffdaaab186465db
1735549984501,"2024-12-30 09:13:00,229 [INFO ] pool-1-thread-3 ACCESS_LOG - /169.254.178.2:44552 ""GET /ping HTTP/1.1"" 200 0",AllTraffic/i-08ffdaaab186465db
1735549989461,"2024-12-30 09:13:05,229 [INFO ] pool-1-thread-3 ACCESS_LOG - /169.254.178.2:44552 ""GET /ping HTTP/1.1"" 200 0",AllTraffic/i-08ffdaaab186465db
1735549994460,"2024-12-30 09:13:10,229 [INFO ] pool-1-thread-3 ACCESS_LOG - /169.254.178.2:44552 ""GET /ping HTTP/1.1"" 200 0",AllTraffic/i-08ffdaaab186465db
1735549999502,"2024-12-30 09:13:15,229 [INFO ] pool-1-thread-3 ACCESS_LOG - /169.254.178.2:44552 ""GET /ping HTTP/1.1"" 200 0",AllTraffic/i-08ffdaaab186465db
1735550004461,"2024-12-30 09:13:20,229 [INFO ] pool-1-thread-3 ACCESS_LOG - /169.254.178.2:44552 ""GET /ping HTTP/1.1"" 200 0",AllTraffic/i-08ffdaaab186465db
1735550009460,"2024-12-30 09:13:25,229 [INFO ] pool-1-thread-3 ACCESS_LOG - /169.254.178.2:44552 ""GET /ping HTTP/1.1"" 200 0",AllTraffic/i-08ffdaaab186465db
1735550014501,"2024-12-30 09:13:30,229 [INFO ] pool-1-thread-3 ACCESS_LOG - /169.254.178.2:44552 ""GET /ping HTTP/1.1"" 200 0",AllTraffic/i-08ffdaaab186465db
1735550019461,"2024-12-30 09:13:35,229 [INFO ] pool-1-thread-3 ACCESS_LOG - /169.254.178.2:44552 ""GET /ping HTTP/1.1"" 200 0",AllTraffic/i-08ffdaaab186465db
1735550024461,"2024-12-30 09:13:40,232 [INFO ] pool-1-thread-3 ACCESS_LOG - /169.254.178.2:44552 ""GET /ping HTTP/1.1"" 200 0",AllTraffic/i-08ffdaaab186465db
1735550030235,"2024-12-30 09:13:45,229 [INFO ] pool-1-thread-3 ACCESS_LOG - /169.254.178.2:44552 ""GET /ping HTTP/1.1"" 200 0",AllTraffic/i-08ffdaaab186465db
1735550034460,"2024-12-30 09:13:50,229 [INFO ] pool-1-thread-3 ACCESS_LOG - /169.254.178.2:44552 ""GET /ping HTTP/1.1"" 200 0",AllTraffic/i-08ffdaaab186465db
1735550039461,"2024-12-30 09:13:55,229 [INFO ] pool-1-thread-3 ACCESS_LOG - /169.254.178.2:44552 ""GET /ping HTTP/1.1"" 200 0",AllTraffic/i-08ffdaaab186465db
1735550044461,"2024-12-30 09:14:00,229 [INFO ] pool-1-thread-3 ACCESS_LOG - /169.254.178.2:44552 ""GET /ping HTTP/1.1"" 200 0",AllTraffic/i-08ffdaaab186465db

This is my notebook code:

model_b_data = "s3://recommendation/model_b.tar.gz"
data = {
    "users": [
        {
            "user_id": 57845789,
            "age": 25,
            "gender": "female",
            "location": "Birmingham",
            "user_lat": 52.4823,
            "user_long": 1.8900,
            "products": [
                {
                    "product_id": 997,
                    "title": "Free movie Ticket Tumbbad",
                    "brand": "NaN",
                    "description": "Free movie Ticket Tumbbad.",
                    "categories": ["Event", "Film"],
                    "final_price": 99.99,
                    "original_price": 149.99,
                    "offer_discount": 33.33,
                    "rating": 5,
                    "interaction": {
                        "interaction_type": "viewed",
                        "timestamp": "2024-12-01T12:00:00Z",
                        "rating": 5
                    }
                },
                {
                    "product_id": 338,
                    "title": "CBZ",
                    "brand": "NaN",
                    "description": "CBZ",
                    "categories": ["Other", "Baby", "Beauty"],
                    "final_price": 99.99,
                    "original_price": 149.99,
                    "offer_discount": 33.33,
                    "rating": 5,
                    "interaction": {
                        "interaction_type": "viewed",
                        "timestamp": "2024-12-10T09:00:00Z",
                        "rating": 5
                    }
                }
            ],
            "avg_sessions_per_day": 3.5,
            "return_rate": 0.75,
            "guild": ["Technology", "Beauty"],
            "tribe": ["Loki"],
            "user_categories": ["Shopping", "Fashion", "Women"],
            "user_geohash": "gcpchgcm13fs"
        },
        {
            "user_id": 578457810,
            "age": 30,
            "gender": "male",
            "location": "Westminster",
            "user_lat": 51.4975,
            "user_long": 0.1357,
            "products": [
                {
                    "product_id": 338,
                    "title": "CBZ",
                    "brand": "NaN",
                    "description": "CBZ",
                    "categories": ["Other", "Baby", "Beauty"],
                    "final_price": 99.99,
                    "original_price": 149.99,
                    "offer_discount": 33.33,
                    "rating": 5,
                    "interaction": {
                        "interaction_type": "viewed",
                        "timestamp": "2024-12-10T09:00:00Z",
                        "rating": 5
                    }
                }
            ],
            "avg_sessions_per_day": 2.0,
            "return_rate": 0.50,
            "guild": ["Technology"],
            "tribe": ["Odin"],
            "user_categories": ["Shopping"],
            "user_geohash": "wtw1k99rrxx7"
        }
    ]
}
pytorch_model = PyTorchModel(
    model_data=model_b_data,
    role=role,
    entry_point="pytorch_script/inference.py",
    framework_version="1.5.0",
    py_version="py3",
)
from sagemaker.serializers import JSONSerializer
from sagemaker.deserializers import JSONDeserializer

local_mode = False

if local_mode:
    instance_type = "local"
else:
    instance_type = 'ml.m5.xlarge'

predictor = pytorch_model.deploy(
    initial_instance_count=1,
    instance_type=instance_type,
    serializer=JSONSerializer(),
    deserializer=JSONDeserializer(),
)
predictor.predict(data)

This is my inference code:

logger = logging.getLogger(__name__)
logger.setLevel(logging.INFO)

s3_client = boto3.client('s3')

COLUMNS = [...]

def install_package(package_name):
    """Install a package using pip."""
    try:
        subprocess.check_call([sys.executable, "-m", "pip", "install", package_name])
        logger.info(f"Package '{package_name}' installed successfully.")
    except subprocess.CalledProcessError as e:
        logger.error(f"Failed to install package '{package_name}'. Error: {e}")


required_packages = ["boto3", "numpy", "transformers", "nltk", "pandas", "geohash2"]
for package in required_packages:
    install_package(package)


def load_file(file_name, from_s3=True, bucket_name="recommendation-tara"):
    """Load a file from S3 bucket."""
    if from_s3:
        try:
            logger.info(f"Attempting to load file: {file_name} from S3.")
            obj = s3_client.get_object(Bucket=bucket_name, Key=file_name)
            file_content = obj['Body'].read()
            logger.info(f"Successfully loaded file: {file_name}.")

            if file_name.endswith(('.pt', '.pth')):
                buffer = BytesIO(file_content)
                return torch.load(buffer, map_location='cpu')

            elif file_name.endswith('.csv'):
                return pd.read_csv(BytesIO(file_content))

            elif file_name.endswith('.parquet'):
                return pd.read_parquet(BytesIO(file_content))

            elif file_name.endswith('.json'):
                return json.loads(file_content)

            elif file_name.endswith('.joblib'):
                buffer = BytesIO(file_content)
                return joblib.load(buffer)

            elif file_name.endswith(('.pkl', '.pickle')):
                buffer = BytesIO(file_content)
                return torch.load(buffer, weights_only=False)

            else:
                logger.warning(f"Unsupported file type for '{file_name}'.")
                return None

        except Exception as e:
            logger.error(f"Error loading file from S3: {e}")
            return None

    logger.error(f"Unsupported source for '{file_name}'.")
    return None


def process_text_and_get_embeddings(text_series):
    """Process text and return embeddings as a tensor."""
    clean_texts = [process_text(text) for text in text_series]
    embeddings = [sentence_embedding(text) for text in clean_texts]
    logger.info(f"Processed {len(text_series)} texts to get embeddings.")
    return torch.tensor(np.array(embeddings), dtype=torch.float32).view(-1, 768)


def pad_tensor(tensor, target_columns):
    """Pad the tensor to ensure it has the target number of columns."""
    if tensor.shape[1] != target_columns:
        logger.info(f"Padding tensor from {tensor.shape[1]} to {target_columns} columns.")
        return torch.nn.functional.pad(tensor, (0, target_columns - tensor.shape[1]), "constant", 0)
    return tensor


def create_product_tensor(data, encoder):
    products_interacted = data["products"]
    products = pd.json_normalize(products_interacted)
    products = products[['id', 'title', 'description', 'categories']]
    products['categories'] = products['categories'].astype(str)

    products[['level1', 'level2', 'level3']] = products['categories'].str.split(',', n=2, expand=True)
    products[['level1', 'level2', 'level3']] = products[['level1', 'level2', 'level3']].fillna('empty')

    title_embeddings = process_text_and_get_embeddings(products['title'])
    description_embeddings = process_text_and_get_embeddings(products['description'])

    logger.info(f"Creating product tensor with {len(products)} products.")
    encoding_df = products[['id', 'title', 'description', 'categories', 'level1', 'level2', 'level3']]
    encoded_sentence = encoder.transform(encoding_df).drop(columns=['id', 'title', 'description', 'categories'])
    categories_tensor = torch.tensor(encoded_sentence.fillna(0).values, dtype=torch.float32)

    combined_embeddings = torch.cat((title_embeddings, description_embeddings, categories_tensor), dim=1)
    target_columns = 2323
    return pad_tensor(combined_embeddings, target_columns)


def reduce_products_tensor(data_tensor, product_pca, product_scaler):
    logger.info(f"Reducing product tensor with shape: {data_tensor.shape}.")
    data_scaled = product_scaler.transform(data_tensor.numpy())
    X_reduced = product_pca.transform(data_scaled)
    return torch.tensor(X_reduced, dtype=torch.float32)


def create_and_reduce_user_tensor(data, user_pca, user_scaler, columns):
    new_users = data['users']
    events_subset = pd.json_normalize(new_users)

    events_subset = events_subset[['user_id', 'user_categories']]
    events_subset['user_categories'] = events_subset['user_categories'].apply(
        lambda x: ', '.join(x) if isinstance(x, list) else x
    )
    
    category_encoded = multilabel_encoding(events_subset, 'user_categories')
    events_encoded = pd.concat([events_subset['user_id'], category_encoded], axis=1)
    events_encoded.columns = events_encoded.columns.str.strip().str.lower()
    events_encoded['user_id'] = events_encoded['user_id'].astype(int)
    events_encoded.fillna(0, inplace=True)

    result_df = pd.DataFrame(columns=columns)
    result_df['user_id'] = events_encoded['user_id'].values
    for col in columns:
        result_df[col] = events_encoded.get(col, 0.0)

    data_numpy = result_df.values
    logger.info(f"Creating user tensor with shape: {data_numpy.shape}.")
    data_scaled = user_scaler.transform(data_numpy)
    X_reduced = user_pca.transform(data_scaled)
    return X_reduced


def update_interactions_like(new_data, existing_interactions):
    existing_interactions = existing_interactions.drop(columns='Unnamed: 0', errors='ignore')
    
    extracted_data = []

    for user in new_data["users"]:
        user_id = user["user_id"]
        
        for product in user["products"]:
            product_id = product["product_id"]
            
            extracted_data.append({
                "user_id": user_id,
                "product_id": product_id,
                "liked": 1
            })

    new_interactions = pd.DataFrame(extracted_data)

    interactions_combined = pd.concat([existing_interactions, new_interactions], ignore_index=True)

    return interactions_combined


# def model_fn(model_dir):
#     model = TwoTowerModel(user_embedding_size=167, product_embedding_size=273, dropout_rate=0.3008733795102767)
#     model.load_state_dict(load_file('model_final_b12.pth'))
#     model.eval()
#     return model

def model_fn(model_dir): 
    model = TwoTowerModel(user_embedding_size=167, product_embedding_size=273, dropout_rate=0.3008733795102767)
    with open(os.path.join(model_dir, "model.pth"), "rb") as f:
        model.load_state_dict(torch.load(f))
    model.to(device).eval()
    return model


def input_fn(request_body, request_content_type):
    assert request_content_type == 'application/json'
    data = json.loads(request_body)
    user_data = create_and_reduce_user_tensor(data, user_pca, user_scaler, COLUMNS)
    product_data = reduce_products_tensor(create_product_tensor(data, encoder), product_pca, product_scaler)
    return user_data, product_data


def predict_fn(input_object, model):
    with torch.no_grad():
        prediction = model.predict(input_object)
    return prediction


def output_fn(predictions, content_type):
    assert content_type == 'application/json'
    res = predictions.cpu().numpy().tolist()
    return json.dumps(res)


interactions = load_file('interactions_like.csv')
encoder = load_file('binary_encoder.joblib')
product_scaler = load_file('scaler_product_model.joblib')
product_pca = load_file('pca_product_model.joblib')
user_pca = load_file('pca_user_model.joblib')
user_scaler = load_file('scaler_user_model.joblib')

I would appreciate any help - is there anything wrong with the _fn functions?