From 6d7ab9581d71455377df5d4cb443fd867dca22b4 Mon Sep 17 00:00:00 2001
From: Manon Blanco <blanco@teklia.com>
Date: Wed, 28 Feb 2024 16:43:17 +0000
Subject: [PATCH] Store model_version details regardless of model version
configuration
---
arkindex_worker/worker/base.py | 15 +++++++++++----
tests/test_elements_worker/test_worker.py | 4 ++++
2 files changed, 15 insertions(+), 4 deletions(-)
diff --git a/arkindex_worker/worker/base.py b/arkindex_worker/worker/base.py
index d638df80..1be5df4f 100644
--- a/arkindex_worker/worker/base.py
+++ b/arkindex_worker/worker/base.py
@@ -148,6 +148,13 @@ class BaseWorker:
# there is at least one available sqlite database either given or in the parent tasks
self.use_cache = False
+ # model_version_id will be updated in configure() using the worker_run's model version
+ # or in configure_for_developers() from the environment
+ self.model_version_id = None
+ # model_details will be updated in configure() using the worker_run's model version
+ # or in configure_for_developers() from the environment
+ self.model_details = {}
+
# task_parents will be updated in configure_cache() if the cache is supported,
# if the task ID is set and if no database is passed as argument
self.task_parents = []
@@ -257,15 +264,15 @@ class BaseWorker:
# Load model version configuration when available
model_version = worker_run.get("model_version")
- if model_version and model_version.get("configuration"):
+ if model_version:
logger.info("Loaded model version configuration from WorkerRun")
- self.model_configuration.update(model_version.get("configuration"))
+ self.model_configuration.update(model_version["configuration"])
# Set model_version ID as worker attribute
- self.model_version_id = model_version.get("id")
+ self.model_version_id = model_version["id"]
# Set model details as worker attribute
- self.model_details = model_version.get("model")
+ self.model_details = model_version["model"]
# Retrieve initial configuration from API
self.config = worker_version["configuration"].get("configuration", {})
diff --git a/tests/test_elements_worker/test_worker.py b/tests/test_elements_worker/test_worker.py
index f84ef247..30f81a49 100644
--- a/tests/test_elements_worker/test_worker.py
+++ b/tests/test_elements_worker/test_worker.py
@@ -467,6 +467,10 @@ def test_worker_config_multiple_source(
"id": "12341234-1234-1234-1234-123412341234",
"name": "Model version 1337",
"configuration": model_config,
+ "model": {
+ "id": "hahahaha-haha-haha-haha-hahahahahaha",
+ "name": "My model",
+ },
},
"process": {
"name": None,
--
GitLab