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 "kueue", 

58 } 

59) 

60 

61 

62class EnablementOverrideError(ValueError): 

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

64 

65 

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

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

68 

69 Accepts the two shapes a caller can produce interchangeably — 

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

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

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

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

74 """ 

75 names: list[str] = [] 

76 for entry in raw: 

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

78 candidate = part.strip() 

79 if candidate: 

80 names.append(candidate) 

81 return names 

82 

83 

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

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

86 

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

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

89 ``feature_enabled_overrides=fsx_lustre`` and 

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

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

92 or re-recorded run's argv stable. 

93 

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

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

96 

97 Raises: 

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

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

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

101 operator asked for. 

102 """ 

103 requested = set(split_override_names(raw)) 

104 unknown = sorted(requested - FEATURE_OVERRIDE_KEYS - HELM_CHART_CONFIG_KEYS) 

105 if unknown: 

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

107 raise EnablementOverrideError( 

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

109 ) 

110 

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

112 features = sorted(requested & FEATURE_OVERRIDE_KEYS) 

113 charts = sorted(requested & HELM_CHART_CONFIG_KEYS) 

114 if features: 

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

116 if charts: 

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

118 return context