Coverage for python/lsst/ctrl/bps/panda/panda_service.py: 9%

228 statements  

« prev     ^ index     » next       coverage.py v7.11.0, created at 2025-11-04 01:58 -0800

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"""Interface between generic workflow to PanDA/iDDS workflow system.""" 

28 

29__all__ = ["PanDAService", "PandaBpsWmsWorkflow"] 

30 

31 

32import json 

33import logging 

34import os 

35import pickle 

36import re 

37 

38from idds.workflowv2.workflow import Workflow as IDDS_client_workflow 

39 

40from lsst.ctrl.bps import ( 

41 DEFAULT_MEM_FMT, 

42 DEFAULT_MEM_UNIT, 

43 BaseWmsService, 

44 BaseWmsWorkflow, 

45 WmsRunReport, 

46 WmsStates, 

47) 

48from lsst.ctrl.bps.panda.constants import PANDA_DEFAULT_MAX_COPY_WORKERS 

49from lsst.ctrl.bps.panda.utils import ( 

50 add_final_idds_work, 

51 add_idds_work, 

52 aggregate_by_basename, 

53 copy_files_for_distribution, 

54 create_idds_build_workflow, 

55 extract_taskname, 

56 get_idds_client, 

57 get_idds_result, 

58 idds_call_with_check, 

59) 

60from lsst.resources import ResourcePath 

61from lsst.utils.timer import time_this 

62 

63_LOG = logging.getLogger(__name__) 

64 

65 

66class PanDAService(BaseWmsService): 

67 """PanDA version of WMS service.""" 

68 

69 def prepare(self, config, generic_workflow, out_prefix=None): 

70 # Docstring inherited from BaseWmsService.prepare. 

71 _LOG.debug("out_prefix = '%s'", out_prefix) 

72 

73 _LOG.info("Starting PanDA prepare stage (creating specific implementation of workflow)") 

74 

75 with time_this( 

76 log=_LOG, 

77 level=logging.INFO, 

78 prefix=None, 

79 msg="PanDA prepare stage completed", 

80 mem_usage=True, 

81 mem_unit=DEFAULT_MEM_UNIT, 

82 mem_fmt=DEFAULT_MEM_FMT, 

83 ): 

84 workflow = PandaBpsWmsWorkflow.from_generic_workflow( 

85 config, generic_workflow, out_prefix, f"{self.__class__.__module__}.{self.__class__.__name__}" 

86 ) 

87 workflow.write(out_prefix) 

88 return workflow 

89 

90 def submit(self, workflow, **kwargs): 

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

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

93 

94 if config and remote_build: 

95 _LOG.info("remote build") 

96 

97 idds_build_workflow = create_idds_build_workflow(**kwargs) 

98 idds_client = get_idds_client(self.config) 

99 ret = idds_client.submit_build(idds_build_workflow, username=None, use_dataset_name=False) 

100 _LOG.debug("iDDS client manager submit returned = %s", ret) 

101 

102 # Check submission success 

103 status, result, error = get_idds_result(ret) 

104 if status: 

105 request_id = int(result) 

106 else: 

107 raise RuntimeError(f"Error submitting to PanDA service: {error}") 

108 

109 _LOG.info("Submitted into iDDs with request id=%s", request_id) 

110 idds_build_workflow.run_id = request_id 

111 return idds_build_workflow 

112 

113 else: 

114 _, max_copy_workers = self.config.search( 

115 "maxCopyWorkers", opt={"default": PANDA_DEFAULT_MAX_COPY_WORKERS} 

116 ) 

117 file_distribution_uri = self.config["fileDistributionEndPoint"] 

118 lsst_temp = "LSST_RUN_TEMP_SPACE" 

119 if lsst_temp in file_distribution_uri and lsst_temp not in os.environ: 

120 file_distribution_uri = self.config["fileDistributionEndPointDefault"] 

121 protocol_pattern = re.compile(r"^[a-zA-Z][a-zA-Z\d+\-.]*://") 

122 if not protocol_pattern.match(file_distribution_uri): 

123 file_distribution_uri = "file://" + file_distribution_uri 

124 

125 submit_cmd = workflow.run_attrs.get("bps_iscustom", False) 

126 if not submit_cmd: 

127 copy_files_for_distribution( 

128 workflow.files_to_pre_stage, 

129 ResourcePath(file_distribution_uri, forceDirectory=True), 

130 max_copy_workers, 

131 ) 

132 

133 idds_client = get_idds_client(self.config) 

134 ret = idds_client.submit(workflow.idds_client_workflow, username=None, use_dataset_name=False) 

135 _LOG.debug("iDDS client manager submit returned = %s", ret) 

136 

137 # Check submission success 

138 status, result, error = get_idds_result(ret) 

139 if status: 

140 request_id = int(result) 

141 else: 

142 raise RuntimeError(f"Error submitting to PanDA service: {error}") 

143 

144 _LOG.info("Submitted into iDDs with request id=%s", request_id) 

145 workflow.run_id = request_id 

146 

147 def restart(self, wms_workflow_id): 

148 # Docstring inherited from BaseWmsService.restart. 

149 idds_client = get_idds_client(self.config) 

150 ret = idds_client.retry(request_id=wms_workflow_id) 

151 _LOG.debug("Restart PanDA workflow returned = %s", ret) 

152 

153 status, result, error = get_idds_result(ret) 

154 if status: 

155 _LOG.info("Restarting PanDA workflow %s", result) 

156 return wms_workflow_id, None, json.dumps(result) 

157 

158 return None, None, f"Error retry PanDA workflow: {error}" 

159 

160 def report( 

161 self, 

162 wms_workflow_id=None, 

163 user=None, 

164 hist=0, 

165 pass_thru=None, 

166 is_global=False, 

167 return_exit_codes=False, 

168 ): 

169 # Docstring inherited from BaseWmsService.report. 

170 message = "" 

171 run_reports = [] 

172 

173 if not wms_workflow_id: 

174 message = "Run summary not implemented yet, use 'bps report --id <workflow_id>' instead" 

175 return run_reports, message 

176 

177 idds_client = get_idds_client(self.config) 

178 ret = idds_call_with_check( 

179 idds_client.get_requests, 

180 func_name="get workflow status", 

181 request_id=wms_workflow_id, 

182 with_detail=True, 

183 ) 

184 

185 tasks = ret[1][1] 

186 if not tasks: 

187 message = f"No records found for workflow id '{wms_workflow_id}'. Hint: double check the id" 

188 return run_reports, message 

189 

190 # Create initial WmsRunReport 

191 head = tasks[0] 

192 wms_report = WmsRunReport( 

193 wms_id=str(head["request_id"]), 

194 operator=head["username"], 

195 project="", 

196 campaign="", 

197 payload="", 

198 run=head["name"], 

199 state=WmsStates.UNKNOWN, 

200 total_number_jobs=0, 

201 job_state_counts=dict.fromkeys(WmsStates, 0), 

202 job_summary={}, 

203 run_summary="", 

204 exit_code_summary={}, 

205 ) 

206 

207 # Define workflow status mapping 

208 workflow_status = head["status"]["attributes"]["_name_"] 

209 if workflow_status in ("Finished", "SubFinished"): 

210 wms_report.state = WmsStates.SUCCEEDED 

211 elif workflow_status in ("Failed", "Expired"): 

212 wms_report.state = WmsStates.FAILED 

213 elif workflow_status == "Cancelled": 

214 wms_report.state = WmsStates.DELETED 

215 elif workflow_status == "Suspended": 

216 wms_report.state = WmsStates.HELD 

217 else: 

218 wms_report.state = WmsStates.RUNNING 

219 

220 # Define state mapping for job aggregation 

221 # The status of a task is taken from the first item of state_map. 

222 # The workflow is in status WmsStates.FAILED when: 

223 # All tasks have failed. 

224 # SubFinished tasks has jobs in 

225 # output_processed_files: Finished 

226 # output_failed_files: Failed 

227 # output_missing_files: Missing 

228 state_map = { 

229 "Finished": [WmsStates.SUCCEEDED], 

230 "SubFinished": [WmsStates.SUCCEEDED, WmsStates.FAILED, WmsStates.PRUNED], 

231 "Transforming": [ 

232 WmsStates.RUNNING, 

233 WmsStates.SUCCEEDED, 

234 WmsStates.FAILED, 

235 # WmsStates.READY, 

236 WmsStates.UNREADY, 

237 WmsStates.PRUNED, 

238 ], 

239 "Failed": [WmsStates.FAILED, WmsStates.PRUNED], 

240 } 

241 

242 file_map = { 

243 WmsStates.SUCCEEDED: "output_processed_files", 

244 WmsStates.RUNNING: "output_processing_files", 

245 WmsStates.FAILED: "output_failed_files", 

246 # WmsStates.READY: "output_activated_files", 

247 WmsStates.UNREADY: "input_new_files", 

248 WmsStates.PRUNED: "output_missing_files", 

249 } 

250 

251 # Sort tasks by workload_id or fallback 

252 try: 

253 tasks.sort(key=lambda x: x["transform_workload_id"]) 

254 except (KeyError, TypeError): 

255 tasks.sort(key=lambda x: x["transform_id"]) 

256 

257 exit_codes_all = {} 

258 

259 # --- Process each task sequentially --- 

260 for task in tasks: 

261 if task.get("transform_id") is None: 

262 # Not created task (It happens because of an outer join 

263 # between requests table and transforms table). 

264 continue 

265 

266 task_name = task.get("transform_name", "") 

267 tasklabel = extract_taskname(task_name) 

268 status = task["transform_status"]["attributes"]["_name_"] 

269 totaljobs = task.get("output_total_files", 0) 

270 wms_report.total_number_jobs += totaljobs 

271 

272 # --- If task failed/subfinished, fetch exit codes --- 

273 if status in ("SubFinished", "Failed") and not task_name.startswith("build_task"): 

274 transform_workload_id = task.get("transform_workload_id") 

275 if transform_workload_id: 

276 # When there are failed jobs, ctrl_bps check 

277 # the number of exit codes 

278 nfailed = task.get("output_failed_files", 0) 

279 exit_codes_all[tasklabel] = [1] * nfailed 

280 if return_exit_codes: 

281 new_ret = idds_call_with_check( 

282 idds_client.get_contents_output_ext, 

283 func_name=f"get task {transform_workload_id} detail", 

284 request_id=wms_workflow_id, 

285 workload_id=transform_workload_id, 

286 ) 

287 # task_info is a dictionary of len 1 that contains 

288 # a list of dicts containing panda job info 

289 task_info = new_ret[1][1] 

290 if len(task_info) == 1: 

291 _, wmsjobs = next(iter(task_info.items())) 

292 exit_codes_all[tasklabel] = [ 

293 j["trans_exit_code"] 

294 for j in wmsjobs 

295 if j.get("trans_exit_code") not in (None, 0, "0") 

296 ] 

297 if nfailed > 0 and len(exit_codes_all[tasklabel]) == 0: 

298 _LOG.debug( 

299 f"No exit codes in iDDS task info for workload {transform_workload_id}" 

300 ) 

301 else: 

302 raise RuntimeError( 

303 f"Unexpected iDDS task info for workload {transform_workload_id}: {task_info}" 

304 ) 

305 

306 # --- Aggregate job states --- 

307 taskstatus = {} 

308 mapped_states = state_map.get(status, []) 

309 for state in WmsStates: 

310 njobs = 0 

311 if state in mapped_states and state in file_map: 

312 val = task.get(file_map[state]) 

313 if val: 

314 njobs = val 

315 if state == WmsStates.RUNNING: 

316 njobs += task.get("output_new_files", 0) - task.get("input_new_files", 0) 

317 if state != WmsStates.UNREADY: 

318 wms_report.job_state_counts[state] += njobs 

319 taskstatus[state] = njobs 

320 

321 # Count UNREADY 

322 unready = WmsStates.UNREADY 

323 taskstatus[unready] = totaljobs - sum( 

324 taskstatus[state] for state in WmsStates if state != unready 

325 ) 

326 wms_report.job_state_counts[unready] += taskstatus[unready] 

327 

328 # Store task summary 

329 wms_report.job_summary[tasklabel] = taskstatus 

330 summary_part = f"{tasklabel}:{totaljobs}" 

331 if wms_report.run_summary: 

332 summary_part = f";{summary_part}" 

333 wms_report.run_summary += summary_part 

334 

335 # Store all exit codes 

336 wms_report.exit_code_summary = exit_codes_all 

337 

338 ( 

339 wms_report.job_summary, 

340 wms_report.exit_code_summary, 

341 wms_report.run_summary, 

342 ) = aggregate_by_basename( 

343 wms_report.job_summary, 

344 wms_report.exit_code_summary, 

345 wms_report.run_summary, 

346 ) 

347 

348 run_reports.append(wms_report) 

349 return run_reports, message 

350 

351 def list_submitted_jobs(self, wms_id=None, user=None, require_bps=True, pass_thru=None, is_global=False): 

352 # Docstring inherited from BaseWmsService.list_submitted_jobs. 

353 if wms_id is None and user is not None: 

354 raise RuntimeError( 

355 "Error to get workflow status report: wms_id is required" 

356 " and filtering workflows with 'user' is not supported." 

357 ) 

358 

359 idds_client = get_idds_client(self.config) 

360 ret = idds_client.get_requests(request_id=wms_id) 

361 _LOG.debug("PanDA get workflows returned = %s", ret) 

362 

363 status, result, error = get_idds_result(ret) 

364 if status: 

365 req_ids = [req["request_id"] for req in result] 

366 return req_ids 

367 

368 raise RuntimeError(f"Error list PanDA workflow requests: {error}") 

369 

370 def cancel(self, wms_id, pass_thru=None): 

371 # Docstring inherited from BaseWmsService.cancel. 

372 idds_client = get_idds_client(self.config) 

373 ret = idds_client.abort(request_id=wms_id) 

374 _LOG.debug("Abort PanDA workflow returned = %s", ret) 

375 

376 status, result, error = get_idds_result(ret) 

377 if status: 

378 _LOG.info("Aborting PanDA workflow %s", result) 

379 return True, json.dumps(result) 

380 

381 return False, f"Error abort PanDA workflow: {error}" 

382 

383 def ping(self, pass_thru=None): 

384 # Docstring inherited from BaseWmsService.ping. 

385 idds_client = get_idds_client(self.config) 

386 ret = idds_client.ping() 

387 _LOG.debug("Ping PanDA service returned = %s", ret) 

388 

389 status, result, error = get_idds_result(ret) 

390 if status: 

391 if "Status" in result and result["Status"] == "OK": 

392 return 0, None 

393 

394 return -1, f"Error ping PanDA service: {result}" 

395 

396 return -1, f"Error ping PanDA service: {error}" 

397 

398 def run_submission_checks(self): 

399 # Docstring inherited from BaseWmsService.run_submission_checks. 

400 for key in ["PANDA_URL"]: 

401 if key not in os.environ: 

402 raise OSError(f"Missing environment variable {key}") 

403 

404 status, message = self.ping() 

405 if status != 0: 

406 raise RuntimeError(message) 

407 

408 def get_status( 

409 self, 

410 wms_workflow_id=None, 

411 hist=0, 

412 is_global=False, 

413 ): 

414 # Docstring inherited from BaseWmsService.get_status. 

415 

416 idds_client = get_idds_client(self.config) 

417 ret = idds_client.get_requests(request_id=wms_workflow_id, with_detail=False) 

418 _LOG.debug("PanDA get workflow status returned = %s", str(ret)) 

419 

420 request_status = ret[0] 

421 if request_status != 0: 

422 state = WmsStates.UNKNOWN 

423 message = f"Error getting workflow status for id {wms_workflow_id}: ret = {ret}" 

424 else: 

425 tasks = ret[1][1] 

426 if not tasks: 

427 state = WmsStates.UNKNOWN 

428 message = f"No records found for workflow id '{wms_workflow_id}'. Hint: double check the id" 

429 elif not isinstance(tasks[0], dict): 

430 state = WmsStates.UNKNOWN 

431 message = f"Error getting workflow status for id {wms_workflow_id}: ret = {ret}" 

432 else: 

433 message = "" 

434 head = tasks[0] 

435 workflow_status = head["status"]["attributes"]["_name_"] 

436 if workflow_status in ["Finished"]: 

437 state = WmsStates.SUCCEEDED 

438 elif workflow_status in ["Failed", "Expired", "SubFinished"]: 

439 state = WmsStates.FAILED 

440 elif workflow_status in ["Cancelled"]: 

441 state = WmsStates.DELETED 

442 elif workflow_status in ["Suspended"]: 

443 state = WmsStates.HELD 

444 else: 

445 state = WmsStates.RUNNING 

446 

447 return state, message 

448 

449 

450class PandaBpsWmsWorkflow(BaseWmsWorkflow): 

451 """A single Panda based workflow. 

452 

453 Parameters 

454 ---------- 

455 name : `str` 

456 Unique name for Workflow. 

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

458 BPS configuration that includes necessary submit/runtime information. 

459 """ 

460 

461 def __init__(self, name, config=None): 

462 super().__init__(name, config) 

463 self.files_to_pre_stage = {} # src, dest 

464 self.idds_client_workflow = IDDS_client_workflow(name=name) 

465 self.run_attrs = {} 

466 

467 @classmethod 

468 def from_generic_workflow(cls, config, generic_workflow, out_prefix, service_class): 

469 # Docstring inherited from BaseWmsWorkflow.from_generic_workflow. 

470 wms_workflow = cls(generic_workflow.name, config) 

471 

472 if generic_workflow.run_attrs: 

473 wms_workflow.run_attrs.update(generic_workflow.run_attrs) 

474 

475 files, dag_sink_work, task_count = add_idds_work( 

476 config, generic_workflow, wms_workflow.idds_client_workflow 

477 ) 

478 wms_workflow.files_to_pre_stage.update(files) 

479 

480 files = add_final_idds_work( 

481 config, generic_workflow, wms_workflow.idds_client_workflow, dag_sink_work, task_count + 1, 1 

482 ) 

483 wms_workflow.files_to_pre_stage.update(files) 

484 

485 return wms_workflow 

486 

487 def write(self, out_prefix): 

488 # Docstring inherited from BaseWmsWorkflow.write. 

489 with open(os.path.join(out_prefix, "panda_workflow.pickle"), "wb") as fh: 

490 pickle.dump(self, fh)