added filter params to get_compute_resources, simplify example

Author: tamirse

datacrunch/containers/containers.py

@@ -868,17 +868,27 @@ def delete_deployment_environment_variables(self, deployment_name: str, containe
                 env_var) for env_var in env_vars]
         return result
 
-    def get_compute_resources(self) -> List[ComputeResource]:
-        """Retrieves available compute resources.
+    def get_compute_resources(self, size: int = None, is_available: bool = None) -> List[ComputeResource]:
+        """Retrieves compute resources, optionally filtered by size and availability.
+
+        Args:
+            size: Optional size to filter resources by (e.g. 8 for 8x GPUs)
+            available: Optional boolean to filter by availability status
 
         Returns:
-            List[ComputeResource]: List of available compute resources.
+            List[ComputeResource]: List of compute resources matching the filters.
+                                 If no filters provided, returns all resources.
         """
         response = self.client.get(SERVERLESS_COMPUTE_RESOURCES_ENDPOINT)
         resources = []
         for resource_group in response.json():
             for resource in resource_group:
                 resources.append(ComputeResource.from_dict(resource))
+        if size:
+            resources = [r for r in resources if r.size == size]
+        if is_available:
+            resources = [
+                r for r in resources if r.is_available == is_available]
         return resources
 
     # Function alias

examples/containers/compute_resources_example.py

@@ -1,76 +1,29 @@
 import os
 from datacrunch import DataCrunchClient
-from typing import List
-from datacrunch.containers.containers import ComputeResource
 
 # Get client secret and id from environment variables
 DATACRUNCH_CLIENT_ID = os.environ.get('DATACRUNCH_CLIENT_ID')
 DATACRUNCH_CLIENT_SECRET = os.environ.get('DATACRUNCH_CLIENT_SECRET')
 
-
-def list_all_compute_resources(client: DataCrunchClient) -> List[ComputeResource]:
-    """List all available compute resources.
-
-    Args:
-        client (DataCrunchClient): The DataCrunch API client.
-
-    Returns:
-        List[ComputeResource]: List of all compute resources.
-    """
-    return client.containers.get_compute_resources()
-
-
-def list_available_compute_resources(client: DataCrunchClient) -> List[ComputeResource]:
-    """List only the available compute resources.
-
-    Args:
-        client (DataCrunchClient): The DataCrunch API client.
-
-    Returns:
-        List[ComputeResource]: List of available compute resources.
-    """
-    all_resources = client.containers.get_compute_resources()
-    return [r for r in all_resources if r.is_available]
-
-
-def list_compute_resources_by_size(client: DataCrunchClient, size: int) -> List[ComputeResource]:
-    """List compute resources filtered by size.
-
-    Args:
-        client (DataCrunchClient): The DataCrunch API client.
-        size (int): The size to filter by.
-
-    Returns:
-        List[ComputeResource]: List of compute resources with the specified size.
-    """
-    all_resources = client.containers.get_compute_resources()
-    return [r for r in all_resources if r.size == size]
-
-
-def main():
-    # Initialize the client with your credentials
-    datacrunch = DataCrunchClient(
-        DATACRUNCH_CLIENT_ID, DATACRUNCH_CLIENT_SECRET)
-
-    # Example 1: List all compute resources
-    print("All compute resources:")
-    all_resources = list_all_compute_resources(datacrunch)
-    for resource in all_resources:
-        print(
-            f"Name: {resource.name}, Size: {resource.size}, Available: {resource.is_available}")
-
-    # Example 2: List available compute resources
-    print("Available compute resources:")
-    available_resources = list_available_compute_resources(datacrunch)
-    for resource in available_resources:
-        print(f"Name: {resource.name}, Size: {resource.size}")
-
-    # Example 3: List compute resources of size 8
-    print("Compute resources with size 8:")
-    size_8_resources = list_compute_resources_by_size(datacrunch, 8)
-    for resource in size_8_resources:
-        print(f"Name: {resource.name}, Available: {resource.is_available}")
-
-
-if __name__ == "__main__":
-    main()
+# Initialize the client with your credentials
+datacrunch = DataCrunchClient(DATACRUNCH_CLIENT_ID, DATACRUNCH_CLIENT_SECRET)
+
+# Example 1: List all compute resources
+print("All compute resources:")
+all_resources = datacrunch.containers.get_compute_resources()
+for resource in all_resources:
+    print(
+        f"Name: {resource.name}, Size: {resource.size}, Available: {resource.is_available}")
+
+# Example 2: List available compute resources
+print("\nAvailable compute resources:")
+available_resources = datacrunch.containers.get_compute_resources(
+    is_available=True)
+for resource in available_resources:
+    print(f"Name: {resource.name}, Size: {resource.size}")
+
+# Example 3: List compute resources of size 8
+print("\nCompute resources with size 8:")
+size_8_resources = datacrunch.containers.get_compute_resources(size=8)
+for resource in size_8_resources:
+    print(f"Name: {resource.name}, Available: {resource.is_available}")