Merge pull request #22 from deeppavlov/f/concurrency

F/concurrency

Author: voorhs

src/mcp_evals/_internal/runner/_domain_runner.py

@@ -4,6 +4,7 @@
 from pathlib import Path
 from typing import Any
 
+from loguru import logger
 from pydantic_ai.agent import Agent
 from pydantic_ai.usage import UsageLimits
 from pydantic_evals.reporting import EvaluationReport
@@ -28,6 +29,7 @@ def __init__(  # noqa: PLR0913
         deps_maker: DepsMaker | None = None,
         *,
         max_tasks: int | None = None,
+        max_concurrency: int = 1,
         use_self_correction: bool = False,
         max_self_correction_retries: int = 3,
         start_training: TrainingTestingCallback | None = None,
@@ -42,6 +44,10 @@ def __init__(  # noqa: PLR0913
         self.agent = agent
         self.deps_maker = deps_maker
         self.max_tasks = max_tasks
+        if max_concurrency < 1:
+            msg = "max_concurrency must be >= 1"
+            raise ValueError(msg)
+        self.max_concurrency = max_concurrency
         self.run_result_processor = run_result_processor
         self.usage_limits = usage_limits
         self.grouper = grouper
@@ -56,6 +62,15 @@ def __init__(  # noqa: PLR0913
 
     async def run(self, domain: Domain[Any], experiment_name: str) -> EvaluationReport:
         deps_maker = self.deps_maker or default_deps_maker()
+        if self.max_concurrency > 1 and not domain.supports_concurrency:
+            msg = (
+                f"Domain '{domain.name}' does not support concurrency "
+                f"(supports_concurrency={domain.supports_concurrency}). "
+                "Set max_concurrency=1, or set the domain flag to True after making tasks parallel-safe."
+            )
+            raise ValueError(msg)
+        if self.max_concurrency > 1 and domain.supports_concurrency:
+            logger.debug(f"[{domain.name}] Running with max_concurrency={self.max_concurrency}")
 
         async with domain:
             if self.use_self_correction:
@@ -147,7 +162,7 @@ async def _run_train_phase(
         train_dataset = tasks_to_dataset(train_tasks)
         await train_dataset.evaluate(
             evaluated_fn,
-            max_concurrency=1,
+            max_concurrency=self.max_concurrency,
             case_context_manager=make_task_lifecycle(state, split_idx, "train"),
             progress=False,
             name=f"{base_name}_train_{split_idx}_",
@@ -175,7 +190,7 @@ async def _run_test_phase(
         test_dataset = tasks_to_dataset(test_tasks)
         return await test_dataset.evaluate(
             evaluated_fn,
-            max_concurrency=1,
+            max_concurrency=self.max_concurrency,
             case_context_manager=make_task_lifecycle(state, split_idx, "test"),
             progress=False,
             name=f"{base_name}_test_{split_idx}" if experiment_name else None,

src/mcp_evals/_internal/runner/_run_state.py

@@ -100,6 +100,7 @@ def __init__(self, path: AnyioPath | str) -> None:
         self._split_finished: set[int] = set()
         self._task_finished: set[TaskFinishedKey] = set()
         self._header_written = False
+        self._write_lock = anyio.Lock()
 
     @classmethod
     async def load(
@@ -192,12 +193,13 @@ async def mark_task_finished(self, split_idx: int, phase: Phase, task_name: str)
         self._task_finished.add(TaskFinishedKey(split_idx, phase, task_name))
 
     async def _append_event(self, event: RunStateEvent) -> None:
-        await self._ensure_header()
-        line = event.model_dump_json(exclude_none=True)
-        async with await anyio.open_file(self._path, "a") as f:
-            await f.write(line + "\n")
+        async with self._write_lock:
+            await self._ensure_header_unlocked()
+            line = event.model_dump_json(exclude_none=True)
+            async with await anyio.open_file(self._path, "a") as f:
+                await f.write(line + "\n")
 
-    async def _ensure_header(self) -> None:
+    async def _ensure_header_unlocked(self) -> None:
         if self._header_written:
             return
         if self._n_tasks is None or self._fingerprint is None:
@@ -211,8 +213,9 @@ async def _ensure_header(self) -> None:
 
     async def clear(self) -> None:
         """Remove the state file if it exists (e.g. after a full run)."""
-        with contextlib.suppress(FileNotFoundError):
-            await self._path.unlink()
+        async with self._write_lock:
+            with contextlib.suppress(FileNotFoundError):
+                await self._path.unlink()
 
 
 async def run_state_path(

src/mcp_evals/contrib/filesystem/domain.py

@@ -7,7 +7,7 @@
 
 import aiofiles
 from loguru import logger
-from pydantic_ai.mcp import MCPServerStdio
+from pydantic_ai.mcp import MCPServer
 
 from mcp_evals import Domain
 from mcp_evals.contrib.filesystem.utils import Fixture
@@ -51,10 +51,14 @@
 class FilesystemDomain(Domain[DomainSecrets]):
     """Domain for filesystem tasks from MCP Universe.
 
-    Provides MCP filesystem server and groups related filesystem tasks.
+    Provides a shared temp root directory and groups related filesystem tasks.
+
+    Each task owns its own MCP filesystem server + workspace directory to allow
+    safe parallel execution.
     """
 
     name = "filesystem"
+    supports_concurrency = True
 
     def __init__(self, tool_retries: int = 1) -> None:
         """Init."""
@@ -66,25 +70,12 @@ async def setup(self, stack: AsyncExitStack[Any]) -> None:
         tmpdir_ctx = aiofiles.tempfile.TemporaryDirectory(prefix="mcp-filesystem-")
         self._tmp_dir = Path(await stack.enter_async_context(tmpdir_ctx))
 
-    def mcp_servers(self) -> Sequence[MCPServerStdio]:
-        """Return MCP filesystem server configuration."""
-        return [
-            MCPServerStdio(
-                "docker",
-                [
-                    "run",
-                    "-i",
-                    "--rm",
-                    "--mount",
-                    f"type=bind,src={self._tmp_dir},dst=/projects",
-                    "-w",
-                    "/projects",
-                    "mcp-filesystem-server:mcp-evals",
-                    "/projects",
-                ],
-                max_retries=self.tool_retries,
-            )
-        ]
+    def mcp_servers(self) -> Sequence[MCPServer]:
+        """Return domain-scoped MCP servers (none for filesystem).
+
+        Filesystem tasks run with per-task MCP servers to support parallelism.
+        """
+        return []
 
     def tasks(self) -> Sequence[FilesystemTask]:
         """Return all filesystem tasks."""

src/mcp_evals/contrib/filesystem/task.py

@@ -1,10 +1,12 @@
 """Base class for all filesystem tasks."""
 
+from collections.abc import Sequence
 from contextlib import AsyncExitStack
 from pathlib import Path
 from typing import Any
 
 from pydantic import BaseModel, Field
+from pydantic_ai.mcp import MCPServerStdio
 
 from mcp_evals.contrib.filesystem.utils import Fixture, download_fixture, prepare_workspace
 from mcp_evals.secrets import TaskSecrets
@@ -26,9 +28,31 @@ def __init__(self, work_dir: Path, fixture: Fixture, tool_retries: int = 1) -> N
         """Init."""
         super().__init__(tool_retries=tool_retries)
 
-        self.work_dir = work_dir
+        self.root_dir = work_dir
+        # Per-task workspace (enables safe parallel execution)
+        self.work_dir = work_dir / self.name
         self.fixture = fixture
 
+    def mcp_servers(self) -> Sequence[MCPServerStdio]:
+        """Return task-scoped MCP filesystem server configuration."""
+        return [
+            MCPServerStdio(
+                "docker",
+                [
+                    "run",
+                    "-i",
+                    "--rm",
+                    "--mount",
+                    f"type=bind,src={self.work_dir},dst=/projects",
+                    "-w",
+                    "/projects",
+                    "mcp-filesystem-server:mcp-evals",
+                    "/projects",
+                ],
+                max_retries=self.tool_retries,
+            )
+        ]
+
     async def setup(self, stack: AsyncExitStack[Any]) -> None:
         """Set up the task environment."""
         # Download fixture

src/mcp_evals/contrib/filesystem/utils.py

@@ -8,6 +8,7 @@
 from enum import StrEnum
 from pathlib import Path
 
+import anyio
 from dotenv import load_dotenv
 from loguru import logger
 
@@ -69,6 +70,16 @@ class Fixture(StrEnum):
 
 load_dotenv()
 
+_fixture_locks: dict[Fixture, anyio.Lock] = {}
+
+
+def _fixture_lock(category: Fixture) -> anyio.Lock:
+    lock = _fixture_locks.get(category)
+    if lock is None:
+        lock = anyio.Lock()
+        _fixture_locks[category] = lock
+    return lock
+
 
 async def download_fixture(category: Fixture) -> Path:
     """Download and cache a filesystem test fixture.
@@ -97,69 +108,70 @@ async def download_fixture(category: Fixture) -> Path:
     cache_dir = Path(user_cache_dir("mcp-evals", "mcp-evals")) / "fixtures"
     fixture_path = cache_dir / category
 
-    # Return cached fixture if it exists
-    if fixture_path.exists() and fixture_path.is_dir():
-        logger.debug(f"Using cached fixture '{category.value}'")
-        return fixture_path
+    async with _fixture_lock(category):
+        # Return cached fixture if it exists
+        if fixture_path.exists() and fixture_path.is_dir():
+            logger.debug(f"Using cached fixture '{category.value}'")
+            return fixture_path
 
-    # Download fixture
-    logger.debug(f"Downloading fixture '{category.value}'")
-    url = FIXTURE_URL_MAPPING[category]
-    zip_path = cache_dir / f"{category}.zip"
+        # Download fixture
+        logger.debug(f"Downloading fixture '{category.value}'")
+        url = FIXTURE_URL_MAPPING[category]
+        zip_path = cache_dir / f"{category}.zip"
 
-    # Ensure cache directory exists
-    cache_dir.mkdir(parents=True, exist_ok=True)
+        # Ensure cache directory exists
+        cache_dir.mkdir(parents=True, exist_ok=True)
 
-    try:
-        # Download using httpx with streaming
-        timeout = httpx.Timeout(connect=5.0, read=5.0, write=10.0, pool=5.0)
-        proxy_url = os.getenv("DOWNLOAD_PROXY")
-        async with (
-            httpx.AsyncClient(timeout=timeout, proxy=proxy_url) as client,
-            client.stream("GET", url, follow_redirects=True) as response,
-        ):
-            response.raise_for_status()
-            total_size = int(response.headers.get("content-length", 0)) or None
-            async with aiofiles.open(zip_path, "wb") as f:
-                with tqdm(
-                    total=total_size,
-                    unit="B",
-                    unit_scale=True,
-                    unit_divisor=1024,
-                    desc=f"Downloading {category}",
-                ) as pbar:
-                    async for chunk in response.aiter_bytes():
-                        await f.write(chunk)
-                        pbar.update(len(chunk))
-
-        # Extract ZIP file
-        with zipfile.ZipFile(zip_path) as zip_file:
-            zip_file.extractall(cache_dir)
-
-        # Clean up macOS metadata if present
-        macosx_path = cache_dir / "__MACOSX"
-        if macosx_path.exists():
-            shutil.rmtree(macosx_path)
-
-        # Clean up ZIP file
-        zip_path.unlink(missing_ok=True)
-
-    except httpx.HTTPError as e:
-        msg = f"Failed to download fixture from {url}: {e}"
-        raise RuntimeError(msg) from e
-    except zipfile.BadZipFile as e:
-        msg = f"Invalid ZIP file for category {category}: {e}"
-        raise RuntimeError(msg) from e
-    except Exception as e:
-        msg = f"Failed to download or extract fixture for category {category}: {e}"
-        raise RuntimeError(msg) from e
+        try:
+            # Download using httpx with streaming
+            timeout = httpx.Timeout(connect=5.0, read=5.0, write=10.0, pool=5.0)
+            proxy_url = os.getenv("DOWNLOAD_PROXY")
+            async with (
+                httpx.AsyncClient(timeout=timeout, proxy=proxy_url) as client,
+                client.stream("GET", url, follow_redirects=True) as response,
+            ):
+                response.raise_for_status()
+                total_size = int(response.headers.get("content-length", 0)) or None
+                async with aiofiles.open(zip_path, "wb") as f:
+                    with tqdm(
+                        total=total_size,
+                        unit="B",
+                        unit_scale=True,
+                        unit_divisor=1024,
+                        desc=f"Downloading {category}",
+                    ) as pbar:
+                        async for chunk in response.aiter_bytes():
+                            await f.write(chunk)
+                            pbar.update(len(chunk))
+
+            # Extract ZIP file
+            with zipfile.ZipFile(zip_path) as zip_file:
+                zip_file.extractall(cache_dir)
+
+            # Clean up macOS metadata if present
+            macosx_path = cache_dir / "__MACOSX"
+            if macosx_path.exists():
+                shutil.rmtree(macosx_path)
+
+            # Clean up ZIP file
+            zip_path.unlink(missing_ok=True)
+
+        except httpx.HTTPError as e:
+            msg = f"Failed to download fixture from {url}: {e}"
+            raise RuntimeError(msg) from e
+        except zipfile.BadZipFile as e:
+            msg = f"Invalid ZIP file for category {category}: {e}"
+            raise RuntimeError(msg) from e
+        except Exception as e:
+            msg = f"Failed to download or extract fixture for category {category}: {e}"
+            raise RuntimeError(msg) from e
 
-    # Verify extraction
-    if not fixture_path.exists():
-        msg = f"Extracted directory not found: {fixture_path}"
-        raise RuntimeError(msg)
+        # Verify extraction
+        if not fixture_path.exists():
+            msg = f"Extracted directory not found: {fixture_path}"
+            raise RuntimeError(msg)
 
-    return fixture_path
+        return fixture_path
 
 
 @asynccontextmanager

src/mcp_evals/contrib/postgres/domain.py

@@ -50,6 +50,7 @@ class PostgresDomain(Domain[PgConfig]):
 
     name = "postgres"
     secrets_type = PgConfig
+    supports_concurrency = True
 
     _container: DockerContainer | None = None
     _docker: aiodocker.Docker | None = None