def invoke_agent_expect_result(self, host, command, args = {}):
from chroma_core.services.job_scheduler.agent_rpc import AgentException
result = self.invoke_agent(host, command, args)
# This case is to deal with upgrades, once every installation is using the new protocol then we should not allow this.
# Once everything is 3.0 or later we will also have version information in the wrapper header.
if (result == None) or \
((type(result) == dict) and ('error' not in result) and ('result' not in result)):
job_log.info("Invalid result %s fixed up on called to %s with args %s" % (result, command, args))
# Prior to 3.0 update_packages returned {'update_packages': data} so fix this up. This code is here so that all
# of the legacy fixups are in one place and can easily be removed.
if command == 'install_packages' and 'scan_packages' in result:
result = agent_result(result['scan_packages'])
else:
result = agent_result(result)
if type(result) != dict:
raise AgentException(host.fqdn, command, args, "Expected a dictionary but got a %s when calling %s" % (type(result), command))
if ('error' not in result) and ('result' not in result):
raise AgentException(host.fqdn, command, args, "Expected a dictionary with 'error' or 'result' in keys but got %s when calling %s" % (result, command))
if 'error' in result:
self.log(result['error'])
raise AgentException(host.fqdn, command, args, result['error'])
return result['result']
def invoke_rust_agent_expect_result(self, host, command, args={}):
from chroma_core.services.job_scheduler.agent_rpc import AgentException
try:
result = self.invoke_rust_agent(host, command, args)
except RustAgentCancellation as e:
raise AgentException(host, command, args,
"Cancelled: {}".format(e))
except Exception as e:
raise AgentException(host, command, args,
"Unexpected error: {}".format(e))
try:
result = json.loads(result)
except ValueError as e:
raise AgentException(
host,
command,
args,
"Error parsing json: {}; result: {}; command: {}; args: {}".
format(e, result, command, args),
)
if "Err" in result:
self.log(json.dumps(result["Err"], indent=2))
raise AgentException(host, command, args, result["Err"])
return result["Ok"]
def _detect_scan_device_plugin(host, command, args = None):
self.assertIn(command, ['detect_scan', 'device_plugin'])
if command == 'detect_scan':
return host_data[host]
raise AgentException(host, command, args, "No device plugin data available in unit tests")
def invoke_rust_agent(self, host, command, args={}):
"""
Talks to the iml-action-runner service
"""
from chroma_core.services.job_scheduler.agent_rpc import AgentException
try:
return invoke_rust_agent(host, command, args, self._cancel_event)
except RustAgentCancellation as e:
raise AgentException(
host, command, args,
"Cancelled: {}; command: {}; args: {}".format(
e, command, args))
except Exception as e:
raise AgentException(
host, command, args,
"Unexpected error: {}; command: {}; args: {}".format(
e, command, args))
def _try_ssh_cmd(self, agent_ssh, auth_args, cmd):
from chroma_core.services.job_scheduler.agent_rpc import AgentException
try:
return agent_ssh.ssh(cmd, auth_args=auth_args)
except (AuthenticationException, SSHException):
raise
except Exception, e:
# Re-raise wrapped in an AgentException
raise AgentException(
agent_ssh.address, "Unhandled exception: %s" % e,
", ".join(auth_args),
'\n'.join(traceback.format_exception(*(sys.exc_info()))))
def _fail(cls, fqdn):
from chroma_core.services.job_scheduler.agent_rpc import AgentException
log.error("Synthetic agent error on host %s" % fqdn)
raise AgentException(fqdn, "cmd", {"foo": "bar"}, "Fake backtrace")
