streamlit-octostar-utils 0.4.2.dev24__tar.gz → 0.5.0.dev1__tar.gz

{streamlit_octostar_utils-0.4.2.dev24 → streamlit_octostar_utils-0.5.0.dev1}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: streamlit-octostar-utils
-Version: 0.4.2.dev24
+Version: 0.5.0.dev1
 Summary: 
 License: MIT
 License-File: LICENSE

{streamlit_octostar_utils-0.4.2.dev24 → streamlit_octostar_utils-0.5.0.dev1}/pyproject.toml

@@ -5,7 +5,7 @@ include = '\.pyi?$'
 
 [tool.poetry]
 name = "streamlit-octostar-utils"
-version = "0.4.2-dev.24"
+version = "0.5.0-dev.1"
 description = ""
 license = "MIT"
 authors = ["Octostar"]

{streamlit_octostar_utils-0.4.2.dev24 → streamlit_octostar_utils-0.5.0.dev1}/streamlit_octostar_utils/api_crafter/celery.py

@@ -9,7 +9,6 @@ import subprocess
 from fastapi import Query
 import time
 import os
-import pickle
 import atexit
 import redis
 import uuid
@@ -68,7 +67,31 @@ class CeleryQueueConfig:
         self.options = options
 
 
+class TaskResult:
+    """Wrapper for task results that include binary parts alongside JSON data.
+    Tasks returning binary data (e.g. images) should return a TaskResult
+    so that serialized_io writes them as multipart parts instead of attempting
+    JSON serialization on bytes."""
+
+    def __init__(self, data, part=None):
+        self.data = data
+        self.part = part
+
+
 class CelerySerialized:
+    """Serializes task data to a boundary-delimited multipart file.
+
+    Format: metadata JSON part followed by optional binary/streamed parts,
+    separated by boundary markers (the task_id). Replaces pickle entirely.
+    """
+
+    BOUNDARY_PREFIX = b"--"
+    BOUNDARY_SUFFIX = b"\r\n"
+    BOUNDARY_END = b"--\r\n"
+    CONTENT_TYPE_JSON = b"Content-Type: application/json\r\n"
+    CONTENT_TYPE_BYTES = b"Content-Type: application/octet-stream\r\n"
+    HEADER_END = b"\r\n"
+
     def __init__(self, folder, redis_client, data=None):
         self.folder = folder
         self.data = data
@@ -77,18 +100,70 @@ class CelerySerialized:
     def set_task_id(self, task_id):
         self.task_id = task_id
 
-    def dump(self):
+    def _boundary(self):
+        return self.task_id.encode()
+
+    def _write_boundary(self, f):
+        f.write(self.BOUNDARY_PREFIX + self._boundary() + self.BOUNDARY_SUFFIX)
+
+    def _write_end_boundary(self, f):
+        f.write(self.BOUNDARY_PREFIX + self._boundary() + self.BOUNDARY_END)
+
+    def dump(self, parts=None, part_is_list=False):
+        """Write metadata + optional parts in multipart format.
+        parts: optional list of bytes objects to write as additional parts.
+        part_is_list: whether the original part was a list (preserves type on read).
+        """
         assert self.task_id
+        if isinstance(self.data, dict):
+            metadata = self.data
+        else:
+            metadata = {"data": self.data}
+        metadata["part_count"] = len(parts) if parts else 0
+        metadata["part_is_list"] = part_is_list
         with RedisFileLock(self.redis_client, os.path.join(self.folder, self.task_id)):
-            with open(os.path.join(self.folder, self.task_id), "wb") as target_file:
-                pickle.dump(self.data, file=target_file, protocol=pickle.HIGHEST_PROTOCOL)
+            with open(os.path.join(self.folder, self.task_id), "wb") as f:
+                self._write_boundary(f)
+                f.write(self.CONTENT_TYPE_JSON)
+                f.write(self.HEADER_END)
+                f.write(json.dumps(metadata).encode())
+                f.write(b"\r\n")
+                if parts:
+                    for part in parts:
+                        self._write_boundary(f)
+                        f.write(self.CONTENT_TYPE_BYTES)
+                        f.write(self.HEADER_END)
+                        f.write(part)
+                        f.write(b"\r\n")
+                self._write_end_boundary(f)
 
     def load(self):
+        """Read multipart file. Returns (metadata_dict, list_of_bytes_parts)."""
         assert self.task_id
