lightning-sdk 0.1.58__py3-none-any.whl → 0.2.1__py3-none-any.whl

lightning_sdk/cli/switch.py

@@ -1,43 +1,60 @@
 from typing import Optional
 
+import click
+
 from lightning_sdk import Machine, Studio
+from lightning_sdk.lightning_cloud.openapi.rest import ApiException
+
+_MACHINE_VALUES = tuple([machine.name for machine in Machine.__dict__.values() if isinstance(machine, Machine)])
 
 
-class _Switch:
+@click.group("switch")
+def switch() -> None:
     """Switch machines for resources on the Lightning AI platform."""
 
-    def __init__(self) -> None:
-        _machine_values = tuple([machine.name for machine in Machine.__dict__.values() if isinstance(machine, Machine)])
-
-        docstr_studio = f"""Switch a studio to a given machine.
-
-        Args:
-            name: The name of the studio to start.
-                If not specified, tries to infer from the environment (e.g. when run from within a Studio.)
-            teamspace: The teamspace the studio is part of. Should be of format <OWNER>/<TEAMSPACE_NAME>.
-                If not specified, tries to infer from the environment (e.g. when run from within a Studio.)
-            machine: The machine type to switch to. One of {", ".join(_machine_values)}.
-                Defaults to the CPU Machine.
-        """
-        self.studio.__func__.__doc__ = docstr_studio
-
-    def studio(self, name: Optional[str] = None, teamspace: Optional[str] = None, machine: str = "CPU") -> None:
-        if teamspace is not None:
-            ts_splits = teamspace.split("/")
-            if len(ts_splits) != 2:
-                raise ValueError(f"Teamspace should be of format <OWNER>/<TEAMSPACE_NAME> but got {teamspace}")
-            owner, teamspace = ts_splits
-        else:
-            owner, teamspace = None, None
-
-        try:
-            studio = Studio(name=name, teamspace=teamspace, org=owner, user=None, create_ok=False)
-        except (RuntimeError, ValueError):
-            studio = Studio(name=name, teamspace=teamspace, org=None, user=owner, create_ok=False)
-
-        try:
-            resolved_machine = getattr(Machine, machine.upper(), Machine(machine, machine))
-        except KeyError:
-            resolved_machine = machine
-
-        studio.switch_machine(resolved_machine)
+
+@switch.command("studio")
+@click.argument("name")
+@click.option(
+    "--teamspace",
+    default=None,
+    help=(
+        "The teamspace the studio is part of. "
+        "Should be of format <OWNER>/<TEAMSPACE_NAME>. "
+        "If not specified, tries to infer from the environment (e.g. when run from within a Studio.)"
+    ),
+)
+@click.option(
+    "--machine",
+    default="CPU",
+    show_default=True,
+    type=click.Choice(_MACHINE_VALUES),
+    help="The machine type to switch to.",
+)
+def studio(name: str, teamspace: Optional[str] = None, machine: str = "CPU") -> None:
+    """Switch a studio to a given machine.
+
+    Example:
+      lightning switch studio NAME --machine=CPU
+
+    NAME: the name of the studio to switch machine for.
+    """
+    if teamspace is not None:
+        ts_splits = teamspace.split("/")
+        if len(ts_splits) != 2:
+            raise ValueError(f"Teamspace should be of format <OWNER>/<TEAMSPACE_NAME> but got {teamspace}")
+        owner, teamspace = ts_splits
+    else:
+        owner, teamspace = None, None
+
+    try:
+        studio = Studio(name=name, teamspace=teamspace, org=owner, user=None, create_ok=False)
+    except (RuntimeError, ValueError, ApiException):
+        studio = Studio(name=name, teamspace=teamspace, org=None, user=owner, create_ok=False)
+
+    try:
+        resolved_machine = getattr(Machine, machine.upper(), Machine(machine, machine))
+    except KeyError:
+        resolved_machine = machine
+
+    studio.switch_machine(resolved_machine)