deepspeed_setup/dsync/dsync.py

#!/usr/bin/env python3

import sys
import json
import subprocess
from pathlib import Path
import concurrent.futures

# ANSI color codes
class Colors:
    RED = '\033[31m'
    GREEN = '\033[32m'
    ORANGE = '\033[33m'
    YELLOW = '\033[1;33m'
    BLUE = '\033[34m'
    MAGENTA = '\033[35m'
    CYAN = '\033[36m'
    WHITE = '\033[37m'
    RESET = '\033[0m'
    BOLD = '\033[1m'

def colored(text, color):
    """Apply color to text"""
    return f"{color}{text}{Colors.RESET}"

def find_file_up(filename: str, start_dir: Path = Path.cwd()) -> Path:
    current_dir = start_dir
    while current_dir != current_dir.parent:
        candidate = current_dir / filename
        if candidate.exists():
            return candidate
        current_dir = current_dir.parent
    print(colored(f"Error: {filename} not found in this directory or any parent.", Colors.RED))
    sys.exit(1)

def load_config():
    config_path = find_file_up(".dsyncconfig")
    with open(config_path) as f:
        return json.load(f), config_path.parent

def check_node_status(node, project_dir, ignore_file, master_dir):
    """Check status for a single node"""
    remote = f"{node}:{master_dir}"
    
    cmd = [
        "rsync", "-avz", "--dry-run", "--itemize-changes", "--delete",
        f"--exclude-from={ignore_file}",
        "-e", "ssh -o StrictHostKeyChecking=no -o ConnectTimeout=5",
        f"{project_dir}/", remote
    ]
    
    try:
        result = subprocess.run(cmd, capture_output=True, text=True, timeout=30)
        
        if result.returncode == 0:
            # Parse rsync output
            lines = result.stdout.strip().split('\n')
            changes = []
            
            for line in lines:
                if line.startswith('>f') or line.startswith('>d') or line.startswith('*deleting'):
                    changes.append(line)
                elif line.startswith('<f') or line.startswith('<d'):
                    changes.append(line)
                elif line.startswith('cd') or line.startswith('cf'):
                    changes.append(line)
            
            return {
                'node': node,
                'status': 'ok',
                'changes': changes,
                'total_changes': len(changes)
            }
        else:
            return {
                'node': node,
                'status': 'error',
                'error': result.stderr.strip(),
                'changes': [],
                'total_changes': 0
            }
    except subprocess.TimeoutExpired:
        return {
            'node': node,
            'status': 'timeout',
            'error': 'Connection timeout',
            'changes': [],
            'total_changes': 0
        }
    except Exception as e:
        return {
            'node': node,
            'status': 'error',
            'error': str(e),
            'changes': [],
            'total_changes': 0
        }

def status():
    """Show sync status for all nodes"""
    config, project_dir = load_config()
    ignore_file = project_dir / config.get("ignore_file", ".dsyncignore")
    master_dir = config.get("master_dir", str(project_dir))
    
    # Ensure ignore file exists
    if not ignore_file.exists():
        ignore_file.write_text("# Add patterns to ignore during sync\n")
    
    nodes = config.get("nodes", [])
    if not nodes:
        print(colored("No nodes configured", Colors.YELLOW))
        return
    
    print(f"Checking status for {len(nodes)} node(s)...")
    print("-" * 50)
    
    # Check all nodes in parallel for speed
    with concurrent.futures.ThreadPoolExecutor(max_workers=min(len(nodes), 10)) as executor:
        futures = {
            executor.submit(check_node_status, node, project_dir, ignore_file, master_dir): node 
            for node in nodes
        }
        
        results = []
        for future in concurrent.futures.as_completed(futures):
            result = future.result()
            results.append(result)
    
    # Sort results by node name for consistent output
    results.sort(key=lambda x: x['node'])
    
    # Display results
    total_changes = 0
    for result in results:
        node = result['node']
        status = result['status']
        changes = result['changes']
        change_count = result['total_changes']
        
        if status == 'ok':
            if change_count == 0:
                print(colored(f"[OK] {node}: Up to date", Colors.GREEN))
            else:
                print(colored(f"[!] {node}: {change_count} change(s) pending", Colors.YELLOW))
                total_changes += change_count
                
                # Show first few changes
                for i, change in enumerate(changes[:5]):
                    action = parse_rsync_action(change)
                    print(f"  {action}")
                
                if len(changes) > 5:
                    print(f"  ... and {len(changes) - 5} more change(s)")
        else:
            print(colored(f"[X] {node}: {result.get('error', 'Unknown error')}", Colors.RED))
        
        print()
    
    # Summary
    active_nodes = sum(1 for r in results if r['status'] == 'ok')
    error_nodes = sum(1 for r in results if r['status'] != 'ok')
    
    print("-" * 50)
    if total_changes > 0:
        print(colored(f"Summary: {active_nodes} active, {error_nodes} error(s), {total_changes} total changes", Colors.YELLOW))
        print(colored("Run 'dsync sync' to apply changes", Colors.YELLOW))
    else:
        print(colored(f"Summary: {active_nodes} active, {error_nodes} error(s), {total_changes} total changes", Colors.GREEN))
        if error_nodes == 0:
            print(colored("All nodes are up to date!", Colors.GREEN))

