streamlit-octostar-utils 0.5.0.dev15__tar.gz → 0.5.0.dev17__tar.gz

{streamlit_octostar_utils-0.5.0.dev15 → streamlit_octostar_utils-0.5.0.dev17}/PKG-INFO

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

{streamlit_octostar_utils-0.5.0.dev15 → streamlit_octostar_utils-0.5.0.dev17}/pyproject.toml

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

{streamlit_octostar_utils-0.5.0.dev15 → streamlit_octostar_utils-0.5.0.dev17}/streamlit_octostar_utils/api_crafter/celery.py

@@ -258,7 +258,7 @@ class CeleryExecutor(object):
         self._last_stall_check = 0
         self._stall_check_interval = 60
 
-        atexit.register(self.close)
+        self._is_owner = False
         self.set_cleanup_task()
         self.register_state_signals()
 
@@ -323,6 +323,9 @@ class CeleryExecutor(object):
 
     def cleanup_task_results(in_dir, out_dir, redis_host, redis_port, task_expires, result_expires):
         logger.info("Starting cleanup of expired task results...")
+        if not os.path.isdir(in_dir) or not os.path.isdir(out_dir):
+            logger.warning(f"Data directories missing (in={os.path.isdir(in_dir)}, out={os.path.isdir(out_dir)}), skipping cleanup")
+            return
         redis_client = redis.StrictRedis(host=redis_host, port=redis_port)
         for file_name in os.listdir(in_dir):
             file_path = os.path.join(in_dir, file_name)
@@ -392,6 +395,8 @@ class CeleryExecutor(object):
         }
 
     def start(self):
+        self._is_owner = True
+        atexit.register(self.close)
         logger.info("Initializing data folders...")
         shutil.rmtree(self.root_folder, ignore_errors=True)
         for folder in [self.in_folder, self.out_folder]:
@@ -591,6 +596,9 @@ class CeleryExecutor(object):
                 logger.error(f"Error checking stall for queue '{queue_name}': {e}")
 
     def close(self):
+        if not self._is_owner:
+            return
+        self._is_owner = False
         self.stop_event.set()
         if self.worker_health_check_thread and self.worker_health_check_thread.is_alive():
             self.worker_health_check_thread.join(timeout=2)
@@ -883,6 +891,13 @@ class CeleryExecutor(object):
             self.get_thread_pool, _poll_task_state, self.app, task_id
         )
 
+    async def get_task_info(self, task_id):
+        def _get_info(celery_app, task_id):
+            return celery_app.AsyncResult(task_id).info
+        return await asyncio.get_running_loop().run_in_executor(
+            self.get_thread_pool, _get_info, self.app, task_id,
+        )
+
     async def get_task_result(self, task_id, remove=False):
         def _try_get_task_data(celery_app, task_id):
             async_result = celery_app.AsyncResult(task_id)
@@ -1010,10 +1025,12 @@ class FastAPICeleryTaskRoute(Route):
                 data = {"task_state": "UNKNOWN", "task_id": task_id}
             elif state in ["AWAITING", "STARTED"]:
                 data = {"task_state": state, "task_id": task_id}
+                if state == "STARTED":
+                    info = await self.celery_executor.get_task_info(task_id)
+                    if info and isinstance(info, dict):
+                        data["progress"] = info
             elif state == "SUCCESS":
                 data = {"task_state": state, "task_id": task_id, "data": result}
-            elif state == "STARTED":
-                data = {"task_status": state, "task_id": task_id}
             else:
                 raise ValueError(f"Unknown task state {state}!")
             return CommonModels.DataResponseModel(data=data)

{streamlit_octostar_utils-0.5.0.dev15 → streamlit_octostar_utils-0.5.0.dev17}/streamlit_octostar_utils/api_crafter/nifi.py

@@ -1446,6 +1446,7 @@ class NifiEntity(object):
         source_entity_relationship=SOURCE_FRAGMENT_ENTITY_RELATIONSHIP,
     ):
         fields = {
+            **{k: v for k, v in self.record if k.startswith("fragment") and v is not None},
             **fields,
             "os_parent_uid": self.record["os_entity_uid"],
             "source_entity_uid": source_entity_uid,