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