prefect-client 2.18.3__py3-none-any.whl → 2.19.0__py3-none-any.whl

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
Files changed (41) hide show
  1. prefect/__init__.py +1 -15
  2. prefect/_internal/concurrency/cancellation.py +2 -0
  3. prefect/_internal/schemas/validators.py +10 -0
  4. prefect/_vendor/starlette/testclient.py +1 -1
  5. prefect/blocks/notifications.py +6 -6
  6. prefect/client/base.py +244 -1
  7. prefect/client/cloud.py +4 -2
  8. prefect/client/orchestration.py +515 -106
  9. prefect/client/schemas/actions.py +58 -8
  10. prefect/client/schemas/objects.py +15 -1
  11. prefect/client/schemas/responses.py +19 -0
  12. prefect/client/schemas/schedules.py +1 -1
  13. prefect/client/utilities.py +2 -2
  14. prefect/concurrency/asyncio.py +34 -4
  15. prefect/concurrency/sync.py +40 -6
  16. prefect/context.py +2 -2
  17. prefect/engine.py +2 -2
  18. prefect/events/clients.py +2 -2
  19. prefect/flows.py +91 -17
  20. prefect/infrastructure/process.py +0 -17
  21. prefect/logging/formatters.py +1 -4
  22. prefect/new_flow_engine.py +137 -168
  23. prefect/new_task_engine.py +137 -202
  24. prefect/runner/__init__.py +1 -1
  25. prefect/runner/runner.py +2 -107
  26. prefect/settings.py +11 -0
  27. prefect/tasks.py +76 -57
  28. prefect/types/__init__.py +27 -5
  29. prefect/utilities/annotations.py +1 -8
  30. prefect/utilities/asyncutils.py +4 -0
  31. prefect/utilities/engine.py +106 -1
  32. prefect/utilities/schema_tools/__init__.py +6 -1
  33. prefect/utilities/schema_tools/validation.py +25 -8
  34. prefect/utilities/timeout.py +34 -0
  35. prefect/workers/base.py +7 -3
  36. prefect/workers/process.py +0 -17
  37. {prefect_client-2.18.3.dist-info → prefect_client-2.19.0.dist-info}/METADATA +1 -1
  38. {prefect_client-2.18.3.dist-info → prefect_client-2.19.0.dist-info}/RECORD +41 -40
  39. {prefect_client-2.18.3.dist-info → prefect_client-2.19.0.dist-info}/LICENSE +0 -0
  40. {prefect_client-2.18.3.dist-info → prefect_client-2.19.0.dist-info}/WHEEL +0 -0
  41. {prefect_client-2.18.3.dist-info → prefect_client-2.19.0.dist-info}/top_level.txt +0 -0
@@ -33,7 +33,6 @@ from prefect.settings import (
33
33
  PREFECT_API_SERVICES_TRIGGERS_ENABLED,
34
34
  PREFECT_EXPERIMENTAL_EVENTS,
35
35
  )
36
- from prefect.utilities.asyncutils import run_sync
37
36
 
38
37
  if HAS_PYDANTIC_V2:
39
38
  import pydantic.v1 as pydantic
@@ -121,6 +120,7 @@ from prefect.client.schemas.objects import (
121
120
  from prefect.client.schemas.responses import (
122
121
  DeploymentResponse,
123
122
  FlowRunResponse,
123
+ GlobalConcurrencyLimitResponse,
124
124
  WorkerFlowRunResponse,
125
125
  )
126
126
  from prefect.client.schemas.schedules import SCHEDULE_TYPES
@@ -154,7 +154,13 @@ if TYPE_CHECKING:
154
154
  from prefect.flows import Flow as FlowObject
155
155
  from prefect.tasks import Task as TaskObject
156
156
 
157
- from prefect.client.base import ASGIApp, PrefectHttpxClient, app_lifespan_context
157
+ from prefect.client.base import (
158
+ ASGIApp,
159
+ PrefectHttpxAsyncClient,
160
+ PrefectHttpxSyncClient,
161
+ PrefectHttpxSyncEphemeralClient,
162
+ app_lifespan_context,
163
+ )
158
164
 
159
165
  P = ParamSpec("P")
160
166
  R = TypeVar("R")
@@ -174,7 +180,7 @@ class ServerType(AutoEnum):
174
180
 
175
181
  def get_client(
176
182
  httpx_settings: Optional[Dict[str, Any]] = None, sync_client: bool = False
177
- ) -> "PrefectClient":
183
+ ) -> Union["PrefectClient", "SyncPrefectClient"]:
178
184
  """
179
185
  Retrieve a HTTP client for communicating with the Prefect REST API.
180
186
 
@@ -361,7 +367,7 @@ class PrefectClient:
361
367
  and PREFECT_CLIENT_CSRF_SUPPORT_ENABLED.value()
362
368
  )
363
369
 
364
- self._client = PrefectHttpxClient(
370
+ self._client = PrefectHttpxAsyncClient(
365
371
  **httpx_settings, enable_csrf_support=enable_csrf_support
366
372
  )
367
373
  self._loop = None
@@ -1952,7 +1958,7 @@ class PrefectClient:
1952
1958
  kwargs = {}
1953
1959
  if active is not None:
1954
1960
  kwargs["active"] = active
1955
- elif schedule is not None:
1961
+ if schedule is not None:
1956
1962
  kwargs["schedule"] = schedule
1957
1963
 
1958
1964
  deployment_schedule_update = DeploymentScheduleUpdate(**kwargs)
@@ -3054,10 +3060,10 @@ class PrefectClient:
3054
3060
 
3055
3061
  async def read_global_concurrency_limit_by_name(
3056
3062
  self, name: str
3057
- ) -> Dict[str, object]:
3063
+ ) -> GlobalConcurrencyLimitResponse:
3058
3064
  try:
3059
3065
  response = await self._client.get(f"/v2/concurrency_limits/{name}")
3060
- return response.json()
3066
+ return GlobalConcurrencyLimitResponse.parse_obj(response.json())
3061
3067
  except httpx.HTTPStatusError as e:
3062
3068
  if e.response.status_code == status.HTTP_404_NOT_FOUND:
3063
3069
  raise prefect.exceptions.ObjectNotFound(http_exc=e) from e
@@ -3066,7 +3072,7 @@ class PrefectClient:
3066
3072
 
3067
3073
  async def read_global_concurrency_limits(
3068
3074
  self, limit: int = 10, offset: int = 0
3069
- ) -> List[Dict[str, object]]:
3075
+ ) -> List[GlobalConcurrencyLimitResponse]:
3070
3076
  response = await self._client.post(
3071
3077
  "/v2/concurrency_limits/filter",
3072
3078
  json={
@@ -3074,7 +3080,9 @@ class PrefectClient:
3074
3080
  "offset": offset,
3075
3081
  },
3076
3082
  )
3077
- return response.json()
3083
+ return pydantic.parse_obj_as(
3084
+ List[GlobalConcurrencyLimitResponse], response.json()
3085
+ )
3078
3086
 
3079
3087
  async def create_flow_run_input(
3080
3088
  self, flow_run_id: UUID, key: str, value: str, sender: Optional[str] = None
@@ -3365,7 +3373,7 @@ class SyncPrefectClient:
3365
3373
  api_key: An optional API key for authentication.
3366
3374
  api_version: The API version this client is compatible with.
3367
3375
  httpx_settings: An optional dictionary of settings to pass to the underlying
3368
- `httpx.AsyncClient`
3376
+ `httpx.Client`
3369
3377
 
3370
3378
  Examples:
3371
3379
 
@@ -3386,38 +3394,416 @@ class SyncPrefectClient:
3386
3394
  self,
3387
3395
  api: Union[str, ASGIApp],
3388
3396
  *,
3389
- api_key: Optional[str] = None,
3390
- api_version: Optional[str] = None,
3397
+ api_key: str = None,
3398
+ api_version: str = None,
3391
3399
  httpx_settings: Optional[Dict[str, Any]] = None,
3392
3400
  ) -> None:
3393
- self._prefect_client = PrefectClient(
3394
- api=api,
3395
- api_key=api_key,
3396
- api_version=api_version,
3397
- httpx_settings=httpx_settings,
3401
+ httpx_settings = httpx_settings.copy() if httpx_settings else {}
3402
+ httpx_settings.setdefault("headers", {})
3403
+
3404
+ if PREFECT_API_TLS_INSECURE_SKIP_VERIFY:
3405
+ httpx_settings.setdefault("verify", False)
3406
+ else:
3407
+ cert_file = PREFECT_API_SSL_CERT_FILE.value()
3408
+ if not cert_file:
3409
+ cert_file = certifi.where()
3410
+ httpx_settings.setdefault("verify", cert_file)
3411
+
3412
+ if api_version is None:
3413
+ api_version = SERVER_API_VERSION
3414
+ httpx_settings["headers"].setdefault("X-PREFECT-API-VERSION", api_version)
3415
+ if api_key:
3416
+ httpx_settings["headers"].setdefault("Authorization", f"Bearer {api_key}")
3417
+
3418
+ # Context management
3419
+ self._ephemeral_app: Optional[ASGIApp] = None
3420
+ self.manage_lifespan = True
3421
+ self.server_type: ServerType
3422
+
3423
+ self._closed = False
3424
+ self._started = False
3425
+
3426
+ # Connect to an external application
3427
+ if isinstance(api, str):
3428
+ if httpx_settings.get("app"):
3429
+ raise ValueError(
3430
+ "Invalid httpx settings: `app` cannot be set when providing an "
3431
+ "api url. `app` is only for use with ephemeral instances. Provide "
3432
+ "it as the `api` parameter instead."
3433
+ )
3434
+ httpx_settings.setdefault("base_url", api)
3435
+
3436
+ # See https://www.python-httpx.org/advanced/#pool-limit-configuration
3437
+ httpx_settings.setdefault(
3438
+ "limits",
3439
+ httpx.Limits(
3440
+ # We see instability when allowing the client to open many connections at once.
3441
+ # Limiting concurrency results in more stable performance.
3442
+ max_connections=16,
3443
+ max_keepalive_connections=8,
3444
+ # The Prefect Cloud LB will keep connections alive for 30s.
3445
+ # Only allow the client to keep connections alive for 25s.
3446
+ keepalive_expiry=25,
3447
+ ),
3448
+ )
3449
+
3450
+ # See https://www.python-httpx.org/http2/
3451
+ # Enabling HTTP/2 support on the client does not necessarily mean that your requests
3452
+ # and responses will be transported over HTTP/2, since both the client and the server
3453
+ # need to support HTTP/2. If you connect to a server that only supports HTTP/1.1 the
3454
+ # client will use a standard HTTP/1.1 connection instead.
3455
+ httpx_settings.setdefault("http2", PREFECT_API_ENABLE_HTTP2.value())
3456
+
3457
+ self.server_type = (
3458
+ ServerType.CLOUD
3459
+ if api.startswith(PREFECT_CLOUD_API_URL.value())
3460
+ else ServerType.SERVER
3461
+ )
3462
+
3463
+ # Connect to an in-process application
3464
+ elif isinstance(api, ASGIApp):
3465
+ self._ephemeral_app = api
3466
+ self.server_type = ServerType.EPHEMERAL
3467
+
3468
+ else:
3469
+ raise TypeError(
3470
+ f"Unexpected type {type(api).__name__!r} for argument `api`. Expected"
3471
+ " 'str' or 'ASGIApp/FastAPI'"
3472
+ )
3473
+
3474
+ # See https://www.python-httpx.org/advanced/#timeout-configuration
3475
+ httpx_settings.setdefault(
3476
+ "timeout",
3477
+ httpx.Timeout(
3478
+ connect=PREFECT_API_REQUEST_TIMEOUT.value(),
3479
+ read=PREFECT_API_REQUEST_TIMEOUT.value(),
3480
+ write=PREFECT_API_REQUEST_TIMEOUT.value(),
3481
+ pool=PREFECT_API_REQUEST_TIMEOUT.value(),
3482
+ ),
3398
3483
  )
3399
3484
 
3400
- def __enter__(self):
3401
- run_sync(self._prefect_client.__aenter__())
3485
+ if not PREFECT_UNIT_TEST_MODE:
3486
+ httpx_settings.setdefault("follow_redirects", True)
3487
+
3488
+ enable_csrf_support = (
3489
+ self.server_type != ServerType.CLOUD
3490
+ and PREFECT_CLIENT_CSRF_SUPPORT_ENABLED.value()
3491
+ )
3492
+
3493
+ if self.server_type == ServerType.EPHEMERAL:
3494
+ self._client = PrefectHttpxSyncEphemeralClient(
3495
+ api, base_url="http://ephemeral-prefect/api"
3496
+ )
3497
+ else:
3498
+ self._client = PrefectHttpxSyncClient(
3499
+ **httpx_settings, enable_csrf_support=enable_csrf_support
3500
+ )
3501
+
3502
+ # See https://www.python-httpx.org/advanced/#custom-transports
3503
+ #
3504
+ # If we're using an HTTP/S client (not the ephemeral client), adjust the
3505
+ # transport to add retries _after_ it is instantiated. If we alter the transport
3506
+ # before instantiation, the transport will not be aware of proxies unless we
3507
+ # reproduce all of the logic to make it so.
3508
+ #
3509
+ # Only alter the transport to set our default of 3 retries, don't modify any
3510
+ # transport a user may have provided via httpx_settings.
3511
+ #
3512
+ # Making liberal use of getattr and isinstance checks here to avoid any
3513
+ # surprises if the internals of httpx or httpcore change on us
3514
+ if isinstance(api, str) and not httpx_settings.get("transport"):
3515
+ transport_for_url = getattr(self._client, "_transport_for_url", None)
3516
+ if callable(transport_for_url):
3517
+ server_transport = transport_for_url(httpx.URL(api))
3518
+ if isinstance(server_transport, httpx.HTTPTransport):
3519
+ pool = getattr(server_transport, "_pool", None)
3520
+ if isinstance(pool, httpcore.ConnectionPool):
3521
+ pool._retries = 3
3522
+
3523
+ self.logger = get_logger("client")
3524
+
3525
+ @property
3526
+ def api_url(self) -> httpx.URL:
3527
+ """
3528
+ Get the base URL for the API.
3529
+ """
3530
+ return self._client.base_url
3531
+
3532
+ # Context management ----------------------------------------------------------------
3533
+
3534
+ def __enter__(self) -> "SyncPrefectClient":
3535
+ """
3536
+ Start the client.
3537
+
3538
+ If the client is already started, this will raise an exception.
3539
+
3540
+ If the client is already closed, this will raise an exception. Use a new client
3541
+ instance instead.
3542
+ """
3543
+ if self._closed:
3544
+ # httpx.Client does not allow reuse so we will not either.
3545
+ raise RuntimeError(
3546
+ "The client cannot be started again after closing. "
3547
+ "Retrieve a new client with `get_client()` instead."
3548
+ )
3549
+
3550
+ if self._started:
3551
+ # httpx.Client does not allow reentrancy so we will not either.
3552
+ raise RuntimeError("The client cannot be started more than once.")
3553
+ self._client.__enter__()
3554
+ self._started = True
3555
+
3402
3556
  return self
3403
3557
 
3404
- def __exit__(self, *exc_info):
3405
- return run_sync(self._prefect_client.__aexit__(*exc_info))
3558
+ def __exit__(self, *exc_info) -> None:
3559
+ """
3560
+ Shutdown the client.
3561
+ """
3562
+ self._closed = True
3563
+ self._client.__exit__(*exc_info)
3406
3564
 
3407
- async def __aenter__(self):
3408
- raise RuntimeError(
3409
- "The `SyncPrefectClient` must be entered with a sync context. Use '"
3410
- "with SyncPrefectClient(...)' not 'async with SyncPrefectClient(...)'"
3411
- )
3565
+ # API methods ----------------------------------------------------------------------
3566
+
3567
+ def api_healthcheck(self) -> Optional[Exception]:
3568
+ """
3569
+ Attempts to connect to the API and returns the encountered exception if not
3570
+ successful.
3412
3571
 
