Remove image arg. Move CodeBuild compute type choices to arg parse

Author: brightsparc

sagemaker_studio_image_build/cli.py

@@ -25,14 +25,6 @@ def validate_args(args, extra_args):
                     f"The value of the -f/file argument [{file_value}] is outside the working directory [{os.getcwd()}]"
                 )
 
-    # Validate arg compute_type
-    if args.compute_type:
-        if not args.compute_type in ['BUILD_GENERAL1_SMALL', 'BUILD_GENERAL1_MEDIUM',
-        'BUILD_GENERAL1_LARGE', 'BUILD_GENERAL1_LARGE', 'BUILD_GENERAL1_LARGE', 'BUILD_GENERAL1_2XLARGE']:
-            raise ValueError(
-                f'Error parsing reference: "{args.repository}" is not a valid repository/tag'
-            )
-
 
 def get_role(args):
     if args.role:
@@ -78,13 +70,11 @@ def main():
         "--repository",
         help="The ECR repository:tag for the image (default: sagemaker-studio-${domain_id}:latest)",
     )
-    build_parser.add_argument(
-        "--image",
-        help="The ECR repository:tag for the image (default: sagemaker-studio-${domain_id}:latest)",
-    )
     build_parser.add_argument(
         "--compute-type",
-        help="The code build compute type (default: BUILD_GENERAL1_SMALL)",
+        help="The CodeBuild compute type (default: BUILD_GENERAL1_SMALL)",
+        choices=["BUILD_GENERAL1_SMALL", "BUILD_GENERAL1_MEDIUM",
+                 "BUILD_GENERAL1_LARGE", "BUILD_GENERAL1_2XLARGE"]
         default="BUILD_GENERAL1_SMALL"
     )
     build_parser.add_argument(

sagemaker_studio_image_build/codebuild.py

@@ -18,7 +18,7 @@ def __init__(self, s3_location, role, repository=None, compute_type=None):
         self.session = boto3.session.Session()
         self.domain_id, self.user_profile_name = self._get_studio_metadata()
         self.repo_name = None
-        self.compute_type = compute_type or 'BUILD_GENERAL1_SMALL'
+        self.compute_type = compute_type or "BUILD_GENERAL1_SMALL"
 
         if repository:
             self.repo_name, self.tag = repository.split(":", maxsplit=1)