ASV framework refactor proposal #2227
There are no files selected for viewing
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,116 @@ | ||
| import copy | ||
| from abc import ABC, abstractmethod | ||
| from enum import Enum | ||
| import logging | ||
| import os | ||
| import tempfile | ||
| import time | ||
| import re | ||
| import pandas as pd | ||
| import numpy as np | ||
| from typing import List, Union | ||
|
|
||
| from arcticdb.arctic import Arctic | ||
| from arcticdb.options import LibraryOptions | ||
| from arcticdb.storage_fixtures.s3 import BaseS3StorageFixtureFactory, real_s3_from_environment_variables | ||
| from arcticdb.util.utils import DFGenerator, ListGenerators, TimestampNumber | ||
| from arcticdb.version_store.library import Library | ||
|
|
||
| def get_logger_for_asv(bencmhark_cls): | ||
| logLevel = logging.INFO | ||
| logger = logging.getLogger(bencmhark_cls.__name__) | ||
| logger.setLevel(logLevel) | ||
| console_handler = logging.StreamHandler() | ||
| console_handler.setLevel(logLevel) | ||
| formatter = logging.Formatter('%(asctime)s - %(name)s - %(levelname)s - %(message)s') | ||
| console_handler.setFormatter(formatter) | ||
| logger.addHandler(console_handler) | ||
| return logger | ||
|
|
||
|
|
||
| class Storage(Enum): | ||
| AMAZON = 1 | ||
| LMDB = 2 | ||
|
|
||
|
|
||
| class LibraryType(Enum): | ||
| PERSISTENT = "PERMANENT" | ||
| MODIFIABLE = "MODIFIABLE" | ||
|
There was a problem hiding this comment. would advise adding the concept of TEST library type which is persistent space for tests. Setting then throgh a property of the generator class set_test() could automatically make the code execute on separate space and not polute PERSISTENT space. |
||
|
|
||
|
|
||
| # It is quite clear what is this responsible for: only dataframe generation | ||
| # Using such an abstraction can help us deduplicate the dataframe generation code between the different `EnvironmentSetup`s | ||
| # Note: We use a class instead of a generator function to allow caching of dataframes in the state | ||
| class DataFrameGenerator(ABC): | ||
|
There was a problem hiding this comment. Not necessarily opposed to using a class, but if caching is the only thing we need it for, we can use something like this instead. |
||
| @abstractmethod | ||
| def get_dataframe(self, **kwargs): | ||
| pass | ||
|
|
||
| # The population policy is tightly paired with the `populate_library` | ||
| # Using a separate abstraction for the population policy can allow us to be flexible with how we populate libraries. | ||
| # E.g. One benchmark can use one population policy for read operations, a different for updates and a third for appends | ||
| # Note: Using a kwargs generator to be passed to the df_generator allows customization of dataframe generation params (e.g. num_rows, num_cols, use_string_columns?) | ||
| class LibraryPopulationPolicy(ABC): | ||
|
There was a problem hiding this comment. We discussed with Ivo that in previous version the setup environment class was responsible for both library and symbol generation. In this version it will be responsible only for symbol population. Then ASV test will invoke it for each library it needs to populate. Perhaps this would be ammended also with other methods, attributes:
I would suggest that I do not beleive we will have single only way to There was a problem hiding this comment. I aggree as we discussed offline, we need more complex I'd be happy to move the Probably just a single Policy which allows many symbols / many versions / snapshot frequency can serve for most our tests |
||
| def __init__(self, num_symbols: int, df_generator: DataFrameGenerator): | ||
| self.num_symbols = num_symbols | ||
| self.df_generator = df_generator | ||
|
|
||
| @abstractmethod | ||
| def get_symbol_name(self, ind: int): | ||
| pass | ||
|
|
||
| @abstractmethod | ||
| def get_generator_kwargs(self, ind: int) -> Dict[str, Any]: | ||
| pass | ||
|
|
||
|
|
||
| # Currently we're using the same arctic client for both persistant and modifiable libraries. | ||
|
There was a problem hiding this comment. I would suggest having a separate class for those 3 functions : get_arctic_client Perhaps SharedStorageAccess or better name? |
||
| # We might decide that we want different arctic clients (e.g. different buckets) but probably not needed for now. | ||
| def get_arctic_client(storage: Storage) -> Arctic: | ||
| pass | ||
|
|
||
| def get_library_name(library_type, benchmark_cls, lib_name_suffix): | ||
| if library_type == LibraryType.PERSISTENT: | ||
| return f"{library_type}_{benchmark_cls.__name__}_{lib_name_suffix}" | ||
| if library_type == LibraryType.MODIFIABLE: | ||
| # We want the modifiable libraries to be unique per process/ benchmark class. We embed this deep in the name | ||
| return f"{library_type}_{benchmark_cls.__name__}_{os.getpid()}_{lib_name_suffix}" | ||
|
|
||
| def get_library(ac : Arctic, library_type : LibraryType, benchmark_cls, lib_name_suffix : str) -> Library: | ||
| lib_name = get_library_name(library_type, benchmark_cls, lib_name_suffix) | ||
| if library_type == LibraryType.PERSISTENT: | ||
| # TODO: Change to OpenMode=READ for persistent libraries | ||
| # This however is not available in the c++ library manager API (it currently hardcodes OpenMode=DELETE). | ||
| # We can expose it just for these tests if we decide so. but this is probably for a separate PR. | ||
| pass | ||
| return ac.get_library(lib_name, create_if_missing=True) | ||
|
|
||
|
|
||
| def populate_library(lib, population_policy : LibraryPopulationPolicy): | ||
| num_symbols = population_policy.num_symbols | ||
| df_generator = population_policy.df_generator | ||
| for i in range(num_symbols): | ||
| sym = population_policy.get_symbol(i) | ||
| kwargs = population_policy.get_generator_kwargs(i) | ||
| df = df_generator.get_dataframe(**kwargs) | ||
| lib.write(sym, df) | ||
|
|
||
| # This is the only API to populate a persistent library. If we deem useful we can also add a check whether library is valid (e.g. has the correct num_symbols) | ||
| # As discussed, ideally this will be called in completely separate logic from ASV tests to avoid races, but for simplicity | ||
| # for now we can just rely on setup_cache to populate the persistant libraries if they are missing. | ||
| def populate_persistent_library_if_missing(ac, benchmark_cls, lib_name_suffix, population_policy : LibraryPopulationPolicy): | ||
| lib_name = get_library_name(LibraryType.PERSISTENT, benchmark_cls, lib_name_suffix) | ||
| if ac.has_library(lib_name): | ||
| lib = ac.create_library(lib_name) | ||
| populate_library(lib, population_policy) | ||
|
|
||
| # Not sure how useful is this, one can always just keep track of the libraries created and clear them | ||
| def clear_all_modifiable_libs_from_this_process(ac, benchmark_cls): | ||
| lib_names = ac.list_libraries() | ||
| to_deletes = [lib_name for lib_name in lib_names if lib_name.startswith(f"{LibraryType.MODIFIABLE}_{benchmark_cls.__name__}_{os.getpid()}")] | ||
| for to_delete in to_deletes: | ||
| ac.delete_library(to_delete) | ||
|
|
||
|
|
||
|
|
||
|
|
||
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
I think that both versions are ok.
The original version main goal was to separate the creation structure from the ASV class. That is something that needed to be done because of the persostence library. Only this way the code for generating the environment could be safely tested without making changes to the tests - you fix the logic of generating the structure of objects outside of ASV tests then you go to ASV tests and use it wherever you want. So the class although monolythic has this main function. That makes the setup code llittle more hard to understand but easy to reuse outside of ASV as it does not follow ASV guidelines. But the benchamrk tests are mainly because of the tests not about the setup struct, so once developed the structure would be unwise to change due to loss of history for benchamrks.
The proposed version indeed makes at first glance perhaps things more clear during the review process as the code will be inside the ASV test, primarily. Having more decomposed objects perhaps would increase reusability. Still there are things that are missing from the model and the final solution would be more thick and perhaps again not quite clean here and there when missing things are added.
Thus for me it is ok to go to new way of constucting tests, just have to say that during adaptations of this approach to different benchamrks its cleanes might suffer.