Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Model links lazy evaluation in pipeline code #2205

Merged
Show file tree
Hide file tree
Changes from 16 commits
Commits
Show all changes
25 commits
Select commit Hold shift + click to select a range
a9507df
MV data lazy loading in pipelines
avishniakov Dec 28, 2023
b0d288e
add tests
avishniakov Dec 28, 2023
b3250c4
new template ref
avishniakov Dec 28, 2023
f37c039
Merge branch 'develop' into feature/OSS-2713-model-links-lazy-evaluat…
avishniakov Dec 28, 2023
6b36164
use `BaseModel`
avishniakov Dec 28, 2023
58b092c
Auto-update of Starter template
actions-user Dec 28, 2023
820d80b
update to `v7` syntax
avishniakov Dec 28, 2023
78cea79
update test signatures
avishniakov Dec 28, 2023
41acb73
Auto-update of E2E template
actions-user Dec 28, 2023
a819949
update test signatures
avishniakov Dec 28, 2023
4d88dad
Merge branch 'develop' into feature/OSS-2713-model-links-lazy-evaluat…
avishniakov Jan 9, 2024
1044410
Merge branch 'develop' into feature/OSS-2713-model-links-lazy-evaluat…
avishniakov Jan 9, 2024
81ad088
wandb lint
avishniakov Jan 10, 2024
9fb0412
lint
avishniakov Jan 10, 2024
873a6e1
Merge branch 'develop' into feature/OSS-2713-model-links-lazy-evaluat…
avishniakov Jan 11, 2024
0f737f9
Merge branch 'develop' into feature/OSS-2713-model-links-lazy-evaluat…
avishniakov Jan 11, 2024
963d27c
remove leftover
avishniakov Jan 11, 2024
bc3da8b
Apply suggestions from code review
avishniakov Jan 11, 2024
95e6e46
Merge branch 'feature/OSS-2713-model-links-lazy-evaluation-in-pipelin…
avishniakov Jan 11, 2024
e4870dc
renaming
avishniakov Jan 11, 2024
551ef4e
model lazy load in `model`
avishniakov Jan 11, 2024
394977f
metadata lazy load in `metadata`
avishniakov Jan 11, 2024
7b6307e
Merge branch 'develop' into feature/OSS-2713-model-links-lazy-evaluat…
avishniakov Jan 11, 2024
919b14b
implement Michael's suggestions
avishniakov Jan 12, 2024
d098bd2
Merge branch 'develop' into feature/OSS-2713-model-links-lazy-evaluat…
avishniakov Jan 17, 2024
File filter

Filter by extension

Filter by extension


Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
8 changes: 4 additions & 4 deletions .github/workflows/update-templates-to-examples.yml
Original file line number Diff line number Diff line change
Expand Up @@ -50,7 +50,7 @@ jobs:
python-version: ${{ inputs.python-version }}
stack-name: local
ref-zenml: ${{ github.ref }}
ref-template: 2023.12.12 # Make sure it is aligned with ZENML_PROJECT_TEMPLATES from src/zenml/cli/base.py
ref-template: 2023.12.28 # Make sure it is aligned with ZENML_PROJECT_TEMPLATES from src/zenml/cli/base.py
- name: Clean-up
run: |
rm -rf ./local_checkout
Expand Down Expand Up @@ -95,7 +95,7 @@ jobs:
with:
github-token: ${{ secrets.GITHUB_TOKEN }}
script: |
github.issues.createComment({
github.rest.issues.createComment({
issue_number: ${{ github.event.pull_request.number }},
owner: context.repo.owner,
repo: context.repo.repo,
Expand Down Expand Up @@ -165,7 +165,7 @@ jobs:
with:
github-token: ${{ secrets.GITHUB_TOKEN }}
script: |
github.issues.createComment({
github.rest.issues.createComment({
issue_number: ${{ github.event.pull_request.number }},
owner: context.repo.owner,
repo: context.repo.repo,
Expand Down Expand Up @@ -237,7 +237,7 @@ jobs:
with:
github-token: ${{ secrets.GITHUB_TOKEN }}
script: |-
github.issues.createComment({
github.rest.issues.createComment({
issue_number: ${{ github.event.pull_request.number }},
owner: context.repo.owner,
repo: context.repo.repo,
Expand Down
2 changes: 1 addition & 1 deletion examples/e2e/.copier-answers.yml
Original file line number Diff line number Diff line change
@@ -1,5 +1,5 @@
# Changes here will be overwritten by Copier
_commit: 2023.12.12
_commit: 2023.12.12-2-g97eba5e
_src_path: gh:zenml-io/template-e2e-batch
data_quality_checks: true
email: ''
Expand Down
4 changes: 2 additions & 2 deletions examples/e2e/steps/deployment/deployment_deploy.py
Original file line number Diff line number Diff line change
Expand Up @@ -65,9 +65,9 @@ def deployment_deploy() -> (
# deploy predictor service
deployment_service = mlflow_model_registry_deployer_step.entrypoint(
registry_model_name=model_version.name,
registry_model_version=model_version.metadata[
registry_model_version=model_version.run_metadata[
"model_registry_version"
],
].value,
replace_existing=True,
)
else:
Expand Down
14 changes: 8 additions & 6 deletions examples/e2e/steps/promotion/promote_with_metric_compare.py
Original file line number Diff line number Diff line change
Expand Up @@ -90,17 +90,17 @@ def promote_with_metric_compare(
logger.info(f"Current model version was promoted to '{target_env}'.")

# Promote in Model Registry
latest_version_model_registry_number = latest_version.metadata[
latest_version_model_registry_number = latest_version.run_metadata[
"model_registry_version"
]
].value
if current_version_number is None:
current_version_model_registry_number = (
latest_version_model_registry_number
)
else:
current_version_model_registry_number = current_version.metadata[
"model_registry_version"
]
current_version_model_registry_number = (
current_version.run_metadata["model_registry_version"].value
)
promote_in_model_registry(
latest_version=latest_version_model_registry_number,
current_version=current_version_model_registry_number,
Expand All @@ -109,7 +109,9 @@ def promote_with_metric_compare(
)
promoted_version = latest_version_model_registry_number
else:
promoted_version = current_version.metadata["model_registry_version"]
promoted_version = current_version.run_metadata[
"model_registry_version"
].value

logger.info(
f"Current model version in `{target_env}` is `{promoted_version}` registered in Model Registry"
Expand Down
1 change: 1 addition & 0 deletions src/zenml/__init__.py
Original file line number Diff line number Diff line change
Expand Up @@ -46,6 +46,7 @@
from zenml.artifacts.artifact_config import ArtifactConfig
from zenml.artifacts.external_artifact import ExternalArtifact
from zenml.model.model_version import ModelVersion
from zenml.model.utils import log_model_version_metadata
from zenml.new.pipelines.pipeline_context import get_pipeline_context
from zenml.new.pipelines.pipeline_decorator import pipeline
from zenml.new.steps.step_decorator import step
Expand Down
2 changes: 1 addition & 1 deletion src/zenml/cli/base.py
Original file line number Diff line number Diff line change
Expand Up @@ -73,7 +73,7 @@ def copier_github_url(self) -> str:
ZENML_PROJECT_TEMPLATES = dict(
e2e_batch=ZenMLProjectTemplateLocation(
github_url="zenml-io/template-e2e-batch",
github_tag="2023.12.12", # Make sure it is aligned with .github/workflows/update-templates-to-examples.yml
github_tag="2023.12.28", # Make sure it is aligned with .github/workflows/update-templates-to-examples.yml
),
starter=ZenMLProjectTemplateLocation(
github_url="zenml-io/template-starter",
Expand Down
2 changes: 1 addition & 1 deletion src/zenml/cli/model.py
Original file line number Diff line number Diff line change
Expand Up @@ -60,7 +60,7 @@ def _model_version_to_print(
"number": model_version.number,
"description": model_version.description,
"stage": model_version.stage,
"metadata": model_version.to_model_version().metadata,
"metadata": model_version.to_model_version().run_metadata,
"tags": [t.name for t in model_version.tags],
"data_artifacts_count": len(model_version.data_artifact_ids),
"model_artifacts_count": len(model_version.model_artifact_ids),
Expand Down
25 changes: 24 additions & 1 deletion src/zenml/config/pipeline_configurations.py
Original file line number Diff line number Diff line change
Expand Up @@ -12,7 +12,8 @@
# or implied. See the License for the specific language governing
# permissions and limitations under the License.
"""Pipeline configuration classes."""
from typing import TYPE_CHECKING, Any, Dict, Optional
from contextlib import contextmanager
from typing import TYPE_CHECKING, Any, Dict, Iterator, List, Optional

from pydantic import validator

Expand Down Expand Up @@ -53,6 +54,8 @@ class PipelineConfiguration(PipelineConfigurationUpdate):

name: str

_runtime_state: List[bool] = [True]

@validator("name")
def ensure_pipeline_name_allowed(cls, name: str) -> str:
"""Ensures the pipeline name is allowed.
Expand Down Expand Up @@ -86,3 +89,23 @@ def docker_settings(self) -> "DockerSettings":
DOCKER_SETTINGS_KEY, {}
)
return DockerSettings.parse_obj(model_or_dict)

@contextmanager
def in_design_time(self) -> Iterator[None]:
"""Context manager to mark a pipeline as being in design time.

Yields:
Nothing.
"""
self._runtime_state[0] = False
yield
self._runtime_state[0] = True

@property
def is_runtime(self) -> bool:
"""Returns whether the pipeline is in runtime.

Returns:
Whether the pipeline is in runtime.
"""
return self._runtime_state[0]
2 changes: 2 additions & 0 deletions src/zenml/config/step_configurations.py
Original file line number Diff line number Diff line change
Expand Up @@ -33,6 +33,7 @@
from zenml.config.constants import DOCKER_SETTINGS_KEY, RESOURCE_SETTINGS_KEY
from zenml.config.source import Source, convert_source_validator
from zenml.config.strict_base_model import StrictBaseModel
from zenml.lazy_load.model_version import ModelVersionDataLazyLoader
from zenml.logger import get_logger
from zenml.model.model_version import ModelVersion
from zenml.utils import deprecation_utils
Expand Down Expand Up @@ -152,6 +153,7 @@ class PartialStepConfiguration(StepConfigurationUpdate):
name: str
caching_parameters: Mapping[str, Any] = {}
external_input_artifacts: Mapping[str, ExternalArtifactConfiguration] = {}
model_artifacts_or_metadata: Mapping[str, ModelVersionDataLazyLoader] = {}
outputs: Mapping[str, PartialArtifactConfiguration] = {}

# Override the deprecation validator as we do not want to deprecate the
Expand Down
34 changes: 34 additions & 0 deletions src/zenml/lazy_load/model_version.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,34 @@
# Copyright (c) ZenML GmbH 2023. 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.
"""Model Version Data Lazy Loader definition."""

from typing import Optional

from pydantic import BaseModel

from zenml.model.model_version import ModelVersion


class ModelVersionDataLazyLoader(BaseModel):
"""Model Version Data Lazy Loader helper class.

It helps the inner codes to fetch proper artifact,
model version metadata or artifact metadata from the
model version during runtime time of the step.
"""

model_version: ModelVersion
artifact_name: Optional[str] = None
artifact_version: Optional[str] = None
metadata_name: Optional[str] = None
66 changes: 66 additions & 0 deletions src/zenml/lazy_load/run_metadata.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,66 @@
# Copyright (c) ZenML GmbH 2023. 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.
"""Run Metadata Lazy Loader definition."""

from typing import TYPE_CHECKING, Optional
from uuid import uuid4

if TYPE_CHECKING:
from zenml.model.model_version import ModelVersion
from zenml.models import RunMetadataResponse


class RunMetadataLazyGetter:
"""Run Metadata Lazy Getter helper class.

It serves the purpose to feed back to the user the metadata
lazy loader wrapper for any given key, if called inside a pipeline
design time context.
"""

def __init__(
self,
_lazy_load_model_version: "ModelVersion",
_lazy_load_artifact_name: Optional[str],
_lazy_load_artifact_version: Optional[str],
):
"""Initialize a RunMetadataLazyGetter.

Args:
_lazy_load_model_version: The model version.
_lazy_load_artifact_name: The artifact name.
_lazy_load_artifact_version: The artifact version.
"""
self._lazy_load_model_version = _lazy_load_model_version
self._lazy_load_artifact_name = _lazy_load_artifact_name
self._lazy_load_artifact_version = _lazy_load_artifact_version

def __getitem__(self, key: str) -> "RunMetadataResponse":
"""Get the metadata for the given key.

Args:
key: The metadata key.

Returns:
The metadata lazy loader wrapper for the given key.
"""
from zenml.models import RunMetadataResponse

return RunMetadataResponse(
id=uuid4(),
_lazy_load_model_version=self._lazy_load_model_version,
_lazy_load_artifact_name=self._lazy_load_artifact_name,
_lazy_load_artifact_version=self._lazy_load_artifact_version,
_lazy_load_metadata_name=key,
)
5 changes: 4 additions & 1 deletion src/zenml/metadata/metadata_types.py
Original file line number Diff line number Diff line change
Expand Up @@ -13,10 +13,13 @@
# permissions and limitations under the License.
"""Custom types that can be used as metadata of ZenML artifacts."""

from typing import Any, Dict, List, Set, Tuple, Union
from typing import TYPE_CHECKING, Any, Dict, List, Set, Tuple, Union

from zenml.utils.enum_utils import StrEnum

if TYPE_CHECKING:
pass


class Uri(str):
"""Special string class to indicate a URI."""
Expand Down
Loading
Loading