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