Running scheduler in multiple pods in GKE

[kshitij-enable]

hi,

I am running database scheduler along with fastapi in 2 pods. Which results into scheduler scheduling same job twice.
Is there any way to implement a lock mechanism so that only one scheduler is able to schedule jobs at a given point in time. And if the active one crashes, then second should be able to pickup without any manual intervention.

[farahats9]

Yes it is possible to do so and it is done in the redis scheduler Redbeat using a distributed lock mechanism, I will try to work on this some time soon but if you are in a hurry I would recommend using redbeat for your case.

[kshitij-enable]

Yes it is possible to do so and it is done in the redis scheduler Redbeat using a distributed lock mechanism, I will try to work on this some time soon but if you are in a hurry I would recommend using redbeat for your case.

Thanks @farahats9 for the revert. I had a look at RedBeat, but will prefer to stick to the Database scheduler implementation which you have built. Gives us more control in terms of building monitoring and auditing functions around it. I have implemented distributed lock mechanism using redis in the scheduler code. Please have a look and let me know in case of any suggestions.

@property
    def schedule(self) -> dict[str, ModelEntry]:
        """Get the current schedule."""
        if not self.is_owner():
            return {}

def is_owner(self) -> bool:
        """Checks if the current scheduler is the leader."""
        if self.redis_lock.is_owner():
            logger.warning(
                "refreshing lock for tenant %s",
                getattr(self, "tenant_id", "default"),
            )
            self.redis_lock.refresh_lock()
            return True
        else:
            if not self.redis_lock.acquire_lock():
                logger.warning(
                    "lock held by other, skipping schedule for tenant %s",
                    getattr(self, "tenant_id", "default"),
                )
                return False

            logger.warning(
                "acquired lock, scheduling for tenant %s",
                getattr(self, "tenant_id", "default"),
            )
        return True

class DatabaseScheduler(Scheduler):  # type:ignore
def __init__(self, *args: Any, **kwargs: Any) -> None:
        """Initialize the database scheduler."""
        self.tenant_id = kwargs.get("tenant_id") or change_args.get("tenant_id")
        set_tenant(self.tenant_id)
        self.session = get_session()
        self._dirty: set[str] = set()
        self.redis_lock: SchedulerLock = SchedulerLock(
            lock_key=f"{self.tenant_id}_celery_scheduler_lock",
            ttl=120,
        )
        super().__init__(*args, **kwargs)
        self._finalize = Finalize(self, self.sync, exitpriority=5)
        self.max_interval = kwargs.get("max_interval") or self.app.conf.beat_max_loop_interval or DEFAULT_MAX_INTERVAL

import uuid
from typing import Optional

import redis
from enable.common.redis.engine import get_engine

type Redis = redis.Redis[str] | redis.RedisCluster[str]


class SchedulerLock:
    """A simple Redis-based distributed lock."""

    def __init__(
        self,
        lock_key: str = "celery_scheduler_lock",
        ttl: int = 120,
    ) -> None:
        """A simple Redis-based distributed lock."""
        self.redis: "Redis" = get_engine()
        self.lock_key: str = lock_key
        self.ttl: int = ttl
        # Unique identifier for this lock owner
        self.lock_id: str = str(uuid.uuid4())

    def acquire_lock(self) -> bool:
        """Attempt to acquire the lock. Returns True if successful."""
        # SET key value NX EX ttl
        return self.redis.set(self.lock_key, self.lock_id, nx=True, ex=self.ttl) is True

    def refresh_lock(self) -> bool:
        """Refresh the TTL on the lock if we already own it."""
        current: Optional[str] = self.redis.get(self.lock_key)
        if current and current == self.lock_id:
            return self.redis.expire(self.lock_key, self.ttl)
        return False

    def is_owner(self) -> bool:
        """Check if this instance currently holds the lock."""
        current: Optional[str] = self.redis.get(self.lock_key)
        return bool(current and current == self.lock_id)

[farahats9]

That's a good start, I was thinking of making the lock as a database table in order to keep the beat scheduler independent of the broker used, I know it is probably more efficient to keep the key on Redis but this means we have to build the locking for each broker type and some of them will not be possible like AWS SQS for example.

You can tinker around and make a PR and I will try to give a hand when I have time soon, if you can translate this code to be on the database instead of redis we can make some tests and merge it if it works.