dstack 0.0.9__py3-none-any.whl → 0.20.7__py3-none-any.whl

dstack/_internal/cli/commands/login.py

@@ -0,0 +1,352 @@
+import argparse
+import queue
+import sys
+import threading
+import urllib.parse
+import webbrowser
+from http.server import BaseHTTPRequestHandler, HTTPServer
+from typing import Any, Optional
+
+import questionary
+from rich.prompt import Prompt as RichPrompt
+from rich.text import Text
+
+from dstack._internal.cli.commands import BaseCommand
+from dstack._internal.cli.commands.project import select_default_project
+from dstack._internal.cli.utils.common import console, resolve_url
+from dstack._internal.core.errors import ClientError, CLIError
+from dstack._internal.core.models.users import UserWithCreds
+from dstack._internal.utils.logging import get_logger
+from dstack.api._public.runs import ConfigManager
+from dstack.api.server import APIClient
+
+logger = get_logger(__name__)
+
+is_project_menu_supported = sys.stdin.isatty()
+
+
+class UrlPrompt(RichPrompt):
+    def render_default(self, default: Any) -> Text:
+        return Text(f"({default})", style="bold orange1")
+
+
+class LoginCommand(BaseCommand):
+    NAME = "login"
+    DESCRIPTION = "Authorize the CLI using Single Sign-On"
+
+    def _register(self):
+        super()._register()
+        self._parser.add_argument(
+            "--url",
+            help="The server URL, e.g. https://sky.dstack.ai",
+            required=not is_project_menu_supported,
+        )
+        self._parser.add_argument(
+            "-p",
+            "--provider",
+            help=(
+                "The SSO provider name."
+                " Selected automatically if the server supports only one provider."
+            ),
+        )
+        self._parser.add_argument(
+            "-y",
+            "--yes",
+            help="Don't ask for confirmation (e.g. set first project as default)",
+            action="store_true",
+        )
+        self._parser.add_argument(
+            "-n",
+            "--no",
+            help="Don't ask for confirmation (e.g. do not change default project)",
+            action="store_true",
+        )
+
+    def _command(self, args: argparse.Namespace):
+        super()._command(args)
+        url = args.url
+        if url is None:
+            url = self._prompt_url()
+        base_url = _normalize_url_or_error(url)
+        api_client = APIClient(base_url=base_url)
+        provider = self._select_provider_or_error(api_client=api_client, provider=args.provider)
+        server = _LoginServer(api_client=api_client, provider=provider)
+        try:
+            server.start()
+            auth_resp = api_client.auth.authorize(provider=provider, local_port=server.port)
+            opened = webbrowser.open(auth_resp.authorization_url)
+            if opened:
+                console.print(
+                    f"Your browser has been opened to log in with [code]{provider.title()}[/]:\n"
+                )
+            else:
+                console.print(f"Open the URL to log in with [code]{provider.title()}[/]:\n")
+            print(f"{auth_resp.authorization_url}\n")
+            user = server.get_logged_in_user()
+        finally:
+            server.shutdown()
+        if user is None:
+            raise CLIError("CLI authentication failed")
+        console.print(f"Logged in as [code]{user.username}[/]")
+        api_client = APIClient(base_url=base_url, token=user.creds.token)
+        self._configure_projects(api_client=api_client, user=user, args=args)
+
+    def _select_provider_or_error(self, api_client: APIClient, provider: Optional[str]) -> str:
+        providers = api_client.auth.list_providers()
+        available_providers = [p.name for p in providers if p.enabled]
+        if len(available_providers) == 0:
+            raise CLIError("No SSO providers configured on the server.")
+        if provider is None:
+            if len(available_providers) > 1:
+                if is_project_menu_supported:
+                    return self._prompt_provider(available_providers)
+                raise CLIError(
+                    "Specify -p/--provider to choose SSO provider"
+                    f" Available providers: {', '.join(available_providers)}"
+                )
+            return available_providers[0]
+        if provider not in available_providers:
+            raise CLIError(
+                f"Provider {provider} not configured on the server."
+                f" Available providers: {', '.join(available_providers)}"
+            )
+        return provider
+
+    def _prompt_url(self) -> str:
+        try:
+            url = UrlPrompt.ask(
+                "Enter the server URL",
+                default="https://sky.dstack.ai",
+                console=console,
+            )
+        except KeyboardInterrupt:
+            console.print("\nCancelled by user")
+            raise SystemExit(1)
+        if url is None:
+            raise CLIError("URL is required")
+        return url
+
+    def _prompt_provider(self, available_providers: list[str]) -> str:
+        choices = [
+            questionary.Choice(title=provider, value=provider) for provider in available_providers
+        ]
+        selected_provider = questionary.select(
+            message="Select SSO provider:",
+            choices=choices,
+            qmark="",
+            instruction="(↑↓ Enter)",
+        ).ask()
+        if selected_provider is None:
+            raise SystemExit(1)
+        return selected_provider
+
+    def _configure_projects(
+        self, api_client: APIClient, user: UserWithCreds, args: argparse.Namespace
+    ):
+        projects = api_client.projects.list(include_not_joined=False)
+        if len(projects) == 0:
+            console.print(
+                "No projects configured."
+                " Create your own project via the UI or contact a project manager to add you to the project."
+            )
+            return
+        config_manager = ConfigManager()
+        default_project = config_manager.get_project_config()
+        for project in projects:
+            config_manager.configure_project(
+                name=project.project_name,
+                url=api_client.base_url,
+                token=user.creds.token,
+                default=False,
+            )
+        config_manager.save()
+        project_names = ", ".join(f"[code]{p.project_name}[/]" for p in projects)
+        console.print(
+            f"Added {project_names} project{'' if len(projects) == 1 else 's'} at {config_manager.config_filepath}"
+        )
+
+        project_configs = config_manager.list_project_configs()
+
+        if args.no:
+            return
+
+        if args.yes:
+            if len(projects) > 0:
+                first_project_from_server = projects[0]
+                first_project_config = next(
+                    (
+                        pc
+                        for pc in project_configs
+                        if pc.name == first_project_from_server.project_name
+                    ),
+                    None,
+                )
+                if first_project_config is not None:
+                    config_manager.configure_project(
+                        name=first_project_config.name,
+                        url=first_project_config.url,
+                        token=first_project_config.token,
+                        default=True,
+                    )
+                    config_manager.save()
+                    console.print(
+                        f"Set [code]{first_project_config.name}[/] project as default at {config_manager.config_filepath}"
+                    )
+            return
+
+        if len(project_configs) == 1 or not is_project_menu_supported:
+            selected_project = None
+            if len(project_configs) == 1:
+                selected_project = project_configs[0]
+            else:
+                for i, project in enumerate(projects):
+                    set_as_default = (
+                        default_project is None
+                        and i == 0
+                        or default_project is not None
+                        and default_project.name == project.project_name
+                    )
+                    if set_as_default:
+                        selected_project = next(
+                            (pc for pc in project_configs if pc.name == project.project_name),
+                            None,
+                        )
+                        break
+            if selected_project is not None:
+                config_manager.configure_project(
+                    name=selected_project.name,
+                    url=selected_project.url,
+                    token=selected_project.token,
+                    default=True,
+                )
+                config_manager.save()
+                console.print(
+                    f"Set [code]{selected_project.name}[/] project as default at {config_manager.config_filepath}"
+                )
+        else:
+            console.print()
+            selected_project = select_default_project(project_configs, default_project)
+            if selected_project is not None:
+                config_manager.configure_project(
+                    name=selected_project.name,
+                    url=selected_project.url,
+                    token=selected_project.token,
+                    default=True,
+                )
+                config_manager.save()
+
+
+class _BadRequestError(Exception):
+    pass
+
+
+class _LoginServer:
+    def __init__(self, api_client: APIClient, provider: str):
+        self._api_client = api_client
+        self._provider = provider
+        self._result_queue: queue.Queue[Optional[UserWithCreds]] = queue.Queue()
+        # Using built-in HTTP server to avoid extra deps.
+        callback_handler = self._make_callback_handler(
+            result_queue=self._result_queue,
+            api_client=api_client,
+            provider=provider,
+        )
+        self._server = self._create_server(handler=callback_handler)
+
+    def start(self):
+        self._thread = threading.Thread(target=self._server.serve_forever)
+        self._thread.start()
+
+    def shutdown(self):
+        self._server.shutdown()
+
+    def get_logged_in_user(self) -> Optional[UserWithCreds]:
+        return self._result_queue.get()
+
+    @property
+    def port(self) -> int:
+        return self._server.server_port
+
+    def _make_callback_handler(
+        self,
+        result_queue: queue.Queue[Optional[UserWithCreds]],
+        api_client: APIClient,
+        provider: str,
+    ) -> type[BaseHTTPRequestHandler]:
+        class _CallbackHandler(BaseHTTPRequestHandler):
+            def do_GET(self):
+                parsed_path = urllib.parse.urlparse(self.path)
+                if parsed_path.path != "/auth/callback":
+                    self.send_response(404)
+                    self.end_headers()
+                    return
+                try:
+                    self._handle_auth_callback(parsed_path)
+                except _BadRequestError as e:
+                    self.send_error(400, e.args[0])
+                    result_queue.put(None)
+
+            def log_message(self, format: str, *args):
+                # Do not log server requests.
+                pass
+
+            def _handle_auth_callback(self, parsed_path: urllib.parse.ParseResult):
+                try:
+                    params = urllib.parse.parse_qs(parsed_path.query, strict_parsing=True)
+                except ValueError:
+                    raise _BadRequestError("Bad query params")
+                code = params.get("code", [None])[0]
+                state = params.get("state", [None])[0]
+                if code is None or state is None:
+                    raise _BadRequestError("Missing required params")
+                try:
+                    user = api_client.auth.callback(provider=provider, code=code, state=state)
+                except ClientError:
+                    raise _BadRequestError("Authentication failed")
+                self._send_success_html()
+                result_queue.put(user)
+
+            def _send_success_html(self):
+                body = _SUCCESS_HTML.encode()
+                self.send_response(200)
+                self.send_header("Content-Type", "text/html; charset=utf-8")
+                self.send_header("Content-Length", str(len(body)))
+                self.end_headers()
+                self.wfile.write(body)
+
+        return _CallbackHandler
+
+    def _create_server(self, handler: type[BaseHTTPRequestHandler]) -> HTTPServer:
+        server_address = ("127.0.0.1", 0)
+        server = HTTPServer(server_address, handler)
+        return server
+
+
+def _normalize_url_or_error(url: str) -> str:
+    try:
+        # Validate the URL and determine the URL scheme.
+        # Need to resolve the scheme before making first POST request
+        # since for some redirect codes (301), clients change POST to GET.
+        url = resolve_url(url)
+    except ValueError as e:
+        raise CLIError(e.args[0])
+    return url
+
+
+_SUCCESS_HTML = """\
+<!DOCTYPE HTML>
+<html lang="en">
+    <head>
+        <meta charset="utf-8">
+        <title>CLI authenticated</title>
+        <style>
+            body {font-family: system-ui, sans-serif; }
+            h1 { font-weight: 500; }
+        </style>
+    </head>
+    <body>
+        <h1>dstack CLI authenticated</h1>
+        <p>You may close this page.</p>
+    </body>
+</html>
+"""

