diff --git a/src/hive_cli/config.py b/src/hive_cli/config.py index fd34fea..b3d7061 100644 --- a/src/hive_cli/config.py +++ b/src/hive_cli/config.py @@ -10,7 +10,7 @@ class PlatformType(str, Enum): K8S = "k8s" - ON_PREM = "on-prem" + # ON_PREM = "on-prem" class ResourceConfig(BaseModel): @@ -36,6 +36,29 @@ class EnvConfig(BaseModel): 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( default=None, @@ -71,6 +94,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): @@ -171,7 +198,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/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 9fd0ede..bdf76df 100644 --- a/src/hive_cli/platform/k8s.py +++ b/src/hive_cli/platform/k8s.py @@ -241,6 +241,46 @@ 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, + }, + } + 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) + svc_dict["resources"]["limits"].update(svc.resources.extended_resources) + + services.append(svc_dict) + result = { "apiVersion": f"{GROUP}/{VERSION}", "kind": RESOURCE, @@ -257,6 +297,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,