+        boundary = self.BOUNDARY_PREFIX + self._boundary()
+        end_boundary = self.BOUNDARY_PREFIX + self._boundary() + b"--"
         with RedisFileLock(self.redis_client, os.path.join(self.folder, self.task_id)):
-            with open(os.path.join(self.folder, self.task_id), "rb") as source_file:
-                data = pickle.load(source_file)
-        return data
+            with open(os.path.join(self.folder, self.task_id), "rb") as f:
+                raw = f.read()
+        sections = raw.split(boundary)
+        metadata = None
+        parts = []
+        for section in sections:
+            section = section.strip(b"\r\n")
+            if not section or section == b"--":
+                continue
+            header_end = section.find(b"\r\n\r\n")
+            if header_end == -1:
+                continue
+            header = section[:header_end]
+            body = section[header_end + 4:]
+            if body.endswith(b"\r\n"):
+                body = body[:-2]
+            if b"application/json" in header:
+                metadata = json.loads(body)
+            else:
+                parts.append(body)
+        return metadata or {}, parts
 
 
 class CeleryExecutor(object):
@@ -129,7 +204,10 @@ class CeleryExecutor(object):
         self.get_thread_pool = None
         self.set_thread_pool = None
         self.io_thread_pool = None
-        self.queue_threadlocks = {k: threading.Lock() for k in self.queue_config.keys()}
+        self.queue_semaphores = {
+            k: threading.Semaphore(v.max_tasks_in_queue) if v.max_tasks_in_queue else None
+            for k, v in self.queue_config.items()
+        }
 
         # Folder setup
         self.base_folder = Path(base_folder).resolve()
@@ -520,9 +598,16 @@ class CeleryExecutor(object):
             task_id = task.request.id
             serialized_data = CelerySerialized(folder=self.in_folder, redis_client=self.redis_client)
             serialized_data.set_task_id(task_id)
-            data = serialized_data.load()
+            metadata, parts = serialized_data.load()
             del serialized_data
-            args, kwargs = data.get("args", []), data.get("kwargs", {})
+            args, kwargs = metadata.get("args", []), metadata.get("kwargs", {})
+
+            part_count = metadata.get("part_count", 0)
+            if part_count > 0:
+                if metadata.get("part_is_list", part_count > 1):
+                    args = [parts] + args
+                else:
+                    args = [parts[0]] + args
 
             if self.app.conf.task_always_eager:
                 queue = task.request.delivery_info.get("routing_key", self.app.conf.task_default_routing_key)
@@ -533,9 +618,25 @@ class CeleryExecutor(object):
             queue = task.request.delivery_info.get("routing_key", self.app.conf.task_default_routing_key)
             task.request.resources = (self.resource_registry or {}).get(queue, {})
             out_data = task_fn(task, *args, **kwargs)
-            serialized_data = CelerySerialized(folder=self.out_folder, data=out_data, redis_client=self.redis_client)
+            if isinstance(out_data, TaskResult):
+                out_part_is_list = isinstance(out_data.part, list)
+                if out_data.part is None:
+                    out_parts = None
+                elif out_part_is_list:
+                    out_parts = out_data.part
+                else:
+                    out_parts = [out_data.part]
+                serialized_data = CelerySerialized(
+                    folder=self.out_folder, data=out_data.data, redis_client=self.redis_client
+                )
+            else:
+                out_parts = None
+                out_part_is_list = False
+                serialized_data = CelerySerialized(
+                    folder=self.out_folder, data=out_data, redis_client=self.redis_client
+                )
             serialized_data.set_task_id(task_id)
-            serialized_data.dump()
+            serialized_data.dump(parts=out_parts, part_is_list=out_part_is_list)
             del serialized_data
             if os.path.isfile(os.path.join(self.in_folder, task_id)):
                 with RedisFileLock(self.redis_client, os.path.join(self.in_folder, task_id)):
@@ -576,36 +677,23 @@ class CeleryExecutor(object):
 
         return decorator
 
-    async def send_task(self, task_fn, args=[], kwargs={}, **options) -> str:
+    async def send_task(self, task_fn, args=None, kwargs=None, part=None, **options) -> str:
+        args = args if args is not None else []
+        kwargs = kwargs if kwargs is not None else {}
         if self.app.conf.task_always_eager and "dev_preload" not in self.app.conf:
             self.preload_on_worker_init()
             self.app.conf.dev_preload = True
 
-        def _reserve_queue_slot(queue_name):
+        def _check_queue_llen(queue_name):
             if self._queue_stalled.get(queue_name, False):
                 raise CeleryExecutor.QueueStalledException(
                     f"Queue '{queue_name}' is stalled. Service temporarily unavailable."
                 )
