ms2k/remora
SHA256
Watch
0
Code Releases Activity

Separate orchestration by protocol action

Browse source
This commit is contained in:
ms2k 2026-07-19 14:27:59 +09:00
commit 6f0f20f5cc
9 changed files with 6974 additions and 6839 deletions
Download patch file Download diff file Expand all files Collapse all files

1655
crates/remora-worker/src/executor.rs
View file

File diff suppressed because it is too large Load diff

223
crates/remora-worker/src/executor/envelope.rs Normal file
View file

@ -0,0 +1,223 @@
use super::*;
impl AdapterExecutor {
pub(super) async fn pin_resolved<T: Serialize>(
&self,
context: &AttemptContext,
envelope: &ExecutionEnvelopeDto,
payload: &T,
item_count: Option<u64>,
) -> Result<OpaqueId, AdapterFailure> {
context
.require_live_authority()
.await
.map_err(authority_failure)?;
let value = serde_json::to_value(payload)
.map_err(|_| failed(RuntimeFailureClass::Internal, "source identity is invalid"))?;
if let Some(existing) = &envelope.resolved_source_ref.0 {
let stored = self
.descriptors
.load_ref(&envelope.job_id, existing)
.await
.map_err(authority_failure)?
.ok_or_else(|| {
failed(
RuntimeFailureClass::SourceChanged,
"source identity changed",
)
})?;
if stored != value {
return Err(failed(
RuntimeFailureClass::SourceChanged,
"source identity changed",
));
}
return Ok(existing.id.clone());
}
let descriptor = self
.descriptors
.store_resolved(&envelope.attempt_id, &value, item_count)
.await
.map_err(authority_failure)?;
let id = descriptor.id.clone();
context
.pin_resolved_source(descriptor)
.await
.map_err(authority_failure)?;
Ok(id)
}
pub(super) async fn load_source(
&self,
envelope: &ExecutionEnvelopeDto,
) -> Result<SourceSubmissionDto, AdapterFailure> {
type Row = (
Value,
Value,
String,
String,
String,
Json<Vec<String>>,
String,
Option<String>,
Option<String>,
);
let row = sqlx::query_as::<_, Row>(
"SELECT s.source_input, j.source_projection, j.sensitive_source_ref,
j.target_store, j.target_bucket, j.target_prefix,
j.execution_grant_id, g.credential_id, g.credential_version_id
FROM jobs j
JOIN sensitive_source_inputs s
ON s.sensitive_source_ref = j.sensitive_source_ref
JOIN execution_grants g ON g.grant_id = j.execution_grant_id
WHERE j.job_id = $1",
)
.bind(envelope.job_id.as_str())
.fetch_optional(&self.pool)
.await
.map_err(|_| failed(RuntimeFailureClass::Internal, "source input is unavailable"))?
.ok_or_else(|| {
failed(
RuntimeFailureClass::Authorization,
"source input is unavailable",
)
})?;
let source: SourceSubmissionDto = serde_json::from_value(row.0)
.map_err(|_| failed(RuntimeFailureClass::Internal, "source input is invalid"))?;
let envelope_projection = serde_json::to_value(&envelope.source)
.map_err(|_| failed(RuntimeFailureClass::Internal, "source input is invalid"))?;
let grant_credential = match (row.7, row.8) {
(Some(credential_id), Some(version_id)) => Some((credential_id, version_id)),
(None, None) => None,
_ => {
return Err(failed(
RuntimeFailureClass::Internal,
"source grant is invalid",
));
}
};
let envelope_credential = envelope.credential_version.0.as_ref().map(|credential| {
(
credential.credential_id.as_str(),
credential.version_id.as_str(),
)
});
let source_credential = source_credential_id(&source).map(OpaqueId::as_str);
if !projection_matches(&source, &envelope.source)
|| row.1 != envelope_projection
|| row.2 != envelope.sensitive_source_ref.as_str()
|| row.3 != envelope.target.store
|| row.4 != envelope.target.bucket
|| row.5.0.join("/") != envelope.target.prefix
|| row.6 != envelope.execution_grant_ref.as_str()
|| grant_credential
.as_ref()
.map(|(id, version)| (id.as_str(), version.as_str()))
!= envelope_credential
|| grant_credential.as_ref().map(|(id, _)| id.as_str()) != source_credential
{
return Err(failed(
RuntimeFailureClass::Authorization,
"source input does not match its grant",
));
}
Ok(source)
}
pub(super) async fn load_checkpoint(
&self,
envelope: &ExecutionEnvelopeDto,
) -> Result<ResumeState, AdapterFailure> {
let staging_origin = staging_origin(envelope)?;
let Some(checkpoint) = &envelope.checkpoint_ref.0 else {
return Ok(ResumeState::empty());
};
let Some(resolved) = &envelope.resolved_source_ref.0 else {
return Err(failed(
RuntimeFailureClass::SourceChanged,
"checkpoint source identity is unavailable",
));
};
self.descriptors
.load_ref(&envelope.job_id, resolved)
.await
.map_err(authority_failure)?
.ok_or_else(|| {
failed(
RuntimeFailureClass::SourceChanged,
"checkpoint source identity is unavailable",
)
})?;
let payload = self
.descriptors
.load_ref(&envelope.job_id, checkpoint)
.await
.map_err(authority_failure)?
.ok_or_else(|| {
failed(
RuntimeFailureClass::SourceChanged,
"checkpoint resume evidence is unavailable",
)
})?;
let checkpoint: WorkerCheckpoint = serde_json::from_value(payload).map_err(|_| {
failed(
RuntimeFailureClass::SourceChanged,
"checkpoint resume evidence is unavailable",
)
})?;
if checkpoint.staging_origin_attempt_id != staging_origin
|| checkpoint.resolved_source_id != resolved.id
|| (checkpoint.validated_members == 0) != checkpoint.partial.is_some()
|| checkpoint.partial.as_ref().is_some_and(|partial| {
partial
.destination
.parts
.iter()
.try_fold(0_u64, |total, part| total.checked_add(part.size))
!= Some(checkpoint.validated_bytes)
})
{
return Err(failed(
RuntimeFailureClass::SourceChanged,
"checkpoint resume evidence is unavailable",
));
}
Ok(ResumeState {
resolved_source_id: Some(checkpoint.resolved_source_id),
validated_bytes: checkpoint.validated_bytes,
validated_members: checkpoint.validated_members,
partial: checkpoint.partial,
})
}
pub(super) async fn resolve_credential(
&self,
envelope: &ExecutionEnvelopeDto,
) -> Result<Option<SecretMaterial>, AdapterFailure> {
if envelope.credential_version.0.is_none() {
return Ok(None);
}
self.credentials
.resolve_access(&ActiveCredentialGrant::for_execution_grant(
envelope.execution_grant_ref.clone(),
))
.await
.map(Some)
.map_err(credential_failure)
}
pub(super) fn target_for(
&self,
envelope: &ExecutionEnvelopeDto,
) -> Result<TargetStore, AdapterFailure> {
if self.target.target_store_id != envelope.target.store {
return Err(failed(
RuntimeFailureClass::Authorization,
"target store is not configured",
));
}
let mut config = self.target.clone();
config.bucket = envelope.target.bucket.clone();
TargetStore::connect(config).map_err(target_failure)
}
}

