Cache private keys, reuse remote clients, and batch DAG status queries

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
 

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()
 

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

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()