Update ssh

model/matchmaker.py

@@ -18,8 +18,22 @@ def create_ssh_matchmaker_client(server, port, user, password):
     ssh_matchmaker_client.set_missing_host_key_policy(paramiko.AutoAddPolicy())
     ssh_matchmaker_client.connect(server, port, user, password)
 
-    sftp_matchmaker_client = ssh_matchmaker_client.open_sftp()
+    transport = ssh_matchmaker_client.get_transport()
+    transport.set_keepalive(60)
 
+    sftp_matchmaker_client = ssh_matchmaker_client.open_sftp()
+def is_connected():
+    global ssh_matchmaker_client, sftp_matchmaker_client
+    # 检查SSH连接是否正常
+    if not ssh_matchmaker_client.get_transport().is_active():
+        return False
+    # 检查SFTP连接是否正常
+    try:
+        sftp_matchmaker_client.listdir('.')  # 尝试列出根目录
+    except Exception as e:
+        print(f"Error checking SFTP connection: {e}")
+        return False
+    return True
 def ucb_score(trueskill_diff, t, n):
     exploration_term = np.sqrt((2 * np.log(t + 1e-5)) / (n + 1e-5))
     ucb = -trueskill_diff + 1.0 * exploration_term
@@ -37,6 +51,8 @@ def deserialize_rating(rating_dict):
 
 def save_json_via_sftp(ratings, comparison_counts, total_comparisons):
     global sftp_matchmaker_client
