mlflow/mlflow · error · InvalidTypeHintException

Dictionary key type must be str, got {args[0]} in type hint

Error message

Dictionary key type must be str, got {args[0]} in type hint {_type_hint_repr(type_hint)}

What it means

When inferring a ColSpec from a type hint in mlflow.types.type_hints, dict hints must have str as their key type because MLflow Map dtype only supports string keys. dict[int, X] or dict in Python 3.9+ Any other key type raises InvalidTypeHintException.

Source

Thrown at mlflow/types/type_hints.py:208

        )
        return ColSpecType(dtype=AnyType(), required=True)
    if datatype := TYPE_HINTS_TO_DATATYPE_MAPPING.get(type_hint):
        return ColSpecType(dtype=datatype, required=True)
    elif _is_pydantic_type_hint(type_hint):
        dtype = _infer_type_from_pydantic_model(type_hint)
        return ColSpecType(dtype=dtype, required=True)
    elif origin_type := get_origin(type_hint):
        args = get_args(type_hint)
        if origin_type is list:
            internal_type = _get_element_type_of_list_type_hint(type_hint)
            return ColSpecType(
                dtype=Array(_infer_colspec_type_from_type_hint(type_hint=internal_type).dtype),
                required=True,
            )
        if origin_type is dict:
            if len(args) == 2:
                if args[0] != str:
                    raise InvalidTypeHintException(
                        message=f"Dictionary key type must be str, got {args[0]} in type hint "
                        f"{_type_hint_repr(type_hint)}"
                    )
                return ColSpecType(
                    dtype=Map(_infer_colspec_type_from_type_hint(type_hint=args[1]).dtype),
                    required=True,
                )
            raise InvalidTypeHintException(
                message="Dictionary type hint must contain two element types, got "
                f"{_type_hint_repr(type_hint)}"
            )
        if origin_type in UNION_TYPES:
            if NONE_TYPE in args:
                # This case shouldn't happen, but added for completeness
                if len(args) < 2:
                    raise InvalidTypeHintException(
                        message=f"Union type hint must contain at least one non-None type, "
                        f"got {_type_hint_repr(type_hint)}"

View on GitHub (pinned to 6a27f2decc)

Solutions

Change the key type to str (Dict[str, X]) and convert keys (e.g. str(id)) before prediction
If keys must be non-string, restructure as a List of key/value dataclasses instead of a dict
For enum keys, use the enum's string name as the dict key

Example fix

// before
class Input(pydantic.BaseModel):
    scores: Dict[int, float]
// after
class Input(pydantic.BaseModel):
    scores: Dict[str, float]  # Map dtype requires str keys
# caller: {str(k): v for k, v in raw_scores.items()}

Defensive patterns

Strategy: type-guard

Validate before calling

from typing import get_args, get_origin
for name, t in get_type_hints(MyModel).items():
    if get_origin(t) is dict and get_args(t)[0] is not str:
        raise TypeError(f"{name}: dict keys must be str")

Type guard

def is_str_keyged_map(t) -> bool:
    from typing import get_origin, get_args
    return get_origin(t) is dict and len(get_args(t)) == 2 and get_args(t)[0] is str

Try / catch

try:
    infer_signature(train, model_input)
except Exception as e:
    if "Dictionary key type must be str" in str(e):
        logging.error("Convert dict keys to str: %s", e)
    raise

Prevention

Use Dict[str, X] exclusively in signature type hints
Coerce non-string keys (e.g. str(id)) at prediction boundaries
Document the str-key requirement for shared signature models

When it happens

Trigger: Annotating a pydantic model field or predict signature parameter as Dict[int, float], Dict[Enum, str], Dict[bytes, X], or a non-str key type, then inferring a model signature (mlflow.models.infer_signature or set_signature).

Common situations: Models keyed by integer IDs or enum members; porting general Python typing to MLflow input schemas without realizing Map dtype requires string keys.

Related errors

AI-assisted analysis of mlflow/mlflow@6a27f2decc (2026-08-29). Data as JSON: /api/errors/c992fa7b6843e7ba. Report an issue: GitHub.