Skip to content

Commit

Permalink
Deploying to gh-pages from @ a1e5777 🚀
Browse files Browse the repository at this point in the history
  • Loading branch information
@jessicasyu
jessicasyu committed Aug 28, 2024
1 parent a5eef02 commit c00a8e4
Show file tree
Hide file tree
Showing 16 changed files with 297 additions and 26 deletions.
41 changes: 35 additions & 6 deletions _modules/container_collection/batch/check_batch_job.html
View file
Original file line number Diff line number Diff line change
Expand Up @@ -243,16 +243,42 @@ <h1>Source code for container_collection.batch.check_batch_job</h1><div class="h
<span class="kn">from</span> <span class="nn">typing</span> <span class="kn">import</span> <span class="n">Union</span>

<span class="kn">import</span> <span class="nn">boto3</span>
<span class="kn">import</span> <span class="nn">prefect</span>
<span class="kn">from</span> <span class="nn">prefect.server.schemas.states</span> <span class="kn">import</span> <span class="n">Failed</span><span class="p">,</span> <span class="n">State</span>
<span class="kn">from</span> <span class="nn">prefect.context</span> <span class="kn">import</span> <span class="n">TaskRunContext</span>
<span class="kn">from</span> <span class="nn">prefect.states</span> <span class="kn">import</span> <span class="n">Failed</span><span class="p">,</span> <span class="n">State</span>

<span class="n">RETRIES_EXCEEDED_EXIT_CODE</span> <span class="o">=</span> <span class="mi">80</span>
<span class="sd">&quot;&quot;&quot;Exit code used when task run retries exceed the maximum retries.&quot;&quot;&quot;</span>


<div class="viewcode-block" id="check_batch_job"><a class="viewcode-back" href="../../../_summary/container_collection.batch.check_batch_job.html#container_collection.batch.check_batch_job.check_batch_job">[docs]</a><span class="k">def</span> <span class="nf">check_batch_job</span><span class="p">(</span><span class="n">job_arn</span><span class="p">:</span> <span class="nb">str</span><span class="p">,</span> <span class="n">max_retries</span><span class="p">:</span> <span class="nb">int</span><span class="p">)</span> <span class="o">-&gt;</span> <span class="n">Union</span><span class="p">[</span><span class="nb">int</span><span class="p">,</span> <span class="n">State</span><span class="p">]:</span>
<span class="n">task_run</span> <span class="o">=</span> <span class="n">prefect</span><span class="o">.</span><span class="n">context</span><span class="o">.</span><span class="n">get_run_context</span><span class="p">()</span><span class="o">.</span><span class="n">task_run</span> <span class="c1"># type: ignore</span>
<div class="viewcode-block" id="check_batch_job"><a class="viewcode-back" href="../../../_summary/container_collection.batch.check_batch_job.html#container_collection.batch.check_batch_job.check_batch_job">[docs]</a><span class="k">def</span> <span class="nf">check_batch_job</span><span class="p">(</span><span class="n">job_arn</span><span class="p">:</span> <span class="nb">str</span><span class="p">,</span> <span class="n">max_retries</span><span class="p">:</span> <span class="nb">int</span><span class="p">)</span> <span class="o">-&gt;</span> <span class="n">Union</span><span class="p">[</span><span class="nb">int</span><span class="p">,</span> <span class="n">State</span><span class="p">,</span> <span class="nb">bool</span><span class="p">]:</span>
<span class="w"> </span><span class="sd">&quot;&quot;&quot;</span>
<span class="sd"> Check for exit code of an AWS Batch job.</span>

<span class="k">if</span> <span class="n">task_run</span><span class="o">.</span><span class="n">run_count</span> <span class="o">&gt;</span> <span class="n">max_retries</span><span class="p">:</span>
<span class="sd"> If this task is running within a Prefect flow, it will use the task run</span>
<span class="sd"> context to get the current run count. While the run count is below the</span>
<span class="sd"> maximum number of retries, the task will continue to attempt to get the exit</span>
<span class="sd"> code, and can be called with a retry delay to periodically check the status</span>
<span class="sd"> of jobs.</span>

<span class="sd"> If this task is not running within a Prefect flow, the ``max_retries``</span>
<span class="sd"> parameters is ignored. Jobs that are still running will throw an exception.</span>

<span class="sd"> Parameters</span>
<span class="sd"> ----------</span>
<span class="sd"> job_arn</span>
<span class="sd"> Job ARN.</span>
<span class="sd"> max_retries</span>
<span class="sd"> Maximum number of retries.</span>

<span class="sd"> Returns</span>
<span class="sd"> -------</span>
<span class="sd"> :</span>
<span class="sd"> Exit code if the job is complete, otherwise throws an exception.</span>
<span class="sd"> &quot;&quot;&quot;</span>

<span class="n">context</span> <span class="o">=</span> <span class="n">TaskRunContext</span><span class="o">.</span><span class="n">get</span><span class="p">()</span>

<span class="k">if</span> <span class="n">context</span> <span class="ow">is</span> <span class="ow">not</span> <span class="kc">None</span> <span class="ow">and</span> <span class="n">context</span><span class="o">.</span><span class="n">task_run</span><span class="o">.</span><span class="n">run_count</span> <span class="o">&gt;</span> <span class="n">max_retries</span><span class="p">:</span>
<span class="k">return</span> <span class="n">RETRIES_EXCEEDED_EXIT_CODE</span>

<span class="n">client</span> <span class="o">=</span> <span class="n">boto3</span><span class="o">.</span><span class="n">client</span><span class="p">(</span><span class="s2">&quot;batch&quot;</span><span class="p">)</span>
Expand All @@ -271,8 +297,11 @@ <h1>Source code for container_collection.batch.check_batch_job</h1><div class="h
<span class="n">response</span> <span class="o">=</span> <span class="n">client</span><span class="o">.</span><span class="n">describe_jobs</span><span class="p">(</span><span class="n">jobs</span><span class="o">=</span><span class="p">[</span><span class="n">job_arn</span><span class="p">])[</span><span class="s2">&quot;jobs&quot;</span><span class="p">]</span>
<span class="n">status</span> <span class="o">=</span> <span class="n">response</span><span class="p">[</span><span class="mi">0</span><span class="p">][</span><span class="s2">&quot;status&quot;</span><span class="p">]</span>

<span class="k">if</span> <span class="n">status</span> <span class="o">==</span> <span class="s2">&quot;RUNNING&quot;</span><span class="p">:</span>
<span class="c1"># For jobs that are running, throw the appropriate exception.</span>
<span class="k">if</span> <span class="n">context</span> <span class="ow">is</span> <span class="ow">not</span> <span class="kc">None</span> <span class="ow">and</span> <span class="n">status</span> <span class="o">==</span> <span class="s2">&quot;RUNNING&quot;</span><span class="p">:</span>
<span class="k">return</span> <span class="n">Failed</span><span class="p">()</span>
<span class="k">if</span> <span class="n">status</span> <span class="o">==</span> <span class="s2">&quot;RUNNING&quot;</span><span class="p">:</span>
<span class="k">raise</span> <span class="ne">RuntimeError</span><span class="p">(</span><span class="s2">&quot;Job is in RUNNING state and does not have exit code.&quot;</span><span class="p">)</span>

<span class="n">exitcode</span> <span class="o">=</span> <span class="n">response</span><span class="p">[</span><span class="mi">0</span><span class="p">][</span><span class="s2">&quot;attempts&quot;</span><span class="p">][</span><span class="mi">0</span><span class="p">][</span><span class="s2">&quot;container&quot;</span><span class="p">][</span><span class="s2">&quot;exitCode&quot;</span><span class="p">]</span>
<span class="k">return</span> <span class="n">exitcode</span></div>
Expand Down
26 changes: 24 additions & 2 deletions _modules/container_collection/batch/get_batch_logs.html
View file
Original file line number Diff line number Diff line change
Expand Up @@ -241,8 +241,27 @@
<h1>Source code for container_collection.batch.get_batch_logs</h1><div class="highlight"><pre>
<span></span><span class="kn">import</span> <span class="nn">boto3</span>

<span class="n">LOG_GROUP_NAME</span> <span class="o">=</span> <span class="s2">&quot;/aws/batch/job&quot;</span>
<span class="sd">&quot;&quot;&quot;AWS Batch log group name.&quot;&quot;&quot;</span>


<div class="viewcode-block" id="get_batch_logs"><a class="viewcode-back" href="../../../_summary/container_collection.batch.get_batch_logs.html#container_collection.batch.get_batch_logs.get_batch_logs">[docs]</a><span class="k">def</span> <span class="nf">get_batch_logs</span><span class="p">(</span><span class="n">job_arn</span><span class="p">:</span> <span class="nb">str</span><span class="p">,</span> <span class="n">log_filter</span><span class="p">:</span> <span class="nb">str</span><span class="p">)</span> <span class="o">-&gt;</span> <span class="nb">str</span><span class="p">:</span>
<span class="w"> </span><span class="sd">&quot;&quot;&quot;</span>
<span class="sd"> Get logs for AWS Batch job.</span>

<span class="sd"> Parameters</span>
<span class="sd"> ----------</span>
<span class="sd"> job_arn</span>
<span class="sd"> Job ARN.</span>
<span class="sd"> log_filter</span>
<span class="sd"> Filter for log events.</span>

<span class="sd"> Returns</span>
<span class="sd"> -------</span>
<span class="sd"> :</span>
<span class="sd"> All filtered log events.</span>
<span class="sd"> &quot;&quot;&quot;</span>

<span class="n">client</span> <span class="o">=</span> <span class="n">boto3</span><span class="o">.</span><span class="n">client</span><span class="p">(</span><span class="s2">&quot;batch&quot;</span><span class="p">)</span>
<span class="n">response</span> <span class="o">=</span> <span class="n">client</span><span class="o">.</span><span class="n">describe_jobs</span><span class="p">(</span><span class="n">jobs</span><span class="o">=</span><span class="p">[</span><span class="n">job_arn</span><span class="p">])[</span><span class="s2">&quot;jobs&quot;</span><span class="p">][</span><span class="mi">0</span><span class="p">]</span>
<span class="n">log_stream</span> <span class="o">=</span> <span class="n">response</span><span class="p">[</span><span class="s2">&quot;container&quot;</span><span class="p">][</span><span class="s2">&quot;logStreamName&quot;</span><span class="p">]</span>
Expand All @@ -251,21 +270,24 @@ <h1>Source code for container_collection.batch.get_batch_logs</h1><div class="hi
<span class="n">log_events</span><span class="p">:</span> <span class="nb">list</span><span class="p">[</span><span class="nb">str</span><span class="p">]</span> <span class="o">=</span> <span class="p">[]</span>

<span class="n">response</span> <span class="o">=</span> <span class="n">client</span><span class="o">.</span><span class="n">filter_log_events</span><span class="p">(</span>
<span class="n">logGroupName</span><span class="o">=</span><span class="s2">&quot;/aws/batch/job&quot;</span><span class="p">,</span>
<span class="n">logGroupName</span><span class="o">=</span><span class="n">LOG_GROUP_NAME</span><span class="p">,</span>
<span class="n">logStreamNames</span><span class="o">=</span><span class="p">[</span><span class="n">log_stream</span><span class="p">],</span>
<span class="n">filterPattern</span><span class="o">=</span><span class="n">log_filter</span><span class="p">,</span>
<span class="p">)</span>

