Merge pull request #1032 from StanfordVL/docker-torch

Add optional dependencies for OMPL and curobo, include both in Docker image

.github/workflows/examples-as-test.yml

@@ -28,7 +28,7 @@ jobs:
 
       - name: Install
         working-directory: omnigibson-src
-        run: pip install -e .[dev]
+        run: pip install -e .[dev,primitives]
 
       - name: Generate example tests
         working-directory: omnigibson-src
@@ -76,7 +76,7 @@ jobs:
 
       - name: Install
         working-directory: omnigibson-src
-        run: pip install -e .[dev]
+        run: pip install -e .[dev,primitives]
 
       - name: Run tests
         working-directory: omnigibson-src

.github/workflows/profiling.yml

@@ -34,7 +34,7 @@ jobs:
 
       - name: Install
         working-directory: omnigibson-src
-        run: pip install -e .[dev]
+        run: pip install -e .[dev,primitives]
 
       - name: Run performance benchmark
         run: bash scripts/profiling.sh

.github/workflows/tests.yml

@@ -51,7 +51,7 @@ jobs:
 
       - name: Install
         working-directory: omnigibson-src
-        run: pip install -e .[dev]
+        run: pip install -e .[dev,primitives]
 
       - name: Print env
         run: printenv

docker/prod.Dockerfile

@@ -7,6 +7,8 @@ RUN apt-get update && DEBIAN_FRONTEND=noninteractive apt-get install -y \
   && rm -rf /var/lib/apt/lists/*
 
 RUN rm -rf /isaac-sim/exts/omni.isaac.ml_archive/pip_prebundle/gym*
+RUN rm -rf /isaac-sim/exts/omni.isaac.ml_archive/pip_prebundle/torch*
+RUN rm -rf /isaac-sim/exts/omni.isaac.ml_archive/pip_prebundle/functorch*
 RUN rm -rf /isaac-sim/kit/extscore/omni.kit.pip_archive/pip_prebundle/numpy*
 RUN /isaac-sim/python.sh -m pip install click~=8.1.3
 
@@ -23,6 +25,26 @@ ENV MAMBA_ROOT_PREFIX /micromamba
 RUN micromamba create -n omnigibson -c conda-forge python=3.10
 RUN micromamba shell init --shell=bash
 
+# Install torch
+RUN micromamba run -n omnigibson micromamba install \
+  pytorch torchvision pytorch-cuda=11.8 \
+  -c pytorch -c nvidia -c conda-forge
+
+# Install cuda for compiling curobo
+RUN wget -O /cuda.run https://developer.download.nvidia.com/compute/cuda/11.8.0/local_installers/cuda_11.8.0_520.61.05_linux.run && \
+  sh /cuda.run --silent --toolkit && rm /cuda.run
+ENV PATH=/usr/local/cuda-11.8/bin:$PATH
+ENV LD_LIBRARY_PATH=/usr/local/cuda-11.8/lib64:$LD_LIBRARY_PATH
+
+# Install curobo. This can normally be installed when OmniGibson is pip
+# installed, but we need to install it beforehand here so that it doesn't
+# have to happen on every time a CI action is run (otherwise it's just
+# very slow)
+# Here we also compile this such that it is compatible with GPU architectures
+# Turing, Ampere, and Ada; which correspond to 20, 30, and 40 series GPUs.
+RUN TORCH_CUDA_ARCH_LIST='7.5;8.0;8.6+PTX' \
+  micromamba run -n omnigibson pip install git+https://github.com/StanfordVL/curobo@06d8c79b660db60c2881e9319e60899cbde5c5b5#egg=nvidia_curobo --no-build-isolation
+
 # Make sure isaac gets properly sourced every time omnigibson gets called
 ARG CONDA_ACT_FILE="/micromamba/envs/omnigibson/etc/conda/activate.d/env_vars.sh"
 RUN mkdir -p "/micromamba/envs/omnigibson/etc/conda/activate.d"
@@ -33,31 +55,6 @@ RUN echo "source /isaac-sim/setup_conda_env.sh" >> $CONDA_ACT_FILE
 
 RUN echo "micromamba activate omnigibson" >> /root/.bashrc
 
-# Prepare to build OMPL
-ENV CXX="g++"
-ENV MAKEFLAGS="-j `nproc`"
-RUN micromamba run -n omnigibson micromamba install -c conda-forge boost && \
-    micromamba run -n omnigibson pip install pyplusplus && \
-    git clone https://github.com/ompl/ompl.git /ompl && \
-    mkdir -p /ompl/build/Release && \
-    sed -i "s/find_program(PYPY/# find_program(PYPY/g" /ompl/CMakeModules/Findpypy.cmake
-
-# Build and install OMPL 
-RUN micromamba run -n omnigibson /bin/bash --login -c 'source /isaac-sim/setup_conda_env.sh && (which python > /root/PYTHON_EXEC) && (echo $PYTHONPATH > /root/PYTHONPATH)' && \
-    cd /ompl/build/Release && \
-    micromamba run -n omnigibson cmake ../.. \
-      -DCMAKE_INSTALL_PREFIX="$CONDA_PREFIX" \
-      -DBOOST_ROOT="$CONDA_PREFIX" \
-      -DPYTHON_EXEC=$(cat /root/PYTHON_EXEC) \
-      -DPYTHONPATH=$(cat /root/PYTHONPATH) && \
-    micromamba run -n omnigibson make -j 4 update_bindings && \
-    micromamba run -n omnigibson make -j 4 && \
-    cd py-bindings && \
-    micromamba run -n omnigibson make install
-
-# Test OMPL
-RUN micromamba run -n omnigibson python -c "from ompl import base"
-
 # Copy over omnigibson source
 ADD . /omnigibson-src
 WORKDIR /omnigibson-src
@@ -72,7 +69,7 @@ ENV DEV_MODE=${DEV_MODE}
 ARG WORKDIR_PATH=/omnigibson-src
 RUN if [ "$DEV_MODE" != "1" ]; then \
       echo "OMNIGIBSON_NO_OMNIVERSE=1 python omnigibson/download_datasets.py" >> /root/.bashrc; \
-      micromamba run -n omnigibson pip install -e .[dev]; \
+      micromamba run -n omnigibson pip install -e .[dev,primitives]; \
     else \
       WORKDIR_PATH=/; \
       cd / && rm -rf /omnigibson-src; \

setup.py

@@ -62,7 +62,11 @@
             "mkdocs-section-index",
             "mkdocs-literate-nav",
             "telemoma~=0.1.2",
-        ]
+        ],
+        "primitives": [
+            "nvidia-curobo @ git+https://github.com/StanfordVL/curobo@06d8c79b660db60c2881e9319e60899cbde5c5b5",
+            "ompl @ https://storage.googleapis.com/gibson_scenes/ompl-1.6.0-cp310-cp310-manylinux_2_28_x86_64.whl",
+        ],
     },
     tests_require=[],
     python_requires=">=3",