Entrypoint for testing

Dockerfile

@@ -57,6 +57,9 @@ ENV MODEL_NAME=$MODEL_NAME \
     MODEL_BASE_PATH=$MODEL_BASE_PATH \
     HUGGING_FACE_HUB_TOKEN=$HUGGING_FACE_HUB_TOKEN
 
+# Set the entrypoint
+ENTRYPOINT ["/entrypoint.sh"]
+
 # Run the Python script to download the model
 RUN python -u /download_model.py
 

entrypoint.sh

@@ -0,0 +1,26 @@
+#!/bin/bash
+
+# Export specific ENV variables to /etc/rp_environment
+echo "Exporting environment variables..."
+printenv | grep -E '^RUNPOD_|^PATH=|^_=' | sed 's/^\(.*\)=\(.*\)$/export \1="\2"/' >> /etc/rp_environment
+echo 'source /etc/rp_environment' >> ~/.bashrc
+
+if [[ $PUBLIC_KEY ]]
+then
+    mkdir -p ~/.ssh
+    chmod 700 ~/.ssh
+    echo $PUBLIC_KEY >> ~/.ssh/authorized_keys
+    chmod 700 -R ~/.ssh
+    # Start the SSH service in the background
+    service ssh start
+else
+    echo "No PUBLIC_KEY ENV variable provided, not starting openSSH daemon"
+fi
+
+# Conditional execution based on TESTING
+if [ "$TESTING" = "True" ]; then
+    sleep infinity
+else
+    # Execute the passed arguments (CMD)
+    exec "$@"
+fi