zenml-nightly 0.71.0.dev20250107__py3-none-any.whl → 0.71.0.dev20250111__py3-none-any.whl

zenml/zen_stores/schemas/tag_schemas.py

@@ -14,7 +14,7 @@
 """SQLModel implementation of tag tables."""
 
 from datetime import datetime
-from typing import TYPE_CHECKING, Any, List
+from typing import Any, List
 from uuid import UUID
 
 from sqlalchemy import VARCHAR, Column
@@ -33,16 +33,6 @@ from zenml.models import (
 from zenml.zen_stores.schemas.base_schemas import BaseSchema, NamedSchema
 from zenml.zen_stores.schemas.schema_utils import build_foreign_key_field
 
-if TYPE_CHECKING:
-    from zenml.zen_stores.schemas.artifact_schemas import (
-        ArtifactSchema,
-        ArtifactVersionSchema,
-    )
-    from zenml.zen_stores.schemas.model_schemas import (
-        ModelSchema,
-        ModelVersionSchema,
-    )
-
 
 class TagSchema(NamedSchema, table=True):
     """SQL Model for tag."""
@@ -52,7 +42,7 @@ class TagSchema(NamedSchema, table=True):
     color: str = Field(sa_column=Column(VARCHAR(255), nullable=False))
     links: List["TagResourceSchema"] = Relationship(
         back_populates="tag",
-        sa_relationship_kwargs={"cascade": "delete"},
+        sa_relationship_kwargs={"overlaps": "tags", "cascade": "delete"},
     )
 
     @classmethod
@@ -130,37 +120,11 @@ class TagResourceSchema(BaseSchema, table=True):
         ondelete="CASCADE",
         nullable=False,
     )
-    tag: "TagSchema" = Relationship(back_populates="links")
+    tag: "TagSchema" = Relationship(
+        back_populates="links", sa_relationship_kwargs={"overlaps": "tags"}
+    )
     resource_id: UUID
     resource_type: str = Field(sa_column=Column(VARCHAR(255), nullable=False))
-    artifact: List["ArtifactSchema"] = Relationship(
-        back_populates="tags",
-        sa_relationship_kwargs=dict(
-            primaryjoin=f"and_(TagResourceSchema.resource_type=='{TaggableResourceTypes.ARTIFACT.value}', foreign(TagResourceSchema.resource_id)==ArtifactSchema.id)",
-            overlaps="tags,model,artifact_version,model_version",
-        ),
-    )
-    artifact_version: List["ArtifactVersionSchema"] = Relationship(
-        back_populates="tags",
-        sa_relationship_kwargs=dict(
-            primaryjoin=f"and_(TagResourceSchema.resource_type=='{TaggableResourceTypes.ARTIFACT_VERSION.value}', foreign(TagResourceSchema.resource_id)==ArtifactVersionSchema.id)",
-            overlaps="tags,model,artifact,model_version",
-        ),
-    )
-    model: List["ModelSchema"] = Relationship(
-        back_populates="tags",
-        sa_relationship_kwargs=dict(
-            primaryjoin=f"and_(TagResourceSchema.resource_type=='{TaggableResourceTypes.MODEL.value}', foreign(TagResourceSchema.resource_id)==ModelSchema.id)",
-            overlaps="tags,artifact,artifact_version,model_version",
-        ),
-    )
-    model_version: List["ModelVersionSchema"] = Relationship(
-        back_populates="tags",
-        sa_relationship_kwargs=dict(
-            primaryjoin=f"and_(TagResourceSchema.resource_type=='{TaggableResourceTypes.MODEL_VERSION.value}', foreign(TagResourceSchema.resource_id)==ModelVersionSchema.id)",
-            overlaps="tags,model,artifact,artifact_version",
-        ),
-    )
 
     @classmethod
     def from_request(cls, request: TagResourceRequest) -> "TagResourceSchema":

zenml/zen_stores/schemas/utils.py

