ability to list all nodes for orchestrator compute pool

Author: JannikSt

crates/prime-protocol-py/examples/orchestrator_list_nodes.py

@@ -0,0 +1,47 @@
+#!/usr/bin/env python3
+"""
+Example demonstrating how to list nodes for a specific pool using the OrchestratorClient.
+"""
+
+import os
+from primeprotocol import OrchestratorClient
+
+def main():
+    # Replace with your actual RPC URL and private key
+    RPC_URL = "http://localhost:8545"
+    PRIVATE_KEY = os.getenv("ORCHESTRATOR_PRIVATE_KEY")
+    DISCOVERY_URLS = ["http://localhost:8089"]  # Discovery service URLs
+    
+    # Create orchestrator client
+    orchestrator = OrchestratorClient(
+        rpc_url=RPC_URL,
+        private_key=PRIVATE_KEY,
+        discovery_urls=DISCOVERY_URLS
+    )
+    
+    # Initialize the orchestrator (without P2P for this example)
+    orchestrator.start()
+    
+    # List nodes for a specific pool (example pool ID: 0)
+    pool_id = 0
+    pool_nodes = orchestrator.list_nodes_for_pool(pool_id)
+    print(f"Nodes in pool {pool_id}: {len(pool_nodes)}")
+    
+    # Print details of all nodes in the pool
+    for i, node in enumerate(pool_nodes):
+        print(f"\nNode {i+1}:")
+        print(f"  ID: {node.id}")
+        print(f"  Provider Address: {node.provider_address}")
+        print(f"  IP Address: {node.ip_address}")
+        print(f"  Port: {node.port}")
+        print(f"  Pool ID: {node.compute_pool_id}")
+        print(f"  Validated: {node.is_validated}")
+        print(f"  Active: {node.is_active}")
+        if node.worker_p2p_id:
+            print(f"  Worker P2P ID: {node.worker_p2p_id}")
+    
+    # Stop the orchestrator
+    orchestrator.stop()
+
+if __name__ == "__main__":
+    main() 

crates/prime-protocol-py/src/common/mod.rs

@@ -0,0 +1,68 @@
+use pyo3::prelude::*;
+use shared::models::node::DiscoveryNode;
+
+/// Node details structure shared between validator and orchestrator
+#[pyclass]
+#[derive(Clone)]
+pub struct NodeDetails {
+    #[pyo3(get)]
+    pub id: String,
+    #[pyo3(get)]
+    pub provider_address: String,
+    #[pyo3(get)]
+    pub ip_address: String,
+    #[pyo3(get)]
+    pub port: u16,
+    #[pyo3(get)]
+    pub compute_pool_id: u32,
+    #[pyo3(get)]
+    pub is_validated: bool,
+    #[pyo3(get)]
+    pub is_active: bool,
+    #[pyo3(get)]
+    pub is_provider_whitelisted: bool,
+    #[pyo3(get)]
+    pub is_blacklisted: bool,
+    #[pyo3(get)]
+    pub worker_p2p_id: Option<String>,
+    #[pyo3(get)]
+    pub last_updated: Option<String>,
+    #[pyo3(get)]
+    pub created_at: Option<String>,
+}
+
+impl From<DiscoveryNode> for NodeDetails {
+    fn from(node: DiscoveryNode) -> Self {
+        Self {
+            id: node.node.id,
+            provider_address: node.node.provider_address,
+            ip_address: node.node.ip_address,
+            port: node.node.port,
+            compute_pool_id: node.node.compute_pool_id,
+            is_validated: node.is_validated,
+            is_active: node.is_active,
+            is_provider_whitelisted: node.is_provider_whitelisted,
+            is_blacklisted: node.is_blacklisted,
+            worker_p2p_id: node.node.worker_p2p_id,
+            last_updated: node.last_updated.map(|dt| dt.to_rfc3339()),
+            created_at: node.created_at.map(|dt| dt.to_rfc3339()),
+        }
+    }
+}
+
+#[pymethods]
+impl NodeDetails {
+    /// Get compute specifications as a Python dictionary
+    pub fn get_compute_specs(&self, py: Python) -> PyResult<PyObject> {
+        // This method would need access to the original node data
+        // For now, return None since we don't store compute specs in NodeDetails
+        Ok(py.None())
+    }
+
+    /// Get location information as a Python dictionary
+    pub fn get_location(&self, py: Python) -> PyResult<PyObject> {
+        // This method would need access to the original node data
+        // For now, return None since we don't store location in NodeDetails
+        Ok(py.None())
+    }
+}

