bikes.io.services

22class Service(abc.ABC, pdt.BaseModel, strict=True, frozen=True, extra="forbid"):
23    """Base class for a global service.
24
25    Use services to manage global contexts.
26    e.g., logger object, mlflow client, spark context, ...
27    """
28
29    @abc.abstractmethod
30    def start(self) -> None:
31        """Start the service."""
32
33    def stop(self) -> None:
34        """Stop the service."""
35        # does nothing by default

38class LoggerService(Service):
39    """Service for logging messages.
40
41    https://loguru.readthedocs.io/en/stable/api/logger.html
42
43    Parameters:
44        sink (str): logging output.
45        level (str): logging level.
46        format (str): logging format.
47        colorize (bool): colorize output.
48        serialize (bool): convert to JSON.
49        backtrace (bool): enable exception trace.
50        diagnose (bool): enable variable display.
51        catch (bool): catch errors during log handling.
52    """
53
54    sink: str = "stderr"
55    level: str = "DEBUG"
56    format: str = (
57        "<green>[{time:YYYY-MM-DD HH:mm:ss.SSS}]</green>"
58        "<level>[{level}]</level>"
59        "<cyan>[{name}:{function}:{line}]</cyan>"
60        " <level>{message}</level>"
61    )
62    colorize: bool = True
63    serialize: bool = False
64    backtrace: bool = True
65    diagnose: bool = False
66    catch: bool = True
67
68    @T.override
69    def start(self) -> None:
70        loguru.logger.remove()
71        config = self.model_dump()
72        # use standard sinks or keep the original
73        sinks = {"stderr": sys.stderr, "stdout": sys.stdout}
74        config["sink"] = sinks.get(config["sink"], config["sink"])
75        loguru.logger.add(**config)
76
77    def logger(self) -> loguru.Logger:
78        """Return the main logger.
79
80        Returns:
81            loguru.Logger: the main logger.
82        """
83        return loguru.logger

 86class AlertsService(Service):
 87    """Service for sending notifications.
 88
 89    Require libnotify-bin on Linux systems.
 90
 91    In production, use with Slack, Discord, or emails.
 92
 93    https://plyer.readthedocs.io/en/latest/api.html#plyer.facades.Notification
 94
 95    Parameters:
 96        enable (bool): use notifications or print.
 97        app_name (str): name of the application.
 98        timeout (int | None): timeout in secs.
 99    """
100
101    enable: bool = True
102    app_name: str = "Bikes"
103    timeout: int | None = None
104
105    # Plyer's NOTIFYICONDATAW buffers reserve one UTF-16 unit for the terminator.
106    _MAX_APP_NAME_LENGTH: T.ClassVar[int] = 127
107    _MAX_TITLE_LENGTH: T.ClassVar[int] = 63
108    _MAX_MESSAGE_LENGTH: T.ClassVar[int] = 255
109
110    @T.override
111    def start(self) -> None:
112        pass
113
114    def notify(self, title: str, message: str) -> None:
115        """Send a notification to the system.
116
117        Args:
118            title (str): title of the notification.
119            message (str): message of the notification.
120        """
121        if self.enable:
122            notify_title, notify_message, app_name = title, message, self.app_name
123            if sys.platform == "win32":
124                notify_title = self._truncate(title, self._MAX_TITLE_LENGTH)
125                notify_message = self._truncate(message, self._MAX_MESSAGE_LENGTH)
126                app_name = self._truncate(app_name, self._MAX_APP_NAME_LENGTH)
127            try:
128                notification.notify(
129                    title=notify_title,
130                    message=notify_message,
131                    app_name=app_name,
132                    timeout=self.timeout,
133                )
134            except NotImplementedError:
135                print("Notifications are not supported on this system.")  # noqa: T201  # user-facing fallback
136                self._print(title=title, message=message)
137        else:
138            self._print(title=title, message=message)
139
140    @staticmethod
141    def _truncate(value: str, max_length: int) -> str:
142        """Fit a Windows field in UTF-16 units without splitting a surrogate pair."""
143        encoded = value.encode("utf-16-le")
144        if len(encoded) <= max_length * 2:
145            return value
146        prefix = encoded[: (max_length - 1) * 2].decode("utf-16-le", errors="ignore")
147        return f"{prefix}\N{HORIZONTAL ELLIPSIS}"
148
149    def _print(self, title: str, message: str) -> None:
150        """Print a notification to the system.
151
152        Args:
153            title (str): title of the notification.
154            message (str): message of the notification.
155        """
156        print(f"[{self.app_name}] {title}: {message}")  # noqa: T201  # user-facing fallback

159class MlflowService(Service):
160    """Service for Mlflow tracking and registry.
161
162    Parameters:
163        tracking_uri (str): the URI for the Mlflow tracking server.
164        registry_uri (str): the URI for the Mlflow model registry.
165        experiment_name (str): the name of tracking experiment.
166        registry_name (str): the name of model registry.
167        autolog_disable (bool): disable autologging.
168        autolog_disable_for_unsupported_versions (bool): disable autologging for unsupported versions.
169        autolog_exclusive (bool): If True, enables exclusive autologging.
170        autolog_log_input_examples (bool): If True, logs input examples during autologging.
171        autolog_log_model_signatures (bool): If True, logs model signatures during autologging.
172        autolog_log_models (bool): If True, enables logging of models during autologging.
173        autolog_log_datasets (bool): If True, logs datasets used during autologging.
174        autolog_silent (bool): If True, suppresses all Mlflow warnings during autologging.
175    """
176
177    class RunConfig(pdt.BaseModel, strict=True, frozen=True, extra="forbid"):
178        """Run configuration for Mlflow tracking.
179
180        Parameters:
181            name (str): name of the run.
182            description (str | None): description of the run.
183            tags (dict[str, T.Any] | None): tags for the run.
184            log_system_metrics (bool | None): enable system metrics logging.
185        """
186
187        name: str
188        description: str | None = None
189        tags: dict[str, T.Any] | None = None
190        log_system_metrics: bool | None = True
191
192    # server uri
193    # SQLAlchemy backends are the supported store in MLflow 3: SQLite gives the local
194    # setup the same shape as a production database (Postgres, MySQL) with no server to
195    # run, and it is the only local store the model registry is actually designed for.
196    tracking_uri: str = "sqlite:///mlflow.db"
197    registry_uri: str = "sqlite:///mlflow.db"
198    # experiment
199    experiment_name: str = "bikes"
200    # registry
201    registry_name: str = "bikes"
202    # autolog
203    autolog_disable: bool = False
204    autolog_disable_for_unsupported_versions: bool = False
205    autolog_exclusive: bool = False
206    autolog_log_input_examples: bool = True
207    autolog_log_model_signatures: bool = True
208    autolog_log_models: bool = False
209    autolog_log_datasets: bool = False
210    autolog_silent: bool = False
211
212    @T.override
213    def start(self) -> None:
214        # server uri
215        mlflow.set_tracking_uri(uri=self.tracking_uri)
216        mlflow.set_registry_uri(uri=self.registry_uri)
217        # experiment
218        mlflow.set_experiment(experiment_name=self.experiment_name)
219        # autolog
220        mlflow.autolog(
221            disable=self.autolog_disable,
222            disable_for_unsupported_versions=self.autolog_disable_for_unsupported_versions,
223            exclusive=self.autolog_exclusive,
224            log_input_examples=self.autolog_log_input_examples,
225            log_model_signatures=self.autolog_log_model_signatures,
226            log_datasets=self.autolog_log_datasets,
227            silent=self.autolog_silent,
228        )
229
230    @ctx.contextmanager
231    def run_context(self, run_config: RunConfig) -> T.Generator[mlflow.ActiveRun]:
232        """Yield an active Mlflow run and exit it afterwards.
233
234        Args:
235            run_config (RunConfig): mlflow run parameters.
236
237        Yields:
238            T.Generator[mlflow.ActiveRun, None, None]: active run context. Will be closed at the end of context.
239        """
240        with mlflow.start_run(
241            run_name=run_config.name,
242            tags=run_config.tags,
243            description=run_config.description,
244            log_system_metrics=run_config.log_system_metrics,
245        ) as run:
246            yield run
247
248    def client(self) -> mt.MlflowClient:
249        """Return a new Mlflow client.
250
251        Returns:
252            MlflowClient: the mlflow client.
253        """
254        return mt.MlflowClient(tracking_uri=self.tracking_uri, registry_uri=self.registry_uri)

177    class RunConfig(pdt.BaseModel, strict=True, frozen=True, extra="forbid"):
178        """Run configuration for Mlflow tracking.
179
180        Parameters:
181            name (str): name of the run.
182            description (str | None): description of the run.
183            tags (dict[str, T.Any] | None): tags for the run.
184            log_system_metrics (bool | None): enable system metrics logging.
185        """
186
187        name: str
188        description: str | None = None
189        tags: dict[str, T.Any] | None = None
190        log_system_metrics: bool | None = True