Skip to content

Commit

Permalink
[BE] Improve logging for runner-determinator (pytorch#129679)
Browse files Browse the repository at this point in the history 
This lets us be more flexible about what data we output and throwing exceptions. It's also less likely to break when others make changes (e.g. any print statement would have broken this code before since the printed output was expected to only be a json)
Pull Request resolved: pytorch#129679
Approved by: https://github.com/zxiiro, https://github.com/jeanschmidt, https://github.com/Skylion007
  • Loading branch information
@ZainRizvi @pytorchmergebot
ZainRizvi authored and pytorchmergebot committed Jul 1, 2024
1 parent eeef686 commit 9645eaa
Show file tree
Hide file tree
Showing 2 changed files with 154 additions and 73 deletions.
109 changes: 76 additions & 33 deletions .github/scripts/runner_determinator.py
View file
Original file line number Diff line number Diff line change
@@ -1,16 +1,62 @@
import json
# flake8: noqa: G004

import logging
import os
from argparse import ArgumentParser
from typing import Any, Iterable, Tuple
from logging import LogRecord
from typing import Any, Iterable

from github import Auth, Github
from github.Issue import Issue


WORKFLOW_LABEL_META = "" # use meta runners
WORKFLOW_LABEL_LF = "lf." # use runners from the linux foundation
LABEL_TYPE_KEY = "label_type"
MESSAGE_KEY = "message"
MESSAGE = "" # Debug message to return to the caller

GITHUB_OUTPUT = os.getenv("GITHUB_OUTPUT", "")
GH_OUTPUT_KEY_LABEL_TYPE = "label-type"


class ColorFormatter(logging.Formatter):
"""Color codes the log messages based on the log level"""

COLORS = {
"WARNING": "\033[33m", # Yellow
"ERROR": "\033[31m", # Red
"CRITICAL": "\033[31m", # Red
"INFO": "\033[0m", # Reset
"DEBUG": "\033[0m", # Reset
}

def format(self, record: LogRecord) -> str:
log_color = self.COLORS.get(record.levelname, "\033[0m") # Default to reset
record.msg = f"{log_color}{record.msg}\033[0m"
return super().format(record)


handler = logging.StreamHandler()
handler.setFormatter(ColorFormatter(fmt="%(levelname)-8s: %(message)s"))

log = logging.getLogger(os.path.basename(__file__))
log.addHandler(handler)
log.setLevel(logging.INFO)


def set_github_output(key: str, value: str) -> None:
"""
Defines outputs of the github action that invokes this script
"""
if not GITHUB_OUTPUT:
# See https://github.blog/changelog/2022-10-11-github-actions-deprecating-save-state-and-set-output-commands/ for deprecation notice
log.warning(
"No env var found for GITHUB_OUTPUT, you must be running this code locally. Falling back to the deprecated print method."
)
print(f"::set-output name={key}::{value}")
return

with open(GITHUB_OUTPUT, "a") as f:
log.info(f"Setting output: {key}='{value}'")
f.write(f"{key}={value}\n")


def parse_args() -> Any:
Expand Down Expand Up @@ -91,18 +137,16 @@ def is_exception_branch(branch: str) -> bool:
return branch.split("/")[0] in {"main", "nightly", "release", "landchecks"}


def get_workflow_type(
issue: Issue, workflow_requestors: Iterable[str]
) -> Tuple[str, str]:
def get_workflow_type(issue: Issue, workflow_requestors: Iterable[str]) -> str:
try:
first_comment = issue.get_comments()[0].body.strip("\n\t ")

if first_comment[0] == "!":
MESSAGE = "LF Workflows are disabled for everyone. Using meta runners."
return WORKFLOW_LABEL_META, MESSAGE
log.info("LF Workflows are disabled for everyone. Using meta runners.")
return WORKFLOW_LABEL_META
elif first_comment[0] == "*":
MESSAGE = "LF Workflows are enabled for everyone. Using LF runners."
return WORKFLOW_LABEL_LF, MESSAGE
log.info("LF Workflows are enabled for everyone. Using LF runners.")
return WORKFLOW_LABEL_LF
else:
all_opted_in_users = {
usr_raw.strip("\n\t@ ") for usr_raw in first_comment.split()
Expand All @@ -111,25 +155,29 @@ def get_workflow_type(
usr for usr in workflow_requestors if usr in all_opted_in_users
}
if opted_in_requestors:
MESSAGE = f"LF Workflows are enabled for {', '.join(opted_in_requestors)}. Using LF runners."
return WORKFLOW_LABEL_LF, MESSAGE
log.info(
f"LF Workflows are enabled for {', '.join(opted_in_requestors)}. Using LF runners."
)
return WORKFLOW_LABEL_LF
else:
MESSAGE = f"LF Workflows are disabled for {', '.join(workflow_requestors)}. Using meta runners."
return WORKFLOW_LABEL_META, MESSAGE
log.info(
f"LF Workflows are disabled for {', '.join(workflow_requestors)}. Using meta runners."
)
return WORKFLOW_LABEL_META

except Exception as e:
MESSAGE = f"Failed to get determine workflow type. Falling back to meta runners. Exception: {e}"
return WORKFLOW_LABEL_META, MESSAGE
log.error(
f"Failed to get determine workflow type. Falling back to meta runners. Exception: {e}"
)
return WORKFLOW_LABEL_META


def main() -> None:
args = parse_args()

if args.github_ref_type == "branch" and is_exception_branch(args.github_branch):
output = {
LABEL_TYPE_KEY: WORKFLOW_LABEL_META,
MESSAGE_KEY: f"Exception branch: '{args.github_branch}', using meta runners",
}
log.info(f"Exception branch: '{args.github_branch}', using meta runners")
label_type = WORKFLOW_LABEL_META
else:
try:
gh = get_gh_client(args.github_token)
Expand All @@ -142,25 +190,20 @@ def main() -> None:
args.github_ref_type,
args.github_branch,
)
label_type, message = get_workflow_type(
label_type = get_workflow_type(
issue,
(
args.github_issue_owner,
username,
),
)
output = {
LABEL_TYPE_KEY: label_type,
MESSAGE_KEY: message,
}
except Exception as e:
output = {
LABEL_TYPE_KEY: WORKFLOW_LABEL_META,
MESSAGE_KEY: f"Failed to get issue. Falling back to meta runners. Exception: {e}",
}
log.error(
f"Failed to get issue. Falling back to meta runners. Exception: {e}"
)
label_type = WORKFLOW_LABEL_META

json_output = json.dumps(output)
print(json_output)
set_github_output(GH_OUTPUT_KEY_LABEL_TYPE, label_type)


if __name__ == "__main__":
Expand Down
118 changes: 78 additions & 40 deletions .github/workflows/_runner-determinator.yml
View file
Original file line number Diff line number Diff line change
Expand Up @@ -56,19 +56,65 @@ jobs:
- name: Hardcode runner-determinator script
run: |
cat <<EOF > runner_determinator.py
import json
# flake8: noqa: G004
import logging
import os
from argparse import ArgumentParser
from typing import Any, Iterable, Tuple
from logging import LogRecord
from typing import Any, Iterable
from github import Auth, Github
from github.Issue import Issue
WORKFLOW_LABEL_META = "" # use meta runners
WORKFLOW_LABEL_LF = "lf." # use runners from the linux foundation
LABEL_TYPE_KEY = "label_type"
MESSAGE_KEY = "message"
MESSAGE = "" # Debug message to return to the caller
GITHUB_OUTPUT = os.getenv("GITHUB_OUTPUT", "")
GH_OUTPUT_KEY_LABEL_TYPE = "label-type"
class ColorFormatter(logging.Formatter):
"""Color codes the log messages based on the log level"""
COLORS = {
"WARNING": "\033[33m", # Yellow
"ERROR": "\033[31m", # Red
"CRITICAL": "\033[31m", # Red
"INFO": "\033[0m", # Reset
"DEBUG": "\033[0m", # Reset
}
def format(self, record: LogRecord) -> str:
log_color = self.COLORS.get(record.levelname, "\033[0m") # Default to reset
record.msg = f"{log_color}{record.msg}\033[0m"
return super().format(record)
handler = logging.StreamHandler()
handler.setFormatter(ColorFormatter(fmt="%(levelname)-8s: %(message)s"))
log = logging.getLogger(os.path.basename(__file__))
log.addHandler(handler)
log.setLevel(logging.INFO)
def set_github_output(key: str, value: str) -> None:
"""
Defines outputs of the github action that invokes this script
"""
if not GITHUB_OUTPUT:
# See https://github.blog/changelog/2022-10-11-github-actions-deprecating-save-state-and-set-output-commands/ for deprecation notice
log.warning(
"No env var found for GITHUB_OUTPUT, you must be running this code locally. Falling back to the deprecated print method."
)
print(f"::set-output name={key}::{value}")
return
with open(GITHUB_OUTPUT, "a") as f:
log.info(f"Setting output: {key}='{value}'")
f.write(f"{key}={value}\n")
def parse_args() -> Any:
Expand Down Expand Up @@ -149,18 +195,16 @@ jobs:
return branch.split("/")[0] in {"main", "nightly", "release", "landchecks"}
def get_workflow_type(
issue: Issue, workflow_requestors: Iterable[str]
) -> Tuple[str, str]:
def get_workflow_type(issue: Issue, workflow_requestors: Iterable[str]) -> str:
try:
first_comment = issue.get_comments()[0].body.strip("\n\t ")
if first_comment[0] == "!":
MESSAGE = "LF Workflows are disabled for everyone. Using meta runners."
return WORKFLOW_LABEL_META, MESSAGE
log.info("LF Workflows are disabled for everyone. Using meta runners.")
return WORKFLOW_LABEL_META
elif first_comment[0] == "*":
MESSAGE = "LF Workflows are enabled for everyone. Using LF runners."
return WORKFLOW_LABEL_LF, MESSAGE
log.info("LF Workflows are enabled for everyone. Using LF runners.")
return WORKFLOW_LABEL_LF
else:
all_opted_in_users = {
usr_raw.strip("\n\t@ ") for usr_raw in first_comment.split()
Expand All @@ -169,25 +213,29 @@ jobs:
usr for usr in workflow_requestors if usr in all_opted_in_users
}
if opted_in_requestors:
MESSAGE = f"LF Workflows are enabled for {', '.join(opted_in_requestors)}. Using LF runners."
return WORKFLOW_LABEL_LF, MESSAGE
log.info(
f"LF Workflows are enabled for {', '.join(opted_in_requestors)}. Using LF runners."
)
return WORKFLOW_LABEL_LF
else:
MESSAGE = f"LF Workflows are disabled for {', '.join(workflow_requestors)}. Using meta runners."
return WORKFLOW_LABEL_META, MESSAGE
log.info(
f"LF Workflows are disabled for {', '.join(workflow_requestors)}. Using meta runners."
)
return WORKFLOW_LABEL_META
except Exception as e:
MESSAGE = f"Failed to get determine workflow type. Falling back to meta runners. Exception: {e}"
return WORKFLOW_LABEL_META, MESSAGE
log.error(
f"Failed to get determine workflow type. Falling back to meta runners. Exception: {e}"
)
return WORKFLOW_LABEL_META
def main() -> None:
args = parse_args()
if args.github_ref_type == "branch" and is_exception_branch(args.github_branch):
output = {
LABEL_TYPE_KEY: WORKFLOW_LABEL_META,
MESSAGE_KEY: f"Exception branch: '{args.github_branch}', using meta runners",
}
log.info(f"Exception branch: '{args.github_branch}', using meta runners")
label_type = WORKFLOW_LABEL_META
else:
try:
gh = get_gh_client(args.github_token)
Expand All @@ -200,25 +248,20 @@ jobs:
args.github_ref_type,
args.github_branch,
)
label_type, message = get_workflow_type(
label_type = get_workflow_type(
issue,
(
args.github_issue_owner,
username,
),
)
output = {
LABEL_TYPE_KEY: label_type,
MESSAGE_KEY: message,
}
except Exception as e:
output = {
LABEL_TYPE_KEY: WORKFLOW_LABEL_META,
MESSAGE_KEY: f"Failed to get issue. Falling back to meta runners. Exception: {e}",
}
log.error(
f"Failed to get issue. Falling back to meta runners. Exception: {e}"
)
label_type = WORKFLOW_LABEL_META
json_output = json.dumps(output)
print(json_output)
set_github_output(GH_OUTPUT_KEY_LABEL_TYPE, label_type)
if __name__ == "__main__":
Expand All @@ -237,16 +280,11 @@ jobs:
curr_ref_type="${{ inputs.curr_ref_type }}"
echo "Current branch is '$curr_branch'"
output="$(python3 runner_determinator.py \
python3 runner_determinator.py \
--github-token "$GITHUB_TOKEN" \
--github-issue "$ISSUE_NUMBER" \
--github-branch "$curr_branch" \
--github-actor "$TRIGGERING_ACTOR" \
--github-issue-owner "$ISSUE_OWNER" \
--github-ref-type "$curr_ref_type" \
--github-repo "$GITHUB_REPOSITORY")"
echo "Output: '${output}'"
LABEL_TYPE=$(echo "${output}" | jq -r '.label_type')
echo "label-type=$LABEL_TYPE" >> "$GITHUB_OUTPUT"
--github-repo "$GITHUB_REPOSITORY"

0 comments on commit 9645eaa

Please sign in to comment.