Add prometheus extrinsic

bittensor/_neuron/text/core_server/__init__.py

@@ -137,12 +137,6 @@ def __init__(
             logging_dir = config.neuron.full_path,
         )
 
-        # Init prometheus.
-        # By default we pick the prometheus port to be axon.port - 1000 so that we can match port to server.
-        bittensor.prometheus ( 
-            config = config,
-            port = config.prometheus.port if config.axon.port == bittensor.defaults.axon.port else config.axon.port - 1000
-        )
         # --- Setup prometheus summaries.
         # These will not be posted if the user passes --prometheus.level OFF
         registry = CollectorRegistry()
@@ -176,6 +170,14 @@ def __init__(
             )
         self.axon = axon
         self.query_data = {}
+
+        # Init prometheus.
+        # By default we pick the prometheus port to be axon.port - 1000 so that we can match port to server.
+        bittensor.prometheus ( 
+            config = config,
+            wallet = self.wallet,
+            port = config.prometheus.port if config.axon.port == bittensor.defaults.axon.port else config.axon.port - 1000
+        )
 
         # Verify subnet exists
         if self.config.subtensor.network == 'finney' and not self.subtensor.subnet_exists( netuid = self.config.netuid ):

bittensor/_prometheus/__init__.py

@@ -47,41 +47,81 @@ def __str__(self):
 
     def __new__( 
         cls,
+        wallet: 'bittensor.wallet',
         config: 'bittensor.config' = None,
         port: int = None,
-        level: Union[str, "prometheus.level"] = None
+        level: Union[str, "prometheus.level"] = None,
+        netuid: int = None,
+        network: str = None,
+        chain_endpoint: str = None,
+        subtensor: 'bittensor.subtensor' = None,
+        prompt: bool = False, 
     ):
         """ Instantiates a global prometheus DB which can be accessed by other processes.
             Each prometheus DB is designated by a port.
             Args:
+                wallet (:obj: `bittensor.wallet`, `required`):
+                    bittensor wallet object.
                 config (:obj:`bittensor.Config`, `optional`, defaults to bittensor.prometheus.config()):
                     A config namespace object created by calling bittensor.prometheus.config()
                 port (:obj:`int`, `optional`, defaults to bittensor.defaults.prometheus.port ):
                     The port to run the prometheus DB on, this uniquely identifies the prometheus DB.
                 level (:obj:`prometheus.level`, `optional`, defaults to bittensor.defaults.prometheus.level ):
                     Prometheus logging level. If OFF, the prometheus DB is not initialized.
+                netuid (:obj: `int`, `optional`):
+                    network uid to serve on.
+                subtensor (:obj:`bittensor.Subtensor`, `optional`): 
+                    Chain connection through which to serve.
+                network (default='local', type=str)
+                    If subtensor is not set, uses this network flag to create the subtensor connection.
+                chain_endpoint (default=None, type=str)
+                    Overrides the network argument if not set.
+                prompt (:obj:`bool`, `optional`):
+                    If true, the call waits for confirmation from the user before proceeding.
         """
         if config == None:
             config = prometheus.config()
 
         if isinstance(level, prometheus.level):
             level = level.name # Convert ENUM to str.
 
+        if subtensor == None: subtensor = bittensor.subtensor( network = network, chain_endpoint = chain_endpoint) 
+
         config.prometheus.port = port if port != None else config.prometheus.port
         config.prometheus.level = level if level != None else config.prometheus.level
         cls.check_config( config )
-        if config.prometheus.level != prometheus.level.OFF.name:
+
+        return cls.serve(
+            cls,
+            wallet = wallet,
+            netuid = netuid,
+            subtensor = subtensor,
+            port = config.prometheus.port,
+            level = config.prometheus.level,
+            prompt = prompt,
+        )
+
+    def serve(cls, wallet, subtensor, netuid, port, level, prompt):
+        serve_success = subtensor.serve_prometheus(
+            wallet = wallet,
+            port = port,
+            netuid = netuid,
+            prompt = prompt
+        )
+        if serve_success and (level != prometheus.level.OFF.name):
             try:
