Merge pull request #635 from haozturk/update_ddm_quota

Make ddm_quota calculation quadratic and normalized

docker/CMSRucioClient/scripts/updateDDMQuota

@@ -1,6 +1,8 @@
 #! /usr/bin/env python3
 
 from rucio.client import Client
+import numpy as np
+import pprint
 client = Client()
 
 # ddm_quota are weight given to RSEs based on the amount of free space
@@ -10,43 +12,115 @@ client = Client()
 # NOTE: This probably needs to be reviewed after an assement of age of dynamic data at different sites
 # and if this can be used to normalise that
 
+
 DRY_RUN = False
+VERBOSE = False
+MAKE_QUADRATIC = False
+MAX_DDM_QUOTA = 1000
+
+STATIC_WEIGHT = 0.2
+FREE_WEIGHT = 0.5
+EXPIRED_WEIGHT = 0.3
 
 # Adding ddm_quota attribute to all disk RSEs
 # T3s do not have the "static" usage set, they are quasi-static
 RSE_EXPRESSION = "rse_type=DISK&cms_type=real&tier<3&tier>0"
-
 rses = [rse["rse"] for rse in client.list_rses(rse_expression=RSE_EXPRESSION)]
+ddm_quotas = {}
+overridden_ddm_quotas = {}
 
-for rse in rses:
+def get_stats():
+    values = list(ddm_quotas.values())
+    print("  ====== STATISTICS   ======  ")
+    pprint.pprint(sorted(ddm_quotas.items(), key=lambda x:x[1]))
+    print("")
+    print("Mean: ", np.mean(values))
+    print("Std: ", np.std(values))
+
+def get_sum_of_all_rse_statics():
+    result = 0
+    for rse in rses:
+        static, rucio, expired = get_rse_info(rse)
+        result += static
+    return result
+
+def get_rse_info(rse):
     rse_usage = list(client.get_rse_usage(rse))
 
-    static, rucio, unavailable, expired = 0, 0, 0, 0
+    required_fields = {"static", "rucio", "expired"}
+    relevant_info = {}
+
     for source in rse_usage:
-        if source["source"] == "static":
-            static = source["used"]
-        if source["source"] == "rucio":
-            rucio = source["used"]
-        if source["source"] == "unavailable":
-            unavailable = source["used"]
-        if source["source"] == "expired":
-            expired = source["used"]
-
-    # Normalise
-    if static == 0:
-        continue # Skip if static is 0
+        # Assuming source and used keys exist
+        relevant_info[source["source"]] = source["used"]
+
+    if not required_fields.issubset(relevant_info.keys()):
+        print("Skipping {} due to lack of relevant key in rse".format(rse))
+        print("{} is not a subset of {}".format(required_fields, relevant_info.keys()))
+        return 0,0,0
 
     # ddm_quota is set proportional to percentage of (dynamic + free) space
-    # Apprantly, python integers do not overflow, https://docs.python.org/3/library/exceptions.html#OverflowError
-
-    locked = rucio - unavailable - expired
-    ddm_quota = int((max(static - locked, 0)/static)*1e4)
-
-    # Override ddm_quota for operational purposes
-    rse_attributes = client.list_rse_attributes(rse)
-    if "override_ddm_quota" in rse_attributes:
-        ddm_quota = rse_attributes["override_ddm_quota"]
-
-    if not DRY_RUN:
-        client.add_rse_attribute(rse, "ddm_quota", ddm_quota)
-    print("Setting ddm_quota for {} to {}".format(rse, ddm_quota))
+    # Apparently, python integers do not overflow, https://docs.python.org/3/library/exceptions.html#OverflowError
+
+    static, rucio, expired = relevant_info["static"], relevant_info["rucio"], relevant_info["expired"]
+    return static, rucio, expired
+
+def calculate_ddm_quotas():
+
+    total_static = get_sum_of_all_rse_statics()
+
+    for rse in rses:
+        static, rucio, expired = get_rse_info(rse)
+
+        # Normalise
+        if static == 0:
+            continue  # Skip if static is 0
+
+        free = static - rucio
+        ddm_quota = STATIC_WEIGHT *(static/total_static) + FREE_WEIGHT * (free/static) + EXPIRED_WEIGHT * (expired/static)
+        if MAKE_QUADRATIC:
+            ddm_quota = ddm_quota ** 2
+
+        # Override ddm_quota for operational purposes
+        rse_attributes = client.list_rse_attributes(rse)
+        if "override_ddm_quota" in rse_attributes:
+            overridden_ddm_quotas[rse] = rse_attributes["override_ddm_quota"]
+            continue
+
+        ddm_quotas[rse] = ddm_quota
+
+
+def normalize_ddm_quotas():
+    weights = [value for value in ddm_quotas.values()]
+    for rse, weight in ddm_quotas.items():
+        ddm_quotas[rse] = int(((weight - min(weights)) / (max(weights) - min(weights))) * MAX_DDM_QUOTA)
+
+
+def set_ddm_quotas():
+    # Set automatically calculated ddm quotas
+    for rse, ddm_quota in ddm_quotas.items():
+        if DRY_RUN:
+            print("DRY-RUN: Set ddm_quota for {} to {}".format(rse, ddm_quota))
+        else:
+            client.add_rse_attribute(rse, "ddm_quota", ddm_quota)
+            print("Set ddm_quota for {} to {}".format(rse, ddm_quota))
+
+    # Set overriden ddm_quotas
+    for rse, ddm_quota in overridden_ddm_quotas.items():
+        if DRY_RUN:
+            print("DRY-RUN: Override ddm_quota for {} to {}".format(rse, ddm_quota))
+        else:
+            client.add_rse_attribute(rse, "ddm_quota", ddm_quota)
+            print("Override ddm_quota for {} to {}".format(rse, ddm_quota))
+
+def main():
+
+    calculate_ddm_quotas()
+    normalize_ddm_quotas()
+    set_ddm_quotas()
+    if VERBOSE:
+        get_stats()
+
+if __name__ == "__main__":
+    main()
+