feat(deployment): Clean backend IP fetching

backend/pyproject.toml

@@ -13,7 +13,7 @@ package-mode = false
 python = "^3.12"
 fastapi = { extras = ["standard"], version = "^0.115.2" }
 aleph-sdk-python = "^1.1.0"
-libertai-utils = "0.0.8"
+libertai-utils = "0.0.9"
 setuptools = "^75.4.0"
 paramiko = "^3.5.0"
 

backend/src/interfaces/agent.py

@@ -21,10 +21,6 @@ class SetupAgentBody(DeleteAgentBody):
     account: SubscriptionAccount
 
 
-class UpdateAgentResponse(BaseModel):
-    instance_hash: str
-
-
 class PublicAgentData(BaseModel):
     id: str
     subscription_id: str
@@ -47,12 +43,13 @@ class GetAgentSecretMessage(BaseModel):
 
 
 class GetAgentResponse(PublicAgentData):
-    pass
+    instance_ip: str | None
 
 
 class GetAgentSecretResponse(BaseModel):
     secret: str
 
+
 class AgentPythonPackageManager(str, Enum):
     poetry = "poetry"
     pip = "pip"

backend/src/interfaces/aleph.py

@@ -6,12 +6,3 @@ class AlephVolume(BaseModel):
     mount: str
     ref: str
     use_latest: bool
-
-
-class AlephNodeInfo:
-    def __init__(self, **kwargs):
-        self.data = kwargs.get("data", {})
-        self.nodes = self.data.get("corechannel", {}).get("resource_nodes", [])
-        self.nodes.sort(key=lambda x: x.get("score", 0), reverse=True)
-        self.core_node = self.data.get("corechannel", {}).get("nodes", [])
-        self.core_node.sort(key=lambda x: x.get("score", 0), reverse=True)

backend/src/main.py

@@ -12,6 +12,7 @@
 from aleph_message.models.execution.environment import HypervisorType
 from fastapi import FastAPI, File, Form, HTTPException, UploadFile
 from libertai_utils.chains.index import is_signature_valid
+from libertai_utils.interfaces.agent import UpdateAgentResponse
 from libertai_utils.interfaces.subscription import Subscription
 from libertai_utils.utils.crypto import decrypt, encrypt
 from starlette.middleware.cors import CORSMiddleware
@@ -25,7 +26,6 @@
     GetAgentSecretMessage,
     GetAgentSecretResponse,
     SetupAgentBody,
-    UpdateAgentResponse,
 )
 from src.utils.agent import fetch_agents
 from src.utils.aleph import fetch_instance_ip
@@ -116,9 +116,15 @@ async def get_agent_public_info(agent_id: str) -> GetAgentResponse:
         )
     agent = agents[0]
 
+    try:
+        ip_address = await fetch_instance_ip(agent.instance_hash)
+    except ValueError:
+        ip_address = None
+
     return GetAgentResponse(
         id=agent.id,
         instance_hash=agent.instance_hash,
+        instance_ip=ip_address,
         last_update=agent.last_update,
         subscription_id=agent.subscription_id,
     )
@@ -205,11 +211,10 @@ async def update(
 
     try:
         hostname = await fetch_instance_ip(agent.instance_hash)
-        print(hostname)
     except ValueError:
         raise HTTPException(
             status_code=HTTPStatus.INTERNAL_SERVER_ERROR,
-            detail="Instance IPv6 not found, it's probably not allocated yet.",
+            detail="Instance IPv6 address not found, it probably isn't allocated yet. Please try again in a few minutes.",
         )
 
     # Create a Paramiko SSH client
@@ -233,7 +238,7 @@ async def update(
 
     # Execute a command
     _stdin, stdout, stderr = ssh_client.exec_command(
-        f"wget {deploy_script_url} -O /tmp/deploy-agent.sh -q --no-cached && chmod +x /tmp/deploy-agent.sh && /tmp/deploy-agent.sh {python_version} {package_manager.value}"
+        f"wget {deploy_script_url} -O /tmp/deploy-agent.sh -q --no-cache && chmod +x /tmp/deploy-agent.sh && /tmp/deploy-agent.sh {python_version} {package_manager.value}"
     )
 
     output = stdout.read().decode("utf-8")
@@ -265,7 +270,7 @@ async def update(
             channel=config.ALEPH_CHANNEL,
         )
 
-    return UpdateAgentResponse(instance_hash=hostname)
+    return UpdateAgentResponse(instance_ip=hostname)
 
 
 @app.delete("/agent", description="Remove an agent on subscription end")

backend/src/utils/aleph.py

@@ -1,11 +1,7 @@
 import aiohttp
-from aleph.sdk import AlephHttpClient
-from aleph_message.models import InstanceMessage
 
-from src.config import config
 
-
-async def fetch_instance_ip(item_hash) -> str:
+async def fetch_instance_ip(item_hash: str) -> str:
     """
     Fetches IPv6 of an allocated instance given a message hash.
 
@@ -14,16 +10,18 @@ async def fetch_instance_ip(item_hash) -> str:
     Returns:
         IPv6 address
     """
-    async with AlephHttpClient(api_server=config.ALEPH_API_URL) as client:
-        message = await client.get_message(item_hash, InstanceMessage)
 
     async with aiohttp.ClientSession() as session:
         try:
             async with session.get(
-                f"https://scheduler.api.aleph.cloud/api/v0/allocation/{message.item_hash}"
+                f"https://scheduler.api.aleph.cloud/api/v0/allocation/{item_hash}"
             ) as resp:
                 resp.raise_for_status()
                 allocation = await resp.json()
                 return allocation["vm_ipv6"]
-        except (aiohttp.ClientResponseError, aiohttp.ClientConnectorError):
+        except (
+            aiohttp.ClientResponseError,
+            aiohttp.ClientConnectorError,
+            aiohttp.ConnectionTimeoutError,
+        ):
             raise ValueError()