-            limit = self.queue_config[queue_name].max_tasks_in_queue
-            if limit:
-                reservation_key = f"queue:reserved:{queue_name}"
-                with self.queue_threadlocks[queue_name]:
-                    queue_count = self.redis_client.llen(queue_name)
-                    reserved_count = int(self.redis_client.get(reservation_key) or 0)
-                    total_count = queue_count + reserved_count
-                    if total_count >= limit:
-                        raise CeleryExecutor.QueueFullException(
-                            f"Queue '{queue_name}' has reached its limit of {limit} tasks!"
-                        )
-                    self.redis_client.incr(reservation_key)
-                return True
-            return False
-
-        def _release_queue_slot(queue_name):
-            limit = self.queue_config[queue_name].max_tasks_in_queue
-            if limit:
-                reservation_key = f"queue:reserved:{queue_name}"
-                self.redis_client.decr(reservation_key)
+            if self.redis_client.llen(queue_name) >= self.queue_config[queue_name].max_tasks_in_queue:
+                raise CeleryExecutor.QueueFullException(
+                    f"Queue '{queue_name}' has reached its limit of "
+                    f"{self.queue_config[queue_name].max_tasks_in_queue} tasks!"
+                )
 
         def _write_task_data(in_folder, task_args, task_kwargs, task_id):
             serialized_data = CelerySerialized(
@@ -629,19 +717,37 @@ class CeleryExecutor(object):
         queue_name = self.app.conf.task_default_routing_key
         queue_name = getattr(task_fn, "queue", queue_name)
         queue_name = options.get("queue", queue_name)
-        reserved = False
+
+        sem = self.queue_semaphores.get(queue_name)
+        acquired = False
+        if sem is not None:
+            if not sem.acquire(blocking=False):
+                raise CeleryExecutor.QueueFullException(
+                    f"Queue '{queue_name}' has reached its limit of "
+                    f"{self.queue_config[queue_name].max_tasks_in_queue} tasks!"
+                )
+            acquired = True
+
         try:
-            reserved = await asyncio.get_running_loop().run_in_executor(
-                self.set_thread_pool, _reserve_queue_slot, queue_name
-            )
-            await asyncio.get_running_loop().run_in_executor(
-                self.io_thread_pool,
-                _write_task_data,
-                self.in_folder,
-                args,
-                kwargs,
-                task_id,
-            )
+            if acquired:
+                await asyncio.get_running_loop().run_in_executor(
+                    self.set_thread_pool, _check_queue_llen, queue_name
+                )
+
+            if part is not None:
+                await self._write_task_data_with_part(
+                    task_id, args, kwargs, part
+                )
+            else:
+                await asyncio.get_running_loop().run_in_executor(
+                    self.io_thread_pool,
+                    _write_task_data,
+                    self.in_folder,
+                    args,
+                    kwargs,
+                    task_id,
+                )
+
             await asyncio.get_running_loop().run_in_executor(
                 self.set_thread_pool, _send_task, task_fn, task_id, options
             )
@@ -652,11 +758,61 @@ class CeleryExecutor(object):
             logger.info(f"Cancelling task {task_id} due to disconnect!")
             await self.terminate_task(task_id)
             raise
+        except Exception:
+            try:
+                await self.terminate_task(task_id)
+            except Exception:
+                pass
+            raise
         finally:
-            if reserved:
-                await asyncio.get_running_loop().run_in_executor(self.set_thread_pool, _release_queue_slot, queue_name)
+            if acquired:
+                sem.release()
         return task_id
 
+    async def _write_task_data_with_part(self, task_id, args, kwargs, part):
+        """Write task data with a streamed part to the multipart file.
+        The part becomes the first arg on the worker side.
+        """
+        boundary = CelerySerialized.BOUNDARY_PREFIX + task_id.encode()
+        boundary_line = boundary + CelerySerialized.BOUNDARY_SUFFIX
+        end_boundary_line = boundary + CelerySerialized.BOUNDARY_END
+
+        is_list = isinstance(part, list)
+        items = part if is_list else [part]
+        part_count = len(items)
+
+        metadata = {"args": args, "kwargs": kwargs, "part_count": part_count, "part_is_list": is_list}
+        metadata_bytes = json.dumps(metadata).encode()
+
+        file_path = os.path.join(self.in_folder, task_id)
+        with open(file_path, "wb") as f:
+            f.write(boundary_line)
+            f.write(CelerySerialized.CONTENT_TYPE_JSON)
+            f.write(CelerySerialized.HEADER_END)
+            f.write(metadata_bytes)
+            f.write(b"\r\n")
+
+            for item in items:
+                f.write(boundary_line)
+                f.write(CelerySerialized.CONTENT_TYPE_BYTES)
+                f.write(CelerySerialized.HEADER_END)
+                if hasattr(item, "__aiter__"):
+                    async for chunk in item:
+                        f.write(chunk if isinstance(chunk, bytes) else chunk.encode())
+                elif hasattr(item, "read"):
+                    while True:
+                        chunk = await item.read(65536)
+                        if not chunk:
+                            break
+                        f.write(chunk if isinstance(chunk, bytes) else chunk.encode())
+                elif isinstance(item, bytes):
+                    f.write(item)
+                else:
+                    raise TypeError(f"Unsupported part item type: {type(item)}")
+                f.write(b"\r\n")
+
+            f.write(end_boundary_line)
+
     async def terminate_task(self, task_id):
         def _terminate_task(celery_app, task_id):
             celery_app.control.revoke(task_id, terminate=True)
@@ -710,8 +866,15 @@ class CeleryExecutor(object):
         def _read_task_data(out_folder, task_id):
             serialized_data = CelerySerialized(folder=out_folder, redis_client=self.redis_client)
             serialized_data.set_task_id(task_id)
-            result = serialized_data.load()
-            return result
+            metadata, parts = serialized_data.load()
+            data = metadata.get("data", metadata)
+            part_count = metadata.get("part_count", 0)
+            if part_count > 0:
+                if metadata.get("part_is_list", part_count > 1):
+                    return TaskResult(data=data, part=parts)
+                else:
+                    return TaskResult(data=data, part=parts[0])
+            return data
 
         def _remove_task_data(celery_app, in_folder, out_folder, task_id):
             celery_app.AsyncResult(task_id).forget()
@@ -737,8 +900,10 @@ class CeleryExecutor(object):
             )
         return result
 
-    async def send_and_wait_task(self, task_fn, args=[], kwargs={}, timeout=60, **options):
-        task_id = await self.send_task(task_fn, args, kwargs, **options)
+    async def send_and_wait_task(self, task_fn, args=None, kwargs=None, part=None, timeout=60, **options):
+        args = args if args is not None else []
+        kwargs = kwargs if kwargs is not None else {}
+        task_id = await self.send_task(task_fn, args, kwargs, part=part, **options)
         ready = False
         state = None
         start_time = time.time()

{streamlit_octostar_utils-0.4.2.dev24 → streamlit_octostar_utils-0.5.0.dev1}/streamlit_octostar_utils/api_crafter/fastapi.py

@@ -1,5 +1,5 @@
 import asyncio
-from fastapi import Request, Body
+from fastapi import Request
 from fastapi.responses import JSONResponse, StreamingResponse
 from pydantic import BaseModel
 from typing import List, Optional, Literal, Any
@@ -17,8 +17,6 @@ import traceback
 from copy import copy
 import logging
     
-from octostar.client import make_client
-
 MAX_ERROR_MESSAGE_BYTES = 256
 MAX_ERROR_TRACEBACK_BYTES = 10240
 
@@ -99,100 +97,6 @@ class Route(ABC):
         return self
 
 
