Profile

zenml.cli.profile

CLI for migrating legacy ZenML profiles.

LegacyComponentModel (BaseModel) pydantic-model

Legacy stack component model.

Source code in zenml/cli/profile.py

class LegacyComponentModel(BaseModel):
    """Legacy stack component model."""

    name: str
    type: StackComponentType
    flavor: str
    configuration: Dict[str, Any]

    def to_model(self, project_id: UUID, user_id: UUID) -> ComponentModel:
        """Converts to a component model.

        Args:
            project_id: The project id.
            user_id: The user id.

        Returns:
            The component model.
        """
        return ComponentModel(
            name=self.name,
            type=self.type,
            flavor=self.flavor,
            configuration=self.configuration,
            project=project_id,
            user=user_id,
        )

to_model(self, project_id, user_id)

Converts to a component model.

Parameters:

Name	Type	Description	Default
project_id	UUID	The project id.	required
user_id	UUID	The user id.	required

Returns:

Type	Description
ComponentModel	The component model.

Source code in zenml/cli/profile.py

def to_model(self, project_id: UUID, user_id: UUID) -> ComponentModel:
    """Converts to a component model.

    Args:
        project_id: The project id.
        user_id: The user id.

    Returns:
        The component model.
    """
    return ComponentModel(
        name=self.name,
        type=self.type,
        flavor=self.flavor,
        configuration=self.configuration,
        project=project_id,
        user=user_id,
    )

LegacyFlavorModel (BaseModel) pydantic-model

Legacy flavor model.

Source code in zenml/cli/profile.py

class LegacyFlavorModel(BaseModel):
    """Legacy flavor model."""

    name: str
    type: str
    source: str
    integration: Optional[str]

LegacyStackModel (BaseModel) pydantic-model

Legacy stack model.

Source code in zenml/cli/profile.py

class LegacyStackModel(BaseModel):
    """Legacy stack model."""

    name: str
    components: Dict[StackComponentType, str]

    def to_model(
        self, project_id: UUID, user_id: UUID, components: List[ComponentModel]
    ) -> StackModel:
        """Converts to a stack model.

        Args:
            project_id: The project id.
            user_id: The user id.
            components: The components that are part of the stack.

        Returns:
            The stack model.
        """
        return StackModel(
            name=self.name,
            components={c.type: [c.id] for c in components},
            project=project_id,
            user=user_id,
        )

to_model(self, project_id, user_id, components)

Converts to a stack model.

Parameters:

Name	Type	Description	Default
project_id	UUID	The project id.	required
user_id	UUID	The user id.	required
components	List[zenml.models.component_model.ComponentModel]	The components that are part of the stack.	required

Returns:

Type	Description
StackModel	The stack model.

Source code in zenml/cli/profile.py

def to_model(
    self, project_id: UUID, user_id: UUID, components: List[ComponentModel]
) -> StackModel:
    """Converts to a stack model.

    Args:
        project_id: The project id.
        user_id: The user id.
        components: The components that are part of the stack.

    Returns:
        The stack model.
    """
    return StackModel(
        name=self.name,
        components={c.type: [c.id] for c in components},
        project=project_id,
        user=user_id,
    )

LocalStore (BaseModel) pydantic-model

Pydantic object used for serializing a legacy YAML ZenML store.

Attributes:

Name	Type	Description
stacks	Dict[str, Dict[str, str]]	Maps stack names to a configuration object containing the names and flavors of all stack components.
stack_components	DefaultDict[str, Dict[str, str]]	Contains names and flavors of all registered stack components.
stack_component_flavors	List[zenml.cli.profile.LegacyFlavorModel]	Contains the flavor definitions of each stack component type

Source code in zenml/cli/profile.py

