Enhance state management and logging functionality

- Updated `serve` function in `run.py` to utilize the `workers` argument in development mode.
- Improved `executed_state` function in `executed_state.py` to handle multiple new states more efficiently, including error handling for state insertion.
- Modified `LogsManager` in `logs_manager.py` to dynamically set logging level based on the application mode (development or production).
- Added logging for state creation in `create_next_state.py` to improve traceability.

These changes aim to enhance the application's performance and maintainability in different environments.

Author: NiveditJain

state-manager/app/controller/executed_state.py

@@ -1,4 +1,5 @@
 from beanie import PydanticObjectId
+from beanie.operators import In
 from app.models.executed_models import ExecutedRequestModel, ExecutedResponseModel
 
 from fastapi import HTTPException, status, BackgroundTasks
@@ -36,11 +37,11 @@ async def executed_state(namespace_name: str, state_id: PydanticObjectId, body:
             state.parents = {**state.parents, state.identifier: state.id}
 
             await state.save()
-
             background_tasks.add_task(create_next_state, state)
 
+            new_states = []
             for output in body.outputs[1:]:
-                new_state = State(
+                new_states.append(State(
                     node_name=state.node_name,
                     namespace_name=state.namespace_name,
                     identifier=state.identifier,
@@ -54,9 +55,20 @@ async def executed_state(namespace_name: str, state_id: PydanticObjectId, body:
                         **state.parents,
                         state.identifier: state.id
                     }
-                )
-                await new_state.save()
-                background_tasks.add_task(create_next_state, new_state)
+                ))
+
+            if len(new_states) > 0:
+                inserted_ids = (await State.insert_many(new_states)).inserted_ids
+
+                inserted_states = await State.find(
+                    In(State.id, inserted_ids)
+                ).to_list()
+                
+                if len(inserted_states) != len(new_states):
+                    raise Exception(f"Failed to insert all new states. Expected {len(new_states)} states, but only {len(inserted_states)} were inserted")
+
+                for inserted_state in inserted_states:
+                    background_tasks.add_task(create_next_state, inserted_state)
 
         return ExecutedResponseModel(status=StateStatusEnum.EXECUTED)
 

state-manager/app/singletons/logs_manager.py

@@ -1,5 +1,7 @@
 import structlog
 import logging
+import os
+import sys
 from .SingletonDecorator import singleton
 
 
@@ -28,9 +30,37 @@ def __init__(self):
         handler.setFormatter(formatter)
         logger = logging.getLogger()
         logger.addHandler(handler)
-        logger.setLevel(logging.INFO)
+        
+        # Check if running in development mode
+        # Development mode is determined by the --mode argument passed to run.py
+        is_development = self._is_development_mode()
+        
+        if is_development:
+            # In development mode, set level to WARNING to disable INFO logs
+            logger.setLevel(logging.WARNING)
+        else:
+            # In production mode, keep INFO level
+            logger.setLevel(logging.INFO)
 
         self.logger = structlog.get_logger()
 
+    def _is_development_mode(self) -> bool:
+        """
+        Check if the application is running in development mode.
+        Development mode is determined by checking if '--mode' 'development' 
+        is in the command line arguments.
+        """
+        # Check command line arguments for development mode
+        if '--mode' in sys.argv:
+            try:
+                mode_index = sys.argv.index('--mode')
+                if mode_index + 1 < len(sys.argv) and sys.argv[mode_index + 1] == 'development':
+                    return True
+            except (ValueError, IndexError):
+                pass
+        
+        # Fallback: check environment variable
+        return os.getenv('MODE', '').lower() == 'development'
+
     def get_logger(self):
         return self.logger

state-manager/app/tasks/create_next_state.py

@@ -14,6 +14,7 @@
 logger = LogsManager().get_logger()
 
 async def create_next_state(state: State):
+    logger.info(f"Creating next state for {state.identifier}")
     graph_template = None
 
     if state is None or state.id is None:

state-manager/run.py

@@ -12,10 +12,10 @@
 
 def serve():
     mode = args.mode
+    workers = args.workers
     if mode == "development":
-        uvicorn.run("app.main:app", reload=True, host="0.0.0.0", port=8000)
+        uvicorn.run("app.main:app", workers=workers, reload=True, host="0.0.0.0", port=8000)
     elif mode == "production":
-        workers = args.workers
         print(f"Running with {workers} workers")
         uvicorn.run("app.main:app", workers=workers, host="0.0.0.0", port=8000)
     else: