Merge pull request #222 from latchbio/ayush/qol-improvements

config improvements

latch_cli/centromere/ctx.py

@@ -16,7 +16,7 @@
     _construct_ssh_client,
     _import_flyte_objects,
 )
-from latch_cli.config.latch import _LatchConfig
+from latch_cli.config.latch import config
 from latch_cli.utils import (
     account_id_from_token,
     current_workspace,
@@ -25,9 +25,6 @@
     retrieve_or_login,
 )
 
-config = _LatchConfig()
-endpoints = config.sdk_endpoints
-
 
 @dataclass
 class _Container:
@@ -58,11 +55,11 @@ class _CentromereCtx:
     container_map: Dict[str, _Container]
     workflow_name: Optional[str]
 
-    latch_register_api_url = endpoints["register-workflow"]
-    latch_image_api_url = endpoints["initiate-image-upload"]
-    latch_provision_url = endpoints["provision-centromere"]
-    latch_get_image_url = endpoints["get-image-from-task"]
-    latch_check_version_url = endpoints["check-workflow-version"]
+    latch_register_api_url = config.api.workflow.register
+    latch_image_api_url = config.api.workflow.upload_image
+    latch_provision_url = config.api.centromere.provision
+    latch_get_image_url = config.api.workflow.get_image
+    latch_check_version_url = config.api.workflow.check_version
 
     def __init__(
         self,
@@ -84,7 +81,7 @@ def __init__(
                 self.account_id = ws
 
             self.pkg_root = Path(pkg_root).resolve()
-            self.dkr_repo = _LatchConfig.dkr_repo
+            self.dkr_repo = config.dkr_repo
             self.remote = remote
             self.container_map = {}
 
@@ -131,7 +128,6 @@ def __init__(
             )
 
             if remote is True:
-
                 self.ssh_key_path = Path(self.pkg_root) / ".ssh_key"
                 self.public_key = generate_temporary_ssh_credentials(self.ssh_key_path)
 

latch_cli/config/latch.py

@@ -5,53 +5,91 @@
 """
 
 import os
+from dataclasses import dataclass, fields, is_dataclass
+from typing import Type, TypeVar
 
 DOMAIN = os.environ.get("LATCH_SDK_DOMAIN", "latch.bio")
 CONSOLE_URL = f"https://console.{DOMAIN}"
 NUCLEUS_URL = f"https://nucleus.{DOMAIN}"
 
+T = TypeVar("T")
 
-SDK_ENDPOINTS = {
-    "get-test-data-creds": "/sdk/get-test-data-creds",
-    "initiate-multipart-upload": "/sdk/initiate-multipart-upload",
-    "complete-multipart-upload": "/sdk/complete-multipart-upload",
-    "check-workflow-version": "/sdk/check-workflow-version",
-    "download": "/sdk/download",
-    "list-files": "/sdk/list",
-    "initiate-image-upload": "/sdk/initiate-image-upload",
-    "register-workflow": "/sdk/register-workflow",
-    "get-workflow-interface": "/sdk/wf-interface",
-    "access-jwt": "/sdk/access-jwt",
-    "execute-workflow": "/sdk/wf",
-    "get-workflows": "/sdk/get-wf",
-    "verify": "/sdk/verify",
-    "remove": "/sdk/rm",
-    "id": "/sdk/node-id",
-    "mkdir": "/sdk/mkdir",
-    "rmdir": "/sdk/rmdir",
-    "touch": "/sdk/touch",
-    "pod-exec-info": "/sdk/get-pod-exec-info",
-    "provision-centromere": "/sdk/provision-centromere",
-    "get-image-from-task": "/sdk/get-image-from-task",
-    "preview": "/sdk/workflow-ui-preview",
-    "get-ws": "/sdk/get-ws",
-    "get-executions": "/sdk/get-executions",
-    "get-workflow-graph": "/sdk/get-workflow-graph",
-    "get-logs": "/sdk/get-logs-for-node",
-    "abort-execution": "/sdk/abort-execution",
-    "local-development": "/sdk/initiate-local-development-session",
-    "close-local-development": "/sdk/close-local-development-session",
-    "get-latest-version": "/sdk/get-latest-version-new",
-}
 
+@dataclass
+class _DataAPI:
+    begin_upload: str = "/sdk/initiate-multipart-upload"
+    complete_upload: str = "/sdk/complete-multipart-upload"
+    download: str = "/sdk/download"
+    id: str = "/sdk/node-id"
+    list: str = "/sdk/list"
+    remove: str = "/sdk/rm"
+    touch: str = "/sdk/touch"
+    mkdir: str = "/sdk/mkdir"
+    verify: str = "/sdk/verify"
+    test_data: str = "/sdk/get-test-data-creds"
 
+
+@dataclass
+class _WorkflowAPI:
+    upload_image: str = "/sdk/initiate-image-upload"
+    get_image: str = "/sdk/get-image-from-task"
+    register: str = "/sdk/register-workflow"
+    list: str = "/sdk/get-wf"
+    interface: str = "/sdk/wf-interface"
+    graph: str = "/sdk/get-workflow-graph"
+    check_version: str = "/sdk/check-workflow-version"
+    get_latest: str = "/sdk/get-latest-version-new"
+    preview: str = "/sdk/workflow-ui-preview"
+
+
+@dataclass
+class _ExecutionAPI:
+    create: str = "/sdk/wf"
+    list: str = "/sdk/get-executions"
+    abort: str = "/sdk/abort-execution"
+    logs: str = "/sdk/get-logs-for-node"
+    exec: str = "/sdk/get-pod-exec-info"
+
+
+@dataclass
+class _UserAPI:
+    jwt: str = "/sdk/access-jwt"
+    list_workspaces: str = "/sdk/get-ws"
+
+
+@dataclass
+class _CentromereAPI:
+    provision: str = "/sdk/provision-centromere"
+    start_local_dev: str = "/sdk/initiate-local-development-session"
+    stop_local_dev: str = "/sdk/close-local-development-session"
+
+
+@dataclass
+class _API:
+    data: _DataAPI
+    workflow: _WorkflowAPI
+    execution: _ExecutionAPI
+    user: _UserAPI
+    centromere: _CentromereAPI
+
+
+@dataclass
 class _LatchConfig:
-    dkr_repo = "812206152185.dkr.ecr.us-west-2.amazonaws.com"
-
-    def __init__(self):
-        self.console_url = CONSOLE_URL
-        self.nucleus_url = NUCLEUS_URL
-        self.sdk_endpoints = {
-            key: f"{self.nucleus_url}{endpoint}"
-            for key, endpoint in SDK_ENDPOINTS.items()
-        }
+    api: _API
+    dkr_repo: str = "812206152185.dkr.ecr.us-west-2.amazonaws.com"
+    console_url: str = CONSOLE_URL
+    nucleus_url: str = NUCLEUS_URL
+
+
+def build_endpoints(x: Type[T] = _API) -> T:
+    res = {}
+    for field in fields(x):
+        if is_dataclass(field.type):
+            res[field.name] = build_endpoints(field.type)
+        elif field.type is str:
+            res[field.name] = NUCLEUS_URL + field.default
+    return x(**res)
+
+
+# singleton config instance
+config = _LatchConfig(api=build_endpoints())

latch_cli/config/user.py

@@ -10,21 +10,30 @@
 class _UserConfig:
     """User specific configuration persisted in `~/.latch/`."""
 
-    _root = None
-
     def __init__(self):
-        self._root = Path.home().resolve().joinpath(".latch")
+        self._root = Path.home().resolve() / ".latch"
+        self._token_path = None
+        self._workspace_path = None
 
     @property
-    def root_dir(self):
-        if self._root.exists() is False:
+    def root(self):
+        if not self._root.exists():
             self._root.mkdir(parents=True)
         return self._root
 
-    def token_exists(self) -> bool:
-        if self.token != "":
-            return False
-        return True
+    @property
+    def token_path(self):
+        if self._token_path is None:
+            self._token_path = self.root / "token"
+        self._token_path.touch(exist_ok=True)
+        return self._token_path
+
+    @property
+    def workspace_path(self):
+        if self._workspace_path is None:
+            self._workspace_path = self.root / "workspace"
+        self._workspace_path.touch(exist_ok=True)
+        return self._workspace_path
 
     @property
     def token(self) -> str:
@@ -33,24 +42,27 @@ def token(self) -> str:
         Returns: ID token if exists else an empty string.
         """
         try:
-            with open(self.root_dir.joinpath("token"), "r") as f:
+            with open(self.token_path, "r") as f:
                 return f.read().strip()
         except FileNotFoundError:
             return ""
 
     @property
-    def current_workspace(self) -> str:
+    def workspace(self) -> str:
         try:
-            with open(self.root_dir.joinpath("workspace"), "r") as f:
+            with open(self.workspace_path, "r") as f:
                 return f.read().strip()
         except FileNotFoundError:
             return ""
 
     def update_token(self, token: str):
         """Updates user config with new token regardless if one exists."""
-        with open(self.root_dir.joinpath("token"), "w") as f:
+        with open(self.token_path, "w") as f:
             f.write(token)
 
     def update_workspace(self, workspace: str):
-        with open(self.root_dir.joinpath("workspace"), "w") as f:
+        with open(self.workspace_path, "w") as f:
             f.write(workspace)
+
+
+user_config = _UserConfig()

latch_cli/exceptions/cache.py

@@ -1,4 +1,3 @@
 from typing import Dict, List, Tuple
 
-
 _code_cache: Dict[Tuple[str, int], List[str]] = {}

latch_cli/exceptions/handler.py

@@ -1,19 +1,17 @@
 import json
 import os
-from pathlib import Path
 import platform
-import re
 import sys
 import tarfile
 import tempfile
+from pathlib import Path
 from traceback import print_exc, walk_tb
 from types import TracebackType
 from typing import Optional, Type
 
 from latch_cli.constants import FILE_MAX_SIZE, IGNORE_REGEX
-from latch_cli.utils import get_local_package_version
-
 from latch_cli.exceptions.traceback import _Traceback
+from latch_cli.utils import get_local_package_version
 
 
 class CrashHandler:
@@ -49,7 +47,6 @@ def _write_state_to_tarball(self):
         tarball_path.unlink(missing_ok=True)
 
         with tarfile.open(tarball_path, mode="x:gz") as tf:
-
             # If calling stack frame is handling an exception, we want to store
             # the traceback in a log file.
             if sys.exc_info()[0] is not None:

latch_cli/services/cp.py

@@ -10,15 +10,12 @@
 from tqdm.auto import tqdm
 
 import latch_cli.tinyrequests as tinyrequests
-from latch_cli.config.latch import _LatchConfig
+from latch_cli.config.latch import config
 from latch_cli.constants import FILE_CHUNK_SIZE
 from latch_cli.services.deprecated.mkdir import mkdir
 from latch_cli.services.deprecated.touch import touch
 from latch_cli.utils import _normalize_remote_path, current_workspace, retrieve_or_login
 
-config = _LatchConfig()
-endpoints = config.sdk_endpoints
-
 # tqdm progress bars aren't thread safe so restrict so that only one can update at a time
 IO_LOCK = threading.Lock()
 
@@ -29,7 +26,7 @@ def _dir_exists(remote_dir: str) -> bool:
     token = retrieve_or_login()
     headers = {"Authorization": f"Bearer {token}"}
     data = {"filename": remote_dir}
-    response = tinyrequests.post(url=endpoints["verify"], headers=headers, json=data)
+    response = tinyrequests.post(url=config.api.data.verify, headers=headers, json=data)
     try:
         assert response.status_code == 200
     except:
@@ -110,7 +107,7 @@ def _upload_file(
         return
 
     response = tinyrequests.post(
-        endpoints["initiate-multipart-upload"],
+        config.api.data.begin_upload,
         headers={"Authorization": f"Bearer {retrieve_or_login()}"},
         json={
             "ws_account_id": current_workspace(),
@@ -161,7 +158,7 @@ def _upload_file(
     parts.sort(key=lambda res: res["PartNumber"])
 
     response = tinyrequests.post(
-        endpoints["complete-multipart-upload"],
+        config.api.data.complete_upload,
         headers={"Authorization": f"Bearer {retrieve_or_login()}"},
         json={
             "path": path,
@@ -236,7 +233,7 @@ def download(
 
     print("Generating download URLs...")
     response = tinyrequests.post(
-        endpoints["download"],
+        config.api.data.download,
         headers={"Authorization": f"Bearer {retrieve_or_login()}"},
         json={
             "source_path": remote_source,

latch_cli/services/deprecated/mkdir.py

@@ -1,10 +1,7 @@
 import latch_cli.tinyrequests as tinyrequests
-from latch_cli.config.latch import _LatchConfig
+from latch_cli.config.latch import config
 from latch_cli.utils import _normalize_remote_path, current_workspace, retrieve_or_login
 
-config = _LatchConfig()
-endpoints = config.sdk_endpoints
-
 
 def mkdir(remote_directory):
     """Creates an empty directory on Latch
@@ -44,7 +41,7 @@ def mkdir(remote_directory):
     headers = {"Authorization": f"Bearer {token}"}
     data = {"directory": remote_directory, "ws_account_id": current_workspace()}
 
-    response = tinyrequests.post(endpoints["mkdir"], headers=headers, json=data)
+    response = tinyrequests.post(config.api.data.mkdir, headers=headers, json=data)
     json_data = response.json()
 
     if not json_data["success"]:

latch_cli/services/deprecated/rm.py

@@ -1,11 +1,8 @@
 import requests
 
-from latch_cli.config.latch import _LatchConfig
+from latch_cli.config.latch import config
 from latch_cli.utils import _normalize_remote_path, current_workspace, retrieve_or_login
 
-config = _LatchConfig()
-endpoints = config.sdk_endpoints
-
 
 def rm(remote_path: str):
     """Deletes an entity on Latch
@@ -41,7 +38,7 @@ def rm(remote_path: str):
 
     data = {"filename": remote_path, "ws_account_id": current_workspace()}
     headers = {"Authorization": f"Bearer {token}"}
-    response = requests.post(url=endpoints["remove"], headers=headers, json=data)
+    response = requests.post(config.api.data.remove, headers=headers, json=data)
 
     data = response.json()
     if not data["success"]: