Serve sanity checks

.github/workflows/ci-testing.yml

@@ -106,6 +106,11 @@ jobs:
         pip list
       shell: bash
 
+    - name: Install serve test dependencies
+      if: matrix.topic == 'serve'
+      run: |
+        pip install '.[all]' --pre --upgrade
+
     - name: Cache datasets
       uses: actions/cache@v2
       with:
@@ -115,7 +120,8 @@ jobs:
 
     - name: Tests
       env:
-          FIFTYONE_DO_NOT_TRACK: true
+        FLASH_TEST_TOPIC: ${{ matrix.topic }}
+        FIFTYONE_DO_NOT_TRACK: true
       run: |
         # tox --sitepackages
         coverage run --source flash -m pytest flash tests -v --junitxml=junit/test-results-${{ runner.os }}-${{ matrix.python-version }}-${{ matrix.requires }}.xml
@@ -143,3 +149,8 @@ jobs:
         env_vars: OS,PYTHON
         name: codecov-umbrella
         fail_ci_if_error: false
+
+    - name: Uninstall
+      run: |
+        pip uninstall lightning-flash -y
+      shell: bash

CHANGELOG.md

@@ -12,6 +12,7 @@ The format is based on [Keep a Changelog](http://keepachangelog.com/en/1.0.0/).
 - Added support for `torch.jit` to tasks where possible and documented task JIT compatibility ([#389](https://github.com/PyTorchLightning/lightning-flash/pull/389))
 - Added option to provide a `Sampler` to the `DataModule` to use when creating a `DataLoader` ([#390](https://github.com/PyTorchLightning/lightning-flash/pull/390))
 - Added support for multi-label text classification and toxic comments example ([#401](https://github.com/PyTorchLightning/lightning-flash/pull/401))
+- Added a sanity checking feature to flash.serve ([#423](https://github.com/PyTorchLightning/lightning-flash/pull/423))
 
 ### Changed
 

flash/__init__.py

@@ -27,6 +27,7 @@
     from flash.core.trainer import Trainer  # noqa: E402
 
     _PACKAGE_ROOT = os.path.dirname(__file__)
+    ASSETS_ROOT = os.path.join(_PACKAGE_ROOT, "assets")
     PROJECT_ROOT = os.path.dirname(_PACKAGE_ROOT)
     _IS_TESTING = os.getenv("FLASH_TESTING", "0") == "1"
 

flash/core/classification.py

@@ -11,7 +11,7 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
-from typing import Any, Callable, Dict, List, Mapping, Optional, Sequence, Tuple, Union
+from typing import Any, Callable, Dict, List, Mapping, Optional, Sequence, Union
 
 import torch
 import torch.nn.functional as F
@@ -24,7 +24,6 @@
 from flash.core.utilities.imports import _FIFTYONE_AVAILABLE
 
 if _FIFTYONE_AVAILABLE:
-    import fiftyone as fo
     from fiftyone.core.labels import Classification, Classifications
 else:
     Classification, Classifications = None, None
@@ -121,7 +120,8 @@ def __init__(self, multi_label: bool = False, threshold: float = 0.5):
 
     def serialize(self, sample: Any) -> Union[int, List[int]]:
         sample = sample[DefaultDataKeys.PREDS] if isinstance(sample, Dict) else sample
-        sample = torch.tensor(sample)
+        if not torch.is_tensor(sample):
+            sample = torch.tensor(sample)
         if self.multi_label:
             one_hot = (sample.sigmoid() > self.threshold).int().tolist()
             result = []
@@ -153,7 +153,8 @@ def __init__(self, labels: Optional[List[str]] = None, multi_label: bool = False
             self.set_state(LabelsState(labels))
 
     def serialize(self, sample: Any) -> Union[int, List[int], str, List[str]]:
-        sample = sample[DefaultDataKeys.PREDS] if isinstance(sample, Dict) else sample
+        if isinstance(sample, Dict) and DefaultDataKeys.PREDS in sample:
+            sample = sample[DefaultDataKeys.PREDS]
         sample = torch.tensor(sample)
         labels = None
 

flash/core/data/data_pipeline.py

@@ -119,6 +119,10 @@ def initialize(self, data_pipeline_state: Optional[DataPipelineState] = None) ->
         data_pipeline_state._initialized = True  # TODO: Not sure we need this
         return data_pipeline_state
 
+    @property
+    def example_input(self) -> str:
+        return self._deserializer.example_input
+
     @staticmethod
     def _is_overriden(method_name: str, process_obj, super_obj: Any, prefix: Optional[str] = None) -> bool:
         """

flash/core/data/process.py

@@ -13,7 +13,7 @@
 # limitations under the License.
 import os
 from abc import ABC, abstractclassmethod, abstractmethod
-from typing import Any, Callable, Dict, List, Mapping, Optional, Sequence, TYPE_CHECKING
+from typing import Any, Callable, Dict, List, Mapping, Optional, Sequence
 
 import torch
 from pytorch_lightning.trainer.states import RunningStage
@@ -569,6 +569,11 @@ class Deserializer(Properties):
     def deserialize(self, sample: Any) -> Any:  # TODO: Output must be a tensor???
         raise NotImplementedError
 
+    @property
+    @abstractmethod
+    def example_input(self) -> str:
+        pass
+
     def __call__(self, sample: Any) -> Any:
         return self.deserialize(sample)
 

flash/core/model.py

@@ -13,6 +13,7 @@
 # limitations under the License.
 import functools
 import inspect
+import os
 from copy import deepcopy
 from importlib import import_module
 from typing import Any, Callable, Dict, List, Mapping, Optional, Sequence, Tuple, Type, Union
@@ -592,7 +593,9 @@ def configure_callbacks(self):
         if flash._IS_TESTING and torch.cuda.is_available():
             return [BenchmarkConvergenceCI()]
 
-    def serve(self, host: str = "127.0.0.1", port: int = 8000) -> 'Composition':
+    def serve(self, host: str = "127.0.0.1", port: int = 8000, sanity_check: bool = True) -> 'Composition':
+        from fastapi.testclient import TestClient
+
         from flash.core.serve.flash_components import FlashInputs, FlashOutputs
 
         class FlashServeModelComponent(ModelComponent):
@@ -626,7 +629,19 @@ def predict(self, inputs):
                     preds = self.postprocessor(preds)
                     return preds
 
+        if sanity_check:
+            print("Running sanity check")
+            comp = FlashServeModelComponent(self)
+            composition = Composition(predict=comp, TESTING=True, DEBUG=True)
+            app = composition.serve(host="0.0.0.0", port=8000)
+
+            with TestClient(app) as tc:
+                input_str = self.data_pipeline._deserializer.example_input
+                body = {"session": "UUID", "payload": {"inputs": {"data": input_str}}}
+                resp = tc.post("http://0.0.0.0:8000/predict", json=body)
+                print(f"Sanity check response: {resp.json()}")
+
         comp = FlashServeModelComponent(self)
-        composition = Composition(predict=comp)
+        composition = Composition(predict=comp, TESTING=os.environ["FLASH_TESTING"] == "1")
         composition.serve(host=host, port=port)
         return composition

flash/core/utilities/imports.py

@@ -12,9 +12,9 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 """General utilities"""
-
 import importlib
 import operator
+import os
 from importlib.util import find_spec
 
 from pkg_resources import DistributionNotFound
@@ -94,3 +94,19 @@ def _compare_version(package: str, op, version) -> bool:
 _VIDEO_AVAILABLE = _PYTORCHVIDEO_AVAILABLE
 _IMAGE_AVAILABLE = _TORCHVISION_AVAILABLE and _TIMM_AVAILABLE and _KORNIA_AVAILABLE
 _SERVE_AVAILABLE = _FASTAPI_AVAILABLE and _PYDANTIC_AVAILABLE and _CYTOOLZ_AVAILABLE and _UVICORN_AVAILABLE
+
+_IMAGE_TESTING = False
+_VIDEO_TESTING = False
+_TABULAR_TESTING = False
+_TEXT_TESTING = False
+_IMAGE_STLYE_TRANSFER_TESTING = False
+_SERVE_TESTING = False
+
+if "FLASH_TEST_TOPIC" in os.environ:
+    topic = os.environ["FLASH_TEST_TOPIC"]
+    _IMAGE_TESTING = topic == "image"
+    _VIDEO_TESTING = topic == "video"
+    _TABULAR_TESTING = topic == "tabular"
+    _TEXT_TESTING = topic == "text"
+    _IMAGE_STLYE_TRANSFER_TESTING = topic == "image_style_transfer"
+    _SERVE_TESTING = topic == "serve"

flash/image/classification/data.py

@@ -11,8 +11,6 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
-import base64
-from io import BytesIO
 from typing import Any, Callable, Dict, List, Optional, Tuple, Union
 
 import numpy as np
@@ -25,43 +23,29 @@
 from flash.core.data.data_module import DataModule
 from flash.core.data.data_source import DefaultDataKeys, DefaultDataSources
 from flash.core.data.process import Deserializer, Preprocess
-from flash.core.utilities.imports import _IMAGE_AVAILABLE, _MATPLOTLIB_AVAILABLE, _TORCHVISION_AVAILABLE
+from flash.core.utilities.imports import _IMAGE_AVAILABLE, _MATPLOTLIB_AVAILABLE
 from flash.image.classification.transforms import default_transforms, train_default_transforms
-from flash.image.data import ImageFiftyOneDataSource, ImageNumpyDataSource, ImagePathsDataSource, ImageTensorDataSource
+from flash.image.data import (
+    ImageDeserializer,
+    ImageFiftyOneDataSource,
+    ImageNumpyDataSource,
+    ImagePathsDataSource,
+    ImageTensorDataSource,
+)
 
 if _MATPLOTLIB_AVAILABLE:
     import matplotlib.pyplot as plt
 else:
     plt = None
 
-if _TORCHVISION_AVAILABLE:
-    import torchvision
-
 if _IMAGE_AVAILABLE:
     from PIL import Image
-    from PIL import Image as PILImage
 else:
 
     class Image:
         Image = None
 
 
-class ImageClassificationDeserializer(Deserializer):
-
-    def __init__(self):
-
-        self.to_tensor = torchvision.transforms.ToTensor()
-
-    def deserialize(self, data: str) -> Dict:
-        encoded_with_padding = (data + "===").encode("ascii")
-        img = base64.b64decode(encoded_with_padding)
-        buffer = BytesIO(img)
-        img = PILImage.open(buffer, mode="r")
-        return {
-            DefaultDataKeys.INPUT: img,
-        }
-
-
 class ImageClassificationPreprocess(Preprocess):
 
     def __init__(
@@ -88,7 +72,7 @@ def __init__(
                 DefaultDataSources.NUMPY: ImageNumpyDataSource(),
                 DefaultDataSources.TENSORS: ImageTensorDataSource(),
             },
-            deserializer=deserializer or ImageClassificationDeserializer(),
+            deserializer=deserializer or ImageDeserializer(),
             default_data_source=DefaultDataSources.FILES,
         )
 

flash/image/data.py

@@ -11,25 +11,59 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
+import base64
+from io import BytesIO
+from pathlib import Path
 from typing import Any, Dict, Optional
 
 import torch
 
+import flash
 from flash.core.data.data_source import (
     DefaultDataKeys,
     FiftyOneDataSource,
     NumpyDataSource,
     PathsDataSource,
     TensorDataSource,
 )
-from flash.core.utilities.imports import _TORCHVISION_AVAILABLE
+from flash.core.data.process import Deserializer
+from flash.core.utilities.imports import _IMAGE_AVAILABLE, _TORCHVISION_AVAILABLE
 
 if _TORCHVISION_AVAILABLE:
+    import torchvision
     from torchvision.datasets.folder import default_loader, IMG_EXTENSIONS
     from torchvision.transforms.functional import to_pil_image
 else:
     IMG_EXTENSIONS = []
 
+if _IMAGE_AVAILABLE:
+    from PIL import Image as PILImage
+else:
+
+    class Image:
+        Image = None
+
+
+class ImageDeserializer(Deserializer):
+
+    def __init__(self):
+        super().__init__()
+        self.to_tensor = torchvision.transforms.ToTensor()
+
+    def deserialize(self, data: str) -> Dict:
+        encoded_with_padding = (data + "===").encode("ascii")
+        img = base64.b64decode(encoded_with_padding)
+        buffer = BytesIO(img)
+        img = PILImage.open(buffer, mode="r")
+        return {
+            DefaultDataKeys.INPUT: img,
+        }
+
+    @property
+    def example_input(self) -> str:
+        with (Path(flash.ASSETS_ROOT) / "fish.jpg").open("rb") as f:
+            return base64.b64encode(f.read()).decode("UTF-8")
+
 
 class ImagePathsDataSource(PathsDataSource):
 

flash/image/segmentation/data.py

@@ -38,6 +38,7 @@
 )
 from flash.core.data.process import Deserializer, Preprocess
 from flash.core.utilities.imports import _FIFTYONE_AVAILABLE, _IMAGE_AVAILABLE, _MATPLOTLIB_AVAILABLE
+from flash.image.data import ImageDeserializer
 from flash.image.segmentation.serialization import SegmentationLabels
 from flash.image.segmentation.transforms import default_transforms, train_default_transforms
 
@@ -215,19 +216,13 @@ def predict_load_sample(self, sample: Mapping[str, Any]) -> Mapping[str, Any]:
         return sample
 
 
-class SemanticSegmentationDeserializer(Deserializer):
-
-    def __init__(self):
-
-        self.to_tensor = torchvision.transforms.ToTensor()
+class SemanticSegmentationDeserializer(ImageDeserializer):
 
     def deserialize(self, data: str) -> torch.Tensor:
-        encoded_with_padding = (data + "===").encode("ascii")
-        img = base64.b64decode(encoded_with_padding)
-        buffer = BytesIO(img)
-        img = PILImage.open(buffer, mode="r")
-        img = self.to_tensor(img)
-        return {DefaultDataKeys.INPUT: img, DefaultDataKeys.METADATA: {"size": img.shape}}
+        result = super().deserialize(data)
+        result[DefaultDataKeys.INPUT] = self.to_tensor(result[DefaultDataKeys.INPUT])
+        result[DefaultDataKeys.METADATA] = {"size": result[DefaultDataKeys.INPUT].shape}
+        return result
 
 
 class SemanticSegmentationPreprocess(Preprocess):