[Serve] Pluggable Queueing Policy (#6492)

python/ray/experimental/serve/queues.py

@@ -4,6 +4,7 @@ import numpy as np
 
 import ray
 from ray.experimental.serve.utils import logger
+import itertools
 from blist import sortedlist
 import time
 
@@ -148,8 +149,7 @@ class CentralizedQueues:
 
     def link(self, service, backend):
         logger.debug("Link %s with %s", service, backend)
-        self.traffic[service][backend] = 1.0
-        self.flush()
+        self.set_traffic(service, {backend: 1.0})
 
     def set_traffic(self, service, traffic_dict):
         logger.debug("Setting traffic for service %s to %s", service,
@@ -173,24 +173,8 @@ class CentralizedQueues:
         }
         return list(backends_in_policy.intersection(available_workers))
 
-    def _flush(self):
-        # perform traffic splitting for requests
-        for service, queue in self.queues.items():
-            # while there are incoming requests and there are backends
-            while len(queue) and len(self.traffic[service]):
-                backend_names = list(self.traffic[service].keys())
-                backend_weights = list(self.traffic[service].values())
-                # TODO(alind): is random choice good for deadline awareness?
-                # putting query in a buffer of a non available backend may
-                # not be good
-                chosen_backend = np.random.choice(
-                    backend_names, p=backend_weights).squeeze()
-
-                request = queue.popleft()
-                # maintain a sorted list in the buffer queue of the backend
-                self.buffer_queues[chosen_backend].add(request)
-
-        # distach buffer queues to work queues
+    # flushes the buffer queue and assigns work to workers
+    def _flush_buffer(self):
         for service in self.queues.keys():
             ready_backends = self._get_available_backends(service)
             for backend in ready_backends:
@@ -207,9 +191,33 @@ class CentralizedQueues:
                     )
                     work.replica_handle._ray_serve_call.remote(request)
 
+    # selects the backend and puts the service queue query to the buffer
+    # different policies will implement different backend selection policies
+    def _flush_service_queue(self):
+        """
+        Expected Implementation:
+            The implementer is expected to access and manipulate
+            self.queues        : dict[str,Deque]
+            self.buffer_queues : dict[str,sortedlist]
+        For registering the implemented policies register at policy.py!
+        Expected Behavior:
+            the Deque of all services in self.queues linked with
+            atleast one backend must be empty irrespective of whatever
+            backend policy is implemented.
+        """
+        pass
+
+    # _flush function has to flush the service and buffer queues.
+    def _flush(self):
+        self._flush_service_queue()
+        self._flush_buffer()
+
 
-@ray.remote
 class CentralizedQueuesActor(CentralizedQueues):
+    """
+    A wrapper class for converting wrapper policy classes to ray
+    actors. This is needed to make `flush` call asynchronous.
+    """
     self_handle = None
 
     def register_self_handle(self, handle_to_this_actor):
@@ -220,3 +228,167 @@ class CentralizedQueuesActor(CentralizedQueues):
             self.self_handle._flush.remote()
         else:
             self._flush()
+
+
+class RandomPolicyQueue(CentralizedQueues):
+    """
+    A wrapper class for Random policy.This backend selection policy is
+    `Stateless` meaning the current decisions of selecting backend are
+    not dependent on previous decisions. Random policy (randomly) samples
+    backends based on backend weights for every query. This policy uses the
+    weights assigned to backends.
+    """
+
+    def _flush_service_queue(self):
+        # perform traffic splitting for requests
+        for service, queue in self.queues.items():
+            # while there are incoming requests and there are backends
+            while len(queue) and len(self.traffic[service]):
+                backend_names = list(self.traffic[service].keys())
+                backend_weights = list(self.traffic[service].values())
+                # randomly choose a backend for every query
+                chosen_backend = np.random.choice(
+                    backend_names, p=backend_weights).squeeze()
+
+                request = queue.popleft()
+                self.buffer_queues[chosen_backend].add(request)
+
+
+@ray.remote
+class RandomPolicyQueueActor(RandomPolicyQueue, CentralizedQueuesActor):
+    pass
+
+
+class RoundRobinPolicyQueue(CentralizedQueues):
+    """
+    A wrapper class for RoundRobin policy. This backend selection policy
+    is `Stateful` meaning the current decisions of selecting backend are
+    dependent on previous decisions. RoundRobinPolicy assigns queries in
+    an interleaved manner to every backend serving for a service. Consider
+    backend A,B linked to a service. Now queries will be assigned to backends
+    in the following order - [ A, B, A, B ... ] . This policy doesn't use the
+    weights assigned to backends.
+    """
+
+    # Saves the information about last assigned
+    # backend for every service
+    round_robin_iterator_map = {}
+
+    def set_traffic(self, service, traffic_dict):
+        logger.debug("Setting traffic for service %s to %s", service,
+                     traffic_dict)
+        self.traffic[service] = traffic_dict
+        backend_names = list(self.traffic[service].keys())
+        self.round_robin_iterator_map[service] = itertools.cycle(backend_names)
+        self.flush()
+
+    def _flush_service_queue(self):
+        # perform traffic splitting for requests
+        for service, queue in self.queues.items():
+            # if there are incoming requests and there are backends
+            if len(queue) and len(self.traffic[service]):
+                while len(queue):
+                    # choose the next backend available from persistent
+                    # information
+                    chosen_backend = next(
+                        self.round_robin_iterator_map[service])
+                    request = queue.popleft()
+                    self.buffer_queues[chosen_backend].add(request)
+
+
+@ray.remote
+class RoundRobinPolicyQueueActor(RoundRobinPolicyQueue,
+                                 CentralizedQueuesActor):
+    pass
+
+
+class PowerOfTwoPolicyQueue(CentralizedQueues):
+    """
+    A wrapper class for powerOfTwo policy. This backend selection policy is
+    `Stateless` meaning the current decisions of selecting backend are
+    dependent on previous decisions. PowerOfTwo policy (randomly) samples two
+    backends (say Backend A,B among A,B,C) based on the backend weights
+    specified and chooses the backend which is less loaded. This policy uses
+    the weights assigned to backends.
+    """
+
+    def _flush_service_queue(self):
+        # perform traffic splitting for requests
+        for service, queue in self.queues.items():
+            # while there are incoming requests and there are backends
+            while len(queue) and len(self.traffic[service]):
+                backend_names = list(self.traffic[service].keys())
+                backend_weights = list(self.traffic[service].values())
+                if len(self.traffic[service]) >= 2:
+                    # randomly pick 2 backends
+                    backend1, backend2 = np.random.choice(
+                        backend_names, 2, p=backend_weights)
+
+                    # see the length of buffer queues of the two backends
+                    # and pick the one which has less no. of queries
+                    # in the buffer
+                    if (len(self.buffer_queues[backend1]) <= len(
+                            self.buffer_queues[backend2])):
+                        chosen_backend = backend1
+                    else:
+                        chosen_backend = backend2
+                else:
+                    chosen_backend = np.random.choice(
+                        backend_names, p=backend_weights).squeeze()
+                request = queue.popleft()
+                self.buffer_queues[chosen_backend].add(request)
+
+
+@ray.remote
+class PowerOfTwoPolicyQueueActor(PowerOfTwoPolicyQueue,
+                                 CentralizedQueuesActor):
+    pass
+
+
+class FixedPackingPolicyQueue(CentralizedQueues):
+    """
+    A wrapper class for FixedPacking policy. This backend selection policy is
+    `Stateful` meaning the current decisions of selecting backend are dependent
+    on previous decisions. FixedPackingPolicy is k RoundRobin policy where
+    first packing_num queries are handled by 'backend-1' and next k queries are
+    handled by 'backend-2' and so on ... where 'backend-1' and 'backend-2' are
+    served by the same service. This policy doesn't use the weights assigned to
+    backends.
+
+    """
+
+    def __init__(self, packing_num=3):
+        # Saves the information about last assigned
+        # backend for every service
+        self.fixed_packing_iterator_map = {}
+        self.packing_num = packing_num
+        super().__init__()
+
+    def set_traffic(self, service, traffic_dict):
+        logger.debug("Setting traffic for service %s to %s", service,
+                     traffic_dict)
+        self.traffic[service] = traffic_dict
+        backend_names = list(self.traffic[service].keys())
+        self.fixed_packing_iterator_map[service] = itertools.cycle(
+            itertools.chain.from_iterable(
+                itertools.repeat(x, self.packing_num) for x in backend_names))
+        self.flush()
+
+    def _flush_service_queue(self):
+        # perform traffic splitting for requests
+        for service, queue in self.queues.items():
+            # if there are incoming requests and there are backends
+            if len(queue) and len(self.traffic[service]):
+                while len(queue):
+                    # choose the next backend available from persistent
+                    # information
+                    chosen_backend = next(
+                        self.fixed_packing_iterator_map[service])
+                    request = queue.popleft()
+                    self.buffer_queues[chosen_backend].add(request)
+
+
+@ray.remote
+class FixedPackingPolicyQueueActor(FixedPackingPolicyQueue,
+                                   CentralizedQueuesActor):
+    pass