Skip to content

Commit 4639dc0

Browse files
bobyangyffacebook-github-bot
authored andcommitted
expose fbrunner itself with context manager
Summary: Right now, when we `enter` context after `get_fb_runner`, we are using the existing definition in `Runner` which causes type erasure to the super type `Runner`. While this is fine for most cases, we lose the ability to access the `dryrun` api defined in https://www.internalfb.com/code/fbsource/[587e66c70ce546a78a71b1a22125850d1aff866d]/fbcode/torchx/runner/fb/api.py?lines=103 which has extra parameters. For orchestration sdk to translate single job definitions as part of its dag correctly, we wish to use FBRunner's implementation of `dryrun`, so we need to be able to access `FBRunner` after entering the context manager. Here, we add corresponding `enter` and `exit` on the FBRunner to prevent erasing the type. All implementation still delegates to the existing super functions Reviewed By: kiukchung Differential Revision: D72677312
1 parent ae72fa2 commit 4639dc0

File tree

1 file changed

+2
-1
lines changed

1 file changed

+2
-1
lines changed

torchx/runner/api.py

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -43,6 +43,7 @@
4343

4444
from torchx.util.types import none_throws
4545
from torchx.workspace.api import PkgInfo, WorkspaceBuilder, WorkspaceMixin
46+
from typing_extensions import Self
4647

4748
from .config import get_config, get_configs
4849

@@ -120,7 +121,7 @@ def _get_scheduler_params_from_env(self) -> Dict[str, str]:
120121
scheduler_params[lower_case_key.strip("torchx_")] = value
121122
return scheduler_params
122123

123-
def __enter__(self) -> "Runner":
124+
def __enter__(self) -> Self:
124125
return self
125126

126127
def __exit__(

0 commit comments

Comments
 (0)