3 files changed, 43 insertions, 2 deletions

diff --git a/docker_build/Dockerfile b/docker_build/Dockerfile
index 5df57d2..16f74d6 100644
--- a/docker_build/Dockerfile
+++ b/docker_build/Dockerfile
@@ -19,8 +19,8 @@ RUN mkdir -p /var/run/sshd && \
 RUN mkdir -p /root/.ssh && \
     chmod 700 /root/.ssh && \
     usermod -s /bin/bash root
-# Copy SSH public keys from ssh-keys directory into the image
-COPY ssh-keys/*.pub /tmp/ssh-keys/
+# Copy SSH public keys from docker_build/ssh-keys directory into the image
+COPY docker_build/ssh-keys/*.pub /tmp/ssh-keys/
 RUN cat /tmp/ssh-keys/*.pub > /root/.ssh/authorized_keys && \
     chmod 600 /root/.ssh/authorized_keys && \
     rm -rf /tmp/ssh-keys
diff --git a/docker_build/Dockerfile.gpu b/docker_build/Dockerfile.gpu
new file mode 100644
index 0000000..7ed08a5
--- /dev/null
+++ b/docker_build/Dockerfile.gpu
@@ -0,0 +1,40 @@
+# Multi-stage build - GPU version builds on top of the base dev environment
+FROM rocky_dev:latest
+
+# Update and install GPU-specific packages
+RUN dnf update -y && \
+    dnf install -y kernel-headers kernel-devel pciutils && \
+    dnf clean all
+
+# Install NVIDIA container toolkit dependencies
+RUN dnf config-manager --add-repo https://nvidia.github.io/libnvidia-container/stable/rpm/nvidia-container-toolkit.repo && \
+    dnf install -y nvidia-container-toolkit && \
+    dnf clean all
+
+# Set environment variables for NVIDIA
+ENV NVIDIA_VISIBLE_DEVICES=all
+ENV NVIDIA_DRIVER_CAPABILITIES=compute,utility
+
+# Add GPU test script
+RUN echo '#!/bin/bash' > /usr/local/bin/gpu-test.sh && \
+    echo 'echo "=== System Information ==="' >> /usr/local/bin/gpu-test.sh && \
+    echo 'cat /etc/rocky-release' >> /usr/local/bin/gpu-test.sh && \
+    echo 'echo' >> /usr/local/bin/gpu-test.sh && \
+    echo 'echo "=== PCI Devices (GPUs) ==="' >> /usr/local/bin/gpu-test.sh && \
+    echo 'lspci | grep -i nvidia' >> /usr/local/bin/gpu-test.sh && \
+    echo 'echo' >> /usr/local/bin/gpu-test.sh && \
+    echo 'echo "=== NVIDIA SMI ==="' >> /usr/local/bin/gpu-test.sh && \
+    echo 'if command -v nvidia-smi &> /dev/null; then' >> /usr/local/bin/gpu-test.sh && \
+    echo '    nvidia-smi' >> /usr/local/bin/gpu-test.sh && \
+    echo 'else' >> /usr/local/bin/gpu-test.sh && \
+    echo '    echo "nvidia-smi not found. GPU might not be accessible inside container."' >> /usr/local/bin/gpu-test.sh && \
+    echo 'fi' >> /usr/local/bin/gpu-test.sh && \
+    chmod +x /usr/local/bin/gpu-test.sh
+
+# Create workspace directory for GPU workloads
+RUN mkdir -p /workspace
+
+# Keep the same working directory and CMD from base image
+WORKDIR /root
+EXPOSE 22
+CMD ["/usr/sbin/sshd", "-D", "-e"]
\ No newline at end of file
diff --git a/docker_build/ssh-keys/macm4-resident.pub b/docker_build/ssh-keys/macm4-resident.pub
new file mode 100644
index 0000000..fbccb4f
--- /dev/null
+++ b/docker_build/ssh-keys/macm4-resident.pub
@@ -0,0 +1 @@
+sk-ssh-ed25519@openssh.com AAAAGnNrLXNzaC1lZDI1NTE5QG9wZW5zc2guY29tAAAAIFdHP8n64jOV6Ok7U9TDnGW+LUkXP6V7cvXH6xqN0zcNAAAAEnNzaDptYWNtNC1yZXNpZGVudA== ssh:macm4-resident