diff --git a/src/deadline_worker_agent/worker.py b/src/deadline_worker_agent/worker.py index cf8d9131..eae41cb7 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 4c0e724d..fa1f1b0a 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()