diff --git a/mrq/agent.py b/mrq/agent.py index 35a6f311..3910c5cf 100644 --- a/mrq/agent.py +++ b/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 diff --git a/mrq/basetasks/orchestrator.py b/mrq/basetasks/orchestrator.py index 9d917be8..124502b9 100644 --- a/mrq/basetasks/orchestrator.py +++ b/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 diff --git a/mrq/config.py b/mrq/config.py index e3d443dc..e6dad305 100644 --- a/mrq/config.py +++ b/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="", diff --git a/mrq/context.py b/mrq/context.py index dc1342f6..0a77f64d 100644 --- a/mrq/context.py +++ b/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 = { diff --git a/mrq/processes.py b/mrq/processes.py index 22c857e2..6f901af6 100644 --- a/mrq/processes.py +++ b/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)): diff --git a/mrq/utils.py b/mrq/utils.py index f27dd009..72d84a60 100644 --- a/mrq/utils.py +++ b/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 """ diff --git a/mrq/worker.py b/mrq/worker.py index e54279d5..04d8a362 100644 --- a/mrq/worker.py +++ b/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 diff --git a/tests/test_agent.py b/tests/test_agent.py index d41c03d6..2760e79d 100644 --- a/tests/test_agent.py +++ b/tests/test_agent.py @@ -6,7 +6,7 @@ import psutil -def scenario(profiles, agents): +def scenario(commands, agents): connections.mongodb_jobs.mrq_agents.delete_many({}) connections.mongodb_jobs.mrq_workergroups.delete_many({}) @@ -15,7 +15,7 @@ def scenario(profiles, agents): agent["status"] = "started" connections.mongodb_jobs.mrq_agents.insert_many(agents + [{"worker_group": "yy", "status": "started"}, {"worker_group": "zz"}]) - connections.mongodb_jobs.mrq_workergroups.insert_one({"_id": "xx", "profiles": profiles}) + connections.mongodb_jobs.mrq_workergroups.insert_one({"_id": "xx", "commands": commands}) agent = Agent(worker_group="xx") agent.orchestrate() @@ -29,14 +29,8 @@ def test_orchestration_scenarios(worker): worker.start() # Simplest scenario - assert scenario({ - "a": { - "command": "mrq-worker a", - "memory": 1000, - "cpu": 1024, - "min_count": 1 - } - }, [ + assert scenario( + ["mrq-worker a"], [ { "_id": "worker1", "total_cpu": 1024, @@ -44,61 +38,14 @@ def test_orchestration_scenarios(worker): } ]) == { "worker1": [ - "MRQ_WORKER_PROFILE=a mrq-worker a" + "MRQ_WORKER_GROUP=xx mrq-worker a" ] } - # Not enough memory - assert scenario({ - "a": { - "command": "mrq-worker a", - "memory": 1001, - "cpu": 1024, - "min_count": 1 - } - }, [ - { - "_id": "worker1", - "total_cpu": 1024, - "total_memory": 1000 - } - ]) == { - "worker1": [] - } - - # Not enough CPU - assert scenario({ - "a": { - "command": "mrq-worker a", - "memory": 1000, - "cpu": 1025, - "min_count": 1 - } - }, [ - { - "_id": "worker1", - "total_cpu": 1024, - "total_memory": 1000 - } - ]) == { - "worker1": [] - } - # Remove & add workers - assert scenario({ - "a": { - "command": "mrq-worker a", - "memory": 1, - "cpu": 1, - "min_count": 2 - }, - "b": { - "command": "mrq-worker b", - "memory": 1, - "cpu": 1, - "min_count": 1 - } - }, [ + assert scenario( + ["mrq-worker --processes 2 a", "mrq-worker b"] + , [ { "_id": "worker1", "total_cpu": 3, @@ -107,74 +54,74 @@ def test_orchestration_scenarios(worker): } ]) == { "worker1": [ - "MRQ_WORKER_PROFILE=a mrq-worker a", - "MRQ_WORKER_PROFILE=a mrq-worker a", - "MRQ_WORKER_PROFILE=b mrq-worker b" + "MRQ_WORKER_GROUP=xx mrq-worker a", + "MRQ_WORKER_GROUP=xx mrq-worker a", + "MRQ_WORKER_GROUP=xx mrq-worker b" ] } # Worker removal & add priority - assert scenario({ - "a": { - "command": "mrq-worker a", - "memory": 1, - "cpu": 1, - "min_count": 3 - }, - "b": { - "command": "mrq-worker b", - "memory": 1, - "cpu": 1, - "min_count": 1 - } - }, [ - { - "_id": "worker1", - "total_cpu": 11, - "total_memory": 11, - "desired_workers": ["mrq-worker a", "mrq-worker a"] - }, { - "_id": "worker2", - "total_cpu": 5, - "total_memory": 5, - "desired_workers": ["mrq-worker a", "mrq-worker a"] - } - ]) == { - "worker1": [ - "MRQ_WORKER_PROFILE=a mrq-worker a", - "MRQ_WORKER_PROFILE=a mrq-worker a", - "MRQ_WORKER_PROFILE=b mrq-worker b" - ], - "worker2": ["MRQ_WORKER_PROFILE=a mrq-worker a"] - } - - # Worker diversity enforced under constraints - assert scenario({ - "a": { - "command": "mrq-worker a", - "memory": 1, - "cpu": 1, - "min_count": 3 - }, - "b": { - "command": "mrq-worker b", - "memory": 1, - "cpu": 1, - "min_count": 1 - } - }, [ - { - "_id": "worker1", - "total_cpu": 2, - "total_memory": 2, - "desired_workers": ["mrq-worker a", "mrq-worker a"] - } - ]) == { - "worker1": [ - "MRQ_WORKER_PROFILE=a mrq-worker a", - "MRQ_WORKER_PROFILE=b mrq-worker b" - ] - } + # assert scenario({ + # "a": { + # "command": "mrq-worker a", + # "memory": 1, + # "cpu": 1, + # "min_count": 3 + # }, + # "b": { + # "command": "mrq-worker b", + # "memory": 1, + # "cpu": 1, + # "min_count": 1 + # } + # }, [ + # { + # "_id": "worker1", + # "total_cpu": 11, + # "total_memory": 11, + # "desired_workers": ["mrq-worker a", "mrq-worker a"] + # }, { + # "_id": "worker2", + # "total_cpu": 5, + # "total_memory": 5, + # "desired_workers": ["mrq-worker a", "mrq-worker a"] + # } + # ]) == { + # "worker1": [ + # "MRQ_WORKER_PROFILE=a mrq-worker a", + # "MRQ_WORKER_PROFILE=a mrq-worker a", + # "MRQ_WORKER_PROFILE=b mrq-worker b" + # ], + # "worker2": ["MRQ_WORKER_PROFILE=a mrq-worker a"] + # } + + # # Worker diversity enforced under constraints + # assert scenario({ + # "a": { + # "command": "mrq-worker a", + # "memory": 1, + # "cpu": 1, + # "min_count": 3 + # }, + # "b": { + # "command": "mrq-worker b", + # "memory": 1, + # "cpu": 1, + # "min_count": 1 + # } + # }, [ + # { + # "_id": "worker1", + # "total_cpu": 2, + # "total_memory": 2, + # "desired_workers": ["mrq-worker a", "mrq-worker a"] + # } + # ]) == { + # "worker1": [ + # "MRQ_WORKER_PROFILE=a mrq-worker a", + # "MRQ_WORKER_PROFILE=b mrq-worker b" + # ] + # } def test_agent_process(worker): @@ -189,14 +136,10 @@ def test_agent_process(worker): assert connections.mongodb_jobs.mrq_workers.count() == 0 - connections.mongodb_jobs.mrq_workergroups.insert_one({"_id": "xxx", "profiles": { - "a": { - "command": "TEST_ENVVAR='&42' mrq-worker a --report_interval=1", - "memory": 100, - "cpu": 100, - "min_count": 1 - } - }}) + connections.mongodb_jobs.mrq_workergroups.insert_one({ + "_id": "xxx", + "commands": ["TEST_ENVVAR='&42' mrq-worker a --report_interval=1"] + }) time.sleep(7) @@ -208,7 +151,7 @@ def test_agent_process(worker): assert ctx["environ"].get("TEST_ENVVAR") == "&42" - connections.mongodb_jobs.mrq_workergroups.update_one({"_id": "xxx"}, {"$set": {"profiles": {}}}) + connections.mongodb_jobs.mrq_workergroups.update_one({"_id": "xxx"}, {"$set": {"commands": []}}) time.sleep(4) @@ -295,14 +238,11 @@ def test_agent_force_terminate(worker): # First, test interrupting a worker doing only a sleeping process. # - connections.mongodb_jobs.mrq_workergroups.insert_one({"_id": "xxx", "profiles": { - "a": { - "command": "mrq-worker default --report_interval=60", - "memory": 100, - "cpu": 100, - "min_count": 1 - } - }, "process_termination_timeout": 1}) + connections.mongodb_jobs.mrq_workergroups.insert_one({ + "_id": "xxx", + "commands": ["mrq-worker default --report_interval=60"], + "process_termination_timeout": 1 + }) time.sleep(5) @@ -313,14 +253,10 @@ def test_agent_force_terminate(worker): res1 = get_job_result(job1) assert res1["status"] == "started" - connections.mongodb_jobs.mrq_workergroups.update_one({"_id": "xxx"}, {"$set": {"profiles": { - "a": { - "command": "mrq-worker otherqueue --report_interval=60", - "memory": 100, - "cpu": 100, - "min_count": 1 - } - }, "process_termination_timeout": 1}}) + connections.mongodb_jobs.mrq_workergroups.update_one( + {"_id": "xxx"}, + {"$set": {"commands": ["mrq-worker otherqueue --report_interval=60"] + }}) time.sleep(5) @@ -340,14 +276,10 @@ def test_agent_force_terminate(worker): pids_before_sigkill = psutil.pids() - connections.mongodb_jobs.mrq_workergroups.update_one({"_id": "xxx"}, {"$set": {"profiles": { - "a": { - "command": "mrq-worker otherqueue2 --report_interval=60", - "memory": 100, - "cpu": 100, - "min_count": 1 - } - }, "process_termination_timeout": 1}}) + connections.mongodb_jobs.mrq_workergroups.update_one( + {"_id": "xxx"}, + {"$set": {"commands": ["mrq-worker otherqueue2 --report_interval=60"] + }}) # SIGKILL is sent after 5 seconds time.sleep(10) @@ -361,3 +293,18 @@ def test_agent_force_terminate(worker): assert len(pids_after_sigkill) == len(pids_before_sigkill) assert set(pids_after_sigkill) != set(pids_before_sigkill) + +def test_agent_multiple_processes(worker): + worker.start(agent=True, flags="--worker_group xxx --orchestrate_interval=1 --report_interval=1") + pids_before = psutil.pids() + + connections.mongodb_jobs.mrq_workergroups.insert_one({ + "_id": "xxx", + "commands": ["mrq-worker --processes 2 a", "mrq-worker --processes=2 b"], + "process_termination_timeout": 1 + }) + time.sleep(3) + + pids_after = psutil.pids() + # make sure there are 4 workers running + assert len(pids_before) + 4 == len(pids_after)