Skip to content

Commit 8ffc4ff

Browse files
jaychiaJay Chia
and
Jay Chia
authored
ci: add a tool to launch workloads on cluster (#3516)
Example usage: ``` uv run tools/gha_run_cluster_job.py .github/working-dir/simple.py ``` <img width="1589" alt="image" src="https://github.com/user-attachments/assets/26319681-9846-4c37-a630-12b4193478d8"> --------- Co-authored-by: Jay Chia <[email protected]@users.noreply.github.com>
1 parent ad175ae commit 8ffc4ff

File tree

1 file changed

+88
-0
lines changed

1 file changed

+88
-0
lines changed

tools/gha_run_cluster_job.py

+88
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,88 @@
1+
# /// script
2+
# requires-python = ">=3.12"
3+
# dependencies = ["PyGithub"]
4+
# ///
5+
6+
import argparse
7+
import os
8+
import pathlib
9+
10+
from github import Auth, Github, enable_console_debug_logging
11+
12+
GITHUB_OAUTH_TOKEN_ENVVAR = "GITHUB_OAUTH_TOKEN"
13+
14+
15+
def get_oauth_token() -> str:
16+
token = os.getenv(GITHUB_OAUTH_TOKEN_ENVVAR)
17+
if token is None:
18+
raise RuntimeError(
19+
"Unable to find github token at $GITHUB_OAUTH_TOKEN. Please generate one and assign it to the environment variable. "
20+
"We recommend creating a Classic token (the fine-grained ones seem to not really work) and give it at least Workflow as well as Repository access. "
21+
"https://docs.github.com/en/authentication/keeping-your-account-and-data-secure/managing-your-personal-access-tokens"
22+
)
23+
return token
24+
25+
26+
def main(
27+
daft_wheel_url: str | None,
28+
daft_version: str,
29+
path_to_script: pathlib.Path,
30+
script_args: str | None,
31+
branch: str = "main",
32+
):
33+
auth = Auth.Token(get_oauth_token())
34+
g = Github(auth=auth)
35+
36+
daft_package = {"daft_wheel_url": daft_wheel_url} if daft_wheel_url is not None else {"daft_version": daft_version}
37+
entrypoint_args = {"entrypoint_args": script_args} if script_args is not None else {}
38+
39+
inputs = {
40+
**daft_package,
41+
"entrypoint_script": str(path_to_script.name),
42+
"working_dir": str(path_to_script.parent),
43+
**entrypoint_args,
44+
}
45+
46+
repo = g.get_repo("Eventual-Inc/Daft")
47+
workflow = repo.get_workflow("run-cluster.yaml")
48+
49+
created = workflow.create_dispatch(
50+
ref=branch,
51+
inputs={**inputs},
52+
)
53+
54+
if not created:
55+
raise RuntimeError("Could not create workflow, suggestion: run again with --verbose")
56+
57+
print("Workflow created, view it at: https://github.com/Eventual-Inc/Daft/actions/workflows/run-cluster.yaml")
58+
59+
60+
if __name__ == "__main__":
61+
parser = argparse.ArgumentParser(
62+
description="Run Daft cluster workflow on github actions",
63+
usage="uv run gha_run_cluster_job.py [-h] [--branch BRANCH] [--daft-wheel-url DAFT_WHEEL_URL | --daft-version DAFT_VERSION] path_to_script [-- script_args ...]",
64+
)
65+
parser.add_argument("path_to_script", type=pathlib.Path, help="Path to the script to run")
66+
parser.add_argument("--branch", default="main", help="GitHub branch to use")
67+
group = parser.add_mutually_exclusive_group(required=False)
68+
group.add_argument("--daft-wheel-url", help="URL of the Daft wheel")
69+
group.add_argument("--daft-version", default="0.3.15", help="Released version of daft (default: 0.3.15)")
70+
parser.add_argument("--verbose", action="store_true", help="Enable verbose mode")
71+
args, extra = parser.parse_known_args()
72+
73+
if args.verbose:
74+
enable_console_debug_logging()
75+
76+
script_args = None
77+
if extra:
78+
if extra[0] != "--":
79+
raise ValueError(f"Arguments to the script must be passed in after a '--', received: {' '.join(extra)}")
80+
script_args = " ".join(extra[1:])
81+
82+
main(
83+
daft_wheel_url=args.daft_wheel_url,
84+
daft_version=args.daft_version,
85+
path_to_script=args.path_to_script,
86+
script_args=script_args,
87+
branch=args.branch,
88+
)

0 commit comments

Comments
 (0)