only pin the aws-neuronx-dkms package version to 2.17.17.0 for al2inf

aws · Sep 20, 2024 · feaf3bd · feaf3bd
1 parent 36b687d
commit feaf3bd
Showing 1 changed file with 6 additions and 2 deletions.
diff --git a/scripts/enable-ecs-agent-inferentia-support.sh b/scripts/enable-ecs-agent-inferentia-support.sh
@@ -27,8 +27,12 @@ sudo mv /tmp/neuron.repo /etc/yum.repos.d/neuron.repo
 sudo yum install kernel-devel-$(uname -r) kernel-headers-$(uname -r) -y
 
 # Install Neuron Driver
-# Pin the aws-neuronx-dkms package version to 2.17.17.0, since the newest versions of the Neuron SDK are no longer supporting linux kernel 4.14
-sudo yum install -y aws-neuronx-dkms-2.17.17.0
+if [[ $AMI_TYPE == "al2inf" ]]; then
+    # Pin the aws-neuronx-dkms package version to 2.17.17.0 only for al2inf, since the newest versions of the Neuron SDK are no longer supporting linux kernel 4.14
+    sudo yum install -y aws-neuronx-dkms-2.17.17.0
+else
+    sudo yum install -y aws-neuronx-dkms-2.*
+fi
 sudo yum install -y aws-neuronx-oci-hook-2.*
 
 # Install oci-add-hooks