feat(migrations): remove migration rollback

models/docs/cli/migration.rst

@@ -11,11 +11,6 @@
 anemoi-models provides a way for users to migrate old checkpoints so that they can be
 used with recent versions of anemoi-models.
 
-.. caution::
-
-    Migrating your checkpoint is not secure. Only migrate a checkpoint that you trust.
-    Migrating a checkpoint may run rollback functions which can run arbitrary code.
-
 
 ********
  Inspect
@@ -91,9 +86,3 @@ Update your anemoi-models to the desired version and call ``anemoi-models migrat
 
 Note that this should work when updating to a newer version, as well as downgrading to an older
 version, as long as your checkpoint is not too old.
-
-*********
- Rollback
-*********
-If you update to an older version, the checkpoint will be rollbacked to be compatible with this
-older version.

models/docs/create-migrations.rst

@@ -52,24 +52,6 @@ looks like:
        return ckpt
 
 
-   def rollback(ckpt: CkptType) -> CkptType:
-       """
-       Rollback the checkpoint.
-
-
-       Parameters
-       ----------
-       ckpt : CkptType
-           The checkpoint dict.
-
-       Returns
-       -------
-       CkptType
-           The rollbacked checkpoint dict.
-       """
-
-       return ckpt
-
 ``migrate`` receives an old checkpoint (made before your changes), and
 must return a checkpoint compatible with your changes.
 
@@ -108,34 +90,6 @@ The unit test ``test_migration_order`` will check whether the correct
 order is preserved. If you get an error, you can run ``anemoi-models
 migration fix-order`` to update the timestamps of your scripts.
 
-**********
- Rollback
-**********
-
-``rollback`` does the opposite operation and receives a checkpoint
-compatible with your changes and must return a checkpoint usable before
-your change.
-
-.. note::
-
-   We use `cloudpickle <https://github.com/cloudpipe/cloudpickle>`_ to
-   pickle the rollback function by value rather than by inference. In
-   particular, you should follow the recommandations described `here
-   <https://github.com/cloudpipe/cloudpickle/tree/master?tab=readme-ov-file#overriding-pickles-serialization-mechanism-for-importable-constructs>`_.
-
-.. note::
-
-   Rollback functions are not strictly required. However checkpoints
-   will not be able to be rollbacked before your migration script if it
-   does not have a rollback.
-
-To generate a migration script without a rollback use the
-``--no-rollback`` parameter:
-
-.. code:: bash
-
-   anemoi-models migration create migration-name --no-rollback
-
 ****************
  Simple example
 ****************
@@ -175,24 +129,6 @@ migration:
        return ckpt
 
 
-   def rollback(ckpt: CkptType) -> CkptType:
-       """
-       Rollback the checkpoint.
-
-
-       Parameters
-       ----------
-       ckpt : CkptType
-           The checkpoint dict.
-
-       Returns
-       -------
-       CkptType
-           The rollbacked checkpoint dict.
-       """
-       ckpt["state_dict"]["x"] = ckpt["state_dict"].pop("y")
-       return ckpt
-
 ****************
  Setup callback
 ****************
