Skip to content

Registry

Implementation of a registry to track ZenML integrations.

IntegrationRegistry

Bases: object

Registry to keep track of ZenML Integrations.

Source code in src/zenml/integrations/registry.py 
 27
 28
 29
 30
 31
 32
 33
 34
 35
 36
 37
 38
 39
 40
 41
 42
 43
 44
 45
 46
 47
 48
 49
 50
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
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
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
class IntegrationRegistry(object):
    """Registry to keep track of ZenML Integrations."""

    def __init__(self) -> None:
        """Initializing the integration registry."""
        self._integrations: Dict[str, Type["Integration"]] = {}

    @property
    def integrations(self) -> Dict[str, Type["Integration"]]:
        """Method to get integrations dictionary.

        Returns:
            A dict of integration key to type of `Integration`.
        """
        return self._integrations

    @integrations.setter
    def integrations(self, i: Any) -> None:
        """Setter method for the integrations property.

        Args:
            i: Value to set the integrations property to.

        Raises:
            IntegrationError: If you try to manually set the integrations property.
        """
        raise IntegrationError(
            "Please do not manually change the integrations within the "
            "registry. If you would like to register a new integration "
            "manually, please use "
            "`integration_registry.register_integration()`."
        )

    def register_integration(
        self, key: str, type_: Type["Integration"]
    ) -> None:
        """Method to register an integration with a given name.

        Args:
            key: Name of the integration.
            type_: Type of the integration.
        """
        self._integrations[key] = type_

    def activate_integrations(self) -> None:
        """Method to activate the integrations with are registered in the registry."""
        for name, integration in self._integrations.items():
            if integration.check_installation():
                logger.debug(f"Activating integration `{name}`...")
                integration.activate()
                logger.debug(f"Integration `{name}` is activated.")
            else:
                logger.debug(f"Integration `{name}` could not be activated.")

    @property
    def list_integration_names(self) -> List[str]:
        """Get a list of all possible integrations.

        Returns:
            A list of all possible integrations.
        """
        return [name for name in self._integrations]

    def select_integration_requirements(
        self,
        integration_name: Optional[str] = None,
        target_os: Optional[str] = None,
    ) -> List[str]:
        """Select the requirements for a given integration or all integrations.

        Args:
            integration_name: Name of the integration to check.
            target_os: Target OS for the requirements.

        Returns:
            List of requirements for the integration.

        Raises:
            KeyError: If the integration is not found.
        """
        if integration_name:
            if integration_name in self.list_integration_names:
                return self._integrations[integration_name].get_requirements(
                    target_os=target_os
                )
            else:
                raise KeyError(
                    f"Integration {integration_name} does not exist. "
                    f"Currently the following integrations are implemented. "
                    f"{self.list_integration_names}"
                )
        else:
            return [
                requirement
                for name in self.list_integration_names
                for requirement in self._integrations[name].get_requirements(
                    target_os=target_os
                )
            ]

    def select_uninstall_requirements(
        self,
        integration_name: Optional[str] = None,
        target_os: Optional[str] = None,
    ) -> List[str]:
        """Select the uninstall requirements for a given integration or all integrations.

        Args:
            integration_name: Name of the integration to check.
            target_os: Target OS for the requirements.

        Returns:
            List of requirements for the integration uninstall.

        Raises:
            KeyError: If the integration is not found.
        """
        if integration_name:
            if integration_name in self.list_integration_names:
                return self._integrations[
                    integration_name
                ].get_uninstall_requirements(target_os=target_os)
            else:
                raise KeyError(
                    f"Integration {integration_name} does not exist. "
                    f"Currently the following integrations are implemented. "
                    f"{self.list_integration_names}"
                )
        else:
            return [
                requirement
                for name in self.list_integration_names
                for requirement in self._integrations[
                    name
                ].get_uninstall_requirements(target_os=target_os)
            ]

    def is_installed(self, integration_name: Optional[str] = None) -> bool:
        """Checks if all requirements for an integration are installed.

        Args:
            integration_name: Name of the integration to check.

        Returns:
            True if all requirements are installed, False otherwise.

        Raises:
            KeyError: If the integration is not found.
        """
        if integration_name in self.list_integration_names:
            return self._integrations[integration_name].check_installation()
        elif not integration_name:
            all_installed = [
                self._integrations[item].check_installation()
                for item in self.list_integration_names
            ]
            return all(all_installed)
        else:
            raise KeyError(
                f"Integration '{integration_name}' not found. "
                f"Currently the following integrations are available: "
                f"{self.list_integration_names}"
            )

    def get_installed_integrations(self) -> List[str]:
        """Returns list of installed integrations.

        Returns:
            List of installed integrations.
        """
        return [
            name
            for name, integration in integration_registry.integrations.items()
            if integration.check_installation()
        ]

integrations property writable

Method to get integrations dictionary.

Returns:

Type Description
Dict[str, Type[Integration]]

A dict of integration key to type of Integration.

list_integration_names property

Get a list of all possible integrations.

Returns:

Type Description
List[str]

A list of all possible integrations.

__init__()

Initializing the integration registry.

Source code in src/zenml/integrations/registry.py 
30
31
32
def __init__(self) -> None:
    """Initializing the integration registry."""
    self._integrations: Dict[str, Type["Integration"]] = {}

activate_integrations()

Method to activate the integrations with are registered in the registry.

Source code in src/zenml/integrations/registry.py 
71
72
73
74
75
76
77
78
79
def activate_integrations(self) -> None:
    """Method to activate the integrations with are registered in the registry."""
    for name, integration in self._integrations.items():
        if integration.check_installation():
            logger.debug(f"Activating integration `{name}`...")
            integration.activate()
            logger.debug(f"Integration `{name}` is activated.")
        else:
            logger.debug(f"Integration `{name}` could not be activated.")

get_installed_integrations()

Returns list of installed integrations.

Returns:

Type Description
List[str]

List of installed integrations.
Source code in src/zenml/integrations/registry.py 
191
192
193
194
195
196
197
198
199
200
201
def get_installed_integrations(self) -> List[str]:
    """Returns list of installed integrations.

    Returns:
        List of installed integrations.
    """
    return [
        name
        for name, integration in integration_registry.integrations.items()
        if integration.check_installation()
    ]

is_installed(integration_name=None)

Checks if all requirements for an integration are installed.

Parameters:

Name Type Description Default
integration_name Optional[str]

Name of the integration to check.

 None

Returns:

Type Description
bool

True if all requirements are installed, False otherwise.

Raises:

Type Description
KeyError

If the integration is not found.
Source code in src/zenml/integrations/registry.py 
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
def is_installed(self, integration_name: Optional[str] = None) -> bool:
    """Checks if all requirements for an integration are installed.

    Args:
        integration_name: Name of the integration to check.

    Returns:
        True if all requirements are installed, False otherwise.

    Raises:
        KeyError: If the integration is not found.
    """
    if integration_name in self.list_integration_names:
        return self._integrations[integration_name].check_installation()
    elif not integration_name:
        all_installed = [
            self._integrations[item].check_installation()
            for item in self.list_integration_names
        ]
        return all(all_installed)
    else:
        raise KeyError(
            f"Integration '{integration_name}' not found. "
            f"Currently the following integrations are available: "
            f"{self.list_integration_names}"
        )

register_integration(key, type_)

Method to register an integration with a given name.

Parameters:

Name Type Description Default
key str

Name of the integration.

 required
type_ Type[Integration]

Type of the integration.

 required
Source code in src/zenml/integrations/registry.py 
60
61
62
63
64
65
66
67
68
69
def register_integration(
    self, key: str, type_: Type["Integration"]
) -> None:
    """Method to register an integration with a given name.

    Args:
        key: Name of the integration.
        type_: Type of the integration.
    """
    self._integrations[key] = type_

select_integration_requirements(integration_name=None, target_os=None)

Select the requirements for a given integration or all integrations.

Parameters:

Name Type Description Default
integration_name Optional[str]

Name of the integration to check.

 None
target_os Optional[str]

Target OS for the requirements.

 None

Returns:

Type Description
List[str]

List of requirements for the integration.

Raises:

Type Description
KeyError

If the integration is not found.
Source code in src/zenml/integrations/registry.py 
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
def select_integration_requirements(
    self,
    integration_name: Optional[str] = None,
    target_os: Optional[str] = None,
) -> List[str]:
    """Select the requirements for a given integration or all integrations.

    Args:
        integration_name: Name of the integration to check.
        target_os: Target OS for the requirements.

    Returns:
        List of requirements for the integration.

    Raises:
        KeyError: If the integration is not found.
    """
    if integration_name:
        if integration_name in self.list_integration_names:
            return self._integrations[integration_name].get_requirements(
                target_os=target_os
            )
        else:
            raise KeyError(
                f"Integration {integration_name} does not exist. "
                f"Currently the following integrations are implemented. "
                f"{self.list_integration_names}"
            )
    else:
        return [
            requirement
            for name in self.list_integration_names
            for requirement in self._integrations[name].get_requirements(
                target_os=target_os
            )
        ]

select_uninstall_requirements(integration_name=None, target_os=None)

Select the uninstall requirements for a given integration or all integrations.

Parameters:

Name Type Description Default
integration_name Optional[str]

Name of the integration to check.

 None
target_os Optional[str]

Target OS for the requirements.

 None

Returns:

Type Description
List[str]

List of requirements for the integration uninstall.

Raises:

Type Description
KeyError

If the integration is not found.
Source code in src/zenml/integrations/registry.py 
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
def select_uninstall_requirements(
    self,
    integration_name: Optional[str] = None,
    target_os: Optional[str] = None,
) -> List[str]:
    """Select the uninstall requirements for a given integration or all integrations.

    Args:
        integration_name: Name of the integration to check.
        target_os: Target OS for the requirements.

    Returns:
        List of requirements for the integration uninstall.

    Raises:
        KeyError: If the integration is not found.
    """
    if integration_name:
        if integration_name in self.list_integration_names:
            return self._integrations[
                integration_name
            ].get_uninstall_requirements(target_os=target_os)
        else:
            raise KeyError(
                f"Integration {integration_name} does not exist. "
                f"Currently the following integrations are implemented. "
                f"{self.list_integration_names}"
            )
    else:
        return [
            requirement
            for name in self.list_integration_names
            for requirement in self._integrations[
                name
            ].get_uninstall_requirements(target_os=target_os)
        ]