spark-advisor-analyzer 0.1.6__tar.gz

spark_advisor_analyzer-0.1.6/.gitignore

@@ -0,0 +1,38 @@
+CLAUDE.md
+.claude/
+
+# Python
+__pycache__/
+*.py[cod]
+*.pyo
+*.egg-info/
+dist/
+build/
+
+# Virtual environment
+.venv/
+.envrc
+
+# Testing
+.coverage
+.pytest_cache/
+htmlcov/
+
+# Type checking
+.mypy_cache/
+
+# Ruff
+.ruff_cache/
+
+# IDE
+.idea/
+.vscode/
+*.swp
+*.swo
+*~
+
+# OS
+.DS_Store
+Thumbs.db
+/.claude/
+tasks

spark_advisor_analyzer-0.1.6/Dockerfile

@@ -0,0 +1,13 @@
+FROM python:3.12-slim AS base
+
+COPY --from=ghcr.io/astral-sh/uv:latest /uv /usr/local/bin/uv
+
+WORKDIR /app
+COPY pyproject.toml uv.lock ./
+COPY packages/spark-advisor-models/ packages/spark-advisor-models/
+COPY packages/spark-advisor-rules/ packages/spark-advisor-rules/
+COPY packages/spark-advisor-analyzer/ packages/spark-advisor-analyzer/
+
+RUN uv sync --frozen --no-dev --package spark-advisor-analyzer
+
+CMD ["uv", "run", "--package", "spark-advisor-analyzer", "spark-advisor-analyzer"]

spark_advisor_analyzer-0.1.6/PKG-INFO

@@ -0,0 +1,73 @@
+Metadata-Version: 2.4
+Name: spark-advisor-analyzer
+Version: 0.1.6
+Summary: AI-powered Spark job analyzer - NATS worker service
+Project-URL: Homepage, https://github.com/pstysz/spark-advisor
+Project-URL: Repository, https://github.com/pstysz/spark-advisor
+Project-URL: Issues, https://github.com/pstysz/spark-advisor/issues
+Author: Pawel Stysz
+License-Expression: Apache-2.0
+Keywords: ai,analyzer,apache-spark,claude,nats,performance,spark
+Classifier: Development Status :: 4 - Beta
+Classifier: Intended Audience :: Developers
+Classifier: License :: OSI Approved :: Apache Software License
+Classifier: Programming Language :: Python :: 3
+Classifier: Programming Language :: Python :: 3.12
+Classifier: Programming Language :: Python :: 3.13
+Classifier: Topic :: Software Development :: Quality Assurance
+Classifier: Typing :: Typed
+Requires-Python: >=3.12
+Requires-Dist: anthropic>=0.52
+Requires-Dist: faststream[nats]>=0.6
+Requires-Dist: pyyaml>=6.0
+Requires-Dist: spark-advisor-models==0.1.6
+Requires-Dist: spark-advisor-rules==0.1.6
+Description-Content-Type: text/markdown
+
+# spark-advisor-analyzer
+
+AI-powered Spark job analyzer — NATS worker service. Part of the [spark-advisor](https://github.com/pstysz/spark-advisor) ecosystem.
+
+## Install
+
+```bash
+pip install spark-advisor-analyzer
+```
+
+## What it does
+
+Combines deterministic rules engine with optional Claude AI analysis to provide actionable Spark configuration recommendations.
+
+- **Rules engine** — 11 expert rules detecting data skew, GC pressure, disk spill, and more
+- **AI analysis** — Claude API integration for prioritized recommendations with causal chains
+- **Agent mode** — multi-turn Claude tool_use loop for autonomous job exploration
+- **NATS worker** — FastStream subscriber for `analyze.request` and `analyze.agent.request`
+
+## Deployment
+
+As a standalone NATS worker:
+
+```bash
+export SA_ANALYZER_NATS__URL=nats://localhost:4222
+export ANTHROPIC_API_KEY=sk-ant-...
+spark-advisor-analyzer
+```
+
+Or use via `spark-advisor` CLI or `spark-advisor-mcp` server.
+
+## Configuration
+
+| Variable | Default | Description |
+|----------|---------|-------------|
+| `SA_ANALYZER_NATS__URL` | `nats://localhost:4222` | NATS broker URL |
+| `SA_ANALYZER_AI__ENABLED` | `true` | Enable AI analysis |
+| `ANTHROPIC_API_KEY` | — | Claude API key (required if AI enabled) |
+
+## Links
+
+- [Main project](https://github.com/pstysz/spark-advisor)
+- [Contributing](https://github.com/pstysz/spark-advisor/blob/main/CONTRIBUTING.md)
+
+## License
+
+Apache 2.0

spark_advisor_analyzer-0.1.6/README.md

@@ -0,0 +1,47 @@
+# spark-advisor-analyzer
+
+AI-powered Spark job analyzer — NATS worker service. Part of the [spark-advisor](https://github.com/pstysz/spark-advisor) ecosystem.
+
+## Install
+
+```bash
+pip install spark-advisor-analyzer
+```
+
+## What it does
+
+Combines deterministic rules engine with optional Claude AI analysis to provide actionable Spark configuration recommendations.
+
+- **Rules engine** — 11 expert rules detecting data skew, GC pressure, disk spill, and more
+- **AI analysis** — Claude API integration for prioritized recommendations with causal chains
+- **Agent mode** — multi-turn Claude tool_use loop for autonomous job exploration
+- **NATS worker** — FastStream subscriber for `analyze.request` and `analyze.agent.request`
+
+## Deployment
+
+As a standalone NATS worker:
+
+```bash
+export SA_ANALYZER_NATS__URL=nats://localhost:4222
+export ANTHROPIC_API_KEY=sk-ant-...
+spark-advisor-analyzer
+```
+
+Or use via `spark-advisor` CLI or `spark-advisor-mcp` server.
+
+## Configuration
+
+| Variable | Default | Description |
+|----------|---------|-------------|
+| `SA_ANALYZER_NATS__URL` | `nats://localhost:4222` | NATS broker URL |
+| `SA_ANALYZER_AI__ENABLED` | `true` | Enable AI analysis |
+| `ANTHROPIC_API_KEY` | — | Claude API key (required if AI enabled) |
+
+## Links
+
+- [Main project](https://github.com/pstysz/spark-advisor)
+- [Contributing](https://github.com/pstysz/spark-advisor/blob/main/CONTRIBUTING.md)
+
+## License
+
+Apache 2.0

spark_advisor_analyzer-0.1.6/pyproject.toml

@@ -0,0 +1,87 @@
+[project]
+name = "spark-advisor-analyzer"
+version = "0.1.6" # x-release-please-version
+description = "AI-powered Spark job analyzer - NATS worker service"
+readme = "README.md"
+license = "Apache-2.0"
+requires-python = ">=3.12"
+authors = [
+    { name = "Pawel Stysz" },
+]
+keywords = ["spark", "apache-spark", "ai", "claude", "nats", "performance", "analyzer"]
+classifiers = [
+    "Development Status :: 4 - Beta",
+    "Intended Audience :: Developers",
+    "License :: OSI Approved :: Apache Software License",
+    "Programming Language :: Python :: 3",
+    "Programming Language :: Python :: 3.12",
+    "Programming Language :: Python :: 3.13",
+    "Topic :: Software Development :: Quality Assurance",
+    "Typing :: Typed",
+]
+dependencies = [
+    "spark-advisor-models==0.1.6",  # x-release-please-version
+    "spark-advisor-rules==0.1.6",  # x-release-please-version
+    "anthropic>=0.52",
+    "faststream[nats]>=0.6",
+    "pyyaml>=6.0",
+]
+
+[project.urls]
+Homepage = "https://github.com/pstysz/spark-advisor"
+Repository = "https://github.com/pstysz/spark-advisor"
+Issues = "https://github.com/pstysz/spark-advisor/issues"
+
+[project.scripts]
+spark-advisor-analyzer = "spark_advisor_analyzer.app:main"
+
+[build-system]
+requires = ["hatchling"]
+build-backend = "hatchling.build"
+
+[tool.hatch.build.targets.wheel]
+packages = ["src/spark_advisor_analyzer"]
+
+[tool.uv.sources]
+spark-advisor-models = { workspace = true }
+spark-advisor-rules = { workspace = true }
+
+[dependency-groups]
+dev = [
+    "pytest>=8.3",
+    "pytest-asyncio>=1.0",
+    "pytest-cov>=6.1",
+    "mypy>=1.15",
+    "ruff>=0.11",
+]
+
+[tool.pytest.ini_options]
+testpaths = ["tests"]
+pythonpath = ["src", "tests"]
+asyncio_mode = "auto"
+
+[tool.ruff]
+target-version = "py312"
+line-length = 120
+src = ["src", "tests"]
+
+[tool.ruff.lint]
+select = ["E", "W", "F", "I", "UP", "B", "SIM", "TCH", "RUF"]
+
+[tool.ruff.lint.isort]
+known-first-party = ["spark_advisor_analyzer", "spark_advisor_models", "spark_advisor_rules"]
+
+[tool.ruff.lint.flake8-type-checking]
+runtime-evaluated-base-classes = ["pydantic.BaseModel", "pydantic_settings.BaseSettings"]
+
+[tool.mypy]
+python_version = "3.12"
+strict = true
+plugins = ["pydantic.mypy"]
+
+[tool.coverage.run]
+source = ["spark_advisor_analyzer"]
+branch = true
+
+[tool.coverage.report]
+show_missing = true

spark_advisor_analyzer-0.1.6/src/spark_advisor_analyzer/__init__.py

@@ -0,0 +1 @@
+"""AI-powered Spark job analyzer - NATS worker service."""

spark_advisor_analyzer-0.1.6/src/spark_advisor_analyzer/agent/__init__.py

@@ -0,0 +1,4 @@
+from spark_advisor_analyzer.agent.orchestrator import AgentOrchestrator
+from spark_advisor_analyzer.agent.tools import AgentToolName
+
+__all__ = ["AgentOrchestrator", "AgentToolName"]

spark_advisor_analyzer-0.1.6/src/spark_advisor_analyzer/agent/context.py

@@ -0,0 +1,29 @@
+from __future__ import annotations
+
+from dataclasses import dataclass, field
+from typing import TYPE_CHECKING
+
+if TYPE_CHECKING:
+    from spark_advisor_models.model import JobAnalysis, RuleResult
+    from spark_advisor_rules import StaticAnalysisService
+
+
+@dataclass
+class AgentContext:
+    job: JobAnalysis
+    _rule_results: list[RuleResult] = field(default_factory=list, init=False)
+    _rules_executed: bool = field(default=False, init=False)
+
+    @property
+    def rules_executed(self) -> bool:
+        return self._rules_executed
+
+    @property
+    def rule_results(self) -> list[RuleResult]:
+        return self._rule_results
+
+    def get_or_run_rules(self, static: StaticAnalysisService) -> list[RuleResult]:
+        if not self._rules_executed:
+            self._rule_results = static.analyze(self.job)
+            self._rules_executed = True
+        return self._rule_results

spark_advisor_analyzer-0.1.6/src/spark_advisor_analyzer/agent/handlers.py

@@ -0,0 +1,182 @@
+import json
+import math
+from typing import Any
+
+from pydantic import BaseModel, ValidationError
+
+from spark_advisor_analyzer.agent.context import AgentContext
+from spark_advisor_analyzer.agent.tools import (
+    AgentToolName,
+    CalculateOptimalPartitionsInput,
+    CompareConfigsInput,
+    GetStageDetailsInput,
+)
+from spark_advisor_analyzer.ai.tool_config import IMPORTANT_KEYS
+from spark_advisor_models.util import format_bytes
+from spark_advisor_rules import StaticAnalysisService
+
+_TOOL_INPUT_MODELS: dict[AgentToolName, type[BaseModel]] = {
+    AgentToolName.GET_STAGE_DETAILS: GetStageDetailsInput,
+    AgentToolName.CALCULATE_OPTIMAL_PARTITIONS: CalculateOptimalPartitionsInput,
+    AgentToolName.COMPARE_CONFIGS: CompareConfigsInput,
+}
+
+
+class ToolExecutionError(Exception):
+    pass
+
+
+def execute_tool(
+    name: str,
+    input_data: dict[str, Any],
+    context: AgentContext,
+    static_analysis: StaticAnalysisService,
+) -> str:
+    try:
+        tool = AgentToolName(name)
+    except ValueError as err:
+        raise ToolExecutionError(f"Unknown tool: {name}") from err
+
+    if tool == AgentToolName.SUBMIT_FINAL_REPORT:
+        raise ToolExecutionError("submit_final_report is handled by the orchestrator, not execute_tool")
+
+    input_model = _TOOL_INPUT_MODELS.get(tool)
+    validated: BaseModel | None = None
+    if input_model is not None:
+        try:
+            validated = input_model.model_validate(input_data)
+        except ValidationError as e:
+            raise ToolExecutionError(f"Invalid input for {name}: {e}") from e
+
+    match tool:
+        case AgentToolName.GET_JOB_OVERVIEW:
+            return _handle_get_job_overview(context)
+        case AgentToolName.GET_STAGE_DETAILS:
+            assert isinstance(validated, GetStageDetailsInput)
+            return _handle_get_stage_details(context, validated)
+        case AgentToolName.RUN_RULES_ENGINE:
+            return _handle_run_rules_engine(context, static_analysis)
+        case AgentToolName.CALCULATE_OPTIMAL_PARTITIONS:
+            assert isinstance(validated, CalculateOptimalPartitionsInput)
+            return _handle_calculate_optimal_partitions(context, validated)
+        case AgentToolName.COMPARE_CONFIGS:
+            assert isinstance(validated, CompareConfigsInput)
+            return _handle_compare_configs(context, validated)
+        case _:
+            raise ToolExecutionError(f"Unhandled tool: {tool}")
+
+
+def _handle_get_job_overview(context: AgentContext) -> str:
+    job = context.job
+    total_tasks = sum(s.tasks.task_count for s in job.stages)
+    total_shuffle_read = sum(s.total_shuffle_read_bytes for s in job.stages)
+    total_shuffle_write = sum(s.total_shuffle_write_bytes for s in job.stages)
+    total_spill = sum(s.spill_to_disk_bytes for s in job.stages)
+    total_input = sum(s.input_bytes for s in job.stages)
+
+    config_snapshot = {k: job.config.get(k) for k in IMPORTANT_KEYS if job.config.get(k)}
+
+    overview: dict[str, Any] = {
+        "app_id": job.app_id,
+        "app_name": job.app_name,
+        "spark_version": job.spark_version,
+        "duration_ms": job.duration_ms,
+        "duration_min": round(job.duration_ms / 60_000, 1),
+        "stage_count": len(job.stages),
+        "total_tasks": total_tasks,
+        "total_input": format_bytes(total_input),
+        "total_shuffle_read": format_bytes(total_shuffle_read),
+        "total_shuffle_write": format_bytes(total_shuffle_write),
+        "total_spill_to_disk": format_bytes(total_spill),
+        "stages_summary": [
+            {
+                "stage_id": s.stage_id,
+                "name": s.stage_name,
+                "tasks": s.tasks.task_count,
+                "gc_percent": round(s.gc_time_percent, 1) if s.sum_executor_run_time_ms > 0 else 0.0,
+                "has_spill": s.spill_to_disk_bytes > 0,
+                "has_shuffle": s.total_shuffle_read_bytes > 0 or s.total_shuffle_write_bytes > 0,
+                "skew_ratio": round(s.tasks.duration_skew_ratio, 1),
+                "failed_tasks": s.failed_task_count,
+            }
+            for s in job.stages
+        ],
+        "config": config_snapshot,
+    }
+
+    if job.executors:
+        overview["executors"] = {
+            "count": job.executors.executor_count,
+            "total_cores": job.executors.total_cores,
+            "memory_utilization_percent": round(job.executors.memory_utilization_percent, 1),
+            "failed_tasks": job.executors.failed_tasks,
+        }
+
+    return json.dumps(overview)
+
+
+def _handle_get_stage_details(context: AgentContext, input_data: GetStageDetailsInput) -> str:
+    stage = next((s for s in context.job.stages if s.stage_id == input_data.stage_id), None)
+    if stage is None:
+        available = [s.stage_id for s in context.job.stages]
+        raise ToolExecutionError(f"Stage {input_data.stage_id} not found. Available stages: {available}")
+
+    return stage.model_dump_json()
+
+
+def _handle_run_rules_engine(context: AgentContext, static_analysis: StaticAnalysisService) -> str:
+    results = context.get_or_run_rules(static_analysis)
+
+    findings = [
+        {
+            "rule_id": r.rule_id,
+            "severity": r.severity.value,
+            "title": r.title,
+            "message": r.message,
+            "stage_id": r.stage_id,
+            "current_value": r.current_value,
+            "recommended_value": r.recommended_value,
+        }
+        for r in results
+    ]
+    return json.dumps({"findings_count": len(findings), "findings": findings})
+
+
+def _handle_calculate_optimal_partitions(
+    context: AgentContext,
+    input_data: CalculateOptimalPartitionsInput,
+) -> str:
+    total_bytes = input_data.total_shuffle_bytes
+    target_mb = input_data.target_partition_mb
+    target_bytes = target_mb * 1024 * 1024
+
+    optimal = max(1, math.ceil(total_bytes / target_bytes)) if total_bytes > 0 else 1
+    current = context.job.config.shuffle_partitions
+
+    actual_size_mb = round(total_bytes / (optimal * 1024 * 1024), 1) if optimal > 0 else 0
+
+    return json.dumps(
+        {
+            "current_partitions": current,
+            "optimal_partitions": optimal,
+            "total_shuffle_bytes": total_bytes,
+            "total_shuffle_formatted": format_bytes(total_bytes),
+            "target_partition_size_mb": target_mb,
+            "actual_partition_size_mb": actual_size_mb,
+        }
+    )
+
+
+def _handle_compare_configs(context: AgentContext, input_data: CompareConfigsInput) -> str:
+    changes = []
+    for key, new_value in input_data.proposed_changes.items():
+        current = context.job.config.get(key)
+        changes.append(
+            {
+                "parameter": key,
+                "current_value": current or "(not set)",
+                "proposed_value": new_value,
+                "changed": current != new_value,
+            }
+        )
+    return json.dumps({"changes": changes, "total_changes": len(changes)})

spark_advisor_analyzer-0.1.6/src/spark_advisor_analyzer/agent/orchestrator.py

@@ -0,0 +1,167 @@
+import json
+import logging
+from typing import Any
+
+from anthropic.types import (
+    Message,
+    MessageParam,
+    ToolChoiceAutoParam,
+    ToolChoiceToolParam,
+    ToolResultBlockParam,
+    ToolUseBlock,
+)
+from pydantic import ValidationError
+
+from spark_advisor_analyzer.agent.context import AgentContext
+from spark_advisor_analyzer.agent.handlers import ToolExecutionError, execute_tool
+from spark_advisor_analyzer.agent.prompts import build_agent_system_prompt, build_initial_message
+from spark_advisor_analyzer.agent.tools import AGENT_TOOLS, AgentToolName
+from spark_advisor_analyzer.ai.client import AnthropicClient
+from spark_advisor_analyzer.ai.report_builder import build_advisor_report
+from spark_advisor_models.config import AiSettings
+from spark_advisor_models.model import AnalysisResult, AnalysisToolInput, JobAnalysis
+from spark_advisor_rules import StaticAnalysisService
+
+logger = logging.getLogger(__name__)
+
+
+class AgentOrchestrator:
+    def __init__(
+        self,
+        client: AnthropicClient,
+        static_analysis: StaticAnalysisService,
+        ai_settings: AiSettings,
+    ) -> None:
+        self._client = client
+        self._static = static_analysis
+        self._ai = ai_settings
+        self._max_iterations = ai_settings.max_agent_iterations
+        self._system_prompt = build_agent_system_prompt(self._max_iterations)
+
+    def run(self, job: JobAnalysis) -> AnalysisResult:
+        context = AgentContext(job=job)
+        messages: list[MessageParam] = [
+            MessageParam(role="user", content=build_initial_message(job)),
+        ]
+
+        for iteration in range(self._max_iterations):
+            logger.info("Agent iteration %d/%d", iteration + 1, self._max_iterations)
+
+            response = self._call_claude(messages, tool_choice=ToolChoiceAutoParam(type="auto"))
+
+            final_input = self._extract_final_report(response)
+            if final_input is not None:
+                logger.info("Agent completed after %d iterations", iteration + 1)
+                return self._build_result(job, context, final_input)
+
+            tool_uses = [b for b in response.content if isinstance(b, ToolUseBlock)]
+
+            if not tool_uses:
+                messages.append(MessageParam(role="assistant", content=response.content))
+                messages.append(
+                    MessageParam(
+                        role="user",
+                        content=f"Please call {AgentToolName.SUBMIT_FINAL_REPORT} with your analysis now.",
+                    )
+                )
+                continue
+
+            messages.append(MessageParam(role="assistant", content=response.content))
+
+            tool_results: list[ToolResultBlockParam] = []
+            for tool_use in tool_uses:
+                logger.info("Executing tool: %s", tool_use.name)
+                result_content = self._execute_single_tool(tool_use.name, tool_use.input, context)
+                tool_results.append(
+                    {
+                        "type": "tool_result",
+                        "tool_use_id": tool_use.id,
+                        "content": result_content,
+                    }
+                )
+
+            messages.append(MessageParam(role="user", content=tool_results))
+
+        logger.warning("Agent hit max iterations (%d), forcing final report", self._max_iterations)
+        return self._force_final_report(job, context, messages)
+
+    def _call_claude(
+        self,
+        messages: list[MessageParam],
+        *,
+        tool_choice: ToolChoiceAutoParam | ToolChoiceToolParam,
+    ) -> Message:
+        return self._client.create_message(
+            model=self._ai.model,
+            max_tokens=self._ai.max_tokens,
+            system=self._system_prompt,
+            messages=messages,
+            tools=AGENT_TOOLS,
+            tool_choice=tool_choice,
+        )
+
+    def _execute_single_tool(
+        self,
+        name: str,
+        input_data: Any,
+        context: AgentContext,
+    ) -> str:
+        try:
+            return execute_tool(name, input_data, context, self._static)
+        except ToolExecutionError as e:
+            logger.warning("Tool %s failed: %s", name, e)
+            return json.dumps({"error": str(e)})
+        except Exception:
+            logger.exception("Unexpected error in tool %s", name)
+            return json.dumps({"error": f"Internal error executing {name}"})
+
+    def _force_final_report(
+        self,
+        job: JobAnalysis,
+        context: AgentContext,
+        messages: list[MessageParam],
+    ) -> AnalysisResult:
+        messages.append(
+            MessageParam(
+                role="user",
+                content=(
+                    "You have reached the maximum number of tool calls. "
+                    "You MUST call submit_final_report now with your best analysis "
+                    "based on the information gathered so far."
+                ),
+            )
+        )
+
+        response = self._call_claude(
+            messages,
+            tool_choice=ToolChoiceToolParam(type="tool", name=AgentToolName.SUBMIT_FINAL_REPORT),
+        )
+
+        final_input = self._extract_final_report(response)
+        if final_input is not None:
+            return self._build_result(job, context, final_input)
+
+        logger.error("Force-submit failed, returning rules-only result")
+        rule_results = context.get_or_run_rules(self._static)
+        return AnalysisResult(app_id=job.app_id, job=job, rule_results=rule_results)
+
+    def _build_result(
+        self,
+        job: JobAnalysis,
+        context: AgentContext,
+        parsed: AnalysisToolInput,
+    ) -> AnalysisResult:
+        rule_results = context.get_or_run_rules(self._static)
+        report = build_advisor_report(job.app_id, parsed, rule_results)
+        return AnalysisResult(app_id=job.app_id, job=job, rule_results=rule_results, ai_report=report)
+
+    @staticmethod
+    def _extract_final_report(response: Message) -> AnalysisToolInput | None:
+        for block in response.content:
+            if isinstance(block, ToolUseBlock) and block.name == AgentToolName.SUBMIT_FINAL_REPORT:
+                try:
+                    return AnalysisToolInput.model_validate(block.input)
+                except ValidationError:
+                    logger.warning("Failed to validate submit_final_report input, skipping")
+                    return None
+        return None