fix: @benchmark bot fixes (#565)

* Update ci-testing.yml

* update

* update

* Update .github/benchmark/benchmark.py

* feat: Add support for Makefile arguments in benchmark script and CI workflow

* feat: Enhance benchmark script and CI workflow to handle Makefile arguments and improve error reporting

* update

Author: deependujha

.github/benchmark/benchmark.py

@@ -18,22 +18,26 @@ class BenchmarkArgs:
     """Arguments for the LitData benchmark."""
 
     pr_number: int
+    branch: str
     org: Optional[str]
     user: Optional[str]
     teamspace: str
     machine: Machine
+    make_args: str
 
 
 def parse_args() -> BenchmarkArgs:
     """Parse command line arguments."""
     parser = argparse.ArgumentParser(description="Benchmark LitData in CI.")
     parser.add_argument("--pr", type=int, required=True, help="GitHub PR number")
+    parser.add_argument("--branch", type=str, required=True, help="GitHub PR branch name")
     parser.add_argument("--user", default=DEFAULT_USER, type=str, help="Lightning Studio username")
     parser.add_argument("--org", default=DEFAULT_ORG, type=str, help="Lightning Studio org")
     parser.add_argument("--teamspace", default=DEFAULT_TEAMSPACE, type=str, help="Lightning Studio teamspace")
     parser.add_argument(
         "--machine", type=str, default=DEFAULT_MACHINE, choices=["A10G", "T4", "CPU"], help="Machine type"
     )
+    parser.add_argument("--make-args", type=str, default="", help="Makefile variables passed to the script")
 
     args = parser.parse_args()
 
@@ -45,10 +49,12 @@ def parse_args() -> BenchmarkArgs:
 
     return BenchmarkArgs(
         pr_number=args.pr,
+        branch=args.branch,
         user=args.user,
         org=args.org,
         teamspace=args.teamspace,
         machine=machine_map[args.machine],
+        make_args=args.make_args,
     )
 
 
@@ -63,10 +69,12 @@ def __init__(self, config: BenchmarkArgs):
             raise ValueError("Only one of user or org must be provided.")
 
         self.pr = config.pr_number
+        self.branch = config.branch
         self.teamspace = config.teamspace
         self.user = config.user
         self.org = config.org
         self.machine = config.machine
+        self.make_args = config.make_args
         self.studio: Optional[Studio] = None
 
     def run(self) -> None:
@@ -88,15 +96,21 @@ def setup_studio(self) -> None:
             org=self.org,
         )
         self.studio.start(self.machine)
+        cmd = "rm -rf lit*(N)"
+        output, output_code = self.studio.run_with_exit_code(cmd)
+        if output_code != 0:
+            raise RuntimeError(f"Command failed:\n{cmd}\nExit code {output_code}:\n{output}")
 
     def setup_litdata_pr(self) -> None:
         """Set up the LitData PR in the studio."""
         assert self.studio is not None, "Studio is not set up"
+
+        # don't use `gh` cli, as it requires to login
         commands = [
-            "rm -rf lit*",
             "git clone https://github.com/Lightning-AI/litData.git",
             "cd litData",
-            f"gh pr checkout {self.pr}",
+            f"git fetch origin pull/{self.pr}/head:{self.branch}",
+            f"git checkout {self.branch}",
             "make setup",
         ]
         final_command = " && ".join(commands)
@@ -111,7 +125,7 @@ def setup_and_run_litdata_benchmark(self) -> None:
         commands = [
             "git clone https://github.com/bhimrazy/litdata-benchmark.git",
             "cd litdata-benchmark",
-            "make benchmark",
+            f"make benchmark {self.make_args}",
         ]
         final_command = " && ".join(commands)
         print(f"Running command: {final_command}")

.github/workflows/ci-benchmark.yml

@@ -11,7 +11,7 @@ permissions:
 jobs:
   reply-to-benchmark:
     if: |
