From 90920e08cfc70d6e487760ddf06d370e1130c1b3 Mon Sep 17 00:00:00 2001 From: Alexander Tarasov Date: Wed, 7 Oct 2026 13:16:32 +0200 Subject: [PATCH 1/2] feat(client): Add drain command for the worker preStop hook Add `python3 -m taskbroker_client.drain `, which creates the worker drain file and then sleeps. Pods currently do this with `/bin/sh -c "touch /tmp/taskworker-draining && sleep N"`, which does not work on distroless images that have no shell, touch or sleep. The client owns the drain file path, so the hook no longer has to duplicate it. The module imports only the constants, so the hook starts in milliseconds instead of loading the worker. Co-Authored-By: Claude --- clients/python/src/taskbroker_client/drain.py | 37 ++++++++++++++++ clients/python/tests/test_drain.py | 42 +++++++++++++++++++ 2 files changed, 79 insertions(+) create mode 100644 clients/python/src/taskbroker_client/drain.py create mode 100644 clients/python/tests/test_drain.py diff --git a/clients/python/src/taskbroker_client/drain.py b/clients/python/src/taskbroker_client/drain.py new file mode 100644 index 00000000..d3081d0e --- /dev/null +++ b/clients/python/src/taskbroker_client/drain.py @@ -0,0 +1,37 @@ +""" +Mark the worker as draining, then wait. + +Meant to run as the worker pod's preStop hook: + + python3 -m taskbroker_client.drain 15 + +Kubernetes only sends SIGTERM once preStop finishes, so creating the drain file +first lets the worker stop publishing occupancy while brokers stop routing to +the pod. Imports nothing beyond the constants, so it starts fast and needs no +shell in the image. +""" + +import argparse +import time +from pathlib import Path + +from taskbroker_client.constants import DEFAULT_WORKER_DRAIN_FILE_PATH + + +def drain(seconds: float, path: str = DEFAULT_WORKER_DRAIN_FILE_PATH) -> None: + Path(path).touch() + time.sleep(seconds) + + +def main(argv: list[str] | None = None) -> None: + parser = argparse.ArgumentParser( + prog="python3 -m taskbroker_client.drain", + description="Create the worker drain file, then sleep.", + ) + parser.add_argument("seconds", type=float, help="how long to sleep after creating the file") + args = parser.parse_args(argv) + drain(args.seconds) + + +if __name__ == "__main__": + main() diff --git a/clients/python/tests/test_drain.py b/clients/python/tests/test_drain.py new file mode 100644 index 00000000..5dd379c6 --- /dev/null +++ b/clients/python/tests/test_drain.py @@ -0,0 +1,42 @@ +import os +import subprocess +import sys +from pathlib import Path +from unittest import mock + +from taskbroker_client.drain import drain, main + + +def test_drain_creates_file_before_sleeping(tmp_path: Path) -> None: + path = tmp_path / "draining" + + def sleep(seconds: float) -> None: + assert path.exists() + + with mock.patch("taskbroker_client.drain.time.sleep", side_effect=sleep) as sleep_mock: + drain(15, str(path)) + sleep_mock.assert_called_once_with(15) + + +def test_drain_touches_existing_file(tmp_path: Path) -> None: + path = tmp_path / "draining" + path.touch() + os.utime(path, (0, 0)) + with mock.patch("taskbroker_client.drain.time.sleep"): + drain(0, str(path)) + assert path.stat().st_mtime > 0 + + +def test_main_parses_seconds() -> None: + with mock.patch("taskbroker_client.drain.drain") as drain_mock: + main(["15"]) + drain_mock.assert_called_once_with(15.0) + + +def test_import_stays_light() -> None: + code = ( + "import sys, taskbroker_client.drain; " + "print(sorted(m for m in ('grpc', 'taskbroker_client.worker') if m in sys.modules))" + ) + out = subprocess.run([sys.executable, "-c", code], capture_output=True, text=True, check=True) + assert out.stdout.strip() == "[]" From 584a610d2709bf48f758548e0f15c63c5eecdddf Mon Sep 17 00:00:00 2001 From: Alexander Tarasov Date: Thu, 8 Oct 2026 13:52:03 +0200 Subject: [PATCH 2/2] ref(client): Trim drain module docstring and tests Replace the module docstring with a short one on drain(), since no other client module has one. Drop the test for touching an existing file, which the worker never relies on, and the import-weight test, which guarded about 0.5 s inside a 15 s hook. Co-Authored-By: Claude Opus 5.5 --- clients/python/src/taskbroker_client/drain.py | 17 ++++----------- clients/python/tests/test_drain.py | 21 ------------------- 2 files changed, 4 insertions(+), 34 deletions(-) diff --git a/clients/python/src/taskbroker_client/drain.py b/clients/python/src/taskbroker_client/drain.py index d3081d0e..4773c740 100644 --- a/clients/python/src/taskbroker_client/drain.py +++ b/clients/python/src/taskbroker_client/drain.py @@ -1,16 +1,3 @@ -""" -Mark the worker as draining, then wait. - -Meant to run as the worker pod's preStop hook: - - python3 -m taskbroker_client.drain 15 - -Kubernetes only sends SIGTERM once preStop finishes, so creating the drain file -first lets the worker stop publishing occupancy while brokers stop routing to -the pod. Imports nothing beyond the constants, so it starts fast and needs no -shell in the image. -""" - import argparse import time from pathlib import Path @@ -19,6 +6,10 @@ def drain(seconds: float, path: str = DEFAULT_WORKER_DRAIN_FILE_PATH) -> None: + """ + Create the drain file, then sleep. Runs as the worker pod's preStop hook, + which must finish before Kubernetes sends SIGTERM. + """ Path(path).touch() time.sleep(seconds) diff --git a/clients/python/tests/test_drain.py b/clients/python/tests/test_drain.py index 5dd379c6..3a92a6a9 100644 --- a/clients/python/tests/test_drain.py +++ b/clients/python/tests/test_drain.py @@ -1,6 +1,3 @@ -import os -import subprocess -import sys from pathlib import Path from unittest import mock @@ -18,25 +15,7 @@ def sleep(seconds: float) -> None: sleep_mock.assert_called_once_with(15) -def test_drain_touches_existing_file(tmp_path: Path) -> None: - path = tmp_path / "draining" - path.touch() - os.utime(path, (0, 0)) - with mock.patch("taskbroker_client.drain.time.sleep"): - drain(0, str(path)) - assert path.stat().st_mtime > 0 - - def test_main_parses_seconds() -> None: with mock.patch("taskbroker_client.drain.drain") as drain_mock: main(["15"]) drain_mock.assert_called_once_with(15.0) - - -def test_import_stays_light() -> None: - code = ( - "import sys, taskbroker_client.drain; " - "print(sorted(m for m in ('grpc', 'taskbroker_client.worker') if m in sys.modules))" - ) - out = subprocess.run([sys.executable, "-c", code], capture_output=True, text=True, check=True) - assert out.stdout.strip() == "[]"