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