databricks-agent-notebooks 0.2.2__tar.gz → 0.4.0__tar.gz

{databricks_agent_notebooks-0.2.2 → databricks_agent_notebooks-0.4.0}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: databricks-agent-notebooks
-Version: 0.2.2
+Version: 0.4.0
 Summary: Standalone notebook conversion and local runtime tooling for the databricks-agent-notebooks repository.
 Author: Swoop
 License-Expression: MIT

{databricks_agent_notebooks-0.2.2 → databricks_agent_notebooks-0.4.0}/pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta"
 
 [project]
 name = "databricks-agent-notebooks"
-version = "0.2.2"
+version = "0.4.0"
 description = "Standalone notebook conversion and local runtime tooling for the databricks-agent-notebooks repository."
 readme = "README.md"
 requires-python = ">=3.11"
@@ -52,6 +52,13 @@ agent-notebook = "databricks_agent_notebooks.cli:main"
 testpaths = ["tests"]
 addopts = "-q"
 pythonpath = ["src"]
+markers = [
+    "integration: Fast integration tests (CLI smoke, no side effects)",
+    "slow: Slow tests (kernel downloads, package installs, real execution)",
+    "spark: Tests requiring local Spark execution (PySpark or Almond kernel)",
+    "databricks: Tests requiring live Databricks workspace credentials",
+    "scala: Tests requiring Almond/Scala kernel and Java",
+]
 
 [tool.setuptools]
 package-dir = { "" = "src" }

databricks_agent_notebooks-0.4.0/src/databricks_agent_notebooks/__init__.py

@@ -0,0 +1,10 @@
+"""Standalone Databricks notebook conversion and execution tooling."""
+
+from importlib.metadata import version, PackageNotFoundError
+
+try:
+    __version__ = version("databricks-agent-notebooks")
+except PackageNotFoundError:
+    __version__ = "0.0.0-dev"
+
+__all__ = ["__version__"]

{databricks_agent_notebooks-0.2.2 → databricks_agent_notebooks-0.4.0}/src/databricks_agent_notebooks/_constants.py

@@ -103,3 +103,12 @@ def scala_variant_for_dbr(major: int) -> ScalaVariant:
     if major >= 17:
         return SCALA_213
     return SCALA_212
+
+
+# ---------------------------------------------------------------------------
+# LOCAL_SPARK — reserved profile for local PySpark / vanilla Spark execution
+# ---------------------------------------------------------------------------
+
+LOCAL_SPARK_PROFILE = "LOCAL_SPARK"
+LOCAL_SPARK_DEFAULT_MASTER = "local[*]"
+LOCAL_SPARK_DEFAULT_VERSION = "3.5.4"

{databricks_agent_notebooks-0.2.2 → databricks_agent_notebooks-0.4.0}/src/databricks_agent_notebooks/cli.py

@@ -8,6 +8,7 @@ same argparse/subparsers/dispatch-table pattern as ``libs.continuum.cli``.
 from __future__ import annotations
 
 import argparse
+import re
 import shutil
 import subprocess
 import sys
@@ -15,7 +16,8 @@ import tempfile
 from importlib import resources
 from pathlib import Path
 
-from databricks_agent_notebooks.config.frontmatter import DatabricksConfig, merge_config
+from databricks_agent_notebooks import __version__
+from databricks_agent_notebooks.config.frontmatter import DatabricksConfig, is_local_spark, merge_config
 from databricks_agent_notebooks.execution.executor import RawProgressValue, emit_progress_signal, execute_notebook
 from databricks_agent_notebooks.execution.injection import inject_cells
 from databricks_agent_notebooks.execution.rendering import render
@@ -23,7 +25,14 @@ from databricks_agent_notebooks.formats.conversion import to_notebook, validate_
 from databricks_agent_notebooks.integrations.databricks.clusters import ClusterError, default_service
 from databricks_agent_notebooks.runtime.connect import ensure_cluster_runtime, ensure_serverless_runtime
 from databricks_agent_notebooks.runtime.home import resolve_runtime_home
-from databricks_agent_notebooks._constants import DEFAULT_SCALA_VARIANT, SCALA_212, SCALA_VARIANTS
+import os
+
+from databricks_agent_notebooks._constants import (
+    DEFAULT_SCALA_VARIANT,
+    LOCAL_SPARK_DEFAULT_VERSION,
+    SCALA_212,
+    SCALA_VARIANTS,
+)
 from databricks_agent_notebooks.runtime.scala_connect import prefetch_scala_connect, resolve_scala_connect
 from databricks_agent_notebooks.runtime.doctor import Check, doctor_scala_connect_readiness, run_checks
 from databricks_agent_notebooks.runtime.inventory import doctor_installed_runtimes, list_installed_runtimes
