streamlit-octostar-utils 0.5.4.dev4__tar.gz → 0.5.4.dev5__tar.gz

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

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: streamlit-octostar-utils
-Version: 0.5.4.dev4
+Version: 0.5.4.dev5
 Summary: 
 License: MIT
 License-File: LICENSE

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

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

{streamlit_octostar_utils-0.5.4.dev4 → streamlit_octostar_utils-0.5.4.dev5}/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()