chalk-remote-call-python 1.10.5__cp315-cp315-win_amd64.whl

chalk_remote_call/cli.py

@@ -0,0 +1,160 @@
+from __future__ import annotations
+
+import argparse
+import logging
+import os
+import sys
+
+from chalk_remote_call.handler_loader import load_function, load_handler
+from chalk_remote_call.logging import JsonFormatter
+from chalk_remote_call.server import serve
+
+
+def _env_int(name: str) -> int | None:
+    """Read an optional int from the environment; return None if unset/empty."""
+    raw = os.environ.get(name)
+    if raw is None or raw.strip() == "":
+        return None
+    try:
+        return int(raw)
+    except ValueError as e:
+        raise ValueError(f"{name} must be an integer, got: {raw!r}") from e
+
+
+def main() -> None:
+    log_handler = logging.StreamHandler(sys.stderr)
+    log_handler.setFormatter(JsonFormatter())
+    logging.basicConfig(level=logging.INFO, handlers=[log_handler], force=True)
+    try:
+        _run()
+    except Exception:
+        logging.getLogger(__name__).exception("Remote-call server failed")
+        sys.exit(1)
+
+
+def _run() -> None:
+    parser = argparse.ArgumentParser(
+        prog="chalk-remote-call",
+        description="Start a gRPC server implementing the Chalk RemoteCallService.",
+    )
+    parser.add_argument(
+        "--handler",
+        required=True,
+        help="Dotted path to the handler function (e.g. my_module.handler)",
+    )
+    parser.add_argument(
+        "--port",
+        type=int,
+        default=int(os.environ.get("CHALK_REMOTE_CALL_PORT", "6666")),
+        help="Port to listen on (default: 6666, env: CHALK_REMOTE_CALL_PORT)",
+    )
+    parser.add_argument(
+        "--host",
+        default=os.environ.get("CHALK_REMOTE_CALL_HOST", "[::]"),
+        help="Host to bind to (default: [::], env: CHALK_REMOTE_CALL_HOST)",
+    )
+    parser.add_argument(
+        "--workers",
+        type=int,
+        default=int(os.environ.get("CHALK_REMOTE_CALL_WORKERS", "10")),
+        help="Number of worker threads (default: 10, env: CHALK_REMOTE_CALL_WORKERS)",
+    )
+    parser.add_argument(
+        "--on-startup",
+        default=None,
+        help="Dotted path to a startup function (e.g. my_module.setup)",
+    )
+    parser.add_argument(
+        "--on-shutdown",
+        default=None,
+        help="Dotted path to a shutdown function (e.g. my_module.cleanup)",
+    )
+    parser.add_argument(
+        "--max-batching-size",
+        type=int,
+        default=_env_int("CHALK_REMOTE_CALL_MAX_BATCHING_SIZE"),
+        help=(
+            "Max concurrent requests to coalesce into a single handler "
+            "invocation. Unset disables batching. "
+            "(env: CHALK_REMOTE_CALL_MAX_BATCHING_SIZE)"
+        ),
+    )
+    parser.add_argument(
+        "--max-buffer-duration-ms",
+        type=int,
+        default=_env_int("CHALK_REMOTE_CALL_MAX_BUFFER_DURATION_MS"),
+        help=(
+            "Max milliseconds to buffer concurrent requests before flushing, "
+            "even if max-batching-size hasn't been reached. Only effective "
+            "when --max-batching-size is also set. Defaults to 1000ms when "
+            "batching is enabled. "
+            "(env: CHALK_REMOTE_CALL_MAX_BUFFER_DURATION_MS)"
+        ),
+    )
+    parser.add_argument(
+        "--batching-mode",
+        choices=["per_caller", "combined"],
+        default=os.environ.get("CHALK_REMOTE_CALL_BATCHING_MODE") or None,
+        help=(
+            "Coalescing variant when --max-batching-size is set. "
+            "'per_caller' (default) passes events as a list of dicts; "
+            "'combined' passes a single concatenated RecordBatch with "
+            "offsets. (env: CHALK_REMOTE_CALL_BATCHING_MODE)"
+        ),
+    )
+    parser.add_argument(
+        "--log-level",
+        default="INFO",
+        choices=["DEBUG", "INFO", "WARNING", "ERROR"],
+        help="Log level (default: INFO)",
+    )
+    args = parser.parse_args()
+
+    logging.getLogger().setLevel(args.log_level)
+
+    from chalk_remote_call.input_transform import parse_input_args
+
+    # Load handler
+    try:
+        handler, auto_startup, auto_shutdown = load_handler(args.handler)
+    except Exception:
+        logging.getLogger(__name__).exception("Error loading handler")
+        sys.exit(1)
+
+    # Determine startup function
+    on_startup = None
+    if args.on_startup:
+        try:
+            on_startup = load_function(args.on_startup, label="Startup")
+        except Exception:
+            logging.getLogger(__name__).exception("Error loading startup function")
+            sys.exit(1)
+    elif auto_startup is not None:
+        on_startup = auto_startup
+
+    # Determine shutdown function
+    on_shutdown = None
+    if args.on_shutdown:
+        try:
+            on_shutdown = load_function(args.on_shutdown, label="Shutdown")
+        except Exception:
+            logging.getLogger(__name__).exception("Error loading shutdown function")
+            sys.exit(1)
+    elif auto_shutdown is not None:
+        on_shutdown = auto_shutdown
+
+    # Parse input args
+    arg_names = parse_input_args()
+
+    serve(
+        handler=handler,
+        host=args.host,
+        port=args.port,
+        workers=args.workers,
+        on_startup=on_startup,
+        on_shutdown=on_shutdown,
+        arg_names=arg_names,
+        max_batching_size=args.max_batching_size,
+        max_buffer_duration_ms=args.max_buffer_duration_ms,
+        batching_mode=args.batching_mode,
+    )

chalk_remote_call/errors.py

@@ -0,0 +1,118 @@
+"""Versioned, data-only remote exception transport shared by all call modes."""
+
+from __future__ import annotations
+
+import importlib
+import json
+import traceback
+from types import TracebackType
+
+
+def _user_stack(tb: TracebackType | None) -> traceback.StackSummary:
+    # Start at the outermost bridge's callee, so Python callers above the
+    # dispatch boundary cannot leak into user diagnostics or make async runner
+    # frames look like user code. Stop at the first bridge to preserve user
+    # frames if a handler itself invokes another bridge.
+    dispatch_tb = tb
+    while dispatch_tb is not None:
+        frame = dispatch_tb.tb_frame
+        if frame.f_globals.get("__name__") == "chalk_remote_call.servicer" and frame.f_code.co_name in {
+            "process_batches",
+            "process_batches_coalesced",
+            "process_batches_coalesced_combined",
+        }:
+            tb = dispatch_tb.tb_next
+            break
+        dispatch_tb = dispatch_tb.tb_next
+
+    frames = []
+    entered_user_code = False
+    while tb is not None:
+        frame = tb.tb_frame
+        module = str(frame.f_globals.get("__name__", ""))
+        runner_frame = not entered_user_code and module in {
+            "asyncio.base_events",
+            "asyncio.runners",
+            "concurrent.futures._base",
+            "concurrent.futures.thread",
+        }
+        if not (
+            runner_frame
+            or module == "chalk_remote_call.servicer"
+            or module == "chalk_remote_call.errors"
+            or module.startswith("chalkcompute.")
+            or frame.f_globals.get("__chalk_generated_handler__") is True
+        ):
+            entered_user_code = True
+            filename, lineno = frame.f_code.co_filename, tb.tb_lineno
+            source_map = frame.f_globals.get("__chalk_source_map__", {})
+            mapped = source_map.get(lineno) if isinstance(source_map, dict) else None
+            if isinstance(mapped, tuple) and len(mapped) == 3:
+                filename, lineno, source = mapped
+            else:
+                source = None
+            frames.append(traceback.FrameSummary(filename, lineno, frame.f_code.co_name, line=source))
+        tb = tb.tb_next
+    return traceback.StackSummary.from_list(frames)
+
+
+def _filter_traceback(summary: traceback.TracebackException, error: BaseException) -> None:
+    summary.stack = _user_stack(error.__traceback__)
+    if summary.__cause__ is not None and error.__cause__ is not None:
+        _filter_traceback(summary.__cause__, error.__cause__)
+    if summary.__context__ is not None and error.__context__ is not None:
+        _filter_traceback(summary.__context__, error.__context__)
+    children = getattr(error, "exceptions", ())
+    for child_summary, child in zip(getattr(summary, "exceptions", None) or (), children, strict=False):
+        _filter_traceback(child_summary, child)
+
+
+def serialize_exception(error: BaseException) -> str:
+    """Keep chains and groups, without pickling exceptions or their locals."""
+    summary = traceback.TracebackException.from_exception(error, capture_locals=False)
+    internal = "".join(summary.format())
+    _filter_traceback(summary, error)
+    payload: dict[str, object] = {
+        "chalk_remote_error": 1,
+        "exception": {
+            "kind": type(error).__name__[:256],
+            "message": _safe_message(error),
+            "stacktrace": "".join(summary.format()),
+            "internal_stacktrace": internal,
+        },
+    }
+    # Include a trace ID when serialization runs with an active trace context.
+    try:
+        trace = importlib.import_module("opentelemetry.trace")
+        context = trace.get_current_span().get_span_context()
+        if context.is_valid:
+            payload["trace_id"] = format(context.trace_id, "032x")
+    except ImportError:
+        pass
+    # gRPC status messages live in headers. Keep the complete diagnostic in
+    # server logs, and bound the wire payload below common 8 KiB header limits.
+    encoded = json.dumps(payload, ensure_ascii=True)
+    exception = payload["exception"]
+    assert isinstance(exception, dict)
+    while _status_header_size(encoded) > 6000:
+        field = max(exception, key=lambda key: _status_header_size(json.dumps(exception[key], ensure_ascii=True)))
+        text = exception[field]
+        keep = len(text) // 4
+        if keep < 20:
+            break
+        exception[field] = text[:keep] + "\n... remote traceback truncated ...\n" + text[-keep:]
+        encoded = json.dumps(payload, ensure_ascii=True)
+    return encoded
+
+
+def _status_header_size(message: str) -> int:
+    # Tonic percent-encodes these ASCII characters in grpc-message. The JSON
+    # serializer already escapes controls and non-ASCII code points.
+    return len(message) + 2 * sum(char in ' "#%<>`?{}' for char in message)
+
+
+def _safe_message(error: BaseException) -> str:
+    try:
+        return str(error)
+    except Exception:
+        return "<exception message unavailable>"

chalk_remote_call/handler_loader.py

@@ -0,0 +1,57 @@
+from __future__ import annotations
+
+import importlib
+from collections.abc import Callable
+from typing import Any
+
+
+def load_handler(
+    dotted_path: str,
+) -> tuple[Callable[..., Any], Callable[[], None] | None, Callable[[], None] | None]:
+    """Load a handler function from a dotted module path.
+
+    Returns (handler_fn, on_startup_fn_or_None, on_shutdown_fn_or_None).
+    The on_startup/on_shutdown functions are auto-discovered if the handler's
+    module defines top-level callables with those names.
+    """
+    if "." not in dotted_path:
+        raise ValueError(f"Handler path must be a dotted path like 'my_module.handler', got: {dotted_path!r}")
+
+    module_path, func_name = dotted_path.rsplit(".", 1)
+    module = importlib.import_module(module_path)
+    handler = getattr(module, func_name, None)
+    if handler is None:
+        raise AttributeError(f"Module {module_path!r} has no attribute {func_name!r}")
+    if not callable(handler):
+        raise TypeError(f"{dotted_path!r} is not callable")
+
+    on_startup = getattr(module, "on_startup", None)
+    if on_startup is not None and not callable(on_startup):
+        on_startup = None
+
+    on_shutdown = getattr(module, "on_shutdown", None)
+    if on_shutdown is not None and not callable(on_shutdown):
+        on_shutdown = None
+
+    return handler, on_startup, on_shutdown
+
+
+def load_function(dotted_path: str, label: str = "Function") -> Callable[[], None]:
+    """Load a callable from a dotted module path.
+
+    Args:
+        dotted_path: e.g. 'my_module.setup'
+        label: Human-readable label for error messages (e.g. "Startup", "Shutdown").
+    """
+    if "." not in dotted_path:
+        raise ValueError(f"{label} path must be a dotted path like 'my_module.func', got: {dotted_path!r}")
+
+    module_path, func_name = dotted_path.rsplit(".", 1)
+    module = importlib.import_module(module_path)
+    func = getattr(module, func_name, None)
+    if func is None:
+        raise AttributeError(f"Module {module_path!r} has no attribute {func_name!r}")
+    if not callable(func):
+        raise TypeError(f"{dotted_path!r} is not callable")
+
+    return func

chalk_remote_call/input_transform.py

@@ -0,0 +1,50 @@
+from __future__ import annotations
+
+import os
+from collections.abc import Sequence
+
+import pyarrow as pa
+
+# Columns the caller attaches as framework side channels rather than handler arguments.
+# They are identified by name and are optional per request, so they must be held out of the
+# positional CHALK_INPUT_ARGS mapping, which describes only the handler's declared arguments.
+RESERVED_COLUMNS = frozenset({"__chalk_row_metadata__"})
+
+
+def resolve_column_names(field_names: Sequence[str], arg_names: Sequence[str]) -> list[str]:
+    """Map `arg_names` positionally onto the non-reserved columns of `field_names`.
+
+    Reserved columns keep their own name and consume no entry from `arg_names`, so a request
+    that carries one maps its arguments exactly as a request without one does.
+    """
+    positional = [name for name in field_names if name not in RESERVED_COLUMNS]
+    if len(arg_names) != len(positional):
+        raise ValueError(
+            f"CHALK_INPUT_ARGS specifies {len(arg_names)} names but RecordBatch has {len(positional)} columns"
+        )
+    remaining = iter(arg_names)
+    return [name if name in RESERVED_COLUMNS else next(remaining) for name in field_names]
+
+
+def parse_input_args() -> list[str] | None:
+    """Parse CHALK_INPUT_ARGS env var into a list of column names.
+
+    Returns None if the env var is unset or empty.
+    """
+    raw = os.environ.get("CHALK_INPUT_ARGS", "").strip()
+    if not raw:
+        return None
+    return [name.strip() for name in raw.split(",") if name.strip()]
+
+
+def transform(batch: pa.RecordBatch, arg_names: list[str] | None) -> dict[str, pa.Array]:
+    """Transform a RecordBatch into a dict of named arrays.
+
+    If arg_names is provided, columns are renamed by index to the given names
+    (ignoring the original column names), except for reserved columns, which keep
+    their own name. If arg_names is None, the original column names are used.
+    """
+    if arg_names is not None:
+        names = resolve_column_names(batch.schema.names, arg_names)
+        return {name: batch.column(i) for i, name in enumerate(names)}
+    return {batch.schema.field(i).name: batch.column(i) for i in range(batch.num_columns)}

chalk_remote_call/logging.py

@@ -0,0 +1,29 @@
+from __future__ import annotations
+
+import json
+import logging
+from datetime import datetime, timezone
+from typing import Any
+
+_RECORD_FIELDS = frozenset(logging.makeLogRecord({}).__dict__) | {"message", "asctime"}
+
+
+class JsonFormatter(logging.Formatter):
+    """Emit one JSON object per line using fields recognized by Chalk's log collector."""
+
+    def format(self, record: logging.LogRecord) -> str:
+        fields: dict[str, Any] = {key: value for key, value in vars(record).items() if key not in _RECORD_FIELDS}
+        fields.update(
+            timestamp=datetime.fromtimestamp(record.created, tz=timezone.utc).isoformat(),
+            severity=record.levelname,
+            logger=record.name,
+            message=record.getMessage(),
+        )
+        # Keep tracebacks out of the message and escape their newlines so the collector sees one event.
+        if record.exc_info:
+            fields["exception.stacktrace"] = self.formatException(record.exc_info)
+        elif record.exc_text:
+            fields["exception.stacktrace"] = record.exc_text
+        if record.stack_info:
+            fields["stack_info"] = self.formatStack(record.stack_info)
+        return json.dumps(fields, default=str)

chalk_remote_call/server.py

@@ -0,0 +1,117 @@
+from __future__ import annotations
+
+import logging
+import os
+from collections.abc import Callable
+from typing import Any, Literal
+
+logger = logging.getLogger(__name__)
+
+_BATCHING_MODE_ENV = "CHALK_REMOTE_CALL_BATCHING_MODE"
+_VALID_BATCHING_MODES = ("per_caller", "combined")
+
+
+def _resolve_batching_mode(kwarg: str | None) -> str:
+    """Resolve batching mode from kwarg or env var.
+
+    Precedence: explicit kwarg > env var > default "per_caller". The env var
+    is the canonical opt-in for combined mode; the kwarg exists so embedders
+    can override programmatically.
+    """
+    if kwarg is not None:
+        mode = kwarg
+    else:
+        mode = (os.environ.get(_BATCHING_MODE_ENV, "") or "per_caller").strip() or "per_caller"
+    if mode not in _VALID_BATCHING_MODES:
+        raise ValueError(f"batching_mode must be one of {_VALID_BATCHING_MODES}, got {mode!r}")
+    return mode
+
+
+def serve(
+    handler: Callable[..., Any],
+    host: str = "[::]",
+    port: int = 6666,
+    workers: int = 10,
+    on_startup: Callable[[], None] | None = None,
+    on_shutdown: Callable[[], None] | None = None,
+    arg_names: list[str] | None = None,
+    max_batching_size: int | None = None,
+    max_buffer_duration_ms: int | None = None,
+    batching_mode: Literal["per_caller", "combined"] | None = None,
+) -> None:
+    """Start the gRPC server implementing RemoteCallService.
+
+    Args:
+        handler: The user's handler function.
+            Single-request mode (default): (event, context) -> result.
+            Coalesced per-caller mode (when max_batching_size is set and
+                batching_mode == "per_caller"):
+                (events: list[dict], contexts: list[dict]) -> list[result],
+                one result per input.
+            Coalesced combined mode (when max_batching_size is set and
+                batching_mode == "combined"):
+                (combined: pa.RecordBatch, offsets: list[int], contexts: list[dict])
+                    -> pa.RecordBatch with combined.num_rows rows.
+        host: The host to bind to.
+        port: The port to bind to.
+        workers: Number of worker threads for the tokio runtime.
+        on_startup: Optional startup hook called before the server starts accepting requests.
+        on_shutdown: Optional shutdown hook called after the server stops accepting requests.
+        arg_names: Parsed CHALK_INPUT_ARGS column names, or None.
+        max_batching_size: Maximum number of concurrent requests to coalesce
+            before invoking the handler. None/0 disables coalescing.
+        max_buffer_duration_ms: Max time in milliseconds to buffer incoming
+            requests before flushing even if max_batching_size hasn't been
+            reached. Defaults to 1000ms when batching is enabled.
+        batching_mode: "per_caller" (default) or "combined". Only effective
+            when max_batching_size > 0. Overrides the
+            CHALK_REMOTE_CALL_BATCHING_MODE env var when set; otherwise the
+            env var (or "per_caller") wins.
+    """
+    # Keep package imports lightweight so the CLI can configure logging before dependency failures.
+    from chalk_remote_call._native import start_server
+    from chalk_remote_call.servicer import (
+        process_batches,
+        process_batches_coalesced,
+        process_batches_coalesced_combined,
+    )
+
+    if on_startup is not None:
+        logger.info("Running startup hook...")
+        on_startup()
+
+    mode = _resolve_batching_mode(batching_mode)
+
+    if max_batching_size is not None and max_batching_size > 0:
+        process_fn = process_batches_coalesced_combined if mode == "combined" else process_batches_coalesced
+        batch_size = max_batching_size
+        duration_ms = max_buffer_duration_ms if max_buffer_duration_ms is not None else 1000
+        logger.info(
+            "Batching enabled: mode=%s max_batching_size=%d max_buffer_duration_ms=%d",
+            mode,
+            batch_size,
+            duration_ms,
+        )
+    else:
+        if mode == "combined":
+            raise ValueError("batching_mode='combined' requires max_batching_size > 0")
+        process_fn = process_batches
+        batch_size = 0
+        duration_ms = 0
+
+    logger.info("Starting server on %s:%d", host, port)
+    try:
+        start_server(
+            handler=handler,
+            process_fn=process_fn,
+            host=host,
+            port=port,
+            workers=workers,
+            arg_names=arg_names,
+            max_batching_size=batch_size,
+            max_buffer_duration_ms=duration_ms,
+        )
+    finally:
+        if on_shutdown is not None:
+            logger.info("Running shutdown hook...")
+            on_shutdown()