Add torchx session id as Environment variable

yikaiMeta · web-flow · commit 648913c90cd0 · 2024-11-26T17:33:07.000-08:00
Differential Revision: D66387522 Pull Request resolved: #974
diff --git a/torchx/runner/api.py b/torchx/runner/api.py
@@ -39,6 +39,7 @@
     ENV_TORCHX_TRACKERS,
     tracker_config_env_var_name,
 )
+from torchx.util.session import get_session_id_or_create_new, TORCHX_INTERNAL_SESSION_ID
 
 from torchx.util.types import none_throws
 from torchx.workspace.api import PkgInfo, WorkspaceBuilder, WorkspaceMixin
@@ -390,6 +391,7 @@ def dryrun(
             role.env[ENV_TORCHX_JOB_ID] = make_app_handle(
                 scheduler, self._name, macros.app_id
             )
+            role.env[TORCHX_INTERNAL_SESSION_ID] = get_session_id_or_create_new()
 
             if parent_run_id:
                 role.env[ENV_TORCHX_PARENT_RUN_ID] = parent_run_id

Original file line number	Diff line number	Diff line change
`@@ -39,6 +39,7 @@`
`39`	`39`	`ENV_TORCHX_TRACKERS,`
`40`	`40`	`tracker_config_env_var_name,`
`41`	`41`	`)`
	`42`	`+from torchx.util.session import get_session_id_or_create_new, TORCHX_INTERNAL_SESSION_ID`
`42`	`43`
`43`	`44`	`from torchx.util.types import none_throws`
`44`	`45`	`from torchx.workspace.api import PkgInfo, WorkspaceBuilder, WorkspaceMixin`
`@@ -390,6 +391,7 @@ def dryrun(`
`390`	`391`	`role.env[ENV_TORCHX_JOB_ID] = make_app_handle(`
`391`	`392`	`scheduler, self._name, macros.app_id`
`392`	`393`	`)`
	`394`	`+ role.env[TORCHX_INTERNAL_SESSION_ID] = get_session_id_or_create_new()`
`393`	`395`
`394`	`396`	`if parent_run_id:`
`395`	`397`	`role.env[ENV_TORCHX_PARENT_RUN_ID] = parent_run_id`