FROM public.ecr.aws/w6p6i9i7/aws-efa-nccl-rdma:21.11-pt-py3

# download imagenet tiny for data
RUN apt-get -q update && apt-get -q install -y wget unzip; rm -rf examples/ docker-examples/ tutorials/
RUN wget -q http://cs231n.stanford.edu/tiny-imagenet-200.zip && unzip -q tiny-imagenet-200.zip -d /workspace/data && rm tiny-imagenet-200.zip

RUN apt-get update && apt-get install -y git htop
RUN pip3 install python-etcd

RUN mkdir -p /workspace/data
WORKDIR /workspace

RUN git clone https://github.com/pytorch/torchx.git && pushd torchx && git checkout v0.4.0 && popd
RUN pip3 install hydra-core omegaconf torchx
ENV PYTHONPATH=/workspace/torchx:$PYTHONPATH

RUN wget https://raw.githubusercontent.com/pytorch/elastic/master/examples/imagenet/main.py

ADD *.py /workspace/