Merge pull request #3624 from OpenNeuroOrg/export-states

feat(worker): Only fsck remotes on exports and run drop if zero errors occur

Author: nellh

services/datalad/datalad_service/broker/__init__.py

@@ -2,6 +2,7 @@
 
 from taskiq import InMemoryBroker, SmartRetryMiddleware
 from taskiq_redis import RedisAsyncResultBackend, RedisStreamBroker
+from taskiq_pipelines import PipelineMiddleware
 
 
 from datalad_service import config
@@ -37,5 +38,6 @@
                 use_delay_exponent=True,
                 max_delay_exponent=120,
             ),
+            PipelineMiddleware(),
         )
     )

services/datalad/datalad_service/tasks/fsck.py

@@ -57,39 +57,27 @@ def git_annex_fsck_local(dataset_path):
 
 
 @broker.task
-def git_annex_fsck_remote(dataset_path, branch=None, remote='s3-PUBLIC'):
-    """Run incremental fsck for one branch (tag) and remote."""
+def git_annex_fsck_remote(dataset_path, branch, remote='s3-PUBLIC'):
+    """
+    Run incremental fsck for one branch (tag) and remote.
+
+    Returns True if the check ran and no files failed, False otherwise.
+    """
     try:
         repo = pygit2.Repository(dataset_path)
         commit, references = get_head_commit_and_references(repo)
     except pygit2.GitError:
         logging.error(f'Could not open git repository for {dataset_path}')
-        return
-    if not branch:
-        # Find the newest tag chronologically
-        all_tags = sorted(
-            [tag for tag in repo.references if tag.startswith('refs/tags/')],
-            key=lambda tag: repo.lookup_reference(tag).target.commit_time,
-            reverse=True,
-        )
-        if all_tags:
-            branch = all_tags[0]
-        else:
-            logging.info(
-                f'No tags found for dataset: {dataset_path}. Skipping remote fsck.'
-            )
-            return
+        return False
 
-    # Run at most once per month per dataset
     annex_command = (
         'git-annex',
         'fsck',
         f'--branch={branch}',
-        '--from={remote}',
-        '--fast',
+        f'--from={remote}',
         '--json',
         '--json-error-messages',
-        '--incremental-schedule=7d',
+        '--incremental-schedule=30d',
     )
     annex_process = subprocess.Popen(
         annex_command, cwd=dataset_path, stdout=subprocess.PIPE
@@ -99,8 +87,11 @@ def git_annex_fsck_remote(dataset_path, branch=None, remote='s3-PUBLIC'):
         annexed_file = json.loads(annexed_file_json)
         if not annexed_file['success']:
             bad_files.append(annexed_file)
+    update_file_check(dataset_path, commit, references, bad_files, remote)
     if len(bad_files) > 0:
         logging.error(
             f'{dataset_path} remote {remote} has missing or corrupt annexed objects'
         )
-    update_file_check(dataset_path, commit, references, bad_files, remote)
+        return False
+    else:
+        return True

services/datalad/datalad_service/tasks/publish.py

@@ -8,6 +8,7 @@
 import pygit2
 import boto3
 from github import Github
+from taskiq_pipelines import Pipeline
 
 import datalad_service.common.s3
 import datalad_service.common.github
@@ -33,6 +34,7 @@
     update_s3_sibling,
 )
 from datalad_service.broker import broker
+from datalad_service.tasks.fsck import git_annex_fsck_remote
 
 logger = logging.getLogger('datalad_service.' + __name__)
 
@@ -98,16 +100,28 @@ async def export_dataset(
         update_s3_sibling(dataset_path)
         # Push the most recent tag
         if tags:
-            s3_export(dataset_path, get_s3_remote(), tags[-1].name)
+            new_tag = tags[-1].name
+            s3_export(dataset_path, get_s3_remote(), new_tag)
             s3_backup_push(dataset_path)
             # Once all S3 tags are exported, update GitHub
             if github_enabled:
                 # Perform all GitHub export steps
-                github_export(dataset_id, dataset_path, tags[-1].name)
-        # Drop cache once all exports are complete
-        clear_dataset_cache(dataset_id)
-        # Clean local annexed files once export is complete
-        await annex_drop.kiq(dataset_path)
+                github_export(dataset_id, dataset_path, new_tag)
+            # Drop cache once all exports are complete
+            clear_dataset_cache(dataset_id)
+            # Check and clean local annexed files once export is complete
+            pipeline = Pipeline(broker, git_annex_fsck_remote).call_next(
+                annex_drop, dataset_path=dataset_path, branch=new_tag
+            )
+            # Call the pipeline (arguments for git_annex_fsck_remote)
+            await pipeline.kiq(
+                dataset_path,
+                branch=new_tag,  # Check the history from the new tag just exported
+                remote=get_s3_remote(),
+            )
+        else:
+            # Clear the cache even if only sibling updates occurred
+            clear_dataset_cache(dataset_id)
 
 
 def check_remote_has_version(dataset_path, remote, tag):
@@ -214,7 +228,7 @@ def monitor_remote_configs(dataset_path):
 
 
 @broker.task
-async def annex_drop(dataset_path):
+async def annex_drop(fsck_success, dataset_path, branch):
     """Drop local contents from the annex."""
     # Ensure numcopies is set to 2 before running drop
     await run_check(['git-annex', 'numcopies', '2'], dataset_path)
@@ -226,4 +240,5 @@ async def annex_drop(dataset_path):
         # Not an issue if this fails
         pass
     # Drop will only drop successfully exported files present on both remotes
-    await run_check(['git-annex', 'drop'], dataset_path)
+    if fsck_success:
+        await run_check(['git-annex', 'drop', '--branch', branch], dataset_path)

services/datalad/pyproject.toml

@@ -20,6 +20,7 @@ dependencies = [
     "taskiq-redis>=1.0.9",
     "uvicorn[standard]>=0.34.3",
     "httpx>=0.28.1",
+    "taskiq-pipelines>=0.1.4",
 ]
 dynamic = ["version"]