Add env vars to program + improve UX

Author: philogicae

src/aleph_client/commands/help_strings.py

@@ -18,6 +18,8 @@
 IMMUTABLE_VOLUME = """Immutable volumes are pinned on the network and can be used by multiple VMs at the same time. They are read-only and useful for setting up libraries or other dependencies.\n
 Requires at least "name", "ref" (message hash) and "mount" path. "use_latest" is True by default, to use the latest version of the volume, if it has been amended. See the docs for more info: https://docs.aleph.im/computing/volumes/immutable/\n
 Example: --immutable-volume name=libs,ref=25a393222692c2f73489dc6710ae87605a96742ceef7b91de4d7ec34bb688d94,mount=/lib/python3.8/site-packages"""
+SKIP_ENV_VAR = "Skip prompt to set environment variables"
+ENVIRONMENT_VARIABLES = """Environment variables to pass. They will be public and visible in the message, so don't include secrets. Must be a comma separated list. Example: `KEY=value` or `KEY=value,KEY=value`"""
 ASK_FOR_CONFIRMATION = "Prompt user for confirmation"
 IPFS_CATCH_ALL_PATH = "Choose a relative path to catch all unmatched route or a 404 error"
 PAYMENT_TYPE = "Payment method, either holding tokens, NFTs, or Pay-As-You-Go via token streaming"
@@ -58,3 +60,7 @@
 CREATE_ACTIVE = "Loads the new private key after creation"
 PROMPT_CRN_URL = "URL of the CRN (Compute node) on which the instance is running"
 PROMPT_PROGRAM_CRN_URL = "URL of the CRN (Compute node) on which the program is running"
+PROGRAM_PATH = "Path to your source code. Can be a directory, a .squashfs file or a .zip archive"
+PROGRAM_ENTRYPOINT = "Your program entrypoint. Example: `main:app` for Python programs, else `run.sh` for a script containing your launch command"
+PROGRAM_RUNTIME = "Hash of the runtime to use for your program. You can also create your own runtime and pin it. Currently defaults to `{runtime_id}` (Use `aleph program runtime-checker` to inspect it)"
+PROGRAM_BETA = "If true, you will be prompted to add message subscriptions to your program"

src/aleph_client/commands/instance/__init__.py

@@ -369,7 +369,7 @@ async def create(
                 # User has ctrl-c
                 raise typer.Exit(1)
             crn.display_crn_specs()
-            if not Confirm.ask("\nDeploy on this node ?"):
+            if not Confirm.ask("\nDeploy on this node?"):
                 crn = None
                 continue
     elif crn_url or crn_hash:

src/aleph_client/commands/program.py

@@ -33,6 +33,7 @@
 from aleph_client.commands import help_strings
 from aleph_client.commands.utils import (
     filter_only_valid_messages,
+    get_or_prompt_environment_variables,
     get_or_prompt_volumes,
     input_multiline,
     safe_getattr,
@@ -49,28 +50,26 @@
 
 @app.command()
 async def upload(
-    path: Path = typer.Argument(
-        ..., help="Path to your source code. Can be a directory, a .squashfs file or a .zip archive"
-    ),
+    path: Path = typer.Argument(..., help=help_strings.PROGRAM_PATH),
     entrypoint: str = typer.Argument(
         ...,
-        help="Your program entrypoint. E.g. `main:app` for Python programs, or `run.sh` for a script containing your launch command",
+        help=help_strings.PROGRAM_ENTRYPOINT,
     ),
     channel: Optional[str] = typer.Option(default=settings.DEFAULT_CHANNEL, help=help_strings.CHANNEL),
-    memory: int = typer.Option(settings.DEFAULT_VM_MEMORY, help="Maximum memory allocation on vm in MiB"),
-    vcpus: int = typer.Option(settings.DEFAULT_VM_VCPUS, help="Number of virtual cpus to allocate."),
+    memory: int = typer.Option(settings.DEFAULT_VM_MEMORY, help=help_strings.MEMORY),
+    vcpus: int = typer.Option(settings.DEFAULT_VM_VCPUS, help=help_strings.VCPUS),
     timeout_seconds: float = typer.Option(
         settings.DEFAULT_VM_TIMEOUT,
-        help="If vm is not called after [timeout_seconds] it will shutdown",
+        help=help_strings.TIMEOUT_SECONDS,
     ),
     name: Optional[str] = typer.Option(None, help="Name for your program"),
     runtime: str = typer.Option(
         None,
-        help=f"Hash of the runtime to use for your program. You can also create your own runtime and pin it. Currently defaults to `{settings.DEFAULT_RUNTIME_ID}`",
+        help=help_strings.PROGRAM_RUNTIME.format(runtime_id=settings.DEFAULT_RUNTIME_ID),
     ),
     beta: bool = typer.Option(
         False,
-        help="If true, you will be prompted to add message subscriptions to your program",
+        help=help_strings.PROGRAM_BETA,
     ),
     persistent: bool = False,
     skip_volume: bool = typer.Option(False, help=help_strings.SKIP_VOLUME),
@@ -80,14 +79,16 @@ async def upload(
         None,
         help=help_strings.IMMUTABLE_VOLUME,
     ),
+    skip_env_var: bool = typer.Option(False, help=help_strings.SKIP_ENV_VAR),
+    env_vars: Optional[str] = typer.Option(None, help=help_strings.ENVIRONMENT_VARIABLES),
     private_key: Optional[str] = typer.Option(settings.PRIVATE_KEY_STRING, help=help_strings.PRIVATE_KEY),
     private_key_file: Optional[Path] = typer.Option(settings.PRIVATE_KEY_FILE, help=help_strings.PRIVATE_KEY_FILE),
     print_messages: bool = typer.Option(False),
     print_code_message: bool = typer.Option(False),
     print_program_message: bool = typer.Option(False),
     verbose: bool = True,
     debug: bool = False,
-):
+) -> Optional[str]:
     """Register a program to run on aleph.im. For more information, see https://docs.aleph.im/computing/"""
 
     setup_logging(debug)
@@ -106,7 +107,7 @@ async def upload(
     account: AccountFromPrivateKey = _load_account(private_key, private_key_file)
 
     name = name or validated_prompt("Program name", lambda x: len(x) < 65)
-    runtime = runtime or input(f"Ref of runtime ? [{settings.DEFAULT_RUNTIME_ID}] ") or settings.DEFAULT_RUNTIME_ID
+    runtime = runtime or input(f"Ref of runtime? [{settings.DEFAULT_RUNTIME_ID}] ") or settings.DEFAULT_RUNTIME_ID
 
     volumes = []
     if not skip_volume:
@@ -116,8 +117,12 @@ async def upload(
             immutable_volume=immutable_volume,
         )
 
+    environment_variables = None
+    if not skip_env_var:
+        environment_variables = get_or_prompt_environment_variables(env_vars)
+
     subscriptions: Optional[List[Mapping]] = None
-    if beta and yes_no_input("Subscribe to messages ?", default=False):
+    if beta and yes_no_input("Subscribe to messages?", default=False):
         content_raw = input_multiline()
         try:
             subscriptions = json.loads(content_raw)
@@ -163,23 +168,46 @@ async def upload(
             persistent=persistent,
             encoding=encoding,
             volumes=volumes,
+            environment_variables=environment_variables,
             subscriptions=subscriptions,
         )
         logger.debug("Upload finished")
         if print_messages or print_program_message:
             typer.echo(f"{message.json(indent=4)}")
 
         item_hash: ItemHash = message.item_hash
-        hash_base32 = b32encode(b16decode(item_hash.upper())).strip(b"=").lower().decode()
-
         if verbose:
-            typer.echo(
-                f"Your program has been uploaded on aleph.im\n\n"
-                "Available on:\n"
-                f"  {settings.VM_URL_PATH.format(hash=item_hash)}\n"
-                f"  {settings.VM_URL_HOST.format(hash_base32=hash_base32)}\n"
-                "Visualise on:\n  https://explorer.aleph.im/address/"
-                f"{message.chain.value}/{message.sender}/message/PROGRAM/{item_hash}\n"
+            hash_base32 = b32encode(b16decode(item_hash.upper())).strip(b"=").lower().decode()
+            func_url_1 = f"{settings.VM_URL_PATH.format(hash=item_hash)}"
+            func_url_2 = f"{settings.VM_URL_HOST.format(hash_base32=hash_base32)}"
+
+            console = Console()
+            infos = [
+                Text.from_markup(f"Your program [bright_cyan]{item_hash}[/bright_cyan] has been uploaded on aleph.im."),
+                Text.assemble(
+                    "\n\nAvailable on:\n",
+                    Text.from_markup(
+                        f"↳ [bright_yellow][link={func_url_1}]{func_url_1}[/link][/bright_yellow]\n",
+                        style="italic",
+                    ),
+                    Text.from_markup(
+                        f"↳ [dark_olive_green2][link={func_url_2}]{func_url_2}[/link][/dark_olive_green2]",
+                        style="italic",
+                    ),
+                    "\n\nVisualise on:\n",
+                    Text.from_markup(
+                        f"[blue]https://explorer.aleph.im/address/{message.chain.value}/{message.sender}/message/PROGRAM/{item_hash}[/blue]"
+                    ),
+                ),
+            ]
+            console.print(
+                Panel(
+                    Text.assemble(*infos),
+                    title="Program Created",
+                    border_style="green",
+                    expand=False,
+                    title_align="left",
+                )
             )
         return item_hash
 
@@ -499,6 +527,7 @@ async def runtime_checker(
             beta=False,
             persistent=False,
             skip_volume=True,
+            skip_env_var=True,
             private_key=private_key,
             private_key_file=private_key_file,
             print_messages=False,
@@ -507,6 +536,8 @@ async def runtime_checker(
             verbose=verbose,
             debug=debug,
         )
+        if not program_hash:
+            raise Exception("No program hash")
     except Exception as e:
         echo(f"Failed to deploy the runtime checker program: {e}")
         raise typer.Exit(code=1)

src/aleph_client/commands/utils.py

@@ -80,7 +80,7 @@ def yes_no_input(text: str, default: str | bool) -> bool:
 
 
 def prompt_for_volumes():
-    while yes_no_input("Add volume ?", default=False):
+    while yes_no_input("Add volume?", default=False):
         mount = validated_prompt("Mount path (ex: /opt/data): ", lambda text: len(text) > 0)
         name = validated_prompt("Name: ", lambda text: len(text) > 0)
         comment = Prompt.ask("Comment: ")
@@ -96,7 +96,7 @@ def prompt_for_volumes():
             }
         else:
             ref = validated_prompt("Item hash: ", lambda text: len(text) == 64)
-            use_latest = yes_no_input("Use latest version ?", default=True)
+            use_latest = yes_no_input("Use latest version?", default=True)
             yield {
                 "comment": comment,
                 "mount": mount,
@@ -151,6 +151,27 @@ def get_or_prompt_volumes(ephemeral_volume, immutable_volume, persistent_volume)
     return volumes
 
 
+def env_vars_to_dict(env_vars: Optional[str]) -> Dict[str, str]:
+    dict_store: Dict[str, str] = {}
+    if env_vars:
+        for env_var in env_vars.split(","):
+            label, value = env_var.split("=", 1)
+            dict_store[label.strip()] = value.strip()
+    return dict_store
+
+
+def get_or_prompt_environment_variables(env_vars: Optional[str]) -> Optional[Dict[str, str]]:
+    environment_variables: Dict[str, str] = {}
+    if not env_vars:
+        while yes_no_input("Add environment variable?", default=False):
+            label = validated_prompt("Label: ", lambda text: len(text) > 0)
+            value = validated_prompt("Value: ", lambda text: len(text) > 0)
+            environment_variables[label] = value
+    else:
+        environment_variables = env_vars_to_dict(env_vars)
+    return environment_variables if environment_variables else None
+
+
 def str_to_datetime(date: Optional[str]) -> Optional[datetime]:
     """
     Converts a string representation of a date/time to a datetime object.