nvflare.collab.api.module_wrapper module

Wrapper to use a module’s functions as Collab publish/main methods.

class ModuleWrapper(module: ModuleType | str | None = None)

Bases: object

Wrap a module’s publish, main, init, and final functions for Collab.

Collab creates this wrapper automatically whenever a module is supplied as a primary or named app object. This allows callers to pass modules containing standalone functions anywhere they could pass class instances:

# my_module.py from nvflare.collab import collab

@collab.publish def train(weights=None):

…

@collab.main def fed_avg():

…

# main.py import my_module

recipe = CollabRecipe(server=my_module, client=my_module, …)

For a distributed deployment, the module must be importable on all client machines (i.e., part of the installed package or included in job resources).

Initialize wrapper with a module containing @collab.publish/@collab.main functions.

Parameters:

module – A Python module object OR a fully qualified module name string. When a string is passed, the module will be imported. When no argument is passed (None), the wrapper is in an uninitialized state and will be set up by __setstate__ during unpickling.

Note

For job config serialization, we store the module name as self._module which matches the ‘module’ parameter. FLARE’s _get_args() looks for param or _param in __dict__, so _module matches ‘module’.

When running as __main__, we convert to an importable module path so that remote processes can import the same module.

property module_name: str

Get the module name.

get_caller_module(stack_level: int = 2) → ModuleType | None

Get the module of the caller at the requested stack level.

get_importable_module_name(module: ModuleType) → str

Get an importable module name, handling __main__ case.

When a script is run as python script.py, its __name__ is ‘__main__’, which cannot be imported on remote machines. This function converts ‘__main__’ to the actual importable module path based on the file location.

Parameters:

module – A Python module object

Returns:

An importable module name string

Example

# When running: python path/to/my_module.py # module.__name__ = ‘__main__’ # Returns: ‘my_pkg.my_module’ (or ‘my_module’)

resolve_server_client(server, client, stack_level: int = 2)

Use the caller’s module for a missing server or client object.

wrap_if_module(obj)

Return a ModuleWrapper for a module, or leave any other object unchanged.