feat: support cohort targeting for local evaluation

src/amplitude_experiment/cohort/cohort.py

@@ -0,0 +1,13 @@
+from dataclasses import dataclass, field
+from typing import ClassVar, Set
+
+USER_GROUP_TYPE: ClassVar[str] = "User"
+
+
+@dataclass
+class Cohort:
+    id: str
+    last_modified: int
+    size: int
+    member_ids: Set[str]
+    group_type: str = field(default=USER_GROUP_TYPE)

src/amplitude_experiment/cohort/cohort_download_api.py

@@ -0,0 +1,93 @@
+import time
+import logging
+import base64
+import json
+from http.client import HTTPResponse
+from typing import Optional
+
+from .cohort import Cohort
+from ..connection_pool import HTTPConnectionPool
+from ..exception import HTTPErrorResponseException, CohortTooLargeException, CohortNotModifiedException
+
+CDN_COHORT_SYNC_URL = 'https://cohort-v2.lab.amplitude.com'
+
+
+class CohortDownloadApi:
+    def __init__(self):
+        self.cdn_server_url = CDN_COHORT_SYNC_URL
+
+    def get_cohort(self, cohort_id: str, cohort: Cohort) -> Optional[Cohort]:
+        raise NotImplementedError
+
+
+class DirectCohortDownloadApi(CohortDownloadApi):
+    def __init__(self, api_key: str, secret_key: str, max_cohort_size: int, cohort_request_delay_millis: int,
+                 debug: bool):
+        super().__init__()
+        self.api_key = api_key
+        self.secret_key = secret_key
+        self.max_cohort_size = max_cohort_size
+        self.__setup_connection_pool()
+        self.cohort_request_delay_millis = cohort_request_delay_millis
+        self.logger = logging.getLogger("Amplitude")
+        self.logger.addHandler(logging.StreamHandler())
+        if debug:
+            self.logger.setLevel(logging.DEBUG)
+
+    def get_cohort(self, cohort_id: str, cohort: Cohort) -> Optional[Cohort]:
+        self.logger.debug(f"getCohortMembers({cohort_id}): start")
+        errors = 0
+        while True:
+            response = None
+            try:
+                last_modified = None if cohort is None else cohort.last_modified
+                response = self._get_cohort_members_request(cohort_id, last_modified)
+                self.logger.debug(f"getCohortMembers({cohort_id}): status={response.status}")
+                if response.status == 200:
+                    cohort_info = json.loads(response.read().decode("utf8"))
+                    self.logger.debug(f"getCohortMembers({cohort_id}): end - resultSize={cohort_info['size']}")
+                    return Cohort(
+                        id=cohort_info['cohortId'],
+                        last_modified=cohort_info['lastModified'],
+                        size=cohort_info['size'],
+                        member_ids=set(cohort_info['memberIds']),
+                        group_type=cohort_info['groupType'],
+                    )
+                elif response.status == 204:
+                    raise CohortNotModifiedException(f"Cohort not modified: {response.status}")
+                elif response.status == 413:
+                    raise CohortTooLargeException(f"Cohort exceeds max cohort size: {response.status}")
+                elif response.status != 202:
+                    raise HTTPErrorResponseException(response.status,
+                                                     f"Unexpected response code: {response.status}")
+            except Exception as e:
+                if not isinstance(e, HTTPErrorResponseException) and response.status != 429:
+                    errors += 1
+                self.logger.debug(f"getCohortMembers({cohort_id}): request-status error {errors} - {e}")
+                if errors >= 3 or isinstance(e, CohortNotModifiedException) or isinstance(e, CohortTooLargeException):
+                    raise e
+            time.sleep(self.cohort_request_delay_millis/1000)
+
+    def _get_cohort_members_request(self, cohort_id: str, last_modified: int) -> HTTPResponse:
+        headers = {
+            'Authorization': f'Basic {self._get_basic_auth()}',
+        }
+        conn = self._connection_pool.acquire()
+        try:
+            url = f'/sdk/v1/cohort/{cohort_id}?maxCohortSize={self.max_cohort_size}'
+            if last_modified is not None:
+                url += f'&lastModified={last_modified}'
+            response = conn.request('GET', url, headers=headers)
+            return response
+        finally:
+            self._connection_pool.release(conn)
+
+    def _get_basic_auth(self) -> str:
+        credentials = f'{self.api_key}:{self.secret_key}'
+        return base64.b64encode(credentials.encode('utf-8')).decode('utf-8')
+
+    def __setup_connection_pool(self):
+        scheme, _, host = self.cdn_server_url.split('/', 3)
+        timeout = 10
+        self._connection_pool = HTTPConnectionPool(host, max_size=10, idle_timeout=30, read_timeout=timeout,
+                                                   scheme=scheme)

src/amplitude_experiment/cohort/cohort_loader.py

@@ -0,0 +1,42 @@
+from typing import Dict, Set
+from concurrent.futures import ThreadPoolExecutor, Future
+import threading
+
+from .cohort import Cohort
+from .cohort_download_api import CohortDownloadApi
+from .cohort_storage import CohortStorage
+
+
+class CohortLoader:
+    def __init__(self, cohort_download_api: CohortDownloadApi, cohort_storage: CohortStorage):
+        self.cohort_download_api = cohort_download_api
+        self.cohort_storage = cohort_storage
+        self.jobs: Dict[str, Future] = {}
+        self.lock_jobs = threading.Lock()
+        self.executor = ThreadPoolExecutor(
+            max_workers=32,
+            thread_name_prefix='CohortLoaderExecutor'
+        )
+
+    def load_cohort(self, cohort_id: str) -> Future:
+        with self.lock_jobs:
+            if cohort_id not in self.jobs:
+                def task():
+                    try:
+                        cohort = self.download_cohort(cohort_id)
+                        self.cohort_storage.put_cohort(cohort)
+                    except Exception as e:
+                        raise e
+
+                future = self.executor.submit(task)
+                future.add_done_callback(lambda f: self._remove_job(cohort_id))
+                self.jobs[cohort_id] = future
+            return self.jobs[cohort_id]
+
+    def _remove_job(self, cohort_id: str):
+        if cohort_id in self.jobs:
+            del self.jobs[cohort_id]
+
+    def download_cohort(self, cohort_id: str) -> Cohort:
+        cohort = self.cohort_storage.get_cohort(cohort_id)
+        return self.cohort_download_api.get_cohort(cohort_id, cohort)

src/amplitude_experiment/cohort/cohort_storage.py

@@ -0,0 +1,66 @@
+from typing import Dict, Set, Optional
+from threading import RLock
+
+from .cohort import Cohort, USER_GROUP_TYPE
+
+
+class CohortStorage:
+    def get_cohort(self, cohort_id: str):
+        raise NotImplementedError
+
+    def get_cohorts(self):
+        raise NotImplementedError
+
+    def get_cohorts_for_user(self, user_id: str, cohort_ids: Set[str]) -> Set[str]:
+        raise NotImplementedError
+
+    def get_cohorts_for_group(self, group_type: str, group_name: str, cohort_ids: Set[str]) -> Set[str]:
+        raise NotImplementedError
+
+    def put_cohort(self, cohort_description: Cohort):
+        raise NotImplementedError
+
+    def delete_cohort(self, group_type: str, cohort_id: str):
+        raise NotImplementedError
+
+
+class InMemoryCohortStorage(CohortStorage):
+    def __init__(self):
+        self.lock = RLock()
+        self.group_to_cohort_store: Dict[str, Set[str]] = {}
+        self.cohort_store: Dict[str, Cohort] = {}
+
+    def get_cohort(self, cohort_id: str):
+        with self.lock:
+            return self.cohort_store.get(cohort_id)
+
+    def get_cohorts(self):
+        return self.cohort_store.copy()
+
+    def get_cohorts_for_user(self, user_id: str, cohort_ids: Set[str]) -> Set[str]:
+        return self.get_cohorts_for_group(USER_GROUP_TYPE, user_id, cohort_ids)
+
+    def get_cohorts_for_group(self, group_type: str, group_name: str, cohort_ids: Set[str]) -> Set[str]:
+        result = set()
+        with self.lock:
+            group_type_cohorts = self.group_to_cohort_store.get(group_type, {})
+            for cohort_id in group_type_cohorts:
+                members = self.cohort_store.get(cohort_id).member_ids
+                if cohort_id in cohort_ids and group_name in members:
+                    result.add(cohort_id)
+        return result
+
+    def put_cohort(self, cohort: Cohort):
+        with self.lock:
+            if cohort.group_type not in self.group_to_cohort_store:
+                self.group_to_cohort_store[cohort.group_type] = set()
+            self.group_to_cohort_store[cohort.group_type].add(cohort.id)
+            self.cohort_store[cohort.id] = cohort
+
+    def delete_cohort(self, group_type: str, cohort_id: str):
+        with self.lock:
+            group_cohorts = self.group_to_cohort_store.get(group_type, {})
+            if cohort_id in group_cohorts:
+                group_cohorts.remove(cohort_id)
+            if cohort_id in self.cohort_store:
+                del self.cohort_store[cohort_id]

src/amplitude_experiment/cohort/cohort_sync_config.py

@@ -0,0 +1,7 @@
+class CohortSyncConfig:
+    def __init__(self, api_key: str, secret_key: str, max_cohort_size: int = 15000,
+                 cohort_request_delay_millis: int = 5000):
+        self.api_key = api_key
+        self.secret_key = secret_key
+        self.max_cohort_size = max_cohort_size
+        self.cohort_request_delay_millis = cohort_request_delay_millis

src/amplitude_experiment/deployment/deployment_runner.py

@@ -0,0 +1,111 @@
+import logging
+from typing import Optional, Set
+import threading
+
+from ..local.config import LocalEvaluationConfig
+from ..cohort.cohort_loader import CohortLoader
+from ..cohort.cohort_storage import CohortStorage
+from ..flag.flag_config_api import FlagConfigApi
+from ..flag.flag_config_storage import FlagConfigStorage
+from ..local.poller import Poller
+from ..util.flag_config import get_all_cohort_ids_from_flag
+
+
+class DeploymentRunner:
+    def __init__(
+            self,
+            config: LocalEvaluationConfig,
+            flag_config_api: FlagConfigApi,
+            flag_config_storage: FlagConfigStorage,
+            cohort_storage: CohortStorage,
+            cohort_loader: Optional[CohortLoader] = None,
+    ):
+        self.config = config
+        self.flag_config_api = flag_config_api
+        self.flag_config_storage = flag_config_storage
+        self.cohort_storage = cohort_storage
+        self.cohort_loader = cohort_loader
+        self.lock = threading.Lock()
+        self.poller = Poller(self.config.flag_config_polling_interval_millis / 1000, self.__periodic_refresh)
+        self.logger = logging.getLogger("Amplitude")
+        self.logger.addHandler(logging.StreamHandler())
+        if self.config.debug:
+            self.logger.setLevel(logging.DEBUG)
+
+    def start(self):
+        with self.lock:
+            self.refresh()
+            self.poller.start()
+
+    def stop(self):
+        self.poller.stop()
+
+    def __periodic_refresh(self):
+        try:
+            self.refresh()
+        except Exception as e:
+            self.logger.error(f"Refresh flag and cohort configs failed: {e}")
+
+    def refresh(self):
+        self.logger.debug("Refreshing flag configs.")
+        try:
+            flag_configs = self.flag_config_api.get_flag_configs()
+        except Exception as e:
+            self.logger.error(f'Failed to fetch flag configs: {e}')
+            raise Exception
+
+        flag_keys = {flag['key'] for flag in flag_configs}
+        self.flag_config_storage.remove_if(lambda f: f.key not in flag_keys)
+
+        for flag_config in flag_configs:
+            cohort_ids = get_all_cohort_ids_from_flag(flag_config)
+            if not self.cohort_loader or not cohort_ids:
+                self.logger.debug(f"Putting non-cohort flag {flag_config['key']}")
+                self.flag_config_storage.put_flag_config(flag_config)
+                continue
+
+            # Keep track of old flag and cohort for each flag
+            old_flag_config = self.flag_config_storage.get_flag_config(flag_config['key'])
+
+            try:
+                self._load_cohorts(flag_config, cohort_ids)
+                self.flag_config_storage.put_flag_config(flag_config)  # Store new flag config
+                self.logger.debug(f"Stored flag config {flag_config['key']}")
+
+            except Exception as e:
+                self.logger.warning(f"Failed to load all cohorts for flag {flag_config['key']}. "
+                                    f"Using the old flag config.")
+                self.flag_config_storage.put_flag_config(old_flag_config)
+                raise e
+
+        self._delete_unused_cohorts()
+        self.logger.debug(f"Refreshed {len(flag_configs)} flag configs.")
+
+    def _load_cohorts(self, flag_config: dict, cohort_ids: Set[str]):
+        def task():
+            try:
+                for cohort_id in cohort_ids:
+                    future = self.cohort_loader.load_cohort(cohort_id)
+                    future.result()
+                    self.logger.debug(f"Cohort {cohort_id} loaded for flag {flag_config['key']}")
+            except Exception as e:
+                self.logger.error(f"Failed to load cohorts for flag {flag_config['key']}: {e}")
+                raise e
+
+        cohort_fetched = self.cohort_loader.executor.submit(task)
+        # Wait for both flag and cohort loading to complete
+        cohort_fetched.result()
+
+    def _delete_unused_cohorts(self):
+        flag_cohort_ids = set()
+        for flag in self.flag_config_storage.get_flag_configs().values():
+            flag_cohort_ids.update(get_all_cohort_ids_from_flag(flag))
+
+        storage_cohorts = self.cohort_storage.get_cohorts()
+        deleted_cohort_ids = set(storage_cohorts.keys()) - flag_cohort_ids
+
+        for deleted_cohort_id in deleted_cohort_ids:
+            deleted_cohort = storage_cohorts.get(deleted_cohort_id)
+            if deleted_cohort is not None:
+                self.cohort_storage.delete_cohort(deleted_cohort.group_type, deleted_cohort_id)
+

src/amplitude_experiment/exception.py

@@ -2,3 +2,19 @@ class FetchException(Exception):
     def __init__(self, status_code, message):
         super().__init__(message)
         self.status_code = status_code
+
+
+class CohortNotModifiedException(Exception):
+    def __init__(self, message):
+        super().__init__(message)
+
+
+class CohortTooLargeException(Exception):
+    def __init__(self, message):
+        super().__init__(message)
+
+
+class HTTPErrorResponseException(Exception):
+    def __init__(self, status_code, message):
+        super().__init__(message)
+        self.status_code = status_code