diff --git a/.gitignore b/.gitignore index 9ad40cf4..fcd341af 100644 --- a/.gitignore +++ b/.gitignore @@ -41,6 +41,7 @@ share/python-wheels/ .installed.cfg *.egg MANIFEST +.setup_done # PyInstaller # Usually these files are written by a python script from a template diff --git a/Dockerfile b/Dockerfile new file mode 100644 index 00000000..1fd3f3d8 --- /dev/null +++ b/Dockerfile @@ -0,0 +1,43 @@ +FROM pytorch/pytorch:2.6.0-cuda12.6-cudnn9-runtime + +ENV DEBIAN_FRONTEND=noninteractive \ + PYTHONUNBUFFERED=1 \ + LC_ALL=C.UTF-8 \ + LANG=C.UTF-8 + +# Install Git, OpenSSH Server, and OpenGL (PyTorch's base image already includes Conda and Pip) +RUN apt-get update && apt-get install -y \ + git \ + openssh-server \ + libgl1 \ + && rm -rf /var/lib/apt/lists/* + +WORKDIR /workspace + +# Create an empty README.md file and an empty torch_uncertainty module to satisfy flit +RUN touch README.md && mkdir -p torch_uncertainty && touch torch_uncertainty/__init__.py + +# Copy dependency file +COPY pyproject.toml /workspace/ + +# Install dependencies all dependencies +RUN pip install --no-cache-dir -e ".[all]" + +# Always activate Conda when opening a new terminal +RUN echo "source /opt/conda/bin/activate" >> /root/.bashrc + +# Configure SSH server +RUN echo "PermitRootLogin yes" >> /etc/ssh/sshd_config && \ + echo "PubkeyAuthentication yes" >> /etc/ssh/sshd_config && \ + echo "AuthorizedKeysFile .ssh/authorized_keys" >> /etc/ssh/sshd_config + +# Expose port 8888 for TensorBoard and Jupyter Notebook and port 22 for SSH +EXPOSE 8888 22 + +# Entrypoint script (runs every time the container starts) +COPY docker/entrypoint.sh /usr/local/bin/ +RUN chmod +x /usr/local/bin/entrypoint.sh +ENTRYPOINT ["/usr/local/bin/entrypoint.sh"] + +# Note that if /workspace/ is a mounted volume, any files copied to /workspace/ during the build will be overwritten by the mounted volume +# This is why we copy the entrypoint script to /usr/local/bin/ instead of /workspace/ diff --git a/README.md b/README.md index 52b0ed8e..e0f735a4 100644 --- a/README.md +++ b/README.md @@ -43,6 +43,11 @@ pip install torch-uncertainty The installation procedure for contributors is different: have a look at the [contribution page](https://torch-uncertainty.github.io/contributing.html). +### :whale: Docker image for contributors + +For contributors who want to run experiments on cloud GPU instances, we provide a pre-built Docker image that includes all necessary dependencies and configurations and the Dockerfile for building your custom Docker images. +This allows you to quickly launch an experiment-ready container with minimal setup. Please refer to [DOCKER.md](docker/DOCKER.md) for further details. + ## :racehorse: Quickstart We make a quickstart available at [torch-uncertainty.github.io/quickstart](https://torch-uncertainty.github.io/quickstart.html). @@ -94,4 +99,4 @@ The following projects use TorchUncertainty: - *A Symmetry-Aware Exploration of Bayesian Neural Network Posteriors* - [ICLR 2024](https://arxiv.org/abs/2310.08287) -**If you are using TorchUncertainty in your project, please let us know, we will add your project to this list!** +**If you are using TorchUncertainty in your project, please let us know, and we will add your project to this list!** diff --git a/docker/DOCKER.md b/docker/DOCKER.md new file mode 100644 index 00000000..1b53beb4 --- /dev/null +++ b/docker/DOCKER.md @@ -0,0 +1,69 @@ +# :whale: Docker image for contributors +### Pre-built Docker image +1. To pull the pre-built image from Docker Hub, simply run: + ```bash + docker pull docker.io/tonyzamyatin/torch-uncertainty:latest + ``` + + This image includes: + - PyTorch with CUDA support + - OpenGL (for visualization tasks) + - Git, OpenSSH, and all Python dependencies + + Checkout the [registry on Docker Hub](https://hub.docker.com/repository/docker/tonyzamyatin/torch-uncertainty/general) for all available images. + +2. To start a container using this image, set up the necessary environment variables and run: + ```bash + docker run --rm -it --gpus all -p 8888:8888 -p 22:22 \ + -e VM_SSH_PUBLIC_KEY="your-public-key" \ + -e GITHUB_SSH_PRIVATE_KEY="your-github-key" \ + -e GITHUB_USER="your-github-username" \ + -e GIT_USER_EMAIL="your-git-email" \ + -e GIT_USER_NAME="your-git-name" \ + docker.io/tonyzamyatin/torch-uncertainty + ``` + + Optionally, you can also set `-e USER_COMPACT_SHELL_PROMPT="true"` + to make the VM's shell prompts compact and colorized. + + **Note:** Some cloud providers offer templates, in which you can preconfigure + in advance which Docker image to pull and which environment variables to set. + In this case, the provider will pull the image, set all environment variables, + and start the container for you. + +3. Once your cloud provider has deployed the VM, it will display the host address and SSH port. + You can connect to the container via SSH using: + ```bash + ssh -i /path/to/private_key root@ -p + ``` + + Replace `` and `` with the values provided by your cloud provider, + and `/path/to/private_key` with the private key that corresponds to `VM_SSH_PUBLIC_KEY`. + +4. The container exposes port `8888` in case you want to run Jupyter Notebooks or TensorBoard. + + **Note:** The `/workspace` directory is mounted from your local machine or cloud storage, + so changes persist across container restarts. + If using a cloud provider, ensure your network volume is correctly attached to avoid losing data. + +### Modifying and publishing custom Docker image +If you want to make changes to the Dockerfile, follow these steps: +1. Edit the Dockerfile to fit your needs. + +2. Build the modified image: + ``` + docker build -t my-custom-image . + ``` + +3. Push to a Docker registry (if you want to use it on another VM): + ``` + docker tag my-custom-image mydockerhubuser/my-custom-image:tag + docker push mydockerhubuser/my-custom-image:tag + ``` + +4. Pull the custom image onto your VM: + ``` + docker pull mydockerhubuser/my-custom-image + ``` + +5. Run the container using the same docker run command with the new image name. diff --git a/docker/entrypoint.sh b/docker/entrypoint.sh new file mode 100644 index 00000000..1c1af5e5 --- /dev/null +++ b/docker/entrypoint.sh @@ -0,0 +1,131 @@ +#!/bin/bash +set -e # Exit immediately if a command fails + +echo "🚀 Starting container..." + +# Ensure SSH directory exists and has correct permissions +if [ ! -d /root/.ssh ]; then + echo "📂 Creating SSH directory..." + mkdir -p /root/.ssh && chmod 700 /root/.ssh +fi + +# Ensure the VM's public SSH key is added for authentication +if [ -z "$VM_SSH_PUBLIC_KEY" ]; then + echo "❌ Error: Please set the VM_SSH_PUBLIC_KEY environment variable." + exit 1 +fi +if [ ! -f /root/.ssh/authorized_keys ] || ! grep -q "$VM_SSH_PUBLIC_KEY" /root/.ssh/authorized_keys; then + echo "🔑 Adding VM SSH public key..." + echo "$VM_SSH_PUBLIC_KEY" > /root/.ssh/authorized_keys && chmod 600 /root/.ssh/authorized_keys +fi + +# Ensure GitHub SSH private key is set up for authentication +if [ -z "$GITHUB_SSH_PRIVATE_KEY" ]; then + echo "❌ Error: Please set the GITHUB_SSH_PRIVATE_KEY environment variable." + exit 1 +fi +if [ ! -f /root/.ssh/github_rsa ]; then + echo "🔐 Adding GitHub SSH private key..." + echo "$GITHUB_SSH_PRIVATE_KEY" > /root/.ssh/github_rsa && chmod 600 /root/.ssh/github_rsa +fi + +# Configure SSH client for GitHub authentication +if [ ! -f /root/.ssh/config ] || ! grep -q 'Host github.com' /root/.ssh/config; then + echo "⚙️ Configuring SSH client for GitHub authentication..." + cat < /root/.ssh/config +Host github.com + User git + IdentityFile /root/.ssh/github_rsa +EOF + chmod 600 /root/.ssh/config +fi + +# Add GitHub to known hosts (to avoid SSH verification prompts) +echo "📌 Ensuring GitHub is a known host..." +ssh-keygen -F github.com > /dev/null 2>&1 || ssh-keyscan github.com >> /root/.ssh/known_hosts + +# Start SSH agent and add GitHub private key (if not already added) +if ! pgrep -x "ssh-agent" > /dev/null; then + echo "🕵️ Starting SSH agent..." + eval "$(ssh-agent -s)" +fi +if ssh-add -l | grep -q github_rsa; then + echo "✅ GitHub SSH key already added." +else + echo "🔑 Adding GitHub SSH key to agent..." + ssh-add /root/.ssh/github_rsa +fi + +# Set Git user name and email (if provided) +if [ -n "$GIT_USER_NAME" ]; then + echo "👤 Setting Git username: $GIT_USER_NAME" + git config --global user.name "$GIT_USER_NAME" +fi +if [ -n "$GIT_USER_EMAIL" ]; then + echo "📧 Setting Git email: $GIT_USER_EMAIL" + git config --global user.email "$GIT_USER_EMAIL" +fi + +# Ensure first-time setup runs only once +if [ ! -f /workspace/.setup_done ]; then + echo "🛠️ Running first-time setup..." + + # Ensure GitHub username is set + if [ -z "$GITHUB_USER" ]; then + echo "❌ Error: Please set the GITHUB_USER environment variable." + exit 1 + fi + + # Clone GitHub repo if not already cloned + if [ ! -d "/workspace/.git" ]; then + echo "📦 Cloning repository: $GITHUB_USER/torch-uncertainty..." + git clone git@github.com:$GITHUB_USER/torch-uncertainty.git /workspace + fi + + # Mark setup as completed + touch /workspace/.setup_done + echo "✅ First-time setup complete!" +else + echo "⏩ Skipping first-time setup (already done)." +fi + +# Apply compact shell prompt customization (if enabled) +if [ -n "$USE_COMPACT_SHELL_PROMPT" ]; then + echo "🎨 Applying compact shell prompt customization..." + echo 'force_color_prompt=yes' >> /root/.bashrc + echo 'PS1="\[\033[01;34m\]\W\[\033[00m\]\$ "' >> /root/.bashrc + echo 'if [ -x /usr/bin/dircolors ]; then' >> /root/.bashrc + echo ' test -r ~/.dircolors && eval "$(dircolors -b ~/.dircolors)" || eval "$(dircolors -b)"' >> /root/.bashrc + echo ' alias ls="ls --color=auto"' >> /root/.bashrc + echo ' alias grep="grep --color=auto"' >> /root/.bashrc + echo ' alias fgrep="fgrep --color=auto"' >> /root/.bashrc + echo ' alias egrep="egrep --color=auto"' >> /root/.bashrc + echo 'fi' >> /root/.bashrc +fi + +# Ensure /workspace is in PYTHONPATH +if ! echo "$PYTHONPATH" | grep -q "/workspace"; then + echo "📌 Adding /workspace to PYTHONPATH" + export PYTHONPATH="/workspace:$PYTHONPATH" +else + echo "✅ PYTHONPATH is already correctly set." +fi + +# Check if torch_uncertainty is installed in editable mode +if pip show torch_uncertainty | grep -q "Editable project location: /workspace"; then + echo "✅ torch_uncertainty is already installed in editable mode. 🎉" +else + echo "🔄 Reinstalling torch_uncertainty in editable mode..." + pip uninstall -y torch-uncertainty + pip install -e /workspace + echo "✅ torch_uncertainty is now installed in editable mode! 🚀" +fi + +# Activate pre-commit hooks (if enabled) +echo "🔗 Activating pre-commit hooks..." +pre-commit install + +# Ensure SSH server is started +echo "🔑 Starting SSH server..." +mkdir -p /run/sshd && chmod 755 /run/sshd +/usr/sbin/sshd -D