Add support for hardware.list endpoint

Signed-off-by: Mattt Zmuda <mattt@replicate.com>

Author: mattt

replicate/__init__.py

@@ -2,7 +2,8 @@
 
 default_client = Client()
 run = default_client.run
+hardware = default_client.hardware
+deployments = default_client.deployments
 models = default_client.models
 predictions = default_client.predictions
 trainings = default_client.trainings
-deployments = default_client.deployments

replicate/client.py

@@ -17,6 +17,7 @@
 from replicate.__about__ import __version__
 from replicate.deployment import DeploymentCollection
 from replicate.exceptions import ModelError, ReplicateError
+from replicate.hardware import HardwareCollection
 from replicate.model import ModelCollection
 from replicate.prediction import PredictionCollection
 from replicate.schema import make_schema_backwards_compatible
@@ -83,6 +84,20 @@ def _request(self, method: str, path: str, **kwargs) -> httpx.Response:
 
         return resp
 
+    @property
+    def deployments(self) -> DeploymentCollection:
+        """
+        Namespace for operations related to deployments.
+        """
+        return DeploymentCollection(client=self)
+
+    @property
+    def hardware(self) -> HardwareCollection:
+        """
+        Namespace for operations related to hardware.
+        """
+        return HardwareCollection(client=self)
+
     @property
     def models(self) -> ModelCollection:
         """
@@ -104,13 +119,6 @@ def trainings(self) -> TrainingCollection:
         """
         return TrainingCollection(client=self)
 
-    @property
-    def deployments(self) -> DeploymentCollection:
-        """
-        Namespace for operations related to deployments.
-        """
-        return DeploymentCollection(client=self)
-
     def run(self, model_version: str, **kwargs) -> Union[Any, Iterator[Any]]:  # noqa: ANN401
         """
         Run a model and wait for its output.

replicate/hardware.py

@@ -0,0 +1,50 @@
+from typing import Dict, List, Union
+
+from replicate.base_model import BaseModel
+from replicate.collection import Collection
+
+
+class Hardware(BaseModel):
+    """
+    Hardware for running a model on Replicate.
+    """
+
+    sku: str
+    """
+    The SKU of the hardware.
+    """
+
+    name: str
+    """
+    The name of the hardware.
+    """
+
+
+class HardwareCollection(Collection):
+    """
+    Namespace for operations related to hardware.
+    """
+
+    model = Hardware
+
+    def list(self) -> List[Hardware]:
+        """
+        List all public models.
+
+        Returns:
+            A list of models.
+        """
+
+        resp = self._client._request("GET", "/v1/hardware")
+        hardware = resp.json()
+        return [self._prepare_model(obj) for obj in hardware]
+
+    def _prepare_model(self, attrs: Union[Hardware, Dict]) -> Hardware:
+        if isinstance(attrs, BaseModel):
+            attrs.id = attrs.sku
+        elif isinstance(attrs, dict):
+            attrs["id"] = attrs["sku"]
+
+        hardware = super()._prepare_model(attrs)
+
+        return hardware

tests/cassettes/hardware-list.yaml

@@ -0,0 +1,74 @@
+interactions:
+- request:
+    body: ''
+    headers:
+      accept:
+      - '*/*'
+      accept-encoding:
+      - gzip, deflate
+      connection:
+      - keep-alive
+      host:
+      - api.replicate.com
+      user-agent:
+      - replicate-python/0.15.5
+    method: GET
+    uri: https://api.replicate.com/v1/hardware
+  response:
+    content: '[{"sku":"cpu","name":"CPU"},{"sku":"gpu-t4","name":"Nvidia T4 GPU"},{"sku":"gpu-a40-small","name":"Nvidia
+      A40 GPU"},{"sku":"gpu-a40-large","name":"Nvidia A40 (Large) GPU"}]'
+    headers:
+      CF-Cache-Status:
+      - DYNAMIC
+      CF-RAY:
+      - 81fbfed29fe1c58a-SEA
+      Connection:
+      - keep-alive
+      Content-Encoding:
+      - gzip
+      Content-Type:
+      - application/json
+      Date:
+      - Thu, 02 Nov 2023 11:21:41 GMT
+      Server:
+      - cloudflare
+      Strict-Transport-Security:
+      - max-age=15552000
+      Transfer-Encoding:
+      - chunked
+      allow:
+      - OPTIONS, GET
+      content-security-policy-report-only:
+      - 'connect-src ''report-sample'' ''self'' https://replicate.delivery https://*.replicate.delivery
+        https://*.rudderlabs.com https://*.rudderstack.com https://*.mux.com https://*.sentry.io;
+        worker-src ''none''; script-src ''report-sample'' ''self'' https://cdn.rudderlabs.com/v1.1/rudder-analytics.min.js;
+        style-src ''report-sample'' ''self'' ''unsafe-inline''; font-src ''report-sample''
+        ''self'' data:; img-src ''report-sample'' ''self'' data: https://replicate.delivery
+        https://*.replicate.delivery https://*.githubusercontent.com https://github.com;
+        default-src ''self''; media-src ''report-sample'' ''self'' https://replicate.delivery
+        https://*.replicate.delivery https://*.mux.com https://*.sentry.io; report-uri'
+      cross-origin-opener-policy:
+      - same-origin
+      nel:
+      - '{"report_to":"heroku-nel","max_age":3600,"success_fraction":0.005,"failure_fraction":0.05,"response_headers":["Via"]}'
+      ratelimit-remaining:
+      - '2999'
+      ratelimit-reset:
+      - '1'
+      referrer-policy:
+      - same-origin
+      report-to:
+      - '{"group":"heroku-nel","max_age":3600,"endpoints":[{"url":"https://nel.heroku.com/reports?ts=1698924101&sid=1b10b0ff-8a76-4548-befa-353fc6c6c045&s=lMEXEYwO4dAJOZgt0b6ihblK5I4BwDBadrW6odcdYW8%3D"}]}'
+      reporting-endpoints:
+      - heroku-nel=https://nel.heroku.com/reports?ts=1698924101&sid=1b10b0ff-8a76-4548-befa-353fc6c6c045&s=lMEXEYwO4dAJOZgt0b6ihblK5I4BwDBadrW6odcdYW8%3D
+      vary:
+      - Cookie, origin
+      via:
+      - 1.1 vegur, 1.1 google
+      x-content-type-options:
+      - nosniff
+      x-frame-options:
+      - DENY
+    http_version: HTTP/1.1
+    status_code: 200
+version: 1

tests/test_hardware.py

@@ -0,0 +1,15 @@
+import httpx
+import pytest
+import respx
+
+import replicate
+
+
+@pytest.mark.vcr("hardware-list.yaml")
+@pytest.mark.asyncio
+async def test_hardware_list(mock_replicate_api_token):
+    hardware = replicate.hardware.list()
+
+    assert hardware is not None
+    assert isinstance(hardware, list)
+    assert len(hardware) > 0