Trackio

zenml.integrations.trackio

Initialization for the trackio integration.

The Trackio integration currently enables you to use Trackio tracking as a convenient way to visualize your experiment runs within the local Trackio UI and HF space.

Attributes

TRACKIO = 'trackio' module-attribute

TRACKIO_EXPERIMENT_TRACKER_FLAVOR = 'trackio' module-attribute

Classes

Flavor

Class for ZenML Flavors.

Attributes

config_class: Type[StackComponentConfig] abstractmethod property

Returns StackComponentConfig config class.

Returns:

Type	Description
Type[StackComponentConfig]	The config class.

config_schema: Dict[str, Any] property

The config schema for a flavor.

Returns:

Type	Description
Dict[str, Any]	The config schema.

display_name: Optional[str] property

The display name of the flavor.

By default, converts the technical name to a human-readable format. For example, "vm_kubernetes" becomes "VM Kubernetes". Flavors can override this to provide custom display names.

Returns:

Type	Description
Optional[str]	The display name of the flavor.

docs_url: Optional[str] property

A url to point at docs explaining this flavor.

Returns:

Type	Description
Optional[str]	A flavor docs url.

implementation_class: Type[StackComponent] abstractmethod property

Implementation class for this flavor.

Returns:

Type	Description
Type[StackComponent]	The implementation class for this flavor.

logo_url: Optional[str] property

A url to represent the flavor in the dashboard.

Returns:

Type	Description
Optional[str]	The flavor logo.

name: str abstractmethod property

The flavor name.

Returns:

Type	Description
str	The flavor name.

sdk_docs_url: Optional[str] property

A url to point at SDK docs explaining this flavor.

Returns:

Type	Description
Optional[str]	A flavor SDK docs url.

service_connector_requirements: Optional[ServiceConnectorRequirements] property

Service connector resource requirements for service connectors.

Specifies resource requirements that are used to filter the available service connector types that are compatible with this flavor.

Returns:

Type	Description
Optional[ServiceConnectorRequirements]	Requirements for compatible service connectors, if a service
Optional[ServiceConnectorRequirements]	connector is required for this flavor.

type: StackComponentType abstractmethod property

The stack component type.

Returns:

Type	Description
StackComponentType	The stack component type.

Methods:

from_model(flavor_model: FlavorResponse) -> Flavor classmethod

Loads a flavor from a model.

Parameters:

Name	Type	Description	Default
flavor_model	FlavorResponse	The model to load from.	required

Raises:

Type	Description
CustomFlavorImportError	If the custom flavor can't be imported.
ImportError	If the flavor can't be imported.

Returns:

Type	Description
Flavor	The loaded flavor.

Source code in src/zenml/stack/flavor.py

@classmethod
def from_model(cls, flavor_model: FlavorResponse) -> "Flavor":
    """Loads a flavor from a model.

    Args:
        flavor_model: The model to load from.

    Raises:
        CustomFlavorImportError: If the custom flavor can't be imported.
        ImportError: If the flavor can't be imported.

    Returns:
        The loaded flavor.
    """
    try:
        _, flavor = validate_flavor_source(
            source=flavor_model.source,
            component_type=flavor_model.type,
            validate_component_classes=False,
        )
    except (TypeError, ValueError) as err:
        if flavor_model.is_custom:
            flavor_module, _, _ = flavor_model.source.rpartition(".")
            expected_file_path = os.path.join(
                source_utils.get_source_root(),
                flavor_module.replace(".", os.path.sep),
            )
            raise CustomFlavorImportError(
                f"Couldn't import custom flavor {flavor_model.name}: "
                f"{err}. Make sure the custom flavor class "
                f"`{flavor_model.source}` is importable. If it is part of "
                "a library, make sure it is installed. If "
                "it is a local code file, make sure it exists at "
                f"`{expected_file_path}.py`."
            ) from err
        else:
            raise ImportError(
                f"Couldn't import flavor {flavor_model.name}: {err}"
            ) from err
    return flavor

generate_default_docs_url() -> str

Generate the doc urls for all inbuilt and integration flavors.

Note that this method is not going to be useful for custom flavors, which do not have any docs in the main zenml docs.

Returns:

Type	Description
str	The complete url to the zenml documentation

Source code in src/zenml/stack/flavor.py

