From 3ce92869778ecb20375785beeb6e85dda9ab7821 Mon Sep 17 00:00:00 2001
From: fyrestone <fyrestone@outlook.com>
Date: Wed, 9 Dec 2020 22:12:34 +0800
Subject: [PATCH] Fix dashboard agent check ppid is raylet pid (#12256)
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

* Dashboard agent check ppid is raylet pid

* Improve implementation

* Refine code

* Make the RAY_NODE_PID environment required for dashboard agent

Co-authored-by: 刘宝 <po.lb@antfin.com>
---
 dashboard/agent.py                | 5 ++++-
 dashboard/tests/test_dashboard.py | 5 +++++
 src/ray/raylet/agent_manager.cc   | 2 ++
 3 files changed, 11 insertions(+), 1 deletion(-)

diff --git a/dashboard/agent.py b/dashboard/agent.py
index 676312245..ef146e129 100644
--- a/dashboard/agent.py
+++ b/dashboard/agent.py
@@ -91,9 +91,12 @@ class DashboardAgent(object):
         async def _check_parent():
             """Check if raylet is dead."""
             curr_proc = psutil.Process()
+            ppid = int(os.environ["RAY_NODE_PID"])
+            logger.info("Parent pid is %s", ppid)
             while True:
                 parent = curr_proc.parent()
-                if parent is None or parent.pid == 1:
+                if parent is None or parent.pid == 1 or (ppid and
+                                                         ppid != parent.pid):
                     logger.error("raylet is dead, agent will die because "
                                  "it fate-shares with raylet.")
                     sys.exit(0)
diff --git a/dashboard/tests/test_dashboard.py b/dashboard/tests/test_dashboard.py
index a37fd3012..1b9f82120 100644
--- a/dashboard/tests/test_dashboard.py
+++ b/dashboard/tests/test_dashboard.py
@@ -137,6 +137,11 @@ def test_basic(ray_start_with_dashboard):
         assert agent_proc.pid == agent_pid
         time.sleep(1)
 
+    # The agent should be dead if raylet exits.
+    raylet_proc.kill()
+    raylet_proc.wait()
+    agent_proc.wait(5)
+
     # Check redis keys are set.
     logger.info("Check redis keys are set.")
     dashboard_address = client.get(dashboard_consts.REDIS_KEY_DASHBOARD)
diff --git a/src/ray/raylet/agent_manager.cc b/src/ray/raylet/agent_manager.cc
index 23b8769c8..ef9d8c45d 100644
--- a/src/ray/raylet/agent_manager.cc
+++ b/src/ray/raylet/agent_manager.cc
@@ -58,8 +58,10 @@ void AgentManager::StartAgent() {
   }
   argv.push_back(NULL);
   // Set node id to agent.
+  static std::string pid_string = std::to_string(getpid());
   ProcessEnvironment env;
   env.insert({"RAY_NODE_ID", options_.node_id.Hex()});
+  env.insert({"RAY_NODE_PID", pid_string});
   Process child(argv.data(), nullptr, ec, false, env);
   if (!child.IsValid() || ec) {
     // The worker failed to start. This is a fatal error.