137
crates/remora-worker/src/executor/failures.rs Normal file
View file

@ -0,0 +1,137 @@
use super::*;
pub(super) fn failed(class: RuntimeFailureClass, message: &'static str) -> AdapterFailure {
AdapterFailure::Failed(ExecutionFailure::new(class, message))
}
pub(super) fn authority_failure(error: WorkerError) -> AdapterFailure {
match error {
WorkerError::PreconditionFailed => AdapterFailure::Cancelled,
_ => failed(
RuntimeFailureClass::Internal,
"attempt service is unavailable",
),
}
}
pub(super) fn credential_failure(error: CredentialError) -> AdapterFailure {
match error {
CredentialError::SecretStoreUnavailable | CredentialError::Durable(_) => failed(
RuntimeFailureClass::TransientSource,
"source credential is temporarily unavailable",
),
_ => failed(
RuntimeFailureClass::Authorization,
"source credential is unavailable",
),
}
}
pub(super) fn http_failure(error: HttpSourceError) -> AdapterFailure {
let class = match error {
HttpSourceError::Authorization => RuntimeFailureClass::Authorization,
HttpSourceError::NotFound => RuntimeFailureClass::SourceNotFound,
HttpSourceError::SourceChanged => RuntimeFailureClass::SourceChanged,
HttpSourceError::Transient { .. } => RuntimeFailureClass::TransientSource,
HttpSourceError::InvalidRequest | HttpSourceError::InvalidResponse => {
RuntimeFailureClass::Internal
}
};
failed(class, public_message(class))
}
pub(super) fn object_failure(error: ObjectSourceError) -> AdapterFailure {
let class = match error {
ObjectSourceError::SourceAuthorization | ObjectSourceError::DestinationAuthorization => {
RuntimeFailureClass::Authorization
}
ObjectSourceError::SourceNotFound => RuntimeFailureClass::SourceNotFound,
ObjectSourceError::SourceChanged | ObjectSourceError::DestinationChanged => {
RuntimeFailureClass::SourceChanged
}
ObjectSourceError::TransientSource => RuntimeFailureClass::TransientSource,
ObjectSourceError::TransientDestination => RuntimeFailureClass::TransientTarget,
ObjectSourceError::InvalidRequest | ObjectSourceError::Internal => {
RuntimeFailureClass::Internal
}
};
failed(class, public_message(class))
}
pub(super) fn hf_failure(error: HfSourceError) -> AdapterFailure {
let class = match error {
HfSourceError::Authorization => RuntimeFailureClass::Authorization,
HfSourceError::NotFound => RuntimeFailureClass::SourceNotFound,
HfSourceError::SourceChanged => RuntimeFailureClass::SourceChanged,
HfSourceError::Capacity => RuntimeFailureClass::Capacity,
HfSourceError::Transient { .. } => RuntimeFailureClass::TransientSource,
HfSourceError::InvalidRequest
| HfSourceError::UnsupportedEntry
| HfSourceError::InvalidResponse
| HfSourceError::InvalidInventory => RuntimeFailureClass::Internal,
};
failed(class, public_message(class))
}
pub(super) fn target_failure(error: TargetError) -> AdapterFailure {
let class = match error {
TargetError::Authorization => RuntimeFailureClass::Authorization,
TargetError::Capacity => RuntimeFailureClass::Capacity,
TargetError::Transient(_) => RuntimeFailureClass::TransientTarget,
TargetError::Integrity(_) | TargetError::AlreadyExists => RuntimeFailureClass::Integrity,
TargetError::NotFound | TargetError::Validation(_) | TargetError::Internal(_) => {
RuntimeFailureClass::Internal
}
};
failed(class, public_message(class))
}
pub(super) fn stream_failure(error: StreamError) -> AdapterFailure {
let class = error
.runtime_class()
.unwrap_or(RuntimeFailureClass::Internal);
failed(class, public_message(class))
}
pub(super) fn result_failure(error: ResultError) -> AdapterFailure {
let class = match error {
ResultError::Integrity => RuntimeFailureClass::Integrity,
ResultError::Target(error) => return target_failure(error),
_ => RuntimeFailureClass::Internal,
};
failed(class, public_message(class))
}
pub(super) fn publication_dto(
input: PublicationInput,
) -> Result<ResultDescriptorDto, AdapterFailure> {
Ok(ResultDescriptorDto {
kind: ResultKindDto::Manifest,
manifest_id: input.manifest_id,
schema_version: positive(input.schema_version.into())?,
manifest: ObjectIdentityDto {
store: input.manifest.identity.target_store_id,
bucket: input.manifest.identity.bucket,
key: input.manifest.identity.key,
version_id: input.manifest.identity.version_id,
checksum: input.manifest.content_sha256.as_str().to_owned(),
size: safe(input.manifest.size)?,
},
member_count: positive(input.member_count)?,
total_bytes: safe(input.total_bytes)?,
})
}
pub(super) fn positive(value: u64) -> Result<PositiveIntegerDto, AdapterFailure> {
PositiveIntegerDto::new(value)
.map_err(|_| failed(RuntimeFailureClass::Internal, "result exceeds wire limits"))
}
pub(super) fn safe(value: u64) -> Result<SafeIntegerDto, AdapterFailure> {
SafeIntegerDto::new(value).map_err(|_| {
failed(
RuntimeFailureClass::Internal,
"progress exceeds wire limits",
)
})
}

129
crates/remora-worker/src/executor/http_source.rs Normal file
View file

@ -0,0 +1,129 @@
use super::*;
impl AdapterExecutor {
pub(super) async fn execute_http(
&self,
context: &AttemptContext,
envelope: &ExecutionEnvelopeDto,
source: HttpSourceSubmissionDto,
secret: Option<SecretMaterial>,
resources: &ExecutionResources<'_>,
) -> Result<(), AdapterFailure> {
let ExecutionResources {
target,
inventory,
resume,
} = resources;
if resume.validated_members > 1 {
return Err(failed(
RuntimeFailureClass::SourceChanged,
"checkpoint resume evidence is unavailable",
));
}
let headers = http_headers(secret.as_ref())?;
let request = HttpSourceRequest::new(&source.url, &source.result_path)
.and_then(|request| request.with_credential_headers(headers.clone()))
.map_err(http_failure)?;
context
.require_live_authority()
.await
.map_err(authority_failure)?;
let adapter = HttpSource::new_with_root_certificates(
SystemResolver,
self.outbound_policy.clone(),
self.retry_timing,
self.outbound_roots.clone(),
)
.map_err(http_failure)?;
if let Some(partial) = &resume.partial {
let checkpoint = http_resume_checkpoint(partial, &source.url, &source.result_path)?;
let resolved_source_id = resume.resolved_source_id.as_ref().ok_or_else(|| {
failed(
RuntimeFailureClass::SourceChanged,
"checkpoint source identity is unavailable",
)
})?;
context
.require_live_authority()
.await
.map_err(authority_failure)?;
let (session, parts, member_id) =
self.resume_destination(envelope, target, partial).await?;
let total = checkpoint.identity().representation_length();
let prefix = checkpoint.prefix_integrity().clone();
let opened = adapter
.resume(&checkpoint, &prefix, headers)
.await
.map_err(http_failure)?;
self.copy_resumed_member(
context,
envelope,
target,
inventory,
source.result_path,
opened.into_body(),
total,
resolved_source_id,
session,
parts,
member_id,
)
.await?;
return Ok(());
}
let opened = adapter.open(request).await.map_err(http_failure)?;
let total = opened
.identity()
.map(HttpSourceIdentity::representation_length);
let resolved_source_id = match opened.identity() {
Some(identity) => Some(
self.pin_resolved(context, envelope, &http_identity(identity), Some(1))
.await?,
),
None if resume.validated_members == 0 => None,
None => {
return Err(failed(
RuntimeFailureClass::SourceChanged,
"source identity is unavailable",
));
}
};
let path = opened.result_path().canonical().to_owned();
let partial_source = opened.identity().map(http_partial_source);
if let Some(resolved_source_id) = resolved_source_id.as_ref()
&& self
.resume_member(
context,
envelope,
inventory,
target,
resume,
resolved_source_id,
&path,
total,
1,
total.unwrap_or(0),
)
.await?
{
return Ok(());
}
self.copy_member(
context,
envelope,
target,
inventory,
path,
opened.into_body(),
total,
total,
0,
0,
Some(1),
resolved_source_id.as_ref(),
partial_source,
)
.await?;
Ok(())
}
}

123
crates/remora-worker/src/executor/hugging_face.rs Normal file
View file

@ -0,0 +1,123 @@
use super::*;
impl AdapterExecutor {
pub(super) async fn execute_hugging_face(
&self,
context: &AttemptContext,
envelope: &ExecutionEnvelopeDto,
source: HuggingFaceSourceSubmissionDto,
secret: Option<SecretMaterial>,
resources: &ExecutionResources<'_>,
) -> Result<(), AdapterFailure> {
let ExecutionResources {
target,
inventory,
resume,
} = resources;
let headers = hugging_face_headers(secret.as_ref())?;
let selection = selection(source.selection.as_ref())?;
let request = HfSourceRequest::new(
&self.hugging_face_endpoint,
repository_type(source.repository_type),
&source.repository_id,
&source.revision,
selection,
)
.and_then(|request| request.with_credential_headers(headers.clone()))
.map_err(hf_failure)?;
let adapter = HfSource::new_with_root_certificates(
SystemResolver,
self.outbound_policy.clone(),
self.retry_timing,
self.outbound_roots.clone(),
)
.map_err(hf_failure)?;
context
.require_live_authority()
.await
.map_err(authority_failure)?;
let snapshot = adapter.resolve(request).await.map_err(hf_failure)?;
let resolved_source_id = self
.pin_resolved(
context,
envelope,
&snapshot,
Some(snapshot.inventory().member_count),
)
.await?;
let total_members = snapshot.inventory().member_count;
if resume.validated_members > total_members {
return Err(failed(
RuntimeFailureClass::SourceChanged,
"checkpoint resume evidence is unavailable",
));
}
let mut members = adapter
.inventory(&snapshot, headers.clone())
.map_err(hf_failure)?;
let mut completed_bytes = 0_u64;
let mut completed_members = 0_u64;
loop {
context
.require_live_authority()
.await
.map_err(authority_failure)?;
let Some(member) = members.next().await.map_err(hf_failure)? else {
break;
};
context
.require_live_authority()
.await
.map_err(authority_failure)?;
let body = adapter
.open_member(&snapshot, &member, headers.clone())
.await
.map_err(hf_failure)?;
let size = member.size();
if completed_members < resume.validated_members {
let resumed = self
.resume_member(
context,
envelope,
inventory,
target,
resume,
&resolved_source_id,
member.path(),
Some(size),
completed_members.saturating_add(1),
completed_bytes.saturating_add(size),
)
.await?;
if !resumed {
return Err(failed(
RuntimeFailureClass::SourceChanged,
"checkpoint resume evidence is unavailable",
));
}
completed_bytes = completed_bytes.saturating_add(size);
completed_members = completed_members.saturating_add(1);
continue;
}
self.copy_member(
context,
envelope,
target,
inventory,
member.path().to_owned(),
body,
Some(size),
None,
completed_bytes,
completed_members,
Some(total_members),
Some(&resolved_source_id),
None,
)
.await?;
completed_bytes = completed_bytes.saturating_add(size);
completed_members = completed_members.saturating_add(1);
}
Ok(())
}
}