def parse_rsync_action(line):
    """Parse rsync itemize-changes output into human-readable format"""
    if line.startswith('*deleting'):
        filename = line.split(' ', 1)[1].strip()
        return f"DELETE: {filename}"
    
    if len(line) < 11:
        return line.strip()
    
    # Parse rsync itemize format: YXcstpoguax filename
    changes = line[0:11]
    filename = line[11:].strip()
    
    if not filename:
        return line.strip()
    
    # Check file type and direction
    file_type = changes[0]
    if file_type == '<':
        direction = "SEND"
    elif file_type == '>':
        direction = "RECV"
    elif file_type == 'c':
        direction = "CREATE"
    elif file_type == 'h':
        direction = "HARDLINK"
    elif file_type == '.':
        direction = "UPDATE"
    else:
        direction = "CHANGE"
    
    # Parse item type
    change_type = changes[1]
    if change_type == 'f':
        item_type = "file"
    elif change_type == 'd':
        item_type = "directory"
    elif change_type == 'L':
        item_type = "symlink"
    elif change_type == 'D':
        item_type = "device"
    elif change_type == 'S':
        item_type = "special"
    else:
        item_type = "item"
    
    # Parse specific changes
    change_details = changes[2:11]
    details = []
    
    if change_details[0] == '+':
        return f"NEW {item_type.upper()}: {filename}"
    elif change_details[0] == 'c':
        details.append("content")
    
    if change_details[1] == 's':
        details.append("size")
    if change_details[2] == 't':
        details.append("timestamp")
    if change_details[3] == 'p':
        details.append("permissions")
    if change_details[4] == 'o':
        details.append("owner")
    if change_details[5] == 'g':
        details.append("group")
    if change_details[6] == 'u':
        details.append("unused")
    if change_details[7] == 'a':
        details.append("ACL")
    if change_details[8] == 'x':
        details.append("extended attributes")
    
    # Build result with direction and item type
    if details:
        detail_str = ", ".join(details)
        return f"{direction} {item_type.upper()}: {filename} ({detail_str})"
    else:
        return f"{direction} {item_type.upper()}: {filename}"

def sync():
    config, project_dir = load_config()
    ignore_file = project_dir / config.get("ignore_file", ".dsyncignore")
    master_dir = config.get("master_dir", str(project_dir))
    
    # Ensure ignore file exists
    if not ignore_file.exists():
        ignore_file.write_text("# Add patterns to ignore during sync\n")
    
    success_count = 0
    failed_nodes = []
    
    for node in config["nodes"]:
        remote = f"{node}:{master_dir}"
        print(f"Syncing to {remote}...")
        
        cmd = [
            "rsync", "-avz", "--delete",
            f"--exclude-from={ignore_file}",
            "-e", "ssh -o StrictHostKeyChecking=no -o ConnectTimeout=10",
            f"{project_dir}/", remote
        ]
        
        try:
            subprocess.run(cmd, check=True, capture_output=True, text=True)
            print(colored(f"[OK] Successfully synced to {node}", Colors.GREEN))
            success_count += 1
        except subprocess.CalledProcessError as e:
            print(colored(f"[X] Failed to sync to {node}: {e}", Colors.RED))
            failed_nodes.append(node)
    
    print()
    if success_count == len(config['nodes']):
        print(colored(f"Sync completed: {success_count}/{len(config['nodes'])} nodes successful", Colors.GREEN))
    else:
        print(colored(f"Sync completed: {success_count}/{len(config['nodes'])} nodes successful", Colors.YELLOW))
    
    if failed_nodes:
        print(colored(f"Failed nodes: {', '.join(failed_nodes)}", Colors.RED))

