Add Agent, ProcessPool, and our own Supervisor

Author: sylvinus

Dockerfile

@@ -55,6 +55,9 @@ RUN mkdir -p /data/db
 
 RUN ln -s /app/mrq/bin/mrq_run.py /usr/bin/mrq-run
 RUN ln -s /app/mrq/bin/mrq_worker.py /usr/bin/mrq-worker
+RUN ln -s /app/mrq/bin/mrq_agent.py /usr/bin/mrq-agent
+RUN ln -s /app/mrq/dashboard/app.py /usr/bin/mrq-dashboard
+
 ENV PYTHONPATH /app
 
 VOLUME ["/data"]

mrq/agent.py

@@ -0,0 +1,201 @@
+from future.builtins import object
+
+from .context import get_current_config, connections, log
+import time
+import json
+import gevent
+from bson import ObjectId
+from collections import defaultdict
+from .processes import Process, ProcessPool
+
+
+class Agent(Process):
+    """ MRQ Agent manages its local worker pool and takes turns in orchestrating the others in its group. """
+
+    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()
+
+    def work(self):
+
+        self.install_signal_handlers()
+
+        self.greenlets["orchestrate"] = gevent.spawn(self.greenlet_orchestrate)
+        self.greenlets["orchestrate"].start()
+
+        self.greenlets["manage"] = gevent.spawn(self.greenlet_manage)
+        self.greenlets["manage"].start()
+
+        self.pool.start()
+
+        self.pool.wait()
+
+    def shutdown_now(self):
+        self.pool.terminate()
+
+        self.greenlets["orchestrate"].kill()
+        self.greenlets["manage"].kill()
+
+    def shutdown_graceful(self):
+        self.pool.stop(timeout=None)
+
+    def greenlet_manage(self):
+        """ This greenlet always runs in background to update current status
+            in MongoDB every N seconds.
+         """
+
+        while True:
+            try:
+                self.manage()
+            except Exception as e:  # pylint: disable=broad-except
+                self.log.error("When reporting: %s" % e)
+            finally:
+                time.sleep(self.config["report_interval"])
+
+    def manage(self):
+
+        report = self.get_agent_report()
+
+        try:
+            db = self.mongodb_jobs.mrq_agents.find_and_modify({
+                "_id": ObjectId(self.id)
+            }, {"$set": report}, upsert=True)
+        except Exception as e:  # pylint: disable=broad-except
+            self.log.debug("Agent report failed: %s" % e)
+            return
+
+        # If the desired_workers was changed by an orchestrator, apply the changes locally
+        if sorted(db.get("desired_workers", [])) != sorted(self.pool.desired_commands):
+            self.pool.set_commands(db.get("desired_workers", []))
+
+    def get_agent_report(self):
+        report = {
+            "current_workers": [p["command"] for p in self.pool.processes],
+            "available_cpu": get_current_config()["available_cpu"],
+            "available_memory": get_current_config()["available_memory"]
+        }
+        return report
+
+    def greenlet_orchestrate(self):
+
+        while True:
+            with connections.redis.lock(self.redis_agent_orchestrator_key, timeout=60):
+                self.orchestrate()
+                time.sleep(30)
+
+    @property
+    def redis_agent_orchestrator_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)
+
+    def orchestrate(self):
+        """ Executed periodically on one of the agents, to manage the desired workers of *all* the agents in its group """
+
+        group = self.fetch_worker_group_definition()
+        if not group:
+            log.error("Worker group %s has no definition yet. Can't orchestrate!" % self.worker_group)
+            return
+
+        agents = self.fetch_worker_group_agents()
+
+        desired_workers = self.get_desired_workers_for_group(group)
+
+        # Evaluate what workers are currently, rightfully there. They won't be touched.
+        current_workers = defaultdict(int)
+        for agent in agents:
+            agent["free_memory"] = agent["available_memory"]
+            agent["free_cpu"] = agent["available_cpu"]
+            agent["new_desired_workers"] = []
+            for worker in agent.get("desired_workers", []):
+                if worker in desired_workers:
+                    cpu = desired_workers[worker]["cpu"]
+                    memory = desired_workers[worker]["memory"]
+
+                    # If no more memory for currently existing workers: their requirements must have changed.
+                    # We need to schedule it somewhere else
+                    if cpu <= agent["free_cpu"] and memory <= agent["free_memory"]:
+                        current_workers[worker] += 1
+                        agent["free_cpu"] -= cpu
+                        agent["free_memory"] -= memory
+                        agent["new_desired_workers"].append(worker)
+
+        # What changes need to be made in worker count
+        deltas = {
+            worker: (worker_info["desired_count"] - current_workers[worker])
+            for worker, worker_info in desired_workers.items()
+            if worker_info["desired_count"] != current_workers[worker]
+        }
+
+        # Remove workers from the most loaded machines (TODO improve)
+        for worker, delta in deltas.items():
+            if delta >= 0:
+                continue
+
+            for _ in range(delta, 0):
+                found = False
+                for agent in sorted(agents, key=lambda a: float(a["free_cpu"]) / a["available_cpu"]):
+                    for i in range(len(agent["new_desired_workers"])):
+                        if agent["new_desired_workers"][i] == worker:
+                            agent["new_desired_workers"].pop(i)
+                            agent["free_cpu"] += desired_workers[worker]["cpu"]
+                            agent["free_memory"] += desired_workers[worker]["memory"]
+                            found = True
+                            break
+                    if found:
+                        break
+
+                assert found
+
+        # Add new workers to the least loaded machines
+        for worker, delta in deltas.items():
+            if delta <= 0:
+                continue
+
+            cpu = desired_workers[worker]["cpu"]
+            memory = desired_workers[worker]["memory"]
+
+            for _ in range(delta):
+                found = False
+                for agent in sorted(agents, key=lambda a: float(a["free_cpu"]) / a["available_cpu"], reverse=True):
+                    if cpu <= agent["free_cpu"] and memory <= agent["free_memory"]:
+                        agent["new_desired_workers"].append(worker)
+                        agent["free_cpu"] -= cpu
+                        agent["free_memory"] -= memory
+                        found = True
+                        break
+
+                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
+                    break
+
+        for agent in agents:
+            if sorted(agent["new_desired_workers"]) != sorted(agent.get("desired_workers", [])):
+                # Commit the changes in DB
+                connections.mongodb_jobs.mrq_agents.update_one({"_id": agent["_id"]}, {"$set": {
+                    "desired_workers": agent["new_desired_workers"],
+                    "free_cpu": agent["free_cpu"],
+                    "free_memory": agent["free_memory"]
+                }})
+
+    def get_desired_workers_for_group(self, group):
+
+        workers = {}
+
+        for profile in group.get("profiles", []):
+            workers[profile["command"]] = {
+                "desired_count": profile["min_count"],  # TODO!
+                "memory": profile["memory"],
+                "cpu": profile["cpu"]
+            }
+
+        return workers
+
+    def fetch_worker_group_agents(self):
+        return list(connections.mongodb_jobs.mrq_agents.find({"worker_group": self.worker_group}))
+
+    def fetch_worker_group_definition(self):
+        return connections.mongodb_jobs.mrq_workergroups.find_one({"_id": self.worker_group})
+

