From 9e82a6719a099bd6746b0d9ab882af5c2da1c21c Mon Sep 17 00:00:00 2001 From: Nathan Park Date: Tue, 5 Nov 2024 17:25:23 -0800 Subject: [PATCH 1/2] Support BatchDeleteClusterNodes from sagemaker --- src/sagemaker_core/main/resources.py | 50 +++++++++++++++++++ src/sagemaker_core/main/shapes.py | 14 ++++++ .../tools/additional_operations.json | 8 +++ 3 files changed, 72 insertions(+) diff --git a/src/sagemaker_core/main/resources.py b/src/sagemaker_core/main/resources.py index 35dd1ab8..0d1407a5 100644 --- a/src/sagemaker_core/main/resources.py +++ b/src/sagemaker_core/main/resources.py @@ -3855,6 +3855,56 @@ def update_software( response = client.update_cluster_software(**operation_input_args) logger.debug(f"Response: {response}") + @Base.add_validate_call + def batch_delete_nodes( + self, + node_ids: List[str], + session: Optional[Session] = None, + region: Optional[str] = None, + ) -> Optional[BatchDeleteClusterNodesResponse]: + """ + Deletes specific nodes within a SageMaker HyperPod cluster. + + Parameters: + node_ids: A list of node IDs to be deleted from the specified cluster. For SageMaker HyperPod clusters using the Slurm workload manager, you cannot remove instances that are configured as Slurm controller nodes. + session: Boto3 session. + region: Region name. + + Returns: + BatchDeleteClusterNodesResponse + + Raises: + botocore.exceptions.ClientError: This exception is raised for AWS service related errors. + The error message and error code can be parsed from the exception as follows: + ``` + try: + # AWS service call here + except botocore.exceptions.ClientError as e: + error_message = e.response['Error']['Message'] + error_code = e.response['Error']['Code'] + ``` + ResourceNotFound: Resource being access is not found. + """ + + operation_input_args = { + "ClusterName": self.cluster_name, + "NodeIds": node_ids, + } + # serialize the input request + operation_input_args = serialize(operation_input_args) + logger.debug(f"Serialized input request: {operation_input_args}") + + client = Base.get_sagemaker_client( + session=session, region_name=region, service_name="sagemaker" + ) + + logger.debug(f"Calling batch_delete_cluster_nodes API") + response = client.batch_delete_cluster_nodes(**operation_input_args) + logger.debug(f"Response: {response}") + + transformed_response = transform(response, "BatchDeleteClusterNodesResponse") + return BatchDeleteClusterNodesResponse(**transformed_response) + class CodeRepository(Base): """ diff --git a/src/sagemaker_core/main/shapes.py b/src/sagemaker_core/main/shapes.py index d598b489..37840c4e 100644 --- a/src/sagemaker_core/main/shapes.py +++ b/src/sagemaker_core/main/shapes.py @@ -2377,6 +2377,20 @@ class BatchDeleteClusterNodesError(Base): node_id: str +class BatchDeleteClusterNodesResponse(Base): + """ + BatchDeleteClusterNodesResponse + + Attributes + ---------------------- + failed: A list of errors encountered when deleting the specified nodes. + successful: A list of node IDs that were successfully deleted from the specified cluster. + """ + + failed: Optional[List[BatchDeleteClusterNodesError]] = Unassigned() + successful: Optional[List[str]] = Unassigned() + + class BatchDescribeModelPackageError(Base): """ BatchDescribeModelPackageError diff --git a/src/sagemaker_core/tools/additional_operations.json b/src/sagemaker_core/tools/additional_operations.json index 9c7a8f7d..85796128 100644 --- a/src/sagemaker_core/tools/additional_operations.json +++ b/src/sagemaker_core/tools/additional_operations.json @@ -23,6 +23,14 @@ "return_type": "None", "method_type": "object", "service_name": "sagemaker" + }, + "BatchDeleteClusterNodes": { + "operation_name": "BatchDeleteClusterNodes", + "resource_name": "Cluster", + "method_name": "batch_delete_nodes", + "return_type": "BatchDeleteClusterNodesResponse", + "method_type": "object", + "service_name": "sagemaker" } }, "PipelineExecution": { From cfa6e0deea1ed6dc814c042ea48833430c880fc0 Mon Sep 17 00:00:00 2001 From: Nathan Park Date: Tue, 5 Nov 2024 17:27:35 -0800 Subject: [PATCH 2/2] Update api_coverage.json counts --- src/sagemaker_core/tools/api_coverage.json | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/sagemaker_core/tools/api_coverage.json b/src/sagemaker_core/tools/api_coverage.json index beddc599..4d6fdbd0 100644 --- a/src/sagemaker_core/tools/api_coverage.json +++ b/src/sagemaker_core/tools/api_coverage.json @@ -1 +1 @@ -{"SupportedAPIs": 339, "UnsupportedAPIs": 6} \ No newline at end of file +{"SupportedAPIs": 340, "UnsupportedAPIs": 5} \ No newline at end of file