cli.py
1 import json 2 import os 3 import tempfile 4 5 import click 6 7 import mlflow.models.docker_utils 8 import mlflow.sagemaker 9 from mlflow.sagemaker import DEFAULT_IMAGE_NAME as IMAGE 10 from mlflow.utils import cli_args 11 from mlflow.utils import env_manager as em 12 13 14 @click.group("sagemaker") 15 def commands(): 16 """ 17 Serve models on SageMaker. 18 19 To serve a model associated with a run on a tracking server, set the MLFLOW_TRACKING_URI 20 environment variable to the URL of the desired server. 21 """ 22 23 24 @commands.command("deploy-transform-job") 25 @click.option("--job-name", "-n", help="Transform job name", required=True) 26 @cli_args.MODEL_URI 27 @click.option("--input-data-type", help="Input data type for the transform job", required=True) 28 @click.option( 29 "--input-uri", "-u", help="S3 key name prefix or manifest of the input data", required=True 30 ) 31 @click.option( 32 "--content-type", 33 help="The multipurpose internet mail extension (MIME) type of the data", 34 required=True, 35 ) 36 @click.option( 37 "--output-path", 38 "-o", 39 help="The S3 path to store the output results of the Sagemaker transform job", 40 required=True, 41 ) 42 @click.option( 43 "--compression-type", default="None", help="The compression type of the transform data" 44 ) 45 @click.option( 46 "--split-type", 47 "-s", 48 default="Line", 49 help="The method to split the transform job's data files into smaller batches", 50 ) 51 @click.option( 52 "--accept", 53 "-a", 54 default="text/csv", 55 help="The multipurpose internet mail extension (MIME) type of the output data", 56 ) 57 @click.option( 58 "--assemble-with", 59 default="Line", 60 help="The method to assemble the results of the transform job as a single S3 object", 61 ) 62 @click.option( 63 "--input-filter", 64 default="$", 65 help="A JSONPath expression used to select a portion of the input data for the transform job", 66 ) 67 @click.option( 68 "--output-filter", 69 default="$", 70 help="A JSONPath expression used to select a portion of the output data from the transform job", 71 ) 72 @click.option( 73 "--join-resource", 74 "-j", 75 default="None", 76 help="The source of the data to join with the transformed data", 77 ) 78 @click.option("--execution-role-arn", "-e", default=None, help="SageMaker execution role") 79 @click.option("--bucket", "-b", default=None, help="S3 bucket to store model artifacts") 80 @click.option("--image-url", "-i", default=None, help="ECR URL for the Docker image") 81 @click.option( 82 "--region-name", 83 default="us-west-2", 84 help="Name of the AWS region in which to deploy the transform job", 85 ) 86 @click.option( 87 "--instance-type", 88 "-t", 89 default=mlflow.sagemaker.DEFAULT_SAGEMAKER_INSTANCE_TYPE, 90 help="The type of SageMaker ML instance on which to perform the batch transform job." 91 " For a list of supported instance types, see" 92 " https://aws.amazon.com/sagemaker/pricing/instance-types/.", 93 ) 94 @click.option( 95 "--instance-count", 96 "-c", 97 default=mlflow.sagemaker.DEFAULT_SAGEMAKER_INSTANCE_COUNT, 98 help="The number of SageMaker ML instances on which to perform the batch transform job", 99 ) 100 @click.option( 101 "--vpc-config", 102 "-v", 103 help="Path to a file containing a JSON-formatted VPC configuration. This" 104 " configuration will be used when creating the new SageMaker model associated" 105 " with this application. For more information, see" 106 " https://docs.aws.amazon.com/sagemaker/latest/dg/API_VpcConfig.html", 107 ) 108 @click.option( 109 "--flavor", 110 "-f", 111 default=None, 112 help=( 113 "The name of the flavor to use for deployment. Must be one of the following: " 114 f"{mlflow.sagemaker.SUPPORTED_DEPLOYMENT_FLAVORS}. If unspecified, a flavor will be " 115 "automatically selected from the model's available flavors." 116 ), 117 ) 118 @click.option( 119 "--archive", 120 is_flag=True, 121 help=( 122 "If specified, any SageMaker resources that become inactive after the finished" 123 " batch transform job are preserved. These resources may include the associated" 124 " SageMaker models and model artifacts. Otherwise, if `--archive` is unspecified," 125 " these resources are deleted. `--archive` must be specified when deploying" 126 " asynchronously with `--async`." 127 ), 128 ) 129 @click.option( 130 "--async", 131 "asynchronous", 132 is_flag=True, 133 help=( 134 "If specified, this command will return immediately after starting the" 135 " deployment process. It will not wait for the deployment process to complete." 136 " The caller is responsible for monitoring the deployment process via native" 137 " SageMaker APIs or the AWS console." 138 ), 139 ) 140 @click.option( 141 "--timeout", 142 default=1200, 143 help=( 144 "If the command is executed synchronously, the deployment process will return" 145 " after the specified number of seconds if no definitive result (success or" 146 " failure) is achieved. Once the function returns, the caller is responsible" 147 " for monitoring the health and status of the pending deployment via" 148 " native SageMaker APIs or the AWS console. If the command is executed" 149 " asynchronously using the `--async` flag, this value is ignored." 150 ), 151 ) 152 def deploy_transform_job( 153 job_name, 154 model_uri, 155 input_data_type, 156 input_uri, 157 content_type, 158 output_path, 159 compression_type, 160 split_type, 161 accept, 162 assemble_with, 163 input_filter, 164 output_filter, 165 join_resource, 166 execution_role_arn, 167 bucket, 168 image_url, 169 region_name, 170 instance_type, 171 instance_count, 172 vpc_config, 173 flavor, 174 archive, 175 asynchronous, 176 timeout, 177 ): 178 """ 179 Deploy model on Sagemaker as a batch transform job. Current active AWS account needs to have 180 correct permissions setup. 181 182 By default, unless the ``--async`` flag is specified, this command will block until 183 either the batch transform job completes (definitively succeeds or fails) or the specified 184 timeout elapses. 185 """ 186 if vpc_config is not None: 187 with open(vpc_config) as f: 188 vpc_config = json.load(f) 189 190 mlflow.sagemaker.deploy_transform_job( 191 job_name=job_name, 192 model_uri=model_uri, 193 s3_input_data_type=input_data_type, 194 s3_input_uri=input_uri, 195 content_type=content_type, 196 s3_output_path=output_path, 197 compression_type=compression_type, 198 split_type=split_type, 199 accept=accept, 200 assemble_with=assemble_with, 201 input_filter=input_filter, 202 output_filter=output_filter, 203 join_resource=join_resource, 204 execution_role_arn=execution_role_arn, 205 bucket=bucket, 206 image_url=image_url, 207 region_name=region_name, 208 instance_type=instance_type, 209 instance_count=instance_count, 210 vpc_config=vpc_config, 211 flavor=flavor, 212 archive=archive, 213 synchronous=(not asynchronous), 214 timeout_seconds=timeout, 215 ) 216 217 218 @commands.command("terminate-transform-job") 219 @click.option("--job-name", "-n", help="Transform job name", required=True) 220 @click.option( 221 "--region-name", 222 "-r", 223 default="us-west-2", 224 help="Name of the AWS region in which the transform job is deployed", 225 ) 226 @click.option( 227 "--archive", 228 is_flag=True, 229 help=( 230 "If specified, resources associated with the application are preserved." 231 " These resources may include unused SageMaker models and model artifacts." 232 " Otherwise, if `--archive` is unspecified, these resources are deleted." 233 " `--archive` must be specified when deleting asynchronously with `--async`." 234 ), 235 ) 236 @click.option( 237 "--async", 238 "asynchronous", 239 is_flag=True, 240 help=( 241 "If specified, this command will return immediately after starting the" 242 " termination process. It will not wait for the termination process to complete." 243 " The caller is responsible for monitoring the termination process via native" 244 " SageMaker APIs or the AWS console." 245 ), 246 ) 247 @click.option( 248 "--timeout", 249 default=1200, 250 help=( 251 "If the command is executed synchronously, the termination process will return" 252 " after the specified number of seconds if no definitive result (success or" 253 " failure) is achieved. Once the function returns, the caller is responsible" 254 " for monitoring the health and status of the pending termination via" 255 " native SageMaker APIs or the AWS console. If the command is executed" 256 " asynchronously using the `--async` flag, this value is ignored." 257 ), 258 ) 259 def terminate_transform_job(job_name, region_name, archive, asynchronous, timeout): 260 """ 261 Terminate the specified Sagemaker batch transform job. Unless ``--archive`` is specified, 262 all SageMaker resources associated with the batch transform job are deleted as well. 263 264 By default, unless the ``--async`` flag is specified, this command will block until 265 either the termination process completes (definitively succeeds or fails) or the specified 266 timeout elapses. 267 """ 268 mlflow.sagemaker.terminate_transform_job( 269 job_name=job_name, 270 region_name=region_name, 271 archive=archive, 272 synchronous=(not asynchronous), 273 timeout_seconds=timeout, 274 ) 275 276 277 @commands.command("push-model") 278 @click.option("--model-name", "-n", help="Sagemaker model name", required=True) 279 @cli_args.MODEL_URI 280 @click.option("--execution-role-arn", "-e", default=None, help="SageMaker execution role") 281 @click.option("--bucket", "-b", default=None, help="S3 bucket to store model artifacts") 282 @click.option("--image-url", "-i", default=None, help="ECR URL for the Docker image") 283 @click.option( 284 "--region-name", 285 default="us-west-2", 286 help="Name of the AWS region in which to push the Sagemaker model", 287 ) 288 @click.option( 289 "--vpc-config", 290 "-v", 291 help="Path to a file containing a JSON-formatted VPC configuration. This" 292 " configuration will be used when creating the new SageMaker model." 293 " For more information, see" 294 " https://docs.aws.amazon.com/sagemaker/latest/dg/API_VpcConfig.html", 295 ) 296 @click.option( 297 "--flavor", 298 "-f", 299 default=None, 300 help=( 301 "The name of the flavor to use for deployment. Must be one of the following:" 302 " {supported_flavors}. If unspecified, a flavor will be automatically selected" 303 " from the model's available flavors.".format( 304 supported_flavors=mlflow.sagemaker.SUPPORTED_DEPLOYMENT_FLAVORS 305 ) 306 ), 307 ) 308 def push_model_to_sagemaker( 309 model_name, 310 model_uri, 311 execution_role_arn, 312 bucket, 313 image_url, 314 region_name, 315 vpc_config, 316 flavor, 317 ): 318 """ 319 Push an MLflow model to Sagemaker model registry. Current active AWS account needs to have 320 correct permissions setup. 321 """ 322 if vpc_config is not None: 323 with open(vpc_config) as f: 324 vpc_config = json.load(f) 325 326 mlflow.sagemaker.push_model_to_sagemaker( 327 model_name=model_name, 328 model_uri=model_uri, 329 execution_role_arn=execution_role_arn, 330 bucket=bucket, 331 image_url=image_url, 332 region_name=region_name, 333 vpc_config=vpc_config, 334 flavor=flavor, 335 ) 336 337 338 @commands.command("build-and-push-container") 339 @click.option("--build/--no-build", default=True, help="Build the container if set.") 340 @click.option("--push/--no-push", default=True, help="Push the container to AWS ECR if set.") 341 @click.option("--container", "-c", default=IMAGE, help="image name") 342 @cli_args.INSTALL_JAVA 343 @cli_args.ENV_MANAGER 344 @cli_args.MLFLOW_HOME 345 def build_and_push_container(build, push, container, install_java, env_manager, mlflow_home): 346 """ 347 Build new MLflow Sagemaker image, assign it a name, and push to ECR. 348 349 This function builds an MLflow Docker image. 350 The image is built locally and it requires Docker to run. 351 The image is pushed to ECR under current active AWS account and to current active AWS region. 352 """ 353 from mlflow.models import docker_utils 354 355 env_manager = env_manager or em.VIRTUALENV 356 if not (build or push): 357 click.echo("skipping both build and push, have nothing to do!") 358 if build: 359 sagemaker_image_entrypoint = ( 360 "import sys; from mlflow.models import container as C; " 361 f"C._init(sys.argv[1], '{env_manager}')" 362 ) 363 364 setup_container = ( 365 "# Install minimal serving dependencies\n" 366 'RUN python -c "from mlflow.models.container import _install_pyfunc_deps;' 367 '_install_pyfunc_deps(None, False)"' 368 ) 369 370 with tempfile.TemporaryDirectory() as tmp: 371 docker_utils.generate_dockerfile( 372 base_image=mlflow.models.docker_utils.UBUNTU_BASE_IMAGE, 373 output_dir=tmp, 374 entrypoint=sagemaker_image_entrypoint, 375 env_manager=env_manager, 376 mlflow_home=os.path.abspath(mlflow_home) if mlflow_home else None, 377 model_install_steps=setup_container, 378 # Create a conda env or virtualenv at runtime after the model is loaded 379 disable_env_creation_at_runtime=False, 380 install_java=install_java, 381 ) 382 383 docker_utils.build_image_from_context(tmp, image_name=container) 384 if push: 385 mlflow.sagemaker.push_image_to_ecr(container)