Coverage for python/lsst/rucio/register/script.py: 64%

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# (http://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 

22 

23import datetime 

24import itertools 

25import json 

26import logging 

27import os 

28import sys 

29import threading 

30import time 

31import urllib.request 

32from concurrent.futures import ProcessPoolExecutor, as_completed 

33from typing import Any 

34 

35import click 

36import yaml 

37 

38from lsst.daf.butler import Butler, CollectionType, _exceptions 

39from lsst.daf.butler.cli.opt import log_level_option, options_file_option, query_datasets_options 

40from lsst.daf.butler.script.queryDatasets import QueryDatasets 

41from lsst.resources import ResourcePath 

42from lsst.rucio.register.data_type import DataType 

43from lsst.rucio.register.rucio_interface import RucioInterface 

44from lsst.rucio.register.rucio_register_config import RucioRegisterConfig 

45 

46_thread_local = threading.local() 

47 

48logger = logging.getLogger(__name__) 

49_FORMAT = ( 

50 "%(levelname) -10s %(asctime)s.%(msecs)03dZ %(name) -30s %(funcName) -35s %(lineno) -5d: %(message)s" 

51) 

52 

53RUCIO_REGISTER_CONFIG = "RUCIO_REGISTER_CONFIG" 

54_MSG = "environment variable not set, and no configuration was specified on the command line" 

55 

56 

57def register_options(func): 

58 """Decorator for CLI options that are common across commands.""" 

59 func = click.option( 

60 "--backoff-max-tries", 

61 "--max-retries", 

62 "max_retries", 

63 required=False, 

64 type=int, 

65 default=RucioInterface.DEFAULT_RETRIES, 

66 show_default=True, 

67 help="maximum number of times to retry failed registration", 

68 )(func) 

69 func = click.option( 

70 "--backoff-max-value", 

71 required=False, 

72 type=int, 

73 default=RucioInterface.DEFAULT_BACKOFF_MAX_VALUE, 

74 show_default=True, 

75 help="maximum backoff value in seconds", 

76 )(func) 

77 func = click.option( 

78 "--backoff-factor", 

79 required=False, 

80 type=float, 

81 default=RucioInterface.DEFAULT_BACKOFF_FACTOR, 

82 show_default=True, 

83 help="multiplicative factor to backoff by", 

84 )(func) 

85 func = click.option( 

86 "--chunk-size", 

87 required=False, 

88 type=int, 

89 default=500, 

90 help="number of replica requests to make at once", 

91 )(func) 

92 func = click.option( 

93 "--rucio-register-config", required=False, type=str, help="registration configuration file" 

94 )(func) 

95 func = click.option( 

96 "--rucio-dataset", required=False, type=str, help="rucio dataset to register files to" 

97 )(func) 

98 return func 

99 

100 

101def chunks(refs, chunk_size): 

102 it = iter(refs) 

103 while True: 

104 chunk = list(itertools.islice(it, chunk_size)) 

105 if not chunk: 

106 return 

107 yield chunk 

108 

109 

110_rucio_interface_cache = {} 

111 

112 

113def _get_rucio_interface( 

114 repo, 

115 rucio_register_config, 

116 rubin_butler_type, 

117 clear_is_new=False, 

118 retries=RucioInterface.DEFAULT_RETRIES, 

119 kwargs=None, 

120): 

121 if kwargs: 

122 backoff_factor = kwargs.get("backoff_factor", RucioInterface.DEFAULT_BACKOFF_FACTOR) 

123 backoff_max_value = kwargs.get("backoff_max_value", RucioInterface.DEFAULT_BACKOFF_MAX_VALUE) 

124 backoff_max_tries = kwargs.get("backoff_max_tries", kwargs.get("max_retries", retries)) 

125 if backoff_max_tries is not None: 

126 retries = backoff_max_tries 

127 else: 

128 backoff_factor = RucioInterface.DEFAULT_BACKOFF_FACTOR 

129 backoff_max_value = RucioInterface.DEFAULT_BACKOFF_MAX_VALUE 

130 backoff_max_tries = retries 

131 

132 config_file = os.environ.get(RUCIO_REGISTER_CONFIG, rucio_register_config) 

133 cache_key = ( 

134 repo, 

135 config_file, 

136 rubin_butler_type, 

137 clear_is_new, 

138 retries, 

139 backoff_factor, 

140 backoff_max_value, 

141 ) 

142 if cache_key in _rucio_interface_cache: 

143 logger.debug("Reusing cached RucioInterface/Butler for key: %s", cache_key) 

144 return _rucio_interface_cache[cache_key] 

145 

146 if config_file is None: 

147 raise RuntimeError(f"{RUCIO_REGISTER_CONFIG} {_MSG}") 

148 

149 config = RucioRegisterConfig(config_file) 

150 

151 rucio_rse = config.rucio_rse 

152 scope = config.scope 

153 rse_root = config.rse_root 

154 dtn_url = config.dtn_url 

155 

156 butler = None 

157 if repo: 

158 butler = Butler(repo) 

159 

160 # create RucioInterface object used to register replicas into datasets 

161 ri = RucioInterface( 

162 butler=butler, 

163 rucio_rse=rucio_rse, 

164 scope=scope, 

165 rse_root=rse_root, 

166 dtn_url=dtn_url, 

167 rubin_butler_type=rubin_butler_type, 

168 clear_is_new=clear_is_new, 

169 retries=retries, 

170 ) 

171 ri.set_backoff(factor=backoff_factor, max_value=backoff_max_value, max_tries=retries) 

172 

173 _rucio_interface_cache[cache_key] = (ri, butler) 

174 return ri, butler 

175 

176 

177def worker_register_replicas( 

178 repo, rucio_config, butler_type, rucio_dataset, refs_list, retries, clear_is_new=False, kwargs=None 

179): 

180 ri, _ = _get_rucio_interface( 

181 repo, rucio_config, butler_type, clear_is_new=clear_is_new, retries=retries, kwargs=kwargs 

182 ) 

183 return ri.register_as_replicas(rucio_dataset, refs_list) 

184 

185 

186def worker_register_zips( 

187 repo, rucio_config, butler_type, rucio_dataset, zip_files, retries, clear_is_new=False, kwargs=None 

188): 

189 ri, _ = _get_rucio_interface( 

190 repo, rucio_config, butler_type, clear_is_new=clear_is_new, retries=retries, kwargs=kwargs 

191 ) 

192 return ri.register_zips(rucio_dataset, zip_files) 

193 

194 

195def worker_register_dims( 

196 repo, rucio_config, butler_type, rucio_dataset, dim_files, retries, clear_is_new=False, kwargs=None 

197): 

198 ri, _ = _get_rucio_interface( 

199 repo, rucio_config, butler_type, clear_is_new=clear_is_new, retries=retries, kwargs=kwargs 

200 ) 

201 return ri.register_dims(rucio_dataset, dim_files) 

202 

203 

204def _register( 

205 repo, 

206 rucio_register_config, 

207 rubin_butler_type, 

208 dataset_refs, 

209 chunk_size, 

210 rucio_dataset, 

211 max_workers=1, 

212 retries=RucioInterface.DEFAULT_RETRIES, 

213 stats=None, 

214 clear_is_new=False, 

215 kwargs=None, 

216): 

217 failed_uuids = [] 

218 registered_count = 0 

219 

220 total_refs = None 

221 try: 

222 total_refs = len(dataset_refs) 

223 except TypeError: 

224 pass 

225 

226 if total_refs is not None: 226 ↛ 232line 226 didn't jump to line 232 because the condition on line 226 was always true

227 logger.info( 

228 f"Starting registration of {total_refs} datasets to {rucio_dataset} " 

229 f"(chunk_size={chunk_size}, max_workers={max_workers})" 

230 ) 

231 else: 

232 logger.info( 

233 f"Starting registration of datasets to {rucio_dataset} " 

234 f"(chunk_size={chunk_size}, max_workers={max_workers})" 

235 ) 

236 

237 ref_chunks = list(chunks(dataset_refs, chunk_size)) 

238 total_chunks = len(ref_chunks) 

239 

240 if max_workers > 1 and total_chunks > 1: 240 ↛ 241line 240 didn't jump to line 241 because the condition on line 240 was never true

241 with ProcessPoolExecutor(max_workers=max_workers) as executor: 

242 future_to_chunk = { 

243 executor.submit( 

244 worker_register_replicas, 

245 repo, 

246 rucio_register_config, 

247 rubin_butler_type, 

248 rucio_dataset, 

249 refs, 

250 retries, 

251 clear_is_new, 

252 kwargs, 

253 ): (i, refs) 

254 for i, refs in enumerate(ref_chunks, 1) 

255 } 

256 

257 for future in as_completed(future_to_chunk): 

258 chunk_idx, refs = future_to_chunk[future] 

259 try: 

260 count = future.result() 

261 registered_count += count 

262 logger.info(f"Chunk {chunk_idx}/{total_chunks} completed: {count} datasets registered") 

263 except Exception as e: 

264 logger.error(f"Chunk {chunk_idx}/{total_chunks} failed permanently: {e}") 

265 chunk_uuids = [str(r.id) for r in refs] 

266 failed_uuids.extend(chunk_uuids) 

267 else: 

268 for chunk_idx, refs in enumerate(ref_chunks, 1): 

269 refs_list = list(refs) 

270 try: 

271 count = worker_register_replicas( 

272 repo, 

273 rucio_register_config, 

274 rubin_butler_type, 

275 rucio_dataset, 

276 refs_list, 

277 retries, 

278 clear_is_new, 

279 kwargs, 

280 ) 

281 registered_count += count 

282 logger.info(f"Chunk {chunk_idx}/{total_chunks} completed: {count} datasets registered") 

283 except Exception as e: 

284 logger.error(f"Chunk {chunk_idx}/{total_chunks} failed permanently: {e}") 

285 chunk_uuids = [str(r.id) for r in refs_list] 

286 failed_uuids.extend(chunk_uuids) 

287 

288 if stats is not None: 288 ↛ 293line 288 didn't jump to line 293 because the condition on line 288 was always true

289 stats["registered"] = registered_count 

290 stats["failed"] = len(failed_uuids) 

291 stats["failures"] = failed_uuids 

292 

293 if failed_uuids: 293 ↛ 294line 293 didn't jump to line 294 because the condition on line 293 was never true

294 logger.error( 

295 f"Registration summary for {rucio_dataset} - " 

296 f"registered: {registered_count}, failed: {len(failed_uuids)}" 

297 ) 

298 logger.error(f"Failed dataset UUIDs: {failed_uuids}") 

299 else: 

300 logger.info(f"Registration summary for {rucio_dataset} - registered: {registered_count}, failed: 0") 

301 

302 

303def _set_log_level(log_level: Any) -> None: 

304 if isinstance(log_level, dict): 

305 log_level = log_level.get(None) or (next(iter(log_level.values()), None) if log_level else None) 

306 if isinstance(log_level, str): 

307 handler = logging.StreamHandler(sys.stdout) 

308 handler.setFormatter(logging.Formatter(_FORMAT)) 

309 log = logging.getLogger() 

310 log.handlers.clear() 

311 log.addHandler(handler) 

312 log.setLevel(log_level.upper()) 

313 

314 

315def _limit_dataset_types(dataset_types: list[str], max_dataset_types: int | None) -> list[str]: 

316 """Limit the list of dataset types to max_dataset_types if provided.""" 

317 if max_dataset_types is None: 

318 return dataset_types 

319 return list(dataset_types[:max_dataset_types]) 

320 

321 

322@click.group(context_settings={"help_option_names": ["-h", "--help"]}) 

323def main(): 

324 pass 

325 

326 

327@main.command() 

328@click.option("--repo", required=True, type=str, help="butler repository") 

329@register_options 

330@click.option( 

331 "--max-workers", 

332 required=False, 

333 type=int, 

334 default=5, 

335 help="number of concurrent workers for registration", 

336) 

337@click.option( 

338 "--clear-is-new", 

339 is_flag=True, 

340 default=False, 

341 help="Set is_new metadata to None on Rucio dataset DIDs post-registration", 

342) 

343@log_level_option() 

344@options_file_option() 

345@query_datasets_options(repo=False, showUri=True, useArguments=False) 

346def data_products(**kwargs: Any) -> None: 

347 log_level = kwargs.get("log_level", None) 

348 _set_log_level(log_level) 

349 

350 rucio_dataset = kwargs.get("rucio_dataset", None) 

351 if not rucio_dataset: 

352 raise click.UsageError("--rucio-dataset must be provided for data_products.") 

353 

354 rucio_register_config = kwargs.get("rucio_register_config", None) 

355 chunk_size = kwargs.get("chunk_size", 500) 

356 max_workers = kwargs.get("max_workers", 5) 

357 max_retries = kwargs.get("max_retries", kwargs.get("backoff_max_tries", RucioInterface.DEFAULT_RETRIES)) 

358 clear_is_new = kwargs.get("clear_is_new", False) 

359 

360 repo = kwargs.get("repo", None) 

361 collections = kwargs.get("collections", None) 

362 where = kwargs.get("where", None) 

363 find_first = kwargs.get("find_first", None) 

364 limit = kwargs.get("limit", None) 

365 order_by = kwargs.get("order_by", None) 

366 dataset_type = kwargs.get("dataset_type", None) 

367 

368 ri, butler = _get_rucio_interface( 

369 repo, 

370 rucio_register_config, 

371 DataType.DATA_PRODUCT, 

372 clear_is_new=clear_is_new, 

373 retries=max_retries, 

374 kwargs=kwargs, 

375 ) 

376 

377 query = QueryDatasets( 

378 butler=butler, 

379 glob=dataset_type, 

380 collections=collections, 

381 where=where, 

382 find_first=find_first, 

383 limit=limit, 

384 order_by=order_by, 

385 show_uri=False, 

386 with_dimension_records=True, 

387 ) 

388 

389 dataset_refs = itertools.chain(*query.getDatasets()) 

390 

391 stats = {} 

392 _register( 

393 repo, 

394 rucio_register_config, 

395 DataType.DATA_PRODUCT, 

396 dataset_refs, 

397 chunk_size, 

398 rucio_dataset, 

399 max_workers=max_workers, 

400 retries=max_retries, 

401 stats=stats, 

402 clear_is_new=clear_is_new, 

403 kwargs=kwargs, 

404 ) 

405 

406 

407@main.command() 

408@click.option("--repo", required=True, type=str, help="butler repository") 

409@register_options 

410@click.option( 

411 "--max-workers", 

412 required=False, 

413 type=int, 

414 default=5, 

415 help="number of concurrent workers for registration", 

416) 

417@click.option( 

418 "--uuidlist", 

419 required=False, 

420 type=str, 

421 help=( 

422 "filename of a list of butler dataset UUIDs (e.g. generated by " 

423 "'auto-register --dry-run') to be registered to the rucio dataset." 

424 ), 

425) 

426@log_level_option() 

427@options_file_option() 

428def dataset_list(**kwargs: Any) -> None: 

429 log_level = kwargs.get("log_level", None) 

430 _set_log_level(log_level) 

431 

432 rucio_register_config = kwargs.get("rucio_register_config", None) 

433 chunk_size = kwargs.get("chunk_size", 500) 

434 max_workers = kwargs.get("max_workers", 5) 

435 uuidlist = kwargs.get("uuidlist", None) 

436 if not uuidlist: 436 ↛ 437line 436 didn't jump to line 437 because the condition on line 436 was never true

437 raise click.UsageError("uuidlist is None. Please provide a valid file containing UUIDs.") 

438 

439 logger.info(f"Reading UUIDs from {uuidlist}") 

440 max_retries = kwargs.get("max_retries", kwargs.get("backoff_max_tries", RucioInterface.DEFAULT_RETRIES)) 

441 

442 repo = kwargs.get("repo", None) 

443 

444 # Infer dataset from file header if available 

445 uuids = [] 

446 inferred_dataset = None 

447 with open(uuidlist) as f: 

448 content = f.read().strip() 

449 if content.startswith("[") or content.startswith("{"): 449 ↛ 459line 449 didn't jump to line 459 because the condition on line 449 was always true

450 try: 

451 data = json.loads(content) 

452 if isinstance(data, list): 452 ↛ 454line 452 didn't jump to line 454 because the condition on line 452 was always true

453 uuids = [str(item.get("id", item) if isinstance(item, dict) else item) for item in data] 

454 elif isinstance(data, dict): 

455 uuids = [str(k) for k in data.keys()] 

456 except json.JSONDecodeError: 

457 pass 

458 

459 if not uuids: 459 ↛ 460line 459 didn't jump to line 460 because the condition on line 459 was never true

460 for line in content.splitlines(): 

461 stripped = line.strip() 

462 if stripped.startswith("# Dataset/"): 

463 parts = stripped.split() 

464 if len(parts) >= 2: 

465 inferred_dataset = parts[1] 

466 elif not line.lstrip().startswith("#") and stripped: 

467 clean_uuid = stripped.strip('"').strip("'").rstrip(",").strip() 

468 if clean_uuid and clean_uuid not in ("[", "]"): 

469 uuids.append(clean_uuid) 

470 

471 rucio_dataset = kwargs.get("rucio_dataset") or inferred_dataset 

472 if not rucio_dataset: 472 ↛ 473line 472 didn't jump to line 473 because the condition on line 472 was never true

473 raise click.UsageError( 

474 "--rucio-dataset must be provided on CLI or via '# Dataset/...' header in uuidlist." 

475 ) 

476 

477 ri, butler = _get_rucio_interface( 

478 repo, rucio_register_config, DataType.DATA_PRODUCT, retries=max_retries, kwargs=kwargs 

479 ) 

480 if butler is None: 480 ↛ 481line 480 didn't jump to line 481 because the condition on line 480 was never true

481 raise click.UsageError("Butler instance is None. Please check the repository path.") 

482 

483 dataset_refs = butler.get_many_datasets(uuids) 

484 

485 stats = {} 

486 _register( 

487 repo, 

488 rucio_register_config, 

489 DataType.DATA_PRODUCT, 

490 dataset_refs, 

491 chunk_size, 

492 rucio_dataset, 

493 max_workers=max_workers, 

494 retries=max_retries, 

495 stats=stats, 

496 kwargs=kwargs, 

497 ) 

498 logger.info(f"Successfully registered {len(dataset_refs)} datasets from {uuidlist} to {rucio_dataset}") 

499 

500 

501@main.command() 

502@click.option("--repo", required=True, type=str, help="butler repository") 

503@register_options 

504@click.option( 

505 "--max-workers", 

506 required=False, 

507 type=int, 

508 default=5, 

509 help="number of concurrent workers for registration", 

510) 

511@log_level_option() 

512@options_file_option() 

513@query_datasets_options(repo=False, showUri=True) 

514def raws(**kwargs: Any) -> None: 

515 log_level = kwargs.pop("log_level") 

516 _set_log_level(log_level) 

517 

518 rucio_dataset = kwargs.get("rucio_dataset") 

519 if not rucio_dataset: 

520 raise click.UsageError("--rucio-dataset must be provided for raws.") 

521 

522 rucio_register_config = kwargs.get("rucio_register_config") 

523 chunk_size = kwargs.get("chunk_size", 500) 

524 max_workers = kwargs.get("max_workers", 5) 

525 max_retries = kwargs.get("max_retries", kwargs.get("backoff_max_tries", RucioInterface.DEFAULT_RETRIES)) 

526 

527 repo = kwargs["repo"] 

528 

529 ri, butler = _get_rucio_interface( 

530 repo, rucio_register_config, DataType.RAW_FILE, retries=max_retries, kwargs=kwargs 

531 ) 

532 

533 query_kwargs = { 

534 k: v 

535 for k, v in kwargs.items() 

536 if k 

537 not in ( 

538 "backoff_factor", 

539 "backoff_max_value", 

540 "backoff_max_tries", 

541 "max_retries", 

542 "backoff-factor", 

543 "backoff-max-value", 

544 "backoff-max-tries", 

545 "max_workers", 

546 "rucio_dataset", 

547 "rucio_register_config", 

548 "chunk_size", 

549 ) 

550 } 

551 

552 # chain is needed to flatten the list of lists returned by getDatasets() 

553 dataset_refs = itertools.chain.from_iterable(QueryDatasets(**query_kwargs).getDatasets()) 

554 

555 stats = {} 

556 _register( 

557 repo, 

558 rucio_register_config, 

559 DataType.RAW_FILE, 

560 dataset_refs, 

561 chunk_size, 

562 rucio_dataset, 

563 max_workers=max_workers, 

564 retries=max_retries, 

565 stats=stats, 

566 kwargs=kwargs, 

567 ) 

568 

569 

570@main.command() 

571@register_options 

572@click.option("--zip-file", required=True, help="zip file to register") 

573@log_level_option() 

574def zips(**kwargs: Any) -> None: 

575 log_level = kwargs.get("log_level", None) 

576 _set_log_level(log_level) 

577 

578 rucio_dataset = kwargs.get("rucio_dataset") 

579 if not rucio_dataset: 

580 raise click.UsageError("--rucio-dataset must be provided for zips.") 

581 

582 zip_file = kwargs.get("zip_file") 

583 rucio_register_config = kwargs.get("rucio_register_config") 

584 max_retries = kwargs.get("max_retries", kwargs.get("backoff_max_tries", RucioInterface.DEFAULT_RETRIES)) 

585 

586 ri, butler = _get_rucio_interface( 

587 None, rucio_register_config, DataType.ZIP_FILE, retries=max_retries, kwargs=kwargs 

588 ) 

589 

590 try: 

591 count = worker_register_zips( 

592 repo=None, 

593 rucio_config=rucio_register_config, 

594 butler_type=DataType.ZIP_FILE, 

595 rucio_dataset=rucio_dataset, 

596 zip_files=[ResourcePath(zip_file)], 

597 retries=max_retries, 

598 kwargs=kwargs, 

599 ) 

600 logger.info(f"Batch zip registration summary for {rucio_dataset} - registered: {count}, failed: 0") 

601 logger.info(f"Successfully registered zip file {zip_file} to {rucio_dataset}") 

602 except Exception as e: 

603 logger.error(f"Failed to register zip file {zip_file}: {e}", exc_info=True) 

604 

605 

606@main.command() 

607@register_options 

608@click.option("--dimension-file", required=True, help="dimension file to register") 

609@log_level_option() 

610def dimensions(**kwargs: Any) -> None: 

611 log_level = kwargs.get("log_level", None) 

612 _set_log_level(log_level) 

613 

614 rucio_dataset = kwargs.get("rucio_dataset") 

615 if not rucio_dataset: 

616 raise click.UsageError("--rucio-dataset must be provided for dimensions.") 

617 

618 dim_file = kwargs.get("dimension_file") 

619 rucio_register_config = kwargs.get("rucio_register_config") 

620 max_retries = kwargs.get("max_retries", kwargs.get("backoff_max_tries", RucioInterface.DEFAULT_RETRIES)) 

621 

622 ri, butler = _get_rucio_interface( 

623 None, rucio_register_config, DataType.DIM_FILE, retries=max_retries, kwargs=kwargs 

624 ) 

625 

626 try: 

627 count = worker_register_dims( 

628 repo=None, 

629 rucio_config=rucio_register_config, 

630 butler_type=DataType.DIM_FILE, 

631 rucio_dataset=rucio_dataset, 

632 dim_files=[ResourcePath(dim_file)], 

633 retries=max_retries, 

634 kwargs=kwargs, 

635 ) 

636 logger.info( 

637 f"Batch dimension registration summary for {rucio_dataset} - registered: {count}, failed: 0" 

638 ) 

639 logger.info(f"Successfully registered dimension file {dim_file} to {rucio_dataset}") 

640 except Exception as e: 

641 logger.error(f"Failed to register dimension file {dim_file}: {e}", exc_info=True) 

642 

643 

644def _process_refs_chunk( 

645 refs_chunk, 

646 rucio_dataset_name, 

647 dry_run, 

648 out_dir, 

649 repo, 

650 rucio_register_config, 

651 chunk_size, 

652 max_retries, 

653 clear_is_new=False, 

654 kwargs=None, 

655): 

656 """Register or dry-run write a chunk of refs.""" 

657 local_stats = {"registered": 0, "failed": 0, "failures": []} 

658 if dry_run: 

659 lines = [f"# {rucio_dataset_name} {len(refs_chunk)}\n"] 

660 for ref in refs_chunk: 

661 lines.append(f"{str(ref.id)}\n") 

662 

663 # Sanitize repo path for use in output directory structure 

664 sanitized_repo = repo.strip("/").replace("/", "_") 

665 # Sanitize rucio_dataset_name for use as a filename 

666 sanitized_filename = rucio_dataset_name.replace("/", "-") + ".txt" 

667 

668 uuiddir = os.path.join(out_dir, sanitized_repo, "auto_register") 

669 os.makedirs(uuiddir, exist_ok=True) 

670 uuidfile = os.path.join(uuiddir, sanitized_filename) 

671 

672 with open(uuidfile, "w") as f: 

673 f.write("".join(lines)) 

674 logger.info(f"Dry-run: Wrote {len(refs_chunk)} UUIDs to {uuidfile}") 

675 else: 

676 logger.info(f"Registering {len(refs_chunk)} items to dynamic Rucio dataset {rucio_dataset_name}") 

677 # When called from auto_register, we rely on the outer ProcessPool 

678 # for parallelism (`batch_workers`). We set the inner worker count 

679 # to 1 to avoid nested process pools, which can cause deadlocks. 

680 # This is now the default for _register. 

681 _register( 

682 repo, 

683 rucio_register_config, 

684 DataType.DATA_PRODUCT, 

685 refs_chunk, 

686 chunk_size, 

687 rucio_dataset_name, 

688 retries=max_retries, 

689 stats=local_stats, 

690 clear_is_new=clear_is_new, 

691 kwargs=kwargs, 

692 ) 

693 return local_stats 

694 

695 

696def _process_auto_register_batch(*args: Any, **kwargs: Any) -> dict[str, int]: 

697 """Worker function to process a batch of dataset refs for 

698 auto-registration. 

699 """ 

700 local_stats = {"registered": 0, "failed": 0} 

701 

702 # If pre-queried refs are passed directly: 

703 refs = kwargs.get("refs") 

704 ds_name = kwargs.get("ds_name") 

705 clear_is_new = kwargs.get("clear_is_new", False) 

706 if refs is not None: 706 ↛ 731line 706 didn't jump to line 731 because the condition on line 706 was always true

707 repo = kwargs.get("repo") 

708 rucio_register_config = kwargs.get("rucio_register_config") 

709 dry_run = kwargs.get("dry_run") 

710 out_dir = kwargs.get("out_dir") 

711 chunk_size = kwargs.get("chunk_size") 

712 max_retries = kwargs.get("max_retries") 

713 logger.info( 

714 "Worker processing pre-queried sub-batch of size %d for Rucio dataset %s", 

715 len(refs), 

716 ds_name, 

717 ) 

718 return _process_refs_chunk( 

719 refs, 

720 ds_name, 

721 dry_run, 

722 out_dir, 

723 repo, 

724 rucio_register_config, 

725 chunk_size, 

726 max_retries, 

727 clear_is_new=clear_is_new, 

728 kwargs=kwargs, 

729 ) 

730 

731 return local_stats 

732 

733 

734def _load_curated_dataset_types(transfer_list: str) -> set[str]: 

735 content = "" 

736 if transfer_list.startswith(("http://", "https://")): 

737 try: 

738 logger.info("Fetching curated transfer list from URL: %s", transfer_list) 

739 with urllib.request.urlopen(transfer_list, timeout=15.0) as response: 

740 content = response.read().decode("utf-8") 

741 except Exception as e: 

742 logger.error("Failed to fetch curated transfer list from URL: %s", e) 

743 raise 

744 else: 

745 try: 

746 logger.info("Loading curated transfer list from local file: %s", transfer_list) 

747 with open(transfer_list) as f: 

748 content = f.read() 

749 except Exception as e: 

750 logger.error("Failed to read curated transfer list from local file: %s", e) 

751 raise 

752 

753 try: 

754 data = yaml.safe_load(content) 

755 except Exception as e: 

756 logger.error("Failed to parse curated transfer list YAML: %s", e) 

757 raise 

758 

759 curated_types = set() 

760 

761 def _traverse(node): 

762 if isinstance(node, list): 

763 for item in node: 

764 if isinstance(item, str): 764 ↛ 763line 764 didn't jump to line 763 because the condition on line 764 was always true

765 curated_types.add(item) 

766 elif isinstance(node, dict): 766 ↛ exitline 766 didn't return from function '_traverse' because the condition on line 766 was always true

767 for val in node.values(): 

768 _traverse(val) 

769 

770 _traverse(data) 

771 logger.info("Loaded %d curated dataset types from transfer list", len(curated_types)) 

772 return curated_types 

773 

774 

775def _matches_dataset_type(ref: Any, target_type: str) -> bool: 

776 """Check if a DatasetRef matches a target dataset_type name. 

777 

778 Supports MagicMock. 

779 """ 

780 ds_type = getattr(ref, "dataset_type", getattr(ref, "datasetType", None)) 

781 if ds_type is None: 781 ↛ 782line 781 didn't jump to line 782 because the condition on line 781 was never true

782 return True 

783 type_name = getattr(ds_type, "name", None) if not isinstance(ds_type, str) else ds_type 

784 if type_name is None: 784 ↛ 785line 784 didn't jump to line 785 because the condition on line 784 was never true

785 return True 

786 if type(type_name).__name__ == "MagicMock": 

787 return True 

788 return str(type_name) == str(target_type) 

789 

790 

791def _query_butler_task( 

792 repo: str, task: dict[str, Any], butler: Any = None 

793) -> tuple[dict[str, Any], list[Any]]: 

794 collection = task["collection"] 

795 dataset_type = task["dataset_type"] 

796 where = task["where"] 

797 limit = task.get("limit", None) 

798 curated_types = task.get("curated_types") 

799 

800 if curated_types is not None: 

801 limit = None 

802 

803 try: