def test_register_parser() -> None:
class SimpleConstraint(BaseNodeConstraint):
def __init__(self, name: str = "defaultname") -> None:
self.name = name
def satisfied_by_node(self, node: Node) -> SatisfiedResult:
return SatisfiedResult("success")
def to_dict(self) -> Dict:
return {"custom-parser": self.name}
@staticmethod
def from_dict(d: Dict) -> "SimpleConstraint":
assert "custom-parser" in d
return SimpleConstraint(d["custom-parser"])
def __str__(self) -> str:
return "SimpleConstraint({})".format(self.name)
def __eq__(self, other: object) -> bool:
if not isinstance(other, SimpleConstraint):
return False
return self.name == other.name
register_parser("custom-parser", SimpleConstraint.from_dict)
assert get_constraints([{"custom-parser": "a"}]) == [SimpleConstraint("a")]
def __init__(
self,
name: Optional[Union[str, ht.JobId]] = None,
constraints: Any = None,
iterations: int = 1,
node_count: int = 0,
colocated: bool = False,
packing_strategy: Optional[PackingStrategyType] = None,
executing_hostnames: Optional[List[ht.Hostname]] = None,
) -> None:
if not name:
global _default_job_id
_default_job_id = _default_job_id + 1
name = ht.JobId(str(_default_job_id))
_check_type(name=name, type=str)
if packing_strategy is not None:
assert PackingStrategy.is_valid(
packing_strategy), "Invalid packing_strategy {}".format(
packing_strategy)
self.__packing_strategy = packing_strategy
elif node_count > 0:
self.__packing_strategy = PackingStrategy.SCATTER
else:
self.__packing_strategy = PackingStrategy.PACK
self.__name = ht.JobId(name)
_check_type(iterations=iterations, type=int)
self.__iterations = iterations
self.__iterations_remaining = iterations
_check_type(node_count=node_count, type=int)
self.__node_count = node_count
self.__nodes_remaining = node_count
_check_type(colocated=colocated, type=bool)
self.__colocated = colocated
self.__metadata: Dict[str, Any] = {}
if constraints is None:
constraints = []
if not isinstance(constraints, list):
constraints = [constraints]
self._constraints = get_constraints(constraints)
def update_assignment_id(constraints: Iterable[Any]) -> None:
for constraint in constraints:
assert constraint is not None
if hasattr(constraint, "assignment_id"):
constraint.assignment_id = self.name
update_assignment_id(constraint.get_children())
update_assignment_id(self._constraints)
self.__executing_hostnames = executing_hostnames or []
def __init__(
self,
name: str,
constraints: Optional[List[Constraint]] = None,
members: Optional[List[str]] = None,
) -> None:
self.name = name
self.constraints = get_constraints(constraints or [])
self.members: List[str] = members or []
def _mpi_job(job_name="1", nodes=1, placeby="pg", resources=None):
resources = resources or {"ncpus": 2}
constraints = get_constraints([resources])
constraints.append(InAPlacementGroup())
constraints.append(ExclusiveNode())
return Job(
job_name,
constraints=constraints,
node_count=nodes,
colocated=True,
)
def _parse_contraint(constraint_expr: str) -> List[NodeConstraint]:
try:
constraint_parsed = json.loads(constraint_expr)
except Exception as e:
print(
"Could not parse constraint as json '{}' - {}".format(constraint_expr, e),
file=sys.stderr,
)
sys.exit(1)
if not isinstance(constraint_parsed, list):
constraint_parsed = [constraint_parsed]
return get_constraints(constraint_parsed)
def validate_default_hostgroups(
config: Dict, ge_env: GridEngineEnvironment, warn: WarnFunction,
) -> bool:
default_hostgroups = config.get("gridengine", {}).get("default_hostgroups", [])
failure = False
for dh in default_hostgroups:
# ensure it is a {"select": {}, "hostgroups": []} expression
if set(list(dh.keys())) != set(["select", "hostgroups"]):
warn(
" Invalid entry in gridengine.default_hostgroups. Expected select and hostgroups as keys. %s",
dh,
)
failure = True
continue
try:
constraints.get_constraints(dh["select"])
except Exception as e:
warn(" Invalid select constraint: %s - %s", dh["select"], e)
failure = True
continue
hostgroups = dh["hostgroups"]
if isinstance(hostgroups, str):
hostgroups = [hostgroups]
for hg in hostgroups:
if hg not in ge_env.hostgroups:
warn(
" Unknown hostgroup %s. See gridengine.default_hostgroups - %s",
hg,
dh,
)
failure = True
return not failure
def read_hostgroups(autoscale_config: Dict, qbin: QBin) -> List[Hostgroup]:
# map each host (lowercase) to its set of hostgroups
ret: List[Hostgroup] = []
hg_config = autoscale_config.get("gridengine", {}).get("hostgroups", {})
for hg_name in qbin.qconf(["-shgrpl"]).split():
members = qbin.qconf(["-shgrp_resolved", hg_name]).split()
members = [h.split(".")[0].lower() for h in members]
constraints = hg_config.get(hg_name, {}).get("constraints", []) or []
if not isinstance(constraints, list):
constraints = [constraints]
parsed_constraints = constraintslib.get_constraints(constraints)
hostgroup = Hostgroup(hg_name, parsed_constraints, members)
ret.append(hostgroup)
return ret
def get_node_hostgroups(config: Dict, node: Node) -> List[str]:
hostgroups_expr = node.metadata.get("gridengine_hostgroups")
if not hostgroups_expr:
hostgroups_expr = node.software_configuration.get(
"gridengine_hostgroups")
if not hostgroups_expr:
default_hostgroups = config.get("gridengine",
{}).get("default_hostgroups", [])
for dh in default_hostgroups:
if "select" not in dh:
logging.warning(
"Missing key 'select' in gridengine.default_hostgroups %s",
dh)
continue
if "hostgroups" not in dh:
logging.warning(
"Missing key 'hostgroups' in gridengine.default_hostgroups %s",
dh)
continue
constraint_list = constraints.get_constraints(dh["select"])
satisfied = True
for c in constraint_list:
if not c.satisfied_by_node(node):
satisfied = False
break
if satisfied:
hostgroups = dh["hostgroups"]
if isinstance(hostgroups, str):
hostgroups = [hostgroups]
hostgroups_expr = " ".join(hostgroups)
# set it in metadata so we can output it in the cli
node.metadata["gridengine_hostgroups"] = hostgroups_expr
if hostgroups_expr:
return re.split(",| +", hostgroups_expr)
return []
def add_constraint(self, constraint: typing.Any) -> None:
if not isinstance(constraint, list):
constraint = [constraint]
parsed_cons = get_constraints(constraint)
self._constraints.extend(parsed_cons)
def test_memory() -> None:
c = get_constraints([{"memgb": 1}])
node = SchedulerNode("test", {"memgb": 4.0})
m = minimum_space(c, node)
assert isinstance(m, int)
assert m == 4