workspace_ticket handler added

moai/conf/engine/serve/handlers/input/azure/blob.yaml

@@ -4,5 +4,6 @@ _target_: moai.serve.handlers.azure.AzureBlobInputHandler
 connection_string: StorageAccessToken # ${env:AZURE_STORAGE_CONNECTION_STRING} # retrieve automatically from env
 container_name: StorageContainer # ${env:AZURE_STORAGE_CONTAINER_NAME}
 blob_paths: ???
-working_dir: ???
+working_dir: null
+json_key: 'workspace_run'
 alias: ???

moai/conf/engine/serve/handlers/input/workspace_ticket.yaml

@@ -0,0 +1,5 @@
+# @package handlers.preprocess.workspace_ticket
+
+_target_: moai.serve.handlers.workspace_ticket.WorkSpaceTicketHandler
+working_dir: ???
+json_key: 'workspace_run'

moai/conf/engine/serve/handlers/output/azure/blob.yaml

@@ -4,6 +4,7 @@ _target_: moai.serve.handlers.azure.AzureBlobOutputHandler
 connection_string: StorageAccessToken # ${env:AZURE_STORAGE_CONNECTION_STRING} # retrieve automatically from env
 container_name: StorageContainer # ${env:AZURE_STORAGE_CONTAINER_NAME}
 blob_paths: ???
-working_dir: ???
+working_dir: null
+json_key: 'workspace_run'
 alias: ???
 overwrite: true

moai/conf/engine/serve/handlers/output/cleanup.yaml

@@ -1,3 +1,4 @@
-# @package handlers.postprocess._name_
-_target_: moai.serve.handlers.CleanUp
-dirs: ???
+# @package handlers.postprocess.cleanup
+_target_: moai.serve.handlers.cleanup.CleanUp
+dirs: null
+json_key: 'workspace_run'

moai/serve/handlers/azure/blob.py

@@ -17,6 +17,7 @@ def __init__(
         container_name: str,  # name of the container to download data from
         blob_paths: typing.List[str],  # keys to extract resources from json
         working_dir: str,  # path to working dir
+        json_key: str,
         alias: typing.List[str],  # names of files to be saved
     ):
         """
@@ -35,6 +36,7 @@ def __init__(
         self.connection_string = connection_string
         self.container_name = container_name
         self.working_dir = working_dir
+        self.json_key = json_key
         self.blob_paths = blob_paths
         self.blob_acecessors = [_create_accessor(bl_path) for bl_path in blob_paths]
         self.alias = alias
@@ -45,6 +47,8 @@ def __init__(
     def __call__(
         self, json: typing.Mapping[str, typing.Any], void: typing.Any
     ) -> typing.Any:
+        if self.working_dir is None:
+            self.working_dir = json[self.json_key]
         # initialize connection to Azure Blob Storage
         connect_str = json[self.connection_string]
         try:
@@ -79,6 +83,7 @@ def __init__(
         blob_paths: typing.List[str],  # keys to extract resources from json
         working_dir: str,  # path to working dir
         alias: typing.List[str],  # names of files to be uploaded
+        json_key: str,
         overwrite: bool = True,  # overwrite existing files
     ):
         """
@@ -102,6 +107,7 @@ def __init__(
         self.blob_paths = blob_paths
         self.blob_acecessors = [_create_accessor(bl_path) for bl_path in blob_paths]
         self.working_dir = working_dir
+        self.json_key = json_key
         self.alias = alias
         self.overwrite = overwrite
         log.info(
@@ -115,6 +121,8 @@ def __call__(
         # NOTE: void is the input json response
         # TODO: need to check batched inference
         input_json = void[0].get("body") or void[0].get("raw")
+        if self.working_dir is None:
+            self.working_dir = input_json[self.json_key]
         # initialize connection to Azure Blob Storage
         connect_str = input_json[self.connection_string]
         blob_service_client = BlobServiceClient.from_connection_string(

moai/serve/handlers/cleanup.py

@@ -4,11 +4,16 @@
 from collections.abc import Callable
 
 log = logging.getLogger(__name__)
+import shutil
 
 
 class CleanUp(Callable):
-    def __init__(self, dirs: typing.List[str]) -> None:
+    def __init__(
+        self, dirs: typing.List[str], workdir_from_json: str, json_key: str
+    ) -> None:
         self.dirs = [dirs] if type(dirs) is str else dirs  # directories to be deleted
+        self.workdir_from_json = workdir_from_json
+        self.json_key = json_key
         """
         Responsible for deleting data from Docker worker.
 
@@ -19,6 +24,18 @@ def __init__(self, dirs: typing.List[str]) -> None:
     def __call__(
         self, data: typing.Mapping[str, typing.Any], void: typing.Any
     ) -> typing.Any:
+        if self.workdir_from_json:
+            input_json = void[0].get("body") or void[0].get("raw")
+            folder = input_json[self.json_key]
+            assert os.path.isdir(folder)
+            try:
+                shutil.rmtree(folder)
+            except OSError as e:
+                log.error(f"Error: {folder} : {e.strerror}")
+            return [
+                {"is_success": True, "message": "Directories deleted successfully."}
+            ]
+
         for folder in self.dirs:
             try:
                 # check if is folder or file

moai/serve/handlers/workspace_ticket.py

@@ -0,0 +1,49 @@
+import logging
+import os
+import typing
+from collections.abc import Callable
+
+log = logging.getLogger(__name__)
+log.setLevel(logging.DEBUG)
+
+import tempfile
+
+
+class WorkSpaceTicketHandler(Callable):
+    def __init__(
+        self,
+        working_dir: str,  # aboslute path to working dir
+        json_key: str = "workspace_run",
+    ):
+        """
+        Responsible for creating a working dir and adding its corresponfing path to json data
+        """
+        self.working_dir = working_dir
+        self.json_key = json_key
+
+        if not os.path.isabs(working_dir):
+            log.error(f"working_dir ({working_dir}) is not absolute")
+        os.makedirs(self.working_dir, exist_ok=True)
+
+    def __call__(
+        self, json: typing.Mapping[str, typing.Any], void: typing.Any
+    ) -> typing.Any:
+        generated_dir_path = tempfile.mkdtemp(dir=self.working_dir)
+        json[self.json_key] = generated_dir_path
+        return {}  # generated_dir_path
+
+
+if __name__ == "__main__":
+    import os
+    import platform
+
+    if platform.system() == "Windows":
+        working_dir = f"C:/Users/{os.environ['USERNAME']}/AppData/Local/Temp"
+        assert os.path.isabs(working_dir)
+        print(working_dir)
+
+        data = {}
+        ticket_call = WorkSpaceTicketHandler(working_dir)
+        tmp_path = ticket_call(data, None)
+        print(tmp_path)
+        print(data)