142
crates/remora-worker/src/executor/object.rs Normal file
View file

@ -0,0 +1,142 @@
use super::*;
impl AdapterExecutor {
pub(super) async fn execute_object(
&self,
context: &AttemptContext,
envelope: &ExecutionEnvelopeDto,
source: ObjectStoreSourceSubmissionDto,
secret: Option<SecretMaterial>,
resources: &ExecutionResources<'_>,
) -> Result<(), AdapterFailure> {
let ExecutionResources {
target,
inventory,
resume,
} = resources;
if resume.validated_members > 1 {
return Err(failed(
RuntimeFailureClass::SourceChanged,
"checkpoint resume evidence is unavailable",
));
}
let endpoint = self.object_sources.get(&source.store).ok_or_else(|| {
failed(
RuntimeFailureClass::Authorization,
"object source is not configured",
)
})?;
let (access_key_id, secret_access_key, session_token) = object_credentials(secret)?;
let config = ObjectSourceConfig::new_with_session_token(
&endpoint.endpoint,
source.bucket.clone(),
endpoint.region.clone(),
access_key_id,
secret_access_key,
session_token,
)
.map_err(object_failure)?;
let mut request = ObjectSourceRequest::new(source.key.clone(), &source.result_path)
.map_err(object_failure)?;
if let Some(version_id) = source.version_id.clone() {
request = request
.with_version_id(version_id)
.map_err(object_failure)?;
}
context
.require_live_authority()
.await
.map_err(authority_failure)?;
let adapter = ObjectSource::connect(config).map_err(object_failure)?;
if let Some(partial) = &resume.partial {
let checkpoint = object_resume_checkpoint(
partial,
&endpoint.endpoint,
&source.bucket,
&source.key,
source.version_id.as_deref(),
&source.result_path,
)?;
let resolved_source_id = resume.resolved_source_id.as_ref().ok_or_else(|| {
failed(
RuntimeFailureClass::SourceChanged,
"checkpoint source identity is unavailable",
)
})?;
context
.require_live_authority()
.await
.map_err(authority_failure)?;
let prefix = checkpoint.prefix_integrity().clone();
let total = checkpoint.source().size();
let parts = checkpoint_parts(&partial.destination)?;
let member_id = checkpoint_member_id(&partial.destination.key)?;
let resumed = adapter
.resume(&checkpoint, &prefix, target)
.await
.map_err(object_failure)?;
let (opened, session) = resumed.into_parts();
self.copy_resumed_member(
context,
envelope,
target,
inventory,
source.result_path,
opened.into_body(),
total,
resolved_source_id,
session,
parts,
member_id,
)
.await?;
return Ok(());
}
let opened = adapter.open(request).await.map_err(object_failure)?;
let resolved_source_id = self
.pin_resolved(
context,
envelope,
&object_identity(opened.identity()),
Some(1),
)
.await?;
let total = opened.identity().size();
let path = opened.result_path().canonical().to_owned();
let partial_source = object_partial_source(opened.identity());
if self
.resume_member(
context,
envelope,
inventory,
target,
resume,
&resolved_source_id,
&path,
Some(total),
1,
total,
)
.await?
{
return Ok(());
}
self.copy_member(
context,
envelope,
target,
inventory,
path,
opened.into_body(),
Some(total),
Some(total),
0,
0,
Some(1),
Some(&resolved_source_id),
Some(partial_source),
)
.await?;
Ok(())
}
}

366
crates/remora-worker/src/executor/source_helpers.rs Normal file
View file

