feat: implement API and CLI for setting resource quota on workspace

harness/determined/cli/workspace.py

@@ -166,6 +166,18 @@ def add_workspace_namespace_binding(args: argparse.Namespace):
     print("Workspace " + workspace_name + " is bound to " + n)
     return None
 
+def set_quota(args: argparse.Namespace):
+    sess = cli.setup_session(args)
+    w = api.workspace_by_name(sess, args.workspace_name)
+    content = bindings.v1SetResourceQuotaRequest(
+        id=w.id,
+        quota=args.resource_quota,
+        clusterName=args.cluster_name,
+    )
+    ws_name = str(args.workspace_name)
+    bindings.post_SetResourceQuota(sess, body=content, id=w.id)
+    print("Successfully set resource quota " + str(args.resource_quota) + " for workspace " + str(ws_name))
+    return None
 
 def _parse_agent_user_group_args(args: argparse.Namespace) -> Optional[bindings.v1AgentUserGroup]:
     if args.agent_uid or args.agent_gid or args.agent_user or args.agent_group:
@@ -194,12 +206,23 @@ def _parse_checkpoint_storage_args(args: argparse.Namespace) -> Any:
 def create_workspace(args: argparse.Namespace) -> None:
     agent_user_group = _parse_agent_user_group_args(args)
     checkpoint_storage = _parse_checkpoint_storage_args(args)
+
+    if args.namespace and args.auto_create_namespace:
+         raise api.errors.BadRequestException(
+            "must provide either --auto-create-namespace or --namespace NAMESPACE"
+        )
 
-    if args.cluster_name and (not args.namespace and not args.auto_create_namespace):
+    if (args.cluster_name or args.quota) and (not args.namespace and not args.auto_create_namespace):
         raise api.errors.BadRequestException(
             "must provide either --auto-create-namespace or --namespace NAMESPACE"
         )
-
+    if args.namespace and args.quota:
+        raise api.errors.BadRequestException(
+            "you cannot set the quota on a preexisting namesapce. If you would like to set the \
+            quota on the namespace bound to your workspace, please specify the \
+            --auto-create-namesapce flag, and we will create a Kubernetes namesapce for \
+            you!"
+        )
     if (args.namespace or args.auto_create_namespace) and not args.cluster_name:
         raise api.errors.BadRequestException(
             "must specify --cluster-name CLUSTER_NAME"
@@ -213,10 +236,17 @@ def create_workspace(args: argparse.Namespace) -> None:
         defaultComputePool=args.default_compute_pool,
         defaultAuxPool=args.default_aux_pool,
         clusterName=args.cluster_name,
+        quota=args.quota,
         namespaceName=args.namespace,
         autoCreateNamespace=args.auto_create_namespace,
     )
     w = bindings.post_PostWorkspace(sess, body=content).workspace
+
+    if args.namespace or args.auto_create_namespace:
+        ns = args.namespace
+        if args.auto_create_namespace:
+            ns = "det-" + w.name
+        print("bound workspace " + w.name + " to namespace " + ns)
 
     if args.json:
         render.print_json(w.to_json())
@@ -420,6 +450,8 @@ def yaml_file_arg(val: str) -> Any:
                     cli.Arg("--json", action="store_true", help="print as JSON"),
                     cli.Arg("--cluster-name", type=str, help="cluster within which we create the \
                         workspace-namespace binding"),
+                    cli.Arg("-q", "--quota", type=int, help="resource quota to which the \
+                        namespace is limited"),
                     cli.Group(
                          cli.Arg("-n", "--namespace", type=str, help="existing namespace to which \
                             the workspace is bound"),
@@ -483,6 +515,23 @@ def yaml_file_arg(val: str) -> Any:
                                 the workspace is bound."),
                         ],
                     ),
+                    cli.Cmd(
+                        "set",
+                        None,
+                        "manage quotas",
+                        [
+                            cli.Cmd("quota",
+                                    set_quota,
+                                    "set quota",
+                                [
+                                    cli.Arg("workspace_name", type=str, help="name of the workspace"), 
+                                    cli.Arg("resource_quota", type=int, help="quota for the workspace"),
+                                    cli.Arg("cluster_name", type=str, help="cluster within which we create \
+                                    the quota"),
+                                ],
+                            ),
+                        ],
+                    ),
                 ],
             ),
             cli.Cmd(

master/internal/api_workspace.go

@@ -61,11 +61,15 @@ func maskStorageConfigSecrets(w *workspacev1.Workspace) error {
 }
 
 func validatePostWorkspaceRequest(req *apiv1.PostWorkspaceRequest) error {
-	if req.ClusterName != nil &&
+	if (req.ClusterName != nil || req.Quota != nil) &&
 		(req.NamespaceName == nil && !req.AutoCreateNamespace) {
 		return status.Errorf(codes.InvalidArgument,
 			"Must specify either an existing Kubernetes namespace or indicate that you would like a namespace to be auto-created")
 	}
+	if req.NamespaceName != nil && req.Quota != nil {
+		return status.Errorf(codes.InvalidArgument,
+			"You cannot set a quota on an already existing k8s namespace.")
+	}
 	if (req.AutoCreateNamespace || req.NamespaceName != nil) && req.ClusterName == nil {
 		return status.Errorf(codes.InvalidArgument,
 			"You must specify a cluster for the specified namespace that you would like to bind.")
@@ -418,6 +422,17 @@ func (a *apiServer) PostWorkspace(
 		if err != nil {
 			return nil, fmt.Errorf("Failed to create namespace binding: %w", err)
 		}
+		if req.Quota != nil {
+			quotaReq := &apiv1.SetResourceQuotaRequest{
+				Quota:       *req.Quota,
+				ClusterName: *req.ClusterName,
+				Id:          int32(w.ID),
+			}
+			_, err = a.setResourceQuota(ctx, quotaReq, &tx, w)
+			if err != nil {
+				return nil, errors.Wrap(err, "failed to set the resource quota")
+			}
+		}
 	}
 
 	pin := &model.WorkspacePin{WorkspaceID: w.ID, UserID: w.UserID}
@@ -661,6 +676,48 @@ func (a *apiServer) ModifyWorkspaceNamespaceBinding(ctx context.Context,
 	return res, nil
 }
 
+func (a *apiServer) setResourceQuota(ctx context.Context, req *apiv1.SetResourceQuotaRequest,
+	tx *bun.Tx, w *model.Workspace) (*apiv1.SetResourceQuotaResponse, error) {
+	// Get the namespace bound to the workspace from the db. If the workspace is not bound to a namespace,
+	// return an error.
+	var wsns model.WorkspaceNamespace
+	err := tx.NewSelect().Model(&model.WorkspaceNamespace{}).Where("workspace_id = ?", req.Id).Scan(ctx, &wsns)
+	if err != nil {
+		if err == sql.ErrNoRows {
+			return nil, errors.Wrapf(err, "workspace %s has no namespace binding", w.Name)
+		}
+		return nil, errors.Wrap(err, "error getting workspace-namespace binding")
+	}
+
+	namespaceName, err := generateNamespaceName(w.Name)
+	if err != nil {
+		return nil, errors.Wrapf(err, "error getting generated namespace name for workspace %s", w.Name)
+	}
+	if *namespaceName != wsns.NamespaceName {
+		return nil, errors.Wrapf(err, "cannot set quota on a workspace that is not bound to an auto-created namespace")
+	}
+	err = a.m.rm.SetQuota(int(req.Quota), wsns.NamespaceName, req.ClusterName)
+	if err != nil {
+		return nil, errors.Wrap(err, "Failed to create quota in Kubernetes")
+	}
+	resp := apiv1.SetResourceQuotaResponse{}
+	return &resp, nil
+}
+
+func (a *apiServer) SetResourceQuota(ctx context.Context, req *apiv1.SetResourceQuotaRequest) (*apiv1.SetResourceQuotaResponse, error) {
+	license.RequireLicense("set resource quota")
+	tx, err := db.Bun().BeginTx(ctx, nil)
+	if err != nil {
+		return nil, err
+	}
+	var w model.Workspace
+	err = tx.NewSelect().Model(&model.Workspace{}).Where("id = ?", req.Id).Scan(ctx, &w)
+	if err != nil {
+		return nil, errors.Wrapf(err, "workspace with name %s not found", w.Name)
+	}
+	return a.setResourceQuota(ctx, req, &tx, &w)
+}
+
 func (a *apiServer) deleteWorkspace(
 	ctx context.Context, workspaceID int32, projects []*projectv1.Project,
 ) {

master/internal/rm/agentrm/agent_resource_manager.go

@@ -589,6 +589,10 @@ func (a *ResourceManager) DeleteNamespace(namespaceName string) (*string, error)
 	return nil, nil
 }
 
+func (a *ResourceManager) SetQuota(quota int, namespaceName string, clusterName string) error {
+	return fmt.Errorf("Cannot set quota with resource manager type AgentRM.")
+}
+
 func (a *ResourceManager) createResourcePool(
 	db db.DB, config config.ResourcePoolConfig, cert *tls.Certificate,
 ) (*resourcePool, error) {

master/internal/rm/dispatcherrm/dispatcher_resource_manager.go

@@ -675,6 +675,10 @@ func (*DispatcherResourceManager) DeleteNamespace(string) (*string, error) {
 	return nil, rmerrors.ErrNotSupported
 }
 
+func (*DispatcherResourceManager) SetQuota(int, string, string) error {
+	return rmerrors.ErrNotSupported
+}
+
 // ResolveResourcePool returns the resolved slurm partition or an error if it doesn't exist or
 // can't be resolved due to internal errors.
 // Note to developers: this function doesn't acquire a lock and, ideally, we won't make it, since

master/internal/rm/kubernetesrm/kubernetes_resource_manager.go

@@ -409,6 +409,19 @@ func (k *ResourceManager) DeleteNamespace(namespaceName string) (*string, error)
 	return nil, nil
 }
 
+// SetQuota implements rm.ResourceManager.
+func (k *ResourceManager) SetQuota(quota int, namespaceName string, clusterName string) error {
+	configClusterName := rm.ClusterName(k.config.ClusterName)
+	if clusterName != configClusterName.String() {
+		return nil
+	}
+	err := k.podsService.SetQuota(quota, namespaceName)
+	if err != nil {
+		return fmt.Errorf("error creating quota %s: %w", namespaceName, err)
+	}
+	return nil
+}
+
 // getResourcePoolRef gets an actor ref to a resource pool by name.
 func (k ResourceManager) resourcePoolExists(name string) error {
 	resp, err := k.GetResourcePools()