[Object spilling] Add policy to automatically spill objects on OutOfMemory (#11673)

python/ray/_private/services.py

@@ -1,3 +1,4 @@
+import base64
 import collections
 import errno
 import io
@@ -69,6 +70,21 @@ ProcessInfo = collections.namedtuple("ProcessInfo", [
 ])
 
 
+def serialize_config(config):
+    config_pairs = []
+    for key, value in config.items():
+        if isinstance(value, str):
+            value = value.encode("utf-8")
+        if isinstance(value, bytes):
+            value = base64.b64encode(value).decode("utf-8")
+        config_pairs.append((key, value))
+    config_str = ";".join(["{},{}".format(*kv) for kv in config_pairs])
+    assert " " not in config_str, (
+        "Config parameters currently do not support "
+        "spaces:", config_str)
+    return config_str
+
+
 class ConsolePopen(subprocess.Popen):
     if sys.platform == "win32":
 
@@ -1121,7 +1137,7 @@ def start_gcs_server(redis_address,
     """
     gcs_ip_address, gcs_port = redis_address.split(":")
     redis_password = redis_password or ""
-    config_str = ",".join(["{},{}".format(*kv) for kv in config.items()])
+    config_str = serialize_config(config)
     if gcs_server_port is None:
         gcs_server_port = 0
 
@@ -1176,7 +1192,6 @@ def start_raylet(redis_address,
                  socket_to_use=None,
                  head_node=False,
                  start_initial_python_workers_for_first_job=False,
-                 object_spilling_config=None,
                  code_search_path=None):
     """Start a raylet, which is a combined local scheduler and object manager.
 
@@ -1223,7 +1238,7 @@ def start_raylet(redis_address,
     # The caller must provide a node manager port so that we can correctly
     # populate the command to start a worker.
     assert node_manager_port is not None and node_manager_port != 0
-    config_str = ",".join(["{},{}".format(*kv) for kv in config.items()])
+    config_str = serialize_config(config)
 
     if use_valgrind and use_profiler:
         raise ValueError("Cannot use valgrind and profiler at the same time.")
@@ -1315,10 +1330,6 @@ def start_raylet(redis_address,
     if load_code_from_local:
         start_worker_command += ["--load-code-from-local"]
 
-    if object_spilling_config:
-        start_worker_command.append(
-            f"--object-spilling-config={json.dumps(object_spilling_config)}")
-
     # Create agent command
     agent_command = [
         sys.executable,

python/ray/node.py

@@ -734,7 +734,6 @@ class Node:
             head_node=self.head,
             start_initial_python_workers_for_first_job=self._ray_params.
             start_initial_python_workers_for_first_job,
-            object_spilling_config=self._ray_params.object_spilling_config,
             code_search_path=self._ray_params.code_search_path)
         assert ray_constants.PROCESS_TYPE_RAYLET not in self.all_processes
         self.all_processes[ray_constants.PROCESS_TYPE_RAYLET] = [process_info]

python/ray/parameter.py

@@ -1,3 +1,4 @@
+import json
 import logging
 import os
 
@@ -148,7 +149,6 @@ class RayParams:
                  metrics_agent_port=None,
                  metrics_export_port=None,
                  lru_evict=False,
-                 object_spilling_config=None,
                  code_search_path=None):
         self.object_ref_seed = object_ref_seed
         self.redis_address = redis_address
@@ -191,10 +191,9 @@ class RayParams:
         self.metrics_export_port = metrics_export_port
         self.start_initial_python_workers_for_first_job = (
             start_initial_python_workers_for_first_job)
-        self._system_config = _system_config
+        self._system_config = _system_config or {}
         self._lru_evict = lru_evict
         self._enable_object_reconstruction = enable_object_reconstruction
-        self.object_spilling_config = object_spilling_config
         self._check_usage()
         self.code_search_path = code_search_path
         if code_search_path is None:
@@ -322,8 +321,10 @@ class RayParams:
                            "serialization. Upgrade numpy if using with ray.")
 
         # Make sure object spilling configuration is applicable.
-        object_spilling_config = self.object_spilling_config or {}
+        object_spilling_config = self._system_config.get(
+            "object_spilling_config", {})
         if object_spilling_config:
+            object_spilling_config = json.loads(object_spilling_config)
             from ray import external_storage
             # Validate external storage usage.
             external_storage.setup_external_storage(object_spilling_config)

python/ray/tests/test_object_spilling.py

@@ -26,14 +26,16 @@ smart_open_object_spilling_config = {
 
 
 @pytest.fixture(
-    scope="module",
+    scope="function",
     params=[
         file_system_object_spilling_config,
         # TODO(sang): Add a mock dependency to test S3.
         # smart_open_object_spilling_config,
     ])
-def object_spilling_config(request):
-    yield request.param
+def object_spilling_config(request, tmpdir):
+    if request.param["type"] == "filesystem":
+        request.param["params"]["directory_path"] = str(tmpdir)
+    yield json.dumps(request.param)
 
 
 @pytest.mark.skip("This test is for local benchmark.")
@@ -48,10 +50,10 @@ def test_sample_benchmark(object_spilling_config, shutdown_only):
     # Limit our object store to 200 MiB of memory.
     ray.init(
         object_store_memory=object_store_limit,
-        _object_spilling_config=object_spilling_config,
         _system_config={
             "object_store_full_max_retries": 0,
             "max_io_workers": max_io_workers,
+            "object_spilling_config": object_spilling_config,
         })
     arr = np.random.rand(object_size)
     replay_buffer = []
@@ -91,18 +93,26 @@ def test_invalid_config_raises_exception(shutdown_only):
     # it starts processes when invalid object spilling
     # config is given.
     with pytest.raises(ValueError):
-        ray.init(_object_spilling_config={"type": "abc"})
+        ray.init(_system_config={
+            "object_spilling_config": json.dumps({
+                "type": "abc"
+            }),
+        })
 
     with pytest.raises(Exception):
         copied_config = copy.deepcopy(file_system_object_spilling_config)
         # Add invalid params to the config.
         copied_config["params"].update({"random_arg": "abc"})
-        ray.init(_object_spilling_config=copied_config)
+        ray.init(_system_config={
+            "object_spilling_config": json.dumps(copied_config),
+        })
 
     with pytest.raises(ValueError):
         copied_config = copy.deepcopy(file_system_object_spilling_config)
         copied_config["params"].update({"directory_path": "not_exist_path"})
-        ray.init(_object_spilling_config=copied_config)
+        ray.init(_system_config={
+            "object_spilling_config": json.dumps(copied_config),
+        })
 
 
 @pytest.mark.skipif(
@@ -111,10 +121,11 @@ def test_spill_objects_manually(object_spilling_config, shutdown_only):
     # Limit our object store to 75 MiB of memory.
     ray.init(
         object_store_memory=75 * 1024 * 1024,
-        _object_spilling_config=object_spilling_config,
         _system_config={
             "object_store_full_max_retries": 0,
+            "automatic_object_spilling_enabled": False,
             "max_io_workers": 4,
+            "object_spilling_config": object_spilling_config,
         })
     arr = np.random.rand(1024 * 1024)  # 8 MB data
     replay_buffer = []
@@ -161,10 +172,11 @@ def test_spill_objects_manually_from_workers(object_spilling_config,
     # Limit our object store to 100 MiB of memory.
     ray.init(
         object_store_memory=100 * 1024 * 1024,
-        _object_spilling_config=object_spilling_config,
         _system_config={
             "object_store_full_max_retries": 0,
+            "automatic_object_spilling_enabled": False,
             "max_io_workers": 4,
+            "object_spilling_config": object_spilling_config,
         })
 
     @ray.remote
@@ -190,10 +202,11 @@ def test_spill_objects_manually_with_workers(object_spilling_config,
     # Limit our object store to 75 MiB of memory.
     ray.init(
         object_store_memory=100 * 1024 * 1024,
-        _object_spilling_config=object_spilling_config,
         _system_config={
             "object_store_full_max_retries": 0,
+            "automatic_object_spilling_enabled": False,
             "max_io_workers": 4,
+            "object_spilling_config": object_spilling_config,
         })
     arrays = [np.random.rand(100 * 1024) for _ in range(50)]
     objects = [ray.put(arr) for arr in arrays]
@@ -210,23 +223,27 @@ def test_spill_objects_manually_with_workers(object_spilling_config,
 
 @pytest.mark.skipif(
     platform.system() == "Windows", reason="Failing on Windows.")
-def test_spill_remote_object(object_spilling_config, ray_start_cluster):
-    cluster = ray_start_cluster
-    # # Head node.
-    cluster.add_node(
-        num_cpus=0,
-        object_store_memory=75 * 1024 * 1024,
-        object_spilling_config=object_spilling_config,
-        _system_config={
-            "object_store_full_max_retries": 0,
+@pytest.mark.parametrize(
+    "ray_start_cluster_head", [{
+        "num_cpus": 0,
+        "object_store_memory": 75 * 1024 * 1024,
+        "_system_config": {
+            "automatic_object_spilling_enabled": True,
+            "object_store_full_max_retries": 4,
+            "object_store_full_initial_delay_ms": 100,
             "max_io_workers": 4,
-        })
-    # Worker nodes.
-    cluster.add_node(
-        object_store_memory=75 * 1024 * 1024,
-        object_spilling_config=object_spilling_config)
-    cluster.wait_for_nodes()
-    ray.init(address=cluster.address)
+            "object_spilling_config": json.dumps({
+                "type": "filesystem",
+                "params": {
+                    "directory_path": "/tmp"
+                }
+            }),
+        },
+    }],
+    indirect=True)
+def test_spill_remote_object(ray_start_cluster_head):
+    cluster = ray_start_cluster_head
+    cluster.add_node(object_store_memory=75 * 1024 * 1024)
 
     @ray.remote
     def put():
@@ -240,11 +257,7 @@ def test_spill_remote_object(object_spilling_config, ray_start_cluster):
     copy = np.copy(ray.get(ref))
     # Evict local copy.
     ray.put(np.random.rand(5 * 1024 * 1024))  # 40 MB data
-    # Remote copy should not fit.
-    with pytest.raises(ray.exceptions.RayTaskError):
-        ray.get(put.remote())
-    # Spill 1 object. The second should now fit.
-    ray.experimental.force_spill_objects([ref])
+    # Remote copy should cause first remote object to get spilled.
     ray.get(put.remote())
 
     sample = ray.get(ref)
@@ -258,17 +271,17 @@ def test_spill_remote_object(object_spilling_config, ray_start_cluster):
     ray.get(depends.remote(ref))
 
 
-@pytest.mark.skip(reason="Not implemented yet.")
-def test_spill_objects_automatically(shutdown_only):
+def test_spill_objects_automatically(object_spilling_config, shutdown_only):
     # Limit our object store to 75 MiB of memory.
     ray.init(
         object_store_memory=75 * 1024 * 1024,
-        _system_config=json.dumps({
+        _system_config={
             "max_io_workers": 4,
-            "object_store_full_max_retries": 2,
-            "object_store_full_initial_delay_ms": 10,
-            "auto_object_spilling": True,
-        }))
+            "automatic_object_spilling_enabled": True,
+            "object_store_full_max_retries": 4,
+            "object_store_full_initial_delay_ms": 100,
+            "object_spilling_config": object_spilling_config,
+        })
     arr = np.random.rand(1024 * 1024)  # 8 MB data
     replay_buffer = []
 
@@ -291,5 +304,37 @@ def test_spill_objects_automatically(shutdown_only):
         assert np.array_equal(sample, arr)
 
 
+@pytest.mark.skipif(
+    platform.system() == "Windows", reason="Failing on Windows.")
+def test_spill_during_get(object_spilling_config, shutdown_only):
+    ray.init(
+        num_cpus=4,
+        object_store_memory=100 * 1024 * 1024,
+        _system_config={
+            "automatic_object_spilling_enabled": True,
+            # This test will deadlock if only one IO worker is allowed because
+            # the IO worker will try to restore an object, but this requires
+            # another object to be spilled, which also requires an IO worker.
+            "max_io_workers": 2,
+            "object_spilling_config": object_spilling_config,
+        },
+    )
+
+    @ray.remote
+    def f():
+        return np.zeros(10 * 1024 * 1024)
+
+    ids = []
+    for i in range(10):
+        x = f.remote()
+        print(i, x)
+        ids.append(x)
+
+    # Concurrent gets, which require restoring from external storage, while
+    # objects are being created.
+    for x in ids:
+        print(ray.get(x).shape)
+
+
 if __name__ == "__main__":
     sys.exit(pytest.main(["-sv", __file__]))

python/ray/worker.py

@@ -510,7 +510,6 @@ def init(
         _load_code_from_local=False,
         _lru_evict=False,
         _metrics_export_port=None,
-        _object_spilling_config=None,
         _system_config=None):
     """
     Connect to an existing Ray cluster or start one and connect to it.
@@ -609,8 +608,6 @@ def init(
         _metrics_export_port(int): Port number Ray exposes system metrics
             through a Prometheus endpoint. It is currently under active
             development, and the API is subject to change.
-        _object_spilling_config (str): The configuration json string for object
-            spilling I/O worker.
         _system_config (dict): Configuration for overriding
             RayConfig defaults. For testing purposes ONLY.
 
@@ -727,8 +724,7 @@ def init(
             _system_config=_system_config,
             lru_evict=_lru_evict,
             enable_object_reconstruction=_enable_object_reconstruction,
-            metrics_export_port=_metrics_export_port,
-            object_spilling_config=_object_spilling_config)
+            metrics_export_port=_metrics_export_port)
         # Start the Ray processes. We set shutdown_at_exit=False because we
         # shutdown the node in the ray.shutdown call that happens in the atexit
         # handler. We still spawn a reaper process in case the atexit handler

python/ray/workers/default_worker.py

@@ -1,4 +1,5 @@
 import argparse
+import base64
 import json
 import time
 import sys
@@ -125,20 +126,13 @@ if __name__ == "__main__":
     if mode == ray.IO_WORKER_MODE:
         from ray import external_storage
         if args.object_spilling_config:
-            object_spilling_config = json.loads(args.object_spilling_config)
+            object_spilling_config = base64.b64decode(
+                args.object_spilling_config)
+            object_spilling_config = json.loads(object_spilling_config)
         else:
             object_spilling_config = {}
         external_storage.setup_external_storage(object_spilling_config)
 
-    system_config = {}
-    if args.config_list is not None:
-        config_list = args.config_list.split(",")
-        if len(config_list) > 1:
-            i = 0
-            while i < len(config_list):
-                system_config[config_list[i]] = config_list[i + 1]
-                i += 2
-
     raylet_ip_address = args.raylet_ip_address
     if raylet_ip_address is None:
         raylet_ip_address = args.node_ip_address
@@ -161,7 +155,6 @@ if __name__ == "__main__":
         temp_dir=args.temp_dir,
         load_code_from_local=args.load_code_from_local,
         metrics_agent_port=args.metrics_agent_port,
-        _system_config=system_config,
     )
 
     node = ray.node.Node(