benchmax 0.1.2.dev22__tar.gz → 0.1.2.dev23__tar.gz

{benchmax-0.1.2.dev22 → benchmax-0.1.2.dev23}/PKG-INFO

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

{benchmax-0.1.2.dev22 → benchmax-0.1.2.dev23}/pyproject.toml

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

{benchmax-0.1.2.dev22 → benchmax-0.1.2.dev23}/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,12 +150,11 @@ 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:
+    # 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(pickled_args)
+        args_path.write_bytes(bundle.pickled_constructor_args)
         logger.info("[bundling] Wrote pickled constructor_args to %s", args_path)
 
 

{benchmax-0.1.2.dev22 → benchmax-0.1.2.dev23}/src/benchmax/bundle/payload.py

@@ -75,3 +75,4 @@ class BundledEnv:
 
     pickled_class: bytes
     metadata: BundleMetadata
+    pickled_constructor_args: Optional[bytes] = None

{benchmax-0.1.2.dev22 → benchmax-0.1.2.dev23}/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.dev22 → benchmax-0.1.2.dev23}/src/benchmax.egg-info/PKG-INFO

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