tessl/pypi-zenml

ZenML is a unified MLOps framework that extends battle-tested machine learning operations principles to support the entire AI stack, from classical machine learning models to advanced AI agents.

Overview

Eval results

Files

Stack Components

Name: tessl/pypi-zenml
Author: tessl

Base classes and implementations for ZenML stack components. Stack components are the building blocks of a ZenML stack, each providing specific infrastructure functionality.

Component Types

ZenML provides 13 types of stack components, each serving a specific purpose in the ML infrastructure.

Orchestrators

Orchestrators manage the execution of pipeline steps.

class BaseOrchestrator:
    """
    Abstract base class for orchestrators.

    Orchestrators control where and how pipeline steps are executed.
    """

class LocalOrchestrator(BaseOrchestrator):
    """Orchestrator for local execution."""

class LocalDockerOrchestrator(BaseOrchestrator):
    """Orchestrator for local Docker execution."""

Import from:

from zenml.orchestrators import (
    BaseOrchestrator,
    LocalOrchestrator,
    LocalDockerOrchestrator
)

Artifact Stores

Artifact stores handle storage and retrieval of artifacts.

class BaseArtifactStore:
    """
    Abstract base class for artifact stores.

    Provides filesystem-like interface for artifact storage.
    """

class LocalArtifactStore(BaseArtifactStore):
    """Artifact store for local filesystem."""

Import from:

from zenml.artifact_stores import (
    BaseArtifactStore,
    LocalArtifactStore
)

Container Registries

Container registries store Docker images.

class BaseContainerRegistry:
    """Abstract base class for container registries."""

class DefaultContainerRegistryFlavor:
    """Default/generic container registry flavor."""

class AzureContainerRegistryFlavor:
    """Flavor for Azure Container Registry."""

class DockerHubContainerRegistryFlavor:
    """Flavor for Docker Hub."""

class GCPContainerRegistryFlavor:
    """Flavor for Google Container Registry."""

class GitHubContainerRegistryFlavor:
    """Flavor for GitHub Container Registry."""

Import from:

from zenml.container_registries import (
    BaseContainerRegistry,
    DefaultContainerRegistryFlavor
)

Image Builders

Image builders create container images for pipeline execution.

class BaseImageBuilder:
    """Abstract base class for image builders."""

class LocalImageBuilder(BaseImageBuilder):
    """Image builder for local Docker builds."""

class BuildContext:
    """Context information for building container images."""

Import from:

from zenml.image_builders import (
    BaseImageBuilder,
    LocalImageBuilder,
    BuildContext
)

Step Operators

Step operators enable remote execution of individual steps.

class BaseStepOperator:
    """
    Abstract base class for step operators.

    Step operators execute specific steps on remote infrastructure
    (e.g., SageMaker, Vertex AI) while the orchestrator manages
    the overall pipeline flow.
    """

Import from:

from zenml.step_operators import BaseStepOperator

Experiment Trackers

Experiment trackers log metrics, parameters, and artifacts.

class BaseExperimentTracker:
    """
    Abstract base class for experiment trackers.

    Integrates with experiment tracking platforms like MLflow,
    Weights & Biases, Neptune, etc.
    """

Import from:

from zenml.experiment_trackers import BaseExperimentTracker

Model Deployers

Model deployers handle model serving and deployment.

class BaseModelDeployer:
    """
    Abstract base class for model deployers.

    Manages deployment of trained models to serving infrastructure.
    """

Import from:

from zenml.model_deployers import BaseModelDeployer

Model Registries

Model registries store and version trained models.

class BaseModelRegistry:
    """
    Abstract base class for model registries.

    Provides versioned model storage separate from artifact stores.
    """

Import from:

from zenml.model_registries import BaseModelRegistry

Feature Stores

Feature stores manage feature data for ML models.

class BaseFeatureStore:
    """
    Abstract base class for feature stores.

    Integrates with feature stores like Feast for feature management.
    """

Import from:

from zenml.feature_stores import BaseFeatureStore

Data Validators

Data validators perform data quality checks.

class BaseDataValidator:
    """
    Abstract base class for data validators.

    Integrates with data validation frameworks like Great Expectations,
    Deepchecks, Evidently, etc.
    """

Import from:

