remora/crates/remora-contracts/src/dto.rs

use remora_domain::{JobPhase, OpaqueId, RuntimeFailureClass};
use serde::{Deserialize, Serialize};

use crate::{DateTimeDto, PositiveIntegerDto, SafeIntegerDto, SignedSafeIntegerDto};

#[derive(Clone, Debug, Deserialize, PartialEq, Serialize)]
#[serde(deny_unknown_fields)]
pub struct TargetDto {
    pub store: String,
    pub bucket: String,
    pub prefix: String,
}

#[derive(Clone, Debug, Deserialize, PartialEq, Serialize)]
#[serde(deny_unknown_fields)]
pub struct TargetOverrideDto {
    #[serde(
        default,
        deserialize_with = "optional_non_null",
        skip_serializing_if = "Option::is_none"
    )]
    pub bucket: Option<String>,
    #[serde(
        default,
        deserialize_with = "optional_non_null",
        skip_serializing_if = "Option::is_none"
    )]
    pub prefix: Option<String>,
}

#[derive(Clone, Deserialize, PartialEq, Serialize)]
#[serde(tag = "kind", rename_all = "snake_case")]
pub enum SourceSubmissionDto {
    Http(HttpSourceSubmissionDto),
    ObjectStore(ObjectStoreSourceSubmissionDto),
    HuggingFace(HuggingFaceSourceSubmissionDto),
}

#[derive(Clone, Deserialize, PartialEq, Serialize)]
#[serde(deny_unknown_fields)]
pub struct HttpSourceSubmissionDto {
    pub url: String,
    pub result_path: String,
    #[serde(
        default,
        deserialize_with = "optional_non_null",
        skip_serializing_if = "Option::is_none"
    )]
    pub credential_id: Option<OpaqueId>,
}

#[derive(Clone, Deserialize, PartialEq, Serialize)]
#[serde(deny_unknown_fields)]
pub struct ObjectStoreSourceSubmissionDto {
    pub store: String,
    pub bucket: String,
    pub key: String,
    #[serde(
        default,
        deserialize_with = "optional_non_null",
        skip_serializing_if = "Option::is_none"
    )]
    pub version_id: Option<String>,
    pub result_path: String,
    #[serde(
        default,
        deserialize_with = "optional_non_null",
        skip_serializing_if = "Option::is_none"
    )]
    pub credential_id: Option<OpaqueId>,
}

#[derive(Clone, Debug, Deserialize, PartialEq, Serialize)]
#[serde(deny_unknown_fields)]
pub struct HuggingFaceSelectionDto {
    #[serde(
        default,
        deserialize_with = "optional_non_null",
        skip_serializing_if = "Option::is_none"
    )]
    pub paths: Option<Vec<String>>,
    #[serde(
        default,
        deserialize_with = "optional_non_null",
        skip_serializing_if = "Option::is_none"
    )]
    pub prefixes: Option<Vec<String>>,
}

#[derive(Clone, Deserialize, PartialEq, Serialize)]
#[serde(deny_unknown_fields)]
pub struct HuggingFaceSourceSubmissionDto {
    pub repository_type: HuggingFaceRepositoryTypeDto,
    pub repository_id: String,
    pub revision: String,
    #[serde(
        default,
        deserialize_with = "optional_non_null",
        skip_serializing_if = "Option::is_none"
    )]
    pub selection: Option<HuggingFaceSelectionDto>,
    #[serde(
        default,
        deserialize_with = "optional_non_null",
        skip_serializing_if = "Option::is_none"
    )]
    pub credential_id: Option<OpaqueId>,
}

#[derive(Clone, Copy, Debug, Deserialize, Eq, PartialEq, Serialize)]
#[serde(rename_all = "snake_case")]
pub enum HuggingFaceRepositoryTypeDto {
    Model,
    Dataset,
    Space,
}

#[derive(Clone, Deserialize, PartialEq, Serialize)]
#[serde(deny_unknown_fields)]
pub struct JobSubmissionDto {
    pub source: SourceSubmissionDto,
    #[serde(
        default,
        deserialize_with = "optional_non_null",
        skip_serializing_if = "Option::is_none"
    )]
    pub target: Option<TargetOverrideDto>,
    #[serde(
        default,
        deserialize_with = "optional_non_null",
        skip_serializing_if = "Option::is_none"
    )]
    pub label: Option<String>,
}

#[derive(Clone, Debug, Deserialize, Eq, PartialEq, Serialize)]
#[serde(deny_unknown_fields)]
pub struct AttemptAuthorityDto {
    pub attempt_id: OpaqueId,
    pub fence: PositiveIntegerDto,
    pub worker_incarnation: OpaqueId,
}

#[derive(Clone, Debug, Deserialize, Eq, PartialEq, Serialize)]
#[serde(deny_unknown_fields)]
pub struct ProgressDto {
    pub attempt_id: OpaqueId,
    pub fence: PositiveIntegerDto,
    pub completed_bytes: SafeIntegerDto,
    pub total_bytes: Option<SafeIntegerDto>,
    pub completed_members: SafeIntegerDto,
    pub total_members: Option<SafeIntegerDto>,
    pub observed_at: DateTimeDto,
}

#[derive(Clone, Debug, Deserialize, Eq, PartialEq, Serialize)]
#[serde(untagged)]
pub enum JobFailureDto {
    Permanent(PermanentJobFailureDto),
    Retryable(RetryableJobFailureDto),
}

#[derive(Clone, Debug, Deserialize, Eq, PartialEq, Serialize)]
#[serde(deny_unknown_fields)]
pub struct PermanentJobFailureDto {
    pub class: PermanentFailureClassDto,
    pub message: String,
    pub retryable: PermanentDispositionDto,
    #[serde(
        default,
        deserialize_with = "optional_non_null",
        skip_serializing_if = "Option::is_none"
    )]
    pub safe_details: Option<std::collections::BTreeMap<String, SafeDetailValueDto>>,
}

#[derive(Clone, Debug, Deserialize, Eq, PartialEq, Serialize)]
#[serde(deny_unknown_fields)]
pub struct RetryableJobFailureDto {
    pub class: RetryableFailureClassDto,
    pub message: String,
    pub retryable: RetryableDispositionDto,
    #[serde(
        default,
        deserialize_with = "optional_non_null",
        skip_serializing_if = "Option::is_none"
    )]
    pub safe_details: Option<std::collections::BTreeMap<String, SafeDetailValueDto>>,
}

#[derive(Clone, Copy, Debug, Deserialize, Eq, PartialEq, Serialize)]
#[serde(rename_all = "snake_case")]
pub enum PermanentFailureClassDto {
    Validation,
    Authorization,
    SourceNotFound,
    SourceChanged,
    Cancelled,
}

#[derive(Clone, Copy, Debug, Deserialize, Eq, PartialEq, Serialize)]
#[serde(rename_all = "snake_case")]
pub enum RetryableFailureClassDto {
    TransientSource,
    TransientTarget,
    Integrity,
    Capacity,
    Internal,
    LostAttempt,
}

#[derive(Clone, Copy, Debug, Eq, PartialEq)]
pub struct PermanentDispositionDto;

impl Serialize for PermanentDispositionDto {
    fn serialize<S>(&self, serializer: S) -> Result<S::Ok, S::Error>
    where
        S: serde::Serializer,
    {
        serializer.serialize_bool(false)
    }
}

impl<'de> Deserialize<'de> for PermanentDispositionDto {
    fn deserialize<D>(deserializer: D) -> Result<Self, D::Error>
    where
        D: serde::Deserializer<'de>,
    {
        if bool::deserialize(deserializer)? {
            return Err(serde::de::Error::custom(
                "permanent failure must not be retryable",
            ));
        }
        Ok(Self)
    }
}

#[derive(Clone, Copy, Debug, Eq, PartialEq)]
pub struct RetryableDispositionDto;

impl Serialize for RetryableDispositionDto {
    fn serialize<S>(&self, serializer: S) -> Result<S::Ok, S::Error>
    where
        S: serde::Serializer,
    {
        serializer.serialize_bool(true)
    }
}

impl<'de> Deserialize<'de> for RetryableDispositionDto {
    fn deserialize<D>(deserializer: D) -> Result<Self, D::Error>
    where
        D: serde::Deserializer<'de>,
    {
        if !bool::deserialize(deserializer)? {
            return Err(serde::de::Error::custom(
                "retryable failure must be retryable",
            ));
        }
        Ok(Self)
    }
}

#[derive(Clone, Debug, Deserialize, Eq, PartialEq, Serialize)]
#[serde(deny_unknown_fields)]
pub struct AttemptOutcomeRequestDto {
    pub operation_id: OpaqueId,
    pub fence: PositiveIntegerDto,
    pub worker_incarnation: OpaqueId,
    pub outcome: AttemptOutcomeDto,
}

#[derive(Clone, Debug, Deserialize, Eq, PartialEq, Serialize)]
#[serde(tag = "kind", rename_all = "snake_case")]
pub enum AttemptOutcomeDto {
    Success(AttemptSuccessOutcomeDto),
    Failure(AttemptFailureOutcomeDto),
}

#[derive(Clone, Debug, Deserialize, Eq, PartialEq, Serialize)]
#[serde(deny_unknown_fields)]
pub struct AttemptSuccessOutcomeDto {
    pub result: ResultDescriptorDto,
}

#[derive(Clone, Debug, Deserialize, Eq, PartialEq, Serialize)]
#[serde(deny_unknown_fields)]
pub struct AttemptFailureOutcomeDto {
    pub failure: AttemptRuntimeFailureDto,
}

#[derive(Clone, Debug, Deserialize, Eq, PartialEq, Serialize)]
#[serde(deny_unknown_fields)]
pub struct AttemptRuntimeFailureDto {
    pub class: RuntimeFailureClass,
    pub message: String,
    #[serde(
        default,
        deserialize_with = "optional_non_null",
        skip_serializing_if = "Option::is_none"
    )]
    pub safe_details: Option<std::collections::BTreeMap<String, SafeDetailValueDto>>,
}

#[derive(Clone, Debug, Deserialize, Eq, PartialEq, Serialize)]
#[serde(untagged)]
pub enum SafeDetailValueDto {
    String(String),
    Integer(SignedSafeIntegerDto),
    Boolean(bool),
    Null,
}

#[derive(Clone, Debug, Deserialize, Eq, PartialEq, Serialize)]
#[serde(deny_unknown_fields)]
pub struct ResultDescriptorDto {
    pub kind: ResultKindDto,
    pub manifest_id: OpaqueId,
    pub schema_version: PositiveIntegerDto,
    pub manifest: ObjectIdentityDto,
    pub member_count: PositiveIntegerDto,
    pub total_bytes: SafeIntegerDto,
}

#[derive(Clone, Copy, Debug, Deserialize, Eq, PartialEq, Serialize)]
#[serde(rename_all = "snake_case")]
pub enum ResultKindDto {
    Manifest,
}

#[derive(Clone, Debug, Deserialize, Eq, PartialEq, Serialize)]
#[serde(deny_unknown_fields)]
pub struct ObjectIdentityDto {
    pub store: String,
    pub bucket: String,
    pub key: String,
    #[serde(
        default,
        deserialize_with = "optional_non_null",
        skip_serializing_if = "Option::is_none"
    )]
    pub version_id: Option<String>,
    pub checksum: String,
    pub size: SafeIntegerDto,
}

#[derive(Clone, Debug, Deserialize, Eq, PartialEq, Serialize)]
#[serde(deny_unknown_fields)]
pub struct JobPhaseProjectionDto {
    pub phase: JobPhase,
}

pub(crate) fn optional_non_null<'de, D, T>(deserializer: D) -> Result<Option<T>, D::Error>
where
    D: serde::Deserializer<'de>,
    T: Deserialize<'de>,
{
    T::deserialize(deserializer).map(Some)
}