apptainer_def_files/py311_torch_mpi.def

Bootstrap: docker
#From: python:latest
From: python:3.11-slim

%environment
  # Point to OMPI binaries, libraries, man pages
  export OMPI_DIR=/opt/ompi
  export PATH="$OMPI_DIR/bin:$PATH"
  export LD_LIBRARY_PATH="$OMPI_DIR/lib:$LD_LIBRARY_PATH"
  export MANPATH="$OMPI_DIR/share/man:$MANPATH"

%post
  echo "Installing required packages..."
  apt-get update && apt-get install -y wget git bash gcc gfortran g++ make file bzip2

  echo "Installing Open MPI"
  export OMPI_DIR=/opt/ompi
  export OMPI_VERSION=4.1.5
  export OMPI_URL="https://download.open-mpi.org/release/open-mpi/v4.1/openmpi-$OMPI_VERSION.tar.bz2"
  mkdir -p /tmp/ompi
  mkdir -p /opt
  # Download
  cd /tmp/ompi && wget -O openmpi-$OMPI_VERSION.tar.bz2 $OMPI_URL && tar -xjf openmpi-$OMPI_VERSION.tar.bz2
  # Compile and install
  cd /tmp/ompi/openmpi-$OMPI_VERSION && ./configure --prefix=$OMPI_DIR && make -j$(nproc) install

  # Set env variables so we can compile our application
  export PATH=$OMPI_DIR/bin:$PATH
  export LD_LIBRARY_PATH=$OMPI_DIR/lib:$LD_LIBRARY_PATH

  # Clean up
  apt-get clean
  rm -rf /var/lib/apt/lists/*
  rm -rf /tmp/*

  # Set up the virtual environment
  mkdir /app

  python -m venv /app/venv

  # activate the virtual environment
  . /app/venv/bin/activate

  # Install dependencies
  echo "Installing Python dependencies..."
  # pip install --no-cache-dir pycocotools
  pip install --no-cache-dir torch torchvision torchaudio lightning tensorboard pycocotools matplotlib

%runscript
  . /app/venv/bin/activate
  # echo "This is the runscript"
  exec "$@"