Analytics

zenml.analytics

The 'analytics' module of ZenML.

Attributes

source_context: ContextVar[SourceContextTypes] = ContextVar('Source-Context', default=SourceContextTypes.PYTHON) module-attribute

Classes

AnalyticsEvent

Bases: str, Enum

Enum of events to track in segment.

SourceContextTypes

Bases: StrEnum

Enum for event source types.

Functions

alias(user_id: UUID, previous_id: UUID) -> bool

Alias user IDs.

Parameters:

Name	Type	Description	Default
user_id	UUID	The user ID.	required
previous_id	UUID	Previous ID for the alias.	required

Returns:

Type	Description
bool	True if event is sent successfully, False is not.

Source code in src/zenml/analytics/__init__.py

def alias(user_id: UUID, previous_id: UUID) -> bool:  # type: ignore[return]
    """Alias user IDs.

    Args:
        user_id: The user ID.
        previous_id: Previous ID for the alias.

    Returns:
        True if event is sent successfully, False is not.
    """
    from zenml.analytics.context import AnalyticsContext

    with AnalyticsContext() as analytics:
        return analytics.alias(user_id=user_id, previous_id=previous_id)

group(group_id: UUID, group_metadata: Optional[Dict[str, Any]] = None) -> bool

Attach metadata to a segment group.

Parameters:

Name	Type	Description	Default
group_id	UUID	ID of the group.	required
group_metadata	Optional[Dict[str, Any]]	Metadata to attach to the group.	None

Returns:

Type	Description
bool	True if event is sent successfully, False if not.

Source code in src/zenml/analytics/__init__.py

def group(  # type: ignore[return]
    group_id: UUID,
    group_metadata: Optional[Dict[str, Any]] = None,
) -> bool:
    """Attach metadata to a segment group.

    Args:
        group_id: ID of the group.
        group_metadata: Metadata to attach to the group.

    Returns:
        True if event is sent successfully, False if not.
    """
    from zenml.analytics.context import AnalyticsContext

    with AnalyticsContext() as analytics:
        return analytics.group(group_id=group_id, traits=group_metadata)

identify(metadata: Optional[Dict[str, Any]] = None) -> bool

Attach metadata to user directly.

Parameters:

Name	Type	Description	Default
metadata	Optional[Dict[str, Any]]	Dict of metadata to attach to the user.	None

Returns:

Type	Description
bool	True if event is sent successfully, False is not.

Source code in src/zenml/analytics/__init__.py

def identify(  # type: ignore[return]
    metadata: Optional[Dict[str, Any]] = None
) -> bool:
    """Attach metadata to user directly.

    Args:
        metadata: Dict of metadata to attach to the user.

    Returns:
        True if event is sent successfully, False is not.
    """
    from zenml.analytics.context import AnalyticsContext

    if metadata is None:
        return False

    with AnalyticsContext() as analytics:
        return analytics.identify(traits=metadata)

track(event: AnalyticsEvent, metadata: Optional[Dict[str, Any]] = None) -> bool

Track segment event if user opted-in.

Parameters:

Name	Type	Description	Default
event	AnalyticsEvent	Name of event to track in segment.	required
metadata	Optional[Dict[str, Any]]	Dict of metadata to track.	None

Returns:

Type	Description
bool	True if event is sent successfully, False if not.

Source code in src/zenml/analytics/__init__.py

def track(  # type: ignore[return]
    event: "AnalyticsEvent",
    metadata: Optional[Dict[str, Any]] = None,
) -> bool:
    """Track segment event if user opted-in.

    Args:
        event: Name of event to track in segment.
        metadata: Dict of metadata to track.

    Returns:
        True if event is sent successfully, False if not.
    """
    from zenml.analytics.context import AnalyticsContext

    if metadata is None:
        metadata = {}

    metadata.setdefault("event_success", True)

    with AnalyticsContext() as analytics:
        return analytics.track(event=event, properties=metadata)

Modules

client

The analytics client of ZenML.

Classes

Client(send: bool = True, timeout: int = 15)

Bases: object

The client class for ZenML analytics.

Initialization of the client.

Parameters:

Name	Type	Description	Default
send	bool	Flag to determine whether to send the message.	True
timeout	int	Timeout in seconds.	15

Source code in src/zenml/analytics/client.py

def __init__(self, send: bool = True, timeout: int = 15) -> None:
    """Initialization of the client.

    Args:
        send: Flag to determine whether to send the message.
        timeout: Timeout in seconds.
    """
    self.send = send
    self.timeout = timeout

Functions

alias(user_id: UUID, previous_id: UUID) -> Tuple[bool, str]

Method to alias user IDs.

Parameters:

Name	Type	Description	Default
user_id	UUID	The user ID.	required
previous_id	UUID	Previous ID for the alias.	required

Returns:

Type	Description
Tuple[bool, str]	Tuple (success flag, the original message).

Source code in src/zenml/analytics/client.py

def alias(self, user_id: UUID, previous_id: UUID) -> Tuple[bool, str]:
    """Method to alias user IDs.

    Args:
        user_id: The user ID.
        previous_id: Previous ID for the alias.

    Returns:
        Tuple (success flag, the original message).
    """
    msg = {
        "user_id": user_id,
        "previous_id": previous_id,
        "type": "alias",
        "debug": IS_DEBUG_ENV,
    }
    return self._enqueue(json.dumps(msg, cls=AnalyticsEncoder))

group(user_id: UUID, group_id: UUID, traits: Optional[Dict[Any, Any]]) -> Tuple[bool, str]

Method to group users.

Parameters:

Name	Type	Description	Default
user_id	UUID	The user ID.	required
group_id	UUID	The group ID.	required
traits	Optional[Dict[Any, Any]]	Traits to assign to the group.	required

Returns:

Type	Description
Tuple[bool, str]	Tuple (success flag, the original message).

Source code in src/zenml/analytics/client.py

def group(
    self, user_id: UUID, group_id: UUID, traits: Optional[Dict[Any, Any]]
) -> Tuple[bool, str]:
    """Method to group users.

    Args:
        user_id: The user ID.
        group_id: The group ID.
        traits: Traits to assign to the group.

    Returns:
        Tuple (success flag, the original message).
    """
    msg = {
        "user_id": user_id,
        "group_id": group_id,
        "traits": traits or {},
        "type": "group",
        "debug": IS_DEBUG_ENV,
    }
    return self._enqueue(json.dumps(msg, cls=AnalyticsEncoder))

identify(user_id: UUID, traits: Optional[Dict[Any, Any]]) -> Tuple[bool, str]

Method to identify a user with given traits.

Parameters:

Name	Type	Description	Default
user_id	UUID	The user ID.	required
traits	Optional[Dict[Any, Any]]	The traits for the identification process.	required

Returns:

Type	Description
Tuple[bool, str]	Tuple (success flag, the original message).

Source code in src/zenml/analytics/client.py

def identify(
    self, user_id: UUID, traits: Optional[Dict[Any, Any]]
) -> Tuple[bool, str]:
    """Method to identify a user with given traits.

    Args:
        user_id: The user ID.
        traits: The traits for the identification process.

    Returns:
        Tuple (success flag, the original message).
    """
    msg = {
        "user_id": user_id,
        "traits": traits or {},
        "type": "identify",
        "debug": IS_DEBUG_ENV,
    }
    return self._enqueue(json.dumps(msg, cls=AnalyticsEncoder))

track(user_id: UUID, event: AnalyticsEvent, properties: Optional[Dict[Any, Any]]) -> Tuple[bool, str]

Method to track events.

Parameters:

Name	Type	Description	Default
user_id	UUID	The user ID.	required
event	AnalyticsEvent	The type of the event.	required
properties	Optional[Dict[Any, Any]]	Dict of additional properties for the event.	required

Returns:

Type	Description
Tuple[bool, str]	Tuple (success flag, the original message).

Source code in src/zenml/analytics/client.py

def track(
    self,
    user_id: UUID,
    event: "AnalyticsEvent",
    properties: Optional[Dict[Any, Any]],
) -> Tuple[bool, str]:
    """Method to track events.

    Args:
        user_id: The user ID.
        event: The type of the event.
        properties: Dict of additional properties for the event.

    Returns:
        Tuple (success flag, the original message).
    """
    msg = {
        "user_id": user_id,
        "event": event,
        "properties": properties or {},
        "type": "track",
        "debug": IS_DEBUG_ENV,
    }
    return self._enqueue(json.dumps(msg, cls=AnalyticsEncoder))

Functions

context

The analytics module of ZenML.

This module is based on the 'analytics-python' package created by Segment. The base functionalities are adapted to work with the ZenML analytics server.

Classes

AnalyticsContext()

Client class for ZenML Analytics v2.

Initialization.

Use this as a context manager to ensure that analytics are initialized properly, only tracked when configured to do so and that any errors are handled gracefully.

Source code in src/zenml/analytics/context.py

def __init__(self) -> None:
    """Initialization.

    Use this as a context manager to ensure that analytics are initialized
    properly, only tracked when configured to do so and that any errors
    are handled gracefully.
    """
    self.analytics_opt_in: bool = False

    self.user_id: Optional[UUID] = None
    self.external_user_id: Optional[UUID] = None
    self.executed_by_service_account: Optional[bool] = None
    self.client_id: Optional[UUID] = None
    self.server_id: Optional[UUID] = None
    self.external_server_id: Optional[UUID] = None
    self.server_metadata: Optional[Dict[str, str]] = None

    self.database_type: Optional["ServerDatabaseType"] = None
    self.deployment_type: Optional["ServerDeploymentType"] = None

Attributes

in_server: bool property

Flag to check whether the code is running in a ZenML server.

Returns:

Type	Description
bool	True if running in a server, False otherwise.

Functions

alias(user_id: UUID, previous_id: UUID) -> bool

Alias user IDs.

Parameters:

Name	Type	Description	Default
user_id	UUID	The user ID.	required
previous_id	UUID	Previous ID for the alias.	required

Returns:

Type	Description
bool	True if alias information was sent, False otherwise.

Source code in src/zenml/analytics/context.py

def alias(self, user_id: UUID, previous_id: UUID) -> bool:
    """Alias user IDs.

    Args:
        user_id: The user ID.
        previous_id: Previous ID for the alias.

    Returns:
        True if alias information was sent, False otherwise.
    """
    success = False
    if self.analytics_opt_in:
        success, _ = default_client.alias(
            user_id=user_id,
            previous_id=previous_id,
        )

    return success

group(group_id: UUID, traits: Optional[Dict[str, Any]] = None) -> bool

Group the user.

Parameters:

Name	Type	Description	Default
group_id	UUID	Group ID.	required
traits	Optional[Dict[str, Any]]	Traits of the group.	None

Returns:

Type	Description
bool	True if tracking information was sent, False otherwise.

Source code in src/zenml/analytics/context.py

def group(
    self,
    group_id: UUID,
    traits: Optional[Dict[str, Any]] = None,
) -> bool:
    """Group the user.

    Args:
        group_id: Group ID.
        traits: Traits of the group.

    Returns:
        True if tracking information was sent, False otherwise.
    """
    success = False
    if self.analytics_opt_in and self.user_id is not None:
        success, _ = default_client.group(
            user_id=self.user_id,
            group_id=group_id,
            traits=traits or {},
        )

    return success

identify(traits: Optional[Dict[str, Any]] = None) -> bool

Identify the user through segment.

Parameters:

Name	Type	Description	Default
traits	Optional[Dict[str, Any]]	Traits of the user.	None

Returns:

Type	Description
bool	True if tracking information was sent, False otherwise.

Source code in src/zenml/analytics/context.py

def identify(self, traits: Optional[Dict[str, Any]] = None) -> bool:
    """Identify the user through segment.

    Args:
        traits: Traits of the user.

    Returns:
        True if tracking information was sent, False otherwise.
    """
    success = False
    if self.analytics_opt_in and self.user_id is not None:
        success, _ = default_client.identify(
            user_id=self.user_id,
            traits=traits,
        )

    return success

track(event: AnalyticsEvent, properties: Optional[Dict[str, Any]] = None) -> bool

Track an event.

Parameters:

Name	Type	Description	Default
event	AnalyticsEvent	Event to track.	required
properties	Optional[Dict[str, Any]]	Event properties.	None

Returns:

Type	Description
bool	True if tracking information was sent, False otherwise.

Source code in src/zenml/analytics/context.py

def track(
    self,
    event: "AnalyticsEvent",
    properties: Optional[Dict[str, Any]] = None,
) -> bool:
    """Track an event.

    Args:
        event: Event to track.
        properties: Event properties.

    Returns:
        True if tracking information was sent, False otherwise.
    """
    from zenml.analytics.enums import AnalyticsEvent

    if properties is None:
        properties = {}

    if (
        not self.analytics_opt_in
        and event.value
        not in {
            AnalyticsEvent.OPT_OUT_ANALYTICS,
            AnalyticsEvent.OPT_IN_ANALYTICS,
        }
        or self.user_id is None
    ):
        return False

    # add basics
    properties.update(Environment.get_system_info())
    properties.update(
        {
            "environment": get_environment(),
            "python_version": Environment.python_version(),
            "version": __version__,
            "client_id": str(self.client_id),
            "user_id": str(self.user_id),
            "server_id": str(self.server_id),
            "deployment_type": str(self.deployment_type),
            "database_type": str(self.database_type),
            "executed_by_service_account": self.executed_by_service_account,
        }
    )

    try:
        # Timezone as tzdata
        tz = utc_now_tz_aware().astimezone().tzname()
        if tz is not None:
            properties.update({"timezone": tz})

        # Language code such as "en_DE"
        language_code, encoding = locale.getlocale()
        if language_code is not None:
            properties.update({"locale": language_code})
    except Exception:
        pass

    if self.external_user_id:
        properties["external_user_id"] = self.external_user_id

    if self.external_server_id:
        properties["external_server_id"] = self.external_server_id

    if self.server_metadata:
        properties.update(self.server_metadata)

    if self.project_id:
        properties.setdefault("project_id", str(self.project_id))

    for k, v in properties.items():
        if isinstance(v, UUID):
            properties[k] = str(v)

    success, _ = default_client.track(
        user_id=self.user_id,
        event=event,
        properties=properties,
    )

    logger.debug(
        f"Sending analytics: User: {self.user_id}, Event: {event}, "
        f"Metadata: {properties}"
    )

    return success

Functions

enums

Collection of analytics events for ZenML.

Classes

AnalyticsEvent

Bases: str, Enum

Enum of events to track in segment.

models

Helper models for ZenML analytics.

Classes

AnalyticsTrackedModelMixin

Bases: BaseModel

Mixin for models that are tracked through analytics events.

Classes that have information tracked in analytics events can inherit from this mixin and implement the abstract methods. The @track_method decorator will detect function arguments and return values that inherit from this class and will include the ANALYTICS_FIELDS attributes as tracking metadata.

Functions

get_analytics_metadata() -> Dict[str, Any]

Get the analytics metadata for the model.

Returns:

Type	Description
Dict[str, Any]	Dict of analytics metadata.

Source code in src/zenml/analytics/models.py

def get_analytics_metadata(self) -> Dict[str, Any]:
    """Get the analytics metadata for the model.

    Returns:
        Dict of analytics metadata.
    """
    metadata = {}
    for field_name in self.ANALYTICS_FIELDS:
        metadata[field_name] = getattr(self, field_name, None)
    return metadata

request

The 'analytics' module of ZenML.

This module is based on the 'analytics-python' package created by Segment. The base functionalities are adapted to work with the ZenML analytics server.

Classes

Functions

post(batch: List[str], timeout: int = 15) -> requests.Response

Post a batch of messages to the ZenML analytics server.

Parameters:

Name	Type	Description	Default
batch	List[str]	The messages to send.	required
timeout	int	Timeout in seconds.	15

Returns:

Type	Description
Response	The response.

Raises:

Type	Description
AnalyticsAPIError	If the post request has failed.

Source code in src/zenml/analytics/request.py

def post(batch: List[str], timeout: int = 15) -> requests.Response:
    """Post a batch of messages to the ZenML analytics server.

    Args:
        batch: The messages to send.
        timeout: Timeout in seconds.

    Returns:
        The response.

    Raises:
        AnalyticsAPIError: If the post request has failed.
    """
    from zenml.analytics import source_context

    headers = {
        "accept": "application/json",
        "content-type": "application/json",
        source_context.name: source_context.get().value,
    }
    response = requests.post(
        url=ANALYTICS_SERVER_URL + "/batch",
        headers=headers,
        data=f"[{','.join(batch)}]",
        timeout=timeout,
    )

    if response.status_code == 200:
        logger.debug("data uploaded successfully")
        return response

    raise AnalyticsAPIError(response.status_code, response.text)

utils

Utility functions and classes for ZenML analytics.

Classes

AnalyticsAPIError(status: int, message: str)

Bases: Exception

Custom exception class for API-related errors.

Initialization.

Parameters:

Name	Type	Description	Default
status	int	The status code of the response.	required
message	str	The text of the response.	required

Source code in src/zenml/analytics/utils.py

def __init__(self, status: int, message: str) -> None:
    """Initialization.

    Args:
        status: The status code of the response.
        message: The text of the response.
    """
    self.message = message
    self.status = status

Functions

AnalyticsEncoder

Bases: JSONEncoder

Helper encoder class for JSON serialization.

Functions

default(obj: Any) -> Any

The default method to handle UUID and 'AnalyticsEvent' objects.

Parameters:

Name	Type	Description	Default
obj	Any	The object to encode.	required

Returns:

Type	Description
Any	The encoded object.

Source code in src/zenml/analytics/utils.py

def default(self, obj: Any) -> Any:
    """The default method to handle UUID and 'AnalyticsEvent' objects.

    Args:
        obj: The object to encode.

    Returns:
        The encoded object.
    """
    from zenml.analytics.enums import AnalyticsEvent

    # If the object is UUID, we simply return the value of UUID
    if isinstance(obj, UUID):
        return str(obj)

    # If the object is an AnalyticsEvent, return its value
    elif isinstance(obj, AnalyticsEvent):
        return str(obj.value)

    return json.JSONEncoder.default(self, obj)

analytics_disabler()

Context manager which disables ZenML analytics.

Initialization of the context manager.

Source code in src/zenml/analytics/utils.py

def __init__(self) -> None:
    """Initialization of the context manager."""
    self.original_value: Optional[bool] = None

Functions

track_handler(event: AnalyticsEvent, metadata: Optional[Dict[str, Any]] = None)

Bases: object

Context handler to enable tracking the success status of an event.

Initialization of the context manager.

Parameters:

Name	Type	Description	Default
event	AnalyticsEvent	The type of the analytics event	required
metadata	Optional[Dict[str, Any]]	The metadata of the event.	None

Source code in src/zenml/analytics/utils.py

def __init__(
    self,
    event: AnalyticsEvent,
    metadata: Optional[Dict[str, Any]] = None,
):
    """Initialization of the context manager.

    Args:
        event: The type of the analytics event
        metadata: The metadata of the event.
    """
    self.event: AnalyticsEvent = event
    self.metadata: Dict[str, Any] = metadata or {}

Functions

Functions

email_opt_int(opted_in: bool, email: Optional[str], source: str) -> None

Track the event of the users response to the email prompt, identify them.

Parameters:

Name	Type	Description	Default
opted_in	bool	Did the user decide to opt-in	required
email	Optional[str]	The email the user optionally provided	required
source	str	Location when the user replied ["zenml go", "zenml server"]	required

Source code in src/zenml/analytics/utils.py

def email_opt_int(opted_in: bool, email: Optional[str], source: str) -> None:
    """Track the event of the users response to the email prompt, identify them.

    Args:
        opted_in: Did the user decide to opt-in
        email: The email the user optionally provided
        source: Location when the user replied ["zenml go", "zenml server"]
    """
    # If the user opted in, associate email with the anonymous distinct ID
    if opted_in and email is not None and email != "":
        identify(metadata={"email": email, "source": source})

    # Track that the user answered the prompt
    track(
        AnalyticsEvent.OPT_IN_OUT_EMAIL,
        {"opted_in": opted_in, "source": source},
    )

is_analytics_disabled_internally() -> bool

Whether analytics are disabled by an internal helper function.

Returns:

Type	Description
bool	Whether analytics are disabled by an internal helper function.

Source code in src/zenml/analytics/utils.py

def is_analytics_disabled_internally() -> bool:
    """Whether analytics are disabled by an internal helper function.

    Returns:
        Whether analytics are disabled by an internal helper function.
    """
    return _INTERNAL_DISABLE_ANALYTICS

track_decorator(event: AnalyticsEvent) -> Callable[[F], F]

Decorator to track event.

If the decorated function takes in a AnalyticsTrackedModelMixin object as an argument or returns one, it will be called to track the event. The return value takes precedence over the argument when determining which object is called to track the event.

If the decorated function is a method of a class that inherits from AnalyticsTrackerMixin, the parent object will be used to intermediate tracking analytics.

Parameters:

Name	Type	Description	Default
event	AnalyticsEvent	Event string to stamp with.	required

Returns:

Type	Description
Callable[[F], F]	A decorator that applies the analytics tracking to a function.

Source code in src/zenml/analytics/utils.py

def track_decorator(event: AnalyticsEvent) -> Callable[[F], F]:
    """Decorator to track event.

    If the decorated function takes in a `AnalyticsTrackedModelMixin` object as
    an argument or returns one, it will be called to track the event. The return
    value takes precedence over the argument when determining which object is
    called to track the event.

    If the decorated function is a method of a class that inherits from
    `AnalyticsTrackerMixin`, the parent object will be used to intermediate
    tracking analytics.

    Args:
        event: Event string to stamp with.

    Returns:
        A decorator that applies the analytics tracking to a function.
    """

    def inner_decorator(func: F) -> F:
        """Inner decorator function.

        Args:
            func: Function to decorate.

        Returns:
            Decorated function.
        """

        @wraps(func)
        def inner_func(*args: Any, **kwargs: Any) -> Any:
            """Inner function.

            Args:
                *args: Arguments to be passed to the function.
                **kwargs: Keyword arguments to be passed to the function.

            Returns:
                Result of the function.
            """
            with track_handler(event=event) as handler:
                try:
                    for obj in list(args) + list(kwargs.values()):
                        if isinstance(obj, AnalyticsTrackedModelMixin):
                            handler.metadata = obj.get_analytics_metadata()
                            break
                except Exception as e:
                    logger.debug(f"Analytics tracking failure for {func}: {e}")

                result = func(*args, **kwargs)

                try:
                    if isinstance(result, AnalyticsTrackedModelMixin):
                        handler.metadata = result.get_analytics_metadata()
                except Exception as e:
                    logger.debug(f"Analytics tracking failure for {func}: {e}")

                return result

        return cast(F, inner_func)

    return inner_decorator