def generate_default_docs_url(self) -> str:
    """Generate the doc urls for all inbuilt and integration flavors.

    Note that this method is not going to be useful for custom flavors,
    which do not have any docs in the main zenml docs.

    Returns:
        The complete url to the zenml documentation
    """
    component_type = self.type.plural.replace("_", "-")
    name = self.name.replace("_", "-")

    base = "https://docs.zenml.io"
    return f"{base}/stack-components/{component_type}/{name}"

generate_default_sdk_docs_url() -> str

Generate SDK docs url for a flavor.

Returns:

Type	Description
str	The complete url to the zenml SDK docs

Source code in src/zenml/stack/flavor.py

def generate_default_sdk_docs_url(self) -> str:
    """Generate SDK docs url for a flavor.

    Returns:
        The complete url to the zenml SDK docs
    """
    from zenml import __version__

    base = f"https://sdkdocs.zenml.io/{__version__}"

    component_type = self.type.plural

    if "zenml.integrations" in self.__module__:
        # Get integration name out of module path which will look something
        #  like this "zenml.integrations.<integration>....
        integration = self.__module__.split(
            "zenml.integrations.", maxsplit=1
        )[1].split(".")[0]

        # Get the config class name to point to the specific class
        config_class_name = self.config_class.__name__

        return (
            f"{base}/integration_code_docs"
            f"/integrations-{integration}"
            f"#zenml.integrations.{integration}.flavors.{config_class_name}"
        )

    else:
        return (
            f"{base}/core_code_docs/core-{component_type}/"
            f"#{self.__module__}"
        )

to_model(integration: Optional[str] = None, is_custom: bool = True) -> FlavorRequest

Converts a flavor to a model.

Parameters:

Name	Type	Description	Default
integration	Optional[str]	The integration to use for the model.	None
is_custom	bool	Whether the flavor is a custom flavor.	True

Returns:

Type	Description
FlavorRequest	The model.

Source code in src/zenml/stack/flavor.py

def to_model(
    self,
    integration: Optional[str] = None,
    is_custom: bool = True,
) -> FlavorRequest:
    """Converts a flavor to a model.

    Args:
        integration: The integration to use for the model.
        is_custom: Whether the flavor is a custom flavor.

    Returns:
        The model.
    """
    connector_requirements = self.service_connector_requirements
    connector_type = (
        connector_requirements.connector_type
        if connector_requirements
        else None
    )
    resource_type = (
        connector_requirements.resource_type
        if connector_requirements
        else None
    )
    resource_id_attr = (
        connector_requirements.resource_id_attr
        if connector_requirements
        else None
    )

    model = FlavorRequest(
        name=self.name,
        display_name=self.display_name,
        type=self.type,
        source=source_utils.resolve(self.__class__).import_path,
        config_schema=self.config_schema,
        connector_type=connector_type,
        connector_resource_type=resource_type,
        connector_resource_id_attr=resource_id_attr,
        integration=integration,
        logo_url=self.logo_url,
        docs_url=self.docs_url,
        sdk_docs_url=self.sdk_docs_url,
        is_custom=is_custom,
    )
    return model

Integration

Base class for integration in ZenML.

Methods:

activate() -> None classmethod

Abstract method to activate the integration.

Source code in src/zenml/integrations/integration.py

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

check_installation() -> bool 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 src/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.
    """
    for requirement in cls.get_requirements():
        parsed_requirement = Requirement(requirement)

        if not requirement_installed(parsed_requirement):
            logger.debug(
                "Requirement '%s' for integration '%s' is not installed "
                "or installed with the wrong version.",
                requirement,
                cls.NAME,
            )
            return False

        dependencies = get_dependencies(parsed_requirement)

        for dependency in dependencies:
            if not requirement_installed(dependency):
                logger.debug(
                    "Requirement '%s' for integration '%s' is not "
                    "installed or installed with the wrong version.",
                    dependency,
                    cls.NAME,
                )
                return False

    logger.debug(
        f"Integration '{cls.NAME}' is installed correctly with "
        f"requirements {cls.get_requirements()}."
    )
    return True

flavors() -> List[Type[Flavor]] classmethod

Abstract method to declare new stack component flavors.

Returns:

Type	Description
List[Type[Flavor]]	A list of new stack component flavors.

Source code in src/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: Optional[str] = None, python_version: Optional[str] = None) -> List[str] 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
python_version	Optional[str]	The Python version to use for the requirements.	None

Returns:

Type	Description
List[str]	A list of requirements.

Source code in src/zenml/integrations/integration.py

@classmethod
def get_requirements(
    cls,
    target_os: Optional[str] = None,
    python_version: 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.
        python_version: The Python version to use for the requirements.

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

get_uninstall_requirements(target_os: Optional[str] = None) -> List[str] classmethod

Method to get the uninstall 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 src/zenml/integrations/integration.py

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

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

    Returns:
        A list of requirements.
    """
    ret = []
    for each in cls.get_requirements(target_os=target_os):
        is_ignored = False
        for ignored in cls.REQUIREMENTS_IGNORED_ON_UNINSTALL:
            if each.startswith(ignored):
                is_ignored = True
                break
        if not is_ignored:
            ret.append(each)
    return ret

TrackioIntegration

Bases: Integration

Definition of Trackio integration for ZenML.

Methods:

flavors() -> List[Type[Flavor]] classmethod

Declare the stack component flavors for the Trackio integration. Returns: List of stack component flavors for this integration.

Source code in src/zenml/integrations/trackio/__init__.py

@classmethod
def flavors(cls) -> List[Type[Flavor]]:
    """Declare the stack component flavors for the Trackio integration.
    Returns:
        List of stack component flavors for this integration.
    """
    from zenml.integrations.trackio.flavors import (
        TrackioExperimentTrackerFlavor,
    )

    return [TrackioExperimentTrackerFlavor]

Modules

experiment_trackers

Initialization for the trackio experiment tracker.

Classes

TrackioExperimentTracker(name: str, id: UUID, config: StackComponentConfig, flavor: str, type: StackComponentType, user: Optional[UUID], created: datetime, updated: datetime, environment: Optional[Dict[str, str]] = None, secrets: Optional[List[UUID]] = None, labels: Optional[Dict[str, Any]] = None, connector_requirements: Optional[ServiceConnectorRequirements] = None, connector: Optional[UUID] = None, connector_resource_id: Optional[str] = None, *args: Any, **kwargs: Any)

Bases: BaseExperimentTracker

Track experiments using Trackio.

Source code in src/zenml/stack/stack_component.py

def __init__(
    self,
    name: str,
    id: UUID,
    config: StackComponentConfig,
    flavor: str,
    type: StackComponentType,
    user: Optional[UUID],
    created: datetime,
    updated: datetime,
    environment: Optional[Dict[str, str]] = None,
    secrets: Optional[List[UUID]] = None,
    labels: Optional[Dict[str, Any]] = None,
    connector_requirements: Optional[ServiceConnectorRequirements] = None,
    connector: Optional[UUID] = None,
    connector_resource_id: Optional[str] = None,
    *args: Any,
    **kwargs: Any,
):
    """Initializes a StackComponent.

    Args:
        name: The name of the component.
        id: The unique ID of the component.
        config: The config of the component.
        flavor: The flavor of the component.
        type: The type of the component.
        user: The ID of the user who created the component.
        created: The creation time of the component.
        updated: The last update time of the component.
        environment: Environment variables to set when running on this
            component.
        secrets: Secrets to set as environment variables when running on
            this component.
        labels: The labels of the component.
        connector_requirements: The requirements for the connector.
        connector: The ID of a connector linked to the component.
        connector_resource_id: The custom resource ID to access through
            the connector.
        *args: Additional positional arguments.
        **kwargs: Additional keyword arguments.

    Raises:
        ValueError: If a secret reference is passed as name.
    """
    if secret_utils.is_secret_reference(name):
        raise ValueError(
            "Passing the `name` attribute of a stack component as a "
            "secret reference is not allowed."
        )

    self.id = id
    self.name = name
    self._config = config
    self.flavor = flavor
    self.type = type
    self.user = user
    self.created = created
    self.updated = updated
    self.labels = labels
    self.environment = environment or {}
    self.secrets = secrets or []
    self.connector_requirements = connector_requirements
    self.connector = connector
    self.connector_resource_id = connector_resource_id
    self._connector_instance: Optional[ServiceConnector] = None

Attributes

config: TrackioExperimentTrackerConfig property

Returns the Trackio config.

settings_class: Type[TrackioExperimentTrackerSettings] property

Settings class for Trackio.

Methods:

cleanup_step_run(info: StepRunInfo, step_failed: bool) -> None

Finalize the Trackio run.

Source code in src/zenml/integrations/trackio/experiment_trackers/trackio_experiment_tracker.py

