Skip to content

Sliced-only keys on a prep step config crash at workflow construction #408

Description

@wanghan-iapcm

dargs accepts template_slice_config, continue_on_num_success and continue_on_success_ratio on a prep step config, but the superops only sanitise run_config before spreading. Setting any of them under a prep step config therefore crashes at workflow-construction time, on input the schema declares valid.

Reproduction

Both at e45c147ed261d7c43b4b5bc9ada6e7d271e32c51.

normalize accepts both keys on a prep config:

template_slice_config      -> ACCEPTED, key present: True
continue_on_success_ratio  -> ACCEPTED, key present: True

(a) template_slice_config — only run_config ever pops it, so it survives into the Step(...) spread:

PrepRunLmp("x", PrepLmp, RunLmp,
           prep_config=normalize({"template_slice_config": {"group_size": 2}}),
           run_config=normalize({}))
TypeError: Step.__init__() got an unexpected keyword argument 'template_slice_config'
  at dpgen2/superop/prep_run_lmp.py:151

(b) continue_on_success_ratio on a non-sliced step. dflow only assigns total inside the dflow_nslices / with_param / with_sequence branches, then uses it unconditionally:

prep_lmp = Step(
"prep-lmp",
template=PythonOPTemplate(
prep_op,
output_artifact_archive={"task_paths": None},
python_packages=upload_python_packages,
**prep_template_config,
),
parameters={
"lmp_task_grp": prep_run_steps.inputs.parameters["expl_task_grp"],
},
artifacts={},
key=step_keys["prep-lmp"],
executor=prep_executor,
**prep_config,
)

Isolated against dflow directly, a non-sliced step with the key set:

UnboundLocalError: cannot access local variable 'total' where it is not associated with a value
  at dflow/step.py:970

while the same step with with_param=[1, 2, 3] constructs fine — confirming it is specifically the non-sliced case.

Scope

Every prep_run_*.py pops template_slice_config from run_config only, then spreads the unsanitised prep_config into a non-sliced prep step:

file pops from run_config spreads **prep_config into a non-sliced step
prep_run_calypso.py L165 L181
prep_run_lmp.py L149 L165
prep_run_fp.py L143 L162
prep_run_dp_train.py L166 L183

prep_run_diffcsp.py pops at L136 but spreads nothing at all — see the companion issue.

The fix already exists elsewhere

60c8f94 ("fix: continue_on_num_success/continue_on_success_ratio key error", #226, 2024-06-03) solved exactly this for caly_evo_step.py by deriving a sanitised copy before spreading into non-sliced steps:

template_slice_config = run_config.pop("template_slice_config", {})
expl_mode = caly_evo_step_steps.expl_mode
no_slice_run_config = deepcopy(run_config)
no_slice_run_config.pop("continue_on_num_success", None)
no_slice_run_config.pop("continue_on_success_ratio", None)

The prep side of the four prep_run_*.py files never got the same treatment.

Suggested fix

Give prep_config the same handling run_config gets — pop template_slice_config, and drop the sliced-only continuation keys before spreading into a non-sliced step. Alternatively, reject these keys for prep step configs at the dargs layer so the failure happens at submit-time normalization with a clear message rather than as a TypeError / UnboundLocalError from inside dflow.

Found while reviewing #380.

Activity

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Metadata

Metadata

Assignees

No one assigned

    Labels

    bugSomething isn't working

    Type

    No type

    Projects

    No projects

      Milestone

      No milestone

      Relationships

      None yet

      Development

      No branches or pull requests

      Issue actions