decodingtrust-agent-sdk 0.1.0__py3-none-any.whl

utils/task_executor.py

@@ -0,0 +1,904 @@
+import asyncio
+import grp
+import json
+import os
+import time
+from collections import defaultdict
+from dataclasses import dataclass, field
+from enum import Enum
+from pathlib import Path
+from typing import Any, Callable, Coroutine, Dict, FrozenSet, List, Optional, Set, Tuple
+
+import yaml
+
+from .memory_guard import check_memory_before_launch
+from .reset_helpers import reset_environment
+
+
+def _needs_sudo_for_docker() -> bool:
+    """Check if we need sudo to run docker commands."""
+    import subprocess
+    # First, try running docker directly
+    try:
+        result = subprocess.run(
+            ["docker", "ps"],
+            capture_output=True,
+            timeout=5
+        )
+        if result.returncode == 0:
+            return False
+    except (subprocess.SubprocessError, FileNotFoundError):
+        pass
+    
+    # Fallback to group check
+    try:
+        docker_gid = grp.getgrnam("docker").gr_gid
+        if docker_gid in os.getgroups():
+            return False
+    except (KeyError, OSError):
+        pass
+    # Default to using sudo
+    return True
+
+
+# Cache the result since it won't change during execution
+_USE_SUDO = _needs_sudo_for_docker()
+
+
+def get_task_environments(task_dir: Path) -> List[str]:
+    """
+    Determine which environments a task needs based on its config.yaml.
+
+    Uses the 'environment' field in mcp.yaml to directly map MCP servers
+    to their required Docker environments.
+
+    Args:
+        task_dir: Path to the task directory
+
+    Returns:
+        List of environment names (e.g., ["gmail", "slack"])
+    """
+    config_path = task_dir / "config.yaml"
+    if not config_path.exists():
+        return []
+
+    try:
+        config = yaml.safe_load(config_path.read_text()) or {}
+    except yaml.YAMLError as e:
+        raise ValueError(f"Failed to parse {config_path}: {e}")
+
+    agent_cfg = config.get("Agent", {})
+    mcp_servers = agent_cfg.get("mcp_servers", []) or config.get("mcp_servers", [])
+
+    server_names = [
+        srv.get("name", "").lower()
+        for srv in mcp_servers
+        if srv.get("enabled", True)
+    ]
+
+    project_root = Path(__file__).resolve().parent.parent
+    mcp_config_path = project_root / "dt_arena" / "config" / "mcp.yaml"
+
+    if not mcp_config_path.exists():
+        return server_names
+
+    try:
+        mcp_cfg = yaml.safe_load(mcp_config_path.read_text()) or {}
+    except yaml.YAMLError as e:
+        raise ValueError(f"Failed to parse {mcp_config_path}: {e}")
+
+    mcp_servers_cfg = {srv["name"].lower(): srv for srv in mcp_cfg.get("servers", [])}
+
+    needed_envs: Set[str] = set()
+    for srv_name in server_names:
+        srv_cfg = mcp_servers_cfg.get(srv_name, {})
+        env_value = srv_cfg.get("environment")
+        if env_value:
+            # Support both single string and list of environments
+            if isinstance(env_value, list):
+                needed_envs.update(env_value)
+            else:
+                needed_envs.add(env_value)
+
+    return list(needed_envs)
+
+
+class EnvState(Enum):
+    """Environment instance lifecycle states."""
+    STARTING = "starting"
+    AVAILABLE = "available"
+    IN_USE = "in_use"
+    STOPPING = "stopping"
+
+
+@dataclass
+class EnvInstance:
+    """A running instance of a Docker environment."""
+    instance_id: str  # Unique ID like "gmail:pool_gmail_1_12345"
+    env_name: str  # Environment type like "gmail", "slack"
+    project_name: str  # Docker compose project name
+    compose_file: Path
+    state: EnvState = EnvState.STARTING
+    ports: Dict[str, int] = field(default_factory=dict)
+    current_task_id: Optional[str] = None
+    use_count: int = 0
+    last_used: float = field(default_factory=time.time)
+
+
+@dataclass
+class ScheduledTask:
+    """A task with its environment requirements."""
+    task_dir: Path
+    environments: FrozenSet[str]
+    original_index: int
+    domain: Optional[str] = None
+    task_type: Optional[str] = None
+    threat_model: Optional[str] = None
+    risk_category: Optional[str] = None
+    task_id: Optional[str] = None
+
+
+@dataclass
+class RunningTask:
+    """A task currently being executed."""
+    task: ScheduledTask
+    task_id: str
+    instances: List[EnvInstance]
+    start_time: float = field(default_factory=time.time)
+    future: Optional[asyncio.Future] = None
+
+
+@dataclass
+class ExecutorStats:
+    """Statistics about the executor state."""
+    total_tasks: int
+    pending_tasks: int
+    running_tasks: int
+    completed_tasks: int
+    total_instances: int
+    available_instances: int
+    in_use_instances: int
+    instances_by_env: Dict[str, int]
+
+
+class TaskExecutor:
+    """
+    Task-based parallel executor with environment instance pooling.
+
+    Key features:
+    - Controls parallelism by max_parallel_tasks (not max environments)
+    - Creates environment instances on demand
+    - Reuses available instances when possible
+    - Stops instances when no pending task needs them
+    - Respects per-environment instance limits (e.g., max 1 Windows VM)
+    """
+
+    def __init__(self, max_parallel: int = 5):
+        self.max_parallel = max_parallel
+        self._lock = asyncio.Lock()
+
+        # Task tracking
+        self._pending: List[ScheduledTask] = []
+        self._running: Dict[str, RunningTask] = {}  # task_id -> RunningTask
+        self._completed: List[Tuple[ScheduledTask, int]] = []  # (task, return_code)
+
+        # Environment instance tracking
+        self._instances: Dict[str, EnvInstance] = {}  # instance_id -> EnvInstance
+        self._env_instances: Dict[str, List[str]] = defaultdict(list)  # env_name -> [instance_ids]
+        self._task_instances: Dict[str, List[str]] = {}  # task_id -> [instance_ids]
+
+        # Configuration
+        self._env_config: Dict = {}
+        self._env_limits: Dict[str, int] = {}  # env_name -> max_instances
+        self._default_max_instances: Optional[int] = None
+        self._project_counter = 0
+        self._pool_id = str(os.getpid())
+
+        # Load configuration
+        self._load_config()
+
+    def _load_config(self) -> None:
+        """Load environment configuration including instance limits."""
+        project_root = Path(__file__).resolve().parent.parent
+        env_config_path = project_root / "dt_arena" / "config" / "env.yaml"
+
+        if not env_config_path.exists():
+            return
+
+        try:
+            self._env_config = yaml.safe_load(env_config_path.read_text()) or {}
+        except Exception as e:
+            print(f"[EXECUTOR] Warning: Failed to load env config: {e}")
+            return
+
+        # Load default max instances
+        self._default_max_instances = self._env_config.get("default_max_instances")
+
+        # Load per-environment limits
+        environments = self._env_config.get("environments", {})
+        for env_name, env_def in environments.items():
+            if "max_instances" in env_def:
+                self._env_limits[env_name] = env_def["max_instances"]
+
+    def _get_max_instances(self, env_name: str) -> Optional[int]:
+        """Get max instances for an environment (None = unlimited)."""
+        if env_name in self._env_limits:
+            return self._env_limits[env_name]
+        return self._default_max_instances
+
+    def _get_compose_file(self, env_name: str) -> Optional[Path]:
+        """Get the docker-compose file path for an environment."""
+        project_root = Path(__file__).resolve().parent.parent
+        environments = self._env_config.get("environments", {})
+
+        if env_name in environments:
+            compose_rel = environments[env_name].get("docker_compose")
+            if compose_rel:
+                return (project_root / compose_rel).resolve()
+
+        # Fallback: check standard location
+        env_path = project_root / "dt_arena" / "envs" / env_name
+        for name in ["docker-compose.yml", "docker-compose.yaml"]:
+            compose = env_path / name
+            if compose.exists():
+                return compose
+
+        return None
+
+    def _generate_project_name(self, env_name: str) -> str:
+        """Generate unique project name for a new instance."""
+        self._project_counter += 1
+        return f"pool_{env_name}_{self._project_counter}_{self._pool_id}"
+
+    def _allocate_ports(self, env_name: str) -> Dict[str, int]:
+        """Allocate ports for a new instance."""
+        from .resource_manager import ResourceManager
+
+        mgr = ResourceManager.instance()
+        environments = self._env_config.get("environments", {})
+        env_def = environments.get(env_name, {})
+        ports_cfg = env_def.get("ports", {})
+
+        allocated = {}
+        container_id = f"pool_{env_name}_{self._project_counter}"
+
+        for var_name, meta in ports_cfg.items():
+            default = int(meta.get("default", meta.get("container_port", 8000)))
+            port = mgr.allocate_port(container_id, var_name, default=default)
+            allocated[var_name] = port
+
+        return allocated
+
+    def _validate_env_requirements(self, compose_file: Path) -> bool:
+        """
+        Validate environment-specific requirements before starting.
+
+        Looks for a validate.py in the environment directory and calls its
+        validate(env_dir) function if present.
+
+        Returns:
+            True if validation passes, False otherwise
+        """
+        import importlib.util
+
+        env_dir = compose_file.parent
+        validate_script = env_dir / "validate.py"
+
+        if not validate_script.exists():
+            return True
+
+        try:
+            spec = importlib.util.spec_from_file_location("validate", validate_script)
+            if spec is None or spec.loader is None:
+                return True
+
+            module = importlib.util.module_from_spec(spec)
+            spec.loader.exec_module(module)
+
+            if hasattr(module, "validate"):
+                success, message = module.validate(env_dir)
+                if not success:
+                    print(message, flush=True)
+                    return False
+
+        except Exception as e:
+            print(f"[EXECUTOR] Warning: Failed to run {validate_script}: {e}", flush=True)
+
+        return True
+
+    async def _start_instance(self, env_name: str, max_retries: int = 5, wait_time: int = 30) -> Optional[EnvInstance]:
+        """Start a new environment instance."""
+        # Memory guard: wait for sufficient memory before launching
+        for attempt in range(max_retries):
+            try:
+                check_memory_before_launch()
+                break
+            except RuntimeError as e:
+                if attempt == max_retries - 1:
+                    print(f"[EXECUTOR] Memory guard blocked launch after {max_retries} retries: {e}", flush=True)
+                    return None
+                wait = wait_time * (attempt + 1)
+                print(f"[EXECUTOR] Low memory, waiting {wait}s before retry ({attempt+1}/{max_retries})...", flush=True)
+                await asyncio.sleep(wait)
+
+        compose_file = self._get_compose_file(env_name)
+        if not compose_file or not compose_file.exists():
+            print(f"[EXECUTOR] No docker-compose file found for {env_name}", flush=True)
+            return None
+
+        # Pre-start validation (calls validate.py if present in env directory)
+        if not self._validate_env_requirements(compose_file):
+            return None
+
+        project_name = self._generate_project_name(env_name)
+        instance_id = f"{env_name}:{project_name}"
+        ports = self._allocate_ports(env_name)
+
+        instance = EnvInstance(
+            instance_id=instance_id,
+            env_name=env_name,
+            project_name=project_name,
+            compose_file=compose_file,
+            state=EnvState.STARTING,
+            ports=ports,
+        )
+
+        print(f"[EXECUTOR] Starting instance {instance_id}", flush=True)
+        
+        # Track the new instance
+        self._instances[instance.instance_id] = instance
+        self._env_instances[env_name].append(instance.instance_id)
+
+        try:
+            # Build docker compose command
+            if _USE_SUDO:
+                # Use sudo with 'env' command to ensure env vars are properly passed
+                cmd = ["sudo", "env"] + [f"{k}={v}" for k, v in ports.items()] + [
+                    "docker", "compose", "-p", project_name, "-f", str(compose_file), "up", "-d"
+                ]
+                proc = await asyncio.create_subprocess_exec(
+                    *cmd,
+                    cwd=str(compose_file.parent),
+                    stdout=asyncio.subprocess.PIPE,
+                    stderr=asyncio.subprocess.PIPE,
+                )
+            else:
+                # Set port environment variables when not using sudo
+                env = os.environ.copy()
+                for var_name, port in ports.items():
+                    env[var_name] = str(port)
+
+                proc = await asyncio.create_subprocess_exec(
+                    "docker", "compose", "-p", project_name, "-f", str(compose_file), "up", "-d",
+                    cwd=str(compose_file.parent),
+                    env=env,
+                    stdout=asyncio.subprocess.PIPE,
+                    stderr=asyncio.subprocess.PIPE,
+                )
+
+            stdout, stderr = await proc.communicate()
+            if proc.returncode != 0:
+                print(f"[EXECUTOR] Start failed for {instance_id}: {stderr.decode()}", flush=True)
+                return None
+
+            # Wait for containers to be healthy
+            # Get timeout from env config (default 120s, Windows needs 240s+)
+            environments = self._env_config.get("environments", {})
+            env_def = environments.get(env_name, {})
+            health_timeout = env_def.get("health_timeout", 120)
+            await self._wait_for_healthy(project_name, compose_file, timeout=health_timeout)
+
+            instance.state = EnvState.AVAILABLE
+            print(f"[EXECUTOR] Instance {instance_id} started successfully", flush=True)
+            return instance
+
+        except Exception as e:
+            print(f"[EXECUTOR] Failed to start instance {instance_id}: {e}", flush=True)
+            return None
+
+    async def _wait_for_healthy(
+        self,
+        project_name: str,
+        compose_file: Path,
+        timeout: int = 120,
+    ) -> bool:
+        """Wait for all containers in the project to be healthy."""
+        print(f"[EXECUTOR] Waiting for containers to be healthy...", flush=True)
+        start_time = time.time()
+
+        while time.time() - start_time < timeout:
+            # Check container health status
+            cmd = ["docker", "compose", "-p", project_name, "-f", str(compose_file), "ps", "--format", "json"]
+            if _USE_SUDO:
+                cmd = ["sudo"] + cmd
+
+            proc = await asyncio.create_subprocess_exec(
+                *cmd,
+                cwd=str(compose_file.parent),
+                stdout=asyncio.subprocess.PIPE,
+                stderr=asyncio.subprocess.PIPE,
+            )
+            stdout, stderr = await proc.communicate()
+
+            if proc.returncode != 0:
+                await asyncio.sleep(2)
+                continue
+
+            # Parse JSON output (one JSON object per line)
+            all_healthy = True
+            output = stdout.decode().strip()
+            if not output:
+                await asyncio.sleep(2)
+                continue
+
+            for line in output.split('\n'):
+                if not line.strip():
+                    continue
+                try:
+                    container = json.loads(line)
+                    health = container.get("Health", "")
+                    state = container.get("State", "")
+                    # Container is ready if: no healthcheck (health empty) and running, OR healthy
+                    if state != "running":
+                        all_healthy = False
+                        break
+                    if health and health != "healthy":
+                        all_healthy = False
+                        break
+                except json.JSONDecodeError:
+                    continue
+
+            if all_healthy:
+                print(f"[EXECUTOR] All containers healthy", flush=True)
+                return True
+
+            await asyncio.sleep(2)
+
+        print(f"[EXECUTOR] Timeout waiting for containers to be healthy", flush=True)
+        return False
+
+    async def _stop_instance(self, instance: EnvInstance) -> bool:
+        """Stop and remove an environment instance."""
+        instance.state = EnvState.STOPPING
+        print(f"[EXECUTOR] Stopping instance {instance.instance_id}", flush=True)
+
+        try:
+            cmd = ["docker", "compose", "-p", instance.project_name,
+                   "-f", str(instance.compose_file), "down", "--remove-orphans", "--volumes"]
+            if _USE_SUDO:
+                cmd = ["sudo"] + cmd
+
+            proc = await asyncio.create_subprocess_exec(
+                *cmd,
+                cwd=str(instance.compose_file.parent),
+                stdout=asyncio.subprocess.PIPE,
+                stderr=asyncio.subprocess.PIPE,
+            )
+            await asyncio.wait_for(proc.wait(), timeout=60)
+
+            # Remove from tracking
+            if instance.instance_id in self._instances:
+                del self._instances[instance.instance_id]
+            if instance.instance_id in self._env_instances.get(instance.env_name, []):
+                self._env_instances[instance.env_name].remove(instance.instance_id)
+
+            print(f"[EXECUTOR] Instance {instance.instance_id} stopped", flush=True)
+            return True
+
+        except Exception as e:
+            print(f"[EXECUTOR] Error stopping instance {instance.instance_id}: {e}", flush=True)
+            return False
+
+    def _get_disable_reuse_flag(self, env_name: str) -> bool:
+        """Get disable reuse flag from env config."""
+        environments = self._env_config.get("environments", {})
+        env_def = environments.get(env_name, {})
+        return env_def.get("disable_reuse", False)
+
+    def _get_reset_scripts(self, env_name: str) -> Dict[str, str]:
+        """Get reset script paths for each service from env config."""
+        environments = self._env_config.get("environments", {})
+        env_def = environments.get(env_name, {})
+        return env_def.get("reset_scripts", {})
+
+    def _get_reset_endpoints(self, env_name: str) -> Dict[str, Dict[str, Any]]:
+        """
+        Get reset API endpoints from env config.
+
+        Returns a dict mapping endpoint names to their config:
+        {
+            "endpoint_name": {
+                "url": "http://localhost:${PORT}/api/v1/reset",
+                "method": "POST",  # optional, defaults to POST
+                "port_var": "SALESFORCE_API_PORT"  # which port variable to use
+            }
+        }
+        """
+        environments = self._env_config.get("environments", {})
+        env_def = environments.get(env_name, {})
+        return env_def.get("reset_endpoints", {})
+
+    async def _reset_instance(self, instance: EnvInstance) -> None:
+        """
+        Reset an instance's data state.
+
+        Delegates to reset_environment from reset_helpers module which handles:
+        1. API endpoints (reset_endpoints) - preferred, safer
+        2. Docker exec scripts (reset_scripts) - fallback
+
+        Raises:
+            RuntimeError: If reset fails (both endpoints and scripts failed)
+        """
+        print(f"[EXECUTOR] Resetting instance {instance.instance_id}", flush=True)
+
+        # Read per-environment script timeout (e.g. Windows loadvm needs ~60s)
+        environments = self._env_config.get("environments", {})
+        env_def = environments.get(instance.env_name, {})
+        script_timeout = env_def.get("reset_script_timeout", 60)
+
+        await reset_environment(
+            env_name=instance.env_name,
+            ports=instance.ports,
+            env_config=self._env_config,
+            project_name=instance.project_name,
+            compose_file=instance.compose_file,
+            script_timeout=script_timeout,
+        )
+
+    def _count_env_instances(self, env_name: str, exclude_stopping: bool = True) -> int:
+        """Count running instances of an environment type."""
+        count = 0
+        for inst_id in self._env_instances.get(env_name, []):
+            inst = self._instances.get(inst_id)
+            if inst:
+                if exclude_stopping and inst.state == EnvState.STOPPING:
+                    continue
+                count += 1
+        return count
+
+    def _find_available_instance(self, env_name: str) -> Optional[EnvInstance]:
+        """Find an available instance of the given environment type."""
+        for inst_id in self._env_instances.get(env_name, []):
+            inst = self._instances.get(inst_id)
+            if inst and inst.state == EnvState.AVAILABLE:
+                return inst
+        return None
+
+    def _get_needed_envs(self) -> Set[str]:
+        """Get all environment types needed by pending tasks."""
+        needed = set()
+        for task in self._pending:
+            needed.update(task.environments)
+        return needed
+
+    def _can_start_task(self, task: ScheduledTask) -> bool:
+        """Check if a task can start given current env limits."""
+        for env_name in task.environments:
+            # Check if we have an available instance
+            if self._find_available_instance(env_name):
+                continue
+
+            # Check if we can create a new instance
+            max_inst = self._get_max_instances(env_name)
+            if max_inst is not None:
+                current_count = self._count_env_instances(env_name)
+                if current_count >= max_inst:
+                    return False
+
+        return True
+
+    def _calculate_task_priority(
+        self,
+        task: ScheduledTask,
+        available_envs: Set[str],
+    ) -> Tuple[int, int]:
+        """
+        Calculate priority for a task.
+
+        Returns (negative_reuse_count, original_index) for sorting.
+        Lower values = higher priority.
+        """
+        # Count how many environments can be reused
+        reusable = sum(1 for env in task.environments if env in available_envs)
+        return (-reusable, task.original_index)
+
+    def _pick_next_task(self) -> Optional[ScheduledTask]:
+        """Pick the best pending task to run next."""
+        if not self._pending:
+            return None
+
+        # Get currently available environment types
+        available_envs = set()
+        for inst in self._instances.values():
+            if inst.state == EnvState.AVAILABLE:
+                available_envs.add(inst.env_name)
+
+        # Filter to tasks that can actually start
+        candidates = [t for t in self._pending if self._can_start_task(t)]
+        if not candidates:
+            return None
+
+        # Sort by priority (max reuse, then FIFO)
+        candidates.sort(key=lambda t: self._calculate_task_priority(t, available_envs))
+
+        chosen = candidates[0]
+        self._pending.remove(chosen)
+        return chosen
+
+    async def _acquire_instances_for_task(
+        self,
+        task: ScheduledTask,
+        task_id: str,
+    ) -> Optional[List[EnvInstance]]:
+        """Acquire all required instances for a task."""
+        acquired: List[EnvInstance] = []
+
+        for env_name in task.environments:
+            # Try to find an available instance
+            instance = self._find_available_instance(env_name)
+            disable_reuse = self._get_disable_reuse_flag(env_name)
+            if instance and disable_reuse:
+                print(f"[EXECUTOR] Reuse disabled for {env_name}, stopping instance {instance.instance_id}", flush=True)
+                await self._stop_instance(instance)
+                instance = None
+
+            if instance:
+                # Reuse existing instance
+                instance.state = EnvState.IN_USE
+                instance.current_task_id = task_id
+                instance.use_count += 1
+                instance.last_used = time.time()
+                acquired.append(instance)
+
+                # Track early so release works even if reset fails
+                self._task_instances[task_id] = [inst.instance_id for inst in acquired]
+
+                # Reset before reuse (with configurable retries)
+                environments = self._env_config.get("environments", {})
+                env_def = environments.get(env_name, {})
+                max_retries = env_def.get("reset_retries", 1)
+                retry_delay = env_def.get("reset_retry_delay", 10)
+
+                for attempt in range(1 + max_retries):
+                    try:
+                        await self._reset_instance(instance)
+                        break
+                    except Exception as e:
+                        if attempt < max_retries:
+                            print(f"[EXECUTOR] Reset failed for {instance.instance_id} (attempt {attempt + 1}/{1 + max_retries}), retrying in {retry_delay}s: {e}", flush=True)
+                            await asyncio.sleep(retry_delay)
+                        else:
+                            raise
+
+                print(f"[EXECUTOR] Reusing instance {instance.instance_id} for task {task_id} (use #{instance.use_count})", flush=True)
+            else:
+                # Start new instance (release lock during slow operation)
+                self._lock.release()
+                try:
+                    instance = await self._start_instance(env_name)
+                finally:
+                    await self._lock.acquire()
+
+                if not instance:
+                    # Rollback acquired instances
+                    for inst in acquired:
+                        inst.state = EnvState.AVAILABLE
+                        inst.current_task_id = None
+                    return None
+
+                instance.state = EnvState.IN_USE
+                instance.current_task_id = task_id
+                instance.use_count = 1
+                acquired.append(instance)
+
+                # Track early so release works even if reset fails
+                self._task_instances[task_id] = [inst.instance_id for inst in acquired]
+
+                # Reset after start
+                await self._reset_instance(instance)
+
+                print(f"[EXECUTOR] Started new instance {instance.instance_id} for task {task_id}", flush=True)
+
+        # Final update of tracking (in case multiple environments)
+        self._task_instances[task_id] = [inst.instance_id for inst in acquired]
+
+        return acquired
+
+    async def _release_instances_for_task(self, task_id: str) -> Set[str]:
+        """Release instances used by a task. Returns released env names."""
+        instance_ids = self._task_instances.pop(task_id, [])
+        released_envs = set()
+
+        for inst_id in instance_ids:
+            inst = self._instances.get(inst_id)
+            if inst and inst.state == EnvState.IN_USE:
+                inst.state = EnvState.AVAILABLE
+                inst.current_task_id = None
+                inst.last_used = time.time()
+                released_envs.add(inst.env_name)
+                print(f"[EXECUTOR] Released instance {inst_id}", flush=True)
+
+        return released_envs
+
+    async def _cleanup_unused_instances(self, released_envs: Set[str]) -> None:
+        """Stop instances for environments no longer needed by pending tasks."""
+        needed_envs = self._get_needed_envs()
+
+        for env_name in released_envs:
+            if env_name not in needed_envs:
+                # Stop all available instances of this env type
+                for inst_id in list(self._env_instances.get(env_name, [])):
+                    inst = self._instances.get(inst_id)
+                    if inst and inst.state == EnvState.AVAILABLE:
+                        # Release lock during slow stop operation
+                        self._lock.release()
+                        try:
+                            await self._stop_instance(inst)
+                        finally:
+                            await self._lock.acquire()
+
+    async def _run_single_task(
+        self,
+        task: ScheduledTask,
+        task_id: str,
+        run_fn: Callable[[ScheduledTask, Dict[str, EnvInstance]], Coroutine[Any, Any, int]],
+    ) -> int:
+        """Run a single task and handle completion."""
+        instances: Optional[List[EnvInstance]] = None
+
+        try:
+            # Acquire instances
+            async with self._lock:
+                instances = await self._acquire_instances_for_task(task, task_id)
+
+            if not instances:
+                print(f"[EXECUTOR] Failed to acquire instances for task {task_id}", flush=True)
+                return 1
+
+            # Build env_name -> instance mapping for the task
+            env_instances = {inst.env_name: inst for inst in instances}
+
+            # Run the task
+            return_code = await run_fn(task, env_instances)
+            return return_code
+
+        except Exception as e:
+            print(f"[EXECUTOR] Error running task {task_id}: {e}", flush=True)
+            return 1
+
+        finally:
+            # Release instances and cleanup
+            async with self._lock:
+                released_envs = await self._release_instances_for_task(task_id)
+                await self._cleanup_unused_instances(released_envs)
+
+    async def _worker(
+        self,
+        run_fn: Callable[[ScheduledTask, Dict[str, EnvInstance]], Coroutine[Any, Any, int]],
+        results: List[Tuple[ScheduledTask, int]],
+        slot_available: asyncio.Event,
+    ) -> None:
+        """Worker that processes tasks from the pending queue."""
+        while True:
+            task: Optional[ScheduledTask] = None
+            task_id: Optional[str] = None
+
+            async with self._lock:
+                # Check if we should exit
+                if not self._pending and not self._running:
+                    return
+
+                # Try to pick a task
+                if len(self._running) < self.max_parallel:
+                    task = self._pick_next_task()
+                    if task:
+                        task_id = f"{task.task_dir.name}_{id(task)}"
+                        self._running[task_id] = RunningTask(
+                            task=task,
+                            task_id=task_id,
+                            instances=[],
+                        )
+
+            if task and task_id:
+                # Run the task outside lock
+                return_code = await self._run_single_task(task, task_id, run_fn)
+
+                async with self._lock:
+                    # Record result
+                    results.append((task, return_code))
+                    del self._running[task_id]
+
+                    # Signal that a slot is available
+                    slot_available.set()
+            else:
+                # Wait for a slot to become available
+                slot_available.clear()
+                await slot_available.wait()
+
+    async def run_all(
+        self,
+        tasks: List[ScheduledTask],
+        run_fn: Callable[[ScheduledTask, Dict[str, EnvInstance]], Coroutine[Any, Any, int]],
+    ) -> List[Tuple[ScheduledTask, int]]:
+        """
+        Run all tasks with optimal parallelism and environment reuse.
+
+        Args:
+            tasks: List of tasks to run
+            run_fn: Async function that takes (task, env_instances) and returns exit code
+
+        Returns:
+            List of (task, return_code) tuples
+        """
+        if not tasks:
+            return []
+
+        # Initialize pending queue (maintain original order)
+        self._pending = list(tasks)
+        self._running.clear()
+        results: List[Tuple[ScheduledTask, int]] = []
+
+        print(f"[EXECUTOR] Starting {len(tasks)} tasks with max_parallel={self.max_parallel}", flush=True)
+
+        # Create worker coordination
+        slot_available = asyncio.Event()
+        slot_available.set()  # Start with slots available
+
+        # Create workers
+        workers = [
+            asyncio.create_task(self._worker(run_fn, results, slot_available))
+            for _ in range(self.max_parallel)
+        ]
+
+        # Wait for all workers to complete
+        await asyncio.gather(*workers)
+
+        return results
+
+    def stats(self) -> ExecutorStats:
+        """Get current executor statistics."""
+        instances_by_env: Dict[str, int] = defaultdict(int)
+        available = in_use = 0
+
+        for inst in self._instances.values():
+            instances_by_env[inst.env_name] += 1
+            if inst.state == EnvState.AVAILABLE:
+                available += 1
+            elif inst.state == EnvState.IN_USE:
+                in_use += 1
+
+        return ExecutorStats(
+            total_tasks=len(self._pending) + len(self._running) + len(self._completed),
+            pending_tasks=len(self._pending),
+            running_tasks=len(self._running),
+            completed_tasks=len(self._completed),
+            total_instances=len(self._instances),
+            available_instances=available,
+            in_use_instances=in_use,
+            instances_by_env=dict(instances_by_env),
+        )
+
+    async def shutdown(self) -> None:
+        """Shutdown all instances."""
+        async with self._lock:
+            print(f"[EXECUTOR] Shutting down {len(self._instances)} instance(s)...", flush=True)
+
+            # Stop all instances concurrently
+            stop_tasks = []
+            for inst in list(self._instances.values()):
+                if inst.state != EnvState.STOPPING:
+                    stop_tasks.append(self._stop_instance(inst))
+
+            if stop_tasks:
+                # Release lock during shutdown
+                self._lock.release()
+                try:
+                    await asyncio.gather(*stop_tasks, return_exceptions=True)
+                finally:
+                    await self._lock.acquire()
+
+            self._instances.clear()
+            self._env_instances.clear()
+            self._task_instances.clear()
+
+            print("[EXECUTOR] Shutdown complete", flush=True)