Coverage for python/lsst/ctrl/bps/panda/utils.py: 5%

455 statements  

« prev     ^ index     » next       coverage.py v7.10.6, created at 2025-10-11 01:37 -0700

1# This file is part of ctrl_bps_panda. 

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 software is dual licensed under the GNU General Public License and also 

10# under a 3-clause BSD license. Recipients may choose which of these licenses 

11# to use; please see the files gpl-3.0.txt and/or bsd_license.txt, 

12# respectively. If you choose the GPL option then the following text applies 

13# (but note that there is still no warranty even if you opt for BSD instead): 

14# 

15# This program is free software: you can redistribute it and/or modify 

16# it under the terms of the GNU General Public License as published by 

17# the Free Software Foundation, either version 3 of the License, or 

18# (at your option) any later version. 

19# 

20# This program is distributed in the hope that it will be useful, 

21# but WITHOUT ANY WARRANTY; without even the implied warranty of 

22# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the 

23# GNU General Public License for more details. 

24# 

25# You should have received a copy of the GNU General Public License 

26# along with this program. If not, see <https://www.gnu.org/licenses/>. 

27 

28"""Utilities for bps PanDA plugin.""" 

29 

30__all__ = [ 

31 "add_decoder_prefix", 

32 "convert_exec_string_to_hex", 

33 "copy_files_for_distribution", 

34 "get_idds_client", 

35 "get_idds_result", 

36] 

37 

38import binascii 

39import concurrent.futures 

40import json 

41import logging 

42import os 

43import random 

44import tarfile 

45import time 

46import uuid 

47 

48import idds.common.utils as idds_utils 

49import pandaclient.idds_api 

50from idds.doma.workflowv2.domapandawork import DomaPanDAWork 

51from idds.workflowv2.workflow import AndCondition 

52from idds.workflowv2.workflow import Workflow as IDDS_client_workflow 

53 

54from lsst.ctrl.bps import BpsConfig, GenericWorkflow, GenericWorkflowJob 

55from lsst.ctrl.bps.panda.cmd_line_embedder import CommandLineEmbedder 

56from lsst.ctrl.bps.panda.constants import ( 

57 PANDA_DEFAULT_CLOUD, 

58 PANDA_DEFAULT_CORE_COUNT, 

59 PANDA_DEFAULT_MAX_ATTEMPTS, 

60 PANDA_DEFAULT_MAX_JOBS_PER_TASK, 

61 PANDA_DEFAULT_MAX_PAYLOADS_PER_PANDA_JOB, 

62 PANDA_DEFAULT_MAX_WALLTIME, 

63 PANDA_DEFAULT_NAME_LENGTH, 

64 PANDA_DEFAULT_ORDER_ID_MAP_FILE, 

65 PANDA_DEFAULT_PRIORITY, 

66 PANDA_DEFAULT_PROCESSING_TYPE, 

67 PANDA_DEFAULT_PROD_SOURCE_LABEL, 

68 PANDA_DEFAULT_RSS, 

69 PANDA_DEFAULT_RSS_MAX, 

70 PANDA_DEFAULT_TASK_TYPE, 

71 PANDA_DEFAULT_VO, 

72) 

73from lsst.resources import ResourcePath 

74 

75_LOG = logging.getLogger(__name__) 

76 

77 

78def copy_files_for_distribution(files_to_stage, file_distribution_uri, max_copy_workers): 

79 """Brings locally generated files into Cloud for further 

80 utilization them on the edge nodes. 

81 

82 Parameters 

83 ---------- 

84 files_to_stage : `dict` [`str`, `str`] 

85 Files which need to be copied to a workflow staging area. 

86 file_distribution_uri : `ResourcePath` 

87 Path on the edge node accessed storage, 

88 including access protocol, bucket name to place files. 

89 max_copy_workers : `int` 

90 Maximum number of workers for copying files. 

91 

92 Raises 

93 ------ 

94 RuntimeError 

95 Raised when error copying files to the distribution point. 

96 """ 

97 files_to_copy = {} 

98 

99 # In case there are folders we iterate over its content 

100 for local_pfn in files_to_stage.values(): 

101 folder_name = os.path.basename(os.path.normpath(local_pfn)) 

102 if os.path.isdir(local_pfn): 

103 folder_uri = file_distribution_uri.join(folder_name, forceDirectory=True) 

104 files_in_folder = ResourcePath.findFileResources([local_pfn]) 

105 for file in files_in_folder: 

106 file_name = file.basename() 

107 files_to_copy[file] = folder_uri.join(file_name, forceDirectory=False) 

108 else: 

109 folder_uri = file_distribution_uri.join(folder_name, forceDirectory=False) 

110 files_to_copy[ResourcePath(local_pfn, forceDirectory=False)] = folder_uri 

111 

112 copy_executor = concurrent.futures.ThreadPoolExecutor(max_workers=max_copy_workers) 

113 future_file_copy = [] 

114 for src, trgt in files_to_copy.items(): 

115 _LOG.debug("Staging %s to %s", src, trgt) 

116 # S3 clients explicitly instantiate here to overpass this 

117 # https://stackoverflow.com/questions/52820971/is-boto3-client-thread-safe 

118 trgt.exists() 

119 future_file_copy.append(copy_executor.submit(trgt.transfer_from, src, transfer="copy")) 

120 

121 for future in concurrent.futures.as_completed(future_file_copy): 

122 if future.result() is not None: 

123 raise RuntimeError("Error of placing files to the distribution point") 

124 

125 

126def get_idds_client(config): 

