Merge pull request 'master' (#24) from master into prod

Reviewed-on: #24

.deploy/deploy-dev.yaml

@@ -6,9 +6,11 @@ services:
     image: mathwave/sprint-repo:queues
     networks:
       - queues-development
+      - configurator
     environment:
       MONGO_HOST: "mongo.develop.sprinthub.ru"
       MONGO_PASSWORD: $MONGO_PASSWORD_DEV
+      STAGE: "development"
     deploy:
       mode: replicated
       restart_policy:
@@ -20,3 +22,5 @@ services:
 networks:
   queues-development:
     external: true
+  configurator:
+    external: true

.deploy/deploy-prod.yaml

@@ -6,9 +6,11 @@ services:
     image: mathwave/sprint-repo:queues
     networks:
       - queues
+      - configurator
     environment:
       MONGO_HOST: "mongo.sprinthub.ru"
       MONGO_PASSWORD: $MONGO_PASSWORD_PROD
+      STAGE: "production"
     deploy:
       mode: replicated
       restart_policy:
@@ -20,3 +22,5 @@ services:
 networks:
   queues:
     external: true
+  configurator:
+    external: true

app/routers/take.py

@@ -1,10 +1,15 @@
+import asyncio
+import collections
 import fastapi
 import pydantic
 import typing
 
 from app.storage.mongo import tasks
+# from app.utils.configurator import configurator
 
 
+DEFAULT_RETRY_AFTER = 0.2
+locks = collections.defaultdict(asyncio.Lock)
 router = fastapi.APIRouter()
 
 
@@ -18,9 +23,12 @@ class Response(pydantic.BaseModel):
     task: Task|None
 
 
-@router.get('/api/v1/take', responses={404: {'description': 'Not found'}})
+@router.get('/api/v1/take')
 async def execute(queue: typing.Annotated[str, fastapi.Header()]) -> Response:
-    task = await tasks.take_task(queue)
+    async with locks[queue]:
+        task = await tasks.take_task(queue)
     if not task:
         return Response(task=None)
+    # retry_after_config = configurator.get_config('retry_after')
+    # retry_after = retry_after_config.get(queue) or retry_after_config.get('default') or DEFAULT_RETRY_AFTER
     return Response(task=Task(id=str(task._id), attempt=task.attempts, payload=task.payload))

app/utils/configurator.py

@@ -0,0 +1,88 @@
+import json
+import urllib.parse
+import os
+from threading import Thread
+from time import sleep
+
+from requests import get
+
+
+class ConfiguratorClient:
+    def __init__(self, app_name: str, stage: str, need_poll: bool = True):
+        self.app_name = app_name
+        self.stage = stage
+        self.endpoint = 'http://configurator/'
+        self.fetch_url = urllib.parse.urljoin(self.endpoint, '/api/v1/fetch')
+        self.config_storage = {}
+        self.experiment_storage = {}
+        self.staff_storage = {}
+        self.poll_data()
+        if need_poll:
+            self.poll_data_in_thread()
+
+    def poll_data_in_thread(self):
+        def inner():
+            while True:
+                sleep(30)
+                self.fetch()
+
+        Thread(target=inner, daemon=True).start()
+
+    def poll_data(self):
+        self.fetch(with_exception=True)
+
+    def request_with_retries(self, url, params, with_exception=False, retries_count=3):
+        exception_to_throw = None
+        for _ in range(retries_count):
+            try:
+                response = get(
+                    url,
+                    params=params
+                )
+                if response.status_code == 200:
+                    return response.json()
+                print(f'Failed to request {url}, status_code={response.status_code}')
+                exception_to_throw = Exception('Not 200 status')
+            except Exception as exc:
+                print(exc)
+                exception_to_throw = exc
+            sleep(1)
+        print(f'Failed fetching with retries: {url}, {params}')
+        if with_exception:
+            raise exception_to_throw
+
+    def fetch(self, with_exception=False):
+        if self.stage == 'local':
+            local_platform = json.loads(open('local_platform.json', 'r').read())
+            self.config_storage = local_platform['configs']
+            self.experiment_storage = local_platform['experiments']
+            self.staff_storage = {
+                key: set(value)
+                for key, value in local_platform['platform_staff'].items()
+            }
+            return
+        response_data = self.request_with_retries(self.fetch_url, {
+            'project': self.app_name,
+            'stage': self.stage,
+        }, with_exception)
+        self.config_storage = response_data['configs']
+        self.experiment_storage = response_data['experiments']
+        self.staff_storage = {
+            key: set(value)
+            for key, value in response_data['platform_staff'].items()
+        }
+
+    def is_staff(self, **kwargs):
+        for key, value in kwargs.items():
+            if value in self.staff_storage[key]:
+                return True
+        return False
+
+    def get_config(self, name):
+        return self.config_storage[name]
+
+    def get_experiment(self, name):
+        return self.experiment_storage[name]
+    
+
+# configurator = ConfiguratorClient('queues', os.getenv('STAGE', 'local'))

requirements.txt

@@ -1,6 +1,8 @@
 annotated-types==0.7.0
 anyio==4.6.2.post1
 APScheduler==3.10.4
+certifi==2024.12.14
+charset-normalizer==3.4.1
 click==8.1.7
 dnspython==2.7.0
 fastapi==0.115.4
@@ -12,9 +14,11 @@ pydantic_core==2.23.4
 pymongo==4.9.2
 pytz==2024.2
 redis==5.2.0
+requests==2.32.3
 six==1.16.0
 sniffio==1.3.1
 starlette==0.41.2
 typing_extensions==4.12.2
 tzlocal==5.2
+urllib3==2.3.0
 uvicorn==0.32.0