zenml-nightly 0.73.0.dev20250129__py3-none-any.whl → 0.73.0.dev20250131__py3-none-any.whl

zenml/integrations/vllm/services/vllm_deployment.py

@@ -152,9 +152,14 @@ class VLLMDeploymentService(LocalDaemonService, BaseDeploymentService):
             parser: argparse.ArgumentParser = make_arg_parser(
                 FlexibleArgumentParser()
             )
-            args: argparse.Namespace = parser.parse_args()
+            # pass in empty list to get default args
+            # otherwise it will try to get the args from sys.argv
+            # and if there's a --config in there, it will want to use
+            # that file for vLLM configuration, which is not what we want
+            args: argparse.Namespace = parser.parse_args(args=[])
             # Override port with the available port
             self.config.port = self.endpoint.status.port or self.config.port
+
             # Update the arguments in place
             args.__dict__.update(self.config.model_dump())
             uvloop.run(run_server(args=args))

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/rest_zen_store.py

@@ -365,17 +365,16 @@ class RestZenStoreConfiguration(StoreConfiguration):
 
         if os.path.isfile(verify_ssl):
             with open(verify_ssl, "r") as f:
-                verify_ssl = f.read()
+                cert_content = f.read()
 
         fileio.makedirs(str(secret_folder))
         file_path = Path(secret_folder, "ca_bundle.pem")
         with os.fdopen(
             os.open(file_path, flags=os.O_RDWR | os.O_CREAT, mode=0o600), "w"
         ) as f:
-            f.write(verify_ssl)
-        verify_ssl = str(file_path)
+            f.write(cert_content)
 
-        return verify_ssl
+        return str(file_path)
 
     @classmethod
     def supports_url_scheme(cls, url: str) -> bool:
@@ -389,15 +388,6 @@ class RestZenStoreConfiguration(StoreConfiguration):
         """
         return urlparse(url).scheme in ("http", "https")
 
-    def expand_certificates(self) -> None:
-        """Expands the certificates in the verify_ssl field."""
-        # Load the certificate values back into the configuration
-        if isinstance(self.verify_ssl, str) and os.path.isfile(
-            self.verify_ssl
-        ):
-            with open(self.verify_ssl, "r") as f:
-                self.verify_ssl = f.read()
-
     @model_validator(mode="before")
     @classmethod
     @before_validator_handler

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/zen_stores/secrets_stores/sql_secrets_store.py

@@ -37,6 +37,7 @@ from zenml.exceptions import (
     IllegalOperationError,
 )
 from zenml.logger import get_logger
+from zenml.utils.secret_utils import PlainSerializedSecretStr
 from zenml.zen_stores.schemas import (
     SecretSchema,
 )
@@ -62,7 +63,7 @@ class SqlSecretsStoreConfiguration(SecretsStoreConfiguration):
     """
 
     type: SecretsStoreType = SecretsStoreType.SQL
-    encryption_key: Optional[str] = None
+    encryption_key: Optional[PlainSerializedSecretStr] = None
     model_config = ConfigDict(
         # Don't validate attributes when assigning them. This is necessary
         # because the certificate attributes can be expanded to the contents
@@ -159,7 +160,9 @@ class SqlSecretsStore(BaseSecretsStore):
         # Initialize the encryption engine
         if self.config.encryption_key:
             self._encryption_engine = AesGcmEngine()
-            self._encryption_engine._update_key(self.config.encryption_key)
+            self._encryption_engine._update_key(
+                self.config.encryption_key.get_secret_value()
+            )
 
         # Nothing else to do here, the SQL ZenML store back-end is already
         # initialized

zenml/zen_stores/sql_zen_store.py

@@ -304,6 +304,7 @@ from zenml.utils.networking_utils import (
     replace_localhost_with_internal_hostname,
 )
 from zenml.utils.pydantic_utils import before_validator_handler
+from zenml.utils.secret_utils import PlainSerializedSecretStr
 from zenml.utils.string_utils import (
     format_name_template,
     random_str,
@@ -460,11 +461,11 @@ class SqlZenStoreConfiguration(StoreConfiguration):
 
     driver: Optional[SQLDatabaseDriver] = None
     database: Optional[str] = None
-    username: Optional[str] = None
-    password: Optional[str] = None
-    ssl_ca: Optional[str] = None
-    ssl_cert: Optional[str] = None
-    ssl_key: Optional[str] = None
+    username: Optional[PlainSerializedSecretStr] = None
+    password: Optional[PlainSerializedSecretStr] = None
+    ssl_ca: Optional[PlainSerializedSecretStr] = None
+    ssl_cert: Optional[PlainSerializedSecretStr] = None
+    ssl_key: Optional[PlainSerializedSecretStr] = None
     ssl_verify_server_cert: bool = False
     pool_size: int = 20
     max_overflow: int = 20
@@ -611,10 +612,10 @@ class SqlZenStoreConfiguration(StoreConfiguration):
             self.database = sql_url.database
         elif sql_url.drivername == SQLDatabaseDriver.MYSQL:
             if sql_url.username:
-                self.username = sql_url.username
+                self.username = PlainSerializedSecretStr(sql_url.username)
                 sql_url = sql_url._replace(username=None)
             if sql_url.password:
-                self.password = sql_url.password
+                self.password = PlainSerializedSecretStr(sql_url.password)
                 sql_url = sql_url._replace(password=None)
             if sql_url.database:
                 self.database = sql_url.database
@@ -642,13 +643,13 @@ class SqlZenStoreConfiguration(StoreConfiguration):
                 for k, v in sql_url.query.items():
                     if k == "ssl_ca":
                         if r := _get_query_result(v):
-                            self.ssl_ca = r
+                            self.ssl_ca = PlainSerializedSecretStr(r)
                     elif k == "ssl_cert":
                         if r := _get_query_result(v):
-                            self.ssl_cert = r
+                            self.ssl_cert = PlainSerializedSecretStr(r)
                     elif k == "ssl_key":
                         if r := _get_query_result(v):
-                            self.ssl_key = r
+                            self.ssl_key = PlainSerializedSecretStr(r)
                     elif k == "ssl_verify_server_cert":
                         if r := _get_query_result(v):
                             if is_true_string_value(r):
@@ -688,7 +689,7 @@ class SqlZenStoreConfiguration(StoreConfiguration):
             )
             for key in ["ssl_key", "ssl_ca", "ssl_cert"]:
                 content = getattr(self, key)
-                if content and not os.path.isfile(content):
+                if content and not os.path.isfile(content.get_secret_value()):
                     fileio.makedirs(str(secret_folder))
                     file_path = Path(secret_folder, f"{key}.pem")
                     with os.fdopen(
@@ -697,7 +698,7 @@ class SqlZenStoreConfiguration(StoreConfiguration):
                         ),
                         "w",
                     ) as f:
-                        f.write(content)
+                        f.write(content.get_secret_value())
                     setattr(self, key, str(file_path))
 
         self.url = str(sql_url)
@@ -732,7 +733,7 @@ class SqlZenStoreConfiguration(StoreConfiguration):
         # Load the certificate values back into the configuration
         for key in ["ssl_key", "ssl_ca", "ssl_cert"]:
             file_path = getattr(self, key, None)
-            if file_path and os.path.isfile(file_path):
+            if file_path and os.path.isfile(file_path.get_secret_value()):
                 with open(file_path, "r") as f:
                     setattr(self, key, f.read())
 
@@ -780,8 +781,8 @@ class SqlZenStoreConfiguration(StoreConfiguration):
 
             sql_url = sql_url._replace(
                 drivername="mysql+pymysql",
-                username=self.username,
-                password=self.password,
+                username=self.username.get_secret_value(),
+                password=self.password.get_secret_value(),
                 database=database,
             )
 
@@ -792,11 +793,17 @@ class SqlZenStoreConfiguration(StoreConfiguration):
                 ssl_setting = getattr(self, key)
                 if not ssl_setting:
                     continue
-                if not os.path.isfile(ssl_setting):
+                if not os.path.isfile(ssl_setting.get_secret_value()):
                     logger.warning(
                         f"Database SSL setting `{key}` is not a file. "
                     )
-                sqlalchemy_ssl_args[key.removeprefix("ssl_")] = ssl_setting
+                sqlalchemy_ssl_args[key.lstrip("ssl_")] = (
+                    ssl_setting.get_secret_value()
+                )
+                sqlalchemy_ssl_args[key.removeprefix("ssl_")] = (
+                    ssl_setting.get_secret_value()
+                )
+
             if len(sqlalchemy_ssl_args) > 0:
                 sqlalchemy_ssl_args["check_hostname"] = (
                     self.ssl_verify_server_cert