127 """Get the idds client. 

128 

129 Parameters 

130 ---------- 

131 config : `lsst.ctrl.bps.BpsConfig` 

132 BPS configuration. 

133 

134 Returns 

135 ------- 

136 idds_client: `idds.client.clientmanager.ClientManager` 

137 The iDDS ClientManager object. 

138 """ 

139 idds_server = None 

140 if isinstance(config, BpsConfig): 

141 _, idds_server = config.search("iddsServer", opt={"default": None}) 

142 elif isinstance(config, dict) and "iddsServer" in config: 

143 idds_server = config["iddsServer"] 

144 # if idds_server is None, a default value on the panda relay service 

145 # will be used 

146 idds_client = pandaclient.idds_api.get_api( 

147 idds_utils.json_dumps, idds_host=idds_server, compress=True, manager=True 

148 ) 

149 return idds_client 

150 

151 

152def get_idds_result(ret): 

153 """Parse the results returned from iDDS. 

154 

155 Parameters 

156 ---------- 

157 ret : `tuple` [`int`, `tuple` [`bool`, payload ]] 

158 The first part ``ret[0]`` is the status of PanDA relay service. 

159 The part of ``ret[1][0]`` is the status of iDDS service. 

160 The part of ``ret[1][1]`` is the returned payload. 

161 If ``ret[1][0]`` is `False`, ``ret[1][1]`` can be error messages. 

162 

163 Returns 

164 ------- 

165 status: `bool` 

166 The status of iDDS calls. 

167 result: `int` or `list` or `dict` or `None` 

168 The result returned from iDDS. `None` if error state. 

169 error: `str` or `None` 

170 Error messages. `None` if no error state. 

171 """ 

172 # https://panda-wms.readthedocs.io/en/latest/client/rest_idds.html 

173 if not isinstance(ret, list | tuple) or ret[0] != 0: 

174 # Something wrong with the PanDA relay service. 

175 # The call may not be delivered to iDDS. 

176 status = False 

177 result = None 

178 error = f"PanDA relay service returns errors: {ret}" 

179 else: 

180 if ret[1][0]: 

181 status = True 

182 result = ret[1][1] 

183 error = None 

184 if isinstance(result, str) and "Authentication no permission" in result: 

185 status = False 

186 result = None 

187 error = result 

188 else: 

189 # iDDS returns errors 

190 status = False 

191 result = None 

192 error = f"iDDS returns errors: {ret[1][1]}" 

193 return status, result, error 

194 

195 

196def _make_pseudo_filename(config, gwjob): 

197 """Make the job pseudo filename. 

198 

199 Parameters 

200 ---------- 

201 config : `lsst.ctrl.bps.BpsConfig` 

202 BPS configuration. 

203 gwjob : `lsst.ctrl.bps.GenericWorkflowJob` 

204 Job for which to create the pseudo filename. 

205 

206 Returns 

207 ------- 

208 pseudo_filename : `str` 

209 The pseudo filename for the given job. 

210 """ 

211 cmd_line_embedder = CommandLineEmbedder(config) 

212 _, pseudo_filename = cmd_line_embedder.substitute_command_line( 

213 gwjob.executable.src_uri + " " + gwjob.arguments, gwjob.cmdvals, gwjob.name, [] 

214 ) 

215 return pseudo_filename 

216 

217 

218def _make_doma_work( 

219 config, 

220 generic_workflow, 

221 gwjob, 

222 task_count, 

223 task_chunk, 

224 enable_event_service=False, 

225 enable_job_name_map=False, 

226 order_id_map_files=None, 

227 es_label=None, 

228 max_payloads_per_panda_job=PANDA_DEFAULT_MAX_PAYLOADS_PER_PANDA_JOB, 

229 max_wms_job_wall_time=None, 

230 remote_filename=None, 

231 qnode_map_filename=None, 

232): 

233 """Make the DOMA Work object for a PanDA task. 

234 

235 Parameters 

236 ---------- 

237 config : `lsst.ctrl.bps.BpsConfig` 

238 BPS configuration. 

239 gwjob : `lsst.ctrl.bps.GenericWorkflowJob` 

240 Job representing the jobs for the PanDA task. 

241 task_count : `int` 

242 Count of PanDA tasks used when making unique names. 

243 task_chunk : `int` 

244 Count of chunk of a PanDA tasks used when making unique names. 

245 

246 Returns 

247 ------- 

248 work : `idds.doma.workflowv2.domapandawork.DomaPanDAWork` 

249 The client representation of a PanDA task. 

250 local_pfns : `dict` [`str`, `str`] 

251 Files which need to be copied to a workflow staging area. 

252 """ 

253 if order_id_map_files is None: 

254 order_id_map_files = {} 

255 _LOG.debug("Using gwjob %s to create new PanDA task (gwjob=%s)", gwjob.name, gwjob) 

256 cvals = {"curr_cluster": gwjob.label} 

257 _, site = config.search("computeSite", opt={"curvals": cvals, "required": True}) 

258 cvals["curr_site"] = site 

259 cvals["curr_pipetask"] = gwjob.label 

260 _, processing_type = config.search( 

261 "processingType", opt={"curvals": cvals, "default": PANDA_DEFAULT_PROCESSING_TYPE} 

262 ) 

263 if gwjob.label in ["finalJob", "customJob"]: 

264 _, nonpipetask = config.search(gwjob.label) 

265 default_type = "Rubin_Merge" 

266 if gwjob.label == "customJob": 

267 default_type = PANDA_DEFAULT_PROCESSING_TYPE 

268 processing_type = nonpipetask["processingType"] if nonpipetask["processingType"] else default_type 

269 _, task_type = config.search("taskType", opt={"curvals": cvals, "default": PANDA_DEFAULT_TASK_TYPE}) 

270 _, prod_source_label = config.search( 

271 "prodSourceLabel", opt={"curvals": cvals, "default": PANDA_DEFAULT_PROD_SOURCE_LABEL} 

272 ) 

273 _, vo = config.search("vo", opt={"curvals": cvals, "default": PANDA_DEFAULT_VO}) 

274 

275 _, file_distribution_end_point = config.search( 

276 "fileDistributionEndPoint", opt={"curvals": cvals, "default": None} 

277 ) 

278 

279 _, file_distribution_end_point_default = config.search( 

280 "fileDistributionEndPointDefault", opt={"curvals": cvals, "default": None} 

281 ) 

282 

283 task_rss = gwjob.request_memory if gwjob.request_memory else PANDA_DEFAULT_RSS 

284 task_rss_retry_step = task_rss * gwjob.memory_multiplier if gwjob.memory_multiplier else 0 

285 task_rss_retry_offset = 0 if task_rss_retry_step else task_rss 

286 

287 # Assume input files are same across task 

288 local_pfns = {} 

289 direct_io_files = set() 

290 

291 if gwjob.executable.transfer_executable: 

292 local_pfns["job_executable"] = gwjob.executable.src_uri 

293 job_executable = f"./{os.path.basename(gwjob.executable.src_uri)}" 

294 else: 

295 job_executable = gwjob.executable.src_uri 

296 cmd_line_embedder = CommandLineEmbedder(config) 

297 _LOG.debug( 

298 "job %s inputs = %s, outputs = %s", 

299 gwjob.name, 

300 generic_workflow.get_job_inputs(gwjob.name), 

301 generic_workflow.get_job_outputs(gwjob.name), 

302 ) 

303 

304 job_env = "" 

305 if gwjob.environment: 

306 for key, value in gwjob.environment.items(): 

307 try: 

308 sub_value = value.format_map(gwjob.cmdvals) 

309 except (KeyError, TypeError) as exc: 

310 _LOG.error("Could not replace command variables: replacement for %s not provided", str(exc)) 

311 raise 

312 job_env += f"export {key}={sub_value}; " 

313 

314 cmd_line, _ = cmd_line_embedder.substitute_command_line( 

315 job_env + job_executable + " " + gwjob.arguments, 

316 gwjob.cmdvals, 

317 gwjob.name, 

318 generic_workflow.get_job_inputs(gwjob.name) + generic_workflow.get_job_outputs(gwjob.name), 

319 ) 

320 

321 my_log = f"enable_event_service {enable_event_service} for {gwjob.label}" 

322 _LOG.info(my_log) 

323 if enable_event_service: 

324 if gwjob.request_walltime and max_wms_job_wall_time: 

325 my_log = ( 

326 f"requestWalltime({gwjob.request_walltime}) " 

327 f"and maxWmsJobWalltime({max_wms_job_wall_time}) are set, " 

328 "max_payloads_per_panda_job is int(max_wms_job_wall_time / gwjob.request_walltime), " 

329 "ignore maxPayloadsPerPandaJob." 

330 ) 

331 _LOG.info(my_log) 

332 max_payloads_per_panda_job = int(max_wms_job_wall_time / gwjob.request_walltime) 

333 if max_payloads_per_panda_job < 2: 

334 my_log = ( 

335 f"max_payloads_per_panda_job ({max_payloads_per_panda_job}) is too small, " 

336 "disable EventService" 

337 ) 

338 _LOG.info(my_log) 

339 enable_event_service = False 

340 

341 maxwalltime = gwjob.request_walltime if gwjob.request_walltime else PANDA_DEFAULT_MAX_WALLTIME 

342 if enable_event_service: 

343 if gwjob.request_walltime and max_payloads_per_panda_job: 

344 maxwalltime = gwjob.request_walltime * max_payloads_per_panda_job 

345 elif max_wms_job_wall_time: 

346 maxwalltime = max_wms_job_wall_time 

347 

348 if enable_event_service or enable_job_name_map: 

349 for es_name in order_id_map_files: 

350 local_pfns[es_name] = order_id_map_files[es_name] 

351 

352 for gwfile in generic_workflow.get_job_inputs(gwjob.name, transfer_only=True): 

353 local_pfns[gwfile.name] = gwfile.src_uri 

354 if os.path.isdir(gwfile.src_uri): 

355 # this is needed to make isdir function working 

356 # properly in ButlerURL instance on the edge node 

357 local_pfns[gwfile.name] += "/" 

358 

359 if gwfile.job_access_remote: 

360 direct_io_files.add(gwfile.name) 

361 

362 if qnode_map_filename: 

363 local_pfns.update(qnode_map_filename) 

364 

365 submit_cmd = generic_workflow.run_attrs.get("bps_iscustom", False) 

366 

367 if not direct_io_files: 

368 if submit_cmd: 

369 direct_io_files.add(remote_filename) 

370 else: 

371 direct_io_files.add("cmdlineplaceholder") 

372 

373 lsst_temp = "LSST_RUN_TEMP_SPACE" 

374 if lsst_temp in file_distribution_end_point and lsst_temp not in os.environ: 

375 file_distribution_end_point = file_distribution_end_point_default 

376 if submit_cmd and not file_distribution_end_point: 

377 file_distribution_end_point = "FileDistribution" 

378 

379 executable = add_decoder_prefix( 

380 config, cmd_line, file_distribution_end_point, (local_pfns, direct_io_files) 

381 ) 

