Merge pull request #105 from guzman-raphael/command_as_args

Make `Pod`'s command more flexible

Author: Synicix

src/core/orchestrator/docker.rs

@@ -130,12 +130,6 @@ impl LocalDockerOrchestrator {
             ),
             ("org.orcapod.pod_job.hash".to_owned(), pod_job.hash.clone()),
         ]);
-        let command = pod_job
-            .pod
-            .command
-            .split_whitespace()
-            .map(String::from)
-            .collect::<Vec<_>>();
 
         Ok((
             container_name.clone(),
@@ -145,8 +139,8 @@ impl LocalDockerOrchestrator {
             }),
             Config {
                 image: Some(image),
-                entrypoint: Some(command[..1].to_vec()),
-                cmd: Some(command[1..].to_vec()),
+                entrypoint: Some(pod_job.pod.command[..1].to_vec()),
+                cmd: Some(pod_job.pod.command[1..].to_vec()),
                 env: pod_job.env_vars.as_ref().map(|provided_env_vars| {
                     provided_env_vars
                         .iter()
@@ -170,7 +164,6 @@ impl LocalDockerOrchestrator {
         clippy::cast_precision_loss,
         clippy::cast_possible_truncation,
         clippy::indexing_slicing,
-        clippy::too_many_lines,
         reason = r#"
         - Timestamp and memory should always have a value > 0
         - Container will always have a name with more than 1 character
@@ -229,16 +222,11 @@ impl LocalDockerOrchestrator {
                                 .map(|(key, value)| (key.to_owned(), value.to_owned()))
                         })
                         .collect(),
-                    command: format!(
-                        "{} {}",
-                        container_spec
-                            .config
-                            .as_ref()?
-                            .entrypoint
-                            .as_ref()?
-                            .join(" "),
-                        container_spec.config.as_ref()?.cmd.as_ref()?.join(" ")
-                    ),
+                    command: [
+                        container_spec.config.as_ref()?.entrypoint.as_ref()?.clone(),
+                        container_spec.config.as_ref()?.cmd.as_ref()?.clone(),
+                    ]
+                    .concat(),
                     status: match (
                         container_spec.state.as_ref()?.status.as_ref()?,
                         container_spec.state.as_ref()?.exit_code? as i16,

src/lib.rs

@@ -28,4 +28,6 @@ pub mod core;
 ///       values are owned by Rust
 /// 1. No default trait implementations
 /// 1. (Rust limitation) No associated functions in traits e.g. class methods in Python
+/// 1. Hint: Enum variants with named fields offer a better UX (e.g. in Python) as opposed to
+///    unnamed enum fields i.e. will show up in help.
 pub mod uniffi;

src/uniffi/model/pod.rs

@@ -37,8 +37,9 @@ pub struct Pod {
     pub hash: String,
     /// Reproducible environment for compute.
     pub image: String,
-    /// Space-delimited shell command to begin computation.
-    pub command: String,
+    /// Shell command to begin computation. First element is the executable and remaining elements
+    /// are the arguments.
+    pub command: Vec<String>,
     /// Exposed, internal input specification.
     #[serde(serialize_with = "serialize_hashmap")]
     pub input_spec: HashMap<String, PathInfo>,
@@ -68,7 +69,7 @@ impl Pod {
     pub fn new(
         annotation: Option<Annotation>,
         image: String,
-        command: String,
+        command: Vec<String>,
         input_spec: HashMap<String, PathInfo>,
         output_dir: PathBuf,
         output_spec: HashMap<String, PathInfo>,

src/uniffi/orchestrator/agent.rs

@@ -7,6 +7,7 @@ use crate::{
         store::{Store as _, filestore::LocalFileStore},
     },
 };
+use colored::Colorize as _;
 use derive_more::Display;
 use futures_executor::block_on;
 use futures_util::future::join_all;
@@ -70,9 +71,9 @@ impl AgentClient {
             })?,
         })
     }
-    /// Submit many pod jobs to be processed in parallel.
+    /// Start many pod jobs to be processed in parallel.
     /// Return order will match inputs, casting outputs to `String` (since `uniffi` doesn't support sending unwrapped `Result`s).
-    pub async fn submit_pod_jobs(&self, pod_jobs: Vec<Arc<PodJob>>) -> Vec<Response> {
+    pub async fn start_pod_jobs(&self, pod_jobs: Vec<Arc<PodJob>>) -> Vec<Response> {
         join_all(pod_jobs.iter().map(|pod_job| async {
             match self
                 .publish(&format!("request/pod_job/{}", pod_job.hash), pod_job)
@@ -98,7 +99,7 @@ impl AgentClient {
             .context(selector::AgentCommunicationFailure {})?;
         while let Ok(sample) = subscriber.recv_async().await {
             let value = serde_json::from_slice::<Value>(&sample.payload().to_bytes())?;
-            println!("{}: {value:#}", sample.key_expr().as_str());
+            println!("{}: {value:#}", sample.key_expr().as_str().yellow());
         }
         Ok(())
     }
@@ -162,13 +163,19 @@ impl Agent {
                 Ok(pod_result)
             },
             async |client, pod_result| {
-                let response_topic = match &pod_result.status {
-                    Status::Completed => &format!("success/pod_job/{}", pod_result.pod_job.hash),
-                    Status::Running | Status::Failed(_) | Status::Unset => {
-                        &format!("failure/pod_job/{}", pod_result.pod_job.hash)
-                    }
-                };
-                client.publish(response_topic, &pod_result).await
+                client
+                    .publish(
+                        &format!(
+                            "{}/pod_job/{}",
+                            match &pod_result.status {
+                                Status::Completed => "success",
+                                Status::Running | Status::Failed(_) | Status::Unset => "failure",
+                            },
+                            pod_result.pod_job.hash
+                        ),
+                        &pod_result,
+                    )
+                    .await
             },
         ));
         if let Some(store) = available_store {

src/uniffi/orchestrator/mod.rs

@@ -42,7 +42,7 @@ pub struct RunInfo {
     /// Environment variables set in environment.
     pub env_vars: HashMap<String, String>,
     /// Command used to start run.
-    pub command: String,
+    pub command: Vec<String>,
     /// Current run status.
     pub status: Status,
     /// Mounted volume binds to the environment.

tests/agent.rs

@@ -126,7 +126,7 @@ async fn parallel_four_cores() -> Result<()> {
     });
     // submit requests
     client
-        .submit_pod_jobs(pod_jobs_stresser(image_reference, run_duration_secs, 3, 1)?)
+        .start_pod_jobs(pod_jobs_stresser(image_reference, run_duration_secs, 3, 1)?)
         .await;
 
     services

tests/error.rs

@@ -7,7 +7,7 @@
 
 pub mod fixture;
 use dot_parser::ast::Graph as DOTGraph;
-use fixture::{NAMESPACE_LOOKUP_READ_ONLY, pod_custom, pod_job_custom, pod_job_style};
+use fixture::{NAMESPACE_LOOKUP_READ_ONLY, pod_custom, pod_job_custom, pod_job_style, str_to_vec};
 use glob::glob;
 use orcapod::{
     core::crypto::hash_file,
@@ -124,7 +124,7 @@ fn internal_incomplete_packet() -> Result<()> {
         pod_job_custom(
             &pod_custom(
                 "alpine:3.14",
-                "echo",
+                &["echo".into()],
                 HashMap::from([(
                     "key_1".into(),
                     PathInfo {
@@ -159,15 +159,15 @@ fn internal_key_missing() {
 }
 
 #[tokio::test(flavor = "multi_thread", worker_threads = 1)]
-async fn submit_pod_jobs() -> Result<()> {
+async fn start_pod_jobs() -> Result<()> {
     let client = AgentClient::new("error".into(), "host".into())?;
     let mut pod_job = pod_job_custom(
-        &pod_custom("alpine:3.14", "sleep 5", HashMap::new())?,
+        &pod_custom("alpine:3.14", &str_to_vec("sleep 5"), HashMap::new())?,
         HashMap::new(),
         &NAMESPACE_LOOKUP_READ_ONLY,
     )?;
     pod_job.hash = "bad?hash".into();
-    let responses = client.submit_pod_jobs(vec![pod_job.into()]).await;
+    let responses = client.start_pod_jobs(vec![pod_job.into()]).await;
     assert!(
         responses.len() == 1,
         "Client received an unexpected number of pod job request responses."

tests/extra/python/agent_test.py

@@ -31,7 +31,7 @@ def count(sample):
         with session.declare_subscriber(
             f"group/{group}/success/pod_job/**", count
         ) as subscriber:
-            await asyncio.sleep(20)
+            await asyncio.sleep(20)  # wait for results
 
     if counter != pod_job_count:
         raise Exception(f"Unexpected successful pod job count: {counter}.")
@@ -48,7 +48,7 @@ async def main(client, agent, test_dir, namespace_lookup, pod_jobs):
     await asyncio.sleep(5)  # ensure service ready
 
     try:
-        await client.submit_pod_jobs(pod_jobs=pod_jobs)
+        await client.start_pod_jobs(pod_jobs=pod_jobs)
         await verify(client.group(), len(pod_jobs))
     finally:
         shutil.rmtree(test_dir)
@@ -82,7 +82,9 @@ async def main(client, agent, test_dir, namespace_lookup, pod_jobs):
             pod=Pod(
                 annotation=None,
                 image="ghcr.io/colinianking/stress-ng:e2f96874f951a72c1c83ff49098661f0e013ac40",
-                command="stress-ng --cpu 1 --cpu-load 100 --timeout 5 --metrics-brief",
+                command="stress-ng --cpu 1 --cpu-load 100 --timeout 5 --metrics-brief".split(
+                    " "
+                ),
                 input_spec={},
                 output_dir="/tmp/output",
                 output_spec={},

tests/extra/python/smoke_test.py

@@ -28,7 +28,7 @@ def create_pod(data, _):
             version="1.0.0",
         ),
         image="alpine:3.14",
-        command="sleep 1",
+        command="sleep 1".split(" "),
         input_spec={},
         output_dir="/tmp/output",
         output_spec={},

tests/fixture/mod.rs

@@ -41,7 +41,7 @@ pub fn pod_style() -> Result<Pod> {
             version: "1.0.0".to_owned(),
         }),
         "example.server.com/user/style-transfer:1.0.0".to_owned(),
-        "python /run.py".to_owned(),
+        str_to_vec("python /run.py"),
         HashMap::from([
             (
                 "extra-style".to_owned(),
@@ -148,7 +148,7 @@ pub fn pod_result_style(
 
 pub fn pod_custom(
     image_reference: &str,
-    command: &str,
+    command: &[String],
     input_spec: HashMap<String, PathInfo, RandomState>,
 ) -> Result<Pod> {
     Pod::new(
@@ -197,7 +197,7 @@ pub fn pod_jobs_stresser(
                 return Ok(pod_job_custom(
                     &pod_custom(
                         image_reference,
-                        &format!("stress-ng --cpu 1 --cpu-load 100 --timeout {run_duration_secs} --metrics-brief"),
+                        &str_to_vec(&format!("stress-ng --cpu 1 --cpu-load 100 --timeout {run_duration_secs} --metrics-brief")),
                         HashMap::new()
                     )?,
                     HashMap::new(),
@@ -206,7 +206,7 @@ pub fn pod_jobs_stresser(
                 .into());
             }
             Ok(pod_job_custom(
-                &pod_custom(image_reference, "sleep crash", HashMap::new())?,
+                &pod_custom(image_reference, &str_to_vec("sleep crash"), HashMap::new())?,
                 HashMap::new(),
                 &NAMESPACE_LOOKUP_READ_ONLY,
             )?
@@ -272,6 +272,10 @@ pub fn pull_image(reference: &str) -> Result<()> {
 
 // --- util ---
 
+pub fn str_to_vec(v: &str) -> Vec<String> {
+    v.split_whitespace().map(String::from).collect()
+}
+
 pub struct TestDirs(pub HashMap<String, TempDir>);
 
 impl TestDirs {