zenml-nightly 0.73.0.dev20250130__py3-none-any.whl → 0.73.0.dev20250201__py3-none-any.whl

zenml/pipelines/build_utils.py

@@ -517,32 +517,9 @@ def verify_local_repository_context(
                     "changes."
                 )
 
-        if local_repo_context:
-            if local_repo_context.is_dirty:
-                logger.warning(
-                    "Unable to use code repository to download code for this "
-                    "run as there are uncommitted changes."
-                )
-            elif local_repo_context.has_local_changes:
-                logger.warning(
-                    "Unable to use code repository to download code for this "
-                    "run as there are unpushed changes."
-                )
-
     code_repository = None
     if local_repo_context and not local_repo_context.has_local_changes:
-        model = Client().get_code_repository(
-            local_repo_context.code_repository_id
-        )
-        code_repository = BaseCodeRepository.from_model(model)
-
-        if will_download_from_code_repository(
-            deployment=deployment, local_repo_context=local_repo_context
-        ):
-            logger.info(
-                "Using code repository `%s` to download code for this run.",
-                model.name,
-            )
+        code_repository = local_repo_context.code_repository
 
     return code_repository
 
@@ -738,25 +715,17 @@ def should_upload_code(
     return False
 
 
-def will_download_from_code_repository(
+def allows_download_from_code_repository(
     deployment: PipelineDeploymentBase,
-    local_repo_context: "LocalRepositoryContext",
 ) -> bool:
-    """Checks whether a code repository will be used to download code.
+    """Checks whether a code repository can be used to download code.
 
     Args:
         deployment: The deployment.
-        local_repo_context: The local repository context.
 
     Returns:
-        Whether a code repository will be used to download code.
+        Whether a code repository can be used to download code.
     """
-    if not build_required(deployment=deployment):
-        return False
-
-    if local_repo_context.has_local_changes:
-        return False
-
     for step in deployment.step_configurations.values():
         docker_settings = step.config.docker_settings
 
@@ -764,3 +733,41 @@ def will_download_from_code_repository(
             return True
 
     return False
+
+
+def log_code_repository_usage(
+    deployment: PipelineDeploymentBase,
+    local_repo_context: "LocalRepositoryContext",
+) -> None:
+    """Log what the code repository can (not) be used for given a deployment.
+
+    Args:
+        deployment: The deployment.
+        local_repo_context: The local repository context.
+    """
+    if build_required(deployment) and allows_download_from_code_repository(
+        deployment
+    ):
+        if local_repo_context.is_dirty:
+            logger.warning(
+                "Unable to use code repository `%s` to download code or track "
+                "the commit hash as there are uncommitted or untracked files.",
+                local_repo_context.code_repository.name,
+            )
+        elif local_repo_context.has_local_changes:
+            logger.warning(
+                "Unable to use code repository `%s` to download code as there "
+                "are unpushed commits.",
+                local_repo_context.code_repository.name,
+            )
+        else:
+            logger.info(
+                "Using code repository `%s` to download code for this run.",
+                local_repo_context.code_repository.name,
+            )
+    elif local_repo_context.is_dirty:
+        logger.warning(
+            "Unable to use code repository `%s` to track the commit hash as "
+            "there are uncommitted or untracked files.",
+            local_repo_context.code_repository.name,
+        )

zenml/pipelines/pipeline_definition.py

@@ -643,7 +643,6 @@ To avoid this consider setting pipeline parameters only in one place (config or
         pipeline_id = None
         if register_pipeline:
             pipeline_id = self._register().id
-
         else:
             logger.debug(f"Pipeline {self.name} is unlisted.")
 
@@ -702,6 +701,10 @@ To avoid this consider setting pipeline parameters only in one place (config or
             deployment=deployment, local_repo_context=local_repo_context
         )
         can_download_from_code_repository = code_repository is not None
+        if local_repo_context:
+            build_utils.log_code_repository_usage(
+                deployment=deployment, local_repo_context=local_repo_context
+            )
 
         if prevent_build_reuse:
             logger.warning(
@@ -731,7 +734,7 @@ To avoid this consider setting pipeline parameters only in one place (config or
             code_reference = CodeReferenceRequest(
                 commit=local_repo_context.current_commit,
                 subdirectory=subdirectory.as_posix(),
-                code_repository=local_repo_context.code_repository_id,
+                code_repository=local_repo_context.code_repository.id,
             )
 
         code_path = None

zenml/utils/code_repository_utils.py

@@ -79,7 +79,7 @@ def set_custom_local_repository(
 
     path = os.path.abspath(source_utils.get_source_root())
     _CODE_REPOSITORY_CACHE[path] = _DownloadedRepositoryContext(
-        code_repository_id=repo.id, root=root, commit=commit
+        code_repository=repo, root=root, commit=commit
     )
 
 
@@ -106,7 +106,8 @@ def find_active_code_repository(
         return _CODE_REPOSITORY_CACHE[path]
 
     local_context: Optional["LocalRepositoryContext"] = None
-    for model in depaginate(list_method=Client().list_code_repositories):
+    code_repositories = depaginate(list_method=Client().list_code_repositories)
+    for model in code_repositories:
         try:
             repo = BaseCodeRepository.from_model(model)
         except ImportError:
@@ -125,6 +126,14 @@ def find_active_code_repository(
         local_context = repo.get_local_context(path)
         if local_context:
             break
+    else:
+        if code_repositories:
+            # There are registered code repositories, but none was matching the
+            # current path -> We log the path to help in debugging issues
+            # related to the source root.
+            logger.info(
+                "No matching code repository found for path `%s`.", path
+            )
 
     _CODE_REPOSITORY_CACHE[path] = local_context
     return local_context

zenml/utils/downloaded_repository_context.py

@@ -13,9 +13,7 @@
 #  permissions and limitations under the License.
 """Downloaded code repository."""
 
-from uuid import UUID
-
-from zenml.code_repositories import LocalRepositoryContext
+from zenml.code_repositories import BaseCodeRepository, LocalRepositoryContext
 
 
 class _DownloadedRepositoryContext(LocalRepositoryContext):
@@ -27,11 +25,11 @@ class _DownloadedRepositoryContext(LocalRepositoryContext):
 
     def __init__(
         self,
-        code_repository_id: UUID,
+        code_repository: BaseCodeRepository,
         root: str,
         commit: str,
     ):
-        super().__init__(code_repository_id=code_repository_id)
+        super().__init__(code_repository=code_repository)
         self._root = root
         self._commit = commit
 

zenml/utils/source_utils.py

@@ -226,7 +226,7 @@ def resolve(
             subdir = PurePath(source_root).relative_to(local_repo_context.root)
 
             return CodeRepositorySource(
-                repository_id=local_repo_context.code_repository_id,
+                repository_id=local_repo_context.code_repository.id,
                 commit=local_repo_context.current_commit,
                 subdirectory=subdir.as_posix(),
                 module=module_name,
@@ -482,7 +482,7 @@ def _warn_about_potential_source_loading_issues(
             source.repository_id,
             get_source_root(),
         )
-    elif local_repo.code_repository_id != source.repository_id:
+    elif local_repo.code_repository.id != source.repository_id:
         logger.warning(
             "Potential issue when loading the source `%s`: The source "
             "references the code repository `%s` but there is a different "
@@ -492,7 +492,7 @@ def _warn_about_potential_source_loading_issues(
             "source was originally stored.",
             source.import_path,
             source.repository_id,
-            local_repo.code_repository_id,
+            local_repo.code_repository.id,
             get_source_root(),
         )
     elif local_repo.current_commit != source.commit:

zenml/zen_stores/migrations/utils.py

@@ -34,7 +34,7 @@ from sqlalchemy.engine import URL, Engine
 from sqlalchemy.exc import (
     OperationalError,
 )
-from sqlalchemy.schema import CreateTable
+from sqlalchemy.schema import CreateIndex, CreateTable
 from sqlmodel import (
     create_engine,
     select,
@@ -249,6 +249,7 @@ class MigrationUtils(BaseModel):
                 # them to the create table statement.
 
                 # Extract the unique constraints from the table schema
+                index_create_statements = []
                 unique_constraints = []
                 for index in table.indexes:
                     if index.unique:
@@ -258,6 +259,38 @@ class MigrationUtils(BaseModel):
                         unique_constraints.append(
                             f"UNIQUE KEY `{index.name}` ({', '.join(unique_columns)})"
                         )
+                    else:
+                        if index.name in {
+                            fk.name for fk in table.foreign_key_constraints
+                        }:
+                            # Foreign key indices are already handled by the
+                            # table creation statement.
+                            continue
+
+                        index_create = str(CreateIndex(index)).strip()  # type: ignore[no-untyped-call]
+                        index_create = index_create.replace(
+                            f"CREATE INDEX {index.name}",
+                            f"CREATE INDEX `{index.name}`",
+                        )
+                        index_create = index_create.replace(
+                            f"ON {table.name}", f"ON `{table.name}`"
+                        )
+
+                        for column_name in index.columns.keys():
+                            # We need this logic here to avoid the column names
+                            # inside the index name
+                            index_create = index_create.replace(
+                                f"({column_name}", f"(`{column_name}`"
+                            )
+                            index_create = index_create.replace(
+                                f"{column_name},", f"`{column_name}`,"
+                            )
+                            index_create = index_create.replace(
+                                f"{column_name})", f"`{column_name}`)"
+                            )
+
+                        index_create = index_create.replace('"', "") + ";"
+                        index_create_statements.append(index_create)
 
                 # Add the unique constraints to the create table statement
                 if unique_constraints:
@@ -290,6 +323,14 @@ class MigrationUtils(BaseModel):
                     )
                 )
 
+                for stmt in index_create_statements:
+                    store_db_info(
+                        dict(
+                            table=table.name,
+                            index_create_stmt=stmt,
+                        )
+                    )
+
                 # 2. extract the table data in batches
                 order_by = [col for col in table.primary_key]
 
@@ -356,6 +397,12 @@ class MigrationUtils(BaseModel):
                         "self_references", False
                     )
 
+                if "index_create_stmt" in table_dump:
+                    # execute the index creation statement
+                    connection.execute(text(table_dump["index_create_stmt"]))
+                    # Reload the database metadata after creating the index
+                    metadata.reflect(bind=self.engine)
+
                 if "data" in table_dump:
                     # insert the data into the database
                     table = metadata.tables[table_name]

zenml/zen_stores/migrations/versions/4d5524b92a30_add_run_metadata_tag_index.py

@@ -0,0 +1,67 @@
+"""Add run metadata and tag index [4d5524b92a30].
+
+Revision ID: 4d5524b92a30
+Revises: 0.73.0
+Create Date: 2025-01-30 11:30:36.736452
+
+"""
+
+from alembic import op
+from sqlalchemy import inspect
+
+# revision identifiers, used by Alembic.
+revision = "4d5524b92a30"
+down_revision = "0.73.0"
+branch_labels = None
+depends_on = None
+
+
+def upgrade() -> None:
+    """Upgrade database schema and/or data, creating a new revision."""
+    connection = op.get_bind()
+
+    inspector = inspect(connection)
+    for index in inspector.get_indexes("run_metadata_resource"):
+        # This index was manually added to some databases to improve the
+        # speed and cache utilisation. In this case we simply return here and
+        # don't continue with the migration.
+        if (
+            index["name"]
+            == "ix_run_metadata_resource_resource_id_resource_type_run_metadata_"
+        ):
+            return
+
+    # ### commands auto generated by Alembic - please adjust! ###
+    with op.batch_alter_table(
+        "run_metadata_resource", schema=None
+    ) as batch_op:
+        batch_op.create_index(
+            "ix_run_metadata_resource_resource_id_resource_type_run_metadata_",
+            ["resource_id", "resource_type", "run_metadata_id"],
+            unique=False,
+        )
+
+    with op.batch_alter_table("tag_resource", schema=None) as batch_op:
+        batch_op.create_index(
+            "ix_tag_resource_resource_id_resource_type_tag_id",
+            ["resource_id", "resource_type", "tag_id"],
+            unique=False,
+        )
+
+    # ### end Alembic commands ###
+
+
+def downgrade() -> None:
+    """Downgrade database schema and/or data back to the previous revision."""
+    # ### commands auto generated by Alembic - please adjust! ###
+    with op.batch_alter_table("tag_resource", schema=None) as batch_op:
+        batch_op.drop_index("ix_tag_resource_resource_id_resource_type_tag_id")
+
+    with op.batch_alter_table(
+        "run_metadata_resource", schema=None
+    ) as batch_op:
+        batch_op.drop_index(
+            "ix_run_metadata_resource_resource_id_resource_type_run_metadata_"
+        )
+
+    # ### end Alembic commands ###

zenml/zen_stores/schemas/run_metadata_schemas.py

@@ -11,7 +11,7 @@
 #  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express
 #  or implied. See the License for the specific language governing
 #  permissions and limitations under the License.
-"""SQLModel implementation of pipeline run metadata tables."""
+"""SQLModel implementation of run metadata tables."""
 
 from typing import Optional
 from uuid import UUID, uuid4
@@ -21,7 +21,10 @@ from sqlmodel import Field, Relationship, SQLModel
 
 from zenml.zen_stores.schemas.base_schemas import BaseSchema
 from zenml.zen_stores.schemas.component_schemas import StackComponentSchema
-from zenml.zen_stores.schemas.schema_utils import build_foreign_key_field
+from zenml.zen_stores.schemas.schema_utils import (
+    build_foreign_key_field,
+    build_index,
+)
 from zenml.zen_stores.schemas.step_run_schemas import StepRunSchema
 from zenml.zen_stores.schemas.user_schemas import UserSchema
 from zenml.zen_stores.schemas.workspace_schemas import WorkspaceSchema
@@ -82,6 +85,16 @@ class RunMetadataResourceSchema(SQLModel, table=True):
     """Table for linking resources to run metadata entries."""
 
     __tablename__ = "run_metadata_resource"
+    __table_args__ = (
+        build_index(
+            table_name=__tablename__,
+            column_names=[
+                "resource_id",
+                "resource_type",
+                "run_metadata_id",
+            ],
+        ),
+    )
 
     id: UUID = Field(default_factory=uuid4, primary_key=True)
     resource_id: UUID

zenml/zen_stores/schemas/schema_utils.py

@@ -13,9 +13,9 @@
 #  permissions and limitations under the License.
 """Utility functions for SQLModel schemas."""
 
-from typing import Any
+from typing import Any, List
 
-from sqlalchemy import Column, ForeignKey
+from sqlalchemy import Column, ForeignKey, Index
 from sqlmodel import Field
 
 
@@ -84,3 +84,35 @@ def build_foreign_key_field(
             **sa_column_kwargs,
         ),
     )
+
+
+def get_index_name(table_name: str, column_names: List[str]) -> str:
+    """Get the name for an index.
+
+    Args:
+        table_name: The name of the table for which the index will be created.
+        column_names: Names of the columns on which the index will be created.
+
+    Returns:
+        The index name.
+    """
+    columns = "_".join(column_names)
+    # MySQL allows a maximum of 64 characters in identifiers
+    return f"ix_{table_name}_{columns}"[:64]
+
+
+def build_index(
+    table_name: str, column_names: List[str], **kwargs: Any
+) -> Index:
+    """Build an index object.
+
+    Args:
+        table_name: The name of the table for which the index will be created.
+        column_names: Names of the columns on which the index will be created.
+        **kwargs: Additional keyword arguments to pass to the Index.
+
+    Returns:
+        The index.
+    """
+    name = get_index_name(table_name=table_name, column_names=column_names)
+    return Index(name, *column_names, **kwargs)

zenml/zen_stores/schemas/tag_schemas.py

@@ -31,7 +31,10 @@ from zenml.models import (
 )
 from zenml.utils.time_utils import utc_now
 from zenml.zen_stores.schemas.base_schemas import BaseSchema, NamedSchema
-from zenml.zen_stores.schemas.schema_utils import build_foreign_key_field
+from zenml.zen_stores.schemas.schema_utils import (
+    build_foreign_key_field,
+    build_index,
+)
 
 
 class TagSchema(NamedSchema, table=True):
@@ -111,6 +114,16 @@ class TagResourceSchema(BaseSchema, table=True):
     """SQL Model for tag resource relationship."""
 
     __tablename__ = "tag_resource"
+    __table_args__ = (
+        build_index(
+            table_name=__tablename__,
+            column_names=[
+                "resource_id",
+                "resource_type",
+                "tag_id",
+            ],
+        ),
+    )
 
     tag_id: UUID = build_foreign_key_field(
         source=__tablename__,

{zenml_nightly-0.73.0.dev20250130.dist-info → zenml_nightly-0.73.0.dev20250201.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.3
 Name: zenml-nightly
-Version: 0.73.0.dev20250130
+Version: 0.73.0.dev20250201
 Summary: ZenML: Write production-ready ML code.
 License: Apache-2.0
 Keywords: machine learning,production,pipeline,mlops,devops
@@ -104,6 +104,7 @@ Requires-Dist: python-dateutil (>=2.8.1,<3.0.0)
 Requires-Dist: python-multipart (>=0.0.9,<0.1.0) ; extra == "server"
 Requires-Dist: pyyaml (>=6.0.1)
 Requires-Dist: pyyaml-include (<2.0) ; extra == "templates"
+Requires-Dist: requests (>=2.27.11,<3.0.0) ; extra == "connectors-azure"
 Requires-Dist: rich[jupyter] (>=12.0.0)
 Requires-Dist: ruff (>=0.1.7) ; extra == "templates" or extra == "dev"
 Requires-Dist: s3fs (>=2022.11.0) ; extra == "s3fs"