From 1c913308eff2fd02b37e1b740e2fe0339db5c31c Mon Sep 17 00:00:00 2001 From: Bob Yang Date: Tue, 10 Jun 2025 10:51:25 -0700 Subject: [PATCH] expose runner itself with context manager (#1069) 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 apis in sub classes. Here, we expose the self type instead to prevent erasing types Reviewed By: kiukchung Differential Revision: D72677312 Pulled By: bobyangyf --- torchx/runner/api.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/torchx/runner/api.py b/torchx/runner/api.py index 922aca539..1731c4657 100644 --- a/torchx/runner/api.py +++ b/torchx/runner/api.py @@ -43,6 +43,7 @@ from torchx.util.types import none_throws from torchx.workspace.api import PkgInfo, WorkspaceBuilder, WorkspaceMixin +from typing_extensions import Self from .config import get_config, get_configs @@ -120,7 +121,7 @@ def _get_scheduler_params_from_env(self) -> Dict[str, str]: scheduler_params[lower_case_key.strip("torchx_")] = value return scheduler_params - def __enter__(self) -> "Runner": + def __enter__(self) -> Self: return self def __exit__(