From 729740df81dc574365a7bd70d69dc4de8ce87ffe Mon Sep 17 00:00:00 2001
From: Wing Lian <wing.lian@gmail.com>
Date: Thu, 18 Jan 2024 22:04:25 -0500
Subject: [PATCH] Dockerfile cloud ports (#1148)

* explicitly expose ports 8888 and 22

* support for SSH_KEY from latitude
---
 docker/Dockerfile-cloud     |  3 +++
 scripts/cloud-entrypoint.sh | 12 +++++++++++-
 2 files changed, 14 insertions(+), 1 deletion(-)

diff --git a/docker/Dockerfile-cloud b/docker/Dockerfile-cloud
index 16aa389145..a7ad00ff6e 100644
--- a/docker/Dockerfile-cloud
+++ b/docker/Dockerfile-cloud
@@ -7,6 +7,9 @@ ENV TRANSFORMERS_CACHE="/workspace/data/huggingface-cache/hub"
 ENV HF_HOME="/workspace/data/huggingface-cache/hub"
 ENV HF_HUB_ENABLE_HF_TRANSFER="1"
 
+EXPOSE 8888
+EXPOSE 22
+
 COPY scripts/cloud-entrypoint.sh /root/cloud-entrypoint.sh
 
 RUN pip install jupyterlab notebook && \
diff --git a/scripts/cloud-entrypoint.sh b/scripts/cloud-entrypoint.sh
index 21ee88b85a..f7190b1f5c 100755
--- a/scripts/cloud-entrypoint.sh
+++ b/scripts/cloud-entrypoint.sh
@@ -7,14 +7,24 @@ echo 'source /etc/rp_environment' >> ~/.bashrc
 
 if [[ $PUBLIC_KEY ]]
 then
+    # runpod
     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
+elif [[ $SSH_KEY ]]
+then
+    # latitude.sh
+    mkdir -p ~/.ssh
+    chmod 700 ~/.ssh
+    echo $SSH_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"
+    echo "No PUBLIC_KEY or SSH_KEY environment variable provided, not starting openSSH daemon"
 fi
 
 # Check if JUPYTER_PASSWORD is set and not empty