From 9645eaaaecafd50abb5149eaa3b3ba98327b0a85 Mon Sep 17 00:00:00 2001 From: Zain Rizvi Date: Mon, 1 Jul 2024 22:31:35 +0000 Subject: [PATCH] [BE] Improve logging for runner-determinator (#129679) 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: https://github.com/pytorch/pytorch/pull/129679 Approved by: https://github.com/zxiiro, https://github.com/jeanschmidt, https://github.com/Skylion007 --- .github/scripts/runner_determinator.py | 109 +++++++++++++------ .github/workflows/_runner-determinator.yml | 118 ++++++++++++++------- 2 files changed, 154 insertions(+), 73 deletions(-) diff --git a/.github/scripts/runner_determinator.py b/.github/scripts/runner_determinator.py index ffa718c5bc506..b9878ada36b51 100644 --- a/.github/scripts/runner_determinator.py +++ b/.github/scripts/runner_determinator.py @@ -1,6 +1,10 @@ -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 @@ -8,9 +12,51 @@ 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: @@ -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() @@ -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) @@ -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__": diff --git a/.github/workflows/_runner-determinator.yml b/.github/workflows/_runner-determinator.yml index 4e5bcb1fd4c5b..8c44b57b320fe 100644 --- a/.github/workflows/_runner-determinator.yml +++ b/.github/workflows/_runner-determinator.yml @@ -56,9 +56,13 @@ jobs: - name: Hardcode runner-determinator script run: | cat < 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 @@ -66,9 +70,51 @@ jobs: 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: @@ -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() @@ -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) @@ -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__": @@ -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"