benchmax 0.1.2.dev11__tar.gz → 0.1.2.dev13__tar.gz

{benchmax-0.1.2.dev11 → benchmax-0.1.2.dev13}/PKG-INFO

@@ -1,20 +1,24 @@
 Metadata-Version: 2.4
 Name: benchmax
-Version: 0.1.2.dev11
+Version: 0.1.2.dev13
 Summary: Framework-Agnostic RL Environments for LLM Fine-Tuning
 Author: cgft.io
 Classifier: Programming Language :: Python :: 3
 Classifier: Operating System :: OS Independent
-Requires-Python: ==3.12.*
+Requires-Python: >=3.12
 Description-Content-Type: text/markdown
 License-File: LICENSE
 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: fastmcp~=2.12.0
-Requires-Dist: pyjwt>=2.10.1
-Requires-Dist: skypilot~=0.8.1
+Provides-Extra: mcp
+Requires-Dist: fastmcp~=2.12.0; extra == "mcp"
+Requires-Dist: pyjwt>=2.10.1; extra == "mcp"
+Provides-Extra: skypilot
+Requires-Dist: skypilot[aws,gcp]~=0.8.1; extra == "skypilot"
+Requires-Dist: pip>=25.3; extra == "skypilot"
+Requires-Dist: msrestazure>=0.6.4.post1; extra == "skypilot"
 Dynamic: license-file
 
 <picture>

{benchmax-0.1.2.dev11 → benchmax-0.1.2.dev13}/pyproject.toml

@@ -1,18 +1,15 @@
 [project]
 name = "benchmax"
-version = "0.1.2.dev11"
+version = "0.1.2.dev13"
 description = "Framework-Agnostic RL Environments for LLM Fine-Tuning"
 readme = "README.md"
 authors = [{ name = "cgft.io" }]
