Repository: rucio/rucio

Issue:
Speed up list_replicas() implementation
Motivation
----------
DUNE is planning to use ReplicaClient.list_replics() method at the scale of ~1000 files. Currently, the method is too slow to be practical. It hakes from 2.5 to 3 seconds for 100 files. The method needs to be faster to be useful.


Modification
------------
Currently, core.replica.list_replicas() calls _resolve_dids(), which is issuing a query per requested DID, unless the DID type is known and is "file". Unfortunately, when I try to add "did_type":"file" to my list of DIDs on the client side, the server rejects my request because "did_type" is not accepted by the request schema:

```
Traceback (most recent call last):
  File "list_replicas.py", line 7, in <module>
    out = list(client.list_replicas([{"scope":f["namespace"], "name":f["name"], "did_type":"file"} for f in files], 
  File "/home/ivm/miniconda3/lib/python3.8/site-packages/rucio/client/replicaclient.py", line 196, in list_replicas
    raise exc_cls(exc_msg)
rucio.common.exception.InvalidObject: Provided object does not match schema.
Details: Problem validating r_dids : Additional properties are not allowed ('did_type' was unexpected)
...

On instance[0]:
    {'did_type': 'file',
     'name': 'dc4_np02tde_407021731_455_59_cb.test',
     'scope': 'dc4-vd-coldbox-top'}
```

There can be 2 modifications:
1. Accept "did_type" on the server side
2. Optimize _resolve_dids() so that it issues single query to resolve DID types instead of one query per DID at https://github.com/rucio/rucio/blob/8e32fd1527116640e0ed719cb526bb5ff4303a52/lib/rucio/core/replica.py#L649




Discussion:
So, I'm trying to understand what's going on here.

Here's an example from the ATLAS instance:

```
$ time -p rucio list-file-replicas mc15_13TeV:mc15_13TeV.422003.ParticleGunEvtGen_single_B0_flatPt10to1000.simul.HITS.e4566_s2726_tid06828093_00 | wc -l
1300
real 5.43
user 0.97
sys 0.27
```


So roughly 6 seconds to query a dataset with 999 files (this is first resolved into its constituent file DIDs, then replicas are resolved) into 1300 replicas. Interestingly, the time stays rather constant, with a tendency to have a larger baseline cost. I tried for a couple differently-sized datasets (from 1-2 files, up to 10k files), and it's always noticeably under 10 seconds, but never less than 2 seconds.
@mlassnig if I understand correctly, you are listing replicas in a dataset ? This is not the use case I am talking about. In my use case, the replicas are given with scope/name pairs and do not necessarily map into any combination of datasets
here is my code:
```python
import sys, json
from rucio.client.replicaclient import ReplicaClient

files = json.load(open(sys.argv[1], "r"))
print("loaded", len(files), "files")
client = ReplicaClient()
out = list(client.list_replicas([{"scope":f["namespace"], "name":f["name"], "did_type":"file"} for f in files], 
	all_states=True, ignore_availability=True, resolve_archives=False))
print(len(out))
```
It shouldn't really matter. In fact, my variant should be slower than yours, because the dataset resolve happens before.

Here's your script evaluating 999 files directly (~9 seconds)

```
$ python test.py 
('loaded', 999, 'files')
2022-08-04 15:53:57.555723
2022-08-04 15:54:06.258285
999
```
So, I'm certain that there's some improvements that can be made, however I don't really think that there's a linear increase in time depending on the size of the requested number of files.
here is my modified code:
```python
import sys, json, getopt, time
from rucio.client.replicaclient import ReplicaClient

opts, args = getopt.getopt(sys.argv[1:], "n:")
opts = dict(opts)

files = json.load(open(args[0], "r"))
if "-n" in opts:
	files = files[:int(opts["-n"])]
n = len(files)
print("loaded", n, "files")
client = ReplicaClient()
t0 = time.time()
out = list(client.list_replicas([{"scope":f["namespace"], "name":f["name"]} for f in files], 
	all_states=True, ignore_availability=True))
t1 = time.time()
print("time:", t1-t0, "seconds,   rate:", n/(t1-t0), "files per second")
```
```
$ python list_replicas.py -n 100 files.json 
loaded 100 files
time: 0.9582312107086182 seconds,   rate: 104.3589468621559 files per second
$ python list_replicas.py -n 500 files.json 
loaded 500 files
time: 4.661576509475708 seconds,   rate: 107.25985060711477 files per second
$ python list_replicas.py -n 1000 files.json 
loaded 1000 files
time: 8.867684841156006 seconds,   rate: 112.76900542956581 files per second
```
Looks pretty linear to me
listing replicas by multiple file DIDs and by single dataset DID should make a difference because resolve_dids() has to resolve many file DIDs in the loop in first case while in the dataset case, it has to resolve only one DID: https://github.com/rucio/rucio/blob/8e32fd1527116640e0ed719cb526bb5ff4303a52/lib/rucio/core/replica.py#L649
> Looks pretty linear to me

to me this looks like some throughput problem from your database, maybe we should follow this up on the Slack
Looking at the code, I think it can be fixed in 2 different ways:

1. Let the "did_type" through to bypass the resolving DIDs loop
2. Resolve all DIDs in a single query instead of doing in one DID at a time
I'm happy that did_type cannot be set, I didn't know about that. This means that this is a hack we can actually get rid of. I consider that we must _NOT_ make it available for usage. 
I just performed this test 
```python
from rucio.core.config import get as cconfig_get, set as cconfig_set
from rucio.client import Client

def test_many_replicas(replica_client, mock_scope, core_config_mock):
    files = [{'scope': mock_scope.external, 'name': 'file_%s' % generate_uuid(), 'bytes': 1, 'adler32': '0cc737eb', 'meta': {'events': 10}} for _ in range(1000)]

    client = Client()
    rses = [rse['rse'] for rse in client.list_rses('*') if rse['rse'].startswith('MOCK-')]
    print(f"We have {len(rses)} rses")
    for rse in rses:
        client.add_replicas(rse=rse, files=files)

    for config_val in [True, False, True, False, True, False]:
        cconfig_set('core', 'use_temp_tables', config_val)
        start_time = datetime.now()
        list(list_replicas(dids=[{'scope': mock_scope, 'name': f['name']} for f in files]))
        end_time = datetime.now()
        print(f'Took with use_temp_table={cconfig_get("core", "use_temp_tables")}: {end_time - start_time}')
```

To check if https://github.com/rucio/rucio/pull/5712 improves anything. 

Here is the result : 
```
We have 81 rses
Took with use_temp_table=True: 0:00:03.023870
Took with use_temp_table=False: 0:00:19.340087
Took with use_temp_table=True: 0:00:02.829500
Took with use_temp_table=False: 0:00:19.783525
Took with use_temp_table=True: 0:00:02.805302
Took with use_temp_table=False: 0:00:19.289220
```
resolve_dids is not the problem here:
```
resolve_dids took 0:00:00.528797
Took with use_temp_table=False: 0:00:19.360680
```

Has to be investigated, but it may be that the postgresql optimizer does a bad job on this query
https://github.com/rucio/rucio/blob/8e32fd1527116640e0ed719cb526bb5ff4303a52/lib/rucio/core/replica.py#L874
I will try to bypass the resolve_did step in my dev instance and see if that changes anything
I traced the time spent in various fragments of the server code participating in ``get_replicas()``. For 1000 files, in my smaller dev instance using a dedicated database with ~13K replicas, all replicas are in the same RSE:

