forked from equinor/ert
-
Notifications
You must be signed in to change notification settings - Fork 0
Commit
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
Refactor FeatureToggling to FeatureScheduler
We only need to toggle scheduler. If we ever need to toggle some other feature, we can implement that as a separate class when the time comes. This commit removes the complexity associated with supporting multiple such features.
- Loading branch information
Showing
9 changed files
with
106 additions
and
178 deletions.
There are no files selected for viewing
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -1,119 +1,76 @@ | ||
| from __future__ import annotations | ||
|
|
||
| import logging | ||
| import os | ||
| from argparse import ArgumentParser | ||
| from copy import deepcopy | ||
| from typing import TYPE_CHECKING, Dict, Optional, Union | ||
|
|
||
| if TYPE_CHECKING: | ||
| from ert.namespace import Namespace | ||
| from typing import TYPE_CHECKING, Optional | ||
|
|
||
| logger = logging.getLogger() | ||
|
|
||
| class _Feature: | ||
| def __init__( | ||
| self, default: Optional[bool], msg: Optional[str] = None, optional: bool = False | ||
| ) -> None: | ||
| self._value = default | ||
| self.msg = msg | ||
| self.optional = optional | ||
|
|
||
| def validate_value(self, value: Union[bool, str, None]) -> Optional[bool]: | ||
| if type(value) is bool or value is None: | ||
| return value | ||
| elif value.lower() in ["true", "1"]: | ||
| return True | ||
| elif value.lower() in ["false", "0"]: | ||
| return False | ||
| elif self.optional and value.lower() in ["default", ""]: | ||
| return None | ||
| else: | ||
| raise ValueError( | ||
| f"This option can only be set to {'True/1, False/0 or Default/<empty>' if self.optional else 'True/1 or False/0'}" | ||
| ) | ||
|
|
||
| @property | ||
| def value(self) -> Optional[bool]: | ||
| return self._value | ||
| if TYPE_CHECKING: | ||
| from argparse import ArgumentParser | ||
|
|
||
| @value.setter | ||
| def value(self, value: Optional[bool]) -> None: | ||
| self._value = self.validate_value(value) | ||
| from ert.config.parsing.queue_system import QueueSystem | ||
| from ert.namespace import Namespace | ||
|
|
||
|
|
||
| class FeatureToggling: | ||
| _conf_original: Dict[str, _Feature] = { | ||
| "scheduler": _Feature( | ||
| default=None, | ||
| msg="Default value for use of Scheduler has been overridden\n" | ||
| "This is experimental and may cause problems", | ||
| optional=True, | ||
| ), | ||
| class FeatureScheduler: | ||
| _DEFAULTS = { | ||
| "LOCAL": True, | ||
| "LSF": False, | ||
| "SLURM": False, | ||
| "TORQUE": False, | ||
| } | ||
|
|
||
| _conf = deepcopy(_conf_original) | ||
|
|
||
| @staticmethod | ||
| def is_enabled(feature_name: str) -> bool: | ||
| return FeatureToggling._conf[feature_name].value is True | ||
|
|
||
| @staticmethod | ||
| def value(feature_name: str) -> Optional[bool]: | ||
| return FeatureToggling._conf[feature_name].value | ||
| _value: Optional[bool] = None | ||
|
|
||
| @classmethod | ||
| def is_enabled(cls, queue_system: QueueSystem) -> bool: | ||
| if cls._value is not None: | ||
| return cls._value | ||
| return cls._DEFAULTS[queue_system.name] | ||
|
|
||
| @classmethod | ||
| def set_value(cls, args: Namespace) -> None: | ||
| if ((value := cls._get_from_args(args)) is not None) or ( | ||
| (value := cls._get_from_env()) is not None | ||
| ): | ||
| cls._value = value | ||
| else: | ||
| cls._value = None | ||
|
|
||
| @staticmethod | ||
| def add_feature_toggling_args(parser: ArgumentParser) -> None: | ||
| for name, feature in FeatureToggling._conf.items(): | ||
| env_var_name = f"ERT_FEATURE_{name.replace('-', '_').upper()}" | ||
| env_value: Union[bool, str, None] = None | ||
| if env_var_name in os.environ: | ||
| try: | ||
| feature.value = feature.validate_value(os.environ[env_var_name]) | ||
| except ValueError as e: | ||
| # TODO: this is a bit spammy. It will get called 6 times for each incorrect env var. | ||
| logging.getLogger().warning( | ||
| f"Failed to set {env_var_name} to '{os.environ[env_var_name]}'. {e}" | ||
| ) | ||
|
|
||
| if not feature.optional: | ||
| parser.add_argument( | ||
| f"--{'disable' if feature.value else 'enable'}-{name}", | ||
| action="store_false" if feature.value else "store_true", | ||
| help=f"Toggle {name} (Warning: This is experimental)", | ||
| dest=f"feature-{name}", | ||
| default=env_value if env_value is not None else feature.value, | ||
| ) | ||
| else: | ||
| group = parser.add_mutually_exclusive_group() | ||
| group.add_argument( | ||
| f"--enable-{name}", | ||
| action="store_true", | ||
| help=f"Enable {name}", | ||
| dest=f"feature-{name}", | ||
| default=feature.value, | ||
| ) | ||
| group.add_argument( | ||
| f"--disable-{name}", | ||
| action="store_false", | ||
| help=f"Disable {name}", | ||
| dest=f"feature-{name}", | ||
| default=feature.value, | ||
| ) | ||
| def add_to_argparse(parser: ArgumentParser) -> None: | ||
| group = parser.add_mutually_exclusive_group() | ||
| group.add_argument( | ||
| "--enable-scheduler", | ||
| action="store_true", | ||
| help="Enable new scheduler", | ||
| dest="feature_scheduler", | ||
| default=None, | ||
| ) | ||
| group.add_argument( | ||
| "--disable-scheduler", | ||
| action="store_false", | ||
| help="Disable new scheduler", | ||
| dest="feature_scheduler", | ||
| default=None, | ||
| ) | ||
|
|
||
| @staticmethod | ||
| def update_from_args(args: "Namespace") -> None: | ||
| pattern = "feature-" | ||
| feature_args = [arg for arg in vars(args).items() if arg[0].startswith(pattern)] | ||
| for name, value in feature_args: | ||
| name = name[len(pattern) :] | ||
| if name in FeatureToggling._conf: | ||
| FeatureToggling._conf[name].value = value | ||
|
|
||
| # Print warnings for enabled features. | ||
| for name, feature in FeatureToggling._conf.items(): | ||
| if FeatureToggling.is_enabled(name) and feature.msg is not None: | ||
| logging.getLogger().warning( | ||
| f"{feature.msg}\nValue is set to {feature.value}" | ||
| ) | ||
| def _get_from_env() -> Optional[bool]: | ||
| if (value := os.environ.get("ERT_FEATURE_SCHEDULER")) is None: | ||
| return None | ||
| value = value.lower() | ||
| if value in ("true", "1"): | ||
| return True | ||
| elif value in ("false", "0"): | ||
| return False | ||
| elif value in ("auto", "default", ""): | ||
| return None | ||
| raise ValueError( | ||
| "This option can only be set to 'true'/'1', 'false'/'0' or 'auto'/'default'/''" | ||
| ) | ||
|
|
||
| @staticmethod | ||
| def reset() -> None: | ||
| FeatureToggling._conf = deepcopy(FeatureToggling._conf_original) | ||
| def _get_from_args(args: Namespace) -> Optional[bool]: | ||
| return args.feature_scheduler |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -1,14 +1,17 @@ | ||
| from argparse import ArgumentParser | ||
|
|
||
| import pytest | ||
|
|
||
| from ert.__main__ import ert_parser | ||
| from ert.cli.main import run_cli as cli_runner | ||
| from ert.shared.feature_toggling import FeatureToggling | ||
| from ert.shared.feature_toggling import FeatureScheduler | ||
|
|
||
|
|
||
| def run_cli(*args): | ||
| parser = ArgumentParser(prog="test_main") | ||
| parsed = ert_parser(parser, args) | ||
| FeatureToggling.update_from_args(parsed) | ||
| res = cli_runner(parsed) | ||
| FeatureToggling.reset() | ||
| return res | ||
| with pytest.MonkeyPatch.context() as monkeypatch: | ||
| monkeypatch.setattr(FeatureScheduler, "_value", None) | ||
| parser = ArgumentParser(prog="test_main") | ||
| parsed = ert_parser(parser, args) | ||
| FeatureScheduler.set_value(parsed) | ||
| res = cli_runner(parsed) | ||
| return res |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Oops, something went wrong.