Merge pull request #1043 from allenai/nanna-timeout

Improvements to how timeouts are handled.

Author: Lucaweihs

ai2thor/controller.py

@@ -19,13 +19,15 @@
 import shutil
 import subprocess
 import time
+import traceback
 import uuid
 import warnings
 from collections import defaultdict, deque
 from functools import lru_cache
 from itertools import product
 from platform import architecture as platform_architecture
 from platform import system as platform_system
+from typing import Dict, Any, Union, Optional
 
 import numpy as np
 
@@ -392,6 +394,8 @@ def __init__(
         server_class=None,
         gpu_device=None,
         platform=None,
+        server_timeout: Optional[float] = 100.0,
+        server_start_timeout: float = 300.0,
         **unity_initialization_parameters,
     ):
         self.receptacle_nearest_pivot_points = {}
@@ -401,6 +405,11 @@ def __init__(
         self.width = width
         self.height = height
 
+        self.server_timeout = server_timeout
+        self.server_start_timeout = server_start_timeout
+        assert self.server_timeout is None or 0 < self.server_timeout
+        assert 0 < self.server_start_timeout
+
         self.last_event = None
         self.scene = None
         self._scenes_in_build = None
@@ -577,7 +586,7 @@ def __init__(
             init_return = event.metadata["actionReturn"]
             if init_return:
                 self.server.set_init_params(init_return)
-                logging.info("Initialize return: {}".format(init_return))
+                logging.info(f"Initialize return: {init_return}")
 
     def _build_server(self, host, port, width, height):
 
@@ -593,20 +602,22 @@ def _build_server(self, host, port, width, height):
 
         if self.server_class == ai2thor.wsgi_server.WsgiServer:
             self.server = ai2thor.wsgi_server.WsgiServer(
-                host,
+                host=host,
+                timeout=self.server_timeout,
                 port=port,
-                depth_format=self.depth_format,
-                add_depth_noise=self.add_depth_noise,
                 width=width,
                 height=height,
+                depth_format=self.depth_format,
+                add_depth_noise=self.add_depth_noise,
             )
 
         elif self.server_class == ai2thor.fifo_server.FifoServer:
             self.server = ai2thor.fifo_server.FifoServer(
-                depth_format=self.depth_format,
-                add_depth_noise=self.add_depth_noise,
                 width=width,
                 height=height,
+                timeout=self.server_timeout,
+                depth_format=self.depth_format,
+                add_depth_noise=self.add_depth_noise,
             )
 
     def __enter__(self):
@@ -913,9 +924,9 @@ def multi_step_physics(self, action, timeStep=0.05, max_steps=20):
 
         return events
 
-    def step(self, action=None, **action_args):
+    def step(self, action: Union[str, Dict[str, Any]]=None, **action_args):
 
-        if type(action) is dict:
+        if isinstance(action, Dict):
             action = copy.deepcopy(action)  # prevent changes from leaking
         else:
             action = dict(action=action)
@@ -960,18 +971,27 @@ def step(self, action=None, **action_args):
         self.server.send(action)
         try:
             self.last_event = self.server.receive()
-        except UnityCrashException as e:
+        except UnityCrashException:
             self.server.stop()
             self.server = None
             # we don't need to pass port or host, since this Exception
             # is only thrown from the FifoServer, start_unity is also
             # not passed since Unity would have to have been started
             # for this to be thrown
-            message = "Restarting unity due to crash: %s" % e
+            message = (
+                f"Restarting unity due to crash when when running action {action}"
+                f" in scene {self.last_event.metadata['sceneName']}:\n{traceback.format_exc()}"
+            )
             warnings.warn(message)
             self.start(width=self.width, height=self.height, x_display=self.x_display)
             self.reset()
             raise RestartError(message)
+        except Exception as e:
+            self.server.stop()
+            raise (TimeoutError if isinstance(e, TimeoutError) else RuntimeError)(
+                f"Error encountered when running action {action}"
+                f" in scene {self.last_event.metadata['sceneName']}."
+            )
 
         if not self.last_event.metadata["lastActionSuccess"]:
             if self.last_event.metadata["errorCode"] in [
@@ -1056,7 +1076,7 @@ def _start_unity_thread(self, env, width, height, server_params, image_name):
             pass
 
         self.unity_pid = proc.pid
-        atexit.register(lambda: proc.poll() is None and proc.kill())
+        atexit.register(lambda: self.server.stop())
 
     @property
     def tmp_dir(self):
@@ -1336,7 +1356,7 @@ def start(
             self._start_unity_thread(env, width, height, unity_params, image_name)
 
         # receive the first request
-        self.last_event = self.server.receive()
+        self.last_event = self.server.receive(timeout=self.server_start_timeout)
 
         # we should be able to get rid of this since we check the resolution in .reset()
         if self.server.unity_proc is not None and (height < 300 or width < 300):

ai2thor/fifo_server.py

@@ -5,15 +5,22 @@
 Handles all communication with Unity through a Flask service.  Messages
 are sent to the controller using a pair of request/response queues.
 """
-import ai2thor.server
 import json
-import msgpack
 import os
+import select
+import struct
 import tempfile
-from ai2thor.exceptions import UnityCrashException
-from enum import IntEnum, unique
+import time
 from collections import defaultdict
-import struct
+from enum import IntEnum, unique
+from io import TextIOWrapper
+from typing import Optional
+
+import msgpack
+
+import ai2thor.server
+from ai2thor.exceptions import UnityCrashException
+
 
 # FifoFields
 @unique
@@ -45,17 +52,18 @@ class FifoServer(ai2thor.server.Server):
 
     def __init__(
         self,
-        width,
-        height,
+        width: int,
+        height: int,
+        timeout: Optional[float] = 100.0,
         depth_format=ai2thor.server.DepthFormat.Meters,
-        add_depth_noise=False,
+        add_depth_noise: bool = False,
     ):
 
         self.tmp_dir = tempfile.TemporaryDirectory()
         self.server_pipe_path = os.path.join(self.tmp_dir.name, "server.pipe")
         self.client_pipe_path = os.path.join(self.tmp_dir.name, "client.pipe")
-        self.server_pipe = None
-        self.client_pipe = None
+        self.server_pipe: Optional[TextIOWrapper] = None
+        self.client_pipe: Optional[TextIOWrapper] = None
         self.raw_metadata = None
         self.raw_files = None
         self._last_action_message = None
@@ -93,19 +101,52 @@ def __init__(
         }
 
         self.eom_header = self._create_header(FieldType.END_OF_MESSAGE, b"")
-        super().__init__(width, height, depth_format, add_depth_noise)
+        super().__init__(
+            width=width,
+            height=height,
+            timeout=timeout,
+            depth_format=depth_format,
+            add_depth_noise=add_depth_noise
+        )
 
     def _create_header(self, message_type, body):
         return struct.pack(self.header_format, message_type, len(body))
 
-    def _recv_message(self):
+    def _read_with_timeout(self, server_pipe, message_size: int, timeout: Optional[float]):
+        if timeout is None:
+            return server_pipe.read(message_size)
+
+        start_t = time.time()
+        message = b""
+
+        while message_size > 0:
+            r, w, e = select.select([server_pipe], [], [], timeout)
+            if server_pipe in r:
+                part = os.read(server_pipe.fileno(), message_size)
+                message_size -= len(part)
+                message = message + part
+
+            cur_t = time.time()
+            if timeout is not None and cur_t - start_t > timeout:
+                break
+
+        if message_size != 0:
+            raise TimeoutError(f"Reading from AI2-THOR backend timed out (using {timeout}s) timeout.")
+
+        return message
+
+    def _recv_message(self, timeout: Optional[float]):
         if self.server_pipe is None:
             self.server_pipe = open(self.server_pipe_path, "rb")
 
         metadata = None
         files = defaultdict(list)
         while True:
-            header = self.server_pipe.read(self.header_size)  # message type + length
+            header = self._read_with_timeout(
+                server_pipe=self.server_pipe,
+                message_size=self.header_size,
+                timeout=self.timeout if timeout is None else timeout
+            )  # message type + length
             if len(header) == 0:
                 self.unity_proc.wait(timeout=5)
                 returncode = self.unity_proc.returncode
@@ -131,7 +172,13 @@ def _recv_message(self):
             # print("got header %s" % header)
             field_type_int, message_length = struct.unpack(self.header_format, header)
             field_type = self.field_types[field_type_int]
-            body = self.server_pipe.read(message_length)
+
+            body = self._read_with_timeout(
+                server_pipe=self.server_pipe,
+                message_size=message_length,
+                timeout=self.timeout if timeout is None else timeout
+            )
+
             # print("field type")
             # print(field_type)
             if field_type is FieldType.METADATA:
@@ -177,9 +224,11 @@ def _send_message(self, message_type, body):
         self.client_pipe.write(header + body + self.eom_header)
         self.client_pipe.flush()
 
-    def receive(self):
+    def receive(self, timeout: Optional[float] = None):
 
-        metadata, files = self._recv_message()
+        metadata, files = self._recv_message(
+            timeout=self.timeout if timeout is None else timeout
+        )
 
         if metadata is None:
             raise ValueError("no metadata received from recv_message")
@@ -215,5 +264,11 @@ def unity_params(self):
         return params
 
     def stop(self):
-        self.client_pipe.close()
-        self.server_pipe.close()
+        if self.client_pipe is not None:
+            self.client_pipe.close()
+
+        if self.server_pipe is not None:
+            self.server_pipe.close()
+
+        if self.unity_proc is not None and self.unity_proc.poll() is None:
+            self.unity_proc.kill()