Feat/code carbon each node (#175)

Darinochka · github-actions[bot] · web-flow · commit e848aaa128b9 · 2025-04-02T18:34:43.000+03:00
* feat: update codecarbon

* feat: update codecarbon

* feat: added codecarbon

* Update optimizer_config.schema.json

* fix: fixed import mypy

* fix: codecarbon package

* fix: only float\integer log

* fix: codecarbon package

* fix: mypy

* fix: test

* fix: delete emissions

* fix: test

---------

Co-authored-by: github-actions[bot] &lt;github-actions[bot]@users.noreply.github.com&gt;
diff --git a/autointent/nodes/_node_optimizer.py b/autointent/nodes/_node_optimizer.py
@@ -17,6 +17,7 @@
 from autointent import Dataset
 from autointent.context import Context
 from autointent.custom_types import NodeType, SamplerType, SearchSpaceValidationMode
+from autointent.nodes.emissions_tracker import EmissionsTracker
 from autointent.nodes.info import NODES_INFO
 
 
@@ -67,6 +68,7 @@ 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:
@@ -141,8 +143,13 @@ def objective(
         context.callback_handler.start_module(module_name=module_name, num=self._counter, module_kwargs=config)
 
         self._logger.debug("Scoring %s module...", module_name)
-        all_metrics = module.score(context, metrics=self.metrics)
-        target_metric = all_metrics[self.target_metric]
+
+        self.emissions_tracker.start_task("module_scoring")
+        final_metrics = module.score(context, metrics=self.metrics)
+        emissions_metrics = self.emissions_tracker.stop_task()
+        all_metrics = {**final_metrics, **emissions_metrics}
+
+        target_metric = final_metrics[self.target_metric]
 
         context.callback_handler.log_metrics(all_metrics)
         context.callback_handler.end_module()
@@ -161,7 +168,7 @@ def objective(
             config,
             target_metric,
             self.target_metric,
-            all_metrics,
+            final_metrics,
             module.get_assets(),  # retriever name / scores / predictions
             module_dump_dir,
             module=module if not context.is_ram_to_clear() else None,
diff --git a/autointent/nodes/emissions_tracker.py b/autointent/nodes/emissions_tracker.py
@@ -0,0 +1,53 @@
+"""Emissions tracking functionality for monitoring energy consumption and carbon emissions."""
+
+import json
+import logging
+
+from codecarbon import EmissionsTracker as CodeCarbonTracker  # type: ignore[import-untyped]
+from codecarbon.output import EmissionsData  # type: ignore[import-untyped]
+
+logger = logging.getLogger(__name__)
+
+
+class EmissionsTracker:
+    """Class for tracking energy consumption and carbon emissions."""
+
+    def __init__(self, project_name: str, measure_power_secs: int = 1) -> None:
+        """Initialize the emissions tracker.
+
+        Args:
+            project_name: Name of the project to track emissions for.
+            measure_power_secs: How often to measure power consumption in seconds.
+        """
+        self._logger = logger
+        self.tracker = CodeCarbonTracker(project_name=project_name, measure_power_secs=measure_power_secs)
+
+    def start_task(self, task_name: str) -> None:
+        """Start tracking emissions for a specific task.
+
+        Args:
+            task_name: Name of the task to track emissions for.
+        """
+        self.tracker.start_task(task_name)
+
+    def stop_task(self) -> dict[str, float]:
+        """Stop tracking emissions and return the emissions data.
+
+        Returns:
+            Dictionary containing emissions metrics.
+        """
+        emissions_data = self.tracker.stop_task()
+        _ = self.tracker.stop()
+        return self._process_metrics(emissions_data)
+
+    def _process_metrics(self, emissions_data: EmissionsData) -> dict[str, float]:
+        """Process emissions data into metrics with the 'emissions/' prefix.
+
+        Args:
+            emissions_data: Raw emissions data from the tracker.
+
+        Returns:
+            Dictionary of processed emissions metrics with the 'emissions/' prefix.
+        """
+        emissions_data_dict = json.loads(emissions_data.toJSON())
+        return {f"emissions/{k}": v for k, v in emissions_data_dict.items() if isinstance(v, int | float)}
diff --git a/docs/optimizer_config.schema.json b/docs/optimizer_config.schema.json
@@ -66,16 +66,16 @@
                 "validation_size": {
                     "default": 0.2,
                     "description": "Fraction of train samples to allocate for validation (if input dataset doesn't contain validation split).",
-                    "maximum": 1.0,
-                    "minimum": 0.0,
+                    "maximum": 1,
+                    "minimum": 0,
                     "title": "Validation Size",
                     "type": "number"
                 },
                 "separation_ratio": {
                     "anyOf": [
                         {
-                            "maximum": 1.0,
-                            "minimum": 0.0,
+                            "maximum": 1,
+                            "minimum": 0,
                             "type": "number"
                         },
                         {
@@ -342,6 +342,7 @@
         },
         "search_space": {
             "items": {
+                "additionalProperties": true,
                 "type": "object"
             },
             "title": "Search Space",
diff --git a/pyproject.toml b/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]
diff --git a/tests/callback/test_callback.py b/tests/callback/test_callback.py
@@ -26,6 +26,7 @@ def log_value(self, **kwargs: dict[str, Any]) -> None:
 
     def log_metrics(self, **kwargs: dict[str, Any]) -> None:
         metrics = kwargs["metrics"]
+        metrics = {k: v for k, v in metrics.items() if not k.startswith("emissions/")}
         for metric_name, metric_value in metrics.items():
             if not isinstance(metric_value, str) and np.isnan(metric_value):
                 metrics[metric_name] = None
@@ -103,7 +104,14 @@ def test_pipeline_callbacks(dataset):
                 "num": 0,
             },
         ),
-        ("log_metric", {"metrics": {"retrieval_hit_rate": 1.0}}),
+        (
+            "log_metric",
+            {
+                "metrics": {
+                    "retrieval_hit_rate": 1.0,
+                }
+            },
+        ),
         ("end_module", {}),
         (
             "start_module",
@@ -113,7 +121,14 @@ def test_pipeline_callbacks(dataset):
                 "num": 1,
             },
         ),
-        ("log_metric", {"metrics": {"retrieval_hit_rate": 1.0}}),
+        (
+            "log_metric",
+            {
+                "metrics": {
+                    "retrieval_hit_rate": 1.0,
+                }
+            },
+        ),
         ("end_module", {}),
         (
             "start_module",
@@ -139,7 +154,15 @@ def test_pipeline_callbacks(dataset):
                 "num": 0,
             },
         ),
-        ("log_metric", {"metrics": {"scoring_accuracy": 1.0, "scoring_roc_auc": 1.0}}),
+        (
+            "log_metric",
+            {
+                "metrics": {
+                    "scoring_accuracy": 1.0,
+                    "scoring_roc_auc": 1.0,
+                }
+            },
+        ),
         ("end_module", {}),
         (
             "start_module",
@@ -165,7 +188,15 @@ def test_pipeline_callbacks(dataset):
                 "num": 1,
             },
         ),
