fix for cr

Author: CormickKneey

areal/scheduler/rpc/rpc_server.py

@@ -1,7 +1,9 @@
 import argparse
 import importlib
 import traceback
+from collections.abc import Callable
 from concurrent.futures import Future
+from typing import Any
 
 from flask import Flask, jsonify, request
 
@@ -20,6 +22,53 @@
 # Global engine instance - must be TrainEngine or InferenceEngine
 _engine: TrainEngine | InferenceEngine | None = None
 
+
+def _handle_submit(
+    engine: InferenceEngine,
+    args: list[Any],
+    kwargs: dict[str, Any],
+) -> tuple[list[Any], dict[str, Any]]:
+    workflow_path = kwargs["workflow_path"]
+    workflow_kwargs = kwargs["workflow_kwargs"]
+    episode_data = kwargs["data"]
+    should_accept_path = kwargs["should_accept_path"]
+
+    episode_data = deserialize_value(episode_data)
+
+    module_path, class_name = workflow_path.rsplit(".", 1)
+    module = importlib.import_module(module_path)
+    workflow_class = getattr(module, class_name)
+    logger.info(f"Imported workflow class: {workflow_path}")
+
+    workflow_kwargs = deserialize_value(workflow_kwargs)
+    workflow = workflow_class(**workflow_kwargs)
+    logger.info(f"Workflow '{workflow_path}' instantiated successfully")
+
+    should_accept = None
+    if should_accept_path is not None:
+        module_path, fn_name = should_accept_path.rsplit(".", 1)
+        module = importlib.import_module(module_path)
+        should_accept = getattr(module, fn_name)
+        logger.info(f"Imported filtering function: {should_accept_path}")
+
+    new_args: list[Any] = []
+    new_kwargs: dict[str, Any] = dict(
+        data=episode_data,
+        workflow=workflow,
+        should_accept=should_accept,
+    )
+    return new_args, new_kwargs
+
+
+_METHOD_HANDLERS: dict[
+    str,
+    Callable[
+        [InferenceEngine, list[Any], dict[str, Any]], tuple[list[Any], dict[str, Any]]
+    ],
+] = {
+    "submit": _handle_submit,
+}
+
 # Create Flask app
 app = Flask(__name__)
 
@@ -207,42 +256,12 @@ def call_engine_method():
                 500,
             )
 
-        # Special case for `submit` on inference engines
+        # Special handling for some methods on inference engines
         try:
-            if method_name == "submit" and isinstance(_engine, InferenceEngine):
-                workflow_path = kwargs["workflow_path"]
-                workflow_kwargs = kwargs["workflow_kwargs"]
-                episode_data = kwargs["data"]
-                should_accept_path = kwargs["should_accept_path"]
-
-                # Deserialize episode_data (may contain tensors)
-                episode_data = deserialize_value(episode_data)
-
-                # Dynamic import workflow
-                module_path, class_name = workflow_path.rsplit(".", 1)
-                module = importlib.import_module(module_path)
-                workflow_class = getattr(module, class_name)
-                logger.info(f"Imported workflow class: {workflow_path}")
-
-                # Instantiate workflow
-                workflow_kwargs = deserialize_value(workflow_kwargs)
-                workflow = workflow_class(**workflow_kwargs)
-                logger.info(f"Workflow '{workflow_path}' instantiated successfully")
-
-                should_accept = None
-                if should_accept_path is not None:
-                    # Dynamic import filtering function
-                    module_path, fn_name = should_accept_path.rsplit(".", 1)
-                    module = importlib.import_module(module_path)
-                    should_accept = getattr(module, fn_name)
-                    logger.info(f"Imported filtering function: {should_accept_path}")
-
-                args = []
-                kwargs = dict(
-                    data=episode_data,
-                    workflow=workflow,
-                    should_accept=should_accept,
-                )
+            if isinstance(_engine, InferenceEngine):
+                handler = _METHOD_HANDLERS.get(method_name)
+                if handler is not None:
+                    args, kwargs = handler(_engine, args, kwargs)
         except Exception as e:
             logger.error(
                 f"Workflow data conversion failed: {e}\n{traceback.format_exc()}"
@@ -261,7 +280,9 @@ def call_engine_method():
 
             # HACK: handle update weights future
             if isinstance(result, Future):
+                logger.info("Waiting for update weights future")
                 result = result.result()
+                logger.info("Update weights future done")
 
             # Serialize result (convert tensors to SerializedTensor dicts)
             serialized_result = serialize_value(result)
@@ -296,7 +317,11 @@ def export_stats():
             return jsonify({"error": "Engine not initialized"}), 503
 
         # TrainEngine: reduce stats across data_parallel_group
-        assert isinstance(_engine, TrainEngine)
+        if not isinstance(_engine, TrainEngine):
+            return (
+                jsonify({"error": "/export_stats is only available for TrainEngine"}),
+                400,
+            )
         result = stats_tracker.export(reduce_group=_engine.data_parallel_group)
         return jsonify({"status": "success", "result": result})
 

areal/scheduler/rpc/serialization.py

@@ -25,10 +25,14 @@
 import torch
 from pydantic import BaseModel, Field
 
+from areal.utils import logging
+
 TOKENIZER_ARCHIVE_INLINE_THRESHOLD = 512 * 1024
 TOKENIZER_ZSTD_THRESHOLD = 20 * 1024 * 1024
 TokenizerCompression = Literal["zip", "zstd"]
 
+logger = logging.getLogger("SyncRPCServer")
+
 
 class SerializedTensor(BaseModel):
     """Pydantic model for serialized tensor with metadata.
@@ -96,9 +100,6 @@ def to_tensor(self) -> torch.Tensor:
         dtype_str = self.dtype.replace("torch.", "")
         dtype = getattr(torch, dtype_str)
 
-        # Reconstruct tensor from bytes
-        import numpy as np
-
         np_array = np.frombuffer(buffer, dtype=self._torch_dtype_to_numpy(dtype))
         # Copy the array to make it writable before converting to tensor
         np_array = np_array.copy()
@@ -465,33 +466,44 @@ def deserialize_value(value: Any) -> Any:
                 }
                 # Reconstruct the dataclass instance
                 return dataclass_type(**deserialized_data)
-            except Exception:
+            except Exception as e:
                 # If parsing fails, treat as regular dict
+                logger.warning(
+                    f"Failed to deserialize dataclass, treating as regular dict: {e}"
+                )
                 pass
 
         # Check for SerializedTokenizer marker
         if value.get("type") == "tokenizer":
             try:
                 serialized_tokenizer = SerializedTokenizer.model_validate(value)
                 return serialized_tokenizer.to_tokenizer()
-            except Exception:
+            except Exception as e:
+                logger.warning(
+                    f"Failed to deserialize tokenizer, treating as regular dict: {e}"
+                )
                 pass
 
         # Check for SerializedNDArray marker
         if value.get("type") == "ndarray":
             try:
                 serialized_array = SerializedNDArray.model_validate(value)
                 return serialized_array.to_array()
-            except Exception:
+            except Exception as e:
+                logger.warning(
+                    f"Failed to deserialize ndarray, treating as regular dict: {e}"
+                )
                 pass
 
         # Check for SerializedTensor marker
         if value.get("type") == "tensor":
             try:
                 serialized_tensor = SerializedTensor.model_validate(value)
                 return serialized_tensor.to_tensor()
-            except Exception:
-                # If parsing fails, treat as regular dict
+            except Exception as e:
+                logger.warning(
+                    f"Failed to deserialize tensor, treating as regular dict: {e}"
+                )
                 pass
 
         # Regular dict - recursively deserialize values