zenml-nightly 0.66.0.dev20240928__py3-none-any.whl → 0.68.1.dev20241102__py3-none-any.whl

zenml/integrations/bentoml/services/{bentoml_deployment.py → bentoml_local_deployment.py}

@@ -11,16 +11,17 @@
 #  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.
-"""Implementation for the BentoML inference service."""
+"""Implementation for the BentoML local deployment service."""
 
 import os
 from typing import TYPE_CHECKING, Any, Dict, List, Optional, Union
 
-from bentoml.client import Client
+from bentoml import AsyncHTTPClient, SyncHTTPClient
 from pydantic import BaseModel, Field
 
 from zenml.constants import DEFAULT_LOCAL_SERVICE_IP_ADDRESS
 from zenml.integrations.bentoml.constants import (
+    BENTOML_DEFAULT_PORT,
     BENTOML_HEALTHCHECK_URL_PATH,
     BENTOML_PREDICTION_URL_PATH,
 )
@@ -44,6 +45,8 @@ if TYPE_CHECKING:
 
 logger = get_logger(__name__)
 
+BENTOML_LOCAL_DEPLOYMENT_SERVICE_NAME = "bentoml-local-deployment"
+
 
 class BentoMLDeploymentEndpointConfig(LocalDaemonServiceEndpointConfig):
     """BentoML deployment service configuration.
@@ -100,14 +103,15 @@ class SSLBentoMLParametersConfig(BaseModel):
     ssl_ciphers: Optional[str] = None
 
 
-class BentoMLDeploymentConfig(LocalDaemonServiceConfig):
+class BentoMLLocalDeploymentConfig(LocalDaemonServiceConfig):
     """BentoML model deployment configuration.
 
     Attributes:
         model_name: name of the model to deploy
         model_uri: URI of the model to deploy
         port: port to expose the service on
-        bento: Bento package to deploy
+        bento_tag: Bento package to deploy. A bento tag is a combination of the
+            name of the bento and its version.
         workers: number of workers to use
         backlog: number of requests to queue
         production: whether to run in production mode
@@ -118,11 +122,11 @@ class BentoMLDeploymentConfig(LocalDaemonServiceConfig):
 
     model_name: str
     model_uri: str
-    bento: str
+    bento_tag: str
     bento_uri: Optional[str] = None
     apis: List[str] = []
     workers: int = 1
-    port: int
+    port: Optional[int] = None
     backlog: int = 2048
     production: bool = False
     working_dir: str
@@ -132,7 +136,7 @@ class BentoMLDeploymentConfig(LocalDaemonServiceConfig):
     )
 
 
-class BentoMLDeploymentService(LocalDaemonService, BaseDeploymentService):
+class BentoMLLocalDeploymentService(LocalDaemonService, BaseDeploymentService):
     """BentoML deployment service used to start a local prediction server for BentoML models.
 
     Attributes:
@@ -143,19 +147,19 @@ class BentoMLDeploymentService(LocalDaemonService, BaseDeploymentService):
     """
 
     SERVICE_TYPE = ServiceType(
-        name="bentoml-deployment",
+        name=BENTOML_LOCAL_DEPLOYMENT_SERVICE_NAME,
         type="model-serving",
         flavor="bentoml",
-        description="BentoML prediction service",
+        description="BentoML local prediction service",
         logo_url="https://public-flavor-logos.s3.eu-central-1.amazonaws.com/model_deployer/bentoml.png",
     )
 
-    config: BentoMLDeploymentConfig
+    config: BentoMLLocalDeploymentConfig
     endpoint: BentoMLDeploymentEndpoint
 
     def __init__(
         self,
-        config: Union[BentoMLDeploymentConfig, Dict[str, Any]],
+        config: Union[BentoMLLocalDeploymentConfig, Dict[str, Any]],
         **attrs: Any,
     ) -> None:
         """Initialize the BentoML deployment service.
@@ -168,13 +172,15 @@ class BentoMLDeploymentService(LocalDaemonService, BaseDeploymentService):
         # TODO [ENG-700]: implement a service factory or builder for BentoML
         #   deployment services
         if (
-            isinstance(config, BentoMLDeploymentConfig)
+            isinstance(config, BentoMLLocalDeploymentConfig)
             and "endpoint" not in attrs
         ):
             endpoint = BentoMLDeploymentEndpoint(
                 config=BentoMLDeploymentEndpointConfig(
                     protocol=ServiceEndpointProtocol.HTTP,
-                    port=config.port,
+                    port=config.port
+                    if config.port is not None
+                    else BENTOML_DEFAULT_PORT,
                     ip_address=config.host or DEFAULT_LOCAL_SERVICE_IP_ADDRESS,
                     prediction_url_path=BENTOML_PREDICTION_URL_PATH,
                 ),
@@ -189,6 +195,9 @@ class BentoMLDeploymentService(LocalDaemonService, BaseDeploymentService):
 
     def run(self) -> None:
         """Start the service."""
+        from bentoml import Service
+        from bentoml._internal.service.loader import load
+
         logger.info(
             "Starting BentoML prediction service as blocking "
             "process... press CTRL+C once to stop it."
@@ -198,26 +207,55 @@ class BentoMLDeploymentService(LocalDaemonService, BaseDeploymentService):
         ssl_params = self.config.ssl_parameters or SSLBentoMLParametersConfig()
         # verify if to deploy in production mode or development mode
         logger.info("Running in production mode.")
-        from bentoml.serve import serve_http_production
-
-        try:
-            serve_http_production(
-                self.config.bento,
-                working_dir=self.config.working_dir,
-                port=self.config.port,
-                api_workers=self.config.workers,
-                host=self.config.host or DEFAULT_LOCAL_SERVICE_IP_ADDRESS,
-                backlog=self.config.backlog,
-                ssl_certfile=ssl_params.ssl_certfile,
-                ssl_keyfile=ssl_params.ssl_keyfile,
-                ssl_keyfile_password=ssl_params.ssl_keyfile_password,
-                ssl_version=ssl_params.ssl_version,
-                ssl_cert_reqs=ssl_params.ssl_cert_reqs,
-                ssl_ca_certs=ssl_params.ssl_ca_certs,
-                ssl_ciphers=ssl_params.ssl_ciphers,
-            )
-        except KeyboardInterrupt:
-            logger.info("Stopping BentoML prediction service...")
+        svc = load(
+            bento_identifier=self.config.bento_tag,
+            working_dir=self.config.working_dir or ".",
+        )
+
+        if isinstance(svc, Service):
+            # bentoml<1.2
+            from bentoml.serving import serve_http_production
+
+            try:
+                serve_http_production(
+                    self.config.bento_tag,
+                    working_dir=self.config.working_dir,
+                    port=self.config.port,
+                    api_workers=self.config.workers,
+                    host=self.config.host or DEFAULT_LOCAL_SERVICE_IP_ADDRESS,
+                    backlog=self.config.backlog,
+                    ssl_certfile=ssl_params.ssl_certfile,
+                    ssl_keyfile=ssl_params.ssl_keyfile,
+                    ssl_keyfile_password=ssl_params.ssl_keyfile_password,
+                    ssl_version=ssl_params.ssl_version,
+                    ssl_cert_reqs=ssl_params.ssl_cert_reqs,
+                    ssl_ca_certs=ssl_params.ssl_ca_certs,
+                    ssl_ciphers=ssl_params.ssl_ciphers,
+                )
+            except KeyboardInterrupt:
+                logger.info("Stopping BentoML prediction service...")
+        else:
+            # bentoml>=1.2
+            from _bentoml_impl.server import serve_http  # type: ignore
+
+            svc.inject_config()
+            try:
+                serve_http(
+                    self.config.bento_tag,
+                    working_dir=self.config.working_dir or ".",
+                    host=self.config.host or DEFAULT_LOCAL_SERVICE_IP_ADDRESS,
+                    port=self.config.port,
+                    backlog=self.config.backlog,
+                    ssl_certfile=ssl_params.ssl_certfile,
+                    ssl_keyfile=ssl_params.ssl_keyfile,
+                    ssl_keyfile_password=ssl_params.ssl_keyfile_password,
+                    ssl_version=ssl_params.ssl_version,
+                    ssl_cert_reqs=ssl_params.ssl_cert_reqs,
+                    ssl_ca_certs=ssl_params.ssl_ca_certs,
+                    ssl_ciphers=ssl_params.ssl_ciphers,
+                )
+            except KeyboardInterrupt:
+                logger.info("Stopping BentoML prediction service...")
 
     @property
     def prediction_url(self) -> Optional[str]:
@@ -249,12 +287,15 @@ class BentoMLDeploymentService(LocalDaemonService, BaseDeploymentService):
             ]
         return None
 
-    def predict(self, api_endpoint: str, data: "Any") -> "Any":
+    def predict(
+        self, api_endpoint: str, data: "Any", sync: bool = True
+    ) -> "Any":
         """Make a prediction using the service.
 
         Args:
             data: data to make a prediction on
             api_endpoint: the api endpoint to make the prediction on
+            sync: if set to False, the prediction will be made asynchronously
 
         Returns:
             The prediction result.
@@ -268,11 +309,11 @@ class BentoMLDeploymentService(LocalDaemonService, BaseDeploymentService):
                 "BentoML prediction service is not running. "
                 "Please start the service before making predictions."
             )
-        if self.endpoint.prediction_url is not None:
-            client = Client.from_url(
-                self.endpoint.prediction_url.replace("http://", "").rstrip("/")
-            )
-            result = client.call(api_endpoint, data)
-        else:
+        if self.endpoint.prediction_url is None:
             raise ValueError("No endpoint known for prediction.")
+        if sync:
+            client = SyncHTTPClient(self.endpoint.prediction_url)
+        else:
+            client = AsyncHTTPClient(self.endpoint.prediction_url)
+        result = client.call(api_endpoint, data)
         return result

zenml/{steps/base_parameters.py → integrations/bentoml/services/deployment_type.py}

@@ -1,4 +1,4 @@
-#  Copyright (c) ZenML GmbH 2021. All Rights Reserved.
+#  Copyright (c) ZenML GmbH 2024. All Rights Reserved.
 #
 #  Licensed under the Apache License, Version 2.0 (the "License");
 #  you may not use this file except in compliance with the License.
@@ -11,10 +11,13 @@
 #  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.
-"""Step parameters."""
+"""BentoML Service Deployment Types."""
 
-from pydantic import BaseModel
+from enum import Enum
 
 
-class BaseParameters(BaseModel):
-    """Base class to pass parameters into a step."""
+class BentoMLDeploymentType(Enum):
+    """BentoML Service Deployment Types."""
+
+    LOCAL = "local"
+    CONTAINER = "container"

zenml/integrations/bentoml/steps/bento_builder.py

