Removes redundant function definitions from rpc and server

luigi/rpc.py

@@ -19,7 +19,6 @@
 rpc.py implements the client side of it, server.py implements the server side.
 See :doc:`/central_scheduler` for more info.
 """
-
 import json
 import logging
 import socket
@@ -30,8 +29,7 @@
 from luigi.six.moves.urllib.error import URLError
 
 from luigi import configuration
-from luigi.scheduler import PENDING, Scheduler
-
+from luigi.scheduler import Scheduler, RPC_METHODS
 
 HAS_UNIX_SOCKET = True
 HAS_REQUESTS = True
@@ -149,92 +147,6 @@ def _request(self, url, data, log_exceptions=True, attempts=3, allow_null=True):
  return response
  raise RPCError("Received null response from remote scheduler %r" % self._url)
 
- def ping(self, worker):
- # just one attempt, keep-alive thread will keep trying anyway
- self._request('/api/ping', {'worker': worker}, attempts=1)
-
- def add_task(self, worker, task_id, status=PENDING, runnable=True,
- deps=None, new_deps=None, expl=None, resources=None, priority=0,
- family='', module=None, params=None, assistant=False,
- tracking_url=None):
- self._request('/api/add_task', {
- 'task_id': task_id,
- 'worker': worker,
- 'status': status,
- 'runnable': runnable,
- 'deps': deps,
- 'new_deps': new_deps,
- 'expl': expl,
- 'resources': resources,
- 'priority': priority,
- 'family': family,
- 'module': module,
- 'params': params,
- 'assistant': assistant,
- 'tracking_url': tracking_url,
- })
-
- def get_work(self, worker, host=None, assistant=False, current_tasks=None):
- return self._request(
- '/api/get_work',
- {
- 'worker': worker,
- 'host': host,
- 'assistant': assistant,
- 'current_tasks': current_tasks,
- },
- allow_null=False,
- )
-
- def graph(self):
- return self._request('/api/graph', {})
-
- def dep_graph(self, task_id, include_done=True):
- return self._request('/api/dep_graph', {'task_id': task_id, 'include_done': include_done})
-
- def inverse_dep_graph(self, task_id, include_done=True):
- return self._request('/api/inverse_dep_graph', {
- 'task_id': task_id, 'include_done': include_done})
-
- def task_list(self, status, upstream_status, search=None):
- return self._request('/api/task_list', {
- 'search': search,
- 'status': status,
- 'upstream_status': upstream_status,
- })
-
- def worker_list(self):
- return self._request('/api/worker_list', {})
-
- def resource_list(self):
- return self._request('/api/resource_list', {})
-
- def task_search(self, task_str):
- return self._request('/api/task_search', {'task_str': task_str})
-
- def fetch_error(self, task_id):
- return self._request('/api/fetch_error', {'task_id': task_id})
-
- def add_worker(self, worker, info):
- return self._request('/api/add_worker', {'worker': worker, 'info': info})
-
- def disable_worker(self, worker):
- return self._request('/api/disable_worker', {'worker': worker})
-
- def update_resources(self, **resources):
- return self._request('/api/update_resources', resources)
-
- def prune(self):
- return self._request('/api/prune', {})
-
- def re_enable_task(self, task_id):
- return self._request('/api/re_enable_task', {'task_id': task_id})
-
- def set_task_status_message(self, task_id, status_message):
- self._request('/api/set_task_status_message', {
- 'task_id': task_id,
- 'status_message': status_message
- })
 
- def get_task_status_message(self, task_id):
-  return self._request('/api/get_task_status_message', {'task_id': task_id})
+for method_name, method in RPC_METHODS.items():
+ setattr(RemoteScheduler, method_name, method)

luigi/scheduler.py

@@ -22,6 +22,8 @@
 """
 
 import collections
+import inspect
+
 try:
  import cPickle as pickle
 except ImportError:
@@ -77,6 +79,36 @@ class Scheduler(object):
 
 TASK_FAMILY_RE = re.compile(r'([^(_]+)[(_]')
 
+RPC_METHODS = {}
+
+
+def rpc_method(**request_args):
+ def _rpc_method(fn):
+ # If request args are passed, return this function again for use as
+ # the decorator function with the request args attached.
+ fn_args = inspect.getargspec(fn)
+
+ assert not fn_args.varargs
+ assert fn_args.args[0] == 'self'
+ all_args = fn_args.args[1:]
+ defaults = dict(zip(reversed(all_args), reversed(fn_args.defaults or ())))
+ required_args = frozenset(arg for arg in all_args if arg not in defaults)
+ fn_name = fn.__name__
+
+ @functools.wraps(fn)
+ def rpc_func(self, *args, **kwargs):
+ actual_args = defaults.copy()
+ actual_args.update(dict(zip(all_args, args)))
+ actual_args.update(kwargs)
+ if not all(arg in actual_args for arg in required_args):
+ raise TypeError('{} takes {} arguments ({} given)'.format(
+ fn_name, len(all_args), len(actual_args)))
+ return self._request('/api/{}'.format(fn_name), actual_args, **request_args)
+
+ RPC_METHODS[fn_name] = rpc_func
+ return fn
+ return _rpc_method
+
 
 class scheduler(Config):
  # TODO(erikbern): the config_path is needed for backwards compatilibity. We
@@ -521,6 +553,7 @@ def load(self):
  def dump(self):
  self._state.dump()
 
+ @rpc_method()
  def prune(self):
  logger.info("Starting pruning of task graph")
  self._prune_workers()
@@ -575,18 +608,20 @@ def _update_priority(self, task, prio, worker):
  if t is not None and prio > t.priority:
  self._update_priority(t, prio, worker)
 
+ @rpc_method()
  def add_task(self, task_id=None, status=PENDING, runnable=True,
  deps=None, new_deps=None, expl=None, resources=None,
  priority=0, family='', module=None, params=None,
- assistant=False, tracking_url=None, **kwargs):
+ assistant=False, tracking_url=None, worker=None, **kwargs):
  """
  * add task identified by task_id if it doesn't exist
  * if deps is not None, update dependency list
  * update status of task
  * add additional workers/stakeholders
  * update priority when needed
  """
- worker_id = kwargs['worker']
+ assert worker is not None
+ worker_id = worker
  worker_enabled = self.update(worker_id)
 
  if worker_enabled:
@@ -655,12 +690,15 @@ def add_task(self, task_id=None, status=PENDING, runnable=True,
  self._state.get_worker(worker_id).tasks.add(task)
  task.runnable = runnable
 
+ @rpc_method()
  def add_worker(self, worker, info, **kwargs):
  self._state.get_worker(worker).add_info(info)
 
+ @rpc_method()
  def disable_worker(self, worker):
  self._state.disable_workers({worker})
 
+ @rpc_method()
  def update_resources(self, **resources):
  if self._resources is None:
  self._resources = {}
@@ -706,7 +744,8 @@ def _schedulable(self, task):
  def _retry_time(self, task, config):
  return time.time() + config.retry_delay
 
- def get_work(self, host=None, assistant=False, current_tasks=None, **kwargs):
+ @rpc_method(allow_null=False)
+ def get_work(self, host=None, assistant=False, current_tasks=None, worker=None, **kwargs):
  # TODO: remove any expired nodes
 
  # Algo: iterate over all nodes, find the highest priority node no dependencies and available
@@ -723,7 +762,8 @@ def get_work(self, host=None, assistant=False, current_tasks=None, **kwargs):
  if self._config.prune_on_get_work:
  self.prune()
 
- worker_id = kwargs['worker']
+ assert worker is not None
+ worker_id = worker
  # Return remaining tasks that have no FAILED descendants
  self.update(worker_id, {'host': host}, get_work=True)
  if assistant:
@@ -818,6 +858,7 @@ def get_work(self, host=None, assistant=False, current_tasks=None, **kwargs):
 
  return reply
 
+ @rpc_method(attempts=1)
  def ping(self, **kwargs):
  worker_id = kwargs['worker']
  self.update(worker_id)
@@ -873,6 +914,7 @@ def _serialize_task(self, task_id, include_deps=True, deps=None):
  ret['deps'] = list(task.deps if deps is None else deps)
  return ret
 
+ @rpc_method()
  def graph(self, **kwargs):
  self.prune()
  serialized = {}
@@ -948,12 +990,14 @@ def dep_func(t):
 
  return serialized
 
+ @rpc_method()
  def dep_graph(self, task_id, include_done=True, **kwargs):
  self.prune()
  if not self._state.has_task(task_id):
  return {}
  return self._traverse_graph(task_id, include_done=include_done)
 
+ @rpc_method()
  def inverse_dep_graph(self, task_id, include_done=True, **kwargs):
  self.prune()
  if not self._state.has_task(task_id):
@@ -965,7 +1009,8 @@ def inverse_dep_graph(self, task_id, include_done=True, **kwargs):
  return self._traverse_graph(
  task_id, dep_func=lambda t: inverse_graph[t.id], include_done=include_done)
 
- def task_list(self, status, upstream_status, limit=True, search=None, **kwargs):
+ @rpc_method()
+ def task_list(self, status='', upstream_status='', limit=True, search=None, **kwargs):
  """
  Query for a subset of tasks by status.
  """
@@ -996,6 +1041,7 @@ def _first_task_display_name(self, worker):
  else:
  return task_id
 
+ @rpc_method()
  def worker_list(self, include_running=True, **kwargs):
  self.prune()
  workers = [
@@ -1027,6 +1073,7 @@ def worker_list(self, include_running=True, **kwargs):
  worker['running'] = tasks
  return workers
 
+ @rpc_method()
  def resource_list(self):
  """
  Resources usage info and their consumers (tasks).
@@ -1062,6 +1109,7 @@ def resources(self):
  ret[resource]['used'] = 0
  return ret
 
+ @rpc_method()
  def task_search(self, task_str, **kwargs):
  """
  Query for a subset of tasks by task_id.
@@ -1077,6 +1125,7 @@ def task_search(self, task_str, **kwargs):
  result[task.status][task.id] = serialized
  return result
 
+ @rpc_method()
  def re_enable_task(self, task_id):
  serialized = {}
  task = self._state.get_task(task_id)
@@ -1085,18 +1134,21 @@ def re_enable_task(self, task_id):
  serialized = self._serialize_task(task_id)
  return serialized
 
+ @rpc_method()
  def fetch_error(self, task_id, **kwargs):
  if self._state.has_task(task_id):
  task = self._state.get_task(task_id)
  return {"taskId": task_id, "error": task.expl, 'displayName': task.pretty_id}
  else:
  return {"taskId": task_id, "error": ""}
 
+ @rpc_method()
  def set_task_status_message(self, task_id, status_message):
  if self._state.has_task(task_id):
  task = self._state.get_task(task_id)
  task.status_message = status_message
 
+ @rpc_method()
  def get_task_status_message(self, task_id):
  if self._state.has_task(task_id):
  task = self._state.get_task(task_id)

luigi/server.py

@@ -53,8 +53,7 @@
 import tornado.netutil
 import tornado.web
 
-from luigi.scheduler import CentralPlannerScheduler
-
+from luigi.scheduler import CentralPlannerScheduler, RPC_METHODS
 
 logger = logging.getLogger("luigi.server")
 
@@ -68,26 +67,7 @@ def initialize(self, scheduler):
  self._scheduler = scheduler
 
  def get(self, method):
- if method not in [
- 'add_task',
- 'add_worker',
- 'dep_graph',
- 'disable_worker',
- 'fetch_error',
- 'get_work',
- 'graph',
- 'inverse_dep_graph',
- 'ping',
- 'prune',
- 're_enable_task',
- 'resource_list',
- 'task_list',
- 'task_search',
- 'update_resources',
- 'worker_list',
- 'set_task_status_message',
- 'get_task_status_message',
- ]:
+ if method not in RPC_METHODS:
  self.send_error(404)
  return
  payload = self.get_argument('data', default="{}")