<span class="k">while</span> <span class="s2">&quot;nextToken&quot;</span> <span class="ow">in</span> <span class="n">response</span><span class="p">:</span>
<span class="k">if</span> <span class="n">response</span><span class="p">[</span><span class="s2">&quot;events&quot;</span><span class="p">]:</span>
<span class="n">log_events</span> <span class="o">=</span> <span class="n">log_events</span> <span class="o">+</span> <span class="p">[</span><span class="n">event</span><span class="p">[</span><span class="s2">&quot;message&quot;</span><span class="p">]</span> <span class="k">for</span> <span class="n">event</span> <span class="ow">in</span> <span class="n">response</span><span class="p">[</span><span class="s2">&quot;events&quot;</span><span class="p">]]</span>

<span class="k">while</span> <span class="s2">&quot;nextToken&quot;</span> <span class="ow">in</span> <span class="n">response</span><span class="p">:</span>
<span class="n">response</span> <span class="o">=</span> <span class="n">client</span><span class="o">.</span><span class="n">filter_log_events</span><span class="p">(</span>
<span class="n">logGroupName</span><span class="o">=</span><span class="s2">&quot;/aws/batch/job&quot;</span><span class="p">,</span>
<span class="n">logStreamNames</span><span class="o">=</span><span class="p">[</span><span class="n">log_stream</span><span class="p">],</span>
<span class="n">filterPattern</span><span class="o">=</span><span class="n">log_filter</span><span class="p">,</span>
<span class="n">nextToken</span><span class="o">=</span><span class="n">response</span><span class="p">[</span><span class="s2">&quot;nextToken&quot;</span><span class="p">],</span>
<span class="p">)</span>

<span class="n">log_events</span> <span class="o">=</span> <span class="n">log_events</span> <span class="o">+</span> <span class="p">[</span><span class="n">event</span><span class="p">[</span><span class="s2">&quot;message&quot;</span><span class="p">]</span> <span class="k">for</span> <span class="n">event</span> <span class="ow">in</span> <span class="n">response</span><span class="p">[</span><span class="s2">&quot;events&quot;</span><span class="p">]]</span>

<span class="k">return</span> <span class="s2">&quot;</span><span class="se">\n</span><span class="s2">&quot;</span><span class="o">.</span><span class="n">join</span><span class="p">(</span><span class="n">log_events</span><span class="p">)</span></div>
</pre></div>
</article>
Expand Down
38 changes: 38 additions & 0 deletions _modules/container_collection/batch/make_batch_job.html
View file
Original file line number Diff line number Diff line change
Expand Up @@ -250,6 +250,44 @@ <h1>Source code for container_collection.batch.make_batch_job</h1><div class="hi
<span class="n">environment</span><span class="p">:</span> <span class="n">Optional</span><span class="p">[</span><span class="nb">list</span><span class="p">[</span><span class="nb">dict</span><span class="p">[</span><span class="nb">str</span><span class="p">,</span> <span class="nb">str</span><span class="p">]]]</span> <span class="o">=</span> <span class="kc">None</span><span class="p">,</span>
<span class="n">job_role_arn</span><span class="p">:</span> <span class="n">Optional</span><span class="p">[</span><span class="nb">str</span><span class="p">]</span> <span class="o">=</span> <span class="kc">None</span><span class="p">,</span>
<span class="p">)</span> <span class="o">-&gt;</span> <span class="nb">dict</span><span class="p">:</span>
<span class="w"> </span><span class="sd">&quot;&quot;&quot;</span>
<span class="sd"> Create batch job definition.</span>

<span class="sd"> Docker images on the Docker Hub registry are available by default, and can</span>
<span class="sd"> be specified using ``image:tag``. Otherwise, use ``repository/image:tag``.</span>

<span class="sd"> Environment variables are passed as key-value pairs using the following</span>
<span class="sd"> structure:</span>

<span class="sd"> .. code-block:: python</span>

<span class="sd"> [</span>
<span class="sd"> { &quot;name&quot; : &quot;envName1&quot;, &quot;value&quot; : &quot;envValue1&quot; },</span>
<span class="sd"> { &quot;name&quot; : &quot;envName2&quot;, &quot;value&quot; : &quot;envValue2&quot; },</span>
<span class="sd"> ...</span>
<span class="sd"> ]</span>

<span class="sd"> Parameters</span>
<span class="sd"> ----------</span>
<span class="sd"> name</span>
<span class="sd"> Job definition name.</span>
<span class="sd"> image</span>
<span class="sd"> Docker image.</span>
<span class="sd"> vcpus</span>
<span class="sd"> Number of vCPUs to reserve for the container.</span>
<span class="sd"> memory</span>
<span class="sd"> Memory limit available to the container</span>
<span class="sd"> environment</span>
<span class="sd"> List of environment variables as key-value pairs.</span>
<span class="sd"> job_role_arn</span>
<span class="sd"> ARN for IAM role for the job container.</span>

<span class="sd"> Returns</span>
<span class="sd"> -------</span>
<span class="sd"> :</span>
<span class="sd"> Job definition.</span>
<span class="sd"> &quot;&quot;&quot;</span>

<span class="n">container_properties</span> <span class="o">=</span> <span class="p">{</span>
<span class="s2">&quot;image&quot;</span><span class="p">:</span> <span class="n">image</span><span class="p">,</span>
<span class="s2">&quot;vcpus&quot;</span><span class="p">:</span> <span class="n">vcpus</span><span class="p">,</span>
Expand Down
19 changes: 19 additions & 0 deletions _modules/container_collection/batch/register_batch_job.html
View file
Original file line number Diff line number Diff line change
Expand Up @@ -244,6 +244,24 @@ <h1>Source code for container_collection.batch.register_batch_job</h1><div class


<div class="viewcode-block" id="register_batch_job"><a class="viewcode-back" href="../../../_summary/container_collection.batch.register_batch_job.html#container_collection.batch.register_batch_job.register_batch_job">[docs]</a><span class="k">def</span> <span class="nf">register_batch_job</span><span class="p">(</span><span class="n">job_definition</span><span class="p">:</span> <span class="nb">dict</span><span class="p">)</span> <span class="o">-&gt;</span> <span class="nb">str</span><span class="p">:</span>
<span class="w"> </span><span class="sd">&quot;&quot;&quot;</span>
<span class="sd"> Register job definition to AWS Batch.</span>

<span class="sd"> If a definition for the given job definition name already exists, and the</span>
<span class="sd"> contents of the definition are not changed, then the method will return the</span>
<span class="sd"> existing job definition ARN rather than creating a new revision.</span>

<span class="sd"> Parameters</span>
<span class="sd"> ----------</span>
<span class="sd"> job_definition</span>
<span class="sd"> Batch job definition.</span>

<span class="sd"> Returns</span>
<span class="sd"> -------</span>
<span class="sd"> :</span>
<span class="sd"> Job definition ARN.</span>
<span class="sd"> &quot;&quot;&quot;</span>

<span class="n">client</span> <span class="o">=</span> <span class="n">boto3</span><span class="o">.</span><span class="n">client</span><span class="p">(</span><span class="s2">&quot;batch&quot;</span><span class="p">)</span>

<span class="n">response</span> <span class="o">=</span> <span class="n">client</span><span class="o">.</span><span class="n">describe_job_definitions</span><span class="p">(</span>
Expand All @@ -258,6 +276,7 @@ <h1>Source code for container_collection.batch.register_batch_job</h1><div class
<span class="k">return</span> <span class="n">existing_definition</span><span class="p">[</span><span class="s2">&quot;jobDefinitionArn&quot;</span><span class="p">]</span>

<span class="n">response</span> <span class="o">=</span> <span class="n">client</span><span class="o">.</span><span class="n">register_job_definition</span><span class="p">(</span><span class="o">**</span><span class="n">job_definition</span><span class="p">)</span>

<span class="k">return</span> <span class="n">response</span><span class="p">[</span><span class="s2">&quot;jobDefinitionArn&quot;</span><span class="p">]</span></div>
</pre></div>
</article>
Expand Down
Loading

0 comments on commit c00a8e4

Please sign in to comment.