Support JSON via stdin for celer-g4 (celeritas-project#1048)

sethrj · Dec 5, 2023 · 1782ddb · 1782ddb
1 parent 2d6d983
commit 1782ddb
Show file tree

Hide file tree

Showing 5 changed files with 82 additions and 17 deletions.
diff --git a/app/celer-g4/GlobalSetup.cc b/app/celer-g4/GlobalSetup.cc
@@ -112,13 +112,26 @@ void GlobalSetup::SetIgnoreProcesses(SetupOptions::VecString ignored)
  */
 void GlobalSetup::ReadInput(std::string const& filename)
 {
-    if (ends_with(filename, ".json"))
+    bool is_json_file = ends_with(filename, ".json");
+    if (is_json_file || filename == "-")
     {
-        CELER_LOG(status) << "Reading JSON input from '" << filename << "'";
-        std::ifstream infile(filename);
-        CELER_VALIDATE(infile, << "failed to open '" << filename << "'");
+        CELER_LOG(status) << "Reading JSON input from '"
+                          << (is_json_file ? filename : "<stdin>") << "'";
+        std::istream* instream{nullptr};
+        std::ifstream infile;
+        if (is_json_file)
+        {
+            instream = &infile;
+            infile.open(filename);
+            CELER_VALIDATE(infile, << "failed to open '" << filename << "'");
+        }
+        else
+        {
+            instream = &std::cin;
+        }
+        CELER_ASSERT(instream);
 #if CELERITAS_USE_JSON
-        nlohmann::json::parse(infile).get_to(input_);
+        nlohmann::json::parse(*instream).get_to(input_);
 #else
         CELER_NOT_CONFIGURED("nlohmann_json");
 #endif

diff --git a/app/celer-g4/celer-g4.cc b/app/celer-g4/celer-g4.cc
@@ -76,6 +76,7 @@ void print_usage(std::string_view exec_name)
 {
     // clang-format off
     std::cerr << "usage: " << exec_name << " {input}.json\n"
+                 "       " << exec_name << " -\n"
                  "       " << exec_name << " {commands}.mac\n"
                  "       " << exec_name << " --interactive\n"
                  "       " << exec_name << " [--help|-h]\n"

diff --git a/app/celer-g4/test-harness.py b/app/celer-g4/test-harness.py
@@ -91,20 +91,64 @@ def strtobool(text):
 with open(inp_file, "w") as f:
     json.dump(inp, f, indent=1)
 
+kwargs = {}
+if use_celeritas:
+    # IO through streams should work with celeritas or g4 as driver, but just
+    # do it here as an example
+    inp_file = "-"
+    inp["output_file"] = "-"
+    env = dict(environ)
+    kwargs = dict(
+        input=json.dumps(inp).encode(),
+        stdout=subprocess.PIPE,
+        env=env,
+    )
+
+
 print("Running", exe, inp_file, file=stderr)
-result = subprocess.run([exe, inp_file])
+result = subprocess.run([exe, inp_file], **kwargs)
+
+if use_celeritas:
+    # Decode the output, fail if the run failed
+    out_text = result.stdout.decode()
+    try:
+        j = json.loads(out_text)
+    except json.decoder.JSONDecodeError as e:
+        print(f"error ({e}): expected a JSON object but got the following stdout:")
+        with open(out_file, 'w') as f:
+            f.write(out_text)
+        print("Wrote text to", out_file)
+        if result.returncode:
+            print("fatal:", str(e))
+            exit(result.returncode)
+    else:
+        with open(out_file, 'w') as f:
+            json.dump(j, f, indent=1)
 
 if result.returncode:
     print("fatal: run failed with error", result.returncode)
-    exit(result.returncode)
+    if use_celeritas:
+        try:
+            j = json.loads(result.stdout.decode())
+        except:
+            pass
+        else:
+            outfilename = f'{run_name}.out.failed.json'
+            with open(outfilename, 'w') as f:
+                json.dump(j, f, indent=1)
+            print("Failure written to", outfilename, file=stderr)
+            j = {}
 
-with open(out_file) as f:
-    result = json.load(f)
+    exit(result.returncode)
 
-pprint(result["result"])
+if not use_celeritas:
+    # Load written file
+    with open(out_file) as f:
+        j = json.load(f)
 
-# Rewrite with indentation
-with open(out_file, 'w') as f:
-    json.dump(result, f, indent=1)
+    # Rewrite with indentation
+    with open(out_file, 'w') as f:
+        json.dump(j, f, indent=1)
 
+pprint(j["result"])
 
diff --git a/app/celer-sim/simple-driver.py b/app/celer-sim/simple-driver.py
@@ -129,8 +129,6 @@ def strtobool(text):
 
 print("Received {} bytes of data".format(len(result.stdout)), file=stderr)
 out_text = result.stdout.decode()
-# Filter out spurious HepMC3 output
-out_text = out_text[out_text.find('\n{') + 1:]
 try:
     j = json.loads(out_text)
 except json.decoder.JSONDecodeError as e:

diff --git a/src/accel/SharedParams.cc b/src/accel/SharedParams.cc
@@ -215,8 +215,17 @@ SharedParams::SharedParams(SetupOptions const& options, SPOutputRegistry oreg)
     // Construct core data
     this->initialize_core(options);
 
-    // Set up output after params are constructed
-    this->try_output();
+    if (output_filename_ != "-")
+    {
+        // Write output after params are constructed before anything can go
+        // wrong
+        this->try_output();
+    }
+    else
+    {
+        CELER_LOG(debug) << "Skipping 'startup' JSON output since writing to "
+                            "stdout";
+    }
 
     if (!options.offload_output_file.empty())
     {