Merge branch 'mpiano/CLUSTERMAN-666_event' into master

Author: piax93

clusterman/autoscaler/toggle.py

@@ -0,0 +1,59 @@
+# Copyright 2019 Yelp Inc.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+import time
+from typing import Union
+
+import staticconf
+
+from clusterman.aws.client import dynamodb
+from clusterman.util import AUTOSCALER_PAUSED
+from clusterman.util import CLUSTERMAN_STATE_TABLE
+from clusterman.util import parse_time_string
+
+
+def disable_autoscaling(cluster: str, pool: str, scheduler: str, until: Union[str, int, float]):
+    """Disable autoscaling for a pool
+
+    :param str cluster: name of the cluster
+    :param str pool: name of the pool
+    :param str scheduler: cluster scheduler
+    :param str until: how long should it remain disabled
+    """
+    expiration = parse_time_string(until).timestamp if isinstance(until, str) else int(until)
+    state = {
+        "state": {"S": AUTOSCALER_PAUSED},
+        "entity": {"S": f"{cluster}.{pool}.{scheduler}"},
+        "timestamp": {"N": str(int(time.time()))},
+        "expiration_timestamp": {"N": str(expiration)},
+    }
+    dynamodb.put_item(
+        TableName=staticconf.read("aws.state_table", default=CLUSTERMAN_STATE_TABLE),
+        Item=state,
+    )
+
+
+def enable_autoscaling(cluster: str, pool: str, scheduler: str):
+    """Re-enable autoscaling for a pool
+
+    :param str cluster: name of the cluster
+    :param str pool: name of the pool
+    :param str scheduler: cluster scheduler
+    """
+    dynamodb.delete_item(
+        TableName=staticconf.read("aws.state_table", default=CLUSTERMAN_STATE_TABLE),
+        Key={
+            "state": {"S": AUTOSCALER_PAUSED},
+            "entity": {"S": f"{cluster}.{pool}.{scheduler}"},
+        },
+    )

clusterman/cli/toggle.py

@@ -22,15 +22,15 @@
 from clusterman.args import add_pool_arg
 from clusterman.args import add_scheduler_arg
 from clusterman.args import subparser
-from clusterman.aws.client import dynamodb
+from clusterman.autoscaler.toggle import disable_autoscaling
+from clusterman.autoscaler.toggle import enable_autoscaling
 from clusterman.aws.client import sts
 from clusterman.cli.util import timeout_wrapper
 from clusterman.exceptions import AccountNumberMistmatchError
-from clusterman.util import AUTOSCALER_PAUSED
 from clusterman.util import autoscaling_is_paused
-from clusterman.util import CLUSTERMAN_STATE_TABLE
 from clusterman.util import parse_time_string
 
+
 logger = colorlog.getLogger(__name__)
 
 
@@ -48,24 +48,11 @@ def ensure_account_id(cluster) -> None:
 @timeout_wrapper
 def disable(args: argparse.Namespace) -> None:
     ensure_account_id(args.cluster)
-
-    state = {
-        "state": {"S": AUTOSCALER_PAUSED},
-        "entity": {"S": f"{args.cluster}.{args.pool}.{args.scheduler}"},
-        "timestamp": {"N": str(int(time.time()))},
-    }
-
     if not args.until:
         print("Default has changed; autoscaler will be re-enabled in 30m since no --until flag was used.")
         args.until = "30m"
 
-    if args.until:
-        state["expiration_timestamp"] = {"N": str(parse_time_string(args.until).timestamp)}
-
-    dynamodb.put_item(
-        TableName=staticconf.read("aws.state_table", default=CLUSTERMAN_STATE_TABLE),
-        Item=state,
-    )
+    disable_autoscaling(args.cluster, args.pool, args.scheduler, args.until)
 
     time.sleep(1)  # Give DynamoDB some time to settle
 
@@ -86,14 +73,7 @@ def disable(args: argparse.Namespace) -> None:
 @timeout_wrapper
 def enable(args: argparse.Namespace) -> None:
     ensure_account_id(args.cluster)
-
-    dynamodb.delete_item(
-        TableName=staticconf.read("aws.state_table", default=CLUSTERMAN_STATE_TABLE),
-        Key={
-            "state": {"S": AUTOSCALER_PAUSED},
-            "entity": {"S": f"{args.cluster}.{args.pool}.{args.scheduler}"},
-        },
-    )
+    enable_autoscaling(args.cluster, args.pool, args.scheduler)
     time.sleep(1)  # Give DynamoDB some time to settle
     now = parse_time_string("now").to("local")
     if autoscaling_is_paused(args.cluster, args.pool, args.scheduler, now):

clusterman/cli/util.py

@@ -1,27 +1,38 @@
+# Copyright 2019 Yelp Inc.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
 import argparse
-import signal
 import socket
+from functools import partial
 
 import colorlog
 
+from clusterman.util import limit_function_runtime
+
+
 logger = colorlog.getLogger(__name__)
 TIMEOUT_TIME_SECONDS = 10
 
 
 def timeout_wrapper(main):
     def wrapper(args: argparse.Namespace):
-
         # After 10s, prints a warning message if the command is running from the wrong place
-        def timeout_handler(signum, frame):
+        def timeout_handler():
             warning_string = "This command is taking a long time to run; are you running from the right account?"
             if "yelpcorp" in socket.getfqdn():
                 warning_string += "\nHINT: try ssh'ing to adhoc-prod or another box in the right account."
             logger.warning(warning_string)
 
-        signal.signal(signal.SIGALRM, timeout_handler)
-        signal.alarm(TIMEOUT_TIME_SECONDS)
-
-        main(args)
-        signal.alarm(0)  # Cancel the alarm if we've gotten here
+        limit_function_runtime(partial(main, args), TIMEOUT_TIME_SECONDS, timeout_handler)
 
     return wrapper

clusterman/kubernetes/kubernetes_cluster_connector.py

@@ -74,6 +74,7 @@ class KubernetesClusterConnector(ClusterConnector):
     _unschedulable_pods: List[KubernetesPod]
     _excluded_pods_by_ip: Mapping[str, List[KubernetesPod]]
     _pods_by_ip: Mapping[str, List[KubernetesPod]]
+    _label_selectors: List[str]
 
     def __init__(self, cluster: str, pool: Optional[str], init_crd: bool = False) -> None:
         super().__init__(cluster, pool)
@@ -84,6 +85,14 @@ def __init__(self, cluster: str, pool: Optional[str], init_crd: bool = False) ->
         )
         self._nodes_by_ip = {}
         self._init_crd_client = init_crd
+        self._label_selectors = []
+        if self.pool:
+            # TODO(CLUSTERMAN-659): Switch to using just pool_label_key once the new node labels are applied everywhere
+            node_label_selector = self.pool_config.read_string(
+                "node_label_key",
+                default=self.pool_config.read_string("pool_label_key", default="clusterman.com/pool"),
+            )
+            self._label_selectors.append(f"{node_label_selector}={self.pool}")
 
     def reload_state(self) -> None:
         logger.info("Reloading nodes")
@@ -113,6 +122,16 @@ def reload_client(self) -> None:
             else None
         )
 
+    def set_label_selectors(self, label_selectors: List[str], add_to_existing: bool = False) -> None:
+        """Set label selectors for node listing purposes
+
+        :param List[str] label_selectors: list of selectors (joined with logic and)
+        :param bool add_to_existing: if set add to existing selectors rather than replacing
+        """
+        self._label_selectors = sorted(
+            (set(self._label_selectors) | set(label_selectors)) if add_to_existing else set(label_selectors)
+        )
+
     def get_num_removed_nodes_before_last_reload(self) -> int:
         previous_nodes = self._prev_nodes_by_ip
         current_nodes = self._nodes_by_ip
