zenml-nightly 0.84.0.dev20250728__py3-none-any.whl → 0.84.0.dev20250730__py3-none-any.whl

zenml/models/v2/core/service_account.py

@@ -14,6 +14,7 @@
 """Models representing service accounts."""
 
 from typing import TYPE_CHECKING, ClassVar, List, Optional, Type, Union
+from uuid import UUID
 
 from pydantic import ConfigDict, Field
 
@@ -43,8 +44,13 @@ class ServiceAccountRequest(BaseRequest):
     ]
 
     name: str = Field(
-        title="The unique name for the service account.",
+        title="The unique username for the service account.",
+        max_length=STR_FIELD_MAX_LENGTH,
+    )
+    full_name: str = Field(
+        default="",
         max_length=STR_FIELD_MAX_LENGTH,
+        title="The display name of the service account.",
     )
     description: Optional[str] = Field(
         default=None,
@@ -52,9 +58,23 @@ class ServiceAccountRequest(BaseRequest):
         max_length=TEXT_FIELD_MAX_LENGTH,
     )
     active: bool = Field(title="Whether the service account is active or not.")
+    avatar_url: Optional[str] = Field(
+        default=None,
+        title="The avatar URL for the account.",
+    )
+
     model_config = ConfigDict(validate_assignment=True, extra="ignore")
 
 
+class ServiceAccountInternalRequest(ServiceAccountRequest):
+    """Internal request model for service accounts."""
+
+    external_user_id: Optional[UUID] = Field(
+        default=None,
+        title="The external user ID associated with the account.",
+    )
+
+
 # ------------------ Update Model ------------------
 
 
@@ -68,6 +88,11 @@ class ServiceAccountUpdate(BaseUpdate):
         max_length=STR_FIELD_MAX_LENGTH,
         default=None,
     )
+    full_name: Optional[str] = Field(
+        title="The display name of the service account.",
+        max_length=STR_FIELD_MAX_LENGTH,
+        default=None,
+    )
     description: Optional[str] = Field(
         title="A description of the service account.",
         max_length=TEXT_FIELD_MAX_LENGTH,
@@ -78,16 +103,38 @@ class ServiceAccountUpdate(BaseUpdate):
         default=None,
     )
 
+    avatar_url: Optional[str] = Field(
+        default=None,
+        title="The avatar URL for the account.",
+    )
+
     model_config = ConfigDict(validate_assignment=True)
 
 
+class ServiceAccountInternalUpdate(ServiceAccountUpdate):
+    """Internal update model for service accounts."""
+
+    external_user_id: Optional[UUID] = Field(
+        default=None,
+        title="The external user ID associated with the account.",
+    )
+
+
 # ------------------ Response Model ------------------
 
 
 class ServiceAccountResponseBody(BaseDatedResponseBody):
     """Response body for service accounts."""
 
+    full_name: str = Field(
+        default="",
+        title="The display name of the service account.",
+    )
     active: bool = Field(default=False, title="Whether the account is active.")
+    avatar_url: Optional[str] = Field(
+        default=None,
+        title="The avatar URL for the account.",
+    )
 
 
 class ServiceAccountResponseMetadata(BaseResponseMetadata):
@@ -99,6 +146,11 @@ class ServiceAccountResponseMetadata(BaseResponseMetadata):
         max_length=TEXT_FIELD_MAX_LENGTH,
     )
 
+    external_user_id: Optional[UUID] = Field(
+        default=None,
+        title="The external user ID associated with the account.",
+    )
+
 
 class ServiceAccountResponseResources(BaseResponseResources):
     """Class for all resource models associated with the service account entity."""
@@ -160,6 +212,8 @@ class ServiceAccountResponse(
                 created=self.created,
                 updated=self.updated,
                 is_admin=False,
+                avatar_url=self.avatar_url,
+                full_name=self.full_name,
             ),
             metadata=UserResponseMetadata(
                 description=self.description,
@@ -167,6 +221,15 @@ class ServiceAccountResponse(
         )
 
     # Body and metadata properties
+    @property
+    def full_name(self) -> str:
+        """The `full_name` property.
+
+        Returns:
+            the value of the property.
+        """
+        return self.get_body().full_name
+
     @property
     def active(self) -> bool:
         """The `active` property.
@@ -185,6 +248,24 @@ class ServiceAccountResponse(
         """
         return self.get_metadata().description
 
+    @property
+    def external_user_id(self) -> Optional[UUID]:
+        """The `external_user_id` property.
+
+        Returns:
+            the value of the property.
+        """
+        return self.get_metadata().external_user_id
+
+    @property
+    def avatar_url(self) -> Optional[str]:
+        """The `avatar_url` property.
+
+        Returns:
+            the value of the property.
+        """
+        return self.get_body().avatar_url
+
 
 # ------------------ Filter Model ------------------
 class ServiceAccountFilter(BaseFilter):
@@ -203,6 +284,11 @@ class ServiceAccountFilter(BaseFilter):
         description="Whether the user is active",
         union_mode="left_to_right",
     )
+    external_user_id: Optional[Union[UUID, str]] = Field(
+        default=None,
+        title="The external user ID associated with the account.",
+        union_mode="left_to_right",
+    )
 
     def apply_filter(
         self,

zenml/models/v2/core/step_run.py

@@ -48,6 +48,7 @@ from zenml.models.v2.base.scoped import (
 )
 from zenml.models.v2.core.artifact_version import ArtifactVersionResponse
 from zenml.models.v2.core.model_version import ModelVersionResponse
+from zenml.models.v2.misc.exception_info import ExceptionInfo
 
 if TYPE_CHECKING:
     from sqlalchemy.sql.elements import ColumnElement
@@ -143,6 +144,10 @@ class StepRunRequest(ProjectScopedRequest):
         title="Logs associated with this step run.",
         default=None,
     )
+    exception_info: Optional[ExceptionInfo] = Field(
+        default=None,
+        title="The exception information of the step run.",
+    )
 
     model_config = ConfigDict(protected_namespaces=())
 
@@ -169,6 +174,10 @@ class StepRunUpdate(BaseUpdate):
         title="The end time of the step run.",
         default=None,
     )
+    exception_info: Optional[ExceptionInfo] = Field(
+        default=None,
+        title="The exception information of the step run.",
+    )
     model_config = ConfigDict(protected_namespaces=())
 
 
@@ -237,6 +246,10 @@ class StepRunResponseMetadata(ProjectScopedResponseMetadata):
         default=None,
         max_length=TEXT_FIELD_MAX_LENGTH,
     )
+    exception_info: Optional[ExceptionInfo] = Field(
+        default=None,
+        title="The exception information of the step run.",
+    )
 
     # References
     logs: Optional["LogsResponse"] = Field(

zenml/models/v2/core/user.py

@@ -81,6 +81,10 @@ class UserBase(BaseModel):
         default=None,
         title="The metadata associated with the user.",
     )
+    avatar_url: Optional[str] = Field(
+        default=None,
+        title="The avatar URL for the account.",
+    )
 
     @classmethod
     def _get_crypt_context(cls) -> "CryptContext":
@@ -156,8 +160,7 @@ class UserRequest(UserBase, BaseRequest):
     )
     full_name: str = Field(
         default="",
-        title="The full name for the account owner. Only relevant for user "
-        "accounts.",
+        title="The display name for the account.",
         max_length=STR_FIELD_MAX_LENGTH,
     )
     is_admin: bool = Field(
@@ -186,8 +189,7 @@ class UserUpdate(UserBase, BaseUpdate):
     )
     full_name: Optional[str] = Field(
         default=None,
-        title="The full name for the account owner. Only relevant for user "
-        "accounts.",
+        title="The display name for the account.",
         max_length=STR_FIELD_MAX_LENGTH,
     )
     is_admin: Optional[bool] = Field(
@@ -266,8 +268,7 @@ class UserResponseBody(BaseDatedResponseBody):
     )
     full_name: str = Field(
         default="",
-        title="The full name for the account owner. Only relevant for user "
-        "accounts.",
+        title="The display name for the account.",
         max_length=STR_FIELD_MAX_LENGTH,
     )
     email_opted_in: Optional[bool] = Field(
@@ -287,6 +288,10 @@ class UserResponseBody(BaseDatedResponseBody):
         default=None,
         title="The default project ID for the user.",
     )
+    avatar_url: Optional[str] = Field(
+        default=None,
+        title="The avatar URL for the account.",
+    )
 
 
 class UserResponseMetadata(BaseResponseMetadata):
@@ -439,6 +444,15 @@ class UserResponse(
         """
         return self.get_body().default_project_id
 
+    @property
+    def avatar_url(self) -> Optional[str]:
+        """The `avatar_url` property.
+
+        Returns:
+            the value of the property.
+        """
+        return self.get_body().avatar_url
+
     # Helper methods
     @classmethod
     def _get_crypt_context(cls) -> "CryptContext":

zenml/models/v2/misc/exception_info.py

@@ -0,0 +1,30 @@
+#  Copyright (c) ZenML GmbH 2025. All Rights Reserved.
+#
+#  Licensed under the Apache License, Version 2.0 (the "License");
+#  you may not use this file except in compliance with the License.
+#  You may obtain a copy of the License at:
+#
+#       https://www.apache.org/licenses/LICENSE-2.0
+#
+#  Unless required by applicable law or agreed to in writing, software
+#  distributed under the License is distributed on an "AS IS" BASIS,
+#  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express
+#  or implied. See the License for the specific language governing
+#  permissions and limitations under the License.
+"""Exception information models."""
+
+from typing import Optional
+
+from pydantic import BaseModel, Field
+
+
+class ExceptionInfo(BaseModel):
+    """Exception information."""
+
+    traceback: str = Field(
+        title="The traceback of the exception.",
+    )
+    step_code_line: Optional[int] = Field(
+        default=None,
+        title="The line number of the step code that raised the exception.",
+    )

zenml/models/v2/misc/external_user.py

@@ -23,8 +23,11 @@ class ExternalUserModel(BaseModel):
     """External user model."""
 
     id: UUID
-    email: str
+    username: str
+    email: Optional[str] = None
     name: Optional[str] = None
     is_admin: bool = False
+    is_service_account: bool = False
+    avatar_url: Optional[str] = None
 
     model_config = ConfigDict(extra="ignore")

zenml/orchestrators/publish_utils.py

@@ -13,12 +13,14 @@
 #  permissions and limitations under the License.
 """Utilities to publish pipeline and step runs."""
 
+from contextvars import ContextVar
 from datetime import datetime
 from typing import TYPE_CHECKING, Dict, List, Optional
 
 from zenml.client import Client
 from zenml.enums import ExecutionStatus, MetadataResourceTypes
 from zenml.models import (
+    ExceptionInfo,
     PipelineRunResponse,
     PipelineRunUpdate,
     RunMetadataResource,
@@ -32,6 +34,10 @@ if TYPE_CHECKING:
 
     from zenml.metadata.metadata_types import MetadataType
 
+step_exception_info: ContextVar[Optional[ExceptionInfo]] = ContextVar(
+    "step_exception_info", default=None
+)
+
 
 def publish_successful_step_run(
     step_run_id: "UUID", output_artifact_ids: Dict[str, List["UUID"]]
@@ -59,6 +65,7 @@ def publish_step_run_status_update(
     step_run_id: "UUID",
     status: "ExecutionStatus",
     end_time: Optional[datetime] = None,
+    exception_info: Optional[ExceptionInfo] = None,
 ) -> "StepRunResponse":
     """Publishes a step run update.
 
@@ -66,6 +73,7 @@ def publish_step_run_status_update(
         step_run_id: ID of the step run.
         status: New status of the step run.
         end_time: New end time of the step run.
+        exception_info: Exception information of the step run.
 
     Returns:
         The updated step run.
@@ -83,6 +91,7 @@ def publish_step_run_status_update(
         step_run_update=StepRunUpdate(
             status=status,
             end_time=end_time,
+            exception_info=exception_info,
         ),
     )
 
@@ -102,6 +111,7 @@ def publish_failed_step_run(step_run_id: "UUID") -> "StepRunResponse":
         step_run_id=step_run_id,
         status=ExecutionStatus.FAILED,
         end_time=utc_now(),
+        exception_info=step_exception_info.get(),
     )
 
 

zenml/orchestrators/step_launcher.py

@@ -16,7 +16,6 @@
 import signal
 import time
 from contextlib import nullcontext
-from functools import partial
 from typing import TYPE_CHECKING, Any, Callable, Dict, Optional, Tuple
 
 from zenml.client import Client
@@ -24,6 +23,7 @@ from zenml.config.step_configurations import Step
 from zenml.config.step_run_info import StepRunInfo
 from zenml.constants import (
     ENV_ZENML_DISABLE_STEP_LOGS_STORAGE,
+    ENV_ZENML_STEP_OPERATOR,
     handle_bool_env_var,
 )
 from zenml.enums import ExecutionStatus
@@ -43,7 +43,7 @@ from zenml.orchestrators import output_utils, publish_utils, step_run_utils
 from zenml.orchestrators import utils as orchestrator_utils
 from zenml.orchestrators.step_runner import StepRunner
 from zenml.stack import Stack
-from zenml.utils import string_utils
+from zenml.utils import exception_utils, string_utils
 from zenml.utils.time_utils import utc_now
 
 if TYPE_CHECKING:
@@ -211,7 +211,9 @@ class StepLauncher:
 
         Raises:
             RunStoppedException: If the pipeline run is stopped by the user.
+            BaseException: If the step preparation or execution fails.
         """
+        publish_utils.step_exception_info.set(None)
         pipeline_run, run_was_created = self._create_or_reuse_run()
 
         # Enable or disable step logs storage
@@ -269,10 +271,13 @@ class StepLauncher:
 
             try:
                 request_factory.populate_request(request=step_run_request)
-            except:
+            except BaseException as e:
                 logger.exception(f"Failed preparing step `{self._step_name}`.")
                 step_run_request.status = ExecutionStatus.FAILED
                 step_run_request.end_time = utc_now()
+                step_run_request.exception_info = (
+                    exception_utils.collect_exception_information(e)
+                )
                 raise
             finally:
                 step_run = Client().zen_store.create_run_step(step_run_request)
@@ -293,9 +298,8 @@ class StepLauncher:
                         logs_context,
                         step_logging.PipelineLogsStorageContext,
                     ):
-                        force_write_logs = partial(
-                            logs_context.storage.save_to_file,
-                            force=True,
+                        force_write_logs = (
+                            logs_context.storage.send_merge_event
                         )
                     else:
 
@@ -453,6 +457,7 @@ class StepLauncher:
         environment = orchestrator_utils.get_config_environment_vars(
             pipeline_run_id=step_run_info.run_id,
         )
+        environment[ENV_ZENML_STEP_OPERATOR] = "True"
         logger.info(
             "Using step operator `%s` to run step `%s`.",
             step_operator.name,

zenml/orchestrators/step_runner.py

@@ -16,6 +16,7 @@
 
 import copy
 import inspect
+import os
 from contextlib import nullcontext
 from typing import (
     TYPE_CHECKING,
@@ -34,6 +35,7 @@ from zenml.config.step_configurations import StepConfiguration
 from zenml.config.step_run_info import StepRunInfo
 from zenml.constants import (
     ENV_ZENML_DISABLE_STEP_LOGS_STORAGE,
+    ENV_ZENML_STEP_OPERATOR,
     handle_bool_env_var,
 )
 from zenml.enums import ArtifactSaveType
@@ -41,10 +43,14 @@ from zenml.exceptions import StepInterfaceError
 from zenml.logger import get_logger
 from zenml.logging.step_logging import PipelineLogsStorageContext, redirected
 from zenml.materializers.base_materializer import BaseMaterializer
-from zenml.models.v2.core.step_run import StepRunInputResponse
+from zenml.models.v2.core.step_run import (
+    StepRunInputResponse,
+    StepRunUpdate,
+)
 from zenml.orchestrators.publish_utils import (
     publish_step_run_metadata,
     publish_successful_step_run,
+    step_exception_info,
 )
 from zenml.orchestrators.utils import (
     is_setting_enabled,
@@ -56,6 +62,7 @@ from zenml.steps.utils import (
     resolve_type_annotation,
 )
 from zenml.utils import (
+    exception_utils,
     materializer_utils,
     source_utils,
     string_utils,
@@ -193,6 +200,24 @@ class StepRunner:
                 )
             except BaseException as step_exception:  # noqa: E722
                 step_failed = True
+
+                exception_info = exception_utils.collect_exception_information(
+                    step_exception, step_instance
+                )
+
+                if ENV_ZENML_STEP_OPERATOR in os.environ:
+                    # We're running in a step operator environment, so we can't
+                    # depend on the step launcher to publish the exception info
+                    Client().zen_store.update_run_step(
+                        step_run_id=step_run_info.step_run_id,
+                        step_run_update=StepRunUpdate(
+                            exception_info=exception_info,
+                        ),
+                    )
+                else:
+                    # This will be published by the step launcher
+                    step_exception_info.set(exception_info)
+
                 if not step_run.is_retriable:
                     if (
                         failure_hook_source

zenml/utils/exception_utils.py

@@ -0,0 +1,90 @@
+#  Copyright (c) ZenML GmbH 2025. All Rights Reserved.
+#
+#  Licensed under the Apache License, Version 2.0 (the "License");
+#  you may not use this file except in compliance with the License.
+#  You may obtain a copy of the License at:
+#
+#       https://www.apache.org/licenses/LICENSE-2.0
+#
+#  Unless required by applicable law or agreed to in writing, software
+#  distributed under the License is distributed on an "AS IS" BASIS,
+#  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express
+#  or implied. See the License for the specific language governing
+#  permissions and limitations under the License.
+"""Exception utilities."""
+
+import inspect
+import os
+import re
+import traceback
+from typing import TYPE_CHECKING, Optional
+
+from zenml.constants import MEDIUMTEXT_MAX_LENGTH
+from zenml.logger import get_logger
+from zenml.models import (
+    ExceptionInfo,
+)
+
+if TYPE_CHECKING:
+    from zenml.steps import BaseStep
+
+logger = get_logger(__name__)
+
+
+def collect_exception_information(
+    exception: BaseException, step_instance: Optional["BaseStep"] = None
+) -> ExceptionInfo:
+    """Collects the exception information.
+
+    Args:
+        exception: The exception to collect information from.
+        step_instance: The step instance that is currently running.
+
+    Returns:
+        The exception information.
+    """
+    tb = traceback.format_tb(exception.__traceback__)
+    line_number = None
+    start_index = None
+
+    if step_instance and (
+        source_file := inspect.getsourcefile(step_instance.entrypoint)
+    ):
+        try:
+            source_file = os.path.abspath(source_file)
+
+            lines, start_line = inspect.getsourcelines(
+                step_instance.entrypoint
+            )
+            end_line = start_line + len(lines)
+
+            line_pattern = re.compile(f'File "{source_file}", line (\d+),')
+
+            for index, line in enumerate(tb):
+                match = line_pattern.search(line)
+                if match:
+                    potential_line_number = int(match.group(1))
+                    if (
+                        potential_line_number >= start_line
+                        and potential_line_number <= end_line
+                    ):
+                        line_number = potential_line_number - start_line
+                        start_index = index
+                        break
+        except Exception as e:
+            logger.debug("Failed to detect step code line: %s", e)
+
+    if start_index is not None:
+        # If the code failed while executing user code, we remove the initial
+        # part of the traceback that is happening in the ZenML code.
+        tb = tb[start_index:]
+
+    tb_bytes = "\n".join(tb).encode()
+    tb_bytes = tb_bytes[:MEDIUMTEXT_MAX_LENGTH]
+
+    return ExceptionInfo(
+        # Ignore errors when decoding in case we cut off in the middle of an
+        # encoded character.
+        traceback=tb_bytes.decode(errors="ignore"),
+        step_code_line=line_number,
+    )