#!/bin/bash

MODEL_REQ="${MODEL_REQ:-requirements.txt}";
WORKER_DIR="${WORKER_DIR:-/}";
SYNAPSE_VERSION="${SYNAPSE_VERSION:-1.17.0}";
EVALUATION_REPO="${EVALUATION_REPO:-${WORKER_DIR}/lm-evaluation-harness}";

cp $WORKER_DIR/launch/limits.conf /etc/security/;

apt-get update -y;
apt-get install pdsh tmux ntp -y;

echo MODEL_REQ=$MODEL_REQ

if [ -f "$MODEL_REQ" ]; then
  pip install -r $MODEL_REQ;
else
  wget -q --output-document - https://raw.githubusercontent.com/HabanaAI/Megatron-DeepSpeed/${SYNAPSE_VERSION}/megatron/core/requirements.txt  | grep -v "^-e" > /tmp/requirements.txt;
  python3 -m pip install -r /tmp/requirements.txt ;
fi

if [ -z "${DEEPSPEED_PATH}" ]; then
  pip install git+https://github.com/HabanaAI/DeepSpeed.git@$SYNAPSE_VERSION;
else
  pip install ${DEEPSPEED_PATH}
fi

pip install -r $WORKER_DIR/launch/requirements.txt

sed -i "s/.*HABANA_LOGS.*//g" /etc/profile.d/habanalabs.sh

sed -i 's/^Include.*conf$//' /etc/ssh/ssh_config

pip install "transformers"
cp -r $EVALUATION_REPO /tmp/lm-evaluation-harness
pushd /tmp/lm-evaluation-harness
pip install .
popd