From fba626dd3a302bd5e104c496603b35eb9f966fd9 Mon Sep 17 00:00:00 2001
From: Mark Lugar <14322382+araglu@users.noreply.github.com>
Date: Tue, 10 Jun 2025 11:00:59 -0500
Subject: [PATCH] Increase default timeout in connect() and allow override

Some HPCs take a minute to complete an initial connection through UIT+.
---
 uit/uit.py | 5 ++++-
 1 file changed, 4 insertions(+), 1 deletion(-)

diff --git a/uit/uit.py b/uit/uit.py
index d2dd0f0..e984b61 100644
--- a/uit/uit.py
+++ b/uit/uit.py
@@ -297,6 +297,7 @@ def connect(
         exclude_login_nodes=(),
         retry_on_failure=None,
         num_retries=3,
+        timeout=70,
     ):
         """Connect this client to the UIT servers.
 
@@ -310,12 +311,13 @@ def connect(
                 False will only attempt one connection.
                 Default of None will automatically pick False if login_node is set, otherwise it will pick True.
             num_retries (int): Number of connection attempts. Requires retry_on_failure=True
+            timeout (int): Number of seconds to wait for the call() to the HPC
         """
         login_node, retry_on_failure = self.prepare_connect(system, login_node, exclude_login_nodes, retry_on_failure)
 
         try:
             # working_dir='.' ends up being the location for UIT+ scripts, not the user's home directory
-            self.call(":", working_dir=".", timeout=35)
+            self.call(":", working_dir=".", timeout=timeout)
         except UITError as e:
             self.connected = False
             msg = f"Error while connecting to node {login_node}: {e}"
@@ -332,6 +334,7 @@ def connect(
                     exclude_login_nodes=exclude_login_nodes,
                     retry_on_failure=retry_on_failure,
                     num_retries=num_retries,
+                    timeout=timeout,
                 )
             else:
                 raise MaxRetriesError(msg)