[k8s] allow use of "k8s" instead of "kubernetes" in the CLI and python API

[cg505]

• sky.K8s() can be used in place of sky.Kubernetes()
• Commands such as sky check k8s or sky launch --cloud k8s will use Kubernetes as intended. Anywhere that CLOUD_REGISTRY.from_str() is used should work.
  • To support this, _CloudRegistry class was refactored.

Tested (run the relevant ones):
Need some advice on a good approach to write tests for this.

• Code formatting: bash format.sh
• Any manual or new tests for this PR (please specify below)
  • New test in test_cli.py for the k8s alias in the CLI
  • New test in test_api.py for sky.K8s() in the Python API
• All smoke tests: pytest tests/test_smoke.py
• Relevant individual smoke tests: pytest tests/test_smoke.py::test_fill_in_the_name
• Backward compatibility tests: conda deactivate; bash -i tests/backward_compatibility_tests.sh

[romilbhardwaj]

Thanks @cg505!

[romilbhardwaj]

Looks like from_name_or_alias and from_str do the same thing. For simplicity, should we just retain from_str() and move this functionality there?

[cg505]

Initially this is what I did, but there are cases where it's assumed that from_str will only work for the "canonical" name of the cloud.

For instance, if I hadn't made changes to show_gpus, updating from_str would cause sky show-gpus --cloud k8s to crash. Although this particular case is fixed by my changes, I'm worried there are other uses of from_str that rely on this behavior that I would be breaking.

[romilbhardwaj]

Can we have register support both syntax:

# Where no aliases required, not a callable
@ clouds.CLOUD_REGISTRY.register

# Callable where aliases are required:
clouds.CLOUD_REGISTRY.register(aliases=['k8s'])

Most python decorators follow this convention (e.g., @ray.remote, @dataclass ...)

[cg505]

Sure, I will look at those examples.

[romilbhardwaj]

Might be a little cleaner to have a reverse lookup in CLOUD_REGISTRY, e.g., CLOUD_REGISTRY.to_str(cloud)?

[romilbhardwaj]

Thanks @cg505! Left some minor comments.

[romilbhardwaj]

Should we retain the previous behavior and have the cloud inherit from dict (keeps it pythonic, can do cloud_registry.CLOUD_REGISTRY.keys()/.values() instead of cloud_registry.CLOUD_REGISTRY.clouds.keys())

class _CloudRegistry(dict):
    """Registry of clouds."""

    def __init__(self) -> None:
        super().__init__()
        self.aliases: Dict[str, str] = {}

    def from_str(self, name: Optional[str]) -> Optional['cloud.Cloud']:
        if name is None:
            return None

        search_name = name.lower()

        if search_name in self:
            return self[search_name]

        if search_name in self.aliases:
            return self[self.aliases[search_name]]

        with ux_utils.print_exception_no_traceback():
            raise ValueError(f'Cloud {name!r} is not a valid cloud among '
                             f'{[*self.keys(), *self.aliases.keys()]}')

    @overload
    def register(self, cloud_cls: Type['cloud.Cloud']) -> Type['cloud.Cloud']:
        ...

    @overload
    def register(
        self,
        cloud_cls: None = None,
        aliases: Optional[List[str]] = None,
    ) -> Callable[[Type['cloud.Cloud']], Type['cloud.Cloud']]:
        ...

    def register(
        self,
        cloud_cls: Optional[Type['cloud.Cloud']] = None,
        aliases: Optional[List[str]] = None,
    ) -> Union[Type['cloud.Cloud'], Callable[[Type['cloud.Cloud']],
                                             Type['cloud.Cloud']]]:

        def _register(cloud_cls: Type['cloud.Cloud']) -> Type['cloud.Cloud']:
            name = cloud_cls.canonical_name()
            assert name not in self, f'{name} already registered'
            self[name] = cloud_cls()

            for alias in aliases or []:
                assert alias not in self.aliases, (
                    f'alias {alias} already registered')
                self.aliases[alias] = name

            return cloud_cls

        if cloud_cls is not None:
            # invocation without parens (e.g. just `@register`)
            return _register(cloud_cls)

        # Invocation with parens (e.g. `@register(aliases=['alias'])`)
        return _register

[romilbhardwaj]

Thanks for adding cloud.canonical_name()! With that, looks like this method is redundant and can be removed since callers can directly call cloud.canonical_name()?

[cg505]

This one supports optional values, but that's the only difference.

[romilbhardwaj]

Should we add a quick docstr to define the exact behavior of from_str?

"""Returns the cloud instance from the canonical name or alias."""

[romilbhardwaj]

nice @cg505! lgtm.