from zenml.data_validators import BaseDataValidator

Alerters

Alerters send notifications about pipeline events.

class BaseAlerter:
    """
    Abstract base class for alerters.

    Sends notifications via Slack, Discord, email, etc.
    """

Import from:

from zenml.alerter import BaseAlerter

Annotators

Annotators integrate data annotation tools.

class BaseAnnotator:
    """
    Abstract base class for data annotation tools.

    Integrates with annotation platforms like Label Studio, Argilla, etc.
    """

Import from:

from zenml.annotators import BaseAnnotator

Deployers

Deployers manage pipeline deployment to production.

class BaseDeployer:
    """
    Abstract base class for pipeline deployers.

    Manages deployment of entire pipelines to production environments.
    """

class DockerDeployer(BaseDeployer):
    """Deployer for Docker-based pipeline deployments."""

Import from:

from zenml.deployers import BaseDeployer, DockerDeployer

Usage Examples

Creating Custom Orchestrator

from zenml.orchestrators import BaseOrchestrator, BaseOrchestratorConfig
from zenml.stack import StackComponentConfig

class MyOrchestratorConfig(BaseOrchestratorConfig):
    """Configuration for custom orchestrator."""
    api_endpoint: str
    api_key: str

class MyOrchestrator(BaseOrchestrator):
    """Custom orchestrator implementation."""

    @property
    def config(self) -> MyOrchestratorConfig:
        """Get typed configuration."""
        return self._config

    def prepare_pipeline_deployment(self, deployment, stack):
        """Prepare pipeline for execution."""
        # Custom preparation logic
        pass

    def run(self, deployment, stack):
        """Execute the pipeline."""
        # Custom execution logic
        print(f"Running on {self.config.api_endpoint}")

Using Stack Components via Client

from zenml.client import Client
from zenml.enums import StackComponentType

client = Client()

# List orchestrators
orchestrators = client.list_stack_components(
    type=StackComponentType.ORCHESTRATOR
)

for orch in orchestrators:
    print(f"Orchestrator: {orch.name} ({orch.flavor})")

# Get artifact store details
artifact_store = client.get_stack_component(
    name_or_id="s3_store",
    component_type=StackComponentType.ARTIFACT_STORE
)

print(f"Artifact Store: {artifact_store.name}")
print(f"Configuration: {artifact_store.configuration}")

Accessing Active Stack Components

from zenml.client import Client

client = Client()
stack = client.active_stack

# Access orchestrator
print(f"Orchestrator: {stack.orchestrator.name}")
print(f"Flavor: {stack.orchestrator.flavor}")

# Access artifact store
print(f"Artifact Store: {stack.artifact_store.name}")

# Check for optional components
if stack.experiment_tracker:
    print(f"Experiment Tracker: {stack.experiment_tracker.name}")

if stack.model_deployer:
    print(f"Model Deployer: {stack.model_deployer.name}")

Creating Stack with Components

from zenml.client import Client
from zenml.enums import StackComponentType

client = Client()

# Create components
client.create_stack_component(
    name="local_orch",
    flavor="local",
    component_type=StackComponentType.ORCHESTRATOR,
    configuration={}
)

client.create_stack_component(
    name="local_store",
    flavor="local",
    component_type=StackComponentType.ARTIFACT_STORE,
    configuration={"path": "/tmp/zenml"}
)

client.create_stack_component(
    name="mlflow_tracker",
    flavor="mlflow",
    component_type=StackComponentType.EXPERIMENT_TRACKER,
    configuration={"tracking_uri": "http://localhost:5000"}
)

# Create stack
stack = client.create_stack(
    name="ml_stack",
    components={
        "orchestrator": "local_orch",
        "artifact_store": "local_store",
        "experiment_tracker": "mlflow_tracker"
    }
)

Component Configuration Types

from zenml.client import Client

client = Client()

# Get component with full config
component = client.get_stack_component(
    name_or_id="sagemaker_orch",
    component_type="orchestrator"
)

# Access configuration
config = component.configuration
print(f"Execution role: {config.get('execution_role')}")
print(f"Region: {config.get('region')}")
print(f"Instance type: {config.get('instance_type')}")

Install with Tessl CLI

npx tessl i tessl/pypi-zenml

tessl/pypi-zenml