Coverage for python/lsst/rucio/register/rucio_interface.py: 15%

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/>. 

21 

22import hashlib 

23import logging 

24import random 

25import time 

26import zlib 

27 

28import rucio.common.exception 

29from rucio.client.didclient import DIDClient 

30from rucio.client.replicaclient import ReplicaClient 

31 

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 

38 

39__all__ = ["RucioInterface"] 

40 

41logger = logging.getLogger(__name__) 

42 

43 

44class RucioInterface: 

45 """Add files as replicas in Rucio, along with metadata, 

46 and attach them to datasets. 

47 

48 Parameters 

49 ---------- 

50 butler : `lsst.daf.butler.Butler` 

51 Butler we're operating upon 

52 rucio_rse : `str` 

53 Name of the RSE that the files live in. 

54 scope : `str` 

55 Rucio scope to register the files in. 

56 rse_root : `str` 

57 Full path to root directory of RSE directory structure 

58 dtn_url : `str` 

59 Base URL of the data transfer node for the Rucio physical filename. 

60 rubin_butler_type: `str` 

61 the type registered in "rubin_butler" metadata for rucio 

62 """ 

63 

64 def __init__( 

65 self, 

66 butler: lsst.daf.butler.Butler, 

67 rucio_rse: str, 

68 scope: str, 

69 rse_root: str, 

70 dtn_url: str, 

71 rubin_butler_type: str, 

72 ): 

73 self.butler = butler 

74 self.rse = rucio_rse 

75 self.scope = scope 

76 self.rse_root = rse_root 

77 self.dtn_url = dtn_url 

78 self.pfn_base = f"{dtn_url}" 

79 self.replica_client = ReplicaClient() 

80 self.did_client = DIDClient() 

81 self.rubin_butler_type = rubin_butler_type 

82 

83 def _make_dataset_ref_bundle(self, dataset_id: str, dataset_ref: DatasetRef) -> ResourceBundle: 

84 """Make a ResourceBundle 

85 

86 Parameters 

87 ---------- 

88 dataset_id : `str` 

89 Rucio dataset name 

90 dataset_ref : `DatasetRef` 

91 Butler DatasetRef 

92 

93 Returns 

94 ------- 

95 rb : `ResourceBundle` 

96 ResourceBundle consolidating dataset id and DatasetRef 

97 """ 

98 logging.debug("%s", dataset_ref.to_json()) 

99 did = self._make_did(self.butler.getURI(dataset_ref), dataset_ref.to_json()) 

100 rb = ResourceBundle(dataset_id=dataset_id, did=did) 

101 return rb 

102 

103 def _make_zip_bundle(self, dataset_id: str, resource_path: ResourcePath) -> ResourceBundle: 

104 """Make a ResourceBundle 

105 

106 Parameters 

107 ---------- 

108 dataset_id : `str` 

109 Rucio dataset name 

110 resouce_path : `ResourcePath` 

111 ResourcePath to a file 

112 

113 Returns 

114 ------- 

115 rb: ResourceBundle 

116 ResourceBundle consolidating dataset id and ResourcePath 

117 """ 

118 did = self._make_did(resource_path) 

119 rb = ResourceBundle(dataset_id=dataset_id, did=did) 

120 return rb 

121 

122 def _make_dim_bundle(self, dataset_id: str, resource_path: ResourcePath) -> ResourceBundle: 

123 """Make a ResourceBundle 

124 

125 Parameters 

126 ---------- 

127 dataset_id : `str` 

128 Rucio dataset name 

129 resouce_path : `lsst.resource.ResourcePath` 

130 ResourcePath to a file 

131 

132 Returns 

133 ------- 

134 rb: `lsst.rucio.register.rucio_bundle.ResourceBundle` 

135 ResourceBundle consolidating dataset id and ResourcePath 

136 """ 

137 did = self._make_did(resource_path) 

138 rb = ResourceBundle(dataset_id=dataset_id, did=did) 

139 return rb 

140 

141 def _make_did(self, resource_path: ResourcePath, metadata: str = None) -> RucioDID: 

142 """Make a Rucio data identifier dictionary from a resource. 

143 

144 Parameters 

145 ---------- 

146 resource_path: ResourcePath 

147 ResourcePath object 

148 

149 metadata: `str` 

150 String containing Rubin dataset specific metadata 

151 

152 Returns 

153 ------- 

154 did : `dict` [`str`, `str`|`int`] 

155 Rucio data identifier including physical and logical names, 

156 byte length, adler32 and MD5 checksums, meta, and scope. 

157 """ 

158 with resource_path.open("rb") as f: 

159 contents = f.read() 

160 size = len(contents) 

161 md5 = hashlib.md5(contents).hexdigest() 

162 adler32 = f"{zlib.adler32(contents):08x}" 

163 path = resource_path.unquoted_path.removeprefix(self.rse_root) 

164 pfn = self.pfn_base + path 

165 logging.debug("pfn=%s", pfn) 

166 name = path.removeprefix("/" + self.scope + "/") 

167 logging.debug("name=%s", name) 

168 logging.debug("path=%s", path) 

169 

170 if metadata: 

171 meta = RubinMeta(rubin_butler=self.rubin_butler_type, rubin_sidecar=metadata) 

172 else: 

173 meta = RubinMeta(rubin_butler=self.rubin_butler_type, rubin_sidecar="") 

174 d = RucioDID( 

175 pfn=pfn, 

176 bytes=size, 

177 adler32=adler32, 

178 md5=md5, 

179 name=name, 

180 scope=self.scope, 

181 meta=meta, 

182 ) 

183 

184 return d 

185 

186 def _add_replicas(self, bundles: list[ResourceBundle]) -> None: 

187 """Call the Rucio method add_replica for a list of DIDs 

188 

189 Parameters 

190 ---------- 

191 bundles : `list` [`ResourceBundle`] 

192 A list of ResourceBundles 

193 """ 

194 dids = [bundle.get_did() for bundle in bundles] 

195 retries = 0 

196 max_retries = 5 

197 while True: 

198 try: 

199 self.replica_client.add_replicas(rse=self.rse, files=dids) 

200 break 

201 except rucio.common.exception.RucioException: 

202 retries += 1 

203 if retries < max_retries: 

204 seconds = random.randint(10, 20) 

205 logger.debug("failed to add_replicas; sleeping %d seconds", seconds) 

206 time.sleep(seconds) 

207 self.replica_client = ReplicaClient() # XXX not sure we need to do this. 

208 else: 

209 raise Exception(f"Tried {max_retries} times and couldn't add_replicas") 

210 

211 def _add_file_to_dataset_with_retries(self, dataset_id, did): 

212 retries = 0 

213 max_retries = 5 

214 while True: 

215 try: 

216 self.did_client.add_files_to_dataset( 

217 scope=self.scope, name=dataset_id, files=[did], rse=self.rse 

218 ) 

219 break 

220 except rucio.common.exception.FileAlreadyExists: 

221 if "pfn" in did: 

222 logger.debug("file %s already registered in dataset %s", did["pfn"], dataset_id) 

223 return # we can return, because it's already in the dataset 

224 except rucio.common.exception.RucioException: 

225 retries += 1 

226 if retries < max_retries: 

227 seconds = random.randint(10, 20) 

228 logger.debug("failed to register one did to %s; sleeping %d seconds", dataset_id, seconds) 

229 time.sleep(seconds) 

230 self.did_client = DIDClient() # XXX not sure we need to do this. 

231 else: 

232 # we tried max_retries times, and failed, so we'll bail out 

233 raise Exception(f"Couldn't add {did['pfn']} to dataset {dataset_id}") 

234 

235 def _add_files_to_dataset(self, dataset_id: str, dids: list[dict]) -> None: 

236 """Attach a list of files specified by Rucio DIDs to a Rucio dataset. 

237 

238 Ignores already-attached files for idempotency. 

239 

240 Parameters 

241 ---------- 

242 dataset_id : `str` 

243 Logical name of the Rucio dataset. 

244 dids : `list` [`dict` [`str`, `str`|`int`] ] 

245 List of Rucio data identifiers. 

246 """ 

247 retries = 0 

248 max_retries = 5 

249 while True: 

250 try: 

251 self.did_client.add_files_to_dataset( 

252 scope=self.scope, 

253 name=dataset_id, 

254 files=dids, 

255 rse=self.rse, 

256 ) 

257 return 

258 except rucio.common.exception.FileAlreadyExists: 

259 # At least one already is in the dataset. 

260 # This shouldn't happen, but if it does, 

261 # we have to retry each individually. 

262 for did in dids: 

263 self._add_file_to_dataset_with_retries( 

264 dataset_id=dataset_id, 

265 did=did, 

266 ) 

267 return 

268 except rucio.common.exception.DataIdentifierNotFound as e: 

269 raise e 

270 except rucio.common.exception.RucioException: 

271 retries += 1 

272 if retries < max_retries: 

273 seconds = random.randint(10, 20) 

274 logger.debug("failed to register dids to %s; sleeping %d", dataset_id, seconds) 

275 time.sleep(seconds) 

276 continue 

277 else: 

278 raise Exception(f"Couldn't add files to dataset {dataset_id}") 

279 

280 def _add_dataset_with_retries(self, dataset_id: str, statuses: dict) -> None: 

281 retries = 0 

282 max_retries = 5 

283 while True: 

284 try: 

285 self.did_client.add_dataset( 

286 scope=self.scope, 

287 name=dataset_id, 

288 statuses=statuses, 

289 rse=self.rse, 

290 ) 

291 return 

292 except rucio.common.exception.DataIdentifierAlreadyExists as e: 

293 # If someone else created it in the meantime 

294 raise e 

295 except rucio.common.exception.RucioException: 

296 retries += 1 

297 if retries < max_retries: 

298 seconds = random.randint(10, 20) 

299 logger.debug("couldn't register dids to %s; waiting %d", dataset_id, seconds) 

300 time.sleep(seconds) 

301 continue 

302 else: 

303 raise Exception(f"Tried {max_retries} times and couldn't add dataset {dataset_id}") 

304 

305 def register_to_dataset(self, bundles) -> None: 

306 """Register a list of files in Rucio. 

307 

308 Parameters 

309 ---------- 

310 bundles : `list` [`ResourceBundle`] 

311 List of resource bundles 

312 """ 

313 logger.debug("register to dataset") 

314 

315 datasets = dict() 

316 for bundle in bundles: 

317 dataset_id = bundle.dataset_id 

318 datasets.setdefault(dataset_id, []).append(bundle) 

319 

320 for dataset_id, bundles in datasets.items(): 

321 try: 

322 dids = [rb.get_did() for rb in bundles] 

323 names = [did["pfn"] for did in dids] 

324 logger.info("Registering %s in dataset %s, RSE %s", names, dataset_id, self.rse) 

325 self._add_files_to_dataset(dataset_id, dids) 

326 except rucio.common.exception.DataIdentifierNotFound: 

327 # No such dataset, so create it 

328 try: 

329 logger.info("Creating Rucio dataset %s", dataset_id) 

330 self._add_dataset_with_retries( 

331 dataset_id=dataset_id, 

332 statuses={"monotonic": True}, 

333 ) 

334 except rucio.common.exception.DataIdentifierAlreadyExists: 

335 # If someone else created it in the meantime 

336 pass 

337 # And then retry adding DIDs 

338 self._add_files_to_dataset(dataset_id, dids) 

339 

340 logger.debug("Done with Rucio for %s", bundles) 

341 

342 def register_as_replicas(self, dataset_id, dataset_refs) -> None: 

343 """Register a list of DatasetRefs to a Rucio dataset 

344 

345 Parameters 

346 ---------- 

347 dataset_id : `str` 

348 RUCIO dataset id 

349 dataset_refs : `list` [`DatasetRef`] 

350 list of Butler DatasetRefs 

351 """ 

352 bundles = [] 

353 for dataset_ref in dataset_refs: 

354 if type(dataset_ref) is list: 

355 for dsr in dataset_ref: 

356 bundles.append(self._make_dataset_ref_bundle(dataset_id, dsr)) 

357 else: 

358 bundles.append(self._make_dataset_ref_bundle(dataset_id, dataset_ref)) 

359 if len(bundles) == 0: 

360 return 0 

361 self._add_replicas(bundles) 

362 self.register_to_dataset(bundles) 

363 return len(bundles) 

364 

365 def register_zips(self, dataset_id: str, zip_files: list) -> int: 

366 """Register a list of zips to a Rucio Dataset 

367 

368 Parameters 

369 ---------- 

370 dataset_id : `str` 

371 RUCIO dataset id 

372 zip_files : `list` [`ResourcePath`] 

373 list of ResourcePath 

374 

375 Returns 

376 ------- 

377 num : `int` 

378 number of zip files ingested 

379 """ 

380 bundles = [] 

381 for zip_file in zip_files: 

382 bundles.append(self._make_zip_bundle(dataset_id, zip_file)) 

383 self._add_replicas(bundles) 

384 self.register_to_dataset(bundles) 

385 return len(bundles) 

386 

387 def register_dims(self, dataset_id: str, dim_files: list) -> int: 

388 """Register a list of dimension files to a Rucio Dataset 

389 

390 Parameters 

391 ---------- 

392 dataset_id : `str` 

393 RUCIO dataset id 

394 dim_files : `list` [`lsst.resource.ResourcePath`] 

395 list of ResourcePath 

396 

397 Returns 

398 ------- 

399 num : `int` 

400 number of dimension files ingested 

401 """ 

402 bundles = [] 

403 for dim_file in dim_files: 

404 bundles.append(self._make_dim_bundle(dataset_id, dim_file)) 

405 self._add_replicas(bundles) 

406 self.register_to_dataset(bundles) 

407 return len(bundles)