- REST layer: 4.48 seconds
- api.replica.list_replicas: 4.46 seconds
    - in api.replica.list_replicas, the [fragment adding some metadata](https://github.com/rucio/rucio/blob/8e32fd1527116640e0ed719cb526bb5ff4303a52/lib/rucio/api/replica.py#L205): 2.82 seconds
         - [calls to get_rse_name()](https://github.com/rucio/rucio/blob/8e32fd1527116640e0ed719cb526bb5ff4303a52/lib/rucio/api/replica.py#L213) in the same fragment: 2.80 seconds and it is called 2 times per replica
    - core.replica.list_replicas: 1.60 seconds
        - resolve_dids: 1.33 seconds

so it is spending 4.13 seconds out of 4.48 seconds (92% of time) in resolve_dids() and get_rse_name()
`get_rse_name()` should be served by memcached and be super fast, is it running for you?
get_rse_name() should not be called 2000 times for a system with single RSE
This is correct and should be fixed. Nevertheless, rucio relies very heavily on caching in other parts of the code and you'll experience bad performance in many other places too if you don't use memcached
I looked at the new temp table-based core.list_replicas() code. It seems to be quite complicated. I am going to de-assign myself from this issue, hoping one of the two things happen:
1. New code will have acceptable speed
2. Someone else, more familiar with the new code will take this issue
I'll pick it up once https://github.com/rucio/rucio/pull/5747 is merged. The changes @ivmfnal did in the closed PR are worth porting to the new code.

Resolving patch:
diff --git a/lib/rucio/api/replica.py b/lib/rucio/api/replica.py
--- a/lib/rucio/api/replica.py
+++ b/lib/rucio/api/replica.py
@@ -220,20 +220,9 @@ def list_replicas(dids, schemes=None, unavailable=False, request_id=None,
                                      client_location=client_location, domain=domain,
                                      sign_urls=sign_urls, signature_lifetime=signature_lifetime,
                                      resolve_archives=resolve_archives, resolve_parents=resolve_parents,
-                                     nrandom=nrandom, updated_after=updated_after, session=session)
+                                     nrandom=nrandom, updated_after=updated_after, by_rse_name=True, session=session)
 
     for rep in replicas:
-        # 'rses' and 'states' use rse_id as the key. This needs updating to be rse.
-        keys = ['rses', 'states']
-        for k in keys:
-            old_dict = rep.get(k, None)
-            if old_dict is not None:
-                new_dict = {}
-                for rse_id in old_dict:
-                    rse = get_rse_name(rse_id=rse_id, session=session) if rse_id is not None else None
-                    new_dict[rse] = old_dict[rse_id]
-                rep[k] = new_dict
-
         rep['scope'] = rep['scope'].external
         if 'parents' in rep:
             new_parents = []
diff --git a/lib/rucio/core/replica.py b/lib/rucio/core/replica.py
--- a/lib/rucio/core/replica.py
+++ b/lib/rucio/core/replica.py
@@ -725,31 +725,6 @@ def get_multi_cache_prefix(cache_site, filename, logger=logging.log):
     return ''
 
 
-def _sort_replica_file_pfns(file):
-    """
-    Helper function to be executed after analyzing the last replica of a file (same scope:name)
-
-    """
-    if not file.get('pfns'):
-        return
-
-    # set the total order for the priority
-    # --> exploit that L(AN) comes before W(AN) before Z(IP) alphabetically
-    # and use 1-indexing to be compatible with metalink
-    tmp = sorted([(file['pfns'][p]['domain'], file['pfns'][p]['priority'], p) for p in file['pfns']])
-    for i in range(0, len(tmp)):
-        file['pfns'][tmp[i][2]]['priority'] = i + 1
-
-    # also sort the pfns inside the rse structure
-    file['rses'] = {}
-    rse_pfns = sorted((file['pfns'][pfn]['rse_id'], file['pfns'][pfn]['priority'], pfn) for pfn in file['pfns'])
-    for t_rse, t_priority, t_pfn in rse_pfns:
-        if t_rse in file['rses']:
-            file['rses'][t_rse].append(t_pfn)
-        else:
-            file['rses'][t_rse] = [t_pfn]
-
-
 def _get_list_replicas_protocols(
         rse_id: str,
         domain: str,
@@ -867,7 +842,7 @@ def _build_list_replicas_pfn(
 
 
 def _list_replicas(replicas, show_pfns, schemes, files_wo_replica, client_location, domain,
-                   sign_urls, signature_lifetime, resolve_parents, filters, session):
+                   sign_urls, signature_lifetime, resolve_parents, filters, by_rse_name, session):
 
     # the `domain` variable name will be re-used throughout the function with different values
     input_domain = domain
@@ -885,6 +860,7 @@ def _list_replicas(replicas, show_pfns, schemes, files_wo_replica, client_locati
 
     for _, replica_group in groupby(replicas, key=lambda x: (x[0], x[1])):  # Group by scope/name
         file = {}
+        pfns = {}
         for scope, name, archive_scope, archive_name, bytes_, md5, adler32, path, state, rse_id, rse, rse_type, volatile in replica_group:
             if isinstance(archive_scope, str):
                 archive_scope = InternalScope(archive_scope, fromExternal=False)
@@ -893,7 +869,7 @@ def _list_replicas(replicas, show_pfns, schemes, files_wo_replica, client_locati
             if not file:
                 file['scope'], file['name'] = scope, name
                 file['bytes'], file['md5'], file['adler32'] = bytes_, md5, adler32
-                file['pfns'], file['rses'], file['states'] = {}, defaultdict(list), {}
+                file['pfns'], file['rses'], file['states'] = {}, {}, {}
                 if resolve_parents:
                     file['parents'] = ['%s:%s' % (parent['scope'].internal, parent['name'])
                                        for parent in rucio.core.did.list_all_parent_dids(scope, name, session=session)]
@@ -901,8 +877,8 @@ def _list_replicas(replicas, show_pfns, schemes, files_wo_replica, client_locati
             if not rse_id:
                 continue
 
-            file['states'][rse_id] = str(state.name if state else state)
-            file['rses'][rse_id] = []
+            rse_key = rse if by_rse_name else rse_id
+            file['states'][rse_key] = str(state.name if state else state)
 
             if not show_pfns:
                 continue
@@ -972,7 +948,7 @@ def _list_replicas(replicas, show_pfns, schemes, files_wo_replica, client_locati
                         else:
                             client_extract = True
 
-                    file['pfns'][pfn] = {
+                    pfns[pfn] = {
                         'rse_id': rse_id,
                         'rse': rse,
                         'type': str(rse_type.name),
@@ -981,7 +957,6 @@ def _list_replicas(replicas, show_pfns, schemes, files_wo_replica, client_locati
                         'priority': priority,
                         'client_extract': client_extract
                     }
-                    file['rses'][rse_id].append(pfn)
 
                 except Exception:
                     # never end up here
@@ -993,8 +968,22 @@ def _list_replicas(replicas, show_pfns, schemes, files_wo_replica, client_locati
                     except KeyError:
                         file['space_token'] = None
 
+        # fill the 'pfns' and 'rses' dicts in file
+        if pfns:
+            # set the total order for the priority
+            # --> exploit that L(AN) comes before W(AN) before Z(IP) alphabetically
+            # and use 1-indexing to be compatible with metalink
+            sorted_pfns = sorted(pfns.items(), key=lambda item: (item[1]['domain'], item[1]['priority'], item[0]))
+            for i, (pfn, pfn_value) in enumerate(list(sorted_pfns), start=1):
+                pfn_value['priority'] = i
+                file['pfns'][pfn] = pfn_value
+
+            sorted_pfns = sorted(file['pfns'].items(), key=lambda item: (item[1]['rse_id'], item[1]['priority'], item[0]))
+            for pfn, pfn_value in sorted_pfns:
+                rse_key = pfn_value['rse'] if by_rse_name else pfn_value['rse_id']
+                file['rses'].setdefault(rse_key, []).append(pfn)
+
         if file:
-            _sort_replica_file_pfns(file)
             yield file
 
     for scope, name, bytes_, md5, adler32 in _list_files_wo_replicas(files_wo_replica, session):
@@ -1027,6 +1016,7 @@ def list_replicas(
         resolve_parents: bool = False,
         nrandom: "Optional[int]" = None,
         updated_after: "Optional[datetime]" = None,
+        by_rse_name: bool = False,
         session: "Optional[Session]" = None,
 ):
     """
@@ -1046,6 +1036,7 @@ def list_replicas(
     :param resolve_archives: When set to true, find archives which contain the replicas.
     :param resolve_parents: When set to true, find all parent datasets which contain the replicas.
     :param updated_after: datetime (UTC time), only return replicas updated after this time
+    :param by_rse_name: if True, rse information will be returned in dicts indexed by rse name; otherwise: in dicts indexed by rse id
     :param session: The database session in use.
     """
     # For historical reasons:
@@ -1055,12 +1046,12 @@ def list_replicas(
     if use_temp_tables:
         yield from _list_replicas_with_temp_tables(
             dids, schemes, unavailable, request_id, ignore_availability, all_states, pfns, rse_expression, client_location,
-            domain, sign_urls, signature_lifetime, resolve_archives, resolve_parents, nrandom, updated_after, session
+            domain, sign_urls, signature_lifetime, resolve_archives, resolve_parents, nrandom, updated_after, by_rse_name, session,
         )
     else:
         yield from _list_replicas_wo_temp_tables(
             dids, schemes, unavailable, request_id, ignore_availability, all_states, pfns, rse_expression, client_location,
-            domain, sign_urls, signature_lifetime, resolve_archives, resolve_parents, nrandom, updated_after, session
+            domain, sign_urls, signature_lifetime, resolve_archives, resolve_parents, nrandom, updated_after, by_rse_name, session,
         )
 
 
@@ -1081,6 +1072,7 @@ def _list_replicas_with_temp_tables(
         resolve_parents: bool = False,
         nrandom: "Optional[int]" = None,
         updated_after: "Optional[datetime]" = None,
+        by_rse_name: bool = True,
         session: "Optional[Session]" = None,
 ):
 
@@ -1356,7 +1348,7 @@ def _inspect_dids(temp_table, session):
             _pick_n_random(
                 nrandom,
                 _list_replicas(replica_tuples, pfns, schemes, [], client_location, domain,
-                               sign_urls, signature_lifetime, resolve_parents, filter_, session)
+                               sign_urls, signature_lifetime, resolve_parents, filter_, by_rse_name, session)
             )
         )
         if len(random_replicas) == nrandom:
@@ -1385,7 +1377,7 @@ def _inspect_dids(temp_table, session):
     yield from _pick_n_random(
         nrandom,
         _list_replicas(replica_tuples, pfns, schemes, [], client_location, domain,
-                       sign_urls, signature_lifetime, resolve_parents, filter_, session)
+                       sign_urls, signature_lifetime, resolve_parents, filter_, by_rse_name, session)
     )
 
 
@@ -4404,6 +4396,7 @@ def _list_replicas_wo_temp_tables(
         resolve_parents: bool = False,
         nrandom: "Optional[int]" = None,
         updated_after: "Optional[datetime]" = None,
+        by_rse_name: bool = True,
         session: "Optional[Session]" = None,
 ):
     if dids:
@@ -4436,5 +4429,5 @@ def _list_replicas_wo_temp_tables(
     yield from _pick_n_random(
         nrandom,
         _list_replicas(replica_tuples, pfns, schemes, files_wo_replica, client_location, domain,
-                       sign_urls, signature_lifetime, resolve_parents, filter_, session)
+                       sign_urls, signature_lifetime, resolve_parents, filter_, by_rse_name, session)
     )


Test patch: