Fix shutdown of CRD process proxy (and subclasses) (#1234)

* Fix CRD/SparkOperator shutdown
* Add some troubleshooting, prefer using 'namespaced' methods
* Set EG_INHERITED_ENVS to include PATH for bootstrapping kernels
* Refactor k8s-based process proxies regarding status and termination
* Apply labels to spark operator executors
* Minor refactor to reference k8s kind names
* Cap pyzmq < 25 since jupyter_client 6 references removed code

Author: kevin-bates

enterprise_gateway/services/processproxies/crd.py

@@ -15,57 +15,49 @@
 class CustomResourceProcessProxy(KubernetesProcessProxy):
     """A custom resource process proxy."""
 
-    group = version = plural = None
-    custom_resource_template_name = None
-    kernel_resource_name = None
+    # Identifies the kind of object being managed by this process proxy.
+    # For these values we will prefer the values found in the 'kind' field
+    # of the object's metadata.  This attribute is strictly used to provide
+    # context to log messages.
+    object_kind = "CustomResourceDefinition"
 
     def __init__(self, kernel_manager: RemoteKernelManager, proxy_config: dict):
         """Initialize the proxy."""
         super().__init__(kernel_manager, proxy_config)
+        self.group = self.version = self.plural = None
+        self.kernel_resource_name = None
 
     async def launch_process(
         self, kernel_cmd: str, **kwargs: dict[str, Any] | None
     ) -> "CustomResourceProcessProxy":
         """Launch the process for a kernel."""
-        kwargs["env"][
-            "KERNEL_RESOURCE_NAME"
-        ] = self.kernel_resource_name = self._determine_kernel_pod_name(**kwargs)
+        self.kernel_resource_name = self._determine_kernel_pod_name(**kwargs)
+        kwargs["env"]["KERNEL_RESOURCE_NAME"] = self.kernel_resource_name
         kwargs["env"]["KERNEL_CRD_GROUP"] = self.group
         kwargs["env"]["KERNEL_CRD_VERSION"] = self.version
         kwargs["env"]["KERNEL_CRD_PLURAL"] = self.plural
 
         await super().launch_process(kernel_cmd, **kwargs)
         return self
 
-    def terminate_container_resources(self) -> bool | None:
-        """Terminate the resources for the container."""
-        result = None
-
-        if self.kernel_resource_name:
-            if self.delete_kernel_namespace and not self.kernel_manager.restarting:
-                body = client.V1DeleteOptions(
-                    grace_period_seconds=0, propagation_policy="Background"
-                )
-                v1_status = client.CoreV1Api().delete_namespace(
-                    name=self.kernel_namespace, body=body
-                )
-
-                if v1_status and v1_status.status:
-                    termination_status = ["Succeeded", "Failed", "Terminating"]
-                    if any(status in v1_status.status for status in termination_status):
-                        result = True
-            else:
-                try:
-                    delete_status = client.CustomObjectsApi().delete_cluster_custom_object(
-                        self.group, self.version, self.plurals, self.kernel_resource_name
-                    )
-
-                    result = delete_status and delete_status.get("status", None) == "Success"
-
-                except Exception as err:
-                    result = isinstance(err, client.rest.ApiException) and err.status == 404
-
-        if result:
-            self.kernel_resource_name = None
+    def delete_managed_object(self, termination_stati: list[str]) -> bool:
+        """Deletes the object managed by this process-proxy
+
+        A return value of True indicates the object is considered deleted,
+        otherwise a False or None value is returned.
+
+        Note: the caller is responsible for handling exceptions.
+        """
+        delete_status = client.CustomObjectsApi().delete_namespaced_custom_object(
+            self.group,
+            self.version,
+            self.kernel_namespace,
+            self.plural,
+            self.kernel_resource_name,
+            grace_period_seconds=0,
+            propagation_policy="Background",
+        )
+
+        result = delete_status and delete_status.get("status", None) in termination_stati
 
         return result

enterprise_gateway/services/processproxies/k8s.py

@@ -37,6 +37,12 @@ class KubernetesProcessProxy(ContainerProcessProxy):
     Kernel lifecycle management for Kubernetes kernels.
     """
 
+    # Identifies the kind of object being managed by this process proxy.
+    # For these values we will prefer the values found in the 'kind' field
+    # of the object's metadata.  This attribute is strictly used to provide
+    # context to log messages.
+    object_kind = "Pod"
+
     def __init__(self, kernel_manager: RemoteKernelManager, proxy_config: dict):
         """Initialize the proxy."""
         super().__init__(kernel_manager, proxy_config)
@@ -91,19 +97,37 @@ def get_container_status(self, iteration: int | None) -> str | None:
 
         if iteration:  # only log if iteration is not None (otherwise poll() is too noisy)
             self.log.debug(
-                "{}: Waiting to connect to k8s pod in namespace '{}'. "
-                "Name: '{}', Status: '{}', Pod IP: '{}', KernelID: '{}'".format(
-                    iteration,
-                    self.kernel_namespace,
-                    self.container_name,
-                    pod_status,
-                    self.assigned_ip,
-                    self.kernel_id,
-                )
+                f"{iteration}: Waiting to connect to k8s {self.object_kind.lower()} in "
+                f"namespace '{self.kernel_namespace}'. Name: '{self.container_name}', "
+                f"Status: '{pod_status}', Pod IP: '{self.assigned_ip}', KernelID: '{self.kernel_id}'"
             )
 
         return pod_status
 
+    def delete_managed_object(self, termination_stati: list[str]) -> bool:
+        """Deletes the object managed by this process-proxy
+
+        A return value of True indicates the object is considered deleted,
+        otherwise a False or None value is returned.
+
+        Note: the caller is responsible for handling exceptions.
+        """
+        body = client.V1DeleteOptions(grace_period_seconds=0, propagation_policy="Background")
+
+        # Deleting a Pod will return a v1.Pod if found and its status will be a PodStatus containing
+        # a phase string property
+        # https://kubernetes.io/docs/reference/generated/kubernetes-api/v1.21/#podstatus-v1-core
+        v1_pod = client.CoreV1Api().delete_namespaced_pod(
+            namespace=self.kernel_namespace, body=body, name=self.container_name
+        )
+        status = None
+        if v1_pod and v1_pod.status:
+            status = v1_pod.status.phase
+
+        result = status in termination_stati
+
+        return result
+
     def terminate_container_resources(self) -> bool | None:
         """Terminate any artifacts created on behalf of the container's lifetime."""
         # Kubernetes objects don't go away on their own - so we need to tear down the namespace
@@ -114,46 +138,36 @@ def terminate_container_resources(self) -> bool | None:
         # from the pod deletion API, since it's not necessarily reflective of the actual status.
 
         result = False
-        body = client.V1DeleteOptions(grace_period_seconds=0, propagation_policy="Background")
+        termination_stati = ["Succeeded", "Failed", "Terminating", "Success"]
 
-        # Delete the pod then, if applicable, the namespace
+        # Delete the managed object then, if applicable, the namespace
+        object_type = self.object_kind
         try:
-            object_name = "pod"
-            status = None
-            termination_stati = ["Succeeded", "Failed", "Terminating"]
-
-            # Deleting a Pod will return a v1.Pod if found and its status will be a PodStatus containing
-            # a phase string property
-            # https://kubernetes.io/docs/reference/generated/kubernetes-api/v1.21/#podstatus-v1-core
-            v1_pod = client.CoreV1Api().delete_namespaced_pod(
-                namespace=self.kernel_namespace, body=body, name=self.container_name
-            )
-            if v1_pod and v1_pod.status:
-                status = v1_pod.status.phase
-
-            if status in termination_stati:
-                result = True
-
+            result = self.delete_managed_object(termination_stati)
             if not result:
-                # If the status indicates the pod is not terminated, capture its current status.
+                # If the status indicates the object is not terminated, capture its current status.
                 # If None, update the result to True, else issue warning that it is not YET deleted
                 # since we still have the hard termination sequence to occur.
                 cur_status = self.get_container_status(None)
                 if cur_status is None:
                     result = True
                 else:
                     self.log.warning(
-                        f"Pod {self.kernel_namespace}.{self.container_name} is not yet deleted.  "
-                        f"Current status is '{cur_status}'."
+                        f"{object_type} '{self.kernel_namespace}.{self.container_name}'"
+                        f" is not yet deleted.  Current status is '{cur_status}'."
                     )
 
             if self.delete_kernel_namespace and not self.kernel_manager.restarting:
-                object_name = "namespace"
+                object_type = "Namespace"
                 # Status is a return value for calls that don't return other objects.
                 # https://kubernetes.io/docs/reference/generated/kubernetes-api/v1.21/#status-v1-meta
+                body = client.V1DeleteOptions(
+                    grace_period_seconds=0, propagation_policy="Background"
+                )
                 v1_status = client.CoreV1Api().delete_namespace(
                     name=self.kernel_namespace, body=body
                 )
+                status = None
                 if v1_status:
                     status = v1_status.status
 
@@ -169,25 +183,23 @@ def terminate_container_resources(self) -> bool | None:
 
         except Exception as err:
             if isinstance(err, client.rest.ApiException) and err.status == 404:
-                result = True  # okay if its not found
+                result = True  # okay if it's not found
             else:
-                self.log.warning(f"Error occurred deleting {object_name}: {err}")
+                self.log.warning(f"Error occurred deleting {object_type.lower()}: {err}")
 
         if result:
             self.log.debug(
-                "KubernetesProcessProxy.terminate_container_resources, pod: {}.{}, kernel ID: {} has "
-                "been terminated.".format(
-                    self.kernel_namespace, self.container_name, self.kernel_id
-                )
+                f"KubernetesProcessProxy.terminate_container_resources, "
+                f"{self.object_kind}: {self.kernel_namespace}.{self.container_name}, "
+                f"kernel ID: {self.kernel_id} has been terminated."
             )
             self.container_name = None
             result = None  # maintain jupyter contract
         else:
             self.log.warning(
-                "KubernetesProcessProxy.terminate_container_resources, pod: {}.{}, kernel ID: {} has "
-                "not been terminated.".format(
-                    self.kernel_namespace, self.container_name, self.kernel_id
-                )
+                "KubernetesProcessProxy.terminate_container_resources, "
+                f"{self.object_kind}: {self.kernel_namespace}.{self.container_name}, "
+                f"kernel ID: {self.kernel_id} has not been terminated."
             )
 
         # Check if there's a kernel pod template file for this kernel and silently delete it.

enterprise_gateway/services/processproxies/spark_operator.py

@@ -3,48 +3,22 @@
 # Distributed under the terms of the Modified BSD License.
 from __future__ import annotations
 
-import os
-
 from ..kernels.remotemanager import RemoteKernelManager
-from .crd import CustomResourceProcessProxy, client
-
-enterprise_gateway_namespace = os.environ.get("EG_NAMESPACE", "default")
+from .crd import CustomResourceProcessProxy
 
 
 class SparkOperatorProcessProxy(CustomResourceProcessProxy):
     """Spark operator process proxy."""
 
+    # Identifies the kind of object being managed by this process proxy.
+    # For these values we will prefer the values found in the 'kind' field
+    # of the object's metadata.  This attribute is strictly used to provide
+    # context to log messages.
+    object_kind = "SparkApplication"
+
     def __init__(self, kernel_manager: RemoteKernelManager, proxy_config: dict):
         """Initialize the proxy."""
         super().__init__(kernel_manager, proxy_config)
         self.group = "sparkoperator.k8s.io"
         self.version = "v1beta2"
         self.plural = "sparkapplications"
-
-    def get_container_status(self, iteration: int) -> str:
-        """Get the container status for a given iteration."""
-        pod_status = pod_info = None
-
-        try:
-            custom_resource = client.CustomObjectsApi().get_namespaced_custom_object(
-                self.group,
-                self.version,
-                self.kernel_namespace,
-                self.plural,
-                self.kernel_resource_name,
-            )
-
-            if custom_resource:
-                pod_name = custom_resource["status"]["driverInfo"]["podName"]
-                pod_info = client.CoreV1Api().read_namespaced_pod(pod_name, self.kernel_namespace)
-        except Exception:
-            pass
-
-        if pod_info and pod_info.status:
-            pod_status = pod_info.status.phase
-            if pod_status == "Running" and self.assigned_host == "":
-                self.assigned_ip = pod_info.status.pod_ip
-                self.assigned_host = pod_info.metadata.name
-                self.assigned_node_ip = pod_info.status.host_ip
-
-        return pod_status

etc/kernel-launchers/operators/scripts/sparkoperator.k8s.io-v1beta2.yaml.j2

@@ -29,6 +29,10 @@ spec:
     coreLimit: 1000m
     memory: 1g
   executor:
+    labels:
+      kernel_id: "{{ kernel_id }}"
+      app: enterprise-gateway
+      component: worker
     image: {{ kernel_executor_image }}
     instances: 2
     cores: 1

etc/kubernetes/helm/enterprise-gateway/values.yaml

@@ -62,7 +62,7 @@ deployment:
   tolerations: []
   affinity: {}
   nodeSelector: {}
-  extraEnv: {}
+  extraEnv: {
     # SOME_ENV_VAR_WITH_SIMPLE_VALUE: "example"
     # SOME_ENV_VAR_WITH_LONG_VALUE: >
     #   'this example',
@@ -72,6 +72,8 @@ deployment:
     #   this example
     #   will preserve
     #   line breaks
+    EG_INHERITED_ENVS: "PATH"
+  }
 
 # Log output level.
 logLevel: DEBUG

pyproject.toml

@@ -33,7 +33,7 @@ dependencies = [
   "paramiko>=2.11",
   "pexpect>=4.8.0",
   "pycryptodomex>=3.9.7",
-  "pyzmq>=17.0",
+  "pyzmq>=17.0,<25.0",  # Pyzmq 25 removes deprecated code that jupyter_client 6 uses, remove if v6 gets updated
   "requests>=2.14.2",
   "tornado>=6.1",
   "traitlets>=5.3.0",