Coverage for python/lsst/rucio/register/rucio_interface.py: 73%
259 statements
« prev ^ index » next coverage.py v7.17.2, created at 2026-11-07 09:41 +0000
1# This file is part of rucio_register
2#
3# Developed for the LSST Data Management System.
4# This product includes software developed by the LSST Project
5# (https://www.lsst.org).
6# See the COPYRIGHT file at the top-level directory of this distribution
7# for details of code ownership.
8#
9# This program is free software: you can redistribute it and/or modify
10# it under the terms of the GNU General Public License as published by
11# the Free Software Foundation, either version 3 of the License, or
12# (at your option) any later version.
13#
14# This program is distributed in the hope that it will be useful,
15# but WITHOUT ANY WARRANTY; without even the implied warranty of
16# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
17# GNU General Public License for more details.
18#
19# You should have received a copy of the GNU General Public License
20# along with this program. If not, see <http://www.gnu.org/licenses/>.
22import collections
23import logging
24import os
25import threading
26import zlib
28import backoff
29import requests
30import rucio.common.exception
31import urllib3.exceptions
32from rucio.client.didclient import DIDClient
33from rucio.client.replicaclient import ReplicaClient
35import lsst.daf.butler
36from lsst.daf.butler import DatasetRef
37from lsst.resources import ResourcePath
38from lsst.rucio.register.resource_bundle import ResourceBundle
39from lsst.rucio.register.rubin_meta import RubinMeta
40from lsst.rucio.register.rucio_did import RucioDID
42__all__ = ["RucioInterface"]
44RETRYABLE = (
45 urllib3.exceptions.ReadTimeoutError,
46 urllib3.exceptions.ProtocolError,
47 requests.exceptions.ConnectionError,
48 requests.exceptions.Timeout,
49 requests.exceptions.ChunkedEncodingError,
50 rucio.common.exception.ServerConnectionException,
51 rucio.common.exception.DatabaseException,
52)
54logger = logging.getLogger(__name__)
56_FACTOR = "factor"
57_MAX_VALUE = "max_value"
58_MAX_TRIES = "max_tries"
59_BACKOFF = {} # Initialized after RucioInterface class definition
62def _backoff_message(details):
63 logger.info(
64 "Backing off {wait:0.1f} seconds after {tries} tries "
65 "calling function {target} with args {args} and kwargs "
66 "{kwargs}".format(**details)
67 )
70class RucioInterface:
71 """Add files as replicas in Rucio, along with metadata,
72 and attach them to datasets.
74 Parameters
75 ----------
76 butler : `lsst.daf.butler.Butler`
77 Butler we're operating upon
78 rucio_rse : `str`
79 Name of the RSE that the files live in.
80 scope : `str`
81 Rucio scope to register the files in.
82 rse_root : `str`
83 Full path to root directory of RSE directory structure
84 dtn_url : `str`
85 Base URL of the data transfer node for the Rucio physical filename.
86 rubin_butler_type : `str`
87 The type registered in "rubin_butler" metadata for Rucio.
88 clear_is_new : `bool`, optional
89 If `True`, set ``is_new`` metadata to `None` on Rucio dataset DIDs
90 post-registration.
91 retries : `int`, optional
92 Maximum number of times to retry failed Rucio operations.
93 timeout : `float`, optional
94 HTTP request timeout in seconds for Rucio ``ReplicaClient`` and
95 ``DIDClient`` operations.
96 """
98 DEFAULT_BACKOFF_FACTOR = 5.0
99 DEFAULT_BACKOFF_MAX_VALUE = 30
100 DEFAULT_RETRIES = 8
101 DEFAULT_TIMEOUT = 300.0
103 def __init__(
104 self,
105 butler: lsst.daf.butler.Butler,
106 rucio_rse: str,
107 scope: str,
108 rse_root: str,
109 dtn_url: str,
110 rubin_butler_type: str,
111 clear_is_new: bool = False,
112 retries: int = DEFAULT_RETRIES,
113 timeout: float = DEFAULT_TIMEOUT,
114 ):
115 self.butler = butler
116 self.rse = rucio_rse
117 self.scope = scope
118 self.rse_root = rse_root
119 self.dtn_url = dtn_url
120 self.pfn_base = f"{dtn_url}"
121 self.timeout = timeout
122 self.replica_client = ReplicaClient(timeout=timeout)
123 self.did_client = DIDClient(timeout=timeout)
124 self.rubin_butler_type = rubin_butler_type
125 self.clear_is_new = clear_is_new
126 self.retries = retries
127 self._hash_cache: dict[str, tuple[int, str]] = {}
128 self._checksum_sources: collections.Counter = collections.Counter()
129 self._lock = threading.Lock()
130 self._butler_lock = threading.Lock()
132 def set_backoff(self, factor, max_value, max_tries) -> None:
133 """Set backoff values for retries
135 Parameters
136 ----------
137 factor: `float`
138 Multipler for backoff
139 max_value: `int`
140 Maximum seconds to backoff to
141 max_tries: `int`
142 Maximum times to try
143 """
144 logger.debug("factor=%f, max_value=%d, max_tries=%d", factor, max_value, max_tries)
145 if factor is not None: 145 ↛ 147line 145 didn't jump to line 147 because the condition on line 145 was always true
146 _BACKOFF[_FACTOR] = factor
147 if max_value is not None: 147 ↛ 149line 147 didn't jump to line 149 because the condition on line 147 was always true
148 _BACKOFF[_MAX_VALUE] = max_value
149 if max_tries is not None: 149 ↛ exitline 149 didn't return from function 'set_backoff' because the condition on line 149 was always true
150 _BACKOFF[_MAX_TRIES] = max_tries
151 self.retries = max_tries
153 def _make_dataset_ref_bundle(self, dataset_id: str, dataset_ref: DatasetRef) -> ResourceBundle:
154 """Make a ResourceBundle
156 Parameters
157 ----------
158 dataset_id : `str`
159 Rucio dataset name
160 dataset_ref : `DatasetRef`
161 Butler DatasetRef
163 Returns
164 -------
165 rb : `ResourceBundle`
166 ResourceBundle consolidating dataset id and DatasetRef
167 """
168 logging.debug("%s", dataset_ref.to_json())
169 if self.butler: 169 ↛ 173line 169 didn't jump to line 173 because the condition on line 169 was always true
170 with self._butler_lock:
171 uri = self.butler.getURI(dataset_ref)
172 else:
173 uri = None
174 did = self._make_did(uri, dataset_ref.to_json())
175 rb = ResourceBundle(dataset_id=dataset_id, did=did)
176 return rb
178 def _make_zip_bundle(self, dataset_id: str, resource_path: ResourcePath) -> ResourceBundle:
179 """Make a ResourceBundle
181 Parameters
182 ----------
183 dataset_id : `str`
184 Rucio dataset name
185 resouce_path : `ResourcePath`
186 ResourcePath to a file
188 Returns
189 -------
190 rb: ResourceBundle
191 ResourceBundle consolidating dataset id and ResourcePath
192 """
193 did = self._make_did(resource_path)
194 rb = ResourceBundle(dataset_id=dataset_id, did=did)
195 return rb
197 def _make_dim_bundle(self, dataset_id: str, resource_path: ResourcePath) -> ResourceBundle:
198 """Make a ResourceBundle
200 Parameters
201 ----------
202 dataset_id : `str`
203 Rucio dataset name
204 resouce_path : `lsst.resource.ResourcePath`
205 ResourcePath to a file
207 Returns
208 -------
209 rb: `lsst.rucio.register.rucio_bundle.ResourceBundle`
210 ResourceBundle consolidating dataset id and ResourcePath
211 """
212 did = self._make_did(resource_path)
213 rb = ResourceBundle(dataset_id=dataset_id, did=did)
214 return rb
216 def compute_hashes(self, resource_path: ResourcePath) -> tuple[int, str]:
217 """return the length and adler32 hash for a file.
219 Parameters
220 ----------
221 path: `lsst.resources.ResourcePath`
222 Path to the file.
224 Returns
225 -------
226 hashes: `tuple` [ `int`, `str` ]
227 Size in bytes and Adler32 hex hash.
228 """
229 # Translate to local file if available to bypass HTTP/redirect overhead
230 if resource_path.scheme in ("davs", "https", "http"): 230 ↛ 231line 230 didn't jump to line 231 because the condition on line 230 was never true
231 local_path = resource_path.unquoted_path
232 if os.path.exists(local_path):
233 logger.debug("Translating remote URI %s to local path %s", resource_path, local_path)
234 resource_path = ResourcePath(f"file://{local_path}")
236 cache_key = str(resource_path)
237 with self._lock:
238 if cache_key in self._hash_cache:
239 size, adler32 = self._hash_cache[cache_key]
240 logger.debug(
241 "Read adler32 %s from internal hash cache for %s",
242 adler32,
243 resource_path,
244 )
245 self._checksum_sources["internal hash cache"] += 1
246 return size, adler32
248 size = -1
249 adler32 = None
251 # 1. If it is a local file, try to read the extended attributes
252 # (user.XrdCks.adler32) to get the precomputed checksum instantly
253 # from CephFS.
254 if resource_path.scheme == "file":
255 local_path = resource_path.unquoted_path
256 try:
257 # user.XrdCks.adler32 is a binary block containing the
258 # adler32 checksum.
259 attr_val = os.getxattr(local_path, "user.XrdCks.adler32")
260 if len(attr_val) >= 36:
261 # Bytes 32-35 contain the 4-byte big-endian adler32
262 # checksum.
263 adler32 = attr_val[32:36].hex()
264 size = os.path.getsize(local_path)
265 logger.debug(
266 "Read adler32 %s from CephFS extended attribute for %s",
267 adler32,
268 local_path,
269 )
270 with self._lock:
271 self._checksum_sources["CephFS xattr"] += 1
272 except Exception as e:
273 logger.debug("Failed to read xattr for %s: %s", local_path, e)
275 # 2. If it is a remote URL and we haven't found the checksum, try
276 # to query the HTTP Digest header using lsst.resources's own
277 # WebDAV client pool.
278 if adler32 is None and resource_path.scheme in ("davs", "https", "http"): 278 ↛ 279line 278 didn't jump to line 279 because the condition on line 278 was never true
279 try:
280 from lsst.resources.davutils import DavClientPool
282 pool = DavClientPool._instance
283 if pool:
284 client = pool.get_client_for_url(str(resource_path))
285 if client:
286 resp = client.head(str(resource_path), headers={"Want-Digest": "adler32"})
287 digest = resp.headers.get("Digest")
288 if digest:
289 parts = digest.split("=")
290 if len(parts) == 2 and parts[0].lower() == "adler32":
291 adler32 = parts[1].strip().lower()
292 size = int(resp.headers.get("Content-Length", -1))
293 logger.debug(
294 "Read adler32 %s from HTTP Digest header (Want-Digest) for %s",
295 adler32,
296 resource_path,
297 )
298 with self._lock:
299 self._checksum_sources["HTTP Digest header"] += 1
300 except Exception as e:
301 logger.debug("Failed to query HTTP Digest for %s: %s", resource_path, e)
303 # 3. Fallback to standard lsst.resources info/get_info() or download
304 if adler32 is None: 304 ↛ 332line 304 didn't jump to line 332 because the condition on line 304 was always true
305 info = resource_path.get_info()
306 size = info.size
307 checksums = getattr(info, "checksums", None) or {}
309 if isinstance(checksums, dict) and "adler32" in checksums:
310 adler32 = checksums["adler32"]
311 logger.debug(
312 "Read adler32 %s from standard WebDAV metadata for %s",
313 adler32,
314 resource_path,
315 )
316 with self._lock:
317 self._checksum_sources["standard WebDAV metadata"] += 1
318 else:
319 logger.debug(
320 "adler32 metadata not available for %s; computing from file contents",
321 resource_path,
322 )
323 adler32 = self._compute_adler32(resource_path)
324 logger.debug(
325 "Computed adler32 %s from file contents for %s",
326 adler32,
327 resource_path,
328 )
329 with self._lock:
330 self._checksum_sources["computed from file contents"] += 1
332 with self._lock:
333 self._hash_cache[cache_key] = (size, adler32)
334 return size, adler32
336 def _compute_adler32(self, resource_path: ResourcePath) -> str:
337 logger.debug("computing adler32 for %s", resource_path)
338 adler32 = zlib.adler32(b"")
339 buffer_size = 10 * 1024 * 1024
340 with resource_path.open("rb") as f:
341 while buffer := f.read(buffer_size):
342 adler32 = zlib.adler32(buffer, adler32)
343 adler32_digest = f"{adler32:08x}"
344 return adler32_digest
346 def _make_did(self, resource_path: ResourcePath, metadata: str = None) -> RucioDID:
347 """Make a Rucio data identifier dictionary from a resource.
349 Parameters
350 ----------
351 resource_path: ResourcePath
352 ResourcePath object
354 metadata: `str`
355 String containing Rubin dataset specific metadata
357 Returns
358 -------
359 did : `dict` [`str`, `str`|`int`]
360 Rucio data identifier including physical and logical names,
361 byte length, adler32 checksum, meta, and scope.
362 """
364 size, adler32 = self.compute_hashes(resource_path)
365 path = resource_path.unquoted_path.removeprefix(self.rse_root)
366 pfn = self.pfn_base + path
367 logging.debug("pfn=%s", pfn)
368 name = path.removeprefix("/" + self.scope + "/")
369 logging.debug("name=%s", name)
370 logging.debug("path=%s", path)
372 if metadata: 372 ↛ 375line 372 didn't jump to line 375 because the condition on line 372 was always true
373 meta = RubinMeta(rubin_butler=self.rubin_butler_type, rubin_sidecar=metadata)
374 else:
375 meta = RubinMeta(rubin_butler=self.rubin_butler_type, rubin_sidecar="")
376 d = RucioDID(
377 pfn=pfn,
378 bytes=size,
379 adler32=adler32,
380 name=name,
381 scope=self.scope,
382 meta=meta,
383 )
385 return d
387 @backoff.on_exception(
388 backoff.expo,
389 RETRYABLE,
390 factor=lambda: _BACKOFF[_FACTOR],
391 max_value=lambda: _BACKOFF[_MAX_VALUE],
392 max_tries=lambda: _BACKOFF[_MAX_TRIES],
393 jitter=None,
394 on_backoff=_backoff_message,
395 )
396 def _add_replicas(self, bundles: list[ResourceBundle]) -> None:
397 """Call the Rucio method add_replica for a list of DIDs
399 Parameters
400 ----------
401 bundles : `list` [`ResourceBundle`]
402 A list of ResourceBundles
403 """
404 dids = [bundle.get_did() for bundle in bundles]
405 self.replica_client.add_replicas(rse=self.rse, files=dids)
407 @backoff.on_exception(
408 backoff.expo,
409 RETRYABLE,
410 factor=lambda: _BACKOFF[_FACTOR],
411 max_value=lambda: _BACKOFF[_MAX_VALUE],
412 max_tries=lambda: _BACKOFF[_MAX_TRIES],
413 jitter=None,
414 on_backoff=_backoff_message,
415 )
416 def _add_files_to_dataset(self, dataset_id: str, dids: list[dict]) -> None:
417 """Attach a list of files specified by Rucio DIDs to a Rucio dataset.
419 Ignores already-attached files for idempotency.
421 Parameters
422 ----------
423 dataset_id : `str`
424 Logical name of the Rucio dataset.
425 dids : `list` [`dict` [`str`, `str`|`int`] ]
426 List of Rucio data identifiers.
427 """
428 try:
429 self.did_client.add_files_to_datasets(
430 attachments=[
431 {
432 "scope": self.scope,
433 "name": dataset_id,
434 "dids": dids,
435 "rse": self.rse,
436 }
437 ],
438 ignore_duplicate=True,
439 )
440 return
441 except rucio.common.exception.DataIdentifierNotFound as e:
442 raise e
444 @backoff.on_exception(
445 backoff.expo,
446 RETRYABLE,
447 factor=lambda: _BACKOFF[_FACTOR],
448 max_value=lambda: _BACKOFF[_MAX_VALUE],
449 max_tries=lambda: _BACKOFF[_MAX_TRIES],
450 jitter=None,
451 on_backoff=_backoff_message,
452 )
453 def _add_dataset_with_retries(self, dataset_id: str, statuses: dict) -> None:
454 try:
455 self.did_client.add_dataset(
456 scope=self.scope,
457 name=dataset_id,
458 statuses=statuses,
459 rse=self.rse,
460 )
461 return
462 except rucio.common.exception.DataIdentifierAlreadyExists as e:
463 # If someone else created it in the meantime
464 raise e
466 def register_to_dataset(self, bundles) -> None:
467 """Register a list of files in Rucio.
469 Parameters
470 ----------
471 bundles : `list` [`ResourceBundle`]
472 List of resource bundles
473 """
474 logger.debug("register to dataset")
476 datasets = dict()
477 for bundle in bundles:
478 dataset_id = bundle.dataset_id
479 datasets.setdefault(dataset_id, []).append(bundle)
481 for dataset_id, bundles in datasets.items():
482 try:
483 dids = [rb.get_did() for rb in bundles]
484 names = [did["pfn"] for did in dids]
485 logger.info("Registering %s in dataset %s, RSE %s", names, dataset_id, self.rse)
486 self._add_files_to_dataset(dataset_id, dids)
487 except rucio.common.exception.DataIdentifierNotFound:
488 # No such dataset, so create it
489 try:
490 logger.info("Couldn't register because dataset not yet registered")
491 logger.info("Creating Rucio dataset %s", dataset_id)
492 self._add_dataset_with_retries(
493 dataset_id=dataset_id,
494 statuses={"monotonic": True},
495 )
496 except rucio.common.exception.DataIdentifierAlreadyExists:
497 # If someone else created it in the meantime
498 pass
499 # And then retry adding DIDs
500 logger.info("Dataset registered.")
501 logger.info("Retrying registering %s in dataset %s, RSE %s", names, dataset_id, self.rse)
502 self._add_files_to_dataset(dataset_id, dids)
504 if self.clear_is_new:
505 for dataset_id in datasets.keys():
506 self.clear_is_new_metadata(dataset_id)
508 logger.debug("Done with Rucio for %s", bundles)
510 def clear_is_new_metadata(self, dataset_id: str) -> None:
511 """Set is_new metadata to None for a Rucio dataset DID.
513 Parameters
514 ----------
515 dataset_id : `str`
516 Logical name of the Rucio dataset.
517 """
518 try:
519 logger.info("Setting is_new metadata to None for dataset %s:%s", self.scope, dataset_id)
520 self.did_client.set_metadata(scope=self.scope, name=dataset_id, key="is_new", value=None)
521 except Exception as e:
522 logger.warning("Failed to clear is_new metadata for dataset %s: %s", dataset_id, e)
524 def register_as_replicas(self, dataset_id, dataset_refs) -> int:
525 """Register a list of DatasetRefs to a Rucio dataset
527 Parameters
528 ----------
529 dataset_id : `str`
530 RUCIO dataset id
531 dataset_refs : `list` [`DatasetRef`]
532 list of Butler DatasetRefs
533 """
534 flat_refs = []
535 for dataset_ref in dataset_refs:
536 if isinstance(dataset_ref, list): 536 ↛ 537line 536 didn't jump to line 537 because the condition on line 536 was never true
537 flat_refs.extend(dataset_ref)
538 else:
539 flat_refs.append(dataset_ref)
541 if len(flat_refs) == 0: 541 ↛ 542line 541 didn't jump to line 542 because the condition on line 541 was never true
542 return 0
544 # Bulk query URIs from Butler if butler is available
545 uris_dict = None
546 if self.butler: 546 ↛ 554line 546 didn't jump to line 554 because the condition on line 546 was always true
547 try:
548 uris_dict = self.butler.get_many_uris(flat_refs)
549 except Exception as e:
550 logger.debug("butler.get_many_uris failed, falling back to sequential getURI: %s", e)
552 # Thread pool to parallelize metadata and checksum queries
553 # (HEAD requests)
554 from concurrent.futures import ThreadPoolExecutor
556 def make_bundle(ref):
557 if uris_dict and ref in uris_dict: 557 ↛ 558line 557 didn't jump to line 558 because the condition on line 557 was never true
558 uri_obj = uris_dict[ref]
559 primary_uri = uri_obj.primaryURI if hasattr(uri_obj, "primaryURI") else uri_obj[0]
560 logging.debug("%s", ref.to_json())
561 did = self._make_did(primary_uri, ref.to_json())
562 return ResourceBundle(dataset_id=dataset_id, did=did)
563 else:
564 return self._make_dataset_ref_bundle(dataset_id, ref)
566 with ThreadPoolExecutor(max_workers=min(len(flat_refs), 10)) as thread_executor:
567 bundles = list(thread_executor.map(make_bundle, flat_refs))
569 self._add_replicas(bundles)
570 self.register_to_dataset(bundles)
572 # Log summary of checksum sources used in this batch
573 with self._lock:
574 sources = [f"'{src}': {cnt}" for src, cnt in self._checksum_sources.items()]
575 summary_str = ", ".join(sources)
576 self._checksum_sources.clear()
577 logger.info("Batch checksum sources used: {%s}", summary_str)
579 return len(bundles)
581 def register_zips(self, dataset_id: str, zip_files: list) -> int:
582 """Register a list of zips to a Rucio Dataset
584 Parameters
585 ----------
586 dataset_id : `str`
587 RUCIO dataset id
588 zip_files : `list` [`ResourcePath`]
589 list of ResourcePath
591 Returns
592 -------
593 num : `int`
594 number of zip files ingested
595 """
596 bundles = []
597 for zip_file in zip_files:
598 bundles.append(self._make_zip_bundle(dataset_id, zip_file))
599 self._add_replicas(bundles)
600 self.register_to_dataset(bundles)
601 return len(bundles)
603 def register_dims(self, dataset_id: str, dim_files: list) -> int:
604 """Register a list of dimension files to a Rucio Dataset
606 Parameters
607 ----------
608 dataset_id : `str`
609 RUCIO dataset id
610 dim_files : `list` [`lsst.resource.ResourcePath`]
611 list of ResourcePath
613 Returns
614 -------
615 num : `int`
616 number of dimension files ingested
617 """
618 bundles = []
619 for dim_file in dim_files:
620 bundles.append(self._make_dim_bundle(dataset_id, dim_file))
621 self._add_replicas(bundles)
622 self.register_to_dataset(bundles)
623 return len(bundles)
626_BACKOFF[_FACTOR] = RucioInterface.DEFAULT_BACKOFF_FACTOR
627_BACKOFF[_MAX_VALUE] = RucioInterface.DEFAULT_BACKOFF_MAX_VALUE
628_BACKOFF[_MAX_TRIES] = RucioInterface.DEFAULT_RETRIES