[Execution] Implement support for version overrides

errors/errors.yml

@@ -141,6 +141,13 @@
     All extra files must be declared using relative paths (relative to the COND
     file containing the environment() directive).
 
+2011:
+  name: CorruptedVersionIndex
+  message: >-
+    Detected a corrupted version index: task '{task_identifier}' references
+    timestamp '{timestamp}' in version_overrides, but the corresponding
+    version metadata is missing.
+
 
 # Execution errors (error code 3xxx)
 3001:

src/conductor/errors/generated.py

@@ -412,6 +412,22 @@ def _message(self):
         )
 
 
+class CorruptedVersionIndex(ConductorError):
+    error_code = 2011
+
+    def __init__(self, **kwargs):
+        super().__init__()
+        self.kwargs = kwargs
+        self.task_identifier = kwargs["task_identifier"]
+        self.timestamp = kwargs["timestamp"]
+
+    def _message(self):
+        return "Detected a corrupted version index: task '{task_identifier}' references timestamp '{timestamp}' in version_overrides, but the corresponding version metadata is missing.".format(
+            task_identifier=self.task_identifier,
+            timestamp=self.timestamp,
+        )
+
+
 class TaskNonZeroExit(ConductorError):
     error_code = 3001
 
@@ -923,6 +939,7 @@ def _message(self):
     2008: DuplicateEnvName,
     2009: EnvNotEnv,
     2010: EnvExtraFilesNotRelative,
+    2011: CorruptedVersionIndex,
     3001: TaskNonZeroExit,
     3002: TaskFailed,
     3003: OutputDirTaken,
@@ -989,6 +1006,7 @@ def _message(self):
     "DuplicateEnvName",
     "EnvNotEnv",
     "EnvExtraFilesNotRelative",
+    "CorruptedVersionIndex",
     "TaskNonZeroExit",
     "TaskFailed",
     "OutputDirTaken",

src/conductor/execution/version_index.py

@@ -5,7 +5,7 @@
 from typing import Any, Iterable, List, Optional, Tuple, Sequence
 
 from conductor.config import VERSION_INDEX_BACKUP_NAME_TEMPLATE
-from conductor.errors import UnsupportedVersionIndexFormat
+from conductor.errors import UnsupportedVersionIndexFormat, CorruptedVersionIndex
 from conductor.task_identifier import TaskIdentifier
 from conductor.utils.git import Git
 import conductor.execution.version_index_queries as q
@@ -64,7 +64,8 @@ class VersionIndex:
     """
 
     # v0.4.0 and older: FormatVersion = 1
-    FormatVersion = 2
+    # v0.7.0 and older: FormatVersion = 2
+    FormatVersion = 3
 
     def __init__(
         self,
@@ -84,7 +85,14 @@ def create_or_load(cls, path: pathlib.Path) -> "VersionIndex":
             if format_version == 1:
                 # Upgrade the version index to format 2.
                 cls._run_v1_to_v2_migration(conn, path)
-            elif format_version != cls.FormatVersion:
+                format_version = 2
+
+            if format_version == 2:
+                # Upgrade the version index to format 3.
+                cls._run_v2_to_v3_migration(conn, path)
+                format_version = 3
+
+            if format_version != cls.FormatVersion:
                 raise UnsupportedVersionIndexFormat(version=format_version)
 
             # Need to restore the last timestamp used.
@@ -102,6 +110,7 @@ def create_or_load(cls, path: pathlib.Path) -> "VersionIndex":
         conn = sqlite3.connect(path)
         conn.execute(q.set_format_version.format(version=cls.FormatVersion))
         conn.execute(q.create_table)
+        conn.execute(q.create_version_overrides_table)
         conn.commit()
         return VersionIndex(conn, 0, path)
 
@@ -197,6 +206,40 @@ def get_all_unversioned(self) -> List[TaskIdentifier]:
             # We create the unversioned table only when we add unversioned tasks.
             return []
 
+    def set_version_override(
+        self, task_identifier: TaskIdentifier, timestamp: int
+    ) -> None:
+        cursor = self._conn.cursor()
+        cursor.execute(q.upsert_version_override, (str(task_identifier), timestamp))
+
+    def clear_version_override(self, task_identifier: TaskIdentifier) -> None:
+        cursor = self._conn.cursor()
+        cursor.execute(q.delete_version_override, (str(task_identifier),))
+
+    def get_version_override(
+        self, task_identifier: TaskIdentifier
+    ) -> Optional[Version]:
+        cursor = self._conn.cursor()
+        cursor.execute(q.get_version_override, (str(task_identifier),))
+        row = cursor.fetchone()
+        if row is None:
+            return None
+
+        timestamp = int(row[0])
+        # A missing row in version_index indicates a broken reference from
+        # version_overrides, which is a corruption of the version index state.
+        if row[2] is None:
+            raise CorruptedVersionIndex(
+                task_identifier=str(task_identifier),
+                timestamp=timestamp,
+            )
+
+        return Version(
+            timestamp=timestamp,
+            commit_hash=row[1],
+            has_uncommitted_changes=(False if row[2] == 0 else True),
+        )
+
     def get_versioned_tasks(
         self, tasks: Optional[List[TaskIdentifier]], latest_only: bool
     ) -> List[Tuple[TaskIdentifier, Version]]:
@@ -319,6 +362,26 @@ def _run_v1_to_v2_migration(conn: sqlite3.Connection, path: pathlib.Path):
             conn.rollback()
             raise
 
+    @staticmethod
+    def _run_v2_to_v3_migration(conn: sqlite3.Connection, path: pathlib.Path):
+        # Upgrades the version index's persistent format from version 2 to 3.
+        # This adds the `version_overrides` table.
+        backup_copy_path = path.with_name(
+            VERSION_INDEX_BACKUP_NAME_TEMPLATE.format(vfrom=2, vto=3)
+        )
+        if not backup_copy_path.exists():
+            # Back up the version index file first.
+            shutil.copy2(src=path, dst=backup_copy_path)
+
+        # Run the migration.
+        try:
+            conn.execute(q.v2_to_v3_create_version_overrides_table)
+            conn.execute(q.set_format_version.format(version=3))
+            conn.commit()
+        except RuntimeError:
+            conn.rollback()
+            raise
+
     def _version_from_row(self, row: Sequence[Any]) -> Version:
         return Version(
             timestamp=row[0],

src/conductor/execution/version_index_queries.py

@@ -17,6 +17,14 @@
   )
 """
 
