Print attached files (#352)

Adds methods to form a table of files attached to a Model or to an Ensemble.
Convenience methods to print the table to std out are also added.

[ committed by @al-rigazzi ]
[ reviewed by @MattToast @ankona ]

smartsim/entity/ensemble.py

@@ -29,6 +29,8 @@
 from copy import deepcopy
 from os import getcwd
 
+from tabulate import tabulate
+
 from .._core.utils.helpers import init_default
 from ..error import (
     EntityExistsError,
@@ -275,6 +277,29 @@ def attach_generator_files(
                 to_copy=to_copy, to_symlink=to_symlink, to_configure=to_configure
             )
 
+    @property
+    def attached_files_table(self) -> str:
+        """Return a plain-text table with information about files
+        attached to models belonging to this ensemble.
+
+        :returns: A table of all files attached to all models
+        :rtype: str
+        """
+        if not self.models:
+            return "The ensemble is empty, no files to show."
+
+        table = tabulate(
+            [[model.name, model.attached_files_table] for model in self.models],
+            headers=["Model name", "Files"],
+            tablefmt="grid",
+        )
+
+        return table
+
+    def print_attached_files(self) -> None:
+        """Print table of attached files to std out"""
+        print(self.attached_files_table)
+
     @staticmethod
     def _set_strategy(strategy: str) -> StrategyFunction:
         """Set the permutation strategy for generating models within
@@ -335,7 +360,7 @@ def add_ml_model(
         backend: str,
         model: t.Optional[str] = None,
         model_path: t.Optional[str] = None,
-        device: t.Literal["CPU","GPU"] = "CPU",
+        device: t.Literal["CPU", "GPU"] = "CPU",
         devices_per_node: int = 1,
         batch_size: int = 0,
         min_batch_size: int = 0,
@@ -395,7 +420,7 @@ def add_script(
         name: str,
         script: t.Optional[str] = None,
         script_path: t.Optional[str] = None,
-        device: t.Literal["CPU","GPU"] = "CPU",
+        device: t.Literal["CPU", "GPU"] = "CPU",
         devices_per_node: int = 1,
     ) -> None:
         """TorchScript to launch with every entity belonging to this ensemble
@@ -439,7 +464,7 @@ def add_function(
         self,
         name: str,
         function: t.Optional[str] = None,
-        device: t.Literal["CPU","GPU"] = "CPU",
+        device: t.Literal["CPU", "GPU"] = "CPU",
         devices_per_node: int = 1,
     ) -> None:
         """TorchScript function to launch with every entity belonging to this ensemble

smartsim/entity/files.py

@@ -27,6 +27,7 @@
 import typing as t
 
 from os import path
+from tabulate import tabulate
 
 
 class EntityFiles:
@@ -139,6 +140,22 @@ def _check_path(file_path: str) -> str:
             return full_path
         raise FileNotFoundError(f"File or Directory {file_path} not found")
 
+    def __str__(self) -> str:
+        """Return table summarizing attached files."""
+        values = []
+
+        if self.copy:
+            values.append(["Copy", "\n".join(self.copy)])
+        if self.link:
+            values.append(["Symlink", "\n".join(self.link)])
+        if self.tagged:
+            values.append(["Configure", "\n".join(self.tagged)])
+
+        if not values:
+            return "No file attached to this entity."
+
+        return tabulate(values, headers=["Strategy", "Files"], tablefmt="grid")
+
 
 class TaggedFilesHierarchy:
     """The TaggedFilesHierarchy represents a directory

smartsim/entity/model.py

@@ -165,6 +165,22 @@ def attach_generator_files(
         to_configure = init_default([], to_configure, (list, str))
         self.files = EntityFiles(to_configure, to_copy, to_symlink)
 
+    @property
+    def attached_files_table(self) -> str:
+        """Return a list of attached files as a plain text table
+
+        :returns: String version of table
+        :rtype: str
+        """
+        if not self.files:
+            return "No file attached to this model."
+        return str(self.files)
+
+    def print_attached_files(self) -> None:
+        """Print a table of the attached files on std out
+        """
+        print(self.attached_files_table)
+
     def colocate_db(self, *args: t.Any, **kwargs: t.Any) -> None:
         """An alias for ``Model.colocate_db_tcp``"""
         warnings.warn(

tests/test_configs/to_copy_dir/mock.txt

@@ -0,0 +1 @@
+File mockup.

tests/test_configs/to_symlink_dir/mock2.txt

@@ -0,0 +1 @@
+File mockup two.

tests/test_generator.py

@@ -32,6 +32,9 @@
 from smartsim._core.generation import Generator
 from smartsim.database import Orchestrator
 from smartsim.settings import RunSettings
+from tabulate import tabulate
+
+from smartsim.settings import SbatchSettings
 
 rs = RunSettings("python", exe_args="sleep.py")
 
@@ -108,7 +111,6 @@ def test_ensemble_overwrite_error(fileutils):
 
 
 def test_full_exp(fileutils, wlmutils):
-
     test_dir = fileutils.make_test_dir()
     exp = Experiment("gen-test", test_dir, launcher="local")
 
@@ -148,7 +150,7 @@ def test_dir_files(fileutils):
     params = {"THERMO": [10, 20, 30], "STEPS": [10, 20, 30]}
     ensemble = exp.create_ensemble("dir_test", params=params, run_settings=rs)
     conf_dir = fileutils.get_test_dir_path("test_dir")
-    ensemble.attach_generator_files(to_copy=conf_dir)
+    ensemble.attach_generator_files(to_configure=conf_dir)
 
     exp.generate(ensemble)
 
@@ -160,6 +162,87 @@ def test_dir_files(fileutils):
         assert osp.isfile(osp.join(model_path, "test.py"))
 
 
+def test_print_files(fileutils, capsys):
+    """Test the stdout print of files attached to an ensemble"""
+
+    test_dir = fileutils.make_test_dir()
+    exp = Experiment("print-attached-files-test", test_dir, launcher="local")
+
+    ensemble = exp.create_ensemble("dir_test", replicas=1, run_settings=rs)
+    ensemble.entities = []
+
+    ensemble.print_attached_files()
+    captured = capsys.readouterr()
+    assert captured.out == "The ensemble is empty, no files to show.\n"
+
+    params = {"THERMO": [10, 20], "STEPS": [20, 30]}
+    ensemble = exp.create_ensemble("dir_test", params=params, run_settings=rs)
+    gen_dir = fileutils.get_test_dir_path("test_dir")
+    symlink_dir = fileutils.get_test_dir_path("to_symlink_dir")
+    copy_dir = fileutils.get_test_dir_path("to_copy_dir")
+
+    ensemble.print_attached_files()
+    captured = capsys.readouterr()
+    expected_out = (
+        tabulate(
+            [
+                [model.name, "No file attached to this model."]
+                for model in ensemble.models
+            ],
+            headers=["Model name", "Files"],
+            tablefmt="grid",
+        )
+        + "\n"
+    )
+
+    assert captured.out == expected_out
+
+    ensemble.attach_generator_files()
+    ensemble.print_attached_files()
+    captured = capsys.readouterr()
+    expected_out = (
+        tabulate(
+            [
+                [model.name, "No file attached to this entity."]
+                for model in ensemble.models
+            ],
+            headers=["Model name", "Files"],
+            tablefmt="grid",
+        )
+        + "\n"
+    )
+    assert captured.out == expected_out
+
+    ensemble.attach_generator_files(
+        to_configure=[gen_dir, copy_dir], to_copy=copy_dir, to_symlink=symlink_dir
+    )
+
+    expected_out = tabulate(
+        [
+            ["Copy", copy_dir],
+            ["Symlink", symlink_dir],
+            ["Configure", f"{gen_dir}\n{copy_dir}"],
+        ],
+        headers=["Strategy", "Files"],
+        tablefmt="grid",
+    )
+
+    assert all(str(model.files) == expected_out for model in ensemble.models)
+
+    expected_out_multi = (
+        tabulate(
+            [[model.name, expected_out] for model in ensemble.models],
+            headers=["Model name", "Files"],
+            tablefmt="grid",
+        )
+        + "\n"
+    )
+    ensemble.print_attached_files()
+
+    captured = capsys.readouterr()
+    assert captured.out == expected_out_multi
+
+
 def test_multiple_tags(fileutils):
     """Test substitution of multiple tagged parameters on same line"""
     test_dir = fileutils.make_test_dir()
@@ -198,7 +281,7 @@ def test_config_dir(fileutils):
     gen.generate_experiment(ensemble)
 
     assert osp.isdir(osp.join(test_dir, "test"))
-    # assert False
+
     def _check_generated(test_num, param_0, param_1):
         conf_test_dir = osp.join(test_dir, "test", f"test_{test_num}")
         assert osp.isdir(conf_test_dir)