Require Assistance with Custom Plugin Implementation

[sequentialrant]

I recently started using Avalanche Continual Learning library to explore online continual learning techniques. My research involves experimenting with various selection strategies, and I aim to incorporate a custom metric that measures power consumption of CPU and GPU. To achieve this, I developed a plugin. However, when I run the experiments, the plugin's functions meant to calculate power consumption aren't being invoked. Curiously, when I test these functions separately, they work, and the metrics are computed. Additionally, the output is limited to a single value, which could be either CPU power, GPU power, or total power, instead of all three. I've included a screenshot for better understanding.

Here is the code I'm referencing.

class TegrastatsMetric(Metric[float]):
     def __init__(self, output_file, interval, error_file):
        super().__init__()
        self.output_file = output_file
        self.interval = interval
        self.error_file = error_file
        # self.tegrastats_pid = None
        self.gpu_power = 0.0
        self.cpu_power = 0.0
        self.total_power = 0.0

    def start_tegrastats(self):
        print(f"Starting tegrastats with interval {self.interval} and output file {self.output_file}")
        with open(self.error_file, 'w') as err_file:
            p =  subprocess.Popen(['sudo', 'tegrastats', '--interval', str(self.interval)], stdout=open(self.output_file, 'w'), stderr=err_file)
        
        print(f"tegrastats process started with PID: {p.pid}")
        return p.pid

    def stop_tegrastats(self):
        try:
            os.system("sudo pkill tegrastats")
            print("Subprocess terminated successfully.")
        except (psutil.NoSuchProcess, PermissionError) as e:
            print(f"Termination failed: {e}")


    def parse_tegrastats_output(self):
        total_gpu_soc_power = 0
        total_cpu_cv_power = 0

        with open(self.output_file, 'r') as file:
            for line in file:
                if 'VDD_GPU_SOC' in line:
                    try:
                        gpu_soc_power = line.split('VDD_GPU_SOC ')[1].split('mW')[0]
                        total_gpu_soc_power += float(gpu_soc_power) * (self.interval / 1000.0)
                    except (IndexError, ValueError):
                        print("Error parsing GPU SOC power value.")
                if 'VDD_CPU_CV' in line:
                    try:
                        cpu_cv_power = line.split('VDD_CPU_CV ')[1].split('mW')[0]
                        total_cpu_cv_power += float(cpu_cv_power) * (self.interval / 1000.0)
                        #total_cpu_cv_power += int(cpu_cv_power)
                    except (IndexError, ValueError):
                        print("Error parsing CPU CV power value.")
                        
        total_power = total_gpu_soc_power + total_cpu_cv_power

        return total_gpu_soc_power, total_cpu_cv_power, total_power

    def before_training(self, strategy: 'BaseStrategy') -> None:
        self.tegrastats_pid = self.start_tegrastats() # self.interval, self.output_file)

    def after_training(self, strategy: 'BaseStrategy') -> None:
        self.stop_tegrastats(self.tegrastats_pid)

    def after_training_iteration(self, strategy: 'BaseStrategy') -> None:
        gpu_power, cpu_power, total_power = self.parse_tegrastats_output() # self.output_file, self.interval)
        self.update(gpu_power, cpu_power, total_power)

    def update(self, gpu_power, cpu_power, total_power):
        self.gpu_power = gpu_power
        self.cpu_power = cpu_power
        self.total_power = total_power

    def reset(self) -> None:
        self.gpu_power = 0.0
        self.cpu_power = 0.0
        self.total_power = 0.0
        # if os.path.exists(self.output_file):
        #     os.remove(self.output_file)

    def __str__(self):
        return "Tegrastats_Power_Consumption"

class TegrastatsPluginMetric(PluginMetric[float]):
    def __init__(self, tegrastats_metric: TegrastatsMetric):
        super().__init__()
        self.tegrastats_metric = tegrastats_metric

    def after_training_iteration(self, strategy: 'BaseStrategy') -> MetricResult:
        return self._package_result(strategy)

    def reset(self) -> None:
        self.tegrastats_metric.reset()

    def result(self) -> dict:
        return {
            "cpu_power": self.tegrastats_metric.cpu_power,
            "gpu_power": self.tegrastats_metric.gpu_power,
            "total_power": self.tegrastats_metric.total_power
        }

    def _package_result(self, strategy: 'BaseStrategy') -> MetricResult:
        cpu_power = self.tegrastats_metric.cpu_power
        gpu_power = self.tegrastats_metric.gpu_power
        total_power = self.tegrastats_metric.total_power

        print(f"GPU Power: {gpu_power}")
        print(f"CPU Power: {cpu_power}")
        print(f"Total Power: {total_power}")

        cpu_metric_name = get_metric_name(self, strategy, add_experience=True)
        gpu_metric_name = get_metric_name(self, strategy, add_experience=True)
        total_metric_name = get_metric_name(self, strategy, add_experience=True)

        return [
            MetricValue(self, gpu_metric_name, gpu_power, strategy.clock.train_iterations),
            MetricValue(self, cpu_metric_name, cpu_power, strategy.clock.train_iterations),
            MetricValue(self, total_metric_name, total_power, strategy.clock.train_iterations)
        ]

    def __str__(self):
        return "Tegrastats_Power_Consumption_Plugin"

Using the plugin as below:

interval = 500
output_file = "tegrastats_output.txt
error_file = "tegrastats_error.txt
tegrastats_metric = TegrastatsMetric(output_file, interval, error_file)

eval_plugin = EvaluationPlugin(
accuracy_metrics(experience=True, stream=True, trained_experience=True),
forgetting_metrics(experience=True, stream=True),
TegrastatsPluginMetric(tegrastats_metric),
timing_metrics(experience=True, stream=True),
loggers=[interactive_logger, text_logger]
)

Could someone advise me on what might be going wrong? Am I implementing the plugin correctly?

Kind regards,
Srinidhi