Merge pull request #3525 from OpenNeuroOrg/taskiq-queue

Add Redis queue using taskiq to worker

Author: nellh

docker-compose.yml

@@ -1,5 +1,5 @@
 # docker compose versions
-version: "2.3"
+version: '2.3'
 
 services:
   # This dummy service provides shared configuration for all Node deps
@@ -30,7 +30,7 @@ services:
       retries: 10
     ports:
       # HMR port
-      - "9992:9992"
+      - '9992:9992'
     environment:
       - NODE_ENV=development
     depends_on:
@@ -76,7 +76,7 @@ services:
   # mongodb
   mongo:
     image: docker.io/library/mongo:8
-    command: ["--replSet", "rs0", "--bind_ip_all", "--port", "27017"]
+    command: ['--replSet', 'rs0', '--bind_ip_all', '--port', '27017']
     ports:
       - 27017:27017
     healthcheck:
@@ -105,35 +105,62 @@ services:
     init: true
     command:
       [
-        "uvicorn",
-        "--host",
-        "0.0.0.0",
-        "--port",
-        "80",
-        "--reload",
-        "--factory",
-        "datalad_service.app:create_app",
-        "--workers",
-        "8",
-        "--timeout-keep-alive",
-        "30",
-        "--log-level",
-        "debug",
+        'uvicorn',
+        '--host',
+        '0.0.0.0',
+        '--port',
+        '80',
+        '--reload',
+        '--factory',
+        'datalad_service.app:create_app',
+        '--workers',
+        '8',
+        '--timeout-keep-alive',
+        '30',
+        '--log-level',
+        'debug',
       ]
     networks:
       default:
         aliases:
           - datalad-0
           - datalad-1
 
+  taskiq:
+    build:
+      context: services/datalad
+    volumes:
+      - ${PERSISTENT_DIR}/datalad:/datalad:z
+      - ./services/datalad/datalad_service:/srv/datalad_service
+      - ./datalad-key:/datalad-key
+    env_file: ./config.env
+    init: true
+    command: [
+      'taskiq',
+      'worker',
+      'datalad_service.broker:broker',
+      '--tasks-pattern',
+      'datalad_service/tasks/*.py',
+      '--fs-discover',
+      '--reload'
+    ]
+    depends_on:
+      redis:
+        condition: service_started
+    networks:
+      default:
+        aliases:
+          - datalad-0
+          - datalad-1
+
   # nginx + app
   web:
     image: docker.io/library/nginx:1.16.1
     volumes:
       - ./nginx/nginx.dev.conf:/etc/nginx/conf.d/default.conf:ro
     ports:
-      - "8110:8110"
-      - "9876:80"
+      - '8110:8110'
+      - '9876:80'
     depends_on:
       server:
         condition: service_healthy
@@ -147,7 +174,7 @@ services:
     platform: ${ES_PLATFORM}
     environment:
       discovery.type: single-node
-      cluster.routing.allocation.disk.threshold_enabled: "true"
+      cluster.routing.allocation.disk.threshold_enabled: 'true'
       cluster.routing.allocation.disk.watermark.flood_stage: 1gb
       cluster.routing.allocation.disk.watermark.low: 10gb
       cluster.routing.allocation.disk.watermark.high: 5gb
@@ -158,10 +185,10 @@ services:
     security_opt:
       - seccomp=${SECOMP}
     healthcheck:
-      test: "curl -s -f http://localhost:9200 || exit 1"
+      test: 'curl -s -f http://localhost:9200 || exit 1'
       interval: 10s
       timeout: 5s
       retries: 3
     ports:
-      - "9200:9200"
-      - "9300:9300"
+      - '9200:9200'
+      - '9300:9300'

helm/openneuro/Chart.yaml

@@ -1,6 +1,6 @@
 apiVersion: v1
 name: openneuro
-version: 1.4.0
+version: 1.5.0
 description: OpenNeuro production deployment chart
 home: https://openneuro.org
 sources:

helm/openneuro/templates/dataset-worker-stateful-set.yaml

@@ -72,4 +72,22 @@ spec:
           periodSeconds: 30
           httpGet:
             path: '/heartbeat'
-            port: 80
+            port: 80
+      - name: {{ .Release.Name }}-dataset-taskiq-worker
+        image: 'openneuro/datalad-service:v{{ .Chart.AppVersion }}'
+        command: ['taskiq', 'worker', 'datalad_service.broker:broker', '--tasks-pattern', 'datalad_service/tasks/*.py', '--fs-discover']
+        resources:
+          requests:
+            cpu: {{ .Values.workerCpuRequests }}
+            memory: {{ .Values.workerMemoryRequests }}
+        envFrom:
+        - configMapRef:
+            name: {{ .Release.Name }}-configmap
+        - secretRef:
+            name: {{ .Release.Name }}-secret
+        volumeMounts:
+        - name: datasets
+          mountPath: /datasets
+        - name: ssh-key
+          mountPath: /datalad-key
+          subPath: datalad-key

services/datalad/datalad_service/broker/__init__.py

@@ -0,0 +1,20 @@
+import sys
+
+from taskiq import InMemoryBroker
+from taskiq_redis import RedisAsyncResultBackend, RedisStreamBroker
+
+from datalad_service import config
+from datalad_service.broker.get_docker_scale import get_docker_scale
+
+# Use InMemoryBroker to run during pytest
+if 'pytest' in sys.modules:
+    broker = InMemoryBroker()
+else:
+    redis_url = f'redis://{config.REDIS_HOST}:{config.REDIS_PORT}/{get_docker_scale()}'
+    result_backend = RedisAsyncResultBackend(
+        redis_url=redis_url,
+        result_ex_time=5000,
+    )
+    broker = RedisStreamBroker(
+        url=redis_url,
+    ).with_result_backend(result_backend)

services/datalad/datalad_service/broker/get_docker_scale.py

@@ -0,0 +1,24 @@
+#!/usr/bin/env python
+"""Helper script to get the docker-compose offset from within a container."""
+
+import re
+import socket
+import dns.resolver
+import dns.reversename
+
+
+def get_docker_scale():
+    hostname = socket.gethostname()
+    matches = re.match('.*?-([0-9]+)$', hostname)
+    if matches:
+        # Kubernetes
+        return int(matches[1])
+    else:
+        # docker-compose
+        ip_addr = socket.gethostbyname(hostname)
+        answer = dns.resolver.resolve(dns.reversename.from_address(ip_addr), 'PTR')
+        subdomain = str(answer[0]).split('.')[0]
+        # Subtract 1 to make this zero indexed (matching k8s)
+        docker_compose_offset = int(subdomain.split('_')[-1]) - 1
+        # Printed so this can be referenced as a script
+        return docker_compose_offset

services/datalad/datalad_service/config.py

@@ -25,3 +25,7 @@
 
 # Site URL
 CRN_SERVER_URL = os.getenv('CRN_SERVER_URL')
+
+# Redit connection for task queue
+REDIS_HOST = os.getenv('REDIS_HOST')
+REDIS_PORT = os.getenv('REDIS_PORT')

services/datalad/datalad_service/handlers/validation.py

@@ -1,4 +1,5 @@
 import asyncio
+import logging
 
 import falcon
 
@@ -18,13 +19,14 @@ async def on_post(self, req, resp, dataset, hexsha):
             try:
                 dataset_path = self.store.get_dataset_path(dataset)
                 # Run the validator but don't block on the request
-                asyncio.create_task(
-                    validate_dataset(
-                        dataset, dataset_path, hexsha, req.cookies, user=name
-                    )
+                await validate_dataset.kiq(
+                    dataset, dataset_path, hexsha, req.cookies, user=name
                 )
                 resp.status = falcon.HTTP_OK
-            except:
+            except Exception:
+                logging.exception(
+                    'Validation task enqueue failed for dataset %s', dataset
+                )
                 resp.status = falcon.HTTP_INTERNAL_SERVER_ERROR
         else:
             resp.media = {'error': 'Missing or malformed dataset parameter in request.'}

services/datalad/datalad_service/tasks/files.py

@@ -37,7 +37,7 @@ def commit_files(store, dataset, files, name=None, email=None, cookies=None):
     )
     ref = git_commit(repo, files, author)
     # Run the validator but don't block on the request
-    asyncio.create_task(validate_dataset(dataset, dataset_path, str(ref), cookies))
+    asyncio.create_task(validate_dataset.kiq(dataset, dataset_path, str(ref), cookies))
     return ref
 
 

services/datalad/datalad_service/tasks/validator.py

@@ -7,6 +7,7 @@
 import requests
 
 from datalad_service.config import GRAPHQL_ENDPOINT
+from datalad_service.broker import broker
 
 logger = logging.getLogger('datalad_service.' + __name__)
 
@@ -106,6 +107,7 @@ def issues_mutation(dataset_id, ref, issues, validator_metadata):
     }
 
 
+@broker.task
 async def validate_dataset(dataset_id, dataset_path, ref, cookies=None, user=''):
     # New schema validator second in case of issues
     validator_output_deno = await validate_dataset_deno_call(dataset_path, ref)

services/datalad/get_docker_scale.py

@@ -1,21 +1,22 @@
 #!/usr/bin/env python
 """Helper script to get the docker-compose offset from within a container."""
+
 import re
 import socket
 import dns.resolver
 import dns.reversename
 
 hostname = socket.gethostname()
-matches = re.match(".*?-([0-9]+)$", hostname)
+matches = re.match('.*?-([0-9]+)$', hostname)
 if matches:
-  # Kubernetes
-  print(matches[1])
+    # Kubernetes
+    print(matches[1])
 else:
-  # docker-compose
-  ip_addr = socket.gethostbyname(hostname)
-  answer = dns.resolver.query(dns.reversename.from_address(ip_addr), "PTR")
-  subdomain = str(answer[0]).split('.')[0]
-  # Subtract 1 to make this zero indexed (matching k8s)
-  docker_compose_offset = int(subdomain.split('_')[-1]) - 1
-  # Printed so this can be referenced as a script
-  print(docker_compose_offset)
+    # docker-compose
+    ip_addr = socket.gethostbyname(hostname)
+    answer = dns.resolver.resolve(dns.reversename.from_address(ip_addr), 'PTR')
+    subdomain = str(answer[0]).split('.')[0]
+    # Subtract 1 to make this zero indexed (matching k8s)
+    docker_compose_offset = int(subdomain.split('_')[-1]) - 1
+    # Printed so this can be referenced as a script
+    print(docker_compose_offset)