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

[Aiden Grossman via llvm-commits]

https://github.com/boomanaiden154 updated https://github.com/llvm/llvm-project/pull/82137

>From 4f00bba6df676aa3db25088a528e781196646cbe Mon Sep 17 00:00:00 2001
From: Aiden Grossman <agrossman154 at yahoo.com>
Date: Sat, 17 Feb 2024 23:28:57 +0000
Subject: [PATCH] [Github] Add ability to record jobs over tim to job counting
 script

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.
---
 llvm/utils/count_running_jobs.py | 46 +++++++++++++++++++++++++++++++-
 1 file changed, 45 insertions(+), 1 deletion(-)

diff --git a/llvm/utils/count_running_jobs.py b/llvm/utils/count_running_jobs.py
index a53bc02f7b0ae1..90bb4cb9dcfe8c 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):
@@ -33,6 +35,8 @@ def main(token):
 
     print(f"\nFound {in_progress_jobs} running jobs.")
 
+    return in_progress_jobs
+
 
 if __name__ == "__main__":
     parser = argparse.ArgumentParser(
@@ -45,5 +49,45 @@ def main(token):
         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="?",
+    )
     args = parser.parse_args()
-    main(args.token)
+
+    # 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)
+
+            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)