diff --git a/swh/scheduler/simulator/__init__.py b/swh/scheduler/simulator/__init__.py index 5e470c8..483f9a9 100644 --- a/swh/scheduler/simulator/__init__.py +++ b/swh/scheduler/simulator/__init__.py @@ -1,123 +1,126 @@ # Copyright (C) 2021 The Software Heritage developers # See the AUTHORS file at the top-level directory of this distribution # License: GNU General Public License version 3, or any later version # See top-level LICENSE file for more information from datetime import datetime, timedelta, timezone import logging from typing import Dict, Generator, Optional from simpy import Event from swh.scheduler import get_scheduler from swh.scheduler.model import ListedOrigin from . import origin_scheduler, task_scheduler from .common import Environment, Queue, SimulationReport, Task from .origins import load_task_process logger = logging.getLogger(__name__) def worker_process( env: Environment, name: str, task_queue: Queue, status_queue: Queue ) -> Generator[Event, Task, None]: """A worker which consumes tasks from the input task_queue. Tasks themselves send OriginVisitStatus objects to the status_queue.""" logger.debug("%s worker %s: Start", env.time, name) while True: task = yield task_queue.get() logger.debug( "%s worker %s: Run task %s origin=%s", env.time, name, task.visit_type, task.origin, ) yield env.process(load_task_process(env, task, status_queue=status_queue)) def setup( env: Environment, scheduler: str, policy: Optional[str], workers_per_type: Dict[str, int], task_queue_capacity: int, ): # We expect PGHOST, PGPORT, ... set externally task_queues = { visit_type: Queue(env, capacity=task_queue_capacity) for visit_type in workers_per_type } status_queue = Queue(env) if scheduler == "origin_scheduler": if policy is None: raise ValueError("origin_scheduler needs a scheduling policy") env.process(origin_scheduler.scheduler_runner_process(env, task_queues, policy)) env.process( origin_scheduler.scheduler_journal_client_process(env, status_queue) ) elif scheduler == "task_scheduler": if policy is not None: raise ValueError("task_scheduler doesn't support a scheduling policy") env.process(task_scheduler.scheduler_runner_process(env, task_queues)) env.process(task_scheduler.scheduler_listener_process(env, status_queue)) else: raise ValueError(f"Unknown scheduler: {scheduler}") for visit_type, num_workers in workers_per_type.items(): task_queue = task_queues[visit_type] for i in range(num_workers): worker_name = f"worker-{visit_type}-{i}" env.process(worker_process(env, worker_name, task_queue, status_queue)) -def fill_test_data(): +def fill_test_data(num_origins: int = 100000): scheduler = get_scheduler(cls="local", db="") + stored_lister = scheduler.get_or_create_lister(name="example") + assert stored_lister.id is not None + origins = [ ListedOrigin( lister_id=stored_lister.id, url=f"https://example.com/{i:04d}.git", visit_type="git", last_update=datetime(2020, 6, 15, 16, 0, 0, i, tzinfo=timezone.utc), ) - for i in range(100000) + for i in range(num_origins) ] scheduler.record_listed_origins(origins) scheduler.create_tasks( [ { **origin.as_task_dict(), "policy": "recurring", "next_run": origin.last_update, "interval": timedelta(days=64), } for origin in origins ] ) def run(scheduler: str, policy: Optional[str], runtime: Optional[int]): NUM_WORKERS = 48 start_time = datetime.now(tz=timezone.utc) env = Environment(start_time=start_time) env.scheduler = get_scheduler(cls="local", db="") env.report = SimulationReport() setup( env, scheduler=scheduler, policy=policy, workers_per_type={"git": NUM_WORKERS}, task_queue_capacity=10000, ) try: env.run(until=runtime) except KeyboardInterrupt: pass finally: end_time = env.time print("total time:", end_time - start_time) print(env.report.format()) diff --git a/swh/scheduler/tests/test_simulator.py b/swh/scheduler/tests/test_simulator.py new file mode 100644 index 0000000..97eb1bd --- /dev/null +++ b/swh/scheduler/tests/test_simulator.py @@ -0,0 +1,45 @@ +# Copyright (C) 2021 The Software Heritage developers +# See the AUTHORS file at the top-level directory of this distribution +# License: GNU General Public License version 3, or any later version +# See top-level LICENSE file for more information + +import pytest + +import swh.scheduler.simulator as simulator +from swh.scheduler.tests.common import TASK_TYPES + +NUM_ORIGINS = 42 +TEST_RUNTIME = 1000 + + +@pytest.fixture +def monkeypatch_get_scheduler(swh_scheduler, monkeypatch): + def get_scheduler(*args, **kwargs): + return swh_scheduler + + monkeypatch.setattr(simulator, "get_scheduler", get_scheduler) + + for task_type in TASK_TYPES.values(): + swh_scheduler.create_task_type(task_type) + + +def test_fill_test_data(swh_scheduler, monkeypatch_get_scheduler): + simulator.fill_test_data(num_origins=NUM_ORIGINS) + + res = swh_scheduler.get_listed_origins() + assert len(res.origins) == NUM_ORIGINS + assert res.next_page_token is None + + res = swh_scheduler.search_tasks() + assert len(res) == NUM_ORIGINS + + +@pytest.mark.parametrize("policy", ("oldest_scheduled_first",)) +def test_run_origin_scheduler(swh_scheduler, monkeypatch_get_scheduler, policy): + simulator.fill_test_data(num_origins=NUM_ORIGINS) + simulator.run("origin_scheduler", policy=policy, runtime=TEST_RUNTIME) + + +def test_run_task_scheduler(swh_scheduler, monkeypatch_get_scheduler): + simulator.fill_test_data(num_origins=NUM_ORIGINS) + simulator.run("task_scheduler", policy=None, runtime=TEST_RUNTIME)