Update Dockerfile

Dockerfile

@@ -1,17 +1,15 @@
-FROM nvidia/cuda:12.1.1-cudnn8-devel-ubuntu22.04
-
-ARG DEBIAN_FRONTEND=noninteractive
-
-ENV PYTHONUNBUFFERED=1
-
-RUN apt-get update && apt-get install --no-install-recommends -y \
-  build-essential \
-  python3.9 \
-  python3-pip \
-  libpython3.9-dev \
-  git \
-  ffmpeg \
-  && apt-get clean && rm -rf /var/lib/apt/lists/* \ 
+FROM nvidia/cuda:12.1.0-devel-ubuntu22.04 AS dev
+
+RUN apt-get update -y \
+    && apt-get install -y python3-pip git
+
+# Workaround for https://github.com/openai/triton/issues/2507 and
+# https://github.com/pytorch/pytorch/issues/107960 -- hopefully
+# this won't be needed for future versions of this docker image
+# or future versions of triton.
+RUN ldconfig /usr/local/cuda-12.1/compat/
+RUN pip install cmake>=3.21 ninja psutil sentencepiece numpy requests py-cpuinfo transformers >= 4.39.1  fastapi uvicorn[standard] pydantic >= 2.0 prometheus_client >= 0.18.0 tiktoken == 0.6.0 outlines == 0.0.34 
+RUN ray >= 2.9 pynvml==11.5.0 vllm-nccl-cu12>=2.18,<2.19  torch==2.2.1 xformers == 0.0.25 triton>=2.1.0
 
 WORKDIR /code