Incorrect instantiation order with nested dependency injection

[jkrude]

🐛 Bug report

First of all, this is a really cool project, and I am enjoying it a lot 👍 .
This might be quite niche to be honest, but I encountered unexpected behavior when working on nested classes.
In my use case, the user can choose from multiple loss classes (Strategy in the example), which on their own have another nested dependency injection (NestedStrategy).
Furthermore, some arguments of the loss are linked to (StrategyParameter) and the nested strategy has also arguments provided from somewhere else (NestedParameter).

Now in this setting I found that the instantiation order induced by the topological sort of the dependency graph does not enforce that NestedParameter is instantiated before Strategy.
But when instantiating the Strategy object, all actions which target the nested strategy are triggered too.
So in particular the link

parser.link_arguments(
        "nested_param",
        "root.strategy.init_args.nested.init_args.nested_param",
        apply_on="instantiate",
        compute_fn=compute_fn,
    )

is triggered however before NestedStrategy was instantiated and so the compute_fn is called with the Namespace object.

I am not sure whether the graph should contain the dependency or the instantiation logic should not instantiate the actions leading to the nested class?
In ActionLink.apply_instantiation_links there is this predicate is_nested_instantiation_link (returning False) which hints at the second case being true?

This is how the dependency graph looks like as constructed in ActionLink.instantiation_order:

stateDiagram-v2
    nested_param --> root.strategy.init_args.nested,
    param --> root.strategy

Loading

To reproduce

I tried to keep the example as small as possible.
We have two dependency injected classes Strategy and NestedStrategy and two objects which are used for linking StrategyParameter and NestedParameter.

from jsonargparse import ArgumentParser

class NestedStrategy:

    def __init__(self, nested_param: str):
        self.nested_param = nested_param


class Strategy:

    def __init__(self, param: str, nested: NestedStrategy):
        self.param = param
        self.nested = nested


class Root:

    def __init__(self, strategy: Strategy):
        self.strategy = strategy


class NestedParameter:
    def __init__(self, something_else: str):
        self._something_else = something_else


class StrategyParameter:

    def __init__(self, something: str):
        self.something = something


def compute_fn(nested_param: NestedParameter):
    assert isinstance(
        nested_param, NestedParameter
    ), f"Got wrong type {type(nested_param)}"
    return nested_param


if __name__ == "__main__":
    parser = ArgumentParser()

    parser.add_class_arguments(Root, "root")
    parser.add_class_arguments(NestedParameter, "nested_param")
    parser.add_class_arguments(StrategyParameter, "param")
    parser.link_arguments(
        "nested_param",
        "root.strategy.init_args.nested.init_args.nested_param",
        apply_on="instantiate",
        compute_fn=compute_fn,
    )
    parser.link_arguments(
        "param.param", "root.strategy.init_args.param", apply_on="instantiate"
    )
    parser.add_argument("--config", action="config")
    cfg = parser.parse_args(["--config", "config.yaml"])
    init_cfg = parser.instantiate_classes(cfg)

root:
  strategy:
    class_path: Strategy
    init_args:
      nested:
        class_path: NestedStrategy
nested_param:
  something_else: "Something Else"

param:
  something: "something"

Output

ValueError: Call to compute_fn of link 'compute_fn(nested_param) --> root.strategy.init_args.nested.init_args.nested_param' with args (Namespace(something_else='Something Else')) failed: Got wrong type. Expected NestedParameter but got <class 'jsonargparse._namespace.Namespace'>

Expected behavior

The NestedParameter class is instantiated before passed as parameter to the compute_fn.
Note, the compute_fn is only to showcase the problem but not necessary to reproduce the bug.

Environment

• jsonargparse version: jsonargparse 4.34.0
• Python version: Python 3.10.14
• How jsonargparse was installed: pip
• OS: macOS

[mauvilsa]

Thank you for reporting! I will look at it.

[jkrude]

Any updates on this issue? I have run into it multiple times by now. Let me know if I can help to fix it.

[mauvilsa]

I did have a look at this. I noticed that when adding links an directed acyclic graph (DAG) is used to check that the links make sense. But during instantiation a DAG is not used to determine the order. Adding the DAG to the instantiation is a complex task. I haven't had time to work on it. If you want to help, you could take a look at the code and maybe describe here a potential proposal on how to implement it. But as I said, it is not simple.

I do wonder why only you have reported this problem. It is quite an obvious bug, but it seems in practice it doesn't happen often.

[jkrude]

Thanks for the quick response. Correct me if I am wrong, but the topological sort of the constructed DAG is used to determine the instantiation order of components in instantiate_classes here right?
First, to clarify here is how I understand what happens in my example:

In my example, the call to ActionLink.instantiation_order(self) results in the induced action-link-order of
[param, root.strategy, nested_param, root.strategy.init_args.nested],
notice how root.strategy comes before root.strategy.init_args.nested.

Based on the action-link-order, the components are reordered to
['param.something', 'param', 'root.strategy', 'nested_param.something_else', 'nested_param', 'root'].

Following this order, the param.something and param (e.g., the class StrategyParameter) are instantiated first.
Secondly, the component root.strategy is instantiated, e.g., the class Strategy.

Before the class itself is instantiated, apply_instantiation_links is called with target='root.strategy'.
Here, the relevant action-links are filtered out. In my example, those are:
['compute_fn(nested_param) --> root.strategy.init_args.nested.init_args.nested_param', 'param.param --> root.strategy.init_args.param'].

At this point inside apply_instantiation_links, it is assumed that the sources for each action-link are already instantiated.
And thus in line 361 the error is triggered as nested_param is not yet instantiated.

In general, this looks like:

• Instantiate classes instantiate_classes
  • Sort components by length-of split-key -> ensures that root.strategy is instantiated before root
  • Sort by action links via topological sort
  • Instantiate individual components
    • Instantiate all action-links ActionLink.apply_instantiation_links(self, cfg, target=component.dest)
      • Filter for target_key == target or target_key.startswith(f"{target}.") -> 'root.strategy.init_args.nested.init_args.nested_param' starts with 'root.strategy'
    • Instantiate the component itself -> triggers recursive call to instantiate_classes for subcomponents -> instantiates root.strategy.init_args.nested

If I understand this quite complex logic correctly, there are at least three rules we have to consider:

1. Instantiate action-sources before action-targets
2. Instantiate subcomponents before parent-components
3. Instantiate action-link-sources of subcomponents before instantiating parent-components.

The last rule is the critical in our case, which is not enforced by any mechanism.

Looking at the logic described above, we need to make sure that the sources of action-links to root.strategy were instantiated before, e.g., that nested_param was instantiated.
In this case, I would argue for an additional dependency in our action-link-DAG, adding an edge between a classes' init_args and itself.
So the expected graph would look like:

stateDiagram-v2
    nested_param --> root.strategy.init_args.nested
    root.strategy.init_args.nested --> root.strategy
    param --> root.strategy

Loading

Instead of the one posted above.
This should guarantee the component initialization of:
['param.something', 'param', 'nested_param.something_else', 'nested_param', 'root.strategy', 'root']

I hope you made it this far and this rather long description makes some sense :)

[mauvilsa]

Thanks for the quick response. Correct me if I am wrong, but the topological sort of the constructed DAG is used to determine the instantiation order of components in instantiate_classes here right?

Oh, you are right. Not sure why I had this in mind. I must have gotten confused with something else. I will take another look.

[mauvilsa]

I created pull request #662 with hopefully the fix. Please review and test.

The code snippet above had some issues so I had to change it a bit. See the test https://github.com/omni-us/jsonargparse/pull/662/files#diff-911b02ee18fa5c88fe775f9d0de98b387b537757ec95af835f485649ed3a86f7

[jkrude]

Awesome, that fixes my issue 👍. Sorry for the malformed example ^^

[mauvilsa]

Thank you for the great analysis! Certainly made fixing it way easier.