Added caps management

main_backend_lighteval.py

@@ -6,7 +6,7 @@ from huggingface_hub import snapshot_download
 logging.getLogger("openai").setLevel(logging.WARNING)
 
 from src.backend.run_eval_suite_lighteval import run_evaluation
-from src.backend.manage_requests import check_completed_evals, get_eval_requests, set_eval_request
+from src.backend.manage_requests import check_completed_evals, get_eval_requests, set_eval_request, set_requests_seen
 from src.backend.sort_queue import sort_models_by_priority
 
 from src.envs import QUEUE_REPO, EVAL_REQUESTS_PATH_BACKEND, RESULTS_REPO, EVAL_RESULTS_PATH_BACKEND, API, LIMIT, TOKEN, ACCELERATOR, VENDOR, REGION
@@ -40,7 +40,7 @@ def run_auto_eval():
     )
 
     # Get all eval request that are PENDING, if you want to run other evals, change this parameter
-    eval_requests = get_eval_requests(job_status=current_pending_status, hf_repo=QUEUE_REPO, local_dir=EVAL_REQUESTS_PATH_BACKEND)
+    eval_requests, requests_seen = get_eval_requests(job_status=current_pending_status, hf_repo=QUEUE_REPO, local_dir=EVAL_REQUESTS_PATH_BACKEND)
     # Sort the evals by priority (first submitted first run)
     eval_requests = sort_models_by_priority(api=API, models=eval_requests)
 
@@ -52,29 +52,68 @@ def run_auto_eval():
     eval_request = eval_requests[0]
     pp.pprint(eval_request)
 
-    set_eval_request(
-        api=API,
-        eval_request=eval_request,
-        set_to_status=RUNNING_STATUS,
-        hf_repo=QUEUE_REPO,
-        local_dir=EVAL_REQUESTS_PATH_BACKEND,
-    )
-
-    # This needs to be done
-    #instance_size, instance_type = get_instance_for_model(eval_request)
     # For GPU
     if not eval_request or eval_request.params < 0:
         raise ValueError("Couldn't detect number of params, please make sure the metadata is available")
     elif eval_request.params < 4:
-        instance_size, instance_type = "small", "g4dn.xlarge" 
+        instance_size, instance_type, cap = "small", "g4dn.xlarge", 20
     elif eval_request.params < 9:
-        instance_size, instance_type = "medium", "g5.2xlarge" 
+        instance_size, instance_type, cap = "medium", "g5.2xlarge", 35
     elif eval_request.params < 24:
-        instance_size, instance_type = "xxlarge", "g5.12xlarge"
+        instance_size, instance_type, cap = "xxlarge", "g5.12xlarge", 15
     else:
         raise ValueError("Number of params too big, can't run this model")
-    # For CPU
-    # instance_size, instance_type = "medium", "c6i"
+    
+    counter_key = f'count_{instance_type}'
+    if not counter_key in requests_seen:
+        requests_seen[counter_key] = 0
+    if requests_seen[counter_key] >= cap:
+        set_eval_request(
+            api=API,
+            eval_request=eval_request,
+            set_to_status=FAILED_STATUS,
+            hf_repo=QUEUE_REPO,
+            local_dir=EVAL_REQUESTS_PATH_BACKEND,
+        )
+        pp.pprint(dict(message="Reached maximum cap for requests of this instance type this month", counter=counter_key, instance_type=instance_type, cap=cap))
+        return
+
+    # next, check to see who made the last commit to this repo - keep track of that. One person shouldn't commit more 
+    # than 4 models in one month. 
+    commits = API.list_repo_commits(eval_request.model, revision=eval_request.revision)
+    users = commits[0].authors
+    for user in users:
+        if user in requests_seen and len(requests_seen[user]) >= 4:
+            set_eval_request(
+                api=API,
+                eval_request=eval_request,
+                set_to_status=FAILED_STATUS,
+                hf_repo=QUEUE_REPO,
+                local_dir=EVAL_REQUESTS_PATH_BACKEND,
+            )
+            pp.pprint(dict(message="Reached maximum cap for requests for this user this month", counter=counter_key, user=user))
+            return
+        if not user in requests_seen:
+            requests_seen[user] = []
+        requests_seen[user].append(dict(model_id=eval_request.model, revision=eval_request.revision))
+        
+    requests_seen[counter_key] += 1
+    set_requests_seen(
+        api=API,
+        requests_seen=requests_seen,
+        hf_repo=QUEUE_REPO,
+        local_dir=EVAL_REQUESTS_PATH_BACKEND
+    )
+
+    set_eval_request(
+        api=API,
+        eval_request=eval_request,
+        set_to_status=RUNNING_STATUS,
+        hf_repo=QUEUE_REPO,
+        local_dir=EVAL_REQUESTS_PATH_BACKEND,
+    )
+
+    
 
     run_evaluation(
         eval_request=eval_request, 

src/backend/manage_requests.py

@@ -1,7 +1,9 @@
 import glob
 import json
 from dataclasses import dataclass
+import os
 from typing import Optional
+import datetime
 
 from huggingface_hub import HfApi, snapshot_download
 from src.envs import TOKEN
@@ -42,6 +44,18 @@ class EvalRequest:
         
         return model_args
 
+def set_requests_seen(api: HfApi, requests_seen: dict, hf_repo: str, local_dir: str):
+    """Updates a given eval request with its new status on the hub (running, completed, failed, ...)"""
+    json_filepath = get_requests_seen_json_file(local_dir)
+    with open(json_filepath, "w") as f:
+        f.write(json.dumps(requests_seen))
+
+    api.upload_file(
+        path_or_fileobj=json_filepath,
+        path_in_repo=json_filepath.replace(local_dir, ""),
+        repo_id=hf_repo,
+        repo_type="dataset",
+    )
 
 def set_eval_request(api: HfApi, eval_request: EvalRequest, set_to_status: str, hf_repo: str, local_dir: str):
     """Updates a given eval request with its new status on the hub (running, completed, failed, ...)"""
@@ -62,7 +76,6 @@ def set_eval_request(api: HfApi, eval_request: EvalRequest, set_to_status: str,
         repo_type="dataset",
     )
 
-
 def get_eval_requests(job_status: list, local_dir: str, hf_repo: str) -> list[EvalRequest]:
     """Get all pending evaluation requests and return a list in which private
     models appearing first, followed by public models sorted by the number of
@@ -83,8 +96,16 @@ def get_eval_requests(job_status: list, local_dir: str, hf_repo: str) -> list[Ev
             eval_request = EvalRequest(**data)
             eval_requests.append(eval_request)
 
-    return eval_requests
+    requests_seen_json = get_requests_seen_json_file(local_dir)
+    requests_seen = {}
+    if os.path.isfile(requests_seen_json):
+        with open(requests_seen_json, 'r', encoding='utf8') as r:
+            requests_seen = json.load(r)
+
+    return eval_requests, requests_seen
 
+def get_requests_seen_json_file(local_dir):
+    return f"{local_dir}/counters/{datetime.datetime.now().strftime("%B-%Y")}.json"
 
 def check_completed_evals(
     api: HfApi,
@@ -99,7 +120,7 @@ def check_completed_evals(
     """Checks if the currently running evals are completed, if yes, update their status on the hub."""
     snapshot_download(repo_id=hf_repo_results, revision="main", local_dir=local_dir_results, repo_type="dataset", max_workers=60, token=TOKEN)
 
-    running_evals = get_eval_requests(checked_status, hf_repo=hf_repo, local_dir=local_dir)
+    running_evals, _ = get_eval_requests(checked_status, hf_repo=hf_repo, local_dir=local_dir)
 
     for eval_request in running_evals:
         model = eval_request.model