Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
2 changes: 1 addition & 1 deletion pre-compute/src/compute/pre_compute_app.rs
Original file line number Diff line number Diff line change
Expand Up @@ -187,7 +187,7 @@ mod tests {
input_files: urls.into_iter().map(String::from).collect(),
output_dir: output_dir.to_string(),
is_dataset_required: true,
bulk_size: 0,
iexec_bulk_slice_size: 0,
datasets: vec![Dataset {
url: HTTP_DATASET_URL.to_string(),
checksum: DATASET_CHECKSUM.to_string(),
Expand Down
124 changes: 53 additions & 71 deletions pre-compute/src/compute/pre_compute_args.rs
Original file line number Diff line number Diff line change
Expand Up @@ -15,7 +15,7 @@ pub struct PreComputeArgs {
// Input files
pub input_files: Vec<String>,
// Bulk processing
pub bulk_size: usize,
pub iexec_bulk_slice_size: usize,
pub datasets: Vec<Dataset>,
}

Expand All @@ -28,24 +28,24 @@ impl PreComputeArgs {
/// - `IEXEC_PRE_COMPUTE_OUT`: Output directory path
/// - `IEXEC_DATASET_REQUIRED`: Boolean ("true"/"false") indicating dataset requirement
/// - `IEXEC_INPUT_FILES_NUMBER`: Number of input files to load
/// - `BULK_SIZE`: Number of bulk datasets (0 means no bulk processing)
/// - `IEXEC_BULK_SLICE_SIZE`: Number of bulk datasets (0 means no bulk processing)
/// - Required when `IEXEC_DATASET_REQUIRED` = "true":
/// - `IEXEC_DATASET_URL`: Encrypted dataset URL
/// - `IEXEC_DATASET_KEY`: Base64-encoded dataset encryption key
/// - `IEXEC_DATASET_CHECKSUM`: Encrypted dataset checksum
/// - `IEXEC_DATASET_FILENAME`: Decrypted dataset filename
/// - Required when `BULK_SIZE` > 0 (for each dataset index from 1 to BULK_SIZE):
/// - `BULK_DATASET_#_URL`: Dataset URL
/// - `BULK_DATASET_#_CHECKSUM`: Dataset checksum
/// - `BULK_DATASET_#_FILENAME`: Dataset filename
/// - `BULK_DATASET_#_KEY`: Dataset decryption key
/// - Required when `IEXEC_BULK_SLICE_SIZE` > 0 (for each dataset index from 1 to IEXEC_BULK_SLICE_SIZE):
/// - `IEXEC_DATASET_#_URL`: Dataset URL
/// - `IEXEC_DATASET_#_CHECKSUM`: Dataset checksum
/// - `IEXEC_DATASET_#_FILENAME`: Dataset filename
/// - `IEXEC_DATASET_#_KEY`: Dataset decryption key
/// - Input file URLs (`IEXEC_INPUT_FILE_URL_1`, `IEXEC_INPUT_FILE_URL_2`, etc.)
///
/// # Errors
/// Returns `ReplicateStatusCause` error variants for:
/// - Missing required environment variables
/// - Invalid boolean values in `IEXEC_DATASET_REQUIRED`
/// - Invalid numeric format in `IEXEC_INPUT_FILES_NUMBER` or `BULK_SIZE`
/// - Invalid numeric format in `IEXEC_INPUT_FILES_NUMBER` or `IEXEC_BULK_SLICE_SIZE`
/// - Missing dataset parameters when required
/// - Missing input file URLs
/// - Missing bulk dataset parameters when bulk processing is enabled
Expand Down Expand Up @@ -73,52 +73,34 @@ impl PreComputeArgs {
.parse::<bool>()
.map_err(|_| ReplicateStatusCause::PreComputeIsDatasetRequiredMissing)?;

// Read bulk size (defaults to 0 if not present for backward compatibility)
let bulk_size_str = std::env::var(TeeSessionEnvironmentVariable::BulkSize.name())
.unwrap_or("0".to_string());
let bulk_size = bulk_size_str
// Read iexec bulk slice size (defaults to 0 if not present for backward compatibility)
let iexec_bulk_slice_size_str =
std::env::var(TeeSessionEnvironmentVariable::IexecBulkSliceSize.name())
.unwrap_or("0".to_string());
let iexec_bulk_slice_size = iexec_bulk_slice_size_str
.parse::<usize>()
.map_err(|_| ReplicateStatusCause::PreComputeIsDatasetRequiredMissing)?;
.map_err(|_| ReplicateStatusCause::PreComputeFailedUnknownIssue)?; // TODO: replace with a more specific error

let mut datasets = Vec::with_capacity(bulk_size + 1);
let mut datasets = Vec::with_capacity(iexec_bulk_slice_size + 1);

if is_dataset_required {
// Read datasets
let start_index = if is_dataset_required { 0 } else { 1 };
for i in start_index..=iexec_bulk_slice_size {
let url = get_env_var_or_error(
TeeSessionEnvironmentVariable::IexecDatasetUrl,
ReplicateStatusCause::PreComputeDatasetUrlMissing,
)?;
let key = get_env_var_or_error(
TeeSessionEnvironmentVariable::IexecDatasetKey,
ReplicateStatusCause::PreComputeDatasetKeyMissing,
TeeSessionEnvironmentVariable::IexecDatasetUrl(i),
ReplicateStatusCause::PreComputeDatasetUrlMissing, // TODO: replace with a more specific error for bulk dataset
)?;
let checksum = get_env_var_or_error(
TeeSessionEnvironmentVariable::IexecDatasetChecksum,
ReplicateStatusCause::PreComputeDatasetChecksumMissing,
TeeSessionEnvironmentVariable::IexecDatasetChecksum(i),
ReplicateStatusCause::PreComputeDatasetChecksumMissing, // TODO: replace with a more specific error for bulk dataset
)?;
let filename = get_env_var_or_error(
TeeSessionEnvironmentVariable::IexecDatasetFilename,
ReplicateStatusCause::PreComputeDatasetFilenameMissing,
)?;
datasets.push(Dataset::new(url, checksum, filename, key));
}

// Read bulk datasets
for i in 1..=bulk_size {
let url = get_env_var_or_error(
TeeSessionEnvironmentVariable::BulkDatasetUrl(i),
ReplicateStatusCause::PreComputeDatasetUrlMissing,
)?;
let checksum = get_env_var_or_error(
TeeSessionEnvironmentVariable::BulkDatasetChecksum(i),
ReplicateStatusCause::PreComputeDatasetChecksumMissing,
)?;
let filename = get_env_var_or_error(
TeeSessionEnvironmentVariable::BulkDatasetFilename(i),
ReplicateStatusCause::PreComputeDatasetFilenameMissing,
TeeSessionEnvironmentVariable::IexecDatasetFilename(i),
ReplicateStatusCause::PreComputeDatasetFilenameMissing, // TODO: replace with a more specific error for bulk dataset
)?;
let key = get_env_var_or_error(
TeeSessionEnvironmentVariable::BulkDatasetKey(i),
ReplicateStatusCause::PreComputeDatasetKeyMissing,
TeeSessionEnvironmentVariable::IexecDatasetKey(i),
ReplicateStatusCause::PreComputeDatasetKeyMissing, // TODO: replace with a more specific error for bulk dataset
)?;

datasets.push(Dataset::new(url, checksum, filename, key));
Expand All @@ -145,7 +127,7 @@ impl PreComputeArgs {
output_dir,
is_dataset_required,
input_files,
bulk_size,
iexec_bulk_slice_size,
datasets,
})
}
Expand All @@ -169,16 +151,16 @@ mod tests {
vars.insert(IexecPreComputeOut.name(), OUTPUT_DIR.to_string());
vars.insert(IsDatasetRequired.name(), "true".to_string());
vars.insert(IexecInputFilesNumber.name(), "0".to_string());
vars.insert(BulkSize.name(), "0".to_string()); // Default to no bulk processing
vars.insert(IexecBulkSliceSize.name(), "0".to_string()); // Default to no bulk processing
vars
}

fn setup_dataset_env_vars() -> HashMap<String, String> {
let mut vars = HashMap::new();
vars.insert(IexecDatasetUrl.name(), DATASET_URL.to_string());
vars.insert(IexecDatasetKey.name(), DATASET_KEY.to_string());
vars.insert(IexecDatasetChecksum.name(), DATASET_CHECKSUM.to_string());
vars.insert(IexecDatasetFilename.name(), DATASET_FILENAME.to_string());
vars.insert(IexecDatasetUrl(0).name(), DATASET_URL.to_string());
vars.insert(IexecDatasetKey(0).name(), DATASET_KEY.to_string());
vars.insert(IexecDatasetChecksum(0).name(), DATASET_CHECKSUM.to_string());
vars.insert(IexecDatasetFilename(0).name(), DATASET_FILENAME.to_string());
vars
}

Expand All @@ -198,19 +180,19 @@ mod tests {
// TODO: Collect all errors instead of propagating immediately, and return the list of errors
fn setup_bulk_dataset_env_vars(count: usize) -> HashMap<String, String> {
let mut vars = HashMap::new();
vars.insert(BulkSize.name(), count.to_string());
vars.insert(IexecBulkSliceSize.name(), count.to_string());

for i in 1..=count {
vars.insert(
BulkDatasetUrl(i).name(),
IexecDatasetUrl(i).name(),
format!("https://bulk-dataset-{i}.bin"),
);
vars.insert(BulkDatasetChecksum(i).name(), format!("0x{i}23checksum"));
vars.insert(IexecDatasetChecksum(i).name(), format!("0x{i}23checksum"));
vars.insert(
BulkDatasetFilename(i).name(),
IexecDatasetFilename(i).name(),
format!("bulk-dataset-{i}.txt"),
);
vars.insert(BulkDatasetKey(i).name(), format!("bulkKey{i}23"));
vars.insert(IexecDatasetKey(i).name(), format!("bulkKey{i}23"));
}
vars
}
Expand All @@ -235,7 +217,7 @@ mod tests {
assert!(!args.is_dataset_required);
assert_eq!(args.input_files.len(), 1);
assert_eq!(args.input_files[0], "https://input-1.txt");
assert_eq!(args.bulk_size, 0);
assert_eq!(args.iexec_bulk_slice_size, 0);
assert_eq!(args.datasets.len(), 0);
});
}
Expand All @@ -260,7 +242,7 @@ mod tests {
assert_eq!(args.datasets[0].checksum, DATASET_CHECKSUM.to_string());
assert_eq!(args.datasets[0].filename, DATASET_FILENAME.to_string());
assert_eq!(args.input_files.len(), 0);
assert_eq!(args.bulk_size, 0);
assert_eq!(args.iexec_bulk_slice_size, 0);
assert_eq!(args.datasets.len(), 1);
});
}
Expand All @@ -285,7 +267,7 @@ mod tests {
assert_eq!(args.input_files[0], "https://input-1.txt");
assert_eq!(args.input_files[1], "https://input-2.txt");
assert_eq!(args.input_files[2], "https://input-3.txt");
assert_eq!(args.bulk_size, 0);
assert_eq!(args.iexec_bulk_slice_size, 0);
assert_eq!(args.datasets.len(), 0);
});
}
Expand Down Expand Up @@ -359,7 +341,7 @@ mod tests {

assert_eq!(args.output_dir, OUTPUT_DIR);
assert!(!args.is_dataset_required);
assert_eq!(args.bulk_size, 3);
assert_eq!(args.iexec_bulk_slice_size, 3);
assert_eq!(args.datasets.len(), 3);
assert_eq!(args.input_files.len(), 0);

Expand Down Expand Up @@ -398,7 +380,7 @@ mod tests {

assert_eq!(args.output_dir, OUTPUT_DIR);
assert!(args.is_dataset_required);
assert_eq!(args.bulk_size, 2);
assert_eq!(args.iexec_bulk_slice_size, 2);
assert_eq!(args.datasets.len(), 3); // 1 regular + 2 bulk datasets
assert_eq!(args.input_files.len(), 0);

Expand All @@ -415,18 +397,18 @@ mod tests {
}

#[test]
fn read_args_fails_when_invalid_bulk_size_format() {
fn read_args_fails_when_invalid_iexec_bulk_slice_size_format() {
let mut env_vars = setup_basic_env_vars();
env_vars.insert(IsDatasetRequired.name(), "false".to_string());
env_vars.insert(BulkSize.name(), "not-a-number".to_string());
env_vars.insert(IexecBulkSliceSize.name(), "not-a-number".to_string());
env_vars.extend(setup_input_files_env_vars(0));

temp_env::with_vars(to_temp_env_vars(env_vars), || {
let result = PreComputeArgs::read_args();
assert!(result.is_err());
assert_eq!(
result.unwrap_err(),
ReplicateStatusCause::PreComputeIsDatasetRequiredMissing
ReplicateStatusCause::PreComputeFailedUnknownIssue
);
});
}
Expand All @@ -438,7 +420,7 @@ mod tests {
env_vars.extend(setup_input_files_env_vars(0));
env_vars.extend(setup_bulk_dataset_env_vars(2));
// Remove one of the bulk dataset URLs
env_vars.remove(&BulkDatasetUrl(1).name());
env_vars.remove(&IexecDatasetUrl(1).name());

temp_env::with_vars(to_temp_env_vars(env_vars), || {
let result = PreComputeArgs::read_args();
Expand All @@ -457,7 +439,7 @@ mod tests {
env_vars.extend(setup_input_files_env_vars(0));
env_vars.extend(setup_bulk_dataset_env_vars(2));
// Remove one of the bulk dataset checksums
env_vars.remove(&BulkDatasetChecksum(2).name());
env_vars.remove(&IexecDatasetChecksum(2).name());

temp_env::with_vars(to_temp_env_vars(env_vars), || {
let result = PreComputeArgs::read_args();
Expand All @@ -476,7 +458,7 @@ mod tests {
env_vars.extend(setup_input_files_env_vars(0));
env_vars.extend(setup_bulk_dataset_env_vars(3));
// Remove one of the bulk dataset filenames
env_vars.remove(&BulkDatasetFilename(2).name());
env_vars.remove(&IexecDatasetFilename(2).name());

temp_env::with_vars(to_temp_env_vars(env_vars), || {
let result = PreComputeArgs::read_args();
Expand All @@ -495,7 +477,7 @@ mod tests {
env_vars.extend(setup_input_files_env_vars(0));
env_vars.extend(setup_bulk_dataset_env_vars(2));
// Remove one of the bulk dataset keys
env_vars.remove(&BulkDatasetKey(1).name());
env_vars.remove(&IexecDatasetKey(1).name());

temp_env::with_vars(to_temp_env_vars(env_vars), || {
let result = PreComputeArgs::read_args();
Expand Down Expand Up @@ -525,19 +507,19 @@ mod tests {
ReplicateStatusCause::PreComputeInputFilesNumberMissing,
),
(
IexecDatasetUrl,
IexecDatasetUrl(0),
ReplicateStatusCause::PreComputeDatasetUrlMissing,
),
(
IexecDatasetKey,
IexecDatasetKey(0),
ReplicateStatusCause::PreComputeDatasetKeyMissing,
),
(
IexecDatasetChecksum,
IexecDatasetChecksum(0),
ReplicateStatusCause::PreComputeDatasetChecksumMissing,
),
(
IexecDatasetFilename,
IexecDatasetFilename(0),
ReplicateStatusCause::PreComputeDatasetFilenameMissing,
),
(
Expand Down
Loading