ray-project--ray/python/ray/tests/test_autoscaling_policy.py

import collections
import copy
import logging
import shutil
import sys
import tempfile
import unittest
from queue import PriorityQueue
from typing import Callable, Dict, List

import pytest
import yaml

import ray
from ray._private.gcs_utils import PlacementGroupTableData
from ray.autoscaler._private.cli_logger import cli_logger
from ray.autoscaler._private.constants import AUTOSCALER_UPDATE_INTERVAL_S
from ray.autoscaler._private.load_metrics import LoadMetrics
from ray.autoscaler._private.node_launcher import NodeLauncher
from ray.autoscaler._private.providers import (
    _NODE_PROVIDERS,
    _clear_provider_cache,
)
from ray.autoscaler.tags import (
    NODE_KIND_HEAD,
    TAG_RAY_NODE_KIND,
    TAG_RAY_USER_NODE_TYPE,
)
from ray.core.generated.common_pb2 import Bundle, PlacementStrategy
from ray.tests.test_autoscaler import (
    MockAutoscaler,
    MockGcsClient,
    MockProcessRunner,
    MockProvider,
    mock_node_id,
)
from ray.tests.test_resource_demand_scheduler import MULTI_WORKER_CLUSTER


class Task:
    def __init__(
        self,
        duration: float,
        resources: Dict[str, float],
        start_callback: Callable[[None], None] = None,
        done_callback: Callable[[None], None] = None,
    ):
        self.duration = duration
        self.resources = resources
        self.start_callback = start_callback
        self.done_callback = done_callback
        self.start_time = None
        self.end_time = None
        self.node = None


class Actor(Task):
    pass


class PlacementGroup:
    def __init__(
        self,
        duration: float,
        bundles: List[Dict[str, float]],
        strategy: int,
        start_callback: Callable[[None], None] = None,
        done_callback: Callable[[None], None] = None,
    ):
        self.duration = duration
        self.bundles = bundles
        self.strategy = strategy
        self.start_callback = start_callback
        self.done_callback = done_callback
        self.start_time = None
        self.end_time = None
        self.node = None


class Node:
    def __init__(self, resources, in_cluster, node_type, start_time):
        self.total_resources = copy.deepcopy(resources)
        self.available_resources = copy.deepcopy(resources)
        self.in_cluster = in_cluster
        self.node_type = node_type
        self.start_time = start_time
        self.node_id = mock_node_id()

    def bundle_fits(self, bundle):
        if not self.in_cluster:
            return False
        for resource, quantity in bundle.items():
            if self.available_resources.get(resource, -1) < quantity:
                return False
        return True

    def feasible(self, bundle):
        if not self.in_cluster:
            return False
        for resource, quantity in bundle.items():
            if self.total_resources.get(resource, -1) < quantity:
                return False
        return True

    def allocate(self, bundle):
        assert self.bundle_fits(bundle) and self.in_cluster
        for resource, quantity in bundle.items():
            self.available_resources[resource] -= quantity

    def free(self, bundle):
        for resource, quantity in bundle.items():
            self.available_resources[resource] += quantity
        assert self.feasible(self.available_resources)


class Event:
    def __init__(self, time, event_type, data=None):
        self.time = time
        self.event_type = event_type
        self.data = data

    def __lt__(self, other):
        return self.time < other.time

    def __eq__(self, other):

        return self.time == other.time


SIMULATOR_EVENT_AUTOSCALER_UPDATE = 0
SIMULATOR_EVENT_TASK_DONE = 1
SIMULATOR_EVENT_NODE_JOINED = 2
SIMULATOR_EVENT_PG_DONE = 3


