From 3363c2db3549cc534c6430e798e5526e2ca4ce6c Mon Sep 17 00:00:00 2001 From: Chris Cummins Date: Sun, 6 Mar 2022 19:11:23 +0000 Subject: [PATCH 1/3] [service] Add a CompilerGymServiceContext class. Issue #591. --- compiler_gym/service/BUILD | 12 +++ .../service/CompilerGymServiceContext.cc | 21 +++++ .../service/CompilerGymServiceContext.h | 90 +++++++++++++++++++ 3 files changed, 123 insertions(+) create mode 100644 compiler_gym/service/CompilerGymServiceContext.cc create mode 100644 compiler_gym/service/CompilerGymServiceContext.h diff --git a/compiler_gym/service/BUILD b/compiler_gym/service/BUILD index 280d508eb..f527a82e6 100644 --- a/compiler_gym/service/BUILD +++ b/compiler_gym/service/BUILD @@ -37,6 +37,18 @@ cc_library( ], ) +cc_library( + name = "CompilerGymServiceContext", + srcs = ["CompilerGymServiceContext.cc"], + hdrs = ["CompilerGymServiceContext.h"], + visibility = ["//visibility:public"], + deps = [ + "//compiler_gym/service/proto:compiler_gym_service_cc", + "@boost//:filesystem", + "@com_github_grpc_grpc//:grpc++", + ], +) + py_library( name = "connection", srcs = ["connection.py"], diff --git a/compiler_gym/service/CompilerGymServiceContext.cc b/compiler_gym/service/CompilerGymServiceContext.cc new file mode 100644 index 000000000..5d928f482 --- /dev/null +++ b/compiler_gym/service/CompilerGymServiceContext.cc @@ -0,0 +1,21 @@ +// Copyright (c) Facebook, Inc. and its affiliates. +// +// This source code is licensed under the MIT license found in the +// LICENSE file in the root directory of this source tree. +#pragma once + +#include "compiler_gym/service/CompilerGymServiceContext.h" + +using grpc::Status; + +namespace compiler_gym { + +CompilerGymServiceContext::CompilerGymServiceContext( + const boost::filesystem::path& workingDirectory) + : workingDirectory_(workingDirectory) {} + +virtual Status CompilerGymServiceContext::init() { return Status::OK; } + +virtual Status CompilerGymServiceContext::shutdown() { return Status::OK; } + +} // namespace compiler_gym diff --git a/compiler_gym/service/CompilerGymServiceContext.h b/compiler_gym/service/CompilerGymServiceContext.h new file mode 100644 index 000000000..67008a866 --- /dev/null +++ b/compiler_gym/service/CompilerGymServiceContext.h @@ -0,0 +1,90 @@ +// Copyright (c) Facebook, Inc. and its affiliates. +// +// This source code is licensed under the MIT license found in the +// LICENSE file in the root directory of this source tree. +#pragma once + +#include + +#include "boost/filesystem.hpp" + +namespace compiler_gym { + +/** + * Execution context of a compiler gym service. + * + * This class encapsulates mutable state that is shared between all compilation + * sessions. An instance of this class is passed to every new + * CompilationSession. + * + * You may subclass CompilerGymServiceContext to add additional mutable state. + * The subclass . + * + * \code{.cpp} + * + * #include "compiler_gym/service/CompilationSession.h" + * #include "compiler_gym/service/CompilerGymServiceContext.h" + * #include "compiler_gym/service/runtime/Runtime.h" + * + * using namespace compiler_gym; + * + * class MyServiceContext final : public ServiceContext { ... } + * + * class MyCompilationSession final : public CompilationSession { ... } + * + * int main(int argc, char** argv) { + * return runtime::createAndRunCompilerGymService(); + * } + * \endcode + */ +class CompilerGymServiceContext { + public: + CompilerGymServiceContext(const boost::filesystem::path& workingDirectory); + + /** + * Initialize context. + * + * Called before any compilation sessions are created. Use this method to + * initialize any mutable state. If this routine returns an error, the service + * will terminate. + * + * @return A status. + */ + [[nodiscard]] virtual grpc::Status init(); + + /** + * Uninitialize context. + * + * Called after all compilation sessions have ended, before a service + * terminates. Use this method to perform tidying up. This method is always + * called, even if init() fails. If this routine returns an error, the service + * will terminate with a nonzero error code. + * + * @return A status. + */ + [[nodiscard]] virtual grpc::Status shutdown(); + + /** + * Get the working directory. + * + * The working directory is a local filesystem directory that compilation + * sessions can use to store temporary files such as build artifacts. The + * directory is guaranteed to exist. + * + * \note When possible, an in-memory filesystem will be used for the working + * directory. This means that the working directory is not suitable for + * very large files or executables, as some systems prevent execution of + * in-memory files. + * + * \note A single working directory is shared by all of the compilation + * sessions of a service. Do not assume that you have exclusive access. + * + * @return A path. + */ + inline const boost::filesystem::path& workingDirectory() const { return workingDirectory_; }; + + private: + const boost::filesystem::path workingDirectory_; +}; + +} // namespace compiler_gym From 3d8edf693053a775475fc4ebeb03a2e0b49b125e Mon Sep 17 00:00:00 2001 From: Chris Cummins Date: Sun, 6 Mar 2022 19:14:41 +0000 Subject: [PATCH 2/3] WIP: New service --- compiler_gym/envs/llvm/service/BUILD | 15 +++- .../envs/llvm/service/BenchmarkFactory.h | 44 ++++------- .../envs/llvm/service/ComputeObservation.cc | 2 +- .../envs/llvm/service/LlvmServiceContext.cc | 73 +++++++++++++++++++ .../envs/llvm/service/LlvmServiceContext.h | 29 ++++++++ compiler_gym/envs/llvm/service/LlvmSession.cc | 20 +++-- compiler_gym/envs/llvm/service/LlvmSession.h | 2 +- compiler_gym/envs/llvm/service/RunService.cc | 61 +--------------- .../llvm/service/StripOptNoneAttribute.cc | 2 +- compiler_gym/service/BUILD | 1 + compiler_gym/service/CompilationSession.cc | 4 +- compiler_gym/service/CompilationSession.h | 11 ++- .../service/CompilerGymServiceContext.cc | 12 ++- .../service/CompilerGymServiceContext.h | 2 +- .../service/compiler_gym_service_context.py | 60 +++++++++++++++ compiler_gym/service/runtime/BUILD | 1 + .../service/runtime/CompilerGymService.h | 6 +- .../service/runtime/CompilerGymServiceImpl.h | 36 ++++----- .../CreateAndRunCompilerGymServiceImpl.h | 25 ++++++- compiler_gym/service/runtime/Runtime.h | 7 +- 20 files changed, 278 insertions(+), 135 deletions(-) create mode 100644 compiler_gym/envs/llvm/service/LlvmServiceContext.cc create mode 100644 compiler_gym/envs/llvm/service/LlvmServiceContext.h create mode 100644 compiler_gym/service/compiler_gym_service_context.py diff --git a/compiler_gym/envs/llvm/service/BUILD b/compiler_gym/envs/llvm/service/BUILD index 04a9ac89c..18d29bfc5 100644 --- a/compiler_gym/envs/llvm/service/BUILD +++ b/compiler_gym/envs/llvm/service/BUILD @@ -63,7 +63,7 @@ cc_binary( name = "compiler_gym-llvm-service-prelinked", srcs = ["RunService.cc"], deps = [ - ":BenchmarkFactory", + ":LlvmServiceContext", ":LlvmSession", "//compiler_gym/service/runtime:cc_runtime", ], @@ -207,6 +207,18 @@ cc_library( ], ) +cc_library( + name = "LlvmServiceContext", + srcs = ["LlvmServiceContext.cc"], + hdrs = ["LlvmServiceContext.h"], + deps = [ + ":BenchmarkFactory", + "//compiler_gym/service:CompilerGymServiceContext", + "//compiler_gym/util:GrpcStatusMacros", + "@llvm//10.0.0", + ], +) + cc_library( name = "LlvmSession", srcs = ["LlvmSession.cc"], @@ -224,6 +236,7 @@ cc_library( ":Benchmark", ":BenchmarkFactory", ":Cost", + ":LlvmServiceContext", ":Observation", ":ObservationSpaces", "//compiler_gym/service:CompilationSession", diff --git a/compiler_gym/envs/llvm/service/BenchmarkFactory.h b/compiler_gym/envs/llvm/service/BenchmarkFactory.h index 47b817b46..106ef901b 100644 --- a/compiler_gym/envs/llvm/service/BenchmarkFactory.h +++ b/compiler_gym/envs/llvm/service/BenchmarkFactory.h @@ -46,21 +46,22 @@ constexpr size_t kMaxLoadedBenchmarksCount = 128; class BenchmarkFactory { public: /** - * Return the global benchmark factory singleton. + * Construct a benchmark factory. * - * @param workingDirectory The working directory. - * @param rand An optional random number generator. This is used for cache - * evictions. - * @param maxLoadedBenchmarksCount The maximum number of benchmarks to cache. - * @return The benchmark factory singleton instance. + * @param workingDirectory A filesystem directory to use for storing temporary + * files. + * @param rand is a random seed used to control the selection of random + * benchmarks. + * @param maxLoadedBenchmarksCount is the maximum combined size of the bitcodes + * that may be cached in memory. Once this size is reached, benchmarks are + * offloaded so that they must be re-read from disk. */ - static BenchmarkFactory& getSingleton( - const boost::filesystem::path& workingDirectory, - std::optional rand = std::nullopt, - size_t maxLoadedBenchmarksCount = kMaxLoadedBenchmarksCount) { - static BenchmarkFactory instance(workingDirectory, rand, maxLoadedBenchmarksCount); - return instance; - } + BenchmarkFactory(const boost::filesystem::path& workingDirectory, + std::optional rand = std::nullopt, + size_t maxLoadedBenchmarksCount = kMaxLoadedBenchmarksCount); + + BenchmarkFactory(const BenchmarkFactory&) = delete; + BenchmarkFactory& operator=(const BenchmarkFactory&) = delete; ~BenchmarkFactory(); @@ -86,23 +87,6 @@ class BenchmarkFactory { const std::string& uri, const boost::filesystem::path& path, std::optional dynamicConfig = std::nullopt); - /** - * Construct a benchmark factory. - * - * @param workingDirectory A filesystem directory to use for storing temporary - * files. - * @param rand is a random seed used to control the selection of random - * benchmarks. - * @param maxLoadedBenchmarksCount is the maximum combined size of the bitcodes - * that may be cached in memory. Once this size is reached, benchmarks are - * offloaded so that they must be re-read from disk. - */ - BenchmarkFactory(const boost::filesystem::path& workingDirectory, - std::optional rand, size_t maxLoadedBenchmarksCount); - - BenchmarkFactory(const BenchmarkFactory&) = delete; - BenchmarkFactory& operator=(const BenchmarkFactory&) = delete; - /** * A mapping from URI to benchmarks which have been loaded into memory. */ diff --git a/compiler_gym/envs/llvm/service/ComputeObservation.cc b/compiler_gym/envs/llvm/service/ComputeObservation.cc index c0376d695..8e1237f3a 100644 --- a/compiler_gym/envs/llvm/service/ComputeObservation.cc +++ b/compiler_gym/envs/llvm/service/ComputeObservation.cc @@ -39,7 +39,7 @@ int main(int argc, char** argv) { benchmarkMessage.set_uri("user"); benchmarkMessage.mutable_program()->set_uri(fmt::format("file:///{}", argv[2])); - auto& benchmarkFactory = BenchmarkFactory::getSingleton(workingDirectory); + BenchmarkFactory benchmarkFactory{workingDirectory}; std::unique_ptr<::llvm_service::Benchmark> benchmark; { const auto status = benchmarkFactory.getBenchmark(benchmarkMessage, &benchmark); diff --git a/compiler_gym/envs/llvm/service/LlvmServiceContext.cc b/compiler_gym/envs/llvm/service/LlvmServiceContext.cc new file mode 100644 index 000000000..24a49ce71 --- /dev/null +++ b/compiler_gym/envs/llvm/service/LlvmServiceContext.cc @@ -0,0 +1,73 @@ +// Copyright (c) Facebook, Inc. and its affiliates. +// +// 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/LlvmServiceContext.h" + +#include "compiler_gym/util/GrpcStatusMacros.h" +#include "llvm/InitializePasses.h" +#include "llvm/Pass.h" +#include "llvm/Support/TargetSelect.h" + +using grpc::Status; + +namespace { + +void initLlvm() { + llvm::InitializeNativeTarget(); + + // Initialize passes. + llvm::PassRegistry& Registry = *llvm::PassRegistry::getPassRegistry(); + llvm::initializeCore(Registry); + llvm::initializeCoroutines(Registry); + llvm::initializeScalarOpts(Registry); + llvm::initializeObjCARCOpts(Registry); + llvm::initializeVectorization(Registry); + llvm::initializeIPO(Registry); + llvm::initializeAnalysis(Registry); + llvm::initializeTransformUtils(Registry); + llvm::initializeInstCombine(Registry); + llvm::initializeAggressiveInstCombine(Registry); + llvm::initializeInstrumentation(Registry); + llvm::initializeTarget(Registry); + llvm::initializeExpandMemCmpPassPass(Registry); + llvm::initializeScalarizeMaskedMemIntrinPass(Registry); + llvm::initializeCodeGenPreparePass(Registry); + llvm::initializeAtomicExpandPass(Registry); + llvm::initializeRewriteSymbolsLegacyPassPass(Registry); + llvm::initializeWinEHPreparePass(Registry); + llvm::initializeDwarfEHPreparePass(Registry); + llvm::initializeSafeStackLegacyPassPass(Registry); + llvm::initializeSjLjEHPreparePass(Registry); + llvm::initializePreISelIntrinsicLoweringLegacyPassPass(Registry); + llvm::initializeGlobalMergePass(Registry); + llvm::initializeIndirectBrExpandPassPass(Registry); + llvm::initializeInterleavedAccessPass(Registry); + llvm::initializeEntryExitInstrumenterPass(Registry); + llvm::initializePostInlineEntryExitInstrumenterPass(Registry); + llvm::initializeUnreachableBlockElimLegacyPassPass(Registry); + llvm::initializeExpandReductionsPass(Registry); + llvm::initializeWasmEHPreparePass(Registry); + llvm::initializeWriteBitcodePassPass(Registry); +} + +} // anonymous namespace + +namespace compiler_gym::llvm_service { + +LlvmServiceContext::LlvmServiceContext(const boost::filesystem::path& workingDirectory) + : CompilerGymServiceContext(workingDirectory), benchmarkFactory_(workingDirectory) {} + +Status LlvmServiceContext::init() { + RETURN_IF_ERROR(CompilerGymServiceContext::init()); + initLlvm(); + return Status::OK; +} + +Status LlvmServiceContext::shutdown() { + Status status = CompilerGymServiceContext::shutdown(); + benchmarkFactory().close(); + return status; +} + +} // namespace compiler_gym::llvm_service diff --git a/compiler_gym/envs/llvm/service/LlvmServiceContext.h b/compiler_gym/envs/llvm/service/LlvmServiceContext.h new file mode 100644 index 000000000..2f6ed0206 --- /dev/null +++ b/compiler_gym/envs/llvm/service/LlvmServiceContext.h @@ -0,0 +1,29 @@ +// Copyright (c) Facebook, Inc. and its affiliates. +// +// This source code is licensed under the MIT license found in the +// LICENSE file in the root directory of this source tree. +#pragma once + +#include + +#include "boost/filesystem.hpp" +#include "compiler_gym/envs/llvm/service/BenchmarkFactory.h" +#include "compiler_gym/service/CompilerGymServiceContext.h" + +namespace compiler_gym::llvm_service { + +class LlvmServiceContext final : public CompilerGymServiceContext { + public: + LlvmServiceContext(const boost::filesystem::path& workingDirectory); + + [[nodiscard]] virtual grpc::Status init() final override; + + [[nodiscard]] virtual grpc::Status shutdown() final override; + + BenchmarkFactory& benchmarkFactory() { return benchmarkFactory_; } + + private: + BenchmarkFactory benchmarkFactory_; +}; + +} // namespace compiler_gym::llvm_service diff --git a/compiler_gym/envs/llvm/service/LlvmSession.cc b/compiler_gym/envs/llvm/service/LlvmSession.cc index edd4f4df1..c593dcc75 100644 --- a/compiler_gym/envs/llvm/service/LlvmSession.cc +++ b/compiler_gym/envs/llvm/service/LlvmSession.cc @@ -20,6 +20,7 @@ #include "compiler_gym/envs/llvm/service/Benchmark.h" #include "compiler_gym/envs/llvm/service/BenchmarkFactory.h" #include "compiler_gym/envs/llvm/service/Cost.h" +#include "compiler_gym/envs/llvm/service/LlvmServiceContext.h" #include "compiler_gym/envs/llvm/service/Observation.h" #include "compiler_gym/envs/llvm/service/ObservationSpaces.h" #include "compiler_gym/envs/llvm/service/passes/10.0.0/ActionHeaders.h" @@ -76,18 +77,19 @@ std::vector LlvmSession::getObservationSpaces() const { return getLlvmObservationSpaceList(); } -LlvmSession::LlvmSession(const boost::filesystem::path& workingDirectory) - : CompilationSession(workingDirectory), +LlvmSession::LlvmSession(CompilerGymServiceContext* const context) + : CompilationSession(context), observationSpaceNames_(util::createPascalCaseToEnumLookupTable()) { + // TODO: Move CPUInfo initialize to context setup! cpuinfo_initialize(); } Status LlvmSession::init(const ActionSpace& actionSpace, const BenchmarkProto& benchmark) { - BenchmarkFactory& benchmarkFactory = BenchmarkFactory::getSingleton(workingDirectory()); + LlvmServiceContext* const ctx = static_cast(context()); // Get the benchmark or return an error. std::unique_ptr llvmBenchmark; - RETURN_IF_ERROR(benchmarkFactory.getBenchmark(benchmark, &llvmBenchmark)); + RETURN_IF_ERROR(ctx->benchmarkFactory().getBenchmark(benchmark, &llvmBenchmark)); // Verify the benchmark now to catch errors early. RETURN_IF_ERROR(llvmBenchmark->verify_module()); @@ -101,7 +103,8 @@ Status LlvmSession::init(const ActionSpace& actionSpace, const BenchmarkProto& b Status LlvmSession::init(CompilationSession* other) { // TODO: Static cast? auto llvmOther = static_cast(other); - return init(llvmOther->actionSpace(), llvmOther->benchmark().clone(workingDirectory())); + return init(llvmOther->actionSpace(), + llvmOther->benchmark().clone(context()->workingDirectory())); } Status LlvmSession::init(const LlvmActionSpace& actionSpace, std::unique_ptr benchmark) { @@ -156,7 +159,8 @@ Status LlvmSession::computeObservation(const ObservationSpace& observationSpace, } const LlvmObservationSpace observationSpaceEnum = it->second; - return setObservation(observationSpaceEnum, workingDirectory(), benchmark(), observation); + return setObservation(observationSpaceEnum, context()->workingDirectory(), benchmark(), + observation); } Status LlvmSession::handleSessionParameter(const std::string& key, const std::string& value, @@ -256,8 +260,8 @@ bool LlvmSession::runPass(llvm::FunctionPass* pass) { Status LlvmSession::runOptWithArgs(const std::vector& optArgs) { // Create temporary files for `opt` to read from and write to. - const auto before_path = fs::unique_path(workingDirectory() / "module-%%%%%%%%.bc"); - const auto after_path = fs::unique_path(workingDirectory() / "module-%%%%%%%%.bc"); + const auto before_path = fs::unique_path(context()->workingDirectory() / "module-%%%%%%%%.bc"); + const auto after_path = fs::unique_path(context()->workingDirectory() / "module-%%%%%%%%.bc"); RETURN_IF_ERROR(writeBitcodeFile(benchmark().module(), before_path)); // Build a command line invocation: `opt input.bc -o output.bc `. diff --git a/compiler_gym/envs/llvm/service/LlvmSession.h b/compiler_gym/envs/llvm/service/LlvmSession.h index e6febd462..37dd04e23 100644 --- a/compiler_gym/envs/llvm/service/LlvmSession.h +++ b/compiler_gym/envs/llvm/service/LlvmSession.h @@ -37,7 +37,7 @@ namespace compiler_gym::llvm_service { */ class LlvmSession final : public CompilationSession { public: - LlvmSession(const boost::filesystem::path& workingDirectory); + LlvmSession(CompilerGymServiceContext* const context); std::string getCompilerVersion() const final override; diff --git a/compiler_gym/envs/llvm/service/RunService.cc b/compiler_gym/envs/llvm/service/RunService.cc index 7c5687aaf..1cd3d49e3 100644 --- a/compiler_gym/envs/llvm/service/RunService.cc +++ b/compiler_gym/envs/llvm/service/RunService.cc @@ -2,72 +2,15 @@ // // 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/LlvmServiceContext.h" #include "compiler_gym/envs/llvm/service/LlvmSession.h" #include "compiler_gym/service/runtime/Runtime.h" -#include "llvm/InitializePasses.h" -#include "llvm/Support/TargetSelect.h" const char* usage = R"(LLVM CompilerGym service)"; using namespace compiler_gym::runtime; using namespace compiler_gym::llvm_service; -namespace { - -void initLlvm() { - llvm::InitializeNativeTarget(); - - // Initialize passes. - llvm::PassRegistry& Registry = *llvm::PassRegistry::getPassRegistry(); - llvm::initializeCore(Registry); - llvm::initializeCoroutines(Registry); - llvm::initializeScalarOpts(Registry); - llvm::initializeObjCARCOpts(Registry); - llvm::initializeVectorization(Registry); - llvm::initializeIPO(Registry); - llvm::initializeAnalysis(Registry); - llvm::initializeTransformUtils(Registry); - llvm::initializeInstCombine(Registry); - llvm::initializeAggressiveInstCombine(Registry); - llvm::initializeInstrumentation(Registry); - llvm::initializeTarget(Registry); - llvm::initializeExpandMemCmpPassPass(Registry); - llvm::initializeScalarizeMaskedMemIntrinPass(Registry); - llvm::initializeCodeGenPreparePass(Registry); - llvm::initializeAtomicExpandPass(Registry); - llvm::initializeRewriteSymbolsLegacyPassPass(Registry); - llvm::initializeWinEHPreparePass(Registry); - llvm::initializeDwarfEHPreparePass(Registry); - llvm::initializeSafeStackLegacyPassPass(Registry); - llvm::initializeSjLjEHPreparePass(Registry); - llvm::initializePreISelIntrinsicLoweringLegacyPassPass(Registry); - llvm::initializeGlobalMergePass(Registry); - llvm::initializeIndirectBrExpandPassPass(Registry); - llvm::initializeInterleavedAccessPass(Registry); - llvm::initializeEntryExitInstrumenterPass(Registry); - llvm::initializePostInlineEntryExitInstrumenterPass(Registry); - llvm::initializeUnreachableBlockElimLegacyPassPass(Registry); - llvm::initializeExpandReductionsPass(Registry); - llvm::initializeWasmEHPreparePass(Registry); - llvm::initializeWriteBitcodePassPass(Registry); -} - -} // anonymous namespace - int main(int argc, char** argv) { - initLlvm(); - const auto ret = createAndRunCompilerGymService(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; + return createAndRunCompilerGymService(argc, argv, usage); } diff --git a/compiler_gym/envs/llvm/service/StripOptNoneAttribute.cc b/compiler_gym/envs/llvm/service/StripOptNoneAttribute.cc index faed9a88e..1e4929b50 100644 --- a/compiler_gym/envs/llvm/service/StripOptNoneAttribute.cc +++ b/compiler_gym/envs/llvm/service/StripOptNoneAttribute.cc @@ -60,7 +60,7 @@ int main(int argc, char** argv) { google::InitGoogleLogging(argv[0]); const fs::path workingDirectory{"."}; - auto& benchmarkFactory = BenchmarkFactory::getSingleton(workingDirectory); + BenchmarkFactory benchmarkFactory(workingDirectory); for (int i = 1; i < argc; ++i) { stripOptNoneAttributesOrDie(argv[i], benchmarkFactory); diff --git a/compiler_gym/service/BUILD b/compiler_gym/service/BUILD index f527a82e6..307e3e701 100644 --- a/compiler_gym/service/BUILD +++ b/compiler_gym/service/BUILD @@ -31,6 +31,7 @@ cc_library( hdrs = ["CompilationSession.h"], visibility = ["//visibility:public"], deps = [ + ":CompilerGymServiceContext", "//compiler_gym/service/proto:compiler_gym_service_cc", "@boost//:filesystem", "@com_github_grpc_grpc//:grpc++", diff --git a/compiler_gym/service/CompilationSession.cc b/compiler_gym/service/CompilationSession.cc index d9ffbdb0b..bf31895e2 100644 --- a/compiler_gym/service/CompilationSession.cc +++ b/compiler_gym/service/CompilationSession.cc @@ -25,7 +25,7 @@ Status CompilationSession::handleSessionParameter(const std::string& key, const return Status::OK; } -CompilationSession::CompilationSession(const boost::filesystem::path& workingDirectory) - : workingDirectory_(workingDirectory) {} +CompilationSession::CompilationSession(CompilerGymServiceContext* const context) + : context_(context) {} } // namespace compiler_gym diff --git a/compiler_gym/service/CompilationSession.h b/compiler_gym/service/CompilationSession.h index d15d87ece..a9d3f2901 100644 --- a/compiler_gym/service/CompilationSession.h +++ b/compiler_gym/service/CompilationSession.h @@ -10,6 +10,7 @@ #include #include "boost/filesystem.hpp" +#include "compiler_gym/service/CompilerGymServiceContext.h" #include "compiler_gym/service/proto/compiler_gym_service.pb.h" namespace compiler_gym { @@ -115,7 +116,7 @@ class CompilationSession { [[nodiscard]] virtual grpc::Status endOfStep(bool actionHadNoEffect, bool& endOfEpisode, std::optional& newActionSpace); - CompilationSession(const boost::filesystem::path& workingDirectory); + CompilationSession(CompilerGymServiceContext* const context); virtual ~CompilationSession() = default; @@ -155,12 +156,16 @@ class CompilationSession { * CompilationSession instances. Do not assume that you have exclusive * access. * + * \deprecated Please use context()->workingDirectory() instead. + * * @return A path. */ - inline const boost::filesystem::path& workingDirectory() { return workingDirectory_; } + inline const boost::filesystem::path& workingDirectory() { return context()->workingDirectory(); } + + inline CompilerGymServiceContext* const context() const { return context_; } private: - const boost::filesystem::path workingDirectory_; + CompilerGymServiceContext* const context_; }; } // namespace compiler_gym diff --git a/compiler_gym/service/CompilerGymServiceContext.cc b/compiler_gym/service/CompilerGymServiceContext.cc index 5d928f482..827bc8acf 100644 --- a/compiler_gym/service/CompilerGymServiceContext.cc +++ b/compiler_gym/service/CompilerGymServiceContext.cc @@ -2,8 +2,6 @@ // // This source code is licensed under the MIT license found in the // LICENSE file in the root directory of this source tree. -#pragma once - #include "compiler_gym/service/CompilerGymServiceContext.h" using grpc::Status; @@ -14,8 +12,14 @@ CompilerGymServiceContext::CompilerGymServiceContext( const boost::filesystem::path& workingDirectory) : workingDirectory_(workingDirectory) {} -virtual Status CompilerGymServiceContext::init() { return Status::OK; } +Status CompilerGymServiceContext::init() { + VLOG(2) << "Initializing compiler service context"; + return Status::OK; +} -virtual Status CompilerGymServiceContext::shutdown() { return Status::OK; } +Status CompilerGymServiceContext::shutdown() { + VLOG(2) << "Closing compiler service context"; + return Status::OK; +} } // namespace compiler_gym diff --git a/compiler_gym/service/CompilerGymServiceContext.h b/compiler_gym/service/CompilerGymServiceContext.h index 67008a866..a54363194 100644 --- a/compiler_gym/service/CompilerGymServiceContext.h +++ b/compiler_gym/service/CompilerGymServiceContext.h @@ -28,7 +28,7 @@ namespace compiler_gym { * * using namespace compiler_gym; * - * class MyServiceContext final : public ServiceContext { ... } + * class MyServiceContext final : public CompilerGymServiceContext { ... } * * class MyCompilationSession final : public CompilationSession { ... } * diff --git a/compiler_gym/service/compiler_gym_service_context.py b/compiler_gym/service/compiler_gym_service_context.py new file mode 100644 index 000000000..ef04b454f --- /dev/null +++ b/compiler_gym/service/compiler_gym_service_context.py @@ -0,0 +1,60 @@ +# Copyright (c) Facebook, Inc. and its affiliates. +# +# This source code is licensed under the MIT license found in the +# LICENSE file in the root directory of this source tree. +import logging +from pathlib import Path + +logger = logging.getLogger(__name__) + + +class CompilerGymServiceContext: + """ + Execution context of a compiler gym service. + + This class encapsulates mutable state that is shared between all compilation + sessions. An instance of this class is passed to every new + CompilationSession. + + You may subclass CompilerGymServiceContext to add additional mutable state. + The subclass . + + .. code-block:: python + + from compiler_gym.service import CompilationSession + from compiler_gym.service import CompilerGymServiceContext + from compiler_gym.service import runtime + + class MyServiceContext(CompilerGymServiceContext): + ... + + class MyCompilationSession(CompilationSession): + ... + + if __name__ == "__main__": + runtime.create_and_run_compiler_service( + MyCompilationSession, MyServiceContext, + ) + """ + + def __init__(self, working_directory: Path) -> None: + """ + Initialize context. + + Called before any compilation sessions are created. Use this method to + initialize any mutable state. If this routine returns an error, the + service will terminate. + """ + logger.debug("Initializing compiler service context") + self.working_directory: Path = working_directory + + def shutdown(self) -> None: + """ + Uninitialize context. + + Called after all compilation sessions have ended, before a service + terminates. Use this method to perform tidying up. This method is always + called, even if init() fails. If this routine returns an error, the + service will terminate with a nonzero error code. + """ + logger.debug("Closing compiler service context") diff --git a/compiler_gym/service/runtime/BUILD b/compiler_gym/service/runtime/BUILD index 535516fd3..187c500a9 100644 --- a/compiler_gym/service/runtime/BUILD +++ b/compiler_gym/service/runtime/BUILD @@ -105,6 +105,7 @@ cc_library( hdrs = ["CreateAndRunCompilerGymServiceImpl.h"], deps = [ ":CompilerGymService", + "//compiler_gym/service:CompilerGymServiceContext", "//compiler_gym/util:GrpcStatusMacros", "@boost//:filesystem", "@com_github_grpc_grpc//:grpc++", diff --git a/compiler_gym/service/runtime/CompilerGymService.h b/compiler_gym/service/runtime/CompilerGymService.h index fc9b39c88..71196daf6 100644 --- a/compiler_gym/service/runtime/CompilerGymService.h +++ b/compiler_gym/service/runtime/CompilerGymService.h @@ -28,7 +28,7 @@ namespace compiler_gym::runtime { template class CompilerGymService final : public compiler_gym::CompilerGymService::Service { public: - CompilerGymService(const boost::filesystem::path& workingDirectory, + CompilerGymService(CompilerGymServiceContext* const context, std::unique_ptr benchmarks = nullptr); // RPC endpoints. @@ -77,7 +77,7 @@ class CompilerGymService final : public compiler_gym::CompilerGymService::Servic [[nodiscard]] grpc::Status observation_space(const CompilationSession* session, int index, const ObservationSpace** observationSpace) const; - inline const boost::filesystem::path& workingDirectory() const { return workingDirectory_; } + inline CompilerGymServiceContext* const context() { return context_; } // Add the given session and return its ID. uint64_t addSession(std::unique_ptr session); @@ -88,7 +88,7 @@ class CompilerGymService final : public compiler_gym::CompilerGymService::Servic std::optional& reply); private: - const boost::filesystem::path workingDirectory_; + CompilerGymServiceContext* const context_; const std::vector actionSpaces_; const std::vector observationSpaces_; diff --git a/compiler_gym/service/runtime/CompilerGymServiceImpl.h b/compiler_gym/service/runtime/CompilerGymServiceImpl.h index 693e44083..232cc3c5f 100644 --- a/compiler_gym/service/runtime/CompilerGymServiceImpl.h +++ b/compiler_gym/service/runtime/CompilerGymServiceImpl.h @@ -17,27 +17,26 @@ namespace compiler_gym::runtime { template CompilerGymService::CompilerGymService( - const boost::filesystem::path& workingDirectory, std::unique_ptr benchmarks) - : workingDirectory_(workingDirectory), - actionSpaces_(CompilationSessionType(workingDirectory).getActionSpaces()), - observationSpaces_(CompilationSessionType(workingDirectory).getObservationSpaces()), + CompilerGymServiceContext* const context, std::unique_ptr benchmarks) + : context_(context), + actionSpaces_(CompilationSessionType(context).getActionSpaces()), + observationSpaces_(CompilationSessionType(context).getObservationSpaces()), benchmarks_(benchmarks ? std::move(benchmarks) : std::make_unique()), nextSessionId_(0) {} template grpc::Status CompilerGymService::GetVersion( - grpc::ServerContext* context, const GetVersionRequest* request, GetVersionReply* reply) { + grpc::ServerContext* serverContext, const GetVersionRequest* request, GetVersionReply* reply) { VLOG(2) << "GetVersion()"; reply->set_service_version(COMPILER_GYM_VERSION); - CompilationSessionType environment(workingDirectory()); + CompilationSessionType environment(context()); reply->set_compiler_version(environment.getCompilerVersion()); return grpc::Status::OK; } template -grpc::Status CompilerGymService::GetSpaces(grpc::ServerContext* context, - const GetSpacesRequest* request, - GetSpacesReply* reply) { +grpc::Status CompilerGymService::GetSpaces( + grpc::ServerContext* serverContext, const GetSpacesRequest* request, GetSpacesReply* reply) { VLOG(2) << "GetSpaces()"; for (const auto& actionSpace : actionSpaces_) { *reply->add_action_space_list() = actionSpace; @@ -50,7 +49,8 @@ grpc::Status CompilerGymService::GetSpaces(grpc::ServerC template grpc::Status CompilerGymService::StartSession( - grpc::ServerContext* context, const StartSessionRequest* request, StartSessionReply* reply) { + grpc::ServerContext* serverContext, const StartSessionRequest* request, + StartSessionReply* reply) { if (!request->benchmark().uri().size()) { return grpc::Status(grpc::StatusCode::INVALID_ARGUMENT, "No benchmark URI set for StartSession()"); @@ -72,7 +72,7 @@ grpc::Status CompilerGymService::StartSession( } // Construct the new session. - auto environment = std::make_unique(workingDirectory()); + auto environment = std::make_unique(context()); // Resolve the action space. const ActionSpace* actionSpace; @@ -96,7 +96,8 @@ grpc::Status CompilerGymService::StartSession( template grpc::Status CompilerGymService::ForkSession( - grpc::ServerContext* context, const ForkSessionRequest* request, ForkSessionReply* reply) { + grpc::ServerContext* serverContext, const ForkSessionRequest* request, + ForkSessionReply* reply) { const std::lock_guard lock(sessionsMutex_); CompilationSession* baseSession; @@ -104,7 +105,7 @@ grpc::Status CompilerGymService::ForkSession( VLOG(1) << "ForkSession(" << request->session_id() << "), [" << nextSessionId_ << "]"; // Construct the new session. - auto forked = std::make_unique(workingDirectory()); + auto forked = std::make_unique(context()); // Initialize from the base environment. RETURN_IF_ERROR(forked->init(baseSession)); @@ -116,7 +117,7 @@ grpc::Status CompilerGymService::ForkSession( template grpc::Status CompilerGymService::EndSession( - grpc::ServerContext* context, const EndSessionRequest* request, EndSessionReply* reply) { + grpc::ServerContext* serverContext, const EndSessionRequest* request, EndSessionReply* reply) { VLOG(1) << "EndSession(id=" << request->session_id() << "), " << sessionCount() - 1 << " sessions remaining"; @@ -135,7 +136,7 @@ grpc::Status CompilerGymService::EndSession( } template -grpc::Status CompilerGymService::Step(grpc::ServerContext* context, +grpc::Status CompilerGymService::Step(grpc::ServerContext* serverContext, const StepRequest* request, StepReply* reply) { CompilationSession* environment; @@ -184,7 +185,8 @@ grpc::Status CompilerGymService::Step(grpc::ServerContex template grpc::Status CompilerGymService::AddBenchmark( - grpc::ServerContext* context, const AddBenchmarkRequest* request, AddBenchmarkReply* reply) { + grpc::ServerContext* serverContext, const AddBenchmarkRequest* request, + AddBenchmarkReply* reply) { // We need to grab the sessions lock here to ensure thread safe access to the // benchmarks cache. const std::lock_guard lock(sessionsMutex_); @@ -199,7 +201,7 @@ grpc::Status CompilerGymService::AddBenchmark( template grpc::Status CompilerGymService::SendSessionParameter( - grpc::ServerContext* context, const SendSessionParameterRequest* request, + grpc::ServerContext* serverContext, const SendSessionParameterRequest* request, SendSessionParameterReply* reply) { CompilationSession* environment; RETURN_IF_ERROR(session(request->session_id(), &environment)); diff --git a/compiler_gym/service/runtime/CreateAndRunCompilerGymServiceImpl.h b/compiler_gym/service/runtime/CreateAndRunCompilerGymServiceImpl.h index e22d6b85e..ae21d0a0c 100644 --- a/compiler_gym/service/runtime/CreateAndRunCompilerGymServiceImpl.h +++ b/compiler_gym/service/runtime/CreateAndRunCompilerGymServiceImpl.h @@ -19,6 +19,7 @@ #include #include "boost/filesystem.hpp" +#include "compiler_gym/service/CompilerGymServiceContext.h" #include "compiler_gym/service/proto/compiler_gym_service.pb.h" #include "compiler_gym/service/runtime/CompilerGymService.h" @@ -50,7 +51,8 @@ void setGrpcChannelOptions(grpc::ServerBuilder& builder); // int main(int argc, char** argv) { // createAndRunCompilerGymServiceImpl(argc, argv, "usage string"); // } -template +template [[nodiscard]] int createAndRunCompilerGymServiceImpl(int argc, char** argv, const char* usage) { // Register a signal handler for SIGTERM that will set the shutdown_signal // future value. @@ -80,7 +82,19 @@ template google::InitGoogleLogging(argv[0]); - CompilerGymService service{workingDirectory}; + // Create and initialize the service context. + CompilerGymServiceContextType context{workingDirectory}; + Status status = context.init(); + if (!status.ok()) { + LOG(ERROR) << "Error during initialization of service context: " << status.error_message(); + status = context.shutdown(); + if (!status.ok()) { + LOG(ERROR) << "Error during shutdown of service context: " << status.error_message(); + } + return 1; + } + + CompilerGymService service{&context}; grpc::ServerBuilder builder; builder.RegisterService(&service); @@ -129,8 +143,15 @@ template VLOG(2) << "Shutting down the RPC service"; server->Shutdown(); serverThread.join(); + VLOG(2) << "Service closed"; + status = context.shutdown(); + if (!status.ok()) { + LOG(ERROR) << "Error during shutdown of service context: " << status.error_message(); + return 1; + } + if (service.sessionCount()) { LOG(ERROR) << "ERROR: Killing a service with " << service.sessionCount() << (service.sessionCount() > 1 ? " active sessions!" : " active session!") diff --git a/compiler_gym/service/runtime/Runtime.h b/compiler_gym/service/runtime/Runtime.h index 42d162eb2..ded41b7d8 100644 --- a/compiler_gym/service/runtime/Runtime.h +++ b/compiler_gym/service/runtime/Runtime.h @@ -4,6 +4,7 @@ // LICENSE file in the root directory of this source tree. #pragma once +#include "compiler_gym/service/CompilerGymServiceContext.h" #include "compiler_gym/service/runtime/CompilerGymService.h" #include "compiler_gym/service/runtime/CreateAndRunCompilerGymServiceImpl.h" @@ -31,9 +32,11 @@ namespace compiler_gym::runtime { * * @return An integer return code. */ -template +template [[nodiscard]] int createAndRunCompilerGymService(int argc, char** argv, const char* usage) { - return createAndRunCompilerGymServiceImpl(argc, argv, usage); + return createAndRunCompilerGymServiceImpl( + argc, argv, usage); } } // namespace compiler_gym::runtime From 58fd4e375ee004071df77eade59af53445375a04 Mon Sep 17 00:00:00 2001 From: Chris Cummins Date: Sun, 6 Mar 2022 19:14:41 +0000 Subject: [PATCH 3/3] WIP --- .../service/CompilerGymServiceContext.h | 5 +- .../service/compiler_gym_service_context.py | 19 +++- .../create_and_run_compiler_gym_service.py | 102 ++++++++++-------- 3 files changed, 72 insertions(+), 54 deletions(-) diff --git a/compiler_gym/service/CompilerGymServiceContext.h b/compiler_gym/service/CompilerGymServiceContext.h index a54363194..06c3980ac 100644 --- a/compiler_gym/service/CompilerGymServiceContext.h +++ b/compiler_gym/service/CompilerGymServiceContext.h @@ -17,8 +17,9 @@ namespace compiler_gym { * sessions. An instance of this class is passed to every new * CompilationSession. * - * You may subclass CompilerGymServiceContext to add additional mutable state. - * The subclass . + * You may subclass CompilerGymServiceContext to add additional mutable state, + * or startup and shutdown routines. When overriding methods, subclasses should + * call the parent class implementation first. * * \code{.cpp} * diff --git a/compiler_gym/service/compiler_gym_service_context.py b/compiler_gym/service/compiler_gym_service_context.py index ef04b454f..95137791c 100644 --- a/compiler_gym/service/compiler_gym_service_context.py +++ b/compiler_gym/service/compiler_gym_service_context.py @@ -16,14 +16,15 @@ class CompilerGymServiceContext: sessions. An instance of this class is passed to every new CompilationSession. - You may subclass CompilerGymServiceContext to add additional mutable state. - The subclass . + You may subclass CompilerGymServiceContext to add additional mutable state, + or startup and shutdown routines. When overriding methods, subclasses should + call the parent class implementation first. .. code-block:: python - from compiler_gym.service import CompilationSession - from compiler_gym.service import CompilerGymServiceContext - from compiler_gym.service import runtime + from compiler_gym.service import CompilationSession from + compiler_gym.service import CompilerGymServiceContext from + compiler_gym.service import runtime class MyServiceContext(CompilerGymServiceContext): ... @@ -58,3 +59,11 @@ def shutdown(self) -> None: service will terminate with a nonzero error code. """ logger.debug("Closing compiler service context") + + def __enter__(self) -> "CompilerGymServiceContext": + """Support 'with' syntax.""" + return self + + def __exit__(self, *args): + """Support 'with' syntax.""" + self.shutdown() diff --git a/compiler_gym/service/runtime/create_and_run_compiler_gym_service.py b/compiler_gym/service/runtime/create_and_run_compiler_gym_service.py index 3b53acde1..6415f3684 100644 --- a/compiler_gym/service/runtime/create_and_run_compiler_gym_service.py +++ b/compiler_gym/service/runtime/create_and_run_compiler_gym_service.py @@ -18,7 +18,7 @@ import grpc from absl import app, flags, logging -from compiler_gym.service import connection +from compiler_gym.service import CompilerGymServiceContext, connection from compiler_gym.service.compilation_session import CompilationSession from compiler_gym.service.proto import compiler_gym_service_pb2_grpc from compiler_gym.service.runtime.compiler_gym_service import CompilerGymService @@ -51,6 +51,9 @@ def _shutdown_handler(signal_number, stack_frame): # pragma: no cover def create_and_run_compiler_gym_service( compilation_session_type: Type[CompilationSession], + compiler_gym_service_context_type: Type[ + CompilerGymServiceContext + ] = CompilerGymServiceContext, ): # pragma: no cover """Create and run an RPC service for the given compilation session. @@ -92,52 +95,57 @@ def main(argv): logging.get_absl_handler().use_absl_log_file() logging.set_verbosity(dbg.get_logging_level()) - # Create the service. - server = grpc.server( - futures.ThreadPoolExecutor(max_workers=FLAGS.rpc_service_threads), - options=connection.GRPC_CHANNEL_OPTIONS, - ) - service = CompilerGymService( - working_directory=working_dir, - compilation_session_type=compilation_session_type, - ) - compiler_gym_service_pb2_grpc.add_CompilerGymServiceServicer_to_server( - service, server - ) - - address = f"0.0.0.0:{FLAGS.port}" if FLAGS.port else "0.0.0.0:0" - port = server.add_insecure_port(address) - - with atomic_file_write(working_dir / "port.txt", fileobj=True, mode="w") as f: - f.write(str(port)) - - with atomic_file_write(working_dir / "pid.txt", fileobj=True, mode="w") as f: - f.write(str(os.getpid())) - - logging.info( - "Service %s listening on %d, PID = %d", working_dir, port, os.getpid() - ) - - server.start() - - # Block on the RPC service in a separate thread. This enables the - # current thread to handle the shutdown routine. - server_thread = Thread(target=server.wait_for_termination) - server_thread.start() - - # Block until the shutdown signal is received. - shutdown_signal.wait() - logging.info("Shutting down the RPC service") - server.stop(60).wait() - server_thread.join() - logging.info("Service closed") - - if len(service.sessions): - print( - "ERROR: Killing a service with", - plural(len(service.session), "active session", "active sessions"), - file=sys.stderr, + with compiler_gym_service_context_type(working_dir) as context: + # Create the service. + server = grpc.server( + futures.ThreadPoolExecutor(max_workers=FLAGS.rpc_service_threads), + options=connection.GRPC_CHANNEL_OPTIONS, + ) + service = CompilerGymService( + compilation_session_type=compilation_session_type, + context=context, + ) + compiler_gym_service_pb2_grpc.add_CompilerGymServiceServicer_to_server( + service, server + ) + + address = f"0.0.0.0:{FLAGS.port}" if FLAGS.port else "0.0.0.0:0" + port = server.add_insecure_port(address) + + with atomic_file_write( + working_dir / "port.txt", fileobj=True, mode="w" + ) as f: + f.write(str(port)) + + with atomic_file_write( + working_dir / "pid.txt", fileobj=True, mode="w" + ) as f: + f.write(str(os.getpid())) + + logging.info( + "Service %s listening on %d, PID = %d", working_dir, port, os.getpid() ) - sys.exit(6) + + server.start() + + # Block on the RPC service in a separate thread. This enables the + # current thread to handle the shutdown routine. + server_thread = Thread(target=server.wait_for_termination) + server_thread.start() + + # Block until the shutdown signal is received. + shutdown_signal.wait() + logging.info("Shutting down the RPC service") + server.stop(60).wait() + server_thread.join() + logging.info("Service closed") + + if len(service.sessions): + print( + "ERROR: Killing a service with", + plural(len(service.session), "active session", "active sessions"), + file=sys.stderr, + ) + sys.exit(6) app.run(main)