Allow for registration after GlobalInit (pytorch#15876)

Summary:
Pull Request resolved: pytorch#15876

Build changes made it so some .so libraries are now registered after GlobalInit is called. Although this shouldn't be common, it also shouldn't be explicitly excluded. These changes allow for late Caffe2 registration, but also warn in that case.

Reviewed By: kuttas

Differential Revision: D13608186

fbshipit-source-id: 0ca7bcd32516d374077db0c2548cf8c28ccdd5f6

Author: orionr

caffe2/core/init.h

@@ -15,16 +15,32 @@ class CAFFE2_API Caffe2InitializeRegistry {
   // multiple shared libraries loaded with RTLD_LOCAL
   static Caffe2InitializeRegistry* Registry();
 
-  void Register(InitFunction function, bool run_early,
-                const char* description) {
+  void
+  Register(InitFunction function, bool run_early, const char* description) {
     if (run_early) {
-      // Disallow late registration of early init functions
+      // Disallow registration after GlobalInit of early init functions
       CAFFE_ENFORCE(!early_init_functions_run_yet_);
       early_init_functions_.emplace_back(function, description);
     } else {
-      // Disallow late registration of init functions
-      CAFFE_ENFORCE(!init_functions_run_yet_);
-      init_functions_.emplace_back(function, description);
+      if (init_functions_run_yet_) {
+        // Run immediately, since GlobalInit already ran. This should be
+        // rare but we want to allow it in some cases.
+        LOG(WARNING) << "Running init function after GlobalInit: "
+                     << description;
+        // TODO(orionr): Consider removing argc and argv for non-early
+        // registration. Unfortunately that would require a new InitFunction
+        // typedef, so not making the change right now.
+        //
+        // Note that init doesn't receive argc and argv, so the function
+        // might fail and we want to raise an error in that case.
+        int argc = 0;
+        char** argv = nullptr;
+        bool success = (function)(&argc, &argv);
+        CAFFE_ENFORCE(success);
+      } else {
+        // Wait until GlobalInit to run
+        init_functions_.emplace_back(function, description);
+      }
     }
   }
 

caffe2/core/init_test.cc

@@ -8,11 +8,18 @@
 namespace caffe2 {
 namespace {
 bool gTestInitFunctionHasBeenRun = false;
+bool gTestFailInitFunctionHasBeenRun = false;
 
 bool TestInitFunction(int*, char***) {
   gTestInitFunctionHasBeenRun = true;
   return true;
 }
+
+bool TestFailInitFunction(int*, char***) {
+  gTestFailInitFunctionHasBeenRun = true;
+  return false;
+}
+
 REGISTER_CAFFE2_INIT_FUNCTION(
     TestInitFunction,
     &TestInitFunction,
@@ -27,6 +34,7 @@ char** dummy_argv = const_cast<char**>(&dummy_name);
 TEST(InitTest, TestInitFunctionHasRun) {
   caffe2::GlobalInit(&dummy_argc, &dummy_argv);
   EXPECT_TRUE(gTestInitFunctionHasBeenRun);
+  EXPECT_FALSE(gTestFailInitFunctionHasBeenRun);
 }
 
 TEST(InitTest, CanRerunGlobalInit) {
@@ -36,12 +44,26 @@ TEST(InitTest, CanRerunGlobalInit) {
 
 void LateRegisterInitFunction() {
   ::caffe2::InitRegisterer testInitFunc(
-      TestInitFunction, false, "This should fail");
+      TestInitFunction, false, "This should succeed but warn");
+}
+
+void LateRegisterEarlyInitFunction() {
+  ::caffe2::InitRegisterer testSecondInitFunc(
+      TestInitFunction, true, "This should fail for early init");
+}
+
+void LateRegisterFailInitFunction() {
+  ::caffe2::InitRegisterer testSecondInitFunc(
+      TestFailInitFunction, false, "This should fail for failed init");
 }
 
 TEST(InitTest, FailLateRegisterInitFunction) {
   caffe2::GlobalInit(&dummy_argc, &dummy_argv);
-  EXPECT_THROW(LateRegisterInitFunction(), caffe2::EnforceNotMet);
+  LateRegisterInitFunction();
+  EXPECT_THROW(LateRegisterEarlyInitFunction(), ::c10::Error);
+  EXPECT_THROW(LateRegisterFailInitFunction(), ::c10::Error);
+  EXPECT_TRUE(gTestInitFunctionHasBeenRun);
+  EXPECT_TRUE(gTestFailInitFunctionHasBeenRun);
 }
 
 } // namespace caffe2