+create_version_overrides_table = """
+  CREATE TABLE IF NOT EXISTS version_overrides (
+    task_identifier TEXT NOT NULL,
+    timestamp INTEGER NOT NULL,
+    PRIMARY KEY (task_identifier)
+  )
+"""
+
 set_format_version = "PRAGMA user_version = {version:d}"
 
 get_format_version = "PRAGMA user_version"
@@ -128,6 +136,33 @@
   SELECT task_identifier FROM unversioned
 """
 
+upsert_version_override = """
+  INSERT INTO version_overrides (task_identifier, timestamp)
+  VALUES (?, ?)
+  ON CONFLICT(task_identifier)
+  DO UPDATE SET timestamp=excluded.timestamp
+"""
+
+delete_version_override = """
+  DELETE FROM version_overrides WHERE task_identifier = ?
+"""
+
+get_version_override = """
+  SELECT
+    o.timestamp,
+    v.git_commit_hash,
+    v.has_uncommitted_changes
+  FROM
+    version_overrides AS o
+  LEFT JOIN
+    version_index AS v
+  ON
+    o.task_identifier = v.task_identifier
+    AND o.timestamp = v.timestamp
+  WHERE
+    o.task_identifier = ?
+"""
+
 
 # Queries used in format 1 (retained for testing purposes)
 
@@ -160,3 +195,9 @@
 v1_to_v2_drop_old_table = "DROP TABLE version_index"
 
 v1_to_v2_rename_new_table = "ALTER TABLE version_index_new RENAME TO version_index"
+
+
+# Queries used for migrating from format 2 to format 3
+# - Add the `version_overrides` table
+
+v2_to_v3_create_version_overrides_table = create_version_overrides_table

src/conductor/task_types/run.py

@@ -1,5 +1,5 @@
 import pathlib
-from typing import Sequence, Optional, TYPE_CHECKING
+from typing import Sequence, Optional, Tuple, TYPE_CHECKING
 
 import conductor.filename as f
 from conductor.errors import InternalError
@@ -164,7 +164,8 @@ def __init__(
             env=env,
         )
         self._did_retrieve_version = False
-        self._most_relevant_version: Optional[Version] = None
+        # (most_relevant_version, is_override)
+        self._most_relevant_version: Tuple[Optional[Version], bool] = (None, False)
 
     @property
     def archivable(self) -> bool:
@@ -176,17 +177,18 @@ def record_output(self) -> bool:
 
     def get_output_version(self, ctx: "c.Context") -> Optional[Version]:
         self._ensure_most_relevant_existing_version_computed(ctx)
-        return self._most_relevant_version
+        return self._most_relevant_version[0]
 
     def get_output_path(self, ctx: "c.Context") -> Optional[pathlib.Path]:
         self._ensure_most_relevant_existing_version_computed(ctx)
-        if self._most_relevant_version is None:
+        most_relevant_version, _ = self._most_relevant_version
+        if most_relevant_version is None:
             return None
 
         unversioned_path = super().get_output_path(ctx)
         assert unversioned_path is not None
         return unversioned_path.with_name(
-            f.task_output_dir(self.identifier, version=self._most_relevant_version)
+            f.task_output_dir(self.identifier, version=most_relevant_version)
         )
 
     def get_specific_output_path(
@@ -209,18 +211,20 @@ def should_run(self, ctx: "c.Context", at_least_commit: Optional[str]) -> bool:
         whether or not this task needs to execute.
         """
         self._ensure_most_relevant_existing_version_computed(ctx)
