Merge pull request #151 from kbase/dev-service

Fix bug in distributing files to containers

Author: MrCreosote

cdmtaskservice/models.py

@@ -416,18 +416,13 @@ class Cluster(str, Enum):
 
 FilesPerContainer = NamedTuple("FilesPerContainer", [
     ("containers", int),
-    ("files_per_container", int),
-    ("last_container", int),
     ("files", list[list[S3File | str]])
 ])
 """
 Information about splitting files to be processed among containers.
 
-containers - the number of containers expected to be run; this is the mimimum of the specified
+containers - the number of containers expected to be run; this is the minimum of the specified
     container count and the number of files.
-files_per_container - the number of files to be run per container
-last_container - the remainder of files to be run in the last container. Always less than
-    files_per_container.
 files - a list of lists of files, split up per container.
 """
 
@@ -609,12 +604,18 @@ def inputs_are_S3File(self) -> bool:
 
     def get_files_per_container(self) -> FilesPerContainer:
         """
-        Returns the number of files to be run per container and the files, split up by container.
+        Returns the container count and the files split up by container.
         """
         containers = min(self.num_containers, len(self.input_files))
-        fpc = math.ceil(len(self.input_files) / containers)
-        files = [self.input_files[i:i + fpc] for i in range(0, fpc * containers, fpc)]
-        return FilesPerContainer(containers, fpc, len(self.input_files) % fpc, files)
+        fpc, extra_files = divmod(len(self.input_files), containers)
+        infiles = list(self.input_files)
+        files = []
+        for i in range(1, containers + 1):
+            # this seems dumb. It works though, make it pretty later
+            fcount = fpc + (1 if i <= extra_files else 0)
+            files.append(infiles[:fcount])
+            infiles = infiles[fcount:]
+        return FilesPerContainer(containers, files)
 
 
 class JobState(str, Enum):

test/models_test.py

@@ -1,7 +1,46 @@
-# TODO TEST add tests
+# TODO TEST add  more tests
 
 from cdmtaskservice import models  # @UnusedImport
 
 
-def test_noop():
-    pass
+def test_files_per_container():
+    _files_per_container(1, 1, 1, [[1]])
+    _files_per_container(1, 2, 1, [[1, 2]])
+    _files_per_container(2, 1, 1, [[1]])
+    _files_per_container(2, 2, 2, [[1], [2]])
+    _files_per_container(3, 2, 2, [[1], [2]])
+    _files_per_container(3, 3, 3, [[1], [2], [3]])
+    _files_per_container(3, 4, 3, [[1, 2], [3], [4]])
+    _files_per_container(3, 5, 3, [[1, 2], [3, 4], [5]])
+    _files_per_container(3, 6, 3, [[1, 2], [3, 4], [5, 6]])
+    _files_per_container(3, 7, 3, [[1, 2, 3], [4, 5], [6, 7]])
+
+
+def _files_per_container(containers: int, files: int, expcont: int, expfiles: list[list[int]]):
+    ji = models.JobInput(
+        cluster=models.Cluster.PERLMUTTER_JAWS,
+        image="fakeimage",
+        params=models.Parameters(),
+        num_containers=containers,
+        input_files=[f"foo/bar{i}" for i in range(1, files + 1)],
+        output_dir="foo/bar"
+    )
+    fpc = ji.get_files_per_container()
+    exp = []
+    for lst in expfiles:
+        exp.append([f"foo/bar{i}" for i in lst])
+    assert fpc == models.FilesPerContainer(expcont, exp)
+    
+    ji = models.JobInput(
+        cluster=models.Cluster.PERLMUTTER_JAWS,
+        image="fakeimage",
+        params=models.Parameters(),
+        num_containers=containers,
+        input_files=[models.S3File(file=f"foo/bar{i}") for i in range(1, files + 1)],
+        output_dir="foo/bar"
+    )
+    fpc = ji.get_files_per_container()
+    exp = []
+    for lst in expfiles:
+        exp.append([models.S3File(file=f"foo/bar{i}") for i in lst])
+    assert fpc == models.FilesPerContainer(expcont, exp)