Merge pull request #559 from macrocosm-os/SN1-380-hardcode-api-key

Create SN1 Loadbalancer

Author: bkb2135

.env.api.example

@@ -1,6 +1,5 @@
 API_PORT = "42170" # Port for the API server
 API_HOST = "0.0.0.0" # Host for the API server
 SCORING_KEY = "123" # The scoring key for the validator (must match the scoring key in the .env.validator file)
-SCORE_ORGANICS = True # Whether to score organics
 VALIDATOR_API = "0.0.0.0:8094" # The validator API to forward responses to for scoring
 WORKERS=4

.env.validator.example

@@ -26,7 +26,6 @@ HF_TOKEN = "your_huggingface_token_here"
 
 # Scoring API (optional).
 DEPLOY_SCORING_API = true
-SCORING_ADMIN_KEY = "123456"
 SCORING_API_PORT = 8094
 # Scoring key must match the scoring key in the .env.api.
 # SCORING_KEY="..."

data/top100k_domains.csv

@@ -99997,4 +99997,4 @@
 "99996","tankspotter.com","4.51"
 "99997","targetshootingapp.com","4.51"
 "99998","tastytalegame.com","4.51"
-"99999","tbscan.com","4.51"
+"99999","tbscan.com","4.51"

neurons/validator.py

@@ -4,8 +4,13 @@
 import time
 
 import loguru
+import netaddr
+import requests
 import torch
 import wandb
+from bittensor.core.extrinsics.serving import serve_extrinsic
+
+from prompting.rewards.scoring import task_scorer
 
 # ruff: noqa: E402
 from shared import settings
@@ -34,7 +39,6 @@ async def spawn_loops(task_queue, scoring_queue, reward_events):
         # ruff: noqa: E402
         from prompting.llms.model_manager import model_scheduler
         from prompting.miner_availability.miner_availability import availability_checking_loop
-        from prompting.rewards.scoring import task_scorer
         from prompting.tasks.task_creation import task_loop
         from prompting.tasks.task_sending import task_sender
         from prompting.weight_setting.weight_setter import weight_setter
@@ -87,10 +91,25 @@ async def start():
         # TODO: We should not use 2 availability loops for each process, in reality
         # we should only be sharing the miner availability data between processes.
         from prompting.miner_availability.miner_availability import availability_checking_loop
-        from prompting.rewards.scoring import task_scorer
 
         asyncio.create_task(availability_checking_loop.start())
 
+        try:
+            external_ip = requests.get("https://checkip.amazonaws.com").text.strip()
+            netaddr.IPAddress(external_ip)
+
+            serve_success = serve_extrinsic(
+                subtensor=settings.shared_settings.SUBTENSOR,
+                wallet=settings.shared_settings.WALLET,
+                ip=external_ip,
+                port=settings.shared_settings.SCORING_API_PORT,
+                protocol=4,
+                netuid=settings.shared_settings.NETUID,
+            )
+
+            logger.debug(f"Serve success: {serve_success}")
+        except Exception as e:
+            logger.warning(f"Failed to serve scoring api to chain: {e}")
         await start_scoring_api(task_scorer, scoring_queue, reward_events)
 
         while True:

prompting/api/api.py

@@ -4,6 +4,8 @@
 
 from prompting.api.miner_availabilities.api import router as miner_availabilities_router
 from prompting.api.scoring.api import router as scoring_router
+
+# from prompting.rewards.scoring import task_scorer
 from shared import settings
 
 app = FastAPI()

prompting/api/scoring/api.py

@@ -1,7 +1,8 @@
+import time
 import uuid
 from typing import Any
 
-from fastapi import APIRouter, Depends, Header, HTTPException, Request
+from fastapi import APIRouter, Depends, HTTPException, Request
 from loguru import logger
 
 from prompting.datasets.random_website import DDGDatasetEntry
@@ -11,13 +12,38 @@
 from shared import settings
 from shared.base import DatasetEntry
 from shared.dendrite import DendriteResponseEvent
-from shared.epistula import SynapseStreamResult
+from shared.epistula import SynapseStreamResult, verify_signature
+from shared.settings import shared_settings
 
 router = APIRouter()
 
 
-def validate_scoring_key(api_key: str = Header(...)):
-    if api_key != settings.shared_settings.SCORING_KEY:
+async def verify_scoring_signature(request: Request):
+    signed_by = request.headers.get("Epistula-Signed-By")
+    signed_for = request.headers.get("Epistula-Signed-For")
+    if signed_for != shared_settings.WALLET.hotkey.ss58_address:
+        raise HTTPException(status_code=400, detail="Bad Request, message is not intended for self")
+    if signed_by != shared_settings.API_HOTKEY:
+        raise HTTPException(status_code=401, detail="Signer not the expected ss58 address")
+
+    body = await request.body()
+    now = time.time()
+    err = verify_signature(
+        request.headers.get("Epistula-Request-Signature"),
+        body,
+        request.headers.get("Epistula-Timestamp"),
+        request.headers.get("Epistula-Uuid"),
+        signed_for,
+        signed_by,
+        now,
+    )
+    if err:
+        logger.error(err)
+        raise HTTPException(status_code=400, detail=err)
+
+
+def validate_scoring_key(request: Request):
+    if request.headers.api_key != settings.shared_settings.SCORING_KEY:
         raise HTTPException(status_code=403, detail="Invalid API key")
 
 
@@ -27,54 +53,62 @@ def get_task_scorer(request: Request):
 
 @router.post("/scoring")
 async def score_response(
-    request: Request, api_key_data: dict = Depends(validate_scoring_key), task_scorer=Depends(get_task_scorer)
+    request: Request, api_key_data: dict = Depends(verify_scoring_signature), task_scorer=Depends(get_task_scorer)
 ):
+    logger.debug("Scoring Request received!!!!!!!!!!!!!!!!")
     model = None
+    logger.debug("Setted Model to None")
     payload: dict[str, Any] = await request.json()
+    logger.debug(f"Awaited body: {payload}")
     body = payload.get("body")
-    timeout = payload.get("timeout", settings.shared_settings.NEURON_TIMEOUT)
-    uids = payload.get("uid", [])
+    timeout = payload.get("timeout", shared_settings.NEURON_TIMEOUT)
+    uids = payload.get("uids", [])
     chunks = payload.get("chunks", {})
+    timings = payload.get("timings", {})
+    logger.debug("About to check chunks and uids")
     if not uids or not chunks:
         logger.error(f"Either uids: {uids} or chunks: {chunks} is not valid, skipping scoring")
         return
     uids = [int(uid) for uid in uids]
     model = body.get("model")
-    if model:
-        try:
-            llm_model = ModelZoo.get_model_by_id(model)
-        except Exception:
-            logger.warning(
-                f"Organic request with model {body.get('model')} made but the model cannot be found in model zoo. Skipping scoring."
-            )
+    logger.debug("About to check model")
+    if model and model != shared_settings.LLM_MODEL:
+        logger.error(f"Model {model} not available for scoring on this validator.")
         return
-    else:
-        llm_model = None
+    logger.debug("Model has been checked")
+    llm_model = ModelZoo.get_model_by_id(model)
+    logger.debug("Got LLM Model from ModelZoo")
     task_name = body.get("task")
+    logger.debug(f"Task name set: {task_name}")
+    logger.debug(f"Length pre-insertion: {len(task_scorer.scoring_queue)}")
     if task_name == "InferenceTask":
         logger.info(f"Received Organic InferenceTask with body: {body}")
         logger.info(f"With model of type {type(body.get('model'))}")
         organic_task = InferenceTask(
             messages=body.get("messages"),
             llm_model=llm_model,
-            llm_model_id=body.get("model"),
+            llm_model_id=llm_model,
             seed=int(body.get("seed", 0)),
-            sampling_params=body.get("sampling_parameters", settings.shared_settings.SAMPLING_PARAMS),
+            sampling_params=body.get("sampling_parameters", shared_settings.SAMPLING_PARAMS),
             query=body.get("messages"),
+            organic=True,
         )
         logger.info(f"Task created: {organic_task}")
+
         task_scorer.add_to_queue(
             task=organic_task,
             response=DendriteResponseEvent(
                 uids=uids,
                 stream_results=[SynapseStreamResult(accumulated_chunks=chunks.get(str(uid), None)) for uid in uids],
                 timeout=timeout,
+                stream_results_all_chunks_timings=[timings.get(str(uid), None) for uid in uids],
             ),
             dataset_entry=DatasetEntry(),
-            block=settings.shared_settings.METAGRAPH.block,
+            block=shared_settings.METAGRAPH.block,
             step=-1,
             task_id=str(uuid.uuid4()),
         )
+
     elif task_name == "WebRetrievalTask":
         logger.info(f"Received Organic WebRetrievalTask with body: {body}")
         try:
@@ -91,15 +125,14 @@ async def score_response(
                 query=search_term,
             ),
             response=DendriteResponseEvent(
-                uids=[uids],
-                stream_results=[
-                    SynapseStreamResult(accumulated_chunks=[chunk for chunk in chunks if chunk is not None])
-                ],
-                timeout=body.get("timeout", settings.shared_settings.NEURON_TIMEOUT),
+                uids=uids,
+                stream_results=[SynapseStreamResult(accumulated_chunks=chunks.get(str(uid), [])) for uid in uids],
+                timeout=body.get("timeout", shared_settings.NEURON_TIMEOUT),
             ),
             dataset_entry=DDGDatasetEntry(search_term=search_term),
-            block=settings.shared_settings.METAGRAPH.block,
+            block=shared_settings.METAGRAPH.block,
             step=-1,
             task_id=str(uuid.uuid4()),
         )
+    logger.debug(f"Length post-insertion: {len(task_scorer.scoring_queue)}")
     logger.info("Organic task appended to scoring queue")

prompting/llms/apis/sn19_wrapper.py

@@ -1,6 +1,7 @@
 import json
 
 import requests
+from loguru import logger
 from tenacity import retry, stop_after_attempt, wait_exponential
 
 from prompting.llms.apis.llm_messages import LLMMessages
@@ -9,7 +10,6 @@
 shared_settings = settings.shared_settings
 
 