class LocalStore(BaseModel):
    """Pydantic object used for serializing a legacy YAML ZenML store.

    Attributes:
        stacks: Maps stack names to a configuration object containing the
            names and flavors of all stack components.
        stack_components: Contains names and flavors of all registered stack
            components.
        stack_component_flavors: Contains the flavor definitions of each
            stack component type
    """

    stacks: Dict[str, Dict[str, str]] = Field(default_factory=dict)
    stack_components: DefaultDict[str, Dict[str, str]] = Field(
        default=defaultdict(dict)
    )
    stack_component_flavors: List[LegacyFlavorModel] = Field(
        default_factory=list
    )

    _config_file: str

    @validator("stack_components")
    def _construct_stack_components_defaultdict(
        cls, stack_components: Dict[str, Dict[str, str]]
    ) -> DefaultDict[str, Dict[str, str]]:
        """Ensures that `stack_components` is a defaultdict.

        This is so stack components of a new component type can be added without
        issues.

        Args:
            stack_components: the dictionary of stack components

        Returns:
            Stack components dictionary.
        """
        return defaultdict(dict, stack_components)

    def __init__(self, config_file: str) -> None:
        """Create a local store instance initialized from a configuration file on disk.

        Args:
            config_file: configuration file path. If the file exists, the model
                will be initialized with the values from the file.
        """
        config_dict = {}
        if fileio.exists(config_file):
            config_dict = yaml_utils.read_yaml(config_file)

        self._config_file = config_file
        super().__init__(**config_dict)

    def _get_stack_component_config_path(
        self, component_type: str, name: str
    ) -> str:
        """Path to the configuration file of a stack component.

        Args:
            component_type: The type of the stack component.
            name: The name of the stack component.

        Returns:
            The path to the configuration file of the stack component.
        """
        subpath = f"{component_type}s"
        if component_type == StackComponentType.CONTAINER_REGISTRY:
            subpath = "container_registries"

        path = self.root / subpath / f"{name}.yaml"
        return str(path)

    def get_component(
        self,
        component_type: str,
        name: str,
        prefix: str = "",
    ) -> LegacyComponentModel:
        """Fetch the flavor and configuration for a stack component.

        Args:
            component_type: The type of the component to fetch.
            name: The name of the component to fetch.
            prefix: Optional name prefix to use for the returned component.

        Returns:
            A legacy component model for the stack component.

        Raises:
            KeyError: If no stack component exists for the given type and name.
        """
        components: Dict[str, str] = self.stack_components[component_type]
        if name not in components:
            raise KeyError(
                f"Unable to find stack component (type: {component_type}) "
                f"with name '{name}'. Available names: {set(components)}."
            )

        component_config_path = self._get_stack_component_config_path(
            component_type=component_type, name=name
        )
        flavor = components[name]
        config_dict: Dict[str, Any] = yaml_utils.read_yaml(
            component_config_path
        )
        component_name = (prefix + name) if name != "default" else name
        config_dict.pop("uuid")
        config_dict.pop("name")
        # filter out empty values to avoid validation errors
        config_dict = dict(
            filter(lambda x: x[1] is not None, config_dict.items())
        )

        return LegacyComponentModel(
            name=component_name,
            type=component_type,
            flavor=flavor,
            configuration=config_dict,
        )

    def get_components(self, prefix: str = "") -> List[LegacyComponentModel]:
        """Fetch all stack components.

        The default components are expressly excluded from this list.

        Args:
            prefix: Optional name prefix to use for the returned components.

        Returns:
            A list of component models for all stack components.
        """
        components: List[LegacyComponentModel] = []
        for component_type in self.stack_components:
            if component_type == "metadata_store":
                continue
            for name in self.stack_components[component_type]:
                if name == "default" and component_type in [
                    StackComponentType.ARTIFACT_STORE,
                    StackComponentType.ORCHESTRATOR,
                ]:
                    continue
                components.append(
                    self.get_component(
                        component_type,
                        name,
                        prefix=prefix,
                    )
                )
        return components

    def get_stack(self, name: str, prefix: str = "") -> LegacyStackModel:
        """Fetch the configuration for a stack.

        For default stack components, the default component in the current
        store is used instead of the one in the legacy profile.

        Args:
            name: The name of the stack to fetch.
            prefix: Optional name prefix to use for the returned stack and its
                components (unless default).

        Returns:
            A legacy stack model for the stack.

        Raises:
            KeyError: If no stack exists with the given name.
        """
        stack = self.stacks.get(name)
        if not stack:
            raise KeyError(
                f"Unable to find stack with name '{name}'. Available names: "
                f"{set(self.stacks)}."
            )

        components: Dict[StackComponentType, str] = {}
        for component_type, component_name in stack.items():
            if component_type == "metadata_store":
                continue
            components[StackComponentType(component_type)] = (
                (prefix + component_name)
                if component_name != "default"
                else component_name
            )

        return LegacyStackModel(
            name=prefix + name,
            components=components,
        )

    def get_stacks(self, prefix: str = "") -> List[LegacyStackModel]:
        """Fetch all stacks.

        The default stack is expressly excluded from this list.

        Args:
            prefix: Optional name prefix to use for the returned stack and
                stack components.

        Returns:
            A list of stacks.
        """
        return [
            self.get_stack(name, prefix=prefix)
            for name in self.stacks
            if name != "default"
        ]

    @property
    def root(self) -> Path:
        """The root directory of the zen store.

        Returns:
            The root directory of the zen store.
        """
        return Path(self._config_file).parent

    def is_empty(self) -> bool:
        """Check if the store is empty.

        Returns:
            True if the store is empty, False otherwise.
        """
        return (
            not self.contains_stacks()
            and not self.contains_stack_components()
            and not len(self.stack_component_flavors)
        )

    def contains_stacks(self) -> bool:
        """Check if the store contains stacks.

        The default stack is expressly excluded from this check.

        Returns:
            True if the store contains stacks, False otherwise.
        """
        return len(self.get_stacks()) > 0

    def contains_stack_components(self) -> bool:
        """Check if the store contains stack components.

        The default stack components are expressly excluded from this check.

        Returns:
            True if the store contains stack components, False otherwise.
        """
        Client()
        return len(self.get_components()) > 0

    @property
    def config_file(self) -> str:
        """Return the path to the configuration file.

        Returns:
            The path to the configuration file.
        """
        return self._config_file

    class Config:
        """Pydantic configuration class."""

        # Validate attributes when assigning them. We need to set this in order
        # to have a mix of mutable and immutable attributes
        validate_assignment = True
        # Ignore extra attributes from configs of previous ZenML versions
        extra = "ignore"
        # all attributes with leading underscore are private and therefore
        # are mutable and not included in serialization
        underscore_attrs_are_private = True

config_file: str property readonly

Return the path to the configuration file.

Returns:

Type	Description
str	The path to the configuration file.

root: Path property readonly

The root directory of the zen store.

Returns:

Type	Description
Path	The root directory of the zen store.

Config

Pydantic configuration class.

Source code in zenml/cli/profile.py

class Config:
    """Pydantic configuration class."""

    # Validate attributes when assigning them. We need to set this in order
    # to have a mix of mutable and immutable attributes
    validate_assignment = True
    # Ignore extra attributes from configs of previous ZenML versions
    extra = "ignore"
    # all attributes with leading underscore are private and therefore
    # are mutable and not included in serialization
    underscore_attrs_are_private = True

__init__(self, config_file) special

Create a local store instance initialized from a configuration file on disk.

Parameters:

Name	Type	Description	Default
config_file	str	configuration file path. If the file exists, the model will be initialized with the values from the file.	required

Source code in zenml/cli/profile.py

def __init__(self, config_file: str) -> None:
    """Create a local store instance initialized from a configuration file on disk.

    Args:
        config_file: configuration file path. If the file exists, the model
            will be initialized with the values from the file.
    """
    config_dict = {}
    if fileio.exists(config_file):
        config_dict = yaml_utils.read_yaml(config_file)

    self._config_file = config_file
    super().__init__(**config_dict)