-requires-python = "==3.12.*"
+requires-python = ">=3.12"
 dependencies = [
     "aiohttp>=3.13.1",
     "asyncio>=4.0.0",
     "cloudpickle>=3.0.0",
     "datasets>=4.0.0",
-    "fastmcp~=2.12.0",
-    "pyjwt>=2.10.1",
-    "skypilot~=0.8.1",
 ]
 classifiers = [
     "Programming Language :: Python :: 3",
@@ -26,6 +23,17 @@ build-backend = "setuptools.build_meta"
 [tool.setuptools.packages.find]
 where = ["src"]
 
+[project.optional-dependencies]
+mcp = [
+    "fastmcp~=2.12.0",
+    "pyjwt>=2.10.1",
+]
+skypilot = [
+    "skypilot[aws,gcp]~=0.8.1",
+    "pip>=25.3",
+    "msrestazure>=0.6.4.post1",
+]
+
 [dependency-groups]
 dev = [
     "pytest>=8.4.2",
@@ -33,8 +41,12 @@ dev = [
     "python-dotenv>=1.2.1",
     "ruff>=0.14.2",
 ]
+mcp = [
+    "fastmcp~=2.12.0",
+    "pyjwt>=2.10.1",
+]
 skypilot = [
-    "skypilot[aws,gcp,azure]~=0.8.1", # Change this to your cloud provider
+    "skypilot[aws,gcp]~=0.8.1",       # Add azure only in an Azure-specific env/group
     "pip>=25.3",                      # Added as needed for skypilot launch
     "msrestazure>=0.6.4.post1",
 ]
@@ -53,9 +65,6 @@ crm = ["python-dateutil>=2.9.0.post0", "simple-salesforce>=1.12.9"]
 [tool.uv]
 conflicts = [[{ group = "skypilot" }, { group = "skyrl" }]]
 
-[tool.uv.pip]
-extra = ["dev", "skypilot", "skyrl", "excel", "excel-mac-windows", "crm"]
-
 # [tool.uv.extra-build-dependencies]
 # flash-attn = [{ requirement = "torch", match-runtime = true }]
 

{benchmax-0.1.2.dev11 → benchmax-0.1.2.dev13}/src/benchmax/adapters/skyrl/benchmax_data_process.py

@@ -7,14 +7,15 @@ import logging
 from importlib import import_module
 from pathlib import Path
 from types import ModuleType
-from typing import Type
-from datasets import Dataset, DatasetDict, IterableDataset, IterableDatasetDict
-import datasets
+from typing import TYPE_CHECKING, Type
 import asyncio
 import inspect
 
 from benchmax.envs.base_env import BaseEnv
 
+if TYPE_CHECKING:
+    from datasets import Dataset, DatasetDict, IterableDataset, IterableDatasetDict
+
 # Set logging level to WARNING and above
 logging.basicConfig(level=logging.WARNING)
 
@@ -124,6 +125,8 @@ if __name__ == "__main__":
     benchmax_cls: Type[BaseEnv] = load_class(args.env_path)
     raw_dataset, dataset_path = benchmax_cls.load_dataset(args.dataset_name)
 
+    from datasets import Dataset, DatasetDict, IterableDataset, IterableDatasetDict
+
     if isinstance(raw_dataset, (IterableDataset, IterableDatasetDict)):
         raise TypeError(
             f"Iterable datasets are currently not supported. Got {type(raw_dataset).__name__}. "
@@ -178,7 +181,9 @@ if __name__ == "__main__":
         test_dataset = processed_dataset["test"]
     else:
         if isinstance(processed_dataset, DatasetDict):
-            processed_dataset = datasets.concatenate_datasets(
+            from datasets import concatenate_datasets
+
+            processed_dataset = concatenate_datasets(
                 [ds for ds in processed_dataset.values()]
             ).shuffle(seed=42)
 

{benchmax-0.1.2.dev11 → benchmax-0.1.2.dev13}/src/benchmax/bundle/loader.py

@@ -1,26 +1,25 @@
 import logging
-import subprocess
 import sys
 from pathlib import Path
-from typing import Any, Dict, Optional, Type, Union
+from typing import TYPE_CHECKING, Any, Dict, Optional, Type, Union
 
-import cloudpickle
-
-from benchmax.envs.base_env import BaseEnv
 from benchmax.bundle.errors import (
     DependencyError,
     IncompatiblePythonError,
     IncompatibleBenchmaxError,
     BundlingError,
 )
-from benchmax.bundle.payload import BundleMetadata, BundledEnv
+
+if TYPE_CHECKING:
+    from benchmax.envs.base_env import BaseEnv
+    from benchmax.bundle.payload import BundleMetadata, BundledEnv
 
 logger = logging.getLogger(__name__)
 
 
 def load_env(
     pickled_class: bytes,
-    metadata: BundleMetadata | None = None,
+    metadata: "BundleMetadata | None" = None,
     *,
     pip_dependencies: Optional[list[str]] = None,
     python_version: Optional[str] = None,
@@ -30,7 +29,7 @@ def load_env(
     check_benchmax_version: bool = False,
     install_pip_deps: bool = False,
     instantiate: Optional[bool] = None,
-) -> BaseEnv | Type[BaseEnv]:
+) -> "BaseEnv | Type[BaseEnv]":
     """Load a packaged environment class (and optionally instantiate it).
 
     Args:
@@ -52,6 +51,8 @@ def load_env(
         The unpickled BaseEnv subclass (class object), or an instance if
         instantiation is requested.
     """
+    from benchmax.envs.base_env import BaseEnv
+
     resolved_pip_deps = (
         pip_dependencies
         if pip_dependencies is not None
@@ -101,6 +102,8 @@ def load_env(
         _install_dependencies(resolved_pip_deps)
 
     try:
+        import cloudpickle
+
         env_class = cloudpickle.loads(pickled_class)
     except Exception as e:
         raise BundlingError(
@@ -140,7 +143,9 @@ def load_env_from_files(
     check_benchmax_version: bool = False,
     install_pip_deps: bool = False,
     instantiate: Optional[bool] = None,
-) -> BaseEnv | Type[BaseEnv]:
+) -> "BaseEnv | Type[BaseEnv]":
+    from benchmax.bundle.payload import BundleMetadata
+
     pickle_path = Path(pickle_path)
     metadata_path = Path(metadata_path)
     pickled_class = pickle_path.read_bytes()
@@ -160,7 +165,7 @@ def load_env_from_files(
 
 
 def load_env_from_bundle(
-    bundle: BundledEnv,
+    bundle: "BundledEnv",
     *,
     pip_dependencies: Optional[list[str]] = None,
     python_version: Optional[str] = None,
@@ -170,7 +175,7 @@ def load_env_from_bundle(
     check_benchmax_version: bool = False,
     install_pip_deps: bool = False,
     instantiate: Optional[bool] = None,
-) -> BaseEnv | Type[BaseEnv]:
+) -> "BaseEnv | Type[BaseEnv]":
     return load_env(
         bundle.pickled_class,
         bundle.metadata,
@@ -187,6 +192,8 @@ def load_env_from_bundle(
 
 def _install_dependencies(deps: list[str]) -> None:
     """Install pip dependencies in the current environment."""
+    import subprocess
+
     logger.info(f"[bundling] Installing {len(deps)} dependencies: {deps}")
     cmd = [sys.executable, "-m", "pip", "install", "--quiet", *deps]
     result = subprocess.run(cmd, capture_output=True, text=True)

{benchmax-0.1.2.dev11 → benchmax-0.1.2.dev13}/src/benchmax/envs/base_env.py

@@ -1,27 +1,19 @@
 from abc import ABC, abstractmethod
-from typing import Dict, List, Any, Optional, Tuple
+from typing import TYPE_CHECKING, Dict, List, Any, Optional, Tuple
 from pathlib import Path
-from datasets import (
-    DatasetDict,
-    Dataset,
-    IterableDatasetDict,
-    IterableDataset,
-    load_dataset,
-)
 
 from benchmax.envs.types import ToolDefinition, StandardizedExample
 from benchmax.prompts.tools import render_tools_prompt
 
+if TYPE_CHECKING:
+    from datasets import Dataset, DatasetDict, IterableDataset, IterableDatasetDict
+
 
 class BaseEnv(ABC):
     """Base benchmax environment for tool execution and reward computation"""
 
     system_prompt: str = ""
 
-    @abstractmethod
-    async def shutdown(self):
-        pass
-
     # Override this method if your example does not match the default structure
     @classmethod
     def dataset_preprocess(cls, example: Any, **kwargs) -> StandardizedExample:
@@ -45,7 +37,7 @@ class BaseEnv(ABC):
     def load_dataset(
         cls, dataset_name: str, **kwargs
     ) -> Tuple[
-        DatasetDict | Dataset | IterableDatasetDict | IterableDataset, str | None
+        "DatasetDict | Dataset | IterableDatasetDict | IterableDataset", str | None
     ]:
         """
         Download and prepare a dataset for use with this environment.
@@ -63,6 +55,8 @@ class BaseEnv(ABC):
             Dataset: A dataset object (e.g., HuggingFace Dataset or similar) ready for processing.
             str: Optional string pointing to where the dataset is stored locally
         """
+        from datasets import load_dataset
+
         return load_dataset(dataset_name, **kwargs), None
 
     # Methods all environment subclasses must implement
@@ -98,6 +92,9 @@ class BaseEnv(ABC):
 
     # Optional rollout lifecycle management methods
 
+    async def shutdown(self):
+        pass
+
     async def init_rollout(self, rollout_id: str, **rollout_args) -> None:
         """Initialize resources for a new rollout"""
         return None

{benchmax-0.1.2.dev11 → benchmax-0.1.2.dev13}/src/benchmax/envs/crm/crm_env.py

@@ -1,6 +1,5 @@
 from pathlib import Path
-from typing import Any, Dict, List, Optional
-import sky
+from typing import TYPE_CHECKING, Any, Dict, List, Optional
 
 from benchmax.envs.mcp.parallel_mcp_env import ParallelMcpEnv
 from benchmax.envs.mcp.provisioners.base_provisioner import BaseProvisioner
@@ -8,6 +7,9 @@ from benchmax.envs.mcp.provisioners.local_provisioner import LocalProvisioner
 from benchmax.envs.mcp.provisioners.skypilot_provisioner import SkypilotProvisioner
 from benchmax.envs.types import StandardizedExample
 
+if TYPE_CHECKING:
+    import sky
+
 
 SYSTEM_PROMPT = """\
 You are an expert in Salesforce and you have access to a Salesforce instance.
@@ -85,11 +87,19 @@ class CRMEnvSkypilot(CRMEnv):
 
     def __init__(
         self,
-        cloud: sky.clouds.Cloud = sky.Azure(),
+        cloud: "sky.clouds.Cloud | None" = None,
         num_nodes: int = 2,
         servers_per_node: int = 5,
         **kwargs,
     ):
+        if cloud is None:
+            try:
+                import sky
+            except ModuleNotFoundError as e:
+                raise ModuleNotFoundError(
+                    "skypilot is required for CRMEnvSkypilot. Install with: pip install 'benchmax[skypilot]'"
+                ) from e
+            cloud = sky.Azure()
         workdir_path = Path(__file__).parent / "workdir"
         provisioner = SkypilotProvisioner(
             workdir_path=workdir_path,

{benchmax-0.1.2.dev11 → benchmax-0.1.2.dev13}/src/benchmax/envs/excel/data_utils.py

@@ -1,5 +1,4 @@
 import os
-import requests
 import tarfile
 
 
@@ -7,6 +6,8 @@ def download_and_extract(url, output_path):
     """
     Downloads a tar.gz file from the given URL and extracts it into output_path.
     """
+    import requests
+
     # Ensure the output directory exists
     os.makedirs(output_path, exist_ok=True)
 

{benchmax-0.1.2.dev11 → benchmax-0.1.2.dev13}/src/benchmax/envs/excel/excel_env.py

@@ -1,10 +1,7 @@
 import json
 import os
 from pathlib import Path
-from typing import Any, Dict, Tuple, Optional
-
-from datasets import Dataset, DatasetDict, IterableDataset, IterableDatasetDict
-import sky
+from typing import TYPE_CHECKING, Any, Dict, Tuple, Optional
 
 from benchmax.envs.mcp.parallel_mcp_env import ParallelMcpEnv
 from benchmax.envs.mcp.provisioners.base_provisioner import BaseProvisioner
@@ -16,6 +13,10 @@ from .data_utils import download_and_extract
 # Using library shared with mcp workdir
 from .workdir.excel_utils import excel_to_str_repr
 
+if TYPE_CHECKING:
+    import sky
+    from datasets import Dataset, DatasetDict, IterableDataset, IterableDatasetDict
+
 SYSTEM_PROMPT = """You are a spreadsheet expert who can manipulate spreadsheets through Python code.
 
 You need to solve the given spreadsheet manipulation question, which contains six types of information:
@@ -64,8 +65,10 @@ class ExcelEnv(ParallelMcpEnv):
         data_output_path: str = DEFAULT_DATA_OUTPUT_PATH,
         **kwargs,
     ) -> Tuple[
-        DatasetDict | Dataset | IterableDatasetDict | IterableDataset, str | None
+        "DatasetDict | Dataset | IterableDatasetDict | IterableDataset", str | None
     ]:
+        from datasets import Dataset
+
         # Currently only support spreadsheetbench dataset but can be extended to other datasets in the future
         if dataset_name == "spreadsheetbench":
             folder_path = Path(data_output_path) / SPREADSHEET_BENCH_TRAIN_DATA
@@ -206,11 +209,15 @@ class ExcelEnvSkypilot(ExcelEnv):
 
     def __init__(
         self,
-        cloud: sky.clouds.Cloud = sky.Azure(),
+        cloud: "sky.clouds.Cloud | None" = None,
         num_nodes: int = 2,
         servers_per_node: int = 5,
         **kwargs,
     ):
+        if cloud is None:
+            import sky
+
+            cloud = sky.Azure()
         workdir_path = Path(__file__).parent / "workdir"
         provisioner = SkypilotProvisioner(
             workdir_path=workdir_path,

{benchmax-0.1.2.dev11 → benchmax-0.1.2.dev13}/src/benchmax/envs/math/math_env.py

@@ -1,6 +1,5 @@
 from pathlib import Path
-from typing import Any
-import sky
+from typing import TYPE_CHECKING, Any
 
 from benchmax.envs.mcp.parallel_mcp_env import ParallelMcpEnv
 from benchmax.envs.mcp.provisioners.base_provisioner import BaseProvisioner
@@ -8,6 +7,9 @@ from benchmax.envs.mcp.provisioners.local_provisioner import LocalProvisioner
 from benchmax.envs.mcp.provisioners.skypilot_provisioner import SkypilotProvisioner
 from benchmax.envs.types import StandardizedExample
 
+if TYPE_CHECKING:
+    import sky
+
 SYSTEM_PROMPT = """Please use the tools provided to do any computation.
 Write your complete answer on the final line only, within the xml tags <answer></answer>.\n
 """
@@ -46,11 +48,19 @@ class MathEnvSkypilot(MathEnv):
 
     def __init__(
         self,
-        cloud: sky.clouds.Cloud = sky.Azure(),
+        cloud: "sky.clouds.Cloud | None" = None,
         num_nodes: int = 2,
         servers_per_node: int = 5,
         **kwargs,
     ):
+        if cloud is None:
+            try:
+                import sky
+            except ModuleNotFoundError as e:
+                raise ModuleNotFoundError(
+                    "skypilot is required for MathEnvSkypilot. Install with: pip install 'benchmax[skypilot]'"
+                ) from e
+            cloud = sky.Azure()
         workdir_path = Path(__file__).parent / "workdir"
         provisioner = SkypilotProvisioner(
             workdir_path=workdir_path,

{benchmax-0.1.2.dev11 → benchmax-0.1.2.dev13}/src/benchmax/envs/mcp/parallel_mcp_env.py

@@ -11,7 +11,12 @@ from typing import Any, Callable, Dict, List, Optional
 import warnings
 import aiohttp
 from mcp.types import TextContent
-from fastmcp.exceptions import ToolError
+try:
+    from fastmcp.exceptions import ToolError
+except ModuleNotFoundError as e:
+    raise ModuleNotFoundError(
+        "fastmcp is required for MCP environments. Install with: pip install 'benchmax[mcp]'"
+    ) from e
 
 from benchmax.envs.base_env import BaseEnv
 from benchmax.envs.types import ToolDefinition

{benchmax-0.1.2.dev11 → benchmax-0.1.2.dev13}/src/benchmax/envs/mcp/provisioners/__init__.py

@@ -5,11 +5,16 @@ Server provisioning strategies for ParallelMcpEnv.
 from .base_provisioner import BaseProvisioner
 from .manual_provisioner import ManualProvisioner
 from .local_provisioner import LocalProvisioner
-from .skypilot_provisioner import SkypilotProvisioner
+try:
+    from .skypilot_provisioner import SkypilotProvisioner
+except ModuleNotFoundError:
+    SkypilotProvisioner = None
 
 __all__ = [
     "BaseProvisioner",
     "ManualProvisioner",
     "LocalProvisioner",
-    "SkypilotProvisioner",
 ]
+
+if SkypilotProvisioner is not None:
+    __all__.append("SkypilotProvisioner")

{benchmax-0.1.2.dev11 → benchmax-0.1.2.dev13}/src/benchmax/envs/mcp/provisioners/skypilot_provisioner.py

@@ -1,3 +1,5 @@
+from __future__ import annotations
+
 """
 SkyPilot provisioner for launching cloud-based server clusters.
 """
@@ -5,14 +7,16 @@ SkyPilot provisioner for launching cloud-based server clusters.
 import logging
 import uuid
 from pathlib import Path
-from typing import List, Optional
-import sky
+from typing import TYPE_CHECKING, List, Optional
 
 from .base_provisioner import BaseProvisioner
 from .utils import get_run_command, setup_sync_dir, cleanup_dir, get_setup_command
 
 logger = logging.getLogger(__name__)
 
+if TYPE_CHECKING:
+    import sky
+
 
 class SkypilotProvisioner(BaseProvisioner):
     """
@@ -39,7 +43,7 @@ class SkypilotProvisioner(BaseProvisioner):
     def __init__(
         self,
         workdir_path: Path | str,
-        cloud: sky.clouds.Cloud,
+        cloud: "sky.clouds.Cloud",
         num_nodes: int = 1,
         servers_per_node: int = 5,
         cpus: Optional[str | int] = "2+",
@@ -62,8 +66,15 @@ class SkypilotProvisioner(BaseProvisioner):
             raise ValueError("num_nodes must be at least 1")
         if servers_per_node < 1 or servers_per_node > 100:
             raise ValueError("servers_per_node must be between 1 and 100")
+        try:
+            import sky
+        except ModuleNotFoundError as e:
+            raise ModuleNotFoundError(
+                "skypilot is required for SkypilotProvisioner. Install with: pip install 'benchmax[skypilot]'"
+            ) from e
 
         self._workdir_path = Path(workdir_path).absolute()
+        self._sky = sky
         self._cloud = cloud
         self._num_nodes = num_nodes
         self._servers_per_node = servers_per_node
@@ -123,7 +134,7 @@ class SkypilotProvisioner(BaseProvisioner):
         env = None if api_secret is None else {"API_SECRET": api_secret}
 
         # Configure SkyPilot task
-        sky_task = sky.Task(
+        sky_task = self._sky.Task(
             name="mcp-server",
             run=get_run_command(ports=all_ports),
             setup=get_setup_command(),
@@ -133,7 +144,7 @@ class SkypilotProvisioner(BaseProvisioner):
         )
 
         sky_task.set_resources(
-            sky.Resources(
+            self._sky.Resources(
                 cloud=self._cloud,
                 cpus=self._cpus,
                 memory=self._memory,
@@ -148,7 +159,7 @@ class SkypilotProvisioner(BaseProvisioner):
         )
         cluster_handle = None
         try:
-            _, handle = sky.launch(
+            _, handle = self._sky.launch(
                 task=sky_task,
                 cluster_name=self._cluster_name,
                 detach_run=True,
@@ -192,7 +203,7 @@ class SkypilotProvisioner(BaseProvisioner):
 
         logger.info(f"Tearing down SkyPilot cluster '{self._cluster_name}'...")
         try:
-            sky.down(cluster_name=self._cluster_name)
+            self._sky.down(cluster_name=self._cluster_name)
             logger.info(f"Cluster '{self._cluster_name}' torn down successfully")
         except Exception as e:
             logger.error(f"Error tearing down cluster '{self._cluster_name}': {e}")

{benchmax-0.1.2.dev11 → benchmax-0.1.2.dev13}/src/benchmax/envs/mcp/proxy_server.py

@@ -2,7 +2,6 @@ import os
 import sys
 import shutil
 import uuid
-import yaml
 import asyncio
 import argparse
 import psutil
@@ -12,10 +11,15 @@ from typing import Any, Awaitable, Callable, Dict, Union, Optional, Tuple, List
 from pathlib import Path
 from functools import wraps
 
-from fastmcp.server.middleware import Middleware, MiddlewareContext
-from fastmcp.server.auth.providers.jwt import JWTVerifier
-from fastmcp.server.auth import AccessToken
-from fastmcp import FastMCP, Client
+try:
+    from fastmcp.server.middleware import Middleware, MiddlewareContext
+    from fastmcp.server.auth.providers.jwt import JWTVerifier
+    from fastmcp.server.auth import AccessToken
+    from fastmcp import FastMCP, Client
+except ModuleNotFoundError as e:
+    raise ModuleNotFoundError(
+        "fastmcp is required for MCP environments. Install with: pip install 'benchmax[mcp]'"
+    ) from e
 from starlette.requests import Request
 from starlette.responses import PlainTextResponse, FileResponse, JSONResponse, Response
 from starlette.datastructures import UploadFile
@@ -36,6 +40,8 @@ def setup_workspace(base_dir: Path) -> Path:
 
 def load_config(config_path: Path, workspace: Path) -> Dict[str, Any]:
     """Load YAML config and inject workspace paths."""
+    import yaml
+
     with open(config_path, "r") as f:
         content = f.read().replace(
             "${{ sync_workdir }}", str(Path(__file__).resolve().parent)

{benchmax-0.1.2.dev11 → benchmax-0.1.2.dev13}/src/benchmax/envs/mcp/server_pool.py

@@ -8,7 +8,12 @@ from dataclasses import dataclass
 import random
 from typing import Coroutine, Dict, List, Optional, Set
 import aiohttp
-from fastmcp import Client as FastMCPClient
+try:
+    from fastmcp import Client as FastMCPClient
+except ModuleNotFoundError as e:
+    raise ModuleNotFoundError(
+        "fastmcp is required for MCP environments. Install with: pip install 'benchmax[mcp]'"
+    ) from e
 
 from benchmax.envs.mcp.utils import generate_jwt_token, get_auth_headers
 

{benchmax-0.1.2.dev11 → benchmax-0.1.2.dev13}/src/benchmax/envs/mcp/utils.py

@@ -9,7 +9,12 @@ from contextlib import AsyncExitStack
 from pathlib import Path
 from typing import Optional, Dict, List, Any
 import aiohttp
-from fastmcp import Client
+try:
+    from fastmcp import Client
+except ModuleNotFoundError as e:
+    raise ModuleNotFoundError(
+        "fastmcp is required for MCP environments. Install with: pip install 'benchmax[mcp]'"
+    ) from e
 from mcp import Tool
 
 from benchmax.envs.types import ToolDefinition

{benchmax-0.1.2.dev11 → benchmax-0.1.2.dev13}/src/benchmax.egg-info/PKG-INFO

@@ -1,20 +1,24 @@
 Metadata-Version: 2.4
 Name: benchmax
-Version: 0.1.2.dev11
+Version: 0.1.2.dev13
 Summary: Framework-Agnostic RL Environments for LLM Fine-Tuning
 Author: cgft.io
 Classifier: Programming Language :: Python :: 3
 Classifier: Operating System :: OS Independent
-Requires-Python: ==3.12.*
+Requires-Python: >=3.12
 Description-Content-Type: text/markdown
 License-File: LICENSE
 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: fastmcp~=2.12.0
-Requires-Dist: pyjwt>=2.10.1
-Requires-Dist: skypilot~=0.8.1
+Provides-Extra: mcp
+Requires-Dist: fastmcp~=2.12.0; extra == "mcp"
+Requires-Dist: pyjwt>=2.10.1; extra == "mcp"
+Provides-Extra: skypilot
+Requires-Dist: skypilot[aws,gcp]~=0.8.1; extra == "skypilot"
+Requires-Dist: pip>=25.3; extra == "skypilot"
+Requires-Dist: msrestazure>=0.6.4.post1; extra == "skypilot"
 Dynamic: license-file
 
 <picture>

{benchmax-0.1.2.dev11 → benchmax-0.1.2.dev13}/src/benchmax.egg-info/requires.txt

@@ -2,6 +2,12 @@ aiohttp>=3.13.1
 asyncio>=4.0.0
 cloudpickle>=3.0.0
 datasets>=4.0.0
+
+[mcp]
 fastmcp~=2.12.0
 pyjwt>=2.10.1
-skypilot~=0.8.1
+
+[skypilot]
+skypilot[aws,gcp]~=0.8.1
+pip>=25.3
+msrestazure>=0.6.4.post1