diff --git a/condorpy/htcondor_object_base.py b/condorpy/htcondor_object_base.py index 3bb09bf..6a6203a 100644 --- a/condorpy/htcondor_object_base.py +++ b/condorpy/htcondor_object_base.py @@ -81,6 +81,12 @@ def set_scheduler(self, host, username='root', password=None, private_key=None, Returns: An RemoteClient representing the remote scheduler. """ + # reuse an existing client when the connection parameters are unchanged. + existing = getattr(self, '_remote', None) + if existing is not None and existing.matches(host, username, password, private_key, + private_key_pass, port): + self._remote_id = getattr(self, '_remote_id', None) or uuid.uuid4().hex + return self._remote = RemoteClient(host, username, password, private_key, private_key_pass, port=port) self._remote_id = uuid.uuid4().hex diff --git a/condorpy/remote_utils.py b/condorpy/remote_utils.py index e446c9a..c66ce69 100644 --- a/condorpy/remote_utils.py +++ b/condorpy/remote_utils.py @@ -1,9 +1,37 @@ import os +import threading import paramiko import scp +_private_key_cache = {} +_private_key_cache_lock = threading.Lock() + + +def load_private_key(filename, password=None): + """Load an RSA private key, caching the decrypted result. + + Args: + filename (str): path to the private key file. + password (str, optional): passphrase for the private key. + + Returns: + paramiko.RSAKey: the decrypted key. + """ + try: + mtime = os.path.getmtime(filename) + except OSError: + mtime = None + cache_key = (filename, password, mtime) + with _private_key_cache_lock: + key = _private_key_cache.get(cache_key) + if key is None: + key = paramiko.RSAKey.from_private_key_file(filename=filename, password=password) + _private_key_cache[cache_key] = key + return key + + class RemoteClient(object): def __init__(self, host, @@ -15,14 +43,28 @@ def __init__(self, self.host = host self.username = username self.password = password + self.private_key_path = None + self.private_key_pass = private_key_pass if private_key: private_key = os.path.expanduser(private_key) - self.private_key = paramiko.RSAKey.from_private_key_file(filename=private_key, password=private_key_pass) + self.private_key_path = private_key + self.private_key = load_private_key(private_key, private_key_pass) self.port = port self._transport = None self._scp = None self._sftp = None + def matches(self, host, username, password, private_key, private_key_pass, port): + """True if this client already targets the given connection parameters.""" + return ( + self.host == host + and self.username == username + and self.password == password + and self.port == port + and self.private_key_path == (os.path.expanduser(private_key) if private_key else None) + and self.private_key_pass == private_key_pass + ) + def __del__(self): self.close() diff --git a/condorpy/workflow.py b/condorpy/workflow.py index 7f3a859..3855af0 100644 --- a/condorpy/workflow.py +++ b/condorpy/workflow.py @@ -106,7 +106,7 @@ def add_max_jobs_throttle(self, category, max_jobs): def node_set(self): """ """ - if self.cluster_id != self.NULL_CLUSTER_ID: + if self.cluster_id != self.NULL_CLUSTER_ID and not getattr(self, '_node_ids_resolved', False): self.update_node_ids() return self._node_set @@ -177,6 +177,38 @@ def _update_status(self, sub_job_num=None): return key + def node_statuses_by_cluster_id(self, sub_job_num=None): + """ + Get the status of every node in the DAG with a single remote query. + Parameters: + sub_job_num (int, optional): The sub-job number of the DAG. + Returns: + dict: cluster_id (int) -> condor status name (str), e.g. {12: 'Running'} + """ + dag_id = '%s.%s' % (self.cluster_id, sub_job_num) if sub_job_num else str(self.cluster_id) + job_delimiter = '+++' + attr_delimiter = ';;;' + format = [ + '-format', '"%d' + attr_delimiter + '"', 'ClusterId', + '-format', '"%d' + job_delimiter + '"', 'JobStatus', + ] + cmd = ('condor_q -constraint DAGManJobID=={0} {1} && ' + 'condor_history -constraint DAGManJobID=={0} {1}').format(dag_id, ' '.join(format)) + out, err = self._execute([cmd], shell=True, run_in_job_dir=False) + if err: + raise HTCondorError(err) + + statuses = dict() + for record in out.replace('"', '').split(job_delimiter): + parts = [p for p in record.strip().split(attr_delimiter) if p != ''] + if len(parts) < 2: + continue + try: + statuses[int(parts[0])] = CONDOR_JOB_STATUSES[int(parts[1])] + except (ValueError, KeyError): + continue + return statuses + def _update_statuses(self, sub_job_num=None): """ Update statuses of jobs nodes in workflow. @@ -187,10 +219,22 @@ def _update_statuses(self, sub_job_num=None): for val in CONDOR_JOB_STATUSES.values(): status_dict[val] = 0 + # One batched query for the whole DAG; fall back to per-node queries only + # if the batched form fails (e.g. an older schedd). + try: + by_cluster_id = self.node_statuses_by_cluster_id(sub_job_num=sub_job_num) + except (HTCondorError, KeyError, ValueError): + by_cluster_id = None + for node in self.node_set: job = node.job try: - job_status = job.status + if by_cluster_id is not None: + job_status = by_cluster_id.get(job.cluster_id) + if job_status is None: + job_status = 'Unexpanded' if job.cluster_id == job.NULL_CLUSTER_ID else job.status + else: + job_status = job.status status_dict[job_status] += 1 except (KeyError, HTCondorError): status_dict['Unexpanded'] += 1 @@ -265,6 +309,12 @@ def update_node_ids(self, sub_job_num=None): job._cluster_id = int(cluster_id) break + # Every node that could be resolved has been; don't re-query on each + # node_set access (see the node_set property). + self._node_ids_resolved = all( + node.job.cluster_id != node.job.NULL_CLUSTER_ID for node in self._node_set + ) + except ValueError as e: log.warning(str(e)) @@ -273,6 +323,8 @@ def add_node(self, node): """ assert isinstance(node, Node) self._node_set.add(node) + # A new node has no cluster id yet, so ids must be resolved again. + self._node_ids_resolved = False def add_job(self, job): """ diff --git a/tests/test_remote_reuse.py b/tests/test_remote_reuse.py new file mode 100644 index 0000000..9aa7c8a --- /dev/null +++ b/tests/test_remote_reuse.py @@ -0,0 +1,106 @@ +''' +Tests for private key caching, remote client reuse, and node id resolution. +''' +import os +import shutil +import tempfile +import unittest +from unittest import mock + +import paramiko + +from condorpy import Job, Workflow +from condorpy.remote_utils import RemoteClient, load_private_key +import condorpy.remote_utils as remote_utils + + +def load_tests(loader, tests, pattern): + suite = unittest.TestSuite() + suite.addTests(unittest.TestLoader().loadTestsFromTestCase(TestPrivateKeyCache)) + suite.addTests(unittest.TestLoader().loadTestsFromTestCase(TestRemoteClientReuse)) + suite.addTests(unittest.TestLoader().loadTestsFromTestCase(TestNodeIdResolution)) + return suite + + +class TestPrivateKeyCache(unittest.TestCase): + + def setUp(self): + self.dir = tempfile.mkdtemp() + self.key_path = os.path.join(self.dir, 'id_rsa') + paramiko.RSAKey.generate(2048).write_private_key_file(self.key_path) + remote_utils._private_key_cache.clear() + + def tearDown(self): + shutil.rmtree(self.dir, ignore_errors=True) + remote_utils._private_key_cache.clear() + + def test_key_is_decrypted_once(self): + with mock.patch.object(paramiko.RSAKey, 'from_private_key_file', + wraps=paramiko.RSAKey.from_private_key_file) as from_file: + first = load_private_key(self.key_path) + second = load_private_key(self.key_path) + self.assertIs(first, second) + self.assertEqual(1, from_file.call_count) + + def test_modified_key_file_is_reloaded(self): + first = load_private_key(self.key_path) + os.utime(self.key_path, (0, 0)) + second = load_private_key(self.key_path) + self.assertIsNot(first, second) + + def test_client_uses_cached_key(self): + a = RemoteClient('host', 'user', private_key=self.key_path) + b = RemoteClient('host', 'user', private_key=self.key_path) + self.assertIs(a.private_key, b.private_key) + + +class TestRemoteClientReuse(unittest.TestCase): + + def setUp(self): + self.job = Job('test_reuse') + + def test_scheduler_client_is_reused(self): + self.job.set_scheduler('host', 'user', password='pass') + first = self.job.scheduler + self.job.set_scheduler('host', 'user', password='pass') + self.assertIs(first, self.job.scheduler) + + def test_scheduler_client_replaced_when_host_changes(self): + self.job.set_scheduler('host', 'user', password='pass') + first = self.job.scheduler + self.job.set_scheduler('other-host', 'user', password='pass') + self.assertIsNot(first, self.job.scheduler) + + def test_remote_id_preserved_when_client_reused(self): + self.job.set_scheduler('host', 'user', password='pass') + remote_id = self.job._remote_id + self.job.set_scheduler('host', 'user', password='pass') + self.assertEqual(remote_id, self.job._remote_id) + + +class TestNodeIdResolution(unittest.TestCase): + + def setUp(self): + self.workflow = Workflow('test_nodes', config='', max_jobs=None) + self.workflow._cluster_id = 42 + + def test_node_set_resolves_ids_once(self): + with mock.patch.object(Workflow, 'update_node_ids') as update_node_ids: + self.workflow._node_ids_resolved = True + self.workflow.node_set + self.workflow.node_set + update_node_ids.assert_not_called() + + def test_node_set_resolves_ids_when_unresolved(self): + with mock.patch.object(Workflow, 'update_node_ids') as update_node_ids: + self.workflow.node_set + update_node_ids.assert_called_once() + + def test_adding_node_requires_resolution(self): + self.workflow._node_ids_resolved = True + self.workflow.add_job(Job('new_node')) + self.assertFalse(self.workflow._node_ids_resolved) + + +if __name__ == '__main__': + unittest.main()