Custom usage counter support

tests/props/test_com.py

@@ -13,8 +13,8 @@ def test_com_linear_fixed_price():
 
 def test_com_linear_price_for():
     com: ComLinear = ComLinearFactory(linear_coeffs=LINEAR_COEFFS, usage_vector=DEFINED_USAGES)
-    assert com.price_for[Counter.CPU] == LINEAR_COEFFS[0]
-    assert com.price_for[Counter.TIME] == LINEAR_COEFFS[1]
+    assert com.price_for[Counter.CPU.value] == LINEAR_COEFFS[0]
+    assert com.price_for[Counter.TIME.value] == LINEAR_COEFFS[1]
 
 
 @pytest.mark.parametrize(

tests/test_strategy.py

@@ -160,18 +160,6 @@ async def test_same_score(self):
 
         assert await self.strategy.score_offer(offer1) == await self.strategy.score_offer(offer2)
 
-    @pytest.mark.asyncio
-    async def test_score_unknown_price(self):
-        offer = OfferProposalFactory(
-            **{
-                "proposal__proposal__properties__usage_vector": [
-                    Counter.MAXMEM.value,
-                    Counter.TIME.value,
-                ]
-            }
-        )
-        assert await self.strategy.score_offer(offer) == SCORE_REJECTED
-
     @pytest.mark.asyncio
     @pytest.mark.parametrize(
         "coeffs", [[-0.001, 0.002, 0.1], [0.001, -0.002, 0.1], [0.001, 0.002, -0.1]]

yapapi/ctx.py

@@ -575,5 +575,5 @@ def is_streaming(self) -> bool:
 class ActivityUsage:
     """A high-level representation of activity usage record."""
 
-    current_usage: Dict[Counter, float] = field(default_factory=dict)
+    current_usage: Dict[str, float] = field(default_factory=dict)
     timestamp: Optional[datetime] = None

yapapi/props/com.py

@@ -70,15 +70,12 @@ def fixed_price(self) -> float:
         return self.linear_coeffs[-1]
 
     @property
-    def price_for(self) -> Dict[Counter, float]:
-        return {
-            Counter(self.usage_vector[i]): self.linear_coeffs[i]
-            for i in range(len(self.usage_vector))
-        }
+    def price_for(self) -> Dict[str, float]:
+        return {u: self.linear_coeffs[i] for (i, u) in enumerate(self.usage_vector)}
 
     def calculate_cost(self, usage: List):
         usage = usage + [1.0]  # append the "usage" of the fixed component
-        return sum([self.linear_coeffs[i] * usage[i] for i in range(len(self.linear_coeffs))])
+        return sum([c * usage[i] for (i, c) in enumerate(self.linear_coeffs)])
 
-    def usage_as_dict(self, usage: List) -> Dict[Counter, float]:
-        return {Counter(self.usage_vector[i]): usage[i] for i in range(len(usage))}
+    def usage_as_dict(self, usage: List) -> Dict[str, float]:
+        return {self.usage_vector[i]: u for (i, u) in enumerate(usage)}

yapapi/strategy.py

@@ -5,13 +5,14 @@
 from decimal import Decimal
 import logging
 from types import MappingProxyType
-from typing import Dict, Mapping, Optional
+from typing import Dict, Mapping, Optional, Union
 
-from dataclasses import dataclass, field
+from dataclasses import dataclass
 from typing_extensions import Final, Protocol
 
 from yapapi.props import com, Activity
 from yapapi.props.builder import DemandBuilder, DemandDecorator
+from yapapi.props.com import Counter
 from yapapi import rest
 
 
@@ -50,11 +51,21 @@ class DummyMS(MarketStrategy, object):
     For other offers, returns `SCORE_REJECTED`.
     """
 
-    max_for_counter: Mapping[com.Counter, Decimal] = MappingProxyType(
-        {com.Counter.TIME: Decimal("0.002"), com.Counter.CPU: Decimal("0.002") * 10}
-    )
-    max_fixed: Decimal = Decimal("0.05")
-    _activity: Optional[Activity] = field(init=False, repr=False, default=None)
+    def __init__(
+        self,
+        max_fixed_price: Decimal = Decimal("0.05"),
+        max_price_for: Mapping[Union[Counter, str], Decimal] = MappingProxyType({}),
+        activity: Optional[Activity] = None,
+    ):
+        self._max_fixed_price = max_fixed_price
+        self._max_price_for: Dict[str, Decimal] = defaultdict(lambda: Decimal("inf"))
+        self._max_price_for.update(
+            {com.Counter.TIME.value: Decimal("0.002"), com.Counter.CPU.value: Decimal("0.002") * 10}
+        )
+        self._max_price_for.update(
+            {(c.value if isinstance(c, Counter) else c): v for (c, v) in max_price_for.items()}
+        )
+        self._activity = activity
 
     async def decorate_demand(self, demand: DemandBuilder) -> None:
         """Ensure that the offer uses `PriceModel.LINEAR` price model."""
@@ -71,12 +82,12 @@ async def score_offer(
         if linear.scheme != com.BillingScheme.PAYU:
             return SCORE_REJECTED
 
-        if linear.fixed_price > self.max_fixed:
+        if linear.fixed_price > self._max_fixed_price:
             return SCORE_REJECTED
         for counter, price in linear.price_for.items():
-            if counter not in self.max_for_counter:
+            if counter not in self._max_price_for:
                 return SCORE_REJECTED
-            if price > self.max_for_counter[counter]:
+            if price > self._max_price_for[counter]:
                 return SCORE_REJECTED
 
         return SCORE_NEUTRAL
@@ -90,13 +101,15 @@ def __init__(
         self,
         expected_time_secs: int = 60,
         max_fixed_price: Decimal = Decimal("inf"),
-        max_price_for: Mapping[com.Counter, Decimal] = MappingProxyType({}),
+        max_price_for: Mapping[Union[Counter, str], Decimal] = MappingProxyType({}),
     ):
         self._expected_time_secs = expected_time_secs
         self._logger = logging.getLogger(f"{__name__}.{type(self).__name__}")
-        self._max_fixed_price = max_fixed_price if max_fixed_price is not None else Decimal("inf")
-        self._max_price_for: Dict[com.Counter, Decimal] = defaultdict(lambda: Decimal("inf"))
-        self._max_price_for.update(max_price_for)
+        self._max_fixed_price = max_fixed_price
+        self._max_price_for: Dict[str, Decimal] = defaultdict(lambda: Decimal("inf"))
+        self._max_price_for.update(
+            {(c.value if isinstance(c, Counter) else c): v for (c, v) in max_price_for.items()}
+        )
 
     async def decorate_demand(self, demand: DemandBuilder) -> None:
         """Ensure that the offer uses `PriceModel.LINEAR` price model."""
@@ -115,13 +128,6 @@ async def score_offer(
             )
             return SCORE_REJECTED
 
-        known_time_prices = {com.Counter.TIME, com.Counter.CPU}
-
-        for counter in linear.price_for.keys():
-            if counter not in known_time_prices:
-                self._logger.debug("Rejected offer %s: unsupported counter '%s'", offer.id, counter)
-                return SCORE_REJECTED
-
         if linear.fixed_price > self._max_fixed_price:
             self._logger.debug(
                 "Rejected offer %s: fixed price higher than fixed price cap %f.",
@@ -136,7 +142,7 @@ async def score_offer(
 
         expected_usage = []
 
-        for resource in [com.Counter(u) for u in linear.usage_vector]:
+        for resource in linear.usage_vector:
 
             if linear.price_for[resource] > self._max_price_for[resource]:
                 self._logger.debug(