fix

Author: aIbrahiim

sdks/python/apache_beam/examples/inference/pytorch_sentiment.py

@@ -40,6 +40,7 @@
 from apache_beam.options.pipeline_options import SetupOptions
 from apache_beam.options.pipeline_options import StandardOptions
 from apache_beam.runners.runner import PipelineResult
+from apache_beam.runners.runner import PipelineState
 from transformers import DistilBertConfig
 from transformers import DistilBertForSequenceClassification
 from transformers import DistilBertTokenizerFast
@@ -218,6 +219,37 @@ def run_load_pipeline(known_args, pipeline_args):
   return pipeline.run()
 
 
+def _wait_until_dataflow_terminal(
+    result: PipelineResult,
+    *,
+    phase: str,
+    max_wait_secs: float = 3600.0,
+    poll_secs: float = 5.0) -> None:
+  """Poll until the Dataflow job is terminal.
+
+  ``DataflowPipelineResult.wait_until_finish`` can return while the job is still
+  in ``CANCELLING`` (internal poll uses a short countdown). Avoid
+  ``wait_until_finish(duration=None)`` after cancel; poll ``state`` so each
+  iteration refreshes job status from the API.
+  """
+  deadline = time.time() + max_wait_secs
+  while time.time() < deadline:
+    if PipelineState.is_terminal(result.state):
+      logging.info(
+          '%s: job %s finished in terminal state %s',
+          phase,
+          result.job_id(),
+          result.state)
+      return
+    time.sleep(poll_secs)
+  logging.warning(
+      '%s: job %s not terminal after %.0f s (last state %s); continuing.',
+      phase,
+      result.job_id(),
+      max_wait_secs,
+      result.state)
+
+
 def run(
     argv=None, save_main_session=True, test_pipeline=None) -> PipelineResult:
   known_args, pipeline_args = parse_known_args(argv)
@@ -286,7 +318,7 @@ def run(
   result.cancel()
   # Wait until the job is fully cancelled before tearing down Pub/Sub;
   # otherwise Dataflow can still be using the subscription (NOT_FOUND).
-  result.wait_until_finish(duration=None)
+  _wait_until_dataflow_terminal(result, phase='after cancel')
 
   cleanup_pubsub_resources(
       project=known_args.project,