benchmax 0.1.2.dev18__tar.gz → 0.1.2.dev20__tar.gz

{benchmax-0.1.2.dev18 → benchmax-0.1.2.dev20}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: benchmax
-Version: 0.1.2.dev18
+Version: 0.1.2.dev20
 Summary: Framework-Agnostic RL Environments for LLM Fine-Tuning
 Author: cgft.io
 Classifier: Programming Language :: Python :: 3
@@ -12,7 +12,7 @@ Requires-Dist: aiohttp>=3.13.1
 Requires-Dist: asyncio>=4.0.0
 Requires-Dist: cloudpickle>=3.0.0
 Requires-Dist: datasets>=4.0.0
-Requires-Dist: expt-logger>=0.1.0.dev20
+Requires-Dist: expt-logger>=0.1.0.dev22
 Provides-Extra: mcp
 Requires-Dist: fastmcp~=2.12.0; extra == "mcp"
 Requires-Dist: pyjwt>=2.10.1; extra == "mcp"

{benchmax-0.1.2.dev18 → benchmax-0.1.2.dev20}/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "benchmax"
-version = "0.1.2.dev18"
+version = "0.1.2.dev20"
 description = "Framework-Agnostic RL Environments for LLM Fine-Tuning"
 readme = "README.md"
 authors = [{ name = "cgft.io" }]
@@ -10,7 +10,7 @@ dependencies = [
     "asyncio>=4.0.0",
     "cloudpickle>=3.0.0",
     "datasets>=4.0.0",
-    "expt-logger>=0.1.0.dev20",
+    "expt-logger>=0.1.0.dev22",
 ]
 classifiers = [
     "Programming Language :: Python :: 3",

{benchmax-0.1.2.dev18 → benchmax-0.1.2.dev20}/src/benchmax/bundle/bundler.py

@@ -136,10 +136,38 @@ def write_bundle_files(
     pickle_path.write_bytes(bundle.pickled_class)
     metadata_path.write_bytes(bundle.metadata.to_json_bytes())
 
+    # Write pickled constructor_args alongside the class pickle
+    # when they contain non-JSON-serializable objects (e.g. SearchClient).
+    pickled_args = bundle.metadata.pickled_constructor_args_bytes()
+    if pickled_args is not None:
+        args_path = pickle_path.with_suffix(".args.pkl")
+        args_path.write_bytes(pickled_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.dev18 → benchmax-0.1.2.dev20}/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.dev18 → benchmax-0.1.2.dev20}/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)),
         )
 

{benchmax-0.1.2.dev18 → benchmax-0.1.2.dev20}/src/benchmax.egg-info/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: benchmax
-Version: 0.1.2.dev18
+Version: 0.1.2.dev20
 Summary: Framework-Agnostic RL Environments for LLM Fine-Tuning
 Author: cgft.io
 Classifier: Programming Language :: Python :: 3
@@ -12,7 +12,7 @@ Requires-Dist: aiohttp>=3.13.1
 Requires-Dist: asyncio>=4.0.0
 Requires-Dist: cloudpickle>=3.0.0
 Requires-Dist: datasets>=4.0.0
-Requires-Dist: expt-logger>=0.1.0.dev20
+Requires-Dist: expt-logger>=0.1.0.dev22
 Provides-Extra: mcp
 Requires-Dist: fastmcp~=2.12.0; extra == "mcp"
 Requires-Dist: pyjwt>=2.10.1; extra == "mcp"

{benchmax-0.1.2.dev18 → benchmax-0.1.2.dev20}/src/benchmax.egg-info/requires.txt

@@ -2,7 +2,7 @@ aiohttp>=3.13.1
 asyncio>=4.0.0
 cloudpickle>=3.0.0
 datasets>=4.0.0
-expt-logger>=0.1.0.dev20
+expt-logger>=0.1.0.dev22
 
 [mcp]
 fastmcp~=2.12.0