-# TODO: key error in response.json() when response is 500
 @retry(stop=stop_after_attempt(3), wait=wait_exponential(multiplier=1, min=4, max=10))
 def chat_complete(
     messages: LLMMessages,
@@ -38,6 +38,10 @@ def chat_complete(
         "logprobs": logprobs,
     }
     response = requests.post(url, headers=headers, data=json.dumps(data), timeout=30)
+    if not response.status_code == 200:
+        logger.error(f"SN19 API returned status code {response.status_code}")
+        logger.error(f"Response: {response.text}")
+        raise Exception(f"SN19 API returned status code {response.status_code}")
     response_json = response.json()
     try:
         return response_json["choices"][0]["message"].get("content")

prompting/llms/hf_llm.py

@@ -5,8 +5,7 @@
 from loguru import logger
 from transformers import AutoModelForCausalLM, AutoTokenizer, PreTrainedModel, pipeline
 
-from shared import settings
-from shared.timer import Timer
+from shared.settings import shared_settings
 
 
 class ReproducibleHF:
@@ -31,7 +30,7 @@ def __init__(self, model_id="hugging-quants/Meta-Llama-3.1-70B-Instruct-AWQ-INT4
 
         self.llm = pipeline("text-generation", model=self.model, tokenizer=self.tokenizer)
 
-        self.sampling_params = settings.shared_settings.SAMPLING_PARAMS
+        self.sampling_params = shared_settings.SAMPLING_PARAMS
 
     @torch.inference_mode()
     def generate(self, messages: list[str] | list[dict], sampling_params=None, seed=None):
@@ -46,23 +45,22 @@ def generate(self, messages: list[str] | list[dict], sampling_params=None, seed=
             add_generation_prompt=True,
             return_tensors="pt",
             return_dict=True,
-        ).to(settings.shared_settings.NEURON_DEVICE)
+        ).to(shared_settings.NEURON_DEVICE)
 
         params = sampling_params if sampling_params else self.sampling_params
         filtered_params = {k: v for k, v in params.items() if k in self.valid_generation_params}
 
-        with Timer():
-            # Generate with optimized settings
-            outputs = self.model.generate(
-                **inputs.to(settings.shared_settings.NEURON_DEVICE),
-                **filtered_params,
-                eos_token_id=self.tokenizer.eos_token_id,
-            )
-
-            results = self.tokenizer.batch_decode(
-                outputs[:, inputs["input_ids"].shape[1] :],
-                skip_special_tokens=True,
-            )[0]
+        # Generate with optimized settings
+        outputs = self.model.generate(
+            **inputs.to(shared_settings.NEURON_DEVICE),
+            **filtered_params,
+            eos_token_id=self.tokenizer.eos_token_id,
+        )
+
+        results = self.tokenizer.batch_decode(
+            outputs[:, inputs["input_ids"].shape[1] :],
+            skip_special_tokens=True,
+        )[0]
 
         logger.debug(
             f"""{self.__class__.__name__} queried:

prompting/rewards/scoring.py

@@ -94,20 +94,21 @@ async def run_step(self) -> RewardLoggingEvent:
             f"Scored {scoring_config.task.__class__.__name__} {scoring_config.task.task_id} with model "
             f"{scoring_config.task.llm_model_id}"
         )
-        log_event(
-            RewardLoggingEvent(
-                response_event=scoring_config.response,
-                reward_events=reward_events,
-                reference=scoring_config.task.reference,
-                challenge=scoring_config.task.query,
-                task=scoring_config.task.name,
-                block=scoring_config.block,
-                step=scoring_config.step,
-                task_id=scoring_config.task_id,
-                task_dict=scoring_config.task.model_dump(),
-                source=scoring_config.dataset_entry.source,
+        if not scoring_config.task.organic:
+            log_event(
+                RewardLoggingEvent(
+                    response_event=scoring_config.response,
+                    reward_events=reward_events,
+                    reference=scoring_config.task.reference,
+                    challenge=scoring_config.task.query,
+                    task=scoring_config.task.name,
+                    block=scoring_config.block,
+                    step=scoring_config.step,
+                    task_id=scoring_config.task_id,
+                    task_dict=scoring_config.task.model_dump(),
+                    source=scoring_config.dataset_entry.source,
+                )
             )
-        )
         await asyncio.sleep(0.01)
 
 

prompting/tasks/base_task.py

@@ -33,6 +33,7 @@ class BaseTask(BaseModel, ABC):
     query: Any = None
     reference: Any = None
     task_id: str = Field(default_factory=lambda: str(uuid4()), allow_mutation=False)
+    organic: bool = False
 
     model_config = ConfigDict(arbitrary_types_allowed=True)
 
@@ -60,6 +61,7 @@ class BaseTextTask(BaseTask):
     sampling_params: dict[str, float] = settings.shared_settings.SAMPLING_PARAMS
     timeout: int = settings.shared_settings.NEURON_TIMEOUT
     max_tokens: int = settings.shared_settings.NEURON_MAX_TOKENS
+    organic: bool = True
 
     @property
     def task_messages(self) -> list[str] | list[dict]: