[llvm] Make sure that BenchmarkFactory::close() is called

compiler_gym/envs/llvm/service/BUILD

@@ -63,6 +63,7 @@ cc_binary(
     name = "compiler_gym-llvm-service-prelinked",
     srcs = ["RunService.cc"],
     deps = [
+        ":BenchmarkFactory",
         ":LlvmSession",
         "//compiler_gym/service/runtime:cc_runtime",
     ],

compiler_gym/envs/llvm/service/Benchmark.cc

@@ -162,9 +162,12 @@ Benchmark::Benchmark(const std::string& name, std::unique_ptr<llvm::LLVMContext>
       needsRecompile_(true) {}
 
 void Benchmark::close() {
+  VLOG(3) << "Closing benchmark " << name() << " with scratch directory "
+          << scratchDirectory().string();
   sys::error_code ec;
   fs::remove_all(scratchDirectory(), ec);
   CHECK(!ec) << "Failed to delete scratch directory: " << scratchDirectory().string();
+  VLOG(3) << "Closed benchmark " << name();
 }
 
 std::unique_ptr<Benchmark> Benchmark::clone(const fs::path& workingDirectory) const {

compiler_gym/envs/llvm/service/BenchmarkFactory.cc

@@ -48,6 +48,7 @@ void BenchmarkFactory::close() {
   for (auto& entry : benchmarks_) {
     entry.second.close();
   }
+  benchmarks_.clear();
 }
 
 Status BenchmarkFactory::getBenchmark(const BenchmarkProto& benchmarkMessage,

compiler_gym/envs/llvm/service/CMakeLists.txt

@@ -21,6 +21,7 @@ cg_cc_binary(
     "RunService.cc"
   DEPS
     ::LlvmSession
+    ::BenchmarkFactory
     compiler_gym::service::runtime::cc_runtime
 )
 

compiler_gym/envs/llvm/service/RunService.cc

@@ -2,6 +2,7 @@
 //
 // This source code is licensed under the MIT license found in the
 // LICENSE file in the root directory of this source tree.
+#include "compiler_gym/envs/llvm/service/BenchmarkFactory.h"
 #include "compiler_gym/envs/llvm/service/LlvmSession.h"
 #include "compiler_gym/service/runtime/Runtime.h"
 #include "llvm/InitializePasses.h"
@@ -59,5 +60,17 @@ void initLlvm() {
 
 int main(int argc, char** argv) {
   initLlvm();
-  createAndRunCompilerGymService<LlvmSession>(argc, argv, usage);
+  const auto ret = createAndRunCompilerGymService<LlvmSession>(argc, argv, usage);
+
+  // NOTE(github.com/facebookresearch/CompilerGym/issues/582): We need to make
+  // sure that BenchmarkFactory::close() is called on the global singleton
+  // instance, so that the temporary scratch directories are tidied up.
+  //
+  // TODO(github.com/facebookresearch/CompilerGym/issues/591): Once the runtime
+  // has been refactored to support intra-session mutable state, this singleton
+  // can be replaced by a member variable that is closed on
+  // CompilerGymServiceContext::shutdown().
+  BenchmarkFactory::getSingleton(FLAGS_working_dir).close();
+
+  return ret;
 }

compiler_gym/service/runtime/CreateAndRunCompilerGymServiceImpl.h

@@ -51,7 +51,7 @@ void setGrpcChannelOptions(grpc::ServerBuilder& builder);
 //       createAndRunCompilerGymServiceImpl(argc, argv, "usage string");
 //     }
 template <typename CompilationSessionType>
-[[noreturn]] void createAndRunCompilerGymServiceImpl(int argc, char** argv, const char* usage) {
+[[nodiscard]] int createAndRunCompilerGymServiceImpl(int argc, char** argv, const char* usage) {
   // Register a signal handler for SIGTERM that will set the shutdown_signal
   // future value.
   std::signal(SIGTERM, shutdown_handler);
@@ -62,7 +62,7 @@ template <typename CompilationSessionType>
   gflags::ParseCommandLineFlags(&argc, &argv, /*remove_flags=*/true);
   if (argc > 1) {
     std::cerr << "ERROR: unknown command line argument '" << argv[1] << '\'';
-    exit(1);
+    return 1;
   }
 
   // Set up the working and logging directories.
@@ -129,15 +129,16 @@ template <typename CompilationSessionType>
   VLOG(2) << "Shutting down the RPC service";
   server->Shutdown();
   serverThread.join();
+  VLOG(2) << "Service closed";
 
   if (service.sessionCount()) {
     LOG(ERROR) << "ERROR: Killing a service with " << service.sessionCount()
                << (service.sessionCount() > 1 ? " active sessions!" : " active session!")
                << std::endl;
-    exit(6);
+    return 6;
   }
 
-  exit(0);
+  return 0;
 }
 
 }  // namespace compiler_gym::runtime

compiler_gym/service/runtime/Runtime.h

@@ -20,20 +20,20 @@ namespace compiler_gym::runtime {
  *     #include "my_compiler_service/MyCompilationSession.h"
  *
  *     int main(int argc, char** argv) {
- *       createAndRunCompilerGymService<MyCompilationSession>(
+ *       return createAndRunCompilerGymService<MyCompilationSession>(
  *           argc, argc, "My compiler service"
  *       );
  *     }
  * \endcode
  *
- * This function never returns.
- *
  * @tparam CompilationSessionType A sublass of CompilationSession that provides
  *    implementations of the abstract methods.
+ *
+ * @return An integer return code.
  */
 template <typename CompilationSessionType>
-[[noreturn]] void createAndRunCompilerGymService(int argc, char** argv, const char* usage) {
-  createAndRunCompilerGymServiceImpl<CompilationSessionType>(argc, argv, usage);
+[[nodiscard]] int createAndRunCompilerGymService(int argc, char** argv, const char* usage) {
+  return createAndRunCompilerGymServiceImpl<CompilationSessionType>(argc, argv, usage);
 }
 
 }  // namespace compiler_gym::runtime

compiler_gym/service/runtime/create_and_run_compiler_gym_service.py

@@ -130,6 +130,7 @@ def main(argv):
         logging.info("Shutting down the RPC service")
         server.stop(60).wait()
         server_thread.join()
+        logging.info("Service closed")
 
         if len(service.sessions):
             print(

examples/example_compiler_gym_service/service_cc/ExampleService.cc

@@ -133,5 +133,5 @@ class ExampleCompilationSession final : public CompilationSession {
 }  // namespace
 
 int main(int argc, char** argv) {
-  runtime::createAndRunCompilerGymService<ExampleCompilationSession>(argc, argv, usage);
+  return runtime::createAndRunCompilerGymService<ExampleCompilationSession>(argc, argv, usage);
 }