class Simulator:
    """This autoscaler simulator consists of a few components.

    State is stored in 3 main data structures:
        * Resource management state is stored in self.ip_to_nodes
        * The scheduler's work queue is stored in self.work_queue
        * An event queue which acts as the simulation's "timeline" in
          self.event_queue


    The logic is organized into 3 functions (and their helpers):
        * self.run_autoscaler plays the role of `monitor.py` and translates
          resource management state for load_metrics to consume.
        * self.schedule is the only consumer of the work queue. It dispatches
          work to the appropriate schedulers, which mutate cluster state and
          produce events for the event queue.
        * self.process_event is the sole consumer of the event queue. It
          dispatches work to the appropriate event handlers.

    There are 3 main ways of interacting with the simulator:
        * simulator.submit: To submit tasks
        * simulator.step: To go to the next "event"
        * task/actor/placement group start/done callbacks

    """

    def __init__(
        self,
        config_path,
        provider,
        autoscaler_update_interval_s=AUTOSCALER_UPDATE_INTERVAL_S,
        node_startup_delay_s=120,
    ):
        self.config_path = config_path
        self.provider = provider
        self.autoscaler_update_interval_s = autoscaler_update_interval_s
        self.node_startup_delay_s = node_startup_delay_s

        self._setup_autoscaler()
        self._setup_simulator()

    def _setup_autoscaler(self):
        self.runner = MockProcessRunner()
        self.config = yaml.safe_load(open(self.config_path).read())

        self.provider.create_node(
            {},
            {
                TAG_RAY_NODE_KIND: NODE_KIND_HEAD,
                TAG_RAY_USER_NODE_TYPE: self.config["head_node_type"],
            },
            1,
        )
        self.head_ip = self.provider.non_terminated_node_ips({})[0]

        self.load_metrics = LoadMetrics()
        self.autoscaler = MockAutoscaler(
            self.config_path,
            self.load_metrics,
            MockGcsClient(),
            # Don't let the autoscaler start any node launchers. Instead, we
            # will launch nodes ourself after every update call.
            max_concurrent_launches=0,
            max_failures=0,
            process_runner=self.runner,
            update_interval_s=0,
        )

        # Manually create a node launcher. Note that we won't start it as a
        # separate thread.
        self.node_launcher = NodeLauncher(
            provider=self.autoscaler.provider,
            pending=self.autoscaler.pending_launches,
            event_summarizer=self.autoscaler.event_summarizer,
            node_provider_availability_tracker=self.autoscaler.node_provider_availability_tracker,  # noqa: E501 Flake and black disagree how to format this.
            queue=self.autoscaler.launch_queue,
            index=0,
            node_types=self.autoscaler.available_node_types,
        )

    def _setup_simulator(self):
        self.virtual_time = 0
        self.ip_to_nodes = {}
        self._update_cluster_state(join_immediately=True)

        self.work_queue = []
        self.event_queue = PriorityQueue()
        self.event_queue.put(Event(0, SIMULATOR_EVENT_AUTOSCALER_UPDATE))

    def _update_cluster_state(self, join_immediately=False):
        nodes = self.provider.non_terminated_nodes(tag_filters={})
        for node_id in nodes:
            ip = self.provider.internal_ip(node_id)
            if ip in self.ip_to_nodes:
                continue
            node_tags = self.provider.node_tags(node_id)
            if TAG_RAY_USER_NODE_TYPE in node_tags:
                node_type = node_tags[TAG_RAY_USER_NODE_TYPE]
                resources = self.config["available_node_types"][node_type].get(
                    "resources", {}
                )
                node = Node(resources, join_immediately, node_type, self.virtual_time)
                self.ip_to_nodes[ip] = node
                if not join_immediately:
                    join_time = self.virtual_time + self.node_startup_delay_s
                    self.event_queue.put(
                        Event(join_time, SIMULATOR_EVENT_NODE_JOINED, node)
                    )

    def submit(self, work):
        if isinstance(work, list):
            self.work_queue.extend(work)
        else:
            self.work_queue.append(work)

    def _get_node_to_run(self, bundle, nodes):
        for ip, node in nodes.items():
            if node.bundle_fits(bundle):
                return ip, node
        return None, None

    def _schedule_placement_group(self, pg, nodes):
        # This scheduling algorithm is bad, but it is approximately as bad as
        # the real placement group scheduler.
        to_allocate = []
        if (
            pg.strategy == PlacementStrategy.STRICT_PACK
            or pg.strategy == PlacementStrategy.PACK
        ):
            combined = collections.defaultdict(float)
            for bundle in pg.bundles:
                for k, v in bundle.items():
                    combined[k] += v
            ip, node_to_run = self._get_node_to_run(combined, nodes)
            if node_to_run is None:
                return False
            to_allocate.append((combined, ip))
        elif (
            pg.strategy == PlacementStrategy.STRICT_SPREAD
            or pg.strategy == PlacementStrategy.SPREAD
        ):
            # TODO (Alex): More accurate handling of non-STRICT_PACK groups.
            remaining_nodes = nodes.copy()
            for bundle in pg.bundles:
                ip, node_to_run = self._get_node_to_run(bundle, remaining_nodes)
                if node_to_run is None:
                    return False
                del remaining_nodes[ip]
                to_allocate.append((bundle, ip))

        for bundle, ip in to_allocate:
            node = self.ip_to_nodes[ip]
            node.allocate(bundle)
        pg.start_time = self.virtual_time
        end_time = self.virtual_time + pg.duration
        self.event_queue.put(
            Event(end_time, SIMULATOR_EVENT_PG_DONE, (pg, to_allocate))
        )
        if pg.start_callback:
            pg.start_callback()
        return True

    def _schedule_task(self, task, nodes):
        ip, node = self._get_node_to_run(task.resources, nodes)
        if node is None:
            return False

        node.allocate(task.resources)
        task.node = node
        task.start_time = self.virtual_time
        end_time = self.virtual_time + task.duration
        self.event_queue.put(Event(end_time, SIMULATOR_EVENT_TASK_DONE, task))
        if task.start_callback:
            task.start_callback()
        return True

    def schedule(self):
        # TODO (Alex): Implement a more realistic scheduling algorithm.
        new_work_queue = []
        for work in self.work_queue:
            if isinstance(work, Task):
                scheduled = self._schedule_task(work, self.ip_to_nodes)
            elif isinstance(work, PlacementGroup):
                scheduled = self._schedule_placement_group(work, self.ip_to_nodes)
            else:
                assert False, "Unknown work object!"

            if scheduled is False:
                new_work_queue.append(work)
        self.work_queue = new_work_queue

    def _launch_nodes(self):
        """Launch all queued nodes. Since this will be run serially after
        `autoscaler.update` there are no race conditions in checking if the
        queue is empty.
        """
        while not self.node_launcher.queue.empty():
            config, count, node_type = self.node_launcher.queue.get()
            try:
                self.node_launcher._launch_node(config, count, node_type)
            except Exception:
                pass
            finally:
                self.node_launcher.pending.dec(node_type, count)

    def _infeasible(self, bundle):
        for node in self.ip_to_nodes.values():
            if node.feasible(bundle):
                return False
        return True

    def run_autoscaler(self):
        waiting_bundles = []
        infeasible_bundles = []
        placement_groups = []
        for work in self.work_queue:
            if isinstance(work, Task):
                shape = work.resources
                if self._infeasible(shape):
                    infeasible_bundles.append(shape)
                else:
                    waiting_bundles.append(shape)
            if isinstance(work, PlacementGroup):
                placement_groups.append(
                    PlacementGroupTableData(
                        state=PlacementGroupTableData.PENDING,
                        strategy=work.strategy,
                        bundles=[
                            Bundle(unit_resources=bundle) for bundle in work.bundles
                        ],
                    )
                )

        for ip, node in self.ip_to_nodes.items():
            if not node.in_cluster:
                continue
            self.load_metrics.update(
                ip=ip,
                node_id=node.node_id,
                static_resources=node.total_resources,
                dynamic_resources=node.available_resources,
                node_idle_duration_s=0,
                waiting_bundles=waiting_bundles,
                infeasible_bundles=infeasible_bundles,
                pending_placement_groups=placement_groups,
            )

        self.autoscaler.update()
        self._launch_nodes()
        self._update_cluster_state()

    def process_event(self, event):
        if event.event_type == SIMULATOR_EVENT_AUTOSCALER_UPDATE:
            self.run_autoscaler()
            next_update = self.virtual_time + self.autoscaler_update_interval_s
            self.event_queue.put(Event(next_update, SIMULATOR_EVENT_AUTOSCALER_UPDATE))
        elif event.event_type == SIMULATOR_EVENT_TASK_DONE:
            task = event.data
            task.node.free(task.resources)
            if task.done_callback:
                task.done_callback()
        elif event.event_type == SIMULATOR_EVENT_NODE_JOINED:
            node = event.data
            node.in_cluster = True
        elif event.event_type == SIMULATOR_EVENT_PG_DONE:
            pg, allocated = event.data
            for bundle, ip in allocated:
                self.ip_to_nodes[ip].free(bundle)
            if pg.done_callback:
                pg.done_callback()
        else:
            assert False, "Unknown event!"

    def step(self):
        self.virtual_time = self.event_queue.queue[0].time
        while self.event_queue.queue[0].time == self.virtual_time:
            event = self.event_queue.get()
            self.process_event(event)
        self.schedule()
        print(self.info_string())
        return self.virtual_time

    def node_costs(self):
        """Returns the cost of nodes. Cost is measured in terms of cumulative hours of
        runtime per node type.

        """
        costs = collections.defaultdict(float)

        for node in self.ip_to_nodes.values():
            if not node.in_cluster:
                continue
            runtime = self.virtual_time - node.start_time
            costs[node.node_type] += runtime
        return costs

    def info_string(self):
        num_connected_nodes = len(
            [node for node in self.ip_to_nodes.values() if node.in_cluster]
        )
        num_pending_nodes = len(self.ip_to_nodes) - num_connected_nodes
        return (
            f"[t={self.virtual_time}] "
            f"Connected: {num_connected_nodes}, "
            f"Pending: {num_pending_nodes}, "
            f"Remaining: {len(self.work_queue)}"
        )


SAMPLE_CLUSTER_CONFIG = copy.deepcopy(MULTI_WORKER_CLUSTER)
SAMPLE_CLUSTER_CONFIG["min_workers"] = 0
SAMPLE_CLUSTER_CONFIG["max_workers"] = 9999
SAMPLE_CLUSTER_CONFIG["target_utilization_fraction"] = 0.5
SAMPLE_CLUSTER_CONFIG["available_node_types"]["m4.16xlarge"]["max_workers"] = 100
SAMPLE_CLUSTER_CONFIG["available_node_types"]["m4.4xlarge"]["max_workers"] = 10000


class AutoscalingPolicyTest(unittest.TestCase):
    def setUp(self):
        _NODE_PROVIDERS["mock"] = lambda config: self.create_provider
        self.provider = None
        self.tmpdir = tempfile.mkdtemp()
        logging.disable(level=logging.CRITICAL)

        # This seems to be the only way of turning the cli logger off. The
        # expected methods like `cli_logger.configure` don't work.
        def do_nothing(*args, **kwargs):
            pass

        cli_logger._print = type(cli_logger._print)(do_nothing, type(cli_logger))

    def tearDown(self):
        self.provider = None
        del _NODE_PROVIDERS["mock"]
        _clear_provider_cache()
        shutil.rmtree(self.tmpdir)
        ray.shutdown()

    def create_provider(self, config, cluster_name):
        assert self.provider
        return self.provider

    def write_config(self, config):
        path = self.tmpdir + "/simple.yaml"
        with open(path, "w") as f:
            f.write(yaml.dump(config))
        return path

    def testManyTasks(self):
        config = copy.deepcopy(SAMPLE_CLUSTER_CONFIG)
        config_path = self.write_config(config)
        self.provider = MockProvider()
        simulator = Simulator(config_path, self.provider)

        done_count = 0

        def done_callback():
            nonlocal done_count
            done_count += 1

        tasks = [
            Task(duration=200, resources={"CPU": 1}, done_callback=done_callback)
            for _ in range(5000)
        ]
        simulator.submit(tasks)

        time = 0
        while done_count < len(tasks):
            time = simulator.step()

        assert time < 850
        # TODO (Alex): Not clear what's actually worth asserting here.
        assert simulator.node_costs()

        # Check event logs contain add/remove node events.
        assert any(
            "Adding" in x for x in simulator.autoscaler.event_summarizer.summary()
        )
        assert any(
            "Removing" in x for x in simulator.autoscaler.event_summarizer.summary()
        )

    def testManyActors(self):
        config = copy.deepcopy(SAMPLE_CLUSTER_CONFIG)
        config_path = self.write_config(config)
        self.provider = MockProvider()
        simulator = Simulator(config_path, self.provider)

        start_count = 0

        def start_callback():
            nonlocal start_count
            start_count += 1

        tasks = [
            Actor(
                duration=float("inf"),
                resources={"CPU": 1},
                start_callback=start_callback,
            )
            for _ in range(5000)
        ]
        simulator.submit(tasks)

        time = 0
        while start_count < len(tasks):
            time = simulator.step()

        assert time < 650

        # Check event logs contain add/remove node events.
        assert any(
            "Adding" in x for x in simulator.autoscaler.event_summarizer.summary()
        )
        assert any(
            "Removing" in x for x in simulator.autoscaler.event_summarizer.summary()
        )

    def testManyPlacementGroups(self):
        config = copy.deepcopy(SAMPLE_CLUSTER_CONFIG)
        config_path = self.write_config(config)
        self.provider = MockProvider()
        simulator = Simulator(config_path, self.provider)

        start_count = 0

        def start_callback():
            nonlocal start_count
            start_count += 1

        placement_group_requests = []

        for _ in range(500):
            placement_group_requests.append(
                PlacementGroup(
                    duration=float("inf"),
                    bundles=[{"CPU": 1}, {"CPU": 2}],
                    strategy=PlacementStrategy.STRICT_PACK,
                    start_callback=start_callback,
                )
            )

        for _ in range(500):
            placement_group_requests.append(
                PlacementGroup(
                    duration=float("inf"),
                    bundles=[{"CPU": 1}, {"CPU": 2}],
                    strategy=PlacementStrategy.STRICT_SPREAD,
                    start_callback=start_callback,
                )
            )

        # SPREAD and PACK tests fail, but under the real GCS placement group
        # scheduling algorithm we could also be left in a situation in which
        # the autoscaler thinks the placement group is placeable, but the
        # placement group scheduler doesn't know how to schedule it.

        # for _ in range(500):
        #     placement_group_requests.append(PlacementGroup(
        #         duration=float("inf"), bundles=[{"CPU": 1}, {"CPU": 2}],
        #         strategy=PlacementStrategy.PACK,
        #         start_callback=start_callback))

        # for _ in range(500):
        #     placement_group_requests.append(PlacementGroup(
        #         duration=float("inf"),
        #         bundles=[{"CPU": 2}, {"CPU": 1}],
        #         strategy=PlacementStrategy.SPREAD,
        #         start_callback=start_callback))

        simulator.submit(placement_group_requests)

        time = 0
        while start_count < len(placement_group_requests):
            time = simulator.step()

        assert time < 630

        # Check event logs contain add/remove node events.
        assert any(
            "Adding" in x for x in simulator.autoscaler.event_summarizer.summary()
        )
        assert any(
            "Removing" in x for x in simulator.autoscaler.event_summarizer.summary()
        )


if __name__ == "__main__":
    sys.exit(pytest.main(["-sv", __file__]))