Конвейер для быстрых экспериментов в Kubernetes
Ручное создание нового файла конфигурации для каждого нового эксперимента - утомительный процесс. Особенно, если вы хотите быстро развернуть огромное количество заданий в кластере Kubernetes, автоматическая настройка является обязательной. С помощью python легко создать простой сценарий планирования, который считывает конфигурацию эксперимента, такую как размер пакета, записывает ее в файл YAML и создает новое задание. В этом посте мы обсудим, как это сделать. Самое лучшее - это то, что мы не требуем никаких дополнительных пакетов!
Конвейер состоит из четырех файлов: двух скриптов bash (один для создания и один для удаления заданий Kubernetes), одного скрипта python и одного шаблона .yaml-файла. Давайте рассмотрим их более подробно, начиная со скрипта Python. Вы можете найти полный код в этом репозитории GitHub.
Скрипт на python
Код python структурирован в два метода. Первый метод дает экспериментальные конфигурации; они заполняются примерными значениями. Второй метод выполняет фактическое планирование, которое включает в себя синтаксический анализ .yaml-файла и обмен данными с Kubernetes. Давайте начнем с первой, более простой функции:
def get_experiments(experiment_number: int = -1) -> dict[Any, Any]:
"""
Returns a dictionary of experiments to run. Note that these are just arbitrary examples.
:param experiment_number: The experiment number to run. If -1, run all experiments.
:return: A dictionary of experiments to run.
"""
experiments = {1: {"batch_size": 32, "epochs": 100, "model": "resnet18", "dataset": "cifar10"},
2: {"batch_size": 64, "epochs": 100, "model": "resnet18", "dataset": "omniglot"},
3: {"batch_size": 32, "epochs": 50, "model": "resnet152", "dataset": "animals"},
4: {"batch_size": 64, "epochs": 50, "model": "resnet152", "dataset": "flowers"},
}
if experiment_number != -1:
return {experiment_number: experiments[experiment_number]}
return experiments
Метод get_experiments содержит внутренний объект dictionary, который в нашем примере содержит 4 образца экспериментов. Каждому эксперименту присваивается уникальный номер, и каждый эксперимент снова является словарем. Этот словарь содержит конфигурацию эксперимента и включает стандартные параметры машинного обучения, такие как размер пакета, количество эпох и модель. Далее мы обозначаем, на каком наборе данных мы хотим запустить наш эксперимент, например, на наборе данных CIFAR10. Параметры, перечисленные здесь, выбраны в иллюстративных целях, и вы не ограничены. Например, вы могли бы включить дополнительные гиперпараметры (например, скорость обучения), сохранить пути к файлам (например, каталог набора данных) или включить переменные среды. Короче говоря: адаптируйтесь к вашим потребностям.
Если мы вызываем метод без параметров, это равносильно вызову его с настройкой по умолчанию “experiment_number=-1”. С помощью этого параметра возвращаются конфигурации всех экспериментов. Однако, если вы работали с Kubernetes в течение более длительного времени, у вас часто будут некоторые задания, которые по каким-либо причинам завершатся неудачей. Если это произойдет, лучше всего исправить код и перезапустить этот конкретный эксперимент. Таким образом, в дополнение к получению настроек всех экспериментов, мы включили функциональность для выбора конкретных экспериментов для повторного запуска. Этот вариант использования поддерживается, когда мы вызываем метод с номером конкретного эксперимента. Примером может служить 4, что даст конфигурацию для четвертого эксперимента.
Шаблон .yaml
Метод get_experiments вызывается в основной функции скрипта schedule. Прежде чем рассматривать это более подробно, нам нужно взглянуть на наш файл шаблона. Файл .yaml, который мы включили, ориентирован на файлы, которые мы обычно пишем при проведении экспериментов по машинному обучению в суперкомпьютерном кластере. Однако это недопустимый файл в том смысле, что вы можете использовать его как есть.
Вместо этого он предназначен для того, чтобы показать вам вашу гибкость в заполнении шаблона. Говоря о заполнении шаблона, вот оно:
apiVersion: batch/v1
kind: Job
metadata:
name: username-experiment-{0}
namespace: your-namespace
spec:
backoffLimit: 1
template:
metadata:
labels:
jobgroup: group-model-{1}
spec:
containers:
- image: {6}
imagePullPolicy: Always
name: train #name of the container
env:
- name: DATASET_DIR
value: "{7}"
command:
- "python"
args:
- "-u"
- "experiment.py"
- "--model"
- "{1}"
- "--batch_size"
- "{2}"
- "--dataset"
- "{3}"
- "--epochs"
- "{4}"
resources:
requests:
cpu: "4"
memory: 16Gi
nvidia.com/gpu: "1"
limits:
cpu: "4"
memory: 16Gi
nvidia.com/gpu: "1"
volumeMounts:
- mountPath: "/datasets" # directory as seen WITHIN the container
name: datasets # matches volume's name from below
restartPolicy: Never
volumes:
- name: datasets
cephfs: #note that this is a cephfs volume, adapt to your case
monitors:
- some_random_ips
user: your-namespace # <namespace>
path: "{5}"
secretRef:
name: ceph-secret
readOnly: false
- name: cache-volume
emptyDir:
Обратите внимание на отличительные места с двумя фигурными скобками: {}. Каждому из этих мест присваивается номер, начинающийся с нуля. Мы восполним эти пробелы в скрипте python, к которому вскоре вернемся. Наиболее подходящий для нашего варианта использования машинного обучения символ “{}” может быть размещен в любом месте шаблона; они не ограничены определенными местами.
Чтобы продемонстрировать разнообразие, мы разместили маркеры поперек файла .yaml: мы можем использовать их для передачи аргументов командной строки, монтирования каталогов, заполнения переменных окружения или выбора образа нашего модуля. Кроме того, мы можем повторно использовать заполнитель: в файле шаблона мы дважды использовали заполнитель “{1}”; один раз для назначения задания группе заданий (group-model-{1}, строка 11) и один раз для передачи имени модели команде строка (строки 25 и 26).
Заполнение файла шаблона выполняется с помощью метода расписания скрипта Python.
def schedule(args) -> None:
"""
Schedules the experiments to run.
:param args: The commandline arguments.
:return: Nothing.
"""
with open("template.yaml", "r") as f:
config_string = "".join(f.readlines())
experiments = get_experiments(args.specific_experiment)
image = "your-container-registry/your-image:latest"
directory_to_mount = "/some/directory/to/mount"
dataset_directory = "/some/directory/with/datasets"
for k, v in experiments.items():
print(f"---Running experiment {k}---")
study_name = str(k)
res = config_string.format(
study_name, # {0}
v["model"], # {1}
v["batch_size"], # {2}
v["dataset"], # {3}
v["epochs"], # {4}
directory_to_mount, # {5}
image, # {6}
dataset_directory, # {7}
)
print(res) # print the config file
# Note that this is just an example, and no actual job will be started or terminated
if args.delete_only:
# only delete the job (if it exists)
output = run(["bash", "delete.sh", res], stdout=PIPE)
print(output.stdout.decode("utf-8").strip())
else:
# first, delete old job with same name
output = run(["bash", "delete.sh", res], stdout=PIPE)
print(output.stdout.decode("utf-8").strip())
# then, create new job
output = run(["bash", "schedule.sh", res], stdout=PIPE)
print(output.stdout.decode("utf-8").strip())
Внутри функции мы сначала анализируем шаблон как есть. Затем мы собираем все эксперименты, которые хотим запустить, в строке 10. Следующие три переменные, строки с 11 по 13, предназначены для того, чтобы вдохновить вас: вы могли бы автоматизировать нечто большее, чем просто то, что мы предложили. Интересная часть начинается в строке 14, где мы повторяем все эксперименты, которые мы хотим создать. Как мы уже писали, конфигурация эксперимента хранится в объекте dictionary. Это означает, что мы можем запросить словарь при заполнении шаблона в строках с 17 по 26.
Чтобы было легче понять, какой слот заполнен, мы оставили комментарий за каждой строкой. Например, строка 18 заполняет место, отмеченное “{0}”, строка 19, отмеченная “{1}”, и так далее. Чтобы увидеть, что мы создали после заполнения шаблона, мы печатаем заполненную версию в строке 27.
На данный момент мы создали (в памяти) готовый к использованию файл YAML. Следующим шагом является создание соответствующих заданий Kubernetes, начиная со строки 30. Сначала мы проверяем, хотим ли мы удалить только старый эксперимент (например, потому что он по какой-то причине не удался, и сначала нам нужно исправить ошибки). Если это не так, мы удаляем старое задание — не может быть двух заданий с одинаковым именем — перед созданием задания заново.
Если предыдущее задание не существует, сценарий не завершится ошибкой при попытке завершить его, но выведет пустую строку и перейдет к созданию задания.
Сценарии оболочки
Как создание, так и удаление заданий перенаправляются в два небольших сценария bash. Первый, показанный ниже, использует команду kubectl для создания задания на основе того, что было передано (часть echo “$1”). Обратите внимание, что мы настроили kubectl на использование нашего пространства имен по умолчанию. Если вы этого не сделали, то либо напишите kubectl -n your-namespace, либо зарегистрируйте свое пространство имен как пространство по умолчанию:
#!/usr/bin/env bash
echo "$1" | kubectl create -f -
Сценарий для удаления заданий почти идентичен; мы только переключаем флаг “create” на “delete”:
#!/usr/bin/env bash
echo "$1" | kubectl delete -f -
Возвращение к скрипту Python
После объединения различных частей нам требуется один драйвер для запуска кода. Эта задача выполняется с помощью инструкции “main”, как показано в приведенном ниже фрагменте:
if __name__ == "__main__":
argument_parser = argparse.ArgumentParser()
argument_parser.add_argument("--specific-experiment", dest="specific_experiment", type=int, default=-1)
argument_parser.add_argument("--delete-only", dest="delete_only", action="store_true")
args = argument_parser.parse_args()
schedule(args)
При запуске экспериментов, как упоминалось ранее, мы можем выбрать (повторно) создание только определенного эксперимента. По умолчанию мы запускаем все эксперименты; для запуска отдельных экспериментов мы можем передать их идентификатор в командную строку.
Затем мы создаем флаг, сообщающий планировщику только об удалении запуска или запуске запуска заново. По умолчанию для этого флага установлено значение false, что означает, что мы сначала завершаем существующий запуск эксперимента, а затем перезапускаем его. Только если мы явно устанавливаем флаг в командной строке, он устанавливается в значение true, что означает, что мы только завершаем существующее задание, но не создаем его заново. Кроме того, нам вообще не нужно его устанавливать; отсутствие флага равнозначно его ложности. Наконец, мы анализируем аргументы и запускаем планирование.
Заключение
В этом посте мы рассмотрели эксперименты по быстрому развертыванию в кластере Kubernetes. Для достижения этой цели мы использовали четыре файла: файл шаблона .yaml, два сценария bash и один файл python. В файле python мы использовали два метода для сбора конфигурации эксперимента, а затем заполнили шаблон.
После заполнения шаблона задания создаются и удаляются с помощью двух однострочных bash-скриптов. В них мы использовали нативные команды Kubernetes. В конце концов, мы можем вызвать скрипт python (передав необязательные аргументы командной строки), и наши эксперименты будут запланированы автоматически.
Этот автоматизированный процесс удобен при тестировании комбинаций параметров: ручное создание файла .yaml для каждого эксперимента занимает много времени и чревато ошибками. Таким образом, избавьте себя от хлопот и создайте автоматизированный конвейер, как подробно описано в этом посте. Весь код для начала работы доступен в репозитории.