Fix workload token expiration for cached steps/runs

zenml-io · Dec 4, 2024 · 4090c42 · 4090c42
1 parent 61988c0
commit 4090c42
Show file tree

Hide file tree

Showing 2 changed files with 4 additions and 17 deletions.
diff --git a/src/zenml/zen_server/auth.py b/src/zenml/zen_server/auth.py
@@ -413,10 +413,7 @@ def get_pipeline_run_status(
                 logger.error(error)
                 raise CredentialsNotValid(error)
 
-            if pipeline_run_status in [
-                ExecutionStatus.FAILED,
-                ExecutionStatus.COMPLETED,
-            ]:
+            if pipeline_run_status.is_finished:
                 error = (
                     f"The execution of pipeline run "
                     f"{decoded_token.pipeline_run_id} has already concluded and "
@@ -461,10 +458,7 @@ def get_step_run_status(
                 logger.error(error)
                 raise CredentialsNotValid(error)
 
-            if step_run_status in [
-                ExecutionStatus.FAILED,
-                ExecutionStatus.COMPLETED,
-            ]:
+            if step_run_status.is_finished:
                 error = (
                     f"The execution of step run "
                     f"{decoded_token.step_run_id} has already concluded and "

diff --git a/src/zenml/zen_server/routers/auth_endpoints.py b/src/zenml/zen_server/routers/auth_endpoints.py
@@ -41,7 +41,6 @@
 from zenml.enums import (
     APITokenType,
     AuthScheme,
-    ExecutionStatus,
     OAuthDeviceStatus,
     OAuthGrantTypes,
 )
@@ -589,10 +588,7 @@ def api_token(
                 "security reasons."
             )
 
-        if pipeline_run.status in [
-            ExecutionStatus.FAILED,
-            ExecutionStatus.COMPLETED,
-        ]:
+        if pipeline_run.status.is_finished:
             raise ValueError(
                 f"The execution of pipeline run {pipeline_run_id} has already "
                 "concluded and API tokens can no longer be generated for it "
@@ -609,10 +605,7 @@ def api_token(
                 "be generated for non-existent step runs for security reasons."
             )
 
-        if step_run.status in [
-            ExecutionStatus.FAILED,
-            ExecutionStatus.COMPLETED,
-        ]:
+        if step_run.status.is_finished:
             raise ValueError(
                 f"The execution of step run {step_run_id} has already "
                 "concluded and API tokens can no longer be generated for it "