from __future__ import annotations
import copy
from attrs import define, field, asdict
from collections import deque
from analyzer.core.analysis_modules import (
AnalyzerModule,
ModuleAddition,
)
from analyzer.core.run_builders import DEFAULT_RUN_BUILDER, CompleteSysts, RunBuilder
from analyzer.core.results import (
ResultProvenance,
ResultGroup,
ResultBase,
)
from analyzer.core.param_specs import getWithValues
from collections import ChainMap
from analyzer.modules.common.load_columns import LoadColumns
import logging
from analyzer.core.adl import ADLEmitter, ADLBlock, ADLStatement
import re
from analyzer.utils.structure_tools import SimpleCache, freeze, flatten
[docs]
logger = logging.getLogger("analyzer.core")
[docs]
def getPipelineSpecs(pipeline, metadata):
ret = {}
for module in pipeline:
new_specs = module.getParameterSpec(metadata)
if set(ret) & set(new_specs):
raise RuntimeError("Duplicate module parameter names")
ret.update(new_specs)
return ret
@define
[docs]
class Analyzer:
[docs]
all_modules: list = field(factory=list)
[docs]
base_pipelines: dict[str, list[AnalyzerModule]] = field(factory=dict)
[docs]
default_run_builder: RunBuilder = field(factory=CompleteSysts)
_cache: SimpleCache = field(factory=SimpleCache)
[docs]
def __rich_repr__(self):
modules_ids = [(id(x), x) for x in self.all_modules]
pipelines_ids = {
k: [(id(z), z) for z in x] for k, x in self.base_pipelines.items()
}
yield "modules", modules_ids
yield "pipelines", pipelines_ids
[docs]
def initModules(self, metadata):
pass
# for m in self.all_modules:
# m.preloadForMeta(metadata)
[docs]
def clearCaches(self):
self._cache.clear()
for m in self.all_modules:
m.clearCache()
[docs]
def exportAdl(self, metadata, ignore_pattern=None, title=None, config_path=None):
emitter = ADLEmitter(
title=title,
config_path=config_path,
context_name=metadata.get("dataset_name"),
)
for pipeline_name, pipeline in self.base_pipelines.items():
region_statements = []
for module in pipeline:
if ignore_pattern and re.match(
ignore_pattern, module.__class__.__name__
):
continue
if (
hasattr(module, "should_run")
and module.should_run
and not module.should_run.evaluate(metadata)
):
continue
blocks = module.adlExport(metadata)
if blocks:
for block in blocks:
if block.block_type == "region_statement":
if block.comment:
region_statements.append(
ADLStatement("#", block.comment)
)
region_statements.extend(block.statements)
else:
emitter.addBlock(block)
if region_statements:
emitter.addBlock(
ADLBlock(
block_type="region",
name=pipeline_name,
statements=region_statements,
)
)
return emitter.render()
[docs]
def getUniqueModule(self, module):
found = next((x for x in self.all_modules if x == module), None)
if found is not None:
return found
else:
self.all_modules.append(module)
return module
[docs]
def neededResources(self, metadata):
needed_resources = []
for module in self.all_modules:
needed_resources.extend(module.neededResources(metadata))
return needed_resources
[docs]
def addPipeline(self, name, pipeline):
ret = []
ret.append(self.getUniqueModule(LoadColumns()))
for module in pipeline:
ret.append(self.getUniqueModule(module))
self.base_pipelines[name] = ret
[docs]
def runPipelineWithParameters(
self, pipeline, params, freeze_pipeline=False, result_container_name=None
):
module_keys = [x.selfkey for x in pipeline]
key = hash(freeze((module_keys, params)))
logger.debug(f"Pipeline execution key is {key}")
if key in self._cache:
logger.debug(f"Found key {key}, using cached columns")
return self._cache[key], None
else:
logger.debug(f"Did not find key {key}, recomputing")
params = copy.deepcopy(params)
complete_pipeline = []
to_add = deque(pipeline)
current_spec, columns = None, None
if result_container_name is None:
result_container = None
else:
result_container = ResultGroup(
result_container_name, metadata={"pipeline": result_container_name}
)
while to_add:
head = to_add.popleft()
if (
columns is not None
and head.should_run is not None
and not head.should_run.evaluate(columns.metadata)
):
continue
complete_pipeline.append(head)
if columns is not None:
columns = columns.copy()
current_spec = getPipelineSpecs(complete_pipeline, columns.metadata)
columns, results = head(columns, params)
else:
columns, results = head(params), []
if not result_container:
continue
results = deque(results)
while results:
res = results.popleft()
if isinstance(res, ResultBase):
result_container.addResult(res)
elif isinstance(res, ModuleAddition) and not freeze_pipeline:
module = res.analyzer_module
if res.run_builder is None:
raise NotImplementedError()
module = self.getUniqueModule(module)
if module.should_run is None or module.should_run.evaluate(
columns.metadata
):
logger.debug(f"Adding new module {module} to pipeline")
complete_pipeline.append(module)
params = ChainMap(params, res.this_module_parameters)
else:
logger.debug("RUNNING MULTI PARAMETER PIPELINE!!")
if res.run_builder is DEFAULT_RUN_BUILDER:
run_builder = self.default_run_builder
else:
run_builder = res.run_builder
param_dicts = run_builder(current_spec, columns.metadata)
to_run = [
(x, getWithValues(current_spec, params | y))
for x, y in param_dicts
]
everything = []
for name, params_set in to_run:
c, _ = self.runPipelineWithParameters(
complete_pipeline,
params_set,
freeze_pipeline=True,
result_container_name=None,
)
everything.append((name, c))
logger.debug(
f"Running node {module} with {len(everything)} parameter sets"
)
r = module(everything, res.this_module_parameters or {})
results.extendleft(r)
logger.debug("FINISHED RUNNING MULTIPARAMETER PIPELINE!")
else:
raise RuntimeError(
f"Invalid object type returned from analyzer module. {res}"
)
self._cache[key] = columns
return columns, result_container
[docs]
def run(self, chunk, metadata, pipelines=None):
pipelines = pipelines or list(self.base_pipelines)
root_container = ResultGroup("ROOT")
dataset_container = ResultGroup(metadata["dataset_name"])
sample_container = ResultGroup(metadata["sample_name"], metadata=metadata)
pipeline_container = ResultGroup("pipelines")
root_container.addResult(dataset_container)
dataset_container.addResult(sample_container)
sample_container.addResult(ResultProvenance("_provenance", chunk.toFileSet()))
sample_container.addResult(pipeline_container)
metadata = copy.deepcopy(metadata)
metadata["chunk"] = asdict(chunk)
for k, pipeline in self.base_pipelines.items():
if k not in pipelines:
continue
spec = getPipelineSpecs(pipeline, metadata)
vals = getWithValues(spec, {"chunk": chunk, "metadata": metadata})
_, result = self.runPipelineWithParameters(
pipeline,
vals,
result_container_name=k,
)
pipeline_container.addResult(result)
return root_container
@classmethod
def _structure(cls, data: dict, conv) -> Analyzer:
analyzer = cls()
data = data.copy()
builder = data.pop("default_run_builder", None)
if builder is not None:
analyzer.default_run_builder = conv.structure(builder, RunBuilder)
for k in data:
data[k] = list(flatten(data[k]))
pipelines = conv.structure(data, dict[str, list[AnalyzerModule]])
for k, pipeline in pipelines.items():
analyzer.addPipeline(k, pipeline)
return analyzer
def _unstructure(self, conv) -> dict:
return {
x: conv.unstructure([z.analyzer_module for z in y])
for x, y in self.base_pipelines.items()
}