Add neuro admin add-quota / set-quota commands (#1275)

CHANGELOG.D/1142.feature

@@ -0,0 +1 @@
+Add `neuro admin add-user-quota` and `neuro admin set-user-quota` commands to control user quotas

README.md

@@ -10,6 +10,8 @@
 		* [neuro admin get-cluster-users](#neuro-admin-get-cluster-users)
 		* [neuro admin add-cluster-user](#neuro-admin-add-cluster-user)
 		* [neuro admin remove-cluster-user](#neuro-admin-remove-cluster-user)
+		* [neuro admin set-user-quota](#neuro-admin-set-user-quota)
+		* [neuro admin add-user-quota](#neuro-admin-add-user-quota)
 	* [neuro job](#neuro-job)
 		* [neuro job run](#neuro-job-run)
 		* [neuro job submit](#neuro-job-submit)
@@ -185,6 +187,8 @@ Name | Description|
 | _[neuro admin get\-cluster-users](#neuro-admin-get-cluster-users)_| Print the list of all users in the cluster with their assigned role |
 | _[neuro admin add\-cluster-user](#neuro-admin-add-cluster-user)_| Add user access to specified cluster |
 | _[neuro admin remove\-cluster-user](#neuro-admin-remove-cluster-user)_| Remove user access from the cluster |
+| _[neuro admin set\-user-quota](#neuro-admin-set-user-quota)_| Set user quota to given values |
+| _[neuro admin add\-user-quota](#neuro-admin-add-user-quota)_| Add given values to user quota |
 
 
 
@@ -304,6 +308,48 @@ Name | Description|
 
 
 
+### neuro admin set-user-quota
+
+Set user quota to given values
+
+**Usage:**
+
+```bash
+neuro admin set-user-quota [OPTIONS] CLUSTER_NAME USER_NAME
+```
+
+**Options:**
+
+Name | Description|
+|----|------------|
+|_\-g, --gpu AMOUNT_|GPU quota value in hours \(h) or minutes \(m).|
+|_\-n, --non-gpu AMOUNT_|Non-GPU quota value in hours \(h) or minutes \(m).|
+|_--help_|Show this message and exit.|
+
+
+
+
+### neuro admin add-user-quota
+
+Add given values to user quota
+
+**Usage:**
+
+```bash
+neuro admin add-user-quota [OPTIONS] CLUSTER_NAME USER_NAME
+```
+
+**Options:**
+
+Name | Description|
+|----|------------|
+|_\-g, --gpu AMOUNT_|Additional GPU quota value in hours \(h) or minutes \(m).|
+|_\-n, --non-gpu AMOUNT_|Additional non-GPU quota value in hours \(h) or minutes \(m).|
+|_--help_|Show this message and exit.|
+
+
+
+
 ## neuro job
 
 Job operations.

neuromation/api/admin.py

@@ -23,6 +23,17 @@ class _ClusterUser:
     role: _ClusterUserRoleType
 
 
+@dataclass(frozen=True)
+class _Quota:
+    total_gpu_run_time_minutes: Optional[int]
+    total_non_gpu_run_time_minutes: Optional[int]
+
+
+@dataclass(frozen=True)
+class _ClusterUserWithQuota(_ClusterUser):
+    quota: _Quota
+
+
 @dataclass(frozen=True)
 class _NodePool:
     min_size: int
@@ -116,13 +127,86 @@ async def remove_cluster_user(self, cluster_name: str, user_name: str) -> None:
             # No content response
             pass
 
+    async def set_user_quota(
+        self,
+        cluster_name: str,
+        user_name: str,
+        gpu_value_minutes: Optional[float],
+        non_gpu_value_minutes: Optional[float],
+    ) -> _ClusterUserWithQuota:
+        url = (
+            self._config.admin_url
+            / "clusters"
+            / cluster_name
+            / "users"
+            / user_name
+            / "quota"
+        )
+        payload = {
+            "quota": {
+                "total_gpu_run_time_minutes": gpu_value_minutes,
+                "total_non_gpu_run_time_minutes": non_gpu_value_minutes,
+            },
+        }
+        payload["quota"] = {k: v for k, v in payload["quota"].items() if v is not None}
+
+        auth = await self._config._api_auth()
+
+        async with self._core.request("PATCH", url, json=payload, auth=auth) as resp:
+            payload = await resp.json()
+            return _cluster_user_with_quota_from_api(user_name, payload)
+
+    async def add_user_quota(
+        self,
+        cluster_name: str,
+        user_name: str,
+        additional_gpu_value_minutes: Optional[float],
+        additional_non_gpu_value_minutes: Optional[float],
+    ) -> _ClusterUserWithQuota:
+        url = (
+            self._config.admin_url
+            / "clusters"
+            / cluster_name
+            / "users"
+            / user_name
+            / "quota"
+        )
+        payload = {
+            "additional_quota": {
+                "total_gpu_run_time_minutes": additional_gpu_value_minutes,
+                "total_non_gpu_run_time_minutes": additional_non_gpu_value_minutes,
+            },
+        }
+        payload["additional_quota"] = {
+            k: v for k, v in payload["additional_quota"].items() if v is not None
+        }
+        auth = await self._config._api_auth()
+
+        async with self._core.request("PATCH", url, json=payload, auth=auth) as resp:
+            payload = await resp.json()
+            return _cluster_user_with_quota_from_api(user_name, payload)
+
 
 def _cluster_user_from_api(payload: Dict[str, Any]) -> _ClusterUser:
     return _ClusterUser(
         user_name=payload["user_name"], role=_ClusterUserRoleType(payload["role"])
     )
 
 
+def _cluster_user_with_quota_from_api(
+    user_name: str, payload: Dict[str, Any]
+) -> _ClusterUserWithQuota:
+    quota_dict = payload.get("quota", {})
+    return _ClusterUserWithQuota(
+        user_name=user_name,
+        role=_ClusterUserRoleType(payload["role"]),
+        quota=_Quota(
+            quota_dict.get("total_gpu_run_time_minutes"),
+            quota_dict.get("total_non_gpu_run_time_minutes"),
+        ),
+    )
+
+
 def _cluster_from_api(payload: Dict[str, Any]) -> _Cluster:
     if "cloud_provider" in payload:
         cloud_provider = payload["cloud_provider"]

neuromation/cli/admin.py

@@ -10,6 +10,7 @@
 from neuromation.api.admin import _ClusterUserRoleType
 
 from .formatters import ClustersFormatter, ClusterUserFormatter
+from .formatters.config import QuotaFormatter
 from .root import Root
 from .utils import async_cmd, command, group, pager_maybe
 
@@ -209,6 +210,27 @@ async def add_cluster_user(
         )
 
 
+def _parse_quota_value(
+    value: Optional[str], allow_infinity: bool = False
+) -> Optional[int]:
+    if value is None:
+        return None
+    try:
+        if value[-1] not in ("h", "m"):
+            raise ValueError(f"Unable to parse: '{value}'")
+        result = float(value[:-1]) * {"h": 60, "m": 1}[value[-1]]
+        if result < 0:
+            raise ValueError(f"Negative quota values ({value}) are not allowed")
+        if result == float("inf"):
+            if allow_infinity:
+                return None
+            else:
+                raise ValueError("Infinite quota values are not allowed")
+    except (ValueError, LookupError):
+        raise
+    return int(result)
+
+
 @command()
 @click.argument("cluster_name", required=True, type=str)
 @click.argument("user_name", required=True, type=str)
@@ -225,10 +247,98 @@ async def remove_cluster_user(root: Root, cluster_name: str, user_name: str) ->
         )
 
 
+@command()
+@click.argument("cluster_name", required=True, type=str)
+@click.argument("user_name", required=True, type=str)
+@click.option(
+    "-g",
+    "--gpu",
+    metavar="AMOUNT",
+    type=str,
+    help="GPU quota value in hours (h) or minutes (m).",
+)
+@click.option(
+    "-n",
+    "--non-gpu",
+    metavar="AMOUNT",
+    type=str,
+    help="Non-GPU quota value in hours (h) or minutes (m).",
+)
+@async_cmd()
+async def set_user_quota(
+    root: Root,
+    cluster_name: str,
+    user_name: str,
+    gpu: Optional[str],
+    non_gpu: Optional[str],
+) -> None:
+    """
+    Set user quota to given values
+    """
+    gpu_value_minutes = _parse_quota_value(gpu, allow_infinity=True)
+    non_gpu_value_minutes = _parse_quota_value(non_gpu, allow_infinity=True)
+    user_with_quota = await root.client._admin.set_user_quota(
+        cluster_name, user_name, gpu_value_minutes, non_gpu_value_minutes
+    )
+    fmt = QuotaFormatter()
+    click.echo(
+        f"New quotas for {click.style(user_with_quota.user_name, underline=True)} "
+        f"on cluster {click.style(cluster_name, underline=True)}:"
+    )
+    click.echo(fmt(user_with_quota.quota))
+
+
+@command()
+@click.argument("cluster_name", required=True, type=str)
+@click.argument("user_name", required=True, type=str)
+@click.option(
+    "-g",
+    "--gpu",
+    metavar="AMOUNT",
+    type=str,
+    help="Additional GPU quota value in hours (h) or minutes (m).",
+)
+@click.option(
+    "-n",
+    "--non-gpu",
+    metavar="AMOUNT",
+    type=str,
+    help="Additional non-GPU quota value in hours (h) or minutes (m).",
+)
+@async_cmd()
+async def add_user_quota(
+    root: Root,
+    cluster_name: str,
+    user_name: str,
+    gpu: Optional[str],
+    non_gpu: Optional[str],
+) -> None:
+    """
+    Add given values to user quota
+    """
+    additional_gpu_value_minutes = _parse_quota_value(gpu, False)
+    additional_non_gpu_value_minutes = _parse_quota_value(non_gpu, False)
+    user_with_quota = await root.client._admin.add_user_quota(
+        cluster_name,
+        user_name,
+        additional_gpu_value_minutes,
+        additional_non_gpu_value_minutes,
+    )
+    fmt = QuotaFormatter()
+    click.echo(
+        f"New quotas for {click.style(user_with_quota.user_name, underline=True)} "
+        f"on cluster {click.style(cluster_name, underline=True)}:"
+    )
+    click.echo(fmt(user_with_quota.quota))
+
+
 admin.add_command(get_clusters)
 admin.add_command(generate_cluster_config)
 admin.add_command(add_cluster)
 
 admin.add_command(get_cluster_users)
 admin.add_command(add_cluster_user)
 admin.add_command(remove_cluster_user)
+
+admin.add_command(set_user_quota)
+admin.add_command(add_user_quota)

neuromation/cli/config.py

@@ -68,6 +68,20 @@ async def show_quota(root: Root, user: Optional[str]) -> None:
     click.echo(fmt(cluster_quota))
 
 
+@command()
+@async_cmd()
+async def add_quota(root: Root) -> None:
+    """
+    Print instructions for increasing quota for current user
+    """
+    user_name = root.client.config.username
+    cluster_name = root.client.config.cluster_name
+    click.echo(
+        f"In order to increase your quota, please navigate to "
+        f"https://neuro.payments.com/{user_name}/{cluster_name}?pay=usd100"
+    )
+
+
 @command()
 @click.argument("url", required=False, default=DEFAULT_API_URL, type=URL)
 @async_cmd(init_client=False)

neuromation/cli/formatters/config.py

@@ -4,6 +4,7 @@
 from click import style
 
 from neuromation.api import Client, Cluster, Preset
+from neuromation.api.admin import _Quota
 from neuromation.api.quota import _QuotaInfo
 from neuromation.cli.utils import format_size
 
@@ -67,6 +68,29 @@ def _format_time(self, total_seconds: float) -> str:
         return f"{hours:02d}h {minutes:02d}m"
 
 
+class QuotaFormatter:
+    QUOTA_NOT_SET = "infinity"
+
+    def __call__(self, quota: _Quota) -> str:
+        gpu_details = self._format_quota_details(quota.total_gpu_run_time_minutes)
+        non_gpu_details = self._format_quota_details(
+            quota.total_non_gpu_run_time_minutes
+        )
+        return (
+            f"{style('GPU:', bold=True)}"
+            f" {gpu_details}"
+            "\n"
+            f"{style('CPU:', bold=True)}"
+            f" {non_gpu_details}"
+        )
+
+    def _format_quota_details(self, run_time_minutes: Optional[int]) -> str:
+        if run_time_minutes is None:
+            return self.QUOTA_NOT_SET
+        else:
+            return f"{run_time_minutes}m"
+
+
 class ClustersFormatter:
     def __call__(
         self, clusters: Iterable[Cluster], default_name: Optional[str]