[GR-64688] Always apply command mapper hooks (trackers) for the image stages

PullRequest: mx/1931

Author: ayman-moustadrif

src/mx/_impl/mx.py

@@ -13906,7 +13906,7 @@ def apply_command_mapper_hooks(command, hooks):
             for hook in reversed(hooks):
                 hook_name, hook_func, suite = hook[:3]
                 logv(f"Applying command mapper hook '{hook_name}'")
-                new_cmd = hook_func(new_cmd, suite)
+                new_cmd = hook_func.hook(new_cmd, suite)
                 logv(f"New command: {new_cmd}")
     else:
         log("Skipping command mapper hooks as they were disabled explicitly.")

src/mx/_impl/mx_benchmark.py

@@ -132,6 +132,7 @@
 from typing import Callable, Sequence, Iterable, Optional, Dict, Any, List, Collection
 from dataclasses import dataclass
 
+from .mx_util import Stage, MapperHook, FunctionHookAdapter
 from .support.logging import log_deprecation
 
 from . import mx
@@ -513,7 +514,7 @@ def __init__(self, *args, **kwargs):
         super(BenchmarkSuite, self).__init__(*args, **kwargs)
         self._desired_version = None
         self._suite_dimensions: DataPoint = {}
-        self._command_mapper_hooks = {}
+        self._command_mapper_hooks: Dict[str, MapperHook] = {}
         self._tracker = None
         self._currently_running_benchmark = None
         self._execution_context: List[BenchmarkExecutionContext] = []
@@ -586,23 +587,27 @@ def currently_running_benchmark(self):
         """
         return self._currently_running_benchmark
 
-    def register_command_mapper_hook(self, name, func):
-        """Registers a function that takes as input the benchmark suite object and the command to execute and returns
-        a modified command line.
+    def register_command_mapper_hook(self, name: str, hook: Callable | MapperHook):
+        """Registers a command modification hook with stage awareness.
 
         :param string name: Unique name of the hook.
-        :param function func:
+        :param hook: MapperHook instance or callable
         :return: None
         """
-        self._command_mapper_hooks[name] = func
+        if isinstance(hook, MapperHook):
+            self._command_mapper_hooks[name] = hook
+        elif callable(hook):
+            mx.warn(f"Registering a Callable as a command mapper hook '{name}' for the benchmark suite. Please use the MapperHook interface instead of Callable as it will be deprecated in the future!")
+            self._command_mapper_hooks[name] = FunctionHookAdapter(hook)
+        else:
+            raise ValueError(f"Hook must be MapperHook or callable, got {type(hook)}")
+
 
     def register_tracker(self, name, tracker_type):
         tracker = tracker_type(self)
         self._tracker = tracker
+        hook = tracker.get_hook()
 
-        def hook(cmd, suite):
-            assert suite == self
-            return tracker.map_command(cmd)
         self.register_command_mapper_hook(name, hook)
 
     def version(self):
@@ -1764,9 +1769,10 @@ def vmArgs(self, bmSuiteArgs):
 
                     def func(cmd, bmSuite, prefix_command=prefix_command):
                         return prefix_command + cmd
+
                     if self._command_mapper_hooks and any(hook_name != profiler for hook_name in self._command_mapper_hooks):
                         mx.abort(f"Profiler '{profiler}' conflicts with trackers '{', '.join([hook_name for hook_name in self._command_mapper_hooks if hook_name != profiler])}'\nUse --tracker none to disable all trackers")
-                    self._command_mapper_hooks = {profiler: func}
+                    self.register_command_mapper_hook(profiler, FunctionHookAdapter(func))
         return args
 
     def parserNames(self):
@@ -2010,7 +2016,7 @@ def command_mapper_hooks(self, hooks):
         """
         Registers a list of `hooks` (given as a tuple 'name', 'func', 'suite') to manipulate the command line before its
         execution.
-        :param list[tuple] hooks: the list of hooks given as tuples of names and functions
+        :param list[tuple] hooks: the list of hooks given as tuples of names and MapperHook instances
         """
         self._command_mapper_hooks = hooks
 
@@ -3078,6 +3084,14 @@ def disable_tracker():
     global _use_tracker
     _use_tracker = False
 
+class DefaultTrackerHook(MapperHook):
+    def __init__(self, tracker: Tracker):
+        self.tracker = tracker
+
+    def hook(self, cmd, suite=None):
+        assert suite == self.tracker.bmSuite
+        return self.tracker.map_command(cmd)
+
 class Tracker(object):
     def __init__(self, bmSuite):
         self.bmSuite = bmSuite
@@ -3088,6 +3102,9 @@ def map_command(self, cmd):
     def get_rules(self, bmSuiteArgs):
         raise NotImplementedError()
 
+    def get_hook(self) -> MapperHook:
+        return DefaultTrackerHook(self)
+
 class RssTracker(Tracker):
     def __init__(self, bmSuite):
         super().__init__(bmSuite)
@@ -3461,6 +3478,10 @@ def __init__(self, bmSuite):
         # GR-65536
         atexit.register(self.cleanup)
 
+    def get_hook(self) -> MapperHook:
+        """Returns an energy tracking hook"""
+        return EnergyTrackerHook(self)
+
     @property
     def baseline_power(self):
         """Caches the average baseline power value"""
@@ -3559,6 +3580,15 @@ def get_rules(self, bmSuiteArgs):
 
         return rules
 
+class EnergyTrackerHook(DefaultTrackerHook):
+    """Hook for energy consumption tracking."""
+    def __init__(self, tracker: EnergyConsumptionTracker):
+        assert isinstance(tracker, EnergyConsumptionTracker)
+        super().__init__(tracker)
+
+    def should_apply(self, stage: Optional[Stage]) -> bool:
+        return stage.is_final() if stage else False
+
 _available_trackers = {
     "rss": RssTracker,
     "psrecord": PsrecordTracker,

src/mx/_impl/mx_util.py

@@ -24,6 +24,7 @@
 #
 # ----------------------------------------------------------------------------------------------------
 #
+from __future__ import annotations
 
 #
 # Utility functions for use by mx and mx extensions.
@@ -37,14 +38,21 @@
     "get_file_extension",
     "ensure_dirname_exists",
     "ensure_dir_exists",
-    "SafeFileCreation"
+    "SafeFileCreation",
+    "Stage",
+    "StageName",
+    "Layer",
+    "MapperHook",
+    "FunctionHookAdapter"
 ]
 
 import os.path
 import errno
 import sys
 import tempfile
-from typing import Optional
+from dataclasses import dataclass
+from enum import Enum
+from typing import Optional, List, Callable
 from os.path import dirname, exists, join, isdir, basename
 
 min_required_python_version = (3, 8)
@@ -215,3 +223,93 @@ def _create_tmp_files(tmp_dir, num):
             pid = os.getpid()
             with open(sfc.tmpPath, 'w') as out:
                 print(f"file {i} created by process {pid}", file=out)
+
+
+@dataclass(frozen = True)
+class Stage:
+    stage_name: StageName
+    layer_info: Layer = None
+
+    @staticmethod
+    def from_string(s: str) -> Stage:
+        return Stage(StageName(s))
+
+    def is_image(self) -> bool:
+        return self.stage_name.is_image()
+
+    def is_instrument(self) -> bool:
+        return self.stage_name.is_instrument()
+
+    def is_agent(self) -> bool:
+        return self.stage_name.is_agent()
+
+    def is_final(self) -> bool:
+        return self.stage_name.is_final()
+
+    def is_layered(self) -> bool:
+        return self.layer_info is not None
+
+    def is_requested(self, request: str):
+        """Whether the 'request' is equal to either the full name of the stage or it's name without layer info."""
+        return str(self) == request or str(self.stage_name) == request
+
+    def __str__(self):
+        if not self.is_layered():
+            return str(self.stage_name)
+        return f"{self.stage_name}-{self.layer_info}"
+
+class StageName(Enum):
+    AGENT = "agent"
+    INSTRUMENT_IMAGE = "instrument-image"
+    INSTRUMENT_RUN = "instrument-run"
+    IMAGE = "image"
+    RUN = "run"
+
+    def __str__(self):
+        return self.value
+
+    def is_image(self) -> bool:
+        """Whether this is an image stage (a stage that performs an image build)"""
+        return self in [StageName.INSTRUMENT_IMAGE, StageName.IMAGE]
+
+    def is_instrument(self) -> bool:
+        """Whether this is an image stage (a stage that performs an image build)"""
+        return self in [StageName.INSTRUMENT_IMAGE, StageName.INSTRUMENT_RUN]
+
+    def is_agent(self) -> bool:
+        return self == StageName.AGENT
+
+    def is_final(self) -> bool:
+        return self in [StageName.IMAGE, StageName.RUN]
+
+@dataclass(frozen = True)
+class Layer:
+    index: int
+    is_shared_library: bool
+
+    def __str__(self):
+        return f"layer{self.index}"
+
+
+class MapperHook:
+    """Base class for all command mapper hooks."""
+    def hook(self, cmd: List[str], suite=None) -> List[str]:
+        raise NotImplementedError()
+
+    def should_apply(self, stage: Optional[Stage]) -> bool:
+        """Determines if this hook should be applied for the given stage. By default, hooks shouldn't apply to image stages.
+        Args:
+            stage: For NativeImage suites, the current stage
+                   For JIT suites, None
+        Returns:
+            bool: True if the hook should be applied
+        """
+        return not stage.is_image() if stage else True
+
+class FunctionHookAdapter(MapperHook):
+    """Adapter for function hooks."""
+    def __init__(self, func: Callable[[List[str]], List[str]]):
+        self.func = func
+
+    def hook(self, cmd: List[str], suite=None) -> List[str]:
+        return self.func(cmd, suite)

src/mx/mx_version.py

@@ -1,2 +1,2 @@
 # The version must be updated for every PR (checked in CI) and the comment should reflect the PR's issue
-version = "7.58.5"  # GR-66866: Suppress compile warnings in shaded jar projects in Truffle for Intellij.
+version = "7.58.6"  # GR-64688 Always apply command mapper hooks (trackers) for the image stages