dstack/_internal/cli/commands/logs.py

@@ -0,0 +1,62 @@
+import argparse
+import sys
+
+from dstack._internal.cli.commands import APIBaseCommand
+from dstack._internal.cli.services.completion import RunNameCompleter
+from dstack._internal.cli.utils.common import get_start_time
+from dstack._internal.core.errors import CLIError
+from dstack._internal.utils.logging import get_logger
+
+logger = get_logger(__name__)
+
+
+class LogsCommand(APIBaseCommand):
+    NAME = "logs"
+    DESCRIPTION = "Show logs"
+
+    def _register(self):
+        super()._register()
+        self._parser.add_argument(
+            "-d", "--diagnose", action="store_true", help="Show run diagnostic logs"
+        )
+        self._parser.add_argument(
+            "--replica",
+            help="The replica number. Defaults to 0.",
+            type=int,
+            default=0,
+        )
+        self._parser.add_argument(
+            "--job",
+            help="The job number inside the replica. Defaults to 0.",
+            type=int,
+            default=0,
+        )
+        self._parser.add_argument(
+            "--since",
+            help=(
+                "Show only logs newer than the specified date."
+                " Can be a duration (e.g. 10s, 5m, 1d) or an RFC 3339 string (e.g. 2023-09-24T15:30:00Z)."
+            ),
+            type=str,
+        )
+        self._parser.add_argument("run_name").completer = RunNameCompleter(all=True)  # type: ignore[attr-defined]
+
+    def _command(self, args: argparse.Namespace):
+        super()._command(args)
+        run = self.api.runs.get(args.run_name)
+        if run is None:
+            raise CLIError(f"Run {args.run_name} not found")
+
+        start_time = get_start_time(args.since)
+        logs = run.logs(
+            start_time=start_time,
+            diagnose=args.diagnose,
+            replica_num=args.replica,
+            job_num=args.job,
+        )
+        try:
+            for log in logs:
+                sys.stdout.buffer.write(log)
+                sys.stdout.buffer.flush()
+        except KeyboardInterrupt:
+            pass

dstack/_internal/cli/commands/metrics.py

@@ -0,0 +1,153 @@
+import argparse
+import time
+from typing import Any, Dict, List, Optional, Union
+
+from rich.live import Live
+from rich.table import Table
+
+from dstack._internal.cli.commands import APIBaseCommand
+from dstack._internal.cli.services.completion import RunNameCompleter
+from dstack._internal.cli.utils.common import (
+    LIVE_TABLE_PROVISION_INTERVAL_SECS,
+    LIVE_TABLE_REFRESH_RATE_PER_SEC,
+    add_row_from_dict,
+    console,
+)
+from dstack._internal.core.errors import CLIError
+from dstack._internal.core.models.instances import Resources
+from dstack._internal.core.models.metrics import JobMetrics
+from dstack.api._public import Client
+from dstack.api._public.runs import Run
+
+
+class MetricsCommand(APIBaseCommand):
+    NAME = "metrics"
+    DESCRIPTION = "Show run metrics"
+
+    def _register(self):
+        super()._register()
+        self._parser.add_argument("run_name").completer = RunNameCompleter()  # type: ignore[attr-defined]
+        self._parser.add_argument(
+            "-w",
+            "--watch",
+            help="Watch run metrics in realtime",
+            action="store_true",
+        )
+
+    def _command(self, args: argparse.Namespace):
+        super()._command(args)
+        run = self.api.runs.get(run_name=args.run_name)
+        if run is None:
+            raise CLIError(f"Run {args.run_name} not found")
+        metrics = _get_run_jobs_metrics(api=self.api, run=run)
+
+        if not args.watch:
+            console.print(_get_metrics_table(run, metrics))
+            return
+
+        try:
+            with Live(console=console, refresh_per_second=LIVE_TABLE_REFRESH_RATE_PER_SEC) as live:
+                while True:
+                    live.update(_get_metrics_table(run, metrics))
+                    time.sleep(LIVE_TABLE_PROVISION_INTERVAL_SECS)
+                    run = self.api.runs.get(run_name=args.run_name)
+                    if run is None:
+                        raise CLIError(f"Run {args.run_name} not found")
+                    metrics = _get_run_jobs_metrics(api=self.api, run=run)
+        except KeyboardInterrupt:
+            pass
+
+
+def _get_run_jobs_metrics(api: Client, run: Run) -> List[JobMetrics]:
+    metrics = []
+    for job in run._run.jobs:
+        job_metrics = api.client.metrics.get_job_metrics(
+            project_name=api.project,
+            run_name=run.name,
+            replica_num=job.job_spec.replica_num,
+            job_num=job.job_spec.job_num,
+        )
+        metrics.append(job_metrics)
+    return metrics
+
+
+def _get_metrics_table(run: Run, metrics: List[JobMetrics]) -> Table:
+    table = Table(box=None)
+    table.add_column("NAME", style="bold", no_wrap=True)
+    table.add_column("STATUS")
+    table.add_column("CPU")
+    table.add_column("MEMORY")
+    table.add_column("GPU")
+
+    run_row: Dict[Union[str, int], Any] = {"NAME": run.name, "STATUS": run.status.value}
+    if len(run._run.jobs) != 1:
+        add_row_from_dict(table, run_row)
+
+    for job, job_metrics in zip(run._run.jobs, metrics):
+        jrd = job.job_submissions[-1].job_runtime_data
+        jpd = job.job_submissions[-1].job_provisioning_data
+        resources: Optional[Resources] = None
+        if jrd is not None and jrd.offer is not None:
+            resources = jrd.offer.instance.resources
+        elif jpd is not None:
+            resources = jpd.instance_type.resources
+        cpu_usage = _get_metric_value(job_metrics, "cpu_usage_percent")
+        if cpu_usage is not None:
+            if resources is not None:
+                cpu_usage = cpu_usage / resources.cpus
+            cpu_usage = f"{cpu_usage:.0f}%"
+        memory_usage = _get_metric_value(job_metrics, "memory_working_set_bytes")
+        if memory_usage is not None:
+            memory_usage = _format_memory(memory_usage, 2)
+            if resources is not None:
+                memory_usage += f"/{_format_memory(resources.memory_mib * 1024 * 1024, 2)}"
+        gpu_metrics = ""
+        gpus_detected_num = _get_metric_value(job_metrics, "gpus_detected_num")
+        if gpus_detected_num is not None:
+            for i in range(gpus_detected_num):
+                gpu_memory_usage = _get_metric_value(job_metrics, f"gpu_memory_usage_bytes_gpu{i}")
+                gpu_util_percent = _get_metric_value(job_metrics, f"gpu_util_percent_gpu{i}")
+                if gpu_memory_usage is not None:
+                    if i != 0:
+                        gpu_metrics += "\n"
+                    gpu_metrics += f"gpu={i} mem={_format_memory(gpu_memory_usage, 2)}"
+                    if resources is not None:
+                        gpu_metrics += (
+                            f"/{_format_memory(resources.gpus[i].memory_mib * 1024 * 1024, 2)}"
+                        )
+                    gpu_metrics += f" util={gpu_util_percent}%"
+
+        job_row: Dict[Union[str, int], Any] = {
+            "NAME": f"  replica={job.job_spec.replica_num} job={job.job_spec.job_num}",
+            "STATUS": job.job_submissions[-1].status.value,
+            "CPU": cpu_usage or "-",
+            "MEMORY": memory_usage or "-",
+            "GPU": gpu_metrics or "-",
+        }
+        if len(run._run.jobs) == 1:
+            job_row.update(run_row)
+        add_row_from_dict(table, job_row)
+
+    return table
+
+
+def _get_metric_value(job_metrics: JobMetrics, name: str) -> Optional[Any]:
+    for metric in job_metrics.metrics:
+        if metric.name == name:
+            return metric.values[-1]
+    return None
+
+
+def _format_memory(memory_bytes: int, decimal_places: int) -> str:
+    """See test_format_memory in tests/_internal/cli/commands/test_metrics.py for examples."""
+    memory_mb = memory_bytes / 1024 / 1024
+    if memory_mb >= 1024:
+        value = memory_mb / 1024
+        unit = "GB"
+    else:
+        value = memory_mb
+        unit = "MB"
+
+    if decimal_places == 0:
+        return f"{round(value)}{unit}"
+    return f"{value:.{decimal_places}f}".rstrip("0").rstrip(".") + unit