fix tests + torchx CLI (#6)

Summary:
This fixes the tests so they pass on OSS + make it so torchx entry point is correctly setup.

It also fixes a lot of imports to point to the torchx copy so it works in OSS since torchelastic is out of date there.

Pull Request resolved: #6

Test Plan:
Imported from GitHub, without a `Test Plan:` line.

  buck test //torchx/...

run OSS CI

Reviewed By: tierex

Differential Revision: D28614651

Pulled By: d4l3k

fbshipit-source-id: f8cd9577ac8ec2dbadf73d3539a4cbeb31b76bfa

Author: d4l3k

dev-requirements.txt

@@ -1,3 +1,6 @@
 boto3>=1.11.11
 moto>=2.0.6
 kfp==1.4.0
+pyre-extensions>=0.0.21
+black>=21.5b1
+isort>=5.8.0

setup.py

@@ -32,6 +32,9 @@ def get_version():
     with open("requirements.txt") as f:
         reqs = f.read()
 
+    with open("dev-requirements.txt") as f:
+        test_reqs = f.read()
+
     version = get_version()
     print("-- Building version: " + version)
 
@@ -49,9 +52,15 @@ def get_version():
         keywords=["pytorch", "machine learning"],
         python_requires=">=3.8",
         install_requires=reqs.strip().split("\n"),
+        tests_requires=test_reqs.strip().split("\n"),
         include_package_data=True,
         packages=find_packages(exclude=("*.test", "aws*", "*.fb")),
         test_suite="torchx.test.suites.unittests",
+        entry_points={
+            "console_scripts": [
+                "torchx=torchx.cli.main:main",
+            ],
+        },
         # PyPI package information.
         classifiers=[
             "Development Status :: 4 - Beta",

torchx/cli/cmd_describe.py

@@ -9,8 +9,9 @@
 import dataclasses
 import pprint
 
-import torchelastic.tsm.driver as tsm
 from torchx.cli.cmd_base import SubCommand
+from torchx.runner import get_runner
+from torchx.specs import api
 
 
 class CmdDescribe(SubCommand):
@@ -23,9 +24,9 @@ def add_arguments(self, subparser: argparse.ArgumentParser) -> None:
 
     def run(self, args: argparse.Namespace) -> None:
         app_handle = args.app_handle
-        scheduler, session_name, app_id = tsm.parse_app_handle(app_handle)
-        session = tsm.session(name=session_name)
-        app = session.describe(app_handle)
+        scheduler, session_name, app_id = api.parse_app_handle(app_handle)
+        runner = get_runner(name=session_name)
+        app = runner.describe(app_handle)
 
         if app:
             pprint.pprint(dataclasses.asdict(app), indent=2, width=80)

torchx/cli/cmd_runopts.py

@@ -7,8 +7,8 @@
 
 import argparse
 
-import torchelastic.tsm.driver as tsm
 from torchx.cli.cmd_base import SubCommand
+from torchx.runner.api import get_runner
 
 
 class CmdRunopts(SubCommand):
@@ -22,7 +22,7 @@ def add_arguments(self, subparser: argparse.ArgumentParser) -> None:
 
     def run(self, args: argparse.Namespace) -> None:
         scheduler = args.scheduler
-        run_opts = tsm.session(name="default").run_opts()
+        run_opts = get_runner().run_opts()
 
         if not scheduler:
             print(run_opts)

torchx/cli/cmd_status.py

@@ -12,8 +12,9 @@
 from string import Template
 from typing import List, Optional, Pattern
 
-import torchelastic.tsm.driver as tsm
 from torchx.cli.cmd_base import SubCommand
+from torchx.runner import get_runner
+from torchx.specs import api
 from torchx.specs.api import NONE
 
 _APP_STATUS_FORMAT_TEMPLATE = """Application:
@@ -71,7 +72,7 @@ def format_error_message(msg: str, header: str, width: int = 80) -> str:
     return "\n".join(lines)
 
 
-def format_replica_status(replica_status: tsm.ReplicaStatus) -> str:
+def format_replica_status(replica_status: api.ReplicaStatus) -> str:
     if replica_status.structured_error_msg != NONE:
         error_data = json.loads(replica_status.structured_error_msg)
         error_message = format_error_message(
@@ -88,8 +89,8 @@ def format_replica_status(replica_status: tsm.ReplicaStatus) -> str:
     else:
         data = f"{replica_status.state}"
         if replica_status.state in [
-            tsm.ReplicaState.CANCELLED,
-            tsm.ReplicaState.FAILED,
+            api.ReplicaState.CANCELLED,
+            api.ReplicaState.FAILED,
         ]:
             data += " (no reply file)"
 
@@ -102,7 +103,7 @@ def format_replica_status(replica_status: tsm.ReplicaStatus) -> str:
 
 
 def format_role_status(
-    role_status: tsm.RoleStatus,
+    role_status: api.RoleStatus,
 ) -> str:
     replica_data = ""
 
@@ -112,15 +113,15 @@ def format_role_status(
 
 
 def get_roles(
-    roles: List[tsm.RoleStatus], filter_roles: Optional[List[str]] = None
-) -> List[tsm.RoleStatus]:
+    roles: List[api.RoleStatus], filter_roles: Optional[List[str]] = None
+) -> List[api.RoleStatus]:
     if not filter_roles:
         return roles
     return [role_status for role_status in roles if role_status.role in filter_roles]
 
 
 def format_app_status(
-    app_status: tsm.AppStatus,
+    app_status: api.AppStatus,
     filter_roles: Optional[List[str]] = None,
 ) -> str:
     roles_data = ""
@@ -153,9 +154,9 @@ def add_arguments(self, subparser: argparse.ArgumentParser) -> None:
 
     def run(self, args: argparse.Namespace) -> None:
         app_handle = args.app_handle
-        scheduler, session_name, app_id = tsm.parse_app_handle(app_handle)
-        session = tsm.session(name=session_name)
-        app_status = session.status(app_handle)
+        scheduler, session_name, app_id = api.parse_app_handle(app_handle)
+        runner = get_runner(name=session_name)
+        app_status = runner.status(app_handle)
         filter_roles = parse_list_arg(args.roles)
         if app_status:
             print(format_app_status(app_status, filter_roles))

torchx/cli/main.py

@@ -48,7 +48,7 @@ def create_parser() -> ArgumentParser:
     return parser
 
 
-def main(argv: List[str]) -> None:
+def main(argv: List[str] = sys.argv[1:]) -> None:
     parser = create_parser()
     args = parser.parse_args(argv)
     if "func" not in args:
@@ -58,4 +58,4 @@ def main(argv: List[str]) -> None:
 
 
 if __name__ == "__main__":
-    main(sys.argv[1:])
+    main()

torchx/cli/test/__init__.py

@@ -13,7 +13,7 @@
 from torchx.specs.api import Application, Container, ElasticRole, Resource
 
 
-class CmdStatusTest(unittest.TestCase):
+class CmdDescribeTest(unittest.TestCase):
     def get_test_app(self) -> Application:
         resource = Resource(cpu=2, gpu=0, memMB=256)
         trainer = (

torchx/cli/test/cmd_describe_test.py

@@ -6,6 +6,8 @@
 # LICENSE file in the root directory of this source tree.
 
 import argparse
+import os
+import time
 import unittest
 from unittest.mock import patch
 
@@ -71,6 +73,9 @@ def _get_test_app_status(self) -> AppStatus:
         return AppStatus(state=AppState.RUNNING, roles=[role_status])
 
     def test_format_app_status(self) -> None:
+        os.environ["TZ"] = "Europe/London"
+        time.tzset()
+
         app_status = self._get_test_app_status()
         actual_message = format_app_status(app_status)
         print(actual_message)
@@ -79,7 +84,7 @@ def test_format_app_status(self) -> None:
   Num Restarts: 0
 Roles:
  *worker[0]:FAILED (exitcode: -1)
-    timestamp: 1970-01-15 15:13:02
+    timestamp: 1970-01-16 00:13:02
     hostname: localhost
     error_msg: error
   worker[1]:RUNNING"""

torchx/cli/test/cmd_status_test.py

@@ -8,9 +8,9 @@
 import unittest
 from pathlib import Path
 
-import torchelastic.tsm.driver as tsm
 from torchx.cli.cmd_run import _parse_run_config
 from torchx.cli.main import main
+from torchx.specs import api
 
 
 _root: Path = Path(__file__).parent
@@ -48,12 +48,12 @@ def test_run_builtin_config(self) -> None:
         )
 
     def test_run_scheduler_args_empty(self) -> None:
-        self.assertEqual(_parse_run_config(""), tsm.RunConfig())
+        self.assertEqual(_parse_run_config(""), api.RunConfig())
 
     def test_run_scheduler_args_simple(self) -> None:
         self.assertEqual(
             _parse_run_config("a=1,b=2;3;4"),
-            tsm.RunConfig(
+            api.RunConfig(
                 cfgs={
                     "a": "1",
                     "b": ["2", "3", "4"],