mrq/bin/mrq_agent.py

@@ -0,0 +1,37 @@
+#!/usr/bin/env python
+import os
+
+# Needed to make getaddrinfo() work in pymongo on Mac OS X
+# Docs mention it's a better choice for Linux as well.
+# This must be done asap in the worker
+if "GEVENT_RESOLVER" not in os.environ:
+    os.environ["GEVENT_RESOLVER"] = "ares"
+
+from gevent import monkey
+monkey.patch_all(subprocess=False)
+
+import sys
+import argparse
+
+from .config import get_config
+from .agent import Agent
+from .context import set_current_config
+
+
+def main():
+
+    parser = argparse.ArgumentParser(description='Start a MRQ agent')
+
+    cfg = get_config(parser=parser, config_type="agent", sources=("file", "env", "args"))
+
+    set_current_config(cfg)
+
+    agent = Agent()
+
+    agent.work()
+
+    sys.exit(agent.exitcode)
+
+
+if __name__ == "__main__":
+    main()

mrq/bin/mrq_worker.py

@@ -16,6 +16,7 @@
 import signal
 import psutil
 import argparse
+import pipes
 
 try:
     import subprocess32 as subprocess
@@ -31,88 +32,29 @@
 
 def main():
 
-    parser = argparse.ArgumentParser(description='Start a RQ worker')
+    parser = argparse.ArgumentParser(description='Start a MRQ worker')
 
     cfg = config.get_config(parser=parser, config_type="worker", sources=("file", "env", "args"))
 
-    # If we are launching with a --processes option and without the SUPERVISOR_ENABLED env
-    # then we should just call supervisord.
-    if cfg["processes"] > 0 and not os.environ.get("SUPERVISOR_ENABLED"):
+    set_current_config(cfg)
 
