Support flow orchestrator launches in JSON mode (#1108)

Summary:

`OrchestratorRunner` was not being correctly initialized for JSON cases since the JSON input was parsed after Runner initialized.

This diff moves some of the logic around to make sure:
- in FB version, the correct scheduler is logged
- JSON parsing isn't repeated unnecessarily
- if JSON mode is used, the scheduler is parsed before getting the runner to account for `single_orchestrator_runner` scheduler usage

Differential Revision: D81705462

Author: ishachirimar

torchx/cli/cmd_run.py

@@ -206,6 +206,7 @@ def run(self, args: argparse.Namespace) -> None:
 class CmdRun(SubCommand):
     def __init__(self) -> None:
         self._subparser: Optional[argparse.ArgumentParser] = None
+        self.stdin_data_json: Dict[str, Any] = {}
 
     def add_arguments(self, subparser: argparse.ArgumentParser) -> None:
         scheduler_names = get_scheduler_factories().keys()
@@ -419,18 +420,21 @@ def verify_no_extra_args(self, args: argparse.Namespace) -> None:
             )
 
     def _run(self, runner: Runner, args: argparse.Namespace) -> None:
-        # Verify no conflicting arguments when using to loop over the stdin
-        self.verify_no_extra_args(args)
         if args.stdin:
-            stdin_data_json = self.torchx_json_from_stdin()
-            self._run_from_stdin_args(runner, stdin_data_json)
+            self._run_from_stdin_args(runner, self.stdin_data_json)
         else:
             self._run_from_cli_args(runner, args)
 
+    def _validate_and_save_torchx_stdin_args(self, args: argparse.Namespace) -> None:
+        self.verify_no_extra_args(args)
+        if args.stdin:
+            self.stdin_data_json = self.torchx_json_from_stdin()
+
     def run(self, args: argparse.Namespace) -> None:
         os.environ["TORCHX_CONTEXT_NAME"] = os.getenv("TORCHX_CONTEXT_NAME", "cli_run")
         component_defaults = load_sections(prefix="component")
 
+        self._validate_and_save_torchx_stdin_args(args)
         with get_runner(component_defaults=component_defaults) as runner:
             self._run(runner, args)