replicate-python/replicate/client.py at main · hatgit/replicate-python · GitHub

History

358 lines (293 loc) · 10.7 KB

Raw

100

101

102

103

104

105

106

107

108

109

110

111

112

113

114

115

116

117

118

119

120

121

122

123

124

125

126

127

128

129

130

131

132

133

134

135

136

137

138

139

140

141

142

143

144

145

146

147

148

149

150

151

152

153

154

155

156

157

158

159

160

161

162

163

164

165

166

167

168

169

170

171

172

173

174

175

176

177

178

179

180

181

182

183

184

185

186

187

188

189

190

191

192

193

194

195

196

197

198

199

200

201

202

203

204

205

206

207

208

209

210

211

212

213

214

215

216

217

218

219

220

221

222

223

224

225

226

227

228

229

230

231

232

233

234

235

236

237

238

239

240

241

242

243

244

245

246

247

248

249

250

251

252

253

254

255

256

257

258

259

260

261

262

263

264

265

266

267

268

269

270

271

272

273

274

275

276

277

278

279

280

281

282

283

284

285

286

287

288

289

290

291

292

293

294

295

296

297

298

299

300

301

302

303

304

305

306

307

308

309

310

311

312

313

314

315

316

317

318

319

320

321

322

323

324

325

326

327

328

329

330

331

332

333

334

335

336

337

338

339

340

341

342

343

344

345

346

347

348

349

350

351

352

353

354

355

356

357

358

import os

import random

import time

from datetime import datetime

from typing import (

TYPE_CHECKING,

Any,

AsyncIterator,

Dict,

Iterable,

Iterator,

Mapping,

Optional,

Type,

Union,

)

import httpx

from typing_extensions import Unpack

from replicate.__about__ import __version__

from replicate.collection import Collections

from replicate.deployment import Deployments

from replicate.exceptions import ReplicateError

from replicate.hardware import HardwareNamespace as Hardware

from replicate.model import Models

from replicate.prediction import Predictions

from replicate.run import async_run, run

from replicate.stream import async_stream, stream

from replicate.training import Trainings

if TYPE_CHECKING:

from replicate.stream import ServerSentEvent

class Client:

"""A Replicate API client library"""

__client: Optional[httpx.Client] = None

__async_client: Optional[httpx.AsyncClient] = None

def __init__(

self,

api_token: Optional[str] = None,

base_url: Optional[str] = None,

timeout: Optional[httpx.Timeout] = None,

**kwargs,

) -> None:

super().__init__()

self._api_token = api_token

self._base_url = base_url

self._timeout = timeout

self._client_kwargs = kwargs

self.poll_interval = float(os.environ.get("REPLICATE_POLL_INTERVAL", "0.5"))

@property

def _client(self) -> httpx.Client:

if not self.__client:

self.__client = _build_httpx_client(

httpx.Client,

self._api_token,

self._base_url,

self._timeout,

**self._client_kwargs,

) # type: ignore[assignment]

return self.__client # type: ignore[return-value]

@property

def _async_client(self) -> httpx.AsyncClient:

if not self.__async_client:

self.__async_client = _build_httpx_client(

httpx.AsyncClient,

self._api_token,

self._base_url,

self._timeout,

**self._client_kwargs,

) # type: ignore[assignment]

return self.__async_client # type: ignore[return-value]

def _request(self, method: str, path: str, **kwargs) -> httpx.Response:

resp = self._client.request(method, path, **kwargs)

_raise_for_status(resp)

return resp

async def _async_request(self, method: str, path: str, **kwargs) -> httpx.Response:

resp = await self._async_client.request(method, path, **kwargs)

_raise_for_status(resp)

return resp

@property

def collections(self) -> Collections:

"""

Namespace for operations related to collections of models.

"""

return Collections(client=self)

@property

def deployments(self) -> Deployments:

"""

Namespace for operations related to deployments.

"""

return Deployments(client=self)

@property

def hardware(self) -> Hardware:

"""

Namespace for operations related to hardware.

"""

return Hardware(client=self)

@property

def models(self) -> Models:

"""

Namespace for operations related to models.

"""

return Models(client=self)

@property

def predictions(self) -> Predictions:

"""

Namespace for operations related to predictions.

"""

return Predictions(client=self)

@property

def trainings(self) -> Trainings:

"""

Namespace for operations related to trainings.

"""

return Trainings(client=self)

def run(

self,

ref: str,

input: Optional[Dict[str, Any]] = None,

**params: Unpack["Predictions.CreatePredictionParams"],

) -> Union[Any, Iterator[Any]]: # noqa: ANN401

"""

Run a model and wait for its output.

"""

return run(self, ref, input, **params)

async def async_run(

self,

ref: str,

input: Optional[Dict[str, Any]] = None,

**params: Unpack["Predictions.CreatePredictionParams"],

) -> Union[Any, Iterator[Any]]: # noqa: ANN401

"""

Run a model and wait for its output asynchronously.

"""

return await async_run(self, ref, input, **params)

def stream(

self,

ref: str,

input: Optional[Dict[str, Any]] = None,

**params: Unpack["Predictions.CreatePredictionParams"],

) -> Iterator["ServerSentEvent"]:

"""

Stream a model's output.

"""

return stream(self, ref, input, **params)

async def async_stream(

self,

ref: str,

input: Optional[Dict[str, Any]] = None,

**params: Unpack["Predictions.CreatePredictionParams"],

) -> AsyncIterator["ServerSentEvent"]:

"""

Stream a model's output asynchronously.

"""

return async_stream(self, ref, input, **params)

# Adapted from https://github.com/encode/httpx/issues/108#issuecomment-1132753155

class RetryTransport(httpx.AsyncBaseTransport, httpx.BaseTransport):

"""A custom HTTP transport that automatically retries requests using an exponential backoff strategy

for specific HTTP status codes and request methods.

"""

RETRYABLE_METHODS = frozenset(["HEAD", "GET", "PUT", "DELETE", "OPTIONS", "TRACE"])

RETRYABLE_STATUS_CODES = frozenset(

[

429, # Too Many Requests

503, # Service Unavailable

504, # Gateway Timeout

]

)

MAX_BACKOFF_WAIT = 60

def __init__( # pylint: disable=too-many-arguments

self,

wrapped_transport: Union[httpx.BaseTransport, httpx.AsyncBaseTransport],

max_attempts: int = 10,

max_backoff_wait: float = MAX_BACKOFF_WAIT,

backoff_factor: float = 0.1,

jitter_ratio: float = 0.1,

retryable_methods: Optional[Iterable[str]] = None,

retry_status_codes: Optional[Iterable[int]] = None,

) -> None:

self._wrapped_transport = wrapped_transport

if jitter_ratio < 0 or jitter_ratio > 0.5:

raise ValueError(

f"jitter ratio should be between 0 and 0.5, actual {jitter_ratio}"

)

self.max_attempts = max_attempts

self.backoff_factor = backoff_factor

self.retryable_methods = (

frozenset(retryable_methods)

if retryable_methods

else self.RETRYABLE_METHODS

)

self.retry_status_codes = (

frozenset(retry_status_codes)

if retry_status_codes

else self.RETRYABLE_STATUS_CODES

)

self.jitter_ratio = jitter_ratio

self.max_backoff_wait = max_backoff_wait

def _calculate_sleep(

self, attempts_made: int, headers: Union[httpx.Headers, Mapping[str, str]]

) -> float:

retry_after_header = (headers.get("Retry-After") or "").strip()

if retry_after_header:

if retry_after_header.isdigit():

return float(retry_after_header)

try:

parsed_date = datetime.fromisoformat(retry_after_header).astimezone()

diff = (parsed_date - datetime.now().astimezone()).total_seconds()

if diff > 0:

return min(diff, self.max_backoff_wait)

except ValueError:

pass

backoff = self.backoff_factor * (2 ** (attempts_made - 1))

jitter = (backoff * self.jitter_ratio) * random.choice([1, -1]) # noqa: S311

total_backoff = backoff + jitter

return min(total_backoff, self.max_backoff_wait)

def handle_request(self, request: httpx.Request) -> httpx.Response:

response = self._wrapped_transport.handle_request(request) # type: ignore

if request.method not in self.retryable_methods:

return response

remaining_attempts = self.max_attempts - 1

attempts_made = 1

while True:

if (

remaining_attempts < 1

or response.status_code not in self.retry_status_codes

return response

response.close()

sleep_for = self._calculate_sleep(attempts_made, response.headers)

time.sleep(sleep_for)

response = self._wrapped_transport.handle_request(request) # type: ignore

attempts_made += 1

remaining_attempts -= 1

async def handle_async_request(self, request: httpx.Request) -> httpx.Response:

response = await self._wrapped_transport.handle_async_request(request) # type: ignore

if request.method not in self.retryable_methods:

return response

remaining_attempts = self.max_attempts - 1

attempts_made = 1

while True:

if (

remaining_attempts < 1

or response.status_code not in self.retry_status_codes

return response

response.close()

sleep_for = self._calculate_sleep(attempts_made, response.headers)

time.sleep(sleep_for)

response = await self._wrapped_transport.handle_async_request(request) # type: ignore

attempts_made += 1

remaining_attempts -= 1

async def aclose(self) -> None:

await self._wrapped_transport.aclose() # type: ignore

def close(self) -> None:

self._wrapped_transport.close() # type: ignore

def _build_httpx_client(

client_type: Type[Union[httpx.Client, httpx.AsyncClient]],

api_token: Optional[str] = None,

base_url: Optional[str] = None,

timeout: Optional[httpx.Timeout] = None,

**kwargs,

) -> Union[httpx.Client, httpx.AsyncClient]:

headers = {

"User-Agent": f"replicate-python/{__version__}",

}

if (

api_token := api_token or os.environ.get("REPLICATE_API_TOKEN")

) and api_token != "":

headers["Authorization"] = f"Token {api_token}"

base_url = (

base_url or os.environ.get("REPLICATE_BASE_URL") or "https://api.replicate.com"

)

if base_url == "":

base_url = "https://api.replicate.com"

timeout = timeout or httpx.Timeout(

5.0, read=30.0, write=30.0, connect=5.0, pool=10.0

)

transport = kwargs.pop("transport", None) or (

httpx.HTTPTransport()

if client_type is httpx.Client

else httpx.AsyncHTTPTransport()

)

return client_type(

base_url=base_url,

headers=headers,

timeout=timeout,

transport=RetryTransport(wrapped_transport=transport), # type: ignore[arg-type]

**kwargs,

)

def _raise_for_status(resp: httpx.Response) -> None:

if 400 <= resp.status_code < 600:

raise ReplicateError(resp.json()["detail"])

replicate-python/replicate/client.py at main · hatgit/replicate-python · GitHub

Expand file tree

Breadcrumbs

client.py

Latest commit

History

Breadcrumbs

client.py

File metadata and controls