U d^@sddlmZmZddlmZmZddlmZddlmZddl m Z ddZ Gdd d ej Z Gd d d ej ZGd d d eZddZddZGdddejZGdddeZddZGdddeZGdddejZGdddeZdS))corecontext)Fieldfrom_blob_list) defaultdict)copy) viewitemscCs.|dkr |S|dkr|St|}|||SN)rupdate)abcr6/tmp/pip-unpacked-wheel-ua33x9lu/caffe2/python/task.py_merge_node_kwargs s rc@s8eZdZdZddZddZddZdd Zd d Zd S) Clusterz Context that keeps track of all the node names used. Users shouldn't have to use them directly, since a Cluster is automatically generated at the first usage of 'Node'. cCsg|_i|_dSr )_nodes _node_kwargsselfrrr__init__szCluster.__init__cCsFt||jkr|jt|t||jt||jt|<dSr )strrappendrkwargsrget)rnoderrradd_node"s zCluster.add_nodecCs|jS)zQ Returns the list of unique node names used within this context. )rrrrrnodes)sz Cluster.nodescCs|jSr )rrrrr node_kwargs/szCluster.node_kwargscCsd||S)Nz!Cluster(nodes={}, node_kwargs={}))formatrrrrrr__repr__2szCluster.__repr__N) __name__ __module__ __qualname____doc__rrrrr rrrrrs rc@s2eZdZdZd ddZddZddZd d Zd S) Nodea A Node context is used to indicate that all Tasks instantiated within will run on the given node name. (Only the name of the node actually counts.) Example: with TaskGroup() as tg: with Node('node1'): s1 = execution_step(...) Task(step=s1) with Node('node2'): s2 = execution_step(...) with Node('node1'): s3 = execution_step(...) In this example, all three execution steps will run in parallel. Moreover, s1 and s3 will run on the same node, and can see each others blobs. Additionally, a Node can be passed implementation-specific kwargs, in order to specify properties of the node. localcKs"t||_||_t|dSr )r_name_kwargsrcurrentr)rrrrrrrNs z Node.__init__cCs|jSr )r'rrrr__str__Ssz Node.__str__cCsd|j|jS)NzNode(name={}, kwargs={}))rr'r(rrrrr Vsz Node.__repr__cCs|jSr )r(rrrrrYsz Node.kwargsN)r&)r!r"r#r$rr*r rrrrrr%7s  r%c@seZdZdZdZdZdS) WorkspaceTypez Determines whether tasks of a TaskGroup will run directly at the global workspace, which is kept alive across runs, or whether a new child workspace will be created for the run and destroyed afterwards. privateglobalN)r!r"r#r$PRIVATEGLOBALrrrrr+]sr+c Cst|d}t|d}g}g}g}|D]6}t|drJ|||7}q,t|dr,|||7}q,|D]} t| dr| jrqht| dr| j|krqht| dr| |} t| t t fr|| 7}n6t| tjtj fr| | n| dk rt dt| d | _t| d rh| |} t| t t fr,|| 7}n:t| tjtj frL| | n| dk rft dt| d | _qht|jd kr|d |t|jd kr|d |||fS) Nz/initz/exitget_all_attributesget_attributes _setup_used _setup_targetsetupzUnsupported type for setup: %sTexitr)rNethasattrr0r1r2r3r4 isinstancelisttuple ExecutionStepr TypeErrortyper5lenProtoopinsert) keyZ steps_or_netstargetinit_netexit_net init_nets exit_netsobjsZ step_or_netobjnetsrrrget_setup_netsgsL              rKcCsb|s |s |Sg}|r*|td||||t|dkrV|td||t||S)N%s:initr%s:exit)rrexecution_stepr>)steprFrGnamestepsrrradd_setup_stepss  rRc@seZdZdZdZd ddZddZdd Zd d Zd d Z ddZ ddZ d!ddZ d"ddZ d#ddZd$ddZddZddZdS)% TaskGroupa Context that gathers tasks which will run concurrently, potentially on multiple nodes. All tasks in the same node will share the same workspace and thus can share blobs, while tasks running in different nodes won't be able to directly share data. All tasks of the task group will start concurrently, and the task group will finish execution when the last task of the group finishes. Example: # suppose that s1 ... s5 are execution steps or nets. with TaskGroup() as tg: # these tasks go to default node 'local' Task(step=s1) Task(step=s2) with Node('n2'): Task(step=s3) with Node('n1'): Task(step=s4) with Node('n2'): Task(step=s5) # this will run all steps in parallel. # s1 and s2 will run at default node 'local' # s3 and s5 will run at node 'n2' # s4 will run at node 'n1' session.run(tg) Z local_setupNcCs@d|_g|_d|_d|_g|_i|_g|_||_d|_g|_ dSNF) Z _plan_cache_tasks _already_usedZ _prev_active _tasks_to_add _report_nets _report_steps_workspace_type_tasks_by_node _remote_nets)rworkspace_typerrrrszTaskGroup.__init__cCs|j|dSr )r\r)rnetrrradd_remote_netszTaskGroup.add_remote_netcCs|jSr )r\rrrr remote_netsszTaskGroup.remote_netscCst|jrtd|jdks2|jdks2|j|jks2t|jdkrJ|jpFtj|_|jdkr\|j|_||j|dS)Nz-Cannot add Task to an already used TaskGroup.)rVAssertionErrorrZr+r. _notify_usedrUrrtaskrrradds     z TaskGroup.addcCs(|jD]}||qg|_d|_|jSNT)rWrerVrUrcrrrtaskss   zTaskGroup.taskscCst|jt|jSr )r>rWrUrrrrnum_registered_tasksszTaskGroup.num_registered_taskscCs0g}|j|jD]}|j|kr||jq|Sr )rUrWrr)rusedrdrrr used_nodess  zTaskGroup.used_nodescCs6t|}|||jt|p(t||fdS)a6 Add a "report step" to this TaskGroup. This step will run repeatedly every `interval_ms` milliseconds for the duration of the TaskGroup execution on each of the nodes. It is guaranteed that this step will be run at least once after every Task in the node has finished. N)rto_execution_stepRunEveryMillisrYrrr%r))rrOr interval_msrrr report_steps  zTaskGroup.report_stepcCs`t|pt|}|dks(||jks(t||jkrR|r:|n td||f|j|<|j|dS)z6 DEPRECATED. Use report_step instead. Nz %s/reporterr)rr%r)rXrarr6)rr^rZreport_intervalrrr report_nets  zTaskGroup.report_netc Csxi}|D]}|r||jn|j||j<q |jdk rT|j\}}||ksPtd|St|jD] \}\}}|j|||ddq^i|_tt}|D]}||j} ||  |qtt} |j D]\} } | ||  | qt } t|D]z\}}| |}t t j dd|D||\}}|}g}tj}|D]X}|} | |tjk| dk rd| | ||7}|tjkr0tj}q0t|dkr| tdgt|dkr|d} ntjd ||d d } t|dkst|dkrTg}t|dkr| td ||| | t|dkrH| td ||t||} t|| |d| |dq| |f|_| S)Nz)Cannot call tasks_by_node multiple times.rk)rrncSsg|] }|qSr)get_step.0trrr )sz+TaskGroup.tasks_by_node..rempty%s:bodyT)Zconcurrent_substepsrLrMgrouped_by_node)rrOoutputsrPgroupr])rgrr[rarrXrorr9rrYrSrK LOCAL_SETUPr+r.rrSetCreateWorkspacer]r{r/r>rrNTask)rZ node_remapZnode_maprd tasks_by_nodeZ prev_node_maprr^intervalZ mapped_nodeZreport_steps_by_nodeZ original_noderOrzrg report_stepsZ node_initsZ node_exitsrQr{Zgrouped_workspace_typerrrr s                zTaskGroup.tasks_by_nodecs>tt|fdd}t|dkr6tS|dS)NcsSr r)xrrrTz#TaskGroup.to_task..r)rr%r)rrgr>r)rrrgrrrto_taskRs  zTaskGroup.to_taskcCs|jSr rZrrrrr]YszTaskGroup.workspace_typecCsd|j|j||S)Nz6TaskGroup(tasks={}, workspace_type={}, remote_nets={}))rrUrWr]r`rrrrr \s  zTaskGroup.__repr__)N)NNrk)NNrp)N)N)r!r"r#r$r}rr_r`rergrhrjrorqrrr]r rrrrrSs  I rSc@s:eZdZdZddZd ddZddZd d Zd d ZdS) TaskOutputzd Represents the output of a task. An output can be a blob, a list of blob, or a record. cCsTd|_d|_t|tr&||_|j}t|ttfk|_|jrD|g}||_d|_ dSrT) _schema _is_scalarr8rZ field_blobsr=r:r9names_values)rrrrrris  zTaskOutput.__init__NcCs&t|t|jkst||_||_dSr )r>rrar _fetch_func)rvaluesrrrrsetuszTaskOutput.setcCs@|jdk std|jr"|jdS|jr6t|j|jS|jSdS)NzOutput value not set yet.r)rrarrrrrrrrzs  zTaskOutput.getcsNjdk stdfddjD}jr4|dSjrFtj|S|SdS)Nz#Cannot fetch value for this output.csg|]}|qSr)r)rtvrrrrvsz$TaskOutput.fetch..r)rrarrrr)rZ fetched_valsrrrfetchs  zTaskOutput.fetchcCsd|j|jS)NzTaskOutput(names={}, values={}))rrrrrrrr szTaskOutput.__repr__)N) r!r"r#r$rrrrr rrrrrcs    rcCstjddpt}||S)a  Adds an output to the current Task, or if no task is active, create a dummy task that returns the given blob or record to the client. This will return the value of the blob or record when the last task of the TaskGroup for a given node finishes. Frequired)rr) add_output)Zblob_or_recordZcur_taskrrr final_outputsrc@s4eZdZdZd ddZddZd ddZd d ZdS) TaskOutputListz$ Keeps a list of outputs for a task NcCs|pg|_dSr )r{rr{rrrrszTaskOutputList.__init__cCsg}|jD]}||j7}q |S)z[ Retrive the output names. TODO(azzolini): make this schema-based. )r{r)rrorrrrs  zTaskOutputList.namescCsRd}|jD].}t|j}||||||||7}q |t|ksNtddS)NrzWrong number of output values.)r{r>rrra)rrroffsetrnumrrr set_valuess    zTaskOutputList.set_valuescCs d|jS)NzTaskOutputList(outputs={}))rr{rrrrr szTaskOutputList.__repr__)N)N)r!r"r#r$rrrr rrrrrs   rcseZdZdZdZdZdZeZe ddZ d"dd Z fd d Z fd d Z ddZddZddZddZddZddZddZddZddZd d!ZZS)#raA A Task is composed of an execution step and zero or more outputs. Tasks are executed in the context of a TaskGroup, which, in turn, can be run by a Session. Task outputs are fetched by the session at the end of the run. The recommended way of creating a task is by using `net_builder.ops`. Example: from net_builder import ops with Node('trainer'), Task(name='my_task', num_instances=2): with ops.task_init(): globl = ops.Const(0) with ops.task_instance_init(): local = ops.Const(0) with ops.loop(100): ops.Copy(globl, local) with ops.task_instance_exit(): ops.Add([globl, local], [globl]) with ops.task_exit(): ops.Mul([globl, globl], [globl]) The task above will create 2 instances that will run in parallel. Each instance will copy `local` to `globl` 100 times, Then Add `local` to `globl` once. The `Mul` will only execute once, after all the instances of the task have finished. Z task_setupZtask_instance_setuprocCs`t|dt|}|dkr"tjntdd|jD}|}d}||kr\|d7}d||f}q>|S)N/css|] }|jVqdSr )rPrsrrr sz&Task._get_next_name..rrxz%s:%d)rr_global_names_usedrrW)rr|rPbasenameZ names_usedZcur_nameirrr_get_next_nameszTask._get_next_nameNcCs|st|tjr|j}|s"d}tt|dkr4dnt||_t j|dd|_ t |j|j ||_|j dk r~|j j |d|_d|_d|_g|_|dk r|||dk r||d|_d|_||_d|_||_dS)aV Instantiate a Task and add it to the current TaskGroup and Node. Args: step: If provided, this task will run this ExecutionStep. outputs: If provided, the task will return the provided outputs to the client at completion time. node: If provided, force task execution on the given node. name: Name of the Task. num_instances: If provided, this task will be cloned num_instances times at runtime, and all instances will run concurrently. rdNFr)r8rr;r?rPrr%r)rrSr|rrrWrrV_step_step_with_setup_outputsset_step add_outputsZ _pipelineZ_is_pipeline_contextrZZ _report_net_num_instances)rrOr{r]r|rrPZ num_instancesrrrrs,     z Task.__init__csjtt||jdk r&|jj|||jdks@tdddl m }|j |j d|_ |j |S)Nz(This Task already has an execution step.r) net_builder)Z _fullname)superr __enter__r|rWremove_assert_not_usedrra caffe2.pythonrZ NetBuilderrP _net_builder)rr __class__rrrs   zTask.__enter__csZtt|||||j||||dkr8||j|jdk rP|jj|d|_dSr )rr__exit__rrr|rWr)rr=value tracebackrrrr's  z Task.__exit__cCs|jSr rrrrrr]1szTask.workspace_typecCs|jrtddS)Nz1Cannot modify task since it is already been used.)rVrarrrrr4szTask._assert_not_usedcCs.|t|tr|nt|}|j||Sr )rr8rrr)routputrrrr8s  zTask.add_outputcs8t|ttfkr"|Sfdd|DSdS)Ncsg|]}|qSr)r)rtrrrrrvDsz$Task.add_outputs..)rr=r9r:rrrrrr?s zTask.add_outputscCs|t||_dSr )rrrlr)rrOrrrrFsz Task.set_stepc Csf|jdk r|jS|jdkr0t|jg|_|jSdd|jtjD}|D]}d|_| j sL| dqLt tj |jg||\}}t tj|jg||\}}t|jdkrtd|j}||jgdtjjdd|||s|jntd |j||jg}t||||jd } |jrN|jdkrN| dtjd | g|jd } t| |||j|_|jS) NcSsg|]}t|ds|qS)_report_step_used)r7)rtsrrrrvRs z!Task.get_step..Trkrz %s:outputrx)Zdtyperryz :instancez %s:parallel)Znum_concurrent_instances)rrrrNrPr0r REPORT_STEPrr?Z run_every_msrmrK TASK_SETUPTASK_INSTANCE_SETUPr>rr6rZ ConstantFillZDataTypeZINT32rrRrr~) rrrOZtask_init_netsZtask_exit_netsZinstance_init_netsZinstance_exit_netsZ output_netbodyZstep_with_instance_setuprrrrrJsn          z Task.get_stepcCs t|jSr )rrrrrr output_listyszTask.output_listcCs|jSr )rrrrrr{|sz Task.outputscCs|d|_dSrf)rrrVrrrrrbszTask._notify_usedcCsd|j|j|S)Nz"Task(name={}, node={}, outputs={}))rrPrr{rrrrr s z Task.__repr__)NNNNNNN)r!r"r#r$rrrrr staticmethodrrrrr]rrrrrrrr{rbr __classcell__rrrrrs6  . /rc@s2eZdZdZd ddZddZddZd d ZdS) SetupNetsa Allow to register a list of nets to be run at initialization and finalization of Tasks or TaskGroups. For example, let's say you have the following: init_net = core.Net('init') my_val = init_net.ConstantFill([], 'my_val', value=0) net = core.Net('counter') net.Add([my_val, net.Const(1),], [my_val]) with TaskGroup() as task_group: with Node('trainer'): my_task = Task(step=[net]) In order to have `init_net` run once before `net` runs for the first time, you can do one of the following: net.add_attribute(Task.TASK_SETUP, SetupNets([init_net])) or net.add_attribute(TaskGroup.LOCAL_SETUP, SetupNets([init_net])) - With Task.TASK_SETUP, init_net will run once at my_task startup. - With TaskGroup.LOCAL_SETUP, init_net will run once on node 'trainer', before any task of the task group is run on that node. The same SetupNets object can be added to multiple nets. It will only run once per Task/TaskGroup run. NcCs||_||_dSr )rFrG)rrFrGrrrrszSetupNets.__init__cCs|jSr )rF)rrDrrrr4szSetupNets.setupcCs|jSr )rG)rrErrrr5szSetupNets.exitcCsd|j|jS)Nz%SetupNets(init_nets={}, exit_nets={}))rrFrGrrrrr szSetupNets.__repr__)NN)r!r"r#r$rr4r5r rrrrrs  rN)rrrZcaffe2.python.schemarr collectionsrrZ future.utilsrrZDefaultManagedrr%objectr+rKrRZManagedrSrrrrrrrrrs$    !& , E/ Q