Skip to content

prefect.orion.models.task_runs

Functions for interacting with task run ORM objects. Intended for internal use by the Orion API.

count_task_runs async

Count task runs.

Parameters:

Name Type Description Default
session sa.orm.Session

a database session

required
flow_filter schemas.filters.FlowFilter

only count task runs whose flows match these filters

None
flow_run_filter schemas.filters.FlowRunFilter

only count task runs whose flow runs match these filters

None
task_run_filter schemas.filters.TaskRunFilter

only count task runs that match these filters

None
deployment_filter schemas.filters.DeploymentFilter

only count task runs whose deployments match these filters

None

Returns:

Name Type Description
int int

count of task runs

Source code in prefect/orion/models/task_runs.py
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
@inject_db
async def count_task_runs(
    session: sa.orm.Session,
    db: OrionDBInterface,
    flow_filter: schemas.filters.FlowFilter = None,
    flow_run_filter: schemas.filters.FlowRunFilter = None,
    task_run_filter: schemas.filters.TaskRunFilter = None,
    deployment_filter: schemas.filters.DeploymentFilter = None,
) -> int:
    """
    Count task runs.

    Args:
        session: a database session
        flow_filter: only count task runs whose flows match these filters
        flow_run_filter: only count task runs whose flow runs match these filters
        task_run_filter: only count task runs that match these filters
        deployment_filter: only count task runs whose deployments match these filters
    Returns:
        int: count of task runs
    """

    query = select(sa.func.count(sa.text("*"))).select_from(db.TaskRun)

    query = await _apply_task_run_filters(
        query,
        flow_filter=flow_filter,
        flow_run_filter=flow_run_filter,
        task_run_filter=task_run_filter,
        deployment_filter=deployment_filter,
        db=db,
    )

    result = await session.execute(query)
    return result.scalar()

create_task_run async

Creates a new task run.

If a task run with the same flow_run_id, task_key, and dynamic_key already exists, the existing task run will be returned. If the provided task run has a state attached, it will also be created.

Parameters:

Name Type Description Default
session sa.orm.Session

a database session

required
task_run schemas.core.TaskRun

a task run model

required

Returns:

Type Description

db.TaskRun: the newly-created or existing task run

Source code in prefect/orion/models/task_runs.py
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
@inject_db
async def create_task_run(
    session: sa.orm.Session,
    task_run: schemas.core.TaskRun,
    db: OrionDBInterface,
    orchestration_parameters: dict = None,
):
    """
    Creates a new task run.

    If a task run with the same flow_run_id, task_key, and dynamic_key already exists,
    the existing task run will be returned. If the provided task run has a state
    attached, it will also be created.

    Args:
        session: a database session
        task_run: a task run model

    Returns:
        db.TaskRun: the newly-created or existing task run
    """

    now = pendulum.now("UTC")

    # if a dynamic key exists, we need to guard against conflicts
    insert_stmt = (
        (await db.insert(db.TaskRun))
        .values(
            created=now,
            **task_run.dict(
                shallow=True, exclude={"state", "created"}, exclude_unset=True
            ),
        )
        .on_conflict_do_nothing(
            index_elements=db.task_run_unique_upsert_columns,
        )
    )
    await session.execute(insert_stmt)

    query = (
        sa.select(db.TaskRun)
        .where(
            sa.and_(
                db.TaskRun.flow_run_id == task_run.flow_run_id,
                db.TaskRun.task_key == task_run.task_key,
                db.TaskRun.dynamic_key == task_run.dynamic_key,
            )
        )
        .limit(1)
        .execution_options(populate_existing=True)
    )
    result = await session.execute(query)
    model = result.scalar()

    if model.created == now and task_run.state:
        await models.task_runs.set_task_run_state(
            session=session,
            task_run_id=model.id,
            state=task_run.state,
            force=True,
            orchestration_parameters=orchestration_parameters,
        )
    return model

delete_task_run async

Delete a task run by id.

Parameters:

Name Type Description Default
session sa.orm.Session

a database session

required
task_run_id UUID

the task run id to delete

required

Returns:

Name Type Description
bool bool

whether or not the task run was deleted

Source code in prefect/orion/models/task_runs.py
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
@inject_db
async def delete_task_run(
    session: sa.orm.Session, task_run_id: UUID, db: OrionDBInterface
) -> bool:
    """
    Delete a task run by id.

    Args:
        session: a database session
        task_run_id: the task run id to delete

    Returns:
        bool: whether or not the task run was deleted
    """

    result = await session.execute(
        delete(db.TaskRun).where(db.TaskRun.id == task_run_id)
    )
    return result.rowcount > 0

read_task_run async

Read a task run by id.

Parameters:

Name Type Description Default
session sa.orm.Session

a database session

required
task_run_id UUID

the task run id

required

Returns:

Type Description

db.TaskRun: the task run

Source code in prefect/orion/models/task_runs.py
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
@inject_db
async def read_task_run(
    session: sa.orm.Session, task_run_id: UUID, db: OrionDBInterface
):
    """
    Read a task run by id.

    Args:
        session: a database session
        task_run_id: the task run id

    Returns:
        db.TaskRun: the task run
    """

    model = await session.get(db.TaskRun, task_run_id)
    return model

read_task_runs async

Read task runs.

Parameters:

Name Type Description Default
session sa.orm.Session

a database session

required
flow_filter schemas.filters.FlowFilter

only select task runs whose flows match these filters

None
flow_run_filter schemas.filters.FlowRunFilter

only select task runs whose flow runs match these filters

None
task_run_filter schemas.filters.TaskRunFilter

only select task runs that match these filters

None
deployment_filter schemas.filters.DeploymentFilter

only select task runs whose deployments match these filters

None
offset int

Query offset

None
limit int

Query limit

None
sort schemas.sorting.TaskRunSort

Query sort

schemas.sorting.TaskRunSort.ID_DESC

Returns:

Type Description

List[db.TaskRun]: the task runs

Source code in prefect/orion/models/task_runs.py
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
@inject_db
async def read_task_runs(
    session: sa.orm.Session,
    db: OrionDBInterface,
    flow_filter: schemas.filters.FlowFilter = None,
    flow_run_filter: schemas.filters.FlowRunFilter = None,
    task_run_filter: schemas.filters.TaskRunFilter = None,
    deployment_filter: schemas.filters.DeploymentFilter = None,
    offset: int = None,
    limit: int = None,
    sort: schemas.sorting.TaskRunSort = schemas.sorting.TaskRunSort.ID_DESC,
):
    """
    Read task runs.

    Args:
        session: a database session
        flow_filter: only select task runs whose flows match these filters
        flow_run_filter: only select task runs whose flow runs match these filters
        task_run_filter: only select task runs that match these filters
        deployment_filter: only select task runs whose deployments match these filters
        offset: Query offset
        limit: Query limit
        sort: Query sort

    Returns:
        List[db.TaskRun]: the task runs
    """

    query = select(db.TaskRun).order_by(sort.as_sql_sort(db))

    query = await _apply_task_run_filters(
        query,
        flow_filter=flow_filter,
        flow_run_filter=flow_run_filter,
        task_run_filter=task_run_filter,
        deployment_filter=deployment_filter,
        db=db,
    )

    if offset is not None:
        query = query.offset(offset)

    if limit is not None:
        query = query.limit(limit)

    result = await session.execute(query)
    return result.scalars().unique().all()

set_task_run_state async

Creates a new orchestrated task run state.

Setting a new state on a run is the one of the principal actions that is governed by Orion's orchestration logic. Setting a new run state will not guarantee creation, but instead trigger orchestration rules to govern the proposed state input. If the state is considered valid, it will be written to the database. Otherwise, a it's possible a different state, or no state, will be created. A force flag is supplied to bypass a subset of orchestration logic.

Parameters:

Name Type Description Default
session sa.orm.Session

a database session

required
task_run_id UUID

the task run id

required
state schemas.states.State

a task run state model

required
force bool

if False, orchestration rules will be applied that may alter or prevent the state transition. If True, orchestration rules are not applied.

False

Returns:

Type Description
OrchestrationResult

OrchestrationResult object

Source code in prefect/orion/models/task_runs.py
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
async def set_task_run_state(
    session: sa.orm.Session,
    task_run_id: UUID,
    state: schemas.states.State,
    force: bool = False,
    task_policy: BaseOrchestrationPolicy = None,
    orchestration_parameters: dict = None,
) -> OrchestrationResult:
    """
    Creates a new orchestrated task run state.

    Setting a new state on a run is the one of the principal actions that is governed by
    Orion's orchestration logic. Setting a new run state will not guarantee creation,
    but instead trigger orchestration rules to govern the proposed `state` input. If
    the state is considered valid, it will be written to the database. Otherwise, a
    it's possible a different state, or no state, will be created. A `force` flag is
    supplied to bypass a subset of orchestration logic.

    Args:
        session: a database session
        task_run_id: the task run id
        state: a task run state model
        force: if False, orchestration rules will be applied that may alter or prevent
            the state transition. If True, orchestration rules are not applied.

    Returns:
        OrchestrationResult object
    """

    # load the task run
    run = await models.task_runs.read_task_run(session=session, task_run_id=task_run_id)

    if not run:
        raise ObjectNotFoundError(f"Task run with id {task_run_id} not found")

    initial_state = run.state.as_state() if run.state else None
    initial_state_type = initial_state.type if initial_state else None
    proposed_state_type = state.type if state else None
    intended_transition = (initial_state_type, proposed_state_type)

    if force or task_policy is None:
        task_policy = MinimalTaskPolicy

    orchestration_rules = task_policy.compile_transition_rules(*intended_transition)
    global_rules = GlobalTaskPolicy.compile_transition_rules(*intended_transition)

    context = TaskOrchestrationContext(
        session=session,
        run=run,
        initial_state=initial_state,
        proposed_state=state,
    )

    if orchestration_parameters is not None:
        context.parameters = orchestration_parameters

    # apply orchestration rules and create the new task run state
    async with contextlib.AsyncExitStack() as stack:
        for rule in orchestration_rules:
            context = await stack.enter_async_context(
                rule(context, *intended_transition)
            )

        for rule in global_rules:
            context = await stack.enter_async_context(
                rule(context, *intended_transition)
            )

        await context.validate_proposed_state()

    if context.orchestration_error is not None:
        raise context.orchestration_error

    result = OrchestrationResult(
        state=context.validated_state,
        status=context.response_status,
        details=context.response_details,
    )

    return result