fn_registry/python/functions/notebook/jupyter_kernel.py

"""CRUD de kernels Jupyter via REST API."""

import json
import urllib.error
import urllib.request


def _make_request(
    method: str,
    url: str,
    token: str = "",
    body: dict | None = None,
) -> dict | list | None:
    """Ejecuta una request HTTP a la API de Jupyter.

    Args:
        method: Metodo HTTP (GET, POST, DELETE).
        url: URL completa del endpoint.
        token: Token de autenticacion de Jupyter. Vacio si no se requiere.
        body: Cuerpo de la request para metodos POST.

    Returns:
        Respuesta deserializada como dict o list, o None si la respuesta esta vacia.

    Raises:
        urllib.error.HTTPError: Si la respuesta HTTP indica un error.
        urllib.error.URLError: Si no se puede conectar al servidor.
    """
    data = json.dumps(body).encode("utf-8") if body is not None else None

    headers = {
        "Accept": "application/json",
        "Content-Type": "application/json",
    }
    if token:
        headers["Authorization"] = f"token {token}"

    req = urllib.request.Request(url, data=data, headers=headers, method=method)

    with urllib.request.urlopen(req) as resp:
        raw = resp.read()
        if not raw:
            return None
        return json.loads(raw.decode("utf-8"))


def jupyter_kernel_list(
    server_url: str = "http://localhost:8888",
    token: str = "",
) -> list[dict]:
    """Lista todos los kernels activos en el servidor Jupyter.

    Args:
        server_url: URL base del servidor Jupyter.
        token: Token de autenticacion. Vacio si el servidor no requiere auth.

    Returns:
        Lista de dicts con los campos: id, name, execution_state,
        last_activity, connections.

    Raises:
        urllib.error.HTTPError: Si la respuesta HTTP indica un error.
        urllib.error.URLError: Si no se puede conectar al servidor.
    """
    url = f"{server_url.rstrip('/')}/api/kernels"
    result = _make_request("GET", url, token)
    return result if isinstance(result, list) else []


def jupyter_kernel_start(
    server_url: str = "http://localhost:8888",
    token: str = "",
    name: str = "python3",
) -> dict:
    """Inicia un kernel nuevo en el servidor Jupyter.

    Args:
        server_url: URL base del servidor Jupyter.
        token: Token de autenticacion. Vacio si el servidor no requiere auth.
        name: Nombre del kernel a iniciar (p.ej. "python3", "ir").

    Returns:
        Dict con los campos: id, name, execution_state.

    Raises:
        urllib.error.HTTPError: Si la respuesta HTTP indica un error.
        urllib.error.URLError: Si no se puede conectar al servidor.
    """
    url = f"{server_url.rstrip('/')}/api/kernels"
    result = _make_request("POST", url, token, body={"name": name})
    return result if isinstance(result, dict) else {}


def jupyter_kernel_restart(
    server_url: str = "http://localhost:8888",
    token: str = "",
    kernel_id: str = "",
) -> dict:
    """Reinicia un kernel existente.

    Args:
        server_url: URL base del servidor Jupyter.
        token: Token de autenticacion. Vacio si el servidor no requiere auth.
        kernel_id: ID del kernel a reiniciar.

    Returns:
        Dict con la informacion actualizada del kernel.

    Raises:
        urllib.error.HTTPError: Si la respuesta HTTP indica un error (p.ej. 404 si no existe).
        urllib.error.URLError: Si no se puede conectar al servidor.
    """
    url = f"{server_url.rstrip('/')}/api/kernels/{kernel_id}/restart"
    result = _make_request("POST", url, token)
    return result if isinstance(result, dict) else {}


def jupyter_kernel_interrupt(
    server_url: str = "http://localhost:8888",
    token: str = "",
    kernel_id: str = "",
) -> None:
    """Interrumpe la ejecucion actual de un kernel.

    Args:
        server_url: URL base del servidor Jupyter.
        token: Token de autenticacion. Vacio si el servidor no requiere auth.
        kernel_id: ID del kernel a interrumpir.

    Raises:
        urllib.error.HTTPError: Si la respuesta HTTP indica un error (p.ej. 404 si no existe).
        urllib.error.URLError: Si no se puede conectar al servidor.
    """
    url = f"{server_url.rstrip('/')}/api/kernels/{kernel_id}/interrupt"
    _make_request("POST", url, token)


def jupyter_kernel_shutdown(
    server_url: str = "http://localhost:8888",
    token: str = "",
    kernel_id: str = "",
) -> None:
    """Apaga y elimina un kernel.

    Args:
        server_url: URL base del servidor Jupyter.
        token: Token de autenticacion. Vacio si el servidor no requiere auth.
        kernel_id: ID del kernel a apagar.

    Raises:
        urllib.error.HTTPError: Si la respuesta HTTP indica un error (p.ej. 404 si no existe).
        urllib.error.URLError: Si no se puede conectar al servidor.
    """
    url = f"{server_url.rstrip('/')}/api/kernels/{kernel_id}"
    _make_request("DELETE", url, token)


def jupyter_kernel_sessions(
    server_url: str = "http://localhost:8888",
    token: str = "",
) -> list[dict]:
    """Lista las sesiones activas del servidor Jupyter.

    Cada sesion mapea un notebook a su kernel y usuario actuales.

    Args:
        server_url: URL base del servidor Jupyter.
        token: Token de autenticacion. Vacio si el servidor no requiere auth.

    Returns:
        Lista de dicts con los campos: id, notebook (path), kernel_id,
        kernel_state, type, name.

    Raises:
        urllib.error.HTTPError: Si la respuesta HTTP indica un error.
        urllib.error.URLError: Si no se puede conectar al servidor.
    """
    url = f"{server_url.rstrip('/')}/api/sessions"
    raw = _make_request("GET", url, token)
    if not isinstance(raw, list):
        return []

    sessions = []
    for s in raw:
        kernel = s.get("kernel") or {}
        sessions.append(
            {
                "id": s.get("id", ""),
                "notebook": s.get("path", ""),
                "kernel_id": kernel.get("id", ""),
                "kernel_state": kernel.get("execution_state", ""),
                "type": s.get("type", ""),
                "name": s.get("name", ""),
            }
        )
    return sessions


def jupyter_kernel_cleanup(
    server_url: str = "http://localhost:8888",
    token: str = "",
    idle_seconds: int = 3600,
) -> list[dict]:
    """Apaga todos los kernels que llevan mas de idle_seconds sin actividad.

    Util para liberar recursos en servidores con muchos notebooks abiertos.
    Por defecto cierra kernels inactivos desde hace mas de 1 hora.

    Args:
        server_url: URL base del servidor Jupyter.
        token: Token de autenticacion. Vacio si el servidor no requiere auth.
        idle_seconds: Segundos de inactividad para considerar un kernel ocioso.

    Returns:
        Lista de dicts con los kernels apagados (id, name, last_activity, idle_seconds).

    Raises:
        urllib.error.HTTPError: Si la respuesta HTTP indica un error.
        urllib.error.URLError: Si no se puede conectar al servidor.
    """
    from datetime import datetime, timezone

    kernels = jupyter_kernel_list(server_url, token)
    now = datetime.now(timezone.utc)
    shutdown_list = []

    for k in kernels:
        last_activity = k.get("last_activity", "")
        if not last_activity:
            continue
        try:
            last_dt = datetime.fromisoformat(last_activity.replace("Z", "+00:00"))
            idle = (now - last_dt).total_seconds()
        except (ValueError, TypeError):
            continue
        if idle >= idle_seconds:
            jupyter_kernel_shutdown(server_url, token, k["id"])
            shutdown_list.append({
                "id": k["id"],
                "name": k.get("name", ""),
                "last_activity": last_activity,
                "idle_seconds": int(idle),
            })

    return shutdown_list


def jupyter_kernel_shutdown_all(
    server_url: str = "http://localhost:8888",
    token: str = "",
) -> list[dict]:
    """Apaga todos los kernels activos del servidor.

    Args:
        server_url: URL base del servidor Jupyter.
        token: Token de autenticacion. Vacio si el servidor no requiere auth.

    Returns:
        Lista de dicts con los kernels apagados (id, name).

    Raises:
        urllib.error.HTTPError: Si la respuesta HTTP indica un error.
        urllib.error.URLError: Si no se puede conectar al servidor.
    """
    kernels = jupyter_kernel_list(server_url, token)
    shutdown_list = []
    for k in kernels:
        jupyter_kernel_shutdown(server_url, token, k["id"])
        shutdown_list.append({"id": k["id"], "name": k.get("name", "")})
    return shutdown_list


# ---------------------------------------------------------------------------
# CLI
# ---------------------------------------------------------------------------

if __name__ == "__main__":
    import argparse
    import sys

    parser = argparse.ArgumentParser(
        description="CRUD de kernels Jupyter via REST API."
    )
    parser.add_argument(
        "--server",
        default="http://localhost:8888",
        help="URL base del servidor Jupyter (default: http://localhost:8888)",
    )
    parser.add_argument(
        "--token",
        default="",
        help="Token de autenticacion de Jupyter (default: vacio)",
    )

    subparsers = parser.add_subparsers(dest="command", required=True)

    # list
    subparsers.add_parser("list", help="Lista todos los kernels activos.")

    # start
    sp_start = subparsers.add_parser("start", help="Inicia un kernel nuevo.")
    sp_start.add_argument(
        "--name",
        default="python3",
        help="Nombre del kernel (default: python3)",
    )

    # restart
    sp_restart = subparsers.add_parser("restart", help="Reinicia un kernel existente.")
    sp_restart.add_argument("kernel_id", help="ID del kernel a reiniciar.")

    # interrupt
    sp_interrupt = subparsers.add_parser(
        "interrupt", help="Interrumpe la ejecucion de un kernel."
    )
    sp_interrupt.add_argument("kernel_id", help="ID del kernel a interrumpir.")

    # shutdown
    sp_shutdown = subparsers.add_parser("shutdown", help="Apaga y elimina un kernel.")
    sp_shutdown.add_argument("kernel_id", help="ID del kernel a apagar.")

    # sessions
    subparsers.add_parser("sessions", help="Lista las sesiones activas.")

    # cleanup
    sp_cleanup = subparsers.add_parser("cleanup", help="Apaga kernels inactivos.")
    sp_cleanup.add_argument(
        "--idle-seconds",
        type=int,
        default=3600,
        help="Segundos de inactividad para considerar ocioso (default: 3600)",
    )

    # shutdown-all
    subparsers.add_parser("shutdown-all", help="Apaga todos los kernels activos.")

    args = parser.parse_args()

    try:
        if args.command == "list":
            result = jupyter_kernel_list(args.server, args.token)
        elif args.command == "start":
            result = jupyter_kernel_start(args.server, args.token, args.name)
        elif args.command == "restart":
            result = jupyter_kernel_restart(args.server, args.token, args.kernel_id)
            if result is None:
                result = {"status": "restarted", "kernel_id": args.kernel_id}
        elif args.command == "interrupt":
            jupyter_kernel_interrupt(args.server, args.token, args.kernel_id)
            result = {"status": "interrupted", "kernel_id": args.kernel_id}
        elif args.command == "shutdown":
            jupyter_kernel_shutdown(args.server, args.token, args.kernel_id)
            result = {"status": "shutdown", "kernel_id": args.kernel_id}
        elif args.command == "sessions":
            result = jupyter_kernel_sessions(args.server, args.token)
        elif args.command == "cleanup":
            result = jupyter_kernel_cleanup(args.server, args.token, args.idle_seconds)
        elif args.command == "shutdown-all":
            result = jupyter_kernel_shutdown_all(args.server, args.token)
        else:
            parser.print_help()
            sys.exit(1)

        print(json.dumps(result, indent=2))
    except urllib.error.HTTPError as e:
        body = e.read().decode("utf-8", errors="replace")
        print(
            json.dumps({"error": f"HTTP {e.code}: {e.reason}", "detail": body}),
            file=sys.stderr,
        )
        sys.exit(1)
    except urllib.error.URLError as e:
        print(
            json.dumps({"error": f"URLError: {e.reason}"}),
            file=sys.stderr,
        )
        sys.exit(1)