# syntax=docker/dockerfile:experimental

ARG FROM_IMAGE_NAME
FROM $FROM_IMAGE_NAME
ENV DEBIAN_FRONTEND=noninteractive

RUN sed -i -e 's/^APT/# APT/' -e 's/^DPkg/# DPkg/' \
      /etc/apt/apt.conf.d/docker-clean

RUN apt-get update && \
      apt-get install -y --no-install-recommends gettext && \
      apt-get clean

RUN wget https://github.com/mikefarah/yq/releases/download/v4.44.1/yq_linux_amd64 -O /usr/local/bin/yq && \
chmod a+x /usr/local/bin/yq

RUN pip3 install --no-cache-dir \
      einops \
      flask-restful \
      nltk \
      pytest \
      pytest-cov \
      pytest_mock \
      sentencepiece \
      wrapt \
      git+https://github.com/fanshiqing/grouped_gemm@v1.1.2 \
      zarr \
      tensorstore==0.1.45

COPY . /workspace/megatron-lm

RUN cp -r /workspace/megatron-lm /opt && \
    pip install /opt/megatron-lm