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

Stacks

Name: tessl/pypi-zenml
Author: tessl

Stack and stack component classes for configuring ML infrastructure. A ZenML stack represents a complete infrastructure configuration combining an orchestrator, artifact store, and optional components like container registries, experiment trackers, and step operators.

Capabilities

Stack Class

class Stack:
    """
    Complete ZenML stack configuration.

    A stack combines multiple stack components to provide
    a complete ML infrastructure configuration.

    Attributes:
    - name: Stack name
    - components: Dict mapping component types to component instances
    - id: Stack UUID
    """

    @property
    def orchestrator(self):
        """Get the orchestrator component."""

    @property
    def artifact_store(self):
        """Get the artifact store component."""

    @property
    def container_registry(self):
        """Get the container registry component (optional)."""

    @property
    def step_operator(self):
        """Get the step operator component (optional)."""

    @property
    def experiment_tracker(self):
        """Get the experiment tracker component (optional)."""

    @property
    def model_deployer(self):
        """Get the model deployer component (optional)."""

    @property
    def feature_store(self):
        """Get the feature store component (optional)."""

    @property
    def model_registry(self):
        """Get the model registry component (optional)."""

    @property
    def alerter(self):
        """Get the alerter component (optional)."""

    @property
    def annotator(self):
        """Get the annotator component (optional)."""

    @property
    def data_validator(self):
        """Get the data validator component (optional)."""

    @property
    def image_builder(self):
        """Get the image builder component (optional)."""

    @property
    def deployer(self):
        """Get the deployer component (optional)."""

    def dict(self) -> dict:
        """Get stack configuration as dict."""

Import from:

from zenml.stack import Stack

Stack Component

class StackComponent:
    """
    Base class for stack components.

    All stack components (orchestrators, artifact stores, etc.)
    inherit from this class.

    Attributes:
    - name: Component name
    - id: Component UUID
    - type: Component type (from StackComponentType enum)
    - flavor: Component flavor
    - config: Component configuration
    """

    @property
    def name(self) -> str:
        """Get component name."""

    @property
    def type(self) -> str:
        """Get component type."""

    @property
    def flavor(self) -> str:
        """Get component flavor."""

    @property
    def config(self):
        """Get component configuration."""

Import from:

from zenml.stack import StackComponent

Stack Component Config

class StackComponentConfig:
    """
    Base configuration for stack components.

    All stack component configurations inherit from this class.
    Configuration classes are Pydantic models with validation.
    """

Import from:

from zenml.stack import StackComponentConfig

Flavor

class Flavor:
    """
    Flavor of a stack component.

    Flavors define different implementations of stack components.
    For example, the artifact_store component has flavors like
    'local', 's3', 'gcs', 'azure'.

    Attributes:
    - name: Flavor name
    - type: Component type
    - config_class: Configuration class
    - implementation_class: Implementation class
    """

    @property
    def name(self) -> str:
        """Get flavor name."""

    @property
    def type(self) -> str:
        """Get component type."""

    @property
    def config_class(self) -> type:
        """Get configuration class."""

    @property
    def implementation_class(self) -> type:
        """Get implementation class."""

Import from:

from zenml.stack import Flavor

Stack Validator

class StackValidator:
    """
    Validator for stack configurations.

    Validates that stack components are compatible with each other
    and that all required components are present.
    """

    def validate(self, stack: Stack):
        """
        Validate stack configuration.

        Parameters:
        - stack: Stack to validate

        Raises:
        StackValidationError: If stack configuration is invalid
        """

Import from:

from zenml.stack import StackValidator

Stack Component Type Enum

class StackComponentType(str, Enum):
    """
    Stack component types.

    Values:
    - ORCHESTRATOR: Pipeline orchestration (required)
    - ARTIFACT_STORE: Artifact storage (required)
    - CONTAINER_REGISTRY: Container image registry
    - IMAGE_BUILDER: Container image builder
    - STEP_OPERATOR: Remote step execution
    - EXPERIMENT_TRACKER: Experiment tracking
    - MODEL_DEPLOYER: Model deployment
    - FEATURE_STORE: Feature store
    - MODEL_REGISTRY: Model registry
    - DATA_VALIDATOR: Data validation
    - ALERTER: Alerting/notifications
    - ANNOTATOR: Data annotation
    - DEPLOYER: Pipeline deployment
    """
    ORCHESTRATOR = "orchestrator"
    ARTIFACT_STORE = "artifact_store"
    CONTAINER_REGISTRY = "container_registry"
    IMAGE_BUILDER = "image_builder"
    STEP_OPERATOR = "step_operator"
    EXPERIMENT_TRACKER = "experiment_tracker"
    MODEL_DEPLOYER = "model_deployer"
    FEATURE_STORE = "feature_store"
    MODEL_REGISTRY = "model_registry"
    DATA_VALIDATOR = "data_validator"
    ALERTER = "alerter"
    ANNOTATOR = "annotator"
    DEPLOYER = "deployer"

Import from:

from zenml.enums import StackComponentType

Usage Examples

Accessing Stack Components

from zenml.client import Client

client = Client()
stack = client.active_stack

