Add Linting

Author: sbak5

src/nvidia_resiliency_ext/attribution/mcp_integration/mcp_client.py

@@ -16,9 +16,7 @@
 from mcp.client.session import ClientSession
 from mcp.client.stdio import StdioServerParameters, stdio_client
 
-from nvidia_resiliency_ext.attribution.mcp_integration.registry import (
-    deserialize_result,
-)
+from nvidia_resiliency_ext.attribution.mcp_integration.registry import deserialize_result
 
 logger = logging.getLogger(__name__)
 
@@ -141,7 +139,7 @@ async def run_module(self, module_name: str, **kwargs) -> Dict[str, Any]:
         result_str = await self.call_tool(module_name, arguments)
         return deserialize_result(result_str)
 
-    async def get_result(self,  result_id: str) -> Dict[str, Any]:
+    async def get_result(self, result_id: str) -> Dict[str, Any]:
         """
         Retrieve a cached result by ID.
 

src/nvidia_resiliency_ext/attribution/mcp_integration/mcp_server.py

@@ -9,7 +9,6 @@
 import asyncio
 import json
 import logging
-import uuid
 from typing import Any, Dict, List, Optional
 
 from mcp.server import Server
@@ -199,7 +198,7 @@ async def _handle_module_execution(
         """Execute a single attribution module."""
         # Apply default values from input schema
         arguments_with_defaults = self.registry.apply_defaults(module_name, arguments)
-        
+
         # Get or create module instance
         if module_name not in self.module_instances:
             # Convert arguments to argparse.Namespace
@@ -235,12 +234,11 @@ async def _handle_module_execution(
 
         return [TextContent(type="text", text=serialize_result(response))]
 
-
     async def run(self):
         """Run the MCP server."""
         import os
 
-        logger.info(f"Starting NVRX Attribution MCP Server")
+        logger.info("Starting NVRX Attribution MCP Server")
         logger.info(f"Registered modules: {self.registry.list_modules()}, pid: {os.getpid()}")
 
         async with stdio_server() as (read_stream, write_stream):
@@ -250,4 +248,4 @@ async def run(self):
 
     def run_sync(self):
         """Run the server synchronously."""
-        asyncio.run(self.run())
+        asyncio.run(self.run())

src/nvidia_resiliency_ext/attribution/mcp_integration/module_definitions.py

@@ -9,6 +9,7 @@
 from nvidia_resiliency_ext.attribution.mcp_integration.registry import global_registry
 from nvidia_resiliency_ext.attribution.trace_analyzer.fr_attribution import CollectiveAnalyzer
 
+
 def register_all_modules():
     """Register all NVRX attribution modules with the global registry."""
 
@@ -71,6 +72,7 @@ def register_all_modules():
         dependencies=[],
     )
 
+
 def create_args_from_dict(module_name: str, config: dict) -> argparse.Namespace:
     """
     Create an argparse.Namespace from a configuration dictionary.

src/nvidia_resiliency_ext/attribution/mcp_integration/registry.py

@@ -4,11 +4,10 @@
 """
 
 import hashlib
-import inspect
 import json
 import logging
 from dataclasses import asdict, dataclass, is_dataclass
-from typing import Any, Callable, Dict, List, Optional, Type
+from typing import Any, Dict, List, Optional, Type
 
 from nvidia_resiliency_ext.attribution.base import NVRxAttribution
 
@@ -110,20 +109,20 @@ def apply_defaults(self, module_name: str, arguments: Dict[str, Any]) -> Dict[st
         metadata = self._modules.get(module_name)
         if not metadata:
             return arguments
-        
+
         # Create a copy to avoid modifying the original
         result = dict(arguments)
-        
+
         # Get the properties from the input schema
         input_schema = metadata.input_schema
         properties = input_schema.get("properties", {})
-        
+
         # Apply defaults for missing arguments
         for param_name, param_schema in properties.items():
             if param_name not in result and "default" in param_schema:
                 result[param_name] = param_schema["default"]
                 logger.debug(f"Applied default for {param_name}: {param_schema['default']}")
-        
+
         return result
 
     def cache_result(self, module_name: str, arguments: Dict[str, Any], result: Any):