zenml-nightly 0.80.2.dev20250415__py3-none-any.whl → 0.80.2.dev20250416__py3-none-any.whl

zenml/orchestrators/utils.py

@@ -105,7 +105,6 @@ def is_setting_enabled(
 def get_config_environment_vars(
     schedule_id: Optional[UUID] = None,
     pipeline_run_id: Optional[UUID] = None,
-    step_run_id: Optional[UUID] = None,
 ) -> Dict[str, str]:
     """Gets environment variables to set for mirroring the active config.
 
@@ -119,7 +118,6 @@ def get_config_environment_vars(
         schedule_id: Optional schedule ID to use to generate a new API token.
         pipeline_run_id: Optional pipeline run ID to use to generate a new API
             token.
-        step_run_id: Optional step run ID to use to generate a new API token.
 
     Returns:
         Environment variable dict.
@@ -138,7 +136,7 @@ def get_config_environment_vars(
         credentials_store = get_credentials_store()
         url = global_config.store_configuration.url
         api_token = credentials_store.get_token(url, allow_expired=False)
-        if schedule_id or pipeline_run_id or step_run_id:
+        if schedule_id or pipeline_run_id:
             assert isinstance(global_config.zen_store, RestZenStore)
 
             # The user has the option to manually set an expiration for the API
@@ -173,7 +171,7 @@ def get_config_environment_vars(
                 # If only a schedule is given, the pipeline run credentials will
                 # be valid for the entire duration of the schedule.
                 api_key = credentials_store.get_api_key(url)
-                if not api_key and not pipeline_run_id and not step_run_id:
+                if not api_key and not pipeline_run_id:
                     logger.warning(
                         "An API token without an expiration time will be generated "
                         "and used to run this pipeline on a schedule. This is very "
@@ -194,7 +192,6 @@ def get_config_environment_vars(
                     token_type=APITokenType.WORKLOAD,
                     schedule_id=schedule_id,
                     pipeline_run_id=pipeline_run_id,
-                    step_run_id=step_run_id,
                 )
 
             environment_vars[ENV_ZENML_STORE_PREFIX + "API_TOKEN"] = (

zenml/pipelines/pipeline_context.py

@@ -109,6 +109,7 @@ class PipelineContext:
             pipeline_configuration.enable_artifact_visualization
         )
         self.enable_step_logs = pipeline_configuration.enable_step_logs
+        self.enable_pipeline_logs = pipeline_configuration.enable_pipeline_logs
         self.settings = pipeline_configuration.settings
         self.extra = pipeline_configuration.extra
         self.model = pipeline_configuration.model

zenml/pipelines/pipeline_decorator.py

@@ -50,6 +50,7 @@ def pipeline(
     enable_cache: Optional[bool] = None,
     enable_artifact_metadata: Optional[bool] = None,
     enable_step_logs: Optional[bool] = None,
+    enable_pipeline_logs: Optional[bool] = None,
     settings: Optional[Dict[str, "SettingsOrDict"]] = None,
     tags: Optional[List[Union[str, "Tag"]]] = None,
     extra: Optional[Dict[str, Any]] = None,
@@ -67,6 +68,7 @@ def pipeline(
     enable_cache: Optional[bool] = None,
     enable_artifact_metadata: Optional[bool] = None,
     enable_step_logs: Optional[bool] = None,
+    enable_pipeline_logs: Optional[bool] = None,
     settings: Optional[Dict[str, "SettingsOrDict"]] = None,
     tags: Optional[List[Union[str, "Tag"]]] = None,
     extra: Optional[Dict[str, Any]] = None,
@@ -84,6 +86,7 @@ def pipeline(
         enable_cache: Whether to use caching or not.
         enable_artifact_metadata: Whether to enable artifact metadata or not.
         enable_step_logs: If step logs should be enabled for this pipeline.
+        enable_pipeline_logs: If pipeline logs should be enabled for this pipeline.
         settings: Settings for this pipeline.
         tags: Tags to apply to runs of the pipeline.
         extra: Extra configurations for this pipeline.
@@ -108,6 +111,7 @@ def pipeline(
             enable_cache=enable_cache,
             enable_artifact_metadata=enable_artifact_metadata,
             enable_step_logs=enable_step_logs,
+            enable_pipeline_logs=enable_pipeline_logs,
             settings=settings,
             tags=tags,
             extra=extra,

zenml/pipelines/pipeline_definition.py

@@ -16,7 +16,7 @@
 import copy
 import hashlib
 import inspect
-from contextlib import contextmanager
+from contextlib import contextmanager, nullcontext
 from pathlib import Path
 from typing import (
     TYPE_CHECKING,
@@ -56,8 +56,13 @@ from zenml.enums import StackComponentType
 from zenml.exceptions import EntityExistsError
 from zenml.hooks.hook_validators import resolve_and_validate_hook
 from zenml.logger import get_logger
+from zenml.logging.step_logging import (
+    PipelineLogsStorageContext,
+    prepare_logs_uri,
+)
 from zenml.models import (
     CodeReferenceRequest,
+    LogsRequest,
     PipelineBuildBase,
     PipelineBuildResponse,
     PipelineDeploymentBase,
@@ -130,6 +135,7 @@ class Pipeline:
         enable_artifact_metadata: Optional[bool] = None,
         enable_artifact_visualization: Optional[bool] = None,
         enable_step_logs: Optional[bool] = None,
+        enable_pipeline_logs: Optional[bool] = None,
         settings: Optional[Mapping[str, "SettingsOrDict"]] = None,
         tags: Optional[List[Union[str, "Tag"]]] = None,
         extra: Optional[Dict[str, Any]] = None,
@@ -149,6 +155,7 @@ class Pipeline:
             enable_artifact_visualization: If artifact visualization should be
                 enabled for this pipeline.
             enable_step_logs: If step logs should be enabled for this pipeline.
+            enable_pipeline_logs: If pipeline logs should be enabled for this pipeline.
             settings: Settings for this pipeline.
             tags: Tags to apply to runs of this pipeline.
             extra: Extra configurations for this pipeline.
@@ -174,6 +181,7 @@ class Pipeline:
                 enable_artifact_metadata=enable_artifact_metadata,
                 enable_artifact_visualization=enable_artifact_visualization,
                 enable_step_logs=enable_step_logs,
+                enable_pipeline_logs=enable_pipeline_logs,
                 settings=settings,
                 tags=tags,
                 extra=extra,
@@ -293,6 +301,7 @@ class Pipeline:
         enable_artifact_metadata: Optional[bool] = None,
         enable_artifact_visualization: Optional[bool] = None,
         enable_step_logs: Optional[bool] = None,
+        enable_pipeline_logs: Optional[bool] = None,
         settings: Optional[Mapping[str, "SettingsOrDict"]] = None,
         tags: Optional[List[Union[str, "Tag"]]] = None,
         extra: Optional[Dict[str, Any]] = None,
@@ -322,6 +331,7 @@ class Pipeline:
             enable_artifact_visualization: If artifact visualization should be
                 enabled for this pipeline.
             enable_step_logs: If step logs should be enabled for this pipeline.
+            enable_pipeline_logs: If pipeline logs should be enabled for this pipeline.
             settings: settings for this pipeline.
             tags: Tags to apply to runs of this pipeline.
             extra: Extra configurations for this pipeline.
@@ -364,6 +374,7 @@ class Pipeline:
                 "enable_artifact_metadata": enable_artifact_metadata,
                 "enable_artifact_visualization": enable_artifact_visualization,
                 "enable_step_logs": enable_step_logs,
+                "enable_pipeline_logs": enable_pipeline_logs,
                 "settings": settings,
                 "tags": tags,
                 "extra": extra,
@@ -588,6 +599,7 @@ To avoid this consider setting pipeline parameters only in one place (config or
         enable_artifact_metadata: Optional[bool] = None,
         enable_artifact_visualization: Optional[bool] = None,
         enable_step_logs: Optional[bool] = None,
+        enable_pipeline_logs: Optional[bool] = None,
         schedule: Optional[Schedule] = None,
         build: Union[str, "UUID", "PipelineBuildBase", None] = None,
         settings: Optional[Mapping[str, "SettingsOrDict"]] = None,
@@ -610,6 +622,8 @@ To avoid this consider setting pipeline parameters only in one place (config or
             enable_artifact_visualization: If artifact visualization should be
                 enabled for this pipeline run.
             enable_step_logs: If step logs should be enabled for this pipeline.
+            enable_pipeline_logs: If pipeline logs should be enabled for this
+                pipeline run.
             schedule: Optional schedule to use for the run.
             build: Optional build to use for the run.
             settings: Settings for this pipeline run.
@@ -641,6 +655,7 @@ To avoid this consider setting pipeline parameters only in one place (config or
             enable_artifact_metadata=enable_artifact_metadata,
             enable_artifact_visualization=enable_artifact_visualization,
             enable_step_logs=enable_step_logs,
+            enable_pipeline_logs=enable_pipeline_logs,
             steps=step_configurations,
             settings=settings,
             schedule=schedule,
@@ -723,7 +738,7 @@ To avoid this consider setting pipeline parameters only in one place (config or
         if prevent_build_reuse:
             logger.warning(
                 "Passing `prevent_build_reuse=True` to "
-                "`pipeline.with_opitions(...)` is deprecated. Use "
+                "`pipeline.with_options(...)` is deprecated. Use "
                 "`DockerSettings.prevent_build_reuse` instead."
             )
 
@@ -806,31 +821,77 @@ To avoid this consider setting pipeline parameters only in one place (config or
 
         with track_handler(AnalyticsEvent.RUN_PIPELINE) as analytics_handler:
             stack = Client().active_stack
-            deployment = self._create_deployment(**self._run_args)
 
-            self.log_pipeline_deployment_metadata(deployment)
-            run = create_placeholder_run(deployment=deployment)
+            # Enable or disable pipeline run logs storage
+            if self._run_args.get("schedule"):
+                # Pipeline runs scheduled to run in the future are not logged
+                # via the client.
+                logging_enabled = False
+            elif constants.handle_bool_env_var(
+                constants.ENV_ZENML_DISABLE_PIPELINE_LOGS_STORAGE, False
+            ):
+                logging_enabled = False
+            else:
+                logging_enabled = self._run_args.get(
+                    "enable_pipeline_logs",
+                    self.configuration.enable_pipeline_logs
+                    if self.configuration.enable_pipeline_logs is not None
+                    else True,
+                )
 
-            analytics_handler.metadata = self._get_pipeline_analytics_metadata(
-                deployment=deployment,
-                stack=stack,
-                run_id=run.id if run else None,
-            )
+            logs_context = nullcontext()
+            logs_model = None
 
-            if run:
-                run_url = dashboard_utils.get_run_url(run)
-                if run_url:
-                    logger.info(f"Dashboard URL for Pipeline Run: {run_url}")
-                else:
-                    logger.info(
-                        "You can visualize your pipeline runs in the `ZenML "
-                        "Dashboard`. In order to try it locally, please run "
-                        "`zenml login --local`."
+            if logging_enabled:
+                # Configure the logs
+                logs_uri = prepare_logs_uri(
+                    stack.artifact_store,
+                )
+
+                logs_context = PipelineLogsStorageContext(
+                    logs_uri=logs_uri,
+                    artifact_store=stack.artifact_store,
+                    prepend_step_name=False,
+                )  # type: ignore[assignment]
+
+                logs_model = LogsRequest(
+                    uri=logs_uri,
+                    artifact_store_id=stack.artifact_store.id,
+                )
+
+            with logs_context:
+                deployment = self._create_deployment(**self._run_args)
+
+                self.log_pipeline_deployment_metadata(deployment)
+                run = create_placeholder_run(
+                    deployment=deployment, logs=logs_model
+                )
+
+                analytics_handler.metadata = (
+                    self._get_pipeline_analytics_metadata(
+                        deployment=deployment,
+                        stack=stack,
+                        run_id=run.id if run else None,
                     )
+                )
+
+                if run:
+                    run_url = dashboard_utils.get_run_url(run)
+                    if run_url:
+                        logger.info(
+                            f"Dashboard URL for Pipeline Run: {run_url}"
+                        )
+                    else:
+                        logger.info(
+                            "You can visualize your pipeline runs in the `ZenML "
+                            "Dashboard`. In order to try it locally, please run "
+                            "`zenml login --local`."
+                        )
+
+                deploy_pipeline(
+                    deployment=deployment, stack=stack, placeholder_run=run
+                )
 
-            deploy_pipeline(
-                deployment=deployment, stack=stack, placeholder_run=run
-            )
             if run:
                 return Client().get_pipeline_run(run.id)
             return None

zenml/pipelines/run_utils.py

@@ -15,6 +15,7 @@ from zenml.enums import ExecutionStatus
 from zenml.logger import get_logger
 from zenml.models import (
     FlavorFilter,
+    LogsRequest,
     PipelineDeploymentBase,
     PipelineDeploymentResponse,
     PipelineRunRequest,
@@ -49,6 +50,7 @@ def get_default_run_name(pipeline_name: str) -> str:
 
 def create_placeholder_run(
     deployment: "PipelineDeploymentResponse",
+    logs: Optional["LogsRequest"] = None,
 ) -> Optional["PipelineRunResponse"]:
     """Create a placeholder run for the deployment.
 
@@ -57,6 +59,7 @@ def create_placeholder_run(
 
     Args:
         deployment: The deployment for which to create the placeholder run.
+        logs: The logs for the run.
 
     Returns:
         The placeholder run or `None` if no run was created.
@@ -86,6 +89,7 @@ def create_placeholder_run(
         pipeline=deployment.pipeline.id if deployment.pipeline else None,
         status=ExecutionStatus.INITIALIZING,
         tags=deployment.pipeline_configuration.tags,
+        logs=logs,
     )
     run, _ = Client().zen_store.get_or_create_run(run_request)
     return run

zenml/steps/utils.py

@@ -553,7 +553,7 @@ def run_as_single_step_pipeline(
     orchestrator = Client().active_stack.orchestrator
 
     pipeline_settings: Any = {}
-    if "synchronous" in orchestrator.config.model_fields:
+    if "synchronous" in type(orchestrator.config).model_fields:
         # Make sure the orchestrator runs sync so we stream the logs
         key = settings_utils.get_stack_component_setting_key(orchestrator)
         pipeline_settings[key] = BaseSettings(synchronous=True)

zenml/zen_server/auth.py

@@ -15,7 +15,7 @@
 
 from contextvars import ContextVar
 from datetime import datetime, timedelta
-from typing import Callable, Optional, Union
+from typing import Callable, Optional, Tuple, Union
 from urllib.parse import urlencode, urlparse
 from uuid import UUID, uuid4
 
@@ -33,9 +33,12 @@ from zenml.analytics.context import AnalyticsContext
 from zenml.constants import (
     API,
     DEFAULT_USERNAME,
+    DEFAULT_ZENML_SERVER_GENERIC_API_TOKEN_LIFETIME,
+    ENV_ZENML_WORKLOAD_TOKEN_EXPIRATION_LEEWAY,
     EXTERNAL_AUTHENTICATOR_TIMEOUT,
     LOGIN,
     VERSION_1,
+    handle_int_env_var,
 )
 from zenml.enums import (
     AuthScheme,
@@ -420,28 +423,32 @@ def authenticate_credentials(
             @cache_result(expiry=30)
             def get_pipeline_run_status(
                 pipeline_run_id: UUID,
-            ) -> Optional[ExecutionStatus]:
+            ) -> Tuple[Optional[ExecutionStatus], Optional[datetime]]:
                 """Get the status of a pipeline run.
 
                 Args:
                     pipeline_run_id: The pipeline run ID.
 
                 Returns:
-                    The pipeline run status or None if the pipeline run does not
-                    exist.
+                    The pipeline run status and end time or None if the pipeline
+                    run does not exist.
                 """
                 try:
                     pipeline_run = zen_store().get_run(
-                        pipeline_run_id, hydrate=False
+                        pipeline_run_id, hydrate=True
                     )
                 except KeyError:
-                    return None
+                    return None, None
 
-                return pipeline_run.status
+                return (
+                    pipeline_run.status,
+                    pipeline_run.end_time,
+                )
 
-            pipeline_run_status = get_pipeline_run_status(
-                decoded_token.pipeline_run_id
-            )
+            (
+                pipeline_run_status,
+                pipeline_run_end_time,
+            ) = get_pipeline_run_status(decoded_token.pipeline_run_id)
             if pipeline_run_status is None:
                 error = (
                     f"Authentication error: error retrieving token pipeline run "
@@ -450,59 +457,35 @@ def authenticate_credentials(
                 logger.error(error)
                 raise CredentialsNotValid(error)
 
+            leeway = handle_int_env_var(
+                ENV_ZENML_WORKLOAD_TOKEN_EXPIRATION_LEEWAY,
+                DEFAULT_ZENML_SERVER_GENERIC_API_TOKEN_LIFETIME,
+            )
             if pipeline_run_status.is_finished:
-                error = (
-                    f"The execution of pipeline run "
-                    f"{decoded_token.pipeline_run_id} has already concluded and "
-                    "API tokens scoped to it are no longer valid."
-                )
-                logger.error(error)
-                raise CredentialsNotValid(error)
-
-        if decoded_token.step_run_id:
-            # If the token contains a step run ID, we need to check if the
-            # step run exists in the database and the step run has not concluded.
-            # We use a cached version of the step run status to avoid unnecessary
-            # database queries.
-
-            @cache_result(expiry=30)
-            def get_step_run_status(
-                step_run_id: UUID,
-            ) -> Optional[ExecutionStatus]:
-                """Get the status of a step run.
-
-                Args:
-                    step_run_id: The step run ID.
-
-                Returns:
-                    The step run status or None if the step run does not exist.
-                """
-                try:
-                    step_run = zen_store().get_run_step(
-                        step_run_id, hydrate=False
+                if leeway < 0:
+                    # The token should never expire, we don't need to check
+                    # the end time.
+                    pass
+                elif (
+                    # We don't know the end time. This should never happen, but
+                    # just in case we always expire the token.
+                    pipeline_run_end_time is None
+                    # Calculate whether the token has expired.
+                    or utc_now(tz_aware=pipeline_run_end_time)
+                    > pipeline_run_end_time + timedelta(seconds=leeway)
+                ):
+                    error = (
+                        f"The pipeline run {decoded_token.pipeline_run_id} has "
+                        "finished and API tokens scoped to it are no longer "
+                        "valid. If you want to increase the expiration time "
+                        "of the token to allow steps to continue for longer "
+                        "after other steps have failed, you can do so by "
+                        "configuring the "
+                        f"`{ENV_ZENML_WORKLOAD_TOKEN_EXPIRATION_LEEWAY}` "
+                        "ZenML server environment variable."
                     )
-                except KeyError:
-                    return None
-
-                return step_run.status
-
-            step_run_status = get_step_run_status(decoded_token.step_run_id)
-            if step_run_status is None:
-                error = (
-                    f"Authentication error: error retrieving token step run "
-                    f"{decoded_token.step_run_id}"
-                )
-                logger.error(error)
-                raise CredentialsNotValid(error)
-
-            if step_run_status.is_finished:
-                error = (
-                    f"The execution of step run "
-                    f"{decoded_token.step_run_id} has already concluded and "
-                    "API tokens scoped to it are no longer valid."
-                )
-                logger.error(error)
-                raise CredentialsNotValid(error)
+                    logger.error(error)
+                    raise CredentialsNotValid(error)
 
         auth_context = AuthContext(
             user=user_model,
@@ -861,7 +844,6 @@ def generate_access_token(
     expires_in: Optional[int] = None,
     schedule_id: Optional[UUID] = None,
     pipeline_run_id: Optional[UUID] = None,
-    step_run_id: Optional[UUID] = None,
 ) -> OAuthTokenResponse:
     """Generates an access token for the given user.
 
@@ -880,7 +862,6 @@ def generate_access_token(
             expire.
         schedule_id: The ID of the schedule to scope the token to.
         pipeline_run_id: The ID of the pipeline run to scope the token to.
-        step_run_id: The ID of the step run to scope the token to.
 
     Returns:
         An authentication response with an access token.
@@ -956,7 +937,6 @@ def generate_access_token(
         api_key_id=api_key.id if api_key else None,
         schedule_id=schedule_id,
         pipeline_run_id=pipeline_run_id,
-        step_run_id=step_run_id,
         # Set the session ID if this is a cross-site request
         session_id=session_id,
     ).encode(expires=expires)

zenml/zen_server/download_utils.py

@@ -18,7 +18,6 @@ import tarfile
 import tempfile
 from typing import (
     TYPE_CHECKING,
-    Optional,
 )
 
 from zenml.artifacts.utils import _load_artifact_store
@@ -78,20 +77,15 @@ def verify_artifact_is_downloadable(
 
 def create_artifact_archive(
     artifact: "ArtifactVersionResponse",
-    archive_path: Optional[str] = None,
 ) -> str:
     """Create an archive of the given artifact.
 
     Args:
         artifact: The artifact to archive.
-        archive_path: The path to which to save the archive.
 
     Returns:
         The path to the created archive.
     """
-    if archive_path is None:
-        archive_path = tempfile.mktemp()
-
     artifact_store = verify_artifact_is_downloadable(artifact)
 
     def _prepare_tarinfo(path: str) -> tarfile.TarInfo:
@@ -101,26 +95,29 @@ def create_artifact_archive(
             tarinfo.size = size
         return tarinfo
 
-    with tarfile.open(name=archive_path, mode="w:gz") as tar:
-        if artifact_store.isdir(artifact.uri):
-            for dir, _, files in artifact_store.walk(artifact.uri):
-                dir = dir.decode() if isinstance(dir, bytes) else dir
-                dir_info = tarfile.TarInfo(
-                    name=os.path.relpath(dir, artifact.uri)
-                )
-                dir_info.type = tarfile.DIRTYPE
-                dir_info.mode = 0o755
-                tar.addfile(dir_info)
-
-                for file in files:
-                    file = file.decode() if isinstance(file, bytes) else file
-                    path = os.path.join(dir, file)
-                    tarinfo = _prepare_tarinfo(path)
-                    with artifact_store.open(path, "rb") as f:
-                        tar.addfile(tarinfo, fileobj=f)
-        else:
-            tarinfo = _prepare_tarinfo(artifact.uri)
-            with artifact_store.open(artifact.uri, "rb") as f:
-                tar.addfile(tarinfo, fileobj=f)
-
-    return archive_path
+    with tempfile.NamedTemporaryFile(delete=False) as temp_file:
+        with tarfile.open(fileobj=temp_file, mode="w:gz") as tar:
+            if artifact_store.isdir(artifact.uri):
+                for dir, _, files in artifact_store.walk(artifact.uri):
+                    dir = dir.decode() if isinstance(dir, bytes) else dir
+                    dir_info = tarfile.TarInfo(
+                        name=os.path.relpath(dir, artifact.uri)
+                    )
+                    dir_info.type = tarfile.DIRTYPE
+                    dir_info.mode = 0o755
+                    tar.addfile(dir_info)
+
+                    for file in files:
+                        file = (
+                            file.decode() if isinstance(file, bytes) else file
+                        )
+                        path = os.path.join(dir, file)
+                        tarinfo = _prepare_tarinfo(path)
+                        with artifact_store.open(path, "rb") as f:
+                            tar.addfile(tarinfo, fileobj=f)
+            else:
+                tarinfo = _prepare_tarinfo(artifact.uri)
+                with artifact_store.open(artifact.uri, "rb") as f:
+                    tar.addfile(tarinfo, fileobj=f)
+
+        return temp_file.name

zenml/zen_server/jwt.py

@@ -54,7 +54,6 @@ class JWTToken(BaseModel):
     api_key_id: Optional[UUID] = None
     schedule_id: Optional[UUID] = None
     pipeline_run_id: Optional[UUID] = None
-    step_run_id: Optional[UUID] = None
     session_id: Optional[UUID] = None
     claims: Dict[str, Any] = {}
 
@@ -148,16 +147,6 @@ class JWTToken(BaseModel):
                     "UUID"
                 )
 
-        step_run_id: Optional[UUID] = None
-        if "step_run_id" in claims:
-            try:
-                step_run_id = UUID(claims.pop("step_run_id"))
-            except ValueError:
-                raise CredentialsNotValid(
-                    "Invalid JWT token: the step_run_id claim is not a valid "
-                    "UUID"
-                )
-
         session_id: Optional[UUID] = None
         if "session_id" in claims:
             try:
@@ -174,7 +163,6 @@ class JWTToken(BaseModel):
             api_key_id=api_key_id,
             schedule_id=schedule_id,
             pipeline_run_id=pipeline_run_id,
-            step_run_id=step_run_id,
             session_id=session_id,
             claims=claims,
         )
@@ -212,8 +200,6 @@ class JWTToken(BaseModel):
             claims["schedule_id"] = str(self.schedule_id)
         if self.pipeline_run_id:
             claims["pipeline_run_id"] = str(self.pipeline_run_id)
-        if self.step_run_id:
-            claims["step_run_id"] = str(self.step_run_id)
         if self.session_id:
             claims["session_id"] = str(self.session_id)