382 work = DomaPanDAWork( 

383 executable=executable, 

384 primary_input_collection={ 

385 "scope": "pseudo_dataset", 

386 "name": f"pseudo_input_collection#{task_count}", 

387 }, 

388 output_collections=[{"scope": "pseudo_dataset", "name": f"pseudo_output_collection#{task_count}"}], 

389 log_collections=[], 

390 dependency_map=[], 

391 task_name=f"{generic_workflow.name}_{task_count:02d}_{gwjob.label}_{task_chunk:02d}", 

392 task_queue=gwjob.queue, 

393 task_log={ 

394 "destination": "local", 

395 "value": "log.tgz", 

396 "dataset": "PandaJob_#{pandaid}/", 

397 "token": "local", 

398 "param_type": "log", 

399 "type": "template", 

400 }, 

401 encode_command_line=True, 

402 task_rss=task_rss, 

403 task_rss_retry_offset=task_rss_retry_offset, 

404 task_rss_retry_step=task_rss_retry_step, 

405 task_rss_max=gwjob.request_memory_max if gwjob.request_memory_max else PANDA_DEFAULT_RSS_MAX, 

406 task_cloud=gwjob.compute_cloud if gwjob.compute_cloud else PANDA_DEFAULT_CLOUD, 

407 task_site=site, 

408 task_priority=int(gwjob.priority) if gwjob.priority else PANDA_DEFAULT_PRIORITY, 

409 core_count=gwjob.request_cpus if gwjob.request_cpus else PANDA_DEFAULT_CORE_COUNT, 

410 working_group=gwjob.accounting_group, 

411 processing_type=processing_type, 

412 task_type=task_type, 

413 prodSourceLabel=prod_source_label, 

414 vo=vo, 

415 es=enable_event_service, 

416 es_label=es_label, 

417 max_events_per_job=max_payloads_per_panda_job, 

418 maxattempt=gwjob.number_of_retries if gwjob.number_of_retries else PANDA_DEFAULT_MAX_ATTEMPTS, 

419 maxwalltime=maxwalltime, 

420 ) 

421 return work, local_pfns 

422 

423 

424def add_final_idds_work( 

425 config, generic_workflow, idds_client_workflow, dag_sink_work, task_count, task_chunk 

426): 

427 """Add the special final PanDA task to the client workflow. 

428 

429 Parameters 

430 ---------- 

431 config : `lsst.ctrl.bps.BpsConfig` 

432 BPS configuration. 

433 generic_workflow : `lsst.ctrl.bps.GenericWorkflow` 

434 Generic workflow in which to find the final job. 

435 idds_client_workflow : `idds.workflowv2.workflow.Workflow` 

436 The iDDS client representation of the workflow to which the final task 

437 is added. 

438 dag_sink_work : `list` [`idds.doma.workflowv2.domapandawork.DomaPanDAWork`] 

439 The work nodes in the client workflow which have no successors. 

440 task_count : `int` 

441 Count of PanDA tasks used when making unique names. 

442 task_chunk : `int` 

443 Count of chunk of a PanDA tasks used when making unique names. 

444 

445 Returns 

446 ------- 

447 files : `dict` [`str`, `str`] 

448 Files which need to be copied to a workflow staging area. 

449 

450 Raises 

451 ------ 

452 NotImplementedError 

453 Raised if final job in GenericWorkflow is itself a workflow. 

454 TypeError 

455 Raised if final job in GenericWorkflow is invalid type. 

456 """ 

457 files = {} 

458 

459 # If final job exists in generic workflow, create DAG final job 

460 final = generic_workflow.get_final() 

461 if final: 

462 if isinstance(final, GenericWorkflow): 

463 raise NotImplementedError("PanDA plugin does not support a workflow as the final job") 

464 

465 if not isinstance(final, GenericWorkflowJob): 

466 raise TypeError(f"Invalid type for GenericWorkflow.get_final() results ({type(final)})") 

467 

468 dag_final_work, files = _make_doma_work( 

469 config, 

470 generic_workflow, 

471 final, 

472 task_count, 

473 task_chunk, 

474 ) 

475 pseudo_filename = "pure_pseudoinput+qgraphNodeId:+qgraphId:" 

476 dag_final_work.dependency_map.append( 

477 {"name": pseudo_filename, "submitted": False, "dependencies": []} 

478 ) 

479 idds_client_workflow.add_work(dag_final_work) 

480 conditions = [] 

481 for work in dag_sink_work: 

482 conditions.append(work.is_terminated) 

483 and_cond = AndCondition(conditions=conditions, true_works=[dag_final_work]) 

484 idds_client_workflow.add_condition(and_cond) 

485 else: 

486 _LOG.debug("No final job in GenericWorkflow") 

487 return files 

488 

489 

490def convert_exec_string_to_hex(cmdline): 

491 """Convert the command line into hex representation. 

492 

493 This step is currently involved because large blocks of command lines 

494 including special symbols passed to the pilot/container. To make sure 

495 the 1 to 1 matching and pass by the special symbol stripping 

496 performed by the Pilot we applied the hexing. 

497 

498 Parameters 

499 ---------- 

500 cmdline : `str` 

501 UTF-8 command line string. 

502 

503 Returns 

504 ------- 

505 hex : `str` 

506 Hex representation of string. 

507 """ 

508 return binascii.hexlify(cmdline.encode()).decode("utf-8") 

509 

510 

511def add_decoder_prefix(config, cmd_line, distribution_path, files): 

512 """Compose the command line sent to the pilot from the functional part 

513 (the actual SW running) and the middleware part (containers invocation). 

514 

515 Parameters 

516 ---------- 

517 config : `lsst.ctrl.bps.BpsConfig` 

518 Configuration information. 

519 cmd_line : `str` 

520 UTF-8 based functional part of the command line. 

521 distribution_path : `str` 

522 URI of path where all files are located for distribution. 

523 files : `tuple` [`dict` [`str`, `str`], `list` [`str`]] 

524 File names needed for a task (copied local, direct access). 

525 

526 Returns 

527 ------- 

528 decoder_prefix : `str` 

529 Full command line to be executed on the edge node. 

530 """ 

531 # Manipulate file paths for placement on cmdline 

532 files_plc_hldr = {} 

533 for key, pfn in files[0].items(): 

534 if pfn.endswith("/"): 

535 files_plc_hldr[key] = os.path.basename(pfn[:-1]) 

536 isdir = True 

537 else: 

538 files_plc_hldr[key] = os.path.basename(pfn) 

539 _, extension = os.path.splitext(pfn) 

540 isdir = os.path.isdir(pfn) or (key == "butlerConfig" and extension != "yaml") 

541 if isdir: 

542 # this is needed to make isdir function working 

543 # properly in ButlerURL instance on the egde node 

544 files_plc_hldr[key] += "/" 

545 _LOG.debug("files_plc_hldr[%s] = %s", key, files_plc_hldr[key]) 

546 

547 cmdline_hex = convert_exec_string_to_hex(cmd_line) 

548 _, runner_command = config.search("runnerCommand", opt={"replaceEnvVars": False, "expandEnvVars": False}) 

549 order_id_map_filename = files[0].get("orderIdMapFilename", None) 

550 if order_id_map_filename: 

551 order_id_map_filename = os.path.basename(order_id_map_filename) 

552 order_id_map_filename = os.path.join(distribution_path, order_id_map_filename) 

553 runner_command = runner_command.replace("orderIdMapFilename", order_id_map_filename) 

554 runner_command = runner_command.replace("\n", " ") 

555 decoder_prefix = runner_command.replace( 

556 "_cmd_line_", 

557 str(cmdline_hex) 

558 + " ${IN/L} " 

559 + distribution_path 

560 + " " 

561 + "+".join(f"{k}:{v}" for k, v in files_plc_hldr.items()) 

562 + " " 

563 + "+".join(files[1]), 

564 ) 

565 return decoder_prefix 

566 

567 

568def add_idds_work(config, generic_workflow, idds_workflow): 

569 """Convert GenericWorkflowJobs to iDDS work and add them to the iDDS 

570 workflow. 

571 

572 Parameters 

573 ---------- 

574 config : `lsst.ctrl.bps.BpsConfig` 

575 BPS configuration. 

576 generic_workflow : `lsst.ctrl.bps.GenericWorkflow` 

577 Generic workflow containing jobs to convert. 

578 idds_workflow : `idds.workflowv2.workflow.Workflow` 

579 The iDDS workflow to which the converted jobs should be added. 

580 

581 Returns 

582 ------- 

583 files_to_pre_stage : `dict` [`str`, `str`] 

584 Files that need to be copied to the staging area before submission. 

585 dag_sink_work : `list` [`idds.doma.workflowv2.domapandawork.DomaPanDAWork`] 

586 The work nodes in the client workflow which have no successors. 

587 task_count : `int` 

588 Number of tasks in iDDS workflow used for unique task names. 

589 

590 Raises 

591 ------ 

592 RuntimeError 

593 If cannot recover from dependency issues after pass through workflow. 

594 """ 

595 # event service 

596 _, enable_event_service = config.search("enableEventService", opt={"default": None}) 

597 _, enable_qnode_map = config.search("enableQnodeMap", opt={"default": None}) 

598 _, max_payloads_per_panda_job = config.search( 

599 "maxPayloadsPerPandaJob", opt={"default": PANDA_DEFAULT_MAX_PAYLOADS_PER_PANDA_JOB} 

600 ) 

601 _, max_wms_job_wall_time = config.search("maxWmsJobWalltime", opt={"default": None}) 

602 my_log = ( 

603 f"enableEventService: {enable_event_service}, maxPayloadsPerPandaJob: {max_payloads_per_panda_job}" 

604 ) 

605 _LOG.info(my_log) 

606 

607 # job name map: Use a short job name to map the long job name 

608 _, enable_job_name_map = config.search("enableJobNameMap", opt={"default": None}) 

609 _LOG.info(f"enable_job_name_map: {enable_job_name_map}, {type(enable_job_name_map)}") 

610 if enable_event_service and not enable_job_name_map: 

611 enable_job_name_map = True 

612 my_log = "enable_event_service is set, set enable_job_name_map True." 

613 _LOG.info(my_log) 

614 

615 # Limit number of jobs in single PanDA task 

616 _, max_jobs_per_task = config.search("maxJobsPerTask", opt={"default": PANDA_DEFAULT_MAX_JOBS_PER_TASK}) 

617 

618 files_to_pre_stage = {} 

619 dag_sink_work = [] # Workflow sink nodes that need to be connected to final task 

620 job_to_task = {} 

621 job_to_pseudo_filename = {} 

622 task_count = 0 # Task number/ID in idds workflow used for unique name 

623 remote_archive_filename = None 

624 

625 submit_path = config["submitPath"] 

626 

627 submit_cmd = generic_workflow.run_attrs.get("bps_iscustom", False) 

628 if submit_cmd: 

629 files = generic_workflow.get_executables(data=False, transfer_only=True) 

630 archive_filename = f"jobO.{uuid.uuid4()}.tar.gz" 

631 archive_filename = create_archive_file(submit_path, archive_filename, files) 

632 remote_archive_filename = copy_files_to_pandacache(archive_filename) 

633 

634 order_id_map_files = {} 

635 name_works = {} 

636 order_id_map = {} 

637 job_name_to_order_id_map = {} 

638 order_id_map_file = None 

639 max_payloads_per_panda_job_by_label = {} 

640 if enable_event_service: 

641 enable_event_service = enable_event_service.split(",") 

