velocitatem/PHANTOM
1
0
Fork 0
mirror of https://github.com/velocitatem/PHANTOM.git synced 2026-05-31 08:33:36 +00:00
Code Issues Packages Projects Releases Wiki Activity

new trainer image

Browse Source
This commit is contained in:
Daniel Rosel
2026-02-19 13:03:25 +01:00
parent 843564eeb0
commit 5912062dc0
3 changed files with 79 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files

43
docker/Trainer.dockerfile Normal file
View File

@@ -0,0 +1,43 @@
# syntax=docker/dockerfile:1.7
FROM pytorch/pytorch:2.5.1-cuda12.4-cudnn9-runtime AS gpu
WORKDIR /app
COPY docker/trainer.requirements.txt /tmp/requirements.txt
RUN pip install --no-cache-dir -r /tmp/requirements.txt
# Optional for JAX-on-GPU workflows.
ARG INSTALL_JAX_GPU=false
RUN if [ "${INSTALL_JAX_GPU}" = "true" ]; then \
pip install --no-cache-dir "jax[cuda12]==0.4.30" -f https://storage.googleapis.com/jax-releases/jax_cuda_releases.html; \
fi
COPY --chmod=755 docker/trainer-agent-entrypoint.sh /usr/local/bin/trainer-agent-entrypoint
COPY engine /app/engine
ENV PYTHONPATH=/app \
XLA_PYTHON_CLIENT_PREALLOCATE=false
ENTRYPOINT ["/usr/local/bin/trainer-agent-entrypoint"]
FROM python:3.11-slim AS tpu
WORKDIR /app
COPY docker/trainer.requirements.txt /tmp/requirements.txt
RUN pip install --no-cache-dir -r /tmp/requirements.txt
RUN pip install --no-cache-dir "jax[tpu]==0.4.30" -f https://storage.googleapis.com/jax-releases/libtpu_releases.html
COPY --chmod=755 docker/trainer-agent-entrypoint.sh /usr/local/bin/trainer-agent-entrypoint
COPY engine /app/engine
ENV PYTHONPATH=/app \
PHANTOM_USE_JAX=1 \
PHANTOM_DEFAULT_AGENT_ARGS="--jax" \
JAX_PLATFORMS=tpu,cpu \
XLA_PYTHON_CLIENT_PREALLOCATE=false
ENTRYPOINT ["/usr/local/bin/trainer-agent-entrypoint"]