allegroai / clearml

ClearML - Auto-Magical CI/CD to streamline your AI workload. Experiment Management, Data Management, Pipeline, Orchestration, Scheduling & Serving in one MLOps/LLMOps solution
https://clear.ml/docs
Apache License 2.0
5.43k stars 643 forks source link

Fix missing component callbacks on multiple step calls #1195

Closed materight closed 5 months ago

materight commented 5 months ago

Patch Description

When a pipeline step is called multiple times, if a pre_execute_callback, pre_execute_callback or status_change_callback is set, it is called only the first time.

Testing Instructions

Minimal example:

from clearml import PipelineDecorator

def pre_callaback(pipeline, node, parameters):
    print(node.name)

@PipelineDecorator.component(name='step', pre_execute_callback=pre_callaback)
def step(i: int):
    return i + 1

@PipelineDecorator.pipeline(name='test', project='test')
def pipeline():
    for i in range(5):
        step(i)

if __name__ == '__main__':
    PipelineDecorator.run_locally()
    pipeline()

This will print step only once, while I would expect to see it 5 times.

With this patch, the callbacks are copied to each node replica, and the output shows:

step
step_1
step_2
step_3
step_4