(torchx/components) Remove dependency to pyre_extensions (#1088)

Summary:
Pull Request resolved: #1088

There was only a single usage of `pyre_extensions` in the entire code base. Get rid of this dependency by using an `assert` statement rather than `pyre_extensions.none_throws`.

Reviewed By: highker

Differential Revision: D77621041

Author: kiukchung

dev-requirements.txt

@@ -30,6 +30,7 @@ torchmetrics==1.6.3
 torchserve>=0.10.0
 torchtext==0.18.0
 torchvision==0.22.0
+typing-extensions
 ts==0.5.1
 ray[default]
 wheel

requirements.txt

@@ -1,4 +1,3 @@
-pyre-extensions
 docstring-parser>=0.8.1
 importlib-metadata
 pyyaml

torchx/components/structured_arg.py

@@ -30,8 +30,6 @@
 from pathlib import Path
 from typing import Optional
 
-from pyre_extensions import none_throws
-
 from torchx import specs
 
 
@@ -148,7 +146,8 @@ def parse_from(
             if m:  # use the last module name
                 run_name = m.rpartition(".")[2]
             else:  # use script name w/ no extension
-                run_name = Path(none_throws(script)).stem
+                assert script, "`script` can't be `None` here due checks above"
+                run_name = Path(script).stem
         return StructuredNameArgument(
             experiment_name or default_experiment_name, run_name
         )

torchx/runner/api.py

@@ -14,7 +14,18 @@
 import warnings
 from datetime import datetime
 from types import TracebackType
-from typing import Any, Dict, Iterable, List, Mapping, Optional, Tuple, Type, TypeVar
+from typing import (
+    Any,
+    Dict,
+    Iterable,
+    List,
+    Mapping,
+    Optional,
+    Tuple,
+    Type,
+    TYPE_CHECKING,
+    TypeVar,
+)
 
 from torchx.runner.events import log_event
 from torchx.schedulers import get_scheduler_factories, SchedulerFactory
@@ -43,7 +54,9 @@
 
 from torchx.util.types import none_throws
 from torchx.workspace.api import PkgInfo, WorkspaceBuilder, WorkspaceMixin
-from typing_extensions import Self
+
+if TYPE_CHECKING:
+    from typing_extensions import Self
 
 from .config import get_config, get_configs
 
@@ -121,7 +134,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) -> Self:
+    def __enter__(self) -> "Self":
         return self
 
     def __exit__(

torchx/schedulers/__init__.py

@@ -8,11 +8,10 @@
 # pyre-strict
 
 import importlib
-from typing import Dict, Mapping
+from typing import Mapping, Protocol
 
 from torchx.schedulers.api import Scheduler
 from torchx.util.entrypoints import load_group
-from typing_extensions import Protocol
 
 DEFAULT_SCHEDULER_MODULES: Mapping[str, str] = {
     "local_docker": "torchx.schedulers.docker_scheduler",
@@ -44,15 +43,15 @@ def run(*args: object, **kwargs: object) -> Scheduler:
 
 def get_scheduler_factories(
     group: str = "torchx.schedulers", skip_defaults: bool = False
-) -> Dict[str, SchedulerFactory]:
+) -> dict[str, SchedulerFactory]:
     """
     get_scheduler_factories returns all the available schedulers names under `group` and the
     method to instantiate them.
 
     The first scheduler in the dictionary is used as the default scheduler.
     """
 
-    default_schedulers: Dict[str, SchedulerFactory] = {}
+    default_schedulers: dict[str, SchedulerFactory] = {}
     for scheduler, path in DEFAULT_SCHEDULER_MODULES.items():
         default_schedulers[scheduler] = _defer_load_scheduler(path)
 

torchx/schedulers/api.py

@@ -16,7 +16,6 @@
 
 from torchx.specs import (
     AppDef,
-    AppDryRunInfo,
     AppState,
     NONE,
     NULL_RESOURCE,

torchx/schedulers/aws_batch_scheduler.py

@@ -53,13 +53,13 @@
     Optional,
     Tuple,
     TYPE_CHECKING,
+    TypedDict,
     TypeVar,
 )
 
 import torchx
 import yaml
 from torchx.schedulers.api import (
-    AppDryRunInfo,
     DescribeAppResponse,
     filter_regex,
     ListAppResponse,
@@ -71,6 +71,7 @@
 from torchx.schedulers.ids import make_unique
 from torchx.specs.api import (
     AppDef,
+    AppDryRunInfo,
     AppState,
     BindMount,
     CfgVal,
@@ -86,7 +87,6 @@
 from torchx.specs.named_resources_aws import instance_type_from_resource
 from torchx.util.types import none_throws
 from torchx.workspace.docker_workspace import DockerWorkspaceMixin
-from typing_extensions import TypedDict
 
 ENV_TORCHX_ROLE_IDX = "TORCHX_ROLE_IDX"
 

torchx/schedulers/aws_sagemaker_scheduler.py

@@ -25,6 +25,7 @@
     OrderedDict,
     Tuple,
     TYPE_CHECKING,
+    TypedDict,
     TypeVar,
 )
 
@@ -34,16 +35,14 @@
 from sagemaker.pytorch import PyTorch
 from torchx.components.structured_arg import StructuredNameArgument
 from torchx.schedulers.api import (
-    AppDryRunInfo,
     DescribeAppResponse,
     ListAppResponse,
     Scheduler,
     Stream,
 )
 from torchx.schedulers.ids import make_unique
-from torchx.specs.api import AppDef, AppState, CfgVal, runopts
+from torchx.specs.api import AppDef, AppDryRunInfo, AppState, CfgVal, runopts
 from torchx.workspace.docker_workspace import DockerWorkspaceMixin
-from typing_extensions import TypedDict
 
 
 if TYPE_CHECKING:

torchx/schedulers/docker_scheduler.py

@@ -13,12 +13,11 @@
 import tempfile
 from dataclasses import dataclass
 from datetime import datetime
-from typing import Any, Dict, Iterable, List, Optional, TYPE_CHECKING, Union
+from typing import Any, Dict, Iterable, List, Optional, TYPE_CHECKING, TypedDict, Union
 
 import torchx
 import yaml
 from torchx.schedulers.api import (
-    AppDryRunInfo,
     DescribeAppResponse,
     filter_regex,
     ListAppResponse,
@@ -30,6 +29,7 @@
 from torchx.schedulers.ids import make_unique
 from torchx.specs.api import (
     AppDef,
+    AppDryRunInfo,
     AppState,
     BindMount,
     DeviceMount,
@@ -42,7 +42,6 @@
     VolumeMount,
 )
 from torchx.workspace.docker_workspace import DockerWorkspaceMixin
-from typing_extensions import TypedDict
 
 
 if TYPE_CHECKING:

torchx/schedulers/gcp_batch_scheduler.py

@@ -24,22 +24,28 @@
 
 from dataclasses import dataclass
 from datetime import datetime
-from typing import Any, Dict, Iterable, List, Optional, TYPE_CHECKING
+from typing import Any, Dict, Iterable, List, Optional, TYPE_CHECKING, TypedDict
 
 import torchx
 import yaml
 
 from torchx.schedulers.api import (
-    AppDryRunInfo,
     DescribeAppResponse,
     ListAppResponse,
     Scheduler,
     Stream,
 )
 from torchx.schedulers.ids import make_unique
-from torchx.specs.api import AppDef, AppState, macros, Resource, Role, runopts
+from torchx.specs.api import (
+    AppDef,
+    AppDryRunInfo,
+    AppState,
+    macros,
+    Resource,
+    Role,
+    runopts,
+)
 from torchx.util.strings import normalize_str
-from typing_extensions import TypedDict
 
 
 if TYPE_CHECKING:

torchx/schedulers/kubernetes_mcad_scheduler.py

@@ -17,8 +17,8 @@
 
 TorchX Kubernetes_MCAD scheduler depends on AppWrapper + MCAD.
 
-Install MCAD: 
-See deploying Multi-Cluster-Application-Dispatcher guide 
+Install MCAD:
+See deploying Multi-Cluster-Application-Dispatcher guide
 https://github.com/project-codeflare/multi-cluster-app-dispatcher/blob/main/doc/deploy/deployment.md
 
 This implementation requires MCAD v1.34.1 or higher.
@@ -46,12 +46,12 @@
     Optional,
     Tuple,
     TYPE_CHECKING,
+    TypedDict,
 )
 
 import torchx
 import yaml
 from torchx.schedulers.api import (
-    AppDryRunInfo,
     DescribeAppResponse,
     filter_regex,
     ListAppResponse,
@@ -62,6 +62,7 @@
 from torchx.schedulers.ids import make_unique
 from torchx.specs.api import (
     AppDef,
+    AppDryRunInfo,
     AppState,
     BindMount,
     CfgVal,
@@ -78,7 +79,6 @@
 )
 
 from torchx.workspace.docker_workspace import DockerWorkspaceMixin
-from typing_extensions import TypedDict
 
 if TYPE_CHECKING:
     from docker import DockerClient
@@ -600,7 +600,7 @@ def app_to_resource(
 
     """
     Create Service:
-    The selector will have the key 'appwrapper.workload.codeflare.dev', and the value will be 
+    The selector will have the key 'appwrapper.workload.codeflare.dev', and the value will be
     the appwrapper name
     """
 
@@ -797,7 +797,8 @@ class KubernetesMCADOpts(TypedDict, total=False):
 
 
 class KubernetesMCADScheduler(
-    DockerWorkspaceMixin, Scheduler[KubernetesMCADOpts, AppDef, AppDryRunInfo]
+    DockerWorkspaceMixin,
+    Scheduler[KubernetesMCADOpts, AppDef, AppDryRunInfo[KubernetesMCADJob]],
 ):
     """
     KubernetesMCADScheduler is a TorchX scheduling interface to Kubernetes.
@@ -994,7 +995,7 @@ def _submit_dryrun(
         if image_secret is not None and service_account is not None:
             msg = """Service Account and Image Secret names are both provided.
  Depending on the Service Account configuration, an ImagePullSecret may be defined in your Service Account.
- If this is the case, check service account and image secret configurations to understand the expected behavior for 
+ If this is the case, check service account and image secret configurations to understand the expected behavior for
  patched image push access."""
             warnings.warn(msg)
         namespace = cfg.get("namespace")

torchx/schedulers/kubernetes_scheduler.py

@@ -44,12 +44,12 @@
     Optional,
     Tuple,
     TYPE_CHECKING,
+    TypedDict,
 )
 
 import torchx
 import yaml
 from torchx.schedulers.api import (
-    AppDryRunInfo,
     DescribeAppResponse,
     filter_regex,
     ListAppResponse,
@@ -60,6 +60,7 @@
 from torchx.schedulers.ids import make_unique
 from torchx.specs.api import (
     AppDef,
+    AppDryRunInfo,
     AppState,
     BindMount,
     CfgVal,
@@ -75,7 +76,6 @@
 )
 from torchx.util.strings import normalize_str
 from torchx.workspace.docker_workspace import DockerWorkspaceMixin
-from typing_extensions import TypedDict
 
 
 if TYPE_CHECKING:

torchx/schedulers/local_scheduler.py

@@ -40,10 +40,10 @@
     Protocol,
     TextIO,
     Tuple,
+    TypedDict,
 )
 
 from torchx.schedulers.api import (
-    AppDryRunInfo,
     DescribeAppResponse,
     filter_regex,
     ListAppResponse,
@@ -53,10 +53,10 @@
 )
 from torchx.schedulers.ids import make_unique
 from torchx.schedulers.streams import Tee
+from torchx.specs import AppDryRunInfo
 from torchx.specs.api import AppDef, AppState, is_terminal, macros, NONE, Role, runopts
 
 from torchx.util.types import none_throws
-from typing_extensions import TypedDict
 
 log: logging.Logger = logging.getLogger(__name__)
 

torchx/schedulers/lsf_scheduler.py

@@ -29,11 +29,10 @@
 import tempfile
 from dataclasses import dataclass
 from datetime import datetime
-from typing import Any, Dict, Iterable, List, Optional
+from typing import Any, Dict, Iterable, List, Optional, TypedDict
 
 import torchx
 from torchx.schedulers.api import (
-    AppDryRunInfo,
     DescribeAppResponse,
     filter_regex,
     ListAppResponse,
@@ -45,6 +44,7 @@
 from torchx.schedulers.local_scheduler import LogIterator
 from torchx.specs import (
     AppDef,
+    AppDryRunInfo,
     AppState,
     BindMount,
     DeviceMount,
@@ -57,7 +57,6 @@
     VolumeMount,
 )
 from torchx.util import shlex
-from typing_extensions import TypedDict
 
 JOB_STATE: Dict[str, AppState] = {
     "DONE": AppState.SUCCEEDED,