Unescape string in RPC error message (#49373)

rohan-varma · facebook-github-bot · commit 40d7c1091fff · 2020-12-16T01:40:31.000-08:00
Summary: Pull Request resolved: #49373 Unescaping the string in RPC error message to provide better error msg Test Plan: CI Reviewed By: xush6528 Differential Revision: D25511730 fbshipit-source-id: 054f46d5ffbcb1350012362a023fafb1fe57fca1
diff --git a/torch/distributed/rpc/internal.py b/torch/distributed/rpc/internal.py
@@ -201,7 +201,7 @@ def _run_function(python_udf):
 
 def _handle_exception(result):
     if isinstance(result, RemoteException):
-        raise result.exception_type(result.msg)
+        raise result.exception_type(result.msg.encode("utf-8").decode("unicode_escape"))
 
 
 def _build_rpc_profiling_key(
diff --git a/torch/testing/_internal/distributed/rpc/rpc_test.py b/torch/testing/_internal/distributed/rpc/rpc_test.py
@@ -317,6 +317,10 @@ def my_script_func(tensor):
 def raise_func():
     raise ValueError(expected_err)
 
+expected_err_escape = "\nFirst line of error \n next line of error \n last line of error"
+def raise_func_escape():
+    raise ValueError(expected_err_escape)
+
 
 global_rref = None
 
@@ -1982,6 +1986,20 @@ def test_py_raise_in_user_func(self):
         stderr_lines = err.getvalue()
         self.assertTrue(expected_err in stderr_lines)
 
+    @dist_init
+    def test_py_raise_in_user_func_escaped_str(self):
+        n = self.rank + 1
+        dst_rank = n % self.world_size
+        fut = rpc.rpc_async(worker_name(dst_rank), raise_func_escape)
+        try:
+            fut.wait()
+        except ValueError as e:
+            msg = str(e)
+            # Ensure newlines are unescaped to provide a better repr of error.
+            self.assertEqual(msg, msg.encode("utf-8").decode("unicode_escape"))
+        else:
+            self.assertTrue(False, "expected raise_func_escape to raise ValueError.")
+
     @dist_init
     def test_nested_rpc(self):
         n = self.rank + 1
