[Notes] [Git][BuildGrid/buildgrid][mablanch/117-job-scheduler-refactoring] 3 commits: job.py: Refactor for stronger encapsulation

[Martin Blanchard <gitlab mg gitlab com>]

Title: GitLab

Martin Blanchard pushed to branch mablanch/117-job-scheduler-refactoring at BuildGrid / buildgrid

Commits:

• 1892d2b1
  by Martin Blanchard at 2018-10-11T19:24:56Z

  job.py: Refactor for stronger encapsulation
  

• 6270eac5
  by Martin Blanchard at 2018-10-11T19:30:28Z

  job.py: Implement operation cancellation
  

• e076ce0c
  by Martin Blanchard at 2018-10-11T19:30:30Z

  scheduler.py: Implement operation cancellation
  

5 changed files:

• buildgrid/server/bots/instance.py
• buildgrid/server/execution/instance.py
• buildgrid/server/job.py
• buildgrid/server/scheduler.py
• tests/integration/operations_service.py

Changes:

buildgrid/server/bots/instance.py

---

...	...	@@ -109,7 +109,7 @@ class BotsInterface:
109	109	if server_state == LeaseState.PENDING:
110	110
111	111	if client_state == LeaseState.ACTIVE:
112		- self._scheduler.update_job_lease_state(client_lease.id, client_lease.state)
	112	+ self._scheduler.update_job_lease_state(client_lease.id, client_state)
113	113	elif client_state == LeaseState.COMPLETED:
114	114	# TODO: Lease was rejected
115	115	raise NotImplementedError("'Not Accepted' is unsupported")
...	...	@@ -122,8 +122,7 @@ class BotsInterface:
122	122	pass
123	123
124	124	elif client_state == LeaseState.COMPLETED:
125		- self._scheduler.update_job_lease_state(client_lease.id, client_lease.state)
126		- self._scheduler.job_complete(client_lease.id, client_lease.result, client_lease.status)
	125	+ self._scheduler.update_job_lease_state(client_lease.id, client_state, lease_status=client_lease.status, lease_result=client_lease.result)
127	126	return None
128	127
129	128	else:

buildgrid/server/execution/instance.py

---

...	...	@@ -53,7 +53,7 @@ class ExecutionInstance:
53	53
54	54	self._scheduler.append_job(job, skip_cache_lookup)
55	55
56		- return job.get_operation()
	56	+ return job.operation
57	57
58	58	def register_message_client(self, name, queue):
59	59	try:

buildgrid/server/job.py

---

...	...	@@ -22,6 +22,7 @@ from enum import Enum
22	22	from buildgrid._protos.build.bazel.remote.execution.v2 import remote_execution_pb2
23	23	from buildgrid._protos.google.devtools.remoteworkers.v1test2 import bots_pb2
24	24	from buildgrid._protos.google.longrunning import operations_pb2
	25	+from buildgrid._protos.google.rpc import code_pb2
25	26
26	27
27	28	class OperationStage(Enum):
...	...	@@ -53,18 +54,24 @@ class LeaseState(Enum):
53	54	class Job:
54	55
55	56	def __init__(self, action_digest, do_not_cache=False, message_queue=None):
56		- self.lease = None
57	57	self.logger = logging.getLogger(__name__)
58		- self.n_tries = 0
59		- self.result = None
60		- self.result_cached = False
61	58
	59	+ self._name = str(uuid.uuid4())
62	60	self._action_digest = action_digest
63	61	self._do_not_cache = do_not_cache
64		- self._execute_stage = OperationStage.UNKNOWN
65		- self._name = str(uuid.uuid4())
66		- self._operation = operations_pb2.Operation(name=self._name)
	62	+ self._operation = operations_pb2.Operation()
	63	+ self._lease = None
	64	+ self._n_tries = 0
	65	+
	66	+ self.__operation_metadata = remote_execution_pb2.ExecuteOperationMetadata()
	67	+ self.__operation_canceled = False
	68	+
	69	+ self.__operation_metadata.action_digest.CopyFrom(self._action_digest)
	70	+ self.__operation_metadata.stage = OperationStage.UNKNOWN.value
	71	+
67	72	self._operation_update_queues = []
	73	+ self._operation.name = self._name
	74	+ self._operation.done = False
