def update_states(self, flow_runs: list) -> None: """ After a flow run is grabbed this function sets the state to Submitted so it won't be picked up by any other processes Args: - flow_runs (list): A list of GraphQLResult flow run objects """ for flow_run in flow_runs: # Set flow run state to `Submitted` if it is currently `Scheduled` if state.StateSchema().load(flow_run.serialized_state).is_scheduled(): self.client.set_flow_run_state( flow_run_id=flow_run.id, version=flow_run.version, state=Submitted( message="Submitted for execution", state=state.StateSchema().load(flow_run.serialized_state), ), ) # Set task run states to `Submitted` if they are currently `Scheduled` for task_run in flow_run.task_runs: if state.StateSchema().load(task_run.serialized_state).is_scheduled(): self.client.set_task_run_state( task_run_id=task_run.id, version=task_run.version, state=Submitted( message="Submitted for execution", state=state.StateSchema().load(task_run.serialized_state), ), )
def update_state(self, flow_run: GraphQLResult, deployment_info: str) -> None: """ After a flow run is grabbed this function sets the state to Submitted so it won't be picked up by any other processes Args: - flow_run (GraphQLResult): A GraphQLResult flow run object - deployment_info (str): Identifier information related to the Flow Run deployment """ self.logger.debug( "Updating states for flow run {}".format(flow_run.id) # type: ignore ) # Set flow run state to `Submitted` if it is currently `Scheduled` if state.StateSchema().load(flow_run.serialized_state).is_scheduled(): self.logger.debug( "Flow run {} is in a Scheduled state, updating to Submitted".format( flow_run.id # type: ignore ) ) self.client.set_flow_run_state( flow_run_id=flow_run.id, version=flow_run.version, state=Submitted( message="Submitted for execution. {}".format(deployment_info), state=state.StateSchema().load(flow_run.serialized_state), ), ) # Set task run states to `Submitted` if they are currently `Scheduled` for task_run in flow_run.task_runs: if state.StateSchema().load(task_run.serialized_state).is_scheduled(): self.logger.debug( "Task run {} is in a Scheduled state, updating to Submitted".format( task_run.id # type: ignore ) ) self.client.set_task_run_state( task_run_id=task_run.id, version=task_run.version, state=Submitted( message="Submitted for execution. {}".format(deployment_info), state=state.StateSchema().load(task_run.serialized_state), ), )
def _mark_flow_as_submitted(self, flow_run: GraphQLResult) -> None: """ After a flow run is grabbed this function sets the state to Submitted so it won't be picked up by any other processes Args: - flow_run (GraphQLResult): A GraphQLResult flow run object """ # Set flow run state to `Submitted` if it is currently `Scheduled` if state.StateSchema().load(flow_run.serialized_state).is_scheduled(): self.logger.debug( f"Updating flow run {flow_run.id} state from Scheduled -> Submitted..." ) self.client.set_flow_run_state( flow_run_id=flow_run.id, version=flow_run.version, state=Submitted( message="Submitted for execution", state=state.StateSchema().load(flow_run.serialized_state), ), ) # Set task run states to `Submitted` if they are currently `Scheduled` task_runs_updated = 0 for task_run in flow_run.task_runs: if state.StateSchema().load( task_run.serialized_state).is_scheduled(): task_runs_updated += 1 self.client.set_task_run_state( task_run_id=task_run.id, version=task_run.version, state=Submitted( message="Submitted for execution.", state=state.StateSchema().load( task_run.serialized_state), ), ) if task_runs_updated: self.logger.debug( f"Updated {task_runs_updated} task runs states for flow run " f"{flow_run.id} from Scheduled -> Submitted")