multiply worker group commands by number of processes

Author: maelorn

mrq/agent.py

@@ -4,12 +4,13 @@
 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
 
 
@@ -195,7 +196,8 @@ def fetch_worker_group_definition(self):
         # Prepend all commands by their worker profile.
         commands = []
         for command in definition.get("commands", []):
-            commands.append("MRQ_WORKER_GROUP=%s %s" % (self.worker_group, command))
-        definition["commands"] = 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
@@ -84,8 +85,8 @@ def fetch_worker_group_definitions(self):
             commands = []
             # Prepend all commands by their worker group.
             for command in definition.get("commands", []):
-                commands.append("MRQ_WORKER_GROUP=%s %s" % (definition["_id"], command))
-
+                simplified_command, worker_count = normalize_command(command, definition["_id"])
+                commands.extend([simplified_command] * worker_count)
             definition["commands"] = commands
 
         return definitions

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 """

tests/test_agent.py

@@ -292,3 +292,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()g
+
+    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)