Coverage for python/lsst/rucio/register/rucio_interface.py: 76%
150 statements
« prev ^ index » next coverage.py v7.15.2, created at 2026-07-25 08:38 +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 logging
23import zlib
25import backoff
26import requests
27import rucio.common.exception
28import urllib3.exceptions
29from rucio.client.didclient import DIDClient
30from rucio.client.replicaclient import ReplicaClient
32import lsst.daf.butler
33from lsst.daf.butler import DatasetRef
34from lsst.resources import ResourcePath
35from lsst.rucio.register.resource_bundle import ResourceBundle
36from lsst.rucio.register.rubin_meta import RubinMeta
37from lsst.rucio.register.rucio_did import RucioDID
39__all__ = ["RucioInterface"]
41RETRYABLE = (
42 urllib3.exceptions.ReadTimeoutError,
43 urllib3.exceptions.ProtocolError,
44 requests.exceptions.ConnectionError,
45 requests.exceptions.Timeout,
46 requests.exceptions.ChunkedEncodingError,
47 rucio.common.exception.ServerConnectionException,
48 rucio.common.exception.DatabaseException,
49)
51logger = logging.getLogger(__name__)
53_FACTOR = "factor"
54_MAX_VALUE = "max_value"
55_MAX_TRIES = "max_tries"
56_BACKOFF = {_FACTOR: 5, _MAX_VALUE: 30, _MAX_TRIES: 5}
59def _backoff_message(details):
60 logger.info(
61 "Backing off {wait:0.1f} seconds after {tries} tries "
62 "calling function {target} with args {args} and kwargs "
63 "{kwargs}".format(**details)
64 )
67class RucioInterface:
68 """Add files as replicas in Rucio, along with metadata,
69 and attach them to datasets.
71 Parameters
72 ----------
73 butler : `lsst.daf.butler.Butler`
74 Butler we're operating upon
75 rucio_rse : `str`
76 Name of the RSE that the files live in.
77 scope : `str`
78 Rucio scope to register the files in.
79 rse_root : `str`
80 Full path to root directory of RSE directory structure
81 dtn_url : `str`
82 Base URL of the data transfer node for the Rucio physical filename.
83 rubin_butler_type: `str`
84 the type registered in "rubin_butler" metadata for rucio
85 """
87 def __init__(
88 self,
89 butler: lsst.daf.butler.Butler,
90 rucio_rse: str,
91 scope: str,
92 rse_root: str,
93 dtn_url: str,
94 rubin_butler_type: str,
95 ):
96 self.butler = butler
97 self.rse = rucio_rse
98 self.scope = scope
99 self.rse_root = rse_root
100 self.dtn_url = dtn_url
101 self.pfn_base = f"{dtn_url}"
102 self.replica_client = ReplicaClient()
103 self.did_client = DIDClient()
104 self.rubin_butler_type = rubin_butler_type
106 def _make_dataset_ref_bundle(self, dataset_id: str, dataset_ref: DatasetRef) -> ResourceBundle:
107 """Make a ResourceBundle
109 Parameters
110 ----------
111 dataset_id : `str`
112 Rucio dataset name
113 dataset_ref : `DatasetRef`
114 Butler DatasetRef
116 Returns
117 -------
118 rb : `ResourceBundle`
119 ResourceBundle consolidating dataset id and DatasetRef
120 """
121 logging.debug("%s", dataset_ref.to_json())
122 did = self._make_did(self.butler.getURI(dataset_ref), dataset_ref.to_json())
123 rb = ResourceBundle(dataset_id=dataset_id, did=did)
124 return rb
126 def _make_zip_bundle(self, dataset_id: str, resource_path: ResourcePath) -> ResourceBundle:
127 """Make a ResourceBundle
129 Parameters
130 ----------
131 dataset_id : `str`
132 Rucio dataset name
133 resouce_path : `ResourcePath`
134 ResourcePath to a file
136 Returns
137 -------
138 rb: ResourceBundle
139 ResourceBundle consolidating dataset id and ResourcePath
140 """
141 did = self._make_did(resource_path)
142 rb = ResourceBundle(dataset_id=dataset_id, did=did)
143 return rb
145 def _make_dim_bundle(self, dataset_id: str, resource_path: ResourcePath) -> ResourceBundle:
146 """Make a ResourceBundle
148 Parameters
149 ----------
150 dataset_id : `str`
151 Rucio dataset name
152 resouce_path : `lsst.resource.ResourcePath`
153 ResourcePath to a file
155 Returns
156 -------
157 rb: `lsst.rucio.register.rucio_bundle.ResourceBundle`
158 ResourceBundle consolidating dataset id and ResourcePath
159 """
160 did = self._make_did(resource_path)
161 rb = ResourceBundle(dataset_id=dataset_id, did=did)
162 return rb
164 def compute_hashes(self, resource_path: ResourcePath) -> tuple[int, str]:
165 """return the length and adler32 hash for a file.
167 Parameters
168 ----------
169 path: `lsst.resources.ResourcePath`
170 Path to the file.
172 Returns
173 -------
174 hashes: `tuple` [ `int`, `str` ]
175 Size in bytes and Adler32 hex hash.
176 """
178 info = resource_path.get_info()
179 size = info.size
180 checksums = info.checksums
181 if "adler32" in checksums:
182 adler32 = checksums["adler32"]
183 logger.debug("found adler32 for %s", resource_path)
184 return size, adler32
185 return size, self._compute_adler32(resource_path)
187 def _compute_adler32(self, resource_path: ResourcePath) -> tuple[int, str]:
188 logger.debug("computing adler32 for %s", resource_path)
189 adler32 = zlib.adler32(b"")
190 buffer_size = 10 * 1024 * 1024
191 with resource_path.open("rb") as f:
192 while buffer := f.read(buffer_size):
193 adler32 = zlib.adler32(buffer, adler32)
194 adler32_digest = f"{adler32:08x}"
195 return adler32_digest
197 def _make_did(self, resource_path: ResourcePath, metadata: str = None) -> RucioDID:
198 """Make a Rucio data identifier dictionary from a resource.
200 Parameters
201 ----------
202 resource_path: ResourcePath
203 ResourcePath object
205 metadata: `str`
206 String containing Rubin dataset specific metadata
208 Returns
209 -------
210 did : `dict` [`str`, `str`|`int`]
211 Rucio data identifier including physical and logical names,
212 byte length, adler32 checksum, meta, and scope.
213 """
215 size, adler32 = self.compute_hashes(resource_path)
216 path = resource_path.unquoted_path.removeprefix(self.rse_root)
217 pfn = self.pfn_base + path
218 logging.debug("pfn=%s", pfn)
219 name = path.removeprefix("/" + self.scope + "/")
220 logging.debug("name=%s", name)
221 logging.debug("path=%s", path)
223 if metadata: 223 ↛ 226line 223 didn't jump to line 226 because the condition on line 223 was always true
224 meta = RubinMeta(rubin_butler=self.rubin_butler_type, rubin_sidecar=metadata)
225 else:
226 meta = RubinMeta(rubin_butler=self.rubin_butler_type, rubin_sidecar="")
227 d = RucioDID(
228 pfn=pfn,
229 bytes=size,
230 adler32=adler32,
231 name=name,
232 scope=self.scope,
233 meta=meta,
234 )
236 return d
238 @backoff.on_exception(
239 backoff.expo,
240 RETRYABLE,
241 factor=lambda: _BACKOFF[_FACTOR],
242 max_value=lambda: _BACKOFF[_MAX_VALUE],
243 max_tries=lambda: _BACKOFF[_MAX_TRIES],
244 jitter=None,
245 on_backoff=_backoff_message,
246 )
247 def _add_replicas(self, bundles: list[ResourceBundle]) -> None:
248 """Call the Rucio method add_replica for a list of DIDs
250 Parameters
251 ----------
252 bundles : `list` [`ResourceBundle`]
253 A list of ResourceBundles
254 """
256 dids = [bundle.get_did() for bundle in bundles]
258 self.replica_client.add_replicas(rse=self.rse, files=dids)
260 @backoff.on_exception(
261 backoff.expo,
262 RETRYABLE,
263 factor=lambda: _BACKOFF[_FACTOR],
264 max_value=lambda: _BACKOFF[_MAX_VALUE],
265 max_tries=lambda: _BACKOFF[_MAX_TRIES],
266 jitter=None,
267 on_backoff=_backoff_message,
268 )
269 def _add_files_to_dataset(self, dataset_id: str, dids: list[dict]) -> None:
270 """Attach a list of files specified by Rucio DIDs to a Rucio dataset.
272 Ignores already-attached files for idempotency.
274 Parameters
275 ----------
276 dataset_id : `str`
277 Logical name of the Rucio dataset.
278 dids : `list` [`dict` [`str`, `str`|`int`] ]
279 List of Rucio data identifiers.
280 """
281 try:
282 self.did_client.add_files_to_datasets(
283 attachments=[
284 {
285 "scope": self.scope,
286 "name": dataset_id,
287 "dids": dids,
288 "rse": self.rse,
289 }
290 ],
291 ignore_duplicate=True,
292 )
293 return
294 except rucio.common.exception.DataIdentifierNotFound as e:
295 raise e
297 @backoff.on_exception(
298 backoff.expo,
299 RETRYABLE,
300 factor=lambda: _BACKOFF[_FACTOR],
301 max_value=lambda: _BACKOFF[_MAX_VALUE],
302 max_tries=lambda: _BACKOFF[_MAX_TRIES],
303 jitter=None,
304 on_backoff=_backoff_message,
305 )
306 def _add_dataset_with_retries(self, dataset_id: str, statuses: dict) -> None:
307 try:
308 self.did_client.add_dataset(
309 scope=self.scope,
310 name=dataset_id,
311 statuses=statuses,
312 rse=self.rse,
313 )
314 return
315 except rucio.common.exception.DataIdentifierAlreadyExists as e:
316 # If someone else created it in the meantime
317 raise e
319 def register_to_dataset(self, bundles) -> None:
320 """Register a list of files in Rucio.
322 Parameters
323 ----------
324 bundles : `list` [`ResourceBundle`]
325 List of resource bundles
326 """
327 logger.debug("register to dataset")
329 datasets = dict()
330 for bundle in bundles:
331 dataset_id = bundle.dataset_id
332 datasets.setdefault(dataset_id, []).append(bundle)
334 for dataset_id, bundles in datasets.items():
335 try:
336 dids = [rb.get_did() for rb in bundles]
337 names = [did["pfn"] for did in dids]
338 logger.info("Registering %s in dataset %s, RSE %s", names, dataset_id, self.rse)
339 self._add_files_to_dataset(dataset_id, dids)
340 except rucio.common.exception.DataIdentifierNotFound:
341 # No such dataset, so create it
342 try:
343 logger.info("Couldn't register because dataset not yet registered")
344 logger.info("Creating Rucio dataset %s", dataset_id)
345 self._add_dataset_with_retries(
346 dataset_id=dataset_id,
347 statuses={"monotonic": True},
348 )
349 except rucio.common.exception.DataIdentifierAlreadyExists:
350 # If someone else created it in the meantime
351 pass
352 # And then retry adding DIDs
353 logger.info("Dataset registered.")
354 logger.info("Retrying registering %s in dataset %s, RSE %s", names, dataset_id, self.rse)
355 self._add_files_to_dataset(dataset_id, dids)
357 logger.debug("Done with Rucio for %s", bundles)
359 def set_backoff(self, factor, max_value, max_tries) -> None:
360 """Set backoff values for retries
362 Parameters
363 ----------
364 factor: `float`
365 Multipler for backoff
366 max_value: `int`
367 Maximum seconds to backoff to
368 max_tries: `int`
369 Maximum times to try
370 """
371 logger.debug("factor=%f, max_value=%d, max_tries=%d", factor, max_value, max_tries)
372 _BACKOFF[_FACTOR] = factor
373 _BACKOFF[_MAX_VALUE] = max_value
374 _BACKOFF[_MAX_TRIES] = max_tries
376 def register_as_replicas(self, dataset_id, dataset_refs) -> None:
377 """Register a list of DatasetRefs to a Rucio dataset
379 Parameters
380 ----------
381 dataset_id : `str`
382 RUCIO dataset id
383 dataset_refs : `list` [`DatasetRef`]
384 list of Butler DatasetRefs
385 """
386 bundles = []
387 for dataset_ref in dataset_refs:
388 if type(dataset_ref) is list: 388 ↛ 389line 388 didn't jump to line 389 because the condition on line 388 was never true
389 for dsr in dataset_ref:
390 bundles.append(self._make_dataset_ref_bundle(dataset_id, dsr))
391 else:
392 bundles.append(self._make_dataset_ref_bundle(dataset_id, dataset_ref))
393 if len(bundles) == 0: 393 ↛ 394line 393 didn't jump to line 394 because the condition on line 393 was never true
394 return 0
395 self._add_replicas(bundles)
396 self.register_to_dataset(bundles)
397 return len(bundles)
399 def register_zips(self, dataset_id: str, zip_files: list) -> int:
400 """Register a list of zips to a Rucio Dataset
402 Parameters
403 ----------
404 dataset_id : `str`
405 RUCIO dataset id
406 zip_files : `list` [`ResourcePath`]
407 list of ResourcePath
409 Returns
410 -------
411 num : `int`
412 number of zip files ingested
413 """
414 bundles = []
415 for zip_file in zip_files:
416 bundles.append(self._make_zip_bundle(dataset_id, zip_file))
417 self._add_replicas(bundles)
418 self.register_to_dataset(bundles)
419 return len(bundles)
421 def register_dims(self, dataset_id: str, dim_files: list) -> int:
422 """Register a list of dimension files to a Rucio Dataset
424 Parameters
425 ----------
426 dataset_id : `str`
427 RUCIO dataset id
428 dim_files : `list` [`lsst.resource.ResourcePath`]
429 list of ResourcePath
431 Returns
432 -------
433 num : `int`
434 number of dimension files ingested
435 """
436 bundles = []
437 for dim_file in dim_files:
438 bundles.append(self._make_dim_bundle(dataset_id, dim_file))
439 self._add_replicas(bundles)
440 self.register_to_dataset(bundles)
441 return len(bundles)