dli-mkshards

#!/usr/bin/python3
#
# Copyright (c) 2017 NVIDIA CORPORATION. All rights reserved.
#
# 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.


import argparse
import random
import sys

from dlinputs import gopen, tarrecords, utils

parser = argparse.ArgumentParser("""\
Bundle files into sharded tarfiles.

The list of files is given on the standard input.

The following will find all `.png` and `.cls` files in the current
directory tree, tar them up in shards of size 1 Gbyte, and transfer
them via `ssh` to your home directory on host `mvcdev2` with
names like `myimages-000000.tgz`, ...:

```
find . -name '*.png' -o -name '*.cls' | sort | tarshards -S mvcdev2: myimages
```

For ssh destinations, you should have set up a password-less login.
""")
parser.add_argument("-M", "--maxcount", type=float, default=1e6,
                    help="size of shards to be uploaded")
parser.add_argument("-S", "--maxsize", type=float, default=1e9,
                    help="size of shards to be uploaded")
parser.add_argument("-p", "--progress", action="store_true",
                    help="display progress while uploading")
parser.add_argument("-s", "--shuffle", action="store_true",
                    help="shuffle before writing")
parser.add_argument("output",
                    help="output shard format")
args = parser.parse_args()

assert "%" in args.output

all_files = sys.stdin.readlines()
all_files = [fname.strip() for fname in all_files]

all_records = {}
for fname in all_files:
    base, ext = path.splitallext(fname)
    all_records[base] = all_records.get(base, []) + [fname]
for base in all_records.keys():
    all_records[base] = sorted(all_records[base])


extension_sets = [{path.splitallext(x)[1] for x in bucket} for bucket in all_records.values()]
counter = collections.Counter(extension_sets)
print extension_sets.most_common(20)

keys = sorted(all_records.keys())
if args.shuffle:
    random.shuffle(keys)

sink = tarrecords.ShardWriter(args.output, maxcount=args.maxcount, maxsize=args.maxsize)

for i, key in enumerate(keys):
    bucket = all_records[key]
    sample = [readfile(fname) for fname in bucket]
    sink.write(sample)

sink.close()