[Pools] Prioritize Idle Workers for Scale Down

[lloyd-brown]

This PR addresses a problem where changing the number of workers in a pool leads to us downing workers that are actively running jobs. Here is the working example: we have two workers and there is only a job running on worker 2, we want to scale the number of workers down to 1, ideally we would only tear down worker 1 and leave worker 2 to finish its job, this does not work, currently we tear down both workers in order to update to the new service version that has 2 workers.

Changes made

• The first change we made is a simple modification to the autoscaler to also consider the number of jobs running on a worker when scaling down the number of replicas. This PR changes the selection logic when scaling down workers to prioritize idle workers. We add another sort criteria so that we first sort by replica status, then version, then number of running jobs, then replica ID.
  • This alone does not suffice, with only this change we will still update the version and scale down worker 2 because it is “outdated”
• Instead of creating a new version for --workers updates (which do not allow a new YAML so the only thing that can be changing is the number of workers) we introduce a new call in the serve controller to directly change the number of replicas in the autoscaler.
  • Why do we do this?
    • By changing the version the autoscaler will terminate the existing replicas and schedule new ones which will remove any existing jobs.
  • A side effect of this approach is that changing the number of workers is now faster
  • I could also see this being useful for SkyServe if a user just wants to quickly set the number of replicas manually via CLI

Backwards compatibility:

If we don’t have support for updating replicas directly we will fall back to the old logic of creating a new version with the updated number of workers, losing the ability to scale without cancelling jobs.

I added several unit tests to make sure

• We prioritize replicas without running jobs
• We don't change the existing logic when the job counts are the same
• We value Status and Version more than the job count

I added a smoke test that makes sure when we scale down a pool that has 2 workers and only 1 of them is running a job we always scale down the one that is idle.

Tested (run the relevant ones):

• Code formatting: install pre-commit (auto-check on commit) or bash format.sh
• Any manual or new tests for this PR (please specify below)
• All smoke tests: /smoke-test (CI) or pytest tests/test_smoke.py (local)
• Relevant individual tests: /smoke-test -k test_name (CI) or pytest tests/test_smoke.py::test_name (local)
• Backward compatibility: /quicktest-core (CI) or pytest tests/smoke_tests/test_backward_compat.py (local)

[gemini-code-assist]

Summary of Changes

Hello @lloyd-brown, I'm Gemini Code Assist¹! I'm currently reviewing this pull request and will post my feedback shortly. In the meantime, here's a summary to help you and other reviewers quickly get up to speed!

This pull request enhances the worker pool autoscaling mechanism by introducing a more intelligent selection process for scaling down workers. The primary goal is to minimize disruption and optimize resource utilization by ensuring that workers with active jobs are retained over idle ones when a scale-down event occurs. This change refines the existing sorting criteria to explicitly account for the number of running jobs on each worker, making the autoscaler more efficient and job-aware.

Highlights

• Prioritized Scale Down: The worker selection logic for scaling down pools now prioritizes idle workers (those with fewer running jobs).
• New Sorting Criteria: The sorting mechanism for selecting workers to scale down has been updated to consider replica status, then version, then the number of running jobs, and finally replica ID.
• Comprehensive Testing: New unit tests were added to verify the job count prioritization, ensure existing logic for equal job counts is preserved, and confirm that status and version priorities are still respected. A smoke test was also introduced to validate the idle worker prioritization in a real-world scenario.

Using Gemini Code Assist

The full guide for Gemini Code Assist can be found on our documentation page, here are some quick tips.

Invoking Gemini

You can request assistance from Gemini at any point by creating a comment using either /gemini <command> or @gemini-code-assist <command>. Below is a summary of the supported commands on the current page.

Feature	Command	Description
Code Review	/gemini review	Performs a code review for the current pull request in its current state.
Pull Request Summary	/gemini summary	Provides a summary of the current pull request in its current state.
Comment	@gemini-code-assist	Responds in comments when explicitly tagged, both in pull request comments and review comments.
Help	/gemini help	Displays a list of available commands.

Customization

To customize Gemini Code Assist for GitHub experience, repository maintainers can create a configuration file and/or provide a custom code review style guide (such as PEP-8 for Python) by creating and adding files to a .gemini/ folder in the base of the repository. Detailed instructions can be found here.

Limitations & Feedback

Gemini Code Assist may make mistakes. Please leave feedback on any instances where its feedback is incorrect or counter productive. You can react with 👍 and 👎 on @gemini-code-assist comments. If you're interested in giving your feedback about your experience with Gemini Code Assist for Github and other Google products, sign up here.

You can also get AI-powered code generation, chat, as well as code reviews directly in the IDE at no cost with the Gemini Code Assist IDE Extension.

Footnotes

1. Review the Privacy Notices, Generative AI Prohibited Use Policy, Terms of Service, and learn how to configure Gemini Code Assist in GitHub here. Gemini can make mistakes, so double check it and use code with caution. ↩

[gemini-code-assist]

Code Review

The pull request introduces a significant improvement to the autoscaling logic by prioritizing idle workers during scale-down operations. This ensures that active jobs are not prematurely terminated, leading to a more robust and user-friendly service. The changes are well-documented, and the new sorting criteria are clearly explained in the docstrings. Comprehensive unit tests and a smoke test have been added to validate the new behavior, covering various scenarios including job count prioritization, equal job counts, and the preservation of existing status and version priorities.

[lloyd-brown]

/smoke-test

[lloyd-brown]

/smoke-test -k test_pool_scale_down_with_job_count_priority --kubernetes
/smoke-test -k test_pool_scale_down_with_job_count_priority --kubernetes --grpc
/smoke-test -k test_pool_scale_down_with_job_count_priority --gcp

[lloyd-brown]

/smoke-test

[lloyd-brown]

Tests only failing due to issue fixed in #8256 and resource unavailability.

[cg505]

The autoscaler change looks good.
However, I think a new version that keeps the same config (except for worker count) should re-use the same replicas without needing to add the update_replicas RPC.

skypilot/sky/serve/replica_managers.py

Lines 1476 to 1480 in afb7768

	# Reuse all replicas that have the same config as the new version
	# (except for the `service` field) by directly setting the version to be
	# the latest version. This can significantly improve the speed
	# for updating an existing service with only config changes to the
	# service specs, e.g. scale down the service.

If this isn't working correctly, we should try to understand why instead of just going around it, to avoid having two implementations of the same thing.

[cg505]

Naming: It's confusing that this doesn't call task.validate(). Maybe just _apply_admin_policy?

[cg505]

Is this just a no-op for non-pools?