[service] Enable intra-session mutable state

🚀 Feature

Motivation

The CompilationSession interface allows mutable state only for the duration of a single compilation session, but there are cases where we may want to have mutable state that outlives a single session and is shared between sessions over the lifetime of a service. For example, the LLM service implements a BenchmarkFactory class which is an in-memory cache of parsed bitcodes that is shared across all services.

Pitch

Add a new class that encapsulates all intra-session state of a compiler service in a "Context" object:

class CompilerGymServiceContext {
 public:
  // Called first. User setup routines go here.
  [[nodiscard]] virtual grpc::Status init();

  // Called last. User shutdown routines go here.
  [[nodiscard]] virtual grpc::Status shutdown();

  const boost::filesystem::path& workingDirectory() const;
}

Then change the CompilationSession interface so that it takes a reference to this context:

class CompilationSession {
 public:
  CompilationSession(CompilerGymServiceContext& ctx);
}

The helper function for creating and running a CompilerGymService then needs to be parametrized by this context type:

template <typename CompilationSessionType, typename CompilerGymServiceContext = DefaultCompilerGymServiceContext>
[[noreturn]] void createAndRunCompilerGymService(int argc, char** argv, const char* usage) {...}

Importantly, implementing this new context should be optional. We should provide a default implementation for users who don't need to use it. That way, the only breaking change will be if users decided to overload the default constructor for CompilationSession, since we are changing its signature.

Considerations

The user needs to make sure that any custom logic added to CompilerServiceContext is thread-safe, since a single context object will be shared by all sessions.

Alternatives

Don't provide a context object, but allow users to run their own shutdown code after createAndRunCompilerGymService() completes. The downside to this is that the user needs to remember to return the appropriate return code from createAndRunCompilerGymService().

facebookresearch / CompilerGym