Merge branch 'main' into patch-1

Author: omrishiv

torchx/cli/cmd_status.py

@@ -8,6 +8,7 @@
 # pyre-strict
 
 import argparse
+import json
 import logging
 import sys
 from typing import List, Optional
@@ -46,6 +47,11 @@ def add_arguments(self, subparser: argparse.ArgumentParser) -> None:
         subparser.add_argument(
             "--roles", type=str, default="", help="comma separated roles to filter"
         )
+        subparser.add_argument(
+            "--json",
+            action="store_true",
+            help="output the status in JSON format",
+        )
 
     def run(self, args: argparse.Namespace) -> None:
         app_handle = args.app_handle
@@ -54,7 +60,10 @@ def run(self, args: argparse.Namespace) -> None:
         app_status = runner.status(app_handle)
         filter_roles = parse_list_arg(args.roles)
         if app_status:
-            print(app_status.format(filter_roles))
+            if args.json:
+                print(json.dumps(app_status.to_json(filter_roles)))
+            else:
+                print(app_status.format(filter_roles))
         else:
             logger.error(
                 f"AppDef: {app_id},"

torchx/cli/test/cmd_run_test.py

@@ -21,7 +21,6 @@
 from unittest.mock import MagicMock, patch
 
 from torchx.cli.argparse_util import ArgOnceAction, torchxconfig
-
 from torchx.cli.cmd_run import _parse_component_name_and_args, CmdBuiltins, CmdRun
 from torchx.schedulers.local_scheduler import SignalException
 
@@ -216,38 +215,41 @@ def test_store_experiment_id(self, mock_runner_run: MagicMock) -> None:
         self.assertEqual(call_kwargs["parent_run_id"], "experiment_1")
 
     def test_parse_component_name_and_args_no_default(self) -> None:
+        # set dirs to test tmpdir so tests don't accidentally pick up user's $HOME/.torchxconfig
+        dirs = [str(self.tmpdir)]
+
         sp = argparse.ArgumentParser(prog="test")
         self.assertEqual(
             ("utils.echo", []),
-            _parse_component_name_and_args(["utils.echo"], sp),
+            _parse_component_name_and_args(["utils.echo"], sp, dirs),
         )
         self.assertEqual(
             ("utils.echo", []),
-            _parse_component_name_and_args(["--", "utils.echo"], sp),
+            _parse_component_name_and_args(["--", "utils.echo"], sp, dirs),
         )
         self.assertEqual(
             ("utils.echo", ["--msg", "hello"]),
-            _parse_component_name_and_args(["utils.echo", "--msg", "hello"], sp),
+            _parse_component_name_and_args(["utils.echo", "--msg", "hello"], sp, dirs),
         )
 
         self.assertEqual(
             ("utils.echo", ["--msg", "hello", "--", "--"]),
             _parse_component_name_and_args(
-                ["utils.echo", "--msg", "hello", "--", "--"], sp
+                ["utils.echo", "--msg", "hello", "--", "--"], sp, dirs
             ),
         )
 
         self.assertEqual(
             ("utils.echo", ["--msg", "hello", "-", "-"]),
             _parse_component_name_and_args(
-                ["utils.echo", "--msg", "hello", "-", "-"], sp
+                ["utils.echo", "--msg", "hello", "-", "-"], sp, dirs
             ),
         )
 
         self.assertEqual(
             ("utils.echo", ["--msg", "hello", "-  ", "-  "]),
             _parse_component_name_and_args(
-                ["utils.echo", "--msg", "hello", "-  ", "-  "], sp
+                ["utils.echo", "--msg", "hello", "-  ", "-  "], sp, dirs
             ),
         )
 
@@ -274,32 +276,35 @@ def test_parse_component_name_and_args_no_default(self) -> None:
                     "-m",
                 ],
                 sp,
+                dirs,
             ),
         )
 
         with self.assertRaises(SystemExit):
-            _parse_component_name_and_args(["--"], sp)
+            _parse_component_name_and_args(["--"], sp, dirs)
 
         with self.assertRaises(SystemExit):
-            _parse_component_name_and_args(["--msg", "hello"], sp)
+            _parse_component_name_and_args(["--msg", "hello"], sp, dirs)
 
         with self.assertRaises(SystemExit):
-            _parse_component_name_and_args(["-m", "hello"], sp)
+            _parse_component_name_and_args(["-m", "hello"], sp, dirs)
 
         with self.assertRaises(SystemExit):
-            _parse_component_name_and_args(["-m", "hello", "-m", "repeate"], sp)
+            _parse_component_name_and_args(["-m", "hello", "-m", "repeate"], sp, dirs)
 
         with self.assertRaises(SystemExit):
-            _parse_component_name_and_args(["--msg", "hello", "--msg", "repeate"], sp)
+            _parse_component_name_and_args(
+                ["--msg", "hello", "--msg", "repeate"], sp, dirs
+            )
 
         with self.assertRaises(SystemExit):
             _parse_component_name_and_args(
-                ["--msg  ", "hello", "--msg     ", "repeate"], sp
+                ["--msg  ", "hello", "--msg     ", "repeate"], sp, dirs
             )
 
         with self.assertRaises(SystemExit):
             _parse_component_name_and_args(
-                ["--m", "hello", "--", "--msg", "msg", "--msg", "repeate"], sp
+                ["--m", "hello", "--", "--msg", "msg", "--msg", "repeate"], sp, dirs
             )
 
     def test_parse_component_name_and_args_with_default(self) -> None:

torchx/specs/api.py

@@ -538,6 +538,15 @@ class RoleStatus:
     role: str
     replicas: List[ReplicaStatus]
 
+    def to_json(self) -> Dict[str, Any]:
+        """
+        Convert the RoleStatus to a json object.
+        """
+        return {
+            "role": self.role,
+            "replicas": [asdict(replica) for replica in self.replicas],
+        }
+
 
 @dataclass
 class AppStatus:
@@ -657,6 +666,21 @@ def _format_role_status(
             replica_data += self._format_replica_status(replica)
         return f"{replica_data}"
 
+    def to_json(self, filter_roles: Optional[List[str]] = None) -> Dict[str, Any]:
+        """
+        Convert the AppStatus to a json object, including RoleStatus.
+        """
+        roles = self._get_role_statuses(self.roles, filter_roles)
+
+        return {
+            "state": str(self.state),
+            "num_restarts": self.num_restarts,
+            "roles": [role_status.to_json() for role_status in roles],
+            "msg": self.msg,
+            "structured_error_msg": self.structured_error_msg,
+            "url": self.ui_url,
+        }
+
     def format(
         self,
         filter_roles: Optional[List[str]] = None,
@@ -672,6 +696,7 @@ def format(
         """
         roles_data = ""
         roles = self._get_role_statuses(self.roles, filter_roles)
+
         for role_status in roles:
             roles_data += self._format_role_status(role_status)
         return Template(_APP_STATUS_FORMAT_TEMPLATE).substitute(

torchx/specs/test/api_test.py

@@ -176,6 +176,42 @@ def test_format_app_status(self) -> None:
         # Split and compare to aviod AssertionError.
         self.assertEqual(expected_message.split(), actual_message.split())
 
+    def test_app_status_in_json(self) -> None:
+        app_status = self._get_test_app_status()
+        result = app_status.to_json()
+        error_msg = '{"message":{"message":"error","errorCode":-1,"extraInfo":{"timestamp":1293182}}}'
+        self.assertDictEqual(
+            result,
+            {
+                "state": "RUNNING",
+                "num_restarts": 0,
+                "roles": [
+                    {
+                        "role": "worker",
+                        "replicas": [
+                            {
+                                "id": 0,
+                                "state": 5,
+                                "role": "worker",
+                                "hostname": "localhost",
+                                "structured_error_msg": error_msg,
+                            },
+                            {
+                                "id": 1,
+                                "state": 3,
+                                "role": "worker",
+                                "hostname": "localhost",
+                                "structured_error_msg": "<NONE>",
+                            },
+                        ],
+                    }
+                ],
+                "msg": "",
+                "structured_error_msg": "<NONE>",
+                "url": None,
+            },
+        )
+
 
 class ResourceTest(unittest.TestCase):
     def test_copy_resource(self) -> None:

torchx/util/modules.py

@@ -8,16 +8,13 @@
 
 import importlib
 from types import ModuleType
-from typing import Callable, Optional, Union
+from typing import Callable, Optional, TypeVar, Union
 
 
 def load_module(path: str) -> Union[ModuleType, Optional[Callable[..., object]]]:
     """
     Loads and returns the module/module attr represented by the ``path``: ``full.module.path:optional_attr``
 
-    ::
-
-
     1. ``load_module("this.is.a_module:fn")`` -> equivalent to ``this.is.a_module.fn``
     1. ``load_module("this.is.a_module")`` -> equivalent to ``this.is.a_module``
     """
@@ -33,3 +30,36 @@ def load_module(path: str) -> Union[ModuleType, Optional[Callable[..., object]]]
         return getattr(module, method) if method else module
     except Exception:
         return None
+
+
+T = TypeVar("T")
+
+
+def import_attr(name: str, attr: str, default: T) -> T:
+    """
+    Imports ``name.attr`` and returns it if the module is found.
+    Otherwise, returns the specified ``default``.
+    Useful when getting an attribute from an optional dependency.
+
+    Note that the ``default`` parameter is intentionally not an optional
+    since this function is intended to be used with modules that may not be
+    installed as a dependency. Therefore the caller must ALWAYS provide a
+    sensible default.
+
+    Usage:
+
+    .. code-block:: python
+
+        aws_resources = import_attr("torchx.specs.named_resources_aws", "NAMED_RESOURCES", default={})
+        all_resources.update(aws_resources)
+
+    Raises:
+        AttributeError: If the module exists (e.g. can be imported)
+            but does not have an attribute with name ``attr``.
+    """
+    try:
+        mod = importlib.import_module(name)
+    except ModuleNotFoundError:
+        return default
+    else:
+        return getattr(mod, attr)

torchx/util/test/modules_test.py

@@ -4,9 +4,11 @@
 # This source code is licensed under the BSD-style license found in the
 # LICENSE file in the root directory of this source tree.
 
+# pyre-strict
+
 import unittest
 
-from torchx.util.modules import load_module
+from torchx.util.modules import import_attr, load_module
 
 
 class ModulesTest(unittest.TestCase):
@@ -21,3 +23,23 @@ def test_load_module_method(self) -> None:
         import os
 
         self.assertEqual(result, os.path.join)
+
+    def test_try_import(self) -> None:
+        def _join(_0: str, *_1: str) -> str:
+            return ""  # should never be called
+
+        os_path_join = import_attr("os.path", "join", default=_join)
+        import os
+
+        self.assertEqual(os.path.join, os_path_join)
+
+    def test_try_import_non_existent_module(self) -> None:
+        should_default = import_attr("non.existent", "foo", default="bar")
+        self.assertEqual("bar", should_default)
+
+    def test_try_import_non_existent_attr(self) -> None:
+        def _join(_0: str, *_1: str) -> str:
+            return ""  # should never be called
+
+        with self.assertRaises(AttributeError):
+            import_attr("os.path", "joyin", default=_join)