@@ -224,11 +160,6 @@ define the ``migrate_setup`` callback:
           A MigrationContext instance
        """
 
-.. note::
-
-   The setup is only defined for migrate. The setup required for the
-   rollback pass is automatically inferred.
-
 To generate your script with the setup callbacks, use the
 ``--with-setup`` argument:
 
@@ -297,23 +228,6 @@ For example, if you renamed the module
        return ckpt
 
 
-   def rollback(ckpt: CkptType) -> CkptType:
-       """
-       Rollback the checkpoint.
-
-
-       Parameters
-       ----------
-       ckpt : CkptType
-           The checkpoint dict.
-
-       Returns
-       -------
-       CkptType
-           The rollbacked checkpoint dict.
-       """
-       return ckpt
-
 Similarly, if you moved the class ``NormalizerSchema`` from
 ``anemoi.training.schemas.data`` to
 ``anemoi.models.schemas.data_processor``, the setup callback might look
@@ -403,37 +317,6 @@ Here is a full example of a migration to fix `PR 433
        return ckpt
 
 
-   def rollback_setup(context: MigrationContext) -> None:
-       """
-       Rollback setup callback to be run before loading the checkpoint.
-
-       Parameters
-       ----------
-       context : MigrationContext
-          A MigrationContext instance
-       """
-       context.move_attribute(
-           "anemoi.models.schemas.data_processor.NormalizerSchema", "anemoi.training.schemas.data.NormalizerSchema"
-       )
-
-
-   def rollback(ckpt: CkptType) -> CkptType:
-       """
-       Rollback the checkpoint.
-
-
-       Parameters
-       ----------
-       ckpt : CkptType
-           The checkpoint dict.
-
-       Returns
-       -------
-       CkptType
-           The rollbacked checkpoint dict.
-       """
-       return ckpt
-
 ****************
  Best practices
 ****************

models/docs/modules/migrations.rst

@@ -28,7 +28,6 @@ ordered list of scripts. Each script contains:
 -  some metadata information, such as the version of the migration
    system, or the version of anemoi-models,
 -  a ``migrate`` function to migrate checkpoints
--  optionnally a ``rollback`` function to reverse the migration function
 -  optionnally a ``migrate_setup`` function to fix import issues.
 
 Similarly, the checkpoint contain some migration information that
@@ -44,9 +43,6 @@ informs on its migration state:
    For now, it only logs a warning, but a more complex behavior could be
    added in the future,
 
--  the ``rollback`` function: the same as in the script, in case you
-   need the checkpoint in an older version,
-
 **********************
  Compatibility groups
 **********************
@@ -103,24 +99,20 @@ algorithm. To follow along, here is an example:
 +----------------+-----------------+
 | migration 2    | migration 2     |
 +----------------+-----------------+
-| migration 5    | migration 3     |
+| migration 5    |                 |
 +----------------+-----------------+
-| migration 6    | migration 4     |
+| migration 6    |                 |
 +----------------+-----------------+
 | migration 7    |                 |
 +----------------+-----------------+
 
--  First, we rollback any additional migrations in the checkpoint,
-   starting from the end (here, migration 4 and 3 need to be
-   rollbacked).
+-  First, we check if there are extra migrations in the checkpoint. If so, fail.
 
 -  Then, we migrate any missing migrations in the checkpoint, starting
    from the start (here migration 5, 6 and 7).
 
 In the example, it will produce:
 
--  ROLLBACK migration 4
--  ROLLBACK migration 3
 -  MIGRATE migration 5
 -  MIGRATE migration 6
 -  MIGRATE migration 7
@@ -137,7 +129,6 @@ checkpoint. It can be accessed through:
    >>> history = metadata.get("migrations", {}).get("history", [])
    >>> for executed_migration in history:
    ...    print(executed_migration)
-   { "type": "rollback", "name": "migration_name1.py", "signature": "[...]" }
    { "type": "migrate", "name": "migration_name2.py", "signature": "[...]" }
 
 **********

models/pyproject.toml

@@ -65,7 +65,6 @@ optional-dependencies.docs = [
 ]
 
 optional-dependencies.migrations = [
-  "cloudpickle>=3.1",
   "jinja2>=3.1",
   "rich",
 ]

models/src/anemoi/models/commands/migration.py

@@ -21,9 +21,7 @@
 
 from ..migrations import MIGRATION_PATH
 from ..migrations import IncompatibleCheckpointException
-from ..migrations import MigrationOp
 from ..migrations import Migrator
-from ..migrations import RollbackOp
 from ..migrations.migrator import LOGGER as migrator_logger
 from . import Command
 
@@ -144,25 +142,6 @@ def migrate(ckpt: CkptType) -> CkptType:
         The migrated checkpoint dict.
     \"""
     return ckpt
-{% if not no_rollback %}
-
-
-def rollback(ckpt: CkptType) -> CkptType:
-    \"""Rollback the checkpoint.
-
-
-    Parameters
-    ----------
-    ckpt : CkptType
-        The checkpoint dict.
-
-    Returns
-    -------
-    CkptType
-        The rollbacked checkpoint dict.
-    \"""
-    return ckpt
-{% endif %}
 {% endif %}
 """
 
@@ -196,12 +175,6 @@ def add_arguments(self, command_parser: ArgumentParser) -> None:
             default=False,
             help="Set this if need the migrate_setup and rollback_setup callback.",
         )
