From f069907c6cd096e6ddb023492834cd65259cbf2e Mon Sep 17 00:00:00 2001
From: Gert Mertes <hey@gertmertes.be>
Date: Tue, 12 Mar 2024 18:06:22 +0000
Subject: [PATCH] Improve remote request logging

---
 ai_models/remote.py | 41 ++++++++++++++++++++++++-----------------
 1 file changed, 24 insertions(+), 17 deletions(-)

diff --git a/ai_models/remote.py b/ai_models/remote.py
index 1f468f5..98dc9aa 100644
--- a/ai_models/remote.py
+++ b/ai_models/remote.py
@@ -112,7 +112,7 @@ def __init__(
 
         if url is None:
             url = os.getenv("AI_MODELS_REMOTE_URL", "https://ai-models.ecmwf.int")
-            LOG.info("Using remote %s", url)
+            LOG.info("Using remote server %s", url)
 
         token = token or os.getenv("AI_MODELS_REMOTE_TOKEN", None)
 
@@ -132,29 +132,32 @@ def __init__(
 
     def run(self, cfg: dict):
         # upload file
-        with open(self.input_file, "rb") as f:
-            LOG.info("Uploading input file to remote")
-            status, href = self._request(requests.post, "upload", data=f)
+        with open(self.input_file, "rb") as file:
+            LOG.info("Uploading input file to remote server")
+            _, status, href = self._request(requests.post, "upload", data=file)
 
         if status != "success":
             LOG.error(status)
             sys.exit(1)
 
-        # submit job
-        status, href = self._request(requests.post, href, json=cfg)
+        # submit task
+        id, status, href = self._request(requests.post, href, json=cfg)
+
+        LOG.info("Request submitted")
+        LOG.info("Request id: %s", id)
 
         if status != "queued":
             LOG.error(status)
             sys.exit(1)
 
-        LOG.info("Job status: queued")
+        LOG.info("Request is queued")
         last_status = status
 
         while True:
-            status, href = self._request(requests.get, href)
+            _, status, href = self._request(requests.get, href)
 
             if status != last_status:
-                LOG.info("Job status: %s", status)
+                LOG.info("Request is %s", status)
                 last_status = status
 
             if status == "failed":
@@ -180,29 +183,33 @@ def get_param(self, model, param):
             )
 
     def _request(self, type, href, data=None, json=None, auth=None, with_status=True):
-        r = robust(type, retry_after=30)(
+        response = robust(type, retry_after=30)(
             urljoin(self.url, href),
             json=json,
             data=data,
             auth=self.auth,
             timeout=self._timeout,
         )
+
+        if response.status_code == 401:
+            LOG.error("Unauthorized Access. Check your token.")
+            sys.exit(1)
+
         if with_status:
-            status, href = self._update_state(r)
-            return status, href
+            id, status, href = self._update_state(response)
+            return id, status, href
         else:
-            return r.json()
+            return response.json()
 
     def _update_state(self, response: requests.Response):
-        if response.status_code == 401:
-            return "Unauthorized Access", None
-
         try:
             data = response.json()
             href = data["href"]
             status = data["status"].lower()
+            id = data["id"]
         except Exception:
             status = f"{response.status_code} {response.url} {response.text}"
             href = None
+            id = None
 
-        return status, href
+        return id, status, href