jiang
4b47ae00d1
Full-network pressure comparisons exceeded the worker limit before saved-result previews could run. Preserve complete calculation results and keep read-time summaries and pagination. Incremental graph publication accepted isolated nodes after flow-authoring rules were lost. Require new steps to connect explicitly while retaining historical snapshots and status-only patches. Regression tests reproduce both failures on the previous code.
347 lines
13 KiB
Python
347 lines
13 KiB
Python
"""Worker boundaries use real managed subprocesses, without external databases."""
|
|
import asyncio
|
|
import json
|
|
from functools import wraps
|
|
from types import SimpleNamespace
|
|
from contextlib import asynccontextmanager
|
|
import os
|
|
from pathlib import Path
|
|
import signal
|
|
import sys
|
|
|
|
import pytest
|
|
|
|
from app.services.calculation_worker import CalculationWorker, WorkerSettings, stop_process_group
|
|
|
|
|
|
@pytest.fixture(autouse=True)
|
|
def no_external_cleanup(monkeypatch):
|
|
async def cleanup(*args):
|
|
pass
|
|
monkeypatch.setattr('app.services.calculation_worker.cleanup_job_resources', cleanup)
|
|
|
|
|
|
def run_async(function):
|
|
@wraps(function)
|
|
def wrapper(*args, **kwargs):
|
|
return asyncio.run(function(*args, **kwargs))
|
|
return wrapper
|
|
|
|
|
|
@run_async
|
|
async def test_child_keeps_complete_result_above_eight_mib(monkeypatch, tmp_path):
|
|
from app.mcp import calculation_jobs
|
|
from app.services import calculation_worker
|
|
|
|
data = {"comparisons": "x" * (8 * 1024 * 1024), "map_payload": {"features": ["coverage"]}}
|
|
|
|
@asynccontextmanager
|
|
async def scope(row):
|
|
yield object()
|
|
|
|
async def calculate(scope, row):
|
|
return SimpleNamespace(model_dump=lambda **kwargs: data)
|
|
|
|
monkeypatch.setattr(calculation_worker, "execution_scope", scope)
|
|
monkeypatch.setattr(calculation_jobs, "execute_calculation_job", calculate)
|
|
input_path, output_path = tmp_path / "input.json", tmp_path / "result.json"
|
|
input_path.write_text("{}")
|
|
await calculation_worker.execute_child(str(input_path), str(output_path))
|
|
assert output_path.stat().st_size > 8 * 1024 * 1024
|
|
assert json.loads(output_path.read_text()) == {"status": "succeeded", "result": data}
|
|
|
|
|
|
@run_async
|
|
async def test_supervisor_persists_complete_result_above_eight_mib(monkeypatch, tmp_path):
|
|
monkeypatch.setenv("CALCULATION_WORKER_STATE_DIR", str(tmp_path))
|
|
original = asyncio.create_subprocess_exec
|
|
|
|
async def spawn(*args, **kwargs):
|
|
# A real child writes the same file protocol as the calculation process.
|
|
code = ('import json,pathlib,sys; pathlib.Path(sys.argv[1]).write_text('
|
|
'json.dumps({"status":"succeeded","result":{"comparisons":"x"*(8*1024*1024),'
|
|
'"map_payload":{"features":["coverage"]}}}))')
|
|
return await original(sys.executable, "-c", code, args[-1], start_new_session=True)
|
|
|
|
monkeypatch.setattr(asyncio, "create_subprocess_exec", spawn)
|
|
results = []
|
|
|
|
class Repository:
|
|
async def heartbeat(self, *args):
|
|
return "running"
|
|
|
|
async def finish(self, row, owner, **terminal):
|
|
results.append(terminal)
|
|
return True
|
|
|
|
@asynccontextmanager
|
|
async def repository(project):
|
|
yield Repository()
|
|
|
|
worker = CalculationWorker(WorkerSettings(heartbeat_seconds=.02), repository=repository)
|
|
await asyncio.wait_for(worker.run_job({"project_id": "project", "job_id": "large-result",
|
|
"timeout_seconds": 10}), 15)
|
|
assert len(results) == 1
|
|
assert results[0]["status"] == "succeeded"
|
|
assert results[0]["result"] == {
|
|
"comparisons": "x" * (8 * 1024 * 1024), "map_payload": {"features": ["coverage"]},
|
|
}
|
|
|
|
|
|
def test_worker_settings_require_valid_concurrency(monkeypatch):
|
|
settings = WorkerSettings.from_environment()
|
|
assert settings.max_concurrency == 2
|
|
monkeypatch.setenv('CALCULATION_WORKER_HEARTBEAT_SECONDS', '30')
|
|
with pytest.raises(ValueError):
|
|
WorkerSettings.from_environment()
|
|
|
|
|
|
@run_async
|
|
async def test_process_group_cleanup_kills_native_style_descendant(tmp_path):
|
|
pid_file = tmp_path / 'pid'
|
|
code = ('import subprocess,time,pathlib,sys; '
|
|
'p=subprocess.Popen([sys.executable,"-c","import time; time.sleep(60)"]); '
|
|
'pathlib.Path(sys.argv[1]).write_text(str(p.pid)); time.sleep(60)')
|
|
process = await asyncio.create_subprocess_exec(sys.executable, '-c', code,
|
|
str(pid_file), start_new_session=True)
|
|
try:
|
|
for _ in range(100):
|
|
if pid_file.exists():
|
|
break
|
|
await asyncio.sleep(.02)
|
|
child_pid = int(pid_file.read_text())
|
|
await stop_process_group(process)
|
|
assert process.returncode is not None
|
|
# Linux may retain a dead grandchild as a zombie until PID 1 reaps it.
|
|
stat_path = Path(f'/proc/{child_pid}/stat')
|
|
for _ in range(100):
|
|
if not stat_path.exists() or stat_path.read_text().split()[2] == 'Z':
|
|
break
|
|
await asyncio.sleep(.01)
|
|
assert not stat_path.exists() or stat_path.read_text().split()[2] == 'Z'
|
|
finally:
|
|
await stop_process_group(process)
|
|
|
|
|
|
@run_async
|
|
@pytest.mark.parametrize('lease_status, expected', [('cancel_requested', 'cancelled'), (None, None)])
|
|
async def test_cancel_and_lost_lease_stop_process_before_persisting(monkeypatch, lease_status, expected):
|
|
processes = []
|
|
original = asyncio.create_subprocess_exec
|
|
async def spawn(*args, **kwargs):
|
|
process = await original(sys.executable, '-c', 'import time; time.sleep(60)',
|
|
start_new_session=True)
|
|
processes.append(process)
|
|
return process
|
|
monkeypatch.setattr(asyncio, 'create_subprocess_exec', spawn)
|
|
finishes = []
|
|
class Repository:
|
|
async def heartbeat(self, *args):
|
|
return lease_status
|
|
async def finish(self, row, owner, status, **kwargs):
|
|
assert processes[0].returncode is not None
|
|
finishes.append(status)
|
|
return True
|
|
@asynccontextmanager
|
|
async def repository(project):
|
|
yield Repository()
|
|
worker = CalculationWorker(WorkerSettings(heartbeat_seconds=.05), repository=repository)
|
|
await worker.run_job({'project_id': 'project', 'job_id': 'job', 'timeout_seconds': 60})
|
|
assert finishes == ([expected] if expected else [])
|
|
assert processes[0].returncode is not None
|
|
|
|
|
|
@run_async
|
|
async def test_timeout_stops_process_before_recording(monkeypatch):
|
|
original = asyncio.create_subprocess_exec
|
|
processes = []
|
|
async def spawn(*args, **kwargs):
|
|
p = await original(sys.executable, '-c', 'import time; time.sleep(60)', start_new_session=True)
|
|
processes.append(p)
|
|
return p
|
|
monkeypatch.setattr(asyncio, 'create_subprocess_exec', spawn)
|
|
statuses = []
|
|
class Repository:
|
|
async def heartbeat(self, *args):
|
|
return 'running'
|
|
async def finish(self, row, owner, status, **kwargs):
|
|
assert processes[0].returncode is not None
|
|
statuses.append(status)
|
|
return True
|
|
@asynccontextmanager
|
|
async def repository(project):
|
|
yield Repository()
|
|
worker = CalculationWorker(WorkerSettings(heartbeat_seconds=.02), repository=repository)
|
|
await worker.run_job({'project_id': 'project', 'timeout_seconds': 0})
|
|
assert statuses == ['timed_out']
|
|
|
|
|
|
@run_async
|
|
async def test_duplicate_supervisor_never_claims(monkeypatch):
|
|
from app.infra.db.metadb import database
|
|
class Lock:
|
|
async def scalar(self, statement):
|
|
assert 'pg_try_advisory_lock' in str(statement)
|
|
return False
|
|
async def commit(self):
|
|
pass
|
|
@asynccontextmanager
|
|
async def connection():
|
|
yield Lock()
|
|
monkeypatch.setattr(database, 'engine', SimpleNamespace(connect=connection))
|
|
@asynccontextmanager
|
|
async def repository(project):
|
|
pytest.fail('A duplicate supervisor must never open the queue')
|
|
yield
|
|
with pytest.raises(RuntimeError, match='already active'):
|
|
await CalculationWorker(WorkerSettings(), repository=repository).run()
|
|
|
|
|
|
@run_async
|
|
async def test_unavailable_project_does_not_block_other_projects(monkeypatch):
|
|
from app.infra.db import calculation_worker as infrastructure
|
|
|
|
class Lease:
|
|
async def ping(self):
|
|
pass
|
|
|
|
@asynccontextmanager
|
|
async def supervisor_lease():
|
|
yield Lease()
|
|
|
|
checked = asyncio.Event()
|
|
|
|
class Repository:
|
|
async def recover(self):
|
|
pass
|
|
|
|
async def claim(self, owner, lease_seconds):
|
|
checked.set()
|
|
return None
|
|
|
|
@asynccontextmanager
|
|
async def repository(project):
|
|
if project == 'unavailable':
|
|
raise RuntimeError('offline')
|
|
yield Repository()
|
|
|
|
async def projects():
|
|
return ('unavailable', 'healthy')
|
|
|
|
monkeypatch.setattr(infrastructure, 'calculation_supervisor_lease', supervisor_lease)
|
|
worker = CalculationWorker(
|
|
WorkerSettings(poll_seconds=.01, lease_seconds=1, heartbeat_seconds=.1),
|
|
repository=repository,
|
|
project_provider=projects,
|
|
)
|
|
task = asyncio.create_task(worker.run())
|
|
try:
|
|
await asyncio.wait_for(checked.wait(), 2)
|
|
finally:
|
|
task.cancel()
|
|
await asyncio.gather(task, return_exceptions=True)
|
|
|
|
|
|
@run_async
|
|
async def test_repository_terminal_updates_are_fenced():
|
|
from app.infra.db.postgresql.calculation_jobs import CalculationJobRepository
|
|
statements = []
|
|
class Cursor:
|
|
async def fetchone(self):
|
|
return None
|
|
class Connection:
|
|
async def execute(self, sql, parameters):
|
|
statements.append((sql, parameters))
|
|
return Cursor()
|
|
repository = CalculationJobRepository(Connection(), 'project')
|
|
row = {'job_id': 'job', 'lease_generation': 7}
|
|
assert await repository.finish(row, 'owner', 'succeeded', result={'original': True}) is False
|
|
sql, parameters = statements[-1]
|
|
assert 'lease_generation=%s' in sql and 'lease_owner=%s' in sql
|
|
assert 'lease_expires_at>clock_timestamp()' in sql
|
|
assert "status!='cancel_requested'" in sql
|
|
assert 7 in parameters
|
|
|
|
|
|
@run_async
|
|
async def test_guard_parent_death_pipe_kills_entire_group(tmp_path):
|
|
pid_file = tmp_path / 'guard-child'
|
|
child = 'import os,time,pathlib,sys; pathlib.Path(sys.argv[1]).write_text(str(os.getpid())); time.sleep(60)'
|
|
launcher = '''import asyncio,sys
|
|
from app.services.calculation_worker import guard_child
|
|
original = asyncio.create_subprocess_exec
|
|
async def spawn(*args, **kwargs):
|
|
return await original(sys.executable, '-c', sys.argv[2], sys.argv[1])
|
|
asyncio.create_subprocess_exec = spawn
|
|
asyncio.run(guard_child('unused','unused'))
|
|
'''
|
|
guard = await asyncio.create_subprocess_exec(sys.executable, '-c', launcher,
|
|
str(pid_file), child, stdin=asyncio.subprocess.PIPE, start_new_session=True)
|
|
try:
|
|
for _ in range(100):
|
|
if pid_file.exists():
|
|
break
|
|
await asyncio.sleep(.02)
|
|
child_pid = int(pid_file.read_text())
|
|
guard.stdin.close()
|
|
await asyncio.wait_for(guard.wait(), 3)
|
|
assert guard.returncode == -signal.SIGKILL
|
|
stat = Path(f'/proc/{child_pid}/stat')
|
|
for _ in range(100):
|
|
if not stat.exists() or stat.read_text().split()[2] == 'Z':
|
|
break
|
|
await asyncio.sleep(.01)
|
|
assert not stat.exists() or stat.read_text().split()[2] == 'Z'
|
|
finally:
|
|
await stop_process_group(guard)
|
|
|
|
|
|
@run_async
|
|
async def test_cleanup_failure_is_not_reported_cancelled(monkeypatch, tmp_path):
|
|
monkeypatch.setenv('CALCULATION_WORKER_STATE_DIR', str(tmp_path))
|
|
async def cleanup(row, directory):
|
|
(Path(directory) / 'cleanup-failed').touch()
|
|
raise RuntimeError('external resource unavailable')
|
|
monkeypatch.setattr('app.services.calculation_worker.cleanup_job_resources', cleanup)
|
|
original = asyncio.create_subprocess_exec
|
|
async def spawn(*args, **kwargs):
|
|
return await original(sys.executable, '-c', 'import time; time.sleep(60)', start_new_session=True)
|
|
monkeypatch.setattr(asyncio, 'create_subprocess_exec', spawn)
|
|
results = []
|
|
class Repository:
|
|
async def heartbeat(self, *args):
|
|
return 'cancel_requested'
|
|
async def finish(self, row, owner, status, **kwargs):
|
|
results.append((status, kwargs['error']['code']))
|
|
return True
|
|
@asynccontextmanager
|
|
async def repository(project):
|
|
yield Repository()
|
|
worker = CalculationWorker(WorkerSettings(), repository=repository)
|
|
await worker.run_job({'project_id': 'project', 'timeout_seconds': 60})
|
|
assert results == [('failed', 'CLEANUP_FAILED')]
|
|
assert len(list(tmp_path.glob('*/cleanup-failed'))) == 1
|
|
|
|
|
|
@run_async
|
|
async def test_startup_recovery_checks_persisted_generation(monkeypatch, tmp_path):
|
|
import json
|
|
from app.services.calculation_worker import recover_job_resources
|
|
project = '6b39b128-2ba8-454f-9bbe-e9553902ab53'
|
|
job = '2d3a6231-a8ee-4acb-bb13-58328d29b88f'
|
|
directory = tmp_path / (job + '-owned')
|
|
directory.mkdir()
|
|
(directory / 'scope.json').write_text(json.dumps({
|
|
'project_id': project, 'job_id': job, 'lease_generation': 1}))
|
|
monkeypatch.setenv('CALCULATION_WORKER_STATE_DIR', str(tmp_path))
|
|
class Repository:
|
|
async def get(self, job_id):
|
|
return {'lease_generation': 2}
|
|
@asynccontextmanager
|
|
async def repository(project):
|
|
yield Repository()
|
|
monkeypatch.setattr('app.services.calculation_worker.queue_repository', repository)
|
|
with pytest.raises(RuntimeError, match='ownership'):
|
|
await recover_job_resources((project,))
|
|
assert directory.exists()
|