supervisely 6.73.357__py3-none-any.whl → 6.73.358__py3-none-any.whl

supervisely/nn/inference/inference_request.py

@@ -0,0 +1,412 @@
+from __future__ import annotations
+
+import threading
+import time
+import traceback
+import uuid
+from concurrent.futures import Future, ThreadPoolExecutor
+from functools import partial, wraps
+from typing import Any, Dict, List, Tuple, Union
+
+from supervisely._utils import rand_str
+from supervisely.nn.utils import get_gpu_usage, get_ram_usage
+from supervisely.sly_logger import logger
+from supervisely.task.progress import Progress
+
+
+def generate_uuid(self) -> str:
+    """
+    Generates a unique UUID for the inference request.
+    """
+    return uuid.uuid5(namespace=uuid.NAMESPACE_URL, name=f"{time.time()}-{rand_str(10)}").hex
+
+
+class InferenceRequest:
+    class Stage:
+        PREPARING = "Preparing model for inference..."
+        INFERENCE = "Running inference..."
+        FINISHED = "Finished"
+        CANCELLED = "Cancelled"
+        ERROR = "Error"
+
+    def __init__(
+        self,
+        uuid_: str = None,
+        ttl: Union[int, None] = 60 * 60,
+        manager: InferenceRequestsManager = None,
+    ):
+        if uuid_ is None:
+            uuid_ = uuid.uuid5(namespace=uuid.NAMESPACE_URL, name=f"{time.time()}").hex
+        self._uuid = uuid_
+        self._ttl = ttl
+        self.manager = manager
+        self.context = {}
+        self._lock = threading.Lock()
+        self._stage = InferenceRequest.Stage.PREPARING
+        self._pending_results = []
+        self._final_result = None
+        self._exception = None
+        self._stopped = threading.Event()
+        self.progress = Progress(
+            message=self._stage,
+            total_cnt=1,
+            need_info_log=True,
+            update_task_progress=False,
+            log_extra={"inference_request_uuid": self._uuid},
+        )
+        self._created_at = time.monotonic()
+        self._updated_at = self._created_at
+        self._finished = False
+
+        self.global_progress = None
+        self.global_progress_total = 1
+        self.global_progress_current = 0
+
+    def _updated(self):
+        self._updated_at = time.monotonic()
+
+    @property
+    def uuid(self):
+        return self._uuid
+
+    @property
+    def created_at(self):
+        return self._created_at
+
+    @property
+    def updated_at(self):
+        return self._updated_at
+
+    @property
+    def stage(self):
+        return self._stage
+
+    @property
+    def final_result(self):
+        with self._lock:
+            return self._final_result
+
+    @final_result.setter
+    def final_result(self, result: Any):
+        with self._lock:
+            self._final_result = result
+            self._updated()
+
+    def add_results(self, results: List[Dict]):
+        with self._lock:
+            self._pending_results.extend(results)
+            self._updated()
+
+    def pop_pending_results(self, n: int = None):
+        with self._lock:
+            if len(self._pending_results) == 0:
+                return []
+            if n is None:
+                n = len(self._pending_results)
+            if n > len(self._pending_results):
+                n = len(self._pending_results)
+            results = self._pending_results[:n]
+            self._pending_results = self._pending_results[n:]
+            self._updated()
+            return results
+
+    def pending_num(self):
+        return len(self._pending_results)
+
+    def set_stage(self, stage: str, current: int = None, total: int = None, is_size: bool = False):
+        with self._lock:
+            self._stage = stage
+            self.progress.message = self._stage
+            if current is not None:
+                self.progress.current = current
+            if total is not None:
+                logger.debug("setting total = %s", total)
+                self.progress.total = total
+            if is_size:
+                self.progress.is_size = True
+            self.progress._refresh_labels()
+            self.progress.report_progress()
+            if self._stage == InferenceRequest.Stage.INFERENCE:
+                self.global_progress_total = total
+                self.global_progress_current = current
+                self.manager.global_progress.inference_started(
+                    current=self.global_progress_current,
+                    total=self.global_progress_total,
+                )
+            self._updated()
+
+    def done(self, n=1):
+        with self._lock:
+            self.progress.iters_done_report(n)
+            if self._stage == InferenceRequest.Stage.INFERENCE:
+                self.global_progress_current += n
+                if self.manager is not None:
+                    self.manager.done(n)
+            self._updated()
+
+    @property
+    def exception(self):
+        return self._exception
+
+    @exception.setter
+    def exception(self, exc: Exception):
+        self._exception = exc
+        self.set_stage(InferenceRequest.Stage.ERROR)
+        self._updated()
+
+    def is_inferring(self):
+        return self.stage == InferenceRequest.Stage.INFERENCE
+
+    def stop(self):
+        self._stopped.set()
+        self._updated()
+
+    def is_stopped(self):
+        return self._stopped.is_set()
+
+    def is_finished(self):
+        return self._finished
+
+    def is_expired(self):
+        if self._ttl is None:
+            return False
+        return time.monotonic() - self._updated_at > self._ttl
+
+    def progress_json(self):
+        return {
+            "message": self.progress.message,
+            "status": self.progress.message,
+            "current": self.progress.current,
+            "total": self.progress.total,
+            "is_size": self.progress.is_size,
+        }
+
+    def exception_json(self):
+        if self._exception is None:
+            return None
+        return {
+            "type": type(self._exception).__name__,
+            "message": str(self._exception),
+            "traceback": str(traceback.format_exc()),
+        }
+
+    def to_json(self):
+        return {
+            "uuid": self._uuid,
+            "stage": str(self._stage),
+            "progress": self.progress_json(),
+            "pending_results": self.pending_num(),
+            "final_result": self._final_result is not None,
+            "exception": self.exception_json(),
+            "is_inferring": self.is_inferring(),
+            "stopped": self.is_stopped(),
+            "finished": self._finished,
+            "created_at": self._created_at,
+            "updated_at": self._updated_at,
+        }
+
+    def on_inference_end(self):
+        if self._stage not in (
+            InferenceRequest.Stage.FINISHED,
+            InferenceRequest.Stage.CANCELLED,
+            InferenceRequest.Stage.ERROR,
+        ):
+            if self.is_stopped():
+                self.set_stage(InferenceRequest.Stage.CANCELLED)
+            else:
+                self.set_stage(InferenceRequest.Stage.FINISHED)
+        self._finished = True
+        self._updated()
+
+    def get_usage(self):
+        ram_allocated, ram_total = get_ram_usage()
+        gpu_allocated, gpu_total = get_gpu_usage()
+        return {
+            "gpu_memory": {
+                "allocated": gpu_allocated,
+                "total": gpu_total,
+            },
+            "ram_memory": {
+                "allocated": ram_allocated,
+                "total": ram_total,
+            },
+        }
+
+    def status(self):
+        status_data = self.to_json()
+        for key in ["pending_results", "final_result", "created_at", "updated_at"]:
+            status_data.pop(key, None)
+        status_data.update(self.get_usage())
+        return status_data
+
+
+class GlobalProgress:
+    def __init__(self):
+        self.progress = Progress(message="Ready", total_cnt=1)
+        self._lock = threading.Lock()
+
+    def set_message(self, message: str):
+        with self._lock:
+            if self.progress.message != message:
+                self.progress.message = message
+                self.progress.report_progress()
+
+    def set_ready(self):
+        with self._lock:
+            self.progress.message = "Ready"
+            self.progress.current = 0
+            self.progress.total = 1
+            self.progress.report_progress()
+
+    def done(self, n=1):
+        with self._lock:
+            self.progress.iters_done_report(n)
+        if self.progress.current >= self.progress.total:
+            self.set_ready()
+
+    def inference_started(self, current: int, total: int):
+        with self._lock:
+            if self.progress.message == "Ready":
+                self.progress.total = total
+                self.progress.current = current
+            else:
+                self.progress.total += total
+                self.progress.current += current
+        self.set_message("Inferring model...")
+
+    def inference_finished(self, current: int, total: int):
+        with self._lock:
+            if self.progress.message == "Ready":
+                return
+            self.progress.current = self.progress.current - current
+            self.progress.total = self.progress.total - total
+        if self.progress.current >= self.progress.total:
+            self.set_ready()
+
+    def to_json(self):
+        return {
+            "message": self.progress.message,
+            "current": self.progress.current,
+            "total": self.progress.total,
+        }
+
+
+class InferenceRequestsManager:
+
+    def __init__(self, executor: ThreadPoolExecutor = None):
+        if executor is None:
+            executor = ThreadPoolExecutor(max_workers=1)
+        self._executor = executor
+        self._inference_requests: Dict[str, InferenceRequest] = {}
+        self._lock = threading.Lock()
+        self._stop_event = threading.Event()
+        self._monitor_thread = threading.Thread(target=self.monitor, daemon=True)
+        self._monitor_thread.start()
+        self.global_progress = GlobalProgress()
+
+    def __del__(self):
+        try:
+            self._executor.shutdown(wait=False)
+            self._stop_event.set()
+            self._monitor_thread.join(timeout=5)
+        finally:
+            logger.debug("InferenceRequestsManager was deleted")
+
+    def add(self, inference_request: InferenceRequest):
+        with self._lock:
+            self._inference_requests[inference_request.uuid] = inference_request
+
+    def remove(self, inference_request_uuid: str):
+        with self._lock:
+            inference_request = self._inference_requests.get(inference_request_uuid)
+            if inference_request is not None:
+                inference_request.stop()
+                del self._inference_requests[inference_request_uuid]
+
+    def remove_after(self, inference_request_uuid, wait_time=0):
+        with self._lock:
+            inference_request = self._inference_requests.get(inference_request_uuid)
+            if inference_request is not None:
+                inference_request.stop()
+                inference_request._ttl = wait_time
+                inference_request._updated()
+
+    def get(self, inference_request_uuid: str):
+        if inference_request_uuid is None:
+            return None
+        try:
+            return self._inference_requests[inference_request_uuid]
+        except Exception as ex:
+            raise RuntimeError(
+                f"inference_request_uuid {inference_request_uuid} was given, "
+                f"but there is no such uuid in 'self._inference_requests' ({len(self._inference_requests)} items)"
+            ) from ex
+
+    def create(self, inference_request_uuid: str = None) -> InferenceRequest:
+        inference_request = InferenceRequest(uuid_=inference_request_uuid, manager=self)
+        self.add(inference_request)
+        return inference_request
+
+    def monitor(self):
+        while self._stop_event.is_set() is False:
+            for inference_request_uuid in list(self._inference_requests.keys()):
+                inference_request = self._inference_requests.get(inference_request_uuid)
+                if inference_request is None:
+                    continue
+                if inference_request.is_expired():
+                    self.remove(inference_request_uuid)
+                    logger.debug(f"Expired inference request {inference_request_uuid} was deleted")
+            time.sleep(30)
+
+    def done(self, n=1):
+        with self._lock:
+            self.global_progress.done(n)
+
+    def _on_inference_start(self, inference_request: InferenceRequest):
+        if inference_request.uuid not in self._inference_requests:
+            self.add(inference_request)
+
+    def _on_inference_end(self, future, inference_request_uuid: str):
+        logger.debug("callback: on_inference_end()")
+        inference_request = self._inference_requests.get(inference_request_uuid)
+        if inference_request is not None:
+            inference_request.on_inference_end()
+
+            self.global_progress.inference_finished(
+                current=inference_request.global_progress_current,
+                total=inference_request.global_progress_total,
+            )
+
+    def _handle_error_in_async(self, inference_request_uuid: str, func, args, kwargs):
+        try:
+            return func(*args, **kwargs)
+        except Exception as e:
+            inference_request = self._inference_requests.get(inference_request_uuid, None)
+            if inference_request is not None:
+                inference_request.exception = e
+            logger.error(f"Error in {func.__name__} function: {e}", exc_info=True)
+
+    def schedule_task(self, func, *args, **kwargs) -> Tuple[InferenceRequest, Future]:
+        inference_request = kwargs.get("inference_request", None)
+        if inference_request is None:
+            inference_request = self.create()
+            kwargs["inference_request"] = inference_request
+        self._on_inference_start(inference_request)
+        future = self._executor.submit(
+            self._handle_error_in_async,
+            inference_request.uuid,
+            func,
+            args,
+            kwargs,
+        )
+        end_callback = partial(
+            self._on_inference_end, inference_request_uuid=inference_request.uuid
+        )
+        future.add_done_callback(end_callback)
+        logger.debug("Scheduled task.", extra={"inference_request_uuid": inference_request.uuid})
+        return inference_request, future
+
+    def run(self, func, *args, **kwargs):
+        inference_request, future = self.schedule_task(func, *args, **kwargs)
+        future.result()
+        return inference_request.pop_pending_results()

supervisely/nn/inference/object_detection_3d/object_detection_3d.py

@@ -1,19 +1,22 @@
-from typing import Dict, List, Any
-from supervisely.geometry.cuboid_3d import Cuboid3d
-from supervisely.nn.prediction_dto import PredictionCuboid3d
-from supervisely.annotation.label import Label
-from supervisely.annotation.tag import Tag
-from supervisely.nn.inference.inference import Inference
-from fastapi import Response, Request, status
-from supervisely.sly_logger import logger
 import os
-from supervisely import Api, PointcloudAnnotation, PointcloudObject, PointcloudFigure
-from supervisely.io.fs import silent_remove
+from typing import Any, Dict, List
+
+from fastapi import Request, Response, status
+
+from supervisely import Api, PointcloudAnnotation, PointcloudFigure, PointcloudObject
 from supervisely._utils import rand_str
+from supervisely.annotation.label import Label
+from supervisely.annotation.tag import Tag
 from supervisely.app.content import get_data_dir
+from supervisely.geometry.cuboid_3d import Cuboid3d
+from supervisely.io.fs import silent_remove
+from supervisely.nn.inference.inference import Inference
+from supervisely.nn.prediction_dto import PredictionCuboid3d
 from supervisely.pointcloud_annotation.pointcloud_object_collection import (
     PointcloudObjectCollection,
 )
+from supervisely.sly_logger import logger
+
 
 class ObjectDetection3D(Inference):
     def get_info(self) -> dict:
@@ -23,7 +26,7 @@ class ObjectDetection3D(Inference):
         info["async_video_inference_support"] = False
         info["tracking_on_videos_support"] = False
         info["async_image_inference_support"] = False
-        
+
         # recommended parameters:
         # info["model_name"] = ""
         # info["checkpoint_name"] = ""
@@ -44,7 +47,7 @@ class ObjectDetection3D(Inference):
         raise NotImplementedError(
             "Have to be implemented in child class If sliding_window_mode is 'advanced'."
         )
-    
+
     def _inference_pointcloud_id(self, api: Api, pointcloud_id: int, settings: Dict[str, Any]):
         # 1. download pointcloud
         pcd_path = os.path.join(get_data_dir(), rand_str(10) + ".pcd")
@@ -73,7 +76,9 @@ class ObjectDetection3D(Inference):
         annotation = PointcloudAnnotation(objects, figures)
         return annotation
 
-    def raw_results_from_prediction(self, prediction: List[PredictionCuboid3d]) -> List[Dict[str, Any]]:
+    def raw_results_from_prediction(
+        self, prediction: List[PredictionCuboid3d]
+    ) -> List[Dict[str, Any]]:
         results = []
         for pred in prediction:
             detection_name = pred.class_name
