diff --git a/diracx-tasks/pyproject.toml b/diracx-tasks/pyproject.toml index def809d82..c70716a1e 100644 --- a/diracx-tasks/pyproject.toml +++ b/diracx-tasks/pyproject.toml @@ -37,6 +37,8 @@ TaskDB = "diracx.tasks.plumbing.persistence:TaskDB" [project.entry-points."diracx.tasks.jobs"] CleanSandboxStoreTask = "diracx.tasks.jobs:CleanSandboxStoreTask" +DummyJobExecutorMonitorTask = "diracx.tasks.jobs:DummyJobExecutorMonitorTask" +DummyJobExecutorTask = "diracx.tasks.jobs:DummyJobExecutorTask" [build-system] requires = ["hatchling", "hatch-vcs"] diff --git a/diracx-tasks/src/diracx/tasks/jobs/__init__.py b/diracx-tasks/src/diracx/tasks/jobs/__init__.py index 3b09465d8..60c37cf41 100644 --- a/diracx-tasks/src/diracx/tasks/jobs/__init__.py +++ b/diracx-tasks/src/diracx/tasks/jobs/__init__.py @@ -2,6 +2,9 @@ __all__ = [ "CleanSandboxStoreTask", + "DummyJobExecutorMonitorTask", + "DummyJobExecutorTask", ] from .clean_sandbox_store import CleanSandboxStoreTask +from .dummy_job_executor import DummyJobExecutorMonitorTask, DummyJobExecutorTask diff --git a/diracx-tasks/src/diracx/tasks/jobs/dummy_job_executor.py b/diracx-tasks/src/diracx/tasks/jobs/dummy_job_executor.py new file mode 100644 index 000000000..cd9a72991 --- /dev/null +++ b/diracx-tasks/src/diracx/tasks/jobs/dummy_job_executor.py @@ -0,0 +1,166 @@ +"""Tasks that simulate the execution of jobs, for use in demo deployments.""" + +from __future__ import annotations + +import dataclasses +import logging +from datetime import UTC, datetime, timedelta + +from pydantic import PositiveInt + +from diracx.core.models import ( + JobStatus, + JobStatusUpdate, + ScalarSearchOperator, + ScalarSearchSpec, +) +from diracx.core.settings import ServiceSettingsBase +from diracx.db.os import JobParametersDB +from diracx.db.sql import JobDB, JobLoggingDB, TaskQueueDB +from diracx.logic.jobs import set_job_statuses +from diracx.tasks.plumbing.base_task import BaseTask, PeriodicBaseTask +from diracx.tasks.plumbing.depends import Config +from diracx.tasks.plumbing.enums import Priority, Size +from diracx.tasks.plumbing.lock_registry import JOB +from diracx.tasks.plumbing.locks import BaseLock, MutexLock +from diracx.tasks.plumbing.retry_policies import ExponentialBackoff +from diracx.tasks.plumbing.schedules import IntervalSeconds + +logger = logging.getLogger(__name__) + +MINOR_STATUS = "DummyExecutor" + + +class DummyJobExecutorSettings(ServiceSettingsBase): + """Settings controlling automatic dummy job execution.""" + + model_config = ServiceSettingsBase.model_config | { + "env_prefix": "DIRACX_TASKS_DUMMY_JOB_EXECUTOR_", + "use_attribute_docstrings": True, + } + + enabled: bool = False + """Whether the monitor is scheduled automatically.""" + + interval_seconds: PositiveInt = 10 + """How often the enabled monitor searches for received jobs.""" + + +_settings = DummyJobExecutorSettings() + + +@dataclasses.dataclass +class DummyJobExecutorMonitorTask(PeriodicBaseTask): + """Periodically pick up newly received jobs and hand them to the dummy executor. + + When enabled, every job in ``Received`` state is moved to ``Waiting`` and a + one-shot ``DummyJobExecutorTask`` is scheduled for it. + """ + + priority = Priority.BACKGROUND + size = Size.SMALL + _enabled = _settings.enabled + default_schedule = IntervalSeconds(_settings.interval_seconds) + + async def execute( + self, + config: Config, + job_db: JobDB, + job_logging_db: JobLoggingDB, + task_queue_db: TaskQueueDB, + job_parameters_db: JobParametersDB, + ) -> int: + _, jobs = await job_db.search( + ["JobID"], + [ + ScalarSearchSpec( + parameter="Status", + operator=ScalarSearchOperator.EQUAL, + value=JobStatus.RECEIVED, + ) + ], + [], + ) + if not jobs: + return 0 + + job_ids = [job["JobID"] for job in jobs] + logger.info("Moving %d received job(s) to Waiting: %s", len(job_ids), job_ids) + await set_job_statuses( + { + job_id: { + datetime.now(UTC): JobStatusUpdate( + Status=JobStatus.WAITING, + MinorStatus=MINOR_STATUS, + ) + } + for job_id in job_ids + }, + config=config, + job_db=job_db, + job_logging_db=job_logging_db, + task_queue_db=task_queue_db, + job_parameters_db=job_parameters_db, + ) + + for job_id in job_ids: + await DummyJobExecutorTask(job_id=job_id).schedule() + + return len(job_ids) + + +@dataclasses.dataclass +class DummyJobExecutorTask(BaseTask): + """Simulate the execution of a single job. + + The job is walked through the state machine's valid path + ``Waiting -> Matched -> Running -> Done``: a direct jump to ``Done`` would be + silently rejected. ``set_job_statuses`` applies the timestamped updates in + chronological order, so a single call with increasing timestamps walks the + whole chain. + """ + + priority = Priority.NORMAL + size = Size.LARGE + retry_policy = ExponentialBackoff(base_delay_seconds=10, max_retries=3) + + job_id: int + + @property + def execution_locks(self) -> list[BaseLock]: + return [MutexLock(JOB, self.job_id)] + + async def execute( + self, + config: Config, + job_db: JobDB, + job_logging_db: JobLoggingDB, + task_queue_db: TaskQueueDB, + job_parameters_db: JobParametersDB, + ) -> int: + logger.info("Simulating execution of job %d", self.job_id) + now = datetime.now(UTC) + await set_job_statuses( + { + self.job_id: { + now: JobStatusUpdate( + Status=JobStatus.MATCHED, + MinorStatus=MINOR_STATUS, + ), + now + timedelta(seconds=1): JobStatusUpdate( + Status=JobStatus.RUNNING, + MinorStatus=MINOR_STATUS, + ), + now + timedelta(seconds=5): JobStatusUpdate( + Status=JobStatus.DONE, + MinorStatus=MINOR_STATUS, + ), + } + }, + config=config, + job_db=job_db, + job_logging_db=job_logging_db, + task_queue_db=task_queue_db, + job_parameters_db=job_parameters_db, + ) + return self.job_id diff --git a/diracx-tasks/tests/test_dummy_job_executor.py b/diracx-tasks/tests/test_dummy_job_executor.py new file mode 100644 index 000000000..6166a6187 --- /dev/null +++ b/diracx-tasks/tests/test_dummy_job_executor.py @@ -0,0 +1,171 @@ +"""Tests for the dummy job executor tasks.""" + +from __future__ import annotations + +import json +import os +import subprocess +import sys +from unittest.mock import AsyncMock, MagicMock + +import pytest +from pydantic import ValidationError + +from diracx.core.models import JobStatus +from diracx.tasks.jobs import DummyJobExecutorMonitorTask, DummyJobExecutorTask +from diracx.tasks.jobs import dummy_job_executor as dummy_job_executor_module +from diracx.tasks.plumbing.locks import MutexLock + +FEATURE_ENABLED_ENV = "DIRACX_TASKS_DUMMY_JOB_EXECUTOR_ENABLED" +FEATURE_INTERVAL_ENV = "DIRACX_TASKS_DUMMY_JOB_EXECUTOR_INTERVAL_SECONDS" +SCHEDULER_STATE_SCRIPT = """ +import json +from datetime import UTC, datetime +from unittest.mock import MagicMock + +from diracx.tasks.plumbing.factory import load_task_registry +from diracx.tasks.plumbing.scheduler import TaskScheduler + +task_name = "jobs:DummyJobExecutorMonitorTask" +registry = load_task_registry() +task_cls = registry[task_name] +scheduler = TaskScheduler( + broker=MagicMock(), + redis_url="redis://unused", + task_registry=registry, +) +before = datetime.now(tz=UTC) +scheduler._compute_initial_schedules() +next_run = scheduler._next_runs.get((task_name, "")) +print( + json.dumps( + { + "enabled": task_cls._enabled, + "tracked": next_run is not None, + "delay_seconds": ( + (next_run - before).total_seconds() if next_run is not None else None + ), + } + ) +) +""" + + +def make_dependencies(): + return { + "config": MagicMock(name="config"), + "job_db": AsyncMock(name="job_db"), + "job_logging_db": MagicMock(name="job_logging_db"), + "task_queue_db": MagicMock(name="task_queue_db"), + "job_parameters_db": MagicMock(name="job_parameters_db"), + } + + +def get_scheduler_state(feature_env: dict[str, str]) -> dict: + env = os.environ.copy() + env.pop(FEATURE_ENABLED_ENV, None) + env.pop(FEATURE_INTERVAL_ENV, None) + env.update(feature_env) + result = subprocess.run( + [sys.executable, "-c", SCHEDULER_STATE_SCRIPT], + check=True, + capture_output=True, + env=env, + text=True, + ) + return json.loads(result.stdout) + + +def test_monitor_schedule_activation_is_environment_controlled(): + default_state = get_scheduler_state({}) + assert default_state == { + "enabled": False, + "tracked": False, + "delay_seconds": None, + } + + local_state = get_scheduler_state( + { + FEATURE_ENABLED_ENV: "true", + FEATURE_INTERVAL_ENV: "10", + } + ) + assert local_state["enabled"] is True + assert local_state["tracked"] is True + assert 9 <= local_state["delay_seconds"] <= 11 + + +def test_monitor_interval_must_be_positive(): + with pytest.raises(ValidationError): + dummy_job_executor_module.DummyJobExecutorSettings(interval_seconds=0) + + +def test_executor_takes_a_per_job_mutex(): + locks = DummyJobExecutorTask(job_id=42).execution_locks + + assert len(locks) == 1 + assert isinstance(locks[0], MutexLock) + assert locks[0].redis_key == "lock:mutex:job:42" + + +async def test_executor_walks_job_through_the_state_machine(monkeypatch): + set_job_statuses = AsyncMock() + monkeypatch.setattr(dummy_job_executor_module, "set_job_statuses", set_job_statuses) + deps = make_dependencies() + + result = await DummyJobExecutorTask(job_id=42).execute(**deps) + + assert result == 42 + set_job_statuses.assert_awaited_once() + status_changes = set_job_statuses.await_args.args[0] + assert set(status_changes) == {42} + updates = status_changes[42] + assert list(updates) == sorted(updates), "timestamps must be increasing" + assert [update.status for update in updates.values()] == [ + JobStatus.MATCHED, + JobStatus.RUNNING, + JobStatus.DONE, + ] + + +async def test_monitor_moves_received_jobs_and_schedules_executors(monkeypatch): + set_job_statuses = AsyncMock() + monkeypatch.setattr(dummy_job_executor_module, "set_job_statuses", set_job_statuses) + scheduled = [] + + async def fake_schedule(self, **kwargs): + scheduled.append(self.job_id) + return "task-id" + + monkeypatch.setattr(DummyJobExecutorTask, "schedule", fake_schedule) + deps = make_dependencies() + deps["job_db"].search.return_value = (2, [{"JobID": 1}, {"JobID": 2}]) + + result = await DummyJobExecutorMonitorTask().execute(**deps) + + assert result == 2 + deps["job_db"].search.assert_awaited_once() + (search_spec,) = deps["job_db"].search.await_args.args[1] + assert search_spec["parameter"] == "Status" + assert search_spec["value"] == JobStatus.RECEIVED + set_job_statuses.assert_awaited_once() + status_changes = set_job_statuses.await_args.args[0] + assert set(status_changes) == {1, 2} + for updates in status_changes.values(): + assert [update.status for update in updates.values()] == [JobStatus.WAITING] + assert scheduled == [1, 2] + + +async def test_monitor_does_nothing_without_received_jobs(monkeypatch): + set_job_statuses = AsyncMock() + monkeypatch.setattr(dummy_job_executor_module, "set_job_statuses", set_job_statuses) + schedule_executor = AsyncMock() + monkeypatch.setattr(DummyJobExecutorTask, "schedule", schedule_executor) + deps = make_dependencies() + deps["job_db"].search.return_value = (0, []) + + result = await DummyJobExecutorMonitorTask().execute(**deps) + + assert result == 0 + set_job_statuses.assert_not_awaited() + schedule_executor.assert_not_awaited() diff --git a/docs/admin/how-to/tasks/run-task-manually.md b/docs/admin/how-to/tasks/run-task-manually.md index c8422fbb4..32d57db36 100644 --- a/docs/admin/how-to/tasks/run-task-manually.md +++ b/docs/admin/how-to/tasks/run-task-manually.md @@ -23,6 +23,35 @@ diracx-task-run call lollygag:SyncOwnersTask --args '["alice"]' diracx-task-run call lollygag:SyncOwnersTask --args '["alice"]' --kwargs '{}' ``` +## Run the dummy job executor + +The task package includes two tasks that simulate job execution. The periodic +`jobs:DummyJobExecutorMonitorTask` moves every `Received` job to `Waiting` and +schedules a one-shot `jobs:DummyJobExecutorTask` for it, which walks the job +through `Matched` → `Running` → `Done`. + +Automatic monitoring is disabled by default. `run_local.sh` explicitly enables +the monitor every 10 seconds with +`DIRACX_TASKS_DUMMY_JOB_EXECUTOR_ENABLED=true` and +`DIRACX_TASKS_DUMMY_JOB_EXECUTOR_INTERVAL_SECONDS=10`. Demo deployment +enablement requires the coordinated `diracx-charts` values change and is not +part of this commit. + +Both tasks talk to the job databases, so the relevant `DIRACX_DB_URL_*`, +`DIRACX_OS_DB_*`, and `DIRACX_CONFIG_BACKEND_URL` variables must be set. + +Run the monitor once to pick up all `Received` jobs: + +```bash +diracx-task-run call jobs:DummyJobExecutorMonitorTask +``` + +Or simulate the execution of a single job by passing its job ID: + +```bash +diracx-task-run call jobs:DummyJobExecutorTask --args '[42]' +``` + ## Debugging The `--debugger` flag drops into Python's debugger: diff --git a/docs/admin/reference/env-variables.md b/docs/admin/reference/env-variables.md index e2b2b6285..2e2a37956 100644 --- a/docs/admin/reference/env-variables.md +++ b/docs/admin/reference/env-variables.md @@ -363,3 +363,21 @@ Whether to use an insecure gRPC connection for the OpenTelemetry collector. *Optional*, default value: `None` A JSON-encoded dictionary of headers to pass to the OpenTelemetry collector, e.g. {"tenant_id": "lhcbdiracx-cert"}. + +## Tasks + +## DummyJobExecutorSettings + +Settings controlling automatic dummy job execution. + +### `DIRACX_TASKS_DUMMY_JOB_EXECUTOR_ENABLED` + +*Optional*, default value: `False` + +Whether the monitor is scheduled automatically. + +### `DIRACX_TASKS_DUMMY_JOB_EXECUTOR_INTERVAL_SECONDS` + +*Optional*, default value: `10` + +How often the enabled monitor searches for received jobs. diff --git a/docs/admin/reference/env-variables.md.j2 b/docs/admin/reference/env-variables.md.j2 index 8f17665e1..4630da850 100644 --- a/docs/admin/reference/env-variables.md.j2 +++ b/docs/admin/reference/env-variables.md.j2 @@ -18,3 +18,7 @@ _X, where X is a number. The files will be loaded in order. {{ render_class('SandboxStoreSettings') }} {{ render_class('OTELSettings') }} + +## Tasks + +{{ render_class('DummyJobExecutorSettings') }} diff --git a/run_local.sh b/run_local.sh index 51f4f2c50..a6d405c13 100755 --- a/run_local.sh +++ b/run_local.sh @@ -98,6 +98,8 @@ export DIRACX_SANDBOX_STORE_BUCKET_NAME=sandboxes export DIRACX_SANDBOX_STORE_AUTO_CREATE_BUCKET=true export DIRACX_SANDBOX_STORE_S3_CLIENT_KWARGS='{"endpoint_url": "http://localhost:8333", "aws_access_key_id": "console", "aws_secret_access_key": "console123"}' export DIRACX_TASKS_REDIS_URL="redis://localhost:6379" +export DIRACX_TASKS_DUMMY_JOB_EXECUTOR_ENABLED=true +export DIRACX_TASKS_DUMMY_JOB_EXECUTOR_INTERVAL_SECONDS=10 # Write all DIRACX env vars to a sourceable file for use in other terminals script_dir="$(cd "$(dirname "$0")" && pwd)"