Add additional Space parameters for resources including the fractional GPU (#287)

Author: aws-brianxia

src/sagemaker/hyperpod/cli/space_utils.py

@@ -42,22 +42,40 @@ def generate_click_command(
 
     def decorator(func: Callable) -> Callable:
         # build resources from CPU/memory options  
-        def _build_resources(cpu, memory, gpu):
-            if cpu is None and memory is None and gpu is None:
+        def _build_resources(cpu, cpu_limit, memory, memory_limit, gpu, gpu_limit,
+            accelerator_partition_type, accelerator_partition_count):
+            if not any([cpu, cpu_limit, memory, memory_limit, gpu, gpu_limit,
+                accelerator_partition_type, accelerator_partition_count]):
                 return None
+            
+            if (accelerator_partition_type is None) ^ (accelerator_partition_count is None):
+                raise click.UsageError(
+                    "Both accelerator-partition-type and accelerator-partition-count must be specified together"
+                )
 
             # Build requests dictionary
             requests = {}
+            limits = {}
             if cpu is not None:
                 requests["cpu"] = cpu
+            if cpu_limit is not None:
+                limits["cpu"] = cpu_limit
             if memory is not None:
                 requests["memory"] = memory
+            if memory_limit is not None:
+                limits["memory"] = memory_limit
             if gpu is not None:
                 requests["nvidia.com/gpu"] = gpu
+            if gpu_limit is not None:
+                limits["nvidia.com/gpu"] = gpu_limit
+            if accelerator_partition_type is not None and accelerator_partition_count is not None:
+                requests[f"nvidia.com/{accelerator_partition_type}"] = accelerator_partition_count
+                limits[f"nvidia.com/{accelerator_partition_type}"] = accelerator_partition_count
 
             # Return ResourceRequirements structure
             return {
-                "requests": requests
+                "requests": requests,
+                "limits": limits,
             }
 
         def _parse_volume_param(ctx, param, value):
@@ -140,7 +158,16 @@ def wrapped_func(*args, **kwargs):
             if Model is None:
                 raise click.ClickException(f"Unsupported schema version: {version}")
 
-            resources = _build_resources(kwargs.pop("cpu", None), kwargs.pop("memory", None), kwargs.pop("gpu", None))
+            resources = _build_resources(
+                kwargs.pop("cpu", None),
+                kwargs.pop("cpu_limit", None),
+                kwargs.pop("memory", None),
+                kwargs.pop("memory_limit", None),
+                kwargs.pop("gpu", None),
+                kwargs.pop("gpu_limit", None),
+                kwargs.pop("accelerator_partition_type", None),
+                kwargs.pop("accelerator_partition_count", None),
+            )
             if resources is not None:
                 kwargs["resources"] = resources
 
@@ -210,21 +237,56 @@ def wrapped_func(*args, **kwargs):
             "--cpu",
             type=str,
             default=None,
-            help="CPU resource, e.g. '250m'",
+            help="CPU resource request, e.g. '250m'",
+        )(wrapped_func)
+
+        wrapped_func = click.option(
+            "--cpu-limit",
+            type=str,
+            default=None,
+            help="CPU resource limit, e.g. '250m'",
         )(wrapped_func)
 
         wrapped_func = click.option(
             "--memory",
             type=str,
             default=None,
-            help="Memory resource, e.g. '256Mi'",
+            help="Memory resource request, e.g. '256Mi'",
+        )(wrapped_func)
+
+        wrapped_func = click.option(
+            "--memory-limit",
+            type=str,
+            default=None,
+            help="Memory resource limit, e.g. '256Mi'",
         )(wrapped_func)
 
         wrapped_func = click.option(
             "--gpu",
             type=str,
             default=None,
-            help="Gpu resource, e.g. '1'",
+            help="Gpu resource request, e.g. '1'",
+        )(wrapped_func)
+
+        wrapped_func = click.option(
+            "--gpu-limit",
+            type=str,
+            default=None,
+            help="Gpu resource limit, e.g. '1'",
+        )(wrapped_func)
+
+        wrapped_func = click.option(
+            "--accelerator-partition-type",
+            type=str,
+            default=None,
+            help="Fractional GPU parition type",
+        )(wrapped_func)
+
+        wrapped_func = click.option(
+            "--accelerator-partition-count",
+            type=str,
+            default=None,
+            help="Fractional GPU parition count",
         )(wrapped_func)
 
         wrapped_func = click.option(

src/sagemaker/hyperpod/space/hyperpod_space.py

@@ -314,7 +314,7 @@ def list_pods(self) -> List[str]:
         try:
             pods = v1.list_namespaced_pod(
                 namespace=self.config.namespace,
-                label_selector=f"{SPACE_GROUP}/workspaceName={self.config.name}"
+                label_selector=f"{SPACE_GROUP}/workspace-name={self.config.name}"
             )
             return [pod.metadata.name for pod in pods.items]
         except Exception as e:

test/unit_tests/cli/test_space_utils.py

@@ -91,7 +91,7 @@ def cmd(version, domain_config):
 
     @patch('sagemaker.hyperpod.cli.space_utils.load_schema_for_version')
     def test_resources_building(self, mock_load_schema):
-        """Test CPU and memory resource building"""
+        """Test CPU, memory, GPU and fractional GPU resource building"""
         schema = {
             'properties': {
                 'resources': {
@@ -117,26 +117,43 @@ class Config:
         def cmd(version, domain_config):
             click.echo(json.dumps(domain_config.get('resources')))
 
-        # Test with custom CPU and memory
-        result = self.runner.invoke(cmd, ['--cpu', '1000m', '--memory', '1Gi'])
+        # Test with CPU and memory requests and limits
+        result = self.runner.invoke(cmd, ['--cpu', '1000m', '--cpu-limit', '2000m', '--memory', '1Gi', '--memory-limit', '2Gi'])
         assert result.exit_code == 0
         output = json.loads(result.output)
         assert output['requests']['cpu'] == '1000m'
         assert output['requests']['memory'] == '1Gi'
-        assert 'nvidia.com/gpu' not in output['requests']
+        assert output['limits']['cpu'] == '2000m'
+        assert output['limits']['memory'] == '2Gi'
 
-        # Test with only CPU
-        result = self.runner.invoke(cmd, ['--cpu', '750m'])
+        # Test with GPU requests and limits
+        result = self.runner.invoke(cmd, ['--gpu', '1', '--gpu-limit', '2'])
         assert result.exit_code == 0
         output = json.loads(result.output)
-        assert output['requests']['cpu'] == '750m'
-        assert 'memory' not in output['requests']
+        assert output['requests']['nvidia.com/gpu'] == '1'
+        assert output['limits']['nvidia.com/gpu'] == '2'
+
+        # Test with fractional GPU partitioning
+        result = self.runner.invoke(cmd, ['--accelerator-partition-type', 'mig-1g.5gb', '--accelerator-partition-count', '2'])
+        assert result.exit_code == 0
+        output = json.loads(result.output)
+        assert output['requests']['nvidia.com/mig-1g.5gb'] == '2'
+        assert output['limits']['nvidia.com/mig-1g.5gb'] == '2'
 
         # Test with no resources specified
         result = self.runner.invoke(cmd, [])
         assert result.exit_code == 0
         assert result.output.strip() == 'null'
 
+        # Test error when only one accelerator partition parameter is provided
+        result = self.runner.invoke(cmd, ['--accelerator-partition-type', 'mig-1g.5gb'])
+        assert result.exit_code == 2
+        assert 'Both accelerator-partition-type and accelerator-partition-count must be specified together' in result.output
+
+        result = self.runner.invoke(cmd, ['--accelerator-partition-count', '2'])
+        assert result.exit_code == 2
+        assert 'Both accelerator-partition-type and accelerator-partition-count must be specified together' in result.output
+
     @patch('sagemaker.hyperpod.cli.space_utils.load_schema_for_version')
     def test_type_conversion(self, mock_load_schema):
         """Test type conversion for different parameter types"""

test/unit_tests/test_hyperpod_space.py

@@ -504,7 +504,7 @@ def test_list_pods_success(self, mock_verify_config, mock_core_api_class):
         self.assertEqual(result, ["pod1", "pod2"])
         mock_core_api.list_namespaced_pod.assert_called_once_with(
             namespace="test-namespace",
-            label_selector="workspace.jupyter.org/workspaceName=test-space"
+            label_selector="workspace.jupyter.org/workspace-name=test-space"
         )
 
     @patch('sagemaker.hyperpod.space.hyperpod_space.client.CoreV1Api')