feat: allow to specify custom wallet names

vegaprotocol · Sep 14, 2023 · f0ef68b · f0ef68b
1 parent 6c69268
commit f0ef68b
Show file tree

Hide file tree

Showing 2 changed files with 62 additions and 46 deletions.
diff --git a/vega_sim/devops/scenario.py b/vega_sim/devops/scenario.py
@@ -38,14 +38,8 @@
 )
 from vega_sim.scenario.configurable_market.agents import ConfigurableMarketManager
 from vega_sim.api.market import MarketConfig
-from vega_sim.devops.wallet import (
-    MARKET_CREATOR_AGENT,
-    MARKET_SETTLER_AGENT,
-    MARKET_MAKER_AGENT,
-    AUCTION_TRADER_AGENTS,
-    RANDOM_TRADER_AGENTS,
-    SENSITIVE_TRADER_AGENTS,
-)
+from vega_sim.devops.wallet import ScenarioWallet, default_scenario_wallet
+
 from vega_sim.devops.classes import (
     MarketMakerArgs,
     MarketManagerArgs,
@@ -69,6 +63,7 @@ def __init__(
         state_extraction_fn: Optional[
             Callable[[VegaServiceNull, Dict[str, Agent]], Any]
         ] = None,
+        scenario_wallet: Optional[ScenarioWallet] = None,
         step_length_seconds: float = 10,
         market_name: Optional[str] = None,
     ):
@@ -86,6 +81,8 @@ def __init__(
 
         self.step_length_seconds = step_length_seconds
         self.market_name = market_name
+        self.scenario_wallet = scenario_wallet if scenario_wallet is not None else default_scenario_wallet()
+
 
     def _get_historic_price_process(
         self,
@@ -133,13 +130,19 @@ def configure_agents(
         else:
             self.price_process = LivePrice(product=self.binance_code)
 
+        if self.scenario_wallet.market_creator_agent is None:
+            raise ValueError(f"Missing market_creator wallet for the {self.market_name} devops scenario")
+
+        if self.scenario_wallet.market_maker_agent is None:
+            raise ValueError(f"Missing market_maker wallet for the {self.market_name} devops scenario")
+
         if kwargs.get("run_background", True):
             # Setup agent for proposing and settling the market
             market_manager = ConfigurableMarketManager(
-                proposal_wallet_name=MARKET_CREATOR_AGENT.wallet_name,
-                proposal_key_name=MARKET_CREATOR_AGENT.key_name,
-                termination_wallet_name=MARKET_SETTLER_AGENT.wallet_name,
-                termination_key_name=MARKET_SETTLER_AGENT.key_name,
+                proposal_wallet_name=self.scenario_wallet.market_creator_agent.wallet_name,
+                proposal_key_name=self.scenario_wallet.market_creator_agent.key_name,
+                termination_wallet_name=self.scenario_wallet.market_creator_agent.wallet_name,
+                termination_key_name=self.scenario_wallet.market_creator_agent.key_name,
                 market_config=MarketConfig(),
                 market_name=self.market_name
                 if self.market_name is not None
@@ -154,8 +157,8 @@ def configure_agents(
 
             # Setup agent for proving a market for traders
             market_maker = ExponentialShapedMarketMaker(
-                wallet_name=MARKET_MAKER_AGENT.wallet_name,
-                key_name=MARKET_MAKER_AGENT.key_name,
+                wallet_name=self.scenario_wallet.market_maker_agent.wallet_name,
+                key_name=self.scenario_wallet.market_maker_agent.key_name,
                 market_name=self.market_name
                 if self.market_name is not None
                 else self.market_manager_args.market_name,
@@ -190,7 +193,7 @@ def configure_agents(
                     side=["SIDE_BUY", "SIDE_SELL"][i],
                     tag=i,
                 )
-                for i, party in enumerate(AUCTION_TRADER_AGENTS)
+                for i, party in enumerate(self.scenario_wallet.auction_trader_agents)
             ]
 
             # Setup agents for placing random market orders
@@ -209,7 +212,7 @@ def configure_agents(
                     step_bias=self.random_trader_args.step_bias[i],
                     tag=i,
                 )
-                for i, party in enumerate(RANDOM_TRADER_AGENTS)
+                for i, party in enumerate(self.scenario_wallet.random_trader_agents)
             ]
 
             # Setup agents for placing price-sensitive limit orders
@@ -227,7 +230,7 @@ def configure_agents(
                     max_order_size=self.sensitive_trader_args.max_order_size[i],
                     tag=i,
                 )
-                for i, party in enumerate(SENSITIVE_TRADER_AGENTS)
+                for i, party in enumerate(self.scenario_wallet.sensitive_trader_agents)
             ]
 
             agents = (

diff --git a/vega_sim/devops/wallet.py b/vega_sim/devops/wallet.py
@@ -27,6 +27,8 @@
 import dotenv
 
 from dataclasses import dataclass
+from typing import Optional
+
 
 DEFAULT_USER_WALLET_NAME = "vega-market-sim"
 DEFAULT_USER_WALLET_PASS = "passphrase"
@@ -40,32 +42,43 @@ class Agent:
     wallet_pass: str = os.environ.get("VEGA_USER_WALLET_PASS", DEFAULT_USER_WALLET_PASS)
     key_name: str = None
 
-
-# Market proper creates and approves a market creation proposal
-MARKET_CREATOR_AGENT = Agent(key_name="market_creator")
-
-# Market terminator creates and approves a market settle proposal
-MARKET_SETTLER_AGENT = Agent(key_name="market_settler")
-
-# Market maker makes a market by providing liquidity and an order book
-MARKET_MAKER_AGENT = Agent(key_name="market_maker")
-
-# Auction traders provide trades in auction trading modes
-AUCTION_TRADER_AGENTS = [
-    Agent(key_name="auction_trader_a"),
-    Agent(key_name="auction_trader_b"),
-]
-
-# Random traders provide trades in continuous trading modes
-RANDOM_TRADER_AGENTS = [
-    Agent(key_name="random_trader_a"),
-    Agent(key_name="random_trader_b"),
-    Agent(key_name="random_trader_c"),
-]
-
-# Sensitive traders exploit over-exposed positions in continuous trading modes
-SENSITIVE_TRADER_AGENTS = [
-    Agent(key_name="sensitive_trader_a"),
-    Agent(key_name="sensitive_trader_b"),
-    Agent(key_name="sensitive_trader_c"),
-]
+@dataclass
+class ScenarioWallet:
+    market_creator_agent: Optional[Agent]
+    market_settler_agent: Optional[Agent]
+    market_maker_agent: Optional[Agent]
+    auction_trader_agents: list[Agent]
+    random_trader_agents: list[Agent]
+    sensitive_trader_agents: list[Agent]
+
+def default_scenario_wallet() -> ScenarioWallet:
+    return ScenarioWallet(
+        # Market proper creates and approves a market creation proposal
+        market_creator_agent=Agent(key_name="market_creator"),
+
+        # Market terminator creates and approves a market settle proposal
+        market_settler_agent=Agent(key_name="market_settler"),
+
+        # Market maker makes a market by providing liquidity and an order book
+        market_maker_agent=Agent(key_name="market_maker"),
+
+        # Auction traders provide trades in auction trading modes
+        auction_trader_agents=[
+            Agent(key_name="auction_trader_a"),
+            Agent(key_name="auction_trader_b"),
+        ],
+
+        # Random traders provide trades in continuous trading modes
+        random_trader_agents=[
+            Agent(key_name="random_trader_a"),
+            Agent(key_name="random_trader_b"),
+            Agent(key_name="random_trader_c"),
+        ],
+
+        # Sensitive traders exploit over-exposed positions in continuous trading modes
+        sensitive_trader_agents=[
+            Agent(key_name="sensitive_trader_a"),
+            Agent(key_name="sensitive_trader_b"),
+            Agent(key_name="sensitive_trader_c"),
+        ],
+    )