container_service

@dataclass
class InferServerContainer:
    status: str
    id: str
    port: int
    host: str
    startup_time: float
    version: str

    def __init__(self, docker_container, details):
        self.container = docker_container
        self.status = details.get("status")
        self.id = details.get("uuid")
        self.port = details.get("port")
        self.host = details.get("host")
        self.version = details.get("version")
        t = details.get("startup_time_ts").split(".")[0]
        self.startup_time = (
            datetime.strptime(t, "%Y-%m-%dT%H:%M:%S").timestamp()
            if t is not None
            else datetime.now().timestamp()
        )

    def kill(self):
        try:
            self.container.kill()
            return True, None
        except Exception as e:
            logger.error(e)
            return False, None

    def restart(self):
        try:
            self.container.restart()
            return True, None
        except Exception as e:
            logger.error(e)
            return False, None

    def stop(self):
        try:
            self.container.stop()
            return True, None
        except Exception as e:
            logger.error(e)
            return False, None

    def start(self):
        try:
            self.container.start()
            return True, None
        except Exception as e:
            logger.error(e)
            return False, None

    def inspect(self):
        try:
            info = requests.get(f"http://{self.host}:{self.port}/info").json()
            return True, info
        except Exception as e:
            logger.error(e)
            return False, None

    def snapshot(self):
        try:
            snapshot = self.get_latest_inferred_images()
            snapshot.update({"container_id": self.id})
            return True, snapshot
        except Exception as e:
            logger.error(e)
            return False, None

    def get_latest_inferred_images(self, max=4):
        """
        Retrieve the latest inferred images and associated information for this container.

        This method fetches the most recent inferred images within the time interval defined by METRICS_INTERVAL.

        Args:
            max (int, optional): The maximum number of inferred images to retrieve.
                Defaults to 4.

        Returns:
            dict: A dictionary where each key represents a model ID associated with this
            container, and the corresponding value is a list of dictionaries containing
            information about the latest inferred images. Each dictionary has the following keys:
            - "image" (str): The base64-encoded image data.
            - "dimensions" (dict): Image dimensions (width and height).
            - "predictions" (list): A list of predictions or results associated with the image.

        Notes:
            - This method uses the global constant METRICS_INTERVAL to specify the time interval.
        """

        now = time.time()
        start = now - METRICS_INTERVAL
        api_keys = get_cache_model_items().get(self.id, dict()).keys()
        model_ids = []
        for api_key in api_keys:
            mids = get_cache_model_items().get(self.id, dict()).get(api_key, [])
            model_ids.extend(mids)
        num_images = 0
        latest_inferred_images = dict()
        for model_id in model_ids:
            if num_images >= max:
                break
            latest_reqs = cache.zrangebyscore(
                f"inference:{self.id}:{model_id}", min=start, max=now
            )
            for req in latest_reqs:
                images = req["request"]["image"]
                image_dims = req.get("response", {}).get("image", dict())
                predictions = req.get("response", {}).get("predictions", [])
                if images is None or len(images) == 0:
                    continue
                if type(images) is not list:
                    images = [images]
                for image in images:
                    value = None
                    if image["type"] == "base64":
                        value = image["value"]
                    else:
                        loaded_image = load_image_rgb(image)
                        image_bytes = loaded_image.tobytes()
                        image_base64 = base64.b64encode(image_bytes).decode("utf-8")
                        value = image_base64
                    if latest_inferred_images.get(model_id) is None:
                        latest_inferred_images[model_id] = []
                    inference = dict(
                        image=value, dimensions=image_dims, predictions=predictions
                    )
                    latest_inferred_images[model_id].append(inference)
                    num_images += 1
        return latest_inferred_images

    def get_startup_config(self):
        """
        Get the startup configuration for this container.

        Returns:
            dict: A dictionary containing the startup configuration for this container.
        """
        env_vars = self.container.attrs.get("Config", {}).get("Env", {})
        port_bindings = self.container.attrs.get("HostConfig", {}).get(
            "PortBindings", {}
        )
        detached = self.container.attrs.get("HostConfig", {}).get("Detached", False)
        image = self.container.attrs.get("Config", {}).get("Image", "")
        privileged = self.container.attrs.get("HostConfig", {}).get("Privileged", False)
        labels = self.container.attrs.get("Config", {}).get("Labels", {})
        env = []
        for var in env_vars:
            name, value = var.split("=")
            env.append(f"{name}={value}")
        return {
            "env": env,
            "port_bindings": port_bindings,
            "detach": detached,
            "image": image,
            "privileged": privileged,
            "labels": labels,
            # TODO: add device requests
        }