Add PyTorch-based distributed remote execution with CPU communication groups

.gitignore

@@ -2,4 +2,10 @@
 .pytets*
 *.egg-info
 .python-version/
-.vscode/
+.vscode/
+__pycache__/
+*.pyc
+*.pyo
+.pytest_cache/
+dist/
+build/

README.md

@@ -4,10 +4,12 @@
 [![Tests](https://github.com/wZuck/ucuu/actions/workflows/python-app.yml/badge.svg)](https://github.com/wZuck/ucuu/actions/workflows/python-app.yml) [![GitHub Pages](https://github.com/wZuck/ucuu/actions/workflows/gh-pages.yml/badge.svg?branch=master)](https://github.com/wZuck/ucuu/actions/workflows/gh-pages.yml) [![PyPI Package](https://github.com/wZuck/ucuu/actions/workflows/publish.yml/badge.svg)](https://github.com/wZuck/ucuu/actions/workflows/publish.yml)
 > **⚠️ Important:**  
 > The majority of the code in this repository is generated using AI coding tools such as GitHub Copilot (GPT-4o) and TRAE (Doubao 1.5 Pro).  
+> 
+> **Distributed features contributor:** GitHub Copilot (Claude 3.7 Sonnet)  
 
 ## 1. Brief Introduction
 
-**UCUU** is a Python utility library for function wrapping and proxying. It supports adding proxy logic to functions via decorators or patching, making it easy to extend, debug, and test.
+**UCUU** is a Python utility library for function wrapping, proxying, and distributed computing. It supports adding proxy logic to functions via decorators or patching, and provides distributed communication capabilities for remote execution across multiple nodes using PyTorch.
 
 ---
 
@@ -21,12 +23,20 @@ You can install UCUU from PyPI:
 pip install ucuu
 ```
 
+For distributed computing features with PyTorch support:
+
+```bash
+pip install ucuu[distributed]
+```
+
 Or, clone this repository and install locally:
 
 ```bash
 git clone https://github.com/wZuck/ucuu.git
 cd ucuu
 pip install .
+# or with distributed features
+pip install .[distributed]
 ```
 
 ### 2.2 Usage
@@ -74,22 +84,165 @@ def print_ucuu_hello(ending_words=None, *args, **kwargs):
         print(f"Hello, {ending_words}")
 ```
 
+#### 2.2.4 Remote execution with CPU communication groups
+
+> Example: Set up distributed communication and execute functions remotely across nodes.  
+> **Effect:** Functions can be executed on remote peers with automatic tensor device management.
+
+```python
+import torch
+from ucuu.decorator import ucuu
+from ucuu.distributed import initialize_cpu_group
+
+# Initialize CPU communication group on each node
+# Node A (rank 0):
+comm_group = initialize_cpu_group(
+    backend="gloo",
+    init_method="tcp://master_node:29500",
+    world_size=2,
+    rank=0
+)
+
+# Node B (rank 1):
+comm_group = initialize_cpu_group(
+    backend="gloo",
+    init_method="tcp://master_node:29500",
+    world_size=2,
+    rank=1
+)
+
+# Use remote decorator to execute on peer
+# peer_rank is optional - if not specified, uses current rank
+# In typical scenarios, both peers have matching ranks
+@ucuu("package_utils.print_ucuu_hello", remote=True)
+def compute_on_peer(x):
+    """This function will execute on the peer node with the same rank"""
+    return x * 2
+
+# Tensors are automatically moved to CPU for communication
+# and restored to original device after execution
+gpu_tensor = torch.tensor([1.0, 2.0, 3.0]).cuda()
+result = compute_on_peer(gpu_tensor)  # Executed on peer, result back on GPU
+```
+
+#### 2.2.5 Custom preprocessing and postprocessing for remote execution
+
+> Example: Apply custom transformations to inputs and outputs during remote execution.  
+> **Effect:** Allows flexible data handling for distributed computing scenarios.
+
+```python
+from ucuu.decorator import ucuu
+
+def custom_preprocess(input_dict):
+    """Preprocess inputs before sending to remote peer"""
+    if 'x' in input_dict:
+        # Normalize input
+        input_dict['x'] = input_dict['x'] / 255.0
+    return input_dict
+
+def custom_postprocess(output):
+    """Postprocess output received from remote peer"""
+    # Scale output back
+    return output * 255.0
+
+@ucuu(
+    "package_utils.print_ucuu_hello",
+    remote=True,
+    custom_preprocess=custom_preprocess,
+    custom_postprocess=custom_postprocess
+)
+def process_data(x):
+    return x * 2
+```
+
+---
+
+## 3. Distributed Communication Features 🌐
+
+### 3.1 CPU Communication Groups
+
+UCUU provides CPU-based communication groups for distributed computing across multiple nodes using PyTorch's distributed primitives.
+
+#### Key Features:
+- **Cross-node communication**: Establish communication channels between processes on different nodes
+- **CPU-optimized**: Uses `gloo` backend for efficient CPU-based operations
+- **Automatic device management**: Tensors are automatically moved to CPU for communication
+- **Flexible initialization**: Supports environment variables or explicit configuration
+
+#### Basic Usage:
+
+```python
+from ucuu.distributed import initialize_cpu_group, CPUCommunicationGroup
+
+# Method 1: Using environment variables
+# Set MASTER_ADDR, MASTER_PORT, WORLD_SIZE, and RANK
+comm_group = initialize_cpu_group()
+
+# Method 2: Explicit configuration
+comm_group = initialize_cpu_group(
+    backend="gloo",
+    init_method="tcp://192.168.1.100:29500",
+    world_size=4,
+    rank=0
+)
+
+# Create peer-to-peer groups
+peer_group = comm_group.create_peer_group(ranks=[0, 1])
+
+# Send/receive tensors
+if comm_group.get_rank() == 0:
+    tensor = torch.tensor([1.0, 2.0, 3.0])
+    comm_group.send_tensor(tensor, dst=1)
+else:
+    tensor = torch.zeros(3)
+    comm_group.recv_tensor(tensor, src=0)
+
+# Synchronization
+comm_group.barrier()
+
+# Cleanup when done
+comm_group.cleanup()
+```
+
+### 3.2 Remote Decorator
+
+The `@ucuu` decorator supports a `remote` attribute for executing functions on remote peers:
+
+#### Parameters:
+- `remote` (bool): Enable remote execution (default: False)
+- `peer_rank` (int, optional): Rank of the peer to execute on. If not specified, uses the current rank (suitable for scenarios where both peers have matching ranks)
+- `custom_preprocess` (Callable): Function to preprocess inputs before sending
+- `custom_postprocess` (Callable): Function to postprocess outputs after receiving
+
+#### Device Management:
+- Input tensors are automatically converted to CPU before transmission
+- Output tensors are automatically converted back to the original device
+- Custom preprocessing/postprocessing can be applied at each stage
+
+#### Note:
+In typical distributed scenarios, peer nodes have matching ranks (e.g., rank 0 on node A communicates with rank 0 on node B). Therefore, `peer_rank` can usually be omitted and will default to the current rank.
+
 ---
 
-## 3. Demos in Testcases 🧪
+## 4. Demos in Testcases 🧪
 
-- See the `tests/` directory for test cases covering decorator usage, patching, exception handling, argument binding, and more.
+- See the `tests/` directory for test cases covering:
+  - Decorator usage and patching
+  - Exception handling
+  - Argument binding
+  - Distributed communication groups
+  - Remote execution with preprocessing/postprocessing
 
 
 ---
 
-## 4. License 📄
+## 5. License 📄
 
 [License](./LICENSE)
 
 ---
 
-## 5. Contribute 🤝
+## 6. Contribute 🤝
 
 Contributions via PR or issues are welcome!  
 For suggestions or questions, please leave a message at [GitHub Issues](https://github.com/wZuck/ucuu/issues).

examples/__init__.py

@@ -0,0 +1,3 @@
+"""
+Examples demonstrating UCUU features.
+"""

examples/distributed_example.py

@@ -0,0 +1,107 @@
+"""
+Example demonstrating CPU communication group and remote execution features.
+
+This example shows how to use UCUU's distributed features for remote execution.
+Note: For actual distributed execution, you need to run this script on multiple
+nodes with proper distributed setup (MASTER_ADDR, MASTER_PORT, WORLD_SIZE, RANK).
+"""
+
+try:
+    import torch
+
+    TORCH_AVAILABLE = True
+except ImportError:
+    TORCH_AVAILABLE = False
+    print("PyTorch not available. Install with: pip install ucuu[distributed]")
+
+if TORCH_AVAILABLE:
+    from ucuu.decorator import ucuu
+
+    # Example 1: Basic remote execution
+    print("=" * 60)
+    print("Example 1: Basic Remote Execution")
+    print("=" * 60)
+
+    # For demonstration, we'll use local execution
+    # In production, initialize with proper distributed settings
+    # comm_group = initialize_cpu_group(
+    #     backend="gloo",
+    #     init_method="tcp://master_node:29500",
+    #     world_size=2,
+    #     rank=0  # or 1 for the peer node
+    # )
+
+    @ucuu("package_utils.print_ucuu_hello", remote=True)
+    def compute_remotely(x):
+        """This function will execute on the peer node"""
+        return x * 2
+
+    result = compute_remotely(5)
+    print(f"Result: {result}")
+    print()
+
+    # Example 2: Remote execution with tensors
+    print("=" * 60)
+    print("Example 2: Remote Execution with Tensors")
+    print("=" * 60)
+
+    @ucuu("package_utils.print_ucuu_hello", remote=True)
+    def process_tensor(x):
+        """Process tensor on remote peer"""
+        return x * 2 + 1
+
+    tensor_input = torch.tensor([1.0, 2.0, 3.0])
+    print(f"Input tensor: {tensor_input}")
+    result_tensor = process_tensor(tensor_input)
+    print(f"Result tensor: {result_tensor}")
+    print()
+
+    # Example 3: Remote execution with custom preprocessing
+    print("=" * 60)
+    print("Example 3: Remote Execution with Custom Preprocessing")
+    print("=" * 60)
+
+    def normalize_input(input_dict):
+        """Normalize input values"""
+        if "x" in input_dict:
+            print(f"Preprocessing: Normalizing input {input_dict['x']}")
+            input_dict["x"] = input_dict["x"] / 10.0
+        return input_dict
+
+    def scale_output(output):
+        """Scale output values"""
+        print(f"Postprocessing: Scaling output {output}")
+        return output * 10.0
+
+    @ucuu(
+        "package_utils.print_ucuu_hello",
+        remote=True,
+        custom_preprocess=normalize_input,
+        custom_postprocess=scale_output,
+    )
+    def process_with_transforms(x):
+        """Process with preprocessing and postprocessing"""
+        return x * 2
+
+    result = process_with_transforms(50)
+    print(f"Final result: {result}")
+    print()
+
+    # Example 4: Normal (non-remote) execution
+    print("=" * 60)
+    print("Example 4: Normal Execution (remote=False)")
+    print("=" * 60)
+
+    @ucuu("package_utils.print_ucuu_hello", remote=False, ending_words="local mode")
+    def local_function(x):
+        """This executes locally"""
+        print(f"Executing locally with input: {x}")
+        return x * 3
+
+    result = local_function(7)
+    print(f"Result: {result}")
+    print()
+
+else:
+    print("\nTo use distributed features, install PyTorch:")
+    print("pip install ucuu[distributed]")

pyproject.toml

@@ -39,4 +39,7 @@ docs = [
 publish = [
     "build>=1.0.0",
     "twine>=4.0.0"
+]
+distributed = [
+    "torch>=1.10.0"
 ]