Fix launcher for simulator case

apache · Jun 8, 2022 · af22ddd · af22ddd
1 parent f5d0490
commit af22ddd
Show file tree

Hide file tree

Showing 2 changed files with 16 additions and 8 deletions.
diff --git a/python/tvm/contrib/hexagon/build.py b/python/tvm/contrib/hexagon/build.py
@@ -22,6 +22,7 @@
 import multiprocessing as mp
 import os
 import pathlib
+import shutil
 import signal
 import socket
 import stat
@@ -125,6 +126,11 @@ def stop_server(self):
         """Stop the RPC server"""
         ...
 
+    @abc.abstractmethod
+    def cleanup_directory(self):
+        """Cleanup working directory"""
+        ...
+
     @abc.abstractmethod
     def _copy_to_remote(
         self, local_path: Union[str, pathlib.Path], remote_path: Union[str, pathlib.Path]
@@ -474,7 +480,7 @@ def _terminate_remote(self):
         )
 
     def cleanup_directory(self):
-        # Remove workspace directory on remote target
+        """Abstract method implementation. See description in HexagonLauncherRPC."""
         subprocess.Popen(self._adb_device_sub_cmd + ["shell", f"rm -rf {self._workspace}"])
 
     def start_server(self):
@@ -588,7 +594,8 @@ def _start(self):
         self._server_process.start()
 
     def cleanup_directory(self):
-        pass
+        """Abstract method implementation. See description in HexagonLauncherRPC."""
+        shutil.rmtree(self._workspace)
 
     def stop_server(self):
         """Abstract method implementation. See description in HexagonLauncherRPC."""

diff --git a/python/tvm/contrib/hexagon/pytest_plugin.py b/python/tvm/contrib/hexagon/pytest_plugin.py
@@ -160,7 +160,7 @@ def hexagon_server_process(
     """Initials and returns hexagon launcher if ANDROID_SERIAL_NUMBER is defined.
     This launcher is started only once per test session.
     """
-    if android_serial_number is None:
+    if android_serial_number is None or android_serial_number == "simulator":
         yield None
     else:
         # Requesting these fixtures sets up a local tracker, if one
@@ -194,21 +194,22 @@ def hexagon_launcher(
     android_serial_number,
 ) -> HexagonLauncherRPC:
     """Initials and returns hexagon launcher which reuses RPC info and Android serial number."""
-    if hexagon_server_process._serial_number != "simulator":
+    if android_serial_number is None:
+        yield None
+
+    if android_serial_number != "simulator":
         rpc_info = hexagon_server_process._rpc_info
-        serial_number = hexagon_server_process._serial_number
     else:
-        serial_number = android_serial_number
         rpc_info = {
             "rpc_tracker_host": tvm_tracker_host,
             "rpc_tracker_port": tvm_tracker_port,
             "rpc_server_port": rpc_server_port,
             "adb_server_socket": adb_server_socket,
         }
 
-    launcher = HexagonLauncher(serial_number=serial_number, rpc_info=rpc_info)
+    launcher = HexagonLauncher(serial_number=android_serial_number, rpc_info=rpc_info)
     try:
-        if hexagon_server_process._serial_number == "simulator":
+        if android_serial_number == "simulator":
             launcher.start_server()
         yield launcher
     finally: