mplang-nightly 0.1.dev192__py3-none-any.whl → 0.1.dev268__py3-none-any.whl

mplang/v1/ops/sql_cc.py

@@ -0,0 +1,303 @@
+# Copyright 2025 Ant Group Co., Ltd.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from typing import Any
+
+import sqlglot as sg
+from jax.tree_util import PyTreeDef, tree_flatten
+from sqlglot import exp as sge
+from sqlglot.optimizer import annotate_types as opt_annot
+from sqlglot.optimizer import qualify as opt_qualify
+
+from mplang.v1.core import MPObject, PFunction, TableType
+from mplang.v1.core.dtypes import (
+    BINARY,
+    BOOL,
+    DATE,
+    DECIMAL,
+    FLOAT32,
+    FLOAT64,
+    INT8,
+    INT16,
+    INT32,
+    INT64,
+    INTERVAL,
+    JSON,
+    STRING,
+    TIME,
+    TIMESTAMP,
+    UINT8,
+    UINT16,
+    UINT32,
+    UINT64,
+    UUID,
+    DType,
+)
+from mplang.v1.ops.base import stateless_mod
+
+_SQL_MOD = stateless_mod("sql")
+
+
+# Static dtype mappings (MPLang <-> SQL)
+MP_TO_SQL_TYPE: dict[DType, str] = {
+    # Floats
+    FLOAT64: "DOUBLE",
+    FLOAT32: "FLOAT",
+    # Signed ints
+    INT8: "TINYINT",
+    INT16: "SMALLINT",
+    INT32: "INT",
+    INT64: "BIGINT",
+    # Unsigned ints (portable approximations)
+    UINT8: "SMALLINT",
+    UINT16: "INT",
+    UINT32: "BIGINT",
+    UINT64: "DECIMAL(38)",
+    # Booleans & strings
+    BOOL: "BOOLEAN",
+    STRING: "VARCHAR",
+    # Dates / times
+    DATE: "DATE",
+    TIME: "TIME",
+    TIMESTAMP: "TIMESTAMP",
+    # Other table types
+    DECIMAL: "DECIMAL",
+    JSON: "JSON",
+    BINARY: "BLOB",
+    UUID: "UUID",
+    INTERVAL: "INTERVAL",
+}
+
+SQL_TYPE_TO_MP: dict[str, DType] = {
+    # Floats
+    "double": FLOAT64,
+    "double precision": FLOAT64,
+    "float": FLOAT32,
+    "real": FLOAT32,
+    # Signed ints
+    "bigint": INT64,
+    "long": INT64,
+    "int": INT32,
+    "integer": INT32,
+    "int4": INT32,
+    "smallint": INT16,
+    "int2": INT16,
+    "tinyint": INT8,
+    "int1": INT8,
+    # Unsigned (rare in SQL)
+    "uint8": UINT8,
+    "ubyte": UINT8,
+    "uint16": UINT16,
+    "uint32": UINT32,
+    "uint64": UINT64,
+    # Booleans / strings
+    "bool": BOOL,
+    "boolean": BOOL,
+    "char": STRING,
+    "varchar": STRING,
+    "text": STRING,
+    "string": STRING,
+    # Dates / times
+    "date": DATE,
+    "time": TIME,
+    "timestamp": TIMESTAMP,
+    # Decimal / numeric
+    "decimal": DECIMAL,
+    "numeric": DECIMAL,
+    # Others
+    "json": JSON,
+    "binary": BINARY,
+    "varbinary": BINARY,
+    "blob": BINARY,
+    "uuid": UUID,
+    "interval": INTERVAL,
+}
+
+
+def _deduce_out_schema(
+    parsed: sge.Expression,
+    dialect: str,
+    in_schemas: dict[str, TableType],
+) -> TableType:
+    """Deduce output schema using sqlglot's qualify + annotate_types.
+
+    This implementation leverages sqlglot's optimizer to resolve table/column
+    references (including star expansion) and annotate expression types. It then
+    maps sqlglot DataType to mplang DType and returns a TableType.
+    """
+
+    # 1) Build sqlglot schema from MPObject/TableType inputs
+    def _dtype_to_sql(dt: DType) -> str:
+        return MP_TO_SQL_TYPE.get(dt, "VARCHAR")
+
+    sqlglot_schema: dict[str, dict[str, str]] = {
+        tname: {col: _dtype_to_sql(dt) for col, dt in schema.columns}
+        for tname, schema in in_schemas.items()
+    }
+
+    # 2) Parse with read dialect; 3) Qualify (resolve names, expand star); 4) Annotate types
+    qualified = opt_qualify.qualify(parsed, schema=sqlglot_schema, dialect=dialect)
+    typed = opt_annot.annotate_types(qualified, schema=sqlglot_schema)
+
+    # 5) Extract projection names and types
+    select = typed if isinstance(typed, sge.Select) else typed.find(sge.Select)
+    if select is None:
+        raise NotImplementedError(
+            "Only SELECT queries are supported for schema deduction"
+        )
+
+    def _sqlglot_type_to_dtype(tobj: Any) -> DType:
+        ts = str(tobj).lower().replace(" with time zone", "").strip()
+        base = ts.split("(", 1)[0].strip()
+        return SQL_TYPE_TO_MP.get(base, STRING)
+
+    pairs: list[tuple[str, DType]] = []
+    idx = 0
+    used: set[str] = set()
+    for proj in select.expressions:
+        name = getattr(proj, "alias_or_name", None) or getattr(proj, "name", None)
+        if not name:
+            name = f"expr_{idx}"
+            idx += 1
+        t = getattr(proj, "type", None)
+        if t is None:
+            raise NotImplementedError(
+                "Cannot infer type for projection; please provide out_type explicitly"
+            )
+        dtype = _sqlglot_type_to_dtype(t)
+        if name in used:
+            raise ValueError(
+                f"Duplicate output column name '{name}' after qualification"
+            )
+        used.add(name)
+        pairs.append((name, dtype))
+
+    return TableType.from_pairs(pairs)
+
+
+@_SQL_MOD.op_def()
+def run_sql(
+    query: str,
+    *,
+    out_type: TableType | None = None,
+    dialect: str = "duckdb",
+    **in_tables: Any,
+) -> tuple[PFunction, list[MPObject], PyTreeDef]:
+    """Build a sql.run PFunction from a SQL query with optional schema deduction.
+
+    API: run_sql(query: str, *, out_type: TableType | None = None, dialect: str = "duckdb", **in_tables) -> (PFunction, [MPObject], PyTreeDef)
+
+    Semantics:
+    - Parses the SQL and binds only the tables that are actually referenced in the query by name.
+    - If ``out_type`` is not provided, attempts to deduce the output table schema using sqlglot (qualify + annotate types).
+    - Returns a triad consisting of the constructed PFunction (``fn_type='sql.run'``), the ordered list of input MPObjects, and the output PyTreeDef.
+
+    Difference vs ``run_sql_raw``: this op can infer ``out_type`` and will parse the SQL to filter inputs; ``run_sql_raw`` requires an explicit ``out_type`` and does not parse/filter inputs.
+    """
+    # Extract required table names from SQL (order by first appearance)
+    parsed = sg.parse_one(query, read=dialect)
+    required_names: list[str] = []
+    for t in parsed.find_all(sge.Table):
+        # Prefer .name; fallback to str(this) if needed
+        tname = getattr(t, "name", None) or str(t.this)
+        if tname not in required_names:
+            required_names.append(tname)
+
+    # Disallow extras not referenced by the query to avoid surprises
+    extra = set(in_tables.keys()) - set(required_names)
+    if extra:
+        raise ValueError(
+            f"Unexpected tables provided that are not referenced in SQL: {sorted(extra)}"
+        )
+
+    # Validate required tables and require MPObject for runtime registration
+    in_names: list[str] = []
+    ins_info: list[TableType] = []
+    in_vars: list[MPObject] = []
+    for name in required_names:
+        if name not in in_tables:
+            raise KeyError(f"Missing required table '{name}' for SQL query")
+        obj = in_tables[name]
+        if not isinstance(obj, MPObject):
+            raise TypeError(
+                f"Table '{name}' must be an MPObject (for runtime registration), got {type(obj).__name__}"
+            )
+        assert obj.schema is not None, f"Input table '{name}' missing schema"
+        in_vars.append(obj)
+        ins_info.append(obj.schema)
+        in_names.append(name)
+
+    if out_type is None:
+        in_schemas: dict[str, TableType] = {
+            n: in_tables[n].schema for n in required_names
+        }
+        out_type = _deduce_out_schema(parsed, dialect, in_schemas)
+
+    pfn = PFunction(
+        fn_type="sql.run",
+        ins_info=tuple(ins_info),
+        outs_info=(out_type,),
+        fn_name="",
+        fn_text=query,
+        in_names=tuple(in_names),
+        dialect=dialect,
+    )
+    _, treedef = tree_flatten(out_type)
+    return pfn, in_vars, treedef
+
+
+@_SQL_MOD.op_def()
+def run_sql_raw(
+    query: str,
+    out_type: TableType,
+    *,
+    dialect: str = "duckdb",
+    in_tables: dict[str, MPObject] | None = None,
+) -> tuple[PFunction, list[MPObject], PyTreeDef]:
+    """Build a sql.run PFunction from a SQL query with an explicit output schema.
+
+    API: run_sql_raw(query: str, out_type: TableType, *, dialect: str = "duckdb", in_tables: dict[str, MPObject] | None = None) -> (PFunction, [MPObject], PyTreeDef)
+
+    Semantics:
+    - Does not parse the SQL; carries all tables provided via ``in_tables`` in the mapping's iteration order.
+    - Requires an explicit ``out_type``; no schema deduction is attempted.
+    - Returns a triad consisting of the constructed PFunction (``fn_type='sql.run'``), the ordered list of input MPObjects, and the output PyTreeDef.
+
+    Difference vs ``run_sql``: this op requires ``out_type`` and does not parse/filter inputs; ``run_sql`` can infer ``out_type`` and selects only tables referenced by the query.
+    """
+
+    # Collect inputs strictly as provided by caller
+    in_names: list[str] = []
+    ins_info: list[TableType] = []
+    in_vars: list[MPObject] = []
+    if in_tables:
+        for name, tbl in in_tables.items():
+            if not isinstance(tbl, MPObject):
+                raise TypeError(f"Input table '{name}' is not an MPObject {type(tbl)}")
+            assert tbl.schema is not None, f"Input table '{name}' is missing a schema"
+            in_names.append(name)
+            ins_info.append(tbl.schema)
+            in_vars.append(tbl)
+
+    pfn = PFunction(
+        fn_type="sql.run",
+        fn_name="",
+        fn_text=query,
+        ins_info=tuple(ins_info),
+        outs_info=(out_type,),
+        in_names=tuple(in_names),
+        dialect=dialect,
+    )
+    _, treedef = tree_flatten(out_type)
+    return pfn, in_vars, treedef

mplang/{ops → v1/ops}/tee.py

@@ -14,8 +14,8 @@
 
 from __future__ import annotations
 
-from mplang.core import UINT8, TensorType
-from mplang.ops.base import stateless_mod
+from mplang.v1.core import UINT8, TensorType
+from mplang.v1.ops.base import stateless_mod
 
 _TEE_MOD = stateless_mod("tee")
 

mplang/{runtime → v1/runtime}/__init__.py

@@ -20,8 +20,8 @@ This module contains runtime implementations including:
 - Driver for distributed execution
 """
 
-from mplang.runtime.driver import Driver, DriverVar
-from mplang.runtime.simulation import Simulator
+from mplang.v1.runtime.driver import Driver, DriverVar
+from mplang.v1.runtime.simulation import Simulator
 
 __all__ = [
     "Driver",

mplang/v1/runtime/channel.py

@@ -0,0 +1,230 @@
+# Copyright 2025 Ant Group Co., Ltd.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""SPU IChannel implementation that bridges to MPLang CommunicatorBase.
+
+This module provides BaseChannel, which allows SPU to reuse MPLang's
+existing communication layer (ThreadCommunicator/HttpCommunicator) instead
+of creating separate BRPC connections.
+"""
+
+from __future__ import annotations
+
+import logging
+from typing import TYPE_CHECKING
+
+import spu.libspu as libspu
+
+if TYPE_CHECKING:
+    from mplang.v1.core.comm import CommunicatorBase
+
+
+class BaseChannel(libspu.link.IChannel):
+    """Bridge MPLang CommunicatorBase to SPU IChannel interface.
+
+    This adapter allows SPU to use MPLang's existing communication layer
+    (ThreadCommunicator or HttpCommunicator) instead of creating separate
+    BRPC connections.
+
+    Each BaseChannel represents a channel to ONE peer rank.
+
+    Communication Protocol:
+        - SPU calls send(tag, bytes_data) -> MPLang comm.send(peer, key, bytes_data)
+        - SPU calls recv(tag) -> bytes_data <- MPLang comm.recv(peer, key)
+
+    Tag Namespace:
+        All tags are prefixed with "spu:" to avoid collision with other
+        MPLang traffic on the same communicator.
+    """
+
+    def __init__(
+        self,
+        comm: CommunicatorBase,
+        local_rank: int,
+        peer_rank: int,
+        tag_prefix: str = "spu",
+    ):
+        """Initialize channel to a specific peer.
+
+        Args:
+            comm: MPLang communicator instance (Thread/Http)
+            local_rank: Global rank of this party (for logging/debugging)
+            peer_rank: Global rank of the peer party
+            tag_prefix: Prefix for all tags to avoid collision (default: "spu")
+        """
+        super().__init__()
+        self._comm = comm
+        self._local_rank = local_rank
+        self._peer_rank = peer_rank
+        self._tag_prefix = tag_prefix
+
+        logging.debug(
+            f"BaseChannel initialized: local_rank={local_rank}, "
+            f"peer_rank={peer_rank}, tag_prefix={tag_prefix}"
+        )
+
+    def _make_key(self, tag: str) -> str:
+        """Create unique key for MPLang comm.
+
+        Prefixes the tag to avoid collision with non-SPU traffic.
+
+        Args:
+            tag: SPU-provided tag (e.g., "send_0", "recv_0")
+
+        Returns:
+            Prefixed key (e.g., "spu:send_0")
+        """
+        return f"{self._tag_prefix}:{tag}"
+
+    def Send(self, tag: str, data: bytes) -> None:
+        """Send bytes to peer (synchronous in SPU semantics).
+
+        Args:
+            tag: Message tag for matching send/recv pairs
+            data: Raw bytes to send
+        """
+        key = self._make_key(tag)
+        logging.debug(
+            f"BaseChannel.Send: {self._local_rank} -> {self._peer_rank}, "
+            f"tag={tag}, key={key}, size={len(data)}"
+        )
+
+        # Send raw bytes directly
+        # Note: CommunicatorBase.send expects Any type, bytes is acceptable
+        self._comm.send(self._peer_rank, key, data)
+
+    def Recv(self, tag: str) -> bytes:
+        """Receive bytes from peer (blocking).
+
+        Args:
+            tag: Message tag for matching send/recv pairs
+
+        Returns:
+            Raw bytes received
+        """
+        key = self._make_key(tag)
+        logging.debug(
+            f"BaseChannel.Recv: {self._local_rank} <- {self._peer_rank}, "
+            f"tag={tag}, key={key}"
+        )
+
+        # Receive data (should be bytes)
+        data = self._comm.recv(self._peer_rank, key)
+
+        # Validate data type
+        if not isinstance(data, bytes):
+            raise TypeError(
+                f"Expected bytes from communicator, got {type(data).__name__}. "
+                f"Communicator must support raw bytes transmission for SPU channels."
+            )
+
+        logging.debug(
+            f"BaseChannel.Recv complete: {self._local_rank} <- {self._peer_rank}, "
+            f"tag={tag}, size={len(data)}"
+        )
+        return data
+
+    def SendAsync(self, tag: str, data: bytes) -> None:
+        """Async send (MPLang's send is already async at network layer).
+
+        For HttpCommunicator, the underlying httpx.put() is non-blocking
+        at the HTTP client level. For ThreadCommunicator, send is instant
+        (memory transfer).
+
+        Args:
+            tag: Message tag
+            data: Raw bytes to send
+        """
+        # Reuse synchronous send - it's already async underneath
+        self.Send(tag, data)
+
+    def SendAsyncThrottled(self, tag: str, data: bytes) -> None:
+        """Throttled async send.
+
+        Currently maps to regular SendAsync. Future optimization could
+        implement rate limiting if needed.
+
+        Args:
+            tag: Message tag
+            data: Raw bytes to send
+        """
+        self.SendAsync(tag, data)
+
+    def TestSend(self, timeout: int) -> None:
+        """Test if this channel can send a dummy msg to peer.
+
+        Uses fixed 0 seq_id as dummy msg's id to make this function reentrant.
+        ConnectToMesh will retry on this multiple times.
+
+        Args:
+            timeout: Timeout in milliseconds
+        """
+        # Send a handshake message to test connectivity
+        # Use fixed tag "__test__" to make this reentrant (idempotent)
+        test_data = b"\x00"  # Minimal 1-byte message with seq_id=0
+        self.Send("__test__", test_data)
+
+    def TestRecv(self) -> None:
+        """Wait for dummy msg from peer.
+
+        Timeout is controlled by recv_timeout_ms in link descriptor.
+        """
+        # Receive the handshake message from peer
+        # This blocks until message arrives (timeout from desc.recv_timeout_ms)
+        test_data = self.Recv("__test__")
+        # Validate it's the expected handshake message
+        if test_data != b"\x00":
+            logging.warning(
+                f"TestRecv: unexpected handshake data from {self._peer_rank}, "
+                f"expected b'\\x00', got {test_data!r}"
+            )
+
+    def WaitLinkTaskFinish(self) -> None:
+        """Wait for all pending async tasks.
+
+        For MPLang communicators:
+        - ThreadCommunicator: No-op (instant memory transfer)
+        - HttpCommunicator: No explicit wait needed (httpx handles it)
+
+        This is a no-op in current implementation.
+        """
+
+    def Abort(self) -> None:
+        """Abort communication (cleanup resources).
+
+        This could be extended to notify the communicator to drop pending
+        messages for this channel, but currently is a no-op.
+        """
+        logging.warning(
+            f"BaseChannel.Abort called: {self._local_rank} <-> {self._peer_rank}"
+        )
+        # Future: Could call comm.abort_session() if implemented
+
+    def SetThrottleWindowSize(self, size: int) -> None:
+        """Set throttle window size.
+
+        Not applicable to MPLang communicators. No-op.
+
+        Args:
+            size: Window size (ignored)
+        """
+
+    def SetChunkParallelSendSize(self, size: int) -> None:
+        """Set chunk parallel send size.
+
+        Not applicable to MPLang communicators. No-op.
+
+        Args:
+            size: Chunk size (ignored)
+        """

mplang/{runtime → v1/runtime}/cli.py

@@ -26,9 +26,9 @@ from typing import Any
 import uvicorn
 import yaml
 
-from mplang.core import ClusterSpec
-from mplang.runtime.client import HttpExecutorClient
-from mplang.runtime.server import app
+from mplang.v1.core import ClusterSpec
+from mplang.v1.runtime.client import HttpExecutorClient
+from mplang.v1.runtime.server import app
 
 
 def load_config(config_path: str) -> ClusterSpec:

mplang/{runtime → v1/runtime}/client.py

@@ -27,7 +27,7 @@ from typing import Any
 
 import httpx
 
-from mplang.kernels.value import Value, decode_value, encode_value
+from mplang.v1.kernels.value import Value, decode_value, encode_value
 
 
 class ExecutionStatus:

mplang/{runtime → v1/runtime}/communicator.py

@@ -23,8 +23,8 @@ from typing import Any
 
 import httpx
 
-from mplang.core.comm import CommunicatorBase
-from mplang.kernels.value import Value, decode_value, encode_value
+from mplang.v1.core.comm import CommunicatorBase
+from mplang.v1.kernels.value import Value, decode_value, encode_value
 
 
 class HttpCommunicator(CommunicatorBase):
@@ -57,7 +57,12 @@ class HttpCommunicator(CommunicatorBase):
         return str(res)
 
     def send(self, to: int, key: str, data: Any) -> None:
-        """Sends data to a peer party by PUTing to its /comm/{key}/from/{from_rank} endpoint."""
+        """Sends data to a peer party by PUTing to its /comm/{key}/from/{from_rank} endpoint.
+
+        Supports two modes:
+        - SPU channel (key starts with "spu:"): sends raw bytes directly
+        - Normal channel: wraps data in Value envelope
+        """
         target_endpoint = self.endpoints[to]
         url = f"{target_endpoint}/sessions/{self.session_name}/comm/{key}/from/{self._rank}"
         logging.debug(
@@ -65,19 +70,20 @@ class HttpCommunicator(CommunicatorBase):
         )
 
         try:
-            # Serialize data using Value envelope.
-            if not isinstance(data, Value):
+            # SPU channel mode: send raw bytes directly
+            if key.startswith("spu:") and isinstance(data, bytes):
+                data_b64 = base64.b64encode(data).decode("utf-8")
+                request_data = {"data": data_b64, "is_raw_bytes": True}
+            # Normal mode: serialize using Value envelope
+            elif isinstance(data, Value):
+                data_bytes = encode_value(data)
+                data_b64 = base64.b64encode(data_bytes).decode("utf-8")
+                request_data = {"data": data_b64}
+            else:
                 raise TypeError(
                     f"Communicator requires Value instance, got {type(data).__name__}. "
                     "Wrap data in TensorValue or custom Value subclass."
                 )
-            data_bytes = encode_value(data)
-
-            data_b64 = base64.b64encode(data_bytes).decode("utf-8")
-
-            request_data = {
-                "data": data_b64,
-            }
 
             response = httpx.put(url, json=request_data, timeout=60)
             logging.debug(f"Send response: status={response.status_code}")
@@ -91,14 +97,32 @@ class HttpCommunicator(CommunicatorBase):
             raise OSError(f"Failed to send data to rank {to}") from e
 
     def recv(self, frm: int, key: str) -> Any:
-        """Wait until the key is set, returns the value. Override to add logging."""
+        """Wait until the key is set, returns the value.
+
+        Supports two modes:
+        - SPU channel (key starts with "spu:"): returns raw bytes
+        - Normal channel: returns deserialized Value
+        """
         logging.debug(
             f"Waiting to receive: from_rank={frm}, to_rank={self._rank}, key={key}"
         )
-        data_b64 = super().recv(frm, key)
+        received_data = super().recv(frm, key)
 
+        # Check if this is raw bytes (SPU channel)
+        if isinstance(received_data, dict) and received_data.get("is_raw_bytes"):
+            data_bytes = base64.b64decode(received_data["data"])
+            logging.debug(
+                f"Received raw bytes: from_rank={frm}, to_rank={self._rank}, key={key}, size={len(data_bytes)}"
+            )
+            return data_bytes
+
+        # Normal mode: deserialize Value envelope
+        data_b64 = (
+            received_data
+            if isinstance(received_data, str)
+            else received_data.get("data")
+        )
         data_bytes = base64.b64decode(data_b64)
-        # Deserialize using Value envelope
         result = decode_value(data_bytes)
 
         logging.debug(