fix: Simplify module loading by removing redundant error handling and introduce custom exceptions for agent and server errors

Author: shima004

adf_core_python/core/agent/module/module_manager.py

@@ -51,10 +51,7 @@ def get_module(self, module_name: str, default_module_name: str) -> AbstractModu
             )
             class_name = default_module_name
 
-        try:
-            module_class: type = self._load_module(class_name)
-        except (ImportError, AttributeError) as e:
-            raise RuntimeError(f"Failed to load module {class_name}") from e
+        module_class: type = self._load_module(class_name)
 
         instance = self._modules.get(module_name)
         if instance is not None:
@@ -80,10 +77,7 @@ def get_extend_action(
             action_name, default_action_name
         )
 
-        try:
-            action_class: type = self._load_module(class_name)
-        except (ImportError, AttributeError) as e:
-            raise RuntimeError(f"Failed to load action {class_name}") from e
+        action_class: type = self._load_module(class_name)
 
         instance = self._actions.get(action_name)
         if instance is not None:
@@ -109,10 +103,7 @@ def get_channel_subscriber(
             channel_subscriber_name, default_channel_subscriber_name
         )
 
-        try:
-            channel_subscriber_class: type = self._load_module(class_name)
-        except (ImportError, AttributeError) as e:
-            raise RuntimeError(f"Failed to load channel subscriber {class_name}") from e
+        channel_subscriber_class: type = self._load_module(class_name)
 
         instance = self._channel_subscribers.get(channel_subscriber_name)
         if instance is not None:
@@ -134,12 +125,7 @@ def get_message_coordinator(
             message_coordinator_name, default_message_coordinator_name
         )
 
-        try:
-            message_coordinator_class: type = self._load_module(class_name)
-        except (ImportError, AttributeError) as e:
-            raise RuntimeError(
-                f"Failed to load message coordinator {class_name}"
-            ) from e
+        message_coordinator_class: type = self._load_module(class_name)
 
         instance = self._message_coordinators.get(message_coordinator_name)
         if instance is not None:

adf_core_python/core/launcher/connect/component_launcher.py

@@ -4,6 +4,8 @@
 
 from adf_core_python.core.agent.agent import Agent
 from adf_core_python.core.launcher.connect.connection import Connection
+from adf_core_python.core.launcher.connect.error.agent_error import AgentError
+from adf_core_python.core.launcher.connect.error.server_error import ServerError
 
 
 class ComponentLauncher:
@@ -35,11 +37,18 @@ def connect(self, agent: Agent, _request_id: int) -> None:
 
         try:
             connection.parse_message_from_kernel()
-        except Exception as e:
+        except AgentError as e:
             self.logger.exception(
-                f"Agent threw an exception {e}",
+                f"Agent error: {e}",
                 exception=str(e),
             )
+        except ServerError as e:
+            if isinstance(e.__cause__, EOFError):
+                self.logger.info(
+                    f"Connection closed by server (request_id={_request_id})"
+                )
+            else:
+                self.logger.exception("Server error", exception=str(e))
 
     def generate_request_id(self) -> int:
         self.request_id += 1

adf_core_python/core/launcher/connect/connection.py

@@ -3,6 +3,9 @@
 
 import rcrs_core.connection.rcrs_encoding_utils as rcrs_encoding_utils
 
+from adf_core_python.core.launcher.connect.error.agent_error import AgentError
+from adf_core_python.core.launcher.connect.error.server_error import ServerError
+
 
 class Connection:
     def __init__(self, host: str, port: int) -> None:
@@ -34,12 +37,20 @@ def parse_message_from_kernel(self) -> None:
 
         Raises
         ------
-        IOError
+        ServerError
             If there is an error reading from the socket
+        AgentError
+            If there is an error in the agent calculation
         """
         while True:
-            msg = rcrs_encoding_utils.read_msg(self.socket)
-            self.agent_message_received(msg)
+            try:
+                msg = rcrs_encoding_utils.read_msg(self.socket)
+            except Exception as e:
+                raise ServerError(f"Error reading from socket: {e}") from e
+            try:
+                self.agent_message_received(msg)
+            except Exception as e:
+                raise AgentError(f"Error agent calculation: {e}") from e
 
     def message_received(self, agent_message_received: Callable) -> None:
         self.agent_message_received = agent_message_received

adf_core_python/core/launcher/connect/error/agent_error.py

@@ -0,0 +1,2 @@
+class AgentError(Exception):
+    pass

adf_core_python/core/launcher/connect/error/server_error.py

@@ -0,0 +1,2 @@
+class ServerError(Exception):
+    pass