zenml-nightly 0.73.0.dev20250129__tar.gz → 0.73.0.dev20250131__tar.gz

{zenml_nightly-0.73.0.dev20250129 → zenml_nightly-0.73.0.dev20250131}/PKG-INFO

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

{zenml_nightly-0.73.0.dev20250129 → zenml_nightly-0.73.0.dev20250131}/pyproject.toml

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

zenml_nightly-0.73.0.dev20250131/src/zenml/VERSION

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

{zenml_nightly-0.73.0.dev20250129 → zenml_nightly-0.73.0.dev20250131}/src/zenml/cli/code_repository.py

@@ -162,6 +162,32 @@ def register_code_repository(
         cli_utils.declare(f"Successfully registered code repository `{name}`.")
 
 
+@code_repository.command("describe", help="Describe a code repository.")
+@click.argument(
+    "name_id_or_prefix",
+    type=str,
+    required=True,
+)
+def describe_code_repository(name_id_or_prefix: str) -> None:
+    """Describe a code repository.
+
+    Args:
+        name_id_or_prefix: Name, ID or prefix of the code repository.
+    """
+    client = Client()
+    try:
+        code_repository = client.get_code_repository(
+            name_id_or_prefix=name_id_or_prefix,
+        )
+    except KeyError as err:
+        cli_utils.error(str(err))
+    else:
+        cli_utils.print_pydantic_model(
+            title=f"Code repository '{code_repository.name}'",
+            model=code_repository,
+        )
+
+
 @code_repository.command("list", help="List all connected code repositories.")
 @list_options(CodeRepositoryFilter)
 def list_code_repositories(**kwargs: Any) -> None:

{zenml_nightly-0.73.0.dev20250129 → zenml_nightly-0.73.0.dev20250131}/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.dev20250129 → zenml_nightly-0.73.0.dev20250131}/src/zenml/client.py

@@ -34,7 +34,7 @@ from typing import (
     Union,
     cast,
 )
-from uuid import UUID, uuid4
+from uuid import UUID
 
 from pydantic import ConfigDict, SecretStr
 
@@ -4980,12 +4980,7 @@ class Client(metaclass=ClientMetaClass):
             )
         )
         try:
-            # This does a login to verify the credentials
-            code_repo_class(id=uuid4(), config=config)
-
-            # Explicitly access the config for pydantic validation, in case
-            # the login for some reason did not do that.
-            _ = code_repo_class.config
+            code_repo_class.validate_config(config)
         except Exception as e:
             raise RuntimeError(
                 "Failed to validate code repository config."

{zenml_nightly-0.73.0.dev20250129 → zenml_nightly-0.73.0.dev20250131}/src/zenml/code_repositories/base_code_repository.py

@@ -15,7 +15,7 @@
 
 from abc import ABC, abstractmethod
 from typing import TYPE_CHECKING, Any, Dict, Optional, Set, Type
-from uuid import UUID
+from uuid import UUID, uuid4
 
 from zenml.config.secret_reference_mixin import SecretReferenceMixin
 from zenml.logger import get_logger
@@ -44,15 +44,18 @@ class BaseCodeRepository(ABC):
     def __init__(
         self,
         id: UUID,
+        name: str,
         config: Dict[str, Any],
     ) -> None:
         """Initializes a code repository.
 
         Args:
             id: The ID of the code repository.
+            name: The name of the code repository.
             config: The config of the code repository.
         """
         self._id = id
+        self._name = name
         self._config = config
         self.login()
 
@@ -80,7 +83,23 @@ class BaseCodeRepository(ABC):
                 source=model.source, expected_class=BaseCodeRepository
             )
         )
-        return class_(id=model.id, config=model.config)
+        return class_(id=model.id, name=model.name, config=model.config)
+
+    @classmethod
+    def validate_config(cls, config: Dict[str, Any]) -> None:
+        """Validate the code repository config.
+
+        This method should check that the config/credentials are valid and
+        the configured repository exists.
+
+        Args:
+            config: The configuration.
+        """
+        # The initialization calls the login to verify the credentials
+        code_repo = cls(id=uuid4(), name="", config=config)
+
+        # Explicitly access the config for pydantic validation
+        _ = code_repo.config
 
     @property
     def id(self) -> UUID:
@@ -91,6 +110,15 @@ class BaseCodeRepository(ABC):
         """
         return self._id
 
+    @property
+    def name(self) -> str:
+        """Name of the code repository.
+
+        Returns:
+            The name of the code repository.
+        """
+        return self._name
+
     @property
     def requirements(self) -> Set[str]:
         """Set of PyPI requirements for the repository.

{zenml_nightly-0.73.0.dev20250129 → zenml_nightly-0.73.0.dev20250131}/src/zenml/code_repositories/git/local_git_repository_context.py

@@ -14,11 +14,14 @@
 """Implementation of the Local git repository context."""
 
 from typing import TYPE_CHECKING, Callable, Optional, cast
-from uuid import UUID
 
 from zenml.code_repositories import (
     LocalRepositoryContext,
 )
+from zenml.constants import (
+    ENV_ZENML_CODE_REPOSITORY_IGNORE_UNTRACKED_FILES,
+    handle_bool_env_var,
+)
 from zenml.logger import get_logger
 
 if TYPE_CHECKING:
@@ -26,6 +29,8 @@ if TYPE_CHECKING:
     from git.remote import Remote
     from git.repo.base import Repo
 
+    from zenml.code_repositories import BaseCodeRepository
+
 logger = get_logger(__name__)
 
 
@@ -33,16 +38,19 @@ class LocalGitRepositoryContext(LocalRepositoryContext):
     """Local git repository context."""
 
     def __init__(
-        self, code_repository_id: UUID, git_repo: "Repo", remote_name: str
+        self,
+        code_repository: "BaseCodeRepository",
+        git_repo: "Repo",
+        remote_name: str,
     ):
         """Initializes a local git repository context.
 
         Args:
-            code_repository_id: The ID of the code repository.
+            code_repository: The code repository.
             git_repo: The git repo.
             remote_name: Name of the remote.
         """
-        super().__init__(code_repository_id=code_repository_id)
+        super().__init__(code_repository=code_repository)
         self._git_repo = git_repo
         self._remote = git_repo.remote(name=remote_name)
 
@@ -50,14 +58,14 @@ class LocalGitRepositoryContext(LocalRepositoryContext):
     def at(
         cls,
         path: str,
-        code_repository_id: UUID,
+        code_repository: "BaseCodeRepository",
         remote_url_validation_callback: Callable[[str], bool],
     ) -> Optional["LocalGitRepositoryContext"]:
         """Returns a local git repository at the given path.
 
         Args:
             path: The path to the local git repository.
-            code_repository_id: The ID of the code repository.
+            code_repository: The code repository.
             remote_url_validation_callback: A callback that validates the
                 remote URL of the git repository.
 
@@ -70,11 +78,13 @@ class LocalGitRepositoryContext(LocalRepositoryContext):
             from git.exc import InvalidGitRepositoryError
             from git.repo.base import Repo
         except ImportError:
+            logger.debug("Failed to import git library.")
             return None
 
         try:
             git_repo = Repo(path=path, search_parent_directories=True)
         except InvalidGitRepositoryError:
+            logger.debug("No git repository exists at path %s.", path)
             return None
 
         remote_name = None
@@ -87,7 +97,7 @@ class LocalGitRepositoryContext(LocalRepositoryContext):
             return None
 
         return LocalGitRepositoryContext(
-            code_repository_id=code_repository_id,
+            code_repository=code_repository,
             git_repo=git_repo,
             remote_name=remote_name,
         )
@@ -124,13 +134,19 @@ class LocalGitRepositoryContext(LocalRepositoryContext):
     def is_dirty(self) -> bool:
         """Whether the git repo is dirty.
 
-        A repository counts as dirty if it has any untracked or uncommitted
-        changes.
+        By default, a repository counts as dirty if it has any untracked or
+        uncommitted changes. Users can use an environment variable to ignore
+        untracked files.
 
         Returns:
             True if the git repo is dirty, False otherwise.
         """
-        return self.git_repo.is_dirty(untracked_files=True)
+        ignore_untracked_files = handle_bool_env_var(
+            ENV_ZENML_CODE_REPOSITORY_IGNORE_UNTRACKED_FILES, default=False
+        )
+        return self.git_repo.is_dirty(
+            untracked_files=not ignore_untracked_files
+        )
 
     @property
     def has_local_changes(self) -> bool:

{zenml_nightly-0.73.0.dev20250129 → zenml_nightly-0.73.0.dev20250131}/src/zenml/code_repositories/local_repository_context.py

@@ -14,10 +14,13 @@
 """Base class for local code repository contexts."""
 
 from abc import ABC, abstractmethod
-from uuid import UUID
+from typing import TYPE_CHECKING
 
 from zenml.logger import get_logger
 
+if TYPE_CHECKING:
+    from zenml.code_repositories import BaseCodeRepository
+
 logger = get_logger(__name__)
 
 
@@ -30,22 +33,22 @@ class LocalRepositoryContext(ABC):
     commit, and whether the repository is dirty.
     """
 
-    def __init__(self, code_repository_id: UUID) -> None:
+    def __init__(self, code_repository: "BaseCodeRepository") -> None:
         """Initializes a local repository context.
 
         Args:
-            code_repository_id: The ID of the code repository.
+            code_repository: The code repository.
         """
