Skip to content

fix: normalize request pressure generation domains#679

Merged
eric-tramel merged 1 commit into
scheduling-yolofrom
andreatgretel/fix/request-pressure-domain-key
May 19, 2026
Merged

fix: normalize request pressure generation domains#679
eric-tramel merged 1 commit into
scheduling-yolofrom
andreatgretel/fix/request-pressure-domain-key

Conversation

@andreatgretel

@andreatgretel andreatgretel commented May 19, 2026

Copy link
Copy Markdown
Contributor

Summary

  • Normalize model config generation types before building scheduling metadata so chat completions use the request-admission chat domain.
  • Add scheduling resolver coverage for GenerationType.CHAT_COMPLETION producing a concrete RequestResourceKey.

Why

GenerationType is a string enum, but str(GenerationType.CHAT_COMPLETION) does not produce the request-domain value expected by TaskSchedulingResolver. That left model-backed tasks without a request-pressure key, so queue ordering could not see pressure for those resources.

Validation

  • .venv/bin/ruff check --fix .
  • .venv/bin/ruff format .
  • .venv/bin/pytest packages/data-designer-engine/tests/engine/dataset_builders/scheduling/test_resolver.py packages/data-designer-engine/tests/engine/dataset_builders/test_async_scheduler.py packages/data-designer-engine/tests/engine/models/request_admission/test_controller.py packages/data-designer-engine/tests/engine/models/clients/test_model_request_executor.py

@eric-tramel eric-tramel self-assigned this May 19, 2026
@eric-tramel eric-tramel marked this pull request as ready for review May 19, 2026 18:58
@eric-tramel eric-tramel requested a review from a team as a code owner May 19, 2026 18:58
@greptile-apps

greptile-apps Bot commented May 19, 2026

Copy link
Copy Markdown
Contributor

Greptile Summary

This PR fixes a bug where str(GenerationType.CHAT_COMPLETION) on a (str, Enum) returned the full enum repr (e.g. "GenerationType.CHAT_COMPLETION") instead of its underlying value ("chat-completion"), leaving model-backed tasks without a request_resource_key so queue pressure was invisible to the scheduler.

  • Adds _scheduling_generation_kind() in base.py to extract .value from the generation-type enum and map "chat-completion" → "chat" (matching RequestDomain.CHAT), replacing the broken str() cast in the endpoint tuple.
  • Updates test_resolver.py to use GenerationType.CHAT_COMPLETION in the deduplication fixture (instead of the raw string "chat") and adds assertions that schedulable.request_resource_key resolves to the expected RequestResourceKey in both the model-metadata test and the deduplication test.

Confidence Score: 5/5

Safe to merge — targeted bug fix with direct test coverage for the corrected code path.

The change is small and well-scoped: a single mapping helper replaces a broken str() cast, and the test suite is extended to exercise the exact code path that was previously silently broken. The helper correctly handles all three cases (enum instance, raw matching string, raw non-matching string) and is backward-compatible with any callers that already pass "chat" directly.

No files require special attention.

Important Files Changed

Filename Overview
packages/data-designer-engine/src/data_designer/engine/column_generators/generators/base.py Adds _scheduling_generation_kind() to normalize the generation-type enum value before it enters the endpoint tuple; correctly handles raw strings, enum instances, and the "chat-completion" → "chat" mapping.
packages/data-designer-engine/tests/engine/dataset_builders/scheduling/test_resolver.py Test fixtures updated to use GenerationType.CHAT_COMPLETION instead of the raw "chat" string; new assertions verify the request_resource_key is populated correctly end-to-end.

Sequence Diagram

 
sequenceDiagram
    participant Gen as ColumnGeneratorWithModelRegistry
    participant Helper as _scheduling_generation_kind
    participant Meta as SchedulingMetadata
    participant Resolver as TaskSchedulingResolver
    participant RRK as _request_resource_key

    Gen->>Helper: generation_type (GenerationType.CHAT_COMPLETION)
    Helper->>Helper: extract .value → "chat-completion"
    Helper->>Helper: map to "chat"
    Helper-->>Gen: "chat"
    Gen->>Meta: model(provider, model_id, "chat", weight)
    Meta-->>Resolver: "identity = (model, provider, model_id, "chat")"
    Resolver->>RRK: metadata
    RRK->>RRK: "RequestDomain("chat") = RequestDomain.CHAT"
    RRK-->>Resolver: RequestResourceKey(provider, model_id, CHAT)
Loading

Reviews (1): Last reviewed commit: "fix request pressure domain metadata" | Re-trigger Greptile

@nabinchha nabinchha mentioned this pull request May 19, 2026
Merged
7 tasks
@eric-tramel eric-tramel merged commit 9a8312a into scheduling-yolo May 19, 2026
4 checks passed
@andreatgretel andreatgretel mentioned this pull request May 19, 2026
Merged
20 tasks
@andreatgretel andreatgretel deleted the andreatgretel/fix/request-pressure-domain-key branch May 20, 2026 15:25
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Reviewers

@eric-tramel eric-tramel eric-tramel approved these changes

Assignees

@eric-tramel eric-tramel

Labels

None yet

Projects

None yet

Milestone

No milestone

Development

Successfully merging this pull request may close these issues.

2 participants

@andreatgretel @eric-tramel