benchmax 0.1.2.dev21__tar.gz → 0.1.2.dev22__tar.gz

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

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

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

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

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

@@ -70,18 +70,9 @@ def bundle_env(
                 f"{mod.__name__}"
             )
 
-    # --- Serialize the class AND constructor_args while modules are registered ---
-    pickled_constructor_args: bytes | None = None
+    # --- Serialize the class ---
     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}"
@@ -108,13 +99,8 @@ 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(
@@ -150,11 +136,12 @@ 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:
+    # 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(bundle.pickled_constructor_args)
+        args_path.write_bytes(pickled_args)
         logger.info("[bundling] Wrote pickled constructor_args to %s", args_path)
 
 

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

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

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

@@ -35,8 +35,7 @@ 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,
-        pickled_constructor_args=bundle.pickled_constructor_args,
+        bundle.pickled_class, bundle.metadata, constructor_args
     )
 
 
@@ -111,7 +110,6 @@ 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.
 
@@ -194,30 +192,18 @@ def _run_isolated_validation(
         with open(pickle_path, "wb") as f:
             f.write(pickled_class)
 
-        # 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
+        # 5. Write and run smoke test script using cloudpickle directly
+        constructor_args_json = json.dumps(constructor_args) if constructor_args else "null"
         smoke_script = textwrap.dedent(f"""\
-            import os
+            import json
             import asyncio
             import cloudpickle
 
             with open({pickle_path!r}, "rb") as f:
                 env_class = cloudpickle.load(f)
 
-            args_path = {args_pickle_path!r}
-            if os.path.exists(args_path):
-                with open(args_path, "rb") as f:
-                    constructor_args = cloudpickle.load(f)
+            constructor_args = json.loads({constructor_args_json!r})
+            if constructor_args is not None:
                 instance = env_class(**constructor_args)
                 tools = asyncio.run(instance.list_tools())
                 asyncio.run(instance.shutdown())

{benchmax-0.1.2.dev21 → benchmax-0.1.2.dev22}/src/benchmax/envs/base_env.py

@@ -38,6 +38,15 @@ class BaseEnv(ABC):
 
         self.compute_reward = _tracked
 
+        cls_compute_group_reward = type(self).compute_group_reward
+
+        @wraps(cls_compute_group_reward)
+        async def _tracked_group(*args, **kwargs):
+            with tracking_context(self._tracking_config):
+                return await cls_compute_group_reward(self, *args, **kwargs)
+
+        self.compute_group_reward = _tracked_group
+
     def get_tracking_config(self) -> Optional[TrackingConfig]:
         return self._tracking_config
 

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

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