updated dockerfile

Dockerfile

@@ -1,59 +1,23 @@
-FROM nvidia/cuda:11.8.0-devel-ubuntu22.04
-
-# install python
-ARG DEBIAN_FRONTEND=noninteractive
-ARG PYTHON_VERSION=3.10
-#setting language and locale
-ENV LANG="C.UTF-8" LC_ALL="C.UTF-8"
-
-
-RUN apt-get update && \
-  DEBIAN_FRONTEND=noninteractive apt-get -qq -y install \
-  software-properties-common \
-  build-essential \
-  curl \
-  ffmpeg \
-  git \
-  htop \
-  vim \
-  nano \
-  rsync \
-  wget \
-  && apt-get clean \
-  && rm -rf /var/lib/apt/lists/*
-
-  RUN add-apt-repository ppa:deadsnakes/ppa
-  RUN apt-get update && apt-get install -y -qq python${PYTHON_VERSION} \
-      python${PYTHON_VERSION}-dev \
-      python${PYTHON_VERSION}-distutils
-
-# Set python aliases
-RUN update-alternatives --install /usr/bin/python python /usr/bin/python${PYTHON_VERSION} 1
-RUN update-alternatives --install /usr/bin/python3 python3 /usr/bin/python${PYTHON_VERSION} 1
-RUN curl https://bootstrap.pypa.io/get-pip.py -o get-pip.py && python get-pip.py
+FROM nvcr.io/nvidia/jax:23.10-py3
 
 # default workdir
-WORKDIR /home/workdir
-COPY . .
-
-#jaxmarl from source if needed, all the requirements
-RUN pip install --ignore-installed -e '.[qlearning, dev]'
+WORKDIR /app
 
-# install jax from to enable cuda
-RUN pip install --upgrade "jax[cuda11_pip]" -f https://storage.googleapis.com/jax-releases/jax_cuda_releases.html
+COPY requirements/requirements.txt .
+RUN pip install -r requirements.txt
 
 #disabling preallocation
 RUN export XLA_PYTHON_CLIENT_PREALLOCATE=false
 #safety measures
 RUN export XLA_PYTHON_CLIENT_MEM_FRACTION=0.25 
 RUN export TF_FORCE_GPU_ALLOW_GROWTH=true
 
-#for jupyter
-EXPOSE 9999 
+# if you want jupyter 
+# RUN pip install pip install jupyterlab
+# for secrets and debug
 
-#for secrets and debug
 ENV WANDB_API_KEY=""
 ENV WANDB_ENTITY=""
-RUN git config --global --add safe.directory /home/workdir
+# RUN git config --global --add safe.directory /app
 
-CMD ["/bin/bash"]
+CMD ["/bin/bash"]

compose.yaml

@@ -0,0 +1,11 @@
+version: '3'
+
+services:
+  app:
+    image: jaxmarl:latest
+    stdin_open: true
+    tty: true
+    volumes:
+      - .:/app
+    runtime: nvidia
+