Improved logging - printing a caller script name upon start

Author: ivan-khvostishkov

sagemaker_ssh_helper/env.py

@@ -1,5 +1,6 @@
 import json
 import os
+from typing import List
 
 
 def sm_get_node_rank():
@@ -15,3 +16,15 @@ def sm_get_node_rank():
         rc = json.load(json_file)
 
     return int(rc["hosts"].index(rc["current_host"]))
+
+
+def get_caller_script_name(trace_back=1):
+    import inspect
+    from inspect import FrameInfo
+    from pathlib import Path
+
+    stack: List[FrameInfo] = inspect.stack()
+    if len(stack) < 2:
+        raise AssertionError("Cannot fetch stack trace")
+
+    return Path(stack[trace_back].filename).name