Update generate_click_command inject logic to not expose unwanted flags to hyp-jumpstart-endpoint (#213)

* Update generate_click_command inject logic to not expose unwanted flags to hyp-jumpstart-endpoint

* Update unit tests for bug fix, change --label_selector to --label-selector

Author: mollyheamazon

src/sagemaker/hyperpod/cli/inference_utils.py

@@ -40,45 +40,29 @@ def wrapped_func(*args, **kwargs):
             domain = flat.to_domain()
             return func(namespace, version, domain)
 
-        # 2) inject the special JSON‐env flag before everything else
-        wrapped_func = click.option(
-            "--env",
-            callback=_parse_json_flag,
-            type=str,
-            default=None,
-            help=(
-                "JSON object of environment variables, e.g. "
-                '\'{"VAR1":"foo","VAR2":"bar"}\''
-            ),
-            metavar="JSON",
-        )(wrapped_func)
-
-        wrapped_func = click.option(
-            "--dimensions",
-            callback=_parse_json_flag,
-            type=str,
-            default=None,
-            help=("JSON object of dimensions, e.g. " '\'{"VAR1":"foo","VAR2":"bar"}\''),
-            metavar="JSON",
-        )(wrapped_func)
-
-        wrapped_func = click.option(
-            "--resources-limits",
-            callback=_parse_json_flag,
-            help='JSON object of resource limits, e.g. \'{"cpu":"2","memory":"4Gi"}\'',
-            metavar="JSON",
-        )(wrapped_func)
-
-        wrapped_func = click.option(
-            "--resources-requests",
-            callback=_parse_json_flag,
-            help='JSON object of resource requests, e.g. \'{"cpu":"1","memory":"2Gi"}\'',
-            metavar="JSON",
-        )(wrapped_func)
-
-        # 3) auto-inject all schema.json fields
+        # 2) inject JSON flags only if they exist in the schema
         schema = load_schema_for_version(version, schema_pkg)
         props = schema.get("properties", {})
+        
+        json_flags = {
+            "env": ("JSON object of environment variables, e.g. " '\'{"VAR1":"foo","VAR2":"bar"}\''),
+            "dimensions": ("JSON object of dimensions, e.g. " '\'{"VAR1":"foo","VAR2":"bar"}\''),
+            "resources_limits": ('JSON object of resource limits, e.g. \'{"cpu":"2","memory":"4Gi"}\''),
+            "resources_requests": ('JSON object of resource requests, e.g. \'{"cpu":"1","memory":"2Gi"}\''),
+        }
+        
+        for flag_name, help_text in json_flags.items():
+            if flag_name in props:
+                wrapped_func = click.option(
+                    f"--{flag_name.replace('_', '-')}",
+                    callback=_parse_json_flag,
+                    type=str,
+                    default=None,
+                    help=help_text,
+                    metavar="JSON",
+                )(wrapped_func)
+
+        # 3) auto-inject all schema.json fields
         reqs = set(schema.get("required", []))
 
         for name, spec in reversed(list(props.items())):

src/sagemaker/hyperpod/cli/training_utils.py

@@ -107,7 +107,7 @@ def wrapped_func(*args, **kwargs):
             metavar="JSON",
         )(wrapped_func)
         wrapped_func = click.option(
-            "--label_selector",
+            "--label-selector",
             callback=_parse_json_flag,
             help='JSON object of resource limits, e.g. \'{"cpu":"2","memory":"4Gi"}\'',
             metavar="JSON",

test/unit_tests/cli/test_inference_utils.py

@@ -59,7 +59,15 @@ def cmd(namespace, version, domain):
 
     @patch('sagemaker.hyperpod.cli.inference_utils.load_schema_for_version')
     def test_json_flags(self, mock_load_schema):
-        mock_load_schema.return_value = {'properties': {}, 'required': []}
+        mock_load_schema.return_value = {
+            'properties': {
+                'env': {'type': 'object'},
+                'dimensions': {'type': 'object'},
+                'resources_limits': {'type': 'object'},
+                'resources_requests': {'type': 'object'}
+            }, 
+            'required': []
+        }
         # Domain receives flags as attributes env, dimensions, resources_limits, resources_requests
         class DummyFlat:
             def __init__(self, **kwargs): self.__dict__.update(kwargs)

test/unit_tests/cli/test_training_utils.py

@@ -80,7 +80,7 @@ def cmd(version, debug, config):
         # Test valid JSON input
         result = self.runner.invoke(cmd, [
             '--environment', '{"VAR1":"val1"}',
-            '--label_selector', '{"key":"value"}'
+            '--label-selector', '{"key":"value"}'
         ])
         assert result.exit_code == 0
         output = json.loads(result.output)