compress github action payload (#290)

* compress github action payload

* allow explicit specification of gh branch

* allow backslashes in code content

Author: ngc92

.github/workflows/runner.py

@@ -1,5 +1,7 @@
+import base64
 import json
 import sys
+import zlib
 from dataclasses import asdict
 from datetime import datetime
 from pathlib import Path
@@ -8,8 +10,10 @@
 
 from run_eval import run_config
 
-config = json.loads(Path("payload.json").read_text())
+payload = Path("payload.json").read_text()
 Path("payload.json").unlink()
+payload = zlib.decompress(base64.b64decode(payload)).decode("utf-8")
+config = json.loads(payload)
 
 result = asdict(run_config(config))
 

src/discord-cluster-manager/bot.py

@@ -79,7 +79,9 @@ async def setup_hook(self):
             # Load cogs
             submit_cog = SubmitCog(self)
             submit_cog.register_launcher(ModalLauncher(consts.MODAL_CUDA_INCLUDE_DIRS))
-            submit_cog.register_launcher(GitHubLauncher(env.GITHUB_REPO, env.GITHUB_TOKEN))
+            submit_cog.register_launcher(
+                GitHubLauncher(env.GITHUB_REPO, env.GITHUB_TOKEN, env.GITHUB_WORKFLOW_BRANCH)
+            )
             await self.add_cog(submit_cog)
             await self.add_cog(BotManagerCog(self))
             await self.add_cog(LeaderboardCog(self))

src/discord-cluster-manager/env.py

@@ -1,6 +1,7 @@
 import os
 
 from dotenv import load_dotenv
+from utils import get_github_branch_name
 
 
 def init_environment():
@@ -33,6 +34,7 @@ def init_environment():
 # GitHub-specific constants
 GITHUB_TOKEN = os.getenv("GITHUB_TOKEN")
 GITHUB_REPO = os.getenv("GITHUB_REPO")
+GITHUB_WORKFLOW_BRANCH = os.getenv("GITHUB_WORKFLOW_BRANCH", get_github_branch_name())
 PROBLEMS_REPO = os.getenv("PROBLEMS_REPO")
 
 # Directory that will be used for local problem development.

src/discord-cluster-manager/launchers/github.py

@@ -1,10 +1,12 @@
 import asyncio
+import base64
 import datetime
 import json
 import math
 import pprint
 import tempfile
 import zipfile
+import zlib
 from typing import Awaitable, Callable, Optional
 
 import requests
@@ -20,7 +22,7 @@
 from github import Github, UnknownObjectException, WorkflowRun
 from report import RunProgressReporter
 from run_eval import CompileResult, EvalResult, FullResult, RunResult, SystemInfo
-from utils import get_github_branch_name, setup_logging
+from utils import setup_logging
 
 from .launcher import Launcher
 
@@ -39,10 +41,11 @@ def get_timeout(config: dict) -> int:
 
 
 class GitHubLauncher(Launcher):
-    def __init__(self, repo: str, token: str):
+    def __init__(self, repo: str, token: str, branch: str):
         super().__init__(name="GitHub", gpus=GitHubGPU)
         self.repo = repo
         self.token = token
+        self.branch = branch
         self.trigger_limit = asyncio.Semaphore(1)
 
     async def run_submission(
@@ -71,10 +74,12 @@ async def run_submission(
         lang_name = {"py": "Python", "cu": "CUDA"}[lang]
 
         logger.info(f"Attempting to trigger GitHub action for {lang_name} on {selected_workflow}")
-        run = GitHubRun(self.repo, self.token, selected_workflow)
+        run = GitHubRun(self.repo, self.token, self.branch, selected_workflow)
         logger.info(f"Successfully created GitHub run: {run.run_id}")
 
-        payload = json.dumps(config)
+        payload = base64.b64encode(zlib.compress(json.dumps(config).encode("utf-8"))).decode(
+            "utf-8"
+        )
 
         inputs = {"payload": payload}
         if lang == "py":
@@ -143,10 +148,11 @@ async def wait_callback(self, run: "GitHubRun", status: RunProgressReporter):
 
 
 class GitHubRun:
-    def __init__(self, repo: str, token: str, workflow_file: str):
+    def __init__(self, repo: str, token: str, branch: str, workflow_file: str):
         gh = Github(token)
         self.repo = gh.get_repo(repo)
         self.token = token
+        self.branch = branch
         self.workflow_file = workflow_file
         self.run: Optional[WorkflowRun.WorkflowRun] = None
         self.start_time = None
@@ -189,14 +195,14 @@ async def trigger(self, inputs: dict) -> bool:
             logger.error(f"Could not find workflow {self.workflow_file}", exc_info=e)
             raise ValueError(f"Could not find workflow {self.workflow_file}") from e
 
-        branch_name = get_github_branch_name()
+        logger.info("Dispatching workflow %s on branch %s", self.workflow_file, self.branch)
         logger.debug(
             "Dispatching workflow %s on branch %s with inputs %s",
             self.workflow_file,
-            branch_name,
+            self.branch,
             pprint.pformat(inputs),
         )
-        success = await asyncio.to_thread(workflow.create_dispatch, branch_name, inputs=inputs)
+        success = await asyncio.to_thread(workflow.create_dispatch, self.branch, inputs=inputs)
 
         if success:
             wait_seconds = 5
@@ -248,7 +254,7 @@ async def trigger(self, inputs: dict) -> bool:
                 return False
         else:
             logger.error(
-                f"Failed to dispatch workflow {self.workflow_file} on branch {branch_name}."
+                f"Failed to dispatch workflow {self.workflow_file} on branch {self.branch}."
             )
             return False
 

src/discord-cluster-manager/leaderboard_db.py

@@ -226,9 +226,9 @@ def create_submission(
                 """
                 SELECT id, code
                 FROM leaderboard.code_files
-                WHERE hash = encode(sha256(%s::bytea), 'hex')
+                WHERE hash = encode(sha256(%s), 'hex')
                 """,
-                (code,),
+                (code.encode("utf-8"),),
             )
 
             code_id = None
@@ -245,7 +245,7 @@ def create_submission(
                     VALUES (%s)
                     RETURNING id
                     """,
-                    (code,),
+                    (code.encode("utf-8"),),
                 )
                 code_id = self.cursor.fetchone()
             # Check if user exists in user_info, if not add them