1from pydantic import BaseModel, ConfigDict, Field, field_validator
2from typing import Optional, List, Dict, Literal, Any
3from enum import Enum
4
5# Minimum amazon-sagemaker-spaces addon version required for this template version
6MIN_ADDON_VERSION = "0.1.6"
7
8
[docs]
9class OwnershipType(str, Enum):
10 PUBLIC = "Public"
11 OWNER_ONLY = "OwnerOnly"
12
13
[docs]
14class DesiredStatus(str, Enum):
15 RUNNING = "Running"
16 STOPPED = "Stopped"
17
18
[docs]
19class VolumeSpec(BaseModel):
20 """VolumeSpec defines a volume to mount from an existing PVC"""
21 name: str = Field(
22 description="Name is a unique identifier for this volume within the pod (maps to pod.spec.volumes[].name)",
23 min_length=1
24 )
25 mount_path: str = Field(
26 alias="mountPath",
27 description="MountPath is the path where the volume should be mounted (Unix-style path, e.g. /data)",
28 min_length=1
29 )
30 persistent_volume_claim_name: str = Field(
31 alias="persistentVolumeClaimName",
32 description="PersistentVolumeClaimName is the name of the existing PVC to mount",
33 min_length=1
34 )
35
36
[docs]
37class ContainerConfig(BaseModel):
38 """ContainerConfig defines container command and args configuration"""
39 command: Optional[List[str]] = Field(
40 default=None,
41 description="Command specifies the container command"
42 )
43 args: Optional[List[str]] = Field(
44 default=None,
45 description="Args specifies the container arguments"
46 )
47
48
[docs]
49class AccessStrategyRef(BaseModel):
50 """AccessStrategyRef references a WorkspaceAccessStrategy"""
51 name: str = Field(
52 description="Name of the WorkspaceAccessStrategy"
53 )
54 namespace: Optional[str] = Field(
55 default=None,
56 description="Namespace where the WorkspaceAccessStrategy is located"
57 )
58
59
[docs]
60class TemplateRef(BaseModel):
61 """TemplateRef defines a reference to a WorkspaceTemplate"""
62 name: str = Field(
63 description="Name of the WorkspaceTemplate"
64 )
65 namespace: Optional[str] = Field(
66 default=None,
67 description="Namespace where the WorkspaceTemplate is located"
68 )
69
70
[docs]
71class IdleDetectionSpec(BaseModel):
72 """IdleDetectionSpec defines idle detection methods"""
73 http_get: Optional[Dict[str, Any]] = Field(
74 default=None,
75 alias="httpGet",
76 description="HTTPGet specifies the HTTP request to perform for idle detection"
77 )
78
79
[docs]
80class IdleShutdownSpec(BaseModel):
81 """IdleShutdownSpec defines idle shutdown configuration"""
82 enabled: bool = Field(
83 description="Enabled indicates if idle shutdown is enabled"
84 )
85 idle_timeout_in_minutes: int = Field(
86 alias="idleTimeoutInMinutes",
87 description="IdleTimeoutInMinutes specifies idle timeout in minutes",
88 ge=1
89 )
90 detection: IdleDetectionSpec = Field(
91 description="Detection specifies how to detect idle state"
92 )
93
94
[docs]
95class StorageSpec(BaseModel):
96 """StorageSpec defines the storage configuration for Workspace"""
97 storage_class_name: Optional[str] = Field(
98 default=None,
99 alias="storageClassName",
100 description="StorageClassName specifies the storage class to use for persistent storage"
101 )
102 size: Optional[str] = Field(
103 default="10Gi",
104 description="Size specifies the size of the persistent volume. Supports standard Kubernetes resource quantities (e.g., '10Gi', '500Mi', '1Ti'). Integer values without units are interpreted as bytes"
105 )
106 mount_path: Optional[str] = Field(
107 default="/home",
108 alias="mountPath",
109 description="MountPath specifies where to mount the persistent volume in the container. Default is /home/jovyan (jovyan is the standard user in Jupyter images)"
110 )
111
112
[docs]
113class ResourceRequirements(BaseModel):
114 """ResourceRequirements describes the compute resource requirements"""
115 requests: Optional[Dict[str, Optional[str]]] = Field(
116 default=None,
117 description="Requests describes the minimum amount of compute resources required. If Requests is omitted for a container, it defaults to Limits if that is explicitly specified, otherwise to an implementation-defined value. Requests cannot exceed Limits."
118 )
119 limits: Optional[Dict[str, Optional[str]]] = Field(
120 default=None,
121 description="Limits describes the maximum amount of compute resources allowed."
122 )
123
124
[docs]
125class SpaceConfig(BaseModel):
126 """SpaceConfig defines the desired state of a Space"""
127 model_config = ConfigDict(extra="forbid")
128
129 name: str = Field(
130 description="Space name",
131 min_length=1,
132 max_length=63,
133 pattern=r'^[a-z0-9]([-a-z0-9]*[a-z0-9])?$'
134 )
135 display_name: str = Field(
136 alias="display_name",
137 description="Display Name of the space",
138 min_length=1
139 )
140 namespace: str = Field(
141 default="default",
142 description="Kubernetes namespace",
143 min_length=1
144 )
145 image: Optional[str] = Field(
146 default=None,
147 description="Image specifies the container image to use"
148 )
149 desired_status: Optional[DesiredStatus] = Field(
150 default=None,
151 alias="desired_status",
152 description="DesiredStatus specifies the desired operational status"
153 )
154 ownership_type: Optional[OwnershipType] = Field(
155 default=None,
156 alias="ownership_type",
157 description="OwnershipType specifies who can modify the space. 'Public' means anyone with RBAC permissions can update/delete the space. 'OwnerOnly' means only the creator can update/delete the space."
158 )
159 resources: Optional[ResourceRequirements] = Field(
160 default=None,
161 description="Resources specifies the resource requirements"
162 )
163 storage: Optional[StorageSpec] = Field(
164 default=None,
165 description="Storage specifies the storage configuration"
166 )
167 volumes: Optional[List[VolumeSpec]] = Field(
168 default=None,
169 description="Volumes specifies additional volumes to mount from existing PersistentVolumeClaims"
170 )
171 container_config: Optional[ContainerConfig] = Field(
172 default=None,
173 alias="container_config",
174 description="ContainerConfig specifies container command and args configuration"
175 )
176 node_selector: Optional[Dict[str, str]] = Field(
177 default=None,
178 alias="node_selector",
179 description="NodeSelector specifies node selection constraints for the space pod (JSON string)"
180 )
181 affinity: Optional[Dict[str, Any]] = Field(
182 default=None,
183 description="Affinity specifies node affinity and anti-affinity rules for the space pod (JSON string)"
184 )
185 tolerations: Optional[List[Dict[str, Any]]] = Field(
186 default=None,
187 description="Tolerations specifies tolerations for the space pod to schedule on nodes with matching taints (JSON string)"
188 )
189 lifecycle: Optional[Dict[str, Any]] = Field(
190 default=None,
191 description="Lifecycle specifies actions that the management system should take in response to container lifecycle events (JSON string)"
192 )
193 template_ref: Optional[TemplateRef] = Field(
194 default=None,
195 alias="template_ref",
196 description="TemplateRef references a WorkspaceTemplate to use as base configuration. When set, template provides defaults and workspace spec fields act as overrides"
197 )
198 idle_shutdown: Optional[IdleShutdownSpec] = Field(
199 default=None,
200 alias="idle_shutdown",
201 description="IdleShutdown specifies idle shutdown configuration"
202 )
203 app_type: Optional[str] = Field(
204 default=None,
205 alias="app_type",
206 description="AppType specifies the application type for this workspace"
207 )
208 service_account_name: Optional[str] = Field(
209 default=None,
210 alias="service_account_name",
211 description="ServiceAccountName specifies the name of the ServiceAccount to use for the workspace pod"
212 )
213 access_type: Optional[OwnershipType] = Field(
214 default=None,
215 alias="access_type",
216 description="AccessType specifies who can connect to the workspace. 'Public' means anyone with RBAC permissions can connect. 'OwnerOnly' means only the creator can connect."
217 )
218 env: Optional[List[Dict[str, Any]]] = Field(
219 default=None,
220 description="Environment variables for the workspace container (list of {name, value} objects)"
221 )
222 access_strategy: Optional[AccessStrategyRef] = Field(
223 default=None,
224 alias="access_strategy",
225 description="AccessStrategy references a WorkspaceAccessStrategy to use"
226 )
227 pod_security_context: Optional[Dict[str, Any]] = Field(
228 default=None,
229 alias="pod_security_context",
230 description="Pod-level security context. Overrides template defaults when specified (JSON string)"
231 )
232 container_security_context: Optional[Dict[str, Any]] = Field(
233 default=None,
234 alias="container_security_context",
235 description="Container-level security context for the main workspace container. Overrides template defaults (JSON string)"
236 )
237 init_containers: Optional[List[Dict[str, Any]]] = Field(
238 default=None,
239 alias="init_containers",
240 description="Init containers to run before the workspace container starts (JSON string, max 10)"
241 )
242 queue_name: Optional[str] = Field(
243 default=None,
244 alias="queue_name",
245 description="Queue name for space scheduling",
246 min_length=1,
247 max_length=63,
248 pattern=r'^[a-z0-9]([-a-z0-9]*[a-z0-9])?$'
249 )
250 priority: Optional[str] = Field(
251 default=None,
252 description="Priority class for space scheduling",
253 min_length=1
254 )
255
[docs]
256 @field_validator('volumes')
257 def validate_no_duplicate_volumes(cls, v):
258 """Validate no duplicate volume names or mount paths."""
259 if not v:
260 return v
261
262 # Check for duplicate volume names
263 names = [vol.name for vol in v]
264 if len(names) != len(set(names)):
265 raise ValueError("Duplicate volume names found")
266
267 # Check for duplicate mount paths
268 mount_paths = [vol.mount_path for vol in v]
269 if len(mount_paths) != len(set(mount_paths)):
270 raise ValueError("Duplicate mount paths found")
271
272 return v
273
[docs]
274 def to_domain(self) -> Dict:
275 """
276 Convert flat config to domain model for space creation
277 """
278 # Create the space spec
279 spec = {
280 "displayName": self.display_name
281 }
282
283 # Add optional spec fields
284 if self.image is not None:
285 spec["image"] = self.image
286 if self.desired_status is not None:
287 spec["desiredStatus"] = self.desired_status.value
288 if self.ownership_type is not None:
289 spec["ownershipType"] = self.ownership_type.value
290 if self.resources is not None:
291 spec["resources"] = self.resources.model_dump(exclude_none=True)
292 if self.storage is not None:
293 spec["storage"] = self.storage.model_dump(exclude_none=True, by_alias=True)
294 if self.volumes is not None:
295 spec["volumes"] = [vol.model_dump(exclude_none=True, by_alias=True) for vol in self.volumes]
296 if self.container_config is not None:
297 spec["containerConfig"] = self.container_config.model_dump(exclude_none=True)
298 if self.node_selector is not None:
299 spec["nodeSelector"] = self.node_selector
300 if self.affinity is not None:
301 spec["affinity"] = self.affinity
302 if self.tolerations is not None:
303 spec["tolerations"] = self.tolerations
304 if self.lifecycle is not None:
305 spec["lifecycle"] = self.lifecycle
306 if self.template_ref is not None:
307 spec["templateRef"] = self.template_ref.model_dump(exclude_none=True, by_alias=True)
308 if self.idle_shutdown is not None:
309 spec["idleShutdown"] = self.idle_shutdown.model_dump(exclude_none=True, by_alias=True)
310 if self.app_type is not None:
311 spec["appType"] = self.app_type
312 if self.service_account_name is not None:
313 spec["serviceAccountName"] = self.service_account_name
314 if self.access_type is not None:
315 spec["accessType"] = self.access_type.value
316 if self.env is not None:
317 spec["env"] = self.env
318 if self.access_strategy is not None:
319 spec["accessStrategy"] = self.access_strategy.model_dump(exclude_none=True)
320 if self.pod_security_context is not None:
321 spec["podSecurityContext"] = self.pod_security_context
322 if self.container_security_context is not None:
323 spec["containerSecurityContext"] = self.container_security_context
324 if self.init_containers is not None:
325 spec["initContainers"] = self.init_containers
326
327 # Create metadata
328 metadata = {"name": self.name}
329 if self.namespace is not None:
330 metadata["namespace"] = self.namespace
331
332 # Add kueue labels for task governance
333 labels = {}
334 if self.queue_name is not None:
335 labels["kueue.x-k8s.io/queue-name"] = self.queue_name
336 if self.priority is not None:
337 labels["kueue.x-k8s.io/priority-class"] = self.priority
338 if labels:
339 metadata["labels"] = labels
340
341 # Create the complete space configuration
342 space_config = {
343 "apiVersion": "workspace.jupyter.org/v1alpha1",
344 "kind": "Workspace",
345 "metadata": metadata,
346 "spec": spec
347 }
348
349 return {
350 "name": self.name,
351 "namespace": self.namespace,
352 "space_spec": space_config
353 }
