code/self_healing_finetuning/contracts/provider_adapter.py

from __future__ import annotations

from typing import Any, Literal, Protocol

from n21.types import DeploymentHandle, TrainHandle, ValidationResult


class ProviderAdapter(Protocol):
    def validate_config(self, config: dict[str, Any], *, live: bool = False) -> ValidationResult: ...
    def stage_dataset(self, dataset_manifest: dict[str, Any]) -> dict[str, Any]: ...
    def start_train(self, run_manifest: dict[str, Any]) -> TrainHandle: ...
    def resume_train(self, run_id: str, checkpoint_ptr: dict[str, Any]) -> TrainHandle: ...
    def cancel_train(self, run_id: str) -> dict[str, Any]: ...
    def get_run_status(self, run_id: str) -> dict[str, Any]: ...
    def stream_logs(self, run_id: str, cursor: str | None = None) -> dict[str, Any]: ...
    def export_artifacts(self, run_id: str, export_policy: dict[str, Any]) -> dict[str, Any]: ...
    def deploy_endpoint(self, model_manifest: dict[str, Any], env: Literal["dev", "stage", "prod"]) -> DeploymentHandle: ...
    def invoke(self, deployment_id: str, request: dict[str, Any]) -> dict[str, Any]: ...
    def healthcheck(self, deployment_id: str) -> dict[str, Any]: ...
    def promote(self, deployment_id: str, env: Literal["stage", "prod"]) -> dict[str, Any]: ...
    def rollback(self, env: Literal["stage", "prod"], target: dict[str, Any]) -> dict[str, Any]: ...
    def delete_endpoint(self, deployment_id: str) -> dict[str, Any]: ...
    def estimate_cost(self, plan: dict[str, Any]) -> dict[str, Any]: ...