zenml-nightly 0.82.0.dev20250512__py3-none-any.whl → 0.82.0.dev20250514__py3-none-any.whl

zenml/utils/pipeline_docker_image_builder.py

@@ -68,6 +68,11 @@ PIP_DEFAULT_ARGS = {
 }
 UV_DEFAULT_ARGS = {"no-cache-dir": None}
 
+DEFAULT_PYPROJECT_EXPORT_COMMANDS = [
+    "uv export --format=requirements-txt --directory={directory} --no-hashes --no-emit-project",
+    "poetry export --format=requirements.txt --directory={directory} --without-hashes",
+]
+
 
 class PipelineDockerImageBuilder:
     """Builds Docker images to run a ZenML pipeline."""
@@ -156,10 +161,11 @@ class PipelineDockerImageBuilder:
         requires_zenml_build = any(
             [
                 docker_settings.requirements,
+                docker_settings.pyproject_path,
                 docker_settings.required_integrations,
-                docker_settings.required_hub_plugins,
                 docker_settings.replicate_local_python_environment,
                 docker_settings.install_stack_requirements,
+                docker_settings.local_project_install_command,
                 docker_settings.apt_packages,
                 docker_settings.environment,
                 include_files,
@@ -434,8 +440,56 @@ class PipelineDockerImageBuilder:
             - Packages installed in the local Python environment
             - Requirements defined by stack integrations
             - Requirements defined by user integrations
+            - Requirements exported from a pyproject.toml
             - User-defined requirements
         """
+        implicit_requirements = False
+        pyproject_path = docker_settings.pyproject_path
+        requirements = docker_settings.requirements
+
+        if not any(
+            [
+                docker_settings.replicate_local_python_environment,
+                requirements,
+                pyproject_path,
+            ]
+        ):
+            root = source_utils.get_source_root()
+            requirements_path = os.path.join(root, "requirements.txt")
+            pyproject_file_path = os.path.join(root, "pyproject.toml")
+
+            if os.path.exists(requirements_path):
+                implicit_requirements = True
+                requirements = requirements_path
+            elif os.path.exists(pyproject_file_path):
+                implicit_requirements = True
+                pyproject_path = pyproject_file_path
+
+        if (
+            implicit_requirements
+            and docker_settings.disable_automatic_requirements_detection
+        ):
+            # TODO: This is only temporary to log a warning notifying users
+            # that we will soon switch the default behavior to
+            # `disable_automatic_requirements_detection=False`. Remove and
+            # adjust the logic once we've made this change.
+            if log:
+                logger.warning(
+                    "Detected `requirements.txt` or `pyproject.toml` files in "
+                    "the source root. In future versions of ZenML, these will "
+                    "be automatically picked up and installed in Docker images "
+                    "by default. To disable this behavior and keep the "
+                    "current behavior, set "
+                    "`DockerSettings.disable_automatic_requirements_detection` "
+                    "to `True`. If you want to enable this behavior right away, "
+                    "you can do so by setting "
+                    "`DockerSettings.disable_automatic_requirements_detection` "
+                    "to `False`."
+                )
+            implicit_requirements = False
+            requirements = None
+            pyproject_path = None
+
         requirements_files: List[Tuple[str, str, List[str]]] = []
 
         # Generate requirements file for the local environment if configured
@@ -447,6 +501,10 @@ class PipelineDockerImageBuilder:
                 command = (
                     docker_settings.replicate_local_python_environment.command
                 )
+            elif isinstance(
+                docker_settings.replicate_local_python_environment, bool
+            ):
+                command = PythonEnvironmentExportMethod.PIP_FREEZE.command
             else:
                 command = " ".join(
                     docker_settings.replicate_local_python_environment
@@ -520,9 +578,69 @@ class PipelineDockerImageBuilder:
                     ", ".join(f"`{r}`" for r in integration_requirements_list),
                 )
 
+        if pyproject_path:
+            path = os.path.abspath(pyproject_path)
+
+            if not os.path.exists(path):
+                raise FileNotFoundError(
+                    f"Pyproject file {path} does not exist."
+                )
+
+            def _run_command(command: str) -> str:
+                command = command.format(directory=os.path.dirname(path))
+
+                result = subprocess.run(
+                    command,
+                    capture_output=True,
+                    check=True,
+                    shell=True,  # nosec
+                )
+                return result.stdout.decode()
+
+            if docker_settings.pyproject_export_command:
+                command = " ".join(docker_settings.pyproject_export_command)
+
+                try:
+                    pyproject_requirements = _run_command(command)
+                except subprocess.CalledProcessError as e:
+                    raise RuntimeError(
+                        "Failed to export pyproject dependencies with "
+                        f"command `{command}`: {e.stderr.decode()}"
+                    )
+            else:
+                command_errors = {}
+
+                for command in DEFAULT_PYPROJECT_EXPORT_COMMANDS:
+                    try:
+                        pyproject_requirements = _run_command(command)
+                    except subprocess.CalledProcessError as e:
+                        command_errors[command] = e.stderr.decode()
+                    else:
+                        break
+                else:
+                    raise RuntimeError(
+                        "Failed to export pyproject dependencies with the "
+                        f"following commands: {command_errors}. Please specify "
+                        "a working command to export your pyproject.toml "
+                        "dependencies to a requirements.txt formatted file "
+                        "using `DockerSettings.pyproject_export_command`."
+                    )
+
+            requirements_files.append(
+                (".zenml_pyproject_requirements", pyproject_requirements, [])
+            )
+            if log:
+                logger.info(
+                    "- %s python packages from file `%s`",
+                    "Implicitly including"
+                    if implicit_requirements
+                    else "Including",
+                    path,
+                )
+
         # Generate/Read requirements file for user-defined requirements
-        if isinstance(docker_settings.requirements, str):
-            path = os.path.abspath(docker_settings.requirements)
+        if isinstance(requirements, str):
+            path = os.path.abspath(requirements)
             try:
                 user_requirements = io_utils.read_file_contents_as_string(path)
             except FileNotFoundError as e:
@@ -531,15 +649,18 @@ class PipelineDockerImageBuilder:
                 ) from e
             if log:
                 logger.info(
-                    "- Including user-defined requirements from file `%s`",
+                    "- %s user-defined requirements from file `%s`",
+                    "Implicitly including"
+                    if implicit_requirements
+                    else "Including",
                     path,
                 )
-        elif isinstance(docker_settings.requirements, List):
-            user_requirements = "\n".join(docker_settings.requirements)
+        elif isinstance(requirements, List):
+            user_requirements = "\n".join(requirements)
             if log:
                 logger.info(
                     "- Including user-defined requirements: %s",
-                    ", ".join(f"`{r}`" for r in docker_settings.requirements),
+                    ", ".join(f"`{r}`" for r in requirements),
                 )
         else:
             user_requirements = None
@@ -638,6 +759,11 @@ class PipelineDockerImageBuilder:
         lines.append("COPY . .")
         lines.append("RUN chmod -R a+rw .")
 
+        if docker_settings.local_project_install_command:
+            lines.append(
+                f"RUN {docker_settings.local_project_install_command}"
+            )
+
         if docker_settings.user:
             # Change file ownership to specified user
             lines.append(f"RUN chown -R {docker_settings.user} .")

zenml/utils/tag_utils.py

@@ -348,10 +348,10 @@ def add_tags(
             if isinstance(tag, Tag):
                 tag_model = client.get_tag(tag.name)
 
-                if tag.exclusive != tag_model.exclusive:
+                if bool(tag.exclusive) != tag_model.exclusive:
                     raise ValueError(
-                        f"The tag `{tag.name}` is an "
-                        f"{'exclusive' if tag_model.exclusive else 'non-exclusive'} "
+                        f"The tag `{tag.name}` is "
+                        f"{'an exclusive' if tag_model.exclusive else 'a non-exclusive'} "
                         "tag. Please update it before attaching it to a resource."
                     )
                 if tag.cascade is not None:

zenml/zen_server/feature_gate/endpoint_utils.py

@@ -15,47 +15,44 @@
 
 from uuid import UUID
 
-from zenml.zen_server.rbac.models import ResourceType
 from zenml.zen_server.utils import feature_gate, server_config
 
 
-def check_entitlement(resource_type: ResourceType) -> None:
+def check_entitlement(feature: str) -> None:
     """Queries the feature gate to see if the operation falls within the Pro workspaces entitlements.
 
     Raises an exception if the user is not entitled to create an instance of the
     resource. Otherwise, simply returns.
 
     Args:
-        resource_type: The type of resource to check for.
+        feature: The feature to check for.
     """
     if not server_config().feature_gate_enabled:
         return
-    return feature_gate().check_entitlement(resource=resource_type)
+    return feature_gate().check_entitlement(feature=feature)
 
 
-def report_usage(resource_type: ResourceType, resource_id: UUID) -> None:
+def report_usage(feature: str, resource_id: UUID) -> None:
     """Reports the creation/usage of a feature/resource.
 
     Args:
-        resource_type: The type of resource to report a usage for
+        feature: The feature to report a usage for.
         resource_id: ID of the resource that was created.
     """
     if not server_config().feature_gate_enabled:
         return
-    feature_gate().report_event(
-        resource=resource_type, resource_id=resource_id
-    )
+    feature_gate().report_event(feature=feature, resource_id=resource_id)
 
 
-def report_decrement(resource_type: ResourceType, resource_id: UUID) -> None:
+def report_decrement(feature: str, resource_id: UUID) -> None:
     """Reports the deletion/deactivation of a feature/resource.
 
     Args:
-        resource_type: The type of resource to report a decrement in count for.
+        feature: The feature to report a decrement in count for.
         resource_id: ID of the resource that was deleted.
     """
     if not server_config().feature_gate_enabled:
         return
     feature_gate().report_event(
-        resource=resource_type, resource_id=resource_id, is_decrement=True
+        feature=feature, resource_id=resource_id, is_decrement=True
     )

zenml/zen_server/feature_gate/feature_gate_interface.py

@@ -16,18 +16,16 @@
 from abc import ABC, abstractmethod
 from uuid import UUID
 
-from zenml.zen_server.rbac.models import ResourceType
-
 
 class FeatureGateInterface(ABC):
     """Feature gate interface definition."""
 
     @abstractmethod
-    def check_entitlement(self, resource: ResourceType) -> None:
+    def check_entitlement(self, feature: str) -> None:
         """Checks if a user is entitled to create a resource.
 
         Args:
-            resource: The resource the user wants to create
+            feature: The feature the user wants to use.
 
         Raises:
             UpgradeRequiredError in case a subscription limit is reached
@@ -36,14 +34,14 @@ class FeatureGateInterface(ABC):
     @abstractmethod
     def report_event(
         self,
-        resource: ResourceType,
+        feature: str,
         resource_id: UUID,
         is_decrement: bool = False,
     ) -> None:
         """Reports the usage of a feature to the aggregator backend.
 
         Args:
-            resource: The resource the user created
+            feature: The feature the user used.
             resource_id: ID of the resource that was created/deleted.
             is_decrement: In case this event reports an actual decrement of usage
         """

zenml/zen_server/feature_gate/zenml_cloud_feature_gate.py

@@ -25,7 +25,6 @@ from zenml.zen_server.cloud_utils import cloud_connection
 from zenml.zen_server.feature_gate.feature_gate_interface import (
     FeatureGateInterface,
 )
-from zenml.zen_server.rbac.models import ResourceType
 
 logger = get_logger(__name__)
 
@@ -47,7 +46,7 @@ class RawUsageEvent(BaseModel):
     organization_id: str = Field(
         description="The organization that this usage can be attributed to.",
     )
-    feature: ResourceType = Field(
+    feature: str = Field(
         description="The feature whose usage is being reported.",
     )
     total: int = Field(
@@ -66,22 +65,22 @@ class ZenMLCloudFeatureGateInterface(FeatureGateInterface):
         """Initialize the object."""
         self._connection = cloud_connection()
 
-    def check_entitlement(self, resource: ResourceType) -> None:
+    def check_entitlement(self, feature: str) -> None:
         """Checks if a user is entitled to create a resource.
 
         Args:
-            resource: The resource the user wants to create
+            feature: The feature the user wants to use.
 
         Raises:
             SubscriptionUpgradeRequiredError: in case a subscription limit is reached
         """
         try:
             response = self._connection.get(
-                endpoint=ENTITLEMENT_ENDPOINT + "/" + resource, params=None
+                endpoint=ENTITLEMENT_ENDPOINT + "/" + feature, params=None
             )
         except SubscriptionUpgradeRequiredError:
             raise SubscriptionUpgradeRequiredError(
-                f"Your subscription reached its `{resource}` limit. Please "
+                f"Your subscription reached its `{feature}` limit. Please "
                 f"upgrade your subscription or reach out to us."
             )
 
@@ -94,20 +93,20 @@ class ZenMLCloudFeatureGateInterface(FeatureGateInterface):
 
     def report_event(
         self,
-        resource: ResourceType,
+        feature: str,
         resource_id: UUID,
         is_decrement: bool = False,
     ) -> None:
         """Reports the usage of a feature to the aggregator backend.
 
         Args:
-            resource: The resource the user created
+            feature: The feature the user used.
             resource_id: ID of the resource that was created/deleted.
             is_decrement: In case this event reports an actual decrement of usage
         """
         data = RawUsageEvent(
             organization_id=ORGANIZATION_ID,
-            feature=resource,
+            feature=feature,
             total=1 if not is_decrement else -1,
             metadata={
                 "workspace_id": str(server_config.get_external_server_id()),

zenml/zen_server/rbac/endpoint_utils.py

@@ -182,7 +182,7 @@ def verify_permissions_and_get_or_create_entity(
     def _pre_creation_hook() -> None:
         verify_permission_for_model(model=request_model, action=Action.CREATE)
         if resource_type and needs_usage_increment:
-            check_entitlement(resource_type=resource_type)
+            check_entitlement(feature=resource_type)
 
     model, created = get_or_create_method(request_model, _pre_creation_hook)
 

zenml/zen_server/rbac/rbac_sql_zen_store.py

@@ -67,7 +67,7 @@ class RBACSqlZenStore(SqlZenStore):
                 action=Action.CREATE,
                 project_id=model_request.project,
             )
-            check_entitlement(resource_type=ResourceType.MODEL)
+            check_entitlement(feature=ResourceType.MODEL)
         except Exception as e:
             allow_model_creation = False
             error = e
@@ -92,7 +92,7 @@ class RBACSqlZenStore(SqlZenStore):
 
         if created:
             report_usage(
-                resource_type=ResourceType.MODEL, resource_id=model_response.id
+                feature=ResourceType.MODEL, resource_id=model_response.id
             )
         else:
             verify_permission_for_model(model_response, action=Action.READ)

zenml/zen_server/routers/run_templates_endpoints.py

@@ -25,7 +25,12 @@ from fastapi import (
 from zenml.analytics.enums import AnalyticsEvent
 from zenml.analytics.utils import track_handler
 from zenml.config.pipeline_run_configuration import PipelineRunConfiguration
-from zenml.constants import API, RUN_TEMPLATES, VERSION_1
+from zenml.constants import (
+    API,
+    RUN_TEMPLATE_TRIGGERS_FEATURE_NAME,
+    RUN_TEMPLATES,
+    VERSION_1,
+)
 from zenml.models import (
     Page,
     PipelineRunResponse,
@@ -36,6 +41,9 @@ from zenml.models import (
 )
 from zenml.zen_server.auth import AuthContext, authorize
 from zenml.zen_server.exceptions import error_response
+from zenml.zen_server.feature_gate.endpoint_utils import (
+    check_entitlement,
+)
 from zenml.zen_server.rbac.endpoint_utils import (
     verify_permissions_and_create_entity,
     verify_permissions_and_delete_entity,
@@ -269,6 +277,7 @@ if server_config().workload_manager_enabled:
                 action=Action.CREATE,
                 project_id=template.project.id,
             )
+            check_entitlement(feature=RUN_TEMPLATE_TRIGGERS_FEATURE_NAME)
 
             return run_template(
                 template=template,

zenml/zen_server/template_execution/utils.py

@@ -24,6 +24,7 @@ from zenml.constants import (
     ENV_ZENML_RUNNER_IMAGE_DISABLE_UV,
     ENV_ZENML_RUNNER_PARENT_IMAGE,
     ENV_ZENML_RUNNER_POD_TIMEOUT,
+    RUN_TEMPLATE_TRIGGERS_FEATURE_NAME,
     handle_bool_env_var,
     handle_int_env_var,
 )
@@ -49,7 +50,11 @@ from zenml.pipelines.run_utils import (
 )
 from zenml.stack.flavor import Flavor
 from zenml.utils import pydantic_utils, requirements_utils, settings_utils
+from zenml.utils.time_utils import utc_now
 from zenml.zen_server.auth import AuthContext, generate_access_token
+from zenml.zen_server.feature_gate.endpoint_utils import (
+    report_usage,
+)
 from zenml.zen_server.template_execution.runner_entrypoint_configuration import (
     RunnerEntrypointConfiguration,
 )
@@ -190,6 +195,11 @@ def run_template(
     placeholder_run = create_placeholder_run(deployment=new_deployment)
     assert placeholder_run
 
+    report_usage(
+        feature=RUN_TEMPLATE_TRIGGERS_FEATURE_NAME,
+        resource_id=placeholder_run.id,
+    )
+
     # We create an API token scoped to the pipeline run that never expires
     api_token = generate_access_token(
         user_id=auth_context.user.id,
@@ -300,7 +310,8 @@ def run_template(
                 zen_store().update_run(
                     run_id=placeholder_run.id,
                     run_update=PipelineRunUpdate(
-                        status=ExecutionStatus.FAILED
+                        status=ExecutionStatus.FAILED,
+                        end_time=utc_now(),
                     ),
                 )
                 raise

zenml/zen_stores/migrations/alembic.py

@@ -39,7 +39,7 @@ exclude_tables = ["sqlite_sequence"]
 
 
 def include_object(
-    object: Any, name: str, type_: str, *args: Any, **kwargs: Any
+    object: Any, name: Optional[str], type_: str, *args: Any, **kwargs: Any
 ) -> bool:
     """Function used to exclude tables from the migration scripts.
 
@@ -135,6 +135,7 @@ class Alembic:
             fn_context_args["fn"] = fn
 
         with self.engine.connect() as connection:
+            # Configure the context with our metadata
             self.environment_context.configure(
                 connection=connection,
                 target_metadata=self.metadata,
@@ -180,9 +181,15 @@ class Alembic:
         def do_get_current_rev(rev: _RevIdType, context: Any) -> List[Any]:
             nonlocal current_revisions
 
-            for r in self.script_directory.get_all_current(
-                rev  # type:ignore [arg-type]
-            ):
+            # Handle rev parameter in a way that's compatible with different alembic versions
+            rev_input: Any
+            if isinstance(rev, str):
+                rev_input = rev
+            else:
+                rev_input = tuple(str(r) for r in rev)
+
+            # Get current revision(s)
+            for r in self.script_directory.get_all_current(rev_input):
                 if r is None:
                     continue
                 current_revisions.append(r.revision)
@@ -200,7 +207,13 @@ class Alembic:
         """
 
         def do_stamp(rev: _RevIdType, context: Any) -> List[Any]:
-            return self.script_directory._stamp_revs(revision, rev)
+            # Handle rev parameter in a way that's compatible with different alembic versions
+            if isinstance(rev, str):
+                return self.script_directory._stamp_revs(revision, rev)
+            else:
+                # Convert to tuple for compatibility
+                rev_tuple = tuple(str(r) for r in rev)
+                return self.script_directory._stamp_revs(revision, rev_tuple)
 
         self.run_migrations(do_stamp)
 
@@ -212,10 +225,16 @@ class Alembic:
         """
 
         def do_upgrade(rev: _RevIdType, context: Any) -> List[Any]:
-            return self.script_directory._upgrade_revs(
-                revision,
-                rev,  # type:ignore [arg-type]
-            )
+            # Handle rev parameter in a way that's compatible with different alembic versions
+            if isinstance(rev, str):
+                return self.script_directory._upgrade_revs(revision, rev)
+            else:
+                if rev:
+                    # Use first element or revs for compatibility
+                    return self.script_directory._upgrade_revs(
+                        revision, str(rev[0])
+                    )
+                return []
 
         self.run_migrations(do_upgrade)
 
@@ -227,9 +246,14 @@ class Alembic:
         """
 
         def do_downgrade(rev: _RevIdType, context: Any) -> List[Any]:
-            return self.script_directory._downgrade_revs(
-                revision,
-                rev,  # type:ignore [arg-type]
-            )
+            # Handle rev parameter in a way that's compatible with different alembic versions
+            if isinstance(rev, str):
+                return self.script_directory._downgrade_revs(revision, rev)
+            else:
+                if rev:
+                    return self.script_directory._downgrade_revs(
+                        revision, str(rev[0])
+                    )
+                return self.script_directory._downgrade_revs(revision, None)
 
         self.run_migrations(do_downgrade)

zenml/zen_stores/rest_zen_store.py

@@ -4413,11 +4413,15 @@ class RestZenStore(BaseZenStore):
                         "The current token is no longer valid, and "
                         "it is not possible to generate a new token using the "
                         "configured credentials. Please run "
-                        f"`zenml login --url {self.url}` to re-authenticate to "
-                        "the server or authenticate using an API key. See "
+                        f"`zenml login {self.url}` to "
+                        "re-authenticate to the server or authenticate using "
+                        "an API key. See "
                         "https://docs.zenml.io/how-to/project-setup-and-management/connecting-to-zenml/connect-with-a-service-account "
                         "for more information."
                     )
+                    # Clear the current token from the credentials store to
+                    # force a new authentication flow next time.
+                    get_credentials_store().clear_token(self.url)
                     raise e
                 elif not re_authenticated:
                     # The last request was authenticated with an API token

zenml/zen_stores/sql_zen_store.py

@@ -11354,13 +11354,10 @@ class SqlZenStore(BaseZenStore):
             except EntityExistsError:
                 if isinstance(tag, tag_utils.Tag):
                     tag_schema = self._get_tag_schema(tag.name, session)
-                    if (
-                        tag.exclusive is not None
-                        and tag.exclusive != tag_schema.exclusive
-                    ):
+                    if bool(tag.exclusive) != tag_schema.exclusive:
                         raise ValueError(
-                            f"Tag `{tag_schema.name}` has been defined as a "
-                            f"{'exclusive' if tag_schema.exclusive else 'non-exclusive'} "
+                            f"Tag `{tag_schema.name}` has been defined as "
+                            f"{'an exclusive' if tag_schema.exclusive else 'a non-exclusive'} "
                             "tag. Please update it before attaching it to resources."
                         )
                 else:

{zenml_nightly-0.82.0.dev20250512.dist-info → zenml_nightly-0.82.0.dev20250514.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.3
 Name: zenml-nightly
-Version: 0.82.0.dev20250512
+Version: 0.82.0.dev20250514
 Summary: ZenML: Write production-ready ML code.
 License: Apache-2.0
 Keywords: machine learning,production,pipeline,mlops,devops
@@ -41,7 +41,7 @@ Provides-Extra: terraform
 Provides-Extra: vertex
 Requires-Dist: Jinja2 ; extra == "server"
 Requires-Dist: adlfs (>=2021.10.0) ; extra == "adlfs"
-Requires-Dist: alembic (>=1.8.1,<1.9.0)
+Requires-Dist: alembic (>=1.8.1,<=1.15.2)
 Requires-Dist: aws-profile-manager (>=0.5.0) ; extra == "connectors-aws"
 Requires-Dist: azure-ai-ml (==1.23.1) ; extra == "azureml"
 Requires-Dist: azure-identity (>=1.4.0) ; extra == "secrets-azure" or extra == "connectors-azure"