rucio 37.0.0rc1__py3-none-any.whl

rucio/core/rse_selector.py

@@ -0,0 +1,304 @@
+# Copyright European Organization for Nuclear Research (CERN) since 2012
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#    http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from random import shuffle, uniform
+from typing import TYPE_CHECKING, Optional
+
+from rucio.common.exception import InsufficientAccountLimit, InsufficientTargetRSEs, InvalidRuleWeight, RSEOverQuota
+from rucio.core.account import get_all_rse_usages_per_account, get_usage, has_account_attribute
+from rucio.core.account_limit import get_global_account_limits, get_local_account_limit
+from rucio.core.rse import get_rse_limits, has_rse_attribute, list_rse_attributes
+from rucio.core.rse_counter import get_counter as get_rse_counter
+from rucio.core.rse_expression_parser import parse_expression
+from rucio.db.sqla.session import read_session
+
+if TYPE_CHECKING:
+    from collections.abc import Iterable, Sequence
+
+    from sqlalchemy.orm import Session
+
+    from rucio.common.types import InternalAccount
+
+
+class RSESelector:
+    """
+    Representation of the RSE selector
+    """
+
+    @read_session
+    def __init__(self, account, rses, weight, copies, ignore_account_limit=False, *, session: "Session"):
+        """
+        Initialize the RSE Selector.
+
+        :param account:               Account owning the rule.
+        :param rses:                  List of rse dictionaries.
+        :param weight:                Weighting to use.
+        :param copies:                Number of copies to create.
+        :param ignore_account_limit:  Flag if the quota should be ignored.
+        :param session:               DB Session in use.
+        :raises:                      InvalidRuleWeight, InsufficientAccountLimit, InsufficientTargetRSEs
+        """
+        self.account = account
+        self.rses = []  # [{'rse_id':, 'weight':, 'staging_area'}]
+        self.copies = copies
+        if weight is not None:
+            for rse in rses:
+                attributes = list_rse_attributes(rse_id=rse['id'], session=session)
+                availability_write = True if rse.get('availability_write', True) else False
+                if weight not in attributes:
+                    continue  # The RSE does not have the required weight set, therefore it is ignored
+                try:
+                    self.rses.append({'rse_id': rse['id'],
+                                      'weight': float(attributes[weight]),
+                                      'mock_rse': attributes.get('mock', False),
+                                      'availability_write': availability_write,
+                                      'staging_area': rse['staging_area']})
+                except ValueError:
+                    raise InvalidRuleWeight('The RSE \'%s\' has a non-number specified for the weight \'%s\'' % (rse['rse'], weight))
+        else:
+            for rse in rses:
+                mock_rse = has_rse_attribute(rse['id'], 'mock', session=session)
+                availability_write = True if rse.get('availability_write', True) else False
+                self.rses.append({'rse_id': rse['id'],
+                                  'weight': 1,
+                                  'mock_rse': mock_rse,
+                                  'availability_write': availability_write,
+                                  'staging_area': rse['staging_area']})
+
+        if len(self.rses) < self.copies:
+            raise InsufficientTargetRSEs('Target RSE set not sufficient for number of copies. (%s copies requested, RSE set size %s)' % (self.copies, len(self.rses)))
+
+        rses_with_enough_quota = []
+        if has_account_attribute(account=account, key='admin', session=session) or ignore_account_limit:
+            for rse in self.rses:
+                rse['quota_left'] = float('inf')
+                rse['space_left'] = float('inf')
+                rses_with_enough_quota.append(rse)
+        else:
+            global_quota_limit = get_global_account_limits(account=account, session=session)
+            all_rse_usages = {usage['rse_id']: usage['bytes'] for usage in get_all_rse_usages_per_account(account=account, session=session)}
+            for rse in self.rses:
+                if rse['mock_rse']:
+                    rse['quota_left'] = float('inf')
+                    rse['space_left'] = float('inf')
+                    rses_with_enough_quota.append(rse)
+                else:
+                    # check local quota
+                    local_quota_left = None
+                    quota_limit = get_local_account_limit(account=account, rse_id=rse['rse_id'], session=session)
+                    if quota_limit is None:
+                        local_quota_left = 0
+                    else:
+                        local_quota_left = quota_limit - get_usage(rse_id=rse['rse_id'], account=account, session=session)['bytes']
+
+                    # check global quota
+                    rse['global_quota_left'] = {}
+                    all_global_quota_enough = True
+                    for rse_expression, limit in global_quota_limit.items():
+                        if rse['rse_id'] in limit['resolved_rse_ids']:
+                            quota_limit = limit['limit']
+                            global_quota_left = None
+                            if quota_limit is None:
+                                global_quota_left = 0
+                            else:
+                                rse_expression_usage = 0
+                                for rse_id in limit['resolved_rse_ids']:
+                                    rse_expression_usage += all_rse_usages.get(rse_id, 0)
+                                global_quota_left = quota_limit - rse_expression_usage
+                            if global_quota_left <= 0:
+                                all_global_quota_enough = False
+                                break
+                            else:
+                                rse['global_quota_left'][rse_expression] = global_quota_left
+                    if local_quota_left > 0 and all_global_quota_enough:
+                        rse['quota_left'] = local_quota_left
+                        space_limit = get_rse_limits(name='MaxSpaceAvailable', rse_id=rse['rse_id'], session=session).get('MaxSpaceAvailable')
+                        if space_limit is None or space_limit < 0:
+                            rse['space_left'] = float('inf')
+                        else:
+                            rse['space_left'] = space_limit - get_rse_counter(rse_id=rse['rse_id'], session=session)['bytes']
+                        rses_with_enough_quota.append(rse)
+
+        self.rses = rses_with_enough_quota
+        if len(self.rses) < self.copies:
+            raise InsufficientAccountLimit('There is insufficient quota on any of the target RSE\'s to fulfill the operation.')
+
+        # don't consider removing rses based on the total space here - because files already on the RSE are taken into account
+        # it is possible to have no space but still be able to fulfil the rule
+
+    def select_rse(
+        self,
+        size: int,
+        preferred_rse_ids: "Iterable[str]",
+        copies: int = 0,
+        blocklist: Optional["Sequence[str]"] = None,
+        prioritize_order_over_weight: bool = False,
+        existing_rse_size: Optional[dict[str, int]] = None
+    ) -> list[tuple[str, bool, bool]]:
+        """
+        Select n RSEs to replicate data to.
+
+        :param size:                         Size of the block being replicated.
+        :param preferred_rse_ids:            Ordered list of preferred rses. (If possible replicate to them)
+        :param copies:                       Select this amount of copies, if 0 use the pre-defined rule value.
+        :param blocklist:                    List of blocked rses. (Do not put replicas on these sites)
+        :param prioritze_order_over_weight:  Prioritize the order of the preferred_rse_ids list over the picking done by weight.
+        :existing_rse_size:                  Dictionary of size of files already present at each rse
+        :returns:                            List of (RSE_id, staging_area, availability_write) tuples.
+        :raises:                             InsufficientAccountLimit, InsufficientTargetRSEs
+        """
+
+        blocklist = blocklist or []
+        result = []
+        rses = self.rses
+        count = self.copies if copies == 0 else copies
+
+        # Remove blocklisted rses
+        if blocklist:
+            rses = [rse for rse in self.rses if rse['rse_id'] not in blocklist]
+        if len(rses) < count:
+            raise InsufficientTargetRSEs('There are not enough target RSEs to fulfil the request at this time.')
+
+        # Remove rses which do not have enough space, accounting for the files already at each rse
+        if existing_rse_size is None:
+            existing_rse_size = {}
+        rses = [rse for rse in rses if rse['space_left'] >= size - existing_rse_size.get(rse['rse_id'], 0)]
+        if len(rses) < count:
+            raise RSEOverQuota('There is insufficient space on any of the target RSE\'s to fulfill the operation.')
+
+        # Remove rses which do not have enough local quota
+        rses = [rse for rse in rses if rse['quota_left'] > size]
+        if len(rses) < count:
+            raise InsufficientAccountLimit('There is insufficient quota on any of the target RSE\'s to fulfill the operation.')
+
+        # Remove rses which do not have enough global quota
+        rses_with_enough_quota = []
+        for rse in rses:
+            enough_global_quota = True
+            for rse_expression in rse.get('global_quota_left', []):
+                if rse['global_quota_left'][rse_expression] < size:
+                    enough_global_quota = False
+                    break
+            if enough_global_quota:
+                rses_with_enough_quota.append(rse)
+        rses = rses_with_enough_quota
+        if len(rses) < count:
+            raise InsufficientAccountLimit('There is insufficient quota on any of the target RSE\'s to fulfill the operation.')
+
+        for copy in range(count):
+            # Remove rses already in the result set
+            rses = [rse for rse in rses if rse['rse_id'] not in [item[0] for item in result]]
+            rses_dict = {}
+            for rse in rses:
+                rses_dict[rse['rse_id']] = rse
+            # Prioritize the preferred rses
+            preferred_rses = [rses_dict[rse_id] for rse_id in preferred_rse_ids if rse_id in rses_dict]
+            if prioritize_order_over_weight and preferred_rses:
+                rse = (preferred_rses[0]['rse_id'], preferred_rses[0]['staging_area'], preferred_rses[0]['availability_write'])
+            elif preferred_rses:
+                rse = self.__choose_rse(preferred_rses)
+            else:
+                rse = self.__choose_rse(rses)
+            result.append(rse)
+            self.__update_quota(rse, size)
+        return result
+
+    def get_rse_dictionary(self):
+        """
+        Return the current dictionary of potential RSEs stored in the RSE selector
+
+        :returns:  List of RSE dictionaries
+        """
+        rse_dict = {}
+        for rse in self.rses:
+            rse_dict[rse['rse_id']] = rse
+        return rse_dict
+
+    def __update_quota(self, rse, size):
+        """
+        Update the internal quota value.
+
+        :param rse:      RSE tuple to update.
+        :param size:     Size to subtract.
+        """
+
+        for element in self.rses:
+            if element['rse_id'] == rse[0]:
+                element['quota_left'] -= size
+                for rse_expression in element.get('global_quota_left', []):
+                    element['global_quota_left'][rse_expression] -= size
+                return
+
+    def __choose_rse(self, rses):
+        """
+        Choose an RSE based on weighting.
+
+        :param rses:  The rses to be considered for the choose.
+        :return:      The (rse_id, staging_area) tuple of the chosen RSE.
+        """
+
+        shuffle(rses)
+        pick = uniform(0, sum([rse['weight'] for rse in rses]))  # noqa: S311
+        weight = 0
+        for rse in rses:
+            weight += rse['weight']
+            if pick <= weight:
+                return (rse['rse_id'], rse['staging_area'], rse['availability_write'])
+
+
+@read_session
+def resolve_rse_expression(
+    rse_expression: str,
+    account: "InternalAccount",
+    weight: Optional[int] = None,
+    copies: int = 1,
+    ignore_account_limit: bool = False,
+    size: int = 0,
+    preferred_rses: Optional["Iterable[str]"] = None,
+    blocklist: Optional["Sequence[str]"] = None,
+    prioritize_order_over_weight: bool = False,
+    existing_rse_size: Optional[dict[str, int]] = None,
+    *,
+    session: "Session"
+) -> tuple[list[str], list[str]]:
+    """
+    Resolve a potentially complex RSE expression into `copies` single-RSE expressions. Uses `parse_expression()`
+    to decompose the expression, then `RSESelector.select_rse()` to pick the target RSEs.
+    """
+
+    preferred_rses = preferred_rses or []
+    blocklist = blocklist or []
+    rses = parse_expression(rse_expression, filter_={'vo': account.vo}, session=session)
+
+    rse_to_id = dict((rse_dict['rse'], rse_dict['id']) for rse_dict in rses)
+    id_to_rse = dict((rse_dict['id'], rse_dict['rse']) for rse_dict in rses)
+
+    selector = RSESelector(account=account,
+                           rses=rses,
+                           weight=weight,
+                           copies=copies,
+                           ignore_account_limit=ignore_account_limit,
+                           session=session)
+
+    preferred_rse_ids = [rse_to_id[rse] for rse in preferred_rses if rse in rse_to_id]
+
+    preferred_unmatched = list(set(preferred_rses) - set(rse_dict['rse'] for rse_dict in rses))
+
+    selection_result = selector.select_rse(size=size,
+                                           preferred_rse_ids=preferred_rse_ids,
+                                           blocklist=blocklist,
+                                           prioritize_order_over_weight=prioritize_order_over_weight,
+                                           existing_rse_size=existing_rse_size)
+
+    return [id_to_rse[rse_id] for rse_id, _, _ in selection_result], preferred_unmatched