[llvm] 9c6df7d - [Github] Add ability to record jobs over time to job counting script (#82137)

[via llvm-commits]

Author: Aiden Grossman
Date: 2024-02-19T17:00:21-08:00
New Revision: 9c6df7d90a74c3b98e33d3edf0b5d8982093b4f7

URL: https://github.com/llvm/llvm-project/commit/9c6df7d90a74c3b98e33d3edf0b5d8982093b4f7
DIFF: https://github.com/llvm/llvm-project/commit/9c6df7d90a74c3b98e33d3edf0b5d8982093b4f7.diff

LOG: [Github] Add ability to record jobs over time to job counting script (#82137)

This patch adds new flags to the count_running_jobs.py script to enable
the collection of data over a period of time. Specifically, the
--output_file flag is added to enable writing output data to a file, and
the --data-collection-interval flag is added to configure the frequency
that the script checks the job count.

Added: 
    

Modified: 
    llvm/utils/count_running_jobs.py

Removed: 
    


################################################################################
diff  --git a/llvm/utils/count_running_jobs.py b/llvm/utils/count_running_jobs.py
index dee822a811258e..d4fce9a03dcbe6 100644
--- a/llvm/utils/count_running_jobs.py
+++ b/llvm/utils/count_running_jobs.py
@@ -14,6 +14,8 @@
 
 import argparse
 import github
+import sys
+import time
 
 
 def main(token, filter_gha_runners):
@@ -44,6 +46,8 @@ def main(token, filter_gha_runners):
 
     print(f"\nFound {in_progress_jobs} running jobs.")
 
+    return in_progress_jobs
+
 
 if __name__ == "__main__":
     parser = argparse.ArgumentParser(
@@ -56,6 +60,20 @@ def main(token, filter_gha_runners):
         default=None,
         nargs="?",
     )
+    parser.add_argument(
+        "--output-file",
+        type=str,
+        help="The output file to write time-series data to",
+        default=None,
+        nargs="?",
+    )
+    parser.add_argument(
+        "--data-collection-interval",
+        type=int,
+        help="The number of seconds between data collection intervals",
+        default=None,
+        nargs="?",
+    )
     parser.add_argument(
         "--filter-gha-runners",
         help="Only consider jobs running on hosted Github actions runners",
@@ -68,6 +86,31 @@ def main(token, filter_gha_runners):
         help="Consider all running jobs",
     )
     parser.set_defaults(filter_gha_runners=False)
-
     args = parser.parse_args()
-    main(args.token, args.filter_gha_runners)
+
+    # Perform some basic argument validation
+
+    # If an output file is specified, the user must also specify the data
+    # collection interval.
+    if bool(args.output_file) and not bool(args.data_collection_interval):
+        print("A data collection interval must be specified when --output_file is used")
+        sys.exit(1)
+
+    if args.data_collection_interval:
+        while True:
+            current_time = time.localtime()
+            current_time_string = time.strftime("%Y/%m/%d %H:%M:%S", current_time)
+
+            print(f"Collecting data at {current_time_string}")
+
+            current_job_count = main(args.token, args.filter_gha_runners)
+
+            if args.output_file:
+                with open(args.output_file, "a") as output_file_handle:
+                    output_file_handle.write(
+                        f"{current_time_string},{current_job_count}\n"
+                    )
+
+            time.sleep(args.data_collection_interval)
+    else:
+        main(args.token, args.filter_gha_runners)