zenml-nightly 0.73.0.dev20250128__tar.gz → 0.73.0.dev20250130__tar.gz

{zenml_nightly-0.73.0.dev20250128 → zenml_nightly-0.73.0.dev20250130}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.3
 Name: zenml-nightly
-Version: 0.73.0.dev20250128
+Version: 0.73.0.dev20250130
 Summary: ZenML: Write production-ready ML code.
 License: Apache-2.0
 Keywords: machine learning,production,pipeline,mlops,devops

{zenml_nightly-0.73.0.dev20250128 → zenml_nightly-0.73.0.dev20250130}/pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "zenml-nightly"
-version = "0.73.0.dev20250128"
+version = "0.73.0.dev20250130"
 packages = [{ include = "zenml", from = "src" }]
 description = "ZenML: Write production-ready ML code."
 authors = ["ZenML GmbH <info@zenml.io>"]

zenml_nightly-0.73.0.dev20250130/src/zenml/VERSION

@@ -0,0 +1 @@
+0.73.0.dev20250130

{zenml_nightly-0.73.0.dev20250128 → zenml_nightly-0.73.0.dev20250130}/src/zenml/cli/utils.py

@@ -76,6 +76,7 @@ from zenml.models import (
 from zenml.models.v2.base.filter import FilterGenerator
 from zenml.services import BaseService, ServiceState
 from zenml.stack import StackComponent
+from zenml.stack.flavor import Flavor
 from zenml.stack.stack_component import StackComponentConfig
 from zenml.utils import secret_utils
 from zenml.utils.time_utils import expires_in
@@ -2151,10 +2152,11 @@ def _scrub_secret(config: StackComponentConfig) -> Dict[str, Any]:
     config_dict = {}
     config_fields = config.__class__.model_fields
     for key, value in config_fields.items():
-        if secret_utils.is_secret_field(value):
-            config_dict[key] = "********"
-        else:
-            config_dict[key] = getattr(config, key)
+        if getattr(config, key):
+            if secret_utils.is_secret_field(value):
+                config_dict[key] = "********"
+            else:
+                config_dict[key] = getattr(config, key)
     return config_dict
 
 
@@ -2164,8 +2166,6 @@ def print_debug_stack() -> None:
 
     client = Client()
     stack = client.get_stack()
-    active_stack = client.active_stack
-    components = _get_stack_components(active_stack)
 
     declare("\nCURRENT STACK\n", bold=True)
     console.print(f"Name: {stack.name}")
@@ -2176,7 +2176,8 @@ def print_debug_stack() -> None:
         f"Workspace: {stack.workspace.name} / {str(stack.workspace.id)}"
     )
 
-    for component in components:
+    for component_type, components in stack.components.items():
+        component = components[0]
         component_response = client.get_stack_component(
             name_id_or_prefix=component.id, component_type=component.type
         )
@@ -2186,8 +2187,12 @@ def print_debug_stack() -> None:
         console.print(f"Name: {component.name}")
         console.print(f"ID: {str(component.id)}")
         console.print(f"Type: {component.type.value}")
-        console.print(f"Flavor: {component.flavor}")
-        console.print(f"Configuration: {_scrub_secret(component.config)}")
+        console.print(f"Flavor: {component.flavor_name}")
+
+        flavor = Flavor.from_model(component.flavor)
+        config = flavor.config_class(**component.configuration)
+
+        console.print(f"Configuration: {_scrub_secret(config)}")
         if (
             component_response.user
             and component_response.user.name

{zenml_nightly-0.73.0.dev20250128 → zenml_nightly-0.73.0.dev20250130}/src/zenml/client.py

@@ -2453,7 +2453,9 @@ class Client(metaclass=ClientMetaClass):
     def trigger_pipeline(
         self,
         pipeline_name_or_id: Union[str, UUID, None] = None,
-        run_configuration: Optional[PipelineRunConfiguration] = None,
+        run_configuration: Union[
+            PipelineRunConfiguration, Dict[str, Any], None
+        ] = None,
         config_path: Optional[str] = None,
         template_id: Optional[UUID] = None,
         stack_name_or_id: Union[str, UUID, None] = None,
@@ -2524,6 +2526,11 @@ class Client(metaclass=ClientMetaClass):
         if config_path:
             run_configuration = PipelineRunConfiguration.from_yaml(config_path)
 
+        if isinstance(run_configuration, Dict):
+            run_configuration = PipelineRunConfiguration.model_validate(
+                run_configuration
+            )
+
         if run_configuration:
             validate_run_config_is_runnable_from_server(run_configuration)
 

{zenml_nightly-0.73.0.dev20250128 → zenml_nightly-0.73.0.dev20250130}/src/zenml/constants.py

@@ -138,6 +138,7 @@ ENV_ZENML_USER_ID = "ZENML_USER_ID"
 ENV_ZENML_CONFIG_PATH = "ZENML_CONFIG_PATH"
 ENV_ZENML_DEBUG = "ZENML_DEBUG"
 ENV_ZENML_LOGGING_VERBOSITY = "ZENML_LOGGING_VERBOSITY"
+ENV_ZENML_LOGGING_FORMAT = "ZENML_LOGGING_FORMAT"
 ENV_ZENML_REPOSITORY_PATH = "ZENML_REPOSITORY_PATH"
 ENV_ZENML_PREVENT_PIPELINE_EXECUTION = "ZENML_PREVENT_PIPELINE_EXECUTION"
 ENV_ZENML_ENABLE_RICH_TRACEBACK = "ZENML_ENABLE_RICH_TRACEBACK"
@@ -154,6 +155,7 @@ ENV_ZENML_SECRETS_STORE_PREFIX = "ZENML_SECRETS_STORE_"
 ENV_ZENML_BACKUP_SECRETS_STORE_PREFIX = "ZENML_BACKUP_SECRETS_STORE_"
 ENV_ZENML_SKIP_PIPELINE_REGISTRATION = "ZENML_SKIP_PIPELINE_REGISTRATION"
 ENV_AUTO_OPEN_DASHBOARD = "AUTO_OPEN_DASHBOARD"
+ENV_ZENML_AUTO_OPEN_DASHBOARD = "ZENML_AUTO_OPEN_DASHBOARD"
 ENV_ZENML_DISABLE_DATABASE_MIGRATION = "DISABLE_DATABASE_MIGRATION"
 ENV_ZENML_LOCAL_STORES_PATH = "ZENML_LOCAL_STORES_PATH"
 ENV_ZENML_CONTAINER = "ZENML_CONTAINER"

{zenml_nightly-0.73.0.dev20250128 → zenml_nightly-0.73.0.dev20250130}/src/zenml/entrypoints/base_entrypoint_configuration.py

@@ -22,6 +22,8 @@ from uuid import UUID
 
 from zenml.client import Client
 from zenml.code_repositories import BaseCodeRepository
+from zenml.enums import StackComponentType
+from zenml.exceptions import CustomFlavorImportError
 from zenml.logger import get_logger
 from zenml.utils import (
     code_repository_utils,
@@ -31,6 +33,7 @@ from zenml.utils import (
 )
 
 if TYPE_CHECKING:
+    from zenml.artifact_stores import BaseArtifactStore
     from zenml.models import CodeReferenceResponse, PipelineDeploymentResponse
 
 logger = get_logger(__name__)
@@ -204,6 +207,8 @@ class BaseEntrypointConfiguration(ABC):
                 decision instead.
 
         Raises:
+            CustomFlavorImportError: If the artifact store flavor can't be
+                imported.
             RuntimeError: If the current environment requires code download
                 but the deployment does not have a reference to any code.
         """
@@ -215,7 +220,27 @@ class BaseEntrypointConfiguration(ABC):
             return
 
         if code_path := deployment.code_path:
-            code_utils.download_code_from_artifact_store(code_path=code_path)
+            # Load the artifact store not from the active stack but separately.
+            # This is required in case the stack has custom flavor components
+            # (other than the artifact store) for which the flavor
+            # implementations will only be available once the download finishes.
+            try:
+                artifact_store = self._load_active_artifact_store()
+            except CustomFlavorImportError as e:
+                raise CustomFlavorImportError(
+                    "Failed to import custom artifact store flavor. The "
+                    "artifact store flavor is needed to download your code, "
+                    "but it looks like it might be part of the files "
+                    "that we're trying to download. If this is the case, you "
+                    "should disable downloading code from the artifact store "
+                    "using `DockerSettings(allow_download_from_artifact_store=False)` "
+                    "or make sure the artifact flavor files are included in "
+                    "Docker image by using a custom parent image or installing "
+                    "them as part of a pip dependency."
+                ) from e
+            code_utils.download_code_from_artifact_store(
+                code_path=code_path, artifact_store=artifact_store
+            )
         elif code_reference := deployment.code_reference:
             # TODO: This might fail if the code repository had unpushed changes
             # at the time the pipeline run was started.
@@ -300,6 +325,22 @@ class BaseEntrypointConfiguration(ABC):
 
         return False
 
+    def _load_active_artifact_store(self) -> "BaseArtifactStore":
+        """Load the active artifact store.
+
+        Returns:
+            The active artifact store.
+        """
+        from zenml.artifact_stores import BaseArtifactStore
+
+        artifact_store_model = Client().active_stack_model.components[
+            StackComponentType.ARTIFACT_STORE
+        ][0]
+        artifact_store = BaseArtifactStore.from_model(artifact_store_model)
+        assert isinstance(artifact_store, BaseArtifactStore)
+
+        return artifact_store
+
     @abstractmethod
     def run(self) -> None:
         """Runs the entrypoint configuration."""

{zenml_nightly-0.73.0.dev20250128 → zenml_nightly-0.73.0.dev20250130}/src/zenml/exceptions.py

@@ -314,3 +314,7 @@ class SecretsStoreNotConfiguredError(NotImplementedError):
 
 class BackupSecretsStoreNotConfiguredError(NotImplementedError):
     """Raised when a backup secrets store is not configured."""
+
+
+class CustomFlavorImportError(ImportError):
+    """Raised when failing to import a custom flavor."""

{zenml_nightly-0.73.0.dev20250128 → zenml_nightly-0.73.0.dev20250130}/src/zenml/integrations/gcp/constants.py

@@ -16,7 +16,7 @@
 from google.cloud.aiplatform_v1.types.job_state import JobState
 
 VERTEX_ENDPOINT_SUFFIX = "-aiplatform.googleapis.com"
-POLLING_INTERVAL_IN_SECONDS = 30
+POLLING_INTERVAL_IN_SECONDS = 10
 CONNECTION_ERROR_RETRY_LIMIT = 5
 _VERTEX_JOB_STATE_SUCCEEDED = JobState.JOB_STATE_SUCCEEDED
 _VERTEX_JOB_STATE_FAILED = JobState.JOB_STATE_FAILED

{zenml_nightly-0.73.0.dev20250128 → zenml_nightly-0.73.0.dev20250130}/src/zenml/integrations/gcp/flavors/vertex_step_operator_flavor.py

@@ -51,7 +51,8 @@ class VertexStepOperatorSettings(BaseSettings):
             https://cloud.google.com/vertex-ai/docs/training/configure-compute#boot_disk_options
         boot_disk_type: Type of the boot disk. (Default: pd-ssd)
             https://cloud.google.com/vertex-ai/docs/training/configure-compute#boot_disk_options
-
+        persistent_resource_id: The ID of the persistent resource to use for the job.
+            https://cloud.google.com/vertex-ai/docs/training/persistent-resource-overview
     """
 
     accelerator_type: Optional[str] = None
@@ -59,6 +60,7 @@ class VertexStepOperatorSettings(BaseSettings):
     machine_type: str = "n1-standard-4"
     boot_disk_size_gb: int = 100
     boot_disk_type: str = "pd-ssd"
+    persistent_resource_id: Optional[str] = None
 
 
 class VertexStepOperatorConfig(

{zenml_nightly-0.73.0.dev20250128 → zenml_nightly-0.73.0.dev20250130}/src/zenml/integrations/gcp/step_operators/vertex_step_operator.py

@@ -258,6 +258,7 @@ class VertexStepOperator(BaseStepOperator, GoogleCredentialsMixin):
                     if self.config.reserved_ip_ranges
                     else []
                 ),
+                "persistent_resource_id": settings.persistent_resource_id,
             },
             "labels": job_labels,
             "encryption_spec": {

{zenml_nightly-0.73.0.dev20250128 → zenml_nightly-0.73.0.dev20250130}/src/zenml/integrations/huggingface/services/huggingface_deployment.py

@@ -13,17 +13,18 @@
 #  permissions and limitations under the License.
 """Implementation of the Hugging Face Deployment service."""
 
-from typing import Any, Generator, Optional, Tuple
+from typing import Any, Dict, Generator, Optional, Tuple
 
 from huggingface_hub import (
     InferenceClient,
     InferenceEndpoint,
     InferenceEndpointError,
     InferenceEndpointStatus,
+    InferenceEndpointType,
     create_inference_endpoint,
     get_inference_endpoint,
 )
-from huggingface_hub.utils import HfHubHTTPError
+from huggingface_hub.errors import HfHubHTTPError
 from pydantic import Field
 
 from zenml.client import Client
@@ -138,30 +139,67 @@ class HuggingFaceDeploymentService(BaseDeploymentService):
         """
         return self.hf_endpoint.client
 
+    def _validate_endpoint_configuration(self) -> Dict[str, str]:
+        """Validates the configuration to provision a Huggingface service.
+
+        Raises:
+            ValueError: if there is a missing value in the configuration
+
+        Returns:
+            The validated configuration values.
+        """
+        configuration = {}
+        missing_keys = []
+
+        for k, v in {
+            "repository": self.config.repository,
+            "framework": self.config.framework,
+            "accelerator": self.config.accelerator,
+            "instance_size": self.config.instance_size,
+            "instance_type": self.config.instance_type,
+            "region": self.config.region,
+            "vendor": self.config.vendor,
+            "endpoint_type": self.config.endpoint_type,
+        }.items():
+            if v is None:
+                missing_keys.append(k)
+            else:
+                configuration[k] = v
+
+        if missing_keys:
+            raise ValueError(
+                f"Missing values in the Huggingface Service "
+                f"configuration: {', '.join(missing_keys)}"
+            )
+
+        return configuration
+
     def provision(self) -> None:
         """Provision or update remote Hugging Face deployment instance.
 
         Raises:
-            Exception: If any unexpected error while creating inference endpoint.
+            Exception: If any unexpected error while creating inference
+                endpoint.
         """
         try:
-            # Attempt to create and wait for the inference endpoint
+            validated_config = self._validate_endpoint_configuration()
+
             hf_endpoint = create_inference_endpoint(
                 name=self._generate_an_endpoint_name(),
-                repository=self.config.repository,
-                framework=self.config.framework,
-                accelerator=self.config.accelerator,
-                instance_size=self.config.instance_size,
-                instance_type=self.config.instance_type,
-                region=self.config.region,
-                vendor=self.config.vendor,
+                repository=validated_config["repository"],
+                framework=validated_config["framework"],
+                accelerator=validated_config["accelerator"],
+                instance_size=validated_config["instance_size"],
+                instance_type=validated_config["instance_type"],
+                region=validated_config["region"],
+                vendor=validated_config["vendor"],
                 account_id=self.config.account_id,
                 min_replica=self.config.min_replica,
                 max_replica=self.config.max_replica,
                 revision=self.config.revision,
                 task=self.config.task,
                 custom_image=self.config.custom_image,
-                type=self.config.endpoint_type,
+                type=InferenceEndpointType(validated_config["endpoint_type"]),
                 token=self.get_token(),
                 namespace=self.config.namespace,
             ).wait(timeout=POLLING_TIMEOUT)
@@ -172,21 +210,25 @@ class HuggingFaceDeploymentService(BaseDeploymentService):
             )
             # Catch-all for any other unexpected errors
             raise Exception(
-                f"An unexpected error occurred while provisioning the Hugging Face inference endpoint: {e}"
+                "An unexpected error occurred while provisioning the "
+                f"Hugging Face inference endpoint: {e}"
             )
 
         # Check if the endpoint URL is available after provisioning
         if hf_endpoint.url:
             logger.info(
-                f"Hugging Face inference endpoint successfully deployed and available. Endpoint URL: {hf_endpoint.url}"
+                "Hugging Face inference endpoint successfully deployed "
+                f"and available. Endpoint URL: {hf_endpoint.url}"
             )
         else:
             logger.error(
-                "Failed to start Hugging Face inference endpoint service: No URL available, please check the Hugging Face console for more details."
+                "Failed to start Hugging Face inference endpoint "
+                "service: No URL available, please check the Hugging "
+                "Face console for more details."
             )
 
     def check_status(self) -> Tuple[ServiceState, str]:
-        """Check the the current operational state of the Hugging Face deployment.
+        """Check the current operational state of the Hugging Face deployment.
 
         Returns:
             The operational state of the Hugging Face deployment and a message
@@ -196,26 +238,29 @@ class HuggingFaceDeploymentService(BaseDeploymentService):
         try:
             status = self.hf_endpoint.status
             if status == InferenceEndpointStatus.RUNNING:
-                return (ServiceState.ACTIVE, "")
+                return ServiceState.ACTIVE, ""
 
             elif status == InferenceEndpointStatus.SCALED_TO_ZERO:
                 return (
                     ServiceState.SCALED_TO_ZERO,
-                    "Hugging Face Inference Endpoint is scaled to zero, but still running. It will be started on demand.",
+                    "Hugging Face Inference Endpoint is scaled to zero, but "
+                    "still running. It will be started on demand.",
                 )
 
             elif status == InferenceEndpointStatus.FAILED:
                 return (
                     ServiceState.ERROR,
-                    "Hugging Face Inference Endpoint deployment is inactive or not found",
+                    "Hugging Face Inference Endpoint deployment is inactive "
+                    "or not found",
                 )
             elif status == InferenceEndpointStatus.PENDING:
-                return (ServiceState.PENDING_STARTUP, "")
-            return (ServiceState.PENDING_STARTUP, "")
+                return ServiceState.PENDING_STARTUP, ""
+            return ServiceState.PENDING_STARTUP, ""
         except (InferenceEndpointError, HfHubHTTPError):
             return (
                 ServiceState.INACTIVE,
-                "Hugging Face Inference Endpoint deployment is inactive or not found",
+                "Hugging Face Inference Endpoint deployment is inactive or "
+                "not found",
             )
 
     def deprovision(self, force: bool = False) -> None:
@@ -253,15 +298,13 @@ class HuggingFaceDeploymentService(BaseDeploymentService):
             )
         if self.prediction_url is not None:
             if self.hf_endpoint.task == "text-generation":
-                result = self.inference_client.task_generation(
+                return self.inference_client.text_generation(
                     data, max_new_tokens=max_new_tokens
                 )
-        else:
-            # TODO: Add support for all different supported tasks
-            raise NotImplementedError(
-                "Tasks other than text-generation is not implemented."
-            )
-        return result
+        # TODO: Add support for all different supported tasks
+        raise NotImplementedError(
+            "Tasks other than text-generation is not implemented."
+        )
 
     def get_logs(
         self, follow: bool = False, tail: Optional[int] = None

{zenml_nightly-0.73.0.dev20250128 → zenml_nightly-0.73.0.dev20250130}/src/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_nightly-0.73.0.dev20250128 → zenml_nightly-0.73.0.dev20250130}/src/zenml/logger.py

@@ -24,6 +24,7 @@ from rich.traceback import install as rich_tb_install
 from zenml.constants import (
     ENABLE_RICH_TRACEBACK,
     ENV_ZENML_LOGGING_COLORS_DISABLED,
+    ENV_ZENML_LOGGING_FORMAT,
     ENV_ZENML_SUPPRESS_LOGS,
     ZENML_LOGGING_VERBOSITY,
     handle_bool_env_var,
@@ -144,6 +145,18 @@ def set_root_verbosity() -> None:
         get_logger(__name__).debug("Logging NOTSET")
 
 
+def get_formatter() -> logging.Formatter:
+    """Get a configured logging formatter.
+
+    Returns:
+        The formatter.
+    """
+    if log_format := os.environ.get(ENV_ZENML_LOGGING_FORMAT, None):
+        return logging.Formatter(fmt=log_format)
+    else:
+        return CustomFormatter()
+
+
 def get_console_handler() -> Any:
     """Get console handler for logging.
 
@@ -151,7 +164,7 @@ def get_console_handler() -> Any:
         A console handler.
     """
     console_handler = logging.StreamHandler(sys.stdout)
-    console_handler.setFormatter(CustomFormatter())
+    console_handler.setFormatter(get_formatter())
     return console_handler
 
 
@@ -179,7 +192,7 @@ def init_logging() -> None:
     set_root_verbosity()
 
     console_handler = logging.StreamHandler(sys.stdout)
-    console_handler.setFormatter(CustomFormatter())
+    console_handler.setFormatter(get_formatter())
     logging.root.addHandler(console_handler)
 
     # Enable logs if environment variable SUPPRESS_ZENML_LOGS is not set to True

{zenml_nightly-0.73.0.dev20250128 → zenml_nightly-0.73.0.dev20250130}/src/zenml/models/v2/base/scoped.py

@@ -241,8 +241,9 @@ class UserScopedFilter(BaseFilter):
         if sort_by == "user":
             column = UserSchema.name
 
-            query = query.join(
-                UserSchema, getattr(table, "user_id") == UserSchema.id
+            query = query.outerjoin(
+                UserSchema,
+                getattr(table, "user_id") == UserSchema.id,
             )
 
             query = query.add_columns(UserSchema.name)

{zenml_nightly-0.73.0.dev20250128 → zenml_nightly-0.73.0.dev20250130}/src/zenml/models/v2/core/pipeline_run.py

@@ -950,31 +950,31 @@ class PipelineRunFilter(WorkspaceScopedTaggableFilter):
         sort_by, operand = self.sorting_params
 
         if sort_by == "pipeline":
-            query = query.join(
+            query = query.outerjoin(
                 PipelineSchema,
                 PipelineRunSchema.pipeline_id == PipelineSchema.id,
             )
             column = PipelineSchema.name
         elif sort_by == "stack":
-            query = query.join(
+            query = query.outerjoin(
                 PipelineDeploymentSchema,
                 PipelineRunSchema.deployment_id == PipelineDeploymentSchema.id,
-            ).join(
+            ).outerjoin(
                 StackSchema,
                 PipelineDeploymentSchema.stack_id == StackSchema.id,
             )
             column = StackSchema.name
         elif sort_by == "model":
-            query = query.join(
+            query = query.outerjoin(
                 ModelVersionSchema,
                 PipelineRunSchema.model_version_id == ModelVersionSchema.id,
-            ).join(
+            ).outerjoin(
                 ModelSchema,
                 ModelVersionSchema.model_id == ModelSchema.id,
             )
             column = ModelSchema.name
         elif sort_by == "model_version":
-            query = query.join(
+            query = query.outerjoin(
                 ModelVersionSchema,
                 PipelineRunSchema.model_version_id == ModelVersionSchema.id,
             )

{zenml_nightly-0.73.0.dev20250128 → zenml_nightly-0.73.0.dev20250130}/src/zenml/stack/flavor.py

@@ -13,11 +13,13 @@
 #  permissions and limitations under the License.
 """Base ZenML Flavor implementation."""
 
+import os
 from abc import abstractmethod
 from typing import Any, Dict, Optional, Type, cast
 
 from zenml.client import Client
 from zenml.enums import StackComponentType
+from zenml.exceptions import CustomFlavorImportError
 from zenml.models import (
     FlavorRequest,
     FlavorResponse,
@@ -126,10 +128,34 @@ class Flavor:
         Args:
             flavor_model: The model to load from.
 
+        Raises:
+            CustomFlavorImportError: If the custom flavor can't be imported.
+            ImportError: If the flavor can't be imported.
+
         Returns:
             The loaded flavor.
         """
-        flavor = source_utils.load(flavor_model.source)()
+        try:
+            flavor = source_utils.load(flavor_model.source)()
+        except (ModuleNotFoundError, ImportError, NotImplementedError) as err:
+            if flavor_model.is_custom:
+                flavor_module, _ = flavor_model.source.rsplit(".")
+                expected_file_path = os.path.join(
+                    source_utils.get_source_root(),
+                    flavor_module.replace(".", os.path.sep),
+                )
+                raise CustomFlavorImportError(
+                    f"Couldn't import custom flavor {flavor_model.name}: "
+                    f"{err}. Make sure the custom flavor class "
+                    f"`{flavor_model.source}` is importable. If it is part of "
+                    "a library, make sure it is installed. If "
+                    "it is a local code file, make sure it exists at "
+                    f"`{expected_file_path}.py`."
+                )
+            else:
+                raise ImportError(
+                    f"Couldn't import flavor {flavor_model.name}: {err}"
+                )
         return cast(Flavor, flavor)
 
     def to_model(

{zenml_nightly-0.73.0.dev20250128 → zenml_nightly-0.73.0.dev20250130}/src/zenml/stack/stack_component.py

@@ -402,16 +402,10 @@ class StackComponent:
         Raises:
             ImportError: If the flavor can't be imported.
         """
-        flavor_model = component_model.flavor
-
-        try:
-            from zenml.stack import Flavor
+        from zenml.stack import Flavor
 
-            flavor = Flavor.from_model(flavor_model)
-        except (ModuleNotFoundError, ImportError, NotImplementedError) as err:
-            raise ImportError(
-                f"Couldn't import flavor {flavor_model.name}: {err}"
-            )
+        flavor_model = component_model.flavor
+        flavor = Flavor.from_model(flavor_model)
 
         configuration = flavor.config_class(**component_model.configuration)