# Access required components
print(f"Orchestrator: {stack.orchestrator.name} ({stack.orchestrator.flavor})")
print(f"Artifact Store: {stack.artifact_store.name} ({stack.artifact_store.flavor})")

# Access optional components
if stack.container_registry:
    print(f"Container Registry: {stack.container_registry.name}")

if stack.experiment_tracker:
    print(f"Experiment Tracker: {stack.experiment_tracker.name}")

if stack.step_operator:
    print(f"Step Operator: {stack.step_operator.name}")

Creating a Stack via Client

from zenml.client import Client

client = Client()

# Create components first
client.create_stack_component(
    name="local_orchestrator",
    flavor="local",
    component_type="orchestrator",
    configuration={}
)

client.create_stack_component(
    name="local_artifact_store",
    flavor="local",
    component_type="artifact_store",
    configuration={"path": "/tmp/zenml-artifacts"}
)

# Create stack with components
stack = client.create_stack(
    name="local_stack",
    components={
        "orchestrator": "local_orchestrator",
        "artifact_store": "local_artifact_store"
    },
    description="Local development stack"
)

# Activate the stack
client.activate_stack("local_stack")

Creating Cloud Stack

from zenml.client import Client

client = Client()

# Create S3 artifact store
client.create_stack_component(
    name="s3_store",
    flavor="s3",
    component_type="artifact_store",
    configuration={
        "path": "s3://my-zenml-bucket/artifacts",
        "region": "us-east-1"
    }
)

# Create SageMaker orchestrator
client.create_stack_component(
    name="sagemaker_orchestrator",
    flavor="sagemaker",
    component_type="orchestrator",
    configuration={
        "execution_role": "arn:aws:iam::123456789:role/SageMaker",
        "region": "us-east-1"
    }
)

# Create ECR container registry
client.create_stack_component(
    name="ecr_registry",
    flavor="aws",
    component_type="container_registry",
    configuration={
        "uri": "123456789.dkr.ecr.us-east-1.amazonaws.com",
        "region": "us-east-1"
    }
)

# Create stack
aws_stack = client.create_stack(
    name="aws_production",
    components={
        "orchestrator": "sagemaker_orchestrator",
        "artifact_store": "s3_store",
        "container_registry": "ecr_registry"
    },
    description="AWS production stack"
)

Stack with Experiment Tracking

from zenml.client import Client

client = Client()

# Create MLflow experiment tracker
client.create_stack_component(
    name="mlflow_tracker",
    flavor="mlflow",
    component_type="experiment_tracker",
    configuration={
        "tracking_uri": "http://mlflow-server:5000",
        "tracking_username": "admin",
        "tracking_password": "password"
    }
)

# Add to stack
client.update_stack(
    name_or_id="local_stack",
    components={
        "orchestrator": "local_orchestrator",
        "artifact_store": "local_artifact_store",
        "experiment_tracker": "mlflow_tracker"
    }
)

Listing Available Flavors

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

client = Client()

# List all orchestrator flavors
orchestrator_flavors = client.get_flavors_by_type(
    component_type=StackComponentType.ORCHESTRATOR
)

print("Available orchestrator flavors:")
for flavor in orchestrator_flavors:
    print(f"  - {flavor.name}: {flavor.description}")

# List all artifact store flavors
store_flavors = client.get_flavors_by_type(
    component_type=StackComponentType.ARTIFACT_STORE
)

print("\nAvailable artifact store flavors:")
for flavor in store_flavors:
    print(f"  - {flavor.name}: {flavor.description}")

Stack Component Configuration

from zenml.client import Client

client = Client()

# Get component configuration
component = client.get_stack_component(
    name_or_id="s3_store",
    component_type="artifact_store"
)

print(f"Component: {component.name}")
print(f"Flavor: {component.flavor}")
print(f"Configuration: {component.configuration}")

# Update component configuration
client.update_stack_component(
    name_or_id="s3_store",
    component_type="artifact_store",
    updated_configuration={
        "path": "s3://my-new-bucket/artifacts",
        "region": "us-west-2"
    }
)

Custom Stack Component

from zenml.stack import StackComponent, StackComponentConfig
from zenml.enums import StackComponentType

class MyCustomConfig(StackComponentConfig):
    """Custom component configuration."""
    api_key: str
    endpoint: str

class MyCustomComponent(StackComponent):
    """Custom stack component implementation."""

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

    def custom_method(self):
        """Custom functionality."""
        print(f"Connecting to {self.config.endpoint}")

Install with Tessl CLI

npx tessl i tessl/pypi-zenml

tessl/pypi-zenml

stacks.md.css-3qkkll{font-size:var(--chakra-font-sizes-sm);font-weight:var(--chakra-font-weights-normal);color:var(--chakra-colors-gray-300);}docs/

Stacks

Capabilities

Stack Class

Stack Component

Stack Component Config

Flavor

Stack Validator

Stack Component Type Enum

Usage Examples

Accessing Stack Components

Creating a Stack via Client

Creating Cloud Stack

Stack with Experiment Tracking

Listing Available Flavors

Stack Component Configuration

Custom Stack Component

stacks.mddocs/