Skip to content

Integration

zenml.integrations.integration

Base and meta classes for ZenML integrations.

Integration

Base class for integration in ZenML.

Source code in zenml/integrations/integration.py
class Integration(metaclass=IntegrationMeta):
    """Base class for integration in ZenML."""

    NAME = "base_integration"

    REQUIREMENTS: List[str] = []
    APT_PACKAGES: List[str] = []

    @classmethod
    def check_installation(cls) -> bool:
        """Method to check whether the required packages are installed.

        Returns:
            True if all required packages are installed, False otherwise.
        """
        try:
            for r in cls.get_requirements():
                name, extras = parse_requirement(r)
                if name:
                    dist = pkg_resources.get_distribution(name)
                    # Check if extras are specified and installed
                    if extras:
                        extra_list = extras[1:-1].split(",")
                        # for extra in extra_list:
                        #     try:
                        #         requirements = dist.requires(extras=[extra])  # type: ignore[arg-type]
                        #     except pkg_resources.UnknownExtra as e:
                        #         logger.debug("Unknown extra: " + str(e))
                        #         return False
                        #
                        #     for ri in requirements:
                        #         try:
                        #             pkg_resources.get_distribution(ri)
                        #         except IndexError:
                        #             logger.debug(
                        #                 f"Unable to find required extra '{ri.project_name}' for "
                        #                 f"requirements '{name}' coming from integration '{cls.NAME}'."
                        #             )
                        #             return False
                else:
                    logger.debug(
                        f"Invalid requirement format '{r}' for integration {cls.NAME}."
                    )
                    return False

            logger.debug(
                f"Integration {cls.NAME} is installed correctly with "
                f"requirements {cls.get_requirements()}."
            )
            return True
        except pkg_resources.DistributionNotFound as e:
            logger.debug(
                f"Unable to find required package '{e.req}' for "
                f"integration {cls.NAME}."
            )
            return False
        except pkg_resources.VersionConflict as e:
            logger.debug(
                f"VersionConflict error when loading installation {cls.NAME}: "
                f"{str(e)}"
            )
            return False

    @classmethod
    def get_requirements(cls, target_os: Optional[str] = None) -> List[str]:
        """Method to get the requirements for the integration.

        Args:
            target_os: The target operating system to get the requirements for.

        Returns:
            A list of requirements.
        """
        return cls.REQUIREMENTS

    @classmethod
    def activate(cls) -> None:
        """Abstract method to activate the integration."""

    @classmethod
    def flavors(cls) -> List[Type[Flavor]]:
        """Abstract method to declare new stack component flavors.

        Returns:
            A list of new stack component flavors.
        """
        return []

activate() classmethod

Abstract method to activate the integration.

Source code in zenml/integrations/integration.py
@classmethod
def activate(cls) -> None:
    """Abstract method to activate the integration."""

check_installation() classmethod

Method to check whether the required packages are installed.

Returns:

Type Description
bool

True if all required packages are installed, False otherwise.

Source code in zenml/integrations/integration.py
@classmethod
def check_installation(cls) -> bool:
    """Method to check whether the required packages are installed.

    Returns:
        True if all required packages are installed, False otherwise.
    """
    try:
        for r in cls.get_requirements():
            name, extras = parse_requirement(r)
            if name:
                dist = pkg_resources.get_distribution(name)
                # Check if extras are specified and installed
                if extras:
                    extra_list = extras[1:-1].split(",")
                    # for extra in extra_list:
                    #     try:
                    #         requirements = dist.requires(extras=[extra])  # type: ignore[arg-type]
                    #     except pkg_resources.UnknownExtra as e:
                    #         logger.debug("Unknown extra: " + str(e))
                    #         return False
                    #
                    #     for ri in requirements:
                    #         try:
                    #             pkg_resources.get_distribution(ri)
                    #         except IndexError:
                    #             logger.debug(
                    #                 f"Unable to find required extra '{ri.project_name}' for "
                    #                 f"requirements '{name}' coming from integration '{cls.NAME}'."
                    #             )
                    #             return False
            else:
                logger.debug(
                    f"Invalid requirement format '{r}' for integration {cls.NAME}."
                )
                return False

        logger.debug(
            f"Integration {cls.NAME} is installed correctly with "
            f"requirements {cls.get_requirements()}."
        )
        return True
    except pkg_resources.DistributionNotFound as e:
        logger.debug(
            f"Unable to find required package '{e.req}' for "
            f"integration {cls.NAME}."
        )
        return False
    except pkg_resources.VersionConflict as e:
        logger.debug(
            f"VersionConflict error when loading installation {cls.NAME}: "
            f"{str(e)}"
        )
        return False

flavors() classmethod

Abstract method to declare new stack component flavors.

Returns:

Type Description
List[Type[zenml.stack.flavor.Flavor]]

A list of new stack component flavors.

Source code in zenml/integrations/integration.py
@classmethod
def flavors(cls) -> List[Type[Flavor]]:
    """Abstract method to declare new stack component flavors.

    Returns:
        A list of new stack component flavors.
    """
    return []

get_requirements(target_os=None) classmethod

Method to get the requirements for the integration.

Parameters:

Name Type Description Default
target_os Optional[str]

The target operating system to get the requirements for.

None

Returns:

Type Description
List[str]

A list of requirements.

Source code in zenml/integrations/integration.py
@classmethod
def get_requirements(cls, target_os: Optional[str] = None) -> List[str]:
    """Method to get the requirements for the integration.

    Args:
        target_os: The target operating system to get the requirements for.

    Returns:
        A list of requirements.
    """
    return cls.REQUIREMENTS

IntegrationMeta (type)

Metaclass responsible for registering different Integration subclasses.

Source code in zenml/integrations/integration.py
class IntegrationMeta(type):
    """Metaclass responsible for registering different Integration subclasses."""

    def __new__(
        mcs, name: str, bases: Tuple[Type[Any], ...], dct: Dict[str, Any]
    ) -> "IntegrationMeta":
        """Hook into creation of an Integration class.

        Args:
            name: The name of the class being created.
            bases: The base classes of the class being created.
            dct: The dictionary of attributes of the class being created.

        Returns:
            The newly created class.
        """
        cls = cast(Type["Integration"], super().__new__(mcs, name, bases, dct))
        if name != "Integration":
            integration_registry.register_integration(cls.NAME, cls)
        return cls

__new__(mcs, name, bases, dct) special staticmethod

Hook into creation of an Integration class.

Parameters:

Name Type Description Default
name str

The name of the class being created.

required
bases Tuple[Type[Any], ...]

The base classes of the class being created.

required
dct Dict[str, Any]

The dictionary of attributes of the class being created.

required

Returns:

Type Description
IntegrationMeta

The newly created class.

Source code in zenml/integrations/integration.py
def __new__(
    mcs, name: str, bases: Tuple[Type[Any], ...], dct: Dict[str, Any]
) -> "IntegrationMeta":
    """Hook into creation of an Integration class.

    Args:
        name: The name of the class being created.
        bases: The base classes of the class being created.
        dct: The dictionary of attributes of the class being created.

    Returns:
        The newly created class.
    """
    cls = cast(Type["Integration"], super().__new__(mcs, name, bases, dct))
    if name != "Integration":
        integration_registry.register_integration(cls.NAME, cls)
    return cls