Autoscaling!

Author: sylvinus

mrq/agent.py

@@ -1,13 +1,19 @@
-from future.builtins import object
-
 from .context import get_current_config, connections, log
 import time
 import datetime
 import gevent
+import argparse
+import shlex
+import random
+import math
+import re
+from collections import defaultdict
 from bson import ObjectId
 from redis.lock import LuaLock
-from collections import defaultdict
 from .processes import Process, ProcessPool
+from .config import add_parser_args
+from .utils import MovingETA
+from .queue import Queue
 
 
 class Agent(Process):
@@ -17,7 +23,10 @@ def __init__(self, worker_group=None):
         self.greenlets = {}
         self.id = ObjectId()
         self.worker_group = worker_group or get_current_config()["worker_group"]
-        self.pool = ProcessPool(extra_env={"MRQ_AGENT_ID": str(self.id)})
+        self.pool = ProcessPool(extra_env={
+            "MRQ_AGENT_ID": str(self.id),
+            "MRQ_WORKER_GROUP": self.worker_group
+        })
         self.config = get_current_config()
         self.status = "started"
 
@@ -32,6 +41,9 @@ def work(self):
         self.greenlets["manage"] = gevent.spawn(self.greenlet_manage)
         self.greenlets["manage"].start()
 
+        self.greenlets["queuestats"] = gevent.spawn(self.greenlet_queuestats)
+        self.greenlets["queuestats"].start()
+
         self.pool.start()
 
         try:
@@ -44,8 +56,8 @@ def work(self):
     def shutdown_now(self):
         self.pool.terminate()
 
-        self.greenlets["orchestrate"].kill()
-        self.greenlets["manage"].kill()
+        for g in self.greenlets.values():
+            g.kill()
 
     def shutdown_graceful(self):
         self.pool.stop(timeout=None)
@@ -97,15 +109,15 @@ def get_agent_report(self):
     def greenlet_orchestrate(self):
 
         while True:
