(torchx/components) Remove dependency to pyre_extensions

Differential Revision: D77621041

Pull Request resolved: #1088

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: