From fe1c963eeca99e90c36a97479cfac853dbdd93d9 Mon Sep 17 00:00:00 2001
From: allegroai <>
Date: Sun, 23 Aug 2020 15:40:57 +0300
Subject: [PATCH] Fix internal export utility

---
 server/mongo/initialize/pre_populate.py | 16 +++++++++++++++-
 1 file changed, 15 insertions(+), 1 deletion(-)

diff --git a/server/mongo/initialize/pre_populate.py b/server/mongo/initialize/pre_populate.py
index 521fcc0..3844608 100644
--- a/server/mongo/initialize/pre_populate.py
+++ b/server/mongo/initialize/pre_populate.py
@@ -57,6 +57,10 @@ class PrePopulate:
     metadata_filename = "metadata.json"
     zip_args = dict(mode="w", compression=ZIP_BZIP2)
     artifacts_ext = ".artifacts"
+    img_source_regex = re.compile(
+        r"['\"]source['\"]:\s?['\"](https?://(?:localhost:8081|files.*?)/.*?)['\"]",
+        flags=re.IGNORECASE,
+    )
 
     class JsonLinesWriter:
         def __init__(self, file: BinaryIO):
@@ -360,6 +364,8 @@ class PrePopulate:
     @classmethod
     def update_featured_projects_order(cls):
         featured_order = config.get("services.projects.featured_order", [])
+        if not featured_order:
+            return
 
         def get_index(p: Project):
             for index, entry in enumerate(featured_order):
@@ -502,11 +508,19 @@ class PrePopulate:
                         break
                     scroll_id = res.next_scroll_id
                     for event in res.events:
-                        if event.get("type") == "training_debug_image":
+                        event_type = event.get("type")
+                        if event_type == "training_debug_image":
                             url = cls._get_fixed_url(event.get("url"))
                             if url:
                                 event["url"] = url
                                 artifacts.append(url)
+                        elif event_type == "plot":
+                            plot_str: str = event.get("plot_str", "")
+                            for match in cls.img_source_regex.findall(plot_str):
+                                url = cls._get_fixed_url(match)
+                                if match != url:
+                                    plot_str = plot_str.replace(match, url)
+                                artifacts.append(url)
                         w.write(json.dumps(event))
             data = f.getvalue()
             hash_.update(data)