Coverage for scripts / live_release_validation / checks / schedulers.py: 100.00%

1"""Scheduler enablement resolution for the ``schedulers`` action. 

2 

3The harness has no direct Kubernetes access, so scheduler validation rests on 

4one architectural fact: the default kube-scheduler ignores pods whose 

5``spec.schedulerName`` it does not own, and Kueue's webhook keeps a 

6queue-labeled Job suspended until quota admission. A probe Job that names a 

7scheduler (or carries the Kueue queue label) therefore only ever completes if 

8that scheduler actually did its work — completion *is* the scheduling proof, 

9observable through the deployed manifest API like every other validation Job. 

10 

11This module resolves which schedulers the run must probe: the cdk.json helm 

12block decides, and ``--optional-schedulers`` run overrides force off-by-default 

13schedulers on (threaded to CDK as the ``helm_enabled_overrides`` context by the 

14runner, so the deployed chart set and this resolution share one source). 

15""" 

16 

17from __future__ import annotations 

18 

19from typing import Any 

20 

21from ..models import RunContext 

22 

23#: Schedulers the action can probe, in probe order: name -> cdk.json helm key. 

24#: 

25#: The Kubeflow Trainer is deliberately NOT probed here: it is a workload 

26#: controller (TrainJob -> JobSet compilation), not a pod scheduler, so the 

27#: "completion proves the scheduler did its work" contract this action rests 

28#: on does not apply. Its live proof rides two other rails instead: the 

29#: examples harness runs the kubeflow-trainjob example end to end (gang 

30#: all-reduce through the shipped runtime), and the topology action's 

31#: helmValidation asserts the kubeflow-trainer release converged exactly as 

32#: the deployment config demanded. 

33PROBED_SCHEDULERS: dict[str, str] = { 

34 "volcano": "volcano", 

35 "kueue": "kueue", 

36 "yunikorn": "yunikorn", 

37 "slurm": "slurm", 

38} 

39 

40#: Off-by-default schedulers a run may force-enable with --optional-schedulers. 

41OPTIONAL_SCHEDULERS: tuple[str, ...] = ("yunikorn", "slurm") 

42 

43 

44def effective_scheduler_enablement(ctx: RunContext) -> dict[str, dict[str, Any]]: 

45 """Resolve every probed scheduler's effective enablement and its source. 

46 

47 Mirrors ``gco.stacks.regional_stack._helm_chart_enabled`` semantics for 

48 the probed keys: a run override wins, then the cdk.json toggle, and a 

49 missing key defaults to enabled (the historical chart_map behavior). 

50 """ 

51 helm_config = ctx.cdk_context.get("helm") 

52 if not isinstance(helm_config, dict): 

53 helm_config = {} 

54 overrides = set(ctx.settings.optional_schedulers) 

55 unknown = sorted(overrides - set(OPTIONAL_SCHEDULERS)) 

56 if unknown: 

57 raise RuntimeError( 

58 "optional_schedulers contains names that are not optional schedulers: " 

59 + ", ".join(unknown) 

60 ) 

61 

62 enablement: dict[str, dict[str, Any]] = {} 

63 for scheduler, helm_key in PROBED_SCHEDULERS.items(): 

64 if scheduler in overrides: 

65 enablement[scheduler] = {"enabled": True, "source": "run-override"} 

66 continue 

67 chart_config = helm_config.get(helm_key) 

68 configured = ( 

69 bool(chart_config.get("enabled", True)) if isinstance(chart_config, dict) else True 

70 ) 

71 enablement[scheduler] = {"enabled": configured, "source": f"cdk.json helm.{helm_key}"} 

72 return enablement