buildgrid 0.2.51__tar.gz → 0.2.53__tar.gz

{buildgrid-0.2.51 → buildgrid-0.2.53}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: buildgrid
-Version: 0.2.51
+Version: 0.2.53
 Summary: A remote execution service
 License: Apache License, Version 2.0
 Project-URL: Homepage, https://buildgrid.build

{buildgrid-0.2.51 → buildgrid-0.2.53}/buildgrid/server/scheduler/impl.py

@@ -1048,7 +1048,7 @@ class Scheduler:
                 BotEntry.name.ilike(f"%{name_filter}%"),
                 BotEntry.bot_id.ilike(f"%{name_filter}%"),
             ),
-            BotEntry.instance_name == current_instance(),
+            self._bot_in_instance_pool(),
         )
         stmt = stmt.order_by(BotEntry.bot_id)
 
@@ -1949,7 +1949,9 @@ class Scheduler:
             )
 
     @timed(METRIC.SCHEDULER.ASSIGNMENT_DURATION)
-    def _fetch_job_for_bot(self, session: Session, bot: BotEntry, log_tags: Tags) -> JobEntry | None:
+    def _fetch_job_for_bot(
+        self, session: Session, bot: BotEntry, usage_diffs: InstanceQuotaUsageDiffs, log_tags: Tags
+    ) -> JobEntry | None:
         # Attempt to fetch a new job for a bot to work on.
         # This can help if there are usually more jobs available than bots.
 
@@ -1959,15 +1961,22 @@ class Scheduler:
             if bot.instance_name != "*":
                 job_statement = job_statement.where(self._job_in_instance_pool())
 
-            # Prioritize instance where usage <= max_quota - bot.capacity
-            # `- bot.capacity` to avoid over-assigning jobs to bots when nearing quota limits
-            instances_query = select(InstanceQuota.instance_name).where(
-                InstanceQuota.bot_cohort == bot.cohort,
-                InstanceQuota.current_usage <= InstanceQuota.max_quota - bot.capacity,
-            )
-            instances = session.execute(instances_query).scalars().all()
-            if instances:
-                job_statement = job_statement.where(JobEntry.instance_name.in_(instances))
+            if bot.cohort:
+                # Prioritize instance where usage <= max_quota - bot.capacity
+                # `- bot.capacity` to avoid over-assigning jobs to bots when nearing quota limits
+                instances_query = select(InstanceQuota.instance_name).where(
+                    InstanceQuota.bot_cohort == bot.cohort,
+                    InstanceQuota.current_usage <= InstanceQuota.max_quota - bot.capacity,
+                )
+                instances: set[str] = set()
+                instances.update(session.execute(instances_query).scalars().all())
+                # Always allow scheduling more jobs of an instance if we're returning usage
+                instances.update(
+                    {instance for (cohort, instance), diff in usage_diffs.items() if cohort == bot.cohort and diff < 0}
+                )
+
+                if instances:
+                    job_statement = job_statement.where(JobEntry.instance_name.in_(instances))
 
             if next_job := session.execute(job_statement).scalar_one_or_none():
                 log_tags["db.next_job_name"] = next_job.name
@@ -2321,7 +2330,7 @@ class Scheduler:
 
         for _ in range(fetch_limit):
             # Try to fill up the newly free capacity with new jobs.
-            if new_job := self._fetch_job_for_bot(session, bot, log_tags):
+            if new_job := self._fetch_job_for_bot(session, bot, usage_diffs, log_tags):
                 if bot.cohort:
                     usage_diffs[(bot.cohort, new_job.instance_name)] += 1
                 synchronized_leases.append(new_job.to_lease_proto())

{buildgrid-0.2.51 → buildgrid-0.2.53}/buildgrid/server/version.py

@@ -13,4 +13,4 @@
 # limitations under the License.
 
 
-__version__ = "0.2.51"
+__version__ = "0.2.53"

{buildgrid-0.2.51 → buildgrid-0.2.53}/buildgrid.egg-info/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: buildgrid
-Version: 0.2.51
+Version: 0.2.53
 Summary: A remote execution service
 License: Apache License, Version 2.0
 Project-URL: Homepage, https://buildgrid.build

{buildgrid-0.2.51 → buildgrid-0.2.53}/pyproject.toml

@@ -21,7 +21,7 @@ build-backend = 'setuptools.build_meta'
 
 [project]
 name = "buildgrid"
-version = "0.2.51"
+version = "0.2.53"
 requires-python = ">=3.11"
 description = "A remote execution service"
 readme = "README.rst"

{buildgrid-0.2.51 → buildgrid-0.2.53}/tests/test_scheduler.py

@@ -1341,6 +1341,42 @@ def test_bot_proactively_fetch(controller_postgres):
         assert jobs[1].worker_name == bot_id
 
 
+def test_bot_proactively_fetch_near_max_quota(controller_postgres):
+    scheduler: Scheduler = controller_postgres.execution_instance.scheduler
+    # Create quota configs for two instances, allowing only 1 concurrent job each
+    scheduler.put_instance_quota("sql", "linux", 1, 1)
+    scheduler.put_instance_quota("other", "linux", 1, 1)
+
+    # Enqueue 2 jobs
+    bot_name, bot_id, _, _ = mock_queue_job_action(scheduler, do_not_cache=True, skip_cache_lookup=True, assign=False)
+    _ = mock_queue_job_action(scheduler, do_not_cache=True, skip_cache_lookup=True, assign=False)
+    scheduler.assign_job_by_priority()  # successful
+    scheduler.assign_job_by_priority(failure_backoff=60)  # failure, adding `schedule_after`
+
+    leases = scheduler.synchronize_bot_leases(bot_name, bot_id, BotStatus.OK.value, [])
+    assert len(leases) == 1
+    leases[0].state = LeaseState.COMPLETED.value
+    leases = scheduler.synchronize_bot_leases(bot_name, bot_id, BotStatus.OK.value, leases)
+    assert len(leases)
+
+    with scheduler._sql.session() as session:
+        jobs = session.execute(select(JobEntry)).scalars().all()
+        assert len(jobs) == 2
+        assert jobs[0].stage == OperationStage.COMPLETED.value
+        # Assert that the second job is scheduled
+        assert jobs[1].assigned
+        assert jobs[1].worker_name == bot_id
+
+        # Assert that the bot has no more capacity to fetch more jobs
+        bot_entry = session.execute(select(BotEntry).where(BotEntry.name == bot_name)).scalar_one()
+        assert bot_entry.capacity == 0
+
+    # Check usage is recorded correctly
+    usage = scheduler.get_instance_quota("sql", "linux")
+    assert usage is not None
+    assert usage.current_usage == 1
+
+
 def test_schedule_by_instances(controller_postgres):
     scheduler: Scheduler = controller_postgres.execution_instance.scheduler