feat: expose GPU cluster add-on flags for headlamp and slurm web (#461)

* feat(cli): expose cluster add-on flags

Co-authored-by: Blaine Kasten <blainekasten@gmail.com>

* fix(cli): suppress negative add-on flag variants

Co-authored-by: Blaine Kasten <blainekasten@gmail.com>

* fix(cli): rename cluster add-on flags

Co-authored-by: Blaine Kasten <blainekasten@gmail.com>

---------

Co-authored-by: Cursor Agent <cursoragent@cursor.com>
Co-authored-by: Blaine Kasten <blainekasten@gmail.com>

Author: 3 people

src/together/lib/cli/api/beta/clusters/create.py

@@ -8,7 +8,7 @@
 from together._utils._json import openapi_dumps
 from together.lib.cli.utils.config import CLIConfigParameter
 from together.lib.cli.utils._console import console
-from together.types.beta.cluster_create_params import SharedVolume, ClusterCreateParams
+from together.types.beta.cluster_create_params import AddOn, SharedVolume, ClusterCreateParams
 
 NameParameter = Annotated[Optional[str], Parameter(help="Name of the cluster")]
 NumGpusParameter = Annotated[Optional[int], Parameter(help="Number of GPUs to allocate in the cluster")]
@@ -32,6 +32,10 @@
 AutoScaleMaxGpusParameter = Annotated[Optional[int], Parameter(help="Maximum GPUs for auto-scaling")]
 CapacityPoolIDParameter = Annotated[Optional[str], Parameter(help="Capacity pool ID to use for the cluster")]
 InstallTraefikParameter = Annotated[Optional[bool], Parameter(help="Install Traefik ingress controller")]
+HeadlampParameter = Annotated[
+    Optional[bool], Parameter(help="Enable the Headlamp Kubernetes dashboard add-on", negative=())
+]
+SlurmWebParameter = Annotated[Optional[bool], Parameter(help="Enable the Slurm Web add-on", negative=())]
 NumCapacityPoolGpusParameter = Annotated[
     Optional[int], Parameter(help="Number of GPUs to allocate from a capacity pool")
 ]
@@ -60,6 +64,8 @@ async def create(
     auto_scale_max_gpus: AutoScaleMaxGpusParameter = None,
     capacity_pool_id: CapacityPoolIDParameter = None,
     install_traefik: InstallTraefikParameter = None,
+    headlamp_addon: HeadlampParameter = None,
+    slurm_web_addon: SlurmWebParameter = None,
     num_capacity_pool_gpus: NumCapacityPoolGpusParameter = None,
     num_preemptible_gpus: NumPreemptibleGpusParameter = None,
     num_reserved_gpus: NumReservedGpusParameter = None,
@@ -92,6 +98,25 @@ async def create(
         params["capacity_pool_id"] = capacity_pool_id
     if install_traefik is not None:
         params["install_traefik"] = install_traefik
+    add_ons: list[AddOn] = []
+    if headlamp_addon is not None:
+        add_ons.append(
+            {
+                "add_on_type": "headlamp",
+                "name": "headlamp",
+                "config": {"headlamp": {"enabled": headlamp_addon}},
+            }
+        )
+    if slurm_web_addon is not None:
+        add_ons.append(
+            {
+                "add_on_type": "slurm_web",
+                "name": "slurm_web",
+                "config": {"slurm_web": {"enabled": slurm_web_addon}},
+            }
+        )
+    if add_ons:
+        params["add_ons"] = add_ons
     if num_capacity_pool_gpus is not None:
         params["num_capacity_pool_gpus"] = num_capacity_pool_gpus
     if num_preemptible_gpus is not None:

src/together/lib/cli/api/beta/clusters/update.py

@@ -10,6 +10,7 @@
 from together.lib.cli.utils.config import CLIConfigParameter
 from together.lib.cli.utils._console import console
 from together.lib.cli.components.loader import show_loading_status
+from together.types.beta.cluster_update_params import AddOn
 
 
 async def update(
@@ -34,15 +35,30 @@ async def update(
         Optional[str],
         Parameter(help="Timestamp at which the cluster should be decommissioned"),
     ] = None,
+    headlamp: Annotated[
+        Optional[bool],
+        Parameter(help="Enable or disable the Headlamp Kubernetes dashboard add-on"),
+    ] = None,
+    slurm_web: Annotated[
+        Optional[bool],
+        Parameter(help="Enable or disable the Slurm Web add-on"),
+    ] = None,
     *,
     config: CLIConfigParameter,
 ) -> None:
     """Update a cluster."""
 
+    add_ons: list[AddOn] = []
+    if headlamp is not None:
+        add_ons.append({"name": "headlamp", "config": {"headlamp": {"enabled": headlamp}}})
+    if slurm_web is not None:
+        add_ons.append({"name": "slurm_web", "config": {"slurm_web": {"enabled": slurm_web}}})
+
     await show_loading_status(
         "Updating cluster...",
         config.client.beta.clusters.update(
             cluster_id,
+            add_ons=add_ons if add_ons else omit,
             num_gpus=num_gpus if num_gpus is not None else omit,
             cluster_type=cluster_type if cluster_type is not None else omit,
             num_capacity_pool_gpus=num_capacity_pool_gpus if num_capacity_pool_gpus is not None else omit,

tests/cli/test_beta_clusters.py

@@ -695,6 +695,8 @@ def test_create_accepts_new_cluster_params(self, respx_mock: MockRouter, cli_run
                 "--capacity-pool-id",
                 "pool-1",
                 "--install-traefik",
+                "--headlamp-addon",
+                "--slurm-web-addon",
                 "--num-capacity-pool-gpus",
                 "8",
                 "--num-preemptible-gpus",
@@ -722,6 +724,18 @@ def test_create_accepts_new_cluster_params(self, respx_mock: MockRouter, cli_run
         assert body["capacity_pool_id"] == "pool-1"
         assert "gpu_node_failover_enabled" not in body
         assert body["install_traefik"] is True
+        assert body["add_ons"] == [
+            {
+                "add_on_type": "headlamp",
+                "name": "headlamp",
+                "config": {"headlamp": {"enabled": True}},
+            },
+            {
+                "add_on_type": "slurm_web",
+                "name": "slurm_web",
+                "config": {"slurm_web": {"enabled": True}},
+            },
+        ]
         assert body["num_capacity_pool_gpus"] == 8
         assert body["num_preemptible_gpus"] == 8
         assert body["num_reserved_gpus"] == 8
@@ -768,10 +782,16 @@ def test_update_accepts_new_cluster_params(self, respx_mock: MockRouter, cli_run
                 "16",
                 "--reservation-end-time",
                 "2026-06-02T00:00:00Z",
+                "--no-headlamp",
+                "--slurm-web",
             ],
         )
 
         put_body = json.loads(cast(Call, put.calls[0]).request.content.decode())
+        assert put_body["add_ons"] == [
+            {"name": "headlamp", "config": {"headlamp": {"enabled": False}}},
+            {"name": "slurm_web", "config": {"slurm_web": {"enabled": True}}},
+        ]
         assert put_body["num_preemptible_gpus"] == 8
         assert put_body["num_capacity_pool_gpus"] == 8
         assert put_body["num_reserved_gpus"] == 16