def init(master_dir: Path = Path.cwd()):
    config_file = master_dir / ".dsyncconfig"
    ignore_file = master_dir / ".dsyncignore"
    
    if config_file.exists():
        print(colored(f"dsync already initialized in {master_dir}", Colors.YELLOW))
        return
    
    config = {
        "master_dir": str(master_dir),
        "nodes": [],
        "ignore_file": ".dsyncignore"
    }
    
    config_file.write_text(json.dumps(config, indent=2))
    
    if not ignore_file.exists():
        ignore_content = """# dsync ignore patterns
# Add patterns to ignore during sync
.git/
.gitignore
.dsyncconfig
.dsyncignore
__pycache__/
*.pyc
*.pyo
*.pyd
.Python
build/
develop-eggs/
dist/
downloads/
eggs/
.eggs/
lib/
lib64/
parts/
sdist/
var/
wheels/
*.egg-info/
.installed.cfg
*.egg
MANIFEST
.env
.venv
env/
venv/
ENV/
env.bak/
venv.bak/
.vscode/
.idea/
*.log
node_modules/
.npm
.cache/
"""
        ignore_file.write_text(ignore_content)
    
    print(colored(f"Initialized dsync config in {master_dir}", Colors.GREEN))

def add_nodes(*nodes):
    if not nodes:
        print(colored("No nodes provided", Colors.YELLOW))
        return
    
    config, project_dir = load_config()
    added_nodes = []
    
    for node in nodes:
        if node not in config["nodes"]:
            config["nodes"].append(node)
            added_nodes.append(node)
        else:
            print(colored(f"Node {node} already exists", Colors.YELLOW))
    
    if added_nodes:
        (project_dir / ".dsyncconfig").write_text(json.dumps(config, indent=2))
        print(colored(f"Added nodes: {', '.join(added_nodes)}", Colors.GREEN))
    else:
        print(colored("No new nodes added", Colors.YELLOW))

def del_nodes(*nodes):
    if not nodes:
        print(colored("No nodes provided", Colors.YELLOW))
        return
    
    config, project_dir = load_config()
    removed_nodes = []
    
    for node in nodes:
        if node in config["nodes"]:
            config["nodes"].remove(node)
            removed_nodes.append(node)
        else:
            print(colored(f"Node {node} not found", Colors.YELLOW))
    
    if removed_nodes:
        (project_dir / ".dsyncconfig").write_text(json.dumps(config, indent=2))
        print(colored(f"Removed nodes: {', '.join(removed_nodes)}", Colors.GREEN))
    else:
        print(colored("No nodes removed", Colors.YELLOW))

def list_nodes():
    config, _ = load_config()
    nodes = config.get("nodes", [])
    if nodes:
        print(colored("Configured nodes:", Colors.CYAN))
        for i, node in enumerate(nodes, 1):
            print(f"  {i}. {node}")
    else:
        print(colored("No nodes configured", Colors.YELLOW))

def show_help():
    help_text = """
dsync.py - Directory Synchronization Tool

Usage:
  dsync.py init                           Initialize dsync in current directory
  dsync.py add-node <node1> [node2...]   Add one or more nodes
  dsync.py del-node <node1> [node2...]   Remove one or more nodes
  dsync.py list-nodes                     List all configured nodes
  dsync.py status                         Show sync status for all nodes
  dsync.py sync                           Sync to all configured nodes
  dsync.py help                           Show this help message

Examples:
  dsync.py init
  dsync.py add-node user@server1.com user@server2.com
  dsync.py del-node user@server1.com
  dsync.py status
  dsync.py sync

Notes:
  - Status command shows pending changes without syncing
  - Sync uses rsync with --delete flag to handle file deletions
  - SSH keys should be configured for passwordless access
  - Edit .dsyncignore to exclude files from sync
"""
    print(help_text)

if __name__ == "__main__":
    if len(sys.argv) < 2:
        show_help()
        sys.exit(1)

    cmd = sys.argv[1]
    
    if cmd == "init":
        init()
    elif cmd == "add-node":
        if len(sys.argv) < 3:
            print(colored("Usage: dsync.py add-node <node1> [node2...]", Colors.RED))
            sys.exit(1)
        add_nodes(*sys.argv[2:])
    elif cmd == "del-node":
        if len(sys.argv) < 3:
            print(colored("Usage: dsync.py del-node <node1> [node2...]", Colors.RED))
            sys.exit(1)
        del_nodes(*sys.argv[2:])
    elif cmd == "list-nodes":
        list_nodes()
    elif cmd == "status":
        status()
    elif cmd == "sync":
        sync()
    elif cmd == "help":
        show_help()
    else:
        print(colored(f"Unknown command: {cmd}", Colors.RED))
        show_help()
        sys.exit(1)