@@ -93,6 +93,7 @@ def bento_builder_step(
     # Build the BentoML bundle
     bento = bentos.build(
         service=service,
+        models=[model_name],
         version=version,
         labels=labels,
         description=description,

zenml/integrations/bentoml/steps/bentoml_deployer.py

@@ -13,7 +13,7 @@
 #  permissions and limitations under the License.
 """Implementation of the BentoML model deployer pipeline step."""
 
-from typing import List, Optional, cast
+from typing import List, Optional, Tuple, cast
 
 import bentoml
 from bentoml._internal.bento import bento
@@ -22,12 +22,21 @@ from zenml import get_step_context, step
 from zenml.integrations.bentoml.model_deployers.bentoml_model_deployer import (
     BentoMLModelDeployer,
 )
-from zenml.integrations.bentoml.services.bentoml_deployment import (
-    BentoMLDeploymentConfig,
-    BentoMLDeploymentService,
+from zenml.integrations.bentoml.services.bentoml_container_deployment import (
+    BentoMLContainerDeploymentConfig,
+    BentoMLContainerDeploymentService,
+)
+from zenml.integrations.bentoml.services.bentoml_local_deployment import (
+    BentoMLLocalDeploymentConfig,
+    BentoMLLocalDeploymentService,
     SSLBentoMLParametersConfig,
 )
+from zenml.integrations.bentoml.services.deployment_type import (
+    BentoMLDeploymentType,
+)
 from zenml.logger import get_logger
+from zenml.services.service import BaseService, ServiceConfig
+from zenml.services.service_type import ServiceType
 from zenml.utils import source_utils
 
 logger = get_logger(__name__)
@@ -38,12 +47,16 @@ def bentoml_model_deployer_step(
     bento: bento.Bento,
     model_name: str,
     port: int,
+    deployment_type: BentoMLDeploymentType = BentoMLDeploymentType.LOCAL,
     deploy_decision: bool = True,
     workers: Optional[int] = 1,
     backlog: Optional[int] = 2048,
     production: bool = False,
     working_dir: Optional[str] = None,
     host: Optional[str] = None,
+    image: Optional[str] = None,
+    image_tag: Optional[str] = None,
+    platform: Optional[str] = None,
     ssl_certfile: Optional[str] = None,
     ssl_keyfile: Optional[str] = None,
     ssl_keyfile_password: Optional[str] = None,
@@ -52,7 +65,7 @@ def bentoml_model_deployer_step(
     ssl_ca_certs: Optional[str] = None,
     ssl_ciphers: Optional[str] = None,
     timeout: int = 30,
-) -> BentoMLDeploymentService:
+) -> BaseService:
     """Model deployer pipeline step for BentoML.
 
     This step deploys a given Bento to a local BentoML http prediction server.
@@ -61,12 +74,16 @@ def bentoml_model_deployer_step(
         bento: the bento artifact to deploy
         model_name: the name of the model to deploy.
         port: the port to use for the prediction service.
+        deployment_type: the type of deployment to use. Either "local" or "container".
         deploy_decision: whether to deploy the model or not
         workers: number of workers to use for the prediction service
         backlog: the number of requests to queue up before rejecting requests.
         production: whether to deploy the service in production mode.
         working_dir: the working directory to use for the prediction service.
         host: the host to use for the prediction service.
+        image: the image to use for the prediction service in the container deployment.
+        image_tag: the image tag to use for the prediction service in the container deployment.
+        platform: the platform to use to build the image for the container deployment.
         ssl_certfile: the path to the ssl cert file.
         ssl_keyfile: the path to the ssl key file.
         ssl_keyfile_password: the password for the ssl key file.
@@ -102,40 +119,62 @@ def bentoml_model_deployer_step(
         apis_paths = list(apis.keys())
         return apis_paths
 
-    # create a config for the new model service
-    predictor_cfg = BentoMLDeploymentConfig(
-        model_name=model_name,
-        bento=str(bento.tag),
-        model_uri=bento.info.labels.get("model_uri"),
-        bento_uri=bento.info.labels.get("bento_uri"),
-        apis=service_apis(str(bento.tag)),
-        workers=workers,
-        host=host,
-        backlog=backlog,
-        working_dir=working_dir or source_utils.get_source_root(),
-        port=port,
-        pipeline_name=pipeline_name,
-        pipeline_step_name=step_name,
-        ssl_parameters=SSLBentoMLParametersConfig(
-            ssl_certfile=ssl_certfile,
-            ssl_keyfile=ssl_keyfile,
-            ssl_keyfile_password=ssl_keyfile_password,
-            ssl_version=ssl_version,
-            ssl_cert_reqs=ssl_cert_reqs,
-            ssl_ca_certs=ssl_ca_certs,
-            ssl_ciphers=ssl_ciphers,
-        ),
-    )
-
+    def create_deployment_config(
+        deployment_type: BentoMLDeploymentType,
+    ) -> Tuple[ServiceConfig, ServiceType]:
+        common_config = {
+            "model_name": model_name,
+            "bento_tag": str(bento.tag),
+            "model_uri": bento.info.labels.get("model_uri"),
+            "bento_uri": bento.info.labels.get("bento_uri"),
+            "apis": service_apis(str(bento.tag)),
+            "host": host,
+            "port": port,
+            "pipeline_name": pipeline_name,
+            "pipeline_step_name": step_name,
+            "workers": workers,
+            "backlog": backlog,
+        }
+
+        if deployment_type == BentoMLDeploymentType.CONTAINER:
+            return BentoMLContainerDeploymentConfig(
+                **common_config,
+                image=image,
+                image_tag=image_tag,
+                platform=platform,
+            ), BentoMLContainerDeploymentService.SERVICE_TYPE
+        else:
+            return BentoMLLocalDeploymentConfig(
+                **common_config,
+                working_dir=working_dir or source_utils.get_source_root(),
+                ssl_parameters=SSLBentoMLParametersConfig(
+                    ssl_certfile=ssl_certfile,
+                    ssl_keyfile=ssl_keyfile,
+                    ssl_keyfile_password=ssl_keyfile_password,
+                    ssl_version=ssl_version,
+                    ssl_cert_reqs=ssl_cert_reqs,
+                    ssl_ca_certs=ssl_ca_certs,
+                    ssl_ciphers=ssl_ciphers,
+                ),
+                production=production,
+            ), BentoMLLocalDeploymentService.SERVICE_TYPE
+
+    predictor_cfg, service_type = create_deployment_config(deployment_type)
     # fetch existing services with same pipeline name, step name and model name
     existing_services = model_deployer.find_model_server(
         config=predictor_cfg.model_dump(),
-        service_type=BentoMLDeploymentService.SERVICE_TYPE,
+        service_type=service_type,
     )
 
     # Creating a new service with inactive state and status by default
+    service: Optional[BaseService] = None
     if existing_services:
-        service = cast(BentoMLDeploymentService, existing_services[0])
+        if deployment_type == BentoMLDeploymentType.CONTAINER:
+            service = cast(
+                BentoMLContainerDeploymentService, existing_services[0]
+            )
+        else:
+            service = cast(BentoMLLocalDeploymentService, existing_services[0])
 
     if not deploy_decision and existing_services:
         logger.info(
@@ -144,20 +183,33 @@ def bentoml_model_deployer_step(
             f"'{step_name}' and pipeline '{pipeline_name}' for model "
             f"'{model_name}'..."
         )
+        assert service is not None
         if not service.is_running:
             service.start(timeout=timeout)
         return service
 
     # create a new model deployment and replace an old one if it exists
-    new_service = cast(
-        BentoMLDeploymentService,
-        model_deployer.deploy_model(
-            replace=True,
-            config=predictor_cfg,
-            timeout=timeout,
-            service_type=BentoMLDeploymentService.SERVICE_TYPE,
-        ),
-    )
+    new_service: BaseService
+    if deployment_type == BentoMLDeploymentType.CONTAINER:
+        new_service = cast(
+            BentoMLContainerDeploymentService,
+            model_deployer.deploy_model(
+                replace=True,
+                config=cast(BentoMLContainerDeploymentConfig, predictor_cfg),
+                timeout=timeout,
+                service_type=service_type,
+            ),
+        )
+    else:
+        new_service = cast(
+            BentoMLLocalDeploymentService,
+            model_deployer.deploy_model(
+                replace=True,
+                config=cast(BentoMLLocalDeploymentConfig, predictor_cfg),
+                timeout=timeout,
+                service_type=service_type,
+            ),
+        )
 
     logger.info(
         f"BentoML deployment service started and reachable at:\n"

zenml/integrations/github/code_repositories/github_code_repository.py

@@ -164,7 +164,7 @@ class GitHubCodeRepository(BaseCodeRepository):
                 try:
                     with open(local_path, "wb") as f:
                         f.write(content.decoded_content)
-                except (GithubException, IOError) as e:
+                except (GithubException, IOError, AssertionError) as e:
                     logger.error("Error processing %s: %s", content.path, e)
 
     def get_local_context(self, path: str) -> Optional[LocalRepositoryContext]:

zenml/integrations/gitlab/code_repositories/gitlab_code_repository.py

@@ -159,7 +159,12 @@ class GitLabCodeRepository(BaseCodeRepository):
             return True
 
         ssh_regex = re.compile(
-            f".*{self.config.host}:{self.config.group}/{self.config.project}.git"
+            r"^(?P<scheme_with_delimiter>ssh://)?"
+            r"(?P<userinfo>git)"
+            f"@{self.config.host}:"
+            r"(?P<port>\d+)?"
+            r"(?(scheme_with_delimiter)/|/?)"
+            f"{self.config.group}/{self.config.project}.git$",
         )
         if ssh_regex.fullmatch(url):
             return True

zenml/integrations/huggingface/__init__.py

@@ -57,14 +57,9 @@ class HuggingfaceIntegration(Integration):
             # - https://github.com/huggingface/datasets/issues/6697
             # TODO try relaxing it back going forward
             "fsspec<=2023.12.0",
+            "transformers",
         ]
 
-        # In python 3.8 higher transformers version lead to other packages breaking
-        if sys.version_info.minor > 8:
-            requirements += ["transformers"]
-        else:
-            requirements += ["transformers<=4.31"]
-
         # Add the pandas integration requirements
         from zenml.integrations.pandas import PandasIntegration
 

zenml/integrations/kaniko/image_builders/kaniko_image_builder.py

@@ -199,9 +199,9 @@ class KanikoImageBuilder(BaseImageBuilder):
             "--image-name-with-digest-file=/dev/termination-log",
         ] + self.config.executor_args
 
-        optional_container_args: Dict[str, Any] = {}
+        optional_spec_args: Dict[str, Any] = {}
         if self.config.service_account_name:
-            optional_container_args["serviceAccountName"] = (
+            optional_spec_args["serviceAccountName"] = (
                 self.config.service_account_name
             )
 
@@ -218,10 +218,10 @@ class KanikoImageBuilder(BaseImageBuilder):
                         "env": self.config.env,
                         "envFrom": self.config.env_from,
                         "volumeMounts": self.config.volume_mounts,
-                        **optional_container_args,
                     }
                 ],
                 "volumes": self.config.volumes,
+                **optional_spec_args,
             },
         }
 

zenml/integrations/lightning/orchestrators/lightning_orchestrator.py

@@ -399,7 +399,7 @@ class LightningOrchestrator(WheeledOrchestrator):
                 f"Installing requirements: {pipeline_requirements_to_string}"
             )
             studio.run(f"uv pip install {pipeline_requirements_to_string}")
-            studio.run("pip install zenml -y")
+            studio.run("pip install zenml")
 
             for custom_command in settings.custom_commands or []:
                 studio.run(
@@ -493,7 +493,7 @@ class LightningOrchestrator(WheeledOrchestrator):
             )
             studio.run("pip install uv")
             studio.run(f"uv pip install {requirements}")
-            studio.run("pip install zenml -y")
+            studio.run("pip install zenml")
             # studio.run(f"pip install {wheel_path.rsplit('/', 1)[-1]}")
             for command in settings.custom_commands or []:
                 output = studio.run(
@@ -566,7 +566,7 @@ class LightningOrchestrator(WheeledOrchestrator):
         )
         studio.run("pip install uv")
         studio.run(f"uv pip install {details['requirements']}")
-        studio.run("pip install zenml -y")
+        studio.run("pip install zenml")
         # studio.run(f"pip install {wheel_path.rsplit('/', 1)[-1]}")
         for command in custom_commands or []:
             output = studio.run(

zenml/integrations/lightning/orchestrators/lightning_orchestrator_entrypoint.py

@@ -166,7 +166,7 @@ def main() -> None:
         f"uv pip install {pipeline_requirements_to_string}"
     )
     logger.info(output)
-    output = main_studio.run("pip install zenml -y")
+    output = main_studio.run("pip install zenml")
     logger.info(output)
 
     for command in pipeline_settings.custom_commands or []:
@@ -248,7 +248,7 @@ def main() -> None:
                     f"uv pip install {step_requirements_to_string}"
                 )
                 logger.info(output)
-                output = studio.run("pip install zenml -y")
+                output = studio.run("pip install zenml")
                 logger.info(output)
                 for command in step_settings.custom_commands or []:
                     output = studio.run(

zenml/integrations/mlflow/experiment_trackers/mlflow_experiment_tracker.py

@@ -221,13 +221,16 @@ class MLFlowExperimentTracker(BaseExperimentTracker):
         Returns:
             A dictionary of metadata.
         """
-        return {
+        metadata: Dict[str, Any] = {
             METADATA_EXPERIMENT_TRACKER_URL: Uri(
                 self.get_tracking_uri(as_plain_text=False)
             ),
-            "mlflow_run_id": mlflow.active_run().info.run_id,
-            "mlflow_experiment_id": mlflow.active_run().info.experiment_id,
         }
+        if run := mlflow.active_run():
+            metadata["mlflow_run_id"] = run.info.run_id
+            metadata["mlflow_experiment_id"] = run.info.experiment_id
+
+        return metadata
 
     def disable_autologging(self) -> None:
         """Disables MLflow autologging for all supported frameworks."""

zenml/integrations/mlflow/model_registries/mlflow_model_registry.py

@@ -362,10 +362,16 @@ class MLFlowModelRegistry(BaseModelRegistry):
 
         Raises:
             RuntimeError: If the registered model does not exist.
+            ValueError: If no model source URI was provided.
 
         Returns:
             The registered model version.
         """
+        if not model_source_uri:
+            raise ValueError(
+                "Unable to register model version without model source URI."
+            )
+
         # Check if the model exists, if not create it.
         try:
             self.get_model(name=name)

zenml/integrations/mlflow/services/mlflow_deployment.py

@@ -225,7 +225,7 @@ class MLFlowDeploymentService(LocalDaemonService, BaseDeploymentService):
             # to run the deploy the model on the local running environment
             if int(mlflow_version[0]) >= 2:
                 backend_kwargs["env_manager"] = "local"
-            backend = PyFuncBackend(
+            backend = PyFuncBackend(  # type: ignore[no-untyped-call]
                 config={},
                 no_conda=True,
                 workers=self.config.workers,
@@ -240,7 +240,7 @@ class MLFlowDeploymentService(LocalDaemonService, BaseDeploymentService):
                     "stack."
                 )
             experiment_tracker.configure_mlflow()
-            backend.serve(
+            backend.serve(  # type: ignore[no-untyped-call]
                 model_uri=self.config.model_uri,
                 port=self.endpoint.status.port,
                 host="localhost",

zenml/integrations/mlflow/steps/mlflow_deployer.py

@@ -114,7 +114,7 @@ def mlflow_model_deployer_step(
     # Fetch the model URI from the MLflow artifact store
     model_uri = ""
     if mlflow_run_id and client.list_artifacts(mlflow_run_id, model_name):
-        model_uri = artifact_utils.get_artifact_uri(
+        model_uri = artifact_utils.get_artifact_uri(  # type: ignore[no-untyped-call]
             run_id=mlflow_run_id, artifact_path=model_name
         )
 

zenml/integrations/mlflow/steps/mlflow_registry.py

@@ -124,7 +124,7 @@ def mlflow_register_model_step(
     if not model_source_uri and client.list_artifacts(
         mlflow_run_id, trained_model_name
     ):
-        model_source_uri = artifact_utils.get_artifact_uri(
+        model_source_uri = artifact_utils.get_artifact_uri(  # type: ignore[no-untyped-call]
             run_id=mlflow_run_id, artifact_path=trained_model_name
         )
     if not model_source_uri: