Experimental support for custom instantiators receiving values applied by instantiation links (Lightning-AI/pytorch-lightning#20311).

Author: mauvilsa

CHANGELOG.rst

@@ -27,6 +27,9 @@ Added
   (`#698 <https://github.com/omni-us/jsonargparse/pull/698>`__).
 - Option to enable validation of default values (`#711
   <https://github.com/omni-us/jsonargparse/pull/711>`__).
+- Experimental support for custom instantiators receiving values applied by
+  instantiation links (`#716
+  <https://github.com/omni-us/jsonargparse/pull/716>`__).
 
 Changed
 ^^^^^^^

jsonargparse/_common.py

@@ -64,6 +64,7 @@ def __call__(self, class_type: Type[ClassType], *args, **kwargs) -> ClassType:
 load_value_mode: ContextVar[Optional[str]] = ContextVar("load_value_mode", default=None)
 class_instantiators: ContextVar[Optional[InstantiatorsDictType]] = ContextVar("class_instantiators", default=None)
 nested_links: ContextVar[List[dict]] = ContextVar("nested_links", default=[])
+applied_instantiation_links: ContextVar[Optional[set]] = ContextVar("applied_instantiation_links", default=None)
 
 
 parser_context_vars = dict(
@@ -74,6 +75,7 @@ def __call__(self, class_type: Type[ClassType], *args, **kwargs) -> ClassType:
     load_value_mode=load_value_mode,
     class_instantiators=class_instantiators,
     nested_links=nested_links,
+    applied_instantiation_links=applied_instantiation_links,
 )
 
 
@@ -270,6 +272,12 @@ def __init__(self, instantiators: InstantiatorsDictType) -> None:
     def __call__(self, class_type: Type[ClassType], *args, **kwargs) -> ClassType:
         for (cls, subclasses), instantiator in self.instantiators.items():
             if class_type is cls or (subclasses and is_subclass(class_type, cls)):
+                param_names = set(inspect.signature(instantiator).parameters.keys())
+                if "applied_instantiation_links" in param_names:
+                    applied_links = applied_instantiation_links.get() or set()
+                    kwargs["applied_instantiation_links"] = {
+                        action.target[0]: action.applied_value for action in applied_links
+                    }
                 return instantiator(class_type, *args, **kwargs)
         return default_class_instantiator(class_type, *args, **kwargs)
 

jsonargparse/_core.py

@@ -1174,6 +1174,14 @@ def add_instantiator(
         For reference, the default instantiator is ``return class_type(*args,
         **kwargs)``.
 
+        In some use cases, the instantiator function might need access to values
+        applied by instantiation links. For this, the instantiator function can
+        have an additional keyword parameter ``applied_instantiation_links:
+        dict``. This parameter will be populated with a dictionary having as
+        keys the targets of the instantiation links and corresponding values
+        that were applied. Support for ``applied_instantiation_links`` parameter
+        is EXPERIMENTAL and subject to change or removal in future versions.
+
         Args:
             instantiator: Function that instantiates a class.
             class_type: The class type to instantiate.
@@ -1246,10 +1254,15 @@ def instantiate_classes(
                             parent_parser=self,
                             nested_links=ActionLink.get_nested_links(self, component),
                             class_instantiators=self._get_instantiators(),
+                            applied_instantiation_links=cfg.get("__applied_instantiation_links__"),
                         ):
                             parent[key] = component.instantiate_classes(value)
             else:
-                with parser_context(load_value_mode=self.parser_mode, class_instantiators=self._get_instantiators()):
+                with parser_context(
+                    load_value_mode=self.parser_mode,
+                    class_instantiators=self._get_instantiators(),
+                    applied_instantiation_links=cfg.get("__applied_instantiation_links__"),
+                ):
                     component.instantiate_class(component, cfg)
 
         ActionLink.apply_instantiation_links(self, cfg, order=order)

jsonargparse/_link_arguments.py

@@ -363,6 +363,7 @@ def apply_instantiation_links(parser, cfg, target=None, order=None):
             else:
                 value = action.call_compute_fn(source_objects)
             ActionLink.set_target_value(action, value, cfg, parser.logger)
+            action.applied_value = value
             applied_links.add(action)
             parser.logger.debug(f"Applied link '{action.option_strings[0]}'.")
 

jsonargparse_tests/test_link_arguments.py

@@ -900,6 +900,44 @@ def test_on_instantiate_targets_share_parent(parser):
     assert init.root.child.attr_child is init.source_b.attr_b
 
 
+class Dataloader:
+    def __init__(self, batch_size: int = 6):
+        self.batch_size = batch_size
+        self.num_classes = 7
+
+
+class CustomOptimizer(Optimizer):
+    def __init__(self, params: List[int], num_classes: int, **kwargs):
+        super().__init__(params, **kwargs)
+
+
+def custom_instantiator(class_type, *args, applied_instantiation_links: dict, **kwargs):
+    init = class_type(*args, **kwargs)
+    init.applied_instantiation_links = applied_instantiation_links
+    return init
+
+
+def test_on_instantiate_targets_passed_to_instantiator(parser):
+    parser.add_argument("--data", type=Dataloader)
+    parser.add_argument("--model", type=Model)
+    parser.link_arguments(
+        "data.num_classes",
+        "model.init_args.optimizer.init_args.num_classes",
+        apply_on="instantiate",
+    )
+    parser.add_instantiator(custom_instantiator, Dataloader, subclasses=True)
+    parser.add_instantiator(custom_instantiator, Model, subclasses=True)
+
+    cfg = parser.parse_args(["--data=Dataloader", "--model=Model", "--model.label=ok"])
+    init = parser.instantiate_classes(cfg)
+
+    assert isinstance(init.data, Dataloader)
+    assert init.data.applied_instantiation_links == {}
+    assert isinstance(init.model, Model)
+    assert callable(init.model.optimizer)
+    assert init.model.applied_instantiation_links == {"model.init_args.optimizer.init_args.num_classes": 7}
+
+
 # link creation failures