-            with LuaLock(connections.redis, self.redis_agent_orchestrator_key,
+            with LuaLock(connections.redis, self.redis_orchestrator_lock_key,
                          timeout=self.config["orchestrate_interval"] + 10, thread_local=False, blocking=False):
                 self.orchestrate()
             time.sleep(self.config["orchestrate_interval"])
 
     @property
-    def redis_agent_orchestrator_key(self):
+    def redis_orchestrator_lock_key(self):
         """ Returns the global redis key used to ensure only one agent orchestrator runs at a time """
-        return "%s:agentorchestrator:%s" % (get_current_config()["redis_prefix"], self.worker_group)
+        return "%s:orchestratorlock:%s" % (get_current_config()["redis_prefix"], self.worker_group)
 
     def orchestrate(self):
         """ Executed periodically on one of the agents, to manage the desired workers of *all* the agents in its group """
@@ -119,7 +131,7 @@ def orchestrate(self):
 
         agents = self.fetch_worker_group_agents()
 
-        desired_workers = self.get_desired_workers_for_group(group)
+        desired_workers = self.get_desired_workers_for_group(group, agents)
 
         # Evaluate what workers are currently, rightfully there. They won't be touched.
         current_workers = defaultdict(int)
@@ -167,6 +179,8 @@ def orchestrate(self):
 
                 assert found
 
+        needs_new_agents = False
+
         # Add new workers to the least loaded machines
         for worker, delta in deltas.items():
             if delta <= 0:
@@ -187,9 +201,13 @@ def orchestrate(self):
 
                 if not found:
                     log.debug("Worker orchestration: no agent had enough CPU & memory (%s & %s) to schedule a new worker" % (cpu, memory))
-                    # TODO: communicate the need for new resources
+                    needs_new_agents = True
                     break
 
+        connections.mongodb_jobs.worker_groups.update_one({"_id": self.worker_group}, {"$set": {
+            "needs_new_agents": needs_new_agents
+        }})
+
         for agent in agents:
             if sorted(agent["new_desired_workers"]) != sorted(agent.get("desired_workers", [])):
                 # Commit the changes in DB
@@ -201,21 +219,185 @@ def orchestrate(self):
 
         log.debug("Orchestration finished.")
 
-    def get_desired_workers_for_group(self, group):
+    def get_desired_workers_for_group(self, group, agents):
 
         workers = {}
 
-        for profile in group.get("profiles", []):
+        def unpack(v):
+            s = v.split(" ")
+            return (int(s[0]), None if s[1] == "N" else float(s[1]), int(s[2]))
+
+        # count_jobs, eta, last_time
+        etas = {
+            k: unpack(v)
+            for k, v in connections.redis.hgetall(self.redis_queuestats_key).items()
+        }
+
+        # Compute average usage of each worker profile
+        worker_reports = self.fetch_worker_group_reports(projection=[
+            "_id", "config.worker_profile", "usage_avg", "status", "datestarted"  # , "process.cpu", "process.mem"
+        ])
+
+        worker_warmups_by_profile = defaultdict(int)
+        worker_usages_by_profile = defaultdict(list)
+        for rep in worker_reports:
+            profileid = rep["config"].get("worker_profile")
+            if not profileid or rep.get("status") not in ("wait", "spawn", "full"):
+                continue
+            # Don't take brand new workers into account yet.
+            age = (datetime.datetime.utcnow() - rep["datestarted"]).total_seconds()
+            if age < group.get("profiles", {}).get(profileid, {}).get("warmup", 60):
+                worker_warmups_by_profile[profileid] += 1
+                continue
+            worker_usages_by_profile[profileid].append(rep["usage_avg"])
+
+        worker_count_by_profile = defaultdict(int)
+        for agent in agents:
+            for command in agent.get("desired_workers", []):
+                profile = re.search(r"^MRQ_WORKER_PROFILE=([^\s]+)", command)
+                if profile:
+                    profile = profile.group(1)
+                    worker_count_by_profile[profile] += 1
+
+        # Compute the desired count for each profile
+        # This is the real "autoscaling" part.
+        for profileid, profile in group.get("profiles", {}).items():
+
+            cfg = self.get_config_for_profile(profile)
+            eta_queues = [q for q in cfg.queues if (q in etas and etas[q][1] is not None)]
+            worker_usage = sum(worker_usages_by_profile.get(profileid, []))
+            report_count = len(worker_usages_by_profile.get(profileid, []))
+            current_desired_count = worker_count_by_profile[profileid]
+            total_greenlets = (cfg.greenlets or 1) * (cfg.processes or 1)
+            warmups = worker_warmups_by_profile[profileid]
+
+            # By default, try not to change count
+            desired_count = current_desired_count
+
+            # If the queue is doing OK, are all instances necessary?
+            if warmups == 0 and report_count > 0:
+                desired_count = math.ceil(worker_usage * (float(current_desired_count) / report_count))
+
+            if len(eta_queues) > 0:
+
+                total_jobs = sum(etas[q][0] for q in eta_queues)
+                max_eta = max(etas[q][1] for q in eta_queues)
+
+                max_allowed_eta = profile.get("max_eta", 3600)
+
+                # Queue is taking too long, try to add an instance.
+                if max_eta > max_allowed_eta or any(etas[q][1] < 0 for q in eta_queues):
+
+                    # Don't add a worker if the absolute number of remaining jobs is less than the number of
+                    # greenlets. (They may become available right now)
+                    # Also don't add a worker if the reported, warmed-up worker count is not yet the desired one.
+                    if total_jobs > total_greenlets and report_count == current_desired_count:
+                        desired_count = current_desired_count + 1
+
+            final_count = min(max(desired_count, profile["min_count"]), profile["max_count"])
+
+            if final_count != current_desired_count:
+                log.debug("Autoscaling: Changing worker profile %s count from %s to %s" % (
+                    profileid, current_desired_count, final_count
+                ))
+
             workers[profile["command"]] = {
-                "desired_count": profile["min_count"],  # TODO!
+                "desired_count": final_count,
                 "memory": profile["memory"],
                 "cpu": profile["cpu"]
             }
 
         return workers
 
+    def get_config_for_profile(self, profile):
+        parser = argparse.ArgumentParser()
+        add_parser_args(parser, "worker")
+        parts = shlex.split(profile["command"])
+        if "mrq-worker" in parts:
+            parts = parts[parts.index("mrq-worker") + 1:]
+        return parser.parse_args(parts)
+
+    def greenlet_queuestats(self):
+
+        interval = min(self.config["orchestrate_interval"], 1 * 60)
+        lock_timeout = 5 * 60 + (interval * 2)
+
+        while True:
+            lock = LuaLock(connections.redis, self.redis_queuestats_lock_key,
+                           timeout=lock_timeout, thread_local=False, blocking=False)
+            with lock:
+                lock_expires = time.time() + lock_timeout
+                self.queue_etas = defaultdict(lambda: MovingETA(5))
+
+                while True:
+                    self.queuestats()
+
+                    # Because queue stats can be expensive, we try to keep the lock on the same agent
+                    lock_extend = (time.time() + lock_timeout) - lock_expires
+                    lock_expires += lock_extend
+                    lock.extend(lock_extend)
+
+                    time.sleep(interval)
+
+            time.sleep(interval)
+
+    @property
+    def redis_queuestats_lock_key(self):
+        """ Returns the global redis key used to ensure only one agent orchestrator runs at a time """
+        return "%s:queuestatslock" % (get_current_config()["redis_prefix"])
+
+    @property
+    def redis_queuestats_key(self):
+        """ Returns the global HSET redis key used to store queue stats """
+        return "%s:queuestats" % (get_current_config()["redis_prefix"])
+
+    def queuestats(self):
+        """ Compute ETAs for every known queue & subqueue """
+
+        start_time = time.time()
+        log.debug("Starting queue stats...")
+
+        # Fetch all known queues
+        queues = list(Queue.instanciate_queues(Queue.known_queues.keys()))
+
+        new_queues = {queue.id for queue in queues}
+        old_queues = set(self.queue_etas.keys())
+
+        for deleted_queue in old_queues.difference(new_queues):
+            self.queue_etas.pop(deleted_queue)
+
+        t = time.time()
+        stats = {}
+
+        for queue in queues:
+            cnt = queue.count_jobs_to_dequeue()
+            eta = self.queue_etas[queue.id].next(cnt, t=t)
+
+            # Number of jobs to dequeue, ETA, Time of stats
+            stats[queue.id] = "%d %s %d" % (cnt, eta if eta is not None else "N", int(t))
+
+        with connections.redis.pipeline(transaction=True) as pipe:
+            if random.randint(0, 100) == 0 or len(stats) == 0:
+                pipe.delete(self.redis_queuestats_key)
+            if len(stats) > 0:
+                pipe.hmset(self.redis_queuestats_key, stats)
+            pipe.execute()
+
+        log.debug("... done queue stats in %0.4fs" % (time.time() - start_time))
+
     def fetch_worker_group_agents(self):
         return list(connections.mongodb_jobs.mrq_agents.find({"worker_group": self.worker_group}))
 
+    def fetch_worker_group_reports(self, projection=None):
+        return list(connections.mongodb_jobs.mrq_workers.find({
+            "config.worker_group": self.worker_group
+        }, projection=projection))
+
     def fetch_worker_group_definition(self):
-        return connections.mongodb_jobs.mrq_workergroups.find_one({"_id": self.worker_group})
+        definition = connections.mongodb_jobs.mrq_workergroups.find_one({"_id": self.worker_group})
+
+        # Prepend all commands by their worker profile.
+        for profileid, profile in definition.get("profiles", {}).items():
+            profile["command"] = "MRQ_WORKER_PROFILE=%s %s" % (profileid, profile["command"])
+
+        return definition

mrq/config.py

@@ -358,6 +358,27 @@ def add_parser_args(parser, config_type):
             type=str,
             help='ID of the Agent this worker process is linked to')
 
+        parser.add_argument(
+            '--worker_id',
+            default="",
+            action='store',
+            type=str,
+            help='ID of this worker process. Should be left empty to be autogenerated in most cases.')
+
+        parser.add_argument(
+            '--worker_group',
+            default="",
+            action='store',
+            type=str,
+            help='Worker group of the agent this worker was launched from')
+
+        parser.add_argument(
+            '--worker_profile',
+            default="",
+            action='store',
+            type=str,
+            help='Worker profile used by the agent to launch this worker')
+
         parser.add_argument(
             'queues',
             nargs='*',

mrq/processes.py

@@ -87,7 +87,16 @@ def spawn(self, command):
         env["MRQ_IS_SUBPROCESS"] = "1"
         env.update(self.extra_env or {})
 
-        p = subprocess.Popen(shlex.split(command), shell=False, close_fds=True, env=env, cwd=os.getcwd())
+        # Extract env variables from shell commands.
+        parts = shlex.split(command)
+        for p in list(parts):
+            if "=" in p:
+                env[p.split("=")[0]] = p[len(p.split("=")[0]) + 1:]
+                parts.pop(0)
+            else:
+                break
+
+        p = subprocess.Popen(parts, shell=False, close_fds=True, env=env, cwd=os.getcwd())
 
         self.processes.append({
             "subprocess": p,