Docker updates (#3435)

* Fix installation of pytorch from a branch in the docker container.

* Add docker build trust hosts option and pytorch branch to docker image name.

* Install python after SSL config change.

* Update Levy-stable parameter fit test accuracy threshold from 0.03 to 0.04 (change is probably due to upgrade to pytorch 2.7.0).

---------

Co-authored-by: Ben Zickel <benz@uvisionuav.com>

Author: BenZickel

docker/Dockerfile

@@ -14,6 +14,7 @@ ARG uid=1000
 ARG gid=1000
 ARG ostype=Linux
 ARG pyro_git_url=https://github.com/pyro-ppl/pyro.git
+ARG trust_hosts=no
 
 # Configurable settings
 ENV USER_NAME pyromancer
@@ -37,11 +38,21 @@ RUN bash -c 'if [ ${ostype} == Linux ]; then groupadd -r --gid ${gid} ${USER_NAM
 USER ${USER_NAME}
 
 # Install conda
-RUN wget -O ~/miniconda.sh \
+RUN if [ ${trust_hosts} = yes ] ; then WGET_ARGS="--no-check-certificate" ; fi && \
+    wget ${WGET_ARGS} -O ~/miniconda.sh \
     https://repo.anaconda.com/miniconda/Miniconda${python_version%%.*}-latest-Linux-x86_64.sh && \
     bash ~/miniconda.sh -f -b -p ${CONDA_DIR} && \
-    rm ~/miniconda.sh && \
-    conda install python=${python_version}
+    rm ~/miniconda.sh
+
+# Trust conda and pip hosts if needed
+RUN if [ ${trust_hosts} = yes ] ; \
+    then \
+        pip config set global.trusted-host "pypi.org files.pythonhosted.org download.pytorch.org" && \
+        conda config --set ssl_verify False ; \
+    fi
+
+# Update python version
+RUN conda install python=${python_version}
 
 # Move to home directory; and copy the install script
 WORKDIR ${WORK_DIR}

docker/Makefile

@@ -21,6 +21,7 @@ python_version?=3.12
 pytorch_branch?=release
 pyro_branch?=release
 cmd?=bash
+trust_hosts?="no"
 
 # Determine name of docker image
 build run notebook: img_prefix=pyro-cpu
@@ -29,12 +30,11 @@ build run lab: img_prefix=pyro-cpu
 build-gpu run-gpu lab-gpu: img_prefix=pyro-gpu
 
 ifeq ($(img), )
-	IMG_NAME=${img_prefix}-${pyro_branch}-${python_version}
+	IMG_NAME=${img_prefix}-${pyro_branch}-${pytorch_branch}-${python_version}
 else
 	IMG_NAME=${img}
 endif
 
-
 help:
 	@fgrep -h "##" ${MAKEFILE_LIST} | fgrep -v fgrep | sed -e 's/##//'
 
@@ -51,6 +51,9 @@ build: ##
 	##     default - latest pytorch version on the torch python package index
 	##   pyro_branch: whether to use the released Pyro wheel or a git branch.
 	##     default - latest pyro-ppl wheel on pypi
+	##   trust_hosts: If set to yes hosts SSL ceritificates will be trusted
+	##   (might be needed when running begind a firewall)
+	##     default - Verify hosts SSL certificates
 	##
 	${DOCKER_CMD} build -t ${IMG_NAME} \
 	--build-arg base_img=${BASE_IMG} \
@@ -60,7 +63,8 @@ build: ##
 	--build-arg python_version=${python_version} \
 	--build-arg pytorch_branch=${pytorch_branch} \
 	--build-arg pyro_git_url=${pyro_git_url} \
-	--build-arg pyro_branch=${pyro_branch} -f ${DOCKER_FILE} .
+	--build-arg pyro_branch=${pyro_branch} \
+	--build-arg trust_hosts=${trust_hosts} -f ${DOCKER_FILE} .
 
 build-gpu: ##
 	## Build a docker image for running Pyro on a GPU.
@@ -71,6 +75,9 @@ build-gpu: ##
 	##     default - latest pytorch version on the torch python package index
 	##   pyro_branch: whether to use the released Pyro wheel or a git branch.
 	##     default - latest pyro-ppl wheel on pypi
+	##   trust_hosts: If set to yes hosts SSL ceritificates will be trusted
+	##   (might be needed when running begind a firewall)
+	##     default - Verify hosts SSL certificates
 	##
 	${DOCKER_CMD} build -t ${IMG_NAME} \
 	--build-arg base_img=${BASE_CUDA_IMG} \
@@ -81,7 +88,8 @@ build-gpu: ##
 	--build-arg python_version=${python_version} \
 	--build-arg pytorch_branch=${pytorch_branch} \
 	--build-arg pyro_git_url=${pyro_git_url} \
-	--build-arg pyro_branch=${pyro_branch} -f ${DOCKER_FILE} .
+	--build-arg pyro_branch=${pyro_branch} \
+	--build-arg trust_hosts=${trust_hosts} -f ${DOCKER_FILE} .
 
 create-host-workspace: ##
 	## Create shared volume on the host for sharing files with the container.

docker/install.sh

@@ -18,7 +18,14 @@ if [ ${pytorch_branch} != "release" ]
 then
     git clone --recursive https://github.com/pytorch/pytorch.git
     pushd pytorch && git checkout ${pytorch_branch}
-    pip uninstall torch
+    pip uninstall -y torch
+    conda install cmake ninja
+    pip install -r requirements.txt
+    pip install mkl-static mkl-include
+    if [ ${pytorch_whl} != "cpu" ]
+    then
+        conda install -c pytorch magma-cuda${pytorch_whl:2}
+    fi
     pip install -e .
     popd
 fi

tests/distributions/test_stable_log_prob.py

@@ -129,7 +129,7 @@ def log_progress():
     train(model, guide)
 
     # Verify fit accuracy
-    assert_close(alpha, pyro.param("alpha").item(), atol=0.03)
+    assert_close(alpha, pyro.param("alpha").item(), atol=0.04)
     assert_close(beta, pyro.param("beta").item(), atol=0.06)
     assert_close(c, pyro.param("c").item(), atol=0.2)
     assert_close(mu, pyro.param("mu").item(), atol=0.2)