-        if self._most_relevant_version is None:
+        most_relevant_version, is_override = self._most_relevant_version
+        if most_relevant_version is None:
             # Must run because no relevant version exists.
             return True
-        if at_least_commit is None:
+        if at_least_commit is None or is_override:
             # There already is a most relevant version and we are not asked to
-            # run for at least some commit.
+            # run for at least some commit. (Or this is an override, in which
+            # case we should not compare commit hashes.)
             return False
-        if self._most_relevant_version.commit_hash is None:
+        if most_relevant_version.commit_hash is None:
             # Must re-run since the most relevant version does not have a commit
             # hash and `at_least_commit` is set to some commit.
             return True
-        if self._most_relevant_version.commit_hash == at_least_commit:
+        if most_relevant_version.commit_hash == at_least_commit:
             # No need to re-run. The most relevant version matches `at_least_commit`.
             return False
 
@@ -231,7 +235,7 @@ def should_run(self, ctx: "c.Context", at_least_commit: Optional[str]) -> bool:
         # then it must be "older" (this is based on how we define the most
         # relevant version).
         most_relevant_is_older = ctx.git.is_ancestor(
-            at_least_commit, self._most_relevant_version.commit_hash
+            at_least_commit, most_relevant_version.commit_hash
         )
         if not most_relevant_is_older:
             # No need to re-run.
@@ -243,17 +247,18 @@ def should_run(self, ctx: "c.Context", at_least_commit: Optional[str]) -> bool:
 
     def create_new_version(self, ctx: "c.Context") -> Version:
         self._create_new_version(ctx)
-        assert self._most_relevant_version is not None
-        return self._most_relevant_version
+        assert self._most_relevant_version[0] is not None
+        return self._most_relevant_version[0]
 
     def _create_new_version(self, ctx: "c.Context") -> None:
         # N.B. If this task fails, the value of `most_relevant_version` will be
         # incorrect. However, any tasks that have this task as a dependency will
         # be skipped, so no incorrectness will occur. The new version will not
         # be committed into the version index.
         self._did_retrieve_version = True
-        self._most_relevant_version = ctx.version_index.generate_new_output_version(
-            commit=ctx.current_commit
+        self._most_relevant_version = (
+            ctx.version_index.generate_new_output_version(commit=ctx.current_commit),
+            False,
         )
 
     def _ensure_most_relevant_existing_version_computed(self, ctx: "c.Context"):
@@ -265,7 +270,7 @@ def _ensure_most_relevant_existing_version_computed(self, ctx: "c.Context"):
 
     def _retrieve_most_relevant_existing_version(
         self, ctx: "c.Context"
-    ) -> Optional[Version]:
+    ) -> Tuple[Optional[Version], bool]:
         """
         Finds the "most relevant" existing version of this task's outputs, if
         one exists. The definition of "most relevant" depends on whether git is
@@ -274,19 +279,27 @@ def _retrieve_most_relevant_existing_version(
 
         This method is meant for internal use.
         """
+        # First check if there is a version override for this task.
+        override = ctx.version_index.get_version_override(self._identifier)
+        if override is not None:
+            return (override, True)
+
         # Simple case. If the project does not use git, the most relevant
         # existing version is the latest (newest) version (if it exists).
         if not ctx.uses_git:
             res = ctx.version_index.get_latest_output_version(self._identifier)
-            return res
+            return (res, False)
 
         # Retrieve the commit hash associated with `HEAD`.
         curr_commit = ctx.current_commit
 
         # This case happens if the repository is bare (no commits). Then the
         # most relevant existing version is the latest version, if it exists.
         if curr_commit is None:
-            return ctx.version_index.get_latest_output_version(self._identifier)
+            return (
+                ctx.version_index.get_latest_output_version(self._identifier),
+                False,
+            )
 
         # Retrieve all existing versions for this task. Filter them into tasks
         # with null commit hashes and ones that are ancestors.
@@ -321,7 +334,7 @@ def _retrieve_most_relevant_existing_version(
                 ):
                     selected_version = v
             assert selected_version is not None
-            return selected_version
+            return (selected_version, False)
 
         # There are no ancestor commits and all existing versions do not have a
         # commit hash. We select the newest version. This maintains the same
@@ -330,9 +343,9 @@ def _retrieve_most_relevant_existing_version(
             len(null_commit_versions) == len(existing_versions)
             and len(null_commit_versions) > 0
         ):
-            return max(null_commit_versions, key=lambda v: v.timestamp)
+            return (max(null_commit_versions, key=lambda v: v.timestamp), False)
 
         # Otherwise, this means there may exist versions with commits that are
         # not ancestors of the current commit. For correctness, we should not
         # depend on the results from any previous version.
-        return None
+        return (None, False)