wandb 0.18.2__py3-none-musllinux_1_2_x86_64.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 (827) hide show
  1. package_readme.md +89 -0
  2. wandb/__init__.py +245 -0
  3. wandb/__init__.pyi +1139 -0
  4. wandb/__main__.py +3 -0
  5. wandb/_globals.py +19 -0
  6. wandb/agents/__init__.py +0 -0
  7. wandb/agents/pyagent.py +363 -0
  8. wandb/analytics/__init__.py +3 -0
  9. wandb/analytics/sentry.py +266 -0
  10. wandb/apis/__init__.py +48 -0
  11. wandb/apis/attrs.py +40 -0
  12. wandb/apis/importers/__init__.py +1 -0
  13. wandb/apis/importers/internals/internal.py +385 -0
  14. wandb/apis/importers/internals/protocols.py +99 -0
  15. wandb/apis/importers/internals/util.py +78 -0
  16. wandb/apis/importers/mlflow.py +254 -0
  17. wandb/apis/importers/validation.py +108 -0
  18. wandb/apis/importers/wandb.py +1603 -0
  19. wandb/apis/internal.py +232 -0
  20. wandb/apis/normalize.py +89 -0
  21. wandb/apis/paginator.py +81 -0
  22. wandb/apis/public/__init__.py +34 -0
  23. wandb/apis/public/api.py +1305 -0
  24. wandb/apis/public/artifacts.py +1090 -0
  25. wandb/apis/public/const.py +4 -0
  26. wandb/apis/public/files.py +195 -0
  27. wandb/apis/public/history.py +149 -0
  28. wandb/apis/public/jobs.py +659 -0
  29. wandb/apis/public/projects.py +154 -0
  30. wandb/apis/public/query_generator.py +166 -0
  31. wandb/apis/public/reports.py +469 -0
  32. wandb/apis/public/runs.py +914 -0
  33. wandb/apis/public/sweeps.py +240 -0
  34. wandb/apis/public/teams.py +198 -0
  35. wandb/apis/public/users.py +136 -0
  36. wandb/apis/reports/__init__.py +1 -0
  37. wandb/apis/reports/v1/__init__.py +8 -0
  38. wandb/apis/reports/v2/__init__.py +8 -0
  39. wandb/apis/workspaces/__init__.py +8 -0
  40. wandb/beta/workflows.py +288 -0
  41. wandb/bin/nvidia_gpu_stats +0 -0
  42. wandb/bin/wandb-core +0 -0
  43. wandb/cli/__init__.py +0 -0
  44. wandb/cli/cli.py +3004 -0
  45. wandb/data_types.py +63 -0
  46. wandb/docker/__init__.py +342 -0
  47. wandb/docker/auth.py +436 -0
  48. wandb/docker/wandb-entrypoint.sh +33 -0
  49. wandb/docker/www_authenticate.py +94 -0
  50. wandb/env.py +514 -0
  51. wandb/errors/__init__.py +17 -0
  52. wandb/errors/errors.py +37 -0
  53. wandb/errors/term.py +103 -0
  54. wandb/errors/util.py +57 -0
  55. wandb/errors/warnings.py +2 -0
  56. wandb/filesync/__init__.py +0 -0
  57. wandb/filesync/dir_watcher.py +403 -0
  58. wandb/filesync/stats.py +100 -0
  59. wandb/filesync/step_checksum.py +142 -0
  60. wandb/filesync/step_prepare.py +179 -0
  61. wandb/filesync/step_upload.py +290 -0
  62. wandb/filesync/upload_job.py +142 -0
  63. wandb/integration/__init__.py +0 -0
  64. wandb/integration/catboost/__init__.py +5 -0
  65. wandb/integration/catboost/catboost.py +178 -0
  66. wandb/integration/cohere/__init__.py +3 -0
  67. wandb/integration/cohere/cohere.py +21 -0
  68. wandb/integration/cohere/resolver.py +347 -0
  69. wandb/integration/diffusers/__init__.py +3 -0
  70. wandb/integration/diffusers/autologger.py +76 -0
  71. wandb/integration/diffusers/pipeline_resolver.py +50 -0
  72. wandb/integration/diffusers/resolvers/__init__.py +9 -0
  73. wandb/integration/diffusers/resolvers/multimodal.py +882 -0
  74. wandb/integration/diffusers/resolvers/utils.py +102 -0
  75. wandb/integration/fastai/__init__.py +249 -0
  76. wandb/integration/gym/__init__.py +105 -0
  77. wandb/integration/huggingface/__init__.py +3 -0
  78. wandb/integration/huggingface/huggingface.py +18 -0
  79. wandb/integration/huggingface/resolver.py +213 -0
  80. wandb/integration/keras/__init__.py +11 -0
  81. wandb/integration/keras/callbacks/__init__.py +5 -0
  82. wandb/integration/keras/callbacks/metrics_logger.py +136 -0
  83. wandb/integration/keras/callbacks/model_checkpoint.py +195 -0
  84. wandb/integration/keras/callbacks/tables_builder.py +226 -0
  85. wandb/integration/keras/keras.py +1091 -0
  86. wandb/integration/kfp/__init__.py +6 -0
  87. wandb/integration/kfp/helpers.py +28 -0
  88. wandb/integration/kfp/kfp_patch.py +324 -0
  89. wandb/integration/kfp/wandb_logging.py +182 -0
  90. wandb/integration/langchain/__init__.py +3 -0
  91. wandb/integration/langchain/wandb_tracer.py +48 -0
  92. wandb/integration/lightgbm/__init__.py +239 -0
  93. wandb/integration/lightning/__init__.py +0 -0
  94. wandb/integration/lightning/fabric/__init__.py +3 -0
  95. wandb/integration/lightning/fabric/logger.py +762 -0
  96. wandb/integration/magic.py +556 -0
  97. wandb/integration/metaflow/__init__.py +3 -0
  98. wandb/integration/metaflow/metaflow.py +383 -0
  99. wandb/integration/openai/__init__.py +3 -0
  100. wandb/integration/openai/fine_tuning.py +480 -0
  101. wandb/integration/openai/openai.py +22 -0
  102. wandb/integration/openai/resolver.py +240 -0
  103. wandb/integration/prodigy/__init__.py +3 -0
  104. wandb/integration/prodigy/prodigy.py +299 -0
  105. wandb/integration/sacred/__init__.py +117 -0
  106. wandb/integration/sagemaker/__init__.py +12 -0
  107. wandb/integration/sagemaker/auth.py +28 -0
  108. wandb/integration/sagemaker/config.py +49 -0
  109. wandb/integration/sagemaker/files.py +3 -0
  110. wandb/integration/sagemaker/resources.py +34 -0
  111. wandb/integration/sb3/__init__.py +3 -0
  112. wandb/integration/sb3/sb3.py +153 -0
  113. wandb/integration/sklearn/__init__.py +37 -0
  114. wandb/integration/sklearn/calculate/__init__.py +32 -0
  115. wandb/integration/sklearn/calculate/calibration_curves.py +125 -0
  116. wandb/integration/sklearn/calculate/class_proportions.py +68 -0
  117. wandb/integration/sklearn/calculate/confusion_matrix.py +93 -0
  118. wandb/integration/sklearn/calculate/decision_boundaries.py +40 -0
  119. wandb/integration/sklearn/calculate/elbow_curve.py +55 -0
  120. wandb/integration/sklearn/calculate/feature_importances.py +67 -0
  121. wandb/integration/sklearn/calculate/learning_curve.py +64 -0
  122. wandb/integration/sklearn/calculate/outlier_candidates.py +69 -0
  123. wandb/integration/sklearn/calculate/residuals.py +86 -0
  124. wandb/integration/sklearn/calculate/silhouette.py +118 -0
  125. wandb/integration/sklearn/calculate/summary_metrics.py +62 -0
  126. wandb/integration/sklearn/plot/__init__.py +35 -0
  127. wandb/integration/sklearn/plot/classifier.py +329 -0
  128. wandb/integration/sklearn/plot/clusterer.py +146 -0
  129. wandb/integration/sklearn/plot/regressor.py +121 -0
  130. wandb/integration/sklearn/plot/shared.py +91 -0
  131. wandb/integration/sklearn/utils.py +183 -0
  132. wandb/integration/tensorboard/__init__.py +10 -0
  133. wandb/integration/tensorboard/log.py +355 -0
  134. wandb/integration/tensorboard/monkeypatch.py +185 -0
  135. wandb/integration/tensorflow/__init__.py +5 -0
  136. wandb/integration/tensorflow/estimator_hook.py +54 -0
  137. wandb/integration/torch/__init__.py +0 -0
  138. wandb/integration/torch/wandb_torch.py +554 -0
  139. wandb/integration/ultralytics/__init__.py +11 -0
  140. wandb/integration/ultralytics/bbox_utils.py +208 -0
  141. wandb/integration/ultralytics/callback.py +524 -0
  142. wandb/integration/ultralytics/classification_utils.py +83 -0
  143. wandb/integration/ultralytics/mask_utils.py +202 -0
  144. wandb/integration/ultralytics/pose_utils.py +103 -0
  145. wandb/integration/xgboost/__init__.py +11 -0
  146. wandb/integration/xgboost/xgboost.py +189 -0
  147. wandb/integration/yolov8/__init__.py +0 -0
  148. wandb/integration/yolov8/yolov8.py +284 -0
  149. wandb/jupyter.py +515 -0
  150. wandb/magic.py +3 -0
  151. wandb/mpmain/__init__.py +0 -0
  152. wandb/mpmain/__main__.py +1 -0
  153. wandb/old/__init__.py +0 -0
  154. wandb/old/core.py +53 -0
  155. wandb/old/settings.py +173 -0
  156. wandb/old/summary.py +440 -0
  157. wandb/plot/__init__.py +19 -0
  158. wandb/plot/bar.py +45 -0
  159. wandb/plot/confusion_matrix.py +100 -0
  160. wandb/plot/histogram.py +39 -0
  161. wandb/plot/line.py +43 -0
  162. wandb/plot/line_series.py +88 -0
  163. wandb/plot/pr_curve.py +136 -0
  164. wandb/plot/roc_curve.py +118 -0
  165. wandb/plot/scatter.py +32 -0
  166. wandb/plot/utils.py +183 -0
  167. wandb/plot/viz.py +123 -0
  168. wandb/proto/__init__.py +0 -0
  169. wandb/proto/v3/__init__.py +0 -0
  170. wandb/proto/v3/wandb_base_pb2.py +55 -0
  171. wandb/proto/v3/wandb_internal_pb2.py +1608 -0
  172. wandb/proto/v3/wandb_server_pb2.py +208 -0
  173. wandb/proto/v3/wandb_settings_pb2.py +112 -0
  174. wandb/proto/v3/wandb_telemetry_pb2.py +106 -0
  175. wandb/proto/v4/__init__.py +0 -0
  176. wandb/proto/v4/wandb_base_pb2.py +30 -0
  177. wandb/proto/v4/wandb_internal_pb2.py +360 -0
  178. wandb/proto/v4/wandb_server_pb2.py +63 -0
  179. wandb/proto/v4/wandb_settings_pb2.py +45 -0
  180. wandb/proto/v4/wandb_telemetry_pb2.py +41 -0
  181. wandb/proto/v5/wandb_base_pb2.py +31 -0
  182. wandb/proto/v5/wandb_internal_pb2.py +361 -0
  183. wandb/proto/v5/wandb_server_pb2.py +64 -0
  184. wandb/proto/v5/wandb_settings_pb2.py +46 -0
  185. wandb/proto/v5/wandb_telemetry_pb2.py +42 -0
  186. wandb/proto/wandb_base_pb2.py +10 -0
  187. wandb/proto/wandb_deprecated.py +53 -0
  188. wandb/proto/wandb_generate_deprecated.py +34 -0
  189. wandb/proto/wandb_generate_proto.py +49 -0
  190. wandb/proto/wandb_internal_pb2.py +16 -0
  191. wandb/proto/wandb_server_pb2.py +10 -0
  192. wandb/proto/wandb_settings_pb2.py +10 -0
  193. wandb/proto/wandb_telemetry_pb2.py +10 -0
  194. wandb/py.typed +0 -0
  195. wandb/sdk/__init__.py +37 -0
  196. wandb/sdk/artifacts/__init__.py +0 -0
  197. wandb/sdk/artifacts/_validators.py +90 -0
  198. wandb/sdk/artifacts/artifact.py +2389 -0
  199. wandb/sdk/artifacts/artifact_download_logger.py +43 -0
  200. wandb/sdk/artifacts/artifact_file_cache.py +253 -0
  201. wandb/sdk/artifacts/artifact_instance_cache.py +17 -0
  202. wandb/sdk/artifacts/artifact_manifest.py +74 -0
  203. wandb/sdk/artifacts/artifact_manifest_entry.py +249 -0
  204. wandb/sdk/artifacts/artifact_manifests/__init__.py +0 -0
  205. wandb/sdk/artifacts/artifact_manifests/artifact_manifest_v1.py +92 -0
  206. wandb/sdk/artifacts/artifact_saver.py +269 -0
  207. wandb/sdk/artifacts/artifact_state.py +11 -0
  208. wandb/sdk/artifacts/artifact_ttl.py +7 -0
  209. wandb/sdk/artifacts/exceptions.py +57 -0
  210. wandb/sdk/artifacts/staging.py +25 -0
  211. wandb/sdk/artifacts/storage_handler.py +62 -0
  212. wandb/sdk/artifacts/storage_handlers/__init__.py +0 -0
  213. wandb/sdk/artifacts/storage_handlers/azure_handler.py +208 -0
  214. wandb/sdk/artifacts/storage_handlers/gcs_handler.py +228 -0
  215. wandb/sdk/artifacts/storage_handlers/http_handler.py +114 -0
  216. wandb/sdk/artifacts/storage_handlers/local_file_handler.py +141 -0
  217. wandb/sdk/artifacts/storage_handlers/multi_handler.py +56 -0
  218. wandb/sdk/artifacts/storage_handlers/s3_handler.py +300 -0
  219. wandb/sdk/artifacts/storage_handlers/tracking_handler.py +72 -0
  220. wandb/sdk/artifacts/storage_handlers/wb_artifact_handler.py +135 -0
  221. wandb/sdk/artifacts/storage_handlers/wb_local_artifact_handler.py +74 -0
  222. wandb/sdk/artifacts/storage_layout.py +6 -0
  223. wandb/sdk/artifacts/storage_policies/__init__.py +4 -0
  224. wandb/sdk/artifacts/storage_policies/register.py +1 -0
  225. wandb/sdk/artifacts/storage_policies/wandb_storage_policy.py +378 -0
  226. wandb/sdk/artifacts/storage_policy.py +72 -0
  227. wandb/sdk/backend/__init__.py +0 -0
  228. wandb/sdk/backend/backend.py +222 -0
  229. wandb/sdk/data_types/__init__.py +0 -0
  230. wandb/sdk/data_types/_dtypes.py +914 -0
  231. wandb/sdk/data_types/_private.py +10 -0
  232. wandb/sdk/data_types/audio.py +165 -0
  233. wandb/sdk/data_types/base_types/__init__.py +0 -0
  234. wandb/sdk/data_types/base_types/json_metadata.py +55 -0
  235. wandb/sdk/data_types/base_types/media.py +315 -0
  236. wandb/sdk/data_types/base_types/wb_value.py +272 -0
  237. wandb/sdk/data_types/bokeh.py +70 -0
  238. wandb/sdk/data_types/graph.py +405 -0
  239. wandb/sdk/data_types/helper_types/__init__.py +0 -0
  240. wandb/sdk/data_types/helper_types/bounding_boxes_2d.py +295 -0
  241. wandb/sdk/data_types/helper_types/classes.py +159 -0
  242. wandb/sdk/data_types/helper_types/image_mask.py +235 -0
  243. wandb/sdk/data_types/histogram.py +96 -0
  244. wandb/sdk/data_types/html.py +115 -0
  245. wandb/sdk/data_types/image.py +845 -0
  246. wandb/sdk/data_types/molecule.py +241 -0
  247. wandb/sdk/data_types/object_3d.py +474 -0
  248. wandb/sdk/data_types/plotly.py +82 -0
  249. wandb/sdk/data_types/saved_model.py +446 -0
  250. wandb/sdk/data_types/table.py +1204 -0
  251. wandb/sdk/data_types/trace_tree.py +438 -0
  252. wandb/sdk/data_types/utils.py +229 -0
  253. wandb/sdk/data_types/video.py +247 -0
  254. wandb/sdk/integration_utils/__init__.py +0 -0
  255. wandb/sdk/integration_utils/auto_logging.py +239 -0
  256. wandb/sdk/integration_utils/data_logging.py +475 -0
  257. wandb/sdk/interface/__init__.py +0 -0
  258. wandb/sdk/interface/constants.py +4 -0
  259. wandb/sdk/interface/interface.py +972 -0
  260. wandb/sdk/interface/interface_queue.py +59 -0
  261. wandb/sdk/interface/interface_relay.py +53 -0
  262. wandb/sdk/interface/interface_shared.py +537 -0
  263. wandb/sdk/interface/interface_sock.py +61 -0
  264. wandb/sdk/interface/message_future.py +27 -0
  265. wandb/sdk/interface/message_future_poll.py +50 -0
  266. wandb/sdk/interface/router.py +118 -0
  267. wandb/sdk/interface/router_queue.py +44 -0
  268. wandb/sdk/interface/router_relay.py +39 -0
  269. wandb/sdk/interface/router_sock.py +36 -0
  270. wandb/sdk/interface/summary_record.py +67 -0
  271. wandb/sdk/internal/__init__.py +0 -0
  272. wandb/sdk/internal/context.py +89 -0
  273. wandb/sdk/internal/datastore.py +297 -0
  274. wandb/sdk/internal/file_pusher.py +181 -0
  275. wandb/sdk/internal/file_stream.py +695 -0
  276. wandb/sdk/internal/flow_control.py +263 -0
  277. wandb/sdk/internal/handler.py +901 -0
  278. wandb/sdk/internal/internal.py +417 -0
  279. wandb/sdk/internal/internal_api.py +4358 -0
  280. wandb/sdk/internal/internal_util.py +100 -0
  281. wandb/sdk/internal/job_builder.py +629 -0
  282. wandb/sdk/internal/profiler.py +78 -0
  283. wandb/sdk/internal/progress.py +83 -0
  284. wandb/sdk/internal/run.py +25 -0
  285. wandb/sdk/internal/sample.py +70 -0
  286. wandb/sdk/internal/sender.py +1686 -0
  287. wandb/sdk/internal/sender_config.py +197 -0
  288. wandb/sdk/internal/settings_static.py +90 -0
  289. wandb/sdk/internal/system/__init__.py +0 -0
  290. wandb/sdk/internal/system/assets/__init__.py +27 -0
  291. wandb/sdk/internal/system/assets/aggregators.py +37 -0
  292. wandb/sdk/internal/system/assets/asset_registry.py +20 -0
  293. wandb/sdk/internal/system/assets/cpu.py +163 -0
  294. wandb/sdk/internal/system/assets/disk.py +210 -0
  295. wandb/sdk/internal/system/assets/gpu.py +416 -0
  296. wandb/sdk/internal/system/assets/gpu_amd.py +239 -0
  297. wandb/sdk/internal/system/assets/gpu_apple.py +177 -0
  298. wandb/sdk/internal/system/assets/interfaces.py +207 -0
  299. wandb/sdk/internal/system/assets/ipu.py +177 -0
  300. wandb/sdk/internal/system/assets/memory.py +166 -0
  301. wandb/sdk/internal/system/assets/network.py +125 -0
  302. wandb/sdk/internal/system/assets/open_metrics.py +299 -0
  303. wandb/sdk/internal/system/assets/tpu.py +154 -0
  304. wandb/sdk/internal/system/assets/trainium.py +399 -0
  305. wandb/sdk/internal/system/env_probe_helpers.py +13 -0
  306. wandb/sdk/internal/system/system_info.py +249 -0
  307. wandb/sdk/internal/system/system_monitor.py +229 -0
  308. wandb/sdk/internal/tb_watcher.py +518 -0
  309. wandb/sdk/internal/thread_local_settings.py +18 -0
  310. wandb/sdk/internal/writer.py +206 -0
  311. wandb/sdk/launch/__init__.py +14 -0
  312. wandb/sdk/launch/_launch.py +330 -0
  313. wandb/sdk/launch/_launch_add.py +255 -0
  314. wandb/sdk/launch/_project_spec.py +566 -0
  315. wandb/sdk/launch/agent/__init__.py +5 -0
  316. wandb/sdk/launch/agent/agent.py +924 -0
  317. wandb/sdk/launch/agent/config.py +296 -0
  318. wandb/sdk/launch/agent/job_status_tracker.py +53 -0
  319. wandb/sdk/launch/agent/run_queue_item_file_saver.py +45 -0
  320. wandb/sdk/launch/builder/__init__.py +0 -0
  321. wandb/sdk/launch/builder/abstract.py +156 -0
  322. wandb/sdk/launch/builder/build.py +297 -0
  323. wandb/sdk/launch/builder/context_manager.py +235 -0
  324. wandb/sdk/launch/builder/docker_builder.py +177 -0
  325. wandb/sdk/launch/builder/kaniko_builder.py +595 -0
  326. wandb/sdk/launch/builder/noop.py +58 -0
  327. wandb/sdk/launch/builder/templates/_wandb_bootstrap.py +188 -0
  328. wandb/sdk/launch/builder/templates/dockerfile.py +92 -0
  329. wandb/sdk/launch/create_job.py +528 -0
  330. wandb/sdk/launch/environment/abstract.py +29 -0
  331. wandb/sdk/launch/environment/aws_environment.py +322 -0
  332. wandb/sdk/launch/environment/azure_environment.py +105 -0
  333. wandb/sdk/launch/environment/gcp_environment.py +335 -0
  334. wandb/sdk/launch/environment/local_environment.py +66 -0
  335. wandb/sdk/launch/errors.py +19 -0
  336. wandb/sdk/launch/git_reference.py +109 -0
  337. wandb/sdk/launch/inputs/files.py +148 -0
  338. wandb/sdk/launch/inputs/internal.py +315 -0
  339. wandb/sdk/launch/inputs/manage.py +113 -0
  340. wandb/sdk/launch/inputs/schema.py +39 -0
  341. wandb/sdk/launch/loader.py +249 -0
  342. wandb/sdk/launch/registry/abstract.py +48 -0
  343. wandb/sdk/launch/registry/anon.py +29 -0
  344. wandb/sdk/launch/registry/azure_container_registry.py +124 -0
  345. wandb/sdk/launch/registry/elastic_container_registry.py +192 -0
  346. wandb/sdk/launch/registry/google_artifact_registry.py +219 -0
  347. wandb/sdk/launch/registry/local_registry.py +67 -0
  348. wandb/sdk/launch/runner/__init__.py +0 -0
  349. wandb/sdk/launch/runner/abstract.py +195 -0
  350. wandb/sdk/launch/runner/kubernetes_monitor.py +474 -0
  351. wandb/sdk/launch/runner/kubernetes_runner.py +963 -0
  352. wandb/sdk/launch/runner/local_container.py +301 -0
  353. wandb/sdk/launch/runner/local_process.py +78 -0
  354. wandb/sdk/launch/runner/sagemaker_runner.py +426 -0
  355. wandb/sdk/launch/runner/vertex_runner.py +230 -0
  356. wandb/sdk/launch/sweeps/__init__.py +39 -0
  357. wandb/sdk/launch/sweeps/scheduler.py +742 -0
  358. wandb/sdk/launch/sweeps/scheduler_sweep.py +91 -0
  359. wandb/sdk/launch/sweeps/utils.py +316 -0
  360. wandb/sdk/launch/utils.py +746 -0
  361. wandb/sdk/launch/wandb_reference.py +138 -0
  362. wandb/sdk/lib/__init__.py +5 -0
  363. wandb/sdk/lib/_settings_toposort_generate.py +159 -0
  364. wandb/sdk/lib/_settings_toposort_generated.py +250 -0
  365. wandb/sdk/lib/_wburls_generate.py +25 -0
  366. wandb/sdk/lib/_wburls_generated.py +22 -0
  367. wandb/sdk/lib/apikey.py +273 -0
  368. wandb/sdk/lib/capped_dict.py +26 -0
  369. wandb/sdk/lib/config_util.py +101 -0
  370. wandb/sdk/lib/credentials.py +141 -0
  371. wandb/sdk/lib/deprecate.py +42 -0
  372. wandb/sdk/lib/disabled.py +29 -0
  373. wandb/sdk/lib/exit_hooks.py +54 -0
  374. wandb/sdk/lib/file_stream_utils.py +118 -0
  375. wandb/sdk/lib/filenames.py +64 -0
  376. wandb/sdk/lib/filesystem.py +372 -0
  377. wandb/sdk/lib/fsm.py +174 -0
  378. wandb/sdk/lib/gitlib.py +239 -0
  379. wandb/sdk/lib/gql_request.py +65 -0
  380. wandb/sdk/lib/handler_util.py +21 -0
  381. wandb/sdk/lib/hashutil.py +84 -0
  382. wandb/sdk/lib/import_hooks.py +275 -0
  383. wandb/sdk/lib/ipython.py +146 -0
  384. wandb/sdk/lib/json_util.py +80 -0
  385. wandb/sdk/lib/lazyloader.py +63 -0
  386. wandb/sdk/lib/mailbox.py +460 -0
  387. wandb/sdk/lib/module.py +69 -0
  388. wandb/sdk/lib/paths.py +106 -0
  389. wandb/sdk/lib/preinit.py +42 -0
  390. wandb/sdk/lib/printer.py +313 -0
  391. wandb/sdk/lib/proto_util.py +90 -0
  392. wandb/sdk/lib/redirect.py +845 -0
  393. wandb/sdk/lib/reporting.py +99 -0
  394. wandb/sdk/lib/retry.py +289 -0
  395. wandb/sdk/lib/run_moment.py +78 -0
  396. wandb/sdk/lib/runid.py +12 -0
  397. wandb/sdk/lib/server.py +52 -0
  398. wandb/sdk/lib/service_connection.py +216 -0
  399. wandb/sdk/lib/service_token.py +94 -0
  400. wandb/sdk/lib/sock_client.py +295 -0
  401. wandb/sdk/lib/sparkline.py +45 -0
  402. wandb/sdk/lib/telemetry.py +100 -0
  403. wandb/sdk/lib/timed_input.py +133 -0
  404. wandb/sdk/lib/timer.py +19 -0
  405. wandb/sdk/lib/tracelog.py +255 -0
  406. wandb/sdk/lib/wburls.py +46 -0
  407. wandb/sdk/service/__init__.py +0 -0
  408. wandb/sdk/service/_startup_debug.py +22 -0
  409. wandb/sdk/service/port_file.py +53 -0
  410. wandb/sdk/service/server.py +116 -0
  411. wandb/sdk/service/server_sock.py +276 -0
  412. wandb/sdk/service/service.py +242 -0
  413. wandb/sdk/service/streams.py +417 -0
  414. wandb/sdk/verify/__init__.py +0 -0
  415. wandb/sdk/verify/verify.py +501 -0
  416. wandb/sdk/wandb_alerts.py +12 -0
  417. wandb/sdk/wandb_config.py +322 -0
  418. wandb/sdk/wandb_helper.py +54 -0
  419. wandb/sdk/wandb_init.py +1266 -0
  420. wandb/sdk/wandb_login.py +349 -0
  421. wandb/sdk/wandb_metric.py +110 -0
  422. wandb/sdk/wandb_require.py +97 -0
  423. wandb/sdk/wandb_require_helpers.py +44 -0
  424. wandb/sdk/wandb_run.py +4236 -0
  425. wandb/sdk/wandb_settings.py +2001 -0
  426. wandb/sdk/wandb_setup.py +409 -0
  427. wandb/sdk/wandb_summary.py +150 -0
  428. wandb/sdk/wandb_sweep.py +119 -0
  429. wandb/sdk/wandb_sync.py +81 -0
  430. wandb/sdk/wandb_watch.py +144 -0
  431. wandb/sklearn.py +35 -0
  432. wandb/sync/__init__.py +3 -0
  433. wandb/sync/sync.py +443 -0
  434. wandb/trigger.py +29 -0
  435. wandb/util.py +1956 -0
  436. wandb/vendor/__init__.py +0 -0
  437. wandb/vendor/gql-0.2.0/setup.py +40 -0
  438. wandb/vendor/gql-0.2.0/tests/__init__.py +0 -0
  439. wandb/vendor/gql-0.2.0/tests/starwars/__init__.py +0 -0
  440. wandb/vendor/gql-0.2.0/tests/starwars/fixtures.py +96 -0
  441. wandb/vendor/gql-0.2.0/tests/starwars/schema.py +146 -0
  442. wandb/vendor/gql-0.2.0/tests/starwars/test_dsl.py +293 -0
  443. wandb/vendor/gql-0.2.0/tests/starwars/test_query.py +355 -0
  444. wandb/vendor/gql-0.2.0/tests/starwars/test_validation.py +171 -0
  445. wandb/vendor/gql-0.2.0/tests/test_client.py +31 -0
  446. wandb/vendor/gql-0.2.0/tests/test_transport.py +89 -0
  447. wandb/vendor/gql-0.2.0/wandb_gql/__init__.py +4 -0
  448. wandb/vendor/gql-0.2.0/wandb_gql/client.py +75 -0
  449. wandb/vendor/gql-0.2.0/wandb_gql/dsl.py +152 -0
  450. wandb/vendor/gql-0.2.0/wandb_gql/gql.py +10 -0
  451. wandb/vendor/gql-0.2.0/wandb_gql/transport/__init__.py +0 -0
  452. wandb/vendor/gql-0.2.0/wandb_gql/transport/http.py +6 -0
  453. wandb/vendor/gql-0.2.0/wandb_gql/transport/local_schema.py +15 -0
  454. wandb/vendor/gql-0.2.0/wandb_gql/transport/requests.py +46 -0
  455. wandb/vendor/gql-0.2.0/wandb_gql/utils.py +21 -0
  456. wandb/vendor/graphql-core-1.1/setup.py +86 -0
  457. wandb/vendor/graphql-core-1.1/wandb_graphql/__init__.py +287 -0
  458. wandb/vendor/graphql-core-1.1/wandb_graphql/error/__init__.py +6 -0
  459. wandb/vendor/graphql-core-1.1/wandb_graphql/error/base.py +42 -0
  460. wandb/vendor/graphql-core-1.1/wandb_graphql/error/format_error.py +11 -0
  461. wandb/vendor/graphql-core-1.1/wandb_graphql/error/located_error.py +29 -0
  462. wandb/vendor/graphql-core-1.1/wandb_graphql/error/syntax_error.py +36 -0
  463. wandb/vendor/graphql-core-1.1/wandb_graphql/execution/__init__.py +26 -0
  464. wandb/vendor/graphql-core-1.1/wandb_graphql/execution/base.py +311 -0
  465. wandb/vendor/graphql-core-1.1/wandb_graphql/execution/executor.py +398 -0
  466. wandb/vendor/graphql-core-1.1/wandb_graphql/execution/executors/__init__.py +0 -0
  467. wandb/vendor/graphql-core-1.1/wandb_graphql/execution/executors/asyncio.py +53 -0
  468. wandb/vendor/graphql-core-1.1/wandb_graphql/execution/executors/gevent.py +22 -0
  469. wandb/vendor/graphql-core-1.1/wandb_graphql/execution/executors/process.py +32 -0
  470. wandb/vendor/graphql-core-1.1/wandb_graphql/execution/executors/sync.py +7 -0
  471. wandb/vendor/graphql-core-1.1/wandb_graphql/execution/executors/thread.py +35 -0
  472. wandb/vendor/graphql-core-1.1/wandb_graphql/execution/executors/utils.py +6 -0
  473. wandb/vendor/graphql-core-1.1/wandb_graphql/execution/experimental/__init__.py +0 -0
  474. wandb/vendor/graphql-core-1.1/wandb_graphql/execution/experimental/executor.py +66 -0
  475. wandb/vendor/graphql-core-1.1/wandb_graphql/execution/experimental/fragment.py +252 -0
  476. wandb/vendor/graphql-core-1.1/wandb_graphql/execution/experimental/resolver.py +151 -0
  477. wandb/vendor/graphql-core-1.1/wandb_graphql/execution/experimental/utils.py +7 -0
  478. wandb/vendor/graphql-core-1.1/wandb_graphql/execution/middleware.py +57 -0
  479. wandb/vendor/graphql-core-1.1/wandb_graphql/execution/values.py +145 -0
  480. wandb/vendor/graphql-core-1.1/wandb_graphql/graphql.py +60 -0
  481. wandb/vendor/graphql-core-1.1/wandb_graphql/language/__init__.py +0 -0
  482. wandb/vendor/graphql-core-1.1/wandb_graphql/language/ast.py +1349 -0
  483. wandb/vendor/graphql-core-1.1/wandb_graphql/language/base.py +19 -0
  484. wandb/vendor/graphql-core-1.1/wandb_graphql/language/lexer.py +435 -0
  485. wandb/vendor/graphql-core-1.1/wandb_graphql/language/location.py +30 -0
  486. wandb/vendor/graphql-core-1.1/wandb_graphql/language/parser.py +779 -0
  487. wandb/vendor/graphql-core-1.1/wandb_graphql/language/printer.py +193 -0
  488. wandb/vendor/graphql-core-1.1/wandb_graphql/language/source.py +18 -0
  489. wandb/vendor/graphql-core-1.1/wandb_graphql/language/visitor.py +222 -0
  490. wandb/vendor/graphql-core-1.1/wandb_graphql/language/visitor_meta.py +82 -0
  491. wandb/vendor/graphql-core-1.1/wandb_graphql/pyutils/__init__.py +0 -0
  492. wandb/vendor/graphql-core-1.1/wandb_graphql/pyutils/cached_property.py +17 -0
  493. wandb/vendor/graphql-core-1.1/wandb_graphql/pyutils/contain_subset.py +28 -0
  494. wandb/vendor/graphql-core-1.1/wandb_graphql/pyutils/default_ordered_dict.py +40 -0
  495. wandb/vendor/graphql-core-1.1/wandb_graphql/pyutils/ordereddict.py +8 -0
  496. wandb/vendor/graphql-core-1.1/wandb_graphql/pyutils/pair_set.py +43 -0
  497. wandb/vendor/graphql-core-1.1/wandb_graphql/pyutils/version.py +78 -0
  498. wandb/vendor/graphql-core-1.1/wandb_graphql/type/__init__.py +67 -0
  499. wandb/vendor/graphql-core-1.1/wandb_graphql/type/definition.py +619 -0
  500. wandb/vendor/graphql-core-1.1/wandb_graphql/type/directives.py +132 -0
  501. wandb/vendor/graphql-core-1.1/wandb_graphql/type/introspection.py +440 -0
  502. wandb/vendor/graphql-core-1.1/wandb_graphql/type/scalars.py +131 -0
  503. wandb/vendor/graphql-core-1.1/wandb_graphql/type/schema.py +100 -0
  504. wandb/vendor/graphql-core-1.1/wandb_graphql/type/typemap.py +145 -0
  505. wandb/vendor/graphql-core-1.1/wandb_graphql/utils/__init__.py +0 -0
  506. wandb/vendor/graphql-core-1.1/wandb_graphql/utils/assert_valid_name.py +9 -0
  507. wandb/vendor/graphql-core-1.1/wandb_graphql/utils/ast_from_value.py +65 -0
  508. wandb/vendor/graphql-core-1.1/wandb_graphql/utils/ast_to_code.py +49 -0
  509. wandb/vendor/graphql-core-1.1/wandb_graphql/utils/ast_to_dict.py +24 -0
  510. wandb/vendor/graphql-core-1.1/wandb_graphql/utils/base.py +75 -0
  511. wandb/vendor/graphql-core-1.1/wandb_graphql/utils/build_ast_schema.py +291 -0
  512. wandb/vendor/graphql-core-1.1/wandb_graphql/utils/build_client_schema.py +250 -0
  513. wandb/vendor/graphql-core-1.1/wandb_graphql/utils/concat_ast.py +9 -0
  514. wandb/vendor/graphql-core-1.1/wandb_graphql/utils/extend_schema.py +357 -0
  515. wandb/vendor/graphql-core-1.1/wandb_graphql/utils/get_field_def.py +27 -0
  516. wandb/vendor/graphql-core-1.1/wandb_graphql/utils/get_operation_ast.py +21 -0
  517. wandb/vendor/graphql-core-1.1/wandb_graphql/utils/introspection_query.py +90 -0
  518. wandb/vendor/graphql-core-1.1/wandb_graphql/utils/is_valid_literal_value.py +67 -0
  519. wandb/vendor/graphql-core-1.1/wandb_graphql/utils/is_valid_value.py +66 -0
  520. wandb/vendor/graphql-core-1.1/wandb_graphql/utils/quoted_or_list.py +21 -0
  521. wandb/vendor/graphql-core-1.1/wandb_graphql/utils/schema_printer.py +168 -0
  522. wandb/vendor/graphql-core-1.1/wandb_graphql/utils/suggestion_list.py +56 -0
  523. wandb/vendor/graphql-core-1.1/wandb_graphql/utils/type_comparators.py +69 -0
  524. wandb/vendor/graphql-core-1.1/wandb_graphql/utils/type_from_ast.py +21 -0
  525. wandb/vendor/graphql-core-1.1/wandb_graphql/utils/type_info.py +149 -0
  526. wandb/vendor/graphql-core-1.1/wandb_graphql/utils/value_from_ast.py +69 -0
  527. wandb/vendor/graphql-core-1.1/wandb_graphql/validation/__init__.py +4 -0
  528. wandb/vendor/graphql-core-1.1/wandb_graphql/validation/rules/__init__.py +79 -0
  529. wandb/vendor/graphql-core-1.1/wandb_graphql/validation/rules/arguments_of_correct_type.py +24 -0
  530. wandb/vendor/graphql-core-1.1/wandb_graphql/validation/rules/base.py +8 -0
  531. wandb/vendor/graphql-core-1.1/wandb_graphql/validation/rules/default_values_of_correct_type.py +44 -0
  532. wandb/vendor/graphql-core-1.1/wandb_graphql/validation/rules/fields_on_correct_type.py +113 -0
  533. wandb/vendor/graphql-core-1.1/wandb_graphql/validation/rules/fragments_on_composite_types.py +33 -0
  534. wandb/vendor/graphql-core-1.1/wandb_graphql/validation/rules/known_argument_names.py +70 -0
  535. wandb/vendor/graphql-core-1.1/wandb_graphql/validation/rules/known_directives.py +97 -0
  536. wandb/vendor/graphql-core-1.1/wandb_graphql/validation/rules/known_fragment_names.py +19 -0
  537. wandb/vendor/graphql-core-1.1/wandb_graphql/validation/rules/known_type_names.py +43 -0
  538. wandb/vendor/graphql-core-1.1/wandb_graphql/validation/rules/lone_anonymous_operation.py +23 -0
  539. wandb/vendor/graphql-core-1.1/wandb_graphql/validation/rules/no_fragment_cycles.py +59 -0
  540. wandb/vendor/graphql-core-1.1/wandb_graphql/validation/rules/no_undefined_variables.py +36 -0
  541. wandb/vendor/graphql-core-1.1/wandb_graphql/validation/rules/no_unused_fragments.py +38 -0
  542. wandb/vendor/graphql-core-1.1/wandb_graphql/validation/rules/no_unused_variables.py +37 -0
  543. wandb/vendor/graphql-core-1.1/wandb_graphql/validation/rules/overlapping_fields_can_be_merged.py +529 -0
  544. wandb/vendor/graphql-core-1.1/wandb_graphql/validation/rules/possible_fragment_spreads.py +44 -0
  545. wandb/vendor/graphql-core-1.1/wandb_graphql/validation/rules/provided_non_null_arguments.py +46 -0
  546. wandb/vendor/graphql-core-1.1/wandb_graphql/validation/rules/scalar_leafs.py +33 -0
  547. wandb/vendor/graphql-core-1.1/wandb_graphql/validation/rules/unique_argument_names.py +32 -0
  548. wandb/vendor/graphql-core-1.1/wandb_graphql/validation/rules/unique_fragment_names.py +28 -0
  549. wandb/vendor/graphql-core-1.1/wandb_graphql/validation/rules/unique_input_field_names.py +33 -0
  550. wandb/vendor/graphql-core-1.1/wandb_graphql/validation/rules/unique_operation_names.py +31 -0
  551. wandb/vendor/graphql-core-1.1/wandb_graphql/validation/rules/unique_variable_names.py +27 -0
  552. wandb/vendor/graphql-core-1.1/wandb_graphql/validation/rules/variables_are_input_types.py +21 -0
  553. wandb/vendor/graphql-core-1.1/wandb_graphql/validation/rules/variables_in_allowed_position.py +53 -0
  554. wandb/vendor/graphql-core-1.1/wandb_graphql/validation/validation.py +158 -0
  555. wandb/vendor/promise-2.3.0/conftest.py +30 -0
  556. wandb/vendor/promise-2.3.0/setup.py +64 -0
  557. wandb/vendor/promise-2.3.0/tests/__init__.py +0 -0
  558. wandb/vendor/promise-2.3.0/tests/conftest.py +8 -0
  559. wandb/vendor/promise-2.3.0/tests/test_awaitable.py +32 -0
  560. wandb/vendor/promise-2.3.0/tests/test_awaitable_35.py +47 -0
  561. wandb/vendor/promise-2.3.0/tests/test_benchmark.py +116 -0
  562. wandb/vendor/promise-2.3.0/tests/test_complex_threads.py +23 -0
  563. wandb/vendor/promise-2.3.0/tests/test_dataloader.py +452 -0
  564. wandb/vendor/promise-2.3.0/tests/test_dataloader_awaitable_35.py +99 -0
  565. wandb/vendor/promise-2.3.0/tests/test_dataloader_extra.py +65 -0
  566. wandb/vendor/promise-2.3.0/tests/test_extra.py +670 -0
  567. wandb/vendor/promise-2.3.0/tests/test_issues.py +132 -0
  568. wandb/vendor/promise-2.3.0/tests/test_promise_list.py +70 -0
  569. wandb/vendor/promise-2.3.0/tests/test_spec.py +584 -0
  570. wandb/vendor/promise-2.3.0/tests/test_thread_safety.py +115 -0
  571. wandb/vendor/promise-2.3.0/tests/utils.py +3 -0
  572. wandb/vendor/promise-2.3.0/wandb_promise/__init__.py +38 -0
  573. wandb/vendor/promise-2.3.0/wandb_promise/async_.py +135 -0
  574. wandb/vendor/promise-2.3.0/wandb_promise/compat.py +32 -0
  575. wandb/vendor/promise-2.3.0/wandb_promise/dataloader.py +326 -0
  576. wandb/vendor/promise-2.3.0/wandb_promise/iterate_promise.py +12 -0
  577. wandb/vendor/promise-2.3.0/wandb_promise/promise.py +848 -0
  578. wandb/vendor/promise-2.3.0/wandb_promise/promise_list.py +151 -0
  579. wandb/vendor/promise-2.3.0/wandb_promise/pyutils/__init__.py +0 -0
  580. wandb/vendor/promise-2.3.0/wandb_promise/pyutils/version.py +83 -0
  581. wandb/vendor/promise-2.3.0/wandb_promise/schedulers/__init__.py +0 -0
  582. wandb/vendor/promise-2.3.0/wandb_promise/schedulers/asyncio.py +22 -0
  583. wandb/vendor/promise-2.3.0/wandb_promise/schedulers/gevent.py +21 -0
  584. wandb/vendor/promise-2.3.0/wandb_promise/schedulers/immediate.py +27 -0
  585. wandb/vendor/promise-2.3.0/wandb_promise/schedulers/thread.py +18 -0
  586. wandb/vendor/promise-2.3.0/wandb_promise/utils.py +56 -0
  587. wandb/vendor/pygments/__init__.py +90 -0
  588. wandb/vendor/pygments/cmdline.py +568 -0
  589. wandb/vendor/pygments/console.py +74 -0
  590. wandb/vendor/pygments/filter.py +74 -0
  591. wandb/vendor/pygments/filters/__init__.py +350 -0
  592. wandb/vendor/pygments/formatter.py +95 -0
  593. wandb/vendor/pygments/formatters/__init__.py +153 -0
  594. wandb/vendor/pygments/formatters/_mapping.py +85 -0
  595. wandb/vendor/pygments/formatters/bbcode.py +109 -0
  596. wandb/vendor/pygments/formatters/html.py +851 -0
  597. wandb/vendor/pygments/formatters/img.py +600 -0
  598. wandb/vendor/pygments/formatters/irc.py +182 -0
  599. wandb/vendor/pygments/formatters/latex.py +482 -0
  600. wandb/vendor/pygments/formatters/other.py +160 -0
  601. wandb/vendor/pygments/formatters/rtf.py +147 -0
  602. wandb/vendor/pygments/formatters/svg.py +153 -0
  603. wandb/vendor/pygments/formatters/terminal.py +136 -0
  604. wandb/vendor/pygments/formatters/terminal256.py +309 -0
  605. wandb/vendor/pygments/lexer.py +871 -0
  606. wandb/vendor/pygments/lexers/__init__.py +329 -0
  607. wandb/vendor/pygments/lexers/_asy_builtins.py +1645 -0
  608. wandb/vendor/pygments/lexers/_cl_builtins.py +232 -0
  609. wandb/vendor/pygments/lexers/_cocoa_builtins.py +72 -0
  610. wandb/vendor/pygments/lexers/_csound_builtins.py +1346 -0
  611. wandb/vendor/pygments/lexers/_lasso_builtins.py +5327 -0
  612. wandb/vendor/pygments/lexers/_lua_builtins.py +295 -0
  613. wandb/vendor/pygments/lexers/_mapping.py +500 -0
  614. wandb/vendor/pygments/lexers/_mql_builtins.py +1172 -0
  615. wandb/vendor/pygments/lexers/_openedge_builtins.py +2547 -0
  616. wandb/vendor/pygments/lexers/_php_builtins.py +4756 -0
  617. wandb/vendor/pygments/lexers/_postgres_builtins.py +621 -0
  618. wandb/vendor/pygments/lexers/_scilab_builtins.py +3094 -0
  619. wandb/vendor/pygments/lexers/_sourcemod_builtins.py +1163 -0
  620. wandb/vendor/pygments/lexers/_stan_builtins.py +532 -0
  621. wandb/vendor/pygments/lexers/_stata_builtins.py +419 -0
  622. wandb/vendor/pygments/lexers/_tsql_builtins.py +1004 -0
  623. wandb/vendor/pygments/lexers/_vim_builtins.py +1939 -0
  624. wandb/vendor/pygments/lexers/actionscript.py +240 -0
  625. wandb/vendor/pygments/lexers/agile.py +24 -0
  626. wandb/vendor/pygments/lexers/algebra.py +221 -0
  627. wandb/vendor/pygments/lexers/ambient.py +76 -0
  628. wandb/vendor/pygments/lexers/ampl.py +87 -0
  629. wandb/vendor/pygments/lexers/apl.py +101 -0
  630. wandb/vendor/pygments/lexers/archetype.py +318 -0
  631. wandb/vendor/pygments/lexers/asm.py +641 -0
  632. wandb/vendor/pygments/lexers/automation.py +374 -0
  633. wandb/vendor/pygments/lexers/basic.py +500 -0
  634. wandb/vendor/pygments/lexers/bibtex.py +160 -0
  635. wandb/vendor/pygments/lexers/business.py +612 -0
  636. wandb/vendor/pygments/lexers/c_cpp.py +252 -0
  637. wandb/vendor/pygments/lexers/c_like.py +541 -0
  638. wandb/vendor/pygments/lexers/capnproto.py +78 -0
  639. wandb/vendor/pygments/lexers/chapel.py +102 -0
  640. wandb/vendor/pygments/lexers/clean.py +288 -0
  641. wandb/vendor/pygments/lexers/compiled.py +34 -0
  642. wandb/vendor/pygments/lexers/configs.py +833 -0
  643. wandb/vendor/pygments/lexers/console.py +114 -0
  644. wandb/vendor/pygments/lexers/crystal.py +393 -0
  645. wandb/vendor/pygments/lexers/csound.py +366 -0
  646. wandb/vendor/pygments/lexers/css.py +689 -0
  647. wandb/vendor/pygments/lexers/d.py +251 -0
  648. wandb/vendor/pygments/lexers/dalvik.py +125 -0
  649. wandb/vendor/pygments/lexers/data.py +555 -0
  650. wandb/vendor/pygments/lexers/diff.py +165 -0
  651. wandb/vendor/pygments/lexers/dotnet.py +691 -0
  652. wandb/vendor/pygments/lexers/dsls.py +878 -0
  653. wandb/vendor/pygments/lexers/dylan.py +289 -0
  654. wandb/vendor/pygments/lexers/ecl.py +125 -0
  655. wandb/vendor/pygments/lexers/eiffel.py +65 -0
  656. wandb/vendor/pygments/lexers/elm.py +121 -0
  657. wandb/vendor/pygments/lexers/erlang.py +533 -0
  658. wandb/vendor/pygments/lexers/esoteric.py +277 -0
  659. wandb/vendor/pygments/lexers/ezhil.py +69 -0
  660. wandb/vendor/pygments/lexers/factor.py +344 -0
  661. wandb/vendor/pygments/lexers/fantom.py +250 -0
  662. wandb/vendor/pygments/lexers/felix.py +273 -0
  663. wandb/vendor/pygments/lexers/forth.py +177 -0
  664. wandb/vendor/pygments/lexers/fortran.py +205 -0
  665. wandb/vendor/pygments/lexers/foxpro.py +428 -0
  666. wandb/vendor/pygments/lexers/functional.py +21 -0
  667. wandb/vendor/pygments/lexers/go.py +101 -0
  668. wandb/vendor/pygments/lexers/grammar_notation.py +213 -0
  669. wandb/vendor/pygments/lexers/graph.py +80 -0
  670. wandb/vendor/pygments/lexers/graphics.py +553 -0
  671. wandb/vendor/pygments/lexers/haskell.py +843 -0
  672. wandb/vendor/pygments/lexers/haxe.py +936 -0
  673. wandb/vendor/pygments/lexers/hdl.py +382 -0
  674. wandb/vendor/pygments/lexers/hexdump.py +103 -0
  675. wandb/vendor/pygments/lexers/html.py +602 -0
  676. wandb/vendor/pygments/lexers/idl.py +270 -0
  677. wandb/vendor/pygments/lexers/igor.py +288 -0
  678. wandb/vendor/pygments/lexers/inferno.py +96 -0
  679. wandb/vendor/pygments/lexers/installers.py +322 -0
  680. wandb/vendor/pygments/lexers/int_fiction.py +1343 -0
  681. wandb/vendor/pygments/lexers/iolang.py +63 -0
  682. wandb/vendor/pygments/lexers/j.py +146 -0
  683. wandb/vendor/pygments/lexers/javascript.py +1525 -0
  684. wandb/vendor/pygments/lexers/julia.py +333 -0
  685. wandb/vendor/pygments/lexers/jvm.py +1573 -0
  686. wandb/vendor/pygments/lexers/lisp.py +2621 -0
  687. wandb/vendor/pygments/lexers/make.py +202 -0
  688. wandb/vendor/pygments/lexers/markup.py +595 -0
  689. wandb/vendor/pygments/lexers/math.py +21 -0
  690. wandb/vendor/pygments/lexers/matlab.py +663 -0
  691. wandb/vendor/pygments/lexers/ml.py +769 -0
  692. wandb/vendor/pygments/lexers/modeling.py +358 -0
  693. wandb/vendor/pygments/lexers/modula2.py +1561 -0
  694. wandb/vendor/pygments/lexers/monte.py +204 -0
  695. wandb/vendor/pygments/lexers/ncl.py +894 -0
  696. wandb/vendor/pygments/lexers/nimrod.py +159 -0
  697. wandb/vendor/pygments/lexers/nit.py +64 -0
  698. wandb/vendor/pygments/lexers/nix.py +136 -0
  699. wandb/vendor/pygments/lexers/oberon.py +105 -0
  700. wandb/vendor/pygments/lexers/objective.py +504 -0
  701. wandb/vendor/pygments/lexers/ooc.py +85 -0
  702. wandb/vendor/pygments/lexers/other.py +41 -0
  703. wandb/vendor/pygments/lexers/parasail.py +79 -0
  704. wandb/vendor/pygments/lexers/parsers.py +835 -0
  705. wandb/vendor/pygments/lexers/pascal.py +644 -0
  706. wandb/vendor/pygments/lexers/pawn.py +199 -0
  707. wandb/vendor/pygments/lexers/perl.py +620 -0
  708. wandb/vendor/pygments/lexers/php.py +267 -0
  709. wandb/vendor/pygments/lexers/praat.py +294 -0
  710. wandb/vendor/pygments/lexers/prolog.py +306 -0
  711. wandb/vendor/pygments/lexers/python.py +939 -0
  712. wandb/vendor/pygments/lexers/qvt.py +152 -0
  713. wandb/vendor/pygments/lexers/r.py +453 -0
  714. wandb/vendor/pygments/lexers/rdf.py +270 -0
  715. wandb/vendor/pygments/lexers/rebol.py +431 -0
  716. wandb/vendor/pygments/lexers/resource.py +85 -0
  717. wandb/vendor/pygments/lexers/rnc.py +67 -0
  718. wandb/vendor/pygments/lexers/roboconf.py +82 -0
  719. wandb/vendor/pygments/lexers/robotframework.py +560 -0
  720. wandb/vendor/pygments/lexers/ruby.py +519 -0
  721. wandb/vendor/pygments/lexers/rust.py +220 -0
  722. wandb/vendor/pygments/lexers/sas.py +228 -0
  723. wandb/vendor/pygments/lexers/scripting.py +1222 -0
  724. wandb/vendor/pygments/lexers/shell.py +794 -0
  725. wandb/vendor/pygments/lexers/smalltalk.py +195 -0
  726. wandb/vendor/pygments/lexers/smv.py +79 -0
  727. wandb/vendor/pygments/lexers/snobol.py +83 -0
  728. wandb/vendor/pygments/lexers/special.py +103 -0
  729. wandb/vendor/pygments/lexers/sql.py +681 -0
  730. wandb/vendor/pygments/lexers/stata.py +108 -0
  731. wandb/vendor/pygments/lexers/supercollider.py +90 -0
  732. wandb/vendor/pygments/lexers/tcl.py +145 -0
  733. wandb/vendor/pygments/lexers/templates.py +2283 -0
  734. wandb/vendor/pygments/lexers/testing.py +207 -0
  735. wandb/vendor/pygments/lexers/text.py +25 -0
  736. wandb/vendor/pygments/lexers/textedit.py +169 -0
  737. wandb/vendor/pygments/lexers/textfmts.py +297 -0
  738. wandb/vendor/pygments/lexers/theorem.py +458 -0
  739. wandb/vendor/pygments/lexers/trafficscript.py +54 -0
  740. wandb/vendor/pygments/lexers/typoscript.py +226 -0
  741. wandb/vendor/pygments/lexers/urbi.py +133 -0
  742. wandb/vendor/pygments/lexers/varnish.py +190 -0
  743. wandb/vendor/pygments/lexers/verification.py +111 -0
  744. wandb/vendor/pygments/lexers/web.py +24 -0
  745. wandb/vendor/pygments/lexers/webmisc.py +988 -0
  746. wandb/vendor/pygments/lexers/whiley.py +116 -0
  747. wandb/vendor/pygments/lexers/x10.py +69 -0
  748. wandb/vendor/pygments/modeline.py +44 -0
  749. wandb/vendor/pygments/plugin.py +68 -0
  750. wandb/vendor/pygments/regexopt.py +92 -0
  751. wandb/vendor/pygments/scanner.py +105 -0
  752. wandb/vendor/pygments/sphinxext.py +158 -0
  753. wandb/vendor/pygments/style.py +155 -0
  754. wandb/vendor/pygments/styles/__init__.py +80 -0
  755. wandb/vendor/pygments/styles/abap.py +29 -0
  756. wandb/vendor/pygments/styles/algol.py +63 -0
  757. wandb/vendor/pygments/styles/algol_nu.py +63 -0
  758. wandb/vendor/pygments/styles/arduino.py +98 -0
  759. wandb/vendor/pygments/styles/autumn.py +65 -0
  760. wandb/vendor/pygments/styles/borland.py +51 -0
  761. wandb/vendor/pygments/styles/bw.py +49 -0
  762. wandb/vendor/pygments/styles/colorful.py +81 -0
  763. wandb/vendor/pygments/styles/default.py +73 -0
  764. wandb/vendor/pygments/styles/emacs.py +72 -0
  765. wandb/vendor/pygments/styles/friendly.py +72 -0
  766. wandb/vendor/pygments/styles/fruity.py +42 -0
  767. wandb/vendor/pygments/styles/igor.py +29 -0
  768. wandb/vendor/pygments/styles/lovelace.py +97 -0
  769. wandb/vendor/pygments/styles/manni.py +75 -0
  770. wandb/vendor/pygments/styles/monokai.py +106 -0
  771. wandb/vendor/pygments/styles/murphy.py +80 -0
  772. wandb/vendor/pygments/styles/native.py +65 -0
  773. wandb/vendor/pygments/styles/paraiso_dark.py +125 -0
  774. wandb/vendor/pygments/styles/paraiso_light.py +125 -0
  775. wandb/vendor/pygments/styles/pastie.py +75 -0
  776. wandb/vendor/pygments/styles/perldoc.py +69 -0
  777. wandb/vendor/pygments/styles/rainbow_dash.py +89 -0
  778. wandb/vendor/pygments/styles/rrt.py +33 -0
  779. wandb/vendor/pygments/styles/sas.py +44 -0
  780. wandb/vendor/pygments/styles/stata.py +40 -0
  781. wandb/vendor/pygments/styles/tango.py +141 -0
  782. wandb/vendor/pygments/styles/trac.py +63 -0
  783. wandb/vendor/pygments/styles/vim.py +63 -0
  784. wandb/vendor/pygments/styles/vs.py +38 -0
  785. wandb/vendor/pygments/styles/xcode.py +51 -0
  786. wandb/vendor/pygments/token.py +213 -0
  787. wandb/vendor/pygments/unistring.py +217 -0
  788. wandb/vendor/pygments/util.py +388 -0
  789. wandb/vendor/pynvml/__init__.py +0 -0
  790. wandb/vendor/pynvml/pynvml.py +4779 -0
  791. wandb/vendor/watchdog_0_9_0/wandb_watchdog/__init__.py +17 -0
  792. wandb/vendor/watchdog_0_9_0/wandb_watchdog/events.py +615 -0
  793. wandb/vendor/watchdog_0_9_0/wandb_watchdog/observers/__init__.py +98 -0
  794. wandb/vendor/watchdog_0_9_0/wandb_watchdog/observers/api.py +369 -0
  795. wandb/vendor/watchdog_0_9_0/wandb_watchdog/observers/fsevents.py +172 -0
  796. wandb/vendor/watchdog_0_9_0/wandb_watchdog/observers/fsevents2.py +239 -0
  797. wandb/vendor/watchdog_0_9_0/wandb_watchdog/observers/inotify.py +218 -0
  798. wandb/vendor/watchdog_0_9_0/wandb_watchdog/observers/inotify_buffer.py +81 -0
  799. wandb/vendor/watchdog_0_9_0/wandb_watchdog/observers/inotify_c.py +575 -0
  800. wandb/vendor/watchdog_0_9_0/wandb_watchdog/observers/kqueue.py +730 -0
  801. wandb/vendor/watchdog_0_9_0/wandb_watchdog/observers/polling.py +145 -0
  802. wandb/vendor/watchdog_0_9_0/wandb_watchdog/observers/read_directory_changes.py +133 -0
  803. wandb/vendor/watchdog_0_9_0/wandb_watchdog/observers/winapi.py +348 -0
  804. wandb/vendor/watchdog_0_9_0/wandb_watchdog/patterns.py +265 -0
  805. wandb/vendor/watchdog_0_9_0/wandb_watchdog/tricks/__init__.py +174 -0
  806. wandb/vendor/watchdog_0_9_0/wandb_watchdog/utils/__init__.py +151 -0
  807. wandb/vendor/watchdog_0_9_0/wandb_watchdog/utils/bricks.py +249 -0
  808. wandb/vendor/watchdog_0_9_0/wandb_watchdog/utils/compat.py +29 -0
  809. wandb/vendor/watchdog_0_9_0/wandb_watchdog/utils/decorators.py +198 -0
  810. wandb/vendor/watchdog_0_9_0/wandb_watchdog/utils/delayed_queue.py +88 -0
  811. wandb/vendor/watchdog_0_9_0/wandb_watchdog/utils/dirsnapshot.py +293 -0
  812. wandb/vendor/watchdog_0_9_0/wandb_watchdog/utils/echo.py +157 -0
  813. wandb/vendor/watchdog_0_9_0/wandb_watchdog/utils/event_backport.py +41 -0
  814. wandb/vendor/watchdog_0_9_0/wandb_watchdog/utils/importlib2.py +40 -0
  815. wandb/vendor/watchdog_0_9_0/wandb_watchdog/utils/platform.py +57 -0
  816. wandb/vendor/watchdog_0_9_0/wandb_watchdog/utils/unicode_paths.py +64 -0
  817. wandb/vendor/watchdog_0_9_0/wandb_watchdog/utils/win32stat.py +123 -0
  818. wandb/vendor/watchdog_0_9_0/wandb_watchdog/version.py +28 -0
  819. wandb/vendor/watchdog_0_9_0/wandb_watchdog/watchmedo.py +577 -0
  820. wandb/wandb_agent.py +588 -0
  821. wandb/wandb_controller.py +721 -0
  822. wandb/wandb_run.py +9 -0
  823. wandb-0.18.2.dist-info/METADATA +213 -0
  824. wandb-0.18.2.dist-info/RECORD +827 -0
  825. wandb-0.18.2.dist-info/WHEEL +5 -0
  826. wandb-0.18.2.dist-info/entry_points.txt +3 -0
  827. wandb-0.18.2.dist-info/licenses/LICENSE +21 -0
