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    @T.override
106    def start(self) -> None:
107        pass
108
109    def notify(self, title: str, message: str) -> None:
110        """Send a notification to the system.
111
112        Args:
113            title (str): title of the notification.
114            message (str): message of the notification.
115        """
116        if self.enable:
117            notification.notify(
118                title=title, message=message, app_name=self.app_name, timeout=self.timeout
119            )
120        else:
121            print(f"[{self.app_name}] {title}: {message}")

124class MlflowService(Service):
125    """Service for Mlflow tracking and registry.
126
127    Parameters:
128        tracking_uri (str): the URI for the Mlflow tracking server.
129        registry_uri (str): the URI for the Mlflow model registry.
130        experiment_name (str): the name of tracking experiment.
131        registry_name (str): the name of model registry.
132        autolog_disable (bool): disable autologging.
133        autolog_disable_for_unsupported_versions (bool): disable autologging for unsupported versions.
134        autolog_exclusive (bool): If True, enables exclusive autologging.
135        autolog_log_input_examples (bool): If True, logs input examples during autologging.
136        autolog_log_model_signatures (bool): If True, logs model signatures during autologging.
137        autolog_log_models (bool): If True, enables logging of models during autologging.
138        autolog_log_datasets (bool): If True, logs datasets used during autologging.
139        autolog_silent (bool): If True, suppresses all Mlflow warnings during autologging.
140    """
141
142    class RunConfig(pdt.BaseModel, strict=True, frozen=True, extra="forbid"):
143        """Run configuration for Mlflow tracking.
144
145        Parameters:
146            name (str): name of the run.
147            description (str | None): description of the run.
148            tags (dict[str, T.Any] | None): tags for the run.
149            log_system_metrics (bool | None): enable system metrics logging.
150        """
151
152        name: str
153        description: str | None = None
154        tags: dict[str, T.Any] | None = None
155        log_system_metrics: bool | None = True
156
157    # server uri
158    tracking_uri: str = "./mlruns"
159    registry_uri: str = "./mlruns"
160    # experiment
161    experiment_name: str = "bikes"
162    # registry
163    registry_name: str = "bikes"
164    # autolog
165    autolog_disable: bool = False
166    autolog_disable_for_unsupported_versions: bool = False
167    autolog_exclusive: bool = False
168    autolog_log_input_examples: bool = True
169    autolog_log_model_signatures: bool = True
170    autolog_log_models: bool = False
171    autolog_log_datasets: bool = False
172    autolog_silent: bool = False
173
174    @T.override
175    def start(self) -> None:
176        # server uri
177        mlflow.set_tracking_uri(uri=self.tracking_uri)
178        mlflow.set_registry_uri(uri=self.registry_uri)
179        # experiment
180        mlflow.set_experiment(experiment_name=self.experiment_name)
181        # autolog
182        mlflow.autolog(
183            disable=self.autolog_disable,
184            disable_for_unsupported_versions=self.autolog_disable_for_unsupported_versions,
185            exclusive=self.autolog_exclusive,
186            log_input_examples=self.autolog_log_input_examples,
187            log_model_signatures=self.autolog_log_model_signatures,
188            log_datasets=self.autolog_log_datasets,
189            silent=self.autolog_silent,
190        )
191
192    @ctx.contextmanager
193    def run_context(self, run_config: RunConfig) -> T.Generator[mlflow.ActiveRun, None, None]:
194        """Yield an active Mlflow run and exit it afterwards.
195
196        Args:
197            run (str): run parameters.
198
199        Yields:
200            T.Generator[mlflow.ActiveRun, None, None]: active run context. Will be closed as the end of context.
201        """
202        with mlflow.start_run(
203            run_name=run_config.name,
204            tags=run_config.tags,
205            description=run_config.description,
206            log_system_metrics=run_config.log_system_metrics,
207        ) as run:
208            yield run
209
210    def client(self) -> mt.MlflowClient:
211        """Return a new Mlflow client.
212
213        Returns:
214            MlflowClient: the mlflow client.
215        """
216        return mt.MlflowClient(tracking_uri=self.tracking_uri, registry_uri=self.registry_uri)

142    class RunConfig(pdt.BaseModel, strict=True, frozen=True, extra="forbid"):
143        """Run configuration for Mlflow tracking.
144
145        Parameters:
146            name (str): name of the run.
147            description (str | None): description of the run.
148            tags (dict[str, T.Any] | None): tags for the run.
149            log_system_metrics (bool | None): enable system metrics logging.
150        """
151
152        name: str
153        description: str | None = None
154        tags: dict[str, T.Any] | None = None
155        log_system_metrics: bool | None = True