deploy: 318da40

Author: chenyushuo

_modules/trinity/common/models.html

@@ -123,7 +123,7 @@ <h1>Source code for trinity.common.models</h1><div class="highlight"><pre>
 <a class="viewcode-back" href="../../../build_api/trinity.common.models.html#trinity.common.models.create_inference_models">[docs]</a>
 <span class="k">def</span><span class="w"> </span><span class="nf">create_inference_models</span><span class="p">(</span>
     <span class="n">config</span><span class="p">:</span> <span class="n">Config</span><span class="p">,</span>
-<span class="p">)</span> <span class="o">-&gt;</span> <span class="n">Tuple</span><span class="p">[</span><span class="n">List</span><span class="p">[</span><span class="n">InferenceModel</span><span class="p">],</span> <span class="n">List</span><span class="p">[</span><span class="n">InferenceModel</span><span class="p">]]:</span>
+<span class="p">)</span> <span class="o">-&gt;</span> <span class="n">Tuple</span><span class="p">[</span><span class="n">List</span><span class="p">[</span><span class="n">InferenceModel</span><span class="p">],</span> <span class="n">List</span><span class="p">[</span><span class="n">List</span><span class="p">[</span><span class="n">InferenceModel</span><span class="p">]]]:</span>
 <span class="w">    </span><span class="sd">&quot;&quot;&quot;Create `engine_num` rollout models.</span>
 
 <span class="sd">    Each model has `tensor_parallel_size` workers.</span>
@@ -205,11 +205,12 @@ <h1>Source code for trinity.common.models</h1><div class="highlight"><pre>
 
     <span class="c1"># create auxiliary models</span>
     <span class="k">for</span> <span class="n">model_config</span> <span class="ow">in</span> <span class="n">config</span><span class="o">.</span><span class="n">explorer</span><span class="o">.</span><span class="n">auxiliary_models</span><span class="p">:</span>
+        <span class="n">engines</span> <span class="o">=</span> <span class="p">[]</span>
         <span class="k">for</span> <span class="n">_</span> <span class="ow">in</span> <span class="nb">range</span><span class="p">(</span><span class="n">model_config</span><span class="o">.</span><span class="n">engine_num</span><span class="p">):</span>
             <span class="n">bundles_for_engine</span> <span class="o">=</span> <span class="n">allocator</span><span class="o">.</span><span class="n">allocate</span><span class="p">(</span><span class="n">model_config</span><span class="o">.</span><span class="n">tensor_parallel_size</span><span class="p">)</span>
             <span class="n">model_config</span><span class="o">.</span><span class="n">enable_openai_api</span> <span class="o">=</span> <span class="kc">True</span>
             <span class="n">model_config</span><span class="o">.</span><span class="n">engine_type</span> <span class="o">=</span> <span class="s2">&quot;vllm_async&quot;</span>
-            <span class="n">auxiliary_engines</span><span class="o">.</span><span class="n">append</span><span class="p">(</span>
+            <span class="n">engines</span><span class="o">.</span><span class="n">append</span><span class="p">(</span>
                 <span class="n">ray</span><span class="o">.</span><span class="n">remote</span><span class="p">(</span><span class="n">vLLMAysncRolloutModel</span><span class="p">)</span>
                 <span class="o">.</span><span class="n">options</span><span class="p">(</span>
                     <span class="n">num_cpus</span><span class="o">=</span><span class="mi">0</span><span class="p">,</span>
@@ -221,9 +222,11 @@ <h1>Source code for trinity.common.models</h1><div class="highlight"><pre>
                 <span class="p">)</span>
                 <span class="o">.</span><span class="n">remote</span><span class="p">(</span><span class="n">config</span><span class="o">=</span><span class="n">model_config</span><span class="p">)</span>
             <span class="p">)</span>
+        <span class="n">auxiliary_engines</span><span class="o">.</span><span class="n">append</span><span class="p">(</span><span class="n">engines</span><span class="p">)</span>
     <span class="c1"># all auxiliary engines run api server</span>
-    <span class="k">for</span> <span class="n">engine</span> <span class="ow">in</span> <span class="n">auxiliary_engines</span><span class="p">:</span>
-        <span class="n">engine</span><span class="o">.</span><span class="n">run_api_server</span><span class="o">.</span><span class="n">remote</span><span class="p">()</span>
+    <span class="k">for</span> <span class="n">engines</span> <span class="ow">in</span> <span class="n">auxiliary_engines</span><span class="p">:</span>
+        <span class="k">for</span> <span class="n">engine</span> <span class="ow">in</span> <span class="n">engines</span><span class="p">:</span>
+            <span class="n">engine</span><span class="o">.</span><span class="n">run_api_server</span><span class="o">.</span><span class="n">remote</span><span class="p">()</span>
 
     <span class="k">return</span> <span class="n">rollout_engines</span><span class="p">,</span> <span class="n">auxiliary_engines</span></div>
 

_modules/trinity/explorer/runner_pool.html

@@ -87,11 +87,12 @@
   <h1>Source code for trinity.explorer.runner_pool</h1><div class="highlight"><pre>
 <span></span><span class="sd">&quot;&quot;&quot;Runner pool for running tasks in parallel. Modified from ray.util.actor_pool.ActorPool.&quot;&quot;&quot;</span>
 <span class="kn">import</span><span class="w"> </span><span class="nn">random</span>
-<span class="kn">from</span><span class="w"> </span><span class="nn">typing</span><span class="w"> </span><span class="kn">import</span> <span class="n">List</span><span class="p">,</span> <span class="n">Tuple</span><span class="p">,</span> <span class="n">Union</span>
+<span class="kn">from</span><span class="w"> </span><span class="nn">typing</span><span class="w"> </span><span class="kn">import</span> <span class="n">List</span><span class="p">,</span> <span class="n">Optional</span><span class="p">,</span> <span class="n">Tuple</span><span class="p">,</span> <span class="n">Union</span>
 
 <span class="kn">import</span><span class="w"> </span><span class="nn">ray</span>
 
 <span class="kn">from</span><span class="w"> </span><span class="nn">trinity.common.config</span><span class="w"> </span><span class="kn">import</span> <span class="n">Config</span>
+<span class="kn">from</span><span class="w"> </span><span class="nn">trinity.common.models.model</span><span class="w"> </span><span class="kn">import</span> <span class="n">InferenceModel</span>
 <span class="kn">from</span><span class="w"> </span><span class="nn">trinity.common.workflows</span><span class="w"> </span><span class="kn">import</span> <span class="n">Task</span>
 <span class="kn">from</span><span class="w"> </span><span class="nn">trinity.explorer.workflow_runner</span><span class="w"> </span><span class="kn">import</span> <span class="n">Status</span><span class="p">,</span> <span class="n">WorkflowRunner</span>
 <span class="kn">from</span><span class="w"> </span><span class="nn">trinity.utils.log</span><span class="w"> </span><span class="kn">import</span> <span class="n">get_logger</span>
@@ -110,11 +111,17 @@ <h1>Source code for trinity.explorer.runner_pool</h1><div class="highlight"><pre
 
 <div class="viewcode-block" id="RunnerPool.__init__">
 <a class="viewcode-back" href="../../../build_api/trinity.explorer.html#trinity.explorer.runner_pool.RunnerPool.__init__">[docs]</a>
-    <span class="k">def</span><span class="w"> </span><span class="fm">__init__</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">config</span><span class="p">:</span> <span class="n">Config</span><span class="p">,</span> <span class="n">models</span><span class="p">:</span> <span class="n">List</span><span class="p">):</span>
+    <span class="k">def</span><span class="w"> </span><span class="fm">__init__</span><span class="p">(</span>
+        <span class="bp">self</span><span class="p">,</span>
+        <span class="n">config</span><span class="p">:</span> <span class="n">Config</span><span class="p">,</span>
+        <span class="n">models</span><span class="p">:</span> <span class="n">List</span><span class="p">[</span><span class="n">InferenceModel</span><span class="p">],</span>
+        <span class="n">auxiliary_models</span><span class="p">:</span> <span class="n">Optional</span><span class="p">[</span><span class="n">List</span><span class="p">[</span><span class="n">List</span><span class="p">[</span><span class="n">InferenceModel</span><span class="p">]]]</span> <span class="o">=</span> <span class="kc">None</span><span class="p">,</span>
+    <span class="p">):</span>
         <span class="c1"># actors to be used</span>
         <span class="bp">self</span><span class="o">.</span><span class="n">logger</span> <span class="o">=</span> <span class="n">get_logger</span><span class="p">(</span><span class="vm">__name__</span><span class="p">)</span>
         <span class="bp">self</span><span class="o">.</span><span class="n">config</span> <span class="o">=</span> <span class="n">config</span>
         <span class="bp">self</span><span class="o">.</span><span class="n">models</span> <span class="o">=</span> <span class="n">models</span>
+        <span class="bp">self</span><span class="o">.</span><span class="n">auxiliary_models</span> <span class="o">=</span> <span class="n">auxiliary_models</span> <span class="ow">or</span> <span class="p">[]</span>
         <span class="bp">self</span><span class="o">.</span><span class="n">timeout</span> <span class="o">=</span> <span class="n">config</span><span class="o">.</span><span class="n">explorer</span><span class="o">.</span><span class="n">max_timeout</span>
         <span class="bp">self</span><span class="o">.</span><span class="n">max_retry_times</span> <span class="o">=</span> <span class="n">config</span><span class="o">.</span><span class="n">explorer</span><span class="o">.</span><span class="n">max_retry_times</span>
 
@@ -135,6 +142,9 @@ <h1>Source code for trinity.explorer.runner_pool</h1><div class="highlight"><pre
 
         <span class="c1"># create new actors</span>
         <span class="bp">self</span><span class="o">.</span><span class="n">engine_status</span> <span class="o">=</span> <span class="p">[</span><span class="mi">0</span><span class="p">]</span> <span class="o">*</span> <span class="n">config</span><span class="o">.</span><span class="n">explorer</span><span class="o">.</span><span class="n">rollout_model</span><span class="o">.</span><span class="n">engine_num</span>
+        <span class="bp">self</span><span class="o">.</span><span class="n">auxiliary_engine_status_list</span> <span class="o">=</span> <span class="p">[</span>
+            <span class="p">[</span><span class="mi">0</span><span class="p">]</span> <span class="o">*</span> <span class="n">cfg</span><span class="o">.</span><span class="n">engine_num</span> <span class="k">for</span> <span class="n">cfg</span> <span class="ow">in</span> <span class="n">config</span><span class="o">.</span><span class="n">explorer</span><span class="o">.</span><span class="n">auxiliary_models</span>
+        <span class="p">]</span>
         <span class="bp">self</span><span class="o">.</span><span class="n">_idle_actors</span> <span class="o">=</span> <span class="nb">list</span><span class="p">()</span>
         <span class="bp">self</span><span class="o">.</span><span class="n">actor_to_engine_index</span> <span class="o">=</span> <span class="p">{}</span>
         <span class="bp">self</span><span class="o">.</span><span class="n">_create_actors</span><span class="p">(</span><span class="n">config</span><span class="o">.</span><span class="n">explorer</span><span class="o">.</span><span class="n">runner_num</span><span class="p">)</span></div>
@@ -144,7 +154,15 @@ <h1>Source code for trinity.explorer.runner_pool</h1><div class="highlight"><pre
         <span class="n">new_actors</span> <span class="o">=</span> <span class="p">[]</span>
         <span class="k">for</span> <span class="n">_</span> <span class="ow">in</span> <span class="nb">range</span><span class="p">(</span><span class="n">num</span><span class="p">):</span>
             <span class="n">engine_index</span> <span class="o">=</span> <span class="bp">self</span><span class="o">.</span><span class="n">engine_status</span><span class="o">.</span><span class="n">index</span><span class="p">(</span><span class="nb">min</span><span class="p">(</span><span class="bp">self</span><span class="o">.</span><span class="n">engine_status</span><span class="p">))</span>
-            <span class="n">new_actor</span> <span class="o">=</span> <span class="n">WorkflowRunner</span><span class="o">.</span><span class="n">remote</span><span class="p">(</span><span class="bp">self</span><span class="o">.</span><span class="n">config</span><span class="p">,</span> <span class="bp">self</span><span class="o">.</span><span class="n">models</span><span class="p">[</span><span class="n">engine_index</span><span class="p">])</span>
+            <span class="n">selected_auxiliary_models</span> <span class="o">=</span> <span class="p">[</span>
+                <span class="n">models</span><span class="p">[</span><span class="n">engine_status</span><span class="o">.</span><span class="n">index</span><span class="p">(</span><span class="nb">min</span><span class="p">(</span><span class="n">engine_status</span><span class="p">))]</span>
+                <span class="k">for</span> <span class="n">models</span><span class="p">,</span> <span class="n">engine_status</span> <span class="ow">in</span> <span class="nb">zip</span><span class="p">(</span>
+                    <span class="bp">self</span><span class="o">.</span><span class="n">auxiliary_models</span><span class="p">,</span> <span class="bp">self</span><span class="o">.</span><span class="n">auxiliary_engine_status_list</span>
+                <span class="p">)</span>
+            <span class="p">]</span>
+            <span class="n">new_actor</span> <span class="o">=</span> <span class="n">WorkflowRunner</span><span class="o">.</span><span class="n">remote</span><span class="p">(</span>
+                <span class="bp">self</span><span class="o">.</span><span class="n">config</span><span class="p">,</span> <span class="bp">self</span><span class="o">.</span><span class="n">models</span><span class="p">[</span><span class="n">engine_index</span><span class="p">],</span> <span class="n">selected_auxiliary_models</span>
+            <span class="p">)</span>
             <span class="n">new_actors</span><span class="o">.</span><span class="n">append</span><span class="p">(</span><span class="n">new_actor</span><span class="p">)</span>
             <span class="bp">self</span><span class="o">.</span><span class="n">engine_status</span><span class="p">[</span><span class="n">engine_index</span><span class="p">]</span> <span class="o">+=</span> <span class="mi">1</span>
             <span class="bp">self</span><span class="o">.</span><span class="n">actor_to_engine_index</span><span class="p">[</span><span class="n">new_actor</span><span class="p">]</span> <span class="o">=</span> <span class="n">engine_index</span>

build_api/trinity.common.models.html

@@ -660,7 +660,7 @@ <h2>Submodules<a class="headerlink" href="#submodules" title="Link to this headi
 <span id="module-contents"></span><h2>Module contents<a class="headerlink" href="#module-trinity.common.models" title="Link to this heading"></a></h2>
 <dl class="py function">
 <dt class="sig sig-object py" id="trinity.common.models.create_inference_models">
-<span class="sig-prename descclassname"><span class="pre">trinity.common.models.</span></span><span class="sig-name descname"><span class="pre">create_inference_models</span></span><span class="sig-paren">(</span><em class="sig-param"><span class="n"><span class="pre">config</span></span><span class="p"><span class="pre">:</span></span><span class="w"> </span><span class="n"><a class="reference internal" href="trinity.common.html#trinity.common.config.Config" title="trinity.common.config.Config"><span class="pre">Config</span></a></span></em><span class="sig-paren">)</span> <span class="sig-return"><span class="sig-return-icon">&#x2192;</span> <span class="sig-return-typehint"><span class="pre">Tuple</span><span class="p"><span class="pre">[</span></span><span class="pre">List</span><span class="p"><span class="pre">[</span></span><a class="reference internal" href="#trinity.common.models.model.InferenceModel" title="trinity.common.models.model.InferenceModel"><span class="pre">InferenceModel</span></a><span class="p"><span class="pre">]</span></span><span class="p"><span class="pre">,</span></span><span class="w"> </span><span class="pre">List</span><span class="p"><span class="pre">[</span></span><a class="reference internal" href="#trinity.common.models.model.InferenceModel" title="trinity.common.models.model.InferenceModel"><span class="pre">InferenceModel</span></a><span class="p"><span class="pre">]</span></span><span class="p"><span class="pre">]</span></span></span></span><a class="reference internal" href="../_modules/trinity/common/models.html#create_inference_models"><span class="viewcode-link"><span class="pre">[source]</span></span></a><a class="headerlink" href="#trinity.common.models.create_inference_models" title="Link to this definition"></a></dt>
+<span class="sig-prename descclassname"><span class="pre">trinity.common.models.</span></span><span class="sig-name descname"><span class="pre">create_inference_models</span></span><span class="sig-paren">(</span><em class="sig-param"><span class="n"><span class="pre">config</span></span><span class="p"><span class="pre">:</span></span><span class="w"> </span><span class="n"><a class="reference internal" href="trinity.common.html#trinity.common.config.Config" title="trinity.common.config.Config"><span class="pre">Config</span></a></span></em><span class="sig-paren">)</span> <span class="sig-return"><span class="sig-return-icon">&#x2192;</span> <span class="sig-return-typehint"><span class="pre">Tuple</span><span class="p"><span class="pre">[</span></span><span class="pre">List</span><span class="p"><span class="pre">[</span></span><a class="reference internal" href="#trinity.common.models.model.InferenceModel" title="trinity.common.models.model.InferenceModel"><span class="pre">InferenceModel</span></a><span class="p"><span class="pre">]</span></span><span class="p"><span class="pre">,</span></span><span class="w"> </span><span class="pre">List</span><span class="p"><span class="pre">[</span></span><span class="pre">List</span><span class="p"><span class="pre">[</span></span><a class="reference internal" href="#trinity.common.models.model.InferenceModel" title="trinity.common.models.model.InferenceModel"><span class="pre">InferenceModel</span></a><span class="p"><span class="pre">]</span></span><span class="p"><span class="pre">]</span></span><span class="p"><span class="pre">]</span></span></span></span><a class="reference internal" href="../_modules/trinity/common/models.html#create_inference_models"><span class="viewcode-link"><span class="pre">[source]</span></span></a><a class="headerlink" href="#trinity.common.models.create_inference_models" title="Link to this definition"></a></dt>
 <dd><p>Create <cite>engine_num</cite> rollout models.</p>
 <p>Each model has <cite>tensor_parallel_size</cite> workers.</p>
 </dd></dl>