model-runner-client 0.1.0__tar.gz

model_runner_client-0.1.0/PKG-INFO

@@ -0,0 +1,107 @@
+Metadata-Version: 2.1
+Name: model-runner-client
+Version: 0.1.0
+Summary: Model Runner Client is a Python library that allows Coordinators to manage real-time model synchronization and perform concurrent predictions on distributed model nodes within a Crunch.
+Author: boutrig abdennour
+Author-email: abdennour.boutrig@crunchdao.com
+Requires-Python: >=3.11,<4.0
+Classifier: Programming Language :: Python :: 3
+Classifier: Programming Language :: Python :: 3.11
+Classifier: Programming Language :: Python :: 3.12
+Classifier: Programming Language :: Python :: 3.13
+Requires-Dist: grpcio (>=1.70.0,<2.0.0)
+Requires-Dist: pandas (>=2.2.3,<3.0.0)
+Requires-Dist: protobuf (>=5.29.3,<6.0.0)
+Requires-Dist: pyarrow (>=19.0.0,<20.0.0)
+Requires-Dist: websockets (>=14.2,<15.0)
+Description-Content-Type: text/markdown
+
+# Model Runner Client
+
+**Model Runner Client** is a Python library that allows you, as a Coordinator, to interact with models participating in your crunch. It tracks which models join or leave through a WebSocket connection to the model nodes.
+
+- **Real-Time Model Sync**: Each model participating in your crunch is an instance of `ModelRunner`, maintained via WebSocket in the `ModelCluster`.
+- **Concurrent Predictions (with Timeout Handling)**: Use the `ModelConcurrentRunner` to request predictions from all models simultaneously. Define a timeout to avoid blocking if a model takes too long to predict.
+
+## Installation
+
+```bash
+pip install model-runner-client
+```
+
+> **Note**: Adjust this command (e.g., `pip3` or virtual environments) depending on your setup.
+
+## Usage
+
+Below is a quick example focusing on the ModelConcurrentRunner. It handles concurrent predictions for you and returns all results in one go.
+
+```python
+import asyncio
+from model_runner_client.model_concurrent_runner import ModelConcurrentRunner
+from model_runner_client.protos.model_runner_pb2 import DataType
+from model_runner_client.utils.datatype_transformer import encode_data
+
+async def main():
+    # crunch_id, host, and port are values provided by crunchdao
+    concurrent_runner = ModelConcurrentRunner(
+        timeout=10,
+        crunch_id="bird-game",
+        host="localhost",
+        port=8000
+    )
+
+    # Initialize communication with the model nodes to fetch 
+    # models that want to predict and set up the model cluster
+    await concurrent_runner.init()
+
+    async def prediction_call():
+        while True:
+            # Your data to be predicted (X)
+            value = {
+                'falcon_location': 21.179864629354732,
+                'time': 230.96231205799998,
+                'dove_location': 19.164986723324326,
+                'falcon_id': 1
+            }
+
+            # Encode data as binary and request a prediction
+            result = await concurrent_runner.predict(
+                DataType.JSON,
+                encode_data(DataType.JSON, value)
+            )
+
+            # You receive a dictionary of predictions
+            for model_runner, model_predict_result in result.items():
+                print(f"{model_runner.model_id}: {model_predict_result}")
+
+            # This pause (30s) simulates other work 
+            # the Coordinator might perform between predictions
+            await asyncio.sleep(30)
+
+    # Keep the cluster updated with `concurrent_runner.sync()`, 
+    # which maintains a permanent WebSocket connection.
+    # Then run our prediction process.
+    await asyncio.gather(
+        asyncio.create_task(concurrent_runner.sync()),
+        asyncio.create_task(prediction_call())
+    )
+
+if __name__ == "__main__":
+    try:
+        asyncio.run(main())
+    except KeyboardInterrupt:
+        print("\nReceived exit signal, shutting down gracefully.")
+```
+
+### Important Notes
+
+- **Prediction Failures & Timeouts**: A prediction may fail or exceed the defined timeout, so be sure to handle these cases appropriately. Refer to `ModelPredictResult.Status` for details.
+- **Custom Implementations**: If you need more control over your workflow, you can manage each model individually. Instead of using `ModelConcurrentRunner`, you can directly leverage `ModelRunner` instances from the `ModelCluster`, customizing how you schedule predictions and handle results.
+
+## Contributing
+
+Contributions are welcome! Feel free to open issues or submit pull requests if you encounter any bugs or want to suggest improvements.
+
+## License
+
+This project is distributed under the [MIT License](https://choosealicense.com/licenses/mit/). See the LICENSE file for details.

model_runner_client-0.1.0/README.md

@@ -0,0 +1,89 @@
+# Model Runner Client
+
+**Model Runner Client** is a Python library that allows you, as a Coordinator, to interact with models participating in your crunch. It tracks which models join or leave through a WebSocket connection to the model nodes.
+
+- **Real-Time Model Sync**: Each model participating in your crunch is an instance of `ModelRunner`, maintained via WebSocket in the `ModelCluster`.
+- **Concurrent Predictions (with Timeout Handling)**: Use the `ModelConcurrentRunner` to request predictions from all models simultaneously. Define a timeout to avoid blocking if a model takes too long to predict.
+
+## Installation
+
+```bash
+pip install model-runner-client
+```
+
+> **Note**: Adjust this command (e.g., `pip3` or virtual environments) depending on your setup.
+
+## Usage
+
+Below is a quick example focusing on the ModelConcurrentRunner. It handles concurrent predictions for you and returns all results in one go.
+
+```python
+import asyncio
+from model_runner_client.model_concurrent_runner import ModelConcurrentRunner
+from model_runner_client.protos.model_runner_pb2 import DataType
+from model_runner_client.utils.datatype_transformer import encode_data
+
+async def main():
+    # crunch_id, host, and port are values provided by crunchdao
+    concurrent_runner = ModelConcurrentRunner(
+        timeout=10,
+        crunch_id="bird-game",
+        host="localhost",
+        port=8000
+    )
+
+    # Initialize communication with the model nodes to fetch 
+    # models that want to predict and set up the model cluster
+    await concurrent_runner.init()
+
+    async def prediction_call():
+        while True:
+            # Your data to be predicted (X)
+            value = {
+                'falcon_location': 21.179864629354732,
+                'time': 230.96231205799998,
+                'dove_location': 19.164986723324326,
+                'falcon_id': 1
+            }
+
+            # Encode data as binary and request a prediction
+            result = await concurrent_runner.predict(
+                DataType.JSON,
+                encode_data(DataType.JSON, value)
+            )
+
+            # You receive a dictionary of predictions
+            for model_runner, model_predict_result in result.items():
+                print(f"{model_runner.model_id}: {model_predict_result}")
+
+            # This pause (30s) simulates other work 
+            # the Coordinator might perform between predictions
+            await asyncio.sleep(30)
+
+    # Keep the cluster updated with `concurrent_runner.sync()`, 
+    # which maintains a permanent WebSocket connection.
+    # Then run our prediction process.
+    await asyncio.gather(
+        asyncio.create_task(concurrent_runner.sync()),
+        asyncio.create_task(prediction_call())
+    )
+
+if __name__ == "__main__":
+    try:
+        asyncio.run(main())
+    except KeyboardInterrupt:
+        print("\nReceived exit signal, shutting down gracefully.")
+```
+
+### Important Notes
+
+- **Prediction Failures & Timeouts**: A prediction may fail or exceed the defined timeout, so be sure to handle these cases appropriately. Refer to `ModelPredictResult.Status` for details.
+- **Custom Implementations**: If you need more control over your workflow, you can manage each model individually. Instead of using `ModelConcurrentRunner`, you can directly leverage `ModelRunner` instances from the `ModelCluster`, customizing how you schedule predictions and handle results.
+
+## Contributing
+
+Contributions are welcome! Feel free to open issues or submit pull requests if you encounter any bugs or want to suggest improvements.
+
+## License
+
+This project is distributed under the [MIT License](https://choosealicense.com/licenses/mit/). See the LICENSE file for details.

model_runner_client-0.1.0/model_runner_client/model_cluster.py

@@ -0,0 +1,127 @@
+import logging
+
+from model_runner_client.model_runner import ModelRunner
+from model_runner_client.websocket_client import WebsocketClient
+
+logger = logging.getLogger("model_runner_client")
+import asyncio
+
+
+class ModelCluster:
+    def __init__(self, crunch_id: str, ws_host: str, ws_port: int):
+        """
+        ModelCluster constructor.
+
+        :param crunch_id: The Crunch ID that this cluster is responsible for.
+        :param ws_host: The WebSocket server's host.
+        :param ws_port: The WebSocket server's port.
+        """
+        self.crunch_id = crunch_id
+        self.models_run = {}
+        logger.debug(f"Initializing ModelCluster with Crunch ID: {crunch_id}")
+        self.ws_client = WebsocketClient(ws_host, ws_port, crunch_id, event_handler=self.handle_event)
+
+    async def init(self):
+        await self.ws_client.connect()
+        await self.ws_client.init()
+
+        logger.debug("WebSocket client initialized.")
+
+    async def sync(self):
+        await self.ws_client.listen()
+
+    async def handle_event(self, event_type: str, data: list[dict]):
+        """
+        Handle WebSocket events (`init` and `update`) and update the cluster's state.
+
+        :param event_type: The type of the event (`init` or `update`).
+        :param data: The event data.
+        """
+        try:
+            if event_type == "init":
+                logger.debug(f"**ModelCluster** Processing event type: {event_type}")
+                await self.handle_init_event(data)
+            elif event_type == "update":
+                logger.debug(f"**ModelCluster** Processing event type: {event_type}")
+                await self.handle_update_event(data)
+            else:
+                logger.warning(f"**ModelCluster** Unknown event type received: {event_type}")
+        except Exception as e:
+            logger.error(f"**ModelCluster** Error processing event {event_type}: {e}", exc_info=True)
+            raise e
+
+    async def handle_init_event(self, data: list[dict]):
+        """
+        Process the `init` event to initialize models run.
+        
+        :param data: List of models with their initial states.
+        """
+        logger.debug("**ModelCluster** Handling 'init' event.")
+        await self.update_model_runs(data)
+
+    async def handle_update_event(self, data: list[dict]):
+        """
+        Process the `update` event to update model states.
+        
+        :param data: List of models with their updated states.
+        """
+        logger.debug("**ModelCluster** Handling 'update' event.")
+        await self.update_model_runs(data)
+
+    async def update_model_runs(self, data):
+        tasks = []
+        for model_update in data:
+            model_id = model_update.get("model_id")
+            model_name = model_update.get("model_name")
+            state = model_update.get("state")
+            ip = model_update.get("ip")
+            port = model_update.get("port")
+            logger.debug(f"**ModelCluster** Updating model with ID: {model_id}")
+
+            # Find the model in the current state
+            model_runner = self.models_run.get(model_id)
+
+            if model_runner:
+                if state == "STOPPED":
+                    # Remove model if state is "stopped"
+                    tasks.append(self.remove_model_runner(model_runner))
+                    logger.debug(f"**ModelCluster** Model with ID {model_id} removed due to 'stopped' state.")
+                elif state == "RUNNING":
+                    logger.debug(f"**ModelCluster** Model with ID {model_id} is already running in the cluster. Skipping update for '{model_name}' with state: {state}.")
+                else:
+                    logger.warning(f"**ModelCluster** Model updated: {model_id}, with state: {state} => This state is not handled...")
+            else:
+                if state == "STOPPED":
+                    logger.debug(f"**ModelCluster** Model with ID {model_id} is not found in the cluster state, and its state is 'STOPPED'. No action is required.")
+                elif state == "RUNNING":
+                    logger.debug(f"**ModelCluster** New model with ID {model_id} is running, we add it to the cluster state.")
+                    model_runner = ModelRunner(model_id, model_name, ip, port)
+                    tasks.append(self.add_model_runner(model_runner))
+                else:
+                    logger.warning(f"**ModelCluster** Model updated: {model_id}, with state: {state} => This state is not handled...")
+
+        await asyncio.gather(*tasks)
+
+    async def add_model_runner(self, model_runner: ModelRunner):
+        """
+        Asynchronously initialize a model_runner and add it to the cluster state.
+        """
+        is_initialized = await model_runner.init()
+        if is_initialized:
+            self.models_run[model_runner.model_id] = model_runner
+
+    async def remove_model_runner(self, model_runner: ModelRunner):
+        """
+        Asynchronously initialize a model_runner and add it to the cluster state.
+        """
+        await model_runner.close()
+        del self.models_run[model_runner.model_id]
+
+    async def start(self):
+        """
+        Start the WebSocket client and handle events.
+        """
+        try:
+            await self.ws_client.connect()
+        except Exception as e:
+            logger.error(f"**ModelCluster** Failed to start WebSocket client: {e}", exc_info=True)

model_runner_client-0.1.0/model_runner_client/model_concurrent_runner.py

@@ -0,0 +1,65 @@
+import asyncio
+import logging
+from enum import Enum
+
+from model_runner_client.model_cluster import ModelCluster
+from model_runner_client.model_runner import ModelRunner
+from model_runner_client.protos.model_runner_pb2 import DataType
+
+logger = logging.getLogger("model_runner_client")
+
+
+class ModelPredictResult:
+    class Status(Enum):
+        SUCCESS = "SUCCESS"
+        FAILED = "FAILED"
+        TIMEOUT = "TIMEOUT"
+
+    def __init__(self, model_runner: ModelRunner, result: any, status: Status):
+        self.model_runner = model_runner
+        self.result = result
+        self.status = status
+
+    def __str__(self):
+        return f"ModelPredictResult(model_runner={self.model_runner}, result={self.result}, status={self.status.name})"
+
+
+class ModelConcurrentRunner:
+    def __init__(self, timeout: int, crunch_id: str, host: str, port: int):
+        self.timeout = timeout
+        self.host = host
+        self.port = port
+        self.model_cluster = ModelCluster(crunch_id, self.host, self.port)
+
+        # TODO: If the model returns failures exceeding max_consecutive_failures, exclude the model. Maybe also inform the orchestrator to STOP the model ?
+        #self.max_consecutive_failures
+
+        # TODO: Implement this. If the option is enabled, allow the model time to recover after a timeout.
+        #self.enable_recovery_mode
+        #self.recovery_time
+
+
+    async def init(self):
+        await self.model_cluster.init()
+
+    async def sync(self):
+        await self.model_cluster.sync()
+
+    async def predict(self, argument_type: DataType, argument_value: bytes) -> dict[ModelRunner, ModelPredictResult]:
+        tasks = [self._predict_with_timeout(model, argument_type, argument_value)
+                 for model in self.model_cluster.models_run.values()]
+        logger.debug(f"**ModelConcurrentRunner** predict tasks: {tasks}")
+        results = await asyncio.gather(*tasks, return_exceptions=True)
+
+        return {result.model_runner: result for result in results}
+
+    async def _predict_with_timeout(self, model: ModelRunner, argument_type: DataType, argument_value: bytes):
+        try:
+            result = await asyncio.wait_for(
+                model.predict(argument_type, argument_value), timeout=self.timeout
+            )
+            return ModelPredictResult(model, result, ModelPredictResult.Status.SUCCESS)
+        except asyncio.TimeoutError:
+            return ModelPredictResult(model, None, ModelPredictResult.Status.TIMEOUT)
+        except Exception as e:
+            return ModelPredictResult(model, str(e), ModelPredictResult.Status.FAILED)

model_runner_client-0.1.0/model_runner_client/model_runner.py

@@ -0,0 +1,76 @@
+import asyncio
+import atexit
+import logging
+
+import grpc
+from google.protobuf import empty_pb2
+from grpc.aio import AioRpcError
+
+from model_runner_client.protos.model_runner_pb2 import DataType, InferRequest
+from model_runner_client.protos.model_runner_pb2_grpc import ModelRunnerStub
+from model_runner_client.utils.datatype_transformer import decode_data
+
+logger = logging.getLogger("model_runner_client")
+
+
+class ModelRunner:
+    def __init__(self, model_id: str, model_name: str, ip: str, port: int):
+
+        self.model_id = model_id
+        self.model_name = model_name
+        self.ip = ip
+        self.port = port
+        logger.info(f"**ModelRunner** New model runner created: {self.model_id}, {self.model_name}, {self.ip}:{self.port}, let's connect it")
+
+        self.grpc_channel = None
+        self.grpc_stub = None
+        self.retry_attempts = 5  # args ?
+        self.min_retry_interval = 2  # 2 seconds
+        self.closed = False
+
+    def __del__(self):
+        logger.debug(f"**ModelRunner** Model runner {self.model_id} is destroyed")
+        atexit.register(self.close_sync)
+
+    async def init(self) -> bool:
+        for attempt in range(1, self.retry_attempts + 1):
+            if self.closed:
+                logger.debug(f"**ModelRunner** Model runner {self.model_id} closed, aborting initialization")
+                return False
+            try:
+                self.grpc_channel = grpc.aio.insecure_channel(f"{self.ip}:{self.port}")
+                self.grpc_stub = ModelRunnerStub(self.grpc_channel)
+                await self.grpc_stub.Setup(empty_pb2.Empty())  # maybe orchestrator has to do that ?
+                logger.info(f"**ModelRunner** model runner: {self.model_id}, {self.model_name}, is connected and ready")
+                return True
+            except (AioRpcError, asyncio.TimeoutError) as e:
+                logger.error(f"**ModelRunner** Model {self.model_id} initialization failed on attempt {attempt}/{self.retry_attempts}: {e}")
+            except Exception as e:
+                logger.error(f"**ModelRunner** Unexpected error during initialization of model {self.model_id}: {e}", exc_info=True)
+
+            if attempt < self.retry_attempts:
+                backoff_time = 2 ** attempt  # Backoff with exponential delay
+                logger.warning(f"**ModelRunner** Retrying in {backoff_time} seconds...")
+                await asyncio.sleep(backoff_time)
+            else:
+                logger.error(f"**ModelRunner** Model {self.model_id} failed to initialize after {self.retry_attempts} attempts.")
+                # todo what is the behavior here ? remove it locally ?
+                return False
+
+    async def predict(self, argument_type: DataType, argument_value: bytes):
+        logger.debug(f"**ModelRunner** Doing prediction of model_id:{self.model_id}, name:{self.model_name}, argument_type:{argument_type}")
+        prediction_request = InferRequest(type=argument_type, argument=argument_value)
+
+        response = await self.grpc_stub.Infer(prediction_request)
+
+        return decode_data(response.prediction, response.type)
+
+    def close_sync(self):
+        loop = asyncio.get_event_loop()
+        loop.run_until_complete(self.close())
+
+    async def close(self):
+        self.closed = True
+        if self.grpc_channel:
+            await self.grpc_channel.close()
+            logger.debug(f"**ModelRunner** Model runner {self.model_id} grpc connection closed")

model_runner_client-0.1.0/model_runner_client/protos/model_runner.proto

@@ -0,0 +1,29 @@
+syntax = "proto3";
+
+import "google/protobuf/empty.proto";
+
+service ModelRunner {
+  rpc InferStream(stream InferRequest) returns (stream InferResponse);
+  rpc Infer(InferRequest) returns (InferResponse);
+  rpc Setup(google.protobuf.Empty) returns (google.protobuf.Empty);
+  rpc Reinitialize(google.protobuf.Empty) returns (google.protobuf.Empty);
+}
+
+enum DataType {
+  DOUBLE = 0;
+  INT = 1;
+  STRING = 2;
+  PARQUET = 3;
+  ARROW = 4;
+  JSON = 5;
+}
+
+message InferRequest {
+  DataType type = 1;
+  bytes argument = 2;
+}
+
+message InferResponse {
+  DataType type = 1;
+  bytes prediction = 2;
+}

model_runner_client-0.1.0/model_runner_client/protos/model_runner_pb2.py

@@ -0,0 +1,43 @@
+# -*- coding: utf-8 -*-
+# Generated by the protocol buffer compiler.  DO NOT EDIT!
+# NO CHECKED-IN PROTOBUF GENCODE
+# source: model_runner.proto
+# Protobuf Python Version: 5.29.0
+"""Generated protocol buffer code."""
+from google.protobuf import descriptor as _descriptor
+from google.protobuf import descriptor_pool as _descriptor_pool
+from google.protobuf import runtime_version as _runtime_version
+from google.protobuf import symbol_database as _symbol_database
+from google.protobuf.internal import builder as _builder
+_runtime_version.ValidateProtobufRuntimeVersion(
+    _runtime_version.Domain.PUBLIC,
+    5,
+    29,
+    0,
+    '',
+    'model_runner.proto'
+)
+# @@protoc_insertion_point(imports)
+
+_sym_db = _symbol_database.Default()
+
+
+from google.protobuf import empty_pb2 as google_dot_protobuf_dot_empty__pb2
+
+
+DESCRIPTOR = _descriptor_pool.Default().AddSerializedFile(b'\n\x12model_runner.proto\x1a\x1bgoogle/protobuf/empty.proto\"9\n\x0cInferRequest\x12\x17\n\x04type\x18\x01 \x01(\x0e\x32\t.DataType\x12\x10\n\x08\x61rgument\x18\x02 \x01(\x0c\"<\n\rInferResponse\x12\x17\n\x04type\x18\x01 \x01(\x0e\x32\t.DataType\x12\x12\n\nprediction\x18\x02 \x01(\x0c*M\n\x08\x44\x61taType\x12\n\n\x06\x44OUBLE\x10\x00\x12\x07\n\x03INT\x10\x01\x12\n\n\x06STRING\x10\x02\x12\x0b\n\x07PARQUET\x10\x03\x12\t\n\x05\x41RROW\x10\x04\x12\x08\n\x04JSON\x10\x05\x32\xe0\x01\n\x0bModelRunner\x12\x30\n\x0bInferStream\x12\r.InferRequest\x1a\x0e.InferResponse(\x01\x30\x01\x12&\n\x05Infer\x12\r.InferRequest\x1a\x0e.InferResponse\x12\x37\n\x05Setup\x12\x16.google.protobuf.Empty\x1a\x16.google.protobuf.Empty\x12>\n\x0cReinitialize\x12\x16.google.protobuf.Empty\x1a\x16.google.protobuf.Emptyb\x06proto3')
+
+_globals = globals()
+_builder.BuildMessageAndEnumDescriptors(DESCRIPTOR, _globals)
+_builder.BuildTopDescriptorsAndMessages(DESCRIPTOR, 'model_runner_pb2', _globals)
+if not _descriptor._USE_C_DESCRIPTORS:
+  DESCRIPTOR._loaded_options = None
+  _globals['_DATATYPE']._serialized_start=172
+  _globals['_DATATYPE']._serialized_end=249
+  _globals['_INFERREQUEST']._serialized_start=51
+  _globals['_INFERREQUEST']._serialized_end=108
+  _globals['_INFERRESPONSE']._serialized_start=110
+  _globals['_INFERRESPONSE']._serialized_end=170
+  _globals['_MODELRUNNER']._serialized_start=252
+  _globals['_MODELRUNNER']._serialized_end=476
+# @@protoc_insertion_point(module_scope)

model_runner_client-0.1.0/model_runner_client/protos/model_runner_pb2.pyi

@@ -0,0 +1,38 @@
+from google.protobuf import empty_pb2 as _empty_pb2
+from google.protobuf.internal import enum_type_wrapper as _enum_type_wrapper
+from google.protobuf import descriptor as _descriptor
+from google.protobuf import message as _message
+from typing import ClassVar as _ClassVar, Optional as _Optional, Union as _Union
+
+DESCRIPTOR: _descriptor.FileDescriptor
+
+class DataType(int, metaclass=_enum_type_wrapper.EnumTypeWrapper):
+    __slots__ = ()
+    DOUBLE: _ClassVar[DataType]
+    INT: _ClassVar[DataType]
+    STRING: _ClassVar[DataType]
+    PARQUET: _ClassVar[DataType]
+    ARROW: _ClassVar[DataType]
+    JSON: _ClassVar[DataType]
+DOUBLE: DataType
+INT: DataType
+STRING: DataType
+PARQUET: DataType
+ARROW: DataType
+JSON: DataType
+
+class InferRequest(_message.Message):
+    __slots__ = ("type", "argument")
+    TYPE_FIELD_NUMBER: _ClassVar[int]
+    ARGUMENT_FIELD_NUMBER: _ClassVar[int]
+    type: DataType
+    argument: bytes
+    def __init__(self, type: _Optional[_Union[DataType, str]] = ..., argument: _Optional[bytes] = ...) -> None: ...
+
+class InferResponse(_message.Message):
+    __slots__ = ("type", "prediction")
+    TYPE_FIELD_NUMBER: _ClassVar[int]
+    PREDICTION_FIELD_NUMBER: _ClassVar[int]
+    type: DataType
+    prediction: bytes
+    def __init__(self, type: _Optional[_Union[DataType, str]] = ..., prediction: _Optional[bytes] = ...) -> None: ...

model_runner_client-0.1.0/model_runner_client/protos/model_runner_pb2_grpc.py

@@ -0,0 +1,227 @@
+# Generated by the gRPC Python protocol compiler plugin. DO NOT EDIT!
+"""Client and server classes corresponding to protobuf-defined services."""
+import grpc
+import warnings
+
+from google.protobuf import empty_pb2 as google_dot_protobuf_dot_empty__pb2
+from . import model_runner_pb2 as model__runner__pb2
+
+GRPC_GENERATED_VERSION = '1.69.0'
+GRPC_VERSION = grpc.__version__
+_version_not_supported = False
+
+try:
+    from grpc._utilities import first_version_is_lower
+    _version_not_supported = first_version_is_lower(GRPC_VERSION, GRPC_GENERATED_VERSION)
+except ImportError:
+    _version_not_supported = True
+
+if _version_not_supported:
+    raise RuntimeError(
+        f'The grpc package installed is at version {GRPC_VERSION},'
+        + f' but the generated code in model_runner_pb2_grpc.py depends on'
+        + f' grpcio>={GRPC_GENERATED_VERSION}.'
+        + f' Please upgrade your grpc module to grpcio>={GRPC_GENERATED_VERSION}'
+        + f' or downgrade your generated code using grpcio-tools<={GRPC_VERSION}.'
+    )
+
+
+class ModelRunnerStub(object):
+    """Missing associated documentation comment in .proto file."""
+
+    def __init__(self, channel):
+        """Constructor.
+
+        Args:
+            channel: A grpc.Channel.
+        """
+        self.InferStream = channel.stream_stream(
+                '/ModelRunner/InferStream',
+                request_serializer=model__runner__pb2.InferRequest.SerializeToString,
+                response_deserializer=model__runner__pb2.InferResponse.FromString,
+                _registered_method=True)
+        self.Infer = channel.unary_unary(
+                '/ModelRunner/Infer',
+                request_serializer=model__runner__pb2.InferRequest.SerializeToString,
+                response_deserializer=model__runner__pb2.InferResponse.FromString,
+                _registered_method=True)
+        self.Setup = channel.unary_unary(
+                '/ModelRunner/Setup',
+                request_serializer=google_dot_protobuf_dot_empty__pb2.Empty.SerializeToString,
+                response_deserializer=google_dot_protobuf_dot_empty__pb2.Empty.FromString,
+                _registered_method=True)
+        self.Reinitialize = channel.unary_unary(
+                '/ModelRunner/Reinitialize',
+                request_serializer=google_dot_protobuf_dot_empty__pb2.Empty.SerializeToString,
+                response_deserializer=google_dot_protobuf_dot_empty__pb2.Empty.FromString,
+                _registered_method=True)
+
+
+class ModelRunnerServicer(object):
+    """Missing associated documentation comment in .proto file."""
+
+    def InferStream(self, request_iterator, context):
+        """Missing associated documentation comment in .proto file."""
+        context.set_code(grpc.StatusCode.UNIMPLEMENTED)
+        context.set_details('Method not implemented!')
+        raise NotImplementedError('Method not implemented!')
+
+    def Infer(self, request, context):
+        """Missing associated documentation comment in .proto file."""
+        context.set_code(grpc.StatusCode.UNIMPLEMENTED)
+        context.set_details('Method not implemented!')
+        raise NotImplementedError('Method not implemented!')
+
+    def Setup(self, request, context):
+        """Missing associated documentation comment in .proto file."""
+        context.set_code(grpc.StatusCode.UNIMPLEMENTED)
+        context.set_details('Method not implemented!')
+        raise NotImplementedError('Method not implemented!')
+
+    def Reinitialize(self, request, context):
+        """Missing associated documentation comment in .proto file."""
+        context.set_code(grpc.StatusCode.UNIMPLEMENTED)
+        context.set_details('Method not implemented!')
+        raise NotImplementedError('Method not implemented!')
+
+
+def add_ModelRunnerServicer_to_server(servicer, server):
+    rpc_method_handlers = {
+            'InferStream': grpc.stream_stream_rpc_method_handler(
+                    servicer.InferStream,
+                    request_deserializer=model__runner__pb2.InferRequest.FromString,
+                    response_serializer=model__runner__pb2.InferResponse.SerializeToString,
+            ),
+            'Infer': grpc.unary_unary_rpc_method_handler(
+                    servicer.Infer,
+                    request_deserializer=model__runner__pb2.InferRequest.FromString,
+                    response_serializer=model__runner__pb2.InferResponse.SerializeToString,
+            ),
+            'Setup': grpc.unary_unary_rpc_method_handler(
+                    servicer.Setup,
+                    request_deserializer=google_dot_protobuf_dot_empty__pb2.Empty.FromString,
+                    response_serializer=google_dot_protobuf_dot_empty__pb2.Empty.SerializeToString,
+            ),
+            'Reinitialize': grpc.unary_unary_rpc_method_handler(
+                    servicer.Reinitialize,
+                    request_deserializer=google_dot_protobuf_dot_empty__pb2.Empty.FromString,
+                    response_serializer=google_dot_protobuf_dot_empty__pb2.Empty.SerializeToString,
+            ),
+    }
+    generic_handler = grpc.method_handlers_generic_handler(
+            'ModelRunner', rpc_method_handlers)
+    server.add_generic_rpc_handlers((generic_handler,))
+    server.add_registered_method_handlers('ModelRunner', rpc_method_handlers)
+
+
+ # This class is part of an EXPERIMENTAL API.
+class ModelRunner(object):
+    """Missing associated documentation comment in .proto file."""
+
+    @staticmethod
+    def InferStream(request_iterator,
+            target,
+            options=(),
+            channel_credentials=None,
+            call_credentials=None,
+            insecure=False,
+            compression=None,
+            wait_for_ready=None,
+            timeout=None,
+            metadata=None):
+        return grpc.experimental.stream_stream(
+            request_iterator,
+            target,
+            '/ModelRunner/InferStream',
+            model__runner__pb2.InferRequest.SerializeToString,
+            model__runner__pb2.InferResponse.FromString,
+            options,
+            channel_credentials,
+            insecure,
+            call_credentials,
+            compression,
+            wait_for_ready,
+            timeout,
+            metadata,
+            _registered_method=True)
+
+    @staticmethod
+    def Infer(request,
+            target,
+            options=(),
+            channel_credentials=None,
+            call_credentials=None,
+            insecure=False,
+            compression=None,
+            wait_for_ready=None,
+            timeout=None,
+            metadata=None):
+        return grpc.experimental.unary_unary(
+            request,
+            target,
+            '/ModelRunner/Infer',
+            model__runner__pb2.InferRequest.SerializeToString,
+            model__runner__pb2.InferResponse.FromString,
+            options,
+            channel_credentials,
+            insecure,
+            call_credentials,
+            compression,
+            wait_for_ready,
+            timeout,
+            metadata,
+            _registered_method=True)
+
+    @staticmethod
+    def Setup(request,
+            target,
+            options=(),
+            channel_credentials=None,
+            call_credentials=None,
+            insecure=False,
+            compression=None,
+            wait_for_ready=None,
+            timeout=None,
+            metadata=None):
+        return grpc.experimental.unary_unary(
+            request,
+            target,
+            '/ModelRunner/Setup',
+            google_dot_protobuf_dot_empty__pb2.Empty.SerializeToString,
+            google_dot_protobuf_dot_empty__pb2.Empty.FromString,
+            options,
+            channel_credentials,
+            insecure,
+            call_credentials,
+            compression,
+            wait_for_ready,
+            timeout,
+            metadata,
+            _registered_method=True)
+
+    @staticmethod
+    def Reinitialize(request,
+            target,
+            options=(),
+            channel_credentials=None,
+            call_credentials=None,
+            insecure=False,
+            compression=None,
+            wait_for_ready=None,
+            timeout=None,
+            metadata=None):
+        return grpc.experimental.unary_unary(
+            request,
+            target,
+            '/ModelRunner/Reinitialize',
+            google_dot_protobuf_dot_empty__pb2.Empty.SerializeToString,
+            google_dot_protobuf_dot_empty__pb2.Empty.FromString,
+            options,
+            channel_credentials,
+            insecure,
+            call_credentials,
+            compression,
+            wait_for_ready,
+            timeout,
+            metadata,
+            _registered_method=True)

model_runner_client-0.1.0/model_runner_client/utils/datatype_transformer.py

@@ -0,0 +1,71 @@
+import json
+import struct
+import io
+import pandas
+import pyarrow.parquet as pq
+
+from model_runner_client.protos.model_runner_pb2 import DataType
+
+
+# Encoder: Converts data to bytes
+def encode_data(data_type: DataType, data) -> bytes:
+    if data_type == DataType.DOUBLE:
+        return struct.pack("d", data)
+    elif data_type == DataType.INT:
+        return data.to_bytes(8, byteorder="little", signed=True)
+    elif data_type == DataType.STRING:
+        return data.encode("utf-8")
+    elif data_type == DataType.PARQUET:
+        table = pandas.Table.from_pandas(data)
+        sink = io.BytesIO()
+        pandas.write_table(table, sink)
+        return sink.getvalue()
+    elif data_type == DataType.JSON:
+        try:
+            json_data = json.dumps(data)  # Convert the object to a JSON string
+            return json_data.encode("utf-8")  # Return the JSON string as bytes
+        except TypeError as e:
+            raise ValueError(f"Data cannot be serialized to JSON: {e}")
+    else:
+        raise ValueError(f"Unsupported data type: {data_type}")
+
+
+# Decoder: Converts bytes to data
+def decode_data(data_bytes: bytes, data_type: DataType):
+    if data_type == DataType.DOUBLE:
+        return struct.unpack("d", data_bytes)[0]
+    elif data_type == DataType.INT:
+        return int.from_bytes(data_bytes, byteorder="little", signed=True)
+    elif data_type == DataType.STRING:
+        return data_bytes.decode("utf-8")
+    elif data_type == DataType.PARQUET:
+        buffer = io.BytesIO(data_bytes)
+        return pq.read_table(buffer).to_pandas()
+    elif data_type == DataType.JSON:
+        try:
+            json_data = data_bytes.decode("utf-8")
+            return json.loads(json_data)
+        except json.JSONDecodeError as e:
+            raise ValueError(f"Failed to decode JSON data: {e}")
+
+    else:
+        raise ValueError(f"Unsupported data type: {data_type}")
+
+
+def detect_data_type(data) -> DataType:
+    """
+    Detects the data type based on the Python object and returns
+    the corresponding DataType enum.
+    """
+    if isinstance(data, float):
+        return DataType.DOUBLE
+    elif isinstance(data, int):
+        return DataType.INT
+    elif isinstance(data, str):
+        return DataType.STRING
+    elif isinstance(data, pandas.DataFrame):
+        return DataType.PARQUET
+    elif isinstance(data, dict) or isinstance(data, list):
+        return DataType.JSON
+    else:
+        raise ValueError(f"Unsupported data type: {type(data)}")

model_runner_client-0.1.0/model_runner_client/websocket_client.py

@@ -0,0 +1,134 @@
+import asyncio
+import atexit
+import json
+import logging
+
+import websockets
+from websockets import State
+
+logger = logging.getLogger("model_runner_client")
+
+
+class WebsocketClient:
+    def __init__(self, host, port, crunch_id, event_handler=None):
+        """
+        WebsocketClient constructor.
+
+        :param host: WebSocket server host.
+        :param port: WebSocket server port.
+        :param crunch_id: Crunch ID used to connect to the server.
+        :param event_handler: Optional handler for WebSocket events (delegated to ModelCluster).
+        """
+        self.retry_interval = 10
+        self.max_retries = 5
+        self.listening = None
+        self.host = host
+        self.port = port
+        self.crunch_id = crunch_id
+        self.websocket = None
+        self.event_handler = event_handler  # Delegate to ModelCluster
+
+    def __del__(self):
+        atexit.register(self.disconnect_sync)
+
+    async def connect(self):
+        """
+        Establish a WebSocket connection.
+        """
+        retry_count = 0
+        while self.max_retries > retry_count:
+            try:
+                uri = f"ws://{self.host}:{self.port}/{self.crunch_id}"
+                logger.info(f"Connecting to WebSocket server at {uri}")
+                self.websocket = await websockets.connect(uri)
+                logger.info(f"Connected to WebSocket server at {uri}")
+                break
+            except (websockets.exceptions.ConnectionClosed, ConnectionRefusedError, OSError, asyncio.TimeoutError) as e:
+                logger.warning(f"Connection error ({e.__class__.__name__}): Retrying in {self.retry_interval} seconds...")
+            except Exception as e:
+                logger.error(f"Unexpected error ({e.__class__.__name__}): {e}", exc_info=True)
+                logger.warning(f"Retrying in {self.retry_interval} seconds...")
+
+            await asyncio.sleep(self.retry_interval)
+            retry_count += 1
+
+    async def init(self):
+        """
+        Listen first message who is init from the WebSocket server.
+        """
+        # retry here doesn't make sens, it comme after connection and connection handle retries
+        try:
+            message = await self.websocket.recv()
+            await self.handle_event(message)
+        except websockets.exceptions.ConnectionClosed:
+            logger.warning("WebSocket connection closed by the server.")
+        except Exception as e:
+            logger.error(f"Error while listening to WebSocket messages: {e}", exc_info=True)
+
+    async def listen(self):
+        """
+        Listen for messages from the WebSocket server.
+        """
+        self.listening = True
+        while self.listening:
+            try:
+                if not await self.is_connected():
+                    await self.connect()
+
+                logger.info("Listening for messages...")
+                async for message in self.websocket:
+                    await self.handle_event(message)
+            except (websockets.exceptions.ConnectionClosed, asyncio.TimeoutError):
+                logger.warning(f"Connection lost. Retrying in {self.retry_interval} seconds...")
+                await asyncio.sleep(self.retry_interval)
+            except Exception as e:
+                logger.error(f"An unexpected error occurred: {e}", exc_info=True)
+                logger.warning(f"Retrying in {self.retry_interval} seconds...")
+                await asyncio.sleep(self.retry_interval)
+
+            finally:
+                # Cleanup or reset if needed after the loop iteration
+                # For example, clear websocket references if disconnected
+                if self.websocket and not self.websocket.state:
+                    self.websocket = None
+
+    async def is_connected(self):
+        """
+        Check if the WebSocket connection is open.
+        :return: True if the client is connected, False otherwise.
+        """
+        return self.websocket is not None and self.websocket.state == State.OPEN
+
+    async def handle_event(self, message):
+        """
+        Handle incoming WebSocket messages and forward to the event handler.
+
+        :param message: Message received from the server.
+        """
+        try:
+            event = json.loads(message)
+            event_type = event.get("event")
+            data = event.get("data")
+
+            logger.debug(f"Received event: {event_type}, data: {data}")
+
+            # Delegate to the event handler, if available
+            if self.event_handler:
+                await self.event_handler(event_type, data)
+            else:
+                logger.warning("No event handler defined. Event will be ignored.")
+        except json.JSONDecodeError:
+            logger.error(f"Failed to decode WebSocket message: {message}")
+
+    def disconnect_sync(self):
+        loop = asyncio.get_event_loop()
+        loop.run_until_complete(self.disconnect())
+
+    async def disconnect(self):
+        """
+        Close the WebSocket connection.
+        """
+        self.listening = False
+        if self.websocket:
+            await self.websocket.close()
+            logger.info("WebSocket connection closed.")

model_runner_client-0.1.0/pyproject.toml

@@ -0,0 +1,27 @@
+[tool.poetry]
+name = "model-runner-client"
+version = "0.1.0"
+description = "Model Runner Client is a Python library that allows Coordinators to manage real-time model synchronization and perform concurrent predictions on distributed model nodes within a Crunch."
+authors = ["boutrig abdennour <abdennour.boutrig@crunchdao.com>"]
+readme = "README.md"
+
+[tool.poetry.dependencies]
+python = "^3.11"
+grpcio = "^1.70.0"
+pyarrow = "^19.0.0"
+protobuf = "^5.29.3"
+pandas = "^2.2.3"
+websockets = "^14.2"
+
+
+[tool.poetry.group.dev.dependencies]
+pytest = "^8.3.4"
+pytest-asyncio = "^0.25.3"
+
+[build-system]
+requires = ["poetry-core"]
+build-backend = "poetry.core.masonry.api"
+
+[project.urls]
+"Homepage" = "https://github.com/crunchdao/model-runner-client"
+"Bug Tracker" = "https://github.com/crunchdao/model-runner-client/issues"