5011: Add --replication argument to arv-put.
[arvados.git] / sdk / python / arvados / commands / ws.py
1 #!/usr/bin/env python
2
3 import sys
4 import logging
5 import argparse
6 import arvados
7 import json
8 from arvados.events import subscribe
9 import signal
10
11 def main(arguments=None):
12     logger = logging.getLogger('arvados.arv-ws')
13
14     parser = argparse.ArgumentParser()
15     parser.add_argument('-u', '--uuid', type=str, default="", help="Filter events on object_uuid")
16     parser.add_argument('-f', '--filters', type=str, default="", help="Arvados query filter to apply to log events (JSON encoded)")
17
18     group = parser.add_mutually_exclusive_group()
19     group.add_argument('--poll-interval', default=15, type=int, help="If websockets is not available, specify the polling interval, default is every 15 seconds")
20     group.add_argument('--no-poll', action='store_false', dest='poll_interval', help="Do not poll if websockets are not available, just fail")
21
22     group = parser.add_mutually_exclusive_group()
23     group.add_argument('-p', '--pipeline', type=str, default="", help="Supply pipeline uuid, print log output from pipeline and its jobs")
24     group.add_argument('-j', '--job', type=str, default="", help="Supply job uuid, print log output from jobs")
25
26     args = parser.parse_args(arguments)
27
28     global filters
29     global known_component_jobs
30     global ws
31
32     filters = []
33     known_component_jobs = set()
34     ws = None
35
36     def update_subscribed_components(components):
37         global known_component_jobs
38         global filters
39         pipeline_jobs = set()
40         for c in components:
41             if "job" in components[c]:
42                 pipeline_jobs.add(components[c]["job"]["uuid"])
43         if known_component_jobs != pipeline_jobs:
44             ws.unsubscribe(filters)
45             filters = [['object_uuid', 'in', [args.pipeline] + list(pipeline_jobs)]]
46             ws.subscribe([['object_uuid', 'in', [args.pipeline] + list(pipeline_jobs)]])
47             known_component_jobs = pipeline_jobs
48
49     api = arvados.api('v1')
50
51     if args.uuid:
52         filters += [ ['object_uuid', '=', args.uuid] ]
53
54     if args.filters:
55         filters += json.loads(args.filters)
56
57     if args.job:
58         filters += [ ['object_uuid', '=', args.job] ]
59
60     if args.pipeline:
61         filters += [ ['object_uuid', '=', args.pipeline] ]
62
63     def on_message(ev):
64         global filters
65         global ws
66
67         logger.debug(ev)
68         if 'event_type' in ev and (args.pipeline or args.job):
69             if ev['event_type'] in ('stderr', 'stdout'):
70                 sys.stdout.write(ev["properties"]["text"])
71             elif ev["event_type"] in ("create", "update"):
72                 if ev["object_kind"] == "arvados#pipelineInstance":
73                     update_subscribed_components(ev["properties"]["new_attributes"]["components"])
74
75                 if ev["object_kind"] == "arvados#pipelineInstance" and args.pipeline:
76                     if ev["properties"]["new_attributes"]["state"] in ("Complete", "Failed", "Paused"):
77                         ws.close()
78
79                 if ev["object_kind"] == "arvados#job" and args.job:
80                     if ev["properties"]["new_attributes"]["state"] in ("Complete", "Failed", "Cancelled"):
81                         ws.close()
82         elif 'status' in ev and ev['status'] == 200:
83             pass
84         else:
85             print json.dumps(ev)
86
87     try:
88         ws = subscribe(arvados.api('v1'), filters, on_message, poll_fallback=args.poll_interval)
89         if ws:
90             if args.pipeline:
91                 c = api.pipeline_instances().get(uuid=args.pipeline).execute()
92                 update_subscribed_components(c["components"])
93                 if c["state"] in ("Complete", "Failed", "Paused"):
94                     ws.close()
95             ws.run_forever()
96     except KeyboardInterrupt:
97         pass
98     except Exception as e:
99         logger.error(e)
100     finally:
101         if ws:
102             ws.close()