def _post_signal_tokens(self): """Add some signal tokens to the master.""" request = ModifyRequest(updates=[]) signal = Signal(action=Signal.EXIT) name = Name(signal='exit') signal_token = Token(name=name.get_signal_token_name()) signal_token.data = pickle.dumps(signal) request.updates.append(signal_token) signal = Signal(action=Signal.DRAIN) name.signal = 'drain' name.workflow = 'some_workflow' signal_token = Token(name=name.get_signal_token_name()) signal_token.data = pickle.dumps(signal) request.updates.append(signal_token) name.instance = '123' signal_token = Token(name=name.get_signal_token_name()) signal_token.data = pickle.dumps(signal) request.updates.append(signal_token) signal = Signal(action=Signal.ABORT) name.signal = 'abort' signal_token = Token(name=name.get_signal_token_name()) signal_token.data = pickle.dumps(signal) request.updates.append(signal_token) client = self._factory.get_client() client.modify(request)
def _post_signal_tokens(self): """Add some signal tokens to the master.""" request = ModifyRequest(updates=[]) signal = Signal(action=Signal.EXIT) name = Name(signal='exit') signal_token = Token(name=name.get_signal_token_name()) signal_token.data = pickle.dumps(signal) request.updates.append(signal_token) signal = Signal(action=Signal.DRAIN) name.signal = 'drain' name.workflow = 'some_workflow' signal_token = Token(name=name.get_signal_token_name()) signal_token.data = pickle.dumps(signal) request.updates.append(signal_token) name.instance = '123' signal_token = Token(name=name.get_signal_token_name()) signal_token.data = pickle.dumps(signal) request.updates.append(signal_token) signal = Signal(action=Signal.ABORT) name.signal = 'abort' signal_token = Token(name=name.get_signal_token_name()) signal_token.data = pickle.dumps(signal) request.updates.append(signal_token) client = self._factory.get_client() client.modify(request)
def is_signal_set(self, workflow, instance, action): """Check if a signal is set. Args: workflow: The workflow whose signal should be checked. If None, signals at the global level are checked. instance: The workflow instance whose signal should be checked. If not None, a matching workflow name must be provided. If None, signals at the workflow and the global level are checked. action: The signal action to check. Returns: True iff the signal exists in the specified context. """ for (workflow_name, instance_name) in [(workflow, instance), (workflow, None), (None, None)]: name = Name(workflow=workflow_name, instance=instance_name, signal=Signal.action_to_string(action)) token_name = name.get_signal_token_name() tokens = self._store.read_tokens(token_name) assert len(tokens) <= 1 if tokens: return True return False
def set_action(self, action): """Send a signal with a specific action to the master. Local signal store gets updated with the new action if it is successfully submitted to the master. If the communication with the master fails, locally stored signals get refreshed. Args: action: The action to set. """ attributes = {} if action == Signal.ABORT: attributes[Signal.TIMESTAMP_ATTR] = time.time() elif action == Signal.EXIT: attributes[Signal.GENERATION_ATTR] = PinballConfig.GENERATION signal = self._signals.get(action) if signal and signal.attributes == attributes: return # A signal with the same action but different data may already exist # in the master. signal_token = self._get_signal_token(action) if not signal_token: name = Name(workflow=self._workflow, instance=self._instance, signal=Signal.action_to_string(action)) signal_token = Token(name=name.get_signal_token_name()) signal = Signal(action, attributes) signal_token.data = pickle.dumps(signal) request = ModifyRequest(updates=[signal_token]) if self._send_request(request): self._signals[action] = signal
def set_action(self, action): """Send a signal with a specific action to the master. Local signal store gets updated with the new action if it is successfully submitted to the master. If the communication with the master fails, locally stored signals get refreshed. Args: action: The action to set. """ attributes = {} if action == Signal.ABORT: attributes[Signal.TIMESTAMP_ATTR] = time.time() elif action == Signal.EXIT: attributes[Signal.GENERATION_ATTR] = PinballConfig.GENERATION signal = self._signals.get(action) if signal and signal.attributes == attributes: return # A signal with the same action but different data may already exist # in the master. signal_token = self._get_signal_token(action) if not signal_token: name = Name(workflow=self._workflow, instance=self._instance, signal=Signal.action_to_string(action)) signal_token = Token(name=name.get_signal_token_name()) signal = Signal(action, attributes) signal_token.data = pickle.dumps(signal) request = ModifyRequest(updates=[signal_token]) if self._send_request(request): self._signals[action] = signal
def is_signal_set(self, workflow, instance, action): """Check if a signal is set. Args: workflow: The workflow whose signal should be checked. If None, signals at the global level are checked. instance: The workflow instance whose signal should be checked. If not None, a matching workflow name must be provided. If None, signals at the workflow and the global level are checked. action: The signal action to check. Returns: True iff the signal exists in the specified context. """ for (workflow_name, instance_name) in [(workflow, instance), (workflow, None), (None, None)]: name = Name(workflow=workflow_name, instance=instance_name, signal=Signal.action_to_string(action)) token_name = name.get_signal_token_name() tokens = self._store.read_tokens(token_name) assert len(tokens) <= 1 if tokens: return True return False
def _generate_signal_tokens(workflows): result = [] for w in range(0, workflows, 2): workflow = 'workflow_%d' % w signal = Signal(Signal.DRAIN) name = Name(workflow=workflow, signal=Signal.action_to_string(signal.action)) result.append(Token(name=name.get_signal_token_name(), version=10000000000 * w, data=pickle.dumps(signal))) return result
def _generate_signal_tokens(workflows): result = [] for w in range(0, workflows, 2): workflow = 'workflow_%d' % w signal = Signal(Signal.DRAIN) name = Name(workflow=workflow, signal=Signal.action_to_string(signal.action)) result.append( Token(name=name.get_signal_token_name(), version=10000000000 * w, data=pickle.dumps(signal))) return result
def _get_signal(self, workflow, instance, action, active): signal = Signal.action_to_string(action) name = Name(workflow=workflow, instance=instance, signal=signal) signal_token_name = name.get_signal_token_name() if active: tokens = self._store.read_active_tokens( name_prefix=signal_token_name) else: tokens = self._store.read_archived_tokens( name_prefix=signal_token_name) if not tokens: return None assert len(tokens) == 1 assert tokens[0].name == signal_token_name return pickle.loads(tokens[0].data)
def _get_signal(self, workflow, instance, action, active): signal = Signal.action_to_string(action) name = Name(workflow=workflow, instance=instance, signal=signal) signal_token_name = name.get_signal_token_name() if active: tokens = self._store.read_active_tokens( name_prefix=signal_token_name) else: tokens = self._store.read_archived_tokens( name_prefix=signal_token_name) if not tokens: return None assert len(tokens) == 1 assert tokens[0].name == signal_token_name return pickle.loads(tokens[0].data)
def _has_abort_token(self, tokens): """Check if a list of tokens contains an abort token. Args: tokens: The list of tokens to check. Returns: True iff the list contains an abort token. """ abort_signal = Signal.action_to_string(Signal.ABORT) abort_name = Name(workflow=self._workflow, instance=self._instance, signal=abort_signal) abort_token_name = abort_name.get_signal_token_name() for token in tokens: if token.name == abort_token_name: return True return False
def _move_job_token_to_waiting(self, job, succeeded): """Move the owned job token to the waiting group. If the job succeeded, also post events to job outputs. If the job failed or it is the final job (a job with no outputs), post an archive signal to finish the workflow. Args: job: The job that should be stored in the data field of the waiting job token. succeeded: True if the job succeeded, otherwise False. """ assert self._owned_job_token name = Name.from_job_token_name(self._owned_job_token.name) name.job_state = Name.WAITING_STATE waiting_job_token = Token(name=name.get_job_token_name(), priority=self._owned_job_token.priority, data=pickle.dumps(job)) request = ModifyRequest(deletes=[self._owned_job_token], updates=[waiting_job_token]) if succeeded: request.updates.extend(self._get_output_event_tokens(job)) if not job.outputs or not succeeded: # This is either the only job in the workflow with no outputs or a # failed job. In either case, the workflow is done. signaller = Signaller(self._client, workflow=name.workflow, instance=name.instance) if not signaller.is_action_set(Signal.ARCHIVE): signal_name = Name(workflow=name.workflow, instance=name.instance, signal=Signal.action_to_string( Signal.ARCHIVE)) signal = Signal(Signal.ARCHIVE) signal_token = Token(name=signal_name.get_signal_token_name()) signal_token.data = pickle.dumps(signal) request.updates.append(signal_token) self._send_request(request)
def _move_job_token_to_waiting(self, job, succeeded): """Move the owned job token to the waiting group. If the job succeeded, also post events to job outputs. If the job failed or it is the final job (a job with no outputs), post an archive signal to finish the workflow. Args: job: The job that should be stored in the data field of the waiting job token. succeeded: True if the job succeeded, otherwise False. """ assert self._owned_job_token name = Name.from_job_token_name(self._owned_job_token.name) name.job_state = Name.WAITING_STATE waiting_job_token = Token(name=name.get_job_token_name(), priority=self._owned_job_token.priority, data=pickle.dumps(job)) request = ModifyRequest(deletes=[self._owned_job_token], updates=[waiting_job_token]) if succeeded: request.updates.extend(self._get_output_event_tokens(job)) if not job.outputs or not succeeded: # This is either the only job in the workflow with no outputs or a # failed job. In either case, the workflow is done. signaller = Signaller(self._client, workflow=name.workflow, instance=name.instance) if not signaller.is_action_set(Signal.ARCHIVE): signal_name = Name( workflow=name.workflow, instance=name.instance, signal=Signal.action_to_string(Signal.ARCHIVE)) signal = Signal(Signal.ARCHIVE) signal_token = Token(name=signal_name.get_signal_token_name()) signal_token.data = pickle.dumps(signal) request.updates.append(signal_token) self._send_request(request)
def _get_signal_token(self, action): """Retrieve signal for a specific action from the master. Args: action: The action to get signal for. Returns: The signal token if found, otherwise None. """ request = QueryRequest(queries=[]) name = Name(workflow=self._workflow, instance=self._instance, signal=Signal.action_to_string(action)) query = Query() query.namePrefix = name.get_signal_token_name() request.queries.append(query) response = self._client.query(request) assert len(response.tokens) == 1 tokens = response.tokens[0] if not tokens: return None assert len(tokens) == 1 return tokens[0]
def _get_signal_token(self, action): """Retrieve signal for a specific action from the master. Args: action: The action to get signal for. Returns: The signal token if found, otherwise None. """ request = QueryRequest(queries=[]) name = Name(workflow=self._workflow, instance=self._instance, signal=Signal.action_to_string(action)) query = Query() query.namePrefix = name.get_signal_token_name() request.queries.append(query) response = self._client.query(request) assert len(response.tokens) == 1 tokens = response.tokens[0] if not tokens: return None assert len(tokens) == 1 return tokens[0]