y1dZddlmZddlmZddlmZddlmZddlmZddlm Z ddl m Z dd l m Z dd l mZdd l mZdd l mZGd de j"ZGdde j"ZddZddZGddeZy)z'Utilities for dealing with ML jobs API.)absolute_import)division)unicode_literals)encoding) list_pager)apis) arg_utils) exceptions) properties) resources)yamlceZdZdZy)NoFieldsSpecifiedErrorzEError indicating that no updates were requested in a Patch operation.N__name__ __module__ __qualname____doc__,lib/googlecloudsdk/api_lib/ml_engine/jobs.pyrrsMrrceZdZdZy)NoPackagesSpecifiedErrorz>Error that no packages were specified for non-custom training.Nrrrrrr#sFrrc.tjd|S)Nml)rGetMessagesModule)versions rrr's  g ..rc2tjd||S)Nr)no_http)rGetClientInstance)rrs rr r +s  gw ??rceZdZdZ ddZdZedZdZedZ edZ ed Z dd Z d Z d Zd ZdZ ddZdZdZ ddZy) JobsClientz3Client for jobs service in the Cloud ML Engine API.Ncx|xs t||_|xs|jj|_||_yN)r clientMESSAGES_MODULEmessages_short_message_prefix)selfr%r'short_message_prefixclient_versions r__init__zJobsClient.__init__2s2=-n=DK; ; ;DM!5Drcft|jdj|j|dS)Nz{prefix}{name})prefixname)getattrr'formatr()r)short_message_names rGetShortMessagezJobsClient.GetShortMessage8s9 4==#**$2L2L0B+DEI KKrcB|jjjSr$)r'GoogleCloudMlV1JobStateValueValuesEnumr)s r state_enumzJobsClient.state_enum=s == + + @ @@rc|jj|j}tj|j j |ddS)N)parentjobspageSize)fieldbatch_size_attribute)r'MlProjectsJobsListRequest RelativeNamer YieldFromListr% projects_jobs)r) project_refreqs rListzJobsClient.ListAsN -- 1 1'') 2 +C  # # !!3f' ))rc.|jjSr$)r'r5r7s r job_classzJobsClient.job_classHs == + ++rc.|jjSr$)r'GoogleCloudMlV1TrainingInputr7s rtraining_input_classzJobsClient.training_input_classLs == 5 55rc.|jjSr$)r'GoogleCloudMlV1PredictionInputr7s rprediction_input_classz!JobsClient.prediction_input_classPs == 7 77rc<|jj||S)N)r:googleCloudMlV1Job)r'MlProjectsJobsCreateRequest)r)r:jobs r_MakeCreateRequestzJobsClient._MakeCreateRequestTs$ == 4 4 5  rc|jjj|j|j |S)N)r:rQ)r%rBCreaterRr@)r)rCrQs rrTzJobsClient.CreateYsA ;; $ $ + + ++-   rc|jj|j}|jjj |S)zCancels given job.r/)r'MlProjectsJobsCancelRequestr@r%rBCancelr)job_refrDs rrXzJobsClient.Cancel_s= -- 3 39M9M9O 3 PC ;; $ $ + +C 00rc|jj|j}|jjj |S)NrV)r'MlProjectsJobsGetRequestr@r%rBGetrYs rr]zJobsClient.Getds= -- 0 0g6J6J6L 0 MC ;; $ $ ( ( --rc\|j}g}|jr"|j|_|jd|s t d|j j |j|dj|}|jjj|S)z Update a job.labelszNo updates requested.,)r/rO updateMask) rG needs_updater_appendrr'MlProjectsJobsPatchRequestr@joinr%rBPatch)r)rZ labels_updaterQ update_maskrDs rrfzJobsClient.Patchhs .. CK!! ''cj"  "#: ;; -- 2 2  ! ! #88K( 3 C ;; $ $ * *3 //rc p|j}|r7tj|}|r tj||j}|r||_| | |_|js|j|_||||||| | | | |d }|jD]\}}| t|j||!|rtj|d||r0|jjD]\}}| |jdr|jds|dr)tj|dj||d|ddr,tj|d j||dd|dd r,tj|d j||dd |d k(s|d stj|dj||d t|j|||j!|s|j#|s t%d|S)a~Builds a Cloud ML Engine Job from a config file and/or flag values. Args: path: path to a yaml configuration file module_name: value to set for moduleName field (overrides yaml file) job_name: value to set for jobName field (overrides yaml file) trainer_uri: List of values to set for trainerUri field (overrides yaml file) region: compute region in which to run the job (overrides yaml file) job_dir: Cloud Storage working directory for the job (overrides yaml file) scale_tier: ScaleTierValueValuesEnum the scale tier for the job (overrides yaml file) user_args: [str]. A list of arguments to pass through to the job. (overrides yaml file) runtime_version: the runtime version in which to run the job (overrides yaml file) python_version: the Python version in which to run the job (overrides yaml file) network: user network to which the job should be peered with (overrides yaml file) service_account: A service account (email address string) to use for the job. labels: Job.LabelsValue, the Cloud labels for the job kms_key: A customer-managed encryption key to use for the job. custom_train_server_config: jobs_util.CustomTrainingInputServerConfig, configuration object for custom server parameters. enable_web_access: whether to enable the interactive shell for the job. Raises: NoPackagesSpecifiedError: if a non-custom job was specified without any trainer_uris. Returns: A constructed Job object. ) pythonModuleargs packageUrisregionjobDir scaleTierruntimeVersion pythonVersionnetworkserviceAccountenableWebAccessz)trainingInput.encryptionConfig.kmsKeyNameConfigTfConfigimageUriztrainingInput.{}.imageUriacceleratorConfigtypez'trainingInput.{}.acceleratorConfig.typecountz(trainingInput.{}.acceleratorConfig.count workerConfig tpuTfVersionztrainingInput.{}.tpuTfVersionz#Non-custom jobs must have packages.)rGr load_pathr DictToMessagejobIdr_ trainingInputrJitemssetattrr SetFieldInMessage GetFieldMapendswithr1HasPackageURIsIsCustomContainerTrainingr)r)path module_namejob_name trainer_urirmjob_dir scale_tier user_argsruntime_versionpython_versionrrservice_accountr_kms_keycustom_train_server_configenable_web_accessrQdataadditional_fields field_namevalues rBuildTrainingJobzJobsClient.BuildTrainingJobxsPf .. C ^^D !d $$T4>>:ci cj   335c#")'), /446 E  !!:u57!!#"M")+"9EEGMMO *e  !!(+%%j1Z ))-44Z@ #%()&1));BB:N+,V46()'2))<CCJO+,W57^+n0E))188D') C%%z5 93 P6   s #D,J,J3,O $%J KK Jrc@t|jjSr$)boolrrlr)rQs rrzJobsClient.HasPackageURIss !!-- ..rct|jjxr |jjjSr$)rr masterConfigrwrs rrz$JobsClient.IsCustomContainerTrainings7 !!..8!!..77 99rc <tjjjj }|r|j d}|||}nd}|j |||| | | |}|j||_|r||_ nv|r;tjj|d||d}|j|_n9tjj|dd |i}|j|_| r| |_|j#||| S) avBuilds a Cloud ML Engine Job for batch prediction from flag values. Args: job_name: value to set for jobName field model_dir: str, Google Cloud Storage location of the model files model_name: str, value to set for modelName field version_name: str, value to set for versionName field input_paths: list of input files data_format: format of the input files output_path: single value for the output location region: compute region in which to run the job runtime_version: the runtime version in which to run the job max_worker_count: int, the maximum number of workers to use batch_size: int, the number of records per batch sent to Tensorflow signature_name: str, name of input/output signature in the TF meta graph labels: Job.LabelsValue, the Cloud labels for the job accelerator_count: int, The number of accelerators to attach to the machines accelerator_type: AcceleratorsValueListEntryValuesEnum, The type of accelerator to add to machine. Returns: A constructed Job object. AcceleratorConfig)rzryN) inputPaths outputPathrmrpmaxWorkerCount batchSize acceleratorzml.projects.models.versions)modelsId projectsId) collectionparamszml.projects.modelsr)rpredictionInputr_)r VALUEScoreproject GetOrFailr3rMDataFormatValueValuesEnum dataFormaturir REGISTRYParser@ versionName modelName signatureNamerG)r)r model_dir model_name version_name input_paths data_format output_pathrmrmax_worker_count batch_sizesignature_namer_accelerator_countaccelerator_type project_idaccelerator_config_msgaccelerator_configprediction_input version_ref model_refs rBuildBatchPredictionJobz"JobsClient.BuildBatchPredictionJobsMN""''//99;J#334GH18I7GI 22&'&3#3"L"L#& &&,, #@( C-Ek&1%=%=%?"$$** !5 ++-i$-#9#9#; '5$ >>(  r)NNGoogleCloudMlV1v1)NN)NNNNNNNNNNNNNNNN)NNNNNNNNNNNNNNN)rrrrr,r3propertyr8rErGrJrMrRrTrXr]rfrrrrrrrr"r"/s;+/FJ6 K  A A) , , 6 6 8 8  1 .0"!#' $#'"#"&!%'+&*#'+"#26)-!xt/9 (,(,)-+/*.*.*.%).2/3)--1%)04/3Prr"N)r)rF)r __future__rrrapitools.base.pyrrgooglecloudsdk.api_lib.utilr$googlecloudsdk.command_lib.util.apisr googlecloudsdk.corer r r r Errorrrrr objectr"rrrrsl.&'%',:**)$NZ--NGz//G/@ZZr