zenml-nightly 0.58.2.dev20240626__py3-none-any.whl → 0.61.0.dev20240710__py3-none-any.whl

zenml/integrations/gcp/service_connectors/gcp_service_connector.py

@@ -43,7 +43,8 @@ from google.auth._default import (
     _get_external_account_credentials,
 )
 from google.auth.transport.requests import Request
-from google.cloud import container_v1, storage
+from google.cloud import artifactregistry_v1, container_v1, storage
+from google.cloud.location import locations_pb2
 from google.oauth2 import credentials as gcp_credentials
 from google.oauth2 import service_account as gcp_service_account
 from pydantic import Field, field_validator, model_validator
@@ -351,24 +352,72 @@ class GCPOAuth2Token(AuthenticationConfig):
 class GCPBaseConfig(AuthenticationConfig):
     """GCP base configuration."""
 
+    @property
+    def gcp_project_id(self) -> str:
+        """Get the GCP project ID.
+
+        This method must be implemented by subclasses to ensure that the GCP
+        project ID is always available.
+
+        Raises:
+            NotImplementedError: If the method is not implemented.
+        """
+        raise NotImplementedError
+
+
+class GCPBaseProjectIDConfig(GCPBaseConfig):
+    """GCP base configuration with included project ID."""
+
     project_id: str = Field(
         title="GCP Project ID where the target resource is located.",
     )
 
+    @property
+    def gcp_project_id(self) -> str:
+        """Get the GCP project ID.
+
+        Returns:
+            The GCP project ID.
+        """
+        return self.project_id
 
-class GCPUserAccountConfig(GCPBaseConfig, GCPUserAccountCredentials):
+
+class GCPUserAccountConfig(GCPBaseProjectIDConfig, GCPUserAccountCredentials):
     """GCP user account configuration."""
 
 
 class GCPServiceAccountConfig(GCPBaseConfig, GCPServiceAccountCredentials):
     """GCP service account configuration."""
 
+    _project_id: Optional[str] = None
+
+    @property
+    def gcp_project_id(self) -> str:
+        """Get the GCP project ID.
+
+        When a service account JSON is provided, the project ID can be extracted
+        from it instead of being provided explicitly.
+
+        Returns:
+            The GCP project ID.
+        """
+        if self._project_id is None:
+            self._project_id = json.loads(
+                self.service_account_json.get_secret_value()
+            )["project_id"]
+            # Guaranteed by the field validator
+            assert self._project_id is not None
+
+        return self._project_id
+
 
-class GCPExternalAccountConfig(GCPBaseConfig, GCPExternalAccountCredentials):
+class GCPExternalAccountConfig(
+    GCPBaseProjectIDConfig, GCPExternalAccountCredentials
+):
     """GCP external account configuration."""
 
 
-class GCPOAuth2TokenConfig(GCPBaseConfig, GCPOAuth2Token):
+class GCPOAuth2TokenConfig(GCPBaseProjectIDConfig, GCPOAuth2Token):
     """GCP OAuth 2.0 configuration."""
 
     service_account_email: Optional[str] = Field(
@@ -540,7 +589,7 @@ resources in the specified project. When used remotely in a GCP workload, the
 configured project has to be the same as the project of the attached service
 account.
 """,
-            config_class=GCPBaseConfig,
+            config_class=GCPBaseProjectIDConfig,
         ),
         AuthenticationMethodModel(
             name="GCP User Account",
@@ -786,14 +835,53 @@ GKE clusters in the GCP project that it is configured to use.
             emoji=":cyclone:",
         ),
         ResourceTypeModel(
-            name="GCP GCR container registry",
+            name="GCP GAR container registry",
             resource_type=DOCKER_REGISTRY_RESOURCE_TYPE,
             description="""
-Allows Stack Components to access a GCR registry as a standard
+Allows Stack Components to access a Google Artifact Registry as a standard
 Docker registry resource. When used by Stack Components, they are provided a
 pre-authenticated Python Docker client instance.
 
-The configured credentials must have at least the following [GCP permissions](https://cloud.google.com/iam/docs/permissions-reference):
+The configured credentials must have at least the following [GCP permissions](https://cloud.google.com/iam/docs/understanding-roles#artifact-registry-roles):
+
+- `artifactregistry.repositories.createOnPush`
+- `artifactregistry.repositories.downloadArtifacts`
+- `artifactregistry.repositories.get`
+- `artifactregistry.repositories.list`
+- `artifactregistry.repositories.readViaVirtualRepository`
+- `artifactregistry.repositories.uploadArtifacts`
+- `artifactregistry.locations.list`
+
+The Artifact Registry Create-on-Push Writer role includes all of the above
+permissions.
+
+This resource type also includes legacy GCR container registry support.
+
+**Important Notice: Google Container Registry** [**is being replaced by Artifact Registry**](https://cloud.google.com/artifact-registry/docs/transition/transition-from-gcr).
+Please start using Artifact Registry for your containers. As per Google's
+documentation, *"after May 15, 2024, Artifact Registry will host images for the
+gcr.io domain in Google Cloud projects without previous Container Registry
+usage. After March 18, 2025, Container Registry will be shut down."*.
+
+Support for legacy GCR registries is still included in the GCP service
+connector. Users that already have GCP service connectors configured to access
+GCR registries may continue to use them without taking any action. However, it
+is recommended to transition to Google Artifact Registries as soon as possible
+by following [the GCP guide on this subject](https://cloud.google.com/artifact-registry/docs/transition/transition-from-gcr)
+and making the following updates to ZenML GCP Service Connectors that are used
+to access GCR resources:
+
+* add the IAM permissions documented here to the GCP Service Connector
+credentials to enable them to access the Artifact Registries.
+* users may keep the gcr.io GCR URLs already configured in the GCP Service
+Connectors as well as those used in linked Container Registry stack components
+given that these domains are redirected by Google to GAR as covered in the GCR
+transition guide. Alternatively, users may update the GCP Service Connector
+configuration and/or the Container Registry stack components to use the
+replacement Artifact Registry URLs.
+
+When used with GCR registries, the configured credentials must have at least the
+following [GCP permissions](https://cloud.google.com/iam/docs/understanding-roles#cloud-storage-roles):
 
 - `storage.buckets.get`
 - `storage.multipartUploads.abort`
@@ -807,17 +895,21 @@ The configured credentials must have at least the following [GCP permissions](ht
 The Storage Legacy Bucket Writer role includes all of the above permissions
 while at the same time restricting access to only the GCR buckets.
 
-The resource name associated with this resource type identifies the GCR
-container registry associated with the GCP project (the repository name is
-optional):
+If set, the resource name must identify a GAR or GCR registry using one of the
+following formats:
+
+- Google Artifact Registry repository URI: `[https://]<region>-docker.pkg.dev/<project-id>/<registry-id>[/<repository-name>]`
+- Google Artifact Registry name: `projects/<project-id>/locations/<location>/repositories/<repository-id>`
+- (legacy) GCR repository URI: `[https://][us.|eu.|asia.]gcr.io/<project-id>[/<repository-name>]`
 
-- GCR repository URI: `[https://]gcr.io/{project-id}[/{repository-name}]
+The connector can only be used to access GAR and GCR registries in the GCP
+project that it is configured to use.
 """,
             auth_methods=GCPAuthenticationMethods.values(),
-            # Does not support instances, given that the connector
-            # provides access to the entire GCR container registry
-            # for the configured GCP project.
-            supports_instances=False,
+            # The connector provides access to the entire GCR container registry
+            # for the configured GCP project as well as any number of artifact
+            # registry repositories.
+            supports_instances=True,
             logo_url="https://public-flavor-logos.s3.eu-central-1.amazonaws.com/container_registry/docker.png",
             emoji=":whale:",
         ),
@@ -1006,6 +1098,7 @@ class GCPServiceConnector(ServiceConnector):
                 # service account authentication)
 
                 assert isinstance(cfg, GCPServiceAccountConfig)
+
                 credentials = (
                     gcp_service_account.Credentials.from_service_account_info(
                         json.loads(
@@ -1093,60 +1186,101 @@ class GCPServiceConnector(ServiceConnector):
 
         return bucket_name
 
-    def _parse_gcr_resource_id(
+    def _parse_gar_resource_id(
         self,
         resource_id: str,
-    ) -> str:
-        """Validate and convert an GCR resource ID to an GCR registry ID.
+    ) -> Tuple[str, Optional[str]]:
+        """Validate and convert a GAR resource ID to a Google Artifact Registry ID and name.
 
         Args:
             resource_id: The resource ID to convert.
 
         Returns:
-            The GCR registry ID.
+            The Google Artifact Registry ID and name. The name is omitted if the
+            resource ID is a GCR repository URI.
 
         Raises:
-            ValueError: If the provided resource ID is not a valid GCR
-                repository URI.
+            ValueError: If the provided resource ID is not a valid GAR
+                or GCR repository URI.
         """
         # The resource ID could mean different things:
         #
-        # - an GCR repository URI
+        # - a GAR repository URI
+        # - a GAR repository name
+        # - a GCR repository URI (backwards-compatibility)
         #
         # We need to extract the project ID and registry ID from
         # the provided resource ID
-        config_project_id = self.config.project_id
+        config_project_id = self.config.gcp_project_id
         project_id: Optional[str] = None
-        # A GCR repository URI uses one of several hostnames (gcr.io, us.gcr.io,
-        # eu.gcr.io, asia.gcr.io etc.) and the project ID is the first part of
-        # the URL path
-        if re.match(
-            r"^(https://)?([a-z]+.)*gcr.io/[a-z0-9-]+(/.+)*$",
+        canonical_url: str
+        registry_name: Optional[str] = None
+
+        # A Google Artifact Registry URI uses the <location>-docker-pkg.dev
+        # domain format with the project ID as the first part of the URL path
+        # and the registry name as the second part of the URL path
+        if match := re.match(
+            r"^(https://)?(([a-z0-9-]+)-docker\.pkg\.dev/([a-z0-9-]+)/([a-z0-9-.]+))(/.+)*$",
             resource_id,
         ):
-            # The resource ID is a GCR repository URI
-            if resource_id.startswith("https://"):
-                project_id = resource_id.split("/")[3]
-            else:
-                project_id = resource_id.split("/")[1]
+            # The resource ID is a Google Artifact Registry URI
+            project_id = match[4]
+            location = match[3]
+            repository = match[5]
+
+            # Return the GAR URL without the image name and without the protocol
+            canonical_url = match[2]
+            registry_name = f"projects/{project_id}/locations/{location}/repositories/{repository}"
+
+        # Alternatively, the Google Artifact Registry name uses the
+        # projects/<project-id>/locations/<location>/repositories/<repository-id>
+        # format
+        elif match := re.match(
+            r"^projects/([a-z0-9-]+)/locations/([a-z0-9-]+)/repositories/([a-z0-9-.]+)$",
+            resource_id,
+        ):
+            # The resource ID is a Google Artifact Registry name
+            project_id = match[1]
+            location = match[2]
+            repository = match[3]
+
+            # Return the GAR URL
+            canonical_url = (
+                f"{location}-docker.pkg.dev/{project_id}/{repository}"
+            )
+            registry_name = resource_id
+
+        # A legacy GCR repository URI uses one of several hostnames (gcr.io,
+        # us.gcr.io, eu.gcr.io, asia.gcr.io) and the project ID is the
+        # first part of the URL path
+        elif match := re.match(
+            r"^(https://)?(((us|eu|asia)\.)?gcr\.io/[a-z0-9-]+)(/.+)*$",
+            resource_id,
+        ):
+            # The resource ID is a legacy GCR repository URI.
+            # Return the GAR URL without the image name and without the protocol
+            canonical_url = match[2]
+
         else:
             raise ValueError(
-                f"Invalid resource ID for a GCR registry: {resource_id}. "
-                f"Supported formats are:\n"
+                f"Invalid resource ID for a Google Artifact Registry: "
+                f"{resource_id}. Supported formats are:\n"
+                f"Google Artifact Registry URI: [https://]<region>-docker.pkg.dev/<project-id>/<registry-id>[/<repository-name>]\n"
+                f"Google Artifact Registry name: projects/<project-id>/locations/<location>/repositories/<repository-id>\n"
                 f"GCR repository URI: [https://][us.|eu.|asia.]gcr.io/<project-id>[/<repository-name>]"
             )
 
         # If the connector is configured with a project and the resource ID
-        # is an GCR repository URI that specifies a different project,
+        # is a GAR repository URI that specifies a different project,
         # we raise an error
         if project_id and project_id != config_project_id:
             raise ValueError(
-                f"The GCP project for the {resource_id} GCR repository "
-                f"'{project_id}' does not match the project configured in "
-                f"the connector: '{config_project_id}'."
+                f"The GCP project for the {resource_id} Google Artifact "
+                f"Registry '{project_id}' does not match the project "
+                f"configured in the connector: '{config_project_id}'."
             )
 
-        return f"gcr.io/{project_id}"
+        return canonical_url, registry_name
 
     def _parse_gke_resource_id(self, resource_id: str) -> str:
         """Validate and convert an GKE resource ID to a GKE cluster name.
@@ -1195,7 +1329,7 @@ class GCPServiceConnector(ServiceConnector):
             cluster_name = self._parse_gke_resource_id(resource_id)
             return cluster_name
         elif resource_type == DOCKER_REGISTRY_RESOURCE_TYPE:
-            registry_id = self._parse_gcr_resource_id(
+            registry_id, _ = self._parse_gar_resource_id(
                 resource_id,
             )
             return registry_id
@@ -1219,9 +1353,7 @@ class GCPServiceConnector(ServiceConnector):
                 authorized.
         """
         if resource_type == GCP_RESOURCE_TYPE:
-            return self.config.project_id
-        elif resource_type == DOCKER_REGISTRY_RESOURCE_TYPE:
-            return f"gcr.io/{self.config.project_id}"
+            return self.config.gcp_project_id
 
         raise RuntimeError(
             f"Default resource ID not supported for '{resource_type}' resource "
@@ -1278,7 +1410,7 @@ class GCPServiceConnector(ServiceConnector):
 
             # Create an GCS client for the bucket
             client = storage.Client(
-                project=self.config.project_id, credentials=credentials
+                project=self.config.gcp_project_id, credentials=credentials
             )
             return client
 
@@ -1384,7 +1516,7 @@ class GCPServiceConnector(ServiceConnector):
                             "config",
                             "set",
                             "project",
-                            self.config.project_id,
+                            self.config.gcp_project_id,
                         ],
                         check=True,
                         stderr=subprocess.STDOUT,
@@ -1488,7 +1620,7 @@ class GCPServiceConnector(ServiceConnector):
             )
 
         if auth_method == GCPAuthenticationMethods.IMPLICIT:
-            auth_config = GCPBaseConfig(
+            auth_config = GCPBaseProjectIDConfig(
                 project_id=project_id,
             )
         elif auth_method == GCPAuthenticationMethods.OAUTH2_TOKEN:
@@ -1697,7 +1829,7 @@ class GCPServiceConnector(ServiceConnector):
 
         if resource_type == GCS_RESOURCE_TYPE:
             gcs_client = storage.Client(
-                project=self.config.project_id, credentials=credentials
+                project=self.config.gcp_project_id, credentials=credentials
             )
             if not resource_id:
                 # List all GCS buckets
@@ -1722,11 +1854,87 @@ class GCPServiceConnector(ServiceConnector):
                     raise AuthorizationException(msg) from e
 
         if resource_type == DOCKER_REGISTRY_RESOURCE_TYPE:
-            assert resource_id is not None
+            # Get a GAR client
+            gar_client = artifactregistry_v1.ArtifactRegistryClient(
+                credentials=credentials
+            )
 
-            # No way to verify a GCR registry without attempting to
-            # connect to it via Docker/OCI, so just return the resource ID.
-            return [resource_id]
+            if resource_id:
+                registry_id, registry_name = self._parse_gar_resource_id(
+                    resource_id
+                )
+
+                if registry_name is None:
+                    # This is a legacy GCR repository URI. We can't verify
+                    # the repository access without attempting to connect to it
+                    # via Docker/OCI, so just return the resource ID.
+                    return [registry_id]
+
+                # Check if the specified GAR registry exists
+                try:
+                    repository = gar_client.get_repository(
+                        name=registry_name,
+                    )
+                    if repository.format_.name != "DOCKER":
+                        raise AuthorizationException(
+                            f"Google Artifact Registry '{resource_id}' is not a "
+                            "Docker registry."
+                        )
+                    return [registry_id]
+                except google.api_core.exceptions.GoogleAPIError as e:
+                    msg = f"Failed to fetch Google Artifact Registry '{registry_id}': {e}"
+                    logger.error(msg)
+                    raise AuthorizationException(msg) from e
+
+            # For backwards compatibility, we initialize the list of resource
+            # IDs with all GCR supported registries for the configured GCP
+            # project
+            resource_ids: List[str] = [
+                f"{location}gcr.io/{self.config.gcp_project_id}"
+                for location in ["", "us.", "eu.", "asia."]
+            ]
+
+            # List all Google Artifact Registries
+            try:
+                # First, we need to fetch all the Artifact Registry supported
+                # locations
+                locations = gar_client.list_locations(
+                    request=locations_pb2.ListLocationsRequest(
+                        name=f"projects/{self.config.gcp_project_id}"
+                    )
+                )
+                location_names = [
+                    locations.locations[i].location_id
+                    for i in range(len(locations.locations))
+                ]
+
+                # Then, we need to fetch all the repositories in each location
+                repository_names: List[str] = []
+                for location in location_names:
+                    repositories = gar_client.list_repositories(
+                        parent=f"projects/{self.config.gcp_project_id}/locations/{location}"
+                    )
+                    repository_names.extend(
+                        [
+                            repository.name
+                            for repository in repositories
+                            if repository.format_.name == "DOCKER"
+                        ]
+                    )
+
+                for repository_name in repository_names:
+                    # Convert the repository name to a canonical GAR URL
+                    resource_ids.append(
+                        self._parse_gar_resource_id(repository_name)[0]
+                    )
+
+            except google.api_core.exceptions.GoogleAPIError as e:
+                msg = f"Failed to list Google Artifact Registries: {e}"
+                logger.error(msg)
+                # TODO: enable when GCR is no longer supported:
+                # raise AuthorizationException(msg) from e
+
+            return resource_ids
 
         if resource_type == KUBERNETES_CLUSTER_RESOURCE_TYPE:
             gke_client = container_v1.ClusterManagerClient(
@@ -1736,7 +1944,7 @@ class GCPServiceConnector(ServiceConnector):
             # List all GKE clusters
             try:
                 clusters = gke_client.list_clusters(
-                    parent=f"projects/{self.config.project_id}/locations/-"
+                    parent=f"projects/{self.config.gcp_project_id}/locations/-"
                 )
                 cluster_names = [cluster.name for cluster in clusters.clusters]
             except google.api_core.exceptions.GoogleAPIError as e:
@@ -1810,7 +2018,7 @@ class GCPServiceConnector(ServiceConnector):
             # object
             auth_method: str = GCPAuthenticationMethods.OAUTH2_TOKEN
             config: GCPBaseConfig = GCPOAuth2TokenConfig(
-                project_id=self.config.project_id,
+                project_id=self.config.gcp_project_id,
                 token=credentials.token,
                 service_account_email=credentials.signer_email
                 if hasattr(credentials, "signer_email")
@@ -1855,7 +2063,7 @@ class GCPServiceConnector(ServiceConnector):
         if resource_type == DOCKER_REGISTRY_RESOURCE_TYPE:
             assert resource_id is not None
 
-            registry_id = self._parse_gcr_resource_id(resource_id)
+            registry_id, _ = self._parse_gar_resource_id(resource_id)
 
             # Create a client-side Docker connector instance with the temporary
             # Docker credentials
@@ -1884,7 +2092,7 @@ class GCPServiceConnector(ServiceConnector):
             # List all GKE clusters
             try:
                 clusters = gke_client.list_clusters(
-                    parent=f"projects/{self.config.project_id}/locations/-"
+                    parent=f"projects/{self.config.gcp_project_id}/locations/-"
                 )
                 cluster_map = {
                     cluster.name: cluster for cluster in clusters.clusters
@@ -1928,7 +2136,7 @@ class GCPServiceConnector(ServiceConnector):
                 auth_method=KubernetesAuthenticationMethods.TOKEN,
                 resource_type=resource_type,
                 config=KubernetesTokenConfig(
-                    cluster_name=f"gke_{self.config.project_id}_{cluster_name}",
+                    cluster_name=f"gke_{self.config.gcp_project_id}_{cluster_name}",
                     certificate_authority=cluster_ca_cert,
                     server=f"https://{cluster_server}",
                     token=bearer_token,

zenml/integrations/huggingface/__init__.py

@@ -39,6 +39,7 @@ class HuggingfaceIntegration(Integration):
         # TODO try relaxing it back going forward
         "fsspec<=2023.12.0",
     ]
+    REQUIREMENTS_IGNORED_ON_UNINSTALL = ["fsspec"]
 
     @classmethod
     def activate(cls) -> None:

zenml/integrations/integration.py

@@ -60,6 +60,7 @@ class Integration(metaclass=IntegrationMeta):
 
     REQUIREMENTS: List[str] = []
     APT_PACKAGES: List[str] = []
+    REQUIREMENTS_IGNORED_ON_UNINSTALL: List[str] = []
 
     @classmethod
     def check_installation(cls) -> bool:
@@ -143,6 +144,29 @@ class Integration(metaclass=IntegrationMeta):
         """
         return cls.REQUIREMENTS
 
+    @classmethod
+    def get_uninstall_requirements(
+        cls, target_os: Optional[str] = None
+    ) -> List[str]:
+        """Method to get the uninstall requirements for the integration.
+
+        Args:
+            target_os: The target operating system to get the requirements for.
+
+        Returns:
+            A list of requirements.
+        """
+        ret = []
+        for each in cls.get_requirements(target_os=target_os):
+            is_ignored = False
+            for ignored in cls.REQUIREMENTS_IGNORED_ON_UNINSTALL:
+                if each.startswith(ignored):
+                    is_ignored = True
+                    break
+            if not is_ignored:
+                ret.append(each)
+        return ret
+
     @classmethod
     def activate(cls) -> None:
         """Abstract method to activate the integration."""

zenml/integrations/kubeflow/__init__.py

@@ -31,6 +31,9 @@ class KubeflowIntegration(Integration):
 
     NAME = KUBEFLOW
     REQUIREMENTS = ["kfp>=2.6.0", "kfp-kubernetes>=1.1.0"]  # Only 1.x version that supports pyyaml 6
+    REQUIREMENTS_IGNORED_ON_UNINSTALL = [
+        "kfp", # it is used by GCP as well
+    ]
 
     @classmethod
     def flavors(cls) -> List[Type[Flavor]]:

zenml/integrations/kubeflow/flavors/kubeflow_orchestrator_flavor.py

@@ -149,7 +149,7 @@ class KubeflowOrchestratorConfig(
 
     kubeflow_hostname: Optional[str] = None
     kubeflow_namespace: str = "kubeflow"
-    kubernetes_context: Optional[str]  # TODO: Potential setting
+    kubernetes_context: Optional[str] = None  # TODO: Potential setting
 
     @model_validator(mode="before")
     @classmethod

zenml/integrations/kubeflow/orchestrators/kubeflow_orchestrator.py

@@ -433,7 +433,6 @@ class KubeflowOrchestrator(ContainerizedOrchestrator):
             The dynamic container component.
         """
 
-        @dsl.container_component  # type: ignore[misc]
         def dynamic_container_component() -> dsl.ContainerSpec:
             """Dynamic container component.
 

zenml/integrations/kubernetes/__init__.py

@@ -31,7 +31,9 @@ class KubernetesIntegration(Integration):
 
     NAME = KUBERNETES
     REQUIREMENTS = ["kubernetes>=21.7,<26"]
-
+    REQUIREMENTS_IGNORED_ON_UNINSTALL = [
+        "kfp", # it is used by many others
+    ]
     @classmethod
     def flavors(cls) -> List[Type[Flavor]]:
         """Declare the stack component flavors for the Kubernetes integration.

zenml/integrations/kubernetes/orchestrators/kube_utils.py

@@ -225,8 +225,11 @@ def wait_pod(
             response = core_api.read_namespaced_pod_log(
                 name=pod_name,
                 namespace=namespace,
+                _preload_content=False,
             )
-            logs = response.splitlines()
+            raw_data = response.data
+            decoded_log = raw_data.decode("utf-8", errors="replace")
+            logs = decoded_log.splitlines()
             if len(logs) > logged_lines:
                 for line in logs[logged_lines:]:
                     logger.info(line)

zenml/integrations/label_studio/annotators/label_studio_annotator.py

@@ -527,6 +527,7 @@ class LabelStudioAnnotator(BaseAnnotator, AuthenticationMixin):
                 aws_access_key_id,
                 aws_secret_access_key,
                 aws_session_token,
+                _,
             ) = artifact_store.get_credentials()
 
             if aws_access_key_id and aws_secret_access_key:

zenml/integrations/langchain/__init__.py

@@ -30,6 +30,7 @@ class LangchainIntegration(Integration):
         "pyyaml>=6.0.1",
         "tenacity!=8.4.0",  # https://github.com/jd/tenacity/issues/471
     ]
+    REQUIREMENTS_IGNORED_ON_UNINSTALL = ["pyyaml","tenacity"]
 
     @classmethod
     def activate(cls) -> None:

zenml/integrations/mlflow/__init__.py

@@ -45,9 +45,11 @@ class MlflowIntegration(Integration):
         # This downgrades pydantic to v1 even though mlflow does not have
         # any issues with v2. This is why we have to pin it here so a downgrade
         # will not happen.
-        "pydantic>=2.7.0,<2.8.0"
+        "pydantic>=2.7.0,<2.8.0",
     ]
 
+    REQUIREMENTS_IGNORED_ON_UNINSTALL = ["python-rapidjson", "pydantic"]
+
     @classmethod
     def activate(cls) -> None:
         """Activate the MLflow integration."""

zenml/integrations/neural_prophet/__init__.py

@@ -26,6 +26,7 @@ class NeuralProphetIntegration(Integration):
         "holidays>=0.4.1,<0.25.0",
         "tenacity!=8.4.0",  # https://github.com/jd/tenacity/issues/471
     ]
+    REQUIREMENTS_IGNORED_ON_UNINSTALL = ["tenacity"]
 
     @classmethod
     def activate(cls) -> None:

zenml/integrations/polars/__init__.py

@@ -25,6 +25,7 @@ class PolarsIntegration(Integration):
         "polars>=0.19.5",
         "pyarrow>=12.0.0"
     ]
+    REQUIREMENTS_IGNORED_ON_UNINSTALL = ["pyarrow"]
 
     @classmethod
     def activate(cls) -> None: