4 # --md5sum - display md5 of each file as read from disk
6 import apiclient.errors
22 import arvados.commands._util as arv_cmd
24 CAUGHT_SIGNALS = [signal.SIGINT, signal.SIGQUIT, signal.SIGTERM]
27 upload_opts = argparse.ArgumentParser(add_help=False)
29 upload_opts.add_argument('paths', metavar='path', type=str, nargs='*',
31 Local file or directory. Default: read from standard input.
34 upload_opts.add_argument('--max-manifest-depth', type=int, metavar='N',
36 Maximum depth of directory tree to represent in the manifest
37 structure. A directory structure deeper than this will be represented
38 as a single stream in the manifest. If N=0, the manifest will contain
39 a single stream. Default: -1 (unlimited), i.e., exactly one manifest
40 stream per filesystem directory that contains files.
43 upload_opts.add_argument('--project-uuid', metavar='UUID', help="""
44 Store the collection in the specified project, instead of your Home
48 upload_opts.add_argument('--name', help="""
49 Save the collection with the specified name, rather than the default
50 generic name "Saved at {time} by {username}@{host}".
53 _group = upload_opts.add_mutually_exclusive_group()
55 _group.add_argument('--as-stream', action='store_true', dest='stream',
60 _group.add_argument('--stream', action='store_true',
62 Store the file content and display the resulting manifest on
63 stdout. Do not write the manifest to Keep or save a Collection object
67 _group.add_argument('--as-manifest', action='store_true', dest='manifest',
69 Synonym for --manifest.
72 _group.add_argument('--in-manifest', action='store_true', dest='manifest',
74 Synonym for --manifest.
77 _group.add_argument('--manifest', action='store_true',
79 Store the file data and resulting manifest in Keep, save a Collection
80 object in Arvados, and display the manifest locator (Collection uuid)
81 on stdout. This is the default behavior.
84 _group.add_argument('--as-raw', action='store_true', dest='raw',
89 _group.add_argument('--raw', action='store_true',
91 Store the file content and display the data block locators on stdout,
92 separated by commas, with a trailing newline. Do not store a
96 upload_opts.add_argument('--use-filename', type=str, default=None,
97 dest='filename', help="""
98 Synonym for --filename.
101 upload_opts.add_argument('--filename', type=str, default=None,
103 Use the given filename in the manifest, instead of the name of the
104 local file. This is useful when "-" or "/dev/stdin" is given as an
105 input file. It can be used only if there is exactly one path given and
106 it is not a directory. Implies --manifest.
109 run_opts = argparse.ArgumentParser(add_help=False)
110 _group = run_opts.add_mutually_exclusive_group()
111 _group.add_argument('--progress', action='store_true',
113 Display human-readable progress on stderr (bytes and, if possible,
114 percentage of total data size). This is the default behavior when
118 _group.add_argument('--no-progress', action='store_true',
120 Do not display human-readable progress on stderr, even if stderr is a
124 _group.add_argument('--batch-progress', action='store_true',
126 Display machine-readable progress on stderr (bytes and, if known,
130 _group = run_opts.add_mutually_exclusive_group()
131 _group.add_argument('--resume', action='store_true', default=True,
133 Continue interrupted uploads from cached state (default).
135 _group.add_argument('--no-resume', action='store_false', dest='resume',
137 Do not continue interrupted uploads from cached state.
140 arg_parser = argparse.ArgumentParser(
141 description='Copy data from the local filesystem to Keep.',
142 parents=[upload_opts, run_opts])
144 def parse_arguments(arguments):
145 args = arg_parser.parse_args(arguments)
147 if len(args.paths) == 0:
148 args.paths += ['/dev/stdin']
150 if len(args.paths) != 1 or os.path.isdir(args.paths[0]):
153 --filename argument cannot be used when storing a directory or
157 # Turn on --progress by default if stderr is a tty.
158 if (not (args.batch_progress or args.no_progress)
159 and os.isatty(sys.stderr.fileno())):
162 if args.paths == ['-']:
163 args.paths = ['/dev/stdin']
164 if not args.filename:
169 class ResumeCacheConflict(Exception):
173 class ResumeCache(object):
174 CACHE_DIR = '.cache/arvados/arv-put'
176 def __init__(self, file_spec):
177 self.cache_file = open(file_spec, 'a+')
178 self._lock_file(self.cache_file)
179 self.filename = self.cache_file.name
182 def make_path(cls, args):
184 md5.update(arvados.config.get('ARVADOS_API_HOST', '!nohost'))
185 realpaths = sorted(os.path.realpath(path) for path in args.paths)
186 md5.update('\0'.join(realpaths))
187 if any(os.path.isdir(path) for path in realpaths):
188 md5.update(str(max(args.max_manifest_depth, -1)))
190 md5.update(args.filename)
192 arv_cmd.make_home_conf_dir(cls.CACHE_DIR, 0o700, 'raise'),
195 def _lock_file(self, fileobj):
197 fcntl.flock(fileobj, fcntl.LOCK_EX | fcntl.LOCK_NB)
199 raise ResumeCacheConflict("{} locked".format(fileobj.name))
202 self.cache_file.seek(0)
203 return json.load(self.cache_file)
205 def save(self, data):
207 new_cache_fd, new_cache_name = tempfile.mkstemp(
208 dir=os.path.dirname(self.filename))
209 self._lock_file(new_cache_fd)
210 new_cache = os.fdopen(new_cache_fd, 'r+')
211 json.dump(data, new_cache)
212 os.rename(new_cache_name, self.filename)
213 except (IOError, OSError, ResumeCacheConflict) as error:
215 os.unlink(new_cache_name)
216 except NameError: # mkstemp failed.
219 self.cache_file.close()
220 self.cache_file = new_cache
223 self.cache_file.close()
227 os.unlink(self.filename)
228 except OSError as error:
229 if error.errno != errno.ENOENT: # That's what we wanted anyway.
235 self.__init__(self.filename)
238 class ArvPutCollectionWriter(arvados.ResumableCollectionWriter):
239 STATE_PROPS = (arvados.ResumableCollectionWriter.STATE_PROPS +
240 ['bytes_written', '_seen_inputs'])
242 def __init__(self, cache=None, reporter=None, bytes_expected=None,
244 self.bytes_written = 0
245 self._seen_inputs = []
247 self.reporter = reporter
248 self.bytes_expected = bytes_expected
249 super(ArvPutCollectionWriter, self).__init__(api_client)
252 def from_cache(cls, cache, reporter=None, bytes_expected=None):
255 state['_data_buffer'] = [base64.decodestring(state['_data_buffer'])]
256 writer = cls.from_state(state, cache, reporter, bytes_expected)
257 except (TypeError, ValueError,
258 arvados.errors.StaleWriterStateError) as error:
259 return cls(cache, reporter, bytes_expected)
263 def cache_state(self):
264 if self.cache is None:
266 state = self.dump_state()
267 # Transform attributes for serialization.
268 for attr, value in state.items():
269 if attr == '_data_buffer':
270 state[attr] = base64.encodestring(''.join(value))
271 elif hasattr(value, 'popleft'):
272 state[attr] = list(value)
273 self.cache.save(state)
275 def report_progress(self):
276 if self.reporter is not None:
277 self.reporter(self.bytes_written, self.bytes_expected)
279 def flush_data(self):
280 start_buffer_len = self._data_buffer_len
281 start_block_count = self.bytes_written / self.KEEP_BLOCK_SIZE
282 super(ArvPutCollectionWriter, self).flush_data()
283 if self._data_buffer_len < start_buffer_len: # We actually PUT data.
284 self.bytes_written += (start_buffer_len - self._data_buffer_len)
285 self.report_progress()
286 if (self.bytes_written / self.KEEP_BLOCK_SIZE) > start_block_count:
289 def _record_new_input(self, input_type, source_name, dest_name):
290 # The key needs to be a list because that's what we'll get back
291 # from JSON deserialization.
292 key = [input_type, source_name, dest_name]
293 if key in self._seen_inputs:
295 self._seen_inputs.append(key)
298 def write_file(self, source, filename=None):
299 if self._record_new_input('file', source, filename):
300 super(ArvPutCollectionWriter, self).write_file(source, filename)
302 def write_directory_tree(self,
303 path, stream_name='.', max_manifest_depth=-1):
304 if self._record_new_input('directory', path, stream_name):
305 super(ArvPutCollectionWriter, self).write_directory_tree(
306 path, stream_name, max_manifest_depth)
309 def expected_bytes_for(pathlist):
310 # Walk the given directory trees and stat files, adding up file sizes,
311 # so we can display progress as percent
313 for path in pathlist:
314 if os.path.isdir(path):
315 for filename in arvados.util.listdir_recursive(path):
316 bytesum += os.path.getsize(os.path.join(path, filename))
317 elif not os.path.isfile(path):
320 bytesum += os.path.getsize(path)
323 _machine_format = "{} {}: {{}} written {{}} total\n".format(sys.argv[0],
325 def machine_progress(bytes_written, bytes_expected):
326 return _machine_format.format(
327 bytes_written, -1 if (bytes_expected is None) else bytes_expected)
329 def human_progress(bytes_written, bytes_expected):
331 return "\r{}M / {}M {:.1%} ".format(
332 bytes_written >> 20, bytes_expected >> 20,
333 float(bytes_written) / bytes_expected)
335 return "\r{} ".format(bytes_written)
337 def progress_writer(progress_func, outfile=sys.stderr):
338 def write_progress(bytes_written, bytes_expected):
339 outfile.write(progress_func(bytes_written, bytes_expected))
340 return write_progress
342 def exit_signal_handler(sigcode, frame):
345 def check_project_exists(project_uuid):
347 api_client.groups().get(uuid=project_uuid).execute()
348 except (apiclient.errors.Error, arvados.errors.NotFoundError) as error:
349 raise ValueError("Project {} not found ({})".format(project_uuid,
354 def prep_project_link(args, stderr, project_exists=check_project_exists):
355 # Given the user's command line arguments, return a dictionary with data
356 # to create the desired project link for this Collection, or None.
357 # Raises ValueError if the arguments request something impossible.
358 making_collection = not (args.raw or args.stream)
359 if not making_collection:
360 if args.name or args.project_uuid:
361 raise ValueError("Requested a Link without creating a Collection")
363 link = {'tail_uuid': args.project_uuid,
364 'link_class': 'name',
366 if not link['tail_uuid']:
367 link['tail_uuid'] = api_client.users().current().execute()['uuid']
368 elif not project_exists(link['tail_uuid']):
369 raise ValueError("Project {} not found".format(args.project_uuid))
371 link['name'] = "Saved at {} by {}@{}".format(
372 datetime.datetime.utcnow().strftime("%Y-%m-%d %H:%M:%S UTC"),
373 pwd.getpwuid(os.getuid()).pw_name,
374 socket.gethostname())
376 "arv-put: No --name specified. Saving as \"%s\"\n" % link['name'])
377 link['owner_uuid'] = link['tail_uuid']
380 def create_project_link(locator, link):
381 link['head_uuid'] = locator
382 return api_client.links().create(body=link).execute()
384 def main(arguments=None, stdout=sys.stdout, stderr=sys.stderr):
386 if api_client is None:
387 api_client = arvados.api('v1')
390 args = parse_arguments(arguments)
392 project_link = prep_project_link(args, stderr)
393 except ValueError as error:
394 print >>stderr, "arv-put: {}.".format(error)
398 reporter = progress_writer(human_progress)
399 elif args.batch_progress:
400 reporter = progress_writer(machine_progress)
403 bytes_expected = expected_bytes_for(args.paths)
408 resume_cache = ResumeCache(ResumeCache.make_path(args))
409 except (IOError, OSError, ValueError):
410 pass # Couldn't open cache directory/file. Continue without it.
411 except ResumeCacheConflict:
412 print >>stderr, "\n".join([
413 "arv-put: Another process is already uploading this data.",
414 " Use --no-resume if this is really what you want."])
417 if resume_cache is None:
418 writer = ArvPutCollectionWriter(resume_cache, reporter, bytes_expected)
420 writer = ArvPutCollectionWriter.from_cache(
421 resume_cache, reporter, bytes_expected)
423 # Install our signal handler for each code in CAUGHT_SIGNALS, and save
425 orig_signal_handlers = {sigcode: signal.signal(sigcode, exit_signal_handler)
426 for sigcode in CAUGHT_SIGNALS}
428 if writer.bytes_written > 0: # We're resuming a previous upload.
429 print >>stderr, "\n".join([
430 "arv-put: Resuming previous upload from last checkpoint.",
431 " Use the --no-resume option to start over."])
433 writer.report_progress()
434 writer.do_queued_work() # Do work resumed from cache.
435 for path in args.paths: # Copy file data to Keep.
436 if os.path.isdir(path):
437 writer.write_directory_tree(
438 path, max_manifest_depth=args.max_manifest_depth)
440 writer.start_new_stream()
441 writer.write_file(path, args.filename or os.path.basename(path))
442 writer.finish_current_stream()
444 if args.progress: # Print newline to split stderr from stdout for humans.
448 output = writer.manifest_text()
450 output = ','.join(writer.data_locators())
452 # Register the resulting collection in Arvados.
453 collection = api_client.collections().create(
455 'manifest_text': writer.manifest_text(),
456 'owner_uuid': project_link['tail_uuid']
460 if 'portable_data_hash' in collection and collection['portable_data_hash']:
461 output = collection['portable_data_hash']
463 output = collection['uuid']
465 if project_link is not None:
466 # Update collection name
468 if 'name' in collection:
469 arvados.api().collections().update(uuid=collection['uuid'],
470 body={"name": project_link["name"]}).execute()
472 create_project_link(output, project_link)
473 except apiclient.errors.Error as error:
475 "arv-put: Error adding Collection to project: {}.".format(
479 # Print the locator (uuid) of the new collection.
481 if not output.endswith('\n'):
484 for sigcode, orig_handler in orig_signal_handlers.items():
485 signal.signal(sigcode, orig_handler)
490 if resume_cache is not None:
491 resume_cache.destroy()
495 if __name__ == '__main__':