Установка переменных окружения Runpod после запуска подов

Содержание

Вопрос или проблема
Ответ или решение
1. Возможна ли установка переменных окружения после запуска подов?
2. Могу ли я выполнить bash-команду на поде вне Docker-образа?
Решение задачи
Заключение

Вопрос или проблема

Я запускаю тренировочный скрипт на облаке GPU Runpod, который требует кросс-связи (это реализация FSDP: https://github.com/AnswerDotAI/fsdp_qlora).

В настоящее время я создаю контейнеры Docker на каждом поде и мне нужен IP-адрес главного пода (может быть любой из подов), что является требованием для скрипта.

Из-за этого я предполагаю, что мне нужно установить переменные окружения после запуска подов, поэтому у меня есть 2 вопроса:

Возможно ли это сделать?
Могу ли я вообще запустить bash-команду на поде вне контейнера Docker? Поскольку я даже не вижу IP, пока образ Docker не запущен.

Также я хочу сделать все через API/SDK – в идеале все должно быть автоматизировано.

Я знаю, что мой весь рабочий процесс не может работать, но интересно, как я могу его исправить, чтобы получить связь.

Вот мой dockerfile:

FROM runpod/pytorch:2.0.1-py3.10-cuda11.8.0-devel-ubuntu22.04

ENV WANDB_API_KEY=""
ENV HUGGINGFACE_API_KEY=""

WORKDIR /workspace

RUN git clone https://github.com/AnswerDotAI/fsdp_qlora
WORKDIR /workspace/fsdp_qlora

RUN pip install llama-recipes fastcore "transformers!=4.38.*,!=4.39.*" --extra-index-url https://download.pytorch.org/whl/test/cu118 && \
    pip install "bitsandbytes>=0.43.0" && \
    pip install wandb runpod

COPY run_training.sh /workspace/

CMD ["/bin/bash", "/workspace/run_training.sh"]

run_training.sh:


export WANDB_API_KEY=$(cat $WANDB_API_KEY_FILE)
export HUGGINGFACE_API_KEY=$(cat $HUGGINGFACE_API_KEY_FILE)

# Вход в huggingface
echo $HUGGINGFACE_API_KEY | huggingface-cli login

echo "Начинаю тренировку с RANK ${RANK} из WORLD_SIZE ${WORLD_SIZE}"

python train.py \
    --model_name "meta-llama/Meta-Llama-Guard-2-8B" \
    --batch_size 2 \
    --context_length 512 \
    --precision bf16 \
    --train_type qlora \
    --use_gradient_checkpointing true \
    --use_cpu_offload true \
    --dataset alpaca_sample \
    --reentrant_checkpointing true \
    --log_to wandb \
    --project_name fsdp \
    --world_size ${WORLD_SIZE} \
    --rank ${RANK}

Создание подов выглядит так:

pod = runpod.create_pod(pod_config["name"], gpu_type_id=pod_config['gpu_type_id'], image_name=pod_config["image_name"], support_public_ip=True, volume_mount_path=pod_config["volume_mount_path"], ports=pod_config['ports'])

# код для получения ip-адресов
env_vars = {
        "WORLD_SIZE": "2",
        "MASTER_ADDR": pod_ips[0],
        "MASTER_PORT": "29500",
# как-то установить переменные окружения и запустить bash-скрипт??

Благодарю за любые советы!

изменение:
Одна идея, которую я имею, это запустить скрипт вроде этого в Dockerfile, который бы сказал скрипту запускаться, когда все API-ключи и IP-адреса будут отправлены. Возможно, это не идеально, но я не вижу, почему это не сработало бы.

while [ ! -f /workspace/start_training.text ]
# ждать 5 секунд
done

запустить файл

Ответ или решение

Чтобы решить указанные вами проблемы с настройкой переменных окружения для обучения в облаке Runpod, давайте разберемся с вашими вопросами по порядку.

1. Возможна ли установка переменных окружения после запуска подов?

Да, это возможно. Вы можете установить переменные окружения после того, как ваши поды запущены, используя API или SDK Runpod. Вы можете обновить существующий под и добавить необходимые переменные окружения. Однако, эти изменения обычно не применяются к уже запущенным процессам внутри пода. Поэтому вам нужно будет корректно инициировать или перезапустить процесс обучения с новыми переменными.

2. Могу ли я выполнить bash-команду на поде вне Docker-образа?

К сожалению, если ваш контейнер уже запущен, доступ к его внутренней среде ограничен, и вы не сможете просто выполнить команды bash вне контейнера, пока он работает. Однако, у вас есть возможность использовать docker exec, если у вас есть доступ к Docker на уровне хоста, но в случае с облачными решениями это может быть недоступно.

Решение задачи

Для автоматизации вашего рабочего процесса и настройки переменных окружения с учётом IP-адресов, я предлагаю следующее решение:

Создайте точку запуска: Вместо ожидания создания файла в контейнере, вы можете использовать механизм инициализации, который будет ожидать, пока все поды не будут запущены и не получат свои IP-адреса. После этого вы передадите необходимые переменные окружения через API.
Настройка вашего Dockerfile: В вашем Dockerfile создайте скрипт, который будет ожидать запуска всего кластера:

#!/bin/bash
# wait until master is set
while [ -z "$MASTER_ADDR" ]; do
  sleep 5
done
exec /workspace/run_training.sh

Обновление кода для создания подов: После создания подов и получения их IP-адресов, убедитесь, что вы устанавливаете переменные окружения, используя API:

pod = runpod.create_pod(pod_config["name"], gpu_type_id=pod_config['gpu_type_id'], 
                         image_name=pod_config["image_name"], support_public_ip=True, 
                         volume_mount_path=pod_config["volume_mount_path"], ports=pod_config['ports'])

# Получите IP-адреса подов
pod_ips = get_pod_ips()  # Замените на вашу функцию для получения IP-адресов

# Настройка переменных окружения
env_vars = {
    "WORLD_SIZE": "2",
    "MASTER_ADDR": pod_ips[0],
    "MASTER_PORT": "29500",
}

# Установите переменные окружения через API
set_env_variables(pod.id, env_vars)  # Замените на ваш метод установки переменных окружения

Запуск процесса обучения: После того как переменные окружения установлены, скрипт run_training.sh будет выполнен с правильными параметрами.

Заключение

Таким образом, вы сможете запустить несколько подов, установить необходимые переменные окружения и эффективно организовать процесс обучения. Это решение позволяет минимизировать длительность процесса и обеспечить правильную связь между подами при использовании FSDP. Не забудьте протестировать данное решение в тестовой среде перед использованием в продуктиве, чтобы убедиться, что все работает корректно.