From 84b7240c4bcbd9373ec18db8be483a3e41a4f02a Mon Sep 17 00:00:00 2001
From: Alex Wu <itswu.alex@gmail.com>
Date: Thu, 6 Aug 2020 18:32:01 -0700
Subject: [PATCH] [Core] Read resources from an environment variable (#9831)

---
 python/ray/node.py                  | 28 +++++++++++++++++++++++++---
 python/ray/tests/test_advanced_3.py | 13 +++++++++++++
 2 files changed, 38 insertions(+), 3 deletions(-)

diff --git a/python/ray/node.py b/python/ray/node.py
index 880cdbeca..a92dbf30e 100644
--- a/python/ray/node.py
+++ b/python/ray/node.py
@@ -2,8 +2,9 @@ import atexit
 import collections
 import datetime
 import errno
-import os
+import json
 import logging
+import os
 import random
 import signal
 import socket
@@ -254,12 +255,33 @@ class Node:
 
     def get_resource_spec(self):
         """Resolve and return the current resource spec for the node."""
+
+        def merge_resources(env_dict, params_dict):
+            """Merge two dictionaries, picking from the second in the event of a conflict.
+            Also emit a warning on every conflict.
+            """
+            result = params_dict.copy()
+            result.update(env_dict)
+
+            for key in set(env_dict.keys()).intersection(
+                    set(params_dict.keys())):
+                logger.warning("Autoscaler is overriding your resource:"
+                               "{}: {} with {}.".format(
+                                   key, params_dict[key], env_dict[key]))
+            return result
+
+        env_resources = {}
+        env_string = os.getenv("RAY_OVERRIDE_RESOURCES")
+        if env_string:
+            env_resources = json.loads(env_string)
+
         if not self._resource_spec:
+            resources = merge_resources(env_resources,
+                                        self._ray_params.resources)
             self._resource_spec = ResourceSpec(
                 self._ray_params.num_cpus, self._ray_params.num_gpus,
                 self._ray_params.memory, self._ray_params.object_store_memory,
-                self._ray_params.resources,
-                self._ray_params.redis_max_memory).resolve(
+                resources, self._ray_params.redis_max_memory).resolve(
                     is_head=self.head, node_ip_address=self.node_ip_address)
         return self._resource_spec
 
diff --git a/python/ray/tests/test_advanced_3.py b/python/ray/tests/test_advanced_3.py
index 94fde6b11..7a82fd713 100644
--- a/python/ray/tests/test_advanced_3.py
+++ b/python/ray/tests/test_advanced_3.py
@@ -656,6 +656,19 @@ def test_ray_address_environment_variable(ray_start_cluster):
     ray.shutdown()
 
 
+def test_ray_resources_environment_variable(ray_start_cluster):
+    address = ray_start_cluster.address
+
+    os.environ["RAY_OVERRIDE_RESOURCES"] = "{\"custom1\":1, \"custom2\":2}"
+    ray.init(address=address, resources={"custom1": 3, "custom3": 3})
+
+    cluster_resources = ray.cluster_resources()
+    print(cluster_resources)
+    assert cluster_resources["custom1"] == 1
+    assert cluster_resources["custom2"] == 2
+    assert cluster_resources["custom3"] == 3
+
+
 def test_gpu_info_parsing():
     info_string = """Model:           Tesla V100-SXM2-16GB
 IRQ:             107