Add session manager

Author: xxx7xxxx

src/rest/rest.py

@@ -11,8 +11,6 @@
     UploadFileRequest,
     DeleteFilesRequest,
     ListDirectoryResponse,
-    CreateNamespaceRequest,
-    AudioNamespaceProgressInitial,
 )
 from src.service.audio import AudioService
 from src.service.file import FileService
@@ -64,18 +62,9 @@ async def list_namespaces(self):
         namespaces = self.namespace_service.get_namespaces()
         return {"namespaces": namespaces}
 
-    async def new_namespace(self, new_namespace_request: CreateNamespaceRequest):
+    async def new_namespace(self):
         """Create a new namespace."""
-        args = new_namespace_request.args
-        if new_namespace_request.service_name == "audio":
-            namespace = self.namespace_service.create_namespace(new_namespace_request.service_name, args)
-            namespace.progress = AudioNamespaceProgressInitial()
-            self.namespace_service.submit_namespace(namespace)
-            audio_service = AudioService(args["source_dir"], args["output_dir"], namespace)
-            audio_service.audio_service()
-            return namespace
-
-        namespace = self.namespace_service.create_namespace(new_namespace_request.service_name, new_namespace_request.args)
+        namespace = self.namespace_service.create_namespace()
         return namespace
 
     async def change_namespace(self, namespace_id: str, update_request: UpdateNamespaceRequest):

src/service/namespace.py

@@ -38,7 +38,7 @@ def filter_namespaces(self, fn: Callable[[Namespace], bool]) -> List[Namespace]:
         """
         return sorted(list(filter(fn, self._namespaces.values())), key=lambda t: t.createdAt)
 
-    def create_namespace(self, service_name: str, args: dict) -> Namespace:
+    def create_namespace(self) -> Namespace:
         """Create a new namespace."""
         namespace_id = str(uuid.uuid4())
         namespace_name = f"Namespace-{namespace_id[:8]}"
@@ -51,17 +51,10 @@ def create_namespace(self, service_name: str, args: dict) -> Namespace:
             name=namespace_name,
             createdAt=created_at,
             homePath=home_path,
-            service_name=service_name,
-            args=args,
-            progress=Progress(),
         )
         self._namespaces[namespace_id] = namespace
         return namespace
 
-    def submit_namespace(self, namespace: Namespace):
-        self._save_namespace_metadata(namespace)
-        self._namespaces[namespace.namespaceID] = namespace
-
     def get_namespaces(self) -> List[Namespace]:
         """Get all namespaces."""
         namespaces = []

src/service/session.py

@@ -0,0 +1,95 @@
+import threading
+from functools import wraps
+from typing import Optional, Dict, Any
+from enum import Enum
+
+class Status(Enum):
+    RUNNING = "Running"
+    COMPLETED = "Completed"
+    FAILED = "Failed"
+
+class SessionManager:
+    """Manages training session, ensuring single GPU task execution and tracking task state."""
+
+    _instance = None
+    _lock = threading.Lock()
+
+    def __new__(cls):
+        """Singleton pattern to ensure only one instance of SessionManager exists."""
+        if not cls._instance:
+            with cls._lock:
+                if not cls._instance:
+                    cls._instance = super(SessionManager, cls).__new__(cls)
+                    cls._instance.current_session = None
+        return cls._instance
+
+    def start_session(self, task_name: str):
+        """Attempts to start a new session; rejects if another task is already running."""
+        if self.current_session and self.current_session["status"] == Status.RUNNING:
+            raise RuntimeError(f"A task '{self.current_session['task_name']}' is already running. Cannot submit another task!")
+
+        self.current_session = {
+            "task_name": task_name,
+            "status": Status.RUNNING,
+            "error": None,  # Stores error details if task fails
+        }
+
+    def end_session(self, result: Any):
+        """Marks task as completed successfully."""
+        if self.current_session:
+            self.current_session["status"] = Status.COMPLETED
+            self.current_session["result"] = result
+
+    def fail_session(self, error: str):
+        """Marks task as failed and stores error information."""
+        if self.current_session:
+            self.current_session["status"] = Status.FAILED
+            self.current_session["error"] = error
+
+    def update_session_info(self, info: Dict[str, Any]):
+        """Updates task session with arbitrary info."""
+        if not self.current_session or self.current_session["status"] != Status.RUNNING:
+            raise RuntimeError("No active task to update session info!")
+
+        self.current_session.update(info)
+    def get_session_info(self) -> Optional[Dict[str, Any]]:
+        """Returns current task state information."""
+        return self.current_session
+
+# Decorator to wrap task execution logic
+def session_guard(task_name: str):
+    """Ensures tasks are managed within SessionManager and handles failure states."""
+    def decorator(func):
+        @wraps(func)
+        def wrapper(*args, **kwargs):
+            session_manager = SessionManager()
+
+            try:
+                session_manager.start_session(task_name)
+                result = func(*args, **kwargs)  # Execute the training task
+                session_manager.end_session(result)
+                return result
+            except Exception as e:
+                session_manager.fail_session(str(e))  # Record failure details
+                # NOTICE: No Re-raise exception here,
+                # as we capture the error and record it in session info.
+                # raise e
+        return wrapper
+    return decorator
+
+# Example for using SessionManager and session_guard decorator.
+@session_guard("TrainingModel")
+def train_model():
+    session_manager = SessionManager()
+
+    for epoch in range(1, 6):
+        if epoch == 3:  # Simulate task failure
+            raise RuntimeError("Error occurred at epoch 3!")
+
+        session_manager.update_session_info({
+            "progress": epoch / 5,
+            "loss": 0.05 * (6 - epoch),
+            "epoch": epoch,
+        })
+
+    return "Training Completed"