flyteplugins-union 0.4.2__tar.gz → 0.4.3__tar.gz

{flyteplugins_union-0.4.2/src/flyteplugins_union.egg-info → flyteplugins_union-0.4.3}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: flyteplugins-union
-Version: 0.4.2
+Version: 0.4.3
 Summary: Union SDK - Proprietary extensions for Flyte
 Author-email: unionai <info@union.ai>
 Project-URL: Homepage, https://www.union.ai/
@@ -19,7 +19,7 @@ Description-Content-Type: text/markdown
 License-File: LICENSE
 Requires-Dist: aiosqlite>=0.20.0
 Requires-Dist: connectrpc>=0.9.0
-Requires-Dist: flyte<2.5.0,>=2.3.0
+Requires-Dist: flyte<2.6.0,>=2.5.0
 Requires-Dist: flyteidl2>=2.0.0
 Requires-Dist: kubernetes>=27.2.0
 Requires-Dist: pyyaml>=6.0

{flyteplugins_union-0.4.2 → flyteplugins_union-0.4.3}/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-06-16
+Change Date:          2030-06-23
 
 Change License:       Apache License, Version 2.0
 

{flyteplugins_union-0.4.2 → flyteplugins_union-0.4.3}/pyproject.toml

@@ -12,7 +12,7 @@ requires-python = ">=3.10"
 dependencies = [
     "aiosqlite>=0.20.0",
     "connectrpc>=0.9.0",
-    "flyte>=2.3.0,<2.5.0",
+    "flyte>=2.5.0,<2.6.0",
     "flyteidl2>=2.0.0",
     "kubernetes>=27.2.0",
     "pyyaml>=6.0",
@@ -74,6 +74,10 @@ exclude = ["src/flyteplugins/union/internal"]
 Homepage = "https://www.union.ai/"
 Documentation = "https://docs.union.ai/"
 
+[project.scripts]
+# Stdio<->WebSocket proxy used as the ssh ProxyCommand for SSH-into-task debug.
+flyte-ws-proxy = "flyteplugins.union.ws_proxy:main"
+
 [project.entry-points."flyte.plugins.cli.commands"]
 # CLI plugins for Union commands
 "create.api-key" = "flyteplugins.union.cli.api_key:create_api_key"
@@ -109,6 +113,9 @@ Documentation = "https://docs.union.ai/"
 # than an `explore.volume` entry because the plugin loader won't auto-create a
 # missing `explore` verb group (it warns + skips).
 "explore" = "flyteplugins.union.cli.volume:explore"
+# Top-level `flyte connect` group (a single click.Group with an `ssh`
+# subcommand → `flyte connect ssh`). Same top-level-group rationale as `explore`.
+"connect" = "flyteplugins.union.cli.ssh:connect"
 
 [project.entry-points."flyte.plugins.types"]
 # Custom TypeTransformer that auto-commits a still-mounted Volume at
@@ -177,3 +184,7 @@ select = [
 # VolumeExplore lazy-imports flyte.io / msgpack / boto3 / the Volume model so
 # `import flyteplugins.union.remote` stays light for non-volume callers.
 "src/flyteplugins/union/remote/_volume_explore.py" = ["PLC0415"]
+# ssh-debug client lazy-imports the flyte auth machinery / remote objects so
+# `flyte connect --help` and `import flyteplugins.union.remote` stay light.
+"src/flyteplugins/union/cli/ssh.py" = ["PLC0415"]
+"src/flyteplugins/union/remote/_ssh_debug.py" = ["PLC0415"]

flyteplugins_union-0.4.3/src/flyteplugins/union/cli/ssh.py

@@ -0,0 +1,181 @@
+"""``flyte connect ssh`` — attach your local desktop to a running debug task.
+
+Counterpart to the pod-side ssh + WebSocket bridge (flyte core ``_F_E_SSH``).
+Given a run launched in ssh-debug mode, this resolves the per-pod ``:6060``
+route, mints (and caches) a Bearer token for the OAuth2 ingress, writes the
+headers file, and prints a ready-to-paste ``~/.ssh/config`` block (optionally
+writing it).
+
+Registered as a top-level ``connect`` group (like ``explore``) rather than a
+``connect.ssh`` entry, because the plugin loader won't auto-create a missing
+``connect`` verb group.
+"""
+
+from __future__ import annotations
+
+from typing import Optional
+
+import rich_click as click
+
+# Heavy imports (rich, remote objects) are deferred into the command body so
+# registering the entry point stays fast.
+
+
+@click.group("connect")
+def connect() -> None:
+    """Connect your local machine to a running Flyte task."""
+
+
+@connect.command("ssh")
+@click.argument("run_name", required=True)
+@click.argument("action_name", required=False)
+@click.option("--user", type=str, default="root", show_default=True, help="SSH login user inside the pod.")
+@click.option(
+    "--identity-file",
+    "identity_file",
+    type=str,
+    default=None,
+    help="Private key to authenticate with. Defaults to the auto-managed "
+    "~/.flyte/ssh-debug/id_ed25519 (created for you; no ssh-keygen needed).",
+)
+@click.option(
+    "--host-alias",
+    "--name",
+    "host_alias",
+    type=str,
+    default="flyte-debug",
+    show_default=True,
+    help="Host name to use in the ~/.ssh/config block (use distinct names to keep several runs side by side).",
+)
+@click.option(
+    "--api-key/--no-api-key",
+    "api_key",
+    default=False,
+    help="Authenticate the tunnel with a dedicated, long-lived API key (created/reused "
+    "as `flyte-ssh-debug`) instead of your interactive session token. Survives re-logins "
+    "and won't expire mid-session.",
+)
+@click.option(
+    "--refresh-token",
+    "refresh_token",
+    is_flag=True,
+    help="Force a fresh Bearer instead of reusing the cached one (use if you hit auth errors).",
+)
+@click.option(
+    "--write-config",
+    "write_config",
+    is_flag=True,
+    help="Write the Host block into ~/.ssh/config (replacing any prior block for the same name).",
+)
+@click.option(
+    "--timeout",
+    type=float,
+    default=300.0,
+    show_default=True,
+    help="Seconds to wait for the debug route to become ready.",
+)
+@click.pass_obj
+def connect_ssh(
+    cfg,
+    run_name: str,
+    action_name: Optional[str],
+    user: str,
+    identity_file: str,
+    host_alias: str,
+    api_key: bool,
+    refresh_token: bool,
+    write_config: bool,
+    timeout: float,
+) -> None:
+    """Attach to a running ssh-debug task (launched with ``_F_E_SSH=1``).
+
+    Only RUN_NAME is required; ACTION_NAME defaults to the root action ``a0``.
+    The Bearer token is reused from a local cache between calls (no re-minting
+    every time); pass ``--refresh-token`` to force a new one.
+
+    Examples:
+
+        # Resolve + print the ssh-config for a debug run's root action
+        $ flyte connect ssh my-run
+
+        # A named host + login user, written into ~/.ssh/config
+        $ flyte connect ssh my-run --name my-run-dbg --user flyte --write-config
+
+        # Use a long-lived, user-specific API key for the tunnel auth
+        $ flyte connect ssh my-run --api-key --write-config
+
+        # Then connect (or VS Code -> Remote-SSH -> <name>)
+        $ ssh my-run-dbg
+    """
+    from flyte.cli import _common as common
+
+    from flyteplugins.union.remote import SSHDebug
+
+    action_name = action_name or "a0"
+    cfg.init()
+
+    console = common.get_console()
+    with console.status(f"[bold]Resolving[/] ssh route for {run_name}/{action_name}…", spinner="dots"):
+        try:
+            info = SSHDebug.connect(
+                run_name,
+                action_name,
+                user=user,
+                identity_file=identity_file,
+                host_alias=host_alias,
+                use_api_key=api_key,
+                refresh_token=refresh_token,
+                timeout=timeout,
+            )
+        except TimeoutError as e:
+            raise click.ClickException(str(e)) from e
+
+    _emit(console, info, run_name=run_name, action_name=action_name, write_config=write_config)
+
+
+def _emit(console, info, *, run_name, action_name, write_config):
+    """Write the Host block (optionally) and print the connect panel. Shared by
+    `flyte connect ssh` and `flyte debug`."""
+    import os
+    from pathlib import Path
+
+    from rich.panel import Panel
+    from rich.syntax import Syntax
+
+    if write_config:
+        ssh_config_path = Path(os.path.expanduser("~/.ssh/config"))
+        ssh_config_path.parent.mkdir(parents=True, exist_ok=True)
+        existing = ssh_config_path.read_text() if ssh_config_path.exists() else ""
+        # Drop any prior `Host <alias>` block (header + its indented body) so
+        # re-running replaces stale config instead of shadowing it with a dup
+        # (ssh honors the first matching Host).
+        kept, skipping = [], False
+        for line in existing.splitlines():
+            if line.strip().startswith("Host ") and line.split()[1:] == [info.host_alias]:
+                skipping = True
+                continue
+            if skipping:
+                # Block body is blank or indented; ends at the next non-indented, non-blank line.
+                if line and not line[0].isspace():
+                    skipping = False
+                else:
+                    continue
+            kept.append(line)
+        new_text = "\n".join(kept).rstrip() + "\n\n" + info.ssh_config
+        ssh_config_path.write_text(new_text)
+        console.print(f"[green]Wrote[/] Host [bold]{info.host_alias}[/] to {ssh_config_path}")
+
+    console.print(
+        Panel.fit(
+            Syntax(info.ssh_config, "sshconfig", theme="ansi_dark", background_color="default"),
+            title=f"[bold]ssh -> {run_name}/{action_name}[/bold]",
+            subtitle=f"then: ssh {info.host_alias}   (or VS Code Remote-SSH -> {info.host_alias})",
+            border_style="cyan",
+        )
+    )
+    if info.headers_file:
+        console.print(f"[dim]Bearer header written to {info.headers_file} (0600)[/dim]")
+    console.print(
+        f"[dim]Token expired / 401 on reconnect? Just re-run `flyte connect ssh {run_name}` "
+        f"(refreshes in place — no --write-config needed). For multi-day sessions use --api-key.[/dim]"
+    )

flyteplugins_union-0.4.3/src/flyteplugins/union/debug.py

@@ -0,0 +1,39 @@
+"""SSH-into-task debug helpers.
+
+`with_debugcontext(...)` is `flyte.with_runcontext(...)` preconfigured for
+ssh-into-task debug: it creates/uses the auto-managed debug keypair and sets the
+``_F_E_SSH`` / ``_F_SSH_PK`` / ``_F_E_VS`` env vars, then forwards everything
+else to `flyte.with_runcontext`. After the run starts, resolve the connection
+with `SSHDebug.connect(run.name)` (or the `flyte connect ssh <run>` CLI).
+
+Example:
+```python
+import flyte
+from flyteplugins.union.debug import with_debugcontext, SSHDebug
+
+flyte.init_from_config()
+run = with_debugcontext().run(my_task, x=1)
+print(SSHDebug.connect(run.name).ssh_config)   # paste into ~/.ssh/config, then `ssh flyte-debug`
+```
+"""
+
+from __future__ import annotations
+
+from typing import Any, Dict, Optional
+
+import flyte
+
+from flyteplugins.union.remote._ssh_debug import SSHDebug, launch_env
+
+__all__ = ["SSHDebug", "with_debugcontext"]
+
+
+def with_debugcontext(mode: Any = None, *, env_vars: Optional[Dict[str, str]] = None, **kwargs):
+    """Like `flyte.with_runcontext`, but preconfigured for ssh-into-task debug.
+
+    Ensures the auto-managed debug keypair exists and merges the ssh-debug env
+    (``_F_E_SSH`` / ``_F_SSH_PK`` / ``_F_E_VS``) into *env_vars*; all other
+    arguments are forwarded unchanged. Returns the same runner as
+    `with_runcontext`, so call ``.run(task, ...)`` on it.
+    """
+    return flyte.with_runcontext(mode, env_vars=launch_env(env_vars), **kwargs)

{flyteplugins_union-0.4.2 → flyteplugins_union-0.4.3}/src/flyteplugins/union/errors.py

@@ -53,6 +53,19 @@ class VolumeUnmountError(VolumeError):
     _code = "VolumeUnmountFailed"
 
 
+class VolumeUnmountUnsupportedError(VolumeUnmountError):
+    """The runtime cannot unmount the FUSE client at all because ``fusermount``
+    can't reach its mount table (``/etc/mtab`` is missing, read-only, or
+    unwritable). Unlike a transient :class:`VolumeUnmountError` (EBUSY), retrying
+    won't help — the image is missing the ``/etc/mtab`` → ``/proc/self/mounts``
+    symlink most base images carry. The terminal seal treats this as recoverable:
+    it falls back to a keep-alive (live) commit, leaving the mount in place. A
+    kind of :class:`VolumeUnmountError`, so ``except VolumeUnmountError`` catches
+    it too."""
+
+    _code = "VolumeUnmountUnsupported"
+
+
 class VolumeCommitError(VolumeError):
     """A keep-alive checkpoint could not be made durable: the writeback staging
     queue did not finish uploading to object storage (it timed out, or its