Plugins

zenml.plugins

Modules

base_plugin_flavor

Base implementation for all Plugin Flavors.

Classes

BasePlugin

Bases: ABC

Base Class for all Plugins.

Attributes

config_class: Type[BasePluginConfig] abstractmethod property

Returns the BasePluginConfig config.

Returns:

Type	Description
Type[BasePluginConfig]	The configuration.

flavor_class: Type[BasePluginFlavor] abstractmethod property

Returns the flavor class of the plugin.

Returns:

Type	Description
Type[BasePluginFlavor]	The flavor class of the plugin.

zen_store: BaseZenStore property

Returns the active zen store.

Returns:

Type	Description
BaseZenStore	The active zen store.

BasePluginConfig

Bases: BaseModel, ABC

Allows configuring of Event Source and Filter configuration.

BasePluginFlavor

Bases: ABC

Base Class for all PluginFlavors.

Functions

get_flavor_response_model(hydrate: bool) -> BasePluginFlavorResponse[Any, Any, Any] abstractmethod classmethod

Convert the Flavor into a Response Model.

Parameters:

Name	Type	Description	Default
hydrate	bool	Whether the model should be hydrated.	required

Source code in src/zenml/plugins/base_plugin_flavor.py

@classmethod
@abstractmethod
def get_flavor_response_model(
    cls, hydrate: bool
) -> BasePluginFlavorResponse[Any, Any, Any]:
    """Convert the Flavor into a Response Model.

    Args:
        hydrate: Whether the model should be hydrated.
    """

plugin_flavor_registry

Registry for all plugins.

Classes

PluginFlavorRegistry()

Registry for plugin flavors.

Initialize the event flavor registry.

Source code in src/zenml/plugins/plugin_flavor_registry.py

def __init__(self) -> None:
    """Initialize the event flavor registry."""
    self.plugin_flavors: Dict[
        PluginType, Dict[PluginSubType, Dict[str, RegistryEntry]]
    ] = {}
    self.register_plugin_flavors()

Functions

get_flavor_class(_type: PluginType, subtype: PluginSubType, name: str) -> Type[BasePluginFlavor]

Get a single event_source based on the key.

Parameters:

Name	Type	Description	Default
_type	PluginType	The type of plugin.	required
subtype	PluginSubType	The subtype of plugin.	required
name	str	Indicates the name of the plugin flavor.	required

Returns:

Type	Description
Type[BasePluginFlavor]	BaseEventConfiguration subclass that was registered for this key.

Raises:

Type	Description
KeyError	If there is no entry at this type, subtype, flavor

Source code in src/zenml/plugins/plugin_flavor_registry.py

def get_flavor_class(
    self,
    _type: PluginType,
    subtype: PluginSubType,
    name: str,
) -> Type[BasePluginFlavor]:
    """Get a single event_source based on the key.

    Args:
        _type: The type of plugin.
        subtype: The subtype of plugin.
        name: Indicates the name of the plugin flavor.

    Returns:
        `BaseEventConfiguration` subclass that was registered for this key.

    Raises:
        KeyError: If there is no entry at this type, subtype, flavor
    """
    try:
        return self._get_registry_entry(
            _type=_type,
            subtype=subtype,
            flavor_name=name,
        ).flavor_class
    except KeyError:
        raise KeyError(
            f"No flavor class found for flavor name {name} at type "
            f"{_type} and subtype {subtype}."
        )

get_plugin(_type: PluginType, subtype: PluginSubType, name: str) -> BasePlugin

Get the plugin based on the flavor, type and subtype.

Parameters:

Name	Type	Description	Default
name	str	The name of the plugin flavor.	required
_type	PluginType	The type of plugin.	required
subtype	PluginSubType	The subtype of plugin.	required

Returns:

Type	Description
BasePlugin	Plugin instance associated with the flavor, type and subtype.

Raises:

Type	Description
KeyError	If no plugin is found for the given flavor, type and subtype.
RuntimeError	If the plugin was not initialized.

Source code in src/zenml/plugins/plugin_flavor_registry.py

def get_plugin(
    self,
    _type: PluginType,
    subtype: PluginSubType,
    name: str,
) -> "BasePlugin":
    """Get the plugin based on the flavor, type and subtype.

    Args:
        name: The name of the plugin flavor.
        _type: The type of plugin.
        subtype: The subtype of plugin.

    Returns:
        Plugin instance associated with the flavor, type and subtype.

    Raises:
        KeyError: If no plugin is found for the given flavor, type and
            subtype.
        RuntimeError: If the plugin was not initialized.
    """
    try:
        plugin_entry = self._get_registry_entry(
            _type=_type,
            subtype=subtype,
            flavor_name=name,
        )
        if plugin_entry.plugin_instance is None:
            raise RuntimeError(
                f"Plugin {plugin_entry.flavor_class} was not initialized."
            )
        return plugin_entry.plugin_instance
    except KeyError:
        raise KeyError(
            f"No flavor found for flavor name {name} and type "
            f"{_type} and subtype {subtype}."
        )

initialize_plugins() -> None

Initializes all registered plugins.

Source code in src/zenml/plugins/plugin_flavor_registry.py

def initialize_plugins(self) -> None:
    """Initializes all registered plugins."""
    for _, subtypes_dict in self.plugin_flavors.items():
        for _, flavor_dict in subtypes_dict.items():
            for _, registry_entry in flavor_dict.items():
                # TODO: Only initialize if the integration is active
                registry_entry.plugin_instance = (
                    registry_entry.flavor_class.PLUGIN_CLASS()
                )

list_available_flavor_responses_for_type_and_subtype(_type: PluginType, subtype: PluginSubType, page: int, size: int, hydrate: bool = False) -> Page[BasePluginFlavorResponse[Any, Any, Any]]

Get a list of all subtypes for a specific flavor and type.

Parameters:

Name	Type	Description	Default
_type	PluginType	The type of Plugin	required
subtype	PluginSubType	The subtype of the plugin	required
page	int	Page for pagination (offset +1)	required
size	int	Page size for pagination	required
hydrate	bool	Whether to hydrate the response bodies	False

Returns:

Type	Description
Page[BasePluginFlavorResponse[Any, Any, Any]]	A page of flavors.

Raises:

Type	Description
ValueError	If the page is out of range.

Source code in src/zenml/plugins/plugin_flavor_registry.py

def list_available_flavor_responses_for_type_and_subtype(
    self,
    _type: PluginType,
    subtype: PluginSubType,
    page: int,
    size: int,
    hydrate: bool = False,
) -> Page["BasePluginFlavorResponse[Any, Any, Any]"]:
    """Get a list of all subtypes for a specific flavor and type.

    Args:
        _type: The type of Plugin
        subtype: The subtype of the plugin
        page: Page for pagination (offset +1)
        size: Page size for pagination
        hydrate: Whether to hydrate the response bodies

    Returns:
        A page of flavors.

    Raises:
        ValueError: If the page is out of range.
    """
    flavors = self.list_available_flavors_for_type_and_subtype(
        _type=_type,
        subtype=subtype,
    )
    total = len(flavors)
    if total == 0:
        total_pages = 1
    else:
        total_pages = math.ceil(total / size)

    if page > total_pages:
        raise ValueError(
            f"Invalid page {page}. The requested page size is "
            f"{size} and there are a total of {total} items "
            f"for this query. The maximum page value therefore is "
            f"{total_pages}."
        )

    start = (page - 1) * size
    end = start + size

    page_items = [
        flavor.get_flavor_response_model(hydrate=hydrate)
        for flavor in flavors[start:end]
    ]

    return Page(
        index=page,
        max_size=size,
        total_pages=total_pages,
        total=total,
        items=page_items,
    )

list_available_flavors_for_type_and_subtype(_type: PluginType, subtype: PluginSubType) -> List[Type[BasePluginFlavor]]

Get a list of all subtypes for a specific flavor and type.

Parameters:

Name	Type	Description	Default
_type	PluginType	The type of Plugin	required
subtype	PluginSubType	The subtype of the plugin	required

Returns:

Type	Description
List[Type[BasePluginFlavor]]	List of flavors for the given type/subtype combination.

Source code in src/zenml/plugins/plugin_flavor_registry.py

def list_available_flavors_for_type_and_subtype(
    self,
    _type: PluginType,
    subtype: PluginSubType,
) -> List[Type[BasePluginFlavor]]:
    """Get a list of all subtypes for a specific flavor and type.

    Args:
        _type: The type of Plugin
        subtype: The subtype of the plugin

    Returns:
        List of flavors for the given type/subtype combination.
    """
    flavors = list(
        [
            entry.flavor_class
            for _, entry in self._flavor_entries(
                _type=_type, subtype=subtype
            ).items()
        ]
    )
    return flavors

list_subtypes_within_type(_type: PluginType) -> List[PluginSubType]

Returns all available subtypes for a given type.

Parameters:

Name	Type	Description	Default
_type	PluginType	The type of plugin	required

Returns:

Type	Description
List[PluginSubType]	A list of available plugin subtypes for this plugin type.

Source code in src/zenml/plugins/plugin_flavor_registry.py

def list_subtypes_within_type(
    self, _type: PluginType
) -> List[PluginSubType]:
    """Returns all available subtypes for a given type.

    Args:
        _type: The type of plugin

    Returns:
        A list of available plugin subtypes for this plugin type.
    """
    return list(self.plugin_flavors[_type].keys())

register_plugin_flavor(flavor_class: Type[BasePluginFlavor]) -> None

Registers a new event_source.

Parameters:

Name	Type	Description	Default
flavor_class	Type[BasePluginFlavor]	The flavor to register	required

Source code in src/zenml/plugins/plugin_flavor_registry.py

def register_plugin_flavor(
    self, flavor_class: Type[BasePluginFlavor]
) -> None:
    """Registers a new event_source.

    Args:
        flavor_class: The flavor to register
    """
    try:
        self._get_registry_entry(
            _type=flavor_class.TYPE,
            subtype=flavor_class.SUBTYPE,
            flavor_name=flavor_class.FLAVOR,
        )
    except KeyError:
        (
            self.plugin_flavors.setdefault(flavor_class.TYPE, {})
            .setdefault(flavor_class.SUBTYPE, {})
            .setdefault(
                flavor_class.FLAVOR,
                RegistryEntry(flavor_class=flavor_class),
            )
        )
        logger.debug(
            f"Registered built in plugin {flavor_class.FLAVOR} for "
            f"plugin type {flavor_class.TYPE} and "
            f"subtype {flavor_class.SUBTYPE}: {flavor_class}"
        )
    else:
        logger.debug(
            f"Found existing flavor {flavor_class.FLAVOR} already "
            f"registered for this type {flavor_class.TYPE} and subtype "
            f"{flavor_class.SUBTYPE}. "
            f"Skipping registration for {flavor_class}."
        )

register_plugin_flavors() -> None

Registers all flavors.

Source code in src/zenml/plugins/plugin_flavor_registry.py

def register_plugin_flavors(self) -> None:
    """Registers all flavors."""
    for flavor in self._builtin_flavors:
        self.register_plugin_flavor(flavor_class=flavor)
    for flavor in self._integration_flavors:
        self.register_plugin_flavor(flavor_class=flavor)

RegistryEntry

Bases: BaseModel

Registry Entry Class for the Plugin Registry.

Functions