95 lines
No EOL
3.1 KiB
Python
95 lines
No EOL
3.1 KiB
Python
import requests
|
||
import docker
|
||
import os
|
||
|
||
from ..logger import LoggerFactory
|
||
|
||
manager_logger = LoggerFactory.get_logger("CustomModuleManager")
|
||
|
||
def change_module_status(
|
||
module_id: int,
|
||
status: str
|
||
):
|
||
_ = requests.post(
|
||
f"http://localhost:{int(os.environ.get('PORT', 5010))}/api/v1.0/interact_with_custom_modules",
|
||
data = {"request_type": "change_status", "model_id": module_id, "status": status}
|
||
)
|
||
|
||
def build_image(
|
||
client: docker.DockerClient,
|
||
path: str,
|
||
dockerfile: str,
|
||
tag: str,
|
||
buildargs: dict
|
||
):
|
||
"""Сборка docker контейнеров"""
|
||
_, logs = client.images.build(
|
||
path = path,
|
||
dockerfile = dockerfile,
|
||
tag = tag,
|
||
buildargs = buildargs,
|
||
decode = True
|
||
)
|
||
for chunk in logs:
|
||
if 'stream' in chunk:
|
||
manager_logger.info(chunk['stream'].strip())
|
||
|
||
def start_container(
|
||
module_id: int,
|
||
image_name: str,
|
||
build_args: dict,
|
||
container_name: str,
|
||
network_name: str,
|
||
env_vars: dict,
|
||
gpus = True,
|
||
detach = True
|
||
):
|
||
"""Запуск docker контейнеров"""
|
||
# инициализируем клиента
|
||
client = docker.from_env()
|
||
|
||
# Соберем контейнер
|
||
try:
|
||
change_module_status(module_id, status = "работает")
|
||
build_image(client, **build_args)
|
||
manager_logger.info(f"Контейнер '{container_name}' успешно собран")
|
||
except Exception as e:
|
||
change_module_status(module_id, status = "не создан")
|
||
manager_logger.error(f"Ошибка при сборке контейнера: {e}")
|
||
return
|
||
|
||
# Укажем устройства исполнения
|
||
device_requests = [docker.types.DeviceRequest(count=-1, capabilities=[['gpu']])] if gpus else None
|
||
|
||
# Запустим контейнер
|
||
try:
|
||
container = client.containers.run(
|
||
image = image_name,
|
||
name = container_name,
|
||
network = client.networks.get(network_name),
|
||
environment = env_vars,
|
||
device_requests=device_requests,
|
||
detach = detach,
|
||
)
|
||
manager_logger.info(f"Контейнер '{container_name}' с ID: {container.short_id} запустился")
|
||
except Exception as e:
|
||
change_module_status(module_id, status = "остановлен")
|
||
manager_logger.error(f"Ошибка при запуске контейнера: {e}")
|
||
|
||
def stop_container(
|
||
module_id: int,
|
||
name: str
|
||
):
|
||
"""Удаление docker контейнеров по имени"""
|
||
# инициализируем клиента
|
||
client = docker.from_env()
|
||
|
||
# Остановим контейнер
|
||
try:
|
||
container = client.containers.get(name)
|
||
container.stop()
|
||
container.remove()
|
||
manager_logger.info(f"Контейнер '{name}' был остановлен")
|
||
change_module_status(module_id, status = "остановлен")
|
||
except docker.errors.NotFound:
|
||
manager_logger.warning(f"Контейнер '{name}' не найден") |