-class OctostarRoute(Route):
-    def __init__(self, app, tasks_routes, celery_executor=None, router=None):
-        self.app = app
-        self._router = router
-        self.routed_funcs = []
-        self.tasks_routes = tasks_routes
-        self.celery_executor = celery_executor
-        self.endpoints = {}
-        self.define_routes()
-
-    def register_route(self, op, octostar_task):
-        self.endpoints[op.strip("/")] = octostar_task
-
-    def define_routes(self):
-        if self.celery_executor:
-
-            @Route.route(self, path="/task-state/{task_id}")
-            async def get_task_status(task_id: str) -> JSONResponse:
-                task_status = await self.tasks_routes.get_task(task_id, pop=False)
-                task_status = task_status.model_dump(mode="json")["data"]["task_state"]
-                return JSONResponse(task_status)
-
-            @Route.route(self, path="/task-result/{task_id}")
-            async def get_task_result(task_id: str) -> JSONResponse:
-                return_data = await self.tasks_routes.get_task(task_id, pop=True)
-                return_data = return_data.model_dump(mode="json")["data"]["data"]
-                return JSONResponse(return_data)
-
-            @Route.route(self, path="/{op}", methods=["POST"])
-            async def send_task(
-                op: str,
-                os_context: dict = Body(...),
-                jwt: str = Body(...),
-                params: dict = Body(dict()),
-            ) -> str:
-                """
-                Any request coming from Octostar (e.g. manifest) should enter from here.
-                """
-                path_params = []
-                op = op.split("/")
-                if len(op) > 1:
-                    path_params = op[1:]
-                op = op[0]
-                query_params = params
-                client = make_client(jwt)
-                if op not in self.endpoints.keys():
-                    raise StarletteHTTPException(401, f"Route {op} is forbidden for NiFi.")
-                task_id = await self.celery_executor.send_task(
-                    self.endpoints[op], args=[os_context, client, query_params]
-                )
-                return task_id
-
-        else:
-
-            @Route.route(self, path="/{op}", methods=["POST"])
-            async def call_task(
-                op: str,
-                os_context: dict = Body(...),
-                jwt: str = Body(...),
-                params: dict = Body(dict()),
-            ) -> str:
-                """
-                Any request coming from Octostar (e.g. manifest) should enter from here.
-                """
-                path_params = []
-                op = op.split("/")
-                if len(op) > 1:
-                    path_params = op[1:]
-                op = op[0]
-                query_params = params
-                client = make_client(jwt)
-                if op not in self.endpoints.keys():
-                    raise StarletteHTTPException(401, f"Route {op} is forbidden for NiFi.")
-                result = await self.endpoints[op](os_context, client, query_params)
-                return result
-
-    @staticmethod
-    def octostar_task(celery_executor, *args, **opts):
-        def decorator(func):
-            if celery_executor:
-                serialized_func = celery_executor.serialized_io(func)
-                task_func = celery_executor.app.task(*args, **opts)(serialized_func)
-            else:
-
-                @wraps(func)
-                def octostar_func(*args, **kwargs):
-                    return func(None, *args, **kwargs)
-
-                task_func = octostar_func
-            return task_func
-
-        return decorator
-
-
 class CommonModels(object):
     class OKResponseModel(BaseModel):
         message: str = "OK"

{streamlit_octostar_utils-0.4.2.dev24 → streamlit_octostar_utils-0.5.0.dev1}/streamlit_octostar_utils/api_crafter/nifi.py

@@ -1216,18 +1216,20 @@ class NifiRoute(Route):
             op = op[0]
             query_params = request.query_params
             processor_suffix = query_params["processor_suffix"]
-            body = await request.json()
             processor_name = "processor." + self.processor_name + "." + op + "." + processor_suffix
             if op not in self.endpoints.keys():
                 raise StarletteHTTPException(403, f"Route {op} is forbidden for NiFi.")
-            task_id = await self.celery_executor.send_task(self.endpoints[op], args=[body, processor_name])
+            task_id = await self.celery_executor.send_task(
+                self.endpoints[op], args=[processor_name], part=request.stream()
+            )
             return task_id
 
     @staticmethod
     def nifi_task(celery_executor, *args, **opts):
         def decorator(func):
             @wraps(func)
-            def nifi_func(task, body, processor_name, *args, **kwargs):
+            def nifi_func(task, body_bytes, processor_name, *args, **kwargs):
+                body = json.loads(body_bytes)
                 with NifiContextManager(body) as nifi_context:
                     entity_batches = nifi_context.receive_input(body, processor_name)
                     entity_batches = func(

{streamlit_octostar_utils-0.4.2.dev24 → streamlit_octostar_utils-0.5.0.dev1}/streamlit_octostar_utils/ontology/inheritance.py

@@ -15,13 +15,13 @@ def get_label_keys(type, ontology):
     return list(label_keys.keys())
 
 def get_label(record, ontology):
-    if record.get("os_materialized_label") not in (None, ""):
-        return str(record.get("os_materialized_label"))
-    if record.get("os_custom_label") not in (None, ""):
-        return str(record.get("os_custom_label"))
+    if record.get("os_entity_label_materialized") not in (None, ""):
+        return str(record.get("os_entity_label_materialized")).strip()
+    if record.get("os_entity_label") not in (None, ""):
+        return str(record.get("os_entity_label")).strip()
     label_keys = get_label_keys(record["entity_type"], ontology)
     fields = [record.get(field) for field in label_keys]
     fields = [f for f in fields if f is not None]
     fields = [str(f) for f in fields if f]
     label = " ".join(fields)
-    return label or None
+    return label.strip() or None