Merge pull request #808 from macrocosm-os/release/v3.0.4

v3.0.4 Release

Author: bkb2135

.gitignore

@@ -117,3 +117,5 @@ cython_debug/
 
 # VS Code
 .vscode
+
+wandb/

apex/__init__.py

@@ -44,4 +44,4 @@ def setup_logger(log_file_path: str | Path | None = None, level: str = "INFO") -
     return logger
 
 
-setup_logger(log_file_path="logs.log", level="DEBUG")
+setup_logger(log_file_path="logs.log", level="INFO")

apex/common/config.py

@@ -15,6 +15,7 @@ class Config(BaseModel):
     chain: ConfigClass = Field(default_factory=ConfigClass)
     websearch: ConfigClass = Field(default_factory=ConfigClass)
     logger_db: ConfigClass = Field(default_factory=ConfigClass)
+    logger_wandb: ConfigClass = Field(default_factory=ConfigClass)
     weight_syncer: ConfigClass = Field(default_factory=ConfigClass)
     miner_sampler: ConfigClass = Field(default_factory=ConfigClass)
     miner_scorer: ConfigClass = Field(default_factory=ConfigClass)

apex/common/models.py

@@ -13,12 +13,14 @@ class MinerGeneratorResults(BaseModel):
     query: str
     generator_hotkeys: list[str]
     generator_results: list[str]
+    generator_times: list[float]
 
 
 class MinerDiscriminatorResults(BaseModel):
     query: str
     generator_hotkey: str
     generator_result: str
+    generator_time: float
     generator_score: float
     discriminator_hotkeys: list[str]
     discriminator_results: list[str]

apex/validator/logger_wandb.py

@@ -0,0 +1,69 @@
+from collections.abc import Mapping
+from typing import Any
+
+import wandb
+from loguru import logger
+
+from apex import __version__
+from apex.common.async_chain import AsyncChain
+from apex.common.models import MinerDiscriminatorResults
+
+
+def approximate_tokens(text: str) -> int:
+    """Count the number of tokens in a text."""
+    return len(text) // 4
+
+
+class LoggerWandb:
+    def __init__(
+        self,
+        async_chain: AsyncChain,
+        project: str = "apex-gan-arena",
+        api_key: str | None = None,
+    ):
+        self.run: Any | None = None
+        if project and api_key:
+            try:
+                # Authenticate with W&B, then initialize the run
+                wandb.login(key=api_key)
+                self.run = wandb.init(
+                    entity="macrocosmos",
+                    project=project,
+                    config={
+                        "hotkey": async_chain.wallet.hotkey.ss58_address,
+                        "netuid": async_chain.netuid,
+                        "version": __version__,
+                    },
+                )
+                logger.info(f"Initialized W&B run: {self.run.id}")
+            except Exception as e:
+                logger.error(f"Failed to initialize W&B run: {e}")
+        else:
+            logger.warning("W&B API key not provided, skipping logging to W&B")
+
+    async def log(
+        self,
+        reference: str | None = None,
+        discriminator_results: MinerDiscriminatorResults | None = None,
+        tool_history: list[dict[str, str]] | None = None,
+    ) -> None:
+        """Log an event to wandb."""
+        if self.run:
+            if discriminator_results:
+                processed_event = self.process_event(discriminator_results.model_dump())
+                processed_event["reference"] = reference
+                processed_event["tool_history"] = tool_history
+                self.run.log(processed_event)
+
+    def process_event(self, event: Mapping[str, Any]) -> dict[str, Any]:
+        """Preprocess an event before logging it."""
+        reference = event.get("reference", "")
+        generation = event.get("generation", "")
+        generator_tokens = approximate_tokens(generation)
+        reference_tokens = approximate_tokens(reference)
+
+        processed_event: dict[str, Any] = dict(event)
+        processed_event["generator_tokens"] = generator_tokens
+        processed_event["reference_tokens"] = reference_tokens
+
+        return processed_event

apex/validator/miner_sampler.py

