WIP

Author: ChrisCummins

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}
  *

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()

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)