diff --git a/experimental/recipes/flake.nix b/experimental/recipes/flake.nix new file mode 100644 index 0000000..4ee49dd --- /dev/null +++ b/experimental/recipes/flake.nix @@ -0,0 +1,22 @@ +{ + inputs = { + nixpkgs.url = "github:NixOS/nixpkgs/nixpkgs-unstable"; + systems.url = "github:nix-systems/default"; + }; + outputs = { self, ... } @ inputs: let + forEachSystem = inputs.nixpkgs.lib.genAttrs (import inputs.systems); + in { + packages = forEachSystem (system: let + pkgs = inputs.nixpkgs.legacyPackages.${system}; + in { + graceful-shutdown = pkgs.python3Packages.buildPythonPackage { + pname = "graceful-shutdown"; + version = "0.1.0"; + src = ./src; + propagatedBuildInputs = [ + pkgs.python3Packages.kubernetes + ]; + }; + }); + }; +} \ No newline at end of file diff --git a/experimental/recipes/src/main.py b/experimental/recipes/src/main.py new file mode 100644 index 0000000..45df1a8 --- /dev/null +++ b/experimental/recipes/src/main.py @@ -0,0 +1,111 @@ +#!/usr/bin/env python3 +import argparse +import sys +import time +import subprocess + +from kubernetes import client, config + +def is_control_plane(node): + """Check if node has control-plane/master labels.""" + labels = node.metadata.labels or {} + return any(k in labels for k in [ + "node-role.kubernetes.io/control-plane", + "node-role.kubernetes.io/master" + ]) + +def cordon_node(v1, node_name): + body = {"spec": {"unschedulable": True}} + v1.patch_node(node_name, body) + +def drain_node(v1, node_name, timeout=600, poll_interval=5): + """Evict all pods from node and wait until drained.""" + field_selector = f"spec.nodeName={node_name}" + pods = v1.list_pod_for_all_namespaces(field_selector=field_selector).items + + # Evict pods one by one + policy = client.PolicyV1Api() + for pod in pods: + # Skip mirror pods (static pods) + if pod.metadata.owner_references and any( + owner.kind == "DaemonSet" for owner in pod.metadata.owner_references + ): + continue + if pod.metadata.annotations and "kubernetes.io/config.mirror" in pod.metadata.annotations: + continue + + eviction = client.V1Eviction( + metadata=client.V1ObjectMeta(name=pod.metadata.name, namespace=pod.metadata.namespace), + delete_options=client.V1DeleteOptions(grace_period_seconds=30) + ) + try: + policy.create_namespaced_pod_eviction( + name=pod.metadata.name, + namespace=pod.metadata.namespace, + body=eviction + ) + print(f"Eviction requested: {pod.metadata.namespace}/{pod.metadata.name}") + except Exception as e: + print(f"Failed to evict {pod.metadata.namespace}/{pod.metadata.name}: {e}") + + # Wait until no non-DaemonSet, non-mirror pods remain + start = time.time() + while True: + pods = v1.list_pod_for_all_namespaces(field_selector=field_selector).items + remaining = [ + p for p in pods + if not (p.metadata.owner_references and any(o.kind == "DaemonSet" for o in p.metadata.owner_references)) + and not (p.metadata.annotations and "kubernetes.io/config.mirror" in p.metadata.annotations) + ] + if not remaining: + print("Node successfully drained.") + return True + if time.time() - start > timeout: + print("Timeout waiting for node to drain.") + return False + print(f"Waiting for {len(remaining)} pods to evict...") + time.sleep(poll_interval) + +def shutdown_host(hostname, user="root"): + print(f"Shutting down host {hostname}...") + result = subprocess.run( + ["ssh", f"{user}@{hostname}", "sudo", "shutdown", "-h", "now"], + capture_output=True, text=True + ) + if result.returncode != 0: + print("Failed to shut down host:", result.stderr) + else: + print("Shutdown command issued successfully.") + +def main(): + parser = argparse.ArgumentParser(description="Drain a Kubernetes node and shut it down.") + parser.add_argument("node", help="Node hostname to drain and shut down") + parser.add_argument("--ssh-user", default="root", help="SSH username (default: root)") + args = parser.parse_args() + + # Load kubeconfig + config.load_kube_config() + v1 = client.CoreV1Api() + + # Verify node exists + try: + node = v1.read_node(args.node) + except Exception as e: + print(f"Error: Could not find node {args.node}: {e}") + sys.exit(1) + + # Ensure it's not control-plane + if is_control_plane(node): + print(f"Error: Node {args.node} is a control-plane node. Refusing to shut it down.") + sys.exit(1) + + # Cordon and drain + cordon_node(v1, args.node) + if not drain_node(v1, args.node): + sys.exit(1) + + # Shutdown + shutdown_host(args.node, user=args.ssh_user) + +if __name__ == "__main__": + main() diff --git a/experimental/recipes/src/setup.py b/experimental/recipes/src/setup.py new file mode 100644 index 0000000..7cd4ed4 --- /dev/null +++ b/experimental/recipes/src/setup.py @@ -0,0 +1,9 @@ +from setuptools import setup, find_packages + +setup(name='graceful_shutdown', + version='0.1.0', + # Modules to import from other scripts: + packages=find_packages(), + # Executables + scripts=["main.py"], +) \ No newline at end of file