Spaces:
Sleeping
Sleeping
File size: 2,227 Bytes
89cb5bf | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 | from __future__ import annotations
import math
import random
from dataclasses import dataclass
from typing import Any, Mapping
@dataclass(frozen=True)
class OpDurationSpec:
"""Normal-distribution duration config for one operation type."""
mean: float
variance: float = 0.0
min_value: float = 0.0
def sample(self, rng: random.Random) -> float:
if self.variance < 0:
raise ValueError(f"variance must be non-negative, got {self.variance}")
if self.mean < 0:
raise ValueError(f"mean must be non-negative, got {self.mean}")
if self.variance == 0:
return max(self.min_value, float(self.mean))
value = rng.gauss(float(self.mean), math.sqrt(float(self.variance)))
return max(float(self.min_value), value)
def op_type_value(op_type: Any) -> str:
"""Return a stable string value for pp_scheduler.OpType or a plain string."""
value = getattr(op_type, "value", op_type)
return str(value)
class DurationSampler:
def __init__(
self,
specs: Mapping[Any, OpDurationSpec | Mapping[str, float]],
*,
default_spec: OpDurationSpec | None = None,
seed: int | None = None,
):
self.rng = random.Random(seed)
self.default_spec = default_spec
self.specs = {
op_type_value(op_type): self._coerce_spec(spec)
for op_type, spec in specs.items()
}
def sample(self, op_type: Any, fallback_duration: float | None = None) -> float:
op_value = op_type_value(op_type)
spec = self.specs.get(op_value, self.default_spec)
if spec is not None:
return spec.sample(self.rng)
if fallback_duration is not None:
return max(0.0, float(fallback_duration))
raise KeyError(f"No duration spec for op type {op_value!r}")
@staticmethod
def _coerce_spec(spec: OpDurationSpec | Mapping[str, float]) -> OpDurationSpec:
if isinstance(spec, OpDurationSpec):
return spec
return OpDurationSpec(
mean=float(spec["mean"]),
variance=float(spec.get("variance", 0.0)),
min_value=float(spec.get("min_value", 0.0)),
)
|