Agent fixes

Author: sylvinus

.gitignore

@@ -17,6 +17,7 @@ develop-eggs
 lib
 lib64
 __pycache__
+.cache
 
 # Installer logs
 pip-log.txt

mrq/agent.py

@@ -17,6 +17,7 @@ def __init__(self, worker_group=None):
         self.id = ObjectId()
         self.worker_group = worker_group or get_current_config()["worker_group"]
         self.pool = ProcessPool()
+        self.config = get_current_config()
 
     def work(self):
 
@@ -50,7 +51,7 @@ def greenlet_manage(self):
             try:
                 self.manage()
             except Exception as e:  # pylint: disable=broad-except
-                self.log.error("When reporting: %s" % e)
+                log.error("When reporting: %s" % e)
             finally:
                 time.sleep(self.config["report_interval"])
 
@@ -59,11 +60,13 @@ def manage(self):
         report = self.get_agent_report()
 
         try:
-            db = self.mongodb_jobs.mrq_agents.find_and_modify({
+            db = connections.mongodb_jobs.mrq_agents.find_and_modify({
                 "_id": ObjectId(self.id)
             }, {"$set": report}, upsert=True)
+            if not db:
+                return
         except Exception as e:  # pylint: disable=broad-except
-            self.log.debug("Agent report failed: %s" % e)
+            log.debug("Agent report failed: %s" % e)
             return
 
         # If the desired_workers was changed by an orchestrator, apply the changes locally
@@ -74,16 +77,17 @@ 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"]
+            "available_memory": get_current_config()["available_memory"],
+            "worker_group": self.worker_group
         }
         return report
 
     def greenlet_orchestrate(self):
 
         while True:
-            with connections.redis.lock(self.redis_agent_orchestrator_key, timeout=60):
+            with connections.redis.lock(self.redis_agent_orchestrator_key, timeout=self.config["orchestrate_interval"] + 10):
                 self.orchestrate()
-                time.sleep(30)
+                time.sleep(self.config["orchestrate_interval"])
 
     @property
     def redis_agent_orchestrator_key(self):

mrq/bin/mrq_agent.py

@@ -13,16 +13,18 @@
 import sys
 import argparse
 
-from .config import get_config
-from .agent import Agent
-from .context import set_current_config
+sys.path.insert(0, os.getcwd())
+
+from mrq import config
+from mrq.agent import Agent
+from mrq.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"))
+    cfg = config.get_config(parser=parser, config_type="agent", sources=("file", "env", "args"))
 
     set_current_config(cfg)
 

mrq/config.py

@@ -280,6 +280,20 @@ def add_parser_args(parser, config_type):
             type=int,
             help="How much CPU units this agent's workers can use. We recommend using 1024 per CPU.")
 
+        parser.add_argument(
+            '--orchestrate_interval',
+            default=30,
+            action="store",
+            type=float,
+            help="How much seconds to wait between orchestration runs.")
+
+        parser.add_argument(
+            '--report_interval',
+            default=10,
+            action='store',
+            type=float,
+            help='Seconds between agent reports to MongoDB')
+
     # Worker-specific args
     elif config_type == "worker":
 

mrq/worker.py

@@ -422,6 +422,8 @@ def work(self):
 
         self.work_loop(max_jobs=self.max_jobs)
 
+        self.work_stop()
+
     def work_init(self):
 
         self.connect()

tests/conftest.py

@@ -136,24 +136,30 @@ def __init__(self, request, **kwargs):
 
         self.started = False
 
-    def start(self, flush=True, deps=True, trace=True, **kwargs):
+    def start(self, flush=True, deps=True, trace=True, agent=False, **kwargs):
 
         self.started = True
 
         if deps:
             self.start_deps(flush=flush)
 
         processes = 0
-        m = re.search(r"--processes (\d+)", kwargs.get("flags", ""))
-        if m:
-            processes = int(m.group(1))
-
-        cmdline = "python mrq/bin/mrq_worker.py --mongodb_logs_size 0 %s %s %s %s" % (
-            "--admin_port 20020" if (processes <= 1) else "",
-            "--trace_io --trace_greenlets" if trace else "",
-            kwargs.get("flags", ""),
-            kwargs.get("queues", "high default low")
-        )
+
+        if agent:
+            cmdline = "python mrq/bin/mrq_agent.py %s" % kwargs.get("flags", "")
+
+        else:
+
+            m = re.search(r"--processes (\d+)", kwargs.get("flags", ""))
+            if m:
+                processes = int(m.group(1))
+
+            cmdline = "python mrq/bin/mrq_worker.py --mongodb_logs_size 0 %s %s %s %s" % (
+                "--admin_port 20020" if (processes <= 1) else "",
+                "--trace_io --trace_greenlets" if trace else "",
+                kwargs.get("flags", ""),
+                kwargs.get("queues", "high default low")
+            )
 
         print(cmdline)
         ProcessFixture.start(self, cmdline=cmdline, env=kwargs.get("env"), expected_children=processes)

tests/test_agent.py

@@ -1,5 +1,6 @@
 from mrq.agent import Agent
 from mrq.context import connections
+import time
 
 
 def scenario(profiles, agents):
@@ -136,3 +137,40 @@ def test_orchestration_scenarios(worker):
     }
 
     worker.stop()
+
+
+def test_agent_process(worker):
+
+    worker.start(agent=True, flags="--worker_group xxx --available_memory=500 --available_cpu=500 --orchestrate_interval=1 --report_interval=1")
+
+    time.sleep(3)
+
+    agents = list(connections.mongodb_jobs.mrq_agents.find())
+
+    assert len(agents) == 1
+
+    assert connections.mongodb_jobs.mrq_workers.count() == 0
+
+    connections.mongodb_jobs.mrq_workergroups.insert_one({"_id": "xxx", "profiles": [
+        {
+            "command": "mrq-worker a",
+            "memory": 100,
+            "cpu": 100,
+            "min_count": 1
+        }
+    ]})
+
+    time.sleep(3)
+
+    assert connections.mongodb_jobs.mrq_workers.count() == 1
+    worker = connections.mongodb_jobs.mrq_workers.find_one()
+    assert worker["status"] == "wait"
+
+    connections.mongodb_jobs.mrq_workergroups.update_one({"_id": "xxx"}, {"$set": {"profiles": [
+
+    ]}})
+
+    time.sleep(4)
+
+    worker = connections.mongodb_jobs.mrq_workers.find_one()
+    assert worker["status"] == "stop"