@@ -156,19 +156,32 @@ async def query_miners(
             return ""
         return str(result)
 
+    async def query_miners_with_times(
+        self, body: dict[str, Any], endpoint: str, hotkey: str | None = None, timeout: float = TIMEOUT
+    ) -> tuple[str, float]:
+        """Query the miners for the query."""
+        start_time = time.time()
+        result = await self.query_miners(body, endpoint, hotkey, timeout)
+        return result, time.time() - start_time
+
     async def query_generators(self, query: str) -> MinerGeneratorResults:
         """Query the miners for the query."""
         miner_information = await self._sample_miners(sample_size=self._generator_sample_size)
         body = {"step": "generator", "query": query}
 
         hotkeys: list[str] = []
-        tasks: list[Coroutine[str, str, Any]] = []
+        tasks: list[Coroutine[tuple[str, float], str, Any]] = []
 
         for miner_info in miner_information:
             hotkeys.append(miner_info.hotkey)
-            tasks.append(self.query_miners(body=body, endpoint=miner_info.address, hotkey=miner_info.hotkey))
+            tasks.append(self.query_miners_with_times(body=body, endpoint=miner_info.address, hotkey=miner_info.hotkey))
         generator_results = await asyncio.gather(*tasks)
-        return MinerGeneratorResults(query=query, generator_hotkeys=hotkeys, generator_results=generator_results)
+        return MinerGeneratorResults(
+            query=query,
+            generator_hotkeys=hotkeys,
+            generator_results=[result[0] for result in generator_results],
+            generator_times=[result[1] for result in generator_results],
+        )
 
     async def query_discriminators(
         self,
@@ -181,19 +194,20 @@ async def query_discriminators(
         miner_information = await self._sample_miners(sample_size=self._discriminator_sample_size)
         # Flip the coin for the generator.
         if ground_truth and generator_results:
-            selected_generator: tuple[str, str] = random.choice(
+            selected_generator: tuple[str, str, float] = random.choice(
                 list(
                     zip(
                         generator_results.generator_hotkeys,
                         generator_results.generator_results,
+                        generator_results.generator_times,
                         strict=False,
                     )
                 )
             )
         else:
             if reference is None:
                 raise ValueError("Reference cannot be None when not using miner generator results")
-            selected_generator = (VALIDATOR_REFERENCE_LABEL, reference)
+            selected_generator = (VALIDATOR_REFERENCE_LABEL, reference, 0.0)
 
         body = {
             "step": "discriminator",
@@ -202,7 +216,7 @@ async def query_discriminators(
         }
 
         hotkeys: list[str] = []
-        tasks: list[Coroutine[str, str, Any]] = []
+        tasks: list[Coroutine[tuple[str, float], str, Any]] = []
         for miner_info in miner_information:
             hotkeys.append(miner_info.hotkey)
             tasks.append(self.query_miners(body=body, endpoint=miner_info.address, hotkey=miner_info.hotkey))
@@ -244,6 +258,7 @@ async def query_discriminators(
             generator_hotkey=selected_generator[0],
             generator_result=selected_generator[1],
             generator_score=generator_result_float,
+            generator_time=selected_generator[2],
             discriminator_hotkeys=hotkeys,
             discriminator_results=parsed_discriminator_results,
             discriminator_scores=discriminator_results_float,

apex/validator/pipeline.py

@@ -11,8 +11,8 @@
 from apex.services.llm.llm_base import LLMBase
 from apex.services.websearch.websearch_base import WebSearchBase
 from apex.validator import generate_query, generate_reference
-from apex.validator.logger_apex import LoggerApex
 from apex.validator.logger_local import LoggerLocal
+from apex.validator.logger_wandb import LoggerWandb
 from apex.validator.miner_sampler import MinerSampler
 
 
@@ -23,7 +23,7 @@ def __init__(
         miner_sampler: MinerSampler,
         llm: LLMBase,
         deep_research: DeepResearchBase,
-        logger_apex: LoggerApex | None = None,
+        logger_wandb: LoggerWandb | None = None,
         num_consumers: int = 5,
         timeout_consumer: float = 1200,
         timeout_producer: float = 240,
@@ -36,7 +36,7 @@ def __init__(
         self.miner_registry = miner_sampler
         self.llm = llm
         self.deep_research = deep_research
-        self.logger_apex = logger_apex
+        self.logger_wandb = logger_wandb
         self.num_consumers = num_consumers
         self.timeout_consumer = timeout_consumer
         self.timeout_producer = timeout_producer
@@ -109,8 +109,8 @@ async def run_single(self, task: QueryTask) -> str:
             query=query, generator_results=generator_results, reference=reference, ground_truth=ground_truth
         )
 
-        if self.logger_apex:
-            await self.logger_apex.log(
+        if self.logger_wandb:
+            await self.logger_wandb.log(
                 reference=reference, discriminator_results=discriminator_results, tool_history=tool_history
             )
 

config/mainnet.yaml.example

@@ -12,6 +12,11 @@ websearch:
   kwargs:
     key: "TAVILY_API_KEY"
 
+logger_wandb:
+  kwargs:
+    project: "apex-gan-arena"
+    api_key: "YOUR_WANDB_API_KEY"
+
 llm:
   kwargs:
     key: "CHUTES_API_KEY"

config/testnet.yaml.example

@@ -10,6 +10,11 @@ websearch:
   kwargs:
     key: "TAVILY_API_KEY"
 
+logger_wandb:
+  kwargs:
+    project: "apex-gan-arena"
+    api_key: "YOUR_WANDB_API_KEY"
+
 llm:
   kwargs:
     key: "CHUTES_API_KEY"

pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "apex"
-version = "3.0.3"
+version = "3.0.4"
 description = "Bittensor Subnet 1: Apex"
 readme = "README.md"
 requires-python = "~=3.11"
@@ -34,6 +34,8 @@ dependencies = [
     "types-cachetools>=6.0.0.20250525",
     "dotenv>=0.9.9",
     "pytest-mock>=3.14.1",
+    "wandb>=0.21.1",
+    "ruff>=0.12.5",
 ]