android-env 1.2.2__py3-none-any.whl → 1.2.3__py3-none-any.whl

android_env/components/simulators/emulator/emulator_simulator.py

@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2023 DeepMind Technologies Limited.
+# Copyright 2024 DeepMind Technologies Limited.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.
@@ -17,11 +17,12 @@
 
 import os
 import time
-from typing import Any, Optional
+from typing import Any
 
 from absl import logging
 from android_env.components import adb_controller
 from android_env.components import adb_log_stream
+from android_env.components import config_classes
 from android_env.components import errors
 from android_env.components import log_stream
 from android_env.components.simulators import base_simulator
@@ -38,15 +39,19 @@ from android_env.proto import snapshot_service_pb2_grpc
 from google.protobuf import empty_pb2
 
 
-_DEFAULT_SNAPSHOT_NAME = 'default_snapshot'
+_DEFAULT_SNAPSHOT_NAME: str = 'default_snapshot'
 
 
-def is_existing_emulator_provided(launcher_args: dict[str, Any]) -> bool:
+def _is_existing_emulator_provided(
+    launcher_config: config_classes.EmulatorLauncherConfig,
+) -> bool:
   """Returns true if all necessary args were provided."""
+
   return bool(
-      launcher_args.get('adb_port') and
-      launcher_args.get('emulator_console_port') and
-      launcher_args.get('grpc_port'))
+      launcher_config.adb_port
+      and launcher_config.emulator_console_port
+      and launcher_config.grpc_port
+  )
 
 
 def _pick_adb_port() -> int:
@@ -58,6 +63,7 @@ def _pick_adb_port() -> int:
   Returns:
     port: an available port for adb.
   """
+
   for p in range(5555, 5587, 2):
     if portpicker.is_port_free(p):
       return p
@@ -73,29 +79,13 @@ def _pick_emulator_grpc_port() -> int:
   Returns:
     port: an available port for emulator grpc.
   """
+
   if portpicker.is_port_free(8554):
     return 8554
   else:
     return portpicker.pick_unused_port()
 
 
-def _reconnect_on_grpc_error(func):
-  """Decorator function for reconnecting to emulator upon grpc errors."""
-
-  def wrapper(*args, **kwargs):
-    try:
-      return func(*args, **kwargs)  # pytype: disable=missing-parameter  # always-use-return-annotations
-    except grpc.RpcError:
-      logging.exception('RpcError caught. Reconnecting to emulator...')
-      emu = args[0]  # The first arg of the function is "self"
-      emu._emulator_stub, emu._snapshot_stub = emu._connect_to_emulator(  # pylint: disable=protected-access
-          emu._grpc_port  # pylint: disable=protected-access
-      )
-      return func(*args, **kwargs)  # pytype: disable=missing-parameter  # always-use-return-annotations
-
-  return wrapper
-
-
 class EmulatorBootError(errors.SimulatorError):
   """Raised when an emulator failed to boot."""
 
@@ -107,90 +97,86 @@ class EmulatorCrashError(errors.SimulatorError):
 class EmulatorSimulator(base_simulator.BaseSimulator):
   """Controls an Android Emulator."""
 
-  def __init__(
-      self,
-      emulator_launcher_args: dict[str, Any],
-      adb_controller_args: dict[str, Any],
-      tmp_dir: str = '/tmp/android_env/simulator',
-      logfile_path: Optional[str] = None,
-      launch_n_times_without_reboot: int = 1,
-      launch_n_times_without_reinstall: int = 2,
-      **kwargs,
-  ):
-    """Instantiates an EmulatorSimulator.
+  def __init__(self, config: config_classes.EmulatorConfig):
+    """Instantiates an EmulatorSimulator."""
 
-    Args:
-      emulator_launcher_args: Arguments for EmulatorLauncher.
-      adb_controller_args: Arguments for AdbController.
-      tmp_dir: Temporary directory to hold simulator files.
-      logfile_path: Path to file which holds emulator logs. If not provided, it
-        will be determined by the EmulatorLauncher.
-      launch_n_times_without_reboot: The number of times to try launching the
-        emulator before rebooting (reboot on the n+1-st try).
-      launch_n_times_without_reinstall: The number of times to try launching the
-        emulator before reinstalling (reinstall on the n+1-st try).
-      **kwargs: keyword arguments for base class.
-    """
+    super().__init__(config)
+    self._config = config
 
     # If adb_port, console_port and grpc_port are all already provided,
     # we assume the emulator already exists and there's no need to launch.
