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

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 multiprocessing 

28import os 

29import sys 

30import threading 

31import time 

32import urllib.request 

33from concurrent.futures import ProcessPoolExecutor, as_completed 

34from typing import Any 

35 

36import click 

37import yaml 

38 

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

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

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

42from lsst.resources import ResourcePath 

43from lsst.rucio.register.data_type import DataType 

44from lsst.rucio.register.rucio_interface import RucioInterface 

45from lsst.rucio.register.rucio_register_config import RucioRegisterConfig 

46 

47_thread_local = threading.local() 

48 

49logger = logging.getLogger(__name__) 

50_FORMAT = ( 

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

52) 

53 

54RUCIO_REGISTER_CONFIG = "RUCIO_REGISTER_CONFIG" 

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

56 

57 

58def register_options(func): 

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

60 func = click.option( 

61 "--backoff-max-tries", 

62 "--max-retries", 

63 "max_retries", 

64 required=False, 

65 type=int, 

66 default=RucioInterface.DEFAULT_RETRIES, 

67 show_default=True, 

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

69 )(func) 

70 func = click.option( 

71 "--backoff-max-value", 

72 required=False, 

73 type=int, 

74 default=RucioInterface.DEFAULT_BACKOFF_MAX_VALUE, 

75 show_default=True, 

76 help="maximum backoff value in seconds", 

77 )(func) 

78 func = click.option( 

79 "--backoff-factor", 

80 required=False, 

81 type=float, 

82 default=RucioInterface.DEFAULT_BACKOFF_FACTOR, 

83 show_default=True, 

84 help="multiplicative factor to backoff by", 

85 )(func) 

86 func = click.option( 

87 "--timeout", 

88 required=False, 

89 type=float, 

90 default=RucioInterface.DEFAULT_TIMEOUT, 

91 show_default=True, 

92 help=( 

93 "HTTP request timeout in seconds for Rucio client operations " 

94 "(ReplicaClient and DIDClient bulk replica registration and dataset attachment)" 

95 ), 

96 )(func) 

97 func = click.option( 

98 "--chunk-size", 

99 required=False, 

100 type=int, 

101 default=500, 

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

103 )(func) 

104 func = click.option( 

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

106 )(func) 

107 func = click.option( 

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

109 )(func) 

110 return func 

111 

112 

113def chunks(refs, chunk_size): 

114 it = iter(refs) 

115 while True: 

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

117 if not chunk: 

118 return 

119 yield chunk 

120 

121 

122_rucio_interface_cache = {} 

123 

124 

125def _get_rucio_interface( 

126 repo, 

127 rucio_register_config, 

128 rubin_butler_type, 

129 clear_is_new=False, 

130 retries=RucioInterface.DEFAULT_RETRIES, 

131 kwargs=None, 

132): 

133 if kwargs: 

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

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

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

137 timeout = kwargs.get("timeout", RucioInterface.DEFAULT_TIMEOUT) 

138 if backoff_max_tries is not None: 138 ↛ 146line 138 didn't jump to line 146 because the condition on line 138 was always true

139 retries = backoff_max_tries 

140 else: 

141 backoff_factor = RucioInterface.DEFAULT_BACKOFF_FACTOR 

142 backoff_max_value = RucioInterface.DEFAULT_BACKOFF_MAX_VALUE 

143 backoff_max_tries = retries 

144 timeout = RucioInterface.DEFAULT_TIMEOUT 

145 

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

147 cache_key = ( 

148 os.getpid(), 

149 repo, 

150 config_file, 

151 rubin_butler_type, 

152 clear_is_new, 

153 retries, 

154 backoff_factor, 

155 backoff_max_value, 

156 timeout, 

157 ) 

158 if cache_key in _rucio_interface_cache: 

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

160 return _rucio_interface_cache[cache_key] 

161 

162 if config_file is None: 162 ↛ 163line 162 didn't jump to line 163 because the condition on line 162 was never true

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

164 

165 config = RucioRegisterConfig(config_file) 

166 

167 rucio_rse = config.rucio_rse 

168 scope = config.scope 

169 rse_root = config.rse_root 

170 dtn_url = config.dtn_url 

171 

172 butler = None 

173 if repo: 173 ↛ 177line 173 didn't jump to line 177 because the condition on line 173 was always true

174 butler = Butler(repo) 

175 

176 # create RucioInterface object used to register replicas into datasets 

