# Copyright (C) 2015-2021 Regents of the University of California
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
"""A demonstration of toil. Sorts the lines of a file into ascending order by doing a parallel merge sort.
"""
import codecs
import os
import random
import shutil
from configargparse import ArgumentParser
from toil.common import Toil
from toil.job import Job
from toil.realtimeLogger import RealtimeLogger
defaultLines = 1000
defaultLineLen = 50
sortMemory = "600M"
[docs]
def setup(job, inputFile, N, downCheckpoints, options):
"""
Sets up the sort.
Returns the FileID of the sorted file
"""
RealtimeLogger.info("Starting the merge sort")
return job.addChildJobFn(
down,
inputFile,
N,
"root",
downCheckpoints,
options=options,
preemptible=True,
memory=sortMemory,
).rv()
[docs]
def down(job, inputFileStoreID, N, path, downCheckpoints, options, memory=sortMemory):
"""
Input is a file, a subdivision size N, and a path in the hierarchy of jobs.
If the range is larger than a threshold N the range is divided recursively and
a follow on job is then created which merges back the results else
the file is sorted and placed in the output.
"""
RealtimeLogger.info("Down job starting: %s" % path)
# Read the file
inputFile = job.fileStore.readGlobalFile(inputFileStoreID, cache=False)
length = os.path.getsize(inputFile)
if length > N:
# We will subdivide the file
RealtimeLogger.critical(
"Splitting file: %s of size: %s" % (inputFileStoreID, length)
)
# Split the file into two copies
midPoint = getMidPoint(inputFile, 0, length)
t1 = job.fileStore.getLocalTempFile()
with open(t1, "w") as fH:
fH.write(copySubRangeOfFile(inputFile, 0, midPoint + 1))
t2 = job.fileStore.getLocalTempFile()
with open(t2, "w") as fH:
fH.write(copySubRangeOfFile(inputFile, midPoint + 1, length))
# Call down recursively. By giving the rv() of the two jobs as inputs to the follow-on job, up,
# we communicate the dependency without hindering concurrency.
result = job.addFollowOnJobFn(
up,
job.addChildJobFn(
down,
job.fileStore.writeGlobalFile(t1),
N,
path + "/0",
downCheckpoints,
checkpoint=downCheckpoints,
options=options,
preemptible=True,
memory=options.sortMemory,
).rv(),
job.addChildJobFn(
down,
job.fileStore.writeGlobalFile(t2),
N,
path + "/1",
downCheckpoints,
checkpoint=downCheckpoints,
options=options,
preemptible=True,
memory=options.mergeMemory,
).rv(),
path + "/up",
preemptible=True,
options=options,
memory=options.sortMemory,
).rv()
else:
# We can sort this bit of the file
RealtimeLogger.critical(
"Sorting file: %s of size: %s" % (inputFileStoreID, length)
)
# Sort the copy and write back to the fileStore
shutil.copyfile(inputFile, inputFile + ".sort")
sort(inputFile + ".sort")
result = job.fileStore.writeGlobalFile(inputFile + ".sort")
RealtimeLogger.info("Down job finished: %s" % path)
return result
[docs]
def up(job, inputFileID1, inputFileID2, path, options, memory=sortMemory):
"""
Merges the two files and places them in the output.
"""
RealtimeLogger.info("Up job starting: %s" % path)
with job.fileStore.writeGlobalFileStream() as (fileHandle, outputFileStoreID):
fileHandle = codecs.getwriter("utf-8")(fileHandle)
with job.fileStore.readGlobalFileStream(inputFileID1) as inputFileHandle1:
inputFileHandle1 = codecs.getreader("utf-8")(inputFileHandle1)
with job.fileStore.readGlobalFileStream(inputFileID2) as inputFileHandle2:
inputFileHandle2 = codecs.getreader("utf-8")(inputFileHandle2)
RealtimeLogger.info(
"Merging %s and %s to %s"
% (inputFileID1, inputFileID2, outputFileStoreID)
)
merge(inputFileHandle1, inputFileHandle2, fileHandle)
# Cleanup up the input files - these deletes will occur after the completion is successful.
job.fileStore.deleteGlobalFile(inputFileID1)
job.fileStore.deleteGlobalFile(inputFileID2)
RealtimeLogger.info("Up job finished: %s" % path)
return outputFileStoreID
[docs]
def sort(file):
"""Sorts the given file."""
with open(file) as f:
lines = f.readlines()
lines.sort()
with open(file, "w") as f:
for line in lines:
f.write(line)
[docs]
def merge(fileHandle1, fileHandle2, outputFileHandle):
"""
Merges together two files maintaining sorted order.
All handles must be text-mode streams.
"""
line2 = fileHandle2.readline()
for line1 in fileHandle1.readlines():
while len(line2) != 0 and line2 <= line1:
outputFileHandle.write(line2)
line2 = fileHandle2.readline()
outputFileHandle.write(line1)
while len(line2) != 0:
outputFileHandle.write(line2)
line2 = fileHandle2.readline()
[docs]
def copySubRangeOfFile(inputFile, fileStart, fileEnd):
"""
Copies the range (in bytes) between fileStart and fileEnd to the given
output file handle.
"""
with open(inputFile) as fileHandle:
fileHandle.seek(fileStart)
data = fileHandle.read(fileEnd - fileStart)
assert len(data) == fileEnd - fileStart
return data
[docs]
def getMidPoint(file, fileStart, fileEnd):
"""
Finds the point in the file to split.
Returns an int i such that fileStart <= i < fileEnd
"""
with open(file) as f:
midPoint = (fileStart + fileEnd) // 2
assert midPoint >= fileStart
f.seek(midPoint)
line = f.readline()
assert len(line) >= 1
if len(line) + midPoint < fileEnd:
return midPoint + len(line) - 1
f.seek(fileStart)
line = f.readline()
assert len(line) >= 1
assert len(line) + fileStart <= fileEnd
return len(line) + fileStart - 1
[docs]
def makeFileToSort(fileName, lines=defaultLines, lineLen=defaultLineLen):
with open(fileName, "w") as f:
for _ in range(lines):
line = (
"".join(random.choice("actgACTGNXYZ") for _ in range(lineLen - 1))
+ "\n"
)
f.write(line)
[docs]
def main(options=None):
if not options:
# deal with command line arguments
parser = ArgumentParser()
Job.Runner.addToilOptions(parser)
parser.add_argument(
"--numLines",
default=defaultLines,
help="Number of lines in file to sort.",
type=int,
)
parser.add_argument(
"--lineLength",
default=defaultLineLen,
help="Length of lines in file to sort.",
type=int,
)
parser.add_argument("--fileToSort", help="The file you wish to sort")
parser.add_argument("--outputFile", help="Where the sorted output will go")
parser.add_argument(
"--overwriteOutput",
help="Write over the output file if it already exists.",
default=True,
)
parser.add_argument(
"--N",
dest="N",
help="The threshold below which a serial sort function is used to sort file. "
"All lines must of length less than or equal to N or program will fail",
default=10000,
)
parser.add_argument(
"--downCheckpoints",
action="store_true",
help="If this option is set, the workflow will make checkpoints on its way through"
'the recursive "down" part of the sort',
)
parser.add_argument(
"--sortMemory",
dest="sortMemory",
help="Memory for jobs that sort chunks of the file.",
default=None,
)
parser.add_argument(
"--mergeMemory",
dest="mergeMemory",
help="Memory for jobs that collate results.",
default=None,
)
options = parser.parse_args()
if not hasattr(options, "sortMemory") or not options.sortMemory:
options.sortMemory = sortMemory
if not hasattr(options, "mergeMemory") or not options.mergeMemory:
options.mergeMemory = sortMemory
# do some input verification
sortedFileName = options.outputFile or "sortedFile.txt"
if not options.overwriteOutput and os.path.exists(sortedFileName):
print(
f"Output file {sortedFileName} already exists. "
f"Delete it to run the sort example again or use --overwriteOutput=True"
)
exit()
fileName = options.fileToSort
if options.fileToSort is None:
# make the file ourselves
fileName = "fileToSort.txt"
if os.path.exists(fileName):
print(f"Sorting existing file: {fileName}")
else:
print(
f"No sort file specified. Generating one automatically called: {fileName}."
)
makeFileToSort(
fileName=fileName, lines=options.numLines, lineLen=options.lineLength
)
else:
if not os.path.exists(options.fileToSort):
raise RuntimeError("File to sort does not exist: %s" % options.fileToSort)
if int(options.N) <= 0:
raise RuntimeError("Invalid value of N: %s" % options.N)
# Now we are ready to run
with Toil(options) as workflow:
sortedFileURL = "file://" + os.path.abspath(sortedFileName)
if not workflow.options.restart:
sortFileURL = "file://" + os.path.abspath(fileName)
sortFileID = workflow.importFile(sortFileURL)
sortedFileID = workflow.start(
Job.wrapJobFn(
setup,
sortFileID,
int(options.N),
options.downCheckpoints,
options=options,
memory=sortMemory,
)
)
else:
sortedFileID = workflow.restart()
workflow.exportFile(sortedFileID, sortedFileURL)
if __name__ == "__main__":
main()
