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

Author: bobyangyf

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__(