Not able to install 'pycuda' on HuggingFace container

Amazon SageMaker
1

Hi, I am using HuggingFace SageMaker container for ‘token-classification’ task.
I have fine tuned ‘Bert-base-cased’ model and converted it to onnx format and then to tensorrt engine.
To run inference using tensorrt engine, I want to use PyCuda library.
So I have added PyCuda, tensorrt to requirements.txt file (and custom inference code in inference.py)
But pycuda installation is failing.
Please see the below code and error.

Code:-

from sagemaker.huggingface import HuggingFaceModel

hub = {
‘HF_TASK’: ‘token-classification’
}

huggingface_model = HuggingFaceModel(
model_data = model_artifacts,
transformers_version = ‘4.6.1’,
tensorflow_version = ‘2.4.1’,
py_version = ‘py37’,
env = hub
)

predictor = huggingface_model.deploy(
initial_instance_count = 1,
instance_type = ‘ml.p3.2xlarge’
)

requirements.txt:-

–extra-index-url https://pypi.ngc.nvidia.com
nvidia-tensorrt==8.4.1.5
pycuda==2020.1

But pycuda installation is failing with the below error…

Pycuda installation error:-

In file included from src/cpp/cuda.cpp:4:0:
src/cpp/cuda.hpp:14:10: fatal error: cuda.h: No such file or directory
#include <cuda.h>
^~~~~~~~
compilation terminated.
error: command ‘gcc’ failed with exit status 1

ERROR: Failed building wheel for pycuda

This could be happening due to below two reasons…

  1. PATH variable is not properly set as below
    export PATH=“/usr/local/cuda/bin:$PATH”
    export LD_LIBRARY_PATH=“/usr/local/cuda/lib64:$LD_LIBRARY_PATH”

  2. or CUDA toolkit is not installed

Need help in resolving this issue…