support services in sandbox

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."

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",

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,