xpk 1.1.0__tar.gz → 1.1.1__tar.gz

{xpk-1.1.0/src/xpk.egg-info → xpk-1.1.1}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: xpk
-Version: 1.1.0
+Version: 1.1.1
 Summary: xpk helps Cloud developers to orchestrate training jobs on accelerators on GKE.
 Author-email: XPK team <xpk-code-reviewers@google.com>
 License: Apache-2.0

{xpk-1.1.0 → xpk-1.1.1}/recipes/Workload_create_super-slicing.md

@@ -47,7 +47,7 @@ docker buildx build --platform=linux/amd64 -f 4b6736a12db8ea0f78ce793fd0d4ee0c94
 docker tag dry-run-runner gcr.io/golden-project/dry-run-runner:prefix-current
 [XPK] Task: `Upload Docker Image` is implemented by the following command not running since it is a dry run. 
 docker push gcr.io/golden-project/dry-run-runner:prefix-current
-[XPK] Temp file (608e1382aabe2b0335855e5e99876a2e67de954453ebfa4cf12eb82c966f85da) content: 
+[XPK] Temp file (2eed164577b237a7b764f6adf46b9c249551c2200e440ef0f696c53dfb2d1cb5) content: 
 apiVersion: jobset.x-k8s.io/v1alpha2
 kind: JobSet
 metadata:
@@ -80,13 +80,7 @@ spec:
             - action: FailJob
               onPodConditions: []
               onExitCodes:
-                containerName: jax-tpu-1
-                operator: NotIn
-                values: [42,127,128,129,130,131,132,133,134,135,136,137,138,139,140,141,142,143,144,145,146,147,148,149,150,151,152,153,154,155,156,157,158,159,160,161,162,163,164,165,166,167,168,169,170,171,172,173,174,175,176,177,178,179,180,181,182,183,184,185,186,187,188,189,190,191,192,193,194,195,196,197,198,199,200,201,202,203,204,205,206,207,208,209,210,211,212,213,214,215,216,217,218,219,220,221,222,223,224,225,226,227,228,229,230,231,232,233,234,235,236,237,238,239,240,241,242,243,244,245,246,247,248,249,250,251,252,253,254,255]
-            - action: FailJob
-              onPodConditions: []
-              onExitCodes:
-                containerName: jax-tpu-2
+                containerName: jax-tpu
                 operator: NotIn
                 values: [42,127,128,129,130,131,132,133,134,135,136,137,138,139,140,141,142,143,144,145,146,147,148,149,150,151,152,153,154,155,156,157,158,159,160,161,162,163,164,165,166,167,168,169,170,171,172,173,174,175,176,177,178,179,180,181,182,183,184,185,186,187,188,189,190,191,192,193,194,195,196,197,198,199,200,201,202,203,204,205,206,207,208,209,210,211,212,213,214,215,216,217,218,219,220,221,222,223,224,225,226,227,228,229,230,231,232,233,234,235,236,237,238,239,240,241,242,243,244,245,246,247,248,249,250,251,252,253,254,255]
           template:
@@ -114,42 +108,7 @@ spec:
               terminationGracePeriodSeconds: 30
               containers:
               
-              - name: jax-tpu-1
-                image: gcr.io/golden-project/dry-run-runner:prefix-current
-                
-                env: 
-                securityContext:
-                  privileged: true
-                command:
-                - bash
-                - -c
-                - |
-                  echo XPK Start: $(date);
-                  _sigterm() (kill -SIGTERM $! 2>/dev/null;);
-                  trap _sigterm SIGTERM;
-                  
-                  (bash hello) & PID=$!;
-                  while kill -0 $PID 2>/dev/null;
-                      do sleep 5;
-                  done;
-                  wait $PID;
-                  EXIT_CODE=$?;
-                  
-                  echo XPK End: $(date);
-                  echo EXIT_CODE=$EXIT_CODE;
-                  
-                  
-                  exit $EXIT_CODE
-                resources:
-                  limits:
-                    google.com/tpu: 2
-
-                volumeMounts:
-                - mountPath: /dev/shm
-                  name: dshm-2
-                
-
-              - name: jax-tpu-2
+              - name: jax-tpu
                 image: gcr.io/golden-project/dry-run-runner:prefix-current
                 
                 env: 
@@ -197,7 +156,7 @@ spec:
               
 
 [XPK] Task: `Creating Workload` is implemented by the following command not running since it is a dry run. 
-kubectl apply -f 608e1382aabe2b0335855e5e99876a2e67de954453ebfa4cf12eb82c966f85da
+kubectl apply -f 2eed164577b237a7b764f6adf46b9c249551c2200e440ef0f696c53dfb2d1cb5
 [XPK] Task: `GKE Dashboard List` is implemented by the following command not running since it is a dry run. 
 gcloud monitoring dashboards list --project=golden-project --filter="displayName:'GKE - TPU Monitoring Dashboard'" --format="value(name)" --verbosity=error
 [XPK] Check statistics and outlier mode of GKE metrics here: https://console.cloud.google.com/monitoring/dashboards/builder/0?project=golden-project&f.rlabel.cluster_name.ClusterName=golden-cluster. To view the metric data for your workload, select golden-workload from the JobName filter on the dashboard.

{xpk-1.1.0 → xpk-1.1.1}/src/xpk/commands/workload.py

@@ -481,6 +481,16 @@ def workload_create(args) -> None:
         + lustre_storages
     )
 
+  use_sub_slicing = (
+      workload_scheduling == WorkloadScheduling.SUB_SLICING_AVAILABLE
+  )
+  use_super_slicing = (
+      workload_scheduling == WorkloadScheduling.SUPER_SLICING_AVAILABLE
+  )
+  parallel_containers = workload_system.parallel_containers
+  if args.use_pathways or use_super_slicing:
+    parallel_containers = 1
+
   # Currently failure policy rules are supported for Pathways workloads. b/408465881
   failure_policy_rules = ''
   pod_failure_policy = ''
@@ -497,10 +507,8 @@ def workload_create(args) -> None:
             rules:
           """
     docker_image = get_main_container_docker_image(args, workload_system)
-    for i in range(workload_system.parallel_containers):
-      docker_image_sufix = (
-          f'-{i + 1}' if workload_system.parallel_containers > 1 else ''
-      )
+    for i in range(parallel_containers):
+      docker_image_sufix = f'-{i + 1}' if parallel_containers > 1 else ''
       pod_failure_policy += f"""
             - action: FailJob
               onPodConditions: []
@@ -533,7 +541,7 @@ def workload_create(args) -> None:
   # Create the workload file based on accelerator type or workload type.
   if workload_system.accelerator_type == AcceleratorType.GPU:
     container, debugging_dashboard_id = get_user_workload_container(
-        args, workload_system
+        args, workload_system, parallel_containers=parallel_containers
     )
     gpu_scheduler, return_code = get_gpu_scheduler(
         args, workload_system, autoprovisioning_args
@@ -624,25 +632,21 @@ def workload_create(args) -> None:
         custom_pathways_server=append_custom_pathways_server(args),
         custom_pathways_worker=append_custom_pathways_worker(args),
         colocated_python_sidecar=append_custom_colocated_python_sidecar(args),
-        user_workload=get_user_workload_for_pathways(args, workload_system),
+        user_workload=get_user_workload_for_pathways(
+            args, workload_system, parallel_containers
+        ),
         local_queue_name=LOCAL_QUEUE_NAME,
         autoprovisioning_args=autoprovisioning_args,
         placement_policy_label=placement_policy_label,
     )
   else:
-    use_sub_slicing = (
-        workload_scheduling == WorkloadScheduling.SUB_SLICING_AVAILABLE
-    )
-    use_super_slicing = (
-        workload_scheduling == WorkloadScheduling.SUPER_SLICING_AVAILABLE
-    )
     if use_sub_slicing:
       xpk_print('Workload will be scheduled using the Sub-slicing feature.')
     if use_super_slicing:
       xpk_print('Workload will be scheduled using the Super-slicing feature.')
 
     container, debugging_dashboard_id = get_user_workload_container(
-        args, workload_system
+        args, workload_system, parallel_containers
     )
 
     machine_label = (

{xpk-1.1.0 → xpk-1.1.1}/src/xpk/core/docker_container.py

@@ -30,12 +30,18 @@ from .system_characteristics import (
 )
 
 
-def get_main_and_sidecar_container(args, system, docker_image) -> str:
+def get_main_and_sidecar_container(
+    args,
+    system: SystemCharacteristics,
+    docker_image: str,
+    parallel_containers: int,
+) -> str:
   """Generate yaml for main and sidecar container.
   Args:
     args: user provided arguments for running the command.
     system: system characteristics
     docker_image: docker image
+    parallel_containers: number of containers to run per VM.
 
   Returns:
     str:
@@ -44,7 +50,9 @@ def get_main_and_sidecar_container(args, system, docker_image) -> str:
   resource_type = AcceleratorTypeToAcceleratorCharacteristics[
       system.accelerator_type
   ].resource_type
-  main_container = get_main_container(args, system, docker_image, resource_type)
+  main_container = get_main_container(
+      args, system, docker_image, resource_type, parallel_containers
+  )
   yaml = """- name: stacktrace-explorer
                 image: busybox:1.28
                 args: [/bin/sh, -c, "check_signal() (while [ ! -f /shared-volume/stacktrace_signal ]; do sleep 1; done; pid=$(pidof 'tail'); kill $pid;); check_signal & while [ ! -d /tmp/debugging ]; do sleep 60; done; while [ ! -e /tmp/debugging/* ]; do sleep 60; done; tail -n+1 -f /tmp/debugging/*; exit 0;"]
@@ -59,13 +67,20 @@ def get_main_and_sidecar_container(args, system, docker_image) -> str:
   return yaml.format(main_container=main_container)
 
 
-def get_main_container(args, system, docker_image, resource_type) -> str:
+def get_main_container(
+    args,
+    system: SystemCharacteristics,
+    docker_image: str,
+    resource_type,
+    parallel_containers: int,
+) -> str:
   """Generate yaml for main container including the xpk command.
   Args:
     args: user provided arguments for running the command.
     system: system characteristics
     docker_image: docker image
     resource_type: The label to describe the resource type for TPUs/GPUs/CPUs.
+    parallel_containers: number of containers to run per VM.
 
   Returns:
     str:
@@ -149,14 +164,10 @@ def get_main_container(args, system, docker_image, resource_type) -> str:
                 volumeMounts:
                 {volume_mounts}
 """
-  # pathways job running on 2 parallel containers is not verified yet
-  if args.use_pathways:
-    system.parallel_containers = 1
-
   env = get_env_container(args, system)
   image_pull_policy = add_image_pull_policy_for_pw_or_gpu(args, system)
-  for i in range(system.parallel_containers):
-    docker_name_sufix = f'-{i + 1}' if system.parallel_containers > 1 else ''
+  for i in range(parallel_containers):
+    docker_name_sufix = f'-{i + 1}' if parallel_containers > 1 else ''
     containers.append(
         container_yaml.format(
             args=args,
@@ -177,12 +188,15 @@ def get_main_container(args, system, docker_image, resource_type) -> str:
   return ''.join(containers)
 
 
-def get_user_workload_container(args, system: SystemCharacteristics):
+def get_user_workload_container(
+    args, system: SystemCharacteristics, parallel_containers: int
+):
   """Deploy user workload container
 
   Args:
       args: user provided args.
       system: system characteristics.
+      parallel_containers: number of containers to run per VM.
 
   Returns:
       container: main container
@@ -209,11 +223,15 @@ def get_user_workload_container(args, system: SystemCharacteristics):
         'Sidecar container to display stack traces for TPU workloads will also'
         ' be deployed.'
     )
-    container = get_main_and_sidecar_container(args, system, docker_image)
+    container = get_main_and_sidecar_container(
+        args, system, docker_image, parallel_containers
+    )
     # Get GKE debugging dashboard only when sidecar container is deployed for TPU workloads
     debugging_dashboard_id = get_gke_debugging_dashboard(args)
   else:
-    container = get_main_container(args, system, docker_image, resource_type)
+    container = get_main_container(
+        args, system, docker_image, resource_type, parallel_containers
+    )
   return container, debugging_dashboard_id
 
 

{xpk-1.1.0 → xpk-1.1.1}/src/xpk/core/pathways.py

@@ -245,18 +245,12 @@ def append_custom_colocated_python_sidecar(args) -> str:
 
 
 def get_user_workload_for_pathways(
-    args,
-    system: SystemCharacteristics,
+    args, system: SystemCharacteristics, parallel_containers: int
 ) -> str:
   """
   Create a user workload container for Pathways.
   Don't create one for Pathways headless mode.
 
-  Args:
-    args: user provided args.
-    system: system characteristics.
-
-
   Returns:
     str:
       Pathways server port as a YAML string
@@ -280,7 +274,9 @@ def get_user_workload_for_pathways(
   if args.headless:
     return ''
   else:
-    container, _ = get_user_workload_container(args, system)
+    container, _ = get_user_workload_container(
+        args, system, parallel_containers
+    )
     return user_workload_yaml.format(
         args=args,
         container=container,

{xpk-1.1.0 → xpk-1.1.1}/src/xpk/core/scheduling_test.py

@@ -398,15 +398,23 @@ SUPER_SLICING_CASE = SchedulingTestCase(
             WorkloadScheduling.UNAVAILABLE,
         ),
         (
-            (
-                'Super-slicing should be ignored when a given device is already'
-                ' present in the cluster'
+            'Super-slicing, but one cube',
+            dataclasses.replace(
+                SUPER_SLICING_CASE,
+                workload_system=_get_system_characteristics_or_die('tpu7x-128'),
+                cluster_system=_get_system_characteristics_or_die('tpu7x-128'),
+                resources_config_map={'tpu7x-128': '16'},
             ),
+            WorkloadScheduling.SUPER_SLICING_AVAILABLE,
+        ),
+        (
+            'Super-slicing, but one cube and no super-slicing-topology',
             dataclasses.replace(
                 SUPER_SLICING_CASE,
-                workload_system=_get_system_characteristics_or_die('tpu7x-64'),
-                cluster_system=_get_system_characteristics_or_die('tpu7x-64'),
-                resources_config_map={'tpu7x-64': '16'},
+                workload_system=_get_system_characteristics_or_die('tpu7x-128'),
+                cluster_system=_get_system_characteristics_or_die('tpu7x-128'),
+                resources_config_map={'tpu7x-128': '16'},
+                super_slicing_topology_set=False,
             ),
             WorkloadScheduling.AVAILABLE,
         ),

{xpk-1.1.0 → xpk-1.1.1/src/xpk.egg-info}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: xpk
-Version: 1.1.0
+Version: 1.1.1
 Summary: xpk helps Cloud developers to orchestrate training jobs on accelerators on GKE.
 Author-email: XPK team <xpk-code-reviewers@google.com>
 License: Apache-2.0