From dfa914a90cb88a78426c283494a44eab7f6522f4 Mon Sep 17 00:00:00 2001
From: Colin L Reliability Rice <clr@fb.com>
Date: Wed, 22 Jul 2020 15:33:49 -0700
Subject: [PATCH] Modify lazy_dyndep loading to trigger inside workspace.
 (#41687)

Summary:
Pull Request resolved: https://github.com/pytorch/pytorch/pull/41687

Specifically, this makes a new library (lazy), which can be used from both core
and workspace.

This allows workspace.Createnet to trigger lazy loading of dyndep dependencies.

Test Plan: Added a unit test specifically for workspace.CreateNet

Reviewed By: dzhulgakov

Differential Revision: D22441877

fbshipit-source-id: 3a9d1af9962585d08ea2566c9c85bec7377d39f2
---
 caffe2/python/core.py             | 23 ++++++-----------------
 caffe2/python/lazy.py             | 14 ++++++++++++++
 caffe2/python/lazy_dyndep.py      |  4 ++--
 caffe2/python/lazy_dyndep_test.py | 14 ++++++++++++++
 caffe2/python/workspace.py        |  2 ++
 5 files changed, 38 insertions(+), 19 deletions(-)
 create mode 100644 caffe2/python/lazy.py

diff --git a/caffe2/python/core.py b/caffe2/python/core.py
index ad610ee91dabfd..92d41837429e9c 100644
--- a/caffe2/python/core.py
+++ b/caffe2/python/core.py
@@ -13,6 +13,7 @@
 
 from caffe2.proto import caffe2_pb2
 from caffe2.python import scope, utils, workspace
+from caffe2.python.lazy import TriggerLazyImport
 from caffe2.python.control_ops_grad import \
     gen_do_gradient, gen_if_gradient, gen_while_gradient, disambiguate_grad_if_op_output
 
@@ -49,18 +50,6 @@ def _InitDataType():
 
 _InitDataType()
 
-_import_lazy_calls = []
-
-def RegisterLazyImport(lazy):
-    global _import_lazy_calls
-    _import_lazy_calls += [lazy]
-
-
-def _import_lazy():
-    global _import_lazy_calls
-    for lazy in _import_lazy_calls:
-        lazy()
-
 
 def _GetRegisteredOperators():
     return set(workspace.RegisteredOperators())
@@ -71,7 +60,7 @@ def _GetRegisteredOperators():
 
 def RefreshRegisteredOperators(trigger_lazy=True):
     if trigger_lazy:
-        _import_lazy()
+        TriggerLazyImport()
     global _REGISTERED_OPERATORS
     _REGISTERED_OPERATORS = _GetRegisteredOperators()
 
@@ -80,7 +69,7 @@ def RefreshRegisteredOperators(trigger_lazy=True):
 
 
 def GlobalInit(args):
-    _import_lazy()
+    TriggerLazyImport()
     _GLOBAL_INIT_ARGS.extend(args[1:])
     C.global_init(args)
 
@@ -94,7 +83,7 @@ def IsOperator(op_type):
 
 
 def IsOperatorWithEngine(op_type, engine):
-    _import_lazy()
+    TriggerLazyImport()
     return C.op_registry_key(op_type, engine) in _REGISTERED_OPERATORS
 
 
@@ -294,7 +283,7 @@ def __getattr__(self, op_type):
             op_type, *args, **kwargs)
 
     def __dir__(self):
-        _import_lazy()
+        TriggerLazyImport()
         additional_methods = [
             op
             for op in _REGISTERED_OPERATORS
@@ -2228,7 +2217,7 @@ def __getattr__(self, op_type):
             op_type, *args, **kwargs)
 
     def __dir__(self):
-        _import_lazy()
+        TriggerLazyImport()
         additional_methods = [
             op
             for op in _REGISTERED_OPERATORS
diff --git a/caffe2/python/lazy.py b/caffe2/python/lazy.py
new file mode 100644
index 00000000000000..8e61f41767c867
--- /dev/null
+++ b/caffe2/python/lazy.py
@@ -0,0 +1,14 @@
+## @package workspace
+# Module caffe2.python.lazy
+
+_import_lazy_calls = []
+
+def RegisterLazyImport(lazy):
+    global _import_lazy_calls
+    _import_lazy_calls += [lazy]
+
+
+def TriggerLazyImport():
+    global _import_lazy_calls
+    for lazy in _import_lazy_calls:
+        lazy()
diff --git a/caffe2/python/lazy_dyndep.py b/caffe2/python/lazy_dyndep.py
index 912c3fd6231afe..e1799838f4b2ba 100644
--- a/caffe2/python/lazy_dyndep.py
+++ b/caffe2/python/lazy_dyndep.py
@@ -6,7 +6,7 @@
 from __future__ import unicode_literals
 
 import os
-from caffe2.python import core, dyndep
+from caffe2.python import dyndep, lazy
 
 
 def RegisterOpsLibrary(name):
@@ -81,4 +81,4 @@ def _import_lazy():
         finally:
             _LAZY_IMPORTED_DYNDEPS.remove(name)
 
-core.RegisterLazyImport(_import_lazy)
+lazy.RegisterLazyImport(_import_lazy)
diff --git a/caffe2/python/lazy_dyndep_test.py b/caffe2/python/lazy_dyndep_test.py
index 9a86a0b82d3a03..cbffb4304dadbf 100644
--- a/caffe2/python/lazy_dyndep_test.py
+++ b/caffe2/python/lazy_dyndep_test.py
@@ -113,6 +113,20 @@ def handlernoop(e):
             lazy_dyndep.RegisterOpsLibrary("@/caffe2/caffe2/distributed:file_store_handler_ops")
             core.RefreshRegisteredOperators()
 
+    def test_workspacecreatenet(self):
+        from caffe2.python import workspace, lazy_dyndep
+        import tempfile
+
+        with tempfile.NamedTemporaryFile() as f:
+            lazy_dyndep.RegisterOpsLibrary(f.name)
+            called = False
+
+            def handler(e):
+                raise ValueError("test")
+            lazy_dyndep.SetErrorHandler(handler)
+            with self.assertRaises(ValueError, msg="test"):
+                workspace.CreateNet("fake")
+
 
 if __name__ == "__main__":
     unittest.main()
diff --git a/caffe2/python/workspace.py b/caffe2/python/workspace.py
index da1719643aee8b..cc7f849d5e5b6d 100644
--- a/caffe2/python/workspace.py
+++ b/caffe2/python/workspace.py
@@ -19,6 +19,7 @@
 
 from caffe2.proto import caffe2_pb2
 from caffe2.python import scope, utils
+from caffe2.python.lazy import TriggerLazyImport
 
 import caffe2.python._import_c_extension as C
 
@@ -172,6 +173,7 @@ def ResetWorkspace(root_folder=None):
 
 
 def CreateNet(net, overwrite=False, input_blobs=None):
+    TriggerLazyImport()
     if input_blobs is None:
         input_blobs = []
     for input_blob in input_blobs: