diff --git a/automation/script/experiment.py b/automation/script/experiment.py index ac414fa4c..c5b63e5ed 100644 --- a/automation/script/experiment.py +++ b/automation/script/experiment.py @@ -50,48 +50,29 @@ def experiment_run(self_module, i): exp = {} cur_dir = os.getcwd() - r = self_module.search(i.copy()) + + r = self_module._select_script(i) if r['return'] > 0: return r - lst = r['list'] - if not lst: - return {'return': 1, 'error': 'No scripts were found'} - - # Process each artifact - for artifact in sorted(lst, key=lambda x: x.meta.get('alias', '')): - meta, script_path = artifact.meta, artifact.path - tags, script_alias, script_uid = meta.get( - "tags", []), meta.get( - 'alias', ''), meta.get( - 'uid', '') - - # Execute the experiment script - mlc_script_input = { - 'action': 'run', 'target': 'script' - } - if exp: - for key in exp: - ii = {**mlc_script_input, **run_input} - if isinstance(exp[key], list): - for val in exp[key]: - ii[key] = val - r = run_script_and_tag_experiment( - ii, - self_module.action_object, - experiment_action, - tags, - extra_exp_tags, - meta, - skip_state_save, - logger) - if r['return'] > 0: - return r - elif isinstance(exp[key], dict): - return { - 'return': 1, 'error': 'Dictionary inputs are not supported for mlc experiment script'} - else: - ii[key] = exp[key] + script = r['script'] + + meta, script_path = script.meta, script.path + tags, script_alias, script_uid = meta.get( + "tags", []), meta.get( + 'alias', ''), meta.get( + 'uid', '') + + # Execute the experiment script + mlc_script_input = { + 'action': 'run', 'target': 'script' + } + if exp: + for key in exp: + ii = {**mlc_script_input, **run_input} + if isinstance(exp[key], list): + for val in exp[key]: + ii[key] = val r = run_script_and_tag_experiment( ii, self_module.action_object, @@ -103,6 +84,22 @@ def experiment_run(self_module, i): logger) if r['return'] > 0: return r + elif isinstance(exp[key], dict): + return { + 'return': 1, 'error': 'Dictionary inputs are not supported for mlc experiment script'} + else: + ii[key] = exp[key] + r = run_script_and_tag_experiment( + ii, + self_module.action_object, + experiment_action, + tags, + extra_exp_tags, + meta, + skip_state_save, + logger) + if r['return'] > 0: + return r return {'return': 0} diff --git a/automation/script/module.py b/automation/script/module.py index e44d6db91..a815bdd38 100644 --- a/automation/script/module.py +++ b/automation/script/module.py @@ -4669,7 +4669,7 @@ def _select_script(self, i): sorted_list = sorted(lst, key=lambda x: x.meta.get('alias', '')) # If quiet mode is off, prompt the user - if not i.get('quiet', False): + if not i.get('quiet', False) and len(sorted_list) > 1: print("\nAvailable scripts:") for idx, artifact in enumerate(sorted_list, 1): meta = artifact.meta diff --git a/automation/script/remote_run.py b/automation/script/remote_run.py index f8d7b9ae3..f0ea97aac 100644 --- a/automation/script/remote_run.py +++ b/automation/script/remote_run.py @@ -1,6 +1,5 @@ from collections import defaultdict import os -from mlc.main import ExperimentAction import mlc.utils as utils from mlc import utils from utils import * @@ -39,193 +38,51 @@ def remote_run(self_module, i): return prune_result run_input = prune_result['new_input'] - - r = convert_input(i) + mlc_run_cmd = run_input['mlc_run_cmd'] cur_dir = os.getcwd() - r = self_module.search(i.copy()) - if r['return'] > 0: - return r - - lst = r['list'] - if not lst: - return {'return': 1, 'error': 'No scripts were found'} - - # Process each artifact - for artifact in sorted(lst, key=lambda x: x.meta.get('alias', '')): - meta, script_path = artifact.meta, artifact.path - tags, script_alias, script_uid = meta.get( - "tags", []), meta.get( - 'alias', ''), meta.get( - 'uid', '') - - # Execute the experiment script - mlc_script_input = { - 'action': 'run', 'target': 'script' - } - if exp: - for key in exp: - ii = {**mlc_script_input, **run_input} - if isinstance(exp[key], list): - for val in exp[key]: - ii[key] = val - r = run_script_and_tag_experiment( - ii, - self_module.action_object, - experiment_action, - tags, - extra_exp_tags, - meta, - skip_state_save, - logger) - if r['return'] > 0: - return r - elif isinstance(exp[key], dict): - return { - 'return': 1, 'error': 'Dictionary inputs are not supported for mlc experiment script'} - else: - ii[key] = exp[key] - r = run_script_and_tag_experiment( - ii, - self_module.action_object, - experiment_action, - tags, - extra_exp_tags, - meta, - skip_state_save, - logger) - if r['return'] > 0: - return r - - return {'return': 0} - - -def run_script_and_tag_experiment( - ii, script_action, experiment_action, tags, extra_exp_tags, script_meta, skip_state_save, logger): - - current_path = os.path.abspath(os.getcwd()) - experiment_meta = {} - recursion_spaces = '' - exp_tags = tags + extra_exp_tags - ssi = {'action': 'update', - 'target': 'experiment', - 'script_alias': script_meta['alias'], - 'script_uid': script_meta['uid'], - 'tags': ','.join(exp_tags), - 'extra_tags': ",".join(extra_exp_tags), - 'meta': experiment_meta, - 'force': True} - r = experiment_action.access(ssi) + r = self_module._select_script(i) if r['return'] > 0: return r - experiment = r['list'][0] - - logger.debug( - recursion_spaces + - ' - Changing to {}'.format(experiment.path)) - - os.chdir(experiment.path) - # Get current datetime in YYYY-MM-DD_HH-MM-SS format - timestamp = datetime.now().strftime("%Y-%m-%d_%H-%M-%S") - - # Create a folder name using the timestamp - folder_name = f"run_{timestamp}" - - # Create the directory - os.makedirs(folder_name, exist_ok=True) - os.chdir(folder_name) - - if not skip_state_save: - ssi = {'action': 'run', - 'target': 'script', - 'tags': 'save,system,state', - 'outfile': 'system_state_before.json', - 'quiet': True - } - r = script_action.access(ssi) - if r['return'] > 0: - return r - - start_time = time.time() - r = script_action.access(ii) + script = r['script'] + + meta, script_path = script.meta, script.path + tags, script_alias, script_uid = meta.get( + "tags", []), meta.get( + 'alias', ''), meta.get( + 'uid', '') + + # Execute the experiment script + mlc_script_input = { + 'action': 'run', 'target': 'script' + } + + run_cmds = [] + remote_mlc_python_venv = i.get('remote_python_venv', 'mlcflow') + run_cmds.append(f". {remote_mlc_python_venv}/bin/activate") + if i.get('remote_pull_mlc_repos', False): + run_cmds.append("mlc pull repo") + + script_run_cmd = " ".join(mlc_run_cmd.split(" ")[3:]) + run_cmds.append(f"mlcr {script_run_cmd}") + + remote_inputs = {} + for key in ["host", "port", "user", "client_refresh", + "password", "skip_host_verify", "ssh_key_file"]: + if i.get(f"remote_{key}"): + remote_inputs[key] = i[f"remote_{key}"] + + # Execute the remote command + mlc_remote_input = { + 'action': 'run', 'target': 'script', 'tags': 'remote,run,cmds,ssh', + 'script_tags': i.get('tags'), 'run_cmds': run_cmds, + **remote_inputs + } + + r = self_module.action_object.access(mlc_remote_input) if r['return'] > 0: return r - end_time = time.time() - elapsed = end_time - start_time - time_taken_string = format_elapsed(elapsed) - logger.info(f"Time taken: {time_taken_string}") - - if not skip_state_save: - ssi['outfile'] = 'system_state_after.json' - r = script_action.access(ssi) - if r['return'] > 0: - return r - - ''' - exp_tags = tags - ii = {'action': 'update', - 'target': 'experiment', - 'uid': experiment.meta['uid'], - 'meta': experiment.meta, - 'script_alias': meta['alias'], - 'replace_lists': True, # To replace tags - 'tags': ','.join(exp_tags)} - - r = experiment_action.access(ii) - if r['return'] > 0: - return r - ''' - os.chdir(current_path) - logger.info( - f"Experiment entry saved at: {os.path.join(experiment.path, folder_name)}") - - return {'return': 0, 'experiment': experiment, 'folder_name': folder_name} - - elif seconds < 3600: - mins, secs = divmod(seconds, 60) - return f"{int(mins)} minutes {secs:.1f} seconds" - else: - hours, remainder = divmod(seconds, 3600) - mins, secs = divmod(remainder, 60) - return f"{int(hours)} hours {int(mins)} minutes {secs:.1f} seconds" - - -def parse_value(val): - if isinstance(val, list): - return [parse_value(v) for v in val] - - val = str(val) - - # Handle range inputs like 2:10 or 2:10:2 - if ':' in val: - parts = val.split(':') - try: - parts = list(map(int, parts)) - if len(parts) == 2: - return list(range(parts[0], parts[1] + 1)) - elif len(parts) == 3: - return list(range(parts[0], parts[1] + 1, parts[2])) - except ValueError: - pass # Not a valid range, fall through - - # Convert to int if possible - if val.isdigit(): - return int(val) - - return val - - -def convert_input(input_dict): - output = defaultdict(dict) - - for key, value in input_dict.items(): - if '.' in key: - main_key, sub_key = key.split('.', 1) - output[main_key][sub_key] = parse_value(value) - elif isinstance(value, dict): - output[key].update({k: parse_value(v) for k, v in value.items()}) - - return dict(output) + return {'return': 0}