contains_stack_components(self)

Check if the store contains stack components.

The default stack components are expressly excluded from this check.

Returns:

Type	Description
bool	True if the store contains stack components, False otherwise.

Source code in zenml/cli/profile.py

def contains_stack_components(self) -> bool:
    """Check if the store contains stack components.

    The default stack components are expressly excluded from this check.

    Returns:
        True if the store contains stack components, False otherwise.
    """
    Client()
    return len(self.get_components()) > 0

contains_stacks(self)

Check if the store contains stacks.

The default stack is expressly excluded from this check.

Returns:

Type	Description
bool	True if the store contains stacks, False otherwise.

Source code in zenml/cli/profile.py

def contains_stacks(self) -> bool:
    """Check if the store contains stacks.

    The default stack is expressly excluded from this check.

    Returns:
        True if the store contains stacks, False otherwise.
    """
    return len(self.get_stacks()) > 0

get_component(self, component_type, name, prefix='')

Fetch the flavor and configuration for a stack component.

Parameters:

Name	Type	Description	Default
component_type	str	The type of the component to fetch.	required
name	str	The name of the component to fetch.	required
prefix	str	Optional name prefix to use for the returned component.	''

Returns:

Type	Description
LegacyComponentModel	A legacy component model for the stack component.

Exceptions:

Type	Description
KeyError	If no stack component exists for the given type and name.

Source code in zenml/cli/profile.py

def get_component(
    self,
    component_type: str,
    name: str,
    prefix: str = "",
) -> LegacyComponentModel:
    """Fetch the flavor and configuration for a stack component.

    Args:
        component_type: The type of the component to fetch.
        name: The name of the component to fetch.
        prefix: Optional name prefix to use for the returned component.

    Returns:
        A legacy component model for the stack component.

    Raises:
        KeyError: If no stack component exists for the given type and name.
    """
    components: Dict[str, str] = self.stack_components[component_type]
    if name not in components:
        raise KeyError(
            f"Unable to find stack component (type: {component_type}) "
            f"with name '{name}'. Available names: {set(components)}."
        )

    component_config_path = self._get_stack_component_config_path(
        component_type=component_type, name=name
    )
    flavor = components[name]
    config_dict: Dict[str, Any] = yaml_utils.read_yaml(
        component_config_path
    )
    component_name = (prefix + name) if name != "default" else name
    config_dict.pop("uuid")
    config_dict.pop("name")
    # filter out empty values to avoid validation errors
    config_dict = dict(
        filter(lambda x: x[1] is not None, config_dict.items())
    )

    return LegacyComponentModel(
        name=component_name,
        type=component_type,
        flavor=flavor,
        configuration=config_dict,
    )

get_components(self, prefix='')

Fetch all stack components.

The default components are expressly excluded from this list.

Parameters:

Name	Type	Description	Default
prefix	str	Optional name prefix to use for the returned components.	''

Returns:

Type	Description
List[zenml.cli.profile.LegacyComponentModel]	A list of component models for all stack components.

Source code in zenml/cli/profile.py

def get_components(self, prefix: str = "") -> List[LegacyComponentModel]:
    """Fetch all stack components.

    The default components are expressly excluded from this list.

    Args:
        prefix: Optional name prefix to use for the returned components.

    Returns:
        A list of component models for all stack components.
    """
    components: List[LegacyComponentModel] = []
    for component_type in self.stack_components:
        if component_type == "metadata_store":
            continue
        for name in self.stack_components[component_type]:
            if name == "default" and component_type in [
                StackComponentType.ARTIFACT_STORE,
                StackComponentType.ORCHESTRATOR,
            ]:
                continue
            components.append(
                self.get_component(
                    component_type,
                    name,
                    prefix=prefix,
                )
            )
    return components

get_stack(self, name, prefix='')

Fetch the configuration for a stack.

For default stack components, the default component in the current store is used instead of the one in the legacy profile.

Parameters:

Name	Type	Description	Default
name	str	The name of the stack to fetch.	required
prefix	str	Optional name prefix to use for the returned stack and its components (unless default).	''

Returns:

Type	Description
LegacyStackModel	A legacy stack model for the stack.

Exceptions:

Type	Description
KeyError	If no stack exists with the given name.

Source code in zenml/cli/profile.py

def get_stack(self, name: str, prefix: str = "") -> LegacyStackModel:
    """Fetch the configuration for a stack.

    For default stack components, the default component in the current
    store is used instead of the one in the legacy profile.

    Args:
        name: The name of the stack to fetch.
        prefix: Optional name prefix to use for the returned stack and its
            components (unless default).

    Returns:
        A legacy stack model for the stack.

    Raises:
        KeyError: If no stack exists with the given name.
    """
    stack = self.stacks.get(name)
    if not stack:
        raise KeyError(
            f"Unable to find stack with name '{name}'. Available names: "
            f"{set(self.stacks)}."
        )

    components: Dict[StackComponentType, str] = {}
    for component_type, component_name in stack.items():
        if component_type == "metadata_store":
            continue
        components[StackComponentType(component_type)] = (
            (prefix + component_name)
            if component_name != "default"
            else component_name
        )

    return LegacyStackModel(
        name=prefix + name,
        components=components,
    )

get_stacks(self, prefix='')

Fetch all stacks.

The default stack is expressly excluded from this list.

Parameters:

Name	Type	Description	Default
prefix	str	Optional name prefix to use for the returned stack and stack components.	''

Returns:

Type	Description
List[zenml.cli.profile.LegacyStackModel]	A list of stacks.

Source code in zenml/cli/profile.py

def get_stacks(self, prefix: str = "") -> List[LegacyStackModel]:
    """Fetch all stacks.

    The default stack is expressly excluded from this list.

    Args:
        prefix: Optional name prefix to use for the returned stack and
            stack components.

    Returns:
        A list of stacks.
    """
    return [
        self.get_stack(name, prefix=prefix)
        for name in self.stacks
        if name != "default"
    ]

is_empty(self)

Check if the store is empty.

Returns:

Type	Description
bool	True if the store is empty, False otherwise.

Source code in zenml/cli/profile.py

def is_empty(self) -> bool:
    """Check if the store is empty.

    Returns:
        True if the store is empty, False otherwise.
    """
    return (
        not self.contains_stacks()
        and not self.contains_stack_components()
        and not len(self.stack_component_flavors)
    )

find_profiles(path)

Find all local profiles in a directory.

Point this function to a global config directory or a single profile directory to find and load all profiles stored there.

Parameters:

Name	Type	Description	Default
path	Optional[str]	The path to search for profiles. If None, the global config directory is used.	required

Yields:

Type	Description
Generator[zenml.cli.profile.LocalStore, NoneType, NoneType]	A local profile store.

Source code in zenml/cli/profile.py

def find_profiles(
    path: Optional[str],
) -> Generator[LocalStore, None, None]:
    """Find all local profiles in a directory.

    Point this function to a global config directory or a single profile
    directory to find and load all profiles stored there.

    Args:
        path: The path to search for profiles. If None, the global config
            directory is used.

    Yields:
        A local profile store.
    """
    dirs: List[str] = []
    if path is None:
        path = get_global_config_directory()
    if os.path.isdir(os.path.join(path, "profiles")):
        path = os.path.join(path, "profiles")
        dirs = [os.path.join(path, f) for f in os.listdir(path)]
    else:
        dirs = [path]
    for dir in dirs:
        if os.path.isfile(os.path.join(dir, "stacks.yaml")):
            store = LocalStore(os.path.join(dir, "stacks.yaml"))
            if store.is_empty():
                continue
            yield store