[DataLoader] add repr for WorkerInfo (#39975)
Summary: Pull Request resolved: https://github.com/pytorch/pytorch/pull/39975
Differential Revision: D22039414
Pulled By: ezyang
fbshipit-source-id: 230f68a91fca901bce652fdf88ba88167f39b978
diff --git a/test/test_dataloader.py b/test/test_dataloader.py
index a5fa18d..6b60208 100644
--- a/test/test_dataloader.py
+++ b/test/test_dataloader.py
@@ -320,8 +320,6 @@
set_faulthander_if_available()
# Process `pid` must have called `set_faulthander_if_available`
-
-
def print_traces_of_all_threads(pid):
if HAS_FAULTHANDLER:
if not IS_WINDOWS:
@@ -666,6 +664,8 @@
worker_info.a = 3
except RuntimeError as e:
assert str(e) == "Cannot assign attributes to WorkerInfo objects"
+ for k in ['id', 'num_workers', 'seed', 'dataset']:
+ assert "{}=".format(k) in repr(worker_info)
dataset.value = [worker_id, os.getpid()]
diff --git a/torch/utils/data/_utils/worker.py b/torch/utils/data/_utils/worker.py
index 49d4417..b538129 100644
--- a/torch/utils/data/_utils/worker.py
+++ b/torch/utils/data/_utils/worker.py
@@ -63,6 +63,7 @@
def __init__(self, **kwargs):
for k, v in kwargs.items():
setattr(self, k, v)
+ self.__keys = tuple(kwargs.keys())
self.__initialized = True
def __setattr__(self, key, val):
@@ -70,6 +71,12 @@
raise RuntimeError("Cannot assign attributes to {} objects".format(self.__class__.__name__))
return super(WorkerInfo, self).__setattr__(key, val)
+ def __repr__(self):
+ items = []
+ for k in self.__keys:
+ items.append('{}={}'.format(k, getattr(self, k)))
+ return '{}({})'.format(self.__class__.__name__, ', '.join(items))
+
def get_worker_info():
r"""Returns the information about the current
