# syntax=docker/dockerfile:experimental

ARG FROM_IMAGE_NAME
FROM $FROM_IMAGE_NAME as main
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 python3-venv && \
      apt-get clean && \
      python -m venv /opt/jet

RUN pip3 install --no-cache-dir \
      black==24.4.2 \
      isort==5.13.2 \
      flake8==7.1.0 \
      pylint==3.2.6 \
      mypy

COPY . /opt/megatron-lm

WORKDIR /opt/megatron-lm

##### For NVIDIANS only #####
FROM main as jet
ARG CACHEBUST=0
RUN --mount=type=secret,id=JET_INDEX_URLS \
      JET_INDEX_URLS=$(cat /run/secrets/JET_INDEX_URLS) && \
      pip install jet-client jet-api --upgrade $JET_INDEX_URLS
ENV PATH="$PATH:/opt/jet/bin"
###