Coverage for gco / enablement_overrides.py: 100.00%

1"""Run-scoped enablement overrides shared by the CLI and the recorders. 

2 

3GCO ships every optional add-on **off** in ``cdk.json`` because each one 

4carries real recurring cost (an FSx for Lustre filesystem has a 1.2 TiB 

5provisioned floor; Aurora Serverless v2 keeps a writer and a reader; Valkey 

6Serverless bills storage and ECPUs). Two CDK context keys exist so a single 

7run can force those features on *without rewriting the committed config*: 

8 

9``feature_enabled_overrides`` 

10 Infrastructure blocks whose top-level ``enabled`` flag may be forced on 

11 (see ``gco.config.config_loader.parse_feature_enabled_overrides``). 

12 

13``helm_enabled_overrides`` 

14 ``helm`` block keys whose chart may be forced on (see 

15 ``gco.stacks.regional_stack._parse_helm_enabled_overrides``). 

16 

17Both consumers live in modules that import ``aws_cdk``. The CLI must be able 

18to validate ``--enable`` names and build the ``--context`` pairs *before* the 

19CDK Python toolchain is known to be importable — ``StackManager._run_cdk`` 

20deliberately fails with an actionable message in that case rather than an 

21``ImportError`` at startup. So the canonical name sets live here, in a module 

22with no third-party imports, and ``tests/test_enablement_overrides.py`` pins 

23them in lockstep with the two authoritative parsers. 

24 

25Overrides are one-way: they can only *enable*. A feature an operator turned 

26off stays off unless it is named explicitly. 

27""" 

28 

29from __future__ import annotations 

30 

31from collections.abc import Iterable 

32 

33#: CDK context key carrying forced-on infrastructure feature blocks. 

34FEATURE_OVERRIDE_CONTEXT_KEY = "feature_enabled_overrides" 

35 

36#: CDK context key carrying forced-on Helm chart keys. 

37HELM_OVERRIDE_CONTEXT_KEY = "helm_enabled_overrides" 

38 

39#: cdk.json blocks ``feature_enabled_overrides`` accepts. Lockstep with 

40#: ``gco.config.config_loader.FEATURE_OVERRIDE_KEYS``. 

41FEATURE_OVERRIDE_KEYS = frozenset({"aurora_pgvector", "valkey", "fsx_lustre", "vector_store"}) 

42 

43#: cdk.json ``helm`` block keys ``helm_enabled_overrides`` accepts. Lockstep 

44#: with ``gco.stacks.regional_stack._HELM_CHART_CONFIG_KEYS``. 

45HELM_CHART_CONFIG_KEYS = frozenset( 

46 { 

47 "aws_load_balancer_controller", 

48 "keda", 

49 "aws_efa_device_plugin", 

50 "aws_neuron_device_plugin", 

51 "volcano", 

52 "kuberay", 

53 "cert_manager", 

54 "slurm", 

55 "yunikorn", 

56 "kubeflow_trainer", 

57 "argocd", 

58 "crossplane", 

59 "kueue", 

60 } 

61) 

62 

63 

64class EnablementOverrideError(ValueError): 

65 """Raised when a requested override name is not a known feature or chart.""" 

66 

67 

68def split_override_names(raw: Iterable[str]) -> list[str]: 

69 """Flatten repeated and comma-joined ``--enable`` values into bare names. 

70 

71 Accepts the two shapes a caller can produce interchangeably — 

72 ``--enable valkey --enable slurm`` and ``--enable valkey,slurm`` — and 

73 drops empty segments so a trailing comma is not an error. Order is 

74 preserved for the caller's benefit; duplicates are left intact because 

75 :func:`route_enablement_overrides` de-duplicates when it groups. 

76 """ 

77 names: list[str] = [] 

78 for entry in raw: 

79 for part in entry.split(","): 

80 candidate = part.strip() 

81 if candidate: 

82 names.append(candidate) 

83 return names 

84 

85 

86def route_enablement_overrides(raw: Iterable[str]) -> dict[str, str]: 

87 """Group requested names into the CDK ``--context`` pairs they belong to. 

88 

89 The two namespaces are disjoint, so each name routes unambiguously: a 

90 caller says ``--enable fsx_lustre,slurm`` and gets both 

91 ``feature_enabled_overrides=fsx_lustre`` and 

92 ``helm_enabled_overrides=slurm``. Names are sorted so the resulting 

93 context is byte-identical for any input ordering, which keeps a resumed 

94 or re-recorded run's argv stable. 

95 

96 Returns an empty mapping when nothing was requested, so callers can pass 

97 the result straight to ``StackManager.set_extra_cdk_context``. 

98 

99 Raises: 

100 EnablementOverrideError: if any name is neither a known feature block 

101 nor a known Helm chart key. Failing here — before any AWS call — 

102 means a typo can never silently deploy without the feature the 

103 operator asked for. 

104 """ 

105 requested = set(split_override_names(raw)) 

106 unknown = sorted(requested - FEATURE_OVERRIDE_KEYS - HELM_CHART_CONFIG_KEYS) 

107 if unknown: 

108 valid = ", ".join(sorted(FEATURE_OVERRIDE_KEYS | HELM_CHART_CONFIG_KEYS)) 

109 raise EnablementOverrideError( 

110 f"Unknown --enable name(s): {', '.join(unknown)}. Valid: {valid}" 

111 ) 

112 

113 context: dict[str, str] = {} 

114 features = sorted(requested & FEATURE_OVERRIDE_KEYS) 

115 charts = sorted(requested & HELM_CHART_CONFIG_KEYS) 

116 if features: 

117 context[FEATURE_OVERRIDE_CONTEXT_KEY] = ",".join(features) 

118 if charts: 

119 context[HELM_OVERRIDE_CONTEXT_KEY] = ",".join(charts) 

120 return context