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

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 logging 

23import random 

24import time 

25import zlib 

26 

27import rucio.common.exception 

28from rucio.client.didclient import DIDClient 

29from rucio.client.replicaclient import ReplicaClient 

30 

31import lsst.daf.butler 

32from lsst.daf.butler import DatasetRef 

33from lsst.resources import ResourcePath 

34from lsst.rucio.register.resource_bundle import ResourceBundle 

35from lsst.rucio.register.rubin_meta import RubinMeta 

36from lsst.rucio.register.rucio_did import RucioDID 

37 

38__all__ = ["RucioInterface"] 

39 

40logger = logging.getLogger(__name__) 

41 

42 

43class RucioInterface: 

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

45 and attach them to datasets. 

46 

47 Parameters 

48 ---------- 

49 butler : `lsst.daf.butler.Butler` 

50 Butler we're operating upon 

51 rucio_rse : `str` 

52 Name of the RSE that the files live in. 

53 scope : `str` 

54 Rucio scope to register the files in. 

55 rse_root : `str` 

56 Full path to root directory of RSE directory structure 

57 dtn_url : `str` 

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

59 rubin_butler_type: `str` 

60 the type registered in "rubin_butler" metadata for rucio 

61 """ 

62 

63 def __init__( 

64 self, 

65 butler: lsst.daf.butler.Butler, 

66 rucio_rse: str, 

67 scope: str, 

68 rse_root: str, 

69 dtn_url: str, 

70 rubin_butler_type: str, 

71 ): 

72 self.butler = butler 

73 self.rse = rucio_rse 

74 self.scope = scope 

75 self.rse_root = rse_root 

76 self.dtn_url = dtn_url 

77 self.pfn_base = f"{dtn_url}" 

78 self.replica_client = ReplicaClient() 

79 self.did_client = DIDClient() 

80 self.rubin_butler_type = rubin_butler_type 

81 

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

83 """Make a ResourceBundle 

84 

85 Parameters 

86 ---------- 

87 dataset_id : `str` 

88 Rucio dataset name 

89 dataset_ref : `DatasetRef` 

90 Butler DatasetRef 

91 

92 Returns 

93 ------- 

94 rb : `ResourceBundle` 

95 ResourceBundle consolidating dataset id and DatasetRef 

96 """ 

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

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

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

100 return rb 

101 

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

103 """Make a ResourceBundle 

104 

105 Parameters 

106 ---------- 

107 dataset_id : `str` 

108 Rucio dataset name 

109 resouce_path : `ResourcePath` 

110 ResourcePath to a file 

111 

112 Returns 

113 ------- 

114 rb: ResourceBundle 

115 ResourceBundle consolidating dataset id and ResourcePath 

116 """ 

117 did = self._make_did(resource_path) 

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

119 return rb 

120 

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

122 """Make a ResourceBundle 

123 

124 Parameters 

125 ---------- 

126 dataset_id : `str` 

127 Rucio dataset name 

128 resouce_path : `lsst.resource.ResourcePath` 

129 ResourcePath to a file 

130 

131 Returns 

132 ------- 

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

134 ResourceBundle consolidating dataset id and ResourcePath 

135 """ 

136 did = self._make_did(resource_path) 

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

138 return rb 

139 

140 def compute_hashes(self, resource_path: ResourcePath) -> tuple[int, str]: 

141 """return the length and adler32 hash for a file. 

142 

143 Parameters 

144 ---------- 

145 path: `lsst.resources.ResourcePath` 

146 Path to the file. 

147 

148 Returns 

149 ------- 

150 hashes: `tuple` [ `int`, `str` ] 

151 Size in bytes and Adler32 hex hash. 

152 """ 

153 

154 info = resource_path.get_info() 

155 size = info.size 

156 checksums = info.checksums 

157 if "adler32" in checksums: 

158 adler32 = checksums["adler32"] 

159 logger.debug("found adler32 for %s", resource_path) 

160 return size, adler32 

161 return size, self._compute_adler32(resource_path) 

162 

163 def _compute_adler32(self, resource_path: ResourcePath) -> tuple[int, str]: 

164 logger.debug("computing adler32 for %s", resource_path) 

165 adler32 = zlib.adler32(b"") 

166 buffer_size = 10 * 1024 * 1024 

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

168 while buffer := f.read(buffer_size): 

169 adler32 = zlib.adler32(buffer, adler32) 

170 adler32_digest = f"{adler32:08x}" 

171 return adler32_digest 

172 

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

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

175 

176 Parameters 

177 ---------- 

178 resource_path: ResourcePath 

179 ResourcePath object 

180 

181 metadata: `str` 

182 String containing Rubin dataset specific metadata 

183 

184 Returns 

185 ------- 

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

187 Rucio data identifier including physical and logical names, 

188 byte length, adler32 checksum, meta, and scope. 

189 """ 

190 

191 size, adler32 = self.compute_hashes(resource_path) 

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

193 pfn = self.pfn_base + path 

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

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

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

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

198 

199 if metadata: 199 ↛ 202line 199 didn't jump to line 202 because the condition on line 199 was always true

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

201 else: 

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

203 d = RucioDID( 

204 pfn=pfn, 

205 bytes=size, 

206 adler32=adler32, 

207 name=name, 

208 scope=self.scope, 

209 meta=meta, 

210 ) 

211 

212 return d 

213 

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

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

216 

217 Parameters 

218 ---------- 

219 bundles : `list` [`ResourceBundle`] 

220 A list of ResourceBundles 

221 """ 

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

223 retries = 0 

224 max_retries = 5 

225 while True: 

226 try: 

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

228 break 

229 except rucio.common.exception.RucioException: 

230 retries += 1 

231 if retries < max_retries: 

232 seconds = random.randint(10, 20) 

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

234 time.sleep(seconds) 

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

236 else: 

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

238 

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

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

241 

242 Ignores already-attached files for idempotency. 

243 

244 Parameters 

245 ---------- 

246 dataset_id : `str` 

247 Logical name of the Rucio dataset. 

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

249 List of Rucio data identifiers. 

250 """ 

251 retries = 0 

252 max_retries = 5 

253 while True: 

254 try: 

255 self.did_client.add_files_to_datasets( 

256 attachments=[ 

257 { 

258 "scope": self.scope, 

259 "name": dataset_id, 

260 "dids": dids, 

261 "rse": self.rse, 

262 } 

263 ], 

264 ignore_duplicate=True, 

265 ) 

266 return 

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

268 raise e 

269 except rucio.common.exception.RucioException: 

270 retries += 1 

271 if retries < max_retries: 

272 seconds = random.randint(10, 20) 

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

274 time.sleep(seconds) 

275 continue 

276 else: 

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

278 

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

280 retries = 0 

281 max_retries = 5 

282 while True: 

283 try: 

284 self.did_client.add_dataset( 

285 scope=self.scope, 

286 name=dataset_id, 

287 statuses=statuses, 

288 rse=self.rse, 

289 ) 

290 return 

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

292 # If someone else created it in the meantime 

293 raise e 

294 except rucio.common.exception.RucioException: 

295 retries += 1 

296 if retries < max_retries: 

297 seconds = random.randint(10, 20) 

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

299 time.sleep(seconds) 

300 continue 

301 else: 

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

303 

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

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

306 

307 Parameters 

308 ---------- 

309 bundles : `list` [`ResourceBundle`] 

310 List of resource bundles 

311 """ 

312 logger.debug("register to dataset") 

313 

314 datasets = dict() 

315 for bundle in bundles: 

316 dataset_id = bundle.dataset_id 

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

318 

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

320 try: 

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

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

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

324 self._add_files_to_dataset(dataset_id, dids) 

325 except rucio.common.exception.DataIdentifierNotFound: 

326 # No such dataset, so create it 

327 try: 

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

329 self._add_dataset_with_retries( 

330 dataset_id=dataset_id, 

331 statuses={"monotonic": True}, 

332 ) 

333 except rucio.common.exception.DataIdentifierAlreadyExists: 

334 # If someone else created it in the meantime 

335 pass 

336 # And then retry adding DIDs 

337 self._add_files_to_dataset(dataset_id, dids) 

338 

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

340 

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

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

343 

344 Parameters 

345 ---------- 

346 dataset_id : `str` 

347 RUCIO dataset id 

348 dataset_refs : `list` [`DatasetRef`] 

349 list of Butler DatasetRefs 

350 """ 

351 bundles = [] 

352 for dataset_ref in dataset_refs: 

353 if type(dataset_ref) is list: 353 ↛ 354line 353 didn't jump to line 354 because the condition on line 353 was never true

354 for dsr in dataset_ref: 

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

356 else: 

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

358 if len(bundles) == 0: 358 ↛ 359line 358 didn't jump to line 359 because the condition on line 358 was never true

359 return 0 

360 self._add_replicas(bundles) 

361 self.register_to_dataset(bundles) 

362 return len(bundles) 

363 

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

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

366 

367 Parameters 

368 ---------- 

369 dataset_id : `str` 

370 RUCIO dataset id 

371 zip_files : `list` [`ResourcePath`] 

372 list of ResourcePath 

373 

374 Returns 

375 ------- 

376 num : `int` 

377 number of zip files ingested 

378 """ 

379 bundles = [] 

380 for zip_file in zip_files: 

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

382 self._add_replicas(bundles) 

383 self.register_to_dataset(bundles) 

384 return len(bundles) 

385 

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

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

388 

389 Parameters 

390 ---------- 

391 dataset_id : `str` 

392 RUCIO dataset id 

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

394 list of ResourcePath 

395 

396 Returns 

397 ------- 

398 num : `int` 

399 number of dimension files ingested 

400 """ 

401 bundles = [] 

402 for dim_file in dim_files: 

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

404 self._add_replicas(bundles) 

405 self.register_to_dataset(bundles) 

406 return len(bundles)