Added --files_to_copy support for remote_run

Author: amd-arsuresh

automation/script/module.py

@@ -69,6 +69,8 @@ def __init__(self, action_object, automation_file):
                                'MLC_RENEW_CACHE_ENTRY']
 
         self.host_env_keys = [
+            "USER",
+            "HOME",
             "GH_TOKEN",
             "ftp_proxy",
             "FTP_PROXY",

automation/script/remote_run.py

@@ -54,6 +54,14 @@ def remote_run(self_module, i):
         'alias', ''), meta.get(
         'uid', '')
 
+    #Update meta for selected variation and input
+    r = update_meta_for_selected_variations(self_module, script, i)
+    if r['return'] > 0:
+        return r
+
+    remote_run_settings = r['remote_run_settings']
+    env = r['env']
+
     # Execute the experiment script
     mlc_script_input = {
         'action': 'run', 'target': 'script'
@@ -65,6 +73,12 @@ 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')
+    for key in env_keys_to_copy:
+        if key in env and os.path.exists(env[key]):
+            files_to_copy.append(env[key])
+
     script_run_cmd = " ".join(mlc_run_cmd.split(" ")[1:])
     run_cmds.append(f"mlcr {script_run_cmd}")
 
@@ -73,6 +87,10 @@ def remote_run(self_module, i):
                 "password", "skip_host_verify", "ssh_key_file"]:
         if i.get(f"remote_{key}"):
             remote_inputs[key] = i[f"remote_{key}"]
+    
+    if files_to_copy:
+        remote_inputs['files_to_copy'] = files_to_copy
+    
 
     # Execute the remote command
     mlc_remote_input = {
@@ -86,3 +104,68 @@ def remote_run(self_module, i):
         return r
 
     return {'return': 0}
+
+def update_meta_for_selected_variations(self_module, script, input_params):
+    metadata = script.meta
+    script_directory = script.path
+    script_tags = metadata.get("tags", [])
+    script_alias = metadata.get('alias', '')
+    script_uid = metadata.get('uid', '')
+    tag_values = input_params.get('tags', '').split(",")
+    variation_tags = [tag[1:] for tag in tag_values if tag.startswith("_")]
+
+    run_state = {
+        'deps': [],
+        'fake_deps': [],
+        'parent': None,
+        'script_id': f"{script_alias},{script_uid}",
+        'script_variation_tags': variation_tags
+    }
+    state_data = {}
+    env = input_params.get('env', {})
+    constant_vars = input_params.get('const', {})
+    constant_state = input_params.get('const_state', {})
+
+    remote_run_settings = metadata.get('remote_run', {})
+    remote_run_settings_default_env = remote_run_settings.get('default_env', {})
+    for key in remote_run_settings_default_env:
+        env.setdefault(key, remote_run_settings_default_env[key])
+
+    state_data['remote_run'] = remote_run_settings
+    add_deps_recursive = input_params.get('add_deps_recursive', {})
+
+    # Update state with metadata and variations
+    update_state_result = self_module.update_state_from_meta(
+        metadata, env, state_data, constant_vars, constant_state,
+        deps=[],
+        post_deps=[],
+        prehook_deps=[],
+        posthook_deps=[],
+        new_env_keys=[],
+        new_state_keys=[],
+        run_state=run_state,
+        i=input_params
+    )
+    if update_state_result['return'] > 0:
+        return update_state_result
+
+    update_variations_result = self_module._update_state_from_variations(
+        input_params, metadata, variation_tags, metadata.get(
+            'variations', {}),
+        env, state_data, constant_vars, constant_state,
+        deps=[],  # Add your dependencies if needed
+        post_deps=[],  # Add post dependencies if needed
+        prehook_deps=[],  # Add prehook dependencies if needed
+        posthook_deps=[],  # Add posthook dependencies if needed
+        new_env_keys_from_meta=[],  # Add keys from meta if needed
+        new_state_keys_from_meta=[],  # Add state keys from meta if needed
+        add_deps_recursive=add_deps_recursive,
+        run_state=run_state,
+        recursion_spaces=''
+    )
+    if update_variations_result['return'] > 0:
+        return update_variations_result
+
+    # 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}

script/remote-run-commands/customize.py

@@ -1,6 +1,6 @@
 from mlc import utils
 import os
-
+import subprocess
 
 def preprocess(i):
 
@@ -13,6 +13,8 @@ def preprocess(i):
     # pre_run_cmds = env.get('MLC_SSH_PRE_RUN_CMDS', ['source $HOME/cm/bin/activate'])
     pre_run_cmds = env.get('MLC_SSH_PRE_RUN_CMDS', [])
 
+    files_to_copy = env.get('MLC_SSH_FILES_TO_COPY', [])
+
     run_cmds = env.get('MLC_SSH_RUN_COMMANDS', [])
 
     run_cmds = pre_run_cmds + run_cmds
@@ -24,23 +26,68 @@ def preprocess(i):
             run_cmds[i] = cmd
 
     cmd_string += " ; ".join(run_cmds)
-    user = env.get('MLC_SSH_USER')
+    user = env.get('MLC_SSH_USER', os.environ.get('USER'))
     password = env.get('MLC_SSH_PASSWORD', None)
     host = env.get('MLC_SSH_HOST')
     if password:
         password_string = " -p " + password
     else:
         password_string = ""
-    cmd_extra = ''
+
+
+    ssh_cmd = ["ssh"]
 
     if env.get("MLC_SSH_SKIP_HOST_VERIFY"):
-        cmd_extra += " -o StrictHostKeyChecking=no"
-    if env.get("MLC_SSH_KEY_FILE"):
-        cmd_extra += " -i " + env.get("MLC_SSH_KEY_FILE")
+        ssh_cmd += ["-o", "StrictHostKeyChecking=no", "-o", "UserKnownHostsFile=/dev/null"]
+
+    key_file = env.get("MLC_SSH_KEY_FILE")
+    if key_file:
+        ssh_cmd += ["-i", key_file]
+
+    ssh_cmd_str = " ".join(ssh_cmd)
+
+
+
+    ssh_run_command = ssh_cmd_str + " " +user + "@" + host + \
+        password_string + " '" + cmd_string + "'"
+    env['MLC_SSH_CMD'] = ssh_run_command
+
+
+    # ---- Use sshpass if password is provided ----
+    rsync_base = ["rsync", "-avz"]
+
+    if password:
+        rsync_base = ["sshpass", "-p", password] + rsync_base
+
+    # ---- 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}:",
+        ]
+
+        print("Executing:", " ".join(cmd))
+
+        result = subprocess.run(
+            cmd,
+            env=os.environ,
+            stdout=subprocess.PIPE,
+            stderr=subprocess.PIPE,
+            text=True,
+        )
+
+        if result.returncode != 0:
+            raise RuntimeError(
+                f"❌ rsync failed for {file}\n"
+                f"STDOUT:\n{result.stdout}\n"
+                f"STDERR:\n{result.stderr}"
+            )
+
+        print(f"✅ Copied {file} successfully")
 
-    ssh_command = "ssh " + user + "@" + host + \
-        password_string + cmd_extra + " '" + cmd_string + "'"
-    env['MLC_SSH_CMD'] = ssh_command
 
     return {'return': 0}
 

script/remote-run-commands/meta.yaml

@@ -5,15 +5,15 @@ category: Remote automation
 default_env:
   MLC_SSH_CLIENT_REFRESH: '10'
   MLC_SSH_HOST: localhost
-  MLC_SSH_KEY_FILE: $HOME/.ssh/id_rsa
+  MLC_SSH_KEY_FILE: <<<HOME>>>/.ssh/id_rsa
   MLC_SSH_PORT: '22'
-  MLC_SSH_USER: $USER
 input_mapping:
   client_refresh: MLC_SSH_CLIENT_REFRESH
   host: MLC_SSH_HOST
   password: MLC_SSH_PASSWORD
   port: MLC_SSH_PORT
   run_cmds: MLC_SSH_RUN_COMMANDS
+  files_to_copy: MLC_SSH_FILES_TO_COPY
   skip_host_verify: MLC_SSH_SKIP_HOST_VERIFY
   ssh_key_file: MLC_SSH_KEY_FILE
   user: MLC_SSH_USER