[Model WatchTower] Extend ExternalArtifact

[avishniakov]

Describe changes

Major change:

• ExternalArtifact extended to receive:
  • model_name: Name of a model to search for artifact in (if None - derived from step context).
  • model_version: Version of a model to search for artifact in (if None - derived from step context).
  • model_artifact_name: Name of a model artifact to search for.
  • model_artifact_version: Version of a model artifact to search for.
• Since model config can be also defined on a step level, we cannot resolve the actual artifact ID as before and this part was moved into the execution stage (specifically _prepare of a step). This should also fix https://zenml.atlassian.net/browse/OSS-2452, but this was not tested.

Sample flow can be like:

from typing_extensions import Annotated
from zenml import pipeline, step
from zenml.model import ArtifactConfig, ModelConfig
from zenml.steps.external_artifact import ExternalArtifact


@step
def producer() -> (
    Annotated[int, "predictions", ArtifactConfig(overwrite=False)]
):
    return 1


@step
def consumer(external_artifact: int):
    print("Hello from model artifact:", external_artifact)


@pipeline(
    enable_cache=False,
    model_config=ModelConfig(name="foo", create_new_model_version=True),
)
def producer_pipeline():
    producer()


@pipeline(enable_cache=False, model_config=ModelConfig(name="foo"))
def consumer_pipeline():
    consumer(ExternalArtifact(model_artifact_name="predictions"))


if __name__ == "__main__":
    producer_pipeline()
    consumer_pipeline()

Aside, this one includes a somewhat heavy rework of how we work with Artifacts of any kind inside ModelVersionResponseModel class:

• New keys for artifact dicts look like pipeline_name::step_name::artifact_name. This is done to avoid collisions then 2 steps write the same named artifact to the model version (e.g. not typed output)
• Generic _get_linked_object to retrieve from those collections by full or partial key (user must feed in name, but pipeline_name and step_name remains optional). Next, we construct a lookup key and in case only one or no elements are found - the user will get a response or None, but if there is a collision a RuntimeError is raised.

Pre-requisites

Please ensure you have done the following:

• I have read the CONTRIBUTING.md document.
• If my change requires a change to docs, I have updated the documentation accordingly.
• If I have added an integration, I have updated the integrations table and the corresponding website section.
• I have added tests to cover my changes.
• I have based my new branch on develop and the open PR is targeting develop. If your branch wasn't based on develop read Contribution guide on rebasing branch to develop.

Types of changes

• Bug fix (non-breaking change which fixes an issue)
• New feature (non-breaking change which adds functionality)
• Breaking change (fix or feature that would cause existing functionality to change)
• Other (add details above)

[fa9r]

Looks good aside from the comment you left - I'll look into this as well now, let's see if we can solve it differently

[avishniakov]

There were problems JSONifying complex objects passed as input to ExternalArtifact by value (identified by E2E example), so I broke it down into config and user-facing.
If the user picks the value option - we upload as before during compilation, then the value is substituted by id.
Otherwise - no actions until runtime. In step, run configuration will get resolved into an Artifact object on the fly.

[avishniakov]

Merging this to unblock myself, all fails were caused by the template - this is in progress, and more changes to the template to support ModelWatchTower will follow.