zenml-nightly 0.55.0.dev20240124__py3-none-any.whl → 0.72.0.dev20250116__py3-none-any.whl

zenml/artifacts/utils.py

@@ -14,24 +14,45 @@
 """Utility functions for handling artifacts."""
 
 import base64
+import contextlib
 import os
 import tempfile
-from typing import TYPE_CHECKING, Any, Dict, List, Optional, Type, Union, cast
-from uuid import UUID
+import zipfile
+from pathlib import Path
+from typing import (
+    TYPE_CHECKING,
+    Any,
+    Dict,
+    List,
+    Optional,
+    Type,
+    Union,
+    cast,
+)
+from uuid import UUID, uuid4
 
+from zenml.artifacts.preexisting_data_materializer import (
+    PreexistingDataMaterializer,
+)
 from zenml.client import Client
-from zenml.constants import MODEL_METADATA_YAML_FILE_NAME
+from zenml.constants import (
+    MODEL_METADATA_YAML_FILE_NAME,
+)
 from zenml.enums import (
+    ArtifactSaveType,
+    ArtifactType,
     ExecutionStatus,
-    MetadataResourceTypes,
     StackComponentType,
     VisualizationType,
 )
-from zenml.exceptions import DoesNotExistException, StepContextError
+from zenml.exceptions import (
+    DoesNotExistException,
+    StepContextError,
+)
 from zenml.io import fileio
 from zenml.logger import get_logger
+from zenml.metadata.metadata_types import validate_metadata
 from zenml.models import (
-    ArtifactRequest,
     ArtifactVersionRequest,
     ArtifactVersionResponse,
     ArtifactVisualizationRequest,
@@ -40,8 +61,8 @@ from zenml.models import (
     StepRunResponse,
     StepRunUpdate,
 )
-from zenml.new.steps.step_context import get_step_context
 from zenml.stack import StackComponent
+from zenml.steps.step_context import get_step_context
 from zenml.utils import source_utils
 from zenml.utils.yaml_utils import read_yaml, write_yaml
 
@@ -61,20 +82,134 @@ logger = get_logger(__name__)
 # ----------
 
 
+def _save_artifact_visualizations(
+    data: Any, materializer: "BaseMaterializer"
+) -> List[ArtifactVisualizationRequest]:
+    """Save artifact visualizations.
+
+    Args:
+        data: The data for which to save the visualizations.
+        materializer: The materializer that should be used to generate and
+            save the visualizations.
+
+    Returns:
+        List of requests for the saved visualizations.
+    """
+    try:
+        visualizations = materializer.save_visualizations(data)
+    except Exception as e:
+        logger.warning("Failed to save artifact visualizations: %s", e)
+        return []
+
+    return [
+        ArtifactVisualizationRequest(
+            type=type,
+            uri=uri,
+        )
+        for uri, type in visualizations.items()
+    ]
+
+
+def _store_artifact_data_and_prepare_request(
+    data: Any,
+    name: str,
+    uri: str,
+    materializer_class: Type["BaseMaterializer"],
+    save_type: ArtifactSaveType,
+    version: Optional[Union[int, str]] = None,
+    artifact_type: Optional[ArtifactType] = None,
+    tags: Optional[List[str]] = None,
+    store_metadata: bool = True,
+    store_visualizations: bool = True,
+    has_custom_name: bool = True,
+    metadata: Optional[Dict[str, "MetadataType"]] = None,
+) -> ArtifactVersionRequest:
+    """Store artifact data and prepare a request to the server.
+
+    Args:
+        data: The artifact data.
+        name: The artifact name.
+        uri: The artifact URI.
+        materializer_class: The materializer class to use for storing the
+            artifact data.
+        save_type: Save type of the artifact version.
+        version: The artifact version.
+        artifact_type: The artifact type. If not given, the type will be defined
+            by the materializer that is used to save the artifact.
+        tags: Tags for the artifact version.
+        store_metadata: Whether to store metadata for the artifact version.
+        store_visualizations: Whether to store visualizations for the artifact
+            version.
+        has_custom_name: Whether the artifact has a custom name.
+        metadata: Metadata to store for the artifact version. This will be
+            ignored if `store_metadata` is set to `False`.
+
+    Returns:
+        Artifact version request for the artifact data that was stored.
+    """
+    artifact_store = Client().active_stack.artifact_store
+    artifact_store.makedirs(uri)
+
+    materializer = materializer_class(uri=uri, artifact_store=artifact_store)
+    materializer.uri = materializer.uri.replace("\\", "/")
+
+    data_type = type(data)
+    materializer.validate_save_type_compatibility(data_type)
+    materializer.save(data)
+
+    visualizations = (
+        _save_artifact_visualizations(data=data, materializer=materializer)
+        if store_visualizations
+        else None
+    )
+
+    combined_metadata: Dict[str, "MetadataType"] = {}
+    if store_metadata:
+        try:
+            combined_metadata = materializer.extract_full_metadata(data)
+        except Exception as e:
+            logger.warning("Failed to extract materializer metadata: %s", e)
+
+        # Update with user metadata to potentially overwrite values coming from
+        # the materializer
+        combined_metadata.update(metadata or {})
+
+    artifact_version_request = ArtifactVersionRequest(
+        artifact_name=name,
+        version=version,
+        tags=tags,
+        type=artifact_type or materializer.ASSOCIATED_ARTIFACT_TYPE,
+        uri=materializer.uri,
+        materializer=source_utils.resolve(materializer.__class__),
+        data_type=source_utils.resolve(data_type),
+        user=Client().active_user.id,
+        workspace=Client().active_workspace.id,
+        artifact_store_id=artifact_store.id,
+        visualizations=visualizations,
+        has_custom_name=has_custom_name,
+        save_type=save_type,
+        metadata=validate_metadata(combined_metadata)
+        if combined_metadata
+        else None,
+    )
+
+    return artifact_version_request
+
+
 def save_artifact(
     data: Any,
     name: str,
     version: Optional[Union[int, str]] = None,
+    artifact_type: Optional[ArtifactType] = None,
     tags: Optional[List[str]] = None,
     extract_metadata: bool = True,
     include_visualizations: bool = True,
-    has_custom_name: bool = True,
     user_metadata: Optional[Dict[str, "MetadataType"]] = None,
     materializer: Optional["MaterializerClassOrSource"] = None,
     uri: Optional[str] = None,
-    is_model_artifact: bool = False,
-    is_deployment_artifact: bool = False,
-    manual_save: bool = True,
+    # TODO: remove these once external artifact does not use this function anymore
+    save_type: ArtifactSaveType = ArtifactSaveType.MANUAL,
+    has_custom_name: bool = True,
 ) -> "ArtifactVersionResponse":
     """Upload and publish an artifact.
 
@@ -84,62 +219,45 @@ def save_artifact(
         version: The version of the artifact. If not provided, a new
             auto-incremented version will be used.
         tags: Tags to associate with the artifact.
+        artifact_type: The artifact type. If not given, the type will be defined
+            by the materializer that is used to save the artifact.
         extract_metadata: If artifact metadata should be extracted and returned.
         include_visualizations: If artifact visualizations should be generated.
-        has_custom_name: If the artifact name is custom and should be listed in
-            the dashboard "Artifacts" tab.
         user_metadata: User-provided metadata to store with the artifact.
         materializer: The materializer to use for saving the artifact to the
             artifact store.
         uri: The URI within the artifact store to upload the artifact
             to. If not provided, the artifact will be uploaded to
             `custom_artifacts/{name}/{version}`.
-        is_model_artifact: If the artifact is a model artifact.
-        is_deployment_artifact: If the artifact is a deployment artifact.
-        manual_save: If this function is called manually and should therefore
-            link the artifact to the current step run.
+        save_type: The type of save operation that created the artifact version.
+        has_custom_name: If the artifact name is custom and should be listed in
+            the dashboard "Artifacts" tab.
 
     Returns:
         The saved artifact response.
-
-    Raises:
-        RuntimeError: If artifact URI already exists.
     """
     from zenml.materializers.materializer_registry import (
         materializer_registry,
     )
     from zenml.utils import source_utils
 
-    # TODO: Can we handle this server side? If we leave it empty in the request,
-    # it's an auto-increase?
-    # TODO: This can probably lead to issues when multiple steps request a new
-    # artifact version at the same time?
-    # Get new artifact version if not specified
-    version = version or _get_new_artifact_version(name)
-
-    # Get the current artifact store
     client = Client()
     artifact_store = client.active_stack.artifact_store
 
-    # Build and check the artifact URI
     if not uri:
-        uri = os.path.join("custom_artifacts", name, str(version))
+        uri = os.path.join("custom_artifacts", name, str(uuid4()))
     if not uri.startswith(artifact_store.path):
         uri = os.path.join(artifact_store.path, uri)
 
-    if manual_save and fileio.exists(uri):
+    if save_type == ArtifactSaveType.MANUAL:
         # This check is only necessary for manual saves as we already check
         # it when creating the directory for step output artifacts
-        other_artifacts = client.list_artifact_versions(uri=uri, size=1)
-        if other_artifacts and (other_artifact := other_artifacts[0]):
-            raise RuntimeError(
-                f"Cannot save artifact {name} (version {version}) to URI "
-                f"{uri} because the URI is already used by artifact "
-                f"{other_artifact.name} (version {other_artifact.version})."
-            )
-    fileio.makedirs(uri)
+        _check_if_artifact_with_given_uri_already_registered(
+            artifact_store=artifact_store,
+            uri=uri,
+            name=name,
+        )
 
-    # Find and initialize the right materializer class
     if isinstance(materializer, type):
         materializer_class = materializer
     elif materializer:
@@ -148,106 +266,107 @@ def save_artifact(
         )
     else:
         materializer_class = materializer_registry[type(data)]
-    materializer_object = materializer_class(uri)
 
-    # Force URIs to have forward slashes
-    materializer_object.uri = materializer_object.uri.replace("\\", "/")
+    artifact_version_request = _store_artifact_data_and_prepare_request(
+        data=data,
+        name=name,
+        uri=uri,
+        materializer_class=materializer_class,
+        save_type=save_type,
+        version=version,
+        artifact_type=artifact_type,
+        tags=tags,
+        store_metadata=extract_metadata,
+        store_visualizations=include_visualizations,
+        has_custom_name=has_custom_name,
+        metadata=user_metadata,
+    )
+    artifact_version = client.zen_store.create_artifact_version(
+        artifact_version=artifact_version_request
+    )
 
-    # Save the artifact to the artifact store
-    data_type = type(data)
-    materializer_object.validate_type_compatibility(data_type)
-    materializer_object.save(data)
+    if save_type == ArtifactSaveType.MANUAL:
+        _link_artifact_version_to_the_step_and_model(
+            artifact_version=artifact_version,
+        )
 
-    # Save visualizations of the artifact
-    visualizations: List[ArtifactVisualizationRequest] = []
-    if include_visualizations:
-        try:
-            vis_data = materializer_object.save_visualizations(data)
-            for vis_uri, vis_type in vis_data.items():
-                vis_model = ArtifactVisualizationRequest(
-                    type=vis_type,
-                    uri=vis_uri,
-                )
-                visualizations.append(vis_model)
-        except Exception as e:
-            logger.warning(
-                f"Failed to save visualization for output artifact '{name}': "
-                f"{e}"
-            )
+    return artifact_version
 
-    # Save metadata of the artifact
-    artifact_metadata: Dict[str, "MetadataType"] = {}
-    if extract_metadata:
-        try:
-            artifact_metadata = materializer_object.extract_full_metadata(data)
-            artifact_metadata.update(user_metadata or {})
-        except Exception as e:
-            logger.warning(
-                f"Failed to extract metadata for output artifact '{name}': {e}"
-            )
 
-    # Get or create the artifact
-    try:
-        artifact = client.list_artifacts(name=name)[0]
-    except IndexError:
-        artifact = client.zen_store.create_artifact(
-            ArtifactRequest(
-                name=name,
-                has_custom_name=has_custom_name,
-            )
+def register_artifact(
+    folder_or_file_uri: str,
+    name: str,
+    version: Optional[Union[int, str]] = None,
+    artifact_type: Optional[ArtifactType] = None,
+    tags: Optional[List[str]] = None,
+    has_custom_name: bool = True,
+    artifact_metadata: Dict[str, "MetadataType"] = {},
+) -> "ArtifactVersionResponse":
+    """Register existing data stored in the artifact store as a ZenML Artifact.
+
+    Args:
+        folder_or_file_uri: The full URI within the artifact store to the folder
+            or to the file.
+        name: The name of the artifact.
+        version: The version of the artifact. If not provided, a new
+            auto-incremented version will be used.
+        artifact_type: The artifact type. If not given, the type will default
+            to `data`.
+        tags: Tags to associate with the artifact.
+        has_custom_name: If the artifact name is custom and should be listed in
+            the dashboard "Artifacts" tab.
+        artifact_metadata: Metadata dictionary to attach to the artifact version.
+
+    Returns:
+        The saved artifact response.
+
+    Raises:
+        FileNotFoundError: If the folder URI is outside of the artifact store
+            bounds.
+    """
+    client = Client()
+
+    # Get the current artifact store
+    artifact_store = client.active_stack.artifact_store
+
+    if not folder_or_file_uri.startswith(artifact_store.path):
+        raise FileNotFoundError(
+            f"Folder `{folder_or_file_uri}` is outside of "
+            f"artifact store bounds `{artifact_store.path}`"
         )
 
-    # Create the artifact version
-    artifact_version = ArtifactVersionRequest(
-        artifact_id=artifact.id,
+    _check_if_artifact_with_given_uri_already_registered(
+        artifact_store=artifact_store,
+        uri=folder_or_file_uri,
+        name=name,
+    )
+
+    artifact_version_request = ArtifactVersionRequest(
+        artifact_name=name,
         version=version,
         tags=tags,
-        type=materializer_object.ASSOCIATED_ARTIFACT_TYPE,
-        uri=materializer_object.uri,
-        materializer=source_utils.resolve(materializer_object.__class__),
-        data_type=source_utils.resolve(data_type),
+        type=artifact_type or ArtifactType.DATA,
+        save_type=ArtifactSaveType.PREEXISTING,
+        uri=folder_or_file_uri,
+        materializer=source_utils.resolve(PreexistingDataMaterializer),
+        data_type=source_utils.resolve(Path),
         user=Client().active_user.id,
         workspace=Client().active_workspace.id,
         artifact_store_id=artifact_store.id,
-        visualizations=visualizations,
         has_custom_name=has_custom_name,
+        metadata=validate_metadata(artifact_metadata)
+        if artifact_metadata
+        else None,
     )
-    response = Client().zen_store.create_artifact_version(
-        artifact_version=artifact_version
+    artifact_version = client.zen_store.create_artifact_version(
+        artifact_version=artifact_version_request
     )
-    if artifact_metadata:
-        Client().create_run_metadata(
-            metadata=artifact_metadata,
-            resource_id=response.id,
-            resource_type=MetadataResourceTypes.ARTIFACT_VERSION,
-        )
 
-    if manual_save:
-        try:
-            error_message = "step run"
-            step_context = get_step_context()
-            step_run = step_context.step_run
-            client.zen_store.update_run_step(
-                step_run_id=step_run.id,
-                step_run_update=StepRunUpdate(
-                    saved_artifact_versions={name: response.id}
-                ),
-            )
-            error_message = "model"
-            model = step_context.model
-            if model:
-                from zenml.model.utils import link_artifact_to_model
-
-                link_artifact_to_model(
-                    artifact_version_id=response.id,
-                    model=model,
-                    is_model_artifact=is_model_artifact,
-                    is_deployment_artifact=is_deployment_artifact,
-                )
-        except (RuntimeError, StepContextError):
-            logger.debug(f"Unable to link saved artifact to {error_message}.")
+    _link_artifact_version_to_the_step_and_model(
+        artifact_version=artifact_version,
+    )
 
-    return response
+    return artifact_version
 
 
 def load_artifact(
@@ -265,17 +384,6 @@ def load_artifact(
         The loaded artifact.
     """
     artifact = Client().get_artifact_version(name_or_id, version)
-    try:
-        step_run = get_step_context().step_run
-        client = Client()
-        client.zen_store.update_run_step(
-            step_run_id=step_run.id,
-            step_run_update=StepRunUpdate(
-                loaded_artifact_versions={artifact.name: artifact.id}
-            ),
-        )
-    except RuntimeError:
-        pass  # Cannot link to step run if called outside of a step
     return load_artifact_from_response(artifact)
 
 
@@ -296,46 +404,73 @@ def log_artifact_metadata(
         artifact_version: The version of the artifact to log metadata for. If
             not provided, when being called inside a step that produces an
             artifact named `artifact_name`, the metadata will be associated to
-            the corresponding newly created artifact. Or, if not provided when
-            being called outside of a step, or in a step that does not produce
-            any artifact named `artifact_name`, the metadata will be associated
-            to the latest version of that artifact.
+            the corresponding newly created artifact.
 
     Raises:
         ValueError: If no artifact name is provided and the function is not
             called inside a step with a single output, or, if neither an
             artifact nor an output with the given name exists.
+
     """
-    try:
+    logger.warning(
+        "The `log_artifact_metadata` function is deprecated and will soon be "
+        "removed. Instead, you can consider using: "
+        "`log_metadata(metadata={...}, infer_artifact=True, ...)` instead. For more "
+        "info: https://docs.zenml.io/how-to/model-management-metrics/track-metrics-metadata/attach-metadata-to-an-artifact"
+    )
+
+    from zenml import log_metadata
+
+    if artifact_name and artifact_version:
+        assert artifact_name is not None
+
+        log_metadata(
+            metadata=metadata,
+            artifact_name=artifact_name,
+            artifact_version=artifact_version,
+        )
+
+    step_context = None
+    with contextlib.suppress(RuntimeError):
         step_context = get_step_context()
-        in_step_outputs = (artifact_name in step_context._outputs) or (
-            not artifact_name and len(step_context._outputs) == 1
+
+    if step_context and artifact_name in step_context._outputs.keys():
+        log_metadata(
+            metadata=metadata,
+            artifact_name=artifact_name,
+            infer_artifact=True,
         )
-    except RuntimeError:
-        step_context = None
-        in_step_outputs = False
-
-    if not step_context or not in_step_outputs or artifact_version:
-        if not artifact_name:
-            raise ValueError(
-                "Artifact name must be provided unless the function is called "
-                "inside a step with a single output."
-            )
+    elif step_context and len(step_context._outputs) == 1:
+        single_output_name = list(step_context._outputs.keys())[0]
+
+        log_metadata(
+            metadata=metadata,
+            artifact_name=single_output_name,
+            infer_artifact=True,
+        )
+    elif artifact_name:
         client = Client()
-        response = client.get_artifact_version(artifact_name, artifact_version)
-        client.create_run_metadata(
+        logger.warning(
+            "Deprecation warning! Currently, you are calling "
+            "`log_artifact_metadata` from a context, where we use the "
+            "`artifact_name` to fetch it and link the metadata to its "
+            "latest version. This behavior is deprecated and will be "
+            "removed in the future. To circumvent this, please check"
+            "the `log_metadata` function."
+        )
+        artifact_version_model = client.get_artifact_version(
+            name_id_or_prefix=artifact_name
+        )
+        log_metadata(
             metadata=metadata,
-            resource_id=response.id,
-            resource_type=MetadataResourceTypes.ARTIFACT_VERSION,
+            artifact_version_id=artifact_version_model.id,
         )
-
     else:
-        try:
-            step_context.add_output_metadata(
-                metadata=metadata, output_name=artifact_name
-            )
-        except StepContextError as e:
-            raise ValueError(e)
+        raise ValueError(
+            "You need to call `log_artifact_metadata` either within a step "
+            "(potentially with an artifact name) or outside of a step with an "
+            "artifact name (and/or version)."
+        )
 
 
 # -----------------
@@ -387,18 +522,21 @@ def load_artifact_visualization(
     artifact_store = _load_artifact_store(
         artifact_store_id=artifact.artifact_store_id, zen_store=zen_store
     )
-    mode = "rb" if visualization.type == VisualizationType.IMAGE else "r"
-    value = _load_file_from_artifact_store(
-        uri=visualization.uri,
-        artifact_store=artifact_store,
-        mode=mode,
-    )
+    try:
+        mode = "rb" if visualization.type == VisualizationType.IMAGE else "r"
+        value = _load_file_from_artifact_store(
+            uri=visualization.uri,
+            artifact_store=artifact_store,
+            mode=mode,
+        )
 
-    # Encode image visualizations if requested
-    if visualization.type == VisualizationType.IMAGE and encode_image:
-        value = base64.b64encode(bytes(value))
+        # Encode image visualizations if requested
+        if visualization.type == VisualizationType.IMAGE and encode_image:
+            value = base64.b64encode(bytes(value))
 
-    return LoadedVisualization(type=visualization.type, value=value)
+        return LoadedVisualization(type=visualization.type, value=value)
+    finally:
+        artifact_store.cleanup()
 
 
 def load_artifact_from_response(artifact: "ArtifactVersionResponse") -> Any:
@@ -410,32 +548,73 @@ def load_artifact_from_response(artifact: "ArtifactVersionResponse") -> Any:
     Returns:
         The artifact loaded into memory.
     """
-    artifact_store_loaded = False
-    if artifact.artifact_store_id:
-        try:
-            artifact_store_model = Client().get_stack_component(
-                component_type=StackComponentType.ARTIFACT_STORE,
-                name_id_or_prefix=artifact.artifact_store_id,
-            )
-            _ = StackComponent.from_model(artifact_store_model)
-            artifact_store_loaded = True
-        except (KeyError, ImportError):
-            pass
-
-    if not artifact_store_loaded:
-        logger.warning(
-            "Unable to restore artifact store while trying to load artifact "
-            "`%s`. If this artifact is stored in a remote artifact store, "
-            "this might lead to issues when trying to load the artifact.",
-            artifact.id,
-        )
+    artifact_store = _get_artifact_store_from_response_or_from_active_stack(
+        artifact=artifact
+    )
 
     return _load_artifact_from_uri(
         materializer=artifact.materializer,
         data_type=artifact.data_type,
         uri=artifact.uri,
+        artifact_store=artifact_store,
+    )
+
+
+def download_artifact_files_from_response(
+    artifact: "ArtifactVersionResponse",
+    path: str,
+    overwrite: bool = False,
+) -> None:
+    """Download the given artifact into a file.
+
+    Args:
+        artifact: The artifact to download.
+        path: The path to which to download the artifact.
+        overwrite: Whether to overwrite the file if it already exists.
+
+    Raises:
+        FileExistsError: If the file already exists and `overwrite` is `False`.
+        Exception: If the artifact could not be downloaded to the zip file.
+    """
+    if not overwrite and fileio.exists(path):
+        raise FileExistsError(
+            f"File '{path}' already exists and `overwrite` is set to `False`."
+        )
+
+    artifact_store = _get_artifact_store_from_response_or_from_active_stack(
+        artifact=artifact
     )
 
+    if filepaths := artifact_store.listdir(artifact.uri):
+        # save a zipfile to 'path' containing all the files
+        # in 'filepaths' with compression
+        try:
+            with zipfile.ZipFile(path, "w", zipfile.ZIP_DEFLATED) as zipf:
+                for file in filepaths:
+                    # Ensure 'file' is a string for path operations
+                    # and ZIP entry naming
+                    file_str = (
+                        file.decode() if isinstance(file, bytes) else file
+                    )
+                    file_path = str(Path(artifact.uri) / file_str)
+                    with artifact_store.open(
+                        file_path, mode="rb"
+                    ) as store_file:
+                        # Use a loop to read and write chunks of the file
+                        # instead of reading the entire file into memory
+                        CHUNK_SIZE = 8192
+                        while True:
+                            if file_content := store_file.read(CHUNK_SIZE):
+                                zipf.writestr(file_str, file_content)
+                            else:
+                                break
+        except Exception as e:
+            logger.error(
+                f"Failed to save artifact '{artifact.id}' to zip file "
+                f" '{path}': {e}"
+            )
+            raise
+
 
 def get_producer_step_of_artifact(
     artifact: "ArtifactVersionResponse",
@@ -475,7 +654,8 @@ def get_artifacts_versions_of_pipeline_run(
     artifact_versions: List["ArtifactVersionResponse"] = []
     for step in pipeline_run.steps.values():
         if not only_produced or step.status == ExecutionStatus.COMPLETED:
-            artifact_versions.extend(step.outputs.values())
+            for output in step.outputs.values():
+                artifact_versions.extend(output)
     return artifact_versions
 
 
@@ -484,10 +664,81 @@ def get_artifacts_versions_of_pipeline_run(
 # -------------------------
 
 
+def _check_if_artifact_with_given_uri_already_registered(
+    artifact_store: "BaseArtifactStore",
+    uri: str,
+    name: str,
+) -> None:
+    """Check if the given artifact store already contains an artifact with the given URI.
+
+    Args:
+        artifact_store: The artifact store to check.
+        uri: The uri of the artifact.
+        name: The name of the artifact.
+
+    Raises:
+        RuntimeError: If the artifact store already contains an artifact with
+            the given URI.
+    """
+    if artifact_store.exists(uri):
+        # This check is only necessary for manual saves as we already check
+        # it when creating the directory for step output artifacts
+        other_artifacts = Client().list_artifact_versions(uri=uri, size=1)
+        if other_artifacts and (other_artifact := other_artifacts[0]):
+            raise RuntimeError(
+                f"Cannot create new artifact {name} version with URI "
+                f"{uri} because the URI is already used by artifact "
+                f"{other_artifact.name} (version {other_artifact.version})."
+            )
+
+
+def _link_artifact_version_to_the_step_and_model(
+    artifact_version: ArtifactVersionResponse,
+) -> None:
+    """Link an artifact version to the step and its' context model.
+
+    This function links the AV to:
+        - the step run
+        - the MV from the step context
+
+    Args:
+        artifact_version: The artifact version to link.
+    """
+    client = Client()
+    try:
+        error_message = "step run"
+        step_context = get_step_context()
+        step_run = step_context.step_run
+        client.zen_store.update_run_step(
+            step_run_id=step_run.id,
+            step_run_update=StepRunUpdate(
+                outputs={
+                    artifact_version.artifact.name: [
+                        artifact_version.id,
+                    ]
+                }
+            ),
+        )
+        error_message = "model"
+
+        if step_context.model_version:
+            from zenml.model.utils import (
+                link_artifact_version_to_model_version,
+            )
+
+            link_artifact_version_to_model_version(
+                artifact_version=artifact_version,
+                model_version=step_context.model_version,
+            )
+    except (RuntimeError, StepContextError):
+        logger.debug(f"Unable to link saved artifact to {error_message}.")
+
+
 def _load_artifact_from_uri(
     materializer: Union["Source", str],
     data_type: Union["Source", str],
     uri: str,
+    artifact_store: Optional["BaseArtifactStore"] = None,
 ) -> Any:
     """Load an artifact using the given materializer.
 
@@ -495,6 +746,7 @@ def _load_artifact_from_uri(
         materializer: The source of the materializer class to use.
         data_type: The source of the artifact data type.
         uri: The uri of the artifact.
+        artifact_store: The artifact store used to store this artifact.
 
     Returns:
         The artifact loaded into memory.
@@ -504,6 +756,15 @@ def _load_artifact_from_uri(
     """
     from zenml.materializers.base_materializer import BaseMaterializer
 
+    if not artifact_store:
+        artifact_versions_by_uri = Client().list_artifact_versions(uri=uri)
+        if artifact_versions_by_uri.total == 1:
+            artifact_store = (
+                _get_artifact_store_from_response_or_from_active_stack(
+                    artifact_versions_by_uri.items[0]
+                )
+            )
+
     # Resolve the materializer class
     try:
         materializer_class = source_utils.load(materializer)
@@ -531,7 +792,9 @@ def _load_artifact_from_uri(
         artifact_class.__qualname__,
         uri,
     )
-    materializer_object: BaseMaterializer = materializer_class(uri)
+    materializer_object: BaseMaterializer = materializer_class(
+        uri, artifact_store
+    )
     artifact = materializer_object.load(artifact_class)
     logger.debug("Artifact loaded successfully.")
 
@@ -581,7 +844,7 @@ def _load_artifact_store(
             StackComponent.from_model(artifact_store_model),
         )
     except ImportError:
-        link = "https://docs.zenml.io/stacks-and-components/component-guide/artifact-stores/custom#enabling-artifact-visualizations-with-custom-artifact-stores"
+        link = "https://docs.zenml.io/stack-components/artifact-stores/custom#enabling-artifact-visualizations-with-custom-artifact-stores"
         raise NotImplementedError(
             f"Artifact store '{artifact_store_model.name}' could not be "
             f"instantiated. This is likely because the artifact store's "
@@ -591,34 +854,42 @@ def _load_artifact_store(
     return artifact_store
 
 
-def _get_new_artifact_version(artifact_name: str) -> int:
-    """Get the next auto-incremented version for an artifact name.
-
-    Args:
-        artifact_name: The name of the artifact.
-
-    Returns:
-        The next auto-incremented version.
-    """
-    artifact_versions = Client().list_artifact_versions(
-        name=artifact_name,
-        sort_by="desc:version_number",
-        size=1,
-    )
-
-    # If a numbered version exists, increment it
-    try:
-        return int(artifact_versions[0].version) + 1
-
-    # If no numbered versions exist yet, start at 1
-    except (IndexError, ValueError):
-        return 1
+def _get_artifact_store_from_response_or_from_active_stack(
+    artifact: ArtifactVersionResponse,
+) -> "BaseArtifactStore":
+    if artifact.artifact_store_id:
+        try:
+            artifact_store_model = Client().get_stack_component(
+                component_type=StackComponentType.ARTIFACT_STORE,
+                name_id_or_prefix=artifact.artifact_store_id,
+            )
+            return cast(
+                "BaseArtifactStore",
+                StackComponent.from_model(artifact_store_model),
+            )
+        except KeyError:
+            raise RuntimeError(
+                "Unable to fetch the artifact store with id: "
+                f"'{artifact.artifact_store_id}'. Check whether the artifact "
+                "store still exists and you have the right permissions to "
+                "access it."
+            )
+        except ImportError:
+            raise RuntimeError(
+                "Unable to load the implementation of the artifact store with"
+                f"id: '{artifact.artifact_store_id}'. Please make sure that "
+                "the environment that you are loading this artifact from "
+                "has the right dependencies."
+            )
+    return Client().active_stack.artifact_store
 
 
 def _load_file_from_artifact_store(
     uri: str,
     artifact_store: "BaseArtifactStore",
     mode: str = "rb",
+    offset: int = 0,
+    length: Optional[int] = None,
 ) -> Any:
     """Load the given uri from the given artifact store.
 
@@ -626,6 +897,8 @@ def _load_file_from_artifact_store(
         uri: The uri of the file to load.
         artifact_store: The artifact store from which to load the file.
         mode: The mode in which to open the file.
+        offset: The offset from which to start reading.
+        length: The amount of bytes that should be read.
 
     Returns:
         The loaded file.
@@ -633,18 +906,33 @@ def _load_file_from_artifact_store(
     Raises:
         DoesNotExistException: If the file does not exist in the artifact store.
         NotImplementedError: If the artifact store cannot open the file.
+        IOError: If the artifact store rejects the request.
     """
     try:
         with artifact_store.open(uri, mode) as text_file:
-            return text_file.read()
+            if offset < 0:
+                # If the offset is negative, we seek backwards from the end of
+                # the file
+                try:
+                    text_file.seek(offset, os.SEEK_END)
+                except OSError:
+                    # The negative offset was too large for the file, we seek
+                    # to the start of the file
+                    text_file.seek(0, os.SEEK_SET)
+            elif offset > 0:
+                text_file.seek(offset, os.SEEK_SET)
+
+            return text_file.read(length)
     except FileNotFoundError:
         raise DoesNotExistException(
             f"File '{uri}' does not exist in artifact store "
             f"'{artifact_store.name}'."
         )
+    except IOError as e:
+        raise e
     except Exception as e:
         logger.exception(e)
-        link = "https://docs.zenml.io/stacks-and-components/component-guide/artifact-stores/custom#enabling-artifact-visualizations-with-custom-artifact-stores"
+        link = "https://docs.zenml.io/stack-components/artifact-stores/custom#enabling-artifact-visualizations-with-custom-artifact-stores"
         raise NotImplementedError(
             f"File '{uri}' could not be loaded because the underlying artifact "
             f"store '{artifact_store.name}' could not open the file. This is "
@@ -700,14 +988,27 @@ def load_model_from_metadata(model_uri: str) -> Any:
         The ML model object loaded into memory.
     """
     # Load the model from its metadata
-    with fileio.open(
+    artifact_versions_by_uri = Client().list_artifact_versions(uri=model_uri)
+    if artifact_versions_by_uri.total == 1:
+        artifact_store = (
+            _get_artifact_store_from_response_or_from_active_stack(
+                artifact_versions_by_uri.items[0]
+            )
+        )
+    else:
+        artifact_store = Client().active_stack.artifact_store
+
+    with artifact_store.open(
         os.path.join(model_uri, MODEL_METADATA_YAML_FILE_NAME), "r"
     ) as f:
         metadata = read_yaml(f.name)
     data_type = metadata["datatype"]
     materializer = metadata["materializer"]
     model = _load_artifact_from_uri(
-        materializer=materializer, data_type=data_type, uri=model_uri
+        materializer=materializer,
+        data_type=data_type,
+        uri=model_uri,
+        artifact_store=artifact_store,
     )
 
     # Switch to eval mode if the model is a torch model