@@ -39,6 +48,30 @@ from databricks_agent_notebooks.runtime.kernel import (
 
 import nbformat
 
+_VALID_SCALA_LOCAL_MASTER_RE = re.compile(r"^local(\[(\*|\d+)(,\d+)?\])?$")
+
+
+def _validate_scala_local_spark(master: str, executor_memory: str | None) -> str | None:
+    """Validate LOCAL_SPARK configuration for Scala notebooks.
+
+    Returns an error message string if validation fails, or None if valid.
+    """
+    if not _VALID_SCALA_LOCAL_MASTER_RE.match(master):
+        return (
+            f"error: Spark master URL '{master}' is not supported for Scala notebooks. "
+            "Scala LOCAL_SPARK only supports local[*], local[N], local[N,M], or local. "
+            "The local-cluster mode is incompatible with the Almond kernel's classloader "
+            "(ClassNotFoundException in executors). "
+            "Use local[N] for parallelism control instead."
+        )
+    if executor_memory:
+        return (
+            "error: AGENT_NOTEBOOK_LOCAL_SPARK_EXECUTOR_MEMORY is not supported for Scala notebooks. "
+            "In local mode there are no separate executor processes — the driver IS the executor. "
+            "Use AGENT_NOTEBOOK_LOCAL_SPARK_DRIVER_MEMORY to control the total JVM heap size instead."
+        )
+    return None
+
 
 # ---------------------------------------------------------------------------
 # Parser
@@ -75,6 +108,10 @@ def _build_parser() -> argparse.ArgumentParser:
         epilog=_agent_docs_epilog(),
         formatter_class=argparse.RawDescriptionHelpFormatter,
     )
+    parser.add_argument(
+        "--version", action="version",
+        version=f"%(prog)s {__version__}",
+    )
     subparsers = parser.add_subparsers(dest="command")
 
     # -- run --
@@ -207,9 +244,85 @@ def _cmd_run(args: argparse.Namespace) -> int:
     language = _resolve_execution_language(notebook, config)
     inject_session = not args.no_inject_session
 
+    # Step 2b: LOCAL_SPARK validation and branching
+    local_spark = is_local_spark(config)
+    if local_spark and config.cluster:
+        print("error: --profile LOCAL_SPARK and --cluster are mutually exclusive", file=sys.stderr)
+        return 1
+
     # Step 3: Resolve cluster
     cluster = None
-    if config.cluster and inject_session:
+    if local_spark and inject_session:
+        emit_progress_signal("compute", mode=RawProgressValue("local-spark"))
+        # Pre-flight: verify pyspark is importable for Python LOCAL_SPARK
+        if language != "scala":
+            try:
+                import pyspark  # noqa: F401, PLC0415
+            except ImportError:
+                print(
+                    "error: pyspark is required for Python LOCAL_SPARK but not importable. "
+                    "Install it with: pip install pyspark (or uv pip install pyspark)",
+                    file=sys.stderr,
+                )
+                return 1
+        # Validate and set up Scala LOCAL_SPARK
+        if language == "scala":
+            local_spark_master = os.environ.get("AGENT_NOTEBOOK_LOCAL_SPARK_MASTER", "local[*]")
+            executor_memory = os.environ.get("AGENT_NOTEBOOK_LOCAL_SPARK_EXECUTOR_MEMORY")
+            validation_error = _validate_scala_local_spark(local_spark_master, executor_memory)
+            if validation_error:
+                print(validation_error, file=sys.stderr)
+                return 1
+            spark_ver = os.environ.get(
+                "AGENT_NOTEBOOK_LOCAL_SPARK_VERSION", LOCAL_SPARK_DEFAULT_VERSION,
+            )
+            scala_variant = SCALA_VARIANTS.get(
+                "2.13" if spark_ver.startswith("4.") else "2.12", SCALA_212,
+            )
+            notebook.metadata["kernelspec"] = {
+                "name": scala_variant.kernel_id,
+                "display_name": scala_variant.kernel_display_name,
+                "language": "scala",
+            }
+            # Spark 3.x on Java 17+ needs --add-opens for the Almond kernel JVM.
+            # PySpark handles this in its own launcher; for Scala we propagate via
+            # JDK_JAVA_OPTIONS (standard since Java 9). Unlike JAVA_TOOL_OPTIONS,
+            # JDK_JAVA_OPTIONS does not print "Picked up ..." to stderr.
+            existing_jdk = os.environ.get("JDK_JAVA_OPTIONS", "")
+            existing_jto = os.environ.get("JAVA_TOOL_OPTIONS", "")
+            if spark_ver.startswith("3.") and "--add-opens" not in existing_jdk and "--add-opens" not in existing_jto:
+                _opens = (
+                    "--add-opens=java.base/java.lang=ALL-UNNAMED "
+                    "--add-opens=java.base/java.lang.invoke=ALL-UNNAMED "
+                    "--add-opens=java.base/java.lang.reflect=ALL-UNNAMED "
+                    "--add-opens=java.base/java.io=ALL-UNNAMED "
+                    "--add-opens=java.base/java.net=ALL-UNNAMED "
+                    "--add-opens=java.base/java.nio=ALL-UNNAMED "
+                    "--add-opens=java.base/java.util=ALL-UNNAMED "
+                    "--add-opens=java.base/java.util.concurrent=ALL-UNNAMED "
+                    "--add-opens=java.base/java.util.concurrent.atomic=ALL-UNNAMED "
+                    "--add-opens=java.base/sun.nio.ch=ALL-UNNAMED "
+                    "--add-opens=java.base/sun.nio.cs=ALL-UNNAMED "
+                    "--add-opens=java.base/sun.security.action=ALL-UNNAMED "
+                    "--add-opens=java.base/sun.util.calendar=ALL-UNNAMED "
+                    "--add-opens=java.security.jgss/sun.security.krb5=ALL-UNNAMED"
+                )
+                os.environ["JDK_JAVA_OPTIONS"] = f"{existing_jdk} {_opens}".strip()
+            # Inject driver memory as -Xmx into JDK_JAVA_OPTIONS for Scala.
+            # The SparkConf injection (in injection.py) generates
+            # spark.driver.memory but that is decorative for Scala — the JVM
+            # is already running by the time SparkSession.builder reads it.
+            existing_jdk = os.environ.get("JDK_JAVA_OPTIONS", "")
+            driver_memory = os.environ.get("AGENT_NOTEBOOK_LOCAL_SPARK_DRIVER_MEMORY")
+            if driver_memory and "-Xmx" not in existing_jdk and "-Xmx" not in existing_jto:
+                os.environ["JDK_JAVA_OPTIONS"] = f"{existing_jdk} -Xmx{driver_memory}".strip()
+                print(
+                    f"warning: In local mode, spark.driver.memory controls the total JVM heap "
+                    f"shared by driver and all task threads — there are no separate executor "
+                    f"processes. -Xmx{driver_memory} has been set on the JVM.",
+                    file=sys.stderr,
+                )
+    elif config.cluster and inject_session:
         service = default_service()
         try:
             effective_profile = config.profile or "DEFAULT"
@@ -224,7 +337,7 @@ def _cmd_run(args: argparse.Namespace) -> int:
             emit_progress_signal("failed", error=str(exc))
             print(f"error: {exc}", file=sys.stderr)
             return 1
-    elif not config.cluster:
+    elif not config.cluster and not local_spark:
         emit_progress_signal("compute", mode=RawProgressValue("serverless"))
 
     managed_python_executable = None
@@ -252,7 +365,7 @@ def _cmd_run(args: argparse.Namespace) -> int:
                 "display_name": scala_variant.kernel_display_name,
                 "language": "scala",
             }
-    elif cluster is None and inject_session and language == "python":
+    elif cluster is None and inject_session and not local_spark and language == "python":
         try:
             managed_runtime = ensure_serverless_runtime(
                 profile=config.profile,
@@ -263,7 +376,7 @@ def _cmd_run(args: argparse.Namespace) -> int:
             print(f"error: {exc}", file=sys.stderr)
             return 1
         managed_python_executable = managed_runtime.python_executable
-    elif cluster is None and inject_session and language == "scala":
+    elif cluster is None and inject_session and not local_spark and language == "scala":
         scala_variant = DEFAULT_SCALA_VARIANT
         notebook.metadata["kernelspec"] = {
             "name": scala_variant.kernel_id,
@@ -275,6 +388,7 @@ def _cmd_run(args: argparse.Namespace) -> int:
     if inject_session:
         notebook = inject_cells(
             notebook, config, input_path,
+            local_spark=local_spark,
             scala_connect_version=scala_connect_version,
             scala_variant=scala_variant,
         )
@@ -565,18 +679,6 @@ def _print_doctor_checks(checks: list[Check]) -> list[Check]:
     return [check for check in checks if check.status == "fail"]
 
 
-def _cmd_kernels_doctor(args: argparse.Namespace) -> int:
-    """Run environment validation checks."""
-    checks = _run_kernels_doctor_checks(args)
-    failures = _print_doctor_checks(checks)
-    if failures:
-        print(f"\n{len(failures)} check(s) failed.", file=sys.stderr)
-        return 1
-
-    print("\nAll checks passed.")
-    return 0
-
-
 def _cmd_kernels(args: argparse.Namespace) -> int:
     handler = _KERNEL_HANDLERS.get(args.kernels_command)
     if handler is None:
@@ -601,18 +703,6 @@ def _cmd_runtimes_list(_args: argparse.Namespace) -> int:
     return 0
 
 
-def _cmd_runtimes_doctor(_args: argparse.Namespace) -> int:
-    """Validate managed runtime receipts rooted under runtime-home."""
-    checks = _run_runtimes_doctor_checks()
-    failures = _print_doctor_checks(checks)
-    if failures:
-        print(f"\n{len(failures)} check(s) failed.", file=sys.stderr)
-        return 1
-
-    print("\nAll checks passed.")
-    return 0
-
-
 def _cmd_runtimes(args: argparse.Namespace) -> int:
     handler = _RUNTIME_HANDLERS.get(args.runtimes_command)
     if handler is None:
@@ -640,12 +730,10 @@ _KERNEL_HANDLERS = {
     "install": _cmd_kernels_install,
     "list": _cmd_kernels_list,
     "remove": _cmd_kernels_remove,
-    "doctor": _cmd_kernels_doctor,
 }
 
 _RUNTIME_HANDLERS = {
     "list": _cmd_runtimes_list,
-    "doctor": _cmd_runtimes_doctor,
 }
 
 

{databricks_agent_notebooks-0.2.2 → databricks_agent_notebooks-0.4.0}/src/databricks_agent_notebooks/config/frontmatter.py

@@ -13,6 +13,8 @@ from typing import Any
 
 import yaml
 
+from databricks_agent_notebooks._constants import LOCAL_SPARK_PROFILE
+
 
 @dataclass(frozen=True)
 class DatabricksConfig:
@@ -82,3 +84,12 @@ def merge_config(
         cluster=cli_cluster if cli_cluster is not None else frontmatter.cluster,
         language=cli_language if cli_language is not None else frontmatter.language,
     )
+
+
+def is_local_spark(config: DatabricksConfig) -> bool:
+    """Detect whether *config* specifies the reserved LOCAL_SPARK profile.
+
+    Case-insensitive: ``local_spark``, ``Local_Spark``, ``LOCAL_SPARK``
+    all match.
+    """
+    return config.profile is not None and config.profile.upper() == LOCAL_SPARK_PROFILE

{databricks_agent_notebooks-0.2.2 → databricks_agent_notebooks-0.4.0}/src/databricks_agent_notebooks/execution/executor.py

@@ -19,11 +19,14 @@ from collections.abc import Sequence
 from contextlib import contextmanager
 from dataclasses import dataclass
 from pathlib import Path
+from unittest.mock import patch
 
 from jupyter_client.kernelspec import KernelSpecManager, NoSuchKernel
 from nbconvert.preprocessors import ExecutePreprocessor
 
 import nbformat
+import nbclient.client as _nbclient_mod
+import nbformat.v4.nbbase as _nbbase
 
 from databricks_agent_notebooks.runtime.home import resolve_runtime_home
 
@@ -33,8 +36,7 @@ except ModuleNotFoundError:  # pragma: no cover - exercised via packaging verifi
     install_ipykernel = None
 
 
-HEARTBEAT_INTERVAL_SECONDS = 30.0
-REDACTED_CELL_SNIPPET = "[source redacted]"
+HEARTBEAT_INTERVAL_SECONDS = 60.0
 
 
 class RawProgressValue(str):
@@ -133,18 +135,11 @@ def _build_cell_label(cell: nbformat.NotebookNode, lines: list[str]) -> str:
     return f"[{cell_type} cell]"
 
 
-def _build_cell_snippet(_cell: nbformat.NotebookNode, lines: list[str]) -> str:
-    if not lines:
-        return "[empty cell]"
-    return REDACTED_CELL_SNIPPET
-
-
 def _describe_cell(cell: nbformat.NotebookNode, *, cell_index: int) -> dict[str, object]:
     lines = _iter_meaningful_lines(cell.source)
     return {
         "cell_index": cell_index + 1,
         "cell_label": _build_cell_label(cell, lines),
-        "cell_snippet": _build_cell_snippet(cell, lines),
     }
 
 
@@ -243,6 +238,36 @@ def _patched_execution_environment(runtime_data_dir: str):
             os.environ.pop("JUPYTER_PATH", None)
 
 
+@contextmanager
+def _robust_kernel_error_output():
+    """Ensure kernel error messages have ``traceback``, ``ename``, and ``evalue``.
+
+    Some Jupyter kernels (notably Almond/Scala) omit ``traceback`` from IOPub
+    ``error`` messages on compile errors.  ``nbformat.v4.nbbase.output_from_msg``
+    accesses these fields via direct dict lookup, so a missing key raises
+    ``KeyError("traceback")`` which propagates as a useless error string.
+
+    This context manager monkey-patches ``output_from_msg`` at its call site in
+    ``nbclient.client`` (where the bound reference lives) to default missing
+    fields before the original function sees them, keeping the normal
+    ``CellExecutionError`` path intact so callers get ename/evalue in the error
+    message.
+    """
+    _original = _nbclient_mod.output_from_msg
+
+    def _patched(msg):  # noqa: ANN001, ANN202
+        msg_type = msg.get("header", {}).get("msg_type") or msg.get("msg_type", "")
+        if msg_type == "error":
+            content = msg.get("content", {})
+            content.setdefault("traceback", [])
+            content.setdefault("ename", "UnknownError")
+            content.setdefault("evalue", "")
+        return _original(msg)
+
+    with patch.object(_nbclient_mod, "output_from_msg", _patched):
+        yield
+
+
 def _execute_notebook_local(
     notebook_path: Path,
     *,
@@ -302,7 +327,10 @@ def _execute_notebook_local(
     )
 
     try:
-        with _patched_execution_environment(str(runtime_home.kernels_dir.parent)):
+        with (
+            _patched_execution_environment(str(runtime_home.kernels_dir.parent)),
+            _robust_kernel_error_output(),
+        ):
             executed_notebook, _ = executor.preprocess(notebook, resources)
     except Exception as exc:
         return ExecutionResult(
@@ -366,7 +394,7 @@ def _execute_notebook_subprocess(
         result_path = Path(result_file.name)
 
     try:
-        completed = subprocess.run(
+        process = subprocess.Popen(
             _subprocess_command(
                 python_executable=python_executable,
                 notebook_path=notebook_path,
@@ -376,14 +404,30 @@ def _execute_notebook_subprocess(
                 timeout=timeout,
                 allow_errors=allow_errors,
             ),
-            capture_output=True,
+            stdout=subprocess.PIPE,
+            stderr=subprocess.PIPE,
             text=True,
-            check=False,
         )
-        if completed.stdout:
-            print(completed.stdout, end="", file=sys.stdout)
-        if completed.stderr:
-            print(completed.stderr, end="", file=sys.stderr)
+
+        stderr_lines: list[str] = []
+
+        def _stream_stderr() -> None:
+            assert process.stderr is not None
+            for line in process.stderr:
+                print(line, end="", file=sys.stderr, flush=True)
+                stderr_lines.append(line)
+
+        stderr_thread = threading.Thread(target=_stream_stderr, daemon=True)
+        stderr_thread.start()
+
+        stdout_text = process.stdout.read() if process.stdout else ""
+        process.wait()
+        stderr_thread.join()
+
+        if stdout_text:
+            print(stdout_text, end="", file=sys.stdout)
+
+        stderr_text = "".join(stderr_lines).strip()
 
         if result_path.is_file():
             payload = json.loads(result_path.read_text(encoding="utf-8"))
@@ -399,7 +443,7 @@ def _execute_notebook_subprocess(
             success=False,
             output_path=output_path,
             duration_seconds=time.monotonic() - start,
-            error=completed.stderr.strip() or f"Managed runtime execution failed with exit code {completed.returncode}",
+            error=stderr_text or f"Managed runtime execution failed with exit code {process.returncode}",
         )
     finally:
         result_path.unlink(missing_ok=True)
@@ -417,7 +461,7 @@ def execute_notebook(
     if output_path is None:
         output_path = notebook_path.with_suffix(".executed.ipynb")
 
-    if python_executable is not None and python_executable.resolve() != Path(sys.executable).resolve():
+    if python_executable is not None and str(python_executable.resolve().parent.parent) != sys.prefix:
         return _execute_notebook_subprocess(
             notebook_path,
             output_path=output_path,