from __future__ import annotations
from typing import Literal
from analyzer.core.param_specs import (
ModuleParameterSpec,
ModuleParameterValues,
ParameterSpec, # noqa
)
from analyzer.utils.structure_tools import freeze
import functools as ft
from cattrs.strategies import include_subclasses, configure_tagged_union
from analyzer.core.run_builders import RunBuilder, DEFAULT_RUN_BUILDER
from analyzer.core.datasets import SampleType
from attrs import define, field, make_class, asdict
from analyzer.core.results import ResultBase
from analyzer.utils.structure_tools import SimpleCache
from analyzer.core.columns import TrackedColumns, Column, ColumnCollection
import contextlib
import abc
from typing import Any, TYPE_CHECKING
import logging
from analyzer.core.adl import ADLBlock
from analyzer.core.linting import LintMessage
[docs]
logger = logging.getLogger("analyzer.core")
# This system could probably be unified with the pattern matching system
@define
@define
[docs]
class IsYear(MetadataExpr):
[docs]
def evaluate(self, metadata):
return metadata["era"]["name"] == self.year
@define
[docs]
class IsSampleType(MetadataExpr):
[docs]
sample_type: SampleType = field(converter=SampleType)
[docs]
def evaluate(self, metadata):
return metadata["sample_type"] == self.sample_type
@define
[docs]
class IsRun(MetadataExpr):
[docs]
def evaluate(self, metadata):
is_run_2 = any(x in metadata["era"]["name"] for x in ("2016", "2017", "2018"))
return (self.run == 2) == is_run_2
@define
@define
@define
[docs]
MultiColumns = list[tuple[Any, TrackedColumns]]
[docs]
def moduleExcludeFilter(attribute, value):
return attribute.name not in ("should_run",) and not attribute.name.startswith("_")
@define
[docs]
class BaseAnalyzerModule(abc.ABC):
"""
Base class for a unit of analysis.
Provides a cache to avoid duplicate execution when running systematics.
"""
_cache: SimpleCache = field(
factory=lambda: SimpleCache(max_size=AnalyzerModule.MAX_CACHE_SIZE),
init=False,
repr=False,
)
[docs]
should_run: MetadataExpr | None = field(default=None, kw_only=True)
@abc.abstractmethod
[docs]
def getParameterSpec(self, metadata: dict) -> ModuleParameterSpec:
return {}
[docs]
def filterParams(self, metadata, params):
spec = self.getParameterSpec(metadata)
return {x: y for x, y in params.items() if x in spec}
[docs]
def getFromCache(self, key):
return self._cache[key]
[docs]
def neededResources(self, metadata) -> list[str]:
return []
[docs]
def adlExport(self, metadata) -> list["ADLBlock"] | None:
return None
# Why is this not a property?
@classmethod
[docs]
def name(cls):
return cls.__name__
[docs]
def clearCache(self):
self._cache.clear()
[docs]
def lint(self) -> list[LintMessage]:
return []
# Technically this is dangerous, as the class is not frozen, but the user should never modify the properties of a module once it is created
@ft.cached_property
[docs]
def selfkey(self):
return hash(freeze(asdict(self, filter=moduleExcludeFilter)))
@define
[docs]
class AnalyzerModule(BaseAnalyzerModule):
"""Abstract base class for all analyzer modules.
Subclasses must implement the inputs and run methods.
"""
@abc.abstractmethod
[docs]
def outputs(
self, metadata
) -> ColumnCollection | list[Column] | Literal["EVENTS"]: ...
@abc.abstractmethod
[docs]
def run(
self, columns, params: ModuleParameterValues
) -> tuple[TrackedColumns, list[ResultBase | ModuleAddition]]:
pass
[docs]
def getKey(self, columns, params):
"""
Get a unique key associated with a given run.
Acts a proxy for caching the output of the run method.
Computed by hashing the self key, the parameters, and the key associated with the columns.
"""
logger.debug(f"Params are {params}")
inp = self.inputs(columns.metadata)
if inp == "EVENTS":
k = columns.getKeyForAll()
else:
k = columns.getKeyForColumns(inp)
# The first two should not be necessary but got worried.
ret = hash((self.selfkey, self.name(), freeze(params), k))
# logger.info(f"Module: {self}. Input columns are: {inp}.\nColumn key is {k}\nFinal key is {ret}")
return ret
[docs]
def getKeyNoParams(self, columns):
"""
Why is ths even a different method.
"""
inp = self.inputs(columns.metadata)
if inp == "EVENTS":
k = columns.getKeyForAll()
else:
k = columns.getKeyForColumns(inp)
ret = hash((self.selfkey, self.name(), k))
return ret
def __run(self, columns, params):
orig_columns = columns
columns = columns.copy()
input_key = self.getKey(columns, params)
outputs = self.outputs(columns.metadata)
if outputs == "EVENTS":
cache_key = hash((input_key, columns.getKeyForAll()))
else:
cache_key = input_key
logger.debug(f"Input key is {input_key}, Cache key is {cache_key}")
logger.debug(f"Cached keys are {list(self._cache)}")
if cache_key in self._cache:
logger.debug("Found key, using cached result")
cached_cols, r, internal = self._cache[cache_key]
if outputs == "EVENTS":
return cached_cols, r
outputs += internal
logger.debug(f"Updating columns with cached output from columns: {outputs}")
columns.addColumnsFrom(cached_cols, outputs)
columns.pipeline_data = cached_cols.pipeline_data
return columns, r
logger.debug(f"Did not find cached result, running module {self.name()}")
outputs = self.outputs(columns.metadata)
inputs = self.inputs(columns.metadata)
if outputs == "EVENTS":
output_cx = contextlib.nullcontext()
else:
output_cx = columns.allowedOutputs(outputs)
if inputs == "EVENTS":
inputs_cx = contextlib.nullcontext()
else:
inputs_cx = columns.allowedInputs(inputs)
with (
columns.useKey(input_key),
inputs_cx,
output_cx,
):
_, res = self.run(columns, params)
internal = columns.updatedColumns(orig_columns, Column("INTERNAL_USE"))
self._cache[cache_key] = (columns, res, internal)
return columns, res
[docs]
def __call__(self, columns, params):
try:
logger.debug(f"Running analyzer module {self}")
if self.should_run is not None:
metadata = columns.metadata
should_run = self.should_run.evaluate(metadata)
if not should_run:
return columns, []
params = self.filterParams(columns.metadata, params)
return self.__run(columns, params)
except Exception as e:
logger.exception(f"An exception occurred while running module {self}. {e}")
raise e
@define
[docs]
class EventSourceModule(BaseAnalyzerModule):
"""Abstract base class for event source modules.
Subclasses must implement the outputs and run methods.
"""
@abc.abstractmethod
[docs]
def outputs(self, metadata) -> ColumnCollection | list[Column]: ...
@abc.abstractmethod
[docs]
def run(self, params: ModuleParameterValues) -> TrackedColumns:
pass
[docs]
def getKey(self, params):
ret = hash((self.name(), self.selfkey, freeze(params)))
return ret
[docs]
def __call__(self, params):
try:
spec = self.getParameterSpec(None)
params = {x: y for x, y in params.items() if x in spec}
logger.debug(f"Running analyzer module {self}")
key = self.getKey(params)
logger.debug(f"Execution key is {key}")
logger.debug(f"Cached keys are {list(self._cache)}")
if key in self._cache:
logger.debug("Found key, using cached result")
cached_cols = self._cache[key]
return cached_cols
logger.debug(f"Did not find cached result, running module {self}")
ret = self.run(params)
self._cache[key] = ret
return ret
except Exception as e:
logger.error(f"An exception occurred while running module {self}")
raise e
@define
[docs]
class PureResultModule(BaseAnalyzerModule):
"""Abstract base class for pure result modules.
Subclasses must implement the outputs and run methods.
"""
@abc.abstractmethod
[docs]
def outputs(self, metadata) -> ColumnCollection | list[Column]: ...
@abc.abstractmethod
[docs]
def run(
self, columns: MultiColumns, params: ModuleParameterValues
) -> list[ResultBase]:
pass
[docs]
def getKey(self, columns: MultiColumns, params: ModuleParameterValues):
ret = hash(
(
self.name(),
freeze(params),
self.selfkey,
frozenset(
(x, y.getKeyForColumns(self.inputs(y.metadata)))
for x, y in (columns or [])
),
)
)
return ret
def __run(self, columns: MultiColumns, params):
columns = [(x, y.copy()) for x, y in columns]
just_cols = [x[1] for x in columns]
key = self.getKey(columns, params)
logger.debug(f"Execution key is {key}")
logger.debug(f"Cached keys are {list(self._cache)}")
if key in self._cache:
logger.debug("Found key, using cached result")
ret = self._cache[key]
return ret
logger.debug(f"Did not find cached result, running module {self.name()}")
with contextlib.ExitStack() as stack:
for c in just_cols:
stack.enter_context(c.useKey(key))
stack.enter_context(c.allowedOutputs(self.outputs(c.metadata)))
stack.enter_context(c.allowedInputs(self.inputs(c.metadata)))
ret = self.run(columns, params)
self._cache[key] = ret
return ret
[docs]
def __call__(self, columns: MultiColumns, params):
try:
logger.debug(f"Running analyzer module {self}")
# if self.should_run is not None and columns is not None:
# metadata = columns.metadata
# should_run = self.should_run.evaluate(metadata)
# if not should_run:
# return columns, []
spec = self.getParameterSpec(None)
params = {x: y for x, y in params.items() if x in spec}
return self.__run(columns, params)
except Exception as e:
logger.error(f"An exception occurred while running module {self}")
raise e
[docs]
def defaultCols(columns):
def inner(self, metadata):
return [Column(x) for x in columns]
return inner
[docs]
def defaultParameterSpec(params):
def inner(self, metadata):
return ModuleParameterSpec(params)
return inner
[docs]
def register_module(input_columns, output_columns, configuration=None, params=None):
"""
Attempt to make decorator for constructing modules.
Has problems with pickle.
"""
raise NotImplementedError()
configuration = configuration or {}
params = params or {}
def wrapper(func):
getParameterSpec = defaultParameterSpec(params)
run = func
if callable(input_columns):
inputs = input_columns
else:
inputs = defaultCols(input_columns)
if callable(output_columns):
outputs = output_columns
else:
outputs = defaultCols(output_columns)
cls = make_class(
func.__name__,
configuration,
bases=(AnalyzerModule,),
class_body=dict(
getParameterSpec=getParameterSpec,
run=run,
inputs=inputs,
outputs=outputs,
),
)
return cls
return wrapper
@define
[docs]
class ModuleAddition:
[docs]
analyzer_module: AnalyzerModule | PureResultModule
[docs]
run_builder: RunBuilder | type | None = DEFAULT_RUN_BUILDER
[docs]
this_module_parameters: dict | None = None
# parameter_runs: list[PipelineParameterValues]