Make --log wait automatically in run (#764)

* Wait for app to start with torchx log

Summary:
Right now it silently times out after waiting for 10s and then (depending on the scheduler) fails because the job hasn't actually started; I'm trying to use `--log` with the `run` command and inevitably run into this.

I wasn't sure if we want to explicitly timeout, or what the best ergonomics would be -- I can increase the timeout to 10 minutes and blow up if it doesn't start in that time, or change it to a while loop (as done here) -- wanted to check :).

Differential Revision: D48840617

fbshipit-source-id: 111ec6708e18d83f70a1d946429226f6ab025082

* Make --log wait automatically in run (#764)

Summary:
Pull Request resolved: #764

I generally used to find it confusing and assume that `--log` was broken before realizing I always need to specify `--wait` to make it work -- I think it would be better to have `--log` automatically wait for the job to finish.

Differential Revision: D48874784

fbshipit-source-id: 728325213d695b5077b8b4139e42b763875d8a91

Author: kunalb-meta

torchx/cli/cmd_log.py

@@ -101,12 +101,14 @@ def get_logs(
     if len(path) == 4:
         replica_ids = [(role_name, int(id)) for id in path[3].split(",") if id]
     else:
-        for i in range(10):
+        display_waiting = True
+        while True:
             status = runner.status(app_handle)
             if status and is_started(status.state):
                 break
-            if i == 0:
-                logger.info("Waiting for app to start before logging...")
+            elif display_waiting:
+                logger.info("Waiting for app state response before fetching logs...")
+                display_waiting = False
             time.sleep(1)
 
         app = none_throws(runner.describe(app_handle))

torchx/cli/cmd_run.py

@@ -223,7 +223,7 @@ def _run(self, runner: Runner, args: argparse.Namespace) -> None:
                     app_status = runner.status(app_handle)
                     if app_status:
                         logger.info(app_status.format())
-                    if args.wait:
+                    if args.wait or args.log:
                         self._wait_and_exit(runner, app_handle, log=args.log)
 
         except (ComponentValidationException, ComponentNotFoundException) as e:

torchx/cli/test/cmd_log_test.py

@@ -69,6 +69,18 @@ def log_lines(
         return iter([line for line in log_lines if re.match(regex, line)])
 
 
+class MockPendingRunner(MockRunner):
+    def __init__(self) -> None:
+        self._status_checks = 0
+
+    def status(self, app_handle: str) -> AppStatus:
+        if self._status_checks < 2:
+            self._status_checks += 1
+            return AppStatus(state=AppState.PENDING)
+        else:
+            return super().status(app_handle)
+
+
 class CmdLogTest(unittest.TestCase):
     @patch("sys.exit", side_effect=SentinelError)
     def test_cmd_log_bad_job_identifier(self, exit_mock: MagicMock) -> None:
@@ -164,6 +176,15 @@ def test_cmd_log_some_replicas(self, mock_runner: MagicMock) -> None:
             set(out.getvalue().split("\n")),
         )
 
+    @patch(RUNNER, new_callable=MockPendingRunner)
+    def test_cmd_log_wait_to_start(self, mock_runner: MagicMock) -> None:
+        out = io.StringIO()
+        with self.assertLogs(level="INFO") as cm:
+            get_logs(
+                out, "local_docker://test-session/SparseNNAppdef/trainer", regex=None
+            )
+        self.assertIn("Waiting", "\n".join(cm.output))
+
     @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