
                            d Z ddlmZ ddlmZ ddlmZ ddlmZ ddlmZ ddl	m
Z
 ddlmZ dd	lmZ dd
lmZ ddlmZ ddlmZ ddlmZ ddlmZ ddlmZ ddlmZ d Zd Z ej2                  ej4                  j6                        ej8                   G d dej:                                      Z ej2                  ej4                  j>                  ej4                  j@                         G d de             Z!y)z)Vertex AI endpoints deploy-model command.    )absolute_import)division)unicode_literals)encoding)
operations)client)base)	constants)endpoint_util)endpoints_util)flags)operations_util)region_util)
validation)logc                    t        j                  d      j                  |        t        j                         j                  |        t        j                  | |       t        j
                  | |       t        j                         j                  |        t        j                         j                  |        t        j                  d      j                  |        t        j                         j                  |        t        j                  ||        t        j                  | dt        j                         |t        j                  k7  rt        j                   | d       yy)z*Prepares for the arguments of the command.zdeployed modelzdeployed-modelzto deploy a model to)prompt_funczto co-host a model onN)r   GetDisplayNameArgAddToParserGetTrafficSplitArgAddPredictionResourcesArgsAddScaleToZeroArgsGetEnableAccessLoggingArgGetServiceAccountArgGetUserSpecifiedIdArgGetAutoscalingMetricSpecsArgAddModelIdArgAddEndpointResourceArgr   PromptForOpRegionr
   
GA_VERSIONAddSharedResourcesArg)parserversions     (lib/surface/ai/endpoints/deploy_model.py_AddArgsr%   #   s   *+77?((0""67367+!!#//7**62./;;FC$$&226:gv&//1 	$$$	! %    c                    t        j                  | j                         t        j                  | j                         t        j
                  | j                  | j                         | j                  j                  j                         }|j                         d   | _        t        j                  || j                        5  t        j                   |      }t#        j$                         }|t&        j(                  k(  r|j+                  || j,                  | j                  | j                  | j.                  | j0                  | j                  | j2                  | j                  | j4                  | j                  | j6                  | j8                  | j:                  | j<                  | j>                  | j@                        }n| j                  jB                  j                         }t        jD                  || j.                  | j0                  | j                  | j2                  | j                  | j                         t        jF                  | j                  | jH                  | j2                  | jJ                  | jL                         |jO                  || j,                  | j                  | j                  | j.                  | jP                  | jR                  | j0                  | j                  | j2                  | j                  | j4                  | j                  | j6                  | j8                  | jT                  | j<                  | j>                  | j@                  || jJ                  | jL                  | jH                  | jV                        }tY        jZ                  ||t]        j^                  |j`                              }|ftc        jd                  |      }d|v rMd	|d   v rFtf        jh                  jk                  d
jm                  |j                         d   |d   d	                |cddd       S # 1 sw Y   yxY w)z1Deploy a model to an existing Vertex AI endpoint.locationsId)region)r#   )machine_typeaccelerator_dictmin_replica_countmax_replica_countrequired_replica_countreservation_affinityautoscaling_metric_specsspotenable_access_loggingdisable_container_loggingservice_accounttraffic_splitdeployed_model_id)shared_resources_refr*   r+   r,   r-   r.   r0   )r*   tpu_topologymultihost_gpu_node_countr+   r,   r-   r.   r/   r0   r1   r2   enable_container_loggingr4   r5   r6   r7   min_scaleup_periodidle_scaledown_periodinitial_replica_countgpu_partition_sizeNdeployedModelidzBDeployed a model to the endpoint {}. Id of the deployed model: {}.endpointsId)7r   ValidateDisplayNamedisplay_nameValidateAutoscalingMetricSpecsr0   ValidateRequiredReplicaCountr.   r,   CONCEPTSendpointParseAsDictr)   r   AiplatformEndpointOverridesr   EndpointsClientr   OperationsClientr
   r    DeployModelmodelr*   acceleratorr-   r/   r1   r2   r3   r4   r5   r6   shared_resourcesValidateSharedResourceArgsValidateScaleToZeroArgsr=   r;   r<   DeployModelBetar8   r9   r:   r>   r   WaitForOpMayber   ParseOperationnamer   MessageToPyValuer   statusPrintformat)	argsr#   endpoint_refendpoints_clientoperation_clientopr7   response_msgresponses	            r$   _Runrb   8   s|     !2!23++D,I,IJ))$*E*E*.*@*@B''--/,##%m4$+00M--g>!224)&&&''

**
++


((++ 22 22!%!<!<#88#'#@#@yy $ : :$($B$B..** 22# ( 4b& "]];;AAC++3((++ 22 22!%!<!<#'#@#@B ((

 
 

$
$d&<&<

!
!4#=#=
 ++

**
++


((((#'#@#@++ 22 22!%!<!<#88#'#@#@yy $ : :#'#@#@..** 223!44 $ : : $ : :!441 , 6b2 #11"n;;BGGDFL**<8h	H	$/1J)J

 :;A6*113MB&7=<?	@ O NMMs   M+P;;Qc                   &    e Zd ZdZed        Zd Zy)DeployModelGa   Deploy a model to an existing Vertex AI endpoint.

  ## EXAMPLES

  To deploy a model ``456'' to an endpoint ``123'' under project ``example'' in
  region ``us-central1'', run:

    $ {command} 123 --project=example --region=us-central1 --model=456
    --display-name=my_deployed_model
  c                 ~    t        | t        j                         t        j                         j                  |        y N)r%   r
   r    r   GetDisableContainerLoggingArgr   r"   s    r$   ArgszDeployModelGa.Args   s)    VY))*	'')55f=r&   c                 8    t        |t        j                         y rg   )rb   r
   r    selfr[   s     r$   RunzDeployModelGa.Run   s    y##$r&   N__name__
__module____qualname____doc__staticmethodrj   rn    r&   r$   rd   rd      s     	 > >%r&   rd   c                   &    e Zd ZdZed        Zd Zy)rS   re   c                 ~    t        | t        j                         t        j                         j                  |        y rg   )r%   r
   BETA_VERSIONr   GetEnableContainerLoggingArgr   ri   s    r$   rj   zDeployModelBeta.Args   s)    VY++,	&&(44V<r&   c                 8    t        |t        j                         y rg   )rb   r
   rx   rl   s     r$   rn   zDeployModelBeta.Run   s    y%%&r&   Nro   ru   r&   r$   rS   rS      s     	 = ='r&   rS   N)"rs   
__future__r   r   r   apitools.base.pyr   googlecloudsdk.api_lib.air   #googlecloudsdk.api_lib.ai.endpointsr   googlecloudsdk.callioper	   googlecloudsdk.command_lib.air
   r   r   r   r   r   r   googlecloudsdk.corer   r%   rb   ReleaseTracksReleaseTrackGAUniverseCompatibleCommandrd   BETAALPHArS   ru   r&   r$   <module>r      s    0 &  ' % 0 6 ( 3 7 8 / 9 5 4 #!*Od D%%(()%DLL %  *%* D%%**D,=,=,C,CD'm ' E'r&   