xax 0.2.18__tar.gz → 0.2.20__tar.gz

{xax-0.2.18/xax.egg-info → xax-0.2.20}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: xax
-Version: 0.2.18
+Version: 0.2.20
 Summary: A library for fast Jax experimentation
 Home-page: https://github.com/kscalelabs/xax
 Author: Benjamin Bolte

{xax-0.2.18 → xax-0.2.20}/setup.py

@@ -48,4 +48,9 @@ setup(
             "requirements*.txt",
         ],
     },
+    entry_points={
+        "console_scripts": [
+            "xax-edit-config=xax.cli.edit_config:main",
+        ],
+    },
 )

{xax-0.2.18 → xax-0.2.20}/xax/__init__.py

@@ -12,7 +12,7 @@ and running the update script:
     python -m scripts.update_api --inplace
 """
 
-__version__ = "0.2.18"
+__version__ = "0.2.20"
 
 # This list shouldn't be modified by hand; instead, run the update script.
 __all__ = [

xax-0.2.20/xax/cli/edit_config.py

@@ -0,0 +1,67 @@
+"""Lets you edit a checkpoint config programmatically."""
+
+import argparse
+import difflib
+import io
+import os
+import subprocess
+import tarfile
+import tempfile
+from pathlib import Path
+
+from omegaconf import OmegaConf
+
+from xax.task.mixins.checkpointing import load_ckpt
+from xax.utils.text import colored, show_info
+
+
+def main() -> None:
+    parser = argparse.ArgumentParser()
+    parser.add_argument("ckpt_path", type=Path)
+    args = parser.parse_args()
+
+    # Loads the config from the checkpoint.
+    config = load_ckpt(args.ckpt_path, part="config")
+    config_str = OmegaConf.to_yaml(config)
+
+    # Opens the user's preferred editor to edit the config.
+    with tempfile.NamedTemporaryFile(suffix=".yaml", delete=False) as f:
+        f.write(config_str.encode("utf-8"))
+        f.flush()
+        subprocess.run([os.environ.get("EDITOR", "vim"), f.name], check=True)
+
+    # Loads the edited config.
+    try:
+        edited_config = OmegaConf.load(f.name)
+        edited_config_str = OmegaConf.to_yaml(edited_config, sort_keys=True)
+    finally:
+        os.remove(f.name)
+
+    if edited_config_str == config_str:
+        show_info("No changes were made to the config.")
+        return
+
+    # Diffs the original and edited configs.
+    diff = difflib.ndiff(config_str.splitlines(), edited_config_str.splitlines())
+    for line in diff:
+        if line.startswith("+ "):
+            print(colored(line, "light-green"), flush=True)
+        elif line.startswith("- "):
+            print(colored(line, "light-red"), flush=True)
+        elif line.startswith("? "):
+            print(colored(line, "light-cyan"), flush=True)
+
+    # Saves the edited config to the checkpoint.
+    with tarfile.open(args.ckpt_path, "w:gz") as tar:
+
+        def add_file_bytes(name: str, data: bytes) -> None:  # noqa: ANN401
+            info = tarfile.TarInfo(name=name)
+            info.size = len(data)
+            tar.addfile(info, io.BytesIO(data))
+
+        add_file_bytes("config", edited_config_str.encode())
+
+
+if __name__ == "__main__":
+    # python -m xax.cli.edit_config
+    main()

{xax-0.2.18 → xax-0.2.20}/xax/task/base.py

@@ -92,7 +92,11 @@ class BaseTask(Generic[Config]):
 
     @functools.cached_property
     def task_path(self) -> Path:
-        return Path(inspect.getfile(self.__class__))
+        try:
+            return Path(inspect.getfile(self.__class__))
+        except OSError:
+            logger.warning("Could not resolve task path for %s, returning current working directory")
+            return Path.cwd()
 
     @functools.cached_property
     def task_module(self) -> str:
@@ -172,14 +176,18 @@ class BaseTask(Generic[Config]):
         Returns:
             The merged configs.
         """
-        task_path = Path(inspect.getfile(cls))
+        try:
+            task_path = Path(inspect.getfile(cls))
+        except OSError:
+            logger.warning("Could not resolve task path for %s, returning current working directory", cls.__name__)
+            task_path = Path.cwd()
         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.stdout.write(OmegaConf.to_yaml(cfg, sort_keys=True))
+                sys.stdout.flush()
                 sys.exit(0)
 
             # Attempts to load any paths as configs.
@@ -202,7 +210,7 @@ class BaseTask(Generic[Config]):
 
     @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))
+        return OmegaConf.to_yaml(cls.get_config(*cfgs, use_cli=use_cli), sort_keys=True)
 
     @classmethod
     def get_task(cls, *cfgs: RawConfigType, use_cli: bool | list[str] = True) -> Self:

{xax-0.2.18 → xax-0.2.20}/xax/task/mixins/artifacts.py

@@ -43,8 +43,14 @@ class ArtifactsMixin(BaseTask[Config]):
     def run_dir(self) -> Path:
         run_dir = get_run_dir()
         if run_dir is None:
-            task_file = inspect.getfile(self.__class__)
-            run_dir = Path(task_file).resolve().parent
+            try:
+                task_file = inspect.getfile(self.__class__)
+                run_dir = Path(task_file).resolve().parent
+            except OSError:
+                logger.warning(
+                    "Could not resolve task path for %s, returning current working directory", self.__class__.__name__
+                )
+                run_dir = Path.cwd()
         return run_dir / self.task_name
 
     @property

{xax-0.2.18 → xax-0.2.20}/xax/task/mixins/checkpointing.py

@@ -292,7 +292,7 @@ class CheckpointingMixin(ArtifactsMixin[Config], Generic[Config]):
 
             if state is not None:
                 add_file_bytes("state", json.dumps(state.to_dict(), indent=2).encode())
-            add_file_bytes("config", OmegaConf.to_yaml(self.config).encode())
+            add_file_bytes("config", OmegaConf.to_yaml(self.config, sort_keys=True).encode())
 
         # Updates the symlink to the new checkpoint
         last_ckpt_path.unlink(missing_ok=True)

{xax-0.2.18 → xax-0.2.20/xax.egg-info}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: xax
-Version: 0.2.18
+Version: 0.2.20
 Summary: A library for fast Jax experimentation
 Home-page: https://github.com/kscalelabs/xax
 Author: Benjamin Bolte

{xax-0.2.18 → xax-0.2.20}/xax.egg-info/SOURCES.txt

@@ -11,8 +11,11 @@ xax/requirements.txt
 xax.egg-info/PKG-INFO
 xax.egg-info/SOURCES.txt
 xax.egg-info/dependency_links.txt
+xax.egg-info/entry_points.txt
 xax.egg-info/requires.txt
 xax.egg-info/top_level.txt
+xax/cli/__init__.py
+xax/cli/edit_config.py
 xax/core/__init__.py
 xax/core/conf.py
 xax/core/state.py

xax-0.2.20/xax.egg-info/entry_points.txt

@@ -0,0 +1,2 @@
+[console_scripts]
+xax-edit-config = xax.cli.edit_config:main