change emission tracker to callback (#228)

* change emission tracker to callback

* Update optimizer_config.schema.json

* fix callback

* fix types

* make update_* not abstract

* try print metrics

* fix metrics

* fix test

---------

Co-authored-by: github-actions[bot] <github-actions[bot]@users.noreply.github.com>

Author: Samoed

autointent/_callbacks/__init__.py

@@ -2,10 +2,11 @@
 
 from autointent._callbacks.base import OptimizerCallback
 from autointent._callbacks.callback_handler import CallbackHandler
+from autointent._callbacks.emissions_tracker import EmissionsTrackerCallback
 from autointent._callbacks.tensorboard import TensorBoardCallback
 from autointent._callbacks.wandb import WandbCallback
 
-REPORTERS = {cb.name: cb for cb in [WandbCallback, TensorBoardCallback]}
+REPORTERS = {cb.name: cb for cb in [WandbCallback, TensorBoardCallback, EmissionsTrackerCallback]}
 
 REPORTERS_NAMES = Literal[tuple(REPORTERS.keys())]  # type: ignore[valid-type]
 
@@ -34,6 +35,7 @@ def get_callbacks(reporters: list[str] | None) -> CallbackHandler:
 __all__ = [
     "REPORTERS_NAMES",
     "CallbackHandler",
+    "EmissionsTrackerCallback",
     "OptimizerCallback",
     "TensorBoardCallback",
     "WandbCallback",

autointent/_callbacks/base.py

@@ -66,3 +66,22 @@ def log_final_metrics(self, metrics: dict[str, Any]) -> None:
         Args:
             metrics: Final metrics.
         """
+
+    def update_metrics(self, metrics: dict[str, Any]) -> dict[str, Any]:
+        """Update metrics during training.
+
+        Args:
+            metrics: Metrics to update.
+        """
+        return metrics
+
+    def update_final_metrics(self, metrics: dict[str, Any]) -> dict[str, Any]:
+        """Update final metrics.
+
+        Args:
+            metrics: Final metrics to update.
+
+        Returns:
+            Updated final metrics.
+        """
+        return metrics

autointent/_callbacks/callback_handler.py

@@ -82,3 +82,31 @@ def call_events(self, event: str, **kwargs: Any) -> None:  # noqa: ANN401
         """
         for callback in self.callbacks:
             getattr(callback, event)(**kwargs)
+
+    def update_metrics(self, metrics: dict[str, Any]) -> dict[str, Any]:
+        """Update metrics during training.
+
+        Args:
+            metrics: Metrics to update.
+
+        Returns:
+            Updated metrics.
+        """
+        for callback in self.callbacks:
+            if hasattr(callback, "update_metrics"):
+                metrics = callback.update_metrics(metrics)
+        return metrics
+
+    def update_final_metrics(self, metrics: dict[str, Any]) -> dict[str, Any]:
+        """Update final metrics.
+
+        Args:
+            metrics: Final metrics to update.
+
+        Returns:
+            Updated final metrics.
+        """
+        for callback in self.callbacks:
+            if hasattr(callback, "update_final_metrics"):
+                metrics = callback.update_final_metrics(metrics)
+        return metrics

autointent/_callbacks/emissions_tracker.py

@@ -0,0 +1,98 @@
+"""Emissions tracking functionality for monitoring energy consumption and carbon emissions."""
+
+import json
+import logging
+from pathlib import Path
+from typing import Any
+
+from autointent._callbacks import OptimizerCallback
+
+logger = logging.getLogger(__name__)
+
+
+class EmissionsTrackerCallback(OptimizerCallback):
+    """Class for tracking energy consumption and carbon emissions."""
+
+    name = "emissions_tracker"
+
+    current_module_name: str | None = None
+
+    def __init__(self) -> None:
+        """Initialize the emission tracker."""
+        try:
+            from codecarbon import EmissionsTracker
+        except ImportError as e:
+            msg = (
+                "EmissionsTrackerCallback requires the codecarbon package to be installed. "
+                "Please install it with `pip install autointent[codecarbon]`."
+            )
+            raise ImportError(msg) from e
+        self.emission_tracker = EmissionsTracker
+
+    def start_run(self, run_name: str, dirpath: Path, log_interval_time: float) -> None:  # noqa: ARG002
+        """Start tracking emissions for the entire run.
+
+        Args:
+            run_name: Name of the run.
+            dirpath: Path to the directory where the logs will be saved.
+            log_interval_time: Sampling interval for the system monitor in seconds.
+        """
+        self.tracker = self.emission_tracker(project_name=run_name, measure_power_secs=log_interval_time)
+
+        self.tracker.start()
+        self.current_module_name = None
+
+    def start_module(self, module_name: str, num: int, module_kwargs: dict[str, Any]) -> None:  # noqa: ARG002
+        """Start tracking emissions for a specific task.
+
+        Args:
+            module_name: Name of the task to track emissions for.
+            num: Number of the module.
+            module_kwargs: Module parameters.
+        """
+        self.current_module_name = f"{module_name}_{num}"
+        self.tracker.start_task(self.current_module_name)
+
+    def update_metrics(self, metrics: dict[str, Any]) -> dict[str, float]:
+        """Stop tracking emissions and return the emissions data.
+
+        Returns:
+            Dictionary containing emissions metrics.
+        """
+        emissions_data = self.tracker.stop_task(self.current_module_name)
+        emissions_data_json = json.loads(emissions_data.toJSON())
+        emissions_data_dict = {
+            f"emissions/{k}": v for k, v in emissions_data_json.items() if isinstance(v, int | float)
+        }
+        return emissions_data_dict | metrics
+
+    def update_final_metrics(self, metrics: dict[str, Any]) -> dict[str, Any]:
+        """Update final metrics with emissions data.
+
+        Args:
+            metrics: Final metrics to update.
+
+        Returns:
+            Updated metrics including emissions data.
+        """
+        _ = self.tracker.stop()
+        emissions_data_json = json.loads(self.tracker.final_emissions_data.toJSON())
+        emissions_data_dict = {
+            f"emissions/{k}": v for k, v in emissions_data_json.items() if isinstance(v, int | float)
+        }
+        return {"emissions": emissions_data_dict} | metrics
+
+    def log_value(self, **kwargs: dict[str, Any]) -> None:
+        pass
+
+    def log_metrics(self, metrics: dict[str, Any]) -> None:
+        pass
+
+    def end_module(self) -> None:
+        pass
+
+    def end_run(self) -> None:
+        pass
+
+    def log_final_metrics(self, metrics: dict[str, Any]) -> None:
+        pass

autointent/_callbacks/wandb.py

@@ -112,7 +112,7 @@ def log_final_metrics(self, metrics: dict[str, Any]) -> None:
         }
 
         try:
-            config = metrics["configs"]
+            config = metrics.get("configs")
             self.wandb.init(config=config, **wandb_run_init_args)
             self.wandb.log(metrics)
         except Exception as e:

autointent/_dataset/_dataset.py

@@ -102,7 +102,7 @@ def from_hub(cls, repo_name: str) -> "Dataset":
         """
         from ._reader import DictReader
 
-        splits = load_dataset(repo_name)
+        splits = load_dataset(repo_name, "default")
         mapping = dict(**splits)
         if Split.INTENTS in get_dataset_config_names(repo_name):
             mapping["intents"] = load_dataset(repo_name, Split.INTENTS)[Split.INTENTS].to_list()

autointent/_pipeline/_pipeline.py

@@ -252,7 +252,10 @@ def fit(
                     context.data_handler.test_labels(),
                     predictions,
                 )
-            context.callback_handler.log_final_metrics(context.optimization_info.dump_evaluation_results())
+            all_final_metrics = context.callback_handler.update_final_metrics(
+                context.optimization_info.dump_evaluation_results(),
+            )
+            context.callback_handler.log_final_metrics(all_final_metrics)
 
         return context
 

autointent/modules/scoring/_sklearn/sklearn_scorer.py

@@ -5,6 +5,7 @@
 
 import numpy as np
 import numpy.typing as npt
+from sklearn.base import BaseEstimator
 from sklearn.multioutput import MultiOutputClassifier
 from sklearn.utils import all_estimators
 from typing_extensions import Self
@@ -15,7 +16,7 @@
 from autointent.modules.base import BaseScorer
 
 logger = logging.getLogger(__name__)
-AVAILABLE_CLASSIFIERS = {
+AVAILABLE_CLASSIFIERS: dict[str, type[BaseEstimator]] = {
     name: class_
     for name, class_ in all_estimators(
         type_filter=[
@@ -61,7 +62,7 @@ def __init__(
         self,
         clf_name: str = "LogisticRegression",
         embedder_config: EmbedderConfig | str | dict[str, Any] | None = None,
-        **clf_args: Any,  # noqa: ANN401
+        **clf_args: dict[str, float | str | bool],
     ) -> None:
         """Initialize the SklearnScorer.
 
@@ -89,7 +90,7 @@ def from_context(
         context: Context,
         clf_name: str = "LogisticRegression",
         embedder_config: EmbedderConfig | str | None = None,
-        **clf_args: float | str | bool,
+        **clf_args: dict[str, float | str | bool],
     ) -> Self:
         """Create a SklearnScorer instance using a Context object.
 

autointent/nodes/_node_optimizer.py

@@ -17,7 +17,6 @@
 from autointent import Dataset
 from autointent.context import Context
 from autointent.custom_types import NodeType, SearchSpaceValidationMode
-from autointent.nodes.emissions_tracker import EmissionsTracker
 from autointent.nodes.info import NODES_INFO
 from autointent.schemas.node_validation import ParamSpaceFloat, ParamSpaceInt, ParamSpaceT, SearchSpaceConfig
 
@@ -50,7 +49,6 @@ def __init__(
         self.node_type = node_type
         self.node_info = NODES_INFO[node_type]
         self.target_metric = target_metric
-        self.emissions_tracker = EmissionsTracker(project_name=f"{self.node_info.node_type}")
 
         self.metrics = metrics if metrics is not None else []
         if self.target_metric not in self.metrics:
@@ -143,13 +141,10 @@ def objective(
 
         self._logger.debug("Scoring %s module...", module_name)
 
-        self.emissions_tracker.start_task("module_scoring")
         quality_metrics = module.score(context, metrics=self.metrics)
-        emissions_metrics = self.emissions_tracker.stop_task()
-        all_metrics = {**quality_metrics, **emissions_metrics}
 
         target_metric = quality_metrics[self.target_metric]
-
+        all_metrics = context.callback_handler.update_metrics(quality_metrics)
         context.callback_handler.log_metrics(all_metrics)
         context.callback_handler.end_module()