feat: add spark client module for SparkConnect sessions#51
feat: add spark client module for SparkConnect sessions#51vedhakoushik wants to merge 5 commits into
Conversation
|
[APPROVALNOTIFIER] This PR is NOT APPROVED This pull-request has been approved by: The full list of commands accepted by this bot can be found here. DetailsNeeds approval from an approver in each of these files:Approvers can indicate their approval by writing |
c43dd76 to
5b1352b
Compare
|
🎉 Welcome to the Kubeflow MCP Server! 🎉 Thanks for opening your first PR! We're happy to have you as part of our community 🚀 Here's what happens next:
Join the community:
Feel free to ask questions in the comments if you need any help or clarification! |
There was a problem hiding this comment.
Pull request overview
Adds a new spark client module to the Kubeflow MCP server, exposing SparkConnect session lifecycle + monitoring tools via the same modular architecture used by the existing trainer client.
Changes:
- Introduces
kubeflow_mcp.sparkwith 5 MCP tools (list/get/logs/create/delete) plus resources and serialization helpers. - Wires the spark module into server registration, tool phases, personas, and destructive-tool policy.
- Adds a comprehensive unit-test suite for the spark module and adjusts an existing trainer architecture test to be module-scoped.
Reviewed changes
Copilot reviewed 17 out of 18 changed files in this pull request and generated 2 comments.
Show a summary per file
| File | Description |
|---|---|
| tests/unit/trainer/test_architecture.py | Scopes phase-to-section assertions to phases actually used by trainer tools. |
| tests/unit/spark/test_spark_module.py | New unit tests covering spark tool metadata, wiring, and mocked behavior. |
| tests/unit/spark/init.py | Initializes spark unit test package. |
| README.md | Documents spark as a selectable client module and updates CLI example. |
| kubeflow_mcp/spark/types/init.py | Adds JSON-safe serialization helper for SparkConnectInfo-like objects. |
| kubeflow_mcp/spark/resources/troubleshooting.md | Adds spark-specific troubleshooting guide. |
| kubeflow_mcp/spark/resources/session-patterns.md | Adds agent-facing usage patterns for SparkConnect sessions. |
| kubeflow_mcp/spark/api/sessions.py | Implements create/delete tools with confirm-gating and validation. |
| kubeflow_mcp/spark/api/monitoring.py | Implements bounded log retrieval tool with truncation behavior. |
| kubeflow_mcp/spark/api/discovery.py | Implements list/get discovery tools and “next steps” guidance. |
| kubeflow_mcp/spark/api/init.py | Declares spark API package. |
| kubeflow_mcp/spark/init.py | Defines spark module metadata, tool registration, annotations, resources, and instruction sections. |
| kubeflow_mcp/core/server.py | Registers spark in CLIENT_MODULES. |
| kubeflow_mcp/core/policy.py | Adds spark tools to persona policies and destructive-tool set. |
| kubeflow_mcp/common/utils.py | Adds lazy SparkClient factories with per-namespace caching. |
| kubeflow_mcp/common/constants.py | Adds spark phases/tools mapping and marks spark capability as implemented. |
| kubeflow_mcp/cli.py | Updates --clients help text to include spark. |
| ARCHITECTURE.md | Updates architecture matrix to reflect spark module availability. |
💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.
| def get_spark_client_for_namespace(namespace: str | None = None) -> Any: | ||
| """Return a SparkClient targeting the given namespace. | ||
|
|
||
| When *namespace* is ``None`` the shared singleton (default kubeconfig | ||
| namespace) is returned. When a namespace is explicitly provided a cached | ||
| ``SparkClient`` scoped to that namespace is returned. | ||
| """ | ||
| if namespace is None: | ||
| return get_spark_client() | ||
| with _spark_ns_client_lock: | ||
| if namespace in _spark_ns_client_cache: | ||
| return _spark_ns_client_cache[namespace] | ||
| spark_client_cls = _import_spark_client() | ||
| from kubeflow.common.types import KubernetesBackendConfig | ||
|
|
||
| client = spark_client_cls(backend_config=KubernetesBackendConfig(namespace=namespace)) | ||
| if len(_spark_ns_client_cache) >= _NS_CLIENT_CACHE_MAX: | ||
| oldest = next(iter(_spark_ns_client_cache)) | ||
| del _spark_ns_client_cache[oldest] | ||
| _spark_ns_client_cache[namespace] = client | ||
| return client |
There was a problem hiding this comment.
Fixed in 7d93f9e. get_spark_client_for_namespace(None) now resolves the effective namespace via the same resolver check_namespace_allowed uses (get_trainer_effective_namespace) and scopes the SparkClient to it, so the policy-checked namespace always matches the one Spark operates in — falling back to the unscoped singleton only if resolution fails. Added two unit tests covering the resolve and fallback paths.
| kubeflow-mcp serve \ | ||
| --clients trainer \ # modules: trainer, optimizer (stub), hub (stub) | ||
| --clients trainer,spark \ # modules: trainer, spark, optimizer (stub), hub (stub) | ||
| --persona ml-engineer \ # readonly | data-scientist | ml-engineer | platform-admin | ||
| --mode full \ # full | progressive | semantic | ||
| --instruction-tier full \ # full | compact | minimal |
There was a problem hiding this comment.
Good catch — fixed in 7d93f9e. The kubeflow-mcp serve example is now a clean, copy/paste-safe command (no trailing \ followed by inline # comments), with the per-flag options moved into a table right below it.
|
Thanks for raising the PR @vedhakoushik! Regarding the implementation I had two questions @andreyvelich @abhijeet-dhumal.
|
|
Thanks for the Spark client work. Before review, please link a design doc or KEP reference (Phase 6 / #5) describing the tool surface and scope. XXL PRs need an accepted design per CONTRIBUTING.md — happy to review once that's linked. /hold |
|
Thanks @abhijeet-dhumal! I've added a design doc describing the tool surface and scope, committed to this PR and linked from
It references the SDK design (KEP-107 / kubeflow/sdk#107), enumerates the five SparkConnect tools and their SDK mapping, and records the design decisions (lazy Happy to adjust anything, or reshape it into a KEP-style proposal if you'd prefer that format. Ready for review whenever you get a chance — thanks! |
@vedhakoushik I would suggest to start with a proposal like this : #48 |
|
Thanks @abhijeet-dhumal, that's really helpful — I went through #48 / KEP-34. Since the Spark work here is scoped to SparkConnect and already lines up with the project's compatibility baseline ( Before I push, two quick questions so I match your preference:
Scope stays SparkConnect-only; Spark batch jobs (SparkJob) are a follow-up once kubeflow/sdk#521 lands. Happy to adjust anything — thanks for the guidance! |
|
Opened the proposal as a separate PR, mirroring #48: #52 — KEP-5: SparkConnect (Spark) Client Support. It follows the KEP-34 structure and includes the SDK-compatibility metrics ( This PR (#51) is the reference implementation linked from the proposal. |
All four verified against the actual code/SDK before fixing: - Serialize the released SDK field name, not main. Downloaded and inspected the real kubeflow==0.4.0 and 0.4.1 wheels from PyPI: SparkConnectInfo names the driver-pod field `pod_name` on both released versions; `driver_pod_name` only exists on unreleased SDK main (I had read GitHub's default branch, not a tagged release). Design now specifies pod_name first, driver_pod_name fallback, with a matching test. This also means the reference implementation on kubeflow#51 has a real bug (getattr(..., "driver_pod_name", None) silently returns None against the declared >=0.4.0 baseline) that needs a follow-up fix there. - Ownership guard on delete_spark_session. Confirmed the trainer module's precedent (is_mcp_managed() / MCP_MANAGED_LABEL in common/utils.py, enforced in trainer/api/lifecycle.py, label attached at creation in trainer/api/training.py). The Spark design had no equivalent, so any persona with delete access could remove any SparkConnect session in an allowed namespace, not just MCP-created ones. Design now mirrors the trainer pattern. - get_session_logs() log-bounding is client-side only. Confirmed via the SDK's backend.py: read_namespaced_pod_log() is called without a server-side tail limit, so the full log is materialized before our tail_lines truncation applies. Reworded tool descriptions to stop overclaiming "bounded" and added a mitigation (bypass the SDK wrapper for a true server-side tail, matching trainer's existing direct-CoreV1Api precedent in discovery.py). - Out-of-cluster port-forward leak on create_spark_session. Confirmed via the SDK's KubernetesBackend.connect(): outside the cluster it spawns a kubectl port-forward subprocess that SparkSession.stop() does not terminate. Documented as a known limitation with a recommended mitigation (in-cluster deployment). Updated: SDK Compatibility table + note, Goals, Key Design Decisions (2 new bullets + 1 extended), Non-Goals, MCP Tools tables, Persona Coverage, SDK Compatibility Update mapping, Risks and Mitigations (+3 rows), Testing Plan. Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com> Signed-off-by: Bellam vedha kousik <bellam.vedhakoushik@gmail.com>
|
/assign @vedhakoushik |
Implement the `spark` client module wrapping kubeflow.spark.SparkClient so AI agents can manage SparkConnect sessions over MCP, mirroring the trainer module's architecture. Tools (5): - list_spark_sessions / get_spark_session / get_spark_session_logs (read-only) - create_spark_session / delete_spark_session (write, two-phase confirm gate) Design notes: - Only the public, JSON-serializable SparkClient methods are wrapped. All kubeflow.spark / pyspark imports are lazy, so the module and its metadata load without the optional `kubeflow[spark]` extra. - create_spark_session uses the SDK's only public creation path (connect()), releases the transient driver SparkSession, and returns session metadata; the data plane attaches with PySpark. Documented in resources + Known Limitations (attach-to-existing and follow=True streaming are out of scope). - New spark_* tool phases; personas and DESTRUCTIVE_TOOLS updated; module registered in CLIENT_MODULES and selectable via `--clients trainer,spark`. - Agent-facing guides added under spark/resources (session-patterns, troubleshooting); ARCHITECTURE.md / README updated. Testing: - 33 unit tests (metadata, wiring, serialization, mocked-client behavior). - Updated trainer test_all_phases_mapped: each client module now owns/maps its own phases, so the invariant is scoped to trainer's own tools. - Full unit suite (73 tests) + ruff check/format pass; server loads and registers all 5 tools in full and progressive modes. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com> Signed-off-by: Bellam vedha kousik <bellam.vedhakoushik@gmail.com>
- README: make the `kubeflow-mcp serve` example copy/paste-safe. The prior block ended each line with `\` followed by spaces and an inline `#` comment, so the backslash did not escape the newline and the command broke when pasted. Replace inline comments with a clean runnable command + an options table. - common/utils: when namespace is None, resolve the effective namespace via the same resolver check_namespace_allowed uses (get_trainer_effective_namespace) and scope the SparkClient to it, so the policy-checked namespace always matches the one Spark operates in (closes a potential namespace-allowlist bypass). Falls back to the unscoped singleton if resolution fails. - Add 2 unit tests for the factory namespace-resolution behavior (75 total). Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com> Signed-off-by: Bellam vedha kousik <bellam.vedhakoushik@gmail.com>
Add a design doc describing the spark module's tool surface and scope, as requested in review before merging an XXL PR. Covers the five SparkConnect tools, the SDK mapping (KEP-107), design decisions (lazy imports, create-via- connect semantics, namespace safety), testing, and — importantly — the scope split: this PR is SparkConnect-only; Spark batch jobs (SparkJob) are a follow-up gated on kubeflow/sdk#521. Linked from ARCHITECTURE.md. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com> Signed-off-by: Bellam vedha kousik <bellam.vedhakoushik@gmail.com>
Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com> Signed-off-by: Bellam vedha kousik <bellam.vedhakoushik@gmail.com>
Two Step-A fixes ahead of the KEP-5 design work landing: - kubeflow_mcp/spark/types/__init__.py: session_info_to_dict() read driver_pod_name via getattr, which silently returns None against the declared kubeflow[spark]>=0.4.0 baseline — the released 0.4.0/0.4.1 wheels name the field pod_name; driver_pod_name only exists on unreleased SDK main. Read pod_name first, fall back to driver_pod_name, so session info is correct today and keeps working if the rename ships in a future release. Added regression tests for both field shapes. - Removed docs/design/spark-client.md, superseded by the KEP-5 design proposal (kubeflow#52, docs/proposals/KEP-5.md in that PR). Repointed ARCHITECTURE.md's Spark row to kubeflow#52. Rebased onto upstream/main to resolve the README.md merge conflict (upstream added --otel-endpoint to the serve example; kept our copy/paste-safe table format and folded the new flag into it). Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com> Signed-off-by: Bellam vedha kousik <bellam.vedhakoushik@gmail.com>
973a53f to
bace2f4
Compare
|
Thanks for the assignment! Pushed an update to clear the merge conflict and fix a bug the KEP-5 review surfaced:
Still on hold pending #52. Once that's accepted I'll bring the remaining KEP-5 commitments (ownership guard on |
|
Hi @abhijeet-dhumal — Two quick questions as a first-time contributor here:
Happy to go either way — just want to follow the right process. Thanks! |
What
Implements the
sparkclient module (Closes #5), wrappingkubeflow.spark.SparkClientso AI agents can manage SparkConnect sessions over MCP. Structure mirrors the existingtrainermodule (api/,types/,resources/).Tools (5)
list_spark_sessionsSparkClient.list_sessions()get_spark_sessionSparkClient.get_session()get_spark_session_logsSparkClient.get_session_logs()create_spark_sessionSparkClient.connect()(create mode)delete_spark_sessionSparkClient.delete_session()Both write tools use the repo's two-phase confirm gate (
confirmed=Falsepreviews,confirmed=Trueexecutes).Design notes / decisions for reviewers
SparkClientmethods are wrapped. Allkubeflow.spark/pysparkimports are lazy, so the module and its metadata load even without the optionalkubeflow[spark]extra (missing extra → friendlySDK_ERRORtelling the user to install it). This keeps--clients sparkimportable on any host.create_spark_sessionuses the SDK's only public creation path,connect(), which returns a live pysparkSparkSession. The tool releases that transient session (.stop()) and returns session metadata; the data plane attaches with PySpark using the returned connect info — the MCP server never proxies Spark RPCs. Ifconnect()fails after the CR is created (e.g. the host can't reach the driver), the tool still reports the provisioned session with a warning rather than losing it. A provision-only SDK method would remove the host-reachability requirement — noted as a follow-up inspark://guides/troubleshooting.connect(base_url=...)returns a non-serializable live session) andfollow=Truelog streaming.SparkApplicationbatch jobs are not part of the SDK'sSparkClienttoday.spark_*tool phases; personas updated (readonly → read tools, data-scientist → lifecycle);delete_spark_sessionadded toDESTRUCTIVE_TOOLS; module registered inCLIENT_MODULESand selectable via--clients trainer,spark.tests/unit/trainer/test_architecture.py::test_all_phases_mapped. It previously asserted the trainer module'sPHASE_TO_SECTIONmaps every global phase, an invariant only valid while trainer was the sole client module. The server already merges each module'sPHASE_TO_SECTION, so I scoped the test to trainer's own phases. Flagging explicitly since it touches an existing test.Testing
tests/unit/spark/): metadata/annotation consistency, phase + persona + destructive wiring,SparkConnectInfoserialization (enum/datetime), and mocked-client behavior for all 5 tools (state filter, not-found →RESOURCE_NOT_FOUND, log tailing/truncation, confirm-gate previews, validation).ruff checkandruff formatclean.trainer+sparkand registers all 5 tools in bothfullandprogressivemodes.Docs
ARCHITECTURE.mdandREADME.mdupdated; two agent-facing guides added underspark/resources/(session-patterns.md,troubleshooting.md).🤖 Authored with Claude Code. First-time contributor — happy to adjust scope, the creation approach, or naming to match maintainer preferences.