flyteplugins-union 0.2.1__tar.gz → 0.3.0__tar.gz

{flyteplugins_union-0.2.1/src/flyteplugins_union.egg-info → flyteplugins_union-0.3.0}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: flyteplugins-union
-Version: 0.2.1
+Version: 0.3.0
 Summary: Union SDK - Proprietary extensions for Flyte
 Author-email: unionai <info@union.ai>
 Project-URL: Homepage, https://www.union.ai/
@@ -17,7 +17,8 @@ Classifier: Topic :: Software Development
 Requires-Python: >=3.10
 Description-Content-Type: text/markdown
 License-File: LICENSE
-Requires-Dist: flyte<2.1.0,>=2.0.0
+Requires-Dist: connectrpc>=0.9.0
+Requires-Dist: flyte<2.5.0,>=2.1.0
 Requires-Dist: flyteidl2>=2.0.0
 Requires-Dist: pyyaml>=6.0
 Requires-Dist: rich-click>=1.8.0

{flyteplugins_union-0.2.1 → flyteplugins_union-0.3.0}/licenses/BSL.txt

@@ -17,7 +17,7 @@ Additional Use Grant: You may make use of the Licensed Work, provided that
                       computer programs whose source code are controlled by
                       such third parties.
 
-Change Date:          2030-03-05
+Change Date:          2030-05-06
 
 Change License:       Apache License, Version 2.0
 

{flyteplugins_union-0.2.1 → flyteplugins_union-0.3.0}/pyproject.toml

@@ -10,7 +10,8 @@ description = "Union SDK - Proprietary extensions for Flyte"
 readme = "README.md"
 requires-python = ">=3.10"
 dependencies = [
-    "flyte>=2.0.0,<2.1.0",
+    "connectrpc>=0.9.0",
+    "flyte>=2.1.0,<2.5.0",
     "flyteidl2>=2.0.0",
     "pyyaml>=6.0",
     "rich-click>=1.8.0",
@@ -29,7 +30,6 @@ classifiers = [
 
 [dependency-groups]
 dev = [
-    "grpcio-tools>=1.71.0",
     "build>=1.2.2.post1",
     "mock>=5.2.0",
     "pytest>=8.3.5",
@@ -63,6 +63,10 @@ Documentation = "https://docs.union.ai/"
 "get.user" = "flyteplugins.union.cli.user:get_user"
 "delete.user" = "flyteplugins.union.cli.user:delete_user"
 "get.member" = "flyteplugins.union.cli.member:get_member"
+"get.cluster" = "flyteplugins.union.cli.cluster:get_cluster"
+"create.queue" = "flyteplugins.union.cli.queue:create_queue"
+"get.queue" = "flyteplugins.union.cli.queue:get_queue"
+"update.queue" = "flyteplugins.union.cli.queue:update_queue"
 "create.assignment" = "flyteplugins.union.cli.assignment:create_assignment"
 "get.assignment" = "flyteplugins.union.cli.assignment:get_assignment"
 "delete.assignment" = "flyteplugins.union.cli.assignment:delete_assignment"
@@ -76,9 +80,7 @@ include = ["flyteplugins*"]
 
 [tool.ruff]
 line-length = 120
-extend-exclude = [
-    "src/flyteplugins/union/internal",
-]
+extend-exclude = ["src/flyteplugins/union/internal"]
 
 [tool.ruff.lint]
 select = [

{flyteplugins_union-0.2.1 → flyteplugins_union-0.3.0}/src/flyteplugins/union/cli/api_key.py

@@ -7,8 +7,14 @@ from flyteplugins.union.remote import ApiKey
 
 @click.command("api-key")
 @click.option("--name", type=str, help="Name for API key", required=True)
+@click.option(
+    "--no-default-policies",
+    "skip_default_policy_assignments",
+    is_flag=True,
+    help="Skip attaching the server's default policies. Grant access explicitly via 'flyte create assignment'.",
+)
 @click.pass_obj
-def create_api_key(cfg: common.CLIConfig, name: str):
+def create_api_key(cfg: common.CLIConfig, name: str, skip_default_policy_assignments: bool):
     """
     Create an API key for headless authentication.
 
@@ -22,6 +28,9 @@ def create_api_key(cfg: common.CLIConfig, name: str):
         # Create an API key named "ci-pipeline"
         $ flyte create api-key --name ci-pipeline
 
+        # Create a locked-down key with no default policy attachments
+        $ flyte create api-key --name ci-pipeline --no-default-policies
+
         # The output will include an export command like:
         # export FLYTE_API_KEY="<base64-encoded-credentials>"
     """
@@ -29,7 +38,7 @@ def create_api_key(cfg: common.CLIConfig, name: str):
     cfg.init(project="", domain="")
 
     try:
-        api_key = ApiKey.create(name=name)
+        api_key = ApiKey.create(name=name, skip_default_policy_assignments=skip_default_policy_assignments)
 
         console = common.get_console()
 

flyteplugins_union-0.3.0/src/flyteplugins/union/cli/cluster.py

@@ -0,0 +1,178 @@
+import rich_click as click
+from flyte.cli import _common as common
+
+from flyteplugins.union.remote import Cluster
+
+
+@click.command("cluster")
+@click.argument("name", type=str, required=False)
+@click.option("--limit", type=int, default=100, help="Maximum number of clusters to return.")
+@click.pass_obj
+def get_cluster(cfg: common.CLIConfig, name: str | None, limit: int):
+    """Get a cluster or list all clusters.
+
+    If NAME is provided, fetch that specific cluster and render a detailed view.
+    Otherwise list all clusters.
+
+    Examples:
+
+        $ flyte --org my-org get cluster
+
+        $ flyte --org my-org get cluster my-cluster
+    """
+    cfg.init(project="", domain="")
+    console = common.get_console()
+
+    try:
+        if name:
+            cluster = Cluster.get(name)
+            if cfg.output_format in ("json", "json-raw"):
+                console.print(common.format(f"Cluster {name}", [cluster], cfg.output_format))
+            else:
+                console.print(_render_cluster_detail(cluster))
+        else:
+            clusters = list(Cluster.listall(limit=limit))
+            if not clusters:
+                console.print("[yellow]No clusters found.[/yellow]")
+                return
+            enabled = [c for c in clusters if c.state == "enabled"]
+            disabled = [c for c in clusters if c.state != "enabled"]
+            if enabled:
+                console.print(common.format("Enabled Clusters", enabled, cfg.output_format))
+            if disabled:
+                if enabled:
+                    console.print()
+                console.print(common.format("Disabled Clusters", disabled, cfg.output_format))
+    except Exception as e:
+        raise click.ClickException(f"Unable to get clusters: {e}") from e
+
+
+# ---------------------------------------------------------------------------
+# Detailed single-cluster rendering
+# ---------------------------------------------------------------------------
+
+
+def _render_cluster_detail(cluster: Cluster):
+    """Render a multi-section Panel for a single cluster."""
+    from rich.panel import Panel  # noqa: PLC0415
+    from rich.table import Table  # noqa: PLC0415
+
+    outer = Table.grid(padding=(0, 0))
+    outer.add_column(ratio=1)
+
+    sections = [
+        ("Identity", _kv_table(cluster, [("name", "name"), ("organization", "organization")])),
+        (
+            "State",
+            _kv_table(
+                cluster,
+                [
+                    ("state", "state"),
+                    ("health", "health_display"),
+                    ("unhealthy reasons", "_unhealthy_reasons_str"),
+                ],
+            ),
+        ),
+        (
+            "Cloud",
+            _kv_table(
+                cluster,
+                [
+                    ("host", "cloud_host_name"),
+                    ("storage", "storage_type"),
+                    ("bucket", "bucket_name"),
+                    ("region", "bucket_region"),
+                    ("metadata prefix", "metadata_bucket_prefix"),
+                    ("gcp project", "gcp_project_id"),
+                    ("user role", "user_role"),
+                    ("operator app", "operator_app_id"),
+                ],
+                skip_empty=True,
+            ),
+        ),
+        (
+            "Config",
+            _kv_table(
+                cluster,
+                [
+                    ("assigned", "assigned_config_id"),
+                    ("synced", "synced_config_id"),
+                    ("synced at", "synced_at"),
+                    ("drift", "_drift_display"),
+                ],
+                skip_empty=True,
+            ),
+        ),
+        (
+            "Capacity",
+            _kv_table(cluster, [("resources", "capacity"), ("ingress (dataplane)", "dataplane_ingress_enabled")]),
+        ),
+        (
+            "Tunnel",
+            _kv_table(
+                cluster,
+                [("status", "tunnel_status_display"), ("public url", "tunnel_url")],
+                skip_empty=True,
+            ),
+        ),
+    ]
+
+    for title, body in sections:
+        if body is None:
+            continue
+        outer.add_row(Panel(body, title=f"[bold]{title}[/bold]", border_style="dim"))
+
+    if cluster.queues:
+        outer.add_row(_render_queues_panel(cluster.queues))
+
+    return outer
+
+
+def _kv_table(cluster: Cluster, fields: list[tuple[str, str]], *, skip_empty: bool = False):
+    """Render label / value rows. Returns None when skip_empty and every value is empty."""
+    from rich.table import Table  # noqa: PLC0415
+    from rich.text import Text  # noqa: PLC0415
+
+    rows = []
+    for label, accessor in fields:
+        value = _accessor(cluster, accessor)
+        if skip_empty and not value and value is not False:
+            continue
+        rows.append((label, value))
+
+    if not rows:
+        return None
+
+    t = Table.grid(padding=(0, 2))
+    t.add_column(style="dim", no_wrap=True)
+    t.add_column()
+    for label, value in rows:
+        t.add_row(label, Text.from_markup(str(value)) if value != "" else Text("—", style="dim"))
+    return t
+
+
+def _accessor(cluster: Cluster, name: str):
+    """Pull a property/value from the cluster, or compute a small helper string."""
+    if name == "_unhealthy_reasons_str":
+        return "; ".join(cluster.unhealthy_reasons) if cluster.unhealthy_reasons else ""
+    if name == "_drift_display":
+        if not cluster.assigned_config_id:
+            return ""
+        return "[bright_red]drifted[/bright_red]" if cluster.config_drift else "[green]in sync[/green]"
+    return getattr(cluster, name)
+
+
+def _render_queues_panel(queues: list[dict]):
+    """Render the cluster's bound queues as a small panel."""
+    from rich.panel import Panel  # noqa: PLC0415
+    from rich.table import Table  # noqa: PLC0415
+
+    t = Table.grid(padding=(0, 2))
+    t.add_column(style="dim", no_wrap=True)
+    t.add_column()
+    for q in queues:
+        scope_parts = [s for s in (q.get("project"), q.get("domain")) if s]
+        scope = "/".join(scope_parts) if scope_parts else "—"
+        t.add_row(q["name"], scope)
+
+    return Panel(t, title=f"[bold]Queues ({len(queues)})[/bold]", border_style="dim")

flyteplugins_union-0.3.0/src/flyteplugins/union/cli/queue.py

@@ -0,0 +1,366 @@
+import rich_click as click
+
+# All heavy imports (yaml, rich.*, remote objects) are deferred to inside
+# command functions so that registering CLI entry points stays fast.
+from flyteplugins.union.remote._queue import FAIRNESS_VALUES, PRIORITY_VALUES
+
+_PRIORITY_CHOICES = click.Choice(list(PRIORITY_VALUES), case_sensitive=False)
+_FAIRNESS_CHOICES = click.Choice(list(FAIRNESS_VALUES), case_sensitive=False)
+
+
+# ---------------------------------------------------------------------------
+# create queue
+# ---------------------------------------------------------------------------
+
+
+@click.command("queue")
+@click.argument("name", type=str)
+@click.option("--run-concurrency", type=int, required=True, help="Max concurrent runs (required)")
+@click.option("--action-concurrency", type=int, required=True, help="Max concurrent actions (required)")
+@click.option("--depth", type=int, default=10000, show_default=True, help="Max queue depth")
+@click.option("--priority", type=_PRIORITY_CHOICES, default="medium", show_default=True, help="Queue priority")
+@click.option(
+    "--fairness",
+    type=_FAIRNESS_CHOICES,
+    default="round_robin",
+    show_default=True,
+    help="Fairness algorithm",
+)
+@click.option("--cluster", "clusters", type=str, multiple=True, help="Target cluster(s). Repeat for multiple.")
+@click.option("--project", type=str, default="", help="Scope queue to a project")
+@click.option("--domain", type=str, default="", help="Scope queue to a domain")
+@click.pass_obj
+def create_queue(
+    cfg,
+    name: str,
+    run_concurrency: int,
+    action_concurrency: int,
+    depth: int,
+    priority: str,
+    fairness: str,
+    clusters: tuple[str, ...],
+    project: str,
+    domain: str,
+):
+    """Create a scheduling queue.
+
+    Examples:
+
+        $ flyte create queue my-queue --run-concurrency 100 --action-concurrency 1000
+
+        $ flyte create queue gpu-queue --run-concurrency 50 --action-concurrency 500 \\
+            --priority min --cluster gpu-cluster-1
+
+        $ flyte create queue backfill --run-concurrency 10 --action-concurrency 100 \\
+            --depth 5000 --priority max
+
+        $ flyte create queue team-queue --run-concurrency 100 --action-concurrency 1000 \\
+            --project my-project --domain production
+    """
+    from flyte.cli import _common as common  # noqa: PLC0415
+
+    from flyteplugins.union.remote import Queue  # noqa: PLC0415
+
+    cfg.init(project="", domain="")
+    console = common.get_console()
+
+    try:
+        queue = Queue.create(
+            name,
+            org=cfg.org,
+            domain=domain,
+            project=project,
+            run_concurrency=run_concurrency,
+            action_concurrency=action_concurrency,
+            depth=depth,
+            priority=priority,
+            fairness=fairness,
+            clusters=list(clusters) if clusters else None,
+        )
+        console.print(f"[green]✓[/green] Successfully created queue: [cyan]{queue.name}[/cyan]")
+        console.print(common.format(f"Queue {queue.name}", [queue], "json"))
+    except Exception as e:
+        raise click.ClickException(f"Unable to create queue: {e}") from e
+
+
+# ---------------------------------------------------------------------------
+# get queue
+# ---------------------------------------------------------------------------
+
+
+@click.command("queue")
+@click.argument("name", type=str, required=False)
+@click.option("--project", type=str, default="", help="Scope to a project")
+@click.option("--domain", type=str, default="", help="Scope to a domain")
+@click.option("--limit", type=int, default=100, help="Maximum number of queues to return")
+@click.option("--watch", is_flag=True, help="Stream live queue metrics (requires NAME)")
+@click.pass_obj
+def get_queue(cfg, name: str | None, project: str, domain: str, limit: int, watch: bool):
+    """Get a queue or list all queues.
+
+    If NAME is provided, fetch that specific queue with its current metrics.
+    Use --watch to stream live metrics with progress bars.
+    Otherwise list all queues.
+
+    Examples:
+
+        $ flyte get queue
+
+        $ flyte get queue my-queue
+
+        $ flyte get queue my-queue --watch
+    """
+    if watch and not name:
+        raise click.ClickException("--watch requires a queue NAME")
+
+    from flyte.cli import _common as common  # noqa: PLC0415
+
+    from flyteplugins.union.remote import Queue  # noqa: PLC0415
+
+    cfg.init(project="", domain="")
+    console = common.get_console()
+
+    try:
+        if name and watch:
+            _watch_queue(console, name, cfg.org, domain, project)
+        elif name:
+            queue = Queue.get(name, org=cfg.org, domain=domain, project=project)
+            console.print(common.format(f"Queue {name}", [queue], "json"))
+
+            # Show metrics snapshot
+            try:
+                metrics = Queue.details(name, org=cfg.org, domain=domain, project=project)
+                console.print(_render_metrics(metrics, queue))
+            except Exception:
+                pass  # Metrics are best-effort
+        else:
+            queues = list(Queue.listall(org=cfg.org, limit=limit))
+            if not queues:
+                console.print("[yellow]No queues found.[/yellow]")
+                return
+            console.print(common.format("Queues", queues, cfg.output_format))
+    except Exception as e:
+        raise click.ClickException(f"Unable to get queue(s): {e}") from e
+
+
+# ---------------------------------------------------------------------------
+# Watch display helpers
+# ---------------------------------------------------------------------------
+
+
+def _pct_color(ratio: float) -> str:
+    if ratio >= 0.9:
+        return "bright_red"
+    if ratio >= 0.7:
+        return "yellow"
+    return "medium_purple1"
+
+
+def _bar_style(ratio: float) -> str:
+    if ratio >= 0.9:
+        return "bright_red"
+    if ratio >= 0.7:
+        return "yellow"
+    return "medium_purple1"
+
+
+def _render_metric_block(label: str, used: int, cap: int):
+    """Render a single metric: header row with label/value/pct, then a slim bar."""
+    from rich.progress_bar import ProgressBar  # noqa: PLC0415
+    from rich.table import Table  # noqa: PLC0415
+    from rich.text import Text  # noqa: PLC0415
+
+    ratio = used / cap if cap > 0 else 0
+    pct = ratio * 100
+    color = _pct_color(ratio)
+    style = _bar_style(ratio)
+
+    block = Table.grid(padding=0)
+    block.add_column("label", ratio=1)
+    block.add_column("value", justify="right")
+    block.add_column("pct", justify="right", min_width=6)
+
+    block.add_row(
+        Text(label, style="bold"),
+        Text(f"{used:,} / {cap:,}", style="dim"),
+        Text(f"{pct:.0f}%", style=color),
+    )
+    bar = ProgressBar(
+        total=max(cap, 1),
+        completed=min(used, cap),
+        width=None,
+        style="bar.back",
+        complete_style=style,
+        finished_style=style,
+    )
+    # bar spans all columns
+    block.add_row(bar, "", "")
+    return block
+
+
+def _render_metrics(metrics: dict, queue=None):
+    """Build a rich Panel showing queue metrics with slim progress bars."""
+    from datetime import datetime, timezone  # noqa: PLC0415
+
+    from rich.panel import Panel  # noqa: PLC0415
+    from rich.table import Table  # noqa: PLC0415
+
+    outer = Table.grid(padding=(1, 2))
+    outer.add_column(ratio=1)
+
+    outer.add_row(_render_metric_block("Runs in-flight", metrics["runs_in_flight"], metrics["run_concurrency"]))
+    outer.add_row(
+        _render_metric_block("Actions in-flight", metrics["actions_in_flight"], metrics["action_concurrency"])
+    )
+    outer.add_row(_render_metric_block("Queue depth", metrics["current_depth"], metrics["depth"]))
+
+    ts = datetime.fromtimestamp(metrics["observed_at"], tz=timezone.utc).strftime("%H:%M:%S UTC")
+    subtitle_parts = [f"last updated {ts}"]
+
+    status_str = ""
+    if queue:
+        status_color = {"active": "green", "draining": "yellow", "drained": "dim"}.get(queue.status, "white")
+        status_str = f"  [{status_color}]{queue.status}[/{status_color}]"
+        subtitle_parts.append(f"priority: {queue.priority}")
+        subtitle_parts.append(f"fairness: {queue.fairness}")
+        if queue.clusters:
+            subtitle_parts.append(f"clusters: {', '.join(queue.clusters)}")
+
+    title = f"[bold]{metrics['name']}[/bold]{status_str}"
+    subtitle = "  |  ".join(subtitle_parts)
+    return Panel(outer, title=title, subtitle=f"[dim]{subtitle}[/dim]", border_style="dim")
+
+
+def _watch_queue(console, name: str, org: str, domain: str, project: str) -> None:
+    """Stream live queue metrics using rich Live display."""
+    from rich.console import Console  # noqa: PLC0415
+    from rich.live import Live  # noqa: PLC0415
+
+    from flyteplugins.union.remote import Queue  # noqa: PLC0415
+
+    queue = Queue.get(name, org=org, domain=domain, project=project)
+    console.print(f"Watching queue [cyan]{name}[/cyan]  (Ctrl+C to stop)\n")
+
+    # Use a console that auto-detects the real terminal width so that
+    # Live's cursor-up count matches the actual number of physical lines.
+    live_console = Console(force_terminal=True)
+    try:
+        with Live(console=live_console, refresh_per_second=4) as live:
+            for metrics in Queue.watch(name, org=org, domain=domain, project=project):
+                # Re-fetch queue config occasionally to pick up status changes
+                try:
+                    queue = Queue.get(name, org=org, domain=domain, project=project)
+                except Exception:
+                    pass
+                live.update(_render_metrics(metrics, queue))
+    except KeyboardInterrupt:
+        console.print("\n[dim]Watch stopped.[/dim]")
+
+
+# ---------------------------------------------------------------------------
+# YAML helpers for interactive editing
+# ---------------------------------------------------------------------------
+
+
+def _queue_to_yaml(queue) -> str:
+    """Serialize a Queue to YAML for interactive editing."""
+    import yaml as _yaml  # noqa: PLC0415
+
+    data = {
+        "name": queue.name,
+        "run_concurrency": queue.run_concurrency,
+        "action_concurrency": queue.action_concurrency,
+        "depth": queue.depth,
+        "priority": queue.priority,
+        "fairness": queue.fairness,
+        "clusters": queue.clusters,
+    }
+    if queue.project:
+        data["project"] = queue.project
+    if queue.domain:
+        data["domain"] = queue.domain
+    comment = (
+        "# Edit queue configuration below.\n"
+        "# Priority: min, medium, max\n"
+        "# Fairness: round_robin, shuffle_interleave\n"
+        "# Clusters: list of cluster names, or ['*'] for all\n"
+    )
+    return comment + "\n" + _yaml.safe_dump(data, default_flow_style=False, sort_keys=False)
+
+
+# ---------------------------------------------------------------------------
+# update queue
+# ---------------------------------------------------------------------------
+
+
+@click.command("queue")
+@click.argument("name", type=str)
+@click.option("--project", type=str, default="", help="Scope to a project")
+@click.option("--domain", type=str, default="", help="Scope to a domain")
+@click.option("--drain", is_flag=True, help="Begin draining the queue")
+@click.option("--activate", is_flag=True, help="Re-activate a draining or drained queue")
+@click.option("--edit", "edit_mode", is_flag=True, help="Open an editor to modify queue settings")
+@click.pass_obj
+def update_queue(cfg, name: str, project: str, domain: str, drain: bool, activate: bool, edit_mode: bool):
+    """Update a queue.
+
+    Use --drain to begin draining (stops new submissions).
+    Use --activate to re-activate a draining or drained queue.
+    Use --edit to interactively modify queue configuration.
+
+    Examples:
+
+        $ flyte update queue my-queue --drain
+
+        $ flyte update queue my-queue --activate
+
+        $ flyte update queue my-queue --edit
+    """
+    if drain and activate:
+        raise click.ClickException("Cannot specify both --drain and --activate")
+    if not drain and not activate and not edit_mode:
+        raise click.ClickException("Must specify --drain, --activate, or --edit")
+
+    import yaml as _yaml  # noqa: PLC0415, I001
+    from flyte.cli import _common as common  # noqa: PLC0415
+    from flyteplugins.union.cli import edit_with_retry  # noqa: PLC0415
+    from flyteplugins.union.remote import Queue  # noqa: PLC0415
+
+    cfg.init(project="", domain="")
+    console = common.get_console()
+
+    try:
+        if drain:
+            queue = Queue.drain(name, org=cfg.org, domain=domain, project=project)
+            console.print(f"[green]✓[/green] Queue [cyan]{name}[/cyan] is now draining")
+            console.print(common.format(f"Queue {name}", [queue], "json"))
+        elif activate:
+            queue = Queue.activate(name, org=cfg.org, domain=domain, project=project)
+            console.print(f"[green]✓[/green] Queue [cyan]{name}[/cyan] is now active")
+            console.print(common.format(f"Queue {name}", [queue], "json"))
+        else:
+            # Interactive edit
+            current = Queue.get(name, org=cfg.org, domain=domain, project=project)
+            yaml_text = _queue_to_yaml(current)
+
+            def _apply(edited_yaml):
+                data = _yaml.safe_load(edited_yaml)
+                clusters = data.get("clusters")
+                return Queue.update(
+                    name,
+                    org=cfg.org,
+                    domain=domain,
+                    project=project,
+                    run_concurrency=data.get("run_concurrency"),
+                    action_concurrency=data.get("action_concurrency"),
+                    depth=data.get("depth"),
+                    priority=data.get("priority"),
+                    fairness=data.get("fairness"),
+                    clusters=clusters,
+                )
+
+            result = edit_with_retry(yaml_text, _apply, console=console, noun="queue")
+            if result and result != "no_changes":
+                console.print(f"[green]✓[/green] Successfully updated queue: [cyan]{name}[/cyan]")
+    except Exception as e:
+        raise click.ClickException(f"Unable to update queue: {e}") from e