Implements address space cleanup when a workspace service is uninstalled - #4744
Implements address space cleanup when a workspace service is uninstalled#4744JC-wk wants to merge 30 commits into
Conversation
Unit Test Results732 tests 732 ✅ 10s ⏱️ Results for commit 9a08ecd. ♻️ This comment has been updated with latest results. |
|
I have been testing this for a few days, I am not sure if unit tests are needed and how best to write them if anyone wants to assist. |
There was a problem hiding this comment.
Pull request overview
This PR addresses IP range exhaustion risk by ensuring workspace address spaces allocated by workspace services are freed on successful uninstall, and by triggering a workspace upgrade so downstream infra reflects the removal.
Changes:
- Add post-uninstall cleanup in the service bus deployment status handler to remove a workspace-service
address_spacefrom the parent workspace’saddress_spaces. - Update AzureML and Databricks workspace-service templates to run a workspace
upgradestep after uninstall. - Bump API + template versions and add a changelog entry.
Reviewed changes
Copilot reviewed 7 out of 7 changed files in this pull request and generated 3 comments.
Show a summary per file
| File | Description |
|---|---|
| templates/workspace_services/databricks/template_schema.json | Adds a workspace upgrade step after uninstall (and JSON formatting changes). |
| templates/workspace_services/databricks/porter.yaml | Patch version bump. |
| templates/workspace_services/azureml/template_schema.json | Adds a workspace upgrade step after uninstall. |
| templates/workspace_services/azureml/porter.yaml | Patch version bump. |
| api_app/service_bus/deployment_status_updater.py | Implements address space cleanup after successful uninstall main step. |
| api_app/_version.py | API patch version bump. |
| CHANGELOG.md | Adds an Unreleased entry describing the change. |
There was a problem hiding this comment.
🟡 Not ready to approve
The refactor to receive_messages uses async with credentials.get_credential_async() (a coroutine) which will fail at runtime and prevent the service bus updater from running.
Once you've addressed the issues Copilot identified, you can request another Copilot review.
This review doesn't count toward merge requirements. Sign up for the private preview to control whether Copilot approvals count.
Review details
- Files reviewed: 10/10 changed files
- Comments generated: 1
- Review effort level: Lite
We're testing this review assessment. Please use 👍 or 👎 to tell us if it's correct.
There was a problem hiding this comment.
🟡 Not ready to approve
The deployment status receiver loop should ensure AutoLockRenewer is always closed (e.g., via async with/finally) to avoid leaking renewal tasks if an exception occurs mid-session.
Once you've addressed the issues Copilot identified, you can request another Copilot review.
This review doesn't count toward merge requirements. Sign up for the private preview to control whether Copilot approvals count.
Review details
Suppressed comments (1)
api_app/service_bus/deployment_status_updater.py:67
AutoLockReneweris created and closed only after theasync for msg in receiverloop completes. If an exception is raised while iterating messages or completing/abandoning a message,renewer.close()will be skipped, potentially leaking renewal tasks and leaving the session lock renewal running longer than intended. Useasync with AutoLockRenewer()(or atry/finally) so the renewer is always closed.
async with service_bus_client.get_queue_receiver(queue_name=config.SERVICE_BUS_DEPLOYMENT_STATUS_UPDATE_QUEUE, max_wait_time=config.SERVICE_BUS_MAX_WAIT_TIME, session_id=NEXT_AVAILABLE_SESSION) as receiver:
renewer = AutoLockRenewer()
renewer.register(receiver, receiver.session, max_lock_renewal_duration=60)
async for msg in receiver:
complete_message = await self.process_message(msg)
if complete_message:
await receiver.complete_message(msg)
else:
await receiver.abandon_message(msg)
await renewer.close()
polling_count = 0
- Files reviewed: 10/10 changed files
- Comments generated: 0 new
- Review effort level: Lite
We're testing this review assessment. Please use 👍 or 👎 to tell us if it's correct.
…an asynchronous context manager (async with AutoLockRenewer() as renewer:). • This ensures that renewer.close() is guaranteed to execute upon exiting the block, even if an exception occurs while iterating over messages or completing/abandoning a message.
There was a problem hiding this comment.
🟡 Not ready to approve
The updated service bus receiver loop references an undefined config value and weakens fault-tolerance around credential/client creation, which can break the updater at runtime.
Once you've addressed the issues Copilot identified, you can request another Copilot review.
This review doesn't count toward merge requirements. Sign up for the private preview to control whether Copilot approvals count.
Review details
Suppressed comments (1)
api_app/service_bus/deployment_status_updater.py:88
- This refactor removed the span attributes (
step_id,operation_id,status) that were previously attached to theprocess_messagetracing span. That makes correlating traces/messages materially harder in production; please restore these attributes (and keep the existing logging).
async def process_message(self, msg) -> bool:
complete_message = False
with tracer.start_as_current_span("process_message"):
try:
message = parse_obj_as(DeploymentStatusUpdateMessage, json.loads(str(msg)))
complete_message = await self.update_status_in_database(message)
logger.info(f"Update status in DB for {message.operationId} - {message.status}")
except (json.JSONDecodeError, ValidationError):
- Files reviewed: 10/10 changed files
- Comments generated: 2
- Review effort level: Lite
We're testing this review assessment. Please use 👍 or 👎 to tell us if it's correct.
• Captured the active tracing span context via with tracer.start_as_current_span("process_message") as current_span:.
• Set the attributes step_id, operation_id, and status on current_span after parsing the DeploymentStatusUpdateMessage.
2. **test_deployment_status_update.py**:
• Added a unit test test_process_message_sets_span_attributes to verify that set_attribute is called with the expected values.
There was a problem hiding this comment.
🟡 Not ready to approve
The deployment status receiver references config.SERVICE_BUS_MAX_WAIT_TIME, which is not defined in the API config and would cause a runtime failure.
Once you've addressed the issues Copilot identified, you can request another Copilot review.
This review doesn't count toward merge requirements. Sign up for the private preview to control whether Copilot approvals count.
Review details
Suppressed comments (2)
api_app/service_bus/deployment_status_updater.py:53
- The exception handling only wraps code after entering the
ServiceBusClientcontext. Failures while acquiring the credential context or constructing/entering theServiceBusClientcontext (e.g., credential acquisition errors) will bypass theOperationTimeoutError/ServiceBusConnectionErrorhandlers and terminatereceive_messages(). Consider moving thetry/exceptto wrap the entire loop iteration, including the credential and client context managers, to keep the updater resilient.
complete_message = True
async with credentials.get_credential_async_context() as credential:
async with ServiceBusClient(fully_qualified_namespace=config.SERVICE_BUS_FULLY_QUALIFIED_NAMESPACE, credential=credential) as service_bus_client:
try:
api_app/service_bus/deployment_status_updater.py:55
config.SERVICE_BUS_MAX_WAIT_TIMEis referenced for the Service Bus receivermax_wait_time, but that setting does not exist inapi_app/core/config.py. This will raise anAttributeErrorat runtime and stop the receiver loop.
async with service_bus_client.get_queue_receiver(queue_name=config.SERVICE_BUS_DEPLOYMENT_STATUS_UPDATE_QUEUE, max_wait_time=config.SERVICE_BUS_MAX_WAIT_TIME, session_id=NEXT_AVAILABLE_SESSION) as receiver:
- Files reviewed: 10/10 changed files
- Comments generated: 0 new
- Review effort level: Lite
We're testing this review assessment. Please use 👍 or 👎 to tell us if it's correct.
…g in DeploymentStatusUpdater
There was a problem hiding this comment.
🟡 Not ready to approve
The newly added retry/error-path unit tests use AsyncMock.side_effect with an exception class (not an exception instance), so they won’t reliably raise and are likely to fail or not exercise the intended branches.
Once you've addressed the issues Copilot identified, you can request another Copilot review.
This review doesn't count toward merge requirements. Sign up for the private preview to control whether Copilot approvals count.
Review details
Suppressed comments (2)
api_app/tests_ma/test_service_bus/test_deployment_status_update.py:710
AsyncMock.side_effect = CosmosAccessConditionFailedErrortreats the exception class as a callable and returns an exception instance rather than raising, so this test won’t drive the max-retry error path. Setside_effectto an exception instance instead.
# All attempts raise CosmosAccessConditionFailedError
workspace_repo.patch_workspace.side_effect = CosmosAccessConditionFailedError
workspace_repo_mock.return_value = workspace_repo
api_app/tests_ma/test_service_bus/test_deployment_status_update.py:635
- These mocks won’t actually raise on the first call: when
AsyncMock.side_effectis an iterable, an exception class is treated as a normal return value (not raised). Use an exception instance for the first element so the retry path is exercised reliably.
This issue also appears on line 708 of the same file.
# First attempt raises CosmosAccessConditionFailedError, second succeeds
workspace_repo.patch_workspace.side_effect = [CosmosAccessConditionFailedError, None]
workspace_repo_mock.return_value = workspace_repo
- Files reviewed: 11/11 changed files
- Comments generated: 0 new
- Review effort level: Lite
We're testing this review assessment. Please use 👍 or 👎 to tell us if it's correct.
…ck request and resource patching
There was a problem hiding this comment.
🟢 Ready to approve
Core behavior change (address space cleanup + template trailing upgrade) is implemented and covered by focused unit tests, with only a minor tracing-attribute typing improvement suggested.
This review doesn't count toward merge requirements. Sign up for the private preview to control whether Copilot approvals count.
Review details
Suppressed comments (2)
api_app/service_bus/deployment_status_updater.py:89
message.operationIdis aUUID4(UUID) and is being passed directly to OpenTelemetrySpan.set_attribute. OTEL span attributes are expected to be primitive types (e.g., str/int/bool); non-primitive values like UUIDs may be dropped (or logged as invalid), so the trace can lose the operation id. Consider stringifying UUIDs (and any enums) before setting attributes, and update the unit test accordingly.
message = parse_obj_as(DeploymentStatusUpdateMessage, json.loads(str(msg)))
current_span.set_attribute("step_id", message.stepId)
current_span.set_attribute("operation_id", message.operationId)
current_span.set_attribute("status", message.status)
api_app/tests_ma/test_service_bus/test_deployment_status_update.py:179
- This test currently asserts
operation_idis set as auuid.UUIDinstance. OpenTelemetry span attributes are expected to be primitives; ifprocess_messageis updated to setoperation_idas a string (recommended), this assertion should be updated to match.
mock_span.set_attribute.assert_any_call("step_id", test_sb_message["stepId"])
mock_span.set_attribute.assert_any_call("operation_id", uuid.UUID(test_sb_message["operationId"]))
mock_span.set_attribute.assert_any_call("status", test_sb_message["status"])
- Files reviewed: 13/13 changed files
- Comments generated: 0 new
- Review effort level: Lite
We're testing this review assessment. Please use 👍 or 👎 to tell us if it's correct.
…pdate test assertions
There was a problem hiding this comment.
🟢 Ready to approve
The address space cleanup is implemented in the appropriate post-uninstall path, templates are updated to reconcile workspace state, and the new behavior is covered by focused unit tests (including ETag-retry scenarios).
This review doesn't count toward merge requirements. Sign up for the private preview to control whether Copilot approvals count.
Review details
- Files reviewed: 13/13 changed files
- Comments generated: 0 new
- Review effort level: Lite
We're testing this review assessment. Please use 👍 or 👎 to tell us if it's correct.
There was a problem hiding this comment.
🟡 Not ready to approve
Cleanup failures are swallowed, allowing uninstall operations to succeed while address spaces remain allocated.
Once you've addressed the issues Copilot identified, you can request another Copilot review.
This review doesn't count toward merge requirements. Sign up for the private preview to control whether Copilot approvals count.
Review details
Suppressed comments (1)
api_app/service_bus/deployment_status_updater.py:55
- This substantially changes the receiver/client lifecycle, idle timeout, and complete-versus-abandon flow, but the added
receive_messagestest only exercises failure before a receiver is created. Add an async receiver test covering a processed message and asserting both settlement branches, plus the idle timeout path, so this long-running queue loop is protected from regressions.
async with ServiceBusClient(fully_qualified_namespace=config.SERVICE_BUS_FULLY_QUALIFIED_NAMESPACE, credential=credential) as service_bus_client:
logger.debug("Creating Deployment Status receiver session")
async with service_bus_client.get_queue_receiver(queue_name=config.SERVICE_BUS_DEPLOYMENT_STATUS_UPDATE_QUEUE, max_wait_time=config.SERVICE_BUS_MAX_WAIT_TIME, session_id=NEXT_AVAILABLE_SESSION) as receiver:
- Files reviewed: 13/13 changed files
- Comments generated: 1
- Review effort level: Balanced
We're testing this review assessment. Please use 👍 or 👎 to tell us if it's correct.
| except Exception as e: | ||
| logger.error(f"[ADDRESS_SPACE_CLEANUP_FAILED] Failed to free workspace address space {address_to_free} for workspace {parent_workspace_id} after uninstalling {resource_id}: {e}", exc_info=True) |
Resolves #4727
PR
What is being addressed
Currently address spaces are not cleaned up when a workspace service is uninstalled from a workspace, this could lead to ip range exhaustion.
This PR adds the functionality to delete the
address_spaceused by a workspace-service on uninstall of the service. The address range is then freed from the workspace and can be reused.How is this addressed