arbor-ai 0.1.5__py3-none-any.whl → 0.1.6__py3-none-any.whl

arbor/server/services/comms/comms.py

@@ -0,0 +1,226 @@
+import os
+import queue
+import socket
+import threading
+import time
+
+import zmq
+
+
+class ArborServerCommsHandler:
+    """Handles socket communication between manager and training process"""
+
+    def __init__(self, host="localhost"):
+        self.host = host
+        self.context = zmq.Context()
+
+        # Command socket (REQ/REP pattern)
+        self.command_socket = self.context.socket(zmq.REQ)
+        self.command_port = self.command_socket.bind_to_random_port(f"tcp://{host}")
+
+        # Status socket (PUB/SUB pattern)
+        self.status_socket = self.context.socket(zmq.SUB)
+        self.status_port = self.status_socket.bind_to_random_port(f"tcp://{host}")
+        self.status_socket.setsockopt_string(zmq.SUBSCRIBE, "")
+
+        # Data socket (PUB/SUB pattern)
+        self.data_socket = self.context.socket(zmq.PUB)
+        self.data_port = self.data_socket.bind_to_random_port(f"tcp://{host}")
+
+        self.broadcast_socket = self.context.socket(zmq.PUB)
+        self.broadcast_port = self.broadcast_socket.bind_to_random_port(f"tcp://{host}")
+
+        self.handshake_socket = self.context.socket(zmq.REP)
+        self.handshake_port = self.handshake_socket.bind_to_random_port(f"tcp://{host}")
+
+    def send_command(self, command):
+        self.command_socket.send_json(command)
+        return self.command_socket.recv_json()  # Wait for acknowledgment
+
+    def send_data(self, data):
+        self.data_socket.send_json(data)
+
+    def send_broadcast(self, message):
+        self.broadcast_socket.send_json(message)
+
+    def receive_status(self):
+        while True:
+            status = self.status_socket.recv_json()
+            yield status
+
+    def close(self):
+        self.command_socket.close()
+        self.status_socket.close()
+        self.data_socket.close()
+        self.broadcast_socket.close()
+        self.handshake_socket.close()
+        self.context.term()
+
+    def wait_for_clients(self, expected_count):
+        connected_clients = []
+        while len(connected_clients) < expected_count:
+            print(f"Waiting for {expected_count} clients to connect...")
+            msg = self.handshake_socket.recv_json()
+            if msg.get("type") == "hello":
+                client_id = msg.get("client_id")
+                connected_clients.append(client_id)
+                self.handshake_socket.send_json({"status": "ack"})
+            print(f"Received handshake from {client_id}")
+        print(f"All {expected_count} clients connected!")
+
+
+class ArborScriptCommsHandler:
+    def __init__(
+        self,
+        host,
+        command_port,
+        status_port,
+        data_port,
+        broadcast_port,
+        handshake_port,
+        is_main_process,
+    ):
+        self.context = zmq.Context()
+        self.is_main_process = is_main_process
+
+        # Command socket (main process only)
+        if is_main_process:
+            self.command_socket = self.context.socket(zmq.REP)
+            self.command_socket.connect(f"tcp://{host}:{command_port}")
+
+            self.status_socket = self.context.socket(zmq.PUB)
+            self.status_socket.connect(f"tcp://{host}:{status_port}")
+        else:
+            self.command_socket = None
+            self.status_socket = None
+
+        # Data socket (all processes)
+        self.data_socket = self.context.socket(zmq.SUB)
+        self.data_socket.connect(f"tcp://{host}:{data_port}")
+        self.data_socket.setsockopt_string(zmq.SUBSCRIBE, "")
+        self.data_queue = queue.Queue()
+        self._start_data_receiver()
+
+        # Broadcast socket (all processes)
+        self.broadcast_socket = self.context.socket(zmq.SUB)
+        self.broadcast_socket.connect(f"tcp://{host}:{broadcast_port}")
+        self.broadcast_socket.setsockopt_string(zmq.SUBSCRIBE, "")
+
+        # Handshake socket (all processes)
+        self.handshake_socket = self.context.socket(zmq.REQ)
+        self.handshake_socket.connect(f"tcp://{host}:{handshake_port}")
+        self._send_handshake()
+
+    def send_status(self, status):
+        if self.status_socket is not None:
+            self.status_socket.send_json(status)
+
+    def receive_command(self):
+        if self.command_socket is not None:
+            while True:
+                command = self.command_socket.recv_json()
+                # Send acknowledgment
+                self.command_socket.send_json({"status": "received"})
+                yield command
+
+    def receive_data(self):
+        return self.data_queue.get()
+
+    def _start_data_receiver(self):
+        def _receiver():
+            while True:
+                try:
+                    data = self.data_socket.recv_json()
+                    self.data_queue.put(data)
+                except Exception as e:
+                    print(f"Error receiving data: {e}")
+                    break
+
+        self.receiver_thread = threading.Thread(target=_receiver, daemon=True)
+        self.receiver_thread.start()
+
+    def is_data_queue_empty(self):
+        return self.data_queue.empty()
+
+    def get_data_queue_size(self):
+        return self.data_queue.qsize()
+
+    def receive_broadcast(self):
+        while True:
+            broadcast = self.broadcast_socket.recv_json()
+            yield broadcast
+
+    def close(self):
+        if self.command_socket is not None:
+            self.command_socket.close()
+        if self.status_socket is not None:
+            self.status_socket.close()
+        self.data_socket.close()
+        self.broadcast_socket.close()
+        self.handshake_socket.close()
+        self.context.term()
+
+    def _get_client_id(self):
+        # Return a unique identifier for this client (could be hostname, PID, etc.)
+        return f"{socket.gethostname()}_{os.getpid()}"
+
+    def _send_handshake(self):
+        print(f"Sending handshake to {self.handshake_socket}")
+        self.handshake_socket.send_json(
+            {"type": "hello", "client_id": self._get_client_id()}
+        )
+        self.handshake_socket.recv_json()  # Wait for ack
+
+
+if __name__ == "__main__":
+
+    def _server_thread(server_comms):
+        server_comms.wait_for_clients(expected_count=3)
+        server_comms.send_data({"data": "test"})
+        # server_comms.send_command({"command": "test"})
+        # print("Server sent command")
+
+    def _client_thread(script_comms):
+        for data in script_comms.receive_data():
+            print("Client received data:", data)
+
+    server_comms = ArborServerCommsHandler()
+    t1 = threading.Thread(target=_server_thread, args=(server_comms,))
+    t1.start()
+    print("Server started")
+
+    client_threads = []
+    script_comms_list = []
+    for i in range(3):
+        script_comms = ArborScriptCommsHandler(
+            "localhost",
+            server_comms.command_port,
+            server_comms.status_port,
+            server_comms.data_port,
+            server_comms.broadcast_port,
+            server_comms.handshake_port,
+            False,
+        )
+        t = threading.Thread(target=_client_thread, args=(script_comms,))
+        t.start()
+        script_comms_list.append(script_comms)
+
+    import time
+
+    time.sleep(1)
+    import pdb
+
+    pdb.set_trace()
+
+    try:
+        t1.join()
+        for t in client_threads:
+            t.join()
+    except KeyboardInterrupt:
+        print("Keyboard interrupt")
+    except Exception as e:
+        print(f"Error: {e}")
+    finally:
+        for script_comms in script_comms_list:
+            script_comms.close()
+        server_comms.close()

