From 157596ee48af0ee9bf7b3079d72e18496541f611 Mon Sep 17 00:00:00 2001 From: Shicong Huang Date: Mon, 21 Oct 2024 14:21:32 -0700 Subject: [PATCH] Allow setting torchx session id from cmd (#967) Summary: This will be used by e2e test to locate the test runs and retrieve metrics for comparison. Differential Revision: D64678003 --- torchx/util/session.py | 7 +++++++ 1 file changed, 7 insertions(+) diff --git a/torchx/util/session.py b/torchx/util/session.py index f03b2b10b..6068169fb 100644 --- a/torchx/util/session.py +++ b/torchx/util/session.py @@ -7,9 +7,12 @@ # pyre-strict +import os import uuid from typing import Optional +TORCHX_INTERNAL_SESSION_ID = "TORCHX_INTERNAL_SESSION_ID" + CURRENT_SESSION_ID: Optional[str] = None @@ -22,6 +25,10 @@ def get_session_id_or_create_new() -> str: global CURRENT_SESSION_ID if CURRENT_SESSION_ID: return CURRENT_SESSION_ID + env_session_id = os.getenv(TORCHX_INTERNAL_SESSION_ID) + if env_session_id: + CURRENT_SESSION_ID = env_session_id + return CURRENT_SESSION_ID session_id = str(uuid.uuid4()) CURRENT_SESSION_ID = session_id return session_id