# dobromirpopov/petals-ai :latest full
FROM python:slim

# System dependencies
RUN apt-get update && apt-get install -y git && apt-get clean && rm -rf /var/lib/apt/lists/*

# Upgrade pip
RUN pip install --upgrade pip

# Add ENV for MODEL and HF_TOKEN
ENV MODEL_NAME=stabilityai/StableBeluga2
ENV HF_TOKEN=
ENV DHT_PEERS=
ENV ENV_PARAMS=

# pull petals and bake it into the image. comment for lite image
RUN git clone https://github.com/d-popov/health.petals.dev.git && cd /health.petals.dev && pip install -r requirements.txt


# Create the entrypoint script inside the Dockerfile
RUN echo '#!/bin/sh\n\
INITIAL_PEERS=""\n\
if [ -n "$DHT_PEERS" ]; then\n\
    INITIAL_PEERS="--initial_peers $DHT_PEERS"\n\
fi\n\
if [ ! -f "/initialized" ]; then\n\
    touch /initialized\n\
    pip install --upgrade petals\n\
    case "$MODE" in\n\
        "monitor")\n\
            git clone https://github.com/d-popov/health.petals.dev.git && \
            cd /health.petals.dev && pip install -r requirements.txt\n\
            ;;\n\
        "chat")\n\
            git clone https://github.com/petals-infra/chat.petals.dev.git && \
            cd /chat.petals.dev && pip install -r requirements.txt\n\
            ;;\n\
        "server" | "dht")\n\
            pip install git+https://github.com/bigscience-workshop/petals\n\
            ;;\n\
    esac\n\
fi\n\
if [ "$MODE" = "monitor" ] && [ -n "$DHT_PEERS" ]; then\n\
    echo "INITIAL_PEERS = ['$(echo $DHT_PEERS | sed "s/ /','/g")']" >> /health.petals.dev/config.py;\n\
    echo "setting initial peers to ['$(echo $DHT_PEERS | sed "s/ /','/g")']";\n\
fi\n\
if [ -n "$HF_TOKEN" ]; then\n\
    git config --global credential.helper store && \n\
    huggingface-cli login --token $HF_TOKEN --add-to-git-credential && \n\
    ENV_PARAMS="--use_auth_token $ENV_PARAMS"\n\
fi\n\
case "$MODE" in\n\
    "monitor")\n\
        cd /health.petals.dev\n\
        exec gunicorn app:app --bind 0.0.0.0:8484 --workers 4 --threads 10\n\
        ;;\n\
    "chat")\n\
        cd /chat.petals.dev\n\
        exec gunicorn app:app --bind 0.0.0.0:8484 --workers 4 --threads 10\n\
        ;;\n\
    "dht")\n\
        exec python -m petals.cli.run_dht --host_maddrs /ip4/0.0.0.0/tcp/31337 --identity_path bootstrap1.id $INITIAL_PEERS\n\
        ;;\n\
    "server" | *)\n\
        exec python -m petals.cli.run_server $ENV_PARAMS $MODEL_NAME $INITIAL_PEERS\n\
        ;;\n\
esac' > /entrypoint.sh && chmod +x /entrypoint.sh

# if using separate entrypoint file:
# COPY entrypoint.sh /entrypoint.sh
# RUN chmod +x /entrypoint.sh  TEST

ENTRYPOINT ["/entrypoint.sh"]