Coverage for netbox_data_import/proposal_dispatch.py: 97.85%

1# SPDX-License-Identifier: Apache-2.0 

2# SPDX-FileCopyrightText: 2026 Marcin Zieba <marcinpsk@gmail.com> 

3"""Start Resolution Proposal Jobs only while the active Inference Backend has a free request slot. 

4 

5A requested proposal waits as a queued row without a Job. The dispatcher gives the oldest waiting rows 

6a Job while slots are free. A slot is in use while a proposal's Job can still run, whatever the 

7proposal's status: the worker of a cancelled row can still be inside its backend call. A slot is 

8released only when no worker holds its Job lock and the queue no longer has its task. 

9""" 

10 

11import logging 

12 

13from django.core.cache import cache 

14from django.db import connection, transaction 

15from django.utils import timezone 

16from netbox.constants import RQ_QUEUE_LOW 

17from redis.exceptions import ConnectionError as RedisConnectionError 

18from redis.exceptions import TimeoutError as RedisTimeoutError 

19 

20from .inference_backend import NoActiveInferenceBackend, resolve_active_backend 

21from .inference_credentials import CredentialFailure 

22from .inference_trust import InvalidInferenceConfiguration 

23from .jobs import ResolutionProposalJob, _task_lost, _try_job_lock, import_queue_task 

24from .models import ProposalFailureReason, ProposalStatus, ResolutionProposal 

25from .proposal_jobs import CREDENTIAL_FAILURE_REASONS, proposal_job_timeout 

26from .resolution_proposals import fail_proposal, fail_queued_proposal, record_proposal_job 

27 

28logger = logging.getLogger(__name__) 

29 

30SLOT_TASK_LOST = "The proposal task is no longer queued or running." 

31_DISPATCH_LOCK = "netbox-data-import-proposal-dispatch" 

32POLL_DISPATCH_KEY = "netbox-data-import-proposal-poll-dispatch" 

33POLL_DISPATCH_INTERVAL_SECONDS = 5 

34 

35 

36def dispatch_proposals() -> list[int]: 

37 """Release leaked slots, then give the oldest waiting proposals a Job while slots are free. 

38 

39 Call it outside a transaction. Return the ids of the dispatched proposals. A queue push that 

40 fails after the commit fails its attempt with queue_unavailable, and its error propagates. 

41 """ 

42 committed: list[bool] = [] 

43 dispatched: list[tuple[int, object]] = [] 

44 try: 

45 with transaction.atomic(): 

46 # Registered first, so it runs before a queue push that may raise after the commit. 

47 transaction.on_commit(lambda: committed.append(True)) 

48 # Two dispatchers that both counted the same free slot would exceed the limit. 

49 with connection.cursor() as cursor: 

50 cursor.execute("SELECT pg_advisory_xact_lock(hashtextextended(%s, 0))", [_DISPATCH_LOCK]) 

51 dispatched = _fill_free_slots() 

52 except BaseException: 

53 if committed: 

54 for proposal_id, job in dispatched: 

55 _fail_unpushed(proposal_id, job) 

56 raise 

57 return [proposal_id for proposal_id, _job in dispatched] 

58 

59 

60def dispatch_proposals_on_poll() -> list[int]: 

61 """Dispatch at most once per interval for all pending-card polls together. 

62 

63 The poll is only the backstop for a slot that a killed worker never handed on: a request and the 

64 end of each proposal Job dispatch at once. The cache is shared by every NetBox process. 

65 """ 

66 if not cache.add(POLL_DISPATCH_KEY, True, timeout=POLL_DISPATCH_INTERVAL_SECONDS): 

67 return [] 

68 return dispatch_proposals() 

69 

70 

71def _fill_free_slots() -> list[tuple[int, object]]: 

72 """Enqueue one Job for each waiting proposal that a free slot can take, oldest first.""" 

73 waiting = ResolutionProposal.objects.filter(status=ProposalStatus.QUEUED, job__isnull=True).order_by( 

74 "created", "pk" 

75 ) 

76 in_use = _release_leaked_slots() 

77 try: 

78 backend = resolve_active_backend() 

79 except (NoActiveInferenceBackend, InvalidInferenceConfiguration, CredentialFailure) as exc: 

80 # No Job could call a backend, so each waiting row fails with the reason its Job would record. 

81 reason = ProposalFailureReason.INVALID_CONFIGURATION 

82 if isinstance(exc, CredentialFailure): 

83 reason = CREDENTIAL_FAILURE_REASONS.get(exc.category, ProposalFailureReason.CREDENTIAL_UNAVAILABLE) 

84 for proposal_id in waiting.values_list("pk", flat=True): 

85 fail_queued_proposal(proposal_id, reason=reason) 

86 return [] 

87 job_timeout = proposal_job_timeout(backend) 

88 free = max(backend.max_concurrent_requests - in_use, 0) 

89 dispatched = [] 

90 for proposal in waiting.select_related("requested_by")[:free]: 

91 # The low queue lets a trace sync queued later start before minutes of inference. 

92 job = ResolutionProposalJob.enqueue( 

93 name=ResolutionProposalJob.Meta.name, 

94 user=proposal.requested_by, 

95 queue_name=RQ_QUEUE_LOW, 

96 job_timeout=job_timeout, 

97 proposal_id=proposal.pk, 

98 ) 

99 record_proposal_job(proposal.pk, job) 

100 dispatched.append((proposal.pk, job)) 

101 return dispatched 

102 

103 

104def _release_leaked_slots() -> int: 

105 """Release each slot whose Job can no longer run, and return how many slots stay in use.""" 

106 from core.choices import JobStatusChoices 

107 from core.models import Job 

108 

109 in_use = 0 

110 holders = ResolutionProposal.objects.filter(job__status__in=JobStatusChoices.ENQUEUED_STATE_CHOICES) 

111 for proposal in holders.select_related("job"): 

112 job = proposal.job 

113 # A worker holds this lock through its whole delivery, whatever Redis or the clock say. 

114 if not _try_job_lock(job): 

115 in_use += 1 

116 continue 

117 try: 

118 job.refresh_from_db() 

119 except Job.DoesNotExist: 

120 # The Job was deleted after the holder query, and a deleted Job holds no slot. 

121 continue 

122 if job.status not in JobStatusChoices.ENQUEUED_STATE_CHOICES: 

123 continue 

124 # A free lock is no proof alone: a dropped database session releases it while the call goes on. 

125 if not _task_lost(job, import_queue_task(job)): 

126 in_use += 1 

127 continue 

128 logger.warning("Released the request slot of Resolution Proposal %s: %s", proposal.pk, SLOT_TASK_LOST) 

129 Job.objects.filter(pk=job.pk, status__in=JobStatusChoices.ENQUEUED_STATE_CHOICES).update( 

130 status=JobStatusChoices.STATUS_ERRORED, completed=timezone.now(), error=SLOT_TASK_LOST 

131 ) 

132 # The job system lost the run, not the backend; a terminal row keeps its status. 

133 fail_proposal(proposal.pk, reason=ProposalFailureReason.QUEUE_UNAVAILABLE) 

134 return in_use 

135 

136 

137def _fail_unpushed(proposal_id: int, job) -> None: 

138 """Fail an attempt whose task never reached the queue, so it neither waits nor holds a slot.""" 

139 from core.choices import JobStatusChoices 

140 from core.models import Job 

141 

142 try: 

143 pushed = import_queue_task(job) is not None 

144 except (RedisConnectionError, RedisTimeoutError): 

145 # A queue that cannot answer cannot show the task, so the attempt fails. 

146 pushed = False 

147 # A failed push stops only the later pushes, so a task that reached the queue keeps its attempt. 

148 if pushed: 

149 return 

150 # A worker that took the attempt owns it, so only an attempt still queued fails here. 

151 if fail_queued_proposal(proposal_id, reason=ProposalFailureReason.QUEUE_UNAVAILABLE): 

152 Job.objects.filter(pk=job.pk, status=JobStatusChoices.STATUS_PENDING).update( 

153 status=JobStatusChoices.STATUS_ERRORED 

154 ) 

155 

156 

157__all__ = ("dispatch_proposals", "dispatch_proposals_on_poll")