Find local model directory method

Closes #152 (closed)

arkindex_worker/worker/base.py

@@ -45,6 +45,12 @@ def _is_500_error(exc: Exception) -> bool:
     return 500 <= exc.status_code < 600
 
 
+class ModelNotFoundError(Exception):
+    """
+    Exception raised when the path towards the model is invalid
+    """
+
+
 class BaseWorker(object):
     """
     Base class for Arkindex workers.
@@ -94,6 +100,12 @@ class BaseWorker(object):
             action="store_true",
             default=False,
         )
+        # To load models locally
+        self.parser.add_argument(
+            "--model-dir",
+            help=("The path to a local model's directory (development only)."),
+            type=Path,
+        )
 
         # Call potential extra arguments
         self.add_arguments()
@@ -110,6 +122,10 @@ class BaseWorker(object):
             self.work_dir = os.path.join(xdg_data_home, "arkindex")
             os.makedirs(self.work_dir, exist_ok=True)
 
+        # Store task ID. This is only available when running in production
+        # through a ponos agent
+        self.task_id = os.environ.get("PONOS_TASK")
+
         self.worker_version_id = os.environ.get("WORKER_VERSION_ID")
         if not self.worker_version_id:
             logger.warning(
@@ -245,12 +261,11 @@ class BaseWorker(object):
         """
         Setup the necessary attribute when using the cache system of `Base-Worker`.
         """
-        task_id = os.environ.get("PONOS_TASK")
         paths = None
         if self.support_cache and self.args.database is not None:
             self.use_cache = True
-        elif self.support_cache and task_id:
-            task = self.request("RetrieveTaskFromAgent", id=task_id)
+        elif self.support_cache and self.task_id:
+            task = self.request("RetrieveTaskFromAgent", id=self.task_id)
             paths = retrieve_parents_cache_path(
                 task["parents"],
                 data_dir=os.environ.get("PONOS_DATA", "/data"),
@@ -265,7 +280,9 @@ class BaseWorker(object):
                 ), f"Database in {self.args.database} does not exist"
                 self.cache_path = self.args.database
             else:
-                cache_dir = os.path.join(os.environ.get("PONOS_DATA", "/data"), task_id)
+                cache_dir = os.path.join(
+                    os.environ.get("PONOS_DATA", "/data"), self.task_id
+                )
                 assert os.path.isdir(cache_dir), f"Missing task cache in {cache_dir}"
                 self.cache_path = os.path.join(cache_dir, "db.sqlite")
 
@@ -335,6 +352,35 @@ class BaseWorker(object):
         # By default give raw secret payload
         return secret
 
+    def find_model_directory(self) -> Path:
+        """
+        Find the local path to the model. This supports two modes:
+        - the worker runs in ponos, the model is available at `/data/current`
+        - the worker runs locally, the developer may specify it using either
+           - the `model_dir` configuration parameter
+           - the `--model-dir` CLI parameter
+
+        :return: Path to the model on disk
+        """
+        if self.task_id:
+            # When running in production with ponos, the agent
+            # downloads the model and set it in the current task work dir
+            return Path(self.work_dir)
+        else:
+            model_dir = self.config.get("model_dir", self.args.model_dir)
+            if model_dir is None:
+                raise ModelNotFoundError(
+                    "No path to the model was provided. "
+                    "Please provide model_dir either through configuration "
+                    "or as CLI argument."
+                )
+            model_dir = Path(model_dir)
+            if not model_dir.exists():
+                raise ModelNotFoundError(
+                    f"The path {model_dir} does not link to any directory"
+                )
+            return model_dir
+
     @retry(
         retry=retry_if_exception(_is_500_error),
         wait=wait_exponential(multiplier=2, min=3),