fix: better peek handling for fastx records

Author: sstadick

ishlib/fastx_search_runner.mojo

@@ -21,19 +21,20 @@ struct FastxSearchRunner[M: Matcher]:
     fn run_search[
         W: MovableWriter
     ](mut self, mut writer: BufferedWriter[W]) raises:
-        # Simple thing first?
+        # Peek only the first file to determine fastq or not, assume not-binary
+        var peek = peek_file[record_type = RecordType.FASTX](
+            self.settings.files[0]
+        )
         for file in self.settings.files:
             var f = file[]  # force copy
-            var peek = peek_file[record_type = RecordType.FASTX](f)
-            if peek.is_binary:
-                if self.settings.verbose:
-                    Logger.warn("Skipping binary file:", file[])
-                continue
             Logger.debug("Processing", f)
-            self.run_search_on_file(f, writer)
+            if peek.is_fastq:
+                self.run_search_on_file[is_fastq=True](f, writer)
+            else:
+                self.run_search_on_file[is_fastq=False](f, writer)
 
     fn run_search_on_file[
-        W: MovableWriter
+        W: MovableWriter, *, is_fastq: Bool = False
     ](mut self, file: Path, mut writer: BufferedWriter[W]) raises:
         var reader = FastxReader[read_comment=False](
             BufferedReader(GZFile(String(file), "r"))
@@ -74,7 +75,8 @@ struct FastxSearchRunner[M: Matcher]:
                     writer.write_bytes(reader.seq.as_span())
 
                 # Handle FASTQ Case
-                if len(reader.qual.as_span()) > 0:
+                @parameter
+                if is_fastq:
                     writer.write("\n+\n")
                     writer.write_bytes(reader.qual.as_span())
                 writer.write("\n")

ishlib/matcher/striped_semi_global_matcher.mojo

@@ -223,6 +223,28 @@ struct StripedSemiGlobalMatcher(GpuMatcher):
                 ends_free.target_start,
                 ends_free.target_end,
             )
+        elif matrix_kind == MatrixKind.ACTGN0:
+            gpu_align_coarse[
+                MatrixKind.ACTGN0, max_query_length, max_target_length
+            ](
+                query,
+                ref_buffer,
+                target_ends,
+                score_result_buffer,
+                query_end_result_buffer,
+                ref_end_result_buffer,
+                basic_matrix_values,
+                basic_matrix_len,
+                query_len,
+                target_ends_len,
+                thread_count,
+                gap_open,
+                gap_extend,
+                ends_free.query_start,
+                ends_free.query_end,
+                ends_free.target_start,
+                ends_free.target_end,
+            )
         elif matrix_kind == MatrixKind.BLOSUM62:
             gpu_align_coarse[
                 MatrixKind.BLOSUM62, max_query_length, max_target_length

ishlib/parallel_fastx_search_runner.mojo

@@ -30,7 +30,7 @@ from ishlib.matcher import (
     WhereComputed,
 )
 from ishlib import ByteSpanWriter, RecordType
-from ishlib.peek_file import peek_file
+from ishlib.peek_file import peek_file, PeekFindings
 from ishlib.matcher.alignment.striped_utils import AlignmentResult
 from ishlib.vendor.kseq import BufferedReader, FastxReader
 from ishlib.vendor.zlib import GZFile
@@ -64,15 +64,12 @@ struct ParallelFastxSearchRunner[M: Matcher]:
     fn run_search[
         W: MovableWriter
     ](mut self, mut writer: BufferedWriter[W]) raises:
-        # Simple thing first?
+        # Peek only the first file to determine fastq or not, assume not-binary
+        var peek = peek_file[record_type = RecordType.FASTX](
+            self.settings.files[0]
+        )
         for file in self.settings.files:
             var f = file[]  # force copy
-            var peek = peek_file[record_type = RecordType.FASTX](f)
-            Logger.debug("Suggested length:", peek.suggested_max_length)
-            if peek.is_binary:
-                if self.settings.verbose:
-                    Logger.warn("Skipping binary file:", file[])
-                continue
             Logger.debug("Processing", f)
             if peek.is_fastq:
                 self.run_search_on_file[is_fastq=True](f, writer)
@@ -204,51 +201,39 @@ struct GpuParallelFastxSearchRunner[
         W: MovableWriter
     ](mut self, mut writer: BufferedWriter[W]) raises:
         # Peek the first file to get the suggested size, then use that for all of them.
-        # Still peek each for binary
-
-        # First non-binary file
+        # Assume non-binary
         var files = self.settings.files
         var first_peek = peek_file[
             record_type = RecordType.FASTX, check_record_size=True
         ](files[0])
-        if first_peek.is_binary:
-            for i in range(1, len(files)):
-                first_peek = peek_file[
-                    record_type = RecordType.FASTX, check_record_size=True
-                ](files[i])
 
         Logger.debug("Suggested length of:", first_peek.suggested_max_length)
 
         # Create ctxs
-        @parameter
-        @always_inline
-        fn choose_max_target_length(suggested_max_length: Int) raises:
-            alias MAX_TARGET_LENGTHS = List(128, 256, 512, 1024, 2048, 4096)
+        alias MAX_TARGET_LENGTHS = List(128, 256, 512, 1024, 2048, 4096)
 
-            @parameter
-            for i in range(0, len(MAX_TARGET_LENGTHS)):
-                alias max_target_length = MAX_TARGET_LENGTHS[i]
-                if suggested_max_length <= max_target_length:
-                    var ctxs = self.create_ctxs[
-                        max_query_length, max_target_length
-                    ]()
-                    self.search_files[
-                        W,
-                        max_query_length=max_query_length,
-                        max_target_length=max_target_length,
-                    ](files, ctxs, writer)
-                    return
-            else:
-                Logger.warn(
-                    "Longer line lengths than supported, more work will"
-                    " be sent to CPU, consider running with max-gpus set to 0."
-                )
-                var ctxs = self.create_ctxs[max_query_length, 4096]()
+        @parameter
+        for i in range(0, len(MAX_TARGET_LENGTHS)):
+            alias max_target_length = MAX_TARGET_LENGTHS[i]
+            if first_peek.suggested_max_length <= max_target_length:
+                var ctxs = self.create_ctxs[
+                    max_query_length, max_target_length
+                ]()
                 self.search_files[
-                    W, max_query_length=max_query_length, max_target_length=4096
-                ](files, ctxs, writer)
-
-        choose_max_target_length(first_peek.suggested_max_length)
+                    W,
+                    max_query_length=max_query_length,
+                    max_target_length=max_target_length,
+                ](files, ctxs, writer, first_peek)
+                return
+        else:
+            Logger.warn(
+                "Longer line lengths than supported, more work will"
+                " be sent to CPU, consider running with max-gpus set to 0."
+            )
+            var ctxs = self.create_ctxs[max_query_length, 4096]()
+            self.search_files[
+                W, max_query_length=max_query_length, max_target_length=4096
+            ](files, ctxs, writer, first_peek)
 
     fn create_ctxs[
         max_query_length: UInt = 200, max_target_length: UInt = 1024
@@ -281,27 +266,21 @@ struct GpuParallelFastxSearchRunner[
             ]
         ],
         mut writer: BufferedWriter[W],
+        read peek: PeekFindings,
     ) raises:
         for i in range(0, len(paths)):
             var f = paths[i]  # force copy
             Logger.debug("Processing", f)
-            if i > 0:
-                # Can skip the first peek since we've already checked it.
-                var peek = peek_file[
-                    record_type = RecordType.FASTX, check_record_size=False
-                ](f)
-                if peek.is_binary:
-                    if self.settings.verbose:
-                        Logger.warn("Skipping binary file:", paths[i])
-                    continue
-                if peek.is_fastq:
-                    self.run_search_on_file[
-                        W, max_query_length, max_target_length, is_fastq=True
-                    ](f, ctxs, writer)
-                else:
-                    self.run_search_on_file[
-                        W, max_query_length, max_target_length, is_fastq=False
-                    ](f, ctxs, writer)
+            # Assume that since the user specified FASTX, there are no binary files
+            # Assume that all files are the same record type
+            if peek.is_fastq:
+                self.run_search_on_file[
+                    W, max_query_length, max_target_length, is_fastq=True
+                ](f, ctxs, writer)
+            else:
+                self.run_search_on_file[
+                    W, max_query_length, max_target_length, is_fastq=False
+                ](f, ctxs, writer)
 
     fn run_search_on_file[
         W: MovableWriter,

tests/test_searcher.mojo

@@ -1,7 +1,6 @@
 from pathlib import Path
 from testing import assert_equal
 from tempfile import TemporaryDirectory
-from utils import StringSlice
 
 from ExtraMojo.io.buffered import BufferedWriter
 
@@ -166,7 +165,13 @@ ACTGACTGACGACGACGACTAATAGNNNNACTGANNNATCATCTAGGGGGGGGGGGGGGGGGGGGGGGGGGGGGGGGGGG
         var writer = BufferedWriter(open(output, "w"))
         do_search(setting[], writer^)
         var found = output.read_text()
-        assert_equal(found.upper(), expected.upper())
+        assert_equal(
+            found.upper(),
+            expected.upper(),
+            String("max_gpus {}, num_threads {}, algo {}").format(
+                setting[].max_gpus, setting[].threads, setting[].match_algo
+            ),
+        )
 
 
 def test_fasta_blosum62_search():
@@ -216,7 +221,13 @@ MAFSAEDVLKEYDRRRRMEALLLSLYYPNDRKLLDYKEWSPPRVQVECPKAPVEWNNPPSEKGLIVGHFSGIKYKGEKAQ
     for setting in settings:
         var writer = BufferedWriter(open(output, "w"))
         do_search(setting[], writer^)
-        assert_equal(output.read_text().upper(), expected.upper())
+        assert_equal(
+            output.read_text().upper(),
+            expected.upper(),
+            String("max_gpus {}, num_threads {}, algo {}").format(
+                setting[].max_gpus, setting[].threads, setting[].match_algo
+            ),
+        )
 
 
 def main():