Merge pull request #204 from pricingassistant/refactorworkergroup

Refactorworkergroup

Author: maelorn

mrq/agent.py

@@ -1,15 +1,16 @@
-from .context import get_current_config, connections, log, run_task
+from .context import get_current_config, connections, log, run_task, metric
 import time
 import datetime
 import gevent
 import argparse
 import random
+import shlex
 import traceback
 from collections import defaultdict
 from bson import ObjectId
 from redis.lock import LuaLock
 from .processes import Process, ProcessPool
-from .utils import MovingETA
+from .utils import MovingETA, normalize_command
 from .queue import Queue
 
 
@@ -26,7 +27,7 @@ def __init__(self, worker_group=None):
         })
         self.config = get_current_config()
         self.status = "started"
-
+        metric("agent", data={"worker_group": self.worker_group, "agent_id": self.id})
         self.dateorchestrated = None
 
         # global redis key used to ensure only one agent orchestrator runs at a time
@@ -115,6 +116,7 @@ def get_agent_report(self):
             "datereported": datetime.datetime.utcnow(),
             "dateexpires": datetime.datetime.utcnow() + datetime.timedelta(seconds=(self.config["report_interval"] * 3) + 5)
         }
+        metric("agent", data={"worker_group": self.worker_group, "agent_id": self.id, "worker_count": len(self.pool.processes)})
         return report
 
     def greenlet_orchestrate(self):
@@ -192,7 +194,10 @@ def fetch_worker_group_definition(self):
         definition = connections.mongodb_jobs.mrq_workergroups.find_one({"_id": self.worker_group})
 
         # Prepend all commands by their worker profile.
-        for profileid, profile in (definition or {}).get("profiles", {}).items():
-            profile["command"] = "MRQ_WORKER_PROFILE=%s %s" % (profileid, profile["command"])
+        commands = []
+        for command in definition.get("commands", []):
+            simplified_command, worker_count = normalize_command(command, self.worker_group)
+            commands.extend([simplified_command] * worker_count)
 
+        definition["commands"] = commands
         return definition

mrq/basetasks/orchestrator.py

@@ -8,6 +8,7 @@
 import shlex
 import argparse
 from ..config import add_parser_args
+from ..utils import normalize_command
 import traceback
 import datetime
 import re
@@ -47,143 +48,16 @@ def do_orchestrate(self, group):
 
         agents = self.fetch_worker_group_agents(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)
         for agent in agents:
-            agent["free_memory"] = agent["total_memory"]
-            agent["free_cpu"] = agent["total_cpu"]
+            desired_workers = self.get_desired_workers_for_agent(group, agent)
             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["total_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
-
-        needs_new_agents = False
-
-        # 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["total_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))
-                    needs_new_agents = True
-                    break
-
-        # Worker diversity enforcement: if we are under-scaled, make sure that at least one worker
-        # of each profile is launched. if not, make space forcefully on other workers.
-        if needs_new_agents:
-
-            all_profiles = defaultdict(int)
-            for agent in agents:
-                for w in agent["new_desired_workers"]:
-                    all_profiles[w] += 1
-
-            removable = sorted([list(x) for x in all_profiles.items() if x[1] > 1], key=lambda item: item[1], reverse=True)
-
-            for profile in desired_workers:
-                if desired_workers[profile]["desired_count"] == 0:
-                    continue
-                if profile not in all_profiles:
-                    found = False
-                    # This profile is not currently represented in the workers. Make some space for it.
-                    for rem in removable:
-                        for agent in agents:
-                            for i in range(len(agent["new_desired_workers"])):
-                                if agent["new_desired_workers"][i] == rem[0] and rem[1] > 1:
-                                    agent["new_desired_workers"][i] = None
-                                    agent["free_cpu"] += desired_workers[rem[0]]["cpu"]
-                                    agent["free_memory"] += desired_workers[rem[0]]["memory"]
-                                    rem[1] -= 1
-                                if desired_workers[profile]["cpu"] <= agent["free_cpu"] and desired_workers[profile]["memory"] <= agent["free_memory"]:
-                                    log.debug("Orchestration: enforcing worker diversity with %s => %s" % (rem[0], profile))
-                                    agent["new_desired_workers"].append(profile)
-                                    agent["free_cpu"] -= desired_workers[profile]["cpu"]
-                                    agent["free_memory"] -= desired_workers[profile]["memory"]
-                                    found = True
-                                    break
-
-                            agent["new_desired_workers"] = [x for x in agent["new_desired_workers"] if x is not None]
-                            if found:
-                                break
-
-                        if found:
-                            break
-
-                    if not found:
-                        log.debug("Orchestration couldn't enforce worker diversity for profile '%s'" % profile)
-
-        # User-provided autoscaling task
-        if self.config.get("autoscaling_taskpath"):
-            result = run_task(self.config["autoscaling_taskpath"], {
-                "agents": agents,
-                "needs_new_agents": needs_new_agents,
-                "worker_group": group["_id"]
-            })
-            needs_new_agents = result["needs_new_agents"]
-            agents = result["agents"]
-
-        # Save the new values in the DB. They will be applied by each agent process.
-        connections.mongodb_jobs.worker_groups.update_one({"_id": group["_id"]}, {"$set": {
-            "needs_new_agents": needs_new_agents
-        }})
+            agent["new_desired_workers"] = desired_workers
 
         for agent in agents:
             if sorted(agent["new_desired_workers"]) != sorted(agent.get("desired_workers", [])):
                 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"]
+                    "desired_workers": agent["new_desired_workers"]
                 }})
 
         # Remember the date of the last successful orchestration (will be reported)
@@ -195,95 +69,8 @@ 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 get_desired_workers_for_group(self, group, agents):
-
-        workers = {}
-
-        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(group, 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.get("min_count", 0)), profile.get("max_count", 100))
-
-            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": final_count,
-                "memory": profile["memory"],
-                "cpu": profile["cpu"]
-            }
-
-        return workers
+    def get_desired_workers_for_agent(self, group, agent):
+        return group.get("commands", [])
 
     def fetch_worker_group_reports(self, worker_group, projection=None):
         return list(connections.mongodb_jobs.mrq_workers.find({
@@ -295,9 +82,12 @@ def fetch_worker_group_definitions(self):
         definitions = list(connections.mongodb_jobs.mrq_workergroups.find())
 
         for definition in definitions:
-            # Prepend all commands by their worker profile.
-            for profileid, profile in (definition or {}).get("profiles", {}).items():
-                profile["command"] = "MRQ_WORKER_PROFILE=%s %s" % (profileid, profile["command"])
+            commands = []
+            # Prepend all commands by their worker group.
+            for command in definition.get("commands", []):
+                simplified_command, worker_count = normalize_command(command, definition["_id"])
+                commands.extend([simplified_command] * worker_count)
+            definition["commands"] = commands
 
         return definitions
 

mrq/config.py

@@ -409,13 +409,6 @@ def add_parser_args(parser, config_type):
             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(
             '--task_whitelist',
             default="",

mrq/context.py

@@ -14,7 +14,6 @@
 from .config import get_config
 from .subpool import subpool_map, subpool_imap
 
-
 # This should be MRQ's only Python object shared by all the jobs in the same process
 _GLOBAL_CONTEXT = {
 

mrq/processes.py

@@ -59,7 +59,6 @@ def set_commands(self, commands, timeout=None):
         self.desired_commands = commands
 
         target_commands = list(self.desired_commands)
-
         for process in list(self.processes):
             found = False
             for i in range(len(target_commands)):

mrq/utils.py

@@ -10,12 +10,36 @@
 from collections import deque
 from bson import ObjectId
 import uuid
+import shlex
 
 #
 # Utils are functions that should be independent from the rest of MRQ's codebase
 #
 
 
+def normalize_command(command, worker_group):
+    if "--processes" in command:
+        simplified_command = ""
+        worker_count = 0
+        skip_next = False
+        for part in shlex.split(command):
+            if skip_next:
+                worker_count = part
+                skip_next = False
+                continue
+            if part.startswith("--processes="):
+                worker_count = part.split("=")[1]
+                continue
+            if part == "--processes":
+                skip_next = True
+                continue
+            simplified_command += " %s" % part
+            skip_next = False
+        simplified_command = "MRQ_WORKER_GROUP=%s%s" % (worker_group, simplified_command)
+        return simplified_command, int(worker_count)
+    return "MRQ_WORKER_GROUP=%s %s" % (worker_group, command), 1
+
+
 def get_local_ip():
     """ Returns the local IP. Can be overwritten in the config with --local-ip so don't call
     this function directly, instead get the current value from the config """

mrq/worker.py

@@ -292,8 +292,7 @@ def get_worker_report(self, with_memory=False):
             "local_ip",
             "external_ip",
             "agent_id",
-            "worker_group",
-            "worker_profile"
+            "worker_group"
         ]
 
         io = None