Fixed all broken things

Author: AlexejPenner

src/zenml/container_registries/base_container_registry.py

@@ -273,9 +273,8 @@ def registry_server_uri(self) -> str:
         if not registry_uri.startswith(("http://", "https://")):
             domain = registry_uri.split("/")[0]
             return f"https://{domain}"
-        
-        return registry_uri
 
+        return registry_uri
 
 
 class BaseContainerRegistryFlavor(Flavor):

src/zenml/integrations/kubernetes/orchestrators/kube_utils.py

@@ -386,6 +386,7 @@ def create_secret(
     namespace: str,
     secret_name: str,
     data: Dict[str, Optional[str]],
+    secret_type: str = "Opaque",
 ) -> None:
     """Create a Kubernetes secret.
 
@@ -394,10 +395,13 @@ def create_secret(
         namespace: The namespace in which to create the secret.
         secret_name: The name of the secret to create.
         data: The secret data.
+        secret_type: The secret type.
     """
     core_api.create_namespaced_secret(
         namespace=namespace,
-        body=build_secret_manifest(name=secret_name, data=data),
+        body=build_secret_manifest(
+            name=secret_name, data=data, secret_type=secret_type
+        ),
     )
 
 
@@ -416,10 +420,20 @@ def update_secret(
         data: The secret data. If the value is None, the key will be removed
             from the secret.
     """
+    # For updates, we only patch the data field, not the type (which is immutable)
+    import base64
+
+    encoded_data = {
+        key: base64.b64encode(value.encode()).decode() if value else None
+        for key, value in data.items()
+    }
+
+    patch_body = {"data": encoded_data}
+
     core_api.patch_namespaced_secret(
         namespace=namespace,
         name=secret_name,
-        body=build_secret_manifest(name=secret_name, data=data),
+        body=patch_body,
     )
 
 
@@ -428,6 +442,7 @@ def create_or_update_secret(
     namespace: str,
     secret_name: str,
     data: Dict[str, Optional[str]],
+    secret_type: str = "Opaque",
 ) -> None:
     """Create a Kubernetes secret if it doesn't exist, or update it if it does.
 
@@ -437,17 +452,35 @@ def create_or_update_secret(
         secret_name: The name of the secret to create or update.
         data: The secret data. If the value is None, the key will be removed
             from the secret.
+        secret_type: The secret type.
 
     Raises:
         ApiException: If the secret creation failed for any reason other than
             the secret already existing.
     """
     try:
-        create_secret(core_api, namespace, secret_name, data)
+        create_secret(core_api, namespace, secret_name, data, secret_type)
     except ApiException as e:
         if e.status != 409:
             raise
-        update_secret(core_api, namespace, secret_name, data)
+
+        # Check if the existing secret has a different type
+        try:
+            existing_secret = core_api.read_namespaced_secret(
+                name=secret_name, namespace=namespace
+            )
+            if existing_secret.type != secret_type:
+                # Delete the secret if the type is different (type is immutable)
+                delete_secret(core_api, namespace, secret_name)
+                create_secret(
+                    core_api, namespace, secret_name, data, secret_type
+                )
+            else:
+                # Same type, just update the data
+                update_secret(core_api, namespace, secret_name, data)
+        except ApiException:
+            # If we can't read the existing secret, try to update anyway
+            update_secret(core_api, namespace, secret_name, data)
 
 
 def create_or_update_secret_from_manifest(
@@ -466,30 +499,37 @@ def create_or_update_secret_from_manifest(
     """
     namespace = secret_manifest["metadata"]["namespace"]
     secret_name = secret_manifest["metadata"]["name"]
-    
+
     # Extract data from manifest - handle both 'data' and 'stringData' fields
     secret_data = {}
-    
+
     # Handle base64-encoded 'data' field
     if "data" in secret_manifest:
         import base64
+
         for key, encoded_value in secret_manifest["data"].items():
             if encoded_value is not None:
                 # Decode base64 data back to string
-                secret_data[key] = base64.b64decode(encoded_value).decode('utf-8')
+                secret_data[key] = base64.b64decode(encoded_value).decode(
+                    "utf-8"
+                )
             else:
                 secret_data[key] = None
-    
+
     # Handle plain text 'stringData' field
     if "stringData" in secret_manifest:
         secret_data.update(secret_manifest["stringData"])
-    
+
+    # Extract secret type from manifest, default to "Opaque" if not specified
+    secret_type = secret_manifest.get("type", "Opaque")
+
     # Use the existing create_or_update_secret function
     create_or_update_secret(
         core_api=core_api,
         namespace=namespace,
         secret_name=secret_name,
         data=secret_data,
+        secret_type=secret_type,
     )
 
 

src/zenml/integrations/kubernetes/orchestrators/kubernetes_orchestrator.py

@@ -532,12 +532,12 @@ def submit_pipeline(
                     "schedule. Use `Schedule(cron_expression=...)` instead."
                 )
             cron_expression = deployment.schedule.cron_expression
-            
+
             # Prepare dependencies for pod creation
-            registry_credentials, local_stores_path = self._prepare_pod_dependencies(
-                settings, stack
+            registry_credentials, local_stores_path = (
+                self._prepare_pod_dependencies(settings, stack)
             )
-            
+
             cron_job_manifest, secret_manifests = build_cron_job_manifest(
                 cron_expression=cron_expression,
                 pod_name=pod_name,
@@ -574,13 +574,12 @@ def submit_pipeline(
             return None
         else:
             # Prepare dependencies for pod creation
-            registry_credentials, local_stores_path = self._prepare_pod_dependencies(
-                settings, stack
+            registry_credentials, local_stores_path = (
+                self._prepare_pod_dependencies(settings, stack)
             )
-            
+
             # Create and run the orchestrator pod.
             pod_manifest, secret_manifests = build_pod_manifest(
-                run_name=orchestrator_run_name,
                 pod_name=pod_name,
                 image_name=image,
                 command=command,
@@ -675,16 +674,14 @@ def _get_service_account_name(
             return service_account_name
 
     def _prepare_pod_dependencies(
-        self, 
-        settings: KubernetesOrchestratorSettings,
-        stack: "Stack"
+        self, settings: KubernetesOrchestratorSettings, stack: "Stack"
     ) -> Tuple[Optional[Tuple[str, str, str]], Optional[str]]:
         """Prepare dependencies needed for pod manifest creation.
-        
+
         Args:
             settings: The orchestrator settings.
             stack: The stack the pipeline will run on.
-            
+
         Returns:
             Tuple of (registry_credentials, local_stores_path).
         """
@@ -694,23 +691,23 @@ def _prepare_pod_dependencies(
             registry_credentials = self.get_kubernetes_image_pull_secret_data(
                 stack.container_registry
             )
-        
+
         # Get local stores path if mounting local stores
         local_stores_path = None
         if self.config.is_local:
             from zenml.config.global_config import GlobalConfiguration
-            
+
             stack.check_local_paths()
             local_stores_path = GlobalConfiguration().local_stores_path
-            
+
         return registry_credentials, local_stores_path
 
     def _create_image_pull_secrets(
         self,
         secret_manifests: List[Dict[str, str]],
     ) -> None:
         """Create imagePullSecrets in the cluster.
-        
+
         Args:
             secret_manifests: List of secret manifests for imagePullSecrets.
         """
@@ -750,7 +747,7 @@ def get_kubernetes_image_pull_secret_data(
 
         Returns:
             Tuple of (registry_uri, username, password) if credentials are available,
-            None otherwise. The registry_uri is normalized for use in Kubernetes 
+            None otherwise. The registry_uri is normalized for use in Kubernetes
             imagePullSecrets.
         """
         from zenml.logger import get_logger

src/zenml/integrations/kubernetes/orchestrators/kubernetes_orchestrator_entrypoint.py

@@ -239,6 +239,15 @@ def run_step_on_kubernetes(step_name: str) -> None:
 
         logger.info(f"Container registry: {active_stack.container_registry}")
 
+        # Get registry credentials for imagePullSecrets
+        registry_credentials = None
+        if pipeline_settings.auto_generate_image_pull_secrets:
+            registry_credentials = (
+                orchestrator.get_kubernetes_image_pull_secret_data(
+                    active_stack.container_registry
+                )
+            )
+
         # Define Kubernetes pod manifest and any required secrets.
         pod_manifest, secret_manifests = build_pod_manifest(
             pod_name=pod_name,
@@ -252,20 +261,18 @@ def run_step_on_kubernetes(step_name: str) -> None:
             or settings.service_account_name,
             mount_local_stores=mount_local_stores,
             owner_references=owner_references,
-            namespace=args.kubernetes_namespace,
+            namespace=namespace,
             auto_generate_image_pull_secrets=pipeline_settings.auto_generate_image_pull_secrets,
-            container_registry=active_stack.container_registry,
+            registry_credentials=registry_credentials,
             core_api=core_api,
+            labels=step_pod_labels,
         )
 
         # Step pods should reuse secrets created by the orchestrator pod
         # Only create secrets if they don't already exist
         create_image_pull_secrets_from_manifests(
             secret_manifests=secret_manifests,
             core_api=core_api,
-            namespace=args.kubernetes_namespace,
-            reuse_existing=True,  # Step pods reuse orchestrator-created secrets
-            labels=step_pod_labels,
         )
 
         kube_utils.create_and_wait_for_pod_to_start(
@@ -386,7 +393,7 @@ def finalize_run(node_states: Dict[str, NodeStatus]) -> None:
             logger.info("Cleaning up old imagePullSecrets...")
             cleanup_old_image_pull_secrets(
                 core_api=core_api,
-                namespace=args.kubernetes_namespace,
+                namespace=namespace,
                 max_age_hours=24,  # Keep secrets for 24 hours
             )
 

src/zenml/integrations/kubernetes/orchestrators/manifest_utils.py

@@ -26,6 +26,7 @@
 from zenml.integrations.airflow.orchestrators.dag_generator import (
     ENV_ZENML_LOCAL_STORES_PATH,
 )
+from zenml.integrations.kubernetes.orchestrators import kube_utils
 from zenml.integrations.kubernetes.pod_settings import KubernetesPodSettings
 from zenml.logger import get_logger
 
@@ -143,26 +144,25 @@ def _generate_image_pull_secrets(
             safe_registry_name = safe_registry_name[8:]
         elif safe_registry_name.startswith("http://"):
             safe_registry_name = safe_registry_name[7:]
-        
+
         # Replace invalid characters for Kubernetes names
         safe_registry_name = (
-            safe_registry_name
-            .replace(".", "-")
+            safe_registry_name.replace(".", "-")
             .replace("/", "-")
             .replace(":", "-")
             .replace("_", "-")
             .lower()
         )
-        
+
         # Ensure it starts and ends with alphanumeric character
         safe_registry_name = safe_registry_name.strip("-")
         if not safe_registry_name:
             safe_registry_name = "registry"
-            
+
         secret_name = f"zenml-registry-{safe_registry_name}-{i}"[
             :63
         ]  # K8s name limit
-        
+
         # Final validation: ensure name is valid for Kubernetes
         # Must start and end with alphanumeric character
         if not secret_name[0].isalnum():
@@ -239,7 +239,7 @@ def create_image_pull_secrets_from_manifests(
     """
     for secret_manifest in secret_manifests:
         secret_name = secret_manifest["metadata"]["name"]
-        
+
         try:
             kube_utils.create_or_update_secret_from_manifest(
                 core_api=core_api,
@@ -449,12 +449,10 @@ def build_pod_manifest(
     # Auto-generate imagePullSecrets from container registry credentials
     if auto_generate_image_pull_secrets and registry_credentials:
         try:
-            generated_secrets, generated_refs = (
-                _generate_image_pull_secrets(
-                    namespace=namespace,
-                    registry_credentials=registry_credentials,
-                    core_api=core_api,
-                )
+            generated_secrets, generated_refs = _generate_image_pull_secrets(
+                namespace=namespace,
+                registry_credentials=registry_credentials,
+                core_api=core_api,
             )
             secret_manifests.extend(generated_secrets)
             image_pull_secrets.extend(generated_refs)