nvidia-nat 1.2.0__py3-none-any.whl

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
Files changed (435) hide show
  1. aiq/__init__.py +66 -0
  2. nat/agent/__init__.py +0 -0
  3. nat/agent/base.py +256 -0
  4. nat/agent/dual_node.py +67 -0
  5. nat/agent/react_agent/__init__.py +0 -0
  6. nat/agent/react_agent/agent.py +363 -0
  7. nat/agent/react_agent/output_parser.py +104 -0
  8. nat/agent/react_agent/prompt.py +44 -0
  9. nat/agent/react_agent/register.py +149 -0
  10. nat/agent/reasoning_agent/__init__.py +0 -0
  11. nat/agent/reasoning_agent/reasoning_agent.py +225 -0
  12. nat/agent/register.py +23 -0
  13. nat/agent/rewoo_agent/__init__.py +0 -0
  14. nat/agent/rewoo_agent/agent.py +415 -0
  15. nat/agent/rewoo_agent/prompt.py +110 -0
  16. nat/agent/rewoo_agent/register.py +157 -0
  17. nat/agent/tool_calling_agent/__init__.py +0 -0
  18. nat/agent/tool_calling_agent/agent.py +119 -0
  19. nat/agent/tool_calling_agent/register.py +106 -0
  20. nat/authentication/__init__.py +14 -0
  21. nat/authentication/api_key/__init__.py +14 -0
  22. nat/authentication/api_key/api_key_auth_provider.py +96 -0
  23. nat/authentication/api_key/api_key_auth_provider_config.py +124 -0
  24. nat/authentication/api_key/register.py +26 -0
  25. nat/authentication/exceptions/__init__.py +14 -0
  26. nat/authentication/exceptions/api_key_exceptions.py +38 -0
  27. nat/authentication/http_basic_auth/__init__.py +0 -0
  28. nat/authentication/http_basic_auth/http_basic_auth_provider.py +81 -0
  29. nat/authentication/http_basic_auth/register.py +30 -0
  30. nat/authentication/interfaces.py +93 -0
  31. nat/authentication/oauth2/__init__.py +14 -0
  32. nat/authentication/oauth2/oauth2_auth_code_flow_provider.py +107 -0
  33. nat/authentication/oauth2/oauth2_auth_code_flow_provider_config.py +39 -0
  34. nat/authentication/oauth2/register.py +25 -0
  35. nat/authentication/register.py +21 -0
  36. nat/builder/__init__.py +0 -0
  37. nat/builder/builder.py +285 -0
  38. nat/builder/component_utils.py +316 -0
  39. nat/builder/context.py +270 -0
  40. nat/builder/embedder.py +24 -0
  41. nat/builder/eval_builder.py +161 -0
  42. nat/builder/evaluator.py +29 -0
  43. nat/builder/framework_enum.py +24 -0
  44. nat/builder/front_end.py +73 -0
  45. nat/builder/function.py +344 -0
  46. nat/builder/function_base.py +380 -0
  47. nat/builder/function_info.py +627 -0
  48. nat/builder/intermediate_step_manager.py +174 -0
  49. nat/builder/llm.py +25 -0
  50. nat/builder/retriever.py +25 -0
  51. nat/builder/user_interaction_manager.py +78 -0
  52. nat/builder/workflow.py +148 -0
  53. nat/builder/workflow_builder.py +1117 -0
  54. nat/cli/__init__.py +14 -0
  55. nat/cli/cli_utils/__init__.py +0 -0
  56. nat/cli/cli_utils/config_override.py +231 -0
  57. nat/cli/cli_utils/validation.py +37 -0
  58. nat/cli/commands/__init__.py +0 -0
  59. nat/cli/commands/configure/__init__.py +0 -0
  60. nat/cli/commands/configure/channel/__init__.py +0 -0
  61. nat/cli/commands/configure/channel/add.py +28 -0
  62. nat/cli/commands/configure/channel/channel.py +34 -0
  63. nat/cli/commands/configure/channel/remove.py +30 -0
  64. nat/cli/commands/configure/channel/update.py +30 -0
  65. nat/cli/commands/configure/configure.py +33 -0
  66. nat/cli/commands/evaluate.py +139 -0
  67. nat/cli/commands/info/__init__.py +14 -0
  68. nat/cli/commands/info/info.py +37 -0
  69. nat/cli/commands/info/list_channels.py +32 -0
  70. nat/cli/commands/info/list_components.py +129 -0
  71. nat/cli/commands/info/list_mcp.py +304 -0
  72. nat/cli/commands/registry/__init__.py +14 -0
  73. nat/cli/commands/registry/publish.py +88 -0
  74. nat/cli/commands/registry/pull.py +118 -0
  75. nat/cli/commands/registry/registry.py +36 -0
  76. nat/cli/commands/registry/remove.py +108 -0
  77. nat/cli/commands/registry/search.py +155 -0
  78. nat/cli/commands/sizing/__init__.py +14 -0
  79. nat/cli/commands/sizing/calc.py +297 -0
  80. nat/cli/commands/sizing/sizing.py +27 -0
  81. nat/cli/commands/start.py +246 -0
  82. nat/cli/commands/uninstall.py +81 -0
  83. nat/cli/commands/validate.py +47 -0
  84. nat/cli/commands/workflow/__init__.py +14 -0
  85. nat/cli/commands/workflow/templates/__init__.py.j2 +0 -0
  86. nat/cli/commands/workflow/templates/config.yml.j2 +16 -0
  87. nat/cli/commands/workflow/templates/pyproject.toml.j2 +22 -0
  88. nat/cli/commands/workflow/templates/register.py.j2 +5 -0
  89. nat/cli/commands/workflow/templates/workflow.py.j2 +36 -0
  90. nat/cli/commands/workflow/workflow.py +37 -0
  91. nat/cli/commands/workflow/workflow_commands.py +317 -0
  92. nat/cli/entrypoint.py +135 -0
  93. nat/cli/main.py +57 -0
  94. nat/cli/register_workflow.py +488 -0
  95. nat/cli/type_registry.py +1000 -0
  96. nat/data_models/__init__.py +14 -0
  97. nat/data_models/api_server.py +716 -0
  98. nat/data_models/authentication.py +231 -0
  99. nat/data_models/common.py +171 -0
  100. nat/data_models/component.py +58 -0
  101. nat/data_models/component_ref.py +168 -0
  102. nat/data_models/config.py +410 -0
  103. nat/data_models/dataset_handler.py +169 -0
  104. nat/data_models/discovery_metadata.py +305 -0
  105. nat/data_models/embedder.py +27 -0
  106. nat/data_models/evaluate.py +127 -0
  107. nat/data_models/evaluator.py +26 -0
  108. nat/data_models/front_end.py +26 -0
  109. nat/data_models/function.py +30 -0
  110. nat/data_models/function_dependencies.py +72 -0
  111. nat/data_models/interactive.py +246 -0
  112. nat/data_models/intermediate_step.py +302 -0
  113. nat/data_models/invocation_node.py +38 -0
  114. nat/data_models/llm.py +27 -0
  115. nat/data_models/logging.py +26 -0
  116. nat/data_models/memory.py +27 -0
  117. nat/data_models/object_store.py +44 -0
  118. nat/data_models/profiler.py +54 -0
  119. nat/data_models/registry_handler.py +26 -0
  120. nat/data_models/retriever.py +30 -0
  121. nat/data_models/retry_mixin.py +35 -0
  122. nat/data_models/span.py +190 -0
  123. nat/data_models/step_adaptor.py +64 -0
  124. nat/data_models/streaming.py +33 -0
  125. nat/data_models/swe_bench_model.py +54 -0
  126. nat/data_models/telemetry_exporter.py +26 -0
  127. nat/data_models/ttc_strategy.py +30 -0
  128. nat/embedder/__init__.py +0 -0
  129. nat/embedder/nim_embedder.py +59 -0
  130. nat/embedder/openai_embedder.py +43 -0
  131. nat/embedder/register.py +22 -0
  132. nat/eval/__init__.py +14 -0
  133. nat/eval/config.py +60 -0
  134. nat/eval/dataset_handler/__init__.py +0 -0
  135. nat/eval/dataset_handler/dataset_downloader.py +106 -0
  136. nat/eval/dataset_handler/dataset_filter.py +52 -0
  137. nat/eval/dataset_handler/dataset_handler.py +367 -0
  138. nat/eval/evaluate.py +510 -0
  139. nat/eval/evaluator/__init__.py +14 -0
  140. nat/eval/evaluator/base_evaluator.py +77 -0
  141. nat/eval/evaluator/evaluator_model.py +45 -0
  142. nat/eval/intermediate_step_adapter.py +99 -0
  143. nat/eval/rag_evaluator/__init__.py +0 -0
  144. nat/eval/rag_evaluator/evaluate.py +178 -0
  145. nat/eval/rag_evaluator/register.py +143 -0
  146. nat/eval/register.py +23 -0
  147. nat/eval/remote_workflow.py +133 -0
  148. nat/eval/runners/__init__.py +14 -0
  149. nat/eval/runners/config.py +39 -0
  150. nat/eval/runners/multi_eval_runner.py +54 -0
  151. nat/eval/runtime_event_subscriber.py +52 -0
  152. nat/eval/swe_bench_evaluator/__init__.py +0 -0
  153. nat/eval/swe_bench_evaluator/evaluate.py +215 -0
  154. nat/eval/swe_bench_evaluator/register.py +36 -0
  155. nat/eval/trajectory_evaluator/__init__.py +0 -0
  156. nat/eval/trajectory_evaluator/evaluate.py +75 -0
  157. nat/eval/trajectory_evaluator/register.py +40 -0
  158. nat/eval/tunable_rag_evaluator/__init__.py +0 -0
  159. nat/eval/tunable_rag_evaluator/evaluate.py +245 -0
  160. nat/eval/tunable_rag_evaluator/register.py +52 -0
  161. nat/eval/usage_stats.py +41 -0
  162. nat/eval/utils/__init__.py +0 -0
  163. nat/eval/utils/output_uploader.py +140 -0
  164. nat/eval/utils/tqdm_position_registry.py +40 -0
  165. nat/eval/utils/weave_eval.py +184 -0
  166. nat/experimental/__init__.py +0 -0
  167. nat/experimental/decorators/__init__.py +0 -0
  168. nat/experimental/decorators/experimental_warning_decorator.py +134 -0
  169. nat/experimental/test_time_compute/__init__.py +0 -0
  170. nat/experimental/test_time_compute/editing/__init__.py +0 -0
  171. nat/experimental/test_time_compute/editing/iterative_plan_refinement_editor.py +147 -0
  172. nat/experimental/test_time_compute/editing/llm_as_a_judge_editor.py +204 -0
  173. nat/experimental/test_time_compute/editing/motivation_aware_summarization.py +107 -0
  174. nat/experimental/test_time_compute/functions/__init__.py +0 -0
  175. nat/experimental/test_time_compute/functions/execute_score_select_function.py +105 -0
  176. nat/experimental/test_time_compute/functions/plan_select_execute_function.py +224 -0
  177. nat/experimental/test_time_compute/functions/ttc_tool_orchestration_function.py +205 -0
  178. nat/experimental/test_time_compute/functions/ttc_tool_wrapper_function.py +146 -0
  179. nat/experimental/test_time_compute/models/__init__.py +0 -0
  180. nat/experimental/test_time_compute/models/editor_config.py +132 -0
  181. nat/experimental/test_time_compute/models/scoring_config.py +112 -0
  182. nat/experimental/test_time_compute/models/search_config.py +120 -0
  183. nat/experimental/test_time_compute/models/selection_config.py +154 -0
  184. nat/experimental/test_time_compute/models/stage_enums.py +43 -0
  185. nat/experimental/test_time_compute/models/strategy_base.py +66 -0
  186. nat/experimental/test_time_compute/models/tool_use_config.py +41 -0
  187. nat/experimental/test_time_compute/models/ttc_item.py +48 -0
  188. nat/experimental/test_time_compute/register.py +36 -0
  189. nat/experimental/test_time_compute/scoring/__init__.py +0 -0
  190. nat/experimental/test_time_compute/scoring/llm_based_agent_scorer.py +168 -0
  191. nat/experimental/test_time_compute/scoring/llm_based_plan_scorer.py +168 -0
  192. nat/experimental/test_time_compute/scoring/motivation_aware_scorer.py +111 -0
  193. nat/experimental/test_time_compute/search/__init__.py +0 -0
  194. nat/experimental/test_time_compute/search/multi_llm_planner.py +128 -0
  195. nat/experimental/test_time_compute/search/multi_query_retrieval_search.py +122 -0
  196. nat/experimental/test_time_compute/search/single_shot_multi_plan_planner.py +128 -0
  197. nat/experimental/test_time_compute/selection/__init__.py +0 -0
  198. nat/experimental/test_time_compute/selection/best_of_n_selector.py +63 -0
  199. nat/experimental/test_time_compute/selection/llm_based_agent_output_selector.py +131 -0
  200. nat/experimental/test_time_compute/selection/llm_based_output_merging_selector.py +159 -0
  201. nat/experimental/test_time_compute/selection/llm_based_plan_selector.py +128 -0
  202. nat/experimental/test_time_compute/selection/threshold_selector.py +58 -0
  203. nat/front_ends/__init__.py +14 -0
  204. nat/front_ends/console/__init__.py +14 -0
  205. nat/front_ends/console/authentication_flow_handler.py +233 -0
  206. nat/front_ends/console/console_front_end_config.py +32 -0
  207. nat/front_ends/console/console_front_end_plugin.py +96 -0
  208. nat/front_ends/console/register.py +25 -0
  209. nat/front_ends/cron/__init__.py +14 -0
  210. nat/front_ends/fastapi/__init__.py +14 -0
  211. nat/front_ends/fastapi/auth_flow_handlers/__init__.py +0 -0
  212. nat/front_ends/fastapi/auth_flow_handlers/http_flow_handler.py +27 -0
  213. nat/front_ends/fastapi/auth_flow_handlers/websocket_flow_handler.py +107 -0
  214. nat/front_ends/fastapi/fastapi_front_end_config.py +241 -0
  215. nat/front_ends/fastapi/fastapi_front_end_controller.py +68 -0
  216. nat/front_ends/fastapi/fastapi_front_end_plugin.py +116 -0
  217. nat/front_ends/fastapi/fastapi_front_end_plugin_worker.py +1087 -0
  218. nat/front_ends/fastapi/html_snippets/__init__.py +14 -0
  219. nat/front_ends/fastapi/html_snippets/auth_code_grant_success.py +35 -0
  220. nat/front_ends/fastapi/intermediate_steps_subscriber.py +80 -0
  221. nat/front_ends/fastapi/job_store.py +183 -0
  222. nat/front_ends/fastapi/main.py +72 -0
  223. nat/front_ends/fastapi/message_handler.py +320 -0
  224. nat/front_ends/fastapi/message_validator.py +352 -0
  225. nat/front_ends/fastapi/register.py +25 -0
  226. nat/front_ends/fastapi/response_helpers.py +195 -0
  227. nat/front_ends/fastapi/step_adaptor.py +319 -0
  228. nat/front_ends/mcp/__init__.py +14 -0
  229. nat/front_ends/mcp/mcp_front_end_config.py +36 -0
  230. nat/front_ends/mcp/mcp_front_end_plugin.py +81 -0
  231. nat/front_ends/mcp/mcp_front_end_plugin_worker.py +143 -0
  232. nat/front_ends/mcp/register.py +27 -0
  233. nat/front_ends/mcp/tool_converter.py +241 -0
  234. nat/front_ends/register.py +22 -0
  235. nat/front_ends/simple_base/__init__.py +14 -0
  236. nat/front_ends/simple_base/simple_front_end_plugin_base.py +54 -0
  237. nat/llm/__init__.py +0 -0
  238. nat/llm/aws_bedrock_llm.py +57 -0
  239. nat/llm/nim_llm.py +46 -0
  240. nat/llm/openai_llm.py +46 -0
  241. nat/llm/register.py +23 -0
  242. nat/llm/utils/__init__.py +14 -0
  243. nat/llm/utils/env_config_value.py +94 -0
  244. nat/llm/utils/error.py +17 -0
  245. nat/memory/__init__.py +20 -0
  246. nat/memory/interfaces.py +183 -0
  247. nat/memory/models.py +112 -0
  248. nat/meta/pypi.md +58 -0
  249. nat/object_store/__init__.py +20 -0
  250. nat/object_store/in_memory_object_store.py +76 -0
  251. nat/object_store/interfaces.py +84 -0
  252. nat/object_store/models.py +38 -0
  253. nat/object_store/register.py +20 -0
  254. nat/observability/__init__.py +14 -0
  255. nat/observability/exporter/__init__.py +14 -0
  256. nat/observability/exporter/base_exporter.py +449 -0
  257. nat/observability/exporter/exporter.py +78 -0
  258. nat/observability/exporter/file_exporter.py +33 -0
  259. nat/observability/exporter/processing_exporter.py +322 -0
  260. nat/observability/exporter/raw_exporter.py +52 -0
  261. nat/observability/exporter/span_exporter.py +288 -0
  262. nat/observability/exporter_manager.py +335 -0
  263. nat/observability/mixin/__init__.py +14 -0
  264. nat/observability/mixin/batch_config_mixin.py +26 -0
  265. nat/observability/mixin/collector_config_mixin.py +23 -0
  266. nat/observability/mixin/file_mixin.py +288 -0
  267. nat/observability/mixin/file_mode.py +23 -0
  268. nat/observability/mixin/resource_conflict_mixin.py +134 -0
  269. nat/observability/mixin/serialize_mixin.py +61 -0
  270. nat/observability/mixin/type_introspection_mixin.py +183 -0
  271. nat/observability/processor/__init__.py +14 -0
  272. nat/observability/processor/batching_processor.py +310 -0
  273. nat/observability/processor/callback_processor.py +42 -0
  274. nat/observability/processor/intermediate_step_serializer.py +28 -0
  275. nat/observability/processor/processor.py +71 -0
  276. nat/observability/register.py +96 -0
  277. nat/observability/utils/__init__.py +14 -0
  278. nat/observability/utils/dict_utils.py +236 -0
  279. nat/observability/utils/time_utils.py +31 -0
  280. nat/plugins/.namespace +1 -0
  281. nat/profiler/__init__.py +0 -0
  282. nat/profiler/calc/__init__.py +14 -0
  283. nat/profiler/calc/calc_runner.py +627 -0
  284. nat/profiler/calc/calculations.py +288 -0
  285. nat/profiler/calc/data_models.py +188 -0
  286. nat/profiler/calc/plot.py +345 -0
  287. nat/profiler/callbacks/__init__.py +0 -0
  288. nat/profiler/callbacks/agno_callback_handler.py +295 -0
  289. nat/profiler/callbacks/base_callback_class.py +20 -0
  290. nat/profiler/callbacks/langchain_callback_handler.py +290 -0
  291. nat/profiler/callbacks/llama_index_callback_handler.py +205 -0
  292. nat/profiler/callbacks/semantic_kernel_callback_handler.py +238 -0
  293. nat/profiler/callbacks/token_usage_base_model.py +27 -0
  294. nat/profiler/data_frame_row.py +51 -0
  295. nat/profiler/data_models.py +24 -0
  296. nat/profiler/decorators/__init__.py +0 -0
  297. nat/profiler/decorators/framework_wrapper.py +131 -0
  298. nat/profiler/decorators/function_tracking.py +254 -0
  299. nat/profiler/forecasting/__init__.py +0 -0
  300. nat/profiler/forecasting/config.py +18 -0
  301. nat/profiler/forecasting/model_trainer.py +75 -0
  302. nat/profiler/forecasting/models/__init__.py +22 -0
  303. nat/profiler/forecasting/models/forecasting_base_model.py +40 -0
  304. nat/profiler/forecasting/models/linear_model.py +197 -0
  305. nat/profiler/forecasting/models/random_forest_regressor.py +269 -0
  306. nat/profiler/inference_metrics_model.py +28 -0
  307. nat/profiler/inference_optimization/__init__.py +0 -0
  308. nat/profiler/inference_optimization/bottleneck_analysis/__init__.py +0 -0
  309. nat/profiler/inference_optimization/bottleneck_analysis/nested_stack_analysis.py +460 -0
  310. nat/profiler/inference_optimization/bottleneck_analysis/simple_stack_analysis.py +258 -0
  311. nat/profiler/inference_optimization/data_models.py +386 -0
  312. nat/profiler/inference_optimization/experimental/__init__.py +0 -0
  313. nat/profiler/inference_optimization/experimental/concurrency_spike_analysis.py +468 -0
  314. nat/profiler/inference_optimization/experimental/prefix_span_analysis.py +405 -0
  315. nat/profiler/inference_optimization/llm_metrics.py +212 -0
  316. nat/profiler/inference_optimization/prompt_caching.py +163 -0
  317. nat/profiler/inference_optimization/token_uniqueness.py +107 -0
  318. nat/profiler/inference_optimization/workflow_runtimes.py +72 -0
  319. nat/profiler/intermediate_property_adapter.py +102 -0
  320. nat/profiler/profile_runner.py +473 -0
  321. nat/profiler/utils.py +184 -0
  322. nat/registry_handlers/__init__.py +0 -0
  323. nat/registry_handlers/local/__init__.py +0 -0
  324. nat/registry_handlers/local/local_handler.py +176 -0
  325. nat/registry_handlers/local/register_local.py +37 -0
  326. nat/registry_handlers/metadata_factory.py +60 -0
  327. nat/registry_handlers/package_utils.py +571 -0
  328. nat/registry_handlers/pypi/__init__.py +0 -0
  329. nat/registry_handlers/pypi/pypi_handler.py +251 -0
  330. nat/registry_handlers/pypi/register_pypi.py +40 -0
  331. nat/registry_handlers/register.py +21 -0
  332. nat/registry_handlers/registry_handler_base.py +157 -0
  333. nat/registry_handlers/rest/__init__.py +0 -0
  334. nat/registry_handlers/rest/register_rest.py +56 -0
  335. nat/registry_handlers/rest/rest_handler.py +237 -0
  336. nat/registry_handlers/schemas/__init__.py +0 -0
  337. nat/registry_handlers/schemas/headers.py +42 -0
  338. nat/registry_handlers/schemas/package.py +68 -0
  339. nat/registry_handlers/schemas/publish.py +68 -0
  340. nat/registry_handlers/schemas/pull.py +82 -0
  341. nat/registry_handlers/schemas/remove.py +36 -0
  342. nat/registry_handlers/schemas/search.py +91 -0
  343. nat/registry_handlers/schemas/status.py +47 -0
  344. nat/retriever/__init__.py +0 -0
  345. nat/retriever/interface.py +41 -0
  346. nat/retriever/milvus/__init__.py +14 -0
  347. nat/retriever/milvus/register.py +81 -0
  348. nat/retriever/milvus/retriever.py +228 -0
  349. nat/retriever/models.py +77 -0
  350. nat/retriever/nemo_retriever/__init__.py +14 -0
  351. nat/retriever/nemo_retriever/register.py +60 -0
  352. nat/retriever/nemo_retriever/retriever.py +190 -0
  353. nat/retriever/register.py +22 -0
  354. nat/runtime/__init__.py +14 -0
  355. nat/runtime/loader.py +220 -0
  356. nat/runtime/runner.py +195 -0
  357. nat/runtime/session.py +162 -0
  358. nat/runtime/user_metadata.py +130 -0
  359. nat/settings/__init__.py +0 -0
  360. nat/settings/global_settings.py +318 -0
  361. nat/test/.namespace +1 -0
  362. nat/tool/__init__.py +0 -0
  363. nat/tool/chat_completion.py +74 -0
  364. nat/tool/code_execution/README.md +151 -0
  365. nat/tool/code_execution/__init__.py +0 -0
  366. nat/tool/code_execution/code_sandbox.py +267 -0
  367. nat/tool/code_execution/local_sandbox/.gitignore +1 -0
  368. nat/tool/code_execution/local_sandbox/Dockerfile.sandbox +60 -0
  369. nat/tool/code_execution/local_sandbox/__init__.py +13 -0
  370. nat/tool/code_execution/local_sandbox/local_sandbox_server.py +198 -0
  371. nat/tool/code_execution/local_sandbox/sandbox.requirements.txt +6 -0
  372. nat/tool/code_execution/local_sandbox/start_local_sandbox.sh +50 -0
  373. nat/tool/code_execution/register.py +74 -0
  374. nat/tool/code_execution/test_code_execution_sandbox.py +414 -0
  375. nat/tool/code_execution/utils.py +100 -0
  376. nat/tool/datetime_tools.py +42 -0
  377. nat/tool/document_search.py +141 -0
  378. nat/tool/github_tools/__init__.py +0 -0
  379. nat/tool/github_tools/create_github_commit.py +133 -0
  380. nat/tool/github_tools/create_github_issue.py +87 -0
  381. nat/tool/github_tools/create_github_pr.py +106 -0
  382. nat/tool/github_tools/get_github_file.py +106 -0
  383. nat/tool/github_tools/get_github_issue.py +166 -0
  384. nat/tool/github_tools/get_github_pr.py +256 -0
  385. nat/tool/github_tools/update_github_issue.py +100 -0
  386. nat/tool/mcp/__init__.py +14 -0
  387. nat/tool/mcp/exceptions.py +142 -0
  388. nat/tool/mcp/mcp_client.py +255 -0
  389. nat/tool/mcp/mcp_tool.py +96 -0
  390. nat/tool/memory_tools/__init__.py +0 -0
  391. nat/tool/memory_tools/add_memory_tool.py +79 -0
  392. nat/tool/memory_tools/delete_memory_tool.py +67 -0
  393. nat/tool/memory_tools/get_memory_tool.py +72 -0
  394. nat/tool/nvidia_rag.py +95 -0
  395. nat/tool/register.py +38 -0
  396. nat/tool/retriever.py +94 -0
  397. nat/tool/server_tools.py +66 -0
  398. nat/utils/__init__.py +0 -0
  399. nat/utils/data_models/__init__.py +0 -0
  400. nat/utils/data_models/schema_validator.py +58 -0
  401. nat/utils/debugging_utils.py +43 -0
  402. nat/utils/dump_distro_mapping.py +32 -0
  403. nat/utils/exception_handlers/__init__.py +0 -0
  404. nat/utils/exception_handlers/automatic_retries.py +289 -0
  405. nat/utils/exception_handlers/mcp.py +211 -0
  406. nat/utils/exception_handlers/schemas.py +114 -0
  407. nat/utils/io/__init__.py +0 -0
  408. nat/utils/io/model_processing.py +28 -0
  409. nat/utils/io/yaml_tools.py +119 -0
  410. nat/utils/log_utils.py +37 -0
  411. nat/utils/metadata_utils.py +74 -0
  412. nat/utils/optional_imports.py +142 -0
  413. nat/utils/producer_consumer_queue.py +178 -0
  414. nat/utils/reactive/__init__.py +0 -0
  415. nat/utils/reactive/base/__init__.py +0 -0
  416. nat/utils/reactive/base/observable_base.py +65 -0
  417. nat/utils/reactive/base/observer_base.py +55 -0
  418. nat/utils/reactive/base/subject_base.py +79 -0
  419. nat/utils/reactive/observable.py +59 -0
  420. nat/utils/reactive/observer.py +76 -0
  421. nat/utils/reactive/subject.py +131 -0
  422. nat/utils/reactive/subscription.py +49 -0
  423. nat/utils/settings/__init__.py +0 -0
  424. nat/utils/settings/global_settings.py +197 -0
  425. nat/utils/string_utils.py +38 -0
  426. nat/utils/type_converter.py +290 -0
  427. nat/utils/type_utils.py +484 -0
  428. nat/utils/url_utils.py +27 -0
  429. nvidia_nat-1.2.0.dist-info/METADATA +365 -0
  430. nvidia_nat-1.2.0.dist-info/RECORD +435 -0
  431. nvidia_nat-1.2.0.dist-info/WHEEL +5 -0
  432. nvidia_nat-1.2.0.dist-info/entry_points.txt +21 -0
  433. nvidia_nat-1.2.0.dist-info/licenses/LICENSE-3rd-party.txt +5478 -0
  434. nvidia_nat-1.2.0.dist-info/licenses/LICENSE.md +201 -0
  435. nvidia_nat-1.2.0.dist-info/top_level.txt +2 -0
@@ -0,0 +1,288 @@
1
+ # SPDX-FileCopyrightText: Copyright (c) 2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
2
+ # SPDX-License-Identifier: Apache-2.0
3
+ #
4
+ # Licensed under the Apache License, Version 2.0 (the "License");
5
+ # you may not use this file except in compliance with the License.
6
+ # You may obtain a copy of the License at
7
+ #
8
+ # http://www.apache.org/licenses/LICENSE-2.0
9
+ #
10
+ # Unless required by applicable law or agreed to in writing, software
11
+ # distributed under the License is distributed on an "AS IS" BASIS,
12
+ # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13
+ # See the License for the specific language governing permissions and
14
+ # limitations under the License.
15
+
16
+ import logging
17
+
18
+ import numpy as np
19
+
20
+ from nat.profiler.calc.data_models import FitConfig
21
+ from nat.profiler.calc.data_models import GPUEstimates
22
+ from nat.profiler.calc.data_models import LinearFitResult
23
+
24
+ logger = logging.getLogger(__name__)
25
+
26
+
27
+ def compute_slope(concurrencies: list[float],
28
+ time_metrics: list[float],
29
+ fit_config: FitConfig | None = None) -> LinearFitResult:
30
+ """
31
+ Concurrency is the independent variable (x-axis) and time metric (which can be runtime or latency)
32
+ is the dependent variable (y-axis). This function computes the slope of the linear relationship
33
+ between concurrency and time metric.
34
+
35
+ Args:
36
+ concurrencies: List of concurrency values (x-axis)
37
+ time_metrics: List of time metric values (y-axis)
38
+ fit_config: Configuration for outlier detection and fit validation
39
+
40
+ Returns:
41
+ LinearFitResult containing slope, intercept, R-squared, and outliers removed
42
+
43
+ Raises:
44
+ ValueError: If the relationship is not linear (R² < min_r_squared)
45
+ """
46
+ # Use default config if none provided
47
+ if fit_config is None:
48
+ fit_config = FitConfig()
49
+
50
+ # Convert to numpy arrays for calculations
51
+ x = np.array(concurrencies)
52
+ y = np.array(time_metrics)
53
+
54
+ # Validate input
55
+ if len(x) != len(y):
56
+ raise ValueError("Concurrencies and time_metrics must have the same length")
57
+ if len(x) < 2:
58
+ raise ValueError("Need at least 2 points for linear regression")
59
+
60
+ outliers_removed = []
61
+
62
+ # Remove outliers if requested
63
+ if fit_config.remove_outliers and len(x) > 4: # Need at least 4 points for outlier detection
64
+ x_clean, y_clean, removed_concurrencies = _remove_outliers(x, y, fit_config)
65
+ x, y = x_clean, y_clean
66
+ outliers_removed = removed_concurrencies
67
+
68
+ # Calculate linear regression using least squares
69
+ n = len(x)
70
+ sum_x = x.sum()
71
+ sum_y = y.sum()
72
+ sum_xy = (x * y).sum()
73
+ sum_x2 = (x**2).sum()
74
+
75
+ # Calculate slope and intercept
76
+ slope = (n * sum_xy - sum_x * sum_y) / (n * sum_x2 - sum_x**2)
77
+ intercept = (sum_y - slope * sum_x) / n
78
+
79
+ # Calculate R-squared
80
+ y_pred = slope * x + intercept
81
+ ss_res = ((y - y_pred)**2).sum()
82
+ ss_tot = ((y - y.mean())**2).sum()
83
+ r_squared = 1 - (ss_res / ss_tot) if ss_tot != 0 else 0.0
84
+
85
+ # Validate linearity
86
+ if r_squared < fit_config.min_r_squared:
87
+ raise ValueError(f"Poor linear fit detected (R² = {r_squared:.3f} < {fit_config.min_r_squared}). "
88
+ f"The relationship may not be linear. Consider using non-linear regression.")
89
+
90
+ return LinearFitResult(slope=slope, intercept=intercept, r_squared=r_squared, outliers_removed=outliers_removed)
91
+
92
+
93
+ def _remove_outliers(x: np.ndarray, y: np.ndarray, fit_config: FitConfig) -> tuple[np.ndarray, np.ndarray, list[int]]:
94
+ """
95
+ Remove outliers using the Interquartile Range (IQR) method.
96
+ For small concurrency range (≤ threshold points), also checks raw y-values for extreme outliers.
97
+
98
+ Args:
99
+ x: Input x values (concurrencies)
100
+ y: Input y values (time metrics)
101
+ fit_config: Configuration for outlier detection
102
+
103
+ Returns:
104
+ Tuple of (cleaned_x, cleaned_y, list_of_removed_concurrencies)
105
+ """
106
+ # if the number of concurrency points is less removing outliers can be challenging
107
+ # as extreme outliers can skew the results.
108
+ # We use a threshold to check for extreme outliers in raw y-values first.
109
+ n = len(x)
110
+ all_removed_concurrencies = []
111
+
112
+ # For smaller concurrency ranges, check for extreme outliers in raw y-values first
113
+ if n <= fit_config.small_concurrency_range_threshold:
114
+ # Calculate IQR on raw y-values
115
+ y_q1 = np.percentile(y, 25)
116
+ y_q3 = np.percentile(y, 75)
117
+ y_iqr = y_q3 - y_q1
118
+
119
+ # Use a more aggressive threshold for small datasets
120
+ y_lower_bound = y_q1 - fit_config.extreme_outlier_threshold * y_iqr # More aggressive than 1.5
121
+ y_upper_bound = y_q3 + fit_config.extreme_outlier_threshold * y_iqr
122
+
123
+ # Find extreme outliers in raw values
124
+ extreme_outlier_mask = (y >= y_lower_bound) & (y <= y_upper_bound)
125
+ extreme_outliers_removed = np.sum(~extreme_outlier_mask)
126
+
127
+ if extreme_outliers_removed > 0:
128
+ extreme_removed_concurrencies = x[~extreme_outlier_mask].tolist()
129
+ all_removed_concurrencies.extend(extreme_removed_concurrencies)
130
+ logger.info("Removed %d extreme outliers from raw values: concurrencies %s",
131
+ extreme_outliers_removed,
132
+ extreme_removed_concurrencies)
133
+ # Continue with residual-based detection on the cleaned data
134
+ x = x[extreme_outlier_mask]
135
+ y = y[extreme_outlier_mask]
136
+ n = len(x)
137
+
138
+ # Standard residual-based outlier detection
139
+ # Calculate residuals from a simple linear fit
140
+ if n == 0:
141
+ raise ValueError("No data points remaining after outlier removal. Cannot compute linear fit.")
142
+
143
+ sum_x = x.sum()
144
+ sum_y = y.sum()
145
+ sum_xy = (x * y).sum()
146
+ sum_x2 = (x**2).sum()
147
+
148
+ slope = (n * sum_xy - sum_x * sum_y) / (n * sum_x2 - sum_x**2)
149
+ intercept = (sum_y - slope * sum_x) / n
150
+
151
+ # Calculate residuals
152
+ y_pred = slope * x + intercept
153
+ residuals = y - y_pred
154
+
155
+ # Use IQR method to detect outliers
156
+ q1 = np.percentile(residuals, 25)
157
+ q3 = np.percentile(residuals, 75)
158
+ iqr = q3 - q1
159
+
160
+ # Define outlier bounds (1.5 * IQR rule)
161
+ lower_bound = q1 - fit_config.conservative_outlier_threshold * iqr
162
+ upper_bound = q3 + fit_config.conservative_outlier_threshold * iqr
163
+
164
+ # Find non-outlier indices
165
+ non_outlier_mask = (residuals >= lower_bound) & (residuals <= upper_bound)
166
+
167
+ outliers_removed = np.sum(~non_outlier_mask)
168
+ residual_removed_concurrencies = x[~non_outlier_mask].tolist()
169
+ all_removed_concurrencies.extend(residual_removed_concurrencies)
170
+
171
+ # Add debugging for small datasets
172
+ if len(x) <= fit_config.small_concurrency_range_threshold:
173
+ logger.debug("Outlier detection for small dataset (n=%d):", len(x))
174
+ logger.debug(" Data points: %s", list(zip(x, y)))
175
+ logger.debug(" Residuals: %s", residuals.tolist())
176
+ logger.debug(" Q1=%.3f, Q3=%.3f, IQR=%.3f", q1, q3, iqr)
177
+ logger.debug(" Bounds: [%.3f, %.3f]", lower_bound, upper_bound)
178
+ logger.info(" Outliers removed: %d (concurrencies: %s)", outliers_removed, residual_removed_concurrencies)
179
+
180
+ return x[non_outlier_mask], y[non_outlier_mask], all_removed_concurrencies
181
+
182
+
183
+ def calc_gpu_estimate_based_on_slope(target_time_metric: float,
184
+ target_users: int,
185
+ test_gpu_count: int,
186
+ observed_slope: float,
187
+ observed_intercept: float = 0.0) -> float:
188
+ """
189
+ Calculate the GPU estimate based on the slope of the time metric.
190
+
191
+ This function uses the linear relationship between concurrency and time metrics
192
+ to estimate the required GPU count for a target user load.
193
+
194
+ Args:
195
+ target_time_metric: Target time metric (latency or runtime) in seconds
196
+ observed_slope: Slope from linear regression of time vs concurrency
197
+ target_users: Target number of concurrent users
198
+ test_gpu_count: Number of GPUs used in the test
199
+ observed_intercept: Y-intercept from linear regression (default: 0.0)
200
+
201
+ Returns:
202
+ Estimated number of GPUs required
203
+
204
+ Raises:
205
+ ValueError: If target_time_metric is less than or equal to intercept
206
+ """
207
+ if target_time_metric <= observed_intercept:
208
+ raise ValueError(f"Target time metric ({target_time_metric}) must be greater than "
209
+ f"the intercept ({observed_intercept}) for valid GPU estimation.")
210
+
211
+ # Calculate the concurrency that would achieve the target time metric
212
+ # Using the linear equation: time = slope * concurrency + intercept
213
+ # Solving for concurrency: concurrency = (time - intercept) / slope
214
+ calculated_concurrency = (target_time_metric - observed_intercept) / observed_slope
215
+ logger.info("Calculated concurrency: %f for target time metric: %f, observed intercept: %f, observed slope: %f",
216
+ calculated_concurrency,
217
+ target_time_metric,
218
+ observed_intercept,
219
+ observed_slope)
220
+
221
+ if calculated_concurrency <= 0:
222
+ raise ValueError(f"Calculated target concurrency ({calculated_concurrency}) is not positive. "
223
+ f"This suggests the slope or intercept values may be invalid.")
224
+
225
+ # Estimate GPUs using the ratio of target users to target concurrency
226
+ # scaled by the test GPU count
227
+ gpu_estimate = (target_users / calculated_concurrency) * test_gpu_count
228
+
229
+ return gpu_estimate
230
+
231
+
232
+ def calc_gpu_estimate_for_single_concurrency(target_llm_latency: float,
233
+ target_workflow_runtime: float,
234
+ target_users: int,
235
+ test_concurrency: int,
236
+ test_gpu_count: int,
237
+ observed_latency: float,
238
+ observed_runtime: float) -> GPUEstimates:
239
+ """
240
+ ROUGH ESTIMATE: Calculate GPU count estimate for a single concurrency level.
241
+
242
+ This is a simplified estimate that assumes linear scaling and should be used
243
+ as a baseline only. For more accurate estimates, use slope-based estimation
244
+ with multiple concurrency levels.
245
+
246
+ Formula based on the target latency:
247
+ G_required = (U_target / C_test) * (L_obs / L_target) * G_test
248
+
249
+ Formula based on the target runtime:
250
+ G_required = (U_target / C_test) * (R_obs / R_target) * G_test
251
+
252
+ where:
253
+ - U_target: Target number of users
254
+ - C_test: Test concurrency level
255
+ - L_obs: Observed LLM latency
256
+ - L_target: Target LLM latency
257
+ - R_obs: Observed workflow runtime
258
+ - R_target: Target workflow runtime
259
+ - G_test: Test GPU count
260
+
261
+ WARNING: This is a rough estimate that:
262
+ - Assumes perfect linear scaling (rarely true in practice)
263
+ - Doesn't account for GPU utilization inefficiencies
264
+ - May underestimate GPU requirements for high concurrency
265
+ - Should be validated against slope-based estimates
266
+ """
267
+ use_latency = target_llm_latency > 0
268
+ use_runtime = target_workflow_runtime > 0
269
+
270
+ # If observed latency or runtime exceeds the target, return empty estimates
271
+ if use_latency and observed_latency > target_llm_latency:
272
+ return GPUEstimates()
273
+
274
+ if use_runtime and observed_runtime > target_workflow_runtime:
275
+ return GPUEstimates()
276
+
277
+ # Calculate multipliers (how much faster we need to be)
278
+ llm_latency_multiplier = observed_latency / target_llm_latency if use_latency else 1.0
279
+ wf_runtime_multiplier = observed_runtime / target_workflow_runtime if use_runtime else 1.0
280
+
281
+ # Calculate GPU estimates using the corrected formula
282
+ gpu_estimate_by_wf_runtime = (target_users /
283
+ test_concurrency) * wf_runtime_multiplier * test_gpu_count if use_runtime else None
284
+ gpu_estimate_by_llm_latency = (target_users /
285
+ test_concurrency) * llm_latency_multiplier * test_gpu_count if use_latency else None
286
+
287
+ return GPUEstimates(gpu_estimate_by_wf_runtime=gpu_estimate_by_wf_runtime,
288
+ gpu_estimate_by_llm_latency=gpu_estimate_by_llm_latency)
@@ -0,0 +1,188 @@
1
+ # SPDX-FileCopyrightText: Copyright (c) 2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
2
+ # SPDX-License-Identifier: Apache-2.0
3
+ #
4
+ # Licensed under the Apache License, Version 2.0 (the "License");
5
+ # you may not use this file except in compliance with the License.
6
+ # You may obtain a copy of the License at
7
+ #
8
+ # http://www.apache.org/licenses/LICENSE-2.0
9
+ #
10
+ # Unless required by applicable law or agreed to in writing, software
11
+ # distributed under the License is distributed on an "AS IS" BASIS,
12
+ # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13
+ # See the License for the specific language governing permissions and
14
+ # limitations under the License.
15
+
16
+ import typing
17
+ from pathlib import Path
18
+
19
+ from pydantic import BaseModel
20
+ from pydantic import Field
21
+
22
+
23
+ class FitConfig(BaseModel):
24
+ """
25
+ Configuration parameters for linear fit and outlier detection.
26
+ """
27
+ # Threshold for small concurrency range (≤ 8 points) to check for extreme outliers in raw y-values first
28
+ small_concurrency_range_threshold: int = 8
29
+
30
+ # Extreme outlier threshold is 2.0 times the IQR, extreme outliers are removed
31
+ extreme_outlier_threshold: float = 2.0
32
+
33
+ # Conservative outlier threshold is 1.5 times the IQR, conservative outliers are removed
34
+ conservative_outlier_threshold: float = 1.5
35
+
36
+ # Minimum R-squared value required for a valid linear fit
37
+ min_r_squared: float = 0.7
38
+
39
+ # Whether to remove outliers during linear fit calculation
40
+ remove_outliers: bool = True
41
+
42
+
43
+ class CalcRunnerConfig(BaseModel):
44
+ """
45
+ Parameters used for a calc runner.
46
+ """
47
+ # base config and endpoints (if remote)- not needed in offline mode
48
+ config_file: Path | None = None
49
+ # endpoint to use for the workflow, if not provided the workflow is run locally
50
+ endpoint: str | None = None
51
+ # timeout for the workflow
52
+ endpoint_timeout: int = 300
53
+
54
+ # if true workflow is not run, instead results from previous runs are used to estimate the
55
+ # GPU count
56
+ offline_mode: bool = False
57
+
58
+ # number of passes at each concurrency, if 0 the dataset is adjusted to a multiple of the
59
+ # concurrency
60
+ num_passes: int = 0
61
+ # concurrency values to test
62
+ concurrencies: list[int] = [1, 2, 4, 8]
63
+
64
+ # Targets for GPU estimation
65
+ target_llm_latency_p95: float = 0
66
+ target_workflow_runtime_p95: float = 0
67
+ target_users: int = 0
68
+
69
+ # Test setup information needed for GPU estimation
70
+ test_gpu_count: int = 0
71
+
72
+ # output directory for results
73
+ output_dir: Path | None = None
74
+ # if true, the job is stored in a new subdirectory of the output directory
75
+ append_job: bool = False
76
+ # if true, the data is plotted
77
+ plot_data: bool = True
78
+
79
+ # Configuration for linear fit and outlier detection
80
+ fit_config: FitConfig = Field(default_factory=FitConfig)
81
+
82
+
83
+ # Sizing metrics are gathered from the evaluation runs and used as input by the calculator.
84
+ class SizingMetricPerItem(BaseModel):
85
+ """
86
+ Sizing metrics per dataset entry item.
87
+ """
88
+ # LLM latency
89
+ llm_latency: float
90
+ # workflow runtime
91
+ workflow_runtime: float
92
+
93
+
94
+ class SizingMetricsAlerts(BaseModel):
95
+ """
96
+ Sizing metrics alerts.
97
+ """
98
+ # if true, the workflow was interrupted that concurrency cannot be used
99
+ workflow_interrupted: bool = False
100
+
101
+
102
+ class SizingMetrics(BaseModel):
103
+ """
104
+ Sizing metrics for a single concurrency.
105
+ """
106
+ # alerts associated with the sizing metrics
107
+ alerts: SizingMetricsAlerts = Field(default_factory=SizingMetricsAlerts)
108
+
109
+ # p95 LLM latency
110
+ llm_latency_p95: float = 0.0
111
+ # p95 workflow runtime
112
+ workflow_runtime_p95: float = 0.0
113
+ # total workflow runtime
114
+ total_runtime: float = 0.0
115
+ # per item metrics, key is the dataset entry id
116
+ per_item_metrics: dict[typing.Any, SizingMetricPerItem] = {}
117
+
118
+
119
+ class LinearFitResult(BaseModel):
120
+ """
121
+ Result of linear regression including slope, intercept, and quality metrics.
122
+ """
123
+ slope: float
124
+ intercept: float
125
+ r_squared: float
126
+ outliers_removed: list[int]
127
+
128
+
129
+ class FitResults(BaseModel):
130
+ """
131
+ Linear fit results for both LLM latency and workflow runtime analysis.
132
+ """
133
+ llm_latency_fit: LinearFitResult | None = None
134
+ wf_runtime_fit: LinearFitResult | None = None
135
+
136
+
137
+ # GPU estimates are generated by the calculator.
138
+ class GPUEstimates(BaseModel):
139
+ """
140
+ GPU estimates.
141
+ """
142
+ # GPU estimate based on the workflow runtime
143
+ gpu_estimate_by_wf_runtime: float | None = None
144
+ # GPU estimate based on the LLM latency
145
+ gpu_estimate_by_llm_latency: float | None = None
146
+
147
+
148
+ # Calc runner alerts are generated by the calculator.
149
+ class CalcAlerts(BaseModel):
150
+ """
151
+ Calc runner alerts.
152
+ """
153
+ # if true, the run was identified as an outlier by the workflow runtime linear fit
154
+ outlier_workflow_runtime: bool = False
155
+ # if true, the run was identified as an outlier by the LLM latency linear fit
156
+ outlier_llm_latency: bool = False
157
+
158
+ # number of items that are greater than the target latency
159
+ num_items_greater_than_target_latency: int = 0
160
+ # number of items that are greater than the target runtime
161
+ num_items_greater_than_target_runtime: int = 0
162
+
163
+
164
+ class CalcData(BaseModel):
165
+ """
166
+ Output of the calc runner per concurrency.
167
+ """
168
+ # ROUGH GPU estimates per concurrency: these are not used for the final GPU estimation
169
+ # they are only available for information purposes
170
+ gpu_estimates: GPUEstimates = Field(default_factory=GPUEstimates)
171
+ # Calc runner alerts
172
+ alerts: CalcAlerts = Field(default_factory=CalcAlerts)
173
+ # Sizing metrics
174
+ sizing_metrics: SizingMetrics = Field(default_factory=SizingMetrics)
175
+
176
+
177
+ class CalcRunnerOutput(BaseModel):
178
+ """
179
+ Output of the calc runner.
180
+ """
181
+ # GPU estimates based on the slope of the time vs concurrency, calculated online or offline
182
+ gpu_estimates: GPUEstimates = Field(default_factory=GPUEstimates)
183
+
184
+ # Linear fit results for analysis and debugging
185
+ fit_results: FitResults = Field(default_factory=FitResults)
186
+
187
+ # Per-concurrency data (GPU estimates, out-of-range runs, and sizing metrics)
188
+ calc_data: dict[int, CalcData] = {}