From 21a446f70e09018f12036e3492229ea55613970b Mon Sep 17 00:00:00 2001 From: Louise Fox <208544511+folouiseAWS@users.noreply.github.com> Date: Wed, 15 Jul 2026 16:31:40 -0700 Subject: [PATCH] fix: Add timeout to IMDS requests to prevent shutdown hang The IMDSv2 token, spot instance-action, and ASG lifecycle-state requests had no timeout and only caught ConnectionError. On hosts where the IMDS address (169.254.169.254) is blackholed rather than refused, the token request in Worker.run() blocks the main thread indefinitely. The scheduler keeps running in the executor, so the agent appears healthy, but SIGTERM shutdown can never complete and the worker is never transitioned to STOPPED. Apply the same 0.5 second timeout already used by startup.bootstrap._get_metadata, and treat request timeouts the same as connection errors: IMDS is unavailable. Verified on a macOS host on a network that blackholes the IMDS address: without this change the agent hung indefinitely after SIGTERM; with it, the agent transitions to STOPPED and exits in under one second. Signed-off-by: Louise Fox <208544511+folouiseAWS@users.noreply.github.com> --- src/deadline_worker_agent/worker.py | 17 +++++++++++--- test/unit/test_worker.py | 36 +++++++++++++++++++++++++++++ 2 files changed, 50 insertions(+), 3 deletions(-) diff --git a/src/deadline_worker_agent/worker.py b/src/deadline_worker_agent/worker.py index cf8d9131e..eae41cb7d 100644 --- a/src/deadline_worker_agent/worker.py +++ b/src/deadline_worker_agent/worker.py @@ -61,6 +61,11 @@ class Worker: """The amount of time to allow the Worker to gracefully shutdown after detecting an auto-scaling life-cycle event.""" + _IMDS_REQUEST_TIMEOUT_SECONDS = 0.5 + """Timeout for EC2 instance metadata service requests. Non-EC2 hosts may blackhole the + IMDS address rather than refuse the connection; without a timeout those requests block + forever. Matches the timeout used by startup.bootstrap._get_metadata.""" + _farm_id: str _fleet_id: str _worker_id: str @@ -391,8 +396,12 @@ def _get_ec2_metadata_imdsv2_token(self) -> str | None: response = requests.put( "http://169.254.169.254/latest/api/token", headers={"X-aws-ec2-metadata-token-ttl-seconds": "10"}, + # Non-EC2 hosts may blackhole the IMDS address instead of refusing the + # connection. Without a timeout this request can block forever, which + # prevents the agent's graceful shutdown from ever running. + timeout=Worker._IMDS_REQUEST_TIMEOUT_SECONDS, ) - except requests.ConnectionError: + except (requests.ConnectionError, requests.Timeout): # Could not connect to the metadata service. Either it's not enabled or we're not # on an EC2 instance. return None @@ -422,8 +431,9 @@ def _get_spot_instance_shutdown_action_timeout(self, *, imdsv2_token: str) -> ti response = requests.get( "http://169.254.169.254/latest/meta-data/spot/instance-action", headers={"X-aws-ec2-metadata-token": imdsv2_token}, + timeout=Worker._IMDS_REQUEST_TIMEOUT_SECONDS, ) - except requests.ConnectionError: + except (requests.ConnectionError, requests.Timeout): # Could not connect to the metadata service. Either it's inactive or we're not # on an EC2 instance. return None @@ -473,8 +483,9 @@ def _is_asg_terminated(self, *, imdsv2_token: str) -> bool: response = requests.get( "http://169.254.169.254/latest/meta-data/autoscaling/target-lifecycle-state", headers={"X-aws-ec2-metadata-token": imdsv2_token}, + timeout=Worker._IMDS_REQUEST_TIMEOUT_SECONDS, ) - except requests.ConnectionError: + except (requests.ConnectionError, requests.Timeout): return False if response.status_code == 200: diff --git a/test/unit/test_worker.py b/test/unit/test_worker.py index 4c0e724dd..fa1f1b0a7 100644 --- a/test/unit/test_worker.py +++ b/test/unit/test_worker.py @@ -406,6 +406,7 @@ def test_get_imdsv2_token(self, worker: Worker, requests_put: MagicMock) -> None requests_put.assert_called_once_with( "http://169.254.169.254/latest/api/token", headers={"X-aws-ec2-metadata-token-ttl-seconds": "10"}, + timeout=Worker._IMDS_REQUEST_TIMEOUT_SECONDS, ) def test_get_imdsv2_token_cannot_connect(self, worker: Worker, requests_put: MagicMock) -> None: @@ -418,6 +419,19 @@ def test_get_imdsv2_token_cannot_connect(self, worker: Worker, requests_put: Mag # THEN assert result is None + def test_get_imdsv2_token_request_times_out( + self, worker: Worker, requests_put: MagicMock + ) -> None: + """A blackholed IMDS address must be treated as IMDS-unavailable, not hang.""" + # GIVEN + requests_put.side_effect = worker_mod.requests.Timeout("Error") + + # WHEN + result = worker._get_ec2_metadata_imdsv2_token() + + # THEN + assert result is None + def test_get_imdsv2_token_imds_inactive(self, worker: Worker, requests_put: MagicMock) -> None: # GIVEN response_mock = MagicMock() @@ -463,6 +477,7 @@ def test_spot_shutdown( requests_get.assert_called_once_with( "http://169.254.169.254/latest/meta-data/spot/instance-action", headers={"X-aws-ec2-metadata-token": fake_token}, + timeout=Worker._IMDS_REQUEST_TIMEOUT_SECONDS, ) if not is_interrupt: assert result is None @@ -532,6 +547,16 @@ def test_spot_shutdown_cannot_connect(self, worker: Worker, requests_get: MagicM # THEN assert result is None + def test_spot_shutdown_request_times_out(self, worker: Worker, requests_get: MagicMock) -> None: + # GIVEN + requests_get.side_effect = worker_mod.requests.Timeout("Error") + + # WHEN + result = worker._get_spot_instance_shutdown_action_timeout(imdsv2_token="token") + + # THEN + assert result is None + def test_spot_shutdown_imds_inactive(self, worker: Worker, requests_get: MagicMock) -> None: # GIVEN response_mock = MagicMock() @@ -575,6 +600,7 @@ def test_asg_terminate( requests_get.assert_called_once_with( "http://169.254.169.254/latest/meta-data/autoscaling/target-lifecycle-state", headers={"X-aws-ec2-metadata-token": fake_token}, + timeout=Worker._IMDS_REQUEST_TIMEOUT_SECONDS, ) def test_asg_terminate_cannot_connect(self, worker: Worker, requests_get: MagicMock) -> None: @@ -587,6 +613,16 @@ def test_asg_terminate_cannot_connect(self, worker: Worker, requests_get: MagicM # THEN assert not result + def test_asg_terminate_request_times_out(self, worker: Worker, requests_get: MagicMock) -> None: + # GIVEN + requests_get.side_effect = worker_mod.requests.Timeout("Error") + + # WHEN + result = worker._is_asg_terminated(imdsv2_token="token") + + # THEN + assert not result + def test_asg_terminate_imds_inactive(self, worker: Worker, requests_get: MagicMock) -> None: # GIVEN response_mock = MagicMock()