streamlit-octostar-utils 0.5.4.dev4__tar.gz → 0.5.5__tar.gz

{streamlit_octostar_utils-0.5.4.dev4 → streamlit_octostar_utils-0.5.5}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: streamlit-octostar-utils
-Version: 0.5.4.dev4
+Version: 0.5.5
 Summary: 
 License: MIT
 License-File: LICENSE
@@ -26,6 +26,7 @@ Requires-Dist: numpy (>=1.20.0)
 Requires-Dist: octostar-streamlit (>=0.1.25,<0.2.0)
 Requires-Dist: pottery (>=3.0.0,<4.0.0)
 Requires-Dist: presidio-analyzer (>=2.2.0,<3.0.0) ; extra == "nlp"
+Requires-Dist: psutil (>=5.9.0)
 Requires-Dist: py3langid (>=0.2.0,<0.3.0) ; extra == "nlp"
 Requires-Dist: pydantic (>=2.6.4,<3.0.0)
 Requires-Dist: python-multipart (>=0.0.9,<0.0.10)

{streamlit_octostar_utils-0.5.4.dev4 → streamlit_octostar_utils-0.5.5}/pyproject.toml

@@ -5,7 +5,7 @@ include = '\.pyi?$'
 
 [tool.poetry]
 name = "streamlit-octostar-utils"
-version = "0.5.4-dev.4"
+version = "0.5.5"
 description = ""
 license = "MIT"
 authors = ["Octostar"]
@@ -37,6 +37,7 @@ numpy = ">=1.20.0"
 scipy = "^1.10.0"
 rapidfuzz = "^3.5.0"
 celery = "^5.3.0"
+psutil = ">=5.9.0"
 redis = ">=4.0.0,<5.0.0"
 pottery = "^3.0.0"
 slowapi = "^0.1.9"

{streamlit_octostar_utils-0.5.4.dev4 → streamlit_octostar_utils-0.5.5}/streamlit_octostar_utils/api_crafter/celery.py

@@ -122,7 +122,7 @@ class MemoryMonitor:
 
     _DEFAULTS = {
         "sample_interval": 30,
-        "warning_threshold_mb": None,
+        "warn_memory_per_child": None,
         "gc_collect_on_end": True,
         "malloc_trim_on_end": True,
     }
@@ -138,7 +138,7 @@ class MemoryMonitor:
         if config and isinstance(config, dict):
             cfg.update({k: v for k, v in config.items() if v is not None})
         self._sample_interval = cfg["sample_interval"]
-        self._warning_threshold_mb = cfg["warning_threshold_mb"]
+        self._warn_memory_per_child = cfg["warn_memory_per_child"]
         self._gc_collect_on_end = cfg["gc_collect_on_end"]
         self._malloc_trim_on_end = cfg["malloc_trim_on_end"]
 
@@ -173,15 +173,18 @@ class MemoryMonitor:
             return "?"
 
     def _resolve_threshold(self, queue):
-        """Resolve ``warning_threshold_mb`` for *queue*.
+        """Resolve ``warn_memory_per_child`` (KiB) for *queue*.
 
