sibi-dst 2025.1.13__py3-none-any.whl → 2025.8.2__py3-none-any.whl

sibi_dst/__init__.py

@@ -10,4 +10,10 @@ except version_reader.PackageNotFoundError:
 
 __all__ = [
     "__version__",
-]
+]
+
+from . import df_helper as df_helper
+from . import osmnx_helper as osmnx_helper
+from . import geopy_helper as geopy_helper
+from . import utils as sibi_utils
+

sibi_dst/df_helper/__init__.py

@@ -3,8 +3,9 @@ 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 ArtifactUpdaterMultiWrapperThreaded, ArtifactUpdaterMultiWrapperAsync
-
+#from ._artifact_updater_multi_wrapper import ArtifactUpdaterMultiWrapperThreaded, ArtifactUpdaterMultiWrapperAsync
+from ._artifact_updater_async import ArtifactUpdaterMultiWrapperAsync
+from ._artifact_updater_threaded import ArtifactUpdaterMultiWrapperThreaded
 __all__ = [
     'DfHelper',
     'ParquetArtifact',

sibi_dst/df_helper/_artifact_updater_async.py

@@ -0,0 +1,238 @@
+from __future__ import annotations
+
+import asyncio
+import datetime
+import random
+import time
+from contextlib import ExitStack
+from dataclasses import dataclass
+from typing import Any, Callable, Dict, List, Optional, Sequence, Tuple, Type
+
+from sibi_dst.utils import ManagedResource
+
+
+@dataclass(slots=True)
+class _RetryCfg:
+    attempts: int = 3
+    backoff_base: float = 2.0
+    backoff_max: float = 60.0
+    jitter: float = 0.15
+
+
+_ORCHESTRATOR_KEYS = {
+    "retry_attempts",
+    "backoff_base",
+    "backoff_max",
+    "backoff_jitter",
+    "update_timeout_seconds",
+    "max_workers",
+    "priority_fn",
+    "artifact_class_kwargs",
+}
+
+
+def _default_artifact_kwargs(resource: ManagedResource) -> Dict[str, Any]:
+    return {
+        "logger": resource.logger,
+        "debug": resource.debug,
+        "fs": resource.fs,
+        "verbose": resource.verbose,
+    }
+
+
+class ArtifactUpdaterMultiWrapperAsync(ManagedResource):
+    """
+    Backward-compatible async orchestrator with shutdown-aware scheduling.
+    """
+
+    def __init__(
+        self,
+        wrapped_classes: Dict[str, Sequence[Type]],
+        *,
+        max_workers: int = 3,
+        retry_attempts: int = 3,
+        update_timeout_seconds: int = 600,
+        backoff_base: float = 2.0,
+        backoff_max: float = 60.0,
+        backoff_jitter: float = 0.15,
+        priority_fn: Optional[Callable[[Type], int]] = None,
+        artifact_class_kwargs: Optional[Dict[str, Any]] = None,
+        **kwargs: Any,
+    ) -> None:
+        super().__init__(**kwargs)
+        self.wrapped_classes = wrapped_classes
+        self.max_workers = int(max_workers)
+        self.update_timeout_seconds = int(update_timeout_seconds)
+        self.priority_fn = priority_fn
+
+        self._retry = _RetryCfg(
+            attempts=int(retry_attempts),
+            backoff_base=float(backoff_base),
+            backoff_max=float(backoff_max),
+            jitter=float(backoff_jitter),
+        )
+
+        self.artifact_class_kwargs = {
+            **_default_artifact_kwargs(self),
+            **(artifact_class_kwargs or {}),
+        }
+
+        self.completion_secs: Dict[str, float] = {}
+        self.failed: List[str] = []
+
+        # NEW: async stop gate — tripped on cleanup/cancel
+        self._stop = asyncio.Event()
+
+    # Trip stop gate on close paths
+    def _cleanup(self) -> None:
+        try:
+            loop = asyncio.get_running_loop()
+            loop.call_soon_threadsafe(self._stop.set)
+        except RuntimeError:
+            self._stop.set()
+
+    async def _acleanup(self) -> None:
+        self._stop.set()
+
+    # ---- internals -----------------------------------------------------------
+
+    def _classes_for(self, period: str) -> List[Type]:
+        try:
+            classes = list(self.wrapped_classes[period])
+        except KeyError:
+            raise ValueError(f"Unsupported period '{period}'.")
+        if not classes:
+            raise ValueError(f"No artifact classes configured for period '{period}'.")
+        if self.priority_fn:
+            try:
+                classes.sort(key=self.priority_fn)
+            except Exception as e:
+                self.logger.warning(f"priority_fn failed; using listed order: {e}")
+        return classes
+
+    @staticmethod
+    def _split_kwargs(raw: Dict[str, Any]) -> Tuple[Dict[str, Any], Dict[str, Any]]:
+        orch: Dict[str, Any] = {}
+        art: Dict[str, Any] = {}
+        for k, v in raw.items():
+            if k in _ORCHESTRATOR_KEYS:
+                orch[k] = v
+            else:
+                art[k] = v
+        return orch, art
+
+    async def _run_one(self, cls: Type, period: str, sem: asyncio.Semaphore, artifact_kwargs: Dict[str, Any]) -> None:
+        name = cls.__name__
+        if self._stop.is_set() or self.closed:
+            raise asyncio.CancelledError()
+
+        self.logger.info(f"Running {name} with period '{period}'", extra={"artifact": name, "period": period})
+        async with sem:
+            loop = asyncio.get_running_loop()
+            start = loop.time()
+            for attempt in range(1, self._retry.attempts + 1):
+                if self._stop.is_set() or self.closed:
+                    raise asyncio.CancelledError()
+                try:
+                    def _sync_block() -> None:
+                        with ExitStack() as stack:
+                            inst = cls(**self.artifact_class_kwargs)
+                            inst = stack.enter_context(inst)
+                            inst.update_parquet(period=period, **artifact_kwargs)
+
+                    await asyncio.wait_for(asyncio.to_thread(_sync_block), timeout=self.update_timeout_seconds)
+                    dt_secs = loop.time() - start
+                    self.completion_secs[name] = dt_secs
+                    self.logger.info(f"✅ {name} ({period}) in {dt_secs:.2f}s")
+                    return
+
+                except asyncio.TimeoutError:
+                    self.logger.warning(f"Timeout in {name} attempt {attempt}/{self._retry.attempts}")
+                except asyncio.CancelledError:
+                    raise
+                except Exception as e:
+                    self.logger.error(
+                        f"{name} attempt {attempt}/{self._retry.attempts} failed: {e}",
+                        exc_info=self.debug,
+                    )
+
+                if attempt < self._retry.attempts and not self._stop.is_set():
+                    delay = min(self._retry.backoff_base ** (attempt - 1), self._retry.backoff_max)
+                    delay *= 1 + random.uniform(0, self._retry.jitter)
+                    try:
+                        await asyncio.sleep(delay)
+                    except asyncio.CancelledError:
+                        raise
+
+            self.failed.append(name)
+            self.logger.error(f"✖️  {name} permanently failed")
+
+    # ---- public API ----------------------------------------------------------
+
+    async def update_data(self, period: str, **kwargs: Any) -> None:
+        """
+        Backward-compatible:
+          - Accepts orchestrator knobs in kwargs (we consume them).
+          - Forwards only artifact-friendly kwargs to update_parquet.
+        """
+        _, artifact_kwargs = self._split_kwargs(kwargs)
+
+        self.completion_secs.clear()
+        self.failed.clear()
+
+        classes = self._classes_for(period)
+        self.logger.info(
+            f"Starting update of {len(classes)} artifacts for period '{period}'",
+            extra={
+                "action_module_name": self.__class__.__name__,
+                "date_of_update": time.strftime('%Y-%m-%d'),
+                "start_time": time.strftime('%H:%M:%S'),
+                "period": period,
+            },
+        )
+
+        sem = asyncio.Semaphore(self.max_workers)
+        tasks = [asyncio.create_task(self._run_one(cls, period, sem, dict(artifact_kwargs))) for cls in classes]
+
+        try:
+            for t in asyncio.as_completed(tasks):
+                if self._stop.is_set():
+                    break
+                await t
+        except (asyncio.CancelledError, KeyboardInterrupt):
+            self._stop.set()
+            for t in tasks:
+                t.cancel()
+            raise
+        finally:
+            # Drain/cancel everything deterministically
+            for t in tasks:
+                if not t.done():
+                    t.cancel()
+            await asyncio.gather(*tasks, return_exceptions=True)
+
+            self.logger.info(
+                f"Update completed for period: {period}",
+                extra={
+                    "action_module_name": self.__class__.__name__,
+                    "date_of_update": datetime.date.today().strftime('%Y-%m-%d'),
+                    "end_time": datetime.datetime.now().strftime('%H:%M:%S'),
+                    "period": period,
+                },
+            )
+            self.logger.info(
+                f"Artifacts processed: total={len(classes)}, "
+                f"completed={len(self.completion_secs)}, failed={len(self.failed)}"
+            )
+
+    def get_update_status(self) -> Dict[str, Any]:
+        done = set(self.completion_secs)
+        fail = set(self.failed)
+        all_names = {c.__name__ for v in self.wrapped_classes.values() for c in v}
+        return {
+            "total": len(all_names),
+            "completed": sorted(done),
+            "failed": sorted(fail),
+            "pending": sorted(all_names - done - fail),
+            "completion_times": dict(self.completion_secs),
+        }

sibi_dst/df_helper/_artifact_updater_threaded.py

@@ -0,0 +1,195 @@
+from __future__ import annotations
+
+import datetime
+import time
+import random
+import threading
+from concurrent.futures import ThreadPoolExecutor, as_completed
+from contextlib import ExitStack
+from dataclasses import dataclass
+from typing import Any, Callable, Dict, List, Optional, Sequence, Type, Tuple
+
+from sibi_dst.utils import ManagedResource
+
+
+@dataclass(slots=True)
+class _RetryCfg:
+    attempts: int = 3
+    backoff_base: float = 2.0
+    backoff_max: float = 60.0
+    jitter: float = 0.15
+
+
+_ORCHESTRATOR_KEYS = {
+    "retry_attempts",
+    "backoff_base",
+    "backoff_max",
+    "backoff_jitter",
+    "update_timeout_seconds",  # accepted but unused in pure-threads version
+    "max_workers",
+    "priority_fn",
+    "artifact_class_kwargs",
+}
+
+
+def _default_artifact_kwargs(resource: ManagedResource) -> Dict[str, Any]:
+    return {
+        "logger": resource.logger,
+        "debug": resource.debug,
+        "fs": resource.fs,
+        "verbose": resource.verbose,
+    }
+
+
+class ArtifactUpdaterMultiWrapperThreaded(ManagedResource):
+    """
+    Backward-compatible threaded orchestrator with shutdown-aware scheduling.
+    """
+
+    def __init__(
+        self,
+        wrapped_classes: Dict[str, Sequence[Type]],
+        *,
+        max_workers: int = 4,
+        retry_attempts: int = 3,
+        backoff_base: float = 2.0,
+        backoff_max: float = 60.0,
+        backoff_jitter: float = 0.15,
+        priority_fn: Optional[Callable[[Type], int]] = None,
+        artifact_class_kwargs: Optional[Dict[str, Any]] = None,
+        **kwargs: Any,
+    ) -> None:
+        super().__init__(**kwargs)
+        self.wrapped_classes = wrapped_classes
+        self.max_workers = int(max_workers)
+        self.priority_fn = priority_fn
+        self._retry = _RetryCfg(
+            attempts=int(retry_attempts),
+            backoff_base=float(backoff_base),
+            backoff_max=float(backoff_max),
+            jitter=float(backoff_jitter),
+        )
+        self.artifact_class_kwargs = {
+            **_default_artifact_kwargs(self),
+            **(artifact_class_kwargs or {}),
+        }
+        self.completion_secs: Dict[str, float] = {}
+        self.failed: List[str] = []
+
+        # NEW: stop gate — tripped on cleanup / Ctrl-C to stop scheduling & retries
+        self._stop_event = threading.Event()
+
+    # Trip the stop gate when this wrapper is closed
+    def _cleanup(self) -> None:
+        self._stop_event.set()
+
+    def _classes_for(self, period: str) -> List[Type]:
+        try:
+            classes = list(self.wrapped_classes[period])
+        except KeyError:
+            raise ValueError(f"Unsupported period '{period}'.")
+        if not classes:
+            raise ValueError(f"No artifact classes configured for period '{period}'.")
+        if self.priority_fn:
+            try:
+                classes.sort(key=self.priority_fn)
+            except Exception as e:
+                self.logger.warning(f"priority_fn failed; using listed order: {e}")
+        return classes
+
+    @staticmethod
+    def _split_kwargs(raw: Dict[str, Any]) -> Tuple[Dict[str, Any], Dict[str, Any]]:
+        orch: Dict[str, Any] = {}
+        art: Dict[str, Any] = {}
+        for k, v in raw.items():
+            if k in _ORCHESTRATOR_KEYS:
+                orch[k] = v
+            else:
+                art[k] = v
+        return orch, art
+
+    def _run_one(self, cls: Type, period: str, artifact_kwargs: Dict[str, Any]) -> str:
+        name = cls.__name__
+        start = time.monotonic()
+        for attempt in range(1, self._retry.attempts + 1):
+            if self._stop_event.is_set() or self.closed:
+                raise RuntimeError("shutting_down")
+            try:
+                with ExitStack() as stack:
+                    inst = cls(**self.artifact_class_kwargs)
+                    inst = stack.enter_context(inst)
+                    inst.update_parquet(period=period, **artifact_kwargs)
+                self.completion_secs[name] = time.monotonic() - start
+                return name
+            except Exception as e:
+                if attempt < self._retry.attempts and not self._stop_event.is_set():
+                    # interruptible backoff sleep
+                    delay = min(self._retry.backoff_base ** (attempt - 1), self._retry.backoff_max)
+                    delay *= 1 + random.uniform(0, self._retry.jitter)
+                    end = time.monotonic() + delay
+                    while not self._stop_event.is_set() and time.monotonic() < end:
+                        time.sleep(min(0.1, end - time.monotonic()))
+                    continue
+                raise RuntimeError(f"{name} failed after {self._retry.attempts} attempts: {e}") from e
+
+    def update_data(self, period: str, **kwargs: Any) -> None:
+        # Split kwargs to preserve backward compatibility
+        _, artifact_kwargs = self._split_kwargs(kwargs)
+
+        self.completion_secs.clear()
+        self.failed.clear()
+
+        classes = self._classes_for(period)
+
+        executor = ThreadPoolExecutor(max_workers=self.max_workers, thread_name_prefix="artifact-updater")
+        try:
+            fut2name: Dict[Any, str] = {}
+            for cls in classes:
+                if self._stop_event.is_set() or self.closed:
+                    break
+                try:
+                    fut = executor.submit(self._run_one, cls, period, dict(artifact_kwargs))
+                except RuntimeError as e:
+                    if "cannot schedule new futures after shutdown" in str(e).lower():
+                        self.logger.warning("Executor shutting down; halting new submissions.")
+                        break
+                    raise
+                fut2name[fut] = cls.__name__
+
+            for fut in as_completed(fut2name):
+                name = fut2name[fut]
+                try:
+                    fut.result()
+                    self.logger.info(f"✅ {name} ({period}) in {self.completion_secs[name]:.2f}s")
+                except Exception as e:
+                    self.failed.append(name)
+                    self.logger.error(f"✖️  {name} permanently failed: {e}")
+        except KeyboardInterrupt:
+            self.logger.warning("KeyboardInterrupt — stopping scheduling and shutting down.")
+            self._stop_event.set()
+            raise
+        finally:
+            # Ensure queued-but-not-started tasks are canceled
+            executor.shutdown(wait=True, cancel_futures=True)
+
+        self.logger.info(
+            f"Artifacts processed: total={len(classes)}, "
+            f"completed={len(self.completion_secs)}, failed={len(self.failed)}",
+            extra={
+                "date_of_update": time.strftime('%Y-%m-%d'),
+                "start_time": time.strftime('%H:%M:%S'),
+                "period": period,
+            },
+        )
+
+    def get_update_status(self) -> Dict[str, Any]:
+        done = set(self.completion_secs)
+        fail = set(self.failed)
+        all_names = {c.__name__ for v in self.wrapped_classes.values() for c in v}
+        return {
+            "total": len(all_names),
+            "completed": sorted(done),
+            "failed": sorted(fail),
+            "pending": sorted(all_names - done - fail),
+            "completion_times": dict(self.completion_secs),
+        }