sibi-dst 2025.1.4__py3-none-any.whl → 2025.1.6__py3-none-any.whl

sibi_dst/__init__.py

@@ -1,4 +1,3 @@
-
 try:
     import importlib.metadata as version_reader
 except ImportError:
@@ -8,3 +7,7 @@ try:
     __version__ = version_reader.version("sibi-dst")
 except version_reader.PackageNotFoundError:
     __version__ = "unknown"
+
+__all__ = [
+    "__version__",
+]

sibi_dst/df_helper/__init__.py

@@ -3,11 +3,11 @@ from __future__ import annotations
 from ._df_helper import DfHelper
 from ._parquet_artifact import ParquetArtifact
 from ._parquet_reader import ParquetReader
-from ._artifact_updater_multi_wrapper import ArtifactUpdaterMultiWrapper
+from ._artifact_updater_multi_wrapper import ArtifactUpdaterMultiWrapperThreaded
 
 __all__ = [
     'DfHelper',
     'ParquetArtifact',
     'ParquetReader',
-    'ArtifactUpdaterMultiWrapper',
+    'ArtifactUpdaterMultiWrapperThreaded',
 ]

sibi_dst/df_helper/_artifact_updater_multi_wrapper.py

@@ -1,80 +1,41 @@
-import asyncio
-import logging
 import datetime
+import time
 import random
-from typing import Any, Callable, Dict, List, Optional, Type
+from concurrent.futures import ThreadPoolExecutor, as_completed
+from typing import Any, Callable, Dict, List, Optional, Type, ClassVar
 
-from sibi_dst.utils import Logger
+from sibi_dst.utils import  ManagedResource
 
 
-class ArtifactUpdaterMultiWrapper:
+class ArtifactUpdaterMultiWrapperThreaded(ManagedResource):
     """
-    Simplified wrapper that updates artifacts concurrently using an asyncio.Semaphore.
-
-    Features:
-    - Caps concurrency at max_workers via semaphore
-    - Optionally prioritises tasks via a priority function or static method on artifact classes
-    - Tracks per-artifact completion times
-    - Configurable retry/backoff strategy
-    - Optional metrics integration
-    - Thread-safe within a single asyncio loop
-
-    Usage:
-        wrapper = ArtifactUpdaterMultiWrapper(
-            wrapped_classes={
-                'mydata': [DataArtifactA, DataArtifactB],
-            },
-            max_workers=4,
-            retry_attempts=3,
-            update_timeout_seconds=600,
-            backoff_base=2,
-            backoff_max=60,
-            backoff_jitter=0.1,
-            priority_fn=None,  # or custom
-            metrics_client=None,
-            debug=True,
-            logger=None,
-            artifact_class_kwargs={
-                'fs': my_fs,
-                'parquet_storage_path': 's3://bucket/data',
-                'logger': my_logger,
-                'debug': True,
-            }
-        )
-        await wrapper.update_data('mydata', period='ytd', overwrite=True)
+    Updates artifacts concurrently using a ThreadPoolExecutor.
+
+    This version is refactored for a pure multi-threaded environment, aligning
+    the orchestration model with the underlying threaded workers (DataWrapper).
     """
+    wrapped_classes: Dict[str, List[Type]]
     def __init__(
-        self,
-        wrapped_classes: Dict[str, List[Type]],
-        *,
-        max_workers: int = 3,
-        retry_attempts: int = 3,
-        update_timeout_seconds: int = 600,
-        backoff_base: int = 2,
-        backoff_max: Optional[int] = 60,
-        backoff_jitter: float = 0.1,
-        priority_fn: Optional[Callable[[Type], int]] = None,
-        metrics_client: Any = None,
-        debug: bool = False,
-        logger: Optional[logging.Logger] = None,
-        artifact_class_kwargs: Optional[Dict[str, Any]] = None,
+            self,
+            wrapped_classes: Dict[str, List[Type]],
+            *,
+            max_workers: int = 4,
+            retry_attempts: int = 3,
+            backoff_base: int = 2,
+            backoff_max: int = 60,
+            backoff_jitter: float = 0.1,
+            priority_fn: Optional[Callable[[Type], int]] = None,
+            artifact_class_kwargs: Optional[Dict[str, Any]] = None,
+            **kwargs: Dict[str, Any]
     ) -> None:
+        super().__init__(**kwargs)
         self.wrapped_classes = wrapped_classes
         self.max_workers = max_workers
         self.retry_attempts = retry_attempts
-        self.update_timeout_seconds = update_timeout_seconds
         self.backoff_base = backoff_base
         self.backoff_max = backoff_max
         self.backoff_jitter = backoff_jitter
         self.priority_fn = priority_fn
-        self.metrics_client = metrics_client
-
-        self.debug = debug
-        self.logger = logger or Logger.default_logger(
-            logger_name=self.__class__.__name__,
-            log_level=Logger.DEBUG if debug else Logger.INFO
-        )
-
         # Default artifact init kwargs
         today = datetime.datetime.today() + datetime.timedelta(days=1)
         default_kwargs = {
@@ -82,158 +43,389 @@ class ArtifactUpdaterMultiWrapper:
             'parquet_end_date':   today.strftime('%Y-%m-%d'),
             'logger':             self.logger,
             'debug':              self.debug,
+            'fs':                self.fs,
+            'verbose':            self.verbose,
         }
         self.artifact_class_kwargs = artifact_class_kwargs or default_kwargs.copy()
 
-        # State
+        # State tracking
         self.completion_times: Dict[str, float] = {}
         self.failed: List[str] = []
         self.original_classes: List[Type] = []
 
     def get_artifact_classes(self, data_type: str) -> List[Type]:
-        """
-        Retrieve artifact classes by data type.
-        """
+        """Retrieve artifact classes by data type."""
         self.logger.info(f"Fetching artifact classes for '{data_type}'")
-        if data_type not in self.wrapped_classes:
+        classes = self.wrapped_classes.get(data_type)
+        if not classes:
             raise ValueError(f"Unsupported data type: {data_type}")
-        classes = self.wrapped_classes[data_type]
         self.logger.info(f"Found {len(classes)} artifact classes for '{data_type}'")
         return classes
 
     def estimate_priority(self, artifact_cls: Type) -> int:
         """
-        Determine task priority for ordering. Lower values run first.
+        Determines task priority. Lower values run first.
+        Note: This is a blocking call and will run sequentially before updates start.
         """
         name = artifact_cls.__name__
+        # Custom priority function takes precedence
         if self.priority_fn:
             try:
-                pr = self.priority_fn(artifact_cls)
-                self.logger.debug(f"priority_fn for {name}: {pr}")
-                return pr
+                return self.priority_fn(artifact_cls)
             except Exception as e:
                 self.logger.warning(f"priority_fn error for {name}: {e}")
-        try:
-            fs = self.artifact_class_kwargs.get('fs')
-            path = self.artifact_class_kwargs.get('parquet_storage_path')
-            pr=1
-            if hasattr(artifact_cls, 'get_size_estimate'):
-                pr = artifact_cls.get_size_estimate(fs, path)
-            self.logger.debug(f"Estimated priority for {name}: {pr}")
-            return pr
-        except Exception:
-            return 1
-
-    async def _bounded_update(self, artifact_cls: Type, sem: asyncio.Semaphore, **update_kwargs) -> None:
+
+        # # Fallback to size estimate if available
+        # if hasattr(artifact_cls, 'get_size_estimate'):
+        #     try:
+        #         # This performs blocking I/O
+        #         return artifact_cls(**self.artifact_class_kwargs).get_size_estimate()
+        #
+        #     except Exception as e:
+        #         self.logger.warning(f"get_size_estimate failed for {name}: {e}")
+
+        # Default priority
+        return 999
+
+    def _update_artifact_with_retry(self, artifact_cls: Type, update_kwargs: Dict[str, Any]) -> str:
         """
-        Wrap update_artifact in a semaphore slot to limit concurrency.
+        A blocking worker function that handles instantiation, update, and retries for a single artifact.
+        This function is designed to be run in a ThreadPoolExecutor.
         """
-        async with sem:
-            name = artifact_cls.__name__
-            start = asyncio.get_event_loop().time()
-            self.logger.info(f"Starting update for {name}")
+        name = artifact_cls.__name__
+        self.logger.debug(f"Worker thread starting update for {name}")
+
+        for attempt in range(1, self.retry_attempts + 1):
             try:
-                for attempt in range(1, self.retry_attempts + 1):
-                    try:
-                        artifact = await asyncio.to_thread(
-                            artifact_cls, **self.artifact_class_kwargs
-                        )
-                        await asyncio.wait_for(
-                            asyncio.to_thread(
-                                artifact.update_parquet, **update_kwargs
-                            ),
-                            timeout=self.update_timeout_seconds
-                        )
-                        duration = asyncio.get_event_loop().time() - start
-                        self.completion_times[name] = duration
-                        self.logger.info(f"✅ {name} updated in {duration:.2f}s (attempt {attempt})")
-                        if self.metrics_client:
-                            self.metrics_client.increment('task_succeeded')
-                        return
-                    except asyncio.TimeoutError:
-                        self.logger.warning(f"Timeout on {name}, attempt {attempt}")
-                    except Exception as e:
-                        self.logger.error(f"Error on {name} attempt {attempt}: {e}")
+                # Instantiate and update directly within the worker thread
+                artifact_instance = artifact_cls(**self.artifact_class_kwargs)
+                artifact_instance.update_parquet(**update_kwargs)
+
+                self.logger.info(f"✅ {name} updated successfully on attempt {attempt}")
+                return name  # Return the name on success
 
+            except Exception as e:
+                self.logger.error(f"Error on {name} attempt {attempt}/{self.retry_attempts}: {e}", exc_info=self.debug)
+                if attempt < self.retry_attempts:
                     delay = min(self.backoff_base ** (attempt - 1), self.backoff_max)
                     delay *= 1 + random.uniform(0, self.backoff_jitter)
                     self.logger.info(f"Sleeping {delay:.1f}s before retrying {name}")
-                    await asyncio.sleep(delay)
+                    time.sleep(delay)
 
-            except asyncio.CancelledError:
-                self.logger.warning(f"{name} update cancelled")
-                raise
-
-            # permanent failure
-            self.logger.error(f"✖️  {name} permanently failed after {self.retry_attempts} attempts")
-            if self.metrics_client:
-                self.metrics_client.increment('task_failed')
-            self.failed.append(name)
+        # If all retries fail, raise an exception to be caught by the main loop
+        raise RuntimeError(f"{name} failed after {self.retry_attempts} attempts.")
 
     async def update_data(self, data_type: str, **kwargs: Any) -> None:
         """
-        Entry point to update all artifacts of a given type concurrently.
+        Entry point to update all artifacts of a given type using a ThreadPoolExecutor.
         """
-        self.logger.info(f"Starting update_data for '{data_type}' with kwargs={kwargs}")
+        self.logger.debug(f"Starting multi-threaded update for '{data_type}' with kwargs={kwargs}")
 
-        # RESET STATE
+        # Reset state for this run
         self.completion_times.clear()
         self.failed.clear()
         self.original_classes = self.get_artifact_classes(data_type)
 
-        # NON-DESTRUCTIVE SORTING
-        ordered = sorted(self.original_classes, key=self.estimate_priority)
+        # Sequentially estimate priorities and sort classes before execution
+        self.logger.debug("Estimating priorities to order tasks...")
+        ordered_classes = sorted(self.original_classes, key=self.estimate_priority)
+        self.logger.debug("Priority estimation complete. Submitting tasks to thread pool.")
 
-        sem = asyncio.Semaphore(self.max_workers)
-        tasks = [
-            asyncio.create_task(self._bounded_update(cls, sem, **kwargs))
-            for cls in ordered
-        ]
+        start_time = time.monotonic()
 
-        try:
-            for coro in asyncio.as_completed(tasks):
-                await coro
-        except asyncio.CancelledError:
-            self.logger.warning("update_data was cancelled—aborting remaining retries")
-            for t in tasks:
-                t.cancel()
-            raise
-        finally:
-            total = len(self.original_classes)
-            completed = len(self.completion_times)
-            failed = len(self.failed)
-            self.logger.info(f"All artifacts processed: total={total}, completed={completed}, failed={failed}")
+        with ThreadPoolExecutor(max_workers=self.max_workers) as executor:
+            future_to_class_name = {
+                executor.submit(self._update_artifact_with_retry, cls, kwargs): cls.__name__
+                for cls in ordered_classes
+            }
 
