fix(sdk): cloud=lightning preferring sched-test clusters over baremetal

python/lightning_sdk/api/cloud_account_api.py

@@ -15,6 +15,19 @@
     from lightning_sdk.machine import CloudProvider
     from lightning_sdk.teamspace import ConnectionType
 
+# Preferred cloud-account IDs when multiple accounts share a provider alias.
+# Without this, dict insertion order can make ``--cloud=<provider>`` resolve to an
+# arbitrary account of that provider (e.g. a BYOC/test cluster) instead of the
+# public Lightning Cloud account.
+_CANONICAL_PROVIDER_CLOUD_ACCOUNT_IDS = {
+    "AWS": "lightning-public-prod",
+    "GCP": "gcp-lightning-public-prod",
+    "LAMBDA_LABS": "lightning-lambda-prod",
+    "LIGHTNING": "lightning-baremetal",
+    "NEBIUS": "lightning-nebius-prod",
+    "VOLTAGE_PARK": "lightning-voltagepark-prod",
+}
+
 
 class CloudAccountApi:
     """Internal API client for API requests to cluster endpoints."""
@@ -195,15 +208,54 @@ def get_cloud_account_provider_mapping(
         else:
             res = self.list_cloud_accounts(teamspace_id=teamspace_id)
 
-        cloud_accounts = {cloud_account.id: cloud_account for cloud_account in res}
-        providers = {cloud_account.id: self._get_cloud_account_provider(cloud_account) for cloud_account in res}
-
-        mapping = {}
-        for cloud_account_id, provider in providers.items():
-            if provider is not None:
-                mapping[provider] = cloud_accounts[cloud_account_id]
+        mapping: Dict["CloudProvider", V1ExternalCluster] = {}
+        for cloud_account in res:
+            provider = self._get_cloud_account_provider(cloud_account)
+            if provider is None:
+                continue
+            existing = mapping.get(provider)
+            preferred = (
+                cloud_account
+                if existing is None
+                else self._prefer_cloud_account_for_provider(
+                    provider, existing, cloud_account, prefer_canonical=global_only
+                )
+            )
+            mapping[provider] = preferred
         return mapping
 
+    @staticmethod
+    def _prefer_cloud_account_for_provider(
+        provider: "CloudProvider",
+        current: V1ExternalCluster,
+        candidate: V1ExternalCluster,
+        prefer_canonical: bool = True,
+    ) -> V1ExternalCluster:
+        """Choose which cloud account should represent a provider alias.
+
+        Preference order:
+        1. When ``prefer_canonical`` (``--cloud=<provider>`` / global-only mapping): the
+           canonical public account ID for the provider
+        2. When not preferring canonical (e.g. data connections): BYOC/private over GLOBAL
+        3. Keep the first-seen account (stable; avoids last-writer wins)
+        """
+        if prefer_canonical:
+            canonical_id = _CANONICAL_PROVIDER_CLOUD_ACCOUNT_IDS.get(str(provider))
+            if canonical_id is not None:
+                if candidate.id == canonical_id:
+                    return candidate
+                if current.id == canonical_id:
+                    return current
+            return current
+
+        current_is_byoc = bool(current.spec and current.spec.cluster_type == V1ClusterType.BYOC)
+        candidate_is_byoc = bool(candidate.spec and candidate.spec.cluster_type == V1ClusterType.BYOC)
+        if candidate_is_byoc and not current_is_byoc:
+            return candidate
+        if current_is_byoc and not candidate_is_byoc:
+            return current
+        return current
+
     @staticmethod
     def _get_cloud_account_provider(cloud_account: Optional[V1ExternalCluster]) -> Optional["CloudProvider"]:
         """Determines the cloud provider based on the cloud_account configuration.

python/tests/api/test_cluster_api.py

@@ -1,3 +1,4 @@
+from typing import Optional
 from unittest.mock import MagicMock, patch
 
 import pytest
@@ -9,7 +10,11 @@
     V1ClusterType,
     V1ExternalCluster,
     V1ExternalClusterSpec,
+    V1GoogleCloudDirectV1,
+    V1LambdaLabsDirectV1,
     V1MachineDirectV1,
+    V1NebiusDirectV1,
+    V1VoltageParkDirectV1,
 )
 from lightning_sdk.lightning_cloud.openapi.models.v1_list_cluster_accelerators_response import (
     V1ListClusterAcceleratorsResponse,
@@ -265,3 +270,150 @@ def test_get_cloud_account_provider_distinguishes_lightning_providers(api, clust
     cluster = V1ExternalCluster(id="acc-lightning", spec=cluster_spec)
 
     assert api._get_cloud_account_provider(cluster) == expected_provider
+
+
+def _machine_cluster(cluster_id: str, name: Optional[str] = None) -> V1ExternalCluster:
+    return V1ExternalCluster(
+        id=cluster_id,
+        name=name or cluster_id,
+        spec=V1ExternalClusterSpec(cluster_type=V1ClusterType.GLOBAL, machine_v1=V1MachineDirectV1()),
+    )
+
+
+def test_lightning_provider_prefers_baremetal_over_sched_test_clusters(api):
+    """``--cloud=lightning`` must resolve to Lightning Cloud, not an internal sched-test MACHINE cluster."""
+    sched_test = _machine_cluster(
+        "01ky8n9bdq7jaapk4xx8ww8w51",
+        name="sched-test-cluster-01ky8n9bdq7jaapk4xx8ww8w51",
+    )
+    baremetal = _machine_cluster("lightning-baremetal", name="Lightning Cloud")
+    earlier_sched_test = _machine_cluster(
+        "01ky8n3ttpr9mv9gk2wg6acmx6",
+        name="sched-test-cluster-01ky8n3ttpr9mv9gk2wg6acmx6",
+    )
+
+    # Baremetal in the middle: last-writer-wins would incorrectly pick the final sched-test cluster.
+    api._client.cluster_service_list_project_clusters.return_value = V1ListProjectClustersResponse(clusters=[])
+    api._client.cluster_service_list_clusters.return_value = V1ListClustersResponse(
+        clusters=[earlier_sched_test, baremetal, sched_test]
+    )
+
+    mapping = api.get_cloud_account_provider_mapping(teamspace_id="ts", global_only=True)
+    assert mapping[CloudProvider.LIGHTNING].id == "lightning-baremetal"
+
+    result = api.resolve_cloud_account(teamspace_id="ts", cloud="lightning", default_cloud_account=None)
+    assert result == "lightning-baremetal"
+
+
+@pytest.mark.parametrize(
+    ("provider", "canonical_id", "other_id", "spec_kwargs"),
+    [
+        (CloudProvider.AWS, "lightning-public-prod", "other-aws", {"aws_v1": V1AWSDirectV1()}),
+        (CloudProvider.GCP, "gcp-lightning-public-prod", "other-gcp", {"google_cloud_v1": V1GoogleCloudDirectV1()}),
+        (
+            CloudProvider.LAMBDA_LABS,
+            "lightning-lambda-prod",
+            "other-lambda",
+            {"lambda_labs_v1": V1LambdaLabsDirectV1()},
+        ),
+        (CloudProvider.LIGHTNING, "lightning-baremetal", "other-machine", {"machine_v1": V1MachineDirectV1()}),
+        (CloudProvider.NEBIUS, "lightning-nebius-prod", "other-nebius", {"nebius_v1": V1NebiusDirectV1()}),
+        (
+            CloudProvider.VOLTAGE_PARK,
+            "lightning-voltagepark-prod",
+            "other-voltagepark",
+            {"voltage_park_v1": V1VoltageParkDirectV1()},
+        ),
+    ],
+)
+def test_global_provider_mapping_prefers_canonical_accounts(api, provider, canonical_id, other_id, spec_kwargs):
+    other = V1ExternalCluster(
+        id=other_id,
+        spec=V1ExternalClusterSpec(cluster_type=V1ClusterType.GLOBAL, **spec_kwargs),
+    )
+    canonical = V1ExternalCluster(
+        id=canonical_id,
+        spec=V1ExternalClusterSpec(cluster_type=V1ClusterType.GLOBAL, **spec_kwargs),
+    )
+
+    api._client.cluster_service_list_project_clusters.return_value = V1ListProjectClustersResponse(clusters=[])
+
+    # Non-canonical first: last-writer-wins would keep canonical, first-wins would keep other.
+    api._client.cluster_service_list_clusters.return_value = V1ListClustersResponse(clusters=[other, canonical])
+    mapping = api.get_cloud_account_provider_mapping(teamspace_id="ts", global_only=True)
+    assert mapping[provider].id == canonical_id
+
+    # Reverse order: still canonical.
+    api.list_cloud_accounts.cache_clear()
+    api.list_global_cloud_accounts.cache_clear()
+    api._client.cluster_service_list_clusters.return_value = V1ListClustersResponse(clusters=[canonical, other])
+    mapping = api.get_cloud_account_provider_mapping(teamspace_id="ts", global_only=True)
+    assert mapping[provider].id == canonical_id
+
+
+@pytest.mark.parametrize(
+    ("provider", "cloud", "first_id", "second_id", "spec_kwargs"),
+    [
+        (CloudProvider.AWS, "aws", "staging-aws", "dev-aws", {"aws_v1": V1AWSDirectV1()}),
+        (CloudProvider.GCP, "gcp", "staging-gcp", "dev-gcp", {"google_cloud_v1": V1GoogleCloudDirectV1()}),
+        (
+            CloudProvider.LAMBDA_LABS,
+            "lambda_labs",
+            "staging-lambda",
+            "dev-lambda",
+            {"lambda_labs_v1": V1LambdaLabsDirectV1()},
+        ),
+        (CloudProvider.LIGHTNING, "lightning", "staging-machine", "dev-machine", {"machine_v1": V1MachineDirectV1()}),
+        (CloudProvider.NEBIUS, "nebius", "staging-nebius", "dev-nebius", {"nebius_v1": V1NebiusDirectV1()}),
+        (
+            CloudProvider.VOLTAGE_PARK,
+            "voltage_park",
+            "staging-voltagepark",
+            "dev-voltagepark",
+            {"voltage_park_v1": V1VoltageParkDirectV1()},
+        ),
+    ],
+)
+def test_global_provider_mapping_falls_back_when_canonical_missing(
+    api, provider, cloud, first_id, second_id, spec_kwargs
+):
+    """Staging/local: if the prod canonical ID is absent, use the first available account."""
+    first = V1ExternalCluster(
+        id=first_id,
+        spec=V1ExternalClusterSpec(cluster_type=V1ClusterType.GLOBAL, **spec_kwargs),
+    )
+    second = V1ExternalCluster(
+        id=second_id,
+        spec=V1ExternalClusterSpec(cluster_type=V1ClusterType.GLOBAL, **spec_kwargs),
+    )
+
+    api._client.cluster_service_list_project_clusters.return_value = V1ListProjectClustersResponse(clusters=[])
+    api._client.cluster_service_list_clusters.return_value = V1ListClustersResponse(clusters=[first, second])
+
+    mapping = api.get_cloud_account_provider_mapping(teamspace_id="ts", global_only=True)
+    assert mapping[provider].id == first_id
+
+    result = api.resolve_cloud_account(teamspace_id="ts", cloud=cloud, default_cloud_account=None)
+    assert result == first_id
+
+
+def test_non_global_mapping_prefers_byoc_over_canonical_public(api):
+    """Data-connection mapping should still prefer private/BYOC accounts over public globals."""
+    byoc = V1ExternalCluster(
+        id="acc-byoc-aws",
+        spec=V1ExternalClusterSpec(driver=V1CloudProvider.AWS, cluster_type=V1ClusterType.BYOC, aws_v1=V1AWSDirectV1()),
+    )
+    public = V1ExternalCluster(
+        id="lightning-public-prod",
+        spec=V1ExternalClusterSpec(
+            driver=V1CloudProvider.AWS, cluster_type=V1ClusterType.GLOBAL, aws_v1=V1AWSDirectV1()
+        ),
+    )
+
+    api._client.cluster_service_list_project_clusters.return_value = V1ListProjectClustersResponse(
+        clusters=[public, byoc]
+    )
+    api._client.cluster_service_list_clusters.return_value = V1ListClustersResponse(clusters=[])
+
+    mapping = api.get_cloud_account_provider_mapping(teamspace_id="ts", global_only=False)
+    assert mapping[CloudProvider.AWS].id == "acc-byoc-aws"