wehub-resource-sync
266 lines
11 KiB
Python
266 lines
11 KiB
Python
"""
|
||
环境池 - 管理多个并行环境实例 (async version)。
|
||
|
||
支持三种隔离级别：
|
||
- pages: 多 Page，共享 Context（最轻量，适用于 namespace 隔离的模拟器）
|
||
- contexts: 多 Context，共享 Browser（中等隔离）
|
||
- browsers: 多 Browser 进程（完全隔离）
|
||
|
||
额外支持 num_browsers 参数：
|
||
当 num_browsers > 1 且 isolation=pages 时，创建多个 Browser 进程，
|
||
将 N 个 page 均匀分配到各 browser/context 中，兼顾轻量和多核利用。
|
||
"""
|
||
|
||
from __future__ import annotations
|
||
|
||
import asyncio
|
||
import math
|
||
import os
|
||
from enum import Enum
|
||
from typing import TYPE_CHECKING, Iterator
|
||
|
||
from bench_env.logger import get_logger
|
||
|
||
if TYPE_CHECKING:
|
||
from bench_env.env.mobile_gym import MobileGymEnv
|
||
|
||
logger = get_logger(__name__)
|
||
|
||
|
||
class Isolation(str, Enum):
|
||
"""并行隔离级别"""
|
||
PAGES = "pages"
|
||
CONTEXTS = "contexts"
|
||
BROWSERS = "browsers"
|
||
|
||
|
||
class EnvPool:
|
||
"""
|
||
环境池 - 管理多个 MobileGymEnv 实例 (async)。
|
||
|
||
Usage:
|
||
async with EnvPool(url, n=4, isolation="browsers") as pool:
|
||
await asyncio.gather(*[
|
||
run_task(pool[i], tasks[i])
|
||
for i in range(len(pool))
|
||
])
|
||
"""
|
||
|
||
def __init__(
|
||
self,
|
||
url: str,
|
||
n: int = 1,
|
||
isolation: Isolation | str = Isolation.PAGES,
|
||
num_browsers: int = 0,
|
||
headless: bool = True,
|
||
proxy: str | None = None,
|
||
coord_space: str = "norm_0_1000",
|
||
delay_after_action: float = 0.8,
|
||
viewport_size: tuple[int, int] = (360, 800),
|
||
physical_size: tuple[int, int] = (1080, 2400),
|
||
device_scale_factor: float = 3,
|
||
verbose: bool = True,
|
||
):
|
||
self.url = url
|
||
self.n = max(1, n)
|
||
self.isolation = Isolation(isolation)
|
||
self.num_browsers = max(0, num_browsers)
|
||
self.headless = headless
|
||
self.proxy = proxy
|
||
self.coord_space = coord_space
|
||
self.delay_after_action = delay_after_action
|
||
self.viewport_size = viewport_size
|
||
self.physical_size = physical_size
|
||
self.device_scale_factor = device_scale_factor
|
||
self.verbose = verbose
|
||
|
||
self._pw = None
|
||
self._browsers: list = []
|
||
self._contexts: list = []
|
||
self._envs: list[MobileGymEnv] = []
|
||
|
||
async def __aenter__(self) -> "EnvPool":
|
||
await self._setup()
|
||
return self
|
||
|
||
async def __aexit__(self, *args):
|
||
await self._cleanup()
|
||
|
||
def __iter__(self) -> Iterator["MobileGymEnv"]:
|
||
return iter(self._envs)
|
||
|
||
def __len__(self) -> int:
|
||
return len(self._envs)
|
||
|
||
def __getitem__(self, idx: int) -> "MobileGymEnv":
|
||
return self._envs[idx]
|
||
|
||
@property
|
||
def envs(self) -> list["MobileGymEnv"]:
|
||
return self._envs
|
||
|
||
def _make_env_kwargs(self) -> dict:
|
||
"""生成 MobileGymEnv 的通用参数"""
|
||
return {
|
||
"url": self.url,
|
||
"viewport_size": self.viewport_size,
|
||
"physical_size": self.physical_size,
|
||
"device_scale_factor": self.device_scale_factor,
|
||
"headless": self.headless,
|
||
"proxy": self.proxy,
|
||
"coord_space": self.coord_space,
|
||
"delay_after_action": self.delay_after_action,
|
||
"verbose": self.verbose,
|
||
}
|
||
|
||
def _resolve_num_browsers(self) -> int:
|
||
"""Resolve effective browser count.
|
||
|
||
- num_browsers=0 (auto): pages/contexts → 1, browsers → N
|
||
- num_browsers>0: clamp to [1, N]
|
||
"""
|
||
if self.num_browsers > 0:
|
||
return min(self.num_browsers, self.n)
|
||
# Auto defaults
|
||
if self.isolation == Isolation.BROWSERS:
|
||
return self.n
|
||
return 1
|
||
|
||
async def _setup(self):
|
||
"""创建环境实例
|
||
|
||
拓扑由 (isolation, num_browsers) 共同决定：
|
||
|
||
┌─────────────┬──────────────┬──────────────────────────────┐
|
||
│ isolation │ num_browsers │ 拓扑 │
|
||
├─────────────┼──────────────┼──────────────────────────────┤
|
||
│ pages │ 1 (default) │ 1 browser, 1 context, N pages│
|
||
│ pages │ B>1 │ B browsers, B contexts, │
|
||
│ │ │ N pages 均分到各 context │
|
||
│ contexts │ 1 (default) │ 1 browser, N contexts │
|
||
│ contexts │ B>1 │ B browsers, N contexts 均分 │
|
||
│ browsers │ N (default) │ N browsers, N pages (1:1) │
|
||
└─────────────┴──────────────┴──────────────────────────────┘
|
||
"""
|
||
from playwright.async_api import async_playwright
|
||
from bench_env.env.mobile_gym import MobileGymEnv
|
||
|
||
env_kwargs = self._make_env_kwargs()
|
||
launch_args = MobileGymEnv.get_launch_args(
|
||
self.headless,
|
||
self.proxy,
|
||
browser_type="chromium",
|
||
)
|
||
|
||
nb = self._resolve_num_browsers()
|
||
|
||
if self.isolation == Isolation.BROWSERS:
|
||
# 共享 1 个 Playwright server，每个 env 拿独立 browser
|
||
# 隔离性不变（独立 Chromium 进程），但省掉 N-1 个 Node.js 进程
|
||
self._pw = await async_playwright().start()
|
||
for i in range(self.n):
|
||
browser = await self._pw.chromium.launch(**launch_args)
|
||
self._browsers.append(browser)
|
||
self._envs.append(MobileGymEnv(**env_kwargs, browser=browser, worker_id=i))
|
||
|
||
elif self.isolation == Isolation.CONTEXTS:
|
||
# 多 context（可跨多个 browser）— 每个 Env 自己创建 context + route
|
||
self._pw = await async_playwright().start()
|
||
browsers = [await self._pw.chromium.launch(**launch_args) for _ in range(nb)]
|
||
self._browsers.extend(browsers)
|
||
for i in range(self.n):
|
||
browser = browsers[i % nb]
|
||
self._envs.append(MobileGymEnv(**env_kwargs, browser=browser, worker_id=i))
|
||
|
||
else: # PAGES
|
||
self._pw = await async_playwright().start()
|
||
ctx_args = MobileGymEnv.get_context_args(
|
||
self.viewport_size[0], self.viewport_size[1], self.device_scale_factor
|
||
)
|
||
if nb <= 1:
|
||
# 原有行为：1 browser, 1 context, N pages
|
||
browser = await self._pw.chromium.launch(**launch_args)
|
||
context = await browser.new_context(**ctx_args)
|
||
await MobileGymEnv.setup_context_routes(context)
|
||
self._browsers.append(browser)
|
||
for i in range(self.n):
|
||
self._envs.append(MobileGymEnv(**env_kwargs, context=context, worker_id=i))
|
||
logger.info(
|
||
f"EnvPool: 1 browser, 1 context, {self.n} pages"
|
||
)
|
||
else:
|
||
# 多 Browser 多 Page：B browsers, 每个 browser 1 context,
|
||
# N pages 均匀分到各 context
|
||
browsers = [await self._pw.chromium.launch(**launch_args) for _ in range(nb)]
|
||
self._browsers.extend(browsers)
|
||
self._contexts = []
|
||
for browser in browsers:
|
||
ctx = await browser.new_context(**ctx_args)
|
||
await MobileGymEnv.setup_context_routes(ctx)
|
||
self._contexts.append(ctx)
|
||
for i in range(self.n):
|
||
ctx = self._contexts[i % nb]
|
||
self._envs.append(MobileGymEnv(**env_kwargs, context=ctx, worker_id=i))
|
||
logger.info(
|
||
f"EnvPool: {nb} browsers × ~{math.ceil(self.n / nb)} pages/browser "
|
||
f"= {self.n} pages total"
|
||
)
|
||
|
||
# 分批启动环境,避免连接风暴。browsers iso 高并发时还要躲 inotify 上限:
|
||
# 每个 chromium 进程会占 ~2 个 inotify instances(IDB seed 解压、文件系统层),
|
||
# host 默认 fs.inotify.max_user_instances=128(无 sudo 改不了)。
|
||
# batch_size=8 + sleep=3s 让任意时刻 in-init chromium ≤ ceil(22/3)*8 ≈ 64,
|
||
# 远低于 128。详情见 bench_env/docs/KNOWN_ISSUES.md §3。
|
||
# contexts/pages iso 的 browser 进程数通常很少,无需 3s 长 sleep。
|
||
batch_size = 8
|
||
default_batch_sleep_s = "3.0" if self.isolation == Isolation.BROWSERS and self.n >= 192 else "0.3"
|
||
batch_sleep_s = float(os.environ.get("MOBILE_GYM_POOL_BATCH_SLEEP_S", default_batch_sleep_s))
|
||
start_errors: list[Exception] = []
|
||
for batch_start in range(0, len(self._envs), batch_size):
|
||
batch = self._envs[batch_start:batch_start + batch_size]
|
||
results = await asyncio.gather(
|
||
*[env.start() for env in batch],
|
||
return_exceptions=True,
|
||
)
|
||
start_errors.extend(r for r in results if isinstance(r, Exception))
|
||
if batch_start + batch_size < len(self._envs):
|
||
await asyncio.sleep(batch_sleep_s)
|
||
if start_errors:
|
||
logger.error(
|
||
f"EnvPool start errors: {len(start_errors)}/{len(self._envs)}. "
|
||
f"First error: {type(start_errors[0]).__name__}: {start_errors[0]}"
|
||
)
|
||
|
||
async def _cleanup(self):
|
||
"""清理资源"""
|
||
# 关闭所有环境（page）
|
||
for i, env in enumerate(self._envs):
|
||
try:
|
||
await env.close()
|
||
except Exception as e:
|
||
logger.debug(f"EnvPool: env[{i}].close() failed: {type(e).__name__}: {e}")
|
||
self._envs.clear()
|
||
|
||
# 关闭池创建的 context（多 browser PAGES 模式下）
|
||
for i, ctx in enumerate(self._contexts):
|
||
try:
|
||
await ctx.close()
|
||
except Exception as e:
|
||
logger.debug(f"EnvPool: context[{i}].close() failed: {type(e).__name__}: {e}")
|
||
self._contexts.clear()
|
||
|
||
# 关闭共享的 browser（如果有）
|
||
for i, browser in enumerate(self._browsers):
|
||
try:
|
||
await browser.close()
|
||
except Exception as e:
|
||
logger.debug(f"EnvPool: browser[{i}].close() failed: {type(e).__name__}: {e}")
|
||
self._browsers.clear()
|
||
|
||
if self._pw:
|
||
try:
|
||
await self._pw.stop()
|
||
except Exception as e:
|
||
logger.debug(f"EnvPool: playwright.stop() failed: {type(e).__name__}: {e}")
|
||
self._pw = None
|