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 test_mark_flow_as_submitted(monkeypatch, cloud_api, with_task_runs): agent = Agent() agent.client = MagicMock() agent._mark_flow_as_submitted(flow_run=GraphQLResult({ "id": "id", "serialized_state": Scheduled().serialize(), "version": 1, "task_runs": ([ GraphQLResult({ "id": "task-id", "version": 1, "serialized_state": Scheduled().serialize(), }) ] if with_task_runs else []), })) agent.client.set_flow_run_state.assert_called_once_with( flow_run_id="id", version=1, state=Submitted(message="Submitted for execution")) if with_task_runs: agent.client.set_task_run_state.assert_called_once_with( task_run_id="task-id", version=1, state=Submitted(message="Submitted for execution"), ) else: agent.client.set_task_run_state.assert_not_called()
async def test_set_flow_run_states_submitted_lock_on_same_flow_run_id( self, run_query, flow_run_id ): result = await run_query( query=self.mutation, variables=dict( input=dict( states=[ dict(flow_run_id=flow_run_id, state=Submitted().serialize()) ] ) ), ) bad_result = await run_query( query=self.mutation, variables=dict( input=dict( states=[ dict(flow_run_id=flow_run_id, state=Submitted().serialize()) ] ) ), ) assert result.data.set_flow_run_states.states[0].status == "SUCCESS" assert ( "State update failed: this run has already been submitted." in bad_result.errors[0].message )
async def test_state_submitted_lock_on_same_flow_run_id(self, flow_run_id): await api.states.set_flow_run_state(flow_run_id=flow_run_id, state=Submitted()) with pytest.raises( ValueError, match= "State update failed: this run has already been submitted.", ): await api.states.set_flow_run_state(flow_run_id=flow_run_id, state=Submitted())
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), ), )
class TestCancelFlowRun: mutation = """ mutation($input: cancel_flow_run_input!) { cancel_flow_run(input: $input) { state } } """ @pytest.mark.parametrize( "state,res_state,version", [ (Running(), "Cancelling", 3), (Success(), "Success", 2), (Submitted(), "Cancelled", 3), ], ) async def test_cancel_flow_run( self, run_query, flow_run_id, state, res_state, version ): await api.states.set_flow_run_state( flow_run_id=flow_run_id, version=1, state=state ) result = await run_query( query=self.mutation, variables={"input": {"flow_run_id": flow_run_id}}, ) assert result.data.cancel_flow_run.state == res_state fr = await models.FlowRun.where(id=flow_run_id).first({"state", "version"}) assert fr.version == version assert fr.state == res_state
async def excess_submitted_task_runs(project_id): parameters = {} # pump up the task counter by creating artificial task runs flow = prefect.Flow( name="Test Flow", schedule=prefect.schedules.IntervalSchedule( start_date=pendulum.datetime(2018, 1, 1), interval=datetime.timedelta(days=1), ), ) for i in range(config.queued_runs_returned_limit): flow.add_task(prefect.Parameter(f"x{i}", default=1)) parameters.update({f"x{i}": 1}) flow_id = await api.flows.create_flow(project_id=project_id, serialized_flow=flow.serialize()) flow_run = await api.runs.create_flow_run(flow_id=flow_id, parameters=parameters) tasks = await models.Task.where({"flow_id": {"_eq": flow_id}}).get("id") for task in tasks: task_run = await api.runs.get_or_create_task_run(flow_run_id=flow_run, task_id=task.id, map_index=None) await api.states.set_task_run_state(task_run_id=task_run, state=Submitted())
async def test_cancelling_non_running_flow_run_returns_cancelled(self, flow_run_id): result = await api.states.set_flow_run_state( flow_run_id=flow_run_id, state=Submitted() ) assert result.flow_run_id == flow_run_id assert result.state == "Submitted" result = await api.states.cancel_flow_run(flow_run_id=flow_run_id) assert result.flow_run_id == flow_run_id assert result.state == "Cancelled"
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")
def test_meta_states_dont_nest(): state = Queued(state=Pending()) for i in range(300): if i % 2: state = Queued(state=state) else: state = Submitted(state=state) assert state.state.is_pending() assert not state.state.is_meta_state() new_state = StateSchema().load(state.serialize()) assert new_state.is_meta_state() assert not new_state.state.is_meta_state()
def test_state_type_methods_with_submitted_state(self): state = Submitted() assert not state.is_cached() assert not state.is_retrying() assert not state.is_pending() assert not state.is_running() assert not state.is_finished() assert not state.is_skipped() assert not state.is_scheduled() assert not state.is_successful() assert not state.is_failed() assert not state.is_mapped() assert state.is_meta_state()
dict(state=Failed(), assert_true={"is_finished", "is_failed"}), dict(state=Finished(), assert_true={"is_finished"}), dict(state=Looped(), assert_true={"is_finished", "is_looped"}), dict(state=Mapped(), assert_true={"is_finished", "is_mapped", "is_successful"}), dict(state=Paused(), assert_true={"is_pending", "is_scheduled"}), dict(state=Pending(), assert_true={"is_pending"}), dict(state=Queued(), assert_true={"is_meta_state", "is_queued"}), dict(state=Resume(), assert_true={"is_pending", "is_scheduled"}), dict(state=Retrying(), assert_true={"is_pending", "is_scheduled", "is_retrying"}), dict(state=Running(), assert_true={"is_running"}), dict(state=Scheduled(), assert_true={"is_pending", "is_scheduled"}), dict(state=Skipped(), assert_true={"is_finished", "is_successful", "is_skipped"}), dict(state=Submitted(), assert_true={"is_meta_state", "is_submitted"}), dict(state=Success(), assert_true={"is_finished", "is_successful"}), dict(state=TimedOut(), assert_true={"is_finished", "is_failed"}), dict(state=TriggerFailed(), assert_true={"is_finished", "is_failed"}), ], ) def test_state_is_methods(state_check): """ Iterates over all of the "is_*()" methods of the state, asserting that each one is False, unless the name of that method is provided as `assert_true`. For example, if `state_check == (Pending(), {'is_pending'})`, then this method will assert that `state.is_running()` is False, `state.is_successful()` is False, etc. but `state.is_pending()` is True. """ state = state_check["state"]
class TestFlowRunStates: async def test_set_flow_run_state(self, flow_run_id): result = await api.states.set_flow_run_state( flow_run_id=flow_run_id, state=Running() ) assert result.flow_run_id == flow_run_id query = await models.FlowRun.where(id=flow_run_id).first( {"version", "state", "serialized_state"} ) assert query.version == 3 assert query.state == "Running" assert query.serialized_state["type"] == "Running" @pytest.mark.parametrize("state", [Running(), Success()]) async def test_set_flow_run_state_fails_with_wrong_flow_run_id(self, state): with pytest.raises(ValueError, match="State update failed"): await api.states.set_flow_run_state( flow_run_id=str(uuid.uuid4()), state=state ) async def test_trigger_failed_state_does_not_set_end_time(self, flow_run_id): # there is no logic in Prefect that would create this sequence of # events, but a user could manually do this await api.states.set_flow_run_state( flow_run_id=flow_run_id, state=TriggerFailed() ) flow_run_info = await models.FlowRun.where(id=flow_run_id).first( {"id", "start_time", "end_time"} ) assert not flow_run_info.start_time assert not flow_run_info.end_time @pytest.mark.parametrize( "state", [ s() for s in State.children() if not s().is_running() and not s().is_submitted() ], ) async def test_state_does_not_set_heartbeat_unless_running_or_submitted( self, state, flow_run_id ): flow_run = await models.FlowRun.where(id=flow_run_id).first({"heartbeat"}) assert flow_run.heartbeat is None dt = pendulum.now("UTC") await api.states.set_flow_run_state(flow_run_id=flow_run_id, state=state) flow_run = await models.FlowRun.where(id=flow_run_id).first({"heartbeat"}) assert flow_run.heartbeat is None @pytest.mark.parametrize("state", [Running(), Submitted()]) async def test_running_and_submitted_state_sets_heartbeat(self, state, flow_run_id): """ Both Running and Submitted states need to set heartbeats for services like Lazarus to function properly. """ flow_run = await models.FlowRun.where(id=flow_run_id).first({"heartbeat"}) assert flow_run.heartbeat is None dt = pendulum.now("UTC") await api.states.set_flow_run_state(flow_run_id=flow_run_id, state=state) flow_run = await models.FlowRun.where(id=flow_run_id).first({"heartbeat"}) assert flow_run.heartbeat > dt async def test_setting_flow_run_to_cancelled_state_sets_unfinished_task_runs_to_cancelled( self, flow_run_id ): task_runs = await models.TaskRun.where( {"flow_run_id": {"_eq": flow_run_id}} ).get({"id"}) task_run_ids = [run.id for run in task_runs] # update the state to Running await api.states.set_flow_run_state(flow_run_id=flow_run_id, state=Running()) # Currently this flow_run_id fixture has at least 3 tasks, if this # changes the test will need to be updated assert len(task_run_ids) >= 3, "flow_run_id fixture has changed" # Set one task run to pending, one to running, and the rest to success pending_task_run = task_run_ids[0] running_task_run = task_run_ids[1] rest = task_run_ids[2:] await api.states.set_task_run_state( task_run_id=pending_task_run, state=Pending() ) await api.states.set_task_run_state( task_run_id=running_task_run, state=Running() ) for task_run_id in rest: await api.states.set_task_run_state( task_run_id=task_run_id, state=Success() ) # set the flow run to a cancelled state await api.states.set_flow_run_state(flow_run_id=flow_run_id, state=Cancelled()) # Confirm the unfinished task runs have been marked as cancelled task_runs = await models.TaskRun.where( {"flow_run_id": {"_eq": flow_run_id}} ).get({"id", "state"}) new_states = {run.id: run.state for run in task_runs} assert new_states[pending_task_run] == "Cancelled" assert new_states[running_task_run] == "Cancelled" assert all(new_states[id] == "Success" for id in rest)
class TestExecuteFlowRunInSubprocess: @pytest.fixture() def mocks(self, monkeypatch): class Mocks: subprocess = MagicMock() wait_for_flow_run_start_time = MagicMock() fail_flow_run = MagicMock() mocks = Mocks() monkeypatch.setattr("prefect.backend.execution.subprocess", mocks.subprocess) monkeypatch.setattr( "prefect.backend.execution._wait_for_flow_run_start_time", mocks.wait_for_flow_run_start_time, ) monkeypatch.setattr("prefect.backend.execution._fail_flow_run", mocks.fail_flow_run) # Since we mocked the module this error cannot be used in try/catch without # replacing it with the correct type mocks.subprocess.CalledProcessError = CalledProcessError return mocks def test_creates_subprocess_correctly(self, cloud_mocks, mocks): # Returned a scheduled flow run to start cloud_mocks.FlowRunView.from_flow_run_id().state = Scheduled() # Return a finished flow run after the first iteration cloud_mocks.FlowRunView().get_latest().state = Success() execute_flow_run_in_subprocess("flow-run-id") # Should pass the correct flow run id to wait for mocks.wait_for_flow_run_start_time.assert_called_once_with( "flow-run-id") # Calls the correct command w/ environment variables mocks.subprocess.run.assert_called_once_with( [sys.executable, "-m", "prefect", "execute", "flow-run"], env={ "PREFECT__CLOUD__SEND_FLOW_RUN_LOGS": "True", "PREFECT__LOGGING__LEVEL": "INFO", "PREFECT__LOGGING__FORMAT": "[%(asctime)s] %(levelname)s - %(name)s | %(message)s", "PREFECT__LOGGING__DATEFMT": "%Y-%m-%d %H:%M:%S%z", "PREFECT__BACKEND": "cloud", "PREFECT__CLOUD__API": "https://api.prefect.io", "PREFECT__CLOUD__TENANT_ID": "", "PREFECT__CLOUD__API_KEY": cloud_mocks.Client().api_key, "PREFECT__CONTEXT__FLOW_RUN_ID": "flow-run-id", "PREFECT__CONTEXT__FLOW_ID": cloud_mocks.FlowRunView.from_flow_run_id().flow_id, "PREFECT__ENGINE__FLOW_RUNNER__DEFAULT_CLASS": "prefect.engine.cloud.CloudFlowRunner", "PREFECT__ENGINE__TASK_RUNNER__DEFAULT_CLASS": "prefect.engine.cloud.CloudTaskRunner", }, ) # Return code is checked mocks.subprocess.run().check_returncode.assert_called_once() @pytest.mark.parametrize("start_state", [Submitted(), Running()]) def test_fails_immediately_if_flow_run_is_being_executed_elsewhere( self, cloud_mocks, start_state, mocks): cloud_mocks.FlowRunView.from_flow_run_id().state = start_state with pytest.raises(RuntimeError, match="already in state"): execute_flow_run_in_subprocess("flow-run-id") def test_handles_signal_interrupt(self, cloud_mocks, mocks): cloud_mocks.FlowRunView.from_flow_run_id().state = Scheduled() mocks.subprocess.run.side_effect = KeyboardInterrupt() # Keyboard interrupt should be re-raised with pytest.raises(KeyboardInterrupt): execute_flow_run_in_subprocess("flow-run-id") # Only tried to run once mocks.subprocess.run.assert_called_once() # Flow run is failed with the proper message mocks.fail_flow_run.assert_called_once_with( flow_run_id="flow-run-id", message="Flow run received an interrupt signal.") def test_handles_unexpected_exception(self, cloud_mocks, mocks): cloud_mocks.FlowRunView.from_flow_run_id().state = Scheduled() mocks.subprocess.run.side_effect = Exception("Foobar") # Re-raised as `RuntmeError` with pytest.raises( RuntimeError, match="encountered unexpected exception during execution"): execute_flow_run_in_subprocess("flow-run-id") # Only tried to run once mocks.subprocess.run.assert_called_once() # Flow run is failed with the proper message mocks.fail_flow_run.assert_called_once_with( flow_run_id="flow-run-id", message=( "Flow run encountered unexpected exception during execution: " f"{Exception('Foobar')!r}"), ) def test_handles_bad_subprocess_result(self, cloud_mocks, mocks): cloud_mocks.FlowRunView.from_flow_run_id().state = Scheduled() mocks.subprocess.run.return_value.check_returncode.side_effect = ( CalledProcessError(cmd="foo", returncode=1)) # Re-raised as `RuntmeError` with pytest.raises(RuntimeError, match="flow run process failed"): execute_flow_run_in_subprocess("flow-run-id") # Only tried to run once mocks.subprocess.run.assert_called_once() # Flow run is not failed at this time -- left to the FlowRunner mocks.fail_flow_run.assert_not_called() def test_loops_until_flow_run_is_finished(self, cloud_mocks, mocks): cloud_mocks.FlowRunView.from_flow_run_id().state = Scheduled() cloud_mocks.FlowRunView.from_flow_run_id().get_latest.side_effect = [ MagicMock(state=Running()), MagicMock(state=Running()), MagicMock(state=Success()), ] execute_flow_run_in_subprocess("flow-run-id") # Ran the subprocess twice assert mocks.subprocess.run.call_count == 2 # Waited each time assert mocks.wait_for_flow_run_start_time.call_count == 2
async def run_scheduled(self, flow_id=None): """ Queries for any flow runs that are SCHEDULED, OR any flow runs that have SCHEDULED task runs [if the flow run itself is RUNNING]. Sets all Scheduled runs to Submitted and runs the flow. If a flow_id is provided, only flow runs of that flow are matched. """ now = pendulum.now() flow_runs = await models.FlowRun.where({ "_or": [ { "state_start_time": { "_lte": str(now) } }, { "state": { "_eq": "Running" }, "task_runs": { "state_start_time": { "_lte": str(now) } }, }, ], "flow_id": { "_eq": flow_id }, }).get( selection_set={ "id": True, "version": True, "tenant_id": True, "state": True, "serialized_state": True, "parameters": True, "flow": {"id", "environment", "name", "storage"}, with_args("task_runs", { "where": { "state_start_time": { "_lte": str(now) } } }): {"id", "version", "task_id", "serialized_state"}, }, limit=100, order_by={"state_start_time": EnumValue("asc")}, ) for fr in flow_runs: skip_counter = 0 fr_serialized_state = state_schema.load(fr.serialized_state) # set the flow run state to submitted, if it's scheduled if fr_serialized_state.is_scheduled(): try: await api.states.set_flow_run_state( flow_run_id=fr.id, state=Submitted( message="Submitted for execution", state=fr_serialized_state, ), ) except ValueError as exc: skip_counter += 1 if "Update failed" not in str(exc): raise # set each task run state to submitted, if it's scheduled for tr in fr.task_runs: tr_serialized_state = state_schema.load(tr.serialized_state) try: await api.states.set_task_run_state( task_run_id=tr.id, state=Submitted( message="Submitted for execution", state=tr_serialized_state, ), ) except ValueError as exc: skip_counter += 1 if "Update failed" not in str(exc): raise # none of the states were set, so we shouldn't bother running if skip_counter == 1 + len(fr.task_runs): continue self.logger.info(f'Submitting flow run "{fr.id}" for execution.') # run the flow self.run_flow( flow_name=fr.flow.name, storage=storage_schema.load(fr.flow.storage), environment=environment_schema.load(fr.flow.environment), config={ "cloud.api": f"http://localhost:4200", "cloud.graphql": "http://localhost:4200", "engine.flow_runner.default_class": "prefect.engine.cloud.CloudFlowRunner", "engine.task_runner.default_class": "prefect.engine.cloud.CloudTaskRunner", "engine.executor.default_class": "prefect.engine.executors.LocalExecutor", }, context={"flow_run_id": fr.id}, )
"id": "id-1", "name": "name-1", "flow_id": "flow_id-1", "serialized_state": Success(message="state-1").serialize(), "states": [ { "timestamp": pendulum.now().subtract(seconds=10).isoformat(), "serialized_state": Running(message="past-state").serialize(), }, { "timestamp": pendulum.now().subtract(seconds=20).isoformat(), "serialized_state": Submitted(message="past-state").serialize(), }, ], "parameters": { "param": "value" }, "context": { "foo": "bar" }, "labels": ["label"], "updated": pendulum.now().isoformat(), "run_config": UniversalRun().serialize(), } FLOW_RUN_DATA_2 = {
state=Failed(message="state-1"), states=[], parameters={"param": "value"}, context={"foo": "bar"}, labels=["label"], updated_at=pendulum.now(), run_config=UniversalRun(), ) # On `get_latest` return the same flow run view FAILED_FLOW_RUN_VIEW.get_latest = MagicMock(return_value=FAILED_FLOW_RUN_VIEW) SUBMITTED_FLOW_RUN_VIEW = FlowRunView( flow_run_id="flow-run-id", name="flow-run-name", flow_id="flow-id", state=Submitted(message="state-1"), states=[], parameters={"param": "value"}, context={"foo": "bar"}, labels=["label"], updated_at=pendulum.now(), run_config=UniversalRun(), ) # On `get_latest` return the same flow run view SUBMITTED_FLOW_RUN_VIEW.get_latest = MagicMock( return_value=SUBMITTED_FLOW_RUN_VIEW) TEST_FLOW_RUN_VIEW = FlowRunView( flow_run_id="flow-run-id", name="flow-run-name", flow_id="flow-id",