Support Python 3.10 runtime compatibility (#5366)

Signed-off-by: Xiaodong Ye <xiaodong.ye@mthreads.com>

apps/ComfyUI-vLLM-Omni/comfyui_vllm_omni/utils/types.py

@@ -1,12 +1,8 @@
 from collections.abc import Callable
 from enum import Enum, auto
-from typing import (
-    Any,
-    Literal,
-    NotRequired,
-    TypeAlias,
-    TypedDict,
-)
+from typing import Any, Literal, TypeAlias
+
+from typing_extensions import NotRequired, TypedDict
 
 AudioFormat: TypeAlias = Literal["mp3", "opus", "aac", "flac", "wav", "pcm"]
 

tests/comfyui/test_comfyui_integration.py

@@ -10,7 +10,7 @@ import multiprocessing
 import time
 import traceback
 from collections.abc import Iterable, Sequence
-from enum import StrEnum, auto
+from enum import Enum
 from types import SimpleNamespace
 from typing import Any, NamedTuple
 
@@ -57,15 +57,15 @@ class SamplingCase(NamedTuple):
     lora: dict | None = None
 
 
-class SamplingKind(StrEnum):
-    IMAGE_NONE = auto()
-    IMAGE_DIFFUSION_SINGLE = auto()
-    UNDERSTANDING_NONE = auto()
-    UNDERSTANDING_AR_LIST = auto()
-    TTS_NONE = auto()
-    TTS_DIFFUSION_SINGLE = auto()
-    VIDEO_NONE = auto()
-    VIDEO_DIFFUSION_SINGLE = auto()
+class SamplingKind(str, Enum):
+    IMAGE_NONE = "image_none"
+    IMAGE_DIFFUSION_SINGLE = "image_diffusion_single"
+    UNDERSTANDING_NONE = "understanding_none"
+    UNDERSTANDING_AR_LIST = "understanding_ar_list"
+    TTS_NONE = "tts_none"
+    TTS_DIFFUSION_SINGLE = "tts_diffusion_single"
+    VIDEO_NONE = "video_none"
+    VIDEO_DIFFUSION_SINGLE = "video_diffusion_single"
 
 
 # Pre-defined arguments to be used in function calls during the tests

tests/entrypoints/test_stream_finish_reason.py

@@ -16,20 +16,10 @@ Key invariants tested:
   - voice/speaker parameter compatibility in chat completions
 """
 
-import enum
 import json
 from unittest.mock import MagicMock
 
 import pytest
-
-# Python 3.10 compat: StrEnum was added in 3.11
-if not hasattr(enum, "StrEnum"):
-
-    class _StrEnum(str, enum.Enum):
-        """Minimal StrEnum backport for Python 3.10."""
-
-    enum.StrEnum = _StrEnum  # type: ignore[attr-defined]
-
 from vllm.entrypoints.openai.chat_completion.protocol import (
     ChatCompletionRequest,
     ChatCompletionResponseStreamChoice,

tests/model_tests/diffusion/config_types.py

@@ -3,7 +3,7 @@ Common definitions for controlling what tests run where.
 """
 
 from collections.abc import Callable
-from enum import StrEnum, auto
+from enum import Enum
 from typing import NamedTuple, TypeAlias
 
 from pytest import MarkDecorator
@@ -15,26 +15,26 @@ from vllm_omni.entrypoints.omni import Omni
 TinyDiffusionBuilder: TypeAlias = Callable[[], str]
 
 
-class DiffusionAccs(StrEnum):
+class DiffusionAccs(str, Enum):
     """Supported acceleration types / test settings for Diffusion Models."""
 
-    HSDP = auto()
-    TEA_CACHE = auto()
-    CACHE_DIT = auto()
-    SEQUENCE_PARALLEL = auto()
-    CFG_PARALLEL = auto()
-    TENSOR_PARALLEL = auto()
-    CPU_OFFLOAD = auto()
-    LAYERWISE_OFFLOAD = auto()
-    VAE_PATCH_PARALLEL = auto()
+    HSDP = "hsdp"
+    TEA_CACHE = "tea_cache"
+    CACHE_DIT = "cache_dit"
+    SEQUENCE_PARALLEL = "sequence_parallel"
+    CFG_PARALLEL = "cfg_parallel"
+    TENSOR_PARALLEL = "tensor_parallel"
+    CPU_OFFLOAD = "cpu_offload"
+    LAYERWISE_OFFLOAD = "layerwise_offload"
+    VAE_PATCH_PARALLEL = "vae_patch_parallel"
 
 
-class DiffusionTasks(StrEnum):
+class DiffusionTasks(str, Enum):
     """Supported tasks for Diffusion Models."""
 
-    TEXT_TO_IMAGE = auto()
-    IMAGE_TO_IMAGE = auto()
-    TEXT_TO_VIDEO = auto()
+    TEXT_TO_IMAGE = "text_to_image"
+    IMAGE_TO_IMAGE = "image_to_image"
+    TEXT_TO_VIDEO = "text_to_video"
     # Text to audio, etc should be added here as needed
 
 

vllm_omni/inputs/data.py

@@ -1,24 +1,16 @@
 import copy
 import pprint
 from dataclasses import asdict, dataclass, field
-from typing import Any, TypeAlias, TypedDict
+from typing import Any, TypeAlias
 
-from vllm.inputs import PromptType
+import torch
+from typing_extensions import NotRequired, TypedDict
+from vllm.inputs import EmbedsPrompt, PromptType, TextPrompt, TokensPrompt
+from vllm.inputs.engine import TokensInput
 from vllm.sampling_params import SamplingParams
 
 from vllm_omni.lora.request import LoRARequest
 
-try:
-    from typing import NotRequired
-except ImportError:
-    # Python < 3.11: use typing_extensions
-    from typing_extensions import NotRequired
-
-
-import torch
-from vllm.inputs import EmbedsPrompt, TextPrompt, TokensPrompt
-from vllm.inputs.engine import TokensInput
-
 
 class OmniTextPrompt(TextPrompt):
     """Text prompt with optional embeddings and stage payloads.

vllm_omni/outputs/output_modality.py

@@ -8,23 +8,8 @@ for type-safe multimodal output routing and tensor merging.
 from __future__ import annotations
 
 import re
-import sys
 from enum import Enum, Flag, auto
 
-if sys.version_info >= (3, 11):
-    from enum import StrEnum
-else:
-    # TODO: Remove this Python 3.10 compatibility shim when Python 3.10 support is dropped.
-    class StrEnum(str, Enum):
-        """``enum.StrEnum`` for Python 3.10, which predates it.
-
-        ``__str__`` is taken from ``str`` so ``str(member)`` yields the member
-        value rather than ``"ClassName.MEMBER"``, matching 3.11 semantics.
-        """
-
-        __str__ = str.__str__
-
-
 _MODALITY_ALIASES: dict[str, str] = {
     "speech": "audio",
     "images": "image",
@@ -38,13 +23,17 @@ _MODALITY_ALIASES: dict[str, str] = {
 }
 
 
-class OutputModalityNames(StrEnum):
+class OutputModalityNames(str, Enum):
     """Keys for output modalities.
 
     TODO: (Alex) Integrate this with the big-flag enum below + throughout the code
     for better type safety (currently only used for output processor).
     """
 
+    # Keep the string conversion provided by ``enum.StrEnum`` while using the
+    # Python 3.10-compatible ``str, Enum`` form.
+    __str__ = str.__str__
+
     TEXT = "text"
     IMAGE = "image"
     AUDIO = "audio"