-    def get_update_status(self) -> Dict[str, Any]:
-        """
-        Returns summary status including completion times.
-        """
+            for future in as_completed(future_to_class_name):
+                name = future_to_class_name[future]
+                try:
+                    # result() will re-raise the exception from the worker if one occurred
+                    future.result()
+                    # If no exception, the task succeeded
+                    self.completion_times[name] = time.monotonic() - start_time
+                except Exception as e:
+                    self.logger.error(f"✖️ {name} permanently failed. See error log above.")
+                    self.failed.append(name)
+
+        # Log final status
         total = len(self.original_classes)
-        completed = set(self.completion_times.keys())
-        failed = set(self.failed)
-        pending = {cls.__name__ for cls in self.original_classes} - completed - failed
+        completed = len(self.completion_times)
+        failed_count = len(self.failed)
+        self.logger.info(f"All artifacts processed: total={total}, completed={completed}, failed={failed_count}")
+
+    def get_update_status(self) -> Dict[str, Any]:
+        """Returns a summary status including completion times."""
+        completed_set = set(self.completion_times.keys())
+        failed_set = set(self.failed)
+        pending_set = {cls.__name__ for cls in self.original_classes} - completed_set - failed_set
 
         return {
-            'total': total,
-            'completed': list(completed),
-            'failed':    list(failed),
-            'pending':   list(pending),
+            'total': len(self.original_classes),
+            'completed': list(completed_set),
+            'failed': list(failed_set),
+            'pending': list(pending_set),
             'completion_times': self.completion_times,
         }
 
     @staticmethod
     def format_status_table(status: Dict[str, Any]) -> str:
-        """
-        Formats the status dict into a readable table.
-        """
+        """Formats the status dictionary into a readable table."""
         lines = [
             f"Total: {status['total']}",
-            f"Completed: {len(status['completed'])}  {status['completed']}",
-            f"Failed:    {len(status['failed'])}  {status['failed']}",
-            f"Pending:   {len(status['pending'])}  {status['pending']}",
-            "",
-            "Per-artifact timings:"
+            f"Completed: {len(status['completed'])}",
+            f"Failed:    {len(status['failed'])}",
+            f"Pending:   {len(status['pending'])}",
+            "\nPer-artifact completion times (seconds):"
         ]
-        for name, dur in status['completion_times'].items():
-            lines.append(f"  {name}: {dur:.2f}s")
+        sorted_times = sorted(status['completion_times'].items(), key=lambda item: item[1], reverse=True)
+        for name, duration in sorted_times:
+            lines.append(f"  - {name:<30}: {duration:.2f}s")
+        if status['failed']:
+            lines.append("\nFailed artifacts:")
+            for name in status['failed']:
+                lines.append(f"  - {name}")
         return "\n".join(lines)
