Fix predictor __init__ staggering (#1591)

Author: deliahu

pkg/workloads/cortex/lib/api/predictor.py

@@ -52,9 +52,6 @@
     ModelsTree,  # only when num workers = 1
 )
 
-# concurrency
-from cortex.lib.concurrency import FileLock
-
 # model validation
 from cortex.lib.model import validate_model_paths
 
@@ -245,10 +242,9 @@ def class_impl(self, project_dir):
             validations = PYTHON_CLASS_VALIDATION
 
         try:
-            with FileLock("/run/init_stagger.lock"):
-                predictor_class = self._get_class_impl(
-                    "cortex_predictor", os.path.join(project_dir, self.path), target_class_name
-                )
+            predictor_class = self._get_class_impl(
+                "cortex_predictor", os.path.join(project_dir, self.path), target_class_name
+            )
         except CortexException as e:
             e.wrap("error in " + self.path)
             raise

pkg/workloads/cortex/serve/serve.py

@@ -35,7 +35,7 @@
 from cortex.lib import util
 from cortex.lib.api import API, get_api
 from cortex.lib.log import cx_logger as logger
-from cortex.lib.concurrency import LockedFile
+from cortex.lib.concurrency import FileLock, LockedFile
 from cortex.lib.storage import S3, LocalStorage
 from cortex.lib.exceptions import UserRuntimeException
 
@@ -304,8 +304,10 @@ def start_fn():
         client = api.predictor.initialize_client(
             tf_serving_host=tf_serving_host, tf_serving_port=tf_serving_port
         )
-        logger().info("loading the predictor from {}".format(api.predictor.path))
-        predictor_impl = api.predictor.initialize_impl(project_dir, client)
+
+        with FileLock("/run/init_stagger.lock"):
+            logger().info("loading the predictor from {}".format(api.predictor.path))
+            predictor_impl = api.predictor.initialize_impl(project_dir, client)
 
         local_cache["api"] = api
         local_cache["provider"] = provider