Skip to content

Conversation

@psychedelicious
Copy link
Contributor

@psychedelicious psychedelicious commented Oct 2, 2023

What type of PR is this? (check all applicable)

  • Refactor
  • Feature
  • Bug Fix
  • Optimization
  • Documentation Update
  • Community Node Submission

Have you discussed this change with the InvokeAI team?

  • Yes
  • No, because:

Description

feat(api): chore: pydantic & fastapi upgrade

Upgrade pydantic and fastapi to latest.

  • pydantic~=2.4.2
  • fastapi~=103.2
  • fastapi-events~=0.9.1

Big Changes

There are a number of logic changes needed to support pydantic v2. Most changes are very simple, like using the new methods to serialized and deserialize models, but there are a few more complex changes.

Invocations

The biggest change relates to invocation creation, instantiation and validation.

Because pydantic v2 moves all validation logic into the rust pydantic-core, we may no longer directly stick our fingers into the validation pie.

Previously, we (ab)used models and fields to allow invocation fields to be optional at instantiation, but required when invoke() is called. We directly manipulated the fields and invocation models when calling invoke().

With pydantic v2, this is much more involved. Changes to the python wrapper do not propagate down to the rust validation logic - you have to rebuild the model. This causes problem with concurrent access to the invocation classes and is not a free operation.

This logic has been totally refactored and we do not need to change the model any more. The details are in baseinvocation.py, in the InputField function and BaseInvocation.invoke_internal() method.

In the end, this implementation is cleaner.

Invocation Fields

In pydantic v2, you can no longer directly add or remove fields from a model.

Previously, we did this to add the type field to invocations.

Invocation Decorators

With pydantic v2, we instead use the imperative create_model() API to create a new model with the additional field. This is done in baseinvocation.py in the invocation() wrapper.

A similar technique is used for invocation_output().

Minor Changes

There are a number of minor changes around the pydantic v2 models API.

Protected model_ Namespace

All models' pydantic-provided methods and attributes are prefixed with model_ and this is considered a protected namespace. This causes some conflict, because "model" means something to us, and we have a ton of pydantic models with attributes starting with "model_".

Forunately, there are no direct conflicts. However, in any pydantic model where we define an attribute or method that starts with "model_", we must tell set the protected namespaces to an empty tuple.

class IPAdapterModelField(BaseModel):
    model_name: str = Field(description="Name of the IP-Adapter model")
    base_model: BaseModelType = Field(description="Base model")

    model_config = ConfigDict(protected_namespaces=())

Model Serialization

Pydantic models no longer have Model.dict() or Model.json().

Instead, we use Model.model_dump() or Model.model_dump_json().

Model Deserialization

Pydantic models no longer have Model.parse_obj() or Model.parse_raw(), and there are no parse_raw_as() or parse_obj_as() functions.

Instead, you need to create a TypeAdapter object to parse python objects or JSON into a model.

adapter_graph = TypeAdapter(Graph)
deserialized_graph_from_json = adapter_graph.validate_json(graph_json)
deserialized_graph_from_dict = adapter_graph.validate_python(graph_dict)

Field Customisation

Pydantic Fields no longer accept arbitrary args.

Now, you must put all additional arbitrary args in a json_schema_extra arg on the field.

Schema Customisation

FastAPI and pydantic schema generation now follows the OpenAPI version 3.1 spec.

This necessitates two changes:

  • Our schema customization logic has been revised
  • Schema parsing to build node templates has been revised

The specific aren't important, but this does present additional surface area for bugs.

Performance Improvements

Pydantic v2 is a full rewrite with a rust backend. This offers a substantial performance improvement (pydantic claims 5x to 50x depending on the task). We'll notice this the most during serialization and deserialization of sessions/graphs, which happens very very often - a couple times per node.

I haven't done any benchmarks, but anecdotally, graph execution is much faster. Also, very larges graphs - like with massive iterators - are much, much faster.

Related Tickets & Documents

I'm sure there are some that are indirectly related but can't think of them.

QA Instructions, Screenshots, Recordings

This definitely needs thorough smoke testing.

You'll need to update with pip install -e ".[dev,test]" to test this.

Added/updated tests?

  • Yes
  • No

@psychedelicious psychedelicious force-pushed the feat/organise-services branch 2 times, most recently from df5ad7b to 6248c1d Compare October 10, 2023 04:45
@psychedelicious psychedelicious force-pushed the chore/pydantic-fastapi-upgrade branch 3 times, most recently from 9546e14 to 0374ad3 Compare October 11, 2023 01:10
@psychedelicious psychedelicious force-pushed the chore/pydantic-fastapi-upgrade branch from 97902ae to 0766c2f Compare October 11, 2023 01:19
Base automatically changed from feat/organise-services to main October 12, 2023 16:15
@psychedelicious psychedelicious force-pushed the chore/pydantic-fastapi-upgrade branch 5 times, most recently from b35cf52 to 594a88d Compare October 13, 2023 06:57
@psychedelicious psychedelicious marked this pull request as ready for review October 13, 2023 06:58
@psychedelicious
Copy link
Contributor Author

This PR is ready for review and testing.

The whole application should function exactly as it did before (albeit a bit faster). There are no new features or user-facing behaviour changes, but there is a good amount of internal logic changed to accomodate pydantic v2.

@psychedelicious psychedelicious force-pushed the chore/pydantic-fastapi-upgrade branch from 594a88d to f642135 Compare October 13, 2023 07:12
@psychedelicious psychedelicious changed the title chore: pydantic fastapi upgrade chore: pydantic v2 & fastapi upgrade Oct 13, 2023
@psychedelicious psychedelicious force-pushed the chore/pydantic-fastapi-upgrade branch from f642135 to 0aedd6d Compare October 16, 2023 06:05
@psychedelicious
Copy link
Contributor Author

Please do not update this branch by merging main into it. It will likely break the branch.

Upgrade pydantic and fastapi to latest.

- pydantic~=2.4.2
- fastapi~=103.2
- fastapi-events~=0.9.1

**Big Changes**

There are a number of logic changes needed to support pydantic v2. Most changes are very simple, like using the new methods to serialized and deserialize models, but there are a few more complex changes.

**Invocations**

The biggest change relates to invocation creation, instantiation and validation.

Because pydantic v2 moves all validation logic into the rust pydantic-core, we may no longer directly stick our fingers into the validation pie.

Previously, we (ab)used models and fields to allow invocation fields to be optional at instantiation, but required when `invoke()` is called. We directly manipulated the fields and invocation models when calling `invoke()`.

With pydantic v2, this is much more involved. Changes to the python wrapper do not propagate down to the rust validation logic - you have to rebuild the model. This causes problem with concurrent access to the invocation classes and is not a free operation.

This logic has been totally refactored and we do not need to change the model any more. The details are in `baseinvocation.py`, in the `InputField` function and `BaseInvocation.invoke_internal()` method.

In the end, this implementation is cleaner.

**Invocation Fields**

In pydantic v2, you can no longer directly add or remove fields from a model.

Previously, we did this to add the `type` field to invocations.

**Invocation Decorators**

With pydantic v2, we instead use the imperative `create_model()` API to create a new model with the additional field. This is done in `baseinvocation.py` in the `invocation()` wrapper.

A similar technique is used for `invocation_output()`.

**Minor Changes**

There are a number of minor changes around the pydantic v2 models API.

**Protected `model_` Namespace**

All models' pydantic-provided methods and attributes are prefixed with `model_` and this is considered a protected namespace. This causes some conflict, because "model" means something to us, and we have a ton of pydantic models with attributes starting with "model_".

Forunately, there are no direct conflicts. However, in any pydantic model where we define an attribute or method that starts with "model_", we must tell set the protected namespaces to an empty tuple.

```py
class IPAdapterModelField(BaseModel):
    model_name: str = Field(description="Name of the IP-Adapter model")
    base_model: BaseModelType = Field(description="Base model")

    model_config = ConfigDict(protected_namespaces=())
```

**Model Serialization**

Pydantic models no longer have `Model.dict()` or `Model.json()`.

Instead, we use `Model.model_dump()` or `Model.model_dump_json()`.

**Model Deserialization**

Pydantic models no longer have `Model.parse_obj()` or `Model.parse_raw()`, and there are no `parse_raw_as()` or `parse_obj_as()` functions.

Instead, you need to create a `TypeAdapter` object to parse python objects or JSON into a model.

```py
adapter_graph = TypeAdapter(Graph)
deserialized_graph_from_json = adapter_graph.validate_json(graph_json)
deserialized_graph_from_dict = adapter_graph.validate_python(graph_dict)
```

**Field Customisation**

Pydantic `Field`s no longer accept arbitrary args.

Now, you must put all additional arbitrary args in a `json_schema_extra` arg on the field.

**Schema Customisation**

FastAPI and pydantic schema generation now follows the OpenAPI version 3.1 spec.

This necessitates two changes:
- Our schema customization logic has been revised
- Schema parsing to build node templates has been revised

The specific aren't important, but this does present additional surface area for bugs.

**Performance Improvements**

Pydantic v2 is a full rewrite with a rust backend. This offers a substantial performance improvement (pydantic claims 5x to 50x depending on the task). We'll notice this the most during serialization and deserialization of sessions/graphs, which happens very very often - a couple times per node.

I haven't done any benchmarks, but anecdotally, graph execution is much faster. Also, very larges graphs - like with massive iterators - are much, much faster.
@psychedelicious psychedelicious force-pushed the chore/pydantic-fastapi-upgrade branch from 0aedd6d to d9a2249 Compare October 17, 2023 03:16
@Millu
Copy link
Contributor

Millu commented Oct 17, 2023

Let's get this in!

@psychedelicious psychedelicious merged commit a094f4c into main Oct 17, 2023
@psychedelicious psychedelicious deleted the chore/pydantic-fastapi-upgrade branch October 17, 2023 03:59
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

4 participants