updateDDMQuota

#! /usr/bin/env python3

import math
from rucio.client import Client
from tabulate import tabulate
client = Client()

# See the following link for documentation and please update it if you change the logic
# https://cmsdmops.docs.cern.ch/Operators/ManageDMWeight/

DRY_RUN = False
MAX_DM_WEIGHT = 100
MIN_DM_WEIGHT = 1
stats = {}

def print_results():
    headers = ['RSE', 'PLEDGE (PB)', 'FREE SPACE (PB)', "RELATIVE FREE (%)", "DM WEIGHT COEFFICIENT", "DM_WEIGHT"]
    table_disk = []
    table_tape = []
    for rse, values in stats.items():
        pledge = values["pledge"] / 1e15
        free = values["free"] / 1e15
        relative_free = (free/pledge) * 100
        dm_weight_coefficient = values["dm_weight_coefficient"]
        dm_weight = values["dm_weight"]
        if "Tape" in rse:
            table_tape.append([rse,pledge,free,relative_free,dm_weight_coefficient, dm_weight])
        else: 
            table_disk.append([rse,pledge,free,relative_free,dm_weight_coefficient, dm_weight])

    print(tabulate(table_disk, headers=headers))
    print(tabulate(table_tape, headers=headers))


def get_sum_of_all_rse_statics(rse_expression):
    rses = [rse["rse"] for rse in client.list_rses(rse_expression=rse_expression)]
    result = 0
    for rse in rses:
        static, _, _ = get_rse_info(rse)
        result += static
    return result


def get_rse_info(rse):
    rse_usage = list(client.get_rse_usage(rse))

    required_fields = {"static", "rucio", "expired"}
    relevant_info = {}

    for source in rse_usage:
        # 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

    # 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_dm_weights(rse_expression, static_weight, free_weight, expired_weight, make_quadratic):

    total_static = get_sum_of_all_rse_statics(rse_expression)
    rses = [rse["rse"] for rse in client.list_rses(rse_expression=rse_expression)]
    dm_weights = {}

    for rse in rses:
        static, rucio, expired = get_rse_info(rse)

        # Normalise
        if static == 0:
            continue  # Skip if static is 0

        # Control dm_weight for specially configured rses
        rse_attributes = client.list_rse_attributes(rse)

        free = static - rucio
        dm_weight = static_weight * (static / total_static) + free_weight * \
            (free / static) + expired_weight * (expired / static)
        dm_weight_coefficient = 1
        if "dm_weight_coefficient" in rse_attributes:
            dm_weight_coefficient = float(rse_attributes["dm_weight_coefficient"])
        dm_weight *= dm_weight_coefficient
        if make_quadratic:
            dm_weight = dm_weight ** 2

        stats[rse] = {
            "pledge": static,
            "free": free,
            "dm_weight_coefficient" : dm_weight_coefficient,
        }

        dm_weights[rse] = dm_weight
    return dm_weights


def normalize_dm_weights(dm_weights):
    weights = [value for value in dm_weights.values()]
    min_positive_weight = min([n for n in weights if n > 0])
    for rse, weight in dm_weights.items():

        # Don't take 0 weights into normalization, since they are manually overriden
        if weight == 0:
            dm_weight = weight
        else:
            # This normalization sets the lowest dm_weight to 0
            dm_weight = int(((weight - min_positive_weight) / (max(weights) - min_positive_weight)) * MAX_DM_WEIGHT)
            # Make the lowest dm_weight non-zero
            if dm_weight == 0:
                dm_weight = MIN_DM_WEIGHT
            #dm_weight = int((weight/sum(weights)) * MAX_DM_WEIGHT)
        dm_weights[rse] = dm_weight
        stats[rse]["dm_weight"] = dm_weight

    return dm_weights


def set_dm_weights(dm_weights):
    # Set automatically calculated dm weights
    for rse, dm_weight in dm_weights.items():
        if DRY_RUN:
            print("DRY-RUN: Set dm_weight for {} to {}".format(rse, dm_weight))
        else:
            client.add_rse_attribute(rse, "dm_weight", dm_weight)
            # Keep setting ddm_quota for disk rses until we deprecate it
            if "Tape" not in rse:
                client.add_rse_attribute(rse, "ddm_quota", dm_weight)
            print("Set dm_weight for {} to {}".format(rse, dm_weight))

def run(rse_expression, static_weight, free_weight, expired_weight, make_quadratic):

    # Calculate dm_weights for disk rses
    # T3s do not have the "static" usage set, they are quasi-static
    dm_weights = calculate_dm_weights(rse_expression = rse_expression,
                         static_weight = static_weight,
                         free_weight = free_weight, 
                         expired_weight = expired_weight,
                         make_quadratic = make_quadratic)
    normalized_dm_weights = normalize_dm_weights(dm_weights)
    set_dm_weights(normalized_dm_weights)

def main():

    # Calculate dm_weights for disk rses
    # T3s do not have the "static" usage set, they are quasi-static
    
    run(rse_expression = "rse_type=DISK&cms_type=real&tier<3&tier>0",
                         static_weight = 0.2,
                         free_weight = 0.5, 
                         expired_weight = 0.3,
                         make_quadratic = True)
    

    # Calculate dm_weights for tape rses
    run(rse_expression = "rse_type=TAPE&cms_type=real&wmcore_output_tape=True",
                         static_weight = 0.5,
                         free_weight = 0.5, 
                         expired_weight = 0,
                         make_quadratic = False)

    print_results()



if __name__ == "__main__":
    main()