177 ri = RucioInterface( 

178 butler=butler, 

179 rucio_rse=rucio_rse, 

180 scope=scope, 

181 rse_root=rse_root, 

182 dtn_url=dtn_url, 

183 rubin_butler_type=rubin_butler_type, 

184 clear_is_new=clear_is_new, 

185 retries=retries, 

186 timeout=timeout, 

187 ) 

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

189 

190 _rucio_interface_cache[cache_key] = (ri, butler) 

191 return ri, butler 

192 

193 

194def worker_register_replicas( 

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

196): 

197 ri, _ = _get_rucio_interface( 

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

199 ) 

200 return ri.register_as_replicas(rucio_dataset, refs_list) 

201 

202 

203def worker_register_zips( 

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

205): 

206 ri, _ = _get_rucio_interface( 

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

208 ) 

209 return ri.register_zips(rucio_dataset, zip_files) 

210 

211 

212def worker_register_dims( 

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

214): 

215 ri, _ = _get_rucio_interface( 

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

217 ) 

218 return ri.register_dims(rucio_dataset, dim_files) 

219 

220 

221def _register( 

222 repo, 

223 rucio_register_config, 

224 rubin_butler_type, 

225 dataset_refs, 

226 chunk_size, 

227 rucio_dataset, 

228 max_workers=1, 

229 retries=RucioInterface.DEFAULT_RETRIES, 

230 stats=None, 

231 clear_is_new=False, 

232 kwargs=None, 

233): 

234 failed_uuids = [] 

235 registered_count = 0 

236 

237 total_refs = None 

238 try: 

239 total_refs = len(dataset_refs) 

240 except TypeError: 

241 pass 

242 

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

244 logger.info( 

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

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

247 ) 

248 else: 

249 logger.info( 

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

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

252 ) 

253 

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

255 total_chunks = len(ref_chunks) 

256 

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

258 mp_context = multiprocessing.get_context("spawn") 

259 log_level = kwargs.get("log_level", None) if kwargs else None 

260 with ProcessPoolExecutor( 

261 max_workers=max_workers, 

262 mp_context=mp_context, 

263 initializer=_set_log_level, 

264 initargs=(log_level,), 

265 ) as executor: 

266 future_to_chunk = { 

267 executor.submit( 

268 worker_register_replicas, 

269 repo, 

270 rucio_register_config, 

271 rubin_butler_type, 

272 rucio_dataset, 

273 refs, 

274 retries, 

275 clear_is_new, 

276 kwargs, 

277 ): (i, refs) 

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

279 } 

280 

281 for future in as_completed(future_to_chunk): 

282 chunk_idx, refs = future_to_chunk[future] 

283 try: 

284 count = future.result() 

285 registered_count += count 

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

287 except Exception as e: 

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

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

290 failed_uuids.extend(chunk_uuids) 

291 else: 

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

293 refs_list = list(refs) 

294 try: 

295 count = worker_register_replicas( 

296 repo, 

297 rucio_register_config, 

298 rubin_butler_type, 

299 rucio_dataset, 

300 refs_list, 

301 retries, 

302 clear_is_new, 

303 kwargs, 

304 ) 

305 registered_count += count 

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

307 except Exception as e: 

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

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

310 failed_uuids.extend(chunk_uuids) 

311 

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

313 stats["registered"] = registered_count 

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

315 stats["failures"] = failed_uuids 

316 

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

318 logger.error( 

319 f"Registration summary for {rucio_dataset} - " 

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

321 ) 

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

323 else: 

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

325 

326 

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

328 if isinstance(log_level, dict): 

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

330 if isinstance(log_level, str): 

331 handler = logging.StreamHandler(sys.stdout) 

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

333 log = logging.getLogger() 

334 log.handlers.clear() 

335 log.addHandler(handler) 

336 log.setLevel(log_level.upper()) 

337 

338 

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

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

341 if max_dataset_types is None: 

342 return dataset_types 

343 return list(dataset_types[:max_dataset_types]) 

344 

345 

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

347def main(): 

348 pass 

349 

350 

351@main.command() 

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

353@register_options 

354@click.option( 

355 "--max-workers", 

356 required=False, 

357 type=int, 

358 default=5, 

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

360) 

361@click.option( 

362 "--clear-is-new", 

363 is_flag=True, 

364 default=False, 

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

366) 

367@log_level_option() 

368@options_file_option() 

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

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

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

372 _set_log_level(log_level) 

373 

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

375 if not rucio_dataset: 

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

377 

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

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

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

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

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

383 

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

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

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

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

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

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

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

391 

392 ri, butler = _get_rucio_interface( 

393 repo, 

394 rucio_register_config, 

395 DataType.DATA_PRODUCT, 

396 clear_is_new=clear_is_new, 

397 retries=max_retries, 

398 kwargs=kwargs, 

399 ) 

400 

401 query = QueryDatasets( 

402 butler=butler, 

403 glob=dataset_type, 

404 collections=collections, 

405 where=where, 

406 find_first=find_first, 

407 limit=limit, 

408 order_by=order_by, 

409 show_uri=False, 

410 with_dimension_records=True, 

411 ) 

412 

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

414 

415 stats = {} 

416 _register( 

417 repo, 

418 rucio_register_config, 

419 DataType.DATA_PRODUCT, 

420 dataset_refs, 

421 chunk_size, 

422 rucio_dataset, 

423 max_workers=max_workers, 

424 retries=max_retries, 

425 stats=stats, 

426 clear_is_new=clear_is_new, 

427 kwargs=kwargs, 

428 ) 

429 

430 

431@main.command() 

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

433@register_options 

434@click.option( 

435 "--max-workers", 

436 required=False, 

437 type=int, 

438 default=5, 

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

440) 

441@click.option( 

442 "--uuidlist", 

443 required=False, 

444 type=str, 

445 help=( 

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

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

448 ), 

449) 

450@log_level_option() 

451@options_file_option() 

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

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

454 _set_log_level(log_level) 

455 

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

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

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

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

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

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

462 

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

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

465 

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

467 

468 # Infer dataset from file header if available 

469 uuids = [] 

470 inferred_dataset = None 

471 with open(uuidlist) as f: 

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

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

474 try: 

475 data = json.loads(content) 

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

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

478 elif isinstance(data, dict): 

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

480 except json.JSONDecodeError: 

481 pass 

482 

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

484 for line in content.splitlines(): 

485 stripped = line.strip() 

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

487 parts = stripped.split() 

488 if len(parts) >= 2: 

489 inferred_dataset = parts[1] 

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

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

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

493 uuids.append(clean_uuid) 

494 

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

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

497 raise click.UsageError( 

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

499 ) 

500 

501 ri, butler = _get_rucio_interface( 

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

503 ) 

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

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

506 

507 dataset_refs = butler.get_many_datasets(uuids) 

508 

509 stats = {} 

510 _register( 

511 repo, 

512 rucio_register_config, 

513 DataType.DATA_PRODUCT, 

514 dataset_refs, 

515 chunk_size, 

516 rucio_dataset, 

517 max_workers=max_workers, 

518 retries=max_retries, 

519 stats=stats, 

520 kwargs=kwargs, 

521 ) 

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

523 

524 

525@main.command() 

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

527@register_options 

528@click.option( 

529 "--max-workers", 

530 required=False, 

531 type=int, 

532 default=5, 

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

534) 

535@log_level_option() 

536@options_file_option() 

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

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

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

540 _set_log_level(log_level) 

541 

542 rucio_dataset = kwargs.get("rucio_dataset") 

543 if not rucio_dataset: 

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

545 

546 rucio_register_config = kwargs.get("rucio_register_config") 

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

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

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

550 

551 repo = kwargs["repo"] 

552 

553 ri, butler = _get_rucio_interface( 

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

555 ) 

556 

557 query_kwargs = { 

558 k: v 

559 for k, v in kwargs.items() 

560 if k 

561 not in ( 

562 "backoff_factor", 

563 "backoff_max_value", 

564 "backoff_max_tries", 

565 "max_retries", 

566 "backoff-factor", 

567 "backoff-max-value", 

568 "backoff-max-tries", 

569 "max_workers", 

570 "rucio_dataset", 

571 "rucio_register_config", 

572 "chunk_size", 

573 "timeout", 

574 "log_level", 

575 ) 

576 } 

577 

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

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

580 

581 stats = {} 

582 _register( 

583 repo, 

584 rucio_register_config, 

585 DataType.RAW_FILE, 

586 dataset_refs, 

587 chunk_size, 

588 rucio_dataset, 

589 max_workers=max_workers, 

590 retries=max_retries, 

591 stats=stats, 

592 kwargs=kwargs, 

593 ) 

594 

595 

596@main.command() 

597@register_options 

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

599@log_level_option() 

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

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

602 _set_log_level(log_level) 

603 

604 rucio_dataset = kwargs.get("rucio_dataset") 

605 if not rucio_dataset: 

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

607 

608 zip_file = kwargs.get("zip_file") 

609 rucio_register_config = kwargs.get("rucio_register_config") 

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

611 

612 ri, butler = _get_rucio_interface( 

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

614 ) 

615 

616 try: 

617 count = worker_register_zips( 

618 repo=None, 

619 rucio_config=rucio_register_config, 

620 butler_type=DataType.ZIP_FILE, 

621 rucio_dataset=rucio_dataset, 

622 zip_files=[ResourcePath(zip_file)], 

623 retries=max_retries, 

624 kwargs=kwargs, 

625 ) 

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

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

628 except Exception as e: 

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

630 

631 

632@main.command() 

633@register_options 

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

635@log_level_option() 

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

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

638 _set_log_level(log_level) 

639 

640 rucio_dataset = kwargs.get("rucio_dataset") 

641 if not rucio_dataset: 

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

643 

644 dim_file = kwargs.get("dimension_file") 

645 rucio_register_config = kwargs.get("rucio_register_config") 

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

647 

648 ri, butler = _get_rucio_interface( 

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

650 ) 

651 

652 try: 

653 count = worker_register_dims( 

654 repo=None, 

655 rucio_config=rucio_register_config, 

656 butler_type=DataType.DIM_FILE, 

657 rucio_dataset=rucio_dataset, 

658 dim_files=[ResourcePath(dim_file)], 

659 retries=max_retries, 

660 kwargs=kwargs, 

661 ) 

662 logger.info( 

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

664 ) 

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

666 except Exception as e: 

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

668 

669 

670def _process_refs_chunk( 

671 refs_chunk, 

672 rucio_dataset_name, 

673 dry_run, 

674 out_dir, 

675 repo, 

676 rucio_register_config, 

677 chunk_size, 

678 max_retries, 

679 clear_is_new=False, 

680 kwargs=None, 

681): 

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

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

684 if dry_run: 

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

686 for ref in refs_chunk: 

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

688 

689 # Sanitize repo path for use in output directory structure 

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

691 # Sanitize rucio_dataset_name for use as a filename 

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

693 

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

695 os.makedirs(uuiddir, exist_ok=True) 

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

697 

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

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

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

701 else: 

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

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

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

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

706 # This is now the default for _register. 

707 _register( 

708 repo, 

709 rucio_register_config, 

710 DataType.DATA_PRODUCT, 

711 refs_chunk, 

712 chunk_size, 

713 rucio_dataset_name, 

714 retries=max_retries, 

715 stats=local_stats, 

716 clear_is_new=clear_is_new, 

717 kwargs=kwargs, 

718 ) 

719 return local_stats 

720 

721 

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

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

724 auto-registration. 

725 """ 

726 refs = kwargs.get("refs") 

727 ds_name = kwargs.get("ds_name") 

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

729 repo = kwargs.get("repo") 

730 rucio_register_config = kwargs.get("rucio_register_config") 

731 dry_run = kwargs.get("dry_run") 

732 out_dir = kwargs.get("out_dir") 

733 chunk_size = kwargs.get("chunk_size") 

734 max_retries = kwargs.get("max_retries") 

735 logger.info( 

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

737 len(refs), 

738 ds_name, 

739 ) 

740 return _process_refs_chunk( 

741 refs, 

742 ds_name, 

743 dry_run, 

744 out_dir, 

745 repo, 

746 rucio_register_config, 

747 chunk_size, 

748 max_retries, 

749 clear_is_new=clear_is_new, 

750 kwargs=kwargs, 

751 ) 

752 

753 

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

755 content = "" 

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

757 try: 

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

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

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

761 except Exception as e: 

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

763 raise 

764 else: 

765 try: 

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

767 with open(transfer_list) as f: 

768 content = f.read() 

769 except Exception as e: 

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

771 raise 

772 

773 try: 

774 data = yaml.safe_load(content) 

775 except Exception as e: 

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

777 raise 

778 

779 curated_types = set() 

780 

781 def _traverse(node): 

782 if isinstance(node, list): 

783 for item in node: 

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

785 curated_types.add(item) 

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

787 for val in node.values(): 

788 _traverse(val) 

789 

790 _traverse(data) 

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

792 return curated_types 

793 

794 

795def _query_butler_task(repo: str, task: dict[str, Any]) -> tuple[dict[str, Any], list[Any]]: 

796 collection = task["collection"] 

797 dataset_type = task["dataset_type"] 

798 where = task["where"] 

799 

800 try: 

801 if not hasattr(_thread_local, "butlers"): 

802 _thread_local.butlers = {} 

803 if repo not in _thread_local.butlers: 

804 _thread_local.butlers[repo] = Butler(repo)