Named disks (#1983)

* Add support of named disk to sdk * Add support of disk name to neuro disk commands * Allow disk names in disk URI in volumes * Allow to create named disks * Adopt endpoints url change * Add changelog * Add displaying of disk names * Fix mounting disk volume by name
neuro-inc · Feb 12, 2021 · fadeadc · fadeadc
1 parent 55b603e
commit fadeadc
Show file tree

Hide file tree

Showing 17 changed files with 226 additions and 49 deletions.
diff --git a/CHANGELOG.D/1983.feature b/CHANGELOG.D/1983.feature
@@ -0,0 +1,4 @@
+Added support of named:
+- create new disk by `neuro disk create --name <disk-name> STORAGE` command
+- name can be used to get/delete disk: `neuro disk get <disk-name>` or `neuro disk delete <disk-name>`
+- name can be used to mount disk: `neuro run -v disk:<disk-name>:/mnt/disk ...`
diff --git a/CLI.md b/CLI.md
@@ -2118,6 +2118,7 @@ neuro disk create 500M
 Name | Description|
 |----|------------|
 |_--help_|Show this message and exit.|
+|_--name NAME_|Optional disk name|
 |_\--timeout-unused TIMEDELTA_|Optional disk lifetime limit after last usage in the format '1d2h3m4s' \(some parts may be missing). Set '0' to disable. Default value '1d' can be changed in the user config.|
 
 
@@ -2130,7 +2131,7 @@ Get disk DISK_ID.
 **Usage:**
 
 ```bash
-neuro disk get [OPTIONS] DISK_ID
+neuro disk get [OPTIONS] DISK
 ```
 
 **Options:**
@@ -2150,7 +2151,7 @@ Remove disk DISK_ID.
 **Usage:**
 
 ```bash
-neuro disk rm [OPTIONS] DISK_IDS...
+neuro disk rm [OPTIONS] DISKS...
 ```
 
 **Options:**

diff --git a/neuro-cli/docs/disk.md b/neuro-cli/docs/disk.md
@@ -83,6 +83,7 @@ $ neuro disk create 500M
 | Name | Description |
 | :--- | :--- |
 | _--help_ | Show this message and exit. |
+| _--name NAME_ | Optional disk name |
 | _--timeout-unused TIMEDELTA_ | Optional disk lifetime limit after last usage in the format '1d2h3m4s' \(some parts may be missing\). Set '0' to disable. Default value '1d' can be changed in the user config. |
 
 
@@ -95,7 +96,7 @@ Get disk DISK_ID
 #### Usage
 
 ```bash
-neuro disk get [OPTIONS] DISK_ID
+neuro disk get [OPTIONS] DISK
 ```
 
 Get disk `DISK`_ID.
@@ -117,7 +118,7 @@ Remove disk DISK_ID
 #### Usage
 
 ```bash
-neuro disk rm [OPTIONS] DISK_IDS...
+neuro disk rm [OPTIONS] DISKS...
 ```
 
 Remove disk `DISK`_ID.

diff --git a/neuro-cli/src/neuro_cli/click_types.py b/neuro-cli/src/neuro_cli/click_types.py
@@ -19,6 +19,13 @@
 JOB_NAME_REGEX = re.compile(JOB_NAME_PATTERN)
 JOB_LIMIT_ENV = "NEURO_CLI_JOB_AUTOCOMPLETE_LIMIT"
 
+
+# NOTE: these disk name valdation are taken from `platform_disk_api` file `schema.py`
+DISK_NAME_MIN_LENGTH = 3
+DISK_NAME_MAX_LENGTH = 40
+DISK_NAME_PATTERN = "^[a-z](?:-?[a-z0-9])*$"
+DISK_NAME_REGEX = re.compile(JOB_NAME_PATTERN)
+
 _T = TypeVar("_T")
 
 
@@ -158,6 +165,32 @@ def convert(
 JOB_NAME = JobNameType()
 
 
+class DiskNameType(click.ParamType):
+    name = "disk_name"
+
+    def convert(
+        self, value: str, param: Optional[click.Parameter], ctx: Optional[click.Context]
+    ) -> str:
+        if (
+            len(value) < DISK_NAME_MIN_LENGTH
+            or len(value) > DISK_NAME_MAX_LENGTH
+            or DISK_NAME_REGEX.match(value) is None
+        ):
+            raise ValueError(
+                f"Invalid disk name '{value}'.\n"
+                "The name can only contain lowercase letters, numbers and hyphens "
+                "with the following rules: \n"
+                "  - the first character must be a letter; \n"
+                "  - each hyphen must be surrounded by non-hyphen characters; \n"
+                f"  - total length must be between {DISK_NAME_MIN_LENGTH} and "
+                f"{DISK_NAME_MAX_LENGTH} characters long."
+            )
+        return value
+
+
+DISK_NAME = JobNameType()
+
+
 class JobColumnsType(click.ParamType):
     name = "columns"
 
@@ -245,3 +278,35 @@ async def async_complete(
 
 
 JOB = JobType()
+
+
+class DiskType(AsyncType[str]):
+    name = "disk"
+
+    async def async_convert(
+        self,
+        root: Root,
+        value: str,
+        param: Optional[click.Parameter],
+        ctx: Optional[click.Context],
+    ) -> str:
+        return value
+
+    async def async_complete(
+        self, root: Root, ctx: click.Context, args: Sequence[str], incomplete: str
+    ) -> List[Tuple[str, Optional[str]]]:
+        async with await root.init_client() as client:
+            ret: List[Tuple[str, Optional[str]]] = []
+            async for disk in client.disks.list():
+                disk_name = disk.name or ""
+                for test in (
+                    disk.id,
+                    disk_name,
+                ):
+                    if test.startswith(incomplete):
+                        ret.append((test, disk_name))
+
+            return ret
+
+
+DISK = DiskType()
diff --git a/neuro-cli/src/neuro_cli/disks.py b/neuro-cli/src/neuro_cli/disks.py
@@ -1,6 +1,9 @@
 from datetime import timedelta
 from typing import Optional, Sequence
 
+from neuro_cli.click_types import DISK, DISK_NAME
+from neuro_cli.utils import resolve_disk
+
 from .formatters.disks import (
     BaseDisksFormatter,
     DiskFormatter,
@@ -61,8 +64,18 @@ async def ls(root: Root, full_uri: bool, long_format: bool) -> None:
         "in the user config."
     ),
 )
+@option(
+    "--name",
+    type=DISK_NAME,
+    metavar="NAME",
+    help="Optional disk name",
+    default=None,
+)
 async def create(
-    root: Root, storage: str, timeout_unused: Optional[str] = None
+    root: Root,
+    storage: str,
+    timeout_unused: Optional[str] = None,
+    name: Optional[str] = None,
 ) -> None:
     """
     Create a disk with at least storage amount STORAGE.
@@ -90,21 +103,22 @@ async def create(
         disk_timeout_unused = timedelta(seconds=timeout_unused_seconds)
 
     disk = await root.client.disks.create(
-        parse_memory(storage), timeout_unused=disk_timeout_unused
+        parse_memory(storage), timeout_unused=disk_timeout_unused, name=name
     )
     disk_fmtr = DiskFormatter(str)
     with root.pager():
         root.print(disk_fmtr(disk))
 
 
 @command()
-@argument("disk_id")
+@argument("disk", type=DISK)
 @option("--full-uri", is_flag=True, help="Output full disk URI.")
-async def get(root: Root, disk_id: str, full_uri: bool) -> None:
+async def get(root: Root, disk: str, full_uri: bool) -> None:
     """
     Get disk DISK_ID.
     """
-    disk = await root.client.disks.get(disk_id)
+    disk_id = await resolve_disk(disk, client=root.client)
+    disk_obj = await root.client.disks.get(disk_id)
     if full_uri:
         uri_fmtr: URIFormatter = str
     else:
@@ -113,16 +127,17 @@ async def get(root: Root, disk_id: str, full_uri: bool) -> None:
         )
     disk_fmtr = DiskFormatter(uri_fmtr)
     with root.pager():
-        root.print(disk_fmtr(disk))
+        root.print(disk_fmtr(disk_obj))
 
 
 @command()
-@argument("disk_ids", nargs=-1, required=True)
-async def rm(root: Root, disk_ids: Sequence[str]) -> None:
+@argument("disks", type=DISK, nargs=-1, required=True)
+async def rm(root: Root, disks: Sequence[str]) -> None:
     """
     Remove disk DISK_ID.
     """
-    for disk_id in disk_ids:
+    for disk in disks:
+        disk_id = await resolve_disk(disk, client=root.client)
         await root.client.disks.rm(disk_id)
         if root.verbosity > 0:
             root.print(f"Disk with id '{disk_id}' was successfully removed.")

diff --git a/neuro-cli/src/neuro_cli/formatters/disks.py b/neuro-cli/src/neuro_cli/formatters/disks.py
@@ -38,9 +38,11 @@ def __init__(
 
     def _disk_to_table_row(self, disk: Disk) -> Sequence[str]:
         storage_str = utils.format_size(disk.storage)
+
         used_str = utils.format_size(disk.used_bytes)
         line = [
             disk.id,
+            disk.name or "",
             storage_str,
             used_str,
             self._uri_formatter(disk.uri),
@@ -60,6 +62,7 @@ def __call__(self, disks: Sequence[Disk]) -> RenderableType:
         # make sure that the first column is fully expanded
         width = len("disk-06bed296-8b27-4aa8-9e2a-f3c47b41c807")
         table.add_column("Id", style="bold", width=width)
+        table.add_column("Name")
         table.add_column("Storage")
         table.add_column("Used")
         table.add_column("Uri")
@@ -89,6 +92,8 @@ def __call__(self, disk: Disk) -> RenderableType:
         table.add_row("Storage", utils.format_size(disk.storage))
         table.add_row("Used", utils.format_size(disk.used_bytes))
         table.add_row("Uri", self._uri_formatter(disk.uri))
+        if disk.name:
+            table.add_row("Name", disk.name)
         table.add_row("Status", disk.status.value)
         table.add_row("Created at", format_datetime(disk.created_at))
         table.add_row("Last used", format_datetime(disk.last_usage))

diff --git a/neuro-cli/src/neuro_cli/job.py b/neuro-cli/src/neuro_cli/job.py
@@ -1,5 +1,6 @@
 import asyncio
 import contextlib
+import dataclasses
 import logging
 import shlex
 import sys
@@ -28,6 +29,7 @@
 
 from neuro_cli.formatters.images import DockerImageProgress
 from neuro_cli.formatters.utils import URIFormatter, image_formatter, uri_formatter
+from neuro_cli.utils import resolve_disk
 
 from .ael import process_attach, process_exec, process_logs
 from .click_types import (
@@ -967,7 +969,18 @@ async def run_job(
     volume_parse_result = root.client.parse.volumes(volume)
     volumes = list(volume_parse_result.volumes)
     secret_files = volume_parse_result.secret_files
-    disk_volumes = volume_parse_result.disk_volumes
+
+    # Replace disk names with disk ids
+    async def _force_disk_id(disk_uri: URL) -> URL:
+        disk_id = await resolve_disk(disk_uri.parts[-1], client=root.client)
+        return disk_uri / f"../{disk_id}"
+
+    disk_volumes = [
+        dataclasses.replace(
+            disk_volume, disk_uri=await _force_disk_id(disk_volume.disk_uri)
+        )
+        for disk_volume in volume_parse_result.disk_volumes
+    ]
 
     if pass_config:
         env_name = PASS_CONFIG_ENV_NAME

diff --git a/neuro-cli/src/neuro_cli/utils.py b/neuro-cli/src/neuro_cli/utils.py
@@ -468,6 +468,18 @@ async def resolve_job(
     return id_or_name
 
 
+DISK_ID_PATTERN = r"disk-[0-9a-z]{8}-[0-9a-z]{4}-[0-9a-z]{4}-[0-9a-z]{4}-[0-9a-z]{12}"
+
+
+async def resolve_disk(id_or_name: str, *, client: Client) -> str:
+    # Temporary fast path.
+    if re.fullmatch(DISK_ID_PATTERN, id_or_name):
+        return id_or_name
+
+    disk = await client.disks.get(id_or_name)
+    return disk.id
+
+
 SHARE_SCHEMES = ("storage", "image", "job", "blob", "role", "secret", "disk")
 
 

diff --git a/neuro-cli/tests/e2e/conftest.py b/neuro-cli/tests/e2e/conftest.py
@@ -923,9 +923,12 @@ async def docker(loop: asyncio.AbstractEventLoop) -> AsyncIterator[aiodocker.Doc
 @pytest.fixture
 def disk_factory(helper: Helper) -> Callable[[str], ContextManager[str]]:
     @contextmanager
-    def _make_disk(storage: str) -> Iterator[str]:
+    def _make_disk(storage: str, name: Optional[str] = None) -> Iterator[str]:
         # Create disk
-        cap = helper.run_cli(["disk", "create", storage])
+        args = ["disk", "create", storage]
+        if name:
+            args += ["--name", name]
+        cap = helper.run_cli(args)
         assert cap.err == ""
         disk_id = cap.out.splitlines()[0].split()[1]
         yield disk_id

diff --git a/neuro-cli/tests/e2e/test_e2e_jobs.py b/neuro-cli/tests/e2e/test_e2e_jobs.py
@@ -1187,3 +1187,34 @@ def test_job_disk_volume(
         job_id = match.group(1)
 
         helper.wait_job_change_state_to(job_id, JobStatus.SUCCEEDED, JobStatus.FAILED)
+
+
+@pytest.mark.e2e
+def test_job_disk_volume_named(
+    helper: Helper, disk_factory: Callable[[str, str], ContextManager[str]]
+) -> None:
+    disk_name = f"test-disk-{os.urandom(5).hex()}"
+
+    with disk_factory("1G", disk_name):
+        bash_script = 'echo "test data" > /mnt/disk/file && cat /mnt/disk/file'
+        command = f"bash -c '{bash_script}'"
+        captured = helper.run_cli(
+            [
+                "job",
+                "run",
+                "--life-span",
+                "1m",  # Avoid completed job to block disk from cleanup
+                "-v",
+                f"disk:{disk_name}:/mnt/disk:rw",
+                "--no-wait-start",
+                UBUNTU_IMAGE_NAME,
+                command,
+            ]
+        )
+
+        out = captured.out
+        match = re.match("Job ID: (.+)", out)
+        assert match is not None, captured
+        job_id = match.group(1)
+
+        helper.wait_job_change_state_to(job_id, JobStatus.SUCCEEDED, JobStatus.FAILED)
diff --git a/neuro-cli/tests/unit/formatters/ascii/test_disk_formatter_0.ref b/neuro-cli/tests/unit/formatters/ascii/test_disk_formatter_0.ref
@@ -2,6 +2,7 @@ Id              disk
  Storage         11.9G                    
  Used                                     
  Uri             disk://cluster/user/disk 
+ Name            test-disk                
  Status          Ready                    
  Created at      Mar 04 2017              
  Last used       Apr 04 2017              

diff --git a/neuro-cli/tests/unit/formatters/ascii/test_disks_formatter_long_0.ref b/neuro-cli/tests/unit/formatters/ascii/test_disks_formatter_long_0.ref
@@ -1,6 +1,6 @@
-Id                                          Storage   Used   Uri                          Status    Created at    Last used     Timeout unused  
- ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ 
-  disk-1                                      50.0G            disk://cluster/user/disk-1   Pending   Mar 04 2017   Mar 08 2017                   
-  disk-2                                      50.0M            disk://cluster/user/disk-2   Ready     Apr 04 2017                 2d3h4m5s        
-  disk-3                                      50.0K            disk://cluster/user/disk-3   Ready     May 04 2017                                 
-  disk-4                                      50B              disk://cluster/user/disk-4   Broken    Jun 04 2017
+Id                                          Name   Storage   Used   Uri                          Status    Created at    Last used     Timeout unused  
+ ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ 
+  disk-1                                             50.0G            disk://cluster/user/disk-1   Pending   Mar 04 2017   Mar 08 2017                   
+  disk-2                                             50.0M            disk://cluster/user/disk-2   Ready     Apr 04 2017                 2d3h4m5s        
+  disk-3                                             50.0K            disk://cluster/user/disk-3   Ready     May 04 2017                                 
+  disk-4                                             50B              disk://cluster/user/disk-4   Broken    Jun 04 2017
diff --git a/neuro-cli/tests/unit/formatters/ascii/test_disks_formatter_short_0.ref b/neuro-cli/tests/unit/formatters/ascii/test_disks_formatter_short_0.ref
@@ -1,6 +1,6 @@
-Id                                          Storage   Used   Uri                          Status   
- ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ 
-  disk-1                                      50.0G            disk://cluster/user/disk-1   Pending  
-  disk-2                                      50.0M            disk://cluster/user/disk-2   Ready    
-  disk-3                                      50.0K            disk://cluster/user/disk-3   Ready    
-  disk-4                                      50B              disk://cluster/user/disk-4   Broken
+Id                                          Name   Storage   Used   Uri                          Status   
+ ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ 
+  disk-1                                             50.0G            disk://cluster/user/disk-1   Pending  
+  disk-2                                             50.0M            disk://cluster/user/disk-2   Ready    
+  disk-3                                             50.0K            disk://cluster/user/disk-3   Ready    
+  disk-4                                             50B              disk://cluster/user/disk-4   Broken