crates/prime-protocol-py/src/lib.rs

@@ -1,8 +1,10 @@
+use crate::common::NodeDetails;
 use crate::orchestrator::OrchestratorClient;
 use crate::validator::ValidatorClient;
 use crate::worker::WorkerClient;
 use pyo3::prelude::*;
 
+mod common;
 mod constants;
 mod error;
 mod orchestrator;
@@ -17,5 +19,6 @@ fn primeprotocol(m: &Bound<'_, PyModule>) -> PyResult<()> {
     m.add_class::<WorkerClient>()?;
     m.add_class::<OrchestratorClient>()?;
     m.add_class::<ValidatorClient>()?;
+    m.add_class::<NodeDetails>()?;
     Ok(())
 }

crates/prime-protocol-py/src/orchestrator/mod.rs

@@ -1,3 +1,4 @@
+use crate::common::NodeDetails;
 use crate::p2p_handler::auth::AuthenticationManager;
 use crate::p2p_handler::message_processor::{MessageProcessor, MessageProcessorConfig};
 use crate::p2p_handler::{Message, MessageType, Service as P2PService};
@@ -12,6 +13,9 @@ use tokio::task::JoinHandle;
 use tokio_util::sync::CancellationToken;
 use url::Url;
 
+// Add new imports for discovery functionality
+use shared::discovery::fetch_nodes_from_discovery_urls;
+
 /// Prime Protocol Orchestrator Client - for managing and distributing tasks
 #[pyclass]
 pub struct OrchestratorClient {
@@ -24,13 +28,19 @@ pub struct OrchestratorClient {
     user_message_rx: Option<Arc<Mutex<Receiver<Message>>>>,
     message_processor_handle: Option<JoinHandle<()>>,
     peer_id: Option<PeerId>,
+    // Discovery service URLs
+    discovery_urls: Vec<String>,
 }
 
 #[pymethods]
 impl OrchestratorClient {
     #[new]
-    #[pyo3(signature = (rpc_url, private_key=None))]
-    pub fn new(rpc_url: String, private_key: Option<String>) -> PyResult<Self> {
+    #[pyo3(signature = (rpc_url, private_key=None, discovery_urls=vec!["http://localhost:8089".to_string()]))]
+    pub fn new(
+        rpc_url: String,
+        private_key: Option<String>,
+        discovery_urls: Vec<String>,
+    ) -> PyResult<Self> {
         let runtime = tokio::runtime::Builder::new_multi_thread()
             .enable_all()
             .build()
@@ -60,6 +70,7 @@ impl OrchestratorClient {
             user_message_rx: None,
             message_processor_handle: None,
             peer_id: None,
+            discovery_urls,
         })
     }
 
@@ -201,14 +212,28 @@ impl OrchestratorClient {
         Ok(())
     }
 
-    pub fn list_validated_nodes(&self) -> PyResult<Vec<String>> {
-        // TODO: Implement orchestrator node listing
-        Ok(vec![])
-    }
+    /// List nodes for a specific compute pool
+    pub fn list_nodes_for_pool(&self, py: Python, pool_id: u32) -> PyResult<Vec<NodeDetails>> {
+        let rt = self.get_or_create_runtime()?;
+
+        let wallet = self.wallet.as_ref().ok_or_else(|| {
+            PyErr::new::<pyo3::exceptions::PyRuntimeError, _>(
+                "Wallet not initialized. Provide private_key when creating client.",
+            )
+        })?;
+
+        let discovery_urls = self.discovery_urls.clone();
+        let route = format!("/api/pool/{}", pool_id);
+
+        let nodes = py.allow_threads(|| {
+            rt.block_on(async {
+                fetch_nodes_from_discovery_urls(&discovery_urls, &route, wallet)
+                    .await
+                    .map_err(|e| PyErr::new::<pyo3::exceptions::PyRuntimeError, _>(e.to_string()))
+            })
+        })?;
 
-    pub fn list_nodes_from_chain(&self) -> PyResult<Vec<String>> {
-        // TODO: Implement orchestrator node listing from chain
-        Ok(vec![])
+        Ok(nodes.into_iter().map(NodeDetails::from).collect())
     }
 }
 

crates/prime-protocol-py/src/validator/mod.rs

@@ -1,3 +1,4 @@
+use crate::common::NodeDetails;
 use crate::p2p_handler::auth::AuthenticationManager;
 use crate::p2p_handler::message_processor::{MessageProcessor, MessageProcessorConfig};
 use crate::p2p_handler::{Message, MessageType, Service as P2PService};
@@ -17,75 +18,9 @@ use tokio::task::JoinHandle;
 use tokio_util::sync::CancellationToken;
 use url::Url;
 
-/// Node details for validator operations
-#[pyclass]
-#[derive(Clone)]
-pub(crate) struct NodeDetails {
-    #[pyo3(get)]
-    pub id: String,
-    #[pyo3(get)]
-    pub provider_address: String,
-    #[pyo3(get)]
-    pub ip_address: String,
-    #[pyo3(get)]
-    pub port: u16,
-    #[pyo3(get)]
-    pub compute_pool_id: u32,
-    #[pyo3(get)]
-    pub is_validated: bool,
-    #[pyo3(get)]
-    pub is_active: bool,
-    #[pyo3(get)]
-    pub is_provider_whitelisted: bool,
-    #[pyo3(get)]
-    pub is_blacklisted: bool,
-    #[pyo3(get)]
-    pub worker_p2p_id: Option<String>,
-    #[pyo3(get)]
-    pub last_updated: Option<String>,
-    #[pyo3(get)]
-    pub created_at: Option<String>,
-}
-
-impl From<DiscoveryNode> for NodeDetails {
-    fn from(node: DiscoveryNode) -> Self {
-        Self {
-            id: node.node.id,
-            provider_address: node.node.provider_address,
-            ip_address: node.node.ip_address,
-            port: node.node.port,
-            compute_pool_id: node.node.compute_pool_id,
-            is_validated: node.is_validated,
-            is_active: node.is_active,
-            is_provider_whitelisted: node.is_provider_whitelisted,
-            is_blacklisted: node.is_blacklisted,
-            worker_p2p_id: node.node.worker_p2p_id,
-            last_updated: node.last_updated.map(|dt| dt.to_rfc3339()),
-            created_at: node.created_at.map(|dt| dt.to_rfc3339()),
-        }
-    }
-}
-
-#[pymethods]
-impl NodeDetails {
-    /// Get compute specs as a Python dictionary
-    pub fn get_compute_specs(&self, py: Python) -> PyResult<PyObject> {
-        // This would need access to the original DiscoveryNode's compute_specs
-        // For now returning None
-        Ok(py.None())
-    }
-
-    /// Get location as a Python dictionary
-    pub fn get_location(&self, py: Python) -> PyResult<PyObject> {
-        // This would need access to the original DiscoveryNode's location
-        // For now returning None
-        Ok(py.None())
-    }
-}
-
 /// Prime Protocol Validator Client - for validating nodes and tasks
 #[pyclass]
-pub(crate) struct ValidatorClient {
+pub struct ValidatorClient {
     runtime: Option<tokio::runtime::Runtime>,
     wallet: Option<Wallet>,
     discovery_urls: Vec<String>,