3413
- async def __aexit__(self, *_):
3414
- assert False, "This should never be called but must be defined for __aenter__"
3572
+ If successful, returns `None`.
3573
+ """
3574
+ try:
3575
+ self._client.get("/health")
3576
+ return None
3577
+ except Exception as exc:
3578
+ return exc
3415
3579
 
3416
3580
  def hello(self) -> httpx.Response:
3417
3581
  """
3418
3582
  Send a GET request to /hello for testing purposes.
3419
3583
  """
3420
- return run_sync(self._prefect_client.hello())
3584
+ return self._client.get("/hello")
3585
+
3586
+ def create_flow(self, flow: "FlowObject") -> UUID:
3587
+ """
3588
+ Create a flow in the Prefect API.
3589
+
3590
+ Args:
3591
+ flow: a [Flow][prefect.flows.Flow] object
3592
+
3593
+ Raises:
3594
+ httpx.RequestError: if a flow was not created for any reason
3595
+
3596
+ Returns:
3597
+ the ID of the flow in the backend
3598
+ """
3599
+ return self.create_flow_from_name(flow.name)
3600
+
3601
+ def create_flow_from_name(self, flow_name: str) -> UUID:
3602
+ """
3603
+ Create a flow in the Prefect API.
3604
+
3605
+ Args:
3606
+ flow_name: the name of the new flow
3607
+
3608
+ Raises:
3609
+ httpx.RequestError: if a flow was not created for any reason
3610
+
3611
+ Returns:
3612
+ the ID of the flow in the backend
3613
+ """
3614
+ flow_data = FlowCreate(name=flow_name)
3615
+ response = self._client.post(
3616
+ "/flows/", json=flow_data.dict(json_compatible=True)
3617
+ )
3618
+
3619
+ flow_id = response.json().get("id")
3620
+ if not flow_id:
3621
+ raise httpx.RequestError(f"Malformed response: {response}")
3622
+
3623
+ # Return the id of the created flow
3624
+ return UUID(flow_id)
3625
+
3626
+ def create_flow_run(
3627
+ self,
3628
+ flow: "FlowObject",
3629
+ name: Optional[str] = None,
3630
+ parameters: Optional[Dict[str, Any]] = None,
3631
+ context: Optional[Dict[str, Any]] = None,
3632
+ tags: Optional[Iterable[str]] = None,
3633
+ parent_task_run_id: Optional[UUID] = None,
3634
+ state: Optional["prefect.states.State"] = None,
3635
+ ) -> FlowRun:
3636
+ """
3637
+ Create a flow run for a flow.
3638
+
3639
+ Args:
3640
+ flow: The flow model to create the flow run for
3641
+ name: An optional name for the flow run
3642
+ parameters: Parameter overrides for this flow run.
3643
+ context: Optional run context data
3644
+ tags: a list of tags to apply to this flow run
3645
+ parent_task_run_id: if a subflow run is being created, the placeholder task
3646
+ run identifier in the parent flow
3647
+ state: The initial state for the run. If not provided, defaults to
3648
+ `Scheduled` for now. Should always be a `Scheduled` type.
3649
+
3650
+ Raises:
3651
+ httpx.RequestError: if the Prefect API does not successfully create a run for any reason
3652
+
3653
+ Returns:
3654
+ The flow run model
3655
+ """
3656
+ parameters = parameters or {}
3657
+ context = context or {}
3658
+
3659
+ if state is None:
3660
+ state = prefect.states.Pending()
3661
+
3662
+ # Retrieve the flow id
3663
+ flow_id = self.create_flow(flow)
3664
+
3665
+ flow_run_create = FlowRunCreate(
3666
+ flow_id=flow_id,
3667
+ flow_version=flow.version,
3668
+ name=name,
3669
+ parameters=parameters,
3670
+ context=context,
3671
+ tags=list(tags or []),
3672
+ parent_task_run_id=parent_task_run_id,
3673
+ state=state.to_state_create(),
3674
+ empirical_policy=FlowRunPolicy(
3675
+ retries=flow.retries,
3676
+ retry_delay=flow.retry_delay_seconds,
3677
+ ),
3678
+ )
3679
+
3680
+ flow_run_create_json = flow_run_create.dict(json_compatible=True)
3681
+ response = self._client.post("/flow_runs/", json=flow_run_create_json)
3682
+ flow_run = FlowRun.parse_obj(response.json())
3683
+
3684
+ # Restore the parameters to the local objects to retain expectations about
3685
+ # Python objects
3686
+ flow_run.parameters = parameters
3687
+
3688
+ return flow_run
3689
+
3690
+ def read_flow_run(self, flow_run_id: UUID) -> FlowRun:
3691
+ """
3692
+ Query the Prefect API for a flow run by id.
3693
+
3694
+ Args:
3695
+ flow_run_id: the flow run ID of interest
3696
+
3697
+ Returns:
3698
+ a Flow Run model representation of the flow run
3699
+ """
3700
+ try:
3701
+ response = self._client.get(f"/flow_runs/{flow_run_id}")
3702
+ except httpx.HTTPStatusError as e:
3703
+ if e.response.status_code == 404:
3704
+ raise prefect.exceptions.ObjectNotFound(http_exc=e) from e
3705
+ else:
3706
+ raise
3707
+ return FlowRun.parse_obj(response.json())
3708
+
3709
+ def read_flow_runs(
3710
+ self,
3711
+ *,
3712
+ flow_filter: FlowFilter = None,
3713
+ flow_run_filter: FlowRunFilter = None,
3714
+ task_run_filter: TaskRunFilter = None,
3715
+ deployment_filter: DeploymentFilter = None,
3716
+ work_pool_filter: WorkPoolFilter = None,
3717
+ work_queue_filter: WorkQueueFilter = None,
3718
+ sort: FlowRunSort = None,
3719
+ limit: int = None,
3720
+ offset: int = 0,
3721
+ ) -> List[FlowRun]:
3722
+ """
3723
+ Query the Prefect API for flow runs. Only flow runs matching all criteria will
3724
+ be returned.
3725
+
3726
+ Args:
3727
+ flow_filter: filter criteria for flows
3728
+ flow_run_filter: filter criteria for flow runs
3729
+ task_run_filter: filter criteria for task runs
3730
+ deployment_filter: filter criteria for deployments
3731
+ work_pool_filter: filter criteria for work pools
3732
+ work_queue_filter: filter criteria for work pool queues
3733
+ sort: sort criteria for the flow runs
3734
+ limit: limit for the flow run query
3735
+ offset: offset for the flow run query
3736
+
3737
+ Returns:
3738
+ a list of Flow Run model representations
3739
+ of the flow runs
3740
+ """
3741
+ body = {
3742
+ "flows": flow_filter.dict(json_compatible=True) if flow_filter else None,
3743
+ "flow_runs": (
3744
+ flow_run_filter.dict(json_compatible=True, exclude_unset=True)
3745
+ if flow_run_filter
3746
+ else None
3747
+ ),
3748
+ "task_runs": (
3749
+ task_run_filter.dict(json_compatible=True) if task_run_filter else None
3750
+ ),
3751
+ "deployments": (
3752
+ deployment_filter.dict(json_compatible=True)
3753
+ if deployment_filter
3754
+ else None
3755
+ ),
3756
+ "work_pools": (
3757
+ work_pool_filter.dict(json_compatible=True)
3758
+ if work_pool_filter
3759
+ else None
3760
+ ),
3761
+ "work_pool_queues": (
3762
+ work_queue_filter.dict(json_compatible=True)
3763
+ if work_queue_filter
3764
+ else None
3765
+ ),
3766
+ "sort": sort,
3767
+ "limit": limit,
3768
+ "offset": offset,
3769
+ }
3770
+
3771
+ response = self._client.post("/flow_runs/filter", json=body)
3772
+ return pydantic.parse_obj_as(List[FlowRun], response.json())
3773
+
3774
+ def set_flow_run_state(
3775
+ self,
3776
+ flow_run_id: UUID,
3777
+ state: "prefect.states.State",
3778
+ force: bool = False,
3779
+ ) -> OrchestrationResult:
3780
+ """
3781
+ Set the state of a flow run.
3782
+
3783
+ Args:
3784
+ flow_run_id: the id of the flow run
3785
+ state: the state to set
3786
+ force: if True, disregard orchestration logic when setting the state,
3787
+ forcing the Prefect API to accept the state
3788
+
3789
+ Returns:
3790
+ an OrchestrationResult model representation of state orchestration output
3791
+ """
3792
+ state_create = state.to_state_create()
3793
+ state_create.state_details.flow_run_id = flow_run_id
3794
+ state_create.state_details.transition_id = uuid4()
3795
+ try:
3796
+ response = self._client.post(
3797
+ f"/flow_runs/{flow_run_id}/set_state",
3798
+ json=dict(state=state_create.dict(json_compatible=True), force=force),
3799
+ )
3800
+ except httpx.HTTPStatusError as e:
3801
+ if e.response.status_code == status.HTTP_404_NOT_FOUND:
3802
+ raise prefect.exceptions.ObjectNotFound(http_exc=e) from e
3803
+ else:
3804
+ raise
3805
+
3806
+ return OrchestrationResult.parse_obj(response.json())
3421
3807
 
3422
3808
  def create_task_run(
3423
3809
  self,
@@ -3457,17 +3843,94 @@ class SyncPrefectClient:
3457
3843
  Returns:
3458
3844
  The created task run.
3459
3845
  """
3460
- return run_sync(
3461
- self._prefect_client.create_task_run(
3462
- task=task,
3463
- flow_run_id=flow_run_id,
3464
- dynamic_key=dynamic_key,
3465
- name=name,
3466
- extra_tags=extra_tags,
3467
- state=state,
3468
- task_inputs=task_inputs,
3469
- )
3846
+ tags = set(task.tags).union(extra_tags or [])
3847
+
3848
+ if state is None:
3849
+ state = prefect.states.Pending()
3850
+
3851
+ task_run_data = TaskRunCreate(
3852
+ name=name,
3853
+ flow_run_id=flow_run_id,
3854
+ task_key=task.task_key,
3855
+ dynamic_key=dynamic_key,
3856
+ tags=list(tags),
3857
+ task_version=task.version,
3858
+ empirical_policy=TaskRunPolicy(
3859
+ retries=task.retries,
3860
+ retry_delay=task.retry_delay_seconds,
3861
+ retry_jitter_factor=task.retry_jitter_factor,
3862
+ ),
3863
+ state=state.to_state_create(),
3864
+ task_inputs=task_inputs or {},
3865
+ )
3866
+
3867
+ response = self._client.post(
3868
+ "/task_runs/", json=task_run_data.dict(json_compatible=True)
3470
3869
  )
3870
+ return TaskRun.parse_obj(response.json())
3871
+
3872
+ def read_task_run(self, task_run_id: UUID) -> TaskRun:
3873
+ """
3874
+ Query the Prefect API for a task run by id.
3875
+
3876
+ Args:
3877
+ task_run_id: the task run ID of interest
3878
+
3879
+ Returns:
3880
+ a Task Run model representation of the task run
3881
+ """
3882
+ response = self._client.get(f"/task_runs/{task_run_id}")
3883
+ return TaskRun.parse_obj(response.json())
3884
+
3885
+ def read_task_runs(
3886
+ self,
3887
+ *,
3888
+ flow_filter: FlowFilter = None,
3889
+ flow_run_filter: FlowRunFilter = None,
3890
+ task_run_filter: TaskRunFilter = None,
3891
+ deployment_filter: DeploymentFilter = None,
3892
+ sort: TaskRunSort = None,
3893
+ limit: int = None,
3894
+ offset: int = 0,
3895
+ ) -> List[TaskRun]:
3896
+ """
3897
+ Query the Prefect API for task runs. Only task runs matching all criteria will
3898
+ be returned.
3899
+
3900
+ Args:
3901
+ flow_filter: filter criteria for flows
3902
+ flow_run_filter: filter criteria for flow runs
3903
+ task_run_filter: filter criteria for task runs
3904
+ deployment_filter: filter criteria for deployments
3905
+ sort: sort criteria for the task runs
3906
+ limit: a limit for the task run query
3907
+ offset: an offset for the task run query
3908
+
3909
+ Returns:
3910
+ a list of Task Run model representations
3911
+ of the task runs
3912
+ """
3913
+ body = {
3914
+ "flows": flow_filter.dict(json_compatible=True) if flow_filter else None,
3915
+ "flow_runs": (
3916
+ flow_run_filter.dict(json_compatible=True, exclude_unset=True)
3917
+ if flow_run_filter
3918
+ else None
3919
+ ),
3920
+ "task_runs": (
3921
+ task_run_filter.dict(json_compatible=True) if task_run_filter else None
3922
+ ),
3923
+ "deployments": (
3924
+ deployment_filter.dict(json_compatible=True)
3925
+ if deployment_filter
3926
+ else None
3927
+ ),
3928
+ "sort": sort,
3929
+ "limit": limit,
3930
+ "offset": offset,
3931
+ }
3932
+ response = self._client.post("/task_runs/filter", json=body)
3933
+ return pydantic.parse_obj_as(List[TaskRun], response.json())
3471
3934
 
3472
3935
  def set_task_run_state(
3473
3936
  self,
@@ -3487,79 +3950,25 @@ class SyncPrefectClient:
3487
3950
  Returns:
3488
3951
  an OrchestrationResult model representation of state orchestration output
3489
3952
  """
3490
- return run_sync(
3491
- self._prefect_client.set_task_run_state(
3492
- task_run_id=task_run_id,
3493
- state=state,
3494
- force=force,
3495
- )
3496
- )
3497
-
3498
- def create_flow_run(
3499
- self,
3500
- flow_id: UUID,
3501
- parameters: Optional[Dict[str, Any]] = None,
3502
- context: Optional[Dict[str, Any]] = None,
3503
- scheduled_start_time: Optional[datetime.datetime] = None,
3504
- run_name: Optional[str] = None,
3505
- labels: Optional[List[str]] = None,
3506
- parameters_json: Optional[str] = None,
3507
- run_config: Optional[Dict[str, Any]] = None,
3508
- idempotency_key: Optional[str] = None,
3509
- ) -> FlowRunResponse:
3510
- """
3511
- Create a new flow run.
3512
-
3513
- Args:
3514
- - flow_id (UUID): the ID of the flow to create a run for
3515
- - parameters (Optional[Dict[str, Any]]): a dictionary of parameter values to pass to the flow
3516
- - context (Optional[Dict[str, Any]]): a dictionary of context values to pass to the flow
3517
- - scheduled_start_time (Optional[datetime.datetime]): the scheduled start time for the flow run
3518
- - run_name (Optional[str]): a name to assign to the flow run
3519
- - labels (Optional[List[str]]): a list of labels to assign to the flow run
3520
- - parameters_json (Optional[str]): a JSON string of parameter values to pass to the flow
3521
- - run_config (Optional[Dict[str, Any]]): a dictionary of run configuration options
3522
- - idempotency_key (Optional[str]): a key to ensure idempotency when creating the flow run
3523
-
3524
- Returns:
3525
- - FlowRunResponse: the created flow run
3526
- """
3527
- return run_sync(
3528
- self._prefect_client.create_flow_run(
3529
- flow_id=flow_id,
3530
- parameters=parameters,
3531
- context=context,
3532
- scheduled_start_time=scheduled_start_time,
3533
- run_name=run_name,
3534
- labels=labels,
3535
- parameters_json=parameters_json,
3536
- run_config=run_config,
3537
- idempotency_key=idempotency_key,
3538
- )
3953
+ state_create = state.to_state_create()
3954
+ state_create.state_details.task_run_id = task_run_id
3955
+ response = self._client.post(
3956
+ f"/task_runs/{task_run_id}/set_state",
3957
+ json=dict(state=state_create.dict(json_compatible=True), force=force),
3539
3958
  )
3959
+ return OrchestrationResult.parse_obj(response.json())
3540
3960
 
3541
- async def set_flow_run_state(
3542
- self,
3543
- flow_run_id: UUID,
3544
- state: "prefect.states.State",
3545
- force: bool = False,
3546
- ) -> OrchestrationResult:
3961
+ def read_task_run_states(self, task_run_id: UUID) -> List[prefect.states.State]:
3547
3962
  """
3548
- Set the state of a flow run.
3963
+ Query for the states of a task run
3549
3964
 
3550
3965
  Args:
3551
- flow_run_id: the id of the flow run
3552
- state: the state to set
3553
- force: if True, disregard orchestration logic when setting the state,
3554
- forcing the Prefect API to accept the state
3966
+ task_run_id: the id of the task run
3555
3967
 
3556
3968
  Returns:
3557
- an OrchestrationResult model representation of state orchestration output
3969
+ a list of State model representations of the task run states
3558
3970
  """
3559
- return run_sync(
3560
- self._prefect_client.set_flow_run_state(
3561
- flow_run_id=flow_run_id,
3562
- state=state,
3563
- force=force,
3564
- )
3971
+ response = self._client.get(
3972
+ "/task_run_states/", params=dict(task_run_id=str(task_run_id))
3565
3973
  )
3974
+ return pydantic.parse_obj_as(List[prefect.states.State], response.json())