From 05ade22d3b6312726b9b94a0804992b5602cb71c Mon Sep 17 00:00:00 2001 From: kerthcet Date: Tue, 6 Jan 2026 13:46:43 +0000 Subject: [PATCH 1/3] support services in sandbox Signed-off-by: kerthcet --- src/hive_cli/config.py | 34 ++++++++++++++++++++++++++++++---- src/hive_cli/platform/k8s.py | 34 ++++++++++++++++++++++++++++++++++ 2 files changed, 64 insertions(+), 4 deletions(-) diff --git a/src/hive_cli/config.py b/src/hive_cli/config.py index fd34fea..df72ee6 100644 --- a/src/hive_cli/config.py +++ b/src/hive_cli/config.py @@ -1,3 +1,4 @@ +from importlib import resources import os from enum import Enum from typing import Optional @@ -10,8 +11,7 @@ class PlatformType(str, Enum): K8S = "k8s" - ON_PREM = "on-prem" - + # ON_PREM = "on-prem" class ResourceConfig(BaseModel): cpu: str = Field( @@ -35,6 +35,26 @@ class EnvConfig(BaseModel): name: str value: str +class PortConfig(BaseModel): + port: int = Field( + description="The port number inside the container.", + ) + protocol: Optional[str] = Field( + default="TCP", + description="The protocol for the port. Default to 'TCP'.", + ) + +class ServiceConfig(BaseModel): + name: str + image: str + ports: Optional[list[PortConfig]] = None + envs: Optional[list[EnvConfig]] = None + command: Optional[list[str]] = None + args: Optional[list[str]] = None + resources: ResourceConfig = Field( + default_factory=ResourceConfig, + description="Resource configuration for the service.", + ) class SandboxConfig(BaseModel): image: Optional[str] = Field( @@ -71,7 +91,10 @@ class SandboxConfig(BaseModel): default=None, description="A pre-processor script to run before the experiment. Use the `/data/preprocessor` directory to load/store datasets.", ) - + services: Optional[list[ServiceConfig]] = Field( + default=None, + description="Additional services to run alongside the sandbox.", + ) class PromptConfig(BaseModel): context: Optional[str] = Field( @@ -171,7 +194,10 @@ class HiveConfig(BaseModel): description="The name of the coordinator config to use for the experiment. Default to 'default-coordinator-config'.", ) - platform: PlatformType = PlatformType.K8S + platform: PlatformType = Field( + default=PlatformType.K8S, + description="The platform type to use for the experiment. Default to 'k8s'.", + ) runtime: RuntimeConfig = Field( default_factory=RuntimeConfig, description="Runtime configuration for the experiment." diff --git a/src/hive_cli/platform/k8s.py b/src/hive_cli/platform/k8s.py index 9fd0ede..f6f8cfd 100644 --- a/src/hive_cli/platform/k8s.py +++ b/src/hive_cli/platform/k8s.py @@ -241,6 +241,39 @@ def construct_experiment(name: str, namespace: str, config: HiveConfig) -> dict: resources["requests"].update(config.sandbox.resources.extended_resources) resources["limits"].update(config.sandbox.resources.extended_resources) + services = [] + if config.sandbox.services is not None: + for svc in config.sandbox.services: + svc_dict = { + "name": svc.name, + "image": svc.image, + "envs": [env.model_dump() for env in svc.envs] if svc.envs else None, + "command": svc.command, + "args": svc.args, + "resources": { + "requests": { + "cpu": svc.resources.cpu, + "memory": svc.resources.memory, + }, + "limits": { + "cpu": svc.resources.cpu, + "memory": svc.resources.memory, + }, + "shmSize": svc.resources.shmsize, + "accelerators": svc.resources.accelerators, + } + } + for port in svc.ports: + if "ports" not in svc_dict: + svc_dict["ports"] = [] + svc_dict["ports"].append({"name": svc.name + "-" + str(port.port), "containerPort": port.port, "protocol": port.protocol}) + + if svc.resources.extended_resources is not None: + svc_dict["resources"]["requests"].update(svc.resources.extended_resources) + svc_dict["resources"]["limits"].update(svc.resources.extended_resources) + + services.append(svc_dict) + result = { "apiVersion": f"{GROUP}/{VERSION}", "kind": RESOURCE, @@ -257,6 +290,7 @@ def construct_experiment(name: str, namespace: str, config: HiveConfig) -> dict: "resources": resources, "envs": envs, "preprocessor": config.sandbox.pre_processor or config.sandbox.preprocessor, + "services": services, }, "repo": { "branch": config.repo.branch, From 00f1a1190cca23cc2eb02e1f6e2d750b9090ba3d Mon Sep 17 00:00:00 2001 From: kerthcet Date: Tue, 6 Jan 2026 13:47:06 +0000 Subject: [PATCH 2/3] format Signed-off-by: kerthcet --- src/hive_cli/config.py | 6 +++++- src/hive_cli/main.py | 24 ++++++------------------ src/hive_cli/platform/k8s.py | 10 ++++++++-- 3 files changed, 19 insertions(+), 21 deletions(-) diff --git a/src/hive_cli/config.py b/src/hive_cli/config.py index df72ee6..b3d7061 100644 --- a/src/hive_cli/config.py +++ b/src/hive_cli/config.py @@ -1,4 +1,3 @@ -from importlib import resources import os from enum import Enum from typing import Optional @@ -13,6 +12,7 @@ class PlatformType(str, Enum): K8S = "k8s" # ON_PREM = "on-prem" + class ResourceConfig(BaseModel): cpu: str = Field( default="1", @@ -35,6 +35,7 @@ class EnvConfig(BaseModel): name: str value: str + class PortConfig(BaseModel): port: int = Field( description="The port number inside the container.", @@ -44,6 +45,7 @@ class PortConfig(BaseModel): description="The protocol for the port. Default to 'TCP'.", ) + class ServiceConfig(BaseModel): name: str image: str @@ -56,6 +58,7 @@ class ServiceConfig(BaseModel): description="Resource configuration for the service.", ) + class SandboxConfig(BaseModel): image: Optional[str] = Field( default=None, @@ -96,6 +99,7 @@ class SandboxConfig(BaseModel): description="Additional services to run alongside the sandbox.", ) + class PromptConfig(BaseModel): context: Optional[str] = Field( default=None, diff --git a/src/hive_cli/main.py b/src/hive_cli/main.py index 374aa2f..0b81695 100644 --- a/src/hive_cli/main.py +++ b/src/hive_cli/main.py @@ -62,9 +62,7 @@ def update_experiment_cli(args): platform.update(args.name, config=config) console = Console() - msg = Text( - f"Experiment {args.name} updated successfully.", style="bold green" - ) + msg = Text(f"Experiment {args.name} updated successfully.", style="bold green") console.print(msg) @@ -129,9 +127,7 @@ def show_dashboard_cli(args): return pod_name = pods.items[0].metadata.name - with portforward.forward( - namespace, pod_name, args.port, remote_port, config.token_path - ): + with portforward.forward(namespace, pod_name, args.port, remote_port, config.token_path): event.wait_for_ctrl_c() console.print("\n[bold yellow]Port forwarding stopped.[/]") @@ -251,9 +247,7 @@ def main(): parser_edit_config.set_defaults(func=edit_cli) # dashboard command - parser_dashboard = subparsers.add_parser( - "dashboard", help="Open the Hive dashboard" - ) + parser_dashboard = subparsers.add_parser("dashboard", help="Open the Hive dashboard") parser_dashboard.add_argument( "--port", default=9090, @@ -269,18 +263,12 @@ def main(): parser_dashboard.set_defaults(func=show_dashboard_cli) # version command - parser_version = subparsers.add_parser( - "version", help="Show Hive CLI version" - ) - parser_version.set_defaults( - func=lambda args: print(f"Hive CLI version {__version__}") - ) + parser_version = subparsers.add_parser("version", help="Show Hive CLI version") + parser_version.set_defaults(func=lambda args: print(f"Hive CLI version {__version__}")) # log command parser_log = subparsers.add_parser("log", help="Show Sandbox logs") - parser_log.add_argument( - "sandbox", help="Name of the sandbox to fetch logs for" - ) + parser_log.add_argument("sandbox", help="Name of the sandbox to fetch logs for") parser_log.add_argument( "-f", "--config", diff --git a/src/hive_cli/platform/k8s.py b/src/hive_cli/platform/k8s.py index f6f8cfd..3be90f3 100644 --- a/src/hive_cli/platform/k8s.py +++ b/src/hive_cli/platform/k8s.py @@ -261,12 +261,18 @@ def construct_experiment(name: str, namespace: str, config: HiveConfig) -> dict: }, "shmSize": svc.resources.shmsize, "accelerators": svc.resources.accelerators, - } + }, } for port in svc.ports: if "ports" not in svc_dict: svc_dict["ports"] = [] - svc_dict["ports"].append({"name": svc.name + "-" + str(port.port), "containerPort": port.port, "protocol": port.protocol}) + svc_dict["ports"].append( + { + "name": svc.name + "-" + str(port.port), + "containerPort": port.port, + "protocol": port.protocol, + } + ) if svc.resources.extended_resources is not None: svc_dict["resources"]["requests"].update(svc.resources.extended_resources) From 0ba4f0699b6658fd3a817504957bee3258e5c38a Mon Sep 17 00:00:00 2001 From: kerthcet Date: Tue, 6 Jan 2026 14:17:34 +0000 Subject: [PATCH 3/3] fix nonetype error Signed-off-by: kerthcet --- src/hive_cli/platform/k8s.py | 21 +++++++++++---------- 1 file changed, 11 insertions(+), 10 deletions(-) diff --git a/src/hive_cli/platform/k8s.py b/src/hive_cli/platform/k8s.py index 3be90f3..bdf76df 100644 --- a/src/hive_cli/platform/k8s.py +++ b/src/hive_cli/platform/k8s.py @@ -263,16 +263,17 @@ def construct_experiment(name: str, namespace: str, config: HiveConfig) -> dict: "accelerators": svc.resources.accelerators, }, } - for port in svc.ports: - if "ports" not in svc_dict: - svc_dict["ports"] = [] - svc_dict["ports"].append( - { - "name": svc.name + "-" + str(port.port), - "containerPort": port.port, - "protocol": port.protocol, - } - ) + if svc.ports is not None: + for port in svc.ports: + if "ports" not in svc_dict: + svc_dict["ports"] = [] + svc_dict["ports"].append( + { + "name": svc.name + "-" + str(port.port), + "containerPort": port.port, + "protocol": port.protocol, + } + ) if svc.resources.extended_resources is not None: svc_dict["resources"]["requests"].update(svc.resources.extended_resources)