From f5d25f09f02e1845ea487058f4e08d16153d4f59 Mon Sep 17 00:00:00 2001
From: yikaiMeta <116014337+yikaiMeta@users.noreply.github.com>
Date: Tue, 1 Oct 2024 13:30:38 -0400
Subject: [PATCH] expose torchx session ID to torchx downstream users.

Differential Revision: D63655049

Pull Request resolved: https://github.com/pytorch/torchx/pull/959
---
 torchx/util/session.py | 9 +++++++++
 1 file changed, 9 insertions(+)

diff --git a/torchx/util/session.py b/torchx/util/session.py
index 80774fab7..f03b2b10b 100644
--- a/torchx/util/session.py
+++ b/torchx/util/session.py
@@ -17,6 +17,7 @@ def get_session_id_or_create_new() -> str:
     """
     Returns the current session ID, or creates a new one if none exists.
     The session ID remains the same as long as it is in the same process.
+    Please DO NOT use this function out of torchx codebase.
     """
     global CURRENT_SESSION_ID
     if CURRENT_SESSION_ID:
@@ -24,3 +25,11 @@ def get_session_id_or_create_new() -> str:
     session_id = str(uuid.uuid4())
     CURRENT_SESSION_ID = session_id
     return session_id
+
+
+def get_torchx_session_id() -> Optional[str]:
+    """
+    Returns the torchx session ID.
+    Please use this function to get the session ID out of torchx codebase.
+    """
+    return CURRENT_SESSION_ID