mplang-nightly 0.1.dev271__tar.gz → 0.1.dev273__tar.gz

{mplang_nightly-0.1.dev271 → mplang_nightly-0.1.dev273}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: mplang-nightly
-Version: 0.1.dev271
+Version: 0.1.dev273
 Summary: Multi-Party Programming Language
 Author-email: SecretFlow Team <secretflow-contact@service.alipay.com>
 License:                                  Apache License

{mplang_nightly-0.1.dev271 → mplang_nightly-0.1.dev273}/mplang/__init__.py

@@ -91,6 +91,9 @@ from mplang.libs.device import (
     set_dev_attr,
 )
 from mplang.libs.device import fetch as device_fetch
+
+# Logging configuration
+from mplang.logging_config import disable_logging, get_logger, setup_logging
 from mplang.runtime.interpreter import Interpreter
 
 # =============================================================================
@@ -359,6 +362,10 @@ def compile(
 __all__ = [  # noqa: RUF022
     # Version
     "__version__",
+    # Logging
+    "disable_logging",
+    "setup_logging",
+    "get_logger",
     # Device API
     "ClusterSpec",
     "Device",

{mplang_nightly-0.1.dev271 → mplang_nightly-0.1.dev273}/mplang/backends/channel.py

@@ -21,11 +21,14 @@ infrastructure instead of creating separate BRPC connections.
 
 from __future__ import annotations
 
-import logging
 from typing import Protocol
 
 import spu.libspu as libspu
 
+from mplang.logging_config import get_logger
+
+logger = get_logger(__name__)
+
 
 class CommunicatorProtocol(Protocol):
     """Protocol for v2 communicators (duck typing).
@@ -73,9 +76,11 @@ class BaseChannel(libspu.link.IChannel):
         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}"
+        logger.debug(
+            "BaseChannel initialized: local_rank=%d, peer_rank=%d, tag_prefix=%s",
+            local_rank,
+            peer_rank,
+            tag_prefix,
         )
 
     def _make_key(self, tag: str) -> str:
@@ -97,9 +102,13 @@ class BaseChannel(libspu.link.IChannel):
             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)}"
+        logger.debug(
+            "BaseChannel.Send: %d -> %d, tag=%s, key=%s, size=%d",
+            self._local_rank,
+            self._peer_rank,
+            tag,
+            key,
+            len(data),
         )
 
         # Send raw bytes directly
@@ -119,9 +128,12 @@ class BaseChannel(libspu.link.IChannel):
             TypeError: If received data is not bytes
         """
         key = self._make_key(tag)
-        logging.debug(
-            f"BaseChannel.Recv: {self._local_rank} <- {self._peer_rank}, "
-            f"tag={tag}, key={key}"
+        logger.debug(
+            "BaseChannel.Recv: %d <- %d, tag=%s, key=%s",
+            self._local_rank,
+            self._peer_rank,
+            tag,
+            key,
         )
 
         # Receive data (should be bytes)
@@ -134,9 +146,12 @@ class BaseChannel(libspu.link.IChannel):
                 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)}"
+        logger.debug(
+            "BaseChannel.Recv complete: %d <- %d, tag=%s, size=%d",
+            self._local_rank,
+            self._peer_rank,
+            tag,
+            len(data),
         )
         return data
 
@@ -184,9 +199,10 @@ class BaseChannel(libspu.link.IChannel):
         """
         test_data = self.Recv("__test__")
         if test_data != b"\x00":
-            logging.warning(
-                f"TestRecv: unexpected handshake from {self._peer_rank}, "
-                f"expected b'\\x00', got {test_data!r}"
+            logger.warning(
+                "TestRecv: unexpected handshake from %d, expected b'\\x00', got %r",
+                self._peer_rank,
+                test_data,
             )
 
     def WaitLinkTaskFinish(self) -> None:
@@ -200,7 +216,9 @@ class BaseChannel(libspu.link.IChannel):
 
         Currently a no-op. Could be extended for resource cleanup.
         """
-        logging.warning(f"BaseChannel.Abort: {self._local_rank} <-> {self._peer_rank}")
+        logger.warning(
+            "BaseChannel.Abort: %d <-> %d", self._local_rank, self._peer_rank
+        )
 
     def SetThrottleWindowSize(self, size: int) -> None:
         """Set throttle window size (no-op).

{mplang_nightly-0.1.dev271 → mplang_nightly-0.1.dev273}/mplang/backends/simp_driver/http.py

@@ -25,9 +25,12 @@ import httpx
 
 from mplang.backends.simp_driver.state import SimpDriver
 from mplang.edsl import serde
+from mplang.logging_config import get_logger
 from mplang.runtime.interpreter import Interpreter
 from mplang.runtime.object_store import ObjectStore
 
+logger = get_logger(__name__)
+
 if TYPE_CHECKING:
     from concurrent.futures import Future
 

{mplang_nightly-0.1.dev271 → mplang_nightly-0.1.dev273}/mplang/backends/simp_worker/http.py

@@ -36,7 +36,6 @@ Security:
 from __future__ import annotations
 
 import concurrent.futures
-import logging
 import os
 import pathlib
 import threading
@@ -55,10 +54,11 @@ from mplang.backends.simp_worker import SimpWorker
 from mplang.backends.simp_worker import ops as _simp_worker_ops  # noqa: F401
 from mplang.edsl import serde
 from mplang.edsl.graph import Graph
+from mplang.logging_config import get_logger
 from mplang.runtime.interpreter import ExecutionTracer, Interpreter
 from mplang.runtime.object_store import ObjectStore
 
-logger = logging.getLogger(__name__)
+logger = get_logger(__name__)
 
 
 class HttpCommunicator:
@@ -99,7 +99,7 @@ class HttpCommunicator:
     def _do_send(self, to: int, key: str, data: Any) -> None:
         """Perform the HTTP send."""
         url = f"{self.endpoints[to]}/comm/{key}"
-        logger.debug(f"Rank {self.rank} sending to {to} key={key}")
+        logger.debug(f"Rank {self.rank} sending to {to} key={key}, url={url}")
 
         # Detect SPU channel (tag prefix "spu:") and handle bytes
         if key.startswith("spu:") and isinstance(data, bytes):

{mplang_nightly-0.1.dev271 → mplang_nightly-0.1.dev273}/mplang/backends/spu_impl.py

@@ -31,9 +31,12 @@ from mplang.backends.tensor_impl import TensorValue
 from mplang.dialects import spu
 from mplang.edsl import serde
 from mplang.edsl.graph import Operation
+from mplang.logging_config import get_logger
 from mplang.runtime.interpreter import Interpreter
 from mplang.runtime.value import WrapValue
 
+logger = get_logger(__name__)
+
 # =============================================================================
 # SPU Share Wrapper
 # =============================================================================

{mplang_nightly-0.1.dev271 → mplang_nightly-0.1.dev273}/mplang/backends/table_impl.py

@@ -33,9 +33,12 @@ from mplang.backends.tensor_impl import TensorValue
 from mplang.dialects import table
 from mplang.edsl import serde
 from mplang.edsl.graph import Operation
+from mplang.logging_config import get_logger
 from mplang.runtime.interpreter import Interpreter
 from mplang.runtime.value import WrapValue
 
+logger = get_logger(__name__)
+
 
 class BatchReader(ABC):
     @property

{mplang_nightly-0.1.dev271 → mplang_nightly-0.1.dev273}/mplang/backends/tensor_impl.py

@@ -36,9 +36,12 @@ import mplang.edsl.typing as elt
 from mplang.dialects import dtypes, tensor
 from mplang.edsl import serde
 from mplang.edsl.graph import Operation
+from mplang.logging_config import get_logger
 from mplang.runtime.interpreter import Interpreter
 from mplang.runtime.value import Value, WrapValue
 
+logger = get_logger(__name__)
+
 # =============================================================================
 # TensorValue Wrapper
 # =============================================================================

{mplang_nightly-0.1.dev271 → mplang_nightly-0.1.dev273}/mplang/cli.py

@@ -49,6 +49,10 @@ from typing import Any, cast
 import uvicorn
 import yaml
 
+from mplang.logging_config import get_logger
+
+logger = get_logger(__name__)
+
 
 def run_worker(
     rank: int,
@@ -287,26 +291,33 @@ def cmd_status(args: argparse.Namespace) -> None:
             resp.raise_for_status()
             print(f"OK  {url} -> {resp.json()}")
         except Exception as exc:
+            logger.warning("Health check failed for %s: %s", url, exc)
             print(f"ERR {url} -> {exc}")
 
 
 def load_user_module(path: str) -> ModuleType:
     """Load a Python module from file path."""
+    logger.debug("Loading user module from %s", path)
     if not os.path.exists(path):
+        logger.error("Module file not found: %s", path)
         raise FileNotFoundError(path)
     spec = importlib.util.spec_from_file_location("mp_user_module", path)
     if spec is None or spec.loader is None:
+        logger.error("Cannot import module from %s", path)
         raise ImportError(f"Cannot import module from {path}")
     module = importlib.util.module_from_spec(spec)
     sys.modules[spec.name] = module
     spec.loader.exec_module(module)
+    logger.debug("Successfully loaded module: %s", path)
     return module
 
 
 def resolve_entry(module: ModuleType, name: str) -> Callable[..., Any]:
     entry = getattr(module, name, None)
     if entry is None or not callable(entry):
+        logger.error("Entry function '%s' not found or not callable in module", name)
         raise AttributeError(f"Entry function '{name}' not found or not callable")
+    logger.debug("Resolved entry function: %s", name)
     return cast(Callable[..., Any], entry)
 
 
@@ -421,6 +432,7 @@ def cmd_trace_merge(args: argparse.Namespace) -> None:
                     merged_events.append(event)
 
         except Exception as e:
+            logger.error("Error processing trace file %s: %s", fname, e, exc_info=True)
             print(f"Error processing {fname}: {e}")
 
     # Write merged file
@@ -458,6 +470,7 @@ def cmd_objects(args: argparse.Namespace) -> None:
                     f"{rank:<6} | {endpoints[rank]:<25} | {'Err':<6} | Status {resp.status_code}"
                 )
         except Exception as e:
+            logger.warning("Failed to list objects on rank %d: %s", rank, e)
             print(f"{rank:<6} | {endpoints[rank]:<25} | {'Err':<6} | {e}")
 
 

{mplang_nightly-0.1.dev271 → mplang_nightly-0.1.dev273}/mplang/edsl/jit.py

@@ -25,6 +25,9 @@ from mplang.edsl.context import (
     get_default_context,
 )
 from mplang.edsl.tracer import Tracer
+from mplang.logging_config import get_logger
+
+logger = get_logger(__name__)
 
 
 def jit(fn: Callable) -> Callable:
@@ -42,13 +45,17 @@ def jit(fn: Callable) -> Callable:
     """
 
     def wrapper(*args: Any, **kwargs: Any) -> Any:
+        fn_name = getattr(fn, "__name__", "anonymous")
+
         # If we are already inside a Tracer (e.g. pcall_static), just inline
         # the function to trace it into the current graph.
         cur_ctx = get_current_context()
         if isinstance(cur_ctx, Tracer):
+            logger.debug("JIT: inlining '%s' into current trace context", fn_name)
             return fn(*args, **kwargs)
 
         # otherwise trace for JIT compilation
+        logger.debug("JIT: compiling '%s'", fn_name)
         with Tracer():
             result = fn(*args, **kwargs)
 
@@ -57,6 +64,7 @@ def jit(fn: Callable) -> Callable:
         assert isinstance(cur_ctx, AbstractInterpreter), (
             "JIT execution requires Interpreter context"
         )
+        logger.debug("JIT: executing compiled '%s'", fn_name)
         return tree_map(cur_ctx.lift, result)
 
     return wrapper

{mplang_nightly-0.1.dev271 → mplang_nightly-0.1.dev273}/mplang/edsl/primitive.py

@@ -29,12 +29,15 @@ from jax.tree_util import tree_map
 
 from mplang.edsl.context import get_current_context, get_default_context
 from mplang.edsl.object import Object
+from mplang.logging_config import get_logger
 
 if TYPE_CHECKING:
     from mplang.edsl.typing import BaseType
 
 T_Ret = TypeVar("T_Ret")
 
+logger = get_logger(__name__)
+
 
 class Primitive(Generic[T_Ret]):
     """Atomic operation definition (similar to JAX Primitive).
@@ -89,6 +92,7 @@ class Primitive(Generic[T_Ret]):
         Returns:
             The same function (for decorator pattern)
         """
+        logger.debug("Registering impl for primitive '%s'", self.name)
         self._impl = fn
         # Register with the global interpreter registry
         from mplang.edsl.registry import register_impl
@@ -140,6 +144,7 @@ class Primitive(Generic[T_Ret]):
             >>> # Variable number of inputs
             >>>     return in_types[0]  # Concatenated type
         """
+        logger.debug("Registering abstract_eval for primitive '%s'", self.name)
         self._abstract_eval = fn
         return fn
 
@@ -187,6 +192,7 @@ class Primitive(Generic[T_Ret]):
             >>>     parts = [slice_p.bind(x, i) for i in range(num_splits)]
             >>>     return parts  # Returns list of Objects
         """
+        logger.debug("Registering trace for primitive '%s'", self.name)
         self._trace = fn
         return fn
 

{mplang_nightly-0.1.dev271 → mplang_nightly-0.1.dev273}/mplang/edsl/tracer.py

@@ -1,4 +1,4 @@
-# Copyright 2025 Ant Group Co., Ltd.
+# Copyright 2026 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.
@@ -36,10 +36,13 @@ from mplang.edsl.graph import Graph
 from mplang.edsl.graph import Value as GraphValue
 from mplang.edsl.object import Object
 from mplang.edsl.typing import BaseType
+from mplang.logging_config import get_logger
 
 if TYPE_CHECKING:
     from mplang.edsl.primitive import Primitive
 
+logger = get_logger(__name__)
+
 
 class TraceObject(Object):
     """Trace-time object (during JIT tracing).
@@ -121,6 +124,12 @@ class Tracer(Context):
         Raises:
             RuntimeError: If primitive has neither trace nor abstract_eval defined
         """
+        logger.debug(
+            "Binding primitive '%s' in trace mode (args=%d, kwargs=%s)",
+            primitive.name,
+            len(args),
+            list(kwargs.keys()),
+        )
         if primitive._trace is not None:
             return primitive._trace(*args, **kwargs)
 
@@ -298,6 +307,7 @@ class Tracer(Context):
             raise TypeError(f"fn must be callable, got {type(fn)}")
 
         fn_name = getattr(fn, "__name__", "anonymous")
+        logger.debug("Starting trace of function: %s", fn_name)
         in_flat, in_treedef = tree_flatten((args, kwargs))
         in_imms, in_var_pos, in_vars = _separate_vars_and_imms(in_flat)
 
@@ -329,6 +339,15 @@ class Tracer(Context):
             obj for obj, _ in self._captured_vars.values()
         ]
 
+        logger.debug(
+            "Trace completed: %s - %d inputs, %d outputs, %d ops, %d captured vars",
+            fn_name,
+            len(graph.inputs),
+            len(graph.outputs),
+            len(graph.operations),
+            len(captured_objects),
+        )
+
         return TracedFunction(
             name=fn_name,
             graph=graph,

{mplang_nightly-0.1.dev271 → mplang_nightly-0.1.dev273}/mplang/libs/device/api.py

@@ -32,9 +32,12 @@ from mplang.backends import load_builtins
 from mplang.dialects import crypto, simp, spu, tee
 from mplang.edsl.object import Object
 from mplang.libs.device.cluster import Device
+from mplang.logging_config import get_logger
 
 load_builtins()
 
+logger = get_logger(__name__)
+
 
 def _resolve_cluster() -> Any:
     """Resolve the active ClusterSpec by traversing the context stack.

mplang_nightly-0.1.dev273/mplang/logging_config.py

@@ -0,0 +1,202 @@
+# Copyright 2026 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.
+
+"""
+Logging configuration for MPLang.
+
+This module provides a unified logging setup for the MPLang library.
+When MPLang is used as a library, logging is disabled by default (NullHandler),
+allowing applications to configure logging as needed.
+
+Example usage:
+    >>> import mplang as mp
+    >>> # Enable logging with INFO level
+    >>> mp.setup_logging(level="INFO")
+    >>> # Enable logging with DEBUG level and write to file
+    >>> mp.setup_logging(level="DEBUG", filename="mplang.log")
+    >>> # Customize format
+    >>> mp.setup_logging(
+    ...     level="DEBUG", format="%(asctime)s [%(name)s] %(levelname)s: %(message)s"
+    ... )
+"""
+
+import logging
+import sys
+from typing import Any, Literal
+
+# Root logger for all MPLang components
+MPLANG_LOGGER_NAME = "mplang"
+
+# Default format for log messages
+DEFAULT_FORMAT = "%(asctime)s [%(levelname)s] %(name)s: %(message)s"
+DEFAULT_DATE_FORMAT = "%Y-%m-%d %H:%M:%S"
+
+
+def setup_logging(
+    level: Literal["DEBUG", "INFO", "WARNING", "ERROR", "CRITICAL"] = "INFO",
+    format: str | None = None,
+    date_format: str | None = None,
+    filename: str | None = None,
+    stream: Any = None,  # type: ignore[type-arg]
+    force: bool = False,
+    propagate: bool = False,
+) -> None:
+    """
+    Configure logging for MPLang.
+
+    This function sets up a logger for all MPLang components. By default,
+    MPLang uses a NullHandler to suppress log output when used as a library.
+    Call this function to enable logging with custom settings.
+
+    For unified application logging (application + mplang logs in same file):
+    - Option 1: Configure Python's root logger, then call setup_logging(propagate=True)
+    - Option 2: Just call setup_logging(filename="app.log") to add mplang logs to that file
+
+    Args:
+        level: Log level (DEBUG, INFO, WARNING, ERROR, CRITICAL). Default is INFO.
+        format: Custom log format string. If None, uses default format.
+        date_format: Custom date format string. If None, uses default format.
+        filename: If provided, log to this file in addition to stream output.
+        stream: Stream to log to (e.g., sys.stdout, sys.stderr). Default is sys.stderr.
+                Set to False to disable stream output (only use file or propagation).
+        force: If True, remove existing handlers before adding new ones.
+        propagate: If True, allow logs to propagate to parent loggers (useful for
+                   unified application logging). If False (default), mplang manages
+                   its own handlers independently.
+
+    Example:
+        >>> import mplang as mp
+        >>> # Enable debug logging to stderr
+        >>> mp.setup_logging(level="DEBUG")
+        >>> # Log mplang to file only
+        >>> mp.setup_logging(level="INFO", filename="mplang.log", stream=False)
+        >>>
+        >>> # Unified application logging (recommended for applications)
+        >>> import logging
+        >>> logging.basicConfig(
+        ...     level=logging.INFO,
+        ...     filename="app.log",
+        ...     format="%(asctime)s [%(levelname)s] %(name)s: %(message)s",
+        ... )
+        >>> mp.setup_logging(level="INFO", propagate=True)  # Inherits app config
+    """
+    # Get the root logger for mplang
+    logger = logging.getLogger(MPLANG_LOGGER_NAME)
+
+    # Set log level
+    log_level = getattr(logging, level.upper())
+    logger.setLevel(log_level)
+
+    # If propagate=True, automatically remove NullHandler to avoid confusion
+    # (NullHandler doesn't prevent propagation, but its presence is confusing)
+    if propagate and not force:
+        # Check if there's only a NullHandler
+        if len(logger.handlers) == 1 and isinstance(
+            logger.handlers[0], logging.NullHandler
+        ):
+            force = True  # Auto-enable force to remove NullHandler
+
+    # Remove existing handlers if force=True
+    if force:
+        for handler in logger.handlers[:]:
+            logger.removeHandler(handler)
+
+    # Set propagation behavior
+    logger.propagate = propagate
+
+    # If propagate is True and no handlers specified, just set level and propagate
+    # This allows mplang logs to use the application's logging configuration
+    if propagate and not filename and stream is None:
+        return
+
+    # Use default format if not provided
+    log_format = format or DEFAULT_FORMAT
+    log_date_format = date_format or DEFAULT_DATE_FORMAT
+    formatter = logging.Formatter(log_format, datefmt=log_date_format)
+
+    # Add stream handler (default to stderr, or skip if stream=False)
+    if stream is not False:
+        if stream is None:
+            stream = sys.stderr
+        stream_handler = logging.StreamHandler(stream)
+        stream_handler.setLevel(log_level)
+        stream_handler.setFormatter(formatter)
+        logger.addHandler(stream_handler)
+
+    # Add file handler if filename is provided
+    if filename:
+        file_handler = logging.FileHandler(filename)
+        file_handler.setLevel(log_level)
+        file_handler.setFormatter(formatter)
+        logger.addHandler(file_handler)
+
+
+def disable_logging() -> None:
+    """
+    Disable all MPLang logging by adding a NullHandler.
+
+    This is useful for testing or when you want to completely suppress
+    MPLang log output.
+
+    Example:
+        >>> import mplang as mp
+        >>> mp.disable_logging()
+    """
+    logger = logging.getLogger(MPLANG_LOGGER_NAME)
+    # Remove all existing handlers
+    for handler in logger.handlers[:]:
+        logger.removeHandler(handler)
+    # Add NullHandler to suppress all output
+    logger.addHandler(logging.NullHandler())
+    logger.propagate = False
+
+
+def get_logger(name: str) -> logging.Logger:
+    """
+    Get a logger for a specific MPLang module.
+
+    This function should be used by all MPLang modules to create their loggers.
+    The logger name will be prefixed with 'mplang' automatically if not already.
+
+    Args:
+        name: Module name, typically __name__ from the calling module.
+
+    Returns:
+        A logger instance for the specified module.
+
+    Example:
+        >>> # In mplang/edsl/tracer.py
+        >>> logger = get_logger(__name__)  # Creates 'mplang.edsl.tracer' logger
+    """
+    # Ensure the logger name is under mplang hierarchy
+    if not name.startswith(MPLANG_LOGGER_NAME):
+        # Handle cases where __name__ might be a relative module name
+        if name.startswith("mplang."):
+            pass  # Already correct
+        elif "." in name:
+            # It's a submodule, but not under mplang, so prefix it.
+            name = f"{MPLANG_LOGGER_NAME}.{name}"
+        else:
+            # Relative module name, prefix with mplang
+            name = f"{MPLANG_LOGGER_NAME}.{name}"
+
+    return logging.getLogger(name)
+
+
+# Initialize with NullHandler by default (library mode)
+# Applications using MPLang should call setup_logging() to enable logging
+_root_logger = logging.getLogger(MPLANG_LOGGER_NAME)
+if not _root_logger.handlers:
+    _root_logger.addHandler(logging.NullHandler())
+    _root_logger.propagate = False