Local Pipeline api and worker (#876)

* First commit

* ruff

* tasks

* remap/pedigree paths

* add paths

* add model test

* Finish bad requests tests

* ruff

* a pattern for a queue

* cleanup

* Still hackin

* In progress

* almost there?

* lint

* lint worker

* update worker

* Remove force from local pipeline

* well

* cleanup

* finish pipeline worker

* ruff

* ruff

* liftover

* update worker

* ruff

* add crdqs

* remove force

Author: bpblanken

requirements.in

@@ -5,3 +5,5 @@ luigi>=3.4.0
 gnomad==0.6.4
 google-cloud-storage>=2.14.0
 google-cloud-secret-manager>=2.20.0
+aiofiles==24.1.0
+pydantic==2.8.2

v03_pipeline/api/app.py

@@ -1,17 +1,80 @@
-from aiohttp import web
+import json
+import os
+import traceback
 
-from v03_pipeline.lib.tasks import *  # noqa: F403
+import aiofiles
+import aiofiles.os
+from aiohttp import web, web_exceptions
+
+from v03_pipeline.api.model import LoadingPipelineRequest
+from v03_pipeline.lib.logger import get_logger
+from v03_pipeline.lib.paths import loading_pipeline_queue_path
+
+logger = get_logger(__name__)
+
+
+@web.middleware
+async def error_middleware(request, handler):
+    try:
+        return await handler(request)
+    except web.HTTPError:
+        logger.exception('HTTPError')
+        raise
+    except Exception as e:
+        logger.exception('Unhandled Exception')
+        error_reason = f'{e}: {traceback.format_exc()}'
+        raise web.HTTPInternalServerError(reason=error_reason) from e
+
+
+async def loading_pipeline_enqueue(request: web.Request) -> web.Response:
+    if not request.body_exists:
+        raise web.HTTPUnprocessableEntity
+
+    try:
+        lpr = LoadingPipelineRequest.model_validate(await request.json())
+    except ValueError as e:
+        raise web.HTTPBadRequest from e
+
+    try:
+        async with aiofiles.open(loading_pipeline_queue_path(), 'r') as f:
+            return web.json_response(
+                {
+                    'Failed to queue due to in process request': json.loads(
+                        await f.read(),
+                    ),
+                },
+                #
+                # The 409 (Conflict) status code indicates that the request
+                # could not be completed due to a conflict with the current
+                # state of the target resource.
+                #
+                status=web_exceptions.HTTPConflict.status_code,
+            )
+    except FileNotFoundError:
+        pass
+
+    async with aiofiles.open(loading_pipeline_queue_path(), 'w') as f:
+        await f.write(lpr.model_dump_json())
+    return web.json_response(
+        {'Successfully queued': lpr.model_dump()},
+        status=web_exceptions.HTTPAccepted.status_code,
+    )
 
 
 async def status(_: web.Request) -> web.Response:
     return web.json_response({'success': True})
 
 
 async def init_web_app():
-    app = web.Application()
+    await aiofiles.os.makedirs(
+        os.path.dirname(loading_pipeline_queue_path()),
+        exist_ok=True,
+    )
+    app = web.Application(middlewares=[error_middleware])
     app.add_routes(
         [
             web.get('/status', status),
+            web.post('/loading_pipeline_enqueue', loading_pipeline_enqueue),
         ],
     )
     return app

v03_pipeline/api/app_test.py

@@ -0,0 +1,111 @@
+from aiohttp import web_exceptions
+from aiohttp.test_utils import AioHTTPTestCase
+
+from v03_pipeline.api.app import init_web_app
+from v03_pipeline.lib.model import DatasetType, ReferenceGenome, SampleType
+from v03_pipeline.lib.test.mocked_dataroot_testcase import MockedDatarootTestCase
+
+CALLSET_PATH = 'v03_pipeline/var/test/callsets/1kg_30variants.vcf'
+
+
+class AppTest(AioHTTPTestCase, MockedDatarootTestCase):
+    async def get_application(self):
+        return await init_web_app()
+
+    async def test_status(self):
+        async with self.client.request('GET', '/status') as resp:
+            self.assertEqual(resp.status, 200)
+            resp_json = await resp.json()
+        self.assertDictEqual(resp_json, {'success': True})
+
+    async def test_missing_route(self):
+        with self.assertLogs(level='ERROR') as _:
+            async with self.client.request('GET', '/loading_pip') as resp:
+                self.assertEqual(resp.status, web_exceptions.HTTPNotFound.status_code)
+
+    async def test_loading_pipeline_invalid_requests(self):
+        with self.assertLogs(level='ERROR') as log:
+            async with self.client.request('GET', '/loading_pipeline_enqueue') as resp:
+                self.assertEqual(
+                    resp.status,
+                    web_exceptions.HTTPMethodNotAllowed.status_code,
+                )
+                self.assertTrue(
+                    'HTTPMethodNotAllowed' in log.output[0],
+                )
+
+        with self.assertLogs(level='ERROR') as log:
+            async with self.client.request('POST', '/loading_pipeline_enqueue') as resp:
+                self.assertEqual(
+                    resp.status,
+                    web_exceptions.HTTPUnprocessableEntity.status_code,
+                )
+                self.assertTrue(
+                    'HTTPUnprocessableEntity' in log.output[0],
+                )
+
+        body = {
+            'callset_path': 'missing.vcf',
+            'projects_to_run': ['project_a'],
+            'sample_type': SampleType.WGS.value,
+            'reference_genome': ReferenceGenome.GRCh38.value,
+            'dataset_type': DatasetType.SNV_INDEL.value,
+        }
+        with self.assertLogs(level='ERROR') as log:
+            async with self.client.request(
+                'POST',
+                '/loading_pipeline_enqueue',
+                json=body,
+            ) as resp:
+                self.assertEqual(
+                    resp.status,
+                    web_exceptions.HTTPBadRequest.status_code,
+                )
+                self.assertTrue(
+                    'callset_path must point to a file that exists' in log.output[0],
+                )
+
+    async def test_loading_pipeline_enqueue(self):
+        body = {
+            'callset_path': CALLSET_PATH,
+            'projects_to_run': ['project_a'],
+            'sample_type': SampleType.WGS.value,
+            'reference_genome': ReferenceGenome.GRCh38.value,
+            'dataset_type': DatasetType.SNV_INDEL.value,
+        }
+        async with self.client.request(
+            'POST',
+            '/loading_pipeline_enqueue',
+            json=body,
+        ) as resp:
+            self.assertEqual(
+                resp.status,
+                web_exceptions.HTTPAccepted.status_code,
+            )
+            resp_json = await resp.json()
+        self.assertDictEqual(
+            resp_json,
+            {
+                'Successfully queued': {
+                    'callset_path': 'v03_pipeline/var/test/callsets/1kg_30variants.vcf',
+                    'dataset_type': 'SNV_INDEL',
+                    'ignore_missing_samples_when_remapping': False,
+                    'projects_to_run': ['project_a'],
+                    'reference_genome': 'GRCh38',
+                    'sample_type': 'WGS',
+                    'skip_validation': False,
+                },
+            },
+        )
+
+        # Second request
+        body['projects_to_run'] = ['project_b', 'project_c']
+        async with self.client.request(
+            'POST',
+            '/loading_pipeline_enqueue',
+            json=body,
+        ) as resp:
+            self.assertEqual(
+                resp.status,
+                web_exceptions.HTTPConflict.status_code,
+            )

v03_pipeline/api/model.py

@@ -0,0 +1,27 @@
+import hailtop.fs as hfs
+from pydantic import BaseModel, Field, field_validator
+
+from v03_pipeline.lib.model import DatasetType, ReferenceGenome, SampleType
+
+VALID_FILE_TYPES = ['vcf', 'vcf.gz', 'vcf.bgz', 'mt']
+
+
+class LoadingPipelineRequest(BaseModel):
+    callset_path: str
+    projects_to_run: list[str] = Field(min_length=1, frozen=True)
+    sample_type: SampleType
+    reference_genome: ReferenceGenome
+    dataset_type: DatasetType
+    ignore_missing_samples_when_remapping: bool = False
+    skip_validation: bool = False
+
+    @field_validator('callset_path')
+    @classmethod
+    def check_valid_callset_path(cls, callset_path: str) -> str:
+        if not any(callset_path.endswith(file_type) for file_type in VALID_FILE_TYPES):
+            msg = 'callset_path must be a VCF or a Hail Matrix Table'
+            raise ValueError(msg)
+        if not hfs.exists(callset_path):
+            msg = 'callset_path must point to a file that exists'
+            raise ValueError(msg)
+        return callset_path

v03_pipeline/api/model_test.py

@@ -0,0 +1,35 @@
+import unittest
+
+from v03_pipeline.api.model import LoadingPipelineRequest
+from v03_pipeline.lib.model import DatasetType, ReferenceGenome, SampleType
+
+TEST_VCF = 'v03_pipeline/var/test/callsets/1kg_30variants.vcf'
+
+
+class ModelTest(unittest.TestCase):
+    def test_valid_loading_pipeline_requests(self) -> None:
+        raw_request = {
+            'callset_path': TEST_VCF,
+            'projects_to_run': ['project_a'],
+            'sample_type': SampleType.WGS.value,
+            'reference_genome': ReferenceGenome.GRCh38.value,
+            'dataset_type': DatasetType.SNV_INDEL.value,
+        }
+        lpr = LoadingPipelineRequest.model_validate(raw_request)
+        self.assertEqual(lpr.reference_genome, ReferenceGenome.GRCh38)
+
+    def test_invalid_loading_pipeline_requests(self) -> None:
+        raw_request = {
+            'callset_path': 'a.txt',
+            'projects_to_run': [],
+            'sample_type': 'BLENDED',
+            'reference_genome': ReferenceGenome.GRCh38.value,
+            'dataset_type': DatasetType.SNV_INDEL.value,
+        }
+        with self.assertRaises(ValueError) as cm:
+            LoadingPipelineRequest.model_validate(raw_request)
+        self.assertTrue(
+            str(cm.exception).startswith(
+                '3 validation errors for LoadingPipelineRequest',
+            ),
+        )

v03_pipeline/bin/pipeline_worker.py

@@ -0,0 +1,89 @@
+#!/usr/bin/env python3
+import datetime
+import os
+import time
+
+import luigi
+
+from v03_pipeline.api.model import LoadingPipelineRequest
+from v03_pipeline.lib.logger import get_logger
+from v03_pipeline.lib.paths import (
+    loading_pipeline_queue_path,
+    project_pedigree_path,
+    project_remap_path,
+)
+from v03_pipeline.lib.tasks import (
+    UpdateCachedReferenceDatasetQueries,
+    UpdateVariantAnnotationsTableWithNewSamplesTask,
+    WriteProjectFamilyTablesTask,
+)
+
+logger = get_logger(__name__)
+
+
+def main():
+    while True:
+        try:
+            if not os.path.exists(loading_pipeline_queue_path()):
+                continue
+            with open(loading_pipeline_queue_path()) as f:
+                lpr = LoadingPipelineRequest.model_validate_json(f.read())
+            project_remap_paths = [
+                project_remap_path(
+                    lpr.reference_genome,
+                    lpr.dataset_type,
+                    lpr.sample_type,
+                    project_guid,
+                )
+                for project_guid in lpr.projects_to_run
+            ]
+            project_pedigree_paths = [
+                project_pedigree_path(
+                    lpr.reference_genome,
+                    lpr.dataset_type,
+                    lpr.sample_type,
+                    project_guid,
+                )
+                for project_guid in lpr.projects_to_run
+            ]
+            task_kwargs = {
+                k: v for k, v in lpr.model_dump().items() if k != 'projects_to_run'
+            }
+            tasks = [
+                UpdateCachedReferenceDatasetQueries(
+                    reference_genome=lpr.reference_genome,
+                    dataset_type=lpr.dataset_type,
+                ),
+                UpdateVariantAnnotationsTableWithNewSamplesTask(
+                    project_guids=lpr.projects_to_run,
+                    project_remap_paths=project_remap_paths,
+                    project_pedigree_paths=project_pedigree_paths,
+                    run_id=datetime.datetime.now(datetime.timezone.utc).strftime(
+                        '%Y%m%d-%H%M%S',
+                    ),
+                    force=False,
+                    **task_kwargs,
+                ),
+                *[
+                    WriteProjectFamilyTablesTask(
+                        project_guid=lpr.projects_to_run[i],
+                        project_remap_path=project_remap_paths[i],
+                        project_pedigree_path=project_pedigree_paths[i],
+                        force=False,
+                        **task_kwargs,
+                    )
+                    for i in range(len(lpr.projects_to_run))
+                ],
+            ]
+            luigi.build(tasks)
+        except Exception:
+            logger.exception('Unhandled Exception')
+        finally:
+            if os.path.exists(loading_pipeline_queue_path()):
+                os.remove(loading_pipeline_queue_path())
+            logger.info('Waiting for work')
+            time.sleep(1)
+
+
+if __name__ == '__main__':
+    main()

v03_pipeline/lib/logger.py

@@ -28,6 +28,9 @@
         'urllib3': {
             'level': 'CRITICAL',
         },
+        'asyncio': {
+            'level': 'CRITICAL',
+        },
     },
 }
 

v03_pipeline/lib/model/cached_reference_dataset_query.py

@@ -13,7 +13,7 @@
 )
 
 
-class CachedReferenceDatasetQuery(Enum):
+class CachedReferenceDatasetQuery(str, Enum):
     CLINVAR_PATH_VARIANTS = 'clinvar_path_variants'
     GNOMAD_CODING_AND_NONCODING_VARIANTS = 'gnomad_coding_and_noncoding_variants'
     GNOMAD_QC = 'gnomad_qc'

v03_pipeline/lib/model/dataset_type.py

@@ -10,7 +10,7 @@
 ZERO = 0.0
 
 
-class DatasetType(Enum):
+class DatasetType(str, Enum):
     GCNV = 'GCNV'
     MITO = 'MITO'
     SNV_INDEL = 'SNV_INDEL'