@@ -0,0 +1,1603 @@
1
+ """Tooling for the W&B Importer."""
2
+
3
+ import itertools
4
+ import json
5
+ import logging
6
+ import numbers
7
+ import os
8
+ import re
9
+ import shutil
10
+ from dataclasses import dataclass, field
11
+ from datetime import datetime as dt
12
+ from pathlib import Path
13
+ from typing import Any, Dict, Iterable, Iterator, List, Optional, Tuple
14
+ from unittest.mock import patch
15
+
16
+ import filelock
17
+ import polars as pl
18
+ import requests
19
+ import urllib3
20
+ import wandb_workspaces.reports.v1 as wr
21
+ import yaml
22
+ from wandb_gql import gql
23
+ from wandb_workspaces.reports.v1 import Report
24
+
25
+ import wandb
26
+ from wandb.apis.public import ArtifactCollection, Run
27
+ from wandb.apis.public.files import File
28
+ from wandb.util import coalesce, remove_keys_with_none_values
29
+
30
+ from . import validation
31
+ from .internals import internal
32
+ from .internals.protocols import PathStr, Policy
33
+ from .internals.util import Namespace, for_each
34
+
35
+ Artifact = wandb.Artifact
36
+ Api = wandb.Api
37
+ Project = wandb.apis.public.Project
38
+
39
+ ARTIFACT_ERRORS_FNAME = "artifact_errors.jsonl"
40
+ ARTIFACT_SUCCESSES_FNAME = "artifact_successes.jsonl"
41
+ RUN_ERRORS_FNAME = "run_errors.jsonl"
42
+ RUN_SUCCESSES_FNAME = "run_successes.jsonl"
43
+
44
+ ART_SEQUENCE_DUMMY_PLACEHOLDER = "__ART_SEQUENCE_DUMMY_PLACEHOLDER__"
45
+ RUN_DUMMY_PLACEHOLDER = "__RUN_DUMMY_PLACEHOLDER__"
46
+ ART_DUMMY_PLACEHOLDER_PATH = "__importer_temp__"
47
+ ART_DUMMY_PLACEHOLDER_TYPE = "__temp__"
48
+
49
+ SRC_ART_PATH = "./artifacts/src"
50
+ DST_ART_PATH = "./artifacts/dst"
51
+
52
+
53
+ logger = logging.getLogger(__name__)
54
+ logger.setLevel(logging.INFO)
55
+
56
+ if os.getenv("WANDB_IMPORTER_ENABLE_RICH_LOGGING"):
57
+ from rich.logging import RichHandler
58
+
59
+ logger.addHandler(RichHandler(rich_tracebacks=True, tracebacks_show_locals=True))
60
+ else:
61
+ console_handler = logging.StreamHandler()
62
+ console_handler.setLevel(logging.INFO)
63
+
64
+ formatter = logging.Formatter("%(asctime)s - %(levelname)s - %(message)s")
65
+ console_handler.setFormatter(formatter)
66
+
67
+ logger.addHandler(console_handler)
68
+
69
+
70
+ @dataclass
71
+ class ArtifactSequence:
72
+ artifacts: Iterable[wandb.Artifact]
73
+ entity: str
74
+ project: str
75
+ type_: str
76
+ name: str
77
+
78
+ def __iter__(self) -> Iterator:
79
+ return iter(self.artifacts)
80
+
81
+ def __repr__(self) -> str:
82
+ return f"ArtifactSequence({self.identifier})"
83
+
84
+ @property
85
+ def identifier(self) -> str:
86
+ return "/".join([self.entity, self.project, self.type_, self.name])
87
+
88
+ @classmethod
89
+ def from_collection(cls, collection: ArtifactCollection):
90
+ arts = collection.artifacts()
91
+ arts = sorted(arts, key=lambda a: int(a.version.lstrip("v")))
92
+ return ArtifactSequence(
93
+ arts,
94
+ collection.entity,
95
+ collection.project,
96
+ collection.type,
97
+ collection.name,
98
+ )
99
+
100
+
101
+ class WandbRun:
102
+ def __init__(
103
+ self,
104
+ run: Run,
105
+ *,
106
+ src_base_url: str,
107
+ src_api_key: str,
108
+ dst_base_url: str,
109
+ dst_api_key: str,
110
+ ) -> None:
111
+ self.run = run
112
+ self.api = wandb.Api(
113
+ api_key=src_api_key,
114
+ overrides={"base_url": src_base_url},
115
+ )
116
+ self.dst_api = wandb.Api(
117
+ api_key=dst_api_key,
118
+ overrides={"base_url": dst_base_url},
119
+ )
120
+
121
+ # For caching
122
+ self._files: Optional[Iterable[Tuple[str, str]]] = None
123
+ self._artifacts: Optional[Iterable[Artifact]] = None
124
+ self._used_artifacts: Optional[Iterable[Artifact]] = None
125
+ self._parquet_history_paths: Optional[Iterable[str]] = None
126
+
127
+ def __repr__(self) -> str:
128
+ s = os.path.join(self.entity(), self.project(), self.run_id())
129
+ return f"WandbRun({s})"
130
+
131
+ def run_id(self) -> str:
132
+ return self.run.id
133
+
134
+ def entity(self) -> str:
135
+ return self.run.entity
136
+
137
+ def project(self) -> str:
138
+ return self.run.project
139
+
140
+ def config(self) -> Dict[str, Any]:
141
+ return self.run.config
142
+
143
+ def summary(self) -> Dict[str, float]:
144
+ s = self.run.summary
145
+ return s
146
+
147
+ def metrics(self) -> Iterable[Dict[str, float]]:
148
+ if self._parquet_history_paths is None:
149
+ self._parquet_history_paths = list(self._get_parquet_history_paths())
150
+
151
+ if self._parquet_history_paths:
152
+ rows = self._get_rows_from_parquet_history_paths()
153
+ else:
154
+ logger.warn(
155
+ "No parquet files detected; using scan history (this may not be reliable)"
156
+ )
157
+ rows = self.run.scan_history()
158
+
159
+ for row in rows:
160
+ row = remove_keys_with_none_values(row)
161
+ yield row
162
+
163
+ def run_group(self) -> Optional[str]:
164
+ return self.run.group
165
+
166
+ def job_type(self) -> Optional[str]:
167
+ return self.run.job_type
168
+
169
+ def display_name(self) -> str:
170
+ return self.run.display_name
171
+
172
+ def notes(self) -> Optional[str]:
173
+ # Notes includes the previous notes and serves as a catch-all for things we missed or can't add back
174
+ previous_link = f"Imported from: {self.run.url}"
175
+ previous_author = f"Author: {self.run.user.username}"
176
+
177
+ header = [previous_link, previous_author]
178
+ previous_notes = self.run.notes or ""
179
+
180
+ return "\n".join(header) + "\n---\n" + previous_notes
181
+
182
+ def tags(self) -> Optional[List[str]]:
183
+ return self.run.tags
184
+
185
+ def artifacts(self) -> Optional[Iterable[Artifact]]:
186
+ if self._artifacts is None:
187
+ _artifacts = []
188
+ for art in self.run.logged_artifacts():
189
+ a = _clone_art(art)
190
+ _artifacts.append(a)
191
+ self._artifacts = _artifacts
192
+
193
+ yield from self._artifacts
194
+
195
+ def used_artifacts(self) -> Optional[Iterable[Artifact]]:
196
+ if self._used_artifacts is None:
197
+ _used_artifacts = []
198
+ for art in self.run.used_artifacts():
199
+ a = _clone_art(art)
200
+ _used_artifacts.append(a)
201
+ self._used_artifacts = _used_artifacts
202
+
203
+ yield from self._used_artifacts
204
+
205
+ def os_version(self) -> Optional[str]: ... # pragma: no cover
206
+
207
+ def python_version(self) -> Optional[str]:
208
+ return self._metadata_file().get("python")
209
+
210
+ def cuda_version(self) -> Optional[str]: ... # pragma: no cover
211
+
212
+ def program(self) -> Optional[str]: ... # pragma: no cover
213
+
214
+ def host(self) -> Optional[str]:
215
+ return self._metadata_file().get("host")
216
+
217
+ def username(self) -> Optional[str]: ... # pragma: no cover
218
+
219
+ def executable(self) -> Optional[str]: ... # pragma: no cover
220
+
221
+ def gpus_used(self) -> Optional[str]: ... # pragma: no cover
222
+
223
+ def cpus_used(self) -> Optional[int]: # can we get the model?
224
+ ... # pragma: no cover
225
+
226
+ def memory_used(self) -> Optional[int]: ... # pragma: no cover
227
+
228
+ def runtime(self) -> Optional[int]:
229
+ wandb_runtime = self.run.summary.get("_wandb", {}).get("runtime")
230
+ base_runtime = self.run.summary.get("_runtime")
231
+
232
+ if (t := coalesce(wandb_runtime, base_runtime)) is None:
233
+ return t
234
+ return int(t)
235
+
236
+ def start_time(self) -> Optional[int]:
237
+ t = dt.fromisoformat(self.run.created_at).timestamp() * 1000
238
+ return int(t)
239
+
240
+ def code_path(self) -> Optional[str]:
241
+ path = self._metadata_file().get("codePath", "")
242
+ return f"code/{path}"
243
+
244
+ def cli_version(self) -> Optional[str]:
245
+ return self._config_file().get("_wandb", {}).get("value", {}).get("cli_version")
246
+
247
+ def files(self) -> Optional[Iterable[Tuple[PathStr, Policy]]]:
248
+ if self._files is None:
249
+ files_dir = f"{internal.ROOT_DIR}/{self.run_id()}/files"
250
+ _files = []
251
+ for f in self.run.files():
252
+ f: File
253
+ # These optimizations are intended to avoid rate limiting when importing many runs in parallel
254
+ # Don't carry over empty files
255
+ if f.size == 0:
256
+ continue
257
+ # Skip deadlist to avoid overloading S3
258
+ if "wandb_manifest.json.deadlist" in f.name:
259
+ continue
260
+
261
+ result = f.download(files_dir, exist_ok=True, api=self.api)
262
+ file_and_policy = (result.name, "end")
263
+ _files.append(file_and_policy)
264
+ self._files = _files
265
+
266
+ yield from self._files
267
+
268
+ def logs(self) -> Optional[Iterable[str]]:
269
+ log_files = self._find_all_in_files_regex(r"^.*output\.log$")
270
+ for path in log_files:
271
+ with open(path) as f:
272
+ yield from f.readlines()
273
+
274
+ def _metadata_file(self) -> Dict[str, Any]:
275
+ if (fname := self._find_in_files("wandb-metadata.json")) is None:
276
+ return {}
277
+
278
+ with open(fname) as f:
279
+ return json.loads(f.read())
280
+
281
+ def _config_file(self) -> Dict[str, Any]:
282
+ if (fname := self._find_in_files("config.yaml")) is None:
283
+ return {}
284
+
285
+ with open(fname) as f:
286
+ return yaml.safe_load(f) or {}
287
+
288
+ def _get_rows_from_parquet_history_paths(self) -> Iterable[Dict[str, Any]]:
289
+ # Unfortunately, it's not feasible to validate non-parquet history
290
+ if not (paths := self._get_parquet_history_paths()):
291
+ yield {}
292
+ return
293
+
294
+ # Collect and merge parquet history
295
+ dfs = [
296
+ pl.read_parquet(p) for path in paths for p in Path(path).glob("*.parquet")
297
+ ]
298
+ if "_step" in (df := _merge_dfs(dfs)):
299
+ df = df.with_columns(pl.col("_step").cast(pl.Int64))
300
+ yield from df.iter_rows(named=True)
301
+
302
+ def _get_parquet_history_paths(self) -> Iterable[str]:
303
+ if self._parquet_history_paths is None:
304
+ paths = []
305
+ # self.artifacts() returns a copy of the artifacts; use this to get raw
306
+ for art in self.run.logged_artifacts():
307
+ if art.type != "wandb-history":
308
+ continue
309
+ if (
310
+ path := _download_art(art, root=f"{SRC_ART_PATH}/{art.name}")
311
+ ) is None:
312
+ continue
313
+ paths.append(path)
314
+ self._parquet_history_paths = paths
315
+
316
+ yield from self._parquet_history_paths
317
+
318
+ def _find_in_files(self, name: str) -> Optional[str]:
319
+ if files := self.files():
320
+ for path, _ in files:
321
+ if name in path:
322
+ return path
323
+ return None
324
+
325
+ def _find_all_in_files_regex(self, regex: str) -> Iterable[str]:
326
+ if files := self.files():
327
+ for path, _ in files:
328
+ if re.match(regex, path):
329
+ yield path
330
+
331
+
332
+ class WandbImporter:
333
+ """Transfers runs, reports, and artifact sequences between W&B instances."""
334
+
335
+ def __init__(
336
+ self,
337
+ src_base_url: str,
338
+ src_api_key: str,
339
+ dst_base_url: str,
340
+ dst_api_key: str,
341
+ *,
342
+ custom_api_kwargs: Optional[Dict[str, Any]] = None,
343
+ ) -> None:
344
+ self.src_base_url = src_base_url
345
+ self.src_api_key = src_api_key
346
+ self.dst_base_url = dst_base_url
347
+ self.dst_api_key = dst_api_key
348
+
349
+ if custom_api_kwargs is None:
350
+ custom_api_kwargs = {"timeout": 600}
351
+
352
+ self.src_api = wandb.Api(
353
+ api_key=src_api_key,
354
+ overrides={"base_url": src_base_url},
355
+ **custom_api_kwargs,
356
+ )
357
+ self.dst_api = wandb.Api(
358
+ api_key=dst_api_key,
359
+ overrides={"base_url": dst_base_url},
360
+ **custom_api_kwargs,
361
+ )
362
+
363
+ self.run_api_kwargs = {
364
+ "src_base_url": src_base_url,
365
+ "src_api_key": src_api_key,
366
+ "dst_base_url": dst_base_url,
367
+ "dst_api_key": dst_api_key,
368
+ }
369
+
370
+ def __repr__(self):
371
+ return f"<WandbImporter src={self.src_base_url}, dst={self.dst_base_url}>" # pragma: no cover
372
+
373
+ def _import_run(
374
+ self,
375
+ run: WandbRun,
376
+ *,
377
+ namespace: Optional[Namespace] = None,
378
+ config: Optional[internal.SendManagerConfig] = None,
379
+ ) -> None:
380
+ """Import one WandbRun.
381
+
382
+ Use `namespace` to specify alternate settings like where the run should be uploaded
383
+ """
384
+ if namespace is None:
385
+ namespace = Namespace(run.entity(), run.project())
386
+
387
+ if config is None:
388
+ config = internal.SendManagerConfig(
389
+ metadata=True,
390
+ files=True,
391
+ media=True,
392
+ code=True,
393
+ history=True,
394
+ summary=True,
395
+ terminal_output=True,
396
+ )
397
+
398
+ settings_override = {
399
+ "api_key": self.dst_api_key,
400
+ "base_url": self.dst_base_url,
401
+ "resume": "true",
402
+ "resumed": True,
403
+ }
404
+
405
+ # Send run with base config
406
+ logger.debug(f"Importing run, {run=}")
407
+ internal.send_run(
408
+ run,
409
+ overrides=namespace.send_manager_overrides,
410
+ settings_override=settings_override,
411
+ config=config,
412
+ )
413
+
414
+ if config.history:
415
+ # Send run again with history artifacts in case config history=True, artifacts=False
416
+ # The history artifact must come with the actual history data
417
+
418
+ logger.debug(f"Collecting history artifacts, {run=}")
419
+ history_arts = []
420
+ for art in run.run.logged_artifacts():
421
+ if art.type != "wandb-history":
422
+ continue
423
+ logger.debug(f"Collecting history artifact {art.name=}")
424
+ new_art = _clone_art(art)
425
+ history_arts.append(new_art)
426
+
427
+ logger.debug(f"Importing history artifacts, {run=}")
428
+ internal.send_run(
429
+ run,
430
+ extra_arts=history_arts,
431
+ overrides=namespace.send_manager_overrides,
432
+ settings_override=settings_override,
433
+ config=config,
434
+ )
435
+
436
+ def _delete_collection_in_dst(
437
+ self,
438
+ seq: ArtifactSequence,
439
+ namespace: Optional[Namespace] = None,
440
+ ):
441
+ """Deletes the equivalent artifact collection in destination.
442
+
443
+ Intended to clear the destination when an uploaded artifact does not pass validation.
444
+ """
445
+ entity = coalesce(namespace.entity, seq.entity)
446
+ project = coalesce(namespace.project, seq.project)
447
+ art_type = f"{entity}/{project}/{seq.type_}"
448
+ art_name = seq.name
449
+
450
+ logger.info(
451
+ f"Deleting collection {entity=}, {project=}, {art_type=}, {art_name=}"
452
+ )
453
+ try:
454
+ dst_collection = self.dst_api.artifact_collection(art_type, art_name)
455
+ except (wandb.CommError, ValueError):
456
+ logger.warn(f"Collection doesn't exist {art_type=}, {art_name=}")
457
+ return
458
+
459
+ try:
460
+ dst_collection.delete()
461
+ except (wandb.CommError, ValueError) as e:
462
+ logger.warn(f"Collection can't be deleted, {art_type=}, {art_name=}, {e=}")
463
+ return
464
+
465
+ def _import_artifact_sequence(
466
+ self,
467
+ seq: ArtifactSequence,
468
+ *,
469
+ namespace: Optional[Namespace] = None,
470
+ ) -> None:
471
+ """Import one artifact sequence.
472
+
473
+ Use `namespace` to specify alternate settings like where the artifact sequence should be uploaded
474
+ """
475
+ if not seq.artifacts:
476
+ # The artifact sequence has no versions. This usually means all artifacts versions were deleted intentionally,
477
+ # but it can also happen if the sequence represents run history and that run was deleted.
478
+ logger.warn(f"Artifact {seq=} has no artifacts, skipping.")
479
+ return
480
+
481
+ if namespace is None:
482
+ namespace = Namespace(seq.entity, seq.project)
483
+
484
+ settings_override = {
485
+ "api_key": self.dst_api_key,
486
+ "base_url": self.dst_base_url,
487
+ "resume": "true",
488
+ "resumed": True,
489
+ }
490
+
491
+ send_manager_config = internal.SendManagerConfig(log_artifacts=True)
492
+
493
+ # Delete any existing artifact sequence, otherwise versions will be out of order
494
+ # Unfortunately, you can't delete only part of the sequence because versions are "remembered" even after deletion
495
+ self._delete_collection_in_dst(seq, namespace)
496
+
497
+ # Get a placeholder run for dummy artifacts we'll upload later
498
+ art = seq.artifacts[0]
499
+ run_or_dummy: Optional[Run] = _get_run_or_dummy_from_art(art, self.src_api)
500
+
501
+ # Each `group_of_artifacts` is either:
502
+ # 1. A single "real" artifact in a list; or
503
+ # 2. A list of dummy artifacts that are uploaded together.
504
+ # This guarantees the real artifacts have the correct version numbers while allowing for parallel upload of dummies.
505
+ groups_of_artifacts = list(_make_groups_of_artifacts(seq))
506
+ for i, group in enumerate(groups_of_artifacts, 1):
507
+ art = group[0]
508
+ if art.description == ART_SEQUENCE_DUMMY_PLACEHOLDER:
509
+ run = WandbRun(run_or_dummy, **self.run_api_kwargs)
510
+ else:
511
+ try:
512
+ wandb_run = art.logged_by()
513
+ except ValueError:
514
+ # The run used to exist but has since been deleted
515
+ # wandb_run = None
516
+ pass
517
+
518
+ # Could be logged by None (rare) or ValueError
519
+ if wandb_run is None:
520
+ logger.warn(
521
+ f"Run for {art.name=} does not exist (deleted?), using {run_or_dummy=}"
522
+ )
523
+ wandb_run = run_or_dummy
524
+
525
+ new_art = _clone_art(art)
526
+ group = [new_art]
527
+ run = WandbRun(wandb_run, **self.run_api_kwargs)
528
+
529
+ logger.info(
530
+ f"Uploading partial artifact {seq=}, {i}/{len(groups_of_artifacts)}"
531
+ )
532
+ internal.send_run(
533
+ run,
534
+ extra_arts=group,
535
+ overrides=namespace.send_manager_overrides,
536
+ settings_override=settings_override,
537
+ config=send_manager_config,
538
+ )
539
+ logger.info(f"Finished uploading {seq=}")
540
+
541
+ # query it back and remove placeholders
542
+ self._remove_placeholders(seq)
543
+
544
+ def _remove_placeholders(self, seq: ArtifactSequence) -> None:
545
+ try:
546
+ retry_arts_func = internal.exp_retry(self._dst_api.artifacts)
547
+ dst_arts = list(retry_arts_func(seq.type_, seq.name))
548
+ except wandb.CommError:
549
+ logger.warn(f"{seq=} does not exist in dst. Has it already been deleted?")
550
+ return
551
+ except TypeError as e:
552
+ logger.error(f"Problem getting dst versions (try again later) {e=}")
553
+ return
554
+
555
+ for art in dst_arts:
556
+ if art.description != ART_SEQUENCE_DUMMY_PLACEHOLDER:
557
+ continue
558
+ if art.type in ("wandb-history", "job"):
559
+ continue
560
+
561
+ try:
562
+ art.delete(delete_aliases=True)
563
+ except wandb.CommError as e:
564
+ if "cannot delete system managed artifact" in str(e):
565
+ logger.warn("Cannot delete system managed artifact")
566
+ else:
567
+ raise e
568
+
569
+ def _get_dst_art(
570
+ self, src_art: Run, entity: Optional[str] = None, project: Optional[str] = None
571
+ ) -> Artifact:
572
+ entity = coalesce(entity, src_art.entity)
573
+ project = coalesce(project, src_art.project)
574
+ name = src_art.name
575
+
576
+ return self.dst_api.artifact(f"{entity}/{project}/{name}")
577
+
578
+ def _get_run_problems(
579
+ self, src_run: Run, dst_run: Run, force_retry: bool = False
580
+ ) -> List[dict]:
581
+ problems = []
582
+
583
+ if force_retry:
584
+ problems.append("__force_retry__")
585
+
586
+ if non_matching_metadata := self._compare_run_metadata(src_run, dst_run):
587
+ problems.append("metadata:" + str(non_matching_metadata))
588
+
589
+ if non_matching_summary := self._compare_run_summary(src_run, dst_run):
590
+ problems.append("summary:" + str(non_matching_summary))
591
+
592
+ # TODO: Compare files?
593
+
594
+ return problems
595
+
596
+ def _compare_run_metadata(self, src_run: Run, dst_run: Run) -> dict:
597
+ fname = "wandb-metadata.json"
598
+ # problems = {}
599
+
600
+ src_f = src_run.file(fname)
601
+ if src_f.size == 0:
602
+ # the src was corrupted so no comparisons here will ever work
603
+ return {}
604
+
605
+ dst_f = dst_run.file(fname)
606
+ try:
607
+ contents = wandb.util.download_file_into_memory(
608
+ dst_f.url, self.dst_api.api_key
609
+ )
610
+ except urllib3.exceptions.ReadTimeoutError:
611
+ return {"Error checking": "Timeout"}
612
+ except requests.HTTPError as e:
613
+ if e.response.status_code == 404:
614
+ return {"Bad upload": f"File not found: {fname}"}
615
+ return {"http problem": f"{fname}: ({e})"}
616
+
617
+ dst_meta = wandb.wandb_sdk.lib.json_util.loads(contents)
618
+
619
+ non_matching = {}
620
+ if src_run.metadata:
621
+ for k, src_v in src_run.metadata.items():
622
+ if k not in dst_meta:
623
+ non_matching[k] = {"src": src_v, "dst": "KEY NOT FOUND"}
624
+ continue
625
+ dst_v = dst_meta[k]
626
+ if src_v != dst_v:
627
+ non_matching[k] = {"src": src_v, "dst": dst_v}
628
+
629
+ return non_matching
630
+
631
+ def _compare_run_summary(self, src_run: Run, dst_run: Run) -> dict:
632
+ non_matching = {}
633
+ for k, src_v in src_run.summary.items():
634
+ # These won't match between systems and that's ok
635
+ if isinstance(src_v, str) and src_v.startswith("wandb-client-artifact://"):
636
+ continue
637
+ if k in ("_wandb", "_runtime"):
638
+ continue
639
+
640
+ src_v = _recursive_cast_to_dict(src_v)
641
+
642
+ dst_v = dst_run.summary.get(k)
643
+ dst_v = _recursive_cast_to_dict(dst_v)
644
+
645
+ if isinstance(src_v, dict) and isinstance(dst_v, dict):
646
+ for kk, sv in src_v.items():
647
+ # These won't match between systems and that's ok
648
+ if isinstance(sv, str) and sv.startswith(
649
+ "wandb-client-artifact://"
650
+ ):
651
+ continue
652
+ dv = dst_v.get(kk)
653
+ if not _almost_equal(sv, dv):
654
+ non_matching[f"{k}-{kk}"] = {"src": sv, "dst": dv}
655
+ else:
656
+ if not _almost_equal(src_v, dst_v):
657
+ non_matching[k] = {"src": src_v, "dst": dst_v}
658
+
659
+ return non_matching
660
+
661
+ def _collect_failed_artifact_sequences(self) -> Iterable[ArtifactSequence]:
662
+ if (df := _read_ndjson(ARTIFACT_ERRORS_FNAME)) is None:
663
+ logger.debug(f"{ARTIFACT_ERRORS_FNAME=} is empty, returning nothing")
664
+ return
665
+
666
+ unique_failed_sequences = df[
667
+ ["src_entity", "src_project", "name", "type"]
668
+ ].unique()
669
+
670
+ for row in unique_failed_sequences.iter_rows(named=True):
671
+ entity = row["src_entity"]
672
+ project = row["src_project"]
673
+ name = row["name"]
674
+ _type = row["type"]
675
+
676
+ art_name = f"{entity}/{project}/{name}"
677
+ arts = self.src_api.artifacts(_type, art_name)
678
+ arts = sorted(arts, key=lambda a: int(a.version.lstrip("v")))
679
+ arts = sorted(arts, key=lambda a: a.type)
680
+
681
+ yield ArtifactSequence(arts, entity, project, _type, name)
682
+
683
+ def _cleanup_dummy_runs(
684
+ self,
685
+ *,
686
+ namespaces: Optional[Iterable[Namespace]] = None,
687
+ api: Optional[Api] = None,
688
+ remapping: Optional[Dict[Namespace, Namespace]] = None,
689
+ ) -> None:
690
+ api = coalesce(api, self.dst_api)
691
+ namespaces = coalesce(namespaces, self._all_namespaces())
692
+
693
+ for ns in namespaces:
694
+ if remapping and ns in remapping:
695
+ ns = remapping[ns]
696
+
697
+ logger.debug(f"Cleaning up, {ns=}")
698
+ try:
699
+ runs = list(
700
+ api.runs(ns.path, filters={"displayName": RUN_DUMMY_PLACEHOLDER})
701
+ )
702
+ except ValueError as e:
703
+ if "Could not find project" in str(e):
704
+ logger.error("Could not find project, does it exist?")
705
+ continue
706
+
707
+ for run in runs:
708
+ logger.debug(f"Deleting dummy {run=}")
709
+ run.delete(delete_artifacts=False)
710
+
711
+ def _import_report(
712
+ self, report: Report, *, namespace: Optional[Namespace] = None
713
+ ) -> None:
714
+ """Import one wandb.Report.
715
+
716
+ Use `namespace` to specify alternate settings like where the report should be uploaded
717
+ """
718
+ if namespace is None:
719
+ namespace = Namespace(report.entity, report.project)
720
+
721
+ entity = coalesce(namespace.entity, report.entity)
722
+ project = coalesce(namespace.project, report.project)
723
+ name = report.name
724
+ title = report.title
725
+ description = report.description
726
+
727
+ api = self.dst_api
728
+
729
+ # We shouldn't need to upsert the project for every report
730
+ logger.debug(f"Upserting {entity=}/{project=}")
731
+ try:
732
+ api.create_project(project, entity)
733
+ except requests.exceptions.HTTPError as e:
734
+ if e.response.status_code != 409:
735
+ logger.warn(f"Issue upserting {entity=}/{project=}, {e=}")
736
+
737
+ logger.debug(f"Upserting report {entity=}, {project=}, {name=}, {title=}")
738
+ api.client.execute(
739
+ wr.report.UPSERT_VIEW,
740
+ variable_values={
741
+ "id": None, # Is there any benefit for this to be the same as default report?
742
+ "name": name,
743
+ "entityName": entity,
744
+ "projectName": project,
745
+ "description": description,
746
+ "displayName": title,
747
+ "type": "runs",
748
+ "spec": json.dumps(report.spec),
749
+ },
750
+ )
751
+
752
+ def _use_artifact_sequence(
753
+ self,
754
+ sequence: ArtifactSequence,
755
+ *,
756
+ namespace: Optional[Namespace] = None,
757
+ ):
758
+ if namespace is None:
759
+ namespace = Namespace(sequence.entity, sequence.project)
760
+
761
+ settings_override = {
762
+ "api_key": self.dst_api_key,
763
+ "base_url": self.dst_base_url,
764
+ "resume": "true",
765
+ "resumed": True,
766
+ }
767
+ logger.debug(f"Using artifact sequence with {settings_override=}, {namespace=}")
768
+
769
+ send_manager_config = internal.SendManagerConfig(use_artifacts=True)
770
+
771
+ for art in sequence:
772
+ if (used_by := art.used_by()) is None:
773
+ continue
774
+
775
+ for wandb_run in used_by:
776
+ run = WandbRun(wandb_run, **self.run_api_kwargs)
777
+
778
+ internal.send_run(
779
+ run,
780
+ overrides=namespace.send_manager_overrides,
781
+ settings_override=settings_override,
782
+ config=send_manager_config,
783
+ )
784
+
785
+ def import_runs(
786
+ self,
787
+ *,
788
+ namespaces: Optional[Iterable[Namespace]] = None,
789
+ remapping: Optional[Dict[Namespace, Namespace]] = None,
790
+ parallel: bool = True,
791
+ incremental: bool = True,
792
+ max_workers: Optional[int] = None,
793
+ limit: Optional[int] = None,
794
+ metadata: bool = True,
795
+ files: bool = True,
796
+ media: bool = True,
797
+ code: bool = True,
798
+ history: bool = True,
799
+ summary: bool = True,
800
+ terminal_output: bool = True,
801
+ ):
802
+ logger.info("START: Import runs")
803
+
804
+ logger.info("Setting up for import")
805
+ _create_files_if_not_exists()
806
+ _clear_fname(RUN_ERRORS_FNAME)
807
+
808
+ logger.info("Collecting runs")
809
+ runs = list(self._collect_runs(namespaces=namespaces, limit=limit))
810
+
811
+ logger.info(f"Validating runs, {len(runs)=}")
812
+ self._validate_runs(
813
+ runs,
814
+ skip_previously_validated=incremental,
815
+ remapping=remapping,
816
+ )
817
+
818
+ logger.info("Collecting failed runs")
819
+ runs = list(self._collect_failed_runs())
820
+
821
+ logger.info(f"Importing runs, {len(runs)=}")
822
+
823
+ def _import_run_wrapped(run):
824
+ namespace = Namespace(run.entity(), run.project())
825
+ if remapping is not None and namespace in remapping:
826
+ namespace = remapping[namespace]
827
+
828
+ config = internal.SendManagerConfig(
829
+ metadata=metadata,
830
+ files=files,
831
+ media=media,
832
+ code=code,
833
+ history=history,
834
+ summary=summary,
835
+ terminal_output=terminal_output,
836
+ )
837
+
838
+ logger.debug(f"Importing {run=}, {namespace=}, {config=}")
839
+ self._import_run(run, namespace=namespace, config=config)
840
+ logger.debug(f"Finished importing {run=}, {namespace=}, {config=}")
841
+
842
+ for_each(_import_run_wrapped, runs, max_workers=max_workers, parallel=parallel)
843
+ logger.info("END: Importing runs")
844
+
845
+ def import_reports(
846
+ self,
847
+ *,
848
+ namespaces: Optional[Iterable[Namespace]] = None,
849
+ limit: Optional[int] = None,
850
+ remapping: Optional[Dict[Namespace, Namespace]] = None,
851
+ ):
852
+ logger.info("START: Importing reports")
853
+
854
+ logger.info("Collecting reports")
855
+ reports = self._collect_reports(namespaces=namespaces, limit=limit)
856
+
857
+ logger.info("Importing reports")
858
+
859
+ def _import_report_wrapped(report):
860
+ namespace = Namespace(report.entity, report.project)
861
+ if remapping is not None and namespace in remapping:
862
+ namespace = remapping[namespace]
863
+
864
+ logger.debug(f"Importing {report=}, {namespace=}")
865
+ self._import_report(report, namespace=namespace)
866
+ logger.debug(f"Finished importing {report=}, {namespace=}")
867
+
868
+ for_each(_import_report_wrapped, reports)
869
+
870
+ logger.info("END: Importing reports")
871
+
872
+ def import_artifact_sequences(
873
+ self,
874
+ *,
875
+ namespaces: Optional[Iterable[Namespace]] = None,
876
+ incremental: bool = True,
877
+ max_workers: Optional[int] = None,
878
+ remapping: Optional[Dict[Namespace, Namespace]] = None,
879
+ ):
880
+ """Import all artifact sequences from `namespaces`.
881
+
882
+ Note: There is a known bug with the AWS backend where artifacts > 2048MB will fail to upload. This seems to be related to multipart uploads, but we don't have a fix yet.
883
+ """
884
+ logger.info("START: Importing artifact sequences")
885
+ _clear_fname(ARTIFACT_ERRORS_FNAME)
886
+
887
+ logger.info("Collecting artifact sequences")
888
+ seqs = list(self._collect_artifact_sequences(namespaces=namespaces))
889
+
890
+ logger.info("Validating artifact sequences")
891
+ self._validate_artifact_sequences(
892
+ seqs,
893
+ incremental=incremental,
894
+ remapping=remapping,
895
+ )
896
+
897
+ logger.info("Collecting failed artifact sequences")
898
+ seqs = list(self._collect_failed_artifact_sequences())
899
+
900
+ logger.info(f"Importing artifact sequences, {len(seqs)=}")
901
+
902
+ def _import_artifact_sequence_wrapped(seq):
903
+ namespace = Namespace(seq.entity, seq.project)
904
+ if remapping is not None and namespace in remapping:
905
+ namespace = remapping[namespace]
906
+
907
+ logger.debug(f"Importing artifact sequence {seq=}, {namespace=}")
908
+ self._import_artifact_sequence(seq, namespace=namespace)
909
+ logger.debug(f"Finished importing artifact sequence {seq=}, {namespace=}")
910
+
911
+ for_each(_import_artifact_sequence_wrapped, seqs, max_workers=max_workers)
912
+
913
+ # it's safer to just use artifact on all seqs to make sure we don't miss anything
914
+ # For seqs that have already been used, this is a no-op.
915
+ logger.debug(f"Using artifact sequences, {len(seqs)=}")
916
+
917
+ def _use_artifact_sequence_wrapped(seq):
918
+ namespace = Namespace(seq.entity, seq.project)
919
+ if remapping is not None and namespace in remapping:
920
+ namespace = remapping[namespace]
921
+
922
+ logger.debug(f"Using artifact sequence {seq=}, {namespace=}")
923
+ self._use_artifact_sequence(seq, namespace=namespace)
924
+ logger.debug(f"Finished using artifact sequence {seq=}, {namespace=}")
925
+
926
+ for_each(_use_artifact_sequence_wrapped, seqs, max_workers=max_workers)
927
+
928
+ # Artifacts whose parent runs have been deleted should have that run deleted in the
929
+ # destination as well
930
+
931
+ logger.info("Cleaning up dummy runs")
932
+ self._cleanup_dummy_runs(
933
+ namespaces=namespaces,
934
+ remapping=remapping,
935
+ )
936
+
937
+ logger.info("END: Importing artifact sequences")
938
+
939
+ def import_all(
940
+ self,
941
+ *,
942
+ runs: bool = True,
943
+ artifacts: bool = True,
944
+ reports: bool = True,
945
+ namespaces: Optional[Iterable[Namespace]] = None,
946
+ incremental: bool = True,
947
+ remapping: Optional[Dict[Namespace, Namespace]] = None,
948
+ ):
949
+ logger.info(f"START: Importing all, {runs=}, {artifacts=}, {reports=}")
950
+ if runs:
951
+ self.import_runs(
952
+ namespaces=namespaces,
953
+ incremental=incremental,
954
+ remapping=remapping,
955
+ )
956
+
957
+ if reports:
958
+ self.import_reports(
959
+ namespaces=namespaces,
960
+ remapping=remapping,
961
+ )
962
+
963
+ if artifacts:
964
+ self.import_artifact_sequences(
965
+ namespaces=namespaces,
966
+ incremental=incremental,
967
+ remapping=remapping,
968
+ )
969
+
970
+ logger.info("END: Importing all")
971
+
972
+ def _validate_run(
973
+ self,
974
+ src_run: Run,
975
+ *,
976
+ remapping: Optional[Dict[Namespace, Namespace]] = None,
977
+ ) -> None:
978
+ namespace = Namespace(src_run.entity, src_run.project)
979
+ if remapping is not None and namespace in remapping:
980
+ namespace = remapping[namespace]
981
+
982
+ dst_entity = namespace.entity
983
+ dst_project = namespace.project
984
+ run_id = src_run.id
985
+
986
+ try:
987
+ dst_run = self.dst_api.run(f"{dst_entity}/{dst_project}/{run_id}")
988
+ except wandb.CommError:
989
+ problems = [f"run does not exist in dst at {dst_entity=}/{dst_project=}"]
990
+ else:
991
+ problems = self._get_run_problems(src_run, dst_run)
992
+
993
+ d = {
994
+ "src_entity": src_run.entity,
995
+ "src_project": src_run.project,
996
+ "dst_entity": dst_entity,
997
+ "dst_project": dst_project,
998
+ "run_id": run_id,
999
+ }
1000
+ if problems:
1001
+ d["problems"] = problems
1002
+ fname = RUN_ERRORS_FNAME
1003
+ else:
1004
+ fname = RUN_SUCCESSES_FNAME
1005
+
1006
+ with filelock.FileLock("runs.lock"):
1007
+ with open(fname, "a") as f:
1008
+ f.write(json.dumps(d) + "\n")
1009
+
1010
+ def _filter_previously_checked_runs(
1011
+ self,
1012
+ runs: Iterable[Run],
1013
+ *,
1014
+ remapping: Optional[Dict[Namespace, Namespace]] = None,
1015
+ ) -> Iterable[Run]:
1016
+ if (df := _read_ndjson(RUN_SUCCESSES_FNAME)) is None:
1017
+ logger.debug(f"{RUN_SUCCESSES_FNAME=} is empty, yielding all runs")
1018
+ yield from runs
1019
+ return
1020
+
1021
+ data = []
1022
+ for r in runs:
1023
+ namespace = Namespace(r.entity, r.project)
1024
+ if remapping is not None and namespace in remapping:
1025
+ namespace = remapping[namespace]
1026
+
1027
+ data.append(
1028
+ {
1029
+ "src_entity": r.entity,
1030
+ "src_project": r.project,
1031
+ "dst_entity": namespace.entity,
1032
+ "dst_project": namespace.project,
1033
+ "run_id": r.id,
1034
+ "data": r,
1035
+ }
1036
+ )
1037
+ df2 = pl.DataFrame(data)
1038
+ logger.debug(f"Starting with {len(runs)=} in namespaces")
1039
+
1040
+ results = df2.join(
1041
+ df,
1042
+ how="anti",
1043
+ on=["src_entity", "src_project", "dst_entity", "dst_project", "run_id"],
1044
+ )
1045
+ logger.debug(f"After filtering out already successful runs, {len(results)=}")
1046
+
1047
+ if not results.is_empty():
1048
+ results = results.filter(~results["run_id"].is_null())
1049
+ results = results.unique(
1050
+ ["src_entity", "src_project", "dst_entity", "dst_project", "run_id"]
1051
+ )
1052
+
1053
+ for r in results.iter_rows(named=True):
1054
+ yield r["data"]
1055
+
1056
+ def _validate_artifact(
1057
+ self,
1058
+ src_art: Artifact,
1059
+ dst_entity: str,
1060
+ dst_project: str,
1061
+ download_files_and_compare: bool = False,
1062
+ check_entries_are_downloadable: bool = True,
1063
+ ):
1064
+ problems = []
1065
+
1066
+ # These patterns of artifacts are special and should not be validated
1067
+ ignore_patterns = [
1068
+ r"^job-(.*?)\.py(:v\d+)?$",
1069
+ # r"^run-.*-history(?:\:v\d+)?$$",
1070
+ ]
1071
+ for pattern in ignore_patterns:
1072
+ if re.search(pattern, src_art.name):
1073
+ return (src_art, dst_entity, dst_project, problems)
1074
+
1075
+ try:
1076
+ dst_art = self._get_dst_art(src_art, dst_entity, dst_project)
1077
+ except Exception:
1078
+ problems.append("destination artifact not found")
1079
+ return (src_art, dst_entity, dst_project, problems)
1080
+
1081
+ try:
1082
+ logger.debug("Comparing artifact manifests")
1083
+ except Exception as e:
1084
+ problems.append(
1085
+ f"Problem getting problems! problem with {src_art.entity=}, {src_art.project=}, {src_art.name=} {e=}"
1086
+ )
1087
+ else:
1088
+ problems += validation._compare_artifact_manifests(src_art, dst_art)
1089
+
1090
+ if check_entries_are_downloadable:
1091
+ # validation._check_entries_are_downloadable(src_art)
1092
+ validation._check_entries_are_downloadable(dst_art)
1093
+
1094
+ if download_files_and_compare:
1095
+ logger.debug(f"Downloading {src_art=}")
1096
+ try:
1097
+ src_dir = _download_art(src_art, root=f"{SRC_ART_PATH}/{src_art.name}")
1098
+ except requests.HTTPError as e:
1099
+ problems.append(
1100
+ f"Invalid download link for src {src_art.entity=}, {src_art.project=}, {src_art.name=}, {e}"
1101
+ )
1102
+
1103
+ logger.debug(f"Downloading {dst_art=}")
1104
+ try:
1105
+ dst_dir = _download_art(dst_art, root=f"{DST_ART_PATH}/{dst_art.name}")
1106
+ except requests.HTTPError as e:
1107
+ problems.append(
1108
+ f"Invalid download link for dst {dst_art.entity=}, {dst_art.project=}, {dst_art.name=}, {e}"
1109
+ )
1110
+ else:
1111
+ logger.debug(f"Comparing artifact dirs {src_dir=}, {dst_dir=}")
1112
+ if problem := validation._compare_artifact_dirs(src_dir, dst_dir):
1113
+ problems.append(problem)
1114
+
1115
+ return (src_art, dst_entity, dst_project, problems)
1116
+
1117
+ def _validate_runs(
1118
+ self,
1119
+ runs: Iterable[WandbRun],
1120
+ *,
1121
+ skip_previously_validated: bool = True,
1122
+ remapping: Optional[Dict[Namespace, Namespace]] = None,
1123
+ ):
1124
+ base_runs = [r.run for r in runs]
1125
+ if skip_previously_validated:
1126
+ base_runs = list(
1127
+ self._filter_previously_checked_runs(
1128
+ base_runs,
1129
+ remapping=remapping,
1130
+ )
1131
+ )
1132
+
1133
+ def _validate_run(run):
1134
+ logger.debug(f"Validating {run=}")
1135
+ self._validate_run(run, remapping=remapping)
1136
+ logger.debug(f"Finished validating {run=}")
1137
+
1138
+ for_each(_validate_run, base_runs)
1139
+
1140
+ def _collect_failed_runs(self):
1141
+ if (df := _read_ndjson(RUN_ERRORS_FNAME)) is None:
1142
+ logger.debug(f"{RUN_ERRORS_FNAME=} is empty, returning nothing")
1143
+ return
1144
+
1145
+ unique_failed_runs = df[
1146
+ ["src_entity", "src_project", "dst_entity", "dst_project", "run_id"]
1147
+ ].unique()
1148
+
1149
+ for row in unique_failed_runs.iter_rows(named=True):
1150
+ src_entity = row["src_entity"]
1151
+ src_project = row["src_project"]
1152
+ # dst_entity = row["dst_entity"]
1153
+ # dst_project = row["dst_project"]
1154
+ run_id = row["run_id"]
1155
+
1156
+ run = self.src_api.run(f"{src_entity}/{src_project}/{run_id}")
1157
+ yield WandbRun(run, **self.run_api_kwargs)
1158
+
1159
+ def _filter_previously_checked_artifacts(self, seqs: Iterable[ArtifactSequence]):
1160
+ if (df := _read_ndjson(ARTIFACT_SUCCESSES_FNAME)) is None:
1161
+ logger.info(
1162
+ f"{ARTIFACT_SUCCESSES_FNAME=} is empty, yielding all artifact sequences"
1163
+ )
1164
+ for seq in seqs:
1165
+ yield from seq.artifacts
1166
+ return
1167
+
1168
+ for seq in seqs:
1169
+ for art in seq:
1170
+ try:
1171
+ logged_by = _get_run_or_dummy_from_art(art, self.src_api)
1172
+ except requests.HTTPError as e:
1173
+ logger.error(f"Failed to get run, skipping: {art=}, {e=}")
1174
+ continue
1175
+
1176
+ if art.type == "wandb-history" and isinstance(logged_by, _DummyRun):
1177
+ logger.debug(f"Skipping history artifact {art=}")
1178
+ # We can never upload valid history for a deleted run, so skip it
1179
+ continue
1180
+
1181
+ entity = art.entity
1182
+ project = art.project
1183
+ _type = art.type
1184
+ name, ver = _get_art_name_ver(art)
1185
+
1186
+ filtered_df = df.filter(
1187
+ (df["src_entity"] == entity)
1188
+ & (df["src_project"] == project)
1189
+ & (df["name"] == name)
1190
+ & (df["version"] == ver)
1191
+ & (df["type"] == _type)
1192
+ )
1193
+
1194
+ # not in file, so not verified yet, don't filter out
1195
+ if len(filtered_df) == 0:
1196
+ yield art
1197
+
1198
+ def _validate_artifact_sequences(
1199
+ self,
1200
+ seqs: Iterable[ArtifactSequence],
1201
+ *,
1202
+ incremental: bool = True,
1203
+ download_files_and_compare: bool = False,
1204
+ check_entries_are_downloadable: bool = True,
1205
+ remapping: Optional[Dict[Namespace, Namespace]] = None,
1206
+ ):
1207
+ if incremental:
1208
+ logger.info("Validating in incremental mode")
1209
+
1210
+ def filtered_sequences():
1211
+ for seq in seqs:
1212
+ if not seq.artifacts:
1213
+ continue
1214
+
1215
+ art = seq.artifacts[0]
1216
+ try:
1217
+ logged_by = _get_run_or_dummy_from_art(art, self.src_api)
1218
+ except requests.HTTPError as e:
1219
+ logger.error(
1220
+ f"Validate Artifact http error: {art.entity=}, {art.project=}, {art.name=}, {e=}"
1221
+ )
1222
+ continue
1223
+
1224
+ if art.type == "wandb-history" and isinstance(logged_by, _DummyRun):
1225
+ # We can never upload valid history for a deleted run, so skip it
1226
+ continue
1227
+
1228
+ yield seq
1229
+
1230
+ artifacts = self._filter_previously_checked_artifacts(filtered_sequences())
1231
+ else:
1232
+ logger.info("Validating in non-incremental mode")
1233
+ artifacts = [art for seq in seqs for art in seq.artifacts]
1234
+
1235
+ def _validate_artifact_wrapped(args):
1236
+ art, entity, project = args
1237
+ if (
1238
+ remapping is not None
1239
+ and (namespace := Namespace(entity, project)) in remapping
1240
+ ):
1241
+ remapped_ns = remapping[namespace]
1242
+ entity = remapped_ns.entity
1243
+ project = remapped_ns.project
1244
+
1245
+ logger.debug(f"Validating {art=}, {entity=}, {project=}")
1246
+ result = self._validate_artifact(
1247
+ art,
1248
+ entity,
1249
+ project,
1250
+ download_files_and_compare=download_files_and_compare,
1251
+ check_entries_are_downloadable=check_entries_are_downloadable,
1252
+ )
1253
+ logger.debug(f"Finished validating {art=}, {entity=}, {project=}")
1254
+ return result
1255
+
1256
+ args = ((art, art.entity, art.project) for art in artifacts)
1257
+ art_problems = for_each(_validate_artifact_wrapped, args)
1258
+ for art, dst_entity, dst_project, problems in art_problems:
1259
+ name, ver = _get_art_name_ver(art)
1260
+ d = {
1261
+ "src_entity": art.entity,
1262
+ "src_project": art.project,
1263
+ "dst_entity": dst_entity,
1264
+ "dst_project": dst_project,
1265
+ "name": name,
1266
+ "version": ver,
1267
+ "type": art.type,
1268
+ }
1269
+
1270
+ if problems:
1271
+ d["problems"] = problems
1272
+ fname = ARTIFACT_ERRORS_FNAME
1273
+ else:
1274
+ fname = ARTIFACT_SUCCESSES_FNAME
1275
+
1276
+ with open(fname, "a") as f:
1277
+ f.write(json.dumps(d) + "\n")
1278
+
1279
+ def _collect_runs(
1280
+ self,
1281
+ *,
1282
+ namespaces: Optional[Iterable[Namespace]] = None,
1283
+ limit: Optional[int] = None,
1284
+ skip_ids: Optional[List[str]] = None,
1285
+ start_date: Optional[str] = None,
1286
+ api: Optional[Api] = None,
1287
+ ) -> Iterable[WandbRun]:
1288
+ api = coalesce(api, self.src_api)
1289
+ namespaces = coalesce(namespaces, self._all_namespaces())
1290
+
1291
+ filters: Dict[str, Any] = {}
1292
+ if skip_ids is not None:
1293
+ filters["name"] = {"$nin": skip_ids}
1294
+ if start_date is not None:
1295
+ filters["createdAt"] = {"$gte": start_date}
1296
+
1297
+ def _runs():
1298
+ for ns in namespaces:
1299
+ logger.debug(f"Collecting runs from {ns=}")
1300
+ for run in api.runs(ns.path, filters=filters):
1301
+ yield WandbRun(run, **self.run_api_kwargs)
1302
+
1303
+ runs = itertools.islice(_runs(), limit)
1304
+ yield from runs
1305
+
1306
+ def _all_namespaces(
1307
+ self, *, entity: Optional[str] = None, api: Optional[Api] = None
1308
+ ):
1309
+ api = coalesce(api, self.src_api)
1310
+ entity = coalesce(entity, api.default_entity)
1311
+ projects = api.projects(entity)
1312
+ for p in projects:
1313
+ yield Namespace(p.entity, p.name)
1314
+
1315
+ def _collect_reports(
1316
+ self,
1317
+ *,
1318
+ namespaces: Optional[Iterable[Namespace]] = None,
1319
+ limit: Optional[int] = None,
1320
+ api: Optional[Api] = None,
1321
+ ):
1322
+ api = coalesce(api, self.src_api)
1323
+ namespaces = coalesce(namespaces, self._all_namespaces())
1324
+
1325
+ wandb.login(key=self.src_api_key, host=self.src_base_url)
1326
+
1327
+ def reports():
1328
+ for ns in namespaces:
1329
+ for r in api.reports(ns.path):
1330
+ yield wr.Report.from_url(r.url, api=api)
1331
+
1332
+ yield from itertools.islice(reports(), limit)
1333
+
1334
+ def _collect_artifact_sequences(
1335
+ self,
1336
+ *,
1337
+ namespaces: Optional[Iterable[Namespace]] = None,
1338
+ limit: Optional[int] = None,
1339
+ api: Optional[Api] = None,
1340
+ ):
1341
+ api = coalesce(api, self.src_api)
1342
+ namespaces = coalesce(namespaces, self._all_namespaces())
1343
+
1344
+ def artifact_sequences():
1345
+ for ns in namespaces:
1346
+ logger.debug(f"Collecting artifact sequences from {ns=}")
1347
+ types = []
1348
+ try:
1349
+ types = [t for t in api.artifact_types(ns.path)]
1350
+ except Exception as e:
1351
+ logger.error(f"Failed to get artifact types {e=}")
1352
+
1353
+ for t in types:
1354
+ collections = []
1355
+
1356
+ # Skip history because it's really for run history
1357
+ if t.name == "wandb-history":
1358
+ continue
1359
+
1360
+ try:
1361
+ collections = t.collections()
1362
+ except Exception as e:
1363
+ logger.error(f"Failed to get artifact collections {e=}")
1364
+
1365
+ for c in collections:
1366
+ if c.is_sequence():
1367
+ yield ArtifactSequence.from_collection(c)
1368
+
1369
+ seqs = itertools.islice(artifact_sequences(), limit)
1370
+ unique_sequences = {seq.identifier: seq for seq in seqs}
1371
+ yield from unique_sequences.values()
1372
+
1373
+
1374
+ def _get_art_name_ver(art: Artifact) -> Tuple[str, int]:
1375
+ name, ver = art.name.split(":v")
1376
+ return name, int(ver)
1377
+
1378
+
1379
+ def _make_dummy_art(name: str, _type: str, ver: int):
1380
+ art = Artifact(name, ART_DUMMY_PLACEHOLDER_TYPE)
1381
+ art._type = _type
1382
+ art._description = ART_SEQUENCE_DUMMY_PLACEHOLDER
1383
+
1384
+ p = Path(ART_DUMMY_PLACEHOLDER_PATH)
1385
+ p.mkdir(parents=True, exist_ok=True)
1386
+
1387
+ # dummy file with different name to prevent dedupe
1388
+ fname = p / str(ver)
1389
+ with open(fname, "w"):
1390
+ pass
1391
+ art.add_file(fname)
1392
+
1393
+ return art
1394
+
1395
+
1396
+ def _make_groups_of_artifacts(seq: ArtifactSequence, start: int = 0):
1397
+ prev_ver = start - 1
1398
+ for art in seq:
1399
+ name, ver = _get_art_name_ver(art)
1400
+
1401
+ # If there's a gap between versions, fill with dummy artifacts
1402
+ if ver - prev_ver > 1:
1403
+ yield [_make_dummy_art(name, art.type, v) for v in range(prev_ver + 1, ver)]
1404
+
1405
+ # Then yield the actual artifact
1406
+ # Must always be a list of one artifact to guarantee ordering
1407
+ yield [art]
1408
+ prev_ver = ver
1409
+
1410
+
1411
+ def _recursive_cast_to_dict(obj):
1412
+ if isinstance(obj, list):
1413
+ return [_recursive_cast_to_dict(item) for item in obj]
1414
+ elif isinstance(obj, dict) or hasattr(obj, "items"):
1415
+ new_dict = {}
1416
+ for key, value in obj.items():
1417
+ new_dict[key] = _recursive_cast_to_dict(value)
1418
+ return new_dict
1419
+ else:
1420
+ return obj
1421
+
1422
+
1423
+ def _almost_equal(x, y, eps=1e-6):
1424
+ if isinstance(x, dict) and isinstance(y, dict):
1425
+ if x.keys() != y.keys():
1426
+ return False
1427
+ return all(_almost_equal(x[k], y[k], eps) for k in x)
1428
+
1429
+ if isinstance(x, numbers.Number) and isinstance(y, numbers.Number):
1430
+ return abs(x - y) < eps
1431
+
1432
+ if type(x) is not type(y):
1433
+ return False
1434
+
1435
+ return x == y
1436
+
1437
+
1438
+ @dataclass
1439
+ class _DummyUser:
1440
+ username: str = ""
1441
+
1442
+
1443
+ @dataclass
1444
+ class _DummyRun:
1445
+ entity: str = ""
1446
+ project: str = ""
1447
+ run_id: str = RUN_DUMMY_PLACEHOLDER
1448
+ id: str = RUN_DUMMY_PLACEHOLDER
1449
+ display_name: str = RUN_DUMMY_PLACEHOLDER
1450
+ notes: str = ""
1451
+ url: str = ""
1452
+ group: str = ""
1453
+ created_at: str = "2000-01-01"
1454
+ user: _DummyUser = field(default_factory=_DummyUser)
1455
+ tags: list = field(default_factory=list)
1456
+ summary: dict = field(default_factory=dict)
1457
+ config: dict = field(default_factory=dict)
1458
+
1459
+ def files(self):
1460
+ return []
1461
+
1462
+
1463
+ def _read_ndjson(fname: str) -> Optional[pl.DataFrame]:
1464
+ try:
1465
+ df = pl.read_ndjson(fname)
1466
+ except FileNotFoundError:
1467
+ return None
1468
+ except RuntimeError as e:
1469
+ # No runs previously checked
1470
+ if "empty string is not a valid JSON value" in str(e):
1471
+ return None
1472
+ if "error parsing ndjson" in str(e):
1473
+ return None
1474
+ raise e
1475
+
1476
+ return df
1477
+
1478
+
1479
+ def _get_run_or_dummy_from_art(art: Artifact, api=None):
1480
+ run = None
1481
+
1482
+ try:
1483
+ run = art.logged_by()
1484
+ except ValueError as e:
1485
+ logger.warn(
1486
+ f"Can't log artifact because run does't exist, {art=}, {run=}, {e=}"
1487
+ )
1488
+
1489
+ if run is not None:
1490
+ return run
1491
+
1492
+ query = gql(
1493
+ """
1494
+ query ArtifactCreatedBy(
1495
+ $id: ID!
1496
+ ) {
1497
+ artifact(id: $id) {
1498
+ createdBy {
1499
+ ... on Run {
1500
+ name
1501
+ project {
1502
+ name
1503
+ entityName
1504
+ }
1505
+ }
1506
+ }
1507
+ }
1508
+ }
1509
+ """
1510
+ )
1511
+ response = api.client.execute(query, variable_values={"id": art.id})
1512
+ creator = response.get("artifact", {}).get("createdBy", {})
1513
+ run = _DummyRun(
1514
+ entity=art.entity,
1515
+ project=art.project,
1516
+ run_id=creator.get("name", RUN_DUMMY_PLACEHOLDER),
1517
+ id=creator.get("name", RUN_DUMMY_PLACEHOLDER),
1518
+ )
1519
+ return run
1520
+
1521
+
1522
+ def _clear_fname(fname: str) -> None:
1523
+ old_fname = f"{internal.ROOT_DIR}/{fname}"
1524
+ new_fname = f"{internal.ROOT_DIR}/prev_{fname}"
1525
+
1526
+ logger.debug(f"Moving {old_fname=} to {new_fname=}")
1527
+ try:
1528
+ shutil.copy2(old_fname, new_fname)
1529
+ except FileNotFoundError:
1530
+ # this is just to make a copy of the last iteration, so its ok if the src doesn't exist
1531
+ pass
1532
+
1533
+ with open(fname, "w"):
1534
+ pass
1535
+
1536
+
1537
+ def _download_art(art: Artifact, root: str) -> Optional[str]:
1538
+ try:
1539
+ with patch("click.echo"):
1540
+ return art.download(root=root, skip_cache=True)
1541
+ except Exception as e:
1542
+ logger.error(f"Error downloading artifact {art=}, {e=}")
1543
+
1544
+
1545
+ def _clone_art(art: Artifact, root: Optional[str] = None):
1546
+ if root is None:
1547
+ # Currently, we would only ever clone a src artifact to move it to dst.
1548
+ root = f"{SRC_ART_PATH}/{art.name}"
1549
+
1550
+ if (path := _download_art(art, root=root)) is None:
1551
+ raise ValueError(f"Problem downloading {art=}")
1552
+
1553
+ name, _ = art.name.split(":v")
1554
+
1555
+ # Hack: skip naming validation check for wandb-* types
1556
+ new_art = Artifact(name, ART_DUMMY_PLACEHOLDER_TYPE)
1557
+ new_art._type = art.type
1558
+ new_art._created_at = art.created_at
1559
+
1560
+ new_art._aliases = art.aliases
1561
+ new_art._description = art.description
1562
+
1563
+ with patch("click.echo"):
1564
+ new_art.add_dir(path)
1565
+
1566
+ return new_art
1567
+
1568
+
1569
+ def _create_files_if_not_exists() -> None:
1570
+ fnames = [
1571
+ ARTIFACT_ERRORS_FNAME,
1572
+ ARTIFACT_SUCCESSES_FNAME,
1573
+ RUN_ERRORS_FNAME,
1574
+ RUN_SUCCESSES_FNAME,
1575
+ ]
1576
+
1577
+ for fname in fnames:
1578
+ logger.debug(f"Creating {fname=} if not exists")
1579
+ with open(fname, "a"):
1580
+ pass
1581
+
1582
+
1583
+ def _merge_dfs(dfs: List[pl.DataFrame]) -> pl.DataFrame:
1584
+ # Ensure there are DataFrames in the list
1585
+ if len(dfs) == 0:
1586
+ return pl.DataFrame()
1587
+
1588
+ if len(dfs) == 1:
1589
+ return dfs[0]
1590
+
1591
+ merged_df = dfs[0]
1592
+ for df in dfs[1:]:
1593
+ merged_df = merged_df.join(df, how="outer", on=["_step"])
1594
+ col_pairs = [
1595
+ (c, f"{c}_right")
1596
+ for c in merged_df.columns
1597
+ if f"{c}_right" in merged_df.columns
1598
+ ]
1599
+ for col, right in col_pairs:
1600
+ new_col = merged_df[col].fill_null(merged_df[right])
1601
+ merged_df = merged_df.with_columns(new_col).drop(right)
1602
+
1603
+ return merged_df