Refactor init_app function and add documentation

pier-digital · Mar 3, 2024 · 41f992f · 41f992f
1 parent 2f269b7
commit 41f992f
Show file tree

Hide file tree

Showing 4 changed files with 127 additions and 86 deletions.
diff --git a/README.md b/README.md
@@ -46,7 +46,7 @@ request_model = [
     {"name": "petal width (cm)", "dtype": "float64"},
 ]
 ```
-Alternatively, you can use a pydantic model to define the request model, where the alias field is used to match variable name used in the training dataset:
+Alternatively, you can use a pydantic model to define the request model, where the alias field is used to match the variable names with the column names in the training dataset:
 
 ```python
 class InputData(pydantic.BaseModel):
@@ -82,14 +82,20 @@ pipeline_runner = ml.SklearnPipelineRunner(
 )
 ```
 
-Now you can extend a FastAPI app with the runners:
+Now you can create a FastAPI app with the runners:
+
+```python
+app = ml.init_app(runners=[simple_runner, pipeline_runner])
+```
+
+You can also pass an existing FastAPI app to the `init_app` function:
 
 ```python
 import fastapi
 
 app = fastapi.FastAPI()
 
-app = ml.init_app(app, [simple_runner, pipeline_runner])
+app = ml.init_app(app=app, runners=[simple_runner, pipeline_runner])
 ```
 
 The `init_app` function will add the necessary routes to the FastAPI app to serve the models. You can now start the app with:

diff --git a/example.py b/example.py
@@ -12,9 +12,7 @@ def create_model():
 
     X, y = load_iris(return_X_y=True, as_frame=True)
 
-    X_train, X_test, y_train, y_test = train_test_split(
-        X, y, test_size=0.2, random_state=42
-    )
+    X_train, _, y_train, _ = train_test_split(X, y, test_size=0.2, random_state=42)
 
     model = Pipeline(
         [
@@ -28,12 +26,12 @@ def create_model():
     return model
 
 
-# InputData = [
-#     {"name": "sepal length (cm)", "dtype": "float64"},
-#     {"name": "sepal width (cm)", "dtype": "float64"},
-#     {"name": "petal length (cm)", "dtype": "float64"},
-#     {"name": "petal width (cm)", "dtype": "float64"},
-# ]
+features_metadata = [
+    {"name": "sepal length (cm)", "dtype": "float64"},
+    {"name": "sepal width (cm)", "dtype": "float64"},
+    {"name": "petal length (cm)", "dtype": "float64"},
+    {"name": "petal width (cm)", "dtype": "float64"},
+]
 
 
 class InputData(pydantic.BaseModel):
@@ -49,7 +47,7 @@ class InputData(pydantic.BaseModel):
     name="my simple model",
     predictor=MODEL,
     method_name="predict",
-    request_model=InputData,
+    request_model=InputData,  # OR request_model=features_metadata
 )
 
 pipeline_runner = ml.SklearnPipelineRunner(
@@ -61,7 +59,7 @@ class InputData(pydantic.BaseModel):
 
 app = fastapi.FastAPI()
 
-app = ml.init_app(app, [simple_runner, pipeline_runner])
+app = ml.init_app(app=app, runners=[simple_runner, pipeline_runner])
 
 if __name__ == "__main__":
     import uvicorn

diff --git a/modelib/server/app.py b/modelib/server/app.py
@@ -9,11 +9,23 @@
 
 
 def init_app(
-    app: fastapi.FastAPI,
+    *,
     runners: typing.List[BaseRunner],
+    app: fastapi.FastAPI = fastapi.FastAPI(),
     include_infrastructure: bool = True,
     **runners_router_kwargs,
 ) -> fastapi.FastAPI:
+    """Initialize FastAPI application with Modelib runners.
+
+    Args:
+        runners: List of runners to be included in the application.
+        app: FastAPI application to be initialized. If not provided, a new application will be created.
+        include_infrastructure: Whether to include infrastructure endpoints.
+        **runners_router_kwargs: Additional keyword arguments to be passed to the runners router.
+
+    Returns:
+        FastAPI application with Modelib runners.
+    """
     exceptions.init_app(app)
 
     app.include_router(

diff --git a/tests/test_example.py b/tests/test_example.py
@@ -1,79 +1,12 @@
 import fastapi
 from fastapi.testclient import TestClient
-
-import modelib as ml
+import example
 import pytest
-import pydantic
-
-
-def create_model():
-    from sklearn.datasets import load_iris
-    from sklearn.ensemble import RandomForestClassifier
-    from sklearn.pipeline import Pipeline
-    from sklearn.preprocessing import StandardScaler
-    from sklearn.model_selection import train_test_split
-
-    X, y = load_iris(return_X_y=True, as_frame=True)
-
-    X_train, _, y_train, _ = train_test_split(X, y, test_size=0.2, random_state=42)
-
-    model = Pipeline(
-        [
-            ("scaler", StandardScaler()),
-            ("clf", RandomForestClassifier(random_state=42)),
-        ]
-    ).set_output(transform="pandas")
-
-    model.fit(X_train, y_train)
-
-    return model
-
-
-FEATURES = [
-    {"name": "sepal length (cm)", "dtype": "float64"},
-    {"name": "sepal width (cm)", "dtype": "float64"},
-    {"name": "petal length (cm)", "dtype": "float64"},
-    {"name": "petal width (cm)", "dtype": "float64"},
-]
-
-
-class InputData(pydantic.BaseModel):
-    sepal_length: float = pydantic.Field(alias="sepal length (cm)")
-    sepal_width: float = pydantic.Field(alias="sepal width (cm)")
-    petal_length: float = pydantic.Field(alias="petal length (cm)")
-    petal_width: float = pydantic.Field(alias="petal width (cm)")
-
-
-MODEL = create_model()
-
-
-@pytest.mark.parametrize(
-    "request_model,model",
-    [
-        (FEATURES, MODEL),
-        (InputData, MODEL),
-    ],
-)
-def test_example(request_model, model):
-    simple_runner = ml.SklearnRunner(
-        name="my simple model",
-        predictor=model,
-        method_name="predict",
-        request_model=request_model,
-    )
-
-    pipeline_runner = ml.SklearnPipelineRunner(
-        "Pipeline Model",
-        predictor=model,
-        method_names=["transform", "predict"],
-        request_model=request_model,
-    )
-
-    app = fastapi.FastAPI()
+import modelib as ml
 
-    app = ml.init_app(app, [simple_runner, pipeline_runner])
 
-    client = TestClient(app)
+def test_example():
+    client = TestClient(example.app)
 
     response = client.get("/docs")
 
@@ -118,3 +51,95 @@ def test_example(request_model, model):
             "clf": [0],
         },
     }
+
+
+@pytest.mark.parametrize(
+    "runner,endpoint_path,expected_response",
+    [
+        (
+            ml.SklearnRunner(
+                name="my simple model",
+                predictor=example.MODEL,
+                method_name="predict",
+                request_model=example.InputData,
+            ),
+            "/my-simple-model",
+            {"result": 0},
+        ),
+        (
+            ml.SklearnRunner(
+                name="my 232, simple model",
+                predictor=example.MODEL,
+                method_name="predict",
+                request_model=example.features_metadata,
+            ),
+            "/my-232-simple-model",
+            {"result": 0},
+        ),
+        (
+            ml.SklearnPipelineRunner(
+                "Pipeline Model",
+                predictor=example.MODEL,
+                method_names=["transform", "predict"],
+                request_model=example.InputData,
+            ),
+            "/pipeline-model",
+            {
+                "result": 0,
+                "steps": {
+                    "scaler": [
+                        {
+                            "sepal length (cm)": -7.081194586015879,
+                            "sepal width (cm)": -6.845571885453045,
+                            "petal length (cm)": -2.135591504400147,
+                            "petal width (cm)": -1.5795728805764124,
+                        }
+                    ],
+                    "clf": [0],
+                },
+            },
+        ),
+        (
+            ml.SklearnPipelineRunner(
+                "My SUPER    Pipeline Model",
+                predictor=example.MODEL,
+                method_names=["transform", "predict"],
+                request_model=example.features_metadata,
+            ),
+            "/my-super-pipeline-model",
+            {
+                "result": 0,
+                "steps": {
+                    "scaler": [
+                        {
+                            "sepal length (cm)": -7.081194586015879,
+                            "sepal width (cm)": -6.845571885453045,
+                            "petal length (cm)": -2.135591504400147,
+                            "petal width (cm)": -1.5795728805764124,
+                        }
+                    ],
+                    "clf": [0],
+                },
+            },
+        ),
+    ],
+)
+def test_runners(runner, endpoint_path, expected_response):
+    app = fastapi.FastAPI()
+
+    app = ml.init_app(app=app, runners=[runner])
+
+    client = TestClient(app)
+
+    response = client.post(
+        endpoint_path,
+        json={
+            "sepal length (cm)": 0,
+            "sepal width (cm)": 0,
+            "petal length (cm)": 0,
+            "petal width (cm)": 0,
+        },
+    )
+
+    assert response.status_code == 200
+    assert response.json() == expected_response