shared-compute implementation to support multiple GPU SKUs#92
Open
mgmachado wants to merge 5 commits into
Open
shared-compute implementation to support multiple GPU SKUs#92mgmachado wants to merge 5 commits into
mgmachado wants to merge 5 commits into
Conversation
There was a problem hiding this comment.
Pull request overview
Adds multi-tenant shared Azure Batch GPU pool support (H100/T4 tiers) with capacity-aware pool routing and per-job user-delegation SAS-based blob I/O, plus IaC/templates and documentation/spec updates to enable shared-pool provisioning and tenant isolation.
Changes:
- Introduces shared-pools IaC (
infra/shared-pools.bicep+ params) and expandsbatchPool.bicepto support fixed/autoscale, dedicated/low-priority, and optional VNet injection. - Updates
hastelibAzure Batch runner to select pools at submit time and optionally use per-job SAS instead of pool identity for blob transfers; processors pass candidate pool lists; adds unit tests. - Documents the feature via a full spec (
spec/features/batch-compute-expansion/*), configuration/deployment docs, and changelog entries; bumps hastegeo wheel references.
Reviewed changes
Copilot reviewed 36 out of 37 changed files in this pull request and generated 8 comments.
Show a summary per file
| File | Description |
|---|---|
| spec/features/batch-compute-expansion/user-stories.md | Defines user stories + acceptance criteria for shared pools, routing, isolation, and rollout. |
| spec/features/batch-compute-expansion/test-plan.md | Test strategy + scenarios for routing/SAS and infra validation. |
| spec/features/batch-compute-expansion/rollout.md | Phased rollout plan with flags, monitoring, and rollback. |
| spec/features/batch-compute-expansion/README.md | Feature overview, motivation, constraints, and success criteria. |
| spec/features/batch-compute-expansion/plan.md | Execution plan with phases/milestones and ownership. |
| spec/features/batch-compute-expansion/networking.md | Network design for blob↔batch reach via VNet injection + storage VNet rules. |
| spec/features/batch-compute-expansion/impact-analysis.md | Risk, dependency, and rollback analysis for the change. |
| spec/features/batch-compute-expansion/design.md | Technical design for pool topology, SAS isolation model, and routing logic. |
| setup/README.md | Updates default/example resource prefix guidance. |
| infra/shared-pools.json | Generated ARM template for shared-pools deployment. |
| infra/shared-pools.bicepparam | Parameter file reading env vars for shared-pools deployment. |
| infra/shared-pools.bicep | New standalone deployment creating shared H100+T4 pools via module calls. |
| infra/modules/storage.bicep | Adds optional shared batch-subnet VNet rule for tenant storage allowlisting. |
| infra/modules/functions.bicep | Adds app settings for pool candidate lists and SAS/manage-pools flags. |
| infra/modules/functionApp.bicep | Grants Function App identity Storage Blob Delegator role assignment. |
| infra/modules/batchPool.bicep | Parameterizes pool scale mode/node type/min nodes and optional subnet injection. |
| infra/main.bicepparam | Changes default prefix and wires new batch routing/SAS params + shared subnet id. |
| infra/main.bicep | Wires shared subnet allowlisting into storage and routes new batch settings into functions module. |
| hastelib/tests/core/runners/test_azure_batch_routing.py | Adds unit tests for select_pool and SAS-vs-identity behavior. |
| hastelib/tests/core/runners/init.py | Adds package marker for runners tests. |
| hastelib/src/hastegeo/core/runners/local.py | Extends LocalRunner ctor signature for interface parity (candidate pools). |
| hastelib/src/hastegeo/core/runners/azure_batch.py | Implements submit-time pool selection and per-job SAS URL augmentation + manage_pools gating. |
| hastelib/src/hastegeo/core/processors/train.py | Passes training candidate pool list into runner creation. |
| hastelib/src/hastegeo/core/processors/inference.py | Passes inference candidate pool list into runner creation. |
| hastelib/src/hastegeo/core/processors/imagery.py | Passes imageryprep candidate pool list into runner creation. |
| hastelib/src/hastegeo/core/processors/embedding.py | Passes candidate pool list into runner creation. |
| hastelib/src/hastegeo/core/processors/artifacts.py | Passes imageryprep candidate pool list into runner creation. |
| hastelib/src/hastegeo/core/config.py | Adds parsing for candidate pool-id lists and SAS/manage-pools flags. |
| hastelib/src/hastegeo/about.py | Bumps hastegeo version. |
| hastelib/pyproject.toml | Pins azure-batch==14.2.0 to avoid track-2 API breakage. |
| docs/deployment.md | Adds production note pointing to shared multi-tenant GPU pools docs. |
| docs/configuration.md | Documents shared multi-tenant GPU pools, flags, and networking model. |
| docker/imageryprep/requirements.txt | Updates hastegeo wheel pin version. |
| CHANGELOG.md | Adds unreleased entry describing shared pools, routing, SAS isolation, and IaC changes. |
| api/hastefuncqueues/requirements.txt | Updates hastegeo wheel pin version. |
| api/hastefuncapi/requirements.txt | Updates hastegeo wheel pin version. |
Comment on lines
+287
to
+298
| def _sas_url(self, url, permissions): | ||
| # Append a user-delegation SAS scoped to the URL's container so a Batch | ||
| # ResourceFile/OutputFile can read/write WITHOUT the pool holding any | ||
| # standing data access — the isolation boundary for multi-tenant shared | ||
| # pools. The submitting identity needs `Storage Blob Delegator` on the | ||
| # account. User-delegation keys are cached per account. | ||
| parsed = urlparse(url) | ||
| account_url = f"{parsed.scheme}://{parsed.netloc}" | ||
| account_name = parsed.netloc.split(".")[0] | ||
| container = parsed.path.lstrip("/").split("/", 1)[0] | ||
| if self._sas_credential is None: | ||
| self._sas_credential = DefaultAzureCredential() |
Comment on lines
+56
to
+59
| // Autoscale targets the node bucket matching nodeType; minNodes=0 => scale-to-zero | ||
| // when idle (shared-demo burst preserves scarce GPU quota). | ||
| var scaleTargetVar = nodeType == 'LowPriority' ? '$TargetLowPriorityNodes' : '$TargetDedicatedNodes' | ||
| var autoscaleFormula = '$samples = $ActiveTasks.GetSamplePercent(TimeInterval_Minute * 15);$tasks = $samples < 70 ? max(0, $ActiveTasks.GetSample(1)) : max($ActiveTasks.GetSample(1), avg($ActiveTasks.GetSample(TimeInterval_Minute * 15)));$targetVMs = $tasks > 0 ? $tasks : ${minNodes};${scaleTargetVar} = max(${minNodes}, min($targetVMs, ${maxNodes}));$NodeDeallocationOption = taskcompletion;' |
Comment on lines
+46
to
+50
| @allowed([ | ||
| 'Dedicated' | ||
| 'LowPriority' | ||
| ]) | ||
| param sharedNodeType string = 'Dedicated' |
|
|
||
| param sharedPrefix = readEnvironmentVariable('HASTE_RESOURCE_PREFIX', 'haste') | ||
| param sharedGroup = readEnvironmentVariable('HASTE_SHARED_GROUP', 'dev') | ||
| param sharedNodeType = readEnvironmentVariable('HASTE_SHARED_NODE_TYPE', 'Dedicated') |
Comment on lines
+58
to
+62
| The submitting Function App knows the tenant. For **every** job (all pools) it | ||
| mints a **user-delegation SAS** (signed by the env identity via | ||
| `get_user_delegation_key`, no account key) scoped to that tenant's container and | ||
| prefix, with a short TTL, read for inputs / write for outputs. It embeds the SAS | ||
| in the blob URL and passes it to Batch: |
Comment on lines
+122
to
+126
| Pools are named `<prefix>-shared-<group>-<tier>-pool` (e.g. the dev-group H100 | ||
| pool); `HASTE_SHARED_GROUP` selects the group (`dev`, `demo`, …). They autoscale | ||
| on **dedicated** nodes within a central per-pool core-quota ceiling. A scarce tier | ||
| (H100) can keep a warm floor instead of scaling to zero — under regional GPU | ||
| contention an autoscale floor (or `Fixed` target) **auto-retries allocation each |
Comment on lines
+141
to
+145
| **Data isolation.** Tenants share compute but not data: each job mints a | ||
| short-lived **user-delegation SAS** scoped to its own storage container, so a | ||
| tenant's task can never read another tenant's data. This requires the submitting | ||
| Function App identity to hold **Storage Blob Delegator** on its storage account | ||
| (granted in `functionApp.bicep`). |
Comment on lines
+135
to
+139
| | `HASTE_SHARED_GROUP` | `dev` | Pool group (`dev`, `demo`, …) — drives the pool names. | | ||
| | `HASTE_SHARED_H100_SCALE_MODE` | `Autoscale` | `Autoscale`, or `Fixed` reserved baseline for the scarce H100 tier. | | ||
| | `HASTE_SHARED_H100_MIN_NODES` | `0` | H100 autoscale floor (>0 keeps chasing/holding a warm node). | | ||
| | `HASTE_SHARED_T4_MIN_NODES` / `HASTE_SHARED_T4_MAX_NODES` | `0` / `2` | T4 autoscale floor / cap. | | ||
| | `HASTE_SHARED_BATCH_SUBNET_ID` | — | Shared hub batch-subnet to VNet-inject both pools into (see Blob↔Batch networking below). | |
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Description
Adds multi-tenant shared GPU Batch pools so many environments can share scarce GPU quota instead of fragmenting it across one pool per environment, plus the v2.1.0
hastelibrouting + per-job SAS changes that make it work.batchPool.bicep(fixed/autoscale × dedicated/low-priority × optional VNet) and a new standaloneshared-pools.bicep+shared-pools.bicepparamthat provision an H100 + T4 pool pair per group (dev, demo, …), autoscaling on low-priority nodes and scaling to zero when idle.hastelibnow takes ordered candidate pool lists and binds a pool at submit time (select_pool) — preferred tier first, spillover to a tier with a free node.HASTE_RESOURCE_PREFIX→ neutralhaste, BYO account/ACR), so other partners can deploy without editing the templates.All new behavior is behind per-environment flags that default to the existing single-pool, pool-identity path, so existing environments are unaffected until they explicitly opt in.
Full design:
spec/features/batch-compute-expansion/.Fixes #
Type of change
Checklist
cd hastelib && hatch run test:pytest) and the UI lints clean (cd ui && npm run lint) — new unit tests pass (9/9); full suite to run in CI (conda env)Testing
hastelib/tests/core/runners/test_azure_batch_routing.py— 9 tests coveringselect_pool(single / spillover-to-idle / preferred fallback), the SAS-vs-identity blob-credential toggle, and config candidate-list parsing. All pass.ResourceFile/OutputFileI/O works with noidentity_reference; a cross-tenant SAS was denied (BlobAccessDenied); the ACR-pull-only shared identity authenticated and pulled the training image.batchPool.bicep+shared-pools.bicepcompile clean;what-ifshows no drift against the live pools.rollout.md).Additional context
AZURE_BATCH_USE_SAS=true,AZURE_BATCH_MANAGE_POOLS=false, andAZURE_BATCH_{TRAINING,INFERENCE,IMAGERYPREP}_POOL_IDS; all default to legacy.Storage Blob Delegatoris granted to theapi/queuesidentity (infunctionApp.bicep) so it can mint the SAS.azure-batchpinned to==14.2.0— the 15.x track-2 rewrite restructures the models this code uses; migration tracked separately.README,plan,impact-analysis,user-stories,design,test-plan,rollout.