Improvements to mlc remote_run (#741)

* Pass --quiet for get,docker

* Improve remote_run

Author: arjunsuresh

automation/script/docker.py

@@ -388,7 +388,7 @@ def docker_run(self_module, i):
 
     # Ensure Docker is available
     r = self_module.action_object.access(
-        {'action': 'run', 'automation': 'script', 'tags': "get,docker"})
+        {'action': 'run', 'automation': 'script', 'tags': "get,docker", 'quiet': True})
     if r['return'] > 0:
         return r
 

automation/script/remote_run.py

@@ -30,7 +30,7 @@ def remote_run(self_module, i):
     logger = self_module.logger
     env = i.get('env', {})
     remote_host = i.get('remote_host', 'localhost')
-    remote_port = i.get('remote_port', 22)
+    remote_port = i.get('remote_port', '22')
 
     prune_result = prune_input(
         {'input': i, 'extra_keys_starts_with': ['remote_']})
@@ -40,6 +40,7 @@ def remote_run(self_module, i):
     run_input = prune_result['new_input']
     mlc_run_cmd = run_input['mlc_run_cmd']
 
+    # print(script_cmd)
     cur_dir = os.getcwd()
 
     r = self_module._select_script(i)
@@ -59,10 +60,20 @@ def remote_run(self_module, i):
     if r['return'] > 0:
         return r
 
+    remote_env = {}
+
     remote_run_settings = r['remote_run_settings']
     env = r['env']
+    state = r['state']
+    meta = r['meta']
+
+    r = call_remote_run_prepare(self_module, meta, script, env, state, i)
+    if r['return'] > 0:
+        return r
+
+    files_to_copy = r.get('files_to_copy', [])
+    remote_env = r.get('remote_env', {})
 
-    # Execute the experiment script
     mlc_script_input = {
         'action': 'run', 'target': 'script'
     }
@@ -73,23 +84,50 @@ def remote_run(self_module, i):
     if i.get('remote_pull_mlc_repos', False):
         run_cmds.append("mlc pull repo")
 
-    files_to_copy = []
     env_keys_to_copy = remote_run_settings.get('env_keys_to_copy', [])
+    input_mapping = meta.get('input_mapping', {})
+
     for key in env_keys_to_copy:
         if key in env and os.path.exists(env[key]):
             files_to_copy.append(env[key])
+            remote_env[key] = os.path.join(
+                "mlc-remote-artifacts",
+                os.path.basename(
+                    env[key]))
+
+            for k, value in input_mapping.items():
+                if value == key and k in run_input:
+                    run_input[k] = remote_env[key]
+
+    i_copy = copy.deepcopy(i)
+    i_copy['run_cmd'] = run_input
+
+    r = regenerate_script_cmd(i_copy)
+    if r['return'] > 0:
+        return r
+
+    # " ".join(mlc_run_cmd.split(" ")[1:])
+    script_run_cmd = r['run_cmd_string']
+
+    if remote_env:
+        for key in remote_env:
+            script_run_cmd += f" --env.{key}={remote_env[key]}"
 
-    script_run_cmd = " ".join(mlc_run_cmd.split(" ")[1:])
-    run_cmds.append(f"mlcr {script_run_cmd}")
+    run_cmds.append(f"{script_run_cmd}")
 
     remote_inputs = {}
+
     for key in ["host", "port", "user", "client_refresh",
-                "password", "skip_host_verify", "ssh_key_file"]:
+                "password", "skip_host_verify", "ssh_key_file", "copy_directory"]:
         if i.get(f"remote_{key}"):
             remote_inputs[key] = i[f"remote_{key}"]
 
     if files_to_copy:
+        remote_copy_directory = i.get(
+            "remote_copy_directory",
+            "mlc-remote-artifacts")
         remote_inputs['files_to_copy'] = files_to_copy
+        remote_inputs['copy_directory'] = remote_copy_directory
 
     # Execute the remote command
     mlc_remote_input = {
@@ -170,4 +208,175 @@ def update_meta_for_selected_variations(self_module, script, input_params):
     # Set Docker-specific configurations
     remote_run_settings = state_data['remote_run']
     return {'return': 0, 'remote_run_settings': remote_run_settings,
-            'env': env, 'state': state_data}
+            'env': env, 'state': state_data, 'meta': metadata}
+
+
+def call_remote_run_prepare(self_module, meta, script_item, env, state, i):
+
+    path_to_customize_py = os.path.join(script_item.path, 'customize.py')
+    logger = self_module.logger
+    recursion_spaces = ''
+
+    # Check and run remote_run_prepare in customize.py
+    if os.path.isfile(path_to_customize_py):
+        r = utils.load_python_module(
+            {'path': script_item.path, 'name': 'customize'})
+        if r['return'] > 0:
+            return r
+
+        customize_code = r['code']
+
+        customize_common_input = {
+            'input': i,
+            'automation': self_module,
+            'artifact': script_item,
+            # 'customize': script_item.meta.get('customize', {}),
+            # 'os_info': os_info,
+            # 'recursion_spaces': recursion_spaces,
+            # 'script_tags': script_tags,
+            # 'variation_tags': variation_tags
+        }
+        run_script_input = {}
+        run_script_input['customize_code'] = customize_code
+        run_script_input['customize_common_input'] = customize_common_input
+
+        if 'remote_run_prepare' in dir(customize_code):
+
+            logger.debug(
+                recursion_spaces +
+                '  - Running remote_run_prepare ...')
+
+            run_script_input['run_state'] = {}
+
+            ii = copy.deepcopy(customize_common_input)
+            ii['env'] = env
+            ii['state'] = state
+            ii['meta'] = meta
+            ii['automation'] = self_module
+            # may need to detect versions in multiple paths
+            ii['run_script_input'] = run_script_input
+
+            r = customize_code.remote_run_prepare(ii)
+            return r
+
+    return {'return': 0}
+
+
+def regenerate_script_cmd(i):
+
+    remote_run_settings = i.get('remote_run_settings', {})
+    fake_run = i.get('fake_run', False)
+
+    i_run_cmd = i['run_cmd']
+
+    # Remove environment variables with host path values
+    if 'env' in i_run_cmd:
+        env = i_run_cmd['env']
+        for key in list(env):
+            value = env[key]
+
+            # Check if the value is a string containing the specified paths
+            if isinstance(value, str) and (
+                    os.path.join("local", "cache", "") in value or
+                    os.path.join("MLC", "repos", "") in value or
+                    "<<<" in value
+            ):
+                del env[key]
+
+            # Check if the value is a list and remove matching items
+            elif isinstance(value, list):
+                # Identify values to remove
+                values_to_remove = [
+                    val for val in value
+                    if isinstance(val, str) and (
+                        os.path.join("local", "cache", "") in val or
+                        os.path.join("MLC", "repos", "") in val or
+                        "<<<" in value
+                    )
+                ]
+
+                # Remove key if all values match; otherwise, filter the list
+                if len(values_to_remove) == len(value):
+                    del env[key]
+                else:
+                    env[key] = [
+                        val for val in value if val not in values_to_remove]
+
+    # docker_run_cmd_prefix = i.get('docker_run_cmd_prefix', '')
+
+    # Regenerate command from dictionary input
+    run_cmd = 'mlcr'
+
+    skip_input_for_fake_run = remote_run_settings.get(
+        'skip_input_for_fake_run', [])
+    add_quotes_to_keys = remote_run_settings.get('add_quotes_to_keys', [])
+
+    def rebuild_flags(
+            command_dict,
+            is_fake_run,
+            skip_keys_for_fake_run,
+            quote_keys,
+            prefix
+    ):
+        """
+        Recursively rebuilds command-line flags from a dictionary of inputs.
+
+        :param command_dict: Dictionary containing command-line keys and values.
+        :param is_fake_run: Boolean indicating if this is a fake run.
+        :param skip_keys_for_fake_run: List of keys to skip in fake run mode.
+        :param quote_keys: List of keys that require values to be quoted.
+        :param prefix: String to prepend to keys for hierarchical keys.
+        :return: A reconstructed command-line string.
+        """
+        command_line = ""
+
+        # Sort keys to ensure 'tags' appears first if present.
+        keys = sorted(command_dict.keys(), key=lambda x: x != "tags")
+
+        for key in keys:
+            if key in ["input", "output", "outdirname"]:
+                continue  # We have the corresponding env keys in container env string
+            # Construct the full key with the prefix.
+            full_key = f"{prefix}.{key}" if prefix else key
+
+            # Skip keys marked for exclusion in fake run mode.
+            if is_fake_run and full_key in skip_keys_for_fake_run:
+                continue
+
+            value = command_dict[key]
+            quote = '"' if full_key in quote_keys else ""
+
+            # Recursively process nested dictionaries.
+            if isinstance(value, dict):
+                command_line += rebuild_flags(
+                    value,
+                    is_fake_run,
+                    skip_keys_for_fake_run,
+                    quote_keys,
+                    full_key
+                )
+            # Process lists by concatenating values with commas.
+            elif isinstance(value, list):
+                list_values = ",".join(
+                    quote_if_needed(
+                        item, quote) for item in value)
+                command_line += f" --{full_key},={list_values}"
+            # Process scalar values.
+            else:
+                if full_key in ['s', 'v']:
+                    command_line += f" -{full_key}"
+                else:
+                    command_line += f" --{full_key}={quote_if_needed(value, quote)}"
+
+        return command_line
+
+    run_cmd += rebuild_flags(i_run_cmd,
+                             fake_run,
+                             skip_input_for_fake_run,
+                             add_quotes_to_keys,
+                             '')
+
+    # run_cmd = docker_run_cmd_prefix + ' && ' + \
+    #    run_cmd if docker_run_cmd_prefix != '' else run_cmd
+
+    return {'return': 0, 'run_cmd_string': run_cmd}

script/remote-run-commands/customize.py

@@ -59,14 +59,16 @@ def preprocess(i):
     if password:
         rsync_base = ["sshpass", "-p", password] + rsync_base
 
+    target_directory = env['MLC_SSH_TARGET_COPY_DIRECTORY']
+
     # ---- Execute copy commands ----
     for file in files_to_copy:
         cmd = [
             "rsync",
             "-avz",
             "-e", " ".join(ssh_cmd),   # rsync expects a single string here
             file,
-            f"{user}@{host}:",
+            f"{user}@{host}:{target_directory}/"
         ]
 
         print("Executing:", " ".join(cmd))

script/remote-run-commands/meta.yaml

@@ -7,6 +7,7 @@ default_env:
   MLC_SSH_HOST: localhost
   MLC_SSH_KEY_FILE: <<<HOME>>>/.ssh/id_rsa
   MLC_SSH_PORT: '22'
+  MLC_SSH_TARGET_COPY_DIRECTORY: '.'
 input_mapping:
   client_refresh: MLC_SSH_CLIENT_REFRESH
   host: MLC_SSH_HOST
@@ -17,6 +18,7 @@ input_mapping:
   skip_host_verify: MLC_SSH_SKIP_HOST_VERIFY
   ssh_key_file: MLC_SSH_KEY_FILE
   user: MLC_SSH_USER
+  copy_directory: MLC_SSH_TARGET_COPY_DIRECTORY
 tags:
 - remote
 - run