feat: model special resource requests too (#512)

* add special resource model * derive "special" in "using" section * update tests
martian-lang · Jul 18, 2024 · f7c10bb · f7c10bb
1 parent 7028107
commit f7c10bb
Show file tree

Hide file tree

Showing 8 changed files with 100 additions and 71 deletions.
diff --git a/martian-derive/src/lib.rs b/martian-derive/src/lib.rs
@@ -3,7 +3,6 @@
 /// TODO:
 /// - Handle default values for FileType
 /// - Repo wide reorganization
-extern crate proc_macro;
 use martian::{utils, MartianBlanketType, MartianPrimaryType, StageKind, Volatile, MARTIAN_TOKENS};
 use quote::quote;
 use std::collections::{HashMap, HashSet};
@@ -393,6 +392,7 @@ attr_parse!(
     mem_gb: i16,
     threads: i16,
     vmem_gb: i16,
+    special: String,
     volatile: Volatile,
     stage_name: String
 );

diff --git a/martian-derive/tests/ui_make_mro/attr_unknown_attr.stderr b/martian-derive/tests/ui_make_mro/attr_unknown_attr.stderr
@@ -1,4 +1,4 @@
-error: Expecting a comma separated `key=value` like tokens here. The allowed keys are: [mem_gb, threads, vmem_gb, volatile, stage_name]
+error: Expecting a comma separated `key=value` like tokens here. The allowed keys are: [mem_gb, threads, vmem_gb, special, volatile, stage_name]
  --> $DIR/attr_unknown_attr.rs:7:12
   |
 7 | #[make_mro(foo)]

diff --git a/martian-lab/examples/sum_sq/src/sum_squares.rs b/martian-lab/examples/sum_sq/src/sum_squares.rs
@@ -1,18 +1,16 @@
 //! SumSquares stage code
 
-use serde::{Deserialize, Serialize};
-
 // The prelude brings the following items in scope:
 // - Traits: MartianMain, MartianStage, RawMartianStage, MartianFileType, MartianMakePath
 // - Struct/Enum: MartianRover, Resource, StageDef, MartianVoid,
 //                Error (from anyhow crate), LevelFilter (from log crate)
 // - Macros: martian_stages!
 // - Functions: martian_main, martian_main_with_log_level, martian_make_mro
 use martian::prelude::*;
-
 // Bring the procedural macros in scope:
 // #[derive(MartianStruct)], #[derive(MartianType)], #[make_mro], martian_filetype!
 use martian_derive::{make_mro, MartianStruct};
+use serde::{Deserialize, Serialize};
 
 // NOTE: The following four structs will serve as the associated type for the
 // trait. The struct fields need to be owned and are limited to
@@ -73,7 +71,7 @@ impl MartianStage for SumSquares {
         for value in args.values {
             let chunk_inputs = SumSquaresChunkInputs { value };
             // It is optional to create a chunk with resource. If not specified, default resource will be used
-            stage_def.add_chunk_with_resource(chunk_inputs, chunk_resource);
+            stage_def.add_chunk_with_resource(chunk_inputs, chunk_resource.clone());
         }
         // Return the stage definition
         Ok(stage_def)

diff --git a/martian-lab/tests/error_test/error_test.json b/martian-lab/tests/error_test/error_test.json
@@ -1,17 +1,19 @@
 {
-  "command": [
-    "run_test.sh"
-  ],
-  "expected_return": 1,
-  "output_dir": "pipeline_test",
-  "contains_only_files": [
-    "*"
-  ],
-  "contents_match": [],
-  "file_contains_string": [
-    {
-      "file": "SUM_SQUARES/SUM_SQUARES/fork0/chnk3/_stackvars",
-      "string": "martian-rust/martian-lab/examples/sum_sq/src/sum_squares.rs:94"
-    }
-  ]
-}
+    "command": [
+        "run_test.sh"
+    ],
+    "expected_return": 1,
+    "output_dir": "pipeline_test",
+    "contains_only_files": [
+        "*"
+    ],
+    "contents_match": [
+
+    ],
+    "file_contains_string": [
+        {
+            "file": "SUM_SQUARES/SUM_SQUARES/fork0/chnk3/_stackvars",
+            "string": "martian-rust/martian-lab/examples/sum_sq/src/sum_squares.rs:92"
+        }
+    ]
+}
diff --git a/martian-lab/tests/error_test/expected/SUM_SQUARES/SUM_SQUARES/fork0/chnk3/_stackvars b/martian-lab/tests/error_test/expected/SUM_SQUARES/SUM_SQUARES/fork0/chnk3/_stackvars
@@ -1,44 +1,45 @@
-   0: anyhow::private::format_err
-             at ./cargo/registry/src/github.com-1ecc6299db9ec823/anyhow-1.0.51/src/lib.rs:659:13
-   1: <sum_sq::sum_squares::SumSquares as martian::stage::MartianStage>::main
-             at ./martian-lab/examples/sum_sq/src/sum_squares.rs:94:24
-   2: <T as martian::stage::RawMartianStage>::main
-             at ./martian/src/stage.rs:681:20
-   3: martian::martian_entry_point
-             at ./martian/src/lib.rs:262:9
-   4: martian::MartianAdapter<S>::run_get_error
-             at ./martian/src/lib.rs:164:9
-   5: martian::MartianAdapter<S>::run
-             at ./martian/src/lib.rs:157:9
-   6: sum_sq::main
-             at ./martian-lab/examples/sum_sq/src/main.rs:48:23
-   7: core::ops::function::FnOnce::call_once
-             at /rustc/c8dfcfe046a7680554bf4eb612bad840e7631c4b/library/core/src/ops/function.rs:227:5
-   8: std::sys_common::backtrace::__rust_begin_short_backtrace
-             at /rustc/c8dfcfe046a7680554bf4eb612bad840e7631c4b/library/std/src/sys_common/backtrace.rs:125:18
-   9: std::rt::lang_start::{{closure}}
-             at /rustc/c8dfcfe046a7680554bf4eb612bad840e7631c4b/library/std/src/rt.rs:63:18
-  10: core::ops::function::impls::<impl core::ops::function::FnOnce<A> for &F>::call_once
-             at /rustc/c8dfcfe046a7680554bf4eb612bad840e7631c4b/library/core/src/ops/function.rs:259:13
-      std::panicking::try::do_call
-             at /rustc/c8dfcfe046a7680554bf4eb612bad840e7631c4b/library/std/src/panicking.rs:401:40
-      std::panicking::try
-             at /rustc/c8dfcfe046a7680554bf4eb612bad840e7631c4b/library/std/src/panicking.rs:365:19
-      std::panic::catch_unwind
-             at /rustc/c8dfcfe046a7680554bf4eb612bad840e7631c4b/library/std/src/panic.rs:434:14
-      std::rt::lang_start_internal::{{closure}}
-             at /rustc/c8dfcfe046a7680554bf4eb612bad840e7631c4b/library/std/src/rt.rs:45:48
-      std::panicking::try::do_call
-             at /rustc/c8dfcfe046a7680554bf4eb612bad840e7631c4b/library/std/src/panicking.rs:401:40
-      std::panicking::try
-             at /rustc/c8dfcfe046a7680554bf4eb612bad840e7631c4b/library/std/src/panicking.rs:365:19
-      std::panic::catch_unwind
-             at /rustc/c8dfcfe046a7680554bf4eb612bad840e7631c4b/library/std/src/panic.rs:434:14
-      std::rt::lang_start_internal
-             at /rustc/c8dfcfe046a7680554bf4eb612bad840e7631c4b/library/std/src/rt.rs:45:20
-  11: std::rt::lang_start
-             at /rustc/c8dfcfe046a7680554bf4eb612bad840e7631c4b/library/std/src/rt.rs:62:5
-  12: main
-  13: __libc_start_main
-  14: <unknown>
-             at /sysdeps/x86_64/elf/start.S:103
+   0: anyhow::error::<impl anyhow::Error>::msg
+             at /mnt/bazelbuild/user/preyas.shah/cargo/registry/src/index.crates.io-6f17d22bba15001f/anyhow-1.0.86/src/backtrace.rs:27:14
+   1: anyhow::__private::format_err
+             at /mnt/bazelbuild/user/preyas.shah/cargo/registry/src/index.crates.io-6f17d22bba15001f/anyhow-1.0.86/src/lib.rs:689:13
+   2: <sum_sq::sum_squares::SumSquares as martian::stage::MartianStage>::main
+             at /mnt/home/preyas/Reps/martian-rust/martian-lab/examples/sum_sq/src/sum_squares.rs:92:24
+   3: <T as martian::stage::RawMartianStage>::main
+             at /mnt/home/preyas/Reps/martian-rust/martian/src/stage.rs:778:20
+   4: martian::martian_entry_point
+             at /mnt/home/preyas/Reps/martian-rust/martian/src/lib.rs:261:9
+   5: martian::MartianAdapter<S>::run_get_error
+             at /mnt/home/preyas/Reps/martian-rust/martian/src/lib.rs:165:9
+   6: martian::MartianAdapter<S>::run
+             at /mnt/home/preyas/Reps/martian-rust/martian/src/lib.rs:158:9
+   7: sum_sq::main
+             at /mnt/home/preyas/Reps/martian-rust/martian-lab/examples/sum_sq/src/main.rs:48:23
+   8: core::ops::function::FnOnce::call_once
+             at /rustc/129f3b9964af4d4a709d1383930ade12dfe7c081/library/core/src/ops/function.rs:250:5
+   9: std::sys_common::backtrace::__rust_begin_short_backtrace
+             at /rustc/129f3b9964af4d4a709d1383930ade12dfe7c081/library/std/src/sys_common/backtrace.rs:155:18
+  10: std::rt::lang_start::{{closure}}
+             at /rustc/129f3b9964af4d4a709d1383930ade12dfe7c081/library/std/src/rt.rs:159:18
+  11: core::ops::function::impls::<impl core::ops::function::FnOnce<A> for &F>::call_once
+             at /rustc/129f3b9964af4d4a709d1383930ade12dfe7c081/library/core/src/ops/function.rs:284:13
+  12: std::panicking::try::do_call
+             at /rustc/129f3b9964af4d4a709d1383930ade12dfe7c081/library/std/src/panicking.rs:559:40
+  13: std::panicking::try
+             at /rustc/129f3b9964af4d4a709d1383930ade12dfe7c081/library/std/src/panicking.rs:523:19
+  14: std::panic::catch_unwind
+             at /rustc/129f3b9964af4d4a709d1383930ade12dfe7c081/library/std/src/panic.rs:149:14
+  15: std::rt::lang_start_internal::{{closure}}
+             at /rustc/129f3b9964af4d4a709d1383930ade12dfe7c081/library/std/src/rt.rs:141:48
+  16: std::panicking::try::do_call
+             at /rustc/129f3b9964af4d4a709d1383930ade12dfe7c081/library/std/src/panicking.rs:559:40
+  17: std::panicking::try
+             at /rustc/129f3b9964af4d4a709d1383930ade12dfe7c081/library/std/src/panicking.rs:523:19
+  18: std::panic::catch_unwind
+             at /rustc/129f3b9964af4d4a709d1383930ade12dfe7c081/library/std/src/panic.rs:149:14
+  19: std::rt::lang_start_internal
+             at /rustc/129f3b9964af4d4a709d1383930ade12dfe7c081/library/std/src/rt.rs:141:20
+  20: std::rt::lang_start
+             at /rustc/129f3b9964af4d4a709d1383930ade12dfe7c081/library/std/src/rt.rs:158:17
+  21: main
+  22: __libc_start_main
+  23: <unknown>
diff --git a/martian-lab/tests/split_test/expected/SUM_SQUARES/SUM_SQUARES/fork0/split/_stage_defs b/martian-lab/tests/split_test/expected/SUM_SQUARES/SUM_SQUARES/fork0/split/_stage_defs
@@ -2,25 +2,29 @@
   "chunks": [
     {
       "__mem_gb": 1,
+      "__special": null,
       "__threads": 1,
       "__vmem_gb": null,
       "value": 1.0
     },
     {
       "__mem_gb": 1,
+      "__special": null,
       "__threads": 1,
       "__vmem_gb": null,
       "value": 2.0
     },
     {
       "__mem_gb": 1,
+      "__special": null,
       "__threads": 1,
       "__vmem_gb": null,
       "value": 3.0
     }
   ],
   "join": {
     "__mem_gb": null,
+    "__special": null,
     "__threads": null,
     "__vmem_gb": null
   }

diff --git a/martian/src/mro.rs b/martian/src/mro.rs
@@ -602,7 +602,7 @@ macro_rules! mro_using {
         ///     threads = 16,
         /// )
         /// ```
-        #[derive(Debug, Default, Clone, Copy, Serialize, Deserialize)]
+        #[derive(Debug, Default, Clone, Serialize, Deserialize)]
         pub struct MroUsing {
             $(pub $property: Option<$type>,)*
         }
@@ -622,7 +622,7 @@ macro_rules! mro_using {
                     return Ok(());
                 }
                 $(
-                    if let Some($property) = self.$property {
+                    if let Some($property) = self.$property.clone() {
                         writeln!(
                             f,
                             "{blank:indent$}{key:<width$} = {value},",
@@ -654,7 +654,7 @@ macro_rules! mro_using {
     };
 }
 
-mro_using! {mem_gb: i16, threads: i16, vmem_gb: i16, volatile: Volatile}
+mro_using! {mem_gb: i16, threads: i16, vmem_gb: i16, special: String, volatile: Volatile}
 
 /// Input and outputs fields together
 #[derive(Debug, Default)]

diff --git a/martian/src/stage.rs b/martian/src/stage.rs
@@ -117,14 +117,16 @@ impl<T: MartianFileType> MartianMakePath for T {
 ///
 /// Memory/ thread request can be negative in matrian. See
 /// [http://martian-lang.org/advanced-features/#resource-consumption](http://martian-lang.org/advanced-features/#resource-consumption)
-#[derive(Debug, Serialize, Deserialize, Copy, Clone, Default)]
+#[derive(Debug, Serialize, Deserialize, Clone, Default)]
 pub struct Resource {
     #[serde(rename = "__mem_gb")]
     mem_gb: Option<isize>,
     #[serde(rename = "__threads")]
     threads: Option<isize>,
     #[serde(rename = "__vmem_gb")]
     vmem_gb: Option<isize>,
+    #[serde(rename = "__special")]
+    special: Option<String>,
 }
 
 impl Resource {
@@ -164,6 +166,11 @@ impl Resource {
         self.threads
     }
 
+    /// Get the special resource request
+    pub fn get_special(&self) -> Option<String> {
+        self.special.clone()
+    }
+
     /// Set the mem_gb
     /// ```rust
     /// use martian::Resource;
@@ -206,6 +213,21 @@ impl Resource {
         self
     }
 
+    /// Set the special request
+    /// ```rust
+    /// use martian::Resource;
+    ///
+    /// let resource = Resource::new().mem_gb(2).vmem_gb(4).special("gpu_count1_mem8".to_owned());
+    /// assert_eq!(resource.get_mem_gb(), Some(2));
+    /// assert_eq!(resource.get_vmem_gb(), Some(4));
+    /// assert_eq!(resource.get_threads(), None);
+    /// assert_eq!(resource.get_special(), Some("gpu_count1_mem8".to_owned()));
+    /// ```
+    pub fn special(mut self, special: String) -> Self {
+        self.special = Some(special);
+        self
+    }
+
     /// Create a resource with the specified `mem_gb`. `vmem_gb` and
     /// `threads` are set to None.
     ///
@@ -222,6 +244,7 @@ impl Resource {
             mem_gb: Some(mem_gb),
             threads: None,
             vmem_gb: None,
+            special: None,
         }
     }
 
@@ -241,6 +264,7 @@ impl Resource {
             mem_gb: None,
             threads: Some(threads),
             vmem_gb: None,
+            special: None,
         }
     }
 }
@@ -609,7 +633,7 @@ pub trait MartianStage: MroMaker {
                     fill_defaults(resource),
                 ))
             }
-            let rover = _chunk_prelude(chunk_idx, run_directory, chunk.resource)?;
+            let rover = _chunk_prelude(chunk_idx, run_directory, chunk.resource.clone())?;
             self.main(args.clone(), chunk.inputs.clone(), rover)
         };