RucioUtils.py

""" a small set of utilities to work with Rucio used in various places """
import logging
import time
import traceback
from functools import wraps
from TaskWorker.WorkerExceptions import TaskWorkerException
from rucio.client import Client as NativeClient
from rucio.common.exception import RSENotFound, RuleNotFound, RucioException

RETRIABLE_RUCIO_HTTP_STATUSES = [503]

def _is_rucio_retriable_http_error(exc):
    """True if this RucioException wraps a transient HTTP error we should retry."""
    msg = str(exc).lower()
    return any(f"http status code: {code}" in msg for code in RETRIABLE_RUCIO_HTTP_STATUSES)

def withExponentialBackOffRetry(retryAttempts=5, fatalExceptions=(), retryExceptions=(Exception,), retryPredicate=None):
    """
    Generic Exponential Back-off Retry
    - retryAttempts: The number of retry attempts to perform before giving up.
        Guidances:
            5 attempts -- Industrial Standard for HTTP API calls — (≈ 30 seconds tolerance, default).
            8 attempts -- Standard HTTP around Rucio API calls — (≈ 8.5 mins tolerance).
            9/10 attempts -- Long-running or critical operations such as job submission, task management, and tape recall (≈ 17/34 minutes tolerance).
    - fatalExceptions: A tuple of exception types that should not be retried, raise immediately.
    - retryExceptions: A tuple of exception types that are eligible for retry. Otherwise will be raise right away as well. (But our default is built-in Exception, so all exception will be catch and retry anyhow.
    - retryPredicate: Optional callable(exception) -> bool. When a fatalException is caught,
        if retryPredicate returns True, it is retried instead of raised. Extra sleep tolerance is applied.
    """
    fatalExceptions = tuple(fatalExceptions)
    retryExceptions = tuple(retryExceptions)

    def decorator(func):
        @wraps(func)
        def wrapper(*args, **kwargs):
            # Ensure universally pick the right logger.
            logger = None
            if args: # Try to get a logger from `self`, probably initiailized by Class.
                logger = getattr(args[0], "_logger", None) or getattr(args[0], "logger", None)

            if logger is None: # Else, robustly retrieve from func module, or current modudle.
                logger = logging.getLogger(getattr(func, "__module__", __name__))

            attempt = 0
            name = func.__name__
            while True:
                try:
                    return func(*args, **kwargs)
                except fatalExceptions as e:
                    if retryPredicate and retryPredicate(e):
                        if attempt > retryAttempts:
                            logger.error(f"Operation '{name}' failed after {attempt} retries (retryPredicate matched): {e}")
                            raise
                        sleepTime = 20 + (2 ** attempt)
                        logger.warning(
                            f"Retriable exception in '{name}' (attempt {attempt+1}/{retryAttempts}): "
                            f"{e}, waiting for {sleepTime} seconds..."
                        )
                        time.sleep(sleepTime)
                        attempt += 1
                        continue
                    logger.exception(f"Fatal exception in '{name}' : {e}")
                    logger.exception(f"Type of Exception: {type(e)}")
                    logger.exception(f"repr(): {repr(e)}")
                    traceback.print_exc()
                    raise
                except retryExceptions as e:
                    if attempt > retryAttempts:
                        logger.error(f"Operation '{name}' failed after {attempt} retries: {e}")
                        raise
                    sleepTime = 2 ** attempt # time in seconds
                    logger.warning(f"Retryable exception in '{name}' (attempt {attempt+1}/{retryAttempts}): {e}, waiting for {sleepTime} seconds...")
                    time.sleep(sleepTime)
                    attempt += 1
        return wrapper
    return decorator

class Client:
    # Wraps NativeClient with configurable retry logic and logging
    def __init__(self, *args, logger=None, **kwargs):
        self._client = NativeClient(*args, **kwargs)
        self._logger = logger or logging.getLogger(__name__)

    # Intercepts Native Rucio Callable attributetes, wrap method calls to enable Exponential Retry.
    def __getattr__(self, name):
        attr = getattr(self._client, name)
        if not callable(attr):
            return attr

        @withExponentialBackOffRetry(retryAttempts=10, fatalExceptions=(RucioException,), retryPredicate=_is_rucio_retriable_http_error)
        @wraps(attr)
        def call(*args, **kwargs):
            return attr(*args, **kwargs)

        return call


def getNativeRucioClient(config=None, logger=None):
    """
    instantiates a Rucio python Client for use in CRAB TaskWorker
    :param config: a TaskWorker configuration object in which
                    at least the variables used below are defined
    :param logger: a valid logger instance
    :return: a Rucio Client object
    """
    logger.info("Initializing native Rucio client")

    rucioLogger = logging.getLogger('RucioClient')
    rucioLogger.setLevel(logging.INFO)

    # silence a few noisy components used by rucio
    ul = logging.getLogger('urllib3')
    ul.setLevel(logging.ERROR)
    dl = logging.getLogger('dogpile')
    dl.setLevel(logging.ERROR)
    cl = logging.getLogger('charset_normalizer')
    cl.setLevel(logging.ERROR)

    # allow for both old and new configuration style

    if getattr(config, 'Services', None):
        rucioConfig = config.Services
    else:
        rucioConfig = config

    rucioCert = getattr(rucioConfig, "Rucio_cert")
    rucioKey = getattr(rucioConfig, "Rucio_key")
    logger.debug("Using cert [%s]\n and key [%s] for rucio client.", rucioCert, rucioKey)
    client = Client(
        rucio_host=rucioConfig.Rucio_host,
        auth_host=rucioConfig.Rucio_authUrl,
        ca_cert=rucioConfig.Rucio_caPath,
        account=rucioConfig.Rucio_account,
        creds={"client_cert": rucioCert, "client_key": rucioKey},
        auth_type='x509',
        logger=rucioLogger
    )

    # Initial check: these calls now retry automatically
    ret = client.ping()
    logger.info("Rucio server v.%s contacted", ret['version'])
    ret = client.whoami()
    logger.info("Rucio client initialized for %s in status %s", ret['account'], ret['status'])

    return client


def getWritePFN(rucioClient=None, siteName='', lfn='',  # pylint: disable=dangerous-default-value
                operations=['third_party_copy_write', 'write'], logger=None):
    """
    convert a single LFN into a PFN which can be used for Writing via Rucio
    Rucio supports the possibility that at some point in the future sites may
    require different protocols or hosts for read or write operations
    :param rucioClient: Rucio python client, e.g. the object returned by getNativeRucioClient above
    :param siteName: e.g. 'T2_CH_CERN'
    :param lfn: a CMS-style LFN
    :param logger: a valid logger instance
    :return: a CMS-style PFN
    """

    # add a scope to turn LFN into Rucio DID syntax
    did = 'cms:' + lfn
    # we prefer to do ASO via FTS which uses 3rd party copy, fall back to protocols defined
    # for other operations in case that fails, order matters here !
    # "third_party_copy_write": provides the PFN to be used with FTS
    # "write": provides the PFN to be used with gfal
    # 2022-08: dario checked with felipe that every sane RSE has non-zero value
    # for the third_party_copy_write column, which means that it is available.
    exceptionString = ""
    didDict = None
    for operation in operations:
        try:
            logger.warning('Try Rucio lfn2pn with operation %s', operation)
            didDict = rucioClient.lfns2pfns(siteName, [did], operation=operation)
            break
        except RSENotFound:
            msg = f"Site {siteName} not found in CMS site list"
            raise TaskWorkerException(msg) from RSENotFound
        except Exception as ex:  # pylint: disable=broad-except
            msg = 'Rucio lfn2pfn resolution for %s failed with:\n%s\nTry next one.'
            logger.warning(msg, operation, str(ex))
            exceptionString += f"operation: {operation}, exception: {ex}\n"
    if not didDict:
        msg = f"lfn2pfn resolution with Rucio failed for site: {siteName}  LFN: {lfn}"
        msg += f" with exception(s) :\n{exceptionString}"
        raise TaskWorkerException(msg)

    # lfns2pfns returns a dictionary with did as key and pfn as value:
    #  https://rucio.readthedocs.io/en/latest/api/rse.html
    # {u'cms:/store/user/rucio': u'gsiftp://eoscmsftp.cern.ch:2811/eos/cms/store/user/rucio'}
    pfn = didDict[did]
    logger.info(f"Will use {pfn} as stageout location")

    return pfn

@withExponentialBackOffRetry(retryAttempts=10, fatalExceptions=(RucioException,))
def getRuleQuota(rucioClient=None, ruleId=None):
    """ return quota needed by this rule in Bytes """
    size = 0
    try:
        rule = rucioClient.get_replication_rule(ruleId)
    except RuleNotFound:
        return 0
    files = rucioClient.list_files(scope=rule['scope'], name= rule['name'])
    size = sum(file['bytes'] for file in files)
    return size

@withExponentialBackOffRetry(retryAttempts=10, fatalExceptions=(RucioException,))
def getRucioUsage(rucioClient=None, account=None, activity =None):
    """ size of Rucio usage for this account (if provided) or by activity """
    if activity is None:
        if account is None:
            totalusage = 0
            raise ValueError("Error: Account and Activity both unspecified")
        else:
            usageGenerator = rucioClient.get_local_account_usage(account=account)
            totalBytes = 0
            for usage in usageGenerator:
                used = usage['bytes']
                totalBytes += used
            totalusage = totalBytes
    else:
        filters = {'activity': activity}
        if account is not None:
            filters['account'] = account
        rules = rucioClient.list_replication_rules(filters=filters)
        if activity == 'Analysis Input':
            valid_states = ['OK', 'REPLICATING', 'STUCK', 'SUSPENDED']
        elif activity == 'Analysis TapeRecall':
            valid_states = ['REPLICATING', 'STUCK', 'SUSPENDED']  # Exclude 'OK' state
        else:
            print("Error: Unknown activity selected in quota report.")
            valid_states = []

        # Calculate usage only if valid_states is set
        # Rucio does not keep track by activity internally, so we need to find all rules and sum all files locked by each rule
        if valid_states:
            totalusage = sum(getRuleQuota(rucioClient, rule['id']) for rule in rules if rule['state'] in valid_states)
        else:
            totalusage = 0

    return totalusage