zenml-nightly 0.75.0.dev20250313__py3-none-any.whl → 0.75.0.dev20250315__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 +4 -4
- zenml/analytics/enums.py +2 -2
- zenml/artifacts/utils.py +2 -2
- zenml/cli/__init__.py +8 -9
- zenml/cli/base.py +2 -2
- zenml/cli/code_repository.py +1 -1
- zenml/cli/login.py +21 -18
- zenml/cli/pipeline.py +3 -3
- zenml/cli/project.py +172 -0
- zenml/cli/server.py +5 -5
- zenml/cli/service_accounts.py +0 -1
- zenml/cli/service_connectors.py +15 -16
- zenml/cli/stack.py +0 -2
- zenml/cli/stack_components.py +2 -2
- zenml/cli/utils.py +3 -3
- zenml/client.py +352 -341
- zenml/config/global_config.py +41 -43
- zenml/config/server_config.py +9 -9
- zenml/constants.py +5 -3
- zenml/event_hub/event_hub.py +1 -1
- zenml/integrations/gcp/__init__.py +1 -0
- zenml/integrations/gcp/flavors/vertex_orchestrator_flavor.py +5 -0
- zenml/integrations/gcp/flavors/vertex_step_operator_flavor.py +5 -28
- zenml/integrations/gcp/orchestrators/vertex_orchestrator.py +125 -78
- zenml/integrations/gcp/service_connectors/gcp_service_connector.py +7 -6
- zenml/integrations/gcp/vertex_custom_job_parameters.py +50 -0
- zenml/integrations/mlflow/steps/mlflow_registry.py +3 -3
- zenml/integrations/wandb/__init__.py +1 -1
- zenml/integrations/wandb/experiment_trackers/wandb_experiment_tracker.py +29 -9
- zenml/integrations/wandb/flavors/wandb_experiment_tracker_flavor.py +2 -0
- zenml/login/credentials.py +26 -27
- zenml/login/credentials_store.py +5 -5
- zenml/login/pro/client.py +9 -9
- zenml/login/pro/utils.py +8 -8
- zenml/login/pro/{tenant → workspace}/__init__.py +1 -1
- zenml/login/pro/{tenant → workspace}/client.py +25 -25
- zenml/login/pro/{tenant → workspace}/models.py +27 -28
- zenml/model/model.py +2 -2
- zenml/model_registries/base_model_registry.py +1 -1
- zenml/models/__init__.py +29 -29
- zenml/models/v2/base/filter.py +1 -1
- zenml/models/v2/base/scoped.py +49 -53
- zenml/models/v2/core/action.py +12 -12
- zenml/models/v2/core/artifact.py +15 -15
- zenml/models/v2/core/artifact_version.py +15 -15
- zenml/models/v2/core/code_repository.py +12 -12
- zenml/models/v2/core/event_source.py +12 -12
- zenml/models/v2/core/model.py +26 -18
- zenml/models/v2/core/model_version.py +15 -15
- zenml/models/v2/core/pipeline.py +15 -15
- zenml/models/v2/core/pipeline_build.py +14 -14
- zenml/models/v2/core/pipeline_deployment.py +12 -14
- zenml/models/v2/core/pipeline_run.py +16 -16
- zenml/models/v2/core/project.py +203 -0
- zenml/models/v2/core/run_metadata.py +2 -2
- zenml/models/v2/core/run_template.py +15 -15
- zenml/models/v2/core/schedule.py +12 -12
- zenml/models/v2/core/secret.py +1 -1
- zenml/models/v2/core/service.py +14 -14
- zenml/models/v2/core/step_run.py +13 -13
- zenml/models/v2/core/tag.py +96 -3
- zenml/models/v2/core/trigger.py +13 -13
- zenml/models/v2/core/trigger_execution.py +2 -2
- zenml/models/v2/core/user.py +0 -17
- zenml/models/v2/misc/server_models.py +6 -6
- zenml/models/v2/misc/statistics.py +4 -4
- zenml/orchestrators/cache_utils.py +7 -7
- zenml/orchestrators/input_utils.py +1 -1
- zenml/orchestrators/step_launcher.py +1 -1
- zenml/orchestrators/step_run_utils.py +3 -3
- zenml/orchestrators/utils.py +4 -4
- zenml/pipelines/build_utils.py +2 -2
- zenml/pipelines/pipeline_definition.py +5 -5
- zenml/pipelines/run_utils.py +1 -1
- zenml/service_connectors/service_connector.py +0 -3
- zenml/service_connectors/service_connector_utils.py +0 -1
- zenml/stack/stack.py +0 -1
- zenml/steps/base_step.py +10 -2
- zenml/utils/dashboard_utils.py +1 -1
- zenml/utils/tag_utils.py +0 -12
- zenml/zen_server/cloud_utils.py +3 -3
- zenml/zen_server/feature_gate/endpoint_utils.py +1 -1
- zenml/zen_server/feature_gate/zenml_cloud_feature_gate.py +1 -1
- zenml/zen_server/rbac/endpoint_utils.py +17 -17
- zenml/zen_server/rbac/models.py +47 -22
- zenml/zen_server/rbac/rbac_sql_zen_store.py +3 -3
- zenml/zen_server/rbac/utils.py +23 -25
- zenml/zen_server/rbac/zenml_cloud_rbac.py +7 -74
- zenml/zen_server/routers/artifact_version_endpoints.py +10 -10
- zenml/zen_server/routers/auth_endpoints.py +6 -6
- zenml/zen_server/routers/code_repositories_endpoints.py +12 -14
- zenml/zen_server/routers/model_versions_endpoints.py +13 -15
- zenml/zen_server/routers/models_endpoints.py +7 -9
- zenml/zen_server/routers/pipeline_builds_endpoints.py +14 -16
- zenml/zen_server/routers/pipeline_deployments_endpoints.py +13 -15
- zenml/zen_server/routers/pipelines_endpoints.py +16 -18
- zenml/zen_server/routers/{workspaces_endpoints.py → projects_endpoints.py} +111 -68
- zenml/zen_server/routers/run_metadata_endpoints.py +7 -9
- zenml/zen_server/routers/run_templates_endpoints.py +15 -17
- zenml/zen_server/routers/runs_endpoints.py +12 -14
- zenml/zen_server/routers/schedule_endpoints.py +12 -14
- zenml/zen_server/routers/secrets_endpoints.py +1 -3
- zenml/zen_server/routers/server_endpoints.py +7 -7
- zenml/zen_server/routers/service_connectors_endpoints.py +11 -13
- zenml/zen_server/routers/service_endpoints.py +7 -9
- zenml/zen_server/routers/stack_components_endpoints.py +9 -11
- zenml/zen_server/routers/stacks_endpoints.py +9 -11
- zenml/zen_server/routers/steps_endpoints.py +6 -6
- zenml/zen_server/routers/users_endpoints.py +5 -43
- zenml/zen_server/template_execution/utils.py +4 -4
- zenml/zen_server/utils.py +10 -10
- zenml/zen_server/zen_server_api.py +6 -5
- zenml/zen_stores/base_zen_store.py +38 -42
- zenml/zen_stores/migrations/versions/12eff0206201_rename_workspace_to_project.py +768 -0
- zenml/zen_stores/migrations/versions/41b28cae31ce_make_artifacts_workspace_scoped.py +3 -3
- zenml/zen_stores/migrations/versions/cbc6acd71f92_add_workspace_display_name.py +58 -0
- zenml/zen_stores/rest_zen_store.py +55 -63
- zenml/zen_stores/schemas/__init__.py +2 -2
- zenml/zen_stores/schemas/action_schemas.py +9 -9
- zenml/zen_stores/schemas/artifact_schemas.py +15 -17
- zenml/zen_stores/schemas/code_repository_schemas.py +16 -18
- zenml/zen_stores/schemas/event_source_schemas.py +9 -9
- zenml/zen_stores/schemas/model_schemas.py +15 -17
- zenml/zen_stores/schemas/pipeline_build_schemas.py +7 -7
- zenml/zen_stores/schemas/pipeline_deployment_schemas.py +7 -7
- zenml/zen_stores/schemas/pipeline_run_schemas.py +9 -9
- zenml/zen_stores/schemas/pipeline_schemas.py +9 -9
- zenml/zen_stores/schemas/{workspace_schemas.py → project_schemas.py} +47 -41
- zenml/zen_stores/schemas/run_metadata_schemas.py +5 -5
- zenml/zen_stores/schemas/run_template_schemas.py +9 -9
- zenml/zen_stores/schemas/schedule_schema.py +9 -9
- zenml/zen_stores/schemas/service_schemas.py +7 -7
- zenml/zen_stores/schemas/step_run_schemas.py +7 -7
- zenml/zen_stores/schemas/trigger_schemas.py +9 -9
- zenml/zen_stores/schemas/user_schemas.py +0 -12
- zenml/zen_stores/sql_zen_store.py +318 -275
- zenml/zen_stores/zen_store_interface.py +56 -70
- {zenml_nightly-0.75.0.dev20250313.dist-info → zenml_nightly-0.75.0.dev20250315.dist-info}/METADATA +1 -1
- {zenml_nightly-0.75.0.dev20250313.dist-info → zenml_nightly-0.75.0.dev20250315.dist-info}/RECORD +143 -140
- zenml/cli/workspace.py +0 -160
- zenml/models/v2/core/workspace.py +0 -131
- {zenml_nightly-0.75.0.dev20250313.dist-info → zenml_nightly-0.75.0.dev20250315.dist-info}/LICENSE +0 -0
- {zenml_nightly-0.75.0.dev20250313.dist-info → zenml_nightly-0.75.0.dev20250315.dist-info}/WHEEL +0 -0
- {zenml_nightly-0.75.0.dev20250313.dist-info → zenml_nightly-0.75.0.dev20250315.dist-info}/entry_points.txt +0 -0
@@ -438,7 +438,7 @@ class GCPUserAccountConfig(GCPBaseProjectIDConfig, GCPUserAccountCredentials):
|
|
438
438
|
class GCPServiceAccountConfig(GCPBaseConfig, GCPServiceAccountCredentials):
|
439
439
|
"""GCP service account configuration."""
|
440
440
|
|
441
|
-
|
441
|
+
project_id: Optional[str] = None
|
442
442
|
|
443
443
|
@property
|
444
444
|
def gcp_project_id(self) -> str:
|
@@ -450,14 +450,14 @@ class GCPServiceAccountConfig(GCPBaseConfig, GCPServiceAccountCredentials):
|
|
450
450
|
Returns:
|
451
451
|
The GCP project ID.
|
452
452
|
"""
|
453
|
-
if self.
|
454
|
-
self.
|
453
|
+
if self.project_id is None:
|
454
|
+
self.project_id = json.loads(
|
455
455
|
self.service_account_json.get_secret_value()
|
456
456
|
)["project_id"]
|
457
457
|
# Guaranteed by the field validator
|
458
|
-
assert self.
|
458
|
+
assert self.project_id is not None
|
459
459
|
|
460
|
-
return self.
|
460
|
+
return self.project_id
|
461
461
|
|
462
462
|
|
463
463
|
class GCPExternalAccountConfig(
|
@@ -798,7 +798,8 @@ connector will distribute the service account credentials JSON to clients
|
|
798
798
|
instead (not recommended).
|
799
799
|
|
800
800
|
A GCP project is required and the connector may only be used to access GCP
|
801
|
-
resources in the specified project.
|
801
|
+
resources in the specified project. If the `project_id` is not provided, the
|
802
|
+
connector will use the one extracted from the service account key JSON.
|
802
803
|
|
803
804
|
If you already have the GOOGLE_APPLICATION_CREDENTIALS environment variable
|
804
805
|
configured to point to a service account key JSON file, it will be automatically
|
@@ -0,0 +1,50 @@
|
|
1
|
+
# Copyright (c) ZenML GmbH 2022. 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
|
+
"""Vertex custom job parameter model."""
|
15
|
+
|
16
|
+
from typing import Optional
|
17
|
+
|
18
|
+
from pydantic import BaseModel
|
19
|
+
|
20
|
+
|
21
|
+
class VertexCustomJobParameters(BaseModel):
|
22
|
+
"""Settings for the Vertex custom job parameters.
|
23
|
+
|
24
|
+
Attributes:
|
25
|
+
accelerator_type: Defines which accelerator (GPU, TPU) is used for the
|
26
|
+
job. Check out out this table to see which accelerator
|
27
|
+
type and count are compatible with your chosen machine type:
|
28
|
+
https://cloud.google.com/vertex-ai/docs/training/configure-compute#gpu-compatibility-table.
|
29
|
+
accelerator_count: Defines number of accelerators to be used for the
|
30
|
+
job. Check out out this table to see which accelerator
|
31
|
+
type and count are compatible with your chosen machine type:
|
32
|
+
https://cloud.google.com/vertex-ai/docs/training/configure-compute#gpu-compatibility-table.
|
33
|
+
machine_type: Machine type specified here
|
34
|
+
https://cloud.google.com/vertex-ai/docs/training/configure-compute#machine-types.
|
35
|
+
boot_disk_size_gb: Size of the boot disk in GB. (Default: 100)
|
36
|
+
https://cloud.google.com/vertex-ai/docs/training/configure-compute#boot_disk_options
|
37
|
+
boot_disk_type: Type of the boot disk. (Default: pd-ssd)
|
38
|
+
https://cloud.google.com/vertex-ai/docs/training/configure-compute#boot_disk_options
|
39
|
+
persistent_resource_id: The ID of the persistent resource to use for the job.
|
40
|
+
https://cloud.google.com/vertex-ai/docs/training/persistent-resource-overview
|
41
|
+
service_account: Specifies the service account to be used.
|
42
|
+
"""
|
43
|
+
|
44
|
+
accelerator_type: Optional[str] = None
|
45
|
+
accelerator_count: int = 0
|
46
|
+
machine_type: str = "n1-standard-4"
|
47
|
+
boot_disk_size_gb: int = 100
|
48
|
+
boot_disk_type: str = "pd-ssd"
|
49
|
+
persistent_resource_id: Optional[str] = None
|
50
|
+
service_account: Optional[str] = None
|
@@ -92,7 +92,7 @@ def mlflow_register_model_step(
|
|
92
92
|
pipeline_name = step_context.pipeline.name
|
93
93
|
current_run_name = step_context.pipeline_run.name
|
94
94
|
pipeline_run_uuid = str(step_context.pipeline_run.id)
|
95
|
-
|
95
|
+
zenml_project = str(step_context.pipeline.project.name)
|
96
96
|
|
97
97
|
# Get MLflow run ID either from params or from experiment tracker using
|
98
98
|
# pipeline name and run name
|
@@ -144,8 +144,8 @@ def mlflow_register_model_step(
|
|
144
144
|
metadata.zenml_run_name = run_name
|
145
145
|
if metadata.zenml_pipeline_run_uuid is None:
|
146
146
|
metadata.zenml_pipeline_run_uuid = pipeline_run_uuid
|
147
|
-
if metadata.
|
148
|
-
metadata.
|
147
|
+
if metadata.zenml_project is None:
|
148
|
+
metadata.zenml_project = zenml_project
|
149
149
|
if getattr(metadata, "mlflow_run_id", None) is None:
|
150
150
|
setattr(metadata, "mlflow_run_id", mlflow_run_id)
|
151
151
|
|
@@ -30,7 +30,7 @@ class WandbIntegration(Integration):
|
|
30
30
|
"""Definition of Plotly integration for ZenML."""
|
31
31
|
|
32
32
|
NAME = WANDB
|
33
|
-
REQUIREMENTS = ["wandb>=0.12.12", "Pillow>=9.1.0"]
|
33
|
+
REQUIREMENTS = ["wandb>=0.12.12,<1.0.0", "Pillow>=9.1.0", "weave>=0.51.33,<1.0.0"]
|
34
34
|
REQUIREMENTS_IGNORED_ON_UNINSTALL = ["Pillow"]
|
35
35
|
|
36
36
|
@classmethod
|
@@ -14,7 +14,7 @@
|
|
14
14
|
"""Implementation for the wandb experiment tracker."""
|
15
15
|
|
16
16
|
import os
|
17
|
-
from typing import TYPE_CHECKING,
|
17
|
+
from typing import TYPE_CHECKING, Dict, List, Optional, Type, cast
|
18
18
|
|
19
19
|
import wandb
|
20
20
|
|
@@ -30,8 +30,6 @@ from zenml.logger import get_logger
|
|
30
30
|
from zenml.metadata.metadata_types import Uri
|
31
31
|
|
32
32
|
if TYPE_CHECKING:
|
33
|
-
from wandb import Settings
|
34
|
-
|
35
33
|
from zenml.config.step_run_info import StepRunInfo
|
36
34
|
from zenml.metadata.metadata_types import MetadataType
|
37
35
|
|
@@ -76,9 +74,7 @@ class WandbExperimentTracker(BaseExperimentTracker):
|
|
76
74
|
wandb_run_name = (
|
77
75
|
settings.run_name or f"{info.run_name}_{info.pipeline_step_name}"
|
78
76
|
)
|
79
|
-
self._initialize_wandb(
|
80
|
-
run_name=wandb_run_name, tags=tags, settings=settings.settings
|
81
|
-
)
|
77
|
+
self._initialize_wandb(run_name=wandb_run_name, tags=tags, info=info)
|
82
78
|
|
83
79
|
def get_step_run_metadata(
|
84
80
|
self, info: "StepRunInfo"
|
@@ -131,25 +127,49 @@ class WandbExperimentTracker(BaseExperimentTracker):
|
|
131
127
|
|
132
128
|
def _initialize_wandb(
|
133
129
|
self,
|
130
|
+
info: "StepRunInfo",
|
134
131
|
run_name: str,
|
135
132
|
tags: List[str],
|
136
|
-
settings: Union["Settings", Dict[str, Any], None] = None,
|
137
133
|
) -> None:
|
138
134
|
"""Initializes a wandb run.
|
139
135
|
|
140
136
|
Args:
|
137
|
+
info: Step run information.
|
141
138
|
run_name: Name of the wandb run to create.
|
142
139
|
tags: Tags to attach to the wandb run.
|
143
|
-
settings: Additional settings for the wandb run.
|
144
140
|
"""
|
145
141
|
logger.info(
|
146
142
|
f"Initializing wandb with entity {self.config.entity}, project "
|
147
143
|
f"name: {self.config.project_name}, run_name: {run_name}."
|
148
144
|
)
|
145
|
+
settings = cast(
|
146
|
+
WandbExperimentTrackerSettings, self.get_settings(info)
|
147
|
+
)
|
149
148
|
wandb.init(
|
150
149
|
entity=self.config.entity,
|
151
150
|
project=self.config.project_name,
|
152
151
|
name=run_name,
|
153
152
|
tags=tags,
|
154
|
-
settings=settings,
|
153
|
+
settings=settings.settings,
|
155
154
|
)
|
155
|
+
|
156
|
+
if settings.enable_weave:
|
157
|
+
import weave
|
158
|
+
|
159
|
+
if self.config.project_name:
|
160
|
+
logger.info("Initializing weave")
|
161
|
+
weave.init(project_name=self.config.project_name)
|
162
|
+
else:
|
163
|
+
logger.info(
|
164
|
+
"Weave enabled but no project_name specified. "
|
165
|
+
"Skipping weave initialization."
|
166
|
+
)
|
167
|
+
else:
|
168
|
+
import weave
|
169
|
+
|
170
|
+
if self.config.project_name:
|
171
|
+
logger.info("Disabling weave")
|
172
|
+
weave.init(
|
173
|
+
project_name=self.config.project_name,
|
174
|
+
settings={"disabled": True},
|
175
|
+
)
|
@@ -46,11 +46,13 @@ class WandbExperimentTrackerSettings(BaseSettings):
|
|
46
46
|
run_name: The Wandb run name.
|
47
47
|
tags: Tags for the Wandb run.
|
48
48
|
settings: Settings for the Wandb run.
|
49
|
+
enable_weave: Whether to enable Weave integration.
|
49
50
|
"""
|
50
51
|
|
51
52
|
run_name: Optional[str] = None
|
52
53
|
tags: List[str] = []
|
53
54
|
settings: Dict[str, Any] = {}
|
55
|
+
enable_weave: bool = False
|
54
56
|
|
55
57
|
@field_validator("settings", mode="before")
|
56
58
|
@classmethod
|
zenml/login/credentials.py
CHANGED
@@ -21,7 +21,7 @@ from uuid import UUID
|
|
21
21
|
from pydantic import BaseModel, ConfigDict
|
22
22
|
|
23
23
|
from zenml.login.pro.constants import ZENML_PRO_API_URL, ZENML_PRO_URL
|
24
|
-
from zenml.login.pro.
|
24
|
+
from zenml.login.pro.workspace.models import WorkspaceRead, WorkspaceStatus
|
25
25
|
from zenml.models import ServerModel
|
26
26
|
from zenml.models.v2.misc.server_models import ServerDeploymentType
|
27
27
|
from zenml.services.service_status import ServiceState
|
@@ -99,8 +99,8 @@ class ServerCredentials(BaseModel):
|
|
99
99
|
# Pro server attributes
|
100
100
|
organization_name: Optional[str] = None
|
101
101
|
organization_id: Optional[UUID] = None
|
102
|
-
|
103
|
-
|
102
|
+
workspace_name: Optional[str] = None
|
103
|
+
workspace_id: Optional[UUID] = None
|
104
104
|
pro_api_url: Optional[str] = None
|
105
105
|
pro_dashboard_url: Optional[str] = None
|
106
106
|
|
@@ -128,7 +128,7 @@ class ServerCredentials(BaseModel):
|
|
128
128
|
return ServerType.PRO_API
|
129
129
|
if self.url == self.pro_api_url:
|
130
130
|
return ServerType.PRO_API
|
131
|
-
if self.organization_id or self.
|
131
|
+
if self.organization_id or self.workspace_id:
|
132
132
|
return ServerType.PRO
|
133
133
|
if urlparse(self.url).hostname in [
|
134
134
|
"localhost",
|
@@ -139,9 +139,9 @@ class ServerCredentials(BaseModel):
|
|
139
139
|
return ServerType.REMOTE
|
140
140
|
|
141
141
|
def update_server_info(
|
142
|
-
self, server_info: Union[ServerModel,
|
142
|
+
self, server_info: Union[ServerModel, WorkspaceRead]
|
143
143
|
) -> None:
|
144
|
-
"""Update with server information received from the server itself or from a ZenML Pro
|
144
|
+
"""Update with server information received from the server itself or from a ZenML Pro workspace descriptor.
|
145
145
|
|
146
146
|
Args:
|
147
147
|
server_info: The server information to update with.
|
@@ -152,20 +152,20 @@ class ServerCredentials(BaseModel):
|
|
152
152
|
# All other attributes can change during the lifetime of the server
|
153
153
|
self.deployment_type = server_info.deployment_type
|
154
154
|
server_name = (
|
155
|
-
server_info.
|
156
|
-
or server_info.metadata.get("
|
155
|
+
server_info.pro_workspace_name
|
156
|
+
or server_info.metadata.get("workspace_name")
|
157
157
|
or server_info.name
|
158
158
|
)
|
159
159
|
if server_name:
|
160
160
|
self.server_name = server_name
|
161
161
|
if server_info.pro_organization_id:
|
162
162
|
self.organization_id = server_info.pro_organization_id
|
163
|
-
if server_info.
|
164
|
-
self.server_id = server_info.
|
163
|
+
if server_info.pro_workspace_id:
|
164
|
+
self.server_id = server_info.pro_workspace_id
|
165
165
|
if server_info.pro_organization_name:
|
166
166
|
self.organization_name = server_info.pro_organization_name
|
167
|
-
if server_info.
|
168
|
-
self.
|
167
|
+
if server_info.pro_workspace_name:
|
168
|
+
self.workspace_name = server_info.pro_workspace_name
|
169
169
|
if server_info.pro_api_url:
|
170
170
|
self.pro_api_url = server_info.pro_api_url
|
171
171
|
if server_info.pro_dashboard_url:
|
@@ -180,8 +180,8 @@ class ServerCredentials(BaseModel):
|
|
180
180
|
self.server_name = server_info.name
|
181
181
|
self.organization_name = server_info.organization_name
|
182
182
|
self.organization_id = server_info.organization_id
|
183
|
-
self.
|
184
|
-
self.
|
183
|
+
self.workspace_name = server_info.name
|
184
|
+
self.workspace_id = server_info.id
|
185
185
|
self.status = server_info.status
|
186
186
|
self.version = server_info.version
|
187
187
|
|
@@ -192,7 +192,7 @@ class ServerCredentials(BaseModel):
|
|
192
192
|
Returns:
|
193
193
|
True if the server is available, False otherwise.
|
194
194
|
"""
|
195
|
-
if self.status not in [
|
195
|
+
if self.status not in [WorkspaceStatus.AVAILABLE, ServiceState.ACTIVE]:
|
196
196
|
return False
|
197
197
|
if (
|
198
198
|
self.api_key
|
@@ -270,20 +270,19 @@ class ServerCredentials(BaseModel):
|
|
270
270
|
Returns:
|
271
271
|
The URL to the ZenML dashboard for this server.
|
272
272
|
"""
|
273
|
-
if self.
|
273
|
+
if self.pro_dashboard_url and self.workspace_name:
|
274
274
|
return (
|
275
|
-
|
276
|
-
|
277
|
-
)
|
275
|
+
self.pro_dashboard_url or ZENML_PRO_URL
|
276
|
+
) + f"/workspaces/{str(self.workspace_name)}"
|
278
277
|
|
279
278
|
return self.url
|
280
279
|
|
281
280
|
@property
|
282
281
|
def dashboard_organization_url(self) -> str:
|
283
|
-
"""Get the URL to the ZenML Pro dashboard for this
|
282
|
+
"""Get the URL to the ZenML Pro dashboard for this workspace's organization.
|
284
283
|
|
285
284
|
Returns:
|
286
|
-
The URL to the ZenML Pro dashboard for this
|
285
|
+
The URL to the ZenML Pro dashboard for this workspace's organization.
|
287
286
|
"""
|
288
287
|
if self.organization_id:
|
289
288
|
return (
|
@@ -293,19 +292,19 @@ class ServerCredentials(BaseModel):
|
|
293
292
|
|
294
293
|
@property
|
295
294
|
def dashboard_hyperlink(self) -> str:
|
296
|
-
"""Get the hyperlink to the ZenML dashboard for this
|
295
|
+
"""Get the hyperlink to the ZenML dashboard for this workspace.
|
297
296
|
|
298
297
|
Returns:
|
299
|
-
The hyperlink to the ZenML dashboard for this
|
298
|
+
The hyperlink to the ZenML dashboard for this workspace.
|
300
299
|
"""
|
301
300
|
return f"[link={self.dashboard_url}]{self.dashboard_url}[/link]"
|
302
301
|
|
303
302
|
@property
|
304
303
|
def api_hyperlink(self) -> str:
|
305
|
-
"""Get the hyperlink to the ZenML OpenAPI dashboard for this
|
304
|
+
"""Get the hyperlink to the ZenML OpenAPI dashboard for this workspace.
|
306
305
|
|
307
306
|
Returns:
|
308
|
-
The hyperlink to the ZenML OpenAPI dashboard for this
|
307
|
+
The hyperlink to the ZenML OpenAPI dashboard for this workspace.
|
309
308
|
"""
|
310
309
|
api_url = self.url + "/docs"
|
311
310
|
return f"[link={api_url}]{self.url}[/link]"
|
@@ -360,10 +359,10 @@ class ServerCredentials(BaseModel):
|
|
360
359
|
|
361
360
|
@property
|
362
361
|
def organization_id_hyperlink(self) -> str:
|
363
|
-
"""Get the hyperlink to the ZenML Pro dashboard for this
|
362
|
+
"""Get the hyperlink to the ZenML Pro dashboard for this workspace's organization using its ID.
|
364
363
|
|
365
364
|
Returns:
|
366
|
-
The hyperlink to the ZenML Pro dashboard for this
|
365
|
+
The hyperlink to the ZenML Pro dashboard for this workspace's
|
367
366
|
organization using its ID.
|
368
367
|
"""
|
369
368
|
if self.organization_id is None:
|
zenml/login/credentials_store.py
CHANGED
@@ -25,7 +25,7 @@ from zenml.constants import (
|
|
25
25
|
from zenml.io import fileio
|
26
26
|
from zenml.logger import get_logger
|
27
27
|
from zenml.login.credentials import APIToken, ServerCredentials, ServerType
|
28
|
-
from zenml.login.pro.
|
28
|
+
from zenml.login.pro.workspace.models import WorkspaceRead
|
29
29
|
from zenml.models import OAuthTokenResponse, ServerModel
|
30
30
|
from zenml.utils import yaml_utils
|
31
31
|
from zenml.utils.singleton import SingletonMetaClass
|
@@ -75,7 +75,7 @@ class CredentialsStore(metaclass=SingletonMetaClass):
|
|
75
75
|
|
76
76
|
Alongside credentials, the Credentials Store is also used to store
|
77
77
|
additional server information:
|
78
|
-
* ZenML Pro
|
78
|
+
* ZenML Pro workspace information populated by the `zenml login` command
|
79
79
|
* ZenML server information populated by the REST zen store by fetching
|
80
80
|
the server's information endpoint after authenticating
|
81
81
|
|
@@ -341,7 +341,7 @@ class CredentialsStore(metaclass=SingletonMetaClass):
|
|
341
341
|
self.clear_token(pro_api_url)
|
342
342
|
|
343
343
|
def clear_all_pro_tokens(
|
344
|
-
self, pro_api_url: str
|
344
|
+
self, pro_api_url: Optional[str] = None
|
345
345
|
) -> List[ServerCredentials]:
|
346
346
|
"""Delete all tokens from the store for ZenML Pro servers connected to a given API server.
|
347
347
|
|
@@ -356,7 +356,7 @@ class CredentialsStore(metaclass=SingletonMetaClass):
|
|
356
356
|
if (
|
357
357
|
server.type == ServerType.PRO
|
358
358
|
and server.pro_api_url
|
359
|
-
and server.pro_api_url == pro_api_url
|
359
|
+
and (pro_api_url is None or server.pro_api_url == pro_api_url)
|
360
360
|
):
|
361
361
|
if server.api_key:
|
362
362
|
continue
|
@@ -556,7 +556,7 @@ class CredentialsStore(metaclass=SingletonMetaClass):
|
|
556
556
|
def update_server_info(
|
557
557
|
self,
|
558
558
|
server_url: str,
|
559
|
-
server_info: Union[ServerModel,
|
559
|
+
server_info: Union[ServerModel, WorkspaceRead],
|
560
560
|
) -> None:
|
561
561
|
"""Update the server information stored for a specific server URL.
|
562
562
|
|
zenml/login/pro/client.py
CHANGED
@@ -41,7 +41,7 @@ logger = get_logger(__name__)
|
|
41
41
|
|
42
42
|
if TYPE_CHECKING:
|
43
43
|
from zenml.login.pro.organization.client import OrganizationClient
|
44
|
-
from zenml.login.pro.
|
44
|
+
from zenml.login.pro.workspace.client import WorkspaceClient
|
45
45
|
|
46
46
|
# type alias for possible json payloads (the Anys are recursive Json instances)
|
47
47
|
Json = Union[Dict[str, Any], List[Any], str, int, float, bool, None]
|
@@ -56,7 +56,7 @@ class ZenMLProClient(metaclass=SingletonMetaClass):
|
|
56
56
|
_url: str
|
57
57
|
_api_token: APIToken
|
58
58
|
_session: Optional[requests.Session] = None
|
59
|
-
|
59
|
+
_workspace: Optional["WorkspaceClient"] = None
|
60
60
|
_organization: Optional["OrganizationClient"] = None
|
61
61
|
|
62
62
|
def __init__(self, url: str, api_token: Optional[APIToken] = None) -> None:
|
@@ -89,17 +89,17 @@ class ZenMLProClient(metaclass=SingletonMetaClass):
|
|
89
89
|
self._api_token = api_token
|
90
90
|
|
91
91
|
@property
|
92
|
-
def
|
93
|
-
"""Get the
|
92
|
+
def workspace(self) -> "WorkspaceClient":
|
93
|
+
"""Get the workspace client.
|
94
94
|
|
95
95
|
Returns:
|
96
|
-
The
|
96
|
+
The workspace client.
|
97
97
|
"""
|
98
|
-
if self.
|
99
|
-
from zenml.login.pro.
|
98
|
+
if self._workspace is None:
|
99
|
+
from zenml.login.pro.workspace.client import WorkspaceClient
|
100
100
|
|
101
|
-
self.
|
102
|
-
return self.
|
101
|
+
self._workspace = WorkspaceClient(client=self)
|
102
|
+
return self._workspace
|
103
103
|
|
104
104
|
@property
|
105
105
|
def organization(self) -> "OrganizationClient":
|
zenml/login/pro/utils.py
CHANGED
@@ -18,7 +18,7 @@ from zenml.login.credentials import ServerType
|
|
18
18
|
from zenml.login.credentials_store import get_credentials_store
|
19
19
|
from zenml.login.pro.client import ZenMLProClient
|
20
20
|
from zenml.login.pro.constants import ZENML_PRO_API_URL
|
21
|
-
from zenml.login.pro.
|
21
|
+
from zenml.login.pro.workspace.models import WorkspaceStatus
|
22
22
|
|
23
23
|
logger = get_logger(__name__)
|
24
24
|
|
@@ -44,13 +44,13 @@ def get_troubleshooting_instructions(url: str) -> str:
|
|
44
44
|
client = ZenMLProClient(pro_api_url)
|
45
45
|
|
46
46
|
try:
|
47
|
-
servers = client.
|
47
|
+
servers = client.workspace.list(url=url, member_only=False)
|
48
48
|
except Exception as e:
|
49
|
-
logger.debug(f"Failed to list
|
49
|
+
logger.debug(f"Failed to list workspaces: {e}")
|
50
50
|
else:
|
51
51
|
if servers:
|
52
52
|
server = servers[0]
|
53
|
-
if server.status ==
|
53
|
+
if server.status == WorkspaceStatus.AVAILABLE:
|
54
54
|
return (
|
55
55
|
f"The '{server.name}' ZenML Pro server that the client "
|
56
56
|
"is connected to is currently running but you may not "
|
@@ -58,9 +58,9 @@ def get_troubleshooting_instructions(url: str) -> str:
|
|
58
58
|
"contact your ZenML Pro administrator for more "
|
59
59
|
"information or try to manage the server members "
|
60
60
|
"yourself if you have the necessary permissions by "
|
61
|
-
f"visiting the ZenML Pro
|
61
|
+
f"visiting the ZenML Pro workspace page at {server.dashboard_url}."
|
62
62
|
)
|
63
|
-
if server.status ==
|
63
|
+
if server.status == WorkspaceStatus.DEACTIVATED:
|
64
64
|
return (
|
65
65
|
f"The '{server.name}' ZenML Pro server that the client "
|
66
66
|
"is connected to has been deactivated. "
|
@@ -69,14 +69,14 @@ def get_troubleshooting_instructions(url: str) -> str:
|
|
69
69
|
"you have the necessary permissions by visiting the "
|
70
70
|
f"ZenML Pro Organization page at {server.dashboard_organization_url}."
|
71
71
|
)
|
72
|
-
if server.status ==
|
72
|
+
if server.status == WorkspaceStatus.PENDING:
|
73
73
|
return (
|
74
74
|
f"The '{server.name}' ZenML Pro server that the client "
|
75
75
|
"is connected to is currently undergoing maintenance "
|
76
76
|
"(e.g. being deployed, upgraded or re-activated). "
|
77
77
|
"Please try again later or contact your ZenML Pro "
|
78
78
|
"administrator for more information. You can also "
|
79
|
-
f"visit the ZenML Pro
|
79
|
+
f"visit the ZenML Pro workspace page at {server.dashboard_url}."
|
80
80
|
)
|
81
81
|
return (
|
82
82
|
f"The '{server.name}' ZenML Pro server that the client "
|
@@ -11,80 +11,80 @@
|
|
11
11
|
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express
|
12
12
|
# or implied. See the License for the specific language governing
|
13
13
|
# permissions and limitations under the License.
|
14
|
-
"""ZenML Pro
|
14
|
+
"""ZenML Pro workspace client."""
|
15
15
|
|
16
16
|
from typing import List, Optional
|
17
17
|
from uuid import UUID
|
18
18
|
|
19
19
|
from zenml.logger import get_logger
|
20
20
|
from zenml.login.pro.client import ZenMLProClient
|
21
|
-
from zenml.login.pro.
|
21
|
+
from zenml.login.pro.workspace.models import WorkspaceRead, WorkspaceStatus
|
22
22
|
|
23
23
|
logger = get_logger(__name__)
|
24
24
|
|
25
|
-
|
25
|
+
WORKSPACES_ROUTE = "/workspaces"
|
26
26
|
|
27
27
|
|
28
|
-
class
|
29
|
-
"""
|
28
|
+
class WorkspaceClient:
|
29
|
+
"""Workspace management client."""
|
30
30
|
|
31
31
|
def __init__(
|
32
32
|
self,
|
33
33
|
client: ZenMLProClient,
|
34
34
|
):
|
35
|
-
"""Initialize the
|
35
|
+
"""Initialize the workspace client.
|
36
36
|
|
37
37
|
Args:
|
38
38
|
client: ZenML Pro client.
|
39
39
|
"""
|
40
40
|
self.client = client
|
41
41
|
|
42
|
-
def get(self, id: UUID) ->
|
43
|
-
"""Get a
|
42
|
+
def get(self, id: UUID) -> WorkspaceRead:
|
43
|
+
"""Get a workspace by id.
|
44
44
|
|
45
45
|
Args:
|
46
|
-
id: Id. of the
|
46
|
+
id: Id. of the workspace to retrieve.
|
47
47
|
|
48
48
|
Returns:
|
49
|
-
A
|
49
|
+
A workspace.
|
50
50
|
"""
|
51
51
|
return self.client._get_resource(
|
52
52
|
resource_id=id,
|
53
|
-
route=
|
54
|
-
response_model=
|
53
|
+
route=WORKSPACES_ROUTE,
|
54
|
+
response_model=WorkspaceRead,
|
55
55
|
)
|
56
56
|
|
57
57
|
def list(
|
58
58
|
self,
|
59
59
|
offset: int = 0,
|
60
60
|
limit: int = 20,
|
61
|
-
|
61
|
+
workspace_name: Optional[str] = None,
|
62
62
|
url: Optional[str] = None,
|
63
63
|
organization_id: Optional[UUID] = None,
|
64
|
-
status: Optional[
|
64
|
+
status: Optional[WorkspaceStatus] = None,
|
65
65
|
member_only: bool = False,
|
66
|
-
) -> List[
|
67
|
-
"""List
|
66
|
+
) -> List[WorkspaceRead]:
|
67
|
+
"""List workspaces.
|
68
68
|
|
69
69
|
Args:
|
70
70
|
offset: Offset to use for filtering.
|
71
71
|
limit: Limit used for filtering.
|
72
|
-
|
73
|
-
url:
|
72
|
+
workspace_name: Workspace name to filter by.
|
73
|
+
url: Workspace service URL to filter by.
|
74
74
|
organization_id: Organization ID to filter by.
|
75
|
-
status: Filter for only
|
76
|
-
member_only: If True, only list
|
77
|
-
(i.e. users that can connect to the
|
75
|
+
status: Filter for only workspaces with this status.
|
76
|
+
member_only: If True, only list workspaces where the user is a member
|
77
|
+
(i.e. users that can connect to the workspace).
|
78
78
|
|
79
79
|
Returns:
|
80
|
-
List of
|
80
|
+
List of workspaces.
|
81
81
|
"""
|
82
82
|
return self.client._list_resources(
|
83
|
-
route=
|
84
|
-
response_model=
|
83
|
+
route=WORKSPACES_ROUTE,
|
84
|
+
response_model=WorkspaceRead,
|
85
85
|
offset=offset,
|
86
86
|
limit=limit,
|
87
|
-
|
87
|
+
workspace_name=workspace_name,
|
88
88
|
url=url,
|
89
89
|
organization_id=organization_id,
|
90
90
|
status=status,
|