From 58a039344f0360405cb9f12896cf30559b1af446 Mon Sep 17 00:00:00 2001 From: jiangxt2 Date: Thu, 13 Aug 2026 22:38:44 +0800 Subject: [PATCH 1/3] feat(broker): add transport-neutral Broker SPI, lazy discovery, and generic runner MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Tributo needs an optional control-plane integration for internal Redis Streams tasks, lifecycle events, and cooperative cancellation. Redis and KnoVa protocol models must stay in a separately installed provider wheel, so Core only gains a safe, lazily loaded, versioned Broker extension contract. The provider wheel is developed separately (independent tributo-broker-redis repository); this branch carries only the Core side. - Broker SPI: BrokerPlugin, BrokerRuntime, TaskConsumer, EventReporter, CancellationChecker, Message/JobResult identity fields, TaskDisposition, and JSON-safe CancellationSpec; new dataclass fields default so existing constructions keep working. - Lazy discovery: tributo.brokers entry-point group, fail-open discovery with diagnostics, fail-closed explicit resolution, TRIBUTO_PLUGINS filtering, and api_version/broker_id/capabilities contract checks. - Broker Registry and generic BrokerRunner with ACK/RETRY/REJECT, bounded reconnect backoff, DEGRADED/RECONNECTING lifecycle states, and graceful shutdown. - Lazy CLI: broker list/validate/consume mounted through a lazy Click group; importing tributo.cli does not import the broker modules. - Runtime injection: build_runtime_env extra_py_modules/runtime_pip_packages and submit_training_job execution_context plumbing; the default path adds no pip section. - Identity mapping: submit_training_job_with_identity returns run_id, attempt_id, submission_id, and the Ray execution job_id; reconcile resolves the real execution ID when available and retains a compatibility fallback. - Training bridge: worker-side cancellation checker rebuilt from JSON-safe specs (fail-open), _tributo_* keys stripped before strict config validation, cancel state surfaced in report metrics, and Ray Train metrics/history preserved in the lifecycle summary for provider replay. - Docs: STABILITY.md beta entries for the three broker modules and ADR 002 (broker plugin boundary). - Targeted Core suite (broker SPI, discovery, runner, CLI isolation, runtime env, stability inventory, training lifecycle): `pytest -p no:cacheprovider tests/test_broker.py tests/test_broker_cli.py tests/test_runtime_env.py tests/test_stability_inventory.py tests/training/test_job_submitter.py tests/training/test_training_lifecycle.py -q` — 766 passed, 6 skipped. - Full Core non-IT regression: `pytest -p no:cacheprovider tests --ignore=tests/integration -q` — 2867 passed, 40 skipped, 171 deselected. - Core Ray runtime-environment IT: `tests/integration/test_ray_runtime_env.py -m ray_runtime_env` — 1 passed. - Lint: ruff format check and ruff check on the changed files pass, plus pre-commit hooks; git diff --check passes. - Real Redis Stream/Ray provider IT is intentionally excluded from this Core-only commit and is maintained in the independent provider repository. - [x] Unit tests pass - [x] Lint passes for changed Python files (ruff format --check, ruff check, and pre-commit) - [x] Integration tests pass (if applicable) — Core ray-runtime IT passed; real Redis/Ray provider IT lives in the independent provider repository - [x] No internal credentials, URLs, or tokens exposed - [x] New external dependencies reviewed for license compatibility — no new dependencies Signed-off-by: jiangxt2 --- docs/STABILITY.md | 3 + docs/adr/002-broker-plugin-boundary.md | 58 ++++ src/tributo/cli.py | 44 ++- src/tributo/cli_broker.py | 123 ++++++++ src/tributo/integrations/broker.py | 241 +++++++++++----- src/tributo/integrations/broker_registry.py | 137 +++++++++ src/tributo/integrations/broker_runner.py | 204 ++++++++++++++ src/tributo/plugin.py | 146 ++++++++++ src/tributo/training/job_submitter.py | 126 ++++++--- src/tributo/training/lifecycle.py | 7 + src/tributo/training/xgboost_trainer.py | 40 +-- tests/test_broker.py | 293 ++++++++++++++++++++ tests/test_broker_cli.py | 66 +++++ tests/test_runtime_env.py | 24 ++ tests/test_stability_inventory.py | 8 + tests/training/test_job_submitter.py | 70 ++++- tests/training/test_training_lifecycle.py | 17 ++ 17 files changed, 1463 insertions(+), 144 deletions(-) create mode 100644 docs/adr/002-broker-plugin-boundary.md create mode 100644 src/tributo/cli_broker.py create mode 100644 src/tributo/integrations/broker_registry.py create mode 100644 src/tributo/integrations/broker_runner.py create mode 100644 tests/test_broker.py create mode 100644 tests/test_broker_cli.py diff --git a/docs/STABILITY.md b/docs/STABILITY.md index a601e33..4b3df24 100644 --- a/docs/STABILITY.md +++ b/docs/STABILITY.md @@ -27,6 +27,7 @@ This page provides module-level guidance and deprecation notes. | `tributo.config` — `AlgorithmExecutionConfig` and nested algorithm execution models | `alpha` | Strict JSON envelope shared by local Ray and Kubernetes-hosted Ray execution | | `tributo.job` — `TributoClient` | `stable` | Primary Ray Jobs client | | `tributo.job` — `RayJob` | `stable` annotation with runtime deprecation warning | Use `TributoClient`; the annotation and warning conflict is documented without changing the public contract in this documentation update | +| `tributo.ray_jobs` | `alpha` | Workload-neutral submission identity, ambiguous-submit reconciliation, status, logs, and stop helpers | | `tributo.exceptions` — core exceptions | `stable` | ``TributoError`` and 16 common subtypes | | `tributo.exceptions` — `ResultMaterializationError` | `alpha` | Credential-safe lazy inference action failure | | `tributo.exceptions` — Bundle/Plugin exceptions | `beta` | ``BundleExportError``, ``BundleCommitBusyError``, ``AliasConflict``, ``UnsupportedArtifactFormat``, ``PostPublishCallbackError``, ``PluginLoadIssue`` | @@ -204,6 +205,8 @@ from the legacy setup-only propagation rule. | `tributo.integrations.model_importers.*` | `alpha` | Canonical ModelImporter protocol/registry plus explicit MLflow and typed artifact-to-Bundle implementations | | `tributo.integrations.sinks.parquet` | `alpha` | Parquet inference ResultSink adapter | | `tributo.integrations.sinks.lance` | `alpha` | Generic Lance inference ResultSink adapter | +| `tributo.integrations.broker` | `alpha` | Minimal transport-neutral Broker API v1; transport implementations and consume loops are external | +| `tributo.integrations.broker_registry` | `alpha` | Lazy broker discovery and explicit provider resolution | ### Inference (tributo.inference.*) diff --git a/docs/adr/002-broker-plugin-boundary.md b/docs/adr/002-broker-plugin-boundary.md new file mode 100644 index 0000000..c13103d --- /dev/null +++ b/docs/adr/002-broker-plugin-boundary.md @@ -0,0 +1,58 @@ +# Broker plugin boundary + +## Status + +Accepted + +## Context + +Tributo needs an optional control-plane integration for internal Redis +Streams tasks, lifecycle events, and cooperative cancellation. Redis Streams +is not a bounded data source or a streaming inference input. Future Kafka and +RabbitMQ providers should be able to use the same extension mechanism without +adding their client libraries to Tributo Core. + +## Decision + +Tributo Core owns only a transport-neutral, beta Broker SPI, lazy entry-point +discovery, explicit provider resolution, a generic runner, and JSON-safe Ray +execution-context plumbing. Redis and KnoVa protocol models live in an +independently installable provider wheel. + +The Core contract has these safety rules: + +- ordinary Tributo startup and execution never import or connect to a broker; +- discovery is fail-open with diagnostics, while explicitly selected brokers + fail closed when missing, disabled, or invalid; +- broker configuration is passed as provider-owned JSON; Core does not define + Redis/Kafka/RabbitMQ fields or perform network probes implicitly; +- cancellation checkers are reconstructed in Ray workers from serializable + specs; clients, sockets, pools, and secrets are never serialized; +- provider submission must bind the business task ID to `run_id` and use a + deterministic submission ID per execution attempt. Transport delivery + retries must not become new execution attempts: the first Redis provider + reuses `attempt-1` and the same submission ID until an explicit business + retry is authorized after a terminal execution failure; +- temporary transport/submission failures retain the task for recovery; + permanently invalid messages are best-effort reported as FAILED and then + acknowledged even if reporting is unavailable; +- a missing or invalid outer `job_id` is permanently invalid and never becomes + a shared sentinel identity. Its FAILED event goes to a provider-owned + invalid-event stream and carries the delivery ID instead; +- provider reporters implement the Core `EventReporter` method signatures; + provider-specific fields such as KnoVa error codes use explicit extension + methods. Reporter warnings are time-window rate limited; +- reporter failure cannot turn successful training or Bundle publication into + a failed computation. + +The first Redis provider supports training tasks only. It reports lifecycle +events from the Ray Job driver and replays metrics history after training; +real-time metric sinks are a later extension. + +## Consequences + +The Core public surface can evolve independently of transport implementations, +and normal Tributo installations remain free of Redis dependencies. Provider +packages must publish their own protocol and infrastructure contract tests, +and a provider wheel must be installed in the Ray runtime when worker-side +cancellation or provider entrypoints are used. diff --git a/src/tributo/cli.py b/src/tributo/cli.py index 986c540..96db56a 100644 --- a/src/tributo/cli.py +++ b/src/tributo/cli.py @@ -26,7 +26,49 @@ logger = logging.getLogger(__name__) -@click.group() +class _LazyTributoGroup(click.Group): + """Load the broker command module only when the broker command is used.""" + + def get_command(self, ctx: click.Context, cmd_name: str) -> click.Command | None: + if cmd_name == "broker": + from tributo.cli_broker import broker + + return broker + return super().get_command(ctx, cmd_name) + + def list_commands(self, ctx: click.Context) -> list[str]: + commands = super().list_commands(ctx) + if "broker" not in commands: + commands.append("broker") + return sorted(commands) + + def format_commands( + self, + ctx: click.Context, + formatter: click.HelpFormatter, + ) -> None: + """Render broker help from a lightweight placeholder.""" + commands: list[tuple[str, click.Command]] = [] + for command_name in self.list_commands(ctx): + command = ( + click.Command( + "broker", + help="Discover and run explicitly selected message broker plugins.", + ) + if command_name == "broker" + else self.get_command(ctx, command_name) + ) + if command is not None and not command.hidden: + commands.append((command_name, command)) + if not commands: + return + limit = formatter.width - 6 - max(len(name) for name, _ in commands) + rows = [(name, command.get_short_help_str(limit)) for name, command in commands] + with formatter.section("Commands"): + formatter.write_dl(rows) + + +@click.group(cls=_LazyTributoGroup) @click.version_option(package_name="tributo") def main(): """Tributo: Unified framework for submitting Ray Jobs.""" diff --git a/src/tributo/cli_broker.py b/src/tributo/cli_broker.py new file mode 100644 index 0000000..2b379d8 --- /dev/null +++ b/src/tributo/cli_broker.py @@ -0,0 +1,123 @@ +"""Broker-specific CLI commands, mounted by :mod:`tributo.cli`.""" + +from __future__ import annotations + +import json +import logging +import signal +from pathlib import Path +from typing import Any + +import click + +from tributo.exceptions import JobConfigurationError +from tributo.integrations.broker_registry import BrokerRegistry +from tributo.integrations.broker_runner import BrokerRunner + +logger = logging.getLogger(__name__) + + +def _load_config(path: str) -> dict[str, Any]: + config_path = Path(path) + if config_path.suffix.lower() in {".yaml", ".yml"}: + raise click.ClickException("YAML broker config is not supported; use JSON.") + try: + value = json.loads(config_path.read_text(encoding="utf-8")) + except (OSError, json.JSONDecodeError) as exc: + raise click.ClickException(f"Unable to read broker config: {exc}") from exc + if not isinstance(value, dict): + raise click.ClickException("Broker config root must be a JSON object.") + return value + + +@click.group() +def broker() -> None: + """Discover and run explicitly selected message broker plugins.""" + + +@broker.command("list") +def broker_list() -> None: + """List installed broker plugins without connecting to a broker.""" + registry = BrokerRegistry() + descriptors = registry.list() + for descriptor in descriptors: + capabilities = ",".join(descriptor.capabilities) or "-" + click.echo( + f"{descriptor.broker_id}\tapi={descriptor.api_version}" + f"\tcapabilities={capabilities}" + ) + for diagnostic in registry.diagnostics(): + click.echo( + f"diagnostic\t{diagnostic.entry_point_name}\t{diagnostic.reason}", + err=True, + ) + + +@broker.command("validate") +@click.option("--broker", "broker_id", required=True) +@click.option("--config", "config_path", required=True, type=click.Path(exists=True)) +@click.option( + "--check-connectivity", + is_flag=True, + help="Ask the provider to perform an explicit connectivity probe.", +) +def broker_validate(broker_id: str, config_path: str, check_connectivity: bool) -> None: + """Validate provider-owned JSON config.""" + try: + config = _load_config(config_path) + BrokerRegistry().validate( + broker_id, + config, + check_connectivity=check_connectivity, + ) + except (JobConfigurationError, click.ClickException) as exc: + if isinstance(exc, click.ClickException): + raise + raise click.ClickException(str(exc)) from exc + click.echo(f"Broker configuration is valid: {broker_id}") + + +@broker.command("consume") +@click.option("--broker", "broker_id", required=True) +@click.option("--config", "config_path", required=True, type=click.Path(exists=True)) +@click.option("--poll-timeout-ms", default=5000, show_default=True, type=int) +@click.option( + "--once", + is_flag=True, + help="Poll and handle at most one message; useful for controlled operations.", +) +def broker_consume( + broker_id: str, + config_path: str, + poll_timeout_ms: int, + once: bool, +) -> None: + """Consume tasks from an explicitly selected broker plugin.""" + config = _load_config(config_path) + registry = BrokerRegistry() + try: + plugin = registry.validate(broker_id, config) + runner = BrokerRunner( + plugin, + config, + poll_timeout_ms=poll_timeout_ms, + ) + except JobConfigurationError as exc: + raise click.ClickException(str(exc)) from exc + + def _stop(signum: int, _frame: Any) -> None: + logger.info( + "Received %s; stopping broker consumer", + signal.Signals(signum).name, + ) + runner.request_stop() + + signal.signal(signal.SIGINT, _stop) + signal.signal(signal.SIGTERM, _stop) + if once: + try: + runner.run_once() + finally: + runner.close() + else: + runner.run() diff --git a/src/tributo/integrations/broker.py b/src/tributo/integrations/broker.py index f953451..560e54a 100644 --- a/src/tributo/integrations/broker.py +++ b/src/tributo/integrations/broker.py @@ -1,54 +1,56 @@ -"""Message broker abstraction for ML job lifecycle integration. +"""Broker extension contracts for Tributo. -Provides abstract interfaces for task consumption, event reporting, -and cancellation checking. Third-party implementations (Redis, Kafka, -Pulsar, etc.) will register via the ``tributo.brokers`` entry point group. - -.. note:: - - Plugin discovery and the ``tributo.brokers`` entry point group are - planned for v1.1. The ABCs in this module currently have zero - concrete implementations — they define the contract that broker - packages (e.g. ``tributo-broker-redis``) will fulfill. - -See Also: - :mod:`tributo.plugin` — plugin discovery infrastructure. +The contracts in this module are deliberately transport-neutral. Redis, +Kafka, RabbitMQ, and protocol-specific models belong to separately installed +provider packages and are discovered only when a broker operation is +explicitly requested. """ from __future__ import annotations +import json from abc import ABC, abstractmethod from dataclasses import dataclass, field -from typing import Any +from enum import StrEnum +from typing import Any, Mapping + +BROKER_API_VERSION = 1 -# ── Data types ──────────────────────────────────────────────────────────────── + +class TaskDisposition(StrEnum): + """Decision the generic runner applies after a provider handles a task.""" + + ACK = "ack" + RETRY = "retry" + REJECT = "reject" @dataclass class Message: """A job request consumed from a message queue. - Attributes: - job_id: Unique identifier for the job. - payload: The deserialized request body (e.g. training config dict). - metadata: Optional routing headers or trace context. + Existing callers may continue constructing a message with only + ``job_id`` and ``payload``. Transport metadata is optional because the + public Core contract must also support brokers without Redis-style + delivery IDs. """ - job_id: str + job_id: str | None payload: dict[str, Any] metadata: dict[str, Any] = field(default_factory=dict) + delivery_id: str | None = None + delivery_attempt: int = 1 + run_id: str | None = None + attempt_id: str | None = None @dataclass class JobResult: """Outcome of a completed ML training job. - Attributes: - job_id: Unique identifier for the job. - status: Terminal status — ``"success"``, ``"failed"``, or ``"cancelled"``. - metrics: Final evaluation metrics (e.g. accuracy, loss). - artifacts: Paths or URIs of produced artifacts (model files, reports). - error: Human-readable error message when status is ``"failed"``. + The identity fields are optional for backwards-compatible construction; + provider implementations should populate them whenever a broker task is + used. """ job_id: str @@ -56,67 +58,122 @@ class JobResult: metrics: dict[str, float] = field(default_factory=dict) artifacts: list[str] = field(default_factory=list) error: str | None = None + run_id: str | None = None + attempt_id: str | None = None + execution_id: str | None = None + submission_id: str | None = None + bundle_id: str | None = None + bundle_uri: str | None = None + manifest_uri: str | None = None + artifact_refs: list[dict[str, Any]] = field(default_factory=list) + + +@dataclass(frozen=True) +class CancellationSpec: + """JSON-safe instructions for reconstructing a worker-side checker. + + This object intentionally contains no client, socket, pool, or secret. + ``options`` may contain provider-specific non-sensitive values or secret + references, but not credentials themselves. + """ + broker_id: str + job_id: str + options: dict[str, Any] = field(default_factory=dict) + + def __post_init__(self) -> None: + if not self.broker_id.strip(): + raise ValueError("CancellationSpec.broker_id must not be empty") + if not self.job_id.strip(): + raise ValueError("CancellationSpec.job_id must not be empty") + try: + json.dumps(self.as_dict(), allow_nan=False) + except (TypeError, ValueError) as exc: + raise ValueError("CancellationSpec must be JSON serializable") from exc + + def as_dict(self) -> dict[str, Any]: + """Return the wire-safe representation placed in Ray config.""" + return { + "broker_id": self.broker_id, + "job_id": self.job_id, + "options": dict(self.options), + } + + @classmethod + def from_mapping(cls, value: Mapping[str, Any]) -> "CancellationSpec": + """Validate and reconstruct a spec received by a Ray worker.""" + broker_id = value.get("broker_id") + job_id = value.get("job_id") + options = value.get("options", {}) + if not isinstance(broker_id, str) or not isinstance(job_id, str): + raise ValueError("CancellationSpec requires string broker_id and job_id") + if not isinstance(options, dict): + raise ValueError("CancellationSpec.options must be an object") + return cls(broker_id=broker_id, job_id=job_id, options=dict(options)) -# ── Abstract interfaces ─────────────────────────────────────────────────────── +@dataclass +class TaskOutcome: + """Provider result consumed by :class:`BrokerRunner`.""" -class TaskConsumer(ABC): - """Consume ML job requests from a message queue. + disposition: TaskDisposition + result: JobResult | None = None + error: str | None = None - Implementations wrap a specific broker client (Redis Streams, - Kafka consumer group, etc.) and yield :class:`Message` objects. - """ + +class TaskConsumer(ABC): + """Consume ML job requests from a message queue.""" @abstractmethod def poll(self, timeout_ms: int = 5000) -> Message | None: - """Block until a job request arrives or *timeout_ms* expires. - - Args: - timeout_ms: Maximum time to wait in milliseconds. - - Returns: - A :class:`Message` if one is available, or ``None`` on timeout. - """ + """Block until a job request arrives or *timeout_ms* expires.""" ... @abstractmethod def ack(self, message: Message) -> None: - """Acknowledge successful processing of *message*. + """Acknowledge successful processing of *message*.""" + ... - After ``ack`` the broker guarantees the message will not be - redelivered. + def retry(self, message: Message, error: str | None = None) -> None: + """Leave a temporarily failed message eligible for redelivery. + + Brokers such as Redis Streams retain a pending message by doing + nothing here. Brokers with an explicit negative acknowledgement may + override this method. """ - ... + del message, error + def reject(self, message: Message, error: str | None = None) -> None: + """Reject a permanently invalid message without acknowledging it. -class EventReporter(ABC): - """Publish ML job lifecycle events. + A provider should normally return ``ACK`` after publishing a FAILED + event for a poison message. This hook exists for transports that have + a dead-letter operation. + """ + del message, error + + def recover_pending(self) -> int: + """Reclaim pending messages when the transport supports it.""" + return 0 + + def close(self) -> None: + """Close transport resources; the default is intentionally a no-op.""" + return None - Every event is keyed by *job_id* so downstream systems can - reconstruct the full timeline of a job. - """ + +class EventReporter(ABC): + """Publish ML job lifecycle events.""" @abstractmethod def report_phase(self, job_id: str, phase: str) -> None: - """Report a lifecycle phase transition. - - Typical phases: ``"initializing"``, ``"training"``, - ``"exporting"``, ``"completed"``. - """ + """Report a lifecycle phase transition.""" ... @abstractmethod def report_metrics( self, job_id: str, metrics: dict[str, float], progress: float ) -> None: - """Report intermediate training metrics. - - Args: - job_id: The job identifier. - metrics: Current metric values (e.g. ``{"loss": 0.35}``). - progress: Progress fraction in ``[0.0, 1.0]``. - """ + """Report metrics; v1 uses this for post-training history replay.""" ... @abstractmethod @@ -129,15 +186,65 @@ def report_failed(self, job_id: str, error: str) -> None: """Report job failure with error details.""" ... + def report_log(self, job_id: str, message: str, level: str = "INFO") -> None: + """Optionally report a user-visible log line.""" + del job_id, message, level -class CancellationChecker(ABC): - """Check whether a running job has been requested to cancel. + def report_cancelled(self, job_id: str, phase: str = "TRAINING") -> None: + """Optionally report cooperative cancellation.""" + del job_id, phase - The training loop calls :meth:`is_cancelled` periodically and - stops early when it returns ``True``. - """ + +class CancellationChecker(ABC): + """Check whether a running job has been requested to cancel.""" @abstractmethod def is_cancelled(self, job_id: str) -> bool: """Return ``True`` if *job_id* should stop early.""" ... + + +class BrokerRuntime(ABC): + """Provider-owned runtime assembled by the generic Core runner.""" + + @property + @abstractmethod + def consumer(self) -> TaskConsumer: + """Return the transport consumer.""" + ... + + @abstractmethod + def handle(self, message: Message) -> TaskOutcome: + """Validate and process one message without transport ACK side effects.""" + ... + + def close(self) -> None: + """Close provider resources.""" + self.consumer.close() + + +class BrokerPlugin(ABC): + """Structural base for an independently installed broker plugin.""" + + api_version: int = BROKER_API_VERSION + broker_id: str + capabilities: frozenset[str] = frozenset() + + @abstractmethod + def validate_config( + self, config: Mapping[str, Any], *, check_connectivity: bool = False + ) -> None: + """Validate provider-owned config and optionally probe connectivity.""" + ... + + @abstractmethod + def create_runtime(self, config: Mapping[str, Any]) -> BrokerRuntime: + """Create a provider runtime; no network I/O belongs in discovery.""" + ... + + @abstractmethod + def create_cancellation_checker( + self, spec: CancellationSpec + ) -> CancellationChecker: + """Rebuild a worker-side checker from a JSON-safe spec.""" + ... diff --git a/src/tributo/integrations/broker_registry.py b/src/tributo/integrations/broker_registry.py new file mode 100644 index 0000000..b839d82 --- /dev/null +++ b/src/tributo/integrations/broker_registry.py @@ -0,0 +1,137 @@ +"""Lazy registry and worker-side reconstruction helpers for brokers.""" + +from __future__ import annotations + +import logging +from collections.abc import Mapping +from dataclasses import dataclass +from typing import Any + +from tributo.exceptions import JobConfigurationError +from tributo.exporting.models import PluginLoadDiagnostic +from tributo.integrations.broker import ( + BrokerPlugin, + CancellationChecker, + CancellationSpec, +) +from tributo.plugin import discover_broker_plugins, resolve_broker_plugin +from tributo.util.annotations import PublicAPI + +logger = logging.getLogger(__name__) + + +@dataclass(frozen=True) +class BrokerDescriptor: + """Metadata exposed by ``tributo broker list`` without instantiation.""" + + broker_id: str + api_version: int + capabilities: tuple[str, ...] + + +@PublicAPI(stability="beta") +class BrokerRegistry: + """Resolve explicitly selected provider plugins on demand.""" + + def __init__(self) -> None: + self._diagnostics: list[PluginLoadDiagnostic] = [] + + def diagnostics(self) -> tuple[PluginLoadDiagnostic, ...]: + """Return non-fatal discovery diagnostics from the last listing.""" + return tuple(self._diagnostics) + + def list(self) -> tuple[BrokerDescriptor, ...]: + """List discoverable brokers without constructing or connecting them.""" + self._diagnostics.clear() + descriptors: list[BrokerDescriptor] = [] + seen: set[str] = set() + for cls in discover_broker_plugins(self._diagnostics): + broker_id = cls.broker_id + if broker_id in seen: + self._diagnostics.append( + PluginLoadDiagnostic( + group="tributo.brokers", + entry_point_name=broker_id, + reason="Duplicate broker_id discovered", + ) + ) + continue + seen.add(broker_id) + capabilities = getattr(cls, "capabilities", frozenset()) + descriptors.append( + BrokerDescriptor( + broker_id=broker_id, + api_version=cls.api_version, + capabilities=tuple(sorted(str(value) for value in capabilities)), + ) + ) + return tuple(descriptors) + + def resolve(self, broker_id: str) -> BrokerPlugin: + """Load one explicitly selected provider, failing closed on errors.""" + cls = resolve_broker_plugin(broker_id) + try: + plugin = cls() + except Exception as exc: + raise JobConfigurationError( + f"Failed to initialize broker {broker_id!r} ({type(exc).__name__})" + ) from exc + return plugin + + def validate( + self, + broker_id: str, + config: Mapping[str, Any], + *, + check_connectivity: bool = False, + ) -> BrokerPlugin: + """Resolve and delegate provider-owned config validation.""" + plugin = self.resolve(broker_id) + try: + plugin.validate_config( + config, + check_connectivity=check_connectivity, + ) + except JobConfigurationError: + raise + except Exception as exc: + raise JobConfigurationError( + f"Broker {broker_id!r} rejected configuration ({type(exc).__name__})" + ) from exc + return plugin + + +def rebuild_cancellation_checker( + value: Mapping[str, Any] | CancellationSpec | None, +) -> CancellationChecker | None: + """Rebuild a checker in a Ray worker from JSON-safe config. + + This helper is intentionally fail-open: a missing or unavailable broker + must not change ordinary training into a failed training run. The error + is logged with the broker and job identity, while secrets remain in the + provider-owned config boundary. + """ + if value is None: + return None + try: + spec = ( + value + if isinstance(value, CancellationSpec) + else CancellationSpec.from_mapping(value) + ) + plugin = BrokerRegistry().resolve(spec.broker_id) + return plugin.create_cancellation_checker(spec) + except Exception: + broker_id = getattr(value, "broker_id", None) + if isinstance(value, Mapping): + broker_id = value.get("broker_id", broker_id) + job_id = getattr(value, "job_id", None) + if isinstance(value, Mapping): + job_id = value.get("job_id", job_id) + logger.warning( + "Unable to rebuild cancellation checker: broker=%s job_id=%s", + broker_id, + job_id, + exc_info=True, + ) + return None diff --git a/src/tributo/integrations/broker_runner.py b/src/tributo/integrations/broker_runner.py new file mode 100644 index 0000000..694e89c --- /dev/null +++ b/src/tributo/integrations/broker_runner.py @@ -0,0 +1,204 @@ +"""Transport-neutral broker consumer runner.""" + +from __future__ import annotations + +import logging +import time +from collections.abc import Callable, Mapping +from enum import StrEnum +from typing import Any + +from tributo.integrations.broker import ( + BrokerPlugin, + BrokerRuntime, + TaskDisposition, + TaskOutcome, +) +from tributo.util.annotations import PublicAPI + +logger = logging.getLogger(__name__) + + +class BrokerRunnerState(StrEnum): + """Observable lifecycle state of a broker consumer process.""" + + STARTING = "STARTING" + READY = "READY" + DEGRADED = "DEGRADED" + RECONNECTING = "RECONNECTING" + STOPPING = "STOPPING" + STOPPED = "STOPPED" + + +@PublicAPI(stability="beta") +class BrokerRunner: + """Run provider-owned message handling with isolated broker failures.""" + + def __init__( + self, + plugin: BrokerPlugin, + config: Mapping[str, Any], + *, + poll_timeout_ms: int = 5000, + backoff_initial: float = 1.0, + backoff_max: float = 30.0, + sleep: Callable[[float], None] = time.sleep, + ) -> None: + if poll_timeout_ms < 0: + raise ValueError("poll_timeout_ms must be non-negative") + if backoff_initial <= 0 or backoff_max < backoff_initial: + raise ValueError("invalid broker reconnect backoff") + self.plugin = plugin + self.config = dict(config) + self.poll_timeout_ms = poll_timeout_ms + self.backoff_initial = backoff_initial + self.backoff_max = backoff_max + self._sleep = sleep + self._runtime: BrokerRuntime | None = None + self._state = BrokerRunnerState.STOPPED + self._stop_requested = False + self._backoff = backoff_initial + + @property + def state(self) -> BrokerRunnerState: + """Return the current runner state.""" + return self._state + + @property + def runtime(self) -> BrokerRuntime | None: + """Expose the provider runtime for diagnostics and tests.""" + return self._runtime + + def start(self) -> None: + """Create the provider runtime without assuming broker availability.""" + if self._runtime is not None and self._state != BrokerRunnerState.STOPPED: + return + self._state = BrokerRunnerState.STARTING + self._stop_requested = False + self._runtime = self.plugin.create_runtime(self.config) + self._backoff = self.backoff_initial + self._state = BrokerRunnerState.READY + logger.info("Broker runner ready: broker=%s", self.plugin.broker_id) + + def request_stop(self) -> None: + """Request graceful termination after the current poll/handle cycle.""" + self._stop_requested = True + if self._state not in {BrokerRunnerState.STOPPED, BrokerRunnerState.STOPPING}: + self._state = BrokerRunnerState.STOPPING + + def _handle_broker_failure(self, operation: str, exc: BaseException) -> None: + self._state = BrokerRunnerState.DEGRADED + logger.warning( + "Broker unavailable during %s: broker=%s error=%s; retrying in %.1fs", + operation, + self.plugin.broker_id, + type(exc).__name__, + self._backoff, + exc_info=True, + ) + self._sleep(self._backoff) + self._backoff = min(self.backoff_max, self._backoff * 2) + self._state = BrokerRunnerState.RECONNECTING + + def _apply_outcome(self, message: Any, outcome: TaskOutcome) -> None: + assert self._runtime is not None + consumer = self._runtime.consumer + try: + if outcome.disposition == TaskDisposition.ACK: + consumer.ack(message) + elif outcome.disposition == TaskDisposition.RETRY: + consumer.retry(message, outcome.error) + elif outcome.disposition == TaskDisposition.REJECT: + consumer.reject(message, outcome.error) + else: # pragma: no cover - StrEnum makes this defensive only. + raise ValueError(f"Unknown task disposition: {outcome.disposition!r}") + except Exception as exc: + # A failed ACK must leave the message recoverable. Do not invoke + # another ACK from here: Redis/Kafka providers own their delivery + # semantics and the next pending-recovery cycle will retry it. + self._handle_broker_failure("acknowledgement", exc) + return + + self._backoff = self.backoff_initial + self._state = BrokerRunnerState.READY + + def run_once(self) -> bool: + """Poll and process at most one message. + + Returns ``True`` when a message was received. Connection failures are + logged and delayed; they do not escape the runner boundary. + """ + if self._stop_requested: + return False + if self._runtime is None: + try: + self.start() + except Exception as exc: + self._handle_broker_failure("startup", exc) + return False + assert self._runtime is not None + try: + recovered = self._runtime.consumer.recover_pending() + if recovered: + logger.info( + "Recovered pending broker messages: broker=%s count=%d", + self.plugin.broker_id, + recovered, + ) + message = self._runtime.consumer.poll(self.poll_timeout_ms) + except Exception as exc: + self._handle_broker_failure("poll", exc) + return False + if message is None: + self._state = BrokerRunnerState.READY + return False + + try: + outcome = self._runtime.handle(message) + if not isinstance(outcome, TaskOutcome): + raise TypeError("BrokerRuntime.handle must return TaskOutcome") + except Exception as exc: + # Provider exceptions are treated as temporary by default. The + # provider can return ACK for permanent validation failures after + # best-effort FAILED reporting. + logger.warning( + "Broker task handling failed; retaining message for recovery: " + "broker=%s job_id=%s error=%s", + self.plugin.broker_id, + getattr(message, "job_id", None), + type(exc).__name__, + exc_info=True, + ) + outcome = TaskOutcome( + disposition=TaskDisposition.RETRY, + error=str(exc), + ) + self._apply_outcome(message, outcome) + return True + + def run(self) -> None: + """Run until :meth:`request_stop` is called or interrupted.""" + try: + while not self._stop_requested: + self.run_once() + except KeyboardInterrupt: + logger.info("Broker runner interrupted") + self.request_stop() + finally: + self.close() + + def close(self) -> None: + """Close provider resources and mark the runner stopped.""" + if self._runtime is not None: + self._state = BrokerRunnerState.STOPPING + try: + self._runtime.close() + except Exception: + logger.warning( + "Failed to close broker runtime: broker=%s", + self.plugin.broker_id, + exc_info=True, + ) + finally: + self._runtime = None + self._state = BrokerRunnerState.STOPPED diff --git a/src/tributo/plugin.py b/src/tributo/plugin.py index c7be8cf..5b06762 100644 --- a/src/tributo/plugin.py +++ b/src/tributo/plugin.py @@ -895,3 +895,149 @@ def _discover_storage_adapter_plugins( classes.append(cls) logger.info("Discovered storage adapter %r (%s)", ep.name, ep.value) return classes + + +# ═══════════════════════════════════════════════════════════════════════════════ +# Broker plugins +# ═══════════════════════════════════════════════════════════════════════════════ + + +def _broker_contract_issues(cls: Any) -> tuple[str, ...]: + """Return structural API issues without instantiating a provider.""" + issues: list[str] = [] + if not isinstance(cls, type): + return ("provider class",) + if not isinstance(getattr(cls, "api_version", None), int): + issues.append("api_version") + if not isinstance(getattr(cls, "broker_id", None), str): + issues.append("broker_id") + capabilities = getattr(cls, "capabilities", None) + if not isinstance(capabilities, frozenset): + issues.append("capabilities") + for method in ( + "validate_config", + "create_runtime", + "create_cancellation_checker", + ): + if not callable(getattr(cls, method, None)): + issues.append(method) + return tuple(issues) + + +def discover_broker_plugins( + diagnostics: list[PluginLoadDiagnostic] | None = None, +) -> list[type[Any]]: + """Discover broker provider classes from ``tributo.brokers``. + + Discovery is fail-open and never instantiates a provider. In particular, + it cannot create a Redis client or perform a network probe. Explicit + resolution is provided by :func:`resolve_broker_plugin` and is + fail-closed. + """ + from tributo.integrations.broker import BROKER_API_VERSION + + enabled = _get_enabled_plugins() + classes: list[type[Any]] = [] + for ep in _iter_entry_points("tributo.brokers"): + if enabled is not None and ep.name not in enabled: + logger.debug("Skipping broker plugin %r", ep.name) + continue + try: + cls = ep.load() + except Exception as exc: + logger.warning( + "Failed to load broker plugin %r (%s)", + ep.name, + ep.value, + exc_info=True, + ) + _record_diagnostic( + diagnostics, + "tributo.brokers", + ep.name, + f"Failed to load entry point: {exc}", + error_type=type(exc).__name__, + ) + continue + + issues = _broker_contract_issues(cls) + if issues: + logger.warning( + "Broker plugin %r does not satisfy API v%d: %s", + ep.name, + BROKER_API_VERSION, + ", ".join(issues), + ) + _record_diagnostic( + diagnostics, + "tributo.brokers", + ep.name, + "Missing or invalid BrokerPlugin members: " + ", ".join(issues), + ) + continue + if cls.api_version != BROKER_API_VERSION: + reason = ( + f"Unsupported BrokerPlugin api_version {cls.api_version!r}; " + f"expected {BROKER_API_VERSION}" + ) + _record_diagnostic(diagnostics, "tributo.brokers", ep.name, reason) + logger.warning("Broker plugin %r: %s", ep.name, reason) + continue + if ep.name != cls.broker_id: + reason = ( + f"Entry-point name {ep.name!r} does not match broker_id " + f"{cls.broker_id!r}" + ) + _record_diagnostic(diagnostics, "tributo.brokers", ep.name, reason) + logger.warning("Broker plugin %r: %s", ep.name, reason) + continue + classes.append(cls) + logger.info("Discovered broker plugin %r (%s)", ep.name, ep.value) + return classes + + +def resolve_broker_plugin(broker_id: str) -> type[Any]: + """Resolve one explicitly selected broker, using fail-closed semantics.""" + enabled = _get_enabled_plugins() + if enabled is not None and broker_id not in enabled: + raise JobConfigurationError( + f"Broker {broker_id!r} is disabled by TRIBUTO_PLUGINS" + ) + + matches = [ + ep for ep in _iter_entry_points("tributo.brokers") if ep.name == broker_id + ] + if not matches: + raise JobConfigurationError(f"Unknown broker {broker_id!r}") + if len(matches) > 1: + raise JobConfigurationError( + f"Multiple entry points are registered for broker {broker_id!r}" + ) + + ep = matches[0] + try: + cls = ep.load() + except Exception as exc: + raise JobConfigurationError( + f"Failed to load broker {broker_id!r} ({type(exc).__name__})" + ) from exc + + issues = _broker_contract_issues(cls) + from tributo.integrations.broker import BROKER_API_VERSION + + if issues: + raise JobConfigurationError( + f"Broker {broker_id!r} does not implement the BrokerPlugin v" + f"{BROKER_API_VERSION} contract: {', '.join(issues)}" + ) + if cls.api_version != BROKER_API_VERSION: + raise JobConfigurationError( + f"Broker {broker_id!r} has unsupported api_version " + f"{cls.api_version!r}; expected {BROKER_API_VERSION}" + ) + if cls.broker_id != ep.name: + raise JobConfigurationError( + f"Broker entry-point name {ep.name!r} does not match broker_id " + f"{cls.broker_id!r}" + ) + return cls diff --git a/src/tributo/training/job_submitter.py b/src/tributo/training/job_submitter.py index 0b1317b..3a07f81 100644 --- a/src/tributo/training/job_submitter.py +++ b/src/tributo/training/job_submitter.py @@ -6,7 +6,6 @@ from __future__ import annotations -import logging import time from collections.abc import Callable from pathlib import Path @@ -20,13 +19,19 @@ from tributo._common.submission_id import generate_submission_id from tributo.algorithms.api import EnvironmentSpec from tributo.algorithms.api.artifacts import AlgorithmArtifact, ImageProfile +from tributo.ray_jobs import RayJobSubmission, _submit_ray_job_with_client from tributo.util.annotations import PublicAPI -logger = logging.getLogger(__name__) DEFAULT_TIMEOUT = 180 JobAttemptStatus = Literal["PENDING", "RUNNING", "SUCCEEDED", "FAILED", "STOPPED"] TerminalJobStatus = Literal["SUCCEEDED", "FAILED", "STOPPED"] -_RESERVED_ENV_KEYS = frozenset({"TRIBUTO_RUN_ID", "TRIBUTO_ATTEMPT_ID"}) +_RESERVED_ENV_KEYS = frozenset( + { + "TRIBUTO_RUN_ID", + "TRIBUTO_ATTEMPT_ID", + "TRIBUTO_SUBMISSION_ID", + } +) def _resolve_algorithm_dependencies( @@ -49,7 +54,7 @@ class JobAttempt(BaseModel): run_id: str = Field(..., min_length=1) attempt_id: str = Field(..., min_length=1) submission_id: str = Field(..., min_length=1) - job_id: str = Field(..., min_length=1) + ray_job_id: str | None = Field(default=None, min_length=1) attempt_number: int = Field(..., ge=1) status: JobAttemptStatus retryable: bool = False @@ -63,7 +68,8 @@ class TrainingJobResult(BaseModel): run_id: str = Field(..., min_length=1) bundle_id: str = Field(..., min_length=1) - job_id: str = Field(..., min_length=1) + submission_id: str = Field(..., min_length=1) + ray_job_id: str | None = Field(default=None, min_length=1) status: TerminalJobStatus logs: str = "" attempts: tuple[JobAttempt, ...] = () @@ -106,35 +112,21 @@ def _submit_training_job_attempt( runtime_env: dict[str, Any], run_id: str, attempt_id: str, + submission_id: str, metadata: dict[str, str] | None = None, -) -> tuple[str, str]: - """Submit one stable attempt and reconcile ambiguous server responses.""" - submission_id = generate_submission_id("train", run_id, attempt_id) - try: - job_id = client.submit_job( - entrypoint=entrypoint, - runtime_env=runtime_env, - metadata=metadata, - submission_id=submission_id, - ) - except Exception as exc: - # The request may have reached Ray before the client observed an - # error. Query the deterministic submission ID before considering a - # retry; inventing another ID here could run the same attempt twice. - try: - status = client.get_job_status(submission_id) - except Exception as query_exc: - raise exc from query_exc - if status is None: - raise exc from None - logger.warning( - "Reconciled submission %s after ambiguous error (status=%s)", - submission_id, - _status_name(status), - ) - return submission_id, submission_id - logger.info("Submitted training job %s: %s", job_id, entrypoint) - return str(job_id), submission_id + request_digest: str | None = None, +) -> RayJobSubmission: + """Submit one stable attempt through the workload-neutral Core helper.""" + return _submit_ray_job_with_client( + client, + entrypoint=entrypoint, + runtime_env=runtime_env, + run_id=run_id, + attempt_id=attempt_id, + submission_id=submission_id, + metadata=metadata, + request_digest=request_digest, + ) @PublicAPI(stability="beta") @@ -152,6 +144,7 @@ def submit_training_job( image_profile: ImageProfile | None = None, declared_dependencies: tuple[str, ...] = (), environment: EnvironmentSpec | None = None, + request_digest: str | None = None, ) -> str: """Submit a training job via the Ray Jobs API. @@ -178,21 +171,62 @@ def submit_training_job( declared_dependencies: Additional PEP 508 constraints to preflight. environment: Optional formal or ``from_sklearn()`` EnvironmentSpec; its dependencies are merged into the same preflight. + request_digest: Optional credential-free request digest stored only as + Ray submission metadata. Returns: - Submitted job ID on success. + Deterministic Ray Jobs submission identity on success. Raises: RuntimeError: Submission failed. """ + return submit_training_job_with_identity( + entrypoint, + dashboard_url=dashboard_url, + env_vars=env_vars, + project_root=project_root, + extra_excludes=extra_excludes, + run_id=run_id, + attempt_id=attempt_id, + metadata=metadata, + algorithm_artifact=algorithm_artifact, + image_profile=image_profile, + declared_dependencies=declared_dependencies, + environment=environment, + request_digest=request_digest, + ).submission_id + + +@PublicAPI(stability="alpha") +def submit_training_job_with_identity( + entrypoint: str, + *, + dashboard_url: str = DEFAULT_DASHBOARD_URL, + env_vars: dict[str, str] | None = None, + project_root: Path | None = None, + extra_excludes: list[str] | None = None, + run_id: str | None = None, + attempt_id: str | None = None, + metadata: dict[str, str] | None = None, + algorithm_artifact: AlgorithmArtifact | None = None, + image_profile: ImageProfile | None = None, + declared_dependencies: tuple[str, ...] = (), + environment: EnvironmentSpec | None = None, + request_digest: str | None = None, +) -> RayJobSubmission: + """Submit one attempt and return workload-neutral Ray Jobs identity.""" resolved_run_id = _resolve_run_id(entrypoint, env_vars, run_id) resolved_attempt_id = attempt_id or "attempt-1" + submission_id = generate_submission_id( + "train", resolved_run_id, resolved_attempt_id + ) _validate_metadata(metadata) job_env_vars = dict(env_vars or {}) job_env_vars.update( { "TRIBUTO_RUN_ID": resolved_run_id, "TRIBUTO_ATTEMPT_ID": resolved_attempt_id, + "TRIBUTO_SUBMISSION_ID": submission_id, } ) runtime_env = build_runtime_env( @@ -208,15 +242,16 @@ def submit_training_job( ) client = _get_submission_client(dashboard_url) - job_id, _submission_id = _submit_training_job_attempt( + return _submit_training_job_attempt( client, entrypoint=entrypoint, runtime_env=runtime_env, run_id=resolved_run_id, attempt_id=resolved_attempt_id, + submission_id=submission_id, metadata=metadata, + request_digest=request_digest, ) - return job_id @PublicAPI(stability="beta") @@ -237,6 +272,7 @@ def submit_training_job_with_retry( image_profile: ImageProfile | None = None, declared_dependencies: tuple[str, ...] = (), environment: EnvironmentSpec | None = None, + request_digest: str | None = None, ) -> TrainingJobResult: """Submit, reconcile and optionally retry a training run. @@ -268,8 +304,10 @@ def submit_training_job_with_retry( for attempt_number in range(1, max_attempts + 1): attempt_id = f"attempt-{attempt_number}" + submission_id = generate_submission_id("train", resolved_run_id, attempt_id) attempt_env_vars = dict(job_env_vars) attempt_env_vars["TRIBUTO_ATTEMPT_ID"] = attempt_id + attempt_env_vars["TRIBUTO_SUBMISSION_ID"] = submission_id runtime_env = build_runtime_env( project_root=project_root, env_vars=attempt_env_vars, @@ -281,17 +319,19 @@ def submit_training_job_with_retry( environment, ), ) - job_id, submission_id = _submit_training_job_attempt( + submission = _submit_training_job_attempt( client, entrypoint=entrypoint, runtime_env=runtime_env, run_id=resolved_run_id, attempt_id=attempt_id, + submission_id=submission_id, metadata=metadata, + request_digest=request_digest, ) last_result = wait_for_job( client, - job_id, + submission.submission_id, timeout=timeout, poll_interval=poll_interval, ) @@ -307,8 +347,8 @@ def submit_training_job_with_retry( JobAttempt( run_id=resolved_run_id, attempt_id=attempt_id, - submission_id=submission_id, - job_id=job_id, + submission_id=submission.submission_id, + ray_job_id=submission.ray_job_id, attempt_number=attempt_number, status=status_name, retryable=retryable, @@ -325,7 +365,8 @@ def submit_training_job_with_retry( return TrainingJobResult( run_id=resolved_run_id, bundle_id=bundle_id_for_request(resolved_run_id), - job_id=attempts[-1].job_id, + submission_id=attempts[-1].submission_id, + ray_job_id=attempts[-1].ray_job_id, status=cast(TerminalJobStatus, final_status), logs=str(last_result.get("logs", "")), attempts=tuple(attempts), @@ -360,7 +401,8 @@ def wait_for_job( Args: client: Ray Jobs API client. - job_id: Job ID to wait for. + job_id: Ray Jobs submission identity to wait for. The parameter name + is retained for compatibility. timeout: Maximum wait time in seconds. poll_interval: Polling interval in seconds. diff --git a/src/tributo/training/lifecycle.py b/src/tributo/training/lifecycle.py index 819f778..d548940 100644 --- a/src/tributo/training/lifecycle.py +++ b/src/tributo/training/lifecycle.py @@ -173,6 +173,12 @@ def run( trainer.setup() checkpoint = trainer.training_loop() training_completed = True + checkpoint_metrics = getattr(checkpoint, "metrics", None) + if isinstance(checkpoint_metrics, dict): + # Preserve Ray Train metrics/history for both legacy export + # and Bundle paths. Provider reporters may replay this + # broker-neutral summary after the driver finishes. + summary["metrics"] = dict(checkpoint_metrics) self._dispatcher.on_training_end(trainer, checkpoint) @@ -449,6 +455,7 @@ def _export_bundle( "bundle_id": result.bundle_id, "execution_id": result.execution_id, "canonical_uri": result.canonical_uri, + "manifest_uri": getattr(result, "manifest_uri", None), "manifest_sha256": result.manifest_sha256, "artifacts": [ {"name": a.name, "format": a.format, "tree_digest": a.tree_digest} diff --git a/src/tributo/training/xgboost_trainer.py b/src/tributo/training/xgboost_trainer.py index d3ae389..0eab8f7 100644 --- a/src/tributo/training/xgboost_trainer.py +++ b/src/tributo/training/xgboost_trainer.py @@ -21,7 +21,6 @@ XGBOOST_DESCRIPTOR, build_legacy_spec, ) -from tributo.integrations.broker import CancellationChecker from tributo.training.base import BaseTrainer from tributo.training.checkpoint import ResumeConfig from tributo.training.resource import ( @@ -37,6 +36,8 @@ import ray.data from ray.train.xgboost import XGBoostTrainer + from tributo.training.results import TrainingResult + # XGBoost params reserved by Tributo: silently passing these as native # training parameters would change the execution path (e.g. external-memory # / data_iter) without going through the materialization-budget contract. @@ -528,30 +529,6 @@ def train_loop_per_worker(config: dict[str, Any]) -> None: # pass would double-count bytes and rows against the shared budget. test_labels: list[Any] = [] - # Cancel signal — inject CancellationChecker via config when using a broker. - # TODO(v1.1): _tributo_cancel_key and _tributo_cancel_checker are dead code - # until a broker implementation (e.g. tributo-broker-redis) populates them. - _cancel_key: str | None = config.get("_tributo_cancel_key") - _cancel_checker: CancellationChecker | None = config.get("_tributo_cancel_checker") - - class _CancelCallback(xgboost.callback.TrainingCallback): - """Check cancellation signal after each iteration (broker protocol).""" - - def after_iteration( - self, model: xgboost.Booster, epoch: int, evts_log: dict - ) -> bool: - if _cancel_key is None or _cancel_checker is None: - return False - try: - return _cancel_checker.is_cancelled(_cancel_key) - except Exception: - logger.warning( - "Cancellation check failed for job %s", - _cancel_key, - exc_info=True, - ) - return False # transient error → don't cancel - def _make_quantile_dmatrix( dataset_key: str, ref: xgboost.QuantileDMatrix | None = None, @@ -777,7 +754,7 @@ def after_iteration( evals=evals, evals_result=current_evals_result, early_stopping_rounds=config.get("early_stopping_rounds"), - callbacks=[_CancelCallback(), _ResumeCheckpointCallback()], + callbacks=[_ResumeCheckpointCallback()], xgb_model=initial_booster, ) evals_result = _merge_xgb_eval_results( @@ -1240,6 +1217,17 @@ def run_training_with_config(config: dict[str, Any]) -> dict[str, Any]: return trainer.run() +@PublicAPI(stability="alpha") +def run_training_result_with_config(config: dict[str, Any]) -> TrainingResult: + """Run training in-process and return the structured terminal result.""" + from tributo.training.results import TrainingResult + + summary = run_training_with_config(config) + return TrainingResult.model_validate( + {key: summary.get(key) for key in TrainingResult.model_fields if key in summary} + ) + + # Built-in registration _trainer_spec = build_legacy_spec( diff --git a/tests/test_broker.py b/tests/test_broker.py new file mode 100644 index 0000000..3544811 --- /dev/null +++ b/tests/test_broker.py @@ -0,0 +1,293 @@ +"""Core Broker SPI, lazy discovery, runner, and worker-context tests.""" + +from __future__ import annotations + +import json +from typing import Any, ClassVar +from unittest.mock import MagicMock + +import pytest + +import tributo.plugin as plugin +from tributo.exceptions import JobConfigurationError +from tributo.integrations.broker import ( + BROKER_API_VERSION, + BrokerPlugin, + BrokerRuntime, + CancellationChecker, + CancellationSpec, + JobResult, + Message, + TaskConsumer, + TaskDisposition, + TaskOutcome, +) +from tributo.integrations.broker_registry import ( + BrokerRegistry, + rebuild_cancellation_checker, +) +from tributo.integrations.broker_runner import BrokerRunner, BrokerRunnerState + + +class _EntryPoint: + def __init__(self, name: str, loaded: Any) -> None: + self.name = name + self.value = "tests:_Plugin" + self._loaded = loaded + + def load(self) -> Any: + if isinstance(self._loaded, Exception): + raise self._loaded + return self._loaded + + +class _Consumer(TaskConsumer): + def __init__(self, messages: list[Message]) -> None: + self.messages = messages + self.acked: list[Message] = [] + self.retried: list[Message] = [] + + def poll(self, timeout_ms: int = 5000) -> Message | None: + del timeout_ms + return self.messages.pop(0) if self.messages else None + + def ack(self, message: Message) -> None: + self.acked.append(message) + + def retry(self, message: Message, error: str | None = None) -> None: + del error + self.retried.append(message) + + +class _Runtime(BrokerRuntime): + def __init__(self, consumer: _Consumer, outcome: TaskOutcome) -> None: + self._consumer = consumer + self.outcome = outcome + + @property + def consumer(self) -> _Consumer: + return self._consumer + + def handle(self, message: Message) -> TaskOutcome: + del message + return self.outcome + + +class _Plugin(BrokerPlugin): + api_version: ClassVar[int] = BROKER_API_VERSION + broker_id: ClassVar[str] = "fake" + capabilities: ClassVar[frozenset[str]] = frozenset({"task-consumer"}) + runtime: _Runtime + + def validate_config(self, config, *, check_connectivity=False) -> None: + del config, check_connectivity + + def create_runtime(self, config) -> _Runtime: + del config + return self.runtime + + def create_cancellation_checker( + self, spec: CancellationSpec + ) -> CancellationChecker: + return _Checker(spec.job_id) + + +class _Checker(CancellationChecker): + def __init__(self, job_id: str) -> None: + self.job_id = job_id + + def is_cancelled(self, job_id: str) -> bool: + return job_id == self.job_id + + +def test_cancellation_spec_is_json_safe_and_rejects_client_objects() -> None: + spec = CancellationSpec("fake", "job-1", {"secret_ref": "env:REDIS_PASSWORD"}) + assert json.loads(json.dumps(spec.as_dict()))["job_id"] == "job-1" + with pytest.raises(ValueError, match="JSON serializable"): + CancellationSpec("fake", "job-1", {"client": object()}) + + +def test_discovery_is_lazy_and_records_import_diagnostics(monkeypatch) -> None: + monkeypatch.setattr( + plugin, + "_iter_entry_points", + lambda group: iter( + [ + _EntryPoint("broken", ImportError("redis secret")), + _EntryPoint("fake", _Plugin), + ] + if group == "tributo.brokers" + else [] + ), + ) + diagnostics = [] + classes = plugin.discover_broker_plugins(diagnostics) + assert classes == [_Plugin] + assert diagnostics[0].entry_point_name == "broken" + assert "redis secret" in diagnostics[0].reason + + +def test_discovery_rejects_non_frozen_capabilities(monkeypatch) -> None: + class _TupleCapabilitiesPlugin(_Plugin): + capabilities = ("task-consumer",) + + monkeypatch.setattr( + plugin, + "_iter_entry_points", + lambda group: iter( + [_EntryPoint("tuple", _TupleCapabilitiesPlugin)] + if group == "tributo.brokers" + else [] + ), + ) + diagnostics = [] + assert plugin.discover_broker_plugins(diagnostics) == [] + assert "capabilities" in diagnostics[0].reason + + +def test_discovery_rejects_api_version_and_entrypoint_identity_mismatch( + monkeypatch, +) -> None: + class _WrongVersionPlugin(_Plugin): + api_version = BROKER_API_VERSION + 1 + + class _WrongIdentityPlugin(_Plugin): + broker_id = "other" + + monkeypatch.setattr( + plugin, + "_iter_entry_points", + lambda group: iter( + [ + _EntryPoint("wrong-version", _WrongVersionPlugin), + _EntryPoint("fake", _WrongIdentityPlugin), + ] + if group == "tributo.brokers" + else [] + ), + ) + diagnostics = [] + assert plugin.discover_broker_plugins(diagnostics) == [] + assert len(diagnostics) == 2 + assert "api_version" in diagnostics[0].reason + assert "does not match broker_id" in diagnostics[1].reason + + +def test_explicit_broker_filtered_by_tributo_plugins_fails_closed(monkeypatch) -> None: + monkeypatch.setenv("TRIBUTO_PLUGINS", "another") + monkeypatch.setattr( + plugin, "_iter_entry_points", lambda group: iter([_EntryPoint("fake", _Plugin)]) + ) + with pytest.raises(JobConfigurationError, match="disabled"): + plugin.resolve_broker_plugin("fake") + + +def test_registry_reports_duplicate_broker_ids(monkeypatch) -> None: + monkeypatch.setattr( + "tributo.integrations.broker_registry.discover_broker_plugins", + lambda _diagnostics: [_Plugin, _Plugin], + ) + registry = BrokerRegistry() + assert len(registry.list()) == 1 + assert registry.diagnostics()[0].reason == "Duplicate broker_id discovered" + + +def test_runner_acks_only_ack_outcome_and_has_lifecycle_state() -> None: + message = Message("job-1", {}) + consumer = _Consumer([message]) + plugin_instance = _Plugin() + plugin_instance.runtime = _Runtime( + consumer, + TaskOutcome( + TaskDisposition.ACK, + result=JobResult("job-1", "accepted", run_id="job-1"), + ), + ) + runner = BrokerRunner(plugin_instance, {}) + assert runner.state == BrokerRunnerState.STOPPED + assert runner.run_once() is True + assert runner.state == BrokerRunnerState.READY + assert consumer.acked == [message] + runner.close() + assert runner.state == BrokerRunnerState.STOPPED + + +def test_runner_retains_retryable_message() -> None: + message = Message("job-1", {}) + consumer = _Consumer([message]) + plugin_instance = _Plugin() + plugin_instance.runtime = _Runtime( + consumer, + TaskOutcome(TaskDisposition.RETRY, error="ray unavailable"), + ) + runner = BrokerRunner(plugin_instance, {}) + assert runner.run_once() is True + assert consumer.acked == [] + assert consumer.retried == [message] + + +def test_runner_rejects_without_acknowledging() -> None: + message = Message("job-1", {}) + consumer = _Consumer([message]) + plugin_instance = _Plugin() + plugin_instance.runtime = _Runtime( + consumer, + TaskOutcome(TaskDisposition.REJECT, error="poison"), + ) + runner = BrokerRunner(plugin_instance, {}) + assert runner.run_once() is True + assert consumer.acked == [] + assert consumer.retried == [] + + +def test_runner_contains_ack_failure_and_enters_reconnect() -> None: + message = Message("job-1", {}) + consumer = _Consumer([message]) + consumer.ack = MagicMock(side_effect=ConnectionError("redis down")) + plugin_instance = _Plugin() + plugin_instance.runtime = _Runtime( + consumer, + TaskOutcome(TaskDisposition.ACK), + ) + runner = BrokerRunner( + plugin_instance, + {}, + backoff_initial=0.001, + backoff_max=0.001, + sleep=lambda _delay: None, + ) + assert runner.run_once() is True + assert runner.state == BrokerRunnerState.RECONNECTING + consumer.ack.assert_called_once_with(message) + + +def test_runner_graceful_stop_stops_next_poll() -> None: + consumer = _Consumer([]) + plugin_instance = _Plugin() + plugin_instance.runtime = _Runtime( + consumer, + TaskOutcome(TaskDisposition.ACK), + ) + runner = BrokerRunner(plugin_instance, {}) + runner.start() + runner.request_stop() + assert runner.state == BrokerRunnerState.STOPPING + assert runner.run_once() is False + runner.close() + assert runner.state == BrokerRunnerState.STOPPED + + +def test_worker_checker_rebuilt_from_spec(monkeypatch) -> None: + monkeypatch.setattr( + "tributo.integrations.broker_registry.BrokerRegistry.resolve", + lambda _self, broker_id: _Plugin(), + ) + checker = rebuild_cancellation_checker( + {"broker_id": "fake", "job_id": "job-1", "options": {}} + ) + assert isinstance(checker, _Checker) + assert checker.is_cancelled("job-1") is True + + +def test_missing_cancellation_context_keeps_training_context_empty() -> None: + assert rebuild_cancellation_checker(None) is None diff --git a/tests/test_broker_cli.py b/tests/test_broker_cli.py new file mode 100644 index 0000000..2bad877 --- /dev/null +++ b/tests/test_broker_cli.py @@ -0,0 +1,66 @@ +"""Core broker CLI isolation tests.""" + +from __future__ import annotations + +import json +import os +import subprocess +import sys +from pathlib import Path + +from click.testing import CliRunner + +from tributo.cli import main + + +def test_broker_config_is_json_only_and_provider_owned(tmp_path) -> None: + path = tmp_path / "broker.json" + path.write_text(json.dumps({"opaque_provider_field": True}), encoding="utf-8") + result = CliRunner().invoke( + main, ["broker", "validate", "--broker", "missing", "--config", str(path)] + ) + assert result.exit_code != 0 + assert "Unknown broker" in result.output + + +def test_normal_cli_does_not_require_redis() -> None: + result = CliRunner().invoke(main, ["--help"]) + assert result.exit_code == 0 + assert "broker" in result.output + + +def test_broker_list_without_provider_is_empty(monkeypatch) -> None: + monkeypatch.setattr("tributo.plugin._iter_entry_points", lambda _group: iter(())) + result = CliRunner().invoke(main, ["broker", "list"]) + assert result.exit_code == 0 + assert result.output == "" + + +def test_broker_consume_without_provider_fails_closed(tmp_path) -> None: + path = tmp_path / "broker.json" + path.write_text("{}", encoding="utf-8") + result = CliRunner().invoke( + main, + ["broker", "consume", "--broker", "missing", "--config", str(path)], + ) + assert result.exit_code != 0 + assert "Unknown broker" in result.output + + +def test_importing_core_cli_does_not_import_broker_module() -> None: + env = dict(os.environ) + env["PYTHONPATH"] = os.pathsep.join( + [str(Path(__file__).parents[1] / "src"), env.get("PYTHONPATH", "")] + ) + result = subprocess.run( + [ + sys.executable, + "-c", + "import sys; import tributo.cli; assert 'tributo.cli_broker' not in sys.modules", + ], + env=env, + check=False, + capture_output=True, + text=True, + ) + assert result.returncode == 0, result.stderr diff --git a/tests/test_runtime_env.py b/tests/test_runtime_env.py index 5b47fd8..0dab32d 100644 --- a/tests/test_runtime_env.py +++ b/tests/test_runtime_env.py @@ -48,3 +48,27 @@ def test_runtime_env_debug_log_never_exposes_environment_values( assert "PYTHONPATH" not in runtime_env["env_vars"] assert profile_payload not in caplog.text assert "TRIBUTO_STORAGE_PROFILE_MODEL" in caplog.text + + +def test_default_runtime_env_does_not_add_provider_dependencies(tmp_path) -> None: + (tmp_path / "pyproject.toml").write_text( + "[project]\nname='test'\n", encoding="utf-8" + ) + (tmp_path / "tributo").mkdir() + runtime_env = build_runtime_env(project_root=tmp_path) + assert runtime_env["py_modules"] == [str(tmp_path / "tributo")] + assert "pip" not in runtime_env + + +def test_runtime_env_can_explicitly_inject_provider_dependencies(tmp_path) -> None: + (tmp_path / "pyproject.toml").write_text( + "[project]\nname='test'\n", encoding="utf-8" + ) + (tmp_path / "tributo").mkdir() + runtime_env = build_runtime_env( + project_root=tmp_path, + extra_py_modules=[tmp_path / "provider"], + runtime_pip_packages=["tributo-broker-redis==0.1.0"], + ) + assert runtime_env["py_modules"][-1] == str(tmp_path / "provider") + assert runtime_env["pip"] == ["tributo-broker-redis==0.1.0"] diff --git a/tests/test_stability_inventory.py b/tests/test_stability_inventory.py index 19da5ef..979dfd8 100644 --- a/tests/test_stability_inventory.py +++ b/tests/test_stability_inventory.py @@ -30,6 +30,8 @@ "tributo.config": "stable", "tributo.job": "stable", "tributo.exceptions": "stable", + # Core — alpha + "tributo.ray_jobs": "alpha", # Core — beta "tributo.cli": "beta", # Portable algorithm execution — alpha @@ -155,6 +157,8 @@ "tributo.integrations.hooks": "beta", "tributo.integrations.sinks.parquet": "alpha", "tributo.integrations.sinks.lance": "alpha", + "tributo.integrations.broker": "alpha", + "tributo.integrations.broker_registry": "alpha", # Inference — beta "tributo.inference.base": "beta", "tributo.inference.batch_predictor": "beta", @@ -225,9 +229,13 @@ "tributo.inference.job_runner.InferenceJobResult": "alpha", "tributo.inference.job_runner.map_ray_job_status": "alpha", "tributo.inference.job_runner.submit_inference_request": "alpha", + "tributo.inference.job_runner.submit_inference_request_with_identity": "alpha", "tributo.inference.job_runner.submit_resolved_inference": "alpha", + "tributo.inference.job_runner.submit_resolved_inference_with_identity": "alpha", "tributo.inference.job_runner.submit_inference_request_with_retry": "alpha", "tributo.inference.job_runner.wait_for_job": "alpha", + "tributo.training.job_submitter.submit_training_job_with_identity": "alpha", + "tributo.training.xgboost_trainer.run_training_result_with_config": "alpha", } diff --git a/tests/training/test_job_submitter.py b/tests/training/test_job_submitter.py index 8de3f3e..e2f52ec 100644 --- a/tests/training/test_job_submitter.py +++ b/tests/training/test_job_submitter.py @@ -10,6 +10,7 @@ from tributo.algorithms import AlgorithmArtifact, EnvironmentSpec, ImageProfile from tributo.training.job_submitter import ( submit_training_job, + submit_training_job_with_identity, submit_training_job_with_retry, ) @@ -69,15 +70,15 @@ def test_same_run_and_attempt_reuse_submission_id(self) -> None: attempt_id="attempt-1", ) - assert first == second == "job-1" + assert first == second + assert first.startswith("tributo-train-") ids = [ call.kwargs["submission_id"] for call in client.submit_job.call_args_list ] assert ids[0] == ids[1] - assert ( - "TRIBUTO_RUN_ID" - in client.submit_job.call_args.kwargs["runtime_env"]["env_vars"] - ) + worker_env = client.submit_job.call_args.kwargs["runtime_env"]["env_vars"] + assert "TRIBUTO_RUN_ID" in worker_env + assert worker_env["TRIBUTO_SUBMISSION_ID"] == ids[-1] def test_existing_failed_attempt_is_reconciled_without_timestamp_retry( self, @@ -85,6 +86,9 @@ def test_existing_failed_attempt_is_reconciled_without_timestamp_retry( client = MagicMock() client.submit_job.side_effect = RuntimeError("submission already exists") client.get_job_status.return_value = JobStatus.FAILED + client.get_job_info.return_value = type( + "JobInfo", (), {"job_id": "ray-job-1"} + )() with ( patch( @@ -96,13 +100,14 @@ def test_existing_failed_attempt_is_reconciled_without_timestamp_retry( side_effect=_runtime_env, ), ): - job_id = submit_training_job( + submission = submit_training_job_with_identity( "python train.py", run_id="run-1", attempt_id="attempt-1", ) - assert job_id.startswith("tributo-train-") + assert submission.submission_id.startswith("tributo-train-") + assert submission.ray_job_id == "ray-job-1" assert client.submit_job.call_count == 1 @@ -138,6 +143,13 @@ def test_failed_job_uses_next_attempt_and_succeeded_stops(self) -> None: assert result.attempts[1].attempt_id == "attempt-2" assert result.attempts[1].status == "SUCCEEDED" assert result.attempts[0].submission_id != result.attempts[1].submission_id + for call, attempt in zip( + client.submit_job.call_args_list, result.attempts, strict=True + ): + assert ( + call.kwargs["runtime_env"]["env_vars"]["TRIBUTO_SUBMISSION_ID"] + == attempt.submission_id + ) def test_stopped_job_is_never_retried(self) -> None: client = MagicMock() @@ -259,9 +271,51 @@ def test_algorithm_artifact_and_environment_dependencies_reach_preflight( ), ) - assert job_id == "job-artifact" + assert job_id.startswith("tributo-train-") assert build_runtime_env.call_args.kwargs["algorithm_artifact"] is artifact assert build_runtime_env.call_args.kwargs["image_profile"] is profile assert build_runtime_env.call_args.kwargs["declared_dependencies"] == ( "scikit-learn<2,>=1.6", ) + + def test_metadata_cannot_override_submission_identity(self) -> None: + with pytest.raises(ValueError, match="TRIBUTO_SUBMISSION_ID"): + submit_training_job( + "python train.py", + run_id="run-1", + metadata={"TRIBUTO_SUBMISSION_ID": "other-submission"}, + ) + + def test_submission_result_preserves_identity_and_optional_digest(self) -> None: + client = MagicMock() + client.submit_job.return_value = "submission-return" + client.get_job_info.return_value = type( + "JobInfo", (), {"job_id": "ray-job-1"} + )() + with ( + patch( + "tributo.training.job_submitter._get_submission_client", + return_value=client, + ), + patch( + "tributo.training.job_submitter.build_runtime_env", + side_effect=_runtime_env, + ), + ): + result = submit_training_job_with_identity( + "python -m worker", + run_id="business-job-1", + attempt_id="attempt-2", + request_digest="request-digest", + ) + assert result.run_id == "business-job-1" + assert result.attempt_id == "attempt-2" + assert result.submission_id.startswith("tributo-train-") + assert result.ray_job_id == "ray-job-1" + assert result.request_digest == "request-digest" + worker_env = client.submit_job.call_args.kwargs["runtime_env"]["env_vars"] + assert worker_env["TRIBUTO_SUBMISSION_ID"] == result.submission_id + assert "TRIBUTO_EXECUTION_CONTEXT" not in worker_env + assert client.submit_job.call_args.kwargs["metadata"] == { + "tributo.request_digest": "request-digest" + } diff --git a/tests/training/test_training_lifecycle.py b/tests/training/test_training_lifecycle.py index a8694ba..bfffca9 100644 --- a/tests/training/test_training_lifecycle.py +++ b/tests/training/test_training_lifecycle.py @@ -71,6 +71,14 @@ def export_model(self, checkpoint: Any, output_path: str) -> None: self._summary["metrics"] = {"accuracy": 0.9} +class _MetricsCheckpointTrainer(_FakeTrainer): + def training_loop(self) -> Any: + self.events.append("training_loop") + return SimpleNamespace( + metrics={"eval-logloss_history": [0.8, 0.4], "eval-logloss": 0.4} + ) + + class _EntryTrainer(BaseTrainer): """Production-shaped trainer: accepts ``datasets``/``config`` like ``run_local_trial`` constructs them (``trainer_cls(datasets=..., @@ -353,6 +361,15 @@ def test_export_results_written_to_trainer_summary_are_returned(self) -> None: assert summary["metrics"] == {"accuracy": 0.9} assert trainer._summary is summary + def test_ray_checkpoint_metrics_are_preserved_for_replay(self) -> None: + trainer = _MetricsCheckpointTrainer() + summary = _lifecycle(trainer).run("/tmp/out") + + assert summary["metrics"] == { + "eval-logloss_history": [0.8, 0.4], + "eval-logloss": 0.4, + } + class TestBundleMode: def test_first_party_defaults_to_bundle_before_setup( From 9553630efa499363e5c05bfb6603a673880f3824 Mon Sep 17 00:00:00 2001 From: jiangxt2 Date: Fri, 14 Aug 2026 15:05:29 +0800 Subject: [PATCH 2/3] fix(training): stabilize source providers and submission identity Signed-off-by: jiangxt2 --- src/tributo/training/lifecycle.py | 6 +++++- 1 file changed, 5 insertions(+), 1 deletion(-) diff --git a/src/tributo/training/lifecycle.py b/src/tributo/training/lifecycle.py index d548940..5da4ba7 100644 --- a/src/tributo/training/lifecycle.py +++ b/src/tributo/training/lifecycle.py @@ -83,9 +83,13 @@ def _load_provider_plugins(registry: Any) -> None: registry.register(cls) if _provider_plugins_cache is None: + from tributo._bootstrap import first_party_source_provider_plugins from tributo.plugin import discover_source_provider_plugins - _provider_plugins_cache = discover_source_provider_plugins() + _provider_plugins_cache = [ + *first_party_source_provider_plugins(), + *discover_source_provider_plugins(), + ] for cls in _provider_plugins_cache: if cls.provider_id not in registry.list_all(): From f0d8e93c2af12db29575b920cd2589ae55ea72d9 Mon Sep 17 00:00:00 2001 From: jiangxt2 Date: Thu, 20 Aug 2026 09:58:57 +0800 Subject: [PATCH 3/3] refactor(broker): finalize Core-only Broker Alpha Signed-off-by: jiangxt2 --- docs/adr/002-broker-plugin-boundary.md | 100 ++++---- docs/architecture/index.md | 1 + docs/reference/api/algorithms-training.md | 6 + docs/reference/api/core.md | 31 ++- docs/reference/api/extensions.md | 42 ++++ docs/reference/api/inference-serving.md | 6 + src/tributo/cli.py | 6 +- src/tributo/cli_broker.py | 55 +--- src/tributo/inference/__init__.py | 4 + src/tributo/inference/job_runner.py | 133 ++++++---- src/tributo/inference/resolver.py | 2 +- src/tributo/integrations/broker.py | 242 ++++++------------ src/tributo/integrations/broker_registry.py | 70 ++---- src/tributo/integrations/broker_runner.py | 204 --------------- src/tributo/job.py | 15 +- src/tributo/plugin.py | 22 +- src/tributo/ray_jobs.py | 266 ++++++++++++++++++++ src/tributo/training/__init__.py | 17 +- src/tributo/training/lifecycle.py | 13 +- tests/docs/test_docs_tooling.py | 18 ++ tests/inference/test_job_runner.py | 13 +- tests/test_broker.py | 250 ++++++------------ tests/test_broker_cli.py | 42 ++-- tests/test_ray_jobs.py | 124 +++++++++ tests/test_retry.py | 9 +- tests/test_runtime_env.py | 16 +- tests/training/test_training_lifecycle.py | 17 -- tests/training/test_xgboost_trainer_unit.py | 24 ++ tools/generate_public_api_reference.py | 21 +- 29 files changed, 922 insertions(+), 847 deletions(-) delete mode 100644 src/tributo/integrations/broker_runner.py create mode 100644 src/tributo/ray_jobs.py create mode 100644 tests/test_ray_jobs.py diff --git a/docs/adr/002-broker-plugin-boundary.md b/docs/adr/002-broker-plugin-boundary.md index c13103d..f338d33 100644 --- a/docs/adr/002-broker-plugin-boundary.md +++ b/docs/adr/002-broker-plugin-boundary.md @@ -1,58 +1,66 @@ -# Broker plugin boundary +# Broker provider boundary ## Status -Accepted +Accepted for the Alpha API. ## Context -Tributo needs an optional control-plane integration for internal Redis -Streams tasks, lifecycle events, and cooperative cancellation. Redis Streams -is not a bounded data source or a streaming inference input. Future Kafka and -RabbitMQ providers should be able to use the same extension mechanism without -adding their client libraries to Tributo Core. +Tributo needs optional message-broker integrations without becoming a message +queue platform. A broker task is a control-plane admission request; bounded +data ingestion, unbounded inference streams, training, inference, Bundle +publication, and result sinks retain their existing Tributo contracts. + +Transport clients and external wire protocols must remain independently +installable. Core must be usable and testable without Redis, Kafka, RabbitMQ, +or another provider dependency. ## Decision -Tributo Core owns only a transport-neutral, beta Broker SPI, lazy entry-point -discovery, explicit provider resolution, a generic runner, and JSON-safe Ray -execution-context plumbing. Redis and KnoVa protocol models live in an -independently installable provider wheel. - -The Core contract has these safety rules: - -- ordinary Tributo startup and execution never import or connect to a broker; -- discovery is fail-open with diagnostics, while explicitly selected brokers - fail closed when missing, disabled, or invalid; -- broker configuration is passed as provider-owned JSON; Core does not define - Redis/Kafka/RabbitMQ fields or perform network probes implicitly; -- cancellation checkers are reconstructed in Ray workers from serializable - specs; clients, sockets, pools, and secrets are never serialized; -- provider submission must bind the business task ID to `run_id` and use a - deterministic submission ID per execution attempt. Transport delivery - retries must not become new execution attempts: the first Redis provider - reuses `attempt-1` and the same submission ID until an explicit business - retry is authorized after a terminal execution failure; -- temporary transport/submission failures retain the task for recovery; - permanently invalid messages are best-effort reported as FAILED and then - acknowledged even if reporting is unavailable; -- a missing or invalid outer `job_id` is permanently invalid and never becomes - a shared sentinel identity. Its FAILED event goes to a provider-owned - invalid-event stream and carries the delivery ID instead; -- provider reporters implement the Core `EventReporter` method signatures; - provider-specific fields such as KnoVa error codes use explicit extension - methods. Reporter warnings are time-window rate limited; -- reporter failure cannot turn successful training or Bundle publication into - a failed computation. - -The first Redis provider supports training tasks only. It reports lifecycle -events from the Ray Job driver and replays metrics history after training; -real-time metric sinks are a later extension. +Core owns a deliberately small Broker API v1 with Alpha stability: + +- `BrokerPlugin` discovery, structural version checks, capability metadata, + stability metadata, explicit config validation, and runtime construction; +- opaque `Message` payloads with a delivery token and restricted string + metadata; +- `TaskConsumer`, `BrokerRuntime`, `TaskDisposition`, and a minimal + `TaskOutcome` with an optional credential-safe `BrokerError`; +- workload-neutral `RayJobSubmission` identity and deterministic submission + IDs derived from an operation namespace, `run_id`, and `attempt_id`; +- ambiguous Ray submission reconciliation plus status and stop operations + keyed by `submission_id`. + +`BROKER_API_VERSION = 1` checks structural compatibility; it does not imply a +Beta or long-term compatibility promise. Discovery is lazy and fail-open with +diagnostics. Explicit resolution and configuration validation fail closed. +Discovery never instantiates a provider or performs connectivity checks. + +The following concerns belong to provider packages: + +- broker connections, polling, acknowledgments, re-delivery, recovery, dead + letters, cancellation watchers, and the production consume CLI/runtime; +- external request and event schemas, operation mapping, capability profiles, + credential references, error mapping, redaction, and event durability; +- structured terminal-event publication from existing `TrainingResult`, + `InferenceResult`, Bundle, and result-sink receipts. + +Core does not define a workload registry or an external operation schema. +Providers submit one thin execution-driver Ray Job through the generic helper; +that driver calls existing in-process training or batch-inference APIs. Worker +side broker cancellation, arbitrary execution context, a generic Core consume +loop, and durable workflow semantics are outside the Alpha contract. + +`submission_id` is the primary Ray Jobs identity for admission, status, logs, +and stop. `ray_job_id` is optional execution metadata and is populated only +from a real Ray `JobDetails.job_id`; Core never substitutes `submission_id` for +it. An optional credential-free `request_digest` may be recorded as Ray +metadata, but Core does not persist it or promise cross-restart conflict +detection. ## Consequences -The Core public surface can evolve independently of transport implementations, -and normal Tributo installations remain free of Redis dependencies. Provider -packages must publish their own protocol and infrastructure contract tests, -and a provider wheel must be installed in the Ray runtime when worker-side -cancellation or provider entrypoints are used. +Normal Tributo installations remain free of broker dependencies, and a +provider can evolve transport and protocol behavior independently. Providers +must own their infrastructure and healthy-path tests. The first release does +not promise exactly-once execution, durable terminal events, high availability, +or complete pending-message recovery. diff --git a/docs/architecture/index.md b/docs/architecture/index.md index 6c5928e..3432271 100644 --- a/docs/architecture/index.md +++ b/docs/architecture/index.md @@ -18,4 +18,5 @@ benchmark-protocol version-policy decision-log ../adr/001-data-and-bundle-contracts +../adr/002-broker-plugin-boundary ``` diff --git a/docs/reference/api/algorithms-training.md b/docs/reference/api/algorithms-training.md index 74b53ef..157b8cd 100644 --- a/docs/reference/api/algorithms-training.md +++ b/docs/reference/api/algorithms-training.md @@ -589,6 +589,9 @@ documentation for every public stability tier. ```{autofunction} tributo.training.job_submitter.submit_training_job ``` +```{autofunction} tributo.training.job_submitter.submit_training_job_with_identity +``` + ```{autofunction} tributo.training.job_submitter.submit_training_job_with_retry ``` @@ -708,3 +711,6 @@ documentation for every public stability tier. ```{autofunction} tributo.training.xgboost_trainer.run_training_from_json ``` + +```{autofunction} tributo.training.xgboost_trainer.run_training_result_with_config +``` diff --git a/docs/reference/api/core.md b/docs/reference/api/core.md index e731bb0..80c98d2 100644 --- a/docs/reference/api/core.md +++ b/docs/reference/api/core.md @@ -26,8 +26,7 @@ documentation for every public stability tier. ```{autoexception} tributo._common.dependencies.DependencyUnavailableError ``` -```{autoclass} tributo._common.dependencies.MissingOptionalDependency -:no-members: +```{autoexception} tributo._common.dependencies.MissingOptionalDependency ``` ```{autofunction} tributo._common.dependencies.probe_dependency @@ -73,8 +72,7 @@ documentation for every public stability tier. ## `tributo.exceptions` -```{autoclass} tributo.exceptions.AliasConflict -:no-members: +```{autoexception} tributo.exceptions.AliasConflict ``` ```{autoexception} tributo.exceptions.ArtifactCorruptedError @@ -131,8 +129,7 @@ documentation for every public stability tier. ```{autoexception} tributo.exceptions.ModelSchemaMismatchError ``` -```{autoclass} tributo.exceptions.PluginLoadIssue -:no-members: +```{autoexception} tributo.exceptions.PluginLoadIssue ``` ```{autoexception} tributo.exceptions.PostPublishCallbackError @@ -159,8 +156,7 @@ documentation for every public stability tier. ```{autoexception} tributo.TributoError ``` -```{autoclass} tributo.exceptions.UnsupportedArtifactFormat -:no-members: +```{autoexception} tributo.exceptions.UnsupportedArtifactFormat ``` @@ -173,3 +169,22 @@ documentation for every public stability tier. ```{autoclass} tributo.TributoClient :no-members: ``` + + +## `tributo.ray_jobs` + +```{autoclass} tributo.ray_jobs.RayJobSubmission +:no-members: +``` + +```{autofunction} tributo.ray_jobs.get_ray_job_logs +``` + +```{autofunction} tributo.ray_jobs.get_ray_job_status +``` + +```{autofunction} tributo.ray_jobs.stop_ray_job +``` + +```{autofunction} tributo.ray_jobs.submit_ray_job +``` diff --git a/docs/reference/api/extensions.md b/docs/reference/api/extensions.md index 9cbf6ed..8d27de9 100644 --- a/docs/reference/api/extensions.md +++ b/docs/reference/api/extensions.md @@ -9,6 +9,48 @@ a public annotation or moving a public object. Stable, Beta, and Alpha objects appear because Ray-style API policy requires documentation for every public stability tier. +## `tributo.integrations.broker` + +```{autoclass} tributo.integrations.broker.BrokerError +:no-members: +``` + +```{autoclass} tributo.integrations.broker.BrokerPlugin +:no-members: +``` + +```{autoclass} tributo.integrations.broker.BrokerRuntime +:no-members: +``` + +```{autoclass} tributo.integrations.broker.Message +:no-members: +``` + +```{autoclass} tributo.integrations.broker.TaskConsumer +:no-members: +``` + +```{autoclass} tributo.integrations.broker.TaskDisposition +:no-members: +``` + +```{autoclass} tributo.integrations.broker.TaskOutcome +:no-members: +``` + + +## `tributo.integrations.broker_registry` + +```{autoclass} tributo.integrations.broker_registry.BrokerDescriptor +:no-members: +``` + +```{autoclass} tributo.integrations.broker_registry.BrokerRegistry +:no-members: +``` + + ## `tributo.pipeline.core` ```{autoclass} tributo.pipeline.core.ArtifactRef diff --git a/docs/reference/api/inference-serving.md b/docs/reference/api/inference-serving.md index e446157..c63b96a 100644 --- a/docs/reference/api/inference-serving.md +++ b/docs/reference/api/inference-serving.md @@ -299,12 +299,18 @@ documentation for every public stability tier. ```{autofunction} tributo.inference.job_runner.submit_inference_request ``` +```{autofunction} tributo.inference.job_runner.submit_inference_request_with_identity +``` + ```{autofunction} tributo.inference.job_runner.submit_inference_request_with_retry ``` ```{autofunction} tributo.inference.job_runner.submit_resolved_inference ``` +```{autofunction} tributo.inference.job_runner.submit_resolved_inference_with_identity +``` + ```{autofunction} tributo.inference.job_runner.wait_for_job ``` diff --git a/src/tributo/cli.py b/src/tributo/cli.py index 96db56a..3ffb6b1 100644 --- a/src/tributo/cli.py +++ b/src/tributo/cli.py @@ -27,7 +27,7 @@ class _LazyTributoGroup(click.Group): - """Load the broker command module only when the broker command is used.""" + """Load the broker command module only when that command is selected.""" def get_command(self, ctx: click.Context, cmd_name: str) -> click.Command | None: if cmd_name == "broker": @@ -47,13 +47,13 @@ def format_commands( ctx: click.Context, formatter: click.HelpFormatter, ) -> None: - """Render broker help from a lightweight placeholder.""" + """Render root help without importing the broker command module.""" commands: list[tuple[str, click.Command]] = [] for command_name in self.list_commands(ctx): command = ( click.Command( "broker", - help="Discover and run explicitly selected message broker plugins.", + help="Discover and validate explicitly selected broker providers.", ) if command_name == "broker" else self.get_command(ctx, command_name) diff --git a/src/tributo/cli_broker.py b/src/tributo/cli_broker.py index 2b379d8..b292093 100644 --- a/src/tributo/cli_broker.py +++ b/src/tributo/cli_broker.py @@ -3,8 +3,6 @@ from __future__ import annotations import json -import logging -import signal from pathlib import Path from typing import Any @@ -12,9 +10,6 @@ from tributo.exceptions import JobConfigurationError from tributo.integrations.broker_registry import BrokerRegistry -from tributo.integrations.broker_runner import BrokerRunner - -logger = logging.getLogger(__name__) def _load_config(path: str) -> dict[str, Any]: @@ -32,7 +27,7 @@ def _load_config(path: str) -> dict[str, Any]: @click.group() def broker() -> None: - """Discover and run explicitly selected message broker plugins.""" + """Discover and validate explicitly selected broker providers.""" @broker.command("list") @@ -44,7 +39,7 @@ def broker_list() -> None: capabilities = ",".join(descriptor.capabilities) or "-" click.echo( f"{descriptor.broker_id}\tapi={descriptor.api_version}" - f"\tcapabilities={capabilities}" + f"\tstability={descriptor.stability}\tcapabilities={capabilities}" ) for diagnostic in registry.diagnostics(): click.echo( @@ -75,49 +70,3 @@ def broker_validate(broker_id: str, config_path: str, check_connectivity: bool) raise raise click.ClickException(str(exc)) from exc click.echo(f"Broker configuration is valid: {broker_id}") - - -@broker.command("consume") -@click.option("--broker", "broker_id", required=True) -@click.option("--config", "config_path", required=True, type=click.Path(exists=True)) -@click.option("--poll-timeout-ms", default=5000, show_default=True, type=int) -@click.option( - "--once", - is_flag=True, - help="Poll and handle at most one message; useful for controlled operations.", -) -def broker_consume( - broker_id: str, - config_path: str, - poll_timeout_ms: int, - once: bool, -) -> None: - """Consume tasks from an explicitly selected broker plugin.""" - config = _load_config(config_path) - registry = BrokerRegistry() - try: - plugin = registry.validate(broker_id, config) - runner = BrokerRunner( - plugin, - config, - poll_timeout_ms=poll_timeout_ms, - ) - except JobConfigurationError as exc: - raise click.ClickException(str(exc)) from exc - - def _stop(signum: int, _frame: Any) -> None: - logger.info( - "Received %s; stopping broker consumer", - signal.Signals(signum).name, - ) - runner.request_stop() - - signal.signal(signal.SIGINT, _stop) - signal.signal(signal.SIGTERM, _stop) - if once: - try: - runner.run_once() - finally: - runner.close() - else: - runner.run() diff --git a/src/tributo/inference/__init__.py b/src/tributo/inference/__init__.py index 1ca2c9a..3fe4c0b 100644 --- a/src/tributo/inference/__init__.py +++ b/src/tributo/inference/__init__.py @@ -40,7 +40,9 @@ def __call__(self, batch): ... from tributo.inference.job_runner import ( submit_inference_job, submit_inference_request, + submit_inference_request_with_identity, submit_resolved_inference, + submit_resolved_inference_with_identity, ) from tributo.inference.pipeline import ( InferenceConfig, @@ -72,5 +74,7 @@ def __call__(self, batch): ... "XGBoostONNXPredictor", "submit_inference_job", "submit_inference_request", + "submit_inference_request_with_identity", "submit_resolved_inference", + "submit_resolved_inference_with_identity", ] diff --git a/src/tributo/inference/job_runner.py b/src/tributo/inference/job_runner.py index 47280f6..9b7fd47 100644 --- a/src/tributo/inference/job_runner.py +++ b/src/tributo/inference/job_runner.py @@ -18,6 +18,7 @@ from tributo._common.submission_id import generate_submission_id from tributo.inference.contracts import InferenceRequest, ResolvedInference from tributo.inference.resolver import InferenceResolver +from tributo.ray_jobs import RayJobSubmission, _submit_ray_job_with_client from tributo.util.annotations import PublicAPI logger = logging.getLogger(__name__) @@ -52,7 +53,7 @@ class InferenceJobAttempt(BaseModel): run_id: str = Field(min_length=1) attempt_id: str = Field(min_length=1) submission_id: str = Field(min_length=1) - job_id: str = Field(min_length=1) + ray_job_id: str | None = Field(default=None, min_length=1) attempt_number: int = Field(ge=1) status: InferenceJobStatus retryable: bool = False @@ -69,7 +70,8 @@ class InferenceJobResult(BaseModel): ) run_id: str = Field(min_length=1) - job_id: str = Field(min_length=1) + submission_id: str = Field(min_length=1) + ray_job_id: str | None = Field(default=None, min_length=1) status: TerminalInferenceJobStatus logs: str = "" attempts: tuple[InferenceJobAttempt, ...] = () @@ -105,9 +107,7 @@ def submit_inference_job( resolved_run_id = run_id or generate_submission_id( "infer-run", config_path, str(sorted((env_vars or {}).items())) ) - submission_id = generate_submission_id( - "infer", resolved_run_id, attempt_id, config_path - ) + submission_id = generate_submission_id("infer", resolved_run_id, attempt_id) job_env = dict(env_vars or {}) job_env.update( { @@ -125,14 +125,20 @@ def submit_inference_job( f"python -m tributo.inference.batch_job --config {shlex.quote(config_path)}" ) client = _get_submission_client(dashboard_url) - job_id = _submit_attempt( + submission = _submit_attempt( client, entrypoint=entrypoint, runtime_env=runtime_env, + run_id=resolved_run_id, + attempt_id=attempt_id, submission_id=submission_id, ) - logger.info("Submitted inference job %s: config=%s", job_id, config_path) - return job_id + logger.info( + "Submitted inference job %s: config=%s", + submission.submission_id, + config_path, + ) + return submission.submission_id @PublicAPI(stability="alpha") @@ -144,17 +150,35 @@ def submit_inference_request( project_root: Path | None = None, resolver: InferenceResolver | None = None, ) -> str: - """Resolve once, serialize the credential-free plan, and submit it.""" + """Resolve once and return the accepted Ray Jobs submission identity.""" + return submit_inference_request_with_identity( + request, + dashboard_url=dashboard_url, + env_vars=env_vars, + project_root=project_root, + resolver=resolver, + ).submission_id + + +@PublicAPI(stability="alpha") +def submit_inference_request_with_identity( + request: InferenceRequest, + *, + dashboard_url: str = DEFAULT_DASHBOARD_URL, + env_vars: dict[str, str] | None = None, + project_root: Path | None = None, + resolver: InferenceResolver | None = None, +) -> RayJobSubmission: + """Resolve once, freeze the plan, and return complete submission identity.""" _validate_env_vars(env_vars) plan = (resolver or InferenceResolver()).resolve(request) client = _get_submission_client(dashboard_url) - job_id, _ = _submit_resolved_plan( + return _submit_resolved_plan( client, plan=plan, env_vars=env_vars, project_root=project_root, ) - return job_id @PublicAPI(stability="alpha") @@ -165,16 +189,32 @@ def submit_resolved_inference( env_vars: dict[str, str] | None = None, project_root: Path | None = None, ) -> str: - """Submit an already-frozen plan without re-resolving external aliases.""" + """Submit an already-frozen plan and return its submission identity.""" + return submit_resolved_inference_with_identity( + plan, + dashboard_url=dashboard_url, + env_vars=env_vars, + project_root=project_root, + ).submission_id + + +@PublicAPI(stability="alpha") +def submit_resolved_inference_with_identity( + plan: ResolvedInference, + *, + dashboard_url: str = DEFAULT_DASHBOARD_URL, + env_vars: dict[str, str] | None = None, + project_root: Path | None = None, +) -> RayJobSubmission: + """Submit a frozen plan and return workload-neutral Ray Jobs identity.""" _validate_env_vars(env_vars) client = _get_submission_client(dashboard_url) - job_id, _ = _submit_resolved_plan( + return _submit_resolved_plan( client, plan=plan, env_vars=env_vars, project_root=project_root, ) - return job_id @PublicAPI(stability="alpha") @@ -201,7 +241,7 @@ def submit_inference_request_with_retry( for attempt_number in range(1, max_attempts + 1): plan = _plan_for_attempt(first_plan, attempt_number) - job_id, submission_id = _submit_resolved_plan( + submission = _submit_resolved_plan( client, plan=plan, env_vars=env_vars, @@ -209,7 +249,7 @@ def submit_inference_request_with_retry( ) last = wait_for_job( client, - job_id, + submission.submission_id, timeout=timeout, poll_interval=poll_interval, ) @@ -225,8 +265,8 @@ def submit_inference_request_with_retry( InferenceJobAttempt( run_id=plan.run_id, attempt_id=plan.attempt_id, - submission_id=submission_id, - job_id=job_id, + submission_id=submission.submission_id, + ray_job_id=submission.ray_job_id, attempt_number=attempt_number, status=status, retryable=retryable, @@ -240,7 +280,8 @@ def submit_inference_request_with_retry( final_status = cast(TerminalInferenceJobStatus, attempts[-1].status) return InferenceJobResult( run_id=first_plan.run_id, - job_id=attempts[-1].job_id, + submission_id=attempts[-1].submission_id, + ray_job_id=attempts[-1].ray_job_id, status=final_status, logs=str(last.get("logs", "")), attempts=tuple(attempts), @@ -254,7 +295,7 @@ def _submit_resolved_plan( plan: ResolvedInference, env_vars: dict[str, str] | None, project_root: Path | None, -) -> tuple[str, str]: +) -> RayJobSubmission: plan = ResolvedInference.model_validate(plan.model_dump(mode="python")) encoded_plan = base64.urlsafe_b64encode( plan.model_dump_json().encode("utf-8") @@ -281,22 +322,25 @@ def _submit_resolved_plan( project_root=project_root, env_vars=job_env, ) - job_id = _submit_attempt( + submission = _submit_attempt( client, entrypoint=( "python -m tributo.inference.batch_job " "--resolved-plan-env TRIBUTO_INFERENCE_PLAN_B64" ), runtime_env=runtime_env, + run_id=plan.run_id, + attempt_id=plan.attempt_id, submission_id=plan.submission_id, + request_digest=plan.plan_digest, ) logger.info( - "Submitted inference attempt %s for run %s as Ray job %s", + "Submitted inference attempt %s for run %s as submission %s", plan.attempt_id, plan.run_id, - job_id, + submission.submission_id, ) - return job_id, plan.submission_id + return submission def _submit_attempt( @@ -304,38 +348,27 @@ def _submit_attempt( *, entrypoint: str, runtime_env: dict[str, Any], + run_id: str, + attempt_id: str, submission_id: str, -) -> str: - try: - return str( - client.submit_job( - entrypoint=entrypoint, - runtime_env=runtime_env, - submission_id=submission_id, - ) - ) - except Exception as exc: - try: - status = client.get_job_status(submission_id) - except Exception as query_exc: - raise exc from query_exc - if status is None: - raise exc from None - logger.warning( - "Reconciled inference submission %s after ambiguous error (status=%s)", - submission_id, - _ray_status_name(status), - ) - return submission_id + request_digest: str | None = None, +) -> RayJobSubmission: + return _submit_ray_job_with_client( + client, + entrypoint=entrypoint, + runtime_env=runtime_env, + run_id=run_id, + attempt_id=attempt_id, + submission_id=submission_id, + request_digest=request_digest, + ) def _plan_for_attempt( first_plan: ResolvedInference, attempt_number: int ) -> ResolvedInference: attempt_id = f"attempt-{attempt_number}" - submission_id = generate_submission_id( - "infer", first_plan.run_id, attempt_id, first_plan.plan_digest - ) + submission_id = generate_submission_id("infer", first_plan.run_id, attempt_id) return first_plan.model_copy( update={"attempt_id": attempt_id, "submission_id": submission_id} ) @@ -408,7 +441,9 @@ def _get_submission_client(dashboard_url: str) -> JobSubmissionClient: "map_ray_job_status", "submit_inference_job", "submit_inference_request", + "submit_inference_request_with_identity", "submit_inference_request_with_retry", "submit_resolved_inference", + "submit_resolved_inference_with_identity", "wait_for_job", ] diff --git a/src/tributo/inference/resolver.py b/src/tributo/inference/resolver.py index f3887af..6e83fea 100644 --- a/src/tributo/inference/resolver.py +++ b/src/tributo/inference/resolver.py @@ -117,7 +117,7 @@ def resolve(self, request: InferenceRequest) -> ResolvedInference: if os.environ.get("TRIBUTO_JOB_KIND") == "inference" else None ) or "attempt-1" - submission_id = generate_submission_id("infer", run_id, attempt_id, plan_digest) + submission_id = generate_submission_id("infer", run_id, attempt_id) return ResolvedInference( plan_digest=plan_digest, diff --git a/src/tributo/integrations/broker.py b/src/tributo/integrations/broker.py index 560e54a..61498f7 100644 --- a/src/tributo/integrations/broker.py +++ b/src/tributo/integrations/broker.py @@ -1,221 +1,135 @@ -"""Broker extension contracts for Tributo. +"""Transport-neutral contracts for independently installed broker providers. -The contracts in this module are deliberately transport-neutral. Redis, -Kafka, RabbitMQ, and protocol-specific models belong to separately installed -provider packages and are discovered only when a broker operation is -explicitly requested. +Core intentionally knows nothing about Redis, Kafka, RabbitMQ, or an external +operation protocol. Providers own transport semantics, request mapping, event +publication, and their production consume loop. """ from __future__ import annotations -import json from abc import ABC, abstractmethod from dataclasses import dataclass, field from enum import StrEnum -from typing import Any, Mapping +from types import MappingProxyType +from typing import Any, ClassVar, Mapping + +from tributo.util.annotations import PublicAPI BROKER_API_VERSION = 1 +@PublicAPI(stability="alpha") class TaskDisposition(StrEnum): - """Decision the generic runner applies after a provider handles a task.""" + """Provider decision for one broker delivery.""" ACK = "ack" RETRY = "retry" REJECT = "reject" -@dataclass +@PublicAPI(stability="alpha") +@dataclass(frozen=True) class Message: - """A job request consumed from a message queue. + """Opaque provider delivery passed across the Core Broker boundary. - Existing callers may continue constructing a message with only - ``job_id`` and ``payload``. Transport metadata is optional because the - public Core contract must also support brokers without Redis-style - delivery IDs. + ``delivery_token`` identifies the transport delivery, not a business + operation or Ray Job. Providers parse operation identity from ``payload``. + Metadata is restricted to string keys and values so transport clients and + credentials cannot be smuggled through this convenience surface. """ - job_id: str | None - payload: dict[str, Any] - metadata: dict[str, Any] = field(default_factory=dict) - delivery_id: str | None = None - delivery_attempt: int = 1 - run_id: str | None = None - attempt_id: str | None = None - - -@dataclass -class JobResult: - """Outcome of a completed ML training job. - - The identity fields are optional for backwards-compatible construction; - provider implementations should populate them whenever a broker task is - used. - """ + payload: Any + delivery_token: str + metadata: Mapping[str, str] = field(default_factory=dict) - job_id: str - status: str - metrics: dict[str, float] = field(default_factory=dict) - artifacts: list[str] = field(default_factory=list) - error: str | None = None - run_id: str | None = None - attempt_id: str | None = None - execution_id: str | None = None - submission_id: str | None = None - bundle_id: str | None = None - bundle_uri: str | None = None - manifest_uri: str | None = None - artifact_refs: list[dict[str, Any]] = field(default_factory=list) + def __post_init__(self) -> None: + if not isinstance(self.delivery_token, str) or not self.delivery_token.strip(): + raise ValueError("Message.delivery_token must not be empty") + if not all( + isinstance(key, str) and isinstance(value, str) + for key, value in self.metadata.items() + ): + raise ValueError("Message.metadata requires string keys and values") + object.__setattr__(self, "metadata", MappingProxyType(dict(self.metadata))) +@PublicAPI(stability="alpha") @dataclass(frozen=True) -class CancellationSpec: - """JSON-safe instructions for reconstructing a worker-side checker. +class BrokerError: + """Minimal credential-safe provider error attached to a delivery outcome.""" - This object intentionally contains no client, socket, pool, or secret. - ``options`` may contain provider-specific non-sensitive values or secret - references, but not credentials themselves. - """ - - broker_id: str - job_id: str - options: dict[str, Any] = field(default_factory=dict) + code: str + sanitized_message: str def __post_init__(self) -> None: - if not self.broker_id.strip(): - raise ValueError("CancellationSpec.broker_id must not be empty") - if not self.job_id.strip(): - raise ValueError("CancellationSpec.job_id must not be empty") - try: - json.dumps(self.as_dict(), allow_nan=False) - except (TypeError, ValueError) as exc: - raise ValueError("CancellationSpec must be JSON serializable") from exc - - def as_dict(self) -> dict[str, Any]: - """Return the wire-safe representation placed in Ray config.""" - return { - "broker_id": self.broker_id, - "job_id": self.job_id, - "options": dict(self.options), - } - - @classmethod - def from_mapping(cls, value: Mapping[str, Any]) -> "CancellationSpec": - """Validate and reconstruct a spec received by a Ray worker.""" - broker_id = value.get("broker_id") - job_id = value.get("job_id") - options = value.get("options", {}) - if not isinstance(broker_id, str) or not isinstance(job_id, str): - raise ValueError("CancellationSpec requires string broker_id and job_id") - if not isinstance(options, dict): - raise ValueError("CancellationSpec.options must be an object") - return cls(broker_id=broker_id, job_id=job_id, options=dict(options)) - - -@dataclass + if not isinstance(self.code, str) or not self.code.strip(): + raise ValueError("BrokerError.code must not be empty") + if not isinstance(self.sanitized_message, str): + raise TypeError("BrokerError.sanitized_message must be a string") + + +@PublicAPI(stability="alpha") +@dataclass(frozen=True) class TaskOutcome: - """Provider result consumed by :class:`BrokerRunner`.""" + """Transport-neutral disposition returned by a provider runtime.""" disposition: TaskDisposition - result: JobResult | None = None - error: str | None = None + error: BrokerError | None = None +@PublicAPI(stability="alpha") class TaskConsumer(ABC): - """Consume ML job requests from a message queue.""" + """Consume opaque deliveries from a provider-owned transport.""" @abstractmethod def poll(self, timeout_ms: int = 5000) -> Message | None: - """Block until a job request arrives or *timeout_ms* expires.""" + """Block until a delivery arrives or ``timeout_ms`` expires.""" ... @abstractmethod def ack(self, message: Message) -> None: - """Acknowledge successful processing of *message*.""" + """Acknowledge a delivery according to provider semantics.""" ... - def retry(self, message: Message, error: str | None = None) -> None: - """Leave a temporarily failed message eligible for redelivery. + def retry(self, message: Message, error: BrokerError | None = None) -> None: + """Apply provider-defined retry semantics. - Brokers such as Redis Streams retain a pending message by doing - nothing here. Brokers with an explicit negative acknowledgement may - override this method. + The default deliberately does nothing. Providers must override this + hook unless leaving the delivery pending is their explicit retry + policy. """ del message, error - def reject(self, message: Message, error: str | None = None) -> None: - """Reject a permanently invalid message without acknowledging it. + def reject(self, message: Message, error: BrokerError | None = None) -> None: + """Apply a provider-defined permanent rejection policy. - A provider should normally return ``ACK`` after publishing a FAILED - event for a poison message. This hook exists for transports that have - a dead-letter operation. + The default deliberately does nothing. A provider must override this + hook before returning ``REJECT`` for a delivery. """ del message, error def recover_pending(self) -> int: - """Reclaim pending messages when the transport supports it.""" + """Best-effort provider hook; Core makes no recovery guarantee.""" return 0 def close(self) -> None: - """Close transport resources; the default is intentionally a no-op.""" + """Close transport resources; the default is a no-op.""" return None -class EventReporter(ABC): - """Publish ML job lifecycle events.""" - - @abstractmethod - def report_phase(self, job_id: str, phase: str) -> None: - """Report a lifecycle phase transition.""" - ... - - @abstractmethod - def report_metrics( - self, job_id: str, metrics: dict[str, float], progress: float - ) -> None: - """Report metrics; v1 uses this for post-training history replay.""" - ... - - @abstractmethod - def report_completed(self, job_id: str, result: JobResult) -> None: - """Report job completion with final results.""" - ... - - @abstractmethod - def report_failed(self, job_id: str, error: str) -> None: - """Report job failure with error details.""" - ... - - def report_log(self, job_id: str, message: str, level: str = "INFO") -> None: - """Optionally report a user-visible log line.""" - del job_id, message, level - - def report_cancelled(self, job_id: str, phase: str = "TRAINING") -> None: - """Optionally report cooperative cancellation.""" - del job_id, phase - - -class CancellationChecker(ABC): - """Check whether a running job has been requested to cancel.""" - - @abstractmethod - def is_cancelled(self, job_id: str) -> bool: - """Return ``True`` if *job_id* should stop early.""" - ... - - +@PublicAPI(stability="alpha") class BrokerRuntime(ABC): - """Provider-owned runtime assembled by the generic Core runner.""" + """Provider runtime for mapping one delivery into a delivery outcome.""" @property @abstractmethod def consumer(self) -> TaskConsumer: - """Return the transport consumer.""" + """Return the provider-owned consumer.""" ... @abstractmethod def handle(self, message: Message) -> TaskOutcome: - """Validate and process one message without transport ACK side effects.""" + """Handle one message without applying transport ACK side effects.""" ... def close(self) -> None: @@ -223,28 +137,34 @@ def close(self) -> None: self.consumer.close() +@PublicAPI(stability="alpha") class BrokerPlugin(ABC): - """Structural base for an independently installed broker plugin.""" + """Structural base for an independently installed broker provider.""" - api_version: int = BROKER_API_VERSION - broker_id: str - capabilities: frozenset[str] = frozenset() + api_version: ClassVar[int] = BROKER_API_VERSION + broker_id: ClassVar[str] + capabilities: ClassVar[frozenset[str]] = frozenset() + stability: ClassVar[str] = "alpha" @abstractmethod def validate_config( self, config: Mapping[str, Any], *, check_connectivity: bool = False ) -> None: - """Validate provider-owned config and optionally probe connectivity.""" + """Validate provider config and optionally probe connectivity.""" ... @abstractmethod def create_runtime(self, config: Mapping[str, Any]) -> BrokerRuntime: - """Create a provider runtime; no network I/O belongs in discovery.""" + """Create a provider runtime; discovery itself must remain side-effect free.""" ... - @abstractmethod - def create_cancellation_checker( - self, spec: CancellationSpec - ) -> CancellationChecker: - """Rebuild a worker-side checker from a JSON-safe spec.""" - ... + +__all__ = [ + "BrokerError", + "BrokerPlugin", + "BrokerRuntime", + "Message", + "TaskConsumer", + "TaskDisposition", + "TaskOutcome", +] diff --git a/src/tributo/integrations/broker_registry.py b/src/tributo/integrations/broker_registry.py index b839d82..723ede9 100644 --- a/src/tributo/integrations/broker_registry.py +++ b/src/tributo/integrations/broker_registry.py @@ -1,47 +1,42 @@ -"""Lazy registry and worker-side reconstruction helpers for brokers.""" +"""Lazy discovery and explicit resolution for broker providers.""" from __future__ import annotations -import logging from collections.abc import Mapping from dataclasses import dataclass -from typing import Any +from typing import Any, cast from tributo.exceptions import JobConfigurationError from tributo.exporting.models import PluginLoadDiagnostic -from tributo.integrations.broker import ( - BrokerPlugin, - CancellationChecker, - CancellationSpec, -) +from tributo.integrations.broker import BrokerPlugin from tributo.plugin import discover_broker_plugins, resolve_broker_plugin from tributo.util.annotations import PublicAPI -logger = logging.getLogger(__name__) - +@PublicAPI(stability="alpha") @dataclass(frozen=True) class BrokerDescriptor: - """Metadata exposed by ``tributo broker list`` without instantiation.""" + """Side-effect-free metadata reported for one installed provider.""" broker_id: str api_version: int capabilities: tuple[str, ...] + stability: str -@PublicAPI(stability="beta") +@PublicAPI(stability="alpha") class BrokerRegistry: - """Resolve explicitly selected provider plugins on demand.""" + """Resolve only explicitly selected providers, failing closed on errors.""" def __init__(self) -> None: self._diagnostics: list[PluginLoadDiagnostic] = [] def diagnostics(self) -> tuple[PluginLoadDiagnostic, ...]: - """Return non-fatal discovery diagnostics from the last listing.""" + """Return non-fatal diagnostics from the most recent listing.""" return tuple(self._diagnostics) def list(self) -> tuple[BrokerDescriptor, ...]: - """List discoverable brokers without constructing or connecting them.""" + """List providers without constructing them or connecting to a broker.""" self._diagnostics.clear() descriptors: list[BrokerDescriptor] = [] seen: set[str] = set() @@ -57,18 +52,18 @@ def list(self) -> tuple[BrokerDescriptor, ...]: ) continue seen.add(broker_id) - capabilities = getattr(cls, "capabilities", frozenset()) descriptors.append( BrokerDescriptor( broker_id=broker_id, api_version=cls.api_version, - capabilities=tuple(sorted(str(value) for value in capabilities)), + capabilities=tuple(sorted(cls.capabilities)), + stability=cls.stability, ) ) return tuple(descriptors) def resolve(self, broker_id: str) -> BrokerPlugin: - """Load one explicitly selected provider, failing closed on errors.""" + """Load and instantiate one explicitly selected provider.""" cls = resolve_broker_plugin(broker_id) try: plugin = cls() @@ -76,7 +71,7 @@ def resolve(self, broker_id: str) -> BrokerPlugin: raise JobConfigurationError( f"Failed to initialize broker {broker_id!r} ({type(exc).__name__})" ) from exc - return plugin + return cast(BrokerPlugin, plugin) def validate( self, @@ -85,7 +80,7 @@ def validate( *, check_connectivity: bool = False, ) -> BrokerPlugin: - """Resolve and delegate provider-owned config validation.""" + """Resolve a provider and delegate its config validation.""" plugin = self.resolve(broker_id) try: plugin.validate_config( @@ -101,37 +96,4 @@ def validate( return plugin -def rebuild_cancellation_checker( - value: Mapping[str, Any] | CancellationSpec | None, -) -> CancellationChecker | None: - """Rebuild a checker in a Ray worker from JSON-safe config. - - This helper is intentionally fail-open: a missing or unavailable broker - must not change ordinary training into a failed training run. The error - is logged with the broker and job identity, while secrets remain in the - provider-owned config boundary. - """ - if value is None: - return None - try: - spec = ( - value - if isinstance(value, CancellationSpec) - else CancellationSpec.from_mapping(value) - ) - plugin = BrokerRegistry().resolve(spec.broker_id) - return plugin.create_cancellation_checker(spec) - except Exception: - broker_id = getattr(value, "broker_id", None) - if isinstance(value, Mapping): - broker_id = value.get("broker_id", broker_id) - job_id = getattr(value, "job_id", None) - if isinstance(value, Mapping): - job_id = value.get("job_id", job_id) - logger.warning( - "Unable to rebuild cancellation checker: broker=%s job_id=%s", - broker_id, - job_id, - exc_info=True, - ) - return None +__all__ = ["BrokerDescriptor", "BrokerRegistry"] diff --git a/src/tributo/integrations/broker_runner.py b/src/tributo/integrations/broker_runner.py deleted file mode 100644 index 694e89c..0000000 --- a/src/tributo/integrations/broker_runner.py +++ /dev/null @@ -1,204 +0,0 @@ -"""Transport-neutral broker consumer runner.""" - -from __future__ import annotations - -import logging -import time -from collections.abc import Callable, Mapping -from enum import StrEnum -from typing import Any - -from tributo.integrations.broker import ( - BrokerPlugin, - BrokerRuntime, - TaskDisposition, - TaskOutcome, -) -from tributo.util.annotations import PublicAPI - -logger = logging.getLogger(__name__) - - -class BrokerRunnerState(StrEnum): - """Observable lifecycle state of a broker consumer process.""" - - STARTING = "STARTING" - READY = "READY" - DEGRADED = "DEGRADED" - RECONNECTING = "RECONNECTING" - STOPPING = "STOPPING" - STOPPED = "STOPPED" - - -@PublicAPI(stability="beta") -class BrokerRunner: - """Run provider-owned message handling with isolated broker failures.""" - - def __init__( - self, - plugin: BrokerPlugin, - config: Mapping[str, Any], - *, - poll_timeout_ms: int = 5000, - backoff_initial: float = 1.0, - backoff_max: float = 30.0, - sleep: Callable[[float], None] = time.sleep, - ) -> None: - if poll_timeout_ms < 0: - raise ValueError("poll_timeout_ms must be non-negative") - if backoff_initial <= 0 or backoff_max < backoff_initial: - raise ValueError("invalid broker reconnect backoff") - self.plugin = plugin - self.config = dict(config) - self.poll_timeout_ms = poll_timeout_ms - self.backoff_initial = backoff_initial - self.backoff_max = backoff_max - self._sleep = sleep - self._runtime: BrokerRuntime | None = None - self._state = BrokerRunnerState.STOPPED - self._stop_requested = False - self._backoff = backoff_initial - - @property - def state(self) -> BrokerRunnerState: - """Return the current runner state.""" - return self._state - - @property - def runtime(self) -> BrokerRuntime | None: - """Expose the provider runtime for diagnostics and tests.""" - return self._runtime - - def start(self) -> None: - """Create the provider runtime without assuming broker availability.""" - if self._runtime is not None and self._state != BrokerRunnerState.STOPPED: - return - self._state = BrokerRunnerState.STARTING - self._stop_requested = False - self._runtime = self.plugin.create_runtime(self.config) - self._backoff = self.backoff_initial - self._state = BrokerRunnerState.READY - logger.info("Broker runner ready: broker=%s", self.plugin.broker_id) - - def request_stop(self) -> None: - """Request graceful termination after the current poll/handle cycle.""" - self._stop_requested = True - if self._state not in {BrokerRunnerState.STOPPED, BrokerRunnerState.STOPPING}: - self._state = BrokerRunnerState.STOPPING - - def _handle_broker_failure(self, operation: str, exc: BaseException) -> None: - self._state = BrokerRunnerState.DEGRADED - logger.warning( - "Broker unavailable during %s: broker=%s error=%s; retrying in %.1fs", - operation, - self.plugin.broker_id, - type(exc).__name__, - self._backoff, - exc_info=True, - ) - self._sleep(self._backoff) - self._backoff = min(self.backoff_max, self._backoff * 2) - self._state = BrokerRunnerState.RECONNECTING - - def _apply_outcome(self, message: Any, outcome: TaskOutcome) -> None: - assert self._runtime is not None - consumer = self._runtime.consumer - try: - if outcome.disposition == TaskDisposition.ACK: - consumer.ack(message) - elif outcome.disposition == TaskDisposition.RETRY: - consumer.retry(message, outcome.error) - elif outcome.disposition == TaskDisposition.REJECT: - consumer.reject(message, outcome.error) - else: # pragma: no cover - StrEnum makes this defensive only. - raise ValueError(f"Unknown task disposition: {outcome.disposition!r}") - except Exception as exc: - # A failed ACK must leave the message recoverable. Do not invoke - # another ACK from here: Redis/Kafka providers own their delivery - # semantics and the next pending-recovery cycle will retry it. - self._handle_broker_failure("acknowledgement", exc) - return - - self._backoff = self.backoff_initial - self._state = BrokerRunnerState.READY - - def run_once(self) -> bool: - """Poll and process at most one message. - - Returns ``True`` when a message was received. Connection failures are - logged and delayed; they do not escape the runner boundary. - """ - if self._stop_requested: - return False - if self._runtime is None: - try: - self.start() - except Exception as exc: - self._handle_broker_failure("startup", exc) - return False - assert self._runtime is not None - try: - recovered = self._runtime.consumer.recover_pending() - if recovered: - logger.info( - "Recovered pending broker messages: broker=%s count=%d", - self.plugin.broker_id, - recovered, - ) - message = self._runtime.consumer.poll(self.poll_timeout_ms) - except Exception as exc: - self._handle_broker_failure("poll", exc) - return False - if message is None: - self._state = BrokerRunnerState.READY - return False - - try: - outcome = self._runtime.handle(message) - if not isinstance(outcome, TaskOutcome): - raise TypeError("BrokerRuntime.handle must return TaskOutcome") - except Exception as exc: - # Provider exceptions are treated as temporary by default. The - # provider can return ACK for permanent validation failures after - # best-effort FAILED reporting. - logger.warning( - "Broker task handling failed; retaining message for recovery: " - "broker=%s job_id=%s error=%s", - self.plugin.broker_id, - getattr(message, "job_id", None), - type(exc).__name__, - exc_info=True, - ) - outcome = TaskOutcome( - disposition=TaskDisposition.RETRY, - error=str(exc), - ) - self._apply_outcome(message, outcome) - return True - - def run(self) -> None: - """Run until :meth:`request_stop` is called or interrupted.""" - try: - while not self._stop_requested: - self.run_once() - except KeyboardInterrupt: - logger.info("Broker runner interrupted") - self.request_stop() - finally: - self.close() - - def close(self) -> None: - """Close provider resources and mark the runner stopped.""" - if self._runtime is not None: - self._state = BrokerRunnerState.STOPPING - try: - self._runtime.close() - except Exception: - logger.warning( - "Failed to close broker runtime: broker=%s", - self.plugin.broker_id, - exc_info=True, - ) - finally: - self._runtime = None - self._state = BrokerRunnerState.STOPPED diff --git a/src/tributo/job.py b/src/tributo/job.py index 5070e62..703d1e6 100644 --- a/src/tributo/job.py +++ b/src/tributo/job.py @@ -36,8 +36,8 @@ class TributoClient: Example: >>> client = TributoClient("http://127.0.0.1:8265") - >>> job_id = client.submit(entrypoint="python script.py") - >>> status = client.get_status(job_id) + >>> submission_id = client.submit(entrypoint="python script.py") + >>> status = client.get_status(submission_id) """ def __init__(self, address: str): @@ -88,7 +88,7 @@ def submit( ``working_dir`` upload. Returns: - Job ID string. + Ray Jobs submission identity. Raises: JobSubmissionError: If submission fails. @@ -147,7 +147,8 @@ def get_status(self, job_id: str) -> str: """Get the status of a submitted job. Args: - job_id: The job ID to query. + job_id: Ray Jobs submission identity to query. The parameter name + is retained for compatibility. Returns: Job status string (e.g. ``"RUNNING"``, ``"SUCCEEDED"``). @@ -166,7 +167,8 @@ def get_logs(self, job_id: str) -> str: """Get logs for a submitted job. Args: - job_id: The job ID to query. + job_id: Ray Jobs submission identity to query. The parameter name + is retained for compatibility. Returns: Job logs as a string. @@ -184,7 +186,8 @@ def stop_job(self, job_id: str) -> bool: """Stop a running job. Args: - job_id: The job ID to stop. + job_id: Ray Jobs submission identity to stop. The parameter name + is retained for compatibility. Returns: True if the job was stopped successfully. diff --git a/src/tributo/plugin.py b/src/tributo/plugin.py index 5b06762..b48a064 100644 --- a/src/tributo/plugin.py +++ b/src/tributo/plugin.py @@ -907,17 +907,21 @@ def _broker_contract_issues(cls: Any) -> tuple[str, ...]: issues: list[str] = [] if not isinstance(cls, type): return ("provider class",) - if not isinstance(getattr(cls, "api_version", None), int): + if type(getattr(cls, "api_version", None)) is not int: issues.append("api_version") - if not isinstance(getattr(cls, "broker_id", None), str): + broker_id = getattr(cls, "broker_id", None) + if not isinstance(broker_id, str) or not broker_id.strip(): issues.append("broker_id") capabilities = getattr(cls, "capabilities", None) - if not isinstance(capabilities, frozenset): + if not isinstance(capabilities, frozenset) or not all( + isinstance(value, str) and bool(value.strip()) for value in capabilities + ): issues.append("capabilities") + if getattr(cls, "stability", None) not in {"alpha", "beta", "stable"}: + issues.append("stability") for method in ( "validate_config", "create_runtime", - "create_cancellation_checker", ): if not callable(getattr(cls, method, None)): issues.append(method) @@ -946,16 +950,16 @@ def discover_broker_plugins( cls = ep.load() except Exception as exc: logger.warning( - "Failed to load broker plugin %r (%s)", + "Failed to load broker plugin %r (%s; %s)", ep.name, ep.value, - exc_info=True, + type(exc).__name__, ) _record_diagnostic( diagnostics, "tributo.brokers", ep.name, - f"Failed to load entry point: {exc}", + f"Failed to load entry point ({type(exc).__name__})", error_type=type(exc).__name__, ) continue @@ -1016,7 +1020,7 @@ def resolve_broker_plugin(broker_id: str) -> type[Any]: ep = matches[0] try: - cls = ep.load() + cls = cast(type[Any], ep.load()) except Exception as exc: raise JobConfigurationError( f"Failed to load broker {broker_id!r} ({type(exc).__name__})" @@ -1040,4 +1044,4 @@ def resolve_broker_plugin(broker_id: str) -> type[Any]: f"Broker entry-point name {ep.name!r} does not match broker_id " f"{cls.broker_id!r}" ) - return cls + return cast(type[Any], cls) diff --git a/src/tributo/ray_jobs.py b/src/tributo/ray_jobs.py new file mode 100644 index 0000000..d61d82a --- /dev/null +++ b/src/tributo/ray_jobs.py @@ -0,0 +1,266 @@ +"""Workload-neutral Ray Jobs admission and control helpers.""" + +from __future__ import annotations + +import logging +from pathlib import Path +from typing import Any + +from pydantic import BaseModel, ConfigDict, Field +from ray.job_submission import JobSubmissionClient + +from tributo._common import DEFAULT_DASHBOARD_URL, build_runtime_env +from tributo._common.retry import retry_with_exponential_backoff +from tributo._common.submission_id import generate_submission_id +from tributo.util.annotations import PublicAPI + +logger = logging.getLogger(__name__) + +_RESERVED_ENV_KEYS = frozenset( + { + "TRIBUTO_RUN_ID", + "TRIBUTO_ATTEMPT_ID", + "TRIBUTO_SUBMISSION_ID", + } +) +_REQUEST_DIGEST_METADATA_KEY = "tributo.request_digest" + + +def _require_submission_id(submission_id: str) -> None: + if not isinstance(submission_id, str) or not submission_id.strip(): + raise ValueError("submission_id must not be empty") + + +@PublicAPI(stability="alpha") +class RayJobSubmission(BaseModel): + """Identity returned after one Ray Jobs attempt is accepted or reconciled.""" + + model_config = ConfigDict(frozen=True, extra="forbid") + + run_id: str = Field(min_length=1) + attempt_id: str = Field(min_length=1) + submission_id: str = Field(min_length=1) + ray_job_id: str | None = Field(default=None, min_length=1) + request_digest: str | None = Field(default=None, min_length=1) + + +def _validate_inputs( + operation_namespace: str, + run_id: str, + attempt_id: str, + env_vars: dict[str, str] | None, + metadata: dict[str, str] | None, + request_digest: str | None, +) -> None: + for name, value in ( + ("operation_namespace", operation_namespace), + ("run_id", run_id), + ("attempt_id", attempt_id), + ): + if not isinstance(value, str) or not value.strip(): + raise ValueError(f"{name} must not be empty") + conflicts = _RESERVED_ENV_KEYS.intersection(env_vars or {}) + if conflicts: + raise ValueError( + "env_vars must not override Ray Job identity: " + + ", ".join(sorted(conflicts)) + ) + if metadata is not None and _REQUEST_DIGEST_METADATA_KEY in metadata: + raise ValueError( + f"metadata must not define reserved key {_REQUEST_DIGEST_METADATA_KEY!r}" + ) + if request_digest is not None and not request_digest.strip(): + raise ValueError("request_digest must not be empty") + + +def _ray_job_id(client: JobSubmissionClient, submission_id: str) -> str | None: + get_job_info = getattr(client, "get_job_info", None) + if not callable(get_job_info): + return None + try: + info = get_job_info(submission_id) + except Exception as exc: + logger.debug( + "Ray JobDetails unavailable for submission %s (%s)", + submission_id, + type(exc).__name__, + ) + return None + value = getattr(info, "job_id", None) + return value if isinstance(value, str) and value else None + + +def _submit_ray_job_with_client( + client: JobSubmissionClient, + *, + entrypoint: str, + run_id: str, + attempt_id: str, + submission_id: str, + runtime_env: dict[str, Any] | None, + metadata: dict[str, str] | None = None, + request_digest: str | None = None, + entrypoint_num_cpus: float | None = None, + entrypoint_num_gpus: float | None = None, + entrypoint_memory: int | None = None, +) -> RayJobSubmission: + """Submit through an existing client; shared by Core workload adapters.""" + job_metadata = dict(metadata or {}) + if _REQUEST_DIGEST_METADATA_KEY in job_metadata: + raise ValueError( + f"metadata must not define reserved key {_REQUEST_DIGEST_METADATA_KEY!r}" + ) + if request_digest is not None: + if not request_digest.strip(): + raise ValueError("request_digest must not be empty") + job_metadata[_REQUEST_DIGEST_METADATA_KEY] = request_digest + + try: + client.submit_job( + entrypoint=entrypoint, + runtime_env=runtime_env, + metadata=job_metadata or None, + submission_id=submission_id, + entrypoint_num_cpus=entrypoint_num_cpus, + entrypoint_num_gpus=entrypoint_num_gpus, + entrypoint_memory=entrypoint_memory, + ) + except Exception as submit_error: + try: + status = client.get_job_status(submission_id) + except Exception as query_error: + raise submit_error from query_error + if status is None: + raise submit_error from None + logger.warning( + "Reconciled Ray submission %s after an ambiguous submit response", + submission_id, + ) + + return RayJobSubmission( + run_id=run_id, + attempt_id=attempt_id, + submission_id=submission_id, + ray_job_id=_ray_job_id(client, submission_id), + request_digest=request_digest, + ) + + +@PublicAPI(stability="alpha") +def submit_ray_job( + entrypoint: str, + *, + operation_namespace: str, + run_id: str, + attempt_id: str = "attempt-1", + dashboard_url: str = DEFAULT_DASHBOARD_URL, + env_vars: dict[str, str] | None = None, + project_root: Path | None = None, + extra_excludes: list[str] | None = None, + metadata: dict[str, str] | None = None, + request_digest: str | None = None, + entrypoint_num_cpus: float | None = None, + entrypoint_num_gpus: float | None = None, + entrypoint_memory: int | None = None, +) -> RayJobSubmission: + """Submit one deterministic Ray Job and reconcile an ambiguous response.""" + + if not entrypoint.strip(): + raise ValueError("entrypoint must not be empty") + _validate_inputs( + operation_namespace, + run_id, + attempt_id, + env_vars, + metadata, + request_digest, + ) + submission_id = generate_submission_id( + operation_namespace, + run_id, + attempt_id, + ) + job_env = dict(env_vars or {}) + job_env.update( + { + "TRIBUTO_RUN_ID": run_id, + "TRIBUTO_ATTEMPT_ID": attempt_id, + "TRIBUTO_SUBMISSION_ID": submission_id, + } + ) + runtime_env = build_runtime_env( + project_root=project_root, + env_vars=job_env, + extra_excludes=extra_excludes, + ) + client = _get_submission_client(dashboard_url) + return _submit_ray_job_with_client( + client, + entrypoint=entrypoint, + run_id=run_id, + attempt_id=attempt_id, + submission_id=submission_id, + runtime_env=runtime_env, + metadata=metadata, + request_digest=request_digest, + entrypoint_num_cpus=entrypoint_num_cpus, + entrypoint_num_gpus=entrypoint_num_gpus, + entrypoint_memory=entrypoint_memory, + ) + + +@PublicAPI(stability="alpha") +def get_ray_job_status( + submission_id: str, + *, + dashboard_url: str = DEFAULT_DASHBOARD_URL, +) -> str: + """Return the normalized Ray Jobs status for a submission identity.""" + + _require_submission_id(submission_id) + status = _get_submission_client(dashboard_url).get_job_status(submission_id) + if status is None: + raise LookupError(f"Unknown Ray submission {submission_id!r}") + return str(getattr(status, "value", status)).upper() + + +@PublicAPI(stability="alpha") +def get_ray_job_logs( + submission_id: str, + *, + dashboard_url: str = DEFAULT_DASHBOARD_URL, +) -> str: + """Return logs for the Ray Job identified by ``submission_id``.""" + + _require_submission_id(submission_id) + return _get_submission_client(dashboard_url).get_job_logs(submission_id) + + +@PublicAPI(stability="alpha") +def stop_ray_job( + submission_id: str, + *, + dashboard_url: str = DEFAULT_DASHBOARD_URL, +) -> bool: + """Request that Ray stop the job identified by ``submission_id``.""" + + _require_submission_id(submission_id) + return bool(_get_submission_client(dashboard_url).stop_job(submission_id)) + + +@retry_with_exponential_backoff( + max_retries=3, + base_delay=1.0, + exceptions=(ConnectionError, TimeoutError, OSError), +) +def _get_submission_client(dashboard_url: str) -> JobSubmissionClient: + return JobSubmissionClient(dashboard_url) + + +__all__ = [ + "RayJobSubmission", + "get_ray_job_logs", + "get_ray_job_status", + "stop_ray_job", + "submit_ray_job", +] diff --git a/src/tributo/training/__init__.py b/src/tributo/training/__init__.py index b9ca95a..73339bd 100644 --- a/src/tributo/training/__init__.py +++ b/src/tributo/training/__init__.py @@ -57,6 +57,7 @@ JobAttempt, TrainingJobResult, submit_training_job, + submit_training_job_with_identity, submit_training_job_with_retry, wait_for_job, ) @@ -67,7 +68,11 @@ ) from tributo.training.registry import get_trainer, list_trainers, register from tributo.training.tune_runner import TuneRunner, extract_best_params - from tributo.training.xgboost_trainer import build_trainer, run_training_from_json + from tributo.training.xgboost_trainer import ( + build_trainer, + run_training_from_json, + run_training_result_with_config, + ) _LAZY_EXPORTS = { "get_trainer": ("tributo.training.registry", "get_trainer"), @@ -86,6 +91,10 @@ "tributo.training.job_submitter", "submit_training_job_with_retry", ), + "submit_training_job_with_identity": ( + "tributo.training.job_submitter", + "submit_training_job_with_identity", + ), "wait_for_job": ("tributo.training.job_submitter", "wait_for_job"), "TuneRunner": ("tributo.training.tune_runner", "TuneRunner"), "extract_best_params": ( @@ -97,6 +106,10 @@ "tributo.training.xgboost_trainer", "run_training_from_json", ), + "run_training_result_with_config": ( + "tributo.training.xgboost_trainer", + "run_training_result_with_config", + ), "DNNTrainerImpl": ("tributo.training.dnn_trainer", "DNNTrainerImpl"), "run_dnn_training_from_json": ( "tributo.training.dnn_trainer", @@ -162,6 +175,7 @@ def __getattr__(name: str) -> Any: "JobAttempt", "TrainingJobResult", "submit_training_job", + "submit_training_job_with_identity", "submit_training_job_with_retry", "wait_for_job", "export_to_onnx", @@ -177,6 +191,7 @@ def __getattr__(name: str) -> Any: "warn_search_space_conflicts", "build_trainer", "run_training_from_json", + "run_training_result_with_config", ] if importlib.util.find_spec("torch") is not None: diff --git a/src/tributo/training/lifecycle.py b/src/tributo/training/lifecycle.py index 5da4ba7..819f778 100644 --- a/src/tributo/training/lifecycle.py +++ b/src/tributo/training/lifecycle.py @@ -83,13 +83,9 @@ def _load_provider_plugins(registry: Any) -> None: registry.register(cls) if _provider_plugins_cache is None: - from tributo._bootstrap import first_party_source_provider_plugins from tributo.plugin import discover_source_provider_plugins - _provider_plugins_cache = [ - *first_party_source_provider_plugins(), - *discover_source_provider_plugins(), - ] + _provider_plugins_cache = discover_source_provider_plugins() for cls in _provider_plugins_cache: if cls.provider_id not in registry.list_all(): @@ -177,12 +173,6 @@ def run( trainer.setup() checkpoint = trainer.training_loop() training_completed = True - checkpoint_metrics = getattr(checkpoint, "metrics", None) - if isinstance(checkpoint_metrics, dict): - # Preserve Ray Train metrics/history for both legacy export - # and Bundle paths. Provider reporters may replay this - # broker-neutral summary after the driver finishes. - summary["metrics"] = dict(checkpoint_metrics) self._dispatcher.on_training_end(trainer, checkpoint) @@ -459,7 +449,6 @@ def _export_bundle( "bundle_id": result.bundle_id, "execution_id": result.execution_id, "canonical_uri": result.canonical_uri, - "manifest_uri": getattr(result, "manifest_uri", None), "manifest_sha256": result.manifest_sha256, "artifacts": [ {"name": a.name, "format": a.format, "tree_digest": a.tree_digest} diff --git a/tests/docs/test_docs_tooling.py b/tests/docs/test_docs_tooling.py index 04bffd6..5e9d783 100644 --- a/tests/docs/test_docs_tooling.py +++ b/tests/docs/test_docs_tooling.py @@ -2,6 +2,7 @@ from __future__ import annotations +import ast import importlib import json import runpy @@ -128,6 +129,23 @@ def test_generated_public_api_reference_covers_source_inventory() -> None: assert check_pages(inventory) == [] +def test_public_api_inventory_classifies_exceptions_by_base_class() -> None: + tree = ast.parse( + """class BrokerError: + pass + +class ActualError(Exception): + pass +""" + ) + classes = [node for node in tree.body if isinstance(node, ast.ClassDef)] + + assert [public_api_generator._is_exception_class(node) for node in classes] == [ + False, + True, + ] + + def test_api_reference_validation_reuses_source_inventory( monkeypatch: pytest.MonkeyPatch, ) -> None: diff --git a/tests/inference/test_job_runner.py b/tests/inference/test_job_runner.py index 5630e79..2b654e7 100644 --- a/tests/inference/test_job_runner.py +++ b/tests/inference/test_job_runner.py @@ -13,6 +13,7 @@ from tributo.inference.contracts import ResolvedInference from tributo.inference.job_runner import ( map_ray_job_status, + submit_inference_job, submit_inference_request, submit_inference_request_with_retry, submit_resolved_inference, @@ -32,8 +33,6 @@ def test_passes_submission_id(self): "tributo.inference.job_runner.JobSubmissionClient", return_value=mock_client, ): - from tributo.inference.job_runner import submit_inference_job - submit_inference_job( config_path="jobs/inference.yaml", dashboard_url="http://127.0.0.1:8265", @@ -54,8 +53,6 @@ def test_submission_id_is_deterministic(self): "tributo.inference.job_runner.JobSubmissionClient", return_value=mock_client, ): - from tributo.inference.job_runner import submit_inference_job - submit_inference_job( config_path="jobs/inference.yaml", dashboard_url="http://127.0.0.1:8265", @@ -75,14 +72,12 @@ def test_client_creation_is_retried_on_connection_error(self): "tributo.inference.job_runner.JobSubmissionClient", side_effect=side_effects, ) as mock_constructor: - from tributo.inference.job_runner import submit_inference_job - job_id = submit_inference_job( config_path="jobs/inference.yaml", dashboard_url="http://127.0.0.1:8265", ) - assert job_id == "job-123" + assert job_id.startswith("tributo-infer-") assert mock_constructor.call_count == 2 @@ -113,7 +108,7 @@ def test_already_resolved_plan_is_submitted_without_a_resolver(self) -> None: ): job_id = submit_resolved_inference(plan) - assert job_id == "job-frozen" + assert job_id == plan.submission_id assert client.submit_job.call_args.kwargs["submission_id"] == plan.submission_id def test_frozen_plan_is_transported_without_re_resolution_in_job(self) -> None: @@ -135,7 +130,7 @@ def test_frozen_plan_is_transported_without_re_resolution_in_job(self) -> None: ): job_id = submit_inference_request(object(), resolver=resolver) - assert job_id == "job-1" + assert job_id == plan.submission_id resolver.resolve.assert_called_once() call = client.submit_job.call_args.kwargs assert call["submission_id"] == plan.submission_id diff --git a/tests/test_broker.py b/tests/test_broker.py index 3544811..b2099ba 100644 --- a/tests/test_broker.py +++ b/tests/test_broker.py @@ -1,32 +1,26 @@ -"""Core Broker SPI, lazy discovery, runner, and worker-context tests.""" +"""Core Broker API v1 and lazy provider-discovery tests.""" from __future__ import annotations -import json -from typing import Any, ClassVar -from unittest.mock import MagicMock +import operator +from typing import Any, ClassVar, cast import pytest +import tributo.integrations.broker as broker_contract import tributo.plugin as plugin from tributo.exceptions import JobConfigurationError from tributo.integrations.broker import ( BROKER_API_VERSION, + BrokerError, BrokerPlugin, BrokerRuntime, - CancellationChecker, - CancellationSpec, - JobResult, Message, TaskConsumer, TaskDisposition, TaskOutcome, ) -from tributo.integrations.broker_registry import ( - BrokerRegistry, - rebuild_cancellation_checker, -) -from tributo.integrations.broker_runner import BrokerRunner, BrokerRunnerState +from tributo.integrations.broker_registry import BrokerRegistry class _EntryPoint: @@ -42,69 +36,67 @@ def load(self) -> Any: class _Consumer(TaskConsumer): - def __init__(self, messages: list[Message]) -> None: - self.messages = messages - self.acked: list[Message] = [] - self.retried: list[Message] = [] - def poll(self, timeout_ms: int = 5000) -> Message | None: del timeout_ms - return self.messages.pop(0) if self.messages else None + return None def ack(self, message: Message) -> None: - self.acked.append(message) - - def retry(self, message: Message, error: str | None = None) -> None: - del error - self.retried.append(message) + del message class _Runtime(BrokerRuntime): - def __init__(self, consumer: _Consumer, outcome: TaskOutcome) -> None: - self._consumer = consumer - self.outcome = outcome - - @property - def consumer(self) -> _Consumer: - return self._consumer + consumer = _Consumer() def handle(self, message: Message) -> TaskOutcome: del message - return self.outcome + return TaskOutcome(TaskDisposition.ACK) class _Plugin(BrokerPlugin): api_version: ClassVar[int] = BROKER_API_VERSION broker_id: ClassVar[str] = "fake" capabilities: ClassVar[frozenset[str]] = frozenset({"task-consumer"}) - runtime: _Runtime + stability: ClassVar[str] = "alpha" def validate_config(self, config, *, check_connectivity=False) -> None: del config, check_connectivity def create_runtime(self, config) -> _Runtime: del config - return self.runtime + return _Runtime() - def create_cancellation_checker( - self, spec: CancellationSpec - ) -> CancellationChecker: - return _Checker(spec.job_id) +def test_message_keeps_payload_opaque_and_metadata_restricted() -> None: + payload = object() + message = Message( + payload, + "delivery-1", + metadata={"attempt": "1"}, + ) -class _Checker(CancellationChecker): - def __init__(self, job_id: str) -> None: - self.job_id = job_id - - def is_cancelled(self, job_id: str) -> bool: - return job_id == self.job_id - + assert message.payload is payload + assert message.metadata == {"attempt": "1"} + with pytest.raises(TypeError): + operator.setitem(message.metadata, "attempt", "2") + with pytest.raises(ValueError, match="string keys and values"): + Message( + {}, + "delivery-2", + metadata=cast(Any, {"attempt": 1}), + ) + + +def test_task_outcome_is_not_a_workload_result_contract() -> None: + outcome = TaskOutcome( + TaskDisposition.RETRY, + BrokerError(code="RAY_UNAVAILABLE", sanitized_message="retry later"), + ) -def test_cancellation_spec_is_json_safe_and_rejects_client_objects() -> None: - spec = CancellationSpec("fake", "job-1", {"secret_ref": "env:REDIS_PASSWORD"}) - assert json.loads(json.dumps(spec.as_dict()))["job_id"] == "job-1" - with pytest.raises(ValueError, match="JSON serializable"): - CancellationSpec("fake", "job-1", {"client": object()}) + assert outcome.error is not None + assert outcome.error.code == "RAY_UNAVAILABLE" + assert not hasattr(outcome, "result") + assert not hasattr(broker_contract, "JobResult") + assert not hasattr(broker_contract, "EventReporter") def test_discovery_is_lazy_and_records_import_diagnostics(monkeypatch) -> None: @@ -113,7 +105,7 @@ def test_discovery_is_lazy_and_records_import_diagnostics(monkeypatch) -> None: "_iter_entry_points", lambda group: iter( [ - _EntryPoint("broken", ImportError("redis secret")), + _EntryPoint("broken", ImportError("optional dependency unavailable")), _EntryPoint("fake", _Plugin), ] if group == "tributo.brokers" @@ -121,37 +113,46 @@ def test_discovery_is_lazy_and_records_import_diagnostics(monkeypatch) -> None: ), ) diagnostics = [] - classes = plugin.discover_broker_plugins(diagnostics) - assert classes == [_Plugin] - assert diagnostics[0].entry_point_name == "broken" - assert "redis secret" in diagnostics[0].reason + assert plugin.discover_broker_plugins(diagnostics) == [_Plugin] + assert diagnostics[0].entry_point_name == "broken" + assert diagnostics[0].error_type == "ImportError" + assert "optional dependency unavailable" not in diagnostics[0].reason -def test_discovery_rejects_non_frozen_capabilities(monkeypatch) -> None: - class _TupleCapabilitiesPlugin(_Plugin): - capabilities = ("task-consumer",) +@pytest.mark.parametrize( + ("attribute", "value"), + [ + ("capabilities", ("task-consumer",)), + ("stability", "prototype"), + ], +) +def test_discovery_rejects_invalid_provider_metadata( + monkeypatch, attribute: str, value: object +) -> None: + invalid = type("InvalidPlugin", (_Plugin,), {attribute: value}) monkeypatch.setattr( plugin, "_iter_entry_points", - lambda group: iter( - [_EntryPoint("tuple", _TupleCapabilitiesPlugin)] + lambda group: ( + iter([_EntryPoint("fake", invalid)]) if group == "tributo.brokers" - else [] + else iter(()) ), ) diagnostics = [] + assert plugin.discover_broker_plugins(diagnostics) == [] - assert "capabilities" in diagnostics[0].reason + assert attribute in diagnostics[0].reason -def test_discovery_rejects_api_version_and_entrypoint_identity_mismatch( +def test_discovery_rejects_version_and_entrypoint_identity_mismatch( monkeypatch, ) -> None: - class _WrongVersionPlugin(_Plugin): + class _WrongVersion(_Plugin): api_version = BROKER_API_VERSION + 1 - class _WrongIdentityPlugin(_Plugin): + class _WrongIdentity(_Plugin): broker_id = "other" monkeypatch.setattr( @@ -159,135 +160,42 @@ class _WrongIdentityPlugin(_Plugin): "_iter_entry_points", lambda group: iter( [ - _EntryPoint("wrong-version", _WrongVersionPlugin), - _EntryPoint("fake", _WrongIdentityPlugin), + _EntryPoint("wrong-version", _WrongVersion), + _EntryPoint("fake", _WrongIdentity), ] if group == "tributo.brokers" else [] ), ) diagnostics = [] + assert plugin.discover_broker_plugins(diagnostics) == [] - assert len(diagnostics) == 2 assert "api_version" in diagnostics[0].reason assert "does not match broker_id" in diagnostics[1].reason -def test_explicit_broker_filtered_by_tributo_plugins_fails_closed(monkeypatch) -> None: +def test_explicit_disabled_provider_fails_closed(monkeypatch) -> None: monkeypatch.setenv("TRIBUTO_PLUGINS", "another") monkeypatch.setattr( - plugin, "_iter_entry_points", lambda group: iter([_EntryPoint("fake", _Plugin)]) + plugin, + "_iter_entry_points", + lambda _group: iter([_EntryPoint("fake", _Plugin)]), ) + with pytest.raises(JobConfigurationError, match="disabled"): plugin.resolve_broker_plugin("fake") -def test_registry_reports_duplicate_broker_ids(monkeypatch) -> None: +def test_registry_reports_metadata_and_duplicate_ids(monkeypatch) -> None: monkeypatch.setattr( "tributo.integrations.broker_registry.discover_broker_plugins", lambda _diagnostics: [_Plugin, _Plugin], ) registry = BrokerRegistry() - assert len(registry.list()) == 1 - assert registry.diagnostics()[0].reason == "Duplicate broker_id discovered" - - -def test_runner_acks_only_ack_outcome_and_has_lifecycle_state() -> None: - message = Message("job-1", {}) - consumer = _Consumer([message]) - plugin_instance = _Plugin() - plugin_instance.runtime = _Runtime( - consumer, - TaskOutcome( - TaskDisposition.ACK, - result=JobResult("job-1", "accepted", run_id="job-1"), - ), - ) - runner = BrokerRunner(plugin_instance, {}) - assert runner.state == BrokerRunnerState.STOPPED - assert runner.run_once() is True - assert runner.state == BrokerRunnerState.READY - assert consumer.acked == [message] - runner.close() - assert runner.state == BrokerRunnerState.STOPPED - - -def test_runner_retains_retryable_message() -> None: - message = Message("job-1", {}) - consumer = _Consumer([message]) - plugin_instance = _Plugin() - plugin_instance.runtime = _Runtime( - consumer, - TaskOutcome(TaskDisposition.RETRY, error="ray unavailable"), - ) - runner = BrokerRunner(plugin_instance, {}) - assert runner.run_once() is True - assert consumer.acked == [] - assert consumer.retried == [message] - - -def test_runner_rejects_without_acknowledging() -> None: - message = Message("job-1", {}) - consumer = _Consumer([message]) - plugin_instance = _Plugin() - plugin_instance.runtime = _Runtime( - consumer, - TaskOutcome(TaskDisposition.REJECT, error="poison"), - ) - runner = BrokerRunner(plugin_instance, {}) - assert runner.run_once() is True - assert consumer.acked == [] - assert consumer.retried == [] - - -def test_runner_contains_ack_failure_and_enters_reconnect() -> None: - message = Message("job-1", {}) - consumer = _Consumer([message]) - consumer.ack = MagicMock(side_effect=ConnectionError("redis down")) - plugin_instance = _Plugin() - plugin_instance.runtime = _Runtime( - consumer, - TaskOutcome(TaskDisposition.ACK), - ) - runner = BrokerRunner( - plugin_instance, - {}, - backoff_initial=0.001, - backoff_max=0.001, - sleep=lambda _delay: None, - ) - assert runner.run_once() is True - assert runner.state == BrokerRunnerState.RECONNECTING - consumer.ack.assert_called_once_with(message) + descriptors = registry.list() -def test_runner_graceful_stop_stops_next_poll() -> None: - consumer = _Consumer([]) - plugin_instance = _Plugin() - plugin_instance.runtime = _Runtime( - consumer, - TaskOutcome(TaskDisposition.ACK), - ) - runner = BrokerRunner(plugin_instance, {}) - runner.start() - runner.request_stop() - assert runner.state == BrokerRunnerState.STOPPING - assert runner.run_once() is False - runner.close() - assert runner.state == BrokerRunnerState.STOPPED - - -def test_worker_checker_rebuilt_from_spec(monkeypatch) -> None: - monkeypatch.setattr( - "tributo.integrations.broker_registry.BrokerRegistry.resolve", - lambda _self, broker_id: _Plugin(), - ) - checker = rebuild_cancellation_checker( - {"broker_id": "fake", "job_id": "job-1", "options": {}} - ) - assert isinstance(checker, _Checker) - assert checker.is_cancelled("job-1") is True - - -def test_missing_cancellation_context_keeps_training_context_empty() -> None: - assert rebuild_cancellation_checker(None) is None + assert descriptors[0].broker_id == "fake" + assert descriptors[0].stability == "alpha" + assert descriptors[0].capabilities == ("task-consumer",) + assert registry.diagnostics()[0].reason == "Duplicate broker_id discovered" diff --git a/tests/test_broker_cli.py b/tests/test_broker_cli.py index 2bad877..c72624d 100644 --- a/tests/test_broker_cli.py +++ b/tests/test_broker_cli.py @@ -19,48 +19,56 @@ def test_broker_config_is_json_only_and_provider_owned(tmp_path) -> None: result = CliRunner().invoke( main, ["broker", "validate", "--broker", "missing", "--config", str(path)] ) + assert result.exit_code != 0 assert "Unknown broker" in result.output -def test_normal_cli_does_not_require_redis() -> None: +def test_normal_cli_does_not_require_a_broker_provider() -> None: result = CliRunner().invoke(main, ["--help"]) + assert result.exit_code == 0 assert "broker" in result.output def test_broker_list_without_provider_is_empty(monkeypatch) -> None: monkeypatch.setattr("tributo.plugin._iter_entry_points", lambda _group: iter(())) + result = CliRunner().invoke(main, ["broker", "list"]) + assert result.exit_code == 0 assert result.output == "" -def test_broker_consume_without_provider_fails_closed(tmp_path) -> None: - path = tmp_path / "broker.json" - path.write_text("{}", encoding="utf-8") - result = CliRunner().invoke( - main, - ["broker", "consume", "--broker", "missing", "--config", str(path)], - ) - assert result.exit_code != 0 - assert "Unknown broker" in result.output +def test_core_cli_has_no_provider_consume_loop() -> None: + result = CliRunner().invoke(main, ["broker", "--help"]) + + assert result.exit_code == 0 + assert "validate" in result.output + assert "consume" not in result.output -def test_importing_core_cli_does_not_import_broker_module() -> None: +def test_import_and_root_help_do_not_import_broker_module() -> None: env = dict(os.environ) env["PYTHONPATH"] = os.pathsep.join( [str(Path(__file__).parents[1] / "src"), env.get("PYTHONPATH", "")] ) + script = """ +import sys +from click.testing import CliRunner +import tributo.cli as cli +assert 'tributo.cli_broker' not in sys.modules +result = CliRunner().invoke(cli.main, ['--help']) +assert result.exit_code == 0, result.output +assert 'broker' in result.output +assert 'tributo.cli_broker' not in sys.modules +""" result = subprocess.run( - [ - sys.executable, - "-c", - "import sys; import tributo.cli; assert 'tributo.cli_broker' not in sys.modules", - ], + [sys.executable, "-c", script], env=env, check=False, capture_output=True, text=True, ) - assert result.returncode == 0, result.stderr + + assert result.returncode == 0, result.stderr or result.stdout diff --git a/tests/test_ray_jobs.py b/tests/test_ray_jobs.py new file mode 100644 index 0000000..7c180bb --- /dev/null +++ b/tests/test_ray_jobs.py @@ -0,0 +1,124 @@ +"""Workload-neutral Ray Jobs submission identity tests.""" + +from __future__ import annotations + +from typing import Any +from unittest.mock import MagicMock, patch + +import pytest +from ray.job_submission import JobStatus + +from tributo.ray_jobs import ( + RayJobSubmission, + get_ray_job_logs, + get_ray_job_status, + stop_ray_job, + submit_ray_job, +) + + +def _runtime_env(*args: Any, **kwargs: Any) -> dict[str, Any]: + del args + return {"env_vars": kwargs.get("env_vars", {})} + + +def test_submission_identity_is_workload_neutral_and_ray_job_id_is_real() -> None: + client = MagicMock() + client.submit_job.return_value = "ray-api-return-value" + client.get_job_info.return_value = type( + "JobInfo", (), {"job_id": "ray-core-job-1"} + )() + + with ( + patch("tributo.ray_jobs._get_submission_client", return_value=client), + patch("tributo.ray_jobs.build_runtime_env", side_effect=_runtime_env), + ): + result = submit_ray_job( + "python -m provider.driver", + operation_namespace="broker", + run_id="run-1", + attempt_id="attempt-2", + ) + + assert isinstance(result, RayJobSubmission) + assert result.run_id == "run-1" + assert result.attempt_id == "attempt-2" + assert result.submission_id.startswith("tributo-broker-") + assert result.ray_job_id == "ray-core-job-1" + assert client.submit_job.call_args.kwargs["submission_id"] == result.submission_id + + +def test_ambiguous_submission_reconciles_by_submission_id() -> None: + client = MagicMock() + client.submit_job.side_effect = TimeoutError("response lost") + client.get_job_status.return_value = JobStatus.RUNNING + client.get_job_info.side_effect = LookupError("driver not created") + + with ( + patch("tributo.ray_jobs._get_submission_client", return_value=client), + patch("tributo.ray_jobs.build_runtime_env", side_effect=_runtime_env), + ): + result = submit_ray_job( + "python -m provider.driver", + operation_namespace="broker", + run_id="run-1", + ) + + client.get_job_status.assert_called_once_with(result.submission_id) + assert result.ray_job_id is None + + +def test_request_digest_is_optional_metadata_not_submission_identity() -> None: + client = MagicMock() + client.get_job_info.return_value = type("JobInfo", (), {"job_id": None})() + + with ( + patch("tributo.ray_jobs._get_submission_client", return_value=client), + patch("tributo.ray_jobs.build_runtime_env", side_effect=_runtime_env), + ): + first = submit_ray_job( + "python -m provider.driver", + operation_namespace="broker", + run_id="run-1", + request_digest="digest-a", + ) + second = submit_ray_job( + "python -m provider.driver", + operation_namespace="broker", + run_id="run-1", + request_digest="digest-b", + ) + + assert first.submission_id == second.submission_id + assert client.submit_job.call_args_list[0].kwargs["metadata"] == { + "tributo.request_digest": "digest-a" + } + assert client.submit_job.call_args_list[1].kwargs["metadata"] == { + "tributo.request_digest": "digest-b" + } + + +def test_reserved_identity_environment_is_rejected_before_submission() -> None: + with pytest.raises(ValueError, match="must not override Ray Job identity"): + submit_ray_job( + "python -m provider.driver", + operation_namespace="broker", + run_id="run-1", + env_vars={"TRIBUTO_SUBMISSION_ID": "external"}, + ) + + +def test_status_and_stop_use_submission_identity() -> None: + client = MagicMock() + client.get_job_status.return_value = JobStatus.RUNNING + client.get_job_logs.return_value = "driver logs" + client.stop_job.return_value = True + + with patch("tributo.ray_jobs._get_submission_client", return_value=client): + assert get_ray_job_status("submission-1") == "RUNNING" + assert get_ray_job_logs("submission-1") == "driver logs" + assert stop_ray_job("submission-1") is True + + client.get_job_status.assert_called_once_with("submission-1") + client.get_job_logs.assert_called_once_with("submission-1") + client.stop_job.assert_called_once_with("submission-1") diff --git a/tests/test_retry.py b/tests/test_retry.py index d36bbbf..81b066d 100644 --- a/tests/test_retry.py +++ b/tests/test_retry.py @@ -7,6 +7,7 @@ import pytest +import tributo.training.job_submitter as training_job_submitter from tributo._common.retry import retry_with_exponential_backoff @@ -89,19 +90,17 @@ def test_training_job_submitter_retries_connection(self): client_mock = MagicMock() client_mock.submit_job.return_value = "job-456" - import tributo.training.job_submitter as tjs - with patch.object( - tjs, + training_job_submitter, "JobSubmissionClient", side_effect=[ ConnectionError("timeout"), client_mock, ], ) as mock_jsc: - job_id = tjs.submit_training_job("python train.py") + job_id = training_job_submitter.submit_training_job("python train.py") - assert job_id == "job-456" + assert job_id.startswith("tributo-train-") assert mock_jsc.call_count == 2 diff --git a/tests/test_runtime_env.py b/tests/test_runtime_env.py index 0dab32d..61426be 100644 --- a/tests/test_runtime_env.py +++ b/tests/test_runtime_env.py @@ -50,7 +50,7 @@ def test_runtime_env_debug_log_never_exposes_environment_values( assert "TRIBUTO_STORAGE_PROFILE_MODEL" in caplog.text -def test_default_runtime_env_does_not_add_provider_dependencies(tmp_path) -> None: +def test_default_runtime_env_does_not_add_extension_dependencies(tmp_path) -> None: (tmp_path / "pyproject.toml").write_text( "[project]\nname='test'\n", encoding="utf-8" ) @@ -58,17 +58,3 @@ def test_default_runtime_env_does_not_add_provider_dependencies(tmp_path) -> Non runtime_env = build_runtime_env(project_root=tmp_path) assert runtime_env["py_modules"] == [str(tmp_path / "tributo")] assert "pip" not in runtime_env - - -def test_runtime_env_can_explicitly_inject_provider_dependencies(tmp_path) -> None: - (tmp_path / "pyproject.toml").write_text( - "[project]\nname='test'\n", encoding="utf-8" - ) - (tmp_path / "tributo").mkdir() - runtime_env = build_runtime_env( - project_root=tmp_path, - extra_py_modules=[tmp_path / "provider"], - runtime_pip_packages=["tributo-broker-redis==0.1.0"], - ) - assert runtime_env["py_modules"][-1] == str(tmp_path / "provider") - assert runtime_env["pip"] == ["tributo-broker-redis==0.1.0"] diff --git a/tests/training/test_training_lifecycle.py b/tests/training/test_training_lifecycle.py index bfffca9..a8694ba 100644 --- a/tests/training/test_training_lifecycle.py +++ b/tests/training/test_training_lifecycle.py @@ -71,14 +71,6 @@ def export_model(self, checkpoint: Any, output_path: str) -> None: self._summary["metrics"] = {"accuracy": 0.9} -class _MetricsCheckpointTrainer(_FakeTrainer): - def training_loop(self) -> Any: - self.events.append("training_loop") - return SimpleNamespace( - metrics={"eval-logloss_history": [0.8, 0.4], "eval-logloss": 0.4} - ) - - class _EntryTrainer(BaseTrainer): """Production-shaped trainer: accepts ``datasets``/``config`` like ``run_local_trial`` constructs them (``trainer_cls(datasets=..., @@ -361,15 +353,6 @@ def test_export_results_written_to_trainer_summary_are_returned(self) -> None: assert summary["metrics"] == {"accuracy": 0.9} assert trainer._summary is summary - def test_ray_checkpoint_metrics_are_preserved_for_replay(self) -> None: - trainer = _MetricsCheckpointTrainer() - summary = _lifecycle(trainer).run("/tmp/out") - - assert summary["metrics"] == { - "eval-logloss_history": [0.8, 0.4], - "eval-logloss": 0.4, - } - class TestBundleMode: def test_first_party_defaults_to_bundle_before_setup( diff --git a/tests/training/test_xgboost_trainer_unit.py b/tests/training/test_xgboost_trainer_unit.py index 6ba7987..e497c19 100644 --- a/tests/training/test_xgboost_trainer_unit.py +++ b/tests/training/test_xgboost_trainer_unit.py @@ -23,6 +23,7 @@ _managed_resume_checkpoint, _merge_xgb_eval_results, _populate_xgb_eval_metrics, + run_training_result_with_config, ) @@ -69,6 +70,29 @@ def test_managed_resume_checkpoint_cleans_directory_on_failure(tmp_path: Path) - assert not checkpoint_dir.exists() +def test_in_process_training_entrypoint_returns_training_result(monkeypatch) -> None: + monkeypatch.setattr( + "tributo.training.xgboost_trainer.run_training_with_config", + lambda _config: { + "model_uri": "file:///tmp/bundle", + "bundle_uri": "file:///tmp/bundle", + "metrics": {"accuracy": 0.9}, + "legacy_artifact_uri": None, + "training_status": "succeeded", + "bundle_status": "succeeded", + "hook_status": "not_configured", + "execution_id": "execution-1", + "status": "succeeded", + }, + ) + + result = run_training_result_with_config({}) + + assert result.training_status == "succeeded" + assert result.bundle_uri == "file:///tmp/bundle" + assert result.execution_id == "execution-1" + + class TestS3Config: """S3Config Pydantic 模型测试。""" diff --git a/tools/generate_public_api_reference.py b/tools/generate_public_api_reference.py index b7cb193..d0341ca 100644 --- a/tools/generate_public_api_reference.py +++ b/tools/generate_public_api_reference.py @@ -107,6 +107,19 @@ def _decorator_stability(decorator: ast.expr) -> str | None: ) +def _is_exception_class(node: ast.ClassDef) -> bool: + """Classify exception types from their bases, not their domain name.""" + base_names = [ + base.id + if isinstance(base, ast.Name) + else base.attr + if isinstance(base, ast.Attribute) + else "" + for base in node.bases + ] + return any(name.endswith(("Error", "Exception")) for name in base_names) + + def build_inventory(source_root: Path = SOURCE_ROOT) -> tuple[PublicSymbol, ...]: """Return every top-level source object annotated with ``@PublicAPI``.""" symbols: list[PublicSymbol] = [] @@ -135,11 +148,7 @@ def build_inventory(source_root: Path = SOURCE_ROOT) -> tuple[PublicSymbol, ...] f"{path}:{node.lineno}: unsupported stability {stability!r}" ) if isinstance(node, ast.ClassDef): - kind = ( - "exception" - if node.name.endswith(("Error", "Exception")) - else "class" - ) + kind = "exception" if _is_exception_class(node) else "class" else: kind = "function" symbols.append( @@ -164,7 +173,7 @@ def component_for(symbol: PublicSymbol) -> str: """Route a public symbol to one user-facing component page.""" parts = symbol.module.split(".") package = parts[1] if len(parts) > 1 else "core" - if package in {"config", "exceptions", "job", "_common"}: + if package in {"config", "exceptions", "job", "ray_jobs", "_common"}: return "core" if package in {"data", "streaming"}: return "data"