[FEA]: Update DFPTraining stage to accept both ControlMessage and MultiDFPMessage

[mdemoret-nv]

Is this a new feature, an improvement, or a change to existing functionality?

New Feature

How would you describe the priority of this feature request

High

Please provide a clear description of problem this feature solves

Currently, the DFP pipeline is implemented using both Modules and Stages. The modules exclusively work with ControlMessage while the stages work exclusively with MultiDFPMessage. This makes it difficult for users customizing the DFP pipeline to use some of the advantages of ControlMessage if they are using stages. One particular area where ControlMessage would be more beneficial is the training stage to allow using data loaders.

Describe your ideal solution

As a first step, the DFPTraining stage should be updated to accept both ControlMessage and MultiDFPMessage and work with either. If it accepts a ControlMessage it should return a ControlMessage and the same for MultiDFPMessage.

Completion Criteria:

• The DFPTraining accepts both ControlMessage and MultiDFPMessage
  • The type can be determined in the build phase and remain static
• The DFPTraining returns a ControlMessage if one is passed in or a MultiAEMessage if a DFPMessageMeta is passed in
• The DFPTraining stage should always pass the data to AutoEncoder as a data loader (even if its a MultiDFPMessage)
• The training results should be unchanged

Describe any alternatives you have considered

No response

Additional context

No response

Code of Conduct

• I agree to follow this project's Code of Conduct
• I have searched the open feature requests and have found no duplicates for this feature request

[mdemoret-nv]

Additional context, using a Data loader for AutoEncoder should help with issues: #933, #816 and #724

[drobison00]

Supporting torch datasets on the _fit_centralized path requires a variety of changes; after looking over how we do _fit_distribtued, there are enough similarities in what we're doing and what needs to change in _fit_centralized that it makes sense to combine the two paths. Going to work through this in two phases, one to align _fit_centralized's behavior with _fit_distributed and one to merge the methods and eliminate duplicated code.