From 9441f85e1aac0201345b2c936db22382e1555c68 Mon Sep 17 00:00:00 2001
From: Barak Michener <me@barakmich.com>
Date: Fri, 29 Jan 2021 12:58:41 -0800
Subject: [PATCH] [client] Hook runtime context (#13750)

Change-Id: I701d21e53900b5f3fb0e23e09f59e8316c7ba623
---
 python/ray/runtime_context.py | 2 ++
 1 file changed, 2 insertions(+)

diff --git a/python/ray/runtime_context.py b/python/ray/runtime_context.py
index fa922cfa0..fed3ab132 100644
--- a/python/ray/runtime_context.py
+++ b/python/ray/runtime_context.py
@@ -1,5 +1,6 @@
 import ray.worker
 import logging
+from ray._private.client_mode_hook import client_mode_hook
 
 logger = logging.getLogger(__name__)
 
@@ -149,6 +150,7 @@ class RuntimeContext(object):
 _runtime_context = None
 
 
+@client_mode_hook
 def get_runtime_context():
     global _runtime_context
     if _runtime_context is None: