[Core] Some simplification of WorkerWrapper changes (#4183)

This commit is contained in:
Nick Hill
2024-04-23 00:49:08 -07:00
committed by GitHub
parent 0ae11f78ab
commit 8f2ea22bde
2 changed files with 45 additions and 54 deletions

View File

@@ -108,7 +108,8 @@ class WorkerWrapperBase:
self.worker_class_name = worker_class_name
self.worker = None
def update_environment_variables(self, envs: Dict[str, str]) -> None:
@staticmethod
def update_environment_variables(envs: Dict[str, str]) -> None:
key = 'CUDA_VISIBLE_DEVICES'
if key in envs and key in os.environ:
# overwriting CUDA_VISIBLE_DEVICES is desired behavior
@@ -138,10 +139,8 @@ class WorkerWrapperBase:
def execute_method(self, method, *args, **kwargs):
try:
if hasattr(self, method):
executor = getattr(self, method)
else:
executor = getattr(self.worker, method)
target = self if self.worker is None else self.worker
executor = getattr(target, method)
return executor(*args, **kwargs)
except Exception as e:
# if the driver worker also execute methods,