ucloud-scheduler/helper.py

from collections import Counter
from functools import reduce

from ucloud_common.vm import VmPool, VMStatus
from ucloud_common.host import HostPool, HostStatus
from ucloud_common.request import RequestEntry, RequestPool, RequestType

from decouple import config
from config import etcd_client as client

vm_pool = VmPool(client, config("VM_PREFIX"))
host_pool = HostPool(client, config("HOST_PREFIX"))
request_pool = RequestPool(client, config("REQUEST_PREFIX"))


def accumulated_specs(vms_specs):
    if not vms_specs:
        return {}
    return reduce((lambda x, y: Counter(x) + Counter(y)), vms_specs)


def remaining_resources(host_specs, vms_specs):
    """Return remaining resources host_specs - vms"""
    vms_specs = Counter(vms_specs)
    remaining = Counter(host_specs)
    remaining.subtract(vms_specs)

    return remaining


def get_suitable_host(vm_specs, hosts=None):
    if hosts is None:
        hosts = host_pool.by_status(HostStatus.alive)

    for host in hosts:
        # Filter them by host_name
        vms = vm_pool.by_host(host.key)

        # Filter them by status
        vms = vm_pool.except_status(VMStatus.requested_new, vms)

        running_vms_specs = [vm.specs for vm in vms]
        # Accumulate all of their combined specs
        running_vms_accumulated_specs = accumulated_specs(
            running_vms_specs
        )

        # Find out remaining resources after
        # host_specs - already running vm_specs
        remaining = remaining_resources(
            host.specs, running_vms_accumulated_specs
        )

        # Find out remaining - new_vm_specs
        remaining = remaining_resources(remaining, vm_specs)
        # if remaining resources >= 0 return this host_name
        if all(
                map(
                    lambda x: True if remaining[x] >= 0 else False,
                    remaining,
                )
        ):
            return host.key

    return None


def dead_host_detection():
    # Bring out your dead! - Monty Python and the Holy Grail
    hosts = host_pool.by_status(HostStatus.alive)
    dead_hosts_keys = []

    for host in hosts:
        # Only check those who claims to be alive
        if host.status == HostStatus.alive:
            if not host.is_alive():
                dead_hosts_keys.append(host.key)

    return dead_hosts_keys


def dead_host_mitigation(dead_hosts_keys):
    for host_key in dead_hosts_keys:
        host = host_pool.get(host_key)
        host.declare_dead()

        vms_hosted_on_dead_host = vm_pool.by_host(host_key)
        for vm in vms_hosted_on_dead_host:
            vm.declare_killed()
            vm_pool.put(vm)
        host_pool.put(host)


def assign_host(vm):
    host_name = get_suitable_host(vm.specs)
    if host_name:
        # if vm.status == VMStatus.requested_new:
        #     vm.status = VMStatus.scheduled_deploy
        #
        # if vm.status == VMStatus.killed:
        #     vm.status = VMStatus.requested_start

        vm.hostname = host_name
        vm_pool.put(vm)

        r = RequestEntry.from_scratch(type=RequestType.StartVM,
                                      uuid=vm.uuid,
                                      hostname=vm.hostname)
        request_pool.put(r)

        vm.log.append("VM scheduled for starting")

        return host_name
    return None