xax 0.0.1__py3-none-any.whl → 0.0.5__py3-none-any.whl

xax/task/base.py

@@ -0,0 +1,207 @@
+"""Defines the base task interface.
+
+This interface is built upon by a large number of other interfaces which
+compose various functionality into a single cohesive unit. The base task
+just stores the configuration and provides hooks which are overridden by
+upstream classes.
+"""
+
+import functools
+import inspect
+import logging
+import sys
+from dataclasses import dataclass, is_dataclass
+from pathlib import Path
+from types import TracebackType
+from typing import Generic, Self, TypeVar, cast
+
+from omegaconf import Container, DictConfig, OmegaConf
+
+from xax.core.state import State
+from xax.utils.text import camelcase_to_snakecase
+
+logger = logging.getLogger(__name__)
+
+
+@dataclass
+class BaseConfig:
+    pass
+
+
+Config = TypeVar("Config", bound=BaseConfig)
+
+RawConfigType = BaseConfig | dict | DictConfig | str | Path
+
+
+def _load_as_dict(path: str | Path) -> DictConfig:
+    cfg = OmegaConf.load(path)
+    if not isinstance(cfg, DictConfig):
+        raise TypeError(f"Config file at {path} must be a dictionary, not {type(cfg)}!")
+    return cfg
+
+
+def get_config(cfg: RawConfigType, task_path: Path) -> DictConfig:
+    if isinstance(cfg, (str, Path)):
+        cfg = Path(cfg)
+        if cfg.exists():
+            cfg = _load_as_dict(cfg)
+        elif task_path is not None and len(cfg.parts) == 1 and (other_cfg_path := task_path.parent / cfg).exists():
+            cfg = _load_as_dict(other_cfg_path)
+        else:
+            raise FileNotFoundError(f"Could not find config file at {cfg}!")
+    elif isinstance(cfg, dict):
+        cfg = OmegaConf.create(cfg)
+    elif is_dataclass(cfg):
+        cfg = OmegaConf.structured(cfg)
+    return cast(DictConfig, cfg)
+
+
+class BaseTask(Generic[Config]):
+    config: Config
+
+    def __init__(self, config: Config) -> None:
+        super().__init__()
+
+        self.config = config
+
+        if isinstance(self.config, Container):
+            OmegaConf.resolve(self.config)
+
+    def on_step_start(self, state: State) -> State:
+        return state
+
+    def on_step_end(self, state: State) -> State:
+        return state
+
+    def on_training_start(self, state: State) -> State:
+        return state
+
+    def on_training_end(self, state: State) -> State:
+        return state
+
+    @functools.cached_property
+    def task_class_name(self) -> str:
+        return self.__class__.__name__
+
+    @functools.cached_property
+    def task_name(self) -> str:
+        return camelcase_to_snakecase(self.task_class_name)
+
+    @functools.cached_property
+    def task_path(self) -> Path:
+        return Path(inspect.getfile(self.__class__))
+
+    @functools.cached_property
+    def task_module(self) -> str:
+        if (mod := inspect.getmodule(self.__class__)) is None:
+            raise RuntimeError(f"Could not find module for task {self.__class__}!")
+        if (spec := mod.__spec__) is None:
+            raise RuntimeError(f"Could not find spec for module {mod}!")
+        return spec.name
+
+    @property
+    def task_key(self) -> str:
+        return f"{self.task_module}.{self.task_class_name}"
+
+    @classmethod
+    def from_task_key(cls, task_key: str) -> type[Self]:
+        task_module, task_class_name = task_key.rsplit(".", 1)
+        try:
+            mod = __import__(task_module, fromlist=[task_class_name])
+        except ImportError as e:
+            raise ImportError(f"Could not import module {task_module} for task {task_key}") from e
+        if not hasattr(mod, task_class_name):
+            raise RuntimeError(f"Could not find class {task_class_name} in module {task_module}")
+        task_class = getattr(mod, task_class_name)
+        if not issubclass(task_class, cls):
+            raise RuntimeError(f"Class {task_class_name} in module {task_module} is not a subclass of {cls}")
+        return task_class
+
+    def debug(self) -> bool:
+        return False
+
+    @property
+    def debugging(self) -> bool:
+        return self.debug()
+
+    def __enter__(self) -> Self:
+        return self
+
+    def __exit__(self, _t: type[BaseException] | None, _e: BaseException | None, _tr: TracebackType | None) -> None:
+        pass
+
+    @classmethod
+    def get_config_class(cls) -> type[Config]:
+        """Recursively retrieves the config class from the generic type.
+
+        Returns:
+            The parsed config class.
+
+        Raises:
+            ValueError: If the config class cannot be found, usually meaning
+            that the generic class has not been used correctly.
+        """
+        if hasattr(cls, "__orig_bases__"):
+            for base in cls.__orig_bases__:
+                if hasattr(base, "__args__"):
+                    for arg in base.__args__:
+                        if issubclass(arg, BaseConfig):
+                            return arg
+
+        raise ValueError(
+            "The config class could not be parsed from the generic type, which usually means that the task is not "
+            "being instantiated correctly. Your class should be defined as follows:\n\n"
+            "  class ExampleTask(mlfab.Task[Config]):\n      ...\n\nThis lets the both the task and the type "
+            "checker know what config the task is using."
+        )
+
+    @classmethod
+    def get_config(cls, *cfgs: RawConfigType, use_cli: bool | list[str] = True) -> Config:
+        """Builds the structured config from the provided config classes.
+
+        Args:
+            cfgs: The config classes to merge. If a string or Path is provided,
+                it will be loaded as a YAML file.
+            use_cli: Whether to allow additional overrides from the CLI.
+
+        Returns:
+            The merged configs.
+        """
+        task_path = Path(inspect.getfile(cls))
+        cfg = OmegaConf.structured(cls.get_config_class())
+        cfg = OmegaConf.merge(cfg, *(get_config(other_cfg, task_path) for other_cfg in cfgs))
+        if use_cli:
+            args = use_cli if isinstance(use_cli, list) else sys.argv[1:]
+            if "-h" in args or "--help" in args:
+                sys.stderr.write(OmegaConf.to_yaml(cfg))
+                sys.stderr.flush()
+                sys.exit(0)
+
+            # Attempts to load any paths as configs.
+            is_path = [Path(arg).is_file() or (task_path / arg).is_file() for arg in args]
+            paths = [arg for arg, is_path in zip(args, is_path) if is_path]
+            non_paths = [arg for arg, is_path in zip(args, is_path) if not is_path]
+            if paths:
+                cfg = OmegaConf.merge(cfg, *(get_config(path, task_path) for path in paths))
+            cfg = OmegaConf.merge(cfg, OmegaConf.from_cli(non_paths))
+
+        return cast(Config, cfg)
+
+    @classmethod
+    def config_str(cls, *cfgs: RawConfigType, use_cli: bool | list[str] = True) -> str:
+        return OmegaConf.to_yaml(cls.get_config(*cfgs, use_cli=use_cli))
+
+    @classmethod
+    def get_task(cls, *cfgs: RawConfigType, use_cli: bool | list[str] = True) -> Self:
+        """Builds the task from the provided config classes.
+
+        Args:
+            cfgs: The config classes to merge. If a string or Path is provided,
+                it will be loaded as a YAML file.
+            use_cli: Whether to allow additional overrides from the CLI.
+
+        Returns:
+            The task.
+        """
+        cfg = cls.get_config(*cfgs, use_cli=use_cli)
+        return cls(cfg)

