Remove oom protection

Author: AMOOOMA

sdks/python/apache_beam/utils/multi_process_shared.py

@@ -117,25 +117,6 @@ def __dir__(self):
     return dir
 
 
-def _run_with_oom_protection(func, *args, **kwargs):
-  try:
-    return func(*args, **kwargs)
-  except Exception as e:
-    # Check string to avoid hard import dependency
-    if 'CUDA out of memory' in str(e):
-      logging.warning("Caught CUDA OOM during operation. Cleaning memory.")
-      try:
-        import gc
-        import torch
-        gc.collect()
-        torch.cuda.empty_cache()
-      except ImportError:
-        pass
-      except Exception as cleanup_error:
-        logging.error("Failed to clean up CUDA memory: %s", cleanup_error)
-    raise e
-
-
 class _SingletonEntry:
   """Represents a single, refcounted entry in this process."""
   def __init__(
@@ -145,15 +126,15 @@ def __init__(
     self.refcount = 0
     self.lock = threading.Lock()
     if initialize_eagerly:
-      self.obj = _run_with_oom_protection(constructor)
+      self.obj = constructor()
       self.initialied = True
     else:
       self.initialied = False
 
   def acquire(self):
     with self.lock:
       if not self.initialied:
-        self.obj = _run_with_oom_protection(self.constructor)
+        self.obj = self.constructor()
         self.initialied = True
       self.refcount += 1
       return _SingletonProxy(self)