feat: add gracefully shutdown support

also fixed 2 bug around resetting nodes

Author: lockwobr

.github/workflows/operator-ci.yaml

@@ -49,7 +49,7 @@ on:
 env:
   REGISTRY: ghcr.io
   IMAGE_NAME: ${{ github.repository }}
-  GO_VERSION: 1.23.7
+  GO_VERSION: 1.23.8
   PLATFORMS: linux/amd64,linux/arm64
 
 # There is a single job in this workflow. It's configured to run on the latest available version of Ubuntu.

agent/skyhook-agent/src/skyhook_agent/controller.py

@@ -33,13 +33,26 @@
 import os
 import shutil
 import glob
+import signal
 
 from skyhook_agent.step import Step, UpgradeStep, Idempotence, Mode, CHECK_TO_APPLY
 from skyhook_agent import interrupts, config
 from typing import List
 
 import logging as logger
 
+# Global flag to track if we received SIGTERM
+received_sigterm = False
+
+def sigterm_handler(signum, frame):
+    """Handle SIGTERM by setting a global flag and logging the event"""
+    global received_sigterm
+    received_sigterm = True
+    logger.info("Received SIGTERM signal - initiating graceful shutdown")
+
+# Register the SIGTERM handler
+signal.signal(signal.SIGTERM, sigterm_handler)
+
 class SkyhookValidationError(Exception):
     pass
 
@@ -414,7 +427,11 @@ def remove_flags(step_data: dict[Mode, list[Step|UpgradeStep]], config_data: dic
         if os.path.exists(flag_file):  # Check if the file exists before trying to remove it
             os.remove(flag_file)
 
-def main(mode: Mode, root_mount: str, copy_dir: str, interrupt_data: None|str, always_run_step=False):
+def main(mode: Mode, root_mount: str, copy_dir: str, interrupt_data: None|str, always_run_step=False) -> bool:
+    '''
+    returns True if the there is a failure in the steps, otherwise returns False
+    '''
+
     if mode not in set(map(str, Mode)):
         logger.warning(f"This version of the Agent doesn't support the {mode} mode. Options are: {','.join(map(str, Mode))}.")
         return False
@@ -448,9 +465,19 @@ def main(mode: Mode, root_mount: str, copy_dir: str, interrupt_data: None|str, a
             if not os.path.exists(f"{root_mount}/{copy_dir}/configmaps/{f}"):
                 raise SkyhookValidationError(f"Expected config file {f} not found in configmaps directory.")
 
-    return agent_main(mode, root_mount, copy_dir, config_data, interrupt_data, always_run_step)
+    try:
+        return agent_main(mode, root_mount, copy_dir, config_data, interrupt_data, always_run_step)
+    except Exception as e:
+        if received_sigterm:
+            logger.info("Gracefully shutting down due to SIGTERM")
+            # Perform any cleanup if needed
+            return True
+        raise
 
-def agent_main(mode: Mode, root_mount: str, copy_dir: str, config_data: dict, interrupt_data: None|str, always_run_step=False):
+def agent_main(mode: Mode, root_mount: str, copy_dir: str, config_data: dict, interrupt_data: None|str, always_run_step=False) -> bool:
+    '''
+    returns True if the there is a failure in the steps, otherwise returns False
+    '''
 
     # Pull out step_data so it matches with existing code
     step_data = config_data["modes"]
@@ -464,6 +491,11 @@ def agent_main(mode: Mode, root_mount: str, copy_dir: str, config_data: dict, in
         logger.warning(f" There are no {mode} steps defined. This will be ran as a no-op.")
 
     for step in step_data.get(mode, []):
+        # Check for SIGTERM
+        if received_sigterm:
+            logger.info("SIGTERM received, stopping step execution")
+            return True
+
         # Make the flag file without the host path argument (first one). This is because in operator world
         # the host path is going to change every time the Skyhook Custom Resource changes so it would
         # look like a step hasn't been run when it fact it had.

chart/templates/skyhook-crd.yaml

@@ -349,6 +349,10 @@ spec:
                         - name
                         type: object
                       type: array
+                    gracefulShutdown:
+                      description: GracefulShutdown is the graceful shutdown timeout
+                        for the package, if not set, uses k8s default
+                      type: string
                     image:
                       description: Image is the container image to run. Do not included
                         the tag, that is set in the version.

containers/agentless/entrypoint.sh

@@ -20,6 +20,13 @@
 # LICENSE END
 # 
 
+# Handle SIGTERM gracefully
+cleanup() {
+    echo "Received SIGTERM signal, shutting down gracefully..."
+    sleep 3
+    exit 0
+}
+trap cleanup SIGTERM
 
 SLEEP_LEN=${SLEEP_LEN:-$(($RANDOM % 5 + 5))}
 

containers/ci.Dockerfile

@@ -18,6 +18,10 @@
 # LICENSE END
 # 
 
+## this container is not used in github CI, was from before before we open sourced this project.
+## should move to doing something like this in the github actions workflow to save time not installing all the deps all the time
+## but for now this is just for when we got to that
+
 ARG GO_VERSION
 
 FROM golang:${GO_VERSION}-bookworm as builder

k8s-tests/chainsaw/skyhook/cleanup-pods/assert-cleaned-pods.yaml

@@ -0,0 +1,165 @@
+# 
+# LICENSE START
+#
+#    Copyright (c) NVIDIA CORPORATION.  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.
+#    You may obtain a copy of the License at
+#
+#        http://www.apache.org/licenses/LICENSE-2.0
+#
+#    Unless required by applicable law or agreed to in writing, software
+#    distributed under the License is distributed on an "AS IS" BASIS,
+#    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.
+#
+# LICENSE END
+# 
+---
+kind: Pod
+apiVersion: v1
+metadata:
+  namespace: skyhook
+  labels:
+    skyhook.nvidia.com/name: cleanup-pods
+    skyhook.nvidia.com/package: aa-1.2.3
+  annotations:
+    ("skyhook.nvidia.com/package" && parse_json("skyhook.nvidia.com/package")):
+      {
+        "name": "aa",
+        "version": "1.2.3",
+        "skyhook": "cleanup-pods",
+        "stage": "apply",
+        "image": "ghcr.io/nvidia/skyhook/agentless"
+      }
+  ownerReferences:
+    - apiVersion: skyhook.nvidia.com/v1alpha1
+      kind: Skyhook
+      name: cleanup-pods
+spec:
+  nodeName: kind-worker
+  initContainers:
+    - name: aa-init
+    - name: aa-apply
+      args:
+        ([0]): apply
+        ([1]): /root
+        (length(@)): 3
+    - name: aa-applycheck
+      args:
+        ([0]): apply-check
+        ([1]): /root
+        (length(@)): 3
+---
+kind: Pod
+apiVersion: v1
+metadata:
+  namespace: skyhook
+  labels:
+    skyhook.nvidia.com/name: cleanup-pods
+    skyhook.nvidia.com/package: aa-1.2.3
+  annotations:
+    ("skyhook.nvidia.com/package" && parse_json("skyhook.nvidia.com/package")):
+      {
+        "name": "aa",
+        "version": "1.2.3",
+        "skyhook": "cleanup-pods",
+        "stage": "config",
+        "image": "ghcr.io/nvidia/skyhook/agentless"
+      }
+  ownerReferences:
+    - apiVersion: skyhook.nvidia.com/v1alpha1
+      kind: Skyhook
+      name: cleanup-pods
+spec:
+  nodeName: kind-worker
+  initContainers:
+    - name: aa-init
+    - name: aa-config
+      args:
+        ([0]): config
+        ([1]): /root
+        (length(@)): 3
+    - name: aa-configcheck
+      args:
+        ([0]): config-check
+        ([1]): /root
+        (length(@)): 3
+---
+kind: Pod
+apiVersion: v1
+metadata:
+  namespace: skyhook
+  labels:
+    skyhook.nvidia.com/name: cleanup-pods
+    skyhook.nvidia.com/package: bb-1.2
+  annotations:
+    ("skyhook.nvidia.com/package" && parse_json("skyhook.nvidia.com/package")):
+      {
+        "name": "bb",
+        "version": "1.2",
+        "skyhook": "cleanup-pods",
+        "stage": "apply",
+        "image": "ghcr.io/nvidia/skyhook/agentless"
+      }
+  ownerReferences:
+    - apiVersion: skyhook.nvidia.com/v1alpha1
+      kind: Skyhook
+      name: cleanup-pods
+spec:
+  nodeName: kind-worker
+  initContainers:
+    - name: bb-init
+    - name: bb-apply
+      args:
+        ([0]): apply
+        ([1]): /root
+        (length(@)): 3
+    - name: bb-applycheck
+      args:
+        ([0]): apply-check
+        ([1]): /root
+        (length(@)): 3
+---
+apiVersion: skyhook.nvidia.com/v1alpha1
+kind: Skyhook
+metadata:
+  name: cleanup-pods
+status:
+  status: erroring
+  observedGeneration: 4
+  nodeState:
+    kind-worker:
+      aa|1.2.3:
+        name: aa
+        state: complete
+        version: '1.2.3'
+        image: ghcr.io/nvidia/skyhook/agentless
+        stage: config
+      bb|1.2:
+        name: bb
+        state: erroring
+        version: '1.2'
+        image: ghcr.io/nvidia/skyhook/agentless
+        stage: apply
+    kind-worker2:
+      aa|1.2.3:
+        name: aa
+        state: complete
+        version: '1.2.3'
+        image: ghcr.io/nvidia/skyhook/agentless
+        stage: config
+      bb|1.2:
+        name: bb
+        state: erroring
+        version: '1.2'
+        image: ghcr.io/nvidia/skyhook/agentless
+        stage: config
+    
+  nodeStatus:
+   # grab values should be one and is complete
+    (values(@)):
+      - erroring
+      - erroring

k8s-tests/chainsaw/skyhook/cleanup-pods/assert-config-complete.yaml

@@ -0,0 +1,82 @@
+# 
+# LICENSE START
+#
+#    Copyright (c) NVIDIA CORPORATION.  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.
+#    You may obtain a copy of the License at
+#
+#        http://www.apache.org/licenses/LICENSE-2.0
+#
+#    Unless required by applicable law or agreed to in writing, software
+#    distributed under the License is distributed on an "AS IS" BASIS,
+#    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.
+#
+# LICENSE END
+# 
+
+---
+kind: Pod
+apiVersion: v1
+metadata:
+  namespace: skyhook
+  labels:
+    skyhook.nvidia.com/name: cleanup-pods
+    skyhook.nvidia.com/package: bb-1.2
+  annotations:
+    ("skyhook.nvidia.com/package" && parse_json("skyhook.nvidia.com/package")):
+      {
+        "name": "bb",
+        "version": "1.2",
+        "skyhook": "cleanup-pods",
+        "stage": "config",
+        "image": "ghcr.io/nvidia/skyhook/agentless"
+      }
+  ownerReferences:
+    - apiVersion: skyhook.nvidia.com/v1alpha1
+      kind: Skyhook
+      name: cleanup-pods
+spec:
+  terminationGracePeriodSeconds: 46
+  initContainers:
+    - name: bb-init
+    - name: bb-config
+      args:
+        ([0]): config
+        ([1]): /root
+        (length(@)): 3
+    - name: bb-configcheck
+      args:
+        ([0]): config-check
+        ([1]): /root
+        (length(@)): 3
+---
+apiVersion: skyhook.nvidia.com/v1alpha1
+kind: Skyhook
+metadata:
+  name: cleanup-pods
+status:
+  status: complete
+  observedGeneration: 4
+  nodeState:
+    (values(@)):
+    - aa|1.2.3:
+        image: ghcr.io/nvidia/skyhook/agentless
+        name: aa
+        stage: config
+        state: complete
+        version: 1.2.3
+    - bb|1.2:
+        image: ghcr.io/nvidia/skyhook/agentless
+        name: bb
+        stage: config
+        state: complete
+        version: "1.2"
+  nodeStatus:
+   # grab values should be one and is complete
+    (values(@)):
+      - complete
+      - complete