Typing args (initial)

pyproject.toml

@@ -69,7 +69,7 @@ dev = [
 version = {attr = "xpk.core.config.__version__"}
 
 [tool.setuptools]
-packages = ["xpk", "xpk.parser", "xpk.core", "xpk.commands", "xpk.api", "xpk.templates", "xpk.utils", "xpk.core.blueprint", "xpk.core.remote_state", "xpk.core.workload_decorators"]
+packages = ["xpk", "xpk.parser", "xpk.core", "xpk.commands", "xpk.api", "xpk.args", "xpk.templates", "xpk.utils", "xpk.core.blueprint", "xpk.core.remote_state", "xpk.core.workload_decorators"]
 package-dir = {"" = "src"}
 package-data = {"xpk.api" = ["storage_crd.yaml"], "xpk.templates" = ["storage.yaml"]}
 

src/xpk/args/__init__.py

@@ -0,0 +1,15 @@
+"""
+Copyright 2024 Google LLC
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+     https://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+"""

src/xpk/args/batch.py

@@ -0,0 +1,22 @@
+"""
+Copyright 2024 Google LLC
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+     https://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+"""
+
+from .cluster import ClusterConfig
+from .slurm import SlurmConfig
+
+
+class BatchArgs(ClusterConfig, SlurmConfig):
+  script: str = None

src/xpk/args/cluster.py

@@ -0,0 +1,24 @@
+"""
+Copyright 2025 Google LLC
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+     https://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+"""
+
+from .gcloud_context import GcloudConfig
+
+
+class ClusterConfig(GcloudConfig):
+  """Class representing cluster config"""
+
+  kind_cluster: bool = False
+  cluster: str = None

src/xpk/args/common.py

@@ -0,0 +1,21 @@
+"""
+Copyright 2025 Google LLC
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+     https://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+"""
+
+
+class GlobalConfig:
+  """Class representing global args type"""
+
+  dry_run: bool = False

src/xpk/args/gcloud_context.py

@@ -0,0 +1,128 @@
+"""
+Copyright 2025 Google LLC
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+     https://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+"""
+
+from typing import Optional
+
+from google.api_core.exceptions import PermissionDenied
+from google.cloud import resourcemanager_v3
+
+import subprocess
+import sys
+from ..utils.console import xpk_exit, xpk_print
+from .common import GlobalConfig
+
+
+def get_project():
+  """Get GCE project from `gcloud config get project`.
+
+  Returns:
+     The project name.
+  """
+  completed_command = subprocess.run(
+      ['gcloud', 'config', 'get', 'project'], check=True, capture_output=True
+  )
+  project_outputs = completed_command.stdout.decode().strip().split('\n')
+  if len(project_outputs) < 1 or project_outputs[-1] == '':
+    sys.exit(
+        'You must specify the project in the project flag or set it with'
+        " 'gcloud config set project <project>'"
+    )
+  return project_outputs[
+      -1
+  ]  # The project name lives on the last line of the output
+
+
+def get_zone():
+  """Get GCE zone from `gcloud config get compute/zone`.
+
+  Returns:
+     The zone name.
+  """
+  completed_command = subprocess.run(
+      ['gcloud', 'config', 'get', 'compute/zone'],
+      check=True,
+      capture_output=True,
+  )
+  zone_outputs = completed_command.stdout.decode().strip().split('\n')
+  if len(zone_outputs) < 1 or zone_outputs[-1] == '':
+    sys.exit(
+        "You must specify the zone in the zone flag or set it with 'gcloud"
+        " config set compute/zone <zone>'"
+    )
+  return zone_outputs[-1]  # The zone name lives on the last line of the output
+
+
+class GcloudConfig(GlobalConfig):
+  """Class representing gcloud project config"""
+
+  gke_version: Optional[str] = None
+
+  _zone: Optional[str] = None
+  _project: Optional[str] = None
+  _project_number: Optional[str] = None
+
+  @property
+  def zone(self) -> str:
+    if self._zone is None:
+      self._zone = get_zone()
+      if self._project is None:
+        self._project = get_project()
+      xpk_print(f'Working on {self._project} and {self._zone}')
+    return str(self._zone)
+
+  @zone.setter
+  def zone(self, value: str):
+    self._zone = value
+
+  @property
+  def region(self):
+    return '-'.join(self.zone.split('-')[:2])
+
+  @property
+  def project(self) -> str:
+    if self._project is None:
+      self._project = get_project()
+      if self._zone is None:
+        self._zone = get_zone()
+      xpk_print(f'Working on {self._project} and {self._zone}')
+    return str(self._project)
+
+  @project.setter
+  def project(self, value: str):
+    self._project = value
+
+  @property
+  def project_number(self) -> str:
+    if self._project_number is None:
+      client = resourcemanager_v3.ProjectsClient()
+      request = resourcemanager_v3.GetProjectRequest()
+      request.name = f'projects/{self.project}'
+      try:
+        response = client.get_project(request=request)
+      except PermissionDenied as e:
+        xpk_print(
+            f"Couldn't translate project id: {self.project} to project number."
+            f' Error: {e}'
+        )
+        xpk_exit(1)
+      parts = response.name.split('/', 1)
+      xpk_print(f'Project number for project: {self.project} is {parts[1]}')
+      self._project_number = str(parts[1])
+    return str(self._project_number)
+
+  @project_number.setter
+  def project_number(self, value: str):
+    self._project_number = value

src/xpk/args/slurm.py

@@ -0,0 +1,38 @@
+"""
+Copyright 2025 Google LLC
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+     https://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+"""
+
+from typing import Optional
+
+
+class SlurmConfig:
+  """Class representing slurm args type"""
+
+  ignore_unknown_flags: bool = False
+  array: Optional[str] = None
+  cpus_per_task: Optional[str] = None
+  gpus_per_task: Optional[str] = None
+  mem: Optional[str] = None
+  mem_per_task: Optional[str] = None
+  mem_per_cpu: Optional[str] = None
+  mem_per_gpu: Optional[str] = None
+  nodes: Optional[int] = None
+  ntasks: Optional[int] = None
+  output: Optional[str] = None
+  error: Optional[str] = None
+  input: Optional[str] = None
+  job_name: Optional[str] = None
+  chdir: Optional[str] = None
+  time: Optional[str] = None

src/xpk/args/storage.py

@@ -0,0 +1,85 @@
+"""
+Copyright 2024 Google LLC
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+     https://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+"""
+
+from typing import Optional, Literal, TypeAlias
+from .cluster import ClusterConfig
+
+
+StorageType: TypeAlias = Literal[
+    'gcsfuse', 'gcpfilestore', 'parallelstore', 'pd'
+]
+
+StorageAccessMode: TypeAlias = Literal[
+    'ReadWriteOnce', 'ReadOnlyMany', 'ReadWriteMany'
+]
+
+FilestoreTier: TypeAlias = Literal[
+    'BASIC_HDD', 'BASIC_SSD', 'ZONAL', 'REGIONAL', 'ENTERPRISE'
+]
+
+
+class StorageAttachArgs(ClusterConfig):
+  """Class representing storage attach args type"""
+
+  name: str = None
+  type: StorageType = None
+  auto_mount: bool = None
+  mount_point: str = None
+  readonly: bool = None
+  size: Optional[int] = None
+  bucket: Optional[str] = None
+  vol: Optional[str] = None
+  access_mode: StorageAccessMode = 'ReadWriteMany'
+  instance: Optional[str] = None
+  prefetch_metadata: bool = True
+  manifest: Optional[str] = None
+  mount_options: Optional[str] = 'implicit-dirs'
+
+
+class StorageCreateArgs(ClusterConfig):
+  """Class representing storage create args type"""
+
+  name: str = None
+  access_mode: StorageAccessMode = 'ReadWriteMany'
+  vol: str = 'default'
+  size: int = None
+  tier: FilestoreTier = 'BASIC_HDD'
+  type: Literal['gcpfilestore'] = 'gcpfilestore'
+  auto_mount: bool = None
+  mount_point: str = None
+  readonly: bool = None
+  instance: Optional[str] = None
+  manifest: Optional[str] = None
+  mount_options: Optional[str] = 'implicit-dirs'
+
+
+class StorageDeleteArgs(ClusterConfig):
+  """Class representing storage delete args type"""
+
+  name: str = None
+  force: Optional[bool] = False
+
+
+class StorageDetachArgs(ClusterConfig):
+  """Class representing storage detach args type"""
+
+  name: str = None
+
+
+class StorageListArgs(ClusterConfig):
+  """Class representing storage list args type"""
+
+  pass

src/xpk/args/utils.py

@@ -0,0 +1,35 @@
+"""
+Copyright 2024 Google LLC
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+     https://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+"""
+
+import inspect
+from argparse import Namespace
+from typing import Any
+
+
+def apply_args(main_args: Namespace, annotation: Any) -> Any:
+  args = annotation()
+
+  # getters and setters
+  for param in inspect.get_annotations(annotation):
+    if param in main_args:
+      setattr(args, param, getattr(main_args, param))
+
+  # parameters
+  for param, _ in inspect.getmembers(annotation):
+    if param in main_args:
+      setattr(args, param, getattr(main_args, param))
+
+  return args  # pytype: disable=bad-return-type