feat: Common code changes required for Azure and GCP support#256
feat: Common code changes required for Azure and GCP support#256IkeM-L wants to merge 9 commits into
Conversation
fgogolli
changed the title
This comment was marked as outdated.
This comment was marked as outdated.
Sorry, something went wrong.
There was a problem hiding this comment.
This PR pre-dates #253's merge by some weeks - the split from #240 happened while #253 was still in review, so a rebase and alignment pass is the dominant action item before any of #256-#258 can land. The inline comments flag specific conflicts; this body summarises the recommended shape of work.
Suggested split (3 PRs, not the current single 96-file PR)
The current PR mixes genuinely shared scaffold, AWS-specific feature work, and items that belong in the Azure PR. Splitting it would make each review tractable and avoids holding up the Azure and GCP PRs on AWS-only concerns.
Slimmed #256 (shared scaffold only, ~10-15 files)
Keep: the orb init prompt-flag pattern and _configure_provider_interactively generalisation, GetAttrChecker in quality_check.py, the provider_config field on TemplateDTO (once the type is corrected to Optional[BaseModel] per the inline comment), GetAttrChecker CI wiring, and any other genuinely provider-neutral touch-ups.
New AWS-focused PR
Move: SpotPlacementPlanner, SpotPlacementExecutionService, and AWSSpotPlacementScoreAdapter (these are AWS-only at present - Azure ships its own AzureSpotPlacementScoreAdapter in #258 and never calls SpotPlacementPlanner; GCP has no spot placement score concept at all), the aws_provider_strategy.py refactor, and the check_hosts_status -> CheckHostsStatusResult migration on AWS handlers. The test coverage gaps called out on SpotPlacementPlanner (zero-candidate input, boundary primary_share_percent, single-candidate HYBRID, tie-breaking) should travel with the file to this PR.
#258 absorbs the Azure-specific items
vm_size/vm_sizes and the four placement_* fields should move to AzureTemplate(Template) in providers/azure/domain/template/ rather than polluting the shared base. AzureAllocationStrategy belongs in providers/azure/ - #258 already maps to it via AzureAllocationStrategy.from_core(...), so it already knows about a provider-specific enum. The Azure-specific error fields in error_codes.py should move to providers/azure/infrastructure/ when they travel with #258.
#257 stays in current scope
GCP content is clean. Rebases onto post-#253 main, adopts CheckHostsStatusResult contract, and it is ready to go.
Items that need alignment with #253 after rebase
Six specific items warrant explicit reconciliation when rebasing. provider_config on TemplateDTO must stay Optional[BaseModel] and route through TemplateExtensionRegistry.create_extension_config(provider_type, raw_extra) - the dict approach loses the per-provider Pydantic validation and the @field_serializer. The TemplateExtensionRegistry implementation belongs in src/orb/infrastructure/registry/ (all sibling registries - DefaultsLoaderRegistry, FieldMappingRegistry, CLISpecRegistry - live there); moving it to the domain layer drags Pydantic into a layer that must stay framework-free. The OperationOutcome discriminated union from #253 (Accepted | Completed | RequiresFollowUp | Failed) should be kept - the assert_never exhaustiveness check is precisely what catches missing provider variants at type-check time; deleting it without a typed replacement leaves provider strategies returning stringly-typed provider_metadata dicts again. The config/loader.py dispatch dict is the inverse of #253's DefaultsLoaderRegistry design - see the inline comment for the registry-lookup replacement. check_hosts_status on the AWS handlers now returns CheckHostsStatusResult(instances, fulfilment=ProviderFulfilment(...)) on main; the PR's list[dict] return shape will need rebasing into that contract or RequestStatusService loses the fulfilment verdict. Related: the strategy now offloads check_hosts_status via asyncio.to_thread (the boto3 calls are synchronous and were blocking the event loop pre-#253) - please keep that wrapper in the rebased handlers.
Test coverage
request_follow_up_context.py has three public functions with no direct tests; the merge-priority rule and the None/empty-string filter in with_request_follow_up_context are the specific behaviours that need pinning. collect_provider_error_codes is pure and trivially testable. Unspec'd MagicMock() is used throughout the new service tests - adding spec= to the key mocks will catch the attribute-access typos that spec-less mocks silently swallow.
| tags: dict[str, Any] = Field(default_factory=dict) | ||
| metadata: dict[str, Any] = Field(default_factory=dict) | ||
| provider_config: dict[str, Any] = Field(default_factory=dict) | ||
|
|
There was a problem hiding this comment.
Since this PR pre-dates #253, the typed Optional[BaseModel] contract isn't here yet - that's expected. After rebase, adopt that contract and route from_domain through TemplateExtensionRegistry.create_extension_config(provider_type, raw_extra) so Azure and GCP provider configs land as Pydantic models rather than untyped dicts. The dict bucket approach loses the per-provider validation and the @field_serializer that #253 added; to_template_config() would then call provider_config.model_dump() in the promotion step rather than iterating raw keys.
|
|
||
|
|
||
| def get_request_follow_up_context(request: Request) -> dict[str, Any]: | ||
| """Return durable provider follow-up context for a request.""" |
There was a problem hiding this comment.
The naming collision framing still applies - on main this file is a re-export shim for FollowUpContext types from orb.domain.base.follow_up_context. After rebase, keep these three functions and add the shim re-exports alongside them. Worth also considering moving them to methods on Request directly - the aggregate already has get_provider_data(key) / update_provider_data(updates), and these helpers are thin wrappers over provider_data['follow_up_context']. Keeping the key string at every call site leaks an internal. The architecture review surfaced this seam.
|
|
||
|
|
||
| class SpotPlacementPlanner: | ||
| """Build a placement plan from normalized candidate scores.""" |
There was a problem hiding this comment.
Suggest this lands in its own AWS-focused PR rather than under 'common code' framing. The planner, the execution service, and the score adapter are all AWS-only - Azure ships its own AzureSpotPlacementScoreAdapter in #258 with no Azure use of SpotPlacementPlanner, and GCP has no spot placement score concept at all. The test coverage gaps noted previously (zero-candidate input, boundary primary_share_percent, single-candidate HYBRID, tie-breaking) should travel with the file to that AWS PR.
|
|
||
|
|
||
| def collect_provider_error_codes(errors: list[ProviderErrorEntry]) -> list[str]: | ||
| """Return unique canonical error codes from normalized provider errors.""" |
There was a problem hiding this comment.
The ProviderErrorEntry dict bakes in CycleCloud-specific fields (node_array, cc_state, launch_template_id, instance_requirements) under a 'provider-neutral' name. Only Azure calls collect_provider_error_codes on the PR branch. Suggest moving the file into providers/azure/infrastructure/error_codes.py for now and introducing an ErrorNormalisationPort only when a second provider actually needs the abstraction. Coverage for the function should follow it.
There was a problem hiding this comment.
Supplemental inline comments to accompany the updated top-level review. The first review covered the provider_config type conflict, request_follow_up_context.py naming, the spot placement planner scope, and error_codes.py placement. These four new threads cover: CLISpecRegistry duplication in init_command_handler, AllocationStrategy/PlacementSplitStrategy in base domain, Azure-specific fields on the shared Template aggregate, and the config/loader.py registry regression.
| strategy = _get_provider_strategy(provider_type) | ||
| default_region = strategy.get_default_region() if strategy is not None else "" | ||
| provider_config = strategy.get_cli_provider_config(args) if strategy is not None else {} | ||
| if not provider_config.get("region"): |
There was a problem hiding this comment.
The new get_cli_provider_config and get_cli_infrastructure_defaults methods on ProviderStrategy duplicate the existing CLISpecRegistry introduced in #253. Each provider already declares its CLI surface via a CLISpec class registered from providers/<name>/cli_spec.py - see aws_cli_spec.py for the pattern. After rebase, route through CLISpecRegistry.get(provider_type).extract_config(args) rather than adding parallel methods on the strategy. Both #257 (GCP) and #258 (Azure) already ship their *_cli_spec.py files; keeping the integration surface to one mechanism avoids a split brain between CLISpec and get_cli_provider_config for every future provider.
|
|
||
|
|
||
| class AllocationStrategy(str, Enum): | ||
| """Allocation strategy enumeration.""" |
There was a problem hiding this comment.
AllocationStrategy and PlacementSplitStrategy carry EC2 Fleet / Spot Fleet API vocabulary (lowestPrice, capacityOptimized, spotPlacementScore, etc.). Main already has AWSAllocationStrategy in providers/aws/domain/template/value_objects.py - that's where these belong. Azure can declare its own AzureAllocationStrategy enum in providers/azure/ (and #258 does map to AzureAllocationStrategy.from_core(...), so it already knows about a provider-specific enum). Base domain stays clean of provider-specific identifiers.
| # Instance configuration | ||
| instance_type: Optional[str] = None | ||
| vm_size: Optional[str] = None | ||
| vm_sizes: list[str] = Field(default_factory=list) |
There was a problem hiding this comment.
vm_size, vm_sizes, and the four placement_* fields are Azure-specific - GCP uses instance_type and handles zone distribution via its own zones/mig_scope. Suggest moving these to AzureTemplate(Template) in providers/azure/domain/template/ rather than the shared base, following the same pattern as AWSTemplate. Every future provider currently inherits four placement_* fields it will never use.
| """Load static provider defaults without bootstrapping provider registries.""" | ||
| merged: dict[str, Any] = {} | ||
| provider_default_loaders = { | ||
| "aws": cls._load_aws_provider_defaults, |
There was a problem hiding this comment.
Hardcoding the dispatch dict here is the inverse of #253's DefaultsLoaderRegistry design - that registry lets each provider register its own defaults loader without touching shared code. After rebase, replace this dict with the registry lookup pattern (see DefaultsLoaderRegistry.collect_defaults() for the consumer side). Azure and GCP authors shouldn't need to modify loader.py to add their defaults.
|
Now that #285 has merged, this PR needs a rebase against main with a few adjustments. The provider extension points from #285 already added The
Most of the PR (SpotPlacement services, FollowUpContext, machine sync improvements) is additive and should apply cleanly. |
|
Thanks so much for this — the common-infrastructure groundwork here is exactly what unblocks Azure and GCP cleanly, and it's clear a lot of careful thought went into it. 🙌 A few things to flag, grouped so it's easy to work through: Heads-up on sequencing (not action for you yet): We have an internal prep PR (#287) landing on
We'll give you a heads-up the moment #287 is in so the rebase is as painless as possible. Worth addressing before merge:
Small scope question: the spot-placement files ( Nice-to-haves (non-blocking): a couple of boundary tests for the placement planner ( Really appreciate the contribution — this is a solid foundation. Let us know if any of the above needs more context! |
2 participants
The common code changes from: #240
This PR contains provider agnostic fixes and improvements needed for Azure and GCP support that are not tightly tied to Azure/GCP, for example bug fixes for bugs that Azure/GCP work surfaced but are bugs regardless, a generic spot placement API that Azure implements and support for follow up context which Azure uses for VMSS cleanup.