Add script for interacting with user vault (#117)

V2arK · web-flow · commit 0511fa9d9781 · 2026-01-30T17:16:35.000-05:00
* add script for interacting with user vault
diff --git a/centml/cli/cluster.py b/centml/cli/cluster.py
@@ -6,7 +6,6 @@
 from centml.sdk import DeploymentType, DeploymentStatus, ServiceStatus, ApiException, HardwareInstanceResponse
 from centml.sdk.api import get_centml_client
 
-
 # convert deployment type enum to a user friendly name
 depl_type_to_name_map = {
     DeploymentType.INFERENCE: "inference",
diff --git a/centml/cli/login.py b/centml/cli/login.py
@@ -14,7 +14,6 @@
 from centml.sdk import auth
 from centml.sdk.config import settings
 
-
 CLIENT_ID = settings.CENTML_WORKOS_CLIENT_ID
 SERVER_HOST = "127.0.0.1"
 SERVER_PORT = 57983
@@ -51,18 +50,14 @@ def do_GET(self):
         self.send_response(200)
         self.send_header("Content-type", "text/html")
         self.end_headers()
-        self.wfile.write(
-            """
+        self.wfile.write("""
             <html>
                 <body>
                     <h1>Succesfully logged into CentML CLI</h1>
                     <p>You can now close this tab and continue in the CLI.</p>
                 </body>
             </html>
-            """.encode(
-                "utf-8"
-            )
-        )
+            """.encode("utf-8"))
 
     def log_message(self, format, *args):
         # Override this to suppress logging
diff --git a/examples/sdk/create_cserve.py b/examples/sdk/create_cserve.py
@@ -36,9 +36,7 @@ def get_default_cserve_config(cclient, name, model):
 def main():
     with get_centml_client() as cclient:
         ### Get the configurations for the Qwen model
-        qwen_config = get_fastest_cserve_config(
-            cclient, name="qwen-fastest", model="Qwen/Qwen2-VL-7B-Instruct"
-        )
+        qwen_config = get_fastest_cserve_config(cclient, name="qwen-fastest", model="Qwen/Qwen2-VL-7B-Instruct")
         # qwen_config = get_default_cserve_config(cclient, name="qwen-default", model="Qwen/Qwen2-VL-7B-Instruct")
 
         ### Modify the recipe if necessary
diff --git a/examples/sdk/create_inference.py b/examples/sdk/create_inference.py
@@ -38,5 +38,6 @@ def main():
         cclient.delete(deployment.id)
         '''
 
+
 if __name__ == "__main__":
     main()
diff --git a/examples/sdk/get_deployment_usage.py b/examples/sdk/get_deployment_usage.py
@@ -3,7 +3,6 @@
 from centml.sdk.api import get_centml_client
 from centml.sdk import Metric
 
-
 HOUR_IN_SECONDS = 60 * 60
 DAY_IN_SECONDS = 24 * HOUR_IN_SECONDS
 MAX_DATA_POINTS = 10_000
@@ -17,22 +16,23 @@ def get_step_size(start_time_in_seconds: int, end_time_in_seconds: int) -> int:
     # 2 days to 7 days: 5m
     elif time_delta_in_seconds <= 7 * DAY_IN_SECONDS:
         return 5 * 60
-	#  7 days to 14 days: 10m
+    #  7 days to 14 days: 10m
     elif time_delta_in_seconds <= 14 * DAY_IN_SECONDS:
         return 10 * 60
-	# 14 days to 30 days: 30m
+    # 14 days to 30 days: 30m
     elif time_delta_in_seconds <= 30 * DAY_IN_SECONDS:
         return 30 * 60
-	# 30 days to 60 days: 1 hour
+    # 30 days to 60 days: 1 hour
     elif time_delta_in_seconds <= 60 * DAY_IN_SECONDS:
         return HOUR_IN_SECONDS
-	# 60 days to 90 days: 2 hours
+    # 60 days to 90 days: 2 hours
     elif time_delta_in_seconds <= 90 * DAY_IN_SECONDS:
         return 2 * HOUR_IN_SECONDS
-	# 90+ days: 3 hours (This catches all ranges greater than 90 days)
+    # 90+ days: 3 hours (This catches all ranges greater than 90 days)
     else:
         return 3 * HOUR_IN_SECONDS
 
+
 def main():
     with get_centml_client() as cclient:
         start_time = 1752084581
diff --git a/scripts/data_collection.py b/scripts/data_collection.py
@@ -7,11 +7,10 @@
     AutoModelForCausalLM,
     AutoTokenizer,
     AutoModelForImageClassification,
-    AutoModelForObjectDetection
+    AutoModelForObjectDetection,
 )
 
 
-
 from centml.compiler.prediction.kdtree import KDTreeWithValues
 from centml.compiler.prediction.profiler import Profiler
 from scripts.timer import timed
@@ -26,8 +25,8 @@
 # Different HuggingFace Models + Different Input Sizes
 llm_tests = [
     ("google/gemma-7b", (1, 128)),
-    ("microsoft/phi-2", (1,512)),
-    ("microsoft/phi-2", (2,512)),
+    ("microsoft/phi-2", (1, 512)),
+    ("microsoft/phi-2", (2, 512)),
     ("facebook/bart-large", (1, 1024)),
     ("facebook/bart-large", (2, 512)),
     ("gpt2-xl", (1, 1024)),
@@ -212,6 +211,7 @@ def image_classification_test(model_name, batch_size, custom_backend):
     gc.collect()
     torch.cuda.empty_cache()
 
+
 def object_detection_test(model_name, batch_size, custom_backend):
     global cuda_kernel_time
     global actual_time
@@ -240,6 +240,7 @@ def object_detection_test(model_name, batch_size, custom_backend):
     gc.collect()
     torch.cuda.empty_cache()
 
+
 # for model_name, input_size in large_llm_tests:
 #     llm_test(model_name, input_size, custom_backend)
 
diff --git a/scripts/timer.py b/scripts/timer.py
@@ -1,5 +1,6 @@
 import torch
 
+
 def timed(fn):
     start = torch.cuda.Event(enable_timing=True)
     end = torch.cuda.Event(enable_timing=True)
diff --git a/scripts/user_vault/README.md b/scripts/user_vault/README.md
@@ -0,0 +1,169 @@
+# User Vault Scripts
+
+Scripts for managing user vault (secrets) in CentML Platform.
+
+## Overview
+
+The CentML User Vault is a secure storage system for sensitive information that can be used across your deployments. This includes environment variables, API tokens, SSH keys, and certificates. These scripts allow you to view and manage your vault items from the command line.
+
+## Prerequisites
+
+### 1. Install the centml package
+
+From the repository root directory:
+
+```bash
+pip install -e ./
+```
+
+Or install directly from GitHub:
+
+```bash
+pip install git+https://github.com/CentML/centml-python-client.git@main
+```
+
+### 2. Authenticate with CentML
+
+Login to your CentML account:
+
+```bash
+centml login
+```
+
+This will open a browser window for authentication. Once completed, your credentials will be stored locally.
+
+## Available Scripts
+
+### get_vault_items.py
+
+Retrieves and displays all items stored in your CentML vault.
+
+#### Supported Vault Types
+
+| Type | Description | Example Use Case |
+|------|-------------|------------------|
+| `env_vars` | Environment variables | Database URLs, API endpoints |
+| `ssh_keys` | SSH keys | Git repository access |
+| `bearer_tokens` | Bearer tokens | Service authentication |
+| `access_tokens` | Access tokens | HuggingFace tokens, Weights & Biases API keys |
+| `certificates` | Certificates | TLS/SSL certificates |
+
+#### Usage
+
+Run the script from the `scripts/user_vault` directory:
+
+```bash
+cd scripts/user_vault
+python get_vault_items.py [OPTIONS]
+```
+
+#### Command Line Options
+
+| Option | Description | Default |
+|--------|-------------|---------|
+| `--type TYPE` | Filter results by vault type (see supported types above) | Show all types |
+| `--search QUERY` | Filter items by key name (case-sensitive substring match) | No filter |
+| `--show-values` | Display the actual secret values | Keys only |
+| `--help` | Show help message and exit | - |
+
+#### Examples
+
+**List all vault items (keys only):**
+
+```bash
+python get_vault_items.py
+```
+
+**List only environment variables:**
+
+```bash
+python get_vault_items.py --type env_vars
+```
+
+**List only access tokens (e.g., HuggingFace tokens):**
+
+```bash
+python get_vault_items.py --type access_tokens
+```
+
+**Search for items containing "HF" in the key name:**
+
+```bash
+python get_vault_items.py --search HF
+```
+
+**Show all items with their values:**
+
+```bash
+python get_vault_items.py --show-values
+```
+
+**Combine multiple options:**
+
+```bash
+python get_vault_items.py --type env_vars --show-values --search DATABASE
+```
+
+#### Example Output
+
+Without `--show-values`:
+
+```
+Found 5 vault item(s)
+
+==================================================
+Type: access_tokens (2 item(s))
+==================================================
+  HF_TOKEN
+  WANDB_API_KEY
+
+==================================================
+Type: env_vars (3 item(s))
+==================================================
+  API_KEY
+  DATABASE_URL
+  MY_SECRET
+```
+
+With `--show-values`:
+
+```
+Found 5 vault item(s)
+
+==================================================
+Type: access_tokens (2 item(s))
+==================================================
+  HF_TOKEN: hf_xxxxxxxxxxxxxxxxxxxx
+  WANDB_API_KEY: xxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxx
+
+==================================================
+Type: env_vars (3 item(s))
+==================================================
+  API_KEY: sk-xxxxxxxxxxxxxxxx
+  DATABASE_URL: postgresql://user:pass@host:5432/db
+  MY_SECRET: my-secret-value
+```
+
+## Troubleshooting
+
+### Authentication Error
+
+If you see an authentication error, try logging in again:
+
+```bash
+centml login
+```
+
+### Module Not Found
+
+If you see `ModuleNotFoundError`, ensure you have installed the centml package:
+
+```bash
+pip install -e ./
+```
+
+### No Items Found
+
+If the script returns "No vault items found", verify that:
+1. You are logged into the correct CentML account
+2. You have created vault items in the CentML web UI or via API
diff --git a/scripts/user_vault/get_vault_items.py b/scripts/user_vault/get_vault_items.py
