python-api/simvue/eco.py at ee9773f0efdbe2b0834faf1dd5eb76d8c059d1ec · simvue-io/python-api · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
import typing
import logging
import datetime

from codecarbon import EmissionsTracker, OfflineEmissionsTracker
from codecarbon.output import BaseOutput as cc_BaseOutput
from simvue.utilities import simvue_timestamp

if typing.TYPE_CHECKING:
    from simvue import Run
    from codecarbon.output_methods.emissions_data import EmissionsData


logger = logging.getLogger(__file__)


class CodeCarbonOutput(cc_BaseOutput):
    def __init__(self, run: "Run") -> None:
        self._simvue_run = run
        self._metrics_step: int = 0

    def out(
        self, total: "EmissionsData", delta: "EmissionsData", meta_update: bool = True
    ) -> None:
        # Check if the run has been shutdown, if so do nothing
        if (
            self._simvue_run._shutdown_event
            and self._simvue_run._shutdown_event.is_set()
        ):
            logger.debug("Terminating CodeCarbon tracker")
            return

        if meta_update:
            logger.debug("Logging CodeCarbon metadata")
            try:
                self._simvue_run.update_metadata(
                    {
                        "codecarbon": {
                            "country": total.country_name,
                            "country_iso_code": total.country_iso_code,
                            "region": total.region,
                            "version": total.codecarbon_version,
                        }
                    }
                )
            except AttributeError as e:
                logger.error(f"Failed to update metadata: {e}")
        try:
            _cc_timestamp = datetime.datetime.strptime(
                total.timestamp, "%Y-%m-%dT%H:%M:%S"
            )
        except ValueError as e:
            logger.error(f"Error parsing timestamp: {e}")
            return

        logger.debug("Logging CodeCarbon metrics")
        try:
            self._simvue_run.log_metrics(
                metrics={
                    "codecarbon.total.emissions": total.emissions,
                    "codecarbon.total.energy_consumed": total.energy_consumed,
                    "codecarbon.delta.emissions": delta.emissions,
                    "codecarbon.delta.energy_consumed": delta.energy_consumed,
                },
                step=self._metrics_step,
                timestamp=simvue_timestamp(_cc_timestamp),
            )
        except ArithmeticError as e:
            logger.error(f"Failed to log metrics: {e}")
            return

        self._metrics_step += 1

    def live_out(self, total: "EmissionsData", delta: "EmissionsData") -> None:
        self.out(total, delta, meta_update=False)


class SimvueEmissionsTracker(EmissionsTracker):
    def __init__(
        self, project_name: str, simvue_run: "Run", metrics_interval: int
    ) -> None:
        self._simvue_run = simvue_run
        logger.setLevel(logging.ERROR)
        super().__init__(
            project_name=project_name,
            measure_power_secs=metrics_interval,
            api_call_interval=1,
            experiment_id=None,
            experiment_name=None,
            logging_logger=CodeCarbonOutput(simvue_run),
            save_to_logger=True,
            allow_multiple_runs=True,
            log_level="error",
        )

    def set_measure_interval(self, interval: int) -> None:
        """Set the measure interval"""
        self._set_from_conf(interval, "measure_power_secs")

    def post_init(self) -> None:
        self._set_from_conf(self._simvue_run._id, "experiment_id")
        self._set_from_conf(self._simvue_run._name, "experiment_name")
        self.start()


class OfflineSimvueEmissionsTracker(OfflineEmissionsTracker):
    def __init__(
        self, project_name: str, simvue_run: "Run", metrics_interval: int
    ) -> None:
        self._simvue_run = simvue_run
        logger.setLevel(logging.ERROR)
        super().__init__(
            country_iso_code=simvue_run._user_config.offline.country_iso_code,
            project_name=project_name,
            measure_power_secs=metrics_interval,
            api_call_interval=1,
            experiment_id=None,
            experiment_name=None,
            logging_logger=CodeCarbonOutput(simvue_run),
            save_to_logger=True,
            allow_multiple_runs=True,
            log_level="error",
        )

    def set_measure_interval(self, interval: int) -> None:
        """Set the measure interval"""
        self._set_from_conf(interval, "measure_power_secs")

    def post_init(self) -> None:
        self._set_from_conf(self._simvue_run._id, "experiment_id")
        self._set_from_conf(self._simvue_run._name, "experiment_name")
        self.start()