-        ("log_metric", {"metrics": {"scoring_accuracy": 1.0, "scoring_roc_auc": 1.0}}),
+        (
+            "log_metric",
+            {
+                "metrics": {
+                    "scoring_accuracy": 1.0,
+                    "scoring_roc_auc": 1.0,
+                }
+            },
+        ),
         ("end_module", {}),
         (
             "start_module",
@@ -189,7 +220,15 @@ def test_pipeline_callbacks(dataset):
                 "num": 0,
             },
         ),
-        ("log_metric", {"metrics": {"scoring_accuracy": 0.75, "scoring_roc_auc": 1.0}}),
+        (
+            "log_metric",
+            {
+                "metrics": {
+                    "scoring_accuracy": 0.75,
+                    "scoring_roc_auc": 1.0,
+                }
+            },
+        ),
         ("end_module", {}),
         ("start_module", {"module_kwargs": {"thresh": 0.5}, "module_name": "threshold", "num": 0}),
         (

Original file line number	Diff line number	Diff line change
`@@ -45,6 +45,7 @@ dependencies = [`
`45`	`45`	`"xxhash (>=3.5.0,<4.0.0)",`
`46`	`46`	`"python-dotenv (>=1.0.1,<2.0.0)",`
`47`	`47`	`"transformers[torch] (>=4.49.0,<5.0.0)",`
	`48`	`+ "codecarbon (==2.6)",`
`48`	`49`	`]`
`49`	`50`
`50`	`51`	`[project.urls]`