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?

[datajoely]

This is really cool - analysis, thanks @illia-shkroba !

[noklam]

@illia-shkroba Sorry for the late response. There are definitely rooms to improve to reduce the overhead. Meanwhile, if you pipeline do not involve saving any versioned datasets, using the Runner object directly with pipelines and catalog would mitigate the problem.