Add new option for agentSST : start a specific agent

Alexis Koralewski
1 parent f2843971
Showing 3 changed files with 42 additions and 13 deletions Show diff stats
pyros.py
src/core/pyros_django/agent/Agent.py
src/core/pyros_django/agent/AgentSST.py
@@ -1016,8 +1016,9 @@ def start(agent: str, configfile: str, observatory: str, unit: str, computer_hos
 @click.option('--observatory', '-o', help='the observatory name to be used')
 @click.option('--unit', '-u', help='the unit name to be used')
 @click.option("--computer_hostname","-cp", help="The name of simulated computer hostname")
+@click.option("--agent","-a", help="The name of agent to be launched")
 @click.option("--foreground","-fg", is_flag=True, help="Print stdout and error in terminal")
-def new_start(configfile: str, observatory: str, unit: str, computer_hostname: str,foreground:bool):
+def new_start(configfile: str, observatory: str, unit: str, computer_hostname: str,foreground:bool,agent:str):
     log.debug("Running start command")
     try:
         from config.pyros.config_pyros import ConfigPyros
@@ -1116,6 +1117,8 @@ def new_start(configfile: str, observatory: str, unit: str, computer_hostname: s
         cmd += " -v"
     if computer_hostname:
         cmd += f" --computer {computer_hostname}"
+    if agent:
+        cmd += f" --agent {agent}"
  
     # if not test_mode(): current_processes.append( [execProcessFromVenvAsync(cmd), agent_name, -1] )
     # Append this process ( [process id, agent_name, result=failure] )
@@ -2721,7 +2721,7 @@ def extract_parameters():
  
 #def build_agent(Agent_type:Agent, name="GenericAgent", RUN_IN_THREAD=True):
 #def build_agent(Agent_type:Agent, RUN_IN_THREAD=True):
-def build_agent(Agent_type:Agent) -> Agent : 
+def build_agent(Agent_type:Agent,param_constr=None) -> Agent : 
     #DEBUG_MODE, WITH_SIM, TEST_MODE, VERBOSE_MODE, configfile = extract_parameters()
     DEBUG_MODE, TEST_MODE, VERBOSE_MODE = extract_parameters()
     log.debug("debug mode is" + os.getenv("PYROS_DEBUG"))
@@ -2732,11 +2732,16 @@ def build_agent(Agent_type:Agent) -&gt; Agent :
     #agent = Agent("GenericAgent", configfile, RUN_IN_THREAD=True)
     #agent = Agent_type(configfile, RUN_IN_THREAD, DEBUG_MODE=DEBUG_MODE)
     #agent = Agent_type(RUN_IN_THREAD)
-    
-    agent = Agent_type()
+    if param_constr:
+        if "AgentSST" in str(Agent_type):
+            agent = Agent_type(agent=param_constr.get("agent"))
+    else:
+        agent = Agent_type()
     if agent.name == "AgentSST":
+        # args are passed two times: first in AgentSST script and then to build_agent function, we need to setup again argparse in order to argparse to allow arguments
         parser = argparse.ArgumentParser(description='Start a agentSST.')
         parser.add_argument("--computer",dest="computer",help='Launch agent with simulated computer hostname',action="store")
+        parser.add_argument("--agent",dest="agent",help='Launch an specific agent ',action="store")
         parser.add_argument("-t", action="store_true" )
  
         # parser.add_argument("-d", action="store_true" )
@@ -2,7 +2,8 @@
  
 from pathlib import Path
 import subprocess
-import sys, os
+import sys, os, argparse
+
 from datetime import datetime, timezone, timedelta
 ##import utils.Logger as L
 #import threading #, multiprocessing, os
@@ -33,7 +34,7 @@ class AgentSST(Agent):
     TEST_COMMANDS_LIST = [
     ]
  
-    def __init__(self, name:str=None,sim_computer=None):
+    def __init__(self, name:str=None, sim_computer=None, agent=None):
  
         super().__init__()
         self.PROJECT_ROOT_PATH = os.environ["PROJECT_ROOT_PATH"]
@@ -59,13 +60,20 @@ class AgentSST(Agent):
                 + os.sep
             )
         self.VENV_PYTHON = VENV_BIN + "python3"
+        if agent:
+            self.init_agent = agent
+        else:
+            self.init_agent = None
         #log.info(f"PC hostname is {self.computer}")
         #name_from_config = self.get_config().get_agent_sst_of_current_computer()
  
     def init(self):
         super().init()
         log.info(f"PC hostname is {self.computer}")
-        self.start_agents()
+        if self.init_agent is None:
+            self.start_agents()
+        else:
+            self.start_agents(self.init_agent)
         self.TEST_MODE = False
         time.sleep(10)
         self.set_delay(3)
@@ -84,6 +92,11 @@ class AgentSST(Agent):
  
         if agent_name:
             agent = agent_name
+            agents = obs_config.get_agents_per_computer(obs_config.unit_name).get(self.computer)
+            if agent not in agents:
+                log.info(f"{agent} isn't associated to this computer : {self.computer}")
+                log.info(f"Agents associated to this computer : {agents}")
+                exit(1)
             # Start a specific agent of obs_config (restart)
             agent_informations = obs_config.get_agent_information(obs_config.unit_name,agent)
             protocol = agent_informations.get("protocol")
@@ -94,12 +107,13 @@ class AgentSST(Agent):
                 if os.path.exists(protocol_folder_abs_path + os.sep + protocol_script_name):
                     cmd = self.VENV_PYTHON +" "+ protocol_folder_abs_path + os.sep + protocol_script_name
                     if not agent in self.agent_in_mode_test:
-                        self.agent_in_mode_test = self.TEST_MODE
+                        self.agent_in_mode_test[agent] = self.TEST_MODE
                     if self.agent_in_mode_test[agent]:
                         cmd += test_mode
                     process = subprocess.Popen(f"{cmd}",shell=True)
                     process.poll()
-                    
+                    if agent not in self.subprocess_dict:
+                        self.subprocess_dict[agent] = {}
                     self.subprocess_dict[agent]["process"] = process
                     nb_try_restart = self.subprocess_dict[agent].get("nb_try_restart",0)
                     nb_try_restart += 1
@@ -158,7 +172,7 @@ class AgentSST(Agent):
         # agent = args[0]
         if agent in self.subprocess_dict.keys():
             cmd = self.send_cmd_to(agent,"do_exit")
-            #cmd = self.send_cmd_to(agent,"do_exit","asap")
+            #cmd = self.send_cmd_to(agent,"do_stop","asap")
             return cmd
  
     def do_restart_agent(self, agent:str)->None:
@@ -217,7 +231,9 @@ class AgentSST(Agent):
                 agent = cmd.args[0]
                 self.force_kill_agent(agent)
  
-        # checking status of agent if they are timeout
+        # checking status of agent if they are timeout if in auto mode
+        # soft_mode = Majordome.objects.objects.last().soft_mode
+        # if soft_mode == "AUTO":
         for agent in self.subprocess_dict.keys():
             try:
                 agent_survey = AgentSurvey.objects.get(name=agent)
@@ -275,6 +291,11 @@ class AgentSST(Agent):
  
  
 if __name__ == "__main__":
-
-    agent = build_agent(AgentSST)
+    parser = argparse.ArgumentParser(description='Start a agentSST.')
+    parser.add_argument("--computer",dest="computer",help='Launch agent with simulated computer hostname',action="store")
+    parser.add_argument("--agent",dest="agent",help='Launch an specific agent ',action="store")
+    parser.add_argument("-t", action="store_true" )
+    args = vars(parser.parse_args())
+    agent = build_agent(AgentSST,param_constr=args)
+    # agent = build_agent(AgentSST)
     agent.run()
 \ No newline at end of file