diff --git a/.typos.toml b/.typos.toml index 409fbff2890..1523c45eec4 100644 --- a/.typos.toml +++ b/.typos.toml @@ -47,6 +47,7 @@ registr = "registr" tokenzation = "tokenzation" Conveyer = "Conveyer" froms = "froms" # used by sqlalchemy in a method name +queueing = "queueing" [default] locale = "en-us" diff --git a/docs/book/getting-started/zenml-pro/deploy-workspace-resource-pools.md b/docs/book/getting-started/zenml-pro/deploy-workspace-resource-pools.md index cb3d1e561d1..f4c80c6538c 100644 --- a/docs/book/getting-started/zenml-pro/deploy-workspace-resource-pools.md +++ b/docs/book/getting-started/zenml-pro/deploy-workspace-resource-pools.md @@ -126,7 +126,9 @@ helm upgrade zenml oci://public.ecr.aws/zenml/zenml \ ## Related resources -- [Resource pools](resource-pools.md) — product overview and concepts -- [Resource pools — Core concepts](resource-pools-core-concepts.md) +- [Resource pools](resource-pools.md) — product overview +- [Core concepts](resource-pools-core-concepts.md) +- [Admin guide](resource-pools-admin-guide.md) +- [User guide](resource-pools-user-guide.md) - [Self-hosted Deployment on Kubernetes with Helm](deploy-workspace-k8s.md) - [Helm chart on Artifact Hub](https://artifacthub.io/packages/helm/zenml/zenml) diff --git a/docs/book/getting-started/zenml-pro/resource-pools-admin-guide.md b/docs/book/getting-started/zenml-pro/resource-pools-admin-guide.md new file mode 100644 index 00000000000..29ea5e9b5c1 --- /dev/null +++ b/docs/book/getting-started/zenml-pro/resource-pools-admin-guide.md @@ -0,0 +1,513 @@ +--- +description: >- + How platform admins configure resource descriptors, pools, classes, target + settings, policies, and grants in ZenML Pro. +--- +# Admin guide + +This guide is for platform admins who operate ZenML Pro resource pools. Admins +define the resource vocabulary, map that vocabulary to real infrastructure, +and decide which users, teams, projects, pipelines, and workloads may consume +the shared capacity. + +Configure resource pools in the ZenML Pro UI: + +* **Organization Settings > Resource Pools** for organization-scoped pools. +* **Workspace Settings > Resource Pools** for workspace-scoped pools. + +The two locations use the same forms and concepts. The JSON snippets below use +the Resource Manager REST API shape exactly, so they are suitable as automation +references. + +## Before you start + +You need: + +* A workspace with resource pools enabled. See + [Enable Resource Pools](deploy-workspace-resource-pools.md). +* Stack components that will run pooled workloads, usually orchestrators and + step operators. +* Service connectors if ZenML needs credentials to access the infrastructure. +* A resource model that matches reality: Kubernetes node pools, cloud machine + types, GPU partitions, license pools, or shared storage classes. + +## Step 1: Build the resource vocabulary + +Open the Resource Pools page and use the resource descriptor section to review +or create descriptors. Stock descriptors are available for `CPU`, `memory`, +and `GPU`. Add custom descriptors for hardware-specific resources or platform +constraints. + +Use descriptor `kind` carefully. The built-in kinds `cpu`, `memory`, and `gpu` +are how Resource Manager matches the typed `ResourceSettings` fields that +authors configure on steps: + +* Use `kind: "gpu"` for every GPU-like descriptor, such as `GPU`, `h200`, or + `a10g`, so typed `ResourceSettings(gpu_count=...)` demands can match them. +* Use `kind: "cpu"` for every CPU-like descriptor that should match + `ResourceSettings(cpu_count=...)`. +* Use `kind: "memory"` for every memory-like descriptor that should match + `ResourceSettings(memory=...)`. +* Use a custom kind such as `license` when authors should request the resource + explicitly by name. + +Example API payload for a custom H200 descriptor: + +```json +{ + "name": "h200", + "kind": "gpu", + "description": "NVIDIA H200 GPU", + "attributes": { + "vendor": "nvidia", + "model": "H200", + "vram_gb": 141, + "architecture": "hopper" + }, + "units": [ + {"name": "GPU", "multiplier": 1} + ] +} +``` + +Example API payload for a custom license descriptor: + +```json +{ + "name": "training-license", + "kind": "license", + "description": "Floating training license seat", + "attributes": { + "vendor": "acme", + "scope": "organization" + }, + "units": [ + {"name": "seat", "multiplier": 1} + ] +} +``` + +## Step 2: Create a resource pool + +In the UI, click **Add resource pool**. Configure: + +| UI area | What it controls | +| --- | --- | +| Name and description | Human-readable identity | +| Rank | Which pool wins when multiple pools can satisfy the same request | +| Accounting mode | Whether Resource Manager enforces allocation (`authoritative`), gives hints (`advisory`), or only governs routing (`governance`) | +| Concurrency limit | Optional pool-wide cap on active requests | +| Targets | Which runtime target subjects the pool governs | +| Classes | Resource bundles available inside the pool | +| Target settings | Component or service connector settings applied from the pool | +| Attributes | Metadata used by pool selectors and operations | + +Use `authoritative` unless your infrastructure scheduler should remain the +source of truth. `governance` is suitable with schedulers such as Kueue or +Run:ai that handle allocation and preemption themselves. `advisory` is an +advanced middle ground. + +### Target bindings + +Targets describe how workloads enter the infrastructure. For pools, the main +targets are stack components and service connectors: + +* Stack components, usually orchestrators or step operators, launch workloads. +* Service connectors provide authentication to infrastructure resources. + +The Resource Manager API uses `target_bindings` with a `target_selector`. +Selectors must match the root-first subject chains attached to requests. + +```json +{ + "target_bindings": [ + { + "target_selector": { + "any": [ + { + "subject_type": "organization", + "subject_id": "", + "contains": { + "subject_type": "workspace", + "subject_id": "", + "contains": { + "subject_type": "component", + "subject_id": "", + "attributes": { + "component_type": "step_operator" + } + } + } + }, + { + "subject_type": "organization", + "subject_id": "", + "contains": { + "subject_type": "workspace", + "subject_id": "", + "contains": { + "subject_type": "service_connector", + "subject_id": "", + "attributes": { + "connector_type": "gcp" + } + } + } + } + ] + } + } + ] +} +``` + +See [Core concepts - Subject selectors](resource-pools-core-concepts.md#subject-selectors) +for the full scope convention. + +## Step 3: Add classes as resource bundles + +Classes are resource bundles. Model the shape your infrastructure actually +offers instead of creating separate classes for each resource. For example, +an H200 Kubernetes node pool should usually bundle H200 GPUs, CPU, and memory +in the same class. + +In the class form, configure: + +| UI area | What it controls | +| --- | --- | +| Class name | Pool-local label, such as `h200-reserved` | +| Rank | Preference among classes; higher ranks are tried first | +| Reclaimable | Whether capacity is safe from external reclaim: `never`, `coordinated`, or `unsafe` | +| Concurrency limit | Optional cap on active requests in the class | +| Resources | One or more descriptor quantities in the bundle | +| Target settings | Component or service connector settings applied when the class wins | +| Attributes | Metadata used by class selectors and operators | + +Use finite quantities for resources you want Resource Manager to limit at the +class level. Use unlimited quantity for resources you want to track without a +pool-class limit. Grants can still impose reservations and limits later. + +Example pool with two believable Kubernetes node-pool classes: + +```json +{ + "name": "prod-eu-gpu", + "description": "Production GPU capacity in the EU Kubernetes cluster", + "rank": 100, + "accounting_mode": "authoritative", + "concurrency_limit": 40, + "target_bindings": [ + { + "target_selector": { + "any": [ + { + "subject_type": "organization", + "subject_id": "", + "contains": { + "subject_type": "workspace", + "subject_id": "", + "contains": { + "subject_type": "component", + "subject_id": "", + "attributes": { + "component_type": "step_operator" + } + } + } + }, + { + "subject_type": "organization", + "subject_id": "", + "contains": { + "subject_type": "workspace", + "subject_id": "", + "contains": { + "subject_type": "service_connector", + "subject_id": "", + "attributes": { + "connector_type": "gcp" + } + } + } + } + ] + } + } + ], + "classes": [ + { + "class": "h200-reserved", + "rank": 100, + "reclaimable": "never", + "concurrency_limit": 8, + "resources": [ + {"resource": "h200", "quantity": 8, "unit": "GPU"}, + {"resource": "CPU", "quantity": 128, "unit": "CPU"}, + {"resource": "memory", "quantity": 1024, "unit": "GiB"} + ], + "target_settings": [ + { + "target_type": "component", + "settings": { + "pod_settings": { + "node_selector": { + "cloud.google.com/gke-nodepool": "h200-reserved" + }, + "tolerations": [ + { + "key": "dedicated", + "operator": "Equal", + "value": "h200", + "effect": "NoSchedule" + } + ] + } + } + } + ], + "attributes": { + "accelerator": "h200", + "purchase": "reserved" + } + }, + { + "class": "a10-burst", + "rank": 40, + "reclaimable": "coordinated", + "concurrency_limit": 20, + "resources": [ + {"resource": "GPU", "quantity": 16, "unit": "GPU"}, + {"resource": "CPU", "quantity": null}, + {"resource": "memory", "quantity": null} + ], + "target_settings": [ + { + "target_type": "component", + "settings": { + "pod_settings": { + "node_selector": { + "cloud.google.com/gke-nodepool": "a10-burst" + } + } + } + } + ], + "attributes": { + "accelerator": "a10", + "purchase": "shared-burst" + } + } + ], + "target_settings": [ + { + "target_type": "service_connector", + "settings": { + "connector_id": "", + "resource_type": "kubernetes-cluster", + "resource_id": "prod-eu" + } + } + ], + "attributes": { + "region": "eu-west-1", + "cluster": "prod-eu" + } +} +``` + +## Step 4: Attach policies + +Open a pool detail page and attach a policy. Policies answer "who may use this +pool, with what priority, and under which limits?" + +In the policy form, configure: + +| UI area | What it controls | +| --- | --- | +| Subjects | Who gets access. Common subjects are workspaces, projects, pipelines, accounts, service accounts, and teams | +| Match mode | Whether any selected subject is enough or all must match | +| Time window | Optional one-time, schedule, or cycle policy window | +| Preemption group | Optional related-work selector that prevents self-preemption | +| Priority or priority lane | Contention ordering | +| Concurrency limit | Optional policy-wide active request cap | +| Component target settings | Policy-level settings merged into the selected component | +| Grants | Optional class-level reservations, limits, defaults, and grant concurrency limits | + +### Grant-based policy + +Use grants when you want reservations, per-team limits, defaults for omitted +resources, or grant-level concurrency caps. + +```json +{ + "pool": "prod-eu-gpu", + "subject_selector": { + "any": [ + { + "subject_type": "organization", + "subject_id": "", + "contains": { + "subject_type": "team", + "subject_id": "" + } + } + ] + }, + "priority": 100, + "priority_lane": false, + "concurrency_limit": 12, + "target_settings": [ + { + "target_type": "component", + "settings": { + "pod_settings": { + "labels": { + "zenml.io/resource-policy": "ml-platform-prod" + } + } + } + } + ], + "grants": [ + { + "class": "h200-reserved", + "concurrency_limit": 6, + "resources": [ + { + "resource": "h200", + "reserved": 4, + "limit": 8, + "unit": "GPU", + "missing_action": "reject" + }, + { + "resource": "CPU", + "reserved": 32, + "limit": 128, + "unit": "CPU", + "missing_action": "default", + "default_value": 4 + }, + { + "resource": "memory", + "reserved": 256, + "limit": 1024, + "unit": "GiB", + "missing_action": "default", + "default_value": 16 + } + ] + } + ] +} +``` + +Grant resource behavior: + +* If a request demands a resource that is in the class but missing from the + grant, the grant cannot admit it. +* `reserved` is the protected share for non-reclaimable work. +* `limit: null` means the limit follows the class resource quantity. +* `missing_action: "omit"` does not allocate an omitted resource. +* `missing_action: "default"` allocates `default_value` when the request omits + the resource. +* `missing_action: "reject"` requires the request to mention the resource. + +### Grantless policy + +Use `grants: []` when the subject should access every matching class resource +in the pool without per-grant reservations or limits. The effective reservation +is 0 and the effective limit is the class resource quantity. + +```json +{ + "pool": "prod-eu-gpu", + "subject_selector": { + "any": [ + { + "subject_type": "organization", + "subject_id": "", + "contains": { + "subject_type": "account", + "subject_id": "", + "attributes": { + "is_service_account": true + } + } + } + ] + }, + "priority_lane": true, + "concurrency_limit": 20, + "grants": [] +} +``` + +Priority lane gives the policy the maximum Resource Manager priority. It can +reclaim lower-priority work that opted into reclaim, but it does not bypass +pool target bindings, class matching, or capacity checks in authoritative mode. + +## Step 5: Use selectors intentionally + +Pool target selectors and policy subject selectors match the subjects attached +to resource requests. The same subject can appear at different scopes, so write +selectors root-first. + +Recommended main story: + +| Selector location | Common subject choices | +| --- | --- | +| Pool targets | Stack components and service connectors | +| Policies | Workspaces, projects, pipelines, user accounts, service accounts, and teams | + +Avoid overly broad policies on scarce resources unless the pool is meant to be +shared by everyone. Prefer a broad target binding, such as one step operator, +and narrower policy subjects, such as one team or project. + +## Step 6: Operate pools + +Use the Resource Pools page and pool detail pages to inspect: + +* Pool classes and their bundled resources. +* Policies attached to a pool. +* Grant reservations and limits. +* Active allocations and queued requests. +* Which pool, class, policy, and grant admitted a request. + +The CLI can inspect resource requests: + +```shell +zenml resource-request list --status pending +zenml resource-request list --pool-id +zenml resource-request describe +``` + +Pool, descriptor, class, policy, grant, and allocation configuration is done +through the UI or the Resource Manager API, not the ZenML CLI. + +## Checklist + +Before handing a pool to users: + +* Descriptors use stable names and kinds. +* Classes bundle resources the way infrastructure actually provides them. +* CPU and memory are included when they matter for scheduling or traceability. +* Unlimited class resources are intentional and documented for operators. +* Target bindings match the stack components or service connectors that will + appear on runtime requests. +* Policies use the right subjects for access control. +* Grant-based policies include every class resource that requests may demand, + or use `missing_action: "default"` where a default allocation is desired. +* Priority-lane policies are reserved for critical or external workloads that + genuinely need maximum priority. +* Accounting mode is `authoritative` unless another scheduler owns allocation. + +## See also + +* [Core concepts](resource-pools-core-concepts.md) - definitions and selector + conventions. +* [User guide](resource-pools-user-guide.md) - what pipeline authors configure. +* [External workloads](resource-pools-external-workloads.md) - direct Resource + Manager requests and priority-lane patterns. +* [Examples](resource-pools-examples.md) - end-to-end scenarios. +* [Reconciliation process](resource-pools-reconciliation.md) - runtime behavior. diff --git a/docs/book/getting-started/zenml-pro/resource-pools-core-concepts.md b/docs/book/getting-started/zenml-pro/resource-pools-core-concepts.md index 929188519dd..6af9c02bbcb 100644 --- a/docs/book/getting-started/zenml-pro/resource-pools-core-concepts.md +++ b/docs/book/getting-started/zenml-pro/resource-pools-core-concepts.md @@ -1,227 +1,492 @@ --- description: >- - Precise definitions for ZenML Pro resource pools, subject policies, and - resource requests. + Definitions for ZenML Pro resource descriptors, pools, classes, policies, + requests, selectors, and leases. --- # Core concepts -This page defines **pools**, **policies**, and **requests** for ZenML Pro. +This page defines the building blocks of ZenML Pro resource pools. For +step-by-step setup, see the [Admin guide](resource-pools-admin-guide.md) and +[User guide](resource-pools-user-guide.md). -## Pools +{% hint style="info" %} +Resource pools apply only to [dynamic pipelines](https://docs.zenml.io/how-to/steps-pipelines/dynamic_pipelines). +Static pipelines do not create resource requests or wait for pool allocation. +{% endhint %} -A **resource pool** is a named shared bucket. For each resource key (for -example `gpu`), you set how many units exist in the pool. Policies on that -pool further split that capacity among orchestrators and step operators. +## Resource descriptors -Steps and pools use the same keys and integer amounts. Typical keys are: +A resource descriptor is an entry in the resource vocabulary. It describes what +can be requested or accounted for. It does not say how much capacity exists or +who may use it. -| Key | Meaning | +| Field | Meaning | | --- | --- | -| `gpu` | GPU count (requested by steps through `ResourceSettings.gpu_count`) | -| `mcpu` | Milli-CPU (requested by steps through `ResourceSettings.cpu_count * 1000`, rounded up) | -| `memory_mb` | Memory in megabytes (requested by steps through `ResourceSettings.memory`) | -| `step_run` | One concurrent step run (added automatically by the server for each step) | +| `name` | Stable descriptor identifier used by pool classes, grants, and exact resource demands | +| `kind` | Type label used by kind-based demands. The built-in kinds `cpu`, `memory`, and `gpu` are the convention used by `ResourceSettings`; custom kinds such as `license` are for explicit custom demands. | +| `description` | Optional human-readable explanation | +| `attributes` | Optional key/value metadata used for filtering and documentation | +| `units` | Optional unit catalog, each with `name` and integer `multiplier` | + +Stock descriptors are: + +| Name | Kind | Units | +| --- | --- | --- | +| `CPU` | `cpu` | `mCPU`, `CPU` | +| `memory` | `memory` | `B`, `KB`, `KiB`, `MB`, `MiB`, `GB`, `GiB`, `TB`, `TiB` | +| `GPU` | `gpu` | `GPU` | + +### Built-in kinds and `ResourceSettings` + +The descriptor `kind` is how the built-in `ResourceSettings` fields connect to +Resource Manager descriptors. ZenML does not require authors to know descriptor +names for basic compute resources. Instead, it emits demands by built-in kind: + +| ResourceSettings field | Built-in demand kind | Typical descriptor names | +| --- | --- | --- | +| `cpu_count` | `cpu` | `CPU`, custom CPU flavors | +| `memory` | `memory` | `memory`, custom memory pools | +| `gpu_count` | `gpu` | `GPU`, `h200`, `a10g`, custom accelerator descriptors | + +Use these built-in kinds for every flavor of CPU, memory, or GPU that should +match the corresponding typed `ResourceSettings` field. For example, an H200 +descriptor named `h200` should use `kind: "gpu"`; a larger CPU pool descriptor +should use `kind: "cpu"`; and a special memory-backed resource that should +satisfy `ResourceSettings(memory=...)` should use `kind: "memory"`. + +Use a custom kind, such as `license`, only when the resource should not be +matched by the basic CPU, memory, or GPU fields. Authors request those +resources explicitly through `ResourceSettings.resources`. + +Admins can add descriptors such as `h200`, `a10g`, `training-license`, or +`shared-storage`. The H200 example below uses `kind: "gpu"` so +`ResourceSettings(gpu_count=...)` can match it. + +```json +{ + "name": "h200", + "kind": "gpu", + "description": "NVIDIA H200 accelerator", + "attributes": { + "vendor": "nvidia", + "model": "H200", + "vram_gb": 141 + }, + "units": [ + {"name": "GPU", "multiplier": 1} + ] +} +``` -Custom keys (for example `tpu`) can be set with `pool_resources` on the step. +## Pools -### CLI: pools +A resource pool is a named scheduling and accounting boundary. It has: -```bash -# Create a pool (capacity as JSON or YAML) -zenml resource-pool create training-gpus \ - --capacity '{"gpu": 8, "step_run": 32}' \ - --description "Shared training GPUs for the workspace" +| Field | Meaning | +| --- | --- | +| `name` | Pool identifier | +| `description` | Optional explanation | +| `rank` | Preference when more than one pool can admit the same request; higher ranks win first | +| `accounting_mode` | How strongly Resource Manager enforces allocation decisions | +| `concurrency_limit` | Optional maximum number of active requests in the pool | +| `target_bindings` | Subject selectors that describe which runtime targets this pool governs | +| `classes` | Pool-local resource bundles | +| `target_settings` | UI-supported settings merged into selected targets | +| `attributes` | Optional metadata for filtering and pool selectors | -# List pools with occupied vs total capacity -zenml resource-pool list +The UI exposes pools under **Organization Settings > Resource Pools** and +**Workspace Settings > Resource Pools**. -# Inspect one pool (name, ID prefix, or full ID) -zenml resource-pool describe training-gpus +### Accounting modes -# Shrink or grow capacity (0 removes a key from the pool) -zenml resource-pool update training-gpus --capacity '{"gpu": 4}' +Most examples use `authoritative`. -# Remove a pool (use -y to skip confirmation) -zenml resource-pool delete training-gpus --yes +| Mode | Meaning | +| --- | --- | +| `authoritative` | Resource Manager controls allocation, queueing, limits, reservations, concurrency caps, and preemption. Use this when the infrastructure scheduler does not already provide the required fairness and conflict handling. | +| `advisory` | Resource Manager computes accounting pressure and ordering hints, but infrastructure may still admit work. Treat this as advanced and experimental. | +| `governance` | Policies decide governance and target routing, but allocation and preemption are left to infrastructure-level schedulers such as Kueue, Run:ai, or similar systems. | + +## Pool classes + +A pool class is a pool-local resource bundle. This is the most important model +change: classes are not one-resource buckets. They usually represent the shape +of real infrastructure, such as a Kubernetes node pool or cloud machine type. + +Example class: + +```json +{ + "class": "h200-reserved", + "rank": 100, + "reclaimable": "never", + "concurrency_limit": 8, + "resources": [ + {"resource": "h200", "quantity": 8, "unit": "GPU"}, + {"resource": "CPU", "quantity": 128, "unit": "CPU"}, + {"resource": "memory", "quantity": 1024, "unit": "GiB"} + ], + "target_settings": [ + { + "target_type": "component", + "settings": { + "pod_settings": { + "node_selector": { + "cloud.google.com/gke-accelerator": "nvidia-h200" + }, + "tolerations": [ + { + "key": "dedicated", + "operator": "Equal", + "value": "h200", + "effect": "NoSchedule" + } + ] + } + } + } + ], + "attributes": { + "accelerator": "h200", + "region": "eu-west-1" + } +} ``` -## Policies - -A **policy** connects one stack component—the orchestrator or step operator that -acts as the *resource requester* for a step—to one pool. Think of three knobs -per resource key: - -* **Reserved** — How much of the pool you *label* as belonging to this component -for accounting. Usage up to that amount counts as *in share*; anything above it -(while the pool still has free units) is *borrowed* idle capacity. Reserved is -not a separate pile of hardware: it is the share used to decide who is “in -their rights” versus who is on spare capacity. Across all policies on the same -pool, reserved totals per key cannot exceed the pool capacity. Reserved must -also be ≤ that policy’s limit for the same key. - -* **Limit** — The hard ceiling on how much this component may hold from the pool -at once for that key. Grants never go above the limit, even if the pool is -empty. For **preemptible** workloads, the space between reserved and limit is -where borrowing can happen (subject to pool free capacity). Non-preemptible -work does not use that band: each requested amount per key must be **≤ -reserved**, and a higher **limit** does not raise that ceiling (limit still caps -preemptible burst and total use). - -* **Priority** — A number; higher means that component’s requests are preferred in -the queue. When the reconciler must **preempt** someone, it looks at **lower** -priority preemptible runs first as victims (see below). - -{% hint style="warning" %} -A single orchestrator or step operator may have **several policies** attached, -each pointing at a **different pool**. The server still treats each step as one -**resource request**. Eligibility and allocation are evaluated **per pool** -against the **full** set of requested keys: the step may be queued on more than -one pool, but at most **one** pool ends up owning the active allocation. ZenML -does **not** split a request across pools (for example GPUs from one pool and -`mcpu` from another). Every key in the request must be satisfiable from the -**same** pool and policy that wins. See -[Examples — Multiple pools and multi-key requests](resource-pools-examples.md#multiple-pools-and-multi-key-requests). -{% endhint %} +Class fields: -### CLI: policies +| Field | Meaning | +| --- | --- | +| `class` | Pool-local class label | +| `resources` | One or more descriptor quantities bundled together | +| `rank` | Preference among matching classes; higher ranks are tried first | +| `reclaimable` | External or Resource Manager reclaim behavior: `never`, `coordinated`, or `unsafe` | +| `concurrency_limit` | Optional maximum number of active requests in this class | +| `target_settings` | UI-supported target settings applied when this class is selected | +| `attributes` | Optional metadata for class selectors | -```bash -# Attach an orchestrator to a pool (default component type is orchestrator) -zenml resource-pool attach-policy training-gpus my-k8s-orch \ - --priority 10 \ - --reserved '{"gpu": 2}' \ - --limit '{"gpu": 4}' +A class resource with `"quantity": null` is unlimited at the pool-class level. +It is still selected and accounted for, so it remains visible for traceability. +Policy grants can still add reservations and limits for that resource. -# Same for a step operator stack component -zenml resource-pool attach-policy training-gpus my-remote-operator \ - --component-type step_operator \ - --priority 5 \ - --reserved '{"gpu": 2}' \ - --limit '{"gpu": 4}' +## Target bindings and target settings -# List every policy on a pool -zenml resource-pool list-policies training-gpus +Target bindings decide which runtime targets a pool governs. In the UI, these +are usually stack components, especially orchestrators and step operators. A +service connector can also appear as a target subject when the pool is tied to +authentication for the infrastructure. -# List all pools a given orchestrator is attached to -zenml resource-pool list-policies --component my-k8s-orch +Target settings are merged into the target chosen for a winning allocation. The +UI currently exposes: -# Remove that component’s policy from the pool -zenml resource-pool detach-policy training-gpus my-k8s-orch +| Location | UI-supported target settings | +| --- | --- | +| Pool | Component settings, service connector settings | +| Class | Component settings, service connector settings | +| Policy | Component settings | +| Grant | Component settings | + +Settings are merged in order from broader to narrower scope: pool route, class, +policy, then grant. More specific settings can override the same keys from +broader settings. + +## Subject selectors + +Resource Manager subject selectors are hierarchical filters. They match the +root-first subject chains attached to resource requests. When you write JSON +directly against the Resource Manager API, use the same scope convention the +ZenML client uses. + +Selector node: + +```json +{ + "subject_type": "organization", + "subject_id": "", + "attributes": {}, + "metadata": {}, + "contains": { + "subject_type": "workspace", + "subject_id": "" + } +} ``` -## Resource requests - -For eligible runs, the server builds a **resource request** from the step’s -`ResourceSettings`, records whether the step is preemptible, and tracks status: -queued, allocated, rejected, preempted, or cancelled. - -{% hint style="info" %} -Only dynamic pipelines participate in resource queuing and allocation -waiting: the server creates resource requests and the client blocks until -allocation when the snapshot is dynamic. Static pipelines do not use this -path today. -{% endhint %} - -What users set in `ResourceSettings` becomes a server-side resource request (for -example `gpu_count` → `gpu`, `cpu_count` → `mcpu`, `memory` → `memory_mb`, plus -an implicit concurrent `step_run` slot). The pool must define capacity for each -key requested by the step, except for three built-in types: if the pool has no -row for `mcpu`, `memory_mb`, or the implicit `step_run` key, ZenML Pro treats that -dimension as **effectively unbounded** at the pool layer, so missing rows there -do not by themselves cause rejection. For every other key (including -`gpu` and custom keys from `pool_resources`), a missing pool row means zero -capacity: a positive request is rejected and the step run fails to start. - -{% hint style="warning" %} -If the pool does define a key but the subject policy omits that key, -limits fall back to the pool total and reserved defaults to zero for bounded -keys. Non-preemptible steps must stay within their reserved capacity, so a -positive request for a key not defined in the policy is rejected. -{% endhint %} - -### Step decorators: `ResourceSettings` +Boolean selector expression: -Declare demand on the step; the server turns it into the resource request when -the pipeline is **dynamic** and pooling applies to the stack. +```json +{ + "any": [ + {"subject_type": "team", "subject_id": ""}, + {"subject_type": "account", "subject_id": ""} + ] +} +``` -**Typical GPU / CPU / memory (preemptible by default):** +Use `any` when one selected subject should match. Use `all` when a request must +carry all selected subject chains. Use `not` for exclusion. Policy subject +selectors may also include time selectors through the UI. -```python -from zenml import step -from zenml.config import ResourceSettings +### Scope convention +Common request subjects are shaped like this: -@step( - settings={ - "resources": ResourceSettings( - gpu_count=2, - cpu_count=4, - memory="16GiB", - ) +| Subject kind | Request subject chain | +| --- | --- | +| Organization | `organization` | +| Workspace | `organization -> workspace` | +| Project | `organization -> workspace -> project` | +| Pipeline | `organization -> workspace -> project -> pipeline` | +| Pipeline run | `organization -> workspace -> project -> pipeline_run` | +| Step run | `organization -> workspace -> project -> pipeline_run -> step_run` | +| Stack component | `organization -> workspace -> component` | +| Service connector | `organization -> workspace -> service_connector` | +| User account | `account` | +| Service account | `organization -> account` | +| Team | `organization -> team` | + +The `step_run` entry above is a subject type used to identify requests. It is +not a resource descriptor or resource demand. + +Example selector for one step operator target: + +```json +{ + "subject_type": "organization", + "subject_id": "", + "contains": { + "subject_type": "workspace", + "subject_id": "", + "contains": { + "subject_type": "component", + "subject_id": "", + "attributes": { + "component_type": "step_operator" + } } -) -def train() -> None: - ... + } +} ``` -**Non-preemptible (must stay within policy reserved per key):** +Example policy selector for any project in a workspace: + +```json +{ + "subject_type": "organization", + "subject_id": "", + "contains": { + "subject_type": "workspace", + "subject_id": "", + "contains": { + "subject_type": "project" + } + } +} +``` -```python -@step( - settings={ - "resources": ResourceSettings( - gpu_count=1, - preemptible=False, - ) +Example policy selector for a team or a service account: + +```json +{ + "any": [ + { + "subject_type": "organization", + "subject_id": "", + "contains": { + "subject_type": "team", + "subject_id": "" + } + }, + { + "subject_type": "organization", + "subject_id": "", + "contains": { + "subject_type": "account", + "subject_id": "", + "attributes": { + "is_service_account": true + } + } } -) -def production_train() -> None: - ... + ] +} ``` -**Custom pool keys** (must exist on the pool and policy when non-preemptible): +## Policies and grants + +A policy attaches access rules to one pool. The policy subject selector answers +"who may use this pool?" Typical policy subjects are workspaces, projects, +pipelines, accounts, teams, and service accounts. -```python -@step( - settings={ - "resources": ResourceSettings( - gpu_count=1, - pool_resources={"tensorrt_sessions": 1}, - ) +Policy fields: + +| Field | Meaning | +| --- | --- | +| `pool` | Pool ID or exact pool name | +| `subject_selector` | Subject selector or boolean selector expression | +| `preemption_group` | Optional selector used to avoid preempting related work | +| `priority` | Higher values win contention; omit when `priority_lane` is true | +| `priority_lane` | Maximum internal priority for critical policies | +| `concurrency_limit` | Optional maximum number of active requests admitted by this policy | +| `grants` | Optional per-class grant list | +| `target_settings` | Policy-level component settings | + +A priority-lane policy gets the maximum internal priority. It can preempt +lower-priority reclaimable work, but it still needs a pool, matching class +resources, and capacity according to the pool accounting mode. + +### Grant-based policies + +A grant references exactly one class in the pool. Inside that class, each grant +resource can reserve capacity, limit active usage, and decide what happens when +a request omits that resource. + +```json +{ + "class": "h200-reserved", + "concurrency_limit": 4, + "resources": [ + { + "resource": "h200", + "reserved": 4, + "limit": 8, + "unit": "GPU", + "missing_action": "reject" + }, + { + "resource": "CPU", + "reserved": 32, + "limit": 128, + "unit": "CPU", + "missing_action": "default", + "default_value": 4 + }, + { + "resource": "memory", + "reserved": 256, + "limit": 1024, + "unit": "GiB", + "missing_action": "default", + "default_value": 16 } -) -def infer() -> None: - ... + ] +} ``` -Typed fields override the same keys if both appear in `pool_resources`. See -[step configuration](https://docs.zenml.io/how-to/steps-pipelines/configuration) -in the OSS docs for the full `ResourceSettings` model. +Grant resource fields: + +| Field | Meaning | +| --- | --- | +| `resource` | Descriptor ID or name in the class bundle | +| `reserved` | Guaranteed share for non-reclaimable work | +| `limit` | Maximum active usage for this policy grant; `null` means follow the class resource quantity | +| `unit` | Unit for `reserved`, `limit`, and `default_value` | +| `missing_action` | `omit`, `default`, or `reject` when a request does not include this resource | +| `default_value` | Quantity to allocate when `missing_action` is `default` | + +If a request asks for a resource that is present in the class but missing from +the grant, that grant cannot admit the request. If the grant includes a +resource that the request omits, `missing_action` controls whether the grant +ignores it, rejects the request, or allocates a default. + +### Grantless policies + +A policy with `"grants": []` admits matching subjects to every matching class +in the pool. Each class resource has reservation 0 and an effective limit equal +to the class resource quantity. Unlimited class resources remain unlimited at +the pool-class level. + +Grantless policies are useful for trusted critical workloads, external +workloads that need the full node shape, or simple governance where you do not +need per-team reservations. They still respect pool/class concurrency limits, +pool target bindings, class matching, accounting mode, and reclaimability. -### CLI: resource requests +## Resource requests -Requests are created when dynamic steps run; you **inspect** or **clean them up** -from the CLI (IDs come from list output or the dashboard). +When a dynamic step is ready to launch, ZenML creates a resource request. The +request contains subjects and demands. External clients may create the same API +shape directly. -```bash -# All resource requests visible to your user (see --help for filters) -zenml resource-request list +Demand fields: -# Full detail for one request (status, resources, step run link) -zenml resource-request describe 01234567-89ab-cdef-0123-456789abcdef +| Field | Meaning | +| --- | --- | +| `resource` | Exact descriptor ID or name | +| `kind` | Match any descriptor with this kind | +| `quantity` | Positive quantity | +| `unit` | Optional unit | +| `class` | Optional exact pool class | +| `resource_selector` | Optional selector over descriptor attributes | +| `class_selector` | Optional selector over class attributes | -# Drop a stuck or abandoned request (use with care) -zenml resource-request delete 01234567-89ab-cdef-0123-456789abcdef +`ResourceSettings` produces these common demands: -# Requests tied to a specific pool: queued only, active only, or both -zenml resource-pool requests training-gpus --view queued -zenml resource-pool requests training-gpus --view active -zenml resource-pool requests training-gpus --view all +| Step setting | Demand | +| --- | --- | +| `gpu_count` | `kind: "gpu"`, quantity from `gpu_count`, optional `class` from `gpu_class` | +| `cpu_count` | `kind: "cpu"`, unit `CPU` | +| `memory` | `kind: "memory"`, unit from the memory suffix such as `GiB` | +| `resources` | Custom descriptor demands by name, kind, or selector | + +Example runtime request: + +```json +{ + "subjects": [ + { + "subject_type": "organization", + "subject_id": "", + "child": { + "subject_type": "workspace", + "subject_id": "", + "child": { + "subject_type": "component", + "subject_id": "", + "attributes": { + "component_type": "step_operator" + } + } + } + } + ], + "demands": [ + {"kind": "gpu", "quantity": 2, "class": "h200-reserved"}, + {"kind": "cpu", "quantity": 8, "unit": "CPU"}, + {"kind": "memory", "quantity": 64, "unit": "GiB"} + ], + "reclaim_tolerance": "none", + "allocation_wait_timeout_seconds": 1800, + "metadata": { + "workload": "nightly-training" + } +} ``` -## See also +## Reclaim tolerance + +Each request carries `reclaim_tolerance`: -* [Resource pools](resource-pools.md) — overview and UX story -* [Examples](resource-pools-examples.md) — scenarios and outcomes -* [How preemption works](resource-pools-reconciliation.md#how-preemption-works) — behavior and ordering -* [Resource pool reconciliation](resource-pools-reconciliation.md) — runtime flow +| Value | Meaning | Compatible class `reclaimable` | +| --- | --- | --- | +| `none` | Do not interrupt this request for pool reasons | `never`; in authoritative mode it must fit reserved grant share for grant-based policies, or matching class capacity for grantless policies | +| `coordinated` | The workload can be stopped and retried cleanly | `never`, `coordinated` | +| `any` | Best-effort work; may use unsafe classes | `never`, `coordinated`, `unsafe` | + +Legacy `preemptible=False` maps to `none`; `preemptible=True` maps to +`coordinated`. Prefer `reclaim_tolerance` for new code. + +## Request statuses + +| Status | Meaning | +| --- | --- | +| `pending` | Admitted and queued, waiting for allocation | +| `allocated` | Capacity granted and leased | +| `preempting` | Reconciler asked the owner to stop so capacity can be reclaimed | +| `preempted` | Request was stopped by preemption | +| `rejected` | No policy/class/grant path can satisfy the request | +| `cancelled` | Request was cancelled before allocation | +| `released` | Work finished and capacity was returned | +| `expired` | Lease expired while allocated | + +Inspect resource requests in the UI or with `zenml resource-request list` and +`zenml resource-request describe `. diff --git a/docs/book/getting-started/zenml-pro/resource-pools-examples.md b/docs/book/getting-started/zenml-pro/resource-pools-examples.md index dcf11b19e71..0774cf47fb1 100644 --- a/docs/book/getting-started/zenml-pro/resource-pools-examples.md +++ b/docs/book/getting-started/zenml-pro/resource-pools-examples.md @@ -1,1235 +1,2957 @@ --- description: >- - Step-by-step ZenML Pro resource pool examples: pool JSON, policy JSON, - ResourceSettings, and outcomes for new users. + End-to-end resource pool use cases for CPU queues, GPU node pools, cloud + machine types, licenses, external workloads, governance schedulers, and + policy patterns. --- -# Resource pool examples (workbook) +# Examples + +The examples below show practical resource pool configurations for common +ZenML Pro workload management scenarios. Each example explains the +infrastructure being modeled, the access policies that apply to it, and the +Resource Manager API JSON that corresponds to the UI configuration. + +Configure these objects in the ZenML Pro UI under **Organization Settings > +Resource Pools** or **Workspace Settings > Resource Pools**. The JSON snippets +use the Resource Manager REST API shape for teams that automate the same setup. + +## How to read these examples + +Resource pool examples have two layers: + +1. **Inventory modeling**: descriptors and pool classes describe what + infrastructure exists. A class should usually represent a real bundle such + as a CPU worker group, a Kubernetes node pool, a SageMaker machine type, or a + licensed training slot. +2. **Access modeling**: policies and grants describe who may consume that + inventory, at which priority, with which reservations, limits, defaults, and + concurrency caps. + +That separation matters. If two teams use the same H200 Kubernetes node pool, +do not duplicate the pool just to express two different access rules. Model the +node pool once, then attach separate policies for production, research, +external services, or emergency operators. + +The examples progress from small to more advanced: + +| Example | Focus | +| --- | --- | +| [Simple CPU and memory queue](#example-1-simple-cpu-and-memory-queue) | First pool, first class, basic concurrency, basic team policy | +| [One GPU machine shape](#example-2-one-gpu-machine-shape) | Exact hardware descriptor, class bundle, production reservation, experiment burst | +| [Kubernetes cluster with two node pools](#example-3-kubernetes-cluster-with-two-node-pools) | Multiple classes, target settings, reserved and burst node pools | +| [Kubernetes cluster with per-node classes](#example-4-kubernetes-cluster-with-per-node-classes) | Modeling each node as a class for node-level placement | +| [Multiple Kubernetes clusters as classes](#example-5-multiple-kubernetes-clusters-as-classes) | Class-level service connector settings and cluster burst | +| [Multi-region GPU dispatch](#example-6-multi-region-gpu-dispatch) | Multiple pools for different regions and step operators | +| [Shared organization pool across workspaces](#example-7-shared-organization-pool-across-workspaces) | Organization-scope policies, workspace/project/team access | +| [Cloud machine types](#example-8-cloud-machine-types) | Modeling machine families such as SageMaker or Vertex AI worker types | +| [Training license with traceability](#example-9-training-license-with-traceability) | Custom descriptor, finite license seats, unlimited CPU/memory tracking | +| [External inference priority lane](#example-10-external-inference-priority-lane) | Service account subject, grantless priority-lane policy, direct request | +| [Governance with Kueue or Run:ai](#example-11-governance-with-kueue-or-runai) | Letting infrastructure own allocation while ZenML owns policy routing | +| [Policy cookbook](#example-12-policy-cookbook) | Common access patterns on top of existing inventory | + +## Example 1: Simple CPU and Memory Queue + +Start with a CPU-only workspace. One workspace uses a remote execution backend, +and pipeline authors can launch more CPU-heavy dynamic steps than the workers +should run at the same time. The resource pool defines a queue, a maximum +number of active jobs, and a baseline access policy for the workspace. + +The inventory model is intentionally small: + +* Use the stock `CPU` and `memory` descriptors. +* Create one pool for the remote worker backend. +* Create one class named `standard-workers`. +* Bundle CPU and memory in that class because they are consumed together by the + same worker group. +* Add a pool concurrency limit so excess jobs queue instead of overloading the + worker backend. + +CPU and memory examples are useful even when those resources are less scarce +than GPUs. Resource pools can provide admission control and traceability for +general compute workloads as well as accelerator workloads. + +### Inventory -Read this page like a short course: each section is one self-contained scenario. -You will always see three things—the pool (shared capacity), the policy (how one -orchestrator or step operator may use that pool), and the step -(`ResourceSettings`). Then we spell out what the server does. +```json +{ + "name": "workspace-standard-compute", + "description": "Standard CPU and memory workers for dynamic steps", + "rank": 10, + "accounting_mode": "authoritative", + "concurrency_limit": 12, + "target_bindings": [ + { + "target_selector": { + "subject_type": "organization", + "subject_id": "", + "contains": { + "subject_type": "workspace", + "subject_id": "", + "contains": { + "subject_type": "component", + "subject_id": "", + "attributes": { + "component_type": "orchestrator" + } + } + } + } + } + ], + "classes": [ + { + "class": "standard-workers", + "rank": 100, + "reclaimable": "never", + "concurrency_limit": 12, + "resources": [ + {"resource": "CPU", "quantity": 96, "unit": "CPU"}, + {"resource": "memory", "quantity": 384, "unit": "GiB"} + ], + "attributes": { + "worker_group": "standard", + "environment": "dev" + } + } + ] +} +``` -**Assumptions unless stated otherwise:** +The class does not describe one worker. It describes the total schedulable +envelope of the worker group. Resource Manager accounts against the class total +while the orchestrator still launches the actual tasks. -* Steps are preemptible by default if you omit `preemptible=False`. -* One step run at a time when we say “no other work is running,” so you can - focus on a single decision. -* Every key in a policy’s `reserved` and `limit` must exist on the pool’s - capacity. You cannot meter a resource in policy that the pool does not - define. -* If one orchestrator or step operator has policies to several pools, the step - still receives at most one allocation from one pool. The whole request must be - eligible on that pool; resources are not split across pools for a single step. +### Access pattern: Workspace baseline -For definitions of reserved, limit, and priority, see -[Core concepts](resource-pools-core-concepts.md). For preemption ordering, see -[How preemption works](resource-pools-reconciliation.md#how-preemption-works). +```json +{ + "pool": "workspace-standard-compute", + "subject_selector": { + "subject_type": "organization", + "subject_id": "", + "contains": { + "subject_type": "workspace", + "subject_id": "" + } + }, + "priority": 10, + "priority_lane": false, + "concurrency_limit": 10, + "grants": [ + { + "class": "standard-workers", + "resources": [ + { + "resource": "CPU", + "reserved": 16, + "limit": 80, + "unit": "CPU", + "missing_action": "reject" + }, + { + "resource": "memory", + "reserved": 64, + "limit": 320, + "unit": "GiB", + "missing_action": "reject" + } + ] + } + ] +} +``` -## Primer: from `ResourceSettings` to the resource request +The policy gives the workspace a small protected CPU/memory share and a larger +ceiling. If the workspace launches a burst of dynamic steps, up to ten requests +can be active under the policy. The rest queue even if raw CPU accounting would +still fit, because the policy concurrency limit is a separate active-request +guardrail. -Say a step declares: +### Runtime behavior: baseline queue ```python -from zenml.config import ResourceSettings - ResourceSettings( - gpu_count=2, cpu_count=4, memory="16GiB", - pool_resources={"tensorrt_sessions": 1}, - preemptible=True, + reclaim_tolerance="none", ) ``` -ZenML turns that into one resource request. Roughly: - -| Source | Request key | How the amount is derived | Example value | -| --- | --- | --- | --- | -| `gpu_count` | `gpu` | Same as `gpu_count` | 2 | -| `cpu_count` | `mcpu` | `ceil(cpu_count * 1000)` | 4000 | -| `memory` | `memory_mb` | Converted to megabytes | 17180 (for `"16GiB"`) | -| `pool_resources` | (your names) | Copied as-is, merged with typed fields | 1 | -| Server | `step_run` | Always 1 per step | 1 | +The typed settings produce `cpu` and `memory` demands. Those match the stock +`CPU` and `memory` descriptors by kind. Because the request is non-reclaimable, +it must fit inside the grant's reserved share in authoritative mode. -The pool must define capacity for bounded keys such as `gpu` and -`tensorrt_sessions`. If the pool has no row for `mcpu`, `memory_mb`, or -`step_run`, that dimension is unbounded at the pool layer (see the examples -below). For everything else, a missing pool row means zero capacity. If you want -a policy to set `reserved` / `limit` on a key, that key must appear on the pool -first—policy keys are always a subset of pool keys. - ---- +If the workspace has four active steps that each request 4 CPUs and 16 GiB +memory, the next identical step still fits the reserved share. If twenty such +steps arrive at once, Resource Manager admits up to the policy concurrency +limit and queues the rest. The queue is useful here even when raw CPU and memory +would technically fit, because the concurrency limit protects the worker +backend from too many active processes. -## Warm-up: one pool, one policy, only GPUs +### Access pattern: Interactive work outranks nightly batch -### Preemptible step borrows past reserved - -**Story:** The team has four GPUs “labeled” for them, but the pool is empty. -They ask for six GPUs and allow preemption. They may borrow two idle GPUs. - -**Pool** +The same CPU inventory can support two different behaviors: interactive work +gets a small but responsive lane, while nightly batch gets a wide but low +priority lane. ```json { - "name": "datacenter-gpus", - "capacity": { - "gpu": 8 - } + "pool": "workspace-standard-compute", + "subject_selector": { + "subject_type": "organization", + "subject_id": "", + "contains": { + "subject_type": "workspace", + "subject_id": "", + "contains": { + "subject_type": "project", + "subject_id": "" + } + } + }, + "priority": 50, + "priority_lane": false, + "concurrency_limit": 4, + "grants": [ + { + "class": "standard-workers", + "resources": [ + {"resource": "CPU", "reserved": 16, "limit": 32, "unit": "CPU", "missing_action": "reject"}, + {"resource": "memory", "reserved": 64, "limit": 128, "unit": "GiB", "missing_action": "reject"} + ] + } + ] } ``` -**Policy** (orchestrator `team-ml-orch` attached to this pool) - ```json { - "pool": "datacenter-gpus", - "component": "team-ml-orch", - "component_type": "orchestrator", - "priority": 10, - "reserved": { - "gpu": 4 + "pool": "workspace-standard-compute", + "subject_selector": { + "subject_type": "organization", + "subject_id": "", + "contains": { + "subject_type": "workspace", + "subject_id": "", + "contains": { + "subject_type": "project", + "subject_id": "" + } + } }, - "limit": { - "gpu": 8 - } + "priority": 5, + "priority_lane": false, + "concurrency_limit": 20, + "grants": [ + { + "class": "standard-workers", + "resources": [ + {"resource": "CPU", "reserved": 0, "limit": 80, "unit": "CPU", "missing_action": "reject"}, + {"resource": "memory", "reserved": 0, "limit": 320, "unit": "GiB", "missing_action": "reject"} + ] + } + ] } ``` -**Step** +When both projects submit at the same time, interactive requests are considered +first because their policy priority is higher. Nightly work can fill idle CPU +and memory, but it has no reservation. If an interactive request arrives while +nightly work holds reclaimable capacity, Resource Manager may preempt the +nightly request. If the nightly request used `reclaim_tolerance="none"`, it +would not have been admitted through this zero-reservation policy in the first +place. -```python -from zenml import step -from zenml.config import ResourceSettings +### Access pattern: Keep one account from flooding the queue +Use an account-level policy when one user or automation account can create many +runs. This keeps a burst of runs from one account in a bounded queue. -@step( - settings={ - "resources": ResourceSettings( - gpu_count=6, - preemptible=True, - ) +```json +{ + "pool": "workspace-standard-compute", + "subject_selector": { + "subject_type": "account", + "subject_id": "", + "attributes": { + "is_service_account": false } -) -def train() -> None: - ... + }, + "priority": 15, + "priority_lane": false, + "concurrency_limit": 2, + "grants": [ + { + "class": "standard-workers", + "resources": [ + {"resource": "CPU", "reserved": 0, "limit": 16, "unit": "CPU", "missing_action": "reject"}, + {"resource": "memory", "reserved": 0, "limit": 64, "unit": "GiB", "missing_action": "reject"} + ] + } + ] +} ``` -**Outcome:** Allocated immediately (no queue). Four GPUs count against the -policy reserved share; two are borrowed from free pool capacity (between -reserved and limit, and pool must still have free units). +If the account launches ten pipeline runs and each run has several dynamic +steps, the account policy caps active requests at two. Other users in the +workspace can still be admitted through broader workspace or project policies +if they have matching grants. ---- +## Example 2: One GPU Machine Shape -### Non-preemptible step stays inside reserved +Now add an accelerator while keeping the infrastructure simple. Assume one GPU +worker group made of four identical H200 machines. Each machine provides two +H200 GPUs plus CPU and memory. Pipeline authors request H200 capacity, and +ZenML selects the matching pool class. -**Story:** Same pool and policy. Production wants two GPUs and opts out of -preemption. Two is within the four-GPU reservation. +The strategy is: -**Pool** +* Add a custom descriptor named `h200` with built-in kind `gpu`. +* Create one class that bundles H200 GPUs, CPU, and memory together. +* Add a production policy with reserved capacity. +* Add a lower-priority experiment policy that can use leftover capacity. -```json -{ - "name": "datacenter-gpus", - "capacity": { - "gpu": 8 - } -} -``` +This example shows why descriptor `kind` is important. A pipeline author can +use `gpu_count=1`, and Resource Manager can match `h200` because the descriptor +uses `kind: "gpu"`. -**Policy** +### Descriptor ```json { - "pool": "datacenter-gpus", - "component": "team-ml-orch", - "component_type": "orchestrator", - "priority": 10, - "reserved": { - "gpu": 4 + "name": "h200", + "kind": "gpu", + "description": "NVIDIA H200 accelerator", + "attributes": { + "vendor": "nvidia", + "model": "H200", + "vram_gb": 141 }, - "limit": { - "gpu": 8 - } + "units": [ + {"name": "GPU", "multiplier": 1} + ] } ``` -**Step** +### Inventory -```python -@step( - settings={ - "resources": ResourceSettings( - gpu_count=2, - preemptible=False, - ) +```json +{ + "name": "team-h200-machines", + "description": "Four H200 training machines owned by the ML platform team", + "rank": 50, + "accounting_mode": "authoritative", + "concurrency_limit": 8, + "target_bindings": [ + { + "target_selector": { + "subject_type": "organization", + "subject_id": "", + "contains": { + "subject_type": "workspace", + "subject_id": "", + "contains": { + "subject_type": "component", + "subject_id": "", + "attributes": { + "component_type": "step_operator" + } + } + } + } } -) -def production_train() -> None: - ... + ], + "classes": [ + { + "class": "h200-standard", + "rank": 100, + "reclaimable": "never", + "concurrency_limit": 8, + "resources": [ + {"resource": "h200", "quantity": 8, "unit": "GPU"}, + {"resource": "CPU", "quantity": 192, "unit": "CPU"}, + {"resource": "memory", "quantity": 1536, "unit": "GiB"} + ], + "attributes": { + "accelerator": "h200", + "machine_count": "4", + "owner": "ml-platform" + } + } + ] +} ``` -**Outcome:** Allocated (assuming no other contention). Non-preemptible work must -satisfy requested ≤ reserved per key; `2 ≤ 4` passes. +The class total is the aggregate capacity of the machine group. The class is +not a single resource row; it is the shape Resource Manager can allocate as one +candidate. A request for 2 H200 GPUs, 16 CPUs, and 128 GiB memory must fit in +this one class. ---- +### Access pattern: Production reservation -### Non-preemptible step beyond reserved +```json +{ + "pool": "team-h200-machines", + "subject_selector": { + "subject_type": "organization", + "subject_id": "", + "contains": { + "subject_type": "team", + "subject_id": "" + } + }, + "priority": 100, + "priority_lane": false, + "concurrency_limit": 4, + "grants": [ + { + "class": "h200-standard", + "resources": [ + {"resource": "h200", "reserved": 4, "limit": 8, "unit": "GPU", "missing_action": "reject"}, + {"resource": "CPU", "reserved": 64, "limit": 192, "unit": "CPU", "missing_action": "default", "default_value": 8}, + {"resource": "memory", "reserved": 512, "limit": 1536, "unit": "GiB", "missing_action": "default", "default_value": 64} + ] + } + ] +} +``` -**Story:** Same pool and policy. Production asks for six GPUs but refuses -preemption. Non-preemptible work cannot use the “borrow” band above reserved. +The production team gets half the GPUs as protected reserved share and can +burst to the whole class when it is free. CPU and memory have defaults so a +GPU-only request from an author still accounts for a reasonable CPU/memory +footprint. -**Pool** +### Access pattern: Experiment burst ```json { - "name": "datacenter-gpus", - "capacity": { - "gpu": 8 - } + "pool": "team-h200-machines", + "subject_selector": { + "subject_type": "organization", + "subject_id": "", + "contains": { + "subject_type": "workspace", + "subject_id": "", + "contains": { + "subject_type": "project", + "subject_id": "" + } + } + }, + "priority": 20, + "priority_lane": false, + "concurrency_limit": 6, + "grants": [ + { + "class": "h200-standard", + "resources": [ + {"resource": "h200", "reserved": 0, "limit": 4, "unit": "GPU", "missing_action": "reject"}, + {"resource": "CPU", "reserved": 0, "limit": 96, "unit": "CPU", "missing_action": "default", "default_value": 4}, + {"resource": "memory", "reserved": 0, "limit": 768, "unit": "GiB", "missing_action": "default", "default_value": 32} + ] + } + ] } ``` -**Policy** +The experiment project has no reservation. It can consume idle capacity, but a +non-reclaimable request from that project is rejected because there is no +reserved share. Authors should use `reclaim_tolerance="coordinated"` or +`"any"` for this policy. + +### Runtime behavior: reservation, limit, and reclaim + +For example, production has two active requests using 2 H200 GPUs each. The +production reservation is now full, but the production limit is still 8 H200s. +A third production request for 2 H200 GPUs can still run if idle class capacity +exists, because the limit allows burst above the reservation. + +In another case, experiments already use 4 H200 GPUs through the lower-priority +policy and a production request for 4 H200 GPUs arrives with +`reclaim_tolerance="none"`. Resource Manager first checks whether the request +fits the production reserved share. If it does not, the request queues or is +rejected according to the admission route; it will not depend on interrupting +experiments. If the production request allows coordinated reclaim and has +higher priority, Resource Manager may select lower-priority experiment +allocations as victims. + +### Access pattern: Avoid self-preemption inside one project + +Preemption groups are useful when different policies match related work. The +selector below tells Resource Manager that requests from the same experiments +project should not preempt one another just because one matched a higher +priority policy. ```json { - "pool": "datacenter-gpus", - "component": "team-ml-orch", - "component_type": "orchestrator", - "priority": 10, - "reserved": { - "gpu": 4 + "pool": "team-h200-machines", + "subject_selector": { + "subject_type": "organization", + "subject_id": "", + "contains": { + "subject_type": "workspace", + "subject_id": "", + "contains": { + "subject_type": "project", + "subject_id": "" + } + } }, - "limit": { - "gpu": 8 - } + "preemption_group": { + "subject_type": "organization", + "subject_id": "", + "contains": { + "subject_type": "workspace", + "subject_id": "", + "contains": { + "subject_type": "project", + "subject_id": "" + } + } + }, + "priority": 60, + "priority_lane": false, + "concurrency_limit": 4, + "grants": [ + { + "class": "h200-standard", + "resources": [ + {"resource": "h200", "reserved": 0, "limit": 4, "unit": "GPU", "missing_action": "reject"}, + {"resource": "CPU", "reserved": 0, "limit": 96, "unit": "CPU", "missing_action": "default", "default_value": 8}, + {"resource": "memory", "reserved": 0, "limit": 768, "unit": "GiB", "missing_action": "default", "default_value": 64} + ] + } + ] } ``` -**Step** +If two pipeline runs from the same project are already allocated and a third +run from the same project arrives through this higher-priority policy, the +preemption group prevents Resource Manager from selecting the first two runs as +victims. The group does not protect unrelated projects. This keeps preemption +available for unrelated lower-priority work while preventing related runs from +interrupting each other. -```python -@step( - settings={ - "resources": ResourceSettings( - gpu_count=6, - preemptible=False, - ) - } -) -def too_large_production_train() -> None: - ... -``` +## Example 3: Kubernetes Cluster With Two Node Pools -**Outcome:** Rejected immediately (dynamic run fails fast). Six exceeds reserved -(4) for `gpu`; non-preemptible requests cannot borrow up to limit. +This example models a Kubernetes cluster with two real node pools: ---- +* `h200-reserved`: expensive reserved H200 nodes for production training. +* `a10-burst`: cheaper shared A10 nodes for experiments and sweeps. -## CPU, memory, and step slots: unbounded vs metered +One ZenML resource pool models both Kubernetes node pools as separate classes. +The pool is the shared accounting boundary. The classes are the infrastructure +shapes inside that boundary. -### GPU-only pool—CPU and memory not quota’d +The strategy is: -**Story:** You only modeled GPUs on the pool and policy. The step still sends -`mcpu` and `memory_mb` on the request, but those keys are unbounded at the pool -layer when omitted, and this policy omits them too—so they do not block -non-preemptible work. +* Use class rank to prefer the reserved H200 class when a request can run + there. +* Use `reclaimable: "never"` for reserved nodes and `"coordinated"` for burst + nodes. +* Use class target settings to inject node selectors and tolerations. +* Attach different policies for production and research. -**Pool** +### Inventory ```json { - "name": "training", - "capacity": { - "gpu": 4 + "name": "k8s-prod-eu-gpu", + "description": "Production Kubernetes GPU node pools in prod-eu", + "rank": 100, + "accounting_mode": "authoritative", + "concurrency_limit": 48, + "target_bindings": [ + { + "target_selector": { + "any": [ + { + "subject_type": "organization", + "subject_id": "", + "contains": { + "subject_type": "workspace", + "subject_id": "", + "contains": { + "subject_type": "component", + "subject_id": "", + "attributes": { + "component_type": "step_operator" + } + } + } + }, + { + "subject_type": "organization", + "subject_id": "", + "contains": { + "subject_type": "workspace", + "subject_id": "", + "contains": { + "subject_type": "service_connector", + "subject_id": "", + "attributes": { + "connector_type": "gcp" + } + } + } + } + ] + } + } + ], + "classes": [ + { + "class": "h200-reserved", + "rank": 100, + "reclaimable": "never", + "concurrency_limit": 8, + "resources": [ + {"resource": "h200", "quantity": 8, "unit": "GPU"}, + {"resource": "CPU", "quantity": 128, "unit": "CPU"}, + {"resource": "memory", "quantity": 1024, "unit": "GiB"} + ], + "target_settings": [ + { + "target_type": "component", + "settings": { + "pod_settings": { + "node_selector": { + "cloud.google.com/gke-nodepool": "h200-reserved" + }, + "tolerations": [ + { + "key": "dedicated", + "operator": "Equal", + "value": "h200", + "effect": "NoSchedule" + } + ] + } + } + } + ], + "attributes": { + "accelerator": "h200", + "tier": "tier-a-gpu", + "purchase": "reserved" + } + }, + { + "class": "a10-burst", + "rank": 30, + "reclaimable": "coordinated", + "concurrency_limit": 24, + "resources": [ + {"resource": "GPU", "quantity": 24, "unit": "GPU"}, + {"resource": "CPU", "quantity": null}, + {"resource": "memory", "quantity": null} + ], + "target_settings": [ + { + "target_type": "component", + "settings": { + "pod_settings": { + "node_selector": { + "cloud.google.com/gke-nodepool": "a10-burst" + }, + "labels": { + "workload-tier": "best-effort" + } + } + } + } + ], + "attributes": { + "accelerator": "a10", + "tier": "tier-b-gpu", + "purchase": "shared-burst" + } + } + ], + "target_settings": [ + { + "target_type": "service_connector", + "settings": { + "connector_id": "", + "resource_type": "kubernetes-cluster", + "resource_id": "prod-eu" + } + } + ], + "attributes": { + "cluster": "prod-eu", + "region": "eu-west-1", + "platform": "kubernetes" } } ``` -**Policy** +The `h200-reserved` class has finite CPU and memory because production capacity +is fully accounted. The `a10-burst` class tracks CPU and memory as unlimited: +GPU quantity and concurrency are enforced, while CPU and memory remain visible +on requests for traceability. + +### Access pattern: Production reservation ```json { - "pool": "training", - "component": "k8s-orch", - "component_type": "orchestrator", - "priority": 10, - "reserved": { - "gpu": 2 + "pool": "k8s-prod-eu-gpu", + "subject_selector": { + "subject_type": "organization", + "subject_id": "", + "contains": { + "subject_type": "team", + "subject_id": "" + } }, - "limit": { - "gpu": 4 - } -} -``` - -**Step** - -```python -@step( - settings={ - "resources": ResourceSettings( - gpu_count=1, - cpu_count=32, - memory="64GiB", - preemptible=False, - ) + "priority": 100, + "priority_lane": false, + "concurrency_limit": 8, + "grants": [ + { + "class": "h200-reserved", + "concurrency_limit": 6, + "resources": [ + {"resource": "h200", "reserved": 4, "limit": 8, "unit": "GPU", "missing_action": "reject"}, + {"resource": "CPU", "reserved": 32, "limit": 128, "unit": "CPU", "missing_action": "default", "default_value": 4}, + {"resource": "memory", "reserved": 256, "limit": 1024, "unit": "GiB", "missing_action": "default", "default_value": 16} + ], + "target_settings": [ + { + "target_type": "component", + "settings": { + "pod_settings": { + "labels": { + "zenml.io/resource-policy": "production" + } + } + } + } + ] } -) -def hungry_but_ok_on_gpu() -> None: - ... + ] +} ``` -**Outcome:** Allocated if nothing else is wrong. Only `gpu` is gated here; -`mcpu` / `memory_mb` / `step_run` are not limited by pool or policy in this -pattern. CPU and memory remain informational unless you add rows later. - ---- - -### Non-preemptible CPU inside policy reserved +Production work gets high priority and reserved H200 share. Authors should use +`reclaim_tolerance="none"` for critical steps. Those requests wait for reserved +capacity instead of borrowing interruptible capacity. -**Story:** You cap milli-CPU on the pool, then split it with reserved / limit -on the policy. Non-preemptible CPU demand must fit reserved per key. - -**Pool** +### Access pattern: Research burst tier ```json { - "name": "training", - "capacity": { - "gpu": 4, - "mcpu": 32000 - } + "pool": "k8s-prod-eu-gpu", + "subject_selector": { + "subject_type": "organization", + "subject_id": "", + "contains": { + "subject_type": "workspace", + "subject_id": "", + "contains": { + "subject_type": "project", + "subject_id": "" + } + } + }, + "priority": 30, + "priority_lane": false, + "concurrency_limit": 16, + "grants": [ + { + "class": "a10-burst", + "concurrency_limit": 12, + "resources": [ + {"resource": "GPU", "reserved": 0, "limit": 16, "unit": "GPU", "missing_action": "reject"}, + {"resource": "CPU", "reserved": 0, "limit": null, "unit": "CPU", "missing_action": "default", "default_value": 4}, + {"resource": "memory", "reserved": 0, "limit": null, "unit": "GiB", "missing_action": "default", "default_value": 24} + ] + } + ] } ``` -**Policy** +The research policy deliberately grants only `a10-burst`. Even if H200 capacity +is idle, this project will not consume it through this policy. This expresses +the rule that experiments may use the burst tier, but not reserved production +accelerators. + +### Access pattern: Platform validation can use both classes + +Image, driver, and runtime validation often needs access to both node pools. +That is a different access rule from either production or research, so keep the +same inventory and add a policy with two grants. ```json { - "pool": "training", - "component": "k8s-orch", - "component_type": "orchestrator", - "priority": 10, - "reserved": { - "gpu": 2, - "mcpu": 4000 + "pool": "k8s-prod-eu-gpu", + "subject_selector": { + "subject_type": "organization", + "subject_id": "", + "contains": { + "subject_type": "team", + "subject_id": "" + } }, - "limit": { - "gpu": 4, - "mcpu": 32000 - } + "priority": 80, + "priority_lane": false, + "concurrency_limit": 10, + "grants": [ + { + "class": "h200-reserved", + "concurrency_limit": 2, + "resources": [ + {"resource": "h200", "reserved": 0, "limit": 2, "unit": "GPU", "missing_action": "reject"}, + {"resource": "CPU", "reserved": 0, "limit": 32, "unit": "CPU", "missing_action": "default", "default_value": 8}, + {"resource": "memory", "reserved": 0, "limit": 256, "unit": "GiB", "missing_action": "default", "default_value": 64} + ] + }, + { + "class": "a10-burst", + "concurrency_limit": 8, + "resources": [ + {"resource": "GPU", "reserved": 0, "limit": 8, "unit": "GPU", "missing_action": "reject"}, + {"resource": "CPU", "reserved": 0, "limit": null, "unit": "CPU", "missing_action": "default", "default_value": 4}, + {"resource": "memory", "reserved": 0, "limit": null, "unit": "GiB", "missing_action": "default", "default_value": 24} + ] + } + ] } ``` -**Step** +The grant-level concurrency limits keep validation from occupying many H200 +slots while still allowing broader A10 testing. A request pinned to +`gpu_class="h200-reserved"` can only use the H200 grant. A request without a +class pin can be planned against either class if all of its demands fit. -```python -@step( - settings={ - "resources": ResourceSettings( - gpu_count=1, - cpu_count=2, - preemptible=False, - ) - } -) -def fits_reserved_cpu() -> None: - ... -``` +### Runtime behavior: class choice under pressure -**Outcome:** Allocated. `cpu_count=2` → `mcpu` 2000 ≤ reserved 4000, and `gpu` is -valid. +If a production request asks for `h200` explicitly, only the `h200-reserved` +class can satisfy it because `h200` is a specific descriptor. If a research +request asks for the built-in `GPU` kind, the `a10-burst` class can satisfy it, +and the H200 class will not be used unless a policy grant and descriptor match +allow that route. ---- +If a platform validation request can use either class, class rank prefers H200 +when both are available. If H200 class concurrency is already full, the A10 +grant can still admit the request if the request is compatible with the stock +`GPU` descriptor and the lower-ranked class has room. + +## Example 4: Kubernetes Cluster With Per-Node Classes + +The previous Kubernetes example models node pools. Some clusters require a more +precise model: each node is meaningful because it has local disks, attached +licenses, a special network path, or a workload that must be placed on one +named machine. In that case, model each Kubernetes node as its own class. + +For simplicity, assume one cluster with three GPU nodes: + +* `gpu-node-a`: reserved production node with four H200 GPUs. +* `gpu-node-b`: general shared node with four H200 GPUs. +* `gpu-node-c`: burst node with eight A10 GPUs. + +The strategy is: -### Non-preemptible CPU over policy reserved +* Use one ZenML pool for the cluster. +* Use one class per node. +* Put a node selector in each class's component target settings. +* Use grants to decide which teams can use which node classes. +* Use class rank to express preferred placement when a request is not pinned to + a specific node class. -**Pool** +### Inventory ```json { - "name": "training", - "capacity": { - "gpu": 4, - "mcpu": 32000 + "name": "k8s-prod-node-level", + "description": "Node-level model of three GPU nodes in one Kubernetes cluster", + "rank": 85, + "accounting_mode": "authoritative", + "concurrency_limit": 24, + "target_bindings": [ + { + "target_selector": { + "subject_type": "organization", + "subject_id": "", + "contains": { + "subject_type": "workspace", + "subject_id": "", + "contains": { + "subject_type": "component", + "subject_id": "", + "attributes": { + "component_type": "step_operator" + } + } + } + } + } + ], + "classes": [ + { + "class": "gpu-node-a", + "rank": 100, + "reclaimable": "never", + "concurrency_limit": 4, + "resources": [ + {"resource": "h200", "quantity": 4, "unit": "GPU"}, + {"resource": "CPU", "quantity": 64, "unit": "CPU"}, + {"resource": "memory", "quantity": 512, "unit": "GiB"} + ], + "target_settings": [ + { + "target_type": "component", + "settings": { + "pod_settings": { + "node_selector": { + "kubernetes.io/hostname": "gpu-node-a" + }, + "tolerations": [ + { + "key": "dedicated", + "operator": "Equal", + "value": "production-gpu", + "effect": "NoSchedule" + } + ] + } + } + } + ], + "attributes": { + "node": "gpu-node-a", + "accelerator": "h200", + "tier": "reserved" + } + }, + { + "class": "gpu-node-b", + "rank": 70, + "reclaimable": "coordinated", + "concurrency_limit": 4, + "resources": [ + {"resource": "h200", "quantity": 4, "unit": "GPU"}, + {"resource": "CPU", "quantity": 64, "unit": "CPU"}, + {"resource": "memory", "quantity": 512, "unit": "GiB"} + ], + "target_settings": [ + { + "target_type": "component", + "settings": { + "pod_settings": { + "node_selector": { + "kubernetes.io/hostname": "gpu-node-b" + } + } + } + } + ], + "attributes": { + "node": "gpu-node-b", + "accelerator": "h200", + "tier": "shared" + } + }, + { + "class": "gpu-node-c", + "rank": 30, + "reclaimable": "coordinated", + "concurrency_limit": 8, + "resources": [ + {"resource": "GPU", "quantity": 8, "unit": "GPU"}, + {"resource": "CPU", "quantity": 96, "unit": "CPU"}, + {"resource": "memory", "quantity": 384, "unit": "GiB"} + ], + "target_settings": [ + { + "target_type": "component", + "settings": { + "pod_settings": { + "node_selector": { + "kubernetes.io/hostname": "gpu-node-c" + }, + "labels": { + "workload-tier": "burst" + } + } + } + } + ], + "attributes": { + "node": "gpu-node-c", + "accelerator": "a10", + "tier": "burst" + } + } + ], + "attributes": { + "cluster": "prod-eu", + "modeling": "per-node" } } ``` -**Policy** +Use this level of detail when node identity matters. Every class has its own +pod settings, so Resource Manager can select a class and the target component +can receive the node selector that makes Kubernetes place the pod on the +intended node. + +### Access pattern: Production can use only the reserved node ```json { - "pool": "training", - "component": "k8s-orch", - "component_type": "orchestrator", - "priority": 10, - "reserved": { - "gpu": 2, - "mcpu": 4000 + "pool": "k8s-prod-node-level", + "subject_selector": { + "subject_type": "organization", + "subject_id": "", + "contains": { + "subject_type": "team", + "subject_id": "" + } }, - "limit": { - "gpu": 4, - "mcpu": 32000 - } + "priority": 120, + "priority_lane": false, + "concurrency_limit": 4, + "grants": [ + { + "class": "gpu-node-a", + "resources": [ + {"resource": "h200", "reserved": 4, "limit": 4, "unit": "GPU", "missing_action": "reject"}, + {"resource": "CPU", "reserved": 32, "limit": 64, "unit": "CPU", "missing_action": "default", "default_value": 8}, + {"resource": "memory", "reserved": 256, "limit": 512, "unit": "GiB", "missing_action": "default", "default_value": 64} + ] + } + ] } ``` -**Step** +This policy intentionally does not grant `gpu-node-b` or `gpu-node-c`. +Production training gets deterministic placement and a protected reservation on +the reserved node. If the production team requests 6 H200 GPUs, the request is +rejected because no single granted node class can satisfy that bundle. -```python -@step( - settings={ - "resources": ResourceSettings( - gpu_count=1, - cpu_count=8, - preemptible=False, - ) +### Access pattern: Research can use shared H200 and burst A10 + +```json +{ + "pool": "k8s-prod-node-level", + "subject_selector": { + "subject_type": "organization", + "subject_id": "", + "contains": { + "subject_type": "workspace", + "subject_id": "", + "contains": { + "subject_type": "project", + "subject_id": "" + } } -) -def exceeds_reserved_cpu() -> None: - ... + }, + "priority": 35, + "priority_lane": false, + "concurrency_limit": 10, + "grants": [ + { + "class": "gpu-node-b", + "resources": [ + {"resource": "h200", "reserved": 0, "limit": 4, "unit": "GPU", "missing_action": "reject"}, + {"resource": "CPU", "reserved": 0, "limit": 64, "unit": "CPU", "missing_action": "default", "default_value": 8}, + {"resource": "memory", "reserved": 0, "limit": 512, "unit": "GiB", "missing_action": "default", "default_value": 64} + ] + }, + { + "class": "gpu-node-c", + "resources": [ + {"resource": "GPU", "reserved": 0, "limit": 8, "unit": "GPU", "missing_action": "reject"}, + {"resource": "CPU", "reserved": 0, "limit": 96, "unit": "CPU", "missing_action": "default", "default_value": 4}, + {"resource": "memory", "reserved": 0, "limit": 384, "unit": "GiB", "missing_action": "default", "default_value": 24} + ] + } + ] +} ``` -**Outcome:** Rejected. `cpu_count=8` → `mcpu` 8000 > reserved 4000; -non-preemptible work cannot borrow toward limit on `mcpu`. - ---- - -### Preemptible CPU burst with policy `mcpu` rows +Research has no reserved share. It can use `gpu-node-b` for H200 experiments +and `gpu-node-c` for generic GPU work. Because `gpu-node-c` uses the stock +`GPU` descriptor, it can satisfy a generic `gpu_count` request. It does not +satisfy a request that explicitly demands the `h200` descriptor. -**Pool** and **Policy:** same as *Non-preemptible CPU inside policy reserved* -(pool includes `gpu` and `mcpu`; policy sets both keys). +### Access pattern: Node repair drains one class -**Step** +When a node needs maintenance, keep the pool and other classes alive. Add or +update a temporary policy that prevents new work from using that node while +still allowing an operations team to run one diagnostic job. -```python -@step( - settings={ - "resources": ResourceSettings( - gpu_count=1, - cpu_count=8, - preemptible=True, - ) +```json +{ + "pool": "k8s-prod-node-level", + "subject_selector": { + "subject_type": "organization", + "subject_id": "", + "contains": { + "subject_type": "team", + "subject_id": "" } -) -def preemptible_cpu_burst() -> None: - ... + }, + "priority": 200, + "priority_lane": true, + "concurrency_limit": 1, + "grants": [ + { + "class": "gpu-node-b", + "concurrency_limit": 1, + "resources": [ + {"resource": "h200", "reserved": 0, "limit": 1, "unit": "GPU", "missing_action": "reject"}, + {"resource": "CPU", "reserved": 0, "limit": 8, "unit": "CPU", "missing_action": "default", "default_value": 4}, + {"resource": "memory", "reserved": 0, "limit": 64, "unit": "GiB", "missing_action": "default", "default_value": 16} + ], + "target_settings": [ + { + "target_type": "component", + "settings": { + "pod_settings": { + "labels": { + "maintenance-window": "node-b" + } + } + } + } + ] + } + ] +} ``` -**Outcome:** May allocate using headroom up to limit on `mcpu` (and pool free -capacity), analogous to GPU borrowing. +The operational decision is expressed in access policy, not by deleting +inventory. Existing lower-priority reclaimable work on `gpu-node-b` can be +preempted if the platform diagnostic request needs the slot. New research work +can still use `gpu-node-c`; production work can still use `gpu-node-a`. ---- +### Runtime behavior: node-level placement + +If an author pins `gpu_class="gpu-node-b"`, only grants for `gpu-node-b` can +admit the request. If the author does not pin a class and the demands are +compatible with more than one granted class, Resource Manager uses policy +priority and class rank to order candidates. The selected class contributes the +pod settings that place the workload on a particular Kubernetes node. + +The trade-off is operational overhead. Per-node classes are precise, but class +resources and node selectors must stay aligned with the cluster. For most +clusters, node-pool classes are simpler. Use per-node classes when the node +itself is a scheduling unit that users or operators care about. + +## Example 5: Multiple Kubernetes Clusters as Classes + +Another modeling style is one pool that represents a fleet of reachable +clusters. Each class is a cluster, not a node or node pool. Use this pattern +when the same logical step operator can launch into several clusters, or when +the platform integration chooses a cluster based on class-level target +settings. -### Preemptible when pool lists `mcpu` but policy omits it +The important difference from the multi-region example later is that this is +one pool with several classes. Each class has its own service connector +settings, so the selected class determines which cluster credentials and +cluster identifier are passed to the target. -**Story:** The pool caps total milli-CPU. With no `mcpu` on the policy, reserved -defaults to 0 and limit falls back to the pool total. +The strategy is: -**Pool** +* Use one pool for the fleet. +* Use one class per cluster. +* Put class-level `service_connector` target settings on each class. +* Use class rank to prefer the primary cluster and burst into the secondary + cluster. +* Use policies to give different teams access to different clusters. + +### Inventory ```json { - "name": "training", - "capacity": { - "gpu": 4, - "mcpu": 8000 + "name": "k8s-cross-cluster-training", + "description": "Training capacity across three Kubernetes clusters", + "rank": 75, + "accounting_mode": "authoritative", + "concurrency_limit": 60, + "target_bindings": [ + { + "target_selector": { + "any": [ + { + "subject_type": "organization", + "subject_id": "", + "contains": { + "subject_type": "workspace", + "subject_id": "", + "contains": { + "subject_type": "component", + "subject_id": "", + "attributes": { + "component_type": "step_operator" + } + } + } + }, + { + "subject_type": "organization", + "subject_id": "", + "contains": { + "subject_type": "workspace", + "subject_id": "", + "contains": { + "subject_type": "service_connector", + "attributes": { + "connector_type": "kubernetes" + } + } + } + } + ] + } + } + ], + "classes": [ + { + "class": "stockholm-cluster", + "rank": 100, + "reclaimable": "never", + "concurrency_limit": 16, + "resources": [ + {"resource": "h200", "quantity": 16, "unit": "GPU"}, + {"resource": "CPU", "quantity": 256, "unit": "CPU"}, + {"resource": "memory", "quantity": 2048, "unit": "GiB"} + ], + "target_settings": [ + { + "target_type": "service_connector", + "settings": { + "connector_id": "", + "resource_type": "kubernetes-cluster", + "resource_id": "stockholm-prod" + } + }, + { + "target_type": "component", + "settings": { + "namespace": "training-prod", + "pod_settings": { + "node_selector": { + "zenml.io/gpu-fleet": "h200" + } + } + } + } + ], + "attributes": { + "cluster": "stockholm-prod", + "region": "stockholm", + "tier": "primary" + } + }, + { + "class": "frankfurt-cluster", + "rank": 70, + "reclaimable": "coordinated", + "concurrency_limit": 12, + "resources": [ + {"resource": "h200", "quantity": 12, "unit": "GPU"}, + {"resource": "CPU", "quantity": 192, "unit": "CPU"}, + {"resource": "memory", "quantity": 1536, "unit": "GiB"} + ], + "target_settings": [ + { + "target_type": "service_connector", + "settings": { + "connector_id": "", + "resource_type": "kubernetes-cluster", + "resource_id": "frankfurt-prod" + } + }, + { + "target_type": "component", + "settings": { + "namespace": "training-prod", + "pod_settings": { + "node_selector": { + "zenml.io/gpu-fleet": "h200" + }, + "labels": { + "cluster-preference": "secondary" + } + } + } + } + ], + "attributes": { + "cluster": "frankfurt-prod", + "region": "frankfurt", + "tier": "secondary" + } + }, + { + "class": "madrid-burst-cluster", + "rank": 30, + "reclaimable": "coordinated", + "concurrency_limit": 24, + "resources": [ + {"resource": "GPU", "quantity": 24, "unit": "GPU"}, + {"resource": "CPU", "quantity": null}, + {"resource": "memory", "quantity": null} + ], + "target_settings": [ + { + "target_type": "service_connector", + "settings": { + "connector_id": "", + "resource_type": "kubernetes-cluster", + "resource_id": "madrid-burst" + } + }, + { + "target_type": "component", + "settings": { + "namespace": "training-burst", + "pod_settings": { + "labels": { + "workload-tier": "burst" + } + } + } + } + ], + "attributes": { + "cluster": "madrid-burst", + "region": "madrid", + "tier": "burst" + } + } + ], + "attributes": { + "modeling": "cluster-as-class", + "fleet": "europe-training" } } ``` -**Policy** (only `gpu`) +This model makes class selection carry both accounting and access details. If +Resource Manager selects `frankfurt-cluster`, the target receives the Frankfurt +service connector settings. If it selects `madrid-burst-cluster`, the target +receives the Madrid connector and burst namespace settings. + +### Access pattern: Primary cluster first, secondary cluster as burst ```json { - "pool": "training", - "component": "k8s-orch", - "component_type": "orchestrator", - "priority": 10, - "reserved": { - "gpu": 2 + "pool": "k8s-cross-cluster-training", + "subject_selector": { + "subject_type": "organization", + "subject_id": "", + "contains": { + "subject_type": "team", + "subject_id": "" + } }, - "limit": { - "gpu": 4 - } + "priority": 90, + "priority_lane": false, + "concurrency_limit": 20, + "grants": [ + { + "class": "stockholm-cluster", + "resources": [ + {"resource": "h200", "reserved": 8, "limit": 16, "unit": "GPU", "missing_action": "reject"}, + {"resource": "CPU", "reserved": 128, "limit": 256, "unit": "CPU", "missing_action": "default", "default_value": 16}, + {"resource": "memory", "reserved": 1024, "limit": 2048, "unit": "GiB", "missing_action": "default", "default_value": 128} + ] + }, + { + "class": "frankfurt-cluster", + "resources": [ + {"resource": "h200", "reserved": 0, "limit": 8, "unit": "GPU", "missing_action": "reject"}, + {"resource": "CPU", "reserved": 0, "limit": 128, "unit": "CPU", "missing_action": "default", "default_value": 16}, + {"resource": "memory", "reserved": 0, "limit": 1024, "unit": "GiB", "missing_action": "default", "default_value": 128} + ] + } + ] } ``` -**Step** +The team gets a protected share in Stockholm and burst access in Frankfurt. +Because `stockholm-cluster` has higher class rank, compatible requests prefer +Stockholm when both grants can admit them. When Stockholm is full, Frankfurt can +admit the request through the second grant if the request is reclaimable enough +for that zero-reservation route. -```python -@step( - settings={ - "resources": ResourceSettings( - gpu_count=1, - cpu_count=4, - preemptible=True, - ) +### Access pattern: Research can use only the burst cluster + +```json +{ + "pool": "k8s-cross-cluster-training", + "subject_selector": { + "subject_type": "organization", + "subject_id": "", + "contains": { + "subject_type": "workspace", + "subject_id": "", + "contains": { + "subject_type": "project", + "subject_id": "" + } } -) -def preemptible_with_pool_mcpu() -> None: - ... + }, + "priority": 25, + "priority_lane": false, + "concurrency_limit": 18, + "grants": [ + { + "class": "madrid-burst-cluster", + "concurrency_limit": 18, + "resources": [ + {"resource": "GPU", "reserved": 0, "limit": 24, "unit": "GPU", "missing_action": "reject"}, + {"resource": "CPU", "reserved": 0, "limit": null, "unit": "CPU", "missing_action": "default", "default_value": 4}, + {"resource": "memory", "reserved": 0, "limit": null, "unit": "GiB", "missing_action": "default", "default_value": 24} + ] + } + ] +} ``` -**Outcome:** Allocated or queued then allocated when possible. `mcpu` 4000 ≤ -effective limit 8000 (pool total). +This policy deliberately excludes the H200 clusters. Research sweeps can fill +the burst cluster, but they cannot consume the primary or secondary H200 +clusters through this policy. ---- +### Access pattern: One team gets direct access to one cluster -### Non-preemptible when pool lists `mcpu` but policy omits it +```json +{ + "pool": "k8s-cross-cluster-training", + "subject_selector": { + "subject_type": "organization", + "subject_id": "", + "contains": { + "subject_type": "team", + "subject_id": "" + } + }, + "priority": 110, + "priority_lane": false, + "concurrency_limit": 6, + "grants": [ + { + "class": "frankfurt-cluster", + "resources": [ + {"resource": "h200", "reserved": 4, "limit": 12, "unit": "GPU", "missing_action": "reject"}, + {"resource": "CPU", "reserved": 64, "limit": 192, "unit": "CPU", "missing_action": "default", "default_value": 16}, + {"resource": "memory", "reserved": 512, "limit": 1536, "unit": "GiB", "missing_action": "default", "default_value": 128} + ], + "target_settings": [ + { + "target_type": "component", + "settings": { + "pod_settings": { + "labels": { + "cluster-owner": "frankfurt-ops" + } + } + } + } + ] + } + ] +} +``` -**Pool** and **Policy:** same as the previous example. +The Frankfurt operations team gets high-priority access to its cluster without +receiving any route into Stockholm or Madrid. This pattern works well when +cluster ownership maps to teams and the organization still requires one +pool-level accounting surface. -**Step** +### Runtime behavior: selected class selects the connector -```python -@step( - settings={ - "resources": ResourceSettings( - gpu_count=1, - cpu_count=1, - preemptible=False, - ) - } -) -def non_preemptible_positive_mcpu_zero_reserved() -> None: - ... -``` +If a foundation models run requests 4 H200 GPUs and Stockholm has room, +Resource Manager selects the `stockholm-cluster` grant. The selected class +contributes the Stockholm service connector settings and the target launches in +`stockholm-prod`. -**Outcome:** Rejected. Any positive `mcpu` with reserved 0 fails for -non-preemptible work. Fix: add `mcpu` to the policy with enough reserved, or -remove `mcpu` from the pool if you wanted fully unbounded CPU at the pool layer. +If Stockholm is full and the run is allowed to use reclaimable burst capacity, +the Frankfurt grant becomes the next candidate. The request is still one +allocation plan; it does not split two GPUs into Stockholm and two into +Frankfurt. If the research sweeps project sends the same generic GPU request, +it only has a Madrid route because its policy only grants +`madrid-burst-cluster`. ---- +## Example 6: Multi-Region GPU Dispatch -### Capping concurrent steps with `step_run` +In some deployments, the same GPU model is available in more than one region or +cluster. Pipeline authors can request H200 capacity without choosing the region +manually. -**Story:** You want both GPUs and a ceiling on how many steps from this -orchestrator run at once. Each step always requests one `step_run`. +Model this as multiple pools, one per reachable execution target or region, +with similar classes and different target bindings. Pool rank represents the +placement preference. A higher-rank pool is tried first when both can admit the +request. Attributes make it possible for an advanced request to target one +region explicitly. -**Pool** +### Stockholm pool ```json { - "name": "training", - "capacity": { - "gpu": 16, - "step_run": 10 + "name": "h200-stockholm", + "description": "H200 capacity in the Stockholm Kubernetes cluster", + "rank": 100, + "accounting_mode": "authoritative", + "target_bindings": [ + { + "target_selector": { + "subject_type": "organization", + "subject_id": "", + "contains": { + "subject_type": "workspace", + "subject_id": "", + "contains": { + "subject_type": "component", + "subject_id": "", + "attributes": { + "component_type": "step_operator" + } + } + } + } + } + ], + "classes": [ + { + "class": "h200-reserved", + "rank": 100, + "reclaimable": "never", + "resources": [ + {"resource": "h200", "quantity": 16, "unit": "GPU"}, + {"resource": "CPU", "quantity": 256, "unit": "CPU"}, + {"resource": "memory", "quantity": 2048, "unit": "GiB"} + ], + "attributes": { + "region": "stockholm", + "tier": "tier-a-gpu" + } + } + ], + "attributes": { + "region": "stockholm", + "cluster": "se-prod" } } ``` -**Policy** +### Frankfurt pool ```json { - "pool": "training", - "component": "k8s-orch", - "component_type": "orchestrator", - "priority": 10, - "reserved": { - "gpu": 8, - "step_run": 4 - }, - "limit": { - "gpu": 16, - "step_run": 4 + "name": "h200-frankfurt", + "description": "H200 capacity in the Frankfurt Kubernetes cluster", + "rank": 80, + "accounting_mode": "authoritative", + "target_bindings": [ + { + "target_selector": { + "subject_type": "organization", + "subject_id": "", + "contains": { + "subject_type": "workspace", + "subject_id": "", + "contains": { + "subject_type": "component", + "subject_id": "", + "attributes": { + "component_type": "step_operator" + } + } + } + } + } + ], + "classes": [ + { + "class": "h200-reserved", + "rank": 100, + "reclaimable": "never", + "resources": [ + {"resource": "h200", "quantity": 8, "unit": "GPU"}, + {"resource": "CPU", "quantity": 128, "unit": "CPU"}, + {"resource": "memory", "quantity": 1024, "unit": "GiB"} + ], + "attributes": { + "region": "frankfurt", + "tier": "tier-a-gpu" + } + } + ], + "attributes": { + "region": "frankfurt", + "cluster": "de-prod" } } ``` -**Step** +### Access pattern: Same team policy on both pools -```python -@step( - settings={ - "resources": ResourceSettings( - gpu_count=2, - preemptible=True, - ) +Attach equivalent policies to both pools: + +```json +{ + "pool": "h200-stockholm", + "subject_selector": { + "subject_type": "organization", + "subject_id": "", + "contains": { + "subject_type": "team", + "subject_id": "" } -) -def train() -> None: - ... + }, + "priority": 90, + "priority_lane": false, + "grants": [ + { + "class": "h200-reserved", + "resources": [ + {"resource": "h200", "reserved": 8, "limit": 16, "unit": "GPU", "missing_action": "reject"}, + {"resource": "CPU", "reserved": 64, "limit": 256, "unit": "CPU", "missing_action": "default", "default_value": 8}, + {"resource": "memory", "reserved": 512, "limit": 2048, "unit": "GiB", "missing_action": "default", "default_value": 64} + ] + } + ] +} ``` -**Outcome:** The server grants only when both `gpu` and `step_run` have enough -free units. If GPUs are free but all `step_run` slots are taken, the request -waits in the queue. - ---- +The same policy for Frankfurt would use `"pool": "h200-frankfurt"` and smaller +reservation/limit values if that cluster has less capacity. -## Custom keys from `pool_resources` +### Runtime behavior: regional dispatch -### Custom key fully configured - -**Story:** You track a scarce license or device class with `pool_resources`. +```python +ResourceSettings( + gpu_count=4, + gpu_class="h200-reserved", + cpu_count=32, + memory="256GiB", + reclaim_tolerance="none", +) +``` -**Pool** +The author asks for the hardware class, not the region. If both pools can admit +the request, Stockholm wins because its pool rank is higher. If Stockholm is +full and Frankfurt can admit the request, Frankfurt can win instead, depending +on the policy and current queue state. -```json -{ - "name": "inference", - "capacity": { - "gpu": 8, - "tensorrt_sessions": 4 - } -} -``` +### Access pattern: Pin one project to Frankfurt -**Policy** +The same regional inventory can also enforce residency or cost decisions. Add a +project-specific policy only on the Frankfurt pool when a project must stay +there. ```json { - "pool": "inference", - "component": "gpu-step-op", - "component_type": "step_operator", - "priority": 10, - "reserved": { - "gpu": 2, - "tensorrt_sessions": 2 + "pool": "h200-frankfurt", + "subject_selector": { + "subject_type": "organization", + "subject_id": "", + "contains": { + "subject_type": "workspace", + "subject_id": "", + "contains": { + "subject_type": "project", + "subject_id": "" + } + } }, - "limit": { - "gpu": 8, - "tensorrt_sessions": 4 - } + "priority": 95, + "priority_lane": false, + "concurrency_limit": 4, + "grants": [ + { + "class": "h200-reserved", + "resources": [ + {"resource": "h200", "reserved": 2, "limit": 8, "unit": "GPU", "missing_action": "reject"}, + {"resource": "CPU", "reserved": 32, "limit": 128, "unit": "CPU", "missing_action": "default", "default_value": 8}, + {"resource": "memory", "reserved": 256, "limit": 1024, "unit": "GiB", "missing_action": "default", "default_value": 64} + ] + } + ] } ``` -**Step** +That project is not admitted by any policy on the Stockholm pool, so +Stockholm's higher pool rank does not matter for it. Other projects that match +policies on both pools still use rank-based dispatch. -```python -@step( - settings={ - "resources": ResourceSettings( - gpu_count=1, - pool_resources={"tensorrt_sessions": 1}, - preemptible=False, - ) - } -) -def infer() -> None: - ... -``` +### Runtime behavior: one request, several candidate pools -**Outcome:** Allocated when `1 ≤ reserved` for both `gpu` and -`tensorrt_sessions`. Unbounded defaults do not apply to custom keys—the pool -must list them. +When a request can target both step operators and matches policies on both +pools, Resource Manager creates candidate routes for Stockholm and Frankfurt. +The winning route must satisfy the entire bundle: GPUs, CPU, memory, +concurrency, target binding, and policy grant. It does not split a single step +across regions. ---- +## Example 7: Shared Organization Pool Across Workspaces -### Custom key on pool but missing from policy +Large organizations often buy GPUs centrally while teams work in many ZenML +workspaces. The goal is not "Workspace A owns a pool and Workspace B owns a +different pool." The goal is one organization-level pool governed by policies +that encode which teams, projects, or workspaces may use it. -**Story:** Same pool capacity; policy only defines `gpu`. +This example focuses on policy strategy: -**Pool** +* The inventory is organization-scoped and shared. +* A broad workspace policy lets a whole workspace use a small default share. +* A team policy gives a production team stronger reservations. +* A project policy caps a sandbox project so it cannot consume the whole pool. + +### Shared inventory ```json { - "name": "inference", - "capacity": { - "gpu": 8, - "tensorrt_sessions": 2 + "name": "org-shared-gpu-fleet", + "description": "Centrally purchased GPU fleet shared across workspaces", + "rank": 90, + "accounting_mode": "authoritative", + "concurrency_limit": 64, + "target_bindings": [ + { + "target_selector": { + "any": [ + { + "subject_type": "organization", + "subject_id": "", + "contains": { + "subject_type": "workspace", + "subject_id": "", + "contains": { + "subject_type": "component", + "subject_id": "", + "attributes": { + "component_type": "step_operator" + } + } + } + }, + { + "subject_type": "organization", + "subject_id": "", + "contains": { + "subject_type": "workspace", + "subject_id": "", + "contains": { + "subject_type": "component", + "subject_id": "", + "attributes": { + "component_type": "step_operator" + } + } + } + } + ] + } + } + ], + "classes": [ + { + "class": "tier-a-gpu", + "rank": 100, + "reclaimable": "never", + "concurrency_limit": 24, + "resources": [ + {"resource": "h200", "quantity": 24, "unit": "GPU"}, + {"resource": "CPU", "quantity": 384, "unit": "CPU"}, + {"resource": "memory", "quantity": 3072, "unit": "GiB"} + ], + "attributes": { + "tier": "tier-a-gpu", + "chargeback_group": "central-ml" + } + } + ], + "attributes": { + "scope": "organization", + "cost_center": "central-ml-platform" } } ``` -**Policy** +The target binding explicitly lists the step operators that can launch work +against this shared fleet. That is an infrastructure access decision. The +policies below are separate access decisions about who may consume that +infrastructure. + +### Access pattern: Default workspace access ```json { - "pool": "inference", - "component": "gpu-step-op", - "component_type": "step_operator", - "priority": 10, - "reserved": { - "gpu": 2 + "pool": "org-shared-gpu-fleet", + "subject_selector": { + "subject_type": "organization", + "subject_id": "", + "contains": { + "subject_type": "workspace", + "subject_id": "" + } }, - "limit": { - "gpu": 8 - } + "priority": 20, + "priority_lane": false, + "concurrency_limit": 8, + "grants": [ + { + "class": "tier-a-gpu", + "resources": [ + {"resource": "h200", "reserved": 0, "limit": 4, "unit": "GPU", "missing_action": "reject"}, + {"resource": "CPU", "reserved": 0, "limit": 64, "unit": "CPU", "missing_action": "default", "default_value": 4}, + {"resource": "memory", "reserved": 0, "limit": 512, "unit": "GiB", "missing_action": "default", "default_value": 32} + ] + } + ] } ``` -**Step** +This policy says: any request from Workspace A can use up to four H200s if they +are idle, but Workspace A does not own a protected share by default. -```python -@step( - settings={ - "resources": ResourceSettings( - gpu_count=1, - pool_resources={"tensorrt_sessions": 1}, - preemptible=False, - ) +### Access pattern: Production team reservation + +```json +{ + "pool": "org-shared-gpu-fleet", + "subject_selector": { + "subject_type": "organization", + "subject_id": "", + "contains": { + "subject_type": "team", + "subject_id": "" } -) -def infer() -> None: - ... + }, + "priority": 120, + "priority_lane": false, + "concurrency_limit": 12, + "grants": [ + { + "class": "tier-a-gpu", + "resources": [ + {"resource": "h200", "reserved": 8, "limit": 16, "unit": "GPU", "missing_action": "reject"}, + {"resource": "CPU", "reserved": 128, "limit": 256, "unit": "CPU", "missing_action": "default", "default_value": 8}, + {"resource": "memory", "reserved": 1024, "limit": 2048, "unit": "GiB", "missing_action": "default", "default_value": 64} + ] + } + ] +} ``` -**Outcome:** Rejected. Missing policy row → reserved 0 for `tensorrt_sessions`; -non-preemptible cannot ask for a positive amount. Fix: add `tensorrt_sessions` to -the policy, or mark the step preemptible if borrowing is acceptable. +This higher-priority policy can coexist with the workspace default policy. A +request from a member of the production team may match both policies; Resource +Manager orders candidate plans by priority and rank. The team reservation is +the protected envelope for critical work. ---- - -## Hard rejections (no queue) - -### Request exceeds pool capacity - -**Pool** +### Access pattern: Sandbox project cap ```json { - "name": "small", - "capacity": { - "gpu": 8 - } + "pool": "org-shared-gpu-fleet", + "subject_selector": { + "subject_type": "organization", + "subject_id": "", + "contains": { + "subject_type": "workspace", + "subject_id": "", + "contains": { + "subject_type": "project", + "subject_id": "" + } + } + }, + "priority": 5, + "priority_lane": false, + "concurrency_limit": 4, + "grants": [ + { + "class": "tier-a-gpu", + "resources": [ + {"resource": "h200", "reserved": 0, "limit": 2, "unit": "GPU", "missing_action": "reject"}, + {"resource": "CPU", "reserved": 0, "limit": 32, "unit": "CPU", "missing_action": "default", "default_value": 2}, + {"resource": "memory", "reserved": 0, "limit": 256, "unit": "GiB", "missing_action": "default", "default_value": 16} + ] + } + ] } ``` -**Policy** +This policy prevents a sandbox project from becoming the reason production work +waits. It still lets users submit large sweeps; extra work queues and rolls in +as capacity frees up. + +### Access pattern: Project-specific promotion without changing inventory + +When a sandbox project becomes production-critical for a week, do not edit the +pool class. Add a temporary higher-priority policy for that project and remove +or time-limit it later. ```json { - "pool": "small", - "component": "k8s-orch", - "component_type": "orchestrator", - "priority": 10, - "reserved": { - "gpu": 8 + "pool": "org-shared-gpu-fleet", + "subject_selector": { + "subject_type": "organization", + "subject_id": "", + "contains": { + "subject_type": "workspace", + "subject_id": "", + "contains": { + "subject_type": "project", + "subject_id": "" + } + } }, - "limit": { - "gpu": 8 - } + "priority": 90, + "priority_lane": false, + "concurrency_limit": 6, + "grants": [ + { + "class": "tier-a-gpu", + "resources": [ + {"resource": "h200", "reserved": 2, "limit": 6, "unit": "GPU", "missing_action": "reject"}, + {"resource": "CPU", "reserved": 32, "limit": 96, "unit": "CPU", "missing_action": "default", "default_value": 8}, + {"resource": "memory", "reserved": 256, "limit": 768, "unit": "GiB", "missing_action": "default", "default_value": 64} + ] + } + ] } ``` -**Step** +Now a request from that project may match both the low-priority sandbox cap and +the temporary promotion. Resource Manager evaluates both routes. The higher +priority route gives the project a short-lived reservation and higher ceiling +without changing the broad default policy. -```python -@step( - settings={ - "resources": ResourceSettings( - gpu_count=10, - preemptible=True, - ) - } -) -def too_big_for_planet() -> None: - ... -``` +### Runtime behavior: overlapping subject scopes -**Outcome:** Rejected immediately. Ten exceeds the pool total for `gpu`; the -request does not join a queue. +A single request can carry subjects for account, team, workspace, project, +pipeline, component, and service connector. In this example, a production team +member running a pipeline in Workspace A may match the workspace policy and the +team policy. The team route ranks higher because its policy priority is 120, +so it wins when it can satisfy the request. ---- +If a sandbox project request asks for 4 H200 GPUs, the low-priority sandbox cap +rejects that route because its limit is 2. If the temporary promotion is active, +the same request can be admitted through the promotion because its limit is 6. +If both routes are impossible, the request queues only when Resource Manager +can identify a future route; otherwise it is rejected with the reason visible +on the request. + +## Example 8: Cloud Machine Types -### Request exceeds policy limit (pool could fit) +Cloud backends often expose machine shapes rather than simple resource pools. +A SageMaker, Vertex AI, Azure ML, or Modal-style backend may ask for an +instance type. Resource pool classes can model these shapes: the class name +represents the machine shape, and target settings carry the backend-specific +machine type. -**Pool** +The strategy is: + +* Use classes for machine families. +* Bundle the machine's GPU, CPU, and memory in the class. +* Use target settings to tell the stack component which instance type to use. +* Use rank to prefer reserved or cheaper machine types. + +### Inventory ```json { - "name": "shared", - "capacity": { - "gpu": 8 - } + "name": "aws-training-machines", + "description": "SageMaker training machine families", + "rank": 60, + "accounting_mode": "authoritative", + "target_bindings": [ + { + "target_selector": { + "subject_type": "organization", + "subject_id": "", + "contains": { + "subject_type": "workspace", + "subject_id": "", + "contains": { + "subject_type": "component", + "subject_id": "", + "attributes": { + "component_type": "step_operator" + } + } + } + } + } + ], + "classes": [ + { + "class": "p5-48xlarge", + "rank": 100, + "reclaimable": "never", + "concurrency_limit": 4, + "resources": [ + {"resource": "GPU", "quantity": 32, "unit": "GPU"}, + {"resource": "CPU", "quantity": 768, "unit": "CPU"}, + {"resource": "memory", "quantity": 8192, "unit": "GiB"} + ], + "target_settings": [ + { + "target_type": "component", + "settings": { + "instance_type": "ml.p5.48xlarge", + "volume_size_gb": 1024 + } + } + ], + "attributes": { + "provider": "aws", + "machine_type": "ml.p5.48xlarge", + "tier": "gpu-large" + } + }, + { + "class": "m7i-24xlarge", + "rank": 20, + "reclaimable": "never", + "concurrency_limit": 10, + "resources": [ + {"resource": "CPU", "quantity": 960, "unit": "CPU"}, + {"resource": "memory", "quantity": 3840, "unit": "GiB"} + ], + "target_settings": [ + { + "target_type": "component", + "settings": { + "instance_type": "ml.m7i.24xlarge", + "volume_size_gb": 512 + } + } + ], + "attributes": { + "provider": "aws", + "machine_type": "ml.m7i.24xlarge", + "tier": "cpu-large" + } + } + ] } ``` -**Policy** +### Access pattern: GPU training pipeline ```json { - "pool": "shared", - "component": "team-a-orch", - "component_type": "orchestrator", - "priority": 10, - "reserved": { - "gpu": 2 + "pool": "aws-training-machines", + "subject_selector": { + "subject_type": "organization", + "subject_id": "", + "contains": { + "subject_type": "workspace", + "subject_id": "", + "contains": { + "subject_type": "project", + "subject_id": "", + "contains": { + "subject_type": "pipeline", + "subject_id": "" + } + } + } }, - "limit": { - "gpu": 4 - } + "priority": 80, + "priority_lane": false, + "concurrency_limit": 4, + "grants": [ + { + "class": "p5-48xlarge", + "resources": [ + {"resource": "GPU", "reserved": 8, "limit": 32, "unit": "GPU", "missing_action": "reject"}, + {"resource": "CPU", "reserved": 96, "limit": 768, "unit": "CPU", "missing_action": "default", "default_value": 24}, + {"resource": "memory", "reserved": 1024, "limit": 8192, "unit": "GiB", "missing_action": "default", "default_value": 128} + ] + } + ] } ``` -**Step** +### Access pattern: CPU preprocessing project -```python -@step( - settings={ - "resources": ResourceSettings( - gpu_count=6, - preemptible=True, - ) +```json +{ + "pool": "aws-training-machines", + "subject_selector": { + "subject_type": "organization", + "subject_id": "", + "contains": { + "subject_type": "workspace", + "subject_id": "", + "contains": { + "subject_type": "project", + "subject_id": "" + } } -) -def over_team_limit() -> None: - ... + }, + "priority": 30, + "priority_lane": false, + "concurrency_limit": 10, + "grants": [ + { + "class": "m7i-24xlarge", + "resources": [ + {"resource": "CPU", "reserved": 120, "limit": 720, "unit": "CPU", "missing_action": "reject"}, + {"resource": "memory", "reserved": 480, "limit": 2880, "unit": "GiB", "missing_action": "reject"} + ] + } + ] +} ``` -**Outcome:** Rejected. Six exceeds this component’s limit (4) for `gpu`, even if -eight GPUs exist in the pool. - ---- +The same pool can represent GPU and CPU machine families. Policies decide which +pipelines or projects can use each class. -## Contention: queues and priorities +### Access pattern: Give production first choice of expensive machines -### Two teams, same priority, not enough GPUs - -**Story:** Red and Blue orchestrators share one pool. Policies use the same -priority. Many preemptible steps each want 2 GPUs; the pool cannot satisfy -everyone at once. - -**Pool** +Use a higher-priority pipeline policy when a small number of production +pipelines should get the scarce, expensive shape before exploratory training. ```json { - "name": "shared", - "capacity": { - "gpu": 8 - } -} -``` - -**Policies** - -```json -[ - { - "pool": "shared", - "component": "red-orch", - "component_type": "orchestrator", - "priority": 10, - "reserved": { - "gpu": 4 - }, - "limit": { - "gpu": 8 + "pool": "aws-training-machines", + "subject_selector": { + "subject_type": "organization", + "subject_id": "", + "contains": { + "subject_type": "workspace", + "subject_id": "", + "contains": { + "subject_type": "project", + "subject_id": "", + "contains": { + "subject_type": "pipeline", + "subject_id": "" + } + } } }, - { - "pool": "shared", - "component": "blue-orch", - "component_type": "orchestrator", - "priority": 10, - "reserved": { - "gpu": 4 - }, - "limit": { - "gpu": 8 + "preemption_group": { + "subject_type": "organization", + "subject_id": "", + "contains": { + "subject_type": "workspace", + "subject_id": "", + "contains": { + "subject_type": "project", + "subject_id": "", + "contains": { + "subject_type": "pipeline", + "subject_id": "" + } + } } - } -] + }, + "priority": 140, + "priority_lane": false, + "concurrency_limit": 2, + "grants": [ + { + "class": "p5-48xlarge", + "resources": [ + {"resource": "GPU", "reserved": 16, "limit": 32, "unit": "GPU", "missing_action": "reject"}, + {"resource": "CPU", "reserved": 192, "limit": 768, "unit": "CPU", "missing_action": "default", "default_value": 96}, + {"resource": "memory", "reserved": 2048, "limit": 8192, "unit": "GiB", "missing_action": "default", "default_value": 1024} + ] + } + ] +} ``` -**Step** (typical for either team) +The preemption group prevents two runs of the same production pipeline from +preempting each other. They queue behind one another when the pipeline-level +concurrency limit is full. -```python -@step( - settings={ - "resources": ResourceSettings( - gpu_count=2, - preemptible=True, - ) - } -) -def train() -> None: - ... -``` +### Runtime behavior: machine shape selection -**Outcome:** Requests wait in the pool queue until GPUs free up. Among the same -policy priority, ordering tends to favor older waiters (FIFO-style). The -allocator also prefers a request that still fits entirely in its unused reserved -slice over one that must borrow when both are waiting—so a team with reservation -headroom is not stuck behind another team that is already bursting, if the next -grant can be served from that reserved slice. No preemption until a -higher-priority waiter or reclaim logic forces it. +A preprocessing step with only `cpu_count` and `memory` demands is planned +against `m7i-24xlarge` because the GPU class is not needed and the CPU policy +matches the project. A training step with `gpu_count=8` cannot run on the CPU +class and must match a policy that grants `p5-48xlarge`. ---- +If exploratory training asks for 16 GPUs while the policy limit is 8, that +route is rejected even if the pool has idle P5 capacity. Increase the policy +limit to allow that subject to use more of the existing inventory. -### Higher priority wins; lower may be preempted +## Example 9: Training License With Traceability -**Story:** Sandbox bursts with preemptible work. Production has higher policy -priority and needs GPUs when the pool is full. +Not all scarce resources are hardware. A vendor might sell four floating +training licenses, while CPU and memory come from a large shared backend. The +license is the resource that requires enforcement. CPU and memory are still +worth recording because operators want to understand what each licensed job +consumed. -**Pool** +The strategy is: + +* Create a custom `training-license` descriptor. +* Put finite license seats in the class. +* Put CPU and memory in the same class with unlimited quantities for + traceability. +* Use a pool or grant concurrency limit when the real rule is "four jobs at + once." + +### Descriptor ```json { - "name": "shared", - "capacity": { - "gpu": 8 - } + "name": "training-license", + "kind": "license", + "description": "Floating training license", + "attributes": { + "vendor": "acme", + "contract": "enterprise" + }, + "units": [ + {"name": "seat", "multiplier": 1} + ] } ``` -**Policies** +### Inventory ```json -[ - { - "pool": "shared", - "component": "sandbox-orch", - "component_type": "orchestrator", - "priority": 10, - "reserved": { - "gpu": 4 - }, - "limit": { - "gpu": 8 +{ + "name": "licensed-training", + "description": "Training jobs that require a floating vendor license", + "rank": 50, + "accounting_mode": "authoritative", + "concurrency_limit": 4, + "target_bindings": [ + { + "target_selector": { + "subject_type": "organization", + "subject_id": "", + "contains": { + "subject_type": "workspace", + "subject_id": "", + "contains": { + "subject_type": "component", + "subject_id": "", + "attributes": { + "component_type": "orchestrator" + } + } + } + } } - }, - { - "pool": "shared", - "component": "prod-orch", - "component_type": "orchestrator", - "priority": 100, - "reserved": { - "gpu": 2 - }, - "limit": { - "gpu": 8 + ], + "classes": [ + { + "class": "license-seat", + "rank": 100, + "reclaimable": "never", + "resources": [ + {"resource": "training-license", "quantity": 4, "unit": "seat"}, + {"resource": "CPU", "quantity": null}, + {"resource": "memory", "quantity": null} + ], + "attributes": { + "license": "acme-training" + } } - } -] + ] +} ``` -**Sandbox** (already holding six GPUs, preemptible) +### Access pattern: Project license access -```python -@step( - settings={ - "resources": ResourceSettings( - gpu_count=6, - preemptible=True, - ) +```json +{ + "pool": "licensed-training", + "subject_selector": { + "subject_type": "organization", + "subject_id": "", + "contains": { + "subject_type": "workspace", + "subject_id": "", + "contains": { + "subject_type": "project", + "subject_id": "" + } } -) -def sandbox_experiment() -> None: - ... + }, + "priority": 50, + "priority_lane": false, + "concurrency_limit": 4, + "grants": [ + { + "class": "license-seat", + "resources": [ + {"resource": "training-license", "reserved": 1, "limit": 4, "unit": "seat", "missing_action": "reject"}, + {"resource": "CPU", "reserved": 0, "limit": null, "unit": "CPU", "missing_action": "omit"}, + {"resource": "memory", "reserved": 0, "limit": null, "unit": "GiB", "missing_action": "omit"} + ] + } + ] +} ``` -**Prod** (new, preemptible, needs four) +### Runtime behavior: custom license demand ```python -@step( - settings={ - "resources": ResourceSettings( - gpu_count=4, - preemptible=True, +from zenml.models import ResourceRequestDemand + +ResourceSettings( + resources=[ + ResourceRequestDemand( + resource="training-license", + quantity=1, + unit="seat", + class_name="license-seat", ) - } + ], + cpu_count=8, + memory="32GiB", ) -def prod_train() -> None: - ... ``` -**Outcome:** If four GPUs cannot be granted without reclaiming space, the -reconciler may preempt Sandbox’s preemptible runs (lower policy priority) so Prod -can proceed. See [How preemption works](resource-pools-reconciliation.md#how-preemption-works) for victim -ordering. - ---- - -### Production non-preemptible waits on reserved only +The request explicitly asks for the custom license. CPU and memory are also +recorded because the author set them through `ResourceSettings`. The pool class +does not limit CPU and memory globally, but request inspection still shows the +full footprint of licensed work. -**Story:** Prod uses `preemptible=False` and asks only for what is reserved. If -another non-preemptible job on the same stack component already holds the reserved -GPUs, this step does not borrow from Sandbox’s burst. +### Access pattern: Split license seats between teams -**Pool** +If two teams share four floating seats, model the seats once and express the +split through policy grants. ```json { - "name": "shared", - "capacity": { - "gpu": 8 - } + "pool": "licensed-training", + "subject_selector": { + "subject_type": "organization", + "subject_id": "", + "contains": { + "subject_type": "team", + "subject_id": "" + } + }, + "priority": 60, + "priority_lane": false, + "concurrency_limit": 3, + "grants": [ + { + "class": "license-seat", + "resources": [ + {"resource": "training-license", "reserved": 2, "limit": 3, "unit": "seat", "missing_action": "reject"}, + {"resource": "CPU", "reserved": 0, "limit": null, "unit": "CPU", "missing_action": "omit"}, + {"resource": "memory", "reserved": 0, "limit": null, "unit": "GiB", "missing_action": "omit"} + ] + } + ] } ``` -**Policies** (same as previous example: Sandbox priority 10, Prod priority 100) - -**Prod step** - -```python -@step( - settings={ - "resources": ResourceSettings( - gpu_count=2, - preemptible=False, - ) +```json +{ + "pool": "licensed-training", + "subject_selector": { + "subject_type": "organization", + "subject_id": "", + "contains": { + "subject_type": "team", + "subject_id": "" } -) -def prod_sla_job() -> None: - ... + }, + "priority": 60, + "priority_lane": false, + "concurrency_limit": 3, + "grants": [ + { + "class": "license-seat", + "resources": [ + {"resource": "training-license", "reserved": 2, "limit": 3, "unit": "seat", "missing_action": "reject"}, + {"resource": "CPU", "reserved": 0, "limit": null, "unit": "CPU", "missing_action": "omit"}, + {"resource": "memory", "reserved": 0, "limit": null, "unit": "GiB", "missing_action": "omit"} + ] + } + ] +} ``` -**Outcome:** Waits in the queue if Prod’s reserved `gpu` (2) is already used by -other non-preemptible work on `prod-orch`. It will not take Sandbox’s borrowed -GPUs. Ways out: raise reserved for Prod, wait for the other job to finish, or -use preemptible Prod work if policy allows. +Each team has two protected seats and can borrow a third seat while idle. The +sum of reservations equals the four-seat license contract. The sum of limits is +larger than the pool quantity by design because limits describe the maximum a +policy may use when capacity is idle, not extra physical seats. ---- +### Runtime behavior: traceability without CPU enforcement + +If the forecasting team has three licensed jobs running, its fourth job waits +or is rejected by the policy limit even if CPU workers are free. If the ranking +team is idle, the forecasting team may borrow one extra seat up to its limit of +three. CPU and memory demands appear on the request and allocation records, but +because the class quantities are unlimited and the grant limits are null, CPU +and memory do not block admission in this example. -## Multiple pools and multi-key requests +## Example 10: External Inference Priority Lane -{% hint style="warning" %} -Several policies on the same stack component mean several pools may try -to satisfy the same resource request, but only one pool can win. -Every key in the request must pass that pool’s checks; ZenML does not take -`gpu` from one pool and `mcpu` from another for one step. -{% endhint %} +This example models a shared cluster where ZenML training runs overnight on GPU +nodes that an inference system needs during traffic peaks. The inference system +does not launch through ZenML, but it can still coordinate with Resource +Manager as a trusted external workload. -### Two pools on one orchestrator—primary pool wins +The strategy is: -**Story:** You attach two policies to the same orchestrator pointing at -different pools. The step still produces one resource request, enqueued in every -eligible pool; only one pool may win. +* Use a service account as the policy subject. +* Include a service connector subject in direct requests so pool target + bindings can match the infrastructure target. +* Use `priority_lane: true` for the external inference policy. +* Use `grants: []` when the external system owns the whole matching class + bundle during its claim. -**Pools** +### Access pattern: Priority-lane external claim ```json -[ - { - "name": "eu-west-gpu", - "capacity": { - "gpu": 16 +{ + "pool": "k8s-prod-eu-gpu", + "subject_selector": { + "subject_type": "organization", + "subject_id": "", + "contains": { + "subject_type": "account", + "subject_id": "", + "attributes": { + "is_service_account": true + } } }, - { - "name": "eu-north-gpu", - "capacity": { - "gpu": 16 + "preemption_group": { + "subject_type": "organization", + "subject_id": "", + "contains": { + "subject_type": "account", + "subject_id": "", + "attributes": { + "is_service_account": true + } } - } -] + }, + "priority_lane": true, + "concurrency_limit": 8, + "grants": [] +} ``` -**Policies** +Grantless access is deliberate here. The service account is trusted to claim +real infrastructure usage, and the policy should not restrict it to a small +reserved slice. In authoritative mode, the request still has to fit the class +resources and concurrency limits. + +### Runtime behavior: direct external request ```json -[ - { - "pool": "eu-west-gpu", - "component": "regional-orch", - "component_type": "orchestrator", - "priority": 20, - "reserved": { - "gpu": 8 - }, - "limit": { - "gpu": 16 - } - }, - { - "pool": "eu-north-gpu", - "component": "regional-orch", - "component_type": "orchestrator", - "priority": 10, - "reserved": { - "gpu": 8 +{ + "subjects": [ + { + "subject_type": "organization", + "subject_id": "", + "child": { + "subject_type": "account", + "subject_id": "", + "attributes": { + "is_service_account": true, + "name": "inference-capacity-controller" + } + } }, - "limit": { - "gpu": 16 + { + "subject_type": "organization", + "subject_id": "", + "child": { + "subject_type": "workspace", + "subject_id": "", + "child": { + "subject_type": "service_connector", + "subject_id": "", + "attributes": { + "connector_type": "gcp", + "effective_resource_type": "kubernetes-cluster", + "effective_resource_id": "prod-eu" + } + } + } } + ], + "pool": "k8s-prod-eu-gpu", + "demands": [ + {"resource": "h200", "quantity": 4, "unit": "GPU", "class": "h200-reserved"}, + {"resource": "CPU", "quantity": 48, "unit": "CPU", "class": "h200-reserved"}, + {"resource": "memory", "quantity": 384, "unit": "GiB", "class": "h200-reserved"} + ], + "reclaim_tolerance": "none", + "lease_expires_at": "2026-07-23T18:30:00Z", + "metadata": { + "service": "recommendations", + "traffic_window": "peak" } -] +} ``` -**Step** +The service account subject matches the policy. The service connector subject +matches the pool target. If lower-priority ZenML work is using reclaimable +capacity, the priority-lane request can cause Resource Manager to preempt it +and return capacity to the external service. -```python -@step( - settings={ - "resources": ResourceSettings( - gpu_count=1, - preemptible=True, - ) +### Access pattern: Soft external reservation without priority lane + +Not every external workload needs emergency power. A batch scoring service can +use a normal policy with a reservation and limit. + +```json +{ + "pool": "k8s-prod-eu-gpu", + "subject_selector": { + "subject_type": "organization", + "subject_id": "", + "contains": { + "subject_type": "account", + "subject_id": "", + "attributes": { + "is_service_account": true + } } -) -def train() -> None: - ... + }, + "priority": 45, + "priority_lane": false, + "concurrency_limit": 4, + "grants": [ + { + "class": "a10-burst", + "resources": [ + {"resource": "GPU", "reserved": 4, "limit": 12, "unit": "GPU", "missing_action": "reject"}, + {"resource": "CPU", "reserved": 0, "limit": null, "unit": "CPU", "missing_action": "default", "default_value": 4}, + {"resource": "memory", "reserved": 0, "limit": null, "unit": "GiB", "missing_action": "default", "default_value": 24} + ] + } + ] +} ``` -**Outcome:** The server tries higher policy priority first—eu-west before -eu-north. Whichever pool grants first owns the allocation; the other queue entry -is dropped as stale. Use this for primary/fallback or regional capacity, not for -splitting one step across unrelated quotas. +This policy lets the external service coordinate fairly with ZenML work. It +does not jump to the maximum possible priority, so higher-priority production +training can still outrank it. Use priority lane only for systems that must +reclaim capacity immediately. ---- +### Runtime behavior: external claim during a traffic spike + +At night, research sweeps may hold reclaimable A10 capacity. At 09:00, the +inference controller creates the priority-lane request for 4 H200 GPUs with +`reclaim_tolerance="none"`. Resource Manager compares it against active +allocations. Lower-priority reclaimable allocations that are not protected by +the same preemption group can be selected as victims. Allocations with +`reclaim_tolerance="none"` are not selected as victims. -### One step must satisfy every key in each pool +If the direct request omits the service connector subject, it may match the +policy subject but fail the pool target binding. External clients should send +both identity subjects and target subjects. -**Story:** Eligibility is checked per pool against all keys on the request. If a -pool lacks a key the step needs, that pool treats it as zero capacity—the request -is not eligible there. +## Example 11: Governance With Kueue or Run:ai -**Pool A** (GPUs only) +Some environments already use an infrastructure scheduler such as Kueue or +Run:ai for hard enforcement. In that case, Resource Manager can route work, +apply target settings, and provide a consistent ZenML policy layer while the +infrastructure scheduler remains responsible for runtime allocation. + +The strategy is: + +* Set `accounting_mode: "governance"`. +* Model classes as the queues or workload classes known to the infrastructure + scheduler. +* Put Kueue, Run:ai, or scheduler-specific settings in component target + settings. +* Use policies to decide which ZenML subjects are allowed to target those + infrastructure queues. + +### Inventory ```json { - "name": "gpu-only", - "capacity": { - "gpu": 8 - } + "name": "kueue-governed-gpu", + "description": "GPU workloads governed by Kubernetes Kueue", + "rank": 80, + "accounting_mode": "governance", + "target_bindings": [ + { + "target_selector": { + "subject_type": "organization", + "subject_id": "", + "contains": { + "subject_type": "workspace", + "subject_id": "", + "contains": { + "subject_type": "component", + "subject_id": "", + "attributes": { + "component_type": "orchestrator" + } + } + } + } + } + ], + "classes": [ + { + "class": "kueue-gpu-training", + "rank": 100, + "reclaimable": "coordinated", + "resources": [ + {"resource": "GPU", "quantity": 32, "unit": "GPU"}, + {"resource": "CPU", "quantity": null}, + {"resource": "memory", "quantity": null} + ], + "target_settings": [ + { + "target_type": "component", + "settings": { + "pod_settings": { + "labels": { + "kueue.x-k8s.io/queue-name": "gpu-training" + } + } + } + } + ], + "attributes": { + "scheduler": "kueue", + "queue": "gpu-training" + } + } + ] } ``` -**Pool B** (full bundle) +### Access pattern: Research queue access ```json { - "name": "gpu-and-trt", - "capacity": { - "gpu": 8, - "tensorrt_sessions": 4 - } + "pool": "kueue-governed-gpu", + "subject_selector": { + "subject_type": "organization", + "subject_id": "", + "contains": { + "subject_type": "team", + "subject_id": "" + } + }, + "priority": 40, + "priority_lane": false, + "grants": [ + { + "class": "kueue-gpu-training", + "resources": [ + {"resource": "GPU", "reserved": 0, "limit": 8, "unit": "GPU", "missing_action": "reject"}, + {"resource": "CPU", "reserved": 0, "limit": null, "unit": "CPU", "missing_action": "default", "default_value": 4}, + {"resource": "memory", "reserved": 0, "limit": null, "unit": "GiB", "missing_action": "default", "default_value": 24} + ] + } + ] } ``` -**Policy** (example: only Pool B is attached, or imagine Pool A attached alone) +In governance mode, use this as a routing and policy-control pattern. Runtime +capacity enforcement, queue fairness, and preemption are expected to happen in +the infrastructure scheduler. -**Step** +### Access pattern: Separate Kueue queues by subject -```python -@step( - settings={ - "resources": ResourceSettings( - gpu_count=1, - pool_resources={"tensorrt_sessions": 1}, - preemptible=True, - ) +Because the class target settings can name a Kueue queue, the same inventory +style can expose multiple infrastructure queues as classes. + +```json +{ + "pool": "kueue-governed-gpu", + "subject_selector": { + "subject_type": "organization", + "subject_id": "", + "contains": { + "subject_type": "team", + "subject_id": "" } -) -def infer() -> None: - ... + }, + "priority": 100, + "priority_lane": false, + "concurrency_limit": 12, + "grants": [ + { + "class": "kueue-gpu-training", + "resources": [ + {"resource": "GPU", "reserved": 0, "limit": 16, "unit": "GPU", "missing_action": "reject"}, + {"resource": "CPU", "reserved": 0, "limit": null, "unit": "CPU", "missing_action": "default", "default_value": 8}, + {"resource": "memory", "reserved": 0, "limit": null, "unit": "GiB", "missing_action": "default", "default_value": 64} + ], + "target_settings": [ + { + "target_type": "component", + "settings": { + "pod_settings": { + "labels": { + "kueue.x-k8s.io/queue-name": "gpu-production" + } + } + } + } + ] + } + ] +} ``` -**Outcome:** A pool with only `gpu` cannot satisfy `tensorrt_sessions`—that -dimension is zero there, so the request does not enqueue on that pool. +The policy-level target settings override or supplement the general class +settings for this subject. Resource Manager still records the request, grants, +and selected class. Kueue decides when pods actually run and how preemption is +enforced inside Kubernetes. -**Lesson:** Model every scarce bounded dimension you care about on one pool (or -ensure every candidate pool defines the same key set for those keys). The next -section shows how `mcpu` and `memory_mb` differ: omitting them on one pool keeps -that path eligible even when another pool meters them strictly. +### Runtime behavior: governance is not authoritative allocation ---- +In authoritative mode, Resource Manager accounts for the pool quantity and +preempts lower-priority reclaimable work when needed. In governance mode, the +policy still controls whether a ZenML subject may target a class, but the +infrastructure scheduler is the source of truth for actual capacity, fairness, +and pod preemption. Use governance when your Kubernetes scheduler already owns +those decisions and ZenML should provide policy, routing, and observability. -### Two pools: higher-priority path meters CPU/RAM; GPU-only path still wins +## Example 12: Policy Cookbook -**Story:** One orchestrator has two policies (same pattern as *Two pools on one -orchestrator—primary pool wins*). Pool B’s capacity and policy include `mcpu` and -`memory_mb`, with reserved amounts sized for small non-preemptible jobs. Pool A -only defines `gpu`; it does not list `mcpu` or `memory_mb`, so those dimensions -are unbounded at the pool layer and its policy does not reserve them. A -non-preemptible step asks for one GPU but more CPU and RAM than Pool B’s policy -allows. The higher-priority policy (Pool B) cannot grant that request; the -lower-priority policy (Pool A) can, because the request’s CPU and memory demand is -not quota’d on that path. The allocation is owned by Pool A. +The remaining examples reuse inventory from earlier sections and focus on +policy language. Once a pool models the infrastructure, policies define +priority, access, reservations, limits, and operational guardrails. -**Pool A** (GPUs only—no `mcpu` or `memory_mb` on the pool) +### Give one team a reserved share and a burst ceiling + +Use this when a team pays for a baseline but may borrow idle capacity. ```json { - "name": "gpu-only-fallback", - "capacity": { - "gpu": 8 - } + "pool": "org-shared-gpu-fleet", + "subject_selector": { + "subject_type": "organization", + "subject_id": "", + "contains": { + "subject_type": "team", + "subject_id": "" + } + }, + "priority": 70, + "priority_lane": false, + "grants": [ + { + "class": "tier-a-gpu", + "resources": [ + {"resource": "h200", "reserved": 4, "limit": 10, "unit": "GPU", "missing_action": "reject"}, + {"resource": "CPU", "reserved": 64, "limit": 160, "unit": "CPU", "missing_action": "default", "default_value": 8}, + {"resource": "memory", "reserved": 512, "limit": 1280, "unit": "GiB", "missing_action": "default", "default_value": 64} + ] + } + ] } ``` -**Pool B** (GPUs plus metered CPU and memory) +The reservation protects the team from starvation. The limit prevents the team +from taking the entire class when everyone else is also busy. + +### Make production outrank experiments + +Use two policies against the same class. The production policy has higher +priority and reserved share. The experiment policy has lower priority and no +reservation. ```json { - "name": "metered-cpu-mem", - "capacity": { - "gpu": 8, - "mcpu": 128000, - "memory_mb": 524288 - } + "pool": "k8s-prod-eu-gpu", + "subject_selector": { + "subject_type": "organization", + "subject_id": "", + "contains": { + "subject_type": "workspace", + "subject_id": "", + "contains": { + "subject_type": "project", + "subject_id": "", + "contains": { + "subject_type": "pipeline", + "subject_id": "" + } + } + } + }, + "priority": 150, + "priority_lane": false, + "grants": [ + { + "class": "h200-reserved", + "resources": [ + {"resource": "h200", "reserved": 6, "limit": 8, "unit": "GPU", "missing_action": "reject"}, + {"resource": "CPU", "reserved": 96, "limit": 128, "unit": "CPU", "missing_action": "default", "default_value": 8}, + {"resource": "memory", "reserved": 768, "limit": 1024, "unit": "GiB", "missing_action": "default", "default_value": 64} + ] + } + ] +} +``` + +```json +{ + "pool": "k8s-prod-eu-gpu", + "subject_selector": { + "subject_type": "organization", + "subject_id": "", + "contains": { + "subject_type": "workspace", + "subject_id": "", + "contains": { + "subject_type": "project", + "subject_id": "" + } + } + }, + "priority": 10, + "priority_lane": false, + "grants": [ + { + "class": "h200-reserved", + "resources": [ + {"resource": "h200", "reserved": 0, "limit": 2, "unit": "GPU", "missing_action": "reject"}, + {"resource": "CPU", "reserved": 0, "limit": 32, "unit": "CPU", "missing_action": "default", "default_value": 4}, + {"resource": "memory", "reserved": 0, "limit": 256, "unit": "GiB", "missing_action": "default", "default_value": 32} + ] + } + ] } ``` -**Policies** (same component, different priorities—B is preferred when both can -grant) +If experiments use `reclaim_tolerance="coordinated"` or `"any"`, production can +preempt them when needed. If production uses `reclaim_tolerance="none"`, it +waits for reserved share instead of depending on reclaimable capacity. + +### Put a hard cap on one user or service account + +Use this when a single automation account launches many jobs and should never +occupy more than a small slice. ```json -[ - { - "pool": "metered-cpu-mem", - "component": "k8s-orch", - "component_type": "orchestrator", - "priority": 100, - "reserved": { - "gpu": 4, - "mcpu": 4000, - "memory_mb": 8192 - }, - "limit": { - "gpu": 8, - "mcpu": 64000, - "memory_mb": 131072 +{ + "pool": "workspace-standard-compute", + "subject_selector": { + "subject_type": "account", + "subject_id": "", + "attributes": { + "is_service_account": false } }, - { - "pool": "gpu-only-fallback", - "component": "k8s-orch", - "component_type": "orchestrator", - "priority": 50, - "reserved": { - "gpu": 4 - }, - "limit": { - "gpu": 8 + "priority": 15, + "priority_lane": false, + "concurrency_limit": 2, + "grants": [ + { + "class": "standard-workers", + "resources": [ + {"resource": "CPU", "reserved": 0, "limit": 16, "unit": "CPU", "missing_action": "reject"}, + {"resource": "memory", "reserved": 0, "limit": 64, "unit": "GiB", "missing_action": "reject"} + ] } - } -] + ] +} ``` -**Step** +The concurrency limit is the most important field here. Even if the user's +jobs request only 1 CPU each, no more than two active requests can run through +this policy. -```python -@step( - settings={ - "resources": ResourceSettings( - gpu_count=1, - cpu_count=8, - memory="32GiB", - preemptible=False, - ) +### Use a scheduled policy window + +Use this when burst access is allowed only overnight. The policy selector is an +expression because the time window is attached to the policy match. + +```json +{ + "pool": "k8s-prod-eu-gpu", + "subject_selector": { + "any": [ + { + "subject_type": "organization", + "subject_id": "", + "contains": { + "subject_type": "workspace", + "subject_id": "", + "contains": { + "subject_type": "project", + "subject_id": "" + } + } + } + ], + "time_selector": { + "schedule": { + "timezone": "Europe/Berlin", + "start_time": "20:00:00", + "end_time": "08:00:00", + "days_of_week": ["mon", "tue", "wed", "thu", "fri"], + "days_of_month": [], + "weeks_of_month": [], + "months": [] + } } -) -def train() -> None: - ... + }, + "priority": 25, + "priority_lane": false, + "grants": [ + { + "class": "a10-burst", + "resources": [ + {"resource": "GPU", "reserved": 0, "limit": 16, "unit": "GPU", "missing_action": "reject"}, + {"resource": "CPU", "reserved": 0, "limit": null, "unit": "CPU", "missing_action": "default", "default_value": 4}, + {"resource": "memory", "reserved": 0, "limit": null, "unit": "GiB", "missing_action": "default", "default_value": 24} + ] + } + ] +} ``` -**Outcome:** The request maps to roughly `mcpu` 8000 and tens of thousands of -`memory_mb` for `32GiB`. For non-preemptible work, each key must be ≤ policy -reserved on the path you use. Pool B’s policy reserves only `mcpu` 4000 and -`memory_mb` 8192, so that path cannot satisfy the step. Pool A’s policy has no -`mcpu` or `memory_mb` rows; with those keys absent from the pool, they are not -treated as zero capacity, so the step remains eligible there on `gpu` alone. The -reconciler allocates from Pool A and drops the competing queue row for Pool B. -If you want large non-preemptible jobs to stay on the metered pool, raise -reserved (and capacity) on Pool B for `mcpu` and `memory_mb`, or reduce demand in -`ResourceSettings`—otherwise the GPU-only policy acts as an escape hatch for -heavy CPU/RAM asks. +This pattern keeps daytime capacity available for interactive and production +work while still letting nightly sweeps fill the cluster. ---- +### Admit a trusted system to the whole pool + +Use grantless policies only for trusted subjects. They expose every matching +class resource to the subject. + +```json +{ + "pool": "team-h200-machines", + "subject_selector": { + "subject_type": "organization", + "subject_id": "", + "contains": { + "subject_type": "account", + "subject_id": "", + "attributes": { + "is_service_account": true + } + } + }, + "priority_lane": true, + "concurrency_limit": 8, + "grants": [] +} +``` + +The service account gets the maximum priority and access to every matching +class. It still cannot allocate resources that the pool class does not contain, +and it still respects pool, class, and policy concurrency limits. ## See also -* [Resource pools](resource-pools.md) — overview -* [Core concepts](resource-pools-core-concepts.md) — pools, policies, requests -* [How preemption works](resource-pools-reconciliation.md#how-preemption-works) — preemption - ordering -* [Step configuration](https://docs.zenml.io/how-to/steps-pipelines/configuration) - — full `ResourceSettings` reference +* [Admin guide](resource-pools-admin-guide.md) - UI setup flow and exact API + shape. +* [Core concepts](resource-pools-core-concepts.md) - descriptors, classes, + subjects, policies, and requests. +* [User guide](resource-pools-user-guide.md) - author settings and request + inspection. +* [External workloads](resource-pools-external-workloads.md) - direct request + details. +* [Reconciliation process](resource-pools-reconciliation.md) - runtime + behavior. diff --git a/docs/book/getting-started/zenml-pro/resource-pools-external-workloads.md b/docs/book/getting-started/zenml-pro/resource-pools-external-workloads.md new file mode 100644 index 00000000000..4b8b439849d --- /dev/null +++ b/docs/book/getting-started/zenml-pro/resource-pools-external-workloads.md @@ -0,0 +1,276 @@ +--- +description: >- + How external workloads use ZenML Pro resource pools through service accounts, + direct Resource Manager requests, and priority-lane policies. +--- +# External workloads + +External workloads are jobs that consume the same infrastructure as ZenML +pipelines but are not launched as ZenML steps. Examples include online +inference services, batch jobs managed by another scheduler, or emergency +capacity reservations driven by an internal platform service. + +Resource pools can govern those workloads when they create Resource Manager +resource requests with the same subject and demand conventions as ZenML steps. + +## Common pattern + +1. Create a ZenML Pro service account for the external system. +2. Configure a resource pool in the UI with classes that represent the real + infrastructure bundle the workload consumes. +3. Attach a policy for the service account. Use priority lane only for + workloads that must outrank normal pipeline work. +4. The external service creates, renews, and releases Resource Manager resource + requests while it holds capacity. Include a service-account subject for the + policy and a target subject, such as a service connector, for the pool + target binding. + +## Policy for an external service + +For an inference service that needs to reclaim an H200 node pool during peak +traffic, use a service-account policy. A grantless policy is often appropriate +when the service should access the full matching class bundle. + +```json +{ + "pool": "prod-eu-gpu", + "subject_selector": { + "any": [ + { + "subject_type": "organization", + "subject_id": "", + "contains": { + "subject_type": "account", + "subject_id": "", + "attributes": { + "is_service_account": true + } + } + } + ] + }, + "preemption_group": { + "subject_type": "organization", + "subject_id": "", + "contains": { + "subject_type": "account", + "subject_id": "", + "attributes": { + "is_service_account": true + } + } + }, + "priority_lane": true, + "concurrency_limit": 8, + "grants": [] +} +``` + +`priority_lane: true` gives the policy the maximum Resource Manager priority. +In authoritative mode, it can reclaim lower-priority work that opted into +reclaim. It still respects pool target bindings, class matching, class +capacity, and concurrency limits. + +Use grant-based policies instead when the external service should receive only +a slice of the pool: + +```json +{ + "pool": "prod-eu-gpu", + "subject_selector": { + "subject_type": "organization", + "subject_id": "", + "contains": { + "subject_type": "account", + "subject_id": "", + "attributes": { + "is_service_account": true + } + } + }, + "priority": 200, + "priority_lane": false, + "grants": [ + { + "class": "h200-reserved", + "resources": [ + { + "resource": "h200", + "reserved": 2, + "limit": 4, + "unit": "GPU", + "missing_action": "reject" + }, + { + "resource": "CPU", + "reserved": 16, + "limit": 64, + "unit": "CPU", + "missing_action": "default", + "default_value": 8 + }, + { + "resource": "memory", + "reserved": 128, + "limit": 512, + "unit": "GiB", + "missing_action": "default", + "default_value": 64 + } + ] + } + ] +} +``` + +## Direct request shape + +Direct callers create runtime requests with subjects and demands. The service +account subject should follow the same root-first convention used by ZenML: +`organization -> account`. + +```json +{ + "subjects": [ + { + "subject_type": "organization", + "subject_id": "", + "child": { + "subject_type": "account", + "subject_id": "", + "attributes": { + "is_service_account": true, + "name": "inference-capacity-controller" + } + } + }, + { + "subject_type": "organization", + "subject_id": "", + "child": { + "subject_type": "workspace", + "subject_id": "", + "child": { + "subject_type": "service_connector", + "subject_id": "", + "attributes": { + "connector_type": "gcp", + "effective_resource_type": "kubernetes-cluster", + "effective_resource_id": "prod-eu" + } + } + } + } + ], + "pool": "prod-eu-gpu", + "demands": [ + { + "resource": "h200", + "quantity": 4, + "unit": "GPU", + "class": "h200-reserved" + }, + { + "resource": "CPU", + "quantity": 48, + "unit": "CPU", + "class": "h200-reserved" + }, + { + "resource": "memory", + "quantity": 384, + "unit": "GiB", + "class": "h200-reserved" + } + ], + "reclaim_tolerance": "none", + "lease_expires_at": "2026-07-23T18:30:00Z", + "metadata": { + "workload": "online-inference", + "service": "recommendations" + } +} +``` + +If the request is admitted, Resource Manager returns an allocated or pending +request. The external service must renew the lease while it holds capacity and +release the request when the workload no longer needs it. + +## Pool selectors + +External systems can use `pool` for an exact pool name or `pool_selector` to +target pools by attributes. Pool selectors are useful when you model one pool +per node, cluster, region, or machine-type family. + +```json +{ + "pool_selector": { + "all": [ + {"equals": {"cluster": "prod-eu"}}, + {"equals": {"node_pool": "h200-reserved"}} + ] + }, + "demands": [ + {"kind": "gpu", "quantity": 2, "class": "h200-reserved"} + ] +} +``` + +## Reclaim behavior + +External workloads typically use one of two patterns: + +| Pattern | Policy | Request reclaim tolerance | +| --- | --- | --- | +| Critical service reclaiming shared hardware | Priority lane, often grantless | `none` | +| Opportunistic external batch work | Normal priority, grant-based limits | `coordinated` or `any` | + +In authoritative pools, priority-lane work can preempt lower-priority requests +only when those requests are eligible for reclaim. It does not preempt other +priority-lane work at the same priority. + +In governance pools, Resource Manager records governance decisions and target +settings, but allocation and preemption are expected to happen in the external +infrastructure scheduler. + +## Target settings + +External workloads may rely on pool or class target settings for service +connector selection. The UI currently exposes service connector settings at the +pool and class levels: + +```json +{ + "target_type": "service_connector", + "settings": { + "connector_id": "", + "resource_type": "kubernetes-cluster", + "resource_id": "prod-eu" + } +} +``` + +Component target settings are also available at pool, class, policy, and grant +levels when ZenML will launch work through a stack component. + +## Operational tips + +* Use service accounts for external systems, not personal user accounts. +* Keep direct request subjects scoped exactly like ZenML subjects so policies + match predictably. +* Renew leases before they expire; otherwise Resource Manager returns capacity + to the pool. +* Release requests as soon as the external workload is done. +* Reserve priority lane for workloads that genuinely need maximum priority. +* Include CPU and memory in class bundles when the external workload consumes + them, even if GPU is the scarce resource. + +## See also + +* [Admin guide](resource-pools-admin-guide.md) - configure service-account + policies. +* [Core concepts](resource-pools-core-concepts.md) - selector and request + conventions. +* [Reconciliation process](resource-pools-reconciliation.md) - priority, + preemption, and leases. +* [Examples](resource-pools-examples.md) - external inference example. diff --git a/docs/book/getting-started/zenml-pro/resource-pools-reconciliation.md b/docs/book/getting-started/zenml-pro/resource-pools-reconciliation.md index c44e22216bf..79b7f7c9599 100644 --- a/docs/book/getting-started/zenml-pro/resource-pools-reconciliation.md +++ b/docs/book/getting-started/zenml-pro/resource-pools-reconciliation.md @@ -1,111 +1,212 @@ --- description: >- - How the resource pool reconciliation process works in ZenML Pro. + How resource pool reconciliation, queueing, accounting, preemption, and lease + management work at runtime. --- -# Resource Pool Reconciliation - -## Runtime flow (orchestration) - -1. Request creation: For eligible runs, the server derives requested - resources from the step’s `ResourceSettings` (see below), adds `step_run: 1`, - and stores `preemptible` from the same settings. The resource requester - is the stack’s step operator if the step uses one, otherwise the - orchestrator. -2. Queuing: If capacity is not available immediately, the step can remain - queued until the reconciler allocates it. -3. Client wait: The step launcher polls the resource request until it is allocated - (with backoff). If the request is not allocated, it is rejected, preempted, or cancelled, the client surfaces an error. When allocation succeeds, the step is published as running and execution proceeds. -4. Preemption: If the job at the front of the queue still cannot be granted, - the reconciler may stop other *preemptible* runs in that pool to free units - (see [How preemption works](#how-preemption-works)). Non-preemptible runs - are never stopped this way. They are also constrained so each request’s - per-key demand is **≤ policy reserved** for that key—even when **limit** is - higher—so they never rely on borrowed capacity that could clash with other - non-preemptible use on the same component. -5. Post-preemption retry: after preemption, if the step configuration allows -retries, the step goes back to the queue and is retried again. If the number of -retries is exhausted, the step fails. See [Automatic Step Retries](https://docs.zenml.io/how-to/steps-pipelines/advanced_features#automatic-step-retries) for more information. -6. Deallocation: When the step run completes, the resources are released back to the pool. If the step crashes unexpectedly, the resources are eventually released back to the pool. - -{% hint style="warning" %} -If the resource requester (orchestrator or step operator) has more than one -policy attached to more than one pool, the same logical request may appear in -several pool queues, but reconciliation still grants at most one active -allocation. Eligibility is computed separately per pool against the entire -resource request; the system never assigns part of a step’s demand to one pool -and the rest to another. -{% endhint %} - -## How preemption works - -**When.** Preemption runs only when the next queued request for a pool cannot be -allocated—there is not enough free capacity, or a policy rule blocks the -grant. The reconciler may then mark selected *already running* requests as -preempted, which cancels those step runs and returns their units to the pool. - -**Who can be stopped.** Only steps with `preemptible=True` (the default in -`ResourceSettings`) are candidates. `preemptible=False` means “never pick this -run as the one to kill.” - -**Who gets stopped first (simple picture).** - -1. Among preemptible runs in the same pool, **lower policy priority** is - considered before higher priority. If the waiting job’s priority is *strictly - higher* than a victim’s, that victim can be preempted to make room, as long - as freeing it actually fixes the shortage. -2. **Reserved** adds a second idea: *reclaim*. If the waiting component still - has unused **reserved** headroom on this pool (reserved minus what it is - already using here), the system may preempt preemptible runs that are using - **borrowed** capacity—even when those runs have the same or higher priority - than the waiter. Intuition: your reserved share is “yours to fill”; if - someone else is on the spare capacity you could have used under your - reservation, they can be moved out of the way. -3. **Limit** does not pick victims. It only caps how much a component may hold; - if the waiting request itself is over its own limit, killing other jobs will - not fix that—you need a higher limit or a smaller request. - -Victims are ordered by ascending policy priority, then by allocation time as a -tie-break. - -### Step-level: `preemptible` - -| Setting | Effect | +# Reconciliation process + +This page explains what happens after a dynamic step or external caller creates +a Resource Manager request. For definitions, see +[Core concepts](resource-pools-core-concepts.md). For author-facing settings, +see the [User guide](resource-pools-user-guide.md). + +Most of this page describes `authoritative` pools, where Resource Manager owns +allocation decisions. `advisory` and `governance` modes are covered separately +below. + +## Lifecycle + +1. **Request creation** - ZenML or an external caller sends subjects, demands, + reclaim tolerance, optional pool selectors, optional preemption group, and + lease settings. +2. **Admission discovery** - Resource Manager finds pools whose target bindings + match the request and policies whose subject selectors match the request. +3. **Planning** - For each matching policy, Resource Manager builds candidate + allocation plans against eligible classes and grants. Grantless policies + consider every matching class in the pool. +4. **Queueing** - If a plan is eligible but capacity is not currently free, the + request stays `pending` in the pool queue. +5. **Allocation** - When a plan fits capacity, limits, reservations, and + concurrency caps, Resource Manager marks the request `allocated` and records + allocation lines. +6. **Heartbeat and renewal** - The owner renews the lease while work is active. +7. **Release or expiry** - Completed work releases capacity. Expired leases + return capacity automatically. + +```mermaid +sequenceDiagram + participant W as Workspace or external caller + participant RM as Resource Manager + participant R as Reconciler + participant T as Runtime target + + W->>RM: Create resource request + RM->>RM: Find pools, policies, classes, grants + RM-->>W: pending, allocated, or rejected + R->>RM: Reconcile pending queues + R->>RM: Allocate or preempt + RM-->>W: Allocated request with target settings + W->>T: Launch or continue workload + W->>RM: Renew lease + W->>RM: Release request +``` + +## Admission rules + +A request is statically eligible only when all of these are true: + +* The pool scope is visible to the request. +* The request matches a pool target binding, or the pool has no target + bindings. +* A policy subject selector matches at least one request subject. +* The request can run on the candidate class `reclaimable` value. +* Every demand matches exactly one resource in the class bundle. +* For grant-based policies, every demanded class resource is present in the + grant and the requested quantity fits the grant limit. +* For grantless policies, every demanded resource is present in the class and + the requested quantity fits the class resource quantity. +* The request's exact `class`, `resource`, kind, and selectors can all be + resolved. + +Resource Manager does not split one request across pools. A plan must satisfy +all demands in one pool, one class, one policy, and either one grant or a +grantless policy path. + +## Ordering + +Candidate plans are ordered deterministically: + +1. Lower accounting pressure first for advisory ordering. +2. Priority lane before normal priority. +3. Higher policy priority. +4. Higher pool rank. +5. Higher class rank. +6. Stable subject, target, policy, class, capacity, and grant identifiers. + +Within a pool queue, older requests win after the priority and rank ordering +keys are equal. + +## Accounting modes + +| Mode | Reconciliation behavior | | --- | --- | -| `preemptible=True` (default) | This run may be preempted to help another request. | -| `preemptible=False` | This run is never preempted. Each requested amount per pool key must be ≤ that key’s **reserved** on the policy; **limit** above reserved does not increase what a non-preemptible step may request. | - -Policies do not override `preemptible`; they only affect ordering and reclaim -among runs that are allowed to be preempted. - -### After preemption - -Preempted step runs are stopped and the resources are released back to the pool. -The steps are put back into the queue and are retried again. If the number of -retries is exhausted or the step is not configured to allow retries, the step fails. See [Automatic Step Retries](https://docs.zenml.io/how-to/steps-pipelines/advanced_features#automatic-step-retries) for more information. - -## Policy scenarios (how reserved, limit, and preemptible interact) - -For the problems these patterns solve in everyday terms, see -[Resource pools](resource-pools.md). - -* Fair share plus burst: set **reserved** to the slice you want to account as - “yours” and **limit** to the most that stack may ever hold. **Preemptible** - steps can **borrow** idle capacity between reserved and limit (and up to the - pool) when the pool has room. **Non-preemptible** steps only use up to - **reserved** per requested key, regardless of a higher limit. -* Production vs experiments: higher **priority** on production policies; - experimental steps stay **preemptible** so production can take capacity or - reclaim borrowed slack when it needs its reservation. -* Non-preemptible training: set `preemptible=False` and size **reserved** so - each step’s per-key request (for example `gpu_count`) is ≤ reserved for that - key. **Limit** can be higher for preemptible burst on the same policy, but it - does not raise the ceiling for non-preemptible requests; raise **reserved** - if those jobs need more per step. The reconciler also blocks non-preemptible - grants that would sit on borrowed capacity in ways that conflict with other - non-preemptible use on the component. -* Several pools for one component: multiple policies with different - **priority** values; higher priority is preferred when queuing and allocating, - subject to each pool’s **limit**. - -For preemption rules (priority vs reclaim), see -[How preemption works](#how-preemption-works). +| `authoritative` | Enforces class quantities, grant limits, reservations, pool/class/policy/grant concurrency, queueing, allocation, and preemption. | +| `advisory` | Computes pressure and orders plans, but does not reject a plan solely because current accounting pressure is high. Use with care. | +| `governance` | Keeps policy and target-setting decisions in ZenML while external infrastructure schedulers handle allocation and preemption. | + +Use `authoritative` for the examples in this documentation unless a scenario +explicitly calls out infrastructure-level scheduling. + +## Concurrency limits + +Concurrency limits cap active requests rather than resource quantities: + +| Limit location | Counts active requests in | +| --- | --- | +| Pool | The pool | +| Class | The selected pool class | +| Policy | The selected policy | +| Grant | The selected grant | + +Grantless policies have no grant-level concurrency limit because no grant row +is selected. Pool, class, and policy concurrency still apply. + +Use concurrency limits when the scarce thing is "how many workloads can run at +once" rather than a descriptor quantity. For example, a pool can enforce at +most eight active H200 jobs even when CPU and memory are tracked as unlimited. + +## Reservations and limits + +Grant-based policies use: + +* `reserved` as protected share. +* `limit` as active usage ceiling. `null` follows the class resource quantity. +* `missing_action` to decide what to do when a request omits a resource that is + listed in the grant. + +For `reclaim_tolerance: "none"`, a grant-based request in an authoritative pool +must fit inside the grant's reserved share. This prevents non-reclaimable work +from depending on capacity that another policy may need back. + +Grantless policies have reservation 0 and limits equal to class resources. +They are useful when the policy subject should access the whole matching pool +instead of a reserved slice. + +## Preemption + +Preemption happens when an authoritative pool has an eligible pending request +that cannot currently allocate without reclaiming capacity or concurrency. + +Victims must be eligible: + +* Lower priority than the waiting request, or same-priority borrowed usage that + can be reclaimed for protected reserved share. +* Not protected by the same preemption group. +* Not configured with `reclaim_tolerance: "none"`. +* Currently active and holding capacity or concurrency needed by the waiter. + +Priority-lane policies sit at the maximum internal priority. They can reclaim +lower-priority eligible work, but they do not automatically preempt other +priority-lane work at the same priority. + +Preemption is cooperative where possible: Resource Manager marks a request +`preempting`, the owner stops or retries, and capacity is returned. If leases +expire or forceful termination is required, the request eventually becomes +terminal and the ledger is cleared. + +## Leases and heartbeats + +An allocated request holds a lease. The owner renews it while the workload is +running. ZenML step launchers do this automatically for dynamic steps. + +Leases protect the ledger from stale allocations: + +* If a queued request times out or is cancelled, queue entries are removed. +* If an allocated request stops renewing, it becomes `expired`. +* Released, cancelled, rejected, preempted, and expired requests no longer hold + active capacity. + +For external workloads, renew the lease before `lease_expires_at` and release +the request promptly when the work is done. + +## Target settings merge + +When a plan wins, Resource Manager returns the target settings selected along +the route. ZenML merges settings from broader to narrower scope: + +1. Pool target route and pool target settings. +2. Pool class target settings. +3. Policy target settings. +4. Grant target settings, when a grant was selected. + +The UI currently exposes component settings at all four levels and service +connector settings at pool and class levels. These settings are applied to the +selected stack component or service connector after allocation. + +## Troubleshooting + +| Symptom | Likely cause | +| --- | --- | +| Request rejected with no admitting policy | No pool target matched, no policy subject matched, wrong class, missing descriptor, or grant missing a demanded resource | +| Non-reclaimable request rejected for reserved share | Grant-based policy did not reserve enough of one demanded resource | +| Request stays pending | Capacity, grant limit, reservation pressure, or concurrency limit is currently blocking allocation | +| Lower-priority work was preempted | Higher-priority eligible work needed the capacity or concurrency and the lower-priority request allowed reclaim | +| Target settings did not apply | The winning pool, class, policy, or grant did not include settings for the selected target type | + +Inspect stuck requests in the UI or with: + +```shell +zenml resource-request list --status pending +zenml resource-request describe +``` + +## See also + +* [Core concepts](resource-pools-core-concepts.md) - descriptors, classes, + policies, subjects, and requests. +* [User guide](resource-pools-user-guide.md) - author settings and request + inspection. +* [External workloads](resource-pools-external-workloads.md) - direct requests + and priority-lane policies. +* [Examples](resource-pools-examples.md) - complete scenarios. +* [Resource pools](resource-pools.md) - overview. diff --git a/docs/book/getting-started/zenml-pro/resource-pools-user-guide.md b/docs/book/getting-started/zenml-pro/resource-pools-user-guide.md new file mode 100644 index 00000000000..9eff2d2066a --- /dev/null +++ b/docs/book/getting-started/zenml-pro/resource-pools-user-guide.md @@ -0,0 +1,197 @@ +--- +description: >- + How pipeline authors request pooled resources with ResourceSettings and + inspect Resource Manager requests. +--- +# User guide + +This guide is for pipeline authors. Platform admins configure descriptors, +pools, classes, and policies in the ZenML Pro UI. Authors express resource +intent on dynamic pipeline steps with `ResourceSettings`. + +{% hint style="info" %} +Resource pools apply only to [dynamic pipelines](https://docs.zenml.io/how-to/steps-pipelines/dynamic_pipelines). +{% endhint %} + +## What you configure + +Most steps use the typed resource settings: + +```python +from zenml import pipeline, step +from zenml.config import ResourceSettings + + +@step( + settings={ + "resources": ResourceSettings( + gpu_count=2, + gpu_class="h200-reserved", + cpu_count=8, + memory="64GiB", + reclaim_tolerance="none", + ) + } +) +def train() -> None: + ... + + +@pipeline(dynamic=True) +def training_pipeline() -> None: + train() +``` + +| Setting | Meaning | +| --- | --- | +| `gpu_count` | Number of GPUs requested; becomes a demand with `kind: "gpu"` | +| `gpu_class` | Optional exact pool class, such as `h200-reserved` | +| `cpu_count` | CPU cores requested; becomes a demand with `kind: "cpu"` and unit `CPU` | +| `memory` | Memory requested, such as `64GiB`; becomes a demand with `kind: "memory"` | +| `resources` | Advanced custom resource demands by descriptor name, kind, or selector | +| `reclaim_tolerance` | Whether the step may be interrupted for pool reasons | +| `allocation_wait_timeout_seconds` | How long the launcher waits for allocation before failing | + +ZenML creates the resource request when the dynamic step is ready to launch. +If capacity is available, the request is allocated and the step runs. If not, +the request stays queued. If no policy, class, or grant path can satisfy the +request, ZenML rejects it with a status reason. + +## Demand translation + +Typed settings translate into Resource Manager demands by kind: + +| ResourceSettings field | Demand shape | +| --- | --- | +| `gpu_count=2` | `{"kind": "gpu", "quantity": 2}` | +| `gpu_count=2, gpu_class="h200-reserved"` | `{"kind": "gpu", "quantity": 2, "class": "h200-reserved"}` | +| `cpu_count=8` | `{"kind": "cpu", "quantity": 8, "unit": "CPU"}` | +| `memory="64GiB"` | `{"kind": "memory", "quantity": 64, "unit": "GiB"}` | + +Pool classes may contain resource bundles, so a single class can satisfy GPU, +CPU, and memory together. If your step requests CPU or memory, the winning +class and policy must account for those resources too, unless the policy uses a +grantless access model. + +## Custom resources + +Use custom resources when your admin created descriptors outside the stock +CPU, memory, and GPU language. For example, a training license: + +```python +from zenml.config import ResourceSettings +from zenml.models import ResourceRequestDemand + +license_settings = ResourceSettings( + resources=[ + ResourceRequestDemand( + resource="training-license", + quantity=1, + unit="seat", + ) + ], +) +``` + +Ask your admin for descriptor names, supported units, and class names. Typed +fields match descriptors by kind, while custom demands can use exact names. + +## Reclaim tolerance + +Reclaim tolerance tells ZenML whether your step may be interrupted when higher +priority work needs capacity. + +| Value | Meaning | Typical use | +| --- | --- | --- | +| `none` | Do not interrupt this step for pool reasons | Production training or critical evaluation | +| `coordinated` | Step may be stopped and retried cleanly | Medium-priority experiments | +| `any` | Best effort; may use unsafe or externally reclaimable capacity | Sweeps and opportunistic runs | + +Non-reclaimable requests can only run on classes with `reclaimable: "never"`. +For grant-based policies in authoritative mode, they must also fit the +reserved grant share. If they do not, the request is rejected instead of +waiting forever. + +Legacy `preemptible=False` maps to `reclaim_tolerance="none"`; +`preemptible=True` maps to `"coordinated"`. Prefer `reclaim_tolerance` in new +code. + +## Waiting for allocation + +Use `allocation_wait_timeout_seconds` to control how long the launcher waits +for capacity: + +```python +ResourceSettings( + gpu_count=1, + cpu_count=4, + memory="32GiB", + reclaim_tolerance="coordinated", + allocation_wait_timeout_seconds=1800, +) +``` + +This timeout affects how long the step launcher waits. It does not change pool +capacity, policy priority, or grant limits. + +## Inspect requests + +The workspace UI shows resource requests linked to dynamic step runs, including +status, selected pool, selected class, queue state, and status reason. + +The ZenML CLI can inspect resource requests: + +```shell +zenml resource-request list --status pending +zenml resource-request list --pool-id +zenml resource-request list --step-run-id +zenml resource-request describe +``` + +Supported list filters include user, reclaim tolerance, component ID, step-run +ID, preemption initiator ID, status, pipeline-run ID, and pool ID. + +## Statuses + +| Status | What it means for your step | +| --- | --- | +| `pending` | The request is admitted and queued. The step waits for allocation. | +| `allocated` | Capacity was granted. The step can run with merged target settings. | +| `rejected` | The request cannot be satisfied by matching policies, classes, or grants. | +| `preempting` | ZenML is asking the step owner to stop so capacity can be reclaimed. | +| `preempted` | The step was stopped by resource-pool preemption. Configure retries if it should re-queue. | +| `cancelled` | The request was cancelled before allocation. | +| `released` | The step finished and capacity was returned. | +| `expired` | The allocated request stopped renewing its lease and capacity was returned. | + +## Common issues + +| Symptom | Likely cause | What to do | +| --- | --- | --- | +| `No resource policy admitted this request.` | No matching policy, no matching pool target, wrong class, or a grant omitted a demanded resource | Check stack component, team/project policy, `gpu_class`, and requested resources with your admin | +| `No grant reserved share satisfies this request with reclaim_tolerance 'none'.` | Non-reclaimable request exceeded grant reservations | Lower reclaim tolerance for best-effort work or ask for reserved grants | +| Step queues longer than expected | Capacity is occupied, a higher-priority queue is ahead, or a concurrency limit is reached | Inspect the request status reason and pool detail page | +| CPU or memory request rejected while GPUs look idle | The matching class or grant does not cover CPU or memory | Ask the admin to model CPU and memory in the class bundle or grant defaults | +| Wrong accelerator selected | `gpu_class` is too broad or missing, or descriptors share `kind: "gpu"` | Ask for the correct class name or descriptor selector | + +## Working with admins + +Share these details when asking for pool access: + +* The stack component you run on, usually an orchestrator or step operator. +* The team, project, pipeline, or account that should receive access. +* GPU, CPU, memory, and custom resources your step requests. +* Whether the workload is production (`reclaim_tolerance="none"`) or + best-effort (`"coordinated"` or `"any"`). +* Whether the step needs a specific class, such as `h200-reserved` or + `a10-burst`. + +## See also + +* [Resource pools](resource-pools.md) - overview. +* [Core concepts](resource-pools-core-concepts.md) - descriptors, classes, + policies, subjects, and requests. +* [Examples](resource-pools-examples.md) - scenarios with admin setup and + author settings. +* [Reconciliation process](resource-pools-reconciliation.md) - queueing, + preemption, leases, and accounting modes. diff --git a/docs/book/getting-started/zenml-pro/resource-pools.md b/docs/book/getting-started/zenml-pro/resource-pools.md index 2bdf3ca9c6e..6999d9807b3 100644 --- a/docs/book/getting-started/zenml-pro/resource-pools.md +++ b/docs/book/getting-started/zenml-pro/resource-pools.md @@ -1,7 +1,7 @@ --- description: >- Fair GPU and compute sharing for AI/ML teams: dependable production capacity, - shared pools, idle reuse, and workspace-level quotas. + shared pools, idle reuse, and organization-level resource governance. icon: layer-group --- @@ -16,197 +16,235 @@ see the [pricing page](https://www.zenml.io/pricing). Resource pools are only available for [dynamic pipelines](https://docs.zenml.io/how-to/steps-pipelines/dynamic_pipelines). {% endhint %} -If you run AI or ML work in a shared environment, you have probably seen the -same problems: jobs fighting over GPUs, surprise slowdowns when another team -launches a big run, or expensive hardware sitting idle while someone waits in a -queue. Resource pools are ZenML Pro’s answer. They are aimed at -*platform and team leaders* who need clear rules, and at *practitioners* who -just want their steps to get the right compute without babysitting the cluster. - -Here are typical situations this feature is built for: - -**“Our production jobs must finish—no surprises.”** -You have training, fine-tuning, or inference that cannot vanish because someone -else submitted a heavier workload. You want a clear agreement: this stack or -this team gets a dependable slice of capacity, and critical steps are not stopped -to make room for ad hoc work. - -**“We share one pool of GPUs across many teams.”** -You need one place that describes how much capacity exists, who may use it, and -what happens when everyone wants it at once—without maintaining a separate -spreadsheet or manual booking process for every pipeline. - -**“We paid for the hardware—we should use it when it’s free.”** -When one group is quiet, you want other teams to use spare capacity so machines -do not sit empty. You also want the original team to get their capacity back when -they return, without a long negotiation or a cluster reconfiguration every -time. - -**“Engineers describe needs; ops maps them to reality.”** -Pipeline authors say what each step requires (GPUs, memory, and so on). Platform -or DevOps ties those stacks to the right shared capacity. The same pipeline code -can run in different environment or stages without hard-coding cluster details. - -**“We need to ration more than just GPUs.”** -Alongside standard compute, you may need to track things like licenses, special -hardware, or how many pipeline steps may run at once. Pools let you treat those -as countable resources under the same workspace-level model, as long as your -organization agrees on names and units. - -None of this replaces your orchestrator or cloud provider—it coordinates -demand so teams see fair queuing, optional sharing of idle capacity, and explicit -rules for critical versus best-effort work. When you are ready for how ZenML -models that behavior, continue with [Introduction to capacity management](#introduction-to-capacity-management) -and [Core concepts](#core-concepts). - -## Introduction to capacity management - -ZenML Pro resource pooling separates three concerns: what capacity exists, what -each pipeline step asks for, and which stack components may draw from which pools under -which rules. The subsections below describe the strategy; [Core concepts](#core-concepts) -then defines terms and mechanics precisely. - -### Supply: resource pools - -On one side, **resource pools** describe resources that are available and -consumable when pipeline steps run. They often include standard compute—GPUs, -CPUs, and RAM—that orchestrators understand and that map to real infrastructure, -but pools are not limited to that. Teams can define **custom resource keys** as a -convention between whoever operates the platform and whoever writes pipelines: -capabilities, access to external services or tools, licenses, or any other -scarce thing you want to schedule in integer units. The same resource key names -must appear in pool capacity and in step requests for the convention to work. -This is the contract between the platform and the users. - -The number of **parallel step runs** is also modeled as a resource today so -concurrency can be capped alongside GPU or CPU demand. -A pool can hold **multiple resource types at once** (for example GPUs and a -`step_run` slot). - -**Who defines pools:** the role that owns infrastructure access—IT ops, -platform engineering, DevOps, or similar — defines and maintains resource pools -at workspace scope so every team in that workspace shares the same catalog -of named capacity. - -### Demand: `ResourceSettings` on steps - -On the other side, **pipeline steps** declare what they need through -`ResourceSettings`. That is typically the ML engineer, AI engineer, or data -engineer annotating each step with GPUs, CPUs, memory, and optionally the same -custom keys the admins put in pools. Authors can align requests with published -pool capacity and follow whatever naming convention the platform agreed on for -non-standard resources. - -Steps also declare whether or not they are **preemptible**. Preemption is the -ability to interrupt a running step run before it completes, to free up resources for -other steps demanding them with a higher priority. Steps that are preemptible -are not guaranteed to complete if they go outside of their reserved capacity. -They may be interrupted, re-added to the queue and restarted again later, or -even cancelled if they are not configured to allow retries. - -That is the product-level tradeoff: **preemptible** steps may access more of the -pool — including capacity that others are not using right now — but another -workload can force them off the machine if priorities and policies require it, -so they can fail early when contention is high. **Non-preemptible** steps opt -into a smaller, dependable slice: they only consume what resources are reserved -for them, so they are not evicted for pool reasons, at the cost of not using the -“burst” capacity above that reservation. - -### The link: subject policies - -**Subject policies** connect pools to execution. A policy binds a workload bearing **stack -component** — today an orchestrator or step operator — to a pool and states what -subset of that pool’s resources that component may use. Steps stay decoupled -from a specific pool name: resolution uses the component in the active stack, -so the same step definition can behave differently in different environments -without code changes. - -Policies are usually owned by the same platform or admin function, but ML -teams can share ownership where it makes sense. Each policy carries **priority**, -**reserved** amounts per key, and optionally **limits**: - -* **Reservation** answers “what pool share is accounted exclusively to this component” - and is the only capacity that **non-preemptible** steps may rely on: they - cannot use the gap between reserved and limit. -* **Limits** (and the pool’s own maximum) bound how many resources **preemptible** work - can take when spare capacity exists; without a limit, the effective ceiling is - what the pool still has free. A higher limit does not raise the ceiling for - non-preemptible requests; you need to raise **reserved** if those jobs need more per step. - -When thinking about subject policies, it's helpful to understand that multiple -steps in the same run or even in different runs will share the resources defined -in the same subject policy if they run on the same stack component. So they will compete for the same reserved resources. If they go outside of their reserved capacity, they will also compete with other stack components that access the same pool. - -Together, pools plus policies plus step annotations implement a **shared, -prioritized, optionally elastic** scheduling story: strict guarantees where -needed, elastic sharing where teams accept preemption risk. - -### What this looks like: three surfaces - -**1 — Pool (supply).** -Platform ops create a workspace pool, name it whatever helps the org (say -**datacenter-one**), and record how much of each scarce thing exists there: **10 -GPUs**, **200 CPUs** and **500 GB of memory**. That number is the shared ceiling everyone draws from. - -```shell -zenml resource-pool create datacenter-one \ - --capacity '{"gpu": 10, "mcpu": 200000, "memory_mb": 5120000}' -``` +Resource pools let platform teams describe shared infrastructure once and let +ZenML decide when a dynamic pipeline step may run. A pool is a virtual ledger +over real infrastructure: Kubernetes node pools, GPU machines, cloud machine +types, service connector resources, or any other bounded platform resource. -**2 — Policy (wiring a stack to a pool).** -They attach a **subject policy** so a specific stack component knows which pool -to use and what slice it may claim. If your pipeline runs on a stack called **prod-stack**, its orchestrator (or step operator) is the component named in -the policy: “prod-stack’s orchestrator may pull from **datacenter-one**, with -*this much* reserved and *this much* limit,” and a priority -versus other stacks. - -```shell -zenml resource-pool attach-policy datacenter-one prod-stack \ - --priority 10 \ - --reserved '{"gpu": 4, "mcpu": 8000, "memory_mb": 20480}' \ - --limit '{"gpu": 6, "mcpu": 16000, "memory_mb": 81920}' -``` +The current user-facing configuration flow is in the ZenML Pro UI: + +* Organization-level pools: **Organization Settings > Resource Pools**. +* Workspace-level pools: **Workspace Settings > Resource Pools**. + +Both pages use the same concepts. Admins define descriptors, create pools, +group resources into pool classes, attach policies, and inspect requests from +the UI. The JSON examples in this section use the Resource Manager REST API +shape for teams that automate the same configuration. + +## Who uses resource pools + +| Actor | Role | +| --- | --- | +| Platform admin | Defines resource descriptors, pools, classes, target bindings, target settings, and policies | +| Pipeline author | Annotates dynamic steps with `ResourceSettings` such as CPU, memory, GPU count, custom resources, and reclaim tolerance | +| ZenML and the reconciler | Convert step settings into resource requests, match policies, queue work, allocate capacity, renew leases, and merge target settings into selected stack components | + +## What resource pools solve + +### Production jobs need dependable capacity + +Critical training, fine-tuning, and inference steps should not disappear +because another team launched a sweep. Policies reserve slices of shared +capacity, assign priority, and decide which work can be interrupted when +contention appears. + +### Shared hardware should not sit idle -**3 — Step request (what the run asks for).** -The data scientist opens a step and says, in effect, **“this step needs three -GPUs, 1 CPU and 2 GB of memory”** and can be preempted. +Reserved capacity can be protected for production while lower-priority work +uses idle capacity up to a configured limit. When higher-priority work returns, +the reconciler can reclaim capacity from work that opted into interruption. + +### Infrastructure is bundled in real life + +Modern infrastructure rarely gives you a GPU without CPU, memory, storage, and +placement constraints. Resource pool classes now model those real bundles. A +class can contain H200 GPUs, CPU, memory, and other descriptors together, with +one class rank, one reclaimability setting, and optional target settings. + +### Authors should request intent, not cluster details + +Pipeline authors request resources by kind: ```python -from zenml import step, pipeline +from zenml import pipeline, step from zenml.config import ResourceSettings -resource_settings=ResourceSettings( - gpu_count=3, - cpu_count=1, - memory="2GiB", - preemptible=True, + +@step( + settings={ + "resources": ResourceSettings( + gpu_count=2, + gpu_class="h200-reserved", + cpu_count=8, + memory="64GiB", + reclaim_tolerance="none", + ) + } ) +def train() -> None: + ... -@step(settings={"resources": resource_settings}) -def my_step(input: str) -> None: - print(input) @pipeline(dynamic=True) -def my_pipeline(input: str) -> None: - my_step(input) +def training_pipeline() -> None: + train() +``` + +Admins decide which pool, class, stack component, node selector, toleration, +and service connector satisfy that intent. + +## How it works + +ZenML separates resource governance into three layers: + +1. **Language**: resource descriptors define names, kinds, units, and + attributes. Stock descriptors are `CPU`, `memory`, and `GPU`; admins can add + custom descriptors such as `h200` or `training-license`. +2. **Supply**: resource pools and classes describe the resource bundles that + exist. A Kubernetes H200 node-pool class might include `h200`, `CPU`, and + `memory` together. +3. **Access**: policies and grants decide which subjects may use a pool, at + which priority, with which reservations, limits, defaults, concurrency caps, + and target settings. + +```mermaid +flowchart LR + Admin[Platform admin
Descriptors, pools, classes, policies] + Author[Pipeline author
ResourceSettings] + Workspace[ZenML workspace
Resource requests] + RM[Resource Manager
Admission and reconciliation] + Target[Stack component
Orchestrator or step operator] + + Admin --> Workspace + Author --> Workspace + Workspace --> RM + RM --> Workspace + Workspace --> Target +``` -if __name__ == "__main__": - my_pipeline(input="Hello, World!") +## Minimal UI walkthrough + +1. Open **Organization Settings > Resource Pools** or + **Workspace Settings > Resource Pools**. +2. Add or review resource descriptors. Use the stock `CPU`, `memory`, and + `GPU` descriptors for common requests; add custom descriptors for specific + hardware or licenses. +3. Create a resource pool. Choose the accounting mode, rank, optional pool + concurrency limit, and the target subjects that the pool governs. +4. Add classes to the pool. Each class is a resource bundle, for example + `h200-reserved` with H200 GPUs, CPU, and memory. A class resource can be + finite or unlimited. Unlimited class resources are tracked for traceability + but are not limited at the pool-class level. +5. Configure pool or class target settings. The UI exposes component settings + and service connector settings where they are currently supported. +6. Open the pool detail page and attach policies. Policies select who may use + the pool. Grants optionally reserve and limit resources inside individual + classes. A policy with `grants: []` admits the subject to all matching + classes in the pool with reservation 0 and limits equal to the class + resources. +7. Run dynamic pipelines. ZenML creates resource requests when steps are ready, + queues them if needed, and applies selected settings when capacity is + allocated. + +## Minimal API shape + +The exact Resource Manager REST API uses nested subject selectors and class +resource bundles. This shortened example shows the important shape: + +```json +{ + "name": "prod-h200-pool", + "description": "Production H200 Kubernetes capacity", + "rank": 100, + "accounting_mode": "authoritative", + "concurrency_limit": 24, + "target_bindings": [ + { + "target_selector": { + "any": [ + { + "subject_type": "organization", + "subject_id": "", + "contains": { + "subject_type": "workspace", + "subject_id": "", + "contains": { + "subject_type": "component", + "subject_id": "", + "attributes": { + "component_type": "step_operator" + } + } + } + }, + { + "subject_type": "organization", + "subject_id": "", + "contains": { + "subject_type": "workspace", + "subject_id": "", + "contains": { + "subject_type": "service_connector", + "subject_id": "", + "attributes": { + "connector_type": "gcp" + } + } + } + } + ] + } + } + ], + "classes": [ + { + "class": "h200-reserved", + "rank": 100, + "reclaimable": "never", + "concurrency_limit": 8, + "resources": [ + {"resource": "h200", "quantity": 8, "unit": "GPU"}, + {"resource": "CPU", "quantity": 128, "unit": "CPU"}, + {"resource": "memory", "quantity": 1024, "unit": "GiB"} + ] + } + ], + "target_settings": [ + { + "target_type": "service_connector", + "settings": { + "connector_id": "", + "resource_type": "kubernetes-cluster", + "resource_id": "prod-eu" + } + } + ], + "attributes": { + "region": "eu-west-1", + "platform": "kubernetes" + } +} ``` -When they launch a **dynamic** run on **prod-stack**, ZenML turns that into a **resource -request**: three GPUs, matched against **datacenter-one** through prod-stack’s policy. -If three are free, the step proceeds; if not, it **waits** in line; if the ask -breaks the rules (too many GPUs, non-preemptible without reservation), it -**fails fast** with a clear status. Run and step views show **queued / allocated -/ rejected** so operators can compare ZenML to the real cluster. +Continue with the [Core concepts](resource-pools-core-concepts.md) for the +schema and selector conventions, then use the +[Admin guide](resource-pools-admin-guide.md) for the UI workflow. + +## Where to go next -The subsections explain how queues, borrowing, and preemption implement -that model in ZenML Pro. +| Page | For whom | What you learn | +| --- | --- | --- | +| [Core concepts](resource-pools-core-concepts.md) | Everyone | Descriptors, pool classes as bundles, policies, subjects, requests, leases | +| [Admin guide](resource-pools-admin-guide.md) | Platform admins | Configure descriptors, pools, classes, target settings, policies, and grants in the UI | +| [User guide](resource-pools-user-guide.md) | Pipeline authors | `ResourceSettings`, reclaim tolerance, and request inspection | +| [External workloads](resource-pools-external-workloads.md) | Platform admins, integrators | Service accounts, direct requests, and priority lanes | +| [Reconciliation process](resource-pools-reconciliation.md) | Admins, operators | Queueing, accounting modes, preemption, leases, and heartbeats | +| [Examples](resource-pools-examples.md) | Everyone | Believable end-to-end resource pool setups | ## See also -* [Workspaces](./workspaces.md) — pools are scoped to the workspace. -* [Teams](./teams.md) — organizational context for who owns which stacks and - policies. +* [Workspaces](./workspaces.md) - workspace-scoped pools and settings. +* [Teams](./teams.md) - organizational subjects for policy access. +* [Service accounts](./service-accounts.md) - identity for external workload + integrations. * ZenML OSS: [step and pipeline configuration](https://docs.zenml.io/how-to/steps-pipelines/configuration). diff --git a/docs/book/getting-started/zenml-pro/toc.md b/docs/book/getting-started/zenml-pro/toc.md index 36e80c8b6bb..5ffcd73e4a7 100644 --- a/docs/book/getting-started/zenml-pro/toc.md +++ b/docs/book/getting-started/zenml-pro/toc.md @@ -39,6 +39,9 @@ * [Triggers](triggers.md) * [Resource Pools](resource-pools.md) * [Core Concepts](resource-pools-core-concepts.md) + * [Admin Guide](resource-pools-admin-guide.md) + * [User Guide](resource-pools-user-guide.md) + * [External Workloads](resource-pools-external-workloads.md) * [Reconciliation Process](resource-pools-reconciliation.md) * [Examples](resource-pools-examples.md) diff --git a/pyproject.toml b/pyproject.toml index 85d1b08a783..e5b0505c314 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -55,6 +55,7 @@ dependencies = [ "packaging>=24.1", "psutil>=5.0.0", "pydantic>=2.0,<=2.12.5", + "pydantic-settings>=2.0.0,<3.0.0", "python-dateutil>=2.8.1,<3.0.0", "pyyaml>=6.0.1", "rich>=12.0.0", # Make sure to keep the version aligned with our `jupyter` extra diff --git a/scripts/lint.sh b/scripts/lint.sh index 75db0de15bc..60af6af038b 100755 --- a/scripts/lint.sh +++ b/scripts/lint.sh @@ -39,6 +39,14 @@ if [ "$SKIP_YAMLFIX" = false ]; then yamlfix --check .github tests -e "dependabot.yml" -e "workflows/release_prepare.yml" -e "workflows/release_finalize.yml" -e "workflows/integration-test-fast-services.yml" -e "workflows/integration-test-slow-services.yml" fi +# autoflake replacement: checks for unused imports and variables +ruff check $SRC --select F401,F841 --exclude "__init__.py" --exclude "*.ipynb" --isolated + +ruff format $SRC --check + +# check type annotations +mypy $SRC_NO_TESTS + # checks for GitHub Actions security issues (SHA pinning, version mismatches, etc.) if [ "$SKIP_ZIZMOR" = false ] && [ -d ".github/workflows" ]; then # Resolve GH_TOKEN: use existing env var, or try gh CLI @@ -53,12 +61,4 @@ if [ "$SKIP_ZIZMOR" = false ] && [ -d ".github/workflows" ]; then echo "⚠️ Skipping zizmor check: no GH_TOKEN and gh CLI not authenticated." echo " Run: GH_TOKEN=\$(gh auth token) bash scripts/lint.sh" fi -fi - -# autoflake replacement: checks for unused imports and variables -ruff check $SRC --select F401,F841 --exclude "__init__.py" --exclude "*.ipynb" --isolated - -ruff format $SRC --check - -# check type annotations -mypy $SRC_NO_TESTS \ No newline at end of file +fi \ No newline at end of file diff --git a/src/zenml/cli/__init__.py b/src/zenml/cli/__init__.py index acf4e16dc69..96be5373dd0 100644 --- a/src/zenml/cli/__init__.py +++ b/src/zenml/cli/__init__.py @@ -2578,7 +2578,6 @@ def my_pipeline(...): from zenml.cli.model import * # noqa from zenml.cli.model_registry import * # noqa from zenml.cli.pipeline import * # noqa -from zenml.cli.resource_pool import * # noqa from zenml.cli.resource_request import * # noqa from zenml.cli.secret import * # noqa from zenml.cli.served_model import * # noqa diff --git a/src/zenml/cli/resource_pool.py b/src/zenml/cli/resource_pool.py deleted file mode 100644 index a94eb45235a..00000000000 --- a/src/zenml/cli/resource_pool.py +++ /dev/null @@ -1,608 +0,0 @@ -# Copyright (c) ZenML GmbH 2026. All Rights Reserved. -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at: -# -# https://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express -# or implied. See the License for the specific language governing -# permissions and limitations under the License. -"""Resource pool CLI commands.""" - -from typing import Any, Dict, Optional - -import click - -from zenml.cli import utils as cli_utils -from zenml.cli.cli import TagGroup, cli -from zenml.cli.utils import ( - OutputFormat, - convert_structured_str_to_dict, - get_default_output_format, - list_options, -) -from zenml.client import Client -from zenml.console import console -from zenml.enums import CliCategories, StackComponentType -from zenml.models import ( - Page, - ResourcePoolFilter, - ResourcePoolSubjectPolicyResponse, - ResourceRequestResponse, -) - - -@cli.group(cls=TagGroup, tag=CliCategories.MANAGEMENT_TOOLS) -def resource_pool() -> None: - """Commands for interacting with resource pools.""" - - -def _parse_resources(resources: str) -> Dict[str, int]: - """Parse a JSON/YAML resources string into int values. - - Args: - resources: The resources string to parse. - - Raises: - ValueError: If the resources string is invalid. - - Returns: - A dictionary of resource types and their values. - """ - raw = convert_structured_str_to_dict(resources) - parsed: Dict[str, int] = {} - for key, value in raw.items(): - try: - parsed[key] = int(value) - except (TypeError, ValueError) as e: - raise ValueError( - f"Invalid resource value for key '{key}': {value!r}. " - "Resource values must be integers." - ) from e - return parsed - - -@resource_pool.command("create", help="Create a resource pool.") -@click.argument("name", type=str, required=True) -@click.option( - "--capacity", - "-c", - type=str, - required=True, - help="Resource pool capacity in JSON or YAML. " - 'Example: \'{"gpu": 2, "cpu": 8}\'.', -) -@click.option( - "--description", - "-d", - type=str, - required=False, - help="Optional description for the resource pool.", -) -def create_resource_pool( - name: str, capacity: str, description: Optional[str] = None -) -> None: - """Create a resource pool. - - Args: - name: Name of the resource pool. - capacity: Capacity in JSON or YAML. - description: Optional description. - """ - try: - parsed_capacity = _parse_resources(capacity) - except ValueError as err: - cli_utils.exception(err) - - with console.status(f"Creating resource pool '{name}'...\n"): - try: - resource_pool_ = Client().create_resource_pool( - name=name, capacity=parsed_capacity, description=description - ) - except Exception as err: - cli_utils.exception(err) - else: - cli_utils.declare( - f"Successfully created resource pool `{resource_pool_.name}`." - ) - - -@resource_pool.command("update", help="Update a resource pool.") -@click.argument("name_id_or_prefix", type=str, required=True) -@click.option( - "--capacity", - "-c", - type=str, - required=False, - help=( - "Updated capacity in JSON or YAML. Setting a value to 0 removes the " - "resource from the pool. Example: '{\"gpu\": 0}'." - ), -) -@click.option( - "--description", - "-d", - type=str, - required=False, - help="Updated description for the resource pool.", -) -def update_resource_pool( - name_id_or_prefix: str, - capacity: Optional[str] = None, - description: Optional[str] = None, -) -> None: - """Update a resource pool. - - Args: - name_id_or_prefix: Name, ID or prefix of the resource pool. - capacity: Updated capacity in JSON or YAML. - description: Updated description. - """ - if capacity is None and description is None: - cli_utils.declare("No updates specified.") - return - - parsed_capacity = None - if capacity is not None: - try: - parsed_capacity = _parse_resources(capacity) - except ValueError as err: - cli_utils.exception(err) - return - - with console.status("Updating resource pool...\n"): - try: - resource_pool_ = Client().update_resource_pool( - name_id_or_prefix=name_id_or_prefix, - description=description, - capacity=parsed_capacity, - ) - except Exception as err: - cli_utils.exception(err) - else: - cli_utils.declare( - f"Successfully updated resource pool `{resource_pool_.name}`." - ) - - -@resource_pool.command("describe", help="Describe a resource pool.") -@click.argument("name_id_or_prefix", type=str, required=True) -def describe_resource_pool(name_id_or_prefix: str) -> None: - """Describe a resource pool. - - Args: - name_id_or_prefix: Name, ID or prefix of the resource pool. - """ - try: - resource_pool_ = Client().get_resource_pool( - name_id_or_prefix=name_id_or_prefix, hydrate=True - ) - except KeyError as err: - cli_utils.exception(err) - else: - cli_utils.print_pydantic_model( - title=f"Resource pool '{resource_pool_.name}'", - model=resource_pool_, - ) - - -@resource_pool.command("list", help="List resource pools.") -@list_options( - ResourcePoolFilter, - default_columns=["id", "name", "capacity", "occupied_resources"], -) -def list_resource_pools( - columns: str, output_format: OutputFormat, **kwargs: Any -) -> None: - """List resource pools. - - Args: - columns: Columns to display in output. - output_format: Format for output (table/json/yaml/csv/tsv). - **kwargs: Keyword arguments to filter resource pools by. - """ - with console.status("Listing resource pools...\n"): - pools = Client().list_resource_pools(**kwargs, hydrate=True) - - cli_utils.print_page( - pools, - columns, - output_format, - empty_message="No resource pools found for this filter.", - ) - - -@resource_pool.command("delete", help="Delete a resource pool.") -@click.argument("name_id_or_prefix", type=str, required=True) -@click.option( - "--yes", - "-y", - is_flag=True, - help="Don't ask for confirmation.", -) -def delete_resource_pool(name_id_or_prefix: str, yes: bool = False) -> None: - """Delete a resource pool. - - Args: - name_id_or_prefix: Name, ID or prefix of the resource pool. - yes: If set, don't ask for confirmation. - """ - if not yes: - confirmation = cli_utils.confirmation( - "Are you sure you want to delete resource pool " - f"`{name_id_or_prefix}`?" - ) - if not confirmation: - cli_utils.declare("Resource pool deletion canceled.") - return - - with console.status("Deleting resource pool...\n"): - try: - Client().delete_resource_pool(name_id_or_prefix=name_id_or_prefix) - except Exception as err: - cli_utils.exception(err) - else: - cli_utils.declare(f"Deleted resource pool `{name_id_or_prefix}`.") - - -@resource_pool.command( - "attach-policy", - help="Attach a policy to a resource pool.", -) -@click.argument("resource_pool", type=str, required=True) -@click.argument("component", type=str, required=True) -@click.option( - "--component-type", - "-t", - type=click.Choice( - [ - StackComponentType.ORCHESTRATOR.value, - StackComponentType.STEP_OPERATOR.value, - ] - ), - default=StackComponentType.ORCHESTRATOR.value, - help="Type of the component to attach the policy to.", -) -@click.option( - "--priority", - "-p", - type=int, - required=True, - help="Priority of this pool for the component. Higher means preferred.", -) -@click.option( - "--reserved", - "-r", - type=str, - required=True, - help="Resources that are reserved for the component.", -) -@click.option( - "--limit", - "-l", - type=str, - required=True, - help="Maximum resources that the component can use.", -) -def attach_policy_to_resource_pool( - resource_pool: str, - component: str, - component_type: StackComponentType, - priority: int, - reserved: str, - limit: Optional[str] = None, -) -> None: - """Attach a policy to a resource pool. - - Args: - resource_pool: Name, ID or prefix of the resource pool. - component: Name, ID or prefix of the component. - component_type: Type of the component to attach the policy to. - priority: Priority of the assignment. - reserved: Resources that are reserved for the component. - limit: Maximum resources that the component can use. - """ - try: - parsed_reserved = _parse_resources(reserved) - except ValueError as err: - cli_utils.exception(err) - - parsed_limit = None - if limit is not None: - try: - parsed_limit = _parse_resources(limit) - except ValueError as err: - cli_utils.exception(err) - - with console.status("Attaching pool policy to resource pool...\n"): - try: - pool = Client().get_resource_pool( - name_id_or_prefix=resource_pool, hydrate=False - ) - component_model = Client().get_stack_component( - name_id_or_prefix=component, - component_type=component_type, - allow_name_prefix_match=False, - hydrate=False, - ) - Client().create_resource_pool_subject_policy( - component_id=component_model.id, - pool_id=pool.id, - priority=priority, - reserved=parsed_reserved, - limit=parsed_limit, - ) - except Exception as err: - cli_utils.exception(err) - else: - cli_utils.declare( - f"Attached pool policy for component `{component_model.name}` " - f"to resource pool `{pool.name}` with priority {priority}." - ) - - -@resource_pool.command( - "detach-policy", - help="Detach a policy from a resource pool.", -) -@click.argument("resource_pool", type=str, required=True) -@click.argument("component", type=str, required=True) -@click.option( - "--component-type", - "-t", - type=click.Choice( - [ - StackComponentType.ORCHESTRATOR.value, - StackComponentType.STEP_OPERATOR.value, - ] - ), - default=StackComponentType.ORCHESTRATOR.value, - help="Type of the component to detach the policy from.", -) -def detach_policy_from_resource_pool( - resource_pool: str, - component: str, - component_type: StackComponentType, -) -> None: - """Detach a policy from a resource pool. - - Args: - resource_pool: Name, ID or prefix of the resource pool. - component: Name, ID or prefix of the component. - component_type: Type of the component to detach the policy from. - - Raises: - KeyError: If no policy is found for the component in the resource pool. - """ - with console.status("Detaching pool policy from resource pool...\n"): - try: - pool = Client().get_resource_pool( - name_id_or_prefix=resource_pool, hydrate=False - ) - component_model = Client().get_stack_component( - name_id_or_prefix=component, - component_type=component_type, - allow_name_prefix_match=False, - hydrate=False, - ) - policies = Client().list_resource_pool_subject_policies( - pool_id=pool.id, - component_id=component_model.id, - hydrate=False, - ) - if not policies.items: - raise KeyError( - f"No policy found for component `{component_model.name}` " - f"in resource pool `{pool.name}`." - ) - Client().delete_resource_pool_subject_policy( - policy_id=policies.items[0].id - ) - except Exception as err: - cli_utils.exception(err) - else: - cli_utils.declare( - f"Detached pool policy for component `{component_model.name}` " - f"from resource pool `{pool.name}`." - ) - - -@resource_pool.command( - "list-policies", - help=( - "List subject policies (orchestrator components attached to pools with " - "priority, reserved resources, and limits). Provide a pool and/or a " - "component; at least one is required." - ), -) -@click.argument("pool", type=str, required=False, default=None) -@click.option( - "--component", - "-c", - type=str, - required=False, - help=( - "Orchestrator component name or ID. If the pool is omitted, lists all " - "policies for this component across pools. If the pool is set, only " - "that component's policy for that pool is listed." - ), -) -@click.option( - "--component-type", - "-t", - type=click.Choice( - [ - StackComponentType.ORCHESTRATOR.value, - StackComponentType.STEP_OPERATOR.value, - ] - ), - default=None, - help="Type of the component to list the policies for.", -) -@click.option( - "--columns", - type=str, - default="id,component,pool,priority,reserved,limit,created", - help="Comma-separated list of columns to display, or 'all' for all columns.", -) -@click.option( - "--output", - "-o", - "output_format", - type=click.Choice(["table", "json", "yaml", "tsv", "csv"]), - default=get_default_output_format(), - help="Output format for the list.", -) -def list_resource_pool_policies( - pool: Optional[str], - component: Optional[str], - component_type: Optional[StackComponentType], - columns: str, - output_format: OutputFormat, -) -> None: - """List policies attached to resource pool subject(s). - - Provide a pool (positional or ``--pool``), a ``--component``, or both. - Omitting the pool lists every policy for that orchestrator component. - - Args: - pool: Optional positional pool name, ID, or prefix. - component: Optional orchestrator component name or ID. - component_type: Type of the component to list the policies for. - columns: Columns to print. - output_format: Output format for the table or structured data. - """ - pool_id = None - pool_label = pool - with console.status("Loading resource pool policies...\n"): - try: - if pool is not None: - pool_model = Client().get_resource_pool( - name_id_or_prefix=pool, hydrate=False - ) - pool_id = pool_model.id - pool_label = pool_model.name - component_id = None - if component is not None: - component_model = Client().get_stack_component( - name_id_or_prefix=component, - component_type=component_type - or StackComponentType.ORCHESTRATOR, - allow_name_prefix_match=False, - hydrate=False, - ) - component_id = component_model.id - policies_page = Client().list_resource_pool_subject_policies( - pool_id=pool_id, - component_id=component_id, - hydrate=True, - ) - except Exception as err: - cli_utils.exception(err) - - page: Page[ResourcePoolSubjectPolicyResponse] = policies_page - if pool is not None and component is not None: - empty = ( - f"No policies found for resource pool `{pool_label}` and " - f"component `{component}`." - ) - elif pool is not None: - empty = f"No policies found for resource pool `{pool_label}`." - else: - empty = f"No policies found for orchestrator component `{component}`." - - cli_utils.print_page( - page=page, - columns=columns, - output_format=output_format, - empty_message=empty, - ) - - -@resource_pool.command( - "requests", - help="List queued and/or active resource requests for a resource pool.", -) -@click.argument("pool", type=str, required=True) -@click.option( - "--view", - type=click.Choice(["queued", "active", "all"]), - default="all", - show_default=True, - help="Which requests to show for this pool.", -) -@click.option( - "--columns", - "-c", - type=str, - default="id,state,status,component,step_run,created", - help="Comma-separated list of columns to display, or 'all' for all columns.", -) -@click.option( - "--output", - "-o", - "output_format", - type=click.Choice(["table", "json", "yaml", "tsv", "csv"]), - default=get_default_output_format(), - help="Output format for the list.", -) -def list_resource_pool_requests( - pool: str, - view: str, - columns: str, - output_format: OutputFormat, -) -> None: - """List queued and/or active resource requests for a resource pool. - - Args: - pool: Name, ID or prefix of the resource pool. - view: Which requests to show ("queued", "active", or "all"). - columns: Output columns. - output_format: Output format. - """ - with console.status("Loading resource pool requests...\n"): - pool_model = Client().get_resource_pool( - name_id_or_prefix=pool, hydrate=True - ) - - if view == "queued": - items = list(item.request for item in pool_model.queued_requests) - state_by_id = {r.id: "queued" for r in items} - elif view == "active": - items = list(item.request for item in pool_model.active_requests) - state_by_id = {r.id: "active" for r in items} - else: - active = list(item.request for item in pool_model.active_requests) - queued = list(item.request for item in pool_model.queued_requests) - items = [*active, *queued] - state_by_id = {r.id: "active" for r in active} - state_by_id.update({r.id: "queued" for r in queued}) - - page: Page["ResourceRequestResponse"] = Page( - index=1, - max_size=max(1, len(items)), - total_pages=1, - total=len(items), - items=items, - ) - - def _row_with_state(item: Any, _: OutputFormat) -> dict[str, Any]: - return {"state": state_by_id.get(item.id, "unknown")} - - cli_utils.print_page( - page=page, - columns=columns, - output_format=output_format, - row_formatter=_row_with_state, - empty_message=( - f"No {view} resource requests found for resource pool " - f"`{pool_model.name}`." - if view != "all" - else "No resource requests found for resource pool " - f"`{pool_model.name}`." - ), - ) diff --git a/src/zenml/cli/resource_request.py b/src/zenml/cli/resource_request.py index 53941d88b69..fba755793c4 100644 --- a/src/zenml/cli/resource_request.py +++ b/src/zenml/cli/resource_request.py @@ -43,7 +43,7 @@ def describe_resource_request(resource_request_id: str) -> None: @resource_request.command("list", help="List resource requests.") @list_options( ResourceRequestFilter, - default_columns=["id", "status", "component", "step_run", "created"], + default_columns=["id", "status", "pool_name", "step_name", "created"], ) def list_resource_requests( columns: str, output_format: OutputFormat, **kwargs: Any @@ -65,16 +65,6 @@ def list_resource_requests( page, columns, output_format, + column_aliases={"pool_name": "pool", "step_name": "step"}, empty_message="No resource requests found for this filter.", ) - - -@resource_request.command("delete", help="Delete a resource request.") -@click.argument("resource_request_id", type=str, required=True) -def delete_resource_request(resource_request_id: str) -> None: - """Delete a resource request. - - Args: - resource_request_id: ID of the resource request. - """ - Client().zen_store.delete_resource_request(UUID(resource_request_id)) diff --git a/src/zenml/client.py b/src/zenml/client.py index a5152dfdbdc..8384c54a544 100644 --- a/src/zenml/client.py +++ b/src/zenml/client.py @@ -67,6 +67,8 @@ LogicalOperators, ModelStages, OAuthDeviceStatus, + ResourceRequestReclaimTolerance, + ResourceRequestStatus, RunWaitConditionResolution, ServiceState, SorterOps, @@ -152,14 +154,8 @@ ProjectRequest, ProjectResponse, ProjectUpdate, - ResourcePoolFilter, - ResourcePoolRequest, - ResourcePoolResponse, - ResourcePoolSubjectPolicyFilter, - ResourcePoolSubjectPolicyRequest, - ResourcePoolSubjectPolicyResponse, - ResourcePoolSubjectPolicyUpdate, - ResourcePoolUpdate, + ResourceRequestFilter, + ResourceRequestResponse, RunMetadataRequest, RunMetadataResource, RunTemplateFilter, @@ -2276,235 +2272,68 @@ def delete_stack_component( component.name, ) - # -------------------------------- Resource Pools -------------------------- - - def create_resource_pool( - self, - name: str, - capacity: Dict[str, int], - description: Optional[str] = None, - ) -> ResourcePoolResponse: - """Create a resource pool. - - Args: - name: The name of the resource pool. - capacity: The capacity of the resource pool. - description: The description of the resource pool. - - Returns: - The created resource pool. - """ - request = ResourcePoolRequest( - name=name, - description=description, - capacity=capacity, - ) - return self.zen_store.create_resource_pool(resource_pool=request) - - def get_resource_pool( - self, - name_id_or_prefix: Union[str, UUID], - allow_name_prefix_match: bool = True, - hydrate: bool = True, - ) -> ResourcePoolResponse: - """Get a resource pool. - - Args: - name_id_or_prefix: The name, id or prefix of the resource pool to - get. - allow_name_prefix_match: If True, allow matching by name prefix. - hydrate: Flag deciding whether to hydrate the output model(s) - by including metadata fields in the response. - - Returns: - The resource pool. - """ - return self._get_entity_by_id_or_name_or_prefix( - get_method=self.zen_store.get_resource_pool, - list_method=self.list_resource_pools, - name_id_or_prefix=name_id_or_prefix, - allow_name_prefix_match=allow_name_prefix_match, - hydrate=hydrate, - ) - - def list_resource_pools( - self, - sort_by: str = "created", - page: int = PAGINATION_STARTING_PAGE, - size: int = PAGE_SIZE_DEFAULT, - logical_operator: LogicalOperators = LogicalOperators.AND, - id: UUIDFilterOption = None, - created: DatetimeFilterOption = None, - updated: DatetimeFilterOption = None, - name: StringFilterOption = None, - user: UUIDFilterOption = None, - hydrate: bool = False, - ) -> Page[ResourcePoolResponse]: - """Lists resource pools. - - Args: - sort_by: The column to sort by - page: The page of items - size: The maximum size of all pages - logical_operator: Which logical operator to use [and, or] - id: Use the id of resource pool to filter by. - created: Filter by creation time. - updated: Filter by last updated time. - name: The name of the resource pool to filter by. - user: The ID of name of the user to filter by. - hydrate: Flag deciding whether to hydrate the output model(s) - by including metadata fields in the response. - - Returns: - A page of resource pools. - """ - filter_model = ResourcePoolFilter( - page=page, - size=size, - sort_by=sort_by, - logical_operator=logical_operator, - id=id, - name=name, - created=created, - updated=updated, - user=user, - ) - - return self.zen_store.list_resource_pools( - filter_model=filter_model, hydrate=hydrate - ) - - def update_resource_pool( - self, - name_id_or_prefix: Union[UUID, str], - description: Optional[str] = None, - capacity: Optional[Dict[str, int]] = None, - ) -> ResourcePoolResponse: - """Update a resource pool. - - Args: - name_id_or_prefix: The name, id or prefix of the resource pool to - update. - description: The new description of the resource pool. - capacity: The new capacity of the resource pool. Setting a value to - 0 will remove the resource from the pool. - - Returns: - The updated resource pool. - """ - resource_pool = self.get_resource_pool( - name_id_or_prefix=name_id_or_prefix, - allow_name_prefix_match=False, - ) - - update_model = ResourcePoolUpdate( - description=description, - capacity=capacity, - ) - - return self.zen_store.update_resource_pool( - resource_pool_id=resource_pool.id, - update=update_model, - ) - - def delete_resource_pool( - self, - name_id_or_prefix: Union[str, UUID], - ) -> None: - """Delete a resource pool. - - Args: - name_id_or_prefix: The name, id or prefix of the resource pool to - delete. - """ - resource_pool = self.get_resource_pool( - name_id_or_prefix=name_id_or_prefix, - allow_name_prefix_match=False, - ) - self.zen_store.delete_resource_pool(resource_pool_id=resource_pool.id) - - def create_resource_pool_subject_policy( - self, - component_id: UUID, - pool_id: UUID, - priority: int, - reserved: Optional[Dict[str, int]] = None, - limit: Optional[Dict[str, int]] = None, - ) -> ResourcePoolSubjectPolicyResponse: - """Create a resource pool subject policy. - - Args: - component_id: The component ID this policy applies to. - pool_id: The pool ID this policy belongs to. - priority: The policy priority. - reserved: Optional reserved resources by key. - limit: Optional resource limits by key. - - Returns: - The created policy. - """ - request = ResourcePoolSubjectPolicyRequest( - component_id=component_id, - pool_id=pool_id, - priority=priority, - reserved=reserved, - limit=limit, - ) - return self.zen_store.create_resource_pool_subject_policy( - policy=request - ) - - def get_resource_pool_subject_policy( - self, policy_id: UUID, hydrate: bool = True - ) -> ResourcePoolSubjectPolicyResponse: - """Get a resource pool subject policy by ID. + def get_resource_request( + self, resource_request_id: UUID, hydrate: bool = True + ) -> ResourceRequestResponse: + """Get a resource request by ID. Args: - policy_id: The ID of the policy to fetch. + resource_request_id: The resource request ID. hydrate: Whether to include metadata in the response. Returns: - The requested policy. + The requested resource request. """ - return self.zen_store.get_resource_pool_subject_policy( - policy_id=policy_id, hydrate=hydrate + return self.zen_store.get_resource_request( + resource_request_id=resource_request_id, + hydrate=hydrate, ) - def list_resource_pool_subject_policies( + def list_resource_requests( self, sort_by: str = "created", page: int = PAGINATION_STARTING_PAGE, size: int = PAGE_SIZE_DEFAULT, logical_operator: LogicalOperators = LogicalOperators.AND, - id: UUIDFilterOption = None, - created: DatetimeFilterOption = None, - updated: DatetimeFilterOption = None, - user: UUIDFilterOption = None, - pool_id: UUIDFilterOption = None, - component_id: UUIDFilterOption = None, - priority: IntegerFilterOption = None, + id: Optional[Union[UUID, str]] = None, + created: Optional[datetime] = None, + updated: Optional[datetime] = None, + user: Optional[Union[UUID, str]] = None, + reclaim_tolerance: Optional[ + Union[ResourceRequestReclaimTolerance, str] + ] = None, + component_id: Optional[Union[UUID, str]] = None, + step_run_id: Optional[Union[UUID, str]] = None, + preemption_initiated_by_id: Optional[Union[UUID, str]] = None, + status: Optional[Union[ResourceRequestStatus, str]] = None, + pipeline_run_id: Optional[Union[UUID, str]] = None, + pool_id: Optional[Union[UUID, str]] = None, hydrate: bool = False, - ) -> Page[ResourcePoolSubjectPolicyResponse]: - """List resource pool subject policies. + ) -> Page[ResourceRequestResponse]: + """List resource requests. Args: sort_by: The column to sort by. page: The page of items. size: The maximum size of all pages. logical_operator: Which logical operator to use. - id: Filter by policy ID. + id: Filter by request ID. created: Filter by creation time. updated: Filter by last update time. user: Filter by user name/ID. + reclaim_tolerance: Filter by reclaim tolerance. + component_id: Filter by stack component name/ID. + step_run_id: Filter by step run name/ID. + preemption_initiated_by_id: Filter by preempting request ID. + status: Filter by lifecycle status. + pipeline_run_id: Filter by pipeline run name/ID. pool_id: Filter by resource pool name/ID. - component_id: Filter by component name/ID. - priority: Filter by priority. - hydrate: Whether to include metadata in the response. + hydrate: Whether to include related resources in the response. Returns: - A page of matching policies. + A page of matching resource requests. """ - filter_model = ResourcePoolSubjectPolicyFilter( + filter_model = ResourceRequestFilter( page=page, size=size, sort_by=sort_by, @@ -2513,47 +2342,19 @@ def list_resource_pool_subject_policies( created=created, updated=updated, user=user, - pool_id=pool_id, + reclaim_tolerance=reclaim_tolerance, component_id=component_id, - priority=priority, - ) - return self.zen_store.list_resource_pool_subject_policies( - filter_model=filter_model, hydrate=hydrate - ) - - def update_resource_pool_subject_policy( - self, - policy_id: UUID, - priority: Optional[int] = None, - reserved: Optional[Dict[str, int]] = None, - limit: Optional[Dict[str, int]] = None, - ) -> ResourcePoolSubjectPolicyResponse: - """Update a resource pool subject policy. - - Args: - policy_id: The policy ID. - priority: Updated priority. - reserved: Updated reserved resources by key. - limit: Updated limits by key. - - Returns: - The updated policy. - """ - update = ResourcePoolSubjectPolicyUpdate( - priority=priority, reserved=reserved, limit=limit + step_run_id=step_run_id, + preemption_initiated_by_id=preemption_initiated_by_id, + status=status, + pipeline_run_id=pipeline_run_id, + pool_id=pool_id, ) - return self.zen_store.update_resource_pool_subject_policy( - policy_id=policy_id, update=update + return self.zen_store.list_resource_requests( + filter_model=filter_model, + hydrate=hydrate, ) - def delete_resource_pool_subject_policy(self, policy_id: UUID) -> None: - """Delete a resource pool subject policy. - - Args: - policy_id: The policy ID. - """ - self.zen_store.delete_resource_pool_subject_policy(policy_id=policy_id) - # --------------------------------- Flavors -------------------------------- def create_flavor( diff --git a/src/zenml/config/__init__.py b/src/zenml/config/__init__.py index e1654db9889..2bd5845ee18 100644 --- a/src/zenml/config/__init__.py +++ b/src/zenml/config/__init__.py @@ -12,27 +12,33 @@ # or implied. See the License for the specific language governing # permissions and limitations under the License. """Config classes.""" + +from zenml.config.cache_policy import CachePolicy from zenml.config.deployment_settings import ( - DeploymentSettings, + AppExtensionSpec, + CORSConfig, DeploymentDefaultEndpoints, DeploymentDefaultMiddleware, - EndpointSpec, + DeploymentSettings, EndpointMethod, + EndpointSpec, MiddlewareSpec, - AppExtensionSpec, SecureHeadersConfig, - CORSConfig, ) from zenml.config.docker_settings import ( DockerSettings, - PythonPackageInstaller, PythonEnvironmentExportMethod, + PythonPackageInstaller, +) +from zenml.config.resource_settings import ( + ByteUnit, + PoolResourceDemand, + ResourceInput, + ResourceSettings, ) -from zenml.config.resource_settings import ResourceSettings, ByteUnit from zenml.config.retry_config import StepRetryConfig from zenml.config.schedule import Schedule from zenml.config.store_config import StoreConfiguration -from zenml.config.cache_policy import CachePolicy __all__ = [ "DeploymentSettings", @@ -48,6 +54,8 @@ "PythonPackageInstaller", "PythonEnvironmentExportMethod", "ResourceSettings", + "ResourceInput", + "PoolResourceDemand", "ByteUnit", "StepRetryConfig", "Schedule", diff --git a/src/zenml/config/resource_settings.py b/src/zenml/config/resource_settings.py index 67abb31b9d3..3e237c0adbd 100644 --- a/src/zenml/config/resource_settings.py +++ b/src/zenml/config/resource_settings.py @@ -13,9 +13,11 @@ # permissions and limitations under the License. """Resource settings class used to specify resources for a step.""" +from __future__ import annotations + import math from enum import Enum -from typing import Dict, Literal, Optional, Union +from typing import TYPE_CHECKING, Any, Dict, List, Literal, Optional, Union from pydantic import ( ConfigDict, @@ -27,11 +29,16 @@ ) from zenml.config.base_settings import BaseSettings +from zenml.enums import ResourceRequestReclaimTolerance, StepRuntime + +if TYPE_CHECKING: + from zenml.models import ResourceRequestDemand class ByteUnit(Enum): """Enum for byte units.""" + B = "B" KB = "KB" KIB = "KiB" MB = "MB" @@ -51,6 +58,7 @@ def byte_value(self) -> int: The byte value of this unit. """ return { + ByteUnit.B: 1, ByteUnit.KB: 10**3, ByteUnit.KIB: 1 << 10, ByteUnit.MB: 10**6, @@ -65,6 +73,59 @@ def byte_value(self) -> int: MEMORY_REGEX = r"^[0-9]+(" + "|".join(unit.value for unit in ByteUnit) + r")$" +BASIC_RESOURCE_KINDS = frozenset({"cpu", "gpu", "memory"}) +DEFAULT_ALLOCATION_WAIT_TIMEOUT_SECONDS = 3600 +DEFAULT_INITIALIZATION_LEASE_SECONDS = 3600 + +ResourceInput = Union[ + str, + Dict[str, PositiveInt], + "PoolResourceDemand", + List["PoolResourceDemand"], +] +"""Flexible resource pool demand input accepted by ``ResourceSettings.resources``. + +May be: + +* a resource name string (quantity defaults to 1), +* a name-to-quantity map (for example ``{"tpu": 2}``), +* one ``PoolResourceDemand``, +* a list of demands or demand dictionaries. +""" + + +class PoolResourceDemand(BaseSettings): + """Resource pool demand used by ResourceSettings. + + Attributes: + quantity: Requested quantity. + name: Optional resource descriptor name. + kind: Optional resource kind (for example ``gpu``). + unit: Optional unit for the requested quantity. + class_name: Optional exact capacity class name. + resource_selector: Optional selector over resource descriptor fields and + attributes. + class_selector: Optional selector over capacity class fields and + attributes. + """ + + quantity: PositiveInt + name: Optional[str] = Field(default=None, min_length=1) + kind: Optional[str] = Field(default=None, min_length=1) + unit: Optional[str] = Field(default=None, min_length=1) + class_name: Optional[str] = Field( + default=None, + alias="class", + serialization_alias="class", + ) + resource_selector: Optional[Dict[str, Any]] = None + class_selector: Optional[Dict[str, Any]] = None + + model_config = ConfigDict( + populate_by_name=True, + # prevent extra attributes during model initialization + extra="ignore", + ) class ResourceSettings(BaseSettings): @@ -112,9 +173,21 @@ class ResourceSettings(BaseSettings): Attributes: cpu_count: The amount of CPU cores that should be configured. gpu_count: The amount of GPUs that should be configured. + gpu_class: Optional GPU capacity class for resource pool scheduling. memory: The amount of memory that should be configured. - preemptible: Whether the resources can be preempted. This only applies - when using ZenML resource pools. + resources: Resource pool demands (``ResourceInput``). After validation + this is a list of ``PoolResourceDemand``. Legacy ``pool_resources`` + and ``pool_resource_demands`` inputs are still accepted and folded + into this field. + reclaim_tolerance: Capacity reclaim behavior tolerated by resource pool + requests. Defaults to ``any`` when not set. Legacy + ``preemptible`` booleans are still accepted on input and mapped + into this field. + allocation_wait_timeout_seconds: Maximum time in seconds to wait for + the resource manager to allocate resources before giving up. + initialization_lease_seconds: Maximum time in seconds to keep allocated + resources reserved while the workload is starting but not yet + running. min_replicas: Minimum number of container instances (replicas). Use 0 to allow scale-to-zero on idle. Only relevant to deployed pipelines. @@ -130,27 +203,36 @@ class ResourceSettings(BaseSettings): max_concurrency: Maximum concurrent requests per instance (if supported by the platform). Defines a concurrency limit for each container. Only relevant to deployed pipelines. - pool_resources: Optional map of resource name to positive integer amount - for ZenML resource pools. Use for custom keys - (for example ``tpu``) or to supply ``gpu`` / ``mcpu`` / ``memory_mb`` - without the typed fields. When ``gpu_count``, ``cpu_count``, or - ``memory`` is set, those fields override the same keys in this map. """ cpu_count: Optional[PositiveFloat] = None gpu_count: Optional[NonNegativeInt] = None + gpu_class: Optional[str] = None memory: Optional[str] = Field(pattern=MEMORY_REGEX, default=None) - pool_resources: Optional[Dict[str, PositiveInt]] = Field( - default=None, + resources: ResourceInput = Field( + default_factory=list, + description=( + "Resource pool demands. Accepts a resource name, a name-to-quantity " + "map, one PoolResourceDemand, or a list of demands." + ), + ) + reclaim_tolerance: ResourceRequestReclaimTolerance = ( + ResourceRequestReclaimTolerance.ANY + ) + allocation_wait_timeout_seconds: PositiveInt = Field( + default=DEFAULT_ALLOCATION_WAIT_TIMEOUT_SECONDS, + description=( + "Maximum time in seconds to wait for the resource manager to " + "allocate resources before giving up." + ), + ) + initialization_lease_seconds: PositiveInt = Field( + default=DEFAULT_INITIALIZATION_LEASE_SECONDS, description=( - "Maps resource names to positive integer amounts allocated from " - "ZenML resource pools. Examples: {'memory': 2}, {'tpu': 4}, " - "{'vcpus': 2}, {'gpus': 4}. When gpu_count, cpu_count, or memory " - "is set, those fields override the gpu, mcpu, and memory_mb keys " - "respectively." + "Maximum time in seconds to keep allocated resources reserved " + "while the workload is starting but not yet running." ), ) - preemptible: bool = True # Settings only applicable for deployers and deployed pipelines min_replicas: Optional[NonNegativeInt] = None @@ -161,6 +243,59 @@ class ResourceSettings(BaseSettings): autoscaling_target: Optional[PositiveFloat] = None max_concurrency: Optional[PositiveInt] = None + @model_validator(mode="before") + @classmethod + def normalize_input(cls, data: Any) -> Any: + """Normalize legacy and flexible resource settings input. + + Args: + data: Raw input data supplied to the model. + + Returns: + Input data with ``resources`` and ``reclaim_tolerance`` populated. + """ + if not isinstance(data, dict): + return data + + normalized = dict(data) + + if normalized.get("reclaim_tolerance") is None: + preemptible = normalized.get("preemptible") + if preemptible is True: + normalized["reclaim_tolerance"] = ( + ResourceRequestReclaimTolerance.COORDINATED + ) + elif preemptible is False: + normalized["reclaim_tolerance"] = ( + ResourceRequestReclaimTolerance.NONE + ) + + if normalized.get("resources") is None: + if normalized.get("pool_resources") is not None: + normalized["resources"] = normalized["pool_resources"] + + normalized.pop("preemptible", None) + normalized.pop("pool_resources", None) + return normalized + + @property + def normalized_resources(self) -> List[PoolResourceDemand]: + """Normalized resources field. + + Returns: + List of normalized resource demands. + """ + if isinstance(self.resources, str): + return [PoolResourceDemand(name=self.resources, quantity=1)] + if isinstance(self.resources, PoolResourceDemand): + return [self.resources] + if isinstance(self.resources, dict): + return [ + PoolResourceDemand(name=name, quantity=quantity) + for name, quantity in self.resources.items() + ] + return self.resources + @model_validator(mode="after") def validate_replicas(self) -> "ResourceSettings": """Validate that min_replicas is not greater than max_replicas. @@ -188,32 +323,33 @@ def empty(self) -> bool: """Returns if this object is "empty" or not. A ResourceSettings instance is considered empty if none of the - generic resource-related values are configured. This excludes the - preemptible flag and the pool_resources map, which are only relevant - for resource pool scheduling and currently ignored by workload - scheduling stack components like orchestrators and step operators. + generic resource-related values are configured. This excludes pool + ``resources`` and ``reclaim_tolerance``, which are only relevant for + resource pool scheduling and currently ignored by workload scheduling + stack components like orchestrators and step operators. Returns: `True` if no values were configured, `False` otherwise. """ - # To detect whether this config is empty (= no values specified), we - # check if there are any attributes which are explicitly set to any - # value other than `None`. Exclude preemptible and pool_resources - # because they are not part of the set of unified resource settings - # that are understood and applied by orchestrators and step operators. return ( len( self.model_dump( exclude_unset=True, exclude_none=True, - exclude={"preemptible", "pool_resources"}, + exclude={ + "resources", + "reclaim_tolerance", + "gpu_class", + "allocation_wait_timeout_seconds", + "initialization_lease_seconds", + }, ) ) == 0 ) def get_memory( - self, unit: Union[str, ByteUnit] = ByteUnit.GB + self, unit: Union[str, "ByteUnit"] = ByteUnit.GB ) -> Optional[float]: """Gets the memory configuration in a specific unit. @@ -234,7 +370,9 @@ def get_memory( unit = ByteUnit(unit) memory = self.memory - for memory_unit in ByteUnit: + for memory_unit in sorted( + ByteUnit, key=lambda unit: len(unit.value), reverse=True + ): if memory.endswith(memory_unit.value): memory_value = int(memory[: -len(memory_unit.value)]) return memory_value * memory_unit.byte_value / unit.byte_value @@ -242,30 +380,123 @@ def get_memory( # Should never happen due to the regex validation raise ValueError(f"Unable to parse memory unit from '{memory}'.") - def merged_requested_resources(self) -> Dict[str, int]: - """Build the resource request map for scheduling and pool allocation. + def memory_quantity_and_unit(self) -> Optional[tuple[int, ByteUnit]]: + """Parse the configured memory string into quantity and unit. Returns: - Resource name to amount. + Quantity and ``ByteUnit`` parsed from ``memory``, or None when + memory is not configured. + + Raises: + ValueError: If ``memory`` is set but does not end with a known + byte unit suffix. """ - merged: Dict[str, int] = ( - dict(self.pool_resources) if self.pool_resources else {} - ) + if not self.memory: + return None + + for memory_unit in sorted( + ByteUnit, key=lambda unit: len(unit.value), reverse=True + ): + if self.memory.endswith(memory_unit.value): + quantity = int(self.memory[: -len(memory_unit.value)]) + return quantity, memory_unit + + raise ValueError(f"Unable to parse memory unit from '{self.memory}'.") + + def merged_resource_demands(self) -> list["ResourceRequestDemand"]: + """Build canonical Resource Manager demands for pool allocation. - if self.gpu_count is not None: - if self.gpu_count > 0: - merged["gpu"] = self.gpu_count - else: - merged.pop("gpu", None) + Pool ``resources`` and typed CPU/GPU/memory fields are converted into + separate demands without merging or overriding. + + Returns: + Resource demand entries for pool scheduling. + """ + from zenml.models import ResourceRequestDemand + + demands = [ + ResourceRequestDemand( + resource=demand.name, + quantity=demand.quantity, + kind=demand.kind, + unit=demand.unit, + class_name=demand.class_name, + resource_selector=demand.resource_selector, + class_selector=demand.class_selector, + ) + for demand in self.normalized_resources + ] + + if self.gpu_count is not None and self.gpu_count > 0: + demands.append( + ResourceRequestDemand( + kind="gpu", + quantity=self.gpu_count, + class_name=self.gpu_class, + ) + ) if self.cpu_count is not None: - merged["mcpu"] = math.ceil(self.cpu_count * 1000) + demands.append( + ResourceRequestDemand( + kind="cpu", + quantity=math.ceil(self.cpu_count), + unit="CPU", + ) + ) + + memory_spec = self.memory_quantity_and_unit() + if memory_spec is not None: + memory_quantity, memory_unit = memory_spec + demands.append( + ResourceRequestDemand( + kind="memory", + quantity=memory_quantity, + unit=memory_unit.value, + ) + ) + + return demands + + @property + def has_basic_resource_demands(self) -> bool: + """Check whether settings request ZenML-known basic resources. + + Returns: + Whether the generated demands include a basic resource kind. + """ + return any( + demand.kind in BASIC_RESOURCE_KINDS + for demand in self.merged_resource_demands() + ) + + @property + def reclaim_tolerance_explicitly_set(self) -> bool: + """Check whether reclaim tolerance was provided by the user. + + Returns: + Whether the field was explicitly set, including through the legacy + ``preemptible`` input. + """ + return "reclaim_tolerance" in self.model_fields_set + + def effective_reclaim_tolerance( + self, runtime: StepRuntime + ) -> ResourceRequestReclaimTolerance: + """Reclaim tolerance used when creating resource pool requests. - memory_amount = self.get_memory(unit=ByteUnit.MB) - if memory_amount is not None: - merged["memory_mb"] = math.ceil(memory_amount) + Args: + runtime: Resolved dynamic step runtime. - return merged + Returns: + The configured reclaim tolerance when set by the user; otherwise + ``none`` for inline steps and ``any`` for isolated steps. + """ + if self.reclaim_tolerance_explicitly_set: + return self.reclaim_tolerance + if runtime == StepRuntime.INLINE: + return ResourceRequestReclaimTolerance.NONE + return ResourceRequestReclaimTolerance.ANY model_config = ConfigDict( # public attributes are immutable diff --git a/src/zenml/constants.py b/src/zenml/constants.py index cfdb1cd3477..b090bdf02da 100644 --- a/src/zenml/constants.py +++ b/src/zenml/constants.py @@ -506,8 +506,6 @@ def handle_float_env_var(var: str, default: float = 0.0) -> float: PIPELINE_SNAPSHOTS = "/pipeline_snapshots" PIPELINES = "/pipelines" PIPELINE_SPEC = "/pipeline-spec" -RESOURCE_POOLS = "/resource_pools" -RESOURCE_POOL_SUBJECT_POLICIES = "/resource_pool_subject_policies" RESOURCE_REQUESTS = "/resource_requests" PROJECTS = "/projects" REFRESH = "/refresh" diff --git a/src/zenml/enums.py b/src/zenml/enums.py index 831de6c97fc..c5ee0508e24 100644 --- a/src/zenml/enums.py +++ b/src/zenml/enums.py @@ -669,6 +669,26 @@ class ResourceRequestStatus(StrEnum): CANCELLED = "cancelled" REJECTED = "rejected" RELEASED = "released" + EXPIRED = "expired" + NO_MATCHING_POOL = "no_matching_pool" + + +class ResourceRequestRuntimeState(StrEnum): + """Owner-reported runtime states for resource requests.""" + + UNKNOWN = "unknown" + PENDING = "pending" + SUBMITTED = "submitted" + RUNNING = "running" + IDLE = "idle" + + +class ResourceRequestReclaimTolerance(StrEnum): + """Caller tolerance for reclaiming lower-priority work on a request.""" + + NONE = "none" + COORDINATED = "coordinated" + ANY = "any" class TriggerType(StrEnum): diff --git a/src/zenml/execution/pipeline/dynamic/compilation.py b/src/zenml/execution/pipeline/dynamic/compilation.py index 0e9a56cc948..dab0ef94227 100644 --- a/src/zenml/execution/pipeline/dynamic/compilation.py +++ b/src/zenml/execution/pipeline/dynamic/compilation.py @@ -320,12 +320,12 @@ def get_step_runtime( if not orchestrator.can_run_isolated_steps: return StepRuntime.INLINE - runtime = step_config.runtime + if step_config.resource_settings.has_basic_resource_demands: + return StepRuntime.ISOLATED + runtime = step_config.runtime if runtime is None: - if not step_config.resource_settings.empty: - runtime = StepRuntime.ISOLATED - elif step_config.docker_settings != pipeline_docker_settings: + if step_config.docker_settings != pipeline_docker_settings: runtime = StepRuntime.ISOLATED else: runtime = StepRuntime.INLINE diff --git a/src/zenml/integrations/kubernetes/kube_utils.py b/src/zenml/integrations/kubernetes/kube_utils.py index c14ee5cec27..1130bd298be 100644 --- a/src/zenml/integrations/kubernetes/kube_utils.py +++ b/src/zenml/integrations/kubernetes/kube_utils.py @@ -34,6 +34,7 @@ import enum import functools import json +import math import re import time from collections import defaultdict @@ -46,6 +47,7 @@ List, Optional, Tuple, + Type, TypeVar, Union, cast, @@ -54,6 +56,7 @@ from kubernetes import client as k8s_client from kubernetes import config as k8s_config from kubernetes.client.rest import ApiException +from pydantic import ValidationError from urllib3.exceptions import ReadTimeoutError from zenml.config.resource_settings import ByteUnit @@ -74,11 +77,14 @@ from zenml.utils.time_utils import utc_now if TYPE_CHECKING: + from zenml.config.base_settings import BaseSettings from zenml.config.resource_settings import ResourceSettings + from zenml.models import ResourceRequestResponse logger = get_logger(__name__) R = TypeVar("R") +SettingsT = TypeVar("SettingsT", bound="BaseSettings") # This is to fix a bug in the kubernetes client which has some wrong @@ -1540,6 +1546,215 @@ def apply_default_resource_requests( return pod_settings +def apply_resource_request_component_settings( + settings: SettingsT, + allocated_resource_request: Optional["ResourceRequestResponse"], + settings_class: Type[SettingsT], +) -> SettingsT: + """Apply matching request target settings to stack settings. + + Args: + settings: The stack component settings to update. + allocated_resource_request: The allocated resource request, if any. + settings_class: The settings class used to validate the result. + + Returns: + The updated and validated stack component settings. + + Raises: + ValueError: If the allocated resource request is invalid. + """ + if not allocated_resource_request: + return settings + + component_settings = ( + allocated_resource_request.get_resources().component_settings + ) + + settings_dict = settings.model_dump(exclude_unset=True) + settings_dict.update(component_settings) + + try: + return settings_class.model_validate(settings_dict) + except (ValidationError, TypeError) as e: + resource_pool_name = ( + allocated_resource_request.get_body().pool_name or "" + ) + component_settings_json = json.dumps( + component_settings, + default=str, + indent=2, + sort_keys=True, + ) + raise ValueError( + "Failed to apply the Resource Pool component setting overrides " + "from the allocated resource request " + f"{allocated_resource_request.id} " + f"to `{settings_class.__name__}` stack component settings. " + "The overrides were merged with the existing stack component " + "settings, but the merged configuration is invalid. " + "Component settings provided by the resource request:\n" + f"{component_settings_json}\n" + "Please check the target settings that you configured for the " + f"`{resource_pool_name}` resource pool and its policies." + ) from e + + +def apply_resource_request_allocations_to_pod_settings( + allocated_resource_request: Optional["ResourceRequestResponse"], + pod_settings: Optional[KubernetesPodSettings] = None, +) -> KubernetesPodSettings: + """Apply allocated CPU, memory, and GPU resources to pod settings. + + Args: + allocated_resource_request: The allocated resource request, if any. + pod_settings: The pod settings to update. A new one will be created + if not provided. + + Returns: + The new or updated pod settings. + """ + if not pod_settings: + pod_settings = KubernetesPodSettings() + else: + pod_settings = pod_settings.model_copy(deep=True) + + if not allocated_resource_request: + return pod_settings + + cpu_millicores = 0 + memory_bytes = 0 + gpu_count = 0 + + for allocation in allocated_resource_request.get_resources().allocations: + kind = allocation.resource_kind + unit = allocation.unit + + if kind == "cpu": + cpu_millicores += _cpu_allocation_to_millicores( + quantity=allocation.quantity, + unit=unit, + ) + elif kind == "memory": + memory_bytes += _memory_allocation_to_bytes( + quantity=allocation.quantity, + unit=unit, + ) + elif kind == "gpu": + gpu_count += allocation.quantity + + resources = { + section: dict(values) + for section, values in (pod_settings.resources or {}).items() + } + requests = resources.setdefault("requests", {}) + limits = resources.setdefault("limits", {}) + resource_request_resources: Dict[str, Dict[str, str]] = { + "requests": {}, + "limits": {}, + } + resource_request_requests = resource_request_resources["requests"] + resource_request_limits = resource_request_resources["limits"] + + if cpu_millicores: + if cpu_millicores % 1000 == 0: + cpu = str(cpu_millicores // 1000) + else: + cpu = f"{cpu_millicores}m" + requests["cpu"] = cpu + limits["cpu"] = cpu + resource_request_requests["cpu"] = cpu + resource_request_limits["cpu"] = cpu + + if memory_bytes: + memory = f"{math.ceil(memory_bytes / ByteUnit.MIB.byte_value)}Mi" + requests["memory"] = memory + limits["memory"] = memory + resource_request_requests["memory"] = memory + resource_request_limits["memory"] = memory + + if gpu_count: + gpu = str(gpu_count) + requests["nvidia.com/gpu"] = gpu + limits["nvidia.com/gpu"] = gpu + resource_request_requests["nvidia.com/gpu"] = gpu + resource_request_limits["nvidia.com/gpu"] = gpu + + logger.debug( + "Configured Kubernetes pod resources from allocated resource request " + "`%s`: %s. Final Kubernetes pod resources: %s", + allocated_resource_request.id, + json.dumps( + resource_request_resources, + default=str, + indent=2, + sort_keys=True, + ), + json.dumps( + resources, + default=str, + indent=2, + sort_keys=True, + ), + ) + + return pod_settings.model_copy(update={"resources": resources}) + + +def _cpu_allocation_to_millicores(quantity: int, unit: Optional[str]) -> int: + """Convert a CPU allocation to Kubernetes millicores. + + Args: + quantity: CPU allocation quantity. + unit: CPU allocation unit. + + Returns: + CPU allocation expressed in Kubernetes millicores. + """ + if unit is None: + if quantity > 50: + return quantity + return quantity * 1000 + + normalized_unit = unit.lower() + if normalized_unit == "cpu": + return quantity * 1000 + if normalized_unit in {"m", "mcpu", "millicpu", "millicore", "millicores"}: + return quantity + + logger.warning( + "Interpreting unsupported CPU allocation unit `%s` as full CPU cores.", + unit, + ) + return quantity * 1000 + + +def _memory_allocation_to_bytes(quantity: int, unit: Optional[str]) -> int: + """Convert a memory allocation to bytes. + + Args: + quantity: Memory allocation quantity. + unit: Memory allocation unit. + + Returns: + Memory allocation expressed in bytes. + """ + if unit is None: + logger.warning("Interpreting memory allocation without a unit as MiB.") + return quantity * ByteUnit.MIB.byte_value + + try: + byte_unit = ByteUnit(unit) + except ValueError: + logger.warning( + "Ignoring memory allocation with unsupported unit `%s`.", + unit, + ) + return 0 + + return quantity * byte_unit.byte_value + + # ============================================================================ # Resource Conversion Utilities # ============================================================================ diff --git a/src/zenml/integrations/kubernetes/orchestrators/kubernetes_orchestrator.py b/src/zenml/integrations/kubernetes/orchestrators/kubernetes_orchestrator.py index 1122ac14814..16d12301b2b 100644 --- a/src/zenml/integrations/kubernetes/orchestrators/kubernetes_orchestrator.py +++ b/src/zenml/integrations/kubernetes/orchestrators/kubernetes_orchestrator.py @@ -107,6 +107,7 @@ PipelineRunResponse, PipelineSnapshotBase, PipelineSnapshotResponse, + ResourceRequestResponse, ScheduleResponse, StepRunResponse, ) @@ -253,6 +254,15 @@ def settings_class(self) -> Optional[Type["BaseSettings"]]: """ return KubernetesOrchestratorSettings + @property + def supports_resource_pool_allocation(self) -> bool: + """Whether the orchestrator supports resource pool allocations. + + Returns: + Whether the orchestrator supports resource pool allocations. + """ + return True + def get_kubernetes_contexts(self) -> Tuple[List[str], str]: """Get list of configured Kubernetes contexts and the active context. @@ -880,8 +890,11 @@ def _wait_for_run_to_finish() -> None: f"{body.get('message', '')}" ) - def submit_isolated_step( - self, step_run_info: "StepRunInfo", environment: Dict[str, str] + def submit_isolated_step_with_allocation( + self, + step_run_info: "StepRunInfo", + environment: Dict[str, str], + allocated_resource_request: Optional["ResourceRequestResponse"], ) -> None: """Submit an isolated step. @@ -889,6 +902,8 @@ def submit_isolated_step( step_run_info: The step run information. environment: The environment variables to set in the execution environment. + allocated_resource_request: The allocated resource request for the + step, if any. """ logger.info( "Launching job for step `%s`.", @@ -898,6 +913,11 @@ def submit_isolated_step( settings = cast( KubernetesOrchestratorSettings, self.get_settings(step_run_info) ) + settings = kube_utils.apply_resource_request_component_settings( + settings=settings, + allocated_resource_request=allocated_resource_request, + settings_class=KubernetesOrchestratorSettings, + ) image = step_run_info.get_image(key=ORCHESTRATOR_DOCKER_IMAGE_KEY) command, args = orchestrator_utils.get_step_entrypoint_command( invocation_id=step_run_info.pipeline_step_name, @@ -937,6 +957,13 @@ def submit_isolated_step( step_name=step_run_info.pipeline_step_name, ) + pod_settings = ( + kube_utils.apply_resource_request_allocations_to_pod_settings( + allocated_resource_request=allocated_resource_request, + pod_settings=settings.pod_settings, + ) + ) + job_manifest = self._prepare_job_manifest( name=job_name, command=command, @@ -946,7 +973,7 @@ def submit_isolated_step( labels=labels, annotations=annotations, settings=settings, - pod_settings=settings.pod_settings, + pod_settings=pod_settings, # In the dynamic pipeline case, we can't handle retries at the # orchestrator level because the entrypoint args contain a step # run ID. @@ -980,6 +1007,11 @@ def submit_isolated_step( str(e), ) + logger.debug( + "Launched job for step `%s`.", + step_run_info.pipeline_step_name, + ) + def get_isolated_step_status( self, step_run: "StepRunResponse" ) -> ExecutionStatus: diff --git a/src/zenml/integrations/kubernetes/step_operators/kubernetes_step_operator.py b/src/zenml/integrations/kubernetes/step_operators/kubernetes_step_operator.py index b0997568568..0743282c44a 100644 --- a/src/zenml/integrations/kubernetes/step_operators/kubernetes_step_operator.py +++ b/src/zenml/integrations/kubernetes/step_operators/kubernetes_step_operator.py @@ -14,7 +14,15 @@ """Kubernetes step operator implementation.""" import random -from typing import TYPE_CHECKING, Dict, List, Optional, Tuple, Type, cast +from typing import ( + TYPE_CHECKING, + Dict, + List, + Optional, + Tuple, + Type, + cast, +) from kubernetes import client as k8s_client @@ -41,7 +49,11 @@ if TYPE_CHECKING: from zenml.config.step_run_info import StepRunInfo - from zenml.models import PipelineSnapshotBase, StepRunResponse + from zenml.models import ( + PipelineSnapshotBase, + ResourceRequestResponse, + StepRunResponse, + ) logger = get_logger(__name__) @@ -72,6 +84,15 @@ def settings_class(self) -> Optional[Type["BaseSettings"]]: """ return KubernetesStepOperatorSettings + @property + def supports_resource_pool_allocation(self) -> bool: + """Whether the step operator supports resource pool allocations. + + Returns: + Whether the step operator supports resource pool allocations. + """ + return True + @property def validator(self) -> Optional[StackValidator]: """Validates the stack. @@ -88,7 +109,7 @@ def _validate_remote_components(stack: "Stack") -> Tuple[bool, str]: "needs to write files into the artifact store, but the " f"artifact store `{stack.artifact_store.name}` of the " "active stack is local. Please ensure that your stack " - "contains a remote artifact store when using the Vertex " + "contains a remote artifact store when using the Kubernetes " "step operator." ) @@ -191,11 +212,12 @@ def _k8s_batch_api(self) -> k8s_client.BatchV1Api: """ return k8s_client.BatchV1Api(self.get_kube_client()) - def submit( + def submit_with_allocation( self, info: "StepRunInfo", entrypoint_command: List[str], environment: Dict[str, str], + allocated_resource_request: Optional["ResourceRequestResponse"], ) -> None: """Submits a step run to Kubernetes. @@ -204,10 +226,17 @@ def submit( entrypoint_command: Command that executes the step. environment: Environment variables to set in the step operator environment. + allocated_resource_request: The allocated resource request for the + step, if any. """ settings = cast( KubernetesStepOperatorSettings, self.get_settings(info) ) + settings = kube_utils.apply_resource_request_component_settings( + settings=settings, + allocated_resource_request=allocated_resource_request, + settings_class=KubernetesStepOperatorSettings, + ) image_name = info.get_image( key=KUBERNETES_STEP_OPERATOR_DOCKER_IMAGE_KEY ) @@ -238,9 +267,15 @@ def submit( # some memory resources itself and, if not specified, the pod will be # scheduled on any node regardless of available memory and risk # negatively impacting or even crashing the node due to memory pressure. + pod_settings = ( + kube_utils.apply_resource_request_allocations_to_pod_settings( + allocated_resource_request=allocated_resource_request, + pod_settings=settings.pod_settings, + ) + ) pod_settings = kube_utils.apply_default_resource_requests( memory="400Mi", - pod_settings=settings.pod_settings, + pod_settings=pod_settings, ) pod_manifest = build_pod_manifest( diff --git a/src/zenml/models/__init__.py b/src/zenml/models/__init__.py index c354691aca7..d94d212ec85 100644 --- a/src/zenml/models/__init__.py +++ b/src/zenml/models/__init__.py @@ -16,6 +16,7 @@ # ------------------------------------- V2 ------------------------------------- # V2 Base +from zenml.enums import ResourceRequestReclaimTolerance from zenml.models.v2.base.base import ( BaseDatedResponseBody, BaseIdentifiedResponse, @@ -255,33 +256,19 @@ ProjectResponseMetadata, ProjectUpdate, ) -from zenml.models.v2.core.resource_pool import ( +from zenml.models.v2.core.resource_request import ( ResourcePoolAllocation, - ResourcePoolFilter, + ResourcePoolCapacityComponentSettings, ResourcePoolQueueItem, - ResourcePoolRequest, - ResourcePoolResponse, - ResourcePoolResponseBody, - ResourcePoolResponseMetadata, - ResourcePoolResponseResources, - ResourcePoolUpdate, -) -from zenml.models.v2.core.resource_pool_subject_policy import ( - ResourcePoolSubjectPolicyFilter, - ResourcePoolSubjectPolicyRequest, - ResourcePoolSubjectPolicyResponse, - ResourcePoolSubjectPolicyResponseBody, - ResourcePoolSubjectPolicyResponseMetadata, - ResourcePoolSubjectPolicyResponseResources, - ResourcePoolSubjectPolicyUpdate, -) -from zenml.models.v2.core.resource_request import ( + ResourceRequestDemand, ResourceRequestFilter, + ResourceRequestRenewalRequest, ResourceRequestRequest, ResourceRequestResponse, ResourceRequestResponseBody, ResourceRequestResponseMetadata, ResourceRequestResponseResources, + ResourceRequestServiceConnectorSettings, ) from zenml.models.v2.core.run_metadata import ( RunMetadataRequest, @@ -371,6 +358,7 @@ StackUpdate, ) from zenml.models.v2.core.step_run import ( + StepHeartbeatRequest, StepHeartbeatResponse, StepRunFilter, StepRunRequest, @@ -381,9 +369,9 @@ StepRunUpdate, ) from zenml.models.v2.core.stream_event import ( - StreamEvent, StreamBatchRequest, StreamBatchResponse, + StreamEvent, ) from zenml.models.v2.core.tag import ( TagFilter, @@ -615,21 +603,12 @@ ServiceConnectorInfo.model_rebuild() ServiceConnectorResourcesInfo.model_rebuild() ResourcesInfo.model_rebuild() -ResourcePoolResponseBody.model_rebuild() -ResourcePoolResponseMetadata.model_rebuild() -ResourcePoolResponseResources.model_rebuild() -ResourcePoolResponse.model_rebuild() ResourceRequestResponseBody.model_rebuild() ResourceRequestResponseMetadata.model_rebuild() ResourceRequestResponseResources.model_rebuild() ResourceRequestResponse.model_rebuild() ResourcePoolAllocation.model_rebuild() ResourcePoolQueueItem.model_rebuild() -ResourcePoolSubjectPolicyRequest.model_rebuild() -ResourcePoolSubjectPolicyResponseBody.model_rebuild() -ResourcePoolSubjectPolicyResponseMetadata.model_rebuild() -ResourcePoolSubjectPolicyResponseResources.model_rebuild() -ResourcePoolSubjectPolicyResponse.model_rebuild() TriggerRequest.model_rebuild() TriggerResponseBody.model_rebuild() TriggerResponseMetadata.model_rebuild() @@ -828,28 +807,19 @@ "PipelineRunResponseMetadata", "PipelineRunResponseResources", "PipelineRunTriggerInfo", - "ResourcePoolRequest", - "ResourcePoolUpdate", - "ResourcePoolFilter", - "ResourcePoolResponse", - "ResourcePoolResponseBody", - "ResourcePoolResponseMetadata", - "ResourcePoolResponseResources", + "ResourcePoolCapacityComponentSettings", "ResourcePoolAllocation", "ResourcePoolQueueItem", - "ResourcePoolSubjectPolicyFilter", - "ResourcePoolSubjectPolicyRequest", - "ResourcePoolSubjectPolicyResponse", - "ResourcePoolSubjectPolicyResponseBody", - "ResourcePoolSubjectPolicyResponseMetadata", - "ResourcePoolSubjectPolicyResponseResources", - "ResourcePoolSubjectPolicyUpdate", - "ResourceRequestRequest", + "ResourceRequestDemand", "ResourceRequestFilter", + "ResourceRequestRequest", + "ResourceRequestRenewalRequest", + "ResourceRequestReclaimTolerance", "ResourceRequestResponse", "ResourceRequestResponseBody", "ResourceRequestResponseMetadata", "ResourceRequestResponseResources", + "ResourceRequestServiceConnectorSettings", "RunTemplateRequest", "RunTemplateUpdate", "RunTemplateResponse", @@ -1012,6 +982,7 @@ "StepRunIdentifier", "ArtifactVersionIdentifier", "ModelVersionIdentifier", + "StepHeartbeatRequest", "StepHeartbeatResponse", "ScheduleTriggerRequest", "ScheduleTriggerResponse", diff --git a/src/zenml/models/v2/core/resource_pool.py b/src/zenml/models/v2/core/resource_pool.py deleted file mode 100644 index a57e2c5f099..00000000000 --- a/src/zenml/models/v2/core/resource_pool.py +++ /dev/null @@ -1,256 +0,0 @@ -# Copyright (c) ZenML GmbH 2026. All Rights Reserved. -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at: -# -# https://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express -# or implied. See the License for the specific language governing -# permissions and limitations under the License. -"""Models representing resource pools.""" - -from datetime import datetime -from typing import ( - TYPE_CHECKING, - ClassVar, - Dict, - List, - Optional, - TypeVar, -) - -from pydantic import ( - BaseModel, - Field, - NonNegativeInt, - PositiveInt, - model_validator, -) - -from zenml.constants import STR_FIELD_MAX_LENGTH -from zenml.models.v2.base.base import BaseUpdate -from zenml.models.v2.base.filter import StringFilterOption -from zenml.models.v2.base.scoped import ( - UserScopedFilter, - UserScopedRequest, - UserScopedResponse, - UserScopedResponseBody, - UserScopedResponseMetadata, - UserScopedResponseResources, -) - -if TYPE_CHECKING: - from zenml.models.v2.core.resource_request import ResourceRequestResponse - from zenml.zen_stores.schemas import BaseSchema - - AnySchema = TypeVar("AnySchema", bound=BaseSchema) - - -class ResourcePoolAllocation(BaseModel): - """Resource pool allocation.""" - - request: "ResourceRequestResponse" = Field( - title="The request that is allocated to the resource pool.", - ) - priority: Optional[int] = Field( - title="The priority of the component in the resource pool.", - ) - allocated_at: datetime = Field( - title="The time the resource pool was allocated.", - ) - - -class ResourcePoolQueueItem(BaseModel): - """Resource pool queue item.""" - - request: "ResourceRequestResponse" = Field( - title="The request that is queued for the resource pool.", - ) - priority: int = Field( - title="The priority of the request in the resource pool.", - ) - - -# ------------------ Request Model ------------------ - - -class ResourcePoolRequest(UserScopedRequest): - """Request model for resource pool creation.""" - - name: str = Field( - title="The name of the resource pool.", max_length=STR_FIELD_MAX_LENGTH - ) - description: Optional[str] = Field( - title="The description of the resource pool", - max_length=STR_FIELD_MAX_LENGTH, - default=None, - ) - capacity: Dict[str, PositiveInt] = Field( - title="The capacity of the resource pool.", - ) - - @model_validator(mode="after") - def _validate_capacity(self) -> "ResourcePoolRequest": - if not self.capacity: - raise ValueError( - "Resource pools with no capacity are not allowed." - ) - - return self - - -# ------------------ Update Model ------------------ - - -class ResourcePoolUpdate(BaseUpdate): - """Update model for resource pools.""" - - description: Optional[str] = Field( - title="The description of the resource pool", - max_length=STR_FIELD_MAX_LENGTH, - default=None, - ) - capacity: Optional[Dict[str, NonNegativeInt]] = Field( - title="The capacity of the resource pool.", - description="The capacity of the resource pool. Setting a value to 0 " - "will remove the resource from the pool.", - default=None, - ) - - -# ------------------ Response Model ------------------ - - -class ResourcePoolResponseBody(UserScopedResponseBody): - """Response body for resource pools.""" - - queue_length: NonNegativeInt = Field( - title="The number of queued requests for the resource pool.", - ) - capacity: Dict[str, int] = Field( - title="The capacity of the resource pool.", - ) - occupied_resources: Dict[str, int] = Field( - title="The occupied resources of the resource pool.", - ) - - -class ResourcePoolResponseMetadata(UserScopedResponseMetadata): - """Response metadata for resource pools.""" - - description: Optional[str] = Field( - default="", - title="The description of the resource pool", - max_length=STR_FIELD_MAX_LENGTH, - ) - - active_requests: List["ResourcePoolAllocation"] = Field( - title="The active allocations for the resource pool.", - default=[], - ) - queued_requests: List["ResourcePoolQueueItem"] = Field( - title="The queued requests for the resource pool.", - default=[], - ) - - -class ResourcePoolResponseResources(UserScopedResponseResources): - """Response resources for resource pools.""" - - -class ResourcePoolResponse( - UserScopedResponse[ - ResourcePoolResponseBody, - ResourcePoolResponseMetadata, - ResourcePoolResponseResources, - ] -): - """Response model for resource pools.""" - - name: str = Field( - title="The name of the resource pool.", max_length=STR_FIELD_MAX_LENGTH - ) - - def get_hydrated_version(self) -> "ResourcePoolResponse": - """Get the hydrated version of this resource pool. - - Returns: - The hydrated version of this resource pool. - """ - from zenml.client import Client - - return Client().zen_store.get_resource_pool(self.id) - - @property - def description(self) -> Optional[str]: - """The `description` property. - - Returns: - the value of the property. - """ - return self.get_metadata().description - - @property - def active_requests(self) -> List["ResourcePoolAllocation"]: - """The `active_requests` property. - - Returns: - the value of the property. - """ - return self.get_metadata().active_requests - - @property - def queued_requests(self) -> List["ResourcePoolQueueItem"]: - """The `queued_requests` property. - - Returns: - the value of the property. - """ - return self.get_metadata().queued_requests - - @property - def capacity(self) -> Dict[str, int]: - """The `capacity` property. - - Returns: - the value of the property. - """ - return self.get_body().capacity - - @property - def occupied_resources(self) -> Dict[str, int]: - """The `occupied_resources` property. - - Returns: - the value of the property. - """ - return self.get_body().occupied_resources - - @property - def queue_length(self) -> int: - """The `queue_length` property. - - Returns: - the value of the property. - """ - return self.get_body().queue_length - - -# ------------------ Filter Model ------------------ - - -class ResourcePoolFilter(UserScopedFilter): - """Resource pool filter model.""" - - FILTER_EXCLUDE_FIELDS: ClassVar[List[str]] = [ - *UserScopedFilter.FILTER_EXCLUDE_FIELDS, - ] - - name: StringFilterOption = Field( - default=None, - description="Name of the resource pool", - ) diff --git a/src/zenml/models/v2/core/resource_pool_subject_policy.py b/src/zenml/models/v2/core/resource_pool_subject_policy.py deleted file mode 100644 index 110e0fc744b..00000000000 --- a/src/zenml/models/v2/core/resource_pool_subject_policy.py +++ /dev/null @@ -1,234 +0,0 @@ -# Copyright (c) ZenML GmbH 2026. All Rights Reserved. -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at: -# -# https://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express -# or implied. See the License for the specific language governing -# permissions and limitations under the License. -"""Models representing resource pool subject policies.""" - -from typing import ( - TYPE_CHECKING, - Dict, - Optional, -) -from uuid import UUID - -from pydantic import Field, NonNegativeInt, model_validator - -from zenml.models.v2.base.base import BaseUpdate -from zenml.models.v2.base.filter import IntegerFilterOption, UUIDFilterOption -from zenml.models.v2.base.scoped import ( - UserScopedFilter, - UserScopedRequest, - UserScopedResponse, - UserScopedResponseBody, - UserScopedResponseMetadata, - UserScopedResponseResources, -) - -if TYPE_CHECKING: - from zenml.models import ComponentResponse - from zenml.models.v2.core.resource_pool import ResourcePoolResponse - - -# ------------------ Request Model ------------------ - - -class ResourcePoolSubjectPolicyRequest(UserScopedRequest): - """Resource pool subject policy request.""" - - component_id: UUID = Field( - title="The ID of the component that is the subject of the policy.", - ) - pool_id: UUID = Field( - title="The ID of the resource pool that the policy is attached to.", - ) - priority: NonNegativeInt = Field( - title="The priority of the policy in the resource pool. Higher " - "means preferred.", - ) - reserved: Optional[Dict[str, NonNegativeInt]] = Field( - title="The resources that are reserved for the policy.", - default=None, - ) - limit: Optional[Dict[str, NonNegativeInt]] = Field( - title="The maximum resources that the policy can use.", - default=None, - ) - - @model_validator(mode="after") - def _validate_resources(self) -> "ResourcePoolSubjectPolicyRequest": - if not self.reserved or not self.limit: - return self - - for key, reserved_value in self.reserved.items(): - limit_value = self.limit.get(key) - - if limit_value is not None and reserved_value > limit_value: - raise ValueError( - f"Reserved value for resource `{key}` ({reserved_value}) " - f"must be less than limit `{limit_value}`." - ) - - return self - - -# ------------------ Update Model ------------------ - - -class ResourcePoolSubjectPolicyUpdate(BaseUpdate): - """Update model for resource pool subject policies.""" - - priority: NonNegativeInt = Field( - title="The priority of the policy in the resource pool. Higher " - "means preferred.", - ) - reserved: Optional[Dict[str, NonNegativeInt]] = Field( - title="The resources that are reserved for the policy.", - default=None, - ) - limit: Optional[Dict[str, NonNegativeInt]] = Field( - title="The maximum resources that the policy can use.", - default=None, - ) - - @model_validator(mode="after") - def _validate_resources(self) -> "ResourcePoolSubjectPolicyUpdate": - if not self.reserved or not self.limit: - return self - - for key, reserved_value in self.reserved.items(): - limit_value = self.limit.get(key) - - if limit_value is not None and reserved_value > limit_value: - raise ValueError( - f"Reserved value for resource `{key}` ({reserved_value}) " - f"must be less than limit `{limit_value}`." - ) - - return self - - -# ------------------ Response Model ------------------ - - -class ResourcePoolSubjectPolicyResponseBody(UserScopedResponseBody): - """Response body for resource pool subject policies.""" - - priority: NonNegativeInt = Field( - title="The priority of the policy in the resource pool. Higher " - "means preferred.", - ) - reserved: Dict[str, NonNegativeInt] = Field( - title="The resources that are reserved for the policy.", - ) - limit: Dict[str, NonNegativeInt] = Field( - title="The maximum resources that the policy can use.", - ) - - -class ResourcePoolSubjectPolicyResponseMetadata(UserScopedResponseMetadata): - """Response metadata for resource pool subject policies.""" - - -class ResourcePoolSubjectPolicyResponseResources(UserScopedResponseResources): - """Response resources for resource pool subject policies.""" - - component: "ComponentResponse" = Field( - title="The component that is requesting the resources." - ) - pool: "ResourcePoolResponse" = Field( - title="The resource pool that the policy is attached to.", - ) - - -class ResourcePoolSubjectPolicyResponse( - UserScopedResponse[ - ResourcePoolSubjectPolicyResponseBody, - ResourcePoolSubjectPolicyResponseMetadata, - ResourcePoolSubjectPolicyResponseResources, - ] -): - """Response model for resource pool subject policies.""" - - def get_hydrated_version(self) -> "ResourcePoolSubjectPolicyResponse": - """Get the hydrated version of this resource pool subject policy. - - Returns: - The hydrated version of this resource pool subject policy. - """ - from zenml.client import Client - - return Client().zen_store.get_resource_pool_subject_policy(self.id) - - @property - def priority(self) -> int: - """The `priority` property. - - Returns: - the value of the property. - """ - return self.get_body().priority - - @property - def reserved(self) -> Dict[str, int]: - """The `reserved` property. - - Returns: - the value of the property. - """ - return self.get_body().reserved - - @property - def limit(self) -> Dict[str, int]: - """The `limit` property. - - Returns: - the value of the property. - """ - return self.get_body().limit - - @property - def component(self) -> "ComponentResponse": - """The `component` property. - - Returns: - the value of the property. - """ - return self.get_resources().component - - @property - def pool(self) -> "ResourcePoolResponse": - """The `pool` property. - - Returns: - the value of the property. - """ - return self.get_resources().pool - - -# ------------------ Filter Model ------------------ - - -class ResourcePoolSubjectPolicyFilter(UserScopedFilter): - """Resource pool subject policy filter model.""" - - pool_id: UUIDFilterOption = Field( - default=None, - description="The ID of the resource pool that the policy is attached to.", - ) - component_id: UUIDFilterOption = Field( - default=None, - description="The ID of the component that is the subject of the policy.", - ) - priority: IntegerFilterOption = Field( - default=None, - description="The priority of the policy in the resource pool.", - ) diff --git a/src/zenml/models/v2/core/resource_request.py b/src/zenml/models/v2/core/resource_request.py index 6b4fa996094..b1437f1e962 100644 --- a/src/zenml/models/v2/core/resource_request.py +++ b/src/zenml/models/v2/core/resource_request.py @@ -13,24 +13,23 @@ # permissions and limitations under the License. """Models representing resource requests.""" +from datetime import datetime from typing import ( - TYPE_CHECKING, - ClassVar, - Dict, - List, + Any, Optional, - Type, - TypeVar, ) from uuid import UUID -from pydantic import Field, PositiveInt, model_validator +from pydantic import ConfigDict, Field, PositiveInt, model_validator -from zenml.enums import ResourceRequestStatus -from zenml.models.v2.base.filter import ( - EnumFilterOption, - UUIDFilterOption, +from zenml.constants import STR_FIELD_MAX_LENGTH +from zenml.enums import ( + ResourceRequestReclaimTolerance, + ResourceRequestRuntimeState, + ResourceRequestStatus, + StackComponentType, ) +from zenml.models.v2.base.base import BaseZenModel from zenml.models.v2.base.scoped import ( UserScopedFilter, UserScopedRequest, @@ -40,68 +39,413 @@ UserScopedResponseResources, ) -if TYPE_CHECKING: - from sqlalchemy.sql.elements import ColumnElement - from zenml.models import ( - ComponentResponse, - PipelineRunResponse, - StepRunResponse, +class ResourcePoolCapacityComponentSettings(BaseZenModel): + """Stack component settings applied when a request allocation is granted.""" + + component_type: StackComponentType = Field( + title="The stack component type to apply settings to.", + ) + flavor: str = Field( + title="The stack component flavor to apply settings to.", + max_length=STR_FIELD_MAX_LENGTH, + ) + settings: dict[str, Any] = Field( + default_factory=dict, + title="The stack component settings to apply on allocation.", ) - from zenml.models.v2.core.resource_pool import ResourcePoolResponse - from zenml.zen_stores.schemas import BaseSchema - AnySchema = TypeVar("AnySchema", bound=BaseSchema) +class ResourceRequestServiceConnectorSettings(BaseZenModel): + """Service connector settings selected for a resource request.""" -# ------------------ Request Model ------------------ + connector_id: Optional[UUID] = Field( + default=None, + title="The service connector ID selected for the request.", + ) + resource_type: Optional[str] = Field( + default=None, + title="The service connector resource type selected for the request.", + min_length=1, + max_length=STR_FIELD_MAX_LENGTH, + ) + resource_id: Optional[str] = Field( + default=None, + title="The service connector resource ID selected for the request.", + min_length=1, + max_length=STR_FIELD_MAX_LENGTH, + ) -class ResourceRequestRequest(UserScopedRequest): - """Request model for resource requests.""" +class ResourcePoolQueueItem(BaseZenModel): + """Queue item linked to a resource request.""" + + id: UUID = Field(title="The unique queue item ID.") + request_id: UUID = Field(title="The queued resource request ID.") + pool_id: UUID = Field(title="The resource pool ID.") + pool_name: Optional[str] = Field( + default=None, + title="The resource pool name.", + ) + policy_id: UUID = Field(title="The resource policy ID.") + priority: int = Field(title="The priority snapshot for this queue item.") + priority_lane: bool = Field( + default=False, + title="Whether this queue item uses the priority lane.", + ) + enqueued_at: datetime = Field(title="The queue insertion timestamp.") + created: Optional[datetime] = Field( + default=None, + title="The queue entry creation timestamp.", + ) + - component_id: UUID = Field( - title="The id of the component that is requesting the resources.", +class ResourcePoolAllocation(BaseZenModel): + """Allocation grant linked to a resource request.""" + + id: UUID = Field(title="The unique allocation ID.") + request_id: UUID = Field(title="The allocated resource request ID.") + demand_index: Optional[int] = Field( + default=None, + ge=0, + title=( + "The zero-based request demand index this grant satisfies, or " + "None for grant-default allocations." + ), + ) + pool_id: UUID = Field(title="The resource pool ID.") + pool_name: Optional[str] = Field( + default=None, + title="The resource pool name.", + ) + capacity_entry_id: Optional[UUID] = Field( + default=None, + title="The selected resource pool capacity entry ID.", + ) + capacity_entry_name: Optional[str] = Field( + default=None, + title="The selected resource pool capacity entry name.", + max_length=STR_FIELD_MAX_LENGTH, + ) + resource_id: UUID = Field(title="The allocated resource descriptor ID.") + resource: Optional[str] = Field( + default=None, + title="The allocated resource descriptor name.", + ) + resource_kind: str = Field( + title="The allocated resource descriptor kind.", + min_length=1, + max_length=STR_FIELD_MAX_LENGTH, + ) + class_name: str = Field( + alias="class", + serialization_alias="class", + title="The allocated capacity class.", + ) + quantity: PositiveInt = Field(title="The allocated quantity.") + unit: Optional[str] = Field( + default=None, + title="The optional unit for the allocated quantity.", + min_length=1, + max_length=64, + ) + base_quantity: Optional[PositiveInt] = Field( + default=None, + title="The allocated quantity converted to the descriptor base unit.", + ) + policy_id: UUID = Field(title="The policy that admitted this allocation.") + grant_id: Optional[UUID] = Field( + default=None, + title="The policy grant that matched the demand, if any.", + ) + priority: int = Field(title="The priority snapshot for this allocation.") + priority_lane: bool = Field( + default=False, + title="Whether this allocation uses the priority lane.", + ) + component_id: Optional[UUID] = Field( + default=None, + title="The stack component ID selected for this allocation.", + ) + account_id: Optional[UUID] = Field( + default=None, + title="The external account ID selected for this allocation.", + ) + component_settings: list[ResourcePoolCapacityComponentSettings] = Field( + default_factory=list, + title="Stack component settings applied for this allocation.", ) - step_run_id: UUID = Field( - title="The id of the step run that is requesting the resources.", + preemption_state: str = Field(title="The preemption state.") + preemption_reason: Optional[str] = Field( + default=None, + title="The preemption reason.", ) - requested_resources: Dict[str, PositiveInt] = Field( - title="The resources requested." + released_at: Optional[datetime] = Field( + default=None, + title="The release timestamp.", ) - preemptible: bool = Field( - default=True, - title="Whether this request can be preempted.", + created: Optional[datetime] = Field( + default=None, + title="The allocation creation timestamp.", + ) + updated: Optional[datetime] = Field( + default=None, + title="The allocation last update timestamp.", ) + model_config = ConfigDict(populate_by_name=True) + @model_validator(mode="after") - def _validate_requested_resources(self) -> "ResourceRequestRequest": - if not self.requested_resources: + def _validate_subject_reference(self) -> "ResourcePoolAllocation": + """Validate that the allocation references a subject. + + Returns: + The validated allocation. + + Raises: + ValueError: If neither component ID nor account ID is set. + """ + if self.component_id is None and self.account_id is None: raise ValueError( - "Resource requests with no requested resources are not allowed." + "An allocation requires component_id or account_id." ) + return self + +class ResourceRequestDemand(BaseZenModel): + """Resource demand for a Resource Manager-backed request.""" + + resource_id: Optional[UUID] = Field( + default=None, + title="The exact resource descriptor ID.", + ) + resource: Optional[str] = Field( + default=None, + title="The exact resource descriptor name.", + ) + kind: Optional[str] = Field( + default=None, + title="The optional resource descriptor kind.", + min_length=1, + max_length=64, + ) + quantity: PositiveInt = Field(title="The resource quantity requested.") + unit: Optional[str] = Field( + default=None, + title="The optional unit for the requested quantity.", + min_length=1, + max_length=64, + ) + class_name: Optional[str] = Field( + default=None, + alias="class", + serialization_alias="class", + title="The exact capacity class.", + ) + resource_selector: Optional[dict[str, Any]] = Field( + default=None, + title="Selector over resource descriptor fields and attributes.", + ) + class_selector: Optional[dict[str, Any]] = Field( + default=None, + title="Selector over capacity class fields and attributes.", + ) + + model_config = ConfigDict(populate_by_name=True) + + @model_validator(mode="after") + def _validate_resource_reference(self) -> "ResourceRequestDemand": + """Validate that the demand can resolve to a resource. + + Returns: + The validated demand. + + Raises: + ValueError: If no resource reference or selector is configured. + """ + if ( + self.resource_id is None + and self.resource is None + and self.resource_selector is None + and self.kind is None + ): + raise ValueError( + "Resource demands require a resource ID, resource name, " + "resource kind, or resource selector." + ) return self -# ------------------ Response Model ------------------ +class ResourceRequestRequest(UserScopedRequest): + """Request model for creating resource requests.""" + + component_ids: Optional[list[UUID]] = Field( + default=None, + title="Stack components that may satisfy the request.", + ) + step_run_id: Optional[UUID] = Field( + default=None, + title="The step run that is requesting the resources.", + ) + demands: list[ResourceRequestDemand] = Field( + min_length=1, + title="The resource demands requested.", + ) + pool_id: Optional[UUID] = Field( + default=None, + title="The exact resource pool ID.", + ) + pool: Optional[str] = Field( + default=None, + title="The exact resource pool name.", + max_length=STR_FIELD_MAX_LENGTH, + ) + pool_selector: Optional[dict[str, Any]] = Field( + default=None, + title="Selector over pool attributes.", + ) + preemption_group: Optional[dict[str, Any]] = Field( + default=None, + title="Fallback selector for protected preemption peers.", + ) + reclaim_tolerance: ResourceRequestReclaimTolerance = Field( + default=ResourceRequestReclaimTolerance.ANY, + title="The capacity reclaim behavior tolerated by this request.", + ) + lease_expires_at: Optional[datetime] = Field( + default=None, + title="The optional initial lease expiration timestamp.", + ) + allocation_wait_timeout_seconds: Optional[int] = Field( + default=None, + title="Seconds to wait for future time-gated admission routes.", + ) + metadata: dict[str, Any] = Field( + default_factory=dict, + title="Optional opaque metadata attached to the request.", + ) + + @model_validator(mode="after") + def _validate_step_run_fields(self) -> "ResourceRequestRequest": + """Validate step-run component and step_run_id fields. + + Returns: + The validated resource request. + + Raises: + ValueError: If step-run fields are inconsistent. + """ + if self.component_ids and self.step_run_id is None: + raise ValueError( + "step_run_id is required when component_ids are set." + ) + if self.step_run_id is not None and not self.component_ids: + raise ValueError( + "component_ids are required when step_run_id is set." + ) + return self + + +class ResourceRequestRenewalRequest(BaseZenModel): + """Request model for renewing a resource request lease.""" + + lease_expires_at: datetime = Field( + title="The renewed lease expiration timestamp.", + ) + runtime_state: Optional[ResourceRequestRuntimeState] = Field( + default=None, + title="The optional owner-reported runtime state.", + ) class ResourceRequestResponseBody(UserScopedResponseBody): """Response body for resource requests.""" - requested_resources: Dict[str, int] = Field( - title="The resources requested." + component_ids: list[UUID] = Field( + default_factory=list, + title="Stack components associated with the request.", + ) + step_run_id: Optional[UUID] = Field( + default=None, + title="The step run associated with the resource request.", + ) + pipeline_run_id: Optional[UUID] = Field( + default=None, + title="The pipeline run associated with the resource request.", + ) + pool_id: Optional[UUID] = Field( + default=None, + title="The resource pool selected for the resource request.", + ) + step_name: Optional[str] = Field( + default=None, + title="The pipeline step name associated with the resource request.", + ) + pipeline_run_name: Optional[str] = Field( + default=None, + title="The pipeline run name associated with the resource request.", + ) + project_id: Optional[UUID] = Field( + default=None, + title="The project that owns the pipeline run for this request.", + ) + pool_name: Optional[str] = Field( + default=None, + title="The resource pool name selected for the resource request.", + ) + pool_selector: Optional[dict[str, Any]] = Field( + default=None, + title="Selector over pool attributes requested for the resource request.", + ) + preemption_group: Optional[dict[str, Any]] = Field( + default=None, + title="The request-level fallback preemption group selector.", + ) + demands: list[ResourceRequestDemand] = Field( + default_factory=list, + title="The resource demands requested.", ) status: ResourceRequestStatus = Field( title="The status of the resource request." ) + runtime_state: ResourceRequestRuntimeState = Field( + default=ResourceRequestRuntimeState.UNKNOWN, + title="The owner-reported runtime state of the request.", + ) + reclaim_tolerance: ResourceRequestReclaimTolerance = Field( + title="The capacity reclaim behavior tolerated by this request.", + ) + lease_expires_at: Optional[datetime] = Field( + default=None, + title="The optional lease expiration timestamp.", + ) + allocation_deadline: Optional[datetime] = Field( + default=None, + title="The deadline for waiting on future time-gated admission routes.", + ) + renewed_at: Optional[datetime] = Field( + default=None, + title="The optional lease renewal timestamp.", + ) + allocated_at: Optional[datetime] = Field( + default=None, + title="The timestamp when capacity was first granted.", + ) + released_at: Optional[datetime] = Field( + default=None, + title="The timestamp when the request was released.", + ) + queued_at: Optional[datetime] = Field( + default=None, + title="The timestamp when the request first entered a pool queue.", + ) status_reason: Optional[str] = Field( title="The reason for the status of the resource request.", default=None, ) - preemptible: bool = Field( - title="Whether this request can be preempted.", + preemption_initiated_by_id: Optional[UUID] = Field( + default=None, + title="The request that initiated preemption.", ) @@ -112,20 +456,23 @@ class ResourceRequestResponseMetadata(UserScopedResponseMetadata): class ResourceRequestResponseResources(UserScopedResponseResources): """Response resources for resource requests.""" - component: Optional["ComponentResponse"] = Field( - title="The component that is requesting the resources.", - default=None, + component_settings: dict[str, Any] = Field( + default_factory=dict, + title="Stack component settings selected for this request.", ) - step_run: Optional["StepRunResponse"] = Field( - title="The step run that is requesting the resources.", default=None - ) - pipeline_run: Optional["PipelineRunResponse"] = Field( - title="The pipeline run that is requesting the resources.", + service_connector_settings: Optional[ + ResourceRequestServiceConnectorSettings + ] = Field( default=None, + title="Service connector settings selected for this request.", ) - pool: Optional["ResourcePoolResponse"] = Field( - title="The pool that the resource request is/was running in.", - default=None, + allocations: list[ResourcePoolAllocation] = Field( + default_factory=list, + title="Allocation grants linked to this request.", + ) + queue_entries: list[ResourcePoolQueueItem] = Field( + default_factory=list, + title="Pool queue memberships linked to this request.", ) @@ -142,163 +489,171 @@ def get_hydrated_version(self) -> "ResourceRequestResponse": """Get the hydrated version of this resource request. Returns: - The hydrated version of this resource request. + The current resource request fetched from the configured ZenML + store. """ from zenml.client import Client return Client().zen_store.get_resource_request(self.id) @property - def requested_resources(self) -> Dict[str, int]: - """The `requested_resources` property. + def component_ids(self) -> list[UUID]: + """Resource request component IDs. Returns: - the value of the property. + Stack components associated with the resource request. """ - return self.get_body().requested_resources + return self.get_body().component_ids @property - def status(self) -> ResourceRequestStatus: - """The `status` property. + def step_run_id(self) -> Optional[UUID]: + """Resource request step run ID. Returns: - the value of the property. + The optional step run associated with the resource request. """ - return self.get_body().status + return self.get_body().step_run_id @property - def status_reason(self) -> Optional[str]: - """The `status_reason` property. + def pipeline_run_id(self) -> Optional[UUID]: + """Resource request pipeline run ID. Returns: - the value of the property. + The optional pipeline run associated with the resource request. """ - return self.get_body().status_reason + return self.get_body().pipeline_run_id + + @property + def pool_id(self) -> Optional[UUID]: + """Resource request pool ID. + + Returns: + The optional resource pool selected for the request. + """ + return self.get_body().pool_id + + @property + def demands(self) -> list[ResourceRequestDemand]: + """Requested resource demands. + + Returns: + The resource demands requested. + """ + return self.get_body().demands + + @property + def preemption_group(self) -> Optional[dict[str, Any]]: + """Request-level fallback preemption group selector. + + Returns: + The optional selector configured on the request. + """ + return self.get_body().preemption_group @property - def preemptible(self) -> bool: - """The `preemptible` property. + def reclaim_tolerance(self) -> ResourceRequestReclaimTolerance: + """Resource request reclaim tolerance. Returns: - the value of the property. + The reclaim behavior tolerated by the request. """ - return self.get_body().preemptible + return self.get_body().reclaim_tolerance @property - def component(self) -> Optional["ComponentResponse"]: - """The `component` property. + def lease_expires_at(self) -> Optional[datetime]: + """Resource request lease expiration timestamp. Returns: - the value of the property. + The optional lease expiration timestamp. """ - return self.get_resources().component + return self.get_body().lease_expires_at @property - def step_run(self) -> Optional["StepRunResponse"]: - """The `step_run` property. + def renewed_at(self) -> Optional[datetime]: + """Resource request renewal timestamp. Returns: - the value of the property. + The optional lease renewal timestamp. """ - return self.get_resources().step_run + return self.get_body().renewed_at @property - def pipeline_run(self) -> Optional["PipelineRunResponse"]: - """The `pipeline_run` property. + def status(self) -> ResourceRequestStatus: + """Resource request status. Returns: - the value of the property. + The lifecycle status of the resource request. """ - return self.get_resources().pipeline_run + return self.get_body().status @property - def pool(self) -> Optional["ResourcePoolResponse"]: - """The `pool` property. + def runtime_state(self) -> ResourceRequestRuntimeState: + """Resource request runtime state. Returns: - the value of the property. + The owner-reported runtime state of the resource request. """ - return self.get_resources().pool + return self.get_body().runtime_state + @property + def status_reason(self) -> Optional[str]: + """Resource request status reason. -# ------------------ Filter Model ------------------ + Returns: + The optional status reason. + """ + return self.get_body().status_reason + + @property + def preemption_initiated_by_id(self) -> Optional[UUID]: + """Request that initiated preemption. + + Returns: + The optional ID of the request that initiated preemption. + """ + return self.get_body().preemption_initiated_by_id class ResourceRequestFilter(UserScopedFilter): """Resource request filter model.""" - FILTER_EXCLUDE_FIELDS: ClassVar[List[str]] = [ - *UserScopedFilter.FILTER_EXCLUDE_FIELDS, - "pipeline_run_id", - ] - API_SINGLE_INPUT_PARAMS: ClassVar[List[str]] = [ - *UserScopedFilter.API_SINGLE_INPUT_PARAMS, - "preemptible", - ] - - preemptible: Optional[bool] = Field( + user: UUID | str | None = Field( default=None, - description="Whether the resource request is preemptible.", + description="Name/ID of the user that created the entity.", + union_mode="left_to_right", ) - component_id: UUIDFilterOption = Field( + reclaim_tolerance: ResourceRequestReclaimTolerance | str | None = Field( default=None, - description="The id of the component that is requesting the resources.", + description="The reclaim behavior tolerated by the request.", + union_mode="left_to_right", ) - step_run_id: UUIDFilterOption = Field( + component_id: UUID | str | None = Field( default=None, - description="The id of the step run that is requesting the resources.", + description="The component requesting resources.", + union_mode="left_to_right", + ) + step_run_id: UUID | str | None = Field( + default=None, + description="The step run requesting resources.", + union_mode="left_to_right", ) - preemption_initiated_by_id: UUIDFilterOption = Field( + preemption_initiated_by_id: UUID | str | None = Field( default=None, - description="The id of the request that initiated the preemption of this request.", + description="The request that initiated preemption.", + union_mode="left_to_right", ) - status: EnumFilterOption[ResourceRequestStatus] = Field( + status: ResourceRequestStatus | str | None = Field( default=None, description="The status of the resource request.", union_mode="left_to_right", ) - pipeline_run_id: UUIDFilterOption = Field( + pipeline_run_id: UUID | str | None = Field( default=None, - description="The id of the pipeline run that is requesting the resources.", + description="The pipeline run requesting resources.", + union_mode="left_to_right", + ) + pool_id: UUID | str | None = Field( + default=None, + description="The resource pool linked to the request.", + union_mode="left_to_right", ) - - def get_custom_filters( - self, - table: Type["AnySchema"], - ) -> List["ColumnElement[bool]"]: - """Get custom filters. - - Args: - table: The query table. - - Returns: - A list of custom filters. - """ - custom_filters = super().get_custom_filters(table) - - from sqlmodel import and_ - - from zenml.zen_stores.schemas import ( - ResourceRequestSchema, - StepRunSchema, - ) - - if self.pipeline_run_id: - pipeline_run_filters = ( - self.pipeline_run_id - if isinstance(self.pipeline_run_id, list) - else [self.pipeline_run_id] - ) - for pipeline_run_filter in pipeline_run_filters: - custom_filters.append( - and_( - ResourceRequestSchema.step_run_id == StepRunSchema.id, - self.generate_custom_query_conditions_for_column( - value=pipeline_run_filter, - table=StepRunSchema, - column="pipeline_run_id", - ), - ) - ) - - return custom_filters diff --git a/src/zenml/models/v2/core/step_run.py b/src/zenml/models/v2/core/step_run.py index 7957e143d15..caf42cba110 100644 --- a/src/zenml/models/v2/core/step_run.py +++ b/src/zenml/models/v2/core/step_run.py @@ -34,6 +34,7 @@ ArtifactSaveType, ExecutionStatus, StepRunInputArtifactType, + StepRuntime, StepType, ) from zenml.metadata.metadata_types import MetadataType @@ -182,6 +183,10 @@ class StepRunRequest(ProjectScopedRequest): "step run.", default=None, ) + resource_request_runtime: Optional[StepRuntime] = Field( + title="The runtime mode used to decide resource request semantics.", + default=None, + ) model_config = ConfigDict(protected_namespaces=()) @@ -255,6 +260,11 @@ class StepRunResponseBody(ProjectScopedResponseBody): "configured by this step run explicitly.", default=None, ) + resource_request_id: Optional[UUID] = Field( + title="The ID of an externally managed resource request for the " + "step run.", + default=None, + ) substitutions: Dict[str, str] = Field( title="The substitutions of the step run.", default={}, @@ -740,6 +750,15 @@ def model_version(self) -> Optional[ModelVersionResponse]: """ return self.get_resources().model_version + @property + def resource_request_id(self) -> Optional[UUID]: + """The `resource_request_id` property. + + Returns: + the value of the property. + """ + return self.get_body().resource_request_id + @property def resource_request(self) -> Optional["ResourceRequestResponse"]: """The `resource_request` property. @@ -910,6 +929,19 @@ def get_custom_filters( # ------------------ Heartbeat Model --------------- +class StepHeartbeatRequest(BaseModel): + """Light-weight model for Step Heartbeat requests.""" + + heartbeat_liveness_timeout_seconds: Optional[int] = Field( + default=None, + gt=0, + title=( + "Seconds the server should wait before considering the " + "heartbeat client dead." + ), + ) + + class StepHeartbeatResponse(BaseModel, use_enum_values=True): """Light-weight model for Step Heartbeat responses.""" diff --git a/src/zenml/models/v2/misc/external_user.py b/src/zenml/models/v2/misc/external_user.py index f6a7110d8ce..9b5e990977f 100644 --- a/src/zenml/models/v2/misc/external_user.py +++ b/src/zenml/models/v2/misc/external_user.py @@ -26,7 +26,7 @@ class ExternalUserModel(BaseModel): username: str email: Optional[str] = None name: Optional[str] = None - is_admin: bool = False + is_superuser: bool = False is_service_account: bool = False avatar_url: Optional[str] = None diff --git a/src/zenml/orchestrators/base_orchestrator.py b/src/zenml/orchestrators/base_orchestrator.py index 6cc65d92a41..1e9240b396d 100644 --- a/src/zenml/orchestrators/base_orchestrator.py +++ b/src/zenml/orchestrators/base_orchestrator.py @@ -63,6 +63,7 @@ from zenml.models import ( PipelineRunResponse, PipelineSnapshotResponse, + ResourceRequestResponse, ScheduleResponse, ScheduleUpdate, StepRunResponse, @@ -557,10 +558,25 @@ def can_run_isolated_steps(self) -> bool: return ( getattr(self.submit_isolated_step, "__func__", None) is not BaseOrchestrator.submit_isolated_step + or getattr( + self.submit_isolated_step_with_allocation, "__func__", None + ) + is not BaseOrchestrator.submit_isolated_step_with_allocation ) + @property + def supports_resource_pool_allocation(self) -> bool: + """Whether the orchestrator supports resource pool allocations. + + Returns: + Whether the orchestrator supports resource pool allocations. + """ + return False + def submit_isolated_step( - self, step_run_info: "StepRunInfo", environment: Dict[str, str] + self, + step_run_info: "StepRunInfo", + environment: Dict[str, str], ) -> None: """Submit an isolated step. @@ -582,6 +598,42 @@ def submit_isolated_step( f"the {self.__class__.__name__} orchestrator." ) + def submit_isolated_step_with_allocation( + self, + step_run_info: "StepRunInfo", + environment: Dict[str, str], + allocated_resource_request: Optional["ResourceRequestResponse"], + ) -> None: + """Submit an isolated step with an allocated resource request. + + Implementations should use + `zenml.orchestrators.utils.get_step_entrypoint_command(...)` to get + the command and arguments to run the step. + + Args: + step_run_info: The step run information. + environment: The environment variables to set in the execution + environment. + allocated_resource_request: The allocated resource request for the + step, if any. + + Raises: + NotImplementedError: If the orchestrator declares resource pool + allocation support but does not implement this method. + """ + if self.supports_resource_pool_allocation: + raise NotImplementedError( + f"The {self.__class__.__name__} orchestrator declares " + "resource pool allocation support but does not implement " + "`submit_isolated_step_with_allocation(...)`." + ) + + _ = allocated_resource_request + self.submit_isolated_step( + step_run_info=step_run_info, + environment=environment, + ) + def get_isolated_step_status( self, step_run: "StepRunResponse" ) -> ExecutionStatus: diff --git a/src/zenml/orchestrators/step_launcher.py b/src/zenml/orchestrators/step_launcher.py index ad32c196681..8454707b7d4 100644 --- a/src/zenml/orchestrators/step_launcher.py +++ b/src/zenml/orchestrators/step_launcher.py @@ -15,6 +15,7 @@ import time from contextlib import nullcontext +from datetime import timedelta from typing import ( TYPE_CHECKING, Any, @@ -34,6 +35,7 @@ ) from zenml.enums import ( ExecutionStatus, + ResourceRequestRuntimeState, ResourceRequestStatus, StepRuntime, ) @@ -45,6 +47,8 @@ PipelineRunRequest, PipelineRunResponse, PipelineSnapshotResponse, + ResourceRequestRenewalRequest, + ResourceRequestResponse, StepRunResponse, ) from zenml.models.v2.core.step_run import StepRunInputResponse @@ -410,8 +414,11 @@ def _run_step( skip_artifact_materialization=runtime.should_skip_artifact_materialization(), ) + allocated_resource_request: Optional[ResourceRequestResponse] = None if self._snapshot.is_dynamic: - self._wait_until_resources_acquired(step_run_info) + allocated_resource_request = self._wait_until_resources_acquired( + step_run_info + ) terminal_step_run = None try: @@ -423,6 +430,7 @@ def _run_step( terminal_step_run = self._run_step_with_step_operator( step_operator_name=step_operator_name, step_run_info=step_run_info, + allocated_resource_request=allocated_resource_request, ) elif not self._snapshot.is_dynamic: terminal_step_run = self._run_step_in_current_thread( @@ -463,10 +471,9 @@ def _run_step( output_artifact_uris=output_artifact_uris, ) else: - terminal_step_run = ( - self._run_step_with_dynamic_orchestrator( - step_run_info=step_run_info - ) + terminal_step_run = self._run_step_with_dynamic_orchestrator( + step_run_info=step_run_info, + allocated_resource_request=allocated_resource_request, ) except: # noqa: E722 output_utils.remove_artifact_dirs( @@ -480,12 +487,15 @@ def _run_step_with_step_operator( self, step_operator_name: Optional[str], step_run_info: StepRunInfo, + allocated_resource_request: Optional[ResourceRequestResponse], ) -> Optional[StepRunResponse]: """Runs the current step with a step operator. Args: step_operator_name: The name of the step operator to use. step_run_info: Additional information needed to run the step. + allocated_resource_request: The allocated resource request for the + step, if any. Raises: RuntimeError: If trying to use a step operator that does not support @@ -531,12 +541,16 @@ def _run_step_with_step_operator( ) try: - step_operator.submit( + step_operator.submit_with_allocation( info=step_run_info, entrypoint_command=entrypoint_command, environment=environment, + allocated_resource_request=allocated_resource_request, ) except NotImplementedError: + if step_operator.supports_resource_pool_allocation: + raise + if not self._wait: # We're running in a dynamic pipeline and for the monitoring to # work correctly, we only allow running with step operators that @@ -561,10 +575,11 @@ def _run_step_with_step_operator( step_operator.name, ) try: - step_operator.launch( + step_operator.launch_with_allocation( info=step_run_info, entrypoint_command=entrypoint_command, environment=environment, + allocated_resource_request=allocated_resource_request, ) finally: try: @@ -600,11 +615,14 @@ def _run_step_with_step_operator( def _run_step_with_dynamic_orchestrator( self, step_run_info: StepRunInfo, + allocated_resource_request: Optional[ResourceRequestResponse] = None, ) -> Optional[StepRunResponse]: """Runs the current step with a dynamic orchestrator. Args: step_run_info: Additional information needed to run the step. + allocated_resource_request: The allocated resource request for the + step, if any. Raises: BaseException: If the step run failed. @@ -624,9 +642,10 @@ def _run_step_with_dynamic_orchestrator( stack=self._stack, ) ) - self._stack.orchestrator.submit_isolated_step( + self._stack.orchestrator.submit_isolated_step_with_allocation( step_run_info=step_run_info, environment=environment, + allocated_resource_request=allocated_resource_request, ) if self._wait: try: @@ -676,7 +695,10 @@ def _finalize_remote_step( return None - def _cleanup_remote_step(self, step_run: StepRunResponse) -> None: + def _cleanup_remote_step( + self, + step_run: StepRunResponse, + ) -> None: """Clean up infrastructure after a remote step has finished. Args: @@ -786,76 +808,136 @@ def _get_initial_step_run_status(self) -> ExecutionStatus: def _wait_until_resources_acquired( self, step_run_info: StepRunInfo - ) -> None: + ) -> Optional[ResourceRequestResponse]: """Waits until the resources are acquired. Args: step_run_info: Step run information. + Returns: + The allocated resource request response, if the step has one. + Raises: RuntimeError: If the resource request was not found, or was rejected, preempted, or cancelled. """ + resource_request_id = step_run_info.step_run.resource_request_id resource_request = step_run_info.step_run.resource_request - if not resource_request: - return + if not resource_request_id: + if resource_request is None: + return None + resource_request_id = resource_request.id + elif ( + resource_request is not None + and resource_request.id != resource_request_id + ): + resource_request = None - if resource_request.status == ResourceRequestStatus.ALLOCATED: - return + step_name = step_run_info.pipeline_step_name + zen_store = Client().zen_store + resource_settings = step_run_info.config.resource_settings + allocation_wait_timeout = timedelta( + seconds=resource_settings.allocation_wait_timeout_seconds + ) + initialization_lease = timedelta( + seconds=resource_settings.initialization_lease_seconds + ) + wait_started_at = utc_now() + allocation_wait_deadline = wait_started_at + allocation_wait_timeout + max_poll_delay_seconds = 20.0 + poller_lease_buffer = timedelta(seconds=2 * max_poll_delay_seconds) for delay in exponential_backoff_delays( initial_delay=1.0, - max_delay=20.0, + max_delay=max_poll_delay_seconds, factor=2.0, jitter="equal", ): - try: - resource_request = Client().zen_store.get_resource_request( - resource_request.id, hydrate=False - ) - except KeyError as e: + now = utc_now() + if now >= allocation_wait_deadline: raise RuntimeError( - f"Resource request `{resource_request.id}` for step " - f"`{step_run_info.pipeline_step_name}` not found. This " - "is most likely because someone deleted the resource " - "request." - ) from e + f"Timed out after {resource_settings.allocation_wait_timeout_seconds} " + f"seconds waiting for resource request `{resource_request_id}` " + f"for step `{step_name}` to be allocated. Increase " + f"`ResourceSettings.allocation_wait_timeout_seconds` on this " + "step to wait longer." + ) + + if resource_request is None: + try: + # This lease tracks allocation-poller liveness before the + # step heartbeat worker exists. + lease_duration = min( + timedelta(seconds=delay) + poller_lease_buffer, + allocation_wait_deadline - now, + ) + resource_request = zen_store.renew_resource_request( + resource_request_id, + ResourceRequestRenewalRequest( + lease_expires_at=now + lease_duration, + runtime_state=ResourceRequestRuntimeState.PENDING, + ), + ) + except KeyError as e: + raise RuntimeError( + f"Resource request `{resource_request_id}` for step " + f"`{step_name}` not found. This is most likely because " + "someone deleted the resource request." + ) from e if resource_request.status == ResourceRequestStatus.ALLOCATED: + resource_request = zen_store.renew_resource_request( + resource_request_id, + ResourceRequestRenewalRequest( + lease_expires_at=utc_now() + initialization_lease, + runtime_state=ResourceRequestRuntimeState.SUBMITTED, + ), + ) logger.info( "Resource request `%s` for step `%s` was approved.", resource_request.id, - step_run_info.pipeline_step_name, + step_name, ) publish_utils.publish_step_run_status_update( step_run_id=step_run_info.step_run_id, status=ExecutionStatus.RUNNING, ) - return - elif resource_request.status == ResourceRequestStatus.REJECTED: + return resource_request + if resource_request.status == ResourceRequestStatus.REJECTED: reason = resource_request.status_reason or "Unknown reason" raise RuntimeError( f"Resource request `{resource_request.id}` for step " - f"`{step_run_info.pipeline_step_name}` was rejected: " - f"{reason}" + f"`{step_name}` was rejected: {reason}" ) - elif resource_request.status == ResourceRequestStatus.PREEMPTED: + if resource_request.status in { + ResourceRequestStatus.PREEMPTING, + ResourceRequestStatus.PREEMPTED, + ResourceRequestStatus.RELEASED, + ResourceRequestStatus.EXPIRED, + }: reason = resource_request.status_reason or "Unknown reason" raise RuntimeError( f"Resource request `{resource_request.id}` for step " - f"`{step_run_info.pipeline_step_name}` was preempted: " - f"{reason}" + f"`{step_name}` reached status " + f"`{resource_request.status}`: {reason}" ) - elif resource_request.status == ResourceRequestStatus.CANCELLED: + if resource_request.status == ResourceRequestStatus.CANCELLED: + reason = resource_request.status_reason or "Unknown reason" raise RuntimeError( f"Resource request `{resource_request.id}` for step " - f"`{step_run_info.pipeline_step_name}` was cancelled." + f"`{step_name}` was cancelled: {reason}" ) logger.info( "Waiting for resource request `%s` of step `%s` to be " "approved...", resource_request.id, - step_run_info.pipeline_step_name, + step_name, ) time.sleep(delay) + resource_request = None + + raise RuntimeError( + f"Stopped waiting for resource request `{resource_request_id}` " + f"for step `{step_name}` before it was allocated." + ) diff --git a/src/zenml/orchestrators/step_run_utils.py b/src/zenml/orchestrators/step_run_utils.py index 0f9ade3fd22..17d8b5786db 100644 --- a/src/zenml/orchestrators/step_run_utils.py +++ b/src/zenml/orchestrators/step_run_utils.py @@ -296,6 +296,17 @@ def populate_request(self, request: StepRunRequest) -> None: request.status != ExecutionStatus.CACHED and self.snapshot.is_dynamic ): + from zenml.execution.pipeline.dynamic.compilation import ( + get_step_runtime, + ) + + request.resource_request_runtime = get_step_runtime( + step_config=step.config, + pipeline_docker_settings=( + self.snapshot.pipeline_configuration.docker_settings + ), + orchestrator=self.stack.orchestrator, + ) if step.config.step_operator: assert self.stack.step_operator request.resource_requester = self.stack.step_operator.id diff --git a/src/zenml/step_operators/base_step_operator.py b/src/zenml/step_operators/base_step_operator.py index cf105f205e1..36aac791627 100644 --- a/src/zenml/step_operators/base_step_operator.py +++ b/src/zenml/step_operators/base_step_operator.py @@ -15,7 +15,7 @@ import time from abc import ABC, abstractmethod -from typing import TYPE_CHECKING, Dict, List, Type, cast +from typing import TYPE_CHECKING, Dict, List, Optional, Type, cast from zenml.client import Client from zenml.enums import ExecutionStatus, StackComponentType @@ -29,6 +29,7 @@ if TYPE_CHECKING: from zenml.config.step_run_info import StepRunInfo + from zenml.models import ResourceRequestResponse logger = get_logger(__name__) @@ -64,6 +65,15 @@ def entrypoint_config_class( """ return StepOperatorEntrypointConfiguration + @property + def supports_resource_pool_allocation(self) -> bool: + """Whether the step operator supports resource pool allocations. + + Returns: + Whether the step operator supports resource pool allocations. + """ + return False + def submit( self, info: "StepRunInfo", @@ -90,6 +100,84 @@ def submit( f"the {self.__class__.__name__} step operator." ) + def submit_with_allocation( + self, + info: "StepRunInfo", + entrypoint_command: List[str], + environment: Dict[str, str], + allocated_resource_request: Optional["ResourceRequestResponse"], + ) -> None: + """Submit a step run with an allocated resource request. + + Args: + info: Information about the step run. + entrypoint_command: Command that executes the step. + environment: Environment variables to set in the step operator + environment. + allocated_resource_request: The allocated resource request for the + step, if any. + + Raises: + NotImplementedError: If the step operator declares resource pool + allocation support but does not implement this method. + """ + if self.supports_resource_pool_allocation: + raise NotImplementedError( + f"The {self.__class__.__name__} step operator declares " + "resource pool allocation support but does not implement " + "`submit_with_allocation(...)`." + ) + + _ = allocated_resource_request + self.submit( + info=info, + entrypoint_command=entrypoint_command, + environment=environment, + ) + + def launch_with_allocation( + self, + info: "StepRunInfo", + entrypoint_command: List[str], + environment: Dict[str, str], + allocated_resource_request: Optional["ResourceRequestResponse"], + ) -> None: + """Launch a step run with an allocated resource request. + + Args: + info: Information about the step run. + entrypoint_command: Command that executes the step. + environment: Environment variables to set in the step operator + environment. + allocated_resource_request: The allocated resource request for the + step, if any. + + Raises: + NotImplementedError: If the step operator declares resource pool + allocation support but does not implement this method, or if no + legacy launch method exists. + """ + if self.supports_resource_pool_allocation: + raise NotImplementedError( + f"The {self.__class__.__name__} step operator declares " + "resource pool allocation support but does not implement " + "`launch_with_allocation(...)`." + ) + + launch = getattr(self, "launch", None) + if not callable(launch): + raise NotImplementedError( + f"The {self.__class__.__name__} step operator does not " + "implement the legacy `launch(...)` method." + ) + + _ = allocated_resource_request + launch( + info=info, + entrypoint_command=entrypoint_command, + environment=environment, + ) + def get_status(self, step_run: "StepRunResponse") -> ExecutionStatus: """Get the status of a step run. diff --git a/src/zenml/steps/heartbeat.py b/src/zenml/steps/heartbeat.py index 5c643286080..1308158c92a 100644 --- a/src/zenml/steps/heartbeat.py +++ b/src/zenml/steps/heartbeat.py @@ -40,11 +40,24 @@ class StepHeartBeatTerminationException(Exception): class StepHeartbeatWorker: """Worker class implementing heartbeat polling and remote termination.""" - STEP_HEARTBEAT_INTERVAL_SECONDS = 15 # submit heartbeats every 2 minutes. + STEP_HEARTBEAT_INTERVAL_SECONDS = 15 MAX_HEARTBEAT_INTERVAL_SECONDS = 60 * 5 MAX_CONSECUTIVE_FAILURES = 3 + @classmethod + def heartbeat_liveness_timeout_seconds(cls) -> int: + """Get the liveness timeout reported with heartbeat requests. + + Returns: + Seconds the server should wait for the next heartbeat before + considering the heartbeat client dead. + """ + return ( + cls.MAX_HEARTBEAT_INTERVAL_SECONDS + + cls.STEP_HEARTBEAT_INTERVAL_SECONDS + ) + def __init__(self, step_id: UUID): """Heartbeat worker constructor. @@ -56,6 +69,7 @@ def __init__(self, step_id: UUID): self._thread: threading.Thread | None = None self._running: bool = False self._terminated: bool = False + self._interrupt_signaled: bool = False self._consecutive_failures: int = 0 self._heartbeat_interval_seconds: int = ( @@ -118,6 +132,7 @@ def start(self) -> None: self._running = True self._terminated = False + self._interrupt_signaled = False active_context = contextvars.copy_context() self._thread = threading.Thread( @@ -153,15 +168,8 @@ def _run(self) -> None: try: self._heartbeat() self._consecutive_failures = 0 - # One-shot: signal the main thread and stop the loop. if self._terminated: - logger.info( - "%s is remotely stopped, interrupting main thread", - self.name, - ) - _thread.interrupt_main() # raises KeyboardInterrupt in main thread - self._running = False - # Ensure we stop our own loop as well. + self._signal_main_thread_interrupt() except Exception as exc: self._consecutive_failures += 1 @@ -188,12 +196,40 @@ def _run(self) -> None: finally: logger.debug("%s run() loop exiting", self.name) + def _signal_main_thread_interrupt(self) -> None: + """Request cooperative shutdown of the main thread. + + ``interrupt_main`` only schedules ``KeyboardInterrupt`` for the main + thread; it may not wake a long-running C call such as ``time.sleep`` + until that call returns. Keep heartbeating until ``stop()`` is called + from the main thread's ``finally`` block so leases stay renewed while + shutdown is in progress. + """ + if not self._interrupt_signaled: + logger.info( + "%s is remotely stopped, interrupting main thread", + self.name, + ) + self._interrupt_signaled = True + else: + logger.debug( + "%s still waiting for main thread shutdown, " + "re-sending interrupt", + self.name, + ) + _thread.interrupt_main() + def _heartbeat(self) -> None: from zenml.config.global_config import GlobalConfiguration store = GlobalConfiguration().zen_store - response = store.update_step_heartbeat(step_run_id=self.step_id) + response = store.update_step_heartbeat( + step_run_id=self.step_id, + heartbeat_liveness_timeout_seconds=( + self.heartbeat_liveness_timeout_seconds() + ), + ) if not response.heartbeat_enabled: logger.debug("Heartbeat set to disabled - stopping worker") diff --git a/src/zenml/zen_server/auth.py b/src/zenml/zen_server/auth.py index 587c5e196ee..963fe3c1323 100644 --- a/src/zenml/zen_server/auth.py +++ b/src/zenml/zen_server/auth.py @@ -883,7 +883,7 @@ def authenticate_external_user( email_opted_in=True, active=True, email=external_user.email, - is_admin=external_user.is_admin, + is_admin=external_user.is_superuser, avatar_url=external_user.avatar_url, external_user_id=external_user.id, ), @@ -925,7 +925,7 @@ def authenticate_external_user( email_opted_in=True, active=True, email=external_user.email, - is_admin=external_user.is_admin, + is_admin=external_user.is_superuser, avatar_url=external_user.avatar_url, ) ) diff --git a/src/zenml/zen_server/rbac/models.py b/src/zenml/zen_server/rbac/models.py index 502e0dbd1e6..456dec84b5e 100644 --- a/src/zenml/zen_server/rbac/models.py +++ b/src/zenml/zen_server/rbac/models.py @@ -71,8 +71,6 @@ class ResourceType(StrEnum): SERVICE_CONNECTOR = "service_connector" STACK = "stack" STACK_COMPONENT = "stack_component" - RESOURCE_POOL = "resource_pool" - RESOURCE_POOL_SUBJECT_POLICY = "resource_pool_subject_policy" TAG = "tag" PROJECT = "project" # Deactivated for now @@ -91,8 +89,6 @@ def is_project_scoped(self) -> bool: self.SERVICE_CONNECTOR, self.STACK, self.STACK_COMPONENT, - self.RESOURCE_POOL, - self.RESOURCE_POOL_SUBJECT_POLICY, self.TAG, self.SERVICE_ACCOUNT, self.PROJECT, diff --git a/src/zenml/zen_server/rbac/utils.py b/src/zenml/zen_server/rbac/utils.py index 90eefbb0cba..a01c8feaa53 100644 --- a/src/zenml/zen_server/rbac/utils.py +++ b/src/zenml/zen_server/rbac/utils.py @@ -498,10 +498,6 @@ def get_resource_type_for_model( PlatformEventTriggerResponse, ProjectRequest, ProjectResponse, - ResourcePoolRequest, - ResourcePoolResponse, - ResourcePoolSubjectPolicyRequest, - ResourcePoolSubjectPolicyResponse, RunMetadataRequest, RunTemplateRequest, RunTemplateResponse, @@ -536,8 +532,6 @@ def get_resource_type_for_model( CodeRepositoryResponse: ResourceType.CODE_REPOSITORY, ComponentRequest: ResourceType.STACK_COMPONENT, ComponentResponse: ResourceType.STACK_COMPONENT, - ResourcePoolRequest: ResourceType.RESOURCE_POOL, - ResourcePoolResponse: ResourceType.RESOURCE_POOL, FlavorRequest: ResourceType.FLAVOR, FlavorResponse: ResourceType.FLAVOR, ModelRequest: ResourceType.MODEL, @@ -573,8 +567,6 @@ def get_resource_type_for_model( TagResponse: ResourceType.TAG, ProjectResponse: ResourceType.PROJECT, ProjectRequest: ResourceType.PROJECT, - ResourcePoolSubjectPolicyRequest: ResourceType.RESOURCE_POOL_SUBJECT_POLICY, - ResourcePoolSubjectPolicyResponse: ResourceType.RESOURCE_POOL_SUBJECT_POLICY, # UserResponse: ResourceType.USER, PlatformEventTriggerRequest: ResourceType.TRIGGER, PlatformEventTriggerResponse: ResourceType.TRIGGER, @@ -701,8 +693,6 @@ def _get_resource_type_schema_mapping() -> Dict[ PipelineRunSchema, PipelineSchema, PipelineSnapshotSchema, - ResourcePoolSchema, - ResourcePoolSubjectPolicySchema, RunMetadataSchema, RunTemplateSchema, ScheduleSchema, @@ -720,7 +710,6 @@ def _get_resource_type_schema_mapping() -> Dict[ ResourceType.STACK: StackSchema, ResourceType.FLAVOR: FlavorSchema, ResourceType.STACK_COMPONENT: StackComponentSchema, - ResourceType.RESOURCE_POOL: ResourcePoolSchema, ResourceType.PIPELINE: PipelineSchema, ResourceType.CODE_REPOSITORY: CodeRepositorySchema, ResourceType.MODEL: ModelSchema, @@ -742,7 +731,6 @@ def _get_resource_type_schema_mapping() -> Dict[ ResourceType.SCHEDULE: ScheduleSchema, # ResourceType.USER: UserSchema, ResourceType.TRIGGER: TriggerSchema, - ResourceType.RESOURCE_POOL_SUBJECT_POLICY: ResourcePoolSubjectPolicySchema, } diff --git a/src/zenml/zen_server/routers/resource_pool_subject_policies_endpoints.py b/src/zenml/zen_server/routers/resource_pool_subject_policies_endpoints.py deleted file mode 100644 index 9e57c83a518..00000000000 --- a/src/zenml/zen_server/routers/resource_pool_subject_policies_endpoints.py +++ /dev/null @@ -1,184 +0,0 @@ -# Copyright (c) ZenML GmbH 2026. All Rights Reserved. -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at: -# -# https://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express -# or implied. See the License for the specific language governing -# permissions and limitations under the License. -"""Endpoint definitions for resource pool subject policies.""" - -from uuid import UUID - -from fastapi import APIRouter, Depends, Security - -from zenml.constants import ( - API, - RESOURCE_POOL_FEATURE, - RESOURCE_POOL_SUBJECT_POLICIES, - VERSION_1, -) -from zenml.models import ( - Page, - ResourcePoolSubjectPolicyFilter, - ResourcePoolSubjectPolicyRequest, - ResourcePoolSubjectPolicyResponse, - ResourcePoolSubjectPolicyUpdate, -) -from zenml.zen_server.auth import AuthContext, authorize -from zenml.zen_server.exceptions import error_response -from zenml.zen_server.feature_gate.endpoint_utils import check_entitlement -from zenml.zen_server.rbac.endpoint_utils import ( - verify_permissions_and_create_entity, - verify_permissions_and_delete_entity, - verify_permissions_and_get_entity, - verify_permissions_and_list_entities, - verify_permissions_and_update_entity, -) -from zenml.zen_server.rbac.models import ResourceType -from zenml.zen_server.utils import ( - async_fastapi_endpoint_wrapper, - make_dependable, - zen_store, -) - -router = APIRouter( - prefix=API + VERSION_1 + RESOURCE_POOL_SUBJECT_POLICIES, - tags=["resource_pool_subject_policies"], - responses={401: error_response, 403: error_response}, -) - - -@router.post( - "", - responses={401: error_response, 409: error_response, 422: error_response}, -) -@async_fastapi_endpoint_wrapper -def create_resource_pool_subject_policy( - policy: ResourcePoolSubjectPolicyRequest, - _: AuthContext = Security(authorize), -) -> ResourcePoolSubjectPolicyResponse: - """Create a resource pool subject policy. - - Args: - policy: Policy request payload. - - Returns: - The created policy. - """ - check_entitlement(feature=RESOURCE_POOL_FEATURE) - - return verify_permissions_and_create_entity( - request_model=policy, - create_method=zen_store().create_resource_pool_subject_policy, - ) - - -@router.get( - "", - responses={401: error_response, 404: error_response, 422: error_response}, -) -@async_fastapi_endpoint_wrapper -def list_resource_pool_subject_policies( - policy_filter_model: ResourcePoolSubjectPolicyFilter = Depends( - make_dependable(ResourcePoolSubjectPolicyFilter) - ), - hydrate: bool = False, - _: AuthContext = Security(authorize), -) -> Page[ResourcePoolSubjectPolicyResponse]: - """List resource pool subject policies. - - Args: - policy_filter_model: Filter model used for pagination and filtering. - hydrate: Whether to include metadata in the response. - - Returns: - Matching policies. - """ - return verify_permissions_and_list_entities( - filter_model=policy_filter_model, - resource_type=ResourceType.RESOURCE_POOL_SUBJECT_POLICY, - list_method=zen_store().list_resource_pool_subject_policies, - hydrate=hydrate, - ) - - -@router.get( - "/{policy_id}", - responses={401: error_response, 404: error_response, 422: error_response}, -) -@async_fastapi_endpoint_wrapper -def get_resource_pool_subject_policy( - policy_id: UUID, - hydrate: bool = True, - _: AuthContext = Security(authorize), -) -> ResourcePoolSubjectPolicyResponse: - """Get a resource pool subject policy. - - Args: - policy_id: Policy ID. - hydrate: Whether to include metadata in the response. - - Returns: - The requested policy. - """ - return verify_permissions_and_get_entity( - id=policy_id, - get_method=zen_store().get_resource_pool_subject_policy, - hydrate=hydrate, - ) - - -@router.put( - "/{policy_id}", - responses={401: error_response, 404: error_response, 422: error_response}, -) -@async_fastapi_endpoint_wrapper -def update_resource_pool_subject_policy( - policy_id: UUID, - policy_update: ResourcePoolSubjectPolicyUpdate, - _: AuthContext = Security(authorize), -) -> ResourcePoolSubjectPolicyResponse: - """Update a resource pool subject policy. - - Args: - policy_id: Policy ID. - policy_update: Update payload. - - Returns: - The updated policy. - """ - check_entitlement(feature=RESOURCE_POOL_FEATURE) - - return verify_permissions_and_update_entity( - id=policy_id, - update_model=policy_update, - get_method=zen_store().get_resource_pool_subject_policy, - update_method=zen_store().update_resource_pool_subject_policy, - ) - - -@router.delete( - "/{policy_id}", - responses={401: error_response, 404: error_response, 422: error_response}, -) -@async_fastapi_endpoint_wrapper -def delete_resource_pool_subject_policy( - policy_id: UUID, - _: AuthContext = Security(authorize), -) -> None: - """Delete a resource pool subject policy. - - Args: - policy_id: Policy ID. - """ - verify_permissions_and_delete_entity( - id=policy_id, - get_method=zen_store().get_resource_pool_subject_policy, - delete_method=zen_store().delete_resource_pool_subject_policy, - ) diff --git a/src/zenml/zen_server/routers/resource_pools_endpoints.py b/src/zenml/zen_server/routers/resource_pools_endpoints.py deleted file mode 100644 index df85b7a4213..00000000000 --- a/src/zenml/zen_server/routers/resource_pools_endpoints.py +++ /dev/null @@ -1,187 +0,0 @@ -# Copyright (c) ZenML GmbH 2026. All Rights Reserved. -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at: -# -# https://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express -# or implied. See the License for the specific language governing -# permissions and limitations under the License. -"""Endpoint definitions for resource pools.""" - -from uuid import UUID - -from fastapi import APIRouter, Depends, Security - -from zenml.constants import ( - API, - RESOURCE_POOL_FEATURE, - RESOURCE_POOLS, - VERSION_1, -) -from zenml.models import ( - Page, - ResourcePoolFilter, - ResourcePoolRequest, - ResourcePoolResponse, - ResourcePoolUpdate, -) -from zenml.zen_server.auth import AuthContext, authorize -from zenml.zen_server.exceptions import error_response -from zenml.zen_server.feature_gate.endpoint_utils import check_entitlement -from zenml.zen_server.rbac.endpoint_utils import ( - verify_permissions_and_create_entity, - verify_permissions_and_delete_entity, - verify_permissions_and_get_entity, - verify_permissions_and_list_entities, - verify_permissions_and_update_entity, -) -from zenml.zen_server.rbac.models import ResourceType -from zenml.zen_server.utils import ( - async_fastapi_endpoint_wrapper, - make_dependable, - zen_store, -) - -router = APIRouter( - prefix=API + VERSION_1 + RESOURCE_POOLS, - tags=["resource_pools"], - responses={401: error_response, 403: error_response}, -) - - -@router.post( - "", - responses={401: error_response, 409: error_response, 422: error_response}, -) -@async_fastapi_endpoint_wrapper -def create_resource_pool( - resource_pool: ResourcePoolRequest, - _: AuthContext = Security(authorize), -) -> ResourcePoolResponse: - """Creates a resource pool. - - Args: - resource_pool: Resource pool to register. - - Returns: - The created resource pool. - """ - check_entitlement(feature=RESOURCE_POOL_FEATURE) - - return verify_permissions_and_create_entity( - request_model=resource_pool, - create_method=zen_store().create_resource_pool, - ) - - -@router.get( - "", - responses={401: error_response, 404: error_response, 422: error_response}, -) -@async_fastapi_endpoint_wrapper -def list_resource_pools( - resource_pool_filter_model: ResourcePoolFilter = Depends( - make_dependable(ResourcePoolFilter) - ), - hydrate: bool = False, - _: AuthContext = Security(authorize), -) -> Page[ResourcePoolResponse]: - """Get a list of all resource pools. - - Args: - resource_pool_filter_model: Filter model used for pagination, sorting, - filtering. - hydrate: Flag deciding whether to hydrate the output model(s) - by including metadata fields in the response. - - Returns: - List of resource pools matching the filter criteria. - """ - return verify_permissions_and_list_entities( - filter_model=resource_pool_filter_model, - resource_type=ResourceType.RESOURCE_POOL, - list_method=zen_store().list_resource_pools, - hydrate=hydrate, - ) - - -@router.get( - "/{resource_pool_id}", - responses={401: error_response, 404: error_response, 422: error_response}, -) -@async_fastapi_endpoint_wrapper -def get_resource_pool( - resource_pool_id: UUID, - hydrate: bool = True, - _: AuthContext = Security(authorize), -) -> ResourcePoolResponse: - """Returns the requested resource pool. - - Args: - resource_pool_id: ID of the resource pool. - hydrate: Flag deciding whether to hydrate the output model(s) - by including metadata fields in the response. - - Returns: - The requested resource pool. - """ - return verify_permissions_and_get_entity( - id=resource_pool_id, - get_method=zen_store().get_resource_pool, - hydrate=hydrate, - ) - - -@router.put( - "/{resource_pool_id}", - responses={401: error_response, 404: error_response, 422: error_response}, -) -@async_fastapi_endpoint_wrapper -def update_resource_pool( - resource_pool_id: UUID, - resource_pool_update: ResourcePoolUpdate, - _: AuthContext = Security(authorize), -) -> ResourcePoolResponse: - """Updates a resource pool. - - Args: - resource_pool_id: ID of the resource pool. - resource_pool_update: Resource pool to use to update. - - Returns: - Updated resource pool. - """ - check_entitlement(feature=RESOURCE_POOL_FEATURE) - - return verify_permissions_and_update_entity( - id=resource_pool_id, - update_model=resource_pool_update, - get_method=zen_store().get_resource_pool, - update_method=zen_store().update_resource_pool, - ) - - -@router.delete( - "/{resource_pool_id}", - responses={401: error_response, 404: error_response, 422: error_response}, -) -@async_fastapi_endpoint_wrapper -def delete_resource_pool( - resource_pool_id: UUID, - _: AuthContext = Security(authorize), -) -> None: - """Deletes a resource pool. - - Args: - resource_pool_id: ID of the resource pool. - """ - verify_permissions_and_delete_entity( - id=resource_pool_id, - get_method=zen_store().get_resource_pool, - delete_method=zen_store().delete_resource_pool, - ) diff --git a/src/zenml/zen_server/routers/resource_requests_endpoints.py b/src/zenml/zen_server/routers/resource_requests_endpoints.py index dcb401e3121..391654b49bb 100644 --- a/src/zenml/zen_server/routers/resource_requests_endpoints.py +++ b/src/zenml/zen_server/routers/resource_requests_endpoints.py @@ -21,6 +21,7 @@ from zenml.models import ( Page, ResourceRequestFilter, + ResourceRequestRenewalRequest, ResourceRequestResponse, ) from zenml.zen_server.auth import AuthContext, authorize @@ -53,8 +54,8 @@ def list_resource_requests( """Get a list of all resource requests. Args: - resource_request_filter_model: Filter model used for pagination, sorting, - filtering. + resource_request_filter_model: Filter model used for pagination, + sorting, filtering. hydrate: Flag deciding whether to hydrate the output model(s) by including metadata fields in the response. @@ -92,18 +93,46 @@ def get_resource_request( ) -@router.delete( - "/{resource_request_id}", +@router.post( + "/{resource_request_id}/renew", responses={401: error_response, 404: error_response, 422: error_response}, ) @async_fastapi_endpoint_wrapper -def delete_resource_request( +def renew_resource_request( resource_request_id: UUID, + renewal_request: ResourceRequestRenewalRequest, _: AuthContext = Security(authorize), -) -> None: - """Deletes a resource request. +) -> ResourceRequestResponse: + """Renew a resource request lease. Args: resource_request_id: ID of the resource request. + renewal_request: Renewed lease expiration timestamp. + + Returns: + The renewed resource request. + """ + return zen_store().renew_resource_request( + resource_request_id, + renewal_request, + ) + + +@router.post( + "/{resource_request_id}/release", + responses={401: error_response, 404: error_response, 422: error_response}, +) +@async_fastapi_endpoint_wrapper +def release_resource_request( + resource_request_id: UUID, + _: AuthContext = Security(authorize), +) -> ResourceRequestResponse: + """Release a resource request on behalf of its owner. + + Args: + resource_request_id: ID of the resource request. + + Returns: + The released resource request. """ - zen_store().delete_resource_request(resource_request_id) + return zen_store().release_resource_request(resource_request_id) diff --git a/src/zenml/zen_server/routers/steps_endpoints.py b/src/zenml/zen_server/routers/steps_endpoints.py index 5134097724e..fe95a6061f1 100644 --- a/src/zenml/zen_server/routers/steps_endpoints.py +++ b/src/zenml/zen_server/routers/steps_endpoints.py @@ -36,7 +36,10 @@ StepRunResponse, StepRunUpdate, ) -from zenml.models.v2.core.step_run import StepHeartbeatResponse +from zenml.models.v2.core.step_run import ( + StepHeartbeatRequest, + StepHeartbeatResponse, +) from zenml.utils.logging_utils import ( LogEntry, fetch_logs, @@ -212,21 +215,32 @@ def update_step( @async_fastapi_endpoint_wrapper(deduplicate=False) def update_heartbeat( step_run_id: UUID, + heartbeat_request: Optional[StepHeartbeatRequest] = None, auth_context: AuthContext = Security(authorize), ) -> StepHeartbeatResponse: """Updates a step. Args: step_run_id: ID of the step. + heartbeat_request: Optional heartbeat request body. auth_context: Authorization/Authentication context. Returns: The step heartbeat response (id, status, last_heartbeat). """ + heartbeat_liveness_timeout_seconds = None + if heartbeat_request: + heartbeat_liveness_timeout_seconds = ( + heartbeat_request.heartbeat_liveness_timeout_seconds + ) + return zen_store().validate_and_update_heartbeat( step_run_id=step_run_id, token_run_id=auth_context.access_token.pipeline_run_id, # type: ignore[union-attr] token_schedule_id=auth_context.access_token.schedule_id, # type: ignore[union-attr] + heartbeat_liveness_timeout_seconds=( + heartbeat_liveness_timeout_seconds + ), ) diff --git a/src/zenml/zen_server/zen_server_api.py b/src/zenml/zen_server/zen_server_api.py index c908f862ce7..ed1cf9ce9be 100644 --- a/src/zenml/zen_server/zen_server_api.py +++ b/src/zenml/zen_server/zen_server_api.py @@ -74,8 +74,6 @@ pipeline_snapshot_endpoints, pipelines_endpoints, projects_endpoints, - resource_pool_subject_policies_endpoints, - resource_pools_endpoints, resource_requests_endpoints, run_metadata_endpoints, run_templates_endpoints, @@ -345,8 +343,6 @@ async def dashboard(request: Request) -> Any: app.include_router(users_endpoints.current_user_router) app.include_router(projects_endpoints.workspace_router) app.include_router(projects_endpoints.router) -app.include_router(resource_pools_endpoints.router) -app.include_router(resource_pool_subject_policies_endpoints.router) app.include_router(resource_requests_endpoints.router) app.include_router(trigger_endpoints.router) diff --git a/src/zenml/zen_stores/migrations/versions/9f2b8c7d6e5a_drop_resource_pool_entities.py b/src/zenml/zen_stores/migrations/versions/9f2b8c7d6e5a_drop_resource_pool_entities.py new file mode 100644 index 00000000000..edded23ee99 --- /dev/null +++ b/src/zenml/zen_stores/migrations/versions/9f2b8c7d6e5a_drop_resource_pool_entities.py @@ -0,0 +1,41 @@ +"""Drop resource pool and request entities [9f2b8c7d6e5a]. + +Revision ID: 9f2b8c7d6e5a +Revises: 0.96.2 +Create Date: 2026-06-20 00:00:00.000000 + +""" + +import sqlalchemy as sa +from alembic import op + +# revision identifiers, used by Alembic. +revision = "9f2b8c7d6e5a" +down_revision = "0.96.2" +branch_labels = None +depends_on = None + + +def upgrade() -> None: + """Upgrade database schema and/or data, creating a new revision.""" + op.drop_table("resource_pool_queue") + op.drop_table("resource_pool_allocation") + op.drop_table("resource_pool_subject_policy_resource") + op.drop_table("resource_pool_subject_policy") + op.drop_table("resource_pool_resource") + op.drop_table("resource_pool") + op.drop_table("resource_request_resource") + op.drop_table("resource_request") + with op.batch_alter_table("step_run", schema=None) as batch_op: + batch_op.add_column( + sa.Column( + "resource_request_id", + sa.Uuid(), + nullable=True, + ) + ) + + +def downgrade() -> None: + """Downgrade database schema and/or data back to the previous revision.""" + # Downgrades are not generally supported in ZenML. diff --git a/src/zenml/zen_stores/resource_pools/resource_manager/__init__.py b/src/zenml/zen_stores/resource_pools/resource_manager/__init__.py new file mode 100644 index 00000000000..91991bd99e2 --- /dev/null +++ b/src/zenml/zen_stores/resource_pools/resource_manager/__init__.py @@ -0,0 +1,20 @@ +# Copyright (c) ZenML GmbH 2026. All Rights Reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at: +# +# https://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express +# or implied. See the License for the specific language governing +# permissions and limitations under the License. +"""ZenML Pro Resource Manager-backed resource pool store.""" + +from zenml.zen_stores.resource_pools.resource_manager.store import ( + ResourceManagerResourcePoolsStore, +) + +__all__ = ["ResourceManagerResourcePoolsStore"] diff --git a/src/zenml/zen_stores/resource_pools/resource_manager/client.py b/src/zenml/zen_stores/resource_pools/resource_manager/client.py new file mode 100644 index 00000000000..67684cc1c5b --- /dev/null +++ b/src/zenml/zen_stores/resource_pools/resource_manager/client.py @@ -0,0 +1,589 @@ +# Copyright (c) ZenML GmbH 2026. All Rights Reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at: +# +# https://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express +# or implied. See the License for the specific language governing +# permissions and limitations under the License. +"""Synchronous client for the ZenML Pro Resource Manager service.""" + +from datetime import datetime, timedelta +from threading import RLock +from typing import Any, Optional, Protocol, Type, TypeVar +from uuid import UUID, uuid4 + +import requests +from pydantic import BaseModel +from requests.adapters import HTTPAdapter, Retry + +from zenml.enums import ResourceRequestRuntimeState +from zenml.exceptions import ( + CredentialsNotValid, + EntityExistsError, + IllegalOperationError, +) +from zenml.utils.time_utils import utc_now +from zenml.zen_stores.resource_pools.resource_manager.transport import ( + RMResourceRequestCreate, + RMResourceRequestListResponse, + RMResourceRequestRenewalRequest, + RMResourceRequestResponse, +) + +ModelT = TypeVar("ModelT", bound=BaseModel) + + +class CloudConnection(Protocol): + """Protocol for the Cloud API connection used by this client.""" + + def get( + self, endpoint: str, params: Optional[dict[str, Any]] + ) -> requests.Response: + """Send a GET request to the Cloud API. + + Args: + endpoint: Cloud API endpoint path. + params: Query parameters for the request. + + Returns: + Response returned by the Cloud API. + """ + + +class ResourceManagerAuthorization(BaseModel): + """Cloud API response for direct Resource Manager access.""" + + token: str + token_type: str = "bearer" + resource_manager_id: UUID + api_url: str + organization_id: UUID + workspace_id: Optional[UUID] = None + expires_at: datetime + + +class ResourceManagerLoginRequest(BaseModel): + """Resource Manager login request.""" + + organization_id: Optional[UUID] = None + workspace_id: Optional[UUID] = None + + +class ResourceManagerLoginResponse(BaseModel): + """Resource Manager local session response.""" + + access_token: str + token_type: str = "bearer" + expires_at: datetime + organization_id: UUID + workspace_id: Optional[UUID] = None + user_id: Optional[UUID] = None + resource_manager_id: UUID + + +class ResourceManagerClient: + """Minimal synchronous client for the Resource Manager REST API.""" + + IDEMPOTENCY_HEADER = "X-Idempotency-Key" + + @staticmethod + def error_message_from_response(response: requests.Response) -> str: + """Extract a human-readable error message from an RM response. + + Args: + response: Failed HTTP response from Resource Manager. + + Returns: + Error message text suitable for ZenML exceptions. + """ + try: + payload = response.json() + except requests.JSONDecodeError: + return response.text + + if not isinstance(payload, dict): + return response.text + + detail = payload.get("detail", response.text) + if isinstance(detail, str): + return detail + if isinstance(detail, list): + return ": ".join(str(item) for item in detail) + return response.text + + @staticmethod + def exception_from_rm_response( + response: requests.Response, + ) -> Exception: + """Map a Resource Manager HTTP error to a ZenML store exception. + + Resource Manager returns FastAPI-style ``{"detail": "..."}`` bodies. + These are translated to the same exception types the workspace REST API + already maps to HTTP status codes. + + Args: + response: Failed HTTP response from Resource Manager. + + Returns: + Exception to raise to workspace callers. + """ + message = ResourceManagerClient.error_message_from_response(response) + status_code = response.status_code + + if status_code == 404: + return KeyError(message) + if status_code == 409: + lowered = message.lower() + if ( + "already exists" in lowered + or "already registered" in lowered + or "duplicate" in lowered + ): + return EntityExistsError(message) + return IllegalOperationError(message) + if status_code in {400, 422}: + return ValueError(message) + if status_code == 403: + return IllegalOperationError(message) + if status_code == 401: + return CredentialsNotValid(message) + + return RuntimeError( + f"{status_code} HTTP Error received from Resource Manager: " + f"{message}" + ) + + def __init__( + self, + *, + timeout: int = 30, + headers: Optional[dict[str, str]] = None, + session: Optional[requests.Session] = None, + cloud_connection: Optional[CloudConnection] = None, + ) -> None: + """Initialize the Resource Manager client. + + Args: + timeout: Request timeout in seconds. + headers: Optional headers included with every request. + session: Optional preconfigured requests session. + cloud_connection: Optional Cloud API connection override. + """ + self._base_url: Optional[str] = None + self._timeout = timeout + self._session = session or requests.Session() + self._configure_session(self._session) + self._headers = headers or {} + self._cloud_connection = cloud_connection + self._access_token: Optional[str] = None + self._access_token_expires_at: Optional[datetime] = None + self._resource_manager_id: Optional[UUID] = None + self._organization_id: Optional[UUID] = None + self._workspace_id: Optional[UUID] = None + self._lock = RLock() + + def create_request( + self, request: RMResourceRequestCreate + ) -> RMResourceRequestResponse: + """Create a runtime Resource Manager request. + + Args: + request: Runtime request create payload. + + Returns: + The created runtime request. + """ + return self._request_model( + "POST", + "/v1/resource-requests", + RMResourceRequestResponse, + json=request, + ) + + def get_request(self, request_id: UUID) -> RMResourceRequestResponse: + """Fetch a runtime Resource Manager request. + + Args: + request_id: Runtime request ID. + + Returns: + The requested runtime request. + """ + return self._request_model( + "GET", + f"/v1/resource-requests/{request_id}", + RMResourceRequestResponse, + ) + + def list_requests( + self, + *, + request_id: Optional[UUID] = None, + subject_id: Optional[UUID] = None, + statuses: Optional[list[str]] = None, + pool_id: Optional[UUID] = None, + reclaim_tolerance: Optional[str] = None, + preemption_initiated_by_id: Optional[UUID] = None, + metadata: Optional[dict[str, str]] = None, + ) -> RMResourceRequestListResponse: + """List runtime Resource Manager requests. + + Args: + request_id: When set, return only the request with this id. + subject_id: When set, return only requests that include this + subject id. + statuses: When set, return only requests in any of these states. + pool_id: When set, return only requests linked to this pool via an + allocation or queue entry. + reclaim_tolerance: When set, return only requests with this + reclaim tolerance. + preemption_initiated_by_id: When set, return only requests + preempted by this request id. + metadata: Optional exact-match filters against opaque entity + metadata. Each entry is sent as ``metadata[key]=value``. + + Returns: + A list wrapper containing matching runtime requests. + """ + params = self._list_params( + request_id=request_id, + subject_id=subject_id, + statuses=statuses, + pool=pool_id, + reclaim_tolerance=reclaim_tolerance, + preemption_initiated_by_id=preemption_initiated_by_id, + metadata=metadata, + ) + return self._request_model( + "GET", + "/v1/resource-requests", + RMResourceRequestListResponse, + params=params, + ) + + def release_request(self, request_id: UUID) -> RMResourceRequestResponse: + """Release a runtime Resource Manager request on behalf of its owner. + + Args: + request_id: Runtime request ID. + + Returns: + The released runtime request. + """ + return self._request_model( + "POST", + f"/v1/resource-requests/{request_id}/release", + RMResourceRequestResponse, + ) + + def renew_request( + self, + request_id: UUID, + *, + lease_expires_at: datetime, + runtime_state: Optional[ResourceRequestRuntimeState] = None, + ) -> RMResourceRequestResponse: + """Renew a runtime Resource Manager request lease. + + Args: + request_id: Runtime request ID. + lease_expires_at: Renewed lease expiration timestamp. + runtime_state: Optional owner-reported runtime state. + + Returns: + The renewed runtime request. + """ + return self._request_model( + "POST", + f"/v1/resource-requests/{request_id}/renew", + RMResourceRequestResponse, + json=RMResourceRequestRenewalRequest( + lease_expires_at=lease_expires_at, + runtime_state=runtime_state, + ), + ) + + def _request_model( + self, + method: str, + path: str, + response_model: Type[ModelT], + *, + json: Optional[BaseModel] = None, + params: Optional[dict[str, str | list[str]]] = None, + ) -> ModelT: + """Send an HTTP request and parse a Pydantic response model. + + Args: + method: HTTP method to use. + path: API path relative to the Resource Manager base URL. + response_model: Pydantic model used to parse the response. + json: Optional Pydantic request body. + params: Optional query string parameters. + + Returns: + Parsed response model. + """ + response = self._request(method, path, json=json, params=params) + return response_model.model_validate(response.json()) + + def _request( + self, + method: str, + path: str, + *, + json: Optional[BaseModel] = None, + params: Optional[dict[str, str | list[str]]] = None, + ) -> requests.Response: + """Send an HTTP request to the Resource Manager service. + + Args: + method: HTTP method to use. + path: API path relative to the Resource Manager base URL. + json: Optional Pydantic request body. + params: Optional query string parameters. + + Returns: + Raw requests response object. + + Raises: + KeyError: If the server returns 404. + EntityExistsError: If the server returns 409. + ValueError: If the server returns 400 or 422. + IllegalOperationError: If the server returns 403. + CredentialsNotValid: If the server returns 401. + RuntimeError: For other HTTP error status codes. + """ # noqa: DOC503 + payload: Any = None + if json is not None: + payload = json.model_dump(mode="json", by_alias=True) + + idempotency_key = str(uuid4()) + headers = self._request_headers(idempotency_key=idempotency_key) + response = self._session.request( + method=method, + url=f"{self._resolved_base_url}{path}", + headers=headers, + json=payload, + params=params, + timeout=self._timeout, + ) + if response.status_code == 401: + self._reset_login() + headers = self._request_headers(idempotency_key=idempotency_key) + response = self._session.request( + method=method, + url=f"{self._resolved_base_url}{path}", + headers=headers, + json=payload, + params=params, + timeout=self._timeout, + ) + if response.status_code >= 400: + raise self.exception_from_rm_response(response) + return response + + @property + def _resolved_base_url(self) -> str: + """Return the currently known Resource Manager base URL. + + Raises: + RuntimeError: If no Resource Manager URL is available. + + Returns: + The Resource Manager base URL. + """ + if self._base_url is None: + raise RuntimeError( + "Resource Manager URL has not been discovered through the " + "Cloud API authorization exchange." + ) + return self._base_url + + def _request_headers(self, *, idempotency_key: str) -> dict[str, str]: + """Build headers for a Resource Manager request. + + Args: + idempotency_key: Idempotency key to send with this request. + + Returns: + HTTP headers for the request. + """ + headers = { + **self._headers, + self.IDEMPOTENCY_HEADER: idempotency_key, + } + headers["Authorization"] = f"Bearer {self._fetch_access_token()}" + return headers + + def _fetch_access_token(self) -> str: + """Fetch or refresh the Resource Manager local access token. + + Returns: + A Resource Manager bearer token. + """ + with self._lock: + if ( + self._access_token is not None + and self._access_token_expires_at is not None + and utc_now(tz_aware=self._access_token_expires_at) + + timedelta(minutes=5) + < self._access_token_expires_at + ): + return self._access_token + + authorization = self._fetch_authorization() + self._base_url = authorization.api_url.rstrip("/") + login = self._login_to_resource_manager(authorization) + + self._access_token = login.access_token + self._access_token_expires_at = login.expires_at + self._resource_manager_id = login.resource_manager_id + self._organization_id = login.organization_id + self._workspace_id = login.workspace_id + return self._access_token + + def _reset_login(self) -> None: + """Force a new Resource Manager login on the next request.""" + with self._lock: + self._access_token = None + self._access_token_expires_at = None + + def _fetch_authorization(self) -> ResourceManagerAuthorization: + """Fetch Resource Manager discovery and one-time auth from Cloud API. + + Returns: + Cloud API authorization response. + + Raises: + RuntimeError: If Cloud API returns an invalid authorization + response. + """ + cloud_connection = self._cloud_connection + if cloud_connection is None: + from zenml.zen_server.cloud_utils import ( + cloud_connection as connect, + ) + + cloud_connection = self._cloud_connection = connect() + + response = cloud_connection.get( + "/auth/resource_manager_authorization", + params=None, + ) + try: + return ResourceManagerAuthorization.model_validate(response.json()) + except Exception as e: + raise RuntimeError( + "Could not fetch Resource Manager authorization from the " + f"Cloud API: {e}" + ) + + def _login_to_resource_manager( + self, authorization: ResourceManagerAuthorization + ) -> ResourceManagerLoginResponse: + """Exchange a Cloud API one-time token for an RM-local session. + + Args: + authorization: Cloud API authorization response. + + Raises: + self.exception_from_rm_response: If Resource Manager rejects the + login request. + RuntimeError: If the RM returns an invalid login response. + + Returns: + Resource Manager local session response. + """ + login_request = ResourceManagerLoginRequest( + organization_id=authorization.organization_id, + workspace_id=authorization.workspace_id, + ) + response = self._session.post( + f"{self._resolved_base_url}/v1/auth/login", + headers={ + "Authorization": f"Bearer {authorization.token}", + }, + json=login_request.model_dump(mode="json"), + timeout=self._timeout, + ) + if response.status_code >= 400: + raise self.exception_from_rm_response(response) + + try: + return ResourceManagerLoginResponse.model_validate(response.json()) + except Exception as e: + raise RuntimeError( + f"Could not log in to the Resource Manager API: {e}" + ) + + def _configure_session(self, session: requests.Session) -> None: + """Configure session retry behavior for transient HTTP failures. + + This mirrors the REST ZenStore client behavior so transient network + and upstream availability issues are retried consistently. + + Args: + session: Session to configure. + """ + retries = Retry( + connect=5, + read=8, + redirect=3, + status=10, + allowed_methods=[ + "HEAD", + "GET", + "POST", + "PUT", + "PATCH", + "DELETE", + "OPTIONS", + ], + status_forcelist=[ + 408, # Request Timeout + 429, # Too Many Requests + 502, # Bad Gateway + 503, # Service Unavailable + 504, # Gateway Timeout + ], + other=3, + backoff_factor=1, + ) + http_adapter = HTTPAdapter(max_retries=retries) + session.mount("https://", http_adapter) + session.mount("http://", http_adapter) + + @staticmethod + def _list_params(**values: object) -> Optional[dict[str, str | list[str]]]: + """Build optional query parameters for list endpoints. + + Scalar values become one query parameter. Sequence values become + repeated query parameters. Mapping values become ``field[key]=value`` + bracket query parameters. + + Args: + **values: Optional filter values keyed by query parameter name. + + Returns: + Query parameters with unset values omitted, or ``None`` when empty. + """ + params: dict[str, str | list[str]] = {} + for key, value in values.items(): + if value is None: + continue + if isinstance(value, dict): + for inner_key, inner_value in value.items(): + params[f"{key}[{inner_key}]"] = str(inner_value) + continue + if isinstance(value, (list, tuple)): + params[key] = [str(entry) for entry in value] + continue + params[key] = str(value) + return params or None diff --git a/src/zenml/zen_stores/resource_pools/resource_manager/store.py b/src/zenml/zen_stores/resource_pools/resource_manager/store.py new file mode 100644 index 00000000000..e4f54cff517 --- /dev/null +++ b/src/zenml/zen_stores/resource_pools/resource_manager/store.py @@ -0,0 +1,562 @@ +# Copyright (c) ZenML GmbH 2026. All Rights Reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at: +# +# https://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express +# or implied. See the License for the specific language governing +# permissions and limitations under the License. +"""Resource pool store backed by the ZenML Pro Resource Manager service.""" + +import json +from typing import TYPE_CHECKING, Any, Optional, Sequence, TypeVar +from uuid import UUID + +from pydantic import BaseModel +from pydantic_settings import BaseSettings, SettingsConfigDict + +from zenml.config.server_config import ServerProConfiguration +from zenml.enums import ResourceRequestStatus +from zenml.exceptions import ( + CredentialsNotValid, + EntityExistsError, + IllegalOperationError, +) +from zenml.logger import get_logger +from zenml.models import ( + Page, + ResourceRequestFilter, + ResourceRequestRequest, + ResourceRequestResponse, + UserFilter, + UserResponse, +) +from zenml.models.v2.base.scoped import UserScopedResponse +from zenml.models.v2.core.resource_request import ( + ResourceRequestRenewalRequest, +) +from zenml.zen_stores.resource_pools.resource_manager.client import ( + ResourceManagerClient, +) +from zenml.zen_stores.resource_pools.resource_manager.transport import ( + PIPELINE_RUN_ID_METADATA_KEY, + PIPELINE_RUN_NAME_METADATA_KEY, + PROJECT_ID_METADATA_KEY, + STEP_NAME_METADATA_KEY, + STEP_RUN_ID_METADATA_KEY, + RMResourceRequestCreate, + RMSubject, + RMSubjectSelector, +) +from zenml.zen_stores.resource_pools.store_interface import ( + ResourcePoolsSQLStoreInterface, +) + +if TYPE_CHECKING: + from zenml.zen_stores.sql_zen_store import Session, SqlZenStore + +logger = get_logger(__name__) +UserScopedResponseT = TypeVar( + "UserScopedResponseT", bound=UserScopedResponse[Any, Any, Any] +) +USER_ID_METADATA_KEY = "user_id" +ORGANIZATION_ID_METADATA_KEY = "organization_id" +WORKSPACE_ID_METADATA_KEY = "workspace_id" + + +class ResourceManagerResourcePoolsStoreSettings(BaseSettings): + """Settings used by the Resource Manager resource pool store.""" + + timeout: int = 30 + + model_config = SettingsConfigDict( + env_prefix="ZENML_PRO_RM_", + extra="ignore", + ) + + +class ResourceManagerResourcePoolsStore(ResourcePoolsSQLStoreInterface): + """Resource pool store facade for ZenML Pro Resource Manager.""" + + def __init__( + self, + store: "SqlZenStore", + client: Optional[ResourceManagerClient] = None, + ) -> None: + """Initialize the Resource Manager-backed resource pool store. + + Args: + store: SQL ZenStore that owns this backend. + client: Optional Resource Manager client. If omitted, the client is + built from Resource Manager store settings. + """ + super().__init__(store=store) + settings = ResourceManagerResourcePoolsStoreSettings() + self._client = client or ResourceManagerClient( + timeout=settings.timeout, + ) + + def get_resource_request( + self, resource_request_id: UUID, hydrate: bool = True + ) -> ResourceRequestResponse: + """Get a resource request through Resource Manager. + + Args: + resource_request_id: The request ID. + hydrate: Ignored for Resource Manager-backed requests. + + Returns: + The requested ZenML resource request response. + """ + response = self._client.get_request(resource_request_id) + logger.info( + f"Resource request response: {response.model_dump_json(indent=2)}" + ) + return response.to_model() + + def list_resource_requests( + self, filter_model: ResourceRequestFilter, hydrate: bool = False + ) -> Page[ResourceRequestResponse]: + """List resource requests through Resource Manager. + + Args: + filter_model: ZenML filter model. Pagination fields are currently + ignored because Resource Manager returns a full page. + hydrate: Ignored for Resource Manager-backed requests. + + Returns: + Matching ZenML request responses. + """ + logger.debug( + f"Listing resource requests with filter model: {filter_model}" + ) + list_kwargs: dict[str, Any] = {} + if filter_model.id is not None: + list_kwargs["request_id"] = UUID(str(filter_model.id)) + if filter_model.component_id is not None: + list_kwargs["subject_id"] = UUID(str(filter_model.component_id)) + elif filter_model.user is not None: + owner = self.store.get_user(filter_model.user, hydrate=False) + if owner.external_user_id is not None: + list_kwargs["subject_id"] = owner.external_user_id + if filter_model.status is not None: + if isinstance(filter_model.status, ResourceRequestStatus): + list_kwargs["statuses"] = [filter_model.status.value] + else: + operator, _, value = str(filter_model.status).partition(":") + if operator == "oneof": + list_kwargs["statuses"] = json.loads(value) + else: + list_kwargs["statuses"] = [str(filter_model.status)] + if filter_model.pool_id is not None: + list_kwargs["pool_id"] = UUID(str(filter_model.pool_id)) + if filter_model.reclaim_tolerance is not None: + list_kwargs["reclaim_tolerance"] = str( + filter_model.reclaim_tolerance + ) + if filter_model.preemption_initiated_by_id is not None: + list_kwargs["preemption_initiated_by_id"] = UUID( + str(filter_model.preemption_initiated_by_id) + ) + + metadata: dict[str, str] = {} + if filter_model.step_run_id is not None: + metadata[STEP_RUN_ID_METADATA_KEY] = str(filter_model.step_run_id) + if filter_model.pipeline_run_id is not None: + metadata[PIPELINE_RUN_ID_METADATA_KEY] = str( + filter_model.pipeline_run_id + ) + if metadata: + list_kwargs["metadata"] = metadata + + _ = hydrate + responses = self._client.list_requests(**list_kwargs).items + items = self._responses_with_users( + responses, [response.to_model() for response in responses] + ) + return Page( + index=1, + max_size=max(len(items), 1), + total_pages=1, + total=len(items), + items=items, + ) + + def release_resource_request( + self, + resource_request_id: UUID, + ) -> ResourceRequestResponse: + """Release a resource request through Resource Manager. + + Owner-initiated release marks allocated requests ``released``, + pending requests ``cancelled``, and preempting requests ``preempted``. + + Args: + resource_request_id: The request ID. + + Returns: + The released resource request. + """ + response = self._client.release_request(resource_request_id) + return response.to_model() + + def renew_resource_request( + self, + resource_request_id: UUID, + renewal_request: ResourceRequestRenewalRequest, + ) -> ResourceRequestResponse: + """Renew a resource request lease through Resource Manager. + + Args: + resource_request_id: The request ID. + renewal_request: Renewal payload. + + Returns: + The renewed resource request. + """ + response = self._client.renew_request( + resource_request_id, + lease_expires_at=renewal_request.lease_expires_at, + runtime_state=renewal_request.runtime_state, + ) + logger.debug( + f"Renewed resource request: {response.model_dump_json(indent=2)}" + ) + return response.to_model() + + def release_step_run_resources( + self, session: "Session", step_run_id: UUID + ) -> None: + """Release Resource Manager allocations for a ZenML step run. + + Args: + session: DB session. The Resource Manager backend does not use the + session, but the SQL store callback requires it. + step_run_id: The ZenML step run whose resources should be released. + """ + from zenml.zen_stores.schemas.step_run_schemas import StepRunSchema + + step_run = session.get(StepRunSchema, step_run_id) + if step_run is None or step_run.resource_request_id is None: + return + + self._client.release_request(step_run.resource_request_id) + + def create_resource_request( + self, + session: "Session", + resource_request: ResourceRequestRequest, + ) -> ResourceRequestResponse: + """Create a step-run resource request with session-backed metadata. + + Enriches the request metadata with step and pipeline-run context read + from the in-flight session (the step run is not yet committed and is + only visible through this session), then submits directly to Resource + Manager. + + Args: + session: DB session used to enrich step-run metadata. + resource_request: The ZenML resource request payload. + + Returns: + The created ZenML resource request response. + + Raises: + ValueError: If the request does not contain the required step-run, + pipeline, or user context. + KeyError: If the referenced step run or pipeline run does not + exist. + EntityExistsError: If Resource Manager rejects the request because + a matching entity already exists. + IllegalOperationError: If Resource Manager rejects the request as + invalid for the current state. + CredentialsNotValid: If Resource Manager authentication is no + longer valid. + RuntimeError: If Resource Manager request creation fails. + """ + if resource_request.step_run_id is None: + raise ValueError( + "step_run_id is required when creating Resource Manager " + "resource requests." + ) + + from zenml.zen_stores.schemas.pipeline_run_schemas import ( + PipelineRunSchema, + ) + from zenml.zen_stores.schemas.step_run_schemas import StepRunSchema + + step_run = session.get(StepRunSchema, resource_request.step_run_id) + if step_run is None: + raise KeyError( + f"Step run '{resource_request.step_run_id}' does not exist." + ) + + pipeline_run_schema = session.get( + PipelineRunSchema, step_run.pipeline_run_id + ) + if pipeline_run_schema is None: + raise KeyError( + f"Pipeline run '{step_run.pipeline_run_id}' does not exist." + ) + pipeline_run = pipeline_run_schema.to_model(include_resources=True) + if pipeline_run.pipeline is None or pipeline_run.user is None: + raise ValueError( + "pipeline and user are required when creating Resource Manager " + "resource requests." + ) + + project = pipeline_run.project + user = pipeline_run.user + pipeline = pipeline_run.pipeline + + pro_config = ServerProConfiguration.get_server_config() + + subjects = [ + RMSubject.from_account( + user, + organization_id=pro_config.organization_id, + organization_name=pro_config.organization_name, + ), + RMSubject.from_pipeline( + organization_id=pro_config.organization_id, + workspace_id=pro_config.workspace_id, + organization_name=pro_config.organization_name, + workspace_name=pro_config.workspace_name, + project_id=pipeline_run.project_id, + project_name=project.name, + pipeline_id=pipeline.id, + pipeline_name=pipeline.name, + ), + RMSubject.from_step_run( + organization_id=pro_config.organization_id, + workspace_id=pro_config.workspace_id, + organization_name=pro_config.organization_name, + workspace_name=pro_config.workspace_name, + project_id=pipeline_run.project_id, + project_name=project.name, + pipeline_run_id=step_run.pipeline_run_id, + pipeline_run_name=pipeline_run.name, + pipeline_id=pipeline.id, + pipeline_name=pipeline.name, + step_run_id=step_run.id, + step_name=step_run.name, + ), + ] + if resource_request.component_ids: + connector_ids: set[UUID] = set() + for component_id in resource_request.component_ids: + component = self.store.get_stack_component( + component_id, hydrate=True + ) + subjects.append( + RMSubject.from_component( + component, + organization_id=pro_config.organization_id, + workspace_id=pro_config.workspace_id, + organization_name=pro_config.organization_name, + workspace_name=pro_config.workspace_name, + ) + ) + connector = component.connector + if connector is None or connector.id in connector_ids: + continue + connector_ids.add(connector.id) + effective_resource_type = ( + component.flavor.connector_resource_type + ) + effective_resource_id = ( + component.connector_resource_id or connector.resource_id + ) + subjects.append( + RMSubject.from_service_connector( + connector, + organization_id=pro_config.organization_id, + workspace_id=pro_config.workspace_id, + organization_name=pro_config.organization_name, + workspace_name=pro_config.workspace_name, + effective_resource_type=effective_resource_type, + effective_resource_id=effective_resource_id, + ) + ) + + metadata = dict(resource_request.metadata) + metadata[STEP_NAME_METADATA_KEY] = step_run.name + metadata[PIPELINE_RUN_ID_METADATA_KEY] = str(step_run.pipeline_run_id) + metadata[PIPELINE_RUN_NAME_METADATA_KEY] = pipeline_run.name + metadata[PROJECT_ID_METADATA_KEY] = str(pipeline_run.project_id) + + try: + request = RMResourceRequestCreate.from_model( + resource_request.model_copy(update={"metadata": metadata}), + subjects=subjects, + preemption_group=RMSubjectSelector.from_pipeline_run( + organization_id=pro_config.organization_id, + workspace_id=pro_config.workspace_id, + project_id=pipeline_run.project_id, + pipeline_run_id=step_run.pipeline_run_id, + ), + user_id=user.external_user_id, + ) + request.metadata = self._metadata_with_server_context( + request.metadata, + user=user, + ) + logger.info(f"Creating resource request: {request}") + + response = self._client.create_request(request) + except KeyError as exc: + raise KeyError( + "An error occurred while trying to create a resource request: " + f"{exc}" + ) from exc + except ValueError as exc: + raise ValueError( + "An error occurred while trying to create a resource request: " + f"{exc}" + ) from exc + except EntityExistsError as exc: + raise EntityExistsError( + "An error occurred while trying to create a resource request: " + f"{exc}" + ) from exc + except IllegalOperationError as exc: + raise IllegalOperationError( + "An error occurred while trying to create a resource request: " + f"{exc}" + ) from exc + except CredentialsNotValid as exc: + raise CredentialsNotValid( + "An error occurred while trying to create a resource request: " + f"{exc}" + ) from exc + except RuntimeError as exc: + raise RuntimeError( + "An error occurred while trying to create a resource request: " + f"{exc}" + ) from exc + + logger.debug( + f"Created resource request: {response.model_dump_json(indent=2)}" + ) + return response.to_model() + + def _metadata_with_server_context( + self, + metadata: dict[str, Any], + user: UserResponse | None = None, + ) -> dict[str, Any]: + """Add immutable ZenML Pro provenance metadata to create payloads. + + Args: + metadata: User-supplied metadata to preserve. + user: The ZenML Pro user to use for metadata. If not provided, the + current active user is used. + + Returns: + Metadata enriched with ZenML Pro user, organization, and workspace + identifiers. + """ + pro_config = ServerProConfiguration.get_server_config() + server_metadata = { + ORGANIZATION_ID_METADATA_KEY: str(pro_config.organization_id), + WORKSPACE_ID_METADATA_KEY: str(pro_config.workspace_id), + } + if user is None: + with self.store.get_session() as session: + user = self.store._get_active_user(session) + if user is not None and user.external_user_id is not None: + server_metadata[USER_ID_METADATA_KEY] = str(user.external_user_id) + + return {**metadata, **server_metadata} + + def _responses_with_users( + self, + responses: Sequence[BaseModel], + models: list[UserScopedResponseT], + ) -> list[UserScopedResponseT]: + """Populate ZenML user fields from Resource Manager metadata in bulk. + + Args: + responses: Resource Manager responses that contain provenance + metadata. + models: ZenML response models converted from the Resource Manager + responses. + + Returns: + ZenML response models with user fields populated where possible. + """ + external_user_ids_by_index: dict[int, UUID] = {} + external_user_ids: set[UUID] = set() + for index, response in enumerate(responses): + external_user_id = self._external_user_id_from_metadata( + getattr(response, "metadata", {}) + ) + if external_user_id is None: + continue + external_user_ids_by_index[index] = external_user_id + external_user_ids.add(external_user_id) + + if not external_user_ids: + return models + + users = self._users_by_external_id(external_user_ids) + for index, external_user_id in external_user_ids_by_index.items(): + user = users.get(external_user_id) + if user is None: + continue + models[index].get_body().user_id = user.id + models[index].get_resources().user = user + + return models + + def _users_by_external_id( + self, external_user_ids: set[UUID] + ) -> dict[UUID, UserResponse]: + """Fetch ZenML users keyed by external ZenML Pro user ID. + + Args: + external_user_ids: External ZenML Pro user IDs to resolve. + + Returns: + ZenML users keyed by their external ZenML Pro user ID. + """ + external_user_id_filter = "oneof:" + json.dumps( + sorted(str(user_id) for user_id in external_user_ids) + ) + user_page = self.store.list_users( + UserFilter( + external_user_id=external_user_id_filter, + size=max(len(external_user_ids), 1), + ), + hydrate=True, + ) + return { + user.external_user_id: user + for user in user_page.items + if user.external_user_id is not None + } + + @staticmethod + def _external_user_id_from_metadata( + metadata: dict[str, Any], + ) -> Optional[UUID]: + """Extract an external ZenML Pro user ID from RM metadata. + + Args: + metadata: Resource Manager metadata values. + + Returns: + External ZenML Pro user ID if present and valid. + """ + value = metadata.get(USER_ID_METADATA_KEY) + if value is None: + return None + try: + return UUID(str(value)) + except ValueError: + return None diff --git a/src/zenml/zen_stores/resource_pools/resource_manager/transport.py b/src/zenml/zen_stores/resource_pools/resource_manager/transport.py new file mode 100644 index 00000000000..2d23584357a --- /dev/null +++ b/src/zenml/zen_stores/resource_pools/resource_manager/transport.py @@ -0,0 +1,958 @@ +# Copyright (c) ZenML GmbH 2026. All Rights Reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at: +# +# https://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express +# or implied. See the License for the specific language governing +# permissions and limitations under the License. +"""Internal Resource Manager transport models for runtime requests.""" + +from datetime import datetime +from typing import TYPE_CHECKING, Any, Optional, Union +from uuid import UUID + +from pydantic import BaseModel, ConfigDict, Field + +from zenml.enums import ( + ResourceRequestRuntimeState, + ResourceRequestStatus, +) +from zenml.models import ( + ResourcePoolAllocation, + ResourcePoolQueueItem, + ResourceRequestDemand, + ResourceRequestReclaimTolerance, + ResourceRequestRequest, + ResourceRequestResponse, + ResourceRequestResponseBody, + ResourceRequestResponseMetadata, + ResourceRequestResponseResources, + ResourceRequestServiceConnectorSettings, +) +from zenml.utils.time_utils import utc_now + +if TYPE_CHECKING: + from zenml.models import ( + ComponentResponse, + ServiceConnectorResponse, + UserResponse, + ) + +PRIORITY_LANE_PRIORITY = 2_147_483_647 +ORGANIZATION_SUBJECT_TYPE = "organization" +WORKSPACE_SUBJECT_TYPE = "workspace" +PROJECT_SUBJECT_TYPE = "project" +PIPELINE_SUBJECT_TYPE = "pipeline" +PIPELINE_RUN_SUBJECT_TYPE = "pipeline_run" +STEP_RUN_SUBJECT_TYPE = "step_run" +COMPONENT_SUBJECT_TYPE = "component" +SERVICE_CONNECTOR_SUBJECT_TYPE = "service_connector" +ACCOUNT_SUBJECT_TYPE = "account" +STEP_RUN_ID_METADATA_KEY = "step_run_id" +PIPELINE_RUN_ID_METADATA_KEY = "pipeline_run_id" +STEP_NAME_METADATA_KEY = "step_name" +PIPELINE_RUN_NAME_METADATA_KEY = "pipeline_run_name" +PROJECT_ID_METADATA_KEY = "project_id" + + +def _metadata_uuid(metadata: dict[str, Any], key: str) -> Optional[UUID]: + """Read a UUID metadata value when present. + + Args: + metadata: Metadata values to read from. + key: Metadata key to read. + + Returns: + Metadata value converted to a UUID, or `None` if it is absent or + invalid. + """ + value = metadata.get(key) + if value is None: + return None + try: + return UUID(str(value)) + except ValueError: + return None + + +class RMSubjectSelector(BaseModel): + """Structured selector used in Resource Manager subject settings.""" + + subject_type: Optional[str] = None + subject_id: Optional[UUID] = None + attributes: dict[str, Any] = Field(default_factory=dict) + metadata: dict[str, Any] = Field(default_factory=dict) + contains: Optional[ + Union["RMSubjectSelector", "RMSubjectSelectorExpression"] + ] = None + + @classmethod + def from_pipeline_run( + cls, + *, + organization_id: UUID, + workspace_id: UUID, + project_id: UUID, + pipeline_run_id: UUID, + ) -> "RMSubjectSelector": + """Build a selector for one scoped ZenML pipeline run. + + Args: + organization_id: ZenML Pro organization ID. + workspace_id: ZenML Pro workspace ID. + project_id: ZenML project ID. + pipeline_run_id: ZenML pipeline run ID. + + Returns: + Resource Manager subject selector matching the pipeline run. + """ + return cls( + subject_id=organization_id, + subject_type=ORGANIZATION_SUBJECT_TYPE, + contains=cls( + subject_id=workspace_id, + subject_type=WORKSPACE_SUBJECT_TYPE, + contains=cls( + subject_id=project_id, + subject_type=PROJECT_SUBJECT_TYPE, + contains=cls( + subject_id=pipeline_run_id, + subject_type=PIPELINE_RUN_SUBJECT_TYPE, + ), + ), + ), + ) + + +class RMSubjectSelectorExpression(BaseModel): + """Boolean expression over Resource Manager subject selectors.""" + + all: list["RMSubjectSelectorNode"] = Field(default_factory=list) + any: list["RMSubjectSelectorNode"] = Field(default_factory=list) + not_: Optional["RMSubjectSelectorNode"] = Field(default=None, alias="not") + + model_config = ConfigDict(populate_by_name=True) + + +RMSubjectSelectorNode = Union[RMSubjectSelector, RMSubjectSelectorExpression] + + +class RMSubject(BaseModel): + """Subject on runtime resource request payloads.""" + + subject_id: UUID + subject_type: str + attributes: dict[str, Any] = Field(default_factory=dict) + metadata: dict[str, Any] = Field(default_factory=dict) + child: Optional["RMSubject"] = None + + @property + def leaf(self) -> "RMSubject": + """Return the deepest subject in the root-first chain.""" + if self.child is None: + return self + return self.child.leaf + + def find(self, subject_id: UUID) -> Optional["RMSubject"]: + """Find a subject by ID in this chain. + + Args: + subject_id: Subject ID to find. + + Returns: + Matching subject, or `None` if no subject in the chain matches. + """ + if self.subject_id == subject_id: + return self + if self.child is None: + return None + return self.child.find(subject_id) + + @staticmethod + def _optional_values(**values: Any) -> dict[str, Any]: + """Build Resource Manager payload data without empty optional values. + + Args: + values: Optional values to include in the payload. + + Returns: + Payload data without `None` or empty-string values. + """ + payload: dict[str, Any] = {} + for key, value in values.items(): + if value is None: + continue + if isinstance(value, str) and not value: + continue + payload[key] = str(value) if isinstance(value, UUID) else value + return payload + + @classmethod + def _organization( + cls, + *, + organization_id: UUID, + organization_name: Optional[str] = None, + child: Optional["RMSubject"] = None, + ) -> "RMSubject": + """Build an organization-rooted subject chain. + + Args: + organization_id: ZenML Pro organization ID. + organization_name: Optional ZenML Pro organization name. + child: Optional child subject node. + + Returns: + Organization-rooted Resource Manager subject chain. + """ + return cls( + subject_id=organization_id, + subject_type=ORGANIZATION_SUBJECT_TYPE, + attributes=cls._optional_values(name=organization_name), + child=child, + ) + + @classmethod + def _workspace( + cls, + *, + workspace_id: UUID, + workspace_name: Optional[str] = None, + child: Optional["RMSubject"] = None, + ) -> "RMSubject": + """Build a workspace subject chain node. + + Args: + workspace_id: ZenML Pro workspace ID. + workspace_name: Optional ZenML Pro workspace name. + child: Optional child subject node. + + Returns: + Workspace Resource Manager subject chain node. + """ + return cls( + subject_id=workspace_id, + subject_type=WORKSPACE_SUBJECT_TYPE, + attributes=cls._optional_values(name=workspace_name), + child=child, + ) + + @classmethod + def from_component( + cls, + component: "ComponentResponse", + *, + organization_id: UUID, + workspace_id: UUID, + organization_name: Optional[str] = None, + workspace_name: Optional[str] = None, + ) -> "RMSubject": + """Build an inline subject for a ZenML stack component. + + Args: + component: ZenML stack component response. + organization_id: ZenML Pro organization ID. + workspace_id: ZenML Pro workspace ID. + organization_name: Optional ZenML Pro organization name. + workspace_name: Optional ZenML Pro workspace name. + + Returns: + Resource Manager subject chain for the stack component. + """ + return cls._organization( + organization_id=organization_id, + organization_name=organization_name, + child=cls._workspace( + workspace_id=workspace_id, + workspace_name=workspace_name, + child=cls( + subject_id=component.id, + subject_type=COMPONENT_SUBJECT_TYPE, + attributes=cls._optional_values( + name=component.name, + stack_component_name=component.name, + workspace_id=workspace_id, + workspace_name=workspace_name, + component_type=component.type.value, + flavor=component.flavor_name, + ), + metadata=cls._optional_values(logo_url=component.logo_url), + ), + ), + ) + + @classmethod + def from_service_connector( + cls, + connector: "ServiceConnectorResponse", + *, + organization_id: UUID, + workspace_id: UUID, + organization_name: Optional[str] = None, + workspace_name: Optional[str] = None, + effective_resource_type: Optional[str] = None, + effective_resource_id: Optional[str] = None, + ) -> "RMSubject": + """Build an inline subject for a workspace service connector. + + Args: + connector: ZenML service connector response. + organization_id: ZenML Pro organization ID. + workspace_id: ZenML Pro workspace ID. + organization_name: Optional ZenML Pro organization name. + workspace_name: Optional ZenML Pro workspace name. + effective_resource_type: Optional service connector resource type. + effective_resource_id: Optional service connector resource ID. + + Returns: + Resource Manager subject chain for the service connector. + """ + return cls._organization( + organization_id=organization_id, + organization_name=organization_name, + child=cls._workspace( + workspace_id=workspace_id, + workspace_name=workspace_name, + child=cls( + subject_id=connector.id, + subject_type=SERVICE_CONNECTOR_SUBJECT_TYPE, + attributes=cls._optional_values( + name=connector.name, + service_connector_name=connector.name, + connector_type=connector.type, + resource_types=connector.resource_types, + effective_resource_type=effective_resource_type, + effective_resource_id=effective_resource_id, + ), + ), + ), + ) + + @classmethod + def from_account( + cls, + user: "UserResponse", + *, + organization_id: UUID, + organization_name: Optional[str] = None, + ) -> "RMSubject": + """Build an inline subject for a ZenML account identity. + + Args: + user: ZenML user response. + organization_id: ZenML Pro organization ID. + organization_name: Optional ZenML Pro organization name. + + Returns: + Resource Manager subject for the account identity. + + Raises: + ValueError: If the ZenML user has no external account ID. + """ + if user.external_user_id is None: + raise ValueError( + f"User '{user.id}' has no external account id configured." + ) + account_subject = cls( + subject_id=user.external_user_id, + subject_type=ACCOUNT_SUBJECT_TYPE, + attributes=cls._optional_values( + name=user.full_name or user.name, + email=user.email, + username=user.name, + is_superuser=user.is_admin, + is_service_account=user.is_service_account, + ), + metadata=cls._optional_values(avatar_url=user.avatar_url), + ) + if not user.is_service_account: + return account_subject + + return cls._organization( + organization_id=organization_id, + organization_name=organization_name, + child=account_subject, + ) + + @classmethod + def from_pipeline( + cls, + *, + organization_id: UUID, + workspace_id: UUID, + project_id: UUID, + pipeline_id: UUID, + organization_name: Optional[str] = None, + workspace_name: Optional[str] = None, + project_name: Optional[str] = None, + pipeline_name: Optional[str] = None, + ) -> "RMSubject": + """Build an organization -> workspace -> project -> pipeline subject. + + Args: + organization_id: ZenML Pro organization ID. + workspace_id: ZenML Pro workspace ID. + project_id: ZenML project ID. + pipeline_id: ZenML pipeline ID. + organization_name: Optional ZenML Pro organization name. + workspace_name: Optional ZenML Pro workspace name. + project_name: Optional ZenML project name. + pipeline_name: Optional ZenML pipeline name. + + Returns: + Resource Manager subject chain for the pipeline. + """ + return cls._organization( + organization_id=organization_id, + organization_name=organization_name, + child=cls._workspace( + workspace_id=workspace_id, + workspace_name=workspace_name, + child=cls( + subject_id=project_id, + subject_type=PROJECT_SUBJECT_TYPE, + attributes=cls._optional_values( + name=project_name, + workspace_id=workspace_id, + workspace_name=workspace_name, + ), + child=cls( + subject_id=pipeline_id, + subject_type=PIPELINE_SUBJECT_TYPE, + attributes=cls._optional_values( + name=pipeline_name, + project_id=project_id, + project_name=project_name, + workspace_id=workspace_id, + workspace_name=workspace_name, + ), + ), + ), + ), + ) + + @classmethod + def from_step_run( + cls, + *, + organization_id: UUID, + workspace_id: UUID, + project_id: UUID, + pipeline_run_id: UUID, + step_run_id: UUID, + organization_name: Optional[str] = None, + workspace_name: Optional[str] = None, + project_name: Optional[str] = None, + pipeline_run_name: Optional[str] = None, + pipeline_id: Optional[UUID] = None, + pipeline_name: Optional[str] = None, + step_name: Optional[str] = None, + ) -> "RMSubject": + """Build a scoped pipeline-run -> step-run subject chain. + + Args: + organization_id: ZenML Pro organization ID. + workspace_id: ZenML Pro workspace ID. + project_id: ZenML project ID. + pipeline_run_id: ZenML pipeline run ID. + step_run_id: ZenML step run ID. + organization_name: Optional ZenML Pro organization name. + workspace_name: Optional ZenML Pro workspace name. + project_name: Optional ZenML project name. + pipeline_run_name: Optional ZenML pipeline run name. + pipeline_id: Optional ZenML pipeline ID. + pipeline_name: Optional ZenML pipeline name. + step_name: Optional ZenML step name. + + Returns: + Resource Manager subject chain for the step run. + """ + return cls._organization( + organization_id=organization_id, + organization_name=organization_name, + child=cls._workspace( + workspace_id=workspace_id, + workspace_name=workspace_name, + child=cls( + subject_id=project_id, + subject_type=PROJECT_SUBJECT_TYPE, + attributes=cls._optional_values( + name=project_name, + workspace_id=workspace_id, + workspace_name=workspace_name, + ), + child=cls( + subject_id=pipeline_run_id, + subject_type=PIPELINE_RUN_SUBJECT_TYPE, + attributes=cls._optional_values( + name=pipeline_run_name, + run_name=pipeline_run_name, + pipeline_id=pipeline_id, + pipeline_name=pipeline_name, + project_id=project_id, + project_name=project_name, + workspace_id=workspace_id, + workspace_name=workspace_name, + ), + child=cls( + subject_id=step_run_id, + subject_type=STEP_RUN_SUBJECT_TYPE, + attributes=cls._optional_values( + name=step_name, + step_name=step_name, + pipeline_run_id=pipeline_run_id, + pipeline_run_name=pipeline_run_name, + pipeline_id=pipeline_id, + pipeline_name=pipeline_name, + project_id=project_id, + project_name=project_name, + workspace_id=workspace_id, + workspace_name=workspace_name, + ), + ), + ), + ), + ), + ) + + +class RMRequestDemand(BaseModel): + """Resource demand payload for the Resource Manager API.""" + + resource_id: Optional[UUID] = Field(default=None, exclude=True) + resource: UUID | str | None = None + kind: Optional[str] = None + quantity: int + unit: Optional[str] = None + class_name: Optional[str] = Field( + default=None, + alias="class", + serialization_alias="class", + ) + resource_selector: Optional[dict[str, Any]] = None + class_selector: Optional[dict[str, Any]] = None + + model_config = ConfigDict(populate_by_name=True) + + @classmethod + def from_model(cls, demand: ResourceRequestDemand) -> "RMRequestDemand": + """Build a demand payload from a ZenML resource demand. + + Args: + demand: ZenML resource demand model. + + Returns: + Resource Manager demand payload. + """ + resource = ( + demand.resource_id + if demand.resource_id is not None + else demand.resource + ) + return cls( + resource=resource, + kind=demand.kind, + resource_selector=demand.resource_selector, + quantity=demand.quantity, + unit=demand.unit, + class_name=demand.class_name, + class_selector=demand.class_selector, + ) + + def to_model(self) -> ResourceRequestDemand: + """Convert this demand into a ZenML resource demand. + + Returns: + ZenML resource demand model. + """ + return ResourceRequestDemand( + resource_id=self.resource_id, + resource=self.resource, + kind=self.kind, + quantity=self.quantity, + unit=self.unit, + class_name=self.class_name, + resource_selector=self.resource_selector, + class_selector=self.class_selector, + ) + + +class RMResourceRequestCreate(BaseModel): + """Runtime resource request create payload for the Resource Manager API.""" + + subjects: list[RMSubject] = Field(min_length=1) + demands: list[RMRequestDemand] + pool: UUID | str | None = None + pool_selector: Optional[dict[str, Any]] = None + preemption_group: Optional[RMSubjectSelectorNode] = None + reclaim_tolerance: str = "none" + lease_expires_at: Optional[datetime] = None + allocation_wait_timeout_seconds: Optional[int] = None + user_id: Optional[UUID] = None + metadata: dict[str, Any] = Field(default_factory=dict) + + @classmethod + def from_model( + cls, + resource_request: ResourceRequestRequest, + *, + subjects: list[RMSubject], + preemption_group: Optional[RMSubjectSelectorNode] = None, + user_id: Optional[UUID] = None, + ) -> "RMResourceRequestCreate": + """Build a runtime request create payload from a ZenML request. + + Args: + resource_request: ZenML resource request model. + subjects: Resource Manager subject chains attached to the request. + preemption_group: Optional Resource Manager preemption group + selector. + user_id: Optional ZenML Pro user ID. + + Returns: + Resource Manager request create payload. + + Raises: + ValueError: If component IDs are set without a step run ID. + """ + if ( + resource_request.component_ids + and resource_request.step_run_id is None + ): + raise ValueError( + "step_run_id is required when component_ids are set." + ) + metadata = dict(resource_request.metadata) + if resource_request.component_ids and resource_request.step_run_id: + metadata[STEP_RUN_ID_METADATA_KEY] = str( + resource_request.step_run_id + ) + return cls( + subjects=subjects, + demands=[ + RMRequestDemand.from_model(demand) + for demand in resource_request.demands + ], + pool=( + resource_request.pool_id + if resource_request.pool_id is not None + else resource_request.pool + ), + pool_selector=resource_request.pool_selector, + preemption_group=( + preemption_group or resource_request.preemption_group + ), + reclaim_tolerance=( + resource_request.reclaim_tolerance + or ResourceRequestReclaimTolerance.NONE + ).value, + lease_expires_at=resource_request.lease_expires_at, + allocation_wait_timeout_seconds=( + resource_request.allocation_wait_timeout_seconds + ), + user_id=user_id, + metadata=metadata, + ) + + +class RMResourceRequestRenewalRequest(BaseModel): + """Runtime resource request renewal payload for the Resource Manager API.""" + + lease_expires_at: datetime + runtime_state: Optional[ResourceRequestRuntimeState] = None + + +class RMResourceRequestResponse(BaseModel): + """Runtime resource request response from the Resource Manager API.""" + + id: UUID + organization_id: UUID + subjects: list[RMSubject] = Field(default_factory=list) + demands: list[RMRequestDemand] = Field(default_factory=list) + allocations: list["RMAllocationResponse"] = Field(default_factory=list) + queue_entries: list["RMQueueEntryResponse"] = Field(default_factory=list) + target_settings: list["RMTargetSettingsResponse"] = Field( + default_factory=list + ) + pool_id: Optional[UUID] = None + pool_name: Optional[str] = None + pool_selector: Optional[dict[str, Any]] = None + preemption_group: Optional[RMSubjectSelectorNode] = None + status: str + runtime_state: str = ResourceRequestRuntimeState.UNKNOWN.value + reclaim_tolerance: str + lease_expires_at: Optional[datetime] = None + allocation_deadline: Optional[datetime] = None + renewed_at: Optional[datetime] = None + created: Optional[datetime] = None + updated: Optional[datetime] = None + allocated_at: Optional[datetime] = None + released_at: Optional[datetime] = None + queued_at: Optional[datetime] = None + status_reason: Optional[str] = None + preemption_initiated_by_id: Optional[UUID] = None + metadata: dict[str, Any] = Field(default_factory=dict) + + def to_model(self) -> ResourceRequestResponse: + """Convert this runtime request response into a ZenML response. + + Returns: + ZenML resource request response model. + """ + now = utc_now() + created = self.created or now + updated = self.updated or created + step_name = self.metadata.get(STEP_NAME_METADATA_KEY) + pipeline_run_name = self.metadata.get(PIPELINE_RUN_NAME_METADATA_KEY) + component_settings: dict[str, Any] = {} + service_connector_settings: ( + ResourceRequestServiceConnectorSettings | None + ) = None + for settings in self.target_settings: + if settings.target_type == "component": + component_settings.update(settings.settings) + elif ( + settings.target_type == "service_connector" + and service_connector_settings is None + ): + service_connector_settings = ( + settings.to_service_connector_settings() + ) + + return ResourceRequestResponse( + id=self.id, + body=ResourceRequestResponseBody( + created=created, + updated=updated, + user_id=None, + component_ids=[ + subject.leaf.subject_id + for subject in self.subjects + if subject.leaf.subject_type == COMPONENT_SUBJECT_TYPE + ], + step_run_id=_metadata_uuid( + self.metadata, STEP_RUN_ID_METADATA_KEY + ), + pipeline_run_id=_metadata_uuid( + self.metadata, PIPELINE_RUN_ID_METADATA_KEY + ), + pool_id=self.pool_id, + step_name=None if step_name is None else str(step_name), + pipeline_run_name=( + None + if pipeline_run_name is None + else str(pipeline_run_name) + ), + project_id=_metadata_uuid( + self.metadata, PROJECT_ID_METADATA_KEY + ), + pool_name=self.pool_name, + pool_selector=self.pool_selector, + preemption_group=( + self.preemption_group.model_dump( + mode="json", + by_alias=True, + exclude_none=True, + exclude_defaults=True, + ) + if self.preemption_group is not None + else None + ), + demands=[demand.to_model() for demand in self.demands], + status=ResourceRequestStatus(self.status), + runtime_state=ResourceRequestRuntimeState(self.runtime_state), + reclaim_tolerance=ResourceRequestReclaimTolerance( + self.reclaim_tolerance + ), + lease_expires_at=self.lease_expires_at, + allocation_deadline=self.allocation_deadline, + renewed_at=self.renewed_at, + allocated_at=self.allocated_at, + released_at=self.released_at, + queued_at=self.queued_at, + status_reason=self.status_reason, + preemption_initiated_by_id=self.preemption_initiated_by_id, + ), + metadata=ResourceRequestResponseMetadata(), + resources=ResourceRequestResponseResources( + component_settings=component_settings, + service_connector_settings=service_connector_settings, + allocations=[ + allocation.to_model(request_subjects=self.subjects) + for allocation in self.allocations + ], + queue_entries=[ + entry.to_model() for entry in self.queue_entries + ], + ), + ) + + +class RMResourceRequestListResponse(BaseModel): + """Runtime resource request list response from the Resource Manager API.""" + + items: list[RMResourceRequestResponse] + total: int + + +class RMQueueEntryResponse(BaseModel): + """Pool queue entry response included in runtime request resources.""" + + id: UUID + organization_id: UUID + request_id: UUID + pool_id: UUID + pool_name: str + policy_id: UUID + priority: int + enqueued_at: datetime + created: Optional[datetime] = None + + def to_model(self) -> ResourcePoolQueueItem: + """Convert this queue entry into a ZenML queue item. + + Returns: + ZenML resource pool queue item. + """ + return ResourcePoolQueueItem( + id=self.id, + request_id=self.request_id, + pool_id=self.pool_id, + pool_name=self.pool_name, + policy_id=self.policy_id, + priority=self.priority, + priority_lane=self.priority >= PRIORITY_LANE_PRIORITY, + enqueued_at=self.enqueued_at, + created=self.created, + ) + + +class RMTargetSettingsResponse(BaseModel): + """Pool target settings selected by Resource Manager for a request.""" + + target_type: str + settings: dict[str, Any] = Field(default_factory=dict) + + def to_service_connector_settings( + self, + ) -> ResourceRequestServiceConnectorSettings: + """Convert service connector target settings into a ZenML model. + + Returns: + ZenML service connector settings model. + """ + connector_id = self.settings.get("connector_id") + return ResourceRequestServiceConnectorSettings( + connector_id=UUID(str(connector_id)) if connector_id else None, + resource_type=self.settings.get("resource_type"), + resource_id=self.settings.get("resource_id"), + ) + + +class RMAllocationResponse(BaseModel): + """Allocation response included in runtime request resources.""" + + id: UUID + organization_id: UUID + request_id: UUID + demand_index: Optional[int] = Field(default=None, ge=0) + pool_id: UUID + pool_name: str + capacity_entry_id: Optional[UUID] = None + capacity_entry_name: Optional[str] = None + resource_id: UUID + resource_name: str + resource_kind: str + class_name: str = Field(alias="class", serialization_alias="class") + quantity: int + unit: Optional[str] = None + base_quantity: Optional[int] = None + admitted_by_policy_id: UUID + resolved_grant_id: UUID | None = None + allocation_priority: int + matched_subject_ids: tuple[UUID, ...] + preemption_state: str + preemption_reason: Optional[str] = None + released_at: Optional[datetime] = None + created: Optional[datetime] = None + updated: Optional[datetime] = None + + model_config = ConfigDict(populate_by_name=True) + + def _resolve_subject_ids( + self, + *, + request_subjects: list[RMSubject] | None = None, + ) -> tuple[Optional[UUID], Optional[UUID]]: + """Resolve matched subjects into component and account IDs. + + Args: + request_subjects: Request subject chains used to resolve matched + Resource Manager subject IDs. + + Returns: + Component and account IDs resolved from matched subjects. + """ + for matched_subject_id in self.matched_subject_ids: + for subject in request_subjects or []: + matched_subject = subject.find(matched_subject_id) + if matched_subject is None: + continue + if matched_subject.subject_type == ACCOUNT_SUBJECT_TYPE: + return None, matched_subject_id + if matched_subject.subject_type == COMPONENT_SUBJECT_TYPE: + return matched_subject_id, None + if self.matched_subject_ids: + return self.matched_subject_ids[0], None + return None, None + + def to_model( + self, + *, + request_subjects: list[RMSubject] | None = None, + ) -> ResourcePoolAllocation: + """Convert this allocation into a ZenML allocation. + + Args: + request_subjects: Request subject chains used to resolve component + and account IDs. + + Returns: + ZenML resource pool allocation model. + """ + component_id, account_id = self._resolve_subject_ids( + request_subjects=request_subjects + ) + return ResourcePoolAllocation( + id=self.id, + request_id=self.request_id, + demand_index=self.demand_index, + pool_id=self.pool_id, + pool_name=self.pool_name, + capacity_entry_id=self.capacity_entry_id, + capacity_entry_name=self.capacity_entry_name, + resource_id=self.resource_id, + resource=self.resource_name, + resource_kind=self.resource_kind, + class_name=self.class_name, + quantity=self.quantity, + unit=self.unit, + base_quantity=self.base_quantity, + policy_id=self.admitted_by_policy_id, + grant_id=self.resolved_grant_id, + priority=self.allocation_priority, + priority_lane=self.allocation_priority >= PRIORITY_LANE_PRIORITY, + component_id=component_id, + account_id=account_id, + preemption_state=self.preemption_state, + preemption_reason=self.preemption_reason, + released_at=self.released_at, + created=self.created, + updated=self.updated, + ) diff --git a/src/zenml/zen_stores/resource_pools/store_interface.py b/src/zenml/zen_stores/resource_pools/store_interface.py index 78acea50e39..c698df46545 100644 --- a/src/zenml/zen_stores/resource_pools/store_interface.py +++ b/src/zenml/zen_stores/resource_pools/store_interface.py @@ -11,7 +11,7 @@ # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express # or implied. See the License for the specific language governing # permissions and limitations under the License. -"""Resource pools store interface.""" +"""Resource request store interface.""" from abc import ABC, abstractmethod from typing import TYPE_CHECKING @@ -19,214 +19,109 @@ from zenml.models import ( Page, - ResourcePoolFilter, - ResourcePoolRequest, - ResourcePoolResponse, - ResourcePoolSubjectPolicyFilter, - ResourcePoolSubjectPolicyRequest, - ResourcePoolSubjectPolicyResponse, - ResourcePoolSubjectPolicyUpdate, - ResourcePoolUpdate, ResourceRequestFilter, ResourceRequestRequest, ResourceRequestResponse, ) +from zenml.models.v2.core.resource_request import ( + ResourceRequestRenewalRequest, +) if TYPE_CHECKING: from zenml.zen_stores.sql_zen_store import Session, SqlZenStore class ResourcePoolsStoreInterface(ABC): - """Resource pools store interface.""" - - # -------------------- Resource Pools ------------- + """Resource request store interface.""" @abstractmethod - def create_resource_pool( - self, resource_pool: ResourcePoolRequest - ) -> ResourcePoolResponse: - """Create a resource pool. - - Args: - resource_pool: The resource pool to create. - - Returns: - The created resource pool. - """ - - @abstractmethod - def get_resource_pool( - self, resource_pool_id: UUID, hydrate: bool = True - ) -> ResourcePoolResponse: - """Get a resource pool by ID. - - Args: - resource_pool_id: The ID of the resource pool to get. - hydrate: Flag deciding whether to hydrate the output model(s) - by including metadata fields in the response. - - Returns: - The resource pool. - """ - - @abstractmethod - def list_resource_pools( - self, filter_model: ResourcePoolFilter, hydrate: bool = False - ) -> Page[ResourcePoolResponse]: - """List all resource pools matching the given filter criteria. - - Args: - filter_model: All filter parameters including pagination - params. - hydrate: Flag deciding whether to hydrate the output model(s) - by including metadata fields in the response. - - Returns: - A list of all resource pools matching the filter criteria. - """ - - @abstractmethod - def update_resource_pool( - self, resource_pool_id: UUID, update: ResourcePoolUpdate - ) -> ResourcePoolResponse: - """Update an existing resource pool. + def get_resource_request( + self, resource_request_id: UUID, hydrate: bool = True + ) -> ResourceRequestResponse: + """Get a resource request by ID. Args: - resource_pool_id: The ID of the resource pool to update. - update: The update to be applied to the resource pool. + resource_request_id: The ID of the resource request to get. + hydrate: Whether to hydrate metadata and resources. Returns: - The updated resource pool. - """ - - @abstractmethod - def delete_resource_pool(self, resource_pool_id: UUID) -> None: - """Delete a resource pool. - - Args: - resource_pool_id: The ID of the resource pool to delete. + The resource request. """ @abstractmethod - def create_resource_pool_subject_policy( - self, policy: ResourcePoolSubjectPolicyRequest - ) -> ResourcePoolSubjectPolicyResponse: - """Create a resource pool subject policy. + def list_resource_requests( + self, filter_model: ResourceRequestFilter, hydrate: bool = False + ) -> Page[ResourceRequestResponse]: + """List resource requests matching the given filter criteria. Args: - policy: The policy to create. + filter_model: Filter and pagination parameters. + hydrate: Whether to hydrate metadata and resources. Returns: - The created policy. + Matching resource requests. """ @abstractmethod - def get_resource_pool_subject_policy( - self, policy_id: UUID, hydrate: bool = True - ) -> ResourcePoolSubjectPolicyResponse: - """Get a resource pool subject policy by ID. + def release_resource_request( + self, + resource_request_id: UUID, + ) -> ResourceRequestResponse: + """Release a resource request on behalf of its owner. Args: - policy_id: The ID of the policy to get. - hydrate: Whether to include metadata fields. + resource_request_id: The ID of the resource request to release. Returns: - The requested policy. + The released resource request. """ @abstractmethod - def list_resource_pool_subject_policies( + def renew_resource_request( self, - filter_model: ResourcePoolSubjectPolicyFilter, - hydrate: bool = False, - ) -> Page[ResourcePoolSubjectPolicyResponse]: - """List resource pool subject policies. + resource_request_id: UUID, + renewal_request: ResourceRequestRenewalRequest, + ) -> ResourceRequestResponse: + """Renew a resource request lease. Args: - filter_model: All filter parameters including pagination params. - hydrate: Whether to include metadata fields. + resource_request_id: The ID of the resource request to renew. + renewal_request: The renewed lease expiration timestamp. Returns: - Matching policies. + The renewed resource request. """ - @abstractmethod - def update_resource_pool_subject_policy( - self, policy_id: UUID, update: ResourcePoolSubjectPolicyUpdate - ) -> ResourcePoolSubjectPolicyResponse: - """Update an existing resource pool subject policy. - - Args: - policy_id: The ID of the policy to update. - update: The update model. - Returns: - The updated policy. - """ +class ResourcePoolsSQLStoreInterface(ResourcePoolsStoreInterface): + """Resource request SQL store interface.""" - @abstractmethod - def delete_resource_pool_subject_policy(self, policy_id: UUID) -> None: - """Delete a resource pool subject policy. + def __init__(self, store: "SqlZenStore") -> None: + """Initialize the resource request SQL store. Args: - policy_id: The ID of the policy to delete. + store: The store to use. """ - - # -------------------- Resource Requests ------------- + super().__init__() + self.store = store + store.resource_pools = self @abstractmethod - def get_resource_request( - self, resource_request_id: UUID, hydrate: bool = True + def create_resource_request( + self, + session: "Session", + resource_request: ResourceRequestRequest, ) -> ResourceRequestResponse: - """Get a resource request by ID. + """Create a step-run resource request with SQL session-backed metadata. Args: - resource_request_id: The ID of the resource request to get. - hydrate: Flag deciding whether to hydrate the output model(s) - by including metadata fields in the response. - - Returns: - The resource request. - """ - - @abstractmethod - def list_resource_requests( - self, filter_model: ResourceRequestFilter, hydrate: bool = False - ) -> Page[ResourceRequestResponse]: - """List all resource requests matching the given filter criteria. - - Args: - filter_model: All filter parameters including pagination - params. - hydrate: Flag deciding whether to hydrate the output model(s) - by including metadata fields in the response. + session: DB session used to enrich step-run metadata. + resource_request: The resource request to create. Returns: - A list of all resource requests matching the filter criteria. + The created resource request. """ - @abstractmethod - def delete_resource_request(self, resource_request_id: UUID) -> None: - """Delete a resource request. - - Args: - resource_request_id: The ID of the resource request to delete. - """ - - -class ResourcePoolsSQLStoreInterface(ResourcePoolsStoreInterface): - """Resource pools SQL store interface.""" - - def __init__(self, store: "SqlZenStore") -> None: - """Initialize the resource pools SQL store. - - Args: - store: The store to use. - """ - super().__init__() - self.store = store - store.resource_pools = self - @abstractmethod def release_step_run_resources( self, session: "Session", step_run_id: UUID @@ -237,17 +132,3 @@ def release_step_run_resources( session: DB session. step_run_id: The ID of the step run to release resources for. """ - - @abstractmethod - def create_resource_request( - self, session: "Session", resource_request: ResourceRequestRequest - ) -> ResourceRequestResponse | None: - """Create a resource request. - - Args: - session: DB session. - resource_request: The resource request to create. - - Returns: - The created resource request or None if the feature is not enabled. - """ diff --git a/src/zenml/zen_stores/rest_zen_store.py b/src/zenml/zen_stores/rest_zen_store.py index 80365b8d8b4..4e5b4fed2a0 100644 --- a/src/zenml/zen_stores/rest_zen_store.py +++ b/src/zenml/zen_stores/rest_zen_store.py @@ -93,8 +93,6 @@ PROJECTS, REPLAY, RESOLVE, - RESOURCE_POOL_SUBJECT_POLICIES, - RESOURCE_POOLS, RESOURCE_REQUESTS, RUN_METADATA, RUN_TEMPLATES, @@ -234,15 +232,8 @@ ProjectRequest, ProjectResponse, ProjectUpdate, - ResourcePoolFilter, - ResourcePoolRequest, - ResourcePoolResponse, - ResourcePoolSubjectPolicyFilter, - ResourcePoolSubjectPolicyRequest, - ResourcePoolSubjectPolicyResponse, - ResourcePoolSubjectPolicyUpdate, - ResourcePoolUpdate, ResourceRequestFilter, + ResourceRequestRenewalRequest, ResourceRequestResponse, RunMetadataRequest, RunStatisticsRequest, @@ -287,6 +278,7 @@ StackRequest, StackResponse, StackUpdate, + StepHeartbeatRequest, StepHeartbeatResponse, StepRunFilter, StepRunRequest, @@ -330,6 +322,7 @@ AnyRequest = TypeVar("AnyRequest", bound=BaseRequest) AnyResponse = TypeVar("AnyResponse", bound=BaseIdentifiedResponse) # type: ignore[type-arg] +AnyModel = TypeVar("AnyModel", bound=BaseModel) class RestZenStoreConfiguration(StoreConfiguration): @@ -1328,183 +1321,6 @@ def delete_stack_component(self, component_id: UUID) -> None: route=STACK_COMPONENTS, ) - # -------------------- Resource Pools ------------- - - def create_resource_pool( - self, resource_pool: ResourcePoolRequest - ) -> ResourcePoolResponse: - """Create a resource pool. - - Args: - resource_pool: The resource pool to create. - - Returns: - The created resource pool. - """ - return self._create_resource( - resource=resource_pool, - route=RESOURCE_POOLS, - response_model=ResourcePoolResponse, - ) - - def get_resource_pool( - self, resource_pool_id: UUID, hydrate: bool = True - ) -> ResourcePoolResponse: - """Get a resource pool by ID. - - Args: - resource_pool_id: The ID of the resource pool to get. - hydrate: Flag deciding whether to hydrate the output model(s) - by including metadata fields in the response. - - Returns: - The resource pool. - """ - return self._get_resource( - resource_id=resource_pool_id, - route=RESOURCE_POOLS, - response_model=ResourcePoolResponse, - params={"hydrate": hydrate}, - ) - - def list_resource_pools( - self, filter_model: ResourcePoolFilter, hydrate: bool = False - ) -> Page[ResourcePoolResponse]: - """List all resource pools matching the given filter criteria. - - Args: - filter_model: All filter parameters including pagination - params. - hydrate: Flag deciding whether to hydrate the output model(s) - by including metadata fields in the response. - - Returns: - A list of all resource pools matching the filter criteria. - """ - return self._list_paginated_resources( - route=RESOURCE_POOLS, - response_model=ResourcePoolResponse, - filter_model=filter_model, - params={"hydrate": hydrate}, - ) - - def update_resource_pool( - self, resource_pool_id: UUID, update: ResourcePoolUpdate - ) -> ResourcePoolResponse: - """Update an existing resource pool. - - Args: - resource_pool_id: The ID of the resource pool to update. - update: The update to be applied to the resource pool. - - Returns: - The updated resource pool. - """ - return self._update_resource( - resource_id=resource_pool_id, - resource_update=update, - route=RESOURCE_POOLS, - response_model=ResourcePoolResponse, - ) - - def delete_resource_pool(self, resource_pool_id: UUID) -> None: - """Delete a resource pool. - - Args: - resource_pool_id: The ID of the resource pool to delete. - """ - self._delete_resource( - resource_id=resource_pool_id, - route=RESOURCE_POOLS, - ) - - def create_resource_pool_subject_policy( - self, policy: ResourcePoolSubjectPolicyRequest - ) -> ResourcePoolSubjectPolicyResponse: - """Create a resource pool subject policy. - - Args: - policy: The policy to create. - - Returns: - The created policy. - """ - return self._create_resource( - resource=policy, - route=RESOURCE_POOL_SUBJECT_POLICIES, - response_model=ResourcePoolSubjectPolicyResponse, - ) - - def get_resource_pool_subject_policy( - self, policy_id: UUID, hydrate: bool = True - ) -> ResourcePoolSubjectPolicyResponse: - """Get a resource pool subject policy by ID. - - Args: - policy_id: The ID of the policy to get. - hydrate: Flag deciding whether to hydrate the output model(s). - - Returns: - The policy. - """ - return self._get_resource( - resource_id=policy_id, - route=RESOURCE_POOL_SUBJECT_POLICIES, - response_model=ResourcePoolSubjectPolicyResponse, - params={"hydrate": hydrate}, - ) - - def list_resource_pool_subject_policies( - self, - filter_model: ResourcePoolSubjectPolicyFilter, - hydrate: bool = False, - ) -> Page[ResourcePoolSubjectPolicyResponse]: - """List resource pool subject policies. - - Args: - filter_model: All filter parameters including pagination params. - hydrate: Flag deciding whether to hydrate the output model(s). - - Returns: - Matching policies. - """ - return self._list_paginated_resources( - route=RESOURCE_POOL_SUBJECT_POLICIES, - response_model=ResourcePoolSubjectPolicyResponse, - filter_model=filter_model, - params={"hydrate": hydrate}, - ) - - def update_resource_pool_subject_policy( - self, policy_id: UUID, update: ResourcePoolSubjectPolicyUpdate - ) -> ResourcePoolSubjectPolicyResponse: - """Update an existing resource pool subject policy. - - Args: - policy_id: The ID of the policy to update. - update: The update to be applied to the policy. - - Returns: - The updated policy. - """ - return self._update_resource( - resource_id=policy_id, - resource_update=update, - route=RESOURCE_POOL_SUBJECT_POLICIES, - response_model=ResourcePoolSubjectPolicyResponse, - ) - - def delete_resource_pool_subject_policy(self, policy_id: UUID) -> None: - """Delete a resource pool subject policy. - - Args: - policy_id: The ID of the policy to delete. - """ - self._delete_resource( - resource_id=policy_id, - route=RESOURCE_POOL_SUBJECT_POLICIES, - ) - # -------------------- Resource Requests ------------- def get_resource_request( @@ -1548,16 +1364,46 @@ def list_resource_requests( params={"hydrate": hydrate}, ) - def delete_resource_request(self, resource_request_id: UUID) -> None: - """Delete a resource request. + def release_resource_request( + self, + resource_request_id: UUID, + ) -> ResourceRequestResponse: + """Release a resource request on behalf of its owner. Args: - resource_request_id: The ID of the resource request to delete. + resource_request_id: The ID of the resource request to release. + + Returns: + The released resource request. """ - self._delete_resource( - resource_id=resource_request_id, - route=RESOURCE_REQUESTS, + response_body = self._request( + "POST", + self.url + + API + + VERSION_1 + + f"{RESOURCE_REQUESTS}/{resource_request_id}/release", + ) + return ResourceRequestResponse.model_validate(response_body) + + def renew_resource_request( + self, + resource_request_id: UUID, + renewal_request: ResourceRequestRenewalRequest, + ) -> ResourceRequestResponse: + """Renew a resource request lease. + + Args: + resource_request_id: The ID of the resource request to renew. + renewal_request: The renewed lease expiration timestamp. + + Returns: + The renewed resource request. + """ + response_body = self.post( + path=f"{RESOURCE_REQUESTS}/{resource_request_id}/renew", + body=renewal_request, ) + return ResourceRequestResponse.model_validate(response_body) # ----------------------------- Flavors ----------------------------- @@ -3939,18 +3785,32 @@ def update_run_step( ) def update_step_heartbeat( - self, step_run_id: UUID + self, + step_run_id: UUID, + heartbeat_liveness_timeout_seconds: Optional[int] = None, ) -> StepHeartbeatResponse: """Updates a step run heartbeat. Args: step_run_id: The ID of the step to update. + heartbeat_liveness_timeout_seconds: Optional number of seconds the + server should wait for another heartbeat before considering the + heartbeat client dead. Returns: The step heartbeat response. """ + request = None + if heartbeat_liveness_timeout_seconds is not None: + request = StepHeartbeatRequest( + heartbeat_liveness_timeout_seconds=( + heartbeat_liveness_timeout_seconds + ) + ) + response_body = self.put( path=f"{STEPS}/{str(step_run_id)}{HEARTBEAT}", + body=request, timeout=5, ) @@ -5555,6 +5415,42 @@ def _list_paginated_resources( ] return page_of_items + def _list_paginated_models( + self, + route: str, + response_model: Type[AnyModel], + filter_model: BaseFilter, + params: Optional[Dict[str, Any]] = None, + ) -> Page[AnyModel]: + """Retrieve a page of Pydantic models from a REST endpoint. + + Args: + route: The resource REST API route to use. + response_model: Model to use to serialize the response body. + filter_model: The filter model to use for the list query. + params: Optional query parameters to pass to the endpoint. + + Returns: + Page of retrieved models matching the filter criteria. + + Raises: + ValueError: If the value returned by the server is not a list. + """ + params = params or {} + params.update(filter_model.model_dump(exclude_none=True)) + body = self.get(route, params=params) + if not isinstance(body, dict): + raise ValueError( + f"Bad API Response. Expected list, got {type(body)}" + ) + + page_of_items: Page[AnyModel] = Page.model_validate(body) + page_of_items.items = [ + response_model.model_validate(generic_item) + for generic_item in body["items"] + ] + return page_of_items + def _list_resources( self, route: str, diff --git a/src/zenml/zen_stores/schemas/__init__.py b/src/zenml/zen_stores/schemas/__init__.py index c6a8f52e309..a5045fadf11 100644 --- a/src/zenml/zen_stores/schemas/__init__.py +++ b/src/zenml/zen_stores/schemas/__init__.py @@ -59,20 +59,6 @@ StepConfigurationSchema, ) from zenml.zen_stores.schemas.project_schemas import ProjectSchema -from zenml.zen_stores.schemas.resource_pool_policy_schemas import ( - ResourcePoolSubjectPolicyResourceSchema, - ResourcePoolSubjectPolicySchema, -) -from zenml.zen_stores.schemas.resource_pool_schemas import ( - ResourcePoolAllocationSchema, - ResourcePoolQueueSchema, - ResourcePoolResourceSchema, - ResourcePoolSchema, -) -from zenml.zen_stores.schemas.resource_request_schemas import ( - ResourceRequestResourceSchema, - ResourceRequestSchema, -) from zenml.zen_stores.schemas.run_metadata_schemas import ( RunMetadataResourceSchema, RunMetadataSchema, @@ -161,14 +147,6 @@ "ProjectSchema", "ApiTransactionResultSchema", "ApiTransactionSchema", - "ResourcePoolQueueSchema", - "ResourcePoolAllocationSchema", - "ResourcePoolSubjectPolicySchema", - "ResourcePoolSubjectPolicyResourceSchema", - "ResourcePoolSchema", - "ResourcePoolResourceSchema", - "ResourceRequestSchema", - "ResourceRequestResourceSchema", "TriggerSchema", "TriggerSnapshotSchema", "TriggerExecutionSchema", diff --git a/src/zenml/zen_stores/schemas/resource_pool_policy_schemas.py b/src/zenml/zen_stores/schemas/resource_pool_policy_schemas.py deleted file mode 100644 index 716768d856e..00000000000 --- a/src/zenml/zen_stores/schemas/resource_pool_policy_schemas.py +++ /dev/null @@ -1,259 +0,0 @@ -# Copyright (c) ZenML GmbH 2026. All Rights Reserved. -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at: -# -# https://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express -# or implied. See the License for the specific language governing -# permissions and limitations under the License. -"""Resource pool subject policy schemas.""" - -from typing import TYPE_CHECKING, Any, List, Optional, Sequence -from uuid import UUID - -from sqlalchemy import UniqueConstraint -from sqlalchemy.orm import selectinload -from sqlalchemy.sql.base import ExecutableOption -from sqlmodel import Field, Relationship - -from zenml.models import ( - ResourcePoolSubjectPolicyRequest, - ResourcePoolSubjectPolicyResponse, - ResourcePoolSubjectPolicyResponseBody, - ResourcePoolSubjectPolicyResponseMetadata, - ResourcePoolSubjectPolicyResponseResources, - ResourcePoolSubjectPolicyUpdate, -) -from zenml.utils.time_utils import utc_now -from zenml.zen_stores.schemas.base_schemas import BaseSchema -from zenml.zen_stores.schemas.component_schemas import StackComponentSchema -from zenml.zen_stores.schemas.schema_utils import ( - build_foreign_key_field, - build_index, -) -from zenml.zen_stores.schemas.user_schemas import UserSchema -from zenml.zen_stores.schemas.utils import jl_arg - -if TYPE_CHECKING: - from zenml.zen_stores.schemas.resource_pool_schemas import ( - ResourcePoolSchema, - ) - - -class ResourcePoolSubjectPolicySchema(BaseSchema, table=True): - """Resource pool subject policy schema.""" - - __tablename__ = "resource_pool_subject_policy" - __table_args__ = ( - UniqueConstraint( - "component_id", - "pool_id", - name="unique_component_id_pool_id", - ), - build_index( - table_name=__tablename__, - column_names=["pool_id", "priority", "component_id"], - ), - build_index( - table_name=__tablename__, - column_names=["component_id", "priority", "pool_id"], - ), - build_index( - table_name=__tablename__, - column_names=["pool_id", "component_id"], - ), - ) - - user_id: Optional[UUID] = build_foreign_key_field( - source=__tablename__, - target=UserSchema.__tablename__, - source_column="user_id", - target_column="id", - ondelete="SET NULL", - nullable=True, - ) - user: Optional["UserSchema"] = Relationship() - - component_id: UUID = build_foreign_key_field( - source=__tablename__, - target=StackComponentSchema.__tablename__, - source_column="component_id", - target_column="id", - ondelete="CASCADE", - nullable=False, - ) - component: "StackComponentSchema" = Relationship() - pool_id: UUID = build_foreign_key_field( - source=__tablename__, - target="resource_pool", - source_column="pool_id", - target_column="id", - ondelete="CASCADE", - nullable=False, - ) - pool: "ResourcePoolSchema" = Relationship() - - priority: int - resources: List["ResourcePoolSubjectPolicyResourceSchema"] = Relationship( - back_populates="policy", - sa_relationship_kwargs={ - "passive_deletes": True, - "cascade": "all, delete-orphan", - }, - ) - - @classmethod - def get_query_options( - cls, - include_metadata: bool = False, - include_resources: bool = False, - **kwargs: Any, - ) -> Sequence[ExecutableOption]: - """Gets query options for this schema. - - Args: - include_metadata: If metadata should be included. - include_resources: If resources should be included. - **kwargs: Additional keyword arguments. - - Returns: - The query options. - """ - options: List[ExecutableOption] = [ - selectinload(jl_arg(ResourcePoolSubjectPolicySchema.resources)), - ] - - if include_resources: - options.extend( - [ - selectinload( - jl_arg(ResourcePoolSubjectPolicySchema.component) - ), - selectinload(jl_arg(ResourcePoolSubjectPolicySchema.pool)), - selectinload(jl_arg(ResourcePoolSubjectPolicySchema.user)), - ] - ) - - return options - - @classmethod - def from_request( - cls, request: ResourcePoolSubjectPolicyRequest - ) -> "ResourcePoolSubjectPolicySchema": - """Creates a schema instance from a request model. - - Args: - request: The request model. - - Returns: - The schema instance. - """ - return cls( - user_id=request.user, - component_id=request.component_id, - pool_id=request.pool_id, - priority=request.priority, - ) - - def update( - self, update: ResourcePoolSubjectPolicyUpdate - ) -> "ResourcePoolSubjectPolicySchema": - """Updates this schema from an update model. - - Args: - update: The update model. - - Returns: - The updated schema. - """ - if update.priority is not None: - self.priority = update.priority - - self.updated = utc_now() - return self - - def to_model( - self, - include_metadata: bool = False, - include_resources: bool = False, - **kwargs: Any, - ) -> "ResourcePoolSubjectPolicyResponse": - """Converts this schema to a response model. - - Args: - include_metadata: Whether to include metadata. - include_resources: Whether to include nested resources. - **kwargs: Additional keyword arguments. - - Returns: - The response model. - """ - body = ResourcePoolSubjectPolicyResponseBody( - created=self.created, - updated=self.updated, - user_id=self.user_id, - priority=self.priority, - reserved={ - resource.key: resource.reserved for resource in self.resources - }, - limit={ - resource.key: resource.limit - for resource in self.resources - if resource.limit is not None - }, - ) - - metadata = None - if include_metadata: - metadata = ResourcePoolSubjectPolicyResponseMetadata() - - resources = None - if include_resources: - resources = ResourcePoolSubjectPolicyResponseResources( - user=self.user.to_model() if self.user else None, - component=self.component.to_model(), - pool=self.pool.to_model( - include_metadata=False, include_resources=False - ), - ) - - return ResourcePoolSubjectPolicyResponse( - id=self.id, - body=body, - metadata=metadata, - resources=resources, - ) - - -class ResourcePoolSubjectPolicyResourceSchema(BaseSchema, table=True): - """Resource pool subject policy resource schema.""" - - __tablename__ = "resource_pool_subject_policy_resource" - __table_args__ = ( - UniqueConstraint( - "policy_id", - "key", - name="unique_resource_pool_subject_policy_resource", - ), - ) - policy_id: UUID = build_foreign_key_field( - source=__tablename__, - target=ResourcePoolSubjectPolicySchema.__tablename__, - source_column="policy_id", - target_column="id", - ondelete="CASCADE", - nullable=False, - custom_constraint_name="fk_pool_subject_policy_resource_policy_id", - ) - key: str - reserved: int = Field(default=0, nullable=False) - limit: Optional[int] = Field(default=None, nullable=True) - - policy: Optional["ResourcePoolSubjectPolicySchema"] = Relationship( - back_populates="resources" - ) diff --git a/src/zenml/zen_stores/schemas/resource_pool_schemas.py b/src/zenml/zen_stores/schemas/resource_pool_schemas.py deleted file mode 100644 index b1858240b52..00000000000 --- a/src/zenml/zen_stores/schemas/resource_pool_schemas.py +++ /dev/null @@ -1,362 +0,0 @@ -# Copyright (c) ZenML GmbH 2022. All Rights Reserved. -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at: -# -# https://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express -# or implied. See the License for the specific language governing -# permissions and limitations under the License. -"""Resource Pool schemas.""" - -from datetime import datetime -from typing import Any, List, Optional, Sequence -from uuid import UUID - -from sqlalchemy import UniqueConstraint, desc -from sqlalchemy.orm import load_only, selectinload -from sqlalchemy.sql.base import ExecutableOption -from sqlmodel import Field, Relationship, col - -from zenml.models import ( - ResourcePoolRequest, - ResourcePoolResponse, - ResourcePoolResponseBody, - ResourcePoolResponseMetadata, - ResourcePoolResponseResources, - ResourcePoolUpdate, -) -from zenml.utils.time_utils import utc_now -from zenml.zen_stores.schemas.base_schemas import BaseSchema, NamedSchema -from zenml.zen_stores.schemas.resource_pool_policy_schemas import ( - ResourcePoolSubjectPolicySchema, -) -from zenml.zen_stores.schemas.resource_request_schemas import ( - ResourceRequestSchema, -) -from zenml.zen_stores.schemas.schema_utils import ( - build_foreign_key_field, - build_index, -) -from zenml.zen_stores.schemas.user_schemas import UserSchema -from zenml.zen_stores.schemas.utils import jl_arg - - -class ResourcePoolQueueSchema(BaseSchema, table=True): - """Resource pool queue schema.""" - - __tablename__ = "resource_pool_queue" - __table_args__ = ( - UniqueConstraint( - "pool_id", - "request_id", - name="unique_resource_pool_queue_pool_id_request_id", - ), - build_index( - table_name=__tablename__, - column_names=["request_id"], - ), - build_index( - table_name=__tablename__, - column_names=[ - "pool_id", - "priority", - "request_created", - "request_id", - ], - ), - build_index( - table_name=__tablename__, - column_names=[ - "pool_id", - "claim_expires_at", - "priority", - "request_created", - "request_id", - ], - ), - ) - - pool_id: UUID = build_foreign_key_field( - source=__tablename__, - target="resource_pool", - source_column="pool_id", - target_column="id", - ondelete="CASCADE", - nullable=False, - ) - request_id: UUID = build_foreign_key_field( - source=__tablename__, - target="resource_request", - source_column="request_id", - target_column="id", - ondelete="CASCADE", - nullable=False, - ) - request: "ResourceRequestSchema" = Relationship() - priority: int - request_created: datetime = Field(default_factory=utc_now) - claim_token: Optional[UUID] = Field(default=None, nullable=True) - claim_expires_at: Optional[datetime] = Field(default=None, nullable=True) - - -class ResourcePoolAllocationSchema(BaseSchema, table=True): - """Resource pool allocation schema.""" - - __tablename__ = "resource_pool_allocation" - __table_args__ = ( - UniqueConstraint( - "request_id", - name="unique_resource_pool_allocation_request_id", - ), - build_index( - table_name=__tablename__, - column_names=[ - "pool_id", - "released_at", - "allocated_at", - "request_id", - ], - ), - ) - - pool_id: UUID = build_foreign_key_field( - source=__tablename__, - target="resource_pool", - source_column="pool_id", - target_column="id", - ondelete="CASCADE", - nullable=False, - ) - request_id: UUID = build_foreign_key_field( - source=__tablename__, - target="resource_request", - source_column="request_id", - target_column="id", - ondelete="CASCADE", - nullable=False, - ) - request: "ResourceRequestSchema" = Relationship() - policy: Optional["ResourcePoolSubjectPolicySchema"] = Relationship( - sa_relationship_kwargs={ - "secondary": "resource_request", - "primaryjoin": ( - "ResourcePoolAllocationSchema.request_id == " - "ResourceRequestSchema.id" - ), - "secondaryjoin": ( - "and_(" - "ResourcePoolSubjectPolicySchema.pool_id == " - "ResourcePoolAllocationSchema.pool_id, " - "ResourcePoolSubjectPolicySchema.component_id == " - "ResourceRequestSchema.component_id" - ")" - ), - "uselist": False, - "viewonly": True, - } - ) - allocated_at: datetime = Field(default_factory=utc_now) - released_at: Optional[datetime] = Field(default=None, nullable=True) - - @property - def priority(self) -> Optional[int]: - """Fetch the priority for this allocation. - - Returns: - The matching policy priority, if a policy exists. - """ - return self.policy.priority if self.policy else None - - -class ResourcePoolSchema(NamedSchema, table=True): - """Resource pool schema.""" - - __tablename__ = "resource_pool" - __table_args__ = ( - UniqueConstraint( - "name", - name="unique_resource_pool_name", - ), - ) - - description: Optional[str] = Field(default=None) - resources: List["ResourcePoolResourceSchema"] = Relationship( - back_populates="pool", - sa_relationship_kwargs={ - "passive_deletes": True, - "cascade": "all, delete-orphan", - }, - ) - - user_id: Optional[UUID] = build_foreign_key_field( - source=__tablename__, - target=UserSchema.__tablename__, - source_column="user_id", - target_column="id", - ondelete="SET NULL", - nullable=True, - ) - user: Optional["UserSchema"] = Relationship() - queue_items: List["ResourcePoolQueueSchema"] = Relationship( - sa_relationship_kwargs={ - "primaryjoin": ( - "ResourcePoolSchema.id == ResourcePoolQueueSchema.pool_id" - ), - "viewonly": True, - } - ) - policies: List["ResourcePoolSubjectPolicySchema"] = Relationship( - back_populates="pool", - sa_relationship_kwargs={ - "order_by": desc(col(ResourcePoolSubjectPolicySchema.priority)), - "passive_deletes": True, - "cascade": "all, delete-orphan", - }, - ) - - @classmethod - def get_query_options( - cls, - include_metadata: bool = False, - include_resources: bool = False, - **kwargs: Any, - ) -> Sequence[ExecutableOption]: - """Get the query options for the schema. - - Args: - include_metadata: Whether metadata will be included when converting - the schema to a model. - include_resources: Whether resources will be included when - converting the schema to a model. - **kwargs: Keyword arguments to allow schema specific logic - - Returns: - A list of query options. - """ - options = [ - selectinload(jl_arg(ResourcePoolSchema.resources)), - selectinload(jl_arg(ResourcePoolSchema.queue_items)).options( - load_only(jl_arg(ResourcePoolQueueSchema.request_id)) - ), - ] - - if include_resources: - options.extend( - [ - selectinload(jl_arg(ResourcePoolSchema.user)), - ] - ) - - return options - - @classmethod - def from_request( - cls, - request: "ResourcePoolRequest", - ) -> "ResourcePoolSchema": - """Create a resource pool schema from a request. - - Args: - request: The request from which to create the resource pool. - - Returns: - The resource pool schema. - """ - return cls( - name=request.name, - user_id=request.user, - description=request.description, - ) - - def update( - self, resource_pool_update: "ResourcePoolUpdate" - ) -> "ResourcePoolSchema": - """Updates a `ResourcePoolSchema` from a `ResourcePoolUpdate`. - - Args: - resource_pool_update: The `ResourcePoolUpdate` to update from. - - Returns: - The updated `ResourcePoolSchema`. - """ - if resource_pool_update.description: - self.description = resource_pool_update.description - - self.updated = utc_now() - return self - - def to_model( - self, - include_metadata: bool = False, - include_resources: bool = False, - **kwargs: Any, - ) -> "ResourcePoolResponse": - """Creates a `ResourcePoolResponse` from a `ResourcePoolSchema`. - - Args: - include_metadata: Whether the metadata will be filled. - include_resources: Whether the resources will be filled. - **kwargs: Keyword arguments to allow schema specific logic - - Returns: - A `ResourcePoolResponse` - """ - body = ResourcePoolResponseBody( - created=self.created, - updated=self.updated, - user_id=self.user_id, - queue_length=len(self.queue_items), - capacity={r.key: r.total for r in self.resources}, - occupied_resources={r.key: r.occupied for r in self.resources}, - ) - - metadata = None - if include_metadata: - metadata = ResourcePoolResponseMetadata( - description=self.description, - ) - - resources = None - if include_resources: - resources = ResourcePoolResponseResources( - user=self.user.to_model() if self.user else None, - ) - - return ResourcePoolResponse( - id=self.id, - name=self.name, - body=body, - metadata=metadata, - resources=resources, - ) - - -class ResourcePoolResourceSchema(BaseSchema, table=True): - """Resource pool resource schema.""" - - __tablename__ = "resource_pool_resource" - __table_args__ = ( - UniqueConstraint( - "pool_id", - "key", - name="unique_resource_pool_resource_pool_id_key", - ), - ) - - pool_id: UUID = build_foreign_key_field( - source=__tablename__, - target=ResourcePoolSchema.__tablename__, - source_column="pool_id", - target_column="id", - ondelete="CASCADE", - nullable=False, - ) - pool: "ResourcePoolSchema" = Relationship(back_populates="resources") - key: str - total: int - occupied: int = 0 diff --git a/src/zenml/zen_stores/schemas/resource_request_schemas.py b/src/zenml/zen_stores/schemas/resource_request_schemas.py deleted file mode 100644 index eb6f334f021..00000000000 --- a/src/zenml/zen_stores/schemas/resource_request_schemas.py +++ /dev/null @@ -1,270 +0,0 @@ -# Copyright (c) ZenML GmbH 2026. All Rights Reserved. -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at: -# -# https://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express -# or implied. See the License for the specific language governing -# permissions and limitations under the License. -"""Resource request schemas.""" - -from typing import TYPE_CHECKING, Any, List, Optional, Sequence -from uuid import UUID - -from sqlalchemy import UniqueConstraint -from sqlalchemy.orm import selectinload -from sqlalchemy.sql.base import ExecutableOption -from sqlmodel import Field, Relationship - -from zenml.enums import ResourceRequestStatus -from zenml.models import ( - ResourceRequestRequest, - ResourceRequestResponse, - ResourceRequestResponseBody, - ResourceRequestResponseMetadata, - ResourceRequestResponseResources, -) -from zenml.zen_stores.schemas.base_schemas import BaseSchema -from zenml.zen_stores.schemas.component_schemas import StackComponentSchema -from zenml.zen_stores.schemas.schema_utils import ( - build_foreign_key_field, - build_index, -) -from zenml.zen_stores.schemas.step_run_schemas import StepRunSchema -from zenml.zen_stores.schemas.user_schemas import UserSchema -from zenml.zen_stores.schemas.utils import jl_arg - -if TYPE_CHECKING: - from zenml.zen_stores.schemas.resource_pool_schemas import ( - ResourcePoolSchema, - ) - - -class ResourceRequestSchema(BaseSchema, table=True): - """Resource request schema.""" - - __tablename__ = "resource_request" - __table_args__ = ( - UniqueConstraint( - "step_run_id", - name="unique_resource_request_step_run_id", - ), - build_index( - table_name=__tablename__, - column_names=["step_run_id", "status", "id"], - ), - build_index( - table_name=__tablename__, - column_names=["component_id", "status", "created", "id"], - ), - ) - - component_id: Optional[UUID] = build_foreign_key_field( - source=__tablename__, - target=StackComponentSchema.__tablename__, - source_column="component_id", - target_column="id", - ondelete="SET NULL", - nullable=True, - ) - step_run_id: Optional[UUID] = build_foreign_key_field( - source=__tablename__, - target=StepRunSchema.__tablename__, - source_column="step_run_id", - target_column="id", - ondelete="SET NULL", - nullable=True, - ) - preemption_initiated_by_id: Optional[UUID] = build_foreign_key_field( - source=__tablename__, - target=__tablename__, - source_column="preemption_initiated_by_id", - target_column="id", - ondelete="SET NULL", - nullable=True, - ) - user_id: Optional[UUID] = build_foreign_key_field( - source=__tablename__, - target=UserSchema.__tablename__, - source_column="user_id", - target_column="id", - ondelete="SET NULL", - nullable=True, - ) - user: Optional["UserSchema"] = Relationship() - component: Optional["StackComponentSchema"] = Relationship() - step_run: Optional["StepRunSchema"] = Relationship( - back_populates="resource_request" - ) - preemption_initiated_by: Optional["ResourceRequestSchema"] = Relationship( - sa_relationship_kwargs={ - "foreign_keys": "[ResourceRequestSchema.preemption_initiated_by_id]", - "remote_side": "ResourceRequestSchema.id", - } - ) - - requested_resources: List["ResourceRequestResourceSchema"] = Relationship( - back_populates="request", - sa_relationship_kwargs={ - "passive_deletes": True, - "cascade": "all, delete-orphan", - }, - ) - pool: Optional["ResourcePoolSchema"] = Relationship( - sa_relationship_kwargs={ - "secondary": "resource_pool_allocation", - "primaryjoin": ( - "ResourceRequestSchema.id == ResourcePoolAllocationSchema.request_id" - ), - "secondaryjoin": ( - "ResourcePoolSchema.id == ResourcePoolAllocationSchema.pool_id" - ), - "uselist": False, - "viewonly": True, - } - ) - status: str - status_reason: Optional[str] = Field(default=None, nullable=True) - preemptible: bool - - @classmethod - def get_query_options( - cls, - include_metadata: bool = False, - include_resources: bool = False, - **kwargs: Any, - ) -> Sequence[ExecutableOption]: - """Get the query options for the schema. - - Args: - include_metadata: Whether to include metadata in the response. - include_resources: Whether to include resources in the response. - **kwargs: Additional keyword arguments. - - Returns: - A list of query options. - """ - options = [ - selectinload(jl_arg(ResourceRequestSchema.requested_resources)), - ] - - if include_resources: - options.extend( - [ - selectinload(jl_arg(ResourceRequestSchema.component)), - selectinload( - jl_arg(ResourceRequestSchema.step_run) - ).joinedload(jl_arg(StepRunSchema.pipeline_run)), - selectinload(jl_arg(ResourceRequestSchema.pool)), - selectinload(jl_arg(ResourceRequestSchema.user)), - ] - ) - - return options - - @classmethod - def from_request( - cls, - request: "ResourceRequestRequest", - ) -> "ResourceRequestSchema": - """Create a resource request schema from a request. - - Args: - request: The `ResourceRequestRequest` to create from. - - Returns: - The created `ResourceRequestSchema`. - """ - return cls( - user_id=request.user, - component_id=request.component_id, - step_run_id=request.step_run_id, - status=ResourceRequestStatus.PENDING.value, - preemption_initiated_by_id=None, - preemptible=request.preemptible, - ) - - def to_model( - self, - include_metadata: bool = False, - include_resources: bool = False, - **kwargs: Any, - ) -> "ResourceRequestResponse": - """Creates a `ResourceRequestResponse` from a `ResourceRequestSchema`. - - Args: - include_metadata: Whether to include metadata in the response. - include_resources: Whether to include resources in the response. - **kwargs: Additional keyword arguments. - - Returns: - A `ResourceRequestResponse` object. - """ - body = ResourceRequestResponseBody( - created=self.created, - updated=self.updated, - user_id=self.user_id, - requested_resources={ - r.key: r.amount for r in self.requested_resources - }, - status=ResourceRequestStatus(self.status), - status_reason=self.status_reason, - preemptible=self.preemptible, - ) - - metadata = None - if include_metadata: - metadata = ResourceRequestResponseMetadata() - - resources = None - if include_resources: - resources = ResourceRequestResponseResources( - user=self.user.to_model() if self.user else None, - component=self.component.to_model() - if self.component - else None, - step_run=self.step_run.to_model() if self.step_run else None, - pipeline_run=self.step_run.pipeline_run.to_model() - if self.step_run - else None, - pool=self.pool.to_model() if self.pool else None, - ) - - return ResourceRequestResponse( - id=self.id, - body=body, - metadata=metadata, - resources=resources, - ) - - -class ResourceRequestResourceSchema(BaseSchema, table=True): - """Resource request resource schema.""" - - __tablename__ = "resource_request_resource" - __table_args__ = ( - UniqueConstraint( - "request_id", - "key", - name="unique_resource_request_resource_request_id_key", - ), - ) - - request_id: UUID = build_foreign_key_field( - source=__tablename__, - target=ResourceRequestSchema.__tablename__, - source_column="request_id", - target_column="id", - ondelete="CASCADE", - nullable=False, - ) - request: "ResourceRequestSchema" = Relationship( - back_populates="requested_resources" - ) - key: str - amount: int diff --git a/src/zenml/zen_stores/schemas/step_run_schemas.py b/src/zenml/zen_stores/schemas/step_run_schemas.py index f3311f727e5..848f2b8422d 100644 --- a/src/zenml/zen_stores/schemas/step_run_schemas.py +++ b/src/zenml/zen_stores/schemas/step_run_schemas.py @@ -70,9 +70,6 @@ from zenml.zen_stores.schemas.artifact_schemas import ArtifactVersionSchema from zenml.zen_stores.schemas.logs_schemas import LogsSchema from zenml.zen_stores.schemas.model_schemas import ModelVersionSchema - from zenml.zen_stores.schemas.resource_request_schemas import ( - ResourceRequestSchema, - ) from zenml.zen_stores.schemas.run_metadata_schemas import RunMetadataSchema @@ -182,6 +179,7 @@ class StepRunSchema(NamedSchema, RunMetadataInterface, table=True): ondelete="SET NULL", nullable=True, ) + resource_request_id: Optional[UUID] = Field(nullable=True) # Relationships project: "ProjectSchema" = Relationship(back_populates="step_runs") @@ -254,10 +252,6 @@ class StepRunSchema(NamedSchema, RunMetadataInterface, table=True): nullable=True, ) ) - resource_request: Optional["ResourceRequestSchema"] = Relationship( - back_populates="step_run" - ) - model_config = ConfigDict(protected_namespaces=()) # type: ignore[assignment] @classmethod @@ -318,7 +312,6 @@ def get_query_options( jl_arg(ModelVersionSchema.model), innerjoin=True ), single_loader(jl_arg(StepRunSchema.user)), - single_loader(jl_arg(StepRunSchema.resource_request)), selectinload(jl_arg(StepRunSchema.input_artifacts)) .joinedload( jl_arg(StepRunInputArtifactSchema.artifact_version), @@ -480,6 +473,7 @@ def to_model( created=self.created, updated=self.updated, model_version_id=self.model_version_id, + resource_request_id=self.resource_request_id, substitutions=step.config.substitutions, heartbeat_threshold=self.heartbeat_threshold, ) @@ -544,11 +538,6 @@ def to_model( ], inputs=input_artifacts, outputs=output_artifacts, - resource_request=self.resource_request.to_model( - include_metadata=True, include_resources=False - ) - if self.resource_request - else None, ) return StepRunResponse( diff --git a/src/zenml/zen_stores/sql_zen_store.py b/src/zenml/zen_stores/sql_zen_store.py index 71951997aac..979c199196e 100644 --- a/src/zenml/zen_stores/sql_zen_store.py +++ b/src/zenml/zen_stores/sql_zen_store.py @@ -168,6 +168,8 @@ MetadataResourceTypes, ModelStages, OnboardingStep, + ResourceRequestReclaimTolerance, + ResourceRequestRuntimeState, ResourceRequestStatus, RunWaitConditionLeaseMode, RunWaitConditionResolution, @@ -177,6 +179,7 @@ StackComponentType, StackDeploymentProvider, StepRunInputArtifactType, + StepRuntime, StoreType, TaggableResourceTypes, VisualizationResourceTypes, @@ -293,14 +296,6 @@ ProjectScopedFilter, ProjectScopedRequest, ProjectUpdate, - ResourcePoolFilter, - ResourcePoolRequest, - ResourcePoolResponse, - ResourcePoolSubjectPolicyFilter, - ResourcePoolSubjectPolicyRequest, - ResourcePoolSubjectPolicyResponse, - ResourcePoolSubjectPolicyUpdate, - ResourcePoolUpdate, ResourceRequestFilter, ResourceRequestRequest, ResourceRequestResponse, @@ -375,6 +370,9 @@ UserScopedRequest, UserUpdate, ) +from zenml.models.v2.core.resource_request import ( + ResourceRequestRenewalRequest, +) from zenml.service_connectors.service_connector_registry import ( service_connector_registry, ) @@ -477,6 +475,7 @@ if TYPE_CHECKING: from concurrent.futures import Future + from zenml.config import ResourceSettings from zenml.metadata.metadata_types import MetadataType, MetadataTypeEnum from zenml.models.v2.core.triggers import ( TriggerExecutionInfo, @@ -3970,151 +3969,6 @@ def _fail_if_component_with_name_type_exists( f"component with the same name and type." ) - # -------------------- Resource Pools ------------- - - def create_resource_pool( - self, resource_pool: ResourcePoolRequest - ) -> ResourcePoolResponse: - """Create a resource pool. - - Args: - resource_pool: The resource pool to create. - - Returns: - The created resource pool. - """ - return self.resource_pools.create_resource_pool(resource_pool) - - def get_resource_pool( - self, resource_pool_id: UUID, hydrate: bool = True - ) -> ResourcePoolResponse: - """Get a resource pool by ID. - - Args: - resource_pool_id: The ID of the resource pool to get. - hydrate: Flag deciding whether to hydrate the output model(s) - by including metadata fields in the response. - - Returns: - The resource pool. - """ - return self.resource_pools.get_resource_pool( - resource_pool_id, hydrate=hydrate - ) - - def list_resource_pools( - self, filter_model: ResourcePoolFilter, hydrate: bool = False - ) -> Page[ResourcePoolResponse]: - """List all resource pools matching the given filter criteria. - - Args: - filter_model: All filter parameters including pagination - params. - hydrate: Flag deciding whether to hydrate the output model(s) - by including metadata fields in the response. - - Returns: - A list of all resource pools matching the filter criteria. - """ - return self.resource_pools.list_resource_pools( - filter_model, hydrate=hydrate - ) - - def update_resource_pool( - self, resource_pool_id: UUID, update: ResourcePoolUpdate - ) -> ResourcePoolResponse: - """Update an existing resource pool. - - Args: - resource_pool_id: The ID of the resource pool to update. - update: The update to be applied to the resource pool. - - Returns: - The updated resource pool. - """ - return self.resource_pools.update_resource_pool( - resource_pool_id, update - ) - - def delete_resource_pool(self, resource_pool_id: UUID) -> None: - """Delete a resource pool. - - Args: - resource_pool_id: The ID of the resource pool to delete. - """ - self.resource_pools.delete_resource_pool(resource_pool_id) - - def create_resource_pool_subject_policy( - self, policy: ResourcePoolSubjectPolicyRequest - ) -> ResourcePoolSubjectPolicyResponse: - """Create a resource pool subject policy. - - Args: - policy: The policy to create. - - Returns: - The created policy. - """ - return self.resource_pools.create_resource_pool_subject_policy(policy) - - def get_resource_pool_subject_policy( - self, policy_id: UUID, hydrate: bool = True - ) -> ResourcePoolSubjectPolicyResponse: - """Get a resource pool subject policy by ID. - - Args: - policy_id: The ID of the policy to get. - hydrate: Whether to include metadata fields. - - Returns: - The requested policy. - """ - return self.resource_pools.get_resource_pool_subject_policy( - policy_id, hydrate=hydrate - ) - - def list_resource_pool_subject_policies( - self, - filter_model: ResourcePoolSubjectPolicyFilter, - hydrate: bool = False, - ) -> Page[ResourcePoolSubjectPolicyResponse]: - """List resource pool subject policies. - - Args: - filter_model: All filter parameters including pagination params. - hydrate: Whether to include metadata fields. - - Returns: - Matching policies. - """ - return self.resource_pools.list_resource_pool_subject_policies( - filter_model, hydrate=hydrate - ) - - def update_resource_pool_subject_policy( - self, policy_id: UUID, update: ResourcePoolSubjectPolicyUpdate - ) -> ResourcePoolSubjectPolicyResponse: - """Update an existing resource pool subject policy. - - Args: - policy_id: The ID of the policy to update. - update: The update model. - - Returns: - The updated policy. - """ - return self.resource_pools.update_resource_pool_subject_policy( - policy_id, update - ) - - def delete_resource_pool_subject_policy(self, policy_id: UUID) -> None: - """Delete a resource pool subject policy. - - Args: - policy_id: The ID of the policy to delete. - """ - self.resource_pools.delete_resource_pool_subject_policy(policy_id) - # -------------------- Resource Requests ------------- def get_resource_request( @@ -4152,13 +4006,40 @@ def list_resource_requests( filter_model, hydrate=hydrate ) - def delete_resource_request(self, resource_request_id: UUID) -> None: - """Delete a resource request. + def release_resource_request( + self, + resource_request_id: UUID, + ) -> ResourceRequestResponse: + """Release a resource request on behalf of its owner. + + Args: + resource_request_id: The ID of the resource request to release. + + Returns: + The released resource request. + """ + return self.resource_pools.release_resource_request( + resource_request_id, + ) + + def renew_resource_request( + self, + resource_request_id: UUID, + renewal_request: ResourceRequestRenewalRequest, + ) -> ResourceRequestResponse: + """Renew a resource request lease. Args: - resource_request_id: The ID of the resource request to delete. + resource_request_id: The ID of the resource request to renew. + renewal_request: The renewed lease expiration timestamp. + + Returns: + The renewed resource request. """ - self.resource_pools.delete_resource_request(resource_request_id) + return self.resource_pools.renew_resource_request( + resource_request_id, + renewal_request, + ) # -------------------------- Devices ------------------------- @@ -8324,7 +8205,7 @@ def delete_trigger(self, trigger_id: UUID, soft: bool = True) -> None: delete(TriggerSnapshotSchema).where( col(TriggerSnapshotSchema.trigger_id) == trigger_id ) - ) # type: ignore[call-overload, unused-ignore] + ) session.commit() @@ -11473,6 +11354,127 @@ def get_stack_deployment_stack( "Stack deployments are not supported by local ZenML deployments." ) + @staticmethod + def _validate_reclaim_tolerance_for_resource_request( + resource_settings: "ResourceSettings", + runtime: StepRuntime, + heartbeat_enabled: bool, + step_name: str, + ) -> None: + """Validate that reclaim tolerance matches runtime capabilities. + + Args: + resource_settings: Resource settings from the step configuration. + runtime: Resolved dynamic step runtime. + heartbeat_enabled: Whether heartbeat is enabled for the step run. + step_name: Step name used in error messages. + + Raises: + IllegalOperationError: If the reclaim tolerance cannot be honored. + """ + reclaim_tolerance = resource_settings.effective_reclaim_tolerance( + runtime + ) + if reclaim_tolerance == ResourceRequestReclaimTolerance.NONE: + return + + if ( + resource_settings.reclaim_tolerance_explicitly_set + and runtime == StepRuntime.INLINE + ): + raise IllegalOperationError( + f"Step `{step_name}` is configured with reclaim tolerance " + f"`{resource_settings.reclaim_tolerance}` but will run inline. " + "Inline dynamic steps are not reclaimable. Configure the step " + "to run isolated or set reclaim tolerance to `none`." + ) + + if ( + resource_settings.reclaim_tolerance_explicitly_set + and not heartbeat_enabled + ): + raise IllegalOperationError( + f"Step `{step_name}` is configured with reclaim tolerance " + f"`{reclaim_tolerance}` but heartbeat is disabled. Enable " + "heartbeat or set reclaim tolerance to `none`." + ) + + def _renew_step_resource_request_from_heartbeat( + self, + session: Session, + step_run: StepRunSchema, + heartbeat_liveness_timeout_seconds: int | None = None, + ) -> ExecutionStatus: + """Renew the step resource request lease during heartbeat. + + Args: + session: Active database session. + step_run: Step run schema receiving the heartbeat. + heartbeat_liveness_timeout_seconds: Optional number of seconds the + server should wait for another heartbeat before considering the + heartbeat client dead. + + Returns: + The status that should be returned to the heartbeat caller. + """ + step_status = ExecutionStatus(step_run.status) + if ( + not self.resource_pools_enabled + or step_run.resource_request_id is None + or step_run.heartbeat_threshold is None + ): + return step_status + + if heartbeat_liveness_timeout_seconds is None: + lease_duration = timedelta(minutes=step_run.heartbeat_threshold) + else: + lease_duration = timedelta( + seconds=heartbeat_liveness_timeout_seconds + ) + + try: + request = self.resource_pools.renew_resource_request( + step_run.resource_request_id, + ResourceRequestRenewalRequest( + lease_expires_at=utc_now() + lease_duration, + runtime_state=ResourceRequestRuntimeState.RUNNING, + ), + ) + except KeyError: + logger.warning( + "Resource request `%s` for step `%s` no longer exists. " + "Cancelling the step run.", + step_run.resource_request_id, + step_run.name, + ) + request_status = ResourceRequestStatus.CANCELLED + except Exception as e: + logger.warning( + "Failed to renew resource request `%s` for step `%s`: %s", + step_run.resource_request_id, + step_run.name, + e, + ) + return step_status + else: + request_status = request.status + + if request_status in { + ResourceRequestStatus.PREEMPTING, + ResourceRequestStatus.PREEMPTED, + ResourceRequestStatus.CANCELLED, + ResourceRequestStatus.REJECTED, + ResourceRequestStatus.RELEASED, + ResourceRequestStatus.EXPIRED, + }: + if not step_status.is_finished: + step_run.status = ExecutionStatus.CANCELLING.value + session.add(step_run) + session.commit() + return ExecutionStatus.CANCELLING + + return step_status + # ----------------------------- Step runs ----------------------------- def create_run_step(self, step_run: StepRunRequest) -> StepRunResponse: @@ -11538,6 +11540,31 @@ def create_run_step(self, step_run: StepRunRequest) -> StepRunResponse: step_run.dynamic_config or run.get_step_configuration(step_name=step_run.name) ) + resource_runtime: Optional[StepRuntime] = None + resource_request_heartbeat_enabled: Optional[bool] = None + if ( + self.resource_pools_enabled + and step_run.status + in { + ExecutionStatus.INITIALIZING, + ExecutionStatus.PROVISIONING, + ExecutionStatus.RUNNING, + } + and step_run.resource_requester + ): + resource_settings = step_config.config.resource_settings + resource_runtime = ( + step_run.resource_request_runtime or StepRuntime.ISOLATED + ) + resource_request_heartbeat_enabled = ( + step_config.spec.enable_heartbeat and run.enable_heartbeat + ) + self._validate_reclaim_tolerance_for_resource_request( + resource_settings=resource_settings, + runtime=resource_runtime, + heartbeat_enabled=resource_request_heartbeat_enabled, + step_name=step_run.name, + ) # Release the read locks of the previous two queries before we # try to acquire more exclusive locks @@ -11890,6 +11917,7 @@ def create_run_step(self, step_run: StepRunRequest) -> StepRunResponse: ) session.refresh(step_schema) + created_resource_request: ResourceRequestResponse | None = None if ( self.resource_pools_enabled and step_schema.status @@ -11900,32 +11928,70 @@ def create_run_step(self, step_run: StepRunRequest) -> StepRunResponse: } and step_run.resource_requester ): - requested_resources = step_config.config.resource_settings.merged_requested_resources() - requested_resources["step_run"] = 1 - - request = self.resource_pools.create_resource_request( - session=session, - resource_request=ResourceRequestRequest( - user=step_run.user, - component_id=step_run.resource_requester, - step_run_id=step_schema.id, - requested_resources=requested_resources, - preemptible=step_config.config.resource_settings.preemptible, - ), + resource_settings = step_config.config.resource_settings + resource_runtime = ( + resource_runtime + or step_run.resource_request_runtime + or StepRuntime.ISOLATED ) - if ( - request is not None - and request.status != ResourceRequestStatus.ALLOCATED - ): - step_schema.status = ExecutionStatus.QUEUED.value - session.add(step_schema) - session.commit() + heartbeat_enabled = ( + resource_request_heartbeat_enabled + if resource_request_heartbeat_enabled is not None + else step_schema.heartbeat_threshold is not None + ) + demands = resource_settings.merged_resource_demands() + + if demands: + lease_expires_at = None + if ( + resource_runtime == StepRuntime.ISOLATED + and heartbeat_enabled + and step_schema.heartbeat_threshold is not None + ): + lease_expires_at = utc_now() + timedelta( + minutes=step_schema.heartbeat_threshold + ) + + request = self.resource_pools.create_resource_request( + session, + ResourceRequestRequest( + user=step_run.user, + component_ids=[step_run.resource_requester], + step_run_id=step_schema.id, + demands=demands, + reclaim_tolerance=resource_settings.effective_reclaim_tolerance( + resource_runtime + ), + lease_expires_at=lease_expires_at, + allocation_wait_timeout_seconds=( + resource_settings.allocation_wait_timeout_seconds + ), + ), + ) + if ( + request.status + != ResourceRequestStatus.NO_MATCHING_POOL + ): + step_schema.resource_request_id = request.id + created_resource_request = request + if request.status != ResourceRequestStatus.ALLOCATED: + step_schema.status = ExecutionStatus.QUEUED.value + session.add(step_schema) + session.commit() session.refresh(step_schema) - return step_schema.to_model( + step_run_response = step_schema.to_model( include_metadata=True, include_resources=True ) + if ( + created_resource_request is not None + and step_run_response.resources is not None + ): + step_run_response.resources.resource_request = ( + created_resource_request + ) + return step_run_response def get_run_step( self, step_run_id: UUID, hydrate: bool = True @@ -12160,7 +12226,9 @@ def _set_hook_invocation_output_artifact( ) def update_step_heartbeat( - self, step_run_id: UUID + self, + step_run_id: UUID, + heartbeat_liveness_timeout_seconds: int | None = None, ) -> StepHeartbeatResponse: """Updates a step run heartbeat value. @@ -12168,6 +12236,9 @@ def update_step_heartbeat( Args: step_run_id: ID of the step run. + heartbeat_liveness_timeout_seconds: Optional number of seconds the + server should wait for another heartbeat before considering the + heartbeat client dead. Returns: Step heartbeat response (minimal info, id, status & latest_heartbeat). @@ -12183,10 +12254,19 @@ def update_step_heartbeat( session.commit() session.refresh(existing_step_run) + heartbeat_status = ( + self._renew_step_resource_request_from_heartbeat( + session=session, + step_run=existing_step_run, + heartbeat_liveness_timeout_seconds=( + heartbeat_liveness_timeout_seconds + ), + ) + ) return StepHeartbeatResponse( id=existing_step_run.id, - status=ExecutionStatus(existing_step_run.status), + status=heartbeat_status, latest_heartbeat=existing_step_run.latest_heartbeat, heartbeat_enabled=existing_step_run.heartbeat_threshold is not None, @@ -12197,6 +12277,7 @@ def validate_and_update_heartbeat( step_run_id: UUID, token_run_id: UUID | None = None, token_schedule_id: UUID | None = None, + heartbeat_liveness_timeout_seconds: int | None = None, ) -> StepHeartbeatResponse: """Updates & Validates a step run heartbeat value. @@ -12206,6 +12287,9 @@ def validate_and_update_heartbeat( step_run_id: ID of the step run. token_run_id: Pipeline run id of the auth context token_schedule_id: Schedule id of the auth context + heartbeat_liveness_timeout_seconds: Optional number of seconds the + server should wait for another heartbeat before considering the + heartbeat client dead. Returns: Step heartbeat response (minimal info, id, status & latest_heartbeat). @@ -12249,10 +12333,19 @@ def validate_and_update_heartbeat( latest_heartbeat = datetime.now(timezone.utc) step_run.latest_heartbeat = latest_heartbeat session.commit() + heartbeat_status = ( + self._renew_step_resource_request_from_heartbeat( + session=session, + step_run=step_run, + heartbeat_liveness_timeout_seconds=( + heartbeat_liveness_timeout_seconds + ), + ) + ) return StepHeartbeatResponse( id=step_run_id, - status=ExecutionStatus(step_run.status), + status=heartbeat_status, latest_heartbeat=latest_heartbeat, heartbeat_enabled=step_run.heartbeat_threshold is not None, pipeline_run_status=ExecutionStatus(run.status), diff --git a/src/zenml/zen_stores/zen_store_interface.py b/src/zenml/zen_stores/zen_store_interface.py index f15f7467ba5..175d1b45da5 100644 --- a/src/zenml/zen_stores/zen_store_interface.py +++ b/src/zenml/zen_stores/zen_store_interface.py @@ -2721,12 +2721,17 @@ def update_run_step( @abstractmethod def update_step_heartbeat( - self, step_run_id: UUID + self, + step_run_id: UUID, + heartbeat_liveness_timeout_seconds: Optional[int] = None, ) -> StepHeartbeatResponse: """Updates a step run heartbeat. Args: step_run_id: The ID of the step to update. + heartbeat_liveness_timeout_seconds: Optional number of seconds the + server should wait for another heartbeat before considering the + heartbeat client dead. Returns: The step heartbeat response. diff --git a/tests/unit/config/test_resource_settings.py b/tests/unit/config/test_resource_settings.py index d193bc531ec..cb38eb91b7c 100644 --- a/tests/unit/config/test_resource_settings.py +++ b/tests/unit/config/test_resource_settings.py @@ -18,7 +18,12 @@ import pytest from pydantic import ValidationError -from zenml.config.resource_settings import ByteUnit, ResourceSettings +from zenml.config.resource_settings import ( + ByteUnit, + PoolResourceDemand, + ResourceSettings, +) +from zenml.enums import ResourceRequestReclaimTolerance, StepRuntime def test_unit_byte_value_defined_for_all_values(): @@ -97,56 +102,197 @@ def test_resource_config_memory_conversion(): assert r.get_memory(unit="GiB") == 11 -def test_pool_resources_empty_dict_merges_to_empty() -> None: - """Empty pool_resources is accepted; merged map has no keys from the map.""" - rs = ResourceSettings(pool_resources={}) +def test_memory_quantity_and_unit_parses_configured_suffix() -> None: + """memory_quantity_and_unit returns the literal quantity and ByteUnit.""" + rs = ResourceSettings(memory="512MB") + assert rs.memory_quantity_and_unit() == (512, ByteUnit.MB) + + rs = ResourceSettings(memory="11GiB") + assert rs.memory_quantity_and_unit() == (11, ByteUnit.GIB) + + +def test_resources_empty_dict_normalizes_to_empty_list() -> None: + """Empty resource maps are accepted and normalize to no demands.""" + rs = ResourceSettings(resources={}) assert rs.empty is True - assert rs.merged_requested_resources() == {} + assert rs.resources == {} + assert rs.normalized_resources == [] + assert rs.merged_resource_demands() == [] + + +def test_resources_from_name_to_quantity_map() -> None: + """Name-to-quantity maps normalize into pool demands.""" + rs = ResourceSettings(resources={"tpu": 2, "widgets": 3}) + assert rs.resources == {"tpu": 2, "widgets": 3} + assert [demand.name for demand in rs.normalized_resources] == [ + "tpu", + "widgets", + ] + assert [demand.quantity for demand in rs.normalized_resources] == [2, 3] + +def test_resources_from_string_defaults_quantity_to_one() -> None: + """A bare resource name becomes a single-quantity demand.""" + rs = ResourceSettings(resources="tpu") + assert rs.resources == "tpu" + assert rs.normalized_resources == [ + PoolResourceDemand(name="tpu", quantity=1), + ] -def test_merged_requested_resources_from_pool_resources_only() -> None: - """Custom keys pass through when typed resource fields are unset.""" - rs = ResourceSettings(pool_resources={"tpu": 2, "widgets": 3}) - assert rs.merged_requested_resources() == {"tpu": 2, "widgets": 3} + +def test_legacy_pool_resources_input_is_normalized() -> None: + """Legacy pool_resources input is folded into resources.""" + rs = ResourceSettings(pool_resources={"tpu": 2}) + assert rs.resources == {"tpu": 2} + assert rs.normalized_resources[0].name == "tpu" + assert rs.normalized_resources[0].quantity == 2 + + +def test_resources_from_pool_resource_demand_list() -> None: + """A list of PoolResourceDemand values is accepted via resources.""" + rs = ResourceSettings( + resources=[ + PoolResourceDemand(name="gpu", quantity=2, kind="gpu"), + ] + ) + assert rs.normalized_resources[0].name == "gpu" + assert rs.normalized_resources[0].kind == "gpu" -def test_merged_requested_resources_typed_fields_override_pool_resources() -> ( +def test_merged_resource_demands_keeps_pool_and_typed_fields_separate() -> ( None ): - """gpu_count, cpu_count, and memory override duplicate dict keys.""" + """Pool resources and typed fields become separate kind-based demands.""" rs = ResourceSettings( - pool_resources={"gpu": 9, "mcpu": 1, "memory_mb": 1}, + resources={"tpu": 2}, gpu_count=2, cpu_count=1.0, memory="512MB", + gpu_class="h100", ) - merged = rs.merged_requested_resources() - assert merged["gpu"] == 2 - assert merged["mcpu"] == 1000 - assert merged["memory_mb"] == 512 + demands = rs.merged_resource_demands() + assert demands[0].resource == "tpu" + assert demands[0].quantity == 2 + assert demands[1].kind == "gpu" + assert demands[1].resource is None + assert demands[1].quantity == 2 + assert demands[1].class_name == "h100" + assert demands[2].kind == "cpu" + assert demands[2].unit == "CPU" + assert demands[2].quantity == 1 + assert demands[3].kind == "memory" + assert demands[3].unit == ByteUnit.MB.value + assert demands[3].quantity == 512 + + +def test_merged_resource_demands_gpu_count_zero_omits_gpu() -> None: + """gpu_count zero does not add a GPU demand.""" + rs = ResourceSettings(gpu_count=0) + assert rs.merged_resource_demands() == [] + + +def test_merged_resource_demands_typed_fields_use_kind_convention() -> None: + """Typed fields map to kind/unit conventions instead of resource names.""" + rs = ResourceSettings(cpu_count=2.5, gpu_count=1, memory="1GB") + demands = rs.merged_resource_demands() + assert demands[0].kind == "gpu" + assert demands[0].quantity == 1 + assert demands[1].kind == "cpu" + assert demands[1].unit == "CPU" + assert demands[1].quantity == 3 + assert demands[2].kind == "memory" + assert demands[2].unit == ByteUnit.GB.value + assert demands[2].quantity == 1 + + +def test_empty_property_excludes_resources() -> None: + """empty stays True when only pool resources are configured.""" + rs = ResourceSettings(resources={"x": 1}) + assert rs.empty is True -def test_merged_requested_resources_gpu_count_zero_clears_gpu() -> None: - """gpu_count zero removes gpu from the merged map.""" - rs = ResourceSettings(pool_resources={"gpu": 2}, gpu_count=0) - assert "gpu" not in rs.merged_requested_resources() +def test_default_reclaim_tolerance_stays_any() -> None: + """Default reclaim tolerance remains backward-compatible.""" + rs = ResourceSettings() + assert rs.reclaim_tolerance is ResourceRequestReclaimTolerance.ANY + assert rs.reclaim_tolerance_explicitly_set is False + assert ( + rs.effective_reclaim_tolerance(StepRuntime.INLINE) + is ResourceRequestReclaimTolerance.NONE + ) + assert ( + rs.effective_reclaim_tolerance(StepRuntime.ISOLATED) + is ResourceRequestReclaimTolerance.ANY + ) -def test_merged_requested_resources_legacy_equivalent() -> None: - """Behavior matches the former StackComponent.get_requested_resources logic.""" - rs = ResourceSettings(cpu_count=2.5, gpu_count=1, memory="1GB") - assert rs.merged_requested_resources() == { - "gpu": 1, - "mcpu": 2500, - "memory_mb": 1000, - } +def test_effective_reclaim_tolerance_preserves_explicit_value() -> None: + """Explicit reclaim tolerance is used for runtime resource requests.""" + rs = ResourceSettings( + reclaim_tolerance=ResourceRequestReclaimTolerance.ANY + ) + assert rs.reclaim_tolerance_explicitly_set is True + assert ( + rs.effective_reclaim_tolerance(StepRuntime.INLINE) + is ResourceRequestReclaimTolerance.ANY + ) + assert ( + rs.effective_reclaim_tolerance(StepRuntime.ISOLATED) + is ResourceRequestReclaimTolerance.ANY + ) -def test_empty_property_excludes_pool_resources() -> None: - """empty stays True when only pool_resources is set. +def test_basic_resource_demands_include_typed_fields() -> None: + """Typed CPU/GPU/memory fields are basic resource demands.""" + assert ResourceSettings(cpu_count=1).has_basic_resource_demands is True + assert ResourceSettings(gpu_count=1).has_basic_resource_demands is True + assert ResourceSettings(memory="1GB").has_basic_resource_demands is True - pool_resources are for resource-pool scheduling; they are not part of the - generic resource fields that orchestrators and step operators consume. - """ - rs = ResourceSettings(pool_resources={"x": 1}) - assert rs.empty is True + +def test_basic_resource_demands_include_pool_resource_kinds() -> None: + """Pool resources with known kinds participate in isolation decisions.""" + rs = ResourceSettings( + resources=[ + PoolResourceDemand(name="gpu-slot", quantity=1, kind="gpu"), + ] + ) + assert rs.has_basic_resource_demands is True + + +def test_step_run_pool_resource_kind_is_not_basic_demand() -> None: + """step_run is no longer a stock infrastructure resource kind.""" + rs = ResourceSettings( + resources=[ + PoolResourceDemand(name="step-slot", quantity=1, kind="step_run"), + ] + ) + assert rs.has_basic_resource_demands is False + + +def test_custom_pool_resources_are_not_basic_demands() -> None: + """Opaque resources do not force isolated runtime by themselves.""" + rs = ResourceSettings(resources={"license": 1}) + assert rs.has_basic_resource_demands is False + + +def test_default_resource_pool_timeouts_are_one_hour() -> None: + """Default allocation and initialization timeouts are one hour.""" + rs = ResourceSettings() + assert rs.allocation_wait_timeout_seconds == 3600 + assert rs.initialization_lease_seconds == 3600 + + +def test_resource_pool_timeouts_do_not_affect_empty() -> None: + """Timeout defaults do not make resource settings non-empty.""" + assert ResourceSettings().empty is True + + +def test_legacy_preemptible_maps_to_reclaim_tolerance() -> None: + """Legacy preemptible input is accepted without storing the field.""" + rs = ResourceSettings(preemptible=True) + assert rs.reclaim_tolerance is ResourceRequestReclaimTolerance.COORDINATED + assert rs.reclaim_tolerance_explicitly_set is True + assert "preemptible" not in rs.model_dump() + + rs = ResourceSettings(preemptible=False) + assert rs.reclaim_tolerance is ResourceRequestReclaimTolerance.NONE diff --git a/tests/unit/execution/pipeline/dynamic/test_compilation.py b/tests/unit/execution/pipeline/dynamic/test_compilation.py new file mode 100644 index 00000000000..9eadab3d127 --- /dev/null +++ b/tests/unit/execution/pipeline/dynamic/test_compilation.py @@ -0,0 +1,95 @@ +# Copyright (c) ZenML GmbH 2026. All Rights Reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at: +# +# https://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express +# or implied. See the License for the specific language governing +# permissions and limitations under the License. +"""Tests for dynamic compilation helpers.""" + +from types import SimpleNamespace + +from zenml.config.resource_settings import PoolResourceDemand, ResourceSettings +from zenml.config.step_configurations import Step, StepConfiguration +from zenml.enums import StepRuntime +from zenml.execution.pipeline.dynamic.compilation import get_step_runtime + + +def _step_config(resource_settings: ResourceSettings) -> StepConfiguration: + """Create a minimal step configuration for runtime classification tests.""" + return Step.model_validate( + { + "spec": { + "source": "module.step_class", + "upstream_steps": [], + "inputs": {}, + }, + "config": { + "name": "step_name", + "enable_cache": True, + "settings": {"resources": resource_settings}, + }, + } + ).config + + +def test_basic_resource_demands_force_isolated_runtime() -> None: + """Basic resource pool kinds force dynamic isolation.""" + step_config = _step_config( + ResourceSettings( + resources=[ + PoolResourceDemand(name="gpu-slot", quantity=1, kind="gpu") + ] + ) + ) + + assert ( + get_step_runtime( + step_config=step_config, + pipeline_docker_settings=step_config.docker_settings, + orchestrator=SimpleNamespace(can_run_isolated_steps=True), + ) + is StepRuntime.ISOLATED + ) + + +def test_step_run_resource_demands_do_not_force_isolated_runtime() -> None: + """step_run is no longer treated as a basic infrastructure resource.""" + step_config = _step_config( + ResourceSettings( + resources=[ + PoolResourceDemand( + name="step-slot", quantity=1, kind="step_run" + ) + ] + ) + ) + + assert ( + get_step_runtime( + step_config=step_config, + pipeline_docker_settings=step_config.docker_settings, + orchestrator=SimpleNamespace(can_run_isolated_steps=True), + ) + is StepRuntime.INLINE + ) + + +def test_custom_resource_demands_do_not_force_isolated_runtime() -> None: + """Opaque pool resources keep the default inline runtime.""" + step_config = _step_config(ResourceSettings(resources={"license": 1})) + + assert ( + get_step_runtime( + step_config=step_config, + pipeline_docker_settings=step_config.docker_settings, + orchestrator=SimpleNamespace(can_run_isolated_steps=True), + ) + is StepRuntime.INLINE + ) diff --git a/tests/unit/integrations/kubernetes/test_resource_request_allocations.py b/tests/unit/integrations/kubernetes/test_resource_request_allocations.py new file mode 100644 index 00000000000..1c8a3b5634f --- /dev/null +++ b/tests/unit/integrations/kubernetes/test_resource_request_allocations.py @@ -0,0 +1,370 @@ +"""Tests for Kubernetes resource request allocation helpers.""" + +from datetime import datetime, timezone +from uuid import UUID, uuid4 + +from zenml.enums import ( + ResourceRequestReclaimTolerance, + ResourceRequestStatus, +) +from zenml.integrations.kubernetes import kube_utils +from zenml.integrations.kubernetes.flavors import ( + KubernetesStepOperatorSettings, +) +from zenml.integrations.kubernetes.pod_settings import KubernetesPodSettings +from zenml.models import ( + ResourcePoolAllocation, + ResourcePoolCapacityComponentSettings, + ResourceRequestDemand, + ResourceRequestResponse, + ResourceRequestResponseBody, + ResourceRequestResponseMetadata, + ResourceRequestResponseResources, + ResourceRequestServiceConnectorSettings, +) + + +def _allocation( + *, + request_id: UUID, + demand_index: int | None, + quantity: int, + unit: str | None = None, + resource: str = "resource", + resource_kind: str = "cpu", + component_id: UUID | None = None, + component_settings: ( + list[ResourcePoolCapacityComponentSettings] | None + ) = None, +) -> ResourcePoolAllocation: + return ResourcePoolAllocation( + id=uuid4(), + request_id=request_id, + demand_index=demand_index, + pool_id=uuid4(), + pool_name="pool", + capacity_entry_id=uuid4(), + capacity_entry_name="pool-capacity-entry", + resource_id=uuid4(), + resource=resource, + resource_kind=resource_kind, + class_name="default", + quantity=quantity, + unit=unit, + policy_id=uuid4(), + priority=0, + component_id=component_id or uuid4(), + component_settings=component_settings or [], + preemption_state="active", + ) + + +def _resource_request( + *, + demands: list[ResourceRequestDemand], + allocations: list[ResourcePoolAllocation], + component_settings: dict[str, object] | None = None, + service_connector_settings: ( + ResourceRequestServiceConnectorSettings | None + ) = None, +) -> ResourceRequestResponse: + request_id = allocations[0].request_id if allocations else uuid4() + return ResourceRequestResponse( + id=request_id, + name=str(request_id), + body=ResourceRequestResponseBody( + status=ResourceRequestStatus.ALLOCATED, + user=uuid4(), + created=datetime.now(tz=timezone.utc), + updated=datetime.now(tz=timezone.utc), + version=1, + demands=demands, + reclaim_tolerance=ResourceRequestReclaimTolerance.NONE, + ), + metadata=ResourceRequestResponseMetadata(), + resources=ResourceRequestResponseResources( + component_settings=component_settings or {}, + service_connector_settings=service_connector_settings, + allocations=allocations, + ), + ) + + +def test_unitless_cpu_allocation_threshold() -> None: + """Unitless CPU allocations above 50 are interpreted as milliCPU.""" + request_id = uuid4() + demands = [ + ResourceRequestDemand(kind="cpu", quantity=50), + ResourceRequestDemand(kind="cpu", quantity=51), + ] + + full_cpu_request = _resource_request( + demands=demands, + allocations=[ + _allocation( + request_id=request_id, + demand_index=0, + quantity=50, + ) + ], + ) + millicpu_request = _resource_request( + demands=demands, + allocations=[ + _allocation( + request_id=request_id, + demand_index=1, + quantity=51, + ) + ], + ) + + full_cpu_settings = ( + kube_utils.apply_resource_request_allocations_to_pod_settings( + full_cpu_request + ) + ) + millicpu_settings = ( + kube_utils.apply_resource_request_allocations_to_pod_settings( + millicpu_request + ) + ) + + assert full_cpu_settings.resources["requests"]["cpu"] == "50" + assert full_cpu_settings.resources["limits"]["cpu"] == "50" + assert millicpu_settings.resources["requests"]["cpu"] == "51m" + assert millicpu_settings.resources["limits"]["cpu"] == "51m" + + +def test_cpu_memory_and_gpu_allocations_override_requests_and_limits() -> None: + """Basic allocations overwrite matching pod requests and limits.""" + request_id = uuid4() + request = _resource_request( + demands=[ + ResourceRequestDemand(kind="cpu", quantity=2, unit="CPU"), + ResourceRequestDemand(kind="memory", quantity=1, unit="GiB"), + ResourceRequestDemand(kind="gpu", quantity=1), + ], + allocations=[ + _allocation( + request_id=request_id, + demand_index=0, + quantity=2, + unit="CPU", + ), + _allocation( + request_id=request_id, + demand_index=1, + quantity=1, + unit="GiB", + resource_kind="memory", + ), + _allocation( + request_id=request_id, + demand_index=2, + quantity=1, + resource_kind="gpu", + ), + ], + ) + pod_settings = KubernetesPodSettings( + resources={ + "requests": { + "cpu": "1", + "example.com/custom": "keep-request", + }, + "limits": { + "memory": "1Gi", + "example.com/custom": "keep-limit", + }, + } + ) + + result = kube_utils.apply_resource_request_allocations_to_pod_settings( + request, + pod_settings=pod_settings, + ) + + assert result.resources["requests"]["cpu"] == "2" + assert result.resources["limits"]["cpu"] == "2" + assert result.resources["requests"]["memory"] == "1024Mi" + assert result.resources["limits"]["memory"] == "1024Mi" + assert result.resources["requests"]["nvidia.com/gpu"] == "1" + assert result.resources["limits"]["nvidia.com/gpu"] == "1" + assert result.resources["requests"]["example.com/custom"] == "keep-request" + assert result.resources["limits"]["example.com/custom"] == "keep-limit" + + +def test_multiple_allocations_for_same_kind_are_summed() -> None: + """Multiple allocations for one resource kind are summed.""" + request_id = uuid4() + request = _resource_request( + demands=[ + ResourceRequestDemand(kind="cpu", quantity=250, unit="mCPU"), + ResourceRequestDemand(kind="gpu", quantity=1), + ], + allocations=[ + _allocation( + request_id=request_id, + demand_index=0, + quantity=250, + unit="mCPU", + ), + _allocation( + request_id=request_id, + demand_index=0, + quantity=250, + unit="mCPU", + ), + _allocation( + request_id=request_id, + demand_index=1, + quantity=1, + resource_kind="gpu", + ), + _allocation( + request_id=request_id, + demand_index=1, + quantity=2, + resource_kind="gpu", + ), + ], + ) + + result = kube_utils.apply_resource_request_allocations_to_pod_settings( + request + ) + + assert result.resources["requests"]["cpu"] == "500m" + assert result.resources["limits"]["cpu"] == "500m" + assert result.resources["requests"]["nvidia.com/gpu"] == "3" + assert result.resources["limits"]["nvidia.com/gpu"] == "3" + + +def test_default_allocations_use_resource_kind_without_demand() -> None: + """Grant-default allocations use descriptor kind without demand lookup.""" + request_id = uuid4() + request = _resource_request( + demands=[ResourceRequestDemand(kind="gpu", quantity=1)], + allocations=[ + _allocation( + request_id=request_id, + demand_index=None, + resource="NVIDIA H200", + resource_kind="gpu", + quantity=1, + ), + _allocation( + request_id=request_id, + demand_index=None, + resource="small-cpu-default", + resource_kind="cpu", + quantity=500, + unit="mCPU", + ), + _allocation( + request_id=request_id, + demand_index=None, + resource="small-memory-default", + resource_kind="memory", + quantity=512, + unit="MiB", + ), + ], + ) + + result = kube_utils.apply_resource_request_allocations_to_pod_settings( + request + ) + + assert result.resources["requests"]["nvidia.com/gpu"] == "1" + assert result.resources["limits"]["nvidia.com/gpu"] == "1" + assert result.resources["requests"]["cpu"] == "500m" + assert result.resources["limits"]["cpu"] == "500m" + assert result.resources["requests"]["memory"] == "512Mi" + assert result.resources["limits"]["memory"] == "512Mi" + + +def test_matching_component_settings_override_existing_settings() -> None: + """Resource request component settings override base settings.""" + component_id = uuid4() + request_id = uuid4() + request = _resource_request( + demands=[ResourceRequestDemand(kind="gpu", quantity=1)], + component_settings={ + "pod_settings": { + "node_selectors": { + "accelerator": "h200", + }, + "resources": { + "requests": { + "nvidia.com/gpu": "{{ quantity }}", + }, + }, + }, + }, + allocations=[ + _allocation( + request_id=request_id, + demand_index=0, + quantity=1, + component_id=component_id, + ) + ], + ) + settings = KubernetesStepOperatorSettings( + pod_settings=KubernetesPodSettings( + node_selectors={"accelerator": "old"}, + ) + ) + + result = kube_utils.apply_resource_request_component_settings( + settings=settings, + allocated_resource_request=request, + settings_class=KubernetesStepOperatorSettings, + ) + + assert result.pod_settings is not None + assert result.pod_settings.node_selectors == {"accelerator": "h200"} + assert ( + result.pod_settings.resources["requests"]["nvidia.com/gpu"] + == "{{ quantity }}" + ) + + +def test_service_connector_settings_do_not_override_component_settings() -> ( + None +): + """Service connector settings do not affect component settings.""" + component_id = uuid4() + request_id = uuid4() + request = _resource_request( + demands=[ResourceRequestDemand(kind="gpu", quantity=1)], + service_connector_settings=ResourceRequestServiceConnectorSettings( + connector_id=uuid4(), + resource_type="kubernetes-cluster", + resource_id="cluster", + ), + allocations=[ + _allocation( + request_id=request_id, + demand_index=0, + quantity=1, + component_id=component_id, + ), + ], + ) + settings = KubernetesStepOperatorSettings( + pod_settings=KubernetesPodSettings( + node_selectors={"accelerator": "old"}, + ) + ) + + result = kube_utils.apply_resource_request_component_settings( + settings=settings, + allocated_resource_request=request, + settings_class=KubernetesStepOperatorSettings, + ) + + assert result.pod_settings is not None + assert result.pod_settings.node_selectors == {"accelerator": "old"} diff --git a/tests/unit/orchestrators/test_step_launcher.py b/tests/unit/orchestrators/test_step_launcher.py index 2353c023971..7be5f38f88d 100644 --- a/tests/unit/orchestrators/test_step_launcher.py +++ b/tests/unit/orchestrators/test_step_launcher.py @@ -11,13 +11,15 @@ # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express # or implied. See the License for the specific language governing # permissions and limitations under the License. +"""Tests for step launcher helpers.""" from contextlib import ExitStack as does_not_raise +from types import SimpleNamespace from uuid import uuid4 import pytest -from zenml.enums import ExecutionStatus, StackComponentType +from zenml.enums import ExecutionStatus, StackComponentType, StepRuntime from zenml.orchestrators.step_launcher import ( StepLauncher, _get_step_operator, @@ -25,10 +27,44 @@ from zenml.stack import Stack -def test_step_operator_validation(local_stack, sample_step_operator): - """Tests that the step operator gets correctly extracted and validated - from the stack.""" +def _dynamic_launcher(step_operator=None) -> StepLauncher: + """Create a minimal launcher for dynamic runtime checks.""" + launcher = object.__new__(StepLauncher) + launcher._snapshot = SimpleNamespace( + is_dynamic=True, + pipeline_configuration=SimpleNamespace(docker_settings=None), + ) + launcher._step = SimpleNamespace( + config=SimpleNamespace(step_operator=step_operator) + ) + launcher._stack = SimpleNamespace(orchestrator=object()) + return launcher + + +def test_dynamic_orchestrator_detection_requires_isolated_runtime( + monkeypatch, +) -> None: + """Tests that only dynamic isolated orchestrator steps are detected.""" + monkeypatch.setattr( + "zenml.execution.pipeline.dynamic.compilation.get_step_runtime", + lambda **_: StepRuntime.ISOLATED, + ) + + assert _dynamic_launcher()._runs_with_dynamic_orchestrator() + assert not _dynamic_launcher( + step_operator="operator" + )._runs_with_dynamic_orchestrator() + + monkeypatch.setattr( + "zenml.execution.pipeline.dynamic.compilation.get_step_runtime", + lambda **_: StepRuntime.INLINE, + ) + assert not _dynamic_launcher()._runs_with_dynamic_orchestrator() + + +def test_step_operator_validation(local_stack, sample_step_operator): + """Tests that the step operator gets correctly extracted and validated.""" with pytest.raises(RuntimeError): _get_step_operator( stack=local_stack, step_operator_name="step_operator" @@ -53,6 +89,7 @@ def test_step_operator_validation(local_stack, sample_step_operator): def test_dynamic_command_step_success_publishes_status(mocker): + """Test that successful dynamic command steps publish a success status.""" launcher = object.__new__(StepLauncher) launcher._stack = mocker.Mock() launcher._stack.orchestrator.wait_for_isolated_step.return_value = ( @@ -92,6 +129,7 @@ def test_dynamic_command_step_success_publishes_status(mocker): def test_dynamic_command_step_failure_raises(mocker): + """Test that failed dynamic command steps raise the remote exception.""" launcher = object.__new__(StepLauncher) launcher._stack = mocker.Mock() launcher._stack.orchestrator.wait_for_isolated_step.return_value = ( @@ -160,6 +198,7 @@ def _make_isolated_step_launcher(mocker, status): def test_isolated_step_cleanup_called_on_success(mocker): + """Test that isolated step cleanup runs after successful remote steps.""" launcher = _make_isolated_step_launcher(mocker, ExecutionStatus.COMPLETED) step_run_info = mocker.Mock() @@ -171,6 +210,7 @@ def test_isolated_step_cleanup_called_on_success(mocker): def test_isolated_step_cleanup_called_on_failure(mocker): + """Test that isolated step cleanup runs after failed remote steps.""" launcher = _make_isolated_step_launcher(mocker, ExecutionStatus.FAILED) step_run_info = mocker.Mock() step_run_info.pipeline_step_name = "step_name" @@ -192,6 +232,7 @@ def test_isolated_step_cleanup_called_on_failure(mocker): def test_isolated_step_cleanup_skipped_when_not_waiting(mocker): + """Test that isolated step cleanup is skipped for asynchronous runs.""" launcher = _make_isolated_step_launcher(mocker, ExecutionStatus.COMPLETED) launcher._wait = False step_run_info = mocker.Mock() @@ -203,6 +244,7 @@ def test_isolated_step_cleanup_skipped_when_not_waiting(mocker): def test_cleanup_remote_step_dispatches_to_step_operator(mocker): + """Test that remote step cleanup dispatches to the selected operator.""" launcher = object.__new__(StepLauncher) launcher._stack = mocker.Mock() launcher._invocation_id = "step_name" @@ -226,6 +268,7 @@ def test_cleanup_remote_step_dispatches_to_step_operator(mocker): def test_cleanup_remote_step_swallows_errors(mocker): + """Test that remote cleanup errors are logged and swallowed.""" launcher = object.__new__(StepLauncher) launcher._stack = mocker.Mock() launcher._invocation_id = "step_name" diff --git a/tests/unit/orchestrators/test_step_launcher_resource_request_lease.py b/tests/unit/orchestrators/test_step_launcher_resource_request_lease.py new file mode 100644 index 00000000000..78d0cef7385 --- /dev/null +++ b/tests/unit/orchestrators/test_step_launcher_resource_request_lease.py @@ -0,0 +1,506 @@ +# Copyright (c) ZenML GmbH 2026. All Rights Reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at: +# +# https://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express +# or implied. See the License for the specific language governing +# permissions and limitations under the License. +"""Tests for step launcher resource request lease handling.""" + +from datetime import datetime, timedelta, timezone +from unittest.mock import MagicMock, patch +from uuid import uuid4 + +import pytest + +from zenml.config.resource_settings import ResourceSettings +from zenml.enums import ResourceRequestReclaimTolerance, ResourceRequestStatus +from zenml.models import ( + ResourceRequestRenewalRequest, + ResourceRequestResponse, + ResourceRequestResponseBody, + ResourceRequestResponseMetadata, + ResourceRequestResponseResources, +) +from zenml.orchestrators.step_launcher import StepLauncher + + +def _launcher() -> StepLauncher: + launcher = object.__new__(StepLauncher) + launcher._snapshot = MagicMock( + id=uuid4(), + is_dynamic=False, + pipeline_configuration=MagicMock(docker_settings=None), + ) + launcher._step = MagicMock(config=MagicMock(step_operator=None)) + launcher._stack = MagicMock(orchestrator=MagicMock()) + launcher._invocation_id = "step" + launcher._wait = False + return launcher + + +def _resource_request( + *, + status: ResourceRequestStatus, + lease_expires_at: datetime | None, +) -> ResourceRequestResponse: + request_id = uuid4() + return ResourceRequestResponse( + id=request_id, + name=str(request_id), + body=ResourceRequestResponseBody( + status=status, + user=uuid4(), + created=datetime.now(tz=timezone.utc), + updated=datetime.now(tz=timezone.utc), + version=1, + lease_expires_at=lease_expires_at, + reclaim_tolerance=ResourceRequestReclaimTolerance.NONE, + ), + metadata=ResourceRequestResponseMetadata(), + resources=ResourceRequestResponseResources(), + ) + + +def _step_run_info( + *, + resource_request_id, + resource_request: ResourceRequestResponse | None = None, + resource_settings: ResourceSettings | None = None, +) -> MagicMock: + if resource_request_id is None and resource_request is not None: + resource_request_id = resource_request.id + return MagicMock( + pipeline_step_name="step", + step_run_id=uuid4(), + config=MagicMock( + resource_settings=resource_settings or ResourceSettings(), + ), + step_run=MagicMock( + resource_request=resource_request, + resource_request_id=resource_request_id, + ), + ) + + +def test_wait_only_calls_renew(monkeypatch): + """Resource request polling renews the request instead of fetching it.""" + launcher = _launcher() + request_id = uuid4() + allocated = _resource_request( + status=ResourceRequestStatus.ALLOCATED, + lease_expires_at=datetime.now(tz=timezone.utc), + ) + zen_store = MagicMock( + renew_resource_request=MagicMock(return_value=allocated) + ) + monkeypatch.setattr( + "zenml.orchestrators.step_launcher.Client", + lambda: MagicMock(zen_store=zen_store), + ) + monkeypatch.setattr( + "zenml.orchestrators.step_launcher.publish_utils.publish_step_run_status_update", + MagicMock(), + ) + + result = launcher._wait_until_resources_acquired( + _step_run_info(resource_request_id=request_id) + ) + + assert result == allocated + assert zen_store.renew_resource_request.call_count == 2 + assert not hasattr(zen_store, "get_resource_request") or ( + not zen_store.get_resource_request.called + ) + polling_renewal = zen_store.renew_resource_request.call_args_list[0][0][1] + initialization_renewal = zen_store.renew_resource_request.call_args_list[ + 1 + ][0][1] + assert isinstance(polling_renewal, ResourceRequestRenewalRequest) + assert isinstance(initialization_renewal, ResourceRequestRenewalRequest) + + +@patch( + "zenml.orchestrators.step_launcher.utc_now", + return_value=datetime(2026, 1, 1, tzinfo=timezone.utc), +) +def test_wait_extends_initialization_lease_when_step_run_has_request( + _utc_now_mock, monkeypatch +): + """Hydrated allocated requests are renewed with the initialization lease.""" + launcher = _launcher() + allocated = _resource_request( + status=ResourceRequestStatus.ALLOCATED, + lease_expires_at=datetime.now(tz=timezone.utc), + ) + renewed = _resource_request( + status=ResourceRequestStatus.ALLOCATED, + lease_expires_at=datetime(2026, 1, 1, 1, 0, 0, tzinfo=timezone.utc), + ) + zen_store = MagicMock( + renew_resource_request=MagicMock(return_value=renewed) + ) + monkeypatch.setattr( + "zenml.orchestrators.step_launcher.Client", + lambda: MagicMock(zen_store=zen_store), + ) + monkeypatch.setattr( + "zenml.orchestrators.step_launcher.publish_utils.publish_step_run_status_update", + MagicMock(), + ) + + result = launcher._wait_until_resources_acquired( + _step_run_info( + resource_request=allocated, + resource_request_id=allocated.id, + resource_settings=ResourceSettings( + initialization_lease_seconds=3600 + ), + ), + ) + + assert result == renewed + renewal_request = zen_store.renew_resource_request.call_args[0][1] + assert renewal_request.lease_expires_at == datetime( + 2026, 1, 1, tzinfo=timezone.utc + ) + timedelta(seconds=3600) + + +@patch( + "zenml.orchestrators.step_launcher.utc_now", + return_value=datetime(2026, 1, 1, tzinfo=timezone.utc), +) +@patch( + "zenml.orchestrators.step_launcher.exponential_backoff_delays", + return_value=iter([3.0, 0.0]), +) +def test_wait_renew_lease_includes_backoff_delay( + _delays_mock, _utc_now_mock, monkeypatch +): + """Resource request lease renewal covers the next polling delay.""" + launcher = _launcher() + request_id = uuid4() + pending = _resource_request( + status=ResourceRequestStatus.PENDING, + lease_expires_at=datetime.now(tz=timezone.utc), + ) + allocated = _resource_request( + status=ResourceRequestStatus.ALLOCATED, + lease_expires_at=datetime.now(tz=timezone.utc), + ) + zen_store = MagicMock( + renew_resource_request=MagicMock( + side_effect=[pending, allocated, allocated] + ) + ) + monkeypatch.setattr( + "zenml.orchestrators.step_launcher.Client", + lambda: MagicMock(zen_store=zen_store), + ) + monkeypatch.setattr( + "zenml.orchestrators.step_launcher.time.sleep", MagicMock() + ) + monkeypatch.setattr( + "zenml.orchestrators.step_launcher.publish_utils.publish_step_run_status_update", + MagicMock(), + ) + + launcher._wait_until_resources_acquired( + _step_run_info(resource_request_id=request_id) + ) + + first_renewal = zen_store.renew_resource_request.call_args_list[0][0][1] + assert first_renewal.lease_expires_at == datetime( + 2026, 1, 1, tzinfo=timezone.utc + ) + timedelta(seconds=43.0) + + +@patch( + "zenml.orchestrators.step_launcher.utc_now", + return_value=datetime(2026, 1, 1, tzinfo=timezone.utc), +) +@patch( + "zenml.orchestrators.step_launcher.exponential_backoff_delays", + return_value=iter([3.0, 0.0]), +) +def test_wait_renew_lease_is_capped_by_allocation_timeout( + _delays_mock, _utc_now_mock, monkeypatch +): + """Pending resource request leases do not exceed the wait timeout.""" + launcher = _launcher() + request_id = uuid4() + pending = _resource_request( + status=ResourceRequestStatus.PENDING, + lease_expires_at=datetime.now(tz=timezone.utc), + ) + allocated = _resource_request( + status=ResourceRequestStatus.ALLOCATED, + lease_expires_at=datetime.now(tz=timezone.utc), + ) + zen_store = MagicMock( + renew_resource_request=MagicMock( + side_effect=[pending, allocated, allocated] + ) + ) + monkeypatch.setattr( + "zenml.orchestrators.step_launcher.Client", + lambda: MagicMock(zen_store=zen_store), + ) + monkeypatch.setattr( + "zenml.orchestrators.step_launcher.time.sleep", MagicMock() + ) + monkeypatch.setattr( + "zenml.orchestrators.step_launcher.publish_utils.publish_step_run_status_update", + MagicMock(), + ) + + launcher._wait_until_resources_acquired( + _step_run_info( + resource_request_id=request_id, + resource_settings=ResourceSettings( + allocation_wait_timeout_seconds=10 + ), + ) + ) + + first_renewal = zen_store.renew_resource_request.call_args_list[0][0][1] + assert first_renewal.lease_expires_at == datetime( + 2026, 1, 1, tzinfo=timezone.utc + ) + timedelta(seconds=10) + + +@patch( + "zenml.orchestrators.step_launcher.exponential_backoff_delays", + return_value=iter([0.0, 0.0]), +) +def test_wait_renews_each_poll(_delays_mock, monkeypatch): + """Resource request polling renews the lease on each store read.""" + launcher = _launcher() + request_id = uuid4() + pending = _resource_request( + status=ResourceRequestStatus.PENDING, + lease_expires_at=datetime.now(tz=timezone.utc), + ) + allocated = _resource_request( + status=ResourceRequestStatus.ALLOCATED, + lease_expires_at=datetime.now(tz=timezone.utc), + ) + zen_store = MagicMock( + renew_resource_request=MagicMock( + side_effect=[pending, allocated, allocated] + ) + ) + monkeypatch.setattr( + "zenml.orchestrators.step_launcher.Client", + lambda: MagicMock(zen_store=zen_store), + ) + sleep_mock = MagicMock() + monkeypatch.setattr( + "zenml.orchestrators.step_launcher.time.sleep", sleep_mock + ) + monkeypatch.setattr( + "zenml.orchestrators.step_launcher.publish_utils.publish_step_run_status_update", + MagicMock(), + ) + + launcher._wait_until_resources_acquired( + _step_run_info(resource_request_id=request_id) + ) + + assert zen_store.renew_resource_request.call_count == 3 + sleep_mock.assert_called_once_with(0.0) + + +@patch( + "zenml.orchestrators.step_launcher.exponential_backoff_delays", + return_value=iter([0.0, 0.0]), +) +def test_wait_renews_after_first_cycle_with_cached_request( + _delays_mock, monkeypatch +): + """Pending cached requests are renewed after the first polling cycle.""" + launcher = _launcher() + pending = _resource_request( + status=ResourceRequestStatus.PENDING, + lease_expires_at=datetime.now(tz=timezone.utc), + ) + allocated = _resource_request( + status=ResourceRequestStatus.ALLOCATED, + lease_expires_at=datetime.now(tz=timezone.utc), + ) + zen_store = MagicMock( + renew_resource_request=MagicMock(return_value=allocated) + ) + monkeypatch.setattr( + "zenml.orchestrators.step_launcher.Client", + lambda: MagicMock(zen_store=zen_store), + ) + sleep_mock = MagicMock() + monkeypatch.setattr( + "zenml.orchestrators.step_launcher.time.sleep", sleep_mock + ) + monkeypatch.setattr( + "zenml.orchestrators.step_launcher.publish_utils.publish_step_run_status_update", + MagicMock(), + ) + + launcher._wait_until_resources_acquired( + _step_run_info( + resource_request=pending, + resource_request_id=pending.id, + ), + ) + + assert zen_store.renew_resource_request.call_count == 2 + sleep_mock.assert_called_once_with(0.0) + + +@patch( + "zenml.orchestrators.step_launcher.exponential_backoff_delays", + return_value=iter([0.0, 0.0]), +) +def test_wait_raises_when_allocation_timeout_elapses( + _delays_mock, monkeypatch +): + """Resource request polling stops after the allocation wait timeout.""" + launcher = _launcher() + request_id = uuid4() + pending = _resource_request( + status=ResourceRequestStatus.PENDING, + lease_expires_at=datetime.now(tz=timezone.utc), + ) + zen_store = MagicMock( + renew_resource_request=MagicMock(return_value=pending) + ) + started_at = datetime(2026, 1, 1, tzinfo=timezone.utc) + monkeypatch.setattr( + "zenml.orchestrators.step_launcher.utc_now", + MagicMock( + side_effect=[ + started_at, + started_at + timedelta(seconds=30), + started_at + timedelta(seconds=60), + ] + ), + ) + monkeypatch.setattr( + "zenml.orchestrators.step_launcher.Client", + lambda: MagicMock(zen_store=zen_store), + ) + monkeypatch.setattr( + "zenml.orchestrators.step_launcher.time.sleep", MagicMock() + ) + + with pytest.raises(RuntimeError, match="Timed out after 60 seconds"): + launcher._wait_until_resources_acquired( + _step_run_info( + resource_request_id=request_id, + resource_settings=ResourceSettings( + allocation_wait_timeout_seconds=60 + ), + ), + ) + + +def test_dynamic_orchestrator_receives_allocated_resource_request( + monkeypatch, +): + """Dynamic orchestrator isolated-step submission receives allocations.""" + launcher = _launcher() + allocated = _resource_request( + status=ResourceRequestStatus.ALLOCATED, + lease_expires_at=datetime.now(tz=timezone.utc), + ) + submit_isolated_step_with_allocation = MagicMock() + launcher._stack = MagicMock( + orchestrator=MagicMock( + submit_isolated_step_with_allocation=( + submit_isolated_step_with_allocation + ), + ) + ) + launcher._wait = False + + monkeypatch.setattr( + "zenml.orchestrators.step_launcher.orchestrator_utils.get_config_environment_vars", + MagicMock(return_value=({}, {})), + ) + monkeypatch.setattr( + "zenml.orchestrators.step_launcher.env_utils.get_runtime_environment", + MagicMock(return_value={}), + ) + + launcher._run_step_with_dynamic_orchestrator( + step_run_info=MagicMock(run_id=uuid4(), config=MagicMock()), + allocated_resource_request=allocated, + ) + + assert ( + submit_isolated_step_with_allocation.call_args.kwargs[ + "allocated_resource_request" + ] + == allocated + ) + + +def test_step_operator_receives_allocated_resource_request(monkeypatch): + """Step operator submission receives the allocated resource request.""" + + class EntrypointConfiguration: + @staticmethod + def get_entrypoint_command(): + return ["python", "-m", "zenml"] + + @staticmethod + def get_entrypoint_arguments(step_name, snapshot_id, step_run_id): + return [step_name, str(snapshot_id), step_run_id] + + launcher = _launcher() + allocated = _resource_request( + status=ResourceRequestStatus.ALLOCATED, + lease_expires_at=datetime.now(tz=timezone.utc), + ) + step_operator = MagicMock( + name="step_operator", + entrypoint_config_class=EntrypointConfiguration, + ) + step_operator.submit_with_allocation = MagicMock() + launcher._stack = MagicMock( + step_operator=step_operator, + step_operators={}, + ) + launcher._snapshot = MagicMock(id=uuid4()) + launcher._invocation_id = "step" + launcher._wait = False + + monkeypatch.setattr( + "zenml.orchestrators.step_launcher.orchestrator_utils.get_config_environment_vars", + MagicMock(return_value=({}, {})), + ) + monkeypatch.setattr( + "zenml.orchestrators.step_launcher.env_utils.get_runtime_environment", + MagicMock(return_value={}), + ) + + launcher._run_step_with_step_operator( + step_operator_name=None, + step_run_info=MagicMock( + run_id=uuid4(), + config=MagicMock(), + step_run_id=uuid4(), + ), + allocated_resource_request=allocated, + ) + + assert ( + step_operator.submit_with_allocation.call_args.kwargs[ + "allocated_resource_request" + ] + == allocated + ) diff --git a/tests/unit/stack/test_stack.py b/tests/unit/stack/test_stack.py index 8df70dd40f7..1ef162b6020 100644 --- a/tests/unit/stack/test_stack.py +++ b/tests/unit/stack/test_stack.py @@ -11,8 +11,11 @@ # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express # or implied. See the License for the specific language governing # permissions and limitations under the License. +"""Tests for stack behavior.""" + from contextlib import ExitStack as does_not_raise -from uuid import uuid4 +from datetime import datetime, timezone +from uuid import UUID, uuid4 import pytest @@ -21,9 +24,70 @@ from zenml.config.build_configuration import BuildConfiguration from zenml.config.compiler import Compiler from zenml.config.pipeline_run_configuration import PipelineRunConfiguration -from zenml.enums import StackComponentType +from zenml.enums import ( + ResourceRequestReclaimTolerance, + ResourceRequestStatus, + StackComponentType, +) from zenml.exceptions import StackValidationError +from zenml.models import ( + ResourceRequestResponse, + ResourceRequestResponseBody, + ResourceRequestResponseMetadata, + ResourceRequestResponseResources, + ResourceRequestServiceConnectorSettings, +) from zenml.stack import Stack +from zenml.step_operators import BaseStepOperator, BaseStepOperatorConfig + + +def _step_operator( + *, + name: str = "step-operator", + connector: UUID | None = None, + connector_resource_id: str | None = None, +) -> BaseStepOperator: + """Create a test step operator.""" + return BaseStepOperator( + name=name, + id=uuid4(), + config=BaseStepOperatorConfig(), + flavor="test", + type=StackComponentType.STEP_OPERATOR, + user=uuid4(), + created=datetime.now(tz=timezone.utc), + updated=datetime.now(tz=timezone.utc), + connector=connector, + connector_resource_id=connector_resource_id, + ) + + +def _allocated_resource_request( + *, + connector_id: UUID, + connector_resource_id: str, +) -> ResourceRequestResponse: + """Create an allocated resource request with connector settings.""" + request_id = uuid4() + return ResourceRequestResponse( + id=request_id, + name=str(request_id), + body=ResourceRequestResponseBody( + status=ResourceRequestStatus.ALLOCATED, + user=uuid4(), + created=datetime.now(tz=timezone.utc), + updated=datetime.now(tz=timezone.utc), + version=1, + reclaim_tolerance=ResourceRequestReclaimTolerance.NONE, + ), + metadata=ResourceRequestResponseMetadata(), + resources=ResourceRequestResponseResources( + service_connector_settings=ResourceRequestServiceConnectorSettings( + connector_id=connector_id, + resource_id=connector_resource_id, + ) + ), + ) def test_initializing_a_stack_from_components( @@ -113,6 +177,146 @@ def test_stack_returns_all_its_components( ) +def test_get_stack_component_returns_original_without_connector_override( + local_orchestrator, local_artifact_store +): + """Tests that unchanged service connector settings reuse the stack component.""" + connector_id = uuid4() + step_operator = _step_operator( + connector=connector_id, + connector_resource_id="cluster-a", + ) + stack = Stack( + id=uuid4(), + name="stack", + orchestrator=local_orchestrator, + artifact_store=local_artifact_store, + step_operator=step_operator, + ) + + assert ( + stack.get_stack_component( + StackComponentType.STEP_OPERATOR, + ) + is step_operator + ) + assert ( + stack.get_stack_component( + StackComponentType.STEP_OPERATOR, + allocated_resource_request=_allocated_resource_request( + connector_id=connector_id, + connector_resource_id="cluster-a", + ), + ) + is step_operator + ) + + +def test_get_stack_component_caches_service_connector_overrides( + local_orchestrator, local_artifact_store +): + """Tests that service connector overrides are copied and cached.""" + original_connector_id = uuid4() + override_connector_id = uuid4() + step_operator = _step_operator( + connector=original_connector_id, + connector_resource_id="cluster-a", + ) + stack = Stack( + id=uuid4(), + name="stack", + orchestrator=local_orchestrator, + artifact_store=local_artifact_store, + step_operator=step_operator, + ) + + overridden_step_operator = stack.get_stack_component( + StackComponentType.STEP_OPERATOR, + name=step_operator.name, + allocated_resource_request=_allocated_resource_request( + connector_id=override_connector_id, + connector_resource_id="cluster-b", + ), + ) + cached_step_operator = stack.get_stack_component( + StackComponentType.STEP_OPERATOR, + name=step_operator.name, + allocated_resource_request=_allocated_resource_request( + connector_id=override_connector_id, + connector_resource_id="cluster-b", + ), + ) + + assert overridden_step_operator is cached_step_operator + assert overridden_step_operator is not step_operator + assert overridden_step_operator is not None + assert overridden_step_operator.connector == override_connector_id + assert overridden_step_operator.connector_resource_id == "cluster-b" + assert step_operator.connector == original_connector_id + assert step_operator.connector_resource_id == "cluster-a" + + +def test_get_step_operator_uses_allocated_request_connector_overrides( + local_orchestrator, local_artifact_store +): + """Tests that request connector settings override step operator links.""" + original_connector_id = uuid4() + override_connector_id = uuid4() + step_operator = _step_operator( + name="gpu-operator", + connector=original_connector_id, + connector_resource_id="cluster-a", + ) + stack = Stack( + id=uuid4(), + name="stack", + orchestrator=local_orchestrator, + artifact_store=local_artifact_store, + step_operator=step_operator, + ) + + overridden_step_operator = stack.get_step_operator( + name="gpu-operator", + allocated_resource_request=_allocated_resource_request( + connector_id=override_connector_id, + connector_resource_id="cluster-b", + ), + ) + + assert overridden_step_operator.name == "gpu-operator" + assert overridden_step_operator.connector == override_connector_id + assert overridden_step_operator.connector_resource_id == "cluster-b" + assert step_operator.connector == original_connector_id + assert step_operator.connector_resource_id == "cluster-a" + + +def test_get_step_operator_raises_for_missing_step_operator( + local_orchestrator, local_artifact_store +): + """Tests that missing step operators raise helpful errors.""" + stack = Stack( + id=uuid4(), + name="stack", + orchestrator=local_orchestrator, + artifact_store=local_artifact_store, + ) + + with pytest.raises(RuntimeError, match="No step operators specified"): + stack.get_step_operator() + + step_operator = _step_operator(name="available") + stack = Stack( + id=uuid4(), + name="stack", + orchestrator=local_orchestrator, + artifact_store=local_artifact_store, + step_operator=step_operator, + ) + + with pytest.raises(RuntimeError, match="No step operator named 'missing'"): + stack.get_step_operator(name="missing") + + def test_stack_requirements(stack_with_mock_components): """Tests that the stack returns the requirements of all its components.""" stack_with_mock_components.orchestrator.requirements = {"one_requirement"} @@ -154,9 +358,12 @@ def test_stack_submission( stack_with_mock_components, empty_pipeline, # noqa: F811 ): - """Tests that when a pipeline is deployed on a stack, the stack calls the + """Tests that stack submission delegates to the orchestrator. + + When a pipeline is deployed on a stack, the stack calls the orchestrator to run the pipeline and calls cleanup methods on all of its - components.""" + components. + """ # Mock the pipeline run registering which tries (and fails) to serialize # our mock objects empty_pipeline.prepare() @@ -208,7 +415,6 @@ def test_submission_server_validation( mocker, stack_with_mock_components, sample_snapshot_response_model ): """Tests that the submission validation fails when the stack requires a remote server but the store is local.""" - ######### Remote server ######### mocker.patch( "zenml.zen_stores.base_zen_store.BaseZenStore.is_local_store", @@ -400,8 +606,10 @@ class MockStepInfo: def test_docker_builds_collection( stack_with_mock_components, sample_snapshot_response_model ): - """Tests that the stack collects the required Docker builds from all its - components.""" + """Tests that the stack collects the required Docker builds. + + The stack collects builds from all its components. + """ first_orchestrator_build = BuildConfiguration( key="orchestrator", settings=DockerSettings() ) diff --git a/tests/unit/steps/test_heartbeat_worker.py b/tests/unit/steps/test_heartbeat_worker.py index 4d281e5e722..4a9f5c10683 100644 --- a/tests/unit/steps/test_heartbeat_worker.py +++ b/tests/unit/steps/test_heartbeat_worker.py @@ -30,8 +30,16 @@ def __init__(self, step_id): self.heartbeat_enabled = True self._call_count = 0 self.step_id = step_id - - def update_step_heartbeat(self, step_run_id) -> StepHeartbeatResponse: + self.heartbeat_liveness_timeout_seconds = None + + def update_step_heartbeat( + self, + step_run_id, + heartbeat_liveness_timeout_seconds=None, + ) -> StepHeartbeatResponse: + self.heartbeat_liveness_timeout_seconds = ( + heartbeat_liveness_timeout_seconds + ) if self.step_id == step_run_id: self._call_count += 1 @@ -85,6 +93,10 @@ def mock_interrupt_main(): time.sleep(0.1) assert fake_server._call_count >= 1 + assert ( + fake_server.heartbeat_liveness_timeout_seconds + == StepHeartbeatWorker.heartbeat_liveness_timeout_seconds() + ) fake_server.execution_status = ExecutionStatus.STOPPED @@ -92,9 +104,17 @@ def mock_interrupt_main(): "_thread.interrupt_main was never called" ) - assert not worker.is_running + assert worker.is_running assert worker.is_terminated + calls_after_interrupt = fake_server._call_count + time.sleep(0.25) + assert fake_server._call_count > calls_after_interrupt + + worker.stop() + time.sleep(0.2) + assert not worker.is_running + def test_heartbeat_worker_with_eager_stopping(monkeypatch): step_run_id = uuid.uuid4() diff --git a/tests/unit/zen_stores/test_resource_request_pool_lifecycle.py b/tests/unit/zen_stores/test_resource_request_pool_lifecycle.py deleted file mode 100644 index 08d0dffd68b..00000000000 --- a/tests/unit/zen_stores/test_resource_request_pool_lifecycle.py +++ /dev/null @@ -1,2020 +0,0 @@ -from copy import deepcopy -from datetime import datetime -from typing import Dict, List, Tuple -from uuid import UUID, uuid4 - -import pytest -from sqlmodel import Session, select - -from zenml.config.step_configurations import Step -from zenml.enums import ( - ExecutionStatus, - ResourceRequestStatus, - StackComponentType, -) -from zenml.exceptions import IllegalOperationError -from zenml.models import ( - PipelineRequest, - ResourcePoolRequest, - ResourcePoolSubjectPolicyRequest, - ResourcePoolUpdate, - ResourceRequestRequest, -) -from zenml.zen_stores.schemas.resource_pool_schemas import ( - ResourcePoolQueueSchema, - ResourcePoolResourceSchema, - ResourceRequestSchema, -) -from zenml.zen_stores.sql_zen_store import SqlZenStore - -pytest.skip( - "Resource pool lifecycle tests are disabled.", allow_module_level=True -) - - -def _create_step_run_in_db( - clean_client, - sample_pipeline_snapshot_request_model, - sample_pipeline_run_request_model, - sample_step_request_model, - *, - step_name: str = "sample_step", -) -> Tuple[UUID, UUID]: - """Create a pipeline run and one step run in the DB. - - Args: - clean_client: Test client fixture. - sample_pipeline_snapshot_request_model: Snapshot request fixture. - sample_pipeline_run_request_model: Pipeline run request fixture. - sample_step_request_model: Step run request fixture. - step_name: Name of the created step. - - Returns: - A tuple of `(pipeline_run_id, step_run_id)`. - """ - store = clean_client.zen_store - pipeline = store.create_pipeline( - PipelineRequest( - name=f"test-pipeline-{uuid4()}", - project=clean_client.active_project.id, - ) - ) - - snapshot_request = deepcopy(sample_pipeline_snapshot_request_model) - snapshot_request.project = clean_client.active_project.id - snapshot_request.stack = clean_client.active_stack.id - snapshot_request.pipeline = pipeline.id - snapshot_request.step_configurations = { - step_name: Step.model_validate( - { - "spec": { - "source": "module.step_class", - "upstream_steps": [], - "inputs": {}, - }, - "config": {"name": step_name}, - } - ) - } - snapshot = store.create_snapshot(snapshot_request) - - run_request = deepcopy(sample_pipeline_run_request_model) - run_request.project = clean_client.active_project.id - run_request.snapshot = snapshot.id - run_request.name = f"test-run-{uuid4()}" - run_request.status = ExecutionStatus.RUNNING - run, created = store.get_or_create_run(run_request) - assert created - - step_request = deepcopy(sample_step_request_model) - step_request.project = clean_client.active_project.id - step_request.pipeline_run_id = run.id - step_request.name = step_name - step_request.status = ExecutionStatus.RUNNING - step_request.start_time = datetime.utcnow() - step = store.create_run_step(step_request) - return run.id, step.id - - -def _create_orchestrator_component(clean_client, name: str) -> UUID: - """Create an orchestrator component and return its ID. - - Args: - clean_client: Test client fixture. - name: Prefix for the component name. - - Returns: - The created component ID. - """ - component = clean_client.create_stack_component( - name=f"{name}-{uuid4()}", - flavor="local", - component_type=StackComponentType.ORCHESTRATOR, - configuration={}, - ) - return component.id - - -def _create_pool( - clean_client, - *, - capacity: Dict[str, int], - policies: List[ResourcePoolSubjectPolicyRequest], -): - """Create a resource pool with policies. - - Args: - clean_client: Test client fixture. - capacity: Pool capacity per resource key. - policies: Policies to attach to the pool. - - Returns: - The created pool model. - """ - return clean_client.zen_store.create_resource_pool( - ResourcePoolRequest( - name=f"pool-{uuid4()}", - capacity=capacity, - policies=policies, - ) - ) - - -def _create_resource_request( - clean_client, - *, - component_id: UUID, - step_run_id: UUID, - requested_resources: Dict[str, int], - preemptible: bool = True, -): - """Create a resource request. - - Args: - clean_client: Test client fixture. - component_id: Component requesting resources. - step_run_id: Step run owning the request. - requested_resources: Requested resources per key. - preemptible: Whether the request is preemptible. - - Returns: - The created resource request model. - """ - return clean_client.zen_store.create_resource_request( - ResourceRequestRequest( - component_id=component_id, - step_run_id=step_run_id, - requested_resources=requested_resources, - preemptible=preemptible, - ) - ) - - -def _get_request(clean_client, request_id: UUID): - return clean_client.zen_store.get_resource_request( - request_id, hydrate=False - ) - - -def _get_pool(clean_client, pool_id: UUID): - return clean_client.zen_store.get_resource_pool(pool_id, hydrate=True) - - -def test_request_rejected_if_exceeds_pool_total_capacity( - clean_client, - sample_pipeline_snapshot_request_model, - sample_pipeline_run_request_model, - sample_step_request_model, -): - component_id = clean_client.active_stack.orchestrator.id - _create_pool( - clean_client, - capacity={"gpu": 1}, - policies=[ - ResourcePoolSubjectPolicyRequest( - component_id=component_id, - priority=1, - reserved={"gpu": 1}, - limit={"gpu": 1}, - ) - ], - ) - - _, step_id = _create_step_run_in_db( - clean_client, - sample_pipeline_snapshot_request_model, - sample_pipeline_run_request_model, - sample_step_request_model, - ) - - request = _create_resource_request( - clean_client, - component_id=component_id, - step_run_id=step_id, - requested_resources={"gpu": 2}, - ) - assert ( - _get_request(clean_client, request.id).status - == ResourceRequestStatus.REJECTED - ) - - -def test_request_rejected_if_exceeds_policy_limit( - clean_client, - sample_pipeline_snapshot_request_model, - sample_pipeline_run_request_model, - sample_step_request_model, -): - component_id = clean_client.active_stack.orchestrator.id - _create_pool( - clean_client, - capacity={"gpu": 2}, - policies=[ - ResourcePoolSubjectPolicyRequest( - component_id=component_id, - priority=1, - reserved={"gpu": 1}, - limit={"gpu": 1}, - ) - ], - ) - - _, step_id = _create_step_run_in_db( - clean_client, - sample_pipeline_snapshot_request_model, - sample_pipeline_run_request_model, - sample_step_request_model, - ) - - request = _create_resource_request( - clean_client, - component_id=component_id, - step_run_id=step_id, - requested_resources={"gpu": 2}, - ) - assert ( - _get_request(clean_client, request.id).status - == ResourceRequestStatus.REJECTED - ) - - -def test_non_preemptible_request_rejected_if_exceeds_reserved_share( - clean_client, - sample_pipeline_snapshot_request_model, - sample_pipeline_run_request_model, - sample_step_request_model, -): - component_id = clean_client.active_stack.orchestrator.id - _create_pool( - clean_client, - capacity={"gpu": 2}, - policies=[ - ResourcePoolSubjectPolicyRequest( - component_id=component_id, - priority=1, - reserved={"gpu": 1}, - limit={"gpu": 2}, - ) - ], - ) - - _, step_id = _create_step_run_in_db( - clean_client, - sample_pipeline_snapshot_request_model, - sample_pipeline_run_request_model, - sample_step_request_model, - ) - - request = _create_resource_request( - clean_client, - component_id=component_id, - step_run_id=step_id, - requested_resources={"gpu": 2}, - preemptible=False, - ) - assert ( - _get_request(clean_client, request.id).status - == ResourceRequestStatus.REJECTED - ) - - -def test_request_allocated_when_equal_pool_total_capacity( - clean_client, - sample_pipeline_snapshot_request_model, - sample_pipeline_run_request_model, - sample_step_request_model, -): - component_id = clean_client.active_stack.orchestrator.id - _create_pool( - clean_client, - capacity={"gpu": 2}, - policies=[ - ResourcePoolSubjectPolicyRequest( - component_id=component_id, - priority=1, - reserved={"gpu": 2}, - limit={"gpu": 2}, - ) - ], - ) - - _, step_id = _create_step_run_in_db( - clean_client, - sample_pipeline_snapshot_request_model, - sample_pipeline_run_request_model, - sample_step_request_model, - ) - request = _create_resource_request( - clean_client, - component_id=component_id, - step_run_id=step_id, - requested_resources={"gpu": 2}, - ) - assert ( - _get_request(clean_client, request.id).status - == ResourceRequestStatus.ALLOCATED - ) - - -def test_request_allocated_when_equal_policy_limit( - clean_client, - sample_pipeline_snapshot_request_model, - sample_pipeline_run_request_model, - sample_step_request_model, -): - component_id = clean_client.active_stack.orchestrator.id - _create_pool( - clean_client, - capacity={"gpu": 3}, - policies=[ - ResourcePoolSubjectPolicyRequest( - component_id=component_id, - priority=1, - reserved={"gpu": 1}, - limit={"gpu": 2}, - ) - ], - ) - - _, step_id = _create_step_run_in_db( - clean_client, - sample_pipeline_snapshot_request_model, - sample_pipeline_run_request_model, - sample_step_request_model, - ) - request = _create_resource_request( - clean_client, - component_id=component_id, - step_run_id=step_id, - requested_resources={"gpu": 2}, - ) - assert ( - _get_request(clean_client, request.id).status - == ResourceRequestStatus.ALLOCATED - ) - - -def test_non_preemptible_request_allocated_when_equal_reserved_share( - clean_client, - sample_pipeline_snapshot_request_model, - sample_pipeline_run_request_model, - sample_step_request_model, -): - component_id = clean_client.active_stack.orchestrator.id - _create_pool( - clean_client, - capacity={"gpu": 3}, - policies=[ - ResourcePoolSubjectPolicyRequest( - component_id=component_id, - priority=1, - reserved={"gpu": 2}, - limit={"gpu": 3}, - ) - ], - ) - - _, step_id = _create_step_run_in_db( - clean_client, - sample_pipeline_snapshot_request_model, - sample_pipeline_run_request_model, - sample_step_request_model, - ) - request = _create_resource_request( - clean_client, - component_id=component_id, - step_run_id=step_id, - requested_resources={"gpu": 2}, - preemptible=False, - ) - assert ( - _get_request(clean_client, request.id).status - == ResourceRequestStatus.ALLOCATED - ) - - -def test_missing_unbounded_pool_key_allows_allocation( - clean_client, - sample_pipeline_snapshot_request_model, - sample_pipeline_run_request_model, - sample_step_request_model, -): - component_id = clean_client.active_stack.orchestrator.id - _create_pool( - clean_client, - capacity={"gpu": 1}, - policies=[ - ResourcePoolSubjectPolicyRequest( - component_id=component_id, - priority=1, - reserved={"gpu": 0}, - limit={"gpu": 1}, - ) - ], - ) - - _, cpu_step = _create_step_run_in_db( - clean_client, - sample_pipeline_snapshot_request_model, - sample_pipeline_run_request_model, - sample_step_request_model, - ) - cpu_request = _create_resource_request( - clean_client, - component_id=component_id, - step_run_id=cpu_step, - requested_resources={"cpu": 100}, - ) - - assert ( - _get_request(clean_client, cpu_request.id).status - == ResourceRequestStatus.ALLOCATED - ) - - -def test_missing_bounded_pool_key_rejects_request( - clean_client, - sample_pipeline_snapshot_request_model, - sample_pipeline_run_request_model, - sample_step_request_model, -): - component_id = clean_client.active_stack.orchestrator.id - _create_pool( - clean_client, - capacity={"gpu": 1}, - policies=[ - ResourcePoolSubjectPolicyRequest( - component_id=component_id, - priority=1, - reserved={"gpu": 0}, - limit={"gpu": 1}, - ) - ], - ) - - _, disk_step = _create_step_run_in_db( - clean_client, - sample_pipeline_snapshot_request_model, - sample_pipeline_run_request_model, - sample_step_request_model, - ) - - disk_request = _create_resource_request( - clean_client, - component_id=component_id, - step_run_id=disk_step, - requested_resources={"disk": 1}, - ) - - assert ( - _get_request(clean_client, disk_request.id).status - == ResourceRequestStatus.REJECTED - ) - - -def test_missing_policy_key_inherits_bounded_pool_capacity( - clean_client, - sample_pipeline_snapshot_request_model, - sample_pipeline_run_request_model, - sample_step_request_model, -): - component_id = clean_client.active_stack.orchestrator.id - _create_pool( - clean_client, - capacity={"gpu": 2}, - policies=[ - ResourcePoolSubjectPolicyRequest( - component_id=component_id, - priority=1, - reserved={}, - limit={}, - ) - ], - ) - - _, step_fit = _create_step_run_in_db( - clean_client, - sample_pipeline_snapshot_request_model, - sample_pipeline_run_request_model, - sample_step_request_model, - ) - _, step_too_large = _create_step_run_in_db( - clean_client, - sample_pipeline_snapshot_request_model, - sample_pipeline_run_request_model, - sample_step_request_model, - ) - - fit_request = _create_resource_request( - clean_client, - component_id=component_id, - step_run_id=step_fit, - requested_resources={"gpu": 2}, - ) - too_large_request = _create_resource_request( - clean_client, - component_id=component_id, - step_run_id=step_too_large, - requested_resources={"gpu": 3}, - ) - - assert ( - _get_request(clean_client, fit_request.id).status - == ResourceRequestStatus.ALLOCATED - ) - assert ( - _get_request(clean_client, too_large_request.id).status - == ResourceRequestStatus.REJECTED - ) - - -def test_explicit_policy_limit_zero_overrides_pool_capacity( - clean_client, - sample_pipeline_snapshot_request_model, - sample_pipeline_run_request_model, - sample_step_request_model, -): - component_id = clean_client.active_stack.orchestrator.id - _create_pool( - clean_client, - capacity={"gpu": 2}, - policies=[ - ResourcePoolSubjectPolicyRequest( - component_id=component_id, - priority=1, - reserved={}, - limit={"gpu": 0}, - ) - ], - ) - - _, step_id = _create_step_run_in_db( - clean_client, - sample_pipeline_snapshot_request_model, - sample_pipeline_run_request_model, - sample_step_request_model, - ) - request = _create_resource_request( - clean_client, - component_id=component_id, - step_run_id=step_id, - requested_resources={"gpu": 1}, - ) - - assert ( - _get_request(clean_client, request.id).status - == ResourceRequestStatus.REJECTED - ) - - -def test_policy_without_key_can_use_new_pool_key_after_capacity_update( - clean_client, - sample_pipeline_snapshot_request_model, - sample_pipeline_run_request_model, - sample_step_request_model, -): - component_id = clean_client.active_stack.orchestrator.id - pool = _create_pool( - clean_client, - capacity={"gpu": 1}, - policies=[ - ResourcePoolSubjectPolicyRequest( - component_id=component_id, - priority=1, - reserved={}, - limit={}, - ) - ], - ) - - _, step_before = _create_step_run_in_db( - clean_client, - sample_pipeline_snapshot_request_model, - sample_pipeline_run_request_model, - sample_step_request_model, - ) - before_update = _create_resource_request( - clean_client, - component_id=component_id, - step_run_id=step_before, - requested_resources={"disk": 1}, - ) - assert ( - _get_request(clean_client, before_update.id).status - == ResourceRequestStatus.REJECTED - ) - - clean_client.zen_store.update_resource_pool( - pool.id, ResourcePoolUpdate(capacity={"gpu": 1, "disk": 2}) - ) - - _, step_after = _create_step_run_in_db( - clean_client, - sample_pipeline_snapshot_request_model, - sample_pipeline_run_request_model, - sample_step_request_model, - ) - after_update = _create_resource_request( - clean_client, - component_id=component_id, - step_run_id=step_after, - requested_resources={"disk": 2}, - ) - assert ( - _get_request(clean_client, after_update.id).status - == ResourceRequestStatus.ALLOCATED - ) - - -def test_allocation_respects_non_preemptible_reserved_share( - clean_client, - sample_pipeline_snapshot_request_model, - sample_pipeline_run_request_model, - sample_step_request_model, -): - component_id = clean_client.active_stack.orchestrator.id - _, step_1 = _create_step_run_in_db( - clean_client, - sample_pipeline_snapshot_request_model, - sample_pipeline_run_request_model, - sample_step_request_model, - ) - _, step_2 = _create_step_run_in_db( - clean_client, - sample_pipeline_snapshot_request_model, - sample_pipeline_run_request_model, - sample_step_request_model, - ) - _, step_3 = _create_step_run_in_db( - clean_client, - sample_pipeline_snapshot_request_model, - sample_pipeline_run_request_model, - sample_step_request_model, - ) - - pool = _create_pool( - clean_client, - capacity={"gpu": 2}, - policies=[ - ResourcePoolSubjectPolicyRequest( - component_id=component_id, - priority=1, - reserved={"gpu": 1}, - limit={"gpu": 2}, - ) - ], - ) - - req_1 = _create_resource_request( - clean_client, - component_id=component_id, - step_run_id=step_1, - requested_resources={"gpu": 1}, - preemptible=False, - ) - req_2 = _create_resource_request( - clean_client, - component_id=component_id, - step_run_id=step_2, - requested_resources={"gpu": 1}, - preemptible=False, - ) - req_3 = _create_resource_request( - clean_client, - component_id=component_id, - step_run_id=step_3, - requested_resources={"gpu": 1}, - preemptible=True, - ) - - assert ( - _get_request(clean_client, req_1.id).status - == ResourceRequestStatus.ALLOCATED - ) - assert ( - _get_request(clean_client, req_2.id).status - == ResourceRequestStatus.PENDING - ) - assert ( - _get_request(clean_client, req_3.id).status - == ResourceRequestStatus.ALLOCATED - ) - - hydrated_pool = _get_pool(clean_client, pool.id) - active_ids = {item.request.id for item in hydrated_pool.active_requests} - queued_ids = {item.request.id for item in hydrated_pool.queued_requests} - assert req_1.id in active_ids - assert req_3.id in active_ids - assert req_2.id in queued_ids - - -def test_preemption_skips_non_preemptible_victims( - clean_client, - sample_pipeline_snapshot_request_model, - sample_pipeline_run_request_model, - sample_step_request_model, - monkeypatch, -): - store = clean_client.zen_store - assert isinstance(store, SqlZenStore) - low_priority_component = _create_orchestrator_component( - clean_client, "low" - ) - high_priority_component = _create_orchestrator_component( - clean_client, "high" - ) - - pool = _create_pool( - clean_client, - capacity={"gpu": 1}, - policies=[ - ResourcePoolSubjectPolicyRequest( - component_id=low_priority_component, - priority=1, - reserved={"gpu": 1}, - limit={"gpu": 1}, - ), - ResourcePoolSubjectPolicyRequest( - component_id=high_priority_component, - priority=2, - reserved={"gpu": 0}, - limit={"gpu": 1}, - ), - ], - ) - - _, low_step = _create_step_run_in_db( - clean_client, - sample_pipeline_snapshot_request_model, - sample_pipeline_run_request_model, - sample_step_request_model, - ) - _, high_step = _create_step_run_in_db( - clean_client, - sample_pipeline_snapshot_request_model, - sample_pipeline_run_request_model, - sample_step_request_model, - ) - - victim = _create_resource_request( - clean_client, - component_id=low_priority_component, - step_run_id=low_step, - requested_resources={"gpu": 1}, - preemptible=False, - ) - - evictions: List[UUID] = [] - - def _mock_trigger(self, session, request_id): - evictions.append(request_id) - - monkeypatch.setattr( - SqlZenStore, "_trigger_resource_request_eviction", _mock_trigger - ) - - head = _create_resource_request( - clean_client, - component_id=high_priority_component, - step_run_id=high_step, - requested_resources={"gpu": 1}, - preemptible=True, - ) - - assert ( - _get_request(clean_client, victim.id).status - == ResourceRequestStatus.ALLOCATED - ) - assert ( - _get_request(clean_client, head.id).status - == ResourceRequestStatus.PENDING - ) - assert evictions == [] - assert _get_pool(clean_client, pool.id).queued_requests - - -def test_preemption_triggers_for_preemptible_victims( - clean_client, - sample_pipeline_snapshot_request_model, - sample_pipeline_run_request_model, - sample_step_request_model, - monkeypatch, -): - store = clean_client.zen_store - assert isinstance(store, SqlZenStore) - low_priority_component = _create_orchestrator_component( - clean_client, "low" - ) - high_priority_component = _create_orchestrator_component( - clean_client, "high" - ) - - _create_pool( - clean_client, - capacity={"gpu": 1}, - policies=[ - ResourcePoolSubjectPolicyRequest( - component_id=low_priority_component, - priority=1, - reserved={"gpu": 0}, - limit={"gpu": 1}, - ), - ResourcePoolSubjectPolicyRequest( - component_id=high_priority_component, - priority=2, - reserved={"gpu": 0}, - limit={"gpu": 1}, - ), - ], - ) - - _, low_step = _create_step_run_in_db( - clean_client, - sample_pipeline_snapshot_request_model, - sample_pipeline_run_request_model, - sample_step_request_model, - ) - _, high_step = _create_step_run_in_db( - clean_client, - sample_pipeline_snapshot_request_model, - sample_pipeline_run_request_model, - sample_step_request_model, - ) - - victim = _create_resource_request( - clean_client, - component_id=low_priority_component, - step_run_id=low_step, - requested_resources={"gpu": 1}, - preemptible=True, - ) - - evictions: List[UUID] = [] - - def _mock_trigger(self, session, request_id): - evictions.append(request_id) - - monkeypatch.setattr( - SqlZenStore, "_trigger_resource_request_eviction", _mock_trigger - ) - - _create_resource_request( - clean_client, - component_id=high_priority_component, - step_run_id=high_step, - requested_resources={"gpu": 1}, - preemptible=True, - ) - - assert evictions == [victim.id] - assert ( - _get_request(clean_client, victim.id).status - == ResourceRequestStatus.PREEMPTING - ) - - -def test_orphaned_requests_are_cancelled_before_allocation( - clean_client, - sample_pipeline_snapshot_request_model, - sample_pipeline_run_request_model, - sample_step_request_model, -): - component_id = clean_client.active_stack.orchestrator.id - pool = _create_pool( - clean_client, - capacity={"gpu": 1}, - policies=[ - ResourcePoolSubjectPolicyRequest( - component_id=component_id, - priority=1, - reserved={"gpu": 1}, - limit={"gpu": 1}, - ) - ], - ) - - run_1, step_1 = _create_step_run_in_db( - clean_client, - sample_pipeline_snapshot_request_model, - sample_pipeline_run_request_model, - sample_step_request_model, - ) - orphan_candidate = _create_resource_request( - clean_client, - component_id=component_id, - step_run_id=step_1, - requested_resources={"gpu": 1}, - preemptible=True, - ) - assert ( - _get_request(clean_client, orphan_candidate.id).status - == ResourceRequestStatus.ALLOCATED - ) - - clean_client.zen_store.delete_run(run_1) - - _, step_2 = _create_step_run_in_db( - clean_client, - sample_pipeline_snapshot_request_model, - sample_pipeline_run_request_model, - sample_step_request_model, - ) - next_request = _create_resource_request( - clean_client, - component_id=component_id, - step_run_id=step_2, - requested_resources={"gpu": 1}, - preemptible=True, - ) - - orphan = _get_request(clean_client, orphan_candidate.id) - assert orphan.status == ResourceRequestStatus.CANCELLED - assert ( - orphan.status_reason - == "Cancelled because owning step run no longer exists." - ) - assert ( - _get_request(clean_client, next_request.id).status - == ResourceRequestStatus.ALLOCATED - ) - assert _get_pool(clean_client, pool.id).active_requests - - -def test_attaching_policy_allocates_pending_requests_from_other_pool( - clean_client, - sample_pipeline_snapshot_request_model, - sample_pipeline_run_request_model, - sample_step_request_model, -): - component_id = clean_client.active_stack.orchestrator.id - pool_a = _create_pool( - clean_client, - capacity={"gpu": 1}, - policies=[ - ResourcePoolSubjectPolicyRequest( - component_id=component_id, - priority=1, - reserved={"gpu": 1}, - limit={"gpu": 1}, - ) - ], - ) - pool_b = _create_pool(clean_client, capacity={"gpu": 1}, policies=[]) - - _, step_1 = _create_step_run_in_db( - clean_client, - sample_pipeline_snapshot_request_model, - sample_pipeline_run_request_model, - sample_step_request_model, - ) - _, step_2 = _create_step_run_in_db( - clean_client, - sample_pipeline_snapshot_request_model, - sample_pipeline_run_request_model, - sample_step_request_model, - ) - - _create_resource_request( - clean_client, - component_id=component_id, - step_run_id=step_1, - requested_resources={"gpu": 1}, - ) - pending = _create_resource_request( - clean_client, - component_id=component_id, - step_run_id=step_2, - requested_resources={"gpu": 1}, - ) - assert ( - _get_request(clean_client, pending.id).status - == ResourceRequestStatus.PENDING - ) - - clean_client.zen_store.update_resource_pool( - pool_b.id, - ResourcePoolUpdate( - attach_policies=[ - ResourcePoolSubjectPolicyRequest( - component_id=component_id, - priority=5, - reserved={"gpu": 1}, - limit={"gpu": 1}, - ) - ] - ), - ) - - pending_after = _get_request(clean_client, pending.id) - assert pending_after.status == ResourceRequestStatus.ALLOCATED - assert pending_after.running_in_pool is not None - assert pending_after.running_in_pool.id == pool_b.id - assert _get_pool(clean_client, pool_a.id).queued_requests == [] - - -def test_detaching_policy_blocked_if_component_has_active_requests( - clean_client, - sample_pipeline_snapshot_request_model, - sample_pipeline_run_request_model, - sample_step_request_model, -): - component_id = clean_client.active_stack.orchestrator.id - pool = _create_pool( - clean_client, - capacity={"gpu": 1}, - policies=[ - ResourcePoolSubjectPolicyRequest( - component_id=component_id, - priority=1, - reserved={"gpu": 1}, - limit={"gpu": 1}, - ) - ], - ) - - _, step_1 = _create_step_run_in_db( - clean_client, - sample_pipeline_snapshot_request_model, - sample_pipeline_run_request_model, - sample_step_request_model, - ) - _, step_2 = _create_step_run_in_db( - clean_client, - sample_pipeline_snapshot_request_model, - sample_pipeline_run_request_model, - sample_step_request_model, - ) - - allocated = _create_resource_request( - clean_client, - component_id=component_id, - step_run_id=step_1, - requested_resources={"gpu": 1}, - ) - queued = _create_resource_request( - clean_client, - component_id=component_id, - step_run_id=step_2, - requested_resources={"gpu": 1}, - ) - assert ( - _get_request(clean_client, queued.id).status - == ResourceRequestStatus.PENDING - ) - - with pytest.raises(IllegalOperationError): - clean_client.zen_store.update_resource_pool( - pool.id, - ResourcePoolUpdate(detach_policies=[component_id]), - ) - - assert ( - _get_request(clean_client, allocated.id).status - == ResourceRequestStatus.ALLOCATED - ) - assert ( - _get_request(clean_client, queued.id).status - == ResourceRequestStatus.PENDING - ) - - -def test_decreasing_pool_capacity_keeps_allocated_and_rebuilds_queue( - clean_client, - sample_pipeline_snapshot_request_model, - sample_pipeline_run_request_model, - sample_step_request_model, -): - component_id = clean_client.active_stack.orchestrator.id - pool = _create_pool( - clean_client, - capacity={"gpu": 2}, - policies=[ - ResourcePoolSubjectPolicyRequest( - component_id=component_id, - priority=1, - reserved={"gpu": 0}, - limit={"gpu": 2}, - ) - ], - ) - - _, step_1 = _create_step_run_in_db( - clean_client, - sample_pipeline_snapshot_request_model, - sample_pipeline_run_request_model, - sample_step_request_model, - ) - _, step_2 = _create_step_run_in_db( - clean_client, - sample_pipeline_snapshot_request_model, - sample_pipeline_run_request_model, - sample_step_request_model, - ) - _, step_3 = _create_step_run_in_db( - clean_client, - sample_pipeline_snapshot_request_model, - sample_pipeline_run_request_model, - sample_step_request_model, - ) - - allocated = _create_resource_request( - clean_client, - component_id=component_id, - step_run_id=step_1, - requested_resources={"gpu": 2}, - ) - pending_too_large = _create_resource_request( - clean_client, - component_id=component_id, - step_run_id=step_2, - requested_resources={"gpu": 2}, - ) - pending_small = _create_resource_request( - clean_client, - component_id=component_id, - step_run_id=step_3, - requested_resources={"gpu": 1}, - ) - assert ( - _get_request(clean_client, pending_too_large.id).status - == ResourceRequestStatus.PENDING - ) - assert ( - _get_request(clean_client, pending_small.id).status - == ResourceRequestStatus.PENDING - ) - - clean_client.zen_store.update_resource_pool( - pool.id, ResourcePoolUpdate(capacity={"gpu": 1}) - ) - - assert ( - _get_request(clean_client, allocated.id).status - == ResourceRequestStatus.ALLOCATED - ) - assert ( - _get_request(clean_client, pending_too_large.id).status - == ResourceRequestStatus.REJECTED - ) - assert ( - _get_request(clean_client, pending_small.id).status - == ResourceRequestStatus.PENDING - ) - - -def test_deleting_pool_blocked_if_pool_has_active_requests( - clean_client, - sample_pipeline_snapshot_request_model, - sample_pipeline_run_request_model, - sample_step_request_model, -): - component_id = clean_client.active_stack.orchestrator.id - pool = _create_pool( - clean_client, - capacity={"gpu": 1}, - policies=[ - ResourcePoolSubjectPolicyRequest( - component_id=component_id, - priority=1, - reserved={"gpu": 1}, - limit={"gpu": 1}, - ) - ], - ) - - _, step_1 = _create_step_run_in_db( - clean_client, - sample_pipeline_snapshot_request_model, - sample_pipeline_run_request_model, - sample_step_request_model, - ) - _, step_2 = _create_step_run_in_db( - clean_client, - sample_pipeline_snapshot_request_model, - sample_pipeline_run_request_model, - sample_step_request_model, - ) - - _create_resource_request( - clean_client, - component_id=component_id, - step_run_id=step_1, - requested_resources={"gpu": 1}, - ) - _create_resource_request( - clean_client, - component_id=component_id, - step_run_id=step_2, - requested_resources={"gpu": 1}, - ) - - with pytest.raises(IllegalOperationError): - clean_client.zen_store.delete_resource_pool(pool.id) - - assert _get_pool(clean_client, pool.id).id == pool.id - - -def test_orphan_cleanup_cancels_allocated_and_queued_requests( - clean_client, - sample_pipeline_snapshot_request_model, - sample_pipeline_run_request_model, - sample_step_request_model, -): - component_id = clean_client.active_stack.orchestrator.id - _create_pool( - clean_client, - capacity={"gpu": 1}, - policies=[ - ResourcePoolSubjectPolicyRequest( - component_id=component_id, - priority=1, - reserved={"gpu": 1}, - limit={"gpu": 1}, - ) - ], - ) - - run_1, step_1 = _create_step_run_in_db( - clean_client, - sample_pipeline_snapshot_request_model, - sample_pipeline_run_request_model, - sample_step_request_model, - ) - run_2, step_2 = _create_step_run_in_db( - clean_client, - sample_pipeline_snapshot_request_model, - sample_pipeline_run_request_model, - sample_step_request_model, - ) - _, step_3 = _create_step_run_in_db( - clean_client, - sample_pipeline_snapshot_request_model, - sample_pipeline_run_request_model, - sample_step_request_model, - ) - - allocated = _create_resource_request( - clean_client, - component_id=component_id, - step_run_id=step_1, - requested_resources={"gpu": 1}, - ) - queued = _create_resource_request( - clean_client, - component_id=component_id, - step_run_id=step_2, - requested_resources={"gpu": 1}, - ) - assert ( - _get_request(clean_client, allocated.id).status - == ResourceRequestStatus.ALLOCATED - ) - assert ( - _get_request(clean_client, queued.id).status - == ResourceRequestStatus.PENDING - ) - - clean_client.zen_store.delete_run(run_1) - clean_client.zen_store.delete_run(run_2) - - next_request = _create_resource_request( - clean_client, - component_id=component_id, - step_run_id=step_3, - requested_resources={"gpu": 1}, - ) - assert ( - _get_request(clean_client, allocated.id).status - == ResourceRequestStatus.CANCELLED - ) - assert ( - _get_request(clean_client, queued.id).status - == ResourceRequestStatus.CANCELLED - ) - assert ( - _get_request(clean_client, next_request.id).status - == ResourceRequestStatus.ALLOCATED - ) - - -def test_preemption_candidates_filtered_by_requested_resource_keys( - clean_client, - sample_pipeline_snapshot_request_model, - sample_pipeline_run_request_model, - sample_step_request_model, - monkeypatch, -): - store = clean_client.zen_store - assert isinstance(store, SqlZenStore) - low_priority_component = _create_orchestrator_component( - clean_client, "low-filter" - ) - high_priority_component = _create_orchestrator_component( - clean_client, "high-filter" - ) - - _create_pool( - clean_client, - capacity={"gpu": 1, "cpu": 1}, - policies=[ - ResourcePoolSubjectPolicyRequest( - component_id=low_priority_component, - priority=1, - reserved={"gpu": 0, "cpu": 0}, - limit={"gpu": 1, "cpu": 1}, - ), - ResourcePoolSubjectPolicyRequest( - component_id=high_priority_component, - priority=2, - reserved={"gpu": 0, "cpu": 0}, - limit={"gpu": 1, "cpu": 1}, - ), - ], - ) - - _, low_step_gpu = _create_step_run_in_db( - clean_client, - sample_pipeline_snapshot_request_model, - sample_pipeline_run_request_model, - sample_step_request_model, - ) - _, low_step_cpu = _create_step_run_in_db( - clean_client, - sample_pipeline_snapshot_request_model, - sample_pipeline_run_request_model, - sample_step_request_model, - ) - _, high_step = _create_step_run_in_db( - clean_client, - sample_pipeline_snapshot_request_model, - sample_pipeline_run_request_model, - sample_step_request_model, - ) - - gpu_victim = _create_resource_request( - clean_client, - component_id=low_priority_component, - step_run_id=low_step_gpu, - requested_resources={"gpu": 1}, - preemptible=True, - ) - cpu_victim = _create_resource_request( - clean_client, - component_id=low_priority_component, - step_run_id=low_step_cpu, - requested_resources={"cpu": 1}, - preemptible=True, - ) - assert ( - _get_request(clean_client, gpu_victim.id).status - == ResourceRequestStatus.ALLOCATED - ) - assert ( - _get_request(clean_client, cpu_victim.id).status - == ResourceRequestStatus.ALLOCATED - ) - - evictions: List[UUID] = [] - - def _mock_trigger(self, session, request_id): - evictions.append(request_id) - - monkeypatch.setattr( - SqlZenStore, "_trigger_resource_request_eviction", _mock_trigger - ) - - _create_resource_request( - clean_client, - component_id=high_priority_component, - step_run_id=high_step, - requested_resources={"gpu": 1}, - preemptible=True, - ) - - assert evictions == [gpu_victim.id] - - -def test_capacity_rebuild_keeps_boundary_multikey_requests_eligible( - clean_client, - sample_pipeline_snapshot_request_model, - sample_pipeline_run_request_model, - sample_step_request_model, -): - component_id = clean_client.active_stack.orchestrator.id - pool = _create_pool( - clean_client, - capacity={"gpu": 1, "cpu": 4}, - policies=[ - ResourcePoolSubjectPolicyRequest( - component_id=component_id, - priority=1, - reserved={"gpu": 0, "cpu": 0}, - limit={"gpu": 1, "cpu": 4}, - ) - ], - ) - - _, step_alloc = _create_step_run_in_db( - clean_client, - sample_pipeline_snapshot_request_model, - sample_pipeline_run_request_model, - sample_step_request_model, - ) - _, step_boundary = _create_step_run_in_db( - clean_client, - sample_pipeline_snapshot_request_model, - sample_pipeline_run_request_model, - sample_step_request_model, - ) - _, step_too_large = _create_step_run_in_db( - clean_client, - sample_pipeline_snapshot_request_model, - sample_pipeline_run_request_model, - sample_step_request_model, - ) - - _create_resource_request( - clean_client, - component_id=component_id, - step_run_id=step_alloc, - requested_resources={"gpu": 1, "cpu": 1}, - ) - boundary = _create_resource_request( - clean_client, - component_id=component_id, - step_run_id=step_boundary, - requested_resources={"gpu": 1, "cpu": 3}, - ) - too_large = _create_resource_request( - clean_client, - component_id=component_id, - step_run_id=step_too_large, - requested_resources={"gpu": 1, "cpu": 4}, - ) - assert ( - _get_request(clean_client, boundary.id).status - == ResourceRequestStatus.PENDING - ) - assert ( - _get_request(clean_client, too_large.id).status - == ResourceRequestStatus.PENDING - ) - - clean_client.zen_store.update_resource_pool( - pool.id, ResourcePoolUpdate(capacity={"gpu": 1, "cpu": 3}) - ) - - assert ( - _get_request(clean_client, boundary.id).status - == ResourceRequestStatus.PENDING - ) - assert ( - _get_request(clean_client, too_large.id).status - == ResourceRequestStatus.REJECTED - ) - - -def test_capacity_rebuild_keeps_pending_non_preemptible_if_still_eligible( - clean_client, - sample_pipeline_snapshot_request_model, - sample_pipeline_run_request_model, - sample_step_request_model, -): - component_id = clean_client.active_stack.orchestrator.id - pool = _create_pool( - clean_client, - capacity={"gpu": 2}, - policies=[ - ResourcePoolSubjectPolicyRequest( - component_id=component_id, - priority=1, - reserved={"gpu": 1}, - limit={"gpu": 2}, - ) - ], - ) - - _, step_alloc = _create_step_run_in_db( - clean_client, - sample_pipeline_snapshot_request_model, - sample_pipeline_run_request_model, - sample_step_request_model, - ) - _, step_non_preemptible = _create_step_run_in_db( - clean_client, - sample_pipeline_snapshot_request_model, - sample_pipeline_run_request_model, - sample_step_request_model, - ) - _, step_preemptible = _create_step_run_in_db( - clean_client, - sample_pipeline_snapshot_request_model, - sample_pipeline_run_request_model, - sample_step_request_model, - ) - - _create_resource_request( - clean_client, - component_id=component_id, - step_run_id=step_alloc, - requested_resources={"gpu": 2}, - preemptible=True, - ) - pending_non_preemptible = _create_resource_request( - clean_client, - component_id=component_id, - step_run_id=step_non_preemptible, - requested_resources={"gpu": 1}, - preemptible=False, - ) - pending_preemptible = _create_resource_request( - clean_client, - component_id=component_id, - step_run_id=step_preemptible, - requested_resources={"gpu": 1}, - preemptible=True, - ) - - clean_client.zen_store.update_resource_pool( - pool.id, ResourcePoolUpdate(capacity={"gpu": 1}) - ) - - assert ( - _get_request(clean_client, pending_non_preemptible.id).status - == ResourceRequestStatus.PENDING - ) - assert ( - _get_request(clean_client, pending_preemptible.id).status - == ResourceRequestStatus.PENDING - ) - - -def test_detach_and_delete_pool_succeed_after_requests_are_drained( - clean_client, - sample_pipeline_snapshot_request_model, - sample_pipeline_run_request_model, - sample_step_request_model, -): - component_id = clean_client.active_stack.orchestrator.id - pool = _create_pool( - clean_client, - capacity={"gpu": 1}, - policies=[ - ResourcePoolSubjectPolicyRequest( - component_id=component_id, - priority=1, - reserved={"gpu": 1}, - limit={"gpu": 1}, - ) - ], - ) - - run_1, step_1 = _create_step_run_in_db( - clean_client, - sample_pipeline_snapshot_request_model, - sample_pipeline_run_request_model, - sample_step_request_model, - ) - _create_resource_request( - clean_client, - component_id=component_id, - step_run_id=step_1, - requested_resources={"gpu": 1}, - ) - clean_client.zen_store.delete_run(run_1) - - store = clean_client.zen_store - assert isinstance(store, SqlZenStore) - with Session(store.engine) as session: - store._allocate_queued_requests_for_pool( - session=session, pool_id=pool.id - ) - session.commit() - - clean_client.zen_store.update_resource_pool( - pool.id, ResourcePoolUpdate(detach_policies=[component_id]) - ) - clean_client.zen_store.delete_resource_pool(pool.id) - - with pytest.raises(KeyError): - clean_client.zen_store.get_resource_pool(pool.id) - - -def test_explicit_bounded_capacity_for_unbounded_key_is_enforced( - clean_client, - sample_pipeline_snapshot_request_model, - sample_pipeline_run_request_model, - sample_step_request_model, -): - component_id = clean_client.active_stack.orchestrator.id - _create_pool( - clean_client, - capacity={"cpu": 1}, - policies=[ - ResourcePoolSubjectPolicyRequest( - component_id=component_id, - priority=1, - reserved={}, - limit={}, - ) - ], - ) - - _, step_id = _create_step_run_in_db( - clean_client, - sample_pipeline_snapshot_request_model, - sample_pipeline_run_request_model, - sample_step_request_model, - ) - request = _create_resource_request( - clean_client, - component_id=component_id, - step_run_id=step_id, - requested_resources={"cpu": 2}, - ) - - assert ( - _get_request(clean_client, request.id).status - == ResourceRequestStatus.REJECTED - ) - - -def test_detach_policy_not_blocked_by_other_component_activity( - clean_client, - sample_pipeline_snapshot_request_model, - sample_pipeline_run_request_model, - sample_step_request_model, -): - component_a = _create_orchestrator_component( - clean_client, "detach-scope-a" - ) - component_b = _create_orchestrator_component( - clean_client, "detach-scope-b" - ) - pool = _create_pool( - clean_client, - capacity={"gpu": 1}, - policies=[ - ResourcePoolSubjectPolicyRequest( - component_id=component_a, - priority=1, - reserved={"gpu": 0}, - limit={"gpu": 1}, - ), - ResourcePoolSubjectPolicyRequest( - component_id=component_b, - priority=2, - reserved={"gpu": 0}, - limit={"gpu": 1}, - ), - ], - ) - - _, step_b = _create_step_run_in_db( - clean_client, - sample_pipeline_snapshot_request_model, - sample_pipeline_run_request_model, - sample_step_request_model, - ) - active_b = _create_resource_request( - clean_client, - component_id=component_b, - step_run_id=step_b, - requested_resources={"gpu": 1}, - ) - assert ( - _get_request(clean_client, active_b.id).status - == ResourceRequestStatus.ALLOCATED - ) - - clean_client.zen_store.update_resource_pool( - pool.id, ResourcePoolUpdate(detach_policies=[component_a]) - ) - - -def test_detach_policy_blocked_when_component_has_only_queued_requests( - clean_client, - sample_pipeline_snapshot_request_model, - sample_pipeline_run_request_model, - sample_step_request_model, -): - component_a = _create_orchestrator_component( - clean_client, "detach-queued-a" - ) - component_b = _create_orchestrator_component( - clean_client, "detach-queued-b" - ) - pool = _create_pool( - clean_client, - capacity={"gpu": 1}, - policies=[ - ResourcePoolSubjectPolicyRequest( - component_id=component_a, - priority=1, - reserved={"gpu": 0}, - limit={"gpu": 1}, - ), - ResourcePoolSubjectPolicyRequest( - component_id=component_b, - priority=2, - reserved={"gpu": 0}, - limit={"gpu": 1}, - ), - ], - ) - - _, step_b = _create_step_run_in_db( - clean_client, - sample_pipeline_snapshot_request_model, - sample_pipeline_run_request_model, - sample_step_request_model, - ) - _, step_a = _create_step_run_in_db( - clean_client, - sample_pipeline_snapshot_request_model, - sample_pipeline_run_request_model, - sample_step_request_model, - ) - - _create_resource_request( - clean_client, - component_id=component_b, - step_run_id=step_b, - requested_resources={"gpu": 1}, - ) - queued_a = _create_resource_request( - clean_client, - component_id=component_a, - step_run_id=step_a, - requested_resources={"gpu": 1}, - ) - assert ( - _get_request(clean_client, queued_a.id).status - == ResourceRequestStatus.PENDING - ) - - with pytest.raises(IllegalOperationError) as error: - clean_client.zen_store.update_resource_pool( - pool.id, ResourcePoolUpdate(detach_policies=[component_a]) - ) - - assert "queued=1" in str(error.value) - assert "allocated=0" in str(error.value) - - -def test_delete_pool_blocked_when_only_allocated_requests_exist( - clean_client, - sample_pipeline_snapshot_request_model, - sample_pipeline_run_request_model, - sample_step_request_model, -): - component_id = clean_client.active_stack.orchestrator.id - pool = _create_pool( - clean_client, - capacity={"gpu": 1}, - policies=[ - ResourcePoolSubjectPolicyRequest( - component_id=component_id, - priority=1, - reserved={"gpu": 0}, - limit={"gpu": 1}, - ) - ], - ) - - _, step_1 = _create_step_run_in_db( - clean_client, - sample_pipeline_snapshot_request_model, - sample_pipeline_run_request_model, - sample_step_request_model, - ) - request = _create_resource_request( - clean_client, - component_id=component_id, - step_run_id=step_1, - requested_resources={"gpu": 1}, - ) - assert ( - _get_request(clean_client, request.id).status - == ResourceRequestStatus.ALLOCATED - ) - - with pytest.raises(IllegalOperationError) as error: - clean_client.zen_store.delete_resource_pool(pool.id) - - assert "queued=0" in str(error.value) - assert "allocated=1" in str(error.value) - - -def test_reconciliation_cancels_orphaned_allocations_without_new_requests( - clean_client, - sample_pipeline_snapshot_request_model, - sample_pipeline_run_request_model, - sample_step_request_model, -): - store = clean_client.zen_store - assert isinstance(store, SqlZenStore) - component_id = clean_client.active_stack.orchestrator.id - _create_pool( - clean_client, - capacity={"gpu": 1}, - policies=[ - ResourcePoolSubjectPolicyRequest( - component_id=component_id, - priority=1, - reserved={"gpu": 1}, - limit={"gpu": 1}, - ) - ], - ) - - run_id, step_id = _create_step_run_in_db( - clean_client, - sample_pipeline_snapshot_request_model, - sample_pipeline_run_request_model, - sample_step_request_model, - ) - request = _create_resource_request( - clean_client, - component_id=component_id, - step_run_id=step_id, - requested_resources={"gpu": 1}, - ) - assert ( - _get_request(clean_client, request.id).status - == ResourceRequestStatus.ALLOCATED - ) - - clean_client.zen_store.delete_run(run_id) - store.reconcile_resource_pools() - - assert ( - _get_request(clean_client, request.id).status - == ResourceRequestStatus.CANCELLED - ) - - -def test_reconciliation_repairs_occupied_resource_counters( - clean_client, - sample_pipeline_snapshot_request_model, - sample_pipeline_run_request_model, - sample_step_request_model, -): - store = clean_client.zen_store - assert isinstance(store, SqlZenStore) - component_id = clean_client.active_stack.orchestrator.id - pool = _create_pool( - clean_client, - capacity={"gpu": 1}, - policies=[ - ResourcePoolSubjectPolicyRequest( - component_id=component_id, - priority=1, - reserved={"gpu": 0}, - limit={"gpu": 1}, - ) - ], - ) - - _, step_id = _create_step_run_in_db( - clean_client, - sample_pipeline_snapshot_request_model, - sample_pipeline_run_request_model, - sample_step_request_model, - ) - request = _create_resource_request( - clean_client, - component_id=component_id, - step_run_id=step_id, - requested_resources={"gpu": 1}, - ) - assert ( - _get_request(clean_client, request.id).status - == ResourceRequestStatus.ALLOCATED - ) - - with Session(store.engine) as session: - pool_resource = session.exec( - select(ResourcePoolResourceSchema) - .where(ResourcePoolResourceSchema.pool_id == pool.id) - .where(ResourcePoolResourceSchema.key == "gpu") - ).one() - pool_resource.occupied = 0 - session.add(pool_resource) - session.commit() - - store.reconcile_resource_pools() - - with Session(store.engine) as session: - pool_resource = session.exec( - select(ResourcePoolResourceSchema) - .where(ResourcePoolResourceSchema.pool_id == pool.id) - .where(ResourcePoolResourceSchema.key == "gpu") - ).one() - assert pool_resource.occupied == 1 - - -def test_reconciliation_cleans_stale_queue_entries( - clean_client, - sample_pipeline_snapshot_request_model, - sample_pipeline_run_request_model, - sample_step_request_model, -): - store = clean_client.zen_store - assert isinstance(store, SqlZenStore) - component_id = clean_client.active_stack.orchestrator.id - pool = _create_pool( - clean_client, - capacity={"gpu": 1}, - policies=[ - ResourcePoolSubjectPolicyRequest( - component_id=component_id, - priority=1, - reserved={"gpu": 0}, - limit={"gpu": 1}, - ) - ], - ) - _, step_id = _create_step_run_in_db( - clean_client, - sample_pipeline_snapshot_request_model, - sample_pipeline_run_request_model, - sample_step_request_model, - ) - request = _create_resource_request( - clean_client, - component_id=component_id, - step_run_id=step_id, - requested_resources={"gpu": 1}, - ) - request_model = _get_request(clean_client, request.id) - assert request_model.status == ResourceRequestStatus.ALLOCATED - - with Session(store.engine) as session: - session.add( - ResourcePoolQueueSchema( - pool_id=pool.id, - request_id=request.id, - priority=1, - request_created=request_model.created, - ) - ) - session.commit() - - with Session(store.engine) as session: - stale_queue_entries_before = len( - session.exec( - select(ResourcePoolQueueSchema.id) - .join( - ResourceRequestSchema, - ResourceRequestSchema.id - == ResourcePoolQueueSchema.request_id, - ) - .where(ResourcePoolQueueSchema.pool_id == pool.id) - .where( - ResourceRequestSchema.status - != ResourceRequestStatus.PENDING.value - ) - ).all() - ) - assert stale_queue_entries_before == 1 - - store.reconcile_resource_pools() - - with Session(store.engine) as session: - stale_queue_entries_after = len( - session.exec( - select(ResourcePoolQueueSchema.id) - .join( - ResourceRequestSchema, - ResourceRequestSchema.id - == ResourcePoolQueueSchema.request_id, - ) - .where(ResourcePoolQueueSchema.pool_id == pool.id) - .where( - ResourceRequestSchema.status - != ResourceRequestStatus.PENDING.value - ) - ).all() - ) - assert stale_queue_entries_after == 0 diff --git a/tests/unit/zen_stores/test_step_resource_requests.py b/tests/unit/zen_stores/test_step_resource_requests.py new file mode 100644 index 00000000000..c4e23381384 --- /dev/null +++ b/tests/unit/zen_stores/test_step_resource_requests.py @@ -0,0 +1,175 @@ +# Copyright (c) ZenML GmbH 2026. All Rights Reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at: +# +# https://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express +# or implied. See the License for the specific language governing +# permissions and limitations under the License. +"""Tests for step-run resource request helpers.""" + +from datetime import datetime, timedelta, timezone +from unittest.mock import MagicMock +from uuid import uuid4 + +import pytest + +from zenml.config.resource_settings import ResourceSettings +from zenml.enums import ( + ExecutionStatus, + ResourceRequestReclaimTolerance, + ResourceRequestRuntimeState, + ResourceRequestStatus, + StepRuntime, +) +from zenml.exceptions import IllegalOperationError +from zenml.zen_stores.sql_zen_store import SqlZenStore + + +def _sql_store_with_resource_pools() -> SqlZenStore: + store = object.__new__(SqlZenStore) + resource_pools = MagicMock() + renewed_request = MagicMock() + renewed_request.status = ResourceRequestStatus.ALLOCATED + resource_pools.renew_resource_request.return_value = renewed_request + store._resource_pools = resource_pools + return store + + +def _heartbeat_step_run() -> MagicMock: + return MagicMock( + status=ExecutionStatus.RUNNING.value, + resource_request_id=uuid4(), + heartbeat_threshold=30, + name="step", + ) + + +def test_omitted_inline_reclaim_tolerance_is_none() -> None: + """Inline runtime requests use none when the user did not set tolerance.""" + assert ( + ResourceSettings().effective_reclaim_tolerance(StepRuntime.INLINE) + is ResourceRequestReclaimTolerance.NONE + ) + + +def test_omitted_isolated_reclaim_tolerance_is_any() -> None: + """Isolated runtime requests use any when the user did not set tolerance.""" + assert ( + ResourceSettings().effective_reclaim_tolerance(StepRuntime.ISOLATED) + is ResourceRequestReclaimTolerance.ANY + ) + + +def test_explicit_reclaim_tolerance_is_effective_value() -> None: + """Explicit reclaim tolerance is preserved for resource requests.""" + resource_settings = ResourceSettings( + reclaim_tolerance=ResourceRequestReclaimTolerance.ANY + ) + + assert ( + resource_settings.effective_reclaim_tolerance(StepRuntime.INLINE) + is ResourceRequestReclaimTolerance.ANY + ) + + +def test_implicit_inline_steps_use_none_reclaim_tolerance() -> None: + """Inline dynamic steps treat omitted reclaim tolerance as none.""" + SqlZenStore._validate_reclaim_tolerance_for_resource_request( + resource_settings=ResourceSettings(), + runtime=StepRuntime.INLINE, + heartbeat_enabled=False, + step_name="step", + ) + + +def test_explicit_inline_reclaimable_steps_fail_validation() -> None: + """Inline dynamic steps cannot be explicitly configured as reclaimable.""" + with pytest.raises(IllegalOperationError): + SqlZenStore._validate_reclaim_tolerance_for_resource_request( + resource_settings=ResourceSettings( + reclaim_tolerance=ResourceRequestReclaimTolerance.ANY + ), + runtime=StepRuntime.INLINE, + heartbeat_enabled=True, + step_name="step", + ) + + +def test_isolated_reclaimable_steps_require_heartbeat() -> None: + """Heartbeat is required for reclaimable isolated dynamic steps.""" + with pytest.raises(IllegalOperationError): + SqlZenStore._validate_reclaim_tolerance_for_resource_request( + resource_settings=ResourceSettings( + reclaim_tolerance=ResourceRequestReclaimTolerance.COORDINATED + ), + runtime=StepRuntime.ISOLATED, + heartbeat_enabled=False, + step_name="step", + ) + + +def test_isolated_non_reclaimable_steps_do_not_require_heartbeat() -> None: + """Non-reclaimable isolated steps can run without heartbeat.""" + SqlZenStore._validate_reclaim_tolerance_for_resource_request( + resource_settings=ResourceSettings(), + runtime=StepRuntime.ISOLATED, + heartbeat_enabled=False, + step_name="step", + ) + + +def test_heartbeat_renewal_uses_expected_heartbeat_interval( + monkeypatch, +) -> None: + """Heartbeat renewal leases use the timeout reported by the caller.""" + fixed_now = datetime(2026, 1, 1, tzinfo=timezone.utc) + monkeypatch.setattr( + "zenml.zen_stores.sql_zen_store.utc_now", lambda: fixed_now + ) + store = _sql_store_with_resource_pools() + step_run = _heartbeat_step_run() + + status = store._renew_step_resource_request_from_heartbeat( + session=MagicMock(), + step_run=step_run, + heartbeat_liveness_timeout_seconds=42, + ) + + assert status == ExecutionStatus.RUNNING + renewal_request = store.resource_pools.renew_resource_request.call_args[0][ + 1 + ] + assert renewal_request.lease_expires_at == fixed_now + timedelta( + seconds=42 + ) + assert renewal_request.runtime_state == ResourceRequestRuntimeState.RUNNING + + +def test_heartbeat_renewal_without_expected_interval_uses_step_threshold( + monkeypatch, +) -> None: + """Legacy heartbeat calls renew leases from the stored step threshold.""" + fixed_now = datetime(2026, 1, 1, tzinfo=timezone.utc) + monkeypatch.setattr( + "zenml.zen_stores.sql_zen_store.utc_now", lambda: fixed_now + ) + store = _sql_store_with_resource_pools() + step_run = _heartbeat_step_run() + + store._renew_step_resource_request_from_heartbeat( + session=MagicMock(), + step_run=step_run, + ) + + renewal_request = store.resource_pools.renew_resource_request.call_args[0][ + 1 + ] + assert renewal_request.lease_expires_at == fixed_now + timedelta( + minutes=step_run.heartbeat_threshold + )