Store model_details from RetrieveWorkerRun.model_version payload

arkindex_worker/worker/base.py

@@ -217,6 +217,9 @@ class BaseWorker(object):
         # Define model_version_id from environment
         self.model_version_id = os.environ.get("ARKINDEX_MODEL_VERSION_ID")
 
+        # Define model_details from environment
+        self.model_details = {"id": os.environ.get("ARKINDEX_MODEL_ID")}
+
         # Load all required secrets
         self.secrets = {name: self.load_secret(Path(name)) for name in required_secrets}
 
@@ -259,6 +262,9 @@ class BaseWorker(object):
             # Set model_version ID as worker attribute
             self.model_version_id = model_version.get("id")
 
+            # Set model details as worker attribute
+            self.model_details = model_version.get("model")
+
         # Retrieve initial configuration from API
         self.config = worker_version["configuration"].get("configuration", {})
         if "user_configuration" in worker_version["configuration"]:

docs/contents/implem/configure.md

@@ -115,6 +115,9 @@ Many attributes are set on the worker during at the configuration stage. Here is
 `model_version_id`
 : The ID of the model version linked to the current `WorkerRun` object on Arkindex. You may set it in developer mode via the `ARKINDEX_MODEL_VERSION_ID` environment variable.
 
+`model_details`
+: The details of the model for the model version linked to the current `WorkerRun` object on Arkindex. You may populate it in developer mode via the `ARKINDEX_MODEL_ID` environment variable.
+
 `process_information`
 : The details about the process parent to this worker execution. Only set in Arkindex mode.
 

tests/test_base_worker.py

@@ -454,7 +454,10 @@ def test_configure_load_model_configuration(mocker, monkeypatch, responses):
         "configuration": None,
         "model_version": {
             "id": "12341234-1234-1234-1234-123412341234",
-            "name": "Model version 1337",
+            "model": {
+                "id": "43214321-4321-4321-4321-432143214321",
+                "name": "Model 1337",
+            },
             "configuration": {
                 "param1": "value1",
                 "param2": 2,
@@ -489,6 +492,10 @@ def test_configure_load_model_configuration(mocker, monkeypatch, responses):
         "param3": None,
     }
     assert worker.model_version_id == "12341234-1234-1234-1234-123412341234"
+    assert worker.model_details == {
+        "id": "43214321-4321-4321-4321-432143214321",
+        "name": "Model 1337",
+    }
 
 
 def test_load_missing_secret():