Warning, file /AID2E-framework/tests/test_schedulers/test_pandaidds_runner.py was not indexed
or was modified since last indexation (in which case cross-reference links may be missing, inaccurate or erroneous).
0001 """Mocked PanDA/iDDS runner contract tests."""
0002
0003 import sys
0004 import threading
0005 import types
0006
0007 import pytest
0008
0009 from aid2e.schedulers.PanDAiDDS.config import PanDAiDDSRunnerConfig
0010 from aid2e.schedulers.PanDAiDDS.multistage_graph import PanDAMultiStageGraphBuilder
0011 from aid2e.schedulers.PanDAiDDS.multistage import PanDAMultiStageJob, PanDAMultiStageSpec
0012 from aid2e.utilities.runtime_builders import run_trial_workflow
0013 from aid2e.schedulers.PanDAiDDS.runner import (
0014 PanDAiDDSScheduler,
0015 _normalize_all2one_input_datasets,
0016 _remote_python_callable_entrypoint,
0017 )
0018
0019
0020 def panda_test_evaluator(context=None, **kwargs):
0021 return {"unused": True}
0022
0023
0024 def panda_adapter_test_evaluator(context=None, **kwargs):
0025 context.add_log("remote adapter ran")
0026 context.xcom_push("seen", kwargs.get("label"))
0027 return {
0028 "job_id": context.job_id,
0029 "design": context.design_point,
0030 "label": kwargs.get("label"),
0031 "input_file_names": kwargs.get("input_file_names"),
0032 "xcom": context.xcom,
0033 }
0034
0035
0036 def panda_prepare_evaluator(context=None, **kwargs):
0037 return {"prepared": kwargs.get("scale", 1)}
0038
0039
0040 def panda_evaluate_from_parent(context=None, **kwargs):
0041 parent = kwargs["parent_results"]
0042 return {"f1": parent["prepared"] + 1, "f2": parent["prepared"] + 2}
0043
0044
0045 def panda_chunk_evaluator(context=None, **kwargs):
0046 return {"chunk": kwargs.get("child_key"), "value": kwargs.get("value", 0)}
0047
0048
0049 def panda_analyze_chunks(context=None, **kwargs):
0050 parent = kwargs["parent_results"]
0051 return {"total": sum(item["value"] for item in parent.values())}
0052
0053
0054 def panda_ana_one2one(context=None, **kwargs):
0055 parent = kwargs["simreco_result"]
0056 return {"analyzed": parent["value"] + 1}
0057
0058
0059 def panda_local_final_from_ana(context=None, **kwargs):
0060 ana_result = kwargs["ana_result"]
0061 return {"objective": ana_result.get("total", ana_result.get("analyzed")) + 1}
0062
0063
0064 def panda_strict_local_stage(context=None):
0065 return {"objective": 1}
0066
0067
0068 def panda_no_context_evaluator(value):
0069 return {"value": value}
0070
0071
0072 class DummyContext:
0073 def __init__(self):
0074 self.pushed = {}
0075
0076 def xcom_push(self, key, value):
0077 self.pushed[key] = value
0078
0079
0080 def install_fake_idds(
0081 monkeypatch,
0082 *,
0083 status="finished",
0084 result=None,
0085 details=None,
0086 result_sequence=None,
0087 details_sequence=None,
0088 status_sequence=None,
0089 ):
0090 state = {
0091 "work_kwargs_history": [],
0092 "work_params_history": [],
0093 "result_lookup_history": [],
0094 "work_history": [],
0095 }
0096 result = {"f1": 1.2, "f2": 0.4} if result is None else result
0097 details = {"attempt": 1} if details is None else details
0098
0099 idds_module = types.ModuleType("idds")
0100 iworkflow_module = types.ModuleType("idds.iworkflow")
0101 workflow_module = types.ModuleType("idds.iworkflow.workflow")
0102 work_module = types.ModuleType("idds.iworkflow.work")
0103
0104 class FakeWorkflow:
0105 def pre_run(self):
0106 state["workflow_pre_run"] = True
0107
0108 def prepare(self):
0109 state["workflow_prepare"] = True
0110
0111 def submit(self):
0112 state["workflow_submitted"] = True
0113 return "req-1"
0114
0115 class FakeResult:
0116 def __init__(self, index):
0117 self.index = index
0118
0119 def get_result(self, name, key, verbose=False, with_details=False):
0120 lookup = {
0121 "name": name,
0122 "key": key,
0123 "verbose": verbose,
0124 "with_details": with_details,
0125 }
0126 state["result_lookup"] = lookup
0127 state["result_lookup_history"].append(lookup)
0128 selected_result = (
0129 result_sequence[self.index]
0130 if result_sequence is not None
0131 else result
0132 )
0133 selected_details = (
0134 details_sequence[self.index]
0135 if details_sequence is not None
0136 else details
0137 )
0138 return selected_result, selected_details
0139
0140 class FakeWork:
0141 def __init__(self, name, index):
0142 self.name = name
0143 self.index = index
0144 self.internal_id = f"internal-{index + 1}"
0145 self.core_count = None
0146 self.init_async_result_called = False
0147 self.cancelled = False
0148
0149 def submit(self):
0150 state["work_submitted"] = True
0151 return f"tf-{self.index + 1}"
0152
0153 def init_async_result(self):
0154 self.init_async_result_called = True
0155
0156 def get_status(self):
0157 if status_sequence is not None:
0158 return status_sequence[self.index]
0159 return status
0160
0161 def is_finished(self, value):
0162 return value == "finished"
0163
0164 def is_failed(self, value):
0165 return value == "failed"
0166
0167 def is_terminated(self):
0168 return self.cancelled
0169
0170 def cancel(self):
0171 self.cancelled = True
0172
0173 def get_results(self):
0174 return FakeResult(self.index)
0175
0176 def workflow_def(**kwargs):
0177 state["workflow_kwargs"] = kwargs
0178
0179 def builder():
0180 return FakeWorkflow()
0181
0182 return builder
0183
0184 def work_def(**kwargs):
0185 state["work_kwargs"] = kwargs
0186 state["work_kwargs_history"].append(kwargs)
0187 index = len(state["work_kwargs_history"]) - 1
0188
0189 def builder(**params):
0190 state["work_params"] = params
0191 state["work_params_history"].append(params)
0192 work = FakeWork(kwargs["name"], index)
0193 state["work"] = work
0194 state["work_history"].append(work)
0195 return work
0196
0197 return builder
0198
0199 workflow_module.workflow = workflow_def
0200 work_module.work = work_def
0201
0202 monkeypatch.setitem(sys.modules, "idds", idds_module)
0203 monkeypatch.setitem(sys.modules, "idds.iworkflow", iworkflow_module)
0204 monkeypatch.setitem(sys.modules, "idds.iworkflow.workflow", workflow_module)
0205 monkeypatch.setitem(sys.modules, "idds.iworkflow.work", work_module)
0206
0207 return state
0208
0209
0210 def test_run_stage_preserves_job_id_and_returns_idds_outputs(monkeypatch, tmp_path):
0211 state = install_fake_idds(monkeypatch, result={"f1": 1.2, "f2": 0.4})
0212 context = DummyContext()
0213 scheduler = PanDAiDDSScheduler(
0214 PanDAiDDSRunnerConfig(
0215 name="user.test.panda",
0216 cloud="US",
0217 queue="BNL",
0218 init_env="setup_panda.sh",
0219 source_dir=str(tmp_path),
0220 )
0221 )
0222
0223 result = scheduler.run_stage(
0224 "evaluate",
0225 [
0226 {
0227 "job_id": "evaluate:dtlz2:0001",
0228 "name": "dtlz2",
0229 "function": panda_test_evaluator,
0230 "params": {"context": context, "x": [0.5, 0.5, 0.5]},
0231 "payload": {"label": "dtlz2"},
0232 "job_context": context,
0233 }
0234 ],
0235 parallelism_policy={"poll_interval": 0},
0236 working_dir=str(tmp_path),
0237 )
0238
0239 assert result.success is True
0240 assert result.artifacts == {"objectives": {"f1": 1.2, "f2": 0.4}}
0241 assert len(result.job_statuses) == 1
0242
0243 status = result.job_statuses[0]
0244 assert status.job_id == "evaluate:dtlz2:0001"
0245 assert status.status == "completed"
0246 assert status.return_code == 0
0247 assert status.outputs == {"objectives": {"f1": 1.2, "f2": 0.4}}
0248 assert status.metrics == {"transform_id": "tf-1"}
0249
0250 assert context.pushed["objectives"] == {"f1": 1.2, "f2": 0.4}
0251 assert context.pushed["results_details"] == {"attempt": 1}
0252 assert state["work_kwargs"]["func"] is _remote_python_callable_entrypoint
0253 assert state["work_params"]["callable_ref"] == (
0254 "tests.test_schedulers.test_pandaidds_runner:panda_test_evaluator"
0255 )
0256 assert state["work_params"]["op_kwargs"] == {"x": [0.5, 0.5, 0.5]}
0257 assert state["work_params"]["with_context"] is True
0258 assert state["work_params"]["context_payload"]["job_id"] == "evaluate:dtlz2:0001"
0259 assert state["work_params"]["context_payload"]["payload"] == {"label": "dtlz2"}
0260 assert state["work_kwargs"]["name"] == state["work_kwargs"]["job_key"]
0261 assert state["work_kwargs"]["name"].startswith("user.test.a2e.")
0262 assert len(state["work_kwargs"]["name"]) <= 64
0263 assert state["work_kwargs"]["name"].split(".")[-2] == "000001"
0264 assert state["work_kwargs"]["log_dataset_name"] == f'{state["work_kwargs"]["name"]}.log/'
0265
0266 cached = scheduler.check_status("evaluate:dtlz2:0001")
0267 assert cached.status == "completed"
0268 assert cached.outputs == {"objectives": {"f1": 1.2, "f2": 0.4}}
0269
0270
0271 def test_run_stage_marks_failed_idds_work_failed(monkeypatch, tmp_path):
0272 state = install_fake_idds(monkeypatch, status="failed")
0273 scheduler = PanDAiDDSScheduler(
0274 PanDAiDDSRunnerConfig(name="user.test.panda", source_dir=str(tmp_path))
0275 )
0276
0277 result = scheduler.run_stage(
0278 "evaluate",
0279 [
0280 {
0281 "job_id": "evaluate:failed:0001",
0282 "name": "failed_eval",
0283 "function": panda_test_evaluator,
0284 "params": {},
0285 }
0286 ],
0287 parallelism_policy={"poll_interval": 0},
0288 working_dir=str(tmp_path),
0289 )
0290
0291 assert result.success is False
0292 assert result.error_message == "Some jobs failed in stage 'evaluate'"
0293 assert result.artifacts == {}
0294
0295 status = result.job_statuses[0]
0296 assert status.job_id == "evaluate:failed:0001"
0297 assert status.status == "failed"
0298 assert status.return_code == -1
0299 assert "failed" in status.stderr
0300 assert status.outputs == {}
0301 assert state["work_params"]["with_context"] is False
0302
0303 cached = scheduler.check_status("evaluate:failed:0001")
0304 assert cached.status == "failed"
0305 assert cached.return_code == -1
0306
0307
0308
0309 def test_repeated_logical_job_submissions_get_distinct_panda_work_names(monkeypatch, tmp_path):
0310 state = install_fake_idds(monkeypatch, status="running")
0311 scheduler = PanDAiDDSScheduler(
0312 PanDAiDDSRunnerConfig(name="user.test.panda", source_dir=str(tmp_path))
0313 )
0314 job = {
0315 "job_id": "dtlz2",
0316 "name": "dtlz2",
0317 "function": panda_test_evaluator,
0318 "params": {},
0319 }
0320
0321 scheduler.submit_job("evaluate", job, str(tmp_path))
0322 scheduler.submit_job("evaluate", job, str(tmp_path))
0323 scheduler.submit_job("evaluate", job, str(tmp_path))
0324
0325 names = [kwargs["name"] for kwargs in state["work_kwargs_history"]]
0326 assert len(names) == 3
0327 assert len(set(names)) == 3
0328 assert all(name.startswith("user.test.a2e.") for name in names)
0329 assert all(len(name) <= 64 for name in names)
0330 assert names[0].split(".")[-2] == "000001"
0331 assert names[1].split(".")[-2] == "000002"
0332 assert names[2].split(".")[-2] == "000003"
0333 assert [kwargs["job_key"] for kwargs in state["work_kwargs_history"]] == names
0334 assert [kwargs["log_dataset_name"] for kwargs in state["work_kwargs_history"]] == [
0335 f"{names[0]}.log/",
0336 f"{names[1]}.log/",
0337 f"{names[2]}.log/",
0338 ]
0339
0340
0341 def test_distinct_scheduler_instances_do_not_reuse_panda_work_names(tmp_path):
0342 first = PanDAiDDSScheduler(
0343 PanDAiDDSRunnerConfig(name="user.test.panda", source_dir=str(tmp_path))
0344 )
0345 second = PanDAiDDSScheduler(
0346 PanDAiDDSRunnerConfig(name="user.test.panda", source_dir=str(tmp_path))
0347 )
0348
0349 first_name = first._next_work_name("evaluate", "dtlz2", "panda_test_evaluator")
0350 second_name = second._next_work_name("evaluate", "dtlz2", "panda_test_evaluator")
0351
0352 assert first_name != second_name
0353 assert first_name.startswith("user.test.a2e.")
0354 assert second_name.startswith("user.test.a2e.")
0355 assert first_name.split(".")[-2] == "000001"
0356 assert second_name.split(".")[-2] == "000001"
0357 assert len(first_name) <= 64
0358 assert len(second_name) <= 64
0359
0360
0361 def test_cancel_job_finds_preserved_job_id(tmp_path):
0362 class CancellableWork:
0363 def __init__(self):
0364 self.cancelled = False
0365
0366 def is_terminated(self):
0367 return False
0368
0369 def cancel(self):
0370 self.cancelled = True
0371
0372 scheduler = PanDAiDDSScheduler(
0373 PanDAiDDSRunnerConfig(name="user.test.panda", source_dir=str(tmp_path))
0374 )
0375 work = CancellableWork()
0376 scheduler.running_funcs["evaluate"] = {
0377 "evaluate:dtlz2:0001": {"funcs": {"panda_test_evaluator": {"work": work}}}
0378 }
0379
0380 assert scheduler.cancel_job("evaluate:dtlz2:0001") is True
0381 assert work.cancelled is True
0382
0383
0384 def test_remote_python_callable_entrypoint_runs_only_configured_callable():
0385 result = _remote_python_callable_entrypoint(
0386 "tests.test_schedulers.test_pandaidds_runner:panda_adapter_test_evaluator",
0387 {
0388 "task_id": "evaluate:dtlz2",
0389 "job_id": "dtlz2",
0390 "stage_id": "evaluate",
0391 "workflow_id": "wf",
0392 "design_point": {"x": 0.5},
0393 "xcom": {},
0394 "execution_dir": "/tmp/work",
0395 "output_dir": "/tmp/out",
0396 },
0397 {"label": "callable-only"},
0398 )
0399
0400 assert result["job_id"] == "dtlz2"
0401 assert result["design"] == {"x": 0.5}
0402 assert result["label"] == "callable-only"
0403 assert result["input_file_names"] is None
0404 assert result["xcom"] == {"evaluate:dtlz2:seen": "callable-only"}
0405
0406
0407 def test_remote_python_callable_entrypoint_accepts_idds_injected_kwargs():
0408 result = _remote_python_callable_entrypoint(
0409 "tests.test_schedulers.test_pandaidds_runner:panda_adapter_test_evaluator",
0410 {
0411 "task_id": "dtlz2",
0412 "job_id": "dtlz2",
0413 "stage_id": "evaluate",
0414 "workflow_id": "wf",
0415 "design_point": {},
0416 "xcom": {},
0417 },
0418 {"label": "dataset-ana"},
0419 input_file_names=["a.json", "b.json"],
0420 )
0421
0422 assert result["label"] == "dataset-ana"
0423 assert result["input_file_names"] == ["a.json", "b.json"]
0424
0425
0426 def test_all2one_input_dataset_adapter_expands_and_restores_yaml_list():
0427 normalized, groups = _normalize_all2one_input_datasets(
0428 {"input_file_names": ["user.test.sim.a", "user.test.sim.b"]}
0429 )
0430
0431 assert normalized == {
0432 "input_file_names__aid2e_0": "user.test.sim.a",
0433 "input_file_names__aid2e_1": "user.test.sim.b",
0434 }
0435 assert groups == {
0436 "input_file_names": [
0437 "input_file_names__aid2e_0",
0438 "input_file_names__aid2e_1",
0439 ]
0440 }
0441
0442 result = _remote_python_callable_entrypoint(
0443 "tests.test_schedulers.test_pandaidds_runner:panda_adapter_test_evaluator",
0444 {
0445 "task_id": "dtlz2",
0446 "job_id": "dtlz2",
0447 "stage_id": "ana",
0448 "workflow_id": "wf",
0449 "design_point": {},
0450 "xcom": {},
0451 },
0452 input_dataset_groups=groups,
0453 input_file_names__aid2e_0=["a.json"],
0454 input_file_names__aid2e_1=["b.json", "c.json"],
0455 )
0456
0457 assert result["input_file_names"] == ["a.json", "b.json", "c.json"]
0458
0459
0460 def test_remote_python_callable_entrypoint_omits_unsupplied_context():
0461 result = _remote_python_callable_entrypoint(
0462 "tests.test_schedulers.test_pandaidds_runner:panda_no_context_evaluator",
0463 {
0464 "task_id": "trial",
0465 "job_id": "trial",
0466 "stage_id": "optimizer",
0467 "workflow_id": "wf",
0468 },
0469 {"value": 4},
0470 with_context=False,
0471 )
0472
0473 assert result == {"value": 4}
0474
0475
0476 def test_panda_multistage_spec_rejects_unknown_dependency_map():
0477 with pytest.raises(ValueError, match="expected 'one2one'"):
0478 PanDAMultiStageSpec.from_payload(
0479 {
0480 "stages": [
0481 {
0482 "name": "simreco",
0483 "python_callable": panda_prepare_evaluator,
0484 "dep_map": "scatter",
0485 }
0486 ]
0487 }
0488 )
0489
0490
0491 def test_panda_multistage_spec_rejects_cycles():
0492 with pytest.raises(ValueError, match="cycle"):
0493 PanDAMultiStageSpec.from_payload(
0494 {
0495 "stages": [
0496 {
0497 "name": "first",
0498 "python_callable": panda_prepare_evaluator,
0499 "depends_on": "second",
0500 },
0501 {
0502 "name": "second",
0503 "python_callable": panda_test_evaluator,
0504 "depends_on": "first",
0505 },
0506 ]
0507 }
0508 )
0509
0510
0511 def test_panda_multistage_failed_parent_does_not_release_result_child(tmp_path):
0512 scheduler = PanDAiDDSScheduler(
0513 PanDAiDDSRunnerConfig(name="user.test.panda", source_dir=str(tmp_path))
0514 )
0515 coordinator = PanDAMultiStageJob(
0516 scheduler=scheduler,
0517 stage_name="evaluate",
0518 job_definition={
0519 "payload": {
0520 "stages": [
0521 {"name": "parent", "python_callable": panda_prepare_evaluator},
0522 {
0523 "name": "child",
0524 "python_callable": panda_test_evaluator,
0525 "depends_on": "parent",
0526 "dep_type": "results",
0527 },
0528 ]
0529 }
0530 },
0531 )
0532 coordinator.works_by_stage["parent"].append(
0533 types.SimpleNamespace(status="failed")
0534 )
0535
0536 assert coordinator._stage_ready_for_submission(coordinator.spec.by_name["child"]) is False
0537
0538
0539 def test_panda_multistage_jobs_run_concurrently_within_batch(monkeypatch, tmp_path):
0540 scheduler = PanDAiDDSScheduler(
0541 PanDAiDDSRunnerConfig(name="user.test.panda", source_dir=str(tmp_path))
0542 )
0543 barrier = threading.Barrier(2, timeout=2)
0544
0545 def fake_run(stage_name, job_definition, working_dir, poll_interval):
0546 barrier.wait()
0547 return {
0548 "status": "completed",
0549 "return_code": 0,
0550 "stdout": "",
0551 "stderr": "",
0552 "outputs": {"objectives": {"value": int(job_definition["job_id"]) }},
0553 }
0554
0555 monkeypatch.setattr(scheduler, "_run_multistage_job", fake_run)
0556 result = scheduler.run_stage(
0557 "optimizer_batch_1_trials",
0558 [
0559 {"job_id": "0", "payload": {"evaluator_type": "panda_multistage"}},
0560 {"job_id": "1", "payload": {"evaluator_type": "panda_multistage"}},
0561 ],
0562 parallelism_policy={"max_concurrent": 2, "poll_interval": 0},
0563 working_dir=str(tmp_path),
0564 )
0565
0566 assert result.success is True
0567 assert [status.job_id for status in result.job_statuses] == ["0", "1"]
0568
0569
0570 def test_panda_multistage_does_not_inject_implicit_stage_kwargs(tmp_path):
0571 scheduler = PanDAiDDSScheduler(
0572 PanDAiDDSRunnerConfig(name="user.test.panda", source_dir=str(tmp_path))
0573 )
0574
0575 result = scheduler.run_stage(
0576 "evaluate",
0577 [
0578 {
0579 "job_id": "logical:strict",
0580 "name": "strict_stage_job",
0581 "payload": {
0582 "evaluator_type": "panda_multistage",
0583 "stages": [
0584 {
0585 "name": "final",
0586 "python_callable": panda_strict_local_stage,
0587 "runner": "local",
0588 "produces_objective": True,
0589 }
0590 ],
0591 },
0592 }
0593 ],
0594 parallelism_policy={"poll_interval": 0},
0595 working_dir=str(tmp_path),
0596 )
0597
0598 assert result.success is True
0599 assert result.job_statuses[0].outputs == {"objectives": {"objective": 1}}
0600
0601
0602 def test_panda_multistage_dataset_dependency_is_managed_by_panda(monkeypatch, tmp_path):
0603 state = install_fake_idds(
0604 monkeypatch,
0605 result_sequence=[{"simreco": "done"}, {"f1": 3, "f2": 4}],
0606 details_sequence=[{"stage": "simreco"}, {"stage": "ana"}],
0607 )
0608 context = DummyContext()
0609 scheduler = PanDAiDDSScheduler(
0610 PanDAiDDSRunnerConfig(name="user.test.panda", source_dir=str(tmp_path))
0611 )
0612
0613 result = scheduler.run_stage(
0614 "evaluate",
0615 [
0616 {
0617 "job_id": "logical:multistage",
0618 "name": "dataset_stage_job",
0619 "payload": {
0620 "evaluator_type": "panda_multistage",
0621 "stages": [
0622 {
0623 "name": "simreco",
0624 "python_callable": panda_prepare_evaluator,
0625 "output_file": "my_test.txt",
0626 "output_dataset": "user.test.simreco.dataset",
0627 "num_events": 200,
0628 "num_events_per_job": 100,
0629 },
0630 {
0631 "name": "ana",
0632 "python_callable": panda_test_evaluator,
0633 "depends_on": "simreco",
0634 "dep_type": "datasets",
0635 "dep_map": "one2one",
0636 "input_datasets": {
0637 "input_file_names": "user.test.simreco.dataset"
0638 },
0639 "produces_objective": True,
0640 },
0641 ],
0642 },
0643 "job_context": context,
0644 }
0645 ],
0646 parallelism_policy={"poll_interval": 0},
0647 working_dir=str(tmp_path),
0648 )
0649
0650 assert result.success is True
0651 assert result.artifacts == {"objectives": {"f1": 3, "f2": 4}}
0652 assert result.job_statuses[0].outputs == {"objectives": {"f1": 3, "f2": 4}}
0653 assert context.pushed["objectives"] == {"f1": 3, "f2": 4}
0654 assert context.pushed["results_details"] == {"stage": "ana"}
0655
0656 simreco_kwargs = state["work_kwargs_history"][0]
0657 ana_kwargs = state["work_kwargs_history"][1]
0658 assert simreco_kwargs["output_file_name"] == "my_test.txt"
0659 assert simreco_kwargs["output_dataset_name"] == "user.test.simreco.dataset/"
0660 assert simreco_kwargs["num_events"] == 200
0661 assert simreco_kwargs["num_events_per_job"] == 100
0662 assert ana_kwargs["input_datasets"] == {
0663 "input_file_names": "user.test.simreco.dataset"
0664 }
0665 assert ana_kwargs["parent_internal_id"] == "internal-1"
0666 assert [kwargs["name"] for kwargs in state["work_kwargs_history"]] == [
0667 state["work_kwargs_history"][0]["job_key"],
0668 state["work_kwargs_history"][1]["job_key"],
0669 ]
0670 assert state["result_lookup_history"][-1] == {
0671 "name": state["work_kwargs_history"][1]["name"],
0672 "key": state["work_kwargs_history"][1]["job_key"],
0673 "verbose": True,
0674 "with_details": True,
0675 }
0676
0677 metrics = result.job_statuses[0].metrics["panda_multistage_stages"]
0678 assert [item["stage"] for item in metrics] == ["simreco", "ana"]
0679 assert [item["internal_id"] for item in metrics] == ["internal-1", "internal-2"]
0680
0681
0682 def test_panda_multistage_one2many_local_to_remote(monkeypatch, tmp_path):
0683 state = install_fake_idds(
0684 monkeypatch,
0685 result_sequence=[{"child": "a", "f1": 1}, {"child": "b", "f1": 2}],
0686 )
0687 scheduler = PanDAiDDSScheduler(
0688 PanDAiDDSRunnerConfig(name="user.test.panda", source_dir=str(tmp_path))
0689 )
0690
0691 result = scheduler.run_stage(
0692 "evaluate",
0693 [
0694 {
0695 "job_id": "logical:one2many",
0696 "name": "one2many_stage_job",
0697 "payload": {
0698 "evaluator_type": "panda_multistage",
0699 "stages": [
0700 {
0701 "name": "prepare",
0702 "python_callable": panda_prepare_evaluator,
0703 "runner": "local",
0704 "dep_type": "results",
0705 "op_kwargs": {"scale": 5},
0706 },
0707 {
0708 "name": "simreco",
0709 "python_callable": panda_chunk_evaluator,
0710 "depends_on": "prepare",
0711 "dep_type": "results",
0712 "dep_map": "one2many",
0713 "jobs": [
0714 {"key": "a", "op_kwargs": {"value": 1}},
0715 {"key": "b", "op_kwargs": {"value": 2}},
0716 ],
0717 "parent_result_parameter_name": "prepared_result",
0718 "produces_objective": True,
0719 },
0720 ],
0721 },
0722 }
0723 ],
0724 parallelism_policy={"poll_interval": 0},
0725 working_dir=str(tmp_path),
0726 )
0727
0728 assert result.success is True
0729 assert len(state["work_kwargs_history"]) == 2
0730 assert state["work_params_history"][0]["op_kwargs"] == {"value": 1, "prepared_result": {"prepared": 5}}
0731 assert state["work_params_history"][1]["op_kwargs"] == {"value": 2, "prepared_result": {"prepared": 5}}
0732 metrics = result.job_statuses[0].metrics["panda_multistage_stages"]
0733 assert [item["execution"] for item in metrics] == ["local", "panda", "panda"]
0734
0735
0736 def test_panda_multistage_submits_remote_children_before_parent_completion(monkeypatch, tmp_path):
0737 state = install_fake_idds(monkeypatch, status_sequence=["running", "running", "running"])
0738 scheduler = PanDAiDDSScheduler(
0739 PanDAiDDSRunnerConfig(name="user.test.panda", source_dir=str(tmp_path))
0740 )
0741 coordinator = PanDAMultiStageJob(
0742 scheduler=scheduler,
0743 stage_name="evaluate",
0744 job_definition={
0745 "job_id": "logical:submit-first",
0746 "name": "submit_first",
0747 "payload": {
0748 "evaluator_type": "panda_multistage",
0749 "stages": [
0750 {
0751 "name": "simreco",
0752 "python_callable": panda_prepare_evaluator,
0753 "output_file": "my_test.txt",
0754 "jobs": [
0755 {"key": "a", "output_dataset": "user.test.simreco.a"},
0756 {"key": "b", "output_dataset": "user.test.simreco.b"},
0757 ],
0758 },
0759 {
0760 "name": "ana",
0761 "python_callable": panda_test_evaluator,
0762 "depends_on": "simreco",
0763 "dep_type": "datasets",
0764 "dep_map": "all2one",
0765 "input_datasets": {
0766 "input_file_names": [
0767 "user.test.simreco.a",
0768 "user.test.simreco.b",
0769 ]
0770 },
0771 "produces_objective": True,
0772 },
0773 ],
0774 },
0775 },
0776 working_dir=str(tmp_path),
0777 poll_interval=0,
0778 )
0779
0780 coordinator._submit_stage(coordinator.spec.stages[0])
0781 assert [work.status for work in coordinator.works_by_stage["simreco"]] == ["running", "running"]
0782
0783 coordinator._submit_stage(coordinator.spec.stages[1])
0784
0785 assert len(state["work_kwargs_history"]) == 3
0786 assert state["work_kwargs_history"][2]["parent_internal_id"] == ["internal-1", "internal-2"]
0787 assert state["work_kwargs_history"][2]["input_datasets"] == {
0788 "input_file_names__aid2e_0": "user.test.simreco.a",
0789 "input_file_names__aid2e_1": "user.test.simreco.b",
0790 }
0791 assert state["work_params_history"][2]["input_dataset_groups"] == {
0792 "input_file_names": ["input_file_names__aid2e_0", "input_file_names__aid2e_1"]
0793 }
0794 assert [work.status for work in coordinator.works_by_stage["simreco"]] == ["running", "running"]
0795 assert coordinator.works_by_stage["ana"][0].status == "running"
0796
0797
0798
0799 def test_panda_multistage_all2one_dataset_remote_to_remote(monkeypatch, tmp_path):
0800 state = install_fake_idds(
0801 monkeypatch,
0802 result_sequence=[{"sim": "a"}, {"sim": "b"}, {"f1": 3, "f2": 4}],
0803 )
0804 scheduler = PanDAiDDSScheduler(
0805 PanDAiDDSRunnerConfig(name="user.test.panda", source_dir=str(tmp_path))
0806 )
0807
0808 result = scheduler.run_stage(
0809 "evaluate",
0810 [
0811 {
0812 "job_id": "logical:all2one",
0813 "name": "all2one_stage_job",
0814 "payload": {
0815 "evaluator_type": "panda_multistage",
0816 "stages": [
0817 {
0818 "name": "simreco",
0819 "python_callable": panda_prepare_evaluator,
0820 "output_file": "my_test.txt",
0821 "jobs": [
0822 {"key": "a", "output_dataset": "user.test.simreco.a"},
0823 {"key": "b", "output_dataset": "user.test.simreco.b"},
0824 ],
0825 },
0826 {
0827 "name": "ana",
0828 "python_callable": panda_test_evaluator,
0829 "depends_on": "simreco",
0830 "dep_type": "datasets",
0831 "dep_map": "all2one",
0832 "input_datasets": {
0833 "input_file_names": [
0834 "user.test.simreco.a",
0835 "user.test.simreco.b",
0836 ]
0837 },
0838 "produces_objective": True,
0839 },
0840 ],
0841 },
0842 }
0843 ],
0844 parallelism_policy={"poll_interval": 0},
0845 working_dir=str(tmp_path),
0846 )
0847
0848 assert result.success is True
0849 assert result.artifacts == {"objectives": {"f1": 3, "f2": 4}}
0850 assert len(state["work_kwargs_history"]) == 3
0851 assert state["work_kwargs_history"][0]["output_dataset_name"] == "user.test.simreco.a/"
0852 assert state["work_kwargs_history"][1]["output_dataset_name"] == "user.test.simreco.b/"
0853 assert state["work_kwargs_history"][2]["input_datasets"] == {
0854 "input_file_names__aid2e_0": "user.test.simreco.a",
0855 "input_file_names__aid2e_1": "user.test.simreco.b",
0856 }
0857 assert state["work_kwargs_history"][2]["parent_internal_id"] == ["internal-1", "internal-2"]
0858 metrics = result.job_statuses[0].metrics["panda_multistage_stages"]
0859 assert [item["execution"] for item in metrics] == ["panda", "panda", "panda"]
0860
0861
0862 def test_panda_multistage_one2one_result_local_to_remote(monkeypatch, tmp_path):
0863 state = install_fake_idds(monkeypatch, result_sequence=[{"f1": 6, "f2": 7}])
0864 scheduler = PanDAiDDSScheduler(
0865 PanDAiDDSRunnerConfig(name="user.test.panda", source_dir=str(tmp_path))
0866 )
0867
0868 result = scheduler.run_stage(
0869 "evaluate",
0870 [
0871 {
0872 "job_id": "logical:one2one-local-remote",
0873 "name": "one2one_local_remote",
0874 "payload": {
0875 "evaluator_type": "panda_multistage",
0876 "stages": [
0877 {
0878 "name": "prepare",
0879 "python_callable": panda_prepare_evaluator,
0880 "runner": "local",
0881 "dep_type": "results",
0882 "op_kwargs": {"scale": 4},
0883 },
0884 {
0885 "name": "ana",
0886 "python_callable": panda_evaluate_from_parent,
0887 "depends_on": "prepare",
0888 "dep_type": "results",
0889 "dep_map": "one2one",
0890 "produces_objective": True,
0891 },
0892 ],
0893 },
0894 }
0895 ],
0896 parallelism_policy={"poll_interval": 0},
0897 working_dir=str(tmp_path),
0898 )
0899
0900 assert result.success is True
0901 assert state["work_params_history"][0]["op_kwargs"]["parent_results"] == {"prepared": 4}
0902 metrics = result.job_statuses[0].metrics["panda_multistage_stages"]
0903 assert [item["execution"] for item in metrics] == ["local", "panda"]
0904
0905
0906 def test_panda_multistage_graph_builder_builds_all2one_dataset_payload():
0907 payload = PanDAMultiStageGraphBuilder().build_payload(
0908 [
0909 {
0910 "name": "simreco",
0911 "jobs": [
0912 {
0913 "name": "sim_a",
0914 "payload": {
0915 "python_callable": panda_prepare_evaluator,
0916 "output_file": "my_test.txt",
0917 "output_dataset": "user.test.sim.a",
0918 "op_kwargs": {"particle": "pi+"},
0919 },
0920 },
0921 {
0922 "name": "sim_b",
0923 "payload": {
0924 "python_callable": panda_prepare_evaluator,
0925 "output_file": "my_test.txt",
0926 "output_dataset": "user.test.sim.b",
0927 "op_kwargs": {"particle": "kaon+"},
0928 },
0929 },
0930 ],
0931 },
0932 {
0933 "name": "ana",
0934 "jobs": [
0935 {
0936 "name": "ana_all",
0937 "payload": {
0938 "python_callable": panda_test_evaluator,
0939 "depends_on": {"stage": "simreco", "jobs": "*"},
0940 "dep_type": "datasets",
0941 "dep_map": "many2one",
0942 "input_datasets": {
0943 "input_file_names": ["user.test.sim.a", "user.test.sim.b"]
0944 },
0945 "produces_objective": True,
0946 },
0947 }
0948 ],
0949 },
0950 ],
0951 final_stage="ana",
0952 trial_index=7,
0953 )
0954
0955 assert payload["evaluator_type"] == "panda_multistage"
0956 assert payload["final_stage"] == "ana"
0957 assert payload["trial_index"] == 7
0958 assert [stage["name"] for stage in payload["stages"]] == ["simreco", "ana"]
0959 assert payload["stages"][0]["python_callable"] is panda_prepare_evaluator
0960 assert payload["stages"][0]["jobs"] == [
0961 {
0962 "key": "sim_a",
0963 "op_kwargs": {"particle": "pi+"},
0964 "output_file": "my_test.txt",
0965 "output_dataset": "user.test.sim.a",
0966 },
0967 {
0968 "key": "sim_b",
0969 "op_kwargs": {"particle": "kaon+"},
0970 "output_file": "my_test.txt",
0971 "output_dataset": "user.test.sim.b",
0972 },
0973 ]
0974 assert payload["stages"][1]["depends_on"] == {"stage": "simreco", "jobs": "*"}
0975 assert payload["stages"][1]["dep_map"] == "many2one"
0976 assert payload["stages"][1]["jobs"] == [
0977 {
0978 "key": "ana_all",
0979 "input_datasets": {
0980 "input_file_names": ["user.test.sim.a", "user.test.sim.b"]
0981 },
0982 }
0983 ]
0984
0985
0986
0987 def test_panda_runner_converts_optimizer_trial_wrapper_to_multistage_graph(monkeypatch, tmp_path):
0988 design_file = tmp_path / "design.params"
0989 design_file.write_text(
0990 "design_space:\n"
0991 " design_parameters:\n"
0992 " test_variables:\n"
0993 " parameters:\n"
0994 " x1:\n"
0995 " value: 0.5\n"
0996 " bounds: [0.0, 1.0]\n"
0997 )
0998 config_file = tmp_path / "trial_graph.yml"
0999 config_file.write_text(
1000 f"""
1001 problem:
1002 name: "trial graph smoke"
1003 problem_type: "DTLZ2"
1004 output_location: "{tmp_path / 'out'}"
1005 work_location: "{tmp_path / 'work'}"
1006 design_parameters_file: "{design_file}"
1007 objectives:
1008 - name: "objective"
1009 direction: "minimize"
1010
1011 optimizer:
1012 name: "ax"
1013 type: "bayesian"
1014 parameters:
1015 n_iterations: 1
1016 n_initial_samples: 1
1017 batch_size: 1
1018
1019 scheduler:
1020 runner_type: "PanDAiDDSRunner"
1021 parameters:
1022 name: "user.test.panda"
1023 source_dir: "{tmp_path}"
1024
1025 workflows:
1026 workflows:
1027 - name: "trial_graph_eval"
1028 branches:
1029 - name: "main"
1030 stages:
1031 - name: "simreco"
1032 jobs:
1033 - name: "sim_a"
1034 command: ""
1035 payload:
1036 evaluator_type: "python"
1037 python_callable: "tests.test_schedulers.test_pandaidds_runner:panda_prepare_evaluator"
1038 output_file: "my_test.txt"
1039 output_dataset: "user.test.trial.sim.a"
1040 - name: "sim_b"
1041 command: ""
1042 payload:
1043 evaluator_type: "python"
1044 python_callable: "tests.test_schedulers.test_pandaidds_runner:panda_prepare_evaluator"
1045 output_file: "my_test.txt"
1046 output_dataset: "user.test.trial.sim.b"
1047 - name: "ana"
1048 jobs:
1049 - name: "ana_all"
1050 command: ""
1051 payload:
1052 evaluator_type: "python"
1053 python_callable: "tests.test_schedulers.test_pandaidds_runner:panda_test_evaluator"
1054 depends_on: "simreco"
1055 dep_type: "datasets"
1056 dep_map: "all2one"
1057 input_datasets:
1058 input_file_names:
1059 - "user.test.trial.sim.a"
1060 - "user.test.trial.sim.b"
1061 - name: "final"
1062 scheduler:
1063 runner_type: "JobLibRunner"
1064 parameters:
1065 n_jobs: 1
1066 jobs:
1067 - name: "final_objective"
1068 command: ""
1069 payload:
1070 evaluator_type: "python"
1071 python_callable: "tests.test_schedulers.test_pandaidds_runner:panda_local_final_from_ana"
1072 depends_on: "ana"
1073 dep_type: "results"
1074 dep_map: "one2one"
1075 parent_result_parameter_name: "ana_result"
1076 final: true
1077 """
1078 )
1079 state = install_fake_idds(
1080 monkeypatch,
1081 result_sequence=[{"sim": "a"}, {"sim": "b"}, {"total": 5}],
1082 )
1083 scheduler = PanDAiDDSScheduler(
1084 PanDAiDDSRunnerConfig(name="user.test.panda", source_dir=str(tmp_path))
1085 )
1086
1087 result = scheduler.run_stage(
1088 "optimizer_batch_1_trials",
1089 [
1090 {
1091 "job_id": "0",
1092 "name": "trial_0",
1093 "function": run_trial_workflow,
1094 "params": {
1095 "config_path": str(config_file),
1096 "run_dir": str(tmp_path / "out"),
1097 "run_work_dir": str(tmp_path / "work"),
1098 "trial_index": 0,
1099 "design_point": {"x1": 0.2},
1100 "workflow_name": None,
1101 "log_level": "WARNING",
1102 },
1103 }
1104 ],
1105 parallelism_policy={"poll_interval": 0},
1106 working_dir=str(tmp_path / "work"),
1107 )
1108
1109 assert result.success is True
1110 assert result.job_statuses[0].outputs == {"objectives": {"objective": 6}}
1111 assert len(state["work_kwargs_history"]) == 3
1112 assert state["work_kwargs_history"][0]["output_dataset_name"] == "user.test.trial.sim.a/"
1113 assert state["work_kwargs_history"][1]["output_dataset_name"] == "user.test.trial.sim.b/"
1114 assert state["work_kwargs_history"][2]["parent_internal_id"] == ["internal-1", "internal-2"]
1115 assert state["work_kwargs_history"][2]["input_datasets"] == {
1116 "input_file_names__aid2e_0": "user.test.trial.sim.a",
1117 "input_file_names__aid2e_1": "user.test.trial.sim.b",
1118 }
1119
1120
1121 def test_panda_multistage_graph_job_runs_stage_records(monkeypatch, tmp_path):
1122 state = install_fake_idds(
1123 monkeypatch,
1124 result_sequence=[{"sim": "a"}, {"sim": "b"}, {"f1": 8, "f2": 9}],
1125 )
1126 scheduler = PanDAiDDSScheduler(
1127 PanDAiDDSRunnerConfig(name="user.test.panda", source_dir=str(tmp_path))
1128 )
1129
1130 result = scheduler.run_stage(
1131 "evaluate",
1132 [
1133 {
1134 "job_id": "logical:graph",
1135 "name": "graph_job",
1136 "payload": {
1137 "evaluator_type": "panda_multistage_graph",
1138 "final_stage": "ana",
1139 "stage_records": [
1140 {
1141 "name": "simreco",
1142 "jobs": [
1143 {
1144 "name": "sim_a",
1145 "payload": {
1146 "python_callable": panda_prepare_evaluator,
1147 "output_file": "my_test.txt",
1148 "output_dataset": "user.test.graph.sim.a",
1149 },
1150 },
1151 {
1152 "name": "sim_b",
1153 "payload": {
1154 "python_callable": panda_prepare_evaluator,
1155 "output_file": "my_test.txt",
1156 "output_dataset": "user.test.graph.sim.b",
1157 },
1158 },
1159 ],
1160 },
1161 {
1162 "name": "ana",
1163 "jobs": [
1164 {
1165 "name": "ana_all",
1166 "payload": {
1167 "python_callable": panda_test_evaluator,
1168 "depends_on": "simreco",
1169 "dep_type": "datasets",
1170 "dep_map": "all2one",
1171 "input_datasets": {
1172 "input_file_names": [
1173 "user.test.graph.sim.a",
1174 "user.test.graph.sim.b",
1175 ]
1176 },
1177 "produces_objective": True,
1178 },
1179 }
1180 ],
1181 },
1182 ],
1183 },
1184 }
1185 ],
1186 parallelism_policy={"poll_interval": 0},
1187 working_dir=str(tmp_path),
1188 )
1189
1190 assert result.success is True
1191 assert result.artifacts == {"objectives": {"f1": 8, "f2": 9}}
1192 assert len(state["work_kwargs_history"]) == 3
1193 assert state["work_kwargs_history"][2]["parent_internal_id"] == ["internal-1", "internal-2"]
1194 assert state["work_kwargs_history"][2]["input_datasets"] == {
1195 "input_file_names__aid2e_0": "user.test.graph.sim.a",
1196 "input_file_names__aid2e_1": "user.test.graph.sim.b",
1197 }
1198 metrics = result.job_statuses[0].metrics["panda_multistage_stages"]
1199 assert [item["stage"] for item in metrics] == ["simreco", "simreco", "ana"]
1200
1201
1202 def test_dtlz2_dataset_simreco_writes_message_output_file(tmp_path, monkeypatch):
1203 from examples.evaluators.dtlz2_panda import panda_stage_simreco
1204
1205 context = types.SimpleNamespace(
1206 design_point={
1207 "DTLZ2_variables.x1": 0.2,
1208 "DTLZ2_variables.x2": 0.3,
1209 "DTLZ2_variables.x3": 0.4,
1210 "DTLZ2_variables.x4": 0.5,
1211 "DTLZ2_variables.x5": 0.6,
1212 }
1213 )
1214 monkeypatch.chdir(tmp_path)
1215
1216 result = panda_stage_simreco(
1217 context,
1218 particle="pi+",
1219 eta_point=0.1,
1220 output_file_name="my_test.txt",
1221 output_dataset_name="user.test.aid2e.dtlz2.simreco.eta_0p1_pi",
1222 num_events=200,
1223 num_events_per_job=100,
1224 )
1225
1226 assert result["particle"] == "pi+"
1227 assert "simreco produced xyz=" in result["message"]
1228 output_path = tmp_path / "my_test.txt"
1229 assert output_path.is_file()
1230
1231 import json
1232
1233 data = json.loads(output_path.read_text())
1234 assert data["xyz"] == result["xyz"]
1235 assert data["message"] == result["message"]
1236 assert data["particle"] == "pi+"
1237 assert data["eta_point"] == 0.1
1238 assert not (tmp_path / "user.test.aid2e.dtlz2.simreco.eta_0p1_pi_000001.my_test.txt").exists()
1239 assert not (tmp_path / "user.test.aid2e.dtlz2.simreco.eta_0p1_pi_000002.my_test.txt").exists()
1240
1241
1242 def test_dtlz2_dataset_ana_reads_and_prints_simreco_messages(tmp_path, capsys):
1243 from examples.evaluators.dtlz2_panda import panda_stage_ana
1244
1245 first = tmp_path / "first.json"
1246 second = tmp_path / "second.json"
1247 first.write_text('{"xyz": 1.25, "message": "first simreco message"}\n')
1248 second.write_text('{"xyz": 2.75, "message": "second simreco message"}\n')
1249
1250 result = panda_stage_ana(
1251 types.SimpleNamespace(design_point={}),
1252 input_file_names=[str(first), str(second)],
1253 )
1254
1255 captured = capsys.readouterr()
1256 assert result == {"xyz": 4.0, "n_inputs": 2}
1257 assert f"ana read from {first}: first simreco message" in captured.out
1258 assert f"ana read from {second}: second simreco message" in captured.out