def cleanup_step_run(
    self,
    info: "StepRunInfo",
    step_failed: bool,
) -> None:
    """Finalize the Trackio run."""
    settings = cast(
        TrackioExperimentTrackerSettings,
        self.get_settings(info),
    )

    try:
        if settings.auto_sync:
            logger.info("Syncing Trackio run.")

            sync_kwargs: Dict[str, Any] = {
                "project": self.config.project_name,
            }

            # Publish to Hugging Face Space
            if self.config.space_id and settings.publish_to_space:
                sync_kwargs["space_id"] = self.config.space_id

            # Publish to Hugging Face Dataset
            if self.config.dataset_id and settings.publish_to_dataset:
                sync_kwargs["dataset_id"] = self.config.dataset_id

            # Optional bucket storage
            if self.config.bucket_id:
                sync_kwargs["bucket_id"] = self.config.bucket_id

            trackio_any = cast(Any, trackio)

            valid_params = inspect.signature(trackio_any.sync).parameters

            logger.debug(
                "valid sync params: %s",
                valid_params.keys(),
            )
            filtered_kwargs: Dict[str, Any] = {
                key: value
                for key, value in sync_kwargs.items()
                if key in valid_params and value is not None
            }

            logger.debug("Trackio sync kwargs: %s", filtered_kwargs)
            cast(Any, trackio).sync(**filtered_kwargs)

        if settings.auto_freeze and self.config.space_id:
            logger.info("Freezing Trackio dashboard.")

            freeze_kwargs: Dict[str, Any] = {
                "space_id": self.config.space_id,
                "project": self.config.project_name,
            }

            if self.config.bucket_id:
                freeze_kwargs["bucket_id"] = self.config.bucket_id

            valid_params = inspect.signature(
                cast(Any, trackio).freeze
            ).parameters
            filtered_kwargs = {
                key: value
                for key, value in freeze_kwargs.items()
                if key in valid_params and value is not None
            }

            logger.debug("Trackio freeze kwargs: %s", filtered_kwargs)
            cast(Any, trackio).freeze(**filtered_kwargs)

    except Exception as e:
        logger.warning(
            "Failed to finalize Trackio sync/freeze: %s",
            e,
        )

    finally:
        cast(Any, trackio.finish)()

        if self.config.hf_token:
            previous_token: Optional[str] = getattr(
                info,
                "_trackio_previous_hf_token",
                None,
            )

            if previous_token is not None:
                os.environ[HF_TOKEN_ENV_VAR] = previous_token
            else:
                os.environ.pop(
                    HF_TOKEN_ENV_VAR,
                    None,
                )

get_step_run_metadata(info: StepRunInfo) -> Dict[str, MetadataType]

Get metadata after a step ran.

Source code in src/zenml/integrations/trackio/experiment_trackers/trackio_experiment_tracker.py

def get_step_run_metadata(
    self,
    info: "StepRunInfo",
) -> Dict[str, "MetadataType"]:
    """Get metadata after a step ran."""
    run_url: Optional[str] = None
    run_name: Optional[str] = None

    current_trackio_run = getattr(
        trackio,
        "run",
        None,
    )

    if current_trackio_run:
        run_name = getattr(
            current_trackio_run,
            "name",
            None,
        )

    if self.config.space_id:
        run_url = f"https://huggingface.co/spaces/{self.config.space_id}"
    elif self.config.server_url:
        run_url = self.config.server_url

    default_run_name = f"{info.run_name}_{info.pipeline_step_name}"

    settings = cast(
        TrackioExperimentTrackerSettings,
        self.get_settings(info),
    )

    run_name = run_name or settings.run_name or default_run_name

    metadata: Dict[str, "MetadataType"] = {
        "trackio_run_name": run_name,
    }

    if run_url:
        metadata[METADATA_EXPERIMENT_TRACKER_URL] = Uri(run_url)

    return metadata

prepare_step_run(info: StepRunInfo) -> None

Configure a Trackio run.

Source code in src/zenml/integrations/trackio/experiment_trackers/trackio_experiment_tracker.py

