Skip to content

Commit a593cb7

Browse files
authored
Refactor _rewrite_inner_scalar function for improved step reference handling (#2352)
- Adjust the order of replacement steps to ensure child step renames occur before parameter bindings, preserving parent step references in bindings. - Update unit tests to reflect the new order and add a test for preserving parent step references in parameter bindings.
1 parent e108b3d commit a593cb7

2 files changed

Lines changed: 26 additions & 18 deletions

File tree

  • inference/core/workflows/execution_engine/v1/inner_workflow
  • tests/workflows/unit_tests/execution_engine/inner_workflow

inference/core/workflows/execution_engine/v1/inner_workflow/inline.py

Lines changed: 14 additions & 17 deletions
Original file line numberDiff line numberDiff line change
@@ -256,11 +256,13 @@ def _rewrite_inner_scalar(
256256
257257
Non-strings are returned unchanged. For strings, replacements run in order:
258258
259-
1. ``_replace_inputs_in_string`` — workflow inputs and defaults.
260-
2. If that step produced a non-string (whole ``$inputs.*`` reference), it is returned as-is;
261-
dotted step references are not applied to non-strings.
262-
3. ``_replace_step_prefixes_in_string`` — ``$steps.<child>.`` → ``$steps.<new>.``.
263-
4. ``_replace_bare_child_step_refs_in_string`` — bare ``$steps.<child>`` tokens.
259+
1. ``_replace_step_prefixes_in_string`` — ``$steps.<child>.`` → ``$steps.<new>.``.
260+
2. ``_replace_bare_child_step_refs_in_string`` — bare ``$steps.<child>`` tokens.
261+
3. ``_replace_inputs_in_string`` — workflow inputs and defaults (parent selectors).
262+
263+
Child step renames run before parameter bindings so a binding such as
264+
``$steps.dynamic_crop.crops`` that refers to a parent step is not rewritten when the
265+
child workflow also defines a step named ``dynamic_crop``.
264266
265267
``step_pairs`` maps each original child step name to its unique name after inlining
266268
(typically ``f"{inner_name}__{child_step}"``).
@@ -293,26 +295,21 @@ def _rewrite_inner_scalar(
293295
if not isinstance(value, str):
294296
return value
295297

296-
after_inputs = _replace_inputs_in_string(
297-
value,
298-
bindings=bindings,
299-
input_defaults=input_defaults,
300-
)
301-
302-
if not isinstance(after_inputs, str):
303-
return after_inputs
304-
305298
dotted = _replace_step_prefixes_in_string(
306-
after_inputs,
299+
value,
307300
old_to_new=step_pairs,
308301
)
309302

310-
out = _replace_bare_child_step_refs_in_string(
303+
after_bare = _replace_bare_child_step_refs_in_string(
311304
dotted,
312305
step_pairs=step_pairs,
313306
)
314307

315-
return out
308+
return _replace_inputs_in_string(
309+
after_bare,
310+
bindings=bindings,
311+
input_defaults=input_defaults,
312+
)
316313

317314

318315
def _deep_map_leaves(obj: Any, fn: Callable[[Any], Any]) -> Any:

tests/workflows/unit_tests/execution_engine/inner_workflow/test_inline.py

Lines changed: 12 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -221,7 +221,7 @@ def test_rewrite_inner_scalar_non_string_unchanged() -> None:
221221
) == {"nested": True}
222222

223223

224-
def test_rewrite_inner_scalar_pipeline_inputs_then_dotted_then_bare() -> None:
224+
def test_rewrite_inner_scalar_pipeline_dotted_then_bare_then_inputs() -> None:
225225
out = _rewrite_inner_scalar(
226226
"$inputs.msg and $steps.echo.out and $steps.echo",
227227
step_pairs=[("echo", "inner__echo")],
@@ -231,6 +231,17 @@ def test_rewrite_inner_scalar_pipeline_inputs_then_dotted_then_bare() -> None:
231231
assert out == "hello and $steps.inner__echo.out and $steps.inner__echo"
232232

233233

234+
def test_rewrite_inner_scalar_preserves_parent_step_in_parameter_binding() -> None:
235+
"""Parent ``$steps.<name>`` in bindings must not be renamed to the inlined child step."""
236+
out = _rewrite_inner_scalar(
237+
"$inputs.vehicle_image",
238+
step_pairs=[("dynamic_crop", "inner__dynamic_crop")],
239+
bindings={"vehicle_image": "$steps.dynamic_crop.crops"},
240+
input_defaults={},
241+
)
242+
assert out == "$steps.dynamic_crop.crops"
243+
244+
234245
def test_rewrite_inner_scalar_whole_input_non_string_skips_step_rewrite() -> None:
235246
sentinel = {"selector": "$steps.echo.x"}
236247
out = _rewrite_inner_scalar(

0 commit comments

Comments
 (0)