From c4c39f51c1785c6c498ecbcd0e710b7f61b36e86 Mon Sep 17 00:00:00 2001
From: Raz Rotenberg <raz.rotenberg@gmail.com>
Date: Tue, 7 May 2024 11:24:55 +0300
Subject: [PATCH] Bug fix in "genv remote llm attach"

* Supporting new environment name format

- Showing port in "ps"
---
 genv/cli/remote.py | 24 ++++++++++++++++++++----
 1 file changed, 20 insertions(+), 4 deletions(-)

diff --git a/genv/cli/remote.py b/genv/cli/remote.py
index b323107..61c0d1c 100755
--- a/genv/cli/remote.py
+++ b/genv/cli/remote.py
@@ -3,6 +3,7 @@
 import dataclasses
 import itertools
 import os
+import re
 import shutil
 import sys
 from typing import Iterable, NoReturn, Optional, Tuple
@@ -241,6 +242,14 @@ async def do_envs(
         )
 
 
+def _find_llm_port(env: genv.Env) -> Optional[int]:
+    """Finds any port an LLM server environment listens on."""
+
+    match = re.match(r"^llm/[^/]+/(\d+)$", env.config.name)
+    if match:
+        return int(match.group(1))
+
+
 async def do_llm(config: genv.remote.Config, args: argparse.Namespace):
     """Runs the "genv llm" logic."""
 
@@ -265,9 +274,16 @@ async def do_llm_attach(config: genv.remote.Config, model: str) -> NoReturn:
     hosts, snapshots = await genv.remote.core.envs.snapshot(config)
 
     for host, snapshot in zip(hosts, snapshots):
-        replicas = len(snapshot.filter(name=f"llm/{model}"))
+        for env in snapshot:
+            if not env.config.name:
+                continue
+
+            if not (
+                env.config.name.startswith(f"llm/{model}/")
+                or env.config.name == f"llm/{model}"  # before 1.4.1
+            ):
+                continue
 
-        if replicas > 0:
             # TODO(raz): we currently attach to the first node; we should pick a node in a smarter way.
             _exec_ssh(host, f"genv llm attach {model}")
 
@@ -309,8 +325,8 @@ async def do_llm_ps(
 
             total += 1
 
-            model = env.config.name.split("llm/")[1]
-            port = "N/A"
+            model = env.config.name.split("/")[1]
+            port = _find_llm_port(env) or "N/A"
             created = env.creation if timestamp else env.time_since
             eid = env.eid
             user = env.username or ""