Skip to content

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
79
80
81
82
83
84
85
86
87
88
@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
45
46
47
48
49
50
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
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
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
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
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
334
335
336
337
338
339
340
341
342
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
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
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
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
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
61
62
63
64
65
66
67
68
69
70
71
72
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
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
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
222
223
224
225
226
227
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