CI fixes

Author: schustmi

src/zenml/artifacts/in_memory_cache.py

@@ -13,7 +13,8 @@
 #  permissions and limitations under the License.
 """In-memory artifact cache."""
 
-from typing import Any
+import contextvars
+from typing import Any, Dict
 from uuid import UUID
 
 from zenml.utils import context_utils
@@ -22,12 +23,12 @@
 class InMemoryArtifactCache(context_utils.BaseContext):
     """In-memory artifact cache."""
 
-    __context_var__ = context_utils.ContextVar("in_memory_artifact_cache")
+    __context_var__ = contextvars.ContextVar("in_memory_artifact_cache")
 
     def __init__(self) -> None:
         """Initialize the artifact cache."""
         super().__init__()
-        self._cache = {}
+        self._cache: Dict[UUID, Any] = {}
 
     def clear(self) -> None:
         """Clear the artifact cache."""

src/zenml/execution/pipeline/dynamic/outputs.py

@@ -140,11 +140,17 @@ def get_artifact(self, key: str) -> ArtifactFuture:
         Args:
             key: The key of the artifact future.
 
+        Raises:
+            KeyError: If no artifact for the given name exists.
+
         Returns:
             The artifact future.
         """
         if key not in self._output_keys:
-            raise KeyError(f"Invalid key: {key}")
+            raise KeyError(
+                f"Step run {self._invocation_id} does not have an output with "
+                f"the name: {key}."
+            )
 
         return ArtifactFuture(
             wrapped=self._wrapped,

src/zenml/models/v2/core/artifact_version.py

@@ -455,11 +455,11 @@ def load(self, disable_cache: bool = False) -> Any:
         cache = InMemoryArtifactCache.get()
 
         if cache and (data := cache.get_artifact_data(self.id)):
-            logger.debug(f"Returning artifact data (%s) from cache", self.id)
+            logger.debug("Returning artifact data (%s) from cache", self.id)
             return data
 
         data = load_artifact_from_response(self)
-        if not disable_cache:
+        if cache and not disable_cache:
             cache.set_artifact_data(self.id, data)
         return data