[CI/Build] Fix test_defaults_with_usage_context in AMD CI

[zhewenl]

Purpose

Fixes test_defaults_with_usage_context test failure on AMD MI300X and other high-memory GPUs by aligning test logic with implementation.
The test was failing on AMD CI with:

  assert 1024 == 256
  AssertionError: max_num_seqs expected 256, got 1024

It's becuase est logic didn't match implementation logic for determining default scheduler values in vllm/engine/arg_utils.py:

if device_memory >= 70 * GiB_bytes and "a100" not in device_name:
            # For GPUs like H100 and MI300x, use larger default values.
            default_max_num_batched_tokens = {
                UsageContext.LLM_CLASS: 16384,
                UsageContext.OPENAI_API_SERVER: 8192,
            }
            default_max_num_seqs = {
                UsageContext.LLM_CLASS: 1024,
                UsageContext.OPENAI_API_SERVER: 1024,
            }

Test Plan

pytest -v -s tests/v1/engine/test_engine_args.py::test_defaults_with_usage_context
...
tests/v1/engine/test_engine_args.py::test_defaults_with_usage_context INFO 11-01 14:57:37 [model.py:668] Resolved architecture: OPTForCausalLM
INFO 11-01 14:57:37 [model.py:1773] Using max model len 2048
INFO 11-01 14:57:38 [scheduler.py:211] Chunked prefill is enabled with max_num_batched_tokens=16384.
INFO 11-01 14:57:39 [model.py:668] Resolved architecture: OPTForCausalLM
INFO 11-01 14:57:39 [model.py:1773] Using max model len 2048
INFO 11-01 14:57:39 [scheduler.py:211] Chunked prefill is enabled with max_num_batched_tokens=8192.
PASSED

================================================================== warnings summary ==================================================================
<frozen importlib._bootstrap>:488
  <frozen importlib._bootstrap>:488: DeprecationWarning: builtin type SwigPyPacked has no __module__ attribute

<frozen importlib._bootstrap>:488
  <frozen importlib._bootstrap>:488: DeprecationWarning: builtin type SwigPyObject has no __module__ attribute

-- Docs: https://docs.pytest.org/en/stable/how-to/capture-warnings.html
=========================================================== 1 passed, 2 warnings in 2.52s ============================================================

[gemini-code-assist]

Code Review

This pull request fixes a test failure on high-memory GPUs by aligning the test logic with the implementation. The change is correct and addresses the immediate issue. However, it does so by duplicating logic from the implementation into the test file. This creates a maintainability issue, as the two copies of the logic can diverge in the future, leading to similar bugs. I've left a comment suggesting a refactoring to extract the duplicated logic into a shared utility function to prevent this problem from recurring. Addressing this would significantly improve the long-term health of the codebase.

[gemini-code-assist]

While this change correctly aligns the test logic with the implementation in vllm/engine/arg_utils.py, it introduces significant code duplication. The logic for determining default scheduler values, including this if block and the corresponding else block on lines 65-68, is now present in both the implementation and this test. This makes the code harder to maintain, as any future changes to this logic will need to be updated in two places, increasing the risk of them becoming out of sync again (which is what this PR is fixing).

To improve maintainability, I recommend refactoring this logic into a shared, public utility function within vllm.engine.arg_utils. This function could then be imported and used by both _set_default_args in arg_utils.py and this test, eliminating the duplication.

For example, you could create a function in vllm/engine/arg_utils.py:

def get_default_scheduler_values(device_memory: int, device_name: str) -> tuple[dict, dict]:
    """Gets the default max_num_batched_tokens and max_num_seqs based on hardware."""
    from vllm.usage.usage_lib import UsageContext
    from vllm.utils.mem_constants import GiB_bytes

    if device_memory >= 70 * GiB_bytes and "a100" not in device_name:
        # For GPUs like H100 and MI300x, use larger default values.
        default_max_num_batched_tokens = {
            UsageContext.LLM_CLASS: 16384,
            UsageContext.OPENAI_API_SERVER: 8192,
        }
        default_max_num_seqs = {
            UsageContext.LLM_CLASS: 1024,
            UsageContext.OPENAI_API_SERVER: 1024,
        }
    else:
        default_max_num_batched_tokens = {
            UsageContext.LLM_CLASS: 8192,
            UsageContext.OPENAI_API_SERVER: 2048,
        }
        default_max_num_seqs = {
            UsageContext.LLM_CLASS: 256,
            UsageContext.OPENAI_API_SERVER: 256,
        }
    return default_max_num_batched_tokens, default_max_num_seqs

Then, this test could be simplified to call this new function to get the expected values, which would make the test more robust and easier to maintain.