arbor/server/services/file_manager.py

@@ -0,0 +1,289 @@
+import json
+import os
+import shutil
+import time
+import uuid
+from pathlib import Path
+
+from fastapi import UploadFile
+
+from arbor.server.core.config import Settings
+
+
+class FileValidationError(Exception):
+    """Custom exception for file validation errors"""
+
+    pass
+
+
+class FileManager:
+    def __init__(self, settings: Settings):
+        self.uploads_dir = Path(settings.STORAGE_PATH) / "uploads"
+        self.uploads_dir.mkdir(parents=True, exist_ok=True)
+        self.files = self.load_files_from_uploads()
+
+    def load_files_from_uploads(self):
+        files = {}
+
+        # Scan through all directories in uploads directory
+        for dir_path in self.uploads_dir.glob("*"):
+            if not dir_path.is_dir():
+                continue
+
+            # Check for metadata.json
+            metadata_path = dir_path / "metadata.json"
+            if not metadata_path.exists():
+                continue
+
+            # Load metadata
+            with open(metadata_path) as f:
+                metadata = json.load(f)
+
+            # Find the .jsonl file
+            jsonl_files = list(dir_path.glob("*.jsonl"))
+            if not jsonl_files:
+                continue
+
+            file_path = jsonl_files[0]
+            files[dir_path.name] = {
+                "path": str(file_path),
+                "purpose": metadata.get("purpose", "training"),
+                "bytes": file_path.stat().st_size,
+                "created_at": metadata.get(
+                    "created_at", int(file_path.stat().st_mtime)
+                ),
+                "filename": metadata.get("filename", file_path.name),
+            }
+
+        return files
+
+    def save_uploaded_file(self, file: UploadFile):
+        file_id = f"file-{str(uuid.uuid4())}"
+        dir_path = self.uploads_dir / file_id
+        dir_path.mkdir(exist_ok=True)
+
+        # Save the actual file
+        file_path = dir_path / f"data.jsonl"
+        with open(file_path, "wb") as f:
+            shutil.copyfileobj(file.file, f)
+
+        # Create metadata
+        metadata = {
+            "purpose": "training",
+            "created_at": int(time.time()),
+            "filename": file.filename,
+        }
+
+        # Save metadata
+        with open(dir_path / "metadata.json", "w") as f:
+            json.dump(metadata, f)
+
+        file_data = {
+            "id": file_id,
+            "path": str(file_path),
+            "purpose": metadata["purpose"],
+            "bytes": file.size,
+            "created_at": metadata["created_at"],
+            "filename": metadata["filename"],
+        }
+
+        self.files[file_id] = file_data
+        return file_data
+
+    def get_file(self, file_id: str):
+        return self.files[file_id]
+
+    def delete_file(self, file_id: str):
+        if file_id not in self.files:
+            return
+
+        dir_path = self.uploads_dir / file_id
+        if dir_path.exists():
+            shutil.rmtree(dir_path)
+
+        del self.files[file_id]
+
+    def validate_file_format_sft(self, file_path: str) -> None:
+        """
+        Validates that the file at file_path is properly formatted JSONL with expected structure.
+        Raises FileValidationError if validation fails.
+        """
+        try:
+            with open(file_path, "r", encoding="utf-8") as f:
+                for line_num, line in enumerate(f, 1):
+                    line = line.strip()
+                    if not line:
+                        continue  # skip empty lines
+                    try:
+                        data = json.loads(line)
+
+                        if not isinstance(data, dict):
+                            raise FileValidationError(
+                                f"Line {line_num}: Each line must be a JSON object"
+                            )
+
+                        if "messages" not in data:
+                            raise FileValidationError(
+                                f"Line {line_num}: Missing 'messages' field"
+                            )
+
+                        if not isinstance(data["messages"], list):
+                            raise FileValidationError(
+                                f"Line {line_num}: 'messages' must be an array"
+                            )
+
+                        for msg in data["messages"]:
+                            if not isinstance(msg, dict):
+                                raise FileValidationError(
+                                    f"Line {line_num}: Each message must be an object"
+                                )
+                            if "role" not in msg or "content" not in msg:
+                                raise FileValidationError(
+                                    f"Line {line_num}: Messages must have 'role' and 'content' fields"
+                                )
+                            if not isinstance(msg["role"], str) or not isinstance(
+                                msg["content"], str
+                            ):
+                                raise FileValidationError(
+                                    f"Line {line_num}: Message 'role' and 'content' must be strings"
+                                )
+
+                    except json.JSONDecodeError:
+                        raise FileValidationError(f"Invalid JSON on line {line_num}")
+
+        except Exception as e:
+            raise FileValidationError(f"Failed to read or validate file: {e}")
+
+    def validate_file_format_dpo(self, file_path: str) -> None:
+        """
+        Validates that the file at file_path is properly formatted JSONL with expected structure
+        for tool-use data (input/messages/tools/parallel_tool_calls and outputs).
+        Raises FileValidationError if validation fails.
+        """
+        try:
+            with open(file_path, "r", encoding="utf-8") as f:
+                for line_num, line in enumerate(f, 1):
+                    line = line.strip()
+                    if not line:
+                        continue
+                    try:
+                        data = json.loads(line)
+
+                        if not isinstance(data, dict):
+                            raise FileValidationError(
+                                f"Line {line_num}: Each line must be a JSON object"
+                            )
+
+                        input_data = data.get("input")
+                        if not isinstance(input_data, dict):
+                            raise FileValidationError(
+                                f"Line {line_num}: Missing or invalid 'input' field"
+                            )
+
+                        if "messages" not in input_data or not isinstance(
+                            input_data["messages"], list
+                        ):
+                            raise FileValidationError(
+                                f"Line {line_num}: 'input.messages' must be a list"
+                            )
+                        for msg in input_data["messages"]:
+                            if not isinstance(msg, dict):
+                                raise FileValidationError(
+                                    f"Line {line_num}: Each 'message' must be an object"
+                                )
+                            if "role" not in msg or "content" not in msg:
+                                raise FileValidationError(
+                                    f"Line {line_num}: Each message must have 'role' and 'content'"
+                                )
+                            if not isinstance(msg["role"], str) or not isinstance(
+                                msg["content"], str
+                            ):
+                                raise FileValidationError(
+                                    f"Line {line_num}: 'role' and 'content' must be strings"
+                                )
+
+                        if "tools" not in input_data or not isinstance(
+                            input_data["tools"], list
+                        ):
+                            raise FileValidationError(
+                                f"Line {line_num}: 'input.tools' must be a list"
+                            )
+
+                        if "parallel_tool_calls" not in input_data or not isinstance(
+                            input_data["parallel_tool_calls"], bool
+                        ):
+                            raise FileValidationError(
+                                f"Line {line_num}: 'input.parallel_tool_calls' must be a boolean"
+                            )
+
+                        preferred = data.get("preferred_output")
+                        if not isinstance(preferred, list):
+                            raise FileValidationError(
+                                f"Line {line_num}: 'preferred_output' must be a list"
+                            )
+                        for msg in preferred:
+                            if not isinstance(msg, dict):
+                                raise FileValidationError(
+                                    f"Line {line_num}: Each 'preferred_output' message must be an object"
+                                )
+                            if "role" not in msg or "content" not in msg:
+                                raise FileValidationError(
+                                    f"Line {line_num}: Each preferred_output message must have 'role' and 'content'"
+                                )
+                            if not isinstance(msg["role"], str) or not isinstance(
+                                msg["content"], str
+                            ):
+                                raise FileValidationError(
+                                    f"Line {line_num}: 'role' and 'content' in preferred_output must be strings"
+                                )
+
+                        non_preferred = data.get("non_preferred_output")
+                        if not isinstance(non_preferred, list):
+                            raise FileValidationError(
+                                f"Line {line_num}: 'non_preferred_output' must be a list"
+                            )
+                        for msg in non_preferred:
+                            if not isinstance(msg, dict):
+                                raise FileValidationError(
+                                    f"Line {line_num}: Each 'non_preferred_output' message must be an object"
+                                )
+                            if "role" not in msg or "content" not in msg:
+                                raise FileValidationError(
+                                    f"Line {line_num}: Each non_preferred_output message must have 'role' and 'content'"
+                                )
+                            if not isinstance(msg["role"], str) or not isinstance(
+                                msg["content"], str
+                            ):
+                                raise FileValidationError(
+                                    f"Line {line_num}: 'role' and 'content' in non_preferred_output must be strings"
+                                )
+
+                    except json.JSONDecodeError:
+                        raise FileValidationError(f"Invalid JSON on line {line_num}")
+
+        except Exception as e:
+            raise FileValidationError(f"Failed to validate file: {e}")
+
+        output_path = file_path.replace(".jsonl", "_formatted.jsonl")
+
+        with (
+            open(file_path, "r", encoding="utf-8") as fin,
+            open(output_path, "w", encoding="utf-8") as fout,
+        ):
+            for line_num, line in enumerate(fin, 1):
+                line = line.strip()
+                if not line:
+                    continue
+                try:
+                    data = json.loads(line)
+                    prompt = data["input"]["messages"]
+                    new_line = {
+                        "chosen": data["preferred_output"],
+                        "rejected": data["non_preferred_output"],
+                        "prompt": prompt,
+                    }
+                    fout.write(json.dumps(new_line) + "\n")
+                except Exception as e:
+                    print(f"Error parsing line {line_num}: {e}")
+
+        os.replace(output_path, file_path)