Kedro with FastAPI - improving performance of KedroSession.load_context

[illia-shkroba]

Hello.

I'm trying to build RestAPI with FastAPI that runs Kedro Pipeline under the hood and come up with this solution:

import pathlib
from typing import Any, Iterable

from fastapi import Depends, FastAPI
from kedro.framework.context import KedroContext
from kedro.framework.session import KedroSession
from kedro.framework.startup import bootstrap_project

app = FastAPI(
    title="FastAPI + Kedro",
    version="0.0.1",
    license_info={
        "name": "GNU GENERAL PUBLIC LICENSE",
        "url": "https://www.gnu.org/licenses/gpl-3.0.html",
    },
)


def get_session() -> Iterable[KedroSession]:
    bootstrap_project(pathlib.Path().cwd())
    with KedroSession.create() as session:
        yield session


def get_context(session: KedroSession = Depends(get_session)) -> Iterable[KedroContext]:
    yield session.load_context()


@app.get("/")
def index(
    session: KedroSession = Depends(get_session),
    context: KedroContext = Depends(get_context),
) -> dict[str, Any]:
    session.run("math")
    catalog = context.catalog
    return catalog.load("output")

session.run("math") runs a simple pipeline that calculates a variance for the input: [1, 2, 3].

The solution seems to work as expected, but it takes nearly 2.1 seconds to finish a request:

time curl http://127.0.0.1:8000
# curl http://127.0.0.1:8000  0.00s user 0.00s system 0% cpu 2.097 total

I've noticed that session.load_context() takes about 1 second to finish. Also I've found that load_context() is used by session.run():

        session_id = self.store["session_id"]
        save_version = session_id
        extra_params = self.store.get("extra_params") or {}
        context = self.load_context()

It seems that load_context() is called twice during the request:

1. Inside of get_context().
2. Inside of session.run().

I've tried to cache the result of session.load_context() like this:

def get_context(session: KedroSession = Depends(get_session)) -> Iterable[KedroContext]:
    context = session.load_context()
    session.load_context = lambda: context
    yield context

And by doing that I've decreased the request processing time to 1.06 seconds.

time curl http://127.0.0.1:8000
# curl http://127.0.0.1:8000  0.00s user 0.00s system 0% cpu 1.062 total

Do you have any suggestions on how I can further optimize the session.run()? Should I try a different approach with a plain DataCatalog/SequentialRunner? Or maybe Kedro's implementation of load_context() should be modified to use some caching?

[illia-shkroba]

Duplicate of #2134
Delete this one please.