-        self._code_repository_id = code_repository_id
+        self._code_repository = code_repository
 
     @property
-    def code_repository_id(self) -> UUID:
-        """Returns the ID of the code repository.
+    def code_repository(self) -> "BaseCodeRepository":
+        """Returns the code repository.
 
         Returns:
-            The ID of the code repository.
+            The code repository.
         """
-        return self._code_repository_id
+        return self._code_repository
 
     @property
     @abstractmethod

{zenml_nightly-0.73.0.dev20250129 → zenml_nightly-0.73.0.dev20250131}/src/zenml/constants.py

@@ -175,6 +175,9 @@ ENV_ZENML_WHEEL_PACKAGE_NAME = "ZENML_WHEEL_PACKAGE_NAME"
 ENV_ZENML_PIPELINE_RUN_API_TOKEN_EXPIRATION = (
     "ZENML_PIPELINE_API_TOKEN_EXPIRATION"
 )
+ENV_ZENML_CODE_REPOSITORY_IGNORE_UNTRACKED_FILES = (
+    "ZENML_CODE_REPOSITORY_IGNORE_UNTRACKED_FILES"
+)
 
 # Materializer environment variables
 ENV_ZENML_MATERIALIZER_ALLOW_NON_ASCII_JSON_DUMPS = (

{zenml_nightly-0.73.0.dev20250129 → zenml_nightly-0.73.0.dev20250131}/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.dev20250129 → zenml_nightly-0.73.0.dev20250131}/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.dev20250129 → zenml_nightly-0.73.0.dev20250131}/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.dev20250129 → zenml_nightly-0.73.0.dev20250131}/src/zenml/integrations/github/code_repositories/github_code_repository.py

@@ -15,7 +15,8 @@
 
 import os
 import re
-from typing import List, Optional
+from typing import Any, Dict, List, Optional
+from uuid import uuid4
 
 import requests
 from github import Consts, Github, GithubException
@@ -62,6 +63,20 @@ class GitHubCodeRepositoryConfig(BaseCodeRepositoryConfig):
 class GitHubCodeRepository(BaseCodeRepository):
     """GitHub code repository."""
 
+    @classmethod
+    def validate_config(cls, config: Dict[str, Any]) -> None:
+        """Validate the code repository config.
+
+        This method should check that the config/credentials are valid and
+        the configured repository exists.
+
+        Args:
+            config: The configuration.
+        """
+        code_repo = cls(id=uuid4(), name="", config=config)
+        # Try to access the project to make sure it exists
+        _ = code_repo.github_repo
+
     @property
     def config(self) -> GitHubCodeRepositoryConfig:
         """Returns the `GitHubCodeRepositoryConfig` config.
@@ -190,7 +205,7 @@ class GitHubCodeRepository(BaseCodeRepository):
         """
         return LocalGitRepositoryContext.at(
             path=path,
-            code_repository_id=self.id,
+            code_repository=self,
             remote_url_validation_callback=self.check_remote_url,
         )
 

{zenml_nightly-0.73.0.dev20250129 → zenml_nightly-0.73.0.dev20250131}/src/zenml/integrations/gitlab/code_repositories/gitlab_code_repository.py

@@ -15,7 +15,8 @@
 
 import os
 import re
-from typing import Optional
+from typing import Any, Dict, Optional
+from uuid import uuid4
 
 from gitlab import Gitlab
 from gitlab.v4.objects import Project
@@ -63,6 +64,20 @@ class GitLabCodeRepositoryConfig(BaseCodeRepositoryConfig):
 class GitLabCodeRepository(BaseCodeRepository):
     """GitLab code repository."""
 
+    @classmethod
+    def validate_config(cls, config: Dict[str, Any]) -> None:
+        """Validate the code repository config.
+
+        This method should check that the config/credentials are valid and
+        the configured repository exists.
+
+        Args:
+            config: The configuration.
+        """
+        code_repo = cls(id=uuid4(), name="", config=config)
+        # Try to access the project to make sure it exists
+        _ = code_repo.gitlab_project
+
     @property
     def config(self) -> GitLabCodeRepositoryConfig:
         """Returns the `GitLabCodeRepositoryConfig` config.
@@ -147,7 +162,7 @@ class GitLabCodeRepository(BaseCodeRepository):
         """
         return LocalGitRepositoryContext.at(
             path=path,
-            code_repository_id=self.id,
+            code_repository=self,
             remote_url_validation_callback=self.check_remote_url,
         )
 

{zenml_nightly-0.73.0.dev20250129 → zenml_nightly-0.73.0.dev20250131}/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