Support extra_files in RQ tasks

closes #1695 (closed)

arkindex/ponos/tasks.py

@@ -105,7 +105,7 @@ def download_extra_files(task) -> None:
             resp.raise_for_status()
 
             # Write file to a specific data directory
-            extra_files_dir = settings.PONOS_DATA_DIR / "extra_files"
+            extra_files_dir = Path(settings.PONOS_DATA_DIR) / "extra_files"
             extra_files_dir.mkdir(exist_ok=True)
             with open(extra_files_dir / path_name, "wb") as f:
                 for chunk in resp.iter_content(chunk_size=8192):
@@ -178,7 +178,22 @@ def run_docker_task(client, task, temp_dir):
     task.state = State.Running
     task.save()
 
-    # 4. Read logs
+    # 4. Download extra_files
+    if task.extra_files:
+        logger.info("Downloading extra_files for task {!s}".format(task))
+        try:
+            download_extra_files(task)
+        except Exception as e:
+            logger.warning(
+                "Failed downloading extra_files for task {!s}: {!s}".format(
+                    task, e
+                )
+            )
+            task.state = State.Error
+            task.save()
+            return
+
+    # 5. Read logs
     logger.debug("Reading logs from the docker container")
     previous_logs = b""
     while container.status == "running":
@@ -198,7 +213,7 @@ def run_docker_task(client, task, temp_dir):
     # Upload logs one last time so we do not miss any data
     upload_logs(task, container.logs())
 
-    # 5. Retrieve the state of the container
+    # 6. Retrieve the state of the container
     container.reload()
     exit_code = container.attrs["State"]["ExitCode"]
     if exit_code != 0:
@@ -210,7 +225,7 @@ def run_docker_task(client, task, temp_dir):
     task.state = State.Completed
     task.save()
 
-    # 6. Upload artifacts
+    # 7. Upload artifacts
     logger.info(f"Uploading artifacts for task {task}")
 
     for path in Path(artifacts_dir).glob("**/*"):