register fb schedulers using entry_point (pytorch#12)

Summary:
Pull Request resolved: pytorch#12

NOTE: I cleaned up a few tsm references while I was at it since removing `//torchx/schedulers/fb:registry` buck target ended up with a ton of autodeps changes that touched a lot of TARGETS files. I recommend taking a look at `torchx_fb.dist-info/entry_points.txt` and `torchx/util/entrypoint.py` for the CORE changes I originally intended for this diff.

1. Now that we have `torchx_fb.dist-info/entry_point.txt` setup to hook in plugin points for Facebook, move the scheduler registry to entrypoint so that we don't have to mess with `base_path` and have autodeps hate us.

1. cleaned up a few places were we still have a back reference to tsm.

Reviewed By: tierex

Differential Revision: D28622469

fbshipit-source-id: 907f384d518a695bbd9087a033d771450dcb07c8

Author: Kiuk Chung

torchx/cli/cmd_log.py

@@ -14,10 +14,12 @@
 from typing import Optional
 from urllib.parse import urlparse
 
-import torchx.specs.lib as torchx
 from pyre_extensions import none_throws
+from torchx import specs
 from torchx.cli.cmd_base import SubCommand
-from torchx.runner import Runner
+from torchx.runner import Runner, get_runner
+from torchx.specs.api import make_app_handle
+
 
 GREEN = "\033[32m"
 ENDC = "\033[0m"
@@ -64,10 +66,10 @@ def get_logs(identifier: str, regex: Optional[str], should_tail: bool = False) -
     app_id = path[1]
     role_name = path[2]
 
-    session_ = torchx.run(name=session_name)
-    app_handle = torchx.make_app_handle(scheduler_backend, session_name, app_id)
+    runner = get_runner(name=session_name)
+    app_handle = make_app_handle(scheduler_backend, session_name, app_id)
 
-    app = none_throws(session_.describe(app_handle))
+    app = none_throws(runner.describe(app_handle))
 
     if len(path) == 4:
         replica_ids = [int(id) for id in path[3].split(",") if id]
@@ -98,7 +100,7 @@ def get_logs(identifier: str, regex: Optional[str], should_tail: bool = False) -
         thread = threading.Thread(
             target=print_log_lines,
             args=(
-                session_,
+                runner,
                 app_handle,
                 role_name,
                 replica_id,
@@ -126,7 +128,7 @@ def get_logs(identifier: str, regex: Optional[str], should_tail: bool = False) -
         raise threads_exceptions[0]
 
 
-def find_role_replicas(app: torchx.Application, role_name: str) -> Optional[int]:
+def find_role_replicas(app: specs.Application, role_name: str) -> Optional[int]:
     for role in app.roles:
         if role_name == role.name:
             return role.num_replicas

torchx/cli/cmd_run.py

@@ -100,18 +100,20 @@ def _parse_run_config(arg: str) -> specs.RunConfig:
 _CONFIG_EXT = ".torchx"
 
 
+def get_abspath(relpath: str) -> str:
+    module = __name__.replace(".", path.sep)  # torchx/cli/cmd_run
+    module_path, _ = path.splitext(__file__)  # $root/torchx/cli/cmd_run
+    root = module_path.replace(module, "")
+    return path.join(root, relpath)
+
+
 def get_file_contents(conf_file: str) -> Optional[str]:
     """
     Reads the ``conf_file`` relative to the root of the project.
     Returns ``None`` if ``$root/$conf_file`` does not exist.
     Example: ``get_file("torchx/cli/config/foo.txt")``
     """
-
-    module = __name__.replace(".", path.sep)  # torchx/cli/cmd_run
-    module_path, _ = path.splitext(__file__)  # $root/torchx/cli/cmd_run
-    root = module_path.replace(module, "")
-    abspath = path.join(root, conf_file)
-
+    abspath = get_abspath(conf_file)
     if path.exists(abspath):
         with open(abspath, "r") as f:
             return f.read()
@@ -148,8 +150,12 @@ def read_conf_file(conf_file: str) -> str:
 
 
 def _builtins() -> List[str]:
+    config_dir = entrypoints.load("torchx.file", "get_dir_path", default=get_abspath)(
+        _CONFIG_DIR
+    )
+
     builtins: List[str] = []
-    for f in os.listdir(_CONFIG_DIR):
+    for f in os.listdir(config_dir):
         _, extension = os.path.splitext(f)
         if f.endswith(_CONFIG_EXT):
             builtins.append(f)

torchx/cli/test/cmd_log_test.py

@@ -7,8 +7,8 @@
 
 import io
 import unittest
-from typing import Optional, Iterator
-from unittest.mock import patch, MagicMock
+from typing import Iterator, Optional
+from unittest.mock import MagicMock, patch
 
 from torchx.cli.cmd_log import ENDC, GREEN, get_logs
 from torchx.specs.api import Application, Role, parse_app_handle
@@ -22,11 +22,11 @@ class SentinelError(Exception):
     pass
 
 
-SESSION = "torchx.specs.lib.run"
+RUNNER = "torchx.cli.cmd_log.get_runner"
 
 
-class MockSession:
-    def __call__(self, name: Optional[str] = None) -> "MockSession":
+class MockRunner:
+    def __call__(self, name: Optional[str] = None) -> "MockRunner":
         return self
 
     def describe(self, app_handle: str) -> Application:
@@ -61,10 +61,10 @@ def test_cmd_log_bad_job_identifier(self, exit_mock: MagicMock) -> None:
             get_logs("local:///SparseNNApplication/", "QPS.*")
         exit_mock.assert_called_once_with(1)
 
-    @patch(SESSION, new_callable=MockSession)
+    @patch(RUNNER, new_callable=MockRunner)
     @patch("sys.exit", side_effect=SentinelError)
     def test_cmd_log_unknown_role(
-        self, exit_mock: MagicMock, session_mock: MagicMock
+        self, exit_mock: MagicMock, mock_runner: MagicMock
     ) -> None:
         with self.assertRaises(SentinelError):
             get_logs(
@@ -74,10 +74,10 @@ def test_cmd_log_unknown_role(
 
         exit_mock.assert_called_once_with(1)
 
-    @patch(SESSION, new_callable=MockSession)
+    @patch(RUNNER, new_callable=MockRunner)
     @patch("sys.stdout", new_callable=io.StringIO)
     def test_cmd_log_all_replicas(
-        self, stdout_mock: MagicMock, session_mock: MagicMock
+        self, stdout_mock: MagicMock, mock_runner: MagicMock
     ) -> None:
         get_logs("local://test-session/SparseNNApplication/trainer", regex="INFO")
         self.assertSetEqual(
@@ -92,10 +92,10 @@ def test_cmd_log_all_replicas(
             set(stdout_mock.getvalue().split("\n")),
         )
 
-    @patch(SESSION, new_callable=MockSession)
+    @patch(RUNNER, new_callable=MockRunner)
     @patch("sys.stdout", new_callable=io.StringIO)
     def test_cmd_log_one_replica(
-        self, stdout_mock: MagicMock, session_mock: MagicMock
+        self, stdout_mock: MagicMock, mock_runner: MagicMock
     ) -> None:
         get_logs("local://test-session/SparseNNApplication/trainer/0", regex=None)
         self.assertSetEqual(
@@ -110,10 +110,10 @@ def test_cmd_log_one_replica(
             set(stdout_mock.getvalue().split("\n")),
         )
 
-    @patch(SESSION, new_callable=MockSession)
+    @patch(RUNNER, new_callable=MockRunner)
     @patch("sys.stdout", new_callable=io.StringIO)
     def test_cmd_log_some_replicas(
-        self, stdout_mock: MagicMock, session_mock: MagicMock
+        self, stdout_mock: MagicMock, mock_runner: MagicMock
     ) -> None:
         get_logs("local://test-session/SparseNNApplication/trainer/0,2", regex="WARN")
         self.assertSetEqual(
@@ -127,11 +127,11 @@ def test_cmd_log_some_replicas(
             set(stdout_mock.getvalue().split("\n")),
         )
 
-    @patch(SESSION, new_callable=MockSession)
-    def test_print_log_lines_throws(self, session_mock: MagicMock) -> None:
+    @patch(RUNNER, new_callable=MockRunner)
+    def test_print_log_lines_throws(self, mock_runner: MagicMock) -> None:
         # makes sure that when the function executed in the threadpool
         # errors out; we raise the exception all the way through
-        with patch.object(session_mock, "log_lines") as log_lines_mock:
+        with patch.object(mock_runner, "log_lines") as log_lines_mock:
             log_lines_mock.side_effect = RuntimeError
             with self.assertRaises(RuntimeError):
                 get_logs(

torchx/cli/test/cmd_runopts_test.py

@@ -9,9 +9,7 @@
 import unittest
 
 from torchx.cli.cmd_runopts import CmdRunopts
-
-# @manual=//torchx/schedulers/fb:registry
-from torchx.schedulers.registry import get_schedulers
+from torchx.schedulers import get_schedulers
 
 
 class CmdRunOptsTest(unittest.TestCase):
@@ -22,7 +20,7 @@ def test_run(self) -> None:
         cmd_runopts = CmdRunopts()
         cmd_runopts.add_arguments(parser)
 
-        schedulers = get_schedulers("test").keys()
+        schedulers = get_schedulers(session_name="test").keys()
         test_configs = [[]] + [[scheduler] for scheduler in schedulers]
         for scheduler in test_configs:
             with self.subTest(scheduler=scheduler):

torchx/runner/api.py

@@ -13,8 +13,8 @@
 
 from pyre_extensions import none_throws
 from torchx.runner.events import log_event
+from torchx.schedulers import get_schedulers
 from torchx.schedulers.api import Scheduler
-from torchx.schedulers.registry import get_schedulers
 from torchx.specs.api import (
     NULL_CONTAINER,
     AppDryRunInfo,
@@ -426,6 +426,5 @@ def get_runner(name: Optional[str] = None, **scheduler_params: Any) -> Runner:
     if not name:
         name = f"torchx_{getpass.getuser()}"
 
-    scheduler_params["session_name"] = name
-    schedulers = get_schedulers(**scheduler_params)
+    schedulers = get_schedulers(session_name=name, **scheduler_params)
     return Runner(name, schedulers)

torchx/schedulers/__init__.py

@@ -4,3 +4,31 @@
 #
 # This source code is licensed under the BSD-style license found in the
 # LICENSE file in the root directory of this source tree.
+
+from typing import Dict
+
+import torchx.schedulers.local_scheduler as local_scheduler
+from torchx.schedulers.api import Scheduler
+from torchx.specs.api import SchedulerBackend
+from torchx.util.entrypoints import load_group
+
+
+def get_schedulers(
+    session_name: str,
+    # pyre-ignore[2]
+    **scheduler_params
+) -> Dict[SchedulerBackend, Scheduler]:
+
+    schedulers = load_group(
+        "torchx.schedulers",
+        default={
+            "local": local_scheduler.create_scheduler,
+            "default": local_scheduler.create_scheduler,
+        },
+        ignore_missing=True,
+    )
+
+    return {
+        scheduler_backend: scheduler_factory_method(session_name, **scheduler_params)
+        for scheduler_backend, scheduler_factory_method in schedulers.items()
+    }

torchx/schedulers/registry.py

@@ -6,13 +6,26 @@
 # LICENSE file in the root directory of this source tree.
 
 import unittest
+from typing import Any, Dict, Optional
+from unittest.mock import MagicMock, patch
 
+from torchx.schedulers import get_schedulers
 from torchx.schedulers.local_scheduler import LocalScheduler
-from torchx.schedulers.registry import get_schedulers
+
+
+class spy_load_group:
+    def __call__(
+        self,
+        group: str,
+        default: Dict[str, Any],
+        ignore_missing: Optional[bool] = False,
+    ) -> Dict[str, Any]:
+        return default
 
 
 class SchedulersTest(unittest.TestCase):
-    def test_get_local_schedulers(self) -> None:
+    @patch("torchx.schedulers.load_group", new_callable=spy_load_group)
+    def test_get_local_schedulers(self, mock_load_group: MagicMock) -> None:
         schedulers = get_schedulers(session_name="test_session")
         self.assertTrue(isinstance(schedulers["local"], LocalScheduler))
         self.assertTrue(isinstance(schedulers["default"], LocalScheduler))