diff --git a/src/devtools/deviceinfra/api/deviceinfo/proto/BUILD b/src/devtools/deviceinfra/api/deviceinfo/proto/BUILD
index 0c62a738b..6dbdca564 100644
--- a/src/devtools/deviceinfra/api/deviceinfo/proto/BUILD
+++ b/src/devtools/deviceinfra/api/deviceinfo/proto/BUILD
@@ -17,6 +17,7 @@ package(
     default_applicable_licenses = ["//:license"],
     default_visibility = [
         "//:deviceinfra_all_pkg",
+        "//java/com/google/cloud/test/common/wrangler:__subpackages__",
     ],
 )
 
diff --git a/src/java/com/google/devtools/mobileharness/infra/master/rpc/stub/BUILD b/src/java/com/google/devtools/mobileharness/infra/master/rpc/stub/BUILD
index 0563dc4e0..f1727ed23 100644
--- a/src/java/com/google/devtools/mobileharness/infra/master/rpc/stub/BUILD
+++ b/src/java/com/google/devtools/mobileharness/infra/master/rpc/stub/BUILD
@@ -45,6 +45,7 @@ java_library(
     name = "device_allocation",
     srcs = ["DeviceAllocationStub.java"],
     deps = [
+        "//java/com/google/devtools/deviceinfra/service/scheduler/common",
         "//src/devtools/mobileharness/infra/master/rpc/proto:device_allocation_service_java_proto",
         "//src/java/com/google/devtools/common/metrics/stability/rpc:exception",
         "@maven//:com_google_errorprone_error_prone_annotations",
diff --git a/src/java/com/google/devtools/mobileharness/infra/master/rpc/stub/DeviceAllocationStub.java b/src/java/com/google/devtools/mobileharness/infra/master/rpc/stub/DeviceAllocationStub.java
index 70a224c96..4a0ebf7d7 100644
--- a/src/java/com/google/devtools/mobileharness/infra/master/rpc/stub/DeviceAllocationStub.java
+++ b/src/java/com/google/devtools/mobileharness/infra/master/rpc/stub/DeviceAllocationStub.java
@@ -17,8 +17,8 @@
 package com.google.devtools.mobileharness.infra.master.rpc.stub;
 
 import com.google.devtools.common.metrics.stability.rpc.RpcExceptionWithErrorId;
+import com.google.devtools.deviceinfra.service.scheduler.common.DeviceAllocationCallback;
 import com.google.devtools.mobileharness.infra.master.rpc.proto.DeviceAllocationServiceProto.AllocationRequest;
-import com.google.devtools.mobileharness.infra.master.rpc.proto.DeviceAllocationServiceProto.AllocationResponse;
 import com.google.devtools.mobileharness.infra.master.rpc.proto.DeviceAllocationServiceProto.DeallocationRequest;
 import com.google.devtools.mobileharness.infra.master.rpc.proto.DeviceAllocationServiceProto.DeallocationResponse;
 import com.google.devtools.mobileharness.infra.master.rpc.proto.DeviceAllocationServiceProto.GetAvailableDevicesResponse;
@@ -26,8 +26,8 @@
 
 /** RPC stub for talking to Master DeviceAllocationService. */
 public interface DeviceAllocationStub extends AutoCloseable {
-  /** Allocates devices. */
-  AllocationResponse allocate(AllocationRequest request) throws RpcExceptionWithErrorId;
+  /** Allocates devices asynchronously and handles the response in the callback. */
+  void allocate(AllocationRequest request, DeviceAllocationCallback deviceAllocationCallback);
 
   /** Deallocates devices. */
   @CanIgnoreReturnValue