Merge branch 'zdy'

desktop_env/controllers/setup.py

@@ -6,8 +6,12 @@ import uuid
 import os.path
 
 from typing import Dict, List
-from typing import Any
+from typing import Any, Union
 
+import logging
+logger = logging.getLogger("desktopenv.setup")
+
+import traceback
 
 class SetupController:
     def __init__(self, http_server: str, cache_dir: str):
@@ -47,31 +51,32 @@ class SetupController:
         # self._open_setup(config)
         # can add other setup steps
 
-    def _command_setup(self, command: str):
-        """
-        Directly send a command into the virtual machine os for setting up.
-        """
-        payload = json.dumps({"command": command})
-        headers = {
-            'Content-Type': 'application/json'
-        }
-        timeout = 5
-        timout_whitelist = ["vlc"]
-
-        try:
-
-            response = requests.post(self.http_server + "/execute", headers=headers, data=payload, timeout=timeout)
-            if response.status_code == 200:
-                print("Command executed successfully:", response.text)
-            else:
-                print("Failed to execute command. Status code:", response.status_code)
-        except requests.exceptions.Timeout as e:
-            if command in timout_whitelist:
-                print("Command executed successfully:", command)
-            else:
-                print("An error occurred while trying to execute the command:", e)
-        except requests.exceptions.RequestException as e:
-            print("An error occurred while trying to execute the command:", e)
+    # ZDY_COMMENT: merged with launch
+    #def _command_setup(self, command: str):
+        #"""
+        #Directly send a command into the virtual machine os for setting up.
+        #"""
+        #payload = json.dumps({"command": command})
+        #headers = {
+            #'Content-Type': 'application/json'
+        #}
+        #timeout = 5
+        #timout_whitelist = ["vlc"]
+#
+        #try:
+#
+            #response = requests.post(self.http_server + "/execute", headers=headers, data=payload, timeout=timeout)
+            #if response.status_code == 200:
+                #print("Command executed successfully:", response.text)
+            #else:
+                #print("Failed to execute command. Status code:", response.status_code)
+        #except requests.exceptions.Timeout as e:
+            #if command in timout_whitelist:
+                #print("Command executed successfully:", command)
+            #else:
+                #print("An error occurred while trying to execute the command:", e)
+        #except requests.exceptions.RequestException as e:
+            #print("An error occurred while trying to execute the command:", e)
 
     def _download_setup(self, files: List[Dict[str, str]]):
         """
@@ -110,12 +115,12 @@ class SetupController:
                             for chunk in response.iter_content(chunk_size=8192):
                                 if chunk:
                                     f.write(chunk)
-                        print("File downloaded successfully")
+                        logger.info("File downloaded successfully")
                         downloaded = True
                         break
 
                     except requests.RequestException as e:
-                        print(f"Failed to download {url}. Retrying... ({max_retries - i - 1} attempts left)")
+                        logger.error(f"Failed to download {url}. Retrying... ({max_retries - i - 1} attempts left)")
                 if not downloaded:
                     raise requests.RequestException(f"Failed to download {url}. No retries left. Error: {e}")
 
@@ -129,17 +134,18 @@ class SetupController:
                 "file_data": (os.path.basename(path), open(cache_path, "rb"))
             })
             headers = {"Content-Type": form.content_type}
-            print(form.content_type)
+            logger.debug(form.content_type)
 
             # send request to server to upload file
             try:
+                logger.debug("REQUEST ADDRESS: %s", self.http_server_setup_root + "/upload")
                 response = requests.post(self.http_server_setup_root + "/upload", headers=headers, data=form)
                 if response.status_code == 200:
-                    print("Command executed successfully:", response.text)
+                    logger.info("Command executed successfully: %s", response.text)
                 else:
-                    print("Failed to upload file. Status code:", response.text)
+                    logger.error("Failed to upload file. Status code: %s", response.text)
             except requests.exceptions.RequestException as e:
-                print("An error occurred while trying to send the request:", e)
+                logger.error("An error occurred while trying to send the request: %s", e)
 
     def _change_wallpaper_setup(self, path: str):
         # if not config:
@@ -160,14 +166,14 @@ class SetupController:
         try:
             response = requests.post(self.http_server_setup_root + "/change_wallpaper", headers=headers, data=payload)
             if response.status_code == 200:
-                print("Command executed successfully:", response.text)
+                logger.info("Command executed successfully: %s", response.text)
             else:
-                print("Failed to change wallpaper. Status code:", response.text)
+                logger.error("Failed to change wallpaper. Status code: %s", response.text)
         except requests.exceptions.RequestException as e:
-            print("An error occurred while trying to send the request:", e)
+            logger.error("An error occurred while trying to send the request: %s", e)
 
     def _tidy_desktop_setup(self, **config):
-        raise NotImplementedError
+        raise NotImplementedError()
 
     def _open_setup(self, path: str):
         # if not config:
@@ -187,8 +193,63 @@ class SetupController:
         try:
             response = requests.post(self.http_server_setup_root + "/open_file", headers=headers, data=payload)
             if response.status_code == 200:
-                print("Command executed successfully:", response.text)
+                logger.info("Command executed successfully: %s", response.text)
             else:
-                print("Failed to open file. Status code:", response.text)
+                logger.error("Failed to open file. Status code: %s", response.text)
         except requests.exceptions.RequestException as e:
-            print("An error occurred while trying to send the request:", e)
+            logger.error("An error occurred while trying to send the request: %s", e)
+
+    def _launch_setup(self, command: List[str]):
+        if not command:
+            raise Exception("Empty comman to launch.")
+
+        payload = json.dumps({"command": command})
+        headers = {"Content-Type": "application/json"}
+
+        try:
+            response = requests.post(self.http_server_setup_root + "/launch", headers=headers, data=payload)
+            if response.status_code == 200:
+                logger.info("Command executed successfully: %s", response.text)
+            else:
+                logger.error("Failed to launch application. Status code: %s", response.text)
+        except requests.exceptions.RequestException as e:
+            logger.error("An error occurred while trying to send the request: %s", e)
+
+    def _execute_setup(self, command: List[str], stdout: str = "", stderr: str = ""):
+        if not command:
+            raise Exception("Empty comman to launch.")
+
+        payload = json.dumps({"command": command})
+        headers = {"Content-Type": "application/json"}
+
+        try:
+            response = requests.post(self.http_server_setup_root + "/execute", headers=headers, data=payload)
+            if response.status_code == 200:
+                results: Dict[str, str] = response.json()
+                if stdout:
+                    with open(os.path.join(self.cache_dir, stdout), "w") as f:
+                        f.write(results["output"])
+                if stderr:
+                    with open(os.path.join(self.cache_dir, stderr), "w") as f:
+                        f.write(results["error"])
+                logger.info( "Command executed successfully: %s -> %s"
+                           , " ".join(command)
+                           , response.text
+                           )
+            else:
+                logger.error("Failed to launch application. Status code: %s", response.text)
+        except requests.exceptions.RequestException as e:
+            logger.error("An error occurred while trying to send the request: %s", e)
+            traceback.print_exc()
+
+    def _act_setup(self, action_seq: List[Union[Dict[str, Any], str]]):
+        # TODO
+        raise NotImplementedError()
+    def _replay_setup(self, trajectory: str):
+        """
+        Args:
+            trajectory (str): path to the replay trajectory file
+        """
+
+        # TODO
+        raise NotImplementedError()