Models
zenml.models
special
Pydantic models for the various concepts in ZenML.
artifact_models
Models representing artifacts.
ArtifactBaseModel (BaseModel)
pydantic-model
Base model for artifacts.
Source code in zenml/models/artifact_models.py
class ArtifactBaseModel(BaseModel):
"""Base model for artifacts."""
name: str # Name of the output in the parent step
artifact_store_id: Optional[UUID]
type: ArtifactType
uri: str
materializer: str
data_type: str
ArtifactRequestModel (ArtifactBaseModel, ProjectScopedRequestModel)
pydantic-model
Request model for artifacts.
Source code in zenml/models/artifact_models.py
class ArtifactRequestModel(ArtifactBaseModel, ProjectScopedRequestModel):
"""Request model for artifacts."""
ArtifactResponseModel (ArtifactBaseModel, ProjectScopedResponseModel)
pydantic-model
Response model for artifacts.
Source code in zenml/models/artifact_models.py
class ArtifactResponseModel(ArtifactBaseModel, ProjectScopedResponseModel):
"""Response model for artifacts."""
producer_step_run_id: Optional[UUID]
ArtifactUpdateModel (ArtifactRequestModel)
pydantic-model
Update model for artifacts.
Source code in zenml/models/artifact_models.py
class ArtifactUpdateModel(ArtifactRequestModel):
"""Update model for artifacts."""
base_models
Base domain model definitions.
BaseRequestModel (AnalyticsTrackedModelMixin)
pydantic-model
Base request model.
Used as a base class for all request models.
Source code in zenml/models/base_models.py
class BaseRequestModel(AnalyticsTrackedModelMixin):
"""Base request model.
Used as a base class for all request models.
"""
BaseResponseModel (AnalyticsTrackedModelMixin)
pydantic-model
Base domain model.
Used as a base class for all domain models that have the following common characteristics:
- are uniquely identified by a UUID
- have a creation timestamp and a last modified timestamp
Source code in zenml/models/base_models.py
class BaseResponseModel(AnalyticsTrackedModelMixin):
"""Base domain model.
Used as a base class for all domain models that have the following common
characteristics:
* are uniquely identified by a UUID
* have a creation timestamp and a last modified timestamp
"""
id: UUID = Field(title="The unique resource id.")
created: datetime = Field(title="Time when this resource was created.")
updated: datetime = Field(title="Time when this resource was last updated.")
def __hash__(self) -> int:
"""Implementation of hash magic method.
Returns:
Hash of the UUID.
"""
return hash((type(self),) + tuple([self.id]))
def __eq__(self, other: Any) -> bool:
"""Implementation of equality magic method.
Args:
other: The other object to compare to.
Returns:
True if the other object is of the same type and has the same UUID.
"""
if isinstance(other, BaseResponseModel):
return self.id == other.id
else:
return False
__eq__(self, other)
special
Implementation of equality magic method.
Parameters:
Name | Type | Description | Default |
---|---|---|---|
other |
Any |
The other object to compare to. |
required |
Returns:
Type | Description |
---|---|
bool |
True if the other object is of the same type and has the same UUID. |
Source code in zenml/models/base_models.py
def __eq__(self, other: Any) -> bool:
"""Implementation of equality magic method.
Args:
other: The other object to compare to.
Returns:
True if the other object is of the same type and has the same UUID.
"""
if isinstance(other, BaseResponseModel):
return self.id == other.id
else:
return False
__hash__(self)
special
Implementation of hash magic method.
Returns:
Type | Description |
---|---|
int |
Hash of the UUID. |
Source code in zenml/models/base_models.py
def __hash__(self) -> int:
"""Implementation of hash magic method.
Returns:
Hash of the UUID.
"""
return hash((type(self),) + tuple([self.id]))
ProjectScopedRequestModel (UserScopedRequestModel)
pydantic-model
Base project-scoped request domain model.
Used as a base class for all domain models that are project-scoped.
Source code in zenml/models/base_models.py
class ProjectScopedRequestModel(UserScopedRequestModel):
"""Base project-scoped request domain model.
Used as a base class for all domain models that are project-scoped.
"""
project: UUID = Field(title="The project to which this resource belongs.")
ProjectScopedResponseModel (UserScopedResponseModel)
pydantic-model
Base project-scoped domain model.
Used as a base class for all domain models that are project-scoped.
Source code in zenml/models/base_models.py
class ProjectScopedResponseModel(UserScopedResponseModel):
"""Base project-scoped domain model.
Used as a base class for all domain models that are project-scoped.
"""
project: "ProjectResponseModel" = Field(
title="The project of this resource."
)
def get_analytics_metadata(self) -> Dict[str, Any]:
"""Fetches the analytics metadata for project scoped models.
Returns:
The analytics metadata.
"""
metadata = super().get_analytics_metadata()
metadata["project"] = self.project.id
return metadata
get_analytics_metadata(self)
Fetches the analytics metadata for project scoped models.
Returns:
Type | Description |
---|---|
Dict[str, Any] |
The analytics metadata. |
Source code in zenml/models/base_models.py
def get_analytics_metadata(self) -> Dict[str, Any]:
"""Fetches the analytics metadata for project scoped models.
Returns:
The analytics metadata.
"""
metadata = super().get_analytics_metadata()
metadata["project"] = self.project.id
return metadata
ShareableRequestModel (ProjectScopedRequestModel)
pydantic-model
Base shareable project-scoped domain model.
Used as a base class for all domain models that are project-scoped and are shareable.
Source code in zenml/models/base_models.py
class ShareableRequestModel(ProjectScopedRequestModel):
"""Base shareable project-scoped domain model.
Used as a base class for all domain models that are project-scoped and are
shareable.
"""
is_shared: bool = Field(
default=False,
title=(
"Flag describing if this resource is shared with other users in "
"the same project."
),
)
ShareableResponseModel (ProjectScopedResponseModel)
pydantic-model
Base shareable project-scoped domain model.
Used as a base class for all domain models that are project-scoped and are shareable.
Source code in zenml/models/base_models.py
class ShareableResponseModel(ProjectScopedResponseModel):
"""Base shareable project-scoped domain model.
Used as a base class for all domain models that are project-scoped and are
shareable.
"""
is_shared: bool = Field(
title=(
"Flag describing if this resource is shared with other users in "
"the same project."
),
)
UserScopedRequestModel (BaseRequestModel)
pydantic-model
Base user-owned request model.
Used as a base class for all domain models that are "owned" by a user.
Source code in zenml/models/base_models.py
class UserScopedRequestModel(BaseRequestModel):
"""Base user-owned request model.
Used as a base class for all domain models that are "owned" by a user.
"""
user: UUID = Field(title="The id of the user that created this resource.")
UserScopedResponseModel (BaseResponseModel)
pydantic-model
Base user-owned domain model.
Used as a base class for all domain models that are "owned" by a user.
Source code in zenml/models/base_models.py
class UserScopedResponseModel(BaseResponseModel):
"""Base user-owned domain model.
Used as a base class for all domain models that are "owned" by a user.
"""
user: Optional["UserResponseModel"] = Field(
title="The user that created this resource.", nullable=True
)
def get_analytics_metadata(self) -> Dict[str, Any]:
"""Fetches the analytics metadata for user scoped models.
Returns:
The analytics metadata.
"""
metadata = super().get_analytics_metadata()
if self.user is not None:
metadata["user"] = self.user.id
return metadata
get_analytics_metadata(self)
Fetches the analytics metadata for user scoped models.
Returns:
Type | Description |
---|---|
Dict[str, Any] |
The analytics metadata. |
Source code in zenml/models/base_models.py
def get_analytics_metadata(self) -> Dict[str, Any]:
"""Fetches the analytics metadata for user scoped models.
Returns:
The analytics metadata.
"""
metadata = super().get_analytics_metadata()
if self.user is not None:
metadata["user"] = self.user.id
return metadata
update_model(_cls)
Base update model.
This is used as a decorator on top of request models to convert them into update models where the fields are optional and can be set to None.
Parameters:
Name | Type | Description | Default |
---|---|---|---|
_cls |
Type[~T] |
The class to decorate |
required |
Returns:
Type | Description |
---|---|
Type[~T] |
The decorated class. |
Source code in zenml/models/base_models.py
def update_model(_cls: Type[T]) -> Type[T]:
"""Base update model.
This is used as a decorator on top of request models to convert them
into update models where the fields are optional and can be set to None.
Args:
_cls: The class to decorate
Returns:
The decorated class.
"""
for _, value in _cls.__fields__.items():
value.required = False
value.allow_none = True
return _cls
component_models
Models representing stack components.
ComponentBaseModel (BaseModel)
pydantic-model
Base model for stack components.
Source code in zenml/models/component_models.py
class ComponentBaseModel(BaseModel):
"""Base model for stack components."""
name: str = Field(
title="The name of the stack component.",
max_length=MODEL_NAME_FIELD_MAX_LENGTH,
)
type: StackComponentType = Field(
title="The type of the stack component.",
)
flavor: str = Field(
title="The flavor of the stack component.",
)
configuration: Dict[str, Any] = Field(
title="The stack component configuration.",
)
ComponentRequestModel (ComponentBaseModel, ShareableRequestModel)
pydantic-model
Request model for stack components.
Source code in zenml/models/component_models.py
class ComponentRequestModel(ComponentBaseModel, ShareableRequestModel):
"""Request model for stack components."""
ANALYTICS_FIELDS: ClassVar[List[str]] = [
"type",
"flavor",
"project",
"user",
"is_shared",
]
@validator("name")
def name_cant_be_a_secret_reference(cls, name: str) -> str:
"""Validator to ensure that the given name is not a secret reference.
Args:
name: The name to validate.
Returns:
The name if it is not a secret reference.
Raises:
ValueError: If the name is a secret reference.
"""
if secret_utils.is_secret_reference(name):
raise ValueError(
"Passing the `name` attribute of a stack component as a "
"secret reference is not allowed."
)
return name
name_cant_be_a_secret_reference(name)
classmethod
Validator to ensure that the given name is not a secret reference.
Parameters:
Name | Type | Description | Default |
---|---|---|---|
name |
str |
The name to validate. |
required |
Returns:
Type | Description |
---|---|
str |
The name if it is not a secret reference. |
Exceptions:
Type | Description |
---|---|
ValueError |
If the name is a secret reference. |
Source code in zenml/models/component_models.py
@validator("name")
def name_cant_be_a_secret_reference(cls, name: str) -> str:
"""Validator to ensure that the given name is not a secret reference.
Args:
name: The name to validate.
Returns:
The name if it is not a secret reference.
Raises:
ValueError: If the name is a secret reference.
"""
if secret_utils.is_secret_reference(name):
raise ValueError(
"Passing the `name` attribute of a stack component as a "
"secret reference is not allowed."
)
return name
ComponentResponseModel (ComponentBaseModel, ShareableResponseModel)
pydantic-model
Response model for stack components.
Source code in zenml/models/component_models.py
class ComponentResponseModel(ComponentBaseModel, ShareableResponseModel):
"""Response model for stack components."""
ANALYTICS_FIELDS: ClassVar[List[str]] = [
"id",
"type",
"flavor",
"project",
"user",
"is_shared",
]
ComponentUpdateModel (ComponentRequestModel)
pydantic-model
Update model for stack components.
Source code in zenml/models/component_models.py
class ComponentUpdateModel(ComponentRequestModel):
"""Update model for stack components."""
constants
Constants used by ZenML domain models.
flavor_models
Models representing stack component flavors.
FlavorBaseModel (BaseModel)
pydantic-model
Base model for stack component flavors.
Source code in zenml/models/flavor_models.py
class FlavorBaseModel(BaseModel):
"""Base model for stack component flavors."""
name: str = Field(
title="The name of the Flavor.",
)
type: StackComponentType = Field(
title="The type of the Flavor.",
)
config_schema: str = Field(
title="The JSON schema of this flavor's corresponding configuration.",
max_length=MODEL_CONFIG_SCHEMA_MAX_LENGTH,
)
source: str = Field(
title="The path to the module which contains this Flavor."
)
integration: Optional[str] = Field(
title="The name of the integration that the Flavor belongs to."
)
FlavorRequestModel (FlavorBaseModel, ProjectScopedRequestModel)
pydantic-model
Request model for stack component flavors.
Source code in zenml/models/flavor_models.py
class FlavorRequestModel(FlavorBaseModel, ProjectScopedRequestModel):
"""Request model for stack component flavors."""
ANALYTICS_FIELDS: ClassVar[List[str]] = [
"type",
"integration",
"project",
"user",
]
FlavorResponseModel (FlavorBaseModel, ProjectScopedResponseModel)
pydantic-model
Response model for stack component flavors.
Source code in zenml/models/flavor_models.py
class FlavorResponseModel(FlavorBaseModel, ProjectScopedResponseModel):
"""Response model for stack component flavors."""
ANALYTICS_FIELDS: ClassVar[List[str]] = [
"id",
"type",
"integration",
"project",
"user",
]
pipeline_models
Models representing pipelines.
PipelineBaseModel (BaseModel)
pydantic-model
Base model for pipelines.
Source code in zenml/models/pipeline_models.py
class PipelineBaseModel(BaseModel):
"""Base model for pipelines."""
name: str = Field(
title="The name of the pipeline.",
max_length=MODEL_NAME_FIELD_MAX_LENGTH,
)
docstring: Optional[str]
spec: PipelineSpec
PipelineRequestModel (PipelineBaseModel, ProjectScopedRequestModel)
pydantic-model
Pipeline request model.
Source code in zenml/models/pipeline_models.py
class PipelineRequestModel(PipelineBaseModel, ProjectScopedRequestModel):
"""Pipeline request model."""
ANALYTICS_FIELDS: ClassVar[List[str]] = ["project", "user"]
PipelineResponseModel (PipelineBaseModel, ProjectScopedResponseModel)
pydantic-model
Pipeline response model user, project, runs, and status hydrated.
Source code in zenml/models/pipeline_models.py
class PipelineResponseModel(PipelineBaseModel, ProjectScopedResponseModel):
"""Pipeline response model user, project, runs, and status hydrated."""
ANALYTICS_FIELDS: ClassVar[List[str]] = ["id", "project", "user"]
runs: Optional[List["PipelineRunResponseModel"]] = Field(
title="A list of the last x Pipeline Runs."
)
status: Optional[List[ExecutionStatus]] = Field(
title="The status of the last x Pipeline Runs."
)
PipelineUpdateModel (PipelineRequestModel)
pydantic-model
Pipeline update model.
Source code in zenml/models/pipeline_models.py
class PipelineUpdateModel(PipelineRequestModel):
"""Pipeline update model."""
pipeline_run_models
Models representing pipeline runs.
PipelineRunBaseModel (BaseModel)
pydantic-model
Base model for pipeline runs.
Source code in zenml/models/pipeline_run_models.py
class PipelineRunBaseModel(BaseModel):
"""Base model for pipeline runs."""
name: str = Field(
title="The name of the pipeline run.",
max_length=MODEL_NAME_FIELD_MAX_LENGTH,
)
orchestrator_run_id: Optional[str] = None
enable_cache: Optional[bool]
start_time: Optional[datetime]
end_time: Optional[datetime]
status: ExecutionStatus
pipeline_configuration: Dict[str, Any]
num_steps: Optional[int]
zenml_version: Optional[str] = current_zenml_version
git_sha: Optional[str] = Field(default_factory=get_git_sha)
PipelineRunRequestModel (PipelineRunBaseModel, ProjectScopedRequestModel)
pydantic-model
Pipeline run model with user, project, pipeline, and stack as UUIDs.
Source code in zenml/models/pipeline_run_models.py
class PipelineRunRequestModel(PipelineRunBaseModel, ProjectScopedRequestModel):
"""Pipeline run model with user, project, pipeline, and stack as UUIDs."""
id: UUID
stack: Optional[UUID] # Might become None if the stack is deleted.
pipeline: Optional[UUID] # Unlisted runs have this as None.
PipelineRunResponseModel (PipelineRunBaseModel, ProjectScopedResponseModel)
pydantic-model
Pipeline run model with user, project, pipeline, and stack hydrated.
Source code in zenml/models/pipeline_run_models.py
class PipelineRunResponseModel(
PipelineRunBaseModel, ProjectScopedResponseModel
):
"""Pipeline run model with user, project, pipeline, and stack hydrated."""
pipeline: Optional["PipelineResponseModel"] = Field(
title="The pipeline this run belongs to."
)
stack: Optional["StackResponseModel"] = Field(
title="The stack that was used for this run."
)
PipelineRunUpdateModel (BaseModel)
pydantic-model
Pipeline run update model.
Source code in zenml/models/pipeline_run_models.py
class PipelineRunUpdateModel(BaseModel):
"""Pipeline run update model."""
status: Optional[ExecutionStatus] = None
end_time: Optional[datetime] = None
get_git_sha(clean=True)
Returns the current git HEAD SHA.
If the current working directory is not inside a git repo, this will return
None
.
Parameters:
Name | Type | Description | Default |
---|---|---|---|
clean |
bool |
If |
True |
Returns:
Type | Description |
---|---|
Optional[str] |
The current git HEAD SHA or |
Source code in zenml/models/pipeline_run_models.py
def get_git_sha(clean: bool = True) -> Optional[str]:
"""Returns the current git HEAD SHA.
If the current working directory is not inside a git repo, this will return
`None`.
Args:
clean: If `True` and there any untracked files or files in the index or
working tree, this function will return `None`.
Returns:
The current git HEAD SHA or `None` if the current working directory is
not inside a git repo.
"""
try:
from git.exc import InvalidGitRepositoryError
from git.repo.base import Repo
except ImportError:
return None
try:
repo = Repo(search_parent_directories=True)
except InvalidGitRepositoryError:
return None
if clean and repo.is_dirty(untracked_files=True):
return None
return cast(str, repo.head.object.hexsha)
project_models
Models representing projects.
ProjectBaseModel (BaseModel)
pydantic-model
Base model for projects.
Source code in zenml/models/project_models.py
class ProjectBaseModel(BaseModel):
"""Base model for projects."""
name: str = Field(
title="The unique name of the project.",
max_length=MODEL_NAME_FIELD_MAX_LENGTH,
)
description: str = Field(
default="",
title="The description of the project.",
max_length=MODEL_DESCRIPTIVE_FIELD_MAX_LENGTH,
)
ProjectRequestModel (ProjectBaseModel, BaseRequestModel)
pydantic-model
Request model for projects.
Source code in zenml/models/project_models.py
class ProjectRequestModel(ProjectBaseModel, BaseRequestModel):
"""Request model for projects."""
ProjectResponseModel (ProjectBaseModel, BaseResponseModel)
pydantic-model
Response model for projects.
Source code in zenml/models/project_models.py
class ProjectResponseModel(ProjectBaseModel, BaseResponseModel):
"""Response model for projects."""
ANALYTICS_FIELDS: ClassVar[List[str]] = ["id"]
ProjectUpdateModel (ProjectRequestModel)
pydantic-model
Update model for projects.
Source code in zenml/models/project_models.py
class ProjectUpdateModel(ProjectRequestModel):
"""Update model for projects."""
role_assignment_models
Models representing role assignments.
RoleAssignmentBaseModel (BaseModel)
pydantic-model
Base model for role assignments.
Source code in zenml/models/role_assignment_models.py
class RoleAssignmentBaseModel(BaseModel):
"""Base model for role assignments."""
@root_validator(pre=True)
def check_team_or_user(cls, values: Dict[str, Any]) -> Dict[str, Any]:
"""Check that either a team or a user is set.
Args:
values: The values to validate.
Returns:
The validated values.
Raises:
ValueError: If both or neither team and user are set.
"""
if not values.get("team") and not values.get("user"):
raise ValueError("Either team or user is required")
elif values.get("team") and values.get("user"):
raise ValueError("An assignment can not contain a user and team")
return values
check_team_or_user(values)
classmethod
Check that either a team or a user is set.
Parameters:
Name | Type | Description | Default |
---|---|---|---|
values |
Dict[str, Any] |
The values to validate. |
required |
Returns:
Type | Description |
---|---|
Dict[str, Any] |
The validated values. |
Exceptions:
Type | Description |
---|---|
ValueError |
If both or neither team and user are set. |
Source code in zenml/models/role_assignment_models.py
@root_validator(pre=True)
def check_team_or_user(cls, values: Dict[str, Any]) -> Dict[str, Any]:
"""Check that either a team or a user is set.
Args:
values: The values to validate.
Returns:
The validated values.
Raises:
ValueError: If both or neither team and user are set.
"""
if not values.get("team") and not values.get("user"):
raise ValueError("Either team or user is required")
elif values.get("team") and values.get("user"):
raise ValueError("An assignment can not contain a user and team")
return values
RoleAssignmentRequestModel (RoleAssignmentBaseModel, BaseRequestModel)
pydantic-model
Request model for role assignments using UUIDs for all entities.
Source code in zenml/models/role_assignment_models.py
class RoleAssignmentRequestModel(RoleAssignmentBaseModel, BaseRequestModel):
"""Request model for role assignments using UUIDs for all entities."""
project: Optional[UUID] = Field(
None, title="The project that the role is limited to."
)
team: Optional[UUID] = Field(
None, title="The team that the role is assigned to."
)
user: Optional[UUID] = Field(
None, title="The user that the role is assigned to."
)
role: UUID = Field(title="The role.")
RoleAssignmentResponseModel (RoleAssignmentBaseModel, BaseResponseModel)
pydantic-model
Response model for role assignments with all entities hydrated.
Source code in zenml/models/role_assignment_models.py
class RoleAssignmentResponseModel(RoleAssignmentBaseModel, BaseResponseModel):
"""Response model for role assignments with all entities hydrated."""
project: Optional["ProjectResponseModel"] = Field(
title="The project scope of this role assignment.", default=None
)
team: Optional["TeamResponseModel"] = Field(
title="The team the role is assigned to.", default=None
)
user: Optional["UserResponseModel"] = Field(
title="The team the role is assigned to.", default=None
)
role: "RoleResponseModel" = Field(
title="The team the role is assigned to.", default=None
)
role_models
Models representing roles that can be assigned to users or teams.
RoleBaseModel (BaseModel)
pydantic-model
Base model for roles.
Source code in zenml/models/role_models.py
class RoleBaseModel(BaseModel):
"""Base model for roles."""
name: str = Field(
title="The unique name of the role.",
max_length=MODEL_NAME_FIELD_MAX_LENGTH,
)
permissions: Set[PermissionType]
RoleRequestModel (RoleBaseModel, BaseRequestModel)
pydantic-model
Request model for roles.
Source code in zenml/models/role_models.py
class RoleRequestModel(RoleBaseModel, BaseRequestModel):
"""Request model for roles."""
RoleResponseModel (RoleBaseModel, BaseResponseModel)
pydantic-model
Response model for roles.
Source code in zenml/models/role_models.py
class RoleResponseModel(RoleBaseModel, BaseResponseModel):
"""Response model for roles."""
ANALYTICS_FIELDS: ClassVar[List[str]] = ["id"]
RoleUpdateModel (RoleRequestModel)
pydantic-model
Update model for roles.
Source code in zenml/models/role_models.py
class RoleUpdateModel(RoleRequestModel):
"""Update model for roles."""
server_models
Model definitions for ZenML servers.
ServerDatabaseType (StrEnum)
Enum for server database types.
Source code in zenml/models/server_models.py
class ServerDatabaseType(StrEnum):
"""Enum for server database types."""
SQLITE = "sqlite"
MYSQL = "mysql"
OTHER = "other"
ServerDeploymentType (StrEnum)
Enum for server deployment types.
Source code in zenml/models/server_models.py
class ServerDeploymentType(StrEnum):
"""Enum for server deployment types."""
LOCAL = "local"
DOCKER = "docker"
KUBERNETES = "kubernetes"
AWS = "aws"
GCP = "gcp"
AZURE = "azure"
OTHER = "other"
ServerModel (BaseModel)
pydantic-model
Domain model for ZenML servers.
Source code in zenml/models/server_models.py
class ServerModel(BaseModel):
"""Domain model for ZenML servers."""
id: UUID = Field(default_factory=uuid4, title="The unique server id.")
version: str = Field(
title="The ZenML version that the server is running.",
)
deployment_type: ServerDeploymentType = Field(
ServerDeploymentType.OTHER,
title="The ZenML server deployment type.",
)
database_type: ServerDatabaseType = Field(
ServerDatabaseType.OTHER,
title="The database type that the server is using.",
)
def is_local(self) -> bool:
"""Return whether the server is running locally.
Returns:
True if the server is running locally, False otherwise.
"""
from zenml.config.global_config import GlobalConfiguration
# Local ZenML servers are identifiable by the fact that their
# server ID is the same as the local client (user) ID.
return self.id == GlobalConfiguration().user_id
is_local(self)
Return whether the server is running locally.
Returns:
Type | Description |
---|---|
bool |
True if the server is running locally, False otherwise. |
Source code in zenml/models/server_models.py
def is_local(self) -> bool:
"""Return whether the server is running locally.
Returns:
True if the server is running locally, False otherwise.
"""
from zenml.config.global_config import GlobalConfiguration
# Local ZenML servers are identifiable by the fact that their
# server ID is the same as the local client (user) ID.
return self.id == GlobalConfiguration().user_id
stack_models
Models representing stacks.
StackBaseModel (BaseModel)
pydantic-model
Base model for stacks.
Source code in zenml/models/stack_models.py
class StackBaseModel(BaseModel):
"""Base model for stacks."""
name: str = Field(
title="The name of the stack.", max_length=MODEL_NAME_FIELD_MAX_LENGTH
)
description: str = Field(
default="",
title="The description of the stack",
max_length=MODEL_DESCRIPTIVE_FIELD_MAX_LENGTH,
)
StackRequestModel (StackBaseModel, ShareableRequestModel)
pydantic-model
Stack model with components, user and project as UUIDs.
Source code in zenml/models/stack_models.py
class StackRequestModel(StackBaseModel, ShareableRequestModel):
"""Stack model with components, user and project as UUIDs."""
ANALYTICS_FIELDS: ClassVar[List[str]] = [
"project",
"user",
"is_shared",
]
components: Dict[StackComponentType, List[UUID]] = Field(
title="A mapping of stack component types to the actual"
"instances of components of this type."
)
def get_analytics_metadata(self) -> Dict[str, Any]:
"""Add the stack components to the stack analytics metadata.
Returns:
Dict of analytics metadata.
"""
metadata = super().get_analytics_metadata()
metadata.update({ct: c[0] for ct, c in self.components.items()})
return metadata
@property
def is_valid(self) -> bool:
"""Check if the stack is valid.
Returns:
True if the stack is valid, False otherwise.
"""
return (
StackComponentType.ARTIFACT_STORE in self.components
and StackComponentType.ORCHESTRATOR in self.components
)
is_valid: bool
property
readonly
Check if the stack is valid.
Returns:
Type | Description |
---|---|
bool |
True if the stack is valid, False otherwise. |
get_analytics_metadata(self)
Add the stack components to the stack analytics metadata.
Returns:
Type | Description |
---|---|
Dict[str, Any] |
Dict of analytics metadata. |
Source code in zenml/models/stack_models.py
def get_analytics_metadata(self) -> Dict[str, Any]:
"""Add the stack components to the stack analytics metadata.
Returns:
Dict of analytics metadata.
"""
metadata = super().get_analytics_metadata()
metadata.update({ct: c[0] for ct, c in self.components.items()})
return metadata
StackResponseModel (StackBaseModel, ShareableResponseModel)
pydantic-model
Stack model with Components, User and Project fully hydrated.
Source code in zenml/models/stack_models.py
class StackResponseModel(StackBaseModel, ShareableResponseModel):
"""Stack model with Components, User and Project fully hydrated."""
ANALYTICS_FIELDS: ClassVar[List[str]] = [
"id",
"project",
"user",
"is_shared",
]
components: Dict[StackComponentType, List[ComponentResponseModel]] = Field(
title="A mapping of stack component types to the actual"
"instances of components of this type."
)
def get_analytics_metadata(self) -> Dict[str, Any]:
"""Add the stack components to the stack analytics metadata.
Returns:
Dict of analytics metadata.
"""
metadata = super().get_analytics_metadata()
metadata.update({ct: c[0].id for ct, c in self.components.items()})
return metadata
@property
def is_valid(self) -> bool:
"""Check if the stack is valid.
Returns:
True if the stack is valid, False otherwise.
"""
return (
StackComponentType.ARTIFACT_STORE in self.components
and StackComponentType.ORCHESTRATOR in self.components
)
def to_yaml(self) -> Dict[str, Any]:
"""Create yaml representation of the Stack Model.
Returns:
The yaml representation of the Stack Model.
"""
component_data = {}
for component_type, components_list in self.components.items():
component = components_list[0]
component_dict = json.loads(
component.json(
include={"name", "type", "flavor", "configuration"}
)
)
component_data[component_type.value] = component_dict
# write zenml version and stack dict to YAML
yaml_data = {
"stack_name": self.name,
"components": component_data,
}
return yaml_data
is_valid: bool
property
readonly
Check if the stack is valid.
Returns:
Type | Description |
---|---|
bool |
True if the stack is valid, False otherwise. |
get_analytics_metadata(self)
Add the stack components to the stack analytics metadata.
Returns:
Type | Description |
---|---|
Dict[str, Any] |
Dict of analytics metadata. |
Source code in zenml/models/stack_models.py
def get_analytics_metadata(self) -> Dict[str, Any]:
"""Add the stack components to the stack analytics metadata.
Returns:
Dict of analytics metadata.
"""
metadata = super().get_analytics_metadata()
metadata.update({ct: c[0].id for ct, c in self.components.items()})
return metadata
to_yaml(self)
Create yaml representation of the Stack Model.
Returns:
Type | Description |
---|---|
Dict[str, Any] |
The yaml representation of the Stack Model. |
Source code in zenml/models/stack_models.py
def to_yaml(self) -> Dict[str, Any]:
"""Create yaml representation of the Stack Model.
Returns:
The yaml representation of the Stack Model.
"""
component_data = {}
for component_type, components_list in self.components.items():
component = components_list[0]
component_dict = json.loads(
component.json(
include={"name", "type", "flavor", "configuration"}
)
)
component_data[component_type.value] = component_dict
# write zenml version and stack dict to YAML
yaml_data = {
"stack_name": self.name,
"components": component_data,
}
return yaml_data
StackUpdateModel (StackRequestModel)
pydantic-model
The update model for stacks.
Source code in zenml/models/stack_models.py
class StackUpdateModel(StackRequestModel):
"""The update model for stacks."""
step_run_models
Models representing steps of pipeline runs.
StepRunBaseModel (BaseModel)
pydantic-model
Base model for step runs.
Source code in zenml/models/step_run_models.py
class StepRunBaseModel(BaseModel):
"""Base model for step runs."""
name: str = Field(
title="The name of the pipeline run step.",
max_length=MODEL_NAME_FIELD_MAX_LENGTH,
)
step: Step
pipeline_run_id: UUID
original_step_run_id: Optional[UUID] = None
status: ExecutionStatus
parent_step_ids: List[UUID] = []
cache_key: Optional[str] = None
start_time: Optional[datetime] = None
end_time: Optional[datetime] = None
StepRunRequestModel (StepRunBaseModel, ProjectScopedRequestModel)
pydantic-model
Request model for step runs.
Source code in zenml/models/step_run_models.py
class StepRunRequestModel(StepRunBaseModel, ProjectScopedRequestModel):
"""Request model for step runs."""
input_artifacts: Dict[str, UUID] = {}
output_artifacts: Dict[str, UUID] = {}
StepRunResponseModel (StepRunBaseModel, ProjectScopedResponseModel)
pydantic-model
Response model for step runs.
Source code in zenml/models/step_run_models.py
class StepRunResponseModel(StepRunBaseModel, ProjectScopedResponseModel):
"""Response model for step runs."""
input_artifacts: Dict[str, "ArtifactResponseModel"] = {}
output_artifacts: Dict[str, "ArtifactResponseModel"] = {}
StepRunUpdateModel (BaseModel)
pydantic-model
Update model for step runs.
Source code in zenml/models/step_run_models.py
class StepRunUpdateModel(BaseModel):
"""Update model for step runs."""
output_artifacts: Dict[str, UUID] = {}
status: Optional[ExecutionStatus] = None
end_time: Optional[datetime] = None
team_models
Models representing teams.
TeamBaseModel (BaseModel)
pydantic-model
Base model for teams.
Source code in zenml/models/team_models.py
class TeamBaseModel(BaseModel):
"""Base model for teams."""
name: str = Field(
title="The unique name of the team.",
max_length=MODEL_NAME_FIELD_MAX_LENGTH,
)
TeamRequestModel (TeamBaseModel, BaseRequestModel)
pydantic-model
Request model for teams.
Source code in zenml/models/team_models.py
class TeamRequestModel(TeamBaseModel, BaseRequestModel):
"""Request model for teams."""
users: Optional[List[UUID]] = Field(
title="The list of users within this team."
)
TeamResponseModel (TeamBaseModel, BaseResponseModel)
pydantic-model
Response model for teams.
Source code in zenml/models/team_models.py
class TeamResponseModel(TeamBaseModel, BaseResponseModel):
"""Response model for teams."""
ANALYTICS_FIELDS: ClassVar[List[str]] = ["id"]
users: List["UserResponseModel"] = Field(
title="The list of users within this team."
)
@property
def user_ids(self) -> List[UUID]:
"""Returns a list of user IDs that are part of this team.
Returns:
A list of user IDs.
"""
if self.users:
return [u.id for u in self.users]
else:
return []
@property
def user_names(self) -> List[str]:
"""Returns a list names of users that are part of this team.
Returns:
A list of names of users.
"""
if self.users:
return [u.name for u in self.users]
else:
return []
user_ids: List[uuid.UUID]
property
readonly
Returns a list of user IDs that are part of this team.
Returns:
Type | Description |
---|---|
List[uuid.UUID] |
A list of user IDs. |
user_names: List[str]
property
readonly
Returns a list names of users that are part of this team.
Returns:
Type | Description |
---|---|
List[str] |
A list of names of users. |
TeamUpdateModel (TeamRequestModel)
pydantic-model
Update model for teams.
Source code in zenml/models/team_models.py
class TeamUpdateModel(TeamRequestModel):
"""Update model for teams."""
user_models
Models representing users.
JWTToken (BaseModel)
pydantic-model
Pydantic object representing a JWT token.
Attributes:
Name | Type | Description |
---|---|---|
token_type |
JWTTokenType |
The type of token. |
user_id |
UUID |
The id of the authenticated User |
permissions |
List[str] |
The permissions scope of the authenticated user |
Source code in zenml/models/user_models.py
class JWTToken(BaseModel):
"""Pydantic object representing a JWT token.
Attributes:
token_type: The type of token.
user_id: The id of the authenticated User
permissions: The permissions scope of the authenticated user
"""
JWT_ALGORITHM: ClassVar[str] = "HS256"
token_type: JWTTokenType
user_id: UUID
permissions: List[str]
@classmethod
def decode(cls, token_type: JWTTokenType, token: str) -> "JWTToken":
"""Decodes a JWT access token.
Decodes a JWT access token and returns a `JWTToken` object with the
information retrieved from its subject claim.
Args:
token_type: The type of token.
token: The encoded JWT token.
Returns:
The decoded JWT access token.
Raises:
AuthorizationException: If the token is invalid.
"""
# import here to keep these dependencies out of the client
from jose import JWTError, jwt # type: ignore[import]
try:
payload = jwt.decode(
token,
GlobalConfiguration().jwt_secret_key,
algorithms=[cls.JWT_ALGORITHM],
)
except JWTError as e:
raise AuthorizationException(f"Invalid JWT token: {e}") from e
subject: str = payload.get("sub")
if subject is None:
raise AuthorizationException(
"Invalid JWT token: the subject claim is missing"
)
permissions: List[str] = payload.get("permissions")
if permissions is None:
raise AuthorizationException(
"Invalid JWT token: the permissions scope is missing"
)
try:
return cls(
token_type=token_type,
user_id=UUID(subject),
permissions=set(permissions),
)
except ValueError as e:
raise AuthorizationException(
f"Invalid JWT token: could not decode subject claim: {e}"
) from e
def encode(self, expire_minutes: Optional[int] = None) -> str:
"""Creates a JWT access token.
Generates and returns a JWT access token with the subject claim set to
contain the information in this Pydantic object.
Args:
expire_minutes: Number of minutes the token should be valid. If not
provided, the token will not be set to expire.
Returns:
The generated access token.
"""
# import here to keep these dependencies out of the client
from jose import jwt
claims: Dict[str, Any] = {
"sub": str(self.user_id),
"permissions": list(self.permissions),
}
if expire_minutes:
expire = datetime.utcnow() + timedelta(minutes=expire_minutes)
claims["exp"] = expire
token: str = jwt.encode(
claims,
GlobalConfiguration().jwt_secret_key,
algorithm=self.JWT_ALGORITHM,
)
return token
decode(token_type, token)
classmethod
Decodes a JWT access token.
Decodes a JWT access token and returns a JWTToken
object with the
information retrieved from its subject claim.
Parameters:
Name | Type | Description | Default |
---|---|---|---|
token_type |
JWTTokenType |
The type of token. |
required |
token |
str |
The encoded JWT token. |
required |
Returns:
Type | Description |
---|---|
JWTToken |
The decoded JWT access token. |
Exceptions:
Type | Description |
---|---|
AuthorizationException |
If the token is invalid. |
Source code in zenml/models/user_models.py
@classmethod
def decode(cls, token_type: JWTTokenType, token: str) -> "JWTToken":
"""Decodes a JWT access token.
Decodes a JWT access token and returns a `JWTToken` object with the
information retrieved from its subject claim.
Args:
token_type: The type of token.
token: The encoded JWT token.
Returns:
The decoded JWT access token.
Raises:
AuthorizationException: If the token is invalid.
"""
# import here to keep these dependencies out of the client
from jose import JWTError, jwt # type: ignore[import]
try:
payload = jwt.decode(
token,
GlobalConfiguration().jwt_secret_key,
algorithms=[cls.JWT_ALGORITHM],
)
except JWTError as e:
raise AuthorizationException(f"Invalid JWT token: {e}") from e
subject: str = payload.get("sub")
if subject is None:
raise AuthorizationException(
"Invalid JWT token: the subject claim is missing"
)
permissions: List[str] = payload.get("permissions")
if permissions is None:
raise AuthorizationException(
"Invalid JWT token: the permissions scope is missing"
)
try:
return cls(
token_type=token_type,
user_id=UUID(subject),
permissions=set(permissions),
)
except ValueError as e:
raise AuthorizationException(
f"Invalid JWT token: could not decode subject claim: {e}"
) from e
encode(self, expire_minutes=None)
Creates a JWT access token.
Generates and returns a JWT access token with the subject claim set to contain the information in this Pydantic object.
Parameters:
Name | Type | Description | Default |
---|---|---|---|
expire_minutes |
Optional[int] |
Number of minutes the token should be valid. If not provided, the token will not be set to expire. |
None |
Returns:
Type | Description |
---|---|
str |
The generated access token. |
Source code in zenml/models/user_models.py
def encode(self, expire_minutes: Optional[int] = None) -> str:
"""Creates a JWT access token.
Generates and returns a JWT access token with the subject claim set to
contain the information in this Pydantic object.
Args:
expire_minutes: Number of minutes the token should be valid. If not
provided, the token will not be set to expire.
Returns:
The generated access token.
"""
# import here to keep these dependencies out of the client
from jose import jwt
claims: Dict[str, Any] = {
"sub": str(self.user_id),
"permissions": list(self.permissions),
}
if expire_minutes:
expire = datetime.utcnow() + timedelta(minutes=expire_minutes)
claims["exp"] = expire
token: str = jwt.encode(
claims,
GlobalConfiguration().jwt_secret_key,
algorithm=self.JWT_ALGORITHM,
)
return token
JWTTokenType (StrEnum)
The type of JWT token.
Source code in zenml/models/user_models.py
class JWTTokenType(StrEnum):
"""The type of JWT token."""
ACCESS_TOKEN = "access_token"
UserAuthModel (UserBaseModel, BaseResponseModel)
pydantic-model
Authentication Model for the User.
This model is only used server-side. The server endpoints can use this model to authenticate the user credentials (Token, Password).
Source code in zenml/models/user_models.py
class UserAuthModel(UserBaseModel, BaseResponseModel):
"""Authentication Model for the User.
This model is only used server-side. The server endpoints can use this model
to authenticate the user credentials (Token, Password).
"""
active: bool = Field(default=False, title="Active account.")
activation_token: Optional[SecretStr] = Field(default=None, exclude=True)
password: Optional[SecretStr] = Field(default=None, exclude=True)
teams: Optional[List["TeamResponseModel"]] = Field(
title="The list of teams for this user."
)
def generate_access_token(self, permissions: List[str]) -> str:
"""Generates an access token.
Generates an access token and returns it.
Args:
permissions: Permissions to add to the token
Returns:
The generated access token.
"""
return JWTToken(
token_type=JWTTokenType.ACCESS_TOKEN,
user_id=self.id,
permissions=permissions,
).encode()
@classmethod
def _is_hashed_secret(cls, secret: SecretStr) -> bool:
"""Checks if a secret value is already hashed.
Args:
secret: The secret value to check.
Returns:
True if the secret value is hashed, otherwise False.
"""
return (
re.match(r"^\$2[ayb]\$.{56}$", secret.get_secret_value())
is not None
)
@classmethod
def _get_hashed_secret(cls, secret: Optional[SecretStr]) -> Optional[str]:
"""Hashes the input secret and returns the hash value.
Only applied if supplied and if not already hashed.
Args:
secret: The secret value to hash.
Returns:
The secret hash value, or None if no secret was supplied.
"""
if secret is None:
return None
if cls._is_hashed_secret(secret):
return secret.get_secret_value()
pwd_context = cls._get_crypt_context()
return cast(str, pwd_context.hash(secret.get_secret_value()))
def get_password(self) -> Optional[str]:
"""Get the password.
Returns:
The password as a plain string, if it exists.
"""
if self.password is None:
return None
return self.password.get_secret_value()
def get_hashed_password(self) -> Optional[str]:
"""Returns the hashed password, if configured.
Returns:
The hashed password.
"""
return self._get_hashed_secret(self.password)
def get_hashed_activation_token(self) -> Optional[str]:
"""Returns the hashed activation token, if configured.
Returns:
The hashed activation token.
"""
return self._get_hashed_secret(self.activation_token)
@classmethod
def verify_password(
cls, plain_password: str, user: Optional["UserAuthModel"] = None
) -> bool:
"""Verifies a given plain password against the stored password.
Args:
plain_password: Input password to be verified.
user: User for which the password is to be verified.
Returns:
True if the passwords match.
"""
# even when the user or password is not set, we still want to execute
# the password hash verification to protect against response discrepancy
# attacks (https://cwe.mitre.org/data/definitions/204.html)
password_hash: Optional[str] = None
if user is not None and user.password is not None: # and user.active:
password_hash = user.get_hashed_password()
pwd_context = cls._get_crypt_context()
return cast(bool, pwd_context.verify(plain_password, password_hash))
@classmethod
def verify_access_token(cls, token: str) -> Optional["UserAuthModel"]:
"""Verifies an access token.
Verifies an access token and returns the user that was used to generate
it if the token is valid and None otherwise.
Args:
token: The access token to verify.
Returns:
The user that generated the token if valid, None otherwise.
"""
try:
access_token = JWTToken.decode(
token_type=JWTTokenType.ACCESS_TOKEN, token=token
)
except AuthorizationException:
return None
zen_store = GlobalConfiguration().zen_store
try:
user = zen_store.get_auth_user(user_name_or_id=access_token.user_id)
except KeyError:
return None
else:
if user.active:
return user
return None
@classmethod
def verify_activation_token(
cls, activation_token: str, user: Optional["UserAuthModel"] = None
) -> bool:
"""Verifies a given activation token against the stored token.
Args:
activation_token: Input activation token to be verified.
user: User for which the activation token is to be verified.
Returns:
True if the token is valid.
"""
# even when the user or token is not set, we still want to execute the
# token hash verification to protect against response discrepancy
# attacks (https://cwe.mitre.org/data/definitions/204.html)
token_hash: Optional[str] = None
if (
user is not None
and user.activation_token is not None
and not user.active
):
token_hash = user.get_hashed_activation_token()
pwd_context = cls._get_crypt_context()
return cast(bool, pwd_context.verify(activation_token, token_hash))
generate_access_token(self, permissions)
Generates an access token.
Generates an access token and returns it.
Parameters:
Name | Type | Description | Default |
---|---|---|---|
permissions |
List[str] |
Permissions to add to the token |
required |
Returns:
Type | Description |
---|---|
str |
The generated access token. |
Source code in zenml/models/user_models.py
def generate_access_token(self, permissions: List[str]) -> str:
"""Generates an access token.
Generates an access token and returns it.
Args:
permissions: Permissions to add to the token
Returns:
The generated access token.
"""
return JWTToken(
token_type=JWTTokenType.ACCESS_TOKEN,
user_id=self.id,
permissions=permissions,
).encode()
get_hashed_activation_token(self)
Returns the hashed activation token, if configured.
Returns:
Type | Description |
---|---|
Optional[str] |
The hashed activation token. |
Source code in zenml/models/user_models.py
def get_hashed_activation_token(self) -> Optional[str]:
"""Returns the hashed activation token, if configured.
Returns:
The hashed activation token.
"""
return self._get_hashed_secret(self.activation_token)
get_hashed_password(self)
Returns the hashed password, if configured.
Returns:
Type | Description |
---|---|
Optional[str] |
The hashed password. |
Source code in zenml/models/user_models.py
def get_hashed_password(self) -> Optional[str]:
"""Returns the hashed password, if configured.
Returns:
The hashed password.
"""
return self._get_hashed_secret(self.password)
get_password(self)
Get the password.
Returns:
Type | Description |
---|---|
Optional[str] |
The password as a plain string, if it exists. |
Source code in zenml/models/user_models.py
def get_password(self) -> Optional[str]:
"""Get the password.
Returns:
The password as a plain string, if it exists.
"""
if self.password is None:
return None
return self.password.get_secret_value()
verify_access_token(token)
classmethod
Verifies an access token.
Verifies an access token and returns the user that was used to generate it if the token is valid and None otherwise.
Parameters:
Name | Type | Description | Default |
---|---|---|---|
token |
str |
The access token to verify. |
required |
Returns:
Type | Description |
---|---|
Optional[UserAuthModel] |
The user that generated the token if valid, None otherwise. |
Source code in zenml/models/user_models.py
@classmethod
def verify_access_token(cls, token: str) -> Optional["UserAuthModel"]:
"""Verifies an access token.
Verifies an access token and returns the user that was used to generate
it if the token is valid and None otherwise.
Args:
token: The access token to verify.
Returns:
The user that generated the token if valid, None otherwise.
"""
try:
access_token = JWTToken.decode(
token_type=JWTTokenType.ACCESS_TOKEN, token=token
)
except AuthorizationException:
return None
zen_store = GlobalConfiguration().zen_store
try:
user = zen_store.get_auth_user(user_name_or_id=access_token.user_id)
except KeyError:
return None
else:
if user.active:
return user
return None
verify_activation_token(activation_token, user=None)
classmethod
Verifies a given activation token against the stored token.
Parameters:
Name | Type | Description | Default |
---|---|---|---|
activation_token |
str |
Input activation token to be verified. |
required |
user |
Optional[UserAuthModel] |
User for which the activation token is to be verified. |
None |
Returns:
Type | Description |
---|---|
bool |
True if the token is valid. |
Source code in zenml/models/user_models.py
@classmethod
def verify_activation_token(
cls, activation_token: str, user: Optional["UserAuthModel"] = None
) -> bool:
"""Verifies a given activation token against the stored token.
Args:
activation_token: Input activation token to be verified.
user: User for which the activation token is to be verified.
Returns:
True if the token is valid.
"""
# even when the user or token is not set, we still want to execute the
# token hash verification to protect against response discrepancy
# attacks (https://cwe.mitre.org/data/definitions/204.html)
token_hash: Optional[str] = None
if (
user is not None
and user.activation_token is not None
and not user.active
):
token_hash = user.get_hashed_activation_token()
pwd_context = cls._get_crypt_context()
return cast(bool, pwd_context.verify(activation_token, token_hash))
verify_password(plain_password, user=None)
classmethod
Verifies a given plain password against the stored password.
Parameters:
Name | Type | Description | Default |
---|---|---|---|
plain_password |
str |
Input password to be verified. |
required |
user |
Optional[UserAuthModel] |
User for which the password is to be verified. |
None |
Returns:
Type | Description |
---|---|
bool |
True if the passwords match. |
Source code in zenml/models/user_models.py
@classmethod
def verify_password(
cls, plain_password: str, user: Optional["UserAuthModel"] = None
) -> bool:
"""Verifies a given plain password against the stored password.
Args:
plain_password: Input password to be verified.
user: User for which the password is to be verified.
Returns:
True if the passwords match.
"""
# even when the user or password is not set, we still want to execute
# the password hash verification to protect against response discrepancy
# attacks (https://cwe.mitre.org/data/definitions/204.html)
password_hash: Optional[str] = None
if user is not None and user.password is not None: # and user.active:
password_hash = user.get_hashed_password()
pwd_context = cls._get_crypt_context()
return cast(bool, pwd_context.verify(plain_password, password_hash))
UserBaseModel (BaseModel)
pydantic-model
Base model for users.
Source code in zenml/models/user_models.py
class UserBaseModel(BaseModel):
"""Base model for users."""
name: str = Field(
title="The unique username for the account.",
max_length=MODEL_NAME_FIELD_MAX_LENGTH,
)
full_name: str = Field(
default="",
title="The full name for the account owner.",
max_length=MODEL_NAME_FIELD_MAX_LENGTH,
)
email_opted_in: Optional[bool] = Field(
title="Whether the user agreed to share their email.",
description="`null` if not answered, `true` if agreed, "
"`false` if skipped.",
)
active: bool = Field(default=False, title="Active account.")
@classmethod
def _get_crypt_context(cls) -> "CryptContext":
"""Returns the password encryption context.
Returns:
The password encryption context.
"""
from passlib.context import CryptContext
return CryptContext(schemes=["bcrypt"], deprecated="auto")
email_opted_in: bool
pydantic-field
null
if not answered, true
if agreed, false
if skipped.
UserRequestModel (UserBaseModel, BaseRequestModel)
pydantic-model
Request model for users.
This model is used to create a user. The email field is optional but is more commonly set on the UpdateRequestModel which inherits from this model. Users can also optionally set their password during creation.
Source code in zenml/models/user_models.py
class UserRequestModel(UserBaseModel, BaseRequestModel):
"""Request model for users.
This model is used to create a user. The email field is optional but is
more commonly set on the UpdateRequestModel which inherits from this model.
Users can also optionally set their password during creation.
"""
ANALYTICS_FIELDS: ClassVar[List[str]] = [
"name",
"full_name",
"active",
"email_opted_in",
]
email: Optional[str] = Field(
default=None,
title="The email address associated with the account.",
max_length=MODEL_NAME_FIELD_MAX_LENGTH,
)
password: Optional[str] = Field(default=None)
activation_token: Optional[str] = Field(default=None)
class Config:
"""Pydantic configuration class."""
# Validate attributes when assigning them
validate_assignment = True
# Forbid extra attributes to prevent unexpected behavior
extra = "forbid"
underscore_attrs_are_private = True
@classmethod
def _create_hashed_secret(cls, secret: Optional[str]) -> Optional[str]:
"""Hashes the input secret and returns the hash value.
Only applied if supplied and if not already hashed.
Args:
secret: The secret value to hash.
Returns:
The secret hash value, or None if no secret was supplied.
"""
if secret is None:
return None
pwd_context = cls._get_crypt_context()
return cast(str, pwd_context.hash(secret))
def create_hashed_password(self) -> Optional[str]:
"""Hashes the password.
Returns:
The hashed password.
"""
return self._create_hashed_secret(self.password)
def create_hashed_activation_token(self) -> Optional[str]:
"""Hashes the activation token.
Returns:
The hashed activation token.
"""
return self._create_hashed_secret(self.activation_token)
def generate_activation_token(self) -> str:
"""Generates and stores a new activation token.
Returns:
The generated activation token.
"""
self.activation_token = token_hex(32)
return self.activation_token
Config
Pydantic configuration class.
Source code in zenml/models/user_models.py
class Config:
"""Pydantic configuration class."""
# Validate attributes when assigning them
validate_assignment = True
# Forbid extra attributes to prevent unexpected behavior
extra = "forbid"
underscore_attrs_are_private = True
create_hashed_activation_token(self)
Hashes the activation token.
Returns:
Type | Description |
---|---|
Optional[str] |
The hashed activation token. |
Source code in zenml/models/user_models.py
def create_hashed_activation_token(self) -> Optional[str]:
"""Hashes the activation token.
Returns:
The hashed activation token.
"""
return self._create_hashed_secret(self.activation_token)
create_hashed_password(self)
Hashes the password.
Returns:
Type | Description |
---|---|
Optional[str] |
The hashed password. |
Source code in zenml/models/user_models.py
def create_hashed_password(self) -> Optional[str]:
"""Hashes the password.
Returns:
The hashed password.
"""
return self._create_hashed_secret(self.password)
generate_activation_token(self)
Generates and stores a new activation token.
Returns:
Type | Description |
---|---|
str |
The generated activation token. |
Source code in zenml/models/user_models.py
def generate_activation_token(self) -> str:
"""Generates and stores a new activation token.
Returns:
The generated activation token.
"""
self.activation_token = token_hex(32)
return self.activation_token
UserResponseModel (UserBaseModel, BaseResponseModel)
pydantic-model
Response model for users.
This returns the activation_token (which is required for the user-invitation-flow of the frontend. This also optionally includes the team the user is a part of. The email is returned optionally as well for use by the analytics on the client-side.
Source code in zenml/models/user_models.py
class UserResponseModel(UserBaseModel, BaseResponseModel):
"""Response model for users.
This returns the activation_token (which is required for the
user-invitation-flow of the frontend. This also optionally includes the
team the user is a part of. The email is returned optionally as well
for use by the analytics on the client-side.
"""
ANALYTICS_FIELDS: ClassVar[List[str]] = [
"id",
"name",
"full_name",
"active",
"email_opted_in",
]
activation_token: Optional[str] = Field(default=None)
teams: Optional[List["TeamResponseModel"]] = Field(
title="The list of teams for this user."
)
email: Optional[str] = Field(
default="",
title="The email address associated with the account.",
max_length=MODEL_NAME_FIELD_MAX_LENGTH,
)
def generate_access_token(self, permissions: List[str]) -> str:
"""Generates an access token.
Generates an access token and returns it.
Args:
permissions: Permissions to add to the token
Returns:
The generated access token.
"""
return JWTToken(
token_type=JWTTokenType.ACCESS_TOKEN,
user_id=self.id,
permissions=permissions,
).encode()
generate_access_token(self, permissions)
Generates an access token.
Generates an access token and returns it.
Parameters:
Name | Type | Description | Default |
---|---|---|---|
permissions |
List[str] |
Permissions to add to the token |
required |
Returns:
Type | Description |
---|---|
str |
The generated access token. |
Source code in zenml/models/user_models.py
def generate_access_token(self, permissions: List[str]) -> str:
"""Generates an access token.
Generates an access token and returns it.
Args:
permissions: Permissions to add to the token
Returns:
The generated access token.
"""
return JWTToken(
token_type=JWTTokenType.ACCESS_TOKEN,
user_id=self.id,
permissions=permissions,
).encode()
UserUpdateModel (UserRequestModel)
pydantic-model
Update model for users.
Source code in zenml/models/user_models.py
class UserUpdateModel(UserRequestModel):
"""Update model for users."""
@root_validator
def user_email_updates(cls, values: Dict[str, Any]) -> Dict[str, Any]:
"""Validate that the UserUpdateModel conforms to the email-opt-in-flow.
Args:
values: The values to validate.
Returns:
The validated values.
Raises:
ValueError: If the email was not provided when the email_opted_in
field was set to True.
"""
# When someone sets the email, or updates the email and hasn't
# before explicitly opted out, they are opted in
if values["email"] is not None:
if values["email_opted_in"] is None:
values["email_opted_in"] = True
# It should not be possible to do opt in without an email
if values["email_opted_in"] is True:
if values["email"] is None:
raise ValueError(
"Please provide an email, when you are opting-in with "
"your email."
)
return values
user_email_updates(values)
classmethod
Validate that the UserUpdateModel conforms to the email-opt-in-flow.
Parameters:
Name | Type | Description | Default |
---|---|---|---|
values |
Dict[str, Any] |
The values to validate. |
required |
Returns:
Type | Description |
---|---|
Dict[str, Any] |
The validated values. |
Exceptions:
Type | Description |
---|---|
ValueError |
If the email was not provided when the email_opted_in field was set to True. |
Source code in zenml/models/user_models.py
@root_validator
def user_email_updates(cls, values: Dict[str, Any]) -> Dict[str, Any]:
"""Validate that the UserUpdateModel conforms to the email-opt-in-flow.
Args:
values: The values to validate.
Returns:
The validated values.
Raises:
ValueError: If the email was not provided when the email_opted_in
field was set to True.
"""
# When someone sets the email, or updates the email and hasn't
# before explicitly opted out, they are opted in
if values["email"] is not None:
if values["email_opted_in"] is None:
values["email_opted_in"] = True
# It should not be possible to do opt in without an email
if values["email_opted_in"] is True:
if values["email"] is None:
raise ValueError(
"Please provide an email, when you are opting-in with "
"your email."
)
return values