Remove MLX runtime fallback and debug logging

Author: zeel2104

backends/mlx/examples/llm/run_llm_hf.py

@@ -21,12 +21,8 @@
 """
 
 import argparse
-import ctypes
 import logging
-import os
-import shutil
 import time
-from pathlib import Path
 
 import torch
 from executorch.runtime import Runtime, Verification
@@ -37,73 +33,6 @@
 logger = logging.getLogger(__name__)
 
 
-def _iter_mlx_backend_candidates():
-    env_path = os.environ.get("ET_MLX_BACKEND_DYLIB")
-    if env_path:
-        yield Path(env_path)
-
-    for parent in Path(__file__).resolve().parents:
-        pip_out = parent / "pip-out"
-        if pip_out.exists():
-            yield from sorted(
-                pip_out.glob(
-                    "temp.*/cmake-out/backends/mlx/libmlxdelegate_runtime.dylib"
-                )
-            )
-            yield from sorted(
-                pip_out.glob("temp.*/cmake-out/backends/mlx/libmlxdelegate.dylib")
-            )
-            break
-
-
-def _ensure_mlx_metallib(dylib_path: Path) -> None:
-    metallib_path = dylib_path.with_name("mlx.metallib")
-    if metallib_path.exists():
-        return
-
-    for parent in Path(__file__).resolve().parents:
-        pip_out = parent / "pip-out"
-        if not pip_out.exists():
-            continue
-        matches = sorted(
-            pip_out.glob(
-                "temp.*/cmake-out/backends/mlx/mlx/mlx/backend/metal/kernels/mlx.metallib"
-            )
-        )
-        if not matches:
-            continue
-        shutil.copyfile(matches[0], metallib_path)
-        logger.info(f"Copied MLX metallib next to runtime library: {metallib_path}")
-        return
-
-
-def _ensure_mlx_backend_registered() -> Runtime:
-    runtime = Runtime.get()
-    if runtime.backend_registry.is_available("MLXBackend"):
-        return runtime
-
-    for candidate in _iter_mlx_backend_candidates():
-        if not candidate.is_file():
-            continue
-        try:
-            _ensure_mlx_metallib(candidate)
-            ctypes.CDLL(str(candidate), mode=ctypes.RTLD_GLOBAL)
-        except OSError as exc:
-            logger.info(f"Failed to load MLX backend library {candidate}: {exc}")
-            continue
-
-        runtime = Runtime.get()
-        if runtime.backend_registry.is_available("MLXBackend"):
-            logger.info(f"Loaded MLX backend runtime library: {candidate}")
-            return runtime
-
-    logger.warning(
-        "MLXBackend is not registered. If you built mlxdelegate locally, "
-        "set ET_MLX_BACKEND_DYLIB to the path of libmlxdelegate_runtime.dylib."
-    )
-    return runtime
-
-
 def _get_max_input_seq_len(program) -> int:
     """Inspect the .pte program metadata to determine the max input_ids seq len.
 
@@ -172,7 +101,7 @@ def run_inference(
     text_processor, uses_processor = _load_text_processor(model_id)
 
     logger.info(f"Loading model from {pte_path}...")
-    et_runtime = _ensure_mlx_backend_registered()
+    et_runtime = Runtime.get()
     program = et_runtime.load_program(pte_path, verification=Verification.Minimal)
 
     max_seq_len = _get_max_input_seq_len(program)

backends/mlx/runtime/MLXBackend.cpp

@@ -19,7 +19,6 @@
 #include <mlx/mlx.h>
 
 #include <cstring>
-#include <iostream>
 #include <limits>
 #include <memory>
 #include <mutex>
@@ -209,18 +208,14 @@ class MLXBackend final : public ::executorch::runtime::BackendInterface {
     }
 
     try {
-      std::cerr << "MLX init: constructing handle" << std::endl;
       new (handle) MLXHandle();
-      std::cerr << "MLX init: handle constructed" << std::endl;
 
       if (!processed || !processed->data() || processed->size() == 0) {
         throw std::runtime_error("init: null or empty delegate payload");
       }
 
-      std::cerr << "MLX init: parsing delegate payload" << std::endl;
       handle->program = loader::load_program(
           static_cast<const uint8_t*>(processed->data()), processed->size());
-      std::cerr << "MLX init: delegate payload parsed" << std::endl;
 
       // Validate schema version
       int schema_version = 1;
@@ -248,34 +243,27 @@ class MLXBackend final : public ::executorch::runtime::BackendInterface {
       // runtime
       const runtime::NamedDataMap* named_data_map =
           context.get_named_data_map();
-      std::cerr << "MLX init: loading constants" << std::endl;
       load_constants(
           handle->program,
           named_data_map,
           handle->constants,
           handle->constant_buffers);
-      std::cerr << "MLX init: constants loaded" << std::endl;
 
       // Delegate payload no longer needed after constants are loaded
       processed->Free();
       processed = nullptr;
 
       // Load mutable buffers (e.g., KV cache)
-      std::cerr << "MLX init: loading mutable buffers" << std::endl;
       load_mutable_buffers(handle->program, handle->mutable_buffers);
-      std::cerr << "MLX init: mutable buffers loaded" << std::endl;
 
       // Bind execution state (reused across execute() calls)
-      std::cerr << "MLX init: binding execution state" << std::endl;
       handle->state.bind(
           handle->program, handle->constants, handle->mutable_buffers);
-      std::cerr << "MLX init: execution state bound" << std::endl;
 
       // Run init chain if present.
       // SAFETY: The >= 0 check ensures init_chain_idx is non-negative, so the
       // static_cast<uint32_t> cannot produce UINT32_MAX from a -1 sentinel.
       if (handle->program.init_chain_idx >= 0) {
-        std::cerr << "MLX init: running init chain" << std::endl;
         handle->state.is_init_chain = true;
         handle->interpreter.run_chain(
             handle->program,
@@ -287,21 +275,17 @@ class MLXBackend final : public ::executorch::runtime::BackendInterface {
         // Evaluate any constants written by the init chain so the first
         // execute() doesn't pay the cost of materializing them.
         eval(handle->constants.tensors);
-        std::cerr << "MLX init: init chain complete" << std::endl;
       }
 
     } catch (const std::exception& e) {
       ET_LOG(Error, "Failed to load MLX program: %s", e.what());
-      std::cerr << "Failed to load MLX program: " << e.what() << std::endl;
       handle->~MLXHandle();
       if (processed != nullptr) {
         processed->Free();
       }
       return Error::InvalidProgram;
     } catch (...) {
       ET_LOG(Error, "Failed to load MLX program: unknown non-std exception");
-      std::cerr << "Failed to load MLX program: unknown non-std exception"
-                << std::endl;
       handle->~MLXHandle();
       if (processed != nullptr) {
         processed->Free();
@@ -437,12 +421,9 @@ class MLXBackend final : public ::executorch::runtime::BackendInterface {
       return Error::Ok;
     } catch (const std::exception& e) {
       ET_LOG(Error, "MLX execute failed: %s", e.what());
-      std::cerr << "MLX execute failed: " << e.what() << std::endl;
       return Error::Internal;
     } catch (...) {
       ET_LOG(Error, "MLX execute failed: unknown non-std exception");
-      std::cerr << "MLX execute failed: unknown non-std exception"
-                << std::endl;
       return Error::Internal;
     }
   }