@@ -82,14 +87,16 @@ class ObjectDetection3D(Inference):
             rotation_z = pred.cuboid_3d.rotation.z
             velocity = [0, 0]  # Is not supported now
             detection_score = pred.score
-            results.append({
-                "detection_name": detection_name,
-                "translation": translation,
-                "size": size,
-                "rotation": rotation_z,
-                "velocity": velocity,
-                "detection_score": detection_score,
-            })
+            results.append(
+                {
+                    "detection_name": detection_name,
+                    "translation": translation,
+                    "size": size,
+                    "rotation": rotation_z,
+                    "velocity": velocity,
+                    "detection_score": detection_score,
+                }
+            )
         return results
 
     def serve(self):
@@ -103,7 +110,7 @@ class ObjectDetection3D(Inference):
                 extra={**request.state.state, "api_token": "***"},
             )
             state = request.state.state
-            api : Api = request.state.api
+            api: Api = request.state.api
             settings = self._get_inference_settings(state)
             prediction = self._inference_pointcloud_id(api, state["pointcloud_id"], settings)
             annotation = self.annotation_from_prediction(prediction)
@@ -123,11 +130,11 @@ class ObjectDetection3D(Inference):
                 extra={**request.state.state, "api_token": "***"},
             )
             state = request.state.state
-            api : Api = request.state.api
+            api: Api = request.state.api
             settings = self._get_inference_settings(state)
             annotations = []
             for pcd_id in state["pointcloud_ids"]:
                 prediction = self._inference_pointcloud_id(api, pcd_id, settings)
                 annotation = self.annotation_from_prediction(prediction)
                 annotations.append(annotation.to_json())
-            return annotations
+            return annotations

supervisely/nn/inference/session.py

@@ -459,9 +459,9 @@ class SessionJSON:
                     progress_widget = preparing_cb(
                         message="Downloading infos", total=resp["total"], unit="it"
                     )
-
                 while resp["status"] == "download_info":
                     current = resp["current"]
+                    # pylint: disable=possibly-used-before-assignment
                     progress_widget.update(current - prev_current)
                     prev_current = current
                     resp = self._get_preparing_progress()
@@ -813,7 +813,7 @@ class Session(SessionJSON):
         frames_direction: Literal["forward", "backward"] = None,
         tracker: Literal["bot", "deepsort"] = None,
         batch_size: int = None,
-        preparing_cb = None,
+        preparing_cb=None,
     ) -> AsyncInferenceIterator:
         frame_iterator = super().inference_video_id_async(
             video_id,