bikes.io.services

  1"""Manage global context during execution."""
  2
  3# %% IMPORTS
  4
  5from __future__ import annotations
  6
  7import abc
  8import contextlib as ctx
  9import os
 10import sys
 11import typing as T
 12import urllib.parse
 13import warnings
 14from typing import TYPE_CHECKING
 15
 16import loguru
 17import mlflow
 18import mlflow.tracking as mt
 19import pydantic as pdt
 20from plyer import notification
 21
 22if TYPE_CHECKING:
 23    from loguru import Logger
 24else:
 25    Logger = type(loguru.logger)
 26
 27# Keep the platform notification proxy outside the documented service API.
 28__all__ = ["AlertsService", "LoggerService", "MlflowService", "Service", "redact_uri"]
 29
 30# %% HELPERS
 31
 32
 33def redact_uri(uri: str) -> str:
 34    """Hide the password of a URI before it reaches the logs.
 35
 36    e.g., "postgresql://user:secret@host/db" becomes "postgresql://user:***@host/db".
 37
 38    Args:
 39        uri (str): URI that may embed credentials.
 40
 41    Returns:
 42        str: the same URI without its password.
 43    """
 44    parts = urllib.parse.urlsplit(uri)
 45    if parts.password is None:
 46        return uri
 47    _, _, host = parts.netloc.rpartition("@")
 48    return parts._replace(netloc=f"{parts.username}:***@{host}").geturl()
 49
 50
 51# %% SERVICES
 52
 53
 54class Service(abc.ABC, pdt.BaseModel, strict=True, frozen=True, extra="forbid"):
 55    """Base class for a global service.
 56
 57    Use services to manage global contexts.
 58    e.g., logger object, mlflow client, spark context, ...
 59    """
 60
 61    @abc.abstractmethod
 62    def start(self) -> None:
 63        """Start the service."""
 64
 65    def stop(self) -> None:
 66        """Stop the service."""
 67        # does nothing by default
 68
 69
 70class LoggerService(Service):
 71    """Service for logging messages.
 72
 73    https://loguru.readthedocs.io/en/stable/api/logger.html
 74
 75    Parameters:
 76        sink (str): logging output.
 77        level (str): logging level.
 78        format (str): logging format.
 79        colorize (bool): colorize output.
 80        serialize (bool): convert to JSON.
 81        backtrace (bool): enable exception trace.
 82        diagnose (bool): enable variable display.
 83        catch (bool): catch errors during log handling.
 84    """
 85
 86    sink: str = "stderr"
 87    level: str = "DEBUG"
 88    format: str = (
 89        "<green>[{time:YYYY-MM-DD HH:mm:ss.SSS}]</green>"
 90        "<level>[{level}]</level>"
 91        "<cyan>[{name}:{function}:{line}]</cyan>"
 92        " <level>{message}</level>"
 93    )
 94    colorize: bool = True
 95    serialize: bool = False
 96    backtrace: bool = True
 97    diagnose: bool = False
 98    catch: bool = True
 99
