kalavai-client 0.5.30__py3-none-any.whl → 0.6.1__py3-none-any.whl

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.
kalavai_client/cli.py CHANGED
@@ -6,8 +6,6 @@ import uuid
6
6
  import time
7
7
  import socket
8
8
  from pathlib import Path
9
- from getpass import getpass
10
- from sys import exit
11
9
 
12
10
  import yaml
13
11
 
@@ -21,7 +19,6 @@ from kalavai_client.env import (
21
19
  USER_LOCAL_SERVER_FILE,
22
20
  TEMPLATE_LABEL,
23
21
  KALAVAI_PLATFORM_URL,
24
- DEFAULT_VPN_CONTAINER_NAME,
25
22
  CONTAINER_HOST_PATH,
26
23
  USER_COMPOSE_FILE,
27
24
  USER_HELM_APPS_FILE,
@@ -41,6 +38,7 @@ from kalavai_client.core import (
41
38
  fetch_devices,
42
39
  fetch_job_logs,
43
40
  fetch_gpus,
41
+ generate_worker_package,
44
42
  load_gpu_models,
45
43
  fetch_job_templates,
46
44
  fetch_job_defaults,
@@ -70,18 +68,12 @@ from kalavai_client.utils import (
70
68
  generate_table,
71
69
  request_to_server,
72
70
  safe_remove,
73
- leave_vpn,
74
71
  load_server_info,
75
- user_login,
76
- user_logout,
77
- get_public_vpns,
78
- register_cluster,
79
- unregister_cluster,
80
72
  get_public_seeds,
81
- load_user_session,
73
+ load_user_id,
82
74
  SERVER_IP_KEY,
83
- NODE_NAME_KEY,
84
- CLUSTER_NAME_KEY
75
+ CLUSTER_NAME_KEY,
76
+ KALAVAI_AUTH
85
77
  )
86
78
 
87
79
 
@@ -189,16 +181,19 @@ def select_token_type():
189
181
  break
190
182
  return {"admin": choice == 0, "user": choice == 1, "worker": choice == 2}
191
183
 
192
- def input_gpus():
184
+ def input_gpus(non_interactive=False):
193
185
  num_gpus = 0
194
186
  try:
195
187
  has_gpus = check_gpu_drivers()
196
188
  if has_gpus:
197
189
  max_gpus = int(run_cmd("nvidia-smi -L | wc -l").decode())
198
- num_gpus = user_confirm(
199
- question=f"{max_gpus} NVIDIA GPU(s) detected. How many GPUs would you like to include?",
200
- options=range(max_gpus+1)
201
- )
190
+ if non_interactive:
191
+ num_gpus = max_gpus
192
+ else:
193
+ num_gpus = user_confirm(
194
+ question=f"{max_gpus} NVIDIA GPU(s) detected. How many GPUs would you like to include?",
195
+ options=range(max_gpus+1)
196
+ )
202
197
  except:
203
198
  console.log(f"[red]WARNING: error when fetching NVIDIA GPU info. GPUs will not be used on this local machine")
204
199
  return num_gpus
@@ -208,18 +203,34 @@ def input_gpus():
208
203
  ##################
209
204
 
210
205
  @arguably.command
211
- def gui__start(*others, backend_only=False, gui_frontend_port=3000, gui_backend_port=8000, bridge_port=8001, log_level="critical"):
206
+ def gui__start(
207
+ *others,
208
+ backend_only=False,
209
+ gui_frontend_port=3000,
210
+ gui_backend_port=8000,
211
+ bridge_port=8001,
212
+ log_level="critical",
213
+ protected_access=False
214
+ ):
212
215
  """Run GUI (docker) and kalavai core backend (api)"""
213
216
  if len(set([gui_frontend_port, gui_backend_port, bridge_port])) < 3:
214
217
  console.log("[red]Error: ports must be unique")
215
218
  return
216
219
 
220
+ user_key = None
221
+ if protected_access:
222
+ user_key = load_user_id()
223
+ if user_key is None:
224
+ console.log("[red]Error: user key not found (required for protected access)")
225
+ return
226
+
217
227
  if not backend_only:
218
228
  values = {
219
229
  "gui_frontend_port": gui_frontend_port,
220
230
  "gui_backend_port": gui_backend_port,
221
231
  "bridge_port": bridge_port,
222
- "path": user_path("")
232
+ "path": user_path(""),
233
+ "protected_access": user_key
223
234
  }
224
235
  compose_yaml = load_template(
225
236
  template_path=DOCKER_COMPOSE_GUI,
@@ -239,62 +250,25 @@ def gui__start(*others, backend_only=False, gui_frontend_port=3000, gui_backend_
239
250
  run_cmd(f"docker compose --file {USER_GUI_COMPOSE_FILE} down")
240
251
  console.log("[green]Kalavai GUI has been stopped")
241
252
 
253
+
242
254
  @arguably.command
243
- def login(*others, username: str=None):
255
+ def auth(user_key, *others):
244
256
  """
245
257
  [AUTH] (For public clusters only) Log in to Kalavai server.
246
-
247
- Args:
248
- *others: all the other positional arguments go here
249
258
  """
250
- console.log(f"Kalavai account details. If you don't have an account, create one at [yellow]{KALAVAI_PLATFORM_URL}")
251
- if username is None:
252
- username = input("User email: ")
253
- password = getpass()
254
- user = user_login(
255
- user_cookie=USER_COOKIE,
256
- username=username,
257
- password=password
258
- )
259
-
260
- if user is not None:
261
- console.log(f"[green]{username} logged in successfully")
259
+ KALAVAI_AUTH.save_auth(user_key)
260
+ if KALAVAI_AUTH.is_authenticated():
261
+ console.log(f"[green]User key stored")
262
262
  else:
263
- console.log(f"[red]Invalid credentials for {username}")
264
-
265
- return user is not None
263
+ console.log(f"[red]Invalid user key")
266
264
 
267
265
  @arguably.command
268
266
  def logout(*others):
269
267
  """
270
- [AUTH] (For public clusters only) Log out of Kalavai server.
271
-
272
- Args:
273
- *others: all the other positional arguments go here
268
+ Log out of Kalavai server.
274
269
  """
275
- user_logout(
276
- user_cookie=USER_COOKIE
277
- )
278
- console.log("[green]Log out successfull")
279
-
280
- @arguably.command
281
- def location__list(*others):
282
- """
283
- [AUTH] List public locations on Kalavai
284
- """
285
- try:
286
- seeds = get_public_vpns(user_cookie=USER_COOKIE)
287
- except Exception as e:
288
- console.log(f"[red]Error: {str(e)}")
289
- console.log("Are you authenticated? Try [yellow]kalavai login")
290
- return
291
- columns, rows = [], []
292
- for idx, seed in enumerate(seeds):
293
- columns = seed.keys()
294
- rows.append([str(idx)] + list(seed.values()))
295
- columns = ["VPN"] + list(columns)
296
- table = generate_table(columns=columns, rows=rows)
297
- console.log(table)
270
+ KALAVAI_AUTH.clear_auth()
271
+ console.log(f"[green]User key removed")
298
272
 
299
273
  @arguably.command
300
274
  def pool__publish(*others, description=None, is_private=True):
@@ -305,11 +279,10 @@ def pool__publish(*others, description=None, is_private=True):
305
279
  # - cluster is up and running
306
280
  # - cluster is connected to vpn (has net token)
307
281
  # - user is authenticated
308
- try:
309
- CLUSTER.is_seed_node()
310
- except Exception as e:
311
- console.log(f"[red]Problems with your pool: {str(e)}")
282
+ if not CLUSTER.is_seed_node():
283
+ console.log(f"You can only create workers from a seed node")
312
284
  return
285
+
313
286
  choices = select_token_type()
314
287
  if choices["admin"]:
315
288
  mode = TokenType.ADMIN
@@ -343,10 +316,8 @@ def pool__unpublish(cluster_name=None, *others):
343
316
  # Check for:
344
317
  # - cluster is up and running
345
318
  # - user is authenticated
346
- try:
347
- CLUSTER.is_seed_node()
348
- except Exception as e:
349
- console.log(f"[red]Problems with your pool: {str(e)}")
319
+ if not CLUSTER.is_seed_node():
320
+ console.log(f"You can only create workers from a seed node")
350
321
  return
351
322
 
352
323
  result = unregister_pool()
@@ -357,6 +328,31 @@ def pool__unpublish(cluster_name=None, *others):
357
328
  else:
358
329
  console.log(f"[green]Your cluster has been removed from {KALAVAI_PLATFORM_URL}")
359
330
 
331
+ @arguably.command
332
+ def pool__package_worker(output_file, *others, num_gpus=0, ip_address="0.0.0.0", node_name=None, storage_compatible=True):
333
+ """
334
+ [AUTH]Package a worker for distribution (docker compose only)
335
+ """
336
+
337
+ if not CLUSTER.is_seed_node():
338
+ console.log(f"[red]You can only create workers from a seed node")
339
+ return
340
+
341
+ compose = generate_worker_package(
342
+ num_gpus=num_gpus,
343
+ ip_address=ip_address,
344
+ node_name=node_name,
345
+ storage_compatible=storage_compatible
346
+ )
347
+
348
+ if "error" in compose:
349
+ console.log(f"[red]{compose['error']}")
350
+ else:
351
+ console.log(f"[green]Worker package created: {output_file}")
352
+ with open(output_file, "w") as f:
353
+ f.write(compose)
354
+
355
+
360
356
  @arguably.command
361
357
  def pool__list(*others, user_only=False):
362
358
  """
@@ -382,7 +378,7 @@ def pool__list(*others, user_only=False):
382
378
 
383
379
 
384
380
  @arguably.command
385
- def pool__start(cluster_name, *others, only_registered_users: bool=False, ip_address: str=None, location: str=None, app_values: str=None, pool_config_values: str=None):
381
+ def pool__start(cluster_name, *others, ip_address: str=None, location: str=None, app_values: str=None, pool_config_values: str=None, non_interactive: bool=False):
386
382
  """
387
383
  Start Kalavai pool and start/resume sharing resources.
388
384
 
@@ -395,19 +391,24 @@ def pool__start(cluster_name, *others, only_registered_users: bool=False, ip_ad
395
391
  return
396
392
 
397
393
  # User acknowledgement
398
- option = user_confirm(
399
- question="Kalavai will now create a pool and a local worker using docker. This won't modify your system. Are you happy to proceed?",
400
- options=["no", "yes"]
401
- )
402
- if option == 0:
403
- console.log("Installation was cancelled and did not complete.")
404
- return
394
+ if not non_interactive:
395
+ option = user_confirm(
396
+ question="Kalavai will now create a pool and a local worker using docker. This won't modify your system. Are you happy to proceed?",
397
+ options=["no", "yes"]
398
+ )
399
+ if option == 0:
400
+ console.log("Installation was cancelled and did not complete.")
401
+ return
405
402
 
406
403
  # select IP address (for external discovery)
407
404
  if ip_address is None and location is None:
408
- # local IP
409
- console.log(f"Scanning for valid IPs")
410
- ip_address = select_ip_address()
405
+ if non_interactive:
406
+ ip_address = "0.0.0.0"
407
+ console.log("[yellow]Using [green]0.0.0.0 [yellow]for server address")
408
+ else:
409
+ # local IP
410
+ console.log(f"Scanning for valid IPs")
411
+ ip_address = select_ip_address()
411
412
 
412
413
  console.log(f"Using {ip_address} address for server")
413
414
 
@@ -418,8 +419,7 @@ def pool__start(cluster_name, *others, only_registered_users: bool=False, ip_ad
418
419
  ip_address=ip_address,
419
420
  app_values=app_values,
420
421
  pool_config_values=pool_config_values,
421
- num_gpus=input_gpus(),
422
- only_registered_users=only_registered_users,
422
+ num_gpus=input_gpus(non_interactive=non_interactive),
423
423
  location=location
424
424
  )
425
425
 
@@ -476,7 +476,7 @@ def pool__check_token(token, *others, public=False):
476
476
  return True
477
477
 
478
478
  @arguably.command
479
- def pool__join(token, *others, node_name=None):
479
+ def pool__join(token, *others, node_name=None, auto_accept=False):
480
480
  """
481
481
  Join Kalavai pool and start/resume sharing resources.
482
482
 
@@ -491,28 +491,38 @@ def pool__join(token, *others, node_name=None):
491
491
  return
492
492
 
493
493
  # check that is not attached to another instance
494
- if os.path.exists(USER_LOCAL_SERVER_FILE):
494
+ if not auto_accept:
495
+ if os.path.exists(USER_LOCAL_SERVER_FILE):
496
+ option = user_confirm(
497
+ question="You seem to be connected to an instance already. Are you sure you want to join a new one?",
498
+ options=["no", "yes"]
499
+ )
500
+ if option == 0:
501
+ console.log("[green]Nothing happened.")
502
+ return
503
+
504
+ user_id = load_user_id()
505
+ if user_id is None:
506
+ console.log("You are not authenticated. If you want to authenticate your node, use [yellow]kalavai auth <user_key>")
507
+
508
+ num_gpus = input_gpus(auto_accept=auto_accept)
509
+
510
+ if not auto_accept:
495
511
  option = user_confirm(
496
- question="You seem to be connected to an instance already. Are you sure you want to join a new one?",
512
+ question="Docker compose ready. Would you like Kalavai to deploy it?",
497
513
  options=["no", "yes"]
498
514
  )
499
515
  if option == 0:
500
- console.log("[green]Nothing happened.")
516
+ console.log("[red]Installation aborted")
501
517
  return
502
518
 
503
- num_gpus = input_gpus()
504
-
505
- option = user_confirm(
506
- question="Docker compose ready. Would you like Kalavai to deploy it?",
507
- options=["no", "yes"]
508
- )
509
- if option == 0:
510
- console.log("[red]Installation aborted")
511
- return
512
-
513
519
  # select IP address (for external discovery)
514
520
  console.log(f"Scanning for valid IPs")
515
- ip_address = select_ip_address()
521
+ if auto_accept:
522
+ ip_address = "0.0.0.0"
523
+ console.log("[yellow]Using [green]0.0.0.0 [yellow]for server address")
524
+ else:
525
+ ip_address = select_ip_address()
516
526
 
517
527
  console.log("Connecting worker to the pool...")
518
528
  result = join_pool(
@@ -803,8 +813,6 @@ def storage__list(*other):
803
813
  return
804
814
 
805
815
  try:
806
- user = load_user_session(user_cookie=USER_COOKIE)
807
- username = user["username"] if user is not None else None
808
816
  result = request_to_server(
809
817
  method="post",
810
818
  endpoint="/v1/get_storage_usage",
@@ -817,8 +825,6 @@ def storage__list(*other):
817
825
  rows = []
818
826
  for namespace, storages in result.items():
819
827
  for name, values in storages.items():
820
- if namespace == username:
821
- namespace = f"**{namespace}**"
822
828
  columns = list(values.keys())
823
829
  rows.append([namespace, name] + [f"{v:.2f} MB" if "capacity" in k else str(v) for k, v in values.items()])
824
830