Mass transfer scaling

.gitignore

@@ -129,6 +129,9 @@ node_modules/
 # The Orthanc storage folders
 OrthancStorage*
 
+# Application logs
+logs/
+
 # Redis DB dump
 dump.rdb
 

adit/mass_transfer/tests/test_scale_mass_transfer_worker.py

@@ -0,0 +1,203 @@
+import os
+import signal
+import subprocess
+import sys
+import time
+from pathlib import Path
+from unittest.mock import call
+
+import pytest
+from django.db import connection
+from procrastinate import JobContext
+from procrastinate.contrib.django import app
+from procrastinate.contrib.django.models import ProcrastinateJob
+from pytest_mock import MockerFixture
+
+from adit.core.models import DicomJob, DicomTask
+from adit.mass_transfer.factories import MassTransferJobFactory, MassTransferTaskFactory
+from adit.mass_transfer.tasks import queue_mass_transfer_tasks
+from cli import scale_mass_transfer_worker
+
+ROOT_DIR = Path(__file__).resolve().parents[3]
+GRACEFUL_TASK_NAME = (
+    "adit.mass_transfer.tests.test_scale_mass_transfer_worker.graceful_mass_transfer_test_task"
+)
+
+
+@app.task(queue="mass_transfer", pass_context=True, name=GRACEFUL_TASK_NAME)
+def graceful_mass_transfer_test_task(
+    context: JobContext,
+    sleep_seconds: int = 60,
+    poll_interval: float = 0.1,
+):
+    """Test-only helper that always completes unless forcefully interrupted."""
+    deadline = time.monotonic() + sleep_seconds
+    while time.monotonic() < deadline:
+        _ = context
+        time.sleep(poll_interval)
+
+
+def _build_database_url_from_connection() -> str:
+    db_settings = connection.settings_dict
+    user = db_settings.get("USER") or ""
+    password = db_settings.get("PASSWORD") or ""
+    host = db_settings.get("HOST") or "localhost"
+    port = db_settings.get("PORT") or "5432"
+    name = db_settings["NAME"]
+    return f"postgres://{user}:{password}@{host}:{port}/{name}"
+
+
+def _get_job_status(job_id: int) -> str | None:
+    job = ProcrastinateJob.objects.filter(id=job_id).first()
+    if not job:
+        return None
+    return str(job.status)
+
+
+def _wait_for_status(job_id: int, statuses: set[str], timeout_seconds: int) -> str | None:
+    deadline = time.monotonic() + timeout_seconds
+    while time.monotonic() < deadline:
+        status = _get_job_status(job_id)
+        if status in statuses:
+            return status
+        time.sleep(0.2)
+    return _get_job_status(job_id)
+
+
+@pytest.mark.django_db(transaction=True)
+def test_scale_mass_transfer_worker_scales_up_and_down_without_touching_queued_jobs(
+    mocker: MockerFixture,
+):
+    """Scaling workers must not abort or detach already queued transfer tasks."""
+    job = MassTransferJobFactory.create(status=DicomJob.Status.PENDING)
+    task1 = MassTransferTaskFactory.create(status=DicomTask.Status.PENDING, job=job)
+    task2 = MassTransferTaskFactory.create(status=DicomTask.Status.PENDING, job=job)
+    queue_mass_transfer_tasks(job_id=job.pk)
+
+    task1.refresh_from_db()
+    task2.refresh_from_db()
+    queued_job_ids = {task1.queued_job_id, task2.queued_job_id}
+    assert None not in queued_job_ids
+
+    queued_job_ids_int = {job_id for job_id in queued_job_ids if job_id is not None}
+    assert (
+        set(ProcrastinateJob.objects.filter(id__in=queued_job_ids_int).values_list("id", flat=True))
+        == queued_job_ids_int
+    )
+
+    helper = mocker.Mock()
+    helper.is_production.return_value = True
+    helper.get_stack_name.return_value = "adit-prod"
+    mocker.patch("cli.cli_helper.CommandHelper", return_value=helper)
+    # Ensure load_config_from_env_file returns a mapping so .get()/.strip() work
+    helper.load_config_from_env_file.return_value = {}
+
+    scale_mass_transfer_worker(replicas=3)
+    scale_mass_transfer_worker(replicas=0)
+
+    assert helper.execute_cmd.call_args_list == [
+        call("docker service update --stop-grace-period 10s adit-prod_mass_transfer_worker"),
+        call("docker service scale adit-prod_mass_transfer_worker=3"),
+        call("docker service update --stop-grace-period 10s adit-prod_mass_transfer_worker"),
+        call("docker service scale adit-prod_mass_transfer_worker=0"),
+    ]
+
+    task1.refresh_from_db()
+    task2.refresh_from_db()
+    assert {task1.queued_job_id, task2.queued_job_id} == queued_job_ids
+    assert (
+        set(ProcrastinateJob.objects.filter(id__in=queued_job_ids_int).values_list("id", flat=True))
+        == queued_job_ids_int
+    )
+
+
+@pytest.mark.django_db(transaction=True)
+def test_scale_mass_transfer_worker_scale_down_finishes_current_job_and_blocks_next(
+    mocker: MockerFixture,
+):
+    """Scale-down should let the current task finish and leave queued transfer tasks untouched."""
+    job = MassTransferJobFactory.create(status=DicomJob.Status.PENDING)
+    task1 = MassTransferTaskFactory.create(status=DicomTask.Status.PENDING, job=job)
+    task2 = MassTransferTaskFactory.create(status=DicomTask.Status.PENDING, job=job)
+    queue_mass_transfer_tasks(job_id=job.pk)
+
+    task1.refresh_from_db()
+    task2.refresh_from_db()
+
+    # Run a deterministic long-running task so we can assert graceful shutdown behavior.
+    running_task_seconds = 30
+    graceful_timeout_seconds = 50
+
+    running_job_id = app.configure_task(
+        GRACEFUL_TASK_NAME,
+        allow_unknown=False,
+        priority=10_000,
+    ).defer(sleep_seconds=running_task_seconds, poll_interval=0.05)
+
+    assert task1.queued_job_id is not None
+    assert task2.queued_job_id is not None
+
+    worker_env = os.environ.copy()
+    worker_env["DATABASE_URL"] = _build_database_url_from_connection()
+    worker_cmd = (
+        "import os; "
+        "os.environ.setdefault('DJANGO_SETTINGS_MODULE', 'adit.settings.development'); "
+        "import django; django.setup(); "
+        "import adit.mass_transfer.tests.test_scale_mass_transfer_worker; "
+        "from django.core.management import execute_from_command_line; "
+        "execute_from_command_line(["
+        "'manage.py', 'procrastinate', 'worker', '--queues', 'mass_transfer', "
+        f"'--shutdown-graceful-timeout', '{graceful_timeout_seconds}', '--delete-jobs', 'never'"
+        "])"
+    )
+    worker_process = subprocess.Popen(
+        [sys.executable, "-c", worker_cmd],
+        cwd=ROOT_DIR,
+        env=worker_env,
+        stdout=subprocess.DEVNULL,
+        stderr=subprocess.DEVNULL,
+    )
+
+    try:
+        first_status = _wait_for_status(running_job_id, {"doing"}, timeout_seconds=20)
+        assert first_status == "doing"
+
+        helper = mocker.Mock()
+        helper.is_production.return_value = True
+        helper.get_stack_name.return_value = "adit-prod"
+        mocker.patch("cli.cli_helper.CommandHelper", return_value=helper)
+        # Ensure load_config_from_env_file returns a mapping so .get()/.strip() work
+        helper.load_config_from_env_file.return_value = {}
+
+        def execute_cmd_side_effect(command: str):
+            if command == "docker service scale adit-prod_mass_transfer_worker=0":
+                worker_process.send_signal(signal.SIGTERM)
+
+        helper.execute_cmd.side_effect = execute_cmd_side_effect
+
+        scale_mass_transfer_worker(replicas=1)
+        scale_mass_transfer_worker(replicas=0)
+
+        worker_process.wait(timeout=running_task_seconds + 20)
+
+        final_running_status = _wait_for_status(
+            running_job_id,
+            {"succeeded", "failed", "cancelled", "aborted"},
+            timeout_seconds=30,
+        )
+        final_blocked_status_1 = _get_job_status(task1.queued_job_id)
+        final_blocked_status_2 = _get_job_status(task2.queued_job_id)
+
+        assert final_running_status == "succeeded"
+        assert final_blocked_status_1 == "todo"
+        assert final_blocked_status_2 == "todo"
+        assert helper.execute_cmd.call_args_list == [
+            call("docker service update --stop-grace-period 10s adit-prod_mass_transfer_worker"),
+            call("docker service scale adit-prod_mass_transfer_worker=1"),
+            call("docker service update --stop-grace-period 10s adit-prod_mass_transfer_worker"),
+            call("docker service scale adit-prod_mass_transfer_worker=0"),
+        ]
+    finally:
+        if worker_process.poll() is None:
+            worker_process.terminate()
+            worker_process.wait(timeout=10)

