Improved performance of mars.learn.metrics.{roc_curve, roc_auc_score} (mars-project#2838)

Author: Xuye (Chris) Qin

.github/workflows/benchmark-ci.yml

@@ -49,7 +49,7 @@ jobs:
           git remote add upstream https://github.com/mars-project/mars.git
           git fetch upstream
           asv machine --yes
-          asv continuous -f 1.1 upstream/master HEAD
+          asv continuous -f 1.1 --strict upstream/master HEAD
         if: ${{ steps.build.outcome == 'success' }}
 
       - name: Publish benchmarks artifact

asv_bench/benchmarks/learn/__init__.py

@@ -0,0 +1,13 @@
+# Copyright 1999-2022 Alibaba Group Holding Ltd.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#      http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# 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.

asv_bench/benchmarks/learn/metrics.py

@@ -0,0 +1,46 @@
+# Copyright 1999-2022 Alibaba Group Holding Ltd.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#      http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# 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 mars
+import mars.tensor as mt
+from sklearn.datasets import make_classification
+from sklearn.linear_model import LogisticRegression
+from mars.learn import metrics
+
+
+class MetricsSuite:
+    """
+    Benchmark learn metrics.
+    """
+
+    params = [20_000, 100_000]
+
+    def setup(self, chunk_size: int):
+        X, y = make_classification(100_000, random_state=0)
+        self.raw_X, self.raw_y = X, y
+        clf = LogisticRegression(solver="liblinear", random_state=0).fit(X, y)
+        self.raw_pred_y = clf.predict_proba(X)[:, 1]
+        self._session = mars.new_session()
+        self.y = mt.tensor(self.raw_y, chunk_size=chunk_size)
+        self.pred_y = mt.tensor(self.raw_pred_y, chunk_size=chunk_size)
+
+    def teardown(self, chunk_size: int):
+        self._session.stop_server()
+
+    def time_roc_curve_auc(self, chunk_size: int):
+        fpr, tpr, _ = metrics.roc_curve(self.y, self.pred_y)
+        metrics.auc(fpr, tpr)
+
+    def time_roc_auc_score(self, chunk_size: int):
+        metrics.roc_auc_score(self.y, self.pred_y)

mars/config.py

@@ -385,159 +385,19 @@ def validate(x):
 # deploy
 default_options.register_option("deploy.open_browser", True, validator=is_bool)
 
-# Scheduler
-default_options.register_option(
-    "scheduler.assign_chunk_workers", False, validator=is_bool, serialize=True
-)
-default_options.register_option(
-    "scheduler.enable_active_push", True, validator=is_bool, serialize=True
-)
-default_options.register_option(
-    "scheduler.enable_chunk_relocation", False, validator=is_bool, serialize=True
-)
-default_options.register_option(
-    "scheduler.check_interval", 1, validator=is_integer, serialize=True
-)
-default_options.register_option(
-    "scheduler.default_cpu_usage", 1, validator=(is_integer, is_float), serialize=True
-)
-default_options.register_option(
-    "scheduler.default_cuda_usage",
-    0.5,
-    validator=(is_integer, is_float),
-    serialize=True,
-)
-default_options.register_option(
-    "scheduler.assign_timeout", 600, validator=is_integer, serialize=True
-)
-default_options.register_option(
-    "scheduler.execution_timeout", 600, validator=is_integer, serialize=True
-)
-default_options.register_option(
-    "scheduler.retry_num", 4, validator=is_integer, serialize=True
-)
-default_options.register_option(
-    "scheduler.fetch_limit", 10 * 1024**2, validator=is_integer, serialize=True
-)
-default_options.register_option(
-    "scheduler.retry_delay", 60, validator=is_integer, serialize=True
-)
-
-default_options.register_option("scheduler.dump_graph_data", False, validator=is_bool)
-
-default_options.register_option(
-    "scheduler.enable_failover", True, validator=is_bool, serialize=True
-)
-default_options.register_option(
-    "scheduler.status_timeout", 30, validator=is_numeric, serialize=True
-)
-default_options.register_option(
-    "scheduler.worker_blacklist_time", 3600, validator=is_numeric, serialize=True
-)
-
-# enqueue operands in a batch when creating OperandActors
-default_options.register_option(
-    "scheduler.batch_enqueue_initials", True, validator=is_bool, serialize=True
-)
-# invoke assigning when where there is no ready descendants
-default_options.register_option(
-    "scheduler.aggressive_assign", False, validator=is_bool, serialize=True
-)
-
-# Worker
-default_options.register_option(
-    "worker.spill_directory", None, validator=(is_null, is_string, is_list)
-)
-default_options.register_option(
-    "worker.disk_compression", "lz4", validator=is_string, serialize=True
-)
-default_options.register_option(
-    "worker.min_spill_size", "5%", validator=(is_string, is_integer)
-)
-default_options.register_option(
-    "worker.max_spill_size", "95%", validator=(is_string, is_integer)
-)
-default_options.register_option(
-    "worker.min_cache_mem_size", None, validator=(is_null, is_string, is_integer)
-)
-default_options.register_option(
-    "worker.callback_preserve_time", 3600 * 24, validator=is_integer
-)
-default_options.register_option(
-    "worker.event_preserve_time", 3600 * 24, validator=(is_integer, is_float)
-)
-default_options.register_option(
-    "worker.copy_block_size", 64 * 1024, validator=is_integer
-)
-default_options.register_option("worker.cuda_thread_num", 2, validator=is_integer)
-default_options.register_option(
-    "worker.transfer_block_size", 1 * 1024**2, validator=is_integer
-)
-default_options.register_option(
-    "worker.transfer_compression", "lz4", validator=is_string, serialize=True
-)
-default_options.register_option(
-    "worker.prepare_data_timeout", 1000, validator=is_integer
-)
-default_options.register_option(
-    "worker.peer_blacklist_time", 3600, validator=is_numeric, serialize=True
-)
-default_options.register_option(
-    "worker.io_parallel_num", 1, validator=is_integer, serialize=True
-)
-default_options.register_option(
-    "worker.recover_dead_process", True, validator=is_bool, serialize=True
-)
-default_options.register_option(
-    "worker.write_shuffle_to_disk", False, validator=is_bool, serialize=True
-)
-
-default_options.register_option(
-    "worker.filemerger.enabled", True, validator=is_bool, serialize=True
-)
-default_options.register_option(
-    "worker.filemerger.concurrency", 128, validator=is_integer, serialize=True
-)
-default_options.register_option(
-    "worker.filemerger.max_accept_size",
-    128 * 1024,
-    validator=is_integer,
-    serialize=True,
-)
-default_options.register_option(
-    "worker.filemerger.max_file_size",
-    32 * 1024**2,
-    validator=is_integer,
-    serialize=True,
-)
-
-default_options.register_option(
-    "worker.plasma_dir", None, validator=(is_string, is_null)
-)
-default_options.register_option(
-    "worker.plasma_limit", None, validator=(is_string, is_integer, is_null)
-)
+# optimization
+default_options.register_option("optimize_tileable_graph", True, validator=is_bool)
 
-default_options.register_option(
-    "worker.plasma_socket", "/tmp/plasma", validator=is_string
-)
+# eager mode
+default_options.register_option("eager_mode", False, validator=is_bool)
 
 # optimization
-default_options.register_option("optimize.min_stats_count", 10, validator=is_integer)
-default_options.register_option(
-    "optimize.stats_sufficient_ratio", 0.9, validator=is_float, serialize=True
-)
-default_options.register_option(
-    "optimize.default_disk_io_speed", 10 * 1024**2, validator=is_integer
-)
 default_options.register_option(
     "optimize.head_optimize_threshold", 1000, validator=is_integer
 )
 
-default_options.register_option("optimize_tileable_graph", True, validator=is_bool)
-
-# eager mode
-default_options.register_option("eager_mode", False, validator=is_bool)
+# debug
+default_options.register_option("warn_duplicated_execution", False, validator=is_bool)
 
 # client serialize type
 default_options.register_option("client.serial_type", "arrow", validator=is_string)

mars/core/entity/tileables.py

@@ -21,7 +21,7 @@
 
 import numpy as np
 
-from ...serialization.serializables import FieldTypes, TupleField
+from ...serialization.serializables import FieldTypes, BoolField, TupleField
 from ...typing import OperandType, TileableType, ChunkType
 from ...utils import on_serialize_shape, on_deserialize_shape, on_serialize_nsplits
 from ..base import Base
@@ -273,6 +273,8 @@ class TileableData(EntityData, _ExecutableMixin):
         FieldTypes.tuple(FieldTypes.uint64),
         on_serialize=on_serialize_nsplits,
     )
+    # cache tileable data, if true, this data will be materialized
+    cache = BoolField("cache", default=False)
 
     def __init__(self: TileableType, *args, **kwargs):
         if kwargs.get("_nsplits", None) is not None:

mars/dataframe/base/diff.py

@@ -26,7 +26,7 @@
 
 
 class DataFrameDiff(DataFrameOperandMixin, DataFrameOperand):
-    _op_type_ = opcodes.DATAFRAME_DIFF
+    _op_type_ = opcodes.DIFF
 
     _periods = Int64Field("periods")
     _axis = Int8Field("axis")