template varible support widget

aiidalab_widgets_base/computational_resources.py

@@ -2,6 +2,7 @@
 import os
 import subprocess
 import threading
+from collections import namedtuple
 from copy import copy
 from pathlib import Path
 from uuid import UUID
@@ -15,8 +16,10 @@
 from aiida.transports.plugins.ssh import parse_sshconfig
 from humanfriendly import InvalidSize, parse_size
 from IPython.display import clear_output, display
+from jinja2 import Environment, meta
+from jinja2.nodes import Filter
 
-from .databases import ComputationalResourcesDatabaseWidget
+from .databases import NewComputationalResourcesDatabaseWidget
 from .utils import StatusHTML
 
 STYLE = {"description_width": "140px"}
@@ -1269,6 +1272,111 @@ def _validate_value(self, change):
             return computer_uuid
 
 
+class TemplateVariablesWidget(ipw.VBox):
+    # The input template is a dictionary of keyname and template string.
+    templates = traitlets.Dict(allow_none=True)
+
+    # The output template is a dictionary of keyname and filled string.
+    filled_templates = traitlets.Dict(allow_none=True)
+
+    def __init__(self):
+        # A placeholder for the template variables widget.
+        self.template_variables = ipw.VBox()
+
+        # A dictionary of mapping variables.
+        # the key is the variable name, and the value is a tuple of (template value and widget).
+        self._mapping_variables = {}
+
+        super().__init__(
+            children=[
+                ipw.HTML(
+                    """<div>Please fill in the template variables below.</div>
+                    """
+                ),
+                self.template_variables,
+            ]
+        )
+
+    @traitlets.observe("templates")
+    def _templates_changed(self, _=None):
+        """Render the template variables widget."""
+        self._render()
+
+        # Update the output filled template.
+        self.filled_templates = self.templates.copy()
+
+    def _render(self):
+        """Render the template variables widget."""
+
+        for key, temp_str in self.templates.items():
+            env = Environment()
+            parsed_content = env.parse(temp_str)
+
+            # vars is a set of variables in the template, while filters is a list of all filters in the template
+            # For example, {{ label | default('daint-mc') }} will have vars = {'label'} and filters = [<Filter: default('daint-mc')>]
+            # not all variables have filter.
+            variables = meta.find_undeclared_variables(parsed_content)
+            filters = list(parsed_content.find_all(Filter))
+
+            default_value = {
+                fltr.node.name: fltr.args[0].value
+                for fltr in filters
+                if fltr.name == "default"
+            }
+
+            for var in variables:
+                w = ipw.Text(
+                    description=f"{var}:",
+                    value=default_value.get(var, ""),
+                    layout=LAYOUT,
+                    style=STYLE,
+                )
+                # Every time the value of the widget changes, we update the filled template.
+                # This migth be too much to sync the final filled template every time.
+                w.observe(self._on_template_variable_filled, names="value")
+                # self._mapping_variables[var] = (key, temp_str, w)
+                self._mapping_variables[var] = namedtuple(
+                    "MappingVariable", ["key", "temp_str", "widget", "variables"]
+                )(key, temp_str, w, variables)
+
+        # Render by change the VBox children of placeholder.
+        self.template_variables.children = [
+            mapping_variable.widget
+            for mapping_variable in self._mapping_variables.values()
+        ]
+
+    def _on_template_variable_filled(self, change):
+        """Callback when a template variable is filled."""
+        # Update the changed filled template for the widget that is changed.
+        for var, mapping_variable in self._mapping_variables.items():
+            if mapping_variable.widget is change["owner"]:
+                # See if all variables are set in widget and ready from the mapping
+                variables = mapping_variable.variables
+                for v in variables:
+                    if self._mapping_variables[v][2].value == "":
+                        return
+
+                # If all variables are ready, update the filled template.
+                inp_dict = {v: self._mapping_variables[v][2].value for v in variables}
+
+                # re-render the template
+                env = Environment()
+                filled_str = env.from_string(mapping_variable.temp_str).render(
+                    **inp_dict
+                )
+
+                # Update the filled template.
+                self.filled_templates[mapping_variable.key] = filled_str
+
+                # Update the template partially in the filled template and filled with other variables in the same template string.
+                # The template string can contain multiple variables.
+                self._mapping_variables[var] = namedtuple(
+                    "MappingVariable", ["key", "temp_str", "widget"]
+                )(mapping_variable.key, filled_str, mapping_variable.widget)
+
+                break
+
+
 class QuickSetupWidget(ipw.VBox):
     """The widget that allows to quickly setup a computer and code."""
 
@@ -1284,7 +1392,7 @@ def __init__(self, default_calc_job_plugin=None, **kwargs):
         quick_setup_button.on_click(self._on_quick_setup)
 
         # resource database for setup computer/code.
-        self.comp_resources_database = ComputationalResourcesDatabaseWidget(
+        self.comp_resources_database = NewComputationalResourcesDatabaseWidget(
             default_calc_job_plugin=default_calc_job_plugin
         )
         self.comp_resources_database.observe(
@@ -1314,13 +1422,19 @@ def __init__(self, default_calc_job_plugin=None, **kwargs):
             (self, "message"),
         )
 
+        # The placeholder widget for the template variable of config.
+        self.template_variables_computer = TemplateVariablesWidget()
+        self.template_variables_code = TemplateVariablesWidget()
+
         super().__init__(
             children=[
                 ipw.HTML(
                     """<div>Please select the computer/code from a database to pre-fill the fields below.</div>
                     """
                 ),
                 self.comp_resources_database,
+                self.template_variables_computer,
+                self.template_variables_code,
                 ipw.HTML("""<div>SSH credential to the remote machine</div>"""),
                 self.ssh_computer_setup.username,
                 self.ssh_computer_setup.password_box,
@@ -1338,27 +1452,41 @@ def _on_select_computer(self, change):
         self.reset()
 
         new_setup = change["new"]
+
+        # Read from template and prepare the widgets for the template variables.
+        self.template_variables_computer.template = new_setup
+
+        # pre-set the input fields no matter if the template variables are set.
         self.aiida_computer_setup.computer_setup = new_setup
         self.computer_setup = new_setup
 
         # ssh config need to sync hostname etc to resource database.
         self.ssh_computer_setup.ssh_config = self.comp_resources_database.ssh_config
         self.ssh_config = self.comp_resources_database.ssh_config
 
+        # for temp_str in change["new"].values():
+
     def _on_select_code(self, change):
         """Update the code trait"""
         if change["new"] is None:
             return
 
+        # XXX check if code needs to be reset or not
+
         new_setup = change["new"]
+        self.template_variables_code.template = new_setup
+
         self.aiida_code_setup.code_setup = new_setup
         self.code_setup = new_setup
 
     def _on_quick_setup(self, _=None):
         """Go through all the setup steps automatically."""
-        self.ssh_computer_setup._on_setup_ssh_button_pressed()
-        if self.aiida_computer_setup.on_setup_computer():
-            self.aiida_code_setup.on_setup_code()
+        # Fill text fields with template variables. (computer setup)
+        # XXX
+
+        # self.ssh_computer_setup._on_setup_ssh_button_pressed()
+        # if self.aiida_computer_setup.on_setup_computer():
+        #    self.aiida_code_setup.on_setup_code()
 
     def _on_setup_computer_success(self):
         """Callback that is called when the computer is successfully set up."""
@@ -1393,6 +1521,7 @@ class DetailedSetupWidget(ipw.VBox):
     computer_setup = traitlets.Dict(allow_none=True)
     code_setup = traitlets.Dict(allow_none=True)
 
+    # XXX: distinguish the message from different sub widgets. add prefix of quick setup and detailed setup.
     message = traitlets.Unicode()
     success = traitlets.Bool(False)
 

notebooks/computational_resources.ipynb

@@ -60,7 +60,7 @@
   {
    "cell_type": "code",
    "execution_count": null,
-   "id": "7a302c5f",
+   "id": "ee9ed2e5",
    "metadata": {},
    "outputs": [],
    "source": [
@@ -71,7 +71,7 @@
   {
    "cell_type": "code",
    "execution_count": null,
-   "id": "fdcce7b2",
+   "id": "11b53118",
    "metadata": {},
    "outputs": [],
    "source": []

tests/test_computational_resources.py

@@ -216,6 +216,51 @@ def test_computer_dropdown_widget(aiida_localhost):
     assert widget._dropdown.value is None
 
 
+def test_template_variables_widget():
+    """Test template_variables_widget."""
+    from aiidalab_widgets_base.computational_resources import TemplateVariablesWidget
+
+    w = TemplateVariablesWidget()
+
+    w.templates = {
+        "label": "{{ label | default('daint-mc') }}",
+        "hostname": "daint.cscs.ch",
+        "description": "Piz Daint supercomputer at CSCS Lugano, Switzerland, multicore partition.",
+        "transport": "core.ssh",
+        "scheduler": "core.slurm",
+        "work_dir": "/scratch/snx3000/{username}/aiida_run/",
+        "shebang": "#!/bin/bash",
+        "mpirun_command": "srun -n {tot_num_mpiprocs}",
+        "mpiprocs_per_machine": 36,
+        "prepend_text": "#SBATCH --partition={{ slurm_partition | default('normal') }}\n#SBATCH --account={{ slurm_account }}\n#SBATCH --constraint=mc\n#SBATCH --cpus-per-task=1\n\nexport OMP_NUM_THREADS=$SLURM_CPUS_PER_TASK\nsource $MODULESHOME/init/bash\nulimit -s unlimited",
+    }
+
+    # Fill the template variables
+    for key, value in w._mapping_variables.items():
+        if key == "label":
+            sub_widget = value[2]
+            sub_widget.value = "daint-mc-test"
+
+            # check the filled value is updated in the filled template
+            assert w.filled_templates[key] == "daint-mc-test"
+
+    # Fill two template variables in one template line
+    for key, value in w._mapping_variables.items():
+        if key == "slurm_partition":
+            sub_widget = value[2]
+            sub_widget.value = "normal-test"
+
+        elif key == "slurm_account":
+            sub_widget = value[2]
+            sub_widget.value = "newuser"
+
+    # check the filled value is updated in the filled template
+    assert (
+        w.filled_templates["prepend_text"]
+        == "#SBATCH --partition=normal-test\n#SBATCH --account=newuser\n#SBATCH --constraint=mc\n#SBATCH --cpus-per-task=1\n\nexport OMP_NUM_THREADS=$SLURM_CPUS_PER_TASK\nsource $MODULESHOME/init/bash\nulimit -s unlimited"
+    )
+
+
 @pytest.mark.usefixtures("aiida_profile_clean")
 def test_quick_setup_widget():
     """Test the QuickSetupWidget."""