[Tunix] Update Dockerfile dependencies and installation order.

PiperOrigin-RevId: 904786949

Author: wang2yn84

Dockerfile

@@ -21,27 +21,24 @@ ENV PATH="/opt/venv/bin:$PATH"
 RUN pip install --upgrade pip
 
 RUN pip install git+https://github.com/ayaka14732/jax-smi.git
-RUN pip install git+https://github.com/AI-Hypercomputer/pathways-utils.git
 # If you encounter a checkpoint issue, try using following old version of pathways-utils.
 # RUN pip install git+https://github.com/AI-Hypercomputer/pathways-utils.git@b72729bb152b7b3426299405950b3af300d765a9#egg=pathwaysutils
 RUN pip install gcsfs
-RUN pip install jax[tpu] -f https://storage.googleapis.com/jax-releases/libtpu_releases.html
-
-RUN pip install --upgrade wandb
+RUN pip install wandb
 
 # Set the working directory
 WORKDIR /app
 
 # Copy the project files to the image
 COPY . .
 
+# Install the project in editable mode
+RUN pip install -e .
+
 ENV VLLM_TARGET_DEVICE=tpu
 
 RUN pip install -r /app/requirements/requirements.txt
 RUN pip install --force-reinstall -r /app/requirements/special_requirements.txt
 
-# Install the project in editable mode
-RUN pip install --force-reinstall .
-
 # Set the default command to bash
 CMD ["bash"]