inference/enterprise/device_manager/command_handler.py

from typing import Literal, Optional

import requests
from pydantic import BaseModel

import docker
from inference.core.devices.utils import GLOBAL_DEVICE_ID
from inference.core.env import API_BASE_URL, API_KEY
from inference.core.logger import logger
from inference.core.utils.url_utils import wrap_url
from inference.enterprise.device_manager.container_service import get_container_by_id


class Command(BaseModel):
    id: str
    containerId: str
    command: Literal["restart", "stop", "ping", "snapshot", "update_version"]
    deviceId: str
    requested_on: Optional[int] = None


def fetch_commands():
    url = wrap_url(
        f"{API_BASE_URL}/devices/{GLOBAL_DEVICE_ID}/commands?api_key={API_KEY}"
    )
    resp = requests.get(url).json()
    for cmd in resp.get("data", []):
        handle_command(cmd)


def handle_command(cmd_payload: dict):
    was_processed = False
    container_id = cmd_payload.get("containerId")
    container = get_container_by_id(container_id)
    if not container:
        logger.warn(f"Container with id {container_id} not found")
        ack_command(cmd_payload.get("id"), was_processed)
        return
    cmd = cmd_payload.get("command")
    data = None
    match cmd:
        case "restart":
            was_processed, data = container.restart()
        case "stop":
            was_processed, data = container.stop()
        case "ping":
            was_processed, data = container.ping()
        case "snapshot":
            was_processed, data = container.snapshot()
        case "start":
            was_processed, data = container.start()
        case "update_version":
            was_processed, data = handle_version_update(container)
        case _:
            logger.error("Unknown command: {}".format(cmd))
    return ack_command(cmd_payload.get("id"), was_processed, data=data)


def ack_command(command_id, was_processed, data=None):
    post_body = dict()
    post_body["api_key"] = API_KEY
    post_body["commandId"] = command_id
    post_body["wasProcessed"] = was_processed
    if data:
        post_body["data"] = data
    url = wrap_url(f"{API_BASE_URL}/devices/{GLOBAL_DEVICE_ID}/commands/ack")
    requests.post(url, json=post_body)


def handle_version_update(container):
    try:
        config = container.get_startup_config()
        image_name = config["image"].split(":")[0]
        container.kill()
        client = docker.from_env()
        new_container = client.containers.run(
            image=f"{image_name}:latest",
            detach=config["detach"],
            privileged=config["privileged"],
            labels=config["labels"],
            ports=config["port_bindings"],
            environment=config["env"],
            network="host",
        )
        logger.info(f"New container started {new_container}")
        return True, None
    except Exception as e:
        logger.error(e)
        return False, None