/ mlflow / entities / workspace.py
workspace.py
 1  """Workspace entity shared between server and stores."""
 2  
 3  from __future__ import annotations
 4  
 5  from dataclasses import dataclass
 6  from enum import Enum
 7  from typing import Any
 8  
 9  from mlflow.protos.service_pb2 import Workspace as ProtoWorkspace
10  
11  
12  class WorkspaceDeletionMode(str, Enum):
13      """Controls what happens to resources when a workspace is deleted."""
14  
15      SET_DEFAULT = "SET_DEFAULT"
16      """Reassign all resources in the workspace to the default workspace."""
17  
18      CASCADE = "CASCADE"
19      """Delete all resources in the workspace."""
20  
21      RESTRICT = "RESTRICT"
22      """Refuse to delete the workspace if it still contains resources."""
23  
24  
25  @dataclass(frozen=True, slots=True)
26  class Workspace:
27      """Minimal metadata describing a workspace."""
28  
29      name: str
30      description: str | None = None
31      default_artifact_root: str | None = None
32  
33      def to_dict(self) -> dict[str, str | None]:
34          return {
35              "name": self.name,
36              "description": self.description,
37              "default_artifact_root": self.default_artifact_root,
38          }
39  
40      @classmethod
41      def from_dict(cls, payload: dict[str, Any]) -> "Workspace":
42          return cls(
43              name=payload["name"],
44              description=payload.get("description"),
45              default_artifact_root=payload.get("default_artifact_root"),
46          )
47  
48      def to_proto(self) -> ProtoWorkspace:
49          workspace = ProtoWorkspace()
50          workspace.name = self.name
51          if self.description is not None:
52              workspace.description = self.description
53          if self.default_artifact_root is not None:
54              workspace.default_artifact_root = self.default_artifact_root
55          return workspace
56  
57      @classmethod
58      def from_proto(cls, proto: ProtoWorkspace) -> "Workspace":
59          description = proto.description if proto.HasField("description") else None
60          default_artifact_root = (
61              proto.default_artifact_root if proto.HasField("default_artifact_root") else None
62          )
63          return cls(
64              name=proto.name,
65              description=description,
66              default_artifact_root=default_artifact_root,
67          )