Ensure workload policy exists on workload create (#847)

feat: ensure workload policy exists on workload create

Author: scaliby

src/xpk/commands/workload.py

@@ -54,6 +54,7 @@
 )
 from ..core.resources import get_cluster_capacity_type, get_cluster_system_characteristics, SystemCharacteristics
 from ..core.resources import CLUSTER_METADATA_CONFIGMAP, get_cluster_configmap
+from ..core.nodepool import ensure_resource_policy_exists
 from ..core.scheduling import (
     check_if_workload_can_schedule,
     create_accelerator_label,
@@ -64,6 +65,7 @@
     get_gpu_scheduler,
     create_sub_slicing_annotations,
     create_placement_policy_label,
+    get_placement_policy_name,
     is_placement_policy_supported,
 )
 from ..core.storage import (
@@ -493,6 +495,14 @@ def workload_create(args) -> None:
                 operator: NotIn
                 values: [{restart_on_exit_codes}]"""
 
+  if is_placement_policy_supported(system):
+    ensure_resource_policy_exists(
+        resource_policy_name=get_placement_policy_name(system),
+        project=args.project,
+        zone=args.zone,
+        topology=system.topology,
+    )
+
   placement_policy_label = (
       create_placement_policy_label(system)
       if is_placement_policy_supported(system)

src/xpk/core/nodepool.py

@@ -260,7 +260,12 @@ def run_gke_node_pool_create_command(
   placement_args = ''
   if is_placement_policy_supported(system):
     placement_policy = get_placement_policy_name(system)
-    ensure_resource_policy_exists(placement_policy, args, system.topology)
+    ensure_resource_policy_exists(
+        resource_policy_name=placement_policy,
+        project=args.project,
+        zone=args.zone,
+        topology=system.topology,
+    )
     placement_args = f' --placement-policy={placement_policy}'
 
   create_commands = []
@@ -587,14 +592,14 @@ def get_desired_node_pool_names(
 
 
 def ensure_resource_policy_exists(
-    resource_policy_name: str, args, topology: str
+    resource_policy_name: str, project: str, zone: str, topology: str
 ) -> None:
   return_code, _ = run_command_for_value(
       (
           'gcloud compute resource-policies describe'
           f' {resource_policy_name} '
-          f'--project={args.project} '
-          f'--region={zone_to_region(args.zone)}'
+          f'--project={project} '
+          f'--region={zone_to_region(zone)}'
       ),
       'Retrieve resource policy',
   )
@@ -605,7 +610,7 @@ def ensure_resource_policy_exists(
   return_code, _ = run_command_for_value(
       (
           'gcloud compute resource-policies create workload-policy'
-          f' {resource_policy_name} --project={args.project} --region={zone_to_region(args.zone)} --type=HIGH_THROUGHPUT'
+          f' {resource_policy_name} --project={project} --region={zone_to_region(zone)} --type=HIGH_THROUGHPUT'
           f' --accelerator-topology={topology}'
       ),
       'Create resource policy',

src/xpk/core/nodepool_test.py

@@ -96,7 +96,12 @@ def test_ensure_resource_policy_exists_with_existing_policy_retrieves_existing_p
   mock = mocker.patch(
       "xpk.core.nodepool.run_command_for_value", return_value=(0, "")
   )
-  ensure_resource_policy_exists("resource-policy", args, "2x2x1")
+  ensure_resource_policy_exists(
+      resource_policy_name="resource-policy",
+      project="test-project",
+      zone="us-central1-a",
+      topology="2x2x1",
+  )
   mock.assert_called_once()
 
 
@@ -108,7 +113,12 @@ def test_ensure_resource_policy_exists_without_existing_policy_creates_policy(
   mock = mocker.patch(
       "xpk.core.nodepool.run_command_for_value", side_effect=[(1, ""), (0, "")]
   )
-  ensure_resource_policy_exists("resource-policy", args, "2x2x1")
+  ensure_resource_policy_exists(
+      resource_policy_name="resource-policy",
+      project="test-project",
+      zone="us-central1-a",
+      topology="2x2x1",
+  )
   assert mock.call_count == 2
   assert mock.call_args_list[0].args[1] == "Retrieve resource policy"
 
@@ -125,7 +135,12 @@ def test_ensure_resource_policy_exits_without_existing_policy_throws_when_creati
         "xpk.core.nodepool.run_command_for_value",
         side_effect=[(1, ""), (1, "")],
     )
-    ensure_resource_policy_exists("resource-policy", args, "2x2x1")
+    ensure_resource_policy_exists(
+        resource_policy_name="resource-policy",
+        project="test-project",
+        zone="us-central1-a",
+        topology="2x2x1",
+    )
 
 
 @pytest.fixture