100    @T.override
101    def start(self) -> None:
102        loguru.logger.remove()
103        config = self.model_dump()
104        # use standard sinks or keep the original
105        sinks = {"stderr": sys.stderr, "stdout": sys.stdout}
106        config["sink"] = sinks.get(config["sink"], config["sink"])
107        loguru.logger.add(**config)
108
109    def logger(self) -> Logger:
110        """Return the main logger.
111
112        Returns:
113            Logger: the main logger.
114        """
115        return loguru.logger
116
117
118class AlertsService(Service):
119    """Service for sending notifications.
120
121    Require libnotify-bin on Linux systems.
122
123    In production, use with Slack, Discord, or emails.
124
125    https://plyer.readthedocs.io/en/latest/api.html#plyer.facades.Notification
126
127    Parameters:
128        enable (bool): use notifications or print.
129        app_name (str): name of the application.
130        timeout (int | None): timeout in secs.
131    """
132
133    enable: bool = True
134    app_name: str = "Bikes"
135    timeout: int | None = None
136
137    # Plyer's NOTIFYICONDATAW buffers reserve one UTF-16 unit for the terminator.
138    _MAX_APP_NAME_LENGTH: T.ClassVar[int] = 127
139    _MAX_TITLE_LENGTH: T.ClassVar[int] = 63
140    _MAX_MESSAGE_LENGTH: T.ClassVar[int] = 255
141
142    @T.override
143    def start(self) -> None:
144        pass
145
146    def notify(self, title: str, message: str) -> None:
147        """Send a notification to the system.
148
149        Args:
150            title (str): title of the notification.
151            message (str): message of the notification.
152        """
153        if self.enable:
154            notify_title, notify_message, app_name = title, message, self.app_name
155            if sys.platform == "win32":
156                notify_title = self._truncate(title, self._MAX_TITLE_LENGTH)
157                notify_message = self._truncate(message, self._MAX_MESSAGE_LENGTH)
158                app_name = self._truncate(app_name, self._MAX_APP_NAME_LENGTH)
159            try:
160                # plyer warns while it probes for a backend (dbus, notify-send): the
161                # NotImplementedError fallback below already reports a missing one.
162                with warnings.catch_warnings():
163                    warnings.filterwarnings(action="ignore", category=UserWarning, module="plyer")
164                    notification.notify(
165                        title=notify_title,
166                        message=notify_message,
167                        app_name=app_name,
168                        timeout=self.timeout,
169                    )
170            except NotImplementedError:
171                print("Notifications are not supported on this system.")  # noqa: T201  # user-facing fallback
172                self._print(title=title, message=message)
173        else:
174            self._print(title=title, message=message)
175
176    @staticmethod
177    def _truncate(value: str, max_length: int) -> str:
178        """Fit a Windows field in UTF-16 units without splitting a surrogate pair."""
179        encoded = value.encode("utf-16-le")
180        if len(encoded) <= max_length * 2:
181            return value
182        prefix = encoded[: (max_length - 1) * 2].decode("utf-16-le", errors="ignore")
183        return f"{prefix}\N{HORIZONTAL ELLIPSIS}"
184
185    def _print(self, title: str, message: str) -> None:
186        """Print a notification to the system.
187
188        Args:
189            title (str): title of the notification.
190            message (str): message of the notification.
191        """
192        print(f"[{self.app_name}] {title}: {message}")  # noqa: T201  # user-facing fallback
193
194
195class MlflowService(Service):
196    """Service for Mlflow tracking and registry.
197
198    Parameters:
199        tracking_uri (str): the URI for the Mlflow tracking server.
200        registry_uri (str): the URI for the Mlflow model registry.
201        experiment_name (str): the name of tracking experiment.
202        registry_name (str): the name of model registry.
203        autolog_disable (bool): disable autologging.
204        autolog_disable_for_unsupported_versions (bool): disable autologging for unsupported versions.
205        autolog_exclusive (bool): If True, enables exclusive autologging.
206        autolog_log_input_examples (bool): If True, logs input examples during autologging.
207        autolog_log_model_signatures (bool): If True, logs model signatures during autologging.
208        autolog_log_models (bool): If True, enables logging of models during autologging.
209        autolog_log_datasets (bool): If True, logs datasets used during autologging.
210        autolog_silent (bool): If True, suppresses all Mlflow warnings during autologging.
211    """
212
213    class RunConfig(pdt.BaseModel, strict=True, frozen=True, extra="forbid"):
214        """Run configuration for Mlflow tracking.
215
216        Parameters:
217            name (str): name of the run.
218            description (str | None): description of the run.
219            tags (dict[str, T.Any] | None): tags for the run.
220            log_system_metrics (bool | None): enable system metrics logging.
221        """
222
223        name: str
224        description: str | None = None
225        tags: dict[str, T.Any] | None = None
226        log_system_metrics: bool | None = True
227
228    # server uri
229    # SQLAlchemy backends are the supported store in MLflow 3: SQLite gives the local
230    # setup the same shape as a production database (Postgres, MySQL) with no server to
231    # run. The legacy file store is in maintenance mode and disabled by default.
232    # Explicit YAML wins; otherwise honor MLflow's environment at construction time.
233    tracking_uri: str = pdt.Field(default_factory=lambda: os.environ.get("MLFLOW_TRACKING_URI", "sqlite:///mlflow.db"))
234    registry_uri: str = pdt.Field(
235        default_factory=lambda data: os.environ.get("MLFLOW_REGISTRY_URI", data["tracking_uri"])
236    )
237    # experiment
238    experiment_name: str = "bikes"
239    # registry
240    registry_name: str = "bikes"
241    # autolog
242    autolog_disable: bool = False
243    autolog_disable_for_unsupported_versions: bool = False
244    autolog_exclusive: bool = False
245    autolog_log_input_examples: bool = True
246    autolog_log_model_signatures: bool = True
247    autolog_log_models: bool = False
248    autolog_log_datasets: bool = False
249    autolog_silent: bool = False
250
251    @T.override
252    def __repr_args__(self) -> T.Iterator[tuple[str | None, T.Any]]:
253        # jobs log their services: never print the credentials of a database URI
254        for key, value in super().__repr_args__():
255            yield key, redact_uri(value) if key is not None and key.endswith("_uri") else value
256
257    @T.override
258    def start(self) -> None:
259        # server uri
260        mlflow.set_tracking_uri(uri=self.tracking_uri)
261        mlflow.set_registry_uri(uri=self.registry_uri)
262        # experiment
263        mlflow.set_experiment(experiment_name=self.experiment_name)
264        # autolog
265        mlflow.autolog(
266            disable=self.autolog_disable,
267            disable_for_unsupported_versions=self.autolog_disable_for_unsupported_versions,
268            exclusive=self.autolog_exclusive,
269            log_input_examples=self.autolog_log_input_examples,
270            log_model_signatures=self.autolog_log_model_signatures,
271            log_models=self.autolog_log_models,
272            log_datasets=self.autolog_log_datasets,
273            silent=self.autolog_silent,
274        )
275
276    @ctx.contextmanager
277    def run_context(self, run_config: RunConfig) -> T.Generator[mlflow.ActiveRun]:
278        """Yield an active Mlflow run and exit it afterwards.
279
280        Args:
281            run_config (RunConfig): mlflow run parameters.
282
283        Yields:
284            T.Generator[mlflow.ActiveRun, None, None]: active run context. Will be closed at the end of context.
285        """
286        with mlflow.start_run(
287            run_name=run_config.name,
288            tags=run_config.tags,
289            description=run_config.description,
290            log_system_metrics=run_config.log_system_metrics,
291        ) as run:
292            yield run
293
294    def client(self) -> mt.MlflowClient:
295        """Return a new Mlflow client.
296
297        Returns:
298            MlflowClient: the mlflow client.
299        """
300        return mt.MlflowClient(tracking_uri=self.tracking_uri, registry_uri=self.registry_uri)