def prepare_step_run(
    self,
    info: "StepRunInfo",
) -> None:
    """Configure a Trackio run."""
    if self.config.hf_token:
        previous_token: Optional[str] = os.environ.get(HF_TOKEN_ENV_VAR)

        setattr(
            info,
            "_trackio_previous_hf_token",
            previous_token,
        )

        os.environ[HF_TOKEN_ENV_VAR] = self.config.hf_token

    settings = cast(
        TrackioExperimentTrackerSettings,
        self.get_settings(info),
    )

    trackio_run_name = (
        settings.run_name or f"{info.run_name}_{info.pipeline_step_name}"
    )

    self._initialize_trackio(
        info=info,
        run_name=trackio_run_name,
    )

    tags = settings.tags

    if tags:
        try:
            trackio.log(
                {
                    "zenml_tags": tags,
                }
            )

        except Exception as e:
            logger.warning(
                "Failed to log Trackio tags: %s",
                e,
            )

Modules

trackio_experiment_tracker

Implementation for the trackio experiment tracker.

Classes

TrackioExperimentTracker(name: str, id: UUID, config: StackComponentConfig, flavor: str, type: StackComponentType, user: Optional[UUID], created: datetime, updated: datetime, environment: Optional[Dict[str, str]] = None, secrets: Optional[List[UUID]] = None, labels: Optional[Dict[str, Any]] = None, connector_requirements: Optional[ServiceConnectorRequirements] = None, connector: Optional[UUID] = None, connector_resource_id: Optional[str] = None, *args: Any, **kwargs: Any)

Bases: BaseExperimentTracker

Track experiments using Trackio.

Source code in src/zenml/stack/stack_component.py

def __init__(
    self,
    name: str,
    id: UUID,
    config: StackComponentConfig,
    flavor: str,
    type: StackComponentType,
    user: Optional[UUID],
    created: datetime,
    updated: datetime,
    environment: Optional[Dict[str, str]] = None,
    secrets: Optional[List[UUID]] = None,
    labels: Optional[Dict[str, Any]] = None,
    connector_requirements: Optional[ServiceConnectorRequirements] = None,
    connector: Optional[UUID] = None,
    connector_resource_id: Optional[str] = None,
    *args: Any,
    **kwargs: Any,
):
    """Initializes a StackComponent.

    Args:
        name: The name of the component.
        id: The unique ID of the component.
        config: The config of the component.
        flavor: The flavor of the component.
        type: The type of the component.
        user: The ID of the user who created the component.
        created: The creation time of the component.
        updated: The last update time of the component.
        environment: Environment variables to set when running on this
            component.
        secrets: Secrets to set as environment variables when running on
            this component.
        labels: The labels of the component.
        connector_requirements: The requirements for the connector.
        connector: The ID of a connector linked to the component.
        connector_resource_id: The custom resource ID to access through
            the connector.
        *args: Additional positional arguments.
        **kwargs: Additional keyword arguments.

    Raises:
        ValueError: If a secret reference is passed as name.
    """
    if secret_utils.is_secret_reference(name):
        raise ValueError(
            "Passing the `name` attribute of a stack component as a "
            "secret reference is not allowed."
        )

    self.id = id
    self.name = name
    self._config = config
    self.flavor = flavor
    self.type = type
    self.user = user
    self.created = created
    self.updated = updated
    self.labels = labels
    self.environment = environment or {}
    self.secrets = secrets or []
    self.connector_requirements = connector_requirements
    self.connector = connector
    self.connector_resource_id = connector_resource_id
    self._connector_instance: Optional[ServiceConnector] = None

Attributes

config: TrackioExperimentTrackerConfig property

Returns the Trackio config.

settings_class: Type[TrackioExperimentTrackerSettings] property

Settings class for Trackio.

Methods:

cleanup_step_run(info: StepRunInfo, step_failed: bool) -> None

Finalize the Trackio run.

Source code in src/zenml/integrations/trackio/experiment_trackers/trackio_experiment_tracker.py

def cleanup_step_run(
    self,
    info: "StepRunInfo",
    step_failed: bool,
) -> None:
    """Finalize the Trackio run."""
    settings = cast(
        TrackioExperimentTrackerSettings,
        self.get_settings(info),
    )

    try:
        if settings.auto_sync:
            logger.info("Syncing Trackio run.")

            sync_kwargs: Dict[str, Any] = {
                "project": self.config.project_name,
            }

            # Publish to Hugging Face Space
            if self.config.space_id and settings.publish_to_space:
                sync_kwargs["space_id"] = self.config.space_id

            # Publish to Hugging Face Dataset
            if self.config.dataset_id and settings.publish_to_dataset:
                sync_kwargs["dataset_id"] = self.config.dataset_id

            # Optional bucket storage
            if self.config.bucket_id:
                sync_kwargs["bucket_id"] = self.config.bucket_id

            trackio_any = cast(Any, trackio)

            valid_params = inspect.signature(trackio_any.sync).parameters

            logger.debug(
                "valid sync params: %s",
                valid_params.keys(),
            )
            filtered_kwargs: Dict[str, Any] = {
                key: value
                for key, value in sync_kwargs.items()
                if key in valid_params and value is not None
            }

            logger.debug("Trackio sync kwargs: %s", filtered_kwargs)
            cast(Any, trackio).sync(**filtered_kwargs)

        if settings.auto_freeze and self.config.space_id:
            logger.info("Freezing Trackio dashboard.")

            freeze_kwargs: Dict[str, Any] = {
                "space_id": self.config.space_id,
                "project": self.config.project_name,
            }

            if self.config.bucket_id:
                freeze_kwargs["bucket_id"] = self.config.bucket_id

            valid_params = inspect.signature(
                cast(Any, trackio).freeze
            ).parameters
            filtered_kwargs = {
                key: value
                for key, value in freeze_kwargs.items()
                if key in valid_params and value is not None
            }

            logger.debug("Trackio freeze kwargs: %s", filtered_kwargs)
            cast(Any, trackio).freeze(**filtered_kwargs)

    except Exception as e:
        logger.warning(
            "Failed to finalize Trackio sync/freeze: %s",
            e,
        )

    finally:
        cast(Any, trackio.finish)()

        if self.config.hf_token:
            previous_token: Optional[str] = getattr(
                info,
                "_trackio_previous_hf_token",
                None,
            )

            if previous_token is not None:
                os.environ[HF_TOKEN_ENV_VAR] = previous_token
            else:
                os.environ.pop(
                    HF_TOKEN_ENV_VAR,
                    None,
                )

get_step_run_metadata(info: StepRunInfo) -> Dict[str, MetadataType]

Get metadata after a step ran.

Source code in src/zenml/integrations/trackio/experiment_trackers/trackio_experiment_tracker.py

def get_step_run_metadata(
    self,
    info: "StepRunInfo",
) -> Dict[str, "MetadataType"]:
    """Get metadata after a step ran."""
    run_url: Optional[str] = None
    run_name: Optional[str] = None

    current_trackio_run = getattr(
        trackio,
        "run",
        None,
    )

    if current_trackio_run:
        run_name = getattr(
            current_trackio_run,
            "name",
            None,
        )

    if self.config.space_id:
        run_url = f"https://huggingface.co/spaces/{self.config.space_id}"
    elif self.config.server_url:
        run_url = self.config.server_url

    default_run_name = f"{info.run_name}_{info.pipeline_step_name}"

    settings = cast(
        TrackioExperimentTrackerSettings,
        self.get_settings(info),
    )

    run_name = run_name or settings.run_name or default_run_name

    metadata: Dict[str, "MetadataType"] = {
        "trackio_run_name": run_name,
    }

    if run_url:
        metadata[METADATA_EXPERIMENT_TRACKER_URL] = Uri(run_url)

    return metadata

prepare_step_run(info: StepRunInfo) -> None

Configure a Trackio run.

Source code in src/zenml/integrations/trackio/experiment_trackers/trackio_experiment_tracker.py

def prepare_step_run(
    self,
    info: "StepRunInfo",
) -> None:
    """Configure a Trackio run."""
    if self.config.hf_token:
        previous_token: Optional[str] = os.environ.get(HF_TOKEN_ENV_VAR)

        setattr(
            info,
            "_trackio_previous_hf_token",
            previous_token,
        )

        os.environ[HF_TOKEN_ENV_VAR] = self.config.hf_token

    settings = cast(
        TrackioExperimentTrackerSettings,
        self.get_settings(info),
    )

    trackio_run_name = (
        settings.run_name or f"{info.run_name}_{info.pipeline_step_name}"
    )

    self._initialize_trackio(
        info=info,
        run_name=trackio_run_name,
    )

    tags = settings.tags

    if tags:
        try:
            trackio.log(
                {
                    "zenml_tags": tags,
                }
            )

        except Exception as e:
            logger.warning(
                "Failed to log Trackio tags: %s",
                e,
            )

Functions:

flavors

Trackio integration flavors.

Classes

TrackioExperimentTrackerConfig(warn_about_plain_text_secrets: bool = False, **kwargs: Any)

Bases: BaseExperimentTrackerConfig, TrackioExperimentTrackerSettings

Config for the Trackio experiment tracker.

Source code in src/zenml/stack/stack_component.py

def __init__(
    self, warn_about_plain_text_secrets: bool = False, **kwargs: Any
) -> None:
    """Ensures that secret references don't clash with pydantic validation.

    StackComponents allow the specification of all their string attributes
    using secret references of the form `{{secret_name.key}}`. This however
    is only possible when the stack component does not perform any explicit
    validation of this attribute using pydantic validators. If this were
    the case, the validation would run on the secret reference and would
    fail or in the worst case, modify the secret reference and lead to
    unexpected behavior. This method ensures that no attributes that require
    custom pydantic validation are set as secret references.

    Args:
        warn_about_plain_text_secrets: If true, then warns about using
            plain-text secrets.
        **kwargs: Arguments to initialize this stack component.

    Raises:
        ValueError: If an attribute that requires custom pydantic validation
            is passed as a secret reference, or if the `name` attribute
            was passed as a secret reference.
    """
    for key, value in kwargs.items():
        try:
            field = self.__class__.model_fields[key]
        except KeyError:
            # Value for a private attribute or non-existing field, this
            # will fail during the upcoming pydantic validation
            continue

        if value is None:
            continue

        if not secret_utils.is_secret_reference(value):
            if (
                secret_utils.is_secret_field(field)
                and warn_about_plain_text_secrets
            ):
                logger.warning(
                    "You specified a plain-text value for the sensitive "
                    f"attribute `{key}` for a `{self.__class__.__name__}` "
                    "stack component. This is currently only a warning, "
                    "but future versions of ZenML will require you to pass "
                    "in sensitive information as secrets. Check out the "
                    "documentation on how to configure your stack "
                    "components with secrets here: "
                    "https://docs.zenml.io/deploying-zenml/deploying-zenml/secret-management"
                )
            continue

        if pydantic_utils.has_validators(
            pydantic_class=self.__class__, field_name=key
        ):
            raise ValueError(
                f"Passing the stack component attribute `{key}` as a "
                "secret reference is not allowed as additional validation "
                "is required for this attribute."
            )

    super().__init__(**kwargs)

TrackioExperimentTrackerFlavor

Bases: BaseExperimentTrackerFlavor

Flavor for the Trackio experiment tracker.

Attributes

config_class: Type[TrackioExperimentTrackerConfig] property

Returns Trackio config class.

Returns:

Type	Description
Type[TrackioExperimentTrackerConfig]	The Trackio config class.

docs_url: Optional[str] property

A URL pointing to flavor documentation.

Returns:

Type	Description
Optional[str]	Flavor docs URL.

implementation_class: Type[TrackioExperimentTracker] property

Implementation class for this flavor.

Returns:

Type	Description
Type[TrackioExperimentTracker]	The implementation class.

logo_url: str property

URL for the Trackio logo.

Returns:

Type	Description
str	Logo URL.

name: str property

Name of the flavor.

Returns:

Type	Description
str	The flavor name.

sdk_docs_url: Optional[str] property

A URL pointing to SDK documentation.

Returns:

Type	Description
Optional[str]	Flavor SDK docs URL.

Modules

trackio_experiment_tracker_flavor

Trackio experiment tracker flavor.

Classes

TrackioExperimentTrackerConfig(warn_about_plain_text_secrets: bool = False, **kwargs: Any)

Bases: BaseExperimentTrackerConfig, TrackioExperimentTrackerSettings

Config for the Trackio experiment tracker.

Source code in src/zenml/stack/stack_component.py

def __init__(
    self, warn_about_plain_text_secrets: bool = False, **kwargs: Any
) -> None:
    """Ensures that secret references don't clash with pydantic validation.

    StackComponents allow the specification of all their string attributes
    using secret references of the form `{{secret_name.key}}`. This however
    is only possible when the stack component does not perform any explicit
    validation of this attribute using pydantic validators. If this were
    the case, the validation would run on the secret reference and would
    fail or in the worst case, modify the secret reference and lead to
    unexpected behavior. This method ensures that no attributes that require
    custom pydantic validation are set as secret references.

    Args:
        warn_about_plain_text_secrets: If true, then warns about using
            plain-text secrets.
        **kwargs: Arguments to initialize this stack component.

    Raises:
        ValueError: If an attribute that requires custom pydantic validation
            is passed as a secret reference, or if the `name` attribute
            was passed as a secret reference.
    """
    for key, value in kwargs.items():
        try:
            field = self.__class__.model_fields[key]
        except KeyError:
            # Value for a private attribute or non-existing field, this
            # will fail during the upcoming pydantic validation
            continue

        if value is None:
            continue

        if not secret_utils.is_secret_reference(value):
            if (
                secret_utils.is_secret_field(field)
                and warn_about_plain_text_secrets
            ):
                logger.warning(
                    "You specified a plain-text value for the sensitive "
                    f"attribute `{key}` for a `{self.__class__.__name__}` "
                    "stack component. This is currently only a warning, "
                    "but future versions of ZenML will require you to pass "
                    "in sensitive information as secrets. Check out the "
                    "documentation on how to configure your stack "
                    "components with secrets here: "
                    "https://docs.zenml.io/deploying-zenml/deploying-zenml/secret-management"
                )
            continue

        if pydantic_utils.has_validators(
            pydantic_class=self.__class__, field_name=key
        ):
            raise ValueError(
                f"Passing the stack component attribute `{key}` as a "
                "secret reference is not allowed as additional validation "
                "is required for this attribute."
            )

    super().__init__(**kwargs)

TrackioExperimentTrackerFlavor

Bases: BaseExperimentTrackerFlavor

Flavor for the Trackio experiment tracker.

Attributes

config_class: Type[TrackioExperimentTrackerConfig] property

Returns Trackio config class.

Returns:

Type	Description
Type[TrackioExperimentTrackerConfig]	The Trackio config class.

docs_url: Optional[str] property

A URL pointing to flavor documentation.

Returns:

Type	Description
Optional[str]	Flavor docs URL.

implementation_class: Type[TrackioExperimentTracker] property

Implementation class for this flavor.

Returns:

Type	Description
Type[TrackioExperimentTracker]	The implementation class.

logo_url: str property

URL for the Trackio logo.

Returns:

Type	Description
str	Logo URL.

name: str property

Name of the flavor.

Returns:

Type	Description
str	The flavor name.

sdk_docs_url: Optional[str] property

A URL pointing to SDK documentation.

Returns:

Type	Description
Optional[str]	Flavor SDK docs URL.

TrackioExperimentTrackerSettings(warn_about_plain_text_secrets: bool = False, **kwargs: Any)

Bases: BaseSettings

Settings for the Trackio experiment tracker.

Source code in src/zenml/config/secret_reference_mixin.py

def __init__(
    self, warn_about_plain_text_secrets: bool = False, **kwargs: Any
) -> None:
    """Ensures that secret references are only passed for valid fields.

    This method ensures that secret references are not passed for fields
    that explicitly prevent them or require pydantic validation.

    Args:
        warn_about_plain_text_secrets: If true, then warns about using plain-text secrets.
        **kwargs: Arguments to initialize this object.

    Raises:
        ValueError: If an attribute that requires custom pydantic validation
            or an attribute which explicitly disallows secret references
            is passed as a secret reference.
    """
    for key, value in kwargs.items():
        try:
            field = self.__class__.model_fields[key]
        except KeyError:
            # Value for a private attribute or non-existing field, this
            # will fail during the upcoming pydantic validation
            continue

        if value is None:
            continue

        if not secret_utils.is_secret_reference(value):
            if (
                secret_utils.is_secret_field(field)
                and warn_about_plain_text_secrets
            ):
                logger.warning(
                    "You specified a plain-text value for the sensitive "
                    f"attribute `{key}`. This is currently only a warning, "
                    "but future versions of ZenML will require you to pass "
                    "in sensitive information as secrets. Check out the "
                    "documentation on how to configure values with secrets "
                    "here: https://docs.zenml.io/deploying-zenml/deploying-zenml/secret-management"
                )
            continue

        if secret_utils.is_clear_text_field(field):
            raise ValueError(
                f"Passing the `{key}` attribute as a secret reference is "
                "not allowed."
            )

        requires_validation = has_validators(
            pydantic_class=self.__class__, field_name=key
        )
        if requires_validation:
            raise ValueError(
                f"Passing the attribute `{key}` as a secret reference is "
                "not allowed as additional validation is required for "
                "this attribute."
            )

    super().__init__(**kwargs)

Functions: