Metadata-Version: 2.4
Name: quarkupy
Version: 0.8.2
Summary: Python abstractions for Quark execution in the Hadron distributed execution infrastructure
Author-email: Alex Breshears <apbresh@gmail.com>
Requires-Python: >=3.10
Requires-Dist: grpcio-tools>=1.78.0
Requires-Dist: grpcio>=1.78.0
Requires-Dist: protobuf>=4.25.0
Requires-Dist: psutil>=5.9.0
Requires-Dist: pyarrow>=23.0.0
Requires-Dist: pybars3>=0.9.7
Requires-Dist: pydantic>=2.12.5
Requires-Dist: pyiceberg[adlfs,pyarrow,pyiceberg-core,rest-sigv4,s3fs]>=0.11.0
Requires-Dist: strenum>=0.4.15; python_version < '3.11'
Requires-Dist: tomli>=2.0.1; python_version < '3.11'
Requires-Dist: uuid6>=2024.1.12
Provides-Extra: dev
Requires-Dist: grpcio-testing>=1.78.0; extra == 'dev'
Requires-Dist: mypy>=1.0.0; extra == 'dev'
Requires-Dist: pytest-asyncio>=0.21.0; extra == 'dev'
Requires-Dist: pytest>=7.0.0; extra == 'dev'
Requires-Dist: ruff>=0.1.0; extra == 'dev'
Description-Content-Type: text/markdown

# Quarkupy

Python framework for building quarks that execute within the Hadron distributed execution infrastructure.

## Overview

Quarkupy provides:

- **QuarkRunner** - Base class for implementing Python quarks with typed input/output
- **QuarkContext** - Execution context with task metadata, logging, progress reporting, and cancellation
- **IPCClient** - Arrow IPC client for dataset exchange via gRPC (metadata) and Arrow Flight (data transfer)
- **CallbackClient** - gRPC client for status reporting and heartbeats to the worker
- **HistoryLogClient** - Streaming log client for centralized log aggregation
- **QuarkHTTPClient** - Lightweight REST SDK for cluster, registry, flow, history, dataset, context, vector, thread, IAM, scheduler, and Iceberg APIs
- **Error Classification** - Structured errors with retry hints (RETRYABLE vs FATAL)
- **Testing Framework** - Mocks, assertions, and test harness for unit and integration testing
- **Build System** - Package Python quarks into distributable tarballs with launchers

## Installation

```bash
# Install with uv (recommended)
uv sync

# Or with pip
pip install -e .

# Generate proto code (required before first use)
python build_protos.py
```

## Quick Start

### 1. Define Your Quark

```python
from pydantic import BaseModel
from quarkupy import QuarkRunner, QuarkContext, QuarkInput, QuarkOutput

class MyInput(BaseModel):
    source_file: str
    multiplier: int = 1

class MyOutput(BaseModel):
    line_count: int
    char_count: int

class LineCountQuark(QuarkRunner[MyInput, MyOutput]):
    name = "Line Counter"
    description = "Counts lines and characters in a file"
    version = "1.0.0"

    def process(
        self,
        ctx: QuarkContext,
        input: QuarkInput[MyInput]
    ) -> QuarkOutput[MyOutput]:
        config = input.config

        ctx.log(f"Processing {config.source_file}")

        with open(config.source_file) as f:
            content = f.read()

        lines = content.count('\n') * config.multiplier
        chars = len(content) * config.multiplier

        ctx.report_progress(1.0, "Complete")

        return QuarkOutput(result=MyOutput(
            line_count=lines,
            char_count=chars
        ))
```

### 2. Run Locally

```bash
quark-py-runner \
    --quark-module my_module \
    --quark-class LineCountQuark \
    --input-json '{"source_file": "test.txt", "multiplier": 2}' \
    --local
```

### 3. Run in Hadron

When executed by Hadron workers, environment variables are set automatically. See [Environment Variables](#environment-variables) for the complete list.

## Core Abstractions

### QuarkRunner

Generic base class for all Python quarks. Type parameters specify input/output types:

```python
class MyQuark(QuarkRunner[InputType, OutputType]):
    name = "My Quark"
    description = "Does something useful"
    version = "1.0.0"

    def process(
        self,
        ctx: QuarkContext,
        input: QuarkInput[InputType]
    ) -> QuarkOutput[OutputType]:
        # Your processing logic
        return QuarkOutput(result=...)
```

### QuarkContext

Provides execution context and utilities:

```python
# Task identification
ctx.task_id           # Task UUID
ctx.flow_id           # Flow UUID
ctx.node_id           # Node ID in the flow DAG
ctx.quark_identifier  # Quark QRN (e.g., "qrn:quark:extractor:smart-ocr")

# Execution context
ctx.attempt_id        # Attempt number (0-indexed, increments on retry)
ctx.partition_index   # Current partition (0-indexed)
ctx.total_partitions  # Total number of partitions
ctx.timeout_secs      # Execution timeout in seconds

# Methods
ctx.log("message")                    # Log a message (info level)
ctx.log("error occurred", level="error")  # Log with level
ctx.report_progress(0.5, "Halfway")   # Report progress (0.0 to 1.0)
ctx.check_cancellation()              # Check if cancelled (returns bool)

# Properties
ctx.progress          # Current progress value
ctx.is_cancelled      # Whether execution was cancelled
```

### QuarkInput

Wrapper for input configuration and optional dataset reference:

```python
input.config              # Parsed configuration (typed as InputT)
input.input_dataset_id    # Optional input dataset ID from upstream quark

# Factory methods
QuarkInput.from_json(json_str, config_type)  # Parse from JSON
QuarkInput.from_dict(data_dict, config_type) # Parse from dict
```

### QuarkOutput

Wrapper for output result and optional dataset reference:

```python
QuarkOutput(
    result=MyOutput(...),           # Required: output result
    output_dataset_id="uuid-...",   # Optional: output dataset ID
    metrics={"rows": 100},          # Optional: custom metrics
)

# Serialization
output.to_dict()  # Convert to dictionary
output.to_json()  # Convert to JSON string
```

## Error Classification

Errors are classified for retry decisions:

```python
from quarkupy import (
    QuarkError,
    ErrorClass,
    ConfigurationError,
    ResourceError,
    ExecutionError,
    CancellationError,
    TimeoutError,
    CallbackError,
    IPCError,
    HistoryError,
)

# Configuration error (always fatal)
raise ConfigurationError("Invalid input format")

# Resource error (retryable by default)
raise ResourceError("Network timeout", retryable=True)
raise ResourceError("File not found", retryable=False)

# IPC error (retryable by default)
raise IPCError("Connection refused", retryable=True)
raise IPCError("Dataset not found", retryable=False)

# Generic quark error with explicit classification
raise QuarkError("Something went wrong", ErrorClass.FATAL)
raise QuarkError("Temporary failure", ErrorClass.RETRYABLE)

# Check if error is retryable
if error.is_retryable:
    # Will be retried by Hadron
    pass
```

## Lifecycle Hooks

Override these methods for custom behavior:

```python
class MyQuark(QuarkRunner[InputType, OutputType]):
    def on_start(self, ctx: QuarkContext) -> None:
        """Called before processing starts. Use for initialization."""
        ctx.log("Initializing resources...")

    def on_complete(self, ctx: QuarkContext, output: QuarkOutput) -> None:
        """Called after successful completion. Use for cleanup."""
        ctx.log(f"Processed {output.result.count} items")

    def on_error(self, ctx: QuarkContext, error: Exception) -> None:
        """Called when processing fails. Use for error cleanup."""
        ctx.log(f"Failed: {error}", level="error")

    def on_cancel(self) -> None:
        """Called when cancellation is requested. Use for graceful shutdown."""
        self._cleanup_resources()

    def validate_input(self, config: Any) -> InputType:
        """Override for custom input validation."""
        # Default implementation uses Pydantic/dataclass parsing
        return super().validate_input(config)
```

## Arrow IPC (Data Exchange)

For quarks that process datasets, use `IPCClient` with separate gRPC and Flight addresses:

```python
from quarkupy import IPCClient, generate_dataset_id

class DataQuark(QuarkRunner[MyInput, MyOutput]):
    def process(self, ctx, input):
        # Get addresses from environment or config
        ipc_grpc_addr = os.environ.get("QUARK__IPC_GRPC_ADDR", "localhost:50300")
        ipc_flight_addr = os.environ.get("QUARK__IPC_FLIGHT_ADDR", "localhost:50301")

        with IPCClient(ipc_grpc_addr, ipc_flight_addr) as ipc:
            # Read input dataset (with optional filtering)
            table = ipc.read_dataset(
                input.input_dataset_id,
                partition_filter=[ctx.partition_index],  # Optional: filter by partition
                row_limit=1000,                          # Optional: limit rows
            )

            # Process data...
            result_table = self._process(table)

            # Generate output dataset ID
            output_id = generate_dataset_id()

            # Register, write, and complete output dataset
            ipc.register_dataset(
                dataset_id=output_id,
                expected_chunks=1,
                total_partitions=ctx.total_partitions,
                name=f"my-output-{ctx.task_id}",
                flow_id=ctx.flow_id,
                quark_id=ctx.quark_identifier,
                task_id=ctx.task_id,
            )

            ipc.write_dataset(
                dataset_id=output_id,
                chunk_id=f"{output_id}-chunk-0",
                attempt_id=str(ctx.attempt_id),
                partition_index=ctx.partition_index,
                table=result_table,
            )

            ipc.complete_dataset(output_id, total_rows=result_table.num_rows)

        return QuarkOutput(
            result=MyOutput(processed=result_table.num_rows),
            output_dataset_id=output_id,
        )
```

## Quark REST SDK

Use `QuarkHTTPClient` when a Python quark, notebook, or automation script needs to
talk to the REST gateway instead of the in-process gRPC/Flight runtime clients.

```python
from quarkupy import QuarkHTTPClient

client = QuarkHTTPClient.from_env()

# Cluster and registry
services = client.cluster.services(healthy_only=True)
tool = client.registry.tool("qrn:tool:registry:list-quarks")

# Run a registered compose flow and inspect results
run = client.flows.run_registered(
    "qrn:flow:extract:source-jsonschema",
    variables={"source_id": "019e..."},
    wait=True,
)
flow = client.flows.get(run["flow_id"])

# Source/context and dataset helpers
files = client.context.files(source_id="019e...", limit=50)
rows = client.datasets.query("dataset-id", "select * from dataset limit 10")

# Sigma threads
thread = client.threads.create({"title": "Extraction check"})
client.threads.start_turn(thread["id"], {"message": "Summarize this source"})
```

Configuration is read from `QUARK__REST_URL`, with auth from
`QUARK__AUTH_TOKEN`/`QUARK__API_KEY` and tenant scope from
`QUARK__TENANT_ID`/`QUARK__WORKSPACE_ID`. Legacy `HADRON_*` env vars remain
accepted as fallbacks. You can also pass those values directly to
`QuarkHTTPClient(...)`.

### Pre-registered Output Datasets

When the planner pre-registers an output dataset, use `write_to_output_dataset`:

```python
output_dataset_id = os.environ.get("QUARK__OUTPUT_DATASET_ID")
if output_dataset_id:
    # Dataset already registered by planner
    ipc.write_to_output_dataset(
        output_dataset_id=output_dataset_id,
        partition_index=ctx.partition_index,
        attempt_id=str(ctx.attempt_id),
        table=result_table,
    )
```

## Testing Framework

Quarkupy provides a comprehensive testing framework:

```python
from quarkupy.testing import (
    QuarkTestHarness,
    assert_execution,
    assert_ipc,
    assert_callback,
    ArrowGenerator,
    create_test_pdf,
)

# Create test harness
harness = QuarkTestHarness.local()

# Execute quark with test configuration
execution = (
    harness.execute(MyQuark)
    .with_input({"source": "test.pdf"})
    .with_input_dataset(input_table)
    .with_partition(0, 4)  # partition_index=0, total_partitions=4
    .execute()
)

# Assert execution results
assert_execution(execution).succeeded()
assert_execution(execution).has_field("processed_count", 5)
assert_execution(execution).has_output_dataset()

# Assert IPC interactions
assert_ipc(execution.ipc).wrote_dataset()
assert_ipc(execution.ipc).registered_dataset_with(expected_chunks=1)

# Assert callback interactions
assert_callback(execution.callback).reported_progress()
assert_callback(execution.callback).completed_successfully()

# Generate test data
pdf_bytes = create_test_pdf(pages=3, text="Sample content")
table = ArrowGenerator.file_list([
    {"file_name": "test.pdf", "binary": pdf_bytes},
])
```

### Available Test Utilities

| Module | Description |
|--------|-------------|
| `QuarkTestHarness` | Fluent harness for executing quarks in test mode |
| `MockIPCClient` | Mock IPC client that captures dataset operations |
| `MockCallbackClient` | Mock callback client that captures status updates |
| `MockHistoryLogClient` | Mock history client that captures log messages |
| `ArrowGenerator` | Generate Arrow tables for common test scenarios |
| `SubprocessRunner` | Run quarks as subprocesses for integration tests |
| Assertions | Fluent assertions for execution, IPC, and callback |
| Fixtures | Pytest fixtures for common test setup |

### Pytest Markers

```python
import pytest
from quarkupy.testing import slow, integration, requires_gpu, requires_openai

@slow
def test_large_file():
    """Excluded by default, run with: pytest -m slow"""
    pass

@integration
def test_with_real_ipc():
    """Integration test requiring real services"""
    pass

@requires_gpu
def test_gpu_acceleration():
    """Test requiring GPU hardware"""
    pass

@requires_openai
def test_vision_api():
    """Test requiring OpenAI API key"""
    pass
```

## Build System

Package quarks into distributable tarballs:

```bash
# Build a quark (requires Quark.toml in current directory)
quark-build --verbose

# Output:
#   dist/<quark-name>.tar.gz   # Tarball with venv + source
#   dist/<quark-name>.sh       # Launcher script
#   dist/<quark-name>.sha256   # Checksum file
```

### Quark.toml Configuration

```toml
[quark]
spec = "0.5.0-alpha"
identifier = "qrn:quark:extractor:my-quark"
name = "My Quark"
description = "Description of what it does"
version = "1.0.0"

node_type = "input_output"  # input, output, transform, input_output
category = "extractor"
partitionable = true

[quark.execution]
mode = "python"
module = "my_package.quark"
class = "MyQuark"
python_version = ">=3.13"

[quark.resources]
cpu = 200
memory = 200
min_memory_mb = 4096
min_cpu_cores = 2.0
requires_gpu = false
```

## Environment Variables

All environment variables are prefixed with `QUARK__`:

### Task Identification

| Variable | Description |
|----------|-------------|
| `QUARK__TASK_ID` | Task UUID |
| `QUARK__FLOW_ID` | Flow UUID |
| `QUARK__NODE_ID` | Node ID in the flow DAG |
| `QUARK__QUARK_IDENTIFIER` | Quark QRN (e.g., `qrn:quark:extractor:smart-ocr`) |

### Execution Context

| Variable | Description |
|----------|-------------|
| `QUARK__ATTEMPT_ID` | Attempt number (0-indexed, increments on retry) |
| `QUARK__PARTITION_INDEX` | Current partition index (0-indexed) |
| `QUARK__TOTAL_PARTITIONS` | Total number of partitions |
| `QUARK__TIMEOUT_SECS` | Execution timeout in seconds |

### Service Addresses

| Variable | Description |
|----------|-------------|
| `QUARK__WORKER_CALLBACK` | gRPC address for worker callback service |
| `QUARK__IPC_GRPC_ADDR` | gRPC address for IPC metadata operations |
| `QUARK__IPC_FLIGHT_ADDR` | Arrow Flight address for IPC data transfer |
| `QUARK__IPC_ADDR` | Legacy alias for `IPC_GRPC_ADDR` (deprecated) |
| `QUARK__HISTORY_ADDR` | gRPC address for history/log service |

### Input/Output

| Variable | Description |
|----------|-------------|
| `QUARK__INPUT_JSON` | JSON input configuration |
| `QUARK__INPUT` | Alias for `INPUT_JSON` |
| `QUARK__INPUT_DATASET_ID` | Input dataset ID from upstream quark |
| `QUARK__OUTPUT_DATASET_ID` | Pre-registered output dataset ID |

### Partitioning (Range-based)

| Variable | Description |
|----------|-------------|
| `QUARK__INPUT_PARTITION_START_ROW` | Start row for this partition (inclusive) |
| `QUARK__INPUT_PARTITION_END_ROW` | End row for this partition (exclusive) |

### Module Loading

| Variable | Description |
|----------|-------------|
| `QUARK__QUARK_MODULE` | Python module containing the quark |
| `QUARK__QUARK_CLASS` | Class name of the quark |
| `QUARK__HEARTBEAT_INTERVAL` | Heartbeat interval in seconds (default: 5) |

## CLI Reference

### quark-py-runner

```
quark-py-runner [OPTIONS]

Options:
  --quark-module TEXT       Python module containing the quark
  --quark-class TEXT        Class name of the quark
  --input-json TEXT         JSON input configuration
  --callback-addr TEXT      gRPC address for worker callback service
  --ipc-grpc-addr TEXT      gRPC address for IPC service
  --ipc-flight-addr TEXT    Arrow Flight address for IPC service
  --task-id TEXT            Task UUID
  --flow-id TEXT            Flow UUID
  --attempt-id INT          Attempt number
  --partition-index INT     Partition index
  --total-partitions INT    Total partitions
  --timeout-secs INT        Execution timeout in seconds
  --heartbeat-interval INT  Heartbeat interval in seconds
  --local                   Run in local mode (no callback service)
```

### quark-build

```
quark-build [OPTIONS]

Options:
  --quark-toml PATH    Path to Quark.toml (default: ./Quark.toml)
  --output-dir PATH    Output directory (default: ./dist)
  --verbose            Enable verbose output
```

## Project Structure

```
quarkupy/
  pyproject.toml                 # Project configuration
  build_protos.py                # Proto generation script
  src/
    quarkupy/
      __init__.py                # Public exports
      runner.py                  # QuarkRunner base class
      context.py                 # QuarkContext, QuarkInput, QuarkOutput
      callback.py                # gRPC callback client
      io.py                      # Arrow IPC client (gRPC + Flight)
      config.py                  # Configuration handling
      errors.py                  # Error classes
      metrics.py                 # Resource metrics
      history.py                 # History log client
      cli.py                     # quark-py-runner CLI
      testing/                   # Testing framework
        __init__.py              # Public testing exports
        harness.py               # QuarkTestHarness
        builder.py               # QuarkExecutionBuilder
        mocks.py                 # Mock clients
        assertions.py            # Fluent assertions
        arrow_gen.py             # Arrow table generators
        subprocess_runner.py     # Subprocess test runner
        fixtures.py              # Pytest fixtures
        types.py                 # Test types
      build/                     # Build system
        __init__.py              # Public build exports
        cli.py                   # quark-build CLI
        config.py                # Quark.toml parsing
        builder.py               # QuarkBuilder
        packager.py              # Tarball packaging
        launcher.py              # Launcher script generation
        venv.py                  # Virtual environment creation
      _generated/                # Auto-generated proto code
```

## Development

```bash
# Install dev dependencies
uv sync --dev

# Generate proto code
python build_protos.py

# Run tests
pytest

# Run tests with markers
pytest -m "not slow"           # Skip slow tests (default)
pytest -m "slow"               # Only slow tests
pytest -m "integration"        # Only integration tests

# Type checking
mypy src/quarkupy

# Linting
ruff check src/quarkupy

# Format code
ruff format src/quarkupy
```

## Example: q-ocr Quark

See `quarks/q-ocr/` for a complete example that:

- Converts documents (PDF, DOCX, PPTX, HTML, images) to Markdown using Docling
- Uses RapidOCR for scanned documents and EasyOCR for digital documents
- Supports horizontal scaling via file-level partitioning
- Reports progress during batch processing
- Handles cancellation gracefully with resource cleanup
- Integrates with IPC for dataset exchange
