From e6d162d600c1866fef7454a3807cff3dbec0d225 Mon Sep 17 00:00:00 2001 From: romer8 Date: Tue, 4 Aug 2026 14:52:32 -0600 Subject: [PATCH 1/4] Cache private keys, reuse remote clients, and batch DAG status queries Job status polling re-decrypted the private key and rebuilt the SSH client on every property access, and queried each DAG node individually. --- condorpy/htcondor_object_base.py | 9 +++ condorpy/remote_utils.py | 47 +++++++++++++- condorpy/workflow.py | 66 ++++++++++++++++++- tests/test_remote_reuse.py | 106 +++++++++++++++++++++++++++++++ 4 files changed, 224 insertions(+), 4 deletions(-) create mode 100644 tests/test_remote_reuse.py diff --git a/condorpy/htcondor_object_base.py b/condorpy/htcondor_object_base.py index 3bb09bf..c183b2a 100644 --- a/condorpy/htcondor_object_base.py +++ b/condorpy/htcondor_object_base.py @@ -81,6 +81,15 @@ def set_scheduler(self, host, username='root', password=None, private_key=None, Returns: An RemoteClient representing the remote scheduler. """ + # Callers (e.g. tethys CondorBase.condor_object) re-invoke this on every + # property access. Rebuilding the client each time discarded its cached + # SSH transport, forcing a fresh handshake per remote command; 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..c811416 100644 --- a/condorpy/remote_utils.py +++ b/condorpy/remote_utils.py @@ -1,8 +1,39 @@ import os +import threading import paramiko import scp +# Decrypting a passphrase-protected private key runs a KDF (~0.5s) and was being +# repeated on every RemoteClient construction -- several times per job status poll. +# Key material only changes when the file does, so cache on (path, passphrase, +# mtime); the mtime keeps a rotated key from being served stale. +_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, @@ -15,14 +46,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..9a5d620 100644 --- a/condorpy/workflow.py +++ b/condorpy/workflow.py @@ -104,9 +104,14 @@ def add_max_jobs_throttle(self, category, max_jobs): @property def node_set(self): + """The set of nodes in this workflow. + + Reading this attribute used to run ``update_node_ids()`` -- and therefore a + remote query -- on every access, so simply iterating the nodes cost a round + trip each time. Node ids only need to be resolved once per object, so the + result is remembered; call ``update_node_ids()`` directly to force a refresh. """ - """ - 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 +182,41 @@ 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. + + Querying each node individually costs one round trip per node (two, counting + condor_history), which dominates status-polling cost on large DAGs. The + DAGManJobID constraint already returns every node of this workflow, so one + query is sufficient. + + 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 +227,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 +317,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 +331,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() From 30b1d4db0f14288914f2991076b8c06e38b39843 Mon Sep 17 00:00:00 2001 From: romer8 Date: Tue, 4 Aug 2026 15:10:37 -0600 Subject: [PATCH 2/4] Optimize remote client usage by reusing existing connections when parameters are unchanged --- condorpy/htcondor_object_base.py | 5 +---- 1 file changed, 1 insertion(+), 4 deletions(-) diff --git a/condorpy/htcondor_object_base.py b/condorpy/htcondor_object_base.py index c183b2a..6a6203a 100644 --- a/condorpy/htcondor_object_base.py +++ b/condorpy/htcondor_object_base.py @@ -81,10 +81,7 @@ def set_scheduler(self, host, username='root', password=None, private_key=None, Returns: An RemoteClient representing the remote scheduler. """ - # Callers (e.g. tethys CondorBase.condor_object) re-invoke this on every - # property access. Rebuilding the client each time discarded its cached - # SSH transport, forcing a fresh handshake per remote command; reuse an - # existing client when the connection parameters are unchanged. + # 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): From 98e88d928e54a8572199868951fa4714057b3fbd Mon Sep 17 00:00:00 2001 From: romer8 Date: Tue, 4 Aug 2026 15:19:08 -0600 Subject: [PATCH 3/4] Remove redundant comments about private key caching in remote_utils.py --- condorpy/remote_utils.py | 5 +---- 1 file changed, 1 insertion(+), 4 deletions(-) diff --git a/condorpy/remote_utils.py b/condorpy/remote_utils.py index c811416..c66ce69 100644 --- a/condorpy/remote_utils.py +++ b/condorpy/remote_utils.py @@ -4,10 +4,7 @@ import paramiko import scp -# Decrypting a passphrase-protected private key runs a KDF (~0.5s) and was being -# repeated on every RemoteClient construction -- several times per job status poll. -# Key material only changes when the file does, so cache on (path, passphrase, -# mtime); the mtime keeps a rotated key from being served stale. + _private_key_cache = {} _private_key_cache_lock = threading.Lock() From 64fc9871b8df9440401c5a30993f6afe07f43d4e Mon Sep 17 00:00:00 2001 From: romer8 Date: Wed, 5 Aug 2026 10:52:22 -0600 Subject: [PATCH 4/4] cleaning --- condorpy/workflow.py | 18 +++++------------- 1 file changed, 5 insertions(+), 13 deletions(-) diff --git a/condorpy/workflow.py b/condorpy/workflow.py index 9a5d620..3855af0 100644 --- a/condorpy/workflow.py +++ b/condorpy/workflow.py @@ -104,12 +104,7 @@ def add_max_jobs_throttle(self, category, max_jobs): @property def node_set(self): - """The set of nodes in this workflow. - - Reading this attribute used to run ``update_node_ids()`` -- and therefore a - remote query -- on every access, so simply iterating the nodes cost a round - trip each time. Node ids only need to be resolved once per object, so the - result is remembered; call ``update_node_ids()`` directly to force a refresh. + """ """ if self.cluster_id != self.NULL_CLUSTER_ID and not getattr(self, '_node_ids_resolved', False): self.update_node_ids() @@ -183,13 +178,10 @@ 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. - - Querying each node individually costs one round trip per node (two, counting - condor_history), which dominates status-polling cost on large DAGs. The - DAGManJobID constraint already returns every node of this workflow, so one - query is sufficient. - + """ + 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'} """