Coverage for python/lsst/rucio/register/rucio_interface.py: 14%
150 statements
« prev ^ index » next coverage.py v7.10.2, created at 2025-08-10 11:55 -0700
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 hashlib
23import logging
24import random
25import time
26import zlib
28import rucio.common.exception
29from rucio.client.didclient import DIDClient
30from rucio.client.replicaclient import ReplicaClient
32import lsst.daf.butler
33from lsst.rucio.register.resource_bundle import ResourceBundle
34from lsst.rucio.register.rubin_meta import RubinMeta
35from lsst.rucio.register.rucio_did import RucioDID
37__all__ = ["RucioInterface"]
39logger = logging.getLogger(__name__)
42class RucioInterface:
43 """Add files as replicas in Rucio, along with metadata,
44 and attach them to datasets.
46 Parameters
47 ----------
48 butler: `lsst.daf.butler.Butler`
49 Butler we're operating upon
50 rucio_rse: `str`
51 Name of the RSE that the files live in.
52 scope: `str`
53 Rucio scope to register the files in.
54 rse_root: `str`
55 Full path to root directory of RSE directory structure
56 dtn_url: `str`
57 Base URL of the data transfer node for the Rucio physical filename.
58 rubin_butler_type: `str`
59 the type registered in "rubin_butler" metadata for rucio
60 """
62 def __init__(
63 self,
64 butler: lsst.daf.butler.Butler,
65 rucio_rse: str,
66 scope: str,
67 rse_root: str,
68 dtn_url: str,
69 rubin_butler_type: str,
70 ):
71 self.butler = butler
72 self.rse = rucio_rse
73 self.scope = scope
74 self.rse_root = rse_root
75 self.dtn_url = dtn_url
76 self.pfn_base = f"{dtn_url}"
77 self.replica_client = ReplicaClient()
78 self.did_client = DIDClient()
79 self.rubin_butler_type = rubin_butler_type
81 def _make_bundle(self, dataset_id, dataset_ref) -> ResourceBundle:
82 """Make a ResourceBundle
84 Parameters
85 ----------
86 dataset_id: `str`
87 Rucio dataset name
88 dataset_ref: `DatasetRef`
89 Butler DatasetRef
90 """
91 logging.debug(f"{dataset_ref.to_json()}")
92 did = self._make_did(self.butler.getURI(dataset_ref), dataset_ref.to_json())
93 rb = ResourceBundle(dataset_id=dataset_id, did=did)
94 return rb
96 def _make_did(self, resource_path, metadata: str) -> RucioDID:
97 """Make a Rucio data identifier dictionary from a resource.
99 Parameters
100 ----------
101 resource_path: ResourcePath
102 ResourcePath object
104 metadata: `str`
105 String containing Rubin dataset specific metadata
107 Returns
108 -------
109 did: `dict [ str, str|int ]`
110 Rucio data identifier including physical and logical names,
111 byte length, adler32 and MD5 checksums, meta, and scope.
112 """
113 with resource_path.open("rb") as f:
114 contents = f.read()
115 size = len(contents)
116 md5 = hashlib.md5(contents).hexdigest()
117 adler32 = f"{zlib.adler32(contents):08x}"
118 path = resource_path.unquoted_path.removeprefix(self.rse_root)
119 pfn = self.pfn_base + path
120 logging.debug(f"{pfn=}")
121 name = path.removeprefix("/" + self.scope + "/")
122 logging.debug(f"{name=}")
123 logging.debug(f"{path=}")
125 meta = RubinMeta(rubin_butler=self.rubin_butler_type, rubin_sidecar=metadata)
126 d = RucioDID(
127 pfn=pfn,
128 bytes=size,
129 adler32=adler32,
130 md5=md5,
131 name=name,
132 scope=self.scope,
133 meta=meta,
134 )
136 return d
138 def _add_replicas(self, bundles: list[ResourceBundle]) -> None:
139 """Call the Rucio method add_replica for a list of DIDs
141 Parameters
142 ----------
143 bundles: `list[ResourceBundle]`
144 A list of ResourceBundles
145 """
146 dids = [bundle.get_did() for bundle in bundles]
147 retries = 0
148 max_retries = 5
149 while True:
150 try:
151 self.replica_client.add_replicas(rse=self.rse, files=dids)
152 break
153 except rucio.common.exception.RucioException:
154 retries += 1
155 if retries < max_retries:
156 seconds = random.randint(10, 20)
157 logger.debug(f"failed to add_replicas; sleeping {seconds} seconds")
158 time.sleep(seconds)
159 self.replica_client = ReplicaClient() # XXX not sure we need to do this.
160 else:
161 raise Exception(f"Tried {max_retries} times and couldn't add_replicas")
163 def _add_file_to_dataset_with_retries(self, dataset_id, did):
164 retries = 0
165 max_retries = 5
166 while True:
167 try:
168 self.did_client.add_files_to_dataset(
169 scope=self.scope, name=dataset_id, files=[did], rse=self.rse
170 )
171 break
172 except rucio.common.exception.FileAlreadyExists:
173 logger.debug(f"file {did} already registered in dataset {dataset_id}")
174 return # we can return, because it's already in the dataset
175 except rucio.common.exception.RucioException:
176 retries += 1
177 if retries < max_retries:
178 seconds = random.randint(10, 20)
179 logger.debug(f"failed to register one did to {dataset_id}; sleeping {seconds}")
180 time.sleep(seconds)
181 self.did_client = DIDClient() # XXX not sure we need to do this.
182 else:
183 # we tried max_retries times, and failed, so we'll bail out
184 raise Exception(f"Couldn't add {did['pfn']} to dataset {dataset_id}")
186 def _add_files_to_dataset(self, dataset_id: str, dids: list[dict]) -> None:
187 """Attach a list of files specified by Rucio DIDs to a Rucio dataset.
189 Ignores already-attached files for idempotency.
191 Parameters
192 ----------
193 dataset_id: `str`
194 Logical name of the Rucio dataset.
195 dids: `list [ dict [ str, str|int ] ]`
196 List of Rucio data identifiers.
197 """
198 retries = 0
199 max_retries = 5
200 while True:
201 try:
202 self.did_client.add_files_to_dataset(
203 scope=self.scope,
204 name=dataset_id,
205 files=dids,
206 rse=self.rse,
207 )
208 return
209 except rucio.common.exception.FileAlreadyExists:
210 # At least one already is in the dataset.
211 # This shouldn't happen, but if it does,
212 # we have to retry each individually.
213 for did in dids:
214 self._add_file_to_dataset_with_retries(
215 dataset_id=dataset_id,
216 did=did,
217 )
218 return
219 except rucio.common.exception.DataIdentifierNotFound as e:
220 raise e
221 except rucio.common.exception.RucioException:
222 retries += 1
223 if retries < max_retries:
224 seconds = random.randint(10, 20)
225 logger.debug(f"failed to register dids to {dataset_id}; sleeping {seconds}")
226 time.sleep(seconds)
227 continue
228 else:
229 raise Exception(f"Couldn't add files to dataset {dataset_id}")
231 def _add_dataset_with_retries(self, dataset_id: str, statuses: dict) -> None:
232 retries = 0
233 max_retries = 5
234 while True:
235 try:
236 self.did_client.add_dataset(
237 scope=self.scope,
238 name=dataset_id,
239 statuses=statuses,
240 rse=self.rse,
241 )
242 return
243 except rucio.common.exception.DataIdentifierAlreadyExists as e:
244 # If someone else created it in the meantime
245 raise e
246 except rucio.common.exception.RucioException:
247 retries += 1
248 if retries < max_retries:
249 seconds = random.randint(10, 20)
250 logger.debug(f"couldn't register dids to {dataset_id}; waiting {seconds}")
251 time.sleep(seconds)
252 continue
253 else:
254 raise Exception(f"Tried {max_retries} times and couldn't add dataset {dataset_id}")
256 def register_to_dataset(self, bundles) -> None:
257 """Register a list of files in Rucio.
259 Parameters
260 ----------
261 Bundles: `list [ ResourceBundle ]`
262 List of resource bundles
263 """
264 logger.debug("register to dataset")
266 datasets = dict()
267 for bundle in bundles:
268 dataset_id = bundle.dataset_id
269 datasets.setdefault(dataset_id, []).append(bundle)
271 for dataset_id, bundles in datasets.items():
272 try:
273 dids = [rb.get_did() for rb in bundles]
274 logger.info("Registering %s in dataset %s, RSE %s", dids, dataset_id, self.rse)
275 self._add_files_to_dataset(dataset_id, dids)
276 except rucio.common.exception.DataIdentifierNotFound:
277 # No such dataset, so create it
278 try:
279 logger.info("Creating Rucio dataset %s", dataset_id)
280 self._add_dataset_with_retries(
281 dataset_id=dataset_id,
282 statuses={"monotonic": True},
283 )
284 except rucio.common.exception.DataIdentifierAlreadyExists:
285 # If someone else created it in the meantime
286 pass
287 # And then retry adding DIDs
288 self._add_files_to_dataset(dataset_id, dids)
290 logger.debug("Done with Rucio for %s", bundles)
292 def register_as_replicas(self, dataset_id, dataset_refs) -> None:
293 """Register a list of DatasetRefs to a Rucio dataset
295 Parameters
296 ----------
297 dataset_id: `str`
298 RUCIO dataset id
299 dataset_refs: `list[DatasetRef]`
300 list of Butler DatasetRefs
301 """
302 bundles = []
303 for dataset_ref in dataset_refs:
304 if type(dataset_ref) is list:
305 for dsr in dataset_ref:
306 logging.debug(f"{self.butler.getURI(dsr)=}")
307 logging.debug(f"{dsr.to_json()=}")
308 bundles.append(self._make_bundle(dataset_id, dsr))
309 else:
310 bundles.append(self._make_bundle(dataset_id, dataset_ref))
311 if len(bundles) == 0:
312 return 0
313 self._add_replicas(bundles)
314 self.register_to_dataset(bundles)
315 return len(bundles)