@ -0,0 +1,366 @@
use super::*;
pub(super) fn source_credential_id(source: &SourceSubmissionDto) -> Option<&OpaqueId> {
match source {
SourceSubmissionDto::Http(source) => source.credential_id.as_ref(),
SourceSubmissionDto::ObjectStore(source) => source.credential_id.as_ref(),
SourceSubmissionDto::HuggingFace(source) => source.credential_id.as_ref(),
}
}
pub(super) fn http_headers(secret: Option<&SecretMaterial>) -> Result<HeaderMap, AdapterFailure> {
let mut headers = HeaderMap::new();
match secret {
None => {}
Some(SecretMaterial::BasicAuth { username, password }) => {
let encoded = STANDARD.encode(format!("{username}:{password}"));
headers.insert(
header::AUTHORIZATION,
HeaderValue::from_str(&format!("Basic {encoded}")).map_err(|_| {
failed(
RuntimeFailureClass::Authorization,
"source credential is invalid",
)
})?,
);
}
Some(SecretMaterial::BearerToken { token }) => {
headers.insert(
header::AUTHORIZATION,
HeaderValue::from_str(&format!("Bearer {token}")).map_err(|_| {
failed(
RuntimeFailureClass::Authorization,
"source credential is invalid",
)
})?,
);
}
Some(SecretMaterial::HttpHeaders { headers: values }) => {
for (name, value) in values {
headers.insert(
HeaderName::from_bytes(name.as_bytes()).map_err(|_| {
failed(
RuntimeFailureClass::Authorization,
"source credential is invalid",
)
})?,
HeaderValue::from_str(value).map_err(|_| {
failed(
RuntimeFailureClass::Authorization,
"source credential is invalid",
)
})?,
);
}
}
Some(_) => {
return Err(failed(
RuntimeFailureClass::Authorization,
"source credential kind does not match",
));
}
}
Ok(headers)
}
pub(super) fn hugging_face_headers(
secret: Option<&SecretMaterial>,
) -> Result<HeaderMap, AdapterFailure> {
match secret {
None => Ok(HeaderMap::new()),
Some(SecretMaterial::HuggingFace { token }) => {
let mut headers = HeaderMap::new();
headers.insert(
header::AUTHORIZATION,
HeaderValue::from_str(&format!("Bearer {token}")).map_err(|_| {
failed(
RuntimeFailureClass::Authorization,
"source credential is invalid",
)
})?,
);
Ok(headers)
}
Some(_) => Err(failed(
RuntimeFailureClass::Authorization,
"source credential kind does not match",
)),
}
}
pub(super) fn object_credentials(
secret: Option<SecretMaterial>,
) -> Result<(String, String, Option<String>), AdapterFailure> {
match secret {
Some(SecretMaterial::ObjectStore {
access_key_id,
secret_access_key,
session_token,
}) => Ok((access_key_id, secret_access_key, session_token)),
_ => Err(failed(
RuntimeFailureClass::Authorization,
"object source credential is unavailable",
)),
}
}
pub(super) fn selection(
value: Option<&remora_contracts::HuggingFaceSelectionDto>,
) -> Result<HfSelection, AdapterFailure> {
HfSelection::new(
value
.and_then(|value| value.paths.as_ref())
.into_iter()
.flatten(),
value
.and_then(|value| value.prefixes.as_ref())
.into_iter()
.flatten(),
)
.map_err(hf_failure)
}
pub(super) fn projection_matches(
source: &SourceSubmissionDto,
projection: &SourceProjectionDto,
) -> bool {
match (source, projection) {
(SourceSubmissionDto::Http(source), SourceProjectionDto::Http(projection)) => {
let Ok(mut display) = reqwest::Url::parse(&source.url) else {
return false;
};
display.set_query(None);
display.set_fragment(None);
projection.display_url == display.as_str()
&& projection.result_path == source.result_path
&& projection.credential_id == source.credential_id
}
(
SourceSubmissionDto::ObjectStore(source),
SourceProjectionDto::ObjectStore(projection),
) => {
projection.store == source.store
&& projection.bucket == source.bucket
&& projection.key == source.key
&& projection.version_id == source.version_id
&& projection.result_path == source.result_path
&& projection.credential_id == source.credential_id
}
(
SourceSubmissionDto::HuggingFace(source),
SourceProjectionDto::HuggingFace(projection),
) => {
projection.repository_type == source.repository_type
&& projection.repository_id == source.repository_id
&& projection.revision == source.revision
&& projection.selection == source.selection
&& projection.credential_id == source.credential_id
}
_ => false,
}
}
pub(super) fn http_identity(identity: &HttpSourceIdentity) -> Value {
json!({
"kind": "http",
"requested_url_sha256": hash_text(identity.requested_url().as_url().as_str()),
"final_url_sha256": hash_text(identity.final_url().as_url().as_str()),
"etag_sha256": hash_text(identity.etag().as_str()),
"representation_length": identity.representation_length(),
"content_encoding": "identity"
})
}
pub(super) fn http_partial_source(identity: &HttpSourceIdentity) -> PartialSource {
PartialSource::Http {
requested_url: identity.requested_url().as_url().as_str().to_owned(),
final_url: identity.final_url().as_url().as_str().to_owned(),
etag: identity.etag().as_str().to_owned(),
representation_length: identity.representation_length(),
}
}
pub(super) fn object_identity(identity: &ObjectSourceIdentity) -> Value {
json!({
"kind": "object_store",
"endpoint": identity.endpoint().as_url().as_str(),
"bucket": identity.bucket(),
"key": identity.key(),
"size": identity.size(),
"version_id": identity.version_id(),
"validator": identity.validator().map(|value| value.as_str())
})
}
pub(super) fn object_partial_source(identity: &ObjectSourceIdentity) -> PartialSource {
PartialSource::Object {
endpoint: identity.endpoint().as_url().as_str().to_owned(),
bucket: identity.bucket().to_owned(),
key: identity.key().to_owned(),
size: identity.size(),
version_id: identity.version_id().map(str::to_owned),
validator: identity.validator().map(|value| value.as_str().to_owned()),
}
}
pub(super) fn http_resume_checkpoint(
partial: &PartialCheckpoint,
requested_url: &str,
result_path: &str,
) -> Result<HttpResumeCheckpoint, AdapterFailure> {
let PartialSource::Http {
requested_url: checkpoint_requested,
final_url,
etag,
representation_length,
} = &partial.source
else {
return Err(failed(
RuntimeFailureClass::SourceChanged,
"checkpoint source identity is unavailable",
));
};
let accepted = CanonicalUrl::parse(requested_url).map_err(|_| {
failed(
RuntimeFailureClass::SourceChanged,
"checkpoint source identity is unavailable",
)
})?;
let checkpoint_requested = CanonicalUrl::parse(checkpoint_requested).map_err(|_| {
failed(
RuntimeFailureClass::SourceChanged,
"checkpoint source identity is unavailable",
)
})?;
if accepted != checkpoint_requested || partial.result_path != result_path {
return Err(failed(
RuntimeFailureClass::SourceChanged,
"checkpoint source identity is unavailable",
));
}
let identity = HttpSourceIdentity::from_parts(
checkpoint_requested,
CanonicalUrl::parse(final_url).map_err(|_| {
failed(
RuntimeFailureClass::SourceChanged,
"checkpoint source identity is unavailable",
)
})?,
StrongEtag::parse(etag).map_err(http_failure)?,
*representation_length,
);
HttpResumeCheckpoint::new(
identity,
MemberPath::parse(result_path).map_err(|_| {
failed(
RuntimeFailureClass::SourceChanged,
"checkpoint source identity is unavailable",
)
})?,
partial.destination.parts.iter().map(|part| part.size).sum(),
IntegrityDigest::parse(&partial.prefix_integrity).map_err(|_| {
failed(
RuntimeFailureClass::SourceChanged,
"checkpoint source identity is unavailable",
)
})?,
)
.map_err(http_failure)
}
pub(super) fn object_resume_checkpoint(
partial: &PartialCheckpoint,
endpoint: &str,
bucket: &str,
key: &str,
version_id: Option<&str>,
result_path: &str,
) -> Result<ObjectResumeCheckpoint, AdapterFailure> {
let PartialSource::Object {
endpoint: checkpoint_endpoint,
bucket: checkpoint_bucket,
key: checkpoint_key,
size,
version_id: checkpoint_version,
validator,
} = &partial.source
else {
return Err(failed(
RuntimeFailureClass::SourceChanged,
"checkpoint source identity is unavailable",
));
};
if CanonicalUrl::parse(endpoint).ok() != CanonicalUrl::parse(checkpoint_endpoint).ok()
|| bucket != checkpoint_bucket
|| key != checkpoint_key
|| version_id != checkpoint_version.as_deref()
|| result_path != partial.result_path
{
return Err(failed(
RuntimeFailureClass::SourceChanged,
"checkpoint source identity is unavailable",
));
}
let identity = ObjectSourceIdentity::from_parts(
CanonicalUrl::parse(checkpoint_endpoint).map_err(|_| {
failed(
RuntimeFailureClass::SourceChanged,
"checkpoint source identity is unavailable",
)
})?,
checkpoint_bucket,
checkpoint_key,
*size,
checkpoint_version.clone(),
validator
.as_deref()
.map(ObjectValidator::parse)
.transpose()
.map_err(object_failure)?,
)
.map_err(object_failure)?;
let parts = checkpoint_parts(&partial.destination)?;
let destination = DestinationResumeProof::new(
CanonicalUrl::parse(&partial.destination.endpoint).map_err(|_| {
failed(
RuntimeFailureClass::SourceChanged,
"checkpoint resume evidence is unavailable",
)
})?,
multipart_identity(&partial.destination),
parts,
None,
)
.map_err(object_failure)?;
ObjectResumeCheckpoint::new(
identity,
MemberPath::parse(result_path).map_err(|_| {
failed(
RuntimeFailureClass::SourceChanged,
"checkpoint source identity is unavailable",
)
})?,
partial.destination.parts.iter().map(|part| part.size).sum(),
IntegrityDigest::parse(&partial.prefix_integrity).map_err(|_| {
failed(
RuntimeFailureClass::SourceChanged,
"checkpoint source identity is unavailable",
)
})?,
destination,
)
.map_err(object_failure)
}
pub(super) fn hash_text(value: &str) -> String {
hex_digest(Sha256::digest(value.as_bytes()).as_slice())
}
pub(super) fn hex_digest(bytes: &[u8]) -> String {
let mut value = String::from("sha256:");
for byte in bytes {
use std::fmt::Write as _;
write!(value, "{byte:02x}").expect("writing to String cannot fail");
}
value
}

117
crates/remora-worker/src/executor/staging.rs Normal file
View file

@ -0,0 +1,117 @@
use super::*;
pub(super) fn staging_object(
target: &TargetStoreConfig,
key: &str,
version_id: Option<String>,
) -> StagingObjectIdentityDto {
StagingObjectIdentityDto {
store: target.target_store_id.clone(),
bucket: target.bucket.clone(),
key: key.to_owned(),
version_id,
}
}
pub(super) fn staging_write_intent(target: &TargetStoreConfig, key: &str) -> StagingWriteIntentDto {
StagingWriteIntentDto {
store: target.target_store_id.clone(),
bucket: target.bucket.clone(),
key: key.to_owned(),
}
}
pub(super) fn staging_multipart(identity: &MultipartIdentity) -> StagingMultipartIdentityDto {
StagingMultipartIdentityDto {
store: identity.target_store_id.clone(),
bucket: identity.bucket.clone(),
key: identity.key.clone(),
upload_id: identity.upload_id.clone(),
}
}
pub(super) fn checkpoint_parts(
checkpoint: &MultipartCheckpoint,
) -> Result<Vec<CompletedPart>, AdapterFailure> {
checkpoint
.parts
.iter()
.map(|part| {
if part.part_number == 0 || part.etag.is_empty() || part.size == 0 {
return Err(failed(
RuntimeFailureClass::SourceChanged,
"checkpoint resume evidence is unavailable",
));
}
Ok(CompletedPart {
part_number: part.part_number,
etag: part.etag.clone(),
size: Some(part.size),
})
})
.collect()
}
pub(super) fn multipart_identity(checkpoint: &MultipartCheckpoint) -> MultipartIdentity {
MultipartIdentity {
target_store_id: checkpoint.target_store_id.clone(),
bucket: checkpoint.bucket.clone(),
key: checkpoint.key.clone(),
upload_id: checkpoint.upload_id.clone(),
}
}
pub(super) fn checkpoint_member_id(key: &str) -> Result<OpaqueId, AdapterFailure> {
remora_target::validate_attempt_key(key).map_err(target_failure)?;
let segments = key.split('/').collect::<Vec<_>>();
let reserved = segments
.iter()
.position(|segment| *segment == remora_target::RESERVED_NAMESPACE_SEGMENT)
.ok_or_else(|| {
failed(
RuntimeFailureClass::SourceChanged,
"checkpoint resume evidence is unavailable",
)
})?;
OpaqueId::parse(segments[reserved + 4]).map_err(|_| {
failed(
RuntimeFailureClass::SourceChanged,
"checkpoint resume evidence is unavailable",
)
})
}
pub(super) fn staging_origin(envelope: &ExecutionEnvelopeDto) -> Result<OpaqueId, AdapterFailure> {
let root = if envelope.target.prefix.is_empty() {
format!("{}/a/", remora_target::RESERVED_NAMESPACE_SEGMENT)
} else {
format!(
"{}/{}/a/",
envelope.target.prefix,
remora_target::RESERVED_NAMESPACE_SEGMENT
)
};
envelope
.staging_target
.prefix
.strip_prefix(&root)
.and_then(|value| value.strip_suffix('/'))
.and_then(|value| OpaqueId::parse(value).ok())
.ok_or_else(|| {
failed(
RuntimeFailureClass::Authorization,
"staging target is invalid",
)
})
}
pub(super) fn generated_id(prefix: char) -> OpaqueId {
let mut bytes = [0_u8; 16];
getrandom::fill(&mut bytes).expect("operating system randomness is required");
let mut value = prefix.to_string();
for byte in bytes {
use std::fmt::Write as _;
write!(value, "{byte:02x}").expect("writing to String cannot fail");
}
OpaqueId::parse(value).expect("generated identifier is valid")
}

434
crates/remora-worker/src/executor/transfer.rs Normal file
View file

@ -0,0 +1,434 @@
use super::*;
impl AdapterExecutor {
#[allow(clippy::too_many_arguments)]
pub(super) async fn copy_member<S: SourceBody>(
&self,
context: &AttemptContext,
envelope: &ExecutionEnvelopeDto,
target: &TargetStore,
inventory: &ResultInventory,
path: String,
mut source: S,
member_size: Option<u64>,
aggregate_total_bytes: Option<u64>,
base_completed_bytes: u64,
base_completed_members: u64,
total_members: Option<u64>,
resolved_source_id: Option<&OpaqueId>,
partial_source: Option<PartialSource>,
) -> Result<(), AdapterFailure> {
context
.require_live_authority()
.await
.map_err(authority_failure)?;
let member_id = generated_id('r');
let staging_origin = staging_origin(envelope)?;
let request = CreateObjectRequest::for_target_attempt_object(
&envelope.target.prefix,
&staging_origin,
&member_id,
Some(&path),
);
context
.register_staging_write_intent(staging_write_intent(target.config(), request.key()))
.await
.map_err(authority_failure)?;
let (stored, completed_upload_id) = if member_size == Some(0) {
let mut probe = [0_u8; 1];
match source.read(&mut probe).await {
Ok(0) => {}
Ok(_) => {
return Err(failed(
RuntimeFailureClass::SourceChanged,
"source identity changed",
));
}
Err(error) => return Err(stream_failure(error)),
}
(
target
.put_create_only(&request, Bytes::new())
.await
.map_err(target_failure)?,
None,
)
} else {
let session = target
.begin_multipart(&request)
.await
.map_err(target_failure)?;
let upload_id = session.identity().upload_id.clone();
context
.register_staging_multipart(staging_multipart(session.identity()))
.await
.map_err(authority_failure)?;
let multipart_target =
MultipartTargetBody::tracked(target.clone(), session, Vec::new(), true);
let checkpoint = match (partial_source, resolved_source_id) {
(Some(source), Some(resolved_source_id)) => Some(CheckpointHooks {
descriptors: self.descriptors.clone(),
attempt_id: envelope.attempt_id.clone(),
staging_origin_attempt_id: staging_origin.clone(),
resolved_source_id: resolved_source_id.clone(),
source,
result_path: path.clone(),
destination_endpoint: target.config().endpoint.clone(),
multipart: multipart_target.progress,
last_persisted_bytes: 0,
}),
_ => None,
};
let mut config = self.copy_config.clone();
config.initial_completed_bytes = 0;
config.total_bytes = member_size;
config.completed_members = base_completed_members;
config.total_members = total_members;
let outcome = copy_stream(
source,
multipart_target.body,
config,
context.cancellation(),
ProgressHooks {
context: context.clone(),
base_completed_bytes,
aggregate_total_bytes,
checkpoint,
},
)
.await;
match outcome {
CopyOutcome::Completed(_) => {}
CopyOutcome::Cancelled { .. } => return Err(AdapterFailure::Cancelled),
CopyOutcome::Failed { failure, .. } => {
let class = failure
.runtime_class()
.unwrap_or(RuntimeFailureClass::Internal);
return Err(failed(class, public_message(class)));
}
}
(
multipart_target
.stored
.lock()
.expect("tracked target lock is not poisoned")
.take()
.ok_or_else(|| {
failed(RuntimeFailureClass::Internal, "target completion failed")
})?,
Some(upload_id),
)
};
context
.require_live_authority()
.await
.map_err(authority_failure)?;
context
.register_staging_object(
staging_object(
target.config(),
&stored.identity.key,
stored.identity.version_id.clone(),
),
completed_upload_id,
)
.await
.map_err(authority_failure)?;
inventory
.add_member(
&envelope.attempt_id,
&ResultMember::new(
member_id,
&path,
stored.size,
stored.identity,
stored.content_sha256,
)
.map_err(result_failure)?,
)
.await
.map_err(result_failure)?;
context
.report_progress(ProgressValueDto {
completed_bytes: safe(base_completed_bytes.saturating_add(stored.size))?,
total_bytes: aggregate_total_bytes.map(safe).transpose()?,
completed_members: safe(base_completed_members.saturating_add(1))?,
total_members: total_members.map(safe).transpose()?,
})
.await
.map_err(authority_failure)?;
let validated_bytes = base_completed_bytes.saturating_add(stored.size);
let validated_members = base_completed_members.saturating_add(1);
if let Some(resolved_source_id) = resolved_source_id {
let checkpoint = WorkerCheckpoint {
staging_origin_attempt_id: staging_origin,
resolved_source_id: resolved_source_id.clone(),
validated_bytes,
validated_members,
partial: None,
};
let descriptor = self
.descriptors
.store_checkpoint(&envelope.attempt_id, &checkpoint, Some(validated_members))
.await
.map_err(authority_failure)?;
context
.persist_checkpoint(
descriptor,
resolved_source_id.clone(),
safe(validated_bytes)?,
safe(validated_members)?,
None,
)
.await
.map_err(authority_failure)?;
}
Ok(())
}
pub(super) async fn resume_destination(
&self,
envelope: &ExecutionEnvelopeDto,
target: &TargetStore,
partial: &PartialCheckpoint,
) -> Result<(MultipartSession, Vec<CompletedPart>, OpaqueId), AdapterFailure> {
let expected_endpoint = CanonicalUrl::parse(&target.config().endpoint).map_err(|_| {
failed(
RuntimeFailureClass::SourceChanged,
"checkpoint resume evidence is unavailable",
)
})?;
let observed_endpoint =
CanonicalUrl::parse(&partial.destination.endpoint).map_err(|_| {
failed(
RuntimeFailureClass::SourceChanged,
"checkpoint resume evidence is unavailable",
)
})?;
let identity = multipart_identity(&partial.destination);
let member_id = checkpoint_member_id(&identity.key)?;
let origin = staging_origin(envelope)?;
let expected = CreateObjectRequest::for_target_attempt_object(
&envelope.target.prefix,
&origin,
&member_id,
Some(&partial.result_path),
);
if expected_endpoint != observed_endpoint
|| identity.target_store_id != target.config().target_store_id
|| identity.bucket != target.config().bucket
|| identity.key != expected.key()
{
return Err(failed(
RuntimeFailureClass::SourceChanged,
"checkpoint resume evidence is unavailable",
));
}
let expected_parts = checkpoint_parts(&partial.destination)?;
let (session, live_parts) = target
.resume_multipart(identity, None)
.await
.map_err(target_failure)?;
if live_parts != expected_parts {
return Err(failed(
RuntimeFailureClass::SourceChanged,
"checkpoint resume evidence is unavailable",
));
}
Ok((session, live_parts, member_id))
}
#[allow(clippy::too_many_arguments)]
pub(super) async fn copy_resumed_member<S: SourceBody>(
&self,
context: &AttemptContext,
envelope: &ExecutionEnvelopeDto,
target: &TargetStore,
inventory: &ResultInventory,
path: String,
source: S,
total_bytes: u64,
resolved_source_id: &OpaqueId,
session: MultipartSession,
parts: Vec<CompletedPart>,
member_id: OpaqueId,
) -> Result<(), AdapterFailure> {
let completed_bytes = parts
.iter()
.filter_map(|part| part.size)
.fold(0_u64, u64::saturating_add);
if completed_bytes == 0 || completed_bytes >= total_bytes {
return Err(failed(
RuntimeFailureClass::SourceChanged,
"checkpoint resume evidence is unavailable",
));
}
let upload_id = session.identity().upload_id.clone();
let multipart_target = MultipartTargetBody::tracked(target.clone(), session, parts, false);
let mut config = self.copy_config.clone();
config.initial_completed_bytes = 0;
config.total_bytes = Some(total_bytes - completed_bytes);
config.completed_members = 0;
config.total_members = Some(1);
let outcome = copy_stream(
source,
multipart_target.body,
config,
context.cancellation(),
ProgressHooks {
context: context.clone(),
base_completed_bytes: completed_bytes,
aggregate_total_bytes: Some(total_bytes),
checkpoint: None,
},
)
.await;
match outcome {
CopyOutcome::Completed(_) => {}
CopyOutcome::Cancelled { .. } => return Err(AdapterFailure::Cancelled),
CopyOutcome::Failed { failure, .. } => {
let class = failure
.runtime_class()
.unwrap_or(RuntimeFailureClass::Internal);
return Err(failed(class, public_message(class)));
}
}
let stored = multipart_target
.stored
.lock()
.expect("tracked target lock is not poisoned")
.take()
.ok_or_else(|| failed(RuntimeFailureClass::Internal, "target completion failed"))?;
if stored.size != total_bytes {
return Err(failed(
RuntimeFailureClass::Integrity,
"completed target size is invalid",
));
}
context
.require_live_authority()
.await
.map_err(authority_failure)?;
context
.register_staging_object(
staging_object(
target.config(),
&stored.identity.key,
stored.identity.version_id.clone(),
),
Some(upload_id),
)
.await
.map_err(authority_failure)?;
inventory
.add_member(
&envelope.attempt_id,
&ResultMember::new(
member_id,
&path,
stored.size,
stored.identity,
stored.content_sha256,
)
.map_err(result_failure)?,
)
.await
.map_err(result_failure)?;
context
.report_progress(ProgressValueDto {
completed_bytes: safe(total_bytes)?,
total_bytes: Some(safe(total_bytes)?),
completed_members: safe(1)?,
total_members: Some(safe(1)?),
})
.await
.map_err(authority_failure)?;
let staging_origin = staging_origin(envelope)?;
let checkpoint = WorkerCheckpoint {
staging_origin_attempt_id: staging_origin,
resolved_source_id: resolved_source_id.clone(),
validated_bytes: total_bytes,
validated_members: 1,
partial: None,
};
let descriptor = self
.descriptors
.store_checkpoint(&envelope.attempt_id, &checkpoint, Some(1))
.await
.map_err(authority_failure)?;
context
.persist_checkpoint(
descriptor,
resolved_source_id.clone(),
safe(total_bytes)?,
safe(1)?,
None,
)
.await
.map_err(authority_failure)?;
Ok(())
}
#[allow(clippy::too_many_arguments)]
pub(super) async fn resume_member(
&self,
context: &AttemptContext,
envelope: &ExecutionEnvelopeDto,
inventory: &ResultInventory,
target: &TargetStore,
resume: &ResumeState,
resolved_source_id: &OpaqueId,
path: &str,
expected_size: Option<u64>,
completed_members: u64,
completed_bytes: u64,
) -> Result<bool, AdapterFailure> {
if resume.validated_members < completed_members {
return Ok(false);
}
if resume.resolved_source_id.as_ref() != Some(resolved_source_id) {
return Err(failed(
RuntimeFailureClass::SourceChanged,
"checkpoint source identity is unavailable",
));
}
context
.require_live_authority()
.await
.map_err(authority_failure)?;
let member = inventory
.verify_member_path(&envelope.attempt_id, path, target)
.await
.map_err(result_failure)?
.ok_or_else(|| {
failed(
RuntimeFailureClass::SourceChanged,
"checkpoint resume evidence is unavailable",
)
})?;
if expected_size.is_some_and(|size| size != member.size()) {
return Err(failed(
RuntimeFailureClass::SourceChanged,
"checkpoint resume evidence is unavailable",
));
}
if completed_members == resume.validated_members
&& completed_bytes != resume.validated_bytes
{
return Err(failed(
RuntimeFailureClass::SourceChanged,
"checkpoint resume evidence is unavailable",
));
}
context
.report_progress(ProgressValueDto {
completed_bytes: safe(completed_bytes)?,
total_bytes: None,
completed_members: safe(completed_members)?,
total_members: None,
})
.await
.map_err(authority_failure)?;
Ok(true)
}
}