diff --git a/.env.example b/.env.example index 538ab30..d8a6a95 100644 --- a/.env.example +++ b/.env.example @@ -4,6 +4,11 @@ # For Dev Registry: https://raw.githubusercontent.com/premAI-io/prem-registry/dev/manifests.json PREM_REGISTRY_URL=https://raw.githubusercontent.com/premAI-io/prem-registry/main/manifests.json +# Prem Daemon +# ------------------------------------------------------------------------------------------ +DEFAULT_PORT='8000' +PREMD_IMAGE=ghcr.io/premai-io/premd + # Sentry # ------------------------------------------------------------------------------------------ SENTRY_DSN=https://75592545ad6b472e9ad7c8ff51740b73@o1068608.ingest.sentry.io/4505244431941632 diff --git a/app/core/config.py b/app/core/config.py index 870b137..33df5b1 100644 --- a/app/core/config.py +++ b/app/core/config.py @@ -11,6 +11,8 @@ DEBUG: bool = os.getenv("DEBUG", False) SECRET_KEY: Secret = Secret(os.getenv("SECRET_KEY", "")) PROJECT_NAME: str = os.getenv("PROJECT_NAME", "Prem Daemon") +PREMD_IMAGE: str = os.getenv("PREMD_IMAGE", "ghcr.io/premai-io/premd") +DEFAULT_PORT: int = int(os.getenv("DEFAULT_PORT", "8000")) # PROXY # ------------------------------------------------------------------------------ diff --git a/app/core/events.py b/app/core/events.py index 989d114..6da383c 100644 --- a/app/core/events.py +++ b/app/core/events.py @@ -9,6 +9,23 @@ def create_start_app_handler(app: FastAPI): def start_app() -> None: + # client = utils.get_docker_client() + # if utils.container_exists(container_name): + # container = client.containers.get(container_name) + # host_port = container.ports.get(f"{utils.DEFAULT_PORT}/tcp", [None])[0][ + # "HostPort" + # ] + # if host_port != f"{utils.DEFAULT_PORT}": + # utils.check_host_port_availability(utils.DEFAULT_PORT) + # new_container = utils.create_new_container( + # utils.PREMD_IMAGE, + # "latest", + # new_container_name, + # container_name, + # utils.DEFAULT_PORT, + # ) + # utils.update_and_remove_old_container(container_name) + # new_container.start() for registry in config.PREM_REGISTRY_URL.strip().split(): utils.add_services_from_registry(registry) diff --git a/app/core/utils.py b/app/core/utils.py index a0f33c8..c26cb0c 100644 --- a/app/core/utils.py +++ b/app/core/utils.py @@ -1,15 +1,21 @@ import logging +import re import subprocess +import time import xml.etree.ElementTree as ET import docker import requests import torch +from bs4 import BeautifulSoup +from packaging.version import parse as parse_version from app.core import config logger = logging.getLogger(__name__) +PREMD_IMAGE = config.PREMD_IMAGE +DEFAULT_PORT = config.DEFAULT_PORT SERVICES = [] REGISTRIES = config.PREM_REGISTRY_URL.strip().split() INTERFACES = [ @@ -198,6 +204,128 @@ def get_gpu_info(): return gpu_name, total_memory_value, used_memory_value, mem_percentage +def extract_labels_from_html_file(html_content, class_names): + soup = BeautifulSoup(html_content, "html.parser") + labels = soup.select(class_names) + return (label.get_text() for label in labels) + + +def find_maximum_label(labels): + pattern = re.compile(r"v\d+\.\d+\.\d+$") + return max(filter(pattern.match, labels), default=None, key=parse_version) + + +def get_premd_last_tag(owner, repository, package): + response = requests.get( + f"https://github.com/{owner}/{repository}/pkgs/container/{package}" + ) + try: + labels = extract_labels_from_html_file( + response.content, ".Label.mr-1.mb-2.text-normal" + ) + except Exception as e: + logger.info(f"Unexpected error: {e}") + return "latest" + else: + return find_maximum_label(labels) + + +def get_local_docker_image_tags(owner, repository): + try: + client = get_docker_client() + image = client.images.get(f"ghcr.io/{owner}/{repository}") + return image.tags + except Exception as e: + logger.info(f"Unexpected error: {e}") + return [] + + +def create_new_container(image_name, image_tag, new_container_name, old_container_name): + client = get_docker_client() + old_container = client.containers.get(old_container_name) + + if is_gpu_available(): + device_requests = [ + docker.types.DeviceRequest(device_ids=["all"], capabilities=[["gpu"]]) + ] + else: + device_requests = [] + + volumes = {} + for mount in old_container.attrs["Mounts"]: + source = mount["Source"] + target = mount["Destination"] + mode = mount["Mode"] + volumes[source] = {"bind": target, "mode": mode} + + current_ports = old_container.attrs["HostConfig"]["PortBindings"] + current_port = list(current_ports.items())[0] + + logger.info( + f"Starting new container {new_container_name} with image {image_name}:{image_tag} at port {current_port[0]}" + ) + new_container = client.containers.create( + image=f"{image_name}:{image_tag}", + name=new_container_name, + ports={current_port[0]: current_port[1]}, + volumes=volumes, + environment=old_container.attrs["Config"]["Env"], + device_requests=device_requests, + network_mode=old_container.attrs["HostConfig"]["NetworkMode"], + detach=True, + ) + return new_container + + +def update_and_remove_old_container(old_container_name): + client = get_docker_client() + logger.info(f"Stopping {old_container_name}") + old_container = client.containers.get(old_container_name) + old_container.stop() + + +def update_container(): + new_container = create_new_container( + PREMD_IMAGE, "latest", "new_container", "premd" + ) + update_and_remove_old_container("premd") + new_container.start() + new_container.rename("premd") + + +def check_host_port_availability(host_port, timeout=30): + start_time = time.time() + client = get_docker_client() + + while True: + if time.time() - start_time > timeout: + return False + + containers = client.containers.list() + port_used = any( + f"{host_port}/tcp" in container.ports + for container in containers + if container.status == "running" + ) + + if not port_used: + return True + + time.sleep(1) + + +def container_exists(container_name): + try: + client = get_docker_client() + _ = client.containers.get(container_name) + return True + except docker.errors.NotFound: + return False + except docker.errors.APIError as e: + logging.error(f"Error checking container existence: {e}") + return False + + cached_domain = None diff --git a/app/routes.py b/app/routes.py index 05781b2..31f33b5 100644 --- a/app/routes.py +++ b/app/routes.py @@ -44,6 +44,55 @@ async def health(): return schemas.HealthResponse(status=True) +@router.get( + "/update-available/", + responses={ + 400: { + "model": schemas.ErrorResponse, + "description": "Failed to check update available.", + } + }, + response_model=schemas.UpdateAvailableResponse, +) +async def update_available(): + try: + owner = "premAI-io" + remote_image = f"{utils.PREMD_IMAGE}:{utils.get_premd_last_tag(owner, 'prem-daemon', 'premd')}" + local_tags = utils.get_local_docker_image_tags(owner.lower(), "premd") + return { + "remote_image": remote_image, + "local_images": local_tags, + "update": remote_image not in local_tags, + } + except Exception as error: + logger.error(error) + raise HTTPException( + status_code=400, + detail={"message": f"Failed to check update available {error}."}, + ) from error + + +@router.get( + "/update-daemon/", + responses={ + 400: { + "model": schemas.ErrorResponse, + "description": "Failed to update.", + } + }, + response_model=schemas.UpdateAvailableResponse, +) +async def trigger_update(): + try: + utils.update_container() + except Exception as error: + logger.error(error) + raise HTTPException( + status_code=400, + detail={"message": f"Failed to update {error}."}, + ) from error + + @router.get("/interfaces/", response_model=list[schemas.InterfaceResponse]) async def interfaces(): return utils.get_interfaces() diff --git a/app/schemas.py b/app/schemas.py index 80bd4d1..3d8c71f 100644 --- a/app/schemas.py +++ b/app/schemas.py @@ -17,6 +17,12 @@ class RunServiceInput(BaseModel): id: str +class UpdateAvailableResponse(BaseModel): + remote_image: str + local_images: list[str] + update: bool + + class ServiceInput(BaseModel): id: str name: str diff --git a/requirements.txt b/requirements.txt index ccaa674..9d2548d 100644 --- a/requirements.txt +++ b/requirements.txt @@ -15,3 +15,5 @@ torchvision==0.15.2 torchaudio==2.0.2 sentry-sdk==1.26.0 psutil==5.9.5 +beautifulsoup4==4.12.2 +packaging==23.1