-        ``warning_threshold_mb`` is a dict mapping queue names to
-        per-queue thresholds.  Returns ``None`` when unconfigured.
+        Returns the threshold in MiB (converted from KiB config) for
+        comparison with RSS readings, or ``None`` when unconfigured.
         """
-        threshold = self._warning_threshold_mb
-        if not isinstance(threshold, dict):
+        thresholds = self._warn_memory_per_child
+        if not isinstance(thresholds, dict):
             return None
-        return threshold.get(queue)
+        kib = thresholds.get(queue)
+        if kib is None:
+            return None
+        return kib / 1024
 
     def _try_malloc_trim(self):
         """Ask glibc to return free pages to the OS. No-op on non-Linux."""
@@ -201,8 +204,8 @@ class MemoryMonitor:
 
         logger.info(
             "[MEM] Memory monitor enabled: sample_interval=%s, "
-            "warning_threshold_mb=%s, gc_collect=%s, malloc_trim=%s",
-            self._sample_interval, self._warning_threshold_mb,
+            "gc_collect=%s, malloc_trim=%s",
+            self._sample_interval,
             self._gc_collect_on_end, self._malloc_trim_on_end,
         )
 
@@ -348,12 +351,14 @@ class CeleryQueueConfig:
         max_tasks_in_queue=None,
         max_tasks_per_child=None,
         max_memory_per_child=None,
+        warn_memory_per_child=None,
         **options,
     ):
         self.n_workers = n_workers
         self.max_tasks_in_queue = max_tasks_in_queue
         self.max_tasks_per_child = max_tasks_per_child
-        self.max_memory_per_child = max_memory_per_child  # KiB
+        self.max_memory_per_child = max_memory_per_child  # KiB — hard limit, worker restarted
+        self.warn_memory_per_child = warn_memory_per_child  # KiB — soft limit, logs WARNING
         self.options = options
 
 
@@ -555,9 +560,15 @@ class CeleryExecutor(object):
 
         self._memory_monitor = None
         if memory_monitor:
-            self._memory_monitor = MemoryMonitor(
-                memory_monitor if isinstance(memory_monitor, dict) else {}
-            )
+            mm_cfg = dict(memory_monitor) if isinstance(memory_monitor, dict) else {}
+            thresholds = {
+                q: qc.warn_memory_per_child
+                for q, qc in self.queue_config.items()
+                if qc.warn_memory_per_child is not None
+            }
+            if thresholds:
+                mm_cfg["warn_memory_per_child"] = thresholds
+            self._memory_monitor = MemoryMonitor(mm_cfg)
 
         self._is_owner = False
         self.set_cleanup_task()

{streamlit_octostar_utils-0.5.4.dev4 → streamlit_octostar_utils-0.5.5}/streamlit_octostar_utils/api_crafter/nifi.py

@@ -17,7 +17,6 @@ from starlette.exceptions import HTTPException as StarletteHTTPException
 
 from octostar.utils.workspace import upsert_entities
 from octostar.utils.ontology import fetch_ontology_data
-from octostar.utils.workspace.permissions import get_permissions, PermissionLevel
 from octostar.utils.pipeline import update_processing_status
 
 from octostar.client import make_client
@@ -85,7 +84,6 @@ class NifiEntityModel(BaseModel):
         is_temporary: bool = False
         exception: dict = Field(default_factory=dict)
         last_processor_name: Optional[str] = None
-        fallback_os_workspace: Optional[str] = None
 
     class RecordModel(BaseModel):
         model_config = ConfigDict(extra="allow")
@@ -475,7 +473,6 @@ class NifiContextManager(object):
     def __init__(self, json_data, lazy_sync=True):
         if not json_data:
             raise ValueError("Nifi context manager received list of 0 entities")
-        self.permissions = {}
         self.in_batches = None
         self.out_entities = None
         self.nonlazy_sync_ids = set()
@@ -552,16 +549,6 @@ class NifiContextManager(object):
     def __enter__(self):
         return self
 
-    def get_workspaces_permissions(self, workspace_ids):
-        permissions_to_fetch = list(set(workspace_ids).difference(set(list(self.permissions.keys()))))
-        if permissions_to_fetch:
-            permissions = get_permissions.sync(permissions_to_fetch, client=self.client)
-            self.permissions.update(permissions)
-        permissions = {}
-        for k in workspace_ids:
-            permissions[k] = self.permissions.get(k, PermissionLevel.NONE)
-        return permissions
-
     def request_entity_sync(
         self,
         entity,
@@ -628,16 +615,18 @@ class NifiContextManager(object):
         import logging
         _lock_logger = logging.getLogger(__name__)
 
-        records = []
-        for e in entities:
-            if isinstance(e, dict):
-                records.append(
-                    (e, e.get("entity_timestamp"))
-                )
-            else:
-                records.append(
-                    (e.record, e.request.get("entity_timestamp") if e.request else None)
-                )
+        def _read_entries():
+            entries = []
+            for e in entities:
+                if isinstance(e, dict):
+                    entries.append((e, e.get("entity_timestamp")))
+                else:
+                    entries.append(
+                        (e.record, e.request.get("entity_timestamp") if e.request else None)
+                    )
+            return entries
+
+        records = _read_entries()
 
         long_expiry = (datetime.now(timezone.utc) + timedelta(seconds=timeout)).strftime("%Y-%m-%dT%H:%M:%SZ")
         statuses = [
@@ -671,6 +660,7 @@ class NifiContextManager(object):
         try:
             yield True
         finally:
+            records = _read_entries()
             short_expiry = datetime.now(timezone.utc).strftime("%Y-%m-%dT%H:%M:%SZ")
             statuses = [
                 {
@@ -705,10 +695,8 @@ class NifiContextManager(object):
             self._sync_upsert_entities(entities_to_upsert)
         self._sync_fetch_relationships(entities, fetch_rel_entities, fetch_concept_rels)
 
-        now_ts = datetime.now(timezone.utc).strftime("%Y-%m-%dT%H:%M:%SZ")
         for entity in all_entities_to_modify:
             entity.request["is_temporary"] = False
-            entity.request["entity_timestamp"] = now_ts
 
         for entity in entities:
             entity.sync_params = {}
@@ -970,29 +958,7 @@ class NifiEntity(object):
 
     @property
     def write_os_workspace(self):
-        permissions = self.context.get_workspaces_permissions(
-            [
-                e
-                for e in [
-                    self.record.get("os_workspace"),
-                    self.request.get("fallback_os_workspace"),
-                ]
-                if e
-            ]
-        )
-        if (
-            self.record.get("os_workspace")
-            and (permissions.get(self.record.get("os_workspace")) or PermissionLevel.NONE) >= PermissionLevel.WRITE
-        ):
-            return self.record["os_workspace"]
-        elif (
-            self.request.get("fallback_os_workspace")
-            and (permissions.get(self.request.get("fallback_os_workspace")) or PermissionLevel.NONE)
-            >= PermissionLevel.WRITE
-        ):
-            return self.request["fallback_os_workspace"]
-        else:
-            return None
+        return self.record.get("os_workspace")
 
     @property
     def label(self):
@@ -1210,7 +1176,6 @@ class NifiEntity(object):
             "is_temporary": True,
             "exception": {},
             "last_processor_name": None,
-            "fallback_os_workspace": self.request["fallback_os_workspace"],
         }
         child_entity = NifiEntity(
             self.context,
@@ -1436,6 +1401,7 @@ class NifiEntity(object):
         os_entity_uid=None,
         os_relationship_uid=None,
         os_entity_type=FRAGMENT_ENTITY_NAME,
+        os_parent_uid=None,
         previous_fragment_uid=None,
         previous_fragment_relationship_uid=None,
         previous_fragment_relationship=PREVIOUS_FRAGMENT_RELATIONSHIP,
@@ -1448,7 +1414,7 @@ class NifiEntity(object):
         fields = {
             **{k: v for k, v in self.record.items() if k.startswith("fragment") and v is not None},
             **fields,
-            "os_parent_uid": self.record["os_entity_uid"],
+            "os_parent_uid": os_parent_uid or self.record["os_entity_uid"],
             "source_entity_uid": source_entity_uid,
             "previous_entity_uid": previous_fragment_uid,
             "next_entity_uid": next_fragment_uid,