68	75
69	76	if message_queue is not None:
70	77	self.register_client(message_queue)
...	...	@@ -81,47 +88,87 @@ class Job:
81	88	def do_not_cache(self):
82	89	return self._do_not_cache
83	90
84		- def check_job_finished(self):
85		- if not self._operation_update_queues:
86		- return self._operation.done
87		- return False
	91	+ @property
	92	+ def operation(self):
	93	+ return self._operation
	94	+
	95	+ @property
	96	+ def operation_stage(self):
	97	+ return OperationStage(self.__operation_metadata.state)
	98	+
	99	+ @property
	100	+ def lease(self):
	101	+ return self._lease
	102	+
	103	+ @property
	104	+ def lease_state(self):
	105	+ if self._lease is not None:
	106	+ return LeaseState(self._lease.state)
	107	+ else:
	108	+ return None
	109	+
	110	+ @property
	111	+ def n_tries(self):
	112	+ return self._n_tries
88	113
89	114	def register_client(self, queue):
90	115	self._operation_update_queues.append(queue)
91		- queue.put(self.get_operation())
	116	+ queue.put(self._operation)
92	117
93	118	def unregister_client(self, queue):
94	119	self._operation_update_queues.remove(queue)
95	120
96		- def get_operation(self):
97		- self._operation.metadata.Pack(self.get_operation_meta())
98		- if self.result is not None:
99		- self._operation.done = True
100		- response = remote_execution_pb2.ExecuteResponse(result=self.result,
101		- cached_result=self.result_cached)
	121	+ def create_lease(self):
	122	+ self._lease = bots_pb2.Lease()
	123	+ self._lease.id = self._name
	124	+ self._lease.payload.Pack(self._action_digest)
	125	+ self._lease.state = LeaseState.PENDING.value
	126	+
	127	+ return self._lease
	128	+
	129	+ def update_lease_state(self, state, status=None, result=None):
	130	+ if state.value == self._lease.state:
	131	+ return
102	132
103		- if not self.result_cached:
104		- response.status.CopyFrom(self.lease.status)
	133	+ self._lease.state = state.value
	134	+
	135	+ if self._lease.state == LeaseState.COMPLETED.value:
	136	+ response = remote_execution_pb2.ExecuteResponse()
	137	+ response.result.CopyFrom(result)
	138	+ response.cached_result = False
	139	+ response.status.CopyFrom(status)
105	140
106	141	self._operation.response.Pack(response)
	142	+ self._operation.done = True
107	143
108		- return self._operation
	144	+ def cancel_lease(self):
	145	+ if self._lease is not None:
	146	+ self.update_lease_state(LeaseState.CANCELLED)
109	147
110		- def get_operation_meta(self):
111		- meta = remote_execution_pb2.ExecuteOperationMetadata()
112		- meta.stage = self._execute_stage.value
113		- meta.action_digest.CopyFrom(self._action_digest)
	148	+ def update_operation_stage(self, stage):
	149	+ if stage.value == self.__operation_metadata.stage:
	150	+ return
114	151
115		- return meta
	152	+ self.__operation_metadata.stage = stage.value
116	153
117		- def create_lease(self):
118		- lease = bots_pb2.Lease(id=self.name, state=LeaseState.PENDING.value)
119		- lease.payload.Pack(self._action_digest)
	154	+ if self.__operation_metadata.stage == OperationStage.QUEUED.value:
	155	+ self._n_tries += 1
120	156
121		- self.lease = lease
122		- return lease
	157	+ self._operation.metadata.Pack(self.__operation_metadata)
123	158
124		- def update_operation_stage(self, stage):
125		- self._execute_stage = stage
126	159	for queue in self._operation_update_queues:
127		- queue.put(self.get_operation())
	160	+ queue.put(self._operation)
	161	+
	162	+ def cancel_operation(self):
	163	+ self.__operation_canceled = True
	164	+
	165	+ self.cancel_lease()
	166	+
	167	+ response = remote_execution_pb2.ExecuteResponse()
	168	+ response.status.code = code_pb2.CANCELLED
	169	+ response.status.message = "The operation was cancelled by the caller."
	170	+
	171	+ self._operation.response.Pack(response)
	172	+ self._operation.done = True
	173	+
	174	+ self.update_operation_stage(OperationStage.COMPLETED)

buildgrid/server/scheduler.py

---

...	...	@@ -45,7 +45,7 @@ class Scheduler:
45	45	def unregister_client(self, name, queue):
46	46	job = self.jobs[name]
47	47	job.unregister_client(queue)
48		- if job.check_job_finished():
	48	+ if job.operation.done:
