more detail energy breakdown

carbon/__init__.py

@@ -6,6 +6,7 @@
 
 from carbon.intensity import CarbonIntensity
 from carbon.job import Job, JobFactory
+from carbon.job.job import EnergyBreakdown
 from carbon.node import Node, NodeFactory
 
 with suppress(PackageNotFoundError):
@@ -18,7 +19,7 @@ class RunResult:
 
     node: Node
     emissions: float
-    energy_consumed: float
+    energy_breakdown: EnergyBreakdown
     job: Job
     carbon_intensity: float
 
@@ -78,13 +79,13 @@ def run(
             intensity = carbon_intensity.fetch()
 
         # Calculate emissions
-        emissions = intensity * energy_consumed
+        emissions = intensity * energy_consumed.total
 
         result_list.append(
             RunResult(
                 node=node,
                 emissions=emissions,
-                energy_consumed=energy_consumed,
+                energy_breakdown=energy_consumed,
                 job=job,
                 carbon_intensity=intensity,
             )

carbon/__main__.py

@@ -18,6 +18,7 @@
     JobFactory,
     PBSJobFactory,
 )
+from carbon.job.job import EnergyBreakdown
 from carbon.node.factories import (
     DummyNodeFactory,
     FileNodeFactory,
@@ -223,7 +224,7 @@ def main(
             total_gputime += job.gputime
             total_memtime += job.memtime
 
-            total_energy_consumed += result.energy_consumed
+            total_energy_consumed += result.energy_breakdown.total
             total_emissions += result.emissions
             intensity_list.append(result.carbon_intensity)
 
@@ -244,7 +245,9 @@ def main(
         agg_result = RunResult(
             node=results[0].node,  # Just use first node for now
             emissions=total_emissions,
-            energy_consumed=total_energy_consumed,
+            energy_breakdown=EnergyBreakdown(
+                cpu=0, gpu=0, memory=0, total=total_energy_consumed
+            ),
             job=agg_job,
             carbon_intensity=sum(intensity_list) / len(intensity_list),
         )
@@ -275,7 +278,7 @@ def output_result(
     """
     node = result.node
     emissions = result.emissions
-    energy_consumed = result.energy_consumed
+    total_energy_consumed = result.energy_breakdown.total
     job = result.job
     intensity = result.carbon_intensity
 
@@ -312,7 +315,12 @@ def output_result(
         f"Estimated energy consumed from {job.cputime:.2f} CPU-hours "
         f"and {job.gputime:.2f} GPU-hours "
         f"and {job.memtime:.2f} GB-hours "
-        f"is {energy_consumed:.2f} kWh"
+        f"is {total_energy_consumed:.2f} kWh"
+    )
+    print(
+        f"  CPU: {result.energy_breakdown.cpu:.2f} kWh, "
+        f"GPU: {result.energy_breakdown.gpu:.2f} kWh, "
+        f"Memory: {result.energy_breakdown.memory:.2f} kWh"
     )
     if average_intensity:
         print(f"Using average carbon intensity of {config.average_intensity} gCO2/kWh")

carbon/job/job.py

@@ -15,6 +15,16 @@ class JobState(Enum):
     EXPIRED = "X"
 
 
+@dataclass
+class EnergyBreakdown:
+    """Breakdown of energy consumption by component."""
+
+    cpu: float
+    gpu: float
+    memory: float
+    total: float
+
+
 @dataclass
 class Job:
     """Represents a compute job, including its resource usage and timing information."""
@@ -43,22 +53,19 @@ class Job:
     state: JobState = JobState.FINISHED
     """The state of the job."""
 
-    def calculate_energy(self, node: Node, pue: float) -> float:
+    def calculate_energy(self, node: Node, pue: float) -> EnergyBreakdown:
         """Calculate energy consumption in kilowatt-hours for a compute job.
 
         Args:
             node (Node): The compute node the job was executed on.
             pue (float): Power Usage Effectiveness of the data center.
 
         Returns:
-            float: The energy consumed in kilowatt-hours.
+            EnergyBreakdown: The energy consumed broken down by component.
         """
-        return (
-            (
-                node.per_core_power_watts * self.cputime
-                + node.per_gpu_power_watts * self.gputime
-                + node.per_gb_power_watts * self.memtime
-            )
-            * pue
-            / 1000.0
+        cpu = node.per_core_power_watts * self.cputime * pue / 1000.0
+        gpu = node.per_gpu_power_watts * self.gputime * pue / 1000.0
+        memory = node.per_gb_power_watts * self.memtime * pue / 1000.0
+        return EnergyBreakdown(
+            cpu=cpu, gpu=gpu, memory=memory, total=cpu + gpu + memory
         )

tests/test_cli.py

@@ -8,6 +8,7 @@
 from carbon import RunResult
 from carbon.__main__ import main
 from carbon.job import Job, JobState
+from carbon.job.job import EnergyBreakdown
 from carbon.node import Node
 
 
@@ -37,7 +38,11 @@ def make_result(job_id: str) -> RunResult:
         per_gb_power_watts=0.5,
     )
     return RunResult(
-        node=node, emissions=1.0, energy_consumed=1.0, job=job, carbon_intensity=137.0
+        node=node,
+        emissions=1.0,
+        job=job,
+        energy_breakdown=EnergyBreakdown(cpu=0, gpu=0, memory=0, total=1.0),
+        carbon_intensity=137.0,
     )
 
 

tests/test_job.py

@@ -147,10 +147,16 @@ def test_energy_calculate() -> None:
         per_gb_power_watts=2.0,
     )
 
-    expected = ((10.0 * 2.0) + (200.0 * 2.0) + (32.0 * 2.0)) * 1.5 / 1000.0
+    expected_cpu = (10.0 * 2.0) * 1.5 / 1000.0
+    expected_gpu = (200.0 * 2.0) * 1.5 / 1000.0
+    expected_memory = (32.0 * 2.0) * 1.5 / 1000.0
+    expected_total = expected_cpu + expected_gpu + expected_memory
     result = job.calculate_energy(node, 1.5)
 
-    assert np.isclose(result, expected, atol=1e-9)
+    assert np.isclose(result.cpu, expected_cpu, atol=1e-9)
+    assert np.isclose(result.gpu, expected_gpu, atol=1e-9)
+    assert np.isclose(result.memory, expected_memory, atol=1e-9)
+    assert np.isclose(result.total, expected_total, atol=1e-9)
 
 
 def test_energy_calculate_no_gpu() -> None:
@@ -174,10 +180,16 @@ def test_energy_calculate_no_gpu() -> None:
         per_gb_power_watts=2.0,
     )
 
-    expected = ((10.0 * 2.0) + (32.0 * 2.0)) * 1.5 / 1000.0
+    expected_cpu = (10.0 * 2.0) * 1.5 / 1000.0
+    expected_gpu = 0.0
+    expected_memory = (32.0 * 2.0) * 1.5 / 1000.0
+    expected_total = expected_cpu + expected_gpu + expected_memory
     result = job.calculate_energy(node, 1.5)
 
-    assert np.isclose(result, expected, atol=1e-9)
+    assert np.isclose(result.cpu, expected_cpu, atol=1e-9)
+    assert np.isclose(result.gpu, expected_gpu, atol=1e-9)
+    assert np.isclose(result.memory, expected_memory, atol=1e-9)
+    assert np.isclose(result.total, expected_total, atol=1e-9)
 
 
 def test_file_job_factory_create(tmp_path: Path) -> None: