[core] Update metadata in options properly (#24458)

* implement proper updating of metadata in options

python/ray/_private/ray_option_utils.py

@@ -221,3 +221,28 @@ def validate_actor_options(options: Dict[str, Any], in_options: bool):
     if options.get("get_if_exists") and not options.get("name"):
         raise ValueError("The actor name must be specified to use `get_if_exists`.")
     _check_deprecate_placement_group(options)
+
+
+def update_options(
+    original_options: Dict[str, Any], new_options: Dict[str, Any]
+) -> Dict[str, Any]:
+    """Update original options with new options and return.
+    The returned updated options contain shallow copy of original options.
+    """
+
+    updated_options = {**original_options, **new_options}
+    # Ensure we update each namespace in "_metadata" independently.
+    # "_metadata" is a dict like {namespace1: config1, namespace2: config2}
+    if (
+        original_options.get("_metadata") is not None
+        and new_options.get("_metadata") is not None
+    ):
+        # make a shallow copy to avoid messing up the metadata dict in
+        # the original options.
+        metadata = original_options["_metadata"].copy()
+        for namespace, config in new_options["_metadata"].items():
+            metadata[namespace] = {**metadata.get(namespace, {}), **config}
+
+        updated_options["_metadata"] = metadata
+
+    return updated_options

python/ray/actor.py

@@ -547,7 +547,9 @@ class ActorClass:
         # "concurrency_groups" could not be used in ".options()",
         # we should remove it before merging options from '@ray.remote'.
         default_options.pop("concurrency_groups", None)
-        updated_options = {**default_options, **actor_options}
+        updated_options = ray_option_utils.update_options(
+            default_options, actor_options
+        )
         ray_option_utils.validate_actor_options(updated_options, in_options=True)
 
         # only update runtime_env when ".options()" specifies new runtime_env

python/ray/remote_function.py

@@ -143,7 +143,7 @@ class RemoteFunction:
         # max_calls could not be used in ".options()", we should remove it before
         # merging options from '@ray.remote'.
         default_options.pop("max_calls", None)
-        updated_options = {**default_options, **task_options}
+        updated_options = ray_option_utils.update_options(default_options, task_options)
         ray_option_utils.validate_task_options(updated_options, in_options=True)
 
         # only update runtime_env when ".options()" specifies new runtime_env

python/ray/tests/test_basic.py

@@ -284,6 +284,32 @@ def test_options():
         with pytest.raises(TypeError):
             v.validate(k, unique_object)
 
+    # test updating each namespace of "_metadata" independently
+    assert ray_option_utils.update_options(
+        {
+            "_metadata": {"ns1": {"a1": 1, "b1": 2, "c1": 3}, "ns2": {"a2": 1}},
+            "num_cpus": 1,
+            "xxx": {"x": 2},
+            "zzz": 42,
+        },
+        {
+            "_metadata": {"ns1": {"b1": 22}, "ns3": {"b3": 2}},
+            "num_cpus": 2,
+            "xxx": {"y": 2},
+            "yyy": 3,
+        },
+    ) == {
+        "_metadata": {
+            "ns1": {"a1": 1, "b1": 22, "c1": 3},
+            "ns2": {"a2": 1},
+            "ns3": {"b3": 2},
+        },
+        "num_cpus": 2,
+        "xxx": {"y": 2},
+        "yyy": 3,
+        "zzz": 42,
+    }
+
 
 # https://github.com/ray-project/ray/issues/17842
 def test_disable_cuda_devices():