zenml-nightly 0.73.0.dev20250124__py3-none-any.whl → 0.73.0.dev20250125__py3-none-any.whl
This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
- zenml/VERSION +1 -1
- zenml/analytics/context.py +2 -6
- zenml/cli/annotator.py +1 -1
- zenml/cli/login.py +15 -6
- zenml/cli/service_connectors.py +5 -5
- zenml/cli/stack.py +2 -2
- zenml/cli/utils.py +2 -54
- zenml/config/pipeline_configurations.py +3 -2
- zenml/config/schedule.py +0 -24
- zenml/event_hub/base_event_hub.py +3 -4
- zenml/integrations/airflow/orchestrators/airflow_orchestrator.py +3 -4
- zenml/integrations/aws/orchestrators/sagemaker_orchestrator.py +11 -9
- zenml/integrations/aws/service_connectors/aws_service_connector.py +8 -13
- zenml/integrations/azure/service_connectors/azure_service_connector.py +4 -10
- zenml/integrations/gcp/service_connectors/gcp_service_connector.py +3 -3
- zenml/integrations/kubernetes/orchestrators/kube_utils.py +3 -3
- zenml/integrations/kubernetes/service_connectors/kubernetes_service_connector.py +6 -2
- zenml/integrations/whylogs/data_validators/whylogs_data_validator.py +2 -3
- zenml/logging/step_logging.py +7 -7
- zenml/login/credentials.py +6 -5
- zenml/login/credentials_store.py +4 -3
- zenml/models/v2/core/api_key.py +5 -2
- zenml/models/v2/core/schedule.py +3 -2
- zenml/orchestrators/publish_utils.py +4 -4
- zenml/orchestrators/step_launcher.py +3 -3
- zenml/orchestrators/step_run_utils.py +2 -2
- zenml/pipelines/run_utils.py +2 -2
- zenml/service_connectors/service_connector.py +7 -4
- zenml/stack/stack.py +5 -4
- zenml/stack_deployments/stack_deployment.py +2 -3
- zenml/utils/string_utils.py +2 -2
- zenml/utils/time_utils.py +138 -0
- zenml/zen_server/auth.py +8 -9
- zenml/zen_server/cloud_utils.py +4 -6
- zenml/zen_server/routers/devices_endpoints.py +2 -4
- zenml/zen_server/zen_server_api.py +9 -8
- zenml/zen_stores/migrations/versions/25155145c545_separate_actions_and_triggers.py +3 -2
- zenml/zen_stores/migrations/versions/3dcc5d20e82f_add_last_user_activity.py +3 -3
- zenml/zen_stores/migrations/versions/46506f72f0ed_add_server_settings.py +3 -2
- zenml/zen_stores/migrations/versions/5994f9ad0489_introduce_role_permissions.py +10 -7
- zenml/zen_stores/migrations/versions/7500f434b71c_remove_shared_columns.py +3 -2
- zenml/zen_stores/migrations/versions/a91762e6be36_artifact_version_table.py +5 -3
- zenml/zen_stores/schemas/action_schemas.py +2 -2
- zenml/zen_stores/schemas/api_key_schemas.py +5 -4
- zenml/zen_stores/schemas/artifact_schemas.py +3 -3
- zenml/zen_stores/schemas/base_schemas.py +5 -7
- zenml/zen_stores/schemas/code_repository_schemas.py +2 -2
- zenml/zen_stores/schemas/component_schemas.py +2 -2
- zenml/zen_stores/schemas/device_schemas.py +5 -4
- zenml/zen_stores/schemas/event_source_schemas.py +2 -2
- zenml/zen_stores/schemas/flavor_schemas.py +2 -2
- zenml/zen_stores/schemas/model_schemas.py +3 -3
- zenml/zen_stores/schemas/pipeline_run_schemas.py +4 -3
- zenml/zen_stores/schemas/pipeline_schemas.py +2 -2
- zenml/zen_stores/schemas/run_template_schemas.py +2 -2
- zenml/zen_stores/schemas/schedule_schema.py +3 -2
- zenml/zen_stores/schemas/secret_schemas.py +2 -2
- zenml/zen_stores/schemas/server_settings_schemas.py +6 -9
- zenml/zen_stores/schemas/service_connector_schemas.py +3 -2
- zenml/zen_stores/schemas/service_schemas.py +2 -2
- zenml/zen_stores/schemas/stack_schemas.py +2 -2
- zenml/zen_stores/schemas/step_run_schemas.py +3 -2
- zenml/zen_stores/schemas/tag_schemas.py +2 -2
- zenml/zen_stores/schemas/trigger_schemas.py +2 -2
- zenml/zen_stores/schemas/user_schemas.py +3 -3
- zenml/zen_stores/schemas/workspace_schemas.py +2 -2
- zenml/zen_stores/sql_zen_store.py +6 -14
- {zenml_nightly-0.73.0.dev20250124.dist-info → zenml_nightly-0.73.0.dev20250125.dist-info}/METADATA +1 -1
- {zenml_nightly-0.73.0.dev20250124.dist-info → zenml_nightly-0.73.0.dev20250125.dist-info}/RECORD +72 -71
- {zenml_nightly-0.73.0.dev20250124.dist-info → zenml_nightly-0.73.0.dev20250125.dist-info}/LICENSE +0 -0
- {zenml_nightly-0.73.0.dev20250124.dist-info → zenml_nightly-0.73.0.dev20250125.dist-info}/WHEEL +0 -0
- {zenml_nightly-0.73.0.dev20250124.dist-info → zenml_nightly-0.73.0.dev20250125.dist-info}/entry_points.txt +0 -0
zenml/login/credentials.py
CHANGED
@@ -13,7 +13,7 @@
|
|
13
13
|
# permissions and limitations under the License.
|
14
14
|
"""ZenML login credentials models."""
|
15
15
|
|
16
|
-
from datetime import datetime, timedelta
|
16
|
+
from datetime import datetime, timedelta
|
17
17
|
from typing import Any, Dict, Optional, Union
|
18
18
|
from urllib.parse import urlparse
|
19
19
|
from uuid import UUID
|
@@ -27,6 +27,7 @@ from zenml.models.v2.misc.server_models import ServerDeploymentType
|
|
27
27
|
from zenml.services.service_status import ServiceState
|
28
28
|
from zenml.utils.enum_utils import StrEnum
|
29
29
|
from zenml.utils.string_utils import get_human_readable_time
|
30
|
+
from zenml.utils.time_utils import to_local_tz, utc_now
|
30
31
|
|
31
32
|
|
32
33
|
class ServerType(StrEnum):
|
@@ -71,7 +72,7 @@ class APIToken(BaseModel):
|
|
71
72
|
expires_at = self.expires_at_with_leeway
|
72
73
|
if not expires_at:
|
73
74
|
return False
|
74
|
-
return expires_at <
|
75
|
+
return expires_at < utc_now(tz_aware=expires_at)
|
75
76
|
|
76
77
|
model_config = ConfigDict(
|
77
78
|
# Allow extra attributes to allow backwards compatibility
|
@@ -223,7 +224,7 @@ class ServerCredentials(BaseModel):
|
|
223
224
|
expires_at = self.api_token.expires_at_with_leeway
|
224
225
|
if not expires_at:
|
225
226
|
return "never expires"
|
226
|
-
if expires_at <
|
227
|
+
if expires_at < utc_now(tz_aware=expires_at):
|
227
228
|
return "expired at " + self.expires_at
|
228
229
|
|
229
230
|
return f"valid until {self.expires_at} (in {self.expires_in})"
|
@@ -242,7 +243,7 @@ class ServerCredentials(BaseModel):
|
|
242
243
|
return "never"
|
243
244
|
|
244
245
|
# Convert the date in the local timezone
|
245
|
-
local_expires_at = expires_at
|
246
|
+
local_expires_at = to_local_tz(expires_at)
|
246
247
|
return local_expires_at.strftime("%Y-%m-%d %H:%M:%S %Z")
|
247
248
|
|
248
249
|
@property
|
@@ -259,7 +260,7 @@ class ServerCredentials(BaseModel):
|
|
259
260
|
return "never"
|
260
261
|
|
261
262
|
# Get the time remaining until the token expires
|
262
|
-
expires_in = expires_at -
|
263
|
+
expires_in = expires_at - utc_now(tz_aware=expires_at)
|
263
264
|
return get_human_readable_time(expires_in.total_seconds())
|
264
265
|
|
265
266
|
@property
|
zenml/login/credentials_store.py
CHANGED
@@ -14,7 +14,7 @@
|
|
14
14
|
"""ZenML login credentials store support."""
|
15
15
|
|
16
16
|
import os
|
17
|
-
from datetime import
|
17
|
+
from datetime import timedelta
|
18
18
|
from typing import Dict, List, Optional, Tuple, Union, cast
|
19
19
|
|
20
20
|
from zenml.config.global_config import GlobalConfiguration
|
@@ -29,6 +29,7 @@ from zenml.login.pro.tenant.models import TenantRead
|
|
29
29
|
from zenml.models import OAuthTokenResponse, ServerModel
|
30
30
|
from zenml.utils import yaml_utils
|
31
31
|
from zenml.utils.singleton import SingletonMetaClass
|
32
|
+
from zenml.utils.time_utils import utc_now_tz_aware
|
32
33
|
|
33
34
|
logger = get_logger(__name__)
|
34
35
|
|
@@ -190,7 +191,7 @@ class CredentialsStore(metaclass=SingletonMetaClass):
|
|
190
191
|
not credential.api_token.expires_at
|
191
192
|
or credential.api_token.expires_at
|
192
193
|
+ timedelta(seconds=TOKEN_STORE_EVICTION_TIME)
|
193
|
-
>
|
194
|
+
> utc_now_tz_aware()
|
194
195
|
)
|
195
196
|
}
|
196
197
|
yaml_utils.write_yaml(credentials_file, credentials_store)
|
@@ -477,7 +478,7 @@ class CredentialsStore(metaclass=SingletonMetaClass):
|
|
477
478
|
"""
|
478
479
|
self.check_and_reload_from_file()
|
479
480
|
if token_response.expires_in:
|
480
|
-
expires_at =
|
481
|
+
expires_at = utc_now_tz_aware() + timedelta(
|
481
482
|
seconds=token_response.expires_in
|
482
483
|
)
|
483
484
|
# Best practice to calculate the leeway depending on the token
|
zenml/models/v2/core/api_key.py
CHANGED
@@ -13,7 +13,7 @@
|
|
13
13
|
# permissions and limitations under the License.
|
14
14
|
"""Models representing API keys."""
|
15
15
|
|
16
|
-
from datetime import datetime, timedelta
|
16
|
+
from datetime import datetime, timedelta
|
17
17
|
from typing import TYPE_CHECKING, ClassVar, List, Optional, Type, Union
|
18
18
|
from uuid import UUID
|
19
19
|
|
@@ -35,6 +35,7 @@ from zenml.models.v2.base.base import (
|
|
35
35
|
)
|
36
36
|
from zenml.models.v2.base.filter import AnyQuery, BaseFilter
|
37
37
|
from zenml.utils.string_utils import b64_decode, b64_encode
|
38
|
+
from zenml.utils.time_utils import utc_now
|
38
39
|
|
39
40
|
if TYPE_CHECKING:
|
40
41
|
from zenml.models.v2.base.filter import AnySchema
|
@@ -319,7 +320,9 @@ class APIKeyInternalResponse(APIKeyResponse):
|
|
319
320
|
and self.retain_period_minutes > 0
|
320
321
|
):
|
321
322
|
# check if the previous key is still valid
|
322
|
-
if
|
323
|
+
if utc_now(
|
324
|
+
tz_aware=self.last_rotated
|
325
|
+
) - self.last_rotated < timedelta(
|
323
326
|
minutes=self.retain_period_minutes
|
324
327
|
):
|
325
328
|
key_hash = self.previous_key
|
zenml/models/v2/core/schedule.py
CHANGED
@@ -31,6 +31,7 @@ from zenml.models.v2.base.scoped import (
|
|
31
31
|
WorkspaceScopedResponseMetadata,
|
32
32
|
WorkspaceScopedResponseResources,
|
33
33
|
)
|
34
|
+
from zenml.utils.time_utils import to_utc_timezone
|
34
35
|
|
35
36
|
logger = get_logger(__name__)
|
36
37
|
|
@@ -182,7 +183,7 @@ class ScheduleResponse(
|
|
182
183
|
if not self.start_time:
|
183
184
|
return None
|
184
185
|
|
185
|
-
return self.start_time
|
186
|
+
return to_utc_timezone(self.start_time).isoformat()
|
186
187
|
|
187
188
|
@property
|
188
189
|
def utc_end_time(self) -> Optional[str]:
|
@@ -194,7 +195,7 @@ class ScheduleResponse(
|
|
194
195
|
if not self.end_time:
|
195
196
|
return None
|
196
197
|
|
197
|
-
return self.end_time
|
198
|
+
return to_utc_timezone(self.end_time).isoformat()
|
198
199
|
|
199
200
|
# Body and metadata properties
|
200
201
|
@property
|
@@ -13,7 +13,6 @@
|
|
13
13
|
# permissions and limitations under the License.
|
14
14
|
"""Utilities to publish pipeline and step runs."""
|
15
15
|
|
16
|
-
from datetime import datetime, timezone
|
17
16
|
from typing import TYPE_CHECKING, Dict, List
|
18
17
|
|
19
18
|
from zenml.client import Client
|
@@ -25,6 +24,7 @@ from zenml.models import (
|
|
25
24
|
StepRunResponse,
|
26
25
|
StepRunUpdate,
|
27
26
|
)
|
27
|
+
from zenml.utils.time_utils import utc_now
|
28
28
|
|
29
29
|
if TYPE_CHECKING:
|
30
30
|
from uuid import UUID
|
@@ -48,7 +48,7 @@ def publish_successful_step_run(
|
|
48
48
|
step_run_id=step_run_id,
|
49
49
|
step_run_update=StepRunUpdate(
|
50
50
|
status=ExecutionStatus.COMPLETED,
|
51
|
-
end_time=
|
51
|
+
end_time=utc_now(),
|
52
52
|
outputs=output_artifact_ids,
|
53
53
|
),
|
54
54
|
)
|
@@ -67,7 +67,7 @@ def publish_failed_step_run(step_run_id: "UUID") -> "StepRunResponse":
|
|
67
67
|
step_run_id=step_run_id,
|
68
68
|
step_run_update=StepRunUpdate(
|
69
69
|
status=ExecutionStatus.FAILED,
|
70
|
-
end_time=
|
70
|
+
end_time=utc_now(),
|
71
71
|
),
|
72
72
|
)
|
73
73
|
|
@@ -87,7 +87,7 @@ def publish_failed_pipeline_run(
|
|
87
87
|
run_id=pipeline_run_id,
|
88
88
|
run_update=PipelineRunUpdate(
|
89
89
|
status=ExecutionStatus.FAILED,
|
90
|
-
end_time=
|
90
|
+
end_time=utc_now(),
|
91
91
|
),
|
92
92
|
)
|
93
93
|
|
@@ -16,7 +16,6 @@
|
|
16
16
|
import os
|
17
17
|
import time
|
18
18
|
from contextlib import nullcontext
|
19
|
-
from datetime import datetime, timezone
|
20
19
|
from functools import partial
|
21
20
|
from typing import TYPE_CHECKING, Any, Callable, Dict, Tuple
|
22
21
|
|
@@ -45,6 +44,7 @@ from zenml.orchestrators import utils as orchestrator_utils
|
|
45
44
|
from zenml.orchestrators.step_runner import StepRunner
|
46
45
|
from zenml.stack import Stack
|
47
46
|
from zenml.utils import string_utils
|
47
|
+
from zenml.utils.time_utils import utc_now
|
48
48
|
|
49
49
|
if TYPE_CHECKING:
|
50
50
|
from zenml.step_operators import BaseStepOperator
|
@@ -201,7 +201,7 @@ class StepLauncher:
|
|
201
201
|
f"Failed preparing step `{self._step_name}`."
|
202
202
|
)
|
203
203
|
step_run_request.status = ExecutionStatus.FAILED
|
204
|
-
step_run_request.end_time =
|
204
|
+
step_run_request.end_time = utc_now()
|
205
205
|
raise
|
206
206
|
finally:
|
207
207
|
step_run = Client().zen_store.create_run_step(
|
@@ -305,7 +305,7 @@ class StepLauncher:
|
|
305
305
|
The created or existing pipeline run,
|
306
306
|
and a boolean indicating whether the run was created or reused.
|
307
307
|
"""
|
308
|
-
start_time =
|
308
|
+
start_time = utc_now()
|
309
309
|
run_name = string_utils.format_name_template(
|
310
310
|
name_template=self._deployment.run_name_template,
|
311
311
|
substitutions=self._deployment.pipeline_configuration._get_full_substitutions(
|
@@ -13,7 +13,6 @@
|
|
13
13
|
# permissions and limitations under the License.
|
14
14
|
"""Utilities for creating step runs."""
|
15
15
|
|
16
|
-
from datetime import datetime, timezone
|
17
16
|
from typing import Dict, List, Optional, Set, Tuple
|
18
17
|
|
19
18
|
from zenml.client import Client
|
@@ -31,6 +30,7 @@ from zenml.models import (
|
|
31
30
|
)
|
32
31
|
from zenml.orchestrators import cache_utils, input_utils, utils
|
33
32
|
from zenml.stack import Stack
|
33
|
+
from zenml.utils.time_utils import utc_now
|
34
34
|
|
35
35
|
logger = get_logger(__name__)
|
36
36
|
|
@@ -75,7 +75,7 @@ class StepRunRequestFactory:
|
|
75
75
|
pipeline_run_id=self.pipeline_run.id,
|
76
76
|
deployment=self.deployment.id,
|
77
77
|
status=ExecutionStatus.RUNNING,
|
78
|
-
start_time=
|
78
|
+
start_time=utc_now(),
|
79
79
|
user=Client().active_user.id,
|
80
80
|
workspace=Client().active_workspace.id,
|
81
81
|
)
|
zenml/pipelines/run_utils.py
CHANGED
@@ -1,7 +1,6 @@
|
|
1
1
|
"""Utility functions for running pipelines."""
|
2
2
|
|
3
3
|
import time
|
4
|
-
from datetime import datetime, timezone
|
5
4
|
from typing import TYPE_CHECKING, Any, Dict, List, Optional, Set, Union
|
6
5
|
from uuid import UUID
|
7
6
|
|
@@ -25,6 +24,7 @@ from zenml.models import (
|
|
25
24
|
from zenml.orchestrators.publish_utils import publish_failed_pipeline_run
|
26
25
|
from zenml.stack import Flavor, Stack
|
27
26
|
from zenml.utils import code_utils, notebook_utils, source_utils, string_utils
|
27
|
+
from zenml.utils.time_utils import utc_now
|
28
28
|
from zenml.zen_stores.base_zen_store import BaseZenStore
|
29
29
|
|
30
30
|
if TYPE_CHECKING:
|
@@ -65,7 +65,7 @@ def create_placeholder_run(
|
|
65
65
|
|
66
66
|
if deployment.schedule:
|
67
67
|
return None
|
68
|
-
start_time =
|
68
|
+
start_time = utc_now()
|
69
69
|
run_request = PipelineRunRequest(
|
70
70
|
name=string_utils.format_name_template(
|
71
71
|
name_template=deployment.run_name_template,
|
@@ -55,6 +55,7 @@ from zenml.models import (
|
|
55
55
|
UserResponse,
|
56
56
|
WorkspaceResponse,
|
57
57
|
)
|
58
|
+
from zenml.utils.time_utils import utc_now
|
58
59
|
|
59
60
|
logger = get_logger(__name__)
|
60
61
|
|
@@ -794,13 +795,14 @@ class ServiceConnector(BaseModel, metaclass=ServiceConnectorMeta):
|
|
794
795
|
"connector configuration is not valid: name and ID must be set"
|
795
796
|
)
|
796
797
|
|
798
|
+
now = utc_now()
|
797
799
|
model = ServiceConnectorResponse(
|
798
800
|
id=id,
|
799
801
|
name=name,
|
800
802
|
body=ServiceConnectorResponseBody(
|
801
803
|
user=user,
|
802
|
-
created=
|
803
|
-
updated=
|
804
|
+
created=now,
|
805
|
+
updated=now,
|
804
806
|
description=description,
|
805
807
|
connector_type=self.get_type(),
|
806
808
|
auth_method=self.auth_method,
|
@@ -845,14 +847,15 @@ class ServiceConnector(BaseModel, metaclass=ServiceConnectorMeta):
|
|
845
847
|
if self.expires_skew_tolerance is not None
|
846
848
|
else SERVICE_CONNECTOR_SKEW_TOLERANCE_SECONDS
|
847
849
|
)
|
848
|
-
|
850
|
+
now = utc_now(tz_aware=expires_at)
|
851
|
+
delta = expires_at - now
|
849
852
|
result = delta < timedelta(seconds=0)
|
850
853
|
|
851
854
|
logger.debug(
|
852
855
|
f"Checking if connector {self.name} has expired.\n"
|
853
856
|
f"Expires at: {self.expires_at}\n"
|
854
857
|
f"Expires at (+skew): {expires_at}\n"
|
855
|
-
f"Current UTC time: {
|
858
|
+
f"Current UTC time: {now}\n"
|
856
859
|
f"Delta: {delta}\n"
|
857
860
|
f"Result: {result}\n"
|
858
861
|
)
|
zenml/stack/stack.py
CHANGED
@@ -16,7 +16,7 @@
|
|
16
16
|
import itertools
|
17
17
|
import json
|
18
18
|
import os
|
19
|
-
from datetime import datetime
|
19
|
+
from datetime import datetime
|
20
20
|
from typing import (
|
21
21
|
TYPE_CHECKING,
|
22
22
|
AbstractSet,
|
@@ -45,6 +45,7 @@ from zenml.logger import get_logger
|
|
45
45
|
from zenml.metadata.metadata_types import MetadataType
|
46
46
|
from zenml.models import StackResponse
|
47
47
|
from zenml.utils import pagination_utils, settings_utils
|
48
|
+
from zenml.utils.time_utils import utc_now
|
48
49
|
|
49
50
|
if TYPE_CHECKING:
|
50
51
|
from zenml.alerter import BaseAlerter
|
@@ -732,7 +733,6 @@ class Stack:
|
|
732
733
|
and not skip_default_image_builder
|
733
734
|
and not self.image_builder
|
734
735
|
):
|
735
|
-
from datetime import datetime
|
736
736
|
from uuid import uuid4
|
737
737
|
|
738
738
|
from zenml.image_builders import (
|
@@ -743,6 +743,7 @@ class Stack:
|
|
743
743
|
|
744
744
|
flavor = LocalImageBuilderFlavor()
|
745
745
|
|
746
|
+
now = utc_now()
|
746
747
|
image_builder = LocalImageBuilder(
|
747
748
|
id=uuid4(),
|
748
749
|
name="temporary_default",
|
@@ -751,8 +752,8 @@ class Stack:
|
|
751
752
|
config=LocalImageBuilderConfig(),
|
752
753
|
user=Client().active_user.id,
|
753
754
|
workspace=Client().active_workspace.id,
|
754
|
-
created=
|
755
|
-
updated=
|
755
|
+
created=now,
|
756
|
+
updated=now,
|
756
757
|
)
|
757
758
|
|
758
759
|
self._image_builder = image_builder
|
@@ -26,6 +26,7 @@ from zenml.models import (
|
|
26
26
|
StackDeploymentConfig,
|
27
27
|
StackDeploymentInfo,
|
28
28
|
)
|
29
|
+
from zenml.utils.time_utils import to_utc_timezone
|
29
30
|
|
30
31
|
STACK_DEPLOYMENT_TERRAFORM = "terraform"
|
31
32
|
|
@@ -199,9 +200,7 @@ class ZenMLCloudStackDeployment(BaseModel):
|
|
199
200
|
# Get all stacks created after the start date
|
200
201
|
|
201
202
|
if date_start and date_start.tzinfo:
|
202
|
-
date_start = date_start
|
203
|
-
tzinfo=None
|
204
|
-
)
|
203
|
+
date_start = to_utc_timezone(date_start).replace(tzinfo=None)
|
205
204
|
stacks = client.list_stacks(
|
206
205
|
created=f"gt:{str(date_start.replace(microsecond=0))}"
|
207
206
|
if date_start
|
zenml/utils/string_utils.py
CHANGED
@@ -17,12 +17,12 @@ import base64
|
|
17
17
|
import functools
|
18
18
|
import random
|
19
19
|
import string
|
20
|
-
from datetime import datetime, timezone
|
21
20
|
from typing import Any, Callable, Dict, Optional, TypeVar, cast
|
22
21
|
|
23
22
|
from pydantic import BaseModel
|
24
23
|
|
25
24
|
from zenml.constants import BANNED_NAME_CHARACTERS
|
25
|
+
from zenml.utils.time_utils import utc_now
|
26
26
|
|
27
27
|
V = TypeVar("V", bound=Any)
|
28
28
|
|
@@ -180,7 +180,7 @@ def format_name_template(
|
|
180
180
|
start_time = None
|
181
181
|
|
182
182
|
if start_time is None:
|
183
|
-
start_time =
|
183
|
+
start_time = utc_now()
|
184
184
|
substitutions.setdefault("date", start_time.strftime("%Y_%m_%d"))
|
185
185
|
substitutions.setdefault("time", start_time.strftime("%H_%M_%S_%f"))
|
186
186
|
|
@@ -0,0 +1,138 @@
|
|
1
|
+
# Copyright (c) ZenML GmbH 2025. All Rights Reserved.
|
2
|
+
#
|
3
|
+
# Licensed under the Apache License, Version 2.0 (the "License");
|
4
|
+
# you may not use this file except in compliance with the License.
|
5
|
+
# You may obtain a copy of the License at:
|
6
|
+
#
|
7
|
+
# https://www.apache.org/licenses/LICENSE-2.0
|
8
|
+
#
|
9
|
+
# Unless required by applicable law or agreed to in writing, software
|
10
|
+
# distributed under the License is distributed on an "AS IS" BASIS,
|
11
|
+
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express
|
12
|
+
# or implied. See the License for the specific language governing
|
13
|
+
# permissions and limitations under the License.
|
14
|
+
"""Time utils."""
|
15
|
+
|
16
|
+
from datetime import datetime, timedelta, timezone
|
17
|
+
from typing import Optional, Union
|
18
|
+
|
19
|
+
|
20
|
+
def utc_now(tz_aware: Union[bool, datetime] = False) -> datetime:
|
21
|
+
"""Get the current time in the UTC timezone.
|
22
|
+
|
23
|
+
Args:
|
24
|
+
tz_aware: Use this flag to control whether the returned datetime is
|
25
|
+
timezone-aware or timezone-naive. If a datetime is provided, the
|
26
|
+
returned datetime will be timezone-aware if and only if the input
|
27
|
+
datetime is also timezone-aware.
|
28
|
+
|
29
|
+
Returns:
|
30
|
+
The current UTC time. If tz_aware is a datetime, the returned datetime
|
31
|
+
will be timezone-aware only if the input datetime is also timezone-aware.
|
32
|
+
If tz_aware is a boolean, the returned datetime will be timezone-aware
|
33
|
+
if True, and timezone-naive if False.
|
34
|
+
"""
|
35
|
+
now = datetime.now(timezone.utc)
|
36
|
+
if (
|
37
|
+
isinstance(tz_aware, bool)
|
38
|
+
and tz_aware is False
|
39
|
+
or isinstance(tz_aware, datetime)
|
40
|
+
and tz_aware.tzinfo is None
|
41
|
+
):
|
42
|
+
return now.replace(tzinfo=None)
|
43
|
+
|
44
|
+
return now
|
45
|
+
|
46
|
+
|
47
|
+
def utc_now_tz_aware() -> datetime:
|
48
|
+
"""Get the current timezone-aware UTC time.
|
49
|
+
|
50
|
+
Returns:
|
51
|
+
The current UTC time.
|
52
|
+
"""
|
53
|
+
return utc_now(tz_aware=True)
|
54
|
+
|
55
|
+
|
56
|
+
def to_local_tz(dt: datetime) -> datetime:
|
57
|
+
"""Convert a datetime to the local timezone.
|
58
|
+
|
59
|
+
If the input datetime is timezone-naive, it will be assumed to be in the UTC
|
60
|
+
timezone.
|
61
|
+
|
62
|
+
Args:
|
63
|
+
dt: datetime to convert.
|
64
|
+
|
65
|
+
Returns:
|
66
|
+
Datetime in the local timezone.
|
67
|
+
"""
|
68
|
+
if dt.tzinfo is None:
|
69
|
+
dt = dt.replace(tzinfo=timezone.utc)
|
70
|
+
return dt.astimezone()
|
71
|
+
|
72
|
+
|
73
|
+
def to_utc_timezone(dt: datetime) -> datetime:
|
74
|
+
"""Convert a datetime to the UTC timezone.
|
75
|
+
|
76
|
+
If the input datetime is timezone-naive, it will be assumed to be in the UTC
|
77
|
+
timezone.
|
78
|
+
|
79
|
+
Args:
|
80
|
+
dt: datetime to convert.
|
81
|
+
|
82
|
+
Returns:
|
83
|
+
Datetime in the UTC timezone.
|
84
|
+
"""
|
85
|
+
if dt.tzinfo is None:
|
86
|
+
dt = dt.replace(tzinfo=timezone.utc)
|
87
|
+
return dt.astimezone(timezone.utc)
|
88
|
+
|
89
|
+
|
90
|
+
def seconds_to_human_readable(time_seconds: int) -> str:
|
91
|
+
"""Converts seconds to human-readable format.
|
92
|
+
|
93
|
+
Args:
|
94
|
+
time_seconds: Seconds to convert.
|
95
|
+
|
96
|
+
Returns:
|
97
|
+
Human readable string.
|
98
|
+
"""
|
99
|
+
seconds = time_seconds % 60
|
100
|
+
minutes = (time_seconds // 60) % 60
|
101
|
+
hours = (time_seconds // 3600) % 24
|
102
|
+
days = time_seconds // 86400
|
103
|
+
tokens = []
|
104
|
+
if days:
|
105
|
+
tokens.append(f"{days}d")
|
106
|
+
if hours:
|
107
|
+
tokens.append(f"{hours}h")
|
108
|
+
if minutes:
|
109
|
+
tokens.append(f"{minutes}m")
|
110
|
+
if seconds:
|
111
|
+
tokens.append(f"{seconds}s")
|
112
|
+
|
113
|
+
return "".join(tokens)
|
114
|
+
|
115
|
+
|
116
|
+
def expires_in(
|
117
|
+
expires_at: datetime,
|
118
|
+
expired_str: str,
|
119
|
+
skew_tolerance: Optional[int] = None,
|
120
|
+
) -> str:
|
121
|
+
"""Returns a human-readable string of the time until an expiration.
|
122
|
+
|
123
|
+
Args:
|
124
|
+
expires_at: Expiration time.
|
125
|
+
expired_str: String to return if the expiration is in the past.
|
126
|
+
skew_tolerance: Seconds of skew tolerance to subtract from the
|
127
|
+
expiration time. If the expiration is within the skew tolerance,
|
128
|
+
the function will return the expired string.
|
129
|
+
|
130
|
+
Returns:
|
131
|
+
Human readable string.
|
132
|
+
"""
|
133
|
+
now = utc_now(tz_aware=expires_at)
|
134
|
+
if skew_tolerance:
|
135
|
+
expires_at -= timedelta(seconds=skew_tolerance)
|
136
|
+
if expires_at < now:
|
137
|
+
return expired_str
|
138
|
+
return seconds_to_human_readable(int((expires_at - now).total_seconds()))
|
zenml/zen_server/auth.py
CHANGED
@@ -14,7 +14,7 @@
|
|
14
14
|
"""Authentication module for ZenML server."""
|
15
15
|
|
16
16
|
from contextvars import ContextVar
|
17
|
-
from datetime import datetime, timedelta
|
17
|
+
from datetime import datetime, timedelta
|
18
18
|
from typing import Callable, Optional, Union
|
19
19
|
from urllib.parse import urlencode, urlparse
|
20
20
|
from uuid import UUID, uuid4
|
@@ -62,6 +62,7 @@ from zenml.models import (
|
|
62
62
|
UserResponse,
|
63
63
|
UserUpdate,
|
64
64
|
)
|
65
|
+
from zenml.utils.time_utils import utc_now
|
65
66
|
from zenml.zen_server.cache import cache_result
|
66
67
|
from zenml.zen_server.csrf import CSRFToken
|
67
68
|
from zenml.zen_server.exceptions import http_exception_from_error
|
@@ -350,7 +351,8 @@ def authenticate_credentials(
|
|
350
351
|
|
351
352
|
if (
|
352
353
|
device_model.expires
|
353
|
-
and
|
354
|
+
and utc_now(tz_aware=device_model.expires)
|
355
|
+
>= device_model.expires
|
354
356
|
):
|
355
357
|
error = (
|
356
358
|
f"Authentication error: device {decoded_token.device_id} "
|
@@ -589,7 +591,7 @@ def authenticate_device(client_id: UUID, device_code: str) -> AuthContext:
|
|
589
591
|
|
590
592
|
if (
|
591
593
|
device_model.expires
|
592
|
-
and
|
594
|
+
and utc_now(tz_aware=device_model.expires) >= device_model.expires
|
593
595
|
):
|
594
596
|
error = (
|
595
597
|
f"Authentication error: device for client ID {client_id} has "
|
@@ -892,21 +894,18 @@ def generate_access_token(
|
|
892
894
|
if expires_in == 0:
|
893
895
|
expires_in = None
|
894
896
|
elif expires_in is not None:
|
895
|
-
expires =
|
897
|
+
expires = utc_now() + timedelta(seconds=expires_in)
|
896
898
|
elif device:
|
897
899
|
# If a device was used for authentication, the token will expire
|
898
900
|
# at the same time as the device.
|
899
901
|
expires = device.expires
|
900
902
|
if expires:
|
901
903
|
expires_in = max(
|
902
|
-
int(
|
903
|
-
expires.timestamp()
|
904
|
-
- datetime.now(timezone.utc).timestamp()
|
905
|
-
),
|
904
|
+
int(expires.timestamp() - utc_now().timestamp()),
|
906
905
|
0,
|
907
906
|
)
|
908
907
|
elif config.jwt_token_expire_minutes:
|
909
|
-
expires =
|
908
|
+
expires = utc_now() + timedelta(
|
910
909
|
minutes=config.jwt_token_expire_minutes
|
911
910
|
)
|
912
911
|
expires_in = config.jwt_token_expire_minutes * 60
|
zenml/zen_server/cloud_utils.py
CHANGED
@@ -1,6 +1,6 @@
|
|
1
1
|
"""Utils concerning anything concerning the cloud control plane backend."""
|
2
2
|
|
3
|
-
from datetime import datetime, timedelta
|
3
|
+
from datetime import datetime, timedelta
|
4
4
|
from typing import Any, Dict, Optional
|
5
5
|
|
6
6
|
import requests
|
@@ -8,6 +8,7 @@ from requests.adapters import HTTPAdapter, Retry
|
|
8
8
|
|
9
9
|
from zenml.config.server_config import ServerProConfiguration
|
10
10
|
from zenml.exceptions import SubscriptionUpgradeRequiredError
|
11
|
+
from zenml.utils.time_utils import utc_now
|
11
12
|
from zenml.zen_server.utils import get_zenml_headers, server_config
|
12
13
|
|
13
14
|
_cloud_connection: Optional["ZenMLCloudConnection"] = None
|
@@ -185,8 +186,7 @@ class ZenMLCloudConnection:
|
|
185
186
|
if (
|
186
187
|
self._token is not None
|
187
188
|
and self._token_expires_at is not None
|
188
|
-
and
|
189
|
-
< self._token_expires_at
|
189
|
+
and utc_now() + timedelta(minutes=5) < self._token_expires_at
|
190
190
|
):
|
191
191
|
return self._token
|
192
192
|
|
@@ -227,9 +227,7 @@ class ZenMLCloudConnection:
|
|
227
227
|
)
|
228
228
|
|
229
229
|
self._token = access_token
|
230
|
-
self._token_expires_at =
|
231
|
-
seconds=expires_in
|
232
|
-
)
|
230
|
+
self._token_expires_at = utc_now() + timedelta(seconds=expires_in)
|
233
231
|
|
234
232
|
assert self._token is not None
|
235
233
|
return self._token
|
@@ -13,7 +13,6 @@
|
|
13
13
|
# permissions and limitations under the License.
|
14
14
|
"""Endpoint definitions for code repositories."""
|
15
15
|
|
16
|
-
from datetime import datetime, timezone
|
17
16
|
from typing import Optional
|
18
17
|
from uuid import UUID
|
19
18
|
|
@@ -36,6 +35,7 @@ from zenml.models import (
|
|
36
35
|
OAuthDeviceVerificationRequest,
|
37
36
|
Page,
|
38
37
|
)
|
38
|
+
from zenml.utils.time_utils import utc_now
|
39
39
|
from zenml.zen_server.auth import AuthContext, authorize
|
40
40
|
from zenml.zen_server.exceptions import error_response
|
41
41
|
from zenml.zen_server.utils import (
|
@@ -219,9 +219,7 @@ def verify_authorized_device(
|
|
219
219
|
)
|
220
220
|
|
221
221
|
# Check if the device verification has expired.
|
222
|
-
if device_model.expires and device_model.expires <
|
223
|
-
timezone.utc
|
224
|
-
):
|
222
|
+
if device_model.expires and device_model.expires < utc_now():
|
225
223
|
raise ValueError(
|
226
224
|
"Invalid request: device verification expired.",
|
227
225
|
)
|