[serve][2/n] Gang scheduling -- validation and utilities

[jeffreywang88]

Description

This PR lays the groundwork for gang scheduling in Ray Serve by extending the data model, adding config validation rules, and introducing a utility for querying active placement group IDs from GCS.

Test plan

• test_schema.py: invalid gang_size, mutual exclusivity with max_replicas_per_node and placement_group_strategy
• test_api.py: @serve.deployment rejects gang config combined with max_replicas_per_node / placement_group_strategy
• test_application_state.py: override gang config via .options(), reject invalid gang size multiple
• test_util.py: verifies the GCS query utility

Related issues

RFC: #60873
Original PR: #60802

Additional information

Optional: Add implementation details, API changes, usage examples, screenshots, etc.

[gemini-code-assist]

Code Review

This pull request lays the groundwork for gang scheduling in Ray Serve by extending the data model and adding essential validation rules. The enforcement of mutual exclusivity between gang scheduling and conflicting placement group strategies is well-implemented. However, there is a critical bug in the get_active_placement_group_ids utility function where it attempts to use HasField on a scalar bytes field, which is not supported in proto3 and will cause a crash. Additionally, the utility should filter out Nil IDs to accurately identify active placement groups. There is also a minor validation gap in the deployment decorator regarding replica multiples.

[abrarsheikh]

i am assuming this will be added later?

[jeffreywang88]

Actually no, we'd prefer setting the placement strategy through GangSchedulingConfig.gang_placement_strategy.

[abrarsheikh]

need to learn more about why we need this. We need to figure out a way around this, one it uses private API from serve and this is disallowed at project level. (b) it is making RPC call to GSC from what it looks like, so we cannot be calling this from every controller iteration loop.

[jeffreywang88]

Great callout, I haven't thought about the GCS aspect. Luckily, we will not call this in every controller loop. Starting from #61207, get_active_placement_group_ids will be invoked in _detect_and_remove_leaked_placement_groups only upon controller recovery, so I guess we don't need worry about the performance aspect.

It'll look something like this:

        gang_pg_names_in_cluster = [
            name
            for name in all_current_placement_group_names
            if name.startswith(GANG_PG_NAME_PREFIX)
        ]
        if gang_pg_names_in_cluster:
            pg_table = ray.util.placement_group_table()
            gang_pg_name_to_id: Dict[str, str] = {}
            for pg_id_hex, entry in pg_table.items():
                name = entry.get("name", "")
                if name.startswith(GANG_PG_NAME_PREFIX):
                    gang_pg_name_to_id[name] = pg_id_hex

            try:
                occupied_pg_ids = get_active_placement_group_ids()
            except Exception:
                logger.warning(
                    "Skipping gang PG leak detection due to GCS query failure.",
                    exc_info=True,
                )
            else:
                for gang_pg_name in gang_pg_names_in_cluster:
                    pg_id = gang_pg_name_to_id.get(gang_pg_name)
                    if pg_id is not None and pg_id not in occupied_pg_ids:
                        leaked_pg_names.append(gang_pg_name)

[jeffreywang88]

Which private APIs from serve are you referring to?

[abrarsheikh]

from ray._private.state import state

[jeffreywang88]

gotcha, let me fix this

[jeffreywang88]

Addressed in the latest commit. Also added logic to only detect pgs reference by serve actors to avoid from deleting non-serve PGs. (The deletion logic will be introduced in a follow-up PR.)

[abrarsheikh]

please merge master

[cursor]

Cursor Bugbot has reviewed your changes and found 1 potential issue.

[cursor]

Schema test missing num_replicas causes wrong validation error

Medium Severity

Both test_mutually_exclusive_max_replicas_per_node_and_gang_scheduling_config and test_mutually_exclusive_placement_group_strategy_and_gang_scheduling_config set gang_size=2 without setting num_replicas. If get_minimal_deployment_schema() returns a schema with default num_replicas=1, then 1 % 2 != 0 triggers the "num_replicas must be a multiple of gang_size" root validator before the mutual-exclusivity validators, causing both tests to fail with an unexpected error message. Compare with the nearby test_gang_scheduling_config_invalid_num_replicas which correctly sets num_replicas=4.

Additional Locations (1)

• python/ray/serve/tests/unit/test_schema.py#L497-L511

 

[jeffreywang88]

In utils.py, importing from ray.util.state at the top level causes the linter to discover and scan ray.util.state.common when it traverses ray.serve, and there are a lot of APIs under ray.util.state.common that violate API annotation linter.

Therefore, I moved the imports inside get_active_placement_group_ids. This helper is only called upon controller recovery, so it should be alright + adding API annotations for classes under ray.util.state.common is out of scope for this PR.

Here's the buildkite failure: https://buildkite.com/ray-project/premerge/builds/60808/steps/canvas?sid=019c9261-214a-4c16-b979-950cf4674171&tab=output.

[abrarsheikh]

move imports to top

[abrarsheikh]

In utils.py, importing from ray.util.state at the top level causes the linter to discover and scan ray.util.state.common when it traverses ray.serve, and there are a lot of APIs under ray.util.state.common that violate API annotation linter.

why does docs api annotation discover ray.util.state?

[jeffreywang88]

Ray Train adopts the same pattern of importing ray.util.state.list_actors. The reason why the linter doesn't complain about that is:

• import ray.train (train.__init__.py) does NOT transitively load ray.train.v2._internal.state.util, so ray.util.state never enters sys.modules. The linter never discovers it.
• import ray.serve DOES transitively load ray.serve._private.utils (code), which has the top-level from ray.util.state import list_actors. This pulls ray.util.state into sys.modules. Then when the linter scans ray.util, it finds ray.util.state and flags all its unannotated symbols.