pipekit-sdk 2.1.2__tar.gz → 7.2.0__tar.gz

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
@@ -1,6 +1,6 @@
1
1
  Metadata-Version: 2.4
2
2
  Name: pipekit-sdk
3
- Version: 2.1.2
3
+ Version: 7.2.0
4
4
  Summary: Pipekit Python SDK
5
5
  License: BSD-3-Clause
6
6
  License-File: LICENSE
@@ -97,6 +97,25 @@ pipekit.submit(w, "<cluster-name>")
97
97
  pipekit.print_logs(pipe_run.uuid)
98
98
  ```
99
99
 
100
+ ## Connecting to Pipekit
101
+
102
+ `pipekit_url` is the single base URL for the ID, Users, and UI APIs. This is correct when they sit behind one gateway, which is the default (`https://api.pipekit.io`).
103
+
104
+ When the services are reachable on separate hosts, set `id_url` and `users_url` (or the `PIPEKIT_ID_URL` and `PIPEKIT_USERS_URL` env vars). The SDK logs in against `id_url` and makes every other call against `users_url`.
105
+
106
+ ```py
107
+ pipekit = PipekitService(
108
+ username="<user>",
109
+ password="<password>",
110
+ id_url="http://id.internal:8080",
111
+ users_url="http://users.internal:8080",
112
+ )
113
+ ```
114
+
115
+ `insecure=True` (or `PIPEKIT_INSECURE=true`) skips TLS verification. Use it only for testing against a cluster with a self-signed certificate, never in production.
116
+
117
+ `timeout` (or `PIPEKIT_TIMEOUT`) sets the per-request timeout in seconds, default 10. Raise it for slow links or for the cron lifecycle calls, which can take up to the server's notification timeout to return. The log stream is exempt and stays unbounded.
118
+
100
119
  ## Managing CronWorkflows
101
120
 
102
121
  You can create, update, and delete a `CronWorkflow` from Python. The namespace
@@ -70,6 +70,25 @@ pipekit.submit(w, "<cluster-name>")
70
70
  pipekit.print_logs(pipe_run.uuid)
71
71
  ```
72
72
 
73
+ ## Connecting to Pipekit
74
+
75
+ `pipekit_url` is the single base URL for the ID, Users, and UI APIs. This is correct when they sit behind one gateway, which is the default (`https://api.pipekit.io`).
76
+
77
+ When the services are reachable on separate hosts, set `id_url` and `users_url` (or the `PIPEKIT_ID_URL` and `PIPEKIT_USERS_URL` env vars). The SDK logs in against `id_url` and makes every other call against `users_url`.
78
+
79
+ ```py
80
+ pipekit = PipekitService(
81
+ username="<user>",
82
+ password="<password>",
83
+ id_url="http://id.internal:8080",
84
+ users_url="http://users.internal:8080",
85
+ )
86
+ ```
87
+
88
+ `insecure=True` (or `PIPEKIT_INSECURE=true`) skips TLS verification. Use it only for testing against a cluster with a self-signed certificate, never in production.
89
+
90
+ `timeout` (or `PIPEKIT_TIMEOUT`) sets the per-request timeout in seconds, default 10. Raise it for slow links or for the cron lifecycle calls, which can take up to the server's notification timeout to return. The log stream is exempt and stays unbounded.
91
+
73
92
  ## Managing CronWorkflows
74
93
 
75
94
  You can create, update, and delete a `CronWorkflow` from Python. The namespace
@@ -1,6 +1,6 @@
1
1
  [tool.poetry]
2
2
  name = "pipekit-sdk"
3
- version = "2.1.2"
3
+ version = "7.2.0"
4
4
  description = "Pipekit Python SDK"
5
5
  authors = ["Pipekit <ci@pipekit.io>"]
6
6
  license = "BSD-3-Clause"
@@ -59,4 +59,13 @@ ignore = ["E501"] # ignore line too long (80 chars)
59
59
  "**/__init__.py" = ["F401", "D107"]
60
60
 
