feat: added carboncode for full pipeline with logging

autointent/_callbacks/base.py

@@ -66,3 +66,11 @@ def log_final_metrics(self, metrics: dict[str, Any]) -> None:
         Args:
             metrics: Final metrics.
         """
+
+    @abstractmethod
+    def log_emissions(self, emissions_data: dict[str, Any]) -> None:
+        """Log emissions data.
+
+        Args:
+            emissions_data: Dictionary containing emissions metrics.
+        """

autointent/_callbacks/callback_handler.py

@@ -73,6 +73,14 @@ def log_final_metrics(self, metrics: dict[str, Any]) -> None:
         """
         self.call_events("log_final_metrics", metrics=metrics)
 
+    def log_emissions(self, emissions_data: dict[str, Any]) -> None:
+        """Log emissions data.
+
+        Args:
+            emissions_data: Dictionary containing emissions metrics.
+        """
+        self.call_events("log_emissions", emissions_data=emissions_data)
+
     def call_events(self, event: str, **kwargs: Any) -> None:  # noqa: ANN401
         """Call events for all callbacks.
 

autointent/_callbacks/tensorboard.py

@@ -120,3 +120,29 @@ def end_module(self) -> None:
 
     def end_run(self) -> None:
         """Ends the current run. This method is currently a placeholder."""
+
+    def log_emissions(self, emissions_data: dict[str, Any]) -> None:
+        """Logs emissions data to TensorBoard.
+
+        A new TensorBoard run named `emissions` is created to store the emissions metrics.
+
+        Args:
+            emissions_data: Dictionary containing emissions metrics.
+
+        Raises:
+            RuntimeError: If `start_run` has not been called before logging emissions.
+        """
+        if self.module_writer is None:
+            msg = "start_run must be called before log_emissions."
+            raise RuntimeError(msg)
+
+        log_dir = Path(self.dirpath) / "emissions"
+        self.module_writer = self.writer(log_dir=log_dir)  # type: ignore[no-untyped-call]
+
+        for key, value in emissions_data.items():
+            if isinstance(value, int | float):
+                self.module_writer.add_scalar(key, value)  # type: ignore[no-untyped-call]
+            else:
+                self.module_writer.add_text(key, str(value))  # type: ignore[no-untyped-call]
+
+        self.module_writer.close()  # type: ignore[no-untyped-call]

autointent/_callbacks/wandb.py

@@ -117,3 +117,22 @@ def end_run(self) -> None:
 
         This method is currently a placeholder and does not perform additional operations.
         """
+
+    def log_emissions(self, emissions_data: dict[str, Any]) -> None:
+        """Logs emissions data to W&B.
+
+        A new W&B run named `emissions` is created to store the emissions metrics.
+
+        Args:
+            emissions_data: Dictionary containing emissions metrics.
+        """
+        self.wandb.init(
+            project=self.project_name,
+            group=self.group,
+            name="emissions",
+            config=emissions_data,
+            settings=self.wandb.Settings(x_stats_sampling_interval=self.log_interval_time),
+        )
+
+        self.wandb.log(emissions_data)
+        self.wandb.finish()

autointent/_pipeline/_pipeline.py

@@ -7,6 +7,7 @@
 
 import numpy as np
 import yaml
+from codecarbon import OfflineEmissionsTracker
 from typing_extensions import assert_never
 
 from autointent import Context, Dataset, OptimizationConfig
@@ -144,15 +145,32 @@ def _fit(self, context: Context, sampler: SamplerType) -> None:
         """
         self.context = context
         self._logger.info("starting pipeline optimization...")
+
+        tracker = None
+        if context.logging_config.track_emissions:
+            tracker = OfflineEmissionsTracker(
+                country_iso_code=context.logging_config.country_iso_code or "USA",
+                output_dir=context.logging_config.dirpath,
+            )
+            tracker.start_task()
+
         self.context.callback_handler.start_run(
             run_name=self.context.logging_config.get_run_name(),
             dirpath=self.context.logging_config.dirpath,
             log_interval_time=self.context.logging_config.log_interval_time,
         )
-        for node_type in NodeType:
-            node_optimizer = self.nodes.get(node_type, None)
-            if node_optimizer is not None:
-                node_optimizer.fit(context, sampler)  # type: ignore[union-attr]
+
+        try:
+            for node_type in NodeType:
+                node_optimizer = self.nodes.get(node_type, None)
+                if node_optimizer is not None:
+                    node_optimizer.fit(context, sampler)  # type: ignore[union-attr]
+        finally:
+            if tracker is not None:
+                emissions_data = tracker.stop_task()
+                emissions_data_dict = json.loads(emissions_data.toJSON())
+                context.callback_handler.log_emissions(emissions_data_dict)
+
         self.context.callback_handler.end_run()
 
     def _is_inference(self) -> bool:

autointent/configs/_optimization.py

@@ -55,6 +55,10 @@ class LoggingConfig(BaseModel):
         0.1, description="Sampling interval for the system monitor in seconds for Wandb logger."
     )
     """List of callbacks to report to. If None, no callbacks will be used"""
+    track_emissions: bool = Field(False, description="Whether to track carbon emissions using codecarbon")
+    """Whether to track carbon emissions using codecarbon"""
+    country_iso_code: str | None = Field(None, description="Country ISO code for emissions calculation")
+    """Country ISO code for emissions calculation"""
 
     @property
     def dirpath(self) -> Path:

docs/optimizer_config.schema.json

@@ -281,6 +281,25 @@
                     "description": "Sampling interval for the system monitor in seconds for Wandb logger.",
                     "title": "Log Interval Time",
                     "type": "number"
+                },
+                "track_emissions": {
+                    "default": false,
+                    "description": "Whether to track carbon emissions using codecarbon",
+                    "title": "Track Emissions",
+                    "type": "boolean"
+                },
+                "country_iso_code": {
+                    "anyOf": [
+                        {
+                            "type": "string"
+                        },
+                        {
+                            "type": "null"
+                        }
+                    ],
+                    "default": null,
+                    "description": "Country ISO code for emissions calculation",
+                    "title": "Country Iso Code"
                 }
             },
             "title": "LoggingConfig",
@@ -355,7 +374,9 @@
                 "dump_modules": false,
                 "clear_ram": false,
                 "report_to": null,
-                "log_interval_time": 0.1
+                "log_interval_time": 0.1,
+                "track_emissions": false,
+                "country_iso_code": null
             }
         },
         "embedder_config": {

pyproject.toml

@@ -45,6 +45,7 @@ dependencies = [
     "xxhash (>=3.5.0,<4.0.0)",
     "python-dotenv (>=1.0.1,<2.0.0)",
     "transformers[torch] (>=4.49.0,<5.0.0)",
+    "codecarbon (==2.6)",
 ]
 
 [project.urls]