-        create_parser.add_argument(
-            "--no-rollback",
-            action="store_true",
-            default=False,
-            help="Set this if you do not plan to support rollbacking.",
-        )
 
         help_sync = "Apply migrations to a checkpoint."
         sync_parser = subparsers.add_parser("sync", help=help_sync, description=help_sync)
@@ -314,14 +287,9 @@ def run_sync(self, args: Namespace) -> None:
             if not len(done_ops):
                 console.print("Your checkpoint is already compatible :party_popper:! No missing migration to execute.")
             for op in done_ops:
-                if isinstance(op, RollbackOp):
-                    console.print(
-                        f"  [red]+ ROLLBACK [bold]{op.migration.name}[/bold] \\[v{op.migration.metadata.versions['anemoi-models']}][/red]"
-                    )
-                elif isinstance(op, MigrationOp):
-                    console.print(
-                        f"  [green]+ MIGRATE [bold]{op.migration.name}[/bold] \\[v{op.migration.metadata.versions['anemoi-models']}][/green]"
-                    )
+                console.print(
+                    f"  [green]+ MIGRATE [bold]{op.migration.name}[/bold] \\[v{op.migration.metadata.versions['anemoi-models']}][/green]"
+                )
         except IncompatibleCheckpointException as e:
             print(str(e))
 
@@ -356,17 +324,6 @@ def run_inspect(self, args: Namespace) -> None:
                 console.print(
                     f"  [cyan]* [bold]{migration.name}[/bold] \\[v{migration.metadata.versions['anemoi-models']}][/cyan]"
                 )
-            if len(extra_migrations):
-                print(
-                    len(extra_migrations),
-                    "extra",
-                    maybe_plural(len(extra_migrations), "migration"),
-                    "to rollback:",
-                )
-            for migration in extra_migrations:
-                console.print(
-                    f"  [red]+ [bold]{migration.name}[/bold] \\[v{migration.metadata.versions['anemoi-models']}][/red]"
-                )
             if len(missing_migrations):
                 print(
                     len(missing_migrations),
@@ -379,7 +336,19 @@ def run_inspect(self, args: Namespace) -> None:
                 console.print(
                     f"  [green]+ [bold]{migration.name}[/bold] \\[v{migration.metadata.versions['anemoi-models']}][/green]"
                 )
-            if len(missing_migrations) or len(extra_migrations):
+            if len(extra_migrations):
+                print(
+                    len(extra_migrations),
+                    "extra",
+                    maybe_plural(len(extra_migrations), "migration"),
+                    "in the checkpoint.",
+                )
+                for migration in extra_migrations:
+                    console.print(f"  [red]+ [bold]{migration}[/bold][/red]")
+                console.print(
+                    "\n[red][italic]Your checkpoint cannot be migrated because it contains extra migrations.[/italic][/red]"
+                )
+            if len(missing_migrations) and not len(extra_migrations):
                 console.print("\n[italic]To update your checkpoint, run:[/italic]")
                 console.print(f"  [italic]anemoi-models migration sync {args.ckpt}[/italic]")
         except IncompatibleCheckpointException as e:

models/src/anemoi/models/migrations/__init__.py

@@ -8,7 +8,6 @@
 # nor does it submit to any jurisdiction.
 
 from .migrator import MIGRATION_PATH
-from .migrator import BaseOp
 from .migrator import CkptType
 from .migrator import IncompatibleCheckpointException
 from .migrator import Migration
@@ -17,14 +16,12 @@
 from .migrator import MigrationVersions
 from .migrator import Migrator
 from .migrator import MissingAttribute
-from .migrator import RollbackOp
 from .migrator import SaveCkpt
 from .migrator import SerializedMigration
 from .setup_context import MigrationContext
 
 __all__ = [
     "MIGRATION_PATH",
-    "BaseOp",
     "CkptType",
     "IncompatibleCheckpointException",
     "Migration",
@@ -33,7 +30,6 @@
     "MigrationVersions",
     "Migrator",
     "MissingAttribute",
-    "RollbackOp",
     "SaveCkpt",
     "SerializedMigration",
     "MigrationContext",