Repository: cortexlabs/cortex (language: go)

Issue:
Support other predictor kinds for AsyncAPI
Add support for Tensorflow and ONNX predictor types.

Additional notes:
* Disable serverside batching for TensorFlow predictor on AsyncAPI kind because AsyncAPI doesn't support multiple processes/threads


Discussion / hints:

Fix patch:
diff --git a/docs/workloads/async/configuration.md b/docs/workloads/async/configuration.md
index da953228d5..5120914ae0 100644
--- a/docs/workloads/async/configuration.md
+++ b/docs/workloads/async/configuration.md
@@ -32,6 +32,52 @@ predictor:
   shm_size: <string>  # size of shared memory (/dev/shm) for sharing data between multiple processes, e.g. 64Mi or 1Gi (default: Null)
 ```
 
+### Tensorflow Predictor
+
+<!-- CORTEX_VERSION_BRANCH_STABLE x3 -->
+
+```yaml
+predictor:
+  type: tensorflow
+  path: <string>  # path to a python file with a TensorFlowPredictor class definition, relative to the Cortex root (required)
+  dependencies: # (optional)
+    pip: <string>  # relative path to requirements.txt (default: requirements.txt)
+    conda: <string>  # relative path to conda-packages.txt (default: conda-packages.txt)
+    shell: <string>  # relative path to a shell script for system package installation (default: dependencies.sh)
+  models:  # (required)
+    path: <string> # S3/GCS path to an exported SavedModel directory (e.g. s3://my-bucket/exported_model/) (either this, 'dir', or 'paths' must be provided)
+    signature_key:  # name of the signature def to use for prediction (required if your model has more than one signature def)
+  config: <string: value>  # arbitrary dictionary passed to the constructor of the Predictor (optional)
+  python_path: <string>  # path to the root of your Python folder that will be appended to PYTHONPATH (default: folder containing cortex.yaml)
+  image: <string>  # docker image to use for the Predictor (default: quay.io/cortexlabs/tensorflow-predictor:master)
+  tensorflow_serving_image: <string>  # docker image to use for the TensorFlow Serving container (default: quay.io/cortexlabs/tensorflow-serving-cpu:master, quay.io/cortexlabs/tensorflow-serving-gpu:master, or quay.io/cortexlabs/tensorflow-serving-inf:master based on compute)
+  env: <string: string>  # dictionary of environment variables
+  log_level: <string>  # log level that can be "debug", "info", "warning" or "error" (default: "info")
+  shm_size: <string>  # size of shared memory (/dev/shm) for sharing data between multiple processes, e.g. 64Mi or 1Gi (default: Null)
+```
+
+### ONNX Predictor
+
+<!-- CORTEX_VERSION_BRANCH_STABLE x3 -->
+
+```yaml
+predictor:
+  type: onnx
+  path: <string>  # path to a python file with an ONNXPredictor class definition, relative to the Cortex root (required)
+  dependencies: # (optional)
+    pip: <string>  # relative path to requirements.txt (default: requirements.txt)
+    conda: <string>  # relative path to conda-packages.txt (default: conda-packages.txt)
+    shell: <string>  # relative path to a shell script for system package installation (default: dependencies.sh)
+  models:  # (required)
+    path: <string> # S3/GCS path to an exported model directory (e.g. s3://my-bucket/exported_model/) (either this, 'dir', or 'paths' must be provided)
+  config: <string: value>  # arbitrary dictionary passed to the constructor of the Predictor (optional)
+  python_path: <string>  # path to the root of your Python folder that will be appended to PYTHONPATH (default: folder containing cortex.yaml)
+  image: <string>  # docker image to use for the Predictor (default: quay.io/cortexlabs/onnx-predictor-cpu:master or quay.io/cortexlabs/onnx-predictor-gpu:master based on compute)
+  env: <string: string>  # dictionary of environment variables
+  log_level: <string>  # log level that can be "debug", "info", "warning" or "error" (default: "info")
+  shm_size: <string>  # size of shared memory (/dev/shm) for sharing data between multiple processes, e.g. 64Mi or 1Gi (default: Null)
+```
+
 ## Compute
 
 ```yaml
diff --git a/docs/workloads/async/predictors.md b/docs/workloads/async/predictors.md
index 601d2ebe4e..c56136aada 100644
--- a/docs/workloads/async/predictors.md
+++ b/docs/workloads/async/predictors.md
@@ -88,6 +88,125 @@ learn about how headers can be used to change the type of `payload` that is pass
 At this moment, the AsyncAPI `predict` method can only return `JSON`-parseable objects. Navigate to
 the [API responses](#api-responses) section to learn about how to configure it.
 
+## TensorFlow Predictor
+
+**Uses TensorFlow version 2.3.0 by default**
+
+### Interface
+
+```python
+class TensorFlowPredictor:
+    def __init__(self, config, tensorflow_client, metrics_client):
+        """(Required) Called once before the API becomes available. Performs
+        setup such as downloading/initializing a vocabulary.
+
+        Args:
+            config (required): Dictionary passed from API configuration (if
+                specified).
+            tensorflow_client (required): TensorFlow client which is used to
+                make predictions. This should be saved for use in predict().
+            metrics_client (optional): The cortex metrics client, which allows
+                you to push custom metrics in order to build custom dashboards
+                in grafana.
+        """
+        self.client = tensorflow_client
+        # Additional initialization may be done here
+
+    def predict(self, payload, request_id):
+        """(Required) Called once per request. Preprocesses the request payload
+        (if necessary), runs inference (e.g. by calling
+        self.client.predict(model_input)), and postprocesses the inference
+        output (if necessary).
+
+        Args:
+            payload (optional): The request payload (see below for the possible
+                payload types).
+            request_id (optional): The request id string that identifies a workload
+
+        Returns:
+            Prediction or a batch of predictions.
+        """
+        pass
+```
+
+<!-- CORTEX_VERSION_MINOR -->
+
+Cortex provides a `tensorflow_client` to your Predictor's constructor. `tensorflow_client` is an instance
+of [TensorFlowClient](https://github.com/cortexlabs/cortex/tree/master/pkg/cortex/serve/cortex_internal/lib/client/tensorflow.py)
+that manages a connection to a TensorFlow Serving container to make predictions using your model. It should be saved as
+an instance variable in your Predictor, and your `predict()` function should call `tensorflow_client.predict()` to make
+an inference with your exported TensorFlow model. Preprocessing of the JSON payload and postprocessing of predictions
+can be implemented in your `predict()` function as well.
+
+For proper separation of concerns, it is recommended to use the constructor's `config` parameter for information such as
+from where to download the model and initialization files, or any configurable model parameters. You define `config` in
+your API configuration, and it is passed through to your Predictor's constructor.
+
+Your API can accept requests with different types of payloads. Navigate to the [API requests](#api-requests) section to
+learn about how headers can be used to change the type of `payload` that is passed into your `predict` method.
+
+At this moment, the AsyncAPI `predict` method can only return `JSON`-parseable objects. Navigate to
+the [API responses](#api-responses) section to learn about how to configure it.
+
+## ONNX Predictor
+
+**Uses ONNX Runtime version 1.6.0 by default**
+
+### Interface
+
+```python
+class ONNXPredictor:
+    def __init__(self, config, onnx_client, metrics_client):
+        """(Required) Called once before the API becomes available. Performs
+        setup such as downloading/initializing a vocabulary.
+
+        Args:
+            onnx_client (required): ONNX client which is used to make
+                predictions. This should be saved for use in predict().
+            config (required): Dictionary passed from API configuration (if
+                specified).
+            metrics_client (optional): The cortex metrics client, which allows
+                you to push custom metrics in order to build custom dashboards
+                in grafana.
+        """
+        self.client = onnx_client
+        # Additional initialization may be done here
+
+    def predict(self, payload, request_id):
+        """(Required) Called once per request. Preprocesses the request payload
+        (if necessary), runs inference (e.g. by calling
+        self.client.predict(model_input)), and postprocesses the inference
+        output (if necessary).
+
+        Args:
+            payload (optional): The request payload (see below for the possible
+                payload types).
+            request_id (optional): The request id string that identifies a workload
+
+        Returns:
+            Prediction or a batch of predictions.
+        """
+        pass
+```
+<!-- CORTEX_VERSION_MINOR -->
+
+Cortex provides an `onnx_client` to your Predictor's constructor. `onnx_client` is an instance
+of [ONNXClient](https://github.com/cortexlabs/cortex/tree/master/pkg/cortex/serve/cortex_internal/lib/client/onnx.py)
+that manages an ONNX Runtime session to make predictions using your model. It should be saved as an instance variable in
+your Predictor, and your `predict()` function should call `onnx_client.predict()` to make an inference with your
+exported ONNX model. Preprocessing of the JSON payload and postprocessing of predictions can be implemented in
+your `predict()` function as well.
+
+For proper separation of concerns, it is recommended to use the constructor's `config` parameter for information such as
+from where to download the model and initialization files, or any configurable model parameters. You define `config` in
+your API configuration, and it is passed through to your Predictor's constructor.
+
+Your API can accept requests with different types of payloads. Navigate to the [API requests](#api-requests) section to
+learn about how headers can be used to change the type of `payload` that is passed into your `predict` method.
+
+At this moment, the AsyncAPI `predict` method can only return `JSON`-parseable objects. Navigate to
+the [API responses](#api-responses) section to learn about how to configure it.
+
 ## API requests
 
 The type of the `payload` parameter in `predict(self, payload)` can vary based on the content type of the request.
diff --git a/pkg/cortex/serve/cortex_internal/lib/api/async.py b/pkg/cortex/serve/cortex_internal/lib/api/async.py
index 7e662c7e4c..8a594c6c25 100644
--- a/pkg/cortex/serve/cortex_internal/lib/api/async.py
+++ b/pkg/cortex/serve/cortex_internal/lib/api/async.py
@@ -24,9 +24,19 @@
 import dill
 
 from cortex_internal.lib.api.validations import validate_class_impl
+from cortex_internal.lib.client.onnx import ONNXClient
+from cortex_internal.lib.client.tensorflow import TensorFlowClient
 from cortex_internal.lib.exceptions import CortexException, UserException, UserRuntimeException
 from cortex_internal.lib.metrics import MetricsClient
+from cortex_internal.lib.model import ModelsHolder
 from cortex_internal.lib.storage import S3
+from cortex_internal.lib.type import (
+    predictor_type_from_api_spec,
+    TensorFlowPredictorType,
+    TensorFlowNeuronPredictorType,
+    ONNXPredictorType,
+    PythonPredictorType,
+)
 
 ASYNC_PYTHON_PREDICTOR_VALIDATION = {
     "required": [
@@ -43,6 +53,36 @@
     ],
 }
 
+ASYNC_TENSORFLOW_PREDICTOR_VALIDATION = {
+    "required": [
+        {
+            "name": "__init__",
+            "required_args": ["self", "tensorflow_client", "config"],
+            "optional_args": ["metrics_client"],
+        },
+        {
+            "name": "predict",
+            "required_args": ["self"],
+            "optional_args": ["payload", "request_id"],
+        },
+    ],
+}
+
+ASYNC_ONNX_PREDICTOR_VALIDATION = {
+    "required": [
+        {
+            "name": "__init__",
+            "required_args": ["self", "onnx_client", "config"],
+            "optional_args": ["metrics_client"],
+        },
+        {
+            "name": "predict",
+            "required_args": ["self"],
+            "optional_args": ["payload", "request_id"],
+        },
+    ],
+}
+
 
 class AsyncAPI:
     def __init__(
@@ -50,14 +90,19 @@ def __init__(
         api_spec: Dict[str, Any],
         storage: S3,
         storage_path: str,
+        model_dir: str,
         statsd_host: str,
         statsd_port: int,
+        lock_dir: str = "/run/cron",
     ):
         self.api_spec = api_spec
         self.storage = storage
         self.storage_path = storage_path
         self.path = api_spec["predictor"]["path"]
         self.config = api_spec["predictor"].get("config", {})
+        self.type = predictor_type_from_api_spec(api_spec)
+        self.model_dir = model_dir
+        self.lock_dir = lock_dir
 
         datadog.initialize(statsd_host=statsd_host, statsd_port=statsd_port)
         self.__statsd = datadog.statsd
@@ -116,7 +161,13 @@ def delete_payload(self, request_id: str):
         key = f"{self.storage_path}/{request_id}/payload"
         self.storage.delete(key)
 
-    def initialize_impl(self, project_dir: str, metrics_client: MetricsClient):
+    def initialize_impl(
+        self,
+        project_dir: str,
+        metrics_client: MetricsClient,
+        tf_serving_host: str = None,
+        tf_serving_port: str = None,
+    ):
         predictor_impl = self._get_impl(project_dir)
         constructor_args = inspect.getfullargspec(predictor_impl.__init__).args
         config = deepcopy(self.config)
@@ -127,6 +178,23 @@ def initialize_impl(self, project_dir: str, metrics_client: MetricsClient):
         if "metrics_client" in constructor_args:
             args["metrics_client"] = metrics_client
 
+        if self.type in [TensorFlowPredictorType, TensorFlowNeuronPredictorType]:
+            tf_serving_address = tf_serving_host + ":" + tf_serving_port
+            tf_client = TensorFlowClient(
+                tf_serving_url=tf_serving_address,
+                api_spec=self.api_spec,
+            )
+            tf_client.sync_models(lock_dir=self.lock_dir)
+            args["tensorflow_client"] = tf_client
+        elif self.type == ONNXPredictorType:
+            models = ModelsHolder(self.type, self.model_dir)
+            onnx_client = ONNXClient(
+                api_spec=self.api_spec,
+                models=models,
+                model_dir=self.model_dir,
+            )
+            args["onnx_client"] = onnx_client
+
         try:
             predictor = predictor_impl(**args)
         except Exception as e:
@@ -135,16 +203,28 @@ def initialize_impl(self, project_dir: str, metrics_client: MetricsClient):
         return predictor
 
     def _get_impl(self, project_dir: str):
+        if self.type in [TensorFlowPredictorType, TensorFlowNeuronPredictorType]:
+            target_class_name = "TensorFlowPredictor"
+            validations = ASYNC_TENSORFLOW_PREDICTOR_VALIDATION
+        elif self.type == ONNXPredictorType:
+            target_class_name = "ONNXPredictor"
+            validations = ASYNC_ONNX_PREDICTOR_VALIDATION
+        elif self.type == PythonPredictorType:
+            target_class_name = "PythonPredictor"
+            validations = ASYNC_PYTHON_PREDICTOR_VALIDATION
+        else:
+            raise CortexException(f"invalid predictor type: {self.type}")
+
         try:
             impl = self._read_impl(
-                "cortex_async_predictor", os.path.join(project_dir, self.path), "PythonPredictor"
+                "cortex_async_predictor", os.path.join(project_dir, self.path), target_class_name
             )
         except CortexException as e:
             e.wrap("error in " + self.path)
             raise
 
         try:
-            self._validate_impl(impl)
+            validate_class_impl(impl, validations)
         except CortexException as e:
             e.wrap("error in " + self.path)
             raise
@@ -182,7 +262,3 @@ def _read_impl(module_name: str, impl_path: str, target_class_name):
             raise UserException(f"{target_class_name} class is not defined")
 
         return predictor_class
-
-    @staticmethod
-    def _validate_impl(impl):
-        return validate_class_impl(impl, ASYNC_PYTHON_PREDICTOR_VALIDATION)
diff --git a/pkg/cortex/serve/cortex_internal/lib/api/predictor.py b/pkg/cortex/serve/cortex_internal/lib/api/predictor.py
index 6520173517..aee5aa0031 100644
--- a/pkg/cortex/serve/cortex_internal/lib/api/predictor.py
+++ b/pkg/cortex/serve/cortex_internal/lib/api/predictor.py
@@ -224,7 +224,7 @@ def initialize_client(
                     self.models_tree,
                 )
                 if not self.caching_enabled:
-                    cron = TFSAPIServingThreadUpdater(interval=5.0, client=client._client)
+                    cron = TFSAPIServingThreadUpdater(interval=5.0, client=client)
                     cron.start()
 
             if self.type == ONNXPredictorType:
diff --git a/pkg/cortex/serve/cortex_internal/lib/client/onnx.py b/pkg/cortex/serve/cortex_internal/lib/client/onnx.py
index 76fcd52602..5d505dfb47 100644
--- a/pkg/cortex/serve/cortex_internal/lib/client/onnx.py
+++ b/pkg/cortex/serve/cortex_internal/lib/client/onnx.py
@@ -12,10 +12,9 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-import os
 import datetime
+import os
 import threading as td
-import multiprocessing as mp
 from typing import Any, Tuple, Optional
 
 try:
@@ -55,7 +54,7 @@ def __init__(
         api_spec: dict,
         models: ModelsHolder,
         model_dir: str,
-        models_tree: Optional[ModelsTree],
+        models_tree: Optional[ModelsTree] = None,
         lock_dir: Optional[str] = "/run/cron",
     ):
         """
diff --git a/pkg/cortex/serve/cortex_internal/lib/client/tensorflow.py b/pkg/cortex/serve/cortex_internal/lib/client/tensorflow.py
index ba68b2dc73..25eced9103 100644
--- a/pkg/cortex/serve/cortex_internal/lib/client/tensorflow.py
+++ b/pkg/cortex/serve/cortex_internal/lib/client/tensorflow.py
@@ -11,18 +11,20 @@
 # 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 json
 import os
-import copy
-import grpc
 import threading as td
-from typing import Any, Dict, Optional, List
+from typing import Any, Optional, List
+
+import grpc
 
+from cortex_internal import consts
 from cortex_internal.lib.exceptions import (
     UserRuntimeException,
     UserException,
     WithBreak,
 )
+from cortex_internal.lib.log import configure_logger
 from cortex_internal.lib.model import (
     TensorFlowServingAPI,
     ModelsHolder,
@@ -31,8 +33,6 @@
     LockedModelsTree,
     get_models_from_api_spec,
 )
-from cortex_internal import consts
-from cortex_internal.lib.log import configure_logger
 
 logger = configure_logger("cortex", os.environ["CORTEX_LOG_CONFIG_FILE"])
 
@@ -42,9 +42,9 @@ def __init__(
         self,
         tf_serving_url,
         api_spec: dict,
-        models: Optional[ModelsHolder],
-        model_dir: Optional[str],
-        models_tree: Optional[ModelsTree],
+        models: Optional[ModelsHolder] = None,
+        model_dir: Optional[str] = None,
+        models_tree: Optional[ModelsTree] = None,
     ):
         """
         Setup gRPC connection to TensorFlow Serving container.
@@ -85,6 +85,34 @@ def __init__(
 
         self._client = TensorFlowServingAPI(tf_serving_url)
 
+    def sync_models(self, lock_dir: str = "/run/cron"):
+        resource_models = os.path.join(lock_dir, "models_tfs.json")
+
+        try:
+            with open(resource_models, "r") as f:
+                models = json.load(f)
+        except Exception:
+            return
+
+        resource_ts = os.path.join(lock_dir, "model_timestamps.json")
+        try:
+            with open(resource_ts, "r") as f:
+                timestamps = json.load(f)
+        except Exception:
+            return
+
+        non_intersecting_model_ids = set(models.keys()).symmetric_difference(timestamps.keys())
+        for non_intersecting_model_id in non_intersecting_model_ids:
+            if non_intersecting_model_id in models:
+                del models[non_intersecting_model_id]
+            if non_intersecting_model_id in timestamps:
+                del timestamps[non_intersecting_model_id]
+
+        for model_id in timestamps.keys():
+            models[model_id]["timestamp"] = timestamps[model_id]
+
+        self._client.models = models
+
     def predict(
         self, model_input: Any, model_name: Optional[str] = None, model_version: str = "latest"
     ) -> dict:
diff --git a/pkg/cortex/serve/cortex_internal/lib/model/cron.py b/pkg/cortex/serve/cortex_internal/lib/model/cron.py
index 1f393cd9ba..22cf8372ce 100644
--- a/pkg/cortex/serve/cortex_internal/lib/model/cron.py
+++ b/pkg/cortex/serve/cortex_internal/lib/model/cron.py
@@ -12,30 +12,25 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-import os
-import threading as td
-import multiprocessing as mp
-import time
+import copy
 import datetime
 import glob
-import shutil
 import json
-import grpc
-import copy
+import multiprocessing as mp
+import os
+import shutil
+import threading as td
+import time
 from concurrent.futures import ThreadPoolExecutor
 from typing import Dict, List, Tuple, Any, Union, Callable, Optional
 
+import grpc
+
 from cortex_internal.lib import util
+from cortex_internal.lib.client.tensorflow import TensorFlowClient
 from cortex_internal.lib.concurrency import LockedFile, get_locked_files
-from cortex_internal.lib.storage import S3, GCS
 from cortex_internal.lib.exceptions import CortexException, WithBreak
-from cortex_internal.lib.type import (
-    predictor_type_from_api_spec,
-    PythonPredictorType,
-    TensorFlowPredictorType,
-    TensorFlowNeuronPredictorType,
-)
-
+from cortex_internal.lib.log import configure_logger
 from cortex_internal.lib.model import (
     find_all_cloud_models,
     validate_model_paths,
@@ -48,8 +43,14 @@
     get_models_from_api_spec,
     ModelsTree,
 )
+from cortex_internal.lib.storage import S3, GCS
 from cortex_internal.lib.telemetry import get_default_tags, init_sentry
-from cortex_internal.lib.log import configure_logger
+from cortex_internal.lib.type import (
+    predictor_type_from_api_spec,
+    PythonPredictorType,
+    TensorFlowPredictorType,
+    TensorFlowNeuronPredictorType,
+)
 
 logger = configure_logger("cortex", os.environ["CORTEX_LOG_CONFIG_FILE"])
 
@@ -1236,8 +1237,8 @@ class TFSAPIServingThreadUpdater(AbstractLoopingThread):
 
     def __init__(
         self,
-        interval: int,
-        client: TensorFlowServingAPI,
+        interval: Union[int, float],
+        client: TensorFlowClient,
         lock_dir: str = "/run/cron",
     ):
         AbstractLoopingThread.__init__(self, interval, self._run_tfs)
@@ -1246,32 +1247,7 @@ def __init__(
         self._lock_dir = lock_dir
 
     def _run_tfs(self) -> None:
-        resource_models = os.path.join(self._lock_dir, "models_tfs.json")
-
-        try:
-            with open(resource_models, "r") as f:
-                models = json.load(f)
-        except Exception:
-            return
-
-        resource_ts = os.path.join(self._lock_dir, "model_timestamps.json")
-        try:
-            with open(resource_ts, "r") as f:
-                timestamps = json.load(f)
-        except Exception:
-            return
-
-        non_intersecting_model_ids = set(models.keys()).symmetric_difference(timestamps.keys())
-        for non_intersecting_model_id in non_intersecting_model_ids:
-            if non_intersecting_model_id in models:
-                del models[non_intersecting_model_id]
-            if non_intersecting_model_id in timestamps:
-                del timestamps[non_intersecting_model_id]
-
-        for model_id in timestamps.keys():
-            models[model_id]["timestamp"] = timestamps[model_id]
-
-        self._client.models = models
+        self._client.sync_models(self._lock_dir)
 
 
 def find_ondisk_models(models_dir: str) -> Dict[str, List[str]]:
diff --git a/pkg/cortex/serve/cortex_internal/lib/model/tree.py b/pkg/cortex/serve/cortex_internal/lib/model/tree.py
index ff267c40e1..c34003407d 100644
--- a/pkg/cortex/serve/cortex_internal/lib/model/tree.py
+++ b/pkg/cortex/serve/cortex_internal/lib/model/tree.py
@@ -12,27 +12,23 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-import os
-import itertools
-import time
 import datetime
-import shutil
+import itertools
+import os
 import threading as td
-from typing import List, Dict, Any, Tuple, Callable, AbstractSet
+from typing import List, Dict, Tuple, AbstractSet
 
 from cortex_internal.lib import util
-from cortex_internal.lib.type import PredictorType
 from cortex_internal.lib.concurrency import ReadWriteLock
 from cortex_internal.lib.exceptions import CortexException, WithBreak
-from cortex_internal.lib.storage import S3, GCS
-
+from cortex_internal.lib.log import configure_logger
 from cortex_internal.lib.model.validation import (
     validate_models_dir_paths,
     validate_model_paths,
     ModelVersion,
 )
-
-from cortex_internal.lib.log import configure_logger
+from cortex_internal.lib.storage import S3, GCS
+from cortex_internal.lib.type import PredictorType
 
 logger = configure_logger("cortex", os.environ["CORTEX_LOG_CONFIG_FILE"])
 
@@ -177,7 +173,7 @@ def update_model(
         sub_paths: List[str],
         timestamp: datetime.datetime,
         removable: bool,
-    ) -> None:
+    ) -> bool:
         """
         Updates the model tree with the given model.
 
diff --git a/pkg/cortex/serve/start/async.py b/pkg/cortex/serve/start/async.py
index ec73c233bc..5ee0d10a02 100644
--- a/pkg/cortex/serve/start/async.py
+++ b/pkg/cortex/serve/start/async.py
@@ -13,7 +13,9 @@
 # limitations under the License.
 import inspect
 import os
+import pathlib
 import sys
+import time
 from typing import Dict, Any
 
 import boto3
@@ -93,6 +95,10 @@ def build_predict_args(payload, request_id):
 
 
 def main():
+    while not pathlib.Path("/mnt/workspace/init_script_run.txt").is_file():
+        time.sleep(0.2)
+
+    model_dir = os.getenv("CORTEX_MODEL_DIR")
     cache_dir = os.environ["CORTEX_CACHE_DIR"]
     provider = os.environ["CORTEX_PROVIDER"]
     api_spec_path = os.environ["CORTEX_API_SPEC"]
@@ -103,6 +109,8 @@ def main():
     queue_url = os.environ["CORTEX_QUEUE_URL"]
     statsd_host = os.getenv("HOST_IP")
     statsd_port = os.getenv("CORTEX_STATSD_PORT", "9125")
+    tf_serving_host = os.getenv("CORTEX_TF_SERVING_HOST")
+    tf_serving_port = os.getenv("CORTEX_TF_BASE_SERVING_PORT")
 
     storage, api_spec = get_spec(provider, api_spec_path, cache_dir, region)
     sqs_client = boto3.client("sqs", region_name=region)
@@ -112,12 +120,18 @@ def main():
         storage_path=workload_path,
         statsd_host=statsd_host,
         statsd_port=int(statsd_port),
+        model_dir=model_dir,
     )
 
     try:
-        log.info("loading the predictor from {}".format(api.path))
+        log.info(f"loading the predictor from {api.path}")
         metrics_client = MetricsClient(api.statsd)
-        predictor_impl = api.initialize_impl(project_dir, metrics_client)
+        predictor_impl = api.initialize_impl(
+            project_dir,
+            metrics_client,
+            tf_serving_host=tf_serving_host,
+            tf_serving_port=tf_serving_port,
+        )
     except UserRuntimeException as err:
         err.wrap(f"failed to initialize predictor implementation")
         log.error(str(err), exc_info=True)
diff --git a/pkg/operator/operator/k8s.go b/pkg/operator/operator/k8s.go
index 40887a7c12..873258370a 100644
--- a/pkg/operator/operator/k8s.go
+++ b/pkg/operator/operator/k8s.go
@@ -203,7 +203,15 @@ func TaskContainers(api *spec.API) ([]kcore.Container, []kcore.Volume) {
 }
 
 func AsyncPythonPredictorContainers(api spec.API, queueURL string) ([]kcore.Container, []kcore.Volume) {
-	return pythonPredictorContainers(&api, getAsyncEnvVars(api, APIContainerName, queueURL))
+	return pythonPredictorContainers(&api, getAsyncAPIEnvVars(api, queueURL))
+}
+
+func AsyncTensorflowPredictorContainers(api spec.API, queueURL string) ([]kcore.Container, []kcore.Volume) {
+	return tensorFlowPredictorContainers(&api, getAsyncAPIEnvVars(api, queueURL))
+}
+
+func AsyncONNXPredictorContainers(api spec.API, queueURL string) ([]kcore.Container, []kcore.Volume) {
+	return onnxPredictorContainers(&api, getAsyncAPIEnvVars(api, queueURL))
 }
 
 func AsyncGatewayContainers(api spec.API, queueURL string) kcore.Container {
@@ -367,6 +375,10 @@ func pythonPredictorContainers(api *spec.API, envVars []kcore.EnvVar) ([]kcore.C
 }
 
 func TensorFlowPredictorContainers(api *spec.API) ([]kcore.Container, []kcore.Volume) {
+	return tensorFlowPredictorContainers(api, getEnvVars(api, APIContainerName))
+}
+
+func tensorFlowPredictorContainers(api *spec.API, envVars []kcore.EnvVar) ([]kcore.Container, []kcore.Volume) {
 	apiResourceList := kcore.ResourceList{}
 	tfServingResourceList := kcore.ResourceList{}
 	tfServingLimitsList := kcore.ResourceList{}
@@ -458,7 +470,7 @@ func TensorFlowPredictorContainers(api *spec.API) ([]kcore.Container, []kcore.Vo
 		Name:            APIContainerName,
 		Image:           api.Predictor.Image,
 		ImagePullPolicy: kcore.PullAlways,
-		Env:             getEnvVars(api, APIContainerName),
+		Env:             envVars,
 		EnvFrom:         baseEnvVars(),
 		VolumeMounts:    volumeMounts,
 		ReadinessProbe:  FileExistsProbe(_apiReadinessFile),
@@ -486,6 +498,10 @@ func TensorFlowPredictorContainers(api *spec.API) ([]kcore.Container, []kcore.Vo
 }
 
 func ONNXPredictorContainers(api *spec.API) ([]kcore.Container, []kcore.Volume) {
+	return onnxPredictorContainers(api, getEnvVars(api, APIContainerName))
+}
+
+func onnxPredictorContainers(api *spec.API, envVars []kcore.EnvVar) ([]kcore.Container, []kcore.Volume) {
 	resourceList := kcore.ResourceList{}
 	resourceLimitsList := kcore.ResourceList{}
 	apiPodVolumeMounts := defaultVolumeMounts()
@@ -533,7 +549,7 @@ func ONNXPredictorContainers(api *spec.API) ([]kcore.Container, []kcore.Volume)
 		Name:            APIContainerName,
 		Image:           api.Predictor.Image,
 		ImagePullPolicy: kcore.PullAlways,
-		Env:             getEnvVars(api, APIContainerName),
+		Env:             envVars,
 		EnvFrom:         baseEnvVars(),
 		VolumeMounts:    apiPodVolumeMounts,
 		ReadinessProbe:  FileExistsProbe(_apiReadinessFile),
@@ -637,16 +653,7 @@ func getTaskEnvVars(api *spec.API, container string) []kcore.EnvVar {
 	return envVars
 }
 
-func getAsyncEnvVars(api spec.API, container string, queueURL string) []kcore.EnvVar {
-	if container == _downloaderInitContainerName {
-		return []kcore.EnvVar{
-			{
-				Name:  "CORTEX_LOG_LEVEL",
-				Value: strings.ToUpper(api.Predictor.LogLevel.String()),
-			},
-		}
-	}
-
+func getAsyncAPIEnvVars(api spec.API, queueURL string) []kcore.EnvVar {
 	envVars := []kcore.EnvVar{
 		{
 			Name:  "CORTEX_KIND",
@@ -669,82 +676,102 @@ func getAsyncEnvVars(api spec.API, container string, queueURL string) []kcore.En
 		})
 	}
 
-	if container == APIContainerName {
-		envVars = append(envVars,
-			kcore.EnvVar{
-				Name:  "CORTEX_LOG_LEVEL",
-				Value: strings.ToUpper(api.Predictor.LogLevel.String()),
-			},
-			kcore.EnvVar{
-				Name: "HOST_IP",
-				ValueFrom: &kcore.EnvVarSource{
-					FieldRef: &kcore.ObjectFieldSelector{
-						FieldPath: "status.hostIP",
-					},
+	envVars = append(envVars,
+		kcore.EnvVar{
+			Name:  "CORTEX_LOG_LEVEL",
+			Value: strings.ToUpper(api.Predictor.LogLevel.String()),
+		},
+		kcore.EnvVar{
+			Name: "HOST_IP",
+			ValueFrom: &kcore.EnvVarSource{
+				FieldRef: &kcore.ObjectFieldSelector{
+					FieldPath: "status.hostIP",
 				},
 			},
+		},
+		kcore.EnvVar{
+			Name:  "CORTEX_SERVING_PORT",
+			Value: DefaultPortStr,
+		},
+		kcore.EnvVar{
+			Name:  "CORTEX_CACHE_DIR",
+			Value: _specCacheDir,
+		},
+		kcore.EnvVar{
+			Name:  "CORTEX_PROJECT_DIR",
+			Value: path.Join(_emptyDirMountPath, "project"),
+		},
+		kcore.EnvVar{
+			Name:  "CORTEX_DEPENDENCIES_PIP",
+			Value: api.Predictor.Dependencies.Pip,
+		},
+		kcore.EnvVar{
+			Name:  "CORTEX_DEPENDENCIES_CONDA",
+			Value: api.Predictor.Dependencies.Conda,
+		},
+		kcore.EnvVar{
+			Name:  "CORTEX_DEPENDENCIES_SHELL",
+			Value: api.Predictor.Dependencies.Shell,
+		},
+		kcore.EnvVar{
+			Name:  "CORTEX_QUEUE_URL",
+			Value: queueURL,
+		},
+		kcore.EnvVar{
+			Name:  "CORTEX_ASYNC_WORKLOAD_PATH",
+			Value: fmt.Sprintf("%s/apis/%s/workloads", config.ClusterName(), api.Name),
+		},
+		kcore.EnvVar{
+			Name:  "CORTEX_API_SPEC",
+			Value: config.BucketPath(api.PredictorKey),
+		},
+		kcore.EnvVar{
+			Name:  "CORTEX_PROCESSES_PER_REPLICA",
+			Value: s.Int32(api.Predictor.ProcessesPerReplica),
+		},
+	)
+
+	if api.Predictor.Type == userconfig.TensorFlowPredictorType {
+		envVars = append(envVars,
 			kcore.EnvVar{
-				Name:  "CORTEX_SERVING_PORT",
-				Value: DefaultPortStr,
-			},
-			kcore.EnvVar{
-				Name:  "CORTEX_CACHE_DIR",
-				Value: _specCacheDir,
-			},
-			kcore.EnvVar{
-				Name:  "CORTEX_PROJECT_DIR",
-				Value: path.Join(_emptyDirMountPath, "project"),
-			},
-			kcore.EnvVar{
-				Name:  "CORTEX_DEPENDENCIES_PIP",
-				Value: api.Predictor.Dependencies.Pip,
-			},
-			kcore.EnvVar{
-				Name:  "CORTEX_DEPENDENCIES_CONDA",
-				Value: api.Predictor.Dependencies.Conda,
-			},
-			kcore.EnvVar{
-				Name:  "CORTEX_DEPENDENCIES_SHELL",
-				Value: api.Predictor.Dependencies.Shell,
-			},
-			kcore.EnvVar{
-				Name:  "CORTEX_QUEUE_URL",
-				Value: queueURL,
+				Name:  "CORTEX_TF_BASE_SERVING_PORT",
+				Value: _tfBaseServingPortStr,
 			},
 			kcore.EnvVar{
-				Name:  "CORTEX_ASYNC_WORKLOAD_PATH",
-				Value: fmt.Sprintf("%s/apis/%s/workloads", config.ClusterName(), api.Name),
+				Name:  "CORTEX_TF_SERVING_HOST",
+				Value: _tfServingHost,
 			},
+		)
+	}
+
+	if api.Autoscaling != nil {
+		envVars = append(envVars,
 			kcore.EnvVar{
-				Name:  "CORTEX_API_SPEC",
-				Value: config.BucketPath(api.PredictorKey),
+				Name:  "CORTEX_MAX_REPLICA_CONCURRENCY",
+				Value: s.Int64(api.Autoscaling.MaxReplicaConcurrency),
 			},
+		)
+	}
+
+	if api.Predictor.Type != userconfig.PythonPredictorType {
+		envVars = append(envVars,
 			kcore.EnvVar{
-				Name:  "CORTEX_PROCESSES_PER_REPLICA",
-				Value: s.Int32(api.Predictor.ProcessesPerReplica),
+				Name:  "CORTEX_MODEL_DIR",
+				Value: path.Join(_emptyDirMountPath, "model"),
 			},
 		)
+	}
 
-		if api.Autoscaling != nil {
-			envVars = append(envVars,
-				kcore.EnvVar{
-					Name:  "CORTEX_MAX_REPLICA_CONCURRENCY",
-					Value: s.Int64(api.Autoscaling.MaxReplicaConcurrency),
-				},
-			)
-		}
-
-		cortexPythonPath := path.Join(_emptyDirMountPath, "project")
-		if api.Predictor.PythonPath != nil {
-			cortexPythonPath = path.Join(_emptyDirMountPath, "project", *api.Predictor.PythonPath)
-		}
-		envVars = append(envVars, kcore.EnvVar{
-			Name:  "CORTEX_PYTHON_PATH",
-			Value: cortexPythonPath,
-		})
+	cortexPythonPath := path.Join(_emptyDirMountPath, "project")
+	if api.Predictor.PythonPath != nil {
+		cortexPythonPath = path.Join(_emptyDirMountPath, "project", *api.Predictor.PythonPath)
 	}
+	envVars = append(envVars, kcore.EnvVar{
+		Name:  "CORTEX_PYTHON_PATH",
+		Value: cortexPythonPath,
+	})
 
-	if api.Compute.Inf > 0 && container == APIContainerName {
+	if api.Compute.Inf > 0 {
 		envVars = append(envVars,
 			kcore.EnvVar{
 				Name:  "NEURONCORE_GROUP_SIZES",
@@ -755,6 +782,19 @@ func getAsyncEnvVars(api spec.API, container string, queueURL string) []kcore.En
 				Value: fmt.Sprintf("unix:%s", _neuronRTDSocket),
 			},
 		)
+
+		if api.Predictor.Type == userconfig.TensorFlowPredictorType {
+			envVars = append(envVars,
+				kcore.EnvVar{
+					Name:  "CORTEX_MULTIPLE_TF_SERVERS",
+					Value: "yes",
+				},
+				kcore.EnvVar{
+					Name:  "CORTEX_ACTIVE_NEURON",
+					Value: "yes",
+				},
+			)
+		}
 	}
 
 	return envVars
diff --git a/pkg/operator/resources/asyncapi/k8s_specs.go b/pkg/operator/resources/asyncapi/k8s_specs.go
index 73d7252fbd..cfa69ac274 100644
--- a/pkg/operator/resources/asyncapi/k8s_specs.go
+++ b/pkg/operator/resources/asyncapi/k8s_specs.go
@@ -17,10 +17,13 @@ limitations under the License.
 package asyncapi
 
 import (
+	"fmt"
+
 	"github.com/cortexlabs/cortex/pkg/lib/k8s"
 	"github.com/cortexlabs/cortex/pkg/lib/pointer"
 	"github.com/cortexlabs/cortex/pkg/operator/operator"
 	"github.com/cortexlabs/cortex/pkg/types/spec"
+	"github.com/cortexlabs/cortex/pkg/types/userconfig"
 	"istio.io/client-go/pkg/apis/networking/v1beta1"
 	kapps "k8s.io/api/apps/v1"
 	kcore "k8s.io/api/core/v1"
@@ -121,7 +124,21 @@ func gatewayVirtualServiceSpec(api spec.API) v1beta1.VirtualService {
 }
 
 func apiDeploymentSpec(api spec.API, prevDeployment *kapps.Deployment, queueURL string) kapps.Deployment {
-	containers, volumes := operator.AsyncPythonPredictorContainers(api, queueURL)
+	var (
+		containers []kcore.Container
+		volumes    []kcore.Volume
+	)
+
+	switch api.Predictor.Type {
+	case userconfig.PythonPredictorType:
+		containers, volumes = operator.AsyncPythonPredictorContainers(api, queueURL)
+	case userconfig.TensorFlowPredictorType:
+		containers, volumes = operator.AsyncTensorflowPredictorContainers(api, queueURL)
+	case userconfig.ONNXPredictorType:
+		containers, volumes = operator.AsyncONNXPredictorContainers(api, queueURL)
+	default:
+		panic(fmt.Sprintf("invalid predictor type: %s", api.Predictor.Type))
+	}
 
 	return *k8s.Deployment(&k8s.DeploymentSpec{
 		Name:           operator.K8sName(api.Name),
diff --git a/pkg/types/spec/validations.go b/pkg/types/spec/validations.go
index 0aa4860a74..8da1f835f0 100644
--- a/pkg/types/spec/validations.go
+++ b/pkg/types/spec/validations.go
@@ -883,10 +883,6 @@ func validatePredictor(
 	k8sClient *k8s.Client,
 ) error {
 	predictor := api.Predictor
-	if api.Kind == userconfig.AsyncAPIKind && predictor.Type != userconfig.PythonPredictorType {
-		return ErrorPredictorTypeNotSupportedForKind(predictor.Type, api.Kind)
-	}
-
 	if err := validateMultiModelsFields(api); err != nil {
 		return err
 	}


Test patch:
diff --git a/test/apis/async/onnx/cortex.yaml b/test/apis/async/onnx/cortex.yaml
new file mode 100644
index 0000000000..a675c299ce
--- /dev/null
+++ b/test/apis/async/onnx/cortex.yaml
@@ -0,0 +1,7 @@
+- name: async-iris-classifier
+  kind: AsyncAPI
+  predictor:
+    type: onnx
+    path: predictor.py
+    models:
+      path: s3://cortex-examples/onnx/iris-classifier/
diff --git a/test/apis/async/onnx/predictor.py b/test/apis/async/onnx/predictor.py
new file mode 100644
index 0000000000..a43d2defbf
--- /dev/null
+++ b/test/apis/async/onnx/predictor.py
@@ -0,0 +1,18 @@
+labels = ["setosa", "versicolor", "virginica"]
+
+
+class ONNXPredictor:
+    def __init__(self, onnx_client, config):
+        self.client = onnx_client
+
+    def predict(self, payload):
+        model_input = [
+            payload["sepal_length"],
+            payload["sepal_width"],
+            payload["petal_length"],
+            payload["petal_width"],
+        ]
+
+        prediction = self.client.predict(model_input)
+        predicted_class_id = prediction[0][0]
+        return {"label": labels[predicted_class_id]}
diff --git a/test/apis/async/onnx/sample.json b/test/apis/async/onnx/sample.json
new file mode 100644
index 0000000000..252c666b3a
--- /dev/null
+++ b/test/apis/async/onnx/sample.json
@@ -0,0 +1,6 @@
+{
+  "sepal_length": 5.2,
+  "sepal_width": 3.6,
+  "petal_length": 1.4,
+  "petal_width": 0.3
+}
diff --git a/test/apis/async/tensorflow/cortex.yaml b/test/apis/async/tensorflow/cortex.yaml
new file mode 100644
index 0000000000..fb5d625b48
--- /dev/null
+++ b/test/apis/async/tensorflow/cortex.yaml
@@ -0,0 +1,7 @@
+- name: async-iris-classifier
+  kind: AsyncAPI
+  predictor:
+    type: tensorflow
+    path: predictor.py
+    models:
+      path: s3://cortex-examples/tensorflow/iris-classifier/nn/
diff --git a/test/apis/async/tensorflow/expectations.yaml b/test/apis/async/tensorflow/expectations.yaml
new file mode 100644
index 0000000000..bef65e0e0d
--- /dev/null
+++ b/test/apis/async/tensorflow/expectations.yaml
@@ -0,0 +1,10 @@
+response:
+  content_type: "json"
+  json_schema:
+    type: "object"
+    properties:
+      label:
+        type: "string"
+        const: "setosa"
+    required:
+      - "label"
diff --git a/test/apis/async/tensorflow/predictor.py b/test/apis/async/tensorflow/predictor.py
new file mode 100644
index 0000000000..239767d3c1
--- /dev/null
+++ b/test/apis/async/tensorflow/predictor.py
@@ -0,0 +1,11 @@
+labels = ["setosa", "versicolor", "virginica"]
+
+
+class TensorFlowPredictor:
+    def __init__(self, tensorflow_client, config):
+        self.client = tensorflow_client
+
+    def predict(self, payload):
+        prediction = self.client.predict(payload)
+        predicted_class_id = int(prediction["class_ids"][0])
+        return {"label": labels[predicted_class_id]}
diff --git a/test/apis/async/tensorflow/sample.json b/test/apis/async/tensorflow/sample.json
new file mode 100644
index 0000000000..252c666b3a
--- /dev/null
+++ b/test/apis/async/tensorflow/sample.json
@@ -0,0 +1,6 @@
+{
+  "sepal_length": 5.2,
+  "sepal_width": 3.6,
+  "petal_length": 1.4,
+  "petal_width": 0.3
+}
diff --git a/test/e2e/tests/aws/test_async.py b/test/e2e/tests/aws/test_async.py
index a3a1d1a6b5..d5eb9215c3 100644
--- a/test/e2e/tests/aws/test_async.py
+++ b/test/e2e/tests/aws/test_async.py
@@ -21,6 +21,8 @@
 
 TEST_APIS = [
     {"name": "async/iris-classifier", "poll_retries": 10, "poll_sleep_time": 1},
+    {"name": "async/tensorflow", "poll_retries": 10, "poll_sleep_time": 1},
+    {"name": "async/onnx", "poll_retries": 10, "poll_sleep_time": 1},
 ]