Upgrade Lightning 2.5.1 → 2.6.1

[Copilot]

Lightning ≥2.5.2 introduced stricter jsonargparse validation that breaks link_arguments with apply_on="instantiate". Validation now runs before linked values are applied, causing Union type errors for dynamically computed fields.

Changes

Migration to before_instantiate_classes() hook pattern:

• Instantiate datamodule early to compute num_labels and feature_vector_size before validation
• Update config directly in hook instead of using compute_fn + apply_on="instantiate"
• Simplified link_arguments to link from pre-populated config values

Code structure improvements:

• Extracted helper methods: _update_model_args(), _update_metrics_num_labels(), _update_trainer_callbacks(), _set_callback_num_labels()
• Added documentation explaining the pattern change

Before (2.5.1):

def add_arguments_to_parser(self, parser):
    def call_data_methods(data):
        if data._num_of_labels is None:
            data.prepare_data()
            data.setup()
        return data.num_of_labels
    
    parser.link_arguments(
        "data",
        "model.init_args.out_dim",
        apply_on="instantiate",
        compute_fn=call_data_methods,
    )

After (2.6.1):

def before_instantiate_classes(self):
    data_instance = self._instantiate_datamodule()
    num_labels = data_instance.num_of_labels
    subcommand["model"]["init_args"]["out_dim"] = num_labels

def add_arguments_to_parser(self, parser):
    parser.link_arguments(
        "model.init_args.out_dim",
        "model.init_args.train_metrics.init_args.metrics.*.init_args.num_labels",
    )

Testing

• Existing CLI test passes with Lightning 2.6.1
• No security vulnerabilities (gh-advisory-database, CodeQL)

Original prompt

---

This section details on the original issue you should resolve

<issue_title>Upgrade Lightning from 2.5.1 → latest</issue_title>
<issue_description>

Title:

Description:
Currently the project is pinned to PyTorch Lightning 2.5.1 in f74c8ac because our LightningCLI-based integration fails on 2.5.2 and newer.

Problem

• In Lightning ≥2.5.2, LightningCLI and its underlying jsonargparse validation became stricter.

• Fields that were previously injected dynamically through parser.link_arguments (e.g. num_labels) now must either exist in the config file or be linked at an earlier stage.

• This causes errors like:

  Does not validate against any of the Union subtypes
  Subtypes: [<class 'NoneType'>, <class 'lightning.pytorch.core.module.LightningModule'>]
  Problem with given class_path ...
  Expected key train_metrics.init_args.metrics.macro-f1.init_args.num_labels to be in data
  

• In 2.5.1 this worked fine, but 2.5.2+ fails during config parsing.

</issue_description>

Comments on the Issue (you are @copilot in this section)

@sfluegel05 not directly related, but also concerning the lightning version:

On feature/new-ensemble-models, I added a custom Fit loop that does not run with lightning 2.1, but does run with 2.5.1 (I haven't tried other versions). </comment_new>

• Fixes Upgrade Lightning from 2.5.1 → latest #124

---

💬 We'd love your input! Share your thoughts on Copilot coding agent in our 2 minute survey.