@@ -270,6 +289,15 @@ def create_node_migration_resource(
         except Exception as e:
             logger.error(f"Failed creating migration event resource: {e}")
 
+    def has_enough_capacity_for_pods(self) -> bool:
+        """Checks whether there are unschedulable pods due to insufficient resources
+
+        :return: True if no unschedulable pods are due to resource constraints
+        """
+        return not any(
+            reason == PodUnschedulableReason.InsufficientResources for _, reason in self.get_unschedulable_pods()
+        )
+
     def _evict_or_delete_pods(self, node_name: str, pods: List[KubernetesPod], disable_eviction: bool) -> bool:
         all_done = True
         action_name = "deleted" if disable_eviction else "evicted"
@@ -403,12 +431,8 @@ def _is_node_safe_to_kill(self, node_ip: str) -> bool:
         return True
 
     def _get_nodes_by_ip(self) -> Mapping[str, KubernetesNode]:
-        # TODO(CLUSTERMAN-659): Switch to using just pool_label_key once the new node labels are applied everywhere
-        node_label_selector = self.pool_config.read_string(
-            "node_label_key", default=self.pool_config.read_string("pool_label_key", default="clusterman.com/pool")
-        )
-        label_selector = f"{node_label_selector}={self.pool}"
-        pool_nodes = self._core_api.list_node(label_selector=label_selector).items
+        kwargs = {"label_selector": ",".join(self._label_selectors)} if self._label_selectors else {}
+        pool_nodes = self._core_api.list_node(**kwargs).items
         return {get_node_ip(node): node for node in pool_nodes}
 
     def _get_pods_info(

clusterman/migration/event.py

@@ -22,30 +22,24 @@
 import semver
 
 from clusterman.aws.markets import EC2_INSTANCE_TYPES
+from clusterman.interfaces.types import ClusterNodeMetadata
+from clusterman.migration.event_enums import ComparableConditionTarget
+from clusterman.migration.event_enums import ComparableVersion
 from clusterman.migration.event_enums import CONDITION_OPERATOR_SUPPORT_MATRIX
 from clusterman.migration.event_enums import ConditionOperator
 from clusterman.migration.event_enums import ConditionTrait
 from clusterman.util import parse_time_interval_seconds
 
 
-ComparableVersion = Union[semver.VersionInfo, packaging.version.Version]
-ComparableConditionTarget = Union[str, int, ComparableVersion]
-
-
-def _load_version_target(target: str) -> ComparableVersion:
+def _load_lsbrelease_version_target(target: str) -> ComparableVersion:
     """Validate condition target as a version
 
     :param str target: version from user input
     :return: same value if validates as version
     """
-    try:
-        parsed = semver.parse_version_info(target)
-    except ValueError:
-        # This allows supporting non-semver version, e.g. with less
-        # then 3 numeric components, as long as they comply with PEP440.
-        parsed = packaging.version.parse(target)
-        if not isinstance(parsed, packaging.version.Version):
-            raise ValueError(f"Invalid version string: {target}")
+    parsed = packaging.version.parse(target)
+    if not isinstance(parsed, packaging.version.Version):
+        raise ValueError(f"Invalid version string: {target}")
     return parsed
 
 
@@ -71,8 +65,8 @@ def _load_instance_type_target(target: str) -> str:
 
 
 CONDITION_TARGET_LOADERS: Dict[ConditionTrait, Callable[[str], Union[str, int, List[str]]]] = {
-    ConditionTrait.KERNEL: _load_version_target,
-    ConditionTrait.LSBRELEASE: _load_version_target,
+    ConditionTrait.KERNEL: semver.VersionInfo.parse,
+    ConditionTrait.LSBRELEASE: _load_lsbrelease_version_target,
     ConditionTrait.UPTIME: load_timespan_target,
     ConditionTrait.INSTANCE_TYPE: _load_instance_type_target,
 }
@@ -116,6 +110,14 @@ def to_dict(self) -> dict:
             "target": self.stringify_target(),
         }
 
+    def matches(self, node: ClusterNodeMetadata) -> bool:
+        """Check if condition is met for a node
+
+        :param ClusterNodeMetadata node: node metadata
+        :return: true if it meets the condition
+        """
+        return self.operator.apply(self.trait.get_from(node), self.target)
+
     def __str__(self) -> str:
         return f"{self.trait.name} {self.operator.value} {self.stringify_target()}"
 

clusterman/migration/event_enums.py

@@ -15,6 +15,16 @@
 import operator
 from typing import Any
 from typing import Collection
+from typing import Union
+
+import packaging.version
+import semver
+
+from clusterman.interfaces.types import ClusterNodeMetadata
+
+
+ComparableVersion = Union[semver.VersionInfo, packaging.version.Version]
+ComparableConditionTarget = Union[str, int, ComparableVersion]
 
 
 class MigrationStatus(enum.Enum):
@@ -30,6 +40,14 @@ class ConditionTrait(enum.Enum):
     INSTANCE_TYPE = "instance_type"
     UPTIME = "uptime"
 
+    def get_from(self, node: ClusterNodeMetadata) -> ComparableConditionTarget:
+        """Get trait value from node metadata
+
+        :param ClusterNodeMetadata node: node metadata
+        :return: value
+        """
+        return CONDITION_TRAIT_GETTERS[self](node)
+
 
 class ConditionOperator(enum.Enum):
     GT = "gt"
@@ -71,3 +89,10 @@ def apply(self, left: Any, right: Any) -> bool:
     },
     ConditionTrait.UPTIME: {ConditionOperator.GT, ConditionOperator.GE, ConditionOperator.LT, ConditionOperator.LE},
 }
+
+CONDITION_TRAIT_GETTERS = {
+    ConditionTrait.KERNEL: lambda node: semver.VersionInfo.parse(node.agent.kernel),
+    ConditionTrait.LSBRELEASE: lambda node: packaging.version.parse(node.agent.lsbrelease),
+    ConditionTrait.INSTANCE_TYPE: lambda node: node.instance.market.instance,
+    ConditionTrait.UPTIME: lambda node: node.instance.uptime.total_seconds(),
+}

clusterman/migration/settings.py

@@ -23,7 +23,7 @@
 DEFAULT_POOL_PRESCALING = 0
 DEFAULT_NODE_BOOT_WAIT = "3m"
 DEFAULT_NODE_BOOT_TIMEOUT = "10m"
-DEFAULT_WORKER_TIMEOUT = "1d"
+DEFAULT_WORKER_TIMEOUT = "2h"
 
 
 class MigrationPrecendence(enum.Enum):
@@ -66,6 +66,9 @@ def __eq__(self, other: object) -> bool:
             raise NotImplementedError()
         return self.init_value == other.init_value
 
+    def __bool__(self) -> bool:
+        return self.value > 0
+
 
 class WorkerSetup(NamedTuple):
     rate: PoolPortion