Improvements to remote-run

Author: arjunsuresh

.github/workflows/test-remote-run.yml

@@ -0,0 +1,95 @@
+name: Test remote run features
+
+on:
+  pull_request:
+    branches: [ "main", "dev" ]
+    paths:
+      - '.github/workflows/test-remote-run.yml'
+      - 'automation/script/remote_run.py'
+      - 'script/remote-run-commands/**'
+
+permissions:
+  contents: read
+
+jobs:
+  test_remote_run:
+    runs-on: ubuntu-latest
+
+    steps:
+    - uses: actions/checkout@v4
+
+    - name: Set up Python
+      uses: actions/setup-python@v3
+      with:
+        python-version: "3.12"
+
+    - name: Install mlcflow and pull repo
+      env:
+        REPO: ${{ github.event.pull_request.head.repo.html_url }}
+        BRANCH: ${{ github.event.pull_request.head.ref }}
+      run: |
+        pip install mlcflow
+        mlc pull repo "$REPO" --branch="$BRANCH"
+
+    - name: Set up SSH to localhost
+      run: |
+        ssh-keygen -t ed25519 -f ~/.ssh/id_ed25519 -N ""
+        cat ~/.ssh/id_ed25519.pub >> ~/.ssh/authorized_keys
+        chmod 600 ~/.ssh/authorized_keys
+        ssh-keyscan -H localhost >> ~/.ssh/known_hosts
+        # Verify SSH to localhost works
+        ssh -o BatchMode=yes localhost "echo SSH_OK"
+
+    - name: Test remote pre_run_cmds and post_run_cmds
+      run: |
+        mlcr remote,run,cmds,ssh \
+          --host=localhost \
+          --ssh_key_file=~/.ssh/id_ed25519 \
+          --run_cmds,="echo MAIN_CMD_RAN" \
+          --pre_run_cmds,="echo PRE_RUN_EXECUTED > /tmp/pre_run_test.txt" \
+          --post_run_cmds,="echo POST_RUN_EXECUTED > /tmp/post_run_test.txt"
+
+        # Verify pre_run_cmds executed
+        test -f /tmp/pre_run_test.txt
+        grep -q "PRE_RUN_EXECUTED" /tmp/pre_run_test.txt
+        echo "pre_run_cmds: PASSED"
+
+        # Verify post_run_cmds executed
+        test -f /tmp/post_run_test.txt
+        grep -q "POST_RUN_EXECUTED" /tmp/post_run_test.txt
+        echo "post_run_cmds: PASSED"
+
+    - name: Test remote_pre_run_cmds via remote_run
+      run: |
+        mlcr detect,os \
+          --remote_host=localhost \
+          --remote_ssh_key_file=~/.ssh/id_ed25519 \
+          --remote_pre_run_cmds,="echo REMOTE_PRE_RAN > /tmp/remote_pre_test.txt" \
+          --remote_post_run_cmds,="echo REMOTE_POST_RAN > /tmp/remote_post_test.txt"
+
+        test -f /tmp/remote_pre_test.txt
+        grep -q "REMOTE_PRE_RAN" /tmp/remote_pre_test.txt
+        echo "remote_pre_run_cmds via remote_run: PASSED"
+
+        test -f /tmp/remote_post_test.txt
+        grep -q "REMOTE_POST_RAN" /tmp/remote_post_test.txt
+        echo "remote_post_run_cmds via remote_run: PASSED"
+
+    - name: Test remote_copy_mlc_repos with custom MLC_REPOS path
+      run: |
+        # Use a random-ish path for MLC_REPOS on the "remote"
+        export CUSTOM_REPOS_PATH="/tmp/custom-mlc-repos-$(date +%s)"
+        mkdir -p "$CUSTOM_REPOS_PATH"
+
+        # Create a dummy repo to copy
+        mkdir -p ~/MLC/repos/test-dummy-repo
+        echo "DUMMY_CONTENT" > ~/MLC/repos/test-dummy-repo/marker.txt
+
+        mlcr detect,os \
+          --remote_host=localhost \
+          --remote_ssh_key_file=~/.ssh/id_ed25519 \
+          --remote_copy_mlc_repos,=test-dummy-repo
+
+        # Verify the repo was copied to MLC/repos (default rsync target)
+        test -f ~/MLC/repos/test-dummy-repo/marker.txt
+        echo "remote_copy_mlc_repos: PASSED"

automation/script/remote_run.py

@@ -140,6 +140,7 @@ def remote_run(self_module, i):
             script_run_cmd += f" --env.{key}={remote_env[key]}"
 
     remote_pre_run_cmds = i.get('remote_pre_run_cmds', [])
+    remote_post_run_cmds = i.get('remote_post_run_cmds', [])
 
     run_cmds.append(f"{script_run_cmd}")
 
@@ -157,6 +158,33 @@ def remote_run(self_module, i):
         remote_inputs['files_to_copy'] = files_to_copy
         remote_inputs['copy_directory'] = remote_copy_directory
 
+    # For repo copying, add a separate copy with MLC/repos target
+    if i.get('remote_copy_mlc_repos', False):
+        local_repos_path = os.path.join(os.path.expanduser("~"), "MLC", "repos")
+        repos_to_copy = i.get('remote_copy_mlc_repos', [])
+        if isinstance(repos_to_copy, bool) or repos_to_copy is True:
+            repos_to_copy = [
+                d for d in os.listdir(local_repos_path)
+                if os.path.isdir(os.path.join(local_repos_path, d))
+            ]
+        repo_files = [
+            os.path.join(local_repos_path, repo)
+            for repo in repos_to_copy
+            if os.path.isdir(os.path.join(local_repos_path, repo))
+        ]
+        if repo_files:
+            remote_inputs['files_to_copy'] = remote_inputs.get(
+                'files_to_copy', []) + repo_files
+            remote_mlc_repos_path = i.get("remote_mlc_repos_path", "MLC/repos")
+            remote_inputs['copy_directory'] = remote_mlc_repos_path
+            # On the remote, if MLC_REPOS is set and differs, symlink so
+            # mlcflow finds the copied repos
+            run_cmds.insert(0,
+                f'if [ -n "$MLC_REPOS" ] && [ "$MLC_REPOS" != "{remote_mlc_repos_path}" ]; then '
+                f'mkdir -p "{remote_mlc_repos_path}" && '
+                f'ln -sfn "$(realpath {remote_mlc_repos_path})"/* "$MLC_REPOS/"; '
+                f'fi')
+
     if files_to_copy_back:
         remote_inputs['files_to_copy_back'] = files_to_copy_back
 
@@ -171,6 +199,7 @@ def remote_run(self_module, i):
         'action': 'run', 'target': 'script', 'tags': 'remote,run,cmds,ssh',
         'script_tags': i.get('tags'), 'run_cmds': run_cmds,
         'pre_run_cmds': remote_pre_run_cmds,
+        'post_run_cmds': remote_post_run_cmds,
         **remote_inputs
     }
 

script/README.md

@@ -1,6 +1,6 @@
 # MLCommons Automation Scripts
 
-*Last updated: 2026-05-22 05:08:04*
+*Last updated: 2026-05-23 04:34:02*
 
 This directory contains automation scripts for MLPerf benchmarks, AI/ML workflows, and development operations.
 

script/remote-run-commands/customize.py

@@ -77,7 +77,9 @@ def preprocess(i):
 
     run_cmds = env.get('MLC_SSH_RUN_COMMANDS', [])
 
-    run_cmds = pre_run_cmds + run_cmds
+    post_run_cmds = env.get('MLC_SSH_POST_RUN_CMDS', [])
+
+    run_cmds = pre_run_cmds + run_cmds + post_run_cmds
 
     for i, cmd in enumerate(run_cmds):
         if 'cm ' in cmd:

script/remote-run-commands/meta.yaml

@@ -29,6 +29,7 @@ input_mapping:
   files_to_copy: MLC_SSH_FILES_TO_COPY
   files_to_copy_back: MLC_SSH_FILES_TO_COPY_BACK
   host: MLC_SSH_HOST
+  post_run_cmds: MLC_SSH_POST_RUN_CMDS
   pre_run_cmds: MLC_SSH_PRE_RUN_CMDS
   password: MLC_SSH_PASSWORD
   path_to_copy_back_files: MLC_SSH_PATH_TO_COPY_BACK_FILES