+    if not is_connected():
+        create_ssh_matchmaker_client(SSH_SERVER, SSH_PORT, SSH_USER, SSH_PASSWORD)
     data = {
         'ratings': [serialize_rating(r) for r in ratings],
         'comparison_counts': comparison_counts.tolist(),
@@ -48,6 +64,8 @@ def save_json_via_sftp(ratings, comparison_counts, total_comparisons):
 
 def load_json_via_sftp():
     global sftp_matchmaker_client
+    if not is_connected():
+        create_ssh_matchmaker_client(SSH_SERVER, SSH_PORT, SSH_USER, SSH_PASSWORD)
     with sftp_matchmaker_client.open(SSH_SKILL, 'r') as f:
         data = json.load(f)
     ratings = [deserialize_rating(r) for r in data['ratings']]

serve/gradio_web.py

@@ -375,7 +375,8 @@ def build_side_by_side_ui_anony(models):
     vote_textbox.submit(
         disable_vote,
         inputs=None,
-        outputs=[vote_submit_btn, vote_mode_btn]
+        outputs=[vote_submit_btn, vote_mode_btn, \
+                 A1_btn, A2_btn, A3_btn, A4_btn, B1_btn, B2_btn, B3_btn, B4_btn, C1_btn, C2_btn, C3_btn, C4_btn, D1_btn, D2_btn, D3_btn, D4_btn]
     ).then(
         text_response_rank_igm,
         inputs=[generate_ig0, generate_ig1, generate_ig2, generate_ig3, Top1_text, Top2_text, Top3_text, Top4_text, vote_textbox], 
@@ -391,7 +392,8 @@ def build_side_by_side_ui_anony(models):
     vote_submit_btn.click(
         disable_vote,
         inputs=None,
-        outputs=[vote_submit_btn, vote_mode_btn]
+        outputs=[vote_submit_btn, vote_mode_btn, \
+                 A1_btn, A2_btn, A3_btn, A4_btn, B1_btn, B2_btn, B3_btn, B4_btn, C1_btn, C2_btn, C3_btn, C4_btn, D1_btn, D2_btn, D3_btn, D4_btn]
     ).then(
         text_response_rank_igm,
         inputs=[generate_ig0, generate_ig1, generate_ig2, generate_ig3, Top1_text, Top2_text, Top3_text, Top4_text, vote_textbox], 

serve/update_skill.py

@@ -19,8 +19,22 @@ def create_ssh_skill_client(server, port, user, password):
     ssh_skill_client.set_missing_host_key_policy(paramiko.AutoAddPolicy())
     ssh_skill_client.connect(server, port, user, password)
 
-    sftp_skill_client = ssh_skill_client.open_sftp()
+    transport = ssh_skill_client.get_transport()
+    transport.set_keepalive(60)
 
+    sftp_skill_client = ssh_skill_client.open_sftp()
+def is_connected():
+    global ssh_skill_client, sftp_skill_client
+    # 检查SSH连接是否正常
+    if not ssh_skill_client.get_transport().is_active():
+        return False
+    # 检查SFTP连接是否正常
+    try:
+        sftp_skill_client.listdir('.')  # 尝试列出根目录
+    except Exception as e:
+        print(f"Error checking SFTP connection: {e}")
+        return False
+    return True
 def ucb_score(trueskill_diff, t, n):
     exploration_term = np.sqrt((2 * np.log(t + 1e-5)) / (n + 1e-5))
     ucb = -trueskill_diff + 1.0 * exploration_term
@@ -39,6 +53,8 @@ def deserialize_rating(rating_dict):
 
 def save_json_via_sftp(ratings, comparison_counts, total_comparisons):
     global sftp_skill_client
+    if not is_connected():
+        create_ssh_skill_client(SSH_SERVER, SSH_PORT, SSH_USER, SSH_PASSWORD)
     data = {
         'ratings': [serialize_rating(r) for r in ratings],
         'comparison_counts': comparison_counts.tolist(),
@@ -50,6 +66,8 @@ def save_json_via_sftp(ratings, comparison_counts, total_comparisons):
 
 def load_json_via_sftp():
     global sftp_skill_client
+    if not is_connected():
+        create_ssh_skill_client(SSH_SERVER, SSH_PORT, SSH_USER, SSH_PASSWORD)
     with sftp_skill_client.open(SSH_SKILL, 'r') as f:
         data = json.load(f)
     ratings = [deserialize_rating(r) for r in data['ratings']]

serve/upload.py

@@ -17,7 +17,22 @@ def create_ssh_client(server, port, user, password):
     ssh_client.set_missing_host_key_policy(paramiko.AutoAddPolicy())
     ssh_client.connect(server, port, user, password)
 
+    transport = ssh_client.get_transport()
+    transport.set_keepalive(60)
+
     sftp_client = ssh_client.open_sftp()
+def is_connected():
+    global ssh_client, sftp_client
+    # 检查SSH连接是否正常
+    if not ssh_client.get_transport().is_active():
+        return False
+    # 检查SFTP连接是否正常
+    try:
+        sftp_client.listdir('.')  # 尝试列出根目录
+    except Exception as e:
+        print(f"Error checking SFTP connection: {e}")
+        return False
+    return True
 
 def get_image_from_url(image_url):
     response = requests.get(image_url)
@@ -26,7 +41,8 @@ def get_image_from_url(image_url):
 
 def get_random_mscoco_prompt():
     global sftp_client
-    create_ssh_client(SSH_SERVER, SSH_PORT, SSH_USER, SSH_PASSWORD)
+    if not is_connected():
+        create_ssh_client(SSH_SERVER, SSH_PORT, SSH_USER, SSH_PASSWORD)
     num = random.randint(0, 2999)
     file = "{}.txt".format(num)
 
@@ -40,6 +56,8 @@ def get_random_mscoco_prompt():
 
 def create_remote_directory(remote_directory):
     global ssh_client
+    if not is_connected():
+        create_ssh_client(SSH_SERVER, SSH_PORT, SSH_USER, SSH_PASSWORD)
     stdin, stdout, stderr = ssh_client.exec_command(f'mkdir -p {SSH_LOG}/{remote_directory}')
     error = stderr.read().decode('utf-8')
     if error:
@@ -50,6 +68,8 @@ def create_remote_directory(remote_directory):
     
 def upload_ssh_all(states, output_dir, data, data_path):
     global sftp_client
+    if not is_connected():
+        create_ssh_client(SSH_SERVER, SSH_PORT, SSH_USER, SSH_PASSWORD)
     output_file_list = []
     image_list = []
     for i in range(len(states)):
@@ -71,4 +91,4 @@ def upload_ssh_all(states, output_dir, data, data_path):
     with io.BytesIO(json_data.encode('utf-8')) as json_byte_stream:
         sftp_client.putfo(json_byte_stream, data_path)
     print(f"Successfully uploaded JSON data to {data_path}")
-    create_ssh_client(SSH_SERVER, SSH_PORT, SSH_USER, SSH_PASSWORD)
+    # create_ssh_client(SSH_SERVER, SSH_PORT, SSH_USER, SSH_PASSWORD)

serve/utils.py

@@ -176,7 +176,7 @@ def enable_vote_buttons():
 def disable_vote_buttons():
     return tuple(gr.update(visible=False, interactive=False) for i in range(6))
 def disable_vote():
-    return (gr.update(interactive=False), gr.update(interactive=False))
+    return (gr.update(interactive=False) for i in range(14))
 def enable_vote_mode_buttons(mode):
     print(mode)
     if mode == "Best":