642 enable_event_service_tmp = [] 

643 for es_def in enable_event_service: 

644 if ":" in es_def: 

645 es_label, m_payloads = es_def.split(":") 

646 else: 

647 es_label, m_payloads = es_def, max_payloads_per_panda_job 

648 es_label = es_label.strip() 

649 enable_event_service_tmp.append(es_label) 

650 max_payloads_per_panda_job_by_label[es_label] = int(m_payloads) 

651 enable_event_service = enable_event_service_tmp 

652 if enable_job_name_map: 

653 _, order_id_map_filename = config.search( 

654 "orderIdMapFilename", opt={"default": PANDA_DEFAULT_ORDER_ID_MAP_FILE} 

655 ) 

656 order_id_map_file = os.path.join(submit_path, order_id_map_filename) 

657 order_id_map_files = {"orderIdMapFilename": order_id_map_file} 

658 files_to_pre_stage.update(order_id_map_files) 

659 

660 # To avoid dying due to optimizing number of times through workflow, 

661 # catch dependency issues to loop through again later. 

662 jobs_with_dependency_issues = {} 

663 

664 # Initialize quantum node map 

665 qnode_map = {} 

666 qnode_map_filename = None 

667 if enable_qnode_map: 

668 qnode_map_file = os.path.join(submit_path, "qnode_map.json") 

669 qnode_map_filename = {"qnodemap": qnode_map_file} 

670 files_to_pre_stage.update(qnode_map_filename) 

671 

672 # Assume jobs with same label share config values 

673 for job_label in generic_workflow.labels: 

674 _LOG.debug("job_label = %s", job_label) 

675 

676 if enable_job_name_map: 

677 order_id_map[job_label] = {} 

678 job_name_to_order_id_map[job_label] = {} 

679 

680 # Add each job with a particular label to a corresponding PanDA task 

681 # A PanDA task has a limit on number of jobs, so break into multiple 

682 # PanDA tasks if needed. 

683 job_count = 0 # Number of jobs in idds task used for task chunking 

684 task_chunk = 1 # Task chunk number within job label used for unique name 

685 work = None 

686 order_id = -1 

687 

688 # Instead of changing code to make chunks up front and round-robin 

689 # assign jobs to chunks, for now keeping chunk creation in loop 

690 # but using knowledge of how many chunks there will be to set better 

691 # maximum number of jobs in a chunk for more even distribution. 

692 jobs_by_label = generic_workflow.get_jobs_by_label(job_label) 

693 num_chunks = -(-len(jobs_by_label) // max_jobs_per_task) # ceil 

694 max_jobs_per_task_this_label = -(-len(jobs_by_label) // num_chunks) 

695 _LOG.debug( 

696 "For job_label = %s, num jobs = %s, num_chunks = %s, max_jobs = %s", 

697 job_label, 

698 len(jobs_by_label), 

699 num_chunks, 

700 max_jobs_per_task_this_label, 

701 ) 

702 for gwjob in jobs_by_label: 

703 order_id += 1 

704 pseudo_filename = _make_pseudo_filename(config, gwjob) 

705 job_to_pseudo_filename[gwjob.name] = pseudo_filename 

706 if enable_job_name_map: 

707 order_id_map[job_label][str(order_id)] = pseudo_filename 

708 job_name_to_order_id_map[job_label][gwjob.name] = str(order_id) 

709 

710 job_count += 1 

711 if job_count > max_jobs_per_task_this_label: 

712 job_count = 1 

713 task_chunk += 1 

714 

715 if job_count == 1: 

716 # Create new PanDA task object 

717 task_count += 1 

718 work_enable_event_service = False 

719 if enable_event_service and job_label in enable_event_service: 

720 work_enable_event_service = True 

721 max_payloads_per_panda_job_current = max_payloads_per_panda_job_by_label.get( 

722 job_label, max_payloads_per_panda_job 

723 ) 

724 work, files = _make_doma_work( 

725 config, 

726 generic_workflow, 

727 gwjob, 

728 task_count, 

729 task_chunk, 

730 enable_event_service=work_enable_event_service, 

731 enable_job_name_map=enable_job_name_map, 

732 order_id_map_files=order_id_map_files, 

733 es_label=job_label, 

734 max_payloads_per_panda_job=max_payloads_per_panda_job_current, 

735 max_wms_job_wall_time=max_wms_job_wall_time, 

736 remote_filename=remote_archive_filename, 

737 qnode_map_filename=qnode_map_filename, 

738 ) 

739 work.dependency_tasks = [] 

740 name_works[work.task_name] = work 

741 files_to_pre_stage.update(files) 

742 idds_workflow.add_work(work) 

743 if generic_workflow.out_degree(gwjob.name) == 0: 

744 dag_sink_work.append(work) 

745 

746 if enable_qnode_map: 

747 job_name_PH = "PH:" + gwjob.name 

748 job_to_pseudo_filename[gwjob.name] = job_name_PH 

749 qnode_map[job_name_PH] = pseudo_filename 

750 

751 job_to_task[gwjob.name] = work.get_work_name() 

752 deps = [] 

753 missing_deps = False 

754 for parent_job_name in generic_workflow.predecessors(gwjob.name): 

755 if parent_job_name not in job_to_task: 

756 _LOG.debug("job_to_task.keys() = %s", job_to_task.keys()) 

757 missing_deps = True 

758 break 

759 else: 

760 if enable_job_name_map: 

761 parent_job = generic_workflow.get_job(parent_job_name) 

762 parent_job_label = parent_job.label 

763 parent_order_id = job_name_to_order_id_map[parent_job_label][parent_job_name] 

764 inputname = f"{parent_job_label}:orderIdMap_{parent_order_id}" 

765 else: 

766 inputname = job_to_pseudo_filename[parent_job_name] 

767 

768 parent_task_name = job_to_task[parent_job_name] 

769 deps.append( 

770 { 

771 "task": parent_task_name, 

772 "inputname": inputname, 

773 } 

774 ) 

775 if parent_task_name not in work.dependency_tasks: 

776 work.dependency_tasks.append(parent_task_name) 

777 if not missing_deps: 

778 j_name = job_to_pseudo_filename[gwjob.name] 

779 f_name = f"{job_label}:orderIdMap_{order_id}" if enable_job_name_map else j_name 

780 work.dependency_map.append( 

781 { 

782 "name": f_name, 

783 "order_id": order_id, 

784 "dependencies": deps, 

785 } 

786 ) 

787 else: 

788 jobs_with_dependency_issues[gwjob.name] = { 

789 "work": work, 

790 "order_id": order_id, 

791 "label": job_label, 

792 } 

793 

794 if enable_qnode_map: 

795 with open(qnode_map_file, "w", encoding="utf-8") as f: 

796 json.dump(qnode_map, f, indent=2) 

797 

798 # If there were any issues figuring out dependencies through earlier loop 

799 if jobs_with_dependency_issues: 

800 _LOG.warning("Could not prepare workflow in single pass. Please notify developers.") 

801 _LOG.info("Trying to recover...") 

802 for job_name, work_item in jobs_with_dependency_issues.items(): 

803 deps = [] 

804 work = work_item["work"] 

805 order_id = work_item["order_id"] 

806 job_label = work_item["label"] 

807 

808 for parent_job_name in generic_workflow.predecessors(job_name): 

809 if parent_job_name not in job_to_task: 

810 _LOG.debug("job_to_task.keys() = %s", job_to_task.keys()) 

811 raise RuntimeError( 

812 "Could not recover from dependency issues ({job_name} missing {parent_job_name})." 

813 ) 

814 if enable_job_name_map: 

815 parent_job = generic_workflow.get_job(parent_job_name) 

816 parent_job_label = parent_job.label 

817 parent_order_id = job_name_to_order_id_map[parent_job_label][parent_job_name] 

818 inputname = f"{parent_job_label}:orderIdMap_{parent_order_id}" 

819 else: 

820 inputname = job_to_pseudo_filename[parent_job_name] 

821 

822 parent_task_name = job_to_task[parent_job_name] 

823 deps.append( 

824 { 

825 "task": parent_task_name, 

826 "inputname": inputname, 

827 } 

828 ) 

829 if parent_task_name not in work.dependency_tasks: 

830 work.dependency_tasks.append(parent_task_name) 

831 

832 work.dependency_map.append( 

833 { 

834 "name": f"{job_label}:orderIdMap_{order_id}" if enable_job_name_map else job_name, 

835 "order_id": order_id, 

836 "dependencies": deps, 

837 } 

838 ) 

839 

840 _LOG.info("Successfully recovered.") 

841 

842 for task_name in name_works: 

843 work = name_works[task_name] 

844 # trigger the setter function which will validate the dependency_map: 

845 # 1) check the name length to avoid the the name too long, 

846 # 2) check to avoid duplicated items. 

847 sorted_dep_map = sorted(work.dependency_map, key=lambda x: x["order_id"]) 

848 work.dependency_map = sorted_dep_map 

849 

850 if enable_job_name_map: 

851 with open(order_id_map_file, "w") as f: 

852 json.dump(order_id_map, f) 

853 

854 return files_to_pre_stage, dag_sink_work, task_count 

855 

856 

857def create_archive_file(submit_path, archive_filename, files): 

858 if not archive_filename.startswith("/"): 

859 archive_filename = os.path.join(submit_path, archive_filename) 

860 

861 with tarfile.open(archive_filename, "w:gz", dereference=True) as tar: 

862 for local_file in files: 

863 base_name = os.path.basename(local_file) 

864 tar.add(local_file, arcname=os.path.basename(base_name)) 

865 return archive_filename 

866 

867 

868def copy_files_to_pandacache(filename): 

869 from pandaclient import Client 

870 

871 attempt = 0 

872 max_attempts = 3 

873 done = False 

874 while attempt < max_attempts and not done: 

875 status, out = Client.putFile(filename, True) 

876 if status == 0: 

877 done = True 

878 print(f"copy_files_to_pandacache: status: {status}, out: {out}") 

879 if out.startswith("NewFileName:"): 

880 # found the same input sandbox to reuse 

881 filename = out.split(":")[-1] 

882 elif out != "True": 

883 print(out) 

884 return None 

885 

886 filename = os.path.basename(filename) 

887 cache_path = os.path.join(os.environ["PANDACACHE_URL"], "cache") 

888 filename = os.path.join(cache_path, filename) 

889 return filename 

890 

891 

892def download_extract_archive(filename, prefix=None): 

893 """Download and extract the tarball from pandacache. 

894 

895 Parameters 

896 ---------- 

897 filename : `str` 

898 The filename to download. 

899 prefix : `str`, optional 

900 The target directory the tarball will be downloaded and extracted to. 

901 If None (default), the current directory will be used. 

902 """ 

903 archive_basename = os.path.basename(filename) 

904 target_dir = prefix if prefix is not None else os.getcwd() 

905 full_output_filename = os.path.join(target_dir, archive_basename) 

906 

907 if filename.startswith("https:"): 

908 panda_cache_url = os.path.dirname(os.path.dirname(filename)) 

909 os.environ["PANDACACHE_URL"] = panda_cache_url 

910 elif "PANDACACHE_URL" not in os.environ and "PANDA_URL_SSL" in os.environ: 

911 os.environ["PANDACACHE_URL"] = os.environ["PANDA_URL_SSL"] 

912 panda_cache_url = os.environ.get("PANDACACHE_URL", None) 

913 print(f"PANDACACHE_URL: {panda_cache_url}") 

914 

915 # The import of PanDA client must happen *after* the PANDACACHE_URL is set. 

916 # Otherwise, the PanDA client the environment setting will not be parsed. 

917 from pandaclient import Client 

918 

919 attempt = 0 

920 max_attempts = 3 

921 while attempt < max_attempts: 

922 status, output = Client.getFile(archive_basename, output_path=full_output_filename) 

923 if status == 0: 

924 break 

925 if attempt <= 1: 

926 secs = random.randint(1, 10) 

927 elif attempt <= 2: 

928 secs = random.randint(1, 60) 

929 else: 

930 secs = random.randint(1, 120) 

931 time.sleep(secs) 

932 print(f"Download archive file from pandacache status: {status}, output: {output}") 

933 if status != 0: 

934 raise RuntimeError("Failed to download archive file from pandacache") 

935 with tarfile.open(full_output_filename, "r:gz") as f: 

936 f.extractall(target_dir) 

937 print(f"Extracted {full_output_filename} to {target_dir}") 

938 os.remove(full_output_filename) 

939 print(f"Removed {full_output_filename}") 

940 

941 

942def get_task_parameter(config, remote_build, key): 

943 search_opt = {"replaceVars": True, "expandEnvVars": False, "replaceEnvVars": False, "required": False} 

944 _, value = remote_build.search(key, search_opt) 

945 if not value: 

946 _, value = config.search(key, search_opt) 

947 return value 

948 

949 

950def create_idds_build_workflow(**kwargs): 

951 config = kwargs["config"] if "config" in kwargs else None 

952 remote_build = kwargs["remote_build"] if "remote_build" in kwargs else None 

953 config_file = kwargs["config_file"] if "config_file" in kwargs else None 

954 config_file_base = os.path.basename(config_file) if config_file else None 

955 compute_site = kwargs["compute_site"] if "compute_site" in kwargs else None 

956 _, files = remote_build.search("files", opt={"default": []}) 

957 submit_path = config["submitPath"] 

958 files.append(config_file) 

959 archive_filename = f"jobO.{uuid.uuid4()}.tar.gz" 

960 archive_filename = create_archive_file(submit_path, archive_filename, files) 

961 _LOG.info(f"archive file name: {archive_filename}") 

962 remote_filename = copy_files_to_pandacache(archive_filename) 

963 _LOG.info(f"pandacache file: {remote_filename}") 

964 

965 _LOG.info(type(remote_build)) 

966 search_opt = {"replaceVars": True, "expandEnvVars": False, "replaceEnvVars": False, "required": False} 

967 cvals = {"LSST_VERSION": get_task_parameter(config, remote_build, "LSST_VERSION")} 

968 cvals["custom_lsst_setup"] = get_task_parameter(config, remote_build, "custom_lsst_setup") 

969 max_name_length = PANDA_DEFAULT_NAME_LENGTH 

970 if "IDDS_MAX_NAME_LENGTH" in os.environ: 

971 max_name_length = int(os.environ["IDDS_MAX_NAME_LENGTH"]) 

972 cvals["IDDS_MAX_NAME_LENGTH"] = max_name_length 

973 search_opt["curvals"] = cvals 

974 _, executable = remote_build.search("runnerCommand", opt=search_opt) 

975 executable = executable.replace("_download_cmd_line_", remote_filename) 

976 executable = executable.replace("_build_cmd_line_", config_file_base) 

977 executable = executable.replace("_compute_site_", compute_site or "") 

978 

979 task_cloud = get_task_parameter(config, remote_build, "computeCloud") 

980 task_site = get_task_parameter(config, remote_build, "computeSite") 

981 task_queue = get_task_parameter(config, remote_build, "queue") 

982 task_rss = get_task_parameter(config, remote_build, "requestMemory") 

983 nretries = get_task_parameter(config, remote_build, "numberOfRetries") 

984 processing_type = get_task_parameter(config, remote_build, "processingType") 

985 _LOG.info("requestMemory: %s", task_rss) 

986 _LOG.info("Site: %s", task_site) 

987 # _LOG.info("executable: %s", executable) 

988 # TODO: fill other parameters based on config 

989 build_work = DomaPanDAWork( 

990 executable=executable, 

991 task_type="lsst_build", 

992 primary_input_collection={"scope": "pseudo_dataset", "name": "pseudo_input_collection#1"}, 

993 output_collections=[{"scope": "pseudo_dataset", "name": "pseudo_output_collection#1"}], 

994 log_collections=[], 

995 dependency_map=None, 

996 task_name="build_task", 

997 task_queue=task_queue, 

998 encode_command_line=True, 

999 prodSourceLabel="managed", 

1000 processing_type=processing_type, 

1001 task_log={ 

1002 "dataset": "PandaJob_#{pandaid}/", 

1003 "destination": "local", 

1004 "param_type": "log", 

1005 "token": "local", 

1006 "type": "template", 

1007 "value": "log.tgz", 

1008 }, 

1009 task_rss=task_rss if task_rss else PANDA_DEFAULT_RSS, 

1010 task_cloud=task_cloud, 

1011 task_site=task_site, 

1012 maxattempt=nretries if nretries > 0 else PANDA_DEFAULT_MAX_ATTEMPTS, 

1013 ) 

1014 

1015 workflow = IDDS_client_workflow() 

1016 

1017 workflow.add_work(build_work) 

1018 workflow.name = config["bps_defined"]["uniqProcName"] 

1019 return workflow