xax/task/launchers/base.py

@@ -0,0 +1,28 @@
+"""Defines the base launcher class."""
+
+from abc import ABC, abstractmethod
+from typing import TYPE_CHECKING
+
+from xax.task.base import RawConfigType
+
+if TYPE_CHECKING:
+    from xax.task.mixins.runnable import Config, RunnableMixin
+
+
+class BaseLauncher(ABC):
+    """Defines the base launcher class."""
+
+    @abstractmethod
+    def launch(
+        self,
+        task: "type[RunnableMixin[Config]]",
+        *cfgs: RawConfigType,
+        use_cli: bool | list[str] = True,
+    ) -> None:
+        """Launches the training process.
+
+        Args:
+            task: The task class to train
+            cfgs: The raw configuration to use for training
+            use_cli: Whether to include CLI arguments in the configuration
+        """

xax/task/launchers/cli.py

@@ -0,0 +1,42 @@
+"""Defines a launcher that can be toggled from the command line."""
+
+import argparse
+import sys
+from typing import TYPE_CHECKING, Literal, get_args
+
+from xax.task.base import RawConfigType
+from xax.task.launchers.base import BaseLauncher
+from xax.task.launchers.single_process import SingleProcessLauncher
+
+if TYPE_CHECKING:
+    from xax.task.mixins.runnable import Config, RunnableMixin
+
+
+LauncherChoice = Literal["single"]
+
+
+class CliLauncher(BaseLauncher):
+    def launch(
+        self,
+        task: "type[RunnableMixin[Config]]",
+        *cfgs: RawConfigType,
+        use_cli: bool | list[str] = True,
+    ) -> None:
+        args = use_cli if isinstance(use_cli, list) else sys.argv[1:]
+        parser = argparse.ArgumentParser(add_help=False)
+        parser.add_argument(
+            "-l",
+            "--launcher",
+            choices=get_args(LauncherChoice),
+            default="single",
+            help="The launcher to use",
+        )
+        args, cli_args_rest = parser.parse_known_intermixed_args(args=args)
+        launcher_choice: LauncherChoice = args.launcher
+        use_cli_next: bool | list[str] = False if not use_cli else cli_args_rest
+
+        match launcher_choice:
+            case "single":
+                SingleProcessLauncher().launch(task, *cfgs, use_cli=use_cli_next)
+            case _:
+                raise ValueError(f"Invalid launcher choice: {launcher_choice}")

xax/task/launchers/single_process.py

@@ -0,0 +1,30 @@
+"""Defines a launcher to train a model locally, in a single process."""
+
+from typing import TYPE_CHECKING
+
+from xax.task.base import RawConfigType
+from xax.task.launchers.base import BaseLauncher
+from xax.utils.logging import configure_logging
+
+if TYPE_CHECKING:
+    from xax.task.mixins.runnable import Config, RunnableMixin
+
+
+def run_single_process_training(
+    task: "type[RunnableMixin[Config]]",
+    *cfgs: RawConfigType,
+    use_cli: bool | list[str] = True,
+) -> None:
+    configure_logging()
+    task_obj = task.get_task(*cfgs, use_cli=use_cli)
+    task_obj.run()
+
+
+class SingleProcessLauncher(BaseLauncher):
+    def launch(
+        self,
+        task: "type[RunnableMixin[Config]]",
+        *cfgs: RawConfigType,
+        use_cli: bool | list[str] = True,
+    ) -> None:
+        run_single_process_training(task, *cfgs, use_cli=use_cli)

xax/task/launchers/staged.py

@@ -0,0 +1,29 @@
+"""Defines a base class with utility functions for staged training runs."""
+
+from abc import ABC
+from pathlib import Path
+
+from xax.task.launchers.base import BaseLauncher
+from xax.task.mixins.artifacts import ArtifactsMixin, Config
+
+
+class StagedLauncher(BaseLauncher, ABC):
+    def __init__(self, config_file_name: str = "config.yaml") -> None:
+        super().__init__()
+
+        self.config_file_name = config_file_name
+
+    def get_config_path(self, task: "ArtifactsMixin[Config]", use_cli: bool | list[str] = True) -> Path:
+        config_path = task.exp_dir / self.config_file_name
+        task.config.exp_dir = str(task.exp_dir)
+        with open(config_path, "w", encoding="utf-8") as f:
+            f.write(task.config_str(task.config, use_cli=use_cli))
+        return config_path
+
+    @classmethod
+    def from_components(cls, task_key: str, config_path: Path, use_cli: bool | list[str] = True) -> "ArtifactsMixin":
+        return (
+            ArtifactsMixin.from_task_key(task_key)
+            .get_task(config_path, use_cli=use_cli)
+            .set_exp_dir(config_path.parent)
+        )