src/swf_testbed_decision_box/models.py

0001 from dataclasses import dataclass, field
0002 from types import MappingProxyType
0003 from typing import Any, Iterable, Mapping
0004
0005
0006 @dataclass(frozen=True, order=True)
0007 class Site:
0008     """PanDA site and dataset suffix used by the decision box."""
0009
0010     name: str
0011
0012     def __post_init__(self) -> None:
0013         if not self.name:
0014             raise ValueError("site name cannot be empty")
0015
0016
0017 @dataclass(frozen=True)
0018 class FileDID:
0019     """Rucio file DID in scope:name form."""
0020
0021     scope: str
0022     name: str
0023
0024     @classmethod
0025     def parse(cls, value: str, default_scope: str | None = None) -> "FileDID":
0026         if ":" in value:
0027             scope, name = value.split(":", 1)
0028         elif default_scope:
0029             scope, name = default_scope, value
0030         else:
0031             raise ValueError(f"file DID must be scope:name, got {value!r}")
0032         if not scope or not name:
0033             raise ValueError(f"invalid file DID {value!r}")
0034         return cls(scope=scope, name=name)
0035
0036     def __str__(self) -> str:
0037         return f"{self.scope}:{self.name}"
0038
0039
0040 @dataclass(frozen=True)
0041 class DecisionContext:
0042     """Inputs available to policy evaluation for one STF file."""
0043
0044     run_dataset: str
0045     file_did: FileDID
0046     available_sites: tuple[Site, ...]
0047     sequence: int
0048     run_number: str | None = None
0049     message: Mapping[str, Any] = field(default_factory=lambda: MappingProxyType({}))
0050     run_conditions: Mapping[str, Any] = field(default_factory=lambda: MappingProxyType({}))
0051     metadata: Mapping[str, Any] = field(default_factory=lambda: MappingProxyType({}))
0052
0053
0054 @dataclass(frozen=True)
0055 class SiteAssignment:
0056     """Decision result for one STF file."""
0057
0058     file_did: FileDID
0059     sites: tuple[Site, ...]
0060     reason: str
0061
0062     @classmethod
0063     def from_sites(cls, file_did: FileDID, sites: Iterable[Site], reason: str) -> "SiteAssignment":
0064         unique_sites = tuple(sorted(set(sites), key=lambda site: site.name))
0065         return cls(file_did=file_did, sites=unique_sites, reason=reason)
0066
0067
0068 @dataclass(frozen=True)
0069 class Decision:
0070     """Applied decision, including the datasets that were mutated."""
0071
0072     run_dataset: str
0073     full_dataset: str
0074     file_did: FileDID
0075     site_datasets: tuple[str, ...]
0076     reason: str