From d4acd670bc41d1bbbc7241cdfd19b60ea99192dc Mon Sep 17 00:00:00 2001
From: allegroai <>
Date: Mon, 20 May 2024 15:57:11 +0300
Subject: [PATCH] Fix connect to session created by SSH application

---
 clearml_session/__main__.py                 | 10 ++++++----
 clearml_session/interactive_session_task.py |  2 +-
 2 files changed, 7 insertions(+), 5 deletions(-)

diff --git a/clearml_session/__main__.py b/clearml_session/__main__.py
index f68caa7..25561ed 100644
--- a/clearml_session/__main__.py
+++ b/clearml_session/__main__.py
@@ -741,7 +741,7 @@ def clone_task(state, project_id=None):
     return task
 
 
-def wait_for_machine(state, task):
+def wait_for_machine(state, task, only_wait_for_ssh=False):
     # wait until task is running
     print('Waiting for remote machine allocation [id={}]'.format(task.id))
     last_status = None
@@ -781,9 +781,9 @@ def wait_for_machine(state, task):
     state['vscode_server'] = vscode_server.strip().lower() != 'false'
 
     wait_properties = ['properties/internal_ssh_port']
-    if state.get('jupyter_lab'):
+    if state.get('jupyter_lab') and not only_wait_for_ssh:
         wait_properties += ['properties/jupyter_port']
-    if state.get('vscode_server'):
+    if state.get('vscode_server') and not only_wait_for_ssh:
         wait_properties += ['properties/vscode_port']
 
     last_lines = []
@@ -1467,6 +1467,7 @@ def cli():
     task = _get_previous_session(client, args, state, task_id=args.attach)
 
     delete_old_tasks_callback = None
+    only_wait_for_ssh = False
 
     if task:
         state['task_id'] = task.id
@@ -1475,6 +1476,7 @@ def cli():
             state['username'] = args.username
         if args.password:
             state['password'] = args.password
+        only_wait_for_ssh = True
     else:
         state.pop('task_id', None)
         save_state(state, state_file)
@@ -1511,7 +1513,7 @@ def cli():
 
     # wait for machine to become available
     try:
-        wait_for_machine(state, task)
+        wait_for_machine(state, task, only_wait_for_ssh=only_wait_for_ssh)
     except ValueError as ex:
         print('\nERROR: {}'.format(ex))
         return 1
diff --git a/clearml_session/interactive_session_task.py b/clearml_session/interactive_session_task.py
index c7056e6..5e27cb2 100644
--- a/clearml_session/interactive_session_task.py
+++ b/clearml_session/interactive_session_task.py
@@ -132,7 +132,7 @@ def init_task(param, a_default_ssh_fingerprint):
     except (TypeError, ValueError):
         a_default_ssh_fingerprint.clear()
         a_default_ssh_fingerprint.update(old_default_ssh_fingerprint)
-    if param.get('default_docker'):
+    if param.get('default_docker') and task.running_locally():
         task.set_base_docker("{} --network host".format(param['default_docker']))
     # leave local process, only run remotely
     task.execute_remotely()