Allow setting torchx session id from cmd (#967) (#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,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