49	49	del self.jobs[name]
50	50
51	51	def append_job(self, job, skip_cache_lookup=False):
...	...	@@ -58,8 +58,6 @@ class Scheduler:
58	58	job.update_operation_stage(OperationStage.QUEUED)
59	59
60	60	else:
61		- job.result = cached_result
62		- job.result_cached = True
63	61	job.update_operation_stage(OperationStage.COMPLETED)
64	62
65	63	else:
...	...	@@ -75,44 +73,45 @@ class Scheduler:
75	73	# TODO: Mark these jobs as done
76	74	else:
77	75	job.update_operation_stage(OperationStage.QUEUED)
78		- job.n_tries += 1
79	76	self.queue.appendleft(job)
80	77
81		- def job_complete(self, name, result, status):
82		- job = self.jobs[name]
83		- job.lease.status.CopyFrom(status)
84		- action_result = remote_execution_pb2.ActionResult()
85		- result.Unpack(action_result)
86		- job.result = action_result
87		- if not job.do_not_cache and self._action_cache is not None:
88		- if not job.lease.status.code:
89		- self._action_cache.update_action_result(job.action_digest, action_result)
90		- job.update_operation_stage(OperationStage.COMPLETED)
91		-
92	78	def get_operations(self):
93	79	response = operations_pb2.ListOperationsResponse()
94	80	for v in self.jobs.values():
95		- response.operations.extend([v.get_operation()])
	81	+ response.operations.extend([v.operation])
96	82	return response
97	83
98		- def update_job_lease_state(self, name, state):
99		- job = self.jobs[name]
100		- job.lease.state = state
	84	+ def update_job_lease_state(self, job_name, lease_state, lease_status=None, lease_result=None):
	85	+ job = self.jobs[job_name]
	86	+
	87	+ if lease_state != LeaseState.COMPLETED:
	88	+ job.update_lease_state(lease_state)
	89	+ else:
	90	+ action_result = remote_execution_pb2.ActionResult()
	91	+ lease_result.Unpack(action_result)
	92	+
	93	+ job.update_lease_state(lease_state, status=lease_status, result=action_result)
	94	+
	95	+ if not job.do_not_cache and self._action_cache is not None:
	96	+ if not job.lease.status.code:
	97	+ self._action_cache.update_action_result(job.action_digest, action_result)
	98	+
	99	+ job.update_operation_stage(OperationStage.COMPLETED)
101	100
102	101	def get_job_lease(self, name):
103	102	return self.jobs[name].lease
104	103
105		- def cancel_session(self, name):
106		- job = self.jobs[name]
107		- state = job.lease.state
108		- if state in (LeaseState.PENDING.value, LeaseState.ACTIVE.value):
109		- self.retry_job(name)
	104	+ def cancel_job_lease(self, job_name):
	105	+ self.jobs[job_name].cancel_lease()
	106	+ self.retry_job(job_name)
	107	+
	108	+ def cancel_job_operation(self, job_name):
	109	+ self.jobs[job_name].cancel_operation()
110	110
111	111	def create_lease(self):
112	112	if self.queue:
113	113	job = self.queue.popleft()
114	114	job.update_operation_stage(OperationStage.EXECUTING)
115	115	job.create_lease()
116		- job.lease.state = LeaseState.PENDING.value
117	116	return job.lease
118	117	return None

tests/integration/operations_service.py

---

...	...	@@ -30,6 +30,7 @@ from buildgrid._protos.google.longrunning import operations_pb2
30	30	from buildgrid._protos.google.rpc import status_pb2
31	31	from buildgrid.server.cas.storage import lru_memory_cache
32	32	from buildgrid.server.controller import ExecutionController
	33	+from buildgrid.server.job import OperationStage
33	34	from buildgrid.server.operations import service
34	35	from buildgrid.server.operations.service import OperationsService
35	36	from buildgrid.utils import create_digest
...	...	@@ -131,9 +132,10 @@ def test_list_operations_with_result(instance, controller, execute_request, cont
131	132	action_result.output_files.extend([output_file])
132	133
133	134	controller.operations_instance._scheduler.jobs[response_execute.name].create_lease()
134		- controller.operations_instance._scheduler.job_complete(response_execute.name,
135		- _pack_any(action_result),
136		- status_pb2.Status())
	135	+ controller.operations_instance._scheduler.update_job_lease_state(response_execute.name,
	136	+ OperationStage.COMPLETED,
	137	+ lease_status=status_pb2.Status(),
	138	+ lease_result=_pack_any(action_result))
137	139
138	140	request = operations_pb2.ListOperationsRequest(name=instance_name)
139	141	response = instance.ListOperations(request, context)

—
View it on GitLab.
You're receiving this email because of your account on gitlab.com. If you'd like to receive fewer emails, you can adjust your notification settings.