benchmax 0.1.2.dev19__tar.gz → 0.1.2.dev21__tar.gz

{benchmax-0.1.2.dev19 → benchmax-0.1.2.dev21}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: benchmax
-Version: 0.1.2.dev19
+Version: 0.1.2.dev21
 Summary: Framework-Agnostic RL Environments for LLM Fine-Tuning
 Author: cgft.io
 Classifier: Programming Language :: Python :: 3

{benchmax-0.1.2.dev19 → benchmax-0.1.2.dev21}/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "benchmax"
-version = "0.1.2.dev19"
+version = "0.1.2.dev21"
 description = "Framework-Agnostic RL Environments for LLM Fine-Tuning"
 readme = "README.md"
 authors = [{ name = "cgft.io" }]

{benchmax-0.1.2.dev19 → benchmax-0.1.2.dev21}/src/benchmax/bundle/bundler.py

@@ -70,9 +70,18 @@ def bundle_env(
                 f"{mod.__name__}"
             )
 
-    # --- Serialize the class ---
+    # --- Serialize the class AND constructor_args while modules are registered ---
+    pickled_constructor_args: bytes | None = None
     try:
         pickled_class = cloudpickle.dumps(env_class)
+        # Pickle constructor_args now (while local_modules are registered)
+        # so non-JSON-serializable objects get pickled by value.
+        if constructor_args is not None:
+            try:
+                import json
+                json.dumps(constructor_args)
+            except TypeError:
+                pickled_constructor_args = cloudpickle.dumps(constructor_args)
     except Exception as e:
         raise BundlingError(
             f"Failed to serialize {env_class.__name__} with cloudpickle: {e}"
@@ -99,8 +108,13 @@ def bundle_env(
         python_version=python_version,
         benchmax_version=benchmax_version,
         constructor_args=constructor_args,
+        constructor_args_pickled=pickled_constructor_args is not None,
+    )
+    payload = BundledEnv(
+        pickled_class=pickled_class,
+        metadata=metadata,
+        pickled_constructor_args=pickled_constructor_args,
     )
-    payload = BundledEnv(pickled_class=pickled_class, metadata=metadata)
 
     size_kb = len(pickled_class) / 1024
     logger.info(
@@ -136,10 +150,37 @@ def write_bundle_files(
     pickle_path.write_bytes(bundle.pickled_class)
     metadata_path.write_bytes(bundle.metadata.to_json_bytes())
 
+    # Write pre-pickled constructor_args (pickled during bundling while
+    # local_modules were registered, so module code is inline).
+    if bundle.pickled_constructor_args is not None:
+        args_path = pickle_path.with_suffix(".args.pkl")
+        args_path.write_bytes(bundle.pickled_constructor_args)
+        logger.info("[bundling] Wrote pickled constructor_args to %s", args_path)
+
 
 def read_bundle_files(pickle_path: Path, metadata_path: Path) -> BundledEnv:
     pickle_path = Path(pickle_path)
     metadata_path = Path(metadata_path)
     pickled_class = pickle_path.read_bytes()
     metadata = BundleMetadata.from_json_bytes(metadata_path.read_bytes())
+
+    # Load pickled constructor_args if they were serialized separately.
+    if metadata.constructor_args_pickled:
+        args_pickle_path = pickle_path.with_suffix(".args.pkl")
+        if args_pickle_path.exists():
+            constructor_args = cloudpickle.loads(args_pickle_path.read_bytes())
+            metadata = BundleMetadata(
+                pip_dependencies=metadata.pip_dependencies,
+                python_version=metadata.python_version,
+                benchmax_version=metadata.benchmax_version,
+                constructor_args=constructor_args,
+                format_version=metadata.format_version,
+            )
+            logger.info("[bundling] Loaded pickled constructor_args from %s", args_pickle_path)
+        else:
+            logger.warning(
+                "[bundling] Metadata indicates pickled constructor_args but %s not found",
+                args_pickle_path,
+            )
+
     return BundledEnv(pickled_class=pickled_class, metadata=metadata)

{benchmax-0.1.2.dev19 → benchmax-0.1.2.dev21}/src/benchmax/bundle/loader.py

@@ -150,6 +150,20 @@ def load_env_from_files(
     metadata_path = Path(metadata_path)
     pickled_class = pickle_path.read_bytes()
     metadata = BundleMetadata.from_json_bytes(metadata_path.read_bytes())
+
+    # If constructor_args were pickled separately, load them.
+    if constructor_args is None and metadata.constructor_args_pickled:
+        args_pickle_path = pickle_path.with_suffix(".args.pkl")
+        if args_pickle_path.exists():
+            import cloudpickle
+            constructor_args = cloudpickle.loads(args_pickle_path.read_bytes())
+            logger.info("[bundling] Loaded pickled constructor_args from %s", args_pickle_path)
+        else:
+            logger.warning(
+                "[bundling] Metadata indicates pickled constructor_args but %s not found",
+                args_pickle_path,
+            )
+
     return load_env(
         pickled_class,
         metadata,

{benchmax-0.1.2.dev19 → benchmax-0.1.2.dev21}/src/benchmax/bundle/payload.py

@@ -13,6 +13,7 @@ class BundleMetadata:
     python_version: str
     benchmax_version: str
     constructor_args: Optional[Dict[str, Any]] = None
+    constructor_args_pickled: bool = False
     format_version: int = FORMAT_VERSION
 
     def to_dict(self) -> Dict[str, Any]:
@@ -21,11 +22,36 @@ class BundleMetadata:
             "python_version": self.python_version,
             "benchmax_version": self.benchmax_version,
             "constructor_args": self.constructor_args,
+            "constructor_args_pickled": self.constructor_args_pickled,
             "format_version": self.format_version,
         }
 
     def to_json_bytes(self) -> bytes:
-        return json.dumps(self.to_dict()).encode("utf-8")
+        data = self.to_dict()
+        # constructor_args may contain non-JSON-serializable objects
+        # (e.g. SearchClient instances). Null them out and set the flag
+        # so the loader knows to look for the separate pickled args file.
+        if self.has_pickled_constructor_args():
+            data["constructor_args"] = None
+            data["constructor_args_pickled"] = True
+        return json.dumps(data).encode("utf-8")
+
+    def has_pickled_constructor_args(self) -> bool:
+        """True if constructor_args need separate pickle serialization."""
+        if self.constructor_args is None:
+            return False
+        try:
+            json.dumps(self.constructor_args)
+            return False
+        except TypeError:
+            return True
+
+    def pickled_constructor_args_bytes(self) -> bytes | None:
+        """Pickle constructor_args if they aren't JSON-serializable."""
+        if not self.has_pickled_constructor_args():
+            return None
+        import cloudpickle
+        return cloudpickle.dumps(self.constructor_args)
 
     @classmethod
     def from_dict(cls, data: Dict[str, Any]) -> "BundleMetadata":
@@ -34,6 +60,7 @@ class BundleMetadata:
             python_version=data.get("python_version", "unknown"),
             benchmax_version=data.get("benchmax_version", "unknown"),
             constructor_args=data.get("constructor_args"),
+            constructor_args_pickled=bool(data.get("constructor_args_pickled", False)),
             format_version=int(data.get("format_version", FORMAT_VERSION)),
         )
 
@@ -48,3 +75,4 @@ class BundledEnv:
 
     pickled_class: bytes
     metadata: BundleMetadata
+    pickled_constructor_args: Optional[bytes] = None

{benchmax-0.1.2.dev19 → benchmax-0.1.2.dev21}/src/benchmax/bundle/validator.py

@@ -35,7 +35,8 @@ def validate_bundle(
     if constructor_args is None:
         constructor_args = bundle.metadata.constructor_args
     return _run_isolated_validation(
-        bundle.pickled_class, bundle.metadata, constructor_args
+        bundle.pickled_class, bundle.metadata, constructor_args,
+        pickled_constructor_args=bundle.pickled_constructor_args,
     )
 
 
@@ -110,6 +111,7 @@ def _run_isolated_validation(
     pickled_class: bytes,
     metadata: BundleMetadata,
     constructor_args: Optional[Dict[str, Any]],
+    pickled_constructor_args: Optional[bytes] = None,
 ) -> List[str]:
     """Create a temp venv, install deps, unpickle the bundle, and smoke test.
 
@@ -192,18 +194,30 @@ def _run_isolated_validation(
         with open(pickle_path, "wb") as f:
             f.write(pickled_class)
 
-        # 5. Write and run smoke test script using cloudpickle directly
-        constructor_args_json = json.dumps(constructor_args) if constructor_args else "null"
+        # 5. Write constructor args as pickle (handles non-JSON-serializable objects)
+        args_pickle_path = os.path.join(venv_dir, "constructor_args.pkl")
+        if pickled_constructor_args is not None:
+            # Use pre-pickled args (pickled during bundling with modules registered)
+            with open(args_pickle_path, "wb") as f:
+                f.write(pickled_constructor_args)
+        elif constructor_args is not None:
+            import cloudpickle as _cp
+            with open(args_pickle_path, "wb") as f:
+                _cp.dump(constructor_args, f)
+
+        # 6. Write and run smoke test script using cloudpickle directly
         smoke_script = textwrap.dedent(f"""\
-            import json
+            import os
             import asyncio
             import cloudpickle
 
             with open({pickle_path!r}, "rb") as f:
                 env_class = cloudpickle.load(f)
 
-            constructor_args = json.loads({constructor_args_json!r})
-            if constructor_args is not None:
+            args_path = {args_pickle_path!r}
+            if os.path.exists(args_path):
+                with open(args_path, "rb") as f:
+                    constructor_args = cloudpickle.load(f)
                 instance = env_class(**constructor_args)
                 tools = asyncio.run(instance.list_tools())
                 asyncio.run(instance.shutdown())

{benchmax-0.1.2.dev19 → benchmax-0.1.2.dev21}/src/benchmax.egg-info/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: benchmax
-Version: 0.1.2.dev19
+Version: 0.1.2.dev21
 Summary: Framework-Agnostic RL Environments for LLM Fine-Tuning
 Author: cgft.io
 Classifier: Programming Language :: Python :: 3