Feat: DAH-1614 - gpu splitting

lium/cli/init/auth.py

@@ -2,6 +2,9 @@
 from typing import Optional
 
 from lium.cli import ui
+from lium.sdk.config import Config
+
+config = Config.load()
 
 class quiet_fds:
     """Redirect stdout/stderr to /dev/null (silences child processes)."""
@@ -18,7 +21,7 @@ def __exit__(self, *_):
         self._null.close()
 
 def init_auth():
-    url = "https://lium.io/api/cli-auth/init"
+    url = f"{config.base_url}/cli-auth/init"
     resp = requests.post(url,
         json={"callback_url": "http://localhost:8080/auth/callback"},
         headers={"Content-Type": "application/json"},
@@ -28,7 +31,7 @@ def init_auth():
     return resp.json()["browser_url"], resp.json()["session_id"]
 
 def poll_auth(session_id, max_attempts=6, interval=5) -> Optional[str]:  # 30 seconds timeout (6 * 5)
-    url = f"https://lium.io/api/cli-auth/poll/{session_id}"
+    url = f"{config.base_url}/cli-auth/poll/{session_id}"
     for _ in range(max_attempts):
         try:
             resp = requests.get(url, timeout=5)

lium/cli/ls/display.py

@@ -22,7 +22,7 @@ def _mid_ellipsize(s: str, width: int = 28) -> str:
 
 def _cfg(exe: ExecutorInfo) -> str:
     """Format GPU configuration string."""
-    return f"{exe.gpu_count}×{exe.gpu_type}"
+    return f"{exe.available_gpu_count}×{exe.gpu_type}"
 
 
 def _country_name(loc: Optional[Dict]) -> str:
@@ -117,8 +117,9 @@ def _sort_key_factory(name: str) -> Callable[[ExecutorInfo], Any]:
 def _add_table_columns(t: Table) -> None:
     """Add columns to the table with fixed widths."""
     t.add_column("", justify="right", width=3, no_wrap=True, style="dim")
-    t.add_column("Id", justify="left", ratio=8, min_width=24, overflow="fold")
+    t.add_column("Id", justify="left", ratio=8, min_width=18, overflow="fold")
     t.add_column("Config", justify="left", width=12, no_wrap=True)
+    t.add_column("GPU Splitting", justify="left", min_width=12, no_wrap=True)
     t.add_column("$/GPU·h", justify="right", width=8, no_wrap=True)
     t.add_column("Location", justify="left", ratio=4, min_width=10, overflow="fold")
     t.add_column("VRAM (Gb)", justify="right", width=11, no_wrap=True)
@@ -201,11 +202,13 @@ def build_executors_table(
         huid = _mid_ellipsize(exe.huid)
         huid += " (DinD)" if exe.docker_in_docker else ""
         huid_display = f"{console.get_styled('★', 'success')} {console.get_styled(huid, 'id')}" if is_pareto else f"  {console.get_styled(huid, 'id')}"
+        gpu_splitting_display = f"Min GPUs: {exe.min_gpu_count_for_rental}" if exe.min_gpu_count_for_rental else "—"
 
         table.add_row(
             str(idx),
             huid_display,
             _cfg(exe),
+            gpu_splitting_display,
             console.get_styled(_money(exe.price_per_gpu_hour), 'success'),
             _country_name(exe.location),
             s["VRAM"],

lium/cli/plugin_example.py

@@ -82,11 +82,11 @@ def compose_up(ctx, file: str, detach: bool):
 
         # Find suitable executor
         gpu_type = model_config.get('gpu_type')
-        gpu_count = model_config.get('gpu_count', 1)
+        gpu_count = model_config.get('gpu_count', None)
         template_id = model_config.get('template_id')
 
         # Get available executors
-        executors = lium.ls(gpu_type=gpu_type)
+        executors = lium.ls(gpu_type=gpu_type, gpu_count=gpu_count)
         if not executors:
             click.echo(f"No executors available for {gpu_type}", err=True)
             continue
@@ -96,6 +96,7 @@ def compose_up(ctx, file: str, detach: bool):
         pod = lium.up(
             executor=executor.id,
             name=model_name,
+            gpu_count=gpu_count,
             template=template_id,
         )
 

lium/cli/ps/display.py

@@ -112,13 +112,11 @@ def build_pods_table(pods: List[PodInfo], short: bool = False) -> tuple[Table |
     for pod in pods:
         executor = pod.executor
         if executor:
-            config = f"{executor.gpu_count}×{executor.gpu_type}" if executor.gpu_count > 1 else executor.gpu_type
-            price_str = f"${executor.price_per_hour:.2f}"
-            price_per_hour = executor.price_per_hour
+            config = f"{pod.gpu_count}×{executor.gpu_type}" if pod.gpu_count > 1 else executor.gpu_type
+            price_str = f"${pod.price:.2f}"
         else:
             config = "—"
             price_str = "—"
-            price_per_hour = None
 
         status_color = console.pod_status_color(pod.status)
         status_text = f"[{status_color}]{pod.status.upper()}[/]"
@@ -132,7 +130,7 @@ def build_pods_table(pods: List[PodInfo], short: bool = False) -> tuple[Table |
             config,
             console.get_styled(template_name, 'info'),
             price_str,
-            _format_cost(pod.created_at, price_per_hour),
+            _format_cost(pod.created_at, pod.price),
             _format_uptime(pod.created_at),
         ]
 

lium/cli/rm/display.py

@@ -15,7 +15,7 @@ def calculate_pod_cost(pod: PodInfo) -> float:
     Returns:
         Total cost in dollars
     """
-    if not pod.executor or not pod.executor.price_per_hour or not pod.created_at:
+    if not pod.price or not pod.created_at:
         return 0.0
 
     try:
@@ -28,7 +28,7 @@ def calculate_pod_cost(pod: PodInfo) -> float:
 
         now_utc = datetime.now(timezone.utc)
         hours = (now_utc - dt_created).total_seconds() / 3600
-        return hours * pod.executor.price_per_hour
+        return hours * pod.price
     except Exception:
         return 0.0
 
@@ -48,8 +48,8 @@ def format_pods_for_removal(pods: List[PodInfo], show_cost: bool = True) -> str:
 
     for pod in pods:
         price_info = ""
-        if pod.executor and pod.executor.price_per_hour:
-            price_info = f" (${pod.executor.price_per_hour:.2f}/h)"
+        if pod.price:
+            price_info = f" (${pod.price:.2f}/h)"
             if show_cost:
                 total_cost += calculate_pod_cost(pod)
 
@@ -75,8 +75,8 @@ def format_pods_for_scheduled_removal(pods: List[PodInfo], termination_time: dat
 
     for pod in pods:
         price_info = ""
-        if pod.executor and pod.executor.price_per_hour:
-            price_info = f" (${pod.executor.price_per_hour:.2f}/h)"
+        if pod.price:
+            price_info = f" (${pod.price:.2f}/h)"
         lines.append(f"  {pod.huid} - {pod.status}{price_info}")
 
     # Add scheduled time info

lium/cli/up/actions.py

@@ -32,6 +32,16 @@ def execute(self, ctx: dict) -> ActionResult:
                 executor = lium.get_executor(executor_id)
                 if not executor:
                     return ActionResult(ok=False, data={}, error=f"Executor '{executor_id}' not found")
+
+                if count:
+                    if count > executor.available_gpu_count:
+                        return ActionResult(ok=False, data={}, error=f"Executor {executor.huid} has insufficient GPUs (available: {executor.available_gpu_count}, required: {count})")
+                    if executor.min_gpu_count_for_rental:
+                        if count < executor.min_gpu_count_for_rental:
+                            return ActionResult(ok=False, data={}, error=f"Executor {executor.huid} requires at least {executor.min_gpu_count_for_rental} GPUs for rental.")
+                    else:
+                        if count < executor.available_gpu_count:
+                            return ActionResult(ok=False, data={}, error=f"Executor {executor.huid} doesn't support gpu splitting.")
 
                 if ports and (not executor.available_port_count or executor.available_port_count < ports):
                     available = executor.available_port_count or 0
@@ -41,10 +51,13 @@ def execute(self, ctx: dict) -> ActionResult:
                         error=f"Executor {executor.huid} has insufficient ports (available: {available}, required: {ports})"
                     )
             else:
-                executors = lium.ls(gpu_type=gpu)
+                executors = lium.ls(gpu_type=gpu, gpu_count=count)
 
                 if count:
-                    executors = [e for e in executors if e.gpu_count == count]
+                    executors = [
+                        e for e in executors
+                        if (not e.min_gpu_count_for_rental and e.available_gpu_count == count) or (e.min_gpu_count_for_rental and e.min_gpu_count_for_rental <= count and e.available_gpu_count >= count)
+                    ]
                 if country:
                     executors = [
                         e for e in executors
@@ -170,6 +183,7 @@ def execute(self, ctx: dict) -> ActionResult:
         lium: Lium = ctx["lium"]
         executor: ExecutorInfo = ctx["executor"]
         template: Template = ctx["template"]
+        gpu_count: Optional[int] = ctx.get("gpu_count")
         name: Optional[str] = ctx.get("name")
         volume_id: Optional[str] = ctx.get("volume_id")
         ports: Optional[int] = ctx.get("ports")
@@ -181,6 +195,7 @@ def execute(self, ctx: dict) -> ActionResult:
             pod_info = lium.up(
                 executor_id=executor.id,
                 name=name,
+                gpu_count=gpu_count,
                 template_id=template.id if template else None,
                 volume_id=volume_id,
                 ports=ports,

lium/cli/up/command.py

@@ -1,7 +1,7 @@
 from typing import Optional, Tuple
 import click
 
-from lium.sdk import Lium
+from lium.sdk import Lium, ExecutorInfo
 from lium.cli import ui
 from lium.cli.utils import handle_errors, ensure_config
 from lium.cli.completion import get_gpu_completions
@@ -134,13 +134,13 @@ def up_command(
         ui.error(result.error)
         return
 
-    executor = result.data["executor"]
+    executor: ExecutorInfo = result.data["executor"]
 
     if not yes:
         confirm_msg = (
             f"Acquire pod on {executor.huid} "
-            f"({executor.gpu_count}×{executor.gpu_type}) "
-            f"at ${executor.price_per_hour:.2f}/h?"
+            f"({count or executor.available_gpu_count}×{executor.gpu_type}) "
+            f"at ${(executor.price_per_gpu_hour * (count or executor.available_gpu_count)):.2f}/h?"
         )
         if not ui.confirm(confirm_msg):
             return
@@ -208,6 +208,7 @@ def up_command(
             "lium": lium,
             "executor": executor,
             "template": template,
+            "gpu_count": count or executor.available_gpu_count,
             "name": name,
             "volume_id": volume_id,
             "ports": ports

lium/cli/up/validation.py

@@ -14,8 +14,8 @@ def validate(
     template_id: str | None = None,
 ) -> tuple[bool, str]:
     """Validate up command inputs."""
-    if executor_id and (gpu or count or country):
-        return False, "Cannot use filters (--gpu, --count, --country) when specifying an executor ID"
+    if executor_id and (gpu or country):
+        return False, "Cannot use filters (--gpu, --country) when specifying an executor ID"
 
     if not executor_id and not (gpu or count or country):
         return False, "Must provide either EXECUTOR_ID or filters (--gpu, --count, --country)"

lium/cli/utils.py

@@ -400,8 +400,14 @@ def store_executor_selection(executors: List[ExecutorInfo]) -> None:
             'huid': executor.huid,
             'gpu_type': executor.gpu_type,
             'gpu_count': executor.gpu_count,
-            'price_per_hour': executor.price_per_hour,
-            'location': executor.location.get('country', 'Unknown') if executor.location else 'Unknown'
+            'available_gpu_count': executor.available_gpu_count,
+            'price_per_gpu_hour': executor.price_per_gpu_hour,
+            'min_gpu_count_for_rental': executor.min_gpu_count_for_rental,
+            'location': executor.location.get('country', 'Unknown') if executor.location else 'Unknown',
+            'status': executor.status,
+            'docker_in_docker': executor.docker_in_docker,
+            'ip': executor.ip,
+            'available_port_count': executor.available_port_count,
         })
 
     # Store in config directory

lium/sdk/client.py

@@ -122,6 +122,7 @@ def _dict_to_executor_info(self, executor_dict: Dict) -> Optional[ExecutorInfo]:
         specs = executor_dict.get("specs", {})
         gpu_info = specs.get("gpu", {})
         gpu_count = gpu_info.get("count", 1)
+        available_gpu_count = executor_dict.get("available_gpu_count", 1)
 
         # Extract GPU type from machine_name or specs
         machine_name = executor_dict.get("machine_name", "")
@@ -135,29 +136,29 @@ def _dict_to_executor_info(self, executor_dict: Dict) -> Optional[ExecutorInfo]:
                 if gpu_name:
                     gpu_type = extract_gpu_type(gpu_name)
 
-        price_per_hour = executor_dict.get("price_per_hour", 0)
-
         return ExecutorInfo(
             id=executor_dict.get("id", ""),
             ip=executor_dict.get("executor_ip_address", ""),
             huid=generate_huid(executor_dict.get("id", "")),
             machine_name=machine_name,
             gpu_type=gpu_type,
             gpu_count=gpu_count,
-            price_per_hour=price_per_hour,
-            price_per_gpu_hour=price_per_hour / max(1, gpu_count),
+            available_gpu_count=available_gpu_count,
+            price_per_gpu_hour=executor_dict.get("price_per_gpu", 0),
             location=executor_dict.get("location", {}),
             specs=specs,
             status=executor_dict.get("status", "unknown"),
             docker_in_docker=specs.get("sysbox_runtime", False),
             available_port_count=specs.get("available_port_count"),
+            min_gpu_count_for_rental=executor_dict.get("min_gpu_count_for_rental", None),
         )
 
     def up(
         self,
         *,
         executor_id: str,
-        name: str = "Your Pod",
+        name: Optional[str] = None,
+        gpu_count: Optional[int] = None,
         template_id: Optional[str] = None,
         volume_id: Optional[str] = None,
         ports: Optional[int] = None,
@@ -190,6 +191,7 @@ def up(
 
         payload = {
             "pod_name": name,
+            "gpu_count": gpu_count,
             "template_id": template_id,
             "volume_id": volume_id,
             "user_public_key": ssh_material,
@@ -336,7 +338,7 @@ def ls(
                 params["machine_names"] = gpu_type
         if gpu_count:
             params["gpu_count_gte"] = gpu_count
-            params["gpu_count_lte"] = gpu_count
+            # params["gpu_count_lte"] = gpu_count
         if lat is not None and lon is not None:
             params["lat"] = lat
             params["lon"] = lon
@@ -365,6 +367,8 @@ def ps(self) -> List[PodInfo]:
                 name=d.get("pod_name", ""),
                 status=d.get("status", "unknown"),
                 huid=generate_huid(d.get("id", "")),
+                gpu_count=int(d.get("gpu_count", 0)),
+                price=d.get("price", 0.0),
                 ssh_cmd=d.get("ssh_connect_cmd"),
                 ports=d.get("ports_mapping", {}),
                 created_at=d.get("created_at", ""),
@@ -900,6 +904,8 @@ def switch_template(self, pod: PodInfo, *, template_id: str) -> PodInfo:
             name=response.get("pod_name", pod.name),
             status=response.get("status", "PENDING"),
             huid=pod.huid,  # Keep the original HUID
+            gpu_count=int(response.get("gpu_count", 0)),
+            price=response.get("price", 0.0),
             ssh_cmd=response.get("ssh_connect_cmd"),
             ports=response.get("ports_mapping", {}),
             created_at=response.get("created_at", ""),
@@ -910,7 +916,7 @@ def switch_template(self, pod: PodInfo, *, template_id: str) -> PodInfo:
                 machine_name="",
                 gpu_type=response.get("gpu_name", ""),
                 gpu_count=int(response.get("gpu_count", 0) or 0),
-                price_per_hour=0.0,
+                available_gpu_count=0,
                 price_per_gpu_hour=0.0,
                 location={},
                 specs={},

lium/sdk/models.py

@@ -12,14 +12,15 @@ class ExecutorInfo:
     machine_name: str
     gpu_type: str
     gpu_count: int
-    price_per_hour: float
+    available_gpu_count: int
     price_per_gpu_hour: float
     location: Dict
     specs: Dict
     status: str
     docker_in_docker: bool
     ip: str
     available_port_count: Optional[int] = None
+    min_gpu_count_for_rental: int | None = None
 
     @property
     def driver_version(self) -> str:
@@ -43,6 +44,8 @@ class PodInfo:
     ports: Dict
     created_at: str
     updated_at: str
+    gpu_count: int
+    price: float
     executor: Optional[ExecutorInfo]
     template: Dict
     removal_scheduled_at: Optional[str]