Add readme. migrate_ckpt also returns list of done migrations

README.md

@@ -1 +1,33 @@
 # migrate-ckpt
+
+```python
+import torch
+from migrate_ckpt import Migration, migrate_ckpt
+
+
+def update_some_keys_callback(ckpt):
+    """
+    Define a callback that takes a checkpoints and updates it.
+    """
+    ckpt["some_keys"] = ckpt["some_other_keys"]
+    del ckpt["some_other_keys"]
+    return ckpt
+
+
+# List a set of migrations. Whenever you update your model architecture,
+# you should add one that updates the model starting from the previous
+# state (output of the previous migration)
+model_migrations = [
+    Migration("Update some keys", update_some_keys_callback),
+]
+
+# Will only perform new migrations.
+# done_migrations returns the list of migration objects that were executed.
+ckpt, done_migrations = migrate_ckpt(
+    torch.load("/path/to/some/checkpoint.ckpt"),
+    model_migrations,
+)
+
+# This has no effect, the model was already migrated.
+ckpt_2, _ = migrate_ckpt(ckpt, model_migrations)
+```

migrate_ckpt/migrate.py

@@ -44,7 +44,7 @@ def _mark_ckpt(ckpt: CkptType, migration: Migration) -> CkptType:
 def migrate_ckpt(
     ckpt: CkptType,
     migrations: Sequence[Migration],
-) -> CkptType:
+) -> tuple[CkptType, Sequence[Migration]]:
     """
     Migrate checkpoint using provided migrations
     Args:
@@ -55,4 +55,4 @@ def migrate_ckpt(
     for migration in missing_migrations:
         ckpt = migration.callback(ckpt)
         ckpt = _mark_ckpt(ckpt, migration)
-    return ckpt
+    return ckpt, missing_migrations

tests/test_migration.py

@@ -26,21 +26,21 @@ def update_test_callback(x):
 
 def test_missing_fields():
     ckpt: dict[str, Any] = {}
-    new_ckpt = migrate_ckpt(ckpt, [blank_migration])
+    new_ckpt, _ = migrate_ckpt(ckpt, [blank_migration])
     assert ckpt_migration_key in new_ckpt
     assert isinstance(new_ckpt[ckpt_migration_key], list)
     assert len(new_ckpt[ckpt_migration_key]) == 1
     assert new_ckpt[ckpt_migration_key][0] == "blank"
 
 
 def test_missing_one_migration():
-    ckpt = migrate_ckpt({}, [blank_migration])
-    new_ckpt = migrate_ckpt(ckpt, [blank2_migration])
+    ckpt, _ = migrate_ckpt({}, [blank_migration])
+    new_ckpt, _ = migrate_ckpt(ckpt, [blank2_migration])
     assert new_ckpt[ckpt_migration_key][1] == "blank2"
 
 
 def test_execute_migration():
-    ckpt = migrate_ckpt(
+    ckpt, _ = migrate_ckpt(
         {},
         [add_field_migration],
     )
@@ -49,7 +49,7 @@ def test_execute_migration():
 
 
 def test_execute_related_migrations():
-    ckpt = migrate_ckpt(
+    ckpt, _ = migrate_ckpt(
         {},
         [add_field_migration, update_test_migration],
     )