-    if is_existing_emulator_provided(emulator_launcher_args):
+    if _is_existing_emulator_provided(self._config.emulator_launcher):
       self._existing_emulator_provided = True
-      self._adb_port = emulator_launcher_args['adb_port']
-      self._console_port = emulator_launcher_args['emulator_console_port']
-      self._grpc_port = emulator_launcher_args['grpc_port']
       logging.info('Connecting to existing emulator "%r"',
                    self.adb_device_name())
     else:
       self._existing_emulator_provided = False
-      self._adb_port = _pick_adb_port()
-      self._console_port = portpicker.pick_unused_port()
-      self._grpc_port = _pick_emulator_grpc_port()
+      self._config.emulator_launcher.adb_port = _pick_adb_port()
+      self._config.emulator_launcher.emulator_console_port = (
+          portpicker.pick_unused_port()
+      )
+      self._config.emulator_launcher.grpc_port = _pick_emulator_grpc_port()
 
     self._channel = None
-    self._emulator_stub = None
+    self._emulator_stub: emulator_controller_pb2_grpc.EmulatorControllerStub | None = (
+        None
+    )
     self._snapshot_stub = None
     # Set the image format to RGBA. The width and height of the returned
     # screenshots will use the device's width and height.
     self._image_format = emulator_controller_pb2.ImageFormat(
         format=emulator_controller_pb2.ImageFormat.ImgFormat.RGBA8888)
 
-    if launch_n_times_without_reboot > launch_n_times_without_reinstall:
+    if (
+        self._config.launch_n_times_without_reboot
+        > self._config.launch_n_times_without_reinstall
+    ):
       raise ValueError(
-          f'Number of launch attempts before reboot '
-          f'({launch_n_times_without_reboot}) should not be greater than '
-          f'number of launch attempts before reinstall '
-          f'({launch_n_times_without_reinstall})')
-    self._launch_n_times_without_reboot = launch_n_times_without_reboot
-    self._launch_n_times_without_reinstall = launch_n_times_without_reinstall
-
-    super().__init__(**kwargs)
+          'Number of launch attempts before reboot'
+          f' ({self._config.launch_n_times_without_reboot}) should not be'
+          ' greater than number of launch attempts before reinstall'
+          f' ({self._config.launch_n_times_without_reinstall})'
+      )
 
     # Initialize own ADB controller.
-    self._adb_controller_args = adb_controller_args
+    self._config.adb_controller.device_name = self.adb_device_name()
     self._adb_controller = self.create_adb_controller()
     self._adb_controller.init_server()
-    logging.info('Initialized simulator with ADB server port %r.',
-                 self._adb_controller_args['adb_server_port'])
+    logging.info(
+        'Initialized simulator with ADB server port %r.',
+        self._config.adb_controller.adb_server_port,
+    )
 
     # If necessary, create EmulatorLauncher.
     if self._existing_emulator_provided:
-      self._logfile_path = logfile_path or None
+      self._logfile_path = self._config.logfile_path or None
       self._launcher = None
     else:
-      emulator_launcher_args.update({
-          'adb_path': self._adb_controller_args['adb_path'],
-          'adb_port': self._adb_port,
-          'adb_server_port': self._adb_controller_args['adb_server_port'],
-          'emulator_console_port': self._console_port,
-          'grpc_port': self._grpc_port,
-          'tmp_dir': tmp_dir,
-      })
-      self._emulator_launcher_args = emulator_launcher_args
-      logging.info('emulator_launcher_args: %r', self._emulator_launcher_args)
+      logging.info(
+          'emulator_launcher config: %r', self._config.emulator_launcher
+      )
       self._launcher = emulator_launcher.EmulatorLauncher(
-          **self._emulator_launcher_args)
-      self._logfile_path = logfile_path or self._launcher.logfile_path()
+          config=self._config.emulator_launcher,
+          adb_controller_config=self._config.adb_controller,
+      )
+      self._logfile_path = (
+          self._config.logfile_path or self._launcher.logfile_path()
+      )
+
+  def _reconnect_on_grpc_error(func):
+    """Decorator function for reconnecting to emulator upon grpc errors."""
+
+    def wrapper(self, *args, **kwargs):
+      try:
+        return func(self, *args, **kwargs)
+      except grpc.RpcError:
+        logging.exception('RpcError caught. Reconnecting to emulator...')
+        self._emulator_stub, self._snapshot_stub = self._connect_to_emulator(
+            self._config.emulator_launcher.grpc_port
+        )
+        return func(self, *args, **kwargs)
+
+    return wrapper
 
   def get_logs(self) -> str:
     """Returns logs recorded by the emulator."""
@@ -201,28 +187,29 @@ class EmulatorSimulator(base_simulator.BaseSimulator):
       return f'Logfile does not exist: {self._logfile_path}.'
 
   def adb_device_name(self) -> str:
-    return 'emulator-%s' % (self._adb_port - 1)
+    return 'emulator-%s' % (self._config.emulator_launcher.adb_port - 1)
 
   def create_adb_controller(self):
     """Returns an ADB controller which can communicate with this simulator."""
-    return adb_controller.AdbController(
-        device_name=self.adb_device_name(), **self._adb_controller_args)
+    return adb_controller.AdbController(self._config.adb_controller)
 
   def create_log_stream(self) -> log_stream.LogStream:
     return adb_log_stream.AdbLogStream(
         adb_command_prefix=self._adb_controller.command_prefix(),
-        verbose=self._verbose_logs)
+        verbose=self._config.verbose_logs,
+    )
 
   def _launch_impl(self) -> None:
     """Prepares an Android Emulator for RL interaction.
 
     The behavior depends on `self._num_launch_attempts`'s value:
-      * <= self._launch_n_times_without_reboot   -> Normal boot behavior.
-      * > self._launch_n_times_without_reboot but <=
-          self._launch_n_times_without_reinstall -> reboot (i.e. process is
-          killed and started again).
-      * > self._launch_n_times_without_reinstall -> reinstall (i.e. process is
-          killed, emulator files are deleted and the process started again).
+      * <= self._config.launch_n_times_without_reboot   -> Normal boot behavior.
+      * > self._config.launch_n_times_without_reboot but <=
+          self._config.launch_n_times_without_reinstall -> reboot (i.e. process
+          is killed and started again).
+      * > self._config.launch_n_times_without_reinstall -> reinstall (i.e.
+          process is killed, emulator files are deleted and the process started
+          again).
     """
 
     logging.info('Attempt %r at launching the Android Emulator (%r)',
@@ -230,19 +217,27 @@ class EmulatorSimulator(base_simulator.BaseSimulator):
 
     if self._launcher is not None:
       # If not the first time, then shutdown the emulator first.
-      if (self._emulator_stub is not None and
-          self._num_launch_attempts > self._launch_n_times_without_reboot):
+      if (
+          self._emulator_stub is not None
+          and self._num_launch_attempts
+          > self._config.launch_n_times_without_reboot
+      ):
         self._shutdown_emulator()
         # Subsequent attempts cause the emulator files to be reinstalled.
-        if self._num_launch_attempts > self._launch_n_times_without_reinstall:
+        if (
+            self._num_launch_attempts
+            > self._config.launch_n_times_without_reinstall
+        ):
           logging.info('Closing emulator (%r)', self.adb_device_name())
           self._launcher.close()
           self._launcher = emulator_launcher.EmulatorLauncher(
-              **self._emulator_launcher_args)
+              config=self._config.emulator_launcher,
+              adb_controller_config=self._config.adb_controller,
+          )
       self._launcher.launch_emulator_process()
     # Establish grpc connection to emulator process.
     self._emulator_stub, self._snapshot_stub = self._connect_to_emulator(
-        self._grpc_port
+        self._config.emulator_launcher.grpc_port
     )
 
     # Confirm booted status.
@@ -272,6 +267,7 @@ class EmulatorSimulator(base_simulator.BaseSimulator):
       * If an error occurred during the snapshot loading process, the status
         will be `ERROR` and the `error_message` field will be filled.
     """
+    assert self._snapshot_stub is not None
     snapshot_name = request.args.get('snapshot_name', _DEFAULT_SNAPSHOT_NAME)
     snapshot_list = self._snapshot_stub.ListSnapshots(
         snapshot_service_pb2.SnapshotFilter(
@@ -317,6 +313,7 @@ class EmulatorSimulator(base_simulator.BaseSimulator):
       * If an error occurred during the snapshot saving process, the status
         will be `ERROR` and the `error_message` field will be filled.
     """
+    assert self._snapshot_stub is not None
     snapshot_name = request.args.get('snapshot_name', _DEFAULT_SNAPSHOT_NAME)
     snapshot_result = self._snapshot_stub.SaveSnapshot(
         snapshot_service_pb2.SnapshotPackage(snapshot_id=snapshot_name)
@@ -366,6 +363,9 @@ class EmulatorSimulator(base_simulator.BaseSimulator):
   def _confirm_booted(self, startup_wait_time_sec: int = 300):
     """Waits until the emulator is fully booted."""
 
+    assert (
+        self._emulator_stub is not None
+    ), 'Emulator stub has not been initialized yet.'
     start_time = time.time()
     deadline = start_time + startup_wait_time_sec
     success = False
@@ -403,7 +403,9 @@ class EmulatorSimulator(base_simulator.BaseSimulator):
           3 identifier: Identifies a particular finger in a multitouch event.
     """
 
-    assert self._emulator_stub, 'Emulator stub has not been initialized yet.'
+    assert (
+        self._emulator_stub is not None
+    ), 'Emulator stub has not been initialized yet.'
     touch_events = [
         emulator_controller_pb2.Touch(
             x=t[0], y=t[1], pressure=int(t[2]), identifier=t[3])
@@ -421,11 +423,15 @@ class EmulatorSimulator(base_simulator.BaseSimulator):
         See the emulator_controller_pb2 for details.
       event_type: Type of key event to be sent.
     """
+
     event_types = emulator_controller_pb2.KeyboardEvent.KeyEventType.keys()
     if event_type not in event_types:
       raise ValueError(
           f'Event type must be one of {event_types} but is {event_type}.')
 
+    assert (
+        self._emulator_stub is not None
+    ), 'Emulator stub has not been initialized yet.'
     self._emulator_stub.sendKey(
         emulator_controller_pb2.KeyboardEvent(
             codeType=emulator_controller_pb2.KeyboardEvent.KeyCodeType.XKB,
@@ -437,9 +443,12 @@ class EmulatorSimulator(base_simulator.BaseSimulator):
     )
 
   @_reconnect_on_grpc_error
-  def get_screenshot(self) -> np.ndarray:
+  def _get_screenshot_impl(self) -> np.ndarray:
     """Fetches the latest screenshot from the emulator."""
-    assert self._emulator_stub, 'Emulator stub has not been initialized yet.'
+
+    assert (
+        self._emulator_stub is not None
+    ), 'Emulator stub has not been initialized yet.'
     assert self._image_format, 'ImageFormat has not been initialized yet.'
     image_proto = self._emulator_stub.getScreenshot(self._image_format)
     h, w = image_proto.format.height, image_proto.format.width
@@ -455,6 +464,8 @@ class EmulatorSimulator(base_simulator.BaseSimulator):
       logging.info('Emulator (%r) is not up.', self.adb_device_name())
       return
 
+    assert self._launcher is not None, 'Launcher is already down.'
+
     logging.info('Shutting down the emulator (%r)...', self.adb_device_name())
     self._emulator_stub.setVmState(
         emulator_controller_pb2.VmRunState(
@@ -462,6 +473,8 @@ class EmulatorSimulator(base_simulator.BaseSimulator):
     self._launcher.confirm_shutdown()
 
   def close(self):
+    super().close()
+
     if self._launcher is not None:
       self._shutdown_emulator()
       logging.info('Closing emulator (%r)', self.adb_device_name())