Add mrjob libraries in yql operation specs

commit_hash:91a36559a0c3bb85b06e327a2cc695639a6cf81d

3 files changed, 60 insertions, 1 deletions

diff --git a/yql/essentials/providers/common/proto/gateways_config.proto b/yql/essentials/providers/common/proto/gateways_config.proto
index 977ab2e7b9e..bd813bade95 100644
--- a/yql/essentials/providers/common/proto/gateways_config.proto
+++ b/yql/essentials/providers/common/proto/gateways_config.proto
@@ -103,6 +103,11 @@ message TYtClusterConfig {
     repeated TAttr Settings = 100;
 }
 
+message TFileWithMd5 {
+    optional string File = 1;
+    optional string Md5 = 2;
+}
+
 message TYtGatewayConfig {
     optional uint32 GatewayThreads = 1 [default = 0]; // Number of gateway MtpQueue threads. The adaptive MtpQueue will be used for 0 value
     optional EYtLogLevel YtLogLevel = 2 [default = YL_ERROR]; // The verbosity level of YT log
@@ -117,6 +122,7 @@ message TYtGatewayConfig {
     optional uint64 YtDebugLogSize = 11 [default = 0]; // Max YT debug log size
     optional bool YtDebugLogAlwaysWrite = 12 [default = false]; // Write YT debug log always or on error only
     optional string LocalChainFile = 14; // File to dump table contet in test mode
+    repeated TFileWithMd5 MrJobSystemLibsWithMd5 = 15; // Paths to linux dynamic libraries required for running mrjob.
 
     repeated TRemoteFilePattern RemoteFilePatterns = 100;
     repeated TYtClusterConfig ClusterMapping = 101; // Accepted cluster shortcuts. All names not from this list will be rejected
diff --git a/yt/yql/providers/yt/gateway/native/yql_yt_spec.cpp b/yt/yql/providers/yt/gateway/native/yql_yt_spec.cpp
index 5044596539b..b2bde43a28a 100644
--- a/yt/yql/providers/yt/gateway/native/yql_yt_spec.cpp
+++ b/yt/yql/providers/yt/gateway/native/yql_yt_spec.cpp
@@ -558,7 +558,34 @@ void FillUserJobSpecImpl(NYT::TUserJobSpec& spec,
         }
     }
 
+    TVector<std::pair<TString, TString>> mrJobSystemLibs;
+    if (execCtx.Config_->MrJobSystemLibsWithMd5Size() > 0) {
+        mrJobSystemLibs.reserve(execCtx.Config_->MrJobSystemLibsWithMd5Size());
+
+        for (const auto& systemLib : execCtx.Config_->GetMrJobSystemLibsWithMd5()) {
+            mrJobSystemLibs.push_back({systemLib.GetFile(), systemLib.GetMd5()});
+
+            const auto libSize = TFileStat(systemLib.GetFile()).Size;
+            YQL_ENSURE(libSize != 0);
+            fileMemUsage += libSize;
+        }
+
+        spec.AddEnvironment("LD_LIBRARY_PATH", ".");
+    }
+
+
     if (!localRun) {
+        for (size_t i = 0; i < mrJobSystemLibs.size(); i++) {
+            if (!mrJobSystemLibs[i].second) {
+                if (GetEnv("YQL_LOCAL") == "1") {
+                    // do not calculate heavy md5 in local mode (YQL-15353)
+                    mrJobSystemLibs[i].second = MD5::Calc(mrJobSystemLibs[i].first);
+                } else {
+                    mrJobSystemLibs[i].second = MD5::File(mrJobSystemLibs[i].first);
+                }
+            }
+        }
+
         if (mrJobBin.empty()) {
             mrJobBinMd5 = GetPersistentExecPathMd5();
         } else if (!mrJobBinMd5) {
@@ -594,15 +621,42 @@ void FillUserJobSpecImpl(NYT::TUserJobSpec& spec,
             } else if (!mrJobBin.empty()) {
                 spec.JobBinaryLocalPath(mrJobBin, mrJobBinMd5);
             }
+        }
 
+        for (size_t i = 0; i < mrJobSystemLibs.size(); i++) {
+            bool useBinCache = false;
+            if (binCacheFolder) {
+                if (auto snapshot = entry->GetBinarySnapshotFromCache(binCacheFolder, mrJobSystemLibs[i].second, mrJobSystemLibs[i].first)) {
+                    spec.AddFile(snapshot->first);
+                    useBinCache = true;
+                }
+            }
+            if (!useBinCache) {
+                if (binTmpFolder) {
+                    const TDuration binExpiration = settings->BinaryExpirationInterval.Get().GetOrElse(TDuration());
+                    auto libSnapshot = entry->GetBinarySnapshot(binTmpFolder, mrJobSystemLibs[i].second, mrJobSystemLibs[i].first, binExpiration);
+                    spec.AddFile(libSnapshot.first);
+                } else {
+                    NYT::TAddLocalFileOptions opts;
+                    opts.MD5CheckSum(mrJobSystemLibs[i].second);
+                    spec.AddLocalFile(mrJobSystemLibs[i].first, opts);
+                }
+            }
         }
+
     }
     else if (!mrJobBin.empty()) {
         const auto binSize = TFileStat(mrJobBin).Size;
         YQL_ENSURE(binSize != 0);
         spec.JobBinaryLocalPath(mrJobBin, mrJobBinMd5);
+        for (auto file : mrJobSystemLibs) {
+            NYT::TAddLocalFileOptions opts;
+            opts.MD5CheckSum(file.second);
+            spec.AddLocalFile(file.first, opts);
+        }
         fileMemUsage += binSize;
     }
+
     auto defaultMemoryLimit = settings->DefaultMemoryLimit.Get(cluster).GetOrElse(0);
     ui64 tmpFsSize = settings->UseTmpfs.Get(cluster).GetOrElse(false)
         ? (ui64)settings->ExtraTmpfsSize.Get(cluster).GetOrElse(8_MB)
diff --git a/yt/yql/providers/yt/gateway/native/yql_yt_spec.h b/yt/yql/providers/yt/gateway/native/yql_yt_spec.h
index c40caf9a39d..3cca737bfdb 100644
--- a/yt/yql/providers/yt/gateway/native/yql_yt_spec.h
+++ b/yt/yql/providers/yt/gateway/native/yql_yt_spec.h
@@ -97,7 +97,6 @@ inline void FillOperationSpec(NYT::TUserOperationSpecBase<TDerived>& spec, const
     if (auto val = execCtx->Options_.Config()->CoreDumpPath.Get()) {
         spec.CoreTablePath(*val);
     }
-    
 }
 
 template <class TExecParamsPtr>