+
+
+# import asyncio
+# import logging
+# import datetime
+# import random
+# from typing import Any, Callable, Dict, List, Optional, Type
+#
+# from sibi_dst.utils import Logger
+#
+#
+# class ArtifactUpdaterMultiWrapperAsync:
+#     """
+#     Simplified wrapper that updates artifacts concurrently using an asyncio.Semaphore.
+#
+#     Features:
+#     - Caps concurrency at max_workers via semaphore
+#     - Optionally prioritises tasks via a priority function or static method on artifact classes
+#     - Tracks per-artifact completion times
+#     - Configurable retry/backoff strategy
+#     - Optional metrics integration
+#     - Thread-safe within a single asyncio loop
+#
+#     Usage:
+#         wrapper = ArtifactUpdaterMultiWrapper(
+#             wrapped_classes={
+#                 'mydata': [DataArtifactA, DataArtifactB],
+#             },
+#             max_workers=4,
+#             retry_attempts=3,
+#             update_timeout_seconds=600,
+#             backoff_base=2,
+#             backoff_max=60,
+#             backoff_jitter=0.1,
+#             priority_fn=None,  # or custom
+#             metrics_client=None,
+#             debug=True,
+#             logger=None,
+#             artifact_class_kwargs={
+#                 'fs': my_fs,
+#                 'parquet_storage_path': 's3://bucket/data',
+#                 'logger': my_logger,
+#                 'debug': True,
+#             }
+#         )
+#         await wrapper.update_data('mydata', period='ytd', overwrite=True)
+#     """
+#     def __init__(
+#         self,
+#         wrapped_classes: Dict[str, List[Type]],
+#         *,
+#         max_workers: int = 3,
+#         retry_attempts: int = 3,
+#         update_timeout_seconds: int = 600,
+#         backoff_base: int = 2,
+#         backoff_max: Optional[int] = 60,
+#         backoff_jitter: float = 0.1,
+#         priority_fn: Optional[Callable[[Type], int]] = None,
+#         metrics_client: Any = None,
+#         debug: bool = False,
+#         logger: Optional[logging.Logger] = None,
+#         artifact_class_kwargs: Optional[Dict[str, Any]] = None,
+#     ) -> None:
+#         self.wrapped_classes = wrapped_classes
+#         self.max_workers = max_workers
+#         self.retry_attempts = retry_attempts
+#         self.update_timeout_seconds = update_timeout_seconds
+#         self.backoff_base = backoff_base
+#         self.backoff_max = backoff_max
+#         self.backoff_jitter = backoff_jitter
+#         self.priority_fn = priority_fn
+#         self.metrics_client = metrics_client
+#
+#         self.debug = debug
+#         self.logger = logger or Logger.default_logger(
+#             logger_name=self.__class__.__name__,
+#             log_level=Logger.DEBUG if debug else Logger.INFO
+#         )
+#
+#         # Default artifact init kwargs
+#         today = datetime.datetime.today() + datetime.timedelta(days=1)
+#         default_kwargs = {
+#             'parquet_start_date': today.strftime('%Y-%m-%d'),
+#             'parquet_end_date':   today.strftime('%Y-%m-%d'),
+#             'logger':             self.logger,
+#             'debug':              self.debug,
+#         }
+#         self.artifact_class_kwargs = artifact_class_kwargs or default_kwargs.copy()
+#
+#         # State
+#         self.completion_times: Dict[str, float] = {}
+#         self.failed: List[str] = []
+#         self.original_classes: List[Type] = []
+#
+#     def get_artifact_classes(self, data_type: str) -> List[Type]:
+#         """
+#         Retrieve artifact classes by data type.
+#         """
+#         self.logger.info(f"Fetching artifact classes for '{data_type}'")
+#         if data_type not in self.wrapped_classes:
+#             raise ValueError(f"Unsupported data type: {data_type}")
+#         classes = self.wrapped_classes[data_type]
+#         self.logger.info(f"Found {len(classes)} artifact classes for '{data_type}'")
+#         return classes
+#
+#     def estimate_priority(self, artifact_cls: Type) -> int:
+#         """
+#         Determine task priority for ordering. Lower values run first.
+#         """
+#         name = artifact_cls.__name__
+#         if self.priority_fn:
+#             try:
+#                 pr = self.priority_fn(artifact_cls)
+#                 self.logger.debug(f"priority_fn for {name}: {pr}")
+#                 return pr
+#             except Exception as e:
+#                 self.logger.warning(f"priority_fn error for {name}: {e}")
+#         try:
+#             fs = self.artifact_class_kwargs.get('fs')
+#             path = self.artifact_class_kwargs.get('parquet_storage_path')
+#             pr=1
+#             if hasattr(artifact_cls, 'get_size_estimate'):
+#                 pr = artifact_cls.get_size_estimate(fs, path)
+#             self.logger.debug(f"Estimated priority for {name}: {pr}")
+#             return pr
+#         except Exception:
+#             return 1
+#
+#     async def _bounded_update(self, artifact_cls: Type, sem: asyncio.Semaphore, **update_kwargs) -> None:
+#         """
+#         Wrap update_artifact in a semaphore slot to limit concurrency.
+#         """
+#         async with sem:
+#             name = artifact_cls.__name__
+#             start = asyncio.get_event_loop().time()
+#             self.logger.info(f"Starting update for {name}")
+#             try:
+#                 for attempt in range(1, self.retry_attempts + 1):
+#                     try:
+#                         artifact = await asyncio.to_thread(
+#                             artifact_cls, **self.artifact_class_kwargs
+#                         )
+#                         await asyncio.wait_for(
+#                             asyncio.to_thread(
+#                                 artifact.update_parquet, **update_kwargs
+#                             ),
+#                             timeout=self.update_timeout_seconds
+#                         )
+#                         duration = asyncio.get_event_loop().time() - start
+#                         self.completion_times[name] = duration
+#                         self.logger.info(f"✅ {name} updated in {duration:.2f}s (attempt {attempt})")
+#                         if self.metrics_client:
+#                             self.metrics_client.increment('task_succeeded')
+#                         return
+#                     except asyncio.TimeoutError:
+#                         self.logger.warning(f"Timeout on {name}, attempt {attempt}")
+#                     except Exception as e:
+#                         self.logger.error(f"Error on {name} attempt {attempt}: {e}")
+#
+#                     delay = min(self.backoff_base ** (attempt - 1), self.backoff_max)
+#                     delay *= 1 + random.uniform(0, self.backoff_jitter)
+#                     self.logger.info(f"Sleeping {delay:.1f}s before retrying {name}")
+#                     await asyncio.sleep(delay)
+#
+#             except asyncio.CancelledError:
+#                 self.logger.warning(f"{name} update cancelled")
+#                 raise
+#
+#             # permanent failure
+#             self.logger.error(f"✖️  {name} permanently failed after {self.retry_attempts} attempts")
+#             if self.metrics_client:
+#                 self.metrics_client.increment('task_failed')
+#             self.failed.append(name)
+#
+#     async def update_data(self, data_type: str, **kwargs: Any) -> None:
+#         """
+#         Entry point to update all artifacts of a given type concurrently.
+#         """
+#         self.logger.info(f"Starting update_data for '{data_type}' with kwargs={kwargs}")
+#
+#         # RESET STATE
+#         self.completion_times.clear()
+#         self.failed.clear()
+#         self.original_classes = self.get_artifact_classes(data_type)
+#
+#         # NON-DESTRUCTIVE SORTING
+#         ordered = sorted(self.original_classes, key=self.estimate_priority)
+#
+#         sem = asyncio.Semaphore(self.max_workers)
+#         tasks = [
+#             asyncio.create_task(self._bounded_update(cls, sem, **kwargs))
+#             for cls in ordered
+#         ]
+#
+#         try:
+#             for coro in asyncio.as_completed(tasks):
+#                 await coro
+#         except asyncio.CancelledError:
+#             self.logger.warning("update_data was cancelled—aborting remaining retries")
+#             for t in tasks:
+#                 t.cancel()
+#             raise
+#         finally:
+#             total = len(self.original_classes)
+#             completed = len(self.completion_times)
+#             failed = len(self.failed)
+#             self.logger.info(f"All artifacts processed: total={total}, completed={completed}, failed={failed}")
+#
+#     def get_update_status(self) -> Dict[str, Any]:
+#         """
+#         Returns summary status including completion times.
+#         """
+#         total = len(self.original_classes)
+#         completed = set(self.completion_times.keys())
+#         failed = set(self.failed)
+#         pending = {cls.__name__ for cls in self.original_classes} - completed - failed
+#
+#         return {
+#             'total': total,
+#             'completed': list(completed),
+#             'failed':    list(failed),
+#             'pending':   list(pending),
+#             'completion_times': self.completion_times,
+#         }
+#
+#     @staticmethod
+#     def format_status_table(status: Dict[str, Any]) -> str:
+#         """
+#         Formats the status dict into a readable table.
+#         """
+#         lines = [
+#             f"Total: {status['total']}",
+#             f"Completed: {len(status['completed'])}  {status['completed']}",
+#             f"Failed:    {len(status['failed'])}  {status['failed']}",
+#             f"Pending:   {len(status['pending'])}  {status['pending']}",
+#             "",
+#             "Per-artifact timings:"
+#         ]
+#         for name, dur in status['completion_times'].items():
+#             lines.append(f"  {name}: {dur:.2f}s")
+#         return "\n".join(lines)