diff --git a/ami/jobs/tests/test_jobs.py b/ami/jobs/tests/test_jobs.py index 7a7f8d5ff..4aab76437 100644 --- a/ami/jobs/tests/test_jobs.py +++ b/ami/jobs/tests/test_jobs.py @@ -1215,7 +1215,10 @@ def test_dispatch_mode_filtering(self): def test_ml_job_dispatch_mode_set_on_creation(self): """Test that ML jobs get dispatch_mode set based on project feature flags at creation time.""" - # Without async flag, ML job should default to sync_api + # With the async flag disabled, an ML job is dispatched via sync_api. + self.project.feature_flags.async_pipeline_workers = False + self.project.save() + sync_job = Job.objects.create( job_type_key=MLJob.key, project=self.project, diff --git a/ami/main/migrations/0094_enable_async_pipeline_workers.py b/ami/main/migrations/0094_enable_async_pipeline_workers.py new file mode 100644 index 000000000..109df3451 --- /dev/null +++ b/ami/main/migrations/0094_enable_async_pipeline_workers.py @@ -0,0 +1,59 @@ +""" +Turn on the ``async_pipeline_workers`` feature flag for every existing project. + +This rolls out async ML processing (workers that pull tasks from the NATS queue +instead of the synchronous push API) to all projects at once. New projects keep +whatever the model default is at creation time; this migration only updates rows +that exist at deploy time. + +The flag lives inside the ``feature_flags`` JSONB column (a ``ProjectFeatureFlags`` +pydantic model). The update toggles only the one key server-side with ``jsonb_set`` +rather than reading each project's JSON into Python, mutating it, and writing the +whole value back. Doing it in place leaves the other feature flags untouched even +if another process changes one of them during the deploy, and it runs as a single +statement instead of one save per row. The reverse flips the flag back off for +every project — a blanket disable. It does not restore per-project values from +before the rollout: some projects may have had the flag enabled individually +beforehand, so the reverse returns every project to the off state rather than to +its prior value. +""" + +from django.db import migrations + + +def _set_async_pipeline_workers(apps, schema_editor, *, enabled): + Project = apps.get_model("main", "Project") + table = schema_editor.connection.ops.quote_name(Project._meta.db_table) + with schema_editor.connection.cursor() as cursor: + cursor.execute( + f""" + UPDATE {table} + SET feature_flags = jsonb_set( + COALESCE(feature_flags, '{{}}'::jsonb), + '{{async_pipeline_workers}}', + to_jsonb(%s::boolean), + true + ) + WHERE COALESCE((feature_flags ->> 'async_pipeline_workers')::boolean, false) + IS DISTINCT FROM %s + """, + [enabled, enabled], + ) + + +def enable_async_pipeline_workers(apps, schema_editor): + _set_async_pipeline_workers(apps, schema_editor, enabled=True) + + +def disable_async_pipeline_workers(apps, schema_editor): + _set_async_pipeline_workers(apps, schema_editor, enabled=False) + + +class Migration(migrations.Migration): + dependencies = [ + ("main", "0093_occurrence_project_score_index"), + ] + + operations = [ + migrations.RunPython(enable_async_pipeline_workers, disable_async_pipeline_workers), + ] diff --git a/ami/main/models.py b/ami/main/models.py index 8d65888bb..3662b4107 100644 --- a/ami/main/models.py +++ b/ami/main/models.py @@ -278,7 +278,7 @@ class ProjectFeatureFlags(pydantic.BaseModel): default_filters: bool = False # Whether to show default filters form in UI # Feature flag for jobs to reprocess all images in the project, even if already processed reprocess_all_images: bool = False - async_pipeline_workers: bool = False # Whether to use async pipeline workers that pull tasks from a queue + async_pipeline_workers: bool = True # Whether to use async pipeline workers that pull tasks from a queue def get_default_feature_flags() -> ProjectFeatureFlags: