Merge pull request #268 from allenai/auto_x_display

Improving x-display management for ithor/robothor tasks.

allenact_plugins/ithor_plugin/extra_environment.yml

@@ -5,6 +5,9 @@ dependencies:
   - ai2thor>=2.5.3
   - numba
   - pip
+  - colour
+  - packaging
   - pip:
       - numpy-quaternion
       - pyquaternion>=0.9.9
+      - python-xlib

allenact_plugins/ithor_plugin/extra_requirements.txt

@@ -1,4 +1,7 @@
 ai2thor>=2.5.3
 numpy-quaternion
 pyquaternion>=0.9.9
+colour
 numba
+packaging
+python-xlib

allenact_plugins/ithor_plugin/ithor_util.py

@@ -1,4 +1,11 @@
+import glob
 import math
+import os
+import platform
+from typing import Sequence
+
+import Xlib
+import Xlib.display
 
 
 def vertical_to_horizontal_fov(
@@ -33,3 +40,35 @@ def round_to_factor(num: float, base: int) -> int:
     base: integer base
     """
     return round(num / base) * base
+
+
+def get_open_x_displays(throw_error_if_empty: bool = False) -> Sequence[str]:
+    assert platform.system() == "Linux", "Can only get X-displays for Linux systems."
+
+    displays = []
+
+    open_display_strs = [
+        os.path.basename(s)[1:] for s in glob.glob("/tmp/.X11-unix/X*")
+    ]
+
+    for open_display_str in sorted(open_display_strs):
+        try:
+            open_display_str = str(int(open_display_str))
+        except Exception:
+            continue
+
+        display = Xlib.display.Display(":{}".format(open_display_str))
+
+        displays.extend(
+            [f"{open_display_str}.{i}" for i in range(display.screen_count())]
+        )
+
+    if throw_error_if_empty and len(displays) == 0:
+        raise IOError(
+            "Could not find any open X-displays on which to run AI2-THOR processes. "
+            " Please see the AI2-THOR installation instructions at"
+            " https://allenact.org/installation/installation-framework/#installation-of-ithor-ithor-plugin"
+            " for information as to how to start such displays."
+        )
+
+    return displays

allenact_plugins/robothor_plugin/extra_environment.yml

@@ -10,3 +10,4 @@ dependencies:
   - pip:
       - numpy-quaternion
       - pyquaternion>=0.9.9
+      - python-xlib

allenact_plugins/robothor_plugin/extra_requirements.txt

@@ -4,3 +4,4 @@ pyquaternion>=0.9.9
 colour
 numba
 packaging
+python-xlib

projects/objectnav_baselines/experiments/objectnav_thor_base.py

@@ -1,5 +1,6 @@
 import glob
 import os
+import platform
 from abc import ABC
 from math import ceil
 from typing import Dict, Any, List, Optional, Sequence
@@ -14,13 +15,26 @@
 from allenact.base_abstractions.task import TaskSampler
 from allenact.utils.experiment_utils import evenly_distribute_count_into_bins
 from allenact.utils.system import get_logger
-from allenact_plugins.ithor_plugin.ithor_util import horizontal_to_vertical_fov
+from allenact_plugins.ithor_plugin.ithor_util import (
+    horizontal_to_vertical_fov,
+    get_open_x_displays,
+)
 from allenact_plugins.robothor_plugin.robothor_sensors import DepthSensorThor
 from allenact_plugins.robothor_plugin.robothor_task_samplers import (
     ObjectNavDatasetTaskSampler,
 )
 from allenact_plugins.robothor_plugin.robothor_tasks import ObjectNavTask
 from projects.objectnav_baselines.experiments.objectnav_base import ObjectNavBaseConfig
+import ai2thor
+from packaging import version
+
+if ai2thor.__version__ not in ["0.0.1", None] and version.parse(
+    ai2thor.__version__
+) < version.parse("2.7.2"):
+    raise ImportError(
+        "To run the AI2-THOR ObjectNav baseline experiments you must use"
+        " ai2thor version 2.7.1 or higher."
+    )
 
 
 class ObjectNavThorBaseConfig(ObjectNavBaseConfig, ABC):
@@ -168,6 +182,20 @@ def _get_sampler_args_for_scene_split(
 
         inds = self._partition_inds(len(scenes), total_processes)
 
+        x_display: Optional[str] = None
+        if platform.system() == "Linux":
+            x_displays = get_open_x_displays(throw_error_if_empty=True)
+
+            if len(devices) > len(x_displays):
+                get_logger().warning(
+                    f"More GPU devices found than X-displays (devices: `{x_displays}`, x_displays: `{x_displays}`)."
+                    f" This is not necessarily a bad thing but may mean that you're not using GPU memory as"
+                    f" efficiently as possible. Consider following the instructions here:"
+                    f" https://allenact.org/installation/installation-framework/#installation-of-ithor-ithor-plugin"
+                    f" describing how to start an X-display on every GPU."
+                )
+            x_display = x_displays[process_ind % len(x_displays)]
+
         return {
             "scenes": scenes[inds[process_ind] : inds[process_ind + 1]],
             "object_types": self.TARGET_TYPES,
@@ -183,16 +211,7 @@ def _get_sampler_args_for_scene_split(
             "seed": seeds[process_ind] if seeds is not None else None,
             "deterministic_cudnn": deterministic_cudnn,
             "rewards_config": self.REWARD_CONFIG,
-            "env_args": {
-                **self.env_args(),
-                "x_display": (
-                    f"0.{devices[process_ind % len(devices)]}"
-                    if devices is not None
-                    and len(devices) > 0
-                    and devices[process_ind % len(devices)] >= 0
-                    else None
-                ),
-            },
+            "env_args": {**self.env_args(), "x_display": x_display,},
         }
 
     def train_task_sampler_args(

projects/objectnav_baselines/experiments/robothor/objectnav_robothor_base.py

@@ -5,17 +5,6 @@
     ObjectNavThorBaseConfig,
 )
 
-import ai2thor
-from packaging import version
-
-if ai2thor.__version__ not in ["0.0.1", None] and version.parse(
-    ai2thor.__version__
-) < version.parse("2.7.2"):
-    raise ImportError(
-        "To run the ObjectNavRoboThor baseline experiments you must use"
-        " ai2thor version 2.7.1 or higher."
-    )
-
 
 class ObjectNavRoboThorBaseConfig(ObjectNavThorBaseConfig, ABC):
     """The base config for all RoboTHOR ObjectNav experiments."""

projects/pointnav_baselines/experiments/pointnav_thor_base.py

@@ -1,26 +1,38 @@
 import glob
 import os
+import platform
 from abc import ABC
 from math import ceil
 from typing import Dict, Any, List, Optional, Sequence
 
+import ai2thor
 import gym
 import numpy as np
 import torch
+from packaging import version
 
 from allenact.base_abstractions.experiment_config import MachineParams
 from allenact.base_abstractions.preprocessor import SensorPreprocessorGraph
 from allenact.base_abstractions.sensor import SensorSuite, ExpertActionSensor
 from allenact.base_abstractions.task import TaskSampler
 from allenact.utils.experiment_utils import evenly_distribute_count_into_bins
 from allenact.utils.system import get_logger
+from allenact_plugins.ithor_plugin.ithor_util import get_open_x_displays
 from allenact_plugins.robothor_plugin.robothor_sensors import DepthSensorThor
 from allenact_plugins.robothor_plugin.robothor_task_samplers import (
     PointNavDatasetTaskSampler,
 )
 from allenact_plugins.robothor_plugin.robothor_tasks import ObjectNavTask
 from projects.pointnav_baselines.experiments.pointnav_base import PointNavBaseConfig
 
+if ai2thor.__version__ not in ["0.0.1", None] and version.parse(
+    ai2thor.__version__
+) < version.parse("2.7.2"):
+    raise ImportError(
+        "To run the PointNav baseline experiments you must use"
+        " ai2thor version 2.7.1 or higher."
+    )
+
 
 class PointNavThorBaseConfig(PointNavBaseConfig, ABC):
     """The base config for all iTHOR PointNav experiments."""
@@ -143,6 +155,20 @@ def _get_sampler_args_for_scene_split(
 
         inds = self._partition_inds(len(scenes), total_processes)
 
+        x_display: Optional[str] = None
+        if platform.system() == "Linux":
+            x_displays = get_open_x_displays(throw_error_if_empty=True)
+
+            if len(devices) > len(x_displays):
+                get_logger().warning(
+                    f"More GPU devices found than X-displays (devices: `{x_displays}`, x_displays: `{x_displays}`)."
+                    f" This is not necessarily a bad thing but may mean that you're not using GPU memory as"
+                    f" efficiently as possible. Consider following the instructions here:"
+                    f" https://allenact.org/installation/installation-framework/#installation-of-ithor-ithor-plugin"
+                    f" describing how to start an X-display on every GPU."
+                )
+            x_display = x_displays[process_ind % len(x_displays)]
+
         return {
             "scenes": scenes[inds[process_ind] : inds[process_ind + 1]],
             "object_types": self.TARGET_TYPES,
@@ -158,16 +184,7 @@ def _get_sampler_args_for_scene_split(
             "seed": seeds[process_ind] if seeds is not None else None,
             "deterministic_cudnn": deterministic_cudnn,
             "rewards_config": self.REWARD_CONFIG,
-            "env_args": {
-                **self.ENV_ARGS,
-                "x_display": (
-                    f"0.{devices[process_ind % len(devices)]}"
-                    if devices is not None
-                    and len(devices) > 0
-                    and devices[process_ind % len(devices)] >= 0
-                    else None
-                ),
-            },
+            "env_args": {**self.ENV_ARGS, "x_display": x_display,},
         }
 
     def train_task_sampler_args(