-      github.event.comment.body == '@benchmark' &&
+      startsWith(github.event.comment.body, '@benchmark') &&
       github.event.issue.pull_request &&
       (
         github.event.comment.user.login == 'tchaton' ||
@@ -24,6 +24,16 @@ jobs:
     runs-on: ubuntu-latest
     steps:
       - uses: actions/checkout@v4
+      - name: Extract make args
+        id: extract
+        uses: actions/github-script@v7
+        with:
+          script: |
+            const fullComment = context.payload.comment.body;
+            const parts = fullComment.trim().split(/\s+/);
+            const makeargs = parts.slice(1).join(' '); // remove "@benchmark" and keep rest
+            core.setOutput("make_args", makeargs);
+
       - name: Reply to PR comment and save ID
         id: comment
         uses: actions/github-script@v7
@@ -35,6 +45,13 @@ jobs:
             const username = context.payload.comment.user.login;
             const prNumber = context.payload.issue.number;
 
+            const pr = await github.rest.pulls.get({
+              owner: context.repo.owner,
+              repo: context.repo.repo,
+              pull_number: context.payload.issue.number,
+            });
+            const branch = pr.data.head.ref;
+
             // Replace placeholders
             const reply = replyTemplate
               .replace(/{{username}}/g, username)
@@ -49,12 +66,19 @@ jobs:
 
             core.setOutput("comment_id", response.data.id);
             core.setOutput("prNumber", prNumber);
+            core.setOutput("branch", branch);
             core.setOutput("username", username);
 
       - name: run python
         run: |
           pip install -U lightning-sdk
-          python .github/benchmark/benchmark.py --pr ${{ steps.comment.outputs.prNumber }}
+          echo '"pr number: ${{ steps.comment.outputs.prNumber }}"; "branch: ${{ steps.comment.outputs.branch }}"'
+          echo "make-args: ${{steps.extract.outputs.make_args}}"
+          python .github/benchmark/benchmark.py \
+            --pr "${{ steps.comment.outputs.prNumber }}" \
+            --branch "${{ steps.comment.outputs.branch }}" \
+            --make-args "${{ steps.extract.outputs.make_args }}" \
+            2> error.txt || true
         env: # the following values are parsed from the repository secrets
           LIGHTNING_USER_ID: ${{ secrets.LIGHTNING_USER_ID }}
           LIGHTNING_API_KEY: ${{ secrets.LIGHTNING_API_KEY }}
@@ -67,11 +91,19 @@ jobs:
         with:
           script: |
             const fs = require('fs');
-            const path = 'result.md';
-            const reply = fs.readFileSync(path, 'utf8');
-
             const comment_id = Number(process.env.COMMENT_ID);
-            const updated_body = `Hi @${process.env.USERNAME}!\n\n- Benchmarking output:\n\n${reply}\n\n\ncc: @tchaton @deependujha @bhimrazy`;
+
+            let reply = '';
+            if (fs.existsSync('result.md')) {
+              reply = fs.readFileSync('result.md', 'utf8');
+            } else if (fs.existsSync('error.txt') && fs.readFileSync('error.txt', 'utf8').trim().length > 0) {
+              const err = fs.readFileSync('error.txt', 'utf8');
+              reply = `❌ **Benchmark failed**\n\n\`\`\`\n${err}\n\`\`\``;
+            } else {
+              reply = '❌ Benchmark completed, but can\'t find `result.md` and `error.txt` file. Something went wrong.';
+            }
+
+            const updated_body = `Hi @${process.env.USERNAME}!\n\n${reply}\n\ncc: @tchaton @deependujha @bhimrazy`;
 
             await github.rest.issues.updateComment({
               owner: context.repo.owner,

.github/workflows/ci-testing.yml

@@ -37,6 +37,7 @@ jobs:
       - name: Install package & dependencies
         run: |
           pip --version
+          pip install -U lightning-sdk
           pip install -e ".[extras]" -r requirements/test.txt -U -q --find-links $TORCH_URL
           pip list