ClearML/clearml-session
1
0
Fork 0
mirror of https://github.com/clearml/clearml-session synced 2025-06-11 00:46:43 +00:00
Code Issues Actions Packages Projects Releases Wiki Activity

Add shutdown support

Browse Source
This commit is contained in:
allegroai 2023-07-05 08:44:22 +03:00
parent 9f9881b0ef
commit f99713ce76
1 changed files with 33 additions and 14 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

45
clearml_session/__main__.py
View File

@ -687,7 +687,7 @@ def wait_for_machine(state, task):
def start_ssh_tunnel(username, remote_address, ssh_port, ssh_password, local_remote_pair_list, debug=False): def start_ssh_tunnel(username, remote_address, ssh_port, ssh_password, local_remote_pair_list, debug=False):
print('Starting SSH tunnel') print('Starting SSH tunnel to {}@{}, port {}'.format(username, remote_address, ssh_port))
child = None child = None
args = ['-N', '-C', args = ['-N', '-C',
'{}@{}'.format(username, remote_address), '-p', '{}'.format(ssh_port), '{}@{}'.format(username, remote_address), '-p', '{}'.format(ssh_port),
@ -913,6 +913,8 @@ def setup_parser(parser):
help='Display the clearml-session utility version') help='Display the clearml-session utility version')
parser.add_argument('--attach', default=False, nargs='?', parser.add_argument('--attach', default=False, nargs='?',
help='Attach to running interactive session (default: previous session)') help='Attach to running interactive session (default: previous session)')
parser.add_argument("--shutdown", "-S", default=None, const="", nargs="?",
help="Shut down an active session (default: previous session)")
parser.add_argument('--debugging-session', type=str, default=None, parser.add_argument('--debugging-session', type=str, default=None,
help='Pass existing Task id (experiment), create a copy of the experiment on a remote machine, ' help='Pass existing Task id (experiment), create a copy of the experiment on a remote machine, '
'and launch jupyter/ssh for interactive access. Example --debugging-session <task_id>') 'and launch jupyter/ssh for interactive access. Example --debugging-session <task_id>')
@ -1031,8 +1033,20 @@ def cli():
client = APIClient() client = APIClient()
if args.shutdown is not None:
task = _get_previous_session(
client, args, state, task_id=args.shutdown, verb="Shutting down",
question_verb="Shut down", ask_for_explicit_id=True
)
if not task:
print("No session to shut down, exiting")
return 1
task.mark_stopped()
print("Session #{} shut down, goodbye!".format(task.id))
return 0
# get previous session, if it is running # get previous session, if it is running
task = _check_previous_session(client, args, state) task = _get_previous_session(client, args, state, task_id=args.attach)
if task: if task:
state['task_id'] = task.id state['task_id'] = task.id
@ -1079,12 +1093,12 @@ def cli():
print('Leaving interactive session') print('Leaving interactive session')
def _check_previous_session(client, args, state): def _get_previous_session(
client, args, state, task_id=None, verb="Connecting to", question_verb="Connect to", ask_for_explicit_id=False
):
assume_yes = args.yes assume_yes = args.yes
# now let's see if we have the requested Task if task_id:
if args.attach: print('Checking session #{}'.format(task_id))
task_id = args.attach
print('Checking previous session')
try: try:
task = Task.get_task(task_id=task_id) task = Task.get_task(task_id=task_id)
except ValueError: except ValueError:
@ -1092,9 +1106,13 @@ def _check_previous_session(client, args, state):
status = task.get_status() if task else None status = task.get_status() if task else None
if status == 'in_progress': if status == 'in_progress':
if not args.debugging_session or task.parent == args.debugging_session: if not args.debugging_session or task.parent == args.debugging_session:
# only ask if we were not requested directly if assume_yes or not ask_for_explicit_id:
print('Using active session id={}'.format(task_id)) print("{} active session id={}".format(verb, task_id))
return task return task
choice = input("{} active session id={} [Y]/n? ".format(question_verb, task_id))
if str(choice).strip().lower() in ("y", "yes"):
return task
return None
raise ValueError('Could not connect to requested session id={} - status \'{}\''.format( raise ValueError('Could not connect to requested session id={} - status \'{}\''.format(
task_id, status or 'Not Found')) task_id, status or 'Not Found'))
@ -1113,10 +1131,10 @@ def _check_previous_session(client, args, state):
if len(running_task_ids_created) == 1: if len(running_task_ids_created) == 1:
task_id = running_task_ids_created[0][0] task_id = running_task_ids_created[0][0]
if assume_yes: if assume_yes:
print("Connecting to active session {}".format(task_id)) print("{} active session {}".format(verb, task_id))
else: else:
choice = input("Connect to active session id={} [Y]/n? ".format(task_id)) choice = input("{} active session id={} [Y]/n? ".format(question_verb, task_id))
if str(choice).strip().lower() not in ("n", "no"): if str(choice).strip().lower() in ("y", "yes"):
return Task.get_task(task_id=task_id) return Task.get_task(task_id=task_id)
# multiple sessions running # multiple sessions running
@ -1138,7 +1156,8 @@ def _check_previous_session(client, args, state):
try: try:
choice = input( choice = input(
session_list session_list
+ "\nConnect to session [{}] or 'N' to skip: ".format( + "\n{} session [{}] or 'N' to skip: ".format(
question_verb,
"0" if len(running_task_ids_created) <= 1 else "0-{}".format(len(running_task_ids_created) - 1) "0" if len(running_task_ids_created) <= 1 else "0-{}".format(len(running_task_ids_created) - 1)
) )
) )