Auto wandb tag with benchmark.py

cleanrl_utils/benchmark.py

@@ -1,7 +1,10 @@
 import argparse
+import os
 import shlex
 import subprocess
 
+import requests
+
 
 def parse_args():
     # fmt: off
@@ -16,6 +19,8 @@ def parse_args():
         help="the number of random seeds")
     parser.add_argument('--workers', type=int, default=0,
         help='the number of eval workers to run benchmark experimenets (skips evaluation when set to 0)')
+    parser.add_argument("--auto-tag", type=lambda x: bool(strtobool(x)), default=True, nargs="?", const=True,
+        help="if toggled, the runs will be tagged with the output from `git describe --tags` (e.g., v1.0.0b2-11-g5db4db7)")
     args = parser.parse_args()
     # fmt: on
     return args
@@ -29,8 +34,35 @@ def run_experiment(command: str):
     assert return_code == 0
 
 
+def autotag() -> str:
+    git_tag = subprocess.check_output(["git", "describe", "--tags"]).decode("ascii").strip()
+    git_commit = subprocess.check_output(["git", "rev-parse", "--verify", "HEAD"]).decode("ascii").strip()
+
+    # try finding the pull request number on github
+    prs = requests.get(f"https://api.github.com/repos/vwxyzjn/cleanrl/commits/{git_commit}/pulls")
+    if prs.status_code == 200:
+        prs = prs.json()
+        if len(prs) > 0:
+            pr = prs[0]
+            pr_number = pr["number"]
+            pr["title"]
+            pr["html_url"]
+            wandb_tag = f"{git_tag},pr{pr_number}"
+        else:
+            wandb_tag = f"{git_tag}"
+    return wandb_tag
+
+
 if __name__ == "__main__":
     args = parse_args()
+    if args.auto_tag:
+        if "WANDB_TAGS" in os.environ:
+            raise ValueError(
+                "WANDB_TAGS is already set. Please unset it before running this script or run the script with --auto-tag False"
+            )
+        wandb_tag = autotag()
+        os.environ["WANDB_TAGS"] = wandb_tag
+
     commands = []
     for seed in range(1, args.num_seeds + 1):
         for env_id in args.env_ids: