[RAPTOR-3979][RAPTOR-4016][RAPTOR-4017] Transform endpoint using arrow/mtx (#236)

* do the things

* add transformer

* fix some things

* refactor

* fix test

* mixin comments

* add uwsgi endpoint

* run black

* add nginx test

* update key name, use enum values

* maybe fix a test

* fix enum thing

* do a bunch of things

* do a bunch of things?

* add arrow support

* do a bunch more stuff

* remove a thing

* black

* get rid of duplicate thing

* try to fix some tests

* black

* change key name

* run black

* move some things, address some comments

* unused import

* add sparse indicator to payload

* add csv fallback

* add more things to tests

* run black

* missed a digit

* update other stuff

* deal with arrow versions maybe

Author: jlpolit

custom_model_runner/datarobot_drum/drum/common.py

@@ -15,6 +15,7 @@
 NEGATIVE_CLASS_LABEL_ARG_KEYWORD = "negative_class_label"
 CLASS_LABELS_ARG_KEYWORD = "class_labels"
 TARGET_TYPE_ARG_KEYWORD = "target_type"
+X_TRANSFORM_KEY = "X.transformed"
 
 URL_PREFIX_ENV_VAR_NAME = "URL_PREFIX"
 

custom_model_runner/datarobot_drum/resource/components/Python/prediction_server/prediction_server.py

@@ -5,7 +5,7 @@
 from datarobot_drum.drum.exceptions import DrumCommonException
 from datarobot_drum.profiler.stats_collector import StatsCollector, StatsOperation
 from datarobot_drum.drum.memory_monitor import MemoryMonitor
-from datarobot_drum.drum.common import RunLanguage, TARGET_TYPE_ARG_KEYWORD
+from datarobot_drum.drum.common import RunLanguage, TARGET_TYPE_ARG_KEYWORD, TargetType
 from datarobot_drum.resource.predict_mixin import PredictMixin
 
 from datarobot_drum.drum.server import (
@@ -32,7 +32,7 @@ def configure(self, params):
         super(PredictionServer, self).configure(params)
         self._show_perf = self._params.get("show_perf")
         self._run_language = RunLanguage(params.get("run_language"))
-        self._target_type = params[TARGET_TYPE_ARG_KEYWORD]
+        self._target_type = TargetType(params[TARGET_TYPE_ARG_KEYWORD])
 
         self._stats_collector = StatsCollector(disable_instance=not self._show_perf)
 
@@ -91,6 +91,21 @@ def predict():
                 self._stats_collector.disable()
             return response, response_status
 
+        @model_api.route("/transform/", methods=["POST"])
+        def transform():
+
+            logger.debug("Entering transform() endpoint")
+
+            self._stats_collector.enable()
+            self._stats_collector.mark("start")
+
+            try:
+                response, response_status = self.do_transform(logger=logger)
+            finally:
+                self._stats_collector.mark("finish")
+                self._stats_collector.disable()
+            return response, response_status
+
         @model_api.route("/predictUnstructured/", methods=["POST"])
         def predict_unstructured():
             logger.debug("Entering predict() endpoint")

custom_model_runner/datarobot_drum/resource/components/Python/uwsgi_component/uwsgi_serving.py

@@ -11,6 +11,7 @@
     URL_PREFIX_ENV_VAR_NAME,
     TARGET_TYPE_ARG_KEYWORD,
     make_predictor_capabilities,
+    TargetType,
 )
 from datarobot_drum.profiler.stats_collector import StatsCollector, StatsOperation
 
@@ -61,7 +62,7 @@ def configure(self, params):
         super(UwsgiServing, self).configure(params)
         self._show_perf = self._params.get("show_perf")
         self._run_language = RunLanguage(params.get("run_language"))
-        self._target_type = params[TARGET_TYPE_ARG_KEYWORD]
+        self._target_type = TargetType(params[TARGET_TYPE_ARG_KEYWORD])
 
         self._stats_collector = StatsCollector(disable_instance=not self._show_perf)
 
@@ -179,6 +180,30 @@ def predict_unstructured(self, url_params, form_params):
             self._stats_collector.disable()
         return response_status, response
 
+    @FlaskRoute(
+        "{}/transform/".format(os.environ.get(URL_PREFIX_ENV_VAR_NAME, "")), methods=["POST"]
+    )
+    def transform(self, url_params, form_params):
+        if self._error_response:
+            return HTTP_513_DRUM_PIPELINE_ERROR, self._error_response
+
+        self._stats_collector.enable()
+        self._stats_collector.mark("start")
+
+        try:
+            response, response_status = self.do_transform()
+
+            if response_status == HTTP_200_OK:
+                # this counter is managed by uwsgi
+                self._total_predict_requests.increase()
+                self._predict_calls_count += 1
+        except Exception as ex:
+            response_status, response = self._handle_exception(ex)
+        finally:
+            self._stats_collector.mark("finish")
+            self._stats_collector.disable()
+        return response_status, response
+
     def _handle_exception(self, ex):
         self._logger.error(ex)
         response_status = HTTP_500_INTERNAL_SERVER_ERROR

custom_model_runner/datarobot_drum/resource/predict_mixin.py

@@ -9,6 +9,13 @@
     TargetType,
     UnstructuredDtoKeys,
     PredictionServerMimetypes,
+    X_TRANSFORM_KEY,
+)
+from datarobot_drum.resource.transform_helpers import (
+    make_arrow_payload,
+    is_sparse,
+    make_mtx_payload,
+    make_csv_payload,
 )
 from datarobot_drum.resource.unstructured_helpers import (
     _resolve_incoming_unstructured_data,
@@ -29,12 +36,19 @@ class PredictMixin:
 
     """
 
-    def do_predict(self, logger=None):
+    def _predict_or_transform(self, logger=None):
         response_status = HTTP_200_OK
 
         file_key = "X"
         filestorage = request.files.get(file_key)
 
+        if self._target_type == TargetType.TRANSFORM:
+            arrow_key = "arrow_version"
+            arrow_version = request.files.get(arrow_key)
+            if arrow_version is not None:
+                arrow_version = eval(arrow_version.getvalue())
+            use_arrow = arrow_version is not None
+
         if not filestorage:
             wrong_key_error_message = (
                 "Samples should be provided as a csv, mtx, or arrow file under `{}` key.".format(
@@ -53,10 +67,38 @@ def do_predict(self, logger=None):
         with tempfile.NamedTemporaryFile(suffix=file_ext) as f:
             filestorage.save(f)
             f.flush()
-            out_data = self._predictor.predict(f.name)
+            if self._target_type == TargetType.TRANSFORM:
+                out_data = self._predictor.transform(f.name)
+            else:
+                out_data = self._predictor.predict(f.name)
 
         if self._target_type == TargetType.UNSTRUCTURED:
             response = out_data
+        elif self._target_type == TargetType.TRANSFORM:
+            if is_sparse(out_data):
+                mtx_payload = make_mtx_payload(out_data)
+                response = (
+                    '{{"{transform_key}":{mtx_payload}, "out.format":"{out_format}"}}'.format(
+                        transform_key=X_TRANSFORM_KEY, mtx_payload=mtx_payload, out_format="sparse"
+                    )
+                )
+            else:
+                if use_arrow:
+                    arrow_payload = make_arrow_payload(out_data, arrow_version)
+                    response = (
+                        '{{"{transform_key}":{arrow_payload}, "out.format":"{out_format}"}}'.format(
+                            transform_key=X_TRANSFORM_KEY,
+                            arrow_payload=arrow_payload,
+                            out_format="arrow",
+                        )
+                    )
+                else:
+                    csv_payload = make_csv_payload(out_data)
+                    response = (
+                        '{{"{transform_key}":{csv_payload}, "out.format":"{out_format}"}}'.format(
+                            transform_key=X_TRANSFORM_KEY, csv_payload=csv_payload, out_format="csv"
+                        )
+                    )
         else:
             num_columns = len(out_data.columns)
             # float32 is not JSON serializable, so cast to float, which is float64
@@ -76,6 +118,19 @@ def do_predict(self, logger=None):
 
         return response, response_status
 
+    def do_predict(self, logger=None):
+        if self._target_type == TargetType.TRANSFORM:
+            wrong_target_type_error_message = (
+                "This project has target type {}, "
+                "use the /transform/ endpoint.".format(self._target_type)
+            )
+            if logger is not None:
+                logger.error(wrong_target_type_error_message)
+            response_status = HTTP_422_UNPROCESSABLE_ENTITY
+            return {"message": "ERROR: " + wrong_target_type_error_message}, response_status
+
+        return self._predict_or_transform(logger=logger)
+
     def do_predict_unstructured(self, logger=None):
         def _validate_content_type_header(header):
             ret_mimetype, content_type_params_dict = werkzeug.http.parse_options_header(header)
@@ -115,3 +170,19 @@ def _validate_content_type_header(header):
             response.headers["Content-Type"] = content_type
 
         return response, response_status
+
+    def do_transform(self, logger=None):
+        if self._target_type != TargetType.TRANSFORM:
+            endpoint = (
+                "predictUnstructured" if self._target_type == TargetType.UNSTRUCTURED else "predict"
+            )
+            wrong_target_type_error_message = (
+                "This project has target type {}, "
+                "use the /{}/ endpoint.".format(self._target_type, endpoint)
+            )
+            if logger is not None:
+                logger.error(wrong_target_type_error_message)
+            response_status = HTTP_422_UNPROCESSABLE_ENTITY
+            return {"message": "ERROR: " + wrong_target_type_error_message}, response_status
+
+        return self._predict_or_transform(logger=logger)

custom_model_runner/datarobot_drum/resource/transform_helpers.py

@@ -0,0 +1,70 @@
+import pyarrow as pa
+import pandas as pd
+
+from io import BytesIO, StringIO
+
+from scipy.io import mmwrite, mmread
+from scipy.sparse.csr import csr_matrix
+from scipy.sparse import vstack
+
+from datarobot_drum.drum.common import X_TRANSFORM_KEY
+
+
+def is_sparse(df):
+    return hasattr(df, "sparse") or type(df.iloc[0].values[0]) == csr_matrix
+
+
+def make_arrow_payload(df, arrow_version):
+    if arrow_version != pa.__version__ and arrow_version < 0.2:
+        batch = pa.RecordBatch.from_pandas(df, nthreads=None, preserve_index=False)
+        sink = pa.BufferOutputStream()
+        options = pa.ipc.IpcWriteOptions(
+            metadata_version=pa.MetadataVersion.V4, use_legacy_format=True
+        )
+        with pa.RecordBatchStreamWriter(sink, batch.schema, options=options) as writer:
+            writer.write_batch(batch)
+        return sink.getvalue().to_pybytes()
+    else:
+        return pa.ipc.serialize_pandas(df, preserve_index=False).to_pybytes()
+
+
+def make_csv_payload(df):
+    s_buf = StringIO()
+    df.to_csv(s_buf, index=False)
+    return s_buf.getvalue().encode("utf-8")
+
+
+def read_arrow_payload(response_dict):
+    bytes = response_dict[X_TRANSFORM_KEY]
+    df = pa.ipc.deserialize_pandas(bytes)
+    return df
+
+
+def read_csv_payload(response_dict):
+    bytes = response_dict[X_TRANSFORM_KEY]
+    return pd.read_csv(BytesIO(bytes))
+
+
+def make_mtx_payload(df):
+    if hasattr(df, "sparse"):
+        sparse_mat = csr_matrix(df.sparse.to_coo())
+    else:
+        sparse_mat = vstack(x[0] for x in df.values)
+    sink = BytesIO()
+    mmwrite(sink, sparse_mat)
+    return sink.getvalue()
+
+
+def read_mtx_payload(response_dict):
+    bytes = response_dict[X_TRANSFORM_KEY]
+    sparse_mat = mmread(BytesIO(bytes))
+    return csr_matrix(sparse_mat)
+
+
+def validate_transformed_output(transformed_output, should_be_sparse=False):
+    if should_be_sparse:
+        assert type(transformed_output) == csr_matrix
+        assert transformed_output.shape[1] == 714
+    else:
+        assert type(transformed_output) == pd.DataFrame
+        assert transformed_output.shape[1] == 10

model_templates/training/python3_sklearn_transform/README.md

@@ -30,11 +30,7 @@ Categoricals:
 - Impute missing values with the string "missing"
 - One hot encode the data (ignoring new categorical levels at prediction time)
 
-SVD:
-After all the above is done, run SVD to reduce the dimensionality of the dataset to 10.
-
-This makes a dataset that can be used with basically any sklearn model.  This step could be removed for models that support sparse data.
-
+This makes a dataset that can be used with any sklearn model that supports sparse data.
 
 ### To run locally using 'drum'
 Paths are relative to `datarobot-user-models` root:

model_templates/training/python3_sklearn_transform/create_transform_pipeline.py

@@ -2,7 +2,6 @@
 import pandas as pd
 from sagemaker_sklearn_extension.feature_extraction.text import MultiColumnTfidfVectorizer
 from sklearn.compose import ColumnTransformer, make_column_selector
-from sklearn.decomposition import TruncatedSVD
 from sklearn.impute import SimpleImputer
 from sklearn.pipeline import Pipeline
 from sklearn.preprocessing import OneHotEncoder, StandardScaler
@@ -34,24 +33,6 @@
     ]
 )
 
-# Modified TruncatedSVD that doesn't fail if n_components > ncols
-class MyTruncatedSVD(TruncatedSVD):
-    def fit_transform(self, X, y=None):
-        if X.shape[1] <= self.n_components:
-            self.n_components = X.shape[1] - 1
-        return TruncatedSVD.fit_transform(self, X=X, y=y)
-
-
-# Dense preprocessing pipeline, for models such as XGboost that do not do well with
-# extremely wide, sparse data
-# This preprocessing will work with linear models such as Ridge too
-dense_preprocessing_pipeline = Pipeline(
-    steps=[
-        ("preprocessing", sparse_preprocessing_pipeline),
-        ("SVD", MyTruncatedSVD(n_components=10, random_state=42, algorithm="randomized")),
-    ]
-)
-
 
 def make_pipeline():
-    return dense_preprocessing_pipeline
+    return sparse_preprocessing_pipeline

model_templates/training/python3_sklearn_transform/custom.py

@@ -1,5 +1,6 @@
 import pickle
 import pandas as pd
+from scipy.sparse.csr import csr_matrix
 
 from create_transform_pipeline import make_pipeline
 
@@ -54,4 +55,8 @@ def transform(data, transformer):
     -------
     transformed DataFrame resulting from applying transform to incoming data
     """
-    return pd.DataFrame(transformer.transform(data))
+    transformed = transformer.transform(data)
+    if type(transform) == csr_matrix:
+        return pd.DataFrame.sparse.from_spmatrix(transformed)
+    else:
+        return pd.DataFrame(transformed)