Populate redis with resource-allocation data

Author: Gossty

qiita_db/meta_util.py

@@ -38,7 +38,7 @@
 from json import loads, dump, dumps
 
 from qiita_db.util import create_nested_path, _retrieve_resource_data
-from qiita_db.util import resource_allocation_plot, get_software_commands
+from qiita_db.util import resource_allocation_plot
 from qiita_core.qiita_settings import qiita_config, r_client
 from qiita_core.configuration_manager import ConfigurationManager
 import qiita_db as qdb
@@ -49,6 +49,7 @@
     "parameters", "samples", "columns", "input_size", "extra_info",
     "MaxRSSRaw", "ElapsedRaw", "Start", "node_name", "node_model"]
 
+
 def _get_data_fpids(constructor, object_id):
     """Small function for getting filepath IDS associated with data object
 
@@ -554,40 +555,96 @@ def generate_plugin_releases():
         f(redis_key, v)
 
 
-def initialize_resource_allocations_redis():
-    time = datetime.now().strftime('%m-%d-%y %H:%M:%S')
+def get_software_commands():
+    # TODO change active=False to True. In test, the data is active=False.
+    software_list = [s for s in qdb.software.Software.iter(active=False)]
+    software_commands = dict()
+    for software in software_list:
+        sname = software.name
+        sversion = software.version
+        commands = software.commands
+        if sname not in software_commands:
+            software_commands[sname] = {}
+        if sversion not in software_commands[sname]:
+            software_commands[sname][sversion] = []
+        for command in commands:
+            software_commands[sname][sversion].append(command.name)
+    return software_commands
+
+
+def update_resource_allocation_redis():
+    time = datetime.now().strftime('%m-%d-%y')
     scommands = get_software_commands()
-    for software, versions in scommands.items():
-        for version, commands in versions.items():
-            for command in commands:
-                print("Generating plot for:", software, version, command)
-                update_resource_allocation_redis(command, software,
-                                                 version, time)
     redis_key = 'resources:commands'
     r_client.set(redis_key, str(scommands))
 
+    for sname, versions in scommands.items():
+        for version, commands in versions.items():
+            for cname in commands:
 
-def update_resource_allocation_redis(cname, sname, version, time):
-    col_name = "samples * columns"
-
-    df = _retrieve_resource_data(cname, sname, version, columns)
-    fig, axs = resource_allocation_plot(df, cname, sname, col_name)
-
-    fig.tight_layout()
-    plot = BytesIO()
-    fig.savefig(plot, format='png')
-    plot.seek(0)
-
-    img = 'data:image/png;base64,' + quote(b64encode(plot.getbuffer()).decode('ascii'))
-    plt.close(fig)
-
-    # SID, CID, col_name
-    values = [
-        ("img", img, r_client.set),
-        ('time', time, r_client.set)
-    ]
-
-    for k, v, f in values:
-        redis_key = 'resources$#%s$#%s$#%s:%s' % (cname, sname, col_name, k)
-        r_client.delete(redis_key)
-        f(redis_key, v)
+                print("Generating plot for:", sname, version, cname)
+                col_name = "samples * columns"
+                df = _retrieve_resource_data(cname, sname, version, columns)
+                if len(df) == 0:
+                    print("No available data for", sname, version, cname)
+                    continue
+                    
+                fig, axs = resource_allocation_plot(df, cname, sname, col_name)
+                titles = [0, 0]
+                images = [0, 0]
+
+                # Splitting 1 image plot into 2 separate for better layout.
+                for i, ax in enumerate(axs):
+                    titles[i] = ax.get_title()
+                    ax.set_title("")
+                    # new_fig, new_ax – copy with either only memory plot or
+                    # only time
+                    new_fig = plt.figure()
+                    new_ax = new_fig.add_subplot(111)
+
+                    scatter_data = ax.collections[0]
+                    new_ax.scatter(scatter_data.get_offsets()[:, 0],
+                                   scatter_data.get_offsets()[:, 1],
+                                   s=scatter_data.get_sizes(), label="data")
+
+                    line = ax.lines[0]
+                    new_ax.plot(line.get_xdata(), line.get_ydata(),
+                                linewidth=1, color='orange')
+
+                    if len(ax.collections) > 1:
+                        failure_data = ax.collections[1]
+                        new_ax.scatter(failure_data.get_offsets()[:, 0],
+                                       failure_data.get_offsets()[:, 1],
+                                       color='red', s=3, label="failures")
+
+                    new_ax.set_xscale('log')
+                    new_ax.set_yscale('log')
+                    new_ax.set_xlabel(ax.get_xlabel())
+                    new_ax.set_ylabel(ax.get_ylabel())
+                    new_ax.legend(loc='upper left')
+
+                    new_fig.tight_layout()
+                    plot = BytesIO()
+                    new_fig.savefig(plot, format='png')
+                    plot.seek(0)
+                    img = 'data:image/png;base64,' + quote(
+                          b64encode(plot.getvalue()).decode('ascii'))
+                    images[i] = img
+                    plt.close(new_fig)
+                plt.close(fig)
+
+                # SID, CID, col_name
+                values = [
+                    ("img_mem", images[0], r_client.set),
+                    ("img_time", images[1], r_client.set),
+                    ('time', time, r_client.set),
+                    ("title_mem", titles[0], r_client.set),
+                    ("title_time", titles[1], r_client.set)
+                ]
+                print(time, titles[0], titles[1])
+
+                for k, v, f in values:
+                    redis_key = 'resources$#%s$#%s$#%s$#%s:%s' % (
+                                cname, sname, version, col_name, k)
+                    r_client.delete(redis_key)
+                    f(redis_key, v)

qiita_db/test/test_meta_util.py

@@ -13,6 +13,7 @@
 from os.path import exists, join
 
 import pandas as pd
+from datetime import datetime
 
 from qiita_core.qiita_settings import qiita_config, r_client
 from qiita_core.util import qiita_test_checker
@@ -523,18 +524,29 @@ def test_update_resource_allocation_redis(self):
         cname = "Split libraries FASTQ"
         sname = "QIIMEq2"
         col_name = "samples * columns"
-        qdb.meta_util.update_resource_allocation_redis(cname, sname)
-
-        vals = [
-            ("img", b's', r_client.get),
-            ('time', b'2024-11-11', r_client.get)
-        ]
-
-        for k, exp, f in vals:
-            redis_key = 'resources$#%s$#%s$#%s:%s' % (cname, sname, col_name, k)
-            # checking redis values
-            print(f(redis_key))
-            # self.assertEqual(f(redis_key), exp)
+        version = "1.9.1"
+        qdb.meta_util.update_resource_allocation_redis()
+        # since time is month, day, year, it should be equal unless test is ran
+        # at midnight.
+        time = datetime.now().strftime('%m-%d-%y')
+        title_mem_str = 'resources$#%s$#%s$#%s$#%s:%s' % (
+            cname, sname, version, col_name, 'title_mem')
+        title_mem = str(r_client.get(title_mem_str))
+        self.assertTrue(
+            "model_chosen: "
+            "k * log(x) + "
+            "b * log(x)^2 + "
+            "a * log(x)^2.5" in title_mem
+        )
+
+        title_time_str = 'resources$#%s$#%s$#%s$#%s:%s' % (
+                        cname, sname, version, col_name, 'title_time')
+        title_time = str(r_client.get(title_time_str))
+        self.assertTrue("model_chosen: a + b + log(x) * k" in title_time)
+        time_create_str = 'resources$#%s$#%s$#%s$#%s:%s' % (
+            cname, sname, version, col_name, 'time')
+        time_create = str(r_client.get(time_create_str))
+        self.assertTrue(time in time_create)
 
 
 if __name__ == '__main__':

qiita_db/util.py

@@ -99,6 +99,27 @@
 MODELS_TIME = [time_model1, time_model2, time_model3, time_model4]
 
 
+def get_model_name(model):
+    if model == mem_model1:
+        return "k * log(x) + x * a + b"
+    elif model == mem_model2:
+        return "k * log(x) + b * log(x)^2 + a"
+    elif model == mem_model3:
+        return "k * log(x) + b * log(x)^2 + a * log(x)^3"
+    elif model == mem_model4:
+        return "k * log(x) + b * log(x)^2 + a * log(x)^2.5"
+    elif model == time_model1:
+        return "a + b + log(x) * k"
+    elif model == time_model2:
+        return "a + b * x + log(x) * k"
+    elif model == time_model3:
+        return "a + b * log(x)^2 + log(x) * k"
+    elif model == time_model4:
+        return "a * log(x)^3 + b * log(x)^2 + log(x) * k"
+    else:
+        return "Unknown model"
+
+
 def scrub_data(s):
     r"""Scrubs data fields of characters not allowed by PostgreSQL
 
@@ -2496,7 +2517,6 @@ def _resource_allocation_plot_helper(
     cmin = naturalsize(cmin_value, gnu=True) if curr == "MaxRSSRaw" else \
         str(timedelta(seconds=round(cmin_value, 2))).rstrip('0').rstrip('.')
 
-
     x_plot = np.array(df[col_name])
     failures_df = _resource_allocation_failures(
         df, k, a, b, best_model, col_name, curr)
@@ -2505,7 +2525,10 @@ def _resource_allocation_plot_helper(
     ax.scatter(failures_df[col_name], failures_df[curr], color='red', s=3,
                label="failures")
 
-    ax.set_title(f'{cname}: {sname}\n real: {mini} || {maxi}\n'
+    ax.set_title(
+                 f'k||a||b: {k}||{a}||{b}\n'
+                 f'model_chosen: {get_model_name(best_model)}\n'
+                 f'real: {mini} || {maxi}\n'
                  f'calculated: {cmin} || {cmax}\n'
                  f'failures: {failures}')
     ax.legend(loc='upper left')
@@ -2924,35 +2947,3 @@ def merge_rows(rows):
                 row['node_model']]
             qdb.sql_connection.TRN.add(sql, sql_args=to_insert)
             qdb.sql_connection.TRN.execute()
-
-
-def get_software_commands():
-    res = []
-    with qdb.sql_connection.TRN:
-        sql_command = """
-            SELECT DISTINCT
-                s.name AS sName,
-                s.version AS sVersion,
-                sc.name AS cName
-            FROM
-                qiita.slurm_resource_allocations sra
-            JOIN
-                qiita.processing_job pr
-                ON sra.processing_job_id = pr.processing_job_id
-            JOIN
-                qiita.software_command sc on pr.command_id = sc.command_id
-            JOIN
-                qiita.software s ON sc.software_id = s.software_id
-            """
-        qdb.sql_connection.TRN.add(sql_command)
-        res = qdb.sql_connection.TRN.execute_fetchindex()
-
-    software_commands = dict()
-    for s_name, s_version, c_name in res:
-        if s_name not in software_commands:
-            software_commands[s_name] = {}
-        if s_version not in software_commands[s_name]:
-            software_commands[s_name][s_version] = set()
-        software_commands[s_name][s_version].add(c_name)
-
-    return software_commands

scripts/all-qiita-cron-job

@@ -4,6 +4,6 @@ qiita-cron-job empty-trash-upload-folder
 qiita-cron-job generate-biom-and-metadata-release
 qiita-cron-job purge-filepaths
 qiita-cron-job update-redis-stats
-qiita-cron-job initialize-resource-allocations-redis
+qiita-cron-job update-resource-allocation-redis
 qiita-cron-job generate-plugin-releases
 qiita-cron-job purge-json-web-tokens

scripts/qiita-cron-job

@@ -16,7 +16,7 @@ from qiita_db.util import (
     quick_mounts_purge as qiita_quick_mounts_purge)
 from qiita_db.meta_util import (
     update_redis_stats as qiita_update_redis_stats,
-    initialize_resource_allocations_redis as qiita_initialize_resource_allocations_redis,
+    update_resource_allocation_redis as qiita_update_resource_allocation_redis,
     generate_biom_and_metadata_release as
     qiita_generate_biom_and_metadata_release,
     generate_plugin_releases as qiita_generate_plugin_releases)
@@ -50,8 +50,8 @@ def update_redis_stats():
 
 
 @commands.command()
-def initialize_resource_allocations_redis():
-    qiita_initialize_resource_allocations_redis()
+def update_resource_allocation_redis():
+    qiita_update_resource_allocation_redis()
 
 
 @commands.command()