Reduce severity for state requests cancelled by runner

model/fn-execution/src/main/proto/org/apache/beam/model/fn_execution/v1/beam_fn_api.proto

@@ -930,6 +930,19 @@ message StateResponse {
   // failed.
   string error = 2;
 
+  enum ErrorReason {
+    ERROR_REASON_UNSPECIFIED = 0;
+
+    // Indicates that the state request failed because the runner cancelled or
+    // invalidated the work item (e.g., due to autoscaling, rebalancing, or
+    // hedging). SDKs may choose to reduce logging severity for this error.
+    CANCELLED = 1;
+  }
+
+  // (Optional) If error is specified, provides a structured reason for why the
+  // request failed.
+  ErrorReason error_reason = 3;
+
   // A corresponding response matching the request will be populated.
   oneof response {
     // A response to getting state.

sdks/python/apache_beam/runners/common.py

@@ -1618,11 +1618,18 @@ def _reraise_augmented(self, exn, windowed_value=None):
     _, _, tb = exc_info
 
     new_exn = new_exn.with_traceback(tb)
+    if isinstance(exn, WorkCancelledException):
+      raise new_exn
     self._maybe_sample_exception(exc_info, windowed_value)
     _LOGGER.exception(new_exn)
     raise new_exn
 
 
+class WorkCancelledException(RuntimeError):
+  """Indicates that a state request or work item was cancelled by the runner."""
+  pass
+
+
 class OutputHandler(object):
   def handle_process_outputs(
       self, windowed_input_element, results, watermark_estimator=None):

sdks/python/apache_beam/runners/worker/sdk_worker.py

@@ -52,6 +52,7 @@
 from apache_beam.portability.api import beam_fn_api_pb2
 from apache_beam.portability.api import beam_fn_api_pb2_grpc
 from apache_beam.portability.api import metrics_pb2
+from apache_beam.runners.common import WorkCancelledException
 from apache_beam.runners.worker import bundle_processor
 from apache_beam.runners.worker import data_plane
 from apache_beam.runners.worker import statesampler
@@ -308,6 +309,14 @@ def _execute(
     with statesampler.instruction_id(request.instruction_id):
       try:
         response = task()
+      except WorkCancelledException:
+        traceback_string = traceback.format_exc()
+        _LOGGER.info(
+            'Instruction %s cancelled by runner. Original traceback is\n%s\n',
+            request.instruction_id,
+            traceback_string)
+        response = beam_fn_api_pb2.InstructionResponse(
+            instruction_id=request.instruction_id, error=traceback_string)
       except:  # pylint: disable=bare-except
         traceback_string = traceback.format_exc()
         print(traceback_string, file=sys.stderr)
@@ -1153,6 +1162,9 @@ def _blocking_request(self, request):
         raise RuntimeError()
     response = req_future.get()
     if response.error:
+      if (response.error_reason ==
+          beam_fn_api_pb2.StateResponse.ErrorReason.CANCELLED):
+        raise WorkCancelledException(response.error)
       raise RuntimeError(response.error)
     else:
       return response

sdks/python/apache_beam/runners/worker/sdk_worker_test.py

@@ -432,6 +432,34 @@ def test_process_bundle_passes_data_stream_id(self):
     mock_bundle_processor.process_bundle.assert_called_once_with(
         instruction_id, 'stream_xyz')
 
+  def test_cancelled_state_request_logs_info(self):
+    state_handler = sdk_worker.GrpcStateHandler(mock.MagicMock())
+    state_handler._context.process_instruction_id = 'bundle_1'
+    cancelled_response = beam_fn_api_pb2.StateResponse(
+        id='1',
+        error='Work item cancelled by runner',
+        error_reason=beam_fn_api_pb2.StateResponse.ErrorReason.CANCELLED)
+    state_handler._request = mock.MagicMock(
+        return_value=sdk_worker._Future().set(cancelled_response))
+
+    with self.assertRaises(sdk_worker.WorkCancelledException):
+      state_handler._blocking_request(beam_fn_api_pb2.StateRequest())
+
+    with mock.patch('grpc.channel_ready_future'):
+      harness = sdk_worker.SdkHarness('localhost:0')
+    request = beam_fn_api_pb2.InstructionRequest(instruction_id='bundle_1')
+    with mock.patch.object(sdk_worker._LOGGER, 'info') as mock_info, \
+         mock.patch.object(sdk_worker._LOGGER, 'error') as mock_error:
+      harness._execute(
+          lambda: state_handler._blocking_request(
+              beam_fn_api_pb2.StateRequest()),
+          request)
+      mock_info.assert_called_once()
+      mock_error.assert_not_called()
+    response = harness._responses.get_nowait()
+    self.assertEqual(response.instruction_id, 'bundle_1')
+    self.assertIn('Work item cancelled by runner', response.error)
+
 
 class CachingStateHandlerTest(unittest.TestCase):
   def test_caching(self):