Improved compatiblity with T4 format, added path support to filenames

Author: fjwillemsen

kernel_tuner/file_utils.py

@@ -4,6 +4,8 @@
 import json
 import subprocess
 import xmltodict
+from sys import platform
+from pathlib import Path
 
 from importlib.metadata import requires, version, PackageNotFoundError
 from packaging.requirements import Requirement
@@ -14,7 +16,7 @@
 
 
 def output_file_schema(target):
-    """ Get the requested JSON schema and the version number
+    """Get the requested JSON schema and the version number
 
     :param target: Name of the T4 schema to return, should be any of ['output', 'metadata']
     :type target: string
@@ -25,13 +27,13 @@ def output_file_schema(target):
     """
     current_version = "1.0.0"
     output_file = schema_dir + f"/T4/{current_version}/{target}-schema.json"
-    with open(output_file, 'r') as fh:
+    with open(output_file, "r") as fh:
         json_string = json.load(fh)
     return current_version, json_string
 
 
 def get_configuration_validity(objective) -> str:
-    """ Convert internal Kernel Tuner error to string """
+    """Convert internal Kernel Tuner error to string"""
     errorstring: str
     if not isinstance(objective, util.ErrorConfig):
         errorstring = "correct"
@@ -40,24 +42,33 @@ def get_configuration_validity(objective) -> str:
             errorstring = "compile"
         elif isinstance(objective, util.RuntimeFailedConfig):
             errorstring = "runtime"
-        else:
+        elif isinstance(objective, util.InvalidConfig):
             errorstring = "constraints"
+        else:
+            raise ValueError(f"Unkown objective type {type(objective)}, value {objective}")
     return errorstring
 
 
 def filename_ensure_json_extension(filename: str) -> str:
-    """ Check if the filename has a .json extension, if not, add it """
+    """Check if the filename has a .json extension, if not, add it"""
     if filename[-5:] != ".json":
         filename += ".json"
     return filename
 
 
-def store_output_file(output_filename, results, tune_params, objective="time"):
-    """ Store the obtained auto-tuning results in a JSON output file
+def make_filenamepath(filenamepath: Path):
+    """Create the given path to a filename if the path does not yet exist"""
+    filepath = filenamepath.parents[0]
+    if not filepath.exists():
+        filepath.mkdir()
+
+
+def store_output_file(output_filename: str, results, tune_params, objective="time"):
+    """Store the obtained auto-tuning results in a JSON output file
 
     This function produces a JSON file that adheres to the T4 auto-tuning output JSON schema.
 
-    :param output_filename: Name of the to be created output file
+    :param output_filename: Name or 'path / name' of the to be created output file
     :type output_filename: string
 
     :param results: Results list as return by tune_kernel
@@ -70,26 +81,20 @@ def store_output_file(output_filename, results, tune_params, objective="time"):
     :type objective: string
 
     """
-    output_filename = filename_ensure_json_extension(output_filename)
+    output_filenamepath = Path(filename_ensure_json_extension(output_filename))
+    make_filenamepath(output_filenamepath)
 
-    timing_keys = [
-        "compile_time", "benchmark_time", "framework_time", "strategy_time",
-        "verification_time"
-    ]
-    not_measurement_keys = list(
-        tune_params.keys()) + timing_keys + ["timestamp"] + ["times"]
+    timing_keys = ["compile_time", "benchmark_time", "framework_time", "strategy_time", "verification_time"]
+    not_measurement_keys = list(tune_params.keys()) + timing_keys + ["timestamp"] + ["times"]
 
     output_data = []
 
     for result in results:
-
         out = {}
 
-        out["timestamp"] = result["timestamp"]
-        out["configuration"] = {
-            k: v
-            for k, v in result.items() if k in tune_params
-        }
+        if "timestamp" in result:
+            out["timestamp"] = result["timestamp"]
+        out["configuration"] = {k: v for k, v in result.items() if k in tune_params}
 
         # collect configuration specific timings
         timings = dict()
@@ -98,7 +103,8 @@ def store_output_file(output_filename, results, tune_params, objective="time"):
         timings["framework"] = result["framework_time"]
         timings["search_algorithm"] = result["strategy_time"]
         timings["validation"] = result["verification_time"]
-        timings["runtimes"] = result["times"]
+        if "times" in result:
+            timings["runtimes"] = result["times"]
         out["times"] = timings
 
         # encode the validity of the configuration
@@ -112,10 +118,7 @@ def store_output_file(output_filename, results, tune_params, objective="time"):
         measurements = []
         for key, value in result.items():
             if key not in not_measurement_keys:
-                measurements.append(
-                    dict(name=key,
-                         value=value,
-                         unit="ms" if key.startswith("time") else ""))
+                measurements.append(dict(name=key, value=value, unit="ms" if key.startswith("time") else ""))
         out["measurements"] = measurements
 
         # objectives
@@ -130,12 +133,12 @@ def store_output_file(output_filename, results, tune_params, objective="time"):
     # write output_data to a JSON file
     version, _ = output_file_schema("results")
     output_json = dict(results=output_data, schema_version=version)
-    with open(output_filename, 'w+') as fh:
+    with open(output_filenamepath, "w+") as fh:
         json.dump(output_json, fh)
 
 
-def get_dependencies(package='kernel_tuner'):
-    """ Get the Python dependencies of Kernel Tuner currently installed and their version numbers """
+def get_dependencies(package="kernel_tuner"):
+    """Get the Python dependencies of Kernel Tuner currently installed and their version numbers"""
     requirements = requires(package)
     deps = [Requirement(req).name for req in requirements]
     depends = []
@@ -150,10 +153,9 @@ def get_dependencies(package='kernel_tuner'):
 
 
 def get_device_query(target):
-    """ Get the information about GPUs in the current system, target is any of ['nvidia', 'amd'] """
+    """Get the information about GPUs in the current system, target is any of ['nvidia', 'amd']"""
     if target == "nvidia":
-        nvidia_smi_out = subprocess.run(["nvidia-smi", "--query", "-x"],
-                                        capture_output=True)
+        nvidia_smi_out = subprocess.run(["nvidia-smi", "--query", "-x"], capture_output=True)
         nvidia_smi = xmltodict.parse(nvidia_smi_out.stdout)
         gpu_info = nvidia_smi["nvidia_smi_log"]["gpu"]
         del_key = "processes"
@@ -162,57 +164,81 @@ def get_device_query(target):
             for gpu in gpu_info:
                 del gpu[del_key]
         elif isinstance(gpu_info, dict) and del_key in gpu_info:
-                del gpu_info[del_key]
+            del gpu_info[del_key]
         return nvidia_smi
     elif target == "amd":
-        rocm_smi_out = subprocess.run(["rocm-smi", "--showallinfo", "--json"],
-                                      capture_output=True)
+        rocm_smi_out = subprocess.run(["rocm-smi", "--showallinfo", "--json"], capture_output=True)
         return json.loads(rocm_smi_out.stdout)
     else:
         raise ValueError("get_device_query target not supported")
 
 
-def store_metadata_file(metadata_filename):
-    """ Store the metadata about the current hardware and software environment in a JSON output file
+def store_metadata_file(metadata_filename: str):
+    """Store the metadata about the current hardware and software environment in a JSON output file
 
     This function produces a JSON file that adheres to the T4 auto-tuning metadata JSON schema.
 
-    :param metadata_filename: Name of the to be created metadata file
+    :param metadata_filename: Name or 'path / name' of the to be created metadata file
     :type metadata_filename: string
 
     """
-    metadata_filename = filename_ensure_json_extension(metadata_filename)
+    metadata_filenamepath = Path(filename_ensure_json_extension(metadata_filename))
+    make_filenamepath(metadata_filenamepath)
     metadata = {}
 
-    # lshw only works on Linux, this intentionally raises a FileNotFoundError when ran on systems that do not have it
-    lshw_out = subprocess.run(["lshw", "-json"], capture_output=True)
-
-    # sometimes lshw outputs a list of length 1, sometimes just as a dict, schema wants a list
-    lshw_string = lshw_out.stdout.decode('utf-8').strip()
-    if lshw_string[0] == '{' and lshw_string[-1] == '}':
-        lshw_string = '[' + lshw_string + ']'
+    # differentiate between OSes, possible values: https://docs.python.org/3/library/sys.html#sys.platform
+    if platform == "linux":
+        os_string = "Linux"
+        hardware_description_out = subprocess.run(["lshw", "-json"], capture_output=True)
+    elif platform == "win32":
+        os_string = "Windows"
+        raise NotImplementedError(f"Hardware specification not yet implemented for Windows")
+    elif platform == "darwin":
+        os_string = "Mac"
+        hardware_description_out = subprocess.run(
+            [
+                "system_profiler",
+                "-json",
+                "-detailLevel",
+                "mini",
+                "SPSoftwareDataType",
+                "SPHardwareDataType",
+                "SPiBridgeDataType",
+                "SPPCIDataType",
+                "SPMemoryDataType",
+                "SPNVMeDataType",
+            ],
+            capture_output=True,
+        )
+    else:
+        raise ValueError(f"Platform {platform} not supported for metadata collection")
 
-    metadata["hardware"] = dict(lshw=json.loads(lshw_string))
+    # process the hardware description output
+    hardware_description_string = hardware_description_out.stdout.decode("utf-8").strip()
+    if hardware_description_string[0] == "{" and hardware_description_string[-1] == "}":
+        # sometimes lshw outputs a list of length 1, sometimes just as a dict, schema wants a list
+        hardware_description_string = "[" + hardware_description_string + "]"
+    metadata["hardware"] = dict(hardware_description=json.loads(hardware_description_string))
+    metadata["operating_system"] = os_string
 
     # attempts to use nvidia-smi or rocm-smi if present
     device_query = {}
     try:
-        device_query['nvidia-smi'] = get_device_query("nvidia")
+        device_query["nvidia-smi"] = get_device_query("nvidia")
     except FileNotFoundError:
         # ignore if nvidia-smi is not found
         pass
 
     try:
-        device_query['rocm-smi'] = get_device_query("amd")
+        device_query["rocm-smi"] = get_device_query("amd")
     except FileNotFoundError:
         # ignore if rocm-smi is not found
         pass
 
-    metadata["environment"] = dict(device_query=device_query,
-                                   requirements=get_dependencies())
+    metadata["environment"] = dict(device_query=device_query, requirements=get_dependencies())
 
     # write metadata to JSON file
     version, _ = output_file_schema("metadata")
     metadata_json = dict(metadata=metadata, schema_version=version)
-    with open(metadata_filename, 'w+') as fh:
+    with open(metadata_filenamepath, "w+") as fh:
         json.dump(metadata_json, fh, indent="  ")