Repository: skypilot-org/skypilot (python)

Issue:
[k8s] `lstrip('ssh-')` causes context starting with any of the charactor in `ssh-` will be stripped

Pull request:
[Core] removes lstrip('ssh-')
<!-- Describe the changes in this PR -->

Fixes #5981 

<!-- Describe the tests ran -->
<!-- Unit tests (tests/test_*.py) are part of GitHub CI; below are tests that launch on the cloud. -->

Tested (run the relevant ones):

- [x] Code formatting: install pre-commit (auto-check on commit) or `bash format.sh`
- [ ] Any manual or new tests for this PR (please specify below)
- [ ] All smoke tests: `/smoke-test` (CI) or `pytest tests/test_smoke.py` (local)
- [ ] Relevant individual tests: `/smoke-test -k test_name` (CI) or `pytest tests/test_smoke.py::test_name` (local)
- [ ] Backward compatibility: `/quicktest-core` (CI) or `pytest tests/smoke_tests/test_backward_compat.py` (local)

<!-- CI commands (/-prefixed) can only be triggered by repo members -->


Patch:
diff --git a/sky/adaptors/kubernetes.py b/sky/adaptors/kubernetes.py
index eece26d3d..c5936666b 100644
--- a/sky/adaptors/kubernetes.py
+++ b/sky/adaptors/kubernetes.py
@@ -93,12 +93,13 @@ def _load_config(context: Optional[str] = None):
             context_name = '(current-context)' if context is None else context
             is_ssh_node_pool = False
             if context_name.startswith('ssh-'):
-                context_name = context_name.lstrip('ssh-')
+                context_name = common_utils.removeprefix(context_name, 'ssh-')
                 is_ssh_node_pool = True
             # Check if exception was due to no current-context
             if 'Expected key current-context' in str(e):
                 if is_ssh_node_pool:
-                    context_name = context_name.lstrip('ssh-')
+                    context_name = common_utils.removeprefix(
+                        context_name, 'ssh-')
                     err_str = ('Failed to load SSH Node Pool configuration for '
                                f'{context_name!r}.\n'
                                '    Run `sky ssh up --infra {context_name}` to '
diff --git a/sky/check.py b/sky/check.py
index c530e955a..05d52ff1f 100644
--- a/sky/check.py
+++ b/sky/check.py
@@ -16,6 +16,7 @@ from sky import skypilot_config
 from sky.adaptors import cloudflare
 from sky.clouds import cloud as sky_cloud
 from sky.skylet import constants
+from sky.utils import common_utils
 from sky.utils import registry
 from sky.utils import rich_utils
 from sky.utils import subprocess_utils
@@ -561,7 +562,7 @@ def _format_context_details(cloud: Union[str, sky_clouds.Cloud],
             # TODO: This is a hack to remove the 'ssh-' prefix from the
             # context name. Once we have a separate kubeconfig for SSH,
             # this will not be required.
-            cleaned_context = context.lstrip('ssh-')
+            cleaned_context = common_utils.removeprefix(context, 'ssh-')
         else:
             cleaned_context = context
         symbol = (ux_utils.INDENT_LAST_SYMBOL if i == len(filtered_contexts) -
diff --git a/sky/cli.py b/sky/cli.py
index b8472e0ba..070213e9c 100644
--- a/sky/cli.py
+++ b/sky/cli.py
@@ -3613,7 +3613,7 @@ def show_gpus(
                 if not _filter_ctx(ctx):
                     continue
                 if is_ssh:
-                    display_ctx = ctx.lstrip('ssh-')
+                    display_ctx = common_utils.removeprefix(ctx, 'ssh-')
                 else:
                     display_ctx = ctx
                 num_filtered_contexts += 1
diff --git a/sky/clouds/ssh.py b/sky/clouds/ssh.py
index 6b257cb58..def4fee2e 100644
--- a/sky/clouds/ssh.py
+++ b/sky/clouds/ssh.py
@@ -12,6 +12,7 @@ from sky.adaptors import kubernetes as kubernetes_adaptor
 from sky.clouds import kubernetes
 from sky.provision.kubernetes import utils as kubernetes_utils
 from sky.utils import annotations
+from sky.utils import common_utils
 from sky.utils import registry
 
 if typing.TYPE_CHECKING:
@@ -88,8 +89,10 @@ class SSH(kubernetes.Kubernetes):
         all_contexts = self.existing_allowed_contexts()
 
         if region is not None and region not in all_contexts:
-            region_name = region.lstrip('ssh-')
-            available_contexts = [c.lstrip('ssh-') for c in all_contexts]
+            region_name = common_utils.removeprefix(region, 'ssh-')
+            available_contexts = [
+                common_utils.removeprefix(c, 'ssh-') for c in all_contexts
+            ]
             err_str = (f'SSH Node Pool {region_name!r} is not set up. '
                        'Run `sky check` for more details. ')
             if available_contexts:
@@ -155,7 +158,8 @@ class SSH(kubernetes.Kubernetes):
             if allowed_node_pools is None:
                 return ctxs
             return [
-                ctx for ctx in ctxs if ctx.lstrip('ssh-') in allowed_node_pools
+                ctx for ctx in ctxs
+                if common_utils.removeprefix(ctx, 'ssh-') in allowed_node_pools
             ]
 
         if all_node_pool_contexts:
diff --git a/sky/provision/kubernetes/instance.py b/sky/provision/kubernetes/instance.py
index 98ba11dd9..ecd088d9a 100644
--- a/sky/provision/kubernetes/instance.py
+++ b/sky/provision/kubernetes/instance.py
@@ -1277,7 +1277,8 @@ def query_instances(
     except kubernetes.max_retry_error():
         with ux_utils.print_exception_no_traceback():
             if is_ssh:
-                node_pool = context.lstrip('ssh-') if context else ''
+                node_pool = common_utils.removeprefix(context,
+                                                      'ssh-') if context else ''
                 msg = (
                     f'Cannot connect to SSH Node Pool {node_pool}. '
                     'Please check if the SSH Node Pool is up and accessible. '
diff --git a/sky/provision/kubernetes/utils.py b/sky/provision/kubernetes/utils.py
index 9963186c5..fba12bde1 100644
--- a/sky/provision/kubernetes/utils.py
+++ b/sky/provision/kubernetes/utils.py
@@ -1167,8 +1167,8 @@ def get_accelerator_label_key_values(
 
     is_ssh_node_pool = context.startswith('ssh-') if context else False
     cloud_name = 'SSH Node Pool' if is_ssh_node_pool else 'Kubernetes cluster'
-    context_display_name = context.lstrip('ssh-') if (
-        context and is_ssh_node_pool) else context
+    context_display_name = common_utils.removeprefix(
+        context, 'ssh-') if (context and is_ssh_node_pool) else context
 
     autoscaler_type = get_autoscaler_type()
     if autoscaler_type is not None:
diff --git a/sky/resources.py b/sky/resources.py
index 43cb359ac..c72aec488 100644
--- a/sky/resources.py
+++ b/sky/resources.py
@@ -480,7 +480,7 @@ class Resources:
         if self.region is not None:
             region_name = self.region
             if self.region.startswith('ssh-'):
-                region_name = self.region.lstrip('ssh-')
+                region_name = common_utils.removeprefix(self.region, 'ssh-')
             region_str = f', region={region_name}'
         zone_str = ''
         if self.zone is not None:
diff --git a/sky/utils/common_utils.py b/sky/utils/common_utils.py
index a0ca0e97a..5745c98b0 100644
--- a/sky/utils/common_utils.py
+++ b/sky/utils/common_utils.py
@@ -1017,3 +1017,9 @@ def _get_cgroup_memory_limit() -> Optional[int]:
 def _is_cgroup_v2() -> bool:
     """Return True if the environment is running cgroup v2."""
     return os.path.isfile('/sys/fs/cgroup/cgroup.controllers')
+
+
+def removeprefix(string: str, prefix: str) -> str:
+    if string.startswith(prefix):
+        return string[len(prefix):]
+    return string
diff --git a/sky/utils/infra_utils.py b/sky/utils/infra_utils.py
index 16560b52f..30e9511a2 100644
--- a/sky/utils/infra_utils.py
+++ b/sky/utils/infra_utils.py
@@ -180,7 +180,7 @@ class InfraInfo:
             # Node Pools.
             # TODO(romilb): This is a workaround while we use the global
             # kubeconfig to store the ssh contexts.
-            region_or_zone = self.region.lstrip('ssh-')
+            region_or_zone = common_utils.removeprefix(self.region, 'ssh-')
 
         if region_or_zone is not None and truncate:
             region_or_zone = common_utils.truncate_long_string(


Test patch:
diff --git a/tests/unit_tests/test_sky/clouds/test_ssh.py b/tests/unit_tests/test_sky/clouds/test_ssh.py
index 4d8e9d9a6..b52ce78fe 100644
--- a/tests/unit_tests/test_sky/clouds/test_ssh.py
+++ b/tests/unit_tests/test_sky/clouds/test_ssh.py
@@ -283,10 +283,7 @@ class TestSSHExistingAllowedContexts(unittest.TestCase):
             'regular-ctx', 'ssh-prod', 'ssh-dev', 'ssh-staging', 'ssh-test'
         ]
         mock_get_workspace_cloud.return_value.get.return_value = None
-        # Only allow specific node pools - note: lstrip('ssh-') behavior
-        # 'ssh-prod'.lstrip('ssh-') -> 'prod'
-        # 'ssh-staging'.lstrip('ssh-') -> 'taging' (not 'staging'!)
-        mock_get_nested.return_value = ['prod', 'taging', 'nonexistent']
+        mock_get_nested.return_value = ['prod', 'staging', 'nonexistent']
 
         # Node pools file has some contexts
         with patch.object(ssh.SSH,