
    #Sh                         d dl Z d dlZd dlmZ d dlmZmZ d dlZddlm	Z	 ddl
mZ ddlmZ dd	lmZ  G d
 d          ZdeddfdZdeddfdZdeddfdZdeddfdZdefdZdeddfdZdeddfdZdeddfdZdS )    N)Path)AnyTuple   )utils)Job)JobEnvironment   )debugc                   6    e Zd ZdZdeddfdZdededefdZdS )	CheckFunctionz8Function used for checking that computations are correctnreturnNc                     t          t          |                    | _        t          t          dd|z                       | _        d S )N
   )listrangedata1data2)selfr   s     _/var/www/html/movieo_spanner_bot/venv/lib/python3.11/site-packages/submitit/local/test_debug.py__init__zCheckFunction.__init__   s6    %((^^
%BF++,,


    xyc                 8    || j         v sJ || j        v sJ ||z   S N)r   r   )r   r   r   s      r   __call__zCheckFunction.__call__   s)    DJDJ1ur   )__name__
__module____qualname____doc__intr   floatr    r   r   r   r      se        BB-# -$ - - - -% E e      r   r   tmp_pathr   c                    dt           dt           fd}t          j        |           }|                    |d          }|                                dk    sJ |                                5  |                    |d          }d d d            n# 1 swxY w Y   |                                dk    sJ d |                                _        |                                dk    sJ d S )Npr   c                     d| z  S )Nr   r%   )r(   s    r   funcztest_debug_job.<locals>.func"   s    1ur            r   )r#   r   DebugExecutorsubmitresultbatch
submissionfunction)r&   r*   executorjobjob2s        r   test_debug_jobr7   !   s$         "8,,H
//$
"
"C::<<1			 ( (tQ''( ( ( ( ( ( ( ( ( ( ( ( ( ( (;;==B!%DOO;;==Bs   +BBBc                 `   t          d          }t          j        |           }|                    ||j        |j                  }t          t          |d                              t          |           t          t          ||j        |j                            d |D             k    sJ d S )Nr-   r   c                 6    g | ]}|                                 S r%   r0   .0js     r   
<listcomp>z(test_debug_map_array.<locals>.<listcomp>6        -G-G-GQahhjj-G-G-Gr   )
r   r   r.   	map_arrayr   r   printtyper   map)r&   gr4   jobss       r   test_debug_map_arrayrF   0   s    aA"8,,Ha!'22D	$tAw--	$KKKAqw(())-G-G$-G-G-GGGGGGGr   c                 4   t          d          t          j        |           }fdt          j        j                  D             }|                    |          }t          t          j        j                            d |D             k    sJ d S )Nr-   c                 B    g | ]\  }}t          j        ||          S r%   )	functoolspartial)r<   r   r   rD   s      r   r>   z+test_debug_submit_array.<locals>.<listcomp><   s,    
H
H
H$!Q9Q1%%
H
H
Hr   c                 6    g | ]}|                                 S r%   r:   r;   s     r   r>   z+test_debug_submit_array.<locals>.<listcomp>>   r?   r   )	r   r   r.   zipr   r   submit_arrayr   rC   )r&   r4   fnsrE   rD   s       @r   test_debug_submit_arrayrO   9   s    aA"8,,H
H
H
H
H#agqw2G2G
H
H
HC  %%DAqw(())-G-G$-G-G-GGGGGGGr   c                     dd}t          j        |           }|                    |          }|                                }t	          |t
                    sJ |j        d         }d|v sJ d S )Nr   c                       t          d          )NFailed on purpose)RuntimeErrorr%   r   r   failing_jobz%test_debug_error.<locals>.failing_jobB   s    .///r   r   rR   )r   N)r   r.   r/   	exception
isinstancerS   args)r&   rT   r4   r5   rU   messages         r   test_debug_errorrY   A   s    0 0 0 0 "8,,H
//+
&
&CIi.....nQG'))))))r   c                      dS )N*   r%   r%   r   r   f_42r\   M   s    2r   c                 \   dt           dt          t          t          f         fd}t	          j        |           }dD ]p}|                    t                    }|j        dk    sJ  ||          dk    sJ  t          ||                       |j        dk    sJ  ||          dk    sJ qd S )	Nr5   r   c                 j    t          | t          j                  sJ | j        j        | j        j        fS r   )rV   r   DebugJob_submission_done_result)r5   s    r   
get_resultz(test_debug_triggered.<locals>.get_resultR   s/    #u~.....%s'>??r   )waitdonerU   resultsQUEUED)FNDONE)Tr[   )
r   r   boolr   r   r.   r/   r\   stategetattr)r&   rc   r4   triggerr5   s        r   test_debug_triggeredrm   Q   s    @ @dCi 0 @ @ @ @ "8,,H; - -ood##yH$$$$z#-////WyF""""z#*,,,,,- -r   c                 X   t          j        |           }|                    t                    }|j        dk    sJ |                                 |j        dk    sJ t          j        t          j	        d          5  |
                                 d d d            d S # 1 swxY w Y   d S )Nrg   	CANCELLEDzwas cancelled)match)r   r.   r/   r\   rj   cancelpytestraisesr   UncompletedJobErrorr0   )r&   r4   r5   s      r   test_cancelru   `   s    "8,,H
//$

C9    JJLLL9####	u0	H	H	H  

                 s   =BB#&B#c                     t          j        |           }d }|                    |          }|                                 dt          j        vsJ d S )Nc                      t                      } | j        dk    sJ | j        dk    sJ | j        dk    sJ | j        dk    sJ | j        dk    sJ d| j        v sJ d S )Nr
   r   DEBUG)r	   	num_nodes	num_tasksnodeglobal_rank
local_rankjob_id)envs    r   use_envz%test_job_environment.<locals>.use_envm   s    }!!!!}!!!!x1}}}}!####~""""#*$$$$$$r   SUBMITIT_DEBUG_JOB_ID)r   r.   r/   r0   osenviron)r&   r4   r   r5   s       r   test_job_environmentr   j   sZ    "8,,H% % % //'
"
"CJJLLL""*444444r   )rI   r   pathlibr   typingr   r   rr   corer   	core.corer   core.job_environmentr	    r   r   r7   rF   rO   rY   r#   r\   rm   ru   r   r%   r   r   <module>r      s       				                            1 1 1 1 1 1      
 
 
 
 
 
 
 
T d    H4 HD H H H HHd Ht H H H H	*t 	* 	* 	* 	* 	*c    -4 -D - - - -$ 4    54 5D 5 5 5 5 5 5r   