(torchx/tracker)(bugfix) honor tracker config in the usual .torchxconfig load hierarchy + lazy load tracker object creation in cmd_tracker (#711)

Author: kiukchung

torchx/cli/argparse_util.py

@@ -5,15 +5,11 @@
 # LICENSE file in the root directory of this source tree.
 
 from argparse import Action, ArgumentParser, Namespace
-from pathlib import Path
 from typing import Any, Dict, Optional, Sequence, Text
 
 from torchx.runner import config
 
 
-CONFIG_DIRS = [str(Path.home()), str(Path.cwd())]
-
-
 class _torchxconfig(Action):
     """
     Custom argparse action that loads default torchx CLI options
@@ -40,7 +36,6 @@ def __init__(
             config.get_configs(
                 prefix="cli",
                 name=subcmd,
-                dirs=CONFIG_DIRS,
             ),
         )
 

torchx/cli/cmd_run.py

@@ -15,7 +15,7 @@
 from typing import Dict, List, Optional, Tuple
 
 import torchx.specs as specs
-from torchx.cli.argparse_util import CONFIG_DIRS, torchxconfig_run
+from torchx.cli.argparse_util import torchxconfig_run
 from torchx.cli.cmd_base import SubCommand
 from torchx.cli.cmd_log import get_logs
 from torchx.runner import config, get_runner, Runner
@@ -182,12 +182,11 @@ def _run(self, runner: Runner, args: argparse.Namespace) -> None:
 
         scheduler_opts = runner.scheduler_run_opts(args.scheduler)
         cfg = scheduler_opts.cfg_from_str(args.scheduler_args)
-        config.apply(scheduler=args.scheduler, cfg=cfg, dirs=CONFIG_DIRS)
+        config.apply(scheduler=args.scheduler, cfg=cfg)
 
         component, component_args = _parse_component_name_and_args(
             args.component_name_and_args,
             none_throws(self._subparser),
-            dirs=CONFIG_DIRS,
         )
         try:
             if args.dryrun:
@@ -243,7 +242,8 @@ def _run(self, runner: Runner, args: argparse.Namespace) -> None:
 
     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", dirs=CONFIG_DIRS)
+        component_defaults = load_sections(prefix="component")
+
         with get_runner(component_defaults=component_defaults) as runner:
             self._run(runner, args)
 

torchx/cli/cmd_tracker.py

@@ -6,33 +6,16 @@
 
 import argparse
 import logging
-import sys
-from typing import Callable, Optional
 
 from tabulate import tabulate
 
 from torchx.cli.cmd_base import SubCommand
 from torchx.runner.api import get_configured_trackers
 from torchx.tracker.api import build_trackers, TrackerBase
-from torchx.util.types import none_throws
 
 logger: logging.Logger = logging.getLogger(__name__)
 
 
-def _requires_tracker(
-    command: Callable[["CmdTracker", argparse.Namespace], None]
-) -> Callable[["CmdTracker", argparse.Namespace], None]:
-    """Checks that command has valid tracker setup"""
-
-    def wrapper(self: "CmdTracker", args: argparse.Namespace) -> None:
-        if not self.tracker:
-            logger.error("Exiting since no trackers were configured.")
-            sys.exit(1)
-        command(self, args)
-
-    return wrapper
-
-
 class CmdTracker(SubCommand):
     """
     Prototype TorchX tracker subcommand that allows querying data by
@@ -49,30 +32,28 @@ class CmdTracker(SubCommand):
     def __init__(self) -> None:
         """
         Queries available tracker implementations and uses the first available one.
-
-        Since the instance needs to be available to setup torchx arguments, subcommands
-        utilize `_requires_tracker()` annotation to check that tracker is available
-        when invoked.
         """
-        self.tracker: Optional[TrackerBase] = None
-        configured_trackers = get_configured_trackers()
-        if configured_trackers:
-            trackers = build_trackers(configured_trackers)
-            if trackers:
-                self.tracker = next(iter(trackers))
-                logger.info(f"Using {self.tracker} to query data")
-            else:
-                logger.error("No trackers were configured!")
+
+    @property
+    def tracker(self) -> TrackerBase:
+        trackers = list(build_trackers(get_configured_trackers()))
+        if trackers:
+            logger.info(f"Using `{trackers[0]}` tracker to query data")
+            return trackers[0]
+        else:
+            raise RuntimeError(
+                "No trackers configured."
+                " See: https://pytorch.org/torchx/latest/runtime/tracking.html"
+            )
 
     def add_list_job_arguments(self, subparser: argparse.ArgumentParser) -> None:
         subparser.add_argument(
             "--parent-run-id", type=str, help="Optional job parent run ID"
         )
 
-    @_requires_tracker
     def list_jobs_command(self, args: argparse.Namespace) -> None:
         parent_run_id = args.parent_run_id
-        job_ids = none_throws(self.tracker).run_ids(parent_run_id=parent_run_id)
+        job_ids = self.tracker.run_ids(parent_run_id=parent_run_id)
 
         tabulated_job_ids = [[job_id] for job_id in job_ids]
         print(tabulate(tabulated_job_ids, headers=["JOB ID"]))
@@ -91,17 +72,15 @@ def add_job_lineage_arguments(self, subparser: argparse.ArgumentParser) -> None:
         )
         subparser.add_argument("RUN_ID", type=str, help="Job run ID")
 
-    @_requires_tracker
     def job_lineage_command(self, args: argparse.Namespace) -> None:
         raise NotImplementedError("")
 
     def add_metadata_arguments(self, subparser: argparse.ArgumentParser) -> None:
         subparser.add_argument("RUN_ID", type=str, help="Job run ID")
 
-    @_requires_tracker
     def list_metadata_command(self, args: argparse.Namespace) -> None:
         run_id = args.RUN_ID
-        metadata = none_throws(self.tracker).metadata(run_id)
+        metadata = self.tracker.metadata(run_id)
         print_data = [[k, v] for k, v in metadata.items()]
 
         print(tabulate(print_data, headers=["ID", "VALUE"]))
@@ -113,21 +92,16 @@ def add_artifacts_arguments(self, subparser: argparse.ArgumentParser) -> None:
 
         subparser.add_argument("RUN_ID", type=str, help="Job run ID")
 
-    @_requires_tracker
     def list_artifacts_command(self, args: argparse.Namespace) -> None:
         run_id = args.RUN_ID
         artifact_filter = args.artifact
 
-        artifacts = none_throws(self.tracker).artifacts(run_id)
-        artifacts = artifacts.values()
+        artifacts = list(self.tracker.artifacts(run_id).values())
 
         if artifact_filter:
-            artifacts = [
-                artifact for artifact in artifacts if artifact.name == artifact_filter
-            ]
-        print_data = [
-            [artifact.name, artifact.path, artifact.metadata] for artifact in artifacts
-        ]
+            artifacts = [a for a in artifacts if a.name == artifact_filter]
+
+        print_data = [[a.name, a.path, a.metadata] for a in artifacts]
 
         print(tabulate(print_data, headers=["ARTIFACT", "PATH", "METADATA"]))
 

torchx/cli/main.py

@@ -5,6 +5,7 @@
 # LICENSE file in the root directory of this source tree.
 
 import logging
+import os
 import sys
 from argparse import ArgumentParser
 from typing import Dict, List
@@ -75,9 +76,9 @@ def create_parser(subcmds: Dict[str, SubCommand]) -> ArgumentParser:
     parser = ArgumentParser(description="torchx CLI")
     parser.add_argument(
         "--log_level",
-        type=int,
+        type=str,
         help="Python logging log level",
-        default=logging.INFO,
+        default=os.getenv("LOGLEVEL", "WARNING"),
     )
     parser.add_argument(
         "--version",

torchx/cli/test/argparse_util_test.py

@@ -4,38 +4,24 @@
 # This source code is licensed under the BSD-style license found in the
 # LICENSE file in the root directory of this source tree.
 
-import shutil
-import tempfile
-import unittest
 from argparse import ArgumentParser
-from pathlib import Path
 from unittest import mock
 
 from torchx.cli import argparse_util
 from torchx.cli.argparse_util import torchxconfig_run
+from torchx.test.fixtures import TestWithTmpDir
 
+DEFAULT_CONFIG_DIRS = "torchx.runner.config.DEFAULT_CONFIG_DIRS"
 
-CONFIG_DIRS = "torchx.cli.argparse_util.CONFIG_DIRS"
-
-
-class ArgparseUtilTest(unittest.TestCase):
-    def _write(self, filename: str, content: str) -> Path:
-        f = Path(self.test_dir) / filename
-        f.parent.mkdir(parents=True, exist_ok=True)
-        with open(f, "w") as fp:
-            fp.write(content)
-        return f
 
+class ArgparseUtilTest(TestWithTmpDir):
     def setUp(self) -> None:
-        self.test_dir = tempfile.mkdtemp(prefix="torchx_argparse_util_test")
+        super().setUp()
         argparse_util._torchxconfig._subcmd_configs.clear()
 
-    def tearDown(self) -> None:
-        shutil.rmtree(self.test_dir)
-
     def test_torchxconfig_action(self) -> None:
-        with mock.patch(CONFIG_DIRS, [self.test_dir]):
-            self._write(
+        with mock.patch(DEFAULT_CONFIG_DIRS, [str(self.tmpdir)]):
+            self.write(
                 ".torchxconfig",
                 """
 [cli:run]
@@ -64,8 +50,8 @@ def test_torchxconfig_action(self) -> None:
             self.assertEqual("baz", args.workspace)
 
     def test_torchxconfig_action_argparse_default(self) -> None:
-        with mock.patch(CONFIG_DIRS, [self.test_dir]):
-            self._write(
+        with mock.patch(DEFAULT_CONFIG_DIRS, [str(self.tmpdir)]):
+            self.write(
                 ".torchxconfig",
                 """
 [cli:run]
@@ -89,8 +75,8 @@ def test_torchxconfig_action_argparse_default(self) -> None:
             self.assertEqual("foo", args.workspace)
 
     def test_torchxconfig_action_required(self) -> None:
-        with mock.patch(CONFIG_DIRS, [self.test_dir]):
-            self._write(
+        with mock.patch(DEFAULT_CONFIG_DIRS, [str(self.tmpdir)]):
+            self.write(
                 ".torchxconfig",
                 """
 [cli:run]
@@ -120,8 +106,8 @@ def test_torchxconfig_action_required(self) -> None:
 
     def test_torchxconfig_action_aliases(self) -> None:
         # for aliases, the config file needs to declare the original arg
-        with mock.patch(CONFIG_DIRS, [self.test_dir]):
-            self._write(
+        with mock.patch(DEFAULT_CONFIG_DIRS, [str(self.tmpdir)]):
+            self.write(
                 ".torchxconfig",
                 """
 [cli:run]

torchx/cli/test/cmd_run_test.py

@@ -134,7 +134,9 @@ def test_run_missing(self) -> None:
     def test_conf_file_missing(self) -> None:
         # guards against existing .torchxconfig files
         # in user's $HOME or the CWD where the test is launched from
-        with patch("torchx.cli.cmd_run.CONFIG_DIRS", return_value=[self.tmpdir]):
+        with patch(
+            "torchx.runner.config.DEFAULT_CONFIG_DIRS", return_value=[self.tmpdir]
+        ):
             with self.assertRaises(SystemExit):
                 args = self.parser.parse_args(
                     [