xax 0.2.18__py3-none-any.whl → 0.2.20__py3-none-any.whl

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/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/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/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/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.dist-info → xax-0.2.20.dist-info}/METADATA

@@ -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.dist-info → xax-0.2.20.dist-info}/RECORD

@@ -1,7 +1,9 @@
-xax/__init__.py,sha256=ru-amUXLIpFI1cE7NfkPpff6NnuLbSh6LHbJkmvH2zM,15733
+xax/__init__.py,sha256=yWQcHMlP2cKIpfJCJLXv796F-AAHQhS-1sRxu7871mw,15733
 xax/py.typed,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 xax/requirements-dev.txt,sha256=qkscNkFzWd1S5fump-AKH53rR65v2x5FmboFdy_kKvs,128
 xax/requirements.txt,sha256=6qY-84e-sTmlfJNrSjwONQKqzAn5h8G_oGIhnhmfSr4,302
+xax/cli/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+xax/cli/edit_config.py,sha256=99x_k6aNimbcebi2vSJhln-cv4364h6GQdRccuv_qcs,2069
 xax/core/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 xax/core/conf.py,sha256=d7Dp_GwKnaxtkztlSrJSM_LR0UYJX_FWTtceIWCBkxc,5138
 xax/core/state.py,sha256=KsNMnM_RgsZ2Ntc2pp4Fi6zG4rZb_89-kqmyGxDvyRg,4974
@@ -16,7 +18,7 @@ xax/nn/metrics.py,sha256=OAkeScwhi-wTBIJ59KHUhYbZTq4V4V-LG-mKlxMJ7bY,3238
 xax/nn/parallel.py,sha256=fnTiT7MsG7eQrJvqwjIz2Ifo3P27TuxIJzmpGYSa_dQ,4608
 xax/nn/ssm.py,sha256=8dLAcQ1hBaMT-kkHvwGu_ecxJeTY32WeMYmd4T4KtxA,10745
 xax/task/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-xax/task/base.py,sha256=ymGld72zTSw50PBpX4WwqQDyJgQSagDrLh56aqvLT40,7720
+xax/task/base.py,sha256=TYANmjNcce4_V5ZSYLnE91PXRn7Nn0nT7hN8plW_Au0,8117
 xax/task/logger.py,sha256=W_BpluYvQai1lh1dDCAj-2_mWUC1buhwJncHygDffjc,41125
 xax/task/script.py,sha256=bMMIJoUtpSBvPp6-7bejTrajTXvSg0794sYLKdPIToE,972
 xax/task/task.py,sha256=UHMpnv__gqMcfbC_L-Hhk-DCnUYlFVsgbNf-v8o8B7U,1424
@@ -31,8 +33,8 @@ xax/task/loggers/state.py,sha256=0Jy0NYnY4c0qt0LvNlaTaCKOSqk5SCKln5VdyuQGnIc,140
 xax/task/loggers/stdout.py,sha256=giKSW2R83YkgRefm3BLkE7t8Pbj5Dux4AgsdJxYIbGo,6619
 xax/task/loggers/tensorboard.py,sha256=sRyBbeBeVXDTYhPZIKIapW0JEfL9hqqzhNTeIcSd374,8883
 xax/task/mixins/__init__.py,sha256=D3oU31rB9FeOr9MPLleLt5JFbftUr4sBTwgnwQdc2qA,809
-xax/task/mixins/artifacts.py,sha256=Ma7fwsp-SA1w6GcuBSskszj5TB83yxYJm4Ns_EnqkI4,3018
-xax/task/mixins/checkpointing.py,sha256=w-xsx8tH8PPYRf9eoqvz-4kAdunHMvazcoZAIj6niyI,11468
+xax/task/mixins/artifacts.py,sha256=IBPAQMCGd7PQiZHfSjLakPW5j7cNuL6AsW6QkVSc02E,3277
+xax/task/mixins/checkpointing.py,sha256=v50IZ7j58DWmEu-_6Zh_02R5KUVGhrMkg5n-MYM_J4c,11484
 xax/task/mixins/compile.py,sha256=PG5aF3W9v_xGiImHgUJ7gmwuQQoSQWufdpl2N_mlLX0,3922
 xax/task/mixins/cpu_stats.py,sha256=rO_9a82ZdsNec61ya4FpYE-rWqPhpijRSXsOfc6caFA,9595
 xax/task/mixins/data_loader.py,sha256=Tp7zqPdfH2_JuE6J6EP-fEtCQpq9MjKlGHYK7Zh-goU,6599
@@ -58,8 +60,9 @@ xax/utils/data/collate.py,sha256=Rd9vMomr_S_zCa_Hi4dO-8ntzAfVwndIUtuXFA3iNcc,706
 xax/utils/types/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 xax/utils/types/frozen_dict.py,sha256=ebtHENhyUzSjyJTlbMaLtcckQIJ7EtgJiok_40TJZpo,4689
 xax/utils/types/hashable_array.py,sha256=l5iIcFmkYzfGeaZmcSoeFkthFASqM8xJYK3AXhZQYwc,992
-xax-0.2.18.dist-info/licenses/LICENSE,sha256=HCN2bImAzUOXldAZZI7JZ9PYq6OwMlDAP_PpX1HnuN0,1071
-xax-0.2.18.dist-info/METADATA,sha256=phU79eE74vMM79HouJP9iQgTDONJ59no3ktUl5GHF78,1880
-xax-0.2.18.dist-info/WHEEL,sha256=wXxTzcEDnjrTwFYjLPcsW_7_XihufBwmpiBeiXNBGEA,91
-xax-0.2.18.dist-info/top_level.txt,sha256=g4Au_r2XhvZ-lTybviH-Fh9g0zF4DAYHYxPue1-xbs8,4
-xax-0.2.18.dist-info/RECORD,,
+xax-0.2.20.dist-info/licenses/LICENSE,sha256=HCN2bImAzUOXldAZZI7JZ9PYq6OwMlDAP_PpX1HnuN0,1071
+xax-0.2.20.dist-info/METADATA,sha256=YCcDox7HsHIVUeDeZYCvoaGYhsU7TzspX-v-Xw0-H4g,1880
+xax-0.2.20.dist-info/WHEEL,sha256=0CuiUZ_p9E4cD6NyLD6UG80LBXYyiSYZOKDm5lp32xk,91
+xax-0.2.20.dist-info/entry_points.txt,sha256=uRC6rx5ce0bf-FblJaZSBMxxKFfMyoWTf8OWbBmLSe8,61
+xax-0.2.20.dist-info/top_level.txt,sha256=g4Au_r2XhvZ-lTybviH-Fh9g0zF4DAYHYxPue1-xbs8,4
+xax-0.2.20.dist-info/RECORD,,

{xax-0.2.18.dist-info → xax-0.2.20.dist-info}/WHEEL

@@ -1,5 +1,5 @@
 Wheel-Version: 1.0
-Generator: setuptools (80.1.0)
+Generator: setuptools (80.3.1)
 Root-Is-Purelib: true
 Tag: py3-none-any
 

xax-0.2.20.dist-info/entry_points.txt

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