Coverage for python/lsst/rucio/register/rucio_interface.py: 72%
257 statements
« prev ^ index » next coverage.py v7.16.2, created at 2026-11-03 01:33 -0800
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 """
90 DEFAULT_BACKOFF_FACTOR = 5.0
91 DEFAULT_BACKOFF_MAX_VALUE = 30
92 DEFAULT_RETRIES = 8
94 def __init__(
95 self,
96 butler: lsst.daf.butler.Butler,
97 rucio_rse: str,
98 scope: str,
99 rse_root: str,
100 dtn_url: str,
101 rubin_butler_type: str,
102 clear_is_new: bool = False,
103 retries: int = DEFAULT_RETRIES,
104 ):
105 self.butler = butler
106 self.rse = rucio_rse
107 self.scope = scope
108 self.rse_root = rse_root
109 self.dtn_url = dtn_url
110 self.pfn_base = f"{dtn_url}"
111 self.replica_client = ReplicaClient(timeout=15.0)
112 self.did_client = DIDClient(timeout=15.0)
113 self.rubin_butler_type = rubin_butler_type
114 self.clear_is_new = clear_is_new
115 self.retries = retries
116 self._hash_cache: dict[str, tuple[int, str]] = {}
117 self._checksum_sources: collections.Counter = collections.Counter()
118 self._lock = threading.Lock()
119 self._butler_lock = threading.Lock()
121 def set_backoff(self, factor, max_value, max_tries) -> None:
122 """Set backoff values for retries
124 Parameters
125 ----------
126 factor: `float`
127 Multipler for backoff
128 max_value: `int`
129 Maximum seconds to backoff to
130 max_tries: `int`
131 Maximum times to try
132 """
133 logger.debug("factor=%f, max_value=%d, max_tries=%d", factor, max_value, max_tries)
134 if factor is not None: 134 ↛ 136line 134 didn't jump to line 136 because the condition on line 134 was always true
135 _BACKOFF[_FACTOR] = factor
136 if max_value is not None: 136 ↛ 138line 136 didn't jump to line 138 because the condition on line 136 was always true
137 _BACKOFF[_MAX_VALUE] = max_value
138 if max_tries is not None: 138 ↛ exitline 138 didn't return from function 'set_backoff' because the condition on line 138 was always true
139 _BACKOFF[_MAX_TRIES] = max_tries
140 self.retries = max_tries
142 def _make_dataset_ref_bundle(self, dataset_id: str, dataset_ref: DatasetRef) -> ResourceBundle:
143 """Make a ResourceBundle
145 Parameters
146 ----------
147 dataset_id : `str`
148 Rucio dataset name
149 dataset_ref : `DatasetRef`
150 Butler DatasetRef
152 Returns
153 -------
154 rb : `ResourceBundle`
155 ResourceBundle consolidating dataset id and DatasetRef
156 """
157 logging.debug("%s", dataset_ref.to_json())
158 if self.butler: 158 ↛ 162line 158 didn't jump to line 162 because the condition on line 158 was always true
159 with self._butler_lock:
160 uri = self.butler.getURI(dataset_ref)
161 else:
162 uri = None
163 did = self._make_did(uri, dataset_ref.to_json())
164 rb = ResourceBundle(dataset_id=dataset_id, did=did)
165 return rb
167 def _make_zip_bundle(self, dataset_id: str, resource_path: ResourcePath) -> ResourceBundle:
168 """Make a ResourceBundle
170 Parameters
171 ----------
172 dataset_id : `str`
173 Rucio dataset name
174 resouce_path : `ResourcePath`
175 ResourcePath to a file
177 Returns
178 -------
179 rb: ResourceBundle
180 ResourceBundle consolidating dataset id and ResourcePath
181 """
182 did = self._make_did(resource_path)
183 rb = ResourceBundle(dataset_id=dataset_id, did=did)
184 return rb
186 def _make_dim_bundle(self, dataset_id: str, resource_path: ResourcePath) -> ResourceBundle:
187 """Make a ResourceBundle
189 Parameters
190 ----------
191 dataset_id : `str`
192 Rucio dataset name
193 resouce_path : `lsst.resource.ResourcePath`
194 ResourcePath to a file
196 Returns
197 -------
198 rb: `lsst.rucio.register.rucio_bundle.ResourceBundle`
199 ResourceBundle consolidating dataset id and ResourcePath
200 """
201 did = self._make_did(resource_path)
202 rb = ResourceBundle(dataset_id=dataset_id, did=did)
203 return rb
205 def compute_hashes(self, resource_path: ResourcePath) -> tuple[int, str]:
206 """return the length and adler32 hash for a file.
208 Parameters
209 ----------
210 path: `lsst.resources.ResourcePath`
211 Path to the file.
213 Returns
214 -------
215 hashes: `tuple` [ `int`, `str` ]
216 Size in bytes and Adler32 hex hash.
217 """
218 # Translate to local file if available to bypass HTTP/redirect overhead
219 if resource_path.scheme in ("davs", "https", "http"): 219 ↛ 220line 219 didn't jump to line 220 because the condition on line 219 was never true
220 local_path = resource_path.unquoted_path
221 if os.path.exists(local_path):
222 logger.debug("Translating remote URI %s to local path %s", resource_path, local_path)
223 resource_path = ResourcePath(f"file://{local_path}")
225 cache_key = str(resource_path)
226 with self._lock:
227 if cache_key in self._hash_cache:
228 size, adler32 = self._hash_cache[cache_key]
229 logger.debug(
230 "Read adler32 %s from internal hash cache for %s",
231 adler32,
232 resource_path,
233 )
234 self._checksum_sources["internal hash cache"] += 1
235 return size, adler32
237 size = -1
238 adler32 = None
240 # 1. If it is a local file, try to read the extended attributes
241 # (user.XrdCks.adler32) to get the precomputed checksum instantly
242 # from CephFS.
243 if resource_path.scheme == "file":
244 local_path = resource_path.unquoted_path
245 try:
246 # user.XrdCks.adler32 is a binary block containing the
247 # adler32 checksum.
248 attr_val = os.getxattr(local_path, "user.XrdCks.adler32")
249 if len(attr_val) >= 36:
250 # Bytes 32-35 contain the 4-byte big-endian adler32
251 # checksum.
252 adler32 = attr_val[32:36].hex()
253 size = os.path.getsize(local_path)
254 logger.debug(
255 "Read adler32 %s from CephFS extended attribute for %s",
256 adler32,
257 local_path,
258 )
259 with self._lock:
260 self._checksum_sources["CephFS xattr"] += 1
261 except Exception as e:
262 logger.debug("Failed to read xattr for %s: %s", local_path, e)
264 # 2. If it is a remote URL and we haven't found the checksum, try
265 # to query the HTTP Digest header using lsst.resources's own
266 # WebDAV client pool.
267 if adler32 is None and resource_path.scheme in ("davs", "https", "http"): 267 ↛ 268line 267 didn't jump to line 268 because the condition on line 267 was never true
268 try:
269 from lsst.resources.davutils import DavClientPool
271 pool = DavClientPool._instance
272 if pool:
273 client = pool.get_client_for_url(str(resource_path))
274 if client:
275 resp = client.head(str(resource_path), headers={"Want-Digest": "adler32"})
276 digest = resp.headers.get("Digest")
277 if digest:
278 parts = digest.split("=")
279 if len(parts) == 2 and parts[0].lower() == "adler32":
280 adler32 = parts[1].strip().lower()
281 size = int(resp.headers.get("Content-Length", -1))
282 logger.debug(
283 "Read adler32 %s from HTTP Digest header (Want-Digest) for %s",
284 adler32,
285 resource_path,
286 )
287 with self._lock:
288 self._checksum_sources["HTTP Digest header"] += 1
289 except Exception as e:
290 logger.debug("Failed to query HTTP Digest for %s: %s", resource_path, e)
292 # 3. Fallback to standard lsst.resources info/get_info() or download
293 if adler32 is None: 293 ↛ 321line 293 didn't jump to line 321 because the condition on line 293 was always true
294 info = resource_path.get_info()
295 size = info.size
296 checksums = getattr(info, "checksums", None) or {}
298 if isinstance(checksums, dict) and "adler32" in checksums:
299 adler32 = checksums["adler32"]
300 logger.debug(
301 "Read adler32 %s from standard WebDAV metadata for %s",
302 adler32,
303 resource_path,
304 )
305 with self._lock:
306 self._checksum_sources["standard WebDAV metadata"] += 1
307 else:
308 logger.debug(
309 "adler32 metadata not available for %s; computing from file contents",
310 resource_path,
311 )
312 adler32 = self._compute_adler32(resource_path)
313 logger.debug(
314 "Computed adler32 %s from file contents for %s",
315 adler32,
316 resource_path,
317 )
318 with self._lock:
319 self._checksum_sources["computed from file contents"] += 1
321 with self._lock:
322 self._hash_cache[cache_key] = (size, adler32)
323 return size, adler32
325 def _compute_adler32(self, resource_path: ResourcePath) -> str:
326 logger.debug("computing adler32 for %s", resource_path)
327 adler32 = zlib.adler32(b"")
328 buffer_size = 10 * 1024 * 1024
329 with resource_path.open("rb") as f:
330 while buffer := f.read(buffer_size):
331 adler32 = zlib.adler32(buffer, adler32)
332 adler32_digest = f"{adler32:08x}"
333 return adler32_digest
335 def _make_did(self, resource_path: ResourcePath, metadata: str = None) -> RucioDID:
336 """Make a Rucio data identifier dictionary from a resource.
338 Parameters
339 ----------
340 resource_path: ResourcePath
341 ResourcePath object
343 metadata: `str`
344 String containing Rubin dataset specific metadata
346 Returns
347 -------
348 did : `dict` [`str`, `str`|`int`]
349 Rucio data identifier including physical and logical names,
350 byte length, adler32 checksum, meta, and scope.
351 """
353 size, adler32 = self.compute_hashes(resource_path)
354 path = resource_path.unquoted_path.removeprefix(self.rse_root)
355 pfn = self.pfn_base + path
356 logging.debug("pfn=%s", pfn)
357 name = path.removeprefix("/" + self.scope + "/")
358 logging.debug("name=%s", name)
359 logging.debug("path=%s", path)
361 if metadata: 361 ↛ 364line 361 didn't jump to line 364 because the condition on line 361 was always true
362 meta = RubinMeta(rubin_butler=self.rubin_butler_type, rubin_sidecar=metadata)
363 else:
364 meta = RubinMeta(rubin_butler=self.rubin_butler_type, rubin_sidecar="")
365 d = RucioDID(
366 pfn=pfn,
367 bytes=size,
368 adler32=adler32,
369 name=name,
370 scope=self.scope,
371 meta=meta,
372 )
374 return d
376 @backoff.on_exception(
377 backoff.expo,
378 RETRYABLE,
379 factor=lambda: _BACKOFF[_FACTOR],
380 max_value=lambda: _BACKOFF[_MAX_VALUE],
381 max_tries=lambda: _BACKOFF[_MAX_TRIES],
382 jitter=None,
383 on_backoff=_backoff_message,
384 )
385 def _add_replicas(self, bundles: list[ResourceBundle]) -> None:
386 """Call the Rucio method add_replica for a list of DIDs
388 Parameters
389 ----------
390 bundles : `list` [`ResourceBundle`]
391 A list of ResourceBundles
392 """
393 dids = [bundle.get_did() for bundle in bundles]
394 self.replica_client.add_replicas(rse=self.rse, files=dids)
396 @backoff.on_exception(
397 backoff.expo,
398 RETRYABLE,
399 factor=lambda: _BACKOFF[_FACTOR],
400 max_value=lambda: _BACKOFF[_MAX_VALUE],
401 max_tries=lambda: _BACKOFF[_MAX_TRIES],
402 jitter=None,
403 on_backoff=_backoff_message,
404 )
405 def _add_files_to_dataset(self, dataset_id: str, dids: list[dict]) -> None:
406 """Attach a list of files specified by Rucio DIDs to a Rucio dataset.
408 Ignores already-attached files for idempotency.
410 Parameters
411 ----------
412 dataset_id : `str`
413 Logical name of the Rucio dataset.
414 dids : `list` [`dict` [`str`, `str`|`int`] ]
415 List of Rucio data identifiers.
416 """
417 try:
418 self.did_client.add_files_to_datasets(
419 attachments=[
420 {
421 "scope": self.scope,
422 "name": dataset_id,
423 "dids": dids,
424 "rse": self.rse,
425 }
426 ],
427 ignore_duplicate=True,
428 )
429 return
430 except rucio.common.exception.DataIdentifierNotFound as e:
431 raise e
433 @backoff.on_exception(
434 backoff.expo,
435 RETRYABLE,
436 factor=lambda: _BACKOFF[_FACTOR],
437 max_value=lambda: _BACKOFF[_MAX_VALUE],
438 max_tries=lambda: _BACKOFF[_MAX_TRIES],
439 jitter=None,
440 on_backoff=_backoff_message,
441 )
442 def _add_dataset_with_retries(self, dataset_id: str, statuses: dict) -> None:
443 try:
444 self.did_client.add_dataset(
445 scope=self.scope,
446 name=dataset_id,
447 statuses=statuses,
448 rse=self.rse,
449 )
450 return
451 except rucio.common.exception.DataIdentifierAlreadyExists as e:
452 # If someone else created it in the meantime
453 raise e
455 def register_to_dataset(self, bundles) -> None:
456 """Register a list of files in Rucio.
458 Parameters
459 ----------
460 bundles : `list` [`ResourceBundle`]
461 List of resource bundles
462 """
463 logger.debug("register to dataset")
465 datasets = dict()
466 for bundle in bundles:
467 dataset_id = bundle.dataset_id
468 datasets.setdefault(dataset_id, []).append(bundle)
470 for dataset_id, bundles in datasets.items():
471 try:
472 dids = [rb.get_did() for rb in bundles]
473 names = [did["pfn"] for did in dids]
474 logger.info("Registering %s in dataset %s, RSE %s", names, dataset_id, self.rse)
475 self._add_files_to_dataset(dataset_id, dids)
476 except rucio.common.exception.DataIdentifierNotFound:
477 # No such dataset, so create it
478 try:
479 logger.info("Couldn't register because dataset not yet registered")
480 logger.info("Creating Rucio dataset %s", dataset_id)
481 self._add_dataset_with_retries(
482 dataset_id=dataset_id,
483 statuses={"monotonic": True},
484 )
485 except rucio.common.exception.DataIdentifierAlreadyExists:
486 # If someone else created it in the meantime
487 pass
488 # And then retry adding DIDs
489 logger.info("Dataset registered.")
490 logger.info("Retrying registering %s in dataset %s, RSE %s", names, dataset_id, self.rse)
491 self._add_files_to_dataset(dataset_id, dids)
493 if self.clear_is_new:
494 for dataset_id in datasets.keys():
495 self.clear_is_new_metadata(dataset_id)
497 logger.debug("Done with Rucio for %s", bundles)
499 def clear_is_new_metadata(self, dataset_id: str) -> None:
500 """Set is_new metadata to None for a Rucio dataset DID.
502 Parameters
503 ----------
504 dataset_id : `str`
505 Logical name of the Rucio dataset.
506 """
507 try:
508 logger.info("Setting is_new metadata to None for dataset %s:%s", self.scope, dataset_id)
509 self.did_client.set_metadata(scope=self.scope, name=dataset_id, key="is_new", value=None)
510 except Exception as e:
511 logger.warning("Failed to clear is_new metadata for dataset %s: %s", dataset_id, e)
513 def register_as_replicas(self, dataset_id, dataset_refs) -> int:
514 """Register a list of DatasetRefs to a Rucio dataset
516 Parameters
517 ----------
518 dataset_id : `str`
519 RUCIO dataset id
520 dataset_refs : `list` [`DatasetRef`]
521 list of Butler DatasetRefs
522 """
523 flat_refs = []
524 for dataset_ref in dataset_refs:
525 if isinstance(dataset_ref, list): 525 ↛ 526line 525 didn't jump to line 526 because the condition on line 525 was never true
526 flat_refs.extend(dataset_ref)
527 else:
528 flat_refs.append(dataset_ref)
530 if len(flat_refs) == 0: 530 ↛ 531line 530 didn't jump to line 531 because the condition on line 530 was never true
531 return 0
533 # Bulk query URIs from Butler if butler is available
534 uris_dict = None
535 if self.butler: 535 ↛ 543line 535 didn't jump to line 543 because the condition on line 535 was always true
536 try:
537 uris_dict = self.butler.get_many_uris(flat_refs)
538 except Exception as e:
539 logger.debug("butler.get_many_uris failed, falling back to sequential getURI: %s", e)
541 # Thread pool to parallelize metadata and checksum queries
542 # (HEAD requests)
543 from concurrent.futures import ThreadPoolExecutor
545 def make_bundle(ref):
546 if uris_dict and ref in uris_dict: 546 ↛ 547line 546 didn't jump to line 547 because the condition on line 546 was never true
547 uri_obj = uris_dict[ref]
548 primary_uri = uri_obj.primaryURI if hasattr(uri_obj, "primaryURI") else uri_obj[0]
549 logging.debug("%s", ref.to_json())
550 did = self._make_did(primary_uri, ref.to_json())
551 return ResourceBundle(dataset_id=dataset_id, did=did)
552 else:
553 return self._make_dataset_ref_bundle(dataset_id, ref)
555 with ThreadPoolExecutor(max_workers=min(len(flat_refs), 10)) as thread_executor:
556 bundles = list(thread_executor.map(make_bundle, flat_refs))
558 self._add_replicas(bundles)
559 self.register_to_dataset(bundles)
561 # Log summary of checksum sources used in this batch
562 with self._lock:
563 sources = [f"'{src}': {cnt}" for src, cnt in self._checksum_sources.items()]
564 summary_str = ", ".join(sources)
565 self._checksum_sources.clear()
566 logger.info("Batch checksum sources used: {%s}", summary_str)
568 return len(bundles)
570 def register_zips(self, dataset_id: str, zip_files: list) -> int:
571 """Register a list of zips to a Rucio Dataset
573 Parameters
574 ----------
575 dataset_id : `str`
576 RUCIO dataset id
577 zip_files : `list` [`ResourcePath`]
578 list of ResourcePath
580 Returns
581 -------
582 num : `int`
583 number of zip files ingested
584 """
585 bundles = []
586 for zip_file in zip_files:
587 bundles.append(self._make_zip_bundle(dataset_id, zip_file))
588 self._add_replicas(bundles)
589 self.register_to_dataset(bundles)
590 return len(bundles)
592 def register_dims(self, dataset_id: str, dim_files: list) -> int:
593 """Register a list of dimension files to a Rucio Dataset
595 Parameters
596 ----------
597 dataset_id : `str`
598 RUCIO dataset id
599 dim_files : `list` [`lsst.resource.ResourcePath`]
600 list of ResourcePath
602 Returns
603 -------
604 num : `int`
605 number of dimension files ingested
606 """
607 bundles = []
608 for dim_file in dim_files:
609 bundles.append(self._make_dim_bundle(dataset_id, dim_file))
610 self._add_replicas(bundles)
611 self.register_to_dataset(bundles)
612 return len(bundles)
615_BACKOFF[_FACTOR] = RucioInterface.DEFAULT_BACKOFF_FACTOR
616_BACKOFF[_MAX_VALUE] = RucioInterface.DEFAULT_BACKOFF_MAX_VALUE
617_BACKOFF[_MAX_TRIES] = RucioInterface.DEFAULT_RETRIES