init add sampling codes / checkpoint

docker/Dockerfile

@@ -0,0 +1,52 @@
+ARG BASE_IMAGE=${BASE_IMAGE:-null}
+ARG PUBLIC_IMAGE=${PUBLIC_IMAGE:-null}
+FROM ${BASE_IMAGE} as public
+ARG PYTHON_VERSION=3.9
+
+COPY ./asset /tmp/build
+RUN apt update &&\
+    xargs apt install -y < /tmp/build/apt_packages.txt
+RUN pip install -r /tmp/build/requirements.txt \
+    && rm -rf /tmp/build
+
+##############################################################################################################
+
+FROM ${PUBLIC_IMAGE} as private
+ARG TORCH_HOME
+ARG _USER
+ARG _UID
+ARG _GID
+ARG PW
+ARG _HOME
+
+COPY ./asset /tmp/build
+RUN xargs apt install -y < /tmp/build/apt_packages.txt
+
+RUN pip install -r /tmp/build/requirements.txt
+
+# Option1: Using unencrypted password/ specifying password
+RUN usermod --password $(echo ${PW} | openssl passwd -1 -stdin) root
+RUN useradd -m ${_USER} -d ${_HOME} --uid=${_UID} -s /usr/bin/zsh && echo "${_USER}:${PW}" | chpasswd
+
+# make the color of zsh-autosuggestions right
+ENV TERM xterm-256color
+
+RUN cp -r /root/.oh-my-zsh ${_HOME} && chown ${_USER}:${_USER} -R ${_HOME}/.oh-my-zsh &&\
+    cp /root/.zshrc ${_HOME} && chown ${_USER}:${_USER} -R ${_HOME}/.zshrc &&\
+    cp /root/.tmux.conf ${_HOME} && chown ${_USER}:${_USER} -R ${_HOME}/.tmux.conf && \
+    cp -r /root/.tmux ${_HOME} && chown ${_USER}:${_USER} -R ${_HOME}/.tmux
+
+
+USER ${_UID}:${_GID}
+WORKDIR ${_HOME}
+
+ENV TORCH_HOME ${TORCH_HOME}
+
+COPY --chown=${_USER}:${_USER} ./asset/ssh .ssh
+RUN chmod 0700 .ssh && chmod 600 .ssh/id_rsa && chmod 644 .ssh/id_rsa.pub \
+    &&ssh-keyscan github.com >> .ssh/known_hosts \
+    && cp /tmp/build/gitconfig ${_HOME}/.gitconfig 
+
+CMD "zsh"
+
+

docker/Makefile

@@ -0,0 +1,129 @@
+SHELL=/bin/bash
+BASE_IMAGE:=jerrikeph/base-cuda1162-dev-ubuntu2004-py39-torch-rdkit2023.03.2-openbabel
+
+PROJECT_NAME:=geobfn
+PUBLIC_IMAG:=${PROJECT_NAME}
+HTTPS_PROXY:=${https_proxy}
+HTTP_PROXY:=${http_proxy}
+_USER := $(shell whoami)
+_UID := $(shell id -u)
+_GID := $(shell id -g)
+PW:=rootpass
+
+LOCAL_NAME:=exp/${_USER}_${PROJECT_NAME}
+PRIVATE_IMAGE:=${LOCAL_NAME}
+CONTRAINER_NAME=$(subst /,_,$(LOCAL_NAME))
+
+PROJ_DIR:=${HOME}
+SHARE_DIR:=${HOME}/share
+TORCH_HOME:=${SHARE_DIR}/torch_home
+
+CHOME:= ${HOME}/home
+TO_SHARE_DIR:=${SHARE_DIR}
+TO_PROJ_DIR:=${CHOME}
+TO_TORCH_HOME:=${TO_SHARE_DIR}/torch_home
+
+MAKEFILE_ := $(CURDIR)/$(firstword $(MAKEFILE_LIST))
+
+run: _check_image _instantiate_container _exec
+
+run_root: _check_image _instantiate_container _exec_root
+
+run_refresh: kill run
+
+kill:
+	if docker container ls | grep -w ${CONTRAINER_NAME}; then \
+		echo "container ${CONTRAINER_NAME} exists, killing it first"; \
+		docker container kill  ${CONTRAINER_NAME}; fi
+
+build_private: build_public _setup_sshkey _setup_gitconfig
+	DOCKER_BUILDKIT=1 docker build --no-cache --target private \
+		--build-arg PUBLIC_IMAGE=${PUBLIC_IMAG} \
+--build-arg TORCH_HOME=${TO_TORCH_HOME} \
+		--build-arg https_proxy=${HTTPS_PROXY} \
+		--build-arg http_proxy=${HTTP_PROXY} \
+		--build-arg _USER=${_USER}  \
+--build-arg _HOME=${CHOME}  \
+		--build-arg _UID=${_UID} \
+		--build-arg _GID=${_GID} \
+		--build-arg PW=${PW} \
+		-t ${PRIVATE_IMAGE} .; \
+
+build_public:
+	DOCKER_BUILDKIT=1 docker build --pull --target public --build-arg BASE_IMAGE=${BASE_IMAGE} \
+	--build-arg https_proxy=${HTTPS_PROXY} \
+	--build-arg http_proxy=${HTTP_PROXY} \
+	-t ${PUBLIC_IMAG} .
+
+_instantiate_container: 
+	touch ${HOME}/.netrc
+	touch ${HOME}/.zsh_history
+	if [ ! -d ${HOME}/project ]; then \
+		mkdir -p ${HOME}/project;fi 
+	if [ ! -d ${HOME}/vscode/vscode-server ]; then \
+		mkdir -p ${HOME}/vscode/vscode-server;fi 
+	if [ ! -d ${HOME}/vscode/vscode-remote-containers ]; then \
+		mkdir -p ${HOME}/vscode/vscode-remote-containers; fi 
+	if ! docker container ls | grep -w ${CONTRAINER_NAME}; then \
+		docker run -itd --rm --shm-size=256g \
+			-v ${SHARE_DIR}:${TO_SHARE_DIR} \
+			-v ${HOME}/.netrc:${CHOME}/.netrc \
+			-v ${HOME}/.zsh_history:${CHOME}/.zsh_history \
+			-v ${PROJ_DIR}:${TO_PROJ_DIR} \
+			-v ${HOME}/vscode/vscode-server:${CHOME}/.vscode-server \
+			-v ${HOME}/vscode/vscode-remote-containers:${CHOME}/.vscode-remote-containers \
+			--user ${_USER} \
+			--name ${CONTRAINER_NAME} --gpus all $(PRIVATE_IMAGE); \
+	fi; \
+
+_exec:
+	docker exec -it ${CONTRAINER_NAME} zsh
+
+_exec_root:
+	docker exec -itu 0 ${CONTRAINER_NAME} zsh
+
+# attach is dangerous, exit from attached shell will terminate the container
+_attach:
+	docker attach ${CONTRAINER_NAME}
+
+
+_check_image:
+	if ! docker image ls | grep -w ${PRIVATE_IMAGE}; then \
+		echo "${PRIVATE_IMAGE} image not found, try building private image first"; \
+		${MAKE} -f ${MAKEFILE_} --no-print-directory build_private; \
+	fi
+
+
+_setup_sshkey:
+	touch ../.gitignore
+	if [ -z "$$(cat ../.gitignore| grep ssh/)" ]; then echo "ssh/" >> ../.gitignore; fi 
+	if [ -e asset/ssh/id_rsa ] && [ -e asset/ssh/id_rsa.pub ]; then \
+		echo "sshkey exists, use default key";\
+	else \
+		if [ ! -e ~/.ssh/id_rsa ] || [ ! -e ~/.ssh/id_rsa.pub ]; then \
+			echo "sshkey don't exist, try generating it"; \
+			ssh-keygen -t rsa -f ~/.ssh/id_rsa -q -N ""; \
+			echo new sshkey generated; \
+		fi; \
+		mkdir -p asset/ssh; \
+		cp ~/.ssh/id_rsa asset/ssh/id_rsa; \
+		cp ~/.ssh/id_rsa.pub asset/ssh/id_rsa.pub; \
+		echo To grant access to your github account, make sure to copy the content in [~/.ssh/id_rsa.pub] to github:;\
+		cat ~/.ssh/id_rsa.pub; \
+	fi
+
+_setup_gitconfig:
+	touch ../.gitignore
+	if [ -z "$$(cat ../.gitignore| grep gitconfig)" ]; then echo "gitconfig" >> ../.gitignore; fi 
+	if [ -e asset/gitconfig ]; then \
+		echo "gitconfig exists, use default gitconfig"; \
+	else \
+		if [ ! -e ~/.gitconfig ]; then \
+			echo "gitconfig don't exist, try configure it"; \
+			read -p "(no space and special character allowed) Enter github username: " GIT_NAME; \
+			read -p "(must be your real github account) Enter github email: " GIT_EMAIL; \
+			git config --global user.name $${GIT_NAME:-${PROJECT_NAME}}; \
+			git config --global user.email $${GIT_EMAIL:-${PROJECT_NAME}@gmail.com}; \
+		fi; \
+		cp ~/.gitconfig asset/gitconfig; \
+	fi

docker/asset/apt_packages.txt

@@ -0,0 +1,6 @@
+unzip
+autoconf
+less 
+gpustat 
+tmux
+libxrender1

docker/asset/requirements.txt

@@ -0,0 +1,21 @@
+# Dependencies necessary to execute run_docker.py
+ipykernel
+black
+fire
+absl-py==0.13.0
+tqdm
+matplotlib
+numpy
+scipy
+lxml
+pika
+wandb
+imageio
+msgpack
+zuko
+torchdiffeq
+torch_geometric
+torch_scatter
+pytorch_lightning
+overrides
+gdown

docker/setup_docker_for_host.sh

@@ -0,0 +1,40 @@
+#!/bin/bash
+set -e
+
+if [ "$EUID" -ne 0 ]
+  then echo "Please run as root or sudo"
+  exit
+fi
+
+if ! [ -x "$(command -v docker)" ]; then
+  echo 'docker is not installed. installing docker now...'
+  curl -fsSL https://get.docker.com -o ~/get-docker.sh; bash ~/get-docker.sh
+fi
+
+if ! [ -x "$(command -v nvidia-container-toolkit)" ]; then
+    curl -fsSL https://nvidia.github.io/libnvidia-container/gpgkey | gpg --dearmor -o /usr/share/keyrings/nvidia-container-toolkit-keyring.gpg \
+    && curl -s -L https://nvidia.github.io/libnvidia-container/stable/deb/nvidia-container-toolkit.list | \
+        sed 's#deb https://#deb [signed-by=/usr/share/keyrings/nvidia-container-toolkit-keyring.gpg] https://#g' | \
+        tee /etc/apt/sources.list.d/nvidia-container-toolkit.list \
+    && apt-get update \
+    && apt-get install -y nvidia-container-toolkit \
+    && nvidia-ctk runtime configure --runtime=docker \
+    && systemctl restart docker
+fi
+
+
+if [ -z "$SUDO_USER" ]; then
+    #read user from input
+    read -p "Please enter the username of the user you want to add to the docker group: " SUDO_USER
+fi
+
+if groups "$SUDO_USER" | grep -q "\bdocker\b"; then
+    echo ""
+else
+    usermod -aG docker $SUDO_USER
+    newgrp docker
+fi
+
+
+
+