-        # We wouldn't need to do all that if supervisord supported environment
-        # variables in all its config fields!
-        with open(cfg["supervisord_template"], "r") as f:
-            conf = f.read()
+    # If we are launching with a --processes option and without MRQ_IS_SUBPROCESS, we are a manager process
+    if cfg["processes"] > 0 and not os.environ.get("MRQ_IS_SUBPROCESS"):
 
-        fh, path = tempfile.mkstemp(prefix="mrqsupervisordconfig")
-        f = os.fdopen(fh, "w")
+        from mrq.supervisor import Supervisor
 
-        # We basically relaunch ourselves, but the config will contain the
-        # MRQ_SUPERVISORD_ISWORKER env.
-        conf = conf.replace("{{ SUPERVISORD_COMMAND }}", " ".join(sys.argv))
-        conf = conf.replace(
-            "{{ SUPERVISORD_PROCESSES }}", str(cfg["processes"]))
+        command = " ".join(map(pipes.quote, sys.argv))
+        w = Supervisor(command, numprocs=cfg["processes"])
+        w.work()
+        sys.exit(w.exitcode)
 
-        f.write(conf)
-        f.close()
-
-        try:
-
-            # start_new_session=True avoids sending the current process'
-            # signals to the child.
-            process = subprocess.Popen(
-                ["supervisord", "-c", path], start_new_session=True)
-
-            def sigint_handler(signum, frame):  # pylint: disable=unused-argument
-
-                # At this point we need to send SIGINT to all workers. Unfortunately supervisord
-                # doesn't support this, so we have to find all the children pids and send them the
-                # signal ourselves :-/
-                # https://github.com/Supervisor/supervisor/issues/179
-                #
-                psutil_process = psutil.Process(process.pid)
-                worker_processes = psutil_process.children(recursive=False)
-
-                if len(worker_processes) == 0:
-                    return process.send_signal(signal.SIGTERM)
-
-                for child_process in worker_processes:
-                    child_process.send_signal(signal.SIGINT)
-
-                # Second time sigint is used, we should terminate supervisord itself which
-                # will send SIGTERM to all the processes anyway.
-                signal.signal(signal.SIGINT, sigterm_handler)
-
-                # Wait for all the childs to finish
-                for child_process in worker_processes:
-                    child_process.wait()
-
-                # Then stop supervisord itself.
-                process.send_signal(signal.SIGTERM)
-
-            def sigterm_handler(signum, frame):  # pylint: disable=unused-argument
-                process.send_signal(signal.SIGTERM)
-
-            signal.signal(signal.SIGINT, sigint_handler)
-            signal.signal(signal.SIGTERM, sigterm_handler)
-
-            process.wait()
-
-        finally:
-            os.remove(path)
-
-    # If not, start the actual worker
+    # If not, start an actual worker
     else:
 
         worker_class = load_class_by_path(cfg["worker_class"])
-
-        set_current_config(cfg)
-
         w = worker_class()
-
-        exitcode = w.work()
-
-        sys.exit(exitcode)
+        w.work()
+        sys.exit(w.exitcode)
 
 if __name__ == "__main__":
     main()

mrq/config.py

@@ -4,6 +4,7 @@
 import os
 import sys
 import re
+import psutil
 from .version import VERSION
 from .utils import get_local_ip, DelimiterArgParser
 import atexit
@@ -255,8 +256,31 @@ def add_parser_args(parser, config_type):
             type=str,
             help='Bind the dashboard to this IP. Default is "0.0.0.0", use "127.0.0.1" to restrict access.')
 
-    # Worker-specific args
+    # Agent-specific args
+    elif config_type == "agent":
+
+        parser.add_argument(
+            '--worker_group',
+            default="default",
+            action="store",
+            type=str,
+            help='The name of the worker group to manage')
+
+        parser.add_argument(
+            '--available_memory',
+            default=int(psutil.virtual_memory().total * 0.8),
+            action="store",
+            type=int,
+            help="How much memory MB this agent's workers can use. Used for scheduling, not a hard limit.")
 
+        parser.add_argument(
+            '--available_cpu',
+            default=psutil.cpu_count(logical=True) * 1024,
+            action="store",
+            type=int,
+            help="How much CPU units this agent's workers can use. We recommend using 1024 per CPU.")
+
+    # Worker-specific args
     elif config_type == "worker":
 
         parser.add_argument(