phoenix/litellm-mirror
0
0
Fork 1
mirror of https://github.com/BerriAI/litellm.git synced 2025-04-26 03:04:13 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions

fix schema to work with lock acquisition

Browse source
This commit is contained in:
Ishaan Jaff 2025-03-27 22:54:46 -07:00
parent 758182fc7f
commit 0cb6cb6c66
3 changed files with 42 additions and 25 deletions
Download patch file Download diff file Expand all files Collapse all files

54
litellm/proxy/db/pod_leader_manager.py
View file

@ -19,16 +19,19 @@ class PodLockManager:
Ensures that only one pod can run a cron job at a time. Ensures that only one pod can run a cron job at a time.
""" """
def __init__(self, prisma_client: Optional[PrismaClient], cronjob_id: str): def __init__(self, cronjob_id: str):
self.pod_id = str(uuid.uuid4()) self.pod_id = str(uuid.uuid4())
self.prisma = prisma_client
self.cronjob_id = cronjob_id self.cronjob_id = cronjob_id
async def acquire_lock(self) -> bool: async def acquire_lock(self) -> bool:
""" """
Attempt to acquire the lock for a specific cron job. Attempt to acquire the lock for a specific cron job.
""" """
if not self.prisma: from litellm.proxy.proxy_server import prisma_client
verbose_proxy_logger.debug("acquiring lock for cronjob_id=%s", self.cronjob_id)
if not prisma_client:
verbose_proxy_logger.debug("prisma is None, returning False")
return False return False
try: try:
current_time = datetime.now(timezone.utc) current_time = datetime.now(timezone.utc)
@ -38,21 +41,24 @@ class PodLockManager:
) )
# Attempt to acquire the lock by upserting the record in the `cronjob_locks` table # Attempt to acquire the lock by upserting the record in the `cronjob_locks` table
cronjob_lock = await self.prisma.db.cronJob.upsert( cronjob_lock = await prisma_client.db.cronjob.upsert(
where={"cronjob_id": self.cronjob_id}, where={"cronjob_id": self.cronjob_id},
create={ data={
"cronjob_id": self.cronjob_id, "create": {
"pod_id": self.pod_id, "cronjob_id": self.cronjob_id,
"status": "ACTIVE", "pod_id": self.pod_id,
"last_updated": current_time, "status": "ACTIVE",
"ttl": ttl_expiry, "last_updated": current_time,
}, "ttl": ttl_expiry,
update={ },
"status": "ACTIVE", "update": {
"last_updated": current_time, "status": "ACTIVE",
"ttl": ttl_expiry, "last_updated": current_time,
"ttl": ttl_expiry,
},
}, },
) )
verbose_proxy_logger.debug("cronjob_lock=%s", cronjob_lock)
if cronjob_lock.status == "ACTIVE" and cronjob_lock.pod_id == self.pod_id: if cronjob_lock.status == "ACTIVE" and cronjob_lock.pod_id == self.pod_id:
verbose_proxy_logger.debug( verbose_proxy_logger.debug(
@ -70,16 +76,21 @@ class PodLockManager:
""" """
Renew the lock (update the TTL) for the pod holding the lock. Renew the lock (update the TTL) for the pod holding the lock.
""" """
if not self.prisma: from litellm.proxy.proxy_server import prisma_client
if not prisma_client:
return False return False
try: try:
verbose_proxy_logger.debug(
"renewing lock for cronjob_id=%s", self.cronjob_id
)
current_time = datetime.now(timezone.utc) current_time = datetime.now(timezone.utc)
# Extend the TTL for another DEFAULT_CRON_JOB_LOCK_TTL_SECONDS # Extend the TTL for another DEFAULT_CRON_JOB_LOCK_TTL_SECONDS
ttl_expiry = current_time + timedelta( ttl_expiry = current_time + timedelta(
seconds=DEFAULT_CRON_JOB_LOCK_TTL_SECONDS seconds=DEFAULT_CRON_JOB_LOCK_TTL_SECONDS
) )
await self.prisma.db.cronJob.update( await prisma_client.db.cronjob.update(
where={"cronjob_id": self.cronjob_id, "pod_id": self.pod_id}, where={"cronjob_id": self.cronjob_id, "pod_id": self.pod_id},
data={"ttl": ttl_expiry, "last_updated": current_time}, data={"ttl": ttl_expiry, "last_updated": current_time},
) )
@ -95,10 +106,15 @@ class PodLockManager:
""" """
Release the lock and mark the pod as inactive. Release the lock and mark the pod as inactive.
""" """
if not self.prisma: from litellm.proxy.proxy_server import prisma_client
if not prisma_client:
return False return False
try: try:
await self.prisma.db.cronJob.update( verbose_proxy_logger.debug(
"releasing lock for cronjob_id=%s", self.cronjob_id
)
await prisma_client.db.cronjob.update(
where={"cronjob_id": self.cronjob_id, "pod_id": self.pod_id}, where={"cronjob_id": self.cronjob_id, "pod_id": self.pod_id},
data={"status": "INACTIVE"}, data={"status": "INACTIVE"},
) )

7
litellm/proxy/schema.prisma
View file

@ -340,10 +340,10 @@ model LiteLLM_DailyUserSpend {
// Track the status of cron jobs running. Only allow one pod to run the job at a time // Track the status of cron jobs running. Only allow one pod to run the job at a time
model CronJob { model CronJob {
id String @id @default(cuid()) // Unique ID for the record cronjob_id String @id @default(cuid()) // Unique ID for the record
podId String // Unique identifier for the pod acting as the leader pod_id String // Unique identifier for the pod acting as the leader
status JobStatus @default(INACTIVE) // Status of the cron job (active or inactive) status JobStatus @default(INACTIVE) // Status of the cron job (active or inactive)
lastUpdated DateTime @default(now()) // Timestamp for the last update of the cron job record last_updated DateTime @default(now()) // Timestamp for the last update of the cron job record
ttl DateTime // Time when the leader's lease expires ttl DateTime // Time when the leader's lease expires
} }
@ -352,3 +352,4 @@ enum JobStatus {
INACTIVE INACTIVE
} }

6
schema.prisma
View file

@ -339,10 +339,10 @@ model LiteLLM_DailyUserSpend {
// Track the status of cron jobs running. Only allow one pod to run the job at a time // Track the status of cron jobs running. Only allow one pod to run the job at a time
model CronJob { model CronJob {
id String @id @default(cuid()) // Unique ID for the record cronjob_id String @id @default(cuid()) // Unique ID for the record
podId String // Unique identifier for the pod acting as the leader pod_id String // Unique identifier for the pod acting as the leader
status JobStatus @default(INACTIVE) // Status of the cron job (active or inactive) status JobStatus @default(INACTIVE) // Status of the cron job (active or inactive)
lastUpdated DateTime @default(now()) // Timestamp for the last update of the cron job record last_updated DateTime @default(now()) // Timestamp for the last update of the cron job record
ttl DateTime // Time when the leader's lease expires ttl DateTime // Time when the leader's lease expires
} }