Improve provenance for Omega PRs

Author: xylar

polaris/provenance.py

@@ -3,7 +3,7 @@
 import sys
 
 
-def write(work_dir, tasks, config=None):
+def write(work_dir, tasks, config=None, machine=None, baseline_dir=None):
     """
     Write a file with provenance, such as the git version, conda packages,
     command, and tasks, to the work directory
@@ -16,9 +16,15 @@ def write(work_dir, tasks, config=None):
     tasks : dict
         A dictionary describing all of the tasks and their steps
 
-    config : polaris.config.PolarisConfigParser
+    config : polaris.config.PolarisConfigParser, optional
         Configuration options for this task, a combination of user configs
         and the defaults for the machine and component
+
+    machine : str, optional
+        The machine on which Polaris is being run
+
+    baseline_dir : str, optional
+        The path to the baseline work directory, if any
     """
     polaris_git_version = None
     if os.path.exists('.git'):
@@ -69,6 +75,14 @@ def write(work_dir, tasks, config=None):
             f'component git version: {component_git_version}\n\n'
         )
     provenance_file.write(f'command: {calling_command}\n\n')
+
+    # Add readily parsable, PR-friendly metadata discovered at setup time
+    _write_meta(provenance_file, 'machine', machine)
+    _write_scheduler_metadata(provenance_file, config)
+    _write_meta(provenance_file, 'compiler', _get_compiler(config))
+    _write_meta(provenance_file, 'work directory', work_dir)
+    _write_meta(provenance_file, 'build directory', _get_build_dir(config))
+    _write_meta(provenance_file, 'baseline work directory', baseline_dir)
     provenance_file.write('tasks:\n')
 
     for _, task in tasks.items():
@@ -126,3 +140,129 @@ def _get_component_git_version(config):
     os.chdir(cwd)
 
     return component_git_version
+
+
+def _get_system(config):
+    if config is None:
+        return None
+    if config.has_option('parallel', 'system'):
+        return config.get('parallel', 'system')
+    return None
+
+
+def _resolve_scheduler_fields(config, system):
+    """Resolve partition/qos/constraint (slurm) or just constraint (pbs)."""
+    partition = None
+    qos = None
+    constraint = None
+
+    if config is None:
+        return partition, qos, constraint
+
+    if system == 'pbs':
+        # PBS: only constraint is relevant for our purposes
+        cons_val = (
+            config.get('job', 'constraint')
+            if config.has_option('job', 'constraint')
+            else None
+        )
+        if cons_val and cons_val != '<<<default>>>':
+            constraint = cons_val
+        elif config.has_option('parallel', 'constraints'):
+            cons_list = config.getlist('parallel', 'constraints')
+            if cons_list:
+                constraint = cons_list[0]
+        return partition, qos, constraint
+
+    # Default to slurm-like resolution
+    part_val = (
+        config.get('job', 'partition')
+        if config.has_option('job', 'partition')
+        else None
+    )
+    if part_val and part_val != '<<<default>>>':
+        partition = part_val
+    elif config.has_option('parallel', 'partitions'):
+        parts = config.getlist('parallel', 'partitions')
+        if parts:
+            partition = parts[0]
+
+    qos_val = (
+        config.get('job', 'qos') if config.has_option('job', 'qos') else None
+    )
+    if qos_val and qos_val != '<<<default>>>':
+        qos = qos_val
+    elif config.has_option('parallel', 'qos'):
+        qos_list = config.getlist('parallel', 'qos')
+        if qos_list:
+            qos = qos_list[0]
+
+    cons_val = (
+        config.get('job', 'constraint')
+        if config.has_option('job', 'constraint')
+        else None
+    )
+    if cons_val and cons_val != '<<<default>>>':
+        constraint = cons_val
+    elif config.has_option('parallel', 'constraints'):
+        cons_list = config.getlist('parallel', 'constraints')
+        if cons_list:
+            constraint = cons_list[0]
+
+    return partition, qos, constraint
+
+
+def _get_compiler(config):
+    if config is None:
+        return None
+    if config.has_option('deploy', 'compiler'):
+        val = config.get('deploy', 'compiler')
+        return val or None
+    return None
+
+
+def _get_build_dir(config):
+    if config is None:
+        return None
+    if config.has_option('paths', 'component_path'):
+        val = config.get('paths', 'component_path')
+        return val or None
+    return None
+
+
+def _write_meta(provenance_file, label, value):
+    """Write a simple 'label: value' line if value is provided."""
+    if value is None:
+        return
+    if isinstance(value, str) and value.strip() == '':
+        return
+    provenance_file.write(f'{label}: {value}\n\n')
+
+
+def _write_scheduler_metadata(provenance_file, config):
+    """Write partition/qos/constraint metadata when available."""
+    system = _get_system(config)
+    partition, qos, constraint = _resolve_scheduler_fields(config, system)
+    _write_meta(provenance_file, 'partition', partition)
+    _write_meta(provenance_file, 'qos', qos)
+    _write_meta(provenance_file, 'constraint', constraint)
+    if config.has_option('paths', 'component_path'):
+        component_path = config.get('paths', 'component_path')
+    else:
+        component_path = None
+
+    if component_path is None or not os.path.exists(component_path):
+        return None
+
+    cwd = os.getcwd()
+    os.chdir(component_path)
+
+    try:
+        args = ['git', 'describe', '--tags', '--dirty', '--always']
+        component_git_version = subprocess.check_output(args).decode('utf-8')
+        component_git_version = component_git_version.strip('\n')
+    except subprocess.CalledProcessError:
+        component_git_version = None
+    os.chdir(cwd)
+
+    return component_git_version

polaris/setup.py

@@ -115,7 +115,13 @@ def setup_tasks(
         config_file, machine, component_path, component, model
     )
 
-    provenance.write(work_dir, tasks, config=basic_config)
+    provenance.write(
+        work_dir,
+        tasks,
+        config=basic_config,
+        machine=machine,
+        baseline_dir=baseline_dir,
+    )
 
     if clean:
         print('')