1 files changed, 23 insertions, 2 deletions
diff --git a/lib/api/helpers/runner.rb b/lib/api/helpers/runner.rb
index 6f25cf507bc..9ec9b5e1e35 100644
--- a/lib/api/helpers/runner.rb
+++ b/lib/api/helpers/runner.rb
@@ -25,6 +25,7 @@ module API
         return get_runner_ip unless params['info'].present?
 
         attributes_for_keys(%w(name version revision platform architecture), params['info'])
+          .merge(get_runner_config_from_request)
           .merge(get_runner_ip)
       end
 
@@ -33,8 +34,15 @@ module API
       end
 
       def current_runner
+        token = params[:token]
+
+        if token
+          ::Gitlab::Database::LoadBalancing::RackMiddleware
+            .stick_or_unstick(env, :runner, token)
+        end
+
         strong_memoize(:current_runner) do
-          ::Ci::Runner.find_by_token(params[:token].to_s)
+          ::Ci::Runner.find_by_token(token.to_s)
         end
       end
 
@@ -64,8 +72,15 @@ module API
       end
 
       def current_job
+        id = params[:id]
+
+        if id
+          ::Gitlab::Database::LoadBalancing::RackMiddleware
+            .stick_or_unstick(env, :build, id)
+        end
+
         strong_memoize(:current_job) do
-          ::Ci::Build.find_by_id(params[:id])
+          ::Ci::Build.find_by_id(id)
         end
       end
 
@@ -91,6 +106,12 @@ module API
       def track_ci_minutes_usage!(_build, _runner)
         # noop: overridden in EE
       end
+
+      private
+
+      def get_runner_config_from_request
+        { config: attributes_for_keys(%w(gpus), params.dig('info', 'config')) }
+      end
     end
   end
 end