uncloud/ucloud/scheduler/helper.py

from collections import Counter
from functools import reduce

import bitmath

from ucloud.common.host import HostStatus
from ucloud.common.request import RequestEntry, RequestType
from ucloud.common.vm import VMStatus
from ucloud.shared import shared
from ucloud.settings import settings


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)

    for component in _vms_specs:
        if isinstance(_vms_specs[component], str):
            _vms_specs[component] = int(
                bitmath.parse_string_unsafe(
                    _vms_specs[component]
                ).to_MB()
            )
        elif isinstance(_vms_specs[component], list):
            _vms_specs[component] = map(
                lambda x: int(bitmath.parse_string_unsafe(x).to_MB()),
                _vms_specs[component],
            )
            _vms_specs[component] = reduce(
                lambda x, y: x + y, _vms_specs[component], 0
            )

    for component in _remaining:
        if isinstance(_remaining[component], str):
            _remaining[component] = int(
                bitmath.parse_string_unsafe(
                    _remaining[component]
                ).to_MB()
            )
        elif isinstance(_remaining[component], list):
            _remaining[component] = map(
                lambda x: int(bitmath.parse_string_unsafe(x).to_MB()),
                _remaining[component],
            )
            _remaining[component] = reduce(
                lambda x, y: x + y, _remaining[component], 0
            )

    _remaining.subtract(_vms_specs)

    return _remaining


class NoSuitableHostFound(Exception):
    """Exception when no host found that can host a VM."""


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

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

        # Filter them by status
        vms = shared.vm_pool.by_status(VMStatus.running, 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 all(map(lambda x: x >= 0, remaining.values())):
            return host.key

    raise NoSuitableHostFound


def dead_host_detection():
    # Bring out your dead! - Monty Python and the Holy Grail
    hosts = shared.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 = shared.host_pool.get(host_key)
        host.declare_dead()

        vms_hosted_on_dead_host = shared.vm_pool.by_host(host_key)
        for vm in vms_hosted_on_dead_host:
            vm.status = "UNKNOWN"
            shared.vm_pool.put(vm)
        shared.host_pool.put(host)


def assign_host(vm):
    vm.hostname = get_suitable_host(vm.specs)
    shared.vm_pool.put(vm)

    r = RequestEntry.from_scratch(
        type=RequestType.StartVM,
        uuid=vm.uuid,
        hostname=vm.hostname,
        request_prefix=settings["etcd"]["request_prefix"],
    )
    shared.request_pool.put(r)

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