-                start_http_server( config.prometheus.port )
+                start_http_server( port )
             except OSError:
                 # The singleton process is likely already running.
-                logger.error( "Prometheus:".ljust(20) + "<blue>{}</blue>  <red>already in use</red> ".format( config.prometheus.port ) )
-                return
+                logger.error( "Prometheus:".ljust(20) + "<blue>{}</blue>  <red>already in use</red> ".format( port ) )
             prometheus.started = True
-            prometheus.port = config.prometheus.port
-            logger.success( "Prometheus:".ljust(20) + "<green>ON</green>".ljust(20) + "using: <blue>[::]:{}</blue>".format( config.prometheus.port ))
+            prometheus.port = port
+            logger.success( "Prometheus:".ljust(20) + "<green>ON</green>".ljust(20) + "using: <blue>[::]:{}</blue>".format( port ))
+            return True
         else:
             logger.success('Prometheus:'.ljust(20) + '<red>OFF</red>')
+            raise RuntimeError('Failed to serve neuron.')
+            return False
 
     @classmethod
     def config(cls) -> 'bittensor.Config':

bittensor/_subtensor/extrinsics/prometheus.py

@@ -0,0 +1,99 @@
+# The MIT License (MIT)
+# Copyright © 2021 Yuma Rao
+# Copyright © 2023 Opentensor Foundation
+
+# Permission is hereby granted, free of charge, to any person obtaining a copy of this software and associated 
+# documentation files (the “Software”), to deal in the Software without restriction, including without limitation 
+# the rights to use, copy, modify, merge, publish, distribute, sublicense, and/or sell copies of the Software, 
+# and to permit persons to whom the Software is furnished to do so, subject to the following conditions:
+
+# The above copyright notice and this permission notice shall be included in all copies or substantial portions of 
+# the Software.
+
+# THE SOFTWARE IS PROVIDED “AS IS”, WITHOUT WARRANTY OF ANY KIND, EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO
+# THE WARRANTIES OF MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL 
+# THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION 
+# OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER 
+# DEALINGS IN THE SOFTWARE.
+import bittensor
+
+import json
+from rich.prompt import Confirm
+import bittensor.utils.networking as net
+from ..errors import *
+
+def prometheus_extrinsic(
+    subtensor: 'bittensor.Subtensor',
+    wallet: 'bittensor.wallet',
+    port: int, 
+    netuid: int,
+    ip: int = None,
+    wait_for_inclusion: bool = False,
+    wait_for_finalization = True,
+) -> bool:
+    r""" Subscribes an bittensor endpoint to the substensor chain.
+    Args:
+        subtensor (bittensor.subtensor):
+            bittensor subtensor object.
+        wallet (bittensor.wallet):
+            bittensor wallet object.
+        ip (str):
+            endpoint host port i.e. 192.122.31.4
+        port (int):
+            endpoint port number i.e. 9221
+        netuid (int):
+            network uid to serve on.
+        wait_for_inclusion (bool):
+            if set, waits for the extrinsic to enter a block before returning true, 
+            or returns false if the extrinsic fails to enter the block within the timeout.   
+        wait_for_finalization (bool):
+            if set, waits for the extrinsic to be finalized on the chain before returning true,
+            or returns false if the extrinsic fails to be finalized within the timeout.
+        prompt (bool):
+            If true, the call waits for confirmation from the user before proceeding.
+    Returns:
+        success (bool):
+            flag is true if extrinsic was finalized or uncluded in the block. 
+            If we did not wait for finalization / inclusion, the response is true.
+    """
+
+    # ---- Get external ip ----
+    if ip == None:
+        try:
+            external_ip = net.get_external_ip()
+            bittensor.__console__.print(":white_heavy_check_mark: [green]Found external ip: {}[/green]".format( external_ip ))
+            bittensor.logging.success(prefix = 'External IP', sufix = '<blue>{}</blue>'.format( external_ip ))
+        except Exception as E:
+            raise RuntimeError('Unable to attain your external ip. Check your internet connection. error: {}'.format(E)) from E
+    else:
+        external_ip = ip
+
+    call_params={
+        'version': bittensor.__version_as_int__, 
+        'ip': net.ip_to_int(external_ip), 
+        'port': port, 
+        'ip_type': net.ip_version(external_ip),
+    }
+
+    with bittensor.__console__.status(":satellite: Serving prometheus on: [white]{}:{}[/white] ...".format(subtensor.network, netuid)):
+        with subtensor.substrate as substrate:
+            call = substrate.compose_call(
+                call_module='SubtensorModule',
+                call_function='serve_prometheus',
+                call_params = call_params
+            )
+            extrinsic = substrate.create_signed_extrinsic( call = call, keypair = wallet.hotkey)
+            response = substrate.submit_extrinsic( extrinsic, wait_for_inclusion = wait_for_inclusion, wait_for_finalization = wait_for_finalization )
+            if wait_for_inclusion or wait_for_finalization:
+                response.process_events()
+                if response.is_success:
+                    bittensor.__console__.print(':white_heavy_check_mark: [green]Served prometheus[/green]\n  [bold white]{}[/bold white]'.format(
+                        json.dumps(call_params, indent=4, sort_keys=True)
+                    ))
+                    return True
+                else:
+                    bittensor.__console__.print(':cross_mark: [green]Failed to serve prometheus[/green] error: {}'.format(response.error_message))
+                    return False
+            else:
+                return True
+

bittensor/_subtensor/subtensor_impl.py

@@ -35,6 +35,7 @@
 from .extrinsics.registration import register_extrinsic
 from .extrinsics.transfer import transfer_extrinsic
 from .extrinsics.set_weights import set_weights_extrinsic
+from .extrinsics.prometheus import prometheus_extrinsic
 from .extrinsics.delegation import delegate_extrinsic, nominate_extrinsic,undelegate_extrinsic
 
 # Logging
@@ -251,6 +252,16 @@ def serve_axon (
     ) -> bool:
         return serve_axon_extrinsic( self, axon, use_upnpc, wait_for_inclusion, wait_for_finalization)
 
+    def serve_prometheus (
+        self,
+        wallet: 'bittensor.wallet',
+        port: int,
+        netuid: int,
+        wait_for_inclusion: bool = False,
+        wait_for_finalization: bool = True,
+        prompt: bool = False,
+    ) -> bool:
+        return prometheus_extrinsic( self, wallet = wallet, port = port, netuid = netuid, wait_for_inclusion = wait_for_inclusion, wait_for_finalization = wait_for_finalization, prompt = prompt)
     #################
     #### Staking ####
     #################

tests/integration_tests/test_prometheus.py

@@ -1,4 +1,46 @@
 import bittensor
+import torch
+import unittest
+import pytest 
+from bittensor._subtensor.subtensor_mock import mock_subtensor
 
-def test_init_prometheus():
-    bittensor.prometheus()
+import random
+import time
+import unittest
+from queue import Empty as QueueEmpty
+from unittest.mock import MagicMock, patch
+
+import bittensor
+import pytest
+from bittensor._subtensor.subtensor_mock import mock_subtensor
+from bittensor.utils.balance import Balance
+from substrateinterface import Keypair
+
+class TestPrometheus(unittest.TestCase):
+
+    def setUp(self):
+        class success():
+            def __init__(self):
+                self.is_success = True
+            def process_events(self):
+                return True
+        class fail():
+            def __init__(self):
+                self.is_success = False
+                self.error_message = 'Mock failure'
+            def process_events(self):
+                return True
+
+        self.subtensor = bittensor.subtensor(network = 'finney', chain_endpoint = 'wss://public.finney.opentensor.ai:443')
+        self.wallet = bittensor.wallet.mock()
+        self.success = success()
+        self.fail = fail()
+
+    def test_init_prometheus_success(self):
+        self.subtensor.substrate.submit_extrinsic = MagicMock(return_value = self.success) 
+        assert bittensor.prometheus(wallet = self.wallet, subtensor = self.subtensor)
+
+    def test_init_prometheus_failed(self):
+        self.subtensor.substrate.submit_extrinsic = MagicMock(return_value = self.fail) 
+        assert False == bittensor.prometheus(wallet = self.wallet, subtensor = self.subtensor)
+