Upgrade examples to BioNeMo 2

[holgerroth]

Fixes # .

Description

Upgrade existing bionemo example to use version 2.4 (upcoming release)

Types of changes

• Non-breaking change (fix or new feature that would not break existing functionality).
• Breaking change (fix or new feature that would cause existing functionality to change).
• New tests added to cover the changes.
• Quick tests passed locally by running ./runtest.sh.
• In-line docstrings updated.
• Documentation updated.

[holgerroth]

/build

[farhadrgh]

Suggestion regarding accumulation TB history for the clients. Clients recording TB history locally from trainer:

from tensorboard.backend.event_processing.event_accumulator import EventAccumulator

def get_tensorboard_history(trainer):
    # Get the log directory from trainer's logger
    log_dir = trainer.logger.experiment.log_dir
    
    # Create an EventAccumulator object
    event_acc = EventAccumulator(log_dir)
    event_acc.Reload()  # Load all data
    
    # Get available tags (metrics)
    tags = event_acc.Tags()['scalars']
    
    # Extract history for each metric
    history = {}
    for tag in tags:
        events = event_acc.Scalars(tag)
        # Each event has attributes: wall_time, step, value
        history[tag] = [(e.step, e.value) for e in events]
    
    return history

And then after sharing through Shareable object you create a TB writer to accumulate them all:

tb_writer = SummaryWriter('continuous_logs')
for shareable in shareables:
    metrics_history = shareable['history']
    # Write each metric to the continuous log
    for metric_name, values in metrics_history.items():
        for step, value in values:
            global_step = (current_round * max_steps_per_round) + step
            writer.add_scalar(f'continuous/{metric_name}',  value, global_step=global_step)

[farhadrgh]

LGTM

[holgerroth]

TB accumulation will be addressed in a separate PR. Filed an issue for tracking #3249

[holgerroth]

/build

[holgerroth]

/build

[holgerroth]

/build

[holgerroth]

/build

[YuanTingHsieh]

LGTM

[holgerroth]

/build