Compare commits
35 Commits
master
...
pika_encap
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
251fea4094 | ||
|
|
d00a0aadc0 | ||
|
|
ded29bb8d0 | ||
|
|
de3899c69f | ||
|
|
48eb2bb792 | ||
|
|
38fb073291 | ||
|
|
8744cdeb6d | ||
|
|
0ce30cafbf | ||
|
|
5ac5fbf50b | ||
|
|
4c7c4b466e | ||
|
|
f8b60aad16 | ||
|
|
a40f07139a | ||
|
|
d3605b0913 | ||
|
|
9130b9fc75 | ||
|
|
1205adecdb | ||
|
|
59d0de97e3 | ||
|
|
f78a6894bc | ||
|
|
32740a914f | ||
|
|
f21cbc141e | ||
|
|
c415dcd8e6 | ||
|
|
ccad97992d | ||
|
|
b4b060ac0e | ||
|
|
66a27d26d3 | ||
|
|
4f9c5626b5 | ||
|
|
91980ff3ed | ||
|
|
963132afb1 | ||
|
|
896918570e | ||
|
|
f0ffc8fb8b | ||
|
|
8d6dd9e552 | ||
|
|
eadefc9b14 | ||
|
|
1347fb57f8 | ||
|
|
ce1217c191 | ||
|
|
9aff168f0d | ||
|
|
08f2746afc | ||
|
|
21592342d5 |
184
README.md
184
README.md
@ -6,154 +6,100 @@ The Infrastructure expects to be deployed in the same Pod / local environment as
|
|||||||
|
|
||||||
A configuration is located in `/config.yaml`. All relevant variables can be configured via exporting environment variables.
|
A configuration is located in `/config.yaml`. All relevant variables can be configured via exporting environment variables.
|
||||||
|
|
||||||
| Environment Variable | Default | Description |
|
| Environment Variable | Default | Description |
|
||||||
|-------------------------------|--------------------------------|--------------------------------------------------------------------------------------------------|
|
|-------------------------------|--------------------------------|-----------------------------------------------------------------------|
|
||||||
| _service_ | | |
|
| LOGGING_LEVEL_ROOT | DEBUG | Logging level for service logger |
|
||||||
| LOGGING_LEVEL_ROOT | DEBUG | Logging level for service logger |
|
| PROBING_WEBSERVER_HOST | "0.0.0.0" | Probe webserver address |
|
||||||
| RESPONSE_TYPE | "stream" | Whether the analysis response is stored as file on storage or sent as stream: "file" or "stream" |
|
| PROBING_WEBSERVER_PORT | 8080 | Probe webserver port |
|
||||||
| RESPONSE_FILE_EXTENSION | ".NER_ENTITIES.json.gz" | Extension to the file that stores the analyzed response on storage |
|
| PROBING_WEBSERVER_MODE | production | Webserver mode: {development, production} |
|
||||||
| _probing_webserver_ | | |
|
| RABBITMQ_HOST | localhost | RabbitMQ host address |
|
||||||
| PROBING_WEBSERVER_HOST | "0.0.0.0" | Probe webserver address |
|
| RABBITMQ_PORT | 5672 | RabbitMQ host port |
|
||||||
| PROBING_WEBSERVER_PORT | 8080 | Probe webserver port |
|
| RABBITMQ_USERNAME | user | RabbitMQ username |
|
||||||
| PROBING_WEBSERVER_MODE | production | Webserver mode: {development, production} |
|
| RABBITMQ_PASSWORD | bitnami | RabbitMQ password |
|
||||||
| _rabbitmq_ | | |
|
| RABBITMQ_HEARTBEAT | 7200 | Controls AMQP heartbeat timeout in seconds |
|
||||||
| RABBITMQ_HOST | localhost | RabbitMQ host address |
|
| REQUEST_QUEUE | request_queue | Requests to service |
|
||||||
| RABBITMQ_PORT | 5672 | RabbitMQ host port |
|
| RESPONSE_QUEUE | response_queue | Responses by service |
|
||||||
| RABBITMQ_USERNAME | user | RabbitMQ username |
|
| DEAD_LETTER_QUEUE | dead_letter_queue | Messages that failed to process |
|
||||||
| RABBITMQ_PASSWORD | bitnami | RabbitMQ password |
|
| ANALYSIS_ENDPOINT | "http://127.0.0.1:5000" | Endpoint for analysis container |
|
||||||
| RABBITMQ_HEARTBEAT | 7200 | Controls AMQP heartbeat timeout in seconds |
|
| STORAGE_BACKEND | s3 | The type of storage to use {s3, azure} |
|
||||||
| _queues_ | | |
|
| STORAGE_BUCKET | "pyinfra-test-bucket" | The bucket / container to pull files specified in queue requests from |
|
||||||
| REQUEST_QUEUE | request_queue | Requests to service |
|
| STORAGE_ENDPOINT | "http://127.0.0.1:9000" | Endpoint for s3 storage |
|
||||||
| RESPONSE_QUEUE | response_queue | Responses by service |
|
| STORAGE_KEY | root | User for s3 storage |
|
||||||
| DEAD_LETTER_QUEUE | dead_letter_queue | Messages that failed to process |
|
| STORAGE_SECRET | password | Password for s3 storage |
|
||||||
| _callback_ | | |
|
| STORAGE_AZURECONNECTIONSTRING | "DefaultEndpointsProtocol=..." | Connection string for Azure storage |
|
||||||
| RETRY | False | Toggles retry behaviour |
|
|
||||||
| MAX_ATTEMPTS | 3 | Number of times a message may fail before being published to dead letter queue |
|
|
||||||
| ANALYSIS_ENDPOINT | "http://127.0.0.1:5000" | |
|
|
||||||
| _storage_ | | |
|
|
||||||
| STORAGE_BACKEND | s3 | The type of storage to use {s3, azure} |
|
|
||||||
| STORAGE_BUCKET | "pyinfra-test-bucket" | The bucket / container to pull files specified in queue requests from |
|
|
||||||
| TARGET_FILE_EXTENSION | ".TEXT.json.gz" | Defines type of file to pull from storage: .TEXT.json.gz or .ORIGIN.pdf.gz |
|
|
||||||
| STORAGE_ENDPOINT | "http://127.0.0.1:9000" | |
|
|
||||||
| STORAGE_KEY | | |
|
|
||||||
| STORAGE_SECRET | | |
|
|
||||||
| STORAGE_AZURECONNECTIONSTRING | "DefaultEndpointsProtocol=..." | |
|
|
||||||
|
|
||||||
## Response Format
|
## Response Format
|
||||||
|
|
||||||
### RESPONSE_AS_FILE == False
|
### Expected AMQP input message:
|
||||||
|
|
||||||
Response-Format:
|
|
||||||
|
|
||||||
```json
|
```json
|
||||||
{
|
{
|
||||||
"dossierId": "klaus",
|
"dossierId": "",
|
||||||
"fileId": "1a7fd8ac0da7656a487b68f89188be82",
|
"fileId": "",
|
||||||
"imageMetadata": ANALYSIS_DATA
|
"targetFileExtension": "",
|
||||||
|
"responseFileExtension": ""
|
||||||
}
|
}
|
||||||
```
|
```
|
||||||
|
|
||||||
Response-example for image-prediction
|
Optionally, the input message can contain a field with the key `"operations"`.
|
||||||
|
|
||||||
|
### AMQP output message:
|
||||||
|
|
||||||
```json
|
```json
|
||||||
{
|
{
|
||||||
"dossierId": "klaus",
|
"dossierId": "",
|
||||||
"fileId": "1a7fd8ac0da7656a487b68f89188be82",
|
"fileId": "",
|
||||||
"imageMetadata": [
|
...
|
||||||
{
|
|
||||||
"classification": {
|
|
||||||
"label": "logo",
|
|
||||||
"probabilities": {
|
|
||||||
"formula": 0.0,
|
|
||||||
"logo": 1.0,
|
|
||||||
"other": 0.0,
|
|
||||||
"signature": 0.0
|
|
||||||
}
|
|
||||||
},
|
|
||||||
"filters": {
|
|
||||||
"allPassed": true,
|
|
||||||
"geometry": {
|
|
||||||
"imageFormat": {
|
|
||||||
"quotient": 1.570791527313267,
|
|
||||||
"tooTall": false,
|
|
||||||
"tooWide": false
|
|
||||||
},
|
|
||||||
"imageSize": {
|
|
||||||
"quotient": 0.19059804229011604,
|
|
||||||
"tooLarge": false,
|
|
||||||
"tooSmall": false
|
|
||||||
}
|
|
||||||
},
|
|
||||||
"probability": {
|
|
||||||
"unconfident": false
|
|
||||||
}
|
|
||||||
},
|
|
||||||
"geometry": {
|
|
||||||
"height": 107.63999999999999,
|
|
||||||
"width": 169.08000000000004
|
|
||||||
},
|
|
||||||
"position": {
|
|
||||||
"pageNumber": 1,
|
|
||||||
"x1": 213.12,
|
|
||||||
"x2": 382.20000000000005,
|
|
||||||
"y1": 568.7604,
|
|
||||||
"y2": 676.4004
|
|
||||||
}
|
|
||||||
}
|
|
||||||
]
|
|
||||||
}
|
}
|
||||||
```
|
```
|
||||||
|
|
||||||
### RESPONSE_AS_FILE == True
|
|
||||||
|
|
||||||
Creates a respone file on the request storage, named `dossier_Id / file_Id + RESPONSE_FILE_EXTENSION` with the `ANALYSIS_DATA` as content.
|
|
||||||
|
|
||||||
## Development
|
## Development
|
||||||
|
|
||||||
### Local Setup
|
Either run `src/serve.py` or the built Docker image.
|
||||||
|
|
||||||
You can run the infrastructure either as module via. `src/serve.py` or as Dockercontainer simulating the kubernetes environment
|
### Setup
|
||||||
|
|
||||||
1. Install module / build docker image
|
Install module.
|
||||||
|
|
||||||
```bash
|
```bash
|
||||||
pip install -e .
|
pip install -e .
|
||||||
pip install -r requirements.txt
|
pip install -r requirements.txt
|
||||||
```
|
```
|
||||||
|
|
||||||
```bash
|
or build docker image.
|
||||||
docker build -f Dockerfile -t pyinfra .
|
|
||||||
```
|
|
||||||
|
|
||||||
2. Run rabbitmq & minio
|
```bash
|
||||||
|
docker build -f Dockerfile -t pyinfra .
|
||||||
|
```
|
||||||
|
|
||||||
```bash
|
### Usage
|
||||||
docker-compose up
|
|
||||||
```
|
|
||||||
|
|
||||||
3. Run module
|
**Shell 1:** Start a MinIO and a RabbitMQ docker container.
|
||||||
|
|
||||||
```bash
|
```bash
|
||||||
python src/serve.py
|
docker-compose up
|
||||||
```
|
```
|
||||||
OR as container:
|
|
||||||
```bash
|
|
||||||
docker run --net=host pyinfra
|
|
||||||
```
|
|
||||||
|
|
||||||
Start your prediction container for example ner-prediction or image-prediction (follow their corresponding README for
|
**Shell 2:** Add files to the local minio storage.
|
||||||
building the container).
|
|
||||||
|
|
||||||
To put a file on the queue do:
|
```bash
|
||||||
|
python scripts/manage_minio.py add <MinIO target folder> -d path/to/a/folder/with/PDFs
|
||||||
|
```
|
||||||
|
|
||||||
python src/manage_minio.py add --file path/to/file dossierID
|
**Shell 2:** Run pyinfra-server.
|
||||||
|
|
||||||
To start mock:
|
```bash
|
||||||
|
python src/serve.py
|
||||||
|
```
|
||||||
|
or as container:
|
||||||
|
|
||||||
python src/mock_client.py
|
```bash
|
||||||
|
docker run --net=host pyinfra
|
||||||
|
```
|
||||||
|
|
||||||
### Hints:
|
**Shell 3:** Run analysis-container.
|
||||||
When stopping the docker-compose up, use docker-compose down to remove containers created by up.
|
|
||||||
|
|
||||||
If uploaded files are stuck, clean the minio storage by using ```python src/manage_minio.py purge``` or delete local
|
**Shell 4:** Start a client that sends requests to process PDFs from the MinIO store and annotates these PDFs according to the service responses.
|
||||||
minio data folder in pyinfra with ```sudo rm -rf data```
|
```bash
|
||||||
|
python scripts/mock_client.py
|
||||||
|
```
|
||||||
|
|||||||
14
config.yaml
14
config.yaml
@ -1,9 +1,5 @@
|
|||||||
service:
|
service:
|
||||||
logging_level: $LOGGING_LEVEL_ROOT|DEBUG # Logging level for service logger
|
logging_level: $LOGGING_LEVEL_ROOT|DEBUG # Logging level for service logger
|
||||||
response:
|
|
||||||
type: $RESPONSE_TYPE|"file" # Whether the analysis response is stored as file on storage or sent as stream
|
|
||||||
extension: $RESPONSE_FILE_EXTENSION|".IMAGE_INFO.json.gz" # {.IMAGE_INFO.json.gz | .NER_ENTITIES.json.gz}
|
|
||||||
key: $RESPONSE_KEY|"imageMetadata" # the key of the result {result, imageMetadata}
|
|
||||||
|
|
||||||
probing_webserver:
|
probing_webserver:
|
||||||
host: $PROBING_WEBSERVER_HOST|"0.0.0.0" # Probe webserver address
|
host: $PROBING_WEBSERVER_HOST|"0.0.0.0" # Probe webserver address
|
||||||
@ -22,22 +18,12 @@ rabbitmq:
|
|||||||
output: $RESPONSE_QUEUE|response_queue # Responses by service
|
output: $RESPONSE_QUEUE|response_queue # Responses by service
|
||||||
dead_letter: $DEAD_LETTER_QUEUE|dead_letter_queue # Messages that failed to process
|
dead_letter: $DEAD_LETTER_QUEUE|dead_letter_queue # Messages that failed to process
|
||||||
|
|
||||||
prefetch_count: 1
|
|
||||||
|
|
||||||
callback:
|
callback:
|
||||||
|
|
||||||
retry: # Controls retry behaviour for messages the processing of which failed
|
|
||||||
# TODO: check if this actually works
|
|
||||||
enabled: $RETRY|False # Toggles retry behaviour
|
|
||||||
max_attempts: $MAX_ATTEMPTS|3 # Number of times a message may fail before being published to dead letter queue
|
|
||||||
|
|
||||||
analysis_endpoint: $ANALYSIS_ENDPOINT|"http://127.0.0.1:5000"
|
analysis_endpoint: $ANALYSIS_ENDPOINT|"http://127.0.0.1:5000"
|
||||||
|
|
||||||
storage:
|
storage:
|
||||||
|
|
||||||
backend: $STORAGE_BACKEND|s3 # The type of storage to use {s3, azure}
|
backend: $STORAGE_BACKEND|s3 # The type of storage to use {s3, azure}
|
||||||
bucket: $STORAGE_BUCKET|"pyinfra-test-bucket" # The bucket / container to pull files specified in queue requests from
|
bucket: $STORAGE_BUCKET|"pyinfra-test-bucket" # The bucket / container to pull files specified in queue requests from
|
||||||
target_file_extension: $TARGET_FILE_EXTENSION|".ORIGIN.pdf.gz" # {.TEXT.json.gz | .ORIGIN.pdf.gz} Defines type of file to pull from storage
|
|
||||||
|
|
||||||
s3:
|
s3:
|
||||||
endpoint: $STORAGE_ENDPOINT|"http://127.0.0.1:9000"
|
endpoint: $STORAGE_ENDPOINT|"http://127.0.0.1:9000"
|
||||||
|
|||||||
@ -1,96 +0,0 @@
|
|||||||
import json
|
|
||||||
import logging
|
|
||||||
import tempfile
|
|
||||||
from time import sleep
|
|
||||||
from pyinfra.config import CONFIG
|
|
||||||
|
|
||||||
from pyinfra.exceptions import AnalysisFailure, DataLoadingFailure
|
|
||||||
from pyinfra.rabbitmq import make_connection, make_channel, declare_queue
|
|
||||||
from pyinfra.storage.storages import get_storage
|
|
||||||
from pyinfra.utils.file import upload_compressed_response
|
|
||||||
|
|
||||||
|
|
||||||
def make_retry_callback(republish, max_attempts):
|
|
||||||
def get_n_previous_attempts(props):
|
|
||||||
return 0 if props.headers is None else props.headers.get("x-retry-count", 0)
|
|
||||||
|
|
||||||
def attempts_remain(n_attempts):
|
|
||||||
return n_attempts < max_attempts
|
|
||||||
|
|
||||||
def callback(channel, method, properties, body):
|
|
||||||
|
|
||||||
n_attempts = get_n_previous_attempts(properties) + 1
|
|
||||||
|
|
||||||
logging.error(f"Message failed to process {n_attempts}/{max_attempts} times: {body}")
|
|
||||||
|
|
||||||
if attempts_remain(n_attempts):
|
|
||||||
republish(channel, body, n_attempts)
|
|
||||||
channel.basic_ack(delivery_tag=method.delivery_tag)
|
|
||||||
|
|
||||||
else:
|
|
||||||
logging.exception(f"Adding to dead letter queue: {body}")
|
|
||||||
channel.basic_reject(delivery_tag=method.delivery_tag, requeue=False)
|
|
||||||
|
|
||||||
return callback
|
|
||||||
|
|
||||||
|
|
||||||
def wrap_callback_in_retry_logic(callback, retry_callback):
|
|
||||||
def wrapped_callback(channel, method, properties, body):
|
|
||||||
try:
|
|
||||||
callback(channel, method, properties, body)
|
|
||||||
except (AnalysisFailure, DataLoadingFailure):
|
|
||||||
sleep(5)
|
|
||||||
retry_callback(channel, method, properties, body)
|
|
||||||
|
|
||||||
return wrapped_callback
|
|
||||||
|
|
||||||
|
|
||||||
def json_wrap(body_processor):
|
|
||||||
def inner(payload):
|
|
||||||
return json.dumps(body_processor(json.loads(payload)))
|
|
||||||
|
|
||||||
return inner
|
|
||||||
|
|
||||||
|
|
||||||
def make_callback_for_output_queue(json_wrapped_body_processor, output_queue_name):
|
|
||||||
|
|
||||||
connection = make_connection()
|
|
||||||
channel = make_channel(connection)
|
|
||||||
declare_queue(channel, output_queue_name)
|
|
||||||
|
|
||||||
def callback(channel, method, _, body):
|
|
||||||
"""
|
|
||||||
response is dossier_id, file_id and analysis result, if CONFIG.service.response.type == "file" the response only
|
|
||||||
contains file_id and dossier_id and the analysis result will be written in a json file
|
|
||||||
Args:
|
|
||||||
channel:
|
|
||||||
method:
|
|
||||||
_:
|
|
||||||
body:
|
|
||||||
|
|
||||||
Returns:
|
|
||||||
|
|
||||||
"""
|
|
||||||
|
|
||||||
dossier_id, file_id, result = json_wrapped_body_processor(body)
|
|
||||||
|
|
||||||
result_key = CONFIG.service.response.key if CONFIG.service.response.key else "result"
|
|
||||||
result = json.dumps({"dossierId": dossier_id, "fileId": file_id, result_key: result})
|
|
||||||
|
|
||||||
if CONFIG.service.response.type == "file":
|
|
||||||
upload_compressed_response(
|
|
||||||
get_storage(CONFIG.storage.backend), CONFIG.storage.bucket, dossier_id, file_id, result
|
|
||||||
)
|
|
||||||
result = json.dumps({"dossierId": dossier_id, "fileId": file_id})
|
|
||||||
|
|
||||||
channel.basic_publish(exchange="", routing_key=output_queue_name, body=result)
|
|
||||||
channel.basic_ack(delivery_tag=method.delivery_tag)
|
|
||||||
|
|
||||||
return callback
|
|
||||||
|
|
||||||
|
|
||||||
def make_retry_callback_for_output_queue(json_wrapped_body_processor, output_queue_name, retry_callback):
|
|
||||||
callback = make_callback_for_output_queue(json_wrapped_body_processor, output_queue_name)
|
|
||||||
callback = wrap_callback_in_retry_logic(callback, retry_callback)
|
|
||||||
|
|
||||||
return callback
|
|
||||||
@ -1,11 +1,23 @@
|
|||||||
"""Implements a config object with dot-indexing syntax."""
|
"""Implements a config object with dot-indexing syntax."""
|
||||||
|
|
||||||
|
|
||||||
from envyaml import EnvYAML
|
from envyaml import EnvYAML
|
||||||
|
|
||||||
from pyinfra.locations import CONFIG_FILE
|
from pyinfra.locations import CONFIG_FILE
|
||||||
|
|
||||||
|
|
||||||
|
def make_art():
|
||||||
|
return """
|
||||||
|
______ _____ __
|
||||||
|
| ___ \ |_ _| / _|
|
||||||
|
| |_/ / _ | | _ __ | |_ _ __ __ _
|
||||||
|
| __/ | | || || '_ \| _| '__/ _` |
|
||||||
|
| | | |_| || || | | | | | | | (_| |
|
||||||
|
\_| \__, \___/_| |_|_| |_| \__,_|
|
||||||
|
__/ |
|
||||||
|
|___/
|
||||||
|
"""
|
||||||
|
|
||||||
|
|
||||||
def _get_item_and_maybe_make_dotindexable(container, item):
|
def _get_item_and_maybe_make_dotindexable(container, item):
|
||||||
ret = container[item]
|
ret = container[item]
|
||||||
return DotIndexable(ret) if isinstance(ret, dict) else ret
|
return DotIndexable(ret) if isinstance(ret, dict) else ret
|
||||||
|
|||||||
@ -1,43 +0,0 @@
|
|||||||
import logging
|
|
||||||
from typing import Callable
|
|
||||||
|
|
||||||
import pika
|
|
||||||
from retry import retry
|
|
||||||
|
|
||||||
from pyinfra.exceptions import ProcessingFailure
|
|
||||||
from pyinfra.rabbitmq import make_connection, make_channel, declare_queue
|
|
||||||
|
|
||||||
|
|
||||||
class ConsumerError(Exception):
|
|
||||||
pass
|
|
||||||
|
|
||||||
|
|
||||||
@retry(ConsumerError, tries=3, delay=5, jitter=(1, 3))
|
|
||||||
def consume(queue_name: str, on_message_callback: Callable):
|
|
||||||
|
|
||||||
connection = make_connection()
|
|
||||||
channel = make_channel(connection)
|
|
||||||
declare_queue(channel, queue_name)
|
|
||||||
logging.info("Started infrastructure.")
|
|
||||||
|
|
||||||
while True:
|
|
||||||
try:
|
|
||||||
logging.info("Waiting for messages...")
|
|
||||||
channel.basic_consume(queue=queue_name, auto_ack=False, on_message_callback=on_message_callback)
|
|
||||||
channel.start_consuming()
|
|
||||||
|
|
||||||
except pika.exceptions.ConnectionClosedByBroker as err:
|
|
||||||
logging.critical(f"Caught a channel error: {err}, stopping.")
|
|
||||||
break
|
|
||||||
|
|
||||||
except pika.exceptions.AMQPChannelError as err:
|
|
||||||
logging.critical(f"Caught a channel error: {err}, stopping.")
|
|
||||||
break
|
|
||||||
|
|
||||||
except pika.exceptions.AMQPConnectionError as err:
|
|
||||||
logging.info("No AMPQ-connection found, retrying...")
|
|
||||||
logging.debug(err)
|
|
||||||
continue
|
|
||||||
|
|
||||||
except ProcessingFailure as err:
|
|
||||||
raise ConsumerError(f"Error while consuming {queue_name}.") from err
|
|
||||||
@ -1,70 +0,0 @@
|
|||||||
import gzip
|
|
||||||
import json
|
|
||||||
import logging
|
|
||||||
from operator import itemgetter
|
|
||||||
|
|
||||||
import requests
|
|
||||||
|
|
||||||
from pyinfra.config import CONFIG
|
|
||||||
from pyinfra.exceptions import DataLoadingFailure, AnalysisFailure, ProcessingFailure
|
|
||||||
from pyinfra.utils.file import combine_dossier_id_and_file_id_and_extension
|
|
||||||
|
|
||||||
|
|
||||||
def make_storage_data_loader(storage, bucket_name):
|
|
||||||
def get_object_name(payload: dict) -> str:
|
|
||||||
dossier_id, file_id = itemgetter("dossierId", "fileId")(payload)
|
|
||||||
object_name = combine_dossier_id_and_file_id_and_extension(
|
|
||||||
dossier_id, file_id, CONFIG.storage.target_file_extension
|
|
||||||
)
|
|
||||||
return object_name
|
|
||||||
|
|
||||||
def download(payload):
|
|
||||||
object_name = get_object_name(payload)
|
|
||||||
logging.debug(f"Downloading {object_name}...")
|
|
||||||
data = storage.get_object(bucket_name, object_name)
|
|
||||||
logging.debug(f"Downloaded {object_name}.")
|
|
||||||
return data
|
|
||||||
|
|
||||||
def decompress(data):
|
|
||||||
return gzip.decompress(data)
|
|
||||||
|
|
||||||
def load_data(payload):
|
|
||||||
try:
|
|
||||||
return decompress(download(payload))
|
|
||||||
except Exception as err:
|
|
||||||
logging.warning(f"Loading data from storage failed for {payload}.")
|
|
||||||
raise DataLoadingFailure() from err
|
|
||||||
|
|
||||||
return load_data
|
|
||||||
|
|
||||||
|
|
||||||
def make_analyzer(analysis_endpoint):
|
|
||||||
def analyze(data):
|
|
||||||
try:
|
|
||||||
logging.debug(f"Requesting analysis from {analysis_endpoint}...")
|
|
||||||
analysis_response = requests.post(analysis_endpoint, data=data)
|
|
||||||
analysis_response.raise_for_status()
|
|
||||||
analysis_response = analysis_response.json()
|
|
||||||
logging.debug(f"Received response.")
|
|
||||||
return analysis_response
|
|
||||||
except Exception as err:
|
|
||||||
logging.warning("Exception caught when calling analysis endpoint.")
|
|
||||||
raise AnalysisFailure() from err
|
|
||||||
|
|
||||||
return analyze
|
|
||||||
|
|
||||||
|
|
||||||
def make_payload_processor(load_data, analyze_file):
|
|
||||||
def process(payload: dict):
|
|
||||||
logging.info(f"Processing {payload}...")
|
|
||||||
try:
|
|
||||||
payload = json.loads(payload)
|
|
||||||
dossier_id, file_id = itemgetter("dossierId", "fileId")(payload)
|
|
||||||
data = load_data(payload)
|
|
||||||
predictions = analyze_file(data)
|
|
||||||
return dossier_id, file_id, predictions
|
|
||||||
except (DataLoadingFailure, AnalysisFailure) as err:
|
|
||||||
logging.warning(f"Processing of {payload} failed.")
|
|
||||||
raise ProcessingFailure() from err
|
|
||||||
|
|
||||||
return process
|
|
||||||
@ -20,3 +20,7 @@ class InvalidEndpoint(ValueError):
|
|||||||
|
|
||||||
class UnknownClient(ValueError):
|
class UnknownClient(ValueError):
|
||||||
pass
|
pass
|
||||||
|
|
||||||
|
|
||||||
|
class ConsumerError(Exception):
|
||||||
|
pass
|
||||||
|
|||||||
@ -1,11 +1,12 @@
|
|||||||
|
import requests
|
||||||
from flask import Flask, jsonify
|
from flask import Flask, jsonify
|
||||||
|
from retry import retry
|
||||||
from waitress import serve
|
from waitress import serve
|
||||||
|
|
||||||
from pyinfra.config import CONFIG
|
from pyinfra.config import CONFIG
|
||||||
|
|
||||||
|
|
||||||
def run_probing_webserver(app, host=None, port=None, mode=None):
|
def run_probing_webserver(app, host=None, port=None, mode=None):
|
||||||
|
|
||||||
if not host:
|
if not host:
|
||||||
host = CONFIG.probing_webserver.host
|
host = CONFIG.probing_webserver.host
|
||||||
|
|
||||||
@ -23,7 +24,6 @@ def run_probing_webserver(app, host=None, port=None, mode=None):
|
|||||||
|
|
||||||
|
|
||||||
def set_up_probing_webserver():
|
def set_up_probing_webserver():
|
||||||
|
|
||||||
app = Flask(__name__)
|
app = Flask(__name__)
|
||||||
|
|
||||||
@app.route("/ready", methods=["GET"])
|
@app.route("/ready", methods=["GET"])
|
||||||
@ -38,4 +38,15 @@ def set_up_probing_webserver():
|
|||||||
resp.status_code = 200
|
resp.status_code = 200
|
||||||
return resp
|
return resp
|
||||||
|
|
||||||
|
@app.route("/prometheus", methods=["GET"])
|
||||||
|
def get_analysis_prometheus_endpoint():
|
||||||
|
@retry(requests.exceptions.ConnectionError, tries=3, delay=5, jitter=(1, 3))
|
||||||
|
def inner():
|
||||||
|
prom_endpoint = f"{CONFIG.rabbitmq.callback.analysis_endpoint}/prometheus"
|
||||||
|
metric = requests.get(prom_endpoint)
|
||||||
|
metric.raise_for_status()
|
||||||
|
return metric.text
|
||||||
|
|
||||||
|
return inner()
|
||||||
|
|
||||||
return app
|
return app
|
||||||
|
|||||||
13
pyinfra/queue/consumer.py
Normal file
13
pyinfra/queue/consumer.py
Normal file
@ -0,0 +1,13 @@
|
|||||||
|
from pyinfra.queue.queue_manager.queue_manager import QueueManager
|
||||||
|
|
||||||
|
|
||||||
|
class Consumer:
|
||||||
|
def __init__(self, callback, queue_manager: QueueManager):
|
||||||
|
self.queue_manager = queue_manager
|
||||||
|
self.callback = callback
|
||||||
|
|
||||||
|
def consume_and_publish(self):
|
||||||
|
self.queue_manager.consume_and_publish(self.callback)
|
||||||
|
|
||||||
|
def consume(self):
|
||||||
|
yield from self.queue_manager.consume()
|
||||||
0
pyinfra/queue/queue_manager/__init__.py
Normal file
0
pyinfra/queue/queue_manager/__init__.py
Normal file
142
pyinfra/queue/queue_manager/pika_queue_manager.py
Normal file
142
pyinfra/queue/queue_manager/pika_queue_manager.py
Normal file
@ -0,0 +1,142 @@
|
|||||||
|
import json
|
||||||
|
import logging
|
||||||
|
|
||||||
|
import pika
|
||||||
|
|
||||||
|
from pyinfra.config import CONFIG
|
||||||
|
from pyinfra.exceptions import ProcessingFailure
|
||||||
|
from pyinfra.queue.queue_manager.queue_manager import QueueHandle, QueueManager
|
||||||
|
|
||||||
|
logger = logging.getLogger("pika")
|
||||||
|
logger.setLevel(logging.WARNING)
|
||||||
|
|
||||||
|
logger = logging.getLogger(__name__)
|
||||||
|
logger.setLevel(CONFIG.service.logging_level)
|
||||||
|
|
||||||
|
|
||||||
|
def monkey_patch_queue_handle(channel, queue) -> QueueHandle:
|
||||||
|
|
||||||
|
empty_message = (None, None, None)
|
||||||
|
|
||||||
|
def is_empty_message(message):
|
||||||
|
return message == empty_message
|
||||||
|
|
||||||
|
queue_handle = QueueHandle()
|
||||||
|
queue_handle.empty = lambda: is_empty_message(channel.basic_get(queue))
|
||||||
|
|
||||||
|
def produce_items():
|
||||||
|
|
||||||
|
while True:
|
||||||
|
message = channel.basic_get(queue)
|
||||||
|
|
||||||
|
if is_empty_message(message):
|
||||||
|
break
|
||||||
|
|
||||||
|
method_frame, properties, body = message
|
||||||
|
channel.basic_ack(method_frame.delivery_tag)
|
||||||
|
yield json.loads(body)
|
||||||
|
|
||||||
|
queue_handle.to_list = lambda: list(produce_items())
|
||||||
|
|
||||||
|
return queue_handle
|
||||||
|
|
||||||
|
|
||||||
|
def get_connection():
|
||||||
|
|
||||||
|
credentials = pika.PlainCredentials(username=CONFIG.rabbitmq.user, password=CONFIG.rabbitmq.password)
|
||||||
|
|
||||||
|
kwargs = {"host": CONFIG.rabbitmq.host, "port": CONFIG.rabbitmq.port, "credentials": credentials}
|
||||||
|
|
||||||
|
parameters = pika.ConnectionParameters(**kwargs)
|
||||||
|
|
||||||
|
connection = pika.BlockingConnection(parameters=parameters)
|
||||||
|
|
||||||
|
return connection
|
||||||
|
|
||||||
|
|
||||||
|
def get_n_previous_attempts(props):
|
||||||
|
return 0 if props.headers is None else props.headers.get("x-retry-count", 0)
|
||||||
|
|
||||||
|
|
||||||
|
def attempts_remain(n_attempts, max_attempts):
|
||||||
|
return n_attempts < max_attempts
|
||||||
|
|
||||||
|
|
||||||
|
class PikaQueueManager(QueueManager):
|
||||||
|
def __init__(self, input_queue, output_queue, dead_letter_queue=None):
|
||||||
|
super().__init__(input_queue, output_queue)
|
||||||
|
|
||||||
|
self.connection = get_connection()
|
||||||
|
self.channel = self.connection.channel()
|
||||||
|
|
||||||
|
if not dead_letter_queue:
|
||||||
|
dead_letter_queue = CONFIG.rabbitmq.queues.dead_letter
|
||||||
|
|
||||||
|
args = {"x-dead-letter-exchange": "", "x-dead-letter-routing-key": dead_letter_queue}
|
||||||
|
|
||||||
|
self.channel.queue_declare(input_queue, arguments=args, auto_delete=False)
|
||||||
|
self.channel.queue_declare(output_queue, arguments=args, auto_delete=False)
|
||||||
|
|
||||||
|
def republish(self, body, n_current_attempts, frame):
|
||||||
|
self.channel.basic_publish(
|
||||||
|
exchange="",
|
||||||
|
routing_key=self._input_queue,
|
||||||
|
body=body,
|
||||||
|
properties=pika.BasicProperties(headers={"x-retry-count": n_current_attempts}),
|
||||||
|
)
|
||||||
|
self.channel.basic_ack(delivery_tag=frame.delivery_tag)
|
||||||
|
|
||||||
|
def publish_request(self, request):
|
||||||
|
logger.debug(f"Publishing {request}")
|
||||||
|
self.channel.basic_publish("", self._input_queue, json.dumps(request).encode())
|
||||||
|
|
||||||
|
def reject(self, body, frame):
|
||||||
|
logger.exception(f"Adding to dead letter queue: {body}")
|
||||||
|
self.channel.basic_reject(delivery_tag=frame.delivery_tag, requeue=False)
|
||||||
|
|
||||||
|
def publish_response(self, message, callback, max_attempts=3):
|
||||||
|
|
||||||
|
logger.debug(f"Publishing response for {message}.")
|
||||||
|
|
||||||
|
frame, properties, body = message
|
||||||
|
|
||||||
|
n_attempts = get_n_previous_attempts(properties) + 1
|
||||||
|
|
||||||
|
try:
|
||||||
|
response = json.dumps(callback(json.loads(body)))
|
||||||
|
self.channel.basic_publish("", self._output_queue, response.encode())
|
||||||
|
self.channel.basic_ack(frame.delivery_tag)
|
||||||
|
except ProcessingFailure:
|
||||||
|
|
||||||
|
logger.error(f"Message failed to process {n_attempts}/{max_attempts} times: {body}")
|
||||||
|
|
||||||
|
if attempts_remain(n_attempts, max_attempts):
|
||||||
|
self.republish(body, n_attempts, frame)
|
||||||
|
else:
|
||||||
|
self.reject(body, frame)
|
||||||
|
|
||||||
|
def pull_request(self):
|
||||||
|
return self.channel.basic_get(self._input_queue)
|
||||||
|
|
||||||
|
def consume(self):
|
||||||
|
logger.debug("Consuming")
|
||||||
|
return self.channel.consume(self._input_queue)
|
||||||
|
|
||||||
|
def consume_and_publish(self, visitor):
|
||||||
|
|
||||||
|
logger.info(f"Consuming with callback {visitor.callback.__name__}")
|
||||||
|
|
||||||
|
for message in self.consume():
|
||||||
|
self.publish_response(message, visitor)
|
||||||
|
|
||||||
|
def clear(self):
|
||||||
|
self.channel.queue_purge(self._input_queue)
|
||||||
|
self.channel.queue_purge(self._output_queue)
|
||||||
|
|
||||||
|
@property
|
||||||
|
def input_queue(self) -> QueueHandle:
|
||||||
|
return monkey_patch_queue_handle(self.channel, self._input_queue)
|
||||||
|
|
||||||
|
@property
|
||||||
|
def output_queue(self) -> QueueHandle:
|
||||||
|
return monkey_patch_queue_handle(self.channel, self._output_queue)
|
||||||
47
pyinfra/queue/queue_manager/queue_manager.py
Normal file
47
pyinfra/queue/queue_manager/queue_manager.py
Normal file
@ -0,0 +1,47 @@
|
|||||||
|
import abc
|
||||||
|
|
||||||
|
|
||||||
|
class QueueHandle:
|
||||||
|
def empty(self) -> bool:
|
||||||
|
raise NotImplemented()
|
||||||
|
|
||||||
|
def to_list(self) -> list:
|
||||||
|
raise NotImplemented()
|
||||||
|
|
||||||
|
|
||||||
|
class QueueManager(abc.ABC):
|
||||||
|
def __init__(self, input_queue, output_queue):
|
||||||
|
self._input_queue = input_queue
|
||||||
|
self._output_queue = output_queue
|
||||||
|
|
||||||
|
@abc.abstractmethod
|
||||||
|
def publish_request(self, request):
|
||||||
|
pass
|
||||||
|
|
||||||
|
@abc.abstractmethod
|
||||||
|
def publish_response(self, response, callback):
|
||||||
|
pass
|
||||||
|
|
||||||
|
@abc.abstractmethod
|
||||||
|
def pull_request(self):
|
||||||
|
pass
|
||||||
|
|
||||||
|
@abc.abstractmethod
|
||||||
|
def consume(self):
|
||||||
|
pass
|
||||||
|
|
||||||
|
@abc.abstractmethod
|
||||||
|
def clear(self):
|
||||||
|
pass
|
||||||
|
|
||||||
|
@abc.abstractmethod
|
||||||
|
def input_queue(self) -> QueueHandle:
|
||||||
|
pass
|
||||||
|
|
||||||
|
@abc.abstractmethod
|
||||||
|
def output_queue(self) -> QueueHandle:
|
||||||
|
pass
|
||||||
|
|
||||||
|
@abc.abstractmethod
|
||||||
|
def consume_and_publish(self, callback):
|
||||||
|
pass
|
||||||
@ -1,31 +0,0 @@
|
|||||||
import pika
|
|
||||||
|
|
||||||
from pyinfra.config import CONFIG
|
|
||||||
|
|
||||||
|
|
||||||
def make_channel(connection) -> pika.adapters.blocking_connection.BlockingChannel:
|
|
||||||
channel = connection.channel()
|
|
||||||
channel.basic_qos(prefetch_count=CONFIG.rabbitmq.prefetch_count)
|
|
||||||
return channel
|
|
||||||
|
|
||||||
|
|
||||||
def declare_queue(channel, queue: str):
|
|
||||||
args = {"x-dead-letter-exchange": "", "x-dead-letter-routing-key": CONFIG.rabbitmq.queues.dead_letter}
|
|
||||||
return channel.queue_declare(queue=queue, auto_delete=False, arguments=args, durable=True)
|
|
||||||
|
|
||||||
|
|
||||||
def read_connection_params():
|
|
||||||
credentials = pika.PlainCredentials(CONFIG.rabbitmq.user, CONFIG.rabbitmq.password)
|
|
||||||
parameters = pika.ConnectionParameters(
|
|
||||||
host=CONFIG.rabbitmq.host,
|
|
||||||
port=CONFIG.rabbitmq.port,
|
|
||||||
heartbeat=CONFIG.rabbitmq.heartbeat,
|
|
||||||
credentials=credentials,
|
|
||||||
)
|
|
||||||
return parameters
|
|
||||||
|
|
||||||
|
|
||||||
def make_connection() -> pika.BlockingConnection:
|
|
||||||
parameters = read_connection_params()
|
|
||||||
connection = pika.BlockingConnection(parameters)
|
|
||||||
return connection
|
|
||||||
0
pyinfra/test/queue/__init__.py
Normal file
0
pyinfra/test/queue/__init__.py
Normal file
43
pyinfra/test/queue/queue_manager_mock.py
Normal file
43
pyinfra/test/queue/queue_manager_mock.py
Normal file
@ -0,0 +1,43 @@
|
|||||||
|
from pyinfra.queue.queue_manager.queue_manager import QueueManager, QueueHandle
|
||||||
|
from pyinfra.test.queue.queue_mock import QueueMock
|
||||||
|
|
||||||
|
|
||||||
|
def monkey_patch_queue_handle(queue) -> QueueHandle:
|
||||||
|
queue_handle = QueueHandle()
|
||||||
|
queue_handle.empty = lambda: not queue
|
||||||
|
queue_handle.to_list = lambda: list(queue)
|
||||||
|
return queue_handle
|
||||||
|
|
||||||
|
|
||||||
|
class QueueManagerMock(QueueManager):
|
||||||
|
def __init__(self, input_queue, output_queue):
|
||||||
|
super().__init__(QueueMock(), QueueMock())
|
||||||
|
|
||||||
|
def publish_request(self, request):
|
||||||
|
self._input_queue.append(request)
|
||||||
|
|
||||||
|
def publish_response(self, message, callback):
|
||||||
|
self._output_queue.append(callback(message))
|
||||||
|
|
||||||
|
def pull_request(self):
|
||||||
|
return self._input_queue.popleft()
|
||||||
|
|
||||||
|
def consume(self):
|
||||||
|
while self._input_queue:
|
||||||
|
yield self.pull_request()
|
||||||
|
|
||||||
|
def consume_and_publish(self, callback):
|
||||||
|
for message in self.consume():
|
||||||
|
self.publish_response(message, callback)
|
||||||
|
|
||||||
|
def clear(self):
|
||||||
|
self._input_queue.clear()
|
||||||
|
self._output_queue.clear()
|
||||||
|
|
||||||
|
@property
|
||||||
|
def input_queue(self) -> QueueHandle:
|
||||||
|
return monkey_patch_queue_handle(self._input_queue)
|
||||||
|
|
||||||
|
@property
|
||||||
|
def output_queue(self) -> QueueHandle:
|
||||||
|
return monkey_patch_queue_handle(self._output_queue)
|
||||||
5
pyinfra/test/queue/queue_mock.py
Normal file
5
pyinfra/test/queue/queue_mock.py
Normal file
@ -0,0 +1,5 @@
|
|||||||
|
from collections import deque
|
||||||
|
|
||||||
|
|
||||||
|
class QueueMock(deque):
|
||||||
|
pass
|
||||||
File diff suppressed because it is too large
Load Diff
@ -6,20 +6,24 @@ import pytest
|
|||||||
|
|
||||||
from pyinfra.exceptions import UnknownClient
|
from pyinfra.exceptions import UnknownClient
|
||||||
from pyinfra.locations import TEST_DIR
|
from pyinfra.locations import TEST_DIR
|
||||||
|
from pyinfra.queue.queue_manager.pika_queue_manager import PikaQueueManager
|
||||||
|
from pyinfra.queue.queue_manager.queue_manager import QueueManager
|
||||||
from pyinfra.storage.adapters.azure import AzureStorageAdapter
|
from pyinfra.storage.adapters.azure import AzureStorageAdapter
|
||||||
from pyinfra.storage.adapters.s3 import S3StorageAdapter
|
from pyinfra.storage.adapters.s3 import S3StorageAdapter
|
||||||
from pyinfra.storage.clients.azure import get_azure_client
|
from pyinfra.storage.clients.azure import get_azure_client
|
||||||
from pyinfra.storage.clients.s3 import get_s3_client
|
from pyinfra.storage.clients.s3 import get_s3_client
|
||||||
from pyinfra.storage.storage import Storage
|
from pyinfra.storage.storage import Storage
|
||||||
from pyinfra.test.config import CONFIG
|
from pyinfra.test.config import CONFIG
|
||||||
|
from pyinfra.test.queue.queue_manager_mock import QueueManagerMock
|
||||||
from pyinfra.test.storage.adapter_mock import StorageAdapterMock
|
from pyinfra.test.storage.adapter_mock import StorageAdapterMock
|
||||||
from pyinfra.test.storage.client_mock import StorageClientMock
|
from pyinfra.test.storage.client_mock import StorageClientMock
|
||||||
|
from pyinfra.visitor import StorageStrategy, ForwardingStrategy, QueueVisitor
|
||||||
|
|
||||||
logger = logging.getLogger(__name__)
|
logger = logging.getLogger(__name__)
|
||||||
logger.setLevel(logging.DEBUG)
|
logger.setLevel(logging.DEBUG)
|
||||||
|
|
||||||
|
|
||||||
@pytest.fixture
|
@pytest.fixture(scope="session")
|
||||||
def bucket_name():
|
def bucket_name():
|
||||||
return "pyinfra-test-bucket"
|
return "pyinfra-test-bucket"
|
||||||
|
|
||||||
@ -51,7 +55,7 @@ def mock_make_load_data():
|
|||||||
return load_data
|
return load_data
|
||||||
|
|
||||||
|
|
||||||
@pytest.fixture(params=["minio", "aws"])
|
@pytest.fixture(params=["minio", "aws"], scope="session")
|
||||||
def storage(client_name, bucket_name, request):
|
def storage(client_name, bucket_name, request):
|
||||||
logger.debug("Setup for storage")
|
logger.debug("Setup for storage")
|
||||||
storage = Storage(get_adapter(client_name, request.param))
|
storage = Storage(get_adapter(client_name, request.param))
|
||||||
@ -71,3 +75,47 @@ def get_adapter(client_name, s3_backend):
|
|||||||
return S3StorageAdapter(get_s3_client(CONFIG.storage[s3_backend]))
|
return S3StorageAdapter(get_s3_client(CONFIG.storage[s3_backend]))
|
||||||
else:
|
else:
|
||||||
raise UnknownClient(client_name)
|
raise UnknownClient(client_name)
|
||||||
|
|
||||||
|
|
||||||
|
def get_queue_manager(queue_manager_name) -> QueueManager:
|
||||||
|
if queue_manager_name == "mock":
|
||||||
|
return QueueManagerMock("input", "output")
|
||||||
|
if queue_manager_name == "pika":
|
||||||
|
return PikaQueueManager("input", "output")
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.fixture(scope="session")
|
||||||
|
def queue_manager(queue_manager_name):
|
||||||
|
queue_manager = get_queue_manager(queue_manager_name)
|
||||||
|
yield queue_manager
|
||||||
|
if queue_manager_name == "pika":
|
||||||
|
queue_manager.connection.close()
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.fixture(scope="session")
|
||||||
|
def callback():
|
||||||
|
def inner(request):
|
||||||
|
return request["data"].decode() * 2
|
||||||
|
|
||||||
|
return inner
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.fixture
|
||||||
|
def analysis_callback(callback):
|
||||||
|
def inner(request):
|
||||||
|
return callback(request)
|
||||||
|
|
||||||
|
return inner
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.fixture
|
||||||
|
def response_strategy(response_strategy_name, storage):
|
||||||
|
if response_strategy_name == "storage":
|
||||||
|
return StorageStrategy(storage)
|
||||||
|
if response_strategy_name == "forwarding":
|
||||||
|
return ForwardingStrategy()
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.fixture()
|
||||||
|
def visitor(storage, analysis_callback, response_strategy):
|
||||||
|
return QueueVisitor(storage, analysis_callback, response_strategy)
|
||||||
|
|||||||
122
pyinfra/test/unit_tests/consumer_test.py
Normal file
122
pyinfra/test/unit_tests/consumer_test.py
Normal file
@ -0,0 +1,122 @@
|
|||||||
|
import gzip
|
||||||
|
import json
|
||||||
|
import logging
|
||||||
|
from operator import itemgetter
|
||||||
|
|
||||||
|
import pytest
|
||||||
|
|
||||||
|
from pyinfra.exceptions import ProcessingFailure
|
||||||
|
from pyinfra.queue.consumer import Consumer
|
||||||
|
from pyinfra.visitor import get_object_descriptor, ForwardingStrategy
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.fixture(scope="session")
|
||||||
|
def consumer(queue_manager, callback):
|
||||||
|
return Consumer(callback, queue_manager)
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.fixture(scope="session")
|
||||||
|
def access_callback():
|
||||||
|
return itemgetter("fileId")
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.fixture()
|
||||||
|
def items():
|
||||||
|
def inner():
|
||||||
|
for i in range(3):
|
||||||
|
body = {
|
||||||
|
"dossierId": "folder",
|
||||||
|
"fileId": f"file{i}",
|
||||||
|
"targetFileExtension": "in.gz",
|
||||||
|
"responseFileExtension": "out.gz",
|
||||||
|
}
|
||||||
|
yield f"{i}".encode(), body
|
||||||
|
|
||||||
|
return list(inner())
|
||||||
|
|
||||||
|
|
||||||
|
class TestConsumer:
|
||||||
|
@pytest.mark.parametrize("queue_manager_name", ["mock", "pika"], scope="session")
|
||||||
|
def test_consuming_empty_input_queue_does_not_put_anything_on_output_queue(self, consumer, queue_manager):
|
||||||
|
queue_manager.clear()
|
||||||
|
consumer.consume()
|
||||||
|
assert queue_manager.output_queue.empty()
|
||||||
|
|
||||||
|
@pytest.mark.parametrize("queue_manager_name", ["mock", "pika"], scope="session")
|
||||||
|
def test_consuming_nonempty_input_queue_puts_messages_on_output_queue_in_fifo_order(
|
||||||
|
self, consumer, queue_manager, callback
|
||||||
|
):
|
||||||
|
def produce_items():
|
||||||
|
return map(str, range(3))
|
||||||
|
|
||||||
|
def mock_visitor(callback):
|
||||||
|
def inner(data):
|
||||||
|
return callback({"data": data.encode()})
|
||||||
|
|
||||||
|
return inner
|
||||||
|
|
||||||
|
callback = mock_visitor(callback)
|
||||||
|
|
||||||
|
queue_manager.clear()
|
||||||
|
|
||||||
|
for item in produce_items():
|
||||||
|
queue_manager.publish_request(item)
|
||||||
|
|
||||||
|
requests = consumer.consume()
|
||||||
|
|
||||||
|
for _, r in zip(produce_items(), requests):
|
||||||
|
queue_manager.publish_response(r, callback)
|
||||||
|
|
||||||
|
assert queue_manager.output_queue.to_list() == ["00", "11", "22"]
|
||||||
|
|
||||||
|
@pytest.mark.parametrize("queue_manager_name", ["mock", "pika"], scope="session")
|
||||||
|
@pytest.mark.parametrize("client_name", ["mock", "s3", "azure"], scope="session")
|
||||||
|
@pytest.mark.parametrize("response_strategy_name", ["forwarding", "storage"], scope="session")
|
||||||
|
def test_consuming_nonempty_input_queue_with_visitor_puts_messages_on_output_queue_in_fifo_order(
|
||||||
|
self, consumer, queue_manager, visitor, bucket_name, storage, items
|
||||||
|
):
|
||||||
|
|
||||||
|
visitor.response_strategy = ForwardingStrategy()
|
||||||
|
|
||||||
|
queue_manager.clear()
|
||||||
|
storage.clear_bucket(bucket_name)
|
||||||
|
|
||||||
|
for data, message in items:
|
||||||
|
storage.put_object(**get_object_descriptor(message), data=gzip.compress(data))
|
||||||
|
queue_manager.publish_request(message)
|
||||||
|
|
||||||
|
requests = consumer.consume()
|
||||||
|
|
||||||
|
for _, r in zip(items, requests):
|
||||||
|
queue_manager.publish_response(r, visitor)
|
||||||
|
|
||||||
|
assert list(map(itemgetter("data"), queue_manager.output_queue.to_list())) == ["00", "11", "22"]
|
||||||
|
|
||||||
|
@pytest.mark.parametrize("queue_manager_name", ["pika"], scope="session")
|
||||||
|
def test_message_is_republished_when_callback_raises_processing_failure_exception(
|
||||||
|
self, consumer, queue_manager, bucket_name, items
|
||||||
|
):
|
||||||
|
class DebugError(Exception):
|
||||||
|
pass
|
||||||
|
|
||||||
|
def callback(_):
|
||||||
|
raise ProcessingFailure()
|
||||||
|
|
||||||
|
def reject_patch(*args, **kwargs):
|
||||||
|
raise DebugError()
|
||||||
|
|
||||||
|
queue_manager.reject = reject_patch
|
||||||
|
|
||||||
|
queue_manager.clear()
|
||||||
|
|
||||||
|
for data, message in items:
|
||||||
|
queue_manager.publish_request(message)
|
||||||
|
|
||||||
|
requests = consumer.consume()
|
||||||
|
|
||||||
|
logger = logging.getLogger("pyinfra.queue.queue_manager.pika_queue_manager")
|
||||||
|
logger.addFilter(lambda record: False)
|
||||||
|
|
||||||
|
with pytest.raises(DebugError):
|
||||||
|
while True:
|
||||||
|
queue_manager.publish_response(next(requests), callback)
|
||||||
@ -1,32 +0,0 @@
|
|||||||
import json
|
|
||||||
from unittest.mock import patch
|
|
||||||
|
|
||||||
from pyinfra.core import make_analyzer, make_payload_processor
|
|
||||||
from pyinfra.test.config import CONFIG
|
|
||||||
|
|
||||||
|
|
||||||
@patch("requests.post")
|
|
||||||
def test_analyse_returns_analysis(mock_post, storage_data, mock_response):
|
|
||||||
mock_post.return_value = mock_response
|
|
||||||
|
|
||||||
analyze = make_analyzer(CONFIG.mock_analysis_endpoint)
|
|
||||||
response = analyze(storage_data)
|
|
||||||
|
|
||||||
assert response == storage_data
|
|
||||||
|
|
||||||
|
|
||||||
@patch("requests.post")
|
|
||||||
def test_process_returns_dossier_id_file_id_predictions(
|
|
||||||
mock_post, mock_make_load_data, storage_data, mock_response, mock_payload
|
|
||||||
):
|
|
||||||
mock_post.return_value = mock_response
|
|
||||||
|
|
||||||
analyze = make_analyzer(CONFIG.mock_analysis_endpoint)
|
|
||||||
mock_load_data = mock_make_load_data
|
|
||||||
|
|
||||||
process = make_payload_processor(mock_load_data, analyze)
|
|
||||||
dossier_id, file_id, predictions = process(mock_payload)
|
|
||||||
|
|
||||||
assert dossier_id == json.loads(mock_payload)["dossierId"]
|
|
||||||
assert file_id == json.loads(mock_payload)["fileId"]
|
|
||||||
assert predictions == storage_data
|
|
||||||
38
pyinfra/test/unit_tests/queue_visitor_test.py
Normal file
38
pyinfra/test/unit_tests/queue_visitor_test.py
Normal file
@ -0,0 +1,38 @@
|
|||||||
|
import gzip
|
||||||
|
import json
|
||||||
|
|
||||||
|
import pytest
|
||||||
|
|
||||||
|
from pyinfra.visitor import get_object_descriptor, get_response_object_descriptor
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.fixture()
|
||||||
|
def body():
|
||||||
|
return {"dossierId": "folder", "fileId": "file", "targetFileExtension": "in.gz", "responseFileExtension": "out.gz"}
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.mark.parametrize("client_name", ["mock", "azure", "s3"], scope="session")
|
||||||
|
class TestVisitor:
|
||||||
|
@pytest.mark.parametrize("response_strategy_name", ["forwarding", "storage"], scope="session")
|
||||||
|
def test_given_a_input_queue_message_callback_pulls_the_data_from_storage(
|
||||||
|
self, visitor, body, storage, bucket_name
|
||||||
|
):
|
||||||
|
storage.clear_bucket(bucket_name)
|
||||||
|
storage.put_object(**get_object_descriptor(body), data=gzip.compress(b"content"))
|
||||||
|
data_received = visitor.load_data(body)
|
||||||
|
assert b"content" == data_received
|
||||||
|
|
||||||
|
@pytest.mark.parametrize("response_strategy_name", ["forwarding", "storage"], scope="session")
|
||||||
|
def test_visitor_pulls_and_processes_data(self, visitor, body, storage, bucket_name):
|
||||||
|
storage.clear_bucket(bucket_name)
|
||||||
|
storage.put_object(**get_object_descriptor(body), data=gzip.compress("2".encode()))
|
||||||
|
response_body = visitor.load_and_process(body)
|
||||||
|
assert response_body["data"] == "22"
|
||||||
|
|
||||||
|
@pytest.mark.parametrize("response_strategy_name", ["storage"], scope="session")
|
||||||
|
def test_visitor_puts_response_on_storage(self, visitor, body, storage, bucket_name):
|
||||||
|
storage.clear_bucket(bucket_name)
|
||||||
|
storage.put_object(**get_object_descriptor(body), data=gzip.compress("2".encode()))
|
||||||
|
response_body = visitor(body)
|
||||||
|
assert "data" not in response_body
|
||||||
|
assert json.loads(gzip.decompress(storage.get_object(**get_response_object_descriptor(body))))["data"] == "22"
|
||||||
@ -9,7 +9,7 @@ logger = logging.getLogger(__name__)
|
|||||||
logger.setLevel(logging.DEBUG)
|
logger.setLevel(logging.DEBUG)
|
||||||
|
|
||||||
|
|
||||||
@pytest.mark.parametrize("client_name", ["mock", "azure", "s3"])
|
@pytest.mark.parametrize("client_name", ["mock", "azure", "s3"], scope="session")
|
||||||
class TestStorage:
|
class TestStorage:
|
||||||
def test_clearing_bucket_yields_empty_bucket(self, storage, bucket_name):
|
def test_clearing_bucket_yields_empty_bucket(self, storage, bucket_name):
|
||||||
storage.clear_bucket(bucket_name)
|
storage.clear_bucket(bucket_name)
|
||||||
@ -17,26 +17,31 @@ class TestStorage:
|
|||||||
assert not {*data_received}
|
assert not {*data_received}
|
||||||
|
|
||||||
def test_getting_object_put_in_bucket_is_object(self, storage, bucket_name):
|
def test_getting_object_put_in_bucket_is_object(self, storage, bucket_name):
|
||||||
|
storage.clear_bucket(bucket_name)
|
||||||
storage.put_object(bucket_name, "file", b"content")
|
storage.put_object(bucket_name, "file", b"content")
|
||||||
data_received = storage.get_object(bucket_name, "file")
|
data_received = storage.get_object(bucket_name, "file")
|
||||||
assert b"content" == data_received
|
assert b"content" == data_received
|
||||||
|
|
||||||
def test_getting_nested_object_put_in_bucket_is_nested_object(self, storage, bucket_name):
|
def test_getting_nested_object_put_in_bucket_is_nested_object(self, storage, bucket_name):
|
||||||
|
storage.clear_bucket(bucket_name)
|
||||||
storage.put_object(bucket_name, "folder/file", b"content")
|
storage.put_object(bucket_name, "folder/file", b"content")
|
||||||
data_received = storage.get_object(bucket_name, "folder/file")
|
data_received = storage.get_object(bucket_name, "folder/file")
|
||||||
assert b"content" == data_received
|
assert b"content" == data_received
|
||||||
|
|
||||||
def test_getting_objects_put_in_bucket_are_objects(self, storage, bucket_name):
|
def test_getting_objects_put_in_bucket_are_objects(self, storage, bucket_name):
|
||||||
|
storage.clear_bucket(bucket_name)
|
||||||
storage.put_object(bucket_name, "file1", b"content 1")
|
storage.put_object(bucket_name, "file1", b"content 1")
|
||||||
storage.put_object(bucket_name, "folder/file2", b"content 2")
|
storage.put_object(bucket_name, "folder/file2", b"content 2")
|
||||||
data_received = storage.get_all_objects(bucket_name)
|
data_received = storage.get_all_objects(bucket_name)
|
||||||
assert {b"content 1", b"content 2"} == {*data_received}
|
assert {b"content 1", b"content 2"} == {*data_received}
|
||||||
|
|
||||||
def test_make_bucket_produces_bucket(self, storage, bucket_name):
|
def test_make_bucket_produces_bucket(self, storage, bucket_name):
|
||||||
|
storage.clear_bucket(bucket_name)
|
||||||
storage.make_bucket(bucket_name)
|
storage.make_bucket(bucket_name)
|
||||||
assert storage.has_bucket(bucket_name)
|
assert storage.has_bucket(bucket_name)
|
||||||
|
|
||||||
def test_listing_bucket_files_yields_all_files_in_bucket(self, storage, bucket_name):
|
def test_listing_bucket_files_yields_all_files_in_bucket(self, storage, bucket_name):
|
||||||
|
storage.clear_bucket(bucket_name)
|
||||||
storage.put_object(bucket_name, "file1", b"content 1")
|
storage.put_object(bucket_name, "file1", b"content 1")
|
||||||
storage.put_object(bucket_name, "file2", b"content 2")
|
storage.put_object(bucket_name, "file2", b"content 2")
|
||||||
full_names_received = storage.get_all_object_names(bucket_name)
|
full_names_received = storage.get_all_object_names(bucket_name)
|
||||||
|
|||||||
@ -1,16 +0,0 @@
|
|||||||
"""Defines utilities for different operations on files."""
|
|
||||||
|
|
||||||
import gzip
|
|
||||||
import os
|
|
||||||
|
|
||||||
from pyinfra.config import CONFIG
|
|
||||||
|
|
||||||
|
|
||||||
def combine_dossier_id_and_file_id_and_extension(dossier_id, file_id, extension):
|
|
||||||
return f"{dossier_id}/{file_id}{extension}"
|
|
||||||
|
|
||||||
|
|
||||||
def upload_compressed_response(storage, bucket_name, dossier_id, file_id, result) -> None:
|
|
||||||
data = gzip.compress(result.encode())
|
|
||||||
path_gz = combine_dossier_id_and_file_id_and_extension(dossier_id, file_id, CONFIG.service.response.extension)
|
|
||||||
storage.put_object(bucket_name, path_gz, data)
|
|
||||||
@ -1,96 +0,0 @@
|
|||||||
import logging
|
|
||||||
import time
|
|
||||||
from functools import partial, wraps
|
|
||||||
from math import exp
|
|
||||||
from typing import Tuple, Type, Callable
|
|
||||||
|
|
||||||
|
|
||||||
class NoAttemptsLeft(Exception):
|
|
||||||
pass
|
|
||||||
|
|
||||||
|
|
||||||
class MaxTimeoutReached(Exception):
|
|
||||||
pass
|
|
||||||
|
|
||||||
|
|
||||||
class _MethodDecoratorAdaptor(object):
|
|
||||||
def __init__(self, decorator, func):
|
|
||||||
self.decorator = decorator
|
|
||||||
self.func = func
|
|
||||||
|
|
||||||
def __call__(self, *args, **kwargs):
|
|
||||||
return self.decorator(self.func)(*args, **kwargs)
|
|
||||||
|
|
||||||
def __get__(self, obj, objtype):
|
|
||||||
return partial(self.__call__, obj)
|
|
||||||
|
|
||||||
|
|
||||||
def auto_adapt_to_methods(decorator):
|
|
||||||
"""Allows you to use the same decorator on methods and functions,
|
|
||||||
hiding the self argument from the decorator."""
|
|
||||||
|
|
||||||
def adapt(func):
|
|
||||||
return _MethodDecoratorAdaptor(decorator, func)
|
|
||||||
|
|
||||||
return adapt
|
|
||||||
|
|
||||||
|
|
||||||
def max_attempts(
|
|
||||||
n_attempts: int = 5, exceptions: Tuple[Type[Exception]] = None, timeout: float = 0.1, max_timeout: float = 10
|
|
||||||
) -> Callable:
|
|
||||||
"""Function decorator that attempts to run the wrapped function a certain number of times. Timeouts increase
|
|
||||||
exponentially according to `Tₖ ≔ t eᵏ`, where `t` is the timeout factor `timeout` and `k` is the attempt number.
|
|
||||||
If `∑ᵢ Tᵢ > mₜ` at the `i-th` attempt, where `mₜ` is the maximum timeout, then the function raises
|
|
||||||
MaxTimeoutReached. If `k > mₐ`, where `mₐ` is the maximum number of attempts allowed, then the function
|
|
||||||
raises NoAttemptsLeft.
|
|
||||||
|
|
||||||
Args:
|
|
||||||
n_attempts: Number of times to attempt running the wrapped function.
|
|
||||||
exceptions: Exceptions to catch for a re-attempt.
|
|
||||||
timeout: Timeout factor in seconds.
|
|
||||||
max_timeout: Maximum allowed timeout.
|
|
||||||
|
|
||||||
Raises:
|
|
||||||
MaxTimeoutReached
|
|
||||||
NoAttemptsLeft
|
|
||||||
|
|
||||||
Returns:
|
|
||||||
Wrapped function.
|
|
||||||
"""
|
|
||||||
if not exceptions:
|
|
||||||
exceptions = (Exception,)
|
|
||||||
assert isinstance(exceptions, tuple)
|
|
||||||
|
|
||||||
@auto_adapt_to_methods
|
|
||||||
def inner(func):
|
|
||||||
@wraps(func)
|
|
||||||
def inner(*args, **kwargs):
|
|
||||||
def run_attempt(attempt, timeout_aggr=0):
|
|
||||||
if attempt:
|
|
||||||
try:
|
|
||||||
return func(*args, **kwargs)
|
|
||||||
except exceptions as err:
|
|
||||||
|
|
||||||
attempt_num = n_attempts - attempt + 1
|
|
||||||
next_timeout = timeout * exp(attempt_num - 1) # start with timeout * e^0 = timeout
|
|
||||||
|
|
||||||
logging.warn(f"{func.__name__} failed; attempt {attempt_num} of {n_attempts}")
|
|
||||||
|
|
||||||
time_left = max(0, max_timeout - timeout_aggr)
|
|
||||||
if time_left:
|
|
||||||
sleep_for = min(next_timeout, time_left)
|
|
||||||
time.sleep(sleep_for)
|
|
||||||
return run_attempt(attempt - 1, timeout_aggr + sleep_for)
|
|
||||||
else:
|
|
||||||
logging.exception(err)
|
|
||||||
raise MaxTimeoutReached(
|
|
||||||
f"{func.__name__} reached maximum timeout ({max_timeout}) after {attempt_num} attempts."
|
|
||||||
)
|
|
||||||
else:
|
|
||||||
raise NoAttemptsLeft(f"{func.__name__} failed {n_attempts} times; all attempts expended.")
|
|
||||||
|
|
||||||
return run_attempt(n_attempts)
|
|
||||||
|
|
||||||
return inner
|
|
||||||
|
|
||||||
return inner
|
|
||||||
88
pyinfra/visitor.py
Normal file
88
pyinfra/visitor.py
Normal file
@ -0,0 +1,88 @@
|
|||||||
|
import abc
|
||||||
|
import gzip
|
||||||
|
import json
|
||||||
|
import logging
|
||||||
|
from operator import itemgetter
|
||||||
|
from typing import Callable
|
||||||
|
|
||||||
|
from pyinfra.config import CONFIG
|
||||||
|
from pyinfra.exceptions import DataLoadingFailure
|
||||||
|
from pyinfra.storage.storage import Storage
|
||||||
|
|
||||||
|
|
||||||
|
def get_object_name(body):
|
||||||
|
dossier_id, file_id, target_file_extension = itemgetter("dossierId", "fileId", "targetFileExtension")(body)
|
||||||
|
object_name = f"{dossier_id}/{file_id}.{target_file_extension}"
|
||||||
|
return object_name
|
||||||
|
|
||||||
|
|
||||||
|
def get_response_object_name(body):
|
||||||
|
dossier_id, file_id, response_file_extension = itemgetter("dossierId", "fileId", "responseFileExtension")(body)
|
||||||
|
object_name = f"{dossier_id}/{file_id}.{response_file_extension}"
|
||||||
|
return object_name
|
||||||
|
|
||||||
|
|
||||||
|
def get_object_descriptor(body):
|
||||||
|
return {"bucket_name": CONFIG.storage.bucket, "object_name": get_object_name(body)}
|
||||||
|
|
||||||
|
|
||||||
|
def get_response_object_descriptor(body):
|
||||||
|
return {"bucket_name": CONFIG.storage.bucket, "object_name": get_response_object_name(body)}
|
||||||
|
|
||||||
|
|
||||||
|
class ResponseStrategy(abc.ABC):
|
||||||
|
@abc.abstractmethod
|
||||||
|
def handle_response(self, body):
|
||||||
|
pass
|
||||||
|
|
||||||
|
def __call__(self, body):
|
||||||
|
return self.handle_response(body)
|
||||||
|
|
||||||
|
|
||||||
|
class StorageStrategy(ResponseStrategy):
|
||||||
|
def __init__(self, storage):
|
||||||
|
self.storage = storage
|
||||||
|
|
||||||
|
def handle_response(self, body):
|
||||||
|
self.storage.put_object(**get_response_object_descriptor(body), data=gzip.compress(json.dumps(body).encode()))
|
||||||
|
body.pop("data")
|
||||||
|
return body
|
||||||
|
|
||||||
|
|
||||||
|
class ForwardingStrategy(ResponseStrategy):
|
||||||
|
def handle_response(self, body):
|
||||||
|
return body
|
||||||
|
|
||||||
|
|
||||||
|
class QueueVisitor:
|
||||||
|
def __init__(self, storage: Storage, callback: Callable, response_strategy):
|
||||||
|
self.storage = storage
|
||||||
|
self.callback = callback
|
||||||
|
self.response_strategy = response_strategy
|
||||||
|
|
||||||
|
def load_data(self, body):
|
||||||
|
def download():
|
||||||
|
logging.debug(f"Downloading {object_descriptor}...")
|
||||||
|
data = self.storage.get_object(**object_descriptor)
|
||||||
|
logging.debug(f"Downloaded {object_descriptor}.")
|
||||||
|
return data
|
||||||
|
|
||||||
|
object_descriptor = get_object_descriptor(body)
|
||||||
|
|
||||||
|
try:
|
||||||
|
return gzip.decompress(download())
|
||||||
|
except Exception as err:
|
||||||
|
logging.warning(f"Loading data from storage failed for {object_descriptor}.")
|
||||||
|
raise DataLoadingFailure() from err
|
||||||
|
|
||||||
|
def process_data(self, data, body):
|
||||||
|
return self.callback({**body, "data": data})
|
||||||
|
|
||||||
|
def load_and_process(self, body):
|
||||||
|
data = self.process_data(self.load_data(body), body)
|
||||||
|
result_body = {**body, "data": data}
|
||||||
|
return result_body
|
||||||
|
|
||||||
|
def __call__(self, body):
|
||||||
|
result_body = self.load_and_process(body)
|
||||||
|
return self.response_strategy(result_body)
|
||||||
@ -7,7 +7,6 @@ from tqdm import tqdm
|
|||||||
|
|
||||||
from pyinfra.config import CONFIG
|
from pyinfra.config import CONFIG
|
||||||
from pyinfra.storage.storages import get_s3_storage
|
from pyinfra.storage.storages import get_s3_storage
|
||||||
from pyinfra.utils.file import combine_dossier_id_and_file_id_and_extension
|
|
||||||
|
|
||||||
|
|
||||||
def parse_args():
|
def parse_args():
|
||||||
@ -29,9 +28,7 @@ def parse_args():
|
|||||||
|
|
||||||
def add_file_compressed(storage, bucket_name, dossier_id, path) -> None:
|
def add_file_compressed(storage, bucket_name, dossier_id, path) -> None:
|
||||||
|
|
||||||
path_gz = combine_dossier_id_and_file_id_and_extension(
|
path_gz = combine_dossier_id_and_file_id_and_extension(dossier_id, Path(path).stem, ".ORIGIN.pdf.gz")
|
||||||
dossier_id, Path(path).stem, CONFIG.storage.target_file_extension
|
|
||||||
)
|
|
||||||
|
|
||||||
with open(path, "rb") as f:
|
with open(path, "rb") as f:
|
||||||
data = gzip.compress(f.read())
|
data = gzip.compress(f.read())
|
||||||
@ -60,3 +57,13 @@ if __name__ == "__main__":
|
|||||||
|
|
||||||
elif args.command == "purge":
|
elif args.command == "purge":
|
||||||
storage.clear_bucket(bucket_name)
|
storage.clear_bucket(bucket_name)
|
||||||
|
|
||||||
|
|
||||||
|
def combine_dossier_id_and_file_id_and_extension(dossier_id, file_id, extension):
|
||||||
|
return f"{dossier_id}/{file_id}{extension}"
|
||||||
|
|
||||||
|
|
||||||
|
def upload_compressed_response(storage, bucket_name, dossier_id, file_id, result) -> None:
|
||||||
|
data = gzip.compress(result.encode())
|
||||||
|
path_gz = combine_dossier_id_and_file_id_and_extension(dossier_id, file_id, CONFIG.service.response.extension)
|
||||||
|
storage.put_object(bucket_name, path_gz, data)
|
||||||
|
|||||||
@ -1,7 +1,8 @@
|
|||||||
import json
|
import json
|
||||||
|
|
||||||
|
import pika
|
||||||
|
|
||||||
from pyinfra.config import CONFIG
|
from pyinfra.config import CONFIG
|
||||||
from pyinfra.rabbitmq import make_channel, declare_queue, make_connection
|
|
||||||
from pyinfra.storage.storages import get_s3_storage
|
from pyinfra.storage.storages import get_s3_storage
|
||||||
|
|
||||||
|
|
||||||
@ -10,7 +11,31 @@ def build_message_bodies():
|
|||||||
for bucket_name, pdf_name in storage.get_all_object_names(CONFIG.storage.bucket):
|
for bucket_name, pdf_name in storage.get_all_object_names(CONFIG.storage.bucket):
|
||||||
file_id = pdf_name.split(".")[0]
|
file_id = pdf_name.split(".")[0]
|
||||||
dossier_id, file_id = file_id.split("/")
|
dossier_id, file_id = file_id.split("/")
|
||||||
yield json.dumps({"dossierId": dossier_id, "fileId": file_id}).encode()
|
yield json.dumps(
|
||||||
|
{
|
||||||
|
"dossierId": dossier_id,
|
||||||
|
"fileId": file_id,
|
||||||
|
"targetFileExtension": "ORIGIN.pdf.gz",
|
||||||
|
"responseFileExtension": "detr.json.gz",
|
||||||
|
}
|
||||||
|
).encode()
|
||||||
|
|
||||||
|
|
||||||
|
def make_channel(connection) -> pika.adapters.blocking_connection.BlockingChannel:
|
||||||
|
channel = connection.channel()
|
||||||
|
channel.basic_qos(prefetch_count=CONFIG.rabbitmq.prefetch_count)
|
||||||
|
return channel
|
||||||
|
|
||||||
|
|
||||||
|
def declare_queue(channel, queue: str):
|
||||||
|
args = {"x-dead-letter-exchange": "", "x-dead-letter-routing-key": CONFIG.rabbitmq.queues.dead_letter}
|
||||||
|
return channel.queue_declare(queue=queue, auto_delete=False, arguments=args)
|
||||||
|
|
||||||
|
|
||||||
|
def make_connection() -> pika.BlockingConnection:
|
||||||
|
parameters = read_connection_params()
|
||||||
|
connection = pika.BlockingConnection(parameters)
|
||||||
|
return connection
|
||||||
|
|
||||||
|
|
||||||
if __name__ == "__main__":
|
if __name__ == "__main__":
|
||||||
@ -27,3 +52,14 @@ if __name__ == "__main__":
|
|||||||
for method_frame, _, body in channel.consume(queue=CONFIG.rabbitmq.queues.output):
|
for method_frame, _, body in channel.consume(queue=CONFIG.rabbitmq.queues.output):
|
||||||
print(f"Received {json.loads(body)}")
|
print(f"Received {json.loads(body)}")
|
||||||
channel.basic_ack(method_frame.delivery_tag)
|
channel.basic_ack(method_frame.delivery_tag)
|
||||||
|
|
||||||
|
|
||||||
|
def read_connection_params():
|
||||||
|
credentials = pika.PlainCredentials(CONFIG.rabbitmq.user, CONFIG.rabbitmq.password)
|
||||||
|
parameters = pika.ConnectionParameters(
|
||||||
|
host=CONFIG.rabbitmq.host,
|
||||||
|
port=CONFIG.rabbitmq.port,
|
||||||
|
heartbeat=CONFIG.rabbitmq.heartbeat,
|
||||||
|
credentials=credentials,
|
||||||
|
)
|
||||||
|
return parameters
|
||||||
|
|||||||
79
src/serve.py
79
src/serve.py
@ -1,60 +1,65 @@
|
|||||||
import logging
|
import logging
|
||||||
from multiprocessing import Process
|
from multiprocessing import Process
|
||||||
|
|
||||||
import pika
|
import requests
|
||||||
|
from retry import retry
|
||||||
|
|
||||||
from pyinfra.callback import (
|
from pyinfra.config import CONFIG, make_art
|
||||||
make_retry_callback_for_output_queue,
|
from pyinfra.exceptions import AnalysisFailure, ConsumerError
|
||||||
make_retry_callback,
|
|
||||||
make_callback_for_output_queue,
|
|
||||||
)
|
|
||||||
from pyinfra.config import CONFIG
|
|
||||||
from pyinfra.consume import consume, ConsumerError
|
|
||||||
from pyinfra.core import make_payload_processor, make_storage_data_loader, make_analyzer
|
|
||||||
from pyinfra.flask import run_probing_webserver, set_up_probing_webserver
|
from pyinfra.flask import run_probing_webserver, set_up_probing_webserver
|
||||||
|
from pyinfra.queue.consumer import Consumer
|
||||||
|
from pyinfra.queue.queue_manager.pika_queue_manager import PikaQueueManager
|
||||||
from pyinfra.storage.storages import get_storage
|
from pyinfra.storage.storages import get_storage
|
||||||
|
from pyinfra.visitor import QueueVisitor, StorageStrategy
|
||||||
|
|
||||||
|
|
||||||
def republish(channel, body, n_current_attempts):
|
def make_callback(analysis_endpoint):
|
||||||
channel.basic_publish(
|
def callback(message):
|
||||||
exchange="",
|
def perform_operation(operation):
|
||||||
routing_key=CONFIG.rabbitmq.queues.input,
|
endpoint = f"{analysis_endpoint}/{operation}"
|
||||||
body=body,
|
try:
|
||||||
properties=pika.BasicProperties(headers={"x-retry-count": n_current_attempts}),
|
logging.debug(f"Requesting analysis from {endpoint}...")
|
||||||
)
|
analysis_response = requests.post(endpoint, data=message["data"])
|
||||||
|
analysis_response.raise_for_status()
|
||||||
|
analysis_response = analysis_response.json()
|
||||||
|
logging.debug(f"Received response.")
|
||||||
|
return analysis_response
|
||||||
|
except Exception as err:
|
||||||
|
logging.warning(f"Exception caught when calling analysis endpoint {endpoint}.")
|
||||||
|
raise AnalysisFailure() from err
|
||||||
|
|
||||||
|
operations = message.get("operations", ["/"])
|
||||||
def make_callback():
|
results = map(perform_operation, operations)
|
||||||
|
result = dict(zip(operations, results))
|
||||||
load_data = make_storage_data_loader(get_storage(CONFIG.storage.backend), CONFIG.storage.bucket)
|
return result
|
||||||
analyze_file = make_analyzer(CONFIG.rabbitmq.callback.analysis_endpoint)
|
|
||||||
|
|
||||||
json_wrapped_body_processor = make_payload_processor(load_data, analyze_file)
|
|
||||||
|
|
||||||
if CONFIG.rabbitmq.callback.retry.enabled:
|
|
||||||
retry_callback = make_retry_callback(republish, max_attempts=CONFIG.rabbitmq.callback.retry.max_attempts)
|
|
||||||
|
|
||||||
callback = make_retry_callback_for_output_queue(
|
|
||||||
json_wrapped_body_processor=json_wrapped_body_processor,
|
|
||||||
output_queue_name=CONFIG.rabbitmq.queues.output,
|
|
||||||
retry_callback=retry_callback,
|
|
||||||
)
|
|
||||||
else:
|
|
||||||
callback = make_callback_for_output_queue(
|
|
||||||
json_wrapped_body_processor=json_wrapped_body_processor, output_queue_name=CONFIG.rabbitmq.queues.output
|
|
||||||
)
|
|
||||||
|
|
||||||
return callback
|
return callback
|
||||||
|
|
||||||
|
|
||||||
def main():
|
def main():
|
||||||
# TODO: implement meaningful checks
|
# TODO: implement meaningful checks
|
||||||
|
logging.info(make_art())
|
||||||
webserver = Process(target=run_probing_webserver, args=(set_up_probing_webserver(),))
|
webserver = Process(target=run_probing_webserver, args=(set_up_probing_webserver(),))
|
||||||
logging.info("Starting webserver...")
|
logging.info("Starting webserver...")
|
||||||
webserver.start()
|
webserver.start()
|
||||||
|
|
||||||
|
callback = make_callback(CONFIG.rabbitmq.callback.analysis_endpoint)
|
||||||
|
storage = get_storage(CONFIG.storage.backend)
|
||||||
|
response_strategy = StorageStrategy(storage)
|
||||||
|
visitor = QueueVisitor(storage, callback, response_strategy)
|
||||||
|
|
||||||
|
queue_manager = PikaQueueManager(CONFIG.rabbitmq.queues.input, CONFIG.rabbitmq.queues.output)
|
||||||
|
|
||||||
|
@retry(ConsumerError, tries=3, delay=5, jitter=(1, 3))
|
||||||
|
def consume():
|
||||||
|
try:
|
||||||
|
consumer = Consumer(visitor, queue_manager)
|
||||||
|
consumer.consume_and_publish()
|
||||||
|
except Exception as err:
|
||||||
|
raise ConsumerError from err
|
||||||
|
|
||||||
try:
|
try:
|
||||||
consume(CONFIG.rabbitmq.queues.input, make_callback())
|
consume()
|
||||||
except KeyboardInterrupt:
|
except KeyboardInterrupt:
|
||||||
pass
|
pass
|
||||||
except ConsumerError:
|
except ConsumerError:
|
||||||
|
|||||||
Loading…
x
Reference in New Issue
Block a user