shared-tensor 0.1.0__py3-none-any.whl

shared_tensor/jsonrpc.py

@@ -0,0 +1,163 @@
+"""
+JSON-RPC 2.0 implementation for shared tensor communication
+
+Implements JSON-RPC 2.0 specification: https://www.jsonrpc.org/specification
+"""
+
+import json
+import uuid
+from typing import Any, Dict, Optional, Union
+from dataclasses import dataclass, asdict
+
+
+@dataclass
+class JsonRpcRequest:
+    """JSON-RPC 2.0 Request object"""
+    method: str
+    params: Optional[Union[Dict[str, Any], list]] = None
+    id: Optional[Union[str, int]] = None
+    jsonrpc: str = "2.0"
+    
+    def __post_init__(self):
+        if self.id is None:
+            self.id = str(uuid.uuid4())
+    
+    def to_dict(self) -> Dict[str, Any]:
+        """Convert to dictionary for JSON serialization"""
+        result = asdict(self)
+        # Remove None params to keep the request clean
+        if result['params'] is None:
+            del result['params']
+        return result
+    
+    def to_json(self) -> str:
+        """Convert to JSON string"""
+        return json.dumps(self.to_dict())
+
+
+@dataclass
+class JsonRpcResponse:
+    """JSON-RPC 2.0 Response object"""
+    id: Optional[Union[str, int]]
+    result: Optional[Any] = None
+    error: Optional[Dict[str, Any]] = None
+    jsonrpc: str = "2.0"
+    
+    def to_dict(self) -> Dict[str, Any]:
+        """Convert to dictionary for JSON serialization"""
+        result = asdict(self)
+        # Only include result OR error, not both
+        if self.error is not None:
+            result.pop('result', None)
+        else:
+            result.pop('error', None)
+        return result
+    
+    def to_json(self) -> str:
+        """Convert to JSON string"""
+        return json.dumps(self.to_dict())
+
+
+@dataclass
+class JsonRpcError:
+    """JSON-RPC 2.0 Error object"""
+    code: int
+    message: str
+    data: Optional[Any] = None
+    
+    def to_dict(self) -> Dict[str, Any]:
+        """Convert to dictionary"""
+        result = asdict(self)
+        if result['data'] is None:
+            del result['data']
+        return result
+
+
+# Standard JSON-RPC error codes
+class JsonRpcErrorCodes:
+    PARSE_ERROR = -32700
+    INVALID_REQUEST = -32600
+    METHOD_NOT_FOUND = -32601
+    INVALID_PARAMS = -32602
+    INTERNAL_ERROR = -32603
+    
+    # Custom error codes for shared tensor
+    FUNCTION_IMPORT_ERROR = -32001
+    FUNCTION_EXECUTION_ERROR = -32002
+    SERIALIZATION_ERROR = -32003
+
+
+def create_success_response(request_id: Optional[Union[str, int]], result: Any) -> JsonRpcResponse:
+    """Create a successful JSON-RPC response"""
+    return JsonRpcResponse(id=request_id, result=result)
+
+
+def create_error_response(
+    request_id: Optional[Union[str, int]], 
+    code: int, 
+    message: str, 
+    data: Optional[Any] = None
+) -> JsonRpcResponse:
+    """Create an error JSON-RPC response"""
+    error = JsonRpcError(code=code, message=message, data=data)
+    return JsonRpcResponse(id=request_id, error=error.to_dict())
+
+
+def parse_request(json_str: str) -> JsonRpcRequest:
+    """Parse JSON string into JsonRpcRequest object"""
+    try:
+        data = json.loads(json_str)
+    except json.JSONDecodeError as e:
+        raise ValueError(f"Invalid JSON: {e}")
+    
+    # Validate required fields
+    if not isinstance(data, dict):
+        raise ValueError("Request must be a JSON object")
+    
+    if data.get('jsonrpc') != '2.0':
+        raise ValueError("Invalid jsonrpc version, must be '2.0'")
+    
+    if 'method' not in data:
+        raise ValueError("Missing required field 'method'")
+    
+    return JsonRpcRequest(
+        method=data['method'],
+        params=data.get('params'),
+        id=data.get('id'),
+        jsonrpc=data['jsonrpc']
+    )
+
+
+def parse_response(json_str: str) -> JsonRpcResponse:
+    """Parse JSON string into JsonRpcResponse object"""
+    try:
+        data = json.loads(json_str)
+    except json.JSONDecodeError as e:
+        raise ValueError(f"Invalid JSON: {e}")
+    
+    # Validate required fields
+    if not isinstance(data, dict):
+        raise ValueError("Response must be a JSON object")
+    
+    if data.get('jsonrpc') != '2.0':
+        raise ValueError("Invalid jsonrpc version, must be '2.0'")
+    
+    if 'id' not in data:
+        raise ValueError("Missing required field 'id'")
+    
+    # Must have either result or error
+    has_result = 'result' in data
+    has_error = 'error' in data
+    
+    if not has_result and not has_error:
+        raise ValueError("Response must have either 'result' or 'error'")
+    
+    if has_result and has_error:
+        raise ValueError("Response cannot have both 'result' and 'error'")
+    
+    return JsonRpcResponse(
+        id=data['id'],
+        result=data.get('result'),
+        error=data.get('error'),
+        jsonrpc=data['jsonrpc']
+    )

shared_tensor/provider.py

@@ -0,0 +1,155 @@
+"""
+Shared Tensor Provider
+
+This module provides a provider for sharing functions across processes using JSON-RPC.
+"""
+
+import inspect
+import os
+import logging
+from functools import wraps
+from typing import Any, Dict, Callable, Optional
+from shared_tensor.client import SharedTensorClient
+from shared_tensor.errors import SharedTensorProviderError
+
+
+__all__ = ["SharedTensorProvider"]
+
+logger = logging.getLogger(__name__)
+global_rank = int(os.getenv("RANK", 0))
+
+
+class SharedTensorProvider:
+
+    def __init__(self, server_port: int = 2537 + global_rank, verbose_debug: bool = False):
+        self.server_port: int = server_port
+        self.server_mode = os.getenv("__SHARED_TENSOR_SERVER_MODE__", "false")
+        self.verbose_debug = verbose_debug
+        logger.debug(f"SharedTensorProvider initialized with server port {self.server_port}, server mode {self.server_mode}, and verbose debug {self.verbose_debug}")
+        self._registered_functions: Dict[str, Dict[str, Any]] = {}
+        self._client = None
+
+    def _get_function_path(self, func: Callable) -> str:
+        """Get the importable path of a function in format 'module.submodule:function_name'"""
+        module = inspect.getmodule(func)
+        
+        if module is None:
+            raise SharedTensorProviderError(f"Failed to get full qualified name for function {func.__name__}, function module is missing")
+        
+        module_name = module.__name__
+        
+        if module_name == "__main__":
+            if hasattr(module, '__file__') and module.__file__:
+                file_path = module.__file__
+                file_name = os.path.splitext(os.path.basename(file_path))[0]
+                # For test files, we need to include the full path
+                if 'tests' in file_path:
+                    # Extract the module path from the file path
+                    path_parts = file_path.split(os.sep)
+                    if 'tests' in path_parts:
+                        test_idx = path_parts.index('tests')
+                        module_parts = path_parts[test_idx:]
+                        # Remove .py extension from last part
+                        module_parts[-1] = os.path.splitext(module_parts[-1])[0]
+                        module_name = '.'.join(module_parts)
+                    else:
+                        module_name = file_name
+                else:
+                    module_name = file_name
+            else:
+                raise SharedTensorProviderError(f"Failed to get full qualified name for function {func.__name__}, function module file path is empty")
+        
+        if hasattr(func, '__qualname__'):
+            qualname = func.__qualname__
+            # Only reject true nested functions with <locals>, but allow test functions
+            # Test functions might have qualname like "TestClass.test_method.<locals>.test_function"
+            # but we can still use them by extracting the actual function name
+            if '<locals>' in qualname:
+                # For nested functions, try to extract the innermost function name
+                # This allows test functions defined inside test methods to work
+                func_path = func.__name__
+                logger.warning(f"Function {func.__name__} appears to be nested (qualname: {qualname}), using function name only")
+            else:
+                func_path = qualname
+        else:
+            func_path = func.__name__
+        
+        return f"{module_name}:{func_path}"
+
+    def share(self, name: Optional[str] = None, singleton: bool = True, singleton_key_formatter: Optional[str] = None):
+        """Decorator to register a function for remote sharing
+        
+        Args:
+            name: Optional custom name for the function
+            singleton: Whether to use a singleton instance of the function result
+            singleton_key_formatter: Formatter for cached results
+
+        Returns:
+            Decorator function that registers the function for remote sharing
+        """
+        def decorator(func: Callable):
+            func_name = name or func.__name__
+
+            if self.server_mode == "true":
+                logger.debug(f"Server mode is true, returning function {func_name} without registering")
+                return func
+
+            logger.debug(f"Server mode is false, registering function {func_name}")
+            function_path = self._get_function_path(func)
+            
+            logger.debug(f"Function {func_name} registered with function path {function_path}")
+            options = {
+                'name': func_name,
+                'singleton': singleton,
+                'singleton_key_formatter': singleton_key_formatter,
+            }
+            function_info = {
+                'name': func_name,
+                'function_path': function_path,
+                'options': options,
+            }
+            self._registered_functions[func_name] = function_info
+            
+            @wraps(func)
+            def wrapper(*args, **kwargs):
+                return self._execute_remote_function(func_name, args, kwargs, options)
+            
+            return wrapper
+        return decorator
+
+    def _get_client(self) -> SharedTensorClient:
+        """Get or create JSON-RPC client"""
+        if self._client is None:
+            logger.debug(f"Creating new JSON-RPC client with server port {self.server_port}")
+            self._client = SharedTensorClient(self.server_port, verbose_debug=self.verbose_debug)
+            logger.debug(f"JSON-RPC client created with server port {self.server_port}")
+        return self._client
+
+    def _execute_remote_function(self, func_name: str, args: tuple, kwargs: dict, options: dict) -> Any:
+        """Execute function remotely using JSON-RPC client"""
+        try:
+            if self.verbose_debug:
+                logger.debug(f"Executing remote function {func_name} with args {args} and kwargs {kwargs}")
+            else:
+                logger.debug(f"Executing remote function {func_name}")
+
+            if func_name not in self._registered_functions:
+                raise SharedTensorProviderError(f"Function {func_name} not registered")
+            
+            function_info = self._registered_functions[func_name]
+            function_path = function_info['function_path']
+            client = self._get_client()
+            logger.debug(f"Executing remote function {func_name} with function path {function_path} and options {options}")
+            return client.execute_function(function_path, args, kwargs, options)
+                
+        except Exception as e:
+            logger.warning(f"Failed to execute remote function {func_name}: {str(e)}")
+            raise SharedTensorProviderError(f"Failed to execute remote function {func_name}: {str(e)}")
+    
+    def close(self):
+        """Close the provider and its client connection"""
+        if self._client:
+            logger.debug(f"Closing JSON-RPC client with server port {self.server_port}")
+            self._client.close()
+            logger.debug(f"JSON-RPC client closed with server port {self.server_port}")
+            self._client = None