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.

Sign up for GitHub

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

Jump to bottom

Lazy loading improvements #3140

Draft
wants to merge 1 commit into
base: develop
Choose a base branch
from
Draft

Lazy loading improvements #3140

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
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
3 changes: 3 additions & 0 deletions src/zenml/client_lazy_loader.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -208,6 +208,9 @@ def _inner(*args: Any, **kwargs: Any) -> Any:
with contextlib.suppress(ValueError):
kwargs[k] = ClientLazyLoader(**v).evaluate()

# Why do we check for `isinstance(v, ClientLazyLoader)` for the
Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This works as expected and evaluates at runtime because intermediate it passed as an arg:

@pipeline(enable_cache=False, model=Model(name=MODEL_NAME))
def test_pipeline():
    intermediate = Client().get_model(MODEL_NAME).latest_version_id

    cll = (
        Client()
        .get_model_version(
            MODEL_NAME, intermediate
        )
        .get_artifact(ARTIFACT_NAME)
    )

If we pass it as a keyword argument instead, the intermediate lazy loader already gets evaluated at pipeline compilation time

@pipeline(enable_cache=False, model=Model(name=MODEL_NAME))
def test_pipeline():
    intermediate = Client().get_model(MODEL_NAME).latest_version_id

    cll = (
        Client()
        .get_model_version(
            MODEL_NAME, model_version_name_or_number_or_id=intermediate
        )
        .get_artifact(ARTIFACT_NAME)
    )

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

@avishniakov I tried to understand the reason for the evaluate_all_lazy_load_args_in_client_methods decorator, and the only thing I could come up with were these nested lazy loaders. Then I ran into this issue with the args/kwargs and was wondering what the expected behaviour would be in your opinion

# args but not for `kwargs`

return func(*args_, **kwargs)

return _inner
Expand Down
5 changes: 5 additions & 0 deletions src/zenml/orchestrators/input_utils.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -140,6 +140,11 @@ def resolve_step_inputs(
)
for name, cll_ in step.config.client_lazy_loaders.items():
value_ = cll_.evaluate()
# TODO: If we pass something as parameter here, the result is not stored
# anywhere. This means the user can not see what value was being passed
# into their step, and also not see where exactly this was coming from
# (e.g. which artifact/model version). Same thing applies to the model
# lazy loaders above
if isinstance(value_, ArtifactVersionResponse):
input_artifacts[name] = value_
elif isinstance(value_, RunMetadataResponse):
Expand Down
Loading