@@ -75,35 +75,34 @@ def get_page_from_list(
 class RunMetadataInterface:
     """The interface for entities with run metadata."""
 
-    run_metadata_resources = Relationship()
+    run_metadata = Relationship()
 
     def fetch_metadata_collection(self) -> Dict[str, List[RunMetadataEntry]]:
-        """Fetches all the metadata entries related to the artifact version.
+        """Fetches all the metadata entries related to the entity.
 
         Returns:
-            a dictionary, where the key is the key of the metadata entry
+            A dictionary, where the key is the key of the metadata entry
                 and the values represent the list of entries with this key.
         """
         metadata_collection: Dict[str, List[RunMetadataEntry]] = {}
 
-        # Fetch the metadata related to this step
-        for rm in self.run_metadata_resources:
-            if rm.run_metadata.key not in metadata_collection:
-                metadata_collection[rm.run_metadata.key] = []
-            metadata_collection[rm.run_metadata.key].append(
+        for rm in self.run_metadata:
+            if rm.key not in metadata_collection:
+                metadata_collection[rm.key] = []
+            metadata_collection[rm.key].append(
                 RunMetadataEntry(
-                    value=json.loads(rm.run_metadata.value),
-                    created=rm.run_metadata.created,
+                    value=json.loads(rm.value),
+                    created=rm.created,
                 )
             )
 
         return metadata_collection
 
     def fetch_metadata(self) -> Dict[str, MetadataType]:
-        """Fetches the latest metadata entry related to the artifact version.
+        """Fetches the latest metadata entry related to the entity.
 
         Returns:
-            a dictionary, where the key is the key of the metadata entry
+            A dictionary, where the key is the key of the metadata entry
                 and the values represent the latest entry with this key.
         """
         metadata_collection = self.fetch_metadata_collection()

zenml/zen_stores/sql_zen_store.py

@@ -64,6 +64,13 @@ from sqlalchemy.exc import (
 )
 from sqlalchemy.orm import Mapped, noload
 from sqlalchemy.util import immutabledict
+
+# Important to note: The select function of SQLModel works slightly differently
+# from the select function of sqlalchemy. If you input only one entity on the
+# select function of SQLModel, it automatically maps it to a SelectOfScalar.
+# As a result, it will not return a tuple as a result, but the first entity in
+# the tuple. While this is convenient in most cases, in unique cases like using
+# the "add_columns" functionality, one might encounter unexpected results.
 from sqlmodel import (
     Session,
     SQLModel,
@@ -1362,8 +1369,7 @@ class SqlZenStore(BaseZenStore):
                     os.remove(dump_file)
                 except OSError:
                     logger.warning(
-                        f"Failed to cleanup database dump file "
-                        f"{dump_file}."
+                        f"Failed to cleanup database dump file {dump_file}."
                     )
                 else:
                     logger.info(
@@ -7392,7 +7398,7 @@ class SqlZenStore(BaseZenStore):
                     )
                 except (ValueError, AuthorizationException) as e:
                     error = (
-                        f'Failed to fetch {resource_type or "available"} '
+                        f"Failed to fetch {resource_type or 'available'} "
                         f"resources from service connector {connector.name}/"
                         f"{connector.id}: {e}"
                     )

{zenml_nightly-0.71.0.dev20250107.dist-info → zenml_nightly-0.71.0.dev20250111.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.3
 Name: zenml-nightly
-Version: 0.71.0.dev20250107
+Version: 0.71.0.dev20250111
 Summary: ZenML: Write production-ready ML code.
 Home-page: https://zenml.io
 License: Apache-2.0
@@ -44,7 +44,7 @@ Requires-Dist: Jinja2 ; extra == "server"
 Requires-Dist: adlfs (>=2021.10.0) ; extra == "adlfs"
 Requires-Dist: alembic (>=1.8.1,<1.9.0)
 Requires-Dist: aws-profile-manager (>=0.5.0) ; extra == "connectors-aws"
-Requires-Dist: azure-ai-ml (==1.18.0) ; extra == "azureml"
+Requires-Dist: azure-ai-ml (==1.23.1) ; extra == "azureml"
 Requires-Dist: azure-identity (>=1.4.0) ; extra == "secrets-azure" or extra == "connectors-azure"
 Requires-Dist: azure-keyvault-secrets (>=4.0.0) ; extra == "secrets-azure"
 Requires-Dist: azure-mgmt-containerregistry (>=10.0.0) ; extra == "connectors-azure"