cli.py

@@ -1,7 +1,9 @@
 #! /usr/bin/env python3
 import os
+import shlex
 import shutil
 from glob import glob
+from pathlib import Path
 from typing import Annotated
 
 import typer
@@ -37,6 +39,154 @@
 app.command()(commands.try_github_actions)
 
 
+@app.command()
+def scale_mass_transfer_worker(
+    replicas: Annotated[
+        int,
+        typer.Argument(
+            help=(
+                "Target replica count for the Docker Swarm mass transfer worker service. "
+                "Use 0 to scale down and values > 0 to scale up."
+            )
+        ),
+    ],
+):
+    if replicas < 0:
+        typer.echo("replicas must be >= 0")
+        raise typer.Exit(code=1)
+
+    helper = cli_helper.CommandHelper()
+    helper.prepare_environment()
+
+    if not helper.is_production():
+        typer.echo(
+            "scale-mass-transfer-worker task can only be used in production environment. "
+            "Check ENVIRONMENT setting in .env file."
+        )
+        raise typer.Exit(code=1)
+
+    service_name = f"{helper.get_stack_name()}_mass_transfer_worker"
+    env = helper.load_config_from_env_file()
+    raw_grace = (env.get("MASS_TRANSFER_WORKER_STOP_GRACE_PERIOD") or "").strip()
+    if not raw_grace:
+        grace_period = "10s"
+    else:
+        grace_period = raw_grace
+
+    helper.execute_cmd(
+        f"docker service update --stop-grace-period {shlex.quote(grace_period)} {service_name}"
+    )
+
+    helper.execute_cmd(f"docker service scale {service_name}={replicas}")
+
+
+@app.command()
+def configure_mass_transfer_worker_cron():
+    helper = cli_helper.CommandHelper()
+    helper.prepare_environment()
+
+    if not helper.is_production():
+        typer.echo(
+            "configure-mass-transfer-worker-cron task can only be used in production environment. "
+            "Check ENVIRONMENT setting in .env file."
+        )
+        raise typer.Exit(code=1)
+
+    env = helper.load_config_from_env_file()
+    # Inline cron block construction and validation so this command is self-contained.
+    raw_up = env.get("MASS_TRANSFER_WORKER_REPLICAS") or "-1"
+    raw_down = env.get("MASS_TRANSFER_WORKER_REPLICAS_DOWNSCALED") or "-1"
+    try:
+        up_replicas = int(raw_up)
+        down_replicas = int(raw_down)
+    except ValueError:
+        typer.echo(
+            f"Invalid int for MASS_TRANSFER_WORKER_REPLICAS(_DOWNSCALED): {raw_up} / {raw_down}"
+        )
+        raise typer.Exit(code=1)
+    if up_replicas < 0 or down_replicas < 0:
+        typer.echo("MASS_TRANSFER_WORKER_REPLICAS(_DOWNSCALED) must be >= 0")
+        raise typer.Exit(code=1)
+
+    up_cron = env.get("MASS_TRANSFER_WORKER_SCALE_UP_CRON") or ""
+    down_cron = env.get("MASS_TRANSFER_WORKER_SCALE_DOWN_CRON") or ""
+    if len(up_cron.split()) != 5:
+        typer.echo(f"Invalid cron expression for MASS_TRANSFER_WORKER_SCALE_UP_CRON: {up_cron}")
+        raise typer.Exit(code=1)
+    if len(down_cron.split()) != 5:
+        typer.echo(f"Invalid cron expression for MASS_TRANSFER_WORKER_SCALE_DOWN_CRON: {down_cron}")
+        raise typer.Exit(code=1)
+
+    project_root = shlex.quote(str(Path(helper.root_path)))
+    logs_dir = shlex.quote(str(Path(helper.root_path) / "logs"))
+
+    # Ensure logs directory exists
+    Path(helper.root_path).joinpath("logs").mkdir(exist_ok=True)
+
+    log_file = f"{logs_dir}/mass_transfer_worker_cron.log"
+
+    scale_up_cmd = (
+        f"cd {project_root} && /usr/local/bin/uv run cli scale-mass-transfer-worker {up_replicas}"
+        f" >> {log_file} 2>&1"
+    )
+    scale_down_cmd = (
+        f"cd {project_root} && /usr/local/bin/uv run cli scale-mass-transfer-worker {down_replicas}"
+        f" >> {log_file} 2>&1"
+    )
+
+    cron_marker_start = "# ADIT_MASS_TRANSFER_WORKER_AUTOSCALE_START"
+    cron_marker_end = "# ADIT_MASS_TRANSFER_WORKER_AUTOSCALE_END"
+    cron_block = "\n".join(
+        [
+            cron_marker_start,
+            f"{up_cron} {scale_up_cmd}",
+            f"{down_cron} {scale_down_cmd}",
+            cron_marker_end,
+        ]
+    )
+    typer.echo("Executed: {}".format(cron_block))
+
+    escaped_start = cron_marker_start.replace("/", "\\/")
+    escaped_end = cron_marker_end.replace("/", "\\/")
+    crontab_install_cmd = (
+        "tmpfile=$(mktemp) && "
+        f"(crontab -l 2>/dev/null | sed '/{escaped_start}/,/{escaped_end}/d'; "
+        "cat <<'EOF'\n"
+        f"{cron_block}\n"
+        "EOF\n"
+        ') > "$tmpfile" && '
+        'crontab "$tmpfile" && '
+        'rm "$tmpfile"'
+    )
+    helper.execute_cmd(crontab_install_cmd)
+
+
+@app.command()
+def remove_mass_transfer_worker_cron():
+    helper = cli_helper.CommandHelper()
+    helper.prepare_environment()
+
+    if not helper.is_production():
+        typer.echo(
+            "remove-mass-transfer-worker-cron task can only be used in production environment. "
+            "Check ENVIRONMENT setting in .env file."
+        )
+        raise typer.Exit(code=1)
+
+    cron_marker_start = "# ADIT_MASS_TRANSFER_WORKER_AUTOSCALE_START"
+    cron_marker_end = "# ADIT_MASS_TRANSFER_WORKER_AUTOSCALE_END"
+    escaped_start = cron_marker_start.replace("/", "\\/")
+    escaped_end = cron_marker_end.replace("/", "\\/")
+
+    crontab_remove_cmd = (
+        "tmpfile=$(mktemp) && "
+        f"(crontab -l 2>/dev/null | sed '/{escaped_start}/,/{escaped_end}/d') > \"$tmpfile\" && "
+        'crontab "$tmpfile" && '
+        'rm "$tmpfile"'
+    )
+    helper.execute_cmd(crontab_remove_cmd)
+
+
 @app.command()
 def populate_orthancs(
     reset: Annotated[bool, typer.Option(help="Clear Orthancs before populate")] = False,

docker-compose.dev.yml

@@ -72,7 +72,7 @@ services:
     command: >
       bash -c "
         wait-for-it -s postgres.local:5432 -t 60 &&
-        ./manage.py bg_worker -l debug -q mass_transfer --autoreload
+        ./manage.py bg_worker -l debug -q mass_transfer
       "
 
   receiver:
@@ -92,7 +92,9 @@ services:
   orthanc1:
     ports:
       - "7501:7501"
+      - "6501:6501"
 
   orthanc2:
     ports:
       - "7502:7502"
+      - "6502:6502"