feat: refactor boot environment logic to support async operations

- Updated boot_environment and related functions to be asynchronous, improving responsiveness and performance.
- Changed device connection and command execution to use async methods from async_adbutils.
- Added adbutils dependency to pyproject.toml for async support.
- Enhanced error handling during environment setup and task execution.

Author: rebootyang

eval/cli.py

@@ -86,9 +86,10 @@ def list_tasks(env_url):
 )
 @click.option("--env-serial", default="emulator-5554", help="Device serial to use.")
 def check(env_url, env_serial):
+    import asyncio
     env = AndroidEnvClient(env_url)
     try:
-        boot_environment(env, env_serial)
+        asyncio.run(boot_environment(env, env_serial))
         logger.info("Environment is healthy")
     except Exception as e:
         logger.error(f"Error booting environment: {e}")
@@ -160,7 +161,7 @@ async def run(
     env = AndroidEnvClient(env_url)
 
     try:
-        boot_environment(env, env_serial)
+        await boot_environment(env, env_serial)
     except Exception as e:
         logger.error(f"Error booting environment: {e}")
         logger.info(
@@ -214,7 +215,7 @@ async def run(
 
         for task_idx in range(num_tasks):
             try:
-                boot_environment(env, env_serial)
+                await boot_environment(env, env_serial)
             except Exception as e:
                 logger.error(f"Error booting environment: {e}")
                 logger.info(

eval/env/boot.py

@@ -11,7 +11,7 @@
     setup_keyboard,
     A11Y_SERVICE_NAME as DROIDRUN_A11Y_SERVICE_NAME,
 )
-from adbutils import adb, AdbDevice
+from async_adbutils import adb, AdbDevice
 import time
 
 logger = logging.getLogger(__name__)
@@ -22,7 +22,7 @@
 )
 DEFAULT_OVERLAY_OFFSET = -126
 
-def ensure_connected(serial: str) -> AdbDevice:
+async def ensure_connected(serial: str) -> AdbDevice:
     try:
         # For emulator devices (emulator-*), they're already connected locally
         # Only try network connect for IP addresses
@@ -32,27 +32,27 @@ def ensure_connected(serial: str) -> AdbDevice:
                 raise RuntimeError(f"Failed to connect: {res}")
 
         # Verify device is available
-        device = adb.device(serial)
+        device = await adb.device(serial)
         # Test if device is accessible
-        device.shell("echo test")
+        await device.shell("echo test")
         return device
     except Exception as e:
         raise RuntimeError(f"Device {serial} is not connected: {e}")
 
 
-def install_portal(device: AdbDevice):
+async def install_portal(device: AdbDevice):
     logger.info("Installing portal...")
 
     try:
         with download_portal_apk() as apk_path:
-            device.install(apk_path, uninstall=True, flags=["-g"], silent=False)
+            await device.install(apk_path, uninstall=True, flags=["-g"], silent=False)
             logger.info("Portal APK installed successfully")
     except Exception as e:
         raise RuntimeError(f"Failed to download and install portal APK: {e}")
 
     try:
         logger.info("Enabling portal as accessibility service...")
-        enable_portal_accessibility(
+        await enable_portal_accessibility(
             device, service_name=DROIDRUN_X_GOOGLE_A11Y_SERVICE_NAME
         )
         logger.info("Portal accessibility enabled successfully")
@@ -61,36 +61,36 @@ def install_portal(device: AdbDevice):
 
     try:
         logger.info("Setting up keyboard for environment...")
-        setup_keyboard(device)
+        await setup_keyboard(device)
         logger.info("Keyboard setup completed successfully!")
     except Exception as e:
         raise RuntimeError(f"Failed to setup keyboard: {e}")
 
 
-def check_portal(device: AdbDevice):
-    if not check_portal_accessibility(
+async def check_portal(device: AdbDevice):
+    if not await check_portal_accessibility(
         device, service_name=DROIDRUN_X_GOOGLE_A11Y_SERVICE_NAME
     ):
         raise RuntimeError("Accessibility settings invalid")
 
     try:
-        set_overlay_offset(device, DEFAULT_OVERLAY_OFFSET)
+        await set_overlay_offset(device, DEFAULT_OVERLAY_OFFSET)
         logger.info("Overlay offset set successfully")
     except Exception as e:
         raise RuntimeError(f"Failed to set overlay offset: {e}")
 
     try:
-        ping_portal(device)
+        await ping_portal(device)
     except Exception as e:
         raise RuntimeError(f"Failed to ping portal: {e}")
 
     try:
-        ping_portal_content(device)
+        await ping_portal_content(device)
     except Exception as e:
         raise RuntimeError(f"Failed to ping portal content: {e}")
 
     try:
-        ping_portal_tcp(device)
+        await ping_portal_tcp(device)
     except Exception as e:
         raise RuntimeError(f"Failed to ping portal TCP: {e}")
 
@@ -124,7 +124,7 @@ def wait_ready(env: AndroidEnvClient, timeout: int = 300):
     )
 
 
-def boot_environment(env: AndroidEnvClient, serial: str):
+async def boot_environment(env: AndroidEnvClient, serial: str):
     try:
         logger.info(f"Waiting for environment {env.base_url} to be ready...")
         wait_ready(env, timeout=600)
@@ -134,15 +134,15 @@ def boot_environment(env: AndroidEnvClient, serial: str):
         raise e
 
     try:
-        device = ensure_connected(serial)
+        device = await ensure_connected(serial)
     except Exception as e:
         logger.error(f"Environment {env.base_url} failed to connect via adb: {e}")
         raise e
 
     # check if portal is already installed
     try:
         logger.info(f"Checking portal for environment {env.base_url}...")
-        check_portal(device)
+        await check_portal(device)
         logger.info("Portal is installed and accessible. You're good to go!")
         return
     except Exception as e:
@@ -152,15 +152,15 @@ def boot_environment(env: AndroidEnvClient, serial: str):
 
     try:
         logger.info(f"Installing portal for environment {env.base_url}...")
-        install_portal(device)
+        await install_portal(device)
         logger.info(f"Portal installed successfully for environment {env.base_url}!")
     except Exception as e:
         logger.error(f"Environment {env.base_url} failed to install portal: {e}")
         raise e
 
     try:
         logger.info(f"Checking portal for environment {env.base_url}...")
-        check_portal(device)
+        await check_portal(device)
         logger.info("Portal is installed and accessible. You're good to go!")
     except Exception as e:
         logger.error(f"Environment {env.base_url} failed to check portal: {e}")

eval/portal/keepalive.py

@@ -6,7 +6,7 @@
 """
 
 import logging
-from adbutils import adb
+from async_adbutils import adb
 from droidrun.portal import toggle_overlay
 import threading
 

pyproject.toml

@@ -11,6 +11,7 @@ dependencies = [
     "droidrun",
     "android-world",
     "llama-index-llms-gemini>=0.5.0",
+    "adbutils>=2.12.0",
 ]
 
 [project.scripts]