Provide the local llama service

src/aios_kernel/compute_kernel.py

@@ -65,7 +65,7 @@ class ComputeKernel:
 
     def _schedule(self, task) -> ComputeNode:
         for node in self.compute_nodes.values():
-            if node.is_support(task.task_type) is True:
+            if node.is_support(task) is True:
                 return node
         logger.warning(
             f"task {task.display()} is not support by any compute node")

src/aios_kernel/compute_node.py

@@ -28,7 +28,7 @@ class ComputeNode(ABC):
         pass
 
     @abstractmethod
-    def is_support(self, task_type: ComputeTaskType) -> bool:
+    def is_support(self, task: ComputeTask) -> bool:
         pass
 
     @abstractmethod
@@ -41,7 +41,6 @@ class ComputeNode(ABC):
     def get_fee_type(self) -> str:
         return "free"
 
-
 class LocalComputeNode(ComputeNode):
     def display(self) -> str:
         return super().display()

src/aios_kernel/local_llama_compute_node.py

@@ -0,0 +1,86 @@
+
+import logging
+import requests
+from typing import Optional
+from pydantic import BaseModel
+
+from .compute_task import ComputeTask, ComputeTaskState, ComputeTaskType
+from .queue_compute_node import Queue_ComputeNode
+
+logger = logging.getLogger(__name__)
+
+"""
+This is a custom implementation, it should be redesigned.
+"""
+
+class LocalLlama_ComputeNode(Queue_ComputeNode):
+    async def execute_task(self, task: ComputeTask) -> {
+        "content": str,
+        "message": str,
+        "state": ComputeTaskState,
+        "error": {
+            "code": int,
+            "message": str,
+        }
+    }:
+        class GenerateResponse(BaseModel):
+            error: Optional[int]
+            msg: Optional[str]
+            results: Optional[str]
+
+        try:    
+            body = {
+                "prompts": task.params["prompts"]
+            }
+            
+            response = requests.post("http://aigc:7880/generate", data = body, verify=False)
+            response.close()
+
+            logger.info(f"LocalLlama_ComputeNode task responsed, request: {body}, status-code: {response.status_code}, headers: {response.headers}, content: {response.content}")
+
+            if response.status_code != 200:
+                return {
+                    "state": ComputeTaskState.ERROR,
+                    "error": {
+                        "code": response.status_code,
+                        "message": "http request failed: " + response.status_code
+                    }
+                }
+            else:
+                resp = GenerateResponse.parse_raw(response.content.decode("utf-8"))
+                if resp.error:
+                    return {
+                        "state": ComputeTaskState.ERROR,
+                        "error": {
+                            "code": resp.error,
+                            "message": "local llama failed:" + resp.msg
+                        }
+                    }
+                else:
+                    return {
+                        "content": str(resp.results),
+                        "message": {}
+                    }
+        except Exception as err:
+            import traceback
+            logger.error(f"{traceback.format_exc()}, error: {err}")
+            
+            return {
+                "state": ComputeTaskState.ERROR,
+                "error": {
+                    "code": -1,
+                    "message": "unknown exception: " + str(err)
+                }
+            }
+
+    def display(self) -> str:
+        return f"LocalLlama_ComputeNode: {self.node_id}"
+
+    def get_capacity(self):
+        pass
+
+    def is_support(self, task: ComputeTask) -> bool:
+        return task.task_type == ComputeTaskType.LLM_COMPLETION and (not task.params["model_name"] or task.params["model_name"] == "llama")
+
+    def is_local(self) -> bool:
+        return True

src/aios_kernel/open_ai_node.py

@@ -103,8 +103,8 @@ class OpenAI_ComputeNode(ComputeNode):
     def get_capacity(self):
         pass
 
-    def is_support(self, task_type: ComputeTaskType) -> bool:
+    def is_support(self, task: ComputeTask) -> bool:
-        return task_type == ComputeTaskType.LLM_COMPLETION
+        return task.task_type == ComputeTaskType.LLM_COMPLETION and (not task.params["model_name"] or task.params["model_name"] == "gpt-4-0613")
 
     def is_local(self) -> bool:
         return False

src/aios_kernel/queue_compute_node.py

@@ -0,0 +1,69 @@
+
+import asyncio
+from asyncio import Queue
+import logging
+from abc import abstractmethod
+
+from .compute_task import ComputeTask, ComputeTaskResult, ComputeTaskState, ComputeTaskType
+from .compute_node import ComputeNode
+
+logger = logging.getLogger(__name__)
+
+class Queue_ComputeNode(ComputeNode):
+    def __init__(self):
+        super().__init__()
+        self.task_queue = Queue()
+
+    @abstractmethod
+    async def execute_task(self, task: ComputeTask) -> {
+        "content": str,
+        "message": str,
+        "state": ComputeTaskState,
+        "error": {
+            "code": int,
+            "message": str,
+        }
+    }:
+        pass
+
+    async def push_task(self, task: ComputeTask, proiority: int = 0):
+        logger.info(f"{self.display()} push task: {task.display()}")
+        self.task_queue.put_nowait(task)
+
+    async def remove_task(self, task_id: str):
+        pass
+
+    async def _run_task(self, task: ComputeTask):
+        task.state = ComputeTaskState.RUNNING
+        
+        resp = await self.execute_task(task)
+
+        result = ComputeTaskResult()
+        result.set_from_task(task)
+
+        task.state = resp["state"]
+
+        if task.state == ComputeTaskState.ERROR:
+            task.error_str = resp["error"]["message"]
+
+
+        result.worker_id = self.node_id
+        result.result_str = resp["content"]
+        result.result_message = resp["message"]
+
+        return result
+
+    def start(self):
+        async def _run_task_loop():
+            while True:
+                task = await self.task_queue.get()
+                logger.info(f"{self.display()} get task: {task.display()}")
+                result = await self._run_task(task)
+                if result is not None:
+                    task.result = result
+
+        asyncio.create_task(_run_task_loop())
+
+
+    def get_task_state(self, task_id: str):
+        pass

src/aios_kernel/stability_node.py

@@ -129,8 +129,8 @@ class Stability_ComputeNode(ComputeNode):
     def get_capacity(self):
         pass
 
-    def is_support(self, task_type: ComputeTaskType) -> bool:
+    def is_support(self, task: ComputeTask) -> bool:
-        return task_type == ComputeTaskType.TEXT_2_IMAGE
+        return task.task_type == ComputeTaskType.TEXT_2_IMAGE
 
     def is_local(self) -> bool:
         return False