61
61
  [tool.mypy]
62
- mypy_path = "src"
62
+ mypy_path = "src"
63
+
64
+ [tool.pytest.ini_options]
65
+ # `make test` runs unit tests only. The integration suite needs a live, in-cluster
66
+ # Pipekit and runs via `make integration`, which passes tests/integration explicitly.
67
+ testpaths = ["tests/test_unit"]
68
+ markers = [
69
+ "integration: end-to-end test that needs a live, in-cluster Pipekit",
70
+ ]
71
+
@@ -9,6 +9,7 @@ from typing import Any, Generator, List, Optional, cast
9
9
 
10
10
  import jwt
11
11
  import requests
12
+ import urllib3
12
13
  from hera.workflows import CronWorkflow, Workflow
13
14
  from hera.workflows.models import CronWorkflow as CronWorkflowModel
14
15
  from hera.workflows.models import UpdateCronWorkflowRequest, WorkflowCreateRequest
@@ -19,6 +20,23 @@ from pipekit_sdk.models.model import Cluster, Logs, Pipe, PipeRun
19
20
 
20
21
  _CONTENT_TYPE_JSON = "application/json"
21
22
 
23
+ _SSE_DATA_PREFIX = "data: "
24
+ _SSE_ID_PREFIX = "id: "
25
+ _SSE_PING = "ping"
26
+ _SSE_CLOSE = "close"
27
+
28
+
29
+ class PipekitAPIError(Exception):
30
+ """Raised when a Pipekit API request returns a non-2xx response."""
31
+
32
+
33
+ def _strip_trailing_slash(url: str) -> str:
34
+ return url[:-1] if url.endswith("/") else url
35
+
36
+
37
+ def _env_is_truthy(name: str) -> bool:
38
+ return os.getenv(name, "").strip().lower() in {"1", "true", "yes", "on"}
39
+
22
40
 
23
41
  class _LoginResponse(BaseModel):
24
42
  """LoginResponse is a model for Login responses."""
@@ -48,6 +66,15 @@ class _ValidationResult(Enum):
48
66
  EXPIRED_TOKEN = 2
49
67
 
50
68
 
69
+ class _SSELineKind(Enum):
70
+ """The meaning of one line in the log event stream."""
71
+
72
+ IGNORE = 0
73
+ ID = 1
74
+ LOG = 2
75
+ CLOSE = 3
76
+
77
+
51
78
  class PipekitService:
52
79
  """PipekitService is a wrapper around the Pipekit API. It provides a simple interface to interact with Pipekit."""
53
80
 
@@ -59,7 +86,8 @@ class PipekitService:
59
86
  with open(pipekit_token_file_path, "r") as f:
60
87
  raw_data = f.read()
61
88
  return _LoginResponse.model_validate_json(raw_data)
62
- except ValidationError:
89
+ except (FileNotFoundError, ValidationError):
90
+ # No saved token, or it is unreadable. Fall back to username/password login.
63
91
  return None
64
92
 
65
93
  @staticmethod
@@ -76,18 +104,42 @@ class PipekitService:
76
104
  password: str = "",
77
105
  token: str = "",
78
106
  pipekit_url: str = "https://api.pipekit.io",
107
+ id_url: str = "",
108
+ users_url: str = "",
109
+ insecure: bool = False,
110
+ timeout: float = 10,
79
111
  ):
80
- """Initialise the service connection with the given username and password."""
112
+ """Initialise the service connection.
113
+
114
+ pipekit_url is the single base URL for the ID, Users, and UI APIs, which is right
115
+ when they sit behind one gateway (the default, api.pipekit.io). When the services
116
+ are reachable on separate hosts, set id_url and users_url, or the PIPEKIT_ID_URL
117
+ and PIPEKIT_USERS_URL env vars, to point at each one.
118
+
119
+ insecure, or the PIPEKIT_INSECURE env var, skips TLS verification. Use it only for
120
+ testing against a cluster with a self-signed certificate.
121
+
122
+ timeout, or the PIPEKIT_TIMEOUT env var, sets the per-request timeout in seconds
123
+ (default 10). Raise it for slow links or for the cron lifecycle calls, which can
124
+ take up to the server's notification timeout to return. The log stream is exempt
125
+ and stays unbounded.
126
+ """
81
127
  self.username = username
82
128
  self.password = password
83
129
 
84
- pipekit_url = os.getenv("PIPEKIT_URL", pipekit_url)
85
- if pipekit_url.endswith("/"):
86
- pipekit_url = pipekit_url[:-1]
130
+ base_url = _strip_trailing_slash(os.getenv("PIPEKIT_URL", pipekit_url))
131
+ self.id_url = _strip_trailing_slash(os.getenv("PIPEKIT_ID_URL", id_url) or base_url)
132
+ self.users_url = _strip_trailing_slash(os.getenv("PIPEKIT_USERS_URL", users_url) or base_url)
133
+ self.ui_url = base_url
134
+
135
+ self.verify = not (insecure or _env_is_truthy("PIPEKIT_INSECURE"))
136
+ self._session = requests.Session()
137
+ self._session.verify = self.verify
138
+ if not self.verify:
139
+ # The caller opted out of TLS verification, so silence the per-request warning.
140
+ urllib3.disable_warnings(urllib3.exceptions.InsecureRequestWarning)
87
141
 
88
- self.id_url = pipekit_url
89
- self.users_url = pipekit_url
90
- self.ui_url = pipekit_url
142
+ self.timeout = float(os.getenv("PIPEKIT_TIMEOUT") or timeout)
91
143
 
92
144
  if token != "":
93
145
  self.access_token = token
@@ -106,14 +158,14 @@ class PipekitService:
106
158
  login_url = f"{self.id_url}/api/id/v1/oauth/token"
107
159
  assert self.token_data is not None, "Token data was None, this implies a bug"
108
160
  login_request = {"grant_type": "refresh_token", "refresh_token": self.token_data.refresh_token}
109
- login_response = requests.post(
161
+ login_response = self._session.post(
110
162
  login_url,
111
163
  data=json.dumps(login_request),
112
164
  headers={"Content-Type": _CONTENT_TYPE_JSON},
113
- timeout=10,
165
+ timeout=self.timeout,
114
166
  )
115
167
  if login_response.status_code // 100 != 2:
116
- raise Exception(
168
+ raise PipekitAPIError(
117
169
  "Couldn't refresh access token, obtained non 2xx status code, try logging in again via the cli"
118
170
  )
119
171
 
@@ -124,7 +176,9 @@ class PipekitService:
124
176
  def __validate_token(self) -> _ValidationResult:
125
177
  assert self.token_data is not None
126
178
  try:
127
- decoded_payload = jwt.decode(self.token_data.access_token, options={"verify_signature": False})
179
+ # We read only the exp claim to decide whether to refresh. The client has no
180
+ # key to verify the signature, and the server verifies it on every request.
181
+ decoded_payload = jwt.decode(self.token_data.access_token, options={"verify_signature": False}) # NOSONAR
128
182
  acc_tk = _AccessToken.model_validate(decoded_payload)
129
183
  exp_time = datetime.fromtimestamp(acc_tk.exp)
130
184
  now_time = datetime.now()
@@ -159,15 +213,15 @@ class PipekitService:
159
213
  login_request["username"] = ""
160
214
 
161
215
  login_url = f"{self.id_url}/api/id/v1/oauth/token"
162
- login_response = requests.post(
216
+ login_response = self._session.post(
163
217
  login_url,
164
218
  data=json.dumps(login_request),
165
219
  headers={"Content-Type": _CONTENT_TYPE_JSON},
166
- timeout=10,
220
+ timeout=self.timeout,
167
221
  )
168
222
 
169
223
  if login_response.status_code // 100 != 2:
170
- raise Exception(f"Login failed: {login_response.text}")
224
+ raise PipekitAPIError(f"Login failed: {login_response.text}")
171
225
 
172
226
  self.access_token = cast(str, login_response.json()["access_token"])
173
227
 
@@ -178,14 +232,14 @@ class PipekitService:
178
232
  def list_clusters(self) -> list[Cluster]:
179
233
  """Get a list of all clusters."""
180
234
  clusters_url = f"{self.users_url}/api/users/v1/clusters"
181
- clusters_response = requests.get(
235
+ clusters_response = self._session.get(
182
236
  clusters_url,
183
237
  headers={"Authorization": f"Bearer {self.access_token}"},
184
- timeout=10,
238
+ timeout=self.timeout,
185
239
  )
186
240
 
187
241
  if clusters_response.status_code // 100 != 2:
188
- raise Exception(f"List clusters failed: {clusters_response.text}")
242
+ raise PipekitAPIError(f"List clusters failed: {clusters_response.text}")
189
243
 
190
244
  clusters = [Cluster.model_validate(cluster) for cluster in clusters_response.json()]
191
245
  return clusters
@@ -205,14 +259,14 @@ class PipekitService:
205
259
  pipes_url += f"/clusters/{cluster_uuid}"
206
260
  pipes_url += f"/pipes?page={page}&limit={limit}"
207
261
 
208
- pipes_response = requests.get(
262
+ pipes_response = self._session.get(
209
263
  pipes_url,
210
264
  headers={"Authorization": f"Bearer {self.access_token}"},
211
- timeout=10,
265
+ timeout=self.timeout,
212
266
  )
213
267
 
214
268
  if pipes_response.status_code // 100 != 2:
215
- raise Exception(f"List pipes failed: {pipes_response.text}")
269
+ raise PipekitAPIError(f"List pipes failed: {pipes_response.text}")
216
270
 
217
271
  pipes_page = [Pipe.model_validate(item) for item in pipes_response.json()["items"]]
218
272
  if len(pipes_page) == 0:
@@ -233,14 +287,14 @@ class PipekitService:
233
287
  ) -> Pipe:
234
288
  """Get a single pipe using its uuid."""
235
289
  pipe_url = f"{self.users_url}/api/users/v1/pipes/{pipe_uuid}"
236
- pipe_response = requests.get(
290
+ pipe_response = self._session.get(
237
291
  pipe_url,
238
292
  headers={"Authorization": f"Bearer {self.access_token}"},
239
- timeout=10,
293
+ timeout=self.timeout,
240
294
  )
241
295
 
242
296
  if pipe_response.status_code // 100 != 2:
243
- raise Exception(f"Get pipe failed: {pipe_response.text}")
297
+ raise PipekitAPIError(f"Get pipe failed: {pipe_response.text}")
244
298
 
245
299
  return Pipe.model_validate(pipe_response.json())
246
300
 
@@ -258,14 +312,14 @@ class PipekitService:
258
312
  while True:
259
313
  runs_url = f"{base_url}?page={page}&limit={limit}&status={status}"
260
314
 
261
- runs_response = requests.get(
315
+ runs_response = self._session.get(
262
316
  runs_url,
263
317
  headers={"Authorization": f"Bearer {self.access_token}"},
264
- timeout=10,
318
+ timeout=self.timeout,
265
319
  )
266
320
 
267
321
  if runs_response.status_code // 100 != 2:
268
- raise Exception(f"List runs failed: {runs_response.text}")
322
+ raise PipekitAPIError(f"List runs failed: {runs_response.text}")
269
323
 
270
324
  runs_page = [PipeRun.model_validate(item) for item in runs_response.json()["items"]]
271
325
  if len(runs_page) == 0:
@@ -294,14 +348,14 @@ class PipekitService:
294
348
  while True:
295
349
  runs_url = f"{base_url}?page={page}&limit={limit}&status={status}"
296
350
 
297
- runs_response = requests.get(
351
+ runs_response = self._session.get(
298
352
  runs_url,
299
353
  headers={"Authorization": f"Bearer {self.access_token}"},
300
- timeout=10,
354
+ timeout=self.timeout,
301
355
  )
302
356
 
303
357
  if runs_response.status_code // 100 != 2:
304
- raise Exception(f"List runs failed: {runs_response.text}")
358
+ raise PipekitAPIError(f"List runs failed: {runs_response.text}")
305
359
 
306
360
  runs_page = [PipeRun.model_validate(item) for item in runs_response.json()["items"]]
307
361
  if len(runs_page) == 0:
@@ -322,14 +376,14 @@ class PipekitService:
322
376
  ) -> PipeRun:
323
377
  """Get a PipeRun by its uuid."""
324
378
  run_url = f"{self.users_url}/api/users/v1/runs/{run_uuid}"
325
- run_response = requests.get(
379
+ run_response = self._session.get(
326
380
  run_url,
327
381
  headers={"Authorization": f"Bearer {self.access_token}"},
328
- timeout=10,
382
+ timeout=self.timeout,
329
383
  )
330
384
 
331
385
  if run_response.status_code // 100 != 2:
332
- raise Exception(f"Get run failed: {run_response.text}")
386
+ raise PipekitAPIError(f"Get run failed: {run_response.text}")
333
387
 
334
388
  return PipeRun.model_validate(run_response.json())
335
389
 
@@ -339,14 +393,14 @@ class PipekitService:
339
393
  ) -> Cluster:
340
394
  """Get a single Cluster by its uuid."""
341
395
  cluster_url = f"{self.users_url}/api/users/v1/clusters/{cluster_uuid}"
342
- cluster_response = requests.get(
396
+ cluster_response = self._session.get(
343
397
  cluster_url,
344
398
  headers={"Authorization": f"Bearer {self.access_token}"},
345
- timeout=10,
399
+ timeout=self.timeout,
346
400
  )
347
401
 
348
402
  if cluster_response.status_code // 100 != 2:
349
- raise Exception(f"Get cluster failed: {cluster_response.text}")
403
+ raise PipekitAPIError(f"Get cluster failed: {cluster_response.text}")
350
404
 
351
405
  return Cluster.model_validate(cluster_response.json())
352
406
 
@@ -386,7 +440,7 @@ class PipekitService:
386
440
  model_workflow = WorkflowCreateRequest(workflow=workflow.build()).workflow # type: ignore
387
441
  assert model_workflow is not None
388
442
 
389
- submit_response = requests.post(
443
+ submit_response = self._session.post(
390
444
  submit_url,
391
445
  data=model_workflow.json(
392
446
  exclude_none=True,
@@ -398,11 +452,11 @@ class PipekitService:
398
452
  "Authorization": f"Bearer {self.access_token}",
399
453
  "Content-Type": _CONTENT_TYPE_JSON,
400
454
  },
401
- timeout=10,
455
+ timeout=self.timeout,
402
456
  )
403
457
 
404
458
  if submit_response.status_code // 100 != 2:
405
- raise Exception(f"Submit failed: {submit_response.text}")
459
+ raise PipekitAPIError(f"Submit failed: {submit_response.text}")
406
460
 
407
461
  return PipeRun.model_validate(submit_response.json())
408
462
 
@@ -417,7 +471,7 @@ class PipekitService:
417
471
  model_workflow = cron_workflow.build()
418
472
  assert model_workflow is not None
419
473
 
420
- create_response = requests.post(
474
+ create_response = self._session.post(
421
475
  create_url,
422
476
  data=model_workflow.json(
423
477
  exclude_none=True,
@@ -429,11 +483,11 @@ class PipekitService:
429
483
  "Authorization": f"Bearer {self.access_token}",
430
484
  "Content-Type": _CONTENT_TYPE_JSON,
431
485
  },
432
- timeout=10,
486
+ timeout=self.timeout,
433
487
  )
434
488
 
435
489
  if create_response.status_code // 100 != 2:
436
- raise Exception(f"Create failed: {create_response.text}")
490
+ raise PipekitAPIError(f"Create failed: {create_response.text}")
437
491
 
438
492
  return PipeRun.model_validate(create_response.json())
439
493
 
@@ -460,7 +514,7 @@ class PipekitService:
460
514
  if data is not None:
461
515
  headers["Content-Type"] = _CONTENT_TYPE_JSON
462
516
 
463
- response = requests.request(method, url, data=data, headers=headers, timeout=10)
517
+ response = self._session.request(method, url, data=data, headers=headers, timeout=self.timeout)
464
518
  if response.status_code // 100 != 2:
465
519
  raise RuntimeError(f"{action} cron failed: {response.text}")
466
520
 
@@ -543,7 +597,10 @@ class PipekitService:
543
597
  The namespace must match the CronWorkflow manifest namespace.
544
598
  """
545
599
  url = self.__cron_by_name(cluster, namespace, name) + "/suspend"
546
- response = self.__cron_request("PUT", url, "Suspend")
600
+ # Argo's suspend endpoint expects a request body. Sending it also sets the JSON
601
+ # Content-Type; a bodyless PUT is rejected with 415 Unsupported Media Type.
602
+ body = json.dumps({"name": name, "namespace": namespace})
603
+ response = self.__cron_request("PUT", url, "Suspend", data=body)
547
604
  return CronWorkflowModel.parse_raw(response.text)
548
605
 
549
606
  def resume_cron(
@@ -557,7 +614,10 @@ class PipekitService:
557
614
  The namespace must match the CronWorkflow manifest namespace.
558
615
  """
559
616
  url = self.__cron_by_name(cluster, namespace, name) + "/resume"
560
- response = self.__cron_request("PUT", url, "Resume")
617
+ # Argo's resume endpoint expects a request body. Sending it also sets the JSON
618
+ # Content-Type; a bodyless PUT is rejected with 415 Unsupported Media Type.
619
+ body = json.dumps({"name": name, "namespace": namespace})
620
+ response = self.__cron_request("PUT", url, "Resume", data=body)
561
621
  return CronWorkflowModel.parse_raw(response.text)
562
622
 
563
623
  def __apply_run_action(
@@ -566,17 +626,17 @@ class PipekitService:
566
626
  action: str,
567
627
  ) -> PipeRun:
568
628
  action_url = f"{self.users_url}/api/users/v1/runs/{run_uuid}/workflows/{action}"
569
- action_response = requests.put(
629
+ action_response = self._session.put(
570
630
  action_url,
571
631
  headers={
572
632
  "Authorization": f"Bearer {self.access_token}",
573
633
  "Content-Type": _CONTENT_TYPE_JSON,
574
634
  },
575
- timeout=10,
635
+ timeout=self.timeout,
576
636
  )
577
637
 
578
638
  if action_response.status_code // 100 != 2:
579
- raise Exception(f"Action failed: {action_response.text}")
639
+ raise PipekitAPIError(f"Action failed: {action_response.text}")
580
640
 
581
641
  return PipeRun.model_validate(action_response.json())
582
642
 
@@ -611,17 +671,17 @@ class PipekitService:
611
671
 
612
672
  # empty "" / [] must be omitted, not sent as ?pod-name= (the API filters to a pod named "")
613
673
  params = {k: v for k, v in {"pod-name": pod_name, "container-name": container_name}.items() if v}
614
- logs_response = requests.get(
674
+ logs_response = self._session.get(
615
675
  logs_url,
616
676
  params=params,
617
677
  headers={"Authorization": f"Bearer {self.access_token}"},
618
- timeout=10,
678
+ timeout=self.timeout,
619
679
  )
620
680
 
621
681
  if logs_response.status_code // 100 != 2:
622
- raise Exception(f"Get logs failed: {logs_response.text}")
682
+ raise PipekitAPIError(f"Get logs failed: {logs_response.text}")
623
683
 
624
- return [Logs.model_validate(log) for log in list(logs_response.json())]
684
+ return [Logs.model_validate(log) for log in logs_response.json()]
625
685
 
626
686
  def __get_logs_stream(
627
687
  self,
@@ -634,7 +694,7 @@ class PipekitService:
634
694
 
635
695
  # empty "" / [] must be omitted, not sent as ?pod-name= (the API filters to a pod named "")
636
696
  params = {k: v for k, v in {"pod-name": pod_name, "container-name": container_name}.items() if v}
637
- logs_response = requests.get(
697
+ logs_response = self._session.get(
638
698
  logs_url,
639
699
  params=params,
640
700
  stream=True,
@@ -648,7 +708,7 @@ class PipekitService:
648
708
  )
649
709
 
650
710
  if logs_response.status_code // 100 != 2:
651
- raise Exception(f"Get logs stream failed: {logs_response.text}")
711
+ raise PipekitAPIError(f"Get logs stream failed: {logs_response.text}")
652
712
 
653
713
  return logs_response
654
714
 
@@ -661,6 +721,26 @@ class PipekitService:
661
721
  """Get logs for a given run."""
662
722
  return self.__get_run_container_logs(run_uuid, pod_name, container_name)
663
723
 
724
+ @staticmethod
725
+ def __parse_sse_line(line: str) -> tuple[_SSELineKind, Any]:
726
+ # Classify one Server-Sent Events line from the log stream. The value is the
727
+ # new last-event-id for ID, a parsed Logs for LOG, and None otherwise.
728
+ line = line.strip()
729
+ if line == "":
730
+ return _SSELineKind.IGNORE, None
731
+ if line.startswith(_SSE_ID_PREFIX):
732
+ return _SSELineKind.ID, line[len(_SSE_ID_PREFIX) :]
733
+ if not line.startswith(_SSE_DATA_PREFIX):
734
+ print(f"unknown message: {line}")
735
+ return _SSELineKind.IGNORE, None
736
+
737
+ payload = line[len(_SSE_DATA_PREFIX) :]
738
+ if payload == _SSE_PING:
739
+ return _SSELineKind.IGNORE, None
740
+ if payload == _SSE_CLOSE:
741
+ return _SSELineKind.CLOSE, None
742
+ return _SSELineKind.LOG, Logs.model_validate(json.loads(payload))
743
+
664
744
  def follow_logs(
665
745
  self,
666
746
  run_uuid: str,
@@ -668,39 +748,19 @@ class PipekitService:
668
748
  container_name: List[str] | str = "",
669
749
  ) -> Generator[Logs, Any, None]:
670
750
  """Return a generator that yields pod logs for a given run."""
671
- data_msg_prefix = "data: "
672
- id_msg_prefix = "id: "
673
- ping_msg = "ping"
674
- close_msg = "close"
675
751
  last_event_id = ""
676
752
 
677
753
  while True:
678
754
  try:
679
755
  logs_stream = self.__get_logs_stream(run_uuid, pod_name, container_name, last_event_id)
680
756
  for line in logs_stream.iter_lines(delimiter="\n", decode_unicode=True):
681
- line = line.strip()
682
-
683
- if line == "":
684
- continue
685
-
686
- if line.startswith(data_msg_prefix):
687
- # strip the prefix
688
- line = line[len(data_msg_prefix) :]
689
- elif line.startswith(id_msg_prefix):
690
- last_event_id = line[len(id_msg_prefix) :]
691
- continue
692
- else:
693
- print(f"unknown message: {line}")
694
- continue
695
-
696
- if line == ping_msg:
697
- # pong
698
- continue
699
- elif line == close_msg:
700
- # bye
757
+ kind, value = self.__parse_sse_line(line)
758
+ if kind == _SSELineKind.ID:
759
+ last_event_id = value
760
+ elif kind == _SSELineKind.LOG:
761
+ yield value
762
+ elif kind == _SSELineKind.CLOSE:
701
763
  return
702
- else:
703
- yield Logs.model_validate(json.loads(line))
704
764
 
705
765
  except requests.exceptions.Timeout:
706
766
  pass # we'll ignore timeout errors and reconnect
File without changes