[Python] AWS SageMaker - Custom Inference With HuggingFace Model

For context, I'm currently working in a JupyterLab space in SageMaker studio.

My goal is to deploy a HuggingFace Llama model for batch transform inference. The data I will be passing in to the LLM is in an S3 bucket in the same AWS account. The data does require some custom handling (changing its format to JSON, wrapping it in my LLM prompt, etc.) so I need a custom inference script for the model.

Unfortunately, no matter what I try, it seems the custom inference script just is not running. Also, the notebook from which the code is being run and inference.py are in the same subdirectory. It looks something like this:

~
|---- user-default-efs
----------------|---- notebook.ipynb
----------------|---- inference.py

Code:

import json
import sagemaker
from sagemaker.huggingface import HuggingFaceModel, get_huggingface_llm_image_uri

hub = {
"HF_MODEL_ID": "{model_here}",
"SM_NUM_GPUS": json.dumps(1),
"HF_TOKEN": "{api_key_here}"
}
instance_type, instance_count = "{ec2_instance_here}", 1

hf_model = HuggingFaceModel(
image_uri=get_huggingface_llm_image_uri("huggingface", version="3.2.3"),
env=hub,
role=sagemaker.get_execution_role(),
entry_point="inference.py"
)

predictor = hf_model.deploy(
initial_instance_count=instance_count,
instance_type=instance_type
)

payload = {
"inputs": "hello world!",
"parameters": {
"max_new_tokens": 64,
"temperature": 0.5,
"do_sample": True
}
}
response = predictor.predict(payload)
print(response)


Note that I do eventually want batch transform inference, but for the purposes of testing I have a standard endpoint set up for now.

And here's my inference.py:

def input_fn(input_data, content_type):
print(f"{input_data=}")
print(f"{content_type=}")
return input_data

def predict_fn(processed_data, model):
print(f"{processed_data=}")
print(f"{model=}")
return "This is a test prediction."

def output_fn(prediction, accept):
print(f"{prediction=}")
print(f"{accept=}")
return "This is a test output.", accept


As you can see, there's no substance to inference.py at the moment - it's only for testing if the script is actually running.

Results:
I would expect response to print "This is a test output.". My understanding is that inference.py should be run and completely replace the call to the LLM. This does not happen. Inference runs normally, and the LLM outputs a normal response, as if inference.py does not exist.

I would also expect the print statements to show up somewhere, whether that be in the JupyterLab notebook output or CloudWatch logs. It does not show up in either of those two places, leading me to believe inference.py is not even running.

Is anyone able to help with this issue? I've been struggling with it for many days, so any help would be greatly appreciated. Thank you so much in advance!

Continue reading...