diff --git a/lib/python3.12/site-packages/ray/_private/__init__.py b/lib/python3.12/site-packages/ray/_private/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/lib/python3.12/site-packages/ray/_private/__pycache__/__init__.cpython-312.pyc b/lib/python3.12/site-packages/ray/_private/__pycache__/__init__.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..5f24a529cf7a70acbb2114f3575031dc13e74ec3 Binary files /dev/null and b/lib/python3.12/site-packages/ray/_private/__pycache__/__init__.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/_private/__pycache__/arrow_serialization.cpython-312.pyc b/lib/python3.12/site-packages/ray/_private/__pycache__/arrow_serialization.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..d0b56f8119443adfa6aeed4f8dda4a67158be779 Binary files /dev/null and b/lib/python3.12/site-packages/ray/_private/__pycache__/arrow_serialization.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/_private/__pycache__/arrow_utils.cpython-312.pyc b/lib/python3.12/site-packages/ray/_private/__pycache__/arrow_utils.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..866d04a755d98351fb98ee27db5ba029cac32b96 Binary files /dev/null and b/lib/python3.12/site-packages/ray/_private/__pycache__/arrow_utils.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/_private/__pycache__/async_compat.cpython-312.pyc b/lib/python3.12/site-packages/ray/_private/__pycache__/async_compat.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..251835b4f035f74ad51260bdc0627e11506cdd5d Binary files /dev/null and b/lib/python3.12/site-packages/ray/_private/__pycache__/async_compat.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/_private/__pycache__/async_utils.cpython-312.pyc b/lib/python3.12/site-packages/ray/_private/__pycache__/async_utils.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..ea4758696a32d64abbdc6a5c015bc15f7e675464 Binary files /dev/null and b/lib/python3.12/site-packages/ray/_private/__pycache__/async_utils.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/_private/__pycache__/authentication_test_utils.cpython-312.pyc b/lib/python3.12/site-packages/ray/_private/__pycache__/authentication_test_utils.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..e62e62d6c2bd8e58ebebf404f13fdf0ff29f15cb Binary files /dev/null and b/lib/python3.12/site-packages/ray/_private/__pycache__/authentication_test_utils.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/_private/__pycache__/auto_init_hook.cpython-312.pyc b/lib/python3.12/site-packages/ray/_private/__pycache__/auto_init_hook.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..7bd9e8447d13354c94fb76fbaf931734218e3cd4 Binary files /dev/null and b/lib/python3.12/site-packages/ray/_private/__pycache__/auto_init_hook.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/_private/__pycache__/client_mode_hook.cpython-312.pyc b/lib/python3.12/site-packages/ray/_private/__pycache__/client_mode_hook.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..7aff5eea8985ea2310531df9186981e176c04085 Binary files /dev/null and b/lib/python3.12/site-packages/ray/_private/__pycache__/client_mode_hook.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/_private/__pycache__/collections_utils.cpython-312.pyc b/lib/python3.12/site-packages/ray/_private/__pycache__/collections_utils.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..3d07c3cad728109d79b5b943353fa61deca00d91 Binary files /dev/null and b/lib/python3.12/site-packages/ray/_private/__pycache__/collections_utils.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/_private/__pycache__/compat.cpython-312.pyc b/lib/python3.12/site-packages/ray/_private/__pycache__/compat.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..96d506211803c9d293c45b819f4f9c33b70ae039 Binary files /dev/null and b/lib/python3.12/site-packages/ray/_private/__pycache__/compat.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/_private/__pycache__/conftest_utils.cpython-312.pyc b/lib/python3.12/site-packages/ray/_private/__pycache__/conftest_utils.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..c8d8254bd7837994bf6a42594c75b330c124d3fc Binary files /dev/null and b/lib/python3.12/site-packages/ray/_private/__pycache__/conftest_utils.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/_private/__pycache__/custom_types.cpython-312.pyc b/lib/python3.12/site-packages/ray/_private/__pycache__/custom_types.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..5418f0855f9526fbc6cffb930ae021d0f024013e Binary files /dev/null and b/lib/python3.12/site-packages/ray/_private/__pycache__/custom_types.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/_private/__pycache__/dict.cpython-312.pyc b/lib/python3.12/site-packages/ray/_private/__pycache__/dict.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..09f01d305f5cd2722f5aeb07e35077e77df10309 Binary files /dev/null and b/lib/python3.12/site-packages/ray/_private/__pycache__/dict.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/_private/__pycache__/external_storage.cpython-312.pyc b/lib/python3.12/site-packages/ray/_private/__pycache__/external_storage.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..5c42892f41d93bb29ec3b5f52b6c12c5fcd49c36 Binary files /dev/null and b/lib/python3.12/site-packages/ray/_private/__pycache__/external_storage.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/_private/__pycache__/function_manager.cpython-312.pyc b/lib/python3.12/site-packages/ray/_private/__pycache__/function_manager.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..25f24a6ce7324e56b0231cf5ee79b6cb3a48fe08 Binary files /dev/null and b/lib/python3.12/site-packages/ray/_private/__pycache__/function_manager.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/_private/__pycache__/gc_collect_manager.cpython-312.pyc b/lib/python3.12/site-packages/ray/_private/__pycache__/gc_collect_manager.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..d03a6f3e25dd9065b4ab6483e0288553a4b2ad96 Binary files /dev/null and b/lib/python3.12/site-packages/ray/_private/__pycache__/gc_collect_manager.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/_private/__pycache__/gcs_pubsub.cpython-312.pyc b/lib/python3.12/site-packages/ray/_private/__pycache__/gcs_pubsub.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..4aaad60299ea587cc7b7d63bda4669b8216cf80e Binary files /dev/null and b/lib/python3.12/site-packages/ray/_private/__pycache__/gcs_pubsub.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/_private/__pycache__/gcs_utils.cpython-312.pyc b/lib/python3.12/site-packages/ray/_private/__pycache__/gcs_utils.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..ab2f432eb0607d315b7ae366b1545fc9d67ff20b Binary files /dev/null and b/lib/python3.12/site-packages/ray/_private/__pycache__/gcs_utils.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/_private/__pycache__/grpc_utils.cpython-312.pyc b/lib/python3.12/site-packages/ray/_private/__pycache__/grpc_utils.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..8bd5e0bc77a14094cf8d57fc2482f4e68adf1e02 Binary files /dev/null and b/lib/python3.12/site-packages/ray/_private/__pycache__/grpc_utils.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/_private/__pycache__/inspect_util.cpython-312.pyc b/lib/python3.12/site-packages/ray/_private/__pycache__/inspect_util.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..bfbdf998139b3a4b30e2b05dd59221241603ee75 Binary files /dev/null and b/lib/python3.12/site-packages/ray/_private/__pycache__/inspect_util.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/_private/__pycache__/internal_api.cpython-312.pyc b/lib/python3.12/site-packages/ray/_private/__pycache__/internal_api.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..2bf57e135d6841b41c098ff2a1a9eeca49be7fd0 Binary files /dev/null and b/lib/python3.12/site-packages/ray/_private/__pycache__/internal_api.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/_private/__pycache__/label_utils.cpython-312.pyc b/lib/python3.12/site-packages/ray/_private/__pycache__/label_utils.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..550937208dbc87bc6f4db120840293821651288e Binary files /dev/null and b/lib/python3.12/site-packages/ray/_private/__pycache__/label_utils.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/_private/__pycache__/log.cpython-312.pyc b/lib/python3.12/site-packages/ray/_private/__pycache__/log.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..585bd7b4449bd66dfa095e6aa6746e1e20b2c97d Binary files /dev/null and b/lib/python3.12/site-packages/ray/_private/__pycache__/log.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/_private/__pycache__/log_monitor.cpython-312.pyc b/lib/python3.12/site-packages/ray/_private/__pycache__/log_monitor.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..361b165db87e53a18afafbf4d9fb5d5507ad2531 Binary files /dev/null and b/lib/python3.12/site-packages/ray/_private/__pycache__/log_monitor.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/_private/__pycache__/logging_utils.cpython-312.pyc b/lib/python3.12/site-packages/ray/_private/__pycache__/logging_utils.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..5ba680ab20c2407211772d063bf779765d7f13e6 Binary files /dev/null and b/lib/python3.12/site-packages/ray/_private/__pycache__/logging_utils.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/_private/__pycache__/memory_monitor.cpython-312.pyc b/lib/python3.12/site-packages/ray/_private/__pycache__/memory_monitor.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..962a3f2a169e7143ab75ba845fd1f4006a712c2b Binary files /dev/null and b/lib/python3.12/site-packages/ray/_private/__pycache__/memory_monitor.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/_private/__pycache__/metrics_agent.cpython-312.pyc b/lib/python3.12/site-packages/ray/_private/__pycache__/metrics_agent.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..436fbe4ccf89466301417eda52bec875b45fdc4a Binary files /dev/null and b/lib/python3.12/site-packages/ray/_private/__pycache__/metrics_agent.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/_private/__pycache__/node.cpython-312.pyc b/lib/python3.12/site-packages/ray/_private/__pycache__/node.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..3ffd48905e5e14ae2726c0c3072d2377bc3a616f Binary files /dev/null and b/lib/python3.12/site-packages/ray/_private/__pycache__/node.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/_private/__pycache__/object_ref_generator.cpython-312.pyc b/lib/python3.12/site-packages/ray/_private/__pycache__/object_ref_generator.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..9728a6b00de8da57fe8c5492007896647893de8b Binary files /dev/null and b/lib/python3.12/site-packages/ray/_private/__pycache__/object_ref_generator.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/_private/__pycache__/parameter.cpython-312.pyc b/lib/python3.12/site-packages/ray/_private/__pycache__/parameter.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..1a2f0dc3fdb1fb66e92c3da688408c81f7ae11f7 Binary files /dev/null and b/lib/python3.12/site-packages/ray/_private/__pycache__/parameter.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/_private/__pycache__/path_utils.cpython-312.pyc b/lib/python3.12/site-packages/ray/_private/__pycache__/path_utils.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..b4f3c88552e68b9b3da82a3f37099c1889f25ce1 Binary files /dev/null and b/lib/python3.12/site-packages/ray/_private/__pycache__/path_utils.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/_private/__pycache__/process_watcher.cpython-312.pyc b/lib/python3.12/site-packages/ray/_private/__pycache__/process_watcher.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..3561636cb50cfff4dc10a488a7e0ea4ca8a2d6c7 Binary files /dev/null and b/lib/python3.12/site-packages/ray/_private/__pycache__/process_watcher.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/_private/__pycache__/profiling.cpython-312.pyc b/lib/python3.12/site-packages/ray/_private/__pycache__/profiling.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..a52937b12e00b5e8f51cb9593079adb8d9aff01c Binary files /dev/null and b/lib/python3.12/site-packages/ray/_private/__pycache__/profiling.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/_private/__pycache__/prometheus_exporter.cpython-312.pyc b/lib/python3.12/site-packages/ray/_private/__pycache__/prometheus_exporter.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..13e2eabbd4bf7cc9fa8f0f0cf876c08648b4c8fa Binary files /dev/null and b/lib/python3.12/site-packages/ray/_private/__pycache__/prometheus_exporter.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/_private/__pycache__/protobuf_compat.cpython-312.pyc b/lib/python3.12/site-packages/ray/_private/__pycache__/protobuf_compat.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..7e2be892ded8b129041e41f819f1c4d5a7af45c7 Binary files /dev/null and b/lib/python3.12/site-packages/ray/_private/__pycache__/protobuf_compat.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/_private/__pycache__/ray_client_microbenchmark.cpython-312.pyc b/lib/python3.12/site-packages/ray/_private/__pycache__/ray_client_microbenchmark.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..1c9aa940eb2de3610940fe261c37b874cc75e47c Binary files /dev/null and b/lib/python3.12/site-packages/ray/_private/__pycache__/ray_client_microbenchmark.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/_private/__pycache__/ray_cluster_perf.cpython-312.pyc b/lib/python3.12/site-packages/ray/_private/__pycache__/ray_cluster_perf.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..2c74009b0bdacfb60dfd9857fe51c88f05019a9b Binary files /dev/null and b/lib/python3.12/site-packages/ray/_private/__pycache__/ray_cluster_perf.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/_private/__pycache__/ray_constants.cpython-312.pyc b/lib/python3.12/site-packages/ray/_private/__pycache__/ray_constants.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..dcd88dcaa8d827ac39a90b87d781b92b5dc5e837 Binary files /dev/null and b/lib/python3.12/site-packages/ray/_private/__pycache__/ray_constants.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/_private/__pycache__/ray_experimental_perf.cpython-312.pyc b/lib/python3.12/site-packages/ray/_private/__pycache__/ray_experimental_perf.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..b7058802df4a74e4e5642270afac3cd53e98c93b Binary files /dev/null and b/lib/python3.12/site-packages/ray/_private/__pycache__/ray_experimental_perf.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/_private/__pycache__/ray_microbenchmark_helpers.cpython-312.pyc b/lib/python3.12/site-packages/ray/_private/__pycache__/ray_microbenchmark_helpers.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..0ed0ae33f64581b1b24e9bf29d7367a75c6924e3 Binary files /dev/null and b/lib/python3.12/site-packages/ray/_private/__pycache__/ray_microbenchmark_helpers.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/_private/__pycache__/ray_perf.cpython-312.pyc b/lib/python3.12/site-packages/ray/_private/__pycache__/ray_perf.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..40313b3d1e4cce02bfe0e4789eea73263b6fa041 Binary files /dev/null and b/lib/python3.12/site-packages/ray/_private/__pycache__/ray_perf.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/_private/__pycache__/ray_process_reaper.cpython-312.pyc b/lib/python3.12/site-packages/ray/_private/__pycache__/ray_process_reaper.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..b22ad35d9a2e5f1105c9a546ccce71129e20ff1f Binary files /dev/null and b/lib/python3.12/site-packages/ray/_private/__pycache__/ray_process_reaper.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/_private/__pycache__/resource_and_label_spec.cpython-312.pyc b/lib/python3.12/site-packages/ray/_private/__pycache__/resource_and_label_spec.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..34668f0b945baf673d9bfc5b9ed89f1c7a05d9b1 Binary files /dev/null and b/lib/python3.12/site-packages/ray/_private/__pycache__/resource_and_label_spec.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/_private/__pycache__/resource_isolation_config.cpython-312.pyc b/lib/python3.12/site-packages/ray/_private/__pycache__/resource_isolation_config.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..5eb926daa02abac1312ba3551cb5ff6404368e1c Binary files /dev/null and b/lib/python3.12/site-packages/ray/_private/__pycache__/resource_isolation_config.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/_private/__pycache__/serialization.cpython-312.pyc b/lib/python3.12/site-packages/ray/_private/__pycache__/serialization.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..0f777c1247db746972b18db5abfa2e06569bd6f4 Binary files /dev/null and b/lib/python3.12/site-packages/ray/_private/__pycache__/serialization.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/_private/__pycache__/services.cpython-312.pyc b/lib/python3.12/site-packages/ray/_private/__pycache__/services.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..270aa686577704dd81a69bf4be434ae87e7bbfdb Binary files /dev/null and b/lib/python3.12/site-packages/ray/_private/__pycache__/services.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/_private/__pycache__/state.cpython-312.pyc b/lib/python3.12/site-packages/ray/_private/__pycache__/state.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..be0bc5794b8680f8262b651b0b4623945df8b0ef Binary files /dev/null and b/lib/python3.12/site-packages/ray/_private/__pycache__/state.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/_private/__pycache__/state_api_test_utils.cpython-312.pyc b/lib/python3.12/site-packages/ray/_private/__pycache__/state_api_test_utils.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..9786dcf7f9a19f0fe70870ef73a9ab8033518c6f Binary files /dev/null and b/lib/python3.12/site-packages/ray/_private/__pycache__/state_api_test_utils.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/_private/__pycache__/test_utils.cpython-312.pyc b/lib/python3.12/site-packages/ray/_private/__pycache__/test_utils.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..871c201805513d65e2cb5308afeb273bf7bb1d00 Binary files /dev/null and b/lib/python3.12/site-packages/ray/_private/__pycache__/test_utils.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/_private/__pycache__/tls_utils.cpython-312.pyc b/lib/python3.12/site-packages/ray/_private/__pycache__/tls_utils.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..bb12f8d56f4ffca48b52b594969414e2d9fc77f1 Binary files /dev/null and b/lib/python3.12/site-packages/ray/_private/__pycache__/tls_utils.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/_private/__pycache__/utils.cpython-312.pyc b/lib/python3.12/site-packages/ray/_private/__pycache__/utils.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..0102fa92206e6f5803a7bbd3e9fe2fa5d52bc03c Binary files /dev/null and b/lib/python3.12/site-packages/ray/_private/__pycache__/utils.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/_private/accelerators/__init__.py b/lib/python3.12/site-packages/ray/_private/accelerators/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..4cb14fef7956a2145b56c0cc1cedfd856b7e6583 --- /dev/null +++ b/lib/python3.12/site-packages/ray/_private/accelerators/__init__.py @@ -0,0 +1,84 @@ +from typing import Optional, Set + +from ray._private.accelerators.accelerator import ( + RAY_ACCEL_ENV_VAR_OVERRIDE_ON_ZERO_ENV_VAR, + AcceleratorManager, +) +from ray._private.accelerators.amd_gpu import AMDGPUAcceleratorManager +from ray._private.accelerators.hpu import HPUAcceleratorManager +from ray._private.accelerators.intel_gpu import IntelGPUAcceleratorManager +from ray._private.accelerators.neuron import NeuronAcceleratorManager +from ray._private.accelerators.npu import NPUAcceleratorManager +from ray._private.accelerators.nvidia_gpu import NvidiaGPUAcceleratorManager +from ray._private.accelerators.rbln import RBLNAcceleratorManager +from ray._private.accelerators.tpu import TPUAcceleratorManager + + +def get_all_accelerator_managers() -> Set[AcceleratorManager]: + """Get all accelerator managers supported by Ray.""" + return { + NvidiaGPUAcceleratorManager, + IntelGPUAcceleratorManager, + AMDGPUAcceleratorManager, + TPUAcceleratorManager, + NeuronAcceleratorManager, + HPUAcceleratorManager, + NPUAcceleratorManager, + RBLNAcceleratorManager, + } + + +def get_all_accelerator_resource_names() -> Set[str]: + """Get all resource names for accelerators.""" + return { + accelerator_manager.get_resource_name() + for accelerator_manager in get_all_accelerator_managers() + } + + +def get_accelerator_manager_for_resource( + resource_name: str, +) -> Optional[AcceleratorManager]: + """Get the corresponding accelerator manager for the given + accelerator resource name + + E.g., TPUAcceleratorManager is returned if resource name is "TPU" + """ + try: + return get_accelerator_manager_for_resource._resource_name_to_accelerator_manager.get( # noqa: E501 + resource_name, None + ) + except AttributeError: + # Lazy initialization. + resource_name_to_accelerator_manager = { + accelerator_manager.get_resource_name(): accelerator_manager + for accelerator_manager in get_all_accelerator_managers() + } + # Special handling for GPU resource name since multiple accelerator managers + # have the same GPU resource name. + if AMDGPUAcceleratorManager.get_current_node_num_accelerators() > 0: + resource_name_to_accelerator_manager["GPU"] = AMDGPUAcceleratorManager + elif IntelGPUAcceleratorManager.get_current_node_num_accelerators() > 0: + resource_name_to_accelerator_manager["GPU"] = IntelGPUAcceleratorManager + else: + resource_name_to_accelerator_manager["GPU"] = NvidiaGPUAcceleratorManager + get_accelerator_manager_for_resource._resource_name_to_accelerator_manager = ( + resource_name_to_accelerator_manager + ) + return resource_name_to_accelerator_manager.get(resource_name, None) + + +__all__ = [ + "NvidiaGPUAcceleratorManager", + "IntelGPUAcceleratorManager", + "AMDGPUAcceleratorManager", + "TPUAcceleratorManager", + "NeuronAcceleratorManager", + "HPUAcceleratorManager", + "NPUAcceleratorManager", + "RBLNAcceleratorManager", + "get_all_accelerator_managers", + "get_all_accelerator_resource_names", + "get_accelerator_manager_for_resource", + "RAY_ACCEL_ENV_VAR_OVERRIDE_ON_ZERO_ENV_VAR", +] diff --git a/lib/python3.12/site-packages/ray/_private/accelerators/__pycache__/__init__.cpython-312.pyc b/lib/python3.12/site-packages/ray/_private/accelerators/__pycache__/__init__.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..152c11306f053a9b6f35e9afc1f8620c0369a8a5 Binary files /dev/null and b/lib/python3.12/site-packages/ray/_private/accelerators/__pycache__/__init__.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/_private/accelerators/__pycache__/accelerator.cpython-312.pyc b/lib/python3.12/site-packages/ray/_private/accelerators/__pycache__/accelerator.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..c694dcfe7c27cadb24c756c12585277c2c690f8c Binary files /dev/null and b/lib/python3.12/site-packages/ray/_private/accelerators/__pycache__/accelerator.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/_private/accelerators/__pycache__/amd_gpu.cpython-312.pyc b/lib/python3.12/site-packages/ray/_private/accelerators/__pycache__/amd_gpu.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..f8c441dbcdc5810954f371a9e1c9e27fffbb1444 Binary files /dev/null and b/lib/python3.12/site-packages/ray/_private/accelerators/__pycache__/amd_gpu.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/_private/accelerators/__pycache__/hpu.cpython-312.pyc b/lib/python3.12/site-packages/ray/_private/accelerators/__pycache__/hpu.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..34c995a04417d6b1251844f0db5f69f9580462ca Binary files /dev/null and b/lib/python3.12/site-packages/ray/_private/accelerators/__pycache__/hpu.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/_private/accelerators/__pycache__/intel_gpu.cpython-312.pyc b/lib/python3.12/site-packages/ray/_private/accelerators/__pycache__/intel_gpu.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..6d90cf5b84648ca089172003892165e36ac3c8d5 Binary files /dev/null and b/lib/python3.12/site-packages/ray/_private/accelerators/__pycache__/intel_gpu.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/_private/accelerators/__pycache__/neuron.cpython-312.pyc b/lib/python3.12/site-packages/ray/_private/accelerators/__pycache__/neuron.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..01e6f132ecfa95975fb5b968cb145d667ca2ebbe Binary files /dev/null and b/lib/python3.12/site-packages/ray/_private/accelerators/__pycache__/neuron.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/_private/accelerators/__pycache__/npu.cpython-312.pyc b/lib/python3.12/site-packages/ray/_private/accelerators/__pycache__/npu.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..bb4634f510c2d2854c80d12b91bb4f5a6e5ffae2 Binary files /dev/null and b/lib/python3.12/site-packages/ray/_private/accelerators/__pycache__/npu.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/_private/accelerators/__pycache__/nvidia_gpu.cpython-312.pyc b/lib/python3.12/site-packages/ray/_private/accelerators/__pycache__/nvidia_gpu.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..3776ff2cff7ba2894186474fd4aa71993a9a6e43 Binary files /dev/null and b/lib/python3.12/site-packages/ray/_private/accelerators/__pycache__/nvidia_gpu.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/_private/accelerators/__pycache__/rbln.cpython-312.pyc b/lib/python3.12/site-packages/ray/_private/accelerators/__pycache__/rbln.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..ba8834461f8cd3de894df2f6cd78818b39ddd66d Binary files /dev/null and b/lib/python3.12/site-packages/ray/_private/accelerators/__pycache__/rbln.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/_private/accelerators/__pycache__/tpu.cpython-312.pyc b/lib/python3.12/site-packages/ray/_private/accelerators/__pycache__/tpu.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..46c2cd14b190d38ae8de2b6928b97efbef1b3bf0 Binary files /dev/null and b/lib/python3.12/site-packages/ray/_private/accelerators/__pycache__/tpu.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/_private/accelerators/accelerator.py b/lib/python3.12/site-packages/ray/_private/accelerators/accelerator.py new file mode 100644 index 0000000000000000000000000000000000000000..4b5332cb8a07800e62563b40c325649fb7c8bc7e --- /dev/null +++ b/lib/python3.12/site-packages/ray/_private/accelerators/accelerator.py @@ -0,0 +1,160 @@ +from abc import ABC, abstractmethod +from typing import Dict, List, Optional, Tuple + +# https://github.com/ray-project/ray/issues/54868 +# In the future, ray will avoid overriding the accelerator ids environment variables +# when the number of accelerators is zero. +# For example, when this environment variable is set, if a user sets `num_gpus=0` +# in the `ray.init()` call, the environment variable `CUDA_VISIBLE_DEVICES` will +# not be set to an empty string. +# +# This environment variable is used to disable this behavior temporarily. +# And to avoid breaking changes, this environment variable is set to True by default +# to follow the previous behavior. +# +RAY_ACCEL_ENV_VAR_OVERRIDE_ON_ZERO_ENV_VAR = "RAY_ACCEL_ENV_VAR_OVERRIDE_ON_ZERO" + + +class AcceleratorManager(ABC): + """This class contains all the functions needed for supporting + an accelerator family in Ray.""" + + @staticmethod + @abstractmethod + def get_resource_name() -> str: + """Get the name of the resource representing this accelerator family. + + Returns: + The resource name: e.g., the resource name for NVIDIA GPUs is "GPU" + """ + + @staticmethod + @abstractmethod + def get_visible_accelerator_ids_env_var() -> str: + """Get the env var that sets the ids of visible accelerators of this family. + + Returns: + The env var for setting visible accelerator ids: e.g., + CUDA_VISIBLE_DEVICES for NVIDIA GPUs. + """ + + @staticmethod + @abstractmethod + def get_current_node_num_accelerators() -> int: + """Get the total number of accelerators of this family on the current node. + + Returns: + The detected total number of accelerators of this family. + Return 0 if the current node doesn't contain accelerators of this family. + """ + + @staticmethod + @abstractmethod + def get_current_node_accelerator_type() -> Optional[str]: + """Get the type of the accelerator of this family on the current node. + + Currently Ray only supports single accelerator type of + an accelerator family on each node. + + The result should only be used when get_current_node_num_accelerators() > 0. + + Returns: + The detected accelerator type of this family: e.g., H100 for NVIDIA GPU. + Return None if it's unknown or the node doesn't have + accelerators of this family. + """ + + @staticmethod + @abstractmethod + def get_current_node_additional_resources() -> Optional[Dict[str, float]]: + """Get any additional resources required for the current node. + + In case a particular accelerator type requires considerations for + additional resources (e.g. for TPUs, providing the TPU pod type and + TPU name), this function can be used to provide the + additional logical resources. + + Returns: + A dictionary representing additional resources that may be + necessary for a particular accelerator type. + """ + + @staticmethod + @abstractmethod + def validate_resource_request_quantity( + quantity: float, + ) -> Tuple[bool, Optional[str]]: + """Validate the resource request quantity of this accelerator resource. + + Args: + quantity: The resource request quantity to be validated. + + Returns: + (valid, error_message) tuple: the first element of the tuple + indicates whether the given quantity is valid or not, + the second element is the error message + if the given quantity is invalid. + """ + + @staticmethod + @abstractmethod + def get_current_process_visible_accelerator_ids() -> Optional[List[str]]: + """Get the ids of accelerators of this family that are visible to the current process. + + Returns: + The list of visiable accelerator ids. + Return None if all accelerators are visible. + """ + + @staticmethod + @abstractmethod + def set_current_process_visible_accelerator_ids(ids: List[str]) -> None: + """Set the ids of accelerators of this family that are visible to the current process. + + Args: + ids: The ids of visible accelerators of this family. + """ + + @staticmethod + def get_ec2_instance_num_accelerators( + instance_type: str, instances: dict + ) -> Optional[int]: + """Get the number of accelerators of this family on ec2 instance with given type. + + Args: + instance_type: The ec2 instance type. + instances: Map from ec2 instance type to instance metadata returned by + ec2 `describe-instance-types`. + + Returns: + The number of accelerators of this family on the ec2 instance + with given type. + Return None if it's unknown. + """ + return None + + @staticmethod + def get_ec2_instance_accelerator_type( + instance_type: str, instances: dict + ) -> Optional[str]: + """Get the accelerator type of this family on ec2 instance with given type. + + Args: + instance_type: The ec2 instance type. + instances: Map from ec2 instance type to instance metadata returned by + ec2 `describe-instance-types`. + + Returns: + The accelerator type of this family on the ec2 instance with given type. + Return None if it's unknown. + """ + return None + + @staticmethod + def get_current_node_accelerator_labels() -> Optional[Dict[str, str]]: + """Get accelerator related Ray node labels of the curent node. + + Returns: + A dictionary mapping accelerator related label keys to values. + """ + return None diff --git a/lib/python3.12/site-packages/ray/_private/accelerators/amd_gpu.py b/lib/python3.12/site-packages/ray/_private/accelerators/amd_gpu.py new file mode 100644 index 0000000000000000000000000000000000000000..997d095f24041bc34fe0a502941647ec4afb96cf --- /dev/null +++ b/lib/python3.12/site-packages/ray/_private/accelerators/amd_gpu.py @@ -0,0 +1,162 @@ +import logging +import os +from typing import List, Optional, Tuple + +from ray._private.accelerators.accelerator import AcceleratorManager +from ray._private.accelerators.nvidia_gpu import CUDA_VISIBLE_DEVICES_ENV_VAR + +logger = logging.getLogger(__name__) + +HIP_VISIBLE_DEVICES_ENV_VAR = "HIP_VISIBLE_DEVICES" +NOSET_HIP_VISIBLE_DEVICES_ENV_VAR = "RAY_EXPERIMENTAL_NOSET_HIP_VISIBLE_DEVICES" + +amd_product_dict = { + "0x66a1": "AMD-Instinct-MI50", + "0x738c": "AMD-Instinct-MI100", + "0x7408": "AMD-Instinct-MI250X", + "0x740c": "AMD-Instinct-MI250X-MI250", + "0x740f": "AMD-Instinct-MI210", + "0x74a0": "AMD-Instinct-MI300A", + "0x74a1": "AMD-Instinct-MI300X-OAM", + "0x74a2": "AMD-Instinct-MI308X-OAM", + "0x74a9": "AMD-Instinct-MI300X-HF", + "0x74a5": "AMD-Instinct-MI325X-OAM", + "0x75a0": "AMD-Instinct-MI350X-OAM", + "0x75a3": "AMD-Instinct-MI355X-OAM", + "0x6798": "AMD-Radeon-R9-200-HD-7900", + "0x6799": "AMD-Radeon-HD-7900", + "0x679A": "AMD-Radeon-HD-7900", + "0x679B": "AMD-Radeon-HD-7900", +} + + +class AMDGPUAcceleratorManager(AcceleratorManager): + """AMD GPU accelerators.""" + + @staticmethod + def get_resource_name() -> str: + return "GPU" + + @staticmethod + def get_visible_accelerator_ids_env_var() -> str: + if ( + HIP_VISIBLE_DEVICES_ENV_VAR not in os.environ + and "ROCR_VISIBLE_DEVICES" in os.environ + ): + raise RuntimeError( + f"Please use {HIP_VISIBLE_DEVICES_ENV_VAR} instead of ROCR_VISIBLE_DEVICES" + ) + + env_var = HIP_VISIBLE_DEVICES_ENV_VAR + if (cuda_val := os.environ.get(CUDA_VISIBLE_DEVICES_ENV_VAR, None)) is not None: + if (hip_val := os.environ.get(HIP_VISIBLE_DEVICES_ENV_VAR, None)) is None: + env_var = CUDA_VISIBLE_DEVICES_ENV_VAR + elif hip_val != cuda_val: + raise ValueError( + f"Inconsistent values found. Please use either {HIP_VISIBLE_DEVICES_ENV_VAR} or {CUDA_VISIBLE_DEVICES_ENV_VAR}." + ) + + return env_var + + @staticmethod + def get_current_process_visible_accelerator_ids() -> Optional[List[str]]: + amd_visible_devices = os.environ.get( + AMDGPUAcceleratorManager.get_visible_accelerator_ids_env_var(), None + ) + + if amd_visible_devices is None: + return None + + if amd_visible_devices == "": + return [] + + if amd_visible_devices == "NoDevFiles": + return [] + + return list(amd_visible_devices.split(",")) + + @staticmethod + def get_current_node_num_accelerators() -> int: + import ray._private.thirdparty.pyamdsmi as pyamdsmi + + num_gpus = 0 + + try: + pyamdsmi.smi_initialize() + num_gpus = pyamdsmi.smi_get_device_count() + except Exception: + pass + finally: + try: + pyamdsmi.smi_shutdown() + except Exception: + pass + + return num_gpus + + @staticmethod + def get_current_node_accelerator_type() -> Optional[str]: + try: + device_ids = AMDGPUAcceleratorManager._get_amd_device_ids() + if device_ids is None: + return None + return AMDGPUAcceleratorManager._gpu_name_to_accelerator_type(device_ids[0]) + except Exception: + return None + + @staticmethod + def _gpu_name_to_accelerator_type(name): + if name is None: + return None + try: + match = amd_product_dict[name] + return match + except Exception: + return None + + @staticmethod + def validate_resource_request_quantity( + quantity: float, + ) -> Tuple[bool, Optional[str]]: + return (True, None) + + @staticmethod + def set_current_process_visible_accelerator_ids( + visible_amd_devices: List[str], + ) -> None: + if os.environ.get(NOSET_HIP_VISIBLE_DEVICES_ENV_VAR): + return + + os.environ[ + AMDGPUAcceleratorManager.get_visible_accelerator_ids_env_var() + ] = ",".join([str(i) for i in visible_amd_devices]) + + @staticmethod + def _get_amd_device_ids() -> List[str]: + """Get the list of GPUs IDs + Example: + On a node with 2x MI210 GPUs + pyamdsmi library python bindings + return: ['0x740f', '0x740f'] + Returns: + A list of strings containing GPU IDs + """ + import ray._private.thirdparty.pyamdsmi as pyamdsmi + + device_ids = [] + try: + pyamdsmi.smi_initialize() + num_devices = pyamdsmi.smi_get_device_count() + for i in range(num_devices): + did = pyamdsmi.smi_get_device_id(i) + if did >= 0: + device_ids.append(hex(did)) + except Exception: + return None + finally: + try: + pyamdsmi.pyamdsmi_shutdown() + except Exception: + pass + + return device_ids diff --git a/lib/python3.12/site-packages/ray/_private/accelerators/hpu.py b/lib/python3.12/site-packages/ray/_private/accelerators/hpu.py new file mode 100644 index 0000000000000000000000000000000000000000..0f5d16b2761e57f76d15db70ec4f8696310e28b2 --- /dev/null +++ b/lib/python3.12/site-packages/ray/_private/accelerators/hpu.py @@ -0,0 +1,121 @@ +import logging +import os +from functools import lru_cache +from importlib.util import find_spec +from typing import List, Optional, Tuple + +from ray._private.accelerators.accelerator import AcceleratorManager + +logger = logging.getLogger(__name__) + +HABANA_VISIBLE_DEVICES_ENV_VAR = "HABANA_VISIBLE_MODULES" +NOSET_HABANA_VISIBLE_MODULES_ENV_VAR = "RAY_EXPERIMENTAL_NOSET_HABANA_VISIBLE_MODULES" + + +@lru_cache() +def is_package_present(package_name: str) -> bool: + try: + return find_spec(package_name) is not None + except ModuleNotFoundError: + return False + + +HPU_PACKAGE_AVAILABLE = is_package_present("habana_frameworks") + + +class HPUAcceleratorManager(AcceleratorManager): + """Intel Habana(HPU) accelerators.""" + + @staticmethod + def get_resource_name() -> str: + return "HPU" + + @staticmethod + def get_visible_accelerator_ids_env_var() -> str: + return HABANA_VISIBLE_DEVICES_ENV_VAR + + @staticmethod + def get_current_process_visible_accelerator_ids() -> Optional[List[str]]: + hpu_visible_devices = os.environ.get( + HPUAcceleratorManager.get_visible_accelerator_ids_env_var(), None + ) + + if hpu_visible_devices is None: + return None + + if hpu_visible_devices == "": + return [] + + return list(hpu_visible_devices.split(",")) + + @staticmethod + def get_current_node_num_accelerators() -> int: + """Attempt to detect the number of HPUs on this machine. + Returns: + The number of HPUs if any were detected, otherwise 0. + """ + if HPU_PACKAGE_AVAILABLE: + import habana_frameworks.torch.hpu as torch_hpu + + if torch_hpu.is_available(): + return torch_hpu.device_count() + else: + logging.info("HPU devices not available") + return 0 + else: + return 0 + + @staticmethod + def is_initialized() -> bool: + """Attempt to check if HPU backend is initialized. + Returns: + True if backend initialized else False. + """ + if HPU_PACKAGE_AVAILABLE: + import habana_frameworks.torch.hpu as torch_hpu + + if torch_hpu.is_available() and torch_hpu.is_initialized(): + return True + else: + return False + else: + return False + + @staticmethod + def get_current_node_accelerator_type() -> Optional[str]: + """Attempt to detect the HPU family type. + Returns: + The device name (GAUDI, GAUDI2) if detected else None. + """ + if HPUAcceleratorManager.is_initialized(): + import habana_frameworks.torch.hpu as torch_hpu + + return f"Intel-{torch_hpu.get_device_name()}" + else: + logging.info("HPU type cannot be detected") + return None + + @staticmethod + def validate_resource_request_quantity( + quantity: float, + ) -> Tuple[bool, Optional[str]]: + if isinstance(quantity, float) and not quantity.is_integer(): + return ( + False, + f"{HPUAcceleratorManager.get_resource_name()} resource quantity" + " must be whole numbers. " + f"The specified quantity {quantity} is invalid.", + ) + else: + return (True, None) + + @staticmethod + def set_current_process_visible_accelerator_ids( + visible_hpu_devices: List[str], + ) -> None: + if os.environ.get(NOSET_HABANA_VISIBLE_MODULES_ENV_VAR): + return + + os.environ[ + HPUAcceleratorManager.get_visible_accelerator_ids_env_var() + ] = ",".join([str(i) for i in visible_hpu_devices]) diff --git a/lib/python3.12/site-packages/ray/_private/accelerators/intel_gpu.py b/lib/python3.12/site-packages/ray/_private/accelerators/intel_gpu.py new file mode 100644 index 0000000000000000000000000000000000000000..8bace1f203668f8d23329922673fc131592335c5 --- /dev/null +++ b/lib/python3.12/site-packages/ray/_private/accelerators/intel_gpu.py @@ -0,0 +1,103 @@ +import logging +import os +from typing import List, Optional, Tuple + +from ray._private.accelerators.accelerator import AcceleratorManager + +logger = logging.getLogger(__name__) + +ONEAPI_DEVICE_SELECTOR_ENV_VAR = "ONEAPI_DEVICE_SELECTOR" +NOSET_ONEAPI_DEVICE_SELECTOR_ENV_VAR = "RAY_EXPERIMENTAL_NOSET_ONEAPI_DEVICE_SELECTOR" +ONEAPI_DEVICE_BACKEND_TYPE = "level_zero" +ONEAPI_DEVICE_TYPE = "gpu" + + +class IntelGPUAcceleratorManager(AcceleratorManager): + """Intel GPU accelerators.""" + + @staticmethod + def get_resource_name() -> str: + return "GPU" + + @staticmethod + def get_visible_accelerator_ids_env_var() -> str: + return ONEAPI_DEVICE_SELECTOR_ENV_VAR + + @staticmethod + def get_current_process_visible_accelerator_ids() -> Optional[List[str]]: + oneapi_visible_devices = os.environ.get( + IntelGPUAcceleratorManager.get_visible_accelerator_ids_env_var(), None + ) + if oneapi_visible_devices is None: + return None + + if oneapi_visible_devices == "": + return [] + + if oneapi_visible_devices == "NoDevFiles": + return [] + + prefix = ONEAPI_DEVICE_BACKEND_TYPE + ":" + + return list(oneapi_visible_devices.split(prefix)[1].split(",")) + + @staticmethod + def get_current_node_num_accelerators() -> int: + try: + import dpctl + except ImportError: + dpctl = None + if dpctl is None: + return 0 + + num_gpus = 0 + try: + dev_info = ONEAPI_DEVICE_BACKEND_TYPE + ":" + ONEAPI_DEVICE_TYPE + context = dpctl.SyclContext(dev_info) + num_gpus = context.device_count + except Exception: + num_gpus = 0 + return num_gpus + + @staticmethod + def get_current_node_accelerator_type() -> Optional[str]: + """Get the name of first Intel GPU. (supposed only one GPU type on a node) + Example: + name: 'Intel(R) Data Center GPU Max 1550' + return name: 'Intel-GPU-Max-1550' + Returns: + A string representing the name of Intel GPU type. + """ + try: + import dpctl + except ImportError: + dpctl = None + if dpctl is None: + return None + + accelerator_type = None + try: + dev_info = ONEAPI_DEVICE_BACKEND_TYPE + ":" + ONEAPI_DEVICE_TYPE + ":0" + dev = dpctl.SyclDevice(dev_info) + accelerator_type = "Intel-GPU-" + "-".join(dev.name.split(" ")[-2:]) + except Exception: + accelerator_type = None + return accelerator_type + + @staticmethod + def validate_resource_request_quantity( + quantity: float, + ) -> Tuple[bool, Optional[str]]: + return (True, None) + + @staticmethod + def set_current_process_visible_accelerator_ids( + visible_xpu_devices: List[str], + ) -> None: + if os.environ.get(NOSET_ONEAPI_DEVICE_SELECTOR_ENV_VAR): + return + + prefix = ONEAPI_DEVICE_BACKEND_TYPE + ":" + os.environ[ + IntelGPUAcceleratorManager.get_visible_accelerator_ids_env_var() + ] = prefix + ",".join([str(i) for i in visible_xpu_devices]) diff --git a/lib/python3.12/site-packages/ray/_private/accelerators/neuron.py b/lib/python3.12/site-packages/ray/_private/accelerators/neuron.py new file mode 100644 index 0000000000000000000000000000000000000000..a2799e3e8110b973cb5e78d086f5f4560588570f --- /dev/null +++ b/lib/python3.12/site-packages/ray/_private/accelerators/neuron.py @@ -0,0 +1,132 @@ +import json +import logging +import os +import subprocess +import sys +from typing import List, Optional, Tuple + +from ray._private.accelerators.accelerator import AcceleratorManager + +logger = logging.getLogger(__name__) + +NEURON_RT_VISIBLE_CORES_ENV_VAR = "NEURON_RT_VISIBLE_CORES" +NOSET_AWS_NEURON_RT_VISIBLE_CORES_ENV_VAR = ( + "RAY_EXPERIMENTAL_NOSET_NEURON_RT_VISIBLE_CORES" +) + +# https://awsdocs-neuron.readthedocs-hosted.com/en/latest/general/arch/neuron-hardware/inf2-arch.html#aws-inf2-arch +# https://awsdocs-neuron.readthedocs-hosted.com/en/latest/general/arch/neuron-hardware/trn1-arch.html#aws-trn1-arch +# Subject to removal after the information is available via public API +AWS_NEURON_INSTANCE_MAP = { + "trn1.2xlarge": 2, + "trn1.32xlarge": 32, + "trn1n.32xlarge": 32, + "inf2.xlarge": 2, + "inf2.8xlarge": 2, + "inf2.24xlarge": 12, + "inf2.48xlarge": 24, +} + + +class NeuronAcceleratorManager(AcceleratorManager): + """AWS Inferentia and Trainium accelerators.""" + + @staticmethod + def get_resource_name() -> str: + return "neuron_cores" + + @staticmethod + def get_visible_accelerator_ids_env_var() -> str: + return NEURON_RT_VISIBLE_CORES_ENV_VAR + + @staticmethod + def get_current_process_visible_accelerator_ids() -> Optional[List[str]]: + neuron_visible_cores = os.environ.get( + NeuronAcceleratorManager.get_visible_accelerator_ids_env_var(), None + ) + + if neuron_visible_cores is None: + return None + + if neuron_visible_cores == "": + return [] + + return list(neuron_visible_cores.split(",")) + + @staticmethod + def get_current_node_num_accelerators() -> int: + """ + Attempt to detect the number of Neuron cores on this machine. + + Returns: + The number of Neuron cores if any were detected, otherwise 0. + """ + nc_count: int = 0 + neuron_path = "/opt/aws/neuron/bin/" + if sys.platform.startswith("linux") and os.path.isdir(neuron_path): + result = subprocess.run( + [os.path.join(neuron_path, "neuron-ls"), "--json-output"], + stdout=subprocess.PIPE, + stderr=subprocess.PIPE, + ) + if result.returncode == 0 and result.stdout: + neuron_devices = json.loads(result.stdout) + for neuron_device in neuron_devices: + nc_count += neuron_device.get("nc_count", 0) + return nc_count + + @staticmethod + def get_current_node_accelerator_type() -> Optional[str]: + from ray.util.accelerators import AWS_NEURON_CORE + + return AWS_NEURON_CORE + + @staticmethod + def validate_resource_request_quantity( + quantity: float, + ) -> Tuple[bool, Optional[str]]: + if isinstance(quantity, float) and not quantity.is_integer(): + return ( + False, + f"{NeuronAcceleratorManager.get_resource_name()} resource quantity" + " must be whole numbers. " + f"The specified quantity {quantity} is invalid.", + ) + else: + return (True, None) + + @staticmethod + def set_current_process_visible_accelerator_ids( + visible_neuron_core_ids: List[str], + ) -> None: + """Set the NEURON_RT_VISIBLE_CORES environment variable based on + given visible_neuron_core_ids. + + Args: + visible_neuron_core_ids (List[str]): List of int representing core IDs. + """ + if os.environ.get(NOSET_AWS_NEURON_RT_VISIBLE_CORES_ENV_VAR): + return + + os.environ[ + NeuronAcceleratorManager.get_visible_accelerator_ids_env_var() + ] = ",".join([str(i) for i in visible_neuron_core_ids]) + + @staticmethod + def get_ec2_instance_num_accelerators( + instance_type: str, instances: dict + ) -> Optional[int]: + # TODO: AWS SDK (public API) doesn't yet expose the NeuronCore + # information. It will be available (work-in-progress) + # as xxAcceleratorInfo in InstanceTypeInfo. + # https://docs.aws.amazon.com/AWSEC2/latest/APIReference/API_InstanceTypeInfo.html + # See https://github.com/ray-project/ray/issues/38473 + return AWS_NEURON_INSTANCE_MAP.get(instance_type.lower(), None) + + @staticmethod + def get_ec2_instance_accelerator_type( + instance_type: str, instances: dict + ) -> Optional[str]: + from ray.util.accelerators import AWS_NEURON_CORE + + return AWS_NEURON_CORE diff --git a/lib/python3.12/site-packages/ray/_private/accelerators/npu.py b/lib/python3.12/site-packages/ray/_private/accelerators/npu.py new file mode 100644 index 0000000000000000000000000000000000000000..ab470619e0453c4810b20846f73630d7e62a14fa --- /dev/null +++ b/lib/python3.12/site-packages/ray/_private/accelerators/npu.py @@ -0,0 +1,99 @@ +import glob +import logging +import os +from typing import List, Optional, Tuple + +from ray._private.accelerators.accelerator import AcceleratorManager + +logger = logging.getLogger(__name__) + +ASCEND_RT_VISIBLE_DEVICES_ENV_VAR = "ASCEND_RT_VISIBLE_DEVICES" +NOSET_ASCEND_RT_VISIBLE_DEVICES_ENV_VAR = ( + "RAY_EXPERIMENTAL_NOSET_ASCEND_RT_VISIBLE_DEVICES" +) + + +class NPUAcceleratorManager(AcceleratorManager): + """Ascend NPU accelerators.""" + + @staticmethod + def get_resource_name() -> str: + return "NPU" + + @staticmethod + def get_visible_accelerator_ids_env_var() -> str: + return ASCEND_RT_VISIBLE_DEVICES_ENV_VAR + + @staticmethod + def get_current_process_visible_accelerator_ids() -> Optional[List[str]]: + ascend_visible_devices = os.environ.get( + NPUAcceleratorManager.get_visible_accelerator_ids_env_var(), None + ) + + if ascend_visible_devices is None: + return None + + if ascend_visible_devices == "": + return [] + + if ascend_visible_devices == "NoDevFiles": + return [] + + return list(ascend_visible_devices.split(",")) + + @staticmethod + def get_current_node_num_accelerators() -> int: + """Attempt to detect the number of NPUs on this machine. + + NPU chips are represented as devices within `/dev/`, either as `/dev/davinci?`. + + Returns: + The number of NPUs if any were detected, otherwise 0. + """ + try: + import acl + + device_count, ret = acl.rt.get_device_count() + if ret == 0: + return device_count + except Exception as e: + logger.debug("Could not import AscendCL: %s", e) + + try: + npu_files = glob.glob("/dev/davinci[0-9]*") + return len(npu_files) + except Exception as e: + logger.debug("Failed to detect number of NPUs: %s", e) + return 0 + + @staticmethod + def get_current_node_accelerator_type() -> Optional[str]: + """Get the type of the Ascend NPU on the current node. + + Returns: + A string of the type, such as "Ascend910A", "Ascend910B", "Ascend310P1". + """ + try: + import acl + + return acl.get_soc_name() + except Exception: + logger.exception("Failed to detect NPU type.") + return None + + @staticmethod + def validate_resource_request_quantity( + quantity: float, + ) -> Tuple[bool, Optional[str]]: + return (True, None) + + @staticmethod + def set_current_process_visible_accelerator_ids( + visible_npu_devices: List[str], + ) -> None: + if os.environ.get(NOSET_ASCEND_RT_VISIBLE_DEVICES_ENV_VAR): + return + + os.environ[ + NPUAcceleratorManager.get_visible_accelerator_ids_env_var() + ] = ",".join([str(i) for i in visible_npu_devices]) diff --git a/lib/python3.12/site-packages/ray/_private/accelerators/nvidia_gpu.py b/lib/python3.12/site-packages/ray/_private/accelerators/nvidia_gpu.py new file mode 100644 index 0000000000000000000000000000000000000000..bfd98ab377d6e71bd3123ef84c0db75f05031656 --- /dev/null +++ b/lib/python3.12/site-packages/ray/_private/accelerators/nvidia_gpu.py @@ -0,0 +1,128 @@ +import logging +import os +import re +from typing import List, Optional, Tuple + +from ray._private.accelerators.accelerator import AcceleratorManager + +logger = logging.getLogger(__name__) + +CUDA_VISIBLE_DEVICES_ENV_VAR = "CUDA_VISIBLE_DEVICES" +NOSET_CUDA_VISIBLE_DEVICES_ENV_VAR = "RAY_EXPERIMENTAL_NOSET_CUDA_VISIBLE_DEVICES" + +# TODO(Alex): This pattern may not work for non NVIDIA Tesla GPUs (which have +# the form "Tesla V100-SXM2-16GB" or "Tesla K80"). +NVIDIA_GPU_NAME_PATTERN = re.compile(r"\w+\s+([A-Z0-9]+)") + + +class NvidiaGPUAcceleratorManager(AcceleratorManager): + """NVIDIA GPU accelerators.""" + + @staticmethod + def get_resource_name() -> str: + return "GPU" + + @staticmethod + def get_visible_accelerator_ids_env_var() -> str: + return CUDA_VISIBLE_DEVICES_ENV_VAR + + @staticmethod + def get_current_process_visible_accelerator_ids() -> Optional[List[str]]: + cuda_visible_devices = os.environ.get( + NvidiaGPUAcceleratorManager.get_visible_accelerator_ids_env_var(), None + ) + if cuda_visible_devices is None: + return None + + if cuda_visible_devices == "": + return [] + + if cuda_visible_devices == "NoDevFiles": + return [] + + return list(cuda_visible_devices.split(",")) + + @staticmethod + def get_current_node_num_accelerators() -> int: + import ray._private.thirdparty.pynvml as pynvml + + try: + pynvml.nvmlInit() + except pynvml.NVMLError: + return 0 # pynvml init failed + device_count = pynvml.nvmlDeviceGetCount() + pynvml.nvmlShutdown() + return device_count + + @staticmethod + def get_current_node_accelerator_type() -> Optional[str]: + import ray._private.thirdparty.pynvml as pynvml + + try: + pynvml.nvmlInit() + except pynvml.NVMLError: + return None # pynvml init failed + device_count = pynvml.nvmlDeviceGetCount() + cuda_device_type = None + if device_count > 0: + handle = pynvml.nvmlDeviceGetHandleByIndex(0) + device_name = pynvml.nvmlDeviceGetName(handle) + if isinstance(device_name, bytes): + device_name = device_name.decode("utf-8") + cuda_device_type = ( + NvidiaGPUAcceleratorManager._gpu_name_to_accelerator_type(device_name) + ) + pynvml.nvmlShutdown() + return cuda_device_type + + @staticmethod + def _gpu_name_to_accelerator_type(name): + if name is None: + return None + match = NVIDIA_GPU_NAME_PATTERN.match(name) + return match.group(1) if match else None + + @staticmethod + def validate_resource_request_quantity( + quantity: float, + ) -> Tuple[bool, Optional[str]]: + return (True, None) + + @staticmethod + def set_current_process_visible_accelerator_ids( + visible_cuda_devices: List[str], + ) -> None: + if os.environ.get(NOSET_CUDA_VISIBLE_DEVICES_ENV_VAR): + return + + os.environ[ + NvidiaGPUAcceleratorManager.get_visible_accelerator_ids_env_var() + ] = ",".join([str(i) for i in visible_cuda_devices]) + + @staticmethod + def get_ec2_instance_num_accelerators( + instance_type: str, instances: dict + ) -> Optional[int]: + if instance_type not in instances: + return None + + gpus = instances[instance_type].get("GpuInfo", {}).get("Gpus") + if gpus is not None: + # TODO(ameer): currently we support one gpu type per node. + assert len(gpus) == 1 + return gpus[0]["Count"] + return None + + @staticmethod + def get_ec2_instance_accelerator_type( + instance_type: str, instances: dict + ) -> Optional[str]: + if instance_type not in instances: + return None + + gpus = instances[instance_type].get("GpuInfo", {}).get("Gpus") + if gpus is not None: + # TODO(ameer): currently we support one gpu type per node. + assert len(gpus) == 1 + return gpus[0]["Name"] + return None diff --git a/lib/python3.12/site-packages/ray/_private/accelerators/rbln.py b/lib/python3.12/site-packages/ray/_private/accelerators/rbln.py new file mode 100644 index 0000000000000000000000000000000000000000..ccebd7eedc5f716d0f0f1eb195a05d52eac9853b --- /dev/null +++ b/lib/python3.12/site-packages/ray/_private/accelerators/rbln.py @@ -0,0 +1,78 @@ +import logging +import os +from typing import List, Optional, Tuple + +from ray._private.accelerators.accelerator import AcceleratorManager + +logger = logging.getLogger(__name__) + +RBLN_RT_VISIBLE_DEVICES_ENV_VAR = "RBLN_DEVICES" +NOSET_RBLN_RT_VISIBLE_DEVICES_ENV_VAR = "RAY_EXPERIMENTAL_NOSET_RBLN_RT_VISIBLE_DEVICES" + + +class RBLNAcceleratorManager(AcceleratorManager): + """Rebellions RBLN accelerators.""" + + @staticmethod + def get_resource_name() -> str: + return "RBLN" + + @staticmethod + def get_visible_accelerator_ids_env_var() -> str: + return RBLN_RT_VISIBLE_DEVICES_ENV_VAR + + @staticmethod + def get_current_process_visible_accelerator_ids() -> Optional[List[str]]: + visible_devices = os.environ.get( + RBLNAcceleratorManager.get_visible_accelerator_ids_env_var() + ) + if visible_devices is None: + return None + if visible_devices == "": + return [] + return visible_devices.split(",") + + @staticmethod + def get_current_node_num_accelerators() -> int: + """Detects the number of RBLN devices on the current machine.""" + try: + from rebel import device_count + + return device_count() + except Exception as e: + logger.debug("Could not detect RBLN devices: %s", e) + return 0 + + @staticmethod + def get_current_node_accelerator_type() -> Optional[str]: + """Gets the type of RBLN NPU on the current node.""" + try: + from rebel import get_npu_name + + return get_npu_name() + except Exception as e: + logger.exception("Failed to detect RBLN NPU type: %s", e) + return None + + @staticmethod + def validate_resource_request_quantity( + quantity: float, + ) -> Tuple[bool, Optional[str]]: + if isinstance(quantity, float) and not quantity.is_integer(): + return ( + False, + f"{RBLNAcceleratorManager.get_resource_name()} resource quantity" + " must be whole numbers. " + f"The specified quantity {quantity} is invalid.", + ) + else: + return (True, None) + + @staticmethod + def set_current_process_visible_accelerator_ids( + visible_rbln_devices: List[str], + ) -> None: + if not os.getenv(NOSET_RBLN_RT_VISIBLE_DEVICES_ENV_VAR): + os.environ[ + RBLNAcceleratorManager.get_visible_accelerator_ids_env_var() + ] = ",".join(map(str, visible_rbln_devices)) diff --git a/lib/python3.12/site-packages/ray/_private/accelerators/tpu.py b/lib/python3.12/site-packages/ray/_private/accelerators/tpu.py new file mode 100644 index 0000000000000000000000000000000000000000..2115bf1d7f9ea7866486deeed46dd1b9fb656d9a --- /dev/null +++ b/lib/python3.12/site-packages/ray/_private/accelerators/tpu.py @@ -0,0 +1,672 @@ +import glob +import logging +import os +import re +from functools import lru_cache +from typing import Dict, List, Optional, Set, Tuple + +import requests + +import ray +from ray._private.accelerators.accelerator import AcceleratorManager +from ray.util.scheduling_strategies import PlacementGroupSchedulingStrategy + +logger = logging.getLogger(__name__) + + +TPU_VALID_CHIP_OPTIONS = (1, 2, 4, 8) +GKE_TPU_ACCELERATOR_TYPE_ENV_VAR = "TPU_ACCELERATOR_TYPE" +GKE_TPU_TOPOLOGY_ENV_VAR = "TPU_TOPOLOGY" +GKE_TPU_WORKER_ID_ENV_VAR = "TPU_WORKER_ID" +GKE_TPU_NAME_ENV_VAR = "TPU_NAME" + +# Constants for accessing the `accelerator-type` from TPU VM +# instance metadata. +# See https://cloud.google.com/compute/docs/metadata/overview +# for more details about VM instance metadata. +GCE_TPU_ACCELERATOR_ENDPOINT = ( + "http://metadata.google.internal/computeMetadata/v1/instance/attributes/" +) +GCE_TPU_HEADERS = {"Metadata-Flavor": "Google"} +GCE_TPU_ACCELERATOR_KEY = "accelerator-type" +GCE_TPU_ENV_KEY = "tpu-env" +GCE_TPU_INSTANCE_ID_KEY = "instance-id" +GCE_TPU_WORKER_ID_KEY = "agent-worker-number" + +TPU_VISIBLE_CHIPS_ENV_VAR = "TPU_VISIBLE_CHIPS" + +NOSET_TPU_VISIBLE_CHIPS_ENV_VAR = "RAY_EXPERIMENTAL_NOSET_TPU_VISIBLE_CHIPS" + +# The following defines environment variables that allow +# us to access a subset of TPU visible chips. +# +# See: https://github.com/google/jax/issues/14977 for an example/more details. +TPU_CHIPS_PER_HOST_BOUNDS_ENV_VAR = "TPU_CHIPS_PER_HOST_BOUNDS" +TPU_CHIPS_PER_HOST_BOUNDS_1_CHIP_CONFIG = "1,1,1" +TPU_CHIPS_PER_HOST_BOUNDS_2_CHIP_CONFIG = "1,2,1" + +TPU_HOST_BOUNDS_ENV_VAR = "TPU_HOST_BOUNDS" +TPU_SINGLE_HOST_BOUNDS = "1,1,1" + +# By default TPU VMs come with 4 chips per host and 2 tensorcores per chip. +# For more details: https://cloud.google.com/tpu/docs/system-architecture-tpu-vm +DEFAULT_TPU_NUM_CHIPS_PER_HOST = 4 +DEFAULT_TPU_NUM_CORES_PER_CHIP = 2 + +# Accelerators that are 4 chips per host: v2, v3, v4, v5p +# Accelerators that are 8 chips per host: v5e, v6e +SINGLE_HOST_8_CHIPS_TPU_TYPES = ("v5litepod", "v6e") + +# Accelerators that are 2 cores per chip: v2, v3, v4, v5p +# Accelerators that are 1 core per chip: v5e, v6e +SINGLE_CORE_TPU_TYPES = ("v5litepod", "v6e") + +# The valid TPU types. +VALID_TPU_TYPES = ("v2", "v3", "v4", "v5p", "v5litepod", "v6e") + +# This is only used to construct TPU 3D topologies +def _get_larger_3d_topologies(max_x: int, max_y: int, max_z: int) -> Set[str]: + """Returns a set of larger 3D TPU topologies given the max x,y,z value. Using DEFAULT_TPU_NUM_CHIPS_PER_HOST as increment""" + topologies = set() + for x in range( + DEFAULT_TPU_NUM_CHIPS_PER_HOST, max_x + 1, DEFAULT_TPU_NUM_CHIPS_PER_HOST + ): + for y in range( + DEFAULT_TPU_NUM_CHIPS_PER_HOST, max_y + 1, DEFAULT_TPU_NUM_CHIPS_PER_HOST + ): + for z in range( + DEFAULT_TPU_NUM_CHIPS_PER_HOST, + max_z + 1, + DEFAULT_TPU_NUM_CHIPS_PER_HOST, + ): + topologies.add(f"{x}x{y}x{z}") + + return topologies + + +# The valid TPU topologies for each of the TPU types +VALID_TPU_TOPOLOGY = { + "v2": {"4x4", "4x8", "8x8", "8x16", "16x16"}, + "v3": {"4x4", "4x8", "8x8", "8x16", "16x16", "16x32", "32x32"}, + "v4": {"2x2x1", "2x2x2", "2x2x4", "2x4x4"}.union( + _get_larger_3d_topologies(12, 12, 16) + ), + "v5p": { + "2x2x1", + "2x2x2", + "2x2x4", + "2x4x4", + }.union(_get_larger_3d_topologies(16, 16, 24)), + "v5litepod": {"2x8", "4x4", "4x8", "8x8", "8x16", "16x16"}, + "v6e": {"2x8", "4x4", "4x8", "8x8", "8x16", "16x16"}, +} + + +def _get_tpu_metadata(key: str) -> Optional[str]: + """Poll and get TPU metadata.""" + try: + accelerator_type_request = requests.get( + os.path.join(GCE_TPU_ACCELERATOR_ENDPOINT, key), + headers=GCE_TPU_HEADERS, + ) + if ( + accelerator_type_request.status_code == 200 + and accelerator_type_request.text + ): + return accelerator_type_request.text + else: + logging.debug( + "Unable to poll TPU GCE Metadata. Got " + f"status code: {accelerator_type_request.status_code} and " + f"content: {accelerator_type_request.text}" + ) + except requests.RequestException as e: + logging.debug("Unable to poll the TPU GCE Metadata: %s", e) + return None + + +def _accelerator_type_check(accelerator_type: str): + if not accelerator_type.startswith(VALID_TPU_TYPES): + raise ValueError( + f"Invalid accelerator type: {accelerator_type}. Must start with one of: {VALID_TPU_TYPES}" + ) + + +def get_num_tpu_visible_chips_per_host(accelerator_type: str) -> int: + _accelerator_type_check(accelerator_type) + if accelerator_type.startswith(SINGLE_HOST_8_CHIPS_TPU_TYPES): + return 8 + + return DEFAULT_TPU_NUM_CHIPS_PER_HOST + + +def get_tpu_cores_per_chip(accelerator_type: str) -> int: + _accelerator_type_check(accelerator_type) + if accelerator_type.startswith(SINGLE_CORE_TPU_TYPES): + return 1 + + return DEFAULT_TPU_NUM_CORES_PER_CHIP + + +def infer_tpu_pod_type_from_topology( + topology: str, accelerator_type: str +) -> Optional[str]: + """Infer the TPU pod type (e.g. v4-32) from topology and accelerator type.""" + if not topology or not accelerator_type: + return None + try: + num_chips = 1 + for value in topology.strip().lower().split("x"): + num_chips *= int(value) + generation = accelerator_type.lower().replace("tpu-", "") + return f"{generation}-{num_chips}" + except Exception as e: + raise ValueError( + f"Failed to infer pod type from topology '{topology}' " + f"and type '{accelerator_type}'" + ) from e + + +def fetch_tpu_slice_name_from_pg(pg): + @ray.remote(num_cpus=0) + def _get_tpu_slice_name(): + return TPUAcceleratorManager.get_current_node_tpu_name() + + tpu_name_ref = _get_tpu_slice_name.options( + scheduling_strategy=PlacementGroupSchedulingStrategy( + placement_group=pg, placement_group_bundle_index=0 + ) + ).remote() + + return ray.get(tpu_name_ref) + + +def get_chips_per_host(topology: str, accelerator_version: str) -> int: + """Get the number of chips per host (aka VMs) based on topology and accelerator version. + The current rule is as follows: + Default chips per host is 4. + If accelerator_version is v5e or v6e AND topology product <= 8, the chips per host will just be the proudct. i.e. 1, 4, or 8 + If accelerator_version is v5e or v6e AND topology product > 8, the chips per host will be 4 + If accelerator_version is v5p or other versions, the chips per host will be 4 + + Args: + topology: The TPU topology string (e.g. "2x2x2"). + accelerator_version: The accelerator version of the node (e.g. "V4", "v4"). + + Returns: + A int representing the number of chips per host (aka VM) + """ + chips_per_host = DEFAULT_TPU_NUM_CHIPS_PER_HOST + total_chips = 1 + for value in topology.strip().lower().split("x"): + total_chips *= int(value) + + if ( + total_chips <= 8 + and accelerator_version.strip().lower() in SINGLE_HOST_8_CHIPS_TPU_TYPES + ): + return total_chips + + return chips_per_host + + +def reserve_tpu_slice( + topology: str, + accelerator_type: str, +) -> Optional[str]: + """Reserves a TPU slice using its head resource and returns the slice name. + This enables gang scheduling of training workers with multi-host TPUs. + This is used by JaxTrainer with TPUs in Ray Train. + + Args: + topology: The TPU topology string (e.g. "2x2x2"). + accelerator_type: The accelerator type of the node (e.g. "TPU-V4"). + + Returns: + A string representing a unique TPU slice name. + """ + pod_type = infer_tpu_pod_type_from_topology(topology, accelerator_type) + if pod_type is None: + return None + + # Reserve a slice by creating a placement group on the TPU head. + head_label_selector = { + "ray.io/tpu-worker-id": "0", + "ray.io/tpu-pod-type": pod_type, + } + head_placement_group = ray.util.placement_group( + bundles=[{f"TPU-{pod_type}-head": 1}], + bundle_label_selector=[head_label_selector], + ) + + logger.debug("Waiting to reserve multi-host slice head.") + timeout = 100 + ready, _ = ray.wait([head_placement_group.ready()], timeout=timeout) + + if not ready: + raise TimeoutError( + "Failed to reserve TPU head for slice with shape: {}. " + "Ensure your cluster has sufficient resources. Requesting TPU " + "head node with labels: {}. Current resources: {}".format( + pod_type, head_label_selector, ray.available_resources() + ) + ) + + # Retrieve the unique slice ID. + slice_name = fetch_tpu_slice_name_from_pg(head_placement_group) + if slice_name is None: + raise RuntimeError( + "Failed to retrieve TPU slice name after reserving head placement group. " + "Ensure that TPU slice metadata is available and correctly configured on multi-host nodes." + ) + + # TODO: return both the slice name and reference to the PG reservation. + return slice_name + + +class TPUAcceleratorManager(AcceleratorManager): + """Google TPU accelerators.""" + + @staticmethod + def get_resource_name() -> str: + return "TPU" + + @staticmethod + def get_visible_accelerator_ids_env_var() -> str: + return TPU_VISIBLE_CHIPS_ENV_VAR + + @staticmethod + def get_current_process_visible_accelerator_ids() -> Optional[List[str]]: + tpu_visible_chips = os.environ.get( + TPUAcceleratorManager.get_visible_accelerator_ids_env_var(), None + ) + + if tpu_visible_chips is None: + return None + + if tpu_visible_chips == "": + return [] + + return list(tpu_visible_chips.split(",")) + + @staticmethod + @lru_cache() + def get_current_node_num_accelerators() -> int: + """Attempt to detect the number of TPUs on this machine. + + TPU chips are represented as devices within `/dev/`, either as + `/dev/accel*` or `/dev/vfio/*`. + + Returns: + The number of TPUs if any were detected, otherwise 0. + """ + accel_files = glob.glob("/dev/accel*") + if accel_files: + return len(accel_files) + + try: + vfio_entries = os.listdir("/dev/vfio") + numeric_entries = [int(entry) for entry in vfio_entries if entry.isdigit()] + return len(numeric_entries) + except FileNotFoundError as e: + logger.debug("Failed to detect number of TPUs: %s", e) + return 0 + + @staticmethod + def is_valid_tpu_accelerator_type(tpu_accelerator_type: str) -> bool: + """Check whether the tpu accelerator_type is formatted correctly. + + The accelerator_type field follows a form of v{generation}-{cores/chips}. + + See the following for more information: + https://cloud.google.com/sdk/gcloud/reference/compute/tpus/tpu-vm/accelerator-types/describe + + Args: + tpu_accelerator_type: The string representation of the accelerator type + to be checked for validity. + + Returns: + True if it's valid, false otherwise. + """ + expected_pattern = re.compile(r"^v\d+[a-zA-Z]*-\d+$") + if not expected_pattern.match(tpu_accelerator_type): + return False + return True + + @staticmethod + def is_valid_tpu_accelerator_topology( + tpu_accelerator_version: str, tpu_topology: str + ) -> bool: + """Check whether the tpu topology is valid. + + The accelerator_type field follows a form of v{generation}. + The accelerator_topology field follows either the form {A}x{B} or {A}x{B}x{C} depending on the v{generation} + + Args: + tpu_accelerator_version: The string representation of the accelerator version. (e.g. v6e, V5P) + tpu_topology: The string representation of the accelerator topology + to be checked for validity + + Returns: + True if it's valid topology, false othrwise + """ + tpu_version_formatted = tpu_accelerator_version.strip().lower().split("-")[0] + if ( + tpu_version_formatted.lower() not in VALID_TPU_TOPOLOGY + or tpu_topology.strip().lower() + not in VALID_TPU_TOPOLOGY[tpu_version_formatted] + ): + return False + return True + + @staticmethod + def validate_resource_request_quantity( + quantity: float, + ) -> Tuple[bool, Optional[str]]: + if quantity not in TPU_VALID_CHIP_OPTIONS: + return ( + False, + f"The number of requested 'TPU' was set to {quantity} which " + "is not a supported chip configuration. Supported configs: " + f"{TPU_VALID_CHIP_OPTIONS}", + ) + else: + return (True, None) + + @staticmethod + def set_current_process_visible_accelerator_ids( + visible_tpu_chips: List[str], + ) -> None: + """Set TPU environment variables based on the provided visible_tpu_chips. + + To access a subset of the TPU visible chips, we must use a combination of + environment variables that tells the compiler (via ML framework) the: + - Visible chips + - The physical bounds of chips per host + - The host bounds within the context of a TPU pod. + + See: https://github.com/google/jax/issues/14977 for an example/more details. + + Args: + visible_tpu_chips (List[str]): List of int representing TPU chips. + """ + if os.environ.get(NOSET_TPU_VISIBLE_CHIPS_ENV_VAR): + return + + num_visible_tpu_chips = len(visible_tpu_chips) + num_accelerators_on_node = ( + TPUAcceleratorManager.get_current_node_num_accelerators() + ) + if num_visible_tpu_chips == num_accelerators_on_node: + # Let the ML framework use the defaults + os.environ.pop(TPU_CHIPS_PER_HOST_BOUNDS_ENV_VAR, None) + os.environ.pop(TPU_HOST_BOUNDS_ENV_VAR, None) + return + os.environ[ + TPUAcceleratorManager.get_visible_accelerator_ids_env_var() + ] = ",".join([str(i) for i in visible_tpu_chips]) + if num_visible_tpu_chips == 1: + os.environ[ + TPU_CHIPS_PER_HOST_BOUNDS_ENV_VAR + ] = TPU_CHIPS_PER_HOST_BOUNDS_1_CHIP_CONFIG + os.environ[TPU_HOST_BOUNDS_ENV_VAR] = TPU_SINGLE_HOST_BOUNDS + elif num_visible_tpu_chips == 2: + os.environ[ + TPU_CHIPS_PER_HOST_BOUNDS_ENV_VAR + ] = TPU_CHIPS_PER_HOST_BOUNDS_2_CHIP_CONFIG + os.environ[TPU_HOST_BOUNDS_ENV_VAR] = TPU_SINGLE_HOST_BOUNDS + + @staticmethod + def get_current_node_tpu_pod_type() -> Optional[str]: + """Get the TPU pod type of the current node if applicable. + + Individual TPU VMs within a TPU pod must know what type + of pod it is a part of. This is necessary for the + ML framework to work properly. + + The logic is different if the TPU was provisioned via: + ``` + gcloud tpus tpu-vm create ... + ``` + (i.e. a GCE VM), vs through GKE: + - GCE VMs will always have a metadata server to poll this info + - GKE VMS will have environment variables preset. + + Returns: + A string representing the current TPU pod type, e.g. + v4-16. + + """ + # Start with GKE-based check + accelerator_type = os.getenv(GKE_TPU_ACCELERATOR_TYPE_ENV_VAR, "") + if not accelerator_type: + # GCE-based VM check + accelerator_type = _get_tpu_metadata(key=GCE_TPU_ACCELERATOR_KEY) + if accelerator_type and TPUAcceleratorManager.is_valid_tpu_accelerator_type( + tpu_accelerator_type=accelerator_type + ): + return accelerator_type + logging.debug("Failed to get a valid accelerator type.") + return None + + @staticmethod + def get_current_node_tpu_name() -> Optional[str]: + """Return the name of the TPU pod that this worker node is a part of. + + For instance, if the TPU was created with name "my-tpu", this function + will return "my-tpu". + + If created through the Ray cluster launcher, the + name will typically be something like "ray-my-tpu-cluster-worker-aa946781-tpu". + + In case the TPU was created through KubeRay, we currently expect that the + environment variable TPU_NAME is set per TPU pod slice, in which case + this function will return the value of that environment variable. + + """ + try: + # Start with GKE-based check + tpu_name = os.getenv(GKE_TPU_NAME_ENV_VAR, None) + if not tpu_name: + # GCE-based VM check + tpu_name = _get_tpu_metadata(key=GCE_TPU_INSTANCE_ID_KEY) + return tpu_name + except ValueError as e: + logging.debug("Could not get TPU name: %s", e) + return None + + @staticmethod + def get_current_node_tpu_worker_id() -> Optional[int]: + """Return the worker index of the TPU pod.""" + try: + # Start with GKE-based check + worker_id = os.getenv(GKE_TPU_WORKER_ID_ENV_VAR, None) + if not worker_id: + # GCE-based VM check + worker_id = _get_tpu_metadata(key=GCE_TPU_WORKER_ID_KEY) + if worker_id: + return int(worker_id) + else: + return None + except ValueError as e: + logging.debug("Could not get TPU worker id: %s", e) + return None + + @staticmethod + def get_num_workers_in_current_tpu_pod() -> Optional[int]: + """Return the total number of workers in a TPU pod.""" + tpu_pod_type = TPUAcceleratorManager.get_current_node_tpu_pod_type() + chips_per_host = TPUAcceleratorManager.get_current_node_num_accelerators() + cores_per_chip = get_tpu_cores_per_chip(tpu_pod_type) # Hard-coded map. + cores_per_host = chips_per_host * cores_per_chip + if tpu_pod_type and cores_per_host > 0: + num_cores = int(tpu_pod_type.split("-")[1]) + num_workers = num_cores // cores_per_host + # If the chip count doesn't fill a full host, a sub-host is still treated as a host. + if num_cores % cores_per_host != 0: + num_workers += 1 + return num_workers + else: + logging.debug("Could not get num workers in TPU pod.") + return None + + @staticmethod + def get_current_node_tpu_topology() -> Optional[str]: + try: + # Attempt GKE based lookup first + if topology := os.environ.get(GKE_TPU_TOPOLOGY_ENV_VAR): + return topology + # GCE-based VM check using TPU env string. + tpu_env = _get_tpu_metadata(key=GCE_TPU_ENV_KEY) + if tpu_env: + topology = re.search(r"TOPOLOGY:\s*'([^']+)'", tpu_env) + if topology: + return topology.group(1) + except ValueError as e: + logging.debug("Could not get TPU topology: %s", e) + return None + + @staticmethod + def get_current_node_accelerator_type() -> Optional[str]: + """Attempt to detect the TPU accelerator type. + + The output of this function will return the "ray accelerator type" + resource (e.g. TPU-V4) that indicates the TPU version. + + We also expect that our TPU nodes contain a "TPU pod type" + resource, which indicates information about the topology of + the TPU pod slice. + + We expect that the "TPU pod type" resource to be used when + running multi host workers, i.e. when TPU units are pod slices. + + We expect that the "ray accelerator type" resource to be used when + running single host workers, i.e. when TPU units are single hosts. + + Returns: + A string representing the TPU accelerator type, + e.g. "TPU-V2", "TPU-V3", "TPU-V4" if applicable, else None. + + """ + + def tpu_pod_type_to_ray_accelerator_type( + tpu_pod_type: str, + ) -> Optional[str]: + return "TPU-" + str(tpu_pod_type.split("-")[0].upper()) + + ray_accelerator_type = None + tpu_pod_type = TPUAcceleratorManager.get_current_node_tpu_pod_type() + + if tpu_pod_type is not None: + ray_accelerator_type = tpu_pod_type_to_ray_accelerator_type( + tpu_pod_type=tpu_pod_type + ) + if ray_accelerator_type is None: + logger.info( + "While trying to autodetect a TPU type, " + f"received malformed accelerator_type: {tpu_pod_type}" + ) + + if ray_accelerator_type is None: + logging.info("Failed to auto-detect TPU type.") + + return ray_accelerator_type + + @staticmethod + def get_current_node_additional_resources() -> Optional[Dict[str, float]]: + """Get additional resources required for TPU nodes. + + This will populate the TPU pod type and the TPU name which + is used for TPU pod execution. + + When running workloads on a TPU pod, we need a way to run + the same binary on every worker in the TPU pod. + + See https://jax.readthedocs.io/en/latest/multi_process.html + for more information. + + To do this in ray, we take advantage of custom resources. We + mark worker 0 of the TPU pod as a "coordinator" that identifies + the other workers in the TPU pod. We therefore need: + - worker 0 to be targetable. + - all workers in the TPU pod to have a unique identifier consistent + within a TPU pod. + + So assuming we want to run the following workload: + + @ray.remote + def my_jax_fn(): + import jax + return jax.device_count() + + We could broadcast this on a TPU pod (e.g. a v4-16) as follows: + + @ray.remote(resources={"TPU-v4-16-head"}) + def run_jax_fn(executable): + # Note this will execute on worker 0 + tpu_name = ray.util.tpu.get_tpu_pod_name() + num_workers = ray.util.tpu.get_tpu_num_workers() + tpu_executable = executable.options(resources={"TPU": 4, tpu_name: 1}) + return [tpu_executable.remote() for _ in range(num_workers)] + + Returns: + A dictionary representing additional resources that may be + necessary for a particular accelerator type. + + """ + resources = {} + tpu_name = TPUAcceleratorManager.get_current_node_tpu_name() + worker_id = TPUAcceleratorManager.get_current_node_tpu_worker_id() + tpu_pod_type = TPUAcceleratorManager.get_current_node_tpu_pod_type() + + if tpu_name and worker_id is not None and tpu_pod_type: + pod_head_resource_name = f"TPU-{tpu_pod_type}-head" + # Add the name of the TPU to the resource. + resources[tpu_name] = 1 + # Only add in the TPU pod type resource to worker 0. + if worker_id == 0: + resources[pod_head_resource_name] = 1 + else: + logging.info( + "Failed to configure TPU pod. Got: " + "tpu_name: %s, worker_id: %s, accelerator_type: %s", + tpu_name, + worker_id, + tpu_pod_type, + ) + if resources: + return resources + return None + + @staticmethod + def get_current_node_accelerator_labels() -> Dict[str, str]: + """Get default TPU-specific Ray node labels for the current node. + + For TPUs, these labels include: + - ray.io/tpu-slice-name: the name of the TPU Pod or slice + - ray.io/tpu-worker-id: the integer worker ID within the slice + - ray.io/tpu-topology: the TPU topology (e.g. 4x4) + - ray.io/tpu-pod-type: the TPU pod type (e.g. v4-8) + + Returns: + A dictionary of TPU label keys and resolved values. + """ + tpu_labels = {} + + tpu_name = TPUAcceleratorManager.get_current_node_tpu_name() + if tpu_name: + tpu_labels[ray._raylet.RAY_NODE_TPU_SLICE_NAME_KEY] = tpu_name + + worker_id = TPUAcceleratorManager.get_current_node_tpu_worker_id() + if worker_id is not None: + tpu_labels[ray._raylet.RAY_NODE_TPU_WORKER_ID_KEY] = str(worker_id) + + tpu_topology = TPUAcceleratorManager.get_current_node_tpu_topology() + if tpu_topology: + tpu_labels[ray._raylet.RAY_NODE_TPU_TOPOLOGY_KEY] = tpu_topology + + pod_type = TPUAcceleratorManager.get_current_node_tpu_pod_type() + if pod_type: + tpu_labels[ray._raylet.RAY_NODE_TPU_POD_TYPE_KEY] = pod_type + + return tpu_labels diff --git a/lib/python3.12/site-packages/ray/_private/arrow_serialization.py b/lib/python3.12/site-packages/ray/_private/arrow_serialization.py new file mode 100644 index 0000000000000000000000000000000000000000..9af0f9e7c50c40affbb019b7d6dafcaff6febcc7 --- /dev/null +++ b/lib/python3.12/site-packages/ray/_private/arrow_serialization.py @@ -0,0 +1,774 @@ +# arrow_serialization.py must resides outside of ray.data, otherwise +# it causes circular dependency issues for AsyncActors due to +# ray.data's lazy import. +# see https://github.com/ray-project/ray/issues/30498 for more context. +import logging +import os +import sys +from dataclasses import dataclass +from typing import TYPE_CHECKING, List, Optional, Tuple + +from ray._private.utils import is_in_test + +if TYPE_CHECKING: + import pyarrow + +RAY_DISABLE_CUSTOM_ARROW_JSON_OPTIONS_SERIALIZATION = ( + "RAY_DISABLE_CUSTOM_ARROW_JSON_OPTIONS_SERIALIZATION" +) +RAY_DISABLE_CUSTOM_ARROW_DATA_SERIALIZATION = ( + "RAY_DISABLE_CUSTOM_ARROW_DATA_SERIALIZATION" +) + +logger = logging.getLogger(__name__) + +# Whether we have already warned the user about bloated fallback serialization. +_serialization_fallback_set = set() + + +def _register_custom_datasets_serializers(serialization_context): + try: + import pyarrow as pa # noqa: F401 + except ModuleNotFoundError: + # No pyarrow installed so not using Arrow, so no need for custom serializers. + return + + # Register all custom serializers required by Datasets. + _register_arrow_data_serializer(serialization_context) + _register_arrow_json_readoptions_serializer(serialization_context) + _register_arrow_json_parseoptions_serializer(serialization_context) + + +# Register custom Arrow JSON ReadOptions serializer to workaround it not being picklable +# in Arrow < 8.0.0. +def _register_arrow_json_readoptions_serializer(serialization_context): + if ( + os.environ.get( + RAY_DISABLE_CUSTOM_ARROW_JSON_OPTIONS_SERIALIZATION, + "0", + ) + == "1" + ): + return + + import pyarrow.json as pajson + + serialization_context._register_cloudpickle_serializer( + pajson.ReadOptions, + custom_serializer=lambda opts: (opts.use_threads, opts.block_size), + custom_deserializer=lambda args: pajson.ReadOptions(*args), + ) + + +def _register_arrow_json_parseoptions_serializer(serialization_context): + if ( + os.environ.get( + RAY_DISABLE_CUSTOM_ARROW_JSON_OPTIONS_SERIALIZATION, + "0", + ) + == "1" + ): + return + + import pyarrow.json as pajson + + serialization_context._register_cloudpickle_serializer( + pajson.ParseOptions, + custom_serializer=lambda opts: ( + opts.explicit_schema, + opts.newlines_in_values, + opts.unexpected_field_behavior, + ), + custom_deserializer=lambda args: pajson.ParseOptions(*args), + ) + + +# Register custom Arrow data serializer to work around zero-copy slice pickling bug. +# See https://issues.apache.org/jira/browse/ARROW-10739. +def _register_arrow_data_serializer(serialization_context): + """Custom reducer for Arrow data that works around a zero-copy slicing pickling + bug by using the Arrow IPC format for the underlying serialization. + + Background: + Arrow has both array-level slicing and buffer-level slicing; both are zero-copy, + but the former has a serialization bug where the entire buffer is serialized + instead of just the slice, while the latter's serialization works as expected + and only serializes the slice of the buffer. I.e., array-level slicing doesn't + propagate the slice down to the buffer when serializing the array. + + We work around this by registering a custom cloudpickle reducers for Arrow + Tables that delegates serialization to the Arrow IPC format; thankfully, Arrow's + IPC serialization has fixed this buffer truncation bug. + + See https://issues.apache.org/jira/browse/ARROW-10739. + """ + if os.environ.get(RAY_DISABLE_CUSTOM_ARROW_DATA_SERIALIZATION, "0") == "1": + return + + import pyarrow as pa + + serialization_context._register_cloudpickle_reducer(pa.Table, _arrow_table_reduce) + + +def _arrow_table_reduce(t: "pyarrow.Table"): + """Custom reducer for Arrow Tables that works around a zero-copy slice pickling bug. + Background: + Arrow has both array-level slicing and buffer-level slicing; both are zero-copy, + but the former has a serialization bug where the entire buffer is serialized + instead of just the slice, while the latter's serialization works as expected + and only serializes the slice of the buffer. I.e., array-level slicing doesn't + propagate the slice down to the buffer when serializing the array. + All that these copy methods do is, at serialization time, take the array-level + slicing and translate them to buffer-level slicing, so only the buffer slice is + sent over the wire instead of the entire buffer. + See https://issues.apache.org/jira/browse/ARROW-10739. + """ + global _serialization_fallback_set + + # Reduce the ChunkedArray columns. + reduced_columns = [] + for column_name in t.column_names: + column = t[column_name] + try: + # Delegate to ChunkedArray reducer. + reduced_column = _arrow_chunked_array_reduce(column) + except Exception as e: + if not _is_dense_union(column.type) and is_in_test(): + # If running in a test and the column is not a dense union array + # (which we expect to need a fallback), we want to raise the error, + # not fall back. + raise e from None + if type(column.type) not in _serialization_fallback_set: + logger.warning( + "Failed to complete optimized serialization of Arrow Table, " + f"serialization of column '{column_name}' of type {column.type} " + "failed, so we're falling back to Arrow IPC serialization for the " + "table. Note that this may result in slower serialization and more " + "worker memory utilization. Serialization error:", + exc_info=True, + ) + _serialization_fallback_set.add(type(column.type)) + # Fall back to Arrow IPC-based workaround for the entire table. + return _arrow_table_ipc_reduce(t) + else: + # Column reducer succeeded, add reduced column to list. + reduced_columns.append(reduced_column) + return _reconstruct_table, (reduced_columns, t.schema) + + +def _reconstruct_table( + reduced_columns: List[Tuple[List["pyarrow.Array"], "pyarrow.DataType"]], + schema: "pyarrow.Schema", +) -> "pyarrow.Table": + """Restore a serialized Arrow Table, reconstructing each reduced column.""" + import pyarrow as pa + + # Reconstruct each reduced column. + columns = [] + for chunks_payload, type_ in reduced_columns: + columns.append(_reconstruct_chunked_array(chunks_payload, type_)) + + return pa.Table.from_arrays(columns, schema=schema) + + +def _arrow_chunked_array_reduce( + ca: "pyarrow.ChunkedArray", +) -> Tuple[List["PicklableArrayPayload"], "pyarrow.DataType"]: + """Custom reducer for Arrow ChunkedArrays that works around a zero-copy slice + pickling bug. This reducer does not return a reconstruction function, since it's + expected to be reconstructed by the Arrow Table reconstructor. + """ + # Convert chunks to serialization payloads. + chunk_payloads = [] + for chunk in ca.chunks: + chunk_payload = PicklableArrayPayload.from_array(chunk) + chunk_payloads.append(chunk_payload) + return chunk_payloads, ca.type + + +def _reconstruct_chunked_array( + chunks: List["PicklableArrayPayload"], type_: "pyarrow.DataType" +) -> "pyarrow.ChunkedArray": + """Restore a serialized Arrow ChunkedArray from chunks and type.""" + import pyarrow as pa + + # Reconstruct chunks from serialization payloads. + chunks = [chunk.to_array() for chunk in chunks] + + return pa.chunked_array(chunks, type_) + + +@dataclass +class PicklableArrayPayload: + """Picklable array payload, holding data buffers and array metadata. + + This is a helper container for pickling and reconstructing nested Arrow Arrays while + ensuring that the buffers that underly zero-copy slice views are properly truncated. + """ + + # Array type. + type: "pyarrow.DataType" + # Length of array. + length: int + # Underlying data buffers. + buffers: List["pyarrow.Buffer"] + # Cached null count. + null_count: int + # Slice offset into base array. + offset: int + # Serialized array payloads for nested (child) arrays. + children: List["PicklableArrayPayload"] + + @classmethod + def from_array(self, a: "pyarrow.Array") -> "PicklableArrayPayload": + """Create a picklable array payload from an Arrow Array. + + This will recursively accumulate data buffer and metadata payloads that are + ready for pickling; namely, the data buffers underlying zero-copy slice views + will be properly truncated. + """ + return _array_to_array_payload(a) + + def to_array(self) -> "pyarrow.Array": + """Reconstruct an Arrow Array from this picklable payload.""" + return _array_payload_to_array(self) + + +def _array_payload_to_array(payload: "PicklableArrayPayload") -> "pyarrow.Array": + """Reconstruct an Arrow Array from a possibly nested PicklableArrayPayload.""" + import pyarrow as pa + + children = [child_payload.to_array() for child_payload in payload.children] + + if pa.types.is_dictionary(payload.type): + # Dedicated path for reconstructing a DictionaryArray, since + # Array.from_buffers() doesn't work for DictionaryArrays. + assert len(children) == 2, len(children) + indices, dictionary = children + return pa.DictionaryArray.from_arrays(indices, dictionary) + elif pa.types.is_map(payload.type) and len(children) > 1: + # In pyarrow<7.0.0, the underlying map child array is not exposed, so we work + # with the key and item arrays. + assert len(children) == 3, len(children) + offsets, keys, items = children + return pa.MapArray.from_arrays(offsets, keys, items) + elif isinstance(payload.type, pa.BaseExtensionType): + assert len(children) == 1, len(children) + storage = children[0] + return payload.type.wrap_array(storage) + else: + # Common case: use Array.from_buffers() to construct an array of a certain type. + return pa.Array.from_buffers( + type=payload.type, + length=payload.length, + buffers=payload.buffers, + null_count=payload.null_count, + offset=payload.offset, + children=children, + ) + + +def _array_to_array_payload(a: "pyarrow.Array") -> "PicklableArrayPayload": + """Serialize an Arrow Array to an PicklableArrayPayload for later pickling. + + This function's primary purpose is to dispatch to the handler for the input array + type. + """ + import pyarrow as pa + + if _is_dense_union(a.type): + # Dense unions are not supported. + # TODO(Clark): Support dense unions. + raise NotImplementedError( + "Custom slice view serialization of dense union arrays is not yet " + "supported." + ) + + # Dispatch to handler for array type. + if pa.types.is_null(a.type): + return _null_array_to_array_payload(a) + elif _is_primitive(a.type): + return _primitive_array_to_array_payload(a) + elif _is_binary(a.type): + return _binary_array_to_array_payload(a) + elif pa.types.is_list(a.type) or pa.types.is_large_list(a.type): + return _list_array_to_array_payload(a) + elif pa.types.is_fixed_size_list(a.type): + return _fixed_size_list_array_to_array_payload(a) + elif pa.types.is_struct(a.type): + return _struct_array_to_array_payload(a) + elif pa.types.is_union(a.type): + return _union_array_to_array_payload(a) + elif pa.types.is_dictionary(a.type): + return _dictionary_array_to_array_payload(a) + elif pa.types.is_map(a.type): + return _map_array_to_array_payload(a) + elif isinstance(a.type, pa.BaseExtensionType): + return _extension_array_to_array_payload(a) + else: + raise ValueError("Unhandled Arrow array type:", a.type) + + +def _is_primitive(type_: "pyarrow.DataType") -> bool: + """Whether the provided Array type is primitive (boolean, numeric, temporal or + fixed-size binary).""" + import pyarrow as pa + + return ( + pa.types.is_integer(type_) + or pa.types.is_floating(type_) + or pa.types.is_decimal(type_) + or pa.types.is_boolean(type_) + or pa.types.is_temporal(type_) + or pa.types.is_fixed_size_binary(type_) + ) + + +def _is_binary(type_: "pyarrow.DataType") -> bool: + """Whether the provided Array type is a variable-sized binary type.""" + import pyarrow as pa + + return ( + pa.types.is_string(type_) + or pa.types.is_large_string(type_) + or pa.types.is_binary(type_) + or pa.types.is_large_binary(type_) + ) + + +def _null_array_to_array_payload(a: "pyarrow.NullArray") -> "PicklableArrayPayload": + """Serialize null array to PicklableArrayPayload.""" + # Buffer scheme: [None] + return PicklableArrayPayload( + type=a.type, + length=len(a), + buffers=[None], # Single null buffer is expected. + null_count=a.null_count, + offset=0, + children=[], + ) + + +def _primitive_array_to_array_payload(a: "pyarrow.Array") -> "PicklableArrayPayload": + """Serialize primitive (numeric, temporal, boolean) arrays to + PicklableArrayPayload. + """ + assert _is_primitive(a.type), a.type + # Buffer scheme: [bitmap, data] + buffers = a.buffers() + assert len(buffers) == 2, len(buffers) + + # Copy bitmap buffer, if needed. + bitmap_buf = buffers[0] + if a.null_count > 0: + bitmap_buf = _copy_bitpacked_buffer_if_needed(bitmap_buf, a.offset, len(a)) + else: + bitmap_buf = None + + # Copy data buffer, if needed. + data_buf = buffers[1] + if data_buf is not None: + data_buf = _copy_buffer_if_needed(buffers[1], a.type, a.offset, len(a)) + + return PicklableArrayPayload( + type=a.type, + length=len(a), + buffers=[bitmap_buf, data_buf], + null_count=a.null_count, + offset=0, + children=[], + ) + + +def _binary_array_to_array_payload(a: "pyarrow.Array") -> "PicklableArrayPayload": + """Serialize binary (variable-sized binary, string) arrays to + PicklableArrayPayload. + """ + assert _is_binary(a.type), a.type + # Buffer scheme: [bitmap, value_offsets, data] + buffers = a.buffers() + assert len(buffers) == 3, len(buffers) + + # Copy bitmap buffer, if needed. + if a.null_count > 0: + bitmap_buf = _copy_bitpacked_buffer_if_needed(buffers[0], a.offset, len(a)) + else: + bitmap_buf = None + + # Copy offset buffer, if needed. + offset_buf = buffers[1] + offset_buf, data_offset, data_length = _copy_offsets_buffer_if_needed( + offset_buf, a.type, a.offset, len(a) + ) + data_buf = buffers[2] + data_buf = _copy_buffer_if_needed(data_buf, None, data_offset, data_length) + return PicklableArrayPayload( + type=a.type, + length=len(a), + buffers=[bitmap_buf, offset_buf, data_buf], + null_count=a.null_count, + offset=0, + children=[], + ) + + +def _list_array_to_array_payload(a: "pyarrow.Array") -> "PicklableArrayPayload": + """Serialize list (regular and large) arrays to PicklableArrayPayload.""" + # Dedicated path for ListArrays. These arrays have a nested set of bitmap and + # offset buffers, eventually bottoming out on a data buffer. + # Buffer scheme: + # [bitmap, offsets, bitmap, offsets, ..., bitmap, data] + buffers = a.buffers() + assert len(buffers) > 1, len(buffers) + + # Copy bitmap buffer, if needed. + if a.null_count > 0: + bitmap_buf = _copy_bitpacked_buffer_if_needed(buffers[0], a.offset, len(a)) + else: + bitmap_buf = None + + # Copy offset buffer, if needed. + offset_buf = buffers[1] + offset_buf, child_offset, child_length = _copy_offsets_buffer_if_needed( + offset_buf, a.type, a.offset, len(a) + ) + + # Propagate slice to child. + child = a.values.slice(child_offset, child_length) + + return PicklableArrayPayload( + type=a.type, + length=len(a), + buffers=[bitmap_buf, offset_buf], + null_count=a.null_count, + offset=0, + children=[_array_to_array_payload(child)], + ) + + +def _fixed_size_list_array_to_array_payload( + a: "pyarrow.FixedSizeListArray", +) -> "PicklableArrayPayload": + """Serialize fixed size list arrays to PicklableArrayPayload.""" + # Dedicated path for fixed-size lists. + # Buffer scheme: + # [bitmap, values_bitmap, values_data, values_subbuffers...] + buffers = a.buffers() + assert len(buffers) >= 1, len(buffers) + + # Copy bitmap buffer, if needed. + if a.null_count > 0: + bitmap_buf = _copy_bitpacked_buffer_if_needed(buffers[0], a.offset, len(a)) + else: + bitmap_buf = None + + # Propagate slice to child. + child_offset = a.type.list_size * a.offset + child_length = a.type.list_size * len(a) + child = a.values.slice(child_offset, child_length) + + return PicklableArrayPayload( + type=a.type, + length=len(a), + buffers=[bitmap_buf], + null_count=a.null_count, + offset=0, + children=[_array_to_array_payload(child)], + ) + + +def _struct_array_to_array_payload(a: "pyarrow.StructArray") -> "PicklableArrayPayload": + """Serialize struct arrays to PicklableArrayPayload.""" + # Dedicated path for StructArrays. + # StructArrays have a top-level bitmap buffer and one or more children arrays. + # Buffer scheme: [bitmap, None, child_bitmap, child_data, ...] + buffers = a.buffers() + assert len(buffers) >= 1, len(buffers) + + # Copy bitmap buffer, if needed. + if a.null_count > 0: + bitmap_buf = _copy_bitpacked_buffer_if_needed(buffers[0], a.offset, len(a)) + else: + bitmap_buf = None + + # Get field children payload. + # Offsets and truncations are already propagated to the field arrays, so we can + # serialize them as-is. + children = [_array_to_array_payload(a.field(i)) for i in range(a.type.num_fields)] + return PicklableArrayPayload( + type=a.type, + length=len(a), + buffers=[bitmap_buf], + null_count=a.null_count, + offset=0, + children=children, + ) + + +def _union_array_to_array_payload(a: "pyarrow.UnionArray") -> "PicklableArrayPayload": + """Serialize union arrays to PicklableArrayPayload.""" + import pyarrow as pa + + # Dedicated path for UnionArrays. + # UnionArrays have a top-level bitmap buffer and type code buffer, and one or + # more children arrays. + # Buffer scheme: [None, typecodes, child_bitmap, child_data, ...] + assert not _is_dense_union(a.type) + buffers = a.buffers() + assert len(buffers) > 1, len(buffers) + + bitmap_buf = buffers[0] + assert bitmap_buf is None, bitmap_buf + + # Copy type code buffer, if needed. + type_code_buf = buffers[1] + type_code_buf = _copy_buffer_if_needed(type_code_buf, pa.int8(), a.offset, len(a)) + + # Get field children payload. + # Offsets and truncations are already propagated to the field arrays, so we can + # serialize them as-is. + children = [_array_to_array_payload(a.field(i)) for i in range(a.type.num_fields)] + return PicklableArrayPayload( + type=a.type, + length=len(a), + buffers=[bitmap_buf, type_code_buf], + null_count=a.null_count, + offset=0, + children=children, + ) + + +def _dictionary_array_to_array_payload( + a: "pyarrow.DictionaryArray", +) -> "PicklableArrayPayload": + """Serialize dictionary arrays to PicklableArrayPayload.""" + # Dedicated path for DictionaryArrays. + # Buffer scheme: [indices_bitmap, indices_data] (dictionary stored separately) + indices_payload = _array_to_array_payload(a.indices) + dictionary_payload = _array_to_array_payload(a.dictionary) + return PicklableArrayPayload( + type=a.type, + length=len(a), + buffers=[], + null_count=a.null_count, + offset=0, + children=[indices_payload, dictionary_payload], + ) + + +def _map_array_to_array_payload(a: "pyarrow.MapArray") -> "PicklableArrayPayload": + """Serialize map arrays to PicklableArrayPayload.""" + import pyarrow as pa + + # Dedicated path for MapArrays. + # Buffer scheme: [bitmap, offsets, child_struct_array_buffers, ...] + buffers = a.buffers() + assert len(buffers) > 0, len(buffers) + + # Copy bitmap buffer, if needed. + if a.null_count > 0: + bitmap_buf = _copy_bitpacked_buffer_if_needed(buffers[0], a.offset, len(a)) + else: + bitmap_buf = None + + new_buffers = [bitmap_buf] + + # Copy offsets buffer, if needed. + offset_buf = buffers[1] + offset_buf, data_offset, data_length = _copy_offsets_buffer_if_needed( + offset_buf, a.type, a.offset, len(a) + ) + + if isinstance(a, pa.lib.ListArray): + # Map arrays directly expose the one child struct array in pyarrow>=7.0.0, which + # is easier to work with than the raw buffers. + new_buffers.append(offset_buf) + children = [_array_to_array_payload(a.values.slice(data_offset, data_length))] + else: + # In pyarrow<7.0.0, the child struct array is not exposed, so we work with the + # key and item arrays. + buffers = a.buffers() + assert len(buffers) > 2, len(buffers) + # Reconstruct offsets array. + offsets = pa.Array.from_buffers( + pa.int32(), len(a) + 1, [bitmap_buf, offset_buf] + ) + # Propagate slice to keys. + keys = a.keys.slice(data_offset, data_length) + # Propagate slice to items. + items = a.items.slice(data_offset, data_length) + children = [ + _array_to_array_payload(offsets), + _array_to_array_payload(keys), + _array_to_array_payload(items), + ] + return PicklableArrayPayload( + type=a.type, + length=len(a), + buffers=new_buffers, + null_count=a.null_count, + offset=0, + children=children, + ) + + +def _extension_array_to_array_payload( + a: "pyarrow.ExtensionArray", +) -> "PicklableArrayPayload": + storage_payload = _array_to_array_payload(a.storage) + return PicklableArrayPayload( + type=a.type, + length=len(a), + buffers=[], + null_count=a.null_count, + offset=0, + children=[storage_payload], + ) + + +def _copy_buffer_if_needed( + buf: "pyarrow.Buffer", + type_: Optional["pyarrow.DataType"], + offset: int, + length: int, +) -> "pyarrow.Buffer": + """Copy buffer, if needed.""" + import pyarrow as pa + + if type_ is not None and pa.types.is_boolean(type_): + # Arrow boolean array buffers are bit-packed, with 8 entries per byte, + # and are accessed via bit offsets. + buf = _copy_bitpacked_buffer_if_needed(buf, offset, length) + else: + type_bytewidth = type_.bit_width // 8 if type_ is not None else 1 + buf = _copy_normal_buffer_if_needed(buf, type_bytewidth, offset, length) + return buf + + +def _copy_normal_buffer_if_needed( + buf: "pyarrow.Buffer", + byte_width: int, + offset: int, + length: int, +) -> "pyarrow.Buffer": + """Copy buffer, if needed.""" + byte_offset = offset * byte_width + byte_length = length * byte_width + if offset > 0 or byte_length < buf.size: + # Array is a zero-copy slice, so we need to copy to a new buffer before + # serializing; this slice of the underlying buffer (not the array) will ensure + # that the buffer is properly copied at pickle-time. + buf = buf.slice(byte_offset, byte_length) + return buf + + +def _copy_bitpacked_buffer_if_needed( + buf: "pyarrow.Buffer", + offset: int, + length: int, +) -> "pyarrow.Buffer": + """Copy bit-packed binary buffer, if needed.""" + bit_offset = offset % 8 + byte_offset = offset // 8 + byte_length = _bytes_for_bits(bit_offset + length) // 8 + if offset > 0 or byte_length < buf.size: + buf = buf.slice(byte_offset, byte_length) + if bit_offset != 0: + # Need to manually shift the buffer to eliminate the bit offset. + buf = _align_bit_offset(buf, bit_offset, byte_length) + return buf + + +def _copy_offsets_buffer_if_needed( + buf: "pyarrow.Buffer", + arr_type: "pyarrow.DataType", + offset: int, + length: int, +) -> Tuple["pyarrow.Buffer", int, int]: + """Copy the provided offsets buffer, returning the copied buffer and the + offset + length of the underlying data. + """ + import pyarrow as pa + import pyarrow.compute as pac + + if ( + pa.types.is_large_list(arr_type) + or pa.types.is_large_string(arr_type) + or pa.types.is_large_binary(arr_type) + or pa.types.is_large_unicode(arr_type) + ): + offset_type = pa.int64() + else: + offset_type = pa.int32() + # Copy offset buffer, if needed. + buf = _copy_buffer_if_needed(buf, offset_type, offset, length + 1) + # Reconstruct the offset array so we can determine the offset and length + # of the child array. + offsets = pa.Array.from_buffers(offset_type, length + 1, [None, buf]) + child_offset = offsets[0].as_py() + child_length = offsets[-1].as_py() - child_offset + # Create new offsets aligned to 0 for the copied data buffer slice. + offsets = pac.subtract(offsets, child_offset) + if pa.types.is_int32(offset_type): + # We need to cast the resulting Int64Array back down to an Int32Array. + offsets = offsets.cast(offset_type, safe=False) + buf = offsets.buffers()[1] + return buf, child_offset, child_length + + +def _bytes_for_bits(n: int) -> int: + """Round up n to the nearest multiple of 8. + This is used to get the byte-padded number of bits for n bits. + """ + return (n + 7) & (-8) + + +def _align_bit_offset( + buf: "pyarrow.Buffer", + bit_offset: int, + byte_length: int, +) -> "pyarrow.Buffer": + """Align the bit offset into the buffer with the front of the buffer by shifting + the buffer and eliminating the offset. + """ + import pyarrow as pa + + bytes_ = buf.to_pybytes() + bytes_as_int = int.from_bytes(bytes_, sys.byteorder) + bytes_as_int >>= bit_offset + bytes_ = bytes_as_int.to_bytes(byte_length, sys.byteorder) + return pa.py_buffer(bytes_) + + +def _arrow_table_ipc_reduce(table: "pyarrow.Table"): + """Custom reducer for Arrow Table that works around a zero-copy slicing pickling + bug by using the Arrow IPC format for the underlying serialization. + + This is currently used as a fallback for unsupported types (or unknown bugs) for + the manual buffer truncation workaround, e.g. for dense unions. + """ + from pyarrow.ipc import RecordBatchStreamWriter + from pyarrow.lib import BufferOutputStream + + output_stream = BufferOutputStream() + with RecordBatchStreamWriter(output_stream, schema=table.schema) as wr: + wr.write_table(table) + # NOTE: output_stream.getvalue() materializes the serialized table to a single + # contiguous bytestring, resulting in a few copy. This adds 1-2 extra copies on the + # serialization side, and 1 extra copy on the deserialization side. + return _restore_table_from_ipc, (output_stream.getvalue(),) + + +def _restore_table_from_ipc(buf: bytes) -> "pyarrow.Table": + """Restore an Arrow Table serialized to Arrow IPC format.""" + from pyarrow.ipc import RecordBatchStreamReader + + with RecordBatchStreamReader(buf) as reader: + return reader.read_all() + + +def _is_dense_union(type_: "pyarrow.DataType") -> bool: + """Whether the provided Arrow type is a dense union.""" + import pyarrow as pa + + return pa.types.is_union(type_) and type_.mode == "dense" diff --git a/lib/python3.12/site-packages/ray/_private/async_utils.py b/lib/python3.12/site-packages/ray/_private/async_utils.py new file mode 100644 index 0000000000000000000000000000000000000000..0c2da0a1ec3054effd2596086bc71296874afa49 --- /dev/null +++ b/lib/python3.12/site-packages/ray/_private/async_utils.py @@ -0,0 +1,52 @@ +# Adapted from [aiodebug](https://gitlab.com/quantlane/libs/aiodebug) + +# Copyright 2016-2022 Quantlane s.r.o. + +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at + +# http://www.apache.org/licenses/LICENSE-2.0 + +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +# Modifications: +# - Removed the dependency to `logwood`. +# - Renamed `monitor_loop_lag.enable()` to just `enable_monitor_loop_lag()`. +# - Miscellaneous changes to make it work with Ray. + +import asyncio +import asyncio.events +from typing import Callable, Optional + + +def enable_monitor_loop_lag( + callback: Callable[[float], None], + interval_s: float = 0.25, + loop: Optional[asyncio.AbstractEventLoop] = None, +) -> None: + """ + Start logging event loop lags to the callback. In ideal circumstances they should be + very close to zero. Lags may increase if event loop callbacks block for too long. + + Note: this works for all event loops, including uvloop. + + :param callback: Callback to call with the lag in seconds. + """ + if loop is None: + loop = asyncio.get_running_loop() + if loop is None: + raise ValueError("No provided loop, nor running loop found.") + + async def monitor(): + while loop.is_running(): + t0 = loop.time() + await asyncio.sleep(interval_s) + lag = loop.time() - t0 - interval_s # Should be close to zero. + callback(lag) + + loop.create_task(monitor(), name="async_utils.monitor_loop_lag") diff --git a/lib/python3.12/site-packages/ray/_private/authentication/__init__.py b/lib/python3.12/site-packages/ray/_private/authentication/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/lib/python3.12/site-packages/ray/_private/authentication/__pycache__/__init__.cpython-312.pyc b/lib/python3.12/site-packages/ray/_private/authentication/__pycache__/__init__.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..33e7af8a1c89a2a805c428d45f04ef1646d6228d Binary files /dev/null and b/lib/python3.12/site-packages/ray/_private/authentication/__pycache__/__init__.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/_private/authentication/__pycache__/authentication_constants.cpython-312.pyc b/lib/python3.12/site-packages/ray/_private/authentication/__pycache__/authentication_constants.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..155955bc839310637b5769cc7ae6770d96fe30f0 Binary files /dev/null and b/lib/python3.12/site-packages/ray/_private/authentication/__pycache__/authentication_constants.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/_private/authentication/__pycache__/authentication_token_generator.cpython-312.pyc b/lib/python3.12/site-packages/ray/_private/authentication/__pycache__/authentication_token_generator.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..694f1a90d2f4fcef3f249c76e0dfc00daeec4c33 Binary files /dev/null and b/lib/python3.12/site-packages/ray/_private/authentication/__pycache__/authentication_token_generator.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/_private/authentication/__pycache__/authentication_token_setup.cpython-312.pyc b/lib/python3.12/site-packages/ray/_private/authentication/__pycache__/authentication_token_setup.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..fd12dd0ab1f7367be38e7cb0b6f7369f58d7317d Binary files /dev/null and b/lib/python3.12/site-packages/ray/_private/authentication/__pycache__/authentication_token_setup.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/_private/authentication/__pycache__/authentication_utils.cpython-312.pyc b/lib/python3.12/site-packages/ray/_private/authentication/__pycache__/authentication_utils.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..128713b7df1bcd28dd9cd93d69b6a87365d695ba Binary files /dev/null and b/lib/python3.12/site-packages/ray/_private/authentication/__pycache__/authentication_utils.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/_private/authentication/__pycache__/grpc_authentication_client_interceptor.cpython-312.pyc b/lib/python3.12/site-packages/ray/_private/authentication/__pycache__/grpc_authentication_client_interceptor.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..a45727b2b34f8389490ef6e7b7ab494232de8209 Binary files /dev/null and b/lib/python3.12/site-packages/ray/_private/authentication/__pycache__/grpc_authentication_client_interceptor.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/_private/authentication/__pycache__/grpc_authentication_server_interceptor.cpython-312.pyc b/lib/python3.12/site-packages/ray/_private/authentication/__pycache__/grpc_authentication_server_interceptor.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..7a32a1eb5a520c3c6b90d930e4793b281ca6a0c6 Binary files /dev/null and b/lib/python3.12/site-packages/ray/_private/authentication/__pycache__/grpc_authentication_server_interceptor.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/_private/authentication/__pycache__/http_token_authentication.cpython-312.pyc b/lib/python3.12/site-packages/ray/_private/authentication/__pycache__/http_token_authentication.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..ddd6eaca9f42b9c10c68a87a229196df51b0ac04 Binary files /dev/null and b/lib/python3.12/site-packages/ray/_private/authentication/__pycache__/http_token_authentication.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/_private/authentication/authentication_constants.py b/lib/python3.12/site-packages/ray/_private/authentication/authentication_constants.py new file mode 100644 index 0000000000000000000000000000000000000000..2dce1872cc396ef933ec63d0da4fd2de596e8c2a --- /dev/null +++ b/lib/python3.12/site-packages/ray/_private/authentication/authentication_constants.py @@ -0,0 +1,14 @@ +# Authentication error messages +TOKEN_AUTH_ENABLED_BUT_NO_TOKEN_FOUND_ERROR_MESSAGE = ( + "Token authentication is enabled but no authentication token was found." +) + +TOKEN_INVALID_ERROR_MESSAGE = "Token authentication is enabled but the authentication token is invalid or incorrect." # noqa: E501 + +# HTTP header and cookie constants +AUTHORIZATION_HEADER_NAME = "authorization" +AUTHORIZATION_BEARER_PREFIX = "Bearer " +RAY_AUTHORIZATION_HEADER_NAME = "x-ray-authorization" + +AUTHENTICATION_TOKEN_COOKIE_NAME = "ray-authentication-token" +AUTHENTICATION_TOKEN_COOKIE_MAX_AGE = 30 * 24 * 60 * 60 # 30 days diff --git a/lib/python3.12/site-packages/ray/_private/authentication/authentication_token_generator.py b/lib/python3.12/site-packages/ray/_private/authentication/authentication_token_generator.py new file mode 100644 index 0000000000000000000000000000000000000000..34fcd1eab8a3391ba7f0ba06579a3b0d8ae94324 --- /dev/null +++ b/lib/python3.12/site-packages/ray/_private/authentication/authentication_token_generator.py @@ -0,0 +1,9 @@ +import secrets + + +def generate_new_authentication_token() -> str: + """Generate an authentication token for the cluster. + + 256 bits of entropy is considered sufficient to be durable to brute force attacks. + """ + return secrets.token_hex(32) diff --git a/lib/python3.12/site-packages/ray/_private/authentication/authentication_token_setup.py b/lib/python3.12/site-packages/ray/_private/authentication/authentication_token_setup.py new file mode 100644 index 0000000000000000000000000000000000000000..ad2081a86c44ad0c2a011c6df73da0bcb82760d6 --- /dev/null +++ b/lib/python3.12/site-packages/ray/_private/authentication/authentication_token_setup.py @@ -0,0 +1,103 @@ +"""Authentication token setup for Ray. + +This module provides functions to generate and save authentication tokens +for Ray's token-based authentication system. Token loading and caching is +handled by the C++ AuthenticationTokenLoader. +""" + +import logging +from pathlib import Path +from typing import Any, Dict, Optional + +from ray._private.authentication.authentication_constants import ( + TOKEN_AUTH_ENABLED_BUT_NO_TOKEN_FOUND_ERROR_MESSAGE, +) +from ray._private.authentication.authentication_token_generator import ( + generate_new_authentication_token, +) +from ray._raylet import ( + AuthenticationMode, + AuthenticationTokenLoader, + get_authentication_mode, +) +from ray.exceptions import AuthenticationError + +logger = logging.getLogger(__name__) + + +def generate_and_save_token() -> None: + """Generate a new random token and save it in the default token path. + + Returns: + The newly generated authentication token. + """ + # Generate a UUID-based token + token = generate_new_authentication_token() + + token_path = _get_default_token_path() + try: + # Create directory if it doesn't exist + token_path.parent.mkdir(parents=True, exist_ok=True) + + # Write token to file with explicit flush and fsync + with open(token_path, "w") as f: + f.write(token) + + logger.info(f"Generated new authentication token and saved to {token_path}") + except Exception: + raise + + +def _get_default_token_path() -> Path: + """Get the default token file path (~/.ray/auth_token). + + Returns: + Path object pointing to ~/.ray/auth_token + """ + return Path.home() / ".ray" / "auth_token" + + +def ensure_token_if_auth_enabled( + system_config: Optional[Dict[str, Any]] = None, create_token_if_missing: bool = True +) -> None: + """Check authentication settings and set up token resources if authentication is enabled. + + Ray calls this early during ray.init() to do the following for token-based authentication: + 1. Check whether you enabled token-based authentication. + 2. Make sure a token is available if authentication is enabled. + 3. Generate and save a default token for new local clusters if one doesn't already exist. + + Args: + system_config: Ray raises an error if you set AUTH_MODE in system_config instead of the environment. + create_token_if_missing: Generate a new token if one doesn't already exist. + + Raises: + RuntimeError: Ray raises this error if authentication is enabled but no token is found when connecting + to an existing cluster. + """ + + # Check if you enabled token authentication. + if get_authentication_mode() != AuthenticationMode.TOKEN: + if ( + system_config + and "AUTH_MODE" in system_config + and system_config["AUTH_MODE"] != "disabled" + ): + raise RuntimeError( + "Set authentication mode can only be set with the `RAY_AUTH_MODE` environment variable, not using the system_config." + ) + return + + token_loader = AuthenticationTokenLoader.instance() + + if not token_loader.has_token(): + if create_token_if_missing: + # Generate a new token. + generate_and_save_token() + + # Reload the cache so subsequent calls to token_loader read the new token. + token_loader.reset_cache() + else: + raise AuthenticationError( + TOKEN_AUTH_ENABLED_BUT_NO_TOKEN_FOUND_ERROR_MESSAGE + ) diff --git a/lib/python3.12/site-packages/ray/_private/authentication/authentication_utils.py b/lib/python3.12/site-packages/ray/_private/authentication/authentication_utils.py new file mode 100644 index 0000000000000000000000000000000000000000..5cd3ab522f7f765e04cf8710e8f621b8f3df89fd --- /dev/null +++ b/lib/python3.12/site-packages/ray/_private/authentication/authentication_utils.py @@ -0,0 +1,43 @@ +try: + from ray._raylet import ( + AuthenticationMode, + get_authentication_mode, + validate_authentication_token, + ) + + _RAYLET_AVAILABLE = True +except ImportError: + # ray._raylet not available during doc builds + _RAYLET_AVAILABLE = False + + +def is_token_auth_enabled() -> bool: + """Check if token authentication is enabled. + + Returns: + bool: True if AUTH_MODE is set to "token" or "k8s", False otherwise + """ + if not _RAYLET_AVAILABLE: + return False + + return get_authentication_mode() in { + AuthenticationMode.TOKEN, + AuthenticationMode.K8S, + } + + +def validate_request_token(auth_header: str) -> bool: + """Validate the Authorization header from an HTTP request. + + Args: + auth_header: The Authorization header value (e.g., "Bearer ") + + Returns: + bool: True if token is valid, False otherwise + """ + if not _RAYLET_AVAILABLE or not auth_header: + return False + + # validate_authentication_token expects full "Bearer " format + # and performs equality comparison via C++ layer + return validate_authentication_token(auth_header) diff --git a/lib/python3.12/site-packages/ray/_private/authentication/grpc_authentication_client_interceptor.py b/lib/python3.12/site-packages/ray/_private/authentication/grpc_authentication_client_interceptor.py new file mode 100644 index 0000000000000000000000000000000000000000..01fdd48d7732215cc9b27c76b2875ac8e0907b9e --- /dev/null +++ b/lib/python3.12/site-packages/ray/_private/authentication/grpc_authentication_client_interceptor.py @@ -0,0 +1,122 @@ +"""gRPC client interceptor for token-based authentication.""" + +import logging +from collections import namedtuple +from typing import Tuple + +import grpc +from grpc import aio as aiogrpc + +from ray._raylet import AuthenticationTokenLoader + +logger = logging.getLogger(__name__) + + +# Named tuple to hold client call details +_ClientCallDetails = namedtuple( + "_ClientCallDetails", + ("method", "timeout", "metadata", "credentials", "wait_for_ready", "compression"), +) + + +def _get_authentication_metadata_tuple() -> Tuple[Tuple[str, str], ...]: + """Get gRPC metadata tuple for authentication. Currently only supported for token authentication. + + Returns: + tuple: Empty tuple or ((AUTHORIZATION_HEADER_NAME, "Bearer "),) + """ + token_loader = AuthenticationTokenLoader.instance() + if not token_loader.has_token(): + return () + + headers = token_loader.get_token_for_http_header() + + # Convert HTTP header dict to gRPC metadata tuple + # gRPC expects: (("key", "value"), ...) + return tuple((k, v) for k, v in headers.items()) + + +class SyncAuthenticationMetadataClientInterceptor( + grpc.UnaryUnaryClientInterceptor, + grpc.UnaryStreamClientInterceptor, + grpc.StreamUnaryClientInterceptor, + grpc.StreamStreamClientInterceptor, +): + """Synchronous gRPC client interceptor that adds authentication metadata.""" + + def _intercept_call_details(self, client_call_details): + """Helper method to add authentication metadata to client call details.""" + metadata = list(client_call_details.metadata or []) + metadata.extend(_get_authentication_metadata_tuple()) + + return _ClientCallDetails( + method=client_call_details.method, + timeout=client_call_details.timeout, + metadata=metadata, + credentials=client_call_details.credentials, + wait_for_ready=getattr(client_call_details, "wait_for_ready", None), + compression=getattr(client_call_details, "compression", None), + ) + + def intercept_unary_unary(self, continuation, client_call_details, request): + new_details = self._intercept_call_details(client_call_details) + return continuation(new_details, request) + + def intercept_unary_stream(self, continuation, client_call_details, request): + new_details = self._intercept_call_details(client_call_details) + return continuation(new_details, request) + + def intercept_stream_unary( + self, continuation, client_call_details, request_iterator + ): + new_details = self._intercept_call_details(client_call_details) + return continuation(new_details, request_iterator) + + def intercept_stream_stream( + self, continuation, client_call_details, request_iterator + ): + new_details = self._intercept_call_details(client_call_details) + return continuation(new_details, request_iterator) + + +class AsyncAuthenticationMetadataClientInterceptor( + aiogrpc.UnaryUnaryClientInterceptor, + aiogrpc.UnaryStreamClientInterceptor, + aiogrpc.StreamUnaryClientInterceptor, + aiogrpc.StreamStreamClientInterceptor, +): + """Async gRPC client interceptor that adds authentication metadata.""" + + def _intercept_call_details(self, client_call_details): + """Helper method to add authentication metadata to client call details.""" + metadata = list(client_call_details.metadata or []) + metadata.extend(_get_authentication_metadata_tuple()) + + return _ClientCallDetails( + method=client_call_details.method, + timeout=client_call_details.timeout, + metadata=metadata, + credentials=client_call_details.credentials, + wait_for_ready=getattr(client_call_details, "wait_for_ready", None), + compression=getattr(client_call_details, "compression", None), + ) + + async def intercept_unary_unary(self, continuation, client_call_details, request): + new_details = self._intercept_call_details(client_call_details) + return await continuation(new_details, request) + + async def intercept_unary_stream(self, continuation, client_call_details, request): + new_details = self._intercept_call_details(client_call_details) + return await continuation(new_details, request) + + async def intercept_stream_unary( + self, continuation, client_call_details, request_iterator + ): + new_details = self._intercept_call_details(client_call_details) + return await continuation(new_details, request_iterator) + + async def intercept_stream_stream( + self, continuation, client_call_details, request_iterator + ): + new_details = self._intercept_call_details(client_call_details) + return await continuation(new_details, request_iterator) diff --git a/lib/python3.12/site-packages/ray/_private/authentication/grpc_authentication_server_interceptor.py b/lib/python3.12/site-packages/ray/_private/authentication/grpc_authentication_server_interceptor.py new file mode 100644 index 0000000000000000000000000000000000000000..d48262dc6d5f940d4791cd0983914ce52fc3c215 --- /dev/null +++ b/lib/python3.12/site-packages/ray/_private/authentication/grpc_authentication_server_interceptor.py @@ -0,0 +1,213 @@ +"""gRPC server interceptor for token-based authentication.""" + +import logging +from typing import Awaitable, Callable + +import grpc +from grpc import aio as aiogrpc + +from ray._private.authentication.authentication_constants import ( + AUTHORIZATION_HEADER_NAME, +) +from ray._private.authentication.authentication_utils import ( + is_token_auth_enabled, + validate_request_token, +) + +logger = logging.getLogger(__name__) + + +def _authenticate_request(metadata: tuple) -> bool: + """Authenticate incoming request. currently only supports token authentication. + + Args: + metadata: gRPC metadata tuple of (key, value) pairs + Returns: + True if authentication succeeds or is not required, False otherwise + """ + if not is_token_auth_enabled(): + return True + + # Extract authorization header from metadata + auth_header = None + for key, value in metadata: + if key.lower() == AUTHORIZATION_HEADER_NAME: + auth_header = value + break + + if not auth_header: + logger.warning("Authentication required but no authorization header provided") + return False + + # Validate the token format and value + # validate_request_token returns bool (True if valid, False otherwise) + return validate_request_token(auth_header) + + +class AsyncAuthenticationServerInterceptor(aiogrpc.ServerInterceptor): + """Async gRPC server interceptor that validates authentication tokens. + + This interceptor checks the "authorization" metadata header for a valid + Bearer token when token authentication is enabled via RAY_AUTH_MODE=token. + + If the token is missing or invalid, the request is rejected with UNAUTHENTICATED status. + """ + + async def intercept_service( + self, + continuation: Callable[ + [grpc.HandlerCallDetails], Awaitable[grpc.RpcMethodHandler] + ], + handler_call_details: grpc.HandlerCallDetails, + ) -> grpc.RpcMethodHandler: + """Intercept service calls to validate authentication. + + This method is called once per RPC to get the handler. We wrap the handler + to validate authentication before executing the actual RPC method. + """ + # Get the actual handler + handler = await continuation(handler_call_details) + + if handler is None: + return None + + # Wrap the RPC behavior with authentication check + def wrap_rpc_behavior(behavior): + """Wrap an RPC method to validate authentication first.""" + if behavior is None: + return None + + async def wrapped(request_or_iterator, context): + if not _authenticate_request(context.invocation_metadata()): + await context.abort( + grpc.StatusCode.UNAUTHENTICATED, + "Invalid or missing authentication token", + ) + return await behavior(request_or_iterator, context) + + return wrapped + + # Create a wrapper class that implements RpcMethodHandler interface + class AuthenticatedHandler: + """Wrapper handler that validates authentication.""" + + def __init__(self, original_handler, wrapper_func): + self._original = original_handler + self._wrap = wrapper_func + + @property + def request_streaming(self): + return self._original.request_streaming + + @property + def response_streaming(self): + return self._original.response_streaming + + @property + def request_deserializer(self): + return self._original.request_deserializer + + @property + def response_serializer(self): + return self._original.response_serializer + + @property + def unary_unary(self): + return self._wrap(self._original.unary_unary) + + @property + def unary_stream(self): + return self._wrap(self._original.unary_stream) + + @property + def stream_unary(self): + return self._wrap(self._original.stream_unary) + + @property + def stream_stream(self): + return self._wrap(self._original.stream_stream) + + return AuthenticatedHandler(handler, wrap_rpc_behavior) + + +class SyncAuthenticationServerInterceptor(grpc.ServerInterceptor): + """Synchronous gRPC server interceptor that validates authentication tokens. + + This interceptor checks the "authorization" metadata header for a valid + Bearer token when token authentication is enabled via RAY_AUTH_MODE=token. + If the token is missing or invalid, the request is rejected with UNAUTHENTICATED status. + """ + + def intercept_service( + self, + continuation: Callable[[grpc.HandlerCallDetails], grpc.RpcMethodHandler], + handler_call_details: grpc.HandlerCallDetails, + ) -> grpc.RpcMethodHandler: + """Intercept service calls to validate authentication. + + This method is called once per RPC to get the handler. We wrap the handler + to validate authentication before executing the actual RPC method. + """ + # Get the actual handler + handler = continuation(handler_call_details) + + if handler is None: + return None + + # Wrap the RPC behavior with authentication check + def wrap_rpc_behavior(behavior): + """Wrap an RPC method to validate authentication first.""" + if behavior is None: + return None + + def wrapped(request_or_iterator, context): + if not _authenticate_request(context.invocation_metadata()): + context.abort( + grpc.StatusCode.UNAUTHENTICATED, + "Invalid or missing authentication token", + ) + return behavior(request_or_iterator, context) + + return wrapped + + # Create a wrapper class that implements RpcMethodHandler interface + class AuthenticatedHandler: + """Wrapper handler that validates authentication.""" + + def __init__(self, original_handler, wrapper_func): + self._original = original_handler + self._wrap = wrapper_func + + @property + def request_streaming(self): + return self._original.request_streaming + + @property + def response_streaming(self): + return self._original.response_streaming + + @property + def request_deserializer(self): + return self._original.request_deserializer + + @property + def response_serializer(self): + return self._original.response_serializer + + @property + def unary_unary(self): + return self._wrap(self._original.unary_unary) + + @property + def unary_stream(self): + return self._wrap(self._original.unary_stream) + + @property + def stream_unary(self): + return self._wrap(self._original.stream_unary) + + @property + def stream_stream(self): + return self._wrap(self._original.stream_stream) + + return AuthenticatedHandler(handler, wrap_rpc_behavior) diff --git a/lib/python3.12/site-packages/ray/_private/authentication/http_token_authentication.py b/lib/python3.12/site-packages/ray/_private/authentication/http_token_authentication.py new file mode 100644 index 0000000000000000000000000000000000000000..3d5d89add4411c0b093f4d74cd2a44ff98794b8e --- /dev/null +++ b/lib/python3.12/site-packages/ray/_private/authentication/http_token_authentication.py @@ -0,0 +1,129 @@ +import logging +from types import ModuleType +from typing import Dict, List, Optional + +from ray._private.authentication import ( + authentication_constants, + authentication_utils as auth_utils, +) + +logger = logging.getLogger(__name__) + + +def get_token_auth_middleware( + aiohttp_module: ModuleType, + whitelisted_exact_paths: Optional[List[str]] = None, + whitelisted_path_prefixes: Optional[List[str]] = None, +): + """Internal helper to create token auth middleware with provided modules. + + Args: + aiohttp_module: The aiohttp module to use + whitelisted_exact_paths: List of exact paths that don't require authentication + whitelisted_path_prefixes: List of path prefixes that don't require authentication + Returns: + An aiohttp middleware function + """ + + @aiohttp_module.web.middleware + async def token_auth_middleware(request, handler): + """Middleware to validate bearer tokens when token authentication is enabled. + + In minimal Ray installations (without ray._raylet), this middleware is a no-op + and passes all requests through without authentication. + """ + # No-op if token auth is not enabled or raylet is not available + if not auth_utils.is_token_auth_enabled(): + return await handler(request) + + # skip authentication for whitelisted paths + if (whitelisted_exact_paths and request.path in whitelisted_exact_paths) or ( + whitelisted_path_prefixes + and request.path.startswith(tuple(whitelisted_path_prefixes)) + ): + return await handler(request) + + # Try to get authentication token from multiple sources (in priority order): + # 1. Standard "Authorization" header (for API clients, SDKs) + # 2. Fallback "X-Ray-Authorization" header (for proxies and KubeRay) + # 3. Cookie (for web dashboard sessions) + + auth_header = request.headers.get( + authentication_constants.AUTHORIZATION_HEADER_NAME, "" + ) + + if not auth_header: + auth_header = request.headers.get( + authentication_constants.RAY_AUTHORIZATION_HEADER_NAME, "" + ) + + if not auth_header: + token = request.cookies.get( + authentication_constants.AUTHENTICATION_TOKEN_COOKIE_NAME + ) + if token: + # Format as Bearer token for validation + auth_header = ( + authentication_constants.AUTHORIZATION_BEARER_PREFIX + token + ) + + if not auth_header: + return aiohttp_module.web.Response( + status=401, text="Unauthorized: Missing authentication token" + ) + + if not auth_utils.validate_request_token(auth_header): + return aiohttp_module.web.Response( + status=403, text="Forbidden: Invalid authentication token" + ) + + return await handler(request) + + return token_auth_middleware + + +def get_auth_headers_if_auth_enabled(user_headers: Dict[str, str]) -> Dict[str, str]: + + if not auth_utils.is_token_auth_enabled(): + return {} + + from ray._raylet import AuthenticationTokenLoader + + # Check if user provided their own Authorization header (case-insensitive) + has_user_auth = any( + key.lower() == authentication_constants.AUTHORIZATION_HEADER_NAME + for key in user_headers.keys() + ) + if has_user_auth: + # User has provided their own auth header, don't override + return {} + + token_loader = AuthenticationTokenLoader.instance() + auth_headers = token_loader.get_token_for_http_header() + + if not auth_headers: + # Token auth enabled but no token found + logger.warning( + "Token authentication is enabled but no token was found. " + "Requests to authenticated clusters will fail." + ) + + return auth_headers + + +def format_authentication_http_error(status: int, body: str) -> Optional[str]: + """Return a user-friendly authentication error message, if applicable.""" + + if status == 401: + return "Authentication required: {body}\n\n{details}".format( + body=body, + details=authentication_constants.TOKEN_AUTH_ENABLED_BUT_NO_TOKEN_FOUND_ERROR_MESSAGE, + ) + + if status == 403: + return "Authentication failed: {body}\n\n{details}".format( + body=body, + details=authentication_constants.TOKEN_INVALID_ERROR_MESSAGE, + ) + + return None diff --git a/lib/python3.12/site-packages/ray/_private/authentication_test_utils.py b/lib/python3.12/site-packages/ray/_private/authentication_test_utils.py new file mode 100644 index 0000000000000000000000000000000000000000..ab0119e52426065059f7ccd08c61ab9308b11bc1 --- /dev/null +++ b/lib/python3.12/site-packages/ray/_private/authentication_test_utils.py @@ -0,0 +1,150 @@ +import os +import shutil +import tempfile +from contextlib import contextmanager +from dataclasses import dataclass +from pathlib import Path +from typing import Dict, Optional + +from ray._raylet import AuthenticationTokenLoader, Config + +_AUTH_ENV_VARS = ("RAY_AUTH_MODE", "RAY_AUTH_TOKEN", "RAY_AUTH_TOKEN_PATH") +_DEFAULT_AUTH_TOKEN_RELATIVE_PATH = Path(".ray") / "auth_token" + + +def reset_auth_token_state() -> None: + """Reset authentication token and AUTH_MODE ray config.""" + + AuthenticationTokenLoader.instance().reset_cache() + Config.initialize("") + + +def set_auth_mode(mode: str) -> None: + """Set the authentication mode environment variable.""" + + os.environ["RAY_AUTH_MODE"] = mode + + +def set_env_auth_token(token: str) -> None: + """Configure the authentication token via environment variable.""" + + os.environ["RAY_AUTH_TOKEN"] = token + os.environ.pop("RAY_AUTH_TOKEN_PATH", None) + + +def set_auth_token_path(token: str, path: Path) -> None: + """Write the authentication token to a specific path and point the loader to it.""" + + token_path = Path(path) + token_path.parent.mkdir(parents=True, exist_ok=True) + token_path.write_text(token) + os.environ["RAY_AUTH_TOKEN_PATH"] = str(token_path) + os.environ.pop("RAY_AUTH_TOKEN", None) + + +def set_default_auth_token(token: str) -> Path: + """Write the authentication token to the default ~/.ray/auth_token location.""" + + default_path = Path.home() / _DEFAULT_AUTH_TOKEN_RELATIVE_PATH + default_path.parent.mkdir(parents=True, exist_ok=True) + default_path.write_text(token) + return default_path + + +def clear_auth_token_sources(remove_default: bool = False) -> None: + """Clear authentication-related environment variables and optional default token file.""" + + for var in ("RAY_AUTH_TOKEN", "RAY_AUTH_TOKEN_PATH"): + os.environ.pop(var, None) + + if remove_default: + default_path = Path.home() / _DEFAULT_AUTH_TOKEN_RELATIVE_PATH + default_path.unlink(missing_ok=True) + + +@dataclass +class AuthenticationEnvSnapshot: + original_env: Dict[str, Optional[str]] + original_home: Optional[str] + home_was_set: bool + temp_home: Optional[Path] + default_token_path: Path + default_token_exists: bool + default_token_contents: Optional[str] + + @classmethod + def capture(cls) -> "AuthenticationEnvSnapshot": + """Capture current authentication-related environment state.""" + + original_env = {var: os.environ.get(var) for var in _AUTH_ENV_VARS} + home_was_set = "HOME" in os.environ + original_home = os.environ.get("HOME") + temp_home: Optional[Path] = None + + if not home_was_set: + # in CI $HOME may not be set which can cause issues with tests related to default auth token file. + test_tmpdir = os.environ.get("TEST_TMPDIR") + base_dir = Path(test_tmpdir) if test_tmpdir else Path(tempfile.gettempdir()) + temp_home = base_dir / "ray_test_home" + temp_home.mkdir(parents=True, exist_ok=True) + os.environ["HOME"] = str(temp_home) + + default_token_path = Path.home() / _DEFAULT_AUTH_TOKEN_RELATIVE_PATH + default_token_exists = default_token_path.exists() + default_token_contents = ( + default_token_path.read_text() if default_token_exists else None + ) + + return cls( + original_env=original_env, + original_home=original_home, + home_was_set=home_was_set, + temp_home=temp_home, + default_token_path=default_token_path, + default_token_exists=default_token_exists, + default_token_contents=default_token_contents, + ) + + def clear_default_token(self) -> None: + """Remove the default token file for the current HOME.""" + + self.default_token_path.unlink(missing_ok=True) + + def restore(self) -> None: + """Restore the captured environment, HOME, and default token file state.""" + + for var, value in self.original_env.items(): + if value is None: + os.environ.pop(var, None) + else: + os.environ[var] = value + + if self.home_was_set: + if self.original_home is None: + os.environ.pop("HOME", None) + else: + os.environ["HOME"] = self.original_home + + if self.default_token_exists: + self.default_token_path.parent.mkdir(parents=True, exist_ok=True) + self.default_token_path.write_text(self.default_token_contents or "") + else: + self.default_token_path.unlink(missing_ok=True) + + if not self.home_was_set: + current_home = os.environ.get("HOME") + if self.temp_home is not None and current_home == str(self.temp_home): + os.environ.pop("HOME", None) + if self.temp_home is not None and self.temp_home.exists(): + shutil.rmtree(self.temp_home, ignore_errors=True) + + +@contextmanager +def authentication_env_guard(): + """Context manager that restores authentication environment state on exit.""" + + snapshot = AuthenticationEnvSnapshot.capture() + try: + yield snapshot + finally: + snapshot.restore() diff --git a/lib/python3.12/site-packages/ray/_private/auto_init_hook.py b/lib/python3.12/site-packages/ray/_private/auto_init_hook.py new file mode 100644 index 0000000000000000000000000000000000000000..2013051a3294936a6e15526c8d9efa0ffe25ae12 --- /dev/null +++ b/lib/python3.12/site-packages/ray/_private/auto_init_hook.py @@ -0,0 +1,32 @@ +import os +import threading +from functools import wraps + +import ray + +auto_init_lock = threading.Lock() +enable_auto_connect = os.environ.get("RAY_ENABLE_AUTO_CONNECT", "") != "0" + + +def auto_init_ray(): + if enable_auto_connect and not ray.is_initialized(): + with auto_init_lock: + if not ray.is_initialized(): + ray.init() + + +def wrap_auto_init(fn): + @wraps(fn) + def auto_init_wrapper(*args, **kwargs): + auto_init_ray() + return fn(*args, **kwargs) + + return auto_init_wrapper + + +def wrap_auto_init_for_all_apis(api_names): + """Wrap public APIs with automatic ray.init.""" + for api_name in api_names: + api = getattr(ray, api_name, None) + assert api is not None, api_name + setattr(ray, api_name, wrap_auto_init(api)) diff --git a/lib/python3.12/site-packages/ray/_private/collections_utils.py b/lib/python3.12/site-packages/ray/_private/collections_utils.py new file mode 100644 index 0000000000000000000000000000000000000000..f7512cfb4b987d022e6bea4218f9097ff27ae69a --- /dev/null +++ b/lib/python3.12/site-packages/ray/_private/collections_utils.py @@ -0,0 +1,10 @@ +from typing import Any, List + + +def split(items: List[Any], chunk_size: int): + """Splits provided list into chunks of given size""" + + assert chunk_size > 0, "Chunk size has to be > 0" + + for i in range(0, len(items), chunk_size): + yield items[i : i + chunk_size] diff --git a/lib/python3.12/site-packages/ray/_private/conftest_utils.py b/lib/python3.12/site-packages/ray/_private/conftest_utils.py new file mode 100644 index 0000000000000000000000000000000000000000..2dc02a2d62cffd433dcf35c969c85de295be45bf --- /dev/null +++ b/lib/python3.12/site-packages/ray/_private/conftest_utils.py @@ -0,0 +1,15 @@ +import pytest + +import ray._private.ray_constants as ray_constants + + +@pytest.fixture +def set_override_dashboard_url(monkeypatch, request): + override_url = getattr(request, "param", "https://external_dashboard_url") + with monkeypatch.context() as m: + if override_url: + m.setenv( + ray_constants.RAY_OVERRIDE_DASHBOARD_URL, + override_url, + ) + yield diff --git a/lib/python3.12/site-packages/ray/_private/custom_types.py b/lib/python3.12/site-packages/ray/_private/custom_types.py new file mode 100644 index 0000000000000000000000000000000000000000..fc5e3a5622fbd1a330b043c41cd0a5744a4754a0 --- /dev/null +++ b/lib/python3.12/site-packages/ray/_private/custom_types.py @@ -0,0 +1,167 @@ +from enum import Enum +from typing import Literal + +from ray.core.generated.common_pb2 import ( + ErrorType, + Language, + TaskStatus, + TaskType, + TensorTransport, + WorkerExitType, + WorkerType, +) +from ray.core.generated.gcs_pb2 import ( + ActorTableData, + GcsNodeInfo, + PlacementGroupTableData, +) + +ACTOR_STATUS = [ + "DEPENDENCIES_UNREADY", + "PENDING_CREATION", + "ALIVE", + "RESTARTING", + "DEAD", +] +TypeActorStatus = Literal[tuple(ACTOR_STATUS)] +PLACEMENT_GROUP_STATUS = [ + "PENDING", + "PREPARED", + "CREATED", + "REMOVED", + "RESCHEDULING", +] +TypePlacementGroupStatus = Literal[tuple(PLACEMENT_GROUP_STATUS)] +TASK_STATUS = [ + "NIL", + "PENDING_ARGS_AVAIL", + "PENDING_NODE_ASSIGNMENT", + "PENDING_OBJ_STORE_MEM_AVAIL", + "PENDING_ARGS_FETCH", + "SUBMITTED_TO_WORKER", + "PENDING_ACTOR_TASK_ARGS_FETCH", + "PENDING_ACTOR_TASK_ORDERING_OR_CONCURRENCY", + "RUNNING", + "RUNNING_IN_RAY_GET", + "RUNNING_IN_RAY_WAIT", + "FINISHED", + "FAILED", + "GETTING_AND_PINNING_ARGS", +] +TypeTaskStatus = Literal[tuple(TASK_STATUS)] +NODE_STATUS = ["ALIVE", "DEAD"] +TypeNodeStatus = Literal[tuple(NODE_STATUS)] +WORKER_TYPE = [ + "WORKER", + "DRIVER", + "SPILL_WORKER", + "RESTORE_WORKER", +] +TypeWorkerType = Literal[tuple(WORKER_TYPE)] +WORKER_EXIT_TYPE = [ + "SYSTEM_ERROR", + "INTENDED_SYSTEM_EXIT", + "USER_ERROR", + "INTENDED_USER_EXIT", + "NODE_OUT_OF_MEMORY", +] +TypeWorkerExitType = Literal[tuple(WORKER_EXIT_TYPE)] +TASK_TYPE = [ + "NORMAL_TASK", + "ACTOR_CREATION_TASK", + "ACTOR_TASK", + "DRIVER_TASK", +] +TypeTaskType = Literal[tuple(TASK_TYPE)] +# TODO(kevin85421): `class ReferenceType(Enum)` is defined in +# `dashboard/memory_utils.py` to avoid complex dependencies. I redefined +# it here. Eventually, we should remove the one in `dashboard/memory_utils.py` +# and define it under `ray/_private`. +REFERENCE_TYPE = [ + "ACTOR_HANDLE", + "PINNED_IN_MEMORY", + "LOCAL_REFERENCE", + "USED_BY_PENDING_TASK", + "CAPTURED_IN_OBJECT", + "UNKNOWN_STATUS", +] +TypeReferenceType = Literal[tuple(REFERENCE_TYPE)] +# The ErrorType enum is used in the export API so it is public +# and any modifications must be backward compatible. +ERROR_TYPE = [ + "WORKER_DIED", + "ACTOR_DIED", + "OBJECT_UNRECONSTRUCTABLE", + "TASK_EXECUTION_EXCEPTION", + "OBJECT_IN_PLASMA", + "TASK_CANCELLED", + "ACTOR_CREATION_FAILED", + "RUNTIME_ENV_SETUP_FAILED", + "OBJECT_LOST", + "OWNER_DIED", + "OBJECT_DELETED", + "DEPENDENCY_RESOLUTION_FAILED", + "OBJECT_UNRECONSTRUCTABLE_MAX_ATTEMPTS_EXCEEDED", + "OBJECT_UNRECONSTRUCTABLE_LINEAGE_EVICTED", + "OBJECT_FETCH_TIMED_OUT", + "LOCAL_RAYLET_DIED", + "TASK_PLACEMENT_GROUP_REMOVED", + "ACTOR_PLACEMENT_GROUP_REMOVED", + "TASK_UNSCHEDULABLE_ERROR", + "ACTOR_UNSCHEDULABLE_ERROR", + "OUT_OF_DISK_ERROR", + "OBJECT_FREED", + "OUT_OF_MEMORY", + "NODE_DIED", + "END_OF_STREAMING_GENERATOR", + "ACTOR_UNAVAILABLE", + "GENERATOR_TASK_FAILED_FOR_OBJECT_RECONSTRUCTION", +] +# The Language enum is used in the export API so it is public +# and any modifications must be backward compatible. +LANGUAGE = ["PYTHON", "JAVA", "CPP"] + +# See `common.proto` for more details. +class TensorTransportEnum(Enum): + OBJECT_STORE = TensorTransport.Value("OBJECT_STORE") + NCCL = TensorTransport.Value("NCCL") + GLOO = TensorTransport.Value("GLOO") + NIXL = TensorTransport.Value("NIXL") + + @classmethod + def from_str(cls, name: str) -> "TensorTransportEnum": + name = name.upper() + if name not in cls.__members__: + raise ValueError( + f"Invalid tensor transport {name}, must be one of {list(cls.__members__.keys())}." + ) + return cls[name] + + +def validate_protobuf_enum(grpc_enum, custom_enum): + """Validate the literal contains the correct enum values from protobuf""" + enum_vals = set(grpc_enum.DESCRIPTOR.values_by_name.keys()) + # Sometimes, the grpc enum is mocked, and it + # doesn't include any values in that case. + if len(enum_vals) > 0: + assert enum_vals == set( + custom_enum + ), """Literals in `custom_types.py` and `.proto` files are out of sync. \ +Consider building //:install_py_proto with Bazel or updating `custom_types.py`.""" + + +# Do the enum validation here. +# It is necessary to avoid regression. Alternatively, we can auto generate this +# directly by protobuf. +validate_protobuf_enum(ActorTableData.ActorState, ACTOR_STATUS) +validate_protobuf_enum( + PlacementGroupTableData.PlacementGroupState, PLACEMENT_GROUP_STATUS +) +validate_protobuf_enum(TaskStatus, TASK_STATUS) +validate_protobuf_enum(GcsNodeInfo.GcsNodeState, NODE_STATUS) +validate_protobuf_enum(WorkerType, WORKER_TYPE) +validate_protobuf_enum(WorkerExitType, WORKER_EXIT_TYPE) +validate_protobuf_enum(TaskType, TASK_TYPE) +validate_protobuf_enum(ErrorType, ERROR_TYPE) +validate_protobuf_enum(Language, LANGUAGE) +validate_protobuf_enum(TensorTransport, list(TensorTransportEnum.__members__.keys())) diff --git a/lib/python3.12/site-packages/ray/_private/dict.py b/lib/python3.12/site-packages/ray/_private/dict.py new file mode 100644 index 0000000000000000000000000000000000000000..3d102b32961f034f4e7af469898cb2fb66c5ec93 --- /dev/null +++ b/lib/python3.12/site-packages/ray/_private/dict.py @@ -0,0 +1,247 @@ +import copy +from collections import deque +from collections.abc import Mapping, Sequence +from typing import Dict, List, Optional, TypeVar, Union + +from ray.util.annotations import Deprecated + +T = TypeVar("T") + + +@Deprecated +def merge_dicts(d1: dict, d2: dict) -> dict: + """ + Args: + d1 (dict): Dict 1. + d2 (dict): Dict 2. + + Returns: + dict: A new dict that is d1 and d2 deep merged. + """ + merged = copy.deepcopy(d1) + deep_update(merged, d2, True, []) + return merged + + +@Deprecated +def deep_update( + original: dict, + new_dict: dict, + new_keys_allowed: bool = False, + allow_new_subkey_list: Optional[List[str]] = None, + override_all_if_type_changes: Optional[List[str]] = None, + override_all_key_list: Optional[List[str]] = None, +) -> dict: + """Updates original dict with values from new_dict recursively. + + If new key is introduced in new_dict, then if new_keys_allowed is not + True, an error will be thrown. Further, for sub-dicts, if the key is + in the allow_new_subkey_list, then new subkeys can be introduced. + + Args: + original: Dictionary with default values. + new_dict: Dictionary with values to be updated + new_keys_allowed: Whether new keys are allowed. + allow_new_subkey_list: List of keys that + correspond to dict values where new subkeys can be introduced. + This is only at the top level. + override_all_if_type_changes: List of top level + keys with value=dict, for which we always simply override the + entire value (dict), iff the "type" key in that value dict changes. + override_all_key_list: List of top level keys + for which we override the entire value if the key is in the new_dict. + """ + allow_new_subkey_list = allow_new_subkey_list or [] + override_all_if_type_changes = override_all_if_type_changes or [] + override_all_key_list = override_all_key_list or [] + + for k, value in new_dict.items(): + if k not in original and not new_keys_allowed: + raise Exception("Unknown config parameter `{}` ".format(k)) + + # Both orginal value and new one are dicts. + if ( + isinstance(original.get(k), dict) + and isinstance(value, dict) + and k not in override_all_key_list + ): + # Check old type vs old one. If different, override entire value. + if ( + k in override_all_if_type_changes + and "type" in value + and "type" in original[k] + and value["type"] != original[k]["type"] + ): + original[k] = value + # Allowed key -> ok to add new subkeys. + elif k in allow_new_subkey_list: + deep_update( + original[k], + value, + True, + override_all_key_list=override_all_key_list, + ) + # Non-allowed key. + else: + deep_update( + original[k], + value, + new_keys_allowed, + override_all_key_list=override_all_key_list, + ) + # Original value not a dict OR new value not a dict: + # Override entire value. + else: + original[k] = value + return original + + +@Deprecated +def flatten_dict( + dt: Dict, + delimiter: str = "/", + prevent_delimiter: bool = False, + flatten_list: bool = False, +): + """Flatten dict. + + Output and input are of the same dict type. + Input dict remains the same after the operation. + """ + + def _raise_delimiter_exception(): + raise ValueError( + f"Found delimiter `{delimiter}` in key when trying to flatten " + f"array. Please avoid using the delimiter in your specification." + ) + + dt = copy.copy(dt) + if prevent_delimiter and any(delimiter in key for key in dt): + # Raise if delimiter is any of the keys + _raise_delimiter_exception() + + while_check = (dict, list) if flatten_list else dict + + while any(isinstance(v, while_check) for v in dt.values()): + remove = [] + add = {} + for key, value in dt.items(): + if isinstance(value, dict): + for subkey, v in value.items(): + if prevent_delimiter and delimiter in subkey: + # Raise if delimiter is in any of the subkeys + _raise_delimiter_exception() + + add[delimiter.join([key, str(subkey)])] = v + remove.append(key) + elif flatten_list and isinstance(value, list): + for i, v in enumerate(value): + if prevent_delimiter and delimiter in subkey: + # Raise if delimiter is in any of the subkeys + _raise_delimiter_exception() + + add[delimiter.join([key, str(i)])] = v + remove.append(key) + + dt.update(add) + for k in remove: + del dt[k] + return dt + + +@Deprecated +def unflatten_dict(dt: Dict[str, T], delimiter: str = "/") -> Dict[str, T]: + """Unflatten dict. Does not support unflattening lists.""" + dict_type = type(dt) + out = dict_type() + for key, val in dt.items(): + path = key.split(delimiter) + item = out + for k in path[:-1]: + item = item.setdefault(k, dict_type()) + if not isinstance(item, dict_type): + raise TypeError( + f"Cannot unflatten dict due the key '{key}' " + f"having a parent key '{k}', which value is not " + f"of type {dict_type} (got {type(item)}). " + "Change the key names to resolve the conflict." + ) + item[path[-1]] = val + return out + + +@Deprecated +def unflatten_list_dict(dt: Dict[str, T], delimiter: str = "/") -> Dict[str, T]: + """Unflatten nested dict and list. + + This function now has some limitations: + (1) The keys of dt must be str. + (2) If unflattened dt (the result) contains list, the index order must be + ascending when accessing dt. Otherwise, this function will throw + AssertionError. + (3) The unflattened dt (the result) shouldn't contain dict with number + keys. + + Be careful to use this function. If you want to improve this function, + please also improve the unit test. See #14487 for more details. + + Args: + dt: Flattened dictionary that is originally nested by multiple + list and dict. + delimiter: Delimiter of keys. + + Example: + >>> dt = {"aaa/0/bb": 12, "aaa/1/cc": 56, "aaa/1/dd": 92} + >>> unflatten_list_dict(dt) + {'aaa': [{'bb': 12}, {'cc': 56, 'dd': 92}]} + """ + out_type = list if list(dt)[0].split(delimiter, 1)[0].isdigit() else type(dt) + out = out_type() + for key, val in dt.items(): + path = key.split(delimiter) + + item = out + for i, k in enumerate(path[:-1]): + next_type = list if path[i + 1].isdigit() else dict + if isinstance(item, dict): + item = item.setdefault(k, next_type()) + elif isinstance(item, list): + if int(k) >= len(item): + item.append(next_type()) + assert int(k) == len(item) - 1 + item = item[int(k)] + + if isinstance(item, dict): + item[path[-1]] = val + elif isinstance(item, list): + item.append(val) + assert int(path[-1]) == len(item) - 1 + return out + + +@Deprecated +def unflattened_lookup( + flat_key: str, lookup: Union[Mapping, Sequence], delimiter: str = "/", **kwargs +) -> Union[Mapping, Sequence]: + """ + Unflatten `flat_key` and iteratively look up in `lookup`. E.g. + `flat_key="a/0/b"` will try to return `lookup["a"][0]["b"]`. + """ + if flat_key in lookup: + return lookup[flat_key] + keys = deque(flat_key.split(delimiter)) + base = lookup + while keys: + key = keys.popleft() + try: + if isinstance(base, Mapping): + base = base[key] + elif isinstance(base, Sequence): + base = base[int(key)] + else: + raise KeyError() + except KeyError as e: + if "default" in kwargs: + return kwargs["default"] + raise e + return base diff --git a/lib/python3.12/site-packages/ray/_private/event/__init__.py b/lib/python3.12/site-packages/ray/_private/event/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/lib/python3.12/site-packages/ray/_private/event/__pycache__/__init__.cpython-312.pyc b/lib/python3.12/site-packages/ray/_private/event/__pycache__/__init__.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..a895c79cd07bdac275b728680bb604212eb8780b Binary files /dev/null and b/lib/python3.12/site-packages/ray/_private/event/__pycache__/__init__.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/_private/event/__pycache__/event_logger.cpython-312.pyc b/lib/python3.12/site-packages/ray/_private/event/__pycache__/event_logger.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..63f650523913630c3ff5dc5e03268d713e209163 Binary files /dev/null and b/lib/python3.12/site-packages/ray/_private/event/__pycache__/event_logger.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/_private/event/__pycache__/export_event_logger.cpython-312.pyc b/lib/python3.12/site-packages/ray/_private/event/__pycache__/export_event_logger.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..709deb567b5c14ccdcf3ec910defe1e26e2455b1 Binary files /dev/null and b/lib/python3.12/site-packages/ray/_private/event/__pycache__/export_event_logger.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/_private/event/event_logger.py b/lib/python3.12/site-packages/ray/_private/event/event_logger.py new file mode 100644 index 0000000000000000000000000000000000000000..d8e54505bdc79656b151dae60edc387ce728d021 --- /dev/null +++ b/lib/python3.12/site-packages/ray/_private/event/event_logger.py @@ -0,0 +1,195 @@ +import json +import logging +import os +import pathlib +import random +import socket +import string +import threading +from datetime import datetime +from typing import Dict, Optional + +from google.protobuf.json_format import Parse + +from ray._private.protobuf_compat import message_to_dict +from ray.core.generated.event_pb2 import Event + +global_logger = logging.getLogger(__name__) + + +def get_event_id(): + return "".join([random.choice(string.hexdigits) for _ in range(36)]) + + +class EventLoggerAdapter: + def __init__(self, source: Event.SourceType, logger: logging.Logger): + """Adapter for the Python logger that's used to emit events. + + When events are emitted, they are aggregated and available via + state API and dashboard. + + This class is thread-safe. + """ + self.logger = logger + # Aligned with `event.proto`'s `message Event`` + self.source = source + self.source_hostname = socket.gethostname() + self.source_pid = os.getpid() + + # The below fields must be protected by this lock. + self.lock = threading.Lock() + # {str -> str} typed dict + self.global_context = {} + + def set_global_context(self, global_context: Dict[str, str] = None): + """Set the global metadata. + + This method overwrites the global metadata if it is called more than once. + """ + with self.lock: + self.global_context = {} if not global_context else global_context + + def trace(self, message: str, **kwargs): + self._emit(Event.Severity.TRACE, message, **kwargs) + + def debug(self, message: str, **kwargs): + self._emit(Event.Severity.DEBUG, message, **kwargs) + + def info(self, message: str, **kwargs): + self._emit(Event.Severity.INFO, message, **kwargs) + + def warning(self, message: str, **kwargs): + self._emit(Event.Severity.WARNING, message, **kwargs) + + def error(self, message: str, **kwargs): + self._emit(Event.Severity.ERROR, message, **kwargs) + + def fatal(self, message: str, **kwargs): + self._emit(Event.Severity.FATAL, message, **kwargs) + + def _emit(self, severity: Event.Severity, message: str, **kwargs): + # NOTE: Python logger is thread-safe, + # so we don't need to protect it using locks. + event = Event() + event.event_id = get_event_id() + event.timestamp = int(datetime.now().timestamp()) + event.message = message + event.severity = severity + # TODO(sang): Support event type & schema. + event.label = "" + event.source_type = self.source + event.source_hostname = self.source_hostname + event.source_pid = self.source_pid + custom_fields = event.custom_fields + with self.lock: + for k, v in self.global_context.items(): + if v is not None and k is not None: + custom_fields[k] = v + for k, v in kwargs.items(): + if v is not None and k is not None: + custom_fields[k] = v + + self.logger.info( + json.dumps( + message_to_dict( + event, + always_print_fields_with_no_presence=True, + preserving_proto_field_name=True, + use_integers_for_enums=False, + ) + ) + ) + + # Force flush so that we won't lose events + self.logger.handlers[0].flush() + + +def _build_event_file_logger(source: Event.SourceType, sink_dir: str): + logger = logging.getLogger("_ray_event_logger") + logger.setLevel(logging.INFO) + dir_path = pathlib.Path(sink_dir) / "events" + filepath = dir_path / f"event_{source}.log" + dir_path.mkdir(exist_ok=True) + filepath.touch(exist_ok=True) + # Configure the logger. + handler = logging.FileHandler(filepath) + formatter = logging.Formatter("%(message)s") + handler.setFormatter(formatter) + logger.addHandler(handler) + logger.propagate = False + return logger + + +# This lock must be used when accessing or updating global event logger dict. +_event_logger_lock = threading.Lock() +_event_logger = {} + + +def get_event_logger(source: Event.SourceType, sink_dir: str): + """Get the event logger of the current process. + + There's only 1 event logger per (process, source). + + TODO(sang): Support more impl than file-based logging. + Currently, the interface also ties to the + file-based logging impl. + + Args: + source: The source of the event. + sink_dir: The directory to sink event logs. + """ + with _event_logger_lock: + global _event_logger + source_name = Event.SourceType.Name(source) + if source_name not in _event_logger: + logger = _build_event_file_logger(source_name, sink_dir) + _event_logger[source_name] = EventLoggerAdapter(source, logger) + + return _event_logger[source_name] + + +def parse_event(event_str: str) -> Optional[Event]: + """Parse an event from a string. + + Args: + event_str: The string to parse. Expect to be a JSON serialized + Event protobuf. + + Returns: + The parsed event if parsable, else None + """ + try: + return Parse(event_str, Event()) + except Exception: + global_logger.exception(f"Failed to parse event: {event_str}") + return None + + +def filter_event_by_level(event: Event, filter_event_level: str) -> bool: + """Filter an event based on event level. + + Args: + event: The event to filter. + filter_event_level: The event level string to filter by. Any events + that are lower than this level will be filtered. + + Returns: + True if the event should be filtered, else False. + """ + + event_levels = { + Event.Severity.TRACE: 0, + Event.Severity.DEBUG: 1, + Event.Severity.INFO: 2, + Event.Severity.WARNING: 3, + Event.Severity.ERROR: 4, + Event.Severity.FATAL: 5, + } + + filter_event_level = filter_event_level.upper() + filter_event_level = Event.Severity.Value(filter_event_level) + + if event_levels[event.severity] < event_levels[filter_event_level]: + return True + + return False diff --git a/lib/python3.12/site-packages/ray/_private/event/export_event_logger.py b/lib/python3.12/site-packages/ray/_private/event/export_event_logger.py new file mode 100644 index 0000000000000000000000000000000000000000..4e77ca1421ce7064efa8bfa93ac8b68a0b40b1ae --- /dev/null +++ b/lib/python3.12/site-packages/ray/_private/event/export_event_logger.py @@ -0,0 +1,239 @@ +import json +import logging +import pathlib +import random +import string +import threading +from datetime import datetime +from enum import Enum +from typing import Union + +from ray._private import ray_constants +from ray._private.protobuf_compat import message_to_dict +from ray.core.generated.export_dataset_metadata_pb2 import ( + ExportDatasetMetadata, +) +from ray.core.generated.export_dataset_operator_event_pb2 import ( + ExportDatasetOperatorEventData, +) +from ray.core.generated.export_event_pb2 import ExportEvent +from ray.core.generated.export_submission_job_event_pb2 import ( + ExportSubmissionJobEventData, +) +from ray.core.generated.export_train_state_pb2 import ( + ExportTrainRunAttemptEventData, + ExportTrainRunEventData, +) + +global_logger = logging.getLogger(__name__) + +# This contains the union of export event data types which emit events +# using the python ExportEventLoggerAdapter +ExportEventDataType = Union[ + ExportSubmissionJobEventData, + ExportTrainRunEventData, + ExportTrainRunAttemptEventData, + ExportDatasetMetadata, + ExportDatasetOperatorEventData, +] + + +class EventLogType(Enum): + """Enum class representing different types of export event logs. + + Each enum value contains a log type name and a set of supported event data types. + + Attributes: + TRAIN_STATE: Export events related to training state, supporting train run and attempt events. + SUBMISSION_JOB: Export events related to job submissions. + DATASET_METADATA: Export events related to dataset metadata. + DATASET_OPERATOR_EVENT: Export events related to Ray Data operator. + """ + + TRAIN_STATE = ( + "EXPORT_TRAIN_STATE", + {ExportTrainRunEventData, ExportTrainRunAttemptEventData}, + ) + SUBMISSION_JOB = ("EXPORT_SUBMISSION_JOB", {ExportSubmissionJobEventData}) + DATASET_METADATA = ("EXPORT_DATASET_METADATA", {ExportDatasetMetadata}) + DATASET_OPERATOR_EVENT = ( + "EXPORT_DATASET_OPERATOR_EVENT", + {ExportDatasetOperatorEventData}, + ) + + def __init__(self, log_type_name: str, event_types: set[ExportEventDataType]): + """Initialize an EventLogType enum value. + + Args: + log_type_name: String identifier for the log type. This name is used to construct the log file name. + See `_build_export_event_file_logger` for more details. + event_types: Set of event data types that this log type supports. + """ + self.log_type_name = log_type_name + self.event_types = event_types + + def supports_event_type(self, event_type: ExportEventDataType) -> bool: + """Check if this log type supports the given event data type. + + Args: + event_type: The event data type to check for support. + + Returns: + bool: True if the event type is supported, False otherwise. + """ + return type(event_type) in self.event_types + + +def generate_event_id(): + return "".join([random.choice(string.hexdigits) for _ in range(18)]) + + +class ExportEventLoggerAdapter: + def __init__(self, log_type: EventLogType, logger: logging.Logger): + """Adapter for the Python logger that's used to emit export events.""" + self.logger = logger + self.log_type = log_type + + def send_event(self, event_data: ExportEventDataType): + # NOTE: Python logger is thread-safe, + # so we don't need to protect it using locks. + try: + event = self._create_export_event(event_data) + except TypeError: + global_logger.exception( + "Failed to create ExportEvent from event_data so no " + "event will be written to file." + ) + return + + event_as_str = self._export_event_to_string(event) + + self.logger.info(event_as_str) + # Force flush so that we won't lose events + self.logger.handlers[0].flush() + + def _create_export_event(self, event_data: ExportEventDataType) -> ExportEvent: + event = ExportEvent() + event.event_id = generate_event_id() + event.timestamp = int(datetime.now().timestamp()) + if isinstance(event_data, ExportSubmissionJobEventData): + event.submission_job_event_data.CopyFrom(event_data) + event.source_type = ExportEvent.SourceType.EXPORT_SUBMISSION_JOB + elif isinstance(event_data, ExportTrainRunEventData): + event.train_run_event_data.CopyFrom(event_data) + event.source_type = ExportEvent.SourceType.EXPORT_TRAIN_RUN + elif isinstance(event_data, ExportTrainRunAttemptEventData): + event.train_run_attempt_event_data.CopyFrom(event_data) + event.source_type = ExportEvent.SourceType.EXPORT_TRAIN_RUN_ATTEMPT + elif isinstance(event_data, ExportDatasetMetadata): + event.dataset_metadata.CopyFrom(event_data) + event.source_type = ExportEvent.SourceType.EXPORT_DATASET_METADATA + elif isinstance(event_data, ExportDatasetOperatorEventData): + event.dataset_operator_event_data.CopyFrom(event_data) + event.source_type = ExportEvent.SourceType.EXPORT_DATASET_OPERATOR_EVENT + else: + raise TypeError(f"Invalid event_data type: {type(event_data)}") + if not self.log_type.supports_event_type(event_data): + global_logger.error( + f"event_data has source type {event.source_type}, however " + f"the event was sent to a logger with log type {self.log_type.log_type_name}. " + f"The event will still be written to the file of {self.log_type.log_type_name} " + "but this indicates a bug in the code." + ) + pass + return event + + def _export_event_to_string(self, event: ExportEvent) -> str: + event_data_json = {} + proto_to_dict_options = { + "always_print_fields_with_no_presence": True, + "preserving_proto_field_name": True, + "use_integers_for_enums": False, + } + event_data_field_set = event.WhichOneof("event_data") + if event_data_field_set: + event_data_json = message_to_dict( + getattr(event, event_data_field_set), + **proto_to_dict_options, + ) + else: + global_logger.error( + f"event_data missing from export event with id {event.event_id} " + f"and type {event.source_type}. An empty event will be written, " + "but this indicates a bug in the code. " + ) + pass + event_json = { + "event_id": event.event_id, + "timestamp": event.timestamp, + "source_type": ExportEvent.SourceType.Name(event.source_type), + "event_data": event_data_json, + } + return json.dumps(event_json) + + +def _build_export_event_file_logger( + log_type_name: str, sink_dir: str +) -> logging.Logger: + logger = logging.getLogger("_ray_export_event_logger_" + log_type_name) + logger.setLevel(logging.INFO) + dir_path = pathlib.Path(sink_dir) / "export_events" + filepath = dir_path / f"event_{log_type_name}.log" + dir_path.mkdir(exist_ok=True) + filepath.touch(exist_ok=True) + # Configure the logger. + # Default is 100 MB max file size + handler = logging.handlers.RotatingFileHandler( + filepath, + maxBytes=(ray_constants.RAY_EXPORT_EVENT_MAX_FILE_SIZE_BYTES), + backupCount=ray_constants.RAY_EXPORT_EVENT_MAX_BACKUP_COUNT, + ) + logger.addHandler(handler) + logger.propagate = False + return logger + + +# This lock must be used when accessing or updating global event logger dict. +_export_event_logger_lock = threading.Lock() +_export_event_logger = {} + + +def get_export_event_logger(log_type: EventLogType, sink_dir: str) -> logging.Logger: + """Get the export event logger of the current process. + + There's only one logger per export event source. + + Args: + log_type: The type of the export event. + sink_dir: The directory to sink event logs. + """ + with _export_event_logger_lock: + global _export_event_logger + log_type_name = log_type.log_type_name + if log_type_name not in _export_event_logger: + logger = _build_export_event_file_logger(log_type.log_type_name, sink_dir) + _export_event_logger[log_type_name] = ExportEventLoggerAdapter( + log_type, logger + ) + + return _export_event_logger[log_type_name] + + +def check_export_api_enabled( + source: ExportEvent.SourceType, +) -> bool: + """ + Check RAY_ENABLE_EXPORT_API_WRITE and RAY_ENABLE_EXPORT_API_WRITE_CONFIG environment + variables to verify if export events should be written for the given source type. + + Args: + source: The source of the export event. + """ + if ray_constants.RAY_ENABLE_EXPORT_API_WRITE: + return True + source_name = ExportEvent.SourceType.Name(source) + return ( + source_name in ray_constants.RAY_ENABLE_EXPORT_API_WRITE_CONFIG + if ray_constants.RAY_ENABLE_EXPORT_API_WRITE_CONFIG + else False + ) diff --git a/lib/python3.12/site-packages/ray/_private/external_storage.py b/lib/python3.12/site-packages/ray/_private/external_storage.py new file mode 100644 index 0000000000000000000000000000000000000000..bfe9002f9e17a352db63d3a0258556537bfb47c3 --- /dev/null +++ b/lib/python3.12/site-packages/ray/_private/external_storage.py @@ -0,0 +1,652 @@ +import abc +import logging +import os +import random +import shutil +import time +import urllib +import uuid +from collections import namedtuple +from typing import IO, List, Optional, Tuple, Union + +import ray +from ray._private.ray_constants import DEFAULT_OBJECT_PREFIX +from ray._raylet import ObjectRef + +ParsedURL = namedtuple("ParsedURL", "base_url, offset, size") +logger = logging.getLogger(__name__) + + +def create_url_with_offset(*, url: str, offset: int, size: int) -> str: + """Methods to create a URL with offset. + + When ray spills objects, it fuses multiple objects + into one file to optimize the performance. That says, each object + needs to keep tracking of its own special url to store metadata. + + This method creates an url_with_offset, which is used internally + by Ray. + + Created url_with_offset can be passed to the self._get_base_url method + to parse the filename used to store files. + + Example) file://path/to/file?offset=""&size="" + + Args: + url: url to the object stored in the external storage. + offset: Offset from the beginning of the file to + the first bytes of this object. + size: Size of the object that is stored in the url. + It is used to calculate the last offset. + + Returns: + url_with_offset stored internally to find + objects from external storage. + """ + return f"{url}?offset={offset}&size={size}" + + +def parse_url_with_offset(url_with_offset: str) -> Tuple[str, int, int]: + """Parse url_with_offset to retrieve information. + + base_url is the url where the object ref + is stored in the external storage. + + Args: + url_with_offset: url created by create_url_with_offset. + + Returns: + named tuple of base_url, offset, and size. + """ + parsed_result = urllib.parse.urlparse(url_with_offset) + query_dict = urllib.parse.parse_qs(parsed_result.query) + # Split by ? to remove the query from the url. + base_url = parsed_result.geturl().split("?")[0] + if "offset" not in query_dict or "size" not in query_dict: + raise ValueError(f"Failed to parse URL: {url_with_offset}") + offset = int(query_dict["offset"][0]) + size = int(query_dict["size"][0]) + return ParsedURL(base_url=base_url, offset=offset, size=size) + + +class ExternalStorage(metaclass=abc.ABCMeta): + """The base class for external storage. + + This class provides some useful functions for zero-copy object + put/get from plasma store. Also it specifies the interface for + object spilling. + + When inheriting this class, please make sure to implement validation + logic inside __init__ method. When ray instance starts, it will + instantiating external storage to validate the config. + + Raises: + ValueError: when given configuration for + the external storage is invalid. + """ + + HEADER_LENGTH = 24 + CORE_WORKER_INIT_GRACE_PERIOD_S = 1 + + def __init__(self): + # NOTE(edoakes): do not access this field directly. Use the `core_worker` + # property instead to handle initialization race conditions. + self._core_worker: Optional["ray._raylet.CoreWorker"] = None + + @property + def core_worker(self) -> "ray._raylet.CoreWorker": + """Get the core_worker initialized in this process. + + In rare cases, the core worker may not be fully initialized by the time an I/O + worker begins to execute an operation because there is no explicit flag set to + indicate that the Python layer is ready to execute tasks. + """ + if self._core_worker is None: + worker = ray._private.worker.global_worker + start = time.time() + while not worker.connected: + time.sleep(0.001) + if time.time() - start > self.CORE_WORKER_INIT_GRACE_PERIOD_S: + raise RuntimeError( + "CoreWorker didn't initialize within grace period of " + f"{self.CORE_WORKER_INIT_GRACE_PERIOD_S}s." + ) + + self._core_worker = worker.core_worker + + return self._core_worker + + def _get_objects_from_store(self, object_refs): + # Since the object should always exist in the plasma store before + # spilling, it can directly get the object from the local plasma + # store. + # issue: https://github.com/ray-project/ray/pull/13831 + return self.core_worker.get_if_local(object_refs) + + def _put_object_to_store( + self, metadata, data_size, file_like, object_ref, owner_address + ): + self.core_worker.put_file_like_object( + metadata, data_size, file_like, object_ref, owner_address + ) + + def _write_multiple_objects( + self, f: IO, object_refs: List[ObjectRef], owner_addresses: List[str], url: str + ) -> List[str]: + """Fuse all given objects into a given file handle. + + Args: + f: File handle to fusion all given object refs. + object_refs: Object references to fusion to a single file. + owner_addresses: Owner addresses for the provided objects. + url: url where the object ref is stored + in the external storage. + + Return: + List of urls_with_offset of fused objects. + The order of returned keys are equivalent to the one + with given object_refs. + """ + keys = [] + offset = 0 + ray_object_pairs = self._get_objects_from_store(object_refs) + for ref, (buf, metadata, _), owner_address in zip( + object_refs, ray_object_pairs, owner_addresses + ): + address_len = len(owner_address) + metadata_len = len(metadata) + if buf is None and len(metadata) == 0: + error = f"Object {ref.hex()} does not exist." + raise ValueError(error) + buf_len = 0 if buf is None else len(buf) + payload = ( + address_len.to_bytes(8, byteorder="little") + + metadata_len.to_bytes(8, byteorder="little") + + buf_len.to_bytes(8, byteorder="little") + + owner_address + + metadata + + (memoryview(buf) if buf_len else b"") + ) + # 24 bytes to store owner address, metadata, and buffer lengths. + payload_len = len(payload) + assert ( + self.HEADER_LENGTH + address_len + metadata_len + buf_len == payload_len + ) + written_bytes = f.write(payload) + assert written_bytes == payload_len + url_with_offset = create_url_with_offset( + url=url, offset=offset, size=written_bytes + ) + keys.append(url_with_offset.encode()) + offset += written_bytes + # Necessary because pyarrow.io.NativeFile does not flush() on close(). + f.flush() + return keys + + def _size_check(self, address_len, metadata_len, buffer_len, obtained_data_size): + """Check whether or not the obtained_data_size is as expected. + + Args: + address_len: Length of the address. + metadata_len: Actual metadata length of the object. + buffer_len: Actual buffer length of the object. + obtained_data_size: Data size specified in the url_with_offset. + + Raises: + ValueError: If obtained_data_size is different from + address_len + metadata_len + buffer_len + 24 (first 8 bytes to store length). + """ + data_size_in_bytes = ( + address_len + metadata_len + buffer_len + self.HEADER_LENGTH + ) + if data_size_in_bytes != obtained_data_size: + raise ValueError( + f"Obtained data has a size of {data_size_in_bytes}, " + "although it is supposed to have the " + f"size of {obtained_data_size}." + ) + + @abc.abstractmethod + def spill_objects(self, object_refs, owner_addresses) -> List[str]: + """Spill objects to the external storage. Objects are specified + by their object refs. + + Args: + object_refs: The list of the refs of the objects to be spilled. + owner_addresses: Owner addresses for the provided objects. + Returns: + A list of internal URLs with object offset. + """ + + @abc.abstractmethod + def restore_spilled_objects( + self, object_refs: List[ObjectRef], url_with_offset_list: List[str] + ) -> int: + """Restore objects from the external storage. + + Args: + object_refs: List of object IDs (note that it is not ref). + url_with_offset_list: List of url_with_offset. + + Returns: + The total number of bytes restored. + """ + + @abc.abstractmethod + def delete_spilled_objects(self, urls: List[str]): + """Delete objects that are spilled to the external storage. + + Args: + urls: URLs that store spilled object files. + + NOTE: This function should not fail if some of the urls + do not exist. + """ + + @abc.abstractmethod + def destroy_external_storage(self): + """Destroy external storage when a head node is down. + + NOTE: This is currently working when the cluster is + started by ray.init + """ + + +class NullStorage(ExternalStorage): + """The class that represents an uninitialized external storage.""" + + def spill_objects(self, object_refs, owner_addresses) -> List[str]: + raise NotImplementedError("External storage is not initialized") + + def restore_spilled_objects(self, object_refs, url_with_offset_list): + raise NotImplementedError("External storage is not initialized") + + def delete_spilled_objects(self, urls: List[str]): + raise NotImplementedError("External storage is not initialized") + + def destroy_external_storage(self): + raise NotImplementedError("External storage is not initialized") + + +class FileSystemStorage(ExternalStorage): + """The class for filesystem-like external storage. + + Raises: + ValueError: Raises directory path to + spill objects doesn't exist. + """ + + def __init__( + self, + node_id: str, + directory_path: Union[str, List[str]], + buffer_size: Optional[int] = None, + ): + super().__init__() + + # -- A list of directory paths to spill objects -- + self._directory_paths = [] + # -- Current directory to spill objects -- + self._current_directory_index = 0 + # -- File buffer size to spill objects -- + self._buffer_size = -1 + + # Validation. + assert ( + directory_path is not None + ), "directory_path should be provided to use object spilling." + if isinstance(directory_path, str): + directory_path = [directory_path] + assert isinstance( + directory_path, list + ), "Directory_path must be either a single string or a list of strings" + if buffer_size is not None: + assert isinstance(buffer_size, int), "buffer_size must be an integer." + self._buffer_size = buffer_size + + # Create directories. + for path in directory_path: + full_dir_path = os.path.join(path, f"{DEFAULT_OBJECT_PREFIX}_{node_id}") + os.makedirs(full_dir_path, exist_ok=True) + if not os.path.exists(full_dir_path): + raise ValueError( + "The given directory path to store objects, " + f"{full_dir_path}, could not be created." + ) + self._directory_paths.append(full_dir_path) + assert len(self._directory_paths) == len(directory_path) + # Choose the current directory. + # It chooses a random index to maximize multiple directories that are + # mounted at different point. + self._current_directory_index = random.randrange(0, len(self._directory_paths)) + + def spill_objects(self, object_refs, owner_addresses) -> List[str]: + if len(object_refs) == 0: + return [] + # Choose the current directory path by round robin order. + self._current_directory_index = (self._current_directory_index + 1) % len( + self._directory_paths + ) + directory_path = self._directory_paths[self._current_directory_index] + + filename = _get_unique_spill_filename(object_refs) + url = f"{os.path.join(directory_path, filename)}" + with open(url, "wb", buffering=self._buffer_size) as f: + return self._write_multiple_objects(f, object_refs, owner_addresses, url) + + def restore_spilled_objects( + self, object_refs: List[ObjectRef], url_with_offset_list: List[str] + ): + total = 0 + for i in range(len(object_refs)): + object_ref = object_refs[i] + url_with_offset = url_with_offset_list[i].decode() + # Retrieve the information needed. + parsed_result = parse_url_with_offset(url_with_offset) + base_url = parsed_result.base_url + offset = parsed_result.offset + # Read a part of the file and recover the object. + with open(base_url, "rb") as f: + f.seek(offset) + address_len = int.from_bytes(f.read(8), byteorder="little") + metadata_len = int.from_bytes(f.read(8), byteorder="little") + buf_len = int.from_bytes(f.read(8), byteorder="little") + self._size_check(address_len, metadata_len, buf_len, parsed_result.size) + total += buf_len + owner_address = f.read(address_len) + metadata = f.read(metadata_len) + # read remaining data to our buffer + self._put_object_to_store( + metadata, buf_len, f, object_ref, owner_address + ) + return total + + def delete_spilled_objects(self, urls: List[str]): + for url in urls: + path = parse_url_with_offset(url.decode()).base_url + try: + os.remove(path) + except FileNotFoundError: + # Occurs when the urls are retried during worker crash/failure. + pass + + def destroy_external_storage(self): + for directory_path in self._directory_paths: + self._destroy_external_storage(directory_path) + + def _destroy_external_storage(self, directory_path): + # There's a race condition where IO workers are still + # deleting each objects while we try deleting the + # whole directory. So we should keep trying it until + # The directory is actually deleted. + while os.path.isdir(directory_path): + try: + shutil.rmtree(directory_path) + except (FileNotFoundError): + # If exception occurs when other IO workers are + # deleting the file at the same time. + pass + except Exception: + logger.exception( + "Error cleaning up spill files. " + "You might still have remaining spilled " + "objects inside `ray_spilled_objects` directory." + ) + break + + +class ExternalStorageSmartOpenImpl(ExternalStorage): + """The external storage class implemented by smart_open. + (https://github.com/RaRe-Technologies/smart_open) + + Smart open supports multiple backend with the same APIs. + + To use this implementation, you should pre-create the given uri. + For example, if your uri is a local file path, you should pre-create + the directory. + + Args: + uri: Storage URI used for smart open. + prefix: Prefix of objects that are stored. + override_transport_params: Overriding the default value of + transport_params for smart-open library. + + Raises: + ModuleNotFoundError: If it fails to setup. + For example, if smart open library + is not downloaded, this will fail. + """ + + def __init__( + self, + node_id: str, + uri: str or list, + override_transport_params: dict = None, + buffer_size=1024 * 1024, # For remote spilling, at least 1MB is recommended. + ): + super().__init__() + + try: + from smart_open import open # noqa + except ModuleNotFoundError as e: + raise ModuleNotFoundError( + "Smart open is chosen to be a object spilling " + "external storage, but smart_open and boto3 " + f"is not downloaded. Original error: {e}" + ) + + # Validation + assert uri is not None, "uri should be provided to use object spilling." + if isinstance(uri, str): + uri = [uri] + assert isinstance(uri, list), "uri must be a single string or list of strings." + assert isinstance(buffer_size, int), "buffer_size must be an integer." + + uri_is_s3 = [u.startswith("s3://") for u in uri] + self.is_for_s3 = all(uri_is_s3) + if not self.is_for_s3: + assert not any(uri_is_s3), "all uri's must be s3 or none can be s3." + self._uris = uri + else: + self._uris = [u.strip("/") for u in uri] + assert len(self._uris) == len(uri) + + self._current_uri_index = random.randrange(0, len(self._uris)) + self.prefix = f"{DEFAULT_OBJECT_PREFIX}_{node_id}" + self.override_transport_params = override_transport_params or {} + + if self.is_for_s3: + import boto3 # noqa + + # Setup boto3. It is essential because if we don't create boto + # session, smart_open will create a new session for every + # open call. + self.s3 = boto3.resource(service_name="s3") + + # smart_open always seek to 0 if we don't set this argument. + # This will lead us to call a Object.get when it is not necessary, + # so defer seek and call seek before reading objects instead. + self.transport_params = { + "defer_seek": True, + "resource": self.s3, + "buffer_size": buffer_size, + } + else: + self.transport_params = {} + + self.transport_params.update(self.override_transport_params) + + def spill_objects(self, object_refs, owner_addresses) -> List[str]: + if len(object_refs) == 0: + return [] + from smart_open import open + + # Choose the current uri by round robin order. + self._current_uri_index = (self._current_uri_index + 1) % len(self._uris) + uri = self._uris[self._current_uri_index] + + key = f"{self.prefix}-{_get_unique_spill_filename(object_refs)}" + url = f"{uri}/{key}" + + with open( + url, + mode="wb", + transport_params=self.transport_params, + ) as file_like: + return self._write_multiple_objects( + file_like, object_refs, owner_addresses, url + ) + + def restore_spilled_objects( + self, object_refs: List[ObjectRef], url_with_offset_list: List[str] + ): + from smart_open import open + + total = 0 + for i in range(len(object_refs)): + object_ref = object_refs[i] + url_with_offset = url_with_offset_list[i].decode() + + # Retrieve the information needed. + parsed_result = parse_url_with_offset(url_with_offset) + base_url = parsed_result.base_url + offset = parsed_result.offset + + with open(base_url, "rb", transport_params=self.transport_params) as f: + # smart open seek reads the file from offset-end_of_the_file + # when the seek is called. + f.seek(offset) + address_len = int.from_bytes(f.read(8), byteorder="little") + metadata_len = int.from_bytes(f.read(8), byteorder="little") + buf_len = int.from_bytes(f.read(8), byteorder="little") + self._size_check(address_len, metadata_len, buf_len, parsed_result.size) + owner_address = f.read(address_len) + total += buf_len + metadata = f.read(metadata_len) + # read remaining data to our buffer + self._put_object_to_store( + metadata, buf_len, f, object_ref, owner_address + ) + return total + + def delete_spilled_objects(self, urls: List[str]): + pass + + def destroy_external_storage(self): + pass + + +_external_storage = NullStorage() + + +class UnstableFileStorage(FileSystemStorage): + """This class is for testing with writing failure.""" + + def __init__(self, node_id: str, **kwargs): + super().__init__(node_id, **kwargs) + self._failure_rate = 0.1 + self._partial_failure_ratio = 0.2 + + def spill_objects(self, object_refs, owner_addresses) -> List[str]: + r = random.random() < self._failure_rate + failed = r < self._failure_rate + partial_failed = r < self._partial_failure_ratio + if failed: + raise IOError("Spilling object failed intentionally for testing.") + elif partial_failed: + i = random.choice(range(len(object_refs))) + return super().spill_objects(object_refs[:i], owner_addresses) + else: + return super().spill_objects(object_refs, owner_addresses) + + +class SlowFileStorage(FileSystemStorage): + """This class is for testing slow object spilling.""" + + def __init__(self, node_id: str, **kwargs): + super().__init__(node_id, **kwargs) + self._min_delay = 1 + self._max_delay = 2 + + def spill_objects(self, object_refs, owner_addresses) -> List[str]: + delay = random.random() * (self._max_delay - self._min_delay) + self._min_delay + time.sleep(delay) + return super().spill_objects(object_refs, owner_addresses) + + +def setup_external_storage(config, node_id, session_name): + """Setup the external storage according to the config.""" + assert node_id is not None, "node_id should be provided." + global _external_storage + if config: + storage_type = config["type"] + if storage_type == "filesystem": + _external_storage = FileSystemStorage(node_id, **config["params"]) + elif storage_type == "smart_open": + _external_storage = ExternalStorageSmartOpenImpl( + node_id, **config["params"] + ) + elif storage_type == "mock_distributed_fs": + # This storage is used to unit test distributed external storages. + # TODO(sang): Delete it after introducing the mock S3 test. + _external_storage = FileSystemStorage(node_id, **config["params"]) + elif storage_type == "unstable_fs": + # This storage is used to unit test unstable file system for fault + # tolerance. + _external_storage = UnstableFileStorage(node_id, **config["params"]) + elif storage_type == "slow_fs": + # This storage is used to unit test slow filesystems. + _external_storage = SlowFileStorage(node_id, **config["params"]) + else: + raise ValueError(f"Unknown external storage type: {storage_type}") + else: + _external_storage = NullStorage() + return _external_storage + + +def reset_external_storage(): + global _external_storage + _external_storage = NullStorage() + + +def spill_objects(object_refs, owner_addresses): + """Spill objects to the external storage. Objects are specified + by their object refs. + + Args: + object_refs: The list of the refs of the objects to be spilled. + owner_addresses: The owner addresses of the provided object refs. + Returns: + A list of keys corresponding to the input object refs. + """ + return _external_storage.spill_objects(object_refs, owner_addresses) + + +def restore_spilled_objects( + object_refs: List[ObjectRef], url_with_offset_list: List[str] +): + """Restore objects from the external storage. + + Args: + object_refs: List of object IDs (note that it is not ref). + url_with_offset_list: List of url_with_offset. + """ + return _external_storage.restore_spilled_objects(object_refs, url_with_offset_list) + + +def delete_spilled_objects(urls: List[str]): + """Delete objects that are spilled to the external storage. + + Args: + urls: URLs that store spilled object files. + """ + _external_storage.delete_spilled_objects(urls) + + +def _get_unique_spill_filename(object_refs: List[ObjectRef]): + """Generate a unqiue spill file name. + + Args: + object_refs: objects to be spilled in this file. + """ + return f"{uuid.uuid4().hex}-multi-{len(object_refs)}" diff --git a/lib/python3.12/site-packages/ray/_private/function_manager.py b/lib/python3.12/site-packages/ray/_private/function_manager.py new file mode 100644 index 0000000000000000000000000000000000000000..e53ce1a1d3f820b2bdfff6b3fc75339d474b36b1 --- /dev/null +++ b/lib/python3.12/site-packages/ray/_private/function_manager.py @@ -0,0 +1,703 @@ +import dis +import hashlib +import importlib +import inspect +import json +import logging +import os +import sys +import threading +import time +import traceback +from collections import defaultdict, namedtuple +from typing import Callable, Optional + +import ray +import ray._private.profiling as profiling +from ray import cloudpickle as pickle +from ray._common.serialization import pickle_dumps +from ray._private import ray_constants +from ray._private.inspect_util import ( + is_class_method, + is_function_or_method, + is_static_method, +) +from ray._private.ray_constants import KV_NAMESPACE_FUNCTION_TABLE +from ray._private.utils import ( + check_oversized_function, + ensure_str, + format_error_message, +) +from ray._raylet import ( + WORKER_PROCESS_SETUP_HOOK_KEY_NAME_GCS, + JobID, + PythonFunctionDescriptor, +) +from ray.remote_function import RemoteFunction + +FunctionExecutionInfo = namedtuple( + "FunctionExecutionInfo", ["function", "function_name", "max_calls"] +) +ImportedFunctionInfo = namedtuple( + "ImportedFunctionInfo", + ["job_id", "function_id", "function_name", "function", "module", "max_calls"], +) + +"""FunctionExecutionInfo: A named tuple storing remote function information.""" + +logger = logging.getLogger(__name__) + + +def make_function_table_key(key_type: bytes, job_id: JobID, key: Optional[bytes]): + if key is None: + return b":".join([key_type, job_id.hex().encode()]) + else: + return b":".join([key_type, job_id.hex().encode(), key]) + + +class FunctionActorManager: + """A class used to export/load remote functions and actors. + Attributes: + _worker: The associated worker that this manager related. + _functions_to_export: The remote functions to export when + the worker gets connected. + _actors_to_export: The actors to export when the worker gets + connected. + _function_execution_info: The function_id + and execution_info. + _num_task_executions: The function + execution times. + imported_actor_classes: The set of actor classes keys (format: + ActorClass:function_id) that are already in GCS. + """ + + def __init__(self, worker): + self._worker = worker + self._functions_to_export = [] + self._actors_to_export = [] + # This field is a dictionary that maps function IDs + # to a FunctionExecutionInfo object. This should only be used on + # workers that execute remote functions. + self._function_execution_info = defaultdict(lambda: {}) + self._num_task_executions = defaultdict(lambda: {}) + # A set of all of the actor class keys that have been imported by the + # import thread. It is safe to convert this worker into an actor of + # these types. + self.imported_actor_classes = set() + self._loaded_actor_classes = {} + # Deserialize an ActorHandle will call load_actor_class(). If a + # function closure captured an ActorHandle, the deserialization of the + # function will be: + # -> fetch_and_register_remote_function (acquire lock) + # -> _load_actor_class_from_gcs (acquire lock, too) + # So, the lock should be a reentrant lock. + self.lock = threading.RLock() + + self.execution_infos = {} + # This is the counter to keep track of how many keys have already + # been exported so that we can find next key quicker. + self._num_exported = 0 + # This is to protect self._num_exported when doing exporting + self._export_lock = threading.Lock() + + def increase_task_counter(self, function_descriptor): + function_id = function_descriptor.function_id + self._num_task_executions[function_id] += 1 + + def get_task_counter(self, function_descriptor): + function_id = function_descriptor.function_id + return self._num_task_executions[function_id] + + def compute_collision_identifier(self, function_or_class): + """The identifier is used to detect excessive duplicate exports. + The identifier is used to determine when the same function or class is + exported many times. This can yield false positives. + Args: + function_or_class: The function or class to compute an identifier + for. + Returns: + The identifier. Note that different functions or classes can give + rise to same identifier. However, the same function should + hopefully always give rise to the same identifier. TODO(rkn): + verify if this is actually the case. Note that if the + identifier is incorrect in any way, then we may give warnings + unnecessarily or fail to give warnings, but the application's + behavior won't change. + """ + import io + + string_file = io.StringIO() + dis.dis(function_or_class, file=string_file, depth=2) + collision_identifier = function_or_class.__name__ + ":" + string_file.getvalue() + + # Return a hash of the identifier in case it is too large. + return hashlib.sha1(collision_identifier.encode("utf-8")).digest() + + def load_function_or_class_from_local(self, module_name, function_or_class_name): + """Try to load a function or class in the module from local.""" + module = importlib.import_module(module_name) + parts = [part for part in function_or_class_name.split(".") if part] + object = module + try: + for part in parts: + object = getattr(object, part) + return object + except Exception: + return None + + def export_setup_func( + self, setup_func: Callable, timeout: Optional[int] = None + ) -> bytes: + """Export the setup hook function and return the key.""" + pickled_function = pickle_dumps( + setup_func, + "Cannot serialize the worker_process_setup_hook " f"{setup_func.__name__}", + ) + + function_to_run_id = hashlib.shake_128(pickled_function).digest( + ray_constants.ID_SIZE + ) + key = make_function_table_key( + # This value should match with gcs_function_manager.h. + # Otherwise, it won't be GC'ed. + WORKER_PROCESS_SETUP_HOOK_KEY_NAME_GCS.encode(), + # b"FunctionsToRun", + self._worker.current_job_id.binary(), + function_to_run_id, + ) + + check_oversized_function( + pickled_function, setup_func.__name__, "function", self._worker + ) + + try: + self._worker.gcs_client.internal_kv_put( + key, + pickle.dumps( + { + "job_id": self._worker.current_job_id.binary(), + "function_id": function_to_run_id, + "function": pickled_function, + } + ), + # overwrite + True, + ray_constants.KV_NAMESPACE_FUNCTION_TABLE, + timeout=timeout, + ) + except Exception as e: + logger.exception( + "Failed to export the setup hook " f"{setup_func.__name__}." + ) + raise e + + return key + + def export(self, remote_function): + """Pickle a remote function and export it to redis. + Args: + remote_function: the RemoteFunction object. + """ + if self._worker.load_code_from_local: + function_descriptor = remote_function._function_descriptor + module_name, function_name = ( + function_descriptor.module_name, + function_descriptor.function_name, + ) + # If the function is dynamic, we still export it to GCS + # even if load_code_from_local is set True. + if ( + self.load_function_or_class_from_local(module_name, function_name) + is not None + ): + return + function = remote_function._function + pickled_function = remote_function._pickled_function + + check_oversized_function( + pickled_function, + remote_function._function_name, + "remote function", + self._worker, + ) + key = make_function_table_key( + b"RemoteFunction", + self._worker.current_job_id, + remote_function._function_descriptor.function_id.binary(), + ) + if self._worker.gcs_client.internal_kv_exists(key, KV_NAMESPACE_FUNCTION_TABLE): + return + val = pickle.dumps( + { + "job_id": self._worker.current_job_id.binary(), + "function_id": remote_function._function_descriptor.function_id.binary(), # noqa: E501 + "function_name": remote_function._function_name, + "module": function.__module__, + "function": pickled_function, + "collision_identifier": self.compute_collision_identifier(function), + "max_calls": remote_function._max_calls, + } + ) + self._worker.gcs_client.internal_kv_put( + key, val, True, KV_NAMESPACE_FUNCTION_TABLE + ) + + def fetch_registered_method( + self, key: str, timeout: Optional[int] = None + ) -> Optional[ImportedFunctionInfo]: + vals = self._worker.gcs_client.internal_kv_get( + key, KV_NAMESPACE_FUNCTION_TABLE, timeout=timeout + ) + if vals is None: + return None + else: + vals = pickle.loads(vals) + fields = [ + "job_id", + "function_id", + "function_name", + "function", + "module", + "max_calls", + ] + return ImportedFunctionInfo._make(vals.get(field) for field in fields) + + def fetch_and_register_remote_function(self, key): + """Import a remote function.""" + remote_function_info = self.fetch_registered_method(key) + if not remote_function_info: + return False + ( + job_id_str, + function_id_str, + function_name, + serialized_function, + module, + max_calls, + ) = remote_function_info + + function_id = ray.FunctionID(function_id_str) + job_id = ray.JobID(job_id_str) + max_calls = int(max_calls) + + # This function is called by ImportThread. This operation needs to be + # atomic. Otherwise, there is race condition. Another thread may use + # the temporary function above before the real function is ready. + with self.lock: + self._num_task_executions[function_id] = 0 + + try: + function = pickle.loads(serialized_function) + except Exception: + # If an exception was thrown when the remote function was + # imported, we record the traceback and notify the scheduler + # of the failure. + traceback_str = format_error_message(traceback.format_exc()) + + def f(*args, **kwargs): + raise RuntimeError( + "The remote function failed to import on the " + "worker. This may be because needed library " + "dependencies are not installed in the worker " + "environment or cannot be found from sys.path " + f"{sys.path}:\n\n{traceback_str}" + ) + + # Use a placeholder method when function pickled failed + self._function_execution_info[function_id] = FunctionExecutionInfo( + function=f, function_name=function_name, max_calls=max_calls + ) + + # Log the error message. Log at DEBUG level to avoid overly + # spamming the log on import failure. The user gets the error + # via the RuntimeError message above. + logger.debug( + "Failed to unpickle the remote function " + f"'{function_name}' with " + f"function ID {function_id.hex()}. " + f"Job ID:{job_id}." + f"Traceback:\n{traceback_str}. " + ) + else: + # The below line is necessary. Because in the driver process, + # if the function is defined in the file where the python + # script was started from, its module is `__main__`. + # However in the worker process, the `__main__` module is a + # different module, which is `default_worker.py` + function.__module__ = module + self._function_execution_info[function_id] = FunctionExecutionInfo( + function=function, function_name=function_name, max_calls=max_calls + ) + return True + + def get_execution_info(self, job_id, function_descriptor): + """Get the FunctionExecutionInfo of a remote function. + Args: + job_id: ID of the job that the function belongs to. + function_descriptor: The FunctionDescriptor of the function to get. + Returns: + A FunctionExecutionInfo object. + """ + function_id = function_descriptor.function_id + # If the function has already been loaded, + # There's no need to load again + if function_id in self._function_execution_info: + return self._function_execution_info[function_id] + if self._worker.load_code_from_local: + # Load function from local code. + if not function_descriptor.is_actor_method(): + # If the function is not able to be loaded, + # try to load it from GCS, + # even if load_code_from_local is set True + if self._load_function_from_local(function_descriptor) is True: + return self._function_execution_info[function_id] + # Load function from GCS. + # Wait until the function to be executed has actually been + # registered on this worker. We will push warnings to the user if + # we spend too long in this loop. + # The driver function may not be found in sys.path. Try to load + # the function from GCS. + with profiling.profile("wait_for_function"): + self._wait_for_function(function_descriptor, job_id) + try: + function_id = function_descriptor.function_id + info = self._function_execution_info[function_id] + except KeyError as e: + message = ( + "Error occurs in get_execution_info: " + "job_id: %s, function_descriptor: %s. Message: %s" + % (job_id, function_descriptor, e) + ) + raise KeyError(message) + return info + + def _load_function_from_local(self, function_descriptor): + assert not function_descriptor.is_actor_method() + function_id = function_descriptor.function_id + + module_name, function_name = ( + function_descriptor.module_name, + function_descriptor.function_name, + ) + + object = self.load_function_or_class_from_local(module_name, function_name) + if object is not None: + # Directly importing from local may break function with dynamic ray.remote, + # such as the _start_controller function utilized for the Ray service. + if isinstance(object, RemoteFunction): + function = object._function + else: + function = object + self._function_execution_info[function_id] = FunctionExecutionInfo( + function=function, + function_name=function_name, + max_calls=0, + ) + self._num_task_executions[function_id] = 0 + return True + else: + return False + + def _wait_for_function(self, function_descriptor, job_id: str, timeout=10): + """Wait until the function to be executed is present on this worker. + This method will simply loop until the import thread has imported the + relevant function. If we spend too long in this loop, that may indicate + a problem somewhere and we will push an error message to the user. + If this worker is an actor, then this will wait until the actor has + been defined. + Args: + function_descriptor : The FunctionDescriptor of the function that + we want to execute. + job_id: The ID of the job to push the error message to + if this times out. + """ + start_time = time.time() + # Only send the warning once. + warning_sent = False + while True: + with self.lock: + if self._worker.actor_id.is_nil(): + if function_descriptor.function_id in self._function_execution_info: + break + else: + key = make_function_table_key( + b"RemoteFunction", + job_id, + function_descriptor.function_id.binary(), + ) + if self.fetch_and_register_remote_function(key) is True: + break + else: + assert not self._worker.actor_id.is_nil() + # Actor loading will happen when execute_task is called. + assert self._worker.actor_id in self._worker.actors + break + + if time.time() - start_time > timeout: + warning_message = ( + "This worker was asked to execute a function " + f"that has not been registered ({function_descriptor}, " + f"node={self._worker.node_ip_address}, " + f"worker_id={self._worker.worker_id.hex()}, " + f"pid={os.getpid()}). You may have to restart Ray." + ) + if not warning_sent: + logger.error(warning_message) + ray._private.utils.push_error_to_driver( + self._worker, + ray_constants.WAIT_FOR_FUNCTION_PUSH_ERROR, + warning_message, + job_id=job_id, + ) + warning_sent = True + time.sleep(0.001) + + def export_actor_class( + self, Class, actor_creation_function_descriptor, actor_method_names + ): + if self._worker.load_code_from_local: + module_name, class_name = ( + actor_creation_function_descriptor.module_name, + actor_creation_function_descriptor.class_name, + ) + # If the class is dynamic, we still export it to GCS + # even if load_code_from_local is set True. + if ( + self.load_function_or_class_from_local(module_name, class_name) + is not None + ): + return + + # `current_job_id` shouldn't be NIL, unless: + # 1) This worker isn't an actor; + # 2) And a previous task started a background thread, which didn't + # finish before the task finished, and still uses Ray API + # after that. + assert not self._worker.current_job_id.is_nil(), ( + "You might have started a background thread in a non-actor " + "task, please make sure the thread finishes before the " + "task finishes." + ) + job_id = self._worker.current_job_id + key = make_function_table_key( + b"ActorClass", + job_id, + actor_creation_function_descriptor.function_id.binary(), + ) + serialized_actor_class = pickle_dumps( + Class, + f"Could not serialize the actor class " + f"{actor_creation_function_descriptor.repr}", + ) + actor_class_info = { + "class_name": actor_creation_function_descriptor.class_name.split(".")[-1], + "module": actor_creation_function_descriptor.module_name, + "class": serialized_actor_class, + "job_id": job_id.binary(), + "collision_identifier": self.compute_collision_identifier(Class), + "actor_method_names": json.dumps(list(actor_method_names)), + } + + check_oversized_function( + actor_class_info["class"], + actor_class_info["class_name"], + "actor", + self._worker, + ) + + self._worker.gcs_client.internal_kv_put( + key, pickle.dumps(actor_class_info), True, KV_NAMESPACE_FUNCTION_TABLE + ) + # TODO(rkn): Currently we allow actor classes to be defined + # within tasks. I tried to disable this, but it may be necessary + # because of https://github.com/ray-project/ray/issues/1146. + + def load_actor_class(self, job_id, actor_creation_function_descriptor): + """Load the actor class. + Args: + job_id: job ID of the actor. + actor_creation_function_descriptor: Function descriptor of + the actor constructor. + Returns: + The actor class. + """ + function_id = actor_creation_function_descriptor.function_id + # Check if the actor class already exists in the cache. + actor_class = self._loaded_actor_classes.get(function_id, None) + if actor_class is None: + # Load actor class. + if self._worker.load_code_from_local: + # Load actor class from local code first. + actor_class = self._load_actor_class_from_local( + actor_creation_function_descriptor + ) + # If the actor is unable to be loaded + # from local, try to load it + # from GCS even if load_code_from_local is set True + if actor_class is None: + actor_class = self._load_actor_class_from_gcs( + job_id, actor_creation_function_descriptor + ) + + else: + # Load actor class from GCS. + actor_class = self._load_actor_class_from_gcs( + job_id, actor_creation_function_descriptor + ) + # Save the loaded actor class in cache. + self._loaded_actor_classes[function_id] = actor_class + + # Generate execution info for the methods of this actor class. + module_name = actor_creation_function_descriptor.module_name + actor_class_name = actor_creation_function_descriptor.class_name + actor_methods = inspect.getmembers( + actor_class, predicate=is_function_or_method + ) + for actor_method_name, actor_method in actor_methods: + # Actor creation function descriptor use a unique function + # hash to solve actor name conflict. When constructing an + # actor, the actor creation function descriptor will be the + # key to find __init__ method execution info. So, here we + # use actor creation function descriptor as method descriptor + # for generating __init__ method execution info. + if actor_method_name == "__init__": + method_descriptor = actor_creation_function_descriptor + else: + method_descriptor = PythonFunctionDescriptor( + module_name, actor_method_name, actor_class_name + ) + method_id = method_descriptor.function_id + executor = self._make_actor_method_executor( + actor_method_name, actor_method + ) + self._function_execution_info[method_id] = FunctionExecutionInfo( + function=executor, + function_name=actor_method_name, + max_calls=0, + ) + self._num_task_executions[method_id] = 0 + self._num_task_executions[function_id] = 0 + return actor_class + + def _load_actor_class_from_local(self, actor_creation_function_descriptor): + """Load actor class from local code.""" + module_name, class_name = ( + actor_creation_function_descriptor.module_name, + actor_creation_function_descriptor.class_name, + ) + + object = self.load_function_or_class_from_local(module_name, class_name) + + if object is not None: + if isinstance(object, ray.actor.ActorClass): + return object.__ray_metadata__.modified_class + else: + return object + else: + return None + + def _create_fake_actor_class( + self, actor_class_name, actor_method_names, traceback_str + ): + class TemporaryActor: + async def __dummy_method(self): + """Dummy method for this fake actor class to work for async actors. + Without this method, this temporary actor class fails to initialize + if the original actor class was async.""" + pass + + def temporary_actor_method(*args, **kwargs): + raise RuntimeError( + f"The actor with name {actor_class_name} " + "failed to import on the worker. This may be because " + "needed library dependencies are not installed in the " + f"worker environment:\n\n{traceback_str}" + ) + + for method in actor_method_names: + setattr(TemporaryActor, method, temporary_actor_method) + + return TemporaryActor + + def _load_actor_class_from_gcs(self, job_id, actor_creation_function_descriptor): + """Load actor class from GCS.""" + key = make_function_table_key( + b"ActorClass", + job_id, + actor_creation_function_descriptor.function_id.binary(), + ) + + # Fetch raw data from GCS. + vals = self._worker.gcs_client.internal_kv_get(key, KV_NAMESPACE_FUNCTION_TABLE) + fields = ["job_id", "class_name", "module", "class", "actor_method_names"] + if vals is None: + vals = {} + else: + vals = pickle.loads(vals) + (job_id_str, class_name, module, pickled_class, actor_method_names) = ( + vals.get(field) for field in fields + ) + + class_name = ensure_str(class_name) + module_name = ensure_str(module) + job_id = ray.JobID(job_id_str) + actor_method_names = json.loads(ensure_str(actor_method_names)) + + actor_class = None + try: + with self.lock: + actor_class = pickle.loads(pickled_class) + except Exception: + logger.debug("Failed to load actor class %s.", class_name) + # If an exception was thrown when the actor was imported, we record + # the traceback and notify the scheduler of the failure. + traceback_str = format_error_message(traceback.format_exc()) + # The actor class failed to be unpickled, create a fake actor + # class instead (just to produce error messages and to prevent + # the driver from hanging). + actor_class = self._create_fake_actor_class( + class_name, actor_method_names, traceback_str + ) + + # The below line is necessary. Because in the driver process, + # if the function is defined in the file where the python script + # was started from, its module is `__main__`. + # However in the worker process, the `__main__` module is a + # different module, which is `default_worker.py` + actor_class.__module__ = module_name + return actor_class + + def _make_actor_method_executor(self, method_name: str, method): + """Make an executor that wraps a user-defined actor method. + The wrapped method updates the worker's internal state and performs any + necessary checkpointing operations. + Args: + method_name: The name of the actor method. + method: The actor method to wrap. This should be a + method defined on the actor class and should therefore take an + instance of the actor as the first argument. + Returns: + A function that executes the given actor method on the worker's + stored instance of the actor. The function also updates the + worker's internal state to record the executed method. + """ + + def actor_method_executor(__ray_actor, *args, **kwargs): + # Execute the assigned method. + is_bound = is_class_method(method) or is_static_method( + type(__ray_actor), method_name + ) + if is_bound: + return method(*args, **kwargs) + else: + return method(__ray_actor, *args, **kwargs) + + # Set method_name and method as attributes to the executor closure + # so we can make decision based on these attributes in task executor. + # Precisely, asyncio support requires to know whether: + # - the method is a ray internal method: starts with __ray + # - the method is a coroutine function: defined by async def + actor_method_executor.name = method_name + actor_method_executor.method = method + + return actor_method_executor diff --git a/lib/python3.12/site-packages/ray/_private/gc_collect_manager.py b/lib/python3.12/site-packages/ray/_private/gc_collect_manager.py new file mode 100644 index 0000000000000000000000000000000000000000..d9bb723b88b06fd78cfdcc216e31acd1d56e69ab --- /dev/null +++ b/lib/python3.12/site-packages/ray/_private/gc_collect_manager.py @@ -0,0 +1,65 @@ +import gc +import logging +import threading +import time +from typing import Callable, Optional + +logger = logging.getLogger(__name__) + + +class PythonGCThread(threading.Thread): + """A background thread that triggers Python garbage collection. + + This thread waits for GC events from CoreWorker and triggers `gc.collect()` when + requested, ensuring that collections are spaced out by at least + `min_interval_s` seconds.""" + + def __init__( + self, *, min_interval_s: int = 5, gc_collect_func: Optional[Callable] = None + ): + logger.debug("Starting Python GC thread") + super().__init__(name="PythonGCThread", daemon=True) + self._should_exit = False + self._last_gc_time = float("-inf") + self._min_gc_interval = min_interval_s + self._gc_event = threading.Event() + # Set the gc_collect_func for UT, defaulting to gc.collect if None + self._gc_collect_func = gc_collect_func or gc.collect + + def trigger_gc(self) -> None: + self._gc_event.set() + + def run(self): + while not self._should_exit: + self._gc_event.wait() + self._gc_event.clear() + + if self._should_exit: + break + + time_since_last_gc = time.monotonic() - self._last_gc_time + if time_since_last_gc < self._min_gc_interval: + logger.debug( + f"Skipping GC, only {time_since_last_gc:.2f}s since last GC" + ) + continue + + try: + start = time.monotonic() + num_freed = self._gc_collect_func() + self._last_gc_time = time.monotonic() + if num_freed > 0: + logger.debug( + "gc.collect() freed {} refs in {} seconds".format( + num_freed, self._last_gc_time - start + ) + ) + except Exception as e: + logger.error(f"Error during GC: {e}") + self._last_gc_time = time.monotonic() + + def stop(self): + logger.debug("Stopping Python GC thread") + self._should_exit = True + self._gc_event.set() + self.join() diff --git a/lib/python3.12/site-packages/ray/_private/gcs_pubsub.py b/lib/python3.12/site-packages/ray/_private/gcs_pubsub.py new file mode 100644 index 0000000000000000000000000000000000000000..e40d74b72bec7cf9d2c374db09d4dd44d24eccdf --- /dev/null +++ b/lib/python3.12/site-packages/ray/_private/gcs_pubsub.py @@ -0,0 +1,274 @@ +import asyncio +import logging +import random +from collections import deque +from typing import List, Tuple + +import grpc +from grpc import aio as aiogrpc + +import ray._private.gcs_utils as gcs_utils +from ray._common.utils import get_or_create_event_loop +from ray.core.generated import ( + gcs_pb2, + gcs_service_pb2, + gcs_service_pb2_grpc, + pubsub_pb2, +) + +logger = logging.getLogger(__name__) + + +class _SubscriberBase: + def __init__(self, worker_id: bytes = None): + self._worker_id = worker_id + # self._subscriber_id needs to match the binary format of a random + # SubscriberID / UniqueID, which is 28 (kUniqueIDSize) random bytes. + self._subscriber_id = bytes(bytearray(random.getrandbits(8) for _ in range(28))) + self._last_batch_size = 0 + self._max_processed_sequence_id = 0 + self._publisher_id = b"" + + # Batch size of the result from last poll. Used to indicate whether the + # subscriber can keep up. + @property + def last_batch_size(self): + return self._last_batch_size + + def _subscribe_request(self, channel): + cmd = pubsub_pb2.Command(channel_type=channel, subscribe_message={}) + req = gcs_service_pb2.GcsSubscriberCommandBatchRequest( + subscriber_id=self._subscriber_id, sender_id=self._worker_id, commands=[cmd] + ) + return req + + def _poll_request(self): + return gcs_service_pb2.GcsSubscriberPollRequest( + subscriber_id=self._subscriber_id, + max_processed_sequence_id=self._max_processed_sequence_id, + publisher_id=self._publisher_id, + ) + + def _unsubscribe_request(self, channels): + req = gcs_service_pb2.GcsSubscriberCommandBatchRequest( + subscriber_id=self._subscriber_id, sender_id=self._worker_id, commands=[] + ) + for channel in channels: + req.commands.append( + pubsub_pb2.Command(channel_type=channel, unsubscribe_message={}) + ) + return req + + @staticmethod + def _should_terminate_polling(e: grpc.RpcError) -> None: + # Caller only expects polling to be terminated after deadline exceeded. + if e.code() == grpc.StatusCode.DEADLINE_EXCEEDED: + return True + # Could be a temporary connection issue. Suppress error. + # TODO: reconnect GRPC channel? + if e.code() == grpc.StatusCode.UNAVAILABLE: + return True + return False + + +class _AioSubscriber(_SubscriberBase): + """Async io subscriber to GCS. + + Usage example common to Aio subscribers: + subscriber = GcsAioXxxSubscriber(address="...") + await subscriber.subscribe() + while running: + ...... = await subscriber.poll() + ...... + await subscriber.close() + """ + + def __init__( + self, + pubsub_channel_type, + worker_id: bytes = None, + address: str = None, + channel: aiogrpc.Channel = None, + ): + super().__init__(worker_id) + + if address: + assert channel is None, "address and channel cannot both be specified" + channel = gcs_utils.create_gcs_channel(address, aio=True) + else: + assert channel is not None, "One of address and channel must be specified" + # GRPC stub to GCS pubsub. + self._stub = gcs_service_pb2_grpc.InternalPubSubGcsServiceStub(channel) + + # Type of the channel. + self._channel = pubsub_channel_type + # A queue of received PubMessage. + self._queue = deque() + # Indicates whether the subscriber has closed. + self._close = asyncio.Event() + + async def subscribe(self) -> None: + """Registers a subscription for the subscriber's channel type. + + Before the registration, published messages in the channel will not be + saved for the subscriber. + """ + if self._close.is_set(): + return + req = self._subscribe_request(self._channel) + await self._stub.GcsSubscriberCommandBatch(req, timeout=30) + + async def _poll_call(self, req, timeout=None): + # Wrap GRPC _AioCall as a coroutine. + return await self._stub.GcsSubscriberPoll(req, timeout=timeout) + + async def _poll(self, timeout=None) -> None: + while len(self._queue) == 0: + req = self._poll_request() + poll = get_or_create_event_loop().create_task( + self._poll_call(req, timeout=timeout) + ) + close = get_or_create_event_loop().create_task(self._close.wait()) + done, others = await asyncio.wait( + [poll, close], timeout=timeout, return_when=asyncio.FIRST_COMPLETED + ) + # Cancel the other task if needed to prevent memory leak. + other_task = others.pop() + if not other_task.done(): + other_task.cancel() + if poll not in done or close in done: + # Request timed out or subscriber closed. + break + try: + self._last_batch_size = len(poll.result().pub_messages) + if poll.result().publisher_id != self._publisher_id: + if self._publisher_id != "": + logger.debug( + f"replied publisher_id {poll.result().publisher_id}" + f"different from {self._publisher_id}, this should " + "only happens during gcs failover." + ) + self._publisher_id = poll.result().publisher_id + self._max_processed_sequence_id = 0 + for msg in poll.result().pub_messages: + if msg.sequence_id <= self._max_processed_sequence_id: + logger.warning(f"Ignoring out of order message {msg}") + continue + self._max_processed_sequence_id = msg.sequence_id + self._queue.append(msg) + except grpc.RpcError as e: + if self._should_terminate_polling(e): + return + raise + + async def close(self) -> None: + """Closes the subscriber and its active subscription.""" + + # Mark close to terminate inflight polling and prevent future requests. + if self._close.is_set(): + return + self._close.set() + req = self._unsubscribe_request(channels=[self._channel]) + try: + await self._stub.GcsSubscriberCommandBatch(req, timeout=5) + except Exception: + pass + self._stub = None + + +class GcsAioResourceUsageSubscriber(_AioSubscriber): + def __init__( + self, + worker_id: bytes = None, + address: str = None, + channel: grpc.Channel = None, + ): + super().__init__( + pubsub_pb2.RAY_NODE_RESOURCE_USAGE_CHANNEL, worker_id, address, channel + ) + + async def poll(self, timeout=None) -> Tuple[bytes, str]: + """Polls for new resource usage message. + + Returns: + A tuple of string reporter ID and resource usage json string. + """ + await self._poll(timeout=timeout) + return self._pop_resource_usage(self._queue) + + @staticmethod + def _pop_resource_usage(queue): + if len(queue) == 0: + return None, None + msg = queue.popleft() + return msg.key_id.decode(), msg.node_resource_usage_message.json + + +class GcsAioActorSubscriber(_AioSubscriber): + def __init__( + self, + worker_id: bytes = None, + address: str = None, + channel: grpc.Channel = None, + ): + super().__init__(pubsub_pb2.GCS_ACTOR_CHANNEL, worker_id, address, channel) + + @property + def queue_size(self): + return len(self._queue) + + async def poll( + self, batch_size, timeout=None + ) -> List[Tuple[bytes, gcs_pb2.ActorTableData]]: + """Polls for new actor message. + + Returns: + A list of tuples of binary actor ID and actor table data. + """ + await self._poll(timeout=timeout) + return self._pop_actors(self._queue, batch_size=batch_size) + + @staticmethod + def _pop_actors(queue, batch_size): + if len(queue) == 0: + return [] + popped = 0 + msgs = [] + while len(queue) > 0 and popped < batch_size: + msg = queue.popleft() + msgs.append((msg.key_id, msg.actor_message)) + popped += 1 + return msgs + + +class GcsAioNodeInfoSubscriber(_AioSubscriber): + def __init__( + self, + worker_id: bytes = None, + address: str = None, + channel: grpc.Channel = None, + ): + super().__init__(pubsub_pb2.GCS_NODE_INFO_CHANNEL, worker_id, address, channel) + + async def poll( + self, batch_size, timeout=None + ) -> List[Tuple[bytes, gcs_pb2.GcsNodeInfo]]: + """Polls for new node info message. + + Returns: + A list of tuples of (node_id, GcsNodeInfo). + """ + await self._poll(timeout=timeout) + return self._pop_node_infos(self._queue, batch_size=batch_size) + + @staticmethod + def _pop_node_infos(queue, batch_size): + if len(queue) == 0: + return [] + popped = 0 + msgs = [] + while len(queue) > 0 and popped < batch_size: + msg = queue.popleft() + msgs.append((msg.key_id, msg.node_info_message)) + popped += 1 + return msgs diff --git a/lib/python3.12/site-packages/ray/_private/gcs_utils.py b/lib/python3.12/site-packages/ray/_private/gcs_utils.py new file mode 100644 index 0000000000000000000000000000000000000000..4590e7370a683ead6b6e18c4b91d508abb23c5d1 --- /dev/null +++ b/lib/python3.12/site-packages/ray/_private/gcs_utils.py @@ -0,0 +1,155 @@ +import logging +from typing import Optional + +from ray._private import ray_constants +from ray.core.generated.common_pb2 import ErrorType, JobConfig +from ray.core.generated.gcs_pb2 import ( + ActorTableData, + AvailableResources, + ErrorTableData, + GcsEntry, + GcsNodeInfo, + JobTableData, + PlacementGroupTableData, + PubSubMessage, + ResourceDemand, + ResourceLoad, + ResourcesData, + ResourceUsageBatchData, + TablePrefix, + TablePubsub, + TaskEvents, + TotalResources, + WorkerTableData, +) + +logger = logging.getLogger(__name__) + +__all__ = [ + "ActorTableData", + "GcsNodeInfo", + "AvailableResources", + "TotalResources", + "JobTableData", + "JobConfig", + "ErrorTableData", + "ErrorType", + "GcsEntry", + "ResourceUsageBatchData", + "ResourcesData", + "TablePrefix", + "TablePubsub", + "TaskEvents", + "ResourceDemand", + "ResourceLoad", + "PubSubMessage", + "WorkerTableData", + "PlacementGroupTableData", +] + + +WORKER = 0 +DRIVER = 1 + +# Cap messages at 512MB +_MAX_MESSAGE_LENGTH = 512 * 1024 * 1024 +# Send keepalive every 60s +_GRPC_KEEPALIVE_TIME_MS = 60 * 1000 +# Keepalive should be replied < 60s +_GRPC_KEEPALIVE_TIMEOUT_MS = 60 * 1000 + +# Also relying on these defaults: +# grpc.keepalive_permit_without_calls=0: No keepalive without inflight calls. +# grpc.use_local_subchannel_pool=0: Subchannels are shared. +_GRPC_OPTIONS = [ + *ray_constants.GLOBAL_GRPC_OPTIONS, + ("grpc.max_send_message_length", _MAX_MESSAGE_LENGTH), + ("grpc.max_receive_message_length", _MAX_MESSAGE_LENGTH), + ("grpc.keepalive_time_ms", _GRPC_KEEPALIVE_TIME_MS), + ("grpc.keepalive_timeout_ms", _GRPC_KEEPALIVE_TIMEOUT_MS), +] + + +def create_gcs_channel(address: str, aio=False): + """Returns a GRPC channel to GCS. + + Args: + address: GCS address string, e.g. ip:port + aio: Whether using grpc.aio + Returns: + grpc.Channel or grpc.aio.Channel to GCS + """ + from ray._private.grpc_utils import init_grpc_channel + + return init_grpc_channel(address, options=_GRPC_OPTIONS, asynchronous=aio) + + +class GcsChannel: + def __init__(self, gcs_address: Optional[str] = None, aio: bool = False): + self._gcs_address = gcs_address + self._aio = aio + + @property + def address(self): + return self._gcs_address + + def connect(self): + # GCS server uses a cached port, so it should use the same port after + # restarting. This means GCS address should stay the same for the + # lifetime of the Ray cluster. + self._channel = create_gcs_channel(self._gcs_address, self._aio) + + def channel(self): + return self._channel + + +def cleanup_redis_storage( + host: str, + port: int, + password: str, + use_ssl: bool, + storage_namespace: str, + username: Optional[str] = None, +): + """This function is used to cleanup the GCS storage in Redis. + It supports Redis in cluster and non-cluster modes. + + Args: + host: The Redis host address. + port: The Redis port. + username: The Redis username. + password: The Redis password. + use_ssl: Whether to encrypt the connection. + storage_namespace: The namespace of the storage to be deleted. + """ + + from ray._raylet import del_key_prefix_from_storage # type: ignore + + if not isinstance(host, str): + raise ValueError("Host must be a string") + + if username is None: + username = "" + + if not isinstance(username, str): + raise ValueError("Username must be a string") + + if not isinstance(password, str): + raise ValueError("Password must be a string") + + if port < 0: + raise ValueError(f"Invalid port: {port}") + + if not isinstance(use_ssl, bool): + raise TypeError("use_ssl must be a boolean") + + if not isinstance(storage_namespace, str): + raise ValueError("storage namespace must be a string") + + # Right now, GCS stores all data in multiple hashes with keys prefixed by + # storage_namespace. So we only need to delete the specific key prefix to cleanup + # the cluster's data. + # Note this deletes all keys with prefix `RAY{key_prefix}@`, not `{key_prefix}`. + return del_key_prefix_from_storage( + host, port, username, password, use_ssl, storage_namespace + ) diff --git a/lib/python3.12/site-packages/ray/_private/grpc_utils.py b/lib/python3.12/site-packages/ray/_private/grpc_utils.py new file mode 100644 index 0000000000000000000000000000000000000000..be5a90fee01ab26a8e81c50027b25c04d46f93a9 --- /dev/null +++ b/lib/python3.12/site-packages/ray/_private/grpc_utils.py @@ -0,0 +1,155 @@ +import os +from concurrent import futures +from typing import Any, Optional, Sequence, Tuple + +import grpc +from grpc import aio as aiogrpc + +import ray +from ray._private.authentication import authentication_utils +from ray._private.tls_utils import load_certs_from_env + + +def init_grpc_channel( + address: str, + options: Optional[Sequence[Tuple[str, Any]]] = None, + asynchronous: bool = False, + credentials: Optional[grpc.ChannelCredentials] = None, +): + """Create a gRPC channel with authentication interceptors if token auth is enabled. + + This function handles: + - TLS configuration via RAY_USE_TLS environment variable or custom credentials + - Authentication interceptors when token auth is enabled + - Keepalive settings from Ray config + - Both synchronous and asynchronous channels + + Args: + address: The gRPC server address (host:port) + options: Optional gRPC channel options as sequence of (key, value) tuples + asynchronous: If True, create async channel; otherwise sync + credentials: Optional custom gRPC credentials for TLS. If provided, takes + precedence over RAY_USE_TLS environment variable. + + Returns: + grpc.Channel or grpc.aio.Channel: Configured gRPC channel with interceptors + """ + grpc_module = aiogrpc if asynchronous else grpc + + options = options or [] + options_dict = dict(options) + options_dict["grpc.keepalive_time_ms"] = options_dict.get( + "grpc.keepalive_time_ms", ray._config.grpc_client_keepalive_time_ms() + ) + options_dict["grpc.keepalive_timeout_ms"] = options_dict.get( + "grpc.keepalive_timeout_ms", ray._config.grpc_client_keepalive_timeout_ms() + ) + options = options_dict.items() + + # Build interceptors list + interceptors = [] + if authentication_utils.is_token_auth_enabled(): + from ray._private.authentication.grpc_authentication_client_interceptor import ( + AsyncAuthenticationMetadataClientInterceptor, + SyncAuthenticationMetadataClientInterceptor, + ) + + if asynchronous: + interceptors.append(AsyncAuthenticationMetadataClientInterceptor()) + else: + interceptors.append(SyncAuthenticationMetadataClientInterceptor()) + + # Determine channel type and credentials + if credentials is not None: + # Use provided custom credentials (takes precedence) + channel_creator = grpc_module.secure_channel + base_args = (address, credentials) + elif os.environ.get("RAY_USE_TLS", "0").lower() in ("1", "true"): + # Use TLS from environment variables + server_cert_chain, private_key, ca_cert = load_certs_from_env() + tls_credentials = grpc.ssl_channel_credentials( + certificate_chain=server_cert_chain, + private_key=private_key, + root_certificates=ca_cert, + ) + channel_creator = grpc_module.secure_channel + base_args = (address, tls_credentials) + else: + # Insecure channel + channel_creator = grpc_module.insecure_channel + base_args = (address,) + + # Create channel (async channels get interceptors in constructor, sync via intercept_channel) + if asynchronous: + channel = channel_creator( + *base_args, options=options, interceptors=interceptors + ) + else: + channel = channel_creator(*base_args, options=options) + if interceptors: + channel = grpc.intercept_channel(channel, *interceptors) + + return channel + + +def create_grpc_server_with_interceptors( + max_workers: Optional[int] = None, + thread_name_prefix: str = "grpc_server", + options: Optional[Sequence[Tuple[str, Any]]] = None, + asynchronous: bool = False, +): + """Create a gRPC server with authentication interceptors if token auth is enabled. + + This function handles: + - Authentication interceptors when token auth is enabled + - Both synchronous and asynchronous servers + - Thread pool configuration for sync servers + + Args: + max_workers: Max thread pool workers (required for sync, ignored for async) + thread_name_prefix: Thread name prefix for sync thread pool + options: Optional gRPC server options as sequence of (key, value) tuples + asynchronous: If True, create async server; otherwise sync + + Returns: + grpc.Server or grpc.aio.Server: Configured gRPC server with interceptors + """ + grpc_module = aiogrpc if asynchronous else grpc + + # Build interceptors list + interceptors = [] + if authentication_utils.is_token_auth_enabled(): + if asynchronous: + from ray._private.authentication.grpc_authentication_server_interceptor import ( + AsyncAuthenticationServerInterceptor, + ) + + interceptors.append(AsyncAuthenticationServerInterceptor()) + else: + from ray._private.authentication.grpc_authentication_server_interceptor import ( + SyncAuthenticationServerInterceptor, + ) + + interceptors.append(SyncAuthenticationServerInterceptor()) + + # Create server + if asynchronous: + server = grpc_module.server( + interceptors=interceptors if interceptors else None, + options=options, + ) + else: + if max_workers is None: + raise ValueError("max_workers is required for synchronous gRPC servers") + + executor = futures.ThreadPoolExecutor( + max_workers=max_workers, + thread_name_prefix=thread_name_prefix, + ) + server = grpc_module.server( + executor, + interceptors=interceptors if interceptors else None, + options=options, + ) + + return server diff --git a/lib/python3.12/site-packages/ray/_private/internal_api.py b/lib/python3.12/site-packages/ray/_private/internal_api.py new file mode 100644 index 0000000000000000000000000000000000000000..7a55732169887a0a53257a1aa752398f5c63dfbd --- /dev/null +++ b/lib/python3.12/site-packages/ray/_private/internal_api.py @@ -0,0 +1,264 @@ +import warnings +from typing import List, Tuple + +import ray +import ray._private.profiling as profiling +import ray._private.services as services +import ray._private.worker +from ray._common.network_utils import build_address +from ray._private.state import GlobalState +from ray._raylet import GcsClientOptions +from ray.core.generated import common_pb2 + +__all__ = ["free", "global_gc"] +MAX_MESSAGE_LENGTH = ray._config.max_grpc_message_size() + + +def global_gc(): + """Trigger gc.collect() on all workers in the cluster.""" + + worker = ray._private.worker.global_worker + worker.core_worker.global_gc() + + +def get_state_from_address(address=None): + address = services.canonicalize_bootstrap_address_or_die(address) + + state = GlobalState() + options = GcsClientOptions.create( + address, None, allow_cluster_id_nil=True, fetch_cluster_id_if_nil=False + ) + state._initialize_global_state(options) + return state + + +def memory_summary( + address=None, + group_by="NODE_ADDRESS", + sort_by="OBJECT_SIZE", + units="B", + line_wrap=True, + stats_only=False, + num_entries=None, +): + from ray.dashboard.memory_utils import memory_summary + + state = get_state_from_address(address) + reply = get_memory_info_reply(state) + + if stats_only: + return store_stats_summary(reply) + return memory_summary( + state, group_by, sort_by, line_wrap, units, num_entries + ) + store_stats_summary(reply) + + +def get_memory_info_reply(state, node_manager_address=None, node_manager_port=None): + """Returns global memory info.""" + + from ray._private.grpc_utils import init_grpc_channel + from ray.core.generated import node_manager_pb2, node_manager_pb2_grpc + + # We can ask any Raylet for the global memory info, that Raylet internally + # asks all nodes in the cluster for memory stats. + if node_manager_address is None or node_manager_port is None: + # We should ask for a raylet that is alive. + raylet = None + for node in state.node_table(): + if node["Alive"]: + raylet = node + break + assert raylet is not None, "Every raylet is dead" + raylet_address = build_address( + raylet["NodeManagerAddress"], raylet["NodeManagerPort"] + ) + else: + raylet_address = build_address(node_manager_address, node_manager_port) + + channel = init_grpc_channel( + raylet_address, + options=[ + ("grpc.max_send_message_length", MAX_MESSAGE_LENGTH), + ("grpc.max_receive_message_length", MAX_MESSAGE_LENGTH), + ], + ) + + stub = node_manager_pb2_grpc.NodeManagerServiceStub(channel) + reply = stub.FormatGlobalMemoryInfo( + node_manager_pb2.FormatGlobalMemoryInfoRequest(include_memory_info=False), + timeout=60.0, + ) + return reply + + +def node_stats( + node_manager_address=None, node_manager_port=None, include_memory_info=True +): + """Returns NodeStats object describing memory usage in the cluster.""" + + from ray._private.grpc_utils import init_grpc_channel + from ray.core.generated import node_manager_pb2, node_manager_pb2_grpc + + # We can ask any Raylet for the global memory info. + assert node_manager_address is not None and node_manager_port is not None + raylet_address = build_address(node_manager_address, node_manager_port) + channel = init_grpc_channel( + raylet_address, + options=[ + ("grpc.max_send_message_length", MAX_MESSAGE_LENGTH), + ("grpc.max_receive_message_length", MAX_MESSAGE_LENGTH), + ], + ) + + stub = node_manager_pb2_grpc.NodeManagerServiceStub(channel) + node_stats = stub.GetNodeStats( + node_manager_pb2.GetNodeStatsRequest(include_memory_info=include_memory_info), + timeout=30.0, + ) + return node_stats + + +def store_stats_summary(reply): + """Returns formatted string describing object store stats in all nodes.""" + store_summary = "--- Aggregate object store stats across all nodes ---\n" + # TODO(ekl) it would be nice if we could provide a full memory usage + # breakdown by type (e.g., pinned by worker, primary, etc.) + store_summary += ( + "Plasma memory usage {} MiB, {} objects, {}% full, {}% " + "needed\n".format( + int(reply.store_stats.object_store_bytes_used / (1024 * 1024)), + reply.store_stats.num_local_objects, + round( + 100 + * reply.store_stats.object_store_bytes_used + / reply.store_stats.object_store_bytes_avail, + 2, + ), + round( + 100 + * reply.store_stats.object_store_bytes_primary_copy + / reply.store_stats.object_store_bytes_avail, + 2, + ), + ) + ) + if reply.store_stats.object_store_bytes_fallback > 0: + store_summary += "Plasma filesystem mmap usage: {} MiB\n".format( + int(reply.store_stats.object_store_bytes_fallback / (1024 * 1024)) + ) + if reply.store_stats.spill_time_total_s > 0: + store_summary += ( + "Spilled {} MiB, {} objects, avg write throughput {} MiB/s\n".format( + int(reply.store_stats.spilled_bytes_total / (1024 * 1024)), + reply.store_stats.spilled_objects_total, + int( + reply.store_stats.spilled_bytes_total + / (1024 * 1024) + / reply.store_stats.spill_time_total_s + ), + ) + ) + if reply.store_stats.restore_time_total_s > 0: + store_summary += ( + "Restored {} MiB, {} objects, avg read throughput {} MiB/s\n".format( + int(reply.store_stats.restored_bytes_total / (1024 * 1024)), + reply.store_stats.restored_objects_total, + int( + reply.store_stats.restored_bytes_total + / (1024 * 1024) + / reply.store_stats.restore_time_total_s + ), + ) + ) + if reply.store_stats.object_pulls_queued: + store_summary += "Object fetches queued, waiting for available memory." + + return store_summary + + +def free(object_refs: list, local_only: bool = False): + """ + DeprecationWarning: `free` is a deprecated API and will be + removed in a future version of Ray. If you have a use case + for this API, please open an issue on GitHub. + + Free a list of IDs from the in-process and plasma object stores. + + This function is a low-level API which should be used in restricted + scenarios. + + If local_only is false, the request will be send to all object stores. + + This method will not return any value to indicate whether the deletion is + successful or not. This function is an instruction to the object store. If + some of the objects are in use, the object stores will delete them later + when the ref count is down to 0. + + Examples: + + .. testcode:: + + import ray + + @ray.remote + def f(): + return 0 + + obj_ref = f.remote() + ray.get(obj_ref) # wait for object to be created first + free([obj_ref]) # unpin & delete object globally + + Args: + object_refs (List[ObjectRef]): List of object refs to delete. + local_only: Whether only deleting the list of objects in local + object store or all object stores. + """ + warnings.warn( + "`free` is a deprecated API and will be removed in a future version of Ray. " + "If you have a use case for this API, please open an issue on GitHub.", + DeprecationWarning, + ) + worker = ray._private.worker.global_worker + + if isinstance(object_refs, ray.ObjectRef): + object_refs = [object_refs] + + if not isinstance(object_refs, list): + raise TypeError( + "free() expects a list of ObjectRef, got {}".format(type(object_refs)) + ) + + # Make sure that the values are object refs. + for object_ref in object_refs: + if not isinstance(object_ref, ray.ObjectRef): + raise TypeError( + "Attempting to call `free` on the value {}, " + "which is not an ray.ObjectRef.".format(object_ref) + ) + + worker.check_connected() + with profiling.profile("ray.free"): + if len(object_refs) == 0: + return + + worker.core_worker.free_objects(object_refs, local_only) + + +def get_local_ongoing_lineage_reconstruction_tasks() -> List[ + Tuple[common_pb2.LineageReconstructionTask, int] +]: + """Return the locally submitted ongoing retry tasks + triggered by lineage reconstruction. + + NOTE: for the lineage reconstruction task status, + this method only returns the status known to the submitter + (i.e. it returns SUBMITTED_TO_WORKER instead of RUNNING). + + The return type is a list of pairs where pair.first is the + lineage reconstruction task info and pair.second is the number + of ongoing lineage reconstruction tasks of this type. + """ + + worker = ray._private.worker.global_worker + worker.check_connected() + return worker.core_worker.get_local_ongoing_lineage_reconstruction_tasks() diff --git a/lib/python3.12/site-packages/ray/_private/label_utils.py b/lib/python3.12/site-packages/ray/_private/label_utils.py new file mode 100644 index 0000000000000000000000000000000000000000..6f6de45976244b7ac6e86a1550d6224c87ef3b46 --- /dev/null +++ b/lib/python3.12/site-packages/ray/_private/label_utils.py @@ -0,0 +1,230 @@ +import json +import re +from typing import ( + Any, + Dict, + List, + Optional, +) + +import yaml + +import ray._private.ray_constants as ray_constants + +# Regex patterns used to validate that labels conform to Kubernetes label syntax rules. +# https://kubernetes.io/docs/concepts/overview/working-with-objects/labels/#syntax-and-character-set + +# Regex for mandatory name (DNS label) or value +# Examples: +# Valid matches: "a", "label-name", "a-._b", "123", "this_is_a_valid_label" +# Invalid matches: "-abc", "abc-", "my@label" +LABEL_REGEX = re.compile(r"([a-zA-Z0-9]([a-zA-Z0-9_.-]{0,61}[a-zA-Z0-9])?)") + +# Regex for optional prefix (DNS subdomain) +# Examples: +# Valid matches: "abc", "sub.domain.example", "my-label", "123.456.789" +# Invalid matches: "-abc", "prefix_", "sub..domain", sub.$$.example +LABEL_PREFIX_REGEX = rf"^({LABEL_REGEX.pattern}?(\.{LABEL_REGEX.pattern}?)*)$" + +# Supported operators for label selector conditions. Not (!) conditions are handled separately. +LABEL_OPERATORS = {"in"} +# Create a pattern string dynamically based on the LABEL_OPERATORS +OPERATOR_PATTERN = "|".join([re.escape(operator) for operator in LABEL_OPERATORS]) + +# Regex to match valid label selector operators and values +# Examples: +# Valid matches: "spot", "!GPU", "213521", "in(A123, B456, C789)", "!in(spot, on-demand)", "valid-value" +# Invalid matches: "-spot", "spot_", "in()", "in(spot,", "in(H100, TPU!GPU)", "!!!in(H100, TPU)" +LABEL_SELECTOR_REGEX = re.compile( + rf"^!?(?:{OPERATOR_PATTERN})?\({LABEL_REGEX.pattern}(?:, ?{LABEL_REGEX.pattern})*\)$|^!?{LABEL_REGEX.pattern}$" +) + + +def parse_node_labels_json(labels_json: str) -> Dict[str, str]: + labels = json.loads(labels_json) + if not isinstance(labels, dict): + raise ValueError("The format after deserialization is not a key-value pair map") + for key, value in labels.items(): + if not isinstance(key, str): + raise ValueError("The key is not string type.") + if not isinstance(value, str): + raise ValueError(f'The value of the "{key}" is not string type') + + # Validate parsed custom node labels don't begin with ray.io prefix + validate_node_labels(labels) + + return labels + + +def parse_node_labels_string(labels_str: str) -> Dict[str, str]: + labels = {} + + # Remove surrounding quotes if they exist + if len(labels_str) > 1 and labels_str.startswith('"') and labels_str.endswith('"'): + labels_str = labels_str[1:-1] + + if labels_str == "": + return labels + + # Labels argument should consist of a string of key=value pairs + # separated by commas. Labels follow Kubernetes label syntax. + label_pairs = labels_str.split(",") + for pair in label_pairs: + # Split each pair by `=` + key_value = pair.split("=") + if len(key_value) != 2: + raise ValueError("Label string is not a key-value pair.") + key = key_value[0].strip() + value = key_value[1].strip() + labels[key] = value + + # Validate parsed node labels follow expected Kubernetes label syntax + validate_node_label_syntax(labels) + + return labels + + +def parse_node_labels_from_yaml_file(path: str) -> Dict[str, str]: + if path == "": + return {} + with open(path, "r") as file: + # Expects valid YAML content + labels = yaml.safe_load(file) + if not isinstance(labels, dict): + raise ValueError( + "The format after deserialization is not a key-value pair map." + ) + for key, value in labels.items(): + if not isinstance(key, str): + raise ValueError("The key is not string type.") + if not isinstance(value, str): + raise ValueError(f'The value of "{key}" is not string type.') + + # Validate parsed node labels follow expected Kubernetes label syntax + validate_node_label_syntax(labels) + + return labels + + +# TODO (ryanaoleary@): This function will be removed after the migration to the label +# selector API from NodeLabelSchedulingPolicy is complete. +def validate_node_labels(labels: Dict[str, str]): + if labels is None: + return + for key in labels.keys(): + if key.startswith(ray_constants.RAY_DEFAULT_LABEL_KEYS_PREFIX): + raise ValueError( + f"Custom label keys `{key}` cannot start with the prefix " + f"`{ray_constants.RAY_DEFAULT_LABEL_KEYS_PREFIX}`. " + f"This is reserved for Ray defined labels." + ) + + +def validate_label_key(key: str) -> Optional[str]: + if "/" in key: + prefix, name = key.rsplit("/", 1) + if len(prefix) > 253 or not re.fullmatch(LABEL_PREFIX_REGEX, prefix): + return str( + f"Invalid label key prefix `{prefix}`. Prefix must be a series of DNS labels " + f"separated by dots (.), not longer than 253 characters in total." + ) + else: + name = key + if len(name) > 63 or not re.fullmatch(LABEL_REGEX, name): + return str( + f"Invalid label key name `{name}`. Name must be 63 chars or less beginning and ending " + f"with an alphanumeric character ([a-z0-9A-Z]) with dashes (-), underscores (_)," + f"dots (.), and alphanumerics between." + ) + return None + + +def validate_label_value(value: str): + if value == "": + return + if len(value) > 63 or not re.fullmatch(LABEL_REGEX, value): + raise ValueError( + f"Invalid label key value `{value}`. Value must be 63 chars or less beginning and ending " + f"with an alphanumeric character ([a-z0-9A-Z]) with dashes (-), underscores (_)," + f"dots (.), and alphanumerics between." + ) + + +def validate_label_selector(label_selector: Optional[Dict[str, str]]) -> Optional[str]: + if label_selector is None: + return None + + for key, value in label_selector.items(): + possible_error_message = validate_label_key(key) + if possible_error_message: + return possible_error_message + if value is not None: + possible_error_message = validate_label_selector_value(value) + if possible_error_message: + return possible_error_message + + return None + + +def validate_label_selector_value(selector: str) -> Optional[str]: + if selector == "": + return None + if not re.fullmatch(LABEL_SELECTOR_REGEX, selector): + return str( + f"Invalid label selector value `{selector}`. The label selector value should contain optional operators and a label value. Supported operators are: ! and {LABEL_OPERATORS}. " + f"Value must be 63 chars or less beginning and ending " + f"with an alphanumeric character ([a-z0-9A-Z]) with dashes (-), underscores (_)," + f"dots (.), and alphanumerics between." + ) + + return None + + +# TODO (ryanaoleary@): This function will replace `validate_node_labels` after +# the migration from NodeLabelSchedulingPolicy to the Label Selector API is complete. +def validate_node_label_syntax(labels: Dict[str, str]): + if labels is None: + return + for key, value in labels.items(): + possible_error_message = validate_label_key(key) + if possible_error_message: + raise ValueError(possible_error_message) + if value is not None: + validate_label_value(value) + + +def validate_fallback_strategy( + fallback_strategy: Optional[List[Dict[str, Any]]] +) -> Optional[str]: + if fallback_strategy is None: + return None + + # Supported options for `fallback_strategy` scheduling. + supported_options = {"label_selector"} + + for strategy in fallback_strategy: + if not isinstance(strategy, dict): + return "Each element in fallback_strategy must be a dictionary." + + if not strategy: + return "Empty dictionary found in `fallback_strategy`." + + # Validate `fallback_strategy` only contains supported options. + for option in strategy: + if option not in supported_options: + return ( + f"Unsupported option found: '{option}'. " + f"Only {list(supported_options)} is currently supported." + ) + + # Validate the 'label_selector' dictionary. + label_selector = strategy.get("label_selector") + if label_selector: + if not isinstance(label_selector, dict): + return 'The value of "label_selector" must be a dictionary.' + + error_message = validate_label_selector(label_selector) + if error_message: + return error_message + + return None diff --git a/lib/python3.12/site-packages/ray/_private/log.py b/lib/python3.12/site-packages/ray/_private/log.py new file mode 100644 index 0000000000000000000000000000000000000000..074960c336211a9c6305b3d3b07da2434f180989 --- /dev/null +++ b/lib/python3.12/site-packages/ray/_private/log.py @@ -0,0 +1,150 @@ +import logging +import threading +import time +from typing import Optional, Union + +INTERNAL_TIMESTAMP_LOG_KEY = "_ray_timestamp_ns" + + +def _print_loggers(): + """Print a formatted list of loggers and their handlers for debugging.""" + loggers = {logging.root.name: logging.root} + loggers.update(dict(sorted(logging.root.manager.loggerDict.items()))) + for name, logger in loggers.items(): + if isinstance(logger, logging.Logger): + print(f" {name}: disabled={logger.disabled}, propagate={logger.propagate}") + for handler in logger.handlers: + print(f" {handler}") + + +def clear_logger(logger: Union[str, logging.Logger]): + """Reset a logger, clearing its handlers and enabling propagation. + + Args: + logger: Logger to be cleared + """ + if isinstance(logger, str): + logger = logging.getLogger(logger) + logger.propagate = True + logger.handlers.clear() + + +class PlainRayHandler(logging.StreamHandler): + """A plain log handler. + + This handler writes to whatever sys.stderr points to at emit-time, + not at instantiation time. See docs for logging._StderrHandler. + """ + + def __init__(self): + super().__init__() + self.plain_handler = logging._StderrHandler() + self.plain_handler.level = self.level + self.plain_handler.formatter = logging.Formatter(fmt="%(message)s") + + def emit(self, record: logging.LogRecord): + """Emit the log message. + + If this is a worker, bypass fancy logging and just emit the log record. + If this is the driver, emit the message using the appropriate console handler. + + Args: + record: Log record to be emitted + """ + import ray + + if ( + hasattr(ray, "_private") + and hasattr(ray._private, "worker") + and ray._private.worker.global_worker.mode + == ray._private.worker.WORKER_MODE + ): + self.plain_handler.emit(record) + else: + logging._StderrHandler.emit(self, record) + + +logger_initialized = False +logging_config_lock = threading.Lock() + + +def _setup_log_record_factory(): + """Setup log record factory to add _ray_timestamp_ns to LogRecord.""" + old_factory = logging.getLogRecordFactory() + + def record_factory(*args, **kwargs): + record = old_factory(*args, **kwargs) + # Python logging module starts to use `time.time_ns()` to generate `created` + # from Python 3.13 to avoid the precision loss caused by the float type. + # Here, we generate the `created` for the LogRecord to support older Python + # versions. + ct = time.time_ns() + record.created = ct / 1e9 + + record.__dict__[INTERNAL_TIMESTAMP_LOG_KEY] = ct + + return record + + logging.setLogRecordFactory(record_factory) + + +def generate_logging_config(): + """Generate the default Ray logging configuration.""" + with logging_config_lock: + global logger_initialized + if logger_initialized: + return + logger_initialized = True + + plain_formatter = logging.Formatter( + "%(asctime)s\t%(levelname)s %(filename)s:%(lineno)s -- %(message)s" + ) + + default_handler = PlainRayHandler() + default_handler.setFormatter(plain_formatter) + + ray_logger = logging.getLogger("ray") + ray_logger.setLevel(logging.INFO) + ray_logger.addHandler(default_handler) + ray_logger.propagate = False + + # Special handling for ray.rllib: only warning-level messages passed through + # See https://github.com/ray-project/ray/pull/31858 for related PR + rllib_logger = logging.getLogger("ray.rllib") + rllib_logger.setLevel(logging.WARN) + + # Set up the LogRecord factory. + _setup_log_record_factory() + + +def setup_process_exit_logger( + process_exit_log_path: str, + level: int = logging.INFO, + formatter: Optional[logging.Formatter] = None, +) -> logging.Logger: + """Configure and return the 'ray.process_exit' logger with a FileHandler.""" + logger = logging.getLogger("ray.process_exit") + logger.setLevel(level) + logger.propagate = False + + fh = logging.FileHandler(process_exit_log_path, encoding="utf-8") + if formatter is None: + formatter = logging.Formatter( + "%(asctime)s\t%(levelname)s %(filename)s:%(lineno)s -- %(message)s" + ) + fh.setFormatter(formatter) + logger.addHandler(fh) + return logger + + +def format_returncode(rc: Optional[int]) -> str: + """Return a consistent string for process return code.""" + if rc is None: + return "None" + try: + rc_int = int(rc) + except Exception: + return str(rc) + if rc_int < 0: + return f"{rc_int} (signal {-rc_int})" + return f"{rc_int}" diff --git a/lib/python3.12/site-packages/ray/_private/log_monitor.py b/lib/python3.12/site-packages/ray/_private/log_monitor.py new file mode 100644 index 0000000000000000000000000000000000000000..5a96a3b5b6a1def5f3b310ebf9d867767deadd02 --- /dev/null +++ b/lib/python3.12/site-packages/ray/_private/log_monitor.py @@ -0,0 +1,624 @@ +import argparse +import errno +import glob +import logging +import logging.handlers +import os +import platform +import re +import shutil +import sys +import time +import traceback +from typing import Callable, List, Optional, Set + +import ray._private.ray_constants as ray_constants +import ray._private.services as services +import ray._private.utils +from ray._private import logging_utils +from ray._private.ray_logging import setup_component_logger +from ray._raylet import GcsClient + +# Logger for this module. It should be configured at the entry point +# into the program using Ray. Ray provides a default configuration at +# entry/init points. +logger = logging.getLogger(__name__) + +# The groups are job id, and pid. +WORKER_LOG_PATTERN = re.compile(r".*worker.*-([0-9a-f]+)-(\d+)") +# The groups are job id. +RUNTIME_ENV_SETUP_PATTERN = re.compile(r".*runtime_env_setup-(\d+).log") +# Log name update interval under pressure. +# We need it because log name update is CPU intensive and uses 100% +# of cpu when there are many log files. +LOG_NAME_UPDATE_INTERVAL_S = float(os.getenv("LOG_NAME_UPDATE_INTERVAL_S", 0.5)) +# Once there are more files than this threshold, +# log monitor start giving backpressure to lower cpu usages. +RAY_LOG_MONITOR_MANY_FILES_THRESHOLD = int( + os.getenv("RAY_LOG_MONITOR_MANY_FILES_THRESHOLD", 1000) +) +RAY_RUNTIME_ENV_LOG_TO_DRIVER_ENABLED = int( + os.getenv("RAY_RUNTIME_ENV_LOG_TO_DRIVER_ENABLED", 0) +) + + +class LogFileInfo: + def __init__( + self, + filename=None, + size_when_last_opened=None, + file_position=None, + file_handle=None, + is_err_file=False, + job_id=None, + worker_pid=None, + ): + assert ( + filename is not None + and size_when_last_opened is not None + and file_position is not None + ) + self.filename = filename + self.size_when_last_opened = size_when_last_opened + self.file_position = file_position + self.file_handle = file_handle + self.is_err_file = is_err_file + self.job_id = job_id + self.worker_pid = worker_pid + self.actor_name = None + self.task_name = None + + def reopen_if_necessary(self): + """Check if the file's inode has changed and reopen it if necessary. + + There are a variety of reasons what we would logically consider a file + would have different inodes, such as log rotation or file syncing + semantics. + + If the file is smaller than our recorded file position, we assume it has been + rotated and start reading it from the beginning. + """ + try: + open_inode = None + if self.file_handle and not self.file_handle.closed: + open_inode = os.fstat(self.file_handle.fileno()).st_ino + + new_statinfo = os.stat(self.filename) + if new_statinfo.st_ino != open_inode: + self.file_handle = open(self.filename, "rb") + + # If the new file is smaller than the last read position, assume that + # the file has been rotated and read from the beginning. Else, continue + # from the existing file position. + if new_statinfo.st_size < self.file_position: + self.file_position = 0 + + self.file_handle.seek(self.file_position) + self.size_when_last_opened = new_statinfo.st_size + except Exception: + logger.debug(f"file no longer exists, skip re-opening of {self.filename}") + + def __repr__(self): + return ( + "FileInfo(\n" + f"\tfilename: {self.filename}\n" + f"\tsize_when_last_opened: {self.size_when_last_opened}\n" + f"\tfile_position: {self.file_position}\n" + f"\tfile_handle: {self.file_handle}\n" + f"\tis_err_file: {self.is_err_file}\n" + f"\tjob_id: {self.job_id}\n" + f"\tworker_pid: {self.worker_pid}\n" + f"\tactor_name: {self.actor_name}\n" + f"\ttask_name: {self.task_name}\n" + ")" + ) + + +class LogMonitor: + """A monitor process for monitoring Ray log files. + + This class maintains a list of open files and a list of closed log files. We + can't simply leave all files open because we'll run out of file + descriptors. + + The "run" method of this class will cycle between doing several things: + 1. First, it will check if any new files have appeared in the log + directory. If so, they will be added to the list of closed files. + 2. Then, if we are unable to open any new files, we will close all of the + files. + 3. Then, we will open as many closed files as we can that may have new + lines (judged by an increase in file size since the last time the file + was opened). + 4. Then we will loop through the open files and see if there are any new + lines in the file. If so, we will publish them to Ray pubsub. + + Attributes: + ip: The hostname of this machine, for grouping log messages. + logs_dir: The directory that the log files are in. + log_filenames: This is the set of filenames of all files in + open_file_infos and closed_file_infos. + open_file_infos (list[LogFileInfo]): Info for all of the open files. + closed_file_infos (list[LogFileInfo]): Info for all of the closed + files. + can_open_more_files: True if we can still open more files and + false otherwise. + max_files_open: The maximum number of files that can be open. + """ + + def __init__( + self, + node_ip_address: str, + logs_dir: str, + gcs_client: GcsClient, + is_proc_alive_fn: Callable[[int], bool], + max_files_open: int = ray_constants.LOG_MONITOR_MAX_OPEN_FILES, + gcs_address: Optional[str] = None, + ): + """Initialize the log monitor object.""" + self.ip: str = node_ip_address + self.logs_dir: str = logs_dir + self.gcs_client = gcs_client + self.log_filenames: Set[str] = set() + self.open_file_infos: List[LogFileInfo] = [] + self.closed_file_infos: List[LogFileInfo] = [] + self.can_open_more_files: bool = True + self.max_files_open: int = max_files_open + self.is_proc_alive_fn: Callable[[int], bool] = is_proc_alive_fn + self.is_autoscaler_v2: bool = self.get_is_autoscaler_v2(gcs_address) + + logger.info( + f"Starting log monitor with [max open files={max_files_open}]," + f" [is_autoscaler_v2={self.is_autoscaler_v2}]" + ) + + def get_is_autoscaler_v2(self, gcs_address: Optional[str]) -> bool: + """Check if autoscaler v2 is enabled.""" + if gcs_address is None: + return False + + if not ray.experimental.internal_kv._internal_kv_initialized(): + ray.experimental.internal_kv._initialize_internal_kv(self.gcs_client) + from ray.autoscaler.v2.utils import is_autoscaler_v2 + + return is_autoscaler_v2() + + def _close_all_files(self): + """Close all open files (so that we can open more).""" + while len(self.open_file_infos) > 0: + file_info = self.open_file_infos.pop(0) + file_info.file_handle.close() + file_info.file_handle = None + + proc_alive = True + # Test if the worker process that generated the log file + # is still alive. Only applies to worker processes. + # For all other system components, we always assume they are alive. + if ( + file_info.worker_pid != "raylet" + and file_info.worker_pid != "gcs_server" + and file_info.worker_pid != "autoscaler" + and file_info.worker_pid != "runtime_env" + and file_info.worker_pid is not None + ): + assert not isinstance(file_info.worker_pid, str), ( + "PID should be an int type. " f"Given PID: {file_info.worker_pid}." + ) + proc_alive = self.is_proc_alive_fn(file_info.worker_pid) + if not proc_alive: + # The process is not alive any more, so move the log file + # out of the log directory so glob.glob will not be slowed + # by it. + target = os.path.join( + self.logs_dir, "old", os.path.basename(file_info.filename) + ) + try: + shutil.move(file_info.filename, target) + except (IOError, OSError) as e: + if e.errno == errno.ENOENT: + logger.warning( + f"Warning: The file {file_info.filename} was not found." + ) + else: + raise e + + if proc_alive: + self.closed_file_infos.append(file_info) + + self.can_open_more_files = True + + def update_log_filenames(self): + """Update the list of log files to monitor.""" + monitor_log_paths = [] + # output of user code is written here + monitor_log_paths += glob.glob( + f"{self.logs_dir}/worker*[.out|.err]" + ) + glob.glob(f"{self.logs_dir}/java-worker*.log") + # segfaults and other serious errors are logged here + monitor_log_paths += glob.glob(f"{self.logs_dir}/raylet*.err") + # monitor logs are needed to report autoscaler events + # TODO(rickyx): remove this after migration. + if not self.is_autoscaler_v2: + # We publish monitor logs in autoscaler v1 + monitor_log_paths += glob.glob(f"{self.logs_dir}/monitor.log") + else: + # We publish autoscaler events directly in autoscaler v2 + monitor_log_paths += glob.glob( + f"{self.logs_dir}/events/event_AUTOSCALER.log" + ) + + # If gcs server restarts, there can be multiple log files. + monitor_log_paths += glob.glob(f"{self.logs_dir}/gcs_server*.err") + + # Add libtpu logs if they exist in the Ray container. + tpu_log_dir = f"{self.logs_dir}/tpu_logs" + if os.path.isdir(tpu_log_dir): + monitor_log_paths += glob.glob(f"{self.logs_dir}/tpu_logs/**") + + # runtime_env setup process is logged here + if RAY_RUNTIME_ENV_LOG_TO_DRIVER_ENABLED: + monitor_log_paths += glob.glob(f"{self.logs_dir}/runtime_env*.log") + for file_path in monitor_log_paths: + if os.path.isfile(file_path) and file_path not in self.log_filenames: + worker_match = WORKER_LOG_PATTERN.match(file_path) + if worker_match: + worker_pid = int(worker_match.group(2)) + else: + worker_pid = None + job_id = None + + # Perform existence check first because most file will not be + # including runtime_env. This saves some cpu cycle. + if "runtime_env" in file_path: + runtime_env_job_match = RUNTIME_ENV_SETUP_PATTERN.match(file_path) + if runtime_env_job_match: + job_id = runtime_env_job_match.group(1) + + is_err_file = file_path.endswith("err") + + self.log_filenames.add(file_path) + self.closed_file_infos.append( + LogFileInfo( + filename=file_path, + size_when_last_opened=0, + file_position=0, + file_handle=None, + is_err_file=is_err_file, + job_id=job_id, + worker_pid=worker_pid, + ) + ) + log_filename = os.path.basename(file_path) + logger.info(f"Beginning to track file {log_filename}") + + def open_closed_files(self): + """Open some closed files if they may have new lines. + + Opening more files may require us to close some of the already open + files. + """ + if not self.can_open_more_files: + # If we can't open any more files. Close all of the files. + self._close_all_files() + + files_with_no_updates = [] + while len(self.closed_file_infos) > 0: + if len(self.open_file_infos) >= self.max_files_open: + self.can_open_more_files = False + break + + file_info = self.closed_file_infos.pop(0) + assert file_info.file_handle is None + # Get the file size to see if it has gotten bigger since we last + # opened it. + try: + file_size = os.path.getsize(file_info.filename) + except (IOError, OSError) as e: + # Catch "file not found" errors. + if e.errno == errno.ENOENT: + logger.warning( + f"Warning: The file {file_info.filename} was not found." + ) + self.log_filenames.remove(file_info.filename) + continue + raise e + + # If some new lines have been added to this file, try to reopen the + # file. + if file_size > file_info.size_when_last_opened: + try: + f = open(file_info.filename, "rb") + except (IOError, OSError) as e: + if e.errno == errno.ENOENT: + logger.warning( + f"Warning: The file {file_info.filename} was not found." + ) + self.log_filenames.remove(file_info.filename) + continue + else: + raise e + + f.seek(file_info.file_position) + file_info.size_when_last_opened = file_size + file_info.file_handle = f + self.open_file_infos.append(file_info) + else: + files_with_no_updates.append(file_info) + + if len(self.open_file_infos) >= self.max_files_open: + self.can_open_more_files = False + # Add the files with no changes back to the list of closed files. + self.closed_file_infos += files_with_no_updates + + def check_log_files_and_publish_updates(self): + """Gets updates to the log files and publishes them. + + Returns: + True if anything was published and false otherwise. + """ + anything_published = False + lines_to_publish = [] + + def flush(): + nonlocal lines_to_publish + nonlocal anything_published + if len(lines_to_publish) > 0: + data = { + "ip": self.ip, + "pid": file_info.worker_pid, + "job": file_info.job_id, + "is_err": file_info.is_err_file, + "lines": lines_to_publish, + "actor_name": file_info.actor_name, + "task_name": file_info.task_name, + } + try: + self.gcs_client.publish_logs(data) + except Exception: + logger.exception(f"Failed to publish log messages {data}") + anything_published = True + lines_to_publish = [] + + for file_info in self.open_file_infos: + assert not file_info.file_handle.closed + file_info.reopen_if_necessary() + + max_num_lines_to_read = ray_constants.LOG_MONITOR_NUM_LINES_TO_READ + for _ in range(max_num_lines_to_read): + try: + next_line = file_info.file_handle.readline() + # Replace any characters not in UTF-8 with + # a replacement character, see + # https://stackoverflow.com/a/38565489/10891801 + next_line = next_line.decode("utf-8", "replace") + if next_line == "": + break + next_line = next_line.rstrip("\r\n") + + if next_line.startswith(ray_constants.LOG_PREFIX_ACTOR_NAME): + flush() # Possible change of task/actor name. + file_info.actor_name = next_line.split( + ray_constants.LOG_PREFIX_ACTOR_NAME, 1 + )[1] + file_info.task_name = None + elif next_line.startswith(ray_constants.LOG_PREFIX_TASK_NAME): + flush() # Possible change of task/actor name. + file_info.task_name = next_line.split( + ray_constants.LOG_PREFIX_TASK_NAME, 1 + )[1] + elif next_line.startswith(ray_constants.LOG_PREFIX_JOB_ID): + file_info.job_id = next_line.split( + ray_constants.LOG_PREFIX_JOB_ID, 1 + )[1] + elif next_line.startswith( + "Windows fatal exception: access violation" + ): + # We are suppressing the + # 'Windows fatal exception: access violation' + # message on workers on Windows here. + # As far as we know it is harmless, + # but is frequently popping up if Python + # functions are run inside the core + # worker C extension. See the investigation in + # github.com/ray-project/ray/issues/18944 + # Also skip the following line, which is an + # empty line. + file_info.file_handle.readline() + else: + lines_to_publish.append(next_line) + except Exception: + logger.error( + f"Error: Reading file: {file_info.filename}, " + f"position: {file_info.file_info.file_handle.tell()} " + "failed." + ) + raise + + if file_info.file_position == 0: + # make filename windows-agnostic + filename = file_info.filename.replace("\\", "/") + if "/raylet" in filename: + file_info.worker_pid = "raylet" + elif "/gcs_server" in filename: + file_info.worker_pid = "gcs_server" + elif "/monitor" in filename or "event_AUTOSCALER" in filename: + file_info.worker_pid = "autoscaler" + elif "/runtime_env" in filename: + file_info.worker_pid = "runtime_env" + + # Record the current position in the file. + file_info.file_position = file_info.file_handle.tell() + flush() + + return anything_published + + def should_update_filenames(self, last_file_updated_time: float) -> bool: + """Return true if filenames should be updated. + + This method is used to apply the backpressure on file updates because + that requires heavy glob operations which use lots of CPUs. + + Args: + last_file_updated_time: The last time filenames are updated. + + Returns: + True if filenames should be updated. False otherwise. + """ + elapsed_seconds = float(time.time() - last_file_updated_time) + return ( + len(self.log_filenames) < RAY_LOG_MONITOR_MANY_FILES_THRESHOLD + or elapsed_seconds > LOG_NAME_UPDATE_INTERVAL_S + ) + + def run(self): + """Run the log monitor. + + This will scan the file system once every LOG_NAME_UPDATE_INTERVAL_S to + check if there are new log files to monitor. It will also publish new + log lines. + """ + last_updated = time.time() + while True: + if self.should_update_filenames(last_updated): + self.update_log_filenames() + last_updated = time.time() + + self.open_closed_files() + anything_published = self.check_log_files_and_publish_updates() + # If nothing was published, then wait a little bit before checking + # for logs to avoid using too much CPU. + if not anything_published: + time.sleep(0.1) + + +def is_proc_alive(pid): + # Import locally to make sure the bundled version is used if needed + import psutil + + try: + return psutil.Process(pid).is_running() + except psutil.NoSuchProcess: + # The process does not exist. + return False + + +if __name__ == "__main__": + parser = argparse.ArgumentParser( + description=("Parse GCS server address for the log monitor to connect to.") + ) + parser.add_argument( + "--gcs-address", required=False, type=str, help="The address (ip:port) of GCS." + ) + parser.add_argument( + "--logging-level", + required=False, + type=str, + default=ray_constants.LOGGER_LEVEL, + choices=ray_constants.LOGGER_LEVEL_CHOICES, + help=ray_constants.LOGGER_LEVEL_HELP, + ) + parser.add_argument( + "--logging-format", + required=False, + type=str, + default=ray_constants.LOGGER_FORMAT, + help=ray_constants.LOGGER_FORMAT_HELP, + ) + parser.add_argument( + "--logging-filename", + required=False, + type=str, + default=ray_constants.LOG_MONITOR_LOG_FILE_NAME, + help="Specify the name of log file, " + "log to stderr if set empty, default is " + f'"{ray_constants.LOG_MONITOR_LOG_FILE_NAME}"', + ) + parser.add_argument( + "--session-dir", + required=True, + type=str, + help="Specify the path of the session directory used by Ray processes.", + ) + parser.add_argument( + "--logs-dir", + required=True, + type=str, + help="Specify the path of the log directory used by Ray processes.", + ) + parser.add_argument( + "--logging-rotate-bytes", + required=True, + type=int, + help="Specify the max bytes for rotating log file.", + ) + parser.add_argument( + "--logging-rotate-backup-count", + required=True, + type=int, + help="Specify the backup count of rotated log file.", + ) + parser.add_argument( + "--stdout-filepath", + required=False, + default="", + type=str, + help="The filepath to dump log monitor stdout.", + ) + parser.add_argument( + "--stderr-filepath", + required=False, + default="", + type=str, + help="The filepath to dump log monitor stderr.", + ) + + args = parser.parse_args() + + # Disable log rotation for windows platform. + logging_rotation_bytes = args.logging_rotate_bytes if sys.platform != "win32" else 0 + logging_rotation_backup_count = ( + args.logging_rotate_backup_count if sys.platform != "win32" else 1 + ) + logging_params = dict( + logging_level=args.logging_level, + logging_format=args.logging_format, + log_dir=args.logs_dir, + filename=args.logging_filename, + max_bytes=logging_rotation_bytes, + backup_count=logging_rotation_backup_count, + ) + logger = setup_component_logger(**logging_params) + + # Setup stdout/stderr redirect files if redirection enabled + logging_utils.redirect_stdout_stderr_if_needed( + args.stdout_filepath, + args.stderr_filepath, + logging_rotation_bytes, + logging_rotation_backup_count, + ) + + node_ip = services.get_cached_node_ip_address(args.session_dir) + gcs_client = GcsClient(address=args.gcs_address) + log_monitor = LogMonitor( + node_ip, + args.logs_dir, + gcs_client, + is_proc_alive, + gcs_address=args.gcs_address, + ) + + try: + log_monitor.run() + except Exception as e: + # Something went wrong, so push an error to all drivers. + traceback_str = ray._private.utils.format_error_message(traceback.format_exc()) + message = ( + f"The log monitor on node {platform.node()} " + f"failed with the following error:\n{traceback_str}" + ) + ray._private.utils.publish_error_to_driver( + ray_constants.LOG_MONITOR_DIED_ERROR, + message, + gcs_client=gcs_client, + ) + logger.error(message) + raise e diff --git a/lib/python3.12/site-packages/ray/_private/logging_utils.py b/lib/python3.12/site-packages/ray/_private/logging_utils.py new file mode 100644 index 0000000000000000000000000000000000000000..21566f87c5ea1a479b0163fdfbc1b67324b7c50a --- /dev/null +++ b/lib/python3.12/site-packages/ray/_private/logging_utils.py @@ -0,0 +1,49 @@ +import sys + +from ray._private.utils import open_log +from ray._raylet import StreamRedirector + + +def redirect_stdout_stderr_if_needed( + stdout_filepath: str, + stderr_filepath: str, + rotation_bytes: int, + rotation_backup_count: int, +): + """This function sets up redirection for stdout and stderr if needed, based on the given rotation parameters. + + params: + stdout_filepath: the filepath stdout will be redirected to; if empty, stdout will not be redirected. + stderr_filepath: the filepath stderr will be redirected to; if empty, stderr will not be redirected. + rotation_bytes: number of bytes which triggers file rotation. + rotation_backup_count: the max size of rotation files. + """ + + # Setup redirection for stdout and stderr. + if stdout_filepath: + StreamRedirector.redirect_stdout( + stdout_filepath, + rotation_bytes, + rotation_backup_count, + False, # tee_to_stdout + False, # tee_to_stderr + ) + if stderr_filepath: + StreamRedirector.redirect_stderr( + stderr_filepath, + rotation_bytes, + rotation_backup_count, + False, # tee_to_stdout + False, # tee_to_stderr + ) + + # Setup python system stdout/stderr. + stdout_fileno = sys.stdout.fileno() + stderr_fileno = sys.stderr.fileno() + # We also manually set sys.stdout and sys.stderr because that seems to + # have an effect on the output buffering. Without doing this, stdout + # and stderr are heavily buffered resulting in seemingly lost logging + # statements. We never want to close the stdout file descriptor, dup2 will + # close it when necessary and we don't want python's GC to close it. + sys.stdout = open_log(stdout_fileno, unbuffered=True, closefd=False) + sys.stderr = open_log(stderr_fileno, unbuffered=True, closefd=False) diff --git a/lib/python3.12/site-packages/ray/_private/memory_monitor.py b/lib/python3.12/site-packages/ray/_private/memory_monitor.py new file mode 100644 index 0000000000000000000000000000000000000000..fd3880c9d50441e775d2227c33ba723ac8584771 --- /dev/null +++ b/lib/python3.12/site-packages/ray/_private/memory_monitor.py @@ -0,0 +1,165 @@ +import logging +import os +import platform +import sys +import time + +import ray # noqa F401 + +# Import ray before psutil will make sure we use psutil's bundled version +from ray._common.utils import get_system_memory + +import psutil # noqa E402 + +logger = logging.getLogger(__name__) + + +def get_rss(memory_info): + """Get the estimated non-shared memory usage from psutil memory_info.""" + mem = memory_info.rss + # OSX doesn't have the shared attribute + if hasattr(memory_info, "shared"): + mem -= memory_info.shared + return mem + + +def get_shared(virtual_memory): + """Get the estimated shared memory usage from psutil virtual mem info.""" + # OSX doesn't have the shared attribute + if hasattr(virtual_memory, "shared"): + return virtual_memory.shared + else: + return 0 + + +def get_top_n_memory_usage(n: int = 10): + """Get the top n memory usage of the process + + Params: + n: Number of top n process memory usage to return. + Returns: + (str) The formatted string of top n process memory usage. + """ + proc_stats = [] + for proc in psutil.process_iter(["memory_info", "cmdline"]): + try: + proc_stats.append( + (get_rss(proc.info["memory_info"]), proc.pid, proc.info["cmdline"]) + ) + except psutil.NoSuchProcess: + # We should skip the process that has exited. Refer this + # issue for more detail: + # https://github.com/ray-project/ray/issues/14929 + continue + except psutil.AccessDenied: + # On MacOS, the proc_pidinfo call (used to get per-process + # memory info) fails with a permission denied error when used + # on a process that isn’t owned by the same user. For now, we + # drop the memory info of any such process, assuming that + # processes owned by other users (e.g. root) aren't Ray + # processes and will be of less interest when an OOM happens + # on a Ray node. + # See issue for more detail: + # https://github.com/ray-project/ray/issues/11845#issuecomment-849904019 # noqa: E501 + continue + proc_str = "PID\tMEM\tCOMMAND" + for rss, pid, cmdline in sorted(proc_stats, reverse=True)[:n]: + proc_str += "\n{}\t{}GiB\t{}".format( + pid, round(rss / (1024**3), 2), " ".join(cmdline)[:100].strip() + ) + return proc_str + + +class RayOutOfMemoryError(Exception): + def __init__(self, msg): + Exception.__init__(self, msg) + + @staticmethod + def get_message(used_gb, total_gb, threshold): + proc_str = get_top_n_memory_usage(n=10) + return ( + "More than {}% of the memory on ".format(int(100 * threshold)) + + "node {} is used ({} / {} GB). ".format( + platform.node(), round(used_gb, 2), round(total_gb, 2) + ) + + f"The top 10 memory consumers are:\n\n{proc_str}" + + "\n\nIn addition, up to {} GiB of shared memory is ".format( + round(get_shared(psutil.virtual_memory()) / (1024**3), 2) + ) + + "currently being used by the Ray object store.\n---\n" + "--- Tip: Use the `ray memory` command to list active " + "objects in the cluster.\n" + "--- To disable OOM exceptions, set " + "RAY_DISABLE_MEMORY_MONITOR=1.\n---\n" + ) + + +class MemoryMonitor: + """Helper class for raising errors on low memory. + + This presents a much cleaner error message to users than what would happen + if we actually ran out of memory. + + The monitor tries to use the cgroup memory limit and usage if it is set + and available so that it is more reasonable inside containers. Otherwise, + it uses `psutil` to check the memory usage. + + The environment variable `RAY_MEMORY_MONITOR_ERROR_THRESHOLD` can be used + to overwrite the default error_threshold setting. + + Used by test only. For production code use memory_monitor.cc + """ + + def __init__(self, error_threshold=0.95, check_interval=1): + # Note: it takes ~50us to check the memory usage through psutil, so + # throttle this check at most once a second or so. + self.check_interval = check_interval + self.last_checked = 0 + try: + self.error_threshold = float( + os.getenv("RAY_MEMORY_MONITOR_ERROR_THRESHOLD") + ) + except (ValueError, TypeError): + self.error_threshold = error_threshold + # Try to read the cgroup memory limit if it is available. + try: + with open("/sys/fs/cgroup/memory/memory.limit_in_bytes", "rb") as f: + self.cgroup_memory_limit_gb = int(f.read()) / (1024**3) + except IOError: + self.cgroup_memory_limit_gb = sys.maxsize / (1024**3) + if not psutil: + logger.warning( + "WARNING: Not monitoring node memory since `psutil` " + "is not installed. Install this with " + "`pip install psutil` to enable " + "debugging of memory-related crashes." + ) + self.disabled = ( + "RAY_DEBUG_DISABLE_MEMORY_MONITOR" in os.environ + or "RAY_DISABLE_MEMORY_MONITOR" in os.environ + ) + + def get_memory_usage(self): + from ray._private.utils import get_used_memory + + total_gb = get_system_memory() / (1024**3) + used_gb = get_used_memory() / (1024**3) + + return used_gb, total_gb + + def raise_if_low_memory(self): + if self.disabled: + return + + if time.time() - self.last_checked > self.check_interval: + self.last_checked = time.time() + used_gb, total_gb = self.get_memory_usage() + + if used_gb > total_gb * self.error_threshold: + raise RayOutOfMemoryError( + RayOutOfMemoryError.get_message( + used_gb, total_gb, self.error_threshold + ) + ) + else: + logger.debug(f"Memory usage is {used_gb} / {total_gb}") diff --git a/lib/python3.12/site-packages/ray/_private/metrics_agent.py b/lib/python3.12/site-packages/ray/_private/metrics_agent.py new file mode 100644 index 0000000000000000000000000000000000000000..6116687e391a2229646c833636ed21ee0242623d --- /dev/null +++ b/lib/python3.12/site-packages/ray/_private/metrics_agent.py @@ -0,0 +1,848 @@ +import json +import logging +import os +import re +import threading +import time +import traceback +from collections import defaultdict, namedtuple +from typing import Any, Dict, List, Set, Tuple, Union + +from opencensus.metrics.export.metric_descriptor import MetricDescriptorType +from opencensus.metrics.export.value import ValueDouble +from opencensus.stats import aggregation, measure as measure_module +from opencensus.stats.aggregation_data import ( + CountAggregationData, + DistributionAggregationData, + LastValueAggregationData, + SumAggregationData, +) +from opencensus.stats.base_exporter import StatsExporter +from opencensus.stats.stats_recorder import StatsRecorder +from opencensus.stats.view import View +from opencensus.stats.view_manager import ViewManager +from opencensus.tags import ( + tag_key as tag_key_module, + tag_map as tag_map_module, + tag_value as tag_value_module, +) +from prometheus_client.core import ( + CounterMetricFamily, + GaugeMetricFamily, + HistogramMetricFamily, + Metric as PrometheusMetric, +) + +import ray +from ray._common.network_utils import build_address +from ray._private.ray_constants import env_bool +from ray._private.telemetry.metric_cardinality import ( + WORKER_ID_TAG_KEY, + MetricCardinality, +) +from ray._raylet import GcsClient +from ray.core.generated.metrics_pb2 import Metric +from ray.util.metrics import _is_invalid_metric_name + +logger = logging.getLogger(__name__) + +# Env var key to decide worker timeout. +# If the worker doesn't report for more than +# this time, we treat workers as dead. +RAY_WORKER_TIMEOUT_S = "RAY_WORKER_TIMEOUT_S" +GLOBAL_COMPONENT_KEY = "CORE" +RE_NON_ALPHANUMS = re.compile(r"[^a-zA-Z0-9]") + + +class Gauge(View): + """Gauge representation of opencensus view. + + This class is used to collect process metrics from the reporter agent. + Cpp metrics should be collected in a different way. + """ + + def __init__(self, name, description, unit, tags: List[str]): + if _is_invalid_metric_name(name): + raise ValueError( + f"Invalid metric name: {name}. Metric will be discarded " + "and data will not be collected or published. " + "Metric names can only contain letters, numbers, _, and :. " + "Metric names cannot start with numbers." + ) + self._measure = measure_module.MeasureInt(name, description, unit) + self._description = description + tags = [tag_key_module.TagKey(tag) for tag in tags] + self._view = View( + name, description, tags, self.measure, aggregation.LastValueAggregation() + ) + + @property + def measure(self): + return self._measure + + @property + def view(self): + return self._view + + @property + def name(self): + return self.measure.name + + @property + def description(self): + return self._description + + +Record = namedtuple("Record", ["gauge", "value", "tags"]) + + +def fix_grpc_metric(metric: Metric): + """ + Fix the inbound `opencensus.proto.metrics.v1.Metric` protos to make it acceptable + by opencensus.stats.DistributionAggregationData. + + - metric name: gRPC OpenCensus metrics have names with slashes and dots, e.g. + `grpc.io/client/server_latency`[1]. However Prometheus metric names only take + alphanums,underscores and colons[2]. We santinize the name by replacing non-alphanum + chars to underscore, like the official opencensus prometheus exporter[3]. + - distribution bucket bounds: The Metric proto asks distribution bucket bounds to + be > 0 [4]. However, gRPC OpenCensus metrics have their first bucket bound == 0 [1]. + This makes the `DistributionAggregationData` constructor to raise Exceptions. This + applies to all bytes and milliseconds (latencies). The fix: we update the initial 0 + bounds to be 0.000_000_1. This will not affect the precision of the metrics, since + we don't expect any less-than-1 bytes, or less-than-1-nanosecond times. + + [1] https://github.com/census-instrumentation/opencensus-specs/blob/master/stats/gRPC.md#units # noqa: E501 + [2] https://prometheus.io/docs/concepts/data_model/#metric-names-and-labels + [3] https://github.com/census-instrumentation/opencensus-cpp/blob/50eb5de762e5f87e206c011a4f930adb1a1775b1/opencensus/exporters/stats/prometheus/internal/prometheus_utils.cc#L39 # noqa: E501 + [4] https://github.com/census-instrumentation/opencensus-proto/blob/master/src/opencensus/proto/metrics/v1/metrics.proto#L218 # noqa: E501 + """ + + if not metric.metric_descriptor.name.startswith("grpc.io/"): + return + + metric.metric_descriptor.name = RE_NON_ALPHANUMS.sub( + "_", metric.metric_descriptor.name + ) + + for series in metric.timeseries: + for point in series.points: + if point.HasField("distribution_value"): + dist_value = point.distribution_value + bucket_bounds = dist_value.bucket_options.explicit.bounds + if len(bucket_bounds) > 0 and bucket_bounds[0] == 0: + bucket_bounds[0] = 0.000_000_1 + + +class OpencensusProxyMetric: + def __init__(self, name: str, desc: str, unit: str, label_keys: List[str]): + """Represents the OpenCensus metrics that will be proxy exported.""" + self._name = name + self._desc = desc + self._unit = unit + # -- The label keys of the metric -- + self._label_keys = label_keys + # -- The data that needs to be proxy exported -- + # tuple of label values -> data (OpenCesnsus Aggregation data) + self._data = {} + + @property + def name(self): + return self._name + + @property + def desc(self): + return self._desc + + @property + def unit(self): + return self._unit + + @property + def label_keys(self): + return self._label_keys + + @property + def data(self): + return self._data + + def is_distribution_aggregation_data(self): + """Check if the metric is a distribution aggreation metric.""" + return len(self._data) > 0 and isinstance( + next(iter(self._data.values())), DistributionAggregationData + ) + + def add_data(self, label_values: Tuple, data: Any): + """Add the data to the metric. + + Args: + label_values: The label values of the metric. + data: The data to be added. + """ + self._data[label_values] = data + + def record(self, metric: Metric): + """Parse the Opencensus Protobuf and store the data. + + The data can be accessed via `data` API once recorded. + """ + timeseries = metric.timeseries + + if len(timeseries) == 0: + return + + # Create the aggregation and fill it in the our stats + for series in timeseries: + labels = tuple(val.value for val in series.label_values) + + # Aggregate points. + for point in series.points: + if ( + metric.metric_descriptor.type + == MetricDescriptorType.CUMULATIVE_INT64 + ): + data = CountAggregationData(point.int64_value) + elif ( + metric.metric_descriptor.type + == MetricDescriptorType.CUMULATIVE_DOUBLE + ): + data = SumAggregationData(ValueDouble, point.double_value) + elif metric.metric_descriptor.type == MetricDescriptorType.GAUGE_DOUBLE: + data = LastValueAggregationData(ValueDouble, point.double_value) + elif ( + metric.metric_descriptor.type + == MetricDescriptorType.CUMULATIVE_DISTRIBUTION + ): + dist_value = point.distribution_value + counts_per_bucket = [bucket.count for bucket in dist_value.buckets] + bucket_bounds = dist_value.bucket_options.explicit.bounds + data = DistributionAggregationData( + dist_value.sum / dist_value.count, + dist_value.count, + dist_value.sum_of_squared_deviation, + counts_per_bucket, + bucket_bounds, + ) + else: + raise ValueError("Summary is not supported") + self._data[labels] = data + + +class Component: + def __init__(self, id: str): + """Represent a component that requests to proxy export metrics + + Args: + id: Id of this component. + """ + self.id = id + # -- The time this component reported its metrics last time -- + # It is used to figure out if this component is stale. + self._last_reported_time = time.monotonic() + # -- Metrics requested to proxy export from this component -- + # metrics_name (str) -> metric (OpencensusProxyMetric) + self._metrics = {} + + @property + def metrics(self) -> Dict[str, OpencensusProxyMetric]: + """Return the metrics requested to proxy export from this component.""" + return self._metrics + + @property + def last_reported_time(self): + return self._last_reported_time + + def record(self, metrics: List[Metric]): + """Parse the Opencensus protobuf and store metrics. + + Metrics can be accessed via `metrics` API for proxy export. + + Args: + metrics: A list of Opencensus protobuf for proxy export. + """ + self._last_reported_time = time.monotonic() + for metric in metrics: + fix_grpc_metric(metric) + descriptor = metric.metric_descriptor + name = descriptor.name + label_keys = [label_key.key for label_key in descriptor.label_keys] + + if name not in self._metrics: + self._metrics[name] = OpencensusProxyMetric( + name, descriptor.description, descriptor.unit, label_keys + ) + self._metrics[name].record(metric) + + +class OpenCensusProxyCollector: + def __init__(self, namespace: str, component_timeout_s: int = 60): + """Prometheus collector implementation for opencensus proxy export. + + Prometheus collector requires to implement `collect` which is + invoked whenever Prometheus queries the endpoint. + + The class is thread-safe. + + Args: + namespace: Prometheus namespace. + """ + # -- Protect `self._components` -- + self._components_lock = threading.Lock() + # -- Timeout until the component is marked as stale -- + # Once the component is considered as stale, + # the metrics from that worker won't be exported. + self._component_timeout_s = component_timeout_s + # -- Prometheus namespace -- + self._namespace = namespace + # -- Component that requests to proxy export metrics -- + # Component means core worker, raylet, and GCS. + # component_id -> Components + # For workers, they contain worker ids. + # For other components (raylet, GCS), + # they contain the global key `GLOBAL_COMPONENT_KEY`. + self._components = {} + # Whether we want to export counter as gauge. + # This is for bug compatibility. + # See https://github.com/ray-project/ray/pull/43795. + self._export_counter_as_gauge = env_bool("RAY_EXPORT_COUNTER_AS_GAUGE", True) + + def record(self, metrics: List[Metric], worker_id_hex: str = None): + """Record the metrics reported from the component that reports it. + + Args: + metrics: A list of opencensus protobuf to proxy export metrics. + worker_id_hex: A worker id that reports these metrics. + If None, it means they are reported from Raylet or GCS. + """ + key = GLOBAL_COMPONENT_KEY if not worker_id_hex else worker_id_hex + with self._components_lock: + if key not in self._components: + self._components[key] = Component(key) + self._components[key].record(metrics) + + def clean_stale_components(self): + """Clean up stale components. + + Stale means the component is dead or unresponsive. + + Stale components won't be reported to Prometheus anymore. + """ + with self._components_lock: + stale_components = [] + stale_component_ids = [] + for id, component in self._components.items(): + elapsed = time.monotonic() - component.last_reported_time + if elapsed > self._component_timeout_s: + stale_component_ids.append(id) + logger.info( + "Metrics from a worker ({}) is cleaned up due to " + "timeout. Time since last report {}s".format(id, elapsed) + ) + for id in stale_component_ids: + stale_components.append(self._components.pop(id)) + return stale_components + + # TODO(sang): add start and end timestamp + def to_prometheus_metrics( + self, + metric_name: str, + metric_description: str, + label_keys: List[str], + metric_units: str, + label_values: Tuple[tag_value_module.TagValue], + agg_data: Any, + metrics_map: Dict[str, List[PrometheusMetric]], + ) -> None: + """to_metric translate the data that OpenCensus create + to Prometheus format, using Prometheus Metric object. + + This method is from Opencensus Prometheus Exporter. + + Args: + metric_name: Name of the metric. + metric_description: Description of the metric. + label_keys: The fixed label keys of the metric. + metric_units: Units of the metric. + label_values: The values of `label_keys`. + agg_data: `opencensus.stats.aggregation_data.AggregationData` object. + Aggregated data that needs to be converted as Prometheus samples + metrics_map: The converted metric is added to this map. + + """ + assert self._components_lock.locked() + metric_name = f"{self._namespace}_{metric_name}" + assert len(label_values) == len(label_keys), (label_values, label_keys) + # Prometheus requires that all tag values be strings hence + # the need to cast none to the empty string before exporting. See + # https://github.com/census-instrumentation/opencensus-python/issues/480 + label_values = [tv if tv else "" for tv in label_values] + + if isinstance(agg_data, CountAggregationData): + metrics = metrics_map.get(metric_name) + if not metrics: + metric = CounterMetricFamily( + name=metric_name, + documentation=metric_description, + unit=metric_units, + labels=label_keys, + ) + metrics = [metric] + metrics_map[metric_name] = metrics + metrics[0].add_metric(labels=label_values, value=agg_data.count_data) + return + + if isinstance(agg_data, SumAggregationData): + # This should be emitted as prometheus counter + # but we used to emit it as prometheus gauge. + # To keep the backward compatibility + # (changing from counter to gauge changes the metric name + # since prometheus client will add "_total" suffix to counter + # per OpenMetrics specification), + # we now emit both counter and gauge and in the + # next major Ray release (3.0) we can stop emitting gauge. + # This leaves people enough time to migrate their dashboards. + # See https://github.com/ray-project/ray/pull/43795. + metrics = metrics_map.get(metric_name) + if not metrics: + metric = CounterMetricFamily( + name=metric_name, + documentation=metric_description, + labels=label_keys, + ) + metrics = [metric] + metrics_map[metric_name] = metrics + metrics[0].add_metric(labels=label_values, value=agg_data.sum_data) + + if not self._export_counter_as_gauge: + pass + elif metric_name.endswith("_total"): + # In this case, we only need to emit prometheus counter + # since for metric name already ends with _total suffix + # prometheus client won't change it + # so there is no backward compatibility issue. + # See https://prometheus.github.io/client_python/instrumenting/counter/ + pass + else: + if len(metrics) == 1: + metric = GaugeMetricFamily( + name=metric_name, + documentation=( + f"(DEPRECATED, use {metric_name}_total metric instead) " + f"{metric_description}" + ), + labels=label_keys, + ) + metrics.append(metric) + assert len(metrics) == 2 + metrics[1].add_metric(labels=label_values, value=agg_data.sum_data) + return + + elif isinstance(agg_data, DistributionAggregationData): + + assert agg_data.bounds == sorted(agg_data.bounds) + # buckets are a list of buckets. Each bucket is another list with + # a pair of bucket name and value, or a triple of bucket name, + # value, and exemplar. buckets need to be in order. + buckets = [] + cum_count = 0 # Prometheus buckets expect cumulative count. + for ii, bound in enumerate(agg_data.bounds): + cum_count += agg_data.counts_per_bucket[ii] + bucket = [str(bound), cum_count] + buckets.append(bucket) + # Prometheus requires buckets to be sorted, and +Inf present. + # In OpenCensus we don't have +Inf in the bucket bonds so need to + # append it here. + buckets.append(["+Inf", agg_data.count_data]) + metrics = metrics_map.get(metric_name) + if not metrics: + metric = HistogramMetricFamily( + name=metric_name, + documentation=metric_description, + labels=label_keys, + ) + metrics = [metric] + metrics_map[metric_name] = metrics + metrics[0].add_metric( + labels=label_values, + buckets=buckets, + sum_value=agg_data.sum, + ) + return + + elif isinstance(agg_data, LastValueAggregationData): + metrics = metrics_map.get(metric_name) + if not metrics: + metric = GaugeMetricFamily( + name=metric_name, + documentation=metric_description, + labels=label_keys, + ) + metrics = [metric] + metrics_map[metric_name] = metrics + metrics[0].add_metric(labels=label_values, value=agg_data.value) + return + + else: + raise ValueError(f"unsupported aggregation type {type(agg_data)}") + + def _aggregate_metric_data( + self, + datas: List[ + Union[LastValueAggregationData, CountAggregationData, SumAggregationData] + ], + ) -> Union[LastValueAggregationData, CountAggregationData, SumAggregationData]: + assert len(datas) > 0 + sample = datas[0] + if isinstance(sample, LastValueAggregationData): + return LastValueAggregationData( + ValueDouble, sum([data.value for data in datas]) + ) + if isinstance(sample, CountAggregationData): + return CountAggregationData(sum([data.count_data for data in datas])) + if isinstance(sample, SumAggregationData): + return SumAggregationData( + ValueDouble, sum([data.sum_data for data in datas]) + ) + + raise ValueError( + f"Unsupported aggregation type {type(sample)}. " + "Supported types are " + f"{CountAggregationData}, {LastValueAggregationData}, {SumAggregationData}." + f"Got {datas}." + ) + + def _aggregate_with_recommended_cardinality( + self, + per_worker_metrics: List[OpencensusProxyMetric], + ) -> List[OpencensusProxyMetric]: + """Collect per-worker metrics, aggregate them into per-node metrics and convert + them to Prometheus format. + + Args: + per_worker_metrics: A list of per-worker metrics for the same metric name. + Returns: + A list of per-node metrics for the same metric name, with the high + cardinality labels removed and the values aggregated. + """ + metric = next(iter(per_worker_metrics), None) + if not metric or WORKER_ID_TAG_KEY not in metric.label_keys: + # No high cardinality labels, return the original metrics. + return per_worker_metrics + + worker_id_label_index = metric.label_keys.index(WORKER_ID_TAG_KEY) + # map from the tuple of label values without worker_id to the list of per worker + # task metrics + label_value_to_data: Dict[ + Tuple, + List[ + Union[ + LastValueAggregationData, + CountAggregationData, + SumAggregationData, + ] + ], + ] = defaultdict(list) + for metric in per_worker_metrics: + for label_values, data in metric.data.items(): + # remove the worker_id from the label values + label_value_to_data[ + label_values[:worker_id_label_index] + + label_values[worker_id_label_index + 1 :] + ].append(data) + + aggregated_metric = OpencensusProxyMetric( + name=metric.name, + desc=metric.desc, + unit=metric.unit, + # remove the worker_id from the label keys + label_keys=metric.label_keys[:worker_id_label_index] + + metric.label_keys[worker_id_label_index + 1 :], + ) + for label_values, datas in label_value_to_data.items(): + aggregated_metric.add_data( + label_values, + self._aggregate_metric_data(datas), + ) + + return [aggregated_metric] + + def collect(self): # pragma: NO COVER + """Collect fetches the statistics from OpenCensus + and delivers them as Prometheus Metrics. + Collect is invoked every time a prometheus.Gatherer is run + for example when the HTTP endpoint is invoked by Prometheus. + + This method is required as a Prometheus Collector. + """ + with self._components_lock: + # First construct the list of opencensus metrics to be converted to + # prometheus metrics. For LEGACY cardinality level, this comprises all + # metrics from all components. For RECOMMENDED cardinality level, we need + # to remove the high cardinality labels and aggreate the component metrics. + open_cencus_metrics: List[OpencensusProxyMetric] = [] + # The metrics that need to be aggregated with recommended cardinality. Key + # is the metric name and value is the list of per-worker metrics. + to_lower_cardinality: Dict[str, List[OpencensusProxyMetric]] = defaultdict( + list + ) + cardinality_level = MetricCardinality.get_cardinality_level() + for component in self._components.values(): + for metric in component.metrics.values(): + if ( + cardinality_level == MetricCardinality.RECOMMENDED + and not metric.is_distribution_aggregation_data() + ): + # We reduce the cardinality for all metrics except for histogram + # metrics. The aggregation of histogram metrics from worker + # level to node level is not well defined. In addition, we + # currently have very few histogram metrics in Ray + # (metric_defs.cc) so the impact of them is negligible. + to_lower_cardinality[metric.name].append(metric) + else: + open_cencus_metrics.append(metric) + for per_worker_metrics in to_lower_cardinality.values(): + open_cencus_metrics.extend( + self._aggregate_with_recommended_cardinality( + per_worker_metrics, + ) + ) + + prometheus_metrics_map = {} + for metric in open_cencus_metrics: + for label_values, data in metric.data.items(): + self.to_prometheus_metrics( + metric.name, + metric.desc, + metric.label_keys, + metric.unit, + label_values, + data, + prometheus_metrics_map, + ) + + for metrics in prometheus_metrics_map.values(): + for metric in metrics: + yield metric + + +class MetricsAgent: + def __init__( + self, + view_manager: ViewManager, + stats_recorder: StatsRecorder, + stats_exporter: StatsExporter = None, + ): + """A class to record and export metrics. + + The class exports metrics in 2 different ways. + - Directly record and export metrics using OpenCensus. + - Proxy metrics from other core components + (e.g., raylet, GCS, core workers). + + This class is thread-safe. + """ + # Lock required because gRPC server uses + # multiple threads to process requests. + self._lock = threading.Lock() + + # + # Opencensus components to record metrics. + # + + # Managing views to export metrics + # If the stats_exporter is None, we disable all metrics export. + self.view_manager = view_manager + # A class that's used to record metrics + # emitted from the current process. + self.stats_recorder = stats_recorder + # A class to export metrics. + self.stats_exporter = stats_exporter + # -- A Prometheus custom collector to proxy export metrics -- + # `None` if the prometheus server is not started. + self.proxy_exporter_collector = None + + if self.stats_exporter is None: + # If the exporter is not given, + # we disable metrics collection. + self.view_manager = None + else: + self.view_manager.register_exporter(stats_exporter) + self.proxy_exporter_collector = OpenCensusProxyCollector( + self.stats_exporter.options.namespace, + component_timeout_s=int(os.getenv(RAY_WORKER_TIMEOUT_S, 120)), + ) + + # Registered view names. + self._registered_views: Set[str] = set() + + def record_and_export(self, records: List[Record], global_tags=None): + """Directly record and export stats from the same process.""" + global_tags = global_tags or {} + with self._lock: + if not self.view_manager: + return + + for record in records: + gauge = record.gauge + value = record.value + tags = record.tags + try: + self._record_gauge(gauge, value, {**tags, **global_tags}) + except Exception as e: + logger.error( + f"Failed to record metric {gauge.name} with value {value} with tags {tags!r} and global tags {global_tags!r} due to: {e!r}" + ) + + def _record_gauge(self, gauge: Gauge, value: float, tags: dict): + if gauge.name not in self._registered_views: + self.view_manager.register_view(gauge.view) + self._registered_views.add(gauge.name) + measurement_map = self.stats_recorder.new_measurement_map() + tag_map = tag_map_module.TagMap() + for key, tag_val in tags.items(): + try: + tag_key = tag_key_module.TagKey(key) + except ValueError as e: + logger.error( + f"Failed to create tag key {key} for metric {gauge.name} due to: {e!r}" + ) + raise e + try: + tag_value = tag_value_module.TagValue(tag_val) + except ValueError as e: + logger.error( + f"Failed to create tag value {tag_val} for key {key} for metric {gauge.name} due to: {e!r}" + ) + raise e + tag_map.insert(tag_key, tag_value) + measurement_map.measure_float_put(gauge.measure, value) + # NOTE: When we record this metric, timestamp will be renewed. + measurement_map.record(tag_map) + + def proxy_export_metrics(self, metrics: List[Metric], worker_id_hex: str = None): + """Proxy export metrics specified by a Opencensus Protobuf. + + This API is used to export metrics emitted from + core components. + + Args: + metrics: A list of protobuf Metric defined from OpenCensus. + worker_id_hex: The worker ID it proxies metrics export. None + if the metric is not from a worker (i.e., raylet, GCS). + """ + with self._lock: + if not self.view_manager: + return + + self._proxy_export_metrics(metrics, worker_id_hex) + + def _proxy_export_metrics(self, metrics: List[Metric], worker_id_hex: str = None): + self.proxy_exporter_collector.record(metrics, worker_id_hex) + + def clean_all_dead_worker_metrics(self): + """Clean dead worker's metrics. + + Worker metrics are cleaned up and won't be exported once + it is considered as dead. + + This method has to be periodically called by a caller. + """ + with self._lock: + if not self.view_manager: + return + + self.proxy_exporter_collector.clean_stale_components() + + +class PrometheusServiceDiscoveryWriter(threading.Thread): + """A class to support Prometheus service discovery. + + It supports file-based service discovery. Checkout + https://prometheus.io/docs/guides/file-sd/ for more details. + + Args: + gcs_address: Gcs address for this cluster. + temp_dir: Temporary directory used by + Ray to store logs and metadata. + """ + + def __init__(self, gcs_address, temp_dir): + gcs_client_options = ray._raylet.GcsClientOptions.create( + gcs_address, None, allow_cluster_id_nil=True, fetch_cluster_id_if_nil=False + ) + self.gcs_address = gcs_address + + ray._private.state.state._initialize_global_state(gcs_client_options) + self.temp_dir = temp_dir + self.default_service_discovery_flush_period = 5 + + # The last service discovery content that PrometheusServiceDiscoveryWriter has seen + self.latest_service_discovery_content = [] + self._content_lock = threading.RLock() + + super().__init__() + + def get_latest_service_discovery_content(self): + """Return the latest stored service discovery content.""" + with self._content_lock: + return self.latest_service_discovery_content + + def get_file_discovery_content(self): + """Return the content for Prometheus service discovery.""" + nodes = ray.nodes() + metrics_export_addresses = [ + build_address(node["NodeManagerAddress"], node["MetricsExportPort"]) + for node in nodes + if node["alive"] is True + ] + gcs_client = GcsClient(address=self.gcs_address) + autoscaler_addr = gcs_client.internal_kv_get(b"AutoscalerMetricsAddress", None) + if autoscaler_addr: + metrics_export_addresses.append(autoscaler_addr.decode("utf-8")) + dashboard_addr = gcs_client.internal_kv_get(b"DashboardMetricsAddress", None) + if dashboard_addr: + metrics_export_addresses.append(dashboard_addr.decode("utf-8")) + content = [{"labels": {"job": "ray"}, "targets": metrics_export_addresses}] + with self._content_lock: + self.latest_service_discovery_content = content + return json.dumps(content) + + def write(self): + # Write a file based on https://prometheus.io/docs/guides/file-sd/ + # Write should be atomic. Otherwise, Prometheus raises an error that + # json file format is invalid because it reads a file when + # file is re-written. Note that Prometheus still works although we + # have this error. + temp_file_name = self.get_temp_file_name() + with open(temp_file_name, "w") as json_file: + json_file.write(self.get_file_discovery_content()) + # NOTE: os.replace is atomic on both Linux and Windows, so we won't + # have race condition reading this file. + os.replace(temp_file_name, self.get_target_file_name()) + + def get_target_file_name(self): + return os.path.join( + self.temp_dir, ray._private.ray_constants.PROMETHEUS_SERVICE_DISCOVERY_FILE + ) + + def get_temp_file_name(self): + return os.path.join( + self.temp_dir, + "{}_{}".format( + "tmp", ray._private.ray_constants.PROMETHEUS_SERVICE_DISCOVERY_FILE + ), + ) + + def run(self): + while True: + # This thread won't be broken by exceptions. + try: + self.write() + except Exception as e: + logger.warning( + "Writing a service discovery file, {}," + "failed.".format(self.get_target_file_name()) + ) + logger.warning(traceback.format_exc()) + logger.warning(f"Error message: {e}") + time.sleep(self.default_service_discovery_flush_period) diff --git a/lib/python3.12/site-packages/ray/_private/node.py b/lib/python3.12/site-packages/ray/_private/node.py new file mode 100644 index 0000000000000000000000000000000000000000..084c80fcfbd3e6e4a869cf9c8b3e7c409c7475a8 --- /dev/null +++ b/lib/python3.12/site-packages/ray/_private/node.py @@ -0,0 +1,1882 @@ +import atexit +import collections +import datetime +import errno +import json +import logging +import os +import random +import signal +import socket +import subprocess +import sys +import tempfile +import threading +import time +import traceback +from collections import defaultdict +from typing import IO, AnyStr, Dict, Optional, Tuple + +from filelock import FileLock + +import ray +import ray._private.ray_constants as ray_constants +import ray._private.services +from ray._common.network_utils import ( + build_address, + get_localhost_ip, + is_ipv6, + parse_address, +) +from ray._common.ray_constants import LOGGING_ROTATE_BACKUP_COUNT, LOGGING_ROTATE_BYTES +from ray._common.utils import try_to_create_directory +from ray._private.resource_and_label_spec import ResourceAndLabelSpec +from ray._private.resource_isolation_config import ResourceIsolationConfig +from ray._private.services import get_address, serialize_config +from ray._private.utils import ( + is_in_test, + open_log, + try_to_symlink, + validate_socket_filepath, +) +from ray._raylet import GcsClient, get_session_key_from_storage + +import psutil + +# Logger for this module. It should be configured at the entry point +# into the program using Ray. Ray configures it by default automatically +# using logging.basicConfig in its entry/init points. +logger = logging.getLogger(__name__) + + +class Node: + """An encapsulation of the Ray processes on a single node. + + This class is responsible for starting Ray processes and killing them, + and it also controls the temp file policy. + + Attributes: + all_processes: A mapping from process type (str) to a list of + ProcessInfo objects. All lists have length one except for the Redis + server list, which has multiple. + """ + + def __init__( + self, + ray_params, + head: bool = False, + shutdown_at_exit: bool = True, + spawn_reaper: bool = True, + connect_only: bool = False, + default_worker: bool = False, + ray_init_cluster: bool = False, + ): + """Start a node. + + Args: + ray_params: The RayParams to use to configure the node. + head: True if this is the head node, which means it will + start additional processes like the Redis servers, monitor + processes, and web UI. + shutdown_at_exit: If true, spawned processes will be cleaned + up if this process exits normally. + spawn_reaper: If true, spawns a process that will clean up + other spawned processes if this process dies unexpectedly. + connect_only: If true, connect to the node without starting + new processes. + default_worker: Whether it's running from a ray worker or not + ray_init_cluster: Whether it's a cluster created by ray.init() + """ + if shutdown_at_exit: + if connect_only: + raise ValueError( + "'shutdown_at_exit' and 'connect_only' cannot both be true." + ) + self._register_shutdown_hooks() + self._default_worker = default_worker + self.head = head + self.kernel_fate_share = bool( + spawn_reaper and ray._private.utils.detect_fate_sharing_support() + ) + self.resource_isolation_config: ResourceIsolationConfig = ( + ray_params.resource_isolation_config + ) + self.all_processes: dict = {} + self.removal_lock = threading.Lock() + + self.ray_init_cluster = ray_init_cluster + if ray_init_cluster: + assert head, "ray.init() created cluster only has the head node" + + # Set up external Redis when `RAY_REDIS_ADDRESS` is specified. + redis_address_env = os.environ.get("RAY_REDIS_ADDRESS") + if ray_params.external_addresses is None and redis_address_env is not None: + external_redis = redis_address_env.split(",") + + # Reuse primary Redis as Redis shard when there's only one + # instance provided. + if len(external_redis) == 1: + external_redis.append(external_redis[0]) + ray_params.external_addresses = external_redis + ray_params.num_redis_shards = len(external_redis) - 1 + + if ( + ray_params._system_config + and len(ray_params._system_config) > 0 + and (not head and not connect_only) + ): + raise ValueError( + "System config parameters can only be set on the head node." + ) + + ray_params.update_if_absent( + include_log_monitor=True, + resources={}, + worker_path=os.path.join( + os.path.dirname(os.path.abspath(__file__)), + "workers", + "default_worker.py", + ), + setup_worker_path=os.path.join( + os.path.dirname(os.path.abspath(__file__)), + "workers", + ray_constants.SETUP_WORKER_FILENAME, + ), + ) + + self._resource_and_label_spec = None + self._localhost = get_localhost_ip() + self._ray_params = ray_params + self._config = ray_params._system_config or {} + + self._dashboard_agent_listen_port = ray_params.dashboard_agent_listen_port + + # Configure log rotation parameters. + self.max_bytes = int(os.getenv("RAY_ROTATION_MAX_BYTES", LOGGING_ROTATE_BYTES)) + self.backup_count = int( + os.getenv("RAY_ROTATION_BACKUP_COUNT", LOGGING_ROTATE_BACKUP_COUNT) + ) + + assert self.max_bytes >= 0 + assert self.backup_count >= 0 + + self._redis_address = ray_params.redis_address + if head: + ray_params.update_if_absent(num_redis_shards=1) + self._gcs_address = ray_params.gcs_address + self._gcs_client = None + + if not self.head: + self.validate_ip_port(self.address) + self._init_gcs_client() + + # Register the temp dir. + self._session_name = ray_params.session_name + if self._session_name is None: + if head: + # We expect this the first time we initialize a cluster, but not during + # subsequent restarts of the head node. + maybe_key = self.check_persisted_session_name() + if maybe_key is None: + # date including microsecond + date_str = datetime.datetime.today().strftime( + "%Y-%m-%d_%H-%M-%S_%f" + ) + self._session_name = f"session_{date_str}_{os.getpid()}" + else: + self._session_name = ray._common.utils.decode(maybe_key) + else: + assert not self._default_worker + session_name = ray._private.utils.internal_kv_get_with_retry( + self.get_gcs_client(), + "session_name", + ray_constants.KV_NAMESPACE_SESSION, + num_retries=ray_constants.NUM_REDIS_GET_RETRIES, + ) + self._session_name = ray._common.utils.decode(session_name) + + # Initialize webui url + if head: + self._webui_url = None + else: + if ray_params.webui is None: + assert not self._default_worker + self._webui_url = ray._private.services.get_webui_url_from_internal_kv() + else: + self._webui_url = build_address( + ray_params.dashboard_host, ray_params.dashboard_port + ) + + # It creates a session_dir. + self._init_temp() + + node_ip_address = ray_params.node_ip_address + if node_ip_address is None: + if connect_only: + node_ip_address = self._wait_and_get_for_node_address() + else: + node_ip_address = ray.util.get_node_ip_address() + + assert node_ip_address is not None + ray_params.update_if_absent(node_ip_address=node_ip_address) + self._node_ip_address = node_ip_address + if not connect_only: + ray._private.services.write_node_ip_address( + self.get_session_dir_path(), node_ip_address + ) + + self._object_spilling_config = self._get_object_spilling_config() + logger.debug( + f"Starting node with object spilling config: {self._object_spilling_config}" + ) + + # Obtain the fallback directoy from the object spilling config + # Currently, we set the fallback directory to be the same as the object spilling + # path when the object spills to file system + self._fallback_directory = None + if self._object_spilling_config: + config = json.loads(self._object_spilling_config) + if config.get("type") == "filesystem": + directory_path = config.get("params", {}).get("directory_path") + if isinstance(directory_path, list): + self._fallback_directory = directory_path[0] + elif isinstance(directory_path, str): + self._fallback_directory = directory_path + + # If it is a head node, try validating if external storage is configurable. + if head: + self.validate_external_storage() + + if connect_only: + # Get socket names from the configuration. + self._plasma_store_socket_name = ray_params.plasma_store_socket_name + self._raylet_socket_name = ray_params.raylet_socket_name + self._node_id = ray_params.node_id + + # If user does not provide the socket name, get it from Redis. + if ( + self._plasma_store_socket_name is None + or self._raylet_socket_name is None + or self._ray_params.node_manager_port is None + or self._node_id is None + ): + # Get the address info of the processes to connect to + # from Redis or GCS. + node_info = ray._private.services.get_node_to_connect_for_driver( + self.gcs_address, + self._node_ip_address, + ) + self._plasma_store_socket_name = node_info["object_store_socket_name"] + self._raylet_socket_name = node_info["raylet_socket_name"] + self._ray_params.node_manager_port = node_info["node_manager_port"] + self._node_id = node_info["node_id"] + else: + # If the user specified a socket name, use it. + self._plasma_store_socket_name = self._prepare_socket_file( + self._ray_params.plasma_store_socket_name, default_prefix="plasma_store" + ) + self._raylet_socket_name = self._prepare_socket_file( + self._ray_params.raylet_socket_name, default_prefix="raylet" + ) + if ( + self._ray_params.env_vars is not None + and "RAY_OVERRIDE_NODE_ID_FOR_TESTING" in self._ray_params.env_vars + ): + node_id = self._ray_params.env_vars["RAY_OVERRIDE_NODE_ID_FOR_TESTING"] + logger.debug( + f"Setting node ID to {node_id} " + "based on ray_params.env_vars override" + ) + self._node_id = node_id + elif os.environ.get("RAY_OVERRIDE_NODE_ID_FOR_TESTING"): + node_id = os.environ["RAY_OVERRIDE_NODE_ID_FOR_TESTING"] + logger.debug(f"Setting node ID to {node_id} based on env override") + self._node_id = node_id + else: + node_id = ray.NodeID.from_random().hex() + logger.debug(f"Setting node ID to {node_id}") + self._node_id = node_id + + # The dashboard agent port is assigned first to avoid + # other processes accidentally taking its default port + self._dashboard_agent_listen_port = self._get_cached_port( + "dashboard_agent_listen_port", + default_port=ray_params.dashboard_agent_listen_port, + ) + + self.metrics_agent_port = self._get_cached_port( + "metrics_agent_port", default_port=ray_params.metrics_agent_port + ) + self._metrics_export_port = self._get_cached_port( + "metrics_export_port", default_port=ray_params.metrics_export_port + ) + self._runtime_env_agent_port = self._get_cached_port( + "runtime_env_agent_port", + default_port=ray_params.runtime_env_agent_port, + ) + + ray_params.update_if_absent( + metrics_agent_port=self.metrics_agent_port, + metrics_export_port=self._metrics_export_port, + dashboard_agent_listen_port=self._dashboard_agent_listen_port, + runtime_env_agent_port=self._runtime_env_agent_port, + ) + + # Pick a GCS server port. + if head: + gcs_server_port = os.getenv(ray_constants.GCS_PORT_ENVIRONMENT_VARIABLE) + if gcs_server_port: + ray_params.update_if_absent(gcs_server_port=int(gcs_server_port)) + if ray_params.gcs_server_port is None or ray_params.gcs_server_port == 0: + ray_params.gcs_server_port = self._get_cached_port("gcs_server_port") + + if not connect_only and spawn_reaper and not self.kernel_fate_share: + self.start_reaper_process() + if not connect_only: + self._ray_params.update_pre_selected_port() + + # Start processes. + if head: + self.start_head_processes() + + if not connect_only: + self.start_ray_processes() + # Wait for the node info to be available in the GCS so that + # we know it's started up. + + # Grace period to let the Raylet register with the GCS. + # We retry in a loop in case it takes longer than expected. + time.sleep(0.1) + start_time = time.monotonic() + raylet_start_wait_time_s = 30 + node_info = None + while True: + try: + # Will raise a RuntimeError if the node info is not available. + node_info = ray._private.services.get_node( + self.gcs_address, + self._node_id, + ) + break + except RuntimeError as e: + logger.info(f"Failed to get node info {e}") + if time.monotonic() - start_time > raylet_start_wait_time_s: + raise Exception( + "The current node timed out during startup. This " + "could happen because some of the raylet failed to " + "startup or the GCS has become overloaded." + ) + # Use node info to update port + if self._ray_params.node_manager_port == 0: + self._ray_params.node_manager_port = node_info["node_manager_port"] + + if connect_only: + # Fetch node info to get labels. + node_info = ray._private.services.get_node( + self.gcs_address, + self._node_id, + ) + # Set node labels from GCS if provided at node init. + self._node_labels = node_info.get("labels", {}) + + # Makes sure the Node object has valid addresses after setup. + self.validate_ip_port(self.address) + self.validate_ip_port(self.gcs_address) + + if not connect_only: + self._record_stats() + + def check_persisted_session_name(self): + if self._ray_params.external_addresses is None: + return None + self._redis_address = self._ray_params.external_addresses[0] + redis_ip_address, redis_port, enable_redis_ssl = get_address( + self._redis_address, + ) + # Address is ip:port or redis://ip:port + if int(redis_port) < 0: + raise ValueError( + f"Invalid Redis port provided: {redis_port}." + "The port must be a non-negative integer." + ) + + return get_session_key_from_storage( + redis_ip_address, + int(redis_port), + self._ray_params.redis_username, + self._ray_params.redis_password, + enable_redis_ssl, + serialize_config(self._config), + b"session_name", + ) + + @staticmethod + def validate_ip_port(ip_port): + """Validates the address is in the ip:port format""" + parts = parse_address(ip_port) + if parts is None: + raise ValueError(f"Port is not specified for address {ip_port}") + try: + _ = int(parts[1]) + except ValueError: + raise ValueError( + f"Unable to parse port number from {parts[1]} (full address = {ip_port})" + ) + + def check_version_info(self): + """Check if the Python and Ray version of this process matches that in GCS. + + This will be used to detect if workers or drivers are started using + different versions of Python, or Ray. + + Raises: + Exception: An exception is raised if there is a version mismatch. + """ + import ray._common.usage.usage_lib as ray_usage_lib + + cluster_metadata = ray_usage_lib.get_cluster_metadata(self.get_gcs_client()) + if cluster_metadata is None: + cluster_metadata = ray_usage_lib.get_cluster_metadata(self.get_gcs_client()) + + if not cluster_metadata: + return + node_ip_address = ray._private.services.get_node_ip_address() + ray._private.utils.check_version_info( + cluster_metadata, f"node {node_ip_address}" + ) + + def _register_shutdown_hooks(self): + # Register the atexit handler. In this case, we shouldn't call sys.exit + # as we're already in the exit procedure. + def atexit_handler(*args): + self.kill_all_processes(check_alive=False, allow_graceful=True) + + atexit.register(atexit_handler) + + # Register the handler to be called if we get a SIGTERM. + # In this case, we want to exit with an error code (1) after + # cleaning up child processes. + def sigterm_handler(signum, frame): + self.kill_all_processes(check_alive=False, allow_graceful=True) + sys.exit(1) + + ray._private.utils.set_sigterm_handler(sigterm_handler) + + def _init_temp(self): + # Create a dictionary to store temp file index. + self._incremental_dict = collections.defaultdict(lambda: 0) + + if self.head: + self._ray_params.update_if_absent( + temp_dir=ray._common.utils.get_ray_temp_dir() + ) + self._temp_dir = self._ray_params.temp_dir + else: + if self._ray_params.temp_dir is None: + assert not self._default_worker + temp_dir = ray._private.utils.internal_kv_get_with_retry( + self.get_gcs_client(), + "temp_dir", + ray_constants.KV_NAMESPACE_SESSION, + num_retries=ray_constants.NUM_REDIS_GET_RETRIES, + ) + self._temp_dir = ray._common.utils.decode(temp_dir) + else: + self._temp_dir = self._ray_params.temp_dir + + try_to_create_directory(self._temp_dir) + + if self.head: + self._session_dir = os.path.join(self._temp_dir, self._session_name) + else: + if self._temp_dir is None or self._session_name is None: + assert not self._default_worker + session_dir = ray._private.utils.internal_kv_get_with_retry( + self.get_gcs_client(), + "session_dir", + ray_constants.KV_NAMESPACE_SESSION, + num_retries=ray_constants.NUM_REDIS_GET_RETRIES, + ) + self._session_dir = ray._common.utils.decode(session_dir) + else: + self._session_dir = os.path.join(self._temp_dir, self._session_name) + session_symlink = os.path.join(self._temp_dir, ray_constants.SESSION_LATEST) + + # Send a warning message if the session exists. + try_to_create_directory(self._session_dir) + try_to_symlink(session_symlink, self._session_dir) + # Create a directory to be used for socket files. + self._sockets_dir = os.path.join(self._session_dir, "sockets") + try_to_create_directory(self._sockets_dir) + # Create a directory to be used for process log files. + self._logs_dir = os.path.join(self._session_dir, "logs") + try_to_create_directory(self._logs_dir) + old_logs_dir = os.path.join(self._logs_dir, "old") + try_to_create_directory(old_logs_dir) + # Create a directory to be used for runtime environment. + self._runtime_env_dir = os.path.join( + self._session_dir, self._ray_params.runtime_env_dir_name + ) + try_to_create_directory(self._runtime_env_dir) + # Create a symlink to the libtpu tpu_logs directory if it exists. + user_temp_dir = ray._common.utils.get_user_temp_dir() + tpu_log_dir = f"{user_temp_dir}/tpu_logs" + if os.path.isdir(tpu_log_dir): + tpu_logs_symlink = os.path.join(self._logs_dir, "tpu_logs") + try_to_symlink(tpu_logs_symlink, tpu_log_dir) + + def get_resource_and_label_spec(self): + """Resolve and return the current ResourceAndLabelSpec for the node.""" + if not self._resource_and_label_spec: + self._resource_and_label_spec = ResourceAndLabelSpec( + self._ray_params.num_cpus, + self._ray_params.num_gpus, + self._ray_params.memory, + self._ray_params.object_store_memory, + self._ray_params.resources, + self._ray_params.labels, + ).resolve(is_head=self.head, node_ip_address=self.node_ip_address) + return self._resource_and_label_spec + + @property + def node_id(self): + """Get the node ID.""" + return self._node_id + + @property + def session_name(self): + """Get the current Ray session name.""" + return self._session_name + + @property + def node_ip_address(self): + """Get the IP address of this node.""" + return self._node_ip_address + + @property + def address(self): + """Get the address for bootstrapping, e.g. the address to pass to + `ray start` or `ray.init()` to start worker nodes, that has been + converted to ip:port format. + """ + return self._gcs_address + + @property + def gcs_address(self): + """Get the gcs address.""" + assert self._gcs_address is not None, "Gcs address is not set" + return self._gcs_address + + @property + def redis_address(self): + """Get the cluster Redis address.""" + return self._redis_address + + @property + def redis_username(self): + """Get the cluster Redis username.""" + return self._ray_params.redis_username + + @property + def redis_password(self): + """Get the cluster Redis password.""" + return self._ray_params.redis_password + + @property + def plasma_store_socket_name(self): + """Get the node's plasma store socket name.""" + return self._plasma_store_socket_name + + @property + def unique_id(self): + """Get a unique identifier for this node.""" + return f"{self.node_ip_address}:{self._plasma_store_socket_name}" + + @property + def webui_url(self): + """Get the cluster's web UI url.""" + return self._webui_url + + @property + def raylet_socket_name(self): + """Get the node's raylet socket name.""" + return self._raylet_socket_name + + @property + def node_manager_port(self): + """Get the node manager's port.""" + return self._ray_params.node_manager_port + + @property + def metrics_export_port(self): + """Get the port that exposes metrics""" + return self._metrics_export_port + + @property + def runtime_env_agent_port(self): + """Get the port that exposes runtime env agent as http""" + return self._runtime_env_agent_port + + @property + def runtime_env_agent_address(self): + """Get the address that exposes runtime env agent as http""" + return f"http://{build_address(self._node_ip_address, self._runtime_env_agent_port)}" + + @property + def dashboard_agent_listen_port(self): + """Get the dashboard agent's listen port""" + return self._dashboard_agent_listen_port + + @property + def logging_config(self): + """Get the logging config of the current node.""" + return { + "log_rotation_max_bytes": self.max_bytes, + "log_rotation_backup_count": self.backup_count, + } + + @property + def address_info(self): + """Get a dictionary of addresses.""" + return { + "node_ip_address": self._node_ip_address, + "redis_address": self.redis_address, + "object_store_address": self._plasma_store_socket_name, + "raylet_socket_name": self._raylet_socket_name, + "webui_url": self._webui_url, + "session_dir": self._session_dir, + "metrics_export_port": self._metrics_export_port, + "gcs_address": self.gcs_address, + "address": self.address, + "dashboard_agent_listen_port": self.dashboard_agent_listen_port, + } + + @property + def node_labels(self): + """Get the node labels.""" + return self._node_labels + + def is_head(self): + return self.head + + def get_gcs_client(self): + if self._gcs_client is None: + self._init_gcs_client() + return self._gcs_client + + def _init_gcs_client(self): + if self.head: + gcs_process = self.all_processes[ray_constants.PROCESS_TYPE_GCS_SERVER][ + 0 + ].process + else: + gcs_process = None + + # TODO(ryw) instead of create a new GcsClient, wrap the one from + # CoreWorkerProcess to save a grpc channel. + for _ in range(ray_constants.NUM_REDIS_GET_RETRIES): + gcs_address = None + last_ex = None + try: + gcs_address = self.gcs_address + client = GcsClient( + address=gcs_address, + cluster_id=self._ray_params.cluster_id, # Hex string + ) + self.cluster_id = client.cluster_id + if self.head: + # Send a simple request to make sure GCS is alive + # if it's a head node. + client.internal_kv_get(b"dummy", None) + self._gcs_client = client + break + except Exception: + if gcs_process is not None and gcs_process.poll() is not None: + # GCS has exited. + break + last_ex = traceback.format_exc() + logger.debug(f"Connecting to GCS: {last_ex}") + time.sleep(1) + + if self._gcs_client is None: + if hasattr(self, "_logs_dir"): + with open(os.path.join(self._logs_dir, "gcs_server.err")) as err: + # Use " C " or " E " to exclude the stacktrace. + # This should work for most cases, especitally + # it's when GCS is starting. Only display last 10 lines of logs. + errors = [e for e in err.readlines() if " C " in e or " E " in e][ + -10: + ] + error_msg = "\n" + "".join(errors) + "\n" + raise RuntimeError( + f"Failed to {'start' if self.head else 'connect to'} GCS. " + f" Last {len(errors)} lines of error files:" + f"{error_msg}." + f"Please check {os.path.join(self._logs_dir, 'gcs_server.out')}" + f" for details. Last connection error: {last_ex}" + ) + else: + raise RuntimeError( + f"Failed to {'start' if self.head else 'connect to'} GCS. Last " + f"connection error: {last_ex}" + ) + + ray.experimental.internal_kv._initialize_internal_kv(self._gcs_client) + + def get_temp_dir_path(self): + """Get the path of the temporary directory.""" + return self._temp_dir + + def get_runtime_env_dir_path(self): + """Get the path of the runtime env.""" + return self._runtime_env_dir + + def get_session_dir_path(self): + """Get the path of the session directory.""" + return self._session_dir + + def get_logs_dir_path(self): + """Get the path of the log files directory.""" + return self._logs_dir + + def get_sockets_dir_path(self): + """Get the path of the sockets directory.""" + return self._sockets_dir + + def _make_inc_temp( + self, suffix: str = "", prefix: str = "", directory_name: Optional[str] = None + ): + """Return an incremental temporary file name. The file is not created. + + Args: + suffix: The suffix of the temp file. + prefix: The prefix of the temp file. + directory_name (str) : The base directory of the temp file. + + Returns: + A string of file name. If there existing a file having + the same name, the returned name will look like + "{directory_name}/{prefix}.{unique_index}{suffix}" + """ + if directory_name is None: + directory_name = ray._common.utils.get_ray_temp_dir() + directory_name = os.path.expanduser(directory_name) + index = self._incremental_dict[suffix, prefix, directory_name] + # `tempfile.TMP_MAX` could be extremely large, + # so using `range` in Python2.x should be avoided. + while index < tempfile.TMP_MAX: + if index == 0: + filename = os.path.join(directory_name, prefix + suffix) + else: + filename = os.path.join( + directory_name, prefix + "." + str(index) + suffix + ) + index += 1 + if not os.path.exists(filename): + # Save the index. + self._incremental_dict[suffix, prefix, directory_name] = index + return filename + + raise FileExistsError(errno.EEXIST, "No usable temporary filename found") + + def should_redirect_logs(self): + redirect_output = self._ray_params.redirect_output + if redirect_output is None: + # Fall back to stderr redirect environment variable. + redirect_output = ( + os.environ.get( + ray_constants.LOGGING_REDIRECT_STDERR_ENVIRONMENT_VARIABLE + ) + != "1" + ) + return redirect_output + + # TODO(hjiang): Re-implement the logic in C++, and expose via cython. + def get_log_file_names( + self, + name: str, + unique: bool = False, + create_out: bool = True, + create_err: bool = True, + ) -> Tuple[Optional[str], Optional[str]]: + """Get filename to dump logs for stdout and stderr, with no files opened. + If output redirection has been disabled, no files will + be opened and `(None, None)` will be returned. + + Args: + name: descriptive string for this log file. + unique: if true, a counter will be attached to `name` to + ensure the returned filename is not already used. + create_out: if True, create a .out file. + create_err: if True, create a .err file. + + Returns: + A tuple of two file handles for redirecting optional (stdout, stderr), + or `(None, None)` if output redirection is disabled. + """ + if not self.should_redirect_logs(): + return None, None + + log_stdout = None + log_stderr = None + + if create_out: + log_stdout = self._get_log_file_name(name, "out", unique=unique) + if create_err: + log_stderr = self._get_log_file_name(name, "err", unique=unique) + return log_stdout, log_stderr + + def get_log_file_handles( + self, + name: str, + unique: bool = False, + create_out: bool = True, + create_err: bool = True, + ) -> Tuple[Optional[IO[AnyStr]], Optional[IO[AnyStr]]]: + """Open log files with partially randomized filenames, returning the + file handles. If output redirection has been disabled, no files will + be opened and `(None, None)` will be returned. + + Args: + name: descriptive string for this log file. + unique: if true, a counter will be attached to `name` to + ensure the returned filename is not already used. + create_out: if True, create a .out file. + create_err: if True, create a .err file. + + Returns: + A tuple of two file handles for redirecting optional (stdout, stderr), + or `(None, None)` if output redirection is disabled. + """ + log_stdout_fname, log_stderr_fname = self.get_log_file_names( + name, unique=unique, create_out=create_out, create_err=create_err + ) + log_stdout = None if log_stdout_fname is None else open_log(log_stdout_fname) + log_stderr = None if log_stderr_fname is None else open_log(log_stderr_fname) + return log_stdout, log_stderr + + def _get_log_file_name( + self, + name: str, + suffix: str, + unique: bool = False, + ) -> str: + """Generate partially randomized filenames for log files. + + Args: + name: descriptive string for this log file. + suffix: suffix of the file. Usually it is .out of .err. + unique: if true, a counter will be attached to `name` to + ensure the returned filename is not already used. + + Returns: + A tuple of two file names for redirecting (stdout, stderr). + """ + # strip if the suffix is something like .out. + suffix = suffix.strip(".") + + if unique: + filename = self._make_inc_temp( + suffix=f".{suffix}", prefix=name, directory_name=self._logs_dir + ) + else: + filename = os.path.join(self._logs_dir, f"{name}.{suffix}") + return filename + + def _get_unused_port(self, allocated_ports=None): + if allocated_ports is None: + allocated_ports = set() + + s = socket.socket( + socket.AF_INET6 if is_ipv6(self._node_ip_address) else socket.AF_INET, + socket.SOCK_STREAM, + ) + s.bind(("", 0)) + port = s.getsockname()[1] + + # Try to generate a port that is far above the 'next available' one. + # This solves issue #8254 where GRPC fails because the port assigned + # from this method has been used by a different process. + for _ in range(ray_constants.NUM_PORT_RETRIES): + new_port = random.randint(port, 65535) + if new_port in allocated_ports: + # This port is allocated for other usage already, + # so we shouldn't use it even if it's not in use right now. + continue + new_s = socket.socket( + socket.AF_INET6 if is_ipv6(self._node_ip_address) else socket.AF_INET, + socket.SOCK_STREAM, + ) + try: + new_s.bind(("", new_port)) + except OSError: + new_s.close() + continue + s.close() + new_s.close() + return new_port + logger.error("Unable to succeed in selecting a random port.") + s.close() + return port + + def _prepare_socket_file(self, socket_path: str, default_prefix: str): + """Prepare the socket file for raylet and plasma. + + This method helps to prepare a socket file. + 1. Make the directory if the directory does not exist. + 2. If the socket file exists, do nothing (this just means we aren't the + first worker on the node). + + Args: + socket_path: the socket file to prepare. + """ + result = socket_path + if sys.platform == "win32": + if socket_path is None: + result = ( + f"tcp://{build_address(self._localhost, self._get_unused_port())}" + ) + else: + if socket_path is None: + result = self._make_inc_temp( + prefix=default_prefix, directory_name=self._sockets_dir + ) + else: + try_to_create_directory(os.path.dirname(socket_path)) + + validate_socket_filepath(result.split("://", 1)[-1]) + return result + + def _get_cached_port( + self, port_name: str, default_port: Optional[int] = None + ) -> int: + """Get a port number from a cache on this node. + + Different driver processes on a node should use the same ports for + some purposes, e.g. exporting metrics. This method returns a port + number for the given port name and caches it in a file. If the + port isn't already cached, an unused port is generated and cached. + + Args: + port_name: The name of the port, e.g. metrics_export_port. + default_port: The port to return and cache if no port has already been + cached for the given port_name. If None, an unused port is generated + and cached. + + Returns: + int: The port number. + """ + file_path = os.path.join(self.get_session_dir_path(), "ports_by_node.json") + + # Make sure only the ports in RAY_CACHED_PORTS are cached. + assert port_name in ray_constants.RAY_ALLOWED_CACHED_PORTS + + # Maps a Node.unique_id to a dict that maps port names to port numbers. + ports_by_node: Dict[str, Dict[str, int]] = defaultdict(dict) + + with FileLock(file_path + ".lock"): + if not os.path.exists(file_path): + with open(file_path, "w") as f: + json.dump({}, f) + + with open(file_path, "r") as f: + ports_by_node.update(json.load(f)) + + if ( + self.unique_id in ports_by_node + and port_name in ports_by_node[self.unique_id] + ): + # The port has already been cached at this node, so use it. + port = int(ports_by_node[self.unique_id][port_name]) + else: + # Pick a new port to use and cache it at this node. + allocated_ports = set(ports_by_node[self.unique_id].values()) + + if default_port is not None and default_port in allocated_ports: + # The default port is already in use, so don't use it. + default_port = None + + port = default_port or self._get_unused_port(allocated_ports) + + ports_by_node[self.unique_id][port_name] = port + with open(file_path, "w") as f: + json.dump(ports_by_node, f) + + return port + + def _wait_and_get_for_node_address(self, timeout_s: int = 60) -> str: + """Wait until the RAY_NODE_IP_FILENAME file is avialable. + + RAY_NODE_IP_FILENAME is created when a ray instance is started. + + Args: + timeout_s: If the ip address is not found within this + timeout, it will raise ValueError. + Returns: + The node_ip_address of the current session if it finds it + within timeout_s. + """ + for i in range(timeout_s): + node_ip_address = ray._private.services.get_cached_node_ip_address( + self.get_session_dir_path() + ) + + if node_ip_address is not None: + return node_ip_address + + time.sleep(1) + if i % 10 == 0: + logger.info( + f"Can't find a `{ray_constants.RAY_NODE_IP_FILENAME}` " + f"file from {self.get_session_dir_path()}. " + "Have you started Ray instance using " + "`ray start` or `ray.init`?" + ) + + raise ValueError( + f"Can't find a `{ray_constants.RAY_NODE_IP_FILENAME}` " + f"file from {self.get_session_dir_path()}. " + f"for {timeout_s} seconds. " + "A ray instance hasn't started. " + "Did you do `ray start` or `ray.init` on this host?" + ) + + def start_reaper_process(self): + """ + Start the reaper process. + + This must be the first process spawned and should only be called when + ray processes should be cleaned up if this process dies. + """ + assert ( + not self.kernel_fate_share + ), "a reaper should not be used with kernel fate-sharing" + process_info = ray._private.services.start_reaper(fate_share=False) + assert ray_constants.PROCESS_TYPE_REAPER not in self.all_processes + if process_info is not None: + self.all_processes[ray_constants.PROCESS_TYPE_REAPER] = [ + process_info, + ] + + def start_log_monitor(self): + """Start the log monitor.""" + stdout_log_fname, stderr_log_fname = self.get_log_file_names( + "log_monitor", unique=True, create_out=True, create_err=True + ) + process_info = ray._private.services.start_log_monitor( + self.get_session_dir_path(), + self._logs_dir, + self.gcs_address, + fate_share=self.kernel_fate_share, + max_bytes=self.max_bytes, + backup_count=self.backup_count, + stdout_filepath=stdout_log_fname, + stderr_filepath=stderr_log_fname, + ) + assert ray_constants.PROCESS_TYPE_LOG_MONITOR not in self.all_processes + self.all_processes[ray_constants.PROCESS_TYPE_LOG_MONITOR] = [ + process_info, + ] + + def start_api_server( + self, *, include_dashboard: Optional[bool], raise_on_failure: bool + ): + """Start the dashboard. + + Args: + include_dashboard: If true, this will load all dashboard-related modules + when starting the API server. Otherwise, it will only + start the modules that are not relevant to the dashboard. + raise_on_failure: If true, this will raise an exception + if we fail to start the API server. Otherwise it will print + a warning if we fail to start the API server. + """ + stdout_log_fname, stderr_log_fname = self.get_log_file_names( + "dashboard", unique=True, create_out=True, create_err=True + ) + self._webui_url, process_info = ray._private.services.start_api_server( + include_dashboard, + raise_on_failure, + self._ray_params.dashboard_host, + self.gcs_address, + self.cluster_id.hex(), + self._node_ip_address, + self._temp_dir, + self._logs_dir, + self._session_dir, + port=self._ray_params.dashboard_port, + fate_share=self.kernel_fate_share, + max_bytes=self.max_bytes, + backup_count=self.backup_count, + stdout_filepath=stdout_log_fname, + stderr_filepath=stderr_log_fname, + ) + assert ray_constants.PROCESS_TYPE_DASHBOARD not in self.all_processes + if process_info is not None: + self.all_processes[ray_constants.PROCESS_TYPE_DASHBOARD] = [ + process_info, + ] + self.get_gcs_client().internal_kv_put( + b"webui:url", + self._webui_url.encode(), + True, + ray_constants.KV_NAMESPACE_DASHBOARD, + ) + + def start_gcs_server(self): + """Start the gcs server.""" + gcs_server_port = self._ray_params.gcs_server_port + assert gcs_server_port > 0 + assert self._gcs_address is None, "GCS server is already running." + assert self._gcs_client is None, "GCS client is already connected." + + stdout_log_fname, stderr_log_fname = self.get_log_file_names( + "gcs_server", unique=True, create_out=True, create_err=True + ) + process_info = ray._private.services.start_gcs_server( + self.redis_address, + log_dir=self._logs_dir, + stdout_filepath=stdout_log_fname, + stderr_filepath=stderr_log_fname, + session_name=self.session_name, + redis_username=self._ray_params.redis_username, + redis_password=self._ray_params.redis_password, + config=self._config, + fate_share=self.kernel_fate_share, + gcs_server_port=gcs_server_port, + metrics_agent_port=self._ray_params.metrics_agent_port, + node_ip_address=self._node_ip_address, + ) + assert ray_constants.PROCESS_TYPE_GCS_SERVER not in self.all_processes + self.all_processes[ray_constants.PROCESS_TYPE_GCS_SERVER] = [ + process_info, + ] + # Connecting via non-localhost address may be blocked by firewall rule, + # e.g. https://github.com/ray-project/ray/issues/15780 + # TODO(mwtian): figure out a way to use 127.0.0.1 for local connection + # when possible. + self._gcs_address = build_address(self._node_ip_address, gcs_server_port) + + def start_raylet( + self, + plasma_directory: str, + fallback_directory: str, + object_store_memory: int, + use_valgrind: bool = False, + use_profiler: bool = False, + ): + """Start the raylet. + + Args: + use_valgrind: True if we should start the process in + valgrind. + use_profiler: True if we should start the process in the + valgrind profiler. + """ + raylet_stdout_filepath, raylet_stderr_filepath = self.get_log_file_names( + ray_constants.PROCESS_TYPE_RAYLET, + unique=True, + create_out=True, + create_err=True, + ) + ( + dashboard_agent_stdout_filepath, + dashboard_agent_stderr_filepath, + ) = self.get_log_file_names( + ray_constants.PROCESS_TYPE_DASHBOARD_AGENT, + unique=True, + create_out=True, + create_err=True, + ) + ( + runtime_env_agent_stdout_filepath, + runtime_env_agent_stderr_filepath, + ) = self.get_log_file_names( + ray_constants.PROCESS_TYPE_RUNTIME_ENV_AGENT, + unique=True, + create_out=True, + create_err=True, + ) + + self.resource_isolation_config.add_system_pids( + self._get_system_processes_for_resource_isolation() + ) + + process_info = ray._private.services.start_raylet( + self.redis_address, + self.gcs_address, + self._node_id, + self._node_ip_address, + self._ray_params.node_manager_port, + self._raylet_socket_name, + self._plasma_store_socket_name, + self.cluster_id.hex(), + self._ray_params.worker_path, + self._ray_params.setup_worker_path, + self._temp_dir, + self._session_dir, + self._runtime_env_dir, + self._logs_dir, + self.get_resource_and_label_spec(), + plasma_directory, + fallback_directory, + object_store_memory, + self.session_name, + is_head_node=self.is_head(), + min_worker_port=self._ray_params.min_worker_port, + max_worker_port=self._ray_params.max_worker_port, + worker_port_list=self._ray_params.worker_port_list, + object_manager_port=self._ray_params.object_manager_port, + redis_username=self._ray_params.redis_username, + redis_password=self._ray_params.redis_password, + metrics_agent_port=self._ray_params.metrics_agent_port, + runtime_env_agent_port=self._ray_params.runtime_env_agent_port, + metrics_export_port=self._metrics_export_port, + dashboard_agent_listen_port=self._ray_params.dashboard_agent_listen_port, + use_valgrind=use_valgrind, + use_profiler=use_profiler, + raylet_stdout_filepath=raylet_stdout_filepath, + raylet_stderr_filepath=raylet_stderr_filepath, + dashboard_agent_stdout_filepath=dashboard_agent_stdout_filepath, + dashboard_agent_stderr_filepath=dashboard_agent_stderr_filepath, + runtime_env_agent_stdout_filepath=runtime_env_agent_stdout_filepath, + runtime_env_agent_stderr_filepath=runtime_env_agent_stderr_filepath, + huge_pages=self._ray_params.huge_pages, + fate_share=self.kernel_fate_share, + socket_to_use=None, + max_bytes=self.max_bytes, + backup_count=self.backup_count, + ray_debugger_external=self._ray_params.ray_debugger_external, + env_updates=self._ray_params.env_vars, + node_name=self._ray_params.node_name, + webui=self._webui_url, + resource_isolation_config=self.resource_isolation_config, + ) + assert ray_constants.PROCESS_TYPE_RAYLET not in self.all_processes + self.all_processes[ray_constants.PROCESS_TYPE_RAYLET] = [process_info] + + def start_monitor(self): + """Start the monitor. + + Autoscaling output goes to these monitor.err/out files, and + any modification to these files may break existing + cluster launching commands. + """ + from ray.autoscaler.v2.utils import is_autoscaler_v2 + + stdout_log_fname, stderr_log_fname = self.get_log_file_names( + "monitor", unique=True, create_out=True, create_err=True + ) + process_info = ray._private.services.start_monitor( + self.gcs_address, + self._logs_dir, + stdout_filepath=stdout_log_fname, + stderr_filepath=stderr_log_fname, + autoscaling_config=self._ray_params.autoscaling_config, + fate_share=self.kernel_fate_share, + max_bytes=self.max_bytes, + backup_count=self.backup_count, + monitor_ip=self._node_ip_address, + autoscaler_v2=is_autoscaler_v2(fetch_from_server=True), + ) + assert ray_constants.PROCESS_TYPE_MONITOR not in self.all_processes + self.all_processes[ray_constants.PROCESS_TYPE_MONITOR] = [process_info] + + def start_ray_client_server(self): + """Start the ray client server process.""" + stdout_file, stderr_file = self.get_log_file_handles( + "ray_client_server", unique=True + ) + process_info = ray._private.services.start_ray_client_server( + self.address, + self._node_ip_address, + self._ray_params.ray_client_server_port, + stdout_file=stdout_file, + stderr_file=stderr_file, + redis_username=self._ray_params.redis_username, + redis_password=self._ray_params.redis_password, + fate_share=self.kernel_fate_share, + runtime_env_agent_address=self.runtime_env_agent_address, + ) + assert ray_constants.PROCESS_TYPE_RAY_CLIENT_SERVER not in self.all_processes + self.all_processes[ray_constants.PROCESS_TYPE_RAY_CLIENT_SERVER] = [ + process_info + ] + + def _write_cluster_info_to_kv(self): + """Write the cluster metadata to GCS. + Cluster metadata is always recorded, but they are + not reported unless usage report is enabled. + Check `usage_stats_head.py` for more details. + """ + # Make sure the cluster metadata wasn't reported before. + import ray._common.usage.usage_lib as ray_usage_lib + + ray_usage_lib.put_cluster_metadata( + self.get_gcs_client(), ray_init_cluster=self.ray_init_cluster + ) + # Make sure GCS is up. + added = self.get_gcs_client().internal_kv_put( + b"session_name", + self._session_name.encode(), + False, + ray_constants.KV_NAMESPACE_SESSION, + ) + if not added: + curr_val = self.get_gcs_client().internal_kv_get( + b"session_name", ray_constants.KV_NAMESPACE_SESSION + ) + assert curr_val == self._session_name.encode("utf-8"), ( + f"Session name {self._session_name} does not match " + f"persisted value {curr_val}. Perhaps there was an " + f"error connecting to Redis." + ) + + self.get_gcs_client().internal_kv_put( + b"session_dir", + self._session_dir.encode(), + True, + ray_constants.KV_NAMESPACE_SESSION, + ) + self.get_gcs_client().internal_kv_put( + b"temp_dir", + self._temp_dir.encode(), + True, + ray_constants.KV_NAMESPACE_SESSION, + ) + # Add tracing_startup_hook to redis / internal kv manually + # since internal kv is not yet initialized. + if self._ray_params.tracing_startup_hook: + self.get_gcs_client().internal_kv_put( + b"tracing_startup_hook", + self._ray_params.tracing_startup_hook.encode(), + True, + ray_constants.KV_NAMESPACE_TRACING, + ) + + def start_head_processes(self): + """Start head processes on the node.""" + logger.debug( + f"Process STDOUT and STDERR is being " f"redirected to {self._logs_dir}." + ) + assert self._gcs_address is None + assert self._gcs_client is None + + self.start_gcs_server() + assert self.get_gcs_client() is not None + self._write_cluster_info_to_kv() + + if not self._ray_params.no_monitor: + self.start_monitor() + + if self._ray_params.ray_client_server_port: + self.start_ray_client_server() + + if self._ray_params.include_dashboard is None: + # Default + raise_on_api_server_failure = False + else: + raise_on_api_server_failure = self._ray_params.include_dashboard + + self.start_api_server( + include_dashboard=self._ray_params.include_dashboard, + raise_on_failure=raise_on_api_server_failure, + ) + + def start_ray_processes(self): + """Start all of the processes on the node.""" + logger.debug( + f"Process STDOUT and STDERR is being " f"redirected to {self._logs_dir}." + ) + + if not self.head: + # Get the system config from GCS first if this is a non-head node. + gcs_options = ray._raylet.GcsClientOptions.create( + self.gcs_address, + self.cluster_id.hex(), + allow_cluster_id_nil=False, + fetch_cluster_id_if_nil=False, + ) + global_state = ray._private.state.GlobalState() + global_state._initialize_global_state(gcs_options) + new_config = global_state.get_system_config() + assert self._config.items() <= new_config.items(), ( + "The system config from GCS is not a superset of the local" + " system config. There might be a configuration inconsistency" + " issue between the head node and non-head nodes." + f" Local system config: {self._config}," + f" GCS system config: {new_config}" + ) + self._config = new_config + + # Make sure we don't call `determine_plasma_store_config` multiple + # times to avoid printing multiple warnings. + resource_and_label_spec = self.get_resource_and_label_spec() + if resource_and_label_spec.labels.get( + ray._raylet.RAY_NODE_ACCELERATOR_TYPE_KEY + ): + from ray._common.usage import usage_lib + + usage_lib.record_hardware_usage( + resource_and_label_spec.labels.get( + ray._raylet.RAY_NODE_ACCELERATOR_TYPE_KEY + ) + ) + + ( + plasma_directory, + fallback_directory, + object_store_memory, + ) = ray._private.services.determine_plasma_store_config( + resource_and_label_spec.object_store_memory, + self._temp_dir, + plasma_directory=self._ray_params.plasma_directory, + fallback_directory=self._fallback_directory, + huge_pages=self._ray_params.huge_pages, + ) + + # add plasma store memory to the total system reserved memory + if self.resource_isolation_config.is_enabled(): + self.resource_isolation_config.add_object_store_memory(object_store_memory) + + if self._ray_params.include_log_monitor: + self.start_log_monitor() + + self.start_raylet(plasma_directory, fallback_directory, object_store_memory) + + def _get_system_processes_for_resource_isolation(self) -> str: + """Returns a list of system processes that will be isolated by raylet. + + NOTE: If a new system process is started before the raylet starts up, it needs to be + added to self.all_processes so it can be moved into the raylet's managed cgroup + hierarchy. + """ + system_process_pids = [ + str(p[0].process.pid) for p in self.all_processes.values() + ] + + # If the dashboard api server was started on the head node, then include all of the api server's + # child processes. + if ray_constants.PROCESS_TYPE_DASHBOARD in self.all_processes: + dashboard_pid = self.all_processes[ray_constants.PROCESS_TYPE_DASHBOARD][ + 0 + ].process.pid + dashboard_process = psutil.Process(dashboard_pid) + system_process_pids += [str(p.pid) for p in dashboard_process.children()] + + return ",".join(system_process_pids) + + def _kill_process_type( + self, + process_type, + allow_graceful: bool = False, + check_alive: bool = True, + wait: bool = False, + ): + """Kill a process of a given type. + + If the process type is PROCESS_TYPE_REDIS_SERVER, then we will kill all + of the Redis servers. + + If the process was started in valgrind, then we will raise an exception + if the process has a non-zero exit code. + + Args: + process_type: The type of the process to kill. + allow_graceful: Send a SIGTERM first and give the process + time to exit gracefully. If that doesn't work, then use + SIGKILL. We usually want to do this outside of tests. + check_alive: If true, then we expect the process to be alive + and will raise an exception if the process is already dead. + wait: If true, then this method will not return until the + process in question has exited. + + Raises: + This process raises an exception in the following cases: + 1. The process had already died and check_alive is true. + 2. The process had been started in valgrind and had a non-zero + exit code. + """ + + # Ensure thread safety + with self.removal_lock: + self._kill_process_impl( + process_type, + allow_graceful=allow_graceful, + check_alive=check_alive, + wait=wait, + ) + + def _kill_process_impl( + self, process_type, allow_graceful=False, check_alive=True, wait=False + ): + """See `_kill_process_type`.""" + if process_type not in self.all_processes: + return + process_infos = self.all_processes[process_type] + if process_type != ray_constants.PROCESS_TYPE_REDIS_SERVER: + assert len(process_infos) == 1 + for process_info in process_infos: + process = process_info.process + # Handle the case where the process has already exited. + if process.poll() is not None: + if check_alive: + raise RuntimeError( + "Attempting to kill a process of type " + f"'{process_type}', but this process is already dead." + ) + else: + continue + + if process_info.use_valgrind: + process.terminate() + process.wait() + if process.returncode != 0: + message = ( + "Valgrind detected some errors in process of " + f"type {process_type}. Error code {process.returncode}." + ) + if process_info.stdout_file is not None: + with open(process_info.stdout_file, "r") as f: + message += "\nPROCESS STDOUT:\n" + f.read() + if process_info.stderr_file is not None: + with open(process_info.stderr_file, "r") as f: + message += "\nPROCESS STDERR:\n" + f.read() + raise RuntimeError(message) + continue + + if process_info.use_valgrind_profiler: + # Give process signal to write profiler data. + os.kill(process.pid, signal.SIGINT) + # Wait for profiling data to be written. + time.sleep(0.1) + + if allow_graceful: + process.terminate() + # Allow the process one second to exit gracefully. + timeout_seconds = 1 + try: + process.wait(timeout_seconds) + except subprocess.TimeoutExpired: + pass + + # If the process did not exit, force kill it. + if process.poll() is None: + process.kill() + # The reason we usually don't call process.wait() here is that + # there's some chance we'd end up waiting a really long time. + if wait: + process.wait() + + del self.all_processes[process_type] + + def kill_redis(self, check_alive: bool = True): + """Kill the Redis servers. + + Args: + check_alive: Raise an exception if any of the processes + were already dead. + """ + self._kill_process_type( + ray_constants.PROCESS_TYPE_REDIS_SERVER, check_alive=check_alive + ) + + def kill_raylet(self, check_alive: bool = True): + """Kill the raylet. + + Args: + check_alive: Raise an exception if the process was already + dead. + """ + self._kill_process_type( + ray_constants.PROCESS_TYPE_RAYLET, check_alive=check_alive + ) + + def kill_log_monitor(self, check_alive: bool = True): + """Kill the log monitor. + + Args: + check_alive: Raise an exception if the process was already + dead. + """ + self._kill_process_type( + ray_constants.PROCESS_TYPE_LOG_MONITOR, check_alive=check_alive + ) + + def kill_dashboard(self, check_alive: bool = True): + """Kill the dashboard. + + Args: + check_alive: Raise an exception if the process was already + dead. + """ + self._kill_process_type( + ray_constants.PROCESS_TYPE_DASHBOARD, check_alive=check_alive + ) + + def kill_monitor(self, check_alive: bool = True): + """Kill the monitor. + + Args: + check_alive: Raise an exception if the process was already + dead. + """ + self._kill_process_type( + ray_constants.PROCESS_TYPE_MONITOR, check_alive=check_alive + ) + + def kill_gcs_server(self, check_alive: bool = True): + """Kill the gcs server. + + Args: + check_alive: Raise an exception if the process was already + dead. + """ + self._kill_process_type( + ray_constants.PROCESS_TYPE_GCS_SERVER, check_alive=check_alive, wait=True + ) + # Clear GCS client and address to indicate no GCS server is running. + self._gcs_address = None + self._gcs_client = None + + def kill_reaper(self, check_alive: bool = True): + """Kill the reaper process. + + Args: + check_alive: Raise an exception if the process was already + dead. + """ + self._kill_process_type( + ray_constants.PROCESS_TYPE_REAPER, check_alive=check_alive + ) + + def kill_all_processes(self, check_alive=True, allow_graceful=False, wait=False): + """Kill all of the processes. + + Note that This is slower than necessary because it calls kill, wait, + kill, wait, ... instead of kill, kill, ..., wait, wait, ... + + Args: + check_alive: Raise an exception if any of the processes were + already dead. + wait: If true, then this method will not return until the + process in question has exited. + """ + # Kill the raylet first. This is important for suppressing errors at + # shutdown because we give the raylet a chance to exit gracefully and + # clean up its child worker processes. If we were to kill the plasma + # store (or Redis) first, that could cause the raylet to exit + # ungracefully, leading to more verbose output from the workers. + if ray_constants.PROCESS_TYPE_RAYLET in self.all_processes: + self._kill_process_type( + ray_constants.PROCESS_TYPE_RAYLET, + check_alive=check_alive, + allow_graceful=allow_graceful, + wait=wait, + ) + + if ray_constants.PROCESS_TYPE_GCS_SERVER in self.all_processes: + self._kill_process_type( + ray_constants.PROCESS_TYPE_GCS_SERVER, + check_alive=check_alive, + allow_graceful=allow_graceful, + wait=wait, + ) + + # We call "list" to copy the keys because we are modifying the + # dictionary while iterating over it. + for process_type in list(self.all_processes.keys()): + # Need to kill the reaper process last in case we die unexpectedly + # while cleaning up. + if process_type != ray_constants.PROCESS_TYPE_REAPER: + self._kill_process_type( + process_type, + check_alive=check_alive, + allow_graceful=allow_graceful, + wait=wait, + ) + + if ray_constants.PROCESS_TYPE_REAPER in self.all_processes: + self._kill_process_type( + ray_constants.PROCESS_TYPE_REAPER, + check_alive=check_alive, + allow_graceful=allow_graceful, + wait=wait, + ) + + def live_processes(self): + """Return a list of the live processes. + + Returns: + A list of the live processes. + """ + result = [] + for process_type, process_infos in self.all_processes.items(): + for process_info in process_infos: + if process_info.process.poll() is None: + result.append((process_type, process_info.process)) + return result + + def dead_processes(self): + """Return a list of the dead processes. + + Note that this ignores processes that have been explicitly killed, + e.g., via a command like node.kill_raylet(). + + Returns: + A list of the dead processes ignoring the ones that have been + explicitly killed. + """ + result = [] + for process_type, process_infos in self.all_processes.items(): + for process_info in process_infos: + if process_info.process.poll() is not None: + result.append((process_type, process_info.process)) + return result + + def any_processes_alive(self): + """Return true if any processes are still alive. + + Returns: + True if any process is still alive. + """ + return any(self.live_processes()) + + def remaining_processes_alive(self): + """Return true if all remaining processes are still alive. + + Note that this ignores processes that have been explicitly killed, + e.g., via a command like node.kill_raylet(). + + Returns: + True if any process that wasn't explicitly killed is still alive. + """ + return not any(self.dead_processes()) + + def destroy_external_storage(self): + object_spilling_config = self._config.get("object_spilling_config", {}) + if object_spilling_config: + object_spilling_config = json.loads(object_spilling_config) + from ray._private import external_storage + + storage = external_storage.setup_external_storage( + object_spilling_config, self._node_id, self._session_name + ) + storage.destroy_external_storage() + + def validate_external_storage(self): + """Make sure we can setup the object spilling external storage.""" + + automatic_spilling_enabled = self._config.get( + "automatic_object_spilling_enabled", True + ) + if not automatic_spilling_enabled: + return + + object_spilling_config = self._object_spilling_config + # Try setting up the storage. + # Configure the proper system config. + # We need to set both ray param's system config and self._config + # because they could've been diverged at this point. + deserialized_config = json.loads(object_spilling_config) + self._ray_params._system_config[ + "object_spilling_config" + ] = object_spilling_config + self._config["object_spilling_config"] = object_spilling_config + + is_external_storage_type_fs = deserialized_config["type"] == "filesystem" + self._ray_params._system_config[ + "is_external_storage_type_fs" + ] = is_external_storage_type_fs + self._config["is_external_storage_type_fs"] = is_external_storage_type_fs + + # Validate external storage usage. + from ray._private import external_storage + + # Node ID is available only after GCS is connected. However, + # validate_external_storage() needs to be called before it to + # be able to validate the configs early. Therefore, we use a + # dummy node ID here and make sure external storage can be set + # up based on the provided config. This storage is destroyed + # right after the validation. + dummy_node_id = ray.NodeID.from_random().hex() + storage = external_storage.setup_external_storage( + deserialized_config, dummy_node_id, self._session_name + ) + storage.destroy_external_storage() + external_storage.reset_external_storage() + + def _get_object_spilling_config(self): + """Consolidate the object spilling config from the ray params, environment + variable, and system config. The object spilling directory specified through + ray params will override the one specified through environment variable and + system config.""" + + object_spilling_directory = self._ray_params.object_spilling_directory + if not object_spilling_directory: + object_spilling_directory = self._config.get( + "object_spilling_directory", "" + ) + + if not object_spilling_directory: + object_spilling_directory = os.environ.get( + "RAY_object_spilling_directory", "" + ) + + if object_spilling_directory: + return json.dumps( + { + "type": "filesystem", + "params": {"directory_path": object_spilling_directory}, + } + ) + + object_spilling_config = self._config.get("object_spilling_config", {}) + if not object_spilling_config: + object_spilling_config = os.environ.get("RAY_object_spilling_config", "") + + # If the config is not specified in ray params, system config or environment + # variable, we fill up the default. + if not object_spilling_config: + object_spilling_config = json.dumps( + {"type": "filesystem", "params": {"directory_path": self._session_dir}} + ) + else: + if not is_in_test(): + logger.warning( + "The object spilling config is specified from an unstable " + "API - system config or environment variable. This is " + "subject to change in the future. You can use the stable " + "API - --object-spilling-directory in ray start or " + "object_spilling_directory in ray.init() to specify the " + "object spilling directory instead. If you need more " + "advanced settings, please open a github issue with the " + "Ray team." + ) + + return object_spilling_config + + def _record_stats(self): + # This is only called when a new node is started. + # Initialize the internal kv so that the metrics can be put + from ray._common.usage.usage_lib import ( + TagKey, + record_extra_usage_tag, + record_hardware_usage, + ) + + if not ray.experimental.internal_kv._internal_kv_initialized(): + ray.experimental.internal_kv._initialize_internal_kv(self.get_gcs_client()) + assert ray.experimental.internal_kv._internal_kv_initialized() + if self.head: + # record head node stats + gcs_storage_type = ( + "redis" if os.environ.get("RAY_REDIS_ADDRESS") is not None else "memory" + ) + record_extra_usage_tag(TagKey.GCS_STORAGE, gcs_storage_type) + cpu_model_name = ray._private.utils.get_current_node_cpu_model_name() + if cpu_model_name: + # CPU model name can be an arbitrary long string + # so we truncate it to the first 50 characters + # to avoid any issues. + record_hardware_usage(cpu_model_name[:50]) diff --git a/lib/python3.12/site-packages/ray/_private/parameter.py b/lib/python3.12/site-packages/ray/_private/parameter.py new file mode 100644 index 0000000000000000000000000000000000000000..78a75ac28bf7d6018180502240268cc0a5ceb9b5 --- /dev/null +++ b/lib/python3.12/site-packages/ray/_private/parameter.py @@ -0,0 +1,474 @@ +import logging +import os +import pathlib +from typing import Dict, List, Optional + +import ray._private.ray_constants as ray_constants +from ray._private.resource_isolation_config import ResourceIsolationConfig +from ray._private.utils import get_ray_client_dependency_error + +logger = logging.getLogger(__name__) + + +class RayParams: + """A class used to store the parameters used by Ray. + + Attributes: + redis_address: The address of the Redis server to connect to. If + this address is not provided, then this command will start Redis, a + raylet, a plasma store, a plasma manager, and some workers. + It will also kill these processes when Python exits. + redis_port: The port that the primary Redis shard should listen + to. If None, then it will fall back to + ray._private.ray_constants.DEFAULT_PORT, or a random port if the default is + not available. + redis_shard_ports: A list of the ports to use for the non-primary Redis + shards. If None, then it will fall back to the ports right after + redis_port, or random ports if those are not available. + num_cpus: Number of CPUs to configure the raylet with. + num_gpus: Number of GPUs to configure the raylet with. + resources: A dictionary mapping the name of a resource to the quantity + of that resource available. + labels: The key-value labels of the node. + memory: Total available memory for workers requesting memory. + object_store_memory: The amount of memory (in bytes) to start the + object store with. + object_manager_port int: The port to use for the object manager. + node_manager_port: The port to use for the node manager. + gcs_server_port: The port to use for the GCS server. + node_ip_address: The IP address of the node that we are on. + min_worker_port: The lowest port number that workers will bind + on. If not set or set to 0, random ports will be chosen. + max_worker_port: The highest port number that workers will bind + on. If set, min_worker_port must also be set. + worker_port_list: An explicit list of ports to be used for + workers (comma-separated). Overrides min_worker_port and + max_worker_port. + ray_client_server_port: The port number the ray client server + will bind on. If not set, the ray client server will not + be started. + redirect_output: True if stdout and stderr for non-worker + processes should be redirected to files and false otherwise. + external_addresses: The address of external Redis server to + connect to, in format of "ip1:port1,ip2:port2,...". If this + address is provided, then ray won't start Redis instances in the + head node but use external Redis server(s) instead. + num_redis_shards: The number of Redis shards to start in addition to + the primary Redis shard. + redis_max_clients: If provided, attempt to configure Redis with this + maxclients number. + redis_username: Prevents external clients without the username + from connecting to Redis if provided. + redis_password: Prevents external clients without the password + from connecting to Redis if provided. + plasma_directory: A directory where the Plasma memory mapped files will + be created. + object_spilling_directory: The path to spill objects to. The same path will + be used as the object store fallback directory as well. + worker_path: The path of the source code that will be run by the + worker. + setup_worker_path: The path of the Python file that will set up + the environment for the worker process. + huge_pages: Boolean flag indicating whether to start the Object + Store with hugetlbfs support. Requires plasma_directory. + include_dashboard: Boolean flag indicating whether to start the web + UI, which displays the status of the Ray cluster. If this value is + None, then the UI will be started if the relevant dependencies are + present. + dashboard_host: The host to bind the web UI server to. Can either be + localhost (127.0.0.1) or 0.0.0.0 (available from all interfaces). + By default, this is set to localhost to prevent access from + external machines. + dashboard_port: The port to bind the dashboard server to. + Defaults to 8265. + dashboard_agent_listen_port: The port for dashboard agents to listen on + for HTTP requests. + Defaults to 52365. + runtime_env_agent_port: The port at which the runtime env agent + listens to for HTTP. + Defaults to random available port. + plasma_store_socket_name: If provided, it specifies the socket + name used by the plasma store. + raylet_socket_name: If provided, it specifies the socket path + used by the raylet process. + temp_dir: If provided, it will specify the root temporary + directory for the Ray process. Must be an absolute path. + runtime_env_dir_name: If provided, specifies the directory that + will be created in the session dir to hold runtime_env files. + include_log_monitor: If True, then start a log monitor to + monitor the log files for all processes on this node and push their + contents to Redis. + autoscaling_config: path to autoscaling config file. + metrics_agent_port: The port to bind metrics agent. + metrics_export_port: The port at which metrics are exposed + through a Prometheus endpoint. + no_monitor: If True, the ray autoscaler monitor for this cluster + will not be started. + _system_config: Configuration for overriding RayConfig + defaults. Used to set system configuration and for experimental Ray + core feature flags. + enable_object_reconstruction: Enable plasma reconstruction on + failure. + ray_debugger_external: If true, make the Ray debugger for a + worker available externally to the node it is running on. This will + bind on 0.0.0.0 instead of localhost. + env_vars: Override environment variables for the raylet. + session_name: The current Ray session name. + webui: The url of the UI. + cluster_id: The cluster ID in hex string. + resource_isolation_config: settings for cgroupv2 based isolation of ray + system processes (defaults to no isolation if config not provided) + """ + + def __init__( + self, + redis_address: Optional[str] = None, + gcs_address: Optional[str] = None, + num_cpus: Optional[int] = None, + num_gpus: Optional[int] = None, + resources: Optional[Dict[str, float]] = None, + labels: Optional[Dict[str, str]] = None, + memory: Optional[float] = None, + object_store_memory: Optional[float] = None, + redis_port: Optional[int] = None, + redis_shard_ports: Optional[List[int]] = None, + object_manager_port: Optional[int] = None, + node_manager_port: int = 0, + gcs_server_port: Optional[int] = None, + node_ip_address: Optional[str] = None, + node_name: Optional[str] = None, + min_worker_port: Optional[int] = None, + max_worker_port: Optional[int] = None, + worker_port_list: Optional[List[int]] = None, + ray_client_server_port: Optional[int] = None, + driver_mode=None, + redirect_output: Optional[bool] = None, + external_addresses: Optional[List[str]] = None, + num_redis_shards: Optional[int] = None, + redis_max_clients: Optional[int] = None, + redis_username: Optional[str] = ray_constants.REDIS_DEFAULT_USERNAME, + redis_password: Optional[str] = ray_constants.REDIS_DEFAULT_PASSWORD, + plasma_directory: Optional[str] = None, + object_spilling_directory: Optional[str] = None, + worker_path: Optional[str] = None, + setup_worker_path: Optional[str] = None, + huge_pages: Optional[bool] = False, + include_dashboard: Optional[bool] = None, + dashboard_host: Optional[str] = ray_constants.DEFAULT_DASHBOARD_IP, + dashboard_port: Optional[bool] = ray_constants.DEFAULT_DASHBOARD_PORT, + dashboard_agent_listen_port: Optional[ + int + ] = ray_constants.DEFAULT_DASHBOARD_AGENT_LISTEN_PORT, + runtime_env_agent_port: Optional[int] = None, + plasma_store_socket_name: Optional[str] = None, + raylet_socket_name: Optional[str] = None, + temp_dir: Optional[str] = None, + runtime_env_dir_name: Optional[str] = None, + include_log_monitor: Optional[str] = None, + autoscaling_config: Optional[str] = None, + ray_debugger_external: bool = False, + _system_config: Optional[Dict[str, str]] = None, + enable_object_reconstruction: Optional[bool] = False, + metrics_agent_port: Optional[int] = None, + metrics_export_port: Optional[int] = None, + tracing_startup_hook=None, + no_monitor: Optional[bool] = False, + env_vars: Optional[Dict[str, str]] = None, + session_name: Optional[str] = None, + webui: Optional[str] = None, + cluster_id: Optional[str] = None, + node_id: Optional[str] = None, + resource_isolation_config: Optional[ResourceIsolationConfig] = None, + ): + self.redis_address = redis_address + self.gcs_address = gcs_address + self.num_cpus = num_cpus + self.num_gpus = num_gpus + self.memory = memory + self.object_store_memory = object_store_memory + self.resources = resources + self.redis_port = redis_port + self.redis_shard_ports = redis_shard_ports + self.object_manager_port = object_manager_port + self.node_manager_port = node_manager_port + self.gcs_server_port = gcs_server_port + self.node_ip_address = node_ip_address + self.node_name = node_name + self.min_worker_port = min_worker_port + self.max_worker_port = max_worker_port + self.worker_port_list = worker_port_list + self.ray_client_server_port = ray_client_server_port + self.driver_mode = driver_mode + self.redirect_output = redirect_output + self.external_addresses = external_addresses + self.num_redis_shards = num_redis_shards + self.redis_max_clients = redis_max_clients + self.redis_username = redis_username + self.redis_password = redis_password + self.plasma_directory = plasma_directory + self.object_spilling_directory = object_spilling_directory + self.worker_path = worker_path + self.setup_worker_path = setup_worker_path + self.huge_pages = huge_pages + self.include_dashboard = include_dashboard + self.dashboard_host = dashboard_host + self.dashboard_port = dashboard_port + self.dashboard_agent_listen_port = dashboard_agent_listen_port + self.runtime_env_agent_port = runtime_env_agent_port + self.plasma_store_socket_name = plasma_store_socket_name + self.raylet_socket_name = raylet_socket_name + self.temp_dir = temp_dir + self.runtime_env_dir_name = ( + runtime_env_dir_name or ray_constants.DEFAULT_RUNTIME_ENV_DIR_NAME + ) + self.include_log_monitor = include_log_monitor + self.autoscaling_config = autoscaling_config + self.metrics_agent_port = metrics_agent_port + self.metrics_export_port = metrics_export_port + self.tracing_startup_hook = tracing_startup_hook + self.no_monitor = no_monitor + self.ray_debugger_external = ray_debugger_external + self.env_vars = env_vars + self.session_name = session_name + self.webui = webui + self._system_config = _system_config or {} + self._enable_object_reconstruction = enable_object_reconstruction + self.labels = labels + self._check_usage() + self.cluster_id = cluster_id + self.node_id = node_id + + self.resource_isolation_config = resource_isolation_config + if not self.resource_isolation_config: + self.resource_isolation_config = ResourceIsolationConfig( + enable_resource_isolation=False + ) + + # Set the internal config options for object reconstruction. + if enable_object_reconstruction: + # Turn off object pinning. + if self._system_config is None: + self._system_config = dict() + print(self._system_config) + self._system_config["lineage_pinning_enabled"] = True + + def update(self, **kwargs): + """Update the settings according to the keyword arguments. + + Args: + kwargs: The keyword arguments to set corresponding fields. + """ + for arg in kwargs: + if hasattr(self, arg): + setattr(self, arg, kwargs[arg]) + else: + raise ValueError(f"Invalid RayParams parameter in update: {arg}") + + self._check_usage() + + def update_if_absent(self, **kwargs): + """Update the settings when the target fields are None. + + Args: + kwargs: The keyword arguments to set corresponding fields. + """ + for arg in kwargs: + if hasattr(self, arg): + if getattr(self, arg) is None: + setattr(self, arg, kwargs[arg]) + else: + raise ValueError( + f"Invalid RayParams parameter in update_if_absent: {arg}" + ) + + self._check_usage() + + def update_pre_selected_port(self): + """Update the pre-selected port information + + Returns: + The dictionary mapping of component -> ports. + """ + + def wrap_port(port): + # 0 port means select a random port for the grpc server. + if port is None or port == 0: + return [] + else: + return [port] + + # Create a dictionary of the component -> port mapping. + pre_selected_ports = { + "gcs": wrap_port(self.redis_port), + "object_manager": wrap_port(self.object_manager_port), + "node_manager": wrap_port(self.node_manager_port), + "gcs_server": wrap_port(self.gcs_server_port), + "client_server": wrap_port(self.ray_client_server_port), + "dashboard": wrap_port(self.dashboard_port), + "dashboard_agent_grpc": wrap_port(self.metrics_agent_port), + "dashboard_agent_http": wrap_port(self.dashboard_agent_listen_port), + "runtime_env_agent": wrap_port(self.runtime_env_agent_port), + "metrics_export": wrap_port(self.metrics_export_port), + } + redis_shard_ports = self.redis_shard_ports + if redis_shard_ports is None: + redis_shard_ports = [] + pre_selected_ports["redis_shards"] = redis_shard_ports + if self.worker_port_list is None: + if self.min_worker_port is not None and self.max_worker_port is not None: + pre_selected_ports["worker_ports"] = list( + range(self.min_worker_port, self.max_worker_port + 1) + ) + else: + # The dict is not updated when it requires random ports. + pre_selected_ports["worker_ports"] = [] + else: + pre_selected_ports["worker_ports"] = [ + int(port) for port in self.worker_port_list.split(",") + ] + + # Update the pre selected port set. + self.reserved_ports = set() + for comp, port_list in pre_selected_ports.items(): + for port in port_list: + if port in self.reserved_ports: + raise ValueError( + f"Ray component {comp} is trying to use " + f"a port number {port} that is used by other components.\n" + f"Port information: {self._format_ports(pre_selected_ports)}\n" + "If you allocate ports, please make sure the same port " + "is not used by multiple components." + ) + self.reserved_ports.add(port) + + def _check_usage(self): + if self.worker_port_list is not None: + for port_str in self.worker_port_list.split(","): + try: + port = int(port_str) + except ValueError as e: + raise ValueError( + "worker_port_list must be a comma-separated " + f"list of integers: {e}" + ) from None + + if port < 1024 or port > 65535: + raise ValueError( + "Ports in worker_port_list must be " + f"between 1024 and 65535. Got: {port}" + ) + + # Used primarily for testing. + if os.environ.get("RAY_USE_RANDOM_PORTS", False): + if self.min_worker_port is None and self.max_worker_port is None: + self.min_worker_port = 0 + self.max_worker_port = 0 + + if self.min_worker_port is not None: + if self.min_worker_port != 0 and ( + self.min_worker_port < 1024 or self.min_worker_port > 65535 + ): + raise ValueError( + "min_worker_port must be 0 or an integer between 1024 and 65535." + ) + + if self.max_worker_port is not None: + if self.min_worker_port is None: + raise ValueError( + "If max_worker_port is set, min_worker_port must also be set." + ) + elif self.max_worker_port != 0: + if self.max_worker_port < 1024 or self.max_worker_port > 65535: + raise ValueError( + "max_worker_port must be 0 or an integer between " + "1024 and 65535." + ) + elif self.max_worker_port <= self.min_worker_port: + raise ValueError( + "max_worker_port must be higher than min_worker_port." + ) + if self.ray_client_server_port is not None: + if get_ray_client_dependency_error() is not None: + raise ValueError( + "Ray Client requires pip package `ray[client]`. " + "If you installed the minimal Ray (e.g. `pip install ray`), " + "please reinstall by executing `pip install ray[client]`." + ) + if ( + self.ray_client_server_port < 1024 + or self.ray_client_server_port > 65535 + ): + raise ValueError( + "ray_client_server_port must be an integer " + "between 1024 and 65535." + ) + if self.runtime_env_agent_port is not None: + if ( + self.runtime_env_agent_port < 1024 + or self.runtime_env_agent_port > 65535 + ): + raise ValueError( + "runtime_env_agent_port must be an integer " + "between 1024 and 65535." + ) + + if self.resources is not None: + + def build_error(resource, alternative): + return ( + f"{self.resources} -> `{resource}` cannot be a " + "custom resource because it is one of the default resources " + f"({ray_constants.DEFAULT_RESOURCES}). " + f"Use `{alternative}` instead. For example, use `ray start " + f"--{alternative.replace('_', '-')}=1` instead of " + f"`ray start --resources={{'{resource}': 1}}`" + ) + + assert "CPU" not in self.resources, build_error("CPU", "num_cpus") + assert "GPU" not in self.resources, build_error("GPU", "num_gpus") + assert "memory" not in self.resources, build_error("memory", "memory") + assert "object_store_memory" not in self.resources, build_error( + "object_store_memory", "object_store_memory" + ) + + if self.redirect_output is not None: + raise DeprecationWarning("The redirect_output argument is deprecated.") + + if self.temp_dir is not None and not os.path.isabs(self.temp_dir): + raise ValueError("temp_dir must be absolute path or None.") + + if self.temp_dir is not None and os.getenv("VIRTUAL_ENV"): + is_relative = True + try: + ( + pathlib.Path(self.temp_dir) + .resolve() + .relative_to(pathlib.Path(os.getenv("VIRTUAL_ENV")).resolve()) + ) + except ValueError: + is_relative = False + + if is_relative: + raise ValueError( + "temp_dir must not be child directory of virtualenv root" + ) + + def _format_ports(self, pre_selected_ports): + """Format the pre-selected ports information to be more human-readable.""" + ports = pre_selected_ports.copy() + + for comp, port_list in ports.items(): + if len(port_list) == 1: + ports[comp] = port_list[0] + elif len(port_list) == 0: + # Nothing is selected, meaning it will be randomly selected. + ports[comp] = "random" + elif comp == "worker_ports": + min_port = port_list[0] + max_port = port_list[len(port_list) - 1] + if len(port_list) < 50: + port_range_str = str(port_list) + else: + port_range_str = f"from {min_port} to {max_port}" + ports[comp] = f"{len(port_list)} ports {port_range_str}" + return ports diff --git a/lib/python3.12/site-packages/ray/_private/path_utils.py b/lib/python3.12/site-packages/ray/_private/path_utils.py new file mode 100644 index 0000000000000000000000000000000000000000..0e432a3b67af4554d70add15ff21b474b72f6080 --- /dev/null +++ b/lib/python3.12/site-packages/ray/_private/path_utils.py @@ -0,0 +1,32 @@ +import pathlib +import urllib + +"""Cross-platform utilities for manipulating paths and URIs. + +NOTE: All functions in this file must support POSIX and Windows. +""" + + +def is_path(path_or_uri: str) -> bool: + """Returns True if uri_or_path is a path and False otherwise. + + Windows paths start with a drive name which can be interpreted as + a URI scheme by urlparse and thus needs to be treated differently + form POSIX paths. + + E.g. Creating a directory returns the path 'C:\\Users\\mp5n6ul72w\\working_dir' + will have the scheme 'C:'. + """ + if not isinstance(path_or_uri, str): + raise TypeError(f" path_or_uri must be a string, got {type(path_or_uri)}.") + + parsed_path = pathlib.Path(path_or_uri) + parsed_uri = urllib.parse.urlparse(path_or_uri) + + if isinstance(parsed_path, pathlib.PurePosixPath): + return not parsed_uri.scheme + elif isinstance(parsed_path, pathlib.PureWindowsPath): + return parsed_uri.scheme == parsed_path.drive.strip(":").lower() + else: + # this should never happen. + raise TypeError(f"Unsupported path type: {type(parsed_path).__name__}") diff --git a/lib/python3.12/site-packages/ray/_private/process_watcher.py b/lib/python3.12/site-packages/ray/_private/process_watcher.py new file mode 100644 index 0000000000000000000000000000000000000000..6dd5c7df36bcf33b1c95681d0132d7d21e2fe486 --- /dev/null +++ b/lib/python3.12/site-packages/ray/_private/process_watcher.py @@ -0,0 +1,197 @@ +import asyncio +import io +import logging +import os +import sys +from concurrent.futures import ThreadPoolExecutor + +import ray +import ray._private.ray_constants as ray_constants +import ray.dashboard.consts as dashboard_consts +from ray._common.utils import run_background_task +from ray._raylet import GcsClient +from ray.dashboard.consts import _PARENT_DEATH_THREASHOLD + +# Import psutil after ray so the packaged version is used. +import psutil + +logger = logging.getLogger(__name__) + +# TODO: move all consts from dashboard_consts to ray_constants and rename to remove +# DASHBOARD_ prefixes. + +# Publishes at most this number of lines of Raylet logs, when the Raylet dies +# unexpectedly. +_RAYLET_LOG_MAX_PUBLISH_LINES = 20 + +# Reads at most this amount of Raylet logs from the tail, for publishing and +# checking if the Raylet was terminated gracefully. +_RAYLET_LOG_MAX_TAIL_SIZE = 1 * 1024**2 + +try: + create_task = asyncio.create_task +except AttributeError: + create_task = asyncio.ensure_future + + +def get_raylet_pid(): + # TODO(edoakes): RAY_RAYLET_PID isn't properly set on Windows. This is + # only used for fate-sharing with the raylet and we need a different + # fate-sharing mechanism for Windows anyways. + if sys.platform in ["win32", "cygwin"]: + return None + raylet_pid = int(os.environ["RAY_RAYLET_PID"]) + assert raylet_pid > 0 + logger.info("raylet pid is %s", raylet_pid) + return raylet_pid + + +def create_check_raylet_task(log_dir, gcs_client, parent_dead_callback, loop): + """ + Creates an asyncio task to periodically check if the raylet process is still + running. If raylet is dead for _PARENT_DEATH_THREASHOLD (5) times, prepare to exit + as follows: + + - Write logs about whether the raylet exit is graceful, by looking into the raylet + log and search for term "SIGTERM", + - Flush the logs via GcsClient, + - Exit. + """ + if sys.platform in ["win32", "cygwin"]: + raise RuntimeError("can't check raylet process in Windows.") + raylet_pid = get_raylet_pid() + + if dashboard_consts.PARENT_HEALTH_CHECK_BY_PIPE: + logger.info("check_parent_via_pipe") + check_parent_task = _check_parent_via_pipe( + log_dir, gcs_client, loop, parent_dead_callback + ) + else: + logger.info("_check_parent") + check_parent_task = _check_parent( + raylet_pid, log_dir, gcs_client, parent_dead_callback + ) + + return run_background_task(check_parent_task) + + +def report_raylet_error_logs(log_dir: str, gcs_client: GcsClient): + log_path = os.path.join(log_dir, "raylet.out") + error = False + msg = "Raylet is terminated. " + try: + with open(log_path, "r", encoding="utf-8") as f: + # Seek to _RAYLET_LOG_MAX_TAIL_SIZE from the end if the + # file is larger than that. + f.seek(0, io.SEEK_END) + pos = max(0, f.tell() - _RAYLET_LOG_MAX_TAIL_SIZE) + f.seek(pos, io.SEEK_SET) + # Read remaining logs by lines. + raylet_logs = f.readlines() + # Assume the SIGTERM message must exist within the last + # _RAYLET_LOG_MAX_TAIL_SIZE of the log file. + if any("Raylet received SIGTERM" in line for line in raylet_logs): + msg += "Termination is graceful." + logger.info(msg) + else: + msg += ( + "Termination is unexpected. Possible reasons " + "include: (1) SIGKILL by the user or system " + "OOM killer, (2) Invalid memory access from " + "Raylet causing SIGSEGV or SIGBUS, " + "(3) Other termination signals. " + f"Last {_RAYLET_LOG_MAX_PUBLISH_LINES} lines " + "of the Raylet logs:\n" + ) + msg += " " + " ".join( + raylet_logs[-_RAYLET_LOG_MAX_PUBLISH_LINES:] + ) + error = True + except Exception as e: + msg += f"Failed to read Raylet logs at {log_path}: {e}!" + logger.exception(msg) + error = True + if error: + logger.error(msg) + # TODO: switch to async if necessary. + ray._private.utils.publish_error_to_driver( + ray_constants.RAYLET_DIED_ERROR, + msg, + gcs_client=gcs_client, + ) + else: + logger.info(msg) + + +async def _check_parent_via_pipe( + log_dir: str, gcs_client: GcsClient, loop, parent_dead_callback +): + while True: + try: + # Read input asynchronously. + # The parent (raylet) should have redirected its pipe + # to stdin. If we read 0 bytes from stdin, it means + # the process is dead. + with ThreadPoolExecutor(max_workers=1) as executor: + input_data = await loop.run_in_executor( + executor, lambda: sys.stdin.readline() + ) + if len(input_data) == 0: + # cannot read bytes from parent == parent is dead. + parent_dead_callback("_check_parent_via_pipe: The parent is dead.") + report_raylet_error_logs(log_dir, gcs_client) + sys.exit(0) + except Exception as e: + logger.exception( + "raylet health checking is failed. " + f"The agent process may leak. Exception: {e}" + ) + + +async def _check_parent(raylet_pid, log_dir, gcs_client, parent_dead_callback): + """Check if raylet is dead and fate-share if it is.""" + try: + curr_proc = psutil.Process() + parent_death_cnt = 0 + while True: + parent = curr_proc.parent() + # If the parent is dead, it is None. + parent_gone = parent is None + init_assigned_for_parent = False + parent_changed = False + + if parent: + # Sometimes, the parent is changed to the `init` process. + # In this case, the parent.pid is 1. + init_assigned_for_parent = parent.pid == 1 + # Sometimes, the parent is dead, and the pid is reused + # by other processes. In this case, this condition is triggered. + parent_changed = raylet_pid != parent.pid + + if parent_gone or init_assigned_for_parent or parent_changed: + parent_death_cnt += 1 + logger.warning( + f"Raylet is considered dead {parent_death_cnt} X. " + f"If it reaches to {_PARENT_DEATH_THREASHOLD}, the agent " + f"will kill itself. Parent: {parent}, " + f"parent_gone: {parent_gone}, " + f"init_assigned_for_parent: {init_assigned_for_parent}, " + f"parent_changed: {parent_changed}." + ) + if parent_death_cnt < _PARENT_DEATH_THREASHOLD: + await asyncio.sleep( + dashboard_consts.DASHBOARD_AGENT_CHECK_PARENT_INTERVAL_S + ) + continue + + parent_dead_callback("_check_parent: The parent is dead.") + report_raylet_error_logs(log_dir, gcs_client) + sys.exit(0) + else: + parent_death_cnt = 0 + await asyncio.sleep( + dashboard_consts.DASHBOARD_AGENT_CHECK_PARENT_INTERVAL_S + ) + except Exception: + logger.exception("Failed to check parent PID, exiting.") + sys.exit(1) diff --git a/lib/python3.12/site-packages/ray/_private/profiling.py b/lib/python3.12/site-packages/ray/_private/profiling.py new file mode 100644 index 0000000000000000000000000000000000000000..ca55aad62ed5f2d5a30727b837769942dc5dba0e --- /dev/null +++ b/lib/python3.12/site-packages/ray/_private/profiling.py @@ -0,0 +1,240 @@ +import json +import os +from collections import defaultdict +from dataclasses import asdict, dataclass +from typing import Dict, List, Union + +import ray + + +class _NullLogSpan: + """A log span context manager that does nothing""" + + def __enter__(self): + pass + + def __exit__(self, type, value, tb): + pass + + +PROFILING_ENABLED = "RAY_PROFILING" in os.environ +NULL_LOG_SPAN = _NullLogSpan() + +# Colors are specified at +# https://github.com/catapult-project/catapult/blob/master/tracing/tracing/base/color_scheme.html. # noqa: E501 +_default_color_mapping = defaultdict( + lambda: "generic_work", + { + "worker_idle": "cq_build_abandoned", + "task": "rail_response", + "task:deserialize_arguments": "rail_load", + "task:execute": "rail_animation", + "task:store_outputs": "rail_idle", + "wait_for_function": "detailed_memory_dump", + "ray.get": "good", + "ray.put": "terrible", + "ray.wait": "vsync_highlight_color", + "submit_task": "background_memory_dump", + "fetch_and_run_function": "detailed_memory_dump", + "register_remote_function": "detailed_memory_dump", + }, +) + + +@dataclass(init=True) +class ChromeTracingCompleteEvent: + # https://docs.google.com/document/d/1CvAClvFfyA5R-PhYUmn5OOQtYMH4h6I0nSsKchNAySU/preview#heading=h.lpfof2aylapb # noqa + # The event categories. This is a comma separated list of categories + # for the event. The categories can be used to hide events in + # the Trace Viewer UI. + cat: str + # The string displayed on the event. + name: str + # The identifier for the group of rows that the event + # appears in. + pid: int + # The identifier for the row that the event appears in. + tid: int + # The start time in microseconds. + ts: int + # The duration in microseconds. + dur: int + # This is the name of the color to display the box in. + cname: str + # The extra user-defined data. + args: Dict[str, Union[str, int]] + # The event type (X means the complete event). + ph: str = "X" + + +@dataclass(init=True) +class ChromeTracingMetadataEvent: + # https://docs.google.com/document/d/1CvAClvFfyA5R-PhYUmn5OOQtYMH4h6I0nSsKchNAySU/preview#bookmark=id.iycbnb4z7i9g # noqa + name: str + # Metadata arguments. E.g., name: + args: Dict[str, str] + # The process id of this event. In Ray, pid indicates the node. + pid: int + # The thread id of this event. In Ray, tid indicates each worker. + tid: int = None + # M means the metadata event. + ph: str = "M" + + +def profile(event_type, extra_data=None): + """Profile a span of time so that it appears in the timeline visualization. + + Note that this only works in the raylet code path. + + This function can be used as follows (both on the driver or within a task). + + .. testcode:: + import ray._private.profiling as profiling + + with profiling.profile("custom event", extra_data={'key': 'val'}): + # Do some computation here. + x = 1 * 2 + + Optionally, a dictionary can be passed as the "extra_data" argument, and + it can have keys "name" and "cname" if you want to override the default + timeline display text and box color. Other values will appear at the bottom + of the chrome tracing GUI when you click on the box corresponding to this + profile span. + + Args: + event_type: A string describing the type of the event. + extra_data: This must be a dictionary mapping strings to strings. This + data will be added to the json objects that are used to populate + the timeline, so if you want to set a particular color, you can + simply set the "cname" attribute to an appropriate color. + Similarly, if you set the "name" attribute, then that will set the + text displayed on the box in the timeline. + + Returns: + An object that can profile a span of time via a "with" statement. + """ + if not PROFILING_ENABLED: + return NULL_LOG_SPAN + worker = ray._private.worker.global_worker + if worker.mode == ray._private.worker.LOCAL_MODE: + return NULL_LOG_SPAN + return worker.core_worker.profile_event(event_type.encode("ascii"), extra_data) + + +def chrome_tracing_dump( + tasks: List[dict], +) -> str: + """Generate a chrome/perfetto tracing dump using task events. + + Args: + tasks: List of tasks generated by a state API list_tasks(detail=True). + + Returns: + Json serialized dump to create a chrome/perfetto tracing. + """ + # All events from given tasks. + all_events = [] + + # Chrome tracing doesn't have a concept of "node". Instead, we use + # chrome tracing's pid == ray's node. + # chrome tracing's tid == ray's process. + # Note that pid or tid is usually integer, but ray's node/process has + # ids in string. + # Unfortunately, perfetto doesn't allow to have string as a value of pid/tid. + # To workaround it, we use Metadata event from chrome tracing schema + # (https://docs.google.com/document/d/1CvAClvFfyA5R-PhYUmn5OOQtYMH4h6I0nSsKchNAySU/preview#heading=h.xqopa5m0e28f) # noqa + # which allows pid/tid -> name mapping. In order to use this schema + # we build node_ip/(node_ip, worker_id) -> arbitrary index mapping. + + # node ip address -> node idx. + node_to_index = {} + # Arbitrary index mapped to the ip address. + node_idx = 0 + # (node index, worker id) -> worker idx + worker_to_index = {} + # Arbitrary index mapped to the (node index, worker id). + worker_idx = 0 + + for task in tasks: + profiling_data = task.get("profiling_data", []) + if profiling_data: + node_ip_address = profiling_data["node_ip_address"] + component_events = profiling_data["events"] + component_type = profiling_data["component_type"] + component_id = component_type + ":" + profiling_data["component_id"] + + if component_type not in ["worker", "driver"]: + continue + + for event in component_events: + extra_data = event["extra_data"] + # Propagate extra data. + extra_data["task_id"] = task["task_id"] + extra_data["job_id"] = task["job_id"] + extra_data["attempt_number"] = task["attempt_number"] + extra_data["func_or_class_name"] = task["func_or_class_name"] + extra_data["actor_id"] = task["actor_id"] + event_name = event["event_name"] + + # build a id -> arbitrary index mapping + if node_ip_address not in node_to_index: + node_to_index[node_ip_address] = node_idx + # Whenever new node ip is introduced, we increment the index. + node_idx += 1 + + if ( + node_to_index[node_ip_address], + component_id, + ) not in worker_to_index: # noqa + worker_to_index[ + (node_to_index[node_ip_address], component_id) + ] = worker_idx # noqa + worker_idx += 1 + + # Modify the name with the additional user-defined extra data. + cname = _default_color_mapping[event["event_name"]] + name = event_name + + if "cname" in extra_data: + cname = _default_color_mapping[event["extra_data"]["cname"]] + if "name" in extra_data: + name = extra_data["name"] + + new_event = ChromeTracingCompleteEvent( + cat=event_name, + name=name, + pid=node_to_index[node_ip_address], + tid=worker_to_index[(node_to_index[node_ip_address], component_id)], + ts=event["start_time"] * 1e3, + dur=(event["end_time"] * 1e3) - (event["start_time"] * 1e3), + cname=cname, + args=extra_data, + ) + all_events.append(asdict(new_event)) + + for node, i in node_to_index.items(): + all_events.append( + asdict( + ChromeTracingMetadataEvent( + name="process_name", + pid=i, + args={"name": f"Node {node}"}, + ) + ) + ) + + for worker, i in worker_to_index.items(): + all_events.append( + asdict( + ChromeTracingMetadataEvent( + name="thread_name", + ph="M", + tid=i, + pid=worker[0], + args={"name": worker[1]}, + ) + ) + ) + + # Handle task event disabled. + return json.dumps(all_events) diff --git a/lib/python3.12/site-packages/ray/_private/prometheus_exporter.py b/lib/python3.12/site-packages/ray/_private/prometheus_exporter.py new file mode 100644 index 0000000000000000000000000000000000000000..c60b764a95b1e2803868a372ff855f3490d2ba1a --- /dev/null +++ b/lib/python3.12/site-packages/ray/_private/prometheus_exporter.py @@ -0,0 +1,363 @@ +# NOTE: This file has been copied from OpenCensus Python exporter. +# It is because OpenCensus Prometheus exporter hasn't released for a while +# and the latest version has a compatibility issue with the latest OpenCensus +# library. + +import logging +import re + +from opencensus.common.transports import sync +from opencensus.stats import aggregation_data as aggregation_data_module, base_exporter +from prometheus_client import start_http_server +from prometheus_client.core import ( + REGISTRY, + CounterMetricFamily, + GaugeMetricFamily, + HistogramMetricFamily, + UnknownMetricFamily, +) + +logger = logging.getLogger(__name__) + + +class Options(object): + """Options contains options for configuring the exporter. + The address can be empty as the prometheus client will + assume it's localhost + :type namespace: str + :param namespace: The prometheus namespace to be used. Defaults to ''. + :type port: int + :param port: The Prometheus port to be used. Defaults to 8000. + :type address: str + :param address: The Prometheus address to be used. Defaults to ''. + :type registry: registry + :param registry: The Prometheus address to be used. Defaults to ''. + :type registry: :class:`~prometheus_client.core.CollectorRegistry` + :param registry: A Prometheus collector registry instance. + """ + + def __init__(self, namespace="", port=8000, address="", registry=REGISTRY): + self._namespace = namespace + self._registry = registry + self._port = int(port) + self._address = address + + @property + def registry(self): + """Prometheus Collector Registry instance""" + return self._registry + + @property + def namespace(self): + """Prefix to be used with view name""" + return self._namespace + + @property + def port(self): + """Port number to listen""" + return self._port + + @property + def address(self): + """Endpoint address (default is localhost)""" + return self._address + + +class Collector(object): + """Collector represents the Prometheus Collector object""" + + def __init__(self, options=Options(), view_name_to_data_map=None): + if view_name_to_data_map is None: + view_name_to_data_map = {} + self._options = options + self._registry = options.registry + self._view_name_to_data_map = view_name_to_data_map + self._registered_views = {} + + @property + def options(self): + """Options to be used to configure the exporter""" + return self._options + + @property + def registry(self): + """Prometheus Collector Registry instance""" + return self._registry + + @property + def view_name_to_data_map(self): + """Map with all view data objects + that will be sent to Prometheus + """ + return self._view_name_to_data_map + + @property + def registered_views(self): + """Map with all registered views""" + return self._registered_views + + def register_view(self, view): + """register_view will create the needed structure + in order to be able to sent all data to Prometheus + """ + v_name = get_view_name(self.options.namespace, view) + + if v_name not in self.registered_views: + desc = { + "name": v_name, + "documentation": view.description, + "labels": list(map(sanitize, view.columns)), + "units": view.measure.unit, + } + self.registered_views[v_name] = desc + + def add_view_data(self, view_data): + """Add view data object to be sent to server""" + self.register_view(view_data.view) + v_name = get_view_name(self.options.namespace, view_data.view) + self.view_name_to_data_map[v_name] = view_data + + # TODO: add start and end timestamp + def to_metric(self, desc, tag_values, agg_data, metrics_map): + """to_metric translate the data that OpenCensus create + to Prometheus format, using Prometheus Metric object + :type desc: dict + :param desc: The map that describes view definition + :type tag_values: tuple of :class: + `~opencensus.tags.tag_value.TagValue` + :param object of opencensus.tags.tag_value.TagValue: + TagValue object used as label values + :type agg_data: object of :class: + `~opencensus.stats.aggregation_data.AggregationData` + :param object of opencensus.stats.aggregation_data.AggregationData: + Aggregated data that needs to be converted as Prometheus samples + :rtype: :class:`~prometheus_client.core.CounterMetricFamily` or + :class:`~prometheus_client.core.HistogramMetricFamily` or + :class:`~prometheus_client.core.UnknownMetricFamily` or + :class:`~prometheus_client.core.GaugeMetricFamily` + """ + metric_name = desc["name"] + metric_description = desc["documentation"] + label_keys = desc["labels"] + metric_units = desc["units"] + assert len(tag_values) == len(label_keys), (tag_values, label_keys) + # Prometheus requires that all tag values be strings hence + # the need to cast none to the empty string before exporting. See + # https://github.com/census-instrumentation/opencensus-python/issues/480 + tag_values = [tv if tv else "" for tv in tag_values] + + if isinstance(agg_data, aggregation_data_module.CountAggregationData): + metric = metrics_map.get(metric_name) + if not metric: + metric = CounterMetricFamily( + name=metric_name, + documentation=metric_description, + unit=metric_units, + labels=label_keys, + ) + metrics_map[metric_name] = metric + metric.add_metric(labels=tag_values, value=agg_data.count_data) + return + + elif isinstance(agg_data, aggregation_data_module.DistributionAggregationData): + + assert agg_data.bounds == sorted(agg_data.bounds) + # buckets are a list of buckets. Each bucket is another list with + # a pair of bucket name and value, or a triple of bucket name, + # value, and exemplar. buckets need to be in order. + buckets = [] + cum_count = 0 # Prometheus buckets expect cumulative count. + for ii, bound in enumerate(agg_data.bounds): + cum_count += agg_data.counts_per_bucket[ii] + bucket = [str(bound), cum_count] + buckets.append(bucket) + # Prometheus requires buckets to be sorted, and +Inf present. + # In OpenCensus we don't have +Inf in the bucket bonds so need to + # append it here. + buckets.append(["+Inf", agg_data.count_data]) + metric = metrics_map.get(metric_name) + if not metric: + metric = HistogramMetricFamily( + name=metric_name, + documentation=metric_description, + labels=label_keys, + ) + metrics_map[metric_name] = metric + metric.add_metric( + labels=tag_values, + buckets=buckets, + sum_value=agg_data.sum, + ) + return + + elif isinstance(agg_data, aggregation_data_module.SumAggregationData): + metric = metrics_map.get(metric_name) + if not metric: + metric = UnknownMetricFamily( + name=metric_name, + documentation=metric_description, + labels=label_keys, + ) + metrics_map[metric_name] = metric + metric.add_metric(labels=tag_values, value=agg_data.sum_data) + return + + elif isinstance(agg_data, aggregation_data_module.LastValueAggregationData): + metric = metrics_map.get(metric_name) + if not metric: + metric = GaugeMetricFamily( + name=metric_name, + documentation=metric_description, + labels=label_keys, + ) + metrics_map[metric_name] = metric + metric.add_metric(labels=tag_values, value=agg_data.value) + return + + else: + raise ValueError(f"unsupported aggregation type {type(agg_data)}") + + def collect(self): # pragma: NO COVER + """Collect fetches the statistics from OpenCensus + and delivers them as Prometheus Metrics. + Collect is invoked every time a prometheus.Gatherer is run + for example when the HTTP endpoint is invoked by Prometheus. + """ + # Make a shallow copy of self._view_name_to_data_map, to avoid seeing + # concurrent modifications when iterating through the dictionary. + metrics_map = {} + for v_name, view_data in self._view_name_to_data_map.copy().items(): + if v_name not in self.registered_views: + continue + desc = self.registered_views[v_name] + for tag_values in view_data.tag_value_aggregation_data_map: + agg_data = view_data.tag_value_aggregation_data_map[tag_values] + self.to_metric(desc, tag_values, agg_data, metrics_map) + + for metric in metrics_map.values(): + yield metric + + +class PrometheusStatsExporter(base_exporter.StatsExporter): + """Exporter exports stats to Prometheus, users need + to register the exporter as an HTTP Handler to be + able to export. + :type options: + :class:`~opencensus.ext.prometheus.stats_exporter.Options` + :param options: An options object with the parameters to instantiate the + prometheus exporter. + :type gatherer: :class:`~prometheus_client.core.CollectorRegistry` + :param gatherer: A Prometheus collector registry instance. + :type transport: + :class:`opencensus.common.transports.sync.SyncTransport` or + :class:`opencensus.common.transports.async_.AsyncTransport` + :param transport: An instance of a Transpor to send data with. + :type collector: + :class:`~opencensus.ext.prometheus.stats_exporter.Collector` + :param collector: An instance of the Prometheus Collector object. + """ + + def __init__( + self, options, gatherer, transport=sync.SyncTransport, collector=Collector() + ): + self._options = options + self._gatherer = gatherer + self._collector = collector + self._transport = transport(self) + self.serve_http() + REGISTRY.register(self._collector) + + @property + def transport(self): + """The transport way to be sent data to server + (default is sync). + """ + return self._transport + + @property + def collector(self): + """Collector class instance to be used + to communicate with Prometheus + """ + return self._collector + + @property + def gatherer(self): + """Prometheus Collector Registry instance""" + return self._gatherer + + @property + def options(self): + """Options to be used to configure the exporter""" + return self._options + + def export(self, view_data): + """export send the data to the transport class + in order to be sent to Prometheus in a sync or async way. + """ + if view_data is not None: # pragma: NO COVER + self.transport.export(view_data) + + def on_register_view(self, view): + return NotImplementedError("Not supported by Prometheus") + + def emit(self, view_data): # pragma: NO COVER + """Emit exports to the Prometheus if view data has one or more rows. + Each OpenCensus AggregationData will be converted to + corresponding Prometheus Metric: SumData will be converted + to Untyped Metric, CountData will be a Counter Metric + DistributionData will be a Histogram Metric. + """ + + for v_data in view_data: + if v_data.tag_value_aggregation_data_map is None: + v_data.tag_value_aggregation_data_map = {} + + self.collector.add_view_data(v_data) + + def serve_http(self): + """serve_http serves the Prometheus endpoint.""" + address = str(self.options.address) + kwargs = {"addr": address} if address else {} + start_http_server(port=self.options.port, **kwargs) + + +def new_stats_exporter(option): + """new_stats_exporter returns an exporter + that exports stats to Prometheus. + """ + if option.namespace == "": + raise ValueError("Namespace can not be empty string.") + + collector = new_collector(option) + + exporter = PrometheusStatsExporter( + options=option, gatherer=option.registry, collector=collector + ) + return exporter + + +def new_collector(options): + """new_collector should be used + to create instance of Collector class in order to + prevent the usage of constructor directly + """ + return Collector(options=options) + + +def get_view_name(namespace, view): + """create the name for the view""" + name = "" + if namespace != "": + name = namespace + "_" + return sanitize(name + view.name) + + +_NON_LETTERS_NOR_DIGITS_RE = re.compile(r"[^\w]", re.UNICODE | re.IGNORECASE) + + +def sanitize(key): + """sanitize the given metric name or label according to Prometheus rule. + Replace all characters other than [A-Za-z0-9_] with '_'. + """ + return _NON_LETTERS_NOR_DIGITS_RE.sub("_", key) diff --git a/lib/python3.12/site-packages/ray/_private/protobuf_compat.py b/lib/python3.12/site-packages/ray/_private/protobuf_compat.py new file mode 100644 index 0000000000000000000000000000000000000000..01256a5a82c397946175c771c09b936cb0cd6485 --- /dev/null +++ b/lib/python3.12/site-packages/ray/_private/protobuf_compat.py @@ -0,0 +1,52 @@ +import inspect + +from google.protobuf.json_format import MessageToDict, MessageToJson + +""" +This module provides a compatibility layer for different versions of the protobuf +library. +""" + +_protobuf_has_old_arg_name_cached = None + + +def _protobuf_has_old_arg_name(): + """Cache the inspect result to avoid doing it for every single message.""" + global _protobuf_has_old_arg_name_cached + if _protobuf_has_old_arg_name_cached is None: + params = inspect.signature(MessageToDict).parameters + _protobuf_has_old_arg_name_cached = "including_default_value_fields" in params + return _protobuf_has_old_arg_name_cached + + +def rename_always_print_fields_with_no_presence(kwargs): + """ + Protobuf version 5.26.0rc2 renamed argument for `MessageToDict` and `MessageToJson`: + `including_default_value_fields` -> `always_print_fields_with_no_presence`. + See https://github.com/protocolbuffers/protobuf/commit/06e7caba58ede0220b110b89d08f329e5f8a7537#diff-8de817c14d6a087981503c9aea38730b1b3e98f4e306db5ff9d525c7c304f234L129 # noqa: E501 + + We choose to always use the new argument name. If user used the old arg, we raise an + error. + + If protobuf does not have the new arg name but have the old arg name, we rename our + arg to the old one. + """ + old_arg_name = "including_default_value_fields" + new_arg_name = "always_print_fields_with_no_presence" + if old_arg_name in kwargs: + raise ValueError(f"{old_arg_name} is deprecated, please use {new_arg_name}") + + if new_arg_name in kwargs and _protobuf_has_old_arg_name(): + kwargs[old_arg_name] = kwargs.pop(new_arg_name) + + return kwargs + + +def message_to_dict(*args, **kwargs): + kwargs = rename_always_print_fields_with_no_presence(kwargs) + return MessageToDict(*args, **kwargs) + + +def message_to_json(*args, **kwargs): + kwargs = rename_always_print_fields_with_no_presence(kwargs) + return MessageToJson(*args, **kwargs) diff --git a/lib/python3.12/site-packages/ray/_private/ray_client_microbenchmark.py b/lib/python3.12/site-packages/ray/_private/ray_client_microbenchmark.py new file mode 100644 index 0000000000000000000000000000000000000000..c3cfec3a8645bd935aaa01e8ddb27af7d6205ec0 --- /dev/null +++ b/lib/python3.12/site-packages/ray/_private/ray_client_microbenchmark.py @@ -0,0 +1,117 @@ +import inspect +import logging +import sys + +import numpy as np + +from ray._private.ray_microbenchmark_helpers import timeit +from ray.util.client.ray_client_helpers import ray_start_client_server + + +def benchmark_get_calls(ray, results): + value = ray.put(0) + + def get_small(): + ray.get(value) + + results += timeit("client: get calls", get_small) + + +def benchmark_tasks_and_get_batch(ray, results): + @ray.remote + def small_value(): + return b"ok" + + def small_value_batch(): + submitted = [small_value.remote() for _ in range(1000)] + ray.get(submitted) + return 0 + + results += timeit("client: tasks and get batch", small_value_batch) + + +def benchmark_put_calls(ray, results): + def put_small(): + ray.put(0) + + results += timeit("client: put calls", put_small) + + +def benchmark_remote_put_calls(ray, results): + @ray.remote + def do_put_small(): + for _ in range(100): + ray.put(0) + + def put_multi_small(): + ray.get([do_put_small.remote() for _ in range(10)]) + + results += timeit("client: tasks and put batch", put_multi_small, 1000) + + +def benchmark_put_large(ray, results): + arr = np.zeros(100 * 1024 * 1024, dtype=np.int64) + + def put_large(): + ray.put(arr) + + results += timeit("client: put gigabytes", put_large, 8 * 0.1) + + +def benchmark_simple_actor(ray, results): + @ray.remote(num_cpus=0) + class Actor: + def small_value(self): + return b"ok" + + def small_value_arg(self, x): + return b"ok" + + def small_value_batch(self, n): + ray.get([self.small_value.remote() for _ in range(n)]) + + a = Actor.remote() + + def actor_sync(): + ray.get(a.small_value.remote()) + + results += timeit("client: 1:1 actor calls sync", actor_sync) + + def actor_async(): + ray.get([a.small_value.remote() for _ in range(1000)]) + + results += timeit("client: 1:1 actor calls async", actor_async, 1000) + + a = Actor.options(max_concurrency=16).remote() + + def actor_concurrent(): + ray.get([a.small_value.remote() for _ in range(1000)]) + + results += timeit("client: 1:1 actor calls concurrent", actor_concurrent, 1000) + + +def main(results=None): + results = results or [] + + ray_config = {"logging_level": logging.WARNING} + + def ray_connect_handler(job_config=None, **ray_init_kwargs): + from ray._private.client_mode_hook import disable_client_hook + + with disable_client_hook(): + import ray as real_ray + + if not real_ray.is_initialized(): + real_ray.init(**ray_config) + + for name, obj in inspect.getmembers(sys.modules[__name__]): + if not name.startswith("benchmark_"): + continue + with ray_start_client_server(ray_connect_handler=ray_connect_handler) as ray: + obj(ray, results) + + return results + + +if __name__ == "__main__": + main() diff --git a/lib/python3.12/site-packages/ray/_private/ray_cluster_perf.py b/lib/python3.12/site-packages/ray/_private/ray_cluster_perf.py new file mode 100644 index 0000000000000000000000000000000000000000..11d23fc85350eaf67d0162df23d229fc42525b68 --- /dev/null +++ b/lib/python3.12/site-packages/ray/_private/ray_cluster_perf.py @@ -0,0 +1,51 @@ +"""This is the script for `ray clusterbenchmark`.""" + +import time + +import numpy as np + +import ray +from ray.cluster_utils import Cluster + + +def main(): + cluster = Cluster( + initialize_head=True, + connect=True, + head_node_args={"object_store_memory": 20 * 1024 * 1024 * 1024, "num_cpus": 16}, + ) + cluster.add_node( + object_store_memory=20 * 1024 * 1024 * 1024, num_gpus=1, num_cpus=16 + ) + + object_ref_list = [] + for i in range(0, 10): + object_ref = ray.put(np.random.rand(1024 * 128, 1024)) + object_ref_list.append(object_ref) + + @ray.remote(num_gpus=1) + def f(object_ref_list): + diffs = [] + for object_ref in object_ref_list: + before = time.time() + ray.get(object_ref) + after = time.time() + diffs.append(after - before) + time.sleep(1) + return np.mean(diffs), np.std(diffs) + + time_diff, time_diff_std = ray.get(f.remote(object_ref_list)) + + print( + "latency to get an 1G object over network", + round(time_diff, 2), + "+-", + round(time_diff_std, 2), + ) + + ray.shutdown() + cluster.shutdown() + + +if __name__ == "__main__": + main() diff --git a/lib/python3.12/site-packages/ray/_private/ray_constants.py b/lib/python3.12/site-packages/ray/_private/ray_constants.py new file mode 100644 index 0000000000000000000000000000000000000000..f933c053b7c0c875c847f5d101bbba35558c4213 --- /dev/null +++ b/lib/python3.12/site-packages/ray/_private/ray_constants.py @@ -0,0 +1,587 @@ +"""Ray constants used in the Python code.""" + +import json +import logging +import os +import sys + +logger = logging.getLogger(__name__) + + +def env_integer(key, default): + if key in os.environ: + value = os.environ[key] + if value.isdigit(): + return int(os.environ[key]) + + logger.debug( + f"Found {key} in environment, but value must " + f"be an integer. Got: {value}. Returning " + f"provided default {default}." + ) + return default + return default + + +def env_float(key, default): + if key in os.environ: + value = os.environ[key] + try: + return float(value) + except ValueError: + logger.debug( + f"Found {key} in environment, but value must " + f"be a float. Got: {value}. Returning " + f"provided default {default}." + ) + return default + return default + + +def env_bool(key, default): + if key in os.environ: + return ( + True + if os.environ[key].lower() == "true" or os.environ[key] == "1" + else False + ) + return default + + +def env_set_by_user(key): + return key in os.environ + + +# Whether event logging to driver is enabled. Set to 0 to disable. +AUTOSCALER_EVENTS = env_integer("RAY_SCHEDULER_EVENTS", 1) + +RAY_LOG_TO_DRIVER = env_bool("RAY_LOG_TO_DRIVER", True) + +# Filter level under which events will be filtered out, i.e. not printing to driver +RAY_LOG_TO_DRIVER_EVENT_LEVEL = os.environ.get("RAY_LOG_TO_DRIVER_EVENT_LEVEL", "INFO") + +# Internal kv keys for storing monitor debug status. +DEBUG_AUTOSCALING_ERROR = "__autoscaling_error" +DEBUG_AUTOSCALING_STATUS = "__autoscaling_status" +DEBUG_AUTOSCALING_STATUS_LEGACY = "__autoscaling_status_legacy" + +ID_SIZE = 28 + +# The following constants are used to create default values for +# resource isolation when it is enabled. +# TODO(54703): Link to OSS documentation about the feature once it's available. +DEFAULT_CGROUP_PATH = "/sys/fs/cgroup" +# The default proportion of cpu cores to reserve for ray system processes. +DEFAULT_SYSTEM_RESERVED_CPU_PROPORTION = env_float( + "RAY_DEFAULT_SYSTEM_RESERVED_CPU_PROPORTION", 0.05 +) +# The default minimum number of cpu cores to reserve for ray system processes. +# This value is used if the available_cores * DEFAULT_SYSTEM_RESERVED_CPU_PROPORTION < this value. +DEFAULT_MIN_SYSTEM_RESERVED_CPU_CORES = env_float( + "RAY_DEFAULT_MIN_SYSTEM_RESERVED_CPU_CORES", 1.0 +) +# The default maximum number of cpu cores to reserve for ray system processes. +# This value is used if the available_cores * DEFAULT_SYSTEM_RESERVED_CPU_PROPORTION > this value. +DEFAULT_MAX_SYSTEM_RESERVED_CPU_CORES = env_float( + "RAY_DEFAULT_MAX_SYSTEM_RESERVED_CPU_CORES", 3.0 +) +# The values for SYSTEM_RESERVED_MEMORY do not include the memory reserveed +# for the object store. +# The default proportion available memory to reserve for ray system processes. +DEFAULT_SYSTEM_RESERVED_MEMORY_PROPORTION = env_float( + "RAY_DEFAULT_SYSTEM_RESERVED_MEMORY_PROPORTION", 0.10 +) +# The default minimum number of bytes to reserve for ray system processes. +# This value is used if the available_memory * DEFAULT_SYSTEM_RESERVED_MEMORY_PROPORTION < this value. +DEFAULT_MIN_SYSTEM_RESERVED_MEMORY_BYTES = env_integer( + "RAY_DEFAULT_MIN_SYSTEM_RESERVED_MEMORY_BYTES", (500) * (1024**2) +) +# The default maximum number of bytes to reserve for ray system processes. +# This value is used if the available_memory * DEFAULT_SYSTEM_RESERVED_MEMORY_PROPORTION > this value. +DEFAULT_MAX_SYSTEM_RESERVED_MEMORY_BYTES = env_integer( + "RAY_DEFAULT_MAX_SYSTEM_RESERVED_MEMORY_BYTES", (10) * (1024**3) +) + +# The default maximum number of bytes to allocate to the object store unless +# overridden by the user. +DEFAULT_OBJECT_STORE_MAX_MEMORY_BYTES = env_integer( + "RAY_DEFAULT_OBJECT_STORE_MAX_MEMORY_BYTES", (200) * (10**9) # 200 GB +) +# The default proportion of available memory allocated to the object store +DEFAULT_OBJECT_STORE_MEMORY_PROPORTION = env_float( + "RAY_DEFAULT_OBJECT_STORE_MEMORY_PROPORTION", + 0.3, +) +# The smallest cap on the memory used by the object store that we allow. +# This must be greater than MEMORY_RESOURCE_UNIT_BYTES +OBJECT_STORE_MINIMUM_MEMORY_BYTES = 75 * 1024 * 1024 +# Each ObjectRef currently uses about 3KB of caller memory. +CALLER_MEMORY_USAGE_PER_OBJECT_REF = 3000 +# Above this number of bytes, raise an error by default unless the user sets +# RAY_ALLOW_SLOW_STORAGE=1. This avoids swapping with large object stores. +REQUIRE_SHM_SIZE_THRESHOLD = 10**10 +# Mac with 16GB memory has degraded performance when the object store size is +# greater than 2GB. +# (see https://github.com/ray-project/ray/issues/20388 for details) +# The workaround here is to limit capacity to 2GB for Mac by default, +# and raise error if the capacity is overwritten by user. +MAC_DEGRADED_PERF_MMAP_SIZE_LIMIT = (2) * (2**30) +# If a user does not specify a port for the primary Ray service, +# we attempt to start the service running at this port. +DEFAULT_PORT = 6379 + +RAY_ADDRESS_ENVIRONMENT_VARIABLE = "RAY_ADDRESS" +RAY_API_SERVER_ADDRESS_ENVIRONMENT_VARIABLE = "RAY_API_SERVER_ADDRESS" +RAY_NAMESPACE_ENVIRONMENT_VARIABLE = "RAY_NAMESPACE" +RAY_RUNTIME_ENV_ENVIRONMENT_VARIABLE = "RAY_RUNTIME_ENV" +RAY_RUNTIME_ENV_URI_PIN_EXPIRATION_S_ENV_VAR = ( + "RAY_RUNTIME_ENV_TEMPORARY_REFERENCE_EXPIRATION_S" +) +# Ray populates this env var to the working dir in the creation of a runtime env. +# For example, `pip` and `conda` users can use this environment variable to locate the +# `requirements.txt` file. +RAY_RUNTIME_ENV_CREATE_WORKING_DIR_ENV_VAR = "RAY_RUNTIME_ENV_CREATE_WORKING_DIR" +# Defaults to 10 minutes. This should be longer than the total time it takes for +# the local working_dir and py_modules to be uploaded, or these files might get +# garbage collected before the job starts. +RAY_RUNTIME_ENV_URI_PIN_EXPIRATION_S_DEFAULT = 10 * 60 +# If set to 1, then `.gitignore` files will not be parsed and loaded into "excludes" +# when using a local working_dir or py_modules. +RAY_RUNTIME_ENV_IGNORE_GITIGNORE = "RAY_RUNTIME_ENV_IGNORE_GITIGNORE" +# Hook for running a user-specified runtime-env hook. This hook will be called +# unconditionally given the runtime_env dict passed for ray.init. It must return +# a rewritten runtime_env dict. Example: "your.module.runtime_env_hook". +RAY_RUNTIME_ENV_HOOK = "RAY_RUNTIME_ENV_HOOK" +# Hook that is invoked on `ray start`. It will be given the cluster parameters and +# whether we are the head node as arguments. The function can modify the params class, +# but otherwise returns void. Example: "your.module.ray_start_hook". +RAY_START_HOOK = "RAY_START_HOOK" +# Hook that is invoked on `ray job submit`. It will be given all the same args as the +# job.cli.submit() function gets, passed as kwargs to this function. +RAY_JOB_SUBMIT_HOOK = "RAY_JOB_SUBMIT_HOOK" +# Headers to pass when using the Job CLI. It will be given to +# instantiate a Job SubmissionClient. +RAY_JOB_HEADERS = "RAY_JOB_HEADERS" + +# Timeout waiting for the dashboard to come alive during node startup. +RAY_DASHBOARD_STARTUP_TIMEOUT_S = env_integer("RAY_DASHBOARD_STARTUP_TIMEOUT_S", 60) + +DEFAULT_DASHBOARD_IP = "127.0.0.1" +DEFAULT_DASHBOARD_PORT = 8265 +DASHBOARD_ADDRESS = "dashboard" +DASHBOARD_CLIENT_MAX_SIZE = 100 * 1024**2 +PROMETHEUS_SERVICE_DISCOVERY_FILE = "prom_metrics_service_discovery.json" +DEFAULT_DASHBOARD_AGENT_LISTEN_PORT = 52365 +# Default resource requirements for actors when no resource requirements are +# specified. +DEFAULT_ACTOR_METHOD_CPU_SIMPLE = 1 +DEFAULT_ACTOR_CREATION_CPU_SIMPLE = 0 +# Default resource requirements for actors when some resource requirements are +# specified in . +DEFAULT_ACTOR_METHOD_CPU_SPECIFIED = 0 +DEFAULT_ACTOR_CREATION_CPU_SPECIFIED = 1 +# Default number of return values for each actor method. +DEFAULT_ACTOR_METHOD_NUM_RETURN_VALS = 1 + +# Wait 30 seconds for client to reconnect after unexpected disconnection +DEFAULT_CLIENT_RECONNECT_GRACE_PERIOD = 30 + +# If a remote function or actor (or some other export) has serialized size +# greater than this quantity, print an warning. +FUNCTION_SIZE_WARN_THRESHOLD = 10**7 +FUNCTION_SIZE_ERROR_THRESHOLD = env_integer("FUNCTION_SIZE_ERROR_THRESHOLD", (10**8)) + +# If remote functions with the same source are imported this many times, then +# print a warning. +DUPLICATE_REMOTE_FUNCTION_THRESHOLD = 100 + +# The maximum resource quantity that is allowed. TODO(rkn): This could be +# relaxed, but the current implementation of the node manager will be slower +# for large resource quantities due to bookkeeping of specific resource IDs. +MAX_RESOURCE_QUANTITY = 100e12 + +# Number of units 1 resource can be subdivided into. +MIN_RESOURCE_GRANULARITY = 0.0001 + +# Set this environment variable to populate the dashboard URL with +# an external hosted Ray dashboard URL (e.g. because the +# dashboard is behind a proxy or load balancer). This only overrides +# the dashboard URL when returning or printing to a user through a public +# API, but not in the internal KV store. +RAY_OVERRIDE_DASHBOARD_URL = "RAY_OVERRIDE_DASHBOARD_URL" + + +# Different types of Ray errors that can be pushed to the driver. +# TODO(rkn): These should be defined in flatbuffers and must be synced with +# the existing C++ definitions. +PICKLING_LARGE_OBJECT_PUSH_ERROR = "pickling_large_object" +WAIT_FOR_FUNCTION_PUSH_ERROR = "wait_for_function" +VERSION_MISMATCH_PUSH_ERROR = "version_mismatch" +WORKER_CRASH_PUSH_ERROR = "worker_crash" +WORKER_DIED_PUSH_ERROR = "worker_died" +WORKER_POOL_LARGE_ERROR = "worker_pool_large" +PUT_RECONSTRUCTION_PUSH_ERROR = "put_reconstruction" +RESOURCE_DEADLOCK_ERROR = "resource_deadlock" +REMOVED_NODE_ERROR = "node_removed" +MONITOR_DIED_ERROR = "monitor_died" +LOG_MONITOR_DIED_ERROR = "log_monitor_died" +DASHBOARD_AGENT_DIED_ERROR = "dashboard_agent_died" +DASHBOARD_DIED_ERROR = "dashboard_died" +RAYLET_DIED_ERROR = "raylet_died" +DETACHED_ACTOR_ANONYMOUS_NAMESPACE_ERROR = "detached_actor_anonymous_namespace" +EXCESS_QUEUEING_WARNING = "excess_queueing_warning" + +# Used by autoscaler to set the node custom resources and labels +# from cluster.yaml. +RESOURCES_ENVIRONMENT_VARIABLE = "RAY_OVERRIDE_RESOURCES" +LABELS_ENVIRONMENT_VARIABLE = "RAY_OVERRIDE_LABELS" + +# Temporary flag to disable log processing in the dashboard. This is useful +# if the dashboard is overloaded by logs and failing to process other +# dashboard API requests (e.g. Job Submission). +DISABLE_DASHBOARD_LOG_INFO = env_integer("RAY_DISABLE_DASHBOARD_LOG_INFO", 0) + +LOGGER_FORMAT = "%(asctime)s\t%(levelname)s %(filename)s:%(lineno)s -- %(message)s" +LOGGER_FORMAT_ESCAPE = json.dumps(LOGGER_FORMAT.replace("%", "%%")) +LOGGER_FORMAT_HELP = f"The logging format. default={LOGGER_FORMAT_ESCAPE}" +# Configure the default logging levels for various Ray components. +# TODO (kevin85421): Currently, I don't encourage Ray users to configure +# `RAY_LOGGER_LEVEL` until its scope and expected behavior are clear and +# easy to understand. Now, only Ray developers should use it. +LOGGER_LEVEL = os.environ.get("RAY_LOGGER_LEVEL", "info") +LOGGER_LEVEL_CHOICES = ["debug", "info", "warning", "error", "critical"] +LOGGER_LEVEL_HELP = ( + "The logging level threshold, choices=['debug', 'info'," + " 'warning', 'error', 'critical'], default='info'" +) + +LOGGING_REDIRECT_STDERR_ENVIRONMENT_VARIABLE = "RAY_LOG_TO_STDERR" +# Logging format when logging stderr. This should be formatted with the +# component before setting the formatter, e.g. via +# format = LOGGER_FORMAT_STDERR.format(component="dashboard") +# handler.setFormatter(logging.Formatter(format)) +LOGGER_FORMAT_STDERR = ( + "%(asctime)s\t%(levelname)s ({component}) %(filename)s:%(lineno)s -- %(message)s" +) + +# Constants used to define the different process types. +PROCESS_TYPE_REAPER = "reaper" +PROCESS_TYPE_MONITOR = "monitor" +PROCESS_TYPE_RAY_CLIENT_SERVER = "ray_client_server" +PROCESS_TYPE_LOG_MONITOR = "log_monitor" +PROCESS_TYPE_DASHBOARD = "dashboard" +PROCESS_TYPE_DASHBOARD_AGENT = "dashboard_agent" +PROCESS_TYPE_RUNTIME_ENV_AGENT = "runtime_env_agent" +PROCESS_TYPE_WORKER = "worker" +PROCESS_TYPE_RAYLET = "raylet" +PROCESS_TYPE_REDIS_SERVER = "redis_server" +PROCESS_TYPE_GCS_SERVER = "gcs_server" +PROCESS_TYPE_PYTHON_CORE_WORKER_DRIVER = "python-core-driver" +PROCESS_TYPE_PYTHON_CORE_WORKER = "python-core-worker" + +# Log file names +MONITOR_LOG_FILE_NAME = f"{PROCESS_TYPE_MONITOR}.log" +LOG_MONITOR_LOG_FILE_NAME = f"{PROCESS_TYPE_LOG_MONITOR}.log" + +# Enable log deduplication. +RAY_DEDUP_LOGS = env_bool("RAY_DEDUP_LOGS", True) + +# How many seconds of messages to buffer for log deduplication. +RAY_DEDUP_LOGS_AGG_WINDOW_S = env_integer("RAY_DEDUP_LOGS_AGG_WINDOW_S", 5) + +# Regex for log messages to never deduplicate, or None. This takes precedence over +# the skip regex below. A default pattern is set for testing. +TESTING_NEVER_DEDUP_TOKEN = "__ray_testing_never_deduplicate__" +RAY_DEDUP_LOGS_ALLOW_REGEX = os.environ.get( + "RAY_DEDUP_LOGS_ALLOW_REGEX", TESTING_NEVER_DEDUP_TOKEN +) + +# Regex for log messages to always skip / suppress, or None. +RAY_DEDUP_LOGS_SKIP_REGEX = os.environ.get("RAY_DEDUP_LOGS_SKIP_REGEX") + +WORKER_PROCESS_TYPE_IDLE_WORKER = "ray::IDLE" +WORKER_PROCESS_TYPE_SPILL_WORKER_NAME = "SpillWorker" +WORKER_PROCESS_TYPE_RESTORE_WORKER_NAME = "RestoreWorker" +WORKER_PROCESS_TYPE_SPILL_WORKER_IDLE = ( + f"ray::IDLE_{WORKER_PROCESS_TYPE_SPILL_WORKER_NAME}" +) +WORKER_PROCESS_TYPE_RESTORE_WORKER_IDLE = ( + f"ray::IDLE_{WORKER_PROCESS_TYPE_RESTORE_WORKER_NAME}" +) +WORKER_PROCESS_TYPE_SPILL_WORKER = f"ray::SPILL_{WORKER_PROCESS_TYPE_SPILL_WORKER_NAME}" +WORKER_PROCESS_TYPE_RESTORE_WORKER = ( + f"ray::RESTORE_{WORKER_PROCESS_TYPE_RESTORE_WORKER_NAME}" +) +WORKER_PROCESS_TYPE_SPILL_WORKER_DELETE = ( + f"ray::DELETE_{WORKER_PROCESS_TYPE_SPILL_WORKER_NAME}" +) +WORKER_PROCESS_TYPE_RESTORE_WORKER_DELETE = ( + f"ray::DELETE_{WORKER_PROCESS_TYPE_RESTORE_WORKER_NAME}" +) + +# The number of files the log monitor will open. If more files exist, they will +# be ignored. +LOG_MONITOR_MAX_OPEN_FILES = int( + os.environ.get("RAY_LOG_MONITOR_MAX_OPEN_FILES", "200") +) + +# The maximum batch of lines to be read in a single iteration. We _always_ try +# to read this number of lines even if there aren't any new lines. +LOG_MONITOR_NUM_LINES_TO_READ = int( + os.environ.get("RAY_LOG_MONITOR_NUM_LINES_TO_READ", "1000") +) + +# Autoscaler events are denoted by the ":event_summary:" magic token. +LOG_PREFIX_EVENT_SUMMARY = ":event_summary:" +# Cluster-level info events are denoted by the ":info_message:" magic token. These may +# be emitted in the stderr of Ray components. +LOG_PREFIX_INFO_MESSAGE = ":info_message:" +# Actor names are recorded in the logs with this magic token as a prefix. +LOG_PREFIX_ACTOR_NAME = ":actor_name:" +# Task names are recorded in the logs with this magic token as a prefix. +LOG_PREFIX_TASK_NAME = ":task_name:" +# Job ids are recorded in the logs with this magic token as a prefix. +LOG_PREFIX_JOB_ID = ":job_id:" + +# The object metadata field uses the following format: It is a comma +# separated list of fields. The first field is mandatory and is the +# type of the object (see types below) or an integer, which is interpreted +# as an error value. The second part is optional and if present has the +# form DEBUG:, it is used for implementing the debugger. + +# A constant used as object metadata to indicate the object is cross language. +OBJECT_METADATA_TYPE_CROSS_LANGUAGE = b"XLANG" +# A constant used as object metadata to indicate the object is python specific. +OBJECT_METADATA_TYPE_PYTHON = b"PYTHON" +# A constant used as object metadata to indicate the object is raw bytes. +OBJECT_METADATA_TYPE_RAW = b"RAW" + +# A constant used as object metadata to indicate the object is an actor handle. +# This value should be synchronized with the Java definition in +# ObjectSerializer.java +# TODO(fyrestone): Serialize the ActorHandle via the custom type feature +# of XLANG. +OBJECT_METADATA_TYPE_ACTOR_HANDLE = b"ACTOR_HANDLE" + +# A constant indicating the debugging part of the metadata (see above). +OBJECT_METADATA_DEBUG_PREFIX = b"DEBUG:" + +AUTOSCALER_RESOURCE_REQUEST_CHANNEL = b"autoscaler_resource_request" + +REDIS_DEFAULT_USERNAME = "" + +REDIS_DEFAULT_PASSWORD = "" + +# The default ip address to bind to. +NODE_DEFAULT_IP = "127.0.0.1" + +# The Mach kernel page size in bytes. +MACH_PAGE_SIZE_BYTES = 4096 + +# The max number of bytes for task execution error message. +MAX_APPLICATION_ERROR_LEN = 500 + +# Max 64 bit integer value, which is needed to ensure against overflow +# in C++ when passing integer values cross-language. +MAX_INT64_VALUE = 9223372036854775807 + +# Object Spilling related constants +DEFAULT_OBJECT_PREFIX = "ray_spilled_objects" + +GCS_PORT_ENVIRONMENT_VARIABLE = "RAY_GCS_SERVER_PORT" + +HEALTHCHECK_EXPIRATION_S = os.environ.get("RAY_HEALTHCHECK_EXPIRATION_S", 10) + +# Filename of "shim process" that sets up Python worker environment. +# Should be kept in sync with kSetupWorkerFilename in +# src/ray/common/constants.h. +SETUP_WORKER_FILENAME = "setup_worker.py" + +# Directory name where runtime_env resources will be created & cached. +DEFAULT_RUNTIME_ENV_DIR_NAME = "runtime_resources" + +# The timeout seconds for the creation of runtime env, +# dafault timeout is 10 minutes +DEFAULT_RUNTIME_ENV_TIMEOUT_SECONDS = 600 + +# Used to separate lines when formatting the call stack where an ObjectRef was +# created. +CALL_STACK_LINE_DELIMITER = " | " + +# The default gRPC max message size is 4 MiB, we use a larger number of 512 MiB +# NOTE: This is equal to the C++ limit of (RAY_CONFIG::max_grpc_message_size) +GRPC_CPP_MAX_MESSAGE_SIZE = 512 * 1024 * 1024 + +# The gRPC send & receive max length for "dashboard agent" server. +# NOTE: This is equal to the C++ limit of RayConfig::max_grpc_message_size +# and HAVE TO STAY IN SYNC with it (ie, meaning that both of these values +# have to be set at the same time) +AGENT_GRPC_MAX_MESSAGE_LENGTH = env_integer( + "AGENT_GRPC_MAX_MESSAGE_LENGTH", 20 * 1024 * 1024 # 20MB +) + + +# GRPC options +GRPC_ENABLE_HTTP_PROXY = ( + 1 + if os.environ.get("RAY_grpc_enable_http_proxy", "0").lower() in ("1", "true") + else 0 +) +GLOBAL_GRPC_OPTIONS = (("grpc.enable_http_proxy", GRPC_ENABLE_HTTP_PROXY),) + +# Internal kv namespaces +KV_NAMESPACE_DASHBOARD = b"dashboard" +KV_NAMESPACE_SESSION = b"session" +KV_NAMESPACE_TRACING = b"tracing" +KV_NAMESPACE_PDB = b"ray_pdb" +KV_NAMESPACE_HEALTHCHECK = b"healthcheck" +KV_NAMESPACE_JOB = b"job" +KV_NAMESPACE_CLUSTER = b"cluster" +KV_HEAD_NODE_ID_KEY = b"head_node_id" +# TODO: Set package for runtime env +# We need to update ray client for this since runtime env use ray client +# This might introduce some compatibility issues so leave it here for now. +KV_NAMESPACE_PACKAGE = None +KV_NAMESPACE_FUNCTION_TABLE = b"fun" + +LANGUAGE_WORKER_TYPES = ["python", "java", "cpp"] + +NEURON_CORES = "neuron_cores" +GPU = "GPU" +TPU = "TPU" +NPU = "NPU" +HPU = "HPU" + + +RAY_WORKER_NICENESS = "RAY_worker_niceness" + +# Default max_retries option in @ray.remote for non-actor +# tasks. +DEFAULT_TASK_MAX_RETRIES = 3 + +# Default max_concurrency option in @ray.remote for threaded actors. +DEFAULT_MAX_CONCURRENCY_THREADED = 1 + +# Ray internal flags. These flags should not be set by users, and we strip them on job +# submission. +# This should be consistent with src/ray/common/ray_internal_flag_def.h +RAY_INTERNAL_FLAGS = [ + "RAY_JOB_ID", + "RAY_RAYLET_PID", + "RAY_OVERRIDE_NODE_ID_FOR_TESTING", +] + + +def gcs_actor_scheduling_enabled(): + return os.environ.get("RAY_gcs_actor_scheduling_enabled") == "true" + + +DEFAULT_RESOURCES = {"CPU", "GPU", "memory", "object_store_memory"} + +# Supported Python versions for runtime env's "conda" field. Ray downloads +# Ray wheels into the conda environment, so the Ray wheels for these Python +# versions must be available online. +RUNTIME_ENV_CONDA_PY_VERSIONS = [(3, 9), (3, 10), (3, 11), (3, 12)] + +# Whether to enable Ray clusters (in addition to local Ray). +# Ray clusters are not explicitly supported for Windows and OSX. +IS_WINDOWS_OR_OSX = sys.platform == "darwin" or sys.platform == "win32" +ENABLE_RAY_CLUSTERS_ENV_VAR = "RAY_ENABLE_WINDOWS_OR_OSX_CLUSTER" +ENABLE_RAY_CLUSTER = env_bool( + ENABLE_RAY_CLUSTERS_ENV_VAR, + not IS_WINDOWS_OR_OSX, +) + +SESSION_LATEST = "session_latest" +NUM_PORT_RETRIES = 40 +NUM_REDIS_GET_RETRIES = int(os.environ.get("RAY_NUM_REDIS_GET_RETRIES", "20")) + +# The allowed cached ports in Ray. Refer to Port configuration for more details: +# https://docs.ray.io/en/latest/ray-core/configure.html#ports-configurations +RAY_ALLOWED_CACHED_PORTS = { + "metrics_agent_port", + "metrics_export_port", + "dashboard_agent_listen_port", + "runtime_env_agent_port", + "gcs_server_port", # the `port` option for gcs port. +} + +# Turn this on if actor task log's offsets are expected to be recorded. +# With this enabled, actor tasks' log could be queried with task id. +RAY_ENABLE_RECORD_ACTOR_TASK_LOGGING = env_bool( + "RAY_ENABLE_RECORD_ACTOR_TASK_LOGGING", False +) + +# RuntimeEnv env var to indicate it exports a function +WORKER_PROCESS_SETUP_HOOK_ENV_VAR = "__RAY_WORKER_PROCESS_SETUP_HOOK_ENV_VAR" +RAY_WORKER_PROCESS_SETUP_HOOK_LOAD_TIMEOUT_ENV_VAR = ( + "RAY_WORKER_PROCESS_SETUP_HOOK_LOAD_TIMEOUT" # noqa +) + +RAY_DEFAULT_LABEL_KEYS_PREFIX = "ray.io/" + +RAY_TPU_MAX_CONCURRENT_CONNECTIONS_ENV_VAR = "RAY_TPU_MAX_CONCURRENT_ACTIVE_CONNECTIONS" + +RAY_NODE_IP_FILENAME = "node_ip_address.json" + +RAY_LOGGING_CONFIG_ENCODING = os.environ.get("RAY_LOGGING_CONFIG_ENCODING") + +RAY_BACKEND_LOG_JSON_ENV_VAR = "RAY_BACKEND_LOG_JSON" + +# Write export API event of all resource types to file if enabled. +# RAY_enable_export_api_write_config will not be considered if +# this is enabled. +RAY_ENABLE_EXPORT_API_WRITE = env_bool("RAY_enable_export_api_write", False) + +# Comma separated string containing individual resource +# to write export API events for. This configuration is only used if +# RAY_enable_export_api_write is not enabled. Full list of valid +# resource types in ExportEvent.SourceType enum in +# src/ray/protobuf/export_api/export_event.proto +# Example config: +# `export RAY_enable_export_api_write_config='EXPORT_SUBMISSION_JOB,EXPORT_ACTOR'` +RAY_ENABLE_EXPORT_API_WRITE_CONFIG_STR = os.environ.get( + "RAY_enable_export_api_write_config", "" +) +RAY_ENABLE_EXPORT_API_WRITE_CONFIG = RAY_ENABLE_EXPORT_API_WRITE_CONFIG_STR.split(",") + +RAY_EXPORT_EVENT_MAX_FILE_SIZE_BYTES = env_bool( + "RAY_EXPORT_EVENT_MAX_FILE_SIZE_BYTES", 100 * 1e6 +) + +RAY_EXPORT_EVENT_MAX_BACKUP_COUNT = env_bool("RAY_EXPORT_EVENT_MAX_BACKUP_COUNT", 20) + +# If this flag is set and you run the driver with `uv run`, Ray propagates the `uv run` +# environment to all workers. Ray does this by setting the `py_executable` to the +# `uv run`` command line and by propagating the working directory +# via the `working_dir` plugin so uv finds the pyproject.toml. +# If you enable RAY_ENABLE_UV_RUN_RUNTIME_ENV AND you run the driver +# with `uv run`, Ray deactivates the regular RAY_RUNTIME_ENV_HOOK +# because in most cases the hooks wouldn't work unless you specifically make the code +# for the runtime env hook available in your uv environment and make sure your hook +# is compatible with your uv runtime environment. If you want to combine a custom +# RAY_RUNTIME_ENV_HOOK with `uv run`, you should flag off RAY_ENABLE_UV_RUN_RUNTIME_ENV +# and call ray._private.runtime_env.uv_runtime_env_hook.hook manually in your hook or +# manually set the py_executable in your runtime environment hook. +RAY_ENABLE_UV_RUN_RUNTIME_ENV = env_bool("RAY_ENABLE_UV_RUN_RUNTIME_ENV", True) + +# Prometheus metric cardinality level setting, either "legacy" or "recommended". +# +# Legacy: report all metrics to prometheus with the set of labels that are reported by +# the component, including WorkerId, (task or actor) Name, etc. This is the default. +# Recommended: report only the node level metrics to prometheus. This means that the +# WorkerId will be removed from all metrics. +RAY_METRIC_CARDINALITY_LEVEL = os.environ.get("RAY_metric_cardinality_level", "legacy") + +# Whether enable OpenTelemetry as the metrics collection backend. The default is +# using OpenCensus. +RAY_ENABLE_OPEN_TELEMETRY = env_bool("RAY_enable_open_telemetry", True) + +# How long to wait for a fetch for an RDT object to complete during ray.get before timing out and raising an exception to the user. +# +# NOTE: This is a tenth of `RayConfig::fetch_fail_timeout_milliseconds` by default as RDT transfers are expected to be much faster. +RDT_FETCH_FAIL_TIMEOUT_SECONDS = ( + env_integer("RAY_rdt_fetch_fail_timeout_milliseconds", 60000) / 1000 +) + +RAY_GC_MIN_COLLECT_INTERVAL = env_float("RAY_GC_MIN_COLLECT_INTERVAL_S", 5) diff --git a/lib/python3.12/site-packages/ray/_private/ray_logging/__init__.py b/lib/python3.12/site-packages/ray/_private/ray_logging/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..b24e1754ff8a7222373bf6b588a8f248bf45c1e1 --- /dev/null +++ b/lib/python3.12/site-packages/ray/_private/ray_logging/__init__.py @@ -0,0 +1,379 @@ +import logging +import logging.handlers +import os +import re +import sys +import threading +import time +from dataclasses import dataclass +from typing import Any, Callable, Dict, Iterable, List, Optional, Set, Tuple, Union + +import colorama + +import ray +from ray._private.ray_constants import ( + RAY_DEDUP_LOGS, + RAY_DEDUP_LOGS_AGG_WINDOW_S, + RAY_DEDUP_LOGS_ALLOW_REGEX, + RAY_DEDUP_LOGS_SKIP_REGEX, +) +from ray.experimental.tqdm_ray import RAY_TQDM_MAGIC +from ray.util.debug import log_once + + +def setup_logger( + logging_level: int, + logging_format: str, +): + """Setup default logging for ray.""" + logger = logging.getLogger("ray") + if logging_format: + # Overwrite the formatters for all default handlers. + formatter = logging.Formatter(logging_format) + for handler in logger.handlers: + handler.setFormatter(formatter) + if type(logging_level) is str: + logging_level = logging.getLevelName(logging_level.upper()) + logger.setLevel(logging_level) + + +def setup_component_logger( + *, + logging_level, + logging_format, + log_dir, + filename: Union[str, Iterable[str]], + max_bytes, + backup_count, + logger_name=None, + propagate=True, +): + """Configure the logger that is used for Ray's python components. + + For example, it should be used for monitor, dashboard, and log monitor. + The only exception is workers. They use the different logging config. + + Ray's python components generally should not write to stdout/stderr, because + messages written there will be redirected to the head node. For deployments where + there may be thousands of workers, this would create unacceptable levels of log + spam. For this reason, we disable the "ray" logger's handlers, and enable + propagation so that log messages that actually do need to be sent to the head node + can reach it. + + Args: + logging_level: Logging level in string or logging enum. + logging_format: Logging format string. + log_dir: Log directory path. If empty, logs will go to + stderr. + filename: A single filename or an iterable of filenames to write logs to. + If empty, logs will go to stderr. + max_bytes: Same argument as RotatingFileHandler's maxBytes. + backup_count: Same argument as RotatingFileHandler's backupCount. + logger_name: Used to create or get the correspoding + logger in getLogger call. It will get the root logger by default. + propagate: Whether to propagate the log to the parent logger. + Returns: + the created or modified logger. + """ + ray._private.log.clear_logger("ray") + + logger = logging.getLogger(logger_name) + if isinstance(logging_level, str): + logging_level = logging.getLevelName(logging_level.upper()) + logger.setLevel(logging_level) + + filenames = [filename] if isinstance(filename, str) else filename + + for filename in filenames: + if not filename or not log_dir: + handler = logging.StreamHandler() + else: + handler = logging.handlers.RotatingFileHandler( + os.path.join(log_dir, filename), + maxBytes=max_bytes, + backupCount=backup_count, + ) + handler.setLevel(logging_level) + handler.setFormatter(logging.Formatter(logging_format)) + logger.addHandler(handler) + + logger.propagate = propagate + return logger + + +def run_callback_on_events_in_ipython(event: str, cb: Callable): + """ + Register a callback to be run after each cell completes in IPython. + E.g.: + This is used to flush the logs after each cell completes. + + If IPython is not installed, this function does nothing. + + Args: + cb: The callback to run. + """ + if "IPython" in sys.modules: + from IPython import get_ipython + + ipython = get_ipython() + # Register a callback on cell completion. + if ipython is not None: + ipython.events.register(event, cb) + + +""" +All components underneath here is used specifically for the default_worker.py. +""" + +# It's worth noticing that filepath format should be kept in sync with function +# `GetWorkerOutputFilepath` under file "src/ray/core_worker/core_worker_process.cc". +def get_worker_log_file_name(worker_type, job_id=None): + if job_id is None: + job_id = os.environ.get("RAY_JOB_ID") + if worker_type == "WORKER": + if job_id is None: + job_id = "" + worker_name = "worker" + else: + job_id = "" + worker_name = "io_worker" + + # Make sure these values are set already. + assert ray._private.worker._global_node is not None + assert ray._private.worker.global_worker is not None + filename = f"{worker_name}-{ray.get_runtime_context().get_worker_id()}-" + if job_id: + filename += f"{job_id}-" + filename += f"{os.getpid()}" + return filename + + +def configure_log_file(out_file, err_file): + # If either of the file handles are None, there are no log files to + # configure since we're redirecting all output to stdout and stderr. + if out_file is None or err_file is None: + return + stdout_fileno = sys.stdout.fileno() + stderr_fileno = sys.stderr.fileno() + # C++ logging requires redirecting the stdout file descriptor. Note that + # dup2 will automatically close the old file descriptor before overriding + # it. + os.dup2(out_file.fileno(), stdout_fileno) + os.dup2(err_file.fileno(), stderr_fileno) + # We also manually set sys.stdout and sys.stderr because that seems to + # have an effect on the output buffering. Without doing this, stdout + # and stderr are heavily buffered resulting in seemingly lost logging + # statements. We never want to close the stdout file descriptor, dup2 will + # close it when necessary and we don't want python's GC to close it. + sys.stdout = ray._private.utils.open_log( + stdout_fileno, unbuffered=True, closefd=False + ) + sys.stderr = ray._private.utils.open_log( + stderr_fileno, unbuffered=True, closefd=False + ) + + +class WorkerStandardStreamDispatcher: + def __init__(self): + self.handlers = [] + self._lock = threading.Lock() + + def add_handler(self, name: str, handler: Callable) -> None: + with self._lock: + self.handlers.append((name, handler)) + + def remove_handler(self, name: str) -> None: + with self._lock: + new_handlers = [pair for pair in self.handlers if pair[0] != name] + self.handlers = new_handlers + + def emit(self, data): + with self._lock: + for pair in self.handlers: + _, handle = pair + handle(data) + + +global_worker_stdstream_dispatcher = WorkerStandardStreamDispatcher() + + +# Regex for canonicalizing log lines. +NUMBERS = re.compile(r"(\d+|0x[0-9a-fA-F]+)") + +# Batch of log lines including ip, pid, lines, etc. +LogBatch = Dict[str, Any] + + +def _canonicalise_log_line(line): + # Remove words containing numbers or hex, since those tend to differ between + # workers. + return " ".join(x for x in line.split() if not NUMBERS.search(x)) + + +@dataclass +class DedupState: + # Timestamp of the earliest log message seen of this pattern. + timestamp: int + + # The number of un-printed occurrances for this pattern. + count: int + + # Latest instance of this log pattern. + line: int + + # Latest metadata dict for this log pattern, not including the lines field. + metadata: LogBatch + + # Set of (ip, pid) sources which have emitted this pattern. + sources: Set[Tuple[str, int]] + + # The string that should be printed to stdout. + def formatted(self) -> str: + return self.line + _color( + f" [repeated {self.count}x across cluster]" + _warn_once() + ) + + +class LogDeduplicator: + def __init__( + self, + agg_window_s: int, + allow_re: Optional[str], + skip_re: Optional[str], + *, + _timesource=None, + ): + self.agg_window_s = agg_window_s + if allow_re: + self.allow_re = re.compile(allow_re) + else: + self.allow_re = None + if skip_re: + self.skip_re = re.compile(skip_re) + else: + self.skip_re = None + # Buffer of up to RAY_DEDUP_LOGS_AGG_WINDOW_S recent log patterns. + # This buffer is cleared if the pattern isn't seen within the window. + self.recent: Dict[str, DedupState] = {} + self.timesource = _timesource or (lambda: time.time()) + + run_callback_on_events_in_ipython("post_execute", self.flush) + + def deduplicate(self, batch: LogBatch) -> List[LogBatch]: + """Rewrite a batch of lines to reduce duplicate log messages. + + Args: + batch: The batch of lines from a single source. + + Returns: + List of batches from this and possibly other previous sources to print. + """ + if not RAY_DEDUP_LOGS: + return [batch] + + now = self.timesource() + metadata = batch.copy() + del metadata["lines"] + source = (metadata.get("ip"), metadata.get("pid")) + output: List[LogBatch] = [dict(**metadata, lines=[])] + + # Decide which lines to emit from the input batch. Put the outputs in the + # first output log batch (output[0]). + for line in batch["lines"]: + if RAY_TQDM_MAGIC in line or (self.allow_re and self.allow_re.search(line)): + output[0]["lines"].append(line) + continue + elif self.skip_re and self.skip_re.search(line): + continue + dedup_key = _canonicalise_log_line(line) + + if dedup_key == "": + # Don't dedup messages that are empty after canonicalization. + # Because that's all the information users want to see. + output[0]["lines"].append(line) + continue + + if dedup_key in self.recent: + sources = self.recent[dedup_key].sources + sources.add(source) + # We deduplicate the warnings/error messages from raylet by default. + if len(sources) > 1 or batch["pid"] == "raylet": + state = self.recent[dedup_key] + self.recent[dedup_key] = DedupState( + state.timestamp, + state.count + 1, + line, + metadata, + sources, + ) + else: + # Don't dedup messages from the same source, just print. + output[0]["lines"].append(line) + else: + self.recent[dedup_key] = DedupState(now, 0, line, metadata, {source}) + output[0]["lines"].append(line) + + # Flush patterns from the buffer that are older than the aggregation window. + while self.recent: + if now - next(iter(self.recent.values())).timestamp < self.agg_window_s: + break + dedup_key = next(iter(self.recent)) + state = self.recent.pop(dedup_key) + # we already logged an instance of this line immediately when received, + # so don't log for count == 0 + if state.count > 1: + # (Actor pid=xxxx) [repeated 2x across cluster] ... + output.append(dict(**state.metadata, lines=[state.formatted()])) + # Continue aggregating for this key but reset timestamp and count. + state.timestamp = now + state.count = 0 + self.recent[dedup_key] = state + elif state.count > 0: + # Aggregation wasn't fruitful, print the line and stop aggregating. + output.append(dict(state.metadata, lines=[state.line])) + + return output + + def flush(self) -> List[dict]: + """Return all buffered log messages and clear the buffer. + + Returns: + List of log batches to print. + """ + output = [] + for state in self.recent.values(): + if state.count > 1: + output.append( + dict( + state.metadata, + lines=[state.formatted()], + ) + ) + elif state.count > 0: + output.append(dict(state.metadata, **{"lines": [state.line]})) + self.recent.clear() + return output + + +def _warn_once() -> str: + if log_once("log_dedup_warning"): + return ( + " (Ray deduplicates logs by default. Set RAY_DEDUP_LOGS=0 to " + "disable log deduplication, or see https://docs.ray.io/en/master/" + "ray-observability/user-guides/configure-logging.html#log-deduplication " + "for more options.)" + ) + else: + return "" + + +def _color(msg: str) -> str: + return "{}{}{}".format(colorama.Fore.GREEN, msg, colorama.Style.RESET_ALL) + + +stdout_deduplicator = LogDeduplicator( + RAY_DEDUP_LOGS_AGG_WINDOW_S, RAY_DEDUP_LOGS_ALLOW_REGEX, RAY_DEDUP_LOGS_SKIP_REGEX +) +stderr_deduplicator = LogDeduplicator( + RAY_DEDUP_LOGS_AGG_WINDOW_S, RAY_DEDUP_LOGS_ALLOW_REGEX, RAY_DEDUP_LOGS_SKIP_REGEX +) diff --git a/lib/python3.12/site-packages/ray/_private/ray_logging/__pycache__/__init__.cpython-312.pyc b/lib/python3.12/site-packages/ray/_private/ray_logging/__pycache__/__init__.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..1e8ff77d77466577e73df14aa7936d96423baf0f Binary files /dev/null and b/lib/python3.12/site-packages/ray/_private/ray_logging/__pycache__/__init__.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/_private/ray_logging/__pycache__/constants.cpython-312.pyc b/lib/python3.12/site-packages/ray/_private/ray_logging/__pycache__/constants.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..44262a175af03e513022a70f750532eb8aeb8666 Binary files /dev/null and b/lib/python3.12/site-packages/ray/_private/ray_logging/__pycache__/constants.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/_private/ray_logging/__pycache__/default_impl.cpython-312.pyc b/lib/python3.12/site-packages/ray/_private/ray_logging/__pycache__/default_impl.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..c7f8e0b9432dab88510508736a87448a98072f1e Binary files /dev/null and b/lib/python3.12/site-packages/ray/_private/ray_logging/__pycache__/default_impl.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/_private/ray_logging/__pycache__/logging_config.cpython-312.pyc b/lib/python3.12/site-packages/ray/_private/ray_logging/__pycache__/logging_config.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..7f17f458ef15e512b07e7a3c1aa93497c04625ce Binary files /dev/null and b/lib/python3.12/site-packages/ray/_private/ray_logging/__pycache__/logging_config.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/_private/ray_logging/constants.py b/lib/python3.12/site-packages/ray/_private/ray_logging/constants.py new file mode 100644 index 0000000000000000000000000000000000000000..a5bf5850a70861f4142060aaed6f422937981cd5 --- /dev/null +++ b/lib/python3.12/site-packages/ray/_private/ray_logging/constants.py @@ -0,0 +1,59 @@ +from enum import Enum + +# A set containing the standard attributes of a LogRecord. This is used to +# help us determine which attributes constitute Ray or user-provided context. It is +# also be used to determine whether a attribute is a standard python logging attribute. +# http://docs.python.org/library/logging.html#logrecord-attributes +LOGRECORD_STANDARD_ATTRS = { + "args", + "asctime", + "created", + "exc_info", + "exc_text", + "filename", + "funcName", + "levelname", + "levelno", + "lineno", + "message", + "module", + "msecs", + "msg", + "name", + "pathname", + "process", + "processName", + "relativeCreated", + "stack_info", + "thread", + "threadName", + "taskName", +} + +LOGGER_FLATTEN_KEYS = { + "ray_serve_extra_fields", +} + + +class LogKey(str, Enum): + # Core context + JOB_ID = "job_id" + WORKER_ID = "worker_id" + NODE_ID = "node_id" + ACTOR_ID = "actor_id" + TASK_ID = "task_id" + ACTOR_NAME = "actor_name" + TASK_NAME = "task_name" + TASK_FUNCTION_NAME = "task_func_name" + + # Logger built-in context + ASCTIME = "asctime" + LEVELNAME = "levelname" + MESSAGE = "message" + FILENAME = "filename" + LINENO = "lineno" + EXC_TEXT = "exc_text" + PROCESS = "process" + + # Ray logging context + TIMESTAMP_NS = "timestamp_ns" diff --git a/lib/python3.12/site-packages/ray/_private/ray_logging/default_impl.py b/lib/python3.12/site-packages/ray/_private/ray_logging/default_impl.py new file mode 100644 index 0000000000000000000000000000000000000000..3901baaa956f3fd35cdfd8a3fa89ee5878137893 --- /dev/null +++ b/lib/python3.12/site-packages/ray/_private/ray_logging/default_impl.py @@ -0,0 +1,4 @@ +def get_logging_configurator(): + from ray._private.ray_logging.logging_config import DefaultLoggingConfigurator + + return DefaultLoggingConfigurator() diff --git a/lib/python3.12/site-packages/ray/_private/ray_logging/logging_config.py b/lib/python3.12/site-packages/ray/_private/ray_logging/logging_config.py new file mode 100644 index 0000000000000000000000000000000000000000..67453163a02798669af2f3d47bd5257c9ffdd1dc --- /dev/null +++ b/lib/python3.12/site-packages/ray/_private/ray_logging/logging_config.py @@ -0,0 +1,161 @@ +import logging +from abc import ABC, abstractmethod +from dataclasses import dataclass, field +from typing import Set + +from ray._common.filters import CoreContextFilter +from ray._common.formatters import JSONFormatter, TextFormatter +from ray._private.ray_logging import default_impl +from ray._private.ray_logging.constants import LOGRECORD_STANDARD_ATTRS +from ray.util.annotations import PublicAPI + + +class LoggingConfigurator(ABC): + @abstractmethod + def get_supported_encodings(self) -> Set[str]: + raise NotImplementedError + + @abstractmethod + def configure(self, logging_config: "LoggingConfig"): + raise NotImplementedError + + +class DefaultLoggingConfigurator(LoggingConfigurator): + def __init__(self): + self._encoding_to_formatter = { + "TEXT": TextFormatter(), + "JSON": JSONFormatter(), + } + + def get_supported_encodings(self) -> Set[str]: + return self._encoding_to_formatter.keys() + + def configure(self, logging_config: "LoggingConfig"): + formatter = self._encoding_to_formatter[logging_config.encoding] + formatter.set_additional_log_standard_attrs( + logging_config.additional_log_standard_attrs + ) + + core_context_filter = CoreContextFilter() + handler = logging.StreamHandler() + handler.setLevel(logging_config.log_level) + handler.setFormatter(formatter) + handler.addFilter(core_context_filter) + + root_logger = logging.getLogger() + root_logger.setLevel(logging_config.log_level) + root_logger.addHandler(handler) + + ray_logger = logging.getLogger("ray") + ray_logger.setLevel(logging_config.log_level) + # Remove all existing handlers added by `ray/__init__.py`. + for h in ray_logger.handlers[:]: + ray_logger.removeHandler(h) + ray_logger.addHandler(handler) + ray_logger.propagate = False + + +_logging_configurator: LoggingConfigurator = default_impl.get_logging_configurator() + + +# Class defines the logging configurations for a Ray job. +# To add a new logging configuration: (1) add a new field to this class; (2) Update the +# logic in the __post_init__ method in this class to add the validation logic; +# (3) Update the configure method in the DefaultLoggingConfigurator +# class to use the new field. +@PublicAPI(stability="alpha") +@dataclass +class LoggingConfig: + encoding: str = "TEXT" + log_level: str = "INFO" + # The list of valid attributes are defined as LOGRECORD_STANDARD_ATTRS in + # constants.py. + additional_log_standard_attrs: list = field(default_factory=list) + + def __post_init__(self): + if self.encoding not in _logging_configurator.get_supported_encodings(): + raise ValueError( + f"Invalid encoding type: {self.encoding}. " + "Valid encoding types are: " + f"{list(_logging_configurator.get_supported_encodings())}" + ) + + for attr in self.additional_log_standard_attrs: + if attr not in LOGRECORD_STANDARD_ATTRS: + raise ValueError( + f"Unknown python logging standard attribute: {attr}. " + "The valid attributes are: " + f"{LOGRECORD_STANDARD_ATTRS}" + ) + + def _configure_logging(self): + """Set up the logging configuration for the current process.""" + _logging_configurator.configure(self) + + def _apply(self): + """Set up the logging configuration.""" + self._configure_logging() + + +LoggingConfig.__doc__ = """ + Logging configuration for a Ray job. These configurations are used to set up the + root logger of the driver process and all Ray tasks and actor processes that belong + to the job. + + Examples: 1. Configure the logging to use TEXT encoding. + .. testcode:: + + import ray + import logging + + ray.init( + logging_config=ray.LoggingConfig(encoding="TEXT", log_level="INFO", additional_log_standard_attrs=['name']) + ) + + @ray.remote + def f(): + logger = logging.getLogger(__name__) + logger.info("This is a Ray task") + + ray.get(f.remote()) + ray.shutdown() + + .. testoutput:: + :options: +MOCK + + 2025-02-12 12:25:16,836 INFO test-log-config.py:11 -- This is a Ray task name=__main__ job_id=01000000 worker_id=51188d9448be4664bf2ea26ac410b67acaaa970c4f31c5ad3ae776a5 node_id=f683dfbffe2c69984859bc19c26b77eaf3866c458884c49d115fdcd4 task_id=c8ef45ccd0112571ffffffffffffffffffffffff01000000 task_name=f task_func_name=test-log-config.f timestamp_ns=1739391916836884000 + + 2. Configure the logging to use JSON encoding. + .. testcode:: + + import ray + import logging + + ray.init( + logging_config=ray.LoggingConfig(encoding="JSON", log_level="INFO", additional_log_standard_attrs=['name']) + ) + + @ray.remote + def f(): + logger = logging.getLogger(__name__) + logger.info("This is a Ray task") + + ray.get(f.remote()) + ray.shutdown() + + .. testoutput:: + :options: +MOCK + + {"asctime": "2025-02-12 12:25:48,766", "levelname": "INFO", "message": "This is a Ray task", "filename": "test-log-config.py", "lineno": 11, "name": "__main__", "job_id": "01000000", "worker_id": "6d307578014873fcdada0fa22ea6d49e0fb1f78960e69d61dfe41f5a", "node_id": "69e3a5e68bdc7eb8ac9abb3155326ee3cc9fc63ea1be04d11c0d93c7", "task_id": "c8ef45ccd0112571ffffffffffffffffffffffff01000000", "task_name": "f", "task_func_name": "test-log-config.f", "timestamp_ns": 1739391948766949000} + + Args: + encoding: Encoding type for the logs. The valid values are + {list(_logging_configurator.get_supported_encodings())} + log_level: Log level for the logs. Defaults to 'INFO'. You can set + it to 'DEBUG' to receive more detailed debug logs. + additional_log_standard_attrs: List of additional standard python logger attributes to + include in the log. Defaults to an empty list. The list of already + included standard attributes are: "asctime", "levelname", "message", + "filename", "lineno", "exc_text". The list of valid attributes are specified + here: http://docs.python.org/library/logging.html#logrecord-attributes + """ # noqa: E501 diff --git a/lib/python3.12/site-packages/ray/_private/ray_microbenchmark_helpers.py b/lib/python3.12/site-packages/ray/_private/ray_microbenchmark_helpers.py new file mode 100644 index 0000000000000000000000000000000000000000..f2986ab792bdcc1e2f5f16bda3e88d1e0e087c87 --- /dev/null +++ b/lib/python3.12/site-packages/ray/_private/ray_microbenchmark_helpers.py @@ -0,0 +1,92 @@ +import os +import time +from contextlib import contextmanager +from typing import List, Optional, Tuple + +import numpy as np + +import ray + +# Only run tests matching this filter pattern. + +filter_pattern = os.environ.get("TESTS_TO_RUN", "") +skip_pattern = os.environ.get("TESTS_TO_SKIP", "") + + +def timeit( + name, fn, multiplier=1, warmup_time_sec=10 +) -> List[Optional[Tuple[str, float, float]]]: + if filter_pattern and filter_pattern not in name: + return [None] + if skip_pattern and skip_pattern in name: + return [None] + # sleep for a while to avoid noisy neigbhors. + # related issue: https://github.com/ray-project/ray/issues/22045 + time.sleep(warmup_time_sec) + # warmup + start = time.perf_counter() + count = 0 + while time.perf_counter() - start < 1: + fn() + count += 1 + # real run + step = count // 10 + 1 + stats = [] + for _ in range(4): + start = time.perf_counter() + count = 0 + while time.perf_counter() - start < 2: + for _ in range(step): + fn() + count += step + end = time.perf_counter() + stats.append(multiplier * count / (end - start)) + + mean = np.mean(stats) + sd = np.std(stats) + print(name, "per second", round(mean, 2), "+-", round(sd, 2)) + return [(name, mean, sd)] + + +async def asyncio_timeit( + name, async_fn, multiplier=1, warmup_time_sec=10 +) -> List[Optional[Tuple[str, float, float]]]: + if filter_pattern and filter_pattern not in name: + return [None] + if skip_pattern and skip_pattern in name: + return [None] + # sleep for a while to avoid noisy neigbhors. + # related issue: https://github.com/ray-project/ray/issues/22045 + time.sleep(warmup_time_sec) + # warmup + start = time.perf_counter() + count = 0 + while time.perf_counter() - start < 1: + await async_fn() + count += 1 + # real run + step = count // 10 + 1 + stats = [] + for _ in range(4): + start = time.perf_counter() + count = 0 + while time.perf_counter() - start < 2: + for _ in range(step): + await async_fn() + count += step + end = time.perf_counter() + stats.append(multiplier * count / (end - start)) + + mean = np.mean(stats) + sd = np.std(stats) + print(name, "per second", round(mean, 2), "+-", round(sd, 2)) + return [(name, mean, sd)] + + +@contextmanager +def ray_setup_and_teardown(**init_args): + ray.init(**init_args) + try: + yield None + finally: + ray.shutdown() diff --git a/lib/python3.12/site-packages/ray/_private/ray_process_reaper.py b/lib/python3.12/site-packages/ray/_private/ray_process_reaper.py new file mode 100644 index 0000000000000000000000000000000000000000..5e15a5e35a9c40c719835f5bf15b8ca4da753a75 --- /dev/null +++ b/lib/python3.12/site-packages/ray/_private/ray_process_reaper.py @@ -0,0 +1,60 @@ +import atexit +import os +import signal +import sys +import time + +""" +This is a lightweight "reaper" process used to ensure that ray processes are +cleaned up properly when the main ray process dies unexpectedly (e.g., +segfaults or gets SIGKILLed). Note that processes may not be cleaned up +properly if this process is SIGTERMed or SIGKILLed. + +It detects that its parent has died by reading from stdin, which must be +inherited from the parent process so that the OS will deliver an EOF if the +parent dies. When this happens, the reaper process kills the rest of its +process group (first attempting graceful shutdown with SIGTERM, then escalating +to SIGKILL). +""" + +SIGTERM_GRACE_PERIOD_SECONDS = 1 + + +def reap_process_group(*args): + def sigterm_handler(*args): + # Give a one-second grace period for other processes to clean up. + time.sleep(SIGTERM_GRACE_PERIOD_SECONDS) + # SIGKILL the pgroup (including ourselves) as a last-resort. + if sys.platform == "win32": + atexit.unregister(sigterm_handler) + os.kill(0, signal.CTRL_BREAK_EVENT) + else: + os.killpg(0, signal.SIGKILL) + + # Set a SIGTERM handler to handle SIGTERMing ourselves with the group. + if sys.platform == "win32": + atexit.register(sigterm_handler) + else: + signal.signal(signal.SIGTERM, sigterm_handler) + + # Our parent must have died, SIGTERM the group (including ourselves). + if sys.platform == "win32": + os.kill(0, signal.CTRL_C_EVENT) + else: + os.killpg(0, signal.SIGTERM) + + +def main(): + # Read from stdin forever. Because stdin is a file descriptor + # inherited from our parent process, we will get an EOF if the parent + # dies, which is signaled by an empty return from read(). + # We intentionally don't set any signal handlers here, so a SIGTERM from + # the parent can be used to kill this process gracefully without it killing + # the rest of the process group. + while len(sys.stdin.read()) != 0: + pass + reap_process_group() + + +if __name__ == "__main__": + main() diff --git a/lib/python3.12/site-packages/ray/_private/resource_and_label_spec.py b/lib/python3.12/site-packages/ray/_private/resource_and_label_spec.py new file mode 100644 index 0000000000000000000000000000000000000000..d737b70961da9e90f69e71d4d53d17cc81e3afca --- /dev/null +++ b/lib/python3.12/site-packages/ray/_private/resource_and_label_spec.py @@ -0,0 +1,484 @@ +import json +import logging +import os +import sys +from typing import Dict, Optional, Tuple + +import ray +import ray._private.ray_constants as ray_constants +from ray._common.constants import HEAD_NODE_RESOURCE_NAME, NODE_ID_PREFIX +from ray._common.utils import RESOURCE_CONSTRAINT_PREFIX +from ray._private import accelerators +from ray._private.accelerators import AcceleratorManager + +logger = logging.getLogger(__name__) + + +class ResourceAndLabelSpec: + """Represents the resource and label configuration passed to a raylet. + + All fields can be None. Before starting services, resolve() should be + called to return a ResourceAndLabelSpec with unknown values filled in with + merged values based on the local machine and user specifications. + """ + + def __init__( + self, + num_cpus: Optional[int] = None, + num_gpus: Optional[int] = None, + memory: Optional[float] = None, + object_store_memory: Optional[float] = None, + resources: Optional[Dict[str, float]] = None, + labels: Optional[Dict[str, str]] = None, + ): + """ + Initialize a ResourceAndLabelSpec + + Args: + num_cpus: The CPUs allocated for this raylet. + num_gpus: The GPUs allocated for this raylet. + memory: The memory allocated for this raylet. + object_store_memory: The object store memory allocated for this raylet. + resources: The custom resources allocated for this raylet. + labels: The labels associated with this node. Labels can be used along + with resources for scheduling. + """ + self.num_cpus = num_cpus + self.num_gpus = num_gpus + self.memory = memory + self.object_store_memory = object_store_memory + self.resources = resources + self.labels = labels + self._is_resolved = False + + def resolved(self) -> bool: + """Returns if resolve() has been called for this ResourceAndLabelSpec + and default values are filled out.""" + return self._is_resolved + + def _all_fields_set(self) -> bool: + """Returns whether all fields in this ResourceAndLabelSpec are not None.""" + return all( + v is not None + for v in ( + self.num_cpus, + self.num_gpus, + self.memory, + self.object_store_memory, + self.resources, + self.labels, + ) + ) + + def to_resource_dict(self): + """Returns a dict suitable to pass to raylet initialization. + + This renames num_cpus / num_gpus to "CPU" / "GPU", + and check types and values. + """ + assert self.resolved() + + resources = dict( + self.resources, + CPU=self.num_cpus, + GPU=self.num_gpus, + memory=int(self.memory), + object_store_memory=int(self.object_store_memory), + ) + + resources = { + resource_label: resource_quantity + for resource_label, resource_quantity in resources.items() + if resource_quantity != 0 + } + + # Check types. + for resource_label, resource_quantity in resources.items(): + assert isinstance(resource_quantity, int) or isinstance( + resource_quantity, float + ), ( + f"{resource_label} ({type(resource_quantity)}): " f"{resource_quantity}" + ) + if ( + isinstance(resource_quantity, float) + and not resource_quantity.is_integer() + ): + raise ValueError( + "Resource quantities must all be whole numbers. " + "Violated by resource '{}' in {}.".format(resource_label, resources) + ) + if resource_quantity < 0: + raise ValueError( + "Resource quantities must be nonnegative. " + "Violated by resource '{}' in {}.".format(resource_label, resources) + ) + if resource_quantity > ray_constants.MAX_RESOURCE_QUANTITY: + raise ValueError( + "Resource quantities must be at most {}. " + "Violated by resource '{}' in {}.".format( + ray_constants.MAX_RESOURCE_QUANTITY, resource_label, resources + ) + ) + + return resources + + def resolve( + self, is_head: bool, node_ip_address: Optional[str] = None + ) -> "ResourceAndLabelSpec": + """Fills out this ResourceAndLabelSpec instance with merged values from system defaults and user specification. + + Args: + is_head: Whether this is the head node. + node_ip_address: The IP address of the node that we are on. + This is used to automatically create a node id resource. + + Returns: + ResourceAndLabelSpec: This instance with all fields resolved. + """ + + self._resolve_resources(is_head=is_head, node_ip_address=node_ip_address) + + # Resolve accelerator-specific resources + ( + accelerator_manager, + num_accelerators, + ) = ResourceAndLabelSpec._get_current_node_accelerator( + self.num_gpus, self.resources + ) + self._resolve_accelerator_resources(accelerator_manager, num_accelerators) + + # Default num_gpus value if unset by user and unable to auto-detect. + if self.num_gpus is None: + self.num_gpus = 0 + + # Resolve and merge node labels from all sources (params, env, and default). + self._resolve_labels(accelerator_manager) + + # Resolve memory resources + self._resolve_memory_resources() + + self._is_resolved = True + assert self._all_fields_set() + return self + + @staticmethod + def _load_env_resources() -> Dict[str, float]: + """Load resource overrides from the environment, if present.""" + env_resources = {} + env_string = os.getenv(ray_constants.RESOURCES_ENVIRONMENT_VARIABLE) + if env_string: + try: + env_resources = json.loads(env_string) + except Exception: + logger.exception(f"Failed to load {env_string}") + raise + logger.debug(f"Autoscaler overriding resources: {env_resources}.") + return env_resources + + @staticmethod + def _merge_resources(env_dict: Dict[str, float], params_dict: Dict[str, float]): + """Merge environment and Ray param-provided resources, with env values taking precedence. + Returns separated special case params (CPU/GPU/memory) and the merged resource dict. + """ + num_cpus = env_dict.pop("CPU", None) + num_gpus = env_dict.pop("GPU", None) + memory = env_dict.pop("memory", None) + object_store_memory = env_dict.pop("object_store_memory", None) + + result = params_dict.copy() + result.update(env_dict) + + for key in set(env_dict.keys()).intersection(params_dict or {}): + if params_dict[key] != env_dict[key]: + logger.warning( + f"Autoscaler is overriding your resource: {key}: " + f"{params_dict[key]} with {env_dict[key]}." + ) + + return num_cpus, num_gpus, memory, object_store_memory, result + + def _resolve_resources( + self, is_head: bool, node_ip_address: Optional[str] = None + ) -> None: + """Resolve CPU, GPU, and custom resources. Merges resources from environment, + Ray params, and defaults in that order of precedence.""" + + # Load environment override resources and merge with resources passed + # in from Ray Params. Separates special case params if found in env. + env_resources = ResourceAndLabelSpec._load_env_resources() + ( + num_cpus, + num_gpus, + memory, + object_store_memory, + merged_resources, + ) = ResourceAndLabelSpec._merge_resources(env_resources, self.resources or {}) + + self.num_cpus = self.num_cpus if num_cpus is None else num_cpus + self.num_gpus = self.num_gpus if num_gpus is None else num_gpus + self.memory = self.memory if memory is None else memory + self.object_store_memory = ( + self.object_store_memory + if object_store_memory is None + else object_store_memory + ) + self.resources = merged_resources + + if node_ip_address is None: + node_ip_address = ray.util.get_node_ip_address() + + # Automatically create a node id resource on each node. This is + # queryable with ray._private.state.node_ids() and + # ray._private.state.current_node_id(). + self.resources[NODE_ID_PREFIX + node_ip_address] = 1.0 + + # Automatically create a head node resource. + if HEAD_NODE_RESOURCE_NAME in self.resources: + raise ValueError( + f"{HEAD_NODE_RESOURCE_NAME}" + " is a reserved resource name, use another name instead." + ) + if is_head: + self.resources[HEAD_NODE_RESOURCE_NAME] = 1.0 + + # Auto-detect CPU count if not explicitly set + if self.num_cpus is None: + self.num_cpus = ray._private.utils.get_num_cpus() + + @staticmethod + def _load_env_labels() -> Dict[str, str]: + env_override_labels = {} + env_override_labels_string = os.getenv( + ray_constants.LABELS_ENVIRONMENT_VARIABLE + ) + if env_override_labels_string: + try: + env_override_labels = json.loads(env_override_labels_string) + except Exception: + logger.exception(f"Failed to load {env_override_labels_string}") + raise + logger.info(f"Autoscaler overriding labels: {env_override_labels}.") + + return env_override_labels + + @staticmethod + def _get_default_labels( + accelerator_manager: Optional[AcceleratorManager], + ) -> Dict[str, str]: + default_labels = {} + + # Get environment variables populated from K8s Pod Spec + node_group = os.environ.get(ray._raylet.NODE_TYPE_NAME_ENV, "") + market_type = os.environ.get(ray._raylet.NODE_MARKET_TYPE_ENV, "") + availability_region = os.environ.get(ray._raylet.NODE_REGION_ENV, "") + availability_zone = os.environ.get(ray._raylet.NODE_ZONE_ENV, "") + + # Map environment variables to default ray node labels + if market_type: + default_labels[ray._raylet.RAY_NODE_MARKET_TYPE_KEY] = market_type + if node_group: + default_labels[ray._raylet.RAY_NODE_GROUP_KEY] = node_group + if availability_zone: + default_labels[ray._raylet.RAY_NODE_ZONE_KEY] = availability_zone + if availability_region: + default_labels[ray._raylet.RAY_NODE_REGION_KEY] = availability_region + + # Get accelerator type from AcceleratorManager + if accelerator_manager: + accelerator_type = accelerator_manager.get_current_node_accelerator_type() + if accelerator_type: + default_labels[ + ray._raylet.RAY_NODE_ACCELERATOR_TYPE_KEY + ] = accelerator_type + + # Set TPU specific default labels to enable multi-host scheduling. + if accelerator_manager.get_resource_name() == "TPU": + tpu_labels = accelerator_manager.get_current_node_accelerator_labels() + if tpu_labels: + default_labels.update(tpu_labels) + + return default_labels + + def _resolve_labels( + self, accelerator_manager: Optional[AcceleratorManager] + ) -> None: + """Resolve and merge environment override, user-input from params, and Ray default + labels in that order of precedence.""" + + # Start with a dictionary filled out with Ray default labels + merged = ResourceAndLabelSpec._get_default_labels(accelerator_manager) + + # Merge user-specified labels from Ray params + for key, val in (self.labels or {}).items(): + if key in merged and merged[key] != val: + logger.warning( + f"User label is overriding Ray default label: {key}: " + f"{key}: {merged[key]} to " + f"{key}: {self.labels[key]}." + ) + merged[key] = val + + # Merge autoscaler override labels from environment + env_labels = ResourceAndLabelSpec._load_env_labels() + for key, val in (env_labels or {}).items(): + if key in merged and merged[key] != val: + logger.warning( + "Autoscaler is overriding your label:" + f"{key}: {merged[key]} to " + f"{key}: {env_labels[key]}." + ) + merged[key] = val + + self.labels = merged + + def _resolve_accelerator_resources(self, accelerator_manager, num_accelerators): + """Detect and update accelerator resources on a node.""" + if not accelerator_manager: + return + + accelerator_resource_name = accelerator_manager.get_resource_name() + visible_accelerator_ids = ( + accelerator_manager.get_current_process_visible_accelerator_ids() + ) + + # Check that the number of accelerators that the raylet wants doesn't + # exceed the amount allowed by visible accelerator ids. + if ( + num_accelerators is not None + and visible_accelerator_ids is not None + and num_accelerators > len(visible_accelerator_ids) + ): + raise ValueError( + f"Attempting to start raylet with {num_accelerators} " + f"{accelerator_resource_name}, " + f"but {accelerator_manager.get_visible_accelerator_ids_env_var()} " + f"contains {visible_accelerator_ids}." + ) + + if accelerator_resource_name == "GPU": + self.num_gpus = num_accelerators + else: + self.resources[accelerator_resource_name] = num_accelerators + + accelerator_type = accelerator_manager.get_current_node_accelerator_type() + if accelerator_type: + self.resources[f"{RESOURCE_CONSTRAINT_PREFIX}{accelerator_type}"] = 1 + additional_resources = ( + accelerator_manager.get_current_node_additional_resources() + ) + if additional_resources: + self.resources.update(additional_resources) + + def _resolve_memory_resources(self): + # Choose a default object store size. + system_memory = ray._common.utils.get_system_memory() + avail_memory = ray._private.utils.estimate_available_memory() + object_store_memory = self.object_store_memory + if object_store_memory is None: + object_store_memory = int( + avail_memory * ray_constants.DEFAULT_OBJECT_STORE_MEMORY_PROPORTION + ) + + # Set the object_store_memory size to 2GB on Mac + # to avoid degraded performance. + # (https://github.com/ray-project/ray/issues/20388) + if sys.platform == "darwin": + object_store_memory = min( + object_store_memory, ray_constants.MAC_DEGRADED_PERF_MMAP_SIZE_LIMIT + ) + + object_store_memory_cap = ( + ray_constants.DEFAULT_OBJECT_STORE_MAX_MEMORY_BYTES + ) + + # Cap by shm size by default to avoid low performance, but don't + # go lower than REQUIRE_SHM_SIZE_THRESHOLD. + if sys.platform == "linux" or sys.platform == "linux2": + # Multiple by 0.95 to give a bit of wiggle-room. + # https://github.com/ray-project/ray/pull/23034/files + shm_avail = ray._private.utils.get_shared_memory_bytes() * 0.95 + shm_cap = max(ray_constants.REQUIRE_SHM_SIZE_THRESHOLD, shm_avail) + + object_store_memory_cap = min(object_store_memory_cap, shm_cap) + + # Cap memory to avoid memory waste and perf issues on large nodes + if ( + object_store_memory_cap + and object_store_memory > object_store_memory_cap + ): + logger.debug( + "Warning: Capping object memory store to {}GB. ".format( + object_store_memory_cap // 1e9 + ) + + "To increase this further, specify `object_store_memory` " + "when calling ray.init() or ray start." + ) + object_store_memory = object_store_memory_cap + + memory = self.memory + if memory is None: + memory = avail_memory - object_store_memory + if memory < 100e6 and memory < 0.05 * system_memory: + raise ValueError( + "After taking into account object store and redis memory " + "usage, the amount of memory on this node available for " + "tasks and actors ({} GB) is less than {}% of total. " + "You can adjust these settings with " + "ray.init(memory=, " + "object_store_memory=).".format( + round(memory / 1e9, 2), int(100 * (memory / system_memory)) + ) + ) + + # Set the resolved memory and object_store_memory + self.object_store_memory = object_store_memory + self.memory = memory + + @staticmethod + def _get_current_node_accelerator( + num_gpus: Optional[int], resources: Dict[str, float] + ) -> Tuple[AcceleratorManager, int]: + """ + Returns the AcceleratorManager and accelerator count for the accelerator + associated with this node. This assumes each node has at most one accelerator type. + If no accelerators are present, returns None. + + The resolved accelerator count uses num_gpus (for GPUs) or resources if set, and + otherwise falls back to the count auto-detected by the AcceleratorManager. The + resolved accelerator count is capped by the number of visible accelerators. + + Args: + num_gpus: GPU count (if provided by user). + resources: Resource dictionary containing custom resource keys. + + Returns: + Tuple[Optional[AcceleratorManager], int]: A tuple containing the accelerator + manager (or None) the final resolved accelerator count. + """ + for resource_name in accelerators.get_all_accelerator_resource_names(): + accelerator_manager = accelerators.get_accelerator_manager_for_resource( + resource_name + ) + if accelerator_manager is None: + continue + # Respect configured value for GPUs if set + if resource_name == "GPU": + num_accelerators = num_gpus + else: + num_accelerators = resources.get(resource_name) + if num_accelerators is None: + num_accelerators = ( + accelerator_manager.get_current_node_num_accelerators() + ) + visible_accelerator_ids = ( + accelerator_manager.get_current_process_visible_accelerator_ids() + ) + if visible_accelerator_ids is not None: + num_accelerators = min( + num_accelerators, len(visible_accelerator_ids) + ) + + if num_accelerators > 0: + return accelerator_manager, num_accelerators + + return None, 0 diff --git a/lib/python3.12/site-packages/ray/_private/resource_isolation_config.py b/lib/python3.12/site-packages/ray/_private/resource_isolation_config.py new file mode 100644 index 0000000000000000000000000000000000000000..24837bf23796e447c3625f230aeff3f7f09a3a53 --- /dev/null +++ b/lib/python3.12/site-packages/ray/_private/resource_isolation_config.py @@ -0,0 +1,304 @@ +import logging +from typing import Optional + +import ray._common.utils +import ray._private.ray_constants as ray_constants +import ray._private.utils as utils + +logger = logging.getLogger(__name__) + +# See https://docs.kernel.org/admin-guide/cgroup-v2.html#weights +# for information about cpu weights +_CGROUP_CPU_MAX_WEIGHT: int = 10000 + + +class ResourceIsolationConfig: + """Configuration for enabling resource isolation by reserving memory and cpu for ray system processes through cgroupv2. + + Validates configuration for resource isolation by enforcing types, correct combinations of values, applying default values, + and sanity checking cpu and memory reservations. Also, converts system_reserved_cpu into cpu.weights for cgroupv2. + + Raises: + ValueError: On invalid inputs. + + Attributes: + enable_resource_isolation: True if cgroupv2 based isolation of ray + system processes is enabled. + cgroup_path: The path for the cgroup the raylet should use to enforce + resource isolation. + system_reserved_cpu: The amount of cores reserved for ray system + processes. Must be >= ray_constants.MINIMUM_SYSTEM_RESERVED_CPU_CORES + and < the total number of cores available. + system_reserved_memory: The amount of memory in bytes reserved + for ray system processes. Must be >= ray_constants.MINIMUM_SYSTEM_RESERVED_MEMORY_BYTES + and system_reserved_cpu + object_store_bytes < the total memory available. + + TODO(54703): Link documentation when it's available. + """ + + def __init__( + self, + enable_resource_isolation: bool = False, + cgroup_path: Optional[str] = None, + system_reserved_cpu: Optional[float] = None, + system_reserved_memory: Optional[int] = None, + ): + self._resource_isolation_enabled = enable_resource_isolation + self.cgroup_path = cgroup_path + self.system_reserved_memory = system_reserved_memory + self.system_pids = "" + + # cgroupv2 cpu.weight calculated from system_reserved_cpu assumes ray uses all available cores. + self.system_reserved_cpu_weight: int = None + + # TODO(irabbani): this is used to ensure that object_store_memory is not added twice + # to self._system_reserved_memory. This should be refactored in the future so that ResourceIsolationConfig + # can take object_store_memory as a constructor parameter and be constructed fully by the constructor. + self._constructed = False + + if not enable_resource_isolation: + if self.cgroup_path: + raise ValueError( + "cgroup_path cannot be set when resource isolation is not enabled. " + "Set enable_resource_isolation to True if you're using ray.init or use the " + "--enable-resource-isolation flag if you're using the ray cli." + ) + if system_reserved_cpu: + raise ValueError( + "system_reserved_cpu cannot be set when resource isolation is not enabled. " + "Set enable_resource_isolation to True if you're using ray.init or use the " + "--enable-resource-isolation flag if you're using the ray cli." + ) + + if self.system_reserved_memory: + raise ValueError( + "system_reserved_memory cannot be set when resource isolation is not enabled. " + "Set enable_resource_isolation to True if you're using ray.init or use the " + "--enable-resource-isolation flag if you're using the ray cli." + ) + return + + self.system_reserved_cpu_weight = self._validate_and_get_system_reserved_cpu( + system_reserved_cpu + ) + + self.system_reserved_memory = self._validate_and_get_system_reserved_memory( + system_reserved_memory + ) + + self.cgroup_path = self._validate_and_get_cgroup_path(cgroup_path) + + def is_enabled(self) -> bool: + return self._resource_isolation_enabled + + def add_object_store_memory(self, object_store_memory_bytes: int): + """Adds object_store_memory to the memory reserved for system processes. + + Args: + object_store_memory_bytes: The amount processes. Must be >= ray_constants.MINIMUM_SYSTEM_RESERVED_CPU_CORES + and < the total number of cores available. + + Raises: + AssertionError: If called with resource isolation not enabled or called more than once for the same instance. + ValueError: If the input is not an integer or if the system_reserved_memory + object_store_memory is greater + than the total memory available on the system. + + """ + assert self.is_enabled(), ( + "Cannot add object_store_memory to system_reserved_memory when " + "enable_resource_isolation is False." + ) + assert not self._constructed, ( + "Cannot call add_object_store_memory more than once with an instance " + "ResourceIsolationConfig. This is a bug in the ray code. " + ) + self.system_reserved_memory += object_store_memory_bytes + available_system_memory = ray._common.utils.get_system_memory() + if self.system_reserved_memory > available_system_memory: + raise ValueError( + f"The total requested system_reserved_memory={self.system_reserved_memory}, calculated by " + "object_store_bytes + system_reserved_memory, is greater than the total memory " + f"available={available_system_memory}. Pick a smaller number of bytes for object_store_bytes " + "or system_reserved_memory." + ) + self._constructed = True + + def add_system_pids(self, system_pids: str): + """A comma-separated list of pids to move into the system cgroup.""" + self.system_pids = system_pids + + @staticmethod + def _validate_and_get_cgroup_path(cgroup_path: Optional[str]) -> str: + """Returns the ray_constants.DEFAULT_CGROUP_PATH if cgroup_path is not specified. + + Args: + cgroup_path: The path for the cgroup the raylet should use to enforce + resource isolation. + + Returns: + str: The validated cgroup path. + + Raises: + ValueError: If cgroup_path is not a string. + """ + if not cgroup_path: + cgroup_path = ray_constants.DEFAULT_CGROUP_PATH + + if not isinstance(cgroup_path, str): + raise ValueError( + f"Invalid value={cgroup_path} for cgroup_path. " + "Use a string to represent the path for the cgroup that the raylet should use " + "to enable resource isolation." + ) + + return cgroup_path + + @staticmethod + def _validate_and_get_system_reserved_cpu( + system_reserved_cpu: Optional[float], + ) -> int: + """If system_reserved_cpu is specified, validates it, otherwise returns the default value. + + Validation entails checking the type, ensuring that the value is in range, and converts it + into cpu.weights for cgroupv2. See https://docs.kernel.org/admin-guide/cgroup-v2.html#weights + for more information. + + If system_reserved_cpu is not specified, returns a default value between + [DEFAULT_MIN_SYSTEM_RESERVED_CPU_CORES, DEFAULT_MAX_SYSTEM_RESERVED_CPU_CORES]. + + # TODO(54703): The errors from this method are user-facing and thus need + to be linked the user-facing documentation once it's available. + + Args: + system_reserved_cpu: The amount of cores reserved for ray system + processes. Must be >= ray_constants.MINIMUM_SYSTEM_RESERVED_CPU_CORES + and < the total number of cores available. + + Raises: + ValueError: If system_reserved_cpu is specified, but invalid or if the system + does not have enough available cpus. + + """ + available_system_cpus = utils.get_num_cpus(truncate=False) + + if available_system_cpus < ray_constants.DEFAULT_MIN_SYSTEM_RESERVED_CPU_CORES: + raise ValueError( + f"The available number of cpu cores on this system {available_system_cpus} is less than " + f"the minimum amount that is required for ray's system processes. " + f"Pick a number of cpu cores greater than or equal to {ray_constants.DEFAULT_MIN_SYSTEM_RESERVED_CPU_CORES}" + ) + + if not system_reserved_cpu: + system_reserved_cpu = float( + min( + max( + ray_constants.DEFAULT_MIN_SYSTEM_RESERVED_CPU_CORES, + ray_constants.DEFAULT_SYSTEM_RESERVED_CPU_PROPORTION + * available_system_cpus, + ), + ray_constants.DEFAULT_MAX_SYSTEM_RESERVED_CPU_CORES, + ) + ) + + if not ( + isinstance(system_reserved_cpu, float) + or isinstance(system_reserved_cpu, int) + ): + raise ValueError( + f"Invalid value={system_reserved_cpu} for system_reserved_cpu. " + "Use a float to represent the number of cores that need to be reserved for " + "ray system processes to enable resource isolation." + ) + + system_reserved_cpu = float(system_reserved_cpu) + + if system_reserved_cpu < ray_constants.DEFAULT_MIN_SYSTEM_RESERVED_CPU_CORES: + raise ValueError( + f"The requested system_reserved_cpu={system_reserved_cpu} is less than " + f"the minimum number of cpus that can be used for resource isolation. " + "Pick a number of cpu cores to reserve for ray system processes " + f"greater than or equal to {ray_constants.DEFAULT_MIN_SYSTEM_RESERVED_CPU_CORES}" + ) + + if system_reserved_cpu >= available_system_cpus: + raise ValueError( + f"The requested system_reserved_cpu={system_reserved_cpu} is greater than or equal to " + f"the number of cpus available={available_system_cpus}. " + "Pick a smaller number of cpu cores to reserve for ray system processes." + ) + + # Converting the number of cores the user defined into cpu.weights + # This assumes that ray is allowed to use all available CPU + # cores and distribute them between system, worker and + # user processes + return int( + (system_reserved_cpu / float(available_system_cpus)) + * _CGROUP_CPU_MAX_WEIGHT + ) + + @staticmethod + def _validate_and_get_system_reserved_memory( + system_reserved_memory: Optional[int], + ) -> int: + """If system_reserved_memory is not specified, returns the default value. Otherwise, + checks the type, makes sure that the value is in range. + + Args: + system_reserved_memory: The amount of memory in bytes reserved + for ray system processes. Must be >= ray_constants.MINIMUM_SYSTEM_RESERVED_MEMORY_BYTES + and < the total memory available. + + Returns: + int: The validated system reserved memory in bytes. + + Raises: + ValueError: If system_reserved_memory is specified, but invalid. + """ + available_system_memory = ray._common.utils.get_system_memory() + + if ( + available_system_memory + < ray_constants.DEFAULT_MIN_SYSTEM_RESERVED_MEMORY_BYTES + ): + raise ValueError( + f"The available memory on this system {available_system_memory} is less than " + f"the minimum amount that is required for ray's system processes. " + f"Pick a number of bytes greater than or equal to {ray_constants.DEFAULT_MIN_SYSTEM_RESERVED_MEMORY_BYTES}" + ) + + if not system_reserved_memory: + system_reserved_memory = int( + min( + max( + ray_constants.DEFAULT_MIN_SYSTEM_RESERVED_MEMORY_BYTES, + ray_constants.DEFAULT_SYSTEM_RESERVED_MEMORY_PROPORTION + * available_system_memory, + ), + ray_constants.DEFAULT_MAX_SYSTEM_RESERVED_MEMORY_BYTES, + ) + ) + + if not isinstance(system_reserved_memory, int): + raise ValueError( + f"Invalid value {system_reserved_memory} for system_reserved_memory. " + "Use an integer to represent the number bytes that need to be reserved for " + "ray system processes to enable resource isolation." + ) + + if ( + system_reserved_memory + < ray_constants.DEFAULT_MIN_SYSTEM_RESERVED_MEMORY_BYTES + ): + raise ValueError( + f"The requested system_reserved_memory {system_reserved_memory} is less than " + f"the minimum number of bytes that can be used for resource isolation. " + "Pick a number of bytes to reserve for ray system processes " + f"greater than or equal to {ray_constants.DEFAULT_MIN_SYSTEM_RESERVED_MEMORY_BYTES}" + ) + + if system_reserved_memory > available_system_memory: + raise ValueError( + f"The total requested system_reserved_memory={system_reserved_memory} is greater than " + f"the amount of memory available={available_system_memory}." + ) + return system_reserved_memory diff --git a/lib/python3.12/site-packages/ray/_private/runtime_env/__init__.py b/lib/python3.12/site-packages/ray/_private/runtime_env/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..d8d40c33d67ca8374373a3258a0daa4445480efb --- /dev/null +++ b/lib/python3.12/site-packages/ray/_private/runtime_env/__init__.py @@ -0,0 +1,3 @@ +# List of files to exclude from the Ray directory when using runtime_env for +# Ray development. These are not necessary in the Ray workers. +RAY_WORKER_DEV_EXCLUDES = ["raylet", "gcs_server", "cpp/", "tests/", "core/src"] diff --git a/lib/python3.12/site-packages/ray/_private/runtime_env/_clonevirtualenv.py b/lib/python3.12/site-packages/ray/_private/runtime_env/_clonevirtualenv.py new file mode 100644 index 0000000000000000000000000000000000000000..1f2eab3d1040352a0a673323838004f2f4cc7f2c --- /dev/null +++ b/lib/python3.12/site-packages/ray/_private/runtime_env/_clonevirtualenv.py @@ -0,0 +1,334 @@ +#!/usr/bin/env python + +from __future__ import with_statement + +import logging +import optparse +import os +import os.path +import re +import shutil +import subprocess +import sys +import itertools + +__version__ = "0.5.7" + + +logger = logging.getLogger() + + +env_bin_dir = "bin" +if sys.platform == "win32": + env_bin_dir = "Scripts" + _WIN32 = True +else: + _WIN32 = False + + +class UserError(Exception): + pass + + +def _dirmatch(path, matchwith): + """Check if path is within matchwith's tree. + >>> _dirmatch('/home/foo/bar', '/home/foo/bar') + True + >>> _dirmatch('/home/foo/bar/', '/home/foo/bar') + True + >>> _dirmatch('/home/foo/bar/etc', '/home/foo/bar') + True + >>> _dirmatch('/home/foo/bar2', '/home/foo/bar') + False + >>> _dirmatch('/home/foo/bar2/etc', '/home/foo/bar') + False + """ + matchlen = len(matchwith) + if path.startswith(matchwith) and path[matchlen : matchlen + 1] in [os.sep, ""]: + return True + return False + + +def _virtualenv_sys(venv_path): + """obtain version and path info from a virtualenv.""" + executable = os.path.join(venv_path, env_bin_dir, "python") + if _WIN32: + env = os.environ.copy() + else: + env = {} + # Must use "executable" as the first argument rather than as the + # keyword argument "executable" to get correct value from sys.path + p = subprocess.Popen( + [ + executable, + "-c", + "import sys;" + 'print ("%d.%d" % (sys.version_info.major, sys.version_info.minor));' + 'print ("\\n".join(sys.path));', + ], + env=env, + stdout=subprocess.PIPE, + ) + stdout, err = p.communicate() + assert not p.returncode and stdout + lines = stdout.decode("utf-8").splitlines() + return lines[0], list(filter(bool, lines[1:])) + + +def clone_virtualenv(src_dir, dst_dir): + if not os.path.exists(src_dir): + raise UserError("src dir %r does not exist" % src_dir) + if os.path.exists(dst_dir): + raise UserError("dest dir %r exists" % dst_dir) + # sys_path = _virtualenv_syspath(src_dir) + logger.info("cloning virtualenv '%s' => '%s'..." % (src_dir, dst_dir)) + shutil.copytree( + src_dir, dst_dir, symlinks=True, ignore=shutil.ignore_patterns("*.pyc") + ) + version, sys_path = _virtualenv_sys(dst_dir) + logger.info("fixing scripts in bin...") + fixup_scripts(src_dir, dst_dir, version) + + has_old = lambda s: any(i for i in s if _dirmatch(i, src_dir)) # noqa: E731 + + if has_old(sys_path): + # only need to fix stuff in sys.path if we have old + # paths in the sys.path of new python env. right? + logger.info("fixing paths in sys.path...") + fixup_syspath_items(sys_path, src_dir, dst_dir) + v_sys = _virtualenv_sys(dst_dir) + remaining = has_old(v_sys[1]) + assert not remaining, v_sys + fix_symlink_if_necessary(src_dir, dst_dir) + + +def fix_symlink_if_necessary(src_dir, dst_dir): + # sometimes the source virtual environment has symlinks that point to itself + # one example is $OLD_VIRTUAL_ENV/local/lib points to $OLD_VIRTUAL_ENV/lib + # this function makes sure + # $NEW_VIRTUAL_ENV/local/lib will point to $NEW_VIRTUAL_ENV/lib + # usually this goes unnoticed unless one tries to upgrade a package though pip, + # so this bug is hard to find. + logger.info("scanning for internal symlinks that point to the original virtual env") + for dirpath, dirnames, filenames in os.walk(dst_dir): + for a_file in itertools.chain(filenames, dirnames): + full_file_path = os.path.join(dirpath, a_file) + if os.path.islink(full_file_path): + target = os.path.realpath(full_file_path) + if target.startswith(src_dir): + new_target = target.replace(src_dir, dst_dir) + logger.debug("fixing symlink in %s" % (full_file_path,)) + os.remove(full_file_path) + os.symlink(new_target, full_file_path) + + +def fixup_scripts(old_dir, new_dir, version, rewrite_env_python=False): + bin_dir = os.path.join(new_dir, env_bin_dir) + root, dirs, files = next(os.walk(bin_dir)) + pybinre = re.compile(r"pythonw?([0-9]+(\.[0-9]+(\.[0-9]+)?)?)?$") + for file_ in files: + filename = os.path.join(root, file_) + if file_ in ["python", "python%s" % version, "activate_this.py"]: + continue + elif file_.startswith("python") and pybinre.match(file_): + # ignore other possible python binaries + continue + elif file_.endswith(".pyc"): + # ignore compiled files + continue + elif file_ == "activate" or file_.startswith("activate."): + fixup_activate(os.path.join(root, file_), old_dir, new_dir) + elif os.path.islink(filename): + fixup_link(filename, old_dir, new_dir) + elif os.path.isfile(filename): + fixup_script_( + root, + file_, + old_dir, + new_dir, + version, + rewrite_env_python=rewrite_env_python, + ) + + +def fixup_script_(root, file_, old_dir, new_dir, version, rewrite_env_python=False): + old_shebang = "#!%s/bin/python" % os.path.normcase(os.path.abspath(old_dir)) + new_shebang = "#!%s/bin/python" % os.path.normcase(os.path.abspath(new_dir)) + env_shebang = "#!/usr/bin/env python" + + filename = os.path.join(root, file_) + with open(filename, "rb") as f: + if f.read(2) != b"#!": + # no shebang + return + f.seek(0) + lines = f.readlines() + + if not lines: + # warn: empty script + return + + def rewrite_shebang(version=None): + logger.debug("fixing %s" % filename) + shebang = new_shebang + if version: + shebang = shebang + version + shebang = (shebang + "\n").encode("utf-8") + with open(filename, "wb") as f: + f.write(shebang) + f.writelines(lines[1:]) + + try: + bang = lines[0].decode("utf-8").strip() + except UnicodeDecodeError: + # binary file + return + + # This takes care of the scheme in which shebang is of type + # '#!/venv/bin/python3' while the version of system python + # is of type 3.x e.g. 3.5. + short_version = bang[len(old_shebang) :] + + if not bang.startswith("#!"): + return + elif bang == old_shebang: + rewrite_shebang() + elif bang.startswith(old_shebang) and bang[len(old_shebang) :] == version: + rewrite_shebang(version) + elif ( + bang.startswith(old_shebang) + and short_version + and bang[len(old_shebang) :] == short_version + ): + rewrite_shebang(short_version) + elif rewrite_env_python and bang.startswith(env_shebang): + if bang == env_shebang: + rewrite_shebang() + elif bang[len(env_shebang) :] == version: + rewrite_shebang(version) + else: + # can't do anything + return + + +def fixup_activate(filename, old_dir, new_dir): + logger.debug("fixing %s" % filename) + with open(filename, "rb") as f: + data = f.read().decode("utf-8") + + data = data.replace(old_dir, new_dir) + with open(filename, "wb") as f: + f.write(data.encode("utf-8")) + + +def fixup_link(filename, old_dir, new_dir, target=None): + logger.debug("fixing %s" % filename) + if target is None: + target = os.readlink(filename) + + origdir = os.path.dirname(os.path.abspath(filename)).replace(new_dir, old_dir) + if not os.path.isabs(target): + target = os.path.abspath(os.path.join(origdir, target)) + rellink = True + else: + rellink = False + + if _dirmatch(target, old_dir): + if rellink: + # keep relative links, but don't keep original in case it + # traversed up out of, then back into the venv. + # so, recreate a relative link from absolute. + target = target[len(origdir) :].lstrip(os.sep) + else: + target = target.replace(old_dir, new_dir, 1) + + # else: links outside the venv, replaced with absolute path to target. + _replace_symlink(filename, target) + + +def _replace_symlink(filename, newtarget): + tmpfn = "%s.new" % filename + os.symlink(newtarget, tmpfn) + os.rename(tmpfn, filename) + + +def fixup_syspath_items(syspath, old_dir, new_dir): + for path in syspath: + if not os.path.isdir(path): + continue + path = os.path.normcase(os.path.abspath(path)) + if _dirmatch(path, old_dir): + path = path.replace(old_dir, new_dir, 1) + if not os.path.exists(path): + continue + elif not _dirmatch(path, new_dir): + continue + root, dirs, files = next(os.walk(path)) + for file_ in files: + filename = os.path.join(root, file_) + if filename.endswith(".pth"): + fixup_pth_file(filename, old_dir, new_dir) + elif filename.endswith(".egg-link"): + fixup_egglink_file(filename, old_dir, new_dir) + + +def fixup_pth_file(filename, old_dir, new_dir): + logger.debug("fixup_pth_file %s" % filename) + + with open(filename, "r") as f: + lines = f.readlines() + + has_change = False + + for num, line in enumerate(lines): + line = (line.decode("utf-8") if hasattr(line, "decode") else line).strip() + + if not line or line.startswith("#") or line.startswith("import "): + continue + elif _dirmatch(line, old_dir): + lines[num] = line.replace(old_dir, new_dir, 1) + has_change = True + + if has_change: + with open(filename, "w") as f: + payload = os.linesep.join([line.strip() for line in lines]) + os.linesep + f.write(payload) + + +def fixup_egglink_file(filename, old_dir, new_dir): + logger.debug("fixing %s" % filename) + with open(filename, "rb") as f: + link = f.read().decode("utf-8").strip() + if _dirmatch(link, old_dir): + link = link.replace(old_dir, new_dir, 1) + with open(filename, "wb") as f: + link = (link + "\n").encode("utf-8") + f.write(link) + + +def main(): + parser = optparse.OptionParser( + "usage: %prog [options] /path/to/existing/venv /path/to/cloned/venv" + ) + parser.add_option( + "-v", action="count", dest="verbose", default=False, help="verbosity" + ) + options, args = parser.parse_args() + try: + old_dir, new_dir = args + except ValueError: + print("virtualenv-clone %s" % (__version__,)) + parser.error("not enough arguments given.") + old_dir = os.path.realpath(old_dir) + new_dir = os.path.realpath(new_dir) + loglevel = (logging.WARNING, logging.INFO, logging.DEBUG)[min(2, options.verbose)] + logging.basicConfig(level=loglevel, format="%(message)s") + try: + clone_virtualenv(old_dir, new_dir) + except UserError: + e = sys.exc_info()[1] + parser.error(str(e)) + + +if __name__ == "__main__": + main() diff --git a/lib/python3.12/site-packages/ray/_private/runtime_env/agent/__init__.py b/lib/python3.12/site-packages/ray/_private/runtime_env/agent/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/lib/python3.12/site-packages/ray/_private/runtime_env/agent/main.py b/lib/python3.12/site-packages/ray/_private/runtime_env/agent/main.py new file mode 100644 index 0000000000000000000000000000000000000000..c0d68bdf6e041df0d40a8e03dfde0cc10d7e71ca --- /dev/null +++ b/lib/python3.12/site-packages/ray/_private/runtime_env/agent/main.py @@ -0,0 +1,238 @@ +import argparse +import logging +import os +import sys + +import ray._private.ray_constants as ray_constants +from ray._common.utils import ( + get_or_create_event_loop, +) +from ray._private import logging_utils +from ray._private.authentication.http_token_authentication import ( + get_token_auth_middleware, +) +from ray._private.process_watcher import create_check_raylet_task +from ray._raylet import GcsClient +from ray.core.generated import ( + runtime_env_agent_pb2, +) + + +def import_libs(): + my_dir = os.path.abspath(os.path.dirname(__file__)) + sys.path.insert(0, os.path.join(my_dir, "thirdparty_files")) # for aiohttp + sys.path.insert(0, my_dir) # for runtime_env_agent and runtime_env_consts + + +import_libs() + +import aiohttp # noqa: E402 +import runtime_env_consts # noqa: E402 +from aiohttp import web # noqa: E402 +from runtime_env_agent import RuntimeEnvAgent # noqa: E402 + +if __name__ == "__main__": + parser = argparse.ArgumentParser(description="Runtime env agent.") + parser.add_argument( + "--node-ip-address", + required=True, + type=str, + help="the IP address of this node.", + ) + parser.add_argument( + "--runtime-env-agent-port", + required=True, + type=int, + default=None, + help="The port on which the runtime env agent will receive HTTP requests.", + ) + + parser.add_argument( + "--gcs-address", required=True, type=str, help="The address (ip:port) of GCS." + ) + parser.add_argument( + "--cluster-id-hex", required=True, type=str, help="The cluster id in hex." + ) + parser.add_argument( + "--runtime-env-dir", + required=True, + type=str, + default=None, + help="Specify the path of the resource directory used by runtime_env.", + ) + + parser.add_argument( + "--logging-level", + required=False, + type=lambda s: logging.getLevelName(s.upper()), + default=ray_constants.LOGGER_LEVEL, + choices=ray_constants.LOGGER_LEVEL_CHOICES, + help=ray_constants.LOGGER_LEVEL_HELP, + ) + parser.add_argument( + "--logging-format", + required=False, + type=str, + default=ray_constants.LOGGER_FORMAT, + help=ray_constants.LOGGER_FORMAT_HELP, + ) + parser.add_argument( + "--logging-filename", + required=False, + type=str, + default=runtime_env_consts.RUNTIME_ENV_AGENT_LOG_FILENAME, + help="Specify the name of log file, " + 'log to stdout if set empty, default is "{}".'.format( + runtime_env_consts.RUNTIME_ENV_AGENT_LOG_FILENAME + ), + ) + parser.add_argument( + "--logging-rotate-bytes", + required=True, + type=int, + help="Specify the max bytes for rotating log file", + ) + parser.add_argument( + "--logging-rotate-backup-count", + required=True, + type=int, + help="Specify the backup count of rotated log file", + ) + parser.add_argument( + "--log-dir", + required=True, + type=str, + default=None, + help="Specify the path of log directory.", + ) + parser.add_argument( + "--temp-dir", + required=True, + type=str, + default=None, + help="Specify the path of the temporary directory use by Ray process.", + ) + parser.add_argument( + "--stdout-filepath", + required=False, + type=str, + default="", + help="The filepath to dump runtime env agent stdout.", + ) + parser.add_argument( + "--stderr-filepath", + required=False, + type=str, + default="", + help="The filepath to dump runtime env agent stderr.", + ) + + args = parser.parse_args() + + # Disable log rotation for windows platform. + logging_rotation_bytes = args.logging_rotate_bytes if sys.platform != "win32" else 0 + logging_rotation_backup_count = ( + args.logging_rotate_backup_count if sys.platform != "win32" else 1 + ) + + logging_params = dict( + logging_level=args.logging_level, + logging_format=args.logging_format, + log_dir=args.log_dir, + filename=args.logging_filename, + max_bytes=logging_rotation_bytes, + backup_count=logging_rotation_backup_count, + ) + + # Setup stdout/stderr redirect files if redirection enabled. + logging_utils.redirect_stdout_stderr_if_needed( + args.stdout_filepath, + args.stderr_filepath, + logging_rotation_bytes, + logging_rotation_backup_count, + ) + + gcs_client = GcsClient(address=args.gcs_address, cluster_id=args.cluster_id_hex) + agent = RuntimeEnvAgent( + runtime_env_dir=args.runtime_env_dir, + logging_params=logging_params, + gcs_client=gcs_client, + temp_dir=args.temp_dir, + address=args.node_ip_address, + runtime_env_agent_port=args.runtime_env_agent_port, + ) + + # POST /get_or_create_runtime_env + # body is serialzied protobuf GetOrCreateRuntimeEnvRequest + # reply is serialzied protobuf GetOrCreateRuntimeEnvReply + async def get_or_create_runtime_env(request: web.Request) -> web.Response: + data = await request.read() + request = runtime_env_agent_pb2.GetOrCreateRuntimeEnvRequest() + request.ParseFromString(data) + reply = await agent.GetOrCreateRuntimeEnv(request) + return web.Response( + body=reply.SerializeToString(), content_type="application/octet-stream" + ) + + # POST /delete_runtime_env_if_possible + # body is serialzied protobuf DeleteRuntimeEnvIfPossibleRequest + # reply is serialzied protobuf DeleteRuntimeEnvIfPossibleReply + async def delete_runtime_env_if_possible(request: web.Request) -> web.Response: + data = await request.read() + request = runtime_env_agent_pb2.DeleteRuntimeEnvIfPossibleRequest() + request.ParseFromString(data) + reply = await agent.DeleteRuntimeEnvIfPossible(request) + return web.Response( + body=reply.SerializeToString(), content_type="application/octet-stream" + ) + + # POST /get_runtime_envs_info + # body is serialzied protobuf GetRuntimeEnvsInfoRequest + # reply is serialzied protobuf GetRuntimeEnvsInfoReply + async def get_runtime_envs_info(request: web.Request) -> web.Response: + data = await request.read() + request = runtime_env_agent_pb2.GetRuntimeEnvsInfoRequest() + request.ParseFromString(data) + reply = await agent.GetRuntimeEnvsInfo(request) + return web.Response( + body=reply.SerializeToString(), content_type="application/octet-stream" + ) + + app = web.Application(middlewares=[get_token_auth_middleware(aiohttp)]) + + app.router.add_post("/get_or_create_runtime_env", get_or_create_runtime_env) + app.router.add_post( + "/delete_runtime_env_if_possible", delete_runtime_env_if_possible + ) + app.router.add_post("/get_runtime_envs_info", get_runtime_envs_info) + + loop = get_or_create_event_loop() + check_raylet_task = None + if sys.platform not in ["win32", "cygwin"]: + + def parent_dead_callback(msg): + agent._logger.info( + "Raylet is dead! Exiting Runtime Env Agent. " + f"addr: {args.node_ip_address}, " + f"port: {args.runtime_env_agent_port}\n" + f"{msg}" + ) + + # No need to await this task. + check_raylet_task = create_check_raylet_task( + args.log_dir, gcs_client, parent_dead_callback, loop + ) + try: + web.run_app( + app, + host=args.node_ip_address, + port=args.runtime_env_agent_port, + loop=loop, + ) + except SystemExit as e: + agent._logger.info(f"SystemExit! {e}") + # We have to poke the task exception, or there's an error message + # "task exception was never retrieved". + if check_raylet_task is not None: + check_raylet_task.exception() + sys.exit(e.code) diff --git a/lib/python3.12/site-packages/ray/_private/runtime_env/agent/runtime_env_agent.py b/lib/python3.12/site-packages/ray/_private/runtime_env/agent/runtime_env_agent.py new file mode 100644 index 0000000000000000000000000000000000000000..c771d91b10051d4dbf6be7e412391e6e4886ca5f --- /dev/null +++ b/lib/python3.12/site-packages/ray/_private/runtime_env/agent/runtime_env_agent.py @@ -0,0 +1,608 @@ +import asyncio +import logging +import os +import time +import traceback +from collections import defaultdict +from dataclasses import dataclass +from typing import Callable, Dict, List, Set, Tuple + +import ray +import ray._private.runtime_env.agent.runtime_env_consts as runtime_env_consts +from ray._common.utils import get_or_create_event_loop +from ray._private.ray_constants import ( + DEFAULT_RUNTIME_ENV_TIMEOUT_SECONDS, +) +from ray._private.ray_logging import setup_component_logger +from ray._private.runtime_env.conda import CondaPlugin +from ray._private.runtime_env.context import RuntimeEnvContext +from ray._private.runtime_env.default_impl import get_image_uri_plugin_cls +from ray._private.runtime_env.image_uri import ContainerPlugin +from ray._private.runtime_env.java_jars import JavaJarsPlugin +from ray._private.runtime_env.nsight import NsightPlugin +from ray._private.runtime_env.pip import PipPlugin +from ray._private.runtime_env.plugin import ( + RuntimeEnvPlugin, + RuntimeEnvPluginManager, + create_for_plugin_if_needed, +) +from ray._private.runtime_env.py_executable import PyExecutablePlugin +from ray._private.runtime_env.py_modules import PyModulesPlugin +from ray._private.runtime_env.rocprof_sys import RocProfSysPlugin +from ray._private.runtime_env.uv import UvPlugin +from ray._private.runtime_env.working_dir import WorkingDirPlugin +from ray._raylet import GcsClient +from ray.core.generated import runtime_env_agent_pb2 +from ray.core.generated.runtime_env_common_pb2 import ( + RuntimeEnvState as ProtoRuntimeEnvState, +) +from ray.runtime_env import RuntimeEnv, RuntimeEnvConfig + +default_logger = logging.getLogger(__name__) + +# TODO(edoakes): this is used for unit tests. We should replace it with a +# better pluggability mechanism once available. +SLEEP_FOR_TESTING_S = os.environ.get("RAY_RUNTIME_ENV_SLEEP_FOR_TESTING_S") + + +@dataclass +class CreatedEnvResult: + # Whether or not the env was installed correctly. + success: bool + # If success is True, will be a serialized RuntimeEnvContext + # If success is False, will be an error message. + result: str + # The time to create a runtime env in ms. + creation_time_ms: int + + +# e.g., "working_dir" +UriType = str + + +class ReferenceTable: + """ + The URI reference table which is used for GC. + When the reference count is decreased to zero, + the URI should be removed from this table and + added to cache if needed. + """ + + def __init__( + self, + uris_parser: Callable[[RuntimeEnv], Tuple[str, UriType]], + unused_uris_callback: Callable[[List[Tuple[str, UriType]]], None], + unused_runtime_env_callback: Callable[[str], None], + ): + # Runtime Environment reference table. The key is serialized runtime env and + # the value is reference count. + self._runtime_env_reference: Dict[str, int] = defaultdict(int) + # URI reference table. The key is URI parsed from runtime env and the value + # is reference count. + self._uri_reference: Dict[str, int] = defaultdict(int) + self._uris_parser = uris_parser + self._unused_uris_callback = unused_uris_callback + self._unused_runtime_env_callback = unused_runtime_env_callback + # send the `DeleteRuntimeEnvIfPossible` RPC when the client exits. The URI won't + # be leaked now because the reference count will be reset to zero when the job + # finished. + self._reference_exclude_sources: Set[str] = { + "client_server", + } + + def _increase_reference_for_uris(self, uris): + default_logger.debug(f"Increase reference for uris {uris}.") + for uri, _ in uris: + self._uri_reference[uri] += 1 + + def _decrease_reference_for_uris(self, uris): + default_logger.debug(f"Decrease reference for uris {uris}.") + unused_uris = list() + for uri, uri_type in uris: + if self._uri_reference[uri] > 0: + self._uri_reference[uri] -= 1 + if self._uri_reference[uri] == 0: + unused_uris.append((uri, uri_type)) + del self._uri_reference[uri] + else: + default_logger.warning(f"URI {uri} does not exist.") + if unused_uris: + default_logger.info(f"Unused uris {unused_uris}.") + self._unused_uris_callback(unused_uris) + return unused_uris + + def _increase_reference_for_runtime_env(self, serialized_env: str): + default_logger.debug(f"Increase reference for runtime env {serialized_env}.") + self._runtime_env_reference[serialized_env] += 1 + + def _decrease_reference_for_runtime_env(self, serialized_env: str): + """Decrease reference count for the given [serialized_env]. Throw exception if we cannot decrement reference.""" + default_logger.debug(f"Decrease reference for runtime env {serialized_env}.") + unused = False + if self._runtime_env_reference[serialized_env] > 0: + self._runtime_env_reference[serialized_env] -= 1 + if self._runtime_env_reference[serialized_env] == 0: + unused = True + del self._runtime_env_reference[serialized_env] + else: + default_logger.warning(f"Runtime env {serialized_env} does not exist.") + raise ValueError( + f"{serialized_env} cannot decrement reference since the reference count is 0" + ) + if unused: + default_logger.info(f"Unused runtime env {serialized_env}.") + self._unused_runtime_env_callback(serialized_env) + + def increase_reference( + self, runtime_env: RuntimeEnv, serialized_env: str, source_process: str + ) -> None: + if source_process in self._reference_exclude_sources: + return + self._increase_reference_for_runtime_env(serialized_env) + uris = self._uris_parser(runtime_env) + self._increase_reference_for_uris(uris) + + def decrease_reference( + self, runtime_env: RuntimeEnv, serialized_env: str, source_process: str + ) -> None: + """Decrease reference count for runtime env and uri. Throw exception if decrement reference count fails.""" + if source_process in self._reference_exclude_sources: + return + self._decrease_reference_for_runtime_env(serialized_env) + uris = self._uris_parser(runtime_env) + self._decrease_reference_for_uris(uris) + + @property + def runtime_env_refs(self) -> Dict[str, int]: + """Return the runtime_env -> ref count mapping. + + Returns: + The mapping of serialized runtime env -> ref count. + """ + return self._runtime_env_reference + + +class RuntimeEnvAgent: + """An RPC server to create and delete runtime envs. + + Attributes: + dashboard_agent: The DashboardAgent object contains global config. + """ + + def __init__( + self, + runtime_env_dir, + logging_params, + gcs_client: GcsClient, + temp_dir, + address, + runtime_env_agent_port, + ): + super().__init__() + + self._logger = default_logger + self._logging_params = logging_params + self._logger = setup_component_logger( + logger_name=default_logger.name, **self._logging_params + ) + # Don't propagate logs to the root logger, because these logs + # might contain sensitive information. Instead, these logs should + # be confined to the runtime env agent log file `self.LOG_FILENAME`. + self._logger.propagate = False + + self._logger.info("Starting runtime env agent at pid %s", os.getpid()) + self._logger.info(f"Parent raylet pid is {os.environ.get('RAY_RAYLET_PID')}") + + self._runtime_env_dir = runtime_env_dir + self._per_job_logger_cache = dict() + # Cache the results of creating envs to avoid repeatedly calling into + # conda and other slow calls. + self._env_cache: Dict[str, CreatedEnvResult] = dict() + # Maps a serialized runtime env to a lock that is used + # to prevent multiple concurrent installs of the same env. + self._env_locks: Dict[str, asyncio.Lock] = dict() + self._gcs_client = gcs_client + + self._pip_plugin = PipPlugin(self._runtime_env_dir) + self._uv_plugin = UvPlugin(self._runtime_env_dir) + self._conda_plugin = CondaPlugin(self._runtime_env_dir) + self._py_modules_plugin = PyModulesPlugin( + self._runtime_env_dir, self._gcs_client + ) + self._py_executable_plugin = PyExecutablePlugin() + self._java_jars_plugin = JavaJarsPlugin(self._runtime_env_dir, self._gcs_client) + self._working_dir_plugin = WorkingDirPlugin( + self._runtime_env_dir, self._gcs_client + ) + self._container_plugin = ContainerPlugin(temp_dir) + # TODO(jonathan-anyscale): change the plugin to ProfilerPlugin + # and unify with nsight and other profilers. + self._nsight_plugin = NsightPlugin(self._runtime_env_dir) + self._rocprof_sys_plugin = RocProfSysPlugin(self._runtime_env_dir) + self._image_uri_plugin = get_image_uri_plugin_cls()(temp_dir) + + # TODO(architkulkarni): "base plugins" and third-party plugins should all go + # through the same code path. We should never need to refer to + # self._xxx_plugin, we should just iterate through self._plugins. + self._base_plugins: List[RuntimeEnvPlugin] = [ + self._working_dir_plugin, + self._uv_plugin, + self._pip_plugin, + self._conda_plugin, + self._py_modules_plugin, + self._py_executable_plugin, + self._java_jars_plugin, + self._container_plugin, + self._nsight_plugin, + self._rocprof_sys_plugin, + self._image_uri_plugin, + ] + self._plugin_manager = RuntimeEnvPluginManager() + for plugin in self._base_plugins: + self._plugin_manager.add_plugin(plugin) + + self._reference_table = ReferenceTable( + self.uris_parser, + self.unused_uris_processor, + self.unused_runtime_env_processor, + ) + + self._logger.info( + "Listening to address %s, port %d", address, runtime_env_agent_port + ) + + try: + self._node_ip = ray.util.get_node_ip_address() + self._node_prefix = f"[Node {self._node_ip}] " + except Exception as e: + self._logger.warning(f"Failed to get node IP address, using fallback: {e}") + self._node_prefix = "[Node unknown] " + + def uris_parser(self, runtime_env: RuntimeEnv): + result = list() + for name, plugin_setup_context in self._plugin_manager.plugins.items(): + plugin = plugin_setup_context.class_instance + uris = plugin.get_uris(runtime_env) + for uri in uris: + result.append((uri, UriType(name))) + return result + + def unused_uris_processor(self, unused_uris: List[Tuple[str, UriType]]) -> None: + for uri, uri_type in unused_uris: + self._plugin_manager.plugins[str(uri_type)].uri_cache.mark_unused(uri) + + def unused_runtime_env_processor(self, unused_runtime_env: str) -> None: + def delete_runtime_env(): + del self._env_cache[unused_runtime_env] + self._logger.info( + "Runtime env %s removed from env-level cache.", unused_runtime_env + ) + + if unused_runtime_env in self._env_cache: + if not self._env_cache[unused_runtime_env].success: + loop = get_or_create_event_loop() + # Cache the bad runtime env result by ttl seconds. + loop.call_later( + runtime_env_consts.BAD_RUNTIME_ENV_CACHE_TTL_SECONDS, + delete_runtime_env, + ) + else: + delete_runtime_env() + + def get_or_create_logger(self, job_id: bytes, log_files: List[str]): + job_id = job_id.decode() + if job_id not in self._per_job_logger_cache: + params = self._logging_params.copy() + params["filename"] = [f"runtime_env_setup-{job_id}.log", *log_files] + params["logger_name"] = f"runtime_env_{job_id}" + params["propagate"] = False + per_job_logger = setup_component_logger(**params) + self._per_job_logger_cache[job_id] = per_job_logger + return self._per_job_logger_cache[job_id] + + async def GetOrCreateRuntimeEnv(self, request): + self._logger.debug( + f"Got request from {request.source_process} to increase " + "reference for runtime env: " + f"{request.serialized_runtime_env}." + ) + + async def _setup_runtime_env( + runtime_env: RuntimeEnv, + runtime_env_config: RuntimeEnvConfig, + ): + log_files = runtime_env_config.get("log_files", []) + # Use a separate logger for each job. + per_job_logger = self.get_or_create_logger(request.job_id, log_files) + context = RuntimeEnvContext(env_vars=runtime_env.env_vars()) + + # Warn about unrecognized fields in the runtime env. + for name, _ in runtime_env.plugins(): + if name not in self._plugin_manager.plugins: + per_job_logger.warning( + f"runtime_env field {name} is not recognized by " + "Ray and will be ignored. In the future, unrecognized " + "fields in the runtime_env will raise an exception." + ) + + # Creates each runtime env URI by their priority. `working_dir` is special + # because it needs to be created before other plugins. All other plugins are + # created in the priority order (smaller priority value -> earlier to + # create), with a special environment variable being set to the working dir. + # ${RAY_RUNTIME_ENV_CREATE_WORKING_DIR} + + # First create working dir... + working_dir_ctx = self._plugin_manager.plugins[WorkingDirPlugin.name] + await create_for_plugin_if_needed( + runtime_env, + working_dir_ctx.class_instance, + working_dir_ctx.uri_cache, + context, + per_job_logger, + ) + + # Then within the working dir, create the other plugins. + working_dir_uri_or_none = runtime_env.working_dir_uri() + with self._working_dir_plugin.with_working_dir_env(working_dir_uri_or_none): + """Run setup for each plugin unless it has already been cached.""" + for ( + plugin_setup_context + ) in self._plugin_manager.sorted_plugin_setup_contexts(): + plugin = plugin_setup_context.class_instance + if plugin.name != WorkingDirPlugin.name: + uri_cache = plugin_setup_context.uri_cache + await create_for_plugin_if_needed( + runtime_env, plugin, uri_cache, context, per_job_logger + ) + return context + + async def _create_runtime_env_with_retry( + runtime_env, + setup_timeout_seconds, + runtime_env_config: RuntimeEnvConfig, + ) -> Tuple[bool, str, str]: + """Create runtime env with retry times. This function won't raise exceptions. + + Args: + runtime_env: The instance of RuntimeEnv class. + setup_timeout_seconds: The timeout of runtime environment creation for + each attempt. + runtime_env_config: The configuration for the runtime environment. + + Returns: + Tuple[bool, str, str]: A tuple containing: + - result (bool): Whether the creation was successful + - runtime_env_context (str): The serialized context if successful, None otherwise + - error_message (str): Error message if failed, None otherwise + """ + self._logger.info( + f"Creating runtime env: {serialized_env} with timeout " + f"{setup_timeout_seconds} seconds." + ) + num_retries = runtime_env_consts.RUNTIME_ENV_RETRY_TIMES + error_message = None + serialized_context = None + for i in range(num_retries): + # Only sleep when retrying. + if i != 0: + await asyncio.sleep( + runtime_env_consts.RUNTIME_ENV_RETRY_INTERVAL_MS / 1000 + ) + + try: + runtime_env_setup_task = _setup_runtime_env( + runtime_env, runtime_env_config + ) + runtime_env_context = await asyncio.wait_for( + runtime_env_setup_task, timeout=setup_timeout_seconds + ) + serialized_context = runtime_env_context.serialize() + error_message = None + break + except Exception as e: + err_msg = f"Failed to create runtime env {serialized_env}." + self._logger.exception(err_msg) + error_message = "".join( + traceback.format_exception(type(e), e, e.__traceback__) + ) + if isinstance(e, asyncio.TimeoutError): + hint = ( + f"Failed to install runtime_env within the " + f"timeout of {setup_timeout_seconds} seconds. Consider " + "increasing the timeout in the runtime_env config. " + "For example: \n" + ' runtime_env={"config": {"setup_timeout_seconds":' + " 1800}, ...}\n" + "If not provided, the default timeout is " + f"{DEFAULT_RUNTIME_ENV_TIMEOUT_SECONDS} seconds. " + ) + error_message = hint + error_message + + if error_message: + self._logger.error( + "runtime_env creation failed %d times, giving up.", + num_retries, + ) + return False, None, error_message + else: + self._logger.info( + "Successfully created runtime env: %s, context: %s", + serialized_env, + serialized_context, + ) + return True, serialized_context, None + + try: + serialized_env = request.serialized_runtime_env + runtime_env = RuntimeEnv.deserialize(serialized_env) + except Exception as e: + self._logger.exception( + "[Increase] Failed to parse runtime env: " f"{serialized_env}" + ) + + error_message = "".join( + traceback.format_exception(type(e), e, e.__traceback__) + ) + + return runtime_env_agent_pb2.GetOrCreateRuntimeEnvReply( + status=runtime_env_agent_pb2.AGENT_RPC_STATUS_FAILED, + error_message=f"{self._node_prefix}{error_message}", + ) + + # Increase reference + self._reference_table.increase_reference( + runtime_env, serialized_env, request.source_process + ) + + if serialized_env not in self._env_locks: + # async lock to prevent the same env being concurrently installed + self._env_locks[serialized_env] = asyncio.Lock() + + async with self._env_locks[serialized_env]: + if serialized_env in self._env_cache: + serialized_context = self._env_cache[serialized_env] + result = self._env_cache[serialized_env] + if result.success: + context = result.result + self._logger.info( + "Runtime env already created " + f"successfully. Env: {serialized_env}, " + f"context: {context}" + ) + return runtime_env_agent_pb2.GetOrCreateRuntimeEnvReply( + status=runtime_env_agent_pb2.AGENT_RPC_STATUS_OK, + serialized_runtime_env_context=context, + ) + else: + error_message = result.result + self._logger.info( + "Runtime env already failed. " + f"Env: {serialized_env}, " + f"err: {error_message}" + ) + # Recover the reference. + self._reference_table.decrease_reference( + runtime_env, serialized_env, request.source_process + ) + return runtime_env_agent_pb2.GetOrCreateRuntimeEnvReply( + status=runtime_env_agent_pb2.AGENT_RPC_STATUS_FAILED, + error_message=f"{self._node_prefix}{error_message}", + ) + + if SLEEP_FOR_TESTING_S: + self._logger.info(f"Sleeping for {SLEEP_FOR_TESTING_S}s.") + time.sleep(int(SLEEP_FOR_TESTING_S)) + + runtime_env_config = RuntimeEnvConfig.from_proto(request.runtime_env_config) + + # accroding to the document of `asyncio.wait_for`, + # None means disable timeout logic + setup_timeout_seconds = ( + None + if runtime_env_config["setup_timeout_seconds"] == -1 + else runtime_env_config["setup_timeout_seconds"] + ) + + start = time.perf_counter() + ( + successful, + serialized_context, + error_message, + ) = await _create_runtime_env_with_retry( + runtime_env, + setup_timeout_seconds, + runtime_env_config, + ) + creation_time_ms = int(round((time.perf_counter() - start) * 1000, 0)) + if not successful: + # Recover the reference. + self._reference_table.decrease_reference( + runtime_env, serialized_env, request.source_process + ) + # Add the result to env cache. + self._env_cache[serialized_env] = CreatedEnvResult( + successful, + serialized_context if successful else error_message, + creation_time_ms, + ) + # Reply the RPC + return runtime_env_agent_pb2.GetOrCreateRuntimeEnvReply( + status=runtime_env_agent_pb2.AGENT_RPC_STATUS_OK + if successful + else runtime_env_agent_pb2.AGENT_RPC_STATUS_FAILED, + serialized_runtime_env_context=serialized_context, + error_message=f"{self._node_prefix}{error_message}" + if not successful + else "", + ) + + async def DeleteRuntimeEnvIfPossible(self, request): + self._logger.info( + f"Got request from {request.source_process} to decrease " + "reference for runtime env: " + f"{request.serialized_runtime_env}." + ) + + try: + runtime_env = RuntimeEnv.deserialize(request.serialized_runtime_env) + except Exception as e: + self._logger.exception( + "[Decrease] Failed to parse runtime env: " + f"{request.serialized_runtime_env}" + ) + + error_message = "".join( + traceback.format_exception(type(e), e, e.__traceback__) + ) + + return runtime_env_agent_pb2.GetOrCreateRuntimeEnvReply( + status=runtime_env_agent_pb2.AGENT_RPC_STATUS_FAILED, + error_message=f"{self._node_prefix}{error_message}", + ) + + try: + self._reference_table.decrease_reference( + runtime_env, request.serialized_runtime_env, request.source_process + ) + except Exception as e: + return runtime_env_agent_pb2.DeleteRuntimeEnvIfPossibleReply( + status=runtime_env_agent_pb2.AGENT_RPC_STATUS_FAILED, + error_message=f"{self._node_prefix}Failed to decrement reference for runtime env for {str(e)}", + ) + + return runtime_env_agent_pb2.DeleteRuntimeEnvIfPossibleReply( + status=runtime_env_agent_pb2.AGENT_RPC_STATUS_OK + ) + + async def GetRuntimeEnvsInfo(self, request): + """Return the runtime env information of the node.""" + # TODO(sang): Currently, it only includes runtime_env information. + # We should include the URI information which includes, + # URIs + # Caller + # Ref counts + # Cache information + # Metrics (creation time & success) + # Deleted URIs + limit = request.limit if request.HasField("limit") else -1 + runtime_env_states = defaultdict(ProtoRuntimeEnvState) + runtime_env_refs = self._reference_table.runtime_env_refs + for runtime_env, ref_cnt in runtime_env_refs.items(): + runtime_env_states[runtime_env].runtime_env = runtime_env + runtime_env_states[runtime_env].ref_cnt = ref_cnt + for runtime_env, result in self._env_cache.items(): + runtime_env_states[runtime_env].runtime_env = runtime_env + runtime_env_states[runtime_env].success = result.success + if not result.success: + runtime_env_states[runtime_env].error = result.result + runtime_env_states[runtime_env].creation_time_ms = result.creation_time_ms + + reply = runtime_env_agent_pb2.GetRuntimeEnvsInfoReply() + count = 0 + for runtime_env_state in runtime_env_states.values(): + if limit != -1 and count >= limit: + break + count += 1 + reply.runtime_env_states.append(runtime_env_state) + reply.total = len(runtime_env_states) + return reply diff --git a/lib/python3.12/site-packages/ray/_private/runtime_env/agent/runtime_env_consts.py b/lib/python3.12/site-packages/ray/_private/runtime_env/agent/runtime_env_consts.py new file mode 100644 index 0000000000000000000000000000000000000000..31545913168fa29397995d4d0c6c118fe583365f --- /dev/null +++ b/lib/python3.12/site-packages/ray/_private/runtime_env/agent/runtime_env_consts.py @@ -0,0 +1,20 @@ +import ray._private.ray_constants as ray_constants + +RUNTIME_ENV_RETRY_TIMES = ray_constants.env_integer("RUNTIME_ENV_RETRY_TIMES", 3) + +RUNTIME_ENV_RETRY_INTERVAL_MS = ray_constants.env_integer( + "RUNTIME_ENV_RETRY_INTERVAL_MS", 1000 +) + +# Cache TTL for bad runtime env. After this time, delete the cache and retry to create +# runtime env if needed. +BAD_RUNTIME_ENV_CACHE_TTL_SECONDS = ray_constants.env_integer( + "BAD_RUNTIME_ENV_CACHE_TTL_SECONDS", 60 * 10 +) + +RUNTIME_ENV_LOG_FILENAME = "runtime_env.log" +RUNTIME_ENV_AGENT_PORT_PREFIX = "RUNTIME_ENV_AGENT_PORT_PREFIX:" +RUNTIME_ENV_AGENT_LOG_FILENAME = "runtime_env_agent.log" +RUNTIME_ENV_AGENT_CHECK_PARENT_INTERVAL_S_ENV_NAME = ( + "RAY_RUNTIME_ENV_AGENT_CHECK_PARENT_INTERVAL_S" # noqa +) diff --git a/lib/python3.12/site-packages/ray/_private/runtime_env/agent/thirdparty_files/aiosignal-1.4.0.dist-info/INSTALLER b/lib/python3.12/site-packages/ray/_private/runtime_env/agent/thirdparty_files/aiosignal-1.4.0.dist-info/INSTALLER new file mode 100644 index 0000000000000000000000000000000000000000..a1b589e38a32041e49332e5e81c2d363dc418d68 --- /dev/null +++ b/lib/python3.12/site-packages/ray/_private/runtime_env/agent/thirdparty_files/aiosignal-1.4.0.dist-info/INSTALLER @@ -0,0 +1 @@ +pip diff --git a/lib/python3.12/site-packages/ray/_private/runtime_env/agent/thirdparty_files/aiosignal-1.4.0.dist-info/METADATA b/lib/python3.12/site-packages/ray/_private/runtime_env/agent/thirdparty_files/aiosignal-1.4.0.dist-info/METADATA new file mode 100644 index 0000000000000000000000000000000000000000..03a6f0f7ff91628f60c4a95c4f3acbfa8d654ea8 --- /dev/null +++ b/lib/python3.12/site-packages/ray/_private/runtime_env/agent/thirdparty_files/aiosignal-1.4.0.dist-info/METADATA @@ -0,0 +1,112 @@ +Metadata-Version: 2.4 +Name: aiosignal +Version: 1.4.0 +Summary: aiosignal: a list of registered asynchronous callbacks +Home-page: https://github.com/aio-libs/aiosignal +Maintainer: aiohttp team +Maintainer-email: team@aiohttp.org +License: Apache 2.0 +Project-URL: Chat: Gitter, https://gitter.im/aio-libs/Lobby +Project-URL: CI: GitHub Actions, https://github.com/aio-libs/aiosignal/actions +Project-URL: Coverage: codecov, https://codecov.io/github/aio-libs/aiosignal +Project-URL: Docs: RTD, https://docs.aiosignal.org +Project-URL: GitHub: issues, https://github.com/aio-libs/aiosignal/issues +Project-URL: GitHub: repo, https://github.com/aio-libs/aiosignal +Classifier: License :: OSI Approved :: Apache Software License +Classifier: Intended Audience :: Developers +Classifier: Programming Language :: Python +Classifier: Programming Language :: Python :: 3 +Classifier: Programming Language :: Python :: 3 :: Only +Classifier: Development Status :: 5 - Production/Stable +Classifier: Operating System :: POSIX +Classifier: Operating System :: MacOS :: MacOS X +Classifier: Operating System :: Microsoft :: Windows +Classifier: Framework :: AsyncIO +Requires-Python: >=3.9 +Description-Content-Type: text/x-rst +License-File: LICENSE +Requires-Dist: frozenlist>=1.1.0 +Requires-Dist: typing-extensions>=4.2; python_version < "3.13" +Dynamic: license-file + +========= +aiosignal +========= + +.. image:: https://github.com/aio-libs/aiosignal/workflows/CI/badge.svg + :target: https://github.com/aio-libs/aiosignal/actions?query=workflow%3ACI + :alt: GitHub status for master branch + +.. image:: https://codecov.io/gh/aio-libs/aiosignal/branch/master/graph/badge.svg?flag=pytest + :target: https://codecov.io/gh/aio-libs/aiosignal?flags[0]=pytest + :alt: codecov.io status for master branch + +.. image:: https://badge.fury.io/py/aiosignal.svg + :target: https://pypi.org/project/aiosignal + :alt: Latest PyPI package version + +.. image:: https://readthedocs.org/projects/aiosignal/badge/?version=latest + :target: https://aiosignal.readthedocs.io/ + :alt: Latest Read The Docs + +.. image:: https://img.shields.io/discourse/topics?server=https%3A%2F%2Faio-libs.discourse.group%2F + :target: https://aio-libs.discourse.group/ + :alt: Discourse group for io-libs + +.. image:: https://badges.gitter.im/Join%20Chat.svg + :target: https://gitter.im/aio-libs/Lobby + :alt: Chat on Gitter + +Introduction +============ + +A project to manage callbacks in `asyncio` projects. + +``Signal`` is a list of registered asynchronous callbacks. + +The signal's life-cycle has two stages: after creation its content +could be filled by using standard list operations: ``sig.append()`` +etc. + +After you call ``sig.freeze()`` the signal is *frozen*: adding, removing +and dropping callbacks is forbidden. + +The only available operation is calling the previously registered +callbacks by using ``await sig.send(data)``. + +For concrete usage examples see the `Signals + +section of the `Web Server Advanced +` chapter of the `aiohttp +documentation`_. + + +Installation +------------ + +:: + + $ pip install aiosignal + + +Documentation +============= + +https://aiosignal.readthedocs.io/ + +License +======= + +``aiosignal`` is offered under the Apache 2 license. + +Source code +=========== + +The project is hosted on GitHub_ + +Please file an issue in the `bug tracker +`_ if you have found a bug +or have some suggestions to improve the library. + +.. _GitHub: https://github.com/aio-libs/aiosignal +.. _aiohttp documentation: https://docs.aiohttp.org/ diff --git a/lib/python3.12/site-packages/ray/_private/runtime_env/agent/thirdparty_files/aiosignal-1.4.0.dist-info/RECORD b/lib/python3.12/site-packages/ray/_private/runtime_env/agent/thirdparty_files/aiosignal-1.4.0.dist-info/RECORD new file mode 100644 index 0000000000000000000000000000000000000000..df5ca06f6fe1bbec1a02456b5b001a34771ecced --- /dev/null +++ b/lib/python3.12/site-packages/ray/_private/runtime_env/agent/thirdparty_files/aiosignal-1.4.0.dist-info/RECORD @@ -0,0 +1,9 @@ +aiosignal-1.4.0.dist-info/INSTALLER,sha256=zuuue4knoyJ-UwPPXg8fezS7VCrXJQrAP7zeNuwvFQg,4 +aiosignal-1.4.0.dist-info/METADATA,sha256=CSR-8dqLxpZyjUcTDnAuQwf299EB1sSFv_nzpxznAI0,3662 +aiosignal-1.4.0.dist-info/RECORD,, +aiosignal-1.4.0.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91 +aiosignal-1.4.0.dist-info/licenses/LICENSE,sha256=b9UkPpLdf5jsacesN3co50kFcJ_1J6W_mNbQJjwE9bY,11332 +aiosignal-1.4.0.dist-info/top_level.txt,sha256=z45aNOKGDdrI1roqZY3BGXQ22kJFPHBmVdwtLYLtXC0,10 +aiosignal/__init__.py,sha256=TIkmUG9HTBt4dfq2nISYBiZiRB2xwvFtEZydLP0HPL4,1537 +aiosignal/__pycache__/__init__.cpython-312.pyc,, +aiosignal/py.typed,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0 diff --git a/lib/python3.12/site-packages/ray/_private/runtime_env/agent/thirdparty_files/aiosignal-1.4.0.dist-info/WHEEL b/lib/python3.12/site-packages/ray/_private/runtime_env/agent/thirdparty_files/aiosignal-1.4.0.dist-info/WHEEL new file mode 100644 index 0000000000000000000000000000000000000000..e7fa31b6f3f78deb1022c1f7927f07d4d16da822 --- /dev/null +++ b/lib/python3.12/site-packages/ray/_private/runtime_env/agent/thirdparty_files/aiosignal-1.4.0.dist-info/WHEEL @@ -0,0 +1,5 @@ +Wheel-Version: 1.0 +Generator: setuptools (80.9.0) +Root-Is-Purelib: true +Tag: py3-none-any + diff --git a/lib/python3.12/site-packages/ray/_private/runtime_env/agent/thirdparty_files/aiosignal-1.4.0.dist-info/licenses/LICENSE b/lib/python3.12/site-packages/ray/_private/runtime_env/agent/thirdparty_files/aiosignal-1.4.0.dist-info/licenses/LICENSE new file mode 100644 index 0000000000000000000000000000000000000000..7082a2d5b9047bfc09589f387053e24ea490bc54 --- /dev/null +++ b/lib/python3.12/site-packages/ray/_private/runtime_env/agent/thirdparty_files/aiosignal-1.4.0.dist-info/licenses/LICENSE @@ -0,0 +1,201 @@ +Apache License + Version 2.0, January 2004 + http://www.apache.org/licenses/ + + TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION + + 1. Definitions. + + "License" shall mean the terms and conditions for use, reproduction, + and distribution as defined by Sections 1 through 9 of this document. + + "Licensor" shall mean the copyright owner or entity authorized by + the copyright owner that is granting the License. + + "Legal Entity" shall mean the union of the acting entity and all + other entities that control, are controlled by, or are under common + control with that entity. For the purposes of this definition, + "control" means (i) the power, direct or indirect, to cause the + direction or management of such entity, whether by contract or + otherwise, or (ii) ownership of fifty percent (50%) or more of the + outstanding shares, or (iii) beneficial ownership of such entity. + + "You" (or "Your") shall mean an individual or Legal Entity + exercising permissions granted by this License. + + "Source" form shall mean the preferred form for making modifications, + including but not limited to software source code, documentation + source, and configuration files. + + "Object" form shall mean any form resulting from mechanical + transformation or translation of a Source form, including but + not limited to compiled object code, generated documentation, + and conversions to other media types. + + "Work" shall mean the work of authorship, whether in Source or + Object form, made available under the License, as indicated by a + copyright notice that is included in or attached to the work + (an example is provided in the Appendix below). + + "Derivative Works" shall mean any work, whether in Source or Object + form, that is based on (or derived from) the Work and for which the + editorial revisions, annotations, elaborations, or other modifications + represent, as a whole, an original work of authorship. For the purposes + of this License, Derivative Works shall not include works that remain + separable from, or merely link (or bind by name) to the interfaces of, + the Work and Derivative Works thereof. + + "Contribution" shall mean any work of authorship, including + the original version of the Work and any modifications or additions + to that Work or Derivative Works thereof, that is intentionally + submitted to Licensor for inclusion in the Work by the copyright owner + or by an individual or Legal Entity authorized to submit on behalf of + the copyright owner. For the purposes of this definition, "submitted" + means any form of electronic, verbal, or written communication sent + to the Licensor or its representatives, including but not limited to + communication on electronic mailing lists, source code control systems, + and issue tracking systems that are managed by, or on behalf of, the + Licensor for the purpose of discussing and improving the Work, but + excluding communication that is conspicuously marked or otherwise + designated in writing by the copyright owner as "Not a Contribution." + + "Contributor" shall mean Licensor and any individual or Legal Entity + on behalf of whom a Contribution has been received by Licensor and + subsequently incorporated within the Work. + + 2. Grant of Copyright License. Subject to the terms and conditions of + this License, each Contributor hereby grants to You a perpetual, + worldwide, non-exclusive, no-charge, royalty-free, irrevocable + copyright license to reproduce, prepare Derivative Works of, + publicly display, publicly perform, sublicense, and distribute the + Work and such Derivative Works in Source or Object form. + + 3. Grant of Patent License. Subject to the terms and conditions of + this License, each Contributor hereby grants to You a perpetual, + worldwide, non-exclusive, no-charge, royalty-free, irrevocable + (except as stated in this section) patent license to make, have made, + use, offer to sell, sell, import, and otherwise transfer the Work, + where such license applies only to those patent claims licensable + by such Contributor that are necessarily infringed by their + Contribution(s) alone or by combination of their Contribution(s) + with the Work to which such Contribution(s) was submitted. If You + institute patent litigation against any entity (including a + cross-claim or counterclaim in a lawsuit) alleging that the Work + or a Contribution incorporated within the Work constitutes direct + or contributory patent infringement, then any patent licenses + granted to You under this License for that Work shall terminate + as of the date such litigation is filed. + + 4. Redistribution. You may reproduce and distribute copies of the + Work or Derivative Works thereof in any medium, with or without + modifications, and in Source or Object form, provided that You + meet the following conditions: + + (a) You must give any other recipients of the Work or + Derivative Works a copy of this License; and + + (b) You must cause any modified files to carry prominent notices + stating that You changed the files; and + + (c) You must retain, in the Source form of any Derivative Works + that You distribute, all copyright, patent, trademark, and + attribution notices from the Source form of the Work, + excluding those notices that do not pertain to any part of + the Derivative Works; and + + (d) If the Work includes a "NOTICE" text file as part of its + distribution, then any Derivative Works that You distribute must + include a readable copy of the attribution notices contained + within such NOTICE file, excluding those notices that do not + pertain to any part of the Derivative Works, in at least one + of the following places: within a NOTICE text file distributed + as part of the Derivative Works; within the Source form or + documentation, if provided along with the Derivative Works; or, + within a display generated by the Derivative Works, if and + wherever such third-party notices normally appear. The contents + of the NOTICE file are for informational purposes only and + do not modify the License. You may add Your own attribution + notices within Derivative Works that You distribute, alongside + or as an addendum to the NOTICE text from the Work, provided + that such additional attribution notices cannot be construed + as modifying the License. + + You may add Your own copyright statement to Your modifications and + may provide additional or different license terms and conditions + for use, reproduction, or distribution of Your modifications, or + for any such Derivative Works as a whole, provided Your use, + reproduction, and distribution of the Work otherwise complies with + the conditions stated in this License. + + 5. Submission of Contributions. Unless You explicitly state otherwise, + any Contribution intentionally submitted for inclusion in the Work + by You to the Licensor shall be under the terms and conditions of + this License, without any additional terms or conditions. + Notwithstanding the above, nothing herein shall supersede or modify + the terms of any separate license agreement you may have executed + with Licensor regarding such Contributions. + + 6. Trademarks. This License does not grant permission to use the trade + names, trademarks, service marks, or product names of the Licensor, + except as required for reasonable and customary use in describing the + origin of the Work and reproducing the content of the NOTICE file. + + 7. Disclaimer of Warranty. Unless required by applicable law or + agreed to in writing, Licensor provides the Work (and each + Contributor provides its Contributions) on an "AS IS" BASIS, + WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or + implied, including, without limitation, any warranties or conditions + of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A + PARTICULAR PURPOSE. You are solely responsible for determining the + appropriateness of using or redistributing the Work and assume any + risks associated with Your exercise of permissions under this License. + + 8. Limitation of Liability. In no event and under no legal theory, + whether in tort (including negligence), contract, or otherwise, + unless required by applicable law (such as deliberate and grossly + negligent acts) or agreed to in writing, shall any Contributor be + liable to You for damages, including any direct, indirect, special, + incidental, or consequential damages of any character arising as a + result of this License or out of the use or inability to use the + Work (including but not limited to damages for loss of goodwill, + work stoppage, computer failure or malfunction, or any and all + other commercial damages or losses), even if such Contributor + has been advised of the possibility of such damages. + + 9. Accepting Warranty or Additional Liability. While redistributing + the Work or Derivative Works thereof, You may choose to offer, + and charge a fee for, acceptance of support, warranty, indemnity, + or other liability obligations and/or rights consistent with this + License. However, in accepting such obligations, You may act only + on Your own behalf and on Your sole responsibility, not on behalf + of any other Contributor, and only if You agree to indemnify, + defend, and hold each Contributor harmless for any liability + incurred by, or claims asserted against, such Contributor by reason + of your accepting any such warranty or additional liability. + + END OF TERMS AND CONDITIONS + + APPENDIX: How to apply the Apache License to your work. + + To apply the Apache License to your work, attach the following + boilerplate notice, with the fields enclosed by brackets "{}" + replaced with your own identifying information. (Don't include + the brackets!) The text should be enclosed in the appropriate + comment syntax for the file format. We also recommend that a + file or class name and description of purpose be included on the + same "printed page" as the copyright notice for easier + identification within third-party archives. + + Copyright 2013-2019 Nikolay Kim and Andrew Svetlov + + Licensed under the Apache License, Version 2.0 (the "License"); + you may not use this file except in compliance with the License. + You may obtain a copy of the License at + + http://www.apache.org/licenses/LICENSE-2.0 + + Unless required by applicable law or agreed to in writing, software + distributed under the License is distributed on an "AS IS" BASIS, + WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + See the License for the specific language governing permissions and + limitations under the License. diff --git a/lib/python3.12/site-packages/ray/_private/runtime_env/agent/thirdparty_files/aiosignal-1.4.0.dist-info/top_level.txt b/lib/python3.12/site-packages/ray/_private/runtime_env/agent/thirdparty_files/aiosignal-1.4.0.dist-info/top_level.txt new file mode 100644 index 0000000000000000000000000000000000000000..ac6df3afe74a5fd43afc7ab7f8393571a495fdc5 --- /dev/null +++ b/lib/python3.12/site-packages/ray/_private/runtime_env/agent/thirdparty_files/aiosignal-1.4.0.dist-info/top_level.txt @@ -0,0 +1 @@ +aiosignal diff --git a/lib/python3.12/site-packages/ray/_private/runtime_env/agent/thirdparty_files/attrs/__init__.py b/lib/python3.12/site-packages/ray/_private/runtime_env/agent/thirdparty_files/attrs/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..dc1ce4b974ea1b3971ced44e3e112e897cf29b0a --- /dev/null +++ b/lib/python3.12/site-packages/ray/_private/runtime_env/agent/thirdparty_files/attrs/__init__.py @@ -0,0 +1,72 @@ +# SPDX-License-Identifier: MIT + +from attr import ( + NOTHING, + Attribute, + AttrsInstance, + Converter, + Factory, + NothingType, + _make_getattr, + assoc, + cmp_using, + define, + evolve, + field, + fields, + fields_dict, + frozen, + has, + make_class, + mutable, + resolve_types, + validate, +) +from attr._make import ClassProps +from attr._next_gen import asdict, astuple, inspect + +from . import converters, exceptions, filters, setters, validators + + +__all__ = [ + "NOTHING", + "Attribute", + "AttrsInstance", + "ClassProps", + "Converter", + "Factory", + "NothingType", + "__author__", + "__copyright__", + "__description__", + "__doc__", + "__email__", + "__license__", + "__title__", + "__url__", + "__version__", + "__version_info__", + "asdict", + "assoc", + "astuple", + "cmp_using", + "converters", + "define", + "evolve", + "exceptions", + "field", + "fields", + "fields_dict", + "filters", + "frozen", + "has", + "inspect", + "make_class", + "mutable", + "resolve_types", + "setters", + "validate", + "validators", +] + +__getattr__ = _make_getattr(__name__) diff --git a/lib/python3.12/site-packages/ray/_private/runtime_env/agent/thirdparty_files/attrs/__init__.pyi b/lib/python3.12/site-packages/ray/_private/runtime_env/agent/thirdparty_files/attrs/__init__.pyi new file mode 100644 index 0000000000000000000000000000000000000000..6364bac4ea29f860934803e5f3ba741812c6dfd1 --- /dev/null +++ b/lib/python3.12/site-packages/ray/_private/runtime_env/agent/thirdparty_files/attrs/__init__.pyi @@ -0,0 +1,314 @@ +import sys + +from typing import ( + Any, + Callable, + Mapping, + Sequence, + overload, + TypeVar, +) + +# Because we need to type our own stuff, we have to make everything from +# attr explicitly public too. +from attr import __author__ as __author__ +from attr import __copyright__ as __copyright__ +from attr import __description__ as __description__ +from attr import __email__ as __email__ +from attr import __license__ as __license__ +from attr import __title__ as __title__ +from attr import __url__ as __url__ +from attr import __version__ as __version__ +from attr import __version_info__ as __version_info__ +from attr import assoc as assoc +from attr import Attribute as Attribute +from attr import AttrsInstance as AttrsInstance +from attr import cmp_using as cmp_using +from attr import converters as converters +from attr import Converter as Converter +from attr import evolve as evolve +from attr import exceptions as exceptions +from attr import Factory as Factory +from attr import fields as fields +from attr import fields_dict as fields_dict +from attr import filters as filters +from attr import has as has +from attr import make_class as make_class +from attr import NOTHING as NOTHING +from attr import resolve_types as resolve_types +from attr import setters as setters +from attr import validate as validate +from attr import validators as validators +from attr import attrib, asdict as asdict, astuple as astuple +from attr import NothingType as NothingType + +if sys.version_info >= (3, 11): + from typing import dataclass_transform +else: + from typing_extensions import dataclass_transform + +_T = TypeVar("_T") +_C = TypeVar("_C", bound=type) + +_EqOrderType = bool | Callable[[Any], Any] +_ValidatorType = Callable[[Any, "Attribute[_T]", _T], Any] +_CallableConverterType = Callable[[Any], Any] +_ConverterType = _CallableConverterType | Converter[Any, Any] +_ReprType = Callable[[Any], str] +_ReprArgType = bool | _ReprType +_OnSetAttrType = Callable[[Any, "Attribute[Any]", Any], Any] +_OnSetAttrArgType = _OnSetAttrType | list[_OnSetAttrType] | setters._NoOpType +_FieldTransformer = Callable[ + [type, list["Attribute[Any]"]], list["Attribute[Any]"] +] +# FIXME: in reality, if multiple validators are passed they must be in a list +# or tuple, but those are invariant and so would prevent subtypes of +# _ValidatorType from working when passed in a list or tuple. +_ValidatorArgType = _ValidatorType[_T] | Sequence[_ValidatorType[_T]] + +@overload +def field( + *, + default: None = ..., + validator: None = ..., + repr: _ReprArgType = ..., + hash: bool | None = ..., + init: bool = ..., + metadata: Mapping[Any, Any] | None = ..., + converter: None = ..., + factory: None = ..., + kw_only: bool | None = ..., + eq: bool | None = ..., + order: bool | None = ..., + on_setattr: _OnSetAttrArgType | None = ..., + alias: str | None = ..., + type: type | None = ..., +) -> Any: ... + +# This form catches an explicit None or no default and infers the type from the +# other arguments. +@overload +def field( + *, + default: None = ..., + validator: _ValidatorArgType[_T] | None = ..., + repr: _ReprArgType = ..., + hash: bool | None = ..., + init: bool = ..., + metadata: Mapping[Any, Any] | None = ..., + converter: _ConverterType + | list[_ConverterType] + | tuple[_ConverterType, ...] + | None = ..., + factory: Callable[[], _T] | None = ..., + kw_only: bool | None = ..., + eq: _EqOrderType | None = ..., + order: _EqOrderType | None = ..., + on_setattr: _OnSetAttrArgType | None = ..., + alias: str | None = ..., + type: type | None = ..., +) -> _T: ... + +# This form catches an explicit default argument. +@overload +def field( + *, + default: _T, + validator: _ValidatorArgType[_T] | None = ..., + repr: _ReprArgType = ..., + hash: bool | None = ..., + init: bool = ..., + metadata: Mapping[Any, Any] | None = ..., + converter: _ConverterType + | list[_ConverterType] + | tuple[_ConverterType, ...] + | None = ..., + factory: Callable[[], _T] | None = ..., + kw_only: bool | None = ..., + eq: _EqOrderType | None = ..., + order: _EqOrderType | None = ..., + on_setattr: _OnSetAttrArgType | None = ..., + alias: str | None = ..., + type: type | None = ..., +) -> _T: ... + +# This form covers type=non-Type: e.g. forward references (str), Any +@overload +def field( + *, + default: _T | None = ..., + validator: _ValidatorArgType[_T] | None = ..., + repr: _ReprArgType = ..., + hash: bool | None = ..., + init: bool = ..., + metadata: Mapping[Any, Any] | None = ..., + converter: _ConverterType + | list[_ConverterType] + | tuple[_ConverterType, ...] + | None = ..., + factory: Callable[[], _T] | None = ..., + kw_only: bool | None = ..., + eq: _EqOrderType | None = ..., + order: _EqOrderType | None = ..., + on_setattr: _OnSetAttrArgType | None = ..., + alias: str | None = ..., + type: type | None = ..., +) -> Any: ... +@overload +@dataclass_transform(field_specifiers=(attrib, field)) +def define( + maybe_cls: _C, + *, + these: dict[str, Any] | None = ..., + repr: bool = ..., + unsafe_hash: bool | None = ..., + hash: bool | None = ..., + init: bool = ..., + slots: bool = ..., + frozen: bool = ..., + weakref_slot: bool = ..., + str: bool = ..., + auto_attribs: bool = ..., + kw_only: bool = ..., + cache_hash: bool = ..., + auto_exc: bool = ..., + eq: bool | None = ..., + order: bool | None = ..., + auto_detect: bool = ..., + getstate_setstate: bool | None = ..., + on_setattr: _OnSetAttrArgType | None = ..., + field_transformer: _FieldTransformer | None = ..., + match_args: bool = ..., +) -> _C: ... +@overload +@dataclass_transform(field_specifiers=(attrib, field)) +def define( + maybe_cls: None = ..., + *, + these: dict[str, Any] | None = ..., + repr: bool = ..., + unsafe_hash: bool | None = ..., + hash: bool | None = ..., + init: bool = ..., + slots: bool = ..., + frozen: bool = ..., + weakref_slot: bool = ..., + str: bool = ..., + auto_attribs: bool = ..., + kw_only: bool = ..., + cache_hash: bool = ..., + auto_exc: bool = ..., + eq: bool | None = ..., + order: bool | None = ..., + auto_detect: bool = ..., + getstate_setstate: bool | None = ..., + on_setattr: _OnSetAttrArgType | None = ..., + field_transformer: _FieldTransformer | None = ..., + match_args: bool = ..., +) -> Callable[[_C], _C]: ... + +mutable = define + +@overload +@dataclass_transform(frozen_default=True, field_specifiers=(attrib, field)) +def frozen( + maybe_cls: _C, + *, + these: dict[str, Any] | None = ..., + repr: bool = ..., + unsafe_hash: bool | None = ..., + hash: bool | None = ..., + init: bool = ..., + slots: bool = ..., + frozen: bool = ..., + weakref_slot: bool = ..., + str: bool = ..., + auto_attribs: bool = ..., + kw_only: bool = ..., + cache_hash: bool = ..., + auto_exc: bool = ..., + eq: bool | None = ..., + order: bool | None = ..., + auto_detect: bool = ..., + getstate_setstate: bool | None = ..., + on_setattr: _OnSetAttrArgType | None = ..., + field_transformer: _FieldTransformer | None = ..., + match_args: bool = ..., +) -> _C: ... +@overload +@dataclass_transform(frozen_default=True, field_specifiers=(attrib, field)) +def frozen( + maybe_cls: None = ..., + *, + these: dict[str, Any] | None = ..., + repr: bool = ..., + unsafe_hash: bool | None = ..., + hash: bool | None = ..., + init: bool = ..., + slots: bool = ..., + frozen: bool = ..., + weakref_slot: bool = ..., + str: bool = ..., + auto_attribs: bool = ..., + kw_only: bool = ..., + cache_hash: bool = ..., + auto_exc: bool = ..., + eq: bool | None = ..., + order: bool | None = ..., + auto_detect: bool = ..., + getstate_setstate: bool | None = ..., + on_setattr: _OnSetAttrArgType | None = ..., + field_transformer: _FieldTransformer | None = ..., + match_args: bool = ..., +) -> Callable[[_C], _C]: ... + +class ClassProps: + # XXX: somehow when defining/using enums Mypy starts looking at our own + # (untyped) code and causes tons of errors. + Hashability: Any + KeywordOnly: Any + + is_exception: bool + is_slotted: bool + has_weakref_slot: bool + is_frozen: bool + # kw_only: ClassProps.KeywordOnly + kw_only: Any + collected_fields_by_mro: bool + added_init: bool + added_repr: bool + added_eq: bool + added_ordering: bool + # hashability: ClassProps.Hashability + hashability: Any + added_match_args: bool + added_str: bool + added_pickling: bool + on_setattr_hook: _OnSetAttrType | None + field_transformer: Callable[[Attribute[Any]], Attribute[Any]] | None + + def __init__( + self, + is_exception: bool, + is_slotted: bool, + has_weakref_slot: bool, + is_frozen: bool, + # kw_only: ClassProps.KeywordOnly + kw_only: Any, + collected_fields_by_mro: bool, + added_init: bool, + added_repr: bool, + added_eq: bool, + added_ordering: bool, + # hashability: ClassProps.Hashability + hashability: Any, + added_match_args: bool, + added_str: bool, + added_pickling: bool, + on_setattr_hook: _OnSetAttrType, + field_transformer: Callable[[Attribute[Any]], Attribute[Any]], + ) -> None: ... + @property + def is_hashable(self) -> bool: ... + +def inspect(cls: type) -> ClassProps: ... diff --git a/lib/python3.12/site-packages/ray/_private/runtime_env/agent/thirdparty_files/attrs/__pycache__/__init__.cpython-312.pyc b/lib/python3.12/site-packages/ray/_private/runtime_env/agent/thirdparty_files/attrs/__pycache__/__init__.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..c2562c69d43d64e8586d59076116b8ed9a5a8b72 Binary files /dev/null and b/lib/python3.12/site-packages/ray/_private/runtime_env/agent/thirdparty_files/attrs/__pycache__/__init__.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/_private/runtime_env/agent/thirdparty_files/attrs/__pycache__/converters.cpython-312.pyc b/lib/python3.12/site-packages/ray/_private/runtime_env/agent/thirdparty_files/attrs/__pycache__/converters.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..25ca926141b599239c4ca5dd7f151ec9cebe3517 Binary files /dev/null and b/lib/python3.12/site-packages/ray/_private/runtime_env/agent/thirdparty_files/attrs/__pycache__/converters.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/_private/runtime_env/agent/thirdparty_files/attrs/__pycache__/exceptions.cpython-312.pyc b/lib/python3.12/site-packages/ray/_private/runtime_env/agent/thirdparty_files/attrs/__pycache__/exceptions.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..195700d73c86e90acf7856c39d876b95457dd815 Binary files /dev/null and b/lib/python3.12/site-packages/ray/_private/runtime_env/agent/thirdparty_files/attrs/__pycache__/exceptions.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/_private/runtime_env/agent/thirdparty_files/attrs/__pycache__/filters.cpython-312.pyc b/lib/python3.12/site-packages/ray/_private/runtime_env/agent/thirdparty_files/attrs/__pycache__/filters.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..8b187698bc3904ca0beecefbd61e361bcc806156 Binary files /dev/null and b/lib/python3.12/site-packages/ray/_private/runtime_env/agent/thirdparty_files/attrs/__pycache__/filters.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/_private/runtime_env/agent/thirdparty_files/attrs/__pycache__/setters.cpython-312.pyc b/lib/python3.12/site-packages/ray/_private/runtime_env/agent/thirdparty_files/attrs/__pycache__/setters.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..d6fe2451a1e273a2f214a51917ae1b25e5d745b0 Binary files /dev/null and b/lib/python3.12/site-packages/ray/_private/runtime_env/agent/thirdparty_files/attrs/__pycache__/setters.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/_private/runtime_env/agent/thirdparty_files/attrs/__pycache__/validators.cpython-312.pyc b/lib/python3.12/site-packages/ray/_private/runtime_env/agent/thirdparty_files/attrs/__pycache__/validators.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..80aaee97df4733e7807d7b799d818e9fa51915c4 Binary files /dev/null and b/lib/python3.12/site-packages/ray/_private/runtime_env/agent/thirdparty_files/attrs/__pycache__/validators.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/_private/runtime_env/agent/thirdparty_files/attrs/converters.py b/lib/python3.12/site-packages/ray/_private/runtime_env/agent/thirdparty_files/attrs/converters.py new file mode 100644 index 0000000000000000000000000000000000000000..7821f6c02cca81277d1ecc87b6bdafad886d8b70 --- /dev/null +++ b/lib/python3.12/site-packages/ray/_private/runtime_env/agent/thirdparty_files/attrs/converters.py @@ -0,0 +1,3 @@ +# SPDX-License-Identifier: MIT + +from attr.converters import * # noqa: F403 diff --git a/lib/python3.12/site-packages/ray/_private/runtime_env/agent/thirdparty_files/attrs/exceptions.py b/lib/python3.12/site-packages/ray/_private/runtime_env/agent/thirdparty_files/attrs/exceptions.py new file mode 100644 index 0000000000000000000000000000000000000000..3323f9d2112c54b203763d45b455bd5abbe020f6 --- /dev/null +++ b/lib/python3.12/site-packages/ray/_private/runtime_env/agent/thirdparty_files/attrs/exceptions.py @@ -0,0 +1,3 @@ +# SPDX-License-Identifier: MIT + +from attr.exceptions import * # noqa: F403 diff --git a/lib/python3.12/site-packages/ray/_private/runtime_env/agent/thirdparty_files/attrs/filters.py b/lib/python3.12/site-packages/ray/_private/runtime_env/agent/thirdparty_files/attrs/filters.py new file mode 100644 index 0000000000000000000000000000000000000000..3080f48398e5ed8d3428ca3efeb7500633b0cb0f --- /dev/null +++ b/lib/python3.12/site-packages/ray/_private/runtime_env/agent/thirdparty_files/attrs/filters.py @@ -0,0 +1,3 @@ +# SPDX-License-Identifier: MIT + +from attr.filters import * # noqa: F403 diff --git a/lib/python3.12/site-packages/ray/_private/runtime_env/agent/thirdparty_files/attrs/py.typed b/lib/python3.12/site-packages/ray/_private/runtime_env/agent/thirdparty_files/attrs/py.typed new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/lib/python3.12/site-packages/ray/_private/runtime_env/agent/thirdparty_files/attrs/setters.py b/lib/python3.12/site-packages/ray/_private/runtime_env/agent/thirdparty_files/attrs/setters.py new file mode 100644 index 0000000000000000000000000000000000000000..f3d73bb793dd49c138950961f41943bb26c57fde --- /dev/null +++ b/lib/python3.12/site-packages/ray/_private/runtime_env/agent/thirdparty_files/attrs/setters.py @@ -0,0 +1,3 @@ +# SPDX-License-Identifier: MIT + +from attr.setters import * # noqa: F403 diff --git a/lib/python3.12/site-packages/ray/_private/runtime_env/agent/thirdparty_files/attrs/validators.py b/lib/python3.12/site-packages/ray/_private/runtime_env/agent/thirdparty_files/attrs/validators.py new file mode 100644 index 0000000000000000000000000000000000000000..037e124f29f32d37c1642d159bf828de44f7c349 --- /dev/null +++ b/lib/python3.12/site-packages/ray/_private/runtime_env/agent/thirdparty_files/attrs/validators.py @@ -0,0 +1,3 @@ +# SPDX-License-Identifier: MIT + +from attr.validators import * # noqa: F403 diff --git a/lib/python3.12/site-packages/ray/_private/runtime_env/agent/thirdparty_files/typing_extensions.py b/lib/python3.12/site-packages/ray/_private/runtime_env/agent/thirdparty_files/typing_extensions.py new file mode 100644 index 0000000000000000000000000000000000000000..77f33e1614fd7d46ccd66b394d6c5d663bf8a8c6 --- /dev/null +++ b/lib/python3.12/site-packages/ray/_private/runtime_env/agent/thirdparty_files/typing_extensions.py @@ -0,0 +1,4317 @@ +import abc +import builtins +import collections +import collections.abc +import contextlib +import enum +import functools +import inspect +import io +import keyword +import operator +import sys +import types as _types +import typing +import warnings + +# Breakpoint: https://github.com/python/cpython/pull/119891 +if sys.version_info >= (3, 14): + import annotationlib + +__all__ = [ + # Super-special typing primitives. + 'Any', + 'ClassVar', + 'Concatenate', + 'Final', + 'LiteralString', + 'ParamSpec', + 'ParamSpecArgs', + 'ParamSpecKwargs', + 'Self', + 'Type', + 'TypeVar', + 'TypeVarTuple', + 'Unpack', + + # ABCs (from collections.abc). + 'Awaitable', + 'AsyncIterator', + 'AsyncIterable', + 'Coroutine', + 'AsyncGenerator', + 'AsyncContextManager', + 'Buffer', + 'ChainMap', + + # Concrete collection types. + 'ContextManager', + 'Counter', + 'Deque', + 'DefaultDict', + 'NamedTuple', + 'OrderedDict', + 'TypedDict', + + # Structural checks, a.k.a. protocols. + 'SupportsAbs', + 'SupportsBytes', + 'SupportsComplex', + 'SupportsFloat', + 'SupportsIndex', + 'SupportsInt', + 'SupportsRound', + 'Reader', + 'Writer', + + # One-off things. + 'Annotated', + 'assert_never', + 'assert_type', + 'clear_overloads', + 'dataclass_transform', + 'deprecated', + 'disjoint_base', + 'Doc', + 'evaluate_forward_ref', + 'get_overloads', + 'final', + 'Format', + 'get_annotations', + 'get_args', + 'get_origin', + 'get_original_bases', + 'get_protocol_members', + 'get_type_hints', + 'IntVar', + 'is_protocol', + 'is_typeddict', + 'Literal', + 'NewType', + 'overload', + 'override', + 'Protocol', + 'Sentinel', + 'reveal_type', + 'runtime', + 'runtime_checkable', + 'Text', + 'TypeAlias', + 'TypeAliasType', + 'TypeForm', + 'TypeGuard', + 'TypeIs', + 'TYPE_CHECKING', + 'type_repr', + 'Never', + 'NoReturn', + 'ReadOnly', + 'Required', + 'NotRequired', + 'NoDefault', + 'NoExtraItems', + + # Pure aliases, have always been in typing + 'AbstractSet', + 'AnyStr', + 'BinaryIO', + 'Callable', + 'Collection', + 'Container', + 'Dict', + 'ForwardRef', + 'FrozenSet', + 'Generator', + 'Generic', + 'Hashable', + 'IO', + 'ItemsView', + 'Iterable', + 'Iterator', + 'KeysView', + 'List', + 'Mapping', + 'MappingView', + 'Match', + 'MutableMapping', + 'MutableSequence', + 'MutableSet', + 'Optional', + 'Pattern', + 'Reversible', + 'Sequence', + 'Set', + 'Sized', + 'TextIO', + 'Tuple', + 'Union', + 'ValuesView', + 'cast', + 'no_type_check', + 'no_type_check_decorator', +] + +# for backward compatibility +PEP_560 = True +GenericMeta = type +# Breakpoint: https://github.com/python/cpython/pull/116129 +_PEP_696_IMPLEMENTED = sys.version_info >= (3, 13, 0, "beta") + +# Added with bpo-45166 to 3.10.1+ and some 3.9 versions +_FORWARD_REF_HAS_CLASS = "__forward_is_class__" in typing.ForwardRef.__slots__ + +# The functions below are modified copies of typing internal helpers. +# They are needed by _ProtocolMeta and they provide support for PEP 646. + + +class _Sentinel: + def __repr__(self): + return "" + + +_marker = _Sentinel() + + +# Breakpoint: https://github.com/python/cpython/pull/27342 +if sys.version_info >= (3, 10): + def _should_collect_from_parameters(t): + return isinstance( + t, (typing._GenericAlias, _types.GenericAlias, _types.UnionType) + ) +else: + def _should_collect_from_parameters(t): + return isinstance(t, (typing._GenericAlias, _types.GenericAlias)) + + +NoReturn = typing.NoReturn + +# Some unconstrained type variables. These are used by the container types. +# (These are not for export.) +T = typing.TypeVar('T') # Any type. +KT = typing.TypeVar('KT') # Key type. +VT = typing.TypeVar('VT') # Value type. +T_co = typing.TypeVar('T_co', covariant=True) # Any type covariant containers. +T_contra = typing.TypeVar('T_contra', contravariant=True) # Ditto contravariant. + + +# Breakpoint: https://github.com/python/cpython/pull/31841 +if sys.version_info >= (3, 11): + from typing import Any +else: + + class _AnyMeta(type): + def __instancecheck__(self, obj): + if self is Any: + raise TypeError("typing_extensions.Any cannot be used with isinstance()") + return super().__instancecheck__(obj) + + def __repr__(self): + if self is Any: + return "typing_extensions.Any" + return super().__repr__() + + class Any(metaclass=_AnyMeta): + """Special type indicating an unconstrained type. + - Any is compatible with every type. + - Any assumed to have all methods. + - All values assumed to be instances of Any. + Note that all the above statements are true from the point of view of + static type checkers. At runtime, Any should not be used with instance + checks. + """ + def __new__(cls, *args, **kwargs): + if cls is Any: + raise TypeError("Any cannot be instantiated") + return super().__new__(cls, *args, **kwargs) + + +ClassVar = typing.ClassVar + +# Vendored from cpython typing._SpecialFrom +# Having a separate class means that instances will not be rejected by +# typing._type_check. +class _SpecialForm(typing._Final, _root=True): + __slots__ = ('_name', '__doc__', '_getitem') + + def __init__(self, getitem): + self._getitem = getitem + self._name = getitem.__name__ + self.__doc__ = getitem.__doc__ + + def __getattr__(self, item): + if item in {'__name__', '__qualname__'}: + return self._name + + raise AttributeError(item) + + def __mro_entries__(self, bases): + raise TypeError(f"Cannot subclass {self!r}") + + def __repr__(self): + return f'typing_extensions.{self._name}' + + def __reduce__(self): + return self._name + + def __call__(self, *args, **kwds): + raise TypeError(f"Cannot instantiate {self!r}") + + def __or__(self, other): + return typing.Union[self, other] + + def __ror__(self, other): + return typing.Union[other, self] + + def __instancecheck__(self, obj): + raise TypeError(f"{self} cannot be used with isinstance()") + + def __subclasscheck__(self, cls): + raise TypeError(f"{self} cannot be used with issubclass()") + + @typing._tp_cache + def __getitem__(self, parameters): + return self._getitem(self, parameters) + + +# Note that inheriting from this class means that the object will be +# rejected by typing._type_check, so do not use it if the special form +# is arguably valid as a type by itself. +class _ExtensionsSpecialForm(typing._SpecialForm, _root=True): + def __repr__(self): + return 'typing_extensions.' + self._name + + +Final = typing.Final + +# Breakpoint: https://github.com/python/cpython/pull/30530 +if sys.version_info >= (3, 11): + final = typing.final +else: + # @final exists in 3.8+, but we backport it for all versions + # before 3.11 to keep support for the __final__ attribute. + # See https://bugs.python.org/issue46342 + def final(f): + """This decorator can be used to indicate to type checkers that + the decorated method cannot be overridden, and decorated class + cannot be subclassed. For example: + + class Base: + @final + def done(self) -> None: + ... + class Sub(Base): + def done(self) -> None: # Error reported by type checker + ... + @final + class Leaf: + ... + class Other(Leaf): # Error reported by type checker + ... + + There is no runtime checking of these properties. The decorator + sets the ``__final__`` attribute to ``True`` on the decorated object + to allow runtime introspection. + """ + try: + f.__final__ = True + except (AttributeError, TypeError): + # Skip the attribute silently if it is not writable. + # AttributeError happens if the object has __slots__ or a + # read-only property, TypeError if it's a builtin class. + pass + return f + + +if hasattr(typing, "disjoint_base"): # 3.15 + disjoint_base = typing.disjoint_base +else: + def disjoint_base(cls): + """This decorator marks a class as a disjoint base. + + Child classes of a disjoint base cannot inherit from other disjoint bases that are + not parent classes of the disjoint base. + + For example: + + @disjoint_base + class Disjoint1: pass + + @disjoint_base + class Disjoint2: pass + + class Disjoint3(Disjoint1, Disjoint2): pass # Type checker error + + Type checkers can use knowledge of disjoint bases to detect unreachable code + and determine when two types can overlap. + + See PEP 800.""" + cls.__disjoint_base__ = True + return cls + + +def IntVar(name): + return typing.TypeVar(name) + + +# A Literal bug was fixed in 3.11.0, 3.10.1 and 3.9.8 +# Breakpoint: https://github.com/python/cpython/pull/29334 +if sys.version_info >= (3, 10, 1): + Literal = typing.Literal +else: + def _flatten_literal_params(parameters): + """An internal helper for Literal creation: flatten Literals among parameters""" + params = [] + for p in parameters: + if isinstance(p, _LiteralGenericAlias): + params.extend(p.__args__) + else: + params.append(p) + return tuple(params) + + def _value_and_type_iter(params): + for p in params: + yield p, type(p) + + class _LiteralGenericAlias(typing._GenericAlias, _root=True): + def __eq__(self, other): + if not isinstance(other, _LiteralGenericAlias): + return NotImplemented + these_args_deduped = set(_value_and_type_iter(self.__args__)) + other_args_deduped = set(_value_and_type_iter(other.__args__)) + return these_args_deduped == other_args_deduped + + def __hash__(self): + return hash(frozenset(_value_and_type_iter(self.__args__))) + + class _LiteralForm(_ExtensionsSpecialForm, _root=True): + def __init__(self, doc: str): + self._name = 'Literal' + self._doc = self.__doc__ = doc + + def __getitem__(self, parameters): + if not isinstance(parameters, tuple): + parameters = (parameters,) + + parameters = _flatten_literal_params(parameters) + + val_type_pairs = list(_value_and_type_iter(parameters)) + try: + deduped_pairs = set(val_type_pairs) + except TypeError: + # unhashable parameters + pass + else: + # similar logic to typing._deduplicate on Python 3.9+ + if len(deduped_pairs) < len(val_type_pairs): + new_parameters = [] + for pair in val_type_pairs: + if pair in deduped_pairs: + new_parameters.append(pair[0]) + deduped_pairs.remove(pair) + assert not deduped_pairs, deduped_pairs + parameters = tuple(new_parameters) + + return _LiteralGenericAlias(self, parameters) + + Literal = _LiteralForm(doc="""\ + A type that can be used to indicate to type checkers + that the corresponding value has a value literally equivalent + to the provided parameter. For example: + + var: Literal[4] = 4 + + The type checker understands that 'var' is literally equal to + the value 4 and no other value. + + Literal[...] cannot be subclassed. There is no runtime + checking verifying that the parameter is actually a value + instead of a type.""") + + +_overload_dummy = typing._overload_dummy + + +if hasattr(typing, "get_overloads"): # 3.11+ + overload = typing.overload + get_overloads = typing.get_overloads + clear_overloads = typing.clear_overloads +else: + # {module: {qualname: {firstlineno: func}}} + _overload_registry = collections.defaultdict( + functools.partial(collections.defaultdict, dict) + ) + + def overload(func): + """Decorator for overloaded functions/methods. + + In a stub file, place two or more stub definitions for the same + function in a row, each decorated with @overload. For example: + + @overload + def utf8(value: None) -> None: ... + @overload + def utf8(value: bytes) -> bytes: ... + @overload + def utf8(value: str) -> bytes: ... + + In a non-stub file (i.e. a regular .py file), do the same but + follow it with an implementation. The implementation should *not* + be decorated with @overload. For example: + + @overload + def utf8(value: None) -> None: ... + @overload + def utf8(value: bytes) -> bytes: ... + @overload + def utf8(value: str) -> bytes: ... + def utf8(value): + # implementation goes here + + The overloads for a function can be retrieved at runtime using the + get_overloads() function. + """ + # classmethod and staticmethod + f = getattr(func, "__func__", func) + try: + _overload_registry[f.__module__][f.__qualname__][ + f.__code__.co_firstlineno + ] = func + except AttributeError: + # Not a normal function; ignore. + pass + return _overload_dummy + + def get_overloads(func): + """Return all defined overloads for *func* as a sequence.""" + # classmethod and staticmethod + f = getattr(func, "__func__", func) + if f.__module__ not in _overload_registry: + return [] + mod_dict = _overload_registry[f.__module__] + if f.__qualname__ not in mod_dict: + return [] + return list(mod_dict[f.__qualname__].values()) + + def clear_overloads(): + """Clear all overloads in the registry.""" + _overload_registry.clear() + + +# This is not a real generic class. Don't use outside annotations. +Type = typing.Type + +# Various ABCs mimicking those in collections.abc. +# A few are simply re-exported for completeness. +Awaitable = typing.Awaitable +Coroutine = typing.Coroutine +AsyncIterable = typing.AsyncIterable +AsyncIterator = typing.AsyncIterator +Deque = typing.Deque +DefaultDict = typing.DefaultDict +OrderedDict = typing.OrderedDict +Counter = typing.Counter +ChainMap = typing.ChainMap +Text = typing.Text +TYPE_CHECKING = typing.TYPE_CHECKING + + +# Breakpoint: https://github.com/python/cpython/pull/118681 +if sys.version_info >= (3, 13, 0, "beta"): + from typing import AsyncContextManager, AsyncGenerator, ContextManager, Generator +else: + def _is_dunder(attr): + return attr.startswith('__') and attr.endswith('__') + + + class _SpecialGenericAlias(typing._SpecialGenericAlias, _root=True): + def __init__(self, origin, nparams, *, inst=True, name=None, defaults=()): + super().__init__(origin, nparams, inst=inst, name=name) + self._defaults = defaults + + def __setattr__(self, attr, val): + allowed_attrs = {'_name', '_inst', '_nparams', '_defaults'} + if _is_dunder(attr) or attr in allowed_attrs: + object.__setattr__(self, attr, val) + else: + setattr(self.__origin__, attr, val) + + @typing._tp_cache + def __getitem__(self, params): + if not isinstance(params, tuple): + params = (params,) + msg = "Parameters to generic types must be types." + params = tuple(typing._type_check(p, msg) for p in params) + if ( + self._defaults + and len(params) < self._nparams + and len(params) + len(self._defaults) >= self._nparams + ): + params = (*params, *self._defaults[len(params) - self._nparams:]) + actual_len = len(params) + + if actual_len != self._nparams: + if self._defaults: + expected = f"at least {self._nparams - len(self._defaults)}" + else: + expected = str(self._nparams) + if not self._nparams: + raise TypeError(f"{self} is not a generic class") + raise TypeError( + f"Too {'many' if actual_len > self._nparams else 'few'}" + f" arguments for {self};" + f" actual {actual_len}, expected {expected}" + ) + return self.copy_with(params) + + _NoneType = type(None) + Generator = _SpecialGenericAlias( + collections.abc.Generator, 3, defaults=(_NoneType, _NoneType) + ) + AsyncGenerator = _SpecialGenericAlias( + collections.abc.AsyncGenerator, 2, defaults=(_NoneType,) + ) + ContextManager = _SpecialGenericAlias( + contextlib.AbstractContextManager, + 2, + name="ContextManager", + defaults=(typing.Optional[bool],) + ) + AsyncContextManager = _SpecialGenericAlias( + contextlib.AbstractAsyncContextManager, + 2, + name="AsyncContextManager", + defaults=(typing.Optional[bool],) + ) + + +_PROTO_ALLOWLIST = { + 'collections.abc': [ + 'Callable', 'Awaitable', 'Iterable', 'Iterator', 'AsyncIterable', + 'Hashable', 'Sized', 'Container', 'Collection', 'Reversible', 'Buffer', + ], + 'contextlib': ['AbstractContextManager', 'AbstractAsyncContextManager'], + 'typing_extensions': ['Buffer'], +} + + +_EXCLUDED_ATTRS = frozenset(typing.EXCLUDED_ATTRIBUTES) | { + "__match_args__", "__protocol_attrs__", "__non_callable_proto_members__", + "__final__", +} + + +def _get_protocol_attrs(cls): + attrs = set() + for base in cls.__mro__[:-1]: # without object + if base.__name__ in {'Protocol', 'Generic'}: + continue + annotations = getattr(base, '__annotations__', {}) + for attr in (*base.__dict__, *annotations): + if (not attr.startswith('_abc_') and attr not in _EXCLUDED_ATTRS): + attrs.add(attr) + return attrs + + +def _caller(depth=1, default='__main__'): + try: + return sys._getframemodulename(depth + 1) or default + except AttributeError: # For platforms without _getframemodulename() + pass + try: + return sys._getframe(depth + 1).f_globals.get('__name__', default) + except (AttributeError, ValueError): # For platforms without _getframe() + pass + return None + + +# `__match_args__` attribute was removed from protocol members in 3.13, +# we want to backport this change to older Python versions. +# Breakpoint: https://github.com/python/cpython/pull/110683 +if sys.version_info >= (3, 13): + Protocol = typing.Protocol +else: + def _allow_reckless_class_checks(depth=2): + """Allow instance and class checks for special stdlib modules. + The abc and functools modules indiscriminately call isinstance() and + issubclass() on the whole MRO of a user class, which may contain protocols. + """ + return _caller(depth) in {'abc', 'functools', None} + + def _no_init(self, *args, **kwargs): + if type(self)._is_protocol: + raise TypeError('Protocols cannot be instantiated') + + def _type_check_issubclass_arg_1(arg): + """Raise TypeError if `arg` is not an instance of `type` + in `issubclass(arg, )`. + + In most cases, this is verified by type.__subclasscheck__. + Checking it again unnecessarily would slow down issubclass() checks, + so, we don't perform this check unless we absolutely have to. + + For various error paths, however, + we want to ensure that *this* error message is shown to the user + where relevant, rather than a typing.py-specific error message. + """ + if not isinstance(arg, type): + # Same error message as for issubclass(1, int). + raise TypeError('issubclass() arg 1 must be a class') + + # Inheriting from typing._ProtocolMeta isn't actually desirable, + # but is necessary to allow typing.Protocol and typing_extensions.Protocol + # to mix without getting TypeErrors about "metaclass conflict" + class _ProtocolMeta(type(typing.Protocol)): + # This metaclass is somewhat unfortunate, + # but is necessary for several reasons... + # + # NOTE: DO NOT call super() in any methods in this class + # That would call the methods on typing._ProtocolMeta on Python <=3.11 + # and those are slow + def __new__(mcls, name, bases, namespace, **kwargs): + if name == "Protocol" and len(bases) < 2: + pass + elif {Protocol, typing.Protocol} & set(bases): + for base in bases: + if not ( + base in {object, typing.Generic, Protocol, typing.Protocol} + or base.__name__ in _PROTO_ALLOWLIST.get(base.__module__, []) + or is_protocol(base) + ): + raise TypeError( + f"Protocols can only inherit from other protocols, " + f"got {base!r}" + ) + return abc.ABCMeta.__new__(mcls, name, bases, namespace, **kwargs) + + def __init__(cls, *args, **kwargs): + abc.ABCMeta.__init__(cls, *args, **kwargs) + if getattr(cls, "_is_protocol", False): + cls.__protocol_attrs__ = _get_protocol_attrs(cls) + + def __subclasscheck__(cls, other): + if cls is Protocol: + return type.__subclasscheck__(cls, other) + if ( + getattr(cls, '_is_protocol', False) + and not _allow_reckless_class_checks() + ): + if not getattr(cls, '_is_runtime_protocol', False): + _type_check_issubclass_arg_1(other) + raise TypeError( + "Instance and class checks can only be used with " + "@runtime_checkable protocols" + ) + if ( + # this attribute is set by @runtime_checkable: + cls.__non_callable_proto_members__ + and cls.__dict__.get("__subclasshook__") is _proto_hook + ): + _type_check_issubclass_arg_1(other) + non_method_attrs = sorted(cls.__non_callable_proto_members__) + raise TypeError( + "Protocols with non-method members don't support issubclass()." + f" Non-method members: {str(non_method_attrs)[1:-1]}." + ) + return abc.ABCMeta.__subclasscheck__(cls, other) + + def __instancecheck__(cls, instance): + # We need this method for situations where attributes are + # assigned in __init__. + if cls is Protocol: + return type.__instancecheck__(cls, instance) + if not getattr(cls, "_is_protocol", False): + # i.e., it's a concrete subclass of a protocol + return abc.ABCMeta.__instancecheck__(cls, instance) + + if ( + not getattr(cls, '_is_runtime_protocol', False) and + not _allow_reckless_class_checks() + ): + raise TypeError("Instance and class checks can only be used with" + " @runtime_checkable protocols") + + if abc.ABCMeta.__instancecheck__(cls, instance): + return True + + for attr in cls.__protocol_attrs__: + try: + val = inspect.getattr_static(instance, attr) + except AttributeError: + break + # this attribute is set by @runtime_checkable: + if val is None and attr not in cls.__non_callable_proto_members__: + break + else: + return True + + return False + + def __eq__(cls, other): + # Hack so that typing.Generic.__class_getitem__ + # treats typing_extensions.Protocol + # as equivalent to typing.Protocol + if abc.ABCMeta.__eq__(cls, other) is True: + return True + return cls is Protocol and other is typing.Protocol + + # This has to be defined, or the abc-module cache + # complains about classes with this metaclass being unhashable, + # if we define only __eq__! + def __hash__(cls) -> int: + return type.__hash__(cls) + + @classmethod + def _proto_hook(cls, other): + if not cls.__dict__.get('_is_protocol', False): + return NotImplemented + + for attr in cls.__protocol_attrs__: + for base in other.__mro__: + # Check if the members appears in the class dictionary... + if attr in base.__dict__: + if base.__dict__[attr] is None: + return NotImplemented + break + + # ...or in annotations, if it is a sub-protocol. + annotations = getattr(base, '__annotations__', {}) + if ( + isinstance(annotations, collections.abc.Mapping) + and attr in annotations + and is_protocol(other) + ): + break + else: + return NotImplemented + return True + + class Protocol(typing.Generic, metaclass=_ProtocolMeta): + __doc__ = typing.Protocol.__doc__ + __slots__ = () + _is_protocol = True + _is_runtime_protocol = False + + def __init_subclass__(cls, *args, **kwargs): + super().__init_subclass__(*args, **kwargs) + + # Determine if this is a protocol or a concrete subclass. + if not cls.__dict__.get('_is_protocol', False): + cls._is_protocol = any(b is Protocol for b in cls.__bases__) + + # Set (or override) the protocol subclass hook. + if '__subclasshook__' not in cls.__dict__: + cls.__subclasshook__ = _proto_hook + + # Prohibit instantiation for protocol classes + if cls._is_protocol and cls.__init__ is Protocol.__init__: + cls.__init__ = _no_init + + +# Breakpoint: https://github.com/python/cpython/pull/113401 +if sys.version_info >= (3, 13): + runtime_checkable = typing.runtime_checkable +else: + def runtime_checkable(cls): + """Mark a protocol class as a runtime protocol. + + Such protocol can be used with isinstance() and issubclass(). + Raise TypeError if applied to a non-protocol class. + This allows a simple-minded structural check very similar to + one trick ponies in collections.abc such as Iterable. + + For example:: + + @runtime_checkable + class Closable(Protocol): + def close(self): ... + + assert isinstance(open('/some/file'), Closable) + + Warning: this will check only the presence of the required methods, + not their type signatures! + """ + if not issubclass(cls, typing.Generic) or not getattr(cls, '_is_protocol', False): + raise TypeError(f'@runtime_checkable can be only applied to protocol classes,' + f' got {cls!r}') + cls._is_runtime_protocol = True + + # typing.Protocol classes on <=3.11 break if we execute this block, + # because typing.Protocol classes on <=3.11 don't have a + # `__protocol_attrs__` attribute, and this block relies on the + # `__protocol_attrs__` attribute. Meanwhile, typing.Protocol classes on 3.12.2+ + # break if we *don't* execute this block, because *they* assume that all + # protocol classes have a `__non_callable_proto_members__` attribute + # (which this block sets) + if isinstance(cls, _ProtocolMeta) or sys.version_info >= (3, 12, 2): + # PEP 544 prohibits using issubclass() + # with protocols that have non-method members. + # See gh-113320 for why we compute this attribute here, + # rather than in `_ProtocolMeta.__init__` + cls.__non_callable_proto_members__ = set() + for attr in cls.__protocol_attrs__: + try: + is_callable = callable(getattr(cls, attr, None)) + except Exception as e: + raise TypeError( + f"Failed to determine whether protocol member {attr!r} " + "is a method member" + ) from e + else: + if not is_callable: + cls.__non_callable_proto_members__.add(attr) + + return cls + + +# The "runtime" alias exists for backwards compatibility. +runtime = runtime_checkable + + +# Our version of runtime-checkable protocols is faster on Python <=3.11 +# Breakpoint: https://github.com/python/cpython/pull/112717 +if sys.version_info >= (3, 12): + SupportsInt = typing.SupportsInt + SupportsFloat = typing.SupportsFloat + SupportsComplex = typing.SupportsComplex + SupportsBytes = typing.SupportsBytes + SupportsIndex = typing.SupportsIndex + SupportsAbs = typing.SupportsAbs + SupportsRound = typing.SupportsRound +else: + @runtime_checkable + class SupportsInt(Protocol): + """An ABC with one abstract method __int__.""" + __slots__ = () + + @abc.abstractmethod + def __int__(self) -> int: + pass + + @runtime_checkable + class SupportsFloat(Protocol): + """An ABC with one abstract method __float__.""" + __slots__ = () + + @abc.abstractmethod + def __float__(self) -> float: + pass + + @runtime_checkable + class SupportsComplex(Protocol): + """An ABC with one abstract method __complex__.""" + __slots__ = () + + @abc.abstractmethod + def __complex__(self) -> complex: + pass + + @runtime_checkable + class SupportsBytes(Protocol): + """An ABC with one abstract method __bytes__.""" + __slots__ = () + + @abc.abstractmethod + def __bytes__(self) -> bytes: + pass + + @runtime_checkable + class SupportsIndex(Protocol): + __slots__ = () + + @abc.abstractmethod + def __index__(self) -> int: + pass + + @runtime_checkable + class SupportsAbs(Protocol[T_co]): + """ + An ABC with one abstract method __abs__ that is covariant in its return type. + """ + __slots__ = () + + @abc.abstractmethod + def __abs__(self) -> T_co: + pass + + @runtime_checkable + class SupportsRound(Protocol[T_co]): + """ + An ABC with one abstract method __round__ that is covariant in its return type. + """ + __slots__ = () + + @abc.abstractmethod + def __round__(self, ndigits: int = 0) -> T_co: + pass + + +if hasattr(io, "Reader") and hasattr(io, "Writer"): + Reader = io.Reader + Writer = io.Writer +else: + @runtime_checkable + class Reader(Protocol[T_co]): + """Protocol for simple I/O reader instances. + + This protocol only supports blocking I/O. + """ + + __slots__ = () + + @abc.abstractmethod + def read(self, size: int = ..., /) -> T_co: + """Read data from the input stream and return it. + + If *size* is specified, at most *size* items (bytes/characters) will be + read. + """ + + @runtime_checkable + class Writer(Protocol[T_contra]): + """Protocol for simple I/O writer instances. + + This protocol only supports blocking I/O. + """ + + __slots__ = () + + @abc.abstractmethod + def write(self, data: T_contra, /) -> int: + """Write *data* to the output stream and return the number of items written.""" # noqa: E501 + + +_NEEDS_SINGLETONMETA = ( + not hasattr(typing, "NoDefault") or not hasattr(typing, "NoExtraItems") +) + +if _NEEDS_SINGLETONMETA: + class SingletonMeta(type): + def __setattr__(cls, attr, value): + # TypeError is consistent with the behavior of NoneType + raise TypeError( + f"cannot set {attr!r} attribute of immutable type {cls.__name__!r}" + ) + + +if hasattr(typing, "NoDefault"): + NoDefault = typing.NoDefault +else: + class NoDefaultType(metaclass=SingletonMeta): + """The type of the NoDefault singleton.""" + + __slots__ = () + + def __new__(cls): + return globals().get("NoDefault") or object.__new__(cls) + + def __repr__(self): + return "typing_extensions.NoDefault" + + def __reduce__(self): + return "NoDefault" + + NoDefault = NoDefaultType() + del NoDefaultType + +if hasattr(typing, "NoExtraItems"): + NoExtraItems = typing.NoExtraItems +else: + class NoExtraItemsType(metaclass=SingletonMeta): + """The type of the NoExtraItems singleton.""" + + __slots__ = () + + def __new__(cls): + return globals().get("NoExtraItems") or object.__new__(cls) + + def __repr__(self): + return "typing_extensions.NoExtraItems" + + def __reduce__(self): + return "NoExtraItems" + + NoExtraItems = NoExtraItemsType() + del NoExtraItemsType + +if _NEEDS_SINGLETONMETA: + del SingletonMeta + + +# Update this to something like >=3.13.0b1 if and when +# PEP 728 is implemented in CPython +_PEP_728_IMPLEMENTED = False + +if _PEP_728_IMPLEMENTED: + # The standard library TypedDict in Python 3.9.0/1 does not honour the "total" + # keyword with old-style TypedDict(). See https://bugs.python.org/issue42059 + # The standard library TypedDict below Python 3.11 does not store runtime + # information about optional and required keys when using Required or NotRequired. + # Generic TypedDicts are also impossible using typing.TypedDict on Python <3.11. + # Aaaand on 3.12 we add __orig_bases__ to TypedDict + # to enable better runtime introspection. + # On 3.13 we deprecate some odd ways of creating TypedDicts. + # Also on 3.13, PEP 705 adds the ReadOnly[] qualifier. + # PEP 728 (still pending) makes more changes. + TypedDict = typing.TypedDict + _TypedDictMeta = typing._TypedDictMeta + is_typeddict = typing.is_typeddict +else: + # 3.10.0 and later + _TAKES_MODULE = "module" in inspect.signature(typing._type_check).parameters + + def _get_typeddict_qualifiers(annotation_type): + while True: + annotation_origin = get_origin(annotation_type) + if annotation_origin is Annotated: + annotation_args = get_args(annotation_type) + if annotation_args: + annotation_type = annotation_args[0] + else: + break + elif annotation_origin is Required: + yield Required + annotation_type, = get_args(annotation_type) + elif annotation_origin is NotRequired: + yield NotRequired + annotation_type, = get_args(annotation_type) + elif annotation_origin is ReadOnly: + yield ReadOnly + annotation_type, = get_args(annotation_type) + else: + break + + class _TypedDictMeta(type): + + def __new__(cls, name, bases, ns, *, total=True, closed=None, + extra_items=NoExtraItems): + """Create new typed dict class object. + + This method is called when TypedDict is subclassed, + or when TypedDict is instantiated. This way + TypedDict supports all three syntax forms described in its docstring. + Subclasses and instances of TypedDict return actual dictionaries. + """ + for base in bases: + if type(base) is not _TypedDictMeta and base is not typing.Generic: + raise TypeError('cannot inherit from both a TypedDict type ' + 'and a non-TypedDict base class') + if closed is not None and extra_items is not NoExtraItems: + raise TypeError(f"Cannot combine closed={closed!r} and extra_items") + + if any(issubclass(b, typing.Generic) for b in bases): + generic_base = (typing.Generic,) + else: + generic_base = () + + ns_annotations = ns.pop('__annotations__', None) + + # typing.py generally doesn't let you inherit from plain Generic, unless + # the name of the class happens to be "Protocol" + tp_dict = type.__new__(_TypedDictMeta, "Protocol", (*generic_base, dict), ns) + tp_dict.__name__ = name + if tp_dict.__qualname__ == "Protocol": + tp_dict.__qualname__ = name + + if not hasattr(tp_dict, '__orig_bases__'): + tp_dict.__orig_bases__ = bases + + annotations = {} + own_annotate = None + if ns_annotations is not None: + own_annotations = ns_annotations + elif sys.version_info >= (3, 14): + if hasattr(annotationlib, "get_annotate_from_class_namespace"): + own_annotate = annotationlib.get_annotate_from_class_namespace(ns) + else: + # 3.14.0a7 and earlier + own_annotate = ns.get("__annotate__") + if own_annotate is not None: + own_annotations = annotationlib.call_annotate_function( + own_annotate, Format.FORWARDREF, owner=tp_dict + ) + else: + own_annotations = {} + else: + own_annotations = {} + msg = "TypedDict('Name', {f0: t0, f1: t1, ...}); each t must be a type" + if _TAKES_MODULE: + own_checked_annotations = { + n: typing._type_check(tp, msg, module=tp_dict.__module__) + for n, tp in own_annotations.items() + } + else: + own_checked_annotations = { + n: typing._type_check(tp, msg) + for n, tp in own_annotations.items() + } + required_keys = set() + optional_keys = set() + readonly_keys = set() + mutable_keys = set() + extra_items_type = extra_items + + for base in bases: + base_dict = base.__dict__ + + if sys.version_info <= (3, 14): + annotations.update(base_dict.get('__annotations__', {})) + required_keys.update(base_dict.get('__required_keys__', ())) + optional_keys.update(base_dict.get('__optional_keys__', ())) + readonly_keys.update(base_dict.get('__readonly_keys__', ())) + mutable_keys.update(base_dict.get('__mutable_keys__', ())) + + # This was specified in an earlier version of PEP 728. Support + # is retained for backwards compatibility, but only for Python + # 3.13 and lower. + if (closed and sys.version_info < (3, 14) + and "__extra_items__" in own_checked_annotations): + annotation_type = own_checked_annotations.pop("__extra_items__") + qualifiers = set(_get_typeddict_qualifiers(annotation_type)) + if Required in qualifiers: + raise TypeError( + "Special key __extra_items__ does not support " + "Required" + ) + if NotRequired in qualifiers: + raise TypeError( + "Special key __extra_items__ does not support " + "NotRequired" + ) + extra_items_type = annotation_type + + annotations.update(own_checked_annotations) + for annotation_key, annotation_type in own_checked_annotations.items(): + qualifiers = set(_get_typeddict_qualifiers(annotation_type)) + + if Required in qualifiers: + required_keys.add(annotation_key) + elif NotRequired in qualifiers: + optional_keys.add(annotation_key) + elif total: + required_keys.add(annotation_key) + else: + optional_keys.add(annotation_key) + if ReadOnly in qualifiers: + mutable_keys.discard(annotation_key) + readonly_keys.add(annotation_key) + else: + mutable_keys.add(annotation_key) + readonly_keys.discard(annotation_key) + + # Breakpoint: https://github.com/python/cpython/pull/119891 + if sys.version_info >= (3, 14): + def __annotate__(format): + annos = {} + for base in bases: + if base is Generic: + continue + base_annotate = base.__annotate__ + if base_annotate is None: + continue + base_annos = annotationlib.call_annotate_function( + base_annotate, format, owner=base) + annos.update(base_annos) + if own_annotate is not None: + own = annotationlib.call_annotate_function( + own_annotate, format, owner=tp_dict) + if format != Format.STRING: + own = { + n: typing._type_check(tp, msg, module=tp_dict.__module__) + for n, tp in own.items() + } + elif format == Format.STRING: + own = annotationlib.annotations_to_string(own_annotations) + elif format in (Format.FORWARDREF, Format.VALUE): + own = own_checked_annotations + else: + raise NotImplementedError(format) + annos.update(own) + return annos + + tp_dict.__annotate__ = __annotate__ + else: + tp_dict.__annotations__ = annotations + tp_dict.__required_keys__ = frozenset(required_keys) + tp_dict.__optional_keys__ = frozenset(optional_keys) + tp_dict.__readonly_keys__ = frozenset(readonly_keys) + tp_dict.__mutable_keys__ = frozenset(mutable_keys) + tp_dict.__total__ = total + tp_dict.__closed__ = closed + tp_dict.__extra_items__ = extra_items_type + return tp_dict + + __call__ = dict # static method + + def __subclasscheck__(cls, other): + # Typed dicts are only for static structural subtyping. + raise TypeError('TypedDict does not support instance and class checks') + + __instancecheck__ = __subclasscheck__ + + _TypedDict = type.__new__(_TypedDictMeta, 'TypedDict', (), {}) + + def _create_typeddict( + typename, + fields, + /, + *, + typing_is_inline, + total, + closed, + extra_items, + **kwargs, + ): + if fields is _marker or fields is None: + if fields is _marker: + deprecated_thing = ( + "Failing to pass a value for the 'fields' parameter" + ) + else: + deprecated_thing = "Passing `None` as the 'fields' parameter" + + example = f"`{typename} = TypedDict({typename!r}, {{}})`" + deprecation_msg = ( + f"{deprecated_thing} is deprecated and will be disallowed in " + "Python 3.15. To create a TypedDict class with 0 fields " + "using the functional syntax, pass an empty dictionary, e.g. " + ) + example + "." + warnings.warn(deprecation_msg, DeprecationWarning, stacklevel=2) + # Support a field called "closed" + if closed is not False and closed is not True and closed is not None: + kwargs["closed"] = closed + closed = None + # Or "extra_items" + if extra_items is not NoExtraItems: + kwargs["extra_items"] = extra_items + extra_items = NoExtraItems + fields = kwargs + elif kwargs: + raise TypeError("TypedDict takes either a dict or keyword arguments," + " but not both") + if kwargs: + # Breakpoint: https://github.com/python/cpython/pull/104891 + if sys.version_info >= (3, 13): + raise TypeError("TypedDict takes no keyword arguments") + warnings.warn( + "The kwargs-based syntax for TypedDict definitions is deprecated " + "in Python 3.11, will be removed in Python 3.13, and may not be " + "understood by third-party type checkers.", + DeprecationWarning, + stacklevel=2, + ) + + ns = {'__annotations__': dict(fields)} + module = _caller(depth=4 if typing_is_inline else 2) + if module is not None: + # Setting correct module is necessary to make typed dict classes + # pickleable. + ns['__module__'] = module + + td = _TypedDictMeta(typename, (), ns, total=total, closed=closed, + extra_items=extra_items) + td.__orig_bases__ = (TypedDict,) + return td + + class _TypedDictSpecialForm(_SpecialForm, _root=True): + def __call__( + self, + typename, + fields=_marker, + /, + *, + total=True, + closed=None, + extra_items=NoExtraItems, + **kwargs + ): + return _create_typeddict( + typename, + fields, + typing_is_inline=False, + total=total, + closed=closed, + extra_items=extra_items, + **kwargs, + ) + + def __mro_entries__(self, bases): + return (_TypedDict,) + + @_TypedDictSpecialForm + def TypedDict(self, args): + """A simple typed namespace. At runtime it is equivalent to a plain dict. + + TypedDict creates a dictionary type such that a type checker will expect all + instances to have a certain set of keys, where each key is + associated with a value of a consistent type. This expectation + is not checked at runtime. + + Usage:: + + class Point2D(TypedDict): + x: int + y: int + label: str + + a: Point2D = {'x': 1, 'y': 2, 'label': 'good'} # OK + b: Point2D = {'z': 3, 'label': 'bad'} # Fails type check + + assert Point2D(x=1, y=2, label='first') == dict(x=1, y=2, label='first') + + The type info can be accessed via the Point2D.__annotations__ dict, and + the Point2D.__required_keys__ and Point2D.__optional_keys__ frozensets. + TypedDict supports an additional equivalent form:: + + Point2D = TypedDict('Point2D', {'x': int, 'y': int, 'label': str}) + + By default, all keys must be present in a TypedDict. It is possible + to override this by specifying totality:: + + class Point2D(TypedDict, total=False): + x: int + y: int + + This means that a Point2D TypedDict can have any of the keys omitted. A type + checker is only expected to support a literal False or True as the value of + the total argument. True is the default, and makes all items defined in the + class body be required. + + The Required and NotRequired special forms can also be used to mark + individual keys as being required or not required:: + + class Point2D(TypedDict): + x: int # the "x" key must always be present (Required is the default) + y: NotRequired[int] # the "y" key can be omitted + + See PEP 655 for more details on Required and NotRequired. + """ + # This runs when creating inline TypedDicts: + if not isinstance(args, dict): + raise TypeError( + "TypedDict[...] should be used with a single dict argument" + ) + + return _create_typeddict( + "", + args, + typing_is_inline=True, + total=True, + closed=True, + extra_items=NoExtraItems, + ) + + _TYPEDDICT_TYPES = (typing._TypedDictMeta, _TypedDictMeta) + + def is_typeddict(tp): + """Check if an annotation is a TypedDict class + + For example:: + class Film(TypedDict): + title: str + year: int + + is_typeddict(Film) # => True + is_typeddict(Union[list, str]) # => False + """ + return isinstance(tp, _TYPEDDICT_TYPES) + + +if hasattr(typing, "assert_type"): + assert_type = typing.assert_type + +else: + def assert_type(val, typ, /): + """Assert (to the type checker) that the value is of the given type. + + When the type checker encounters a call to assert_type(), it + emits an error if the value is not of the specified type:: + + def greet(name: str) -> None: + assert_type(name, str) # ok + assert_type(name, int) # type checker error + + At runtime this returns the first argument unchanged and otherwise + does nothing. + """ + return val + + +if hasattr(typing, "ReadOnly"): # 3.13+ + get_type_hints = typing.get_type_hints +else: # <=3.13 + # replaces _strip_annotations() + def _strip_extras(t): + """Strips Annotated, Required and NotRequired from a given type.""" + if isinstance(t, typing._AnnotatedAlias): + return _strip_extras(t.__origin__) + if hasattr(t, "__origin__") and t.__origin__ in (Required, NotRequired, ReadOnly): + return _strip_extras(t.__args__[0]) + if isinstance(t, typing._GenericAlias): + stripped_args = tuple(_strip_extras(a) for a in t.__args__) + if stripped_args == t.__args__: + return t + return t.copy_with(stripped_args) + if hasattr(_types, "GenericAlias") and isinstance(t, _types.GenericAlias): + stripped_args = tuple(_strip_extras(a) for a in t.__args__) + if stripped_args == t.__args__: + return t + return _types.GenericAlias(t.__origin__, stripped_args) + if hasattr(_types, "UnionType") and isinstance(t, _types.UnionType): + stripped_args = tuple(_strip_extras(a) for a in t.__args__) + if stripped_args == t.__args__: + return t + return functools.reduce(operator.or_, stripped_args) + + return t + + def get_type_hints(obj, globalns=None, localns=None, include_extras=False): + """Return type hints for an object. + + This is often the same as obj.__annotations__, but it handles + forward references encoded as string literals, adds Optional[t] if a + default value equal to None is set and recursively replaces all + 'Annotated[T, ...]', 'Required[T]' or 'NotRequired[T]' with 'T' + (unless 'include_extras=True'). + + The argument may be a module, class, method, or function. The annotations + are returned as a dictionary. For classes, annotations include also + inherited members. + + TypeError is raised if the argument is not of a type that can contain + annotations, and an empty dictionary is returned if no annotations are + present. + + BEWARE -- the behavior of globalns and localns is counterintuitive + (unless you are familiar with how eval() and exec() work). The + search order is locals first, then globals. + + - If no dict arguments are passed, an attempt is made to use the + globals from obj (or the respective module's globals for classes), + and these are also used as the locals. If the object does not appear + to have globals, an empty dictionary is used. + + - If one dict argument is passed, it is used for both globals and + locals. + + - If two dict arguments are passed, they specify globals and + locals, respectively. + """ + hint = typing.get_type_hints( + obj, globalns=globalns, localns=localns, include_extras=True + ) + # Breakpoint: https://github.com/python/cpython/pull/30304 + if sys.version_info < (3, 11): + _clean_optional(obj, hint, globalns, localns) + if include_extras: + return hint + return {k: _strip_extras(t) for k, t in hint.items()} + + _NoneType = type(None) + + def _could_be_inserted_optional(t): + """detects Union[..., None] pattern""" + if not isinstance(t, typing._UnionGenericAlias): + return False + # Assume if last argument is not None they are user defined + if t.__args__[-1] is not _NoneType: + return False + return True + + # < 3.11 + def _clean_optional(obj, hints, globalns=None, localns=None): + # reverts injected Union[..., None] cases from typing.get_type_hints + # when a None default value is used. + # see https://github.com/python/typing_extensions/issues/310 + if not hints or isinstance(obj, type): + return + defaults = typing._get_defaults(obj) # avoid accessing __annotations___ + if not defaults: + return + original_hints = obj.__annotations__ + for name, value in hints.items(): + # Not a Union[..., None] or replacement conditions not fullfilled + if (not _could_be_inserted_optional(value) + or name not in defaults + or defaults[name] is not None + ): + continue + original_value = original_hints[name] + # value=NoneType should have caused a skip above but check for safety + if original_value is None: + original_value = _NoneType + # Forward reference + if isinstance(original_value, str): + if globalns is None: + if isinstance(obj, _types.ModuleType): + globalns = obj.__dict__ + else: + nsobj = obj + # Find globalns for the unwrapped object. + while hasattr(nsobj, '__wrapped__'): + nsobj = nsobj.__wrapped__ + globalns = getattr(nsobj, '__globals__', {}) + if localns is None: + localns = globalns + elif localns is None: + localns = globalns + + original_value = ForwardRef( + original_value, + is_argument=not isinstance(obj, _types.ModuleType) + ) + original_evaluated = typing._eval_type(original_value, globalns, localns) + # Compare if values differ. Note that even if equal + # value might be cached by typing._tp_cache contrary to original_evaluated + if original_evaluated != value or ( + # 3.10: ForwardRefs of UnionType might be turned into _UnionGenericAlias + hasattr(_types, "UnionType") + and isinstance(original_evaluated, _types.UnionType) + and not isinstance(value, _types.UnionType) + ): + hints[name] = original_evaluated + +# Python 3.9 has get_origin() and get_args() but those implementations don't support +# ParamSpecArgs and ParamSpecKwargs, so only Python 3.10's versions will do. +# Breakpoint: https://github.com/python/cpython/pull/25298 +if sys.version_info >= (3, 10): + get_origin = typing.get_origin + get_args = typing.get_args +# 3.9 +else: + def get_origin(tp): + """Get the unsubscripted version of a type. + + This supports generic types, Callable, Tuple, Union, Literal, Final, ClassVar + and Annotated. Return None for unsupported types. Examples:: + + get_origin(Literal[42]) is Literal + get_origin(int) is None + get_origin(ClassVar[int]) is ClassVar + get_origin(Generic) is Generic + get_origin(Generic[T]) is Generic + get_origin(Union[T, int]) is Union + get_origin(List[Tuple[T, T]][int]) == list + get_origin(P.args) is P + """ + if isinstance(tp, typing._AnnotatedAlias): + return Annotated + if isinstance(tp, (typing._BaseGenericAlias, _types.GenericAlias, + ParamSpecArgs, ParamSpecKwargs)): + return tp.__origin__ + if tp is typing.Generic: + return typing.Generic + return None + + def get_args(tp): + """Get type arguments with all substitutions performed. + + For unions, basic simplifications used by Union constructor are performed. + Examples:: + get_args(Dict[str, int]) == (str, int) + get_args(int) == () + get_args(Union[int, Union[T, int], str][int]) == (int, str) + get_args(Union[int, Tuple[T, int]][str]) == (int, Tuple[str, int]) + get_args(Callable[[], T][int]) == ([], int) + """ + if isinstance(tp, typing._AnnotatedAlias): + return (tp.__origin__, *tp.__metadata__) + if isinstance(tp, (typing._GenericAlias, _types.GenericAlias)): + res = tp.__args__ + if get_origin(tp) is collections.abc.Callable and res[0] is not Ellipsis: + res = (list(res[:-1]), res[-1]) + return res + return () + + +# 3.10+ +if hasattr(typing, 'TypeAlias'): + TypeAlias = typing.TypeAlias +# 3.9 +else: + @_ExtensionsSpecialForm + def TypeAlias(self, parameters): + """Special marker indicating that an assignment should + be recognized as a proper type alias definition by type + checkers. + + For example:: + + Predicate: TypeAlias = Callable[..., bool] + + It's invalid when used anywhere except as in the example above. + """ + raise TypeError(f"{self} is not subscriptable") + + +def _set_default(type_param, default): + type_param.has_default = lambda: default is not NoDefault + type_param.__default__ = default + + +def _set_module(typevarlike): + # for pickling: + def_mod = _caller(depth=2) + if def_mod != 'typing_extensions': + typevarlike.__module__ = def_mod + + +class _DefaultMixin: + """Mixin for TypeVarLike defaults.""" + + __slots__ = () + __init__ = _set_default + + +# Classes using this metaclass must provide a _backported_typevarlike ClassVar +class _TypeVarLikeMeta(type): + def __instancecheck__(cls, __instance: Any) -> bool: + return isinstance(__instance, cls._backported_typevarlike) + + +if _PEP_696_IMPLEMENTED: + from typing import TypeVar +else: + # Add default and infer_variance parameters from PEP 696 and 695 + class TypeVar(metaclass=_TypeVarLikeMeta): + """Type variable.""" + + _backported_typevarlike = typing.TypeVar + + def __new__(cls, name, *constraints, bound=None, + covariant=False, contravariant=False, + default=NoDefault, infer_variance=False): + if hasattr(typing, "TypeAliasType"): + # PEP 695 implemented (3.12+), can pass infer_variance to typing.TypeVar + typevar = typing.TypeVar(name, *constraints, bound=bound, + covariant=covariant, contravariant=contravariant, + infer_variance=infer_variance) + else: + typevar = typing.TypeVar(name, *constraints, bound=bound, + covariant=covariant, contravariant=contravariant) + if infer_variance and (covariant or contravariant): + raise ValueError("Variance cannot be specified with infer_variance.") + typevar.__infer_variance__ = infer_variance + + _set_default(typevar, default) + _set_module(typevar) + + def _tvar_prepare_subst(alias, args): + if ( + typevar.has_default() + and alias.__parameters__.index(typevar) == len(args) + ): + args += (typevar.__default__,) + return args + + typevar.__typing_prepare_subst__ = _tvar_prepare_subst + return typevar + + def __init_subclass__(cls) -> None: + raise TypeError(f"type '{__name__}.TypeVar' is not an acceptable base type") + + +# Python 3.10+ has PEP 612 +if hasattr(typing, 'ParamSpecArgs'): + ParamSpecArgs = typing.ParamSpecArgs + ParamSpecKwargs = typing.ParamSpecKwargs +# 3.9 +else: + class _Immutable: + """Mixin to indicate that object should not be copied.""" + __slots__ = () + + def __copy__(self): + return self + + def __deepcopy__(self, memo): + return self + + class ParamSpecArgs(_Immutable): + """The args for a ParamSpec object. + + Given a ParamSpec object P, P.args is an instance of ParamSpecArgs. + + ParamSpecArgs objects have a reference back to their ParamSpec: + + P.args.__origin__ is P + + This type is meant for runtime introspection and has no special meaning to + static type checkers. + """ + def __init__(self, origin): + self.__origin__ = origin + + def __repr__(self): + return f"{self.__origin__.__name__}.args" + + def __eq__(self, other): + if not isinstance(other, ParamSpecArgs): + return NotImplemented + return self.__origin__ == other.__origin__ + + class ParamSpecKwargs(_Immutable): + """The kwargs for a ParamSpec object. + + Given a ParamSpec object P, P.kwargs is an instance of ParamSpecKwargs. + + ParamSpecKwargs objects have a reference back to their ParamSpec: + + P.kwargs.__origin__ is P + + This type is meant for runtime introspection and has no special meaning to + static type checkers. + """ + def __init__(self, origin): + self.__origin__ = origin + + def __repr__(self): + return f"{self.__origin__.__name__}.kwargs" + + def __eq__(self, other): + if not isinstance(other, ParamSpecKwargs): + return NotImplemented + return self.__origin__ == other.__origin__ + + +if _PEP_696_IMPLEMENTED: + from typing import ParamSpec + +# 3.10+ +elif hasattr(typing, 'ParamSpec'): + + # Add default parameter - PEP 696 + class ParamSpec(metaclass=_TypeVarLikeMeta): + """Parameter specification.""" + + _backported_typevarlike = typing.ParamSpec + + def __new__(cls, name, *, bound=None, + covariant=False, contravariant=False, + infer_variance=False, default=NoDefault): + if hasattr(typing, "TypeAliasType"): + # PEP 695 implemented, can pass infer_variance to typing.TypeVar + paramspec = typing.ParamSpec(name, bound=bound, + covariant=covariant, + contravariant=contravariant, + infer_variance=infer_variance) + else: + paramspec = typing.ParamSpec(name, bound=bound, + covariant=covariant, + contravariant=contravariant) + paramspec.__infer_variance__ = infer_variance + + _set_default(paramspec, default) + _set_module(paramspec) + + def _paramspec_prepare_subst(alias, args): + params = alias.__parameters__ + i = params.index(paramspec) + if i == len(args) and paramspec.has_default(): + args = [*args, paramspec.__default__] + if i >= len(args): + raise TypeError(f"Too few arguments for {alias}") + # Special case where Z[[int, str, bool]] == Z[int, str, bool] in PEP 612. + if len(params) == 1 and not typing._is_param_expr(args[0]): + assert i == 0 + args = (args,) + # Convert lists to tuples to help other libraries cache the results. + elif isinstance(args[i], list): + args = (*args[:i], tuple(args[i]), *args[i + 1:]) + return args + + paramspec.__typing_prepare_subst__ = _paramspec_prepare_subst + return paramspec + + def __init_subclass__(cls) -> None: + raise TypeError(f"type '{__name__}.ParamSpec' is not an acceptable base type") + +# 3.9 +else: + + # Inherits from list as a workaround for Callable checks in Python < 3.9.2. + class ParamSpec(list, _DefaultMixin): + """Parameter specification variable. + + Usage:: + + P = ParamSpec('P') + + Parameter specification variables exist primarily for the benefit of static + type checkers. They are used to forward the parameter types of one + callable to another callable, a pattern commonly found in higher order + functions and decorators. They are only valid when used in ``Concatenate``, + or s the first argument to ``Callable``. In Python 3.10 and higher, + they are also supported in user-defined Generics at runtime. + See class Generic for more information on generic types. An + example for annotating a decorator:: + + T = TypeVar('T') + P = ParamSpec('P') + + def add_logging(f: Callable[P, T]) -> Callable[P, T]: + '''A type-safe decorator to add logging to a function.''' + def inner(*args: P.args, **kwargs: P.kwargs) -> T: + logging.info(f'{f.__name__} was called') + return f(*args, **kwargs) + return inner + + @add_logging + def add_two(x: float, y: float) -> float: + '''Add two numbers together.''' + return x + y + + Parameter specification variables defined with covariant=True or + contravariant=True can be used to declare covariant or contravariant + generic types. These keyword arguments are valid, but their actual semantics + are yet to be decided. See PEP 612 for details. + + Parameter specification variables can be introspected. e.g.: + + P.__name__ == 'T' + P.__bound__ == None + P.__covariant__ == False + P.__contravariant__ == False + + Note that only parameter specification variables defined in global scope can + be pickled. + """ + + # Trick Generic __parameters__. + __class__ = typing.TypeVar + + @property + def args(self): + return ParamSpecArgs(self) + + @property + def kwargs(self): + return ParamSpecKwargs(self) + + def __init__(self, name, *, bound=None, covariant=False, contravariant=False, + infer_variance=False, default=NoDefault): + list.__init__(self, [self]) + self.__name__ = name + self.__covariant__ = bool(covariant) + self.__contravariant__ = bool(contravariant) + self.__infer_variance__ = bool(infer_variance) + if bound: + self.__bound__ = typing._type_check(bound, 'Bound must be a type.') + else: + self.__bound__ = None + _DefaultMixin.__init__(self, default) + + # for pickling: + def_mod = _caller() + if def_mod != 'typing_extensions': + self.__module__ = def_mod + + def __repr__(self): + if self.__infer_variance__: + prefix = '' + elif self.__covariant__: + prefix = '+' + elif self.__contravariant__: + prefix = '-' + else: + prefix = '~' + return prefix + self.__name__ + + def __hash__(self): + return object.__hash__(self) + + def __eq__(self, other): + return self is other + + def __reduce__(self): + return self.__name__ + + # Hack to get typing._type_check to pass. + def __call__(self, *args, **kwargs): + pass + + +# 3.9 +if not hasattr(typing, 'Concatenate'): + # Inherits from list as a workaround for Callable checks in Python < 3.9.2. + + # 3.9.0-1 + if not hasattr(typing, '_type_convert'): + def _type_convert(arg, module=None, *, allow_special_forms=False): + """For converting None to type(None), and strings to ForwardRef.""" + if arg is None: + return type(None) + if isinstance(arg, str): + if sys.version_info <= (3, 9, 6): + return ForwardRef(arg) + if sys.version_info <= (3, 9, 7): + return ForwardRef(arg, module=module) + return ForwardRef(arg, module=module, is_class=allow_special_forms) + return arg + else: + _type_convert = typing._type_convert + + class _ConcatenateGenericAlias(list): + + # Trick Generic into looking into this for __parameters__. + __class__ = typing._GenericAlias + + def __init__(self, origin, args): + super().__init__(args) + self.__origin__ = origin + self.__args__ = args + + def __repr__(self): + _type_repr = typing._type_repr + return (f'{_type_repr(self.__origin__)}' + f'[{", ".join(_type_repr(arg) for arg in self.__args__)}]') + + def __hash__(self): + return hash((self.__origin__, self.__args__)) + + # Hack to get typing._type_check to pass in Generic. + def __call__(self, *args, **kwargs): + pass + + @property + def __parameters__(self): + return tuple( + tp for tp in self.__args__ if isinstance(tp, (typing.TypeVar, ParamSpec)) + ) + + # 3.9 used by __getitem__ below + def copy_with(self, params): + if isinstance(params[-1], _ConcatenateGenericAlias): + params = (*params[:-1], *params[-1].__args__) + elif isinstance(params[-1], (list, tuple)): + return (*params[:-1], *params[-1]) + elif (not (params[-1] is ... or isinstance(params[-1], ParamSpec))): + raise TypeError("The last parameter to Concatenate should be a " + "ParamSpec variable or ellipsis.") + return self.__class__(self.__origin__, params) + + # 3.9; accessed during GenericAlias.__getitem__ when substituting + def __getitem__(self, args): + if self.__origin__ in (Generic, Protocol): + # Can't subscript Generic[...] or Protocol[...]. + raise TypeError(f"Cannot subscript already-subscripted {self}") + if not self.__parameters__: + raise TypeError(f"{self} is not a generic class") + + if not isinstance(args, tuple): + args = (args,) + args = _unpack_args(*(_type_convert(p) for p in args)) + params = self.__parameters__ + for param in params: + prepare = getattr(param, "__typing_prepare_subst__", None) + if prepare is not None: + args = prepare(self, args) + # 3.9 & typing.ParamSpec + elif isinstance(param, ParamSpec): + i = params.index(param) + if ( + i == len(args) + and getattr(param, '__default__', NoDefault) is not NoDefault + ): + args = [*args, param.__default__] + if i >= len(args): + raise TypeError(f"Too few arguments for {self}") + # Special case for Z[[int, str, bool]] == Z[int, str, bool] + if len(params) == 1 and not _is_param_expr(args[0]): + assert i == 0 + args = (args,) + elif ( + isinstance(args[i], list) + # 3.9 + # This class inherits from list do not convert + and not isinstance(args[i], _ConcatenateGenericAlias) + ): + args = (*args[:i], tuple(args[i]), *args[i + 1:]) + + alen = len(args) + plen = len(params) + if alen != plen: + raise TypeError( + f"Too {'many' if alen > plen else 'few'} arguments for {self};" + f" actual {alen}, expected {plen}" + ) + + subst = dict(zip(self.__parameters__, args)) + # determine new args + new_args = [] + for arg in self.__args__: + if isinstance(arg, type): + new_args.append(arg) + continue + if isinstance(arg, TypeVar): + arg = subst[arg] + if ( + (isinstance(arg, typing._GenericAlias) and _is_unpack(arg)) + or ( + hasattr(_types, "GenericAlias") + and isinstance(arg, _types.GenericAlias) + and getattr(arg, "__unpacked__", False) + ) + ): + raise TypeError(f"{arg} is not valid as type argument") + + elif isinstance(arg, + typing._GenericAlias + if not hasattr(_types, "GenericAlias") else + (typing._GenericAlias, _types.GenericAlias) + ): + subparams = arg.__parameters__ + if subparams: + subargs = tuple(subst[x] for x in subparams) + arg = arg[subargs] + new_args.append(arg) + return self.copy_with(tuple(new_args)) + +# 3.10+ +else: + _ConcatenateGenericAlias = typing._ConcatenateGenericAlias + + # 3.10 + if sys.version_info < (3, 11): + + class _ConcatenateGenericAlias(typing._ConcatenateGenericAlias, _root=True): + # needed for checks in collections.abc.Callable to accept this class + __module__ = "typing" + + def copy_with(self, params): + if isinstance(params[-1], (list, tuple)): + return (*params[:-1], *params[-1]) + if isinstance(params[-1], typing._ConcatenateGenericAlias): + params = (*params[:-1], *params[-1].__args__) + elif not (params[-1] is ... or isinstance(params[-1], ParamSpec)): + raise TypeError("The last parameter to Concatenate should be a " + "ParamSpec variable or ellipsis.") + return super(typing._ConcatenateGenericAlias, self).copy_with(params) + + def __getitem__(self, args): + value = super().__getitem__(args) + if isinstance(value, tuple) and any(_is_unpack(t) for t in value): + return tuple(_unpack_args(*(n for n in value))) + return value + + +# 3.9.2 +class _EllipsisDummy: ... + + +# <=3.10 +def _create_concatenate_alias(origin, parameters): + if parameters[-1] is ... and sys.version_info < (3, 9, 2): + # Hack: Arguments must be types, replace it with one. + parameters = (*parameters[:-1], _EllipsisDummy) + if sys.version_info >= (3, 10, 3): + concatenate = _ConcatenateGenericAlias(origin, parameters, + _typevar_types=(TypeVar, ParamSpec), + _paramspec_tvars=True) + else: + concatenate = _ConcatenateGenericAlias(origin, parameters) + if parameters[-1] is not _EllipsisDummy: + return concatenate + # Remove dummy again + concatenate.__args__ = tuple(p if p is not _EllipsisDummy else ... + for p in concatenate.__args__) + if sys.version_info < (3, 10): + # backport needs __args__ adjustment only + return concatenate + concatenate.__parameters__ = tuple(p for p in concatenate.__parameters__ + if p is not _EllipsisDummy) + return concatenate + + +# <=3.10 +@typing._tp_cache +def _concatenate_getitem(self, parameters): + if parameters == (): + raise TypeError("Cannot take a Concatenate of no types.") + if not isinstance(parameters, tuple): + parameters = (parameters,) + if not (parameters[-1] is ... or isinstance(parameters[-1], ParamSpec)): + raise TypeError("The last parameter to Concatenate should be a " + "ParamSpec variable or ellipsis.") + msg = "Concatenate[arg, ...]: each arg must be a type." + parameters = (*(typing._type_check(p, msg) for p in parameters[:-1]), + parameters[-1]) + return _create_concatenate_alias(self, parameters) + + +# 3.11+; Concatenate does not accept ellipsis in 3.10 +# Breakpoint: https://github.com/python/cpython/pull/30969 +if sys.version_info >= (3, 11): + Concatenate = typing.Concatenate +# <=3.10 +else: + @_ExtensionsSpecialForm + def Concatenate(self, parameters): + """Used in conjunction with ``ParamSpec`` and ``Callable`` to represent a + higher order function which adds, removes or transforms parameters of a + callable. + + For example:: + + Callable[Concatenate[int, P], int] + + See PEP 612 for detailed information. + """ + return _concatenate_getitem(self, parameters) + + +# 3.10+ +if hasattr(typing, 'TypeGuard'): + TypeGuard = typing.TypeGuard +# 3.9 +else: + @_ExtensionsSpecialForm + def TypeGuard(self, parameters): + """Special typing form used to annotate the return type of a user-defined + type guard function. ``TypeGuard`` only accepts a single type argument. + At runtime, functions marked this way should return a boolean. + + ``TypeGuard`` aims to benefit *type narrowing* -- a technique used by static + type checkers to determine a more precise type of an expression within a + program's code flow. Usually type narrowing is done by analyzing + conditional code flow and applying the narrowing to a block of code. The + conditional expression here is sometimes referred to as a "type guard". + + Sometimes it would be convenient to use a user-defined boolean function + as a type guard. Such a function should use ``TypeGuard[...]`` as its + return type to alert static type checkers to this intention. + + Using ``-> TypeGuard`` tells the static type checker that for a given + function: + + 1. The return value is a boolean. + 2. If the return value is ``True``, the type of its argument + is the type inside ``TypeGuard``. + + For example:: + + def is_str(val: Union[str, float]): + # "isinstance" type guard + if isinstance(val, str): + # Type of ``val`` is narrowed to ``str`` + ... + else: + # Else, type of ``val`` is narrowed to ``float``. + ... + + Strict type narrowing is not enforced -- ``TypeB`` need not be a narrower + form of ``TypeA`` (it can even be a wider form) and this may lead to + type-unsafe results. The main reason is to allow for things like + narrowing ``List[object]`` to ``List[str]`` even though the latter is not + a subtype of the former, since ``List`` is invariant. The responsibility of + writing type-safe type guards is left to the user. + + ``TypeGuard`` also works with type variables. For more information, see + PEP 647 (User-Defined Type Guards). + """ + item = typing._type_check(parameters, f'{self} accepts only a single type.') + return typing._GenericAlias(self, (item,)) + + +# 3.13+ +if hasattr(typing, 'TypeIs'): + TypeIs = typing.TypeIs +# <=3.12 +else: + @_ExtensionsSpecialForm + def TypeIs(self, parameters): + """Special typing form used to annotate the return type of a user-defined + type narrower function. ``TypeIs`` only accepts a single type argument. + At runtime, functions marked this way should return a boolean. + + ``TypeIs`` aims to benefit *type narrowing* -- a technique used by static + type checkers to determine a more precise type of an expression within a + program's code flow. Usually type narrowing is done by analyzing + conditional code flow and applying the narrowing to a block of code. The + conditional expression here is sometimes referred to as a "type guard". + + Sometimes it would be convenient to use a user-defined boolean function + as a type guard. Such a function should use ``TypeIs[...]`` as its + return type to alert static type checkers to this intention. + + Using ``-> TypeIs`` tells the static type checker that for a given + function: + + 1. The return value is a boolean. + 2. If the return value is ``True``, the type of its argument + is the intersection of the type inside ``TypeIs`` and the argument's + previously known type. + + For example:: + + def is_awaitable(val: object) -> TypeIs[Awaitable[Any]]: + return hasattr(val, '__await__') + + def f(val: Union[int, Awaitable[int]]) -> int: + if is_awaitable(val): + assert_type(val, Awaitable[int]) + else: + assert_type(val, int) + + ``TypeIs`` also works with type variables. For more information, see + PEP 742 (Narrowing types with TypeIs). + """ + item = typing._type_check(parameters, f'{self} accepts only a single type.') + return typing._GenericAlias(self, (item,)) + + +# 3.14+? +if hasattr(typing, 'TypeForm'): + TypeForm = typing.TypeForm +# <=3.13 +else: + class _TypeFormForm(_ExtensionsSpecialForm, _root=True): + # TypeForm(X) is equivalent to X but indicates to the type checker + # that the object is a TypeForm. + def __call__(self, obj, /): + return obj + + @_TypeFormForm + def TypeForm(self, parameters): + """A special form representing the value that results from the evaluation + of a type expression. This value encodes the information supplied in the + type expression, and it represents the type described by that type expression. + + When used in a type expression, TypeForm describes a set of type form objects. + It accepts a single type argument, which must be a valid type expression. + ``TypeForm[T]`` describes the set of all type form objects that represent + the type T or types that are assignable to T. + + Usage: + + def cast[T](typ: TypeForm[T], value: Any) -> T: ... + + reveal_type(cast(int, "x")) # int + + See PEP 747 for more information. + """ + item = typing._type_check(parameters, f'{self} accepts only a single type.') + return typing._GenericAlias(self, (item,)) + + + + +if hasattr(typing, "LiteralString"): # 3.11+ + LiteralString = typing.LiteralString +else: + @_SpecialForm + def LiteralString(self, params): + """Represents an arbitrary literal string. + + Example:: + + from typing_extensions import LiteralString + + def query(sql: LiteralString) -> ...: + ... + + query("SELECT * FROM table") # ok + query(f"SELECT * FROM {input()}") # not ok + + See PEP 675 for details. + + """ + raise TypeError(f"{self} is not subscriptable") + + +if hasattr(typing, "Self"): # 3.11+ + Self = typing.Self +else: + @_SpecialForm + def Self(self, params): + """Used to spell the type of "self" in classes. + + Example:: + + from typing import Self + + class ReturnsSelf: + def parse(self, data: bytes) -> Self: + ... + return self + + """ + + raise TypeError(f"{self} is not subscriptable") + + +if hasattr(typing, "Never"): # 3.11+ + Never = typing.Never +else: + @_SpecialForm + def Never(self, params): + """The bottom type, a type that has no members. + + This can be used to define a function that should never be + called, or a function that never returns:: + + from typing_extensions import Never + + def never_call_me(arg: Never) -> None: + pass + + def int_or_str(arg: int | str) -> None: + never_call_me(arg) # type checker error + match arg: + case int(): + print("It's an int") + case str(): + print("It's a str") + case _: + never_call_me(arg) # ok, arg is of type Never + + """ + + raise TypeError(f"{self} is not subscriptable") + + +if hasattr(typing, 'Required'): # 3.11+ + Required = typing.Required + NotRequired = typing.NotRequired +else: # <=3.10 + @_ExtensionsSpecialForm + def Required(self, parameters): + """A special typing construct to mark a key of a total=False TypedDict + as required. For example: + + class Movie(TypedDict, total=False): + title: Required[str] + year: int + + m = Movie( + title='The Matrix', # typechecker error if key is omitted + year=1999, + ) + + There is no runtime checking that a required key is actually provided + when instantiating a related TypedDict. + """ + item = typing._type_check(parameters, f'{self._name} accepts only a single type.') + return typing._GenericAlias(self, (item,)) + + @_ExtensionsSpecialForm + def NotRequired(self, parameters): + """A special typing construct to mark a key of a TypedDict as + potentially missing. For example: + + class Movie(TypedDict): + title: str + year: NotRequired[int] + + m = Movie( + title='The Matrix', # typechecker error if key is omitted + year=1999, + ) + """ + item = typing._type_check(parameters, f'{self._name} accepts only a single type.') + return typing._GenericAlias(self, (item,)) + + +if hasattr(typing, 'ReadOnly'): + ReadOnly = typing.ReadOnly +else: # <=3.12 + @_ExtensionsSpecialForm + def ReadOnly(self, parameters): + """A special typing construct to mark an item of a TypedDict as read-only. + + For example: + + class Movie(TypedDict): + title: ReadOnly[str] + year: int + + def mutate_movie(m: Movie) -> None: + m["year"] = 1992 # allowed + m["title"] = "The Matrix" # typechecker error + + There is no runtime checking for this property. + """ + item = typing._type_check(parameters, f'{self._name} accepts only a single type.') + return typing._GenericAlias(self, (item,)) + + +_UNPACK_DOC = """\ +Type unpack operator. + +The type unpack operator takes the child types from some container type, +such as `tuple[int, str]` or a `TypeVarTuple`, and 'pulls them out'. For +example: + + # For some generic class `Foo`: + Foo[Unpack[tuple[int, str]]] # Equivalent to Foo[int, str] + + Ts = TypeVarTuple('Ts') + # Specifies that `Bar` is generic in an arbitrary number of types. + # (Think of `Ts` as a tuple of an arbitrary number of individual + # `TypeVar`s, which the `Unpack` is 'pulling out' directly into the + # `Generic[]`.) + class Bar(Generic[Unpack[Ts]]): ... + Bar[int] # Valid + Bar[int, str] # Also valid + +From Python 3.11, this can also be done using the `*` operator: + + Foo[*tuple[int, str]] + class Bar(Generic[*Ts]): ... + +The operator can also be used along with a `TypedDict` to annotate +`**kwargs` in a function signature. For instance: + + class Movie(TypedDict): + name: str + year: int + + # This function expects two keyword arguments - *name* of type `str` and + # *year* of type `int`. + def foo(**kwargs: Unpack[Movie]): ... + +Note that there is only some runtime checking of this operator. Not +everything the runtime allows may be accepted by static type checkers. + +For more information, see PEP 646 and PEP 692. +""" + + +# PEP 692 changed the repr of Unpack[] +# Breakpoint: https://github.com/python/cpython/pull/104048 +if sys.version_info >= (3, 12): + Unpack = typing.Unpack + + def _is_unpack(obj): + return get_origin(obj) is Unpack + +else: # <=3.11 + class _UnpackSpecialForm(_ExtensionsSpecialForm, _root=True): + def __init__(self, getitem): + super().__init__(getitem) + self.__doc__ = _UNPACK_DOC + + class _UnpackAlias(typing._GenericAlias, _root=True): + if sys.version_info < (3, 11): + # needed for compatibility with Generic[Unpack[Ts]] + __class__ = typing.TypeVar + + @property + def __typing_unpacked_tuple_args__(self): + assert self.__origin__ is Unpack + assert len(self.__args__) == 1 + arg, = self.__args__ + if isinstance(arg, (typing._GenericAlias, _types.GenericAlias)): + if arg.__origin__ is not tuple: + raise TypeError("Unpack[...] must be used with a tuple type") + return arg.__args__ + return None + + @property + def __typing_is_unpacked_typevartuple__(self): + assert self.__origin__ is Unpack + assert len(self.__args__) == 1 + return isinstance(self.__args__[0], TypeVarTuple) + + def __getitem__(self, args): + if self.__typing_is_unpacked_typevartuple__: + return args + return super().__getitem__(args) + + @_UnpackSpecialForm + def Unpack(self, parameters): + item = typing._type_check(parameters, f'{self._name} accepts only a single type.') + return _UnpackAlias(self, (item,)) + + def _is_unpack(obj): + return isinstance(obj, _UnpackAlias) + + +def _unpack_args(*args): + newargs = [] + for arg in args: + subargs = getattr(arg, '__typing_unpacked_tuple_args__', None) + if subargs is not None and (not (subargs and subargs[-1] is ...)): + newargs.extend(subargs) + else: + newargs.append(arg) + return newargs + + +if _PEP_696_IMPLEMENTED: + from typing import TypeVarTuple + +elif hasattr(typing, "TypeVarTuple"): # 3.11+ + + # Add default parameter - PEP 696 + class TypeVarTuple(metaclass=_TypeVarLikeMeta): + """Type variable tuple.""" + + _backported_typevarlike = typing.TypeVarTuple + + def __new__(cls, name, *, default=NoDefault): + tvt = typing.TypeVarTuple(name) + _set_default(tvt, default) + _set_module(tvt) + + def _typevartuple_prepare_subst(alias, args): + params = alias.__parameters__ + typevartuple_index = params.index(tvt) + for param in params[typevartuple_index + 1:]: + if isinstance(param, TypeVarTuple): + raise TypeError( + f"More than one TypeVarTuple parameter in {alias}" + ) + + alen = len(args) + plen = len(params) + left = typevartuple_index + right = plen - typevartuple_index - 1 + var_tuple_index = None + fillarg = None + for k, arg in enumerate(args): + if not isinstance(arg, type): + subargs = getattr(arg, '__typing_unpacked_tuple_args__', None) + if subargs and len(subargs) == 2 and subargs[-1] is ...: + if var_tuple_index is not None: + raise TypeError( + "More than one unpacked " + "arbitrary-length tuple argument" + ) + var_tuple_index = k + fillarg = subargs[0] + if var_tuple_index is not None: + left = min(left, var_tuple_index) + right = min(right, alen - var_tuple_index - 1) + elif left + right > alen: + raise TypeError(f"Too few arguments for {alias};" + f" actual {alen}, expected at least {plen - 1}") + if left == alen - right and tvt.has_default(): + replacement = _unpack_args(tvt.__default__) + else: + replacement = args[left: alen - right] + + return ( + *args[:left], + *([fillarg] * (typevartuple_index - left)), + replacement, + *([fillarg] * (plen - right - left - typevartuple_index - 1)), + *args[alen - right:], + ) + + tvt.__typing_prepare_subst__ = _typevartuple_prepare_subst + return tvt + + def __init_subclass__(self, *args, **kwds): + raise TypeError("Cannot subclass special typing classes") + +else: # <=3.10 + class TypeVarTuple(_DefaultMixin): + """Type variable tuple. + + Usage:: + + Ts = TypeVarTuple('Ts') + + In the same way that a normal type variable is a stand-in for a single + type such as ``int``, a type variable *tuple* is a stand-in for a *tuple* + type such as ``Tuple[int, str]``. + + Type variable tuples can be used in ``Generic`` declarations. + Consider the following example:: + + class Array(Generic[*Ts]): ... + + The ``Ts`` type variable tuple here behaves like ``tuple[T1, T2]``, + where ``T1`` and ``T2`` are type variables. To use these type variables + as type parameters of ``Array``, we must *unpack* the type variable tuple using + the star operator: ``*Ts``. The signature of ``Array`` then behaves + as if we had simply written ``class Array(Generic[T1, T2]): ...``. + In contrast to ``Generic[T1, T2]``, however, ``Generic[*Shape]`` allows + us to parameterise the class with an *arbitrary* number of type parameters. + + Type variable tuples can be used anywhere a normal ``TypeVar`` can. + This includes class definitions, as shown above, as well as function + signatures and variable annotations:: + + class Array(Generic[*Ts]): + + def __init__(self, shape: Tuple[*Ts]): + self._shape: Tuple[*Ts] = shape + + def get_shape(self) -> Tuple[*Ts]: + return self._shape + + shape = (Height(480), Width(640)) + x: Array[Height, Width] = Array(shape) + y = abs(x) # Inferred type is Array[Height, Width] + z = x + x # ... is Array[Height, Width] + x.get_shape() # ... is tuple[Height, Width] + + """ + + # Trick Generic __parameters__. + __class__ = typing.TypeVar + + def __iter__(self): + yield self.__unpacked__ + + def __init__(self, name, *, default=NoDefault): + self.__name__ = name + _DefaultMixin.__init__(self, default) + + # for pickling: + def_mod = _caller() + if def_mod != 'typing_extensions': + self.__module__ = def_mod + + self.__unpacked__ = Unpack[self] + + def __repr__(self): + return self.__name__ + + def __hash__(self): + return object.__hash__(self) + + def __eq__(self, other): + return self is other + + def __reduce__(self): + return self.__name__ + + def __init_subclass__(self, *args, **kwds): + if '_root' not in kwds: + raise TypeError("Cannot subclass special typing classes") + + +if hasattr(typing, "reveal_type"): # 3.11+ + reveal_type = typing.reveal_type +else: # <=3.10 + def reveal_type(obj: T, /) -> T: + """Reveal the inferred type of a variable. + + When a static type checker encounters a call to ``reveal_type()``, + it will emit the inferred type of the argument:: + + x: int = 1 + reveal_type(x) + + Running a static type checker (e.g., ``mypy``) on this example + will produce output similar to 'Revealed type is "builtins.int"'. + + At runtime, the function prints the runtime type of the + argument and returns it unchanged. + + """ + print(f"Runtime type is {type(obj).__name__!r}", file=sys.stderr) + return obj + + +if hasattr(typing, "_ASSERT_NEVER_REPR_MAX_LENGTH"): # 3.11+ + _ASSERT_NEVER_REPR_MAX_LENGTH = typing._ASSERT_NEVER_REPR_MAX_LENGTH +else: # <=3.10 + _ASSERT_NEVER_REPR_MAX_LENGTH = 100 + + +if hasattr(typing, "assert_never"): # 3.11+ + assert_never = typing.assert_never +else: # <=3.10 + def assert_never(arg: Never, /) -> Never: + """Assert to the type checker that a line of code is unreachable. + + Example:: + + def int_or_str(arg: int | str) -> None: + match arg: + case int(): + print("It's an int") + case str(): + print("It's a str") + case _: + assert_never(arg) + + If a type checker finds that a call to assert_never() is + reachable, it will emit an error. + + At runtime, this throws an exception when called. + + """ + value = repr(arg) + if len(value) > _ASSERT_NEVER_REPR_MAX_LENGTH: + value = value[:_ASSERT_NEVER_REPR_MAX_LENGTH] + '...' + raise AssertionError(f"Expected code to be unreachable, but got: {value}") + + +# dataclass_transform exists in 3.11 but lacks the frozen_default parameter +# Breakpoint: https://github.com/python/cpython/pull/99958 +if sys.version_info >= (3, 12): # 3.12+ + dataclass_transform = typing.dataclass_transform +else: # <=3.11 + def dataclass_transform( + *, + eq_default: bool = True, + order_default: bool = False, + kw_only_default: bool = False, + frozen_default: bool = False, + field_specifiers: typing.Tuple[ + typing.Union[typing.Type[typing.Any], typing.Callable[..., typing.Any]], + ... + ] = (), + **kwargs: typing.Any, + ) -> typing.Callable[[T], T]: + """Decorator that marks a function, class, or metaclass as providing + dataclass-like behavior. + + Example: + + from typing_extensions import dataclass_transform + + _T = TypeVar("_T") + + # Used on a decorator function + @dataclass_transform() + def create_model(cls: type[_T]) -> type[_T]: + ... + return cls + + @create_model + class CustomerModel: + id: int + name: str + + # Used on a base class + @dataclass_transform() + class ModelBase: ... + + class CustomerModel(ModelBase): + id: int + name: str + + # Used on a metaclass + @dataclass_transform() + class ModelMeta(type): ... + + class ModelBase(metaclass=ModelMeta): ... + + class CustomerModel(ModelBase): + id: int + name: str + + Each of the ``CustomerModel`` classes defined in this example will now + behave similarly to a dataclass created with the ``@dataclasses.dataclass`` + decorator. For example, the type checker will synthesize an ``__init__`` + method. + + The arguments to this decorator can be used to customize this behavior: + - ``eq_default`` indicates whether the ``eq`` parameter is assumed to be + True or False if it is omitted by the caller. + - ``order_default`` indicates whether the ``order`` parameter is + assumed to be True or False if it is omitted by the caller. + - ``kw_only_default`` indicates whether the ``kw_only`` parameter is + assumed to be True or False if it is omitted by the caller. + - ``frozen_default`` indicates whether the ``frozen`` parameter is + assumed to be True or False if it is omitted by the caller. + - ``field_specifiers`` specifies a static list of supported classes + or functions that describe fields, similar to ``dataclasses.field()``. + + At runtime, this decorator records its arguments in the + ``__dataclass_transform__`` attribute on the decorated object. + + See PEP 681 for details. + + """ + def decorator(cls_or_fn): + cls_or_fn.__dataclass_transform__ = { + "eq_default": eq_default, + "order_default": order_default, + "kw_only_default": kw_only_default, + "frozen_default": frozen_default, + "field_specifiers": field_specifiers, + "kwargs": kwargs, + } + return cls_or_fn + return decorator + + +if hasattr(typing, "override"): # 3.12+ + override = typing.override +else: # <=3.11 + _F = typing.TypeVar("_F", bound=typing.Callable[..., typing.Any]) + + def override(arg: _F, /) -> _F: + """Indicate that a method is intended to override a method in a base class. + + Usage: + + class Base: + def method(self) -> None: + pass + + class Child(Base): + @override + def method(self) -> None: + super().method() + + When this decorator is applied to a method, the type checker will + validate that it overrides a method with the same name on a base class. + This helps prevent bugs that may occur when a base class is changed + without an equivalent change to a child class. + + There is no runtime checking of these properties. The decorator + sets the ``__override__`` attribute to ``True`` on the decorated object + to allow runtime introspection. + + See PEP 698 for details. + + """ + try: + arg.__override__ = True + except (AttributeError, TypeError): + # Skip the attribute silently if it is not writable. + # AttributeError happens if the object has __slots__ or a + # read-only property, TypeError if it's a builtin class. + pass + return arg + + +# Python 3.13.3+ contains a fix for the wrapped __new__ +# Breakpoint: https://github.com/python/cpython/pull/132160 +if sys.version_info >= (3, 13, 3): + deprecated = warnings.deprecated +else: + _T = typing.TypeVar("_T") + + class deprecated: + """Indicate that a class, function or overload is deprecated. + + When this decorator is applied to an object, the type checker + will generate a diagnostic on usage of the deprecated object. + + Usage: + + @deprecated("Use B instead") + class A: + pass + + @deprecated("Use g instead") + def f(): + pass + + @overload + @deprecated("int support is deprecated") + def g(x: int) -> int: ... + @overload + def g(x: str) -> int: ... + + The warning specified by *category* will be emitted at runtime + on use of deprecated objects. For functions, that happens on calls; + for classes, on instantiation and on creation of subclasses. + If the *category* is ``None``, no warning is emitted at runtime. + The *stacklevel* determines where the + warning is emitted. If it is ``1`` (the default), the warning + is emitted at the direct caller of the deprecated object; if it + is higher, it is emitted further up the stack. + Static type checker behavior is not affected by the *category* + and *stacklevel* arguments. + + The deprecation message passed to the decorator is saved in the + ``__deprecated__`` attribute on the decorated object. + If applied to an overload, the decorator + must be after the ``@overload`` decorator for the attribute to + exist on the overload as returned by ``get_overloads()``. + + See PEP 702 for details. + + """ + def __init__( + self, + message: str, + /, + *, + category: typing.Optional[typing.Type[Warning]] = DeprecationWarning, + stacklevel: int = 1, + ) -> None: + if not isinstance(message, str): + raise TypeError( + "Expected an object of type str for 'message', not " + f"{type(message).__name__!r}" + ) + self.message = message + self.category = category + self.stacklevel = stacklevel + + def __call__(self, arg: _T, /) -> _T: + # Make sure the inner functions created below don't + # retain a reference to self. + msg = self.message + category = self.category + stacklevel = self.stacklevel + if category is None: + arg.__deprecated__ = msg + return arg + elif isinstance(arg, type): + import functools + from types import MethodType + + original_new = arg.__new__ + + @functools.wraps(original_new) + def __new__(cls, /, *args, **kwargs): + if cls is arg: + warnings.warn(msg, category=category, stacklevel=stacklevel + 1) + if original_new is not object.__new__: + return original_new(cls, *args, **kwargs) + # Mirrors a similar check in object.__new__. + elif cls.__init__ is object.__init__ and (args or kwargs): + raise TypeError(f"{cls.__name__}() takes no arguments") + else: + return original_new(cls) + + arg.__new__ = staticmethod(__new__) + + original_init_subclass = arg.__init_subclass__ + # We need slightly different behavior if __init_subclass__ + # is a bound method (likely if it was implemented in Python) + if isinstance(original_init_subclass, MethodType): + original_init_subclass = original_init_subclass.__func__ + + @functools.wraps(original_init_subclass) + def __init_subclass__(*args, **kwargs): + warnings.warn(msg, category=category, stacklevel=stacklevel + 1) + return original_init_subclass(*args, **kwargs) + + arg.__init_subclass__ = classmethod(__init_subclass__) + # Or otherwise, which likely means it's a builtin such as + # object's implementation of __init_subclass__. + else: + @functools.wraps(original_init_subclass) + def __init_subclass__(*args, **kwargs): + warnings.warn(msg, category=category, stacklevel=stacklevel + 1) + return original_init_subclass(*args, **kwargs) + + arg.__init_subclass__ = __init_subclass__ + + arg.__deprecated__ = __new__.__deprecated__ = msg + __init_subclass__.__deprecated__ = msg + return arg + elif callable(arg): + import asyncio.coroutines + import functools + import inspect + + @functools.wraps(arg) + def wrapper(*args, **kwargs): + warnings.warn(msg, category=category, stacklevel=stacklevel + 1) + return arg(*args, **kwargs) + + if asyncio.coroutines.iscoroutinefunction(arg): + # Breakpoint: https://github.com/python/cpython/pull/99247 + if sys.version_info >= (3, 12): + wrapper = inspect.markcoroutinefunction(wrapper) + else: + wrapper._is_coroutine = asyncio.coroutines._is_coroutine + + arg.__deprecated__ = wrapper.__deprecated__ = msg + return wrapper + else: + raise TypeError( + "@deprecated decorator with non-None category must be applied to " + f"a class or callable, not {arg!r}" + ) + +# Breakpoint: https://github.com/python/cpython/pull/23702 +if sys.version_info < (3, 10): + def _is_param_expr(arg): + return arg is ... or isinstance( + arg, (tuple, list, ParamSpec, _ConcatenateGenericAlias) + ) +else: + def _is_param_expr(arg): + return arg is ... or isinstance( + arg, + ( + tuple, + list, + ParamSpec, + _ConcatenateGenericAlias, + typing._ConcatenateGenericAlias, + ), + ) + + +# We have to do some monkey patching to deal with the dual nature of +# Unpack/TypeVarTuple: +# - We want Unpack to be a kind of TypeVar so it gets accepted in +# Generic[Unpack[Ts]] +# - We want it to *not* be treated as a TypeVar for the purposes of +# counting generic parameters, so that when we subscript a generic, +# the runtime doesn't try to substitute the Unpack with the subscripted type. +if not hasattr(typing, "TypeVarTuple"): + def _check_generic(cls, parameters, elen=_marker): + """Check correct count for parameters of a generic cls (internal helper). + + This gives a nice error message in case of count mismatch. + """ + # If substituting a single ParamSpec with multiple arguments + # we do not check the count + if (inspect.isclass(cls) and issubclass(cls, typing.Generic) + and len(cls.__parameters__) == 1 + and isinstance(cls.__parameters__[0], ParamSpec) + and parameters + and not _is_param_expr(parameters[0]) + ): + # Generic modifies parameters variable, but here we cannot do this + return + + if not elen: + raise TypeError(f"{cls} is not a generic class") + if elen is _marker: + if not hasattr(cls, "__parameters__") or not cls.__parameters__: + raise TypeError(f"{cls} is not a generic class") + elen = len(cls.__parameters__) + alen = len(parameters) + if alen != elen: + expect_val = elen + if hasattr(cls, "__parameters__"): + parameters = [p for p in cls.__parameters__ if not _is_unpack(p)] + num_tv_tuples = sum(isinstance(p, TypeVarTuple) for p in parameters) + if (num_tv_tuples > 0) and (alen >= elen - num_tv_tuples): + return + + # deal with TypeVarLike defaults + # required TypeVarLikes cannot appear after a defaulted one. + if alen < elen: + # since we validate TypeVarLike default in _collect_type_vars + # or _collect_parameters we can safely check parameters[alen] + if ( + getattr(parameters[alen], '__default__', NoDefault) + is not NoDefault + ): + return + + num_default_tv = sum(getattr(p, '__default__', NoDefault) + is not NoDefault for p in parameters) + + elen -= num_default_tv + + expect_val = f"at least {elen}" + + # Breakpoint: https://github.com/python/cpython/pull/27515 + things = "arguments" if sys.version_info >= (3, 10) else "parameters" + raise TypeError(f"Too {'many' if alen > elen else 'few'} {things}" + f" for {cls}; actual {alen}, expected {expect_val}") +else: + # Python 3.11+ + + def _check_generic(cls, parameters, elen): + """Check correct count for parameters of a generic cls (internal helper). + + This gives a nice error message in case of count mismatch. + """ + if not elen: + raise TypeError(f"{cls} is not a generic class") + alen = len(parameters) + if alen != elen: + expect_val = elen + if hasattr(cls, "__parameters__"): + parameters = [p for p in cls.__parameters__ if not _is_unpack(p)] + + # deal with TypeVarLike defaults + # required TypeVarLikes cannot appear after a defaulted one. + if alen < elen: + # since we validate TypeVarLike default in _collect_type_vars + # or _collect_parameters we can safely check parameters[alen] + if ( + getattr(parameters[alen], '__default__', NoDefault) + is not NoDefault + ): + return + + num_default_tv = sum(getattr(p, '__default__', NoDefault) + is not NoDefault for p in parameters) + + elen -= num_default_tv + + expect_val = f"at least {elen}" + + raise TypeError(f"Too {'many' if alen > elen else 'few'} arguments" + f" for {cls}; actual {alen}, expected {expect_val}") + +if not _PEP_696_IMPLEMENTED: + typing._check_generic = _check_generic + + +def _has_generic_or_protocol_as_origin() -> bool: + try: + frame = sys._getframe(2) + # - Catch AttributeError: not all Python implementations have sys._getframe() + # - Catch ValueError: maybe we're called from an unexpected module + # and the call stack isn't deep enough + except (AttributeError, ValueError): + return False # err on the side of leniency + else: + # If we somehow get invoked from outside typing.py, + # also err on the side of leniency + if frame.f_globals.get("__name__") != "typing": + return False + origin = frame.f_locals.get("origin") + # Cannot use "in" because origin may be an object with a buggy __eq__ that + # throws an error. + return origin is typing.Generic or origin is Protocol or origin is typing.Protocol + + +_TYPEVARTUPLE_TYPES = {TypeVarTuple, getattr(typing, "TypeVarTuple", None)} + + +def _is_unpacked_typevartuple(x) -> bool: + if get_origin(x) is not Unpack: + return False + args = get_args(x) + return ( + bool(args) + and len(args) == 1 + and type(args[0]) in _TYPEVARTUPLE_TYPES + ) + + +# Python 3.11+ _collect_type_vars was renamed to _collect_parameters +if hasattr(typing, '_collect_type_vars'): + def _collect_type_vars(types, typevar_types=None): + """Collect all type variable contained in types in order of + first appearance (lexicographic order). For example:: + + _collect_type_vars((T, List[S, T])) == (T, S) + """ + if typevar_types is None: + typevar_types = typing.TypeVar + tvars = [] + + # A required TypeVarLike cannot appear after a TypeVarLike with a default + # if it was a direct call to `Generic[]` or `Protocol[]` + enforce_default_ordering = _has_generic_or_protocol_as_origin() + default_encountered = False + + # Also, a TypeVarLike with a default cannot appear after a TypeVarTuple + type_var_tuple_encountered = False + + for t in types: + if _is_unpacked_typevartuple(t): + type_var_tuple_encountered = True + elif ( + isinstance(t, typevar_types) and not isinstance(t, _UnpackAlias) + and t not in tvars + ): + if enforce_default_ordering: + has_default = getattr(t, '__default__', NoDefault) is not NoDefault + if has_default: + if type_var_tuple_encountered: + raise TypeError('Type parameter with a default' + ' follows TypeVarTuple') + default_encountered = True + elif default_encountered: + raise TypeError(f'Type parameter {t!r} without a default' + ' follows type parameter with a default') + + tvars.append(t) + if _should_collect_from_parameters(t): + tvars.extend([t for t in t.__parameters__ if t not in tvars]) + elif isinstance(t, tuple): + # Collect nested type_vars + # tuple wrapped by _prepare_paramspec_params(cls, params) + for x in t: + for collected in _collect_type_vars([x]): + if collected not in tvars: + tvars.append(collected) + return tuple(tvars) + + typing._collect_type_vars = _collect_type_vars +else: + def _collect_parameters(args): + """Collect all type variables and parameter specifications in args + in order of first appearance (lexicographic order). + + For example:: + + assert _collect_parameters((T, Callable[P, T])) == (T, P) + """ + parameters = [] + + # A required TypeVarLike cannot appear after a TypeVarLike with default + # if it was a direct call to `Generic[]` or `Protocol[]` + enforce_default_ordering = _has_generic_or_protocol_as_origin() + default_encountered = False + + # Also, a TypeVarLike with a default cannot appear after a TypeVarTuple + type_var_tuple_encountered = False + + for t in args: + if isinstance(t, type): + # We don't want __parameters__ descriptor of a bare Python class. + pass + elif isinstance(t, tuple): + # `t` might be a tuple, when `ParamSpec` is substituted with + # `[T, int]`, or `[int, *Ts]`, etc. + for x in t: + for collected in _collect_parameters([x]): + if collected not in parameters: + parameters.append(collected) + elif hasattr(t, '__typing_subst__'): + if t not in parameters: + if enforce_default_ordering: + has_default = ( + getattr(t, '__default__', NoDefault) is not NoDefault + ) + + if type_var_tuple_encountered and has_default: + raise TypeError('Type parameter with a default' + ' follows TypeVarTuple') + + if has_default: + default_encountered = True + elif default_encountered: + raise TypeError(f'Type parameter {t!r} without a default' + ' follows type parameter with a default') + + parameters.append(t) + else: + if _is_unpacked_typevartuple(t): + type_var_tuple_encountered = True + for x in getattr(t, '__parameters__', ()): + if x not in parameters: + parameters.append(x) + + return tuple(parameters) + + if not _PEP_696_IMPLEMENTED: + typing._collect_parameters = _collect_parameters + +# Backport typing.NamedTuple as it exists in Python 3.13. +# In 3.11, the ability to define generic `NamedTuple`s was supported. +# This was explicitly disallowed in 3.9-3.10, and only half-worked in <=3.8. +# On 3.12, we added __orig_bases__ to call-based NamedTuples +# On 3.13, we deprecated kwargs-based NamedTuples +# Breakpoint: https://github.com/python/cpython/pull/105609 +if sys.version_info >= (3, 13): + NamedTuple = typing.NamedTuple +else: + def _make_nmtuple(name, types, module, defaults=()): + fields = [n for n, t in types] + annotations = {n: typing._type_check(t, f"field {n} annotation must be a type") + for n, t in types} + nm_tpl = collections.namedtuple(name, fields, + defaults=defaults, module=module) + nm_tpl.__annotations__ = nm_tpl.__new__.__annotations__ = annotations + return nm_tpl + + _prohibited_namedtuple_fields = typing._prohibited + _special_namedtuple_fields = frozenset({'__module__', '__name__', '__annotations__'}) + + class _NamedTupleMeta(type): + def __new__(cls, typename, bases, ns): + assert _NamedTuple in bases + for base in bases: + if base is not _NamedTuple and base is not typing.Generic: + raise TypeError( + 'can only inherit from a NamedTuple type and Generic') + bases = tuple(tuple if base is _NamedTuple else base for base in bases) + if "__annotations__" in ns: + types = ns["__annotations__"] + elif "__annotate__" in ns: + # TODO: Use inspect.VALUE here, and make the annotations lazily evaluated + types = ns["__annotate__"](1) + else: + types = {} + default_names = [] + for field_name in types: + if field_name in ns: + default_names.append(field_name) + elif default_names: + raise TypeError(f"Non-default namedtuple field {field_name} " + f"cannot follow default field" + f"{'s' if len(default_names) > 1 else ''} " + f"{', '.join(default_names)}") + nm_tpl = _make_nmtuple( + typename, types.items(), + defaults=[ns[n] for n in default_names], + module=ns['__module__'] + ) + nm_tpl.__bases__ = bases + if typing.Generic in bases: + if hasattr(typing, '_generic_class_getitem'): # 3.12+ + nm_tpl.__class_getitem__ = classmethod(typing._generic_class_getitem) + else: + class_getitem = typing.Generic.__class_getitem__.__func__ + nm_tpl.__class_getitem__ = classmethod(class_getitem) + # update from user namespace without overriding special namedtuple attributes + for key, val in ns.items(): + if key in _prohibited_namedtuple_fields: + raise AttributeError("Cannot overwrite NamedTuple attribute " + key) + elif key not in _special_namedtuple_fields: + if key not in nm_tpl._fields: + setattr(nm_tpl, key, ns[key]) + try: + set_name = type(val).__set_name__ + except AttributeError: + pass + else: + try: + set_name(val, nm_tpl, key) + except BaseException as e: + msg = ( + f"Error calling __set_name__ on {type(val).__name__!r} " + f"instance {key!r} in {typename!r}" + ) + # BaseException.add_note() existed on py311, + # but the __set_name__ machinery didn't start + # using add_note() until py312. + # Making sure exceptions are raised in the same way + # as in "normal" classes seems most important here. + # Breakpoint: https://github.com/python/cpython/pull/95915 + if sys.version_info >= (3, 12): + e.add_note(msg) + raise + else: + raise RuntimeError(msg) from e + + if typing.Generic in bases: + nm_tpl.__init_subclass__() + return nm_tpl + + _NamedTuple = type.__new__(_NamedTupleMeta, 'NamedTuple', (), {}) + + def _namedtuple_mro_entries(bases): + assert NamedTuple in bases + return (_NamedTuple,) + + def NamedTuple(typename, fields=_marker, /, **kwargs): + """Typed version of namedtuple. + + Usage:: + + class Employee(NamedTuple): + name: str + id: int + + This is equivalent to:: + + Employee = collections.namedtuple('Employee', ['name', 'id']) + + The resulting class has an extra __annotations__ attribute, giving a + dict that maps field names to types. (The field names are also in + the _fields attribute, which is part of the namedtuple API.) + An alternative equivalent functional syntax is also accepted:: + + Employee = NamedTuple('Employee', [('name', str), ('id', int)]) + """ + if fields is _marker: + if kwargs: + deprecated_thing = "Creating NamedTuple classes using keyword arguments" + deprecation_msg = ( + "{name} is deprecated and will be disallowed in Python {remove}. " + "Use the class-based or functional syntax instead." + ) + else: + deprecated_thing = "Failing to pass a value for the 'fields' parameter" + example = f"`{typename} = NamedTuple({typename!r}, [])`" + deprecation_msg = ( + "{name} is deprecated and will be disallowed in Python {remove}. " + "To create a NamedTuple class with 0 fields " + "using the functional syntax, " + "pass an empty list, e.g. " + ) + example + "." + elif fields is None: + if kwargs: + raise TypeError( + "Cannot pass `None` as the 'fields' parameter " + "and also specify fields using keyword arguments" + ) + else: + deprecated_thing = "Passing `None` as the 'fields' parameter" + example = f"`{typename} = NamedTuple({typename!r}, [])`" + deprecation_msg = ( + "{name} is deprecated and will be disallowed in Python {remove}. " + "To create a NamedTuple class with 0 fields " + "using the functional syntax, " + "pass an empty list, e.g. " + ) + example + "." + elif kwargs: + raise TypeError("Either list of fields or keywords" + " can be provided to NamedTuple, not both") + if fields is _marker or fields is None: + warnings.warn( + deprecation_msg.format(name=deprecated_thing, remove="3.15"), + DeprecationWarning, + stacklevel=2, + ) + fields = kwargs.items() + nt = _make_nmtuple(typename, fields, module=_caller()) + nt.__orig_bases__ = (NamedTuple,) + return nt + + NamedTuple.__mro_entries__ = _namedtuple_mro_entries + + +if hasattr(collections.abc, "Buffer"): + Buffer = collections.abc.Buffer +else: + class Buffer(abc.ABC): # noqa: B024 + """Base class for classes that implement the buffer protocol. + + The buffer protocol allows Python objects to expose a low-level + memory buffer interface. Before Python 3.12, it is not possible + to implement the buffer protocol in pure Python code, or even + to check whether a class implements the buffer protocol. In + Python 3.12 and higher, the ``__buffer__`` method allows access + to the buffer protocol from Python code, and the + ``collections.abc.Buffer`` ABC allows checking whether a class + implements the buffer protocol. + + To indicate support for the buffer protocol in earlier versions, + inherit from this ABC, either in a stub file or at runtime, + or use ABC registration. This ABC provides no methods, because + there is no Python-accessible methods shared by pre-3.12 buffer + classes. It is useful primarily for static checks. + + """ + + # As a courtesy, register the most common stdlib buffer classes. + Buffer.register(memoryview) + Buffer.register(bytearray) + Buffer.register(bytes) + + +# Backport of types.get_original_bases, available on 3.12+ in CPython +if hasattr(_types, "get_original_bases"): + get_original_bases = _types.get_original_bases +else: + def get_original_bases(cls, /): + """Return the class's "original" bases prior to modification by `__mro_entries__`. + + Examples:: + + from typing import TypeVar, Generic + from typing_extensions import NamedTuple, TypedDict + + T = TypeVar("T") + class Foo(Generic[T]): ... + class Bar(Foo[int], float): ... + class Baz(list[str]): ... + Eggs = NamedTuple("Eggs", [("a", int), ("b", str)]) + Spam = TypedDict("Spam", {"a": int, "b": str}) + + assert get_original_bases(Bar) == (Foo[int], float) + assert get_original_bases(Baz) == (list[str],) + assert get_original_bases(Eggs) == (NamedTuple,) + assert get_original_bases(Spam) == (TypedDict,) + assert get_original_bases(int) == (object,) + """ + try: + return cls.__dict__.get("__orig_bases__", cls.__bases__) + except AttributeError: + raise TypeError( + f'Expected an instance of type, not {type(cls).__name__!r}' + ) from None + + +# NewType is a class on Python 3.10+, making it pickleable +# The error message for subclassing instances of NewType was improved on 3.11+ +# Breakpoint: https://github.com/python/cpython/pull/30268 +if sys.version_info >= (3, 11): + NewType = typing.NewType +else: + class NewType: + """NewType creates simple unique types with almost zero + runtime overhead. NewType(name, tp) is considered a subtype of tp + by static type checkers. At runtime, NewType(name, tp) returns + a dummy callable that simply returns its argument. Usage:: + UserId = NewType('UserId', int) + def name_by_id(user_id: UserId) -> str: + ... + UserId('user') # Fails type check + name_by_id(42) # Fails type check + name_by_id(UserId(42)) # OK + num = UserId(5) + 1 # type: int + """ + + def __call__(self, obj, /): + return obj + + def __init__(self, name, tp): + self.__qualname__ = name + if '.' in name: + name = name.rpartition('.')[-1] + self.__name__ = name + self.__supertype__ = tp + def_mod = _caller() + if def_mod != 'typing_extensions': + self.__module__ = def_mod + + def __mro_entries__(self, bases): + # We defined __mro_entries__ to get a better error message + # if a user attempts to subclass a NewType instance. bpo-46170 + supercls_name = self.__name__ + + class Dummy: + def __init_subclass__(cls): + subcls_name = cls.__name__ + raise TypeError( + f"Cannot subclass an instance of NewType. " + f"Perhaps you were looking for: " + f"`{subcls_name} = NewType({subcls_name!r}, {supercls_name})`" + ) + + return (Dummy,) + + def __repr__(self): + return f'{self.__module__}.{self.__qualname__}' + + def __reduce__(self): + return self.__qualname__ + + # Breakpoint: https://github.com/python/cpython/pull/21515 + if sys.version_info >= (3, 10): + # PEP 604 methods + # It doesn't make sense to have these methods on Python <3.10 + + def __or__(self, other): + return typing.Union[self, other] + + def __ror__(self, other): + return typing.Union[other, self] + + +# Breakpoint: https://github.com/python/cpython/pull/124795 +if sys.version_info >= (3, 14): + TypeAliasType = typing.TypeAliasType +# <=3.13 +else: + # Breakpoint: https://github.com/python/cpython/pull/103764 + if sys.version_info >= (3, 12): + # 3.12-3.13 + def _is_unionable(obj): + """Corresponds to is_unionable() in unionobject.c in CPython.""" + return obj is None or isinstance(obj, ( + type, + _types.GenericAlias, + _types.UnionType, + typing.TypeAliasType, + TypeAliasType, + )) + else: + # <=3.11 + def _is_unionable(obj): + """Corresponds to is_unionable() in unionobject.c in CPython.""" + return obj is None or isinstance(obj, ( + type, + _types.GenericAlias, + _types.UnionType, + TypeAliasType, + )) + + if sys.version_info < (3, 10): + # Copied and pasted from https://github.com/python/cpython/blob/986a4e1b6fcae7fe7a1d0a26aea446107dd58dd2/Objects/genericaliasobject.c#L568-L582, + # so that we emulate the behaviour of `types.GenericAlias` + # on the latest versions of CPython + _ATTRIBUTE_DELEGATION_EXCLUSIONS = frozenset({ + "__class__", + "__bases__", + "__origin__", + "__args__", + "__unpacked__", + "__parameters__", + "__typing_unpacked_tuple_args__", + "__mro_entries__", + "__reduce_ex__", + "__reduce__", + "__copy__", + "__deepcopy__", + }) + + class _TypeAliasGenericAlias(typing._GenericAlias, _root=True): + def __getattr__(self, attr): + if attr in _ATTRIBUTE_DELEGATION_EXCLUSIONS: + return object.__getattr__(self, attr) + return getattr(self.__origin__, attr) + + + class TypeAliasType: + """Create named, parameterized type aliases. + + This provides a backport of the new `type` statement in Python 3.12: + + type ListOrSet[T] = list[T] | set[T] + + is equivalent to: + + T = TypeVar("T") + ListOrSet = TypeAliasType("ListOrSet", list[T] | set[T], type_params=(T,)) + + The name ListOrSet can then be used as an alias for the type it refers to. + + The type_params argument should contain all the type parameters used + in the value of the type alias. If the alias is not generic, this + argument is omitted. + + Static type checkers should only support type aliases declared using + TypeAliasType that follow these rules: + + - The first argument (the name) must be a string literal. + - The TypeAliasType instance must be immediately assigned to a variable + of the same name. (For example, 'X = TypeAliasType("Y", int)' is invalid, + as is 'X, Y = TypeAliasType("X", int), TypeAliasType("Y", int)'). + + """ + + def __init__(self, name: str, value, *, type_params=()): + if not isinstance(name, str): + raise TypeError("TypeAliasType name must be a string") + if not isinstance(type_params, tuple): + raise TypeError("type_params must be a tuple") + self.__value__ = value + self.__type_params__ = type_params + + default_value_encountered = False + parameters = [] + for type_param in type_params: + if ( + not isinstance(type_param, (TypeVar, TypeVarTuple, ParamSpec)) + # <=3.11 + # Unpack Backport passes isinstance(type_param, TypeVar) + or _is_unpack(type_param) + ): + raise TypeError(f"Expected a type param, got {type_param!r}") + has_default = ( + getattr(type_param, '__default__', NoDefault) is not NoDefault + ) + if default_value_encountered and not has_default: + raise TypeError(f"non-default type parameter '{type_param!r}'" + " follows default type parameter") + if has_default: + default_value_encountered = True + if isinstance(type_param, TypeVarTuple): + parameters.extend(type_param) + else: + parameters.append(type_param) + self.__parameters__ = tuple(parameters) + def_mod = _caller() + if def_mod != 'typing_extensions': + self.__module__ = def_mod + # Setting this attribute closes the TypeAliasType from further modification + self.__name__ = name + + def __setattr__(self, name: str, value: object, /) -> None: + if hasattr(self, "__name__"): + self._raise_attribute_error(name) + super().__setattr__(name, value) + + def __delattr__(self, name: str, /) -> Never: + self._raise_attribute_error(name) + + def _raise_attribute_error(self, name: str) -> Never: + # Match the Python 3.12 error messages exactly + if name == "__name__": + raise AttributeError("readonly attribute") + elif name in {"__value__", "__type_params__", "__parameters__", "__module__"}: + raise AttributeError( + f"attribute '{name}' of 'typing.TypeAliasType' objects " + "is not writable" + ) + else: + raise AttributeError( + f"'typing.TypeAliasType' object has no attribute '{name}'" + ) + + def __repr__(self) -> str: + return self.__name__ + + if sys.version_info < (3, 11): + def _check_single_param(self, param, recursion=0): + # Allow [], [int], [int, str], [int, ...], [int, T] + if param is ...: + return ... + if param is None: + return None + # Note in <= 3.9 _ConcatenateGenericAlias inherits from list + if isinstance(param, list) and recursion == 0: + return [self._check_single_param(arg, recursion+1) + for arg in param] + return typing._type_check( + param, f'Subscripting {self.__name__} requires a type.' + ) + + def _check_parameters(self, parameters): + if sys.version_info < (3, 11): + return tuple( + self._check_single_param(item) + for item in parameters + ) + return tuple(typing._type_check( + item, f'Subscripting {self.__name__} requires a type.' + ) + for item in parameters + ) + + def __getitem__(self, parameters): + if not self.__type_params__: + raise TypeError("Only generic type aliases are subscriptable") + if not isinstance(parameters, tuple): + parameters = (parameters,) + # Using 3.9 here will create problems with Concatenate + if sys.version_info >= (3, 10): + return _types.GenericAlias(self, parameters) + type_vars = _collect_type_vars(parameters) + parameters = self._check_parameters(parameters) + alias = _TypeAliasGenericAlias(self, parameters) + # alias.__parameters__ is not complete if Concatenate is present + # as it is converted to a list from which no parameters are extracted. + if alias.__parameters__ != type_vars: + alias.__parameters__ = type_vars + return alias + + def __reduce__(self): + return self.__name__ + + def __init_subclass__(cls, *args, **kwargs): + raise TypeError( + "type 'typing_extensions.TypeAliasType' is not an acceptable base type" + ) + + # The presence of this method convinces typing._type_check + # that TypeAliasTypes are types. + def __call__(self): + raise TypeError("Type alias is not callable") + + # Breakpoint: https://github.com/python/cpython/pull/21515 + if sys.version_info >= (3, 10): + def __or__(self, right): + # For forward compatibility with 3.12, reject Unions + # that are not accepted by the built-in Union. + if not _is_unionable(right): + return NotImplemented + return typing.Union[self, right] + + def __ror__(self, left): + if not _is_unionable(left): + return NotImplemented + return typing.Union[left, self] + + +if hasattr(typing, "is_protocol"): + is_protocol = typing.is_protocol + get_protocol_members = typing.get_protocol_members +else: + def is_protocol(tp: type, /) -> bool: + """Return True if the given type is a Protocol. + + Example:: + + >>> from typing_extensions import Protocol, is_protocol + >>> class P(Protocol): + ... def a(self) -> str: ... + ... b: int + >>> is_protocol(P) + True + >>> is_protocol(int) + False + """ + return ( + isinstance(tp, type) + and getattr(tp, '_is_protocol', False) + and tp is not Protocol + and tp is not typing.Protocol + ) + + def get_protocol_members(tp: type, /) -> typing.FrozenSet[str]: + """Return the set of members defined in a Protocol. + + Example:: + + >>> from typing_extensions import Protocol, get_protocol_members + >>> class P(Protocol): + ... def a(self) -> str: ... + ... b: int + >>> get_protocol_members(P) + frozenset({'a', 'b'}) + + Raise a TypeError for arguments that are not Protocols. + """ + if not is_protocol(tp): + raise TypeError(f'{tp!r} is not a Protocol') + if hasattr(tp, '__protocol_attrs__'): + return frozenset(tp.__protocol_attrs__) + return frozenset(_get_protocol_attrs(tp)) + + +if hasattr(typing, "Doc"): + Doc = typing.Doc +else: + class Doc: + """Define the documentation of a type annotation using ``Annotated``, to be + used in class attributes, function and method parameters, return values, + and variables. + + The value should be a positional-only string literal to allow static tools + like editors and documentation generators to use it. + + This complements docstrings. + + The string value passed is available in the attribute ``documentation``. + + Example:: + + >>> from typing_extensions import Annotated, Doc + >>> def hi(to: Annotated[str, Doc("Who to say hi to")]) -> None: ... + """ + def __init__(self, documentation: str, /) -> None: + self.documentation = documentation + + def __repr__(self) -> str: + return f"Doc({self.documentation!r})" + + def __hash__(self) -> int: + return hash(self.documentation) + + def __eq__(self, other: object) -> bool: + if not isinstance(other, Doc): + return NotImplemented + return self.documentation == other.documentation + + +_CapsuleType = getattr(_types, "CapsuleType", None) + +if _CapsuleType is None: + try: + import _socket + except ImportError: + pass + else: + _CAPI = getattr(_socket, "CAPI", None) + if _CAPI is not None: + _CapsuleType = type(_CAPI) + +if _CapsuleType is not None: + CapsuleType = _CapsuleType + __all__.append("CapsuleType") + + +if sys.version_info >= (3, 14): + from annotationlib import Format, get_annotations +else: + # Available since Python 3.14.0a3 + # PR: https://github.com/python/cpython/pull/124415 + class Format(enum.IntEnum): + VALUE = 1 + VALUE_WITH_FAKE_GLOBALS = 2 + FORWARDREF = 3 + STRING = 4 + + # Available since Python 3.14.0a1 + # PR: https://github.com/python/cpython/pull/119891 + def get_annotations(obj, *, globals=None, locals=None, eval_str=False, + format=Format.VALUE): + """Compute the annotations dict for an object. + + obj may be a callable, class, or module. + Passing in an object of any other type raises TypeError. + + Returns a dict. get_annotations() returns a new dict every time + it's called; calling it twice on the same object will return two + different but equivalent dicts. + + This is a backport of `inspect.get_annotations`, which has been + in the standard library since Python 3.10. See the standard library + documentation for more: + + https://docs.python.org/3/library/inspect.html#inspect.get_annotations + + This backport adds the *format* argument introduced by PEP 649. The + three formats supported are: + * VALUE: the annotations are returned as-is. This is the default and + it is compatible with the behavior on previous Python versions. + * FORWARDREF: return annotations as-is if possible, but replace any + undefined names with ForwardRef objects. The implementation proposed by + PEP 649 relies on language changes that cannot be backported; the + typing-extensions implementation simply returns the same result as VALUE. + * STRING: return annotations as strings, in a format close to the original + source. Again, this behavior cannot be replicated directly in a backport. + As an approximation, typing-extensions retrieves the annotations under + VALUE semantics and then stringifies them. + + The purpose of this backport is to allow users who would like to use + FORWARDREF or STRING semantics once PEP 649 is implemented, but who also + want to support earlier Python versions, to simply write: + + typing_extensions.get_annotations(obj, format=Format.FORWARDREF) + + """ + format = Format(format) + if format is Format.VALUE_WITH_FAKE_GLOBALS: + raise ValueError( + "The VALUE_WITH_FAKE_GLOBALS format is for internal use only" + ) + + if eval_str and format is not Format.VALUE: + raise ValueError("eval_str=True is only supported with format=Format.VALUE") + + if isinstance(obj, type): + # class + obj_dict = getattr(obj, '__dict__', None) + if obj_dict and hasattr(obj_dict, 'get'): + ann = obj_dict.get('__annotations__', None) + if isinstance(ann, _types.GetSetDescriptorType): + ann = None + else: + ann = None + + obj_globals = None + module_name = getattr(obj, '__module__', None) + if module_name: + module = sys.modules.get(module_name, None) + if module: + obj_globals = getattr(module, '__dict__', None) + obj_locals = dict(vars(obj)) + unwrap = obj + elif isinstance(obj, _types.ModuleType): + # module + ann = getattr(obj, '__annotations__', None) + obj_globals = obj.__dict__ + obj_locals = None + unwrap = None + elif callable(obj): + # this includes types.Function, types.BuiltinFunctionType, + # types.BuiltinMethodType, functools.partial, functools.singledispatch, + # "class funclike" from Lib/test/test_inspect... on and on it goes. + ann = getattr(obj, '__annotations__', None) + obj_globals = getattr(obj, '__globals__', None) + obj_locals = None + unwrap = obj + elif hasattr(obj, '__annotations__'): + ann = obj.__annotations__ + obj_globals = obj_locals = unwrap = None + else: + raise TypeError(f"{obj!r} is not a module, class, or callable.") + + if ann is None: + return {} + + if not isinstance(ann, dict): + raise ValueError(f"{obj!r}.__annotations__ is neither a dict nor None") + + if not ann: + return {} + + if not eval_str: + if format is Format.STRING: + return { + key: value if isinstance(value, str) else typing._type_repr(value) + for key, value in ann.items() + } + return dict(ann) + + if unwrap is not None: + while True: + if hasattr(unwrap, '__wrapped__'): + unwrap = unwrap.__wrapped__ + continue + if isinstance(unwrap, functools.partial): + unwrap = unwrap.func + continue + break + if hasattr(unwrap, "__globals__"): + obj_globals = unwrap.__globals__ + + if globals is None: + globals = obj_globals + if locals is None: + locals = obj_locals or {} + + # "Inject" type parameters into the local namespace + # (unless they are shadowed by assignments *in* the local namespace), + # as a way of emulating annotation scopes when calling `eval()` + if type_params := getattr(obj, "__type_params__", ()): + locals = {param.__name__: param for param in type_params} | locals + + return_value = {key: + value if not isinstance(value, str) else eval(value, globals, locals) + for key, value in ann.items() } + return return_value + + +if hasattr(typing, "evaluate_forward_ref"): + evaluate_forward_ref = typing.evaluate_forward_ref +else: + # Implements annotationlib.ForwardRef.evaluate + def _eval_with_owner( + forward_ref, *, owner=None, globals=None, locals=None, type_params=None + ): + if forward_ref.__forward_evaluated__: + return forward_ref.__forward_value__ + if getattr(forward_ref, "__cell__", None) is not None: + try: + value = forward_ref.__cell__.cell_contents + except ValueError: + pass + else: + forward_ref.__forward_evaluated__ = True + forward_ref.__forward_value__ = value + return value + if owner is None: + owner = getattr(forward_ref, "__owner__", None) + + if ( + globals is None + and getattr(forward_ref, "__forward_module__", None) is not None + ): + globals = getattr( + sys.modules.get(forward_ref.__forward_module__, None), "__dict__", None + ) + if globals is None: + globals = getattr(forward_ref, "__globals__", None) + if globals is None: + if isinstance(owner, type): + module_name = getattr(owner, "__module__", None) + if module_name: + module = sys.modules.get(module_name, None) + if module: + globals = getattr(module, "__dict__", None) + elif isinstance(owner, _types.ModuleType): + globals = getattr(owner, "__dict__", None) + elif callable(owner): + globals = getattr(owner, "__globals__", None) + + # If we pass None to eval() below, the globals of this module are used. + if globals is None: + globals = {} + + if locals is None: + locals = {} + if isinstance(owner, type): + locals.update(vars(owner)) + + if type_params is None and owner is not None: + # "Inject" type parameters into the local namespace + # (unless they are shadowed by assignments *in* the local namespace), + # as a way of emulating annotation scopes when calling `eval()` + type_params = getattr(owner, "__type_params__", None) + + # Type parameters exist in their own scope, which is logically + # between the locals and the globals. We simulate this by adding + # them to the globals. + if type_params is not None: + globals = dict(globals) + for param in type_params: + globals[param.__name__] = param + + arg = forward_ref.__forward_arg__ + if arg.isidentifier() and not keyword.iskeyword(arg): + if arg in locals: + value = locals[arg] + elif arg in globals: + value = globals[arg] + elif hasattr(builtins, arg): + return getattr(builtins, arg) + else: + raise NameError(arg) + else: + code = forward_ref.__forward_code__ + value = eval(code, globals, locals) + forward_ref.__forward_evaluated__ = True + forward_ref.__forward_value__ = value + return value + + def evaluate_forward_ref( + forward_ref, + *, + owner=None, + globals=None, + locals=None, + type_params=None, + format=None, + _recursive_guard=frozenset(), + ): + """Evaluate a forward reference as a type hint. + + This is similar to calling the ForwardRef.evaluate() method, + but unlike that method, evaluate_forward_ref() also: + + * Recursively evaluates forward references nested within the type hint. + * Rejects certain objects that are not valid type hints. + * Replaces type hints that evaluate to None with types.NoneType. + * Supports the *FORWARDREF* and *STRING* formats. + + *forward_ref* must be an instance of ForwardRef. *owner*, if given, + should be the object that holds the annotations that the forward reference + derived from, such as a module, class object, or function. It is used to + infer the namespaces to use for looking up names. *globals* and *locals* + can also be explicitly given to provide the global and local namespaces. + *type_params* is a tuple of type parameters that are in scope when + evaluating the forward reference. This parameter must be provided (though + it may be an empty tuple) if *owner* is not given and the forward reference + does not already have an owner set. *format* specifies the format of the + annotation and is a member of the annotationlib.Format enum. + + """ + if format == Format.STRING: + return forward_ref.__forward_arg__ + if forward_ref.__forward_arg__ in _recursive_guard: + return forward_ref + + # Evaluate the forward reference + try: + value = _eval_with_owner( + forward_ref, + owner=owner, + globals=globals, + locals=locals, + type_params=type_params, + ) + except NameError: + if format == Format.FORWARDREF: + return forward_ref + else: + raise + + if isinstance(value, str): + value = ForwardRef(value) + + # Recursively evaluate the type + if isinstance(value, ForwardRef): + if getattr(value, "__forward_module__", True) is not None: + globals = None + return evaluate_forward_ref( + value, + globals=globals, + locals=locals, + type_params=type_params, owner=owner, + _recursive_guard=_recursive_guard, format=format + ) + if sys.version_info < (3, 12, 5) and type_params: + # Make use of type_params + locals = dict(locals) if locals else {} + for tvar in type_params: + if tvar.__name__ not in locals: # lets not overwrite something present + locals[tvar.__name__] = tvar + if sys.version_info < (3, 12, 5): + return typing._eval_type( + value, + globals, + locals, + recursive_guard=_recursive_guard | {forward_ref.__forward_arg__}, + ) + else: + return typing._eval_type( + value, + globals, + locals, + type_params, + recursive_guard=_recursive_guard | {forward_ref.__forward_arg__}, + ) + + +class Sentinel: + """Create a unique sentinel object. + + *name* should be the name of the variable to which the return value shall be assigned. + + *repr*, if supplied, will be used for the repr of the sentinel object. + If not provided, "" will be used. + """ + + def __init__( + self, + name: str, + repr: typing.Optional[str] = None, + ): + self._name = name + self._repr = repr if repr is not None else f'<{name}>' + + def __repr__(self): + return self._repr + + if sys.version_info < (3, 11): + # The presence of this method convinces typing._type_check + # that Sentinels are types. + def __call__(self, *args, **kwargs): + raise TypeError(f"{type(self).__name__!r} object is not callable") + + # Breakpoint: https://github.com/python/cpython/pull/21515 + if sys.version_info >= (3, 10): + def __or__(self, other): + return typing.Union[self, other] + + def __ror__(self, other): + return typing.Union[other, self] + + def __getstate__(self): + raise TypeError(f"Cannot pickle {type(self).__name__!r} object") + + +if sys.version_info >= (3, 14, 0, "beta"): + type_repr = annotationlib.type_repr +else: + def type_repr(value): + """Convert a Python value to a format suitable for use with the STRING format. + + This is intended as a helper for tools that support the STRING format but do + not have access to the code that originally produced the annotations. It uses + repr() for most objects. + + """ + if isinstance(value, (type, _types.FunctionType, _types.BuiltinFunctionType)): + if value.__module__ == "builtins": + return value.__qualname__ + return f"{value.__module__}.{value.__qualname__}" + if value is ...: + return "..." + return repr(value) + + +# Aliases for items that are in typing in all supported versions. +# We use hasattr() checks so this library will continue to import on +# future versions of Python that may remove these names. +_typing_names = [ + "AbstractSet", + "AnyStr", + "BinaryIO", + "Callable", + "Collection", + "Container", + "Dict", + "FrozenSet", + "Hashable", + "IO", + "ItemsView", + "Iterable", + "Iterator", + "KeysView", + "List", + "Mapping", + "MappingView", + "Match", + "MutableMapping", + "MutableSequence", + "MutableSet", + "Optional", + "Pattern", + "Reversible", + "Sequence", + "Set", + "Sized", + "TextIO", + "Tuple", + "Union", + "ValuesView", + "cast", + "no_type_check", + "no_type_check_decorator", + # This is private, but it was defined by typing_extensions for a long time + # and some users rely on it. + "_AnnotatedAlias", +] +globals().update( + {name: getattr(typing, name) for name in _typing_names if hasattr(typing, name)} +) +# These are defined unconditionally because they are used in +# typing-extensions itself. +Generic = typing.Generic +ForwardRef = typing.ForwardRef +Annotated = typing.Annotated diff --git a/lib/python3.12/site-packages/ray/_private/runtime_env/conda.py b/lib/python3.12/site-packages/ray/_private/runtime_env/conda.py new file mode 100644 index 0000000000000000000000000000000000000000..7ca67890243345edbe24ff387718e37b121dbf9c --- /dev/null +++ b/lib/python3.12/site-packages/ray/_private/runtime_env/conda.py @@ -0,0 +1,396 @@ +import hashlib +import json +import logging +import os +import runpy +import shutil +import subprocess +import sys +from pathlib import Path +from typing import Any, Dict, List, Optional + +import yaml +from filelock import FileLock + +import ray +from ray._common.utils import ( + get_or_create_event_loop, + try_to_create_directory, +) +from ray._private.runtime_env.conda_utils import ( + create_conda_env_if_needed, + delete_conda_env, + get_conda_activate_commands, + get_conda_envs, + get_conda_info_json, +) +from ray._private.runtime_env.context import RuntimeEnvContext +from ray._private.runtime_env.packaging import Protocol, parse_uri +from ray._private.runtime_env.plugin import RuntimeEnvPlugin +from ray._private.runtime_env.validation import parse_and_validate_conda +from ray._private.utils import ( + get_directory_size_bytes, + get_master_wheel_url, + get_release_wheel_url, + get_wheel_filename, +) + +default_logger = logging.getLogger(__name__) + +_WIN32 = os.name == "nt" + + +def _resolve_current_ray_path() -> str: + # When ray is built from source with pip install -e, + # ray.__file__ returns .../python/ray/__init__.py and this function returns + # ".../python". + # When ray is installed from a prebuilt binary, ray.__file__ returns + # .../site-packages/ray/__init__.py and this function returns + # ".../site-packages". + return os.path.split(os.path.split(ray.__file__)[0])[0] + + +def _get_ray_setup_spec(): + """Find the Ray setup_spec from the currently running Ray. + + This function works even when Ray is built from source with pip install -e. + """ + ray_source_python_path = _resolve_current_ray_path() + setup_py_path = os.path.join(ray_source_python_path, "setup.py") + return runpy.run_path(setup_py_path)["setup_spec"] + + +def _resolve_install_from_source_ray_dependencies(): + """Find the Ray dependencies when Ray is installed from source.""" + deps = ( + _get_ray_setup_spec().install_requires + _get_ray_setup_spec().extras["default"] + ) + # Remove duplicates + return list(set(deps)) + + +def _inject_ray_to_conda_site( + conda_path, logger: Optional[logging.Logger] = default_logger +): + """Write the current Ray site package directory to a new site""" + if _WIN32: + python_binary = os.path.join(conda_path, "python") + else: + python_binary = os.path.join(conda_path, "bin/python") + site_packages_path = ( + subprocess.check_output( + [ + python_binary, + "-c", + "import sysconfig; print(sysconfig.get_paths()['purelib'])", + ] + ) + .decode() + .strip() + ) + + ray_path = _resolve_current_ray_path() + logger.warning( + f"Injecting {ray_path} to environment site-packages {site_packages_path} " + "because _inject_current_ray flag is on." + ) + + maybe_ray_dir = os.path.join(site_packages_path, "ray") + if os.path.isdir(maybe_ray_dir): + logger.warning(f"Replacing existing ray installation with {ray_path}") + shutil.rmtree(maybe_ray_dir) + + # See usage of *.pth file at + # https://docs.python.org/3/library/site.html + with open(os.path.join(site_packages_path, "ray_shared.pth"), "w") as f: + f.write(ray_path) + + +def _current_py_version(): + return ".".join(map(str, sys.version_info[:3])) # like 3.6.10 + + +def current_ray_pip_specifier( + logger: Optional[logging.Logger] = default_logger, +) -> Optional[str]: + """The pip requirement specifier for the running version of Ray. + + Returns: + A string which can be passed to `pip install` to install the + currently running Ray version, or None if running on a version + built from source locally (likely if you are developing Ray). + + Examples: + Returns "https://s3-us-west-2.amazonaws.com/ray-wheels/[..].whl" + if running a stable release, a nightly or a specific commit + """ + if os.environ.get("RAY_CI_POST_WHEEL_TESTS"): + # Running in Buildkite CI after the wheel has been built. + # Wheels are at in the ray/.whl directory, but use relative path to + # allow for testing locally if needed. + return os.path.join( + Path(ray.__file__).resolve().parents[2], ".whl", get_wheel_filename() + ) + elif ray.__commit__ == "{{RAY_COMMIT_SHA}}": + # Running on a version built from source locally. + if os.environ.get("RAY_RUNTIME_ENV_LOCAL_DEV_MODE") != "1": + logger.warning( + "Current Ray version could not be detected, most likely " + "because you have manually built Ray from source. To use " + "runtime_env in this case, set the environment variable " + "RAY_RUNTIME_ENV_LOCAL_DEV_MODE=1." + ) + return None + elif "dev" in ray.__version__: + # Running on a nightly wheel. + return get_master_wheel_url() + else: + return get_release_wheel_url() + + +def inject_dependencies( + conda_dict: Dict[Any, Any], + py_version: str, + pip_dependencies: Optional[List[str]] = None, +) -> Dict[Any, Any]: + """Add Ray, Python and (optionally) extra pip dependencies to a conda dict. + + Args: + conda_dict: A dict representing the JSON-serialized conda + environment YAML file. This dict will be modified and returned. + py_version: A string representing a Python version to inject + into the conda dependencies, e.g. "3.7.7" + pip_dependencies (List[str]): A list of pip dependencies that + will be prepended to the list of pip dependencies in + the conda dict. If the conda dict does not already have a "pip" + field, one will be created. + Returns: + The modified dict. (Note: the input argument conda_dict is modified + and returned.) + """ + if pip_dependencies is None: + pip_dependencies = [] + if conda_dict.get("dependencies") is None: + conda_dict["dependencies"] = [] + + # Inject Python dependency. + deps = conda_dict["dependencies"] + + # Add current python dependency. If the user has already included a + # python version dependency, conda will raise a readable error if the two + # are incompatible, e.g: + # ResolvePackageNotFound: - python[version='3.5.*,>=3.6'] + deps.append(f"python={py_version}") + + if "pip" not in deps: + deps.append("pip") + + # Insert pip dependencies. + found_pip_dict = False + for dep in deps: + if isinstance(dep, dict) and dep.get("pip") and isinstance(dep["pip"], list): + dep["pip"] = pip_dependencies + dep["pip"] + found_pip_dict = True + break + if not found_pip_dict: + deps.append({"pip": pip_dependencies}) + + return conda_dict + + +def _get_conda_env_hash(conda_dict: Dict) -> str: + # Set `sort_keys=True` so that different orderings yield the same hash. + serialized_conda_spec = json.dumps(conda_dict, sort_keys=True) + hash = hashlib.sha1(serialized_conda_spec.encode("utf-8")).hexdigest() + return hash + + +def get_uri(runtime_env: Dict) -> Optional[str]: + """Return `"conda://"`, or None if no GC required.""" + conda = runtime_env.get("conda") + if conda is not None: + if isinstance(conda, str): + # User-preinstalled conda env. We don't garbage collect these, so + # we don't track them with URIs. + uri = None + elif isinstance(conda, dict): + uri = f"conda://{_get_conda_env_hash(conda_dict=conda)}" + else: + raise TypeError( + "conda field received by RuntimeEnvAgent must be " + f"str or dict, not {type(conda).__name__}." + ) + else: + uri = None + return uri + + +def _get_conda_dict_with_ray_inserted( + runtime_env: "RuntimeEnv", # noqa: F821 + logger: Optional[logging.Logger] = default_logger, +) -> Dict[str, Any]: + """Returns the conda spec with the Ray and `python` dependency inserted.""" + conda_dict = json.loads(runtime_env.conda_config()) + assert conda_dict is not None + + ray_pip = current_ray_pip_specifier(logger=logger) + if ray_pip: + extra_pip_dependencies = [ray_pip, "ray[default]"] + elif runtime_env.get_extension("_inject_current_ray"): + extra_pip_dependencies = _resolve_install_from_source_ray_dependencies() + else: + extra_pip_dependencies = [] + conda_dict = inject_dependencies( + conda_dict, _current_py_version(), extra_pip_dependencies + ) + return conda_dict + + +class CondaPlugin(RuntimeEnvPlugin): + + name = "conda" + + def __init__(self, resources_dir: str): + self._resources_dir = os.path.join(resources_dir, "conda") + try_to_create_directory(self._resources_dir) + + # It is not safe for multiple processes to install conda envs + # concurrently, even if the envs are different, so use a global + # lock for all conda installs and deletions. + # See https://github.com/ray-project/ray/issues/17086 + self._installs_and_deletions_file_lock = os.path.join( + self._resources_dir, "ray-conda-installs-and-deletions.lock" + ) + # A set of named conda environments (instead of yaml or dict) + # that are validated to exist. + # NOTE: It has to be only used within the same thread, which + # is an event loop. + # Also, we don't need to GC this field because it is pretty small. + self._validated_named_conda_env = set() + + def _get_path_from_hash(self, hash: str) -> str: + """Generate a path from the hash of a conda or pip spec. + + The output path also functions as the name of the conda environment + when using the `--prefix` option to `conda create` and `conda remove`. + + Example output: + /tmp/ray/session_2021-11-03_16-33-59_356303_41018/runtime_resources + /conda/ray-9a7972c3a75f55e976e620484f58410c920db091 + """ + return os.path.join(self._resources_dir, hash) + + def get_uris(self, runtime_env: "RuntimeEnv") -> List[str]: # noqa: F821 + """Return the conda URI from the RuntimeEnv if it exists, else return [].""" + conda_uri = runtime_env.conda_uri() + if conda_uri: + return [conda_uri] + return [] + + def delete_uri( + self, uri: str, logger: Optional[logging.Logger] = default_logger + ) -> int: + """Delete URI and return the number of bytes deleted.""" + logger.info(f"Got request to delete URI {uri}") + protocol, hash = parse_uri(uri) + if protocol != Protocol.CONDA: + raise ValueError( + "CondaPlugin can only delete URIs with protocol " + f"conda. Received protocol {protocol}, URI {uri}" + ) + + conda_env_path = self._get_path_from_hash(hash) + local_dir_size = get_directory_size_bytes(conda_env_path) + + with FileLock(self._installs_and_deletions_file_lock): + successful = delete_conda_env(prefix=conda_env_path, logger=logger) + if not successful: + logger.warning(f"Error when deleting conda env {conda_env_path}. ") + return 0 + + return local_dir_size + + async def create( + self, + uri: Optional[str], + runtime_env: "RuntimeEnv", # noqa: F821 + context: RuntimeEnvContext, + logger: logging.Logger = default_logger, + ) -> int: + if not runtime_env.has_conda(): + return 0 + + def _create(): + result = parse_and_validate_conda(runtime_env.get("conda")) + + if isinstance(result, str): + # The conda env name is given. + # In this case, we only verify if the given + # conda env exists. + + # If the env is already validated, do nothing. + if result in self._validated_named_conda_env: + return 0 + + conda_info = get_conda_info_json() + envs = get_conda_envs(conda_info) + + # We accept `result` as a conda name or full path. + if not any(result == env[0] or result == env[1] for env in envs): + raise ValueError( + f"The given conda environment '{result}' " + f"from the runtime env {runtime_env} doesn't " + "exist from the output of `conda info --json`. " + "You can only specify an env that already exists. " + f"Please make sure to create an env {result} " + ) + self._validated_named_conda_env.add(result) + return 0 + + logger.debug( + "Setting up conda for runtime_env: " f"{runtime_env.serialize()}" + ) + protocol, hash = parse_uri(uri) + conda_env_name = self._get_path_from_hash(hash) + + conda_dict = _get_conda_dict_with_ray_inserted(runtime_env, logger=logger) + + logger.info(f"Setting up conda environment with {runtime_env}") + with FileLock(self._installs_and_deletions_file_lock): + try: + conda_yaml_file = os.path.join( + self._resources_dir, "environment.yml" + ) + with open(conda_yaml_file, "w") as file: + yaml.dump(conda_dict, file) + create_conda_env_if_needed( + conda_yaml_file, prefix=conda_env_name, logger=logger + ) + finally: + os.remove(conda_yaml_file) + + if runtime_env.get_extension("_inject_current_ray"): + _inject_ray_to_conda_site(conda_path=conda_env_name, logger=logger) + logger.info(f"Finished creating conda environment at {conda_env_name}") + return get_directory_size_bytes(conda_env_name) + + loop = get_or_create_event_loop() + return await loop.run_in_executor(None, _create) + + def modify_context( + self, + uris: List[str], + runtime_env: "RuntimeEnv", # noqa: F821 + context: RuntimeEnvContext, + logger: Optional[logging.Logger] = default_logger, + ): + if not runtime_env.has_conda(): + return + + if runtime_env.conda_env_name(): + conda_env_name = runtime_env.conda_env_name() + else: + protocol, hash = parse_uri(runtime_env.conda_uri()) + conda_env_name = self._get_path_from_hash(hash) + context.py_executable = "python" + context.command_prefix += get_conda_activate_commands(conda_env_name) diff --git a/lib/python3.12/site-packages/ray/_private/runtime_env/conda_utils.py b/lib/python3.12/site-packages/ray/_private/runtime_env/conda_utils.py new file mode 100644 index 0000000000000000000000000000000000000000..eb2f67cabbe370b09d897a9303ac6cf719f28a37 --- /dev/null +++ b/lib/python3.12/site-packages/ray/_private/runtime_env/conda_utils.py @@ -0,0 +1,279 @@ +import hashlib +import json +import logging +import os +import shutil +import subprocess +from typing import List, Optional, Tuple, Union + +"""Utilities for conda. Adapted from https://github.com/mlflow/mlflow.""" + +# Name of environment variable indicating a path to a conda installation. Ray +# will default to running "conda" if unset. +RAY_CONDA_HOME = "RAY_CONDA_HOME" + +_WIN32 = os.name == "nt" + + +def get_conda_activate_commands(conda_env_name: str) -> List[str]: + """ + Get a list of commands to run to silently activate the given conda env. + """ + # Checking for newer conda versions + if not _WIN32 and ("CONDA_EXE" in os.environ or RAY_CONDA_HOME in os.environ): + conda_path = get_conda_bin_executable("conda") + activate_conda_env = [ + ".", + f"{os.path.dirname(conda_path)}/../etc/profile.d/conda.sh", + "&&", + ] + activate_conda_env += ["conda", "activate", conda_env_name] + + else: + activate_path = get_conda_bin_executable("activate") + if not _WIN32: + # Use bash command syntax + activate_conda_env = ["source", activate_path, conda_env_name] + else: + conda_path = get_conda_bin_executable("conda") + activate_conda_env = [conda_path, "activate", conda_env_name] + return activate_conda_env + ["1>&2", "&&"] + + +def get_conda_bin_executable(executable_name: str) -> str: + """ + Return path to the specified executable, assumed to be discoverable within + a conda installation. + + The conda home directory (expected to contain a 'bin' subdirectory on + linux) is configurable via the ``RAY_CONDA_HOME`` environment variable. If + ``RAY_CONDA_HOME`` is unspecified, try the ``CONDA_EXE`` environment + variable set by activating conda. If neither is specified, this method + returns `executable_name`. + """ + conda_home = os.environ.get(RAY_CONDA_HOME) + if conda_home: + if _WIN32: + candidate = os.path.join(conda_home, "%s.exe" % executable_name) + if os.path.exists(candidate): + return candidate + candidate = os.path.join(conda_home, "%s.bat" % executable_name) + if os.path.exists(candidate): + return candidate + else: + return os.path.join(conda_home, "bin/%s" % executable_name) + else: + conda_home = "." + # Use CONDA_EXE as per https://github.com/conda/conda/issues/7126 + if "CONDA_EXE" in os.environ: + conda_bin_dir = os.path.dirname(os.environ["CONDA_EXE"]) + if _WIN32: + candidate = os.path.join(conda_home, "%s.exe" % executable_name) + if os.path.exists(candidate): + return candidate + candidate = os.path.join(conda_home, "%s.bat" % executable_name) + if os.path.exists(candidate): + return candidate + else: + return os.path.join(conda_bin_dir, executable_name) + if _WIN32: + return executable_name + ".bat" + return executable_name + + +def _get_conda_env_name(conda_env_path: str) -> str: + conda_env_contents = open(conda_env_path).read() + return "ray-%s" % hashlib.sha1(conda_env_contents.encode("utf-8")).hexdigest() + + +def create_conda_env_if_needed( + conda_yaml_file: str, prefix: str, logger: Optional[logging.Logger] = None +) -> None: + """ + Given a conda YAML, creates a conda environment containing the required + dependencies if such a conda environment doesn't already exist. + Args: + conda_yaml_file: The path to a conda `environment.yml` file. + prefix: Directory to install the environment into via + the `--prefix` option to conda create. This also becomes the name + of the conda env; i.e. it can be passed into `conda activate` and + `conda remove` + """ + if logger is None: + logger = logging.getLogger(__name__) + + conda_path = get_conda_bin_executable("conda") + try: + exec_cmd([conda_path, "--help"], throw_on_error=False) + except (EnvironmentError, FileNotFoundError): + raise ValueError( + f"Could not find Conda executable at '{conda_path}'. " + "Ensure Conda is installed as per the instructions at " + "https://conda.io/projects/conda/en/latest/" + "user-guide/install/index.html. " + "You can also configure Ray to look for a specific " + f"Conda executable by setting the {RAY_CONDA_HOME} " + "environment variable to the path of the Conda executable." + ) + + _, stdout, _ = exec_cmd([conda_path, "env", "list", "--json"]) + envs = json.loads(stdout[stdout.index("{") :])["envs"] + + if prefix in envs: + logger.info(f"Conda environment {prefix} already exists.") + return + + create_cmd = [ + conda_path, + "env", + "create", + "--file", + conda_yaml_file, + "--prefix", + prefix, + ] + + logger.info(f"Creating conda environment {prefix}") + exit_code, output = exec_cmd_stream_to_logger(create_cmd, logger) + if exit_code != 0: + if os.path.exists(prefix): + shutil.rmtree(prefix) + raise RuntimeError( + f"Failed to install conda environment {prefix}:\nOutput:\n{output}" + ) + + +def delete_conda_env(prefix: str, logger: Optional[logging.Logger] = None) -> bool: + if logger is None: + logger = logging.getLogger(__name__) + + logger.info(f"Deleting conda environment {prefix}") + + conda_path = get_conda_bin_executable("conda") + delete_cmd = [conda_path, "remove", "-p", prefix, "--all", "-y"] + exit_code, output = exec_cmd_stream_to_logger(delete_cmd, logger) + + if exit_code != 0: + logger.debug(f"Failed to delete conda environment {prefix}:\n{output}") + return False + + return True + + +def get_conda_env_list() -> list: + """ + Get conda env list in full paths. + """ + conda_path = get_conda_bin_executable("conda") + try: + exec_cmd([conda_path, "--help"], throw_on_error=False) + except EnvironmentError: + raise ValueError(f"Could not find Conda executable at {conda_path}.") + _, stdout, _ = exec_cmd([conda_path, "env", "list", "--json"]) + envs = json.loads(stdout)["envs"] + return envs + + +def get_conda_info_json() -> dict: + """ + Get `conda info --json` output. + + Returns dict of conda info. See [1] for more details. We mostly care about these + keys: + + - `conda_prefix`: str The path to the conda installation. + - `envs`: List[str] absolute paths to conda environments. + + [1] https://github.com/conda/conda/blob/main/conda/cli/main_info.py + """ + conda_path = get_conda_bin_executable("conda") + try: + exec_cmd([conda_path, "--help"], throw_on_error=False) + except EnvironmentError: + raise ValueError(f"Could not find Conda executable at {conda_path}.") + _, stdout, _ = exec_cmd([conda_path, "info", "--json"]) + return json.loads(stdout) + + +def get_conda_envs(conda_info: dict) -> List[Tuple[str, str]]: + """ + Gets the conda environments, as a list of (name, path) tuples. + """ + prefix = conda_info["conda_prefix"] + ret = [] + for env in conda_info["envs"]: + if env == prefix: + ret.append(("base", env)) + else: + ret.append((os.path.basename(env), env)) + return ret + + +class ShellCommandException(Exception): + pass + + +def exec_cmd( + cmd: List[str], throw_on_error: bool = True, logger: Optional[logging.Logger] = None +) -> Union[int, Tuple[int, str, str]]: + """ + Runs a command as a child process. + + A convenience wrapper for running a command from a Python script. + + Note on the return value: A tuple of the exit code, + standard output and standard error is returned. + + Args: + cmd: the command to run, as a list of strings + throw_on_error: if true, raises an Exception if the exit code of the + program is nonzero + """ + child = subprocess.Popen( + cmd, + stdout=subprocess.PIPE, + stdin=subprocess.PIPE, + stderr=subprocess.PIPE, + universal_newlines=True, + ) + (stdout, stderr) = child.communicate() + exit_code = child.wait() + if throw_on_error and exit_code != 0: + raise ShellCommandException( + "Non-zero exit code: %s\n\nSTDOUT:\n%s\n\nSTDERR:%s" + % (exit_code, stdout, stderr) + ) + return exit_code, stdout, stderr + + +def exec_cmd_stream_to_logger( + cmd: List[str], logger: logging.Logger, n_lines: int = 50, **kwargs +) -> Tuple[int, str]: + """Runs a command as a child process, streaming output to the logger. + + The last n_lines lines of output are also returned (stdout and stderr). + """ + if "env" in kwargs and _WIN32 and "PATH" not in [x.upper() for x in kwargs.keys]: + raise ValueError("On windows, Popen requires 'PATH' in 'env'") + child = subprocess.Popen( + cmd, + universal_newlines=True, + stdout=subprocess.PIPE, + stderr=subprocess.STDOUT, + **kwargs, + ) + last_n_lines = [] + with child.stdout: + for line in iter(child.stdout.readline, b""): + exit_code = child.poll() + if exit_code is not None: + break + line = line.strip() + if not line: + continue + last_n_lines.append(line.strip()) + last_n_lines = last_n_lines[-n_lines:] + logger.info(line.strip()) + + exit_code = child.wait() + return exit_code, "\n".join(last_n_lines) diff --git a/lib/python3.12/site-packages/ray/_private/runtime_env/constants.py b/lib/python3.12/site-packages/ray/_private/runtime_env/constants.py new file mode 100644 index 0000000000000000000000000000000000000000..3c6096b5993ef73836771ade7966947681bbab89 --- /dev/null +++ b/lib/python3.12/site-packages/ray/_private/runtime_env/constants.py @@ -0,0 +1,28 @@ +# Env var set by job manager to pass runtime env and metadata to subprocess +RAY_JOB_CONFIG_JSON_ENV_VAR = "RAY_JOB_CONFIG_JSON_ENV_VAR" + +# The plugin config which should be loaded when ray cluster starts. +# It is a json formatted config, +# e.g. [{"class": "xxx.xxx.xxx_plugin", "priority": 10}]. +RAY_RUNTIME_ENV_PLUGINS_ENV_VAR = "RAY_RUNTIME_ENV_PLUGINS" + +# The field name of plugin class in the plugin config. +RAY_RUNTIME_ENV_CLASS_FIELD_NAME = "class" + +# The field name of priority in the plugin config. +RAY_RUNTIME_ENV_PRIORITY_FIELD_NAME = "priority" + +# The default priority of runtime env plugin. +RAY_RUNTIME_ENV_PLUGIN_DEFAULT_PRIORITY = 10 + +# The minimum priority of runtime env plugin. +RAY_RUNTIME_ENV_PLUGIN_MIN_PRIORITY = 0 + +# The maximum priority of runtime env plugin. +RAY_RUNTIME_ENV_PLUGIN_MAX_PRIORITY = 100 + +# The schema files or directories of plugins which should be loaded in workers. +RAY_RUNTIME_ENV_PLUGIN_SCHEMAS_ENV_VAR = "RAY_RUNTIME_ENV_PLUGIN_SCHEMAS" + +# The file suffix of runtime env plugin schemas. +RAY_RUNTIME_ENV_PLUGIN_SCHEMA_SUFFIX = ".json" diff --git a/lib/python3.12/site-packages/ray/_private/runtime_env/context.py b/lib/python3.12/site-packages/ray/_private/runtime_env/context.py new file mode 100644 index 0000000000000000000000000000000000000000..2cc58e28625d5a1e73bbbd62052e9fa92fbabecb --- /dev/null +++ b/lib/python3.12/site-packages/ray/_private/runtime_env/context.py @@ -0,0 +1,108 @@ +import json +import logging +import os +import shlex +import subprocess +import sys +from typing import Dict, List, Optional + +from ray._private.services import get_ray_jars_dir +from ray._private.utils import update_envs +from ray.core.generated.common_pb2 import Language +from ray.util.annotations import DeveloperAPI + +logger = logging.getLogger(__name__) + + +@DeveloperAPI +class RuntimeEnvContext: + """A context used to describe the created runtime env.""" + + def __init__( + self, + command_prefix: List[str] = None, + env_vars: Dict[str, str] = None, + py_executable: Optional[str] = None, + override_worker_entrypoint: Optional[str] = None, + java_jars: List[str] = None, + ): + self.command_prefix = command_prefix or [] + self.env_vars = env_vars or {} + self.py_executable = py_executable or sys.executable + self.override_worker_entrypoint: Optional[str] = override_worker_entrypoint + self.java_jars = java_jars or [] + + def serialize(self) -> str: + return json.dumps(self.__dict__) + + @staticmethod + def deserialize(json_string): + return RuntimeEnvContext(**json.loads(json_string)) + + def exec_worker(self, passthrough_args: List[str], language: Language): + update_envs(self.env_vars) + + if language == Language.PYTHON and sys.platform == "win32": + executable = [self.py_executable] + elif language == Language.PYTHON: + executable = ["exec", self.py_executable] + elif language == Language.JAVA: + executable = ["java"] + ray_jars = os.path.join(get_ray_jars_dir(), "*") + + local_java_jars = [] + for java_jar in self.java_jars: + local_java_jars.append(f"{java_jar}/*") + local_java_jars.append(java_jar) + + class_path_args = ["-cp", ray_jars + ":" + str(":".join(local_java_jars))] + passthrough_args = class_path_args + passthrough_args + elif sys.platform == "win32": + executable = [] + else: + executable = ["exec"] + + # By default, raylet uses the path to default_worker.py on host. + # However, the path to default_worker.py inside the container + # can be different. We need the user to specify the path to + # default_worker.py inside the container. + if self.override_worker_entrypoint: + logger.debug( + f"Changing the worker entrypoint from {passthrough_args[0]} to " + f"{self.override_worker_entrypoint}." + ) + passthrough_args[0] = self.override_worker_entrypoint + + if sys.platform == "win32": + + def quote(s): + s = s.replace("&", "%26") + return s + + passthrough_args = [quote(s) for s in passthrough_args] + + cmd = [*self.command_prefix, *executable, *passthrough_args] + logger.debug(f"Exec'ing worker with command: {cmd}") + subprocess.Popen(cmd, shell=True).wait() + else: + # We use shlex to do the necessary shell escape + # of special characters in passthrough_args. + passthrough_args = [shlex.quote(s) for s in passthrough_args] + cmd = [*self.command_prefix, *executable, *passthrough_args] + # TODO(SongGuyang): We add this env to command for macOS because it doesn't + # work for the C++ process of `os.execvp`. We should find a better way to + # fix it. + MACOS_LIBRARY_PATH_ENV_NAME = "DYLD_LIBRARY_PATH" + if MACOS_LIBRARY_PATH_ENV_NAME in os.environ: + cmd.insert( + 0, + f"{MACOS_LIBRARY_PATH_ENV_NAME}=" + f"{os.environ[MACOS_LIBRARY_PATH_ENV_NAME]}", + ) + logger.debug(f"Exec'ing worker with command: {cmd}") + # PyCharm will monkey patch the os.execvp at + # .pycharm_helpers/pydev/_pydev_bundle/pydev_monkey.py + # The monkey patched os.execvp function has a different + # signature. So, we use os.execvp("executable", args=[]) + # instead of os.execvp(file="executable", args=[]) + os.execvp("bash", args=["bash", "-c", " ".join(cmd)]) diff --git a/lib/python3.12/site-packages/ray/_private/runtime_env/default_impl.py b/lib/python3.12/site-packages/ray/_private/runtime_env/default_impl.py new file mode 100644 index 0000000000000000000000000000000000000000..f0d1567530af82ac3f1c3a8d487a96361dae94e7 --- /dev/null +++ b/lib/python3.12/site-packages/ray/_private/runtime_env/default_impl.py @@ -0,0 +1,5 @@ +from ray._private.runtime_env.image_uri import ImageURIPlugin + + +def get_image_uri_plugin_cls(): + return ImageURIPlugin diff --git a/lib/python3.12/site-packages/ray/_private/runtime_env/dependency_utils.py b/lib/python3.12/site-packages/ray/_private/runtime_env/dependency_utils.py new file mode 100644 index 0000000000000000000000000000000000000000..d208a067c02da00906d7387478ccfbb088018d6d --- /dev/null +++ b/lib/python3.12/site-packages/ray/_private/runtime_env/dependency_utils.py @@ -0,0 +1,118 @@ +"""Util functions to manage dependency requirements.""" + +import logging +import os +import tempfile +from contextlib import asynccontextmanager +from typing import List, Optional, Tuple + +from ray._private.runtime_env import virtualenv_utils +from ray._private.runtime_env.utils import check_output_cmd + +INTERNAL_PIP_FILENAME = "ray_runtime_env_internal_pip_requirements.txt" +MAX_INTERNAL_PIP_FILENAME_TRIES = 100 + + +def gen_requirements_txt(requirements_file: str, pip_packages: List[str]): + """Dump [pip_packages] to the given [requirements_file] for later env setup.""" + with open(requirements_file, "w") as file: + for line in pip_packages: + file.write(line + "\n") + + +@asynccontextmanager +async def check_ray(python: str, cwd: str, logger: logging.Logger): + """A context manager to check ray is not overwritten. + + Currently, we only check ray version and path. It works for virtualenv, + - ray is in Python's site-packages. + - ray is overwritten during yield. + - ray is in virtualenv's site-packages. + """ + + async def _get_ray_version_and_path() -> Tuple[str, str]: + with tempfile.TemporaryDirectory( + prefix="check_ray_version_tempfile" + ) as tmp_dir: + ray_version_path = os.path.join(tmp_dir, "ray_version.txt") + check_ray_cmd = [ + python, + "-c", + """ +import ray +with open(r"{ray_version_path}", "wt") as f: + f.write(ray.__version__) + f.write(" ") + f.write(ray.__path__[0]) + """.format( + ray_version_path=ray_version_path + ), + ] + if virtualenv_utils._WIN32: + env = os.environ.copy() + else: + env = {} + output = await check_output_cmd( + check_ray_cmd, logger=logger, cwd=cwd, env=env + ) + logger.info(f"try to write ray version information in: {ray_version_path}") + with open(ray_version_path, "rt") as f: + output = f.read() + # print after import ray may have  endings, so we strip them by *_ + ray_version, ray_path, *_ = [s.strip() for s in output.split()] + return ray_version, ray_path + + version, path = await _get_ray_version_and_path() + yield + actual_version, actual_path = await _get_ray_version_and_path() + if actual_version != version: + raise RuntimeError( + "Changing the ray version is not allowed: \n" + f" current version: {actual_version}, " + f" expect version: {version}, " + f" current path: {actual_path}, " + f" expect path: {path}, " + "Please ensure the dependencies in the runtime_env pip field " + "do not install a different version of Ray." + ) + if actual_path != path: + logger.info( + f"Detected new Ray package with the same version at {actual_path} (vs system {path})." + ) + + +def get_requirements_file(target_dir: str, pip_list: Optional[List[str]]) -> str: + """Returns the path to the requirements file to use for this runtime env. + + If pip_list is not None, we will check if the internal pip filename is in any of + the entries of pip_list. If so, we will append numbers to the end of the + filename until we find one that doesn't conflict. This prevents infinite + recursion if the user specifies the internal pip filename in their pip list. + + Args: + target_dir: The directory to store the requirements file in. + pip_list: A list of pip requirements specified by the user. + + Returns: + The path to the requirements file to use for this runtime env. + """ + + def filename_in_pip_list(filename: str) -> bool: + for pip_entry in pip_list: + if filename in pip_entry: + return True + return False + + filename = INTERNAL_PIP_FILENAME + if pip_list is not None: + i = 1 + while filename_in_pip_list(filename) and i < MAX_INTERNAL_PIP_FILENAME_TRIES: + filename = f"{INTERNAL_PIP_FILENAME}.{i}" + i += 1 + if i == MAX_INTERNAL_PIP_FILENAME_TRIES: + raise RuntimeError( + "Could not find a valid filename for the internal " + "pip requirements file. Please specify a different " + "pip list in your runtime env." + ) + return os.path.join(target_dir, filename) diff --git a/lib/python3.12/site-packages/ray/_private/runtime_env/image_uri.py b/lib/python3.12/site-packages/ray/_private/runtime_env/image_uri.py new file mode 100644 index 0000000000000000000000000000000000000000..cafed9a6f46573525c948ed00969fc47aa5cb3f4 --- /dev/null +++ b/lib/python3.12/site-packages/ray/_private/runtime_env/image_uri.py @@ -0,0 +1,229 @@ +import asyncio +import logging +import os +import tempfile +from typing import List, Optional + +from ray._private.runtime_env.context import RuntimeEnvContext +from ray._private.runtime_env.plugin import RuntimeEnvPlugin + +default_logger = logging.getLogger(__name__) + + +async def _create_impl(image_uri: str, logger: logging.Logger): + # Pull image if it doesn't exist + # Also get path to `default_worker.py` inside the image. + with tempfile.TemporaryDirectory() as tmpdir: + os.chmod(tmpdir, 0o777) + result_file = os.path.join(tmpdir, "worker_path.txt") + get_worker_path_script = """ +import ray._private.workers.default_worker as dw +with open('/shared/worker_path.txt', 'w') as f: + f.write(dw.__file__) +""" + cmd = [ + "podman", + "run", + "--rm", + "-v", + f"{tmpdir}:/shared:Z", + image_uri, + "python", + "-c", + get_worker_path_script, + ] + + logger.info("Pulling image %s", image_uri) + + process = await asyncio.create_subprocess_exec( + *cmd, stdout=asyncio.subprocess.PIPE, stderr=asyncio.subprocess.PIPE + ) + + stdout, stderr = await process.communicate() + + if process.returncode != 0: + raise RuntimeError( + f"Podman command failed: cmd={cmd}, returncode={process.returncode}, stdout={stdout.decode()}, stderr={stderr.decode()}" + ) + + if not os.path.exists(result_file): + raise FileNotFoundError( + f"Worker path file not created when getting worker path for image {image_uri}" + ) + + with open(result_file, "r") as f: + worker_path = f.read().strip() + + if not worker_path.endswith(".py"): + raise ValueError( + f"Invalid worker path inferred in image {image_uri}: {worker_path}" + ) + + logger.info(f"Inferred worker path in image {image_uri}: {worker_path}") + return worker_path + + +def _modify_context_impl( + image_uri: str, + worker_path: str, + run_options: Optional[List[str]], + context: RuntimeEnvContext, + logger: logging.Logger, + ray_tmp_dir: str, +): + context.override_worker_entrypoint = worker_path + + container_driver = "podman" + container_command = [ + container_driver, + "run", + "-v", + ray_tmp_dir + ":" + ray_tmp_dir, + "--cgroup-manager=cgroupfs", + "--network=host", + "--pid=host", + "--ipc=host", + # NOTE(zcin): Mounted volumes in rootless containers are + # owned by the user `root`. The user on host (which will + # usually be `ray` if this is being run in a ray docker + # image) who started the container is mapped using user + # namespaces to the user `root` in a rootless container. In + # order for the Ray Python worker to access the mounted ray + # tmp dir, we need to use keep-id mode which maps the user + # as itself (instead of as `root`) into the container. + # https://www.redhat.com/sysadmin/rootless-podman-user-namespace-modes + "--userns=keep-id", + ] + + # Environment variables to set in container + env_vars = dict() + + # Propagate all host environment variables that have the prefix "RAY_" + # This should include RAY_RAYLET_PID + for env_var_name, env_var_value in os.environ.items(): + if env_var_name.startswith("RAY_"): + env_vars[env_var_name] = env_var_value + + # Support for runtime_env['env_vars'] + env_vars.update(context.env_vars) + + # Set environment variables + for env_var_name, env_var_value in env_vars.items(): + container_command.append("--env") + container_command.append(f"{env_var_name}='{env_var_value}'") + + # The RAY_JOB_ID environment variable is needed for the default worker. + # It won't be set at the time setup() is called, but it will be set + # when worker command is executed, so we use RAY_JOB_ID=$RAY_JOB_ID + # for the container start command + container_command.append("--env") + container_command.append("RAY_JOB_ID=$RAY_JOB_ID") + + if run_options: + container_command.extend(run_options) + # TODO(chenk008): add resource limit + container_command.append("--entrypoint") + container_command.append("python") + container_command.append(image_uri) + + # Example: + # podman run -v /tmp/ray:/tmp/ray + # --cgroup-manager=cgroupfs --network=host --pid=host --ipc=host + # --userns=keep-id --env RAY_RAYLET_PID=23478 --env RAY_JOB_ID=$RAY_JOB_ID + # --entrypoint python rayproject/ray:nightly-py39 + container_command_str = " ".join(container_command) + logger.info(f"Starting worker in container with prefix {container_command_str}") + + context.py_executable = container_command_str + + +class ImageURIPlugin(RuntimeEnvPlugin): + """Starts worker in a container of a custom image.""" + + name = "image_uri" + + @staticmethod + def get_compatible_keys(): + return {"image_uri", "config", "env_vars"} + + def __init__(self, ray_tmp_dir: str): + self._ray_tmp_dir = ray_tmp_dir + + async def create( + self, + uri: Optional[str], + runtime_env: "RuntimeEnv", # noqa: F821 + context: RuntimeEnvContext, + logger: logging.Logger, + ) -> float: + if not runtime_env.image_uri(): + return + + self.worker_path = await _create_impl(runtime_env.image_uri(), logger) + + def modify_context( + self, + uris: List[str], + runtime_env: "RuntimeEnv", # noqa: F821 + context: RuntimeEnvContext, + logger: Optional[logging.Logger] = default_logger, + ): + if not runtime_env.image_uri(): + return + + _modify_context_impl( + runtime_env.image_uri(), + self.worker_path, + [], + context, + logger, + self._ray_tmp_dir, + ) + + +class ContainerPlugin(RuntimeEnvPlugin): + """Starts worker in container.""" + + name = "container" + + def __init__(self, ray_tmp_dir: str): + self._ray_tmp_dir = ray_tmp_dir + + async def create( + self, + uri: Optional[str], + runtime_env: "RuntimeEnv", # noqa: F821 + context: RuntimeEnvContext, + logger: logging.Logger, + ) -> float: + if not runtime_env.has_py_container() or not runtime_env.py_container_image(): + return + + self.worker_path = await _create_impl(runtime_env.py_container_image(), logger) + + def modify_context( + self, + uris: List[str], + runtime_env: "RuntimeEnv", # noqa: F821 + context: RuntimeEnvContext, + logger: Optional[logging.Logger] = default_logger, + ): + if not runtime_env.has_py_container() or not runtime_env.py_container_image(): + return + + if runtime_env.py_container_worker_path(): + logger.warning( + "You are using `container.worker_path`, but the path to " + "`default_worker.py` is now automatically detected from the image. " + "`container.worker_path` is deprecated and will be removed in future " + "versions." + ) + + _modify_context_impl( + runtime_env.py_container_image(), + runtime_env.py_container_worker_path() or self.worker_path, + runtime_env.py_container_run_options(), + context, + logger, + self._ray_tmp_dir, + ) diff --git a/lib/python3.12/site-packages/ray/_private/runtime_env/java_jars.py b/lib/python3.12/site-packages/ray/_private/runtime_env/java_jars.py new file mode 100644 index 0000000000000000000000000000000000000000..83044ba7054b8a53643d8794a62ed2b2ecf97e17 --- /dev/null +++ b/lib/python3.12/site-packages/ray/_private/runtime_env/java_jars.py @@ -0,0 +1,104 @@ +import logging +import os +from typing import Dict, List, Optional + +from ray._common.utils import try_to_create_directory +from ray._private.runtime_env.context import RuntimeEnvContext +from ray._private.runtime_env.packaging import ( + delete_package, + download_and_unpack_package, + get_local_dir_from_uri, + is_jar_uri, +) +from ray._private.runtime_env.plugin import RuntimeEnvPlugin +from ray._private.utils import get_directory_size_bytes +from ray._raylet import GcsClient +from ray.exceptions import RuntimeEnvSetupError + +default_logger = logging.getLogger(__name__) + + +class JavaJarsPlugin(RuntimeEnvPlugin): + + name = "java_jars" + + def __init__(self, resources_dir: str, gcs_client: GcsClient): + self._resources_dir = os.path.join(resources_dir, "java_jars_files") + self._gcs_client = gcs_client + try_to_create_directory(self._resources_dir) + + def _get_local_dir_from_uri(self, uri: str): + return get_local_dir_from_uri(uri, self._resources_dir) + + def delete_uri( + self, uri: str, logger: Optional[logging.Logger] = default_logger + ) -> int: + """Delete URI and return the number of bytes deleted.""" + local_dir = get_local_dir_from_uri(uri, self._resources_dir) + local_dir_size = get_directory_size_bytes(local_dir) + + deleted = delete_package(uri, self._resources_dir) + if not deleted: + logger.warning(f"Tried to delete nonexistent URI: {uri}.") + return 0 + + return local_dir_size + + def get_uris(self, runtime_env: dict) -> List[str]: + return runtime_env.java_jars() + + async def _download_jars( + self, uri: str, logger: Optional[logging.Logger] = default_logger + ): + """Download a jar URI.""" + try: + jar_file = await download_and_unpack_package( + uri, self._resources_dir, self._gcs_client, logger=logger + ) + except Exception as e: + raise RuntimeEnvSetupError( + "Failed to download jar file: {}".format(e) + ) from e + module_dir = self._get_local_dir_from_uri(uri) + logger.debug(f"Succeeded to download jar file {jar_file} .") + return module_dir + + async def create( + self, + uri: str, + runtime_env: "RuntimeEnv", # noqa: F821 + context: RuntimeEnvContext, + logger: Optional[logging.Logger] = default_logger, + ) -> int: + if not uri: + return 0 + if is_jar_uri(uri): + module_dir = await self._download_jars(uri=uri, logger=logger) + else: + try: + module_dir = await download_and_unpack_package( + uri, self._resources_dir, self._gcs_client, logger=logger + ) + except Exception as e: + raise RuntimeEnvSetupError( + "Failed to download jar file: {}".format(e) + ) from e + + return get_directory_size_bytes(module_dir) + + def modify_context( + self, + uris: List[str], + runtime_env_dict: Dict, + context: RuntimeEnvContext, + logger: Optional[logging.Logger] = default_logger, + ): + for uri in uris: + module_dir = self._get_local_dir_from_uri(uri) + if not module_dir.exists(): + raise ValueError( + f"Local directory {module_dir} for URI {uri} does " + "not exist on the cluster. Something may have gone wrong while " + "downloading, unpacking or installing the java jar files." + ) + context.java_jars.append(str(module_dir)) diff --git a/lib/python3.12/site-packages/ray/_private/runtime_env/nsight.py b/lib/python3.12/site-packages/ray/_private/runtime_env/nsight.py new file mode 100644 index 0000000000000000000000000000000000000000..bd1c44beabd7bf7383baeb3abdf80e79b2ef3df3 --- /dev/null +++ b/lib/python3.12/site-packages/ray/_private/runtime_env/nsight.py @@ -0,0 +1,149 @@ +import asyncio +import copy +import logging +import os +import subprocess +import sys +from pathlib import Path +from typing import Dict, List, Optional, Tuple + +from ray._common.utils import ( + try_to_create_directory, +) +from ray._private.runtime_env.context import RuntimeEnvContext +from ray._private.runtime_env.plugin import RuntimeEnvPlugin +from ray.exceptions import RuntimeEnvSetupError + +default_logger = logging.getLogger(__name__) + +# Nsight options used when runtime_env={"_nsight": "default"} +NSIGHT_DEFAULT_CONFIG = { + "t": "cuda,cudnn,cublas,nvtx", + "o": "'worker_process_%p'", + "stop-on-exit": "true", +} + + +def parse_nsight_config(nsight_config: Dict[str, str]) -> List[str]: + """ + Function to convert dictionary of nsight options into + nsight command line + + The function returns: + - List[str]: nsys profile cmd line split into list of str + """ + nsight_cmd = ["nsys", "profile"] + for option, option_val in nsight_config.items(): + # option standard based on + # https://www.gnu.org/software/libc/manual/html_node/Argument-Syntax.html + if len(option) > 1: + nsight_cmd.append(f"--{option}={option_val}") + else: + nsight_cmd += [f"-{option}", option_val] + return nsight_cmd + + +class NsightPlugin(RuntimeEnvPlugin): + name = "_nsight" + + def __init__(self, resources_dir: str): + self.nsight_cmd = [] + + # replace this with better way to get logs dir + session_dir, runtime_dir = os.path.split(resources_dir) + self._nsight_dir = Path(session_dir) / "logs" / "nsight" + try_to_create_directory(self._nsight_dir) + + async def _check_nsight_script( + self, nsight_config: Dict[str, str] + ) -> Tuple[bool, str]: + """ + Function to validate if nsight_config is a valid nsight profile options + Args: + nsight_config: dictionary mapping nsight option to it's value + Returns: + a tuple consists of a boolean indicating if the nsight_config + is valid option and an error message if the nsight_config is invalid + """ + + # use empty as nsight report test filename + nsight_config_copy = copy.deepcopy(nsight_config) + nsight_config_copy["o"] = str(Path(self._nsight_dir) / "empty") + nsight_cmd = parse_nsight_config(nsight_config_copy) + try: + nsight_cmd = nsight_cmd + [sys.executable, "-c", '""'] + process = await asyncio.create_subprocess_exec( + *nsight_cmd, + stdout=subprocess.PIPE, + stderr=subprocess.PIPE, + ) + stdout, stderr = await process.communicate() + error_msg = stderr.strip() if stderr.strip() != "" else stdout.strip() + + # cleanup test.nsys-rep file + clean_up_cmd = ["rm", f"{nsight_config_copy['o']}.nsys-rep"] + cleanup_process = await asyncio.create_subprocess_exec( + *clean_up_cmd, + stdout=subprocess.PIPE, + stderr=subprocess.PIPE, + ) + _, _ = await cleanup_process.communicate() + if process.returncode == 0: + return True, None + else: + return False, error_msg + except FileNotFoundError: + return False, ("nsight is not installed") + + async def create( + self, + uri: Optional[str], + runtime_env: "RuntimeEnv", # noqa: F821 + context: RuntimeEnvContext, + logger: logging.Logger = default_logger, + ) -> int: + nsight_config = runtime_env.nsight() + if not nsight_config: + return 0 + + if nsight_config and sys.platform != "linux": + raise RuntimeEnvSetupError( + "Nsight CLI is only available in Linux.\n" + "More information can be found in " + "https://docs.nvidia.com/nsight-compute/NsightComputeCli/index.html" + ) + + if isinstance(nsight_config, str): + if nsight_config == "default": + nsight_config = NSIGHT_DEFAULT_CONFIG + else: + raise RuntimeEnvSetupError( + f"Unsupported nsight config: {nsight_config}. " + "The supported config is 'default' or " + "Dictionary of nsight options" + ) + + is_valid_nsight_cmd, error_msg = await self._check_nsight_script(nsight_config) + if not is_valid_nsight_cmd: + logger.warning(error_msg) + raise RuntimeEnvSetupError( + "nsight profile failed to run with the following " + f"error message:\n {error_msg}" + ) + # add set output path to logs dir + nsight_config["o"] = str( + Path(self._nsight_dir) / nsight_config.get("o", NSIGHT_DEFAULT_CONFIG["o"]) + ) + + self.nsight_cmd = parse_nsight_config(nsight_config) + return 0 + + def modify_context( + self, + uris: List[str], + runtime_env: "RuntimeEnv", # noqa: F821 + context: RuntimeEnvContext, + logger: Optional[logging.Logger] = default_logger, + ): + logger.info("Running nsight profiler") + context.py_executable = " ".join(self.nsight_cmd) + " python" diff --git a/lib/python3.12/site-packages/ray/_private/runtime_env/packaging.py b/lib/python3.12/site-packages/ray/_private/runtime_env/packaging.py new file mode 100644 index 0000000000000000000000000000000000000000..d2366da4389e5c8c870c9edf2ecac3fad123c237 --- /dev/null +++ b/lib/python3.12/site-packages/ray/_private/runtime_env/packaging.py @@ -0,0 +1,970 @@ +import asyncio +import hashlib +import logging +import os +import shutil +import time +from pathlib import Path +from tempfile import TemporaryDirectory +from typing import Callable, List, Optional, Tuple +from urllib.parse import urlparse +from zipfile import ZipFile + +from filelock import FileLock + +from ray._private.path_utils import is_path +from ray._private.ray_constants import ( + GRPC_CPP_MAX_MESSAGE_SIZE, + RAY_RUNTIME_ENV_IGNORE_GITIGNORE, + RAY_RUNTIME_ENV_URI_PIN_EXPIRATION_S_DEFAULT, + RAY_RUNTIME_ENV_URI_PIN_EXPIRATION_S_ENV_VAR, +) +from ray._private.runtime_env.conda_utils import exec_cmd_stream_to_logger +from ray._private.runtime_env.protocol import Protocol +from ray._private.thirdparty.pathspec import PathSpec +from ray._raylet import GcsClient +from ray.experimental.internal_kv import ( + _internal_kv_exists, + _internal_kv_put, + _pin_runtime_env_uri, +) +from ray.util.annotations import DeveloperAPI + +default_logger = logging.getLogger(__name__) + +# If an individual file is beyond this size, print a warning. +FILE_SIZE_WARNING = 10 * 1024 * 1024 # 10MiB +# The size is bounded by the max gRPC message size. +# Keep in sync with max_grpc_message_size in ray_config_def.h. +GCS_STORAGE_MAX_SIZE = int( + os.environ.get("RAY_max_grpc_message_size", GRPC_CPP_MAX_MESSAGE_SIZE) +) +RAY_PKG_PREFIX = "_ray_pkg_" + +RAY_RUNTIME_ENV_FAIL_UPLOAD_FOR_TESTING_ENV_VAR = ( + "RAY_RUNTIME_ENV_FAIL_UPLOAD_FOR_TESTING" +) +RAY_RUNTIME_ENV_FAIL_DOWNLOAD_FOR_TESTING_ENV_VAR = ( + "RAY_RUNTIME_ENV_FAIL_DOWNLOAD_FOR_TESTING" +) + +# The name of the hidden top-level directory that appears when files are +# zipped on MacOS. +MAC_OS_ZIP_HIDDEN_DIR_NAME = "__MACOSX" + + +def _mib_string(num_bytes: float) -> str: + size_mib = float(num_bytes / 1024**2) + return f"{size_mib:.2f}MiB" + + +class _AsyncFileLock: + """Asyncio version used to prevent blocking event loop.""" + + def __init__(self, lock_file: str): + self.file = FileLock(lock_file) + + async def __aenter__(self): + while True: + try: + self.file.acquire(timeout=0) + return + except TimeoutError: + await asyncio.sleep(0.1) + + async def __aexit__(self, exc_type, exc, tb): + self.file.release() + + +def _xor_bytes(left: bytes, right: bytes) -> bytes: + if left and right: + return bytes(a ^ b for (a, b) in zip(left, right)) + return left or right + + +def _dir_travel( + path: Path, + excludes: List[Callable], + handler: Callable, + logger: Optional[logging.Logger] = default_logger, +): + """Travels the path recursively, calling the handler on each subpath. + + Respects excludes, which will be called to check if this path is skipped. + """ + e = _get_gitignore(path) + + if e is not None: + excludes.append(e) + + skip = any(e(path) for e in excludes) + if not skip: + try: + handler(path) + except Exception as e: + logger.error(f"Issue with path: {path}") + raise e + if path.is_dir(): + for sub_path in path.iterdir(): + _dir_travel(sub_path, excludes, handler, logger=logger) + + if e is not None: + excludes.pop() + + +def _hash_file_content_or_directory_name( + filepath: Path, + relative_path: Path, + logger: Optional[logging.Logger] = default_logger, +) -> bytes: + """Helper function to create hash of a single file or directory. + + This function hashes the path of the file or directory, + and if it's a file, then it hashes its content too. + """ + + BUF_SIZE = 4096 * 1024 + + sha1 = hashlib.sha1() + sha1.update(str(filepath.relative_to(relative_path)).encode()) + if not filepath.is_dir(): + try: + f = filepath.open("rb") + except Exception as e: + logger.debug( + f"Skipping contents of file {filepath} when calculating package hash " + f"because the file couldn't be opened: {e}" + ) + else: + try: + data = f.read(BUF_SIZE) + while len(data) != 0: + sha1.update(data) + data = f.read(BUF_SIZE) + finally: + f.close() + + return sha1.digest() + + +def _hash_file( + filepath: Path, + relative_path: Path, + logger: Optional[logging.Logger] = default_logger, +) -> bytes: + """Helper function to create hash of a single file. + + It hashes the path of the file and its content to create a hash value. + """ + file_hash = _hash_file_content_or_directory_name( + filepath, relative_path, logger=logger + ) + return _xor_bytes(file_hash, b"0" * 8) + + +def _hash_directory( + root: Path, + relative_path: Path, + excludes: Optional[Callable], + logger: Optional[logging.Logger] = default_logger, +) -> bytes: + """Helper function to create hash of a directory. + + It'll go through all the files in the directory and xor + hash(file_name, file_content) to create a hash value. + """ + hash_val = b"0" * 8 + + def handler(path: Path): + file_hash = _hash_file_content_or_directory_name( + path, relative_path, logger=logger + ) + nonlocal hash_val + hash_val = _xor_bytes(hash_val, file_hash) + + excludes = [] if excludes is None else [excludes] + _dir_travel(root, excludes, handler, logger=logger) + return hash_val + + +def parse_path(pkg_path: str) -> None: + """Parse the path to check it is well-formed and exists.""" + path = Path(pkg_path) + try: + path.resolve(strict=True) + except OSError: + raise ValueError(f"{path} is not a valid path.") + + +def parse_uri(pkg_uri: str) -> Tuple[Protocol, str]: + """ + Parse package uri into protocol and package name based on its format. + Note that the output of this function is not for handling actual IO, it's + only for setting up local directory folders by using package name as path. + + >>> parse_uri("https://test.com/file.zip") + (, 'https_test_com_file.zip') + + >>> parse_uri("https://test.com/file.whl") + (, 'file.whl') + + """ + if is_path(pkg_uri): + raise ValueError(f"Expected URI but received path {pkg_uri}") + + uri = urlparse(pkg_uri) + try: + protocol = Protocol(uri.scheme) + except ValueError as e: + raise ValueError( + f'Invalid protocol for runtime_env URI "{pkg_uri}". ' + f"Supported protocols: {Protocol._member_names_}. Original error: {e}" + ) + + if protocol in Protocol.remote_protocols(): + if uri.path.endswith(".whl"): + # Don't modify the .whl filename. See + # https://peps.python.org/pep-0427/#file-name-convention + # for more information. + package_name = uri.path.split("/")[-1] + else: + package_name = f"{protocol.value}_{uri.netloc}{uri.path}" + + disallowed_chars = ["/", ":", "@", "+", " ", "(", ")"] + for disallowed_char in disallowed_chars: + package_name = package_name.replace(disallowed_char, "_") + + # Remove all periods except the last, which is part of the + # file extension + package_name = package_name.replace(".", "_", package_name.count(".") - 1) + else: + package_name = uri.netloc + return (protocol, package_name) + + +def is_zip_uri(uri: str) -> bool: + try: + protocol, path = parse_uri(uri) + except ValueError: + return False + + return Path(path).suffix == ".zip" + + +def is_whl_uri(uri: str) -> bool: + try: + _, path = parse_uri(uri) + except ValueError: + return False + + return Path(path).suffix == ".whl" + + +def is_jar_uri(uri: str) -> bool: + try: + _, path = parse_uri(uri) + except ValueError: + return False + + return Path(path).suffix == ".jar" + + +def _get_excludes(path: Path, excludes: List[str]) -> Callable: + path = path.absolute() + pathspec = PathSpec.from_lines("gitwildmatch", excludes) + + def match(p: Path): + path_str = str(p.absolute().relative_to(path)) + return pathspec.match_file(path_str) + + return match + + +def _get_gitignore(path: Path) -> Optional[Callable]: + """Returns a function that returns True if the path should be excluded. + + Returns None if there is no .gitignore file in the path, or if the + RAY_RUNTIME_ENV_IGNORE_GITIGNORE environment variable is set to 1. + + Args: + path: The path to the directory to check for a .gitignore file. + + Returns: + A function that returns True if the path should be excluded. + """ + ignore_gitignore = os.environ.get(RAY_RUNTIME_ENV_IGNORE_GITIGNORE, "0") == "1" + if ignore_gitignore: + return None + + path = path.absolute() + ignore_file = path / ".gitignore" + if ignore_file.is_file(): + with ignore_file.open("r") as f: + pathspec = PathSpec.from_lines("gitwildmatch", f.readlines()) + + def match(p: Path): + path_str = str(p.absolute().relative_to(path)) + return pathspec.match_file(path_str) + + return match + else: + return None + + +def pin_runtime_env_uri(uri: str, *, expiration_s: Optional[int] = None) -> None: + """Pin a reference to a runtime_env URI in the GCS on a timeout. + + This is used to avoid premature eviction in edge conditions for job + reference counting. See https://github.com/ray-project/ray/pull/24719. + + Packages are uploaded to GCS in order to be downloaded by a runtime env plugin + (e.g. working_dir, py_modules) after the job starts. + + This function adds a temporary reference to the package in the GCS to prevent + it from being deleted before the job starts. (See #23423 for the bug where + this happened.) + + If this reference didn't have an expiration, then if the script exited + (e.g. via Ctrl-C) before the job started, the reference would never be + removed, so the package would never be deleted. + """ + + if expiration_s is None: + expiration_s = int( + os.environ.get( + RAY_RUNTIME_ENV_URI_PIN_EXPIRATION_S_ENV_VAR, + RAY_RUNTIME_ENV_URI_PIN_EXPIRATION_S_DEFAULT, + ) + ) + elif not isinstance(expiration_s, int): + raise ValueError(f"expiration_s must be an int, got {type(expiration_s)}.") + + if expiration_s < 0: + raise ValueError(f"expiration_s must be >= 0, got {expiration_s}.") + elif expiration_s > 0: + _pin_runtime_env_uri(uri, expiration_s=expiration_s) + + +def _store_package_in_gcs( + pkg_uri: str, + data: bytes, + logger: Optional[logging.Logger] = default_logger, +) -> int: + """Stores package data in the Global Control Store (GCS). + + Args: + pkg_uri: The GCS key to store the data in. + data: The serialized package's bytes to store in the GCS. + logger (Optional[logging.Logger]): The logger used by this function. + + Return: + int: Size of data + + Raises: + RuntimeError: If the upload to the GCS fails. + ValueError: If the data's size exceeds GCS_STORAGE_MAX_SIZE. + """ + + file_size = len(data) + size_str = _mib_string(file_size) + if len(data) >= GCS_STORAGE_MAX_SIZE: + raise ValueError( + f"Package size ({size_str}) exceeds the maximum size of " + f"{_mib_string(GCS_STORAGE_MAX_SIZE)}. You can exclude large " + "files using the 'excludes' option to the runtime_env or provide " + "a remote URI of a zip file using protocols such as 's3://', " + "'https://' and so on, refer to " + "https://docs.ray.io/en/latest/ray-core/handling-dependencies.html#api-reference." # noqa + ) + + logger.info(f"Pushing file package '{pkg_uri}' ({size_str}) to Ray cluster...") + try: + if os.environ.get(RAY_RUNTIME_ENV_FAIL_UPLOAD_FOR_TESTING_ENV_VAR): + raise RuntimeError( + "Simulating failure to upload package for testing purposes." + ) + _internal_kv_put(pkg_uri, data) + except Exception as e: + raise RuntimeError( + "Failed to store package in the GCS.\n" + f" - GCS URI: {pkg_uri}\n" + f" - Package data ({size_str}): {data[:15]}...\n" + ) from e + logger.info(f"Successfully pushed file package '{pkg_uri}'.") + return len(data) + + +def _get_local_path(base_directory: str, pkg_uri: str) -> str: + _, pkg_name = parse_uri(pkg_uri) + return os.path.join(base_directory, pkg_name) + + +def _zip_files( + path_str: str, + excludes: List[str], + output_path: str, + include_parent_dir: bool = False, + logger: Optional[logging.Logger] = default_logger, +) -> None: + """Zip the target file or directory and write it to the output_path. + + path_str: The file or directory to zip. + excludes (List(str)): The directories or file to be excluded. + output_path: The output path for the zip file. + include_parent_dir: If true, includes the top-level directory as a + directory inside the zip file. + """ + pkg_file = Path(output_path).absolute() + with ZipFile(pkg_file, "w", strict_timestamps=False) as zip_handler: + # Put all files in the directory into the zip file. + file_path = Path(path_str).absolute() + dir_path = file_path + if file_path.is_file(): + dir_path = file_path.parent + + def handler(path: Path): + # Pack this path if it's an empty directory or it's a file. + if path.is_dir() and next(path.iterdir(), None) is None or path.is_file(): + file_size = path.stat().st_size + if file_size >= FILE_SIZE_WARNING: + logger.warning( + f"File {path} is very large " + f"({_mib_string(file_size)}). Consider adding this " + "file to the 'excludes' list to skip uploading it: " + "`ray.init(..., " + f"runtime_env={{'excludes': ['{path}']}})`" + ) + to_path = path.relative_to(dir_path) + if include_parent_dir: + to_path = dir_path.name / to_path + zip_handler.write(path, to_path) + + excludes = [_get_excludes(file_path, excludes)] + _dir_travel(file_path, excludes, handler, logger=logger) + + +def package_exists(pkg_uri: str) -> bool: + """Check whether the package with given URI exists or not. + + Args: + pkg_uri: The uri of the package + + Return: + True for package existing and False for not. + """ + protocol, pkg_name = parse_uri(pkg_uri) + if protocol == Protocol.GCS: + return _internal_kv_exists(pkg_uri) + else: + raise NotImplementedError(f"Protocol {protocol} is not supported") + + +def get_uri_for_package(package: Path) -> str: + """Get a content-addressable URI from a package's contents.""" + + if package.suffix == ".whl": + # Wheel file names include the Python package name, version + # and tags, so it is already effectively content-addressed. + return "{protocol}://{whl_filename}".format( + protocol=Protocol.GCS.value, whl_filename=package.name + ) + else: + hash_val = hashlib.sha1(package.read_bytes()).hexdigest() + return "{protocol}://{pkg_name}.zip".format( + protocol=Protocol.GCS.value, pkg_name=RAY_PKG_PREFIX + hash_val + ) + + +def get_uri_for_file(file: str) -> str: + """Get a content-addressable URI from a file's content. + + This function generates the name of the package by the file. + The final package name is _ray_pkg_.zip of this package, + where HASH_VAL is the hash value of the file. + For example: _ray_pkg_029f88d5ecc55e1e4d64fc6e388fd103.zip + + Examples: + + >>> get_uri_for_file("/my_file.py") # doctest: +SKIP + _ray_pkg_af2734982a741.zip + + Args: + file: The file. + + Returns: + URI (str) + + Raises: + ValueError: If the file doesn't exist. + """ + filepath = Path(file).absolute() + if not filepath.exists() or not filepath.is_file(): + raise ValueError(f"File {filepath} must be an existing file") + + hash_val = _hash_file(filepath, filepath.parent) + + return "{protocol}://{pkg_name}.zip".format( + protocol=Protocol.GCS.value, pkg_name=RAY_PKG_PREFIX + hash_val.hex() + ) + + +def get_uri_for_directory(directory: str, excludes: Optional[List[str]] = None) -> str: + """Get a content-addressable URI from a directory's contents. + + This function generates the name of the package by the directory. + It'll go through all the files in the directory and hash the contents + of the files to get the hash value of the package. + The final package name is _ray_pkg_.zip of this package. + For example: _ray_pkg_029f88d5ecc55e1e4d64fc6e388fd103.zip + + Examples: + + >>> get_uri_for_directory("/my_directory") # doctest: +SKIP + _ray_pkg_af2734982a741.zip + + Args: + directory: The directory. + excludes (list[str]): The dir or files that should be excluded. + + Returns: + URI (str) + + Raises: + ValueError: If the directory doesn't exist. + """ + if excludes is None: + excludes = [] + + directory = Path(directory).absolute() + if not directory.exists() or not directory.is_dir(): + raise ValueError(f"directory {directory} must be an existing directory") + + hash_val = _hash_directory(directory, directory, _get_excludes(directory, excludes)) + + return "{protocol}://{pkg_name}.zip".format( + protocol=Protocol.GCS.value, pkg_name=RAY_PKG_PREFIX + hash_val.hex() + ) + + +def upload_package_to_gcs(pkg_uri: str, pkg_bytes: bytes) -> None: + """Upload a local package to GCS. + + Args: + pkg_uri: The URI of the package, e.g. gcs://my_package.zip + pkg_bytes: The data to be uploaded. + + Raises: + RuntimeError: If the upload fails. + ValueError: If the pkg_uri is a remote path or if the data's + size exceeds GCS_STORAGE_MAX_SIZE. + NotImplementedError: If the protocol of the URI is not supported. + + """ + protocol, pkg_name = parse_uri(pkg_uri) + if protocol == Protocol.GCS: + _store_package_in_gcs(pkg_uri, pkg_bytes) + elif protocol in Protocol.remote_protocols(): + raise ValueError( + "upload_package_to_gcs should not be called with a remote path." + ) + else: + raise NotImplementedError(f"Protocol {protocol} is not supported") + + +def create_package( + module_path: str, + target_path: Path, + include_parent_dir: bool = False, + excludes: Optional[List[str]] = None, + logger: Optional[logging.Logger] = default_logger, +): + if excludes is None: + excludes = [] + + if logger is None: + logger = default_logger + + if not target_path.exists(): + logger.info(f"Creating a file package for local module '{module_path}'.") + _zip_files( + module_path, + excludes, + str(target_path), + include_parent_dir=include_parent_dir, + logger=logger, + ) + + +def upload_package_if_needed( + pkg_uri: str, + base_directory: str, + module_path: str, + include_parent_dir: bool = False, + excludes: Optional[List[str]] = None, + logger: Optional[logging.Logger] = default_logger, +) -> bool: + """Upload the contents of the directory under the given URI. + + This will first create a temporary zip file under the passed + base_directory. + + If the package already exists in storage, this is a no-op. + + Args: + pkg_uri: URI of the package to upload. + base_directory: Directory where package files are stored. + module_path: The module to be uploaded, either a single .py file or a directory. + include_parent_dir: If true, includes the top-level directory as a + directory inside the zip file. + excludes: List specifying files to exclude. + + Raises: + RuntimeError: If the upload fails. + ValueError: If the pkg_uri is a remote path or if the data's + size exceeds GCS_STORAGE_MAX_SIZE. + NotImplementedError: If the protocol of the URI is not supported. + """ + if excludes is None: + excludes = [] + + if logger is None: + logger = default_logger + + pin_runtime_env_uri(pkg_uri) + + if package_exists(pkg_uri): + return False + + package_file = Path(_get_local_path(base_directory, pkg_uri)) + # Make the temporary zip file name unique so that it doesn't conflict with + # concurrent upload_package_if_needed calls with the same pkg_uri. + # See https://github.com/ray-project/ray/issues/47471. + package_file = package_file.with_name( + f"{time.time_ns()}_{os.getpid()}_{package_file.name}" + ) + + create_package( + module_path, + package_file, + include_parent_dir=include_parent_dir, + excludes=excludes, + ) + package_file_bytes = package_file.read_bytes() + # Remove the local file to avoid accumulating temporary zip files. + package_file.unlink() + + upload_package_to_gcs(pkg_uri, package_file_bytes) + + return True + + +def get_local_dir_from_uri(uri: str, base_directory: str) -> Path: + """Return the local directory corresponding to this URI.""" + pkg_file = Path(_get_local_path(base_directory, uri)) + local_dir = pkg_file.with_suffix("") + return local_dir + + +@DeveloperAPI +async def download_and_unpack_package( + pkg_uri: str, + base_directory: str, + gcs_client: Optional[GcsClient] = None, + logger: Optional[logging.Logger] = default_logger, + overwrite: bool = False, +) -> str: + """Download the package corresponding to this URI and unpack it if zipped. + + Will be written to a file or directory named {base_directory}/{uri}. + Returns the path to this file or directory. + + Args: + pkg_uri: URI of the package to download. + base_directory: Directory to use as the parent directory of the target + directory for the unpacked files. + gcs_client: Client to use for downloading from the GCS. + logger: The logger to use. + overwrite: If True, overwrite the existing package. + + Returns: + Path to the local directory containing the unpacked package files. + + Raises: + IOError: If the download fails. + ImportError: If smart_open is not installed and a remote URI is used. + NotImplementedError: If the protocol of the URI is not supported. + ValueError: If the GCS client is not provided when downloading from GCS, + or if package URI is invalid. + + """ + pkg_file = Path(_get_local_path(base_directory, pkg_uri)) + if pkg_file.suffix == "": + raise ValueError( + f"Invalid package URI: {pkg_uri}." + "URI must have a file extension and the URI must be valid." + ) + + async with _AsyncFileLock(str(pkg_file) + ".lock"): + if logger is None: + logger = default_logger + + logger.debug(f"Fetching package for URI: {pkg_uri}") + + local_dir = get_local_dir_from_uri(pkg_uri, base_directory) + assert local_dir != pkg_file, "Invalid pkg_file!" + + download_package: bool = True + if local_dir.exists() and not overwrite: + download_package = False + assert local_dir.is_dir(), f"{local_dir} is not a directory" + elif local_dir.exists(): + logger.info(f"Removing {local_dir} with pkg_file {pkg_file}") + shutil.rmtree(local_dir) + + if download_package: + protocol, _ = parse_uri(pkg_uri) + logger.info( + f"Downloading package from {pkg_uri} to {pkg_file} " + f"with protocol {protocol}" + ) + if protocol == Protocol.GCS: + if gcs_client is None: + raise ValueError( + "GCS client must be provided to download from GCS." + ) + + # Download package from the GCS. + code = await gcs_client.async_internal_kv_get( + pkg_uri.encode(), namespace=None, timeout=None + ) + if os.environ.get(RAY_RUNTIME_ENV_FAIL_DOWNLOAD_FOR_TESTING_ENV_VAR): + code = None + if code is None: + raise IOError( + f"Failed to download runtime_env file package {pkg_uri} " + "from the GCS to the Ray worker node. The package may " + "have prematurely been deleted from the GCS due to a " + "long upload time or a problem with Ray. Try setting the " + "environment variable " + f"{RAY_RUNTIME_ENV_URI_PIN_EXPIRATION_S_ENV_VAR} " + " to a value larger than the upload time in seconds " + "(the default is " + f"{RAY_RUNTIME_ENV_URI_PIN_EXPIRATION_S_DEFAULT}). " + "If this fails, try re-running " + "after making any change to a file in the file package." + ) + code = code or b"" + pkg_file.write_bytes(code) + + if is_zip_uri(pkg_uri): + unzip_package( + package_path=pkg_file, + target_dir=local_dir, + remove_top_level_directory=False, + unlink_zip=True, + logger=logger, + ) + else: + return str(pkg_file) + elif protocol in Protocol.remote_protocols(): + protocol.download_remote_uri(source_uri=pkg_uri, dest_file=pkg_file) + + if pkg_file.suffix in [".zip", ".jar"]: + unzip_package( + package_path=pkg_file, + target_dir=local_dir, + remove_top_level_directory=True, + unlink_zip=True, + logger=logger, + ) + elif pkg_file.suffix == ".whl": + return str(pkg_file) + else: + raise NotImplementedError( + f"Package format {pkg_file.suffix} is ", + "not supported for remote protocols", + ) + else: + raise NotImplementedError(f"Protocol {protocol} is not supported") + + return str(local_dir) + + +def get_top_level_dir_from_compressed_package(package_path: str): + """ + If compressed package at package_path contains a single top-level + directory, returns the name of the top-level directory. Otherwise, + returns None. + + Ignores a second top-level directory if it is named __MACOSX. + """ + + package_zip = ZipFile(package_path, "r") + top_level_directory = None + + def is_top_level_file(file_name): + return "/" not in file_name + + def base_dir_name(file_name): + return file_name.split("/")[0] + + for file_name in package_zip.namelist(): + if top_level_directory is None: + # Cache the top_level_directory name when checking + # the first file in the zipped package + if is_top_level_file(file_name): + return None + else: + # Top-level directory, or non-top-level file or directory + dir_name = base_dir_name(file_name) + if dir_name == MAC_OS_ZIP_HIDDEN_DIR_NAME: + continue + top_level_directory = dir_name + else: + # Confirm that all other files + # belong to the same top_level_directory + if is_top_level_file(file_name) or base_dir_name(file_name) not in [ + top_level_directory, + MAC_OS_ZIP_HIDDEN_DIR_NAME, + ]: + return None + + return top_level_directory + + +def remove_dir_from_filepaths(base_dir: str, rdir: str): + """ + base_dir: String path of the directory containing rdir + rdir: String path of directory relative to base_dir whose contents should + be moved to its base_dir, its parent directory + + Removes rdir from the filepaths of all files and directories inside it. + In other words, moves all the files inside rdir to the directory that + contains rdir. Assumes base_dir's contents and rdir's contents have no + name conflicts. + """ + + # Move rdir to a temporary directory, so its contents can be moved to + # base_dir without any name conflicts + with TemporaryDirectory() as tmp_dir: + # shutil.move() is used instead of os.rename() in case rdir and tmp_dir + # are located on separate file systems + shutil.move(os.path.join(base_dir, rdir), os.path.join(tmp_dir, rdir)) + + # Shift children out of rdir and into base_dir + rdir_children = os.listdir(os.path.join(tmp_dir, rdir)) + for child in rdir_children: + shutil.move( + os.path.join(tmp_dir, rdir, child), os.path.join(base_dir, child) + ) + + +def unzip_package( + package_path: str, + target_dir: str, + remove_top_level_directory: bool, + unlink_zip: bool, + logger: Optional[logging.Logger] = default_logger, +) -> None: + """ + Unzip the compressed package contained at package_path to target_dir. + + If remove_top_level_directory is True and the top level consists of a + a single directory (or possibly also a second hidden directory named + __MACOSX at the top level arising from macOS's zip command), the function + will automatically remove the top-level directory and store the contents + directly in target_dir. + + Otherwise, if remove_top_level_directory is False or if the top level + consists of multiple files or directories (not counting __MACOS), + the zip contents will be stored in target_dir. + + Args: + package_path: String path of the compressed package to unzip. + target_dir: String path of the directory to store the unzipped contents. + remove_top_level_directory: Whether to remove the top-level directory + from the zip contents. + unlink_zip: Whether to unlink the zip file stored at package_path. + logger: Optional logger to use for logging. + + """ + try: + os.mkdir(target_dir) + except FileExistsError: + logger.info(f"Directory at {target_dir} already exists") + + logger.debug(f"Unpacking {package_path} to {target_dir}") + + with ZipFile(str(package_path), "r") as zip_ref: + zip_ref.extractall(target_dir) + if remove_top_level_directory: + top_level_directory = get_top_level_dir_from_compressed_package(package_path) + if top_level_directory is not None: + # Remove __MACOSX directory if it exists + macos_dir = os.path.join(target_dir, MAC_OS_ZIP_HIDDEN_DIR_NAME) + if os.path.isdir(macos_dir): + shutil.rmtree(macos_dir) + + remove_dir_from_filepaths(target_dir, top_level_directory) + + if unlink_zip: + Path(package_path).unlink() + + +def delete_package(pkg_uri: str, base_directory: str) -> Tuple[bool, int]: + """Deletes a specific URI from the local filesystem. + + Args: + pkg_uri: URI to delete. + + Returns: + bool: True if the URI was successfully deleted, else False. + """ + + deleted = False + path = Path(_get_local_path(base_directory, pkg_uri)) + with FileLock(str(path) + ".lock"): + path = path.with_suffix("") + if path.exists(): + if path.is_dir() and not path.is_symlink(): + shutil.rmtree(str(path)) + else: + path.unlink() + deleted = True + + return deleted + + +async def install_wheel_package( + wheel_uri: str, + target_dir: str, + logger: Optional[logging.Logger] = default_logger, +) -> None: + """Install packages in the wheel URI, and then delete the local wheel file.""" + + pip_install_cmd = [ + "pip", + "install", + wheel_uri, + f"--target={target_dir}", + ] + + logger.info("Running py_modules wheel install command: %s", str(pip_install_cmd)) + try: + # TODO(architkulkarni): Use `await check_output_cmd` or similar. + exit_code, output = exec_cmd_stream_to_logger(pip_install_cmd, logger) + finally: + wheel_uri_path = Path(wheel_uri) + if wheel_uri_path.exists(): + if wheel_uri_path.is_dir(): + shutil.rmtree(wheel_uri) + else: + Path(wheel_uri).unlink() + + if exit_code != 0: + if Path(target_dir).exists(): + shutil.rmtree(target_dir) + raise RuntimeError( + f"Failed to install py_modules wheel {wheel_uri}" + f"to {target_dir}:\n{output}" + ) diff --git a/lib/python3.12/site-packages/ray/_private/runtime_env/pip.py b/lib/python3.12/site-packages/ray/_private/runtime_env/pip.py new file mode 100644 index 0000000000000000000000000000000000000000..903675496cd8444933888c7e6a6e41a6bde62705 --- /dev/null +++ b/lib/python3.12/site-packages/ray/_private/runtime_env/pip.py @@ -0,0 +1,344 @@ +import asyncio +import hashlib +import json +import logging +import os +import shutil +import sys +from asyncio import create_task, get_running_loop +from typing import Dict, List, Optional + +from ray._common.utils import try_to_create_directory +from ray._private.runtime_env import dependency_utils, virtualenv_utils +from ray._private.runtime_env.packaging import Protocol, parse_uri +from ray._private.runtime_env.plugin import RuntimeEnvPlugin +from ray._private.runtime_env.utils import check_output_cmd +from ray._private.utils import get_directory_size_bytes + +default_logger = logging.getLogger(__name__) + + +def _get_pip_hash(pip_dict: Dict) -> str: + serialized_pip_spec = json.dumps(pip_dict, sort_keys=True) + hash_val = hashlib.sha1(serialized_pip_spec.encode("utf-8")).hexdigest() + return hash_val + + +def get_uri(runtime_env: Dict) -> Optional[str]: + """Return `"pip://"`, or None if no GC required.""" + pip = runtime_env.get("pip") + if pip is not None: + if isinstance(pip, dict): + uri = "pip://" + _get_pip_hash(pip_dict=pip) + elif isinstance(pip, list): + uri = "pip://" + _get_pip_hash(pip_dict=dict(packages=pip)) + else: + raise TypeError( + "pip field received by RuntimeEnvAgent must be " + f"list or dict, not {type(pip).__name__}." + ) + else: + uri = None + return uri + + +class PipProcessor: + def __init__( + self, + target_dir: str, + runtime_env: "RuntimeEnv", # noqa: F821 + logger: Optional[logging.Logger] = default_logger, + ): + try: + import virtualenv # noqa: F401 ensure virtualenv exists. + except ImportError: + raise RuntimeError( + f"Please install virtualenv " + f"`{sys.executable} -m pip install virtualenv`" + f"to enable pip runtime env." + ) + logger.debug("Setting up pip for runtime_env: %s", runtime_env) + self._target_dir = target_dir + self._runtime_env = runtime_env + self._logger = logger + + self._pip_config = self._runtime_env.pip_config() + self._pip_env = os.environ.copy() + self._pip_env.update(self._runtime_env.env_vars()) + + @classmethod + async def _ensure_pip_version( + cls, + path: str, + pip_version: Optional[str], + cwd: str, + pip_env: Dict, + logger: logging.Logger, + ): + """Run the pip command to reinstall pip to the specified version.""" + if not pip_version: + return + + python = virtualenv_utils.get_virtualenv_python(path) + # Ensure pip version. + pip_reinstall_cmd = [ + python, + "-m", + "pip", + "install", + "--disable-pip-version-check", + f"pip{pip_version}", + ] + logger.info("Installing pip with version %s", pip_version) + + await check_output_cmd(pip_reinstall_cmd, logger=logger, cwd=cwd, env=pip_env) + + async def _pip_check( + self, + path: str, + pip_check: bool, + cwd: str, + pip_env: Dict, + logger: logging.Logger, + ): + """Run the pip check command to check python dependency conflicts. + If exists conflicts, the exit code of pip check command will be non-zero. + """ + if not pip_check: + logger.info("Skip pip check.") + return + python = virtualenv_utils.get_virtualenv_python(path) + + await check_output_cmd( + [python, "-m", "pip", "check", "--disable-pip-version-check"], + logger=logger, + cwd=cwd, + env=pip_env, + ) + + logger.info("Pip check on %s successfully.", path) + + async def _install_pip_packages( + self, + path: str, + pip_packages: List[str], + cwd: str, + pip_env: Dict, + logger: logging.Logger, + ): + virtualenv_path = virtualenv_utils.get_virtualenv_path(path) + python = virtualenv_utils.get_virtualenv_python(path) + # TODO(fyrestone): Support -i, --no-deps, --no-cache-dir, ... + pip_requirements_file = dependency_utils.get_requirements_file( + path, pip_packages + ) + + # Avoid blocking the event loop. + loop = get_running_loop() + await loop.run_in_executor( + None, + dependency_utils.gen_requirements_txt, + pip_requirements_file, + pip_packages, + ) + + # Install all dependencies + # The default options for pip install are + # + # --disable-pip-version-check + # Don't periodically check PyPI to determine whether a new version + # of pip is available for download. + # + # --no-cache-dir + # Disable the cache, the pip runtime env is a one-time installation, + # and we don't need to handle the pip cache broken. + # + # Allow users to specify their own options to install packages via `pip`. + pip_install_cmd = [ + python, + "-m", + "pip", + "install", + "-r", + pip_requirements_file, + ] + + pip_opt_list = self._pip_config.get( + "pip_install_options", ["--disable-pip-version-check", "--no-cache-dir"] + ) + pip_install_cmd.extend(pip_opt_list) + + logger.info("Installing python requirements to %s", virtualenv_path) + + await check_output_cmd(pip_install_cmd, logger=logger, cwd=cwd, env=pip_env) + + async def _run(self): + path = self._target_dir + logger = self._logger + pip_packages = self._pip_config["packages"] + # We create an empty directory for exec cmd so that the cmd will + # run more stable. e.g. if cwd has ray, then checking ray will + # look up ray in cwd instead of site packages. + exec_cwd = os.path.join(path, "exec_cwd") + os.makedirs(exec_cwd, exist_ok=True) + try: + await virtualenv_utils.create_or_get_virtualenv(path, exec_cwd, logger) + python = virtualenv_utils.get_virtualenv_python(path) + async with dependency_utils.check_ray(python, exec_cwd, logger): + # Ensure pip version. + await self._ensure_pip_version( + path, + self._pip_config.get("pip_version", None), + exec_cwd, + self._pip_env, + logger, + ) + # Install pip packages. + await self._install_pip_packages( + path, + pip_packages, + exec_cwd, + self._pip_env, + logger, + ) + # Check python environment for conflicts. + await self._pip_check( + path, + self._pip_config.get("pip_check", False), + exec_cwd, + self._pip_env, + logger, + ) + except Exception: + logger.info("Delete incomplete virtualenv: %s", path) + shutil.rmtree(path, ignore_errors=True) + logger.exception("Failed to install pip packages.") + raise + + def __await__(self): + return self._run().__await__() + + +class PipPlugin(RuntimeEnvPlugin): + name = "pip" + + def __init__(self, resources_dir: str): + self._pip_resources_dir = os.path.join(resources_dir, "pip") + self._creating_task = {} + # Maps a URI to a lock that is used to prevent multiple concurrent + # installs of the same virtualenv, see #24513 + self._create_locks: Dict[str, asyncio.Lock] = {} + # Key: created hashes. Value: size of the pip dir. + self._created_hash_bytes: Dict[str, int] = {} + try_to_create_directory(self._pip_resources_dir) + + def _get_path_from_hash(self, hash_val: str) -> str: + """Generate a path from the hash of a pip spec. + + Example output: + /tmp/ray/session_2021-11-03_16-33-59_356303_41018/runtime_resources + /pip/ray-9a7972c3a75f55e976e620484f58410c920db091 + """ + return os.path.join(self._pip_resources_dir, hash_val) + + def get_uris(self, runtime_env: "RuntimeEnv") -> List[str]: # noqa: F821 + """Return the pip URI from the RuntimeEnv if it exists, else return [].""" + pip_uri = runtime_env.pip_uri() + if pip_uri: + return [pip_uri] + return [] + + def delete_uri( + self, uri: str, logger: Optional[logging.Logger] = default_logger + ) -> int: + """Delete URI and return the number of bytes deleted.""" + logger.info("Got request to delete pip URI %s", uri) + protocol, hash_val = parse_uri(uri) + if protocol != Protocol.PIP: + raise ValueError( + "PipPlugin can only delete URIs with protocol " + f"pip. Received protocol {protocol}, URI {uri}" + ) + + # Cancel running create task. + task = self._creating_task.pop(hash_val, None) + if task is not None: + task.cancel() + + del self._created_hash_bytes[hash_val] + + pip_env_path = self._get_path_from_hash(hash_val) + local_dir_size = get_directory_size_bytes(pip_env_path) + del self._create_locks[uri] + try: + shutil.rmtree(pip_env_path) + except OSError as e: + logger.warning(f"Error when deleting pip env {pip_env_path}: {str(e)}") + return 0 + + return local_dir_size + + async def create( + self, + uri: str, + runtime_env: "RuntimeEnv", # noqa: F821 + context: "RuntimeEnvContext", # noqa: F821 + logger: Optional[logging.Logger] = default_logger, + ) -> int: + if not runtime_env.has_pip(): + return 0 + + protocol, hash_val = parse_uri(uri) + target_dir = self._get_path_from_hash(hash_val) + + async def _create_for_hash(): + await PipProcessor( + target_dir, + runtime_env, + logger, + ) + + loop = get_running_loop() + return await loop.run_in_executor( + None, get_directory_size_bytes, target_dir + ) + + if uri not in self._create_locks: + # async lock to prevent the same virtualenv being concurrently installed + self._create_locks[uri] = asyncio.Lock() + + async with self._create_locks[uri]: + if hash_val in self._created_hash_bytes: + return self._created_hash_bytes[hash_val] + self._creating_task[hash_val] = task = create_task(_create_for_hash()) + task.add_done_callback(lambda _: self._creating_task.pop(hash_val, None)) + pip_dir_bytes = await task + self._created_hash_bytes[hash_val] = pip_dir_bytes + return pip_dir_bytes + + def modify_context( + self, + uris: List[str], + runtime_env: "RuntimeEnv", # noqa: F821 + context: "RuntimeEnvContext", # noqa: F821 + logger: logging.Logger = default_logger, + ): + if not runtime_env.has_pip(): + return + # PipPlugin only uses a single URI. + uri = uris[0] + # Update py_executable. + protocol, hash_val = parse_uri(uri) + target_dir = self._get_path_from_hash(hash_val) + virtualenv_python = virtualenv_utils.get_virtualenv_python(target_dir) + + if not os.path.exists(virtualenv_python): + raise ValueError( + f"Local directory {target_dir} for URI {uri} does " + "not exist on the cluster. Something may have gone wrong while " + "installing the runtime_env `pip` packages." + ) + context.py_executable = virtualenv_python + context.command_prefix += virtualenv_utils.get_virtualenv_activate_command( + target_dir + ) diff --git a/lib/python3.12/site-packages/ray/_private/runtime_env/plugin.py b/lib/python3.12/site-packages/ray/_private/runtime_env/plugin.py new file mode 100644 index 0000000000000000000000000000000000000000..e5da3ecbeb6c285f0d3ec296f31598fcc58b5d5a --- /dev/null +++ b/lib/python3.12/site-packages/ray/_private/runtime_env/plugin.py @@ -0,0 +1,265 @@ +import json +import logging +import os +from abc import ABC +from typing import Any, Dict, List, Optional, Type + +from ray._common.utils import import_attr +from ray._private.runtime_env.constants import ( + RAY_RUNTIME_ENV_CLASS_FIELD_NAME, + RAY_RUNTIME_ENV_PLUGIN_DEFAULT_PRIORITY, + RAY_RUNTIME_ENV_PLUGIN_MAX_PRIORITY, + RAY_RUNTIME_ENV_PLUGIN_MIN_PRIORITY, + RAY_RUNTIME_ENV_PLUGINS_ENV_VAR, + RAY_RUNTIME_ENV_PRIORITY_FIELD_NAME, +) +from ray._private.runtime_env.context import RuntimeEnvContext +from ray._private.runtime_env.uri_cache import URICache +from ray.util.annotations import DeveloperAPI + +default_logger = logging.getLogger(__name__) + + +@DeveloperAPI +class RuntimeEnvPlugin(ABC): + """Abstract base class for runtime environment plugins.""" + + name: str = None + priority: int = RAY_RUNTIME_ENV_PLUGIN_DEFAULT_PRIORITY + + @staticmethod + def validate(runtime_env_dict: dict) -> None: + """Validate user entry for this plugin. + + The method is invoked upon installation of runtime env. + + Args: + runtime_env_dict: The user-supplied runtime environment dict. + + Raises: + ValueError: If the validation fails. + """ + pass + + def get_uris(self, runtime_env: "RuntimeEnv") -> List[str]: # noqa: F821 + return [] + + async def create( + self, + uri: Optional[str], + runtime_env, + context: RuntimeEnvContext, + logger: logging.Logger, + ) -> float: + """Create and install the runtime environment. + + Gets called in the runtime env agent at install time. The URI can be + used as a caching mechanism. + + Args: + uri: A URI uniquely describing this resource. + runtime_env: The RuntimeEnv object. + context: Auxiliary information supplied by Ray. + logger: A logger to log messages during the context modification. + + Returns: + float: The disk space taken up by this plugin installation for this + environment. e.g. for working_dir, this downloads the files to the + local node. + """ + return 0 + + def modify_context( + self, + uris: List[str], + runtime_env: "RuntimeEnv", # noqa: F821 + context: RuntimeEnvContext, + logger: logging.Logger, + ) -> None: + """Modify context to change worker startup behavior. + + For example, you can use this to prepend "cd " command to worker + startup, or add new environment variables. + + Args: + uris: The URIs used by this resource. + runtime_env: The RuntimeEnv object. + context: Auxiliary information supplied by Ray. + logger: A logger to log messages during the context modification. + """ + return + + def delete_uri(self, uri: str, logger: logging.Logger) -> float: + """Delete the runtime environment given uri. + + Args: + uri: A URI uniquely describing this resource. + logger: The logger used to log messages during the deletion. + + Returns: + float: The amount of space reclaimed by the deletion. + """ + return 0 + + +class PluginSetupContext: + def __init__( + self, + name: str, + class_instance: RuntimeEnvPlugin, + priority: int, + uri_cache: URICache, + ): + self.name = name + self.class_instance = class_instance + self.priority = priority + self.uri_cache = uri_cache + + +class RuntimeEnvPluginManager: + """This manager is used to load plugins in runtime env agent.""" + + def __init__(self): + self.plugins: Dict[str, PluginSetupContext] = {} + plugin_config_str = os.environ.get(RAY_RUNTIME_ENV_PLUGINS_ENV_VAR) + if plugin_config_str: + plugin_configs = json.loads(plugin_config_str) + self.load_plugins(plugin_configs) + + def validate_plugin_class(self, plugin_class: Type[RuntimeEnvPlugin]) -> None: + if not issubclass(plugin_class, RuntimeEnvPlugin): + raise RuntimeError( + f"Invalid runtime env plugin class {plugin_class}. " + "The plugin class must inherit " + "ray._private.runtime_env.plugin.RuntimeEnvPlugin." + ) + if not plugin_class.name: + raise RuntimeError(f"No valid name in runtime env plugin {plugin_class}.") + if plugin_class.name in self.plugins: + raise RuntimeError( + f"The name of runtime env plugin {plugin_class} conflicts " + f"with {self.plugins[plugin_class.name]}.", + ) + + def validate_priority(self, priority: Any) -> None: + if ( + not isinstance(priority, int) + or priority < RAY_RUNTIME_ENV_PLUGIN_MIN_PRIORITY + or priority > RAY_RUNTIME_ENV_PLUGIN_MAX_PRIORITY + ): + raise RuntimeError( + f"Invalid runtime env priority {priority}, " + "it should be an integer between " + f"{RAY_RUNTIME_ENV_PLUGIN_MIN_PRIORITY} " + f"and {RAY_RUNTIME_ENV_PLUGIN_MAX_PRIORITY}." + ) + + def load_plugins(self, plugin_configs: List[Dict]) -> None: + """Load runtime env plugins and create URI caches for them.""" + for plugin_config in plugin_configs: + if ( + not isinstance(plugin_config, dict) + or RAY_RUNTIME_ENV_CLASS_FIELD_NAME not in plugin_config + ): + raise RuntimeError( + f"Invalid runtime env plugin config {plugin_config}, " + "it should be a object which contains the " + f"{RAY_RUNTIME_ENV_CLASS_FIELD_NAME} field." + ) + plugin_class = import_attr(plugin_config[RAY_RUNTIME_ENV_CLASS_FIELD_NAME]) + self.validate_plugin_class(plugin_class) + + # The priority should be an integer between 0 and 100. + # The default priority is 10. A smaller number indicates a + # higher priority and the plugin will be set up first. + if RAY_RUNTIME_ENV_PRIORITY_FIELD_NAME in plugin_config: + priority = plugin_config[RAY_RUNTIME_ENV_PRIORITY_FIELD_NAME] + else: + priority = plugin_class.priority + self.validate_priority(priority) + + class_instance = plugin_class() + self.plugins[plugin_class.name] = PluginSetupContext( + plugin_class.name, + class_instance, + priority, + self.create_uri_cache_for_plugin(class_instance), + ) + + def add_plugin(self, plugin: RuntimeEnvPlugin) -> None: + """Add a plugin to the manager and create a URI cache for it. + + Args: + plugin: The class instance of the plugin. + """ + plugin_class = type(plugin) + self.validate_plugin_class(plugin_class) + self.validate_priority(plugin_class.priority) + self.plugins[plugin_class.name] = PluginSetupContext( + plugin_class.name, + plugin, + plugin_class.priority, + self.create_uri_cache_for_plugin(plugin), + ) + + def create_uri_cache_for_plugin(self, plugin: RuntimeEnvPlugin) -> URICache: + """Create a URI cache for a plugin. + + Args: + plugin_name: The name of the plugin. + + Returns: + The created URI cache for the plugin. + """ + # Set the max size for the cache. Defaults to 10 GB. + cache_size_env_var = f"RAY_RUNTIME_ENV_{plugin.name}_CACHE_SIZE_GB".upper() + cache_size_bytes = int( + (1024**3) * float(os.environ.get(cache_size_env_var, 10)) + ) + return URICache(plugin.delete_uri, cache_size_bytes) + + def sorted_plugin_setup_contexts(self) -> List[PluginSetupContext]: + """Get the sorted plugin setup contexts, sorted by increasing priority. + + Returns: + The sorted plugin setup contexts. + """ + return sorted(self.plugins.values(), key=lambda x: x.priority) + + +async def create_for_plugin_if_needed( + runtime_env: "RuntimeEnv", # noqa: F821 + plugin: RuntimeEnvPlugin, + uri_cache: URICache, + context: RuntimeEnvContext, + logger: logging.Logger = default_logger, +): + """Set up the environment using the plugin if not already set up and cached.""" + if plugin.name not in runtime_env or runtime_env[plugin.name] is None: + return + + plugin.validate(runtime_env) + + uris = plugin.get_uris(runtime_env) + + if not uris: + logger.debug( + f"No URIs for runtime env plugin {plugin.name}; " + "create always without checking the cache." + ) + await plugin.create(None, runtime_env, context, logger=logger) + + for uri in uris: + if uri not in uri_cache: + logger.debug(f"Cache miss for URI {uri}.") + size_bytes = await plugin.create(uri, runtime_env, context, logger=logger) + uri_cache.add(uri, size_bytes, logger=logger) + else: + logger.info( + f"Runtime env {plugin.name} {uri} is already installed " + "and will be reused. Search " + "all runtime_env_setup-*.log to find the corresponding setup log." + ) + uri_cache.mark_used(uri, logger=logger) + + plugin.modify_context(uris, runtime_env, context, logger) diff --git a/lib/python3.12/site-packages/ray/_private/runtime_env/plugin_schema_manager.py b/lib/python3.12/site-packages/ray/_private/runtime_env/plugin_schema_manager.py new file mode 100644 index 0000000000000000000000000000000000000000..5de24113b26fa73247eb47b7be31ffb191c3ac20 --- /dev/null +++ b/lib/python3.12/site-packages/ray/_private/runtime_env/plugin_schema_manager.py @@ -0,0 +1,97 @@ +import json +import logging +import os +from typing import List + +import jsonschema + +from ray._private.runtime_env.constants import ( + RAY_RUNTIME_ENV_PLUGIN_SCHEMA_SUFFIX, + RAY_RUNTIME_ENV_PLUGIN_SCHEMAS_ENV_VAR, +) + +logger = logging.getLogger(__name__) + + +class RuntimeEnvPluginSchemaManager: + """This manager is used to load plugin json schemas.""" + + default_schema_path = os.path.join( + os.path.dirname(__file__), "../../runtime_env/schemas" + ) + schemas = {} + loaded = False + + @classmethod + def _load_schemas(cls, schema_paths: List[str]): + for schema_path in schema_paths: + try: + with open(schema_path) as f: + schema = json.load(f) + except json.decoder.JSONDecodeError: + logger.error("Invalid runtime env schema %s, skip it.", schema_path) + continue + except OSError: + logger.error("Cannot open runtime env schema %s, skip it.", schema_path) + continue + if "title" not in schema: + logger.error( + "No valid title in runtime env schema %s, skip it.", schema_path + ) + continue + if schema["title"] in cls.schemas: + logger.error( + "The 'title' of runtime env schema %s conflicts with %s, skip it.", + schema_path, + cls.schemas[schema["title"]], + ) + continue + cls.schemas[schema["title"]] = schema + + @classmethod + def _load_default_schemas(cls): + schema_json_files = list() + for root, _, files in os.walk(cls.default_schema_path): + for f in files: + if f.endswith(RAY_RUNTIME_ENV_PLUGIN_SCHEMA_SUFFIX): + schema_json_files.append(os.path.join(root, f)) + logger.debug( + f"Loading the default runtime env schemas: {schema_json_files}." + ) + cls._load_schemas(schema_json_files) + + @classmethod + def _load_schemas_from_env_var(cls): + # The format of env var: + # "/path/to/env_1_schema.json,/path/to/env_2_schema.json,/path/to/schemas_dir/" + schema_paths = os.environ.get(RAY_RUNTIME_ENV_PLUGIN_SCHEMAS_ENV_VAR) + if schema_paths: + schema_json_files = list() + for path in schema_paths.split(","): + if path.endswith(RAY_RUNTIME_ENV_PLUGIN_SCHEMA_SUFFIX): + schema_json_files.append(path) + elif os.path.isdir(path): + for root, _, files in os.walk(path): + for f in files: + if f.endswith(RAY_RUNTIME_ENV_PLUGIN_SCHEMA_SUFFIX): + schema_json_files.append(os.path.join(root, f)) + logger.info( + f"Loading the runtime env schemas from env var: {schema_json_files}." + ) + cls._load_schemas(schema_json_files) + + @classmethod + def validate(cls, name, instance): + if not cls.loaded: + # Load the schemas lazily. + cls._load_default_schemas() + cls._load_schemas_from_env_var() + cls.loaded = True + # if no schema matches, skip the validation. + if name in cls.schemas: + jsonschema.validate(instance=instance, schema=cls.schemas[name]) + + @classmethod + def clear(cls): + cls.schemas.clear() + cls.loaded = False diff --git a/lib/python3.12/site-packages/ray/_private/runtime_env/protocol.py b/lib/python3.12/site-packages/ray/_private/runtime_env/protocol.py new file mode 100644 index 0000000000000000000000000000000000000000..663e9d4366dadfa398ff9605c3ab46274c328fce --- /dev/null +++ b/lib/python3.12/site-packages/ray/_private/runtime_env/protocol.py @@ -0,0 +1,268 @@ +import enum +import os +from urllib.parse import urlparse + + +class ProtocolsProvider: + _MISSING_DEPENDENCIES_WARNING = ( + "Note that these must be preinstalled " + "on all nodes in the Ray cluster; it is not " + "sufficient to install them in the runtime_env." + ) + + @classmethod + def get_protocols(cls): + return { + # For packages dynamically uploaded and managed by the GCS. + "gcs", + # For conda environments installed locally on each node. + "conda", + # For pip environments installed locally on each node. + "pip", + # For uv environments install locally on each node. + "uv", + # Remote https path, assumes everything packed in one zip file. + "https", + # Remote s3 path, assumes everything packed in one zip file. + "s3", + # Remote google storage path, assumes everything packed in one zip file. + "gs", + # Remote azure blob storage path, assumes everything packed in one zip file. + "azure", + # Remote Azure Blob File System Secure path, assumes everything packed in one zip file. + "abfss", + # File storage path, assumes everything packed in one zip file. + "file", + } + + @classmethod + def get_remote_protocols(cls): + return {"https", "s3", "gs", "azure", "abfss", "file"} + + @classmethod + def _handle_s3_protocol(cls): + """Set up S3 protocol handling. + + Returns: + tuple: (open_file function, transport_params) + + Raises: + ImportError: If required dependencies are not installed. + """ + try: + import boto3 + from smart_open import open as open_file + except ImportError: + raise ImportError( + "You must `pip install smart_open[s3]` " + "to fetch URIs in s3 bucket. " + cls._MISSING_DEPENDENCIES_WARNING + ) + + # Create S3 client, falling back to unsigned for public buckets + session = boto3.Session() + # session.get_credentials() will return None if no credentials can be found. + if session.get_credentials(): + # If credentials are found, use a standard signed client. + s3_client = session.client("s3") + else: + # No credentials found, fall back to an unsigned client for public buckets. + from botocore import UNSIGNED + from botocore.config import Config + + s3_client = boto3.client("s3", config=Config(signature_version=UNSIGNED)) + + transport_params = {"client": s3_client} + return open_file, transport_params + + @classmethod + def _handle_gs_protocol(cls): + """Set up Google Cloud Storage protocol handling. + + Returns: + tuple: (open_file function, transport_params) + + Raises: + ImportError: If required dependencies are not installed. + """ + try: + from google.cloud import storage # noqa: F401 + from smart_open import open as open_file + except ImportError: + raise ImportError( + "You must `pip install smart_open[gcs]` " + "to fetch URIs in Google Cloud Storage bucket." + + cls._MISSING_DEPENDENCIES_WARNING + ) + + return open_file, None + + @classmethod + def _handle_azure_protocol(cls): + """Set up Azure blob storage protocol handling. + + Returns: + tuple: (open_file function, transport_params) + + Raises: + ImportError: If required dependencies are not installed. + ValueError: If required environment variables are not set. + """ + try: + from azure.identity import DefaultAzureCredential + from azure.storage.blob import BlobServiceClient # noqa: F401 + from smart_open import open as open_file + except ImportError: + raise ImportError( + "You must `pip install azure-storage-blob azure-identity smart_open[azure]` " + "to fetch URIs in Azure Blob Storage. " + + cls._MISSING_DEPENDENCIES_WARNING + ) + + # Define authentication variable + azure_storage_account_name = os.getenv("AZURE_STORAGE_ACCOUNT") + + if not azure_storage_account_name: + raise ValueError( + "Azure Blob Storage authentication requires " + "AZURE_STORAGE_ACCOUNT environment variable to be set." + ) + + account_url = f"https://{azure_storage_account_name}.blob.core.windows.net/" + transport_params = { + "client": BlobServiceClient( + account_url=account_url, credential=DefaultAzureCredential() + ) + } + + return open_file, transport_params + + @classmethod + def _handle_abfss_protocol(cls): + """Set up Azure Blob File System Secure (ABFSS) protocol handling. + + Returns: + tuple: (open_file function, transport_params) + + Raises: + ImportError: If required dependencies are not installed. + ValueError: If the ABFSS URI format is invalid. + """ + try: + import adlfs + from azure.identity import DefaultAzureCredential + except ImportError: + raise ImportError( + "You must `pip install adlfs azure-identity` " + "to fetch URIs in Azure Blob File System Secure. " + + cls._MISSING_DEPENDENCIES_WARNING + ) + + def open_file(uri, mode, *, transport_params=None): + # Parse and validate the ABFSS URI + parsed = urlparse(uri) + + # Validate ABFSS URI format: abfss://container@account.dfs.core.windows.net/path + if not parsed.netloc or "@" not in parsed.netloc: + raise ValueError( + f"Invalid ABFSS URI format - missing container@account: {uri}" + ) + + container_part, hostname_part = parsed.netloc.split("@", 1) + + # Validate container name (must be non-empty) + if not container_part: + raise ValueError( + f"Invalid ABFSS URI format - empty container name: {uri}" + ) + + # Validate hostname format + if not hostname_part or not hostname_part.endswith(".dfs.core.windows.net"): + raise ValueError( + f"Invalid ABFSS URI format - invalid hostname (must end with .dfs.core.windows.net): {uri}" + ) + + # Extract and validate account name + azure_storage_account_name = hostname_part.split(".")[0] + if not azure_storage_account_name: + raise ValueError( + f"Invalid ABFSS URI format - empty account name: {uri}" + ) + + # Handle ABFSS URI with adlfs + filesystem = adlfs.AzureBlobFileSystem( + account_name=azure_storage_account_name, + credential=DefaultAzureCredential(), + ) + return filesystem.open(uri, mode) + + return open_file, None + + @classmethod + def download_remote_uri(cls, protocol: str, source_uri: str, dest_file: str): + """Download file from remote URI to destination file. + + Args: + protocol: The protocol to use for downloading (e.g., 's3', 'https'). + source_uri: The source URI to download from. + dest_file: The destination file path to save to. + + Raises: + ImportError: If required dependencies for the protocol are not installed. + """ + assert protocol in cls.get_remote_protocols() + + tp = None + open_file = None + + if protocol == "file": + source_uri = source_uri[len("file://") :] + + def open_file(uri, mode, *, transport_params=None): + return open(uri, mode) + + elif protocol == "s3": + open_file, tp = cls._handle_s3_protocol() + elif protocol == "gs": + open_file, tp = cls._handle_gs_protocol() + elif protocol == "azure": + open_file, tp = cls._handle_azure_protocol() + elif protocol == "abfss": + open_file, tp = cls._handle_abfss_protocol() + else: + try: + from smart_open import open as open_file + except ImportError: + raise ImportError( + "You must `pip install smart_open` " + f"to fetch {protocol.upper()} URIs. " + + cls._MISSING_DEPENDENCIES_WARNING + ) + + with open_file(source_uri, "rb", transport_params=tp) as fin: + with open(dest_file, "wb") as fout: + fout.write(fin.read()) + + +Protocol = enum.Enum( + "Protocol", + {protocol.upper(): protocol for protocol in ProtocolsProvider.get_protocols()}, +) + + +@classmethod +def _remote_protocols(cls): + # Returns a list of protocols that support remote storage + # These protocols should only be used with paths that end in ".zip" or ".whl" + return [ + cls[protocol.upper()] for protocol in ProtocolsProvider.get_remote_protocols() + ] + + +Protocol.remote_protocols = _remote_protocols + + +def _download_remote_uri(self, source_uri, dest_file): + return ProtocolsProvider.download_remote_uri(self.value, source_uri, dest_file) + + +Protocol.download_remote_uri = _download_remote_uri diff --git a/lib/python3.12/site-packages/ray/_private/runtime_env/py_executable.py b/lib/python3.12/site-packages/ray/_private/runtime_env/py_executable.py new file mode 100644 index 0000000000000000000000000000000000000000..0a57344909042a6237cd6ac0f1f4c760809db1f2 --- /dev/null +++ b/lib/python3.12/site-packages/ray/_private/runtime_env/py_executable.py @@ -0,0 +1,45 @@ +import logging +from typing import List, Optional + +from ray._private.runtime_env.context import RuntimeEnvContext +from ray._private.runtime_env.plugin import RuntimeEnvPlugin + +default_logger = logging.getLogger(__name__) + + +class PyExecutablePlugin(RuntimeEnvPlugin): + """This plugin allows running Ray workers with a custom Python executable. + + You can use it with + `ray.init(runtime_env={"py_executable": " "})`. If you specify + a `working_dir` in the runtime environment, the executable will have access + to the working directory, for example, to a requirements.txt for a package manager, + a script for a debugger, or the executable could be a shell script in the + working directory. You can also use this plugin to run worker processes + in a custom profiler or use a custom Python interpreter or `python` with + custom arguments. + """ + + name = "py_executable" + + def __init__(self): + pass + + async def create( + self, + uri: Optional[str], + runtime_env: "RuntimeEnv", # noqa: F821 + context: RuntimeEnvContext, + logger: logging.Logger = default_logger, + ) -> int: + return 0 + + def modify_context( + self, + uris: List[str], + runtime_env: "RuntimeEnv", # noqa: F821 + context: RuntimeEnvContext, + logger: Optional[logging.Logger] = default_logger, + ): + logger.info("Running py_executable plugin") + context.py_executable = runtime_env.py_executable() diff --git a/lib/python3.12/site-packages/ray/_private/runtime_env/py_modules.py b/lib/python3.12/site-packages/ray/_private/runtime_env/py_modules.py new file mode 100644 index 0000000000000000000000000000000000000000..fc6e03693e02b5fcca564eb7f2280ffd3f44fe81 --- /dev/null +++ b/lib/python3.12/site-packages/ray/_private/runtime_env/py_modules.py @@ -0,0 +1,233 @@ +import logging +import os +from pathlib import Path +from types import ModuleType +from typing import Any, Dict, List, Optional + +from ray._common.utils import try_to_create_directory +from ray._private.runtime_env.context import RuntimeEnvContext +from ray._private.runtime_env.packaging import ( + Protocol, + delete_package, + download_and_unpack_package, + get_local_dir_from_uri, + get_uri_for_directory, + get_uri_for_file, + get_uri_for_package, + install_wheel_package, + is_whl_uri, + package_exists, + parse_uri, + upload_package_if_needed, + upload_package_to_gcs, +) +from ray._private.runtime_env.plugin import RuntimeEnvPlugin +from ray._private.runtime_env.working_dir import set_pythonpath_in_context +from ray._private.utils import get_directory_size_bytes +from ray._raylet import GcsClient +from ray.exceptions import RuntimeEnvSetupError + +default_logger = logging.getLogger(__name__) + + +def _check_is_uri(s: str) -> bool: + try: + protocol, path = parse_uri(s) + except ValueError: + protocol, path = None, None + + if ( + protocol in Protocol.remote_protocols() + and not path.endswith(".zip") + and not path.endswith(".whl") + ): + raise ValueError("Only .zip or .whl files supported for remote URIs.") + + return protocol is not None + + +def upload_py_modules_if_needed( + runtime_env: Dict[str, Any], + scratch_dir: Optional[str] = os.getcwd(), + logger: Optional[logging.Logger] = default_logger, + upload_fn=None, +) -> Dict[str, Any]: + """Uploads the entries in py_modules and replaces them with a list of URIs. + + For each entry that is already a URI, this is a no-op. + """ + py_modules = runtime_env.get("py_modules") + if py_modules is None: + return runtime_env + + if not isinstance(py_modules, list): + raise TypeError( + "py_modules must be a List of local paths, imported modules, or " + f"URIs, got {type(py_modules)}." + ) + + py_modules_uris = [] + for module in py_modules: + if isinstance(module, str): + # module_path is a local path or a URI. + module_path = module + elif isinstance(module, Path): + module_path = str(module) + elif isinstance(module, ModuleType): + if not hasattr(module, "__path__"): + # This is a single-file module. + module_path = module.__file__ + else: + # NOTE(edoakes): Python allows some installed Python packages to + # be split into multiple directories. We could probably handle + # this, but it seems tricky & uncommon. If it's a problem for + # users, we can add this support on demand. + if len(module.__path__) > 1: + raise ValueError( + "py_modules only supports modules whose __path__" + " has length 1 or those who are single-file." + ) + [module_path] = module.__path__ + else: + raise TypeError( + "py_modules must be a list of file paths, URIs, " + f"or imported modules, got {type(module)}." + ) + + if _check_is_uri(module_path): + module_uri = module_path + else: + # module_path is a local path. + if Path(module_path).is_dir() or Path(module_path).suffix == ".py": + is_dir = Path(module_path).is_dir() + excludes = runtime_env.get("excludes", None) + if is_dir: + module_uri = get_uri_for_directory(module_path, excludes=excludes) + else: + module_uri = get_uri_for_file(module_path) + if upload_fn is None: + try: + upload_package_if_needed( + module_uri, + scratch_dir, + module_path, + excludes=excludes, + include_parent_dir=is_dir, + logger=logger, + ) + except Exception as e: + from ray.util.spark.utils import is_in_databricks_runtime + + if is_in_databricks_runtime(): + raise RuntimeEnvSetupError( + f"Failed to upload module {module_path} to the Ray " + f"cluster, please ensure there are only files under " + f"the module path, notebooks under the path are " + f"not allowed, original exception: {e}" + ) from e + raise RuntimeEnvSetupError( + f"Failed to upload module {module_path} to the Ray " + f"cluster: {e}" + ) from e + else: + upload_fn(module_path, excludes=excludes) + elif Path(module_path).suffix == ".whl": + module_uri = get_uri_for_package(Path(module_path)) + if upload_fn is None: + if not package_exists(module_uri): + try: + upload_package_to_gcs( + module_uri, Path(module_path).read_bytes() + ) + except Exception as e: + raise RuntimeEnvSetupError( + f"Failed to upload {module_path} to the Ray " + f"cluster: {e}" + ) from e + else: + upload_fn(module_path, excludes=None, is_file=True) + else: + raise ValueError( + "py_modules entry must be a .py file, " + "a directory, or a .whl file; " + f"got {module_path}" + ) + + py_modules_uris.append(module_uri) + + # TODO(architkulkarni): Expose a single URI for py_modules. This plugin + # should internally handle the "sub-URIs", the individual modules. + + runtime_env["py_modules"] = py_modules_uris + return runtime_env + + +class PyModulesPlugin(RuntimeEnvPlugin): + + name = "py_modules" + + def __init__(self, resources_dir: str, gcs_client: GcsClient): + self._resources_dir = os.path.join(resources_dir, "py_modules_files") + self._gcs_client = gcs_client + try_to_create_directory(self._resources_dir) + + def _get_local_dir_from_uri(self, uri: str): + return get_local_dir_from_uri(uri, self._resources_dir) + + def delete_uri( + self, uri: str, logger: Optional[logging.Logger] = default_logger + ) -> int: + """Delete URI and return the number of bytes deleted.""" + logger.info("Got request to delete pymodule URI %s", uri) + local_dir = get_local_dir_from_uri(uri, self._resources_dir) + local_dir_size = get_directory_size_bytes(local_dir) + + deleted = delete_package(uri, self._resources_dir) + if not deleted: + logger.warning(f"Tried to delete nonexistent URI: {uri}.") + return 0 + + return local_dir_size + + def get_uris(self, runtime_env) -> List[str]: + return runtime_env.py_modules() + + async def create( + self, + uri: str, + runtime_env: "RuntimeEnv", # noqa: F821 + context: RuntimeEnvContext, + logger: Optional[logging.Logger] = default_logger, + ) -> int: + + module_dir = await download_and_unpack_package( + uri, self._resources_dir, self._gcs_client, logger=logger + ) + + if is_whl_uri(uri): + wheel_uri = module_dir + module_dir = self._get_local_dir_from_uri(uri) + await install_wheel_package( + wheel_uri=wheel_uri, target_dir=module_dir, logger=logger + ) + + return get_directory_size_bytes(module_dir) + + def modify_context( + self, + uris: List[str], + runtime_env_dict: Dict, + context: RuntimeEnvContext, + logger: Optional[logging.Logger] = default_logger, + ): + module_dirs = [] + for uri in uris: + module_dir = self._get_local_dir_from_uri(uri) + if not module_dir.exists(): + raise ValueError( + f"Local directory {module_dir} for URI {uri} does " + "not exist on the cluster. Something may have gone wrong while " + "downloading, unpacking or installing the py_modules files." + ) + module_dirs.append(str(module_dir)) + set_pythonpath_in_context(os.pathsep.join(module_dirs), context) diff --git a/lib/python3.12/site-packages/ray/_private/runtime_env/rocprof_sys.py b/lib/python3.12/site-packages/ray/_private/runtime_env/rocprof_sys.py new file mode 100644 index 0000000000000000000000000000000000000000..a05626de1a0e336561cf1c5c6768899dc8c82ca7 --- /dev/null +++ b/lib/python3.12/site-packages/ray/_private/runtime_env/rocprof_sys.py @@ -0,0 +1,173 @@ +import asyncio +import copy +import logging +import os +import subprocess +import sys +from pathlib import Path +from typing import Dict, List, Optional, Tuple + +from ray._common.utils import try_to_create_directory +from ray._private.runtime_env.context import RuntimeEnvContext +from ray._private.runtime_env.plugin import RuntimeEnvPlugin +from ray.exceptions import RuntimeEnvSetupError + +default_logger = logging.getLogger(__name__) + +# rocprof-sys config used when runtime_env={"_rocprof_sys": "default"} +# Refer to the following link for more information on rocprof-sys options +# https://rocm.docs.amd.com/projects/rocprofiler-systems/en/docs-6.4.0/how-to/understanding-rocprof-sys-output.html +ROCPROFSYS_DEFAULT_CONFIG = { + "env": { + "ROCPROFSYS_TIME_OUTPUT": "false", + "ROCPROFSYS_OUTPUT_PREFIX": "worker_process_%p", + }, + "args": { + "F": "true", + }, +} + + +def parse_rocprof_sys_config( + rocprof_sys_config: Dict[str, str] +) -> Tuple[List[str], List[str]]: + """ + Function to convert dictionary of rocprof-sys options into + rocprof-sys-python command line + + The function returns: + - List[str]: rocprof-sys-python cmd line split into list of str + """ + rocprof_sys_cmd = ["rocprof-sys-python"] + rocprof_sys_env = {} + if "args" in rocprof_sys_config: + # Parse rocprof-sys arg options + for option, option_val in rocprof_sys_config["args"].items(): + # option standard based on + # https://www.gnu.org/software/libc/manual/html_node/Argument-Syntax.html + if len(option) > 1: + rocprof_sys_cmd.append(f"--{option}={option_val}") + else: + rocprof_sys_cmd += [f"-{option}", option_val] + if "env" in rocprof_sys_config: + rocprof_sys_env = rocprof_sys_config["env"] + rocprof_sys_cmd.append("--") + return rocprof_sys_cmd, rocprof_sys_env + + +class RocProfSysPlugin(RuntimeEnvPlugin): + name = "_rocprof_sys" + + def __init__(self, resources_dir: str): + self.rocprof_sys_cmd = [] + self.rocprof_sys_env = {} + + # replace this with better way to get logs dir + session_dir, runtime_dir = os.path.split(resources_dir) + self._rocprof_sys_dir = Path(session_dir) / "logs" / "rocprof_sys" + try_to_create_directory(self._rocprof_sys_dir) + + async def _check_rocprof_sys_script( + self, rocprof_sys_config: Dict[str, str] + ) -> Tuple[bool, str]: + """ + Function to validate if rocprof_sys_config is a valid rocprof_sys profile options + Args: + rocprof_sys_config: dictionary mapping rocprof_sys option to it's value + Returns: + a tuple consists of a boolean indicating if the rocprof_sys_config + is valid option and an error message if the rocprof_sys_config is invalid + """ + + # use empty as rocprof_sys report test filename + test_folder = str(Path(self._rocprof_sys_dir) / "test") + rocprof_sys_cmd, rocprof_sys_env = parse_rocprof_sys_config(rocprof_sys_config) + rocprof_sys_env_copy = copy.deepcopy(rocprof_sys_env) + rocprof_sys_env_copy["ROCPROFSYS_OUTPUT_PATH"] = test_folder + rocprof_sys_env_copy.update(os.environ) + try_to_create_directory(test_folder) + + # Create a test python file to run rocprof_sys + with open(f"{test_folder}/test.py", "w") as f: + f.write("import time\n") + try: + rocprof_sys_cmd = rocprof_sys_cmd + [f"{test_folder}/test.py"] + process = await asyncio.create_subprocess_exec( + *rocprof_sys_cmd, + env=rocprof_sys_env_copy, + stdout=subprocess.PIPE, + stderr=subprocess.PIPE, + ) + stdout, stderr = await process.communicate() + error_msg = stderr.strip() if stderr.strip() != "" else stdout.strip() + + # cleanup temp file + clean_up_cmd = ["rm", "-r", test_folder] + cleanup_process = await asyncio.create_subprocess_exec( + *clean_up_cmd, + stdout=subprocess.PIPE, + stderr=subprocess.PIPE, + ) + _, _ = await cleanup_process.communicate() + if process.returncode == 0: + return True, None + else: + return False, error_msg + except FileNotFoundError: + return False, ("rocprof_sys is not installed") + + async def create( + self, + uri: Optional[str], + runtime_env: "RuntimeEnv", # noqa: F821 + context: RuntimeEnvContext, + logger: logging.Logger = default_logger, + ) -> int: + rocprof_sys_config = runtime_env.rocprof_sys() + if not rocprof_sys_config: + return 0 + + if rocprof_sys_config and sys.platform != "linux": + raise RuntimeEnvSetupError("rocprof-sys CLI is only available in Linux.\n") + + if isinstance(rocprof_sys_config, str): + if rocprof_sys_config == "default": + rocprof_sys_config = ROCPROFSYS_DEFAULT_CONFIG + else: + raise RuntimeEnvSetupError( + f"Unsupported rocprof_sys config: {rocprof_sys_config}. " + "The supported config is 'default' or " + "Dictionary of rocprof_sys options" + ) + + is_valid_rocprof_sys_config, error_msg = await self._check_rocprof_sys_script( + rocprof_sys_config + ) + if not is_valid_rocprof_sys_config: + logger.warning(error_msg) + raise RuntimeEnvSetupError( + "rocprof-sys profile failed to run with the following " + f"error message:\n {error_msg}" + ) + # add set output path to logs dir + if "env" not in rocprof_sys_config: + rocprof_sys_config["env"] = {} + rocprof_sys_config["env"]["ROCPROFSYS_OUTPUT_PATH"] = str( + Path(self._rocprof_sys_dir) + ) + + self.rocprof_sys_cmd, self.rocprof_sys_env = parse_rocprof_sys_config( + rocprof_sys_config + ) + return 0 + + def modify_context( + self, + uris: List[str], + runtime_env: "RuntimeEnv", # noqa: F821 + context: RuntimeEnvContext, + logger: Optional[logging.Logger] = default_logger, + ): + logger.info("Running rocprof-sys profiler") + context.py_executable = " ".join(self.rocprof_sys_cmd) + context.env_vars.update(self.rocprof_sys_env) diff --git a/lib/python3.12/site-packages/ray/_private/runtime_env/setup_hook.py b/lib/python3.12/site-packages/ray/_private/runtime_env/setup_hook.py new file mode 100644 index 0000000000000000000000000000000000000000..e6c1283df274cb4e3937abaffdd64a1dc6b1e2f5 --- /dev/null +++ b/lib/python3.12/site-packages/ray/_private/runtime_env/setup_hook.py @@ -0,0 +1,193 @@ +import base64 +import logging +import os +import traceback +from typing import Any, Callable, Dict, Optional, Union + +import ray +import ray._private.ray_constants as ray_constants +import ray.cloudpickle as pickle +from ray._common.utils import load_class +from ray.runtime_env import RuntimeEnv + +logger = logging.getLogger(__name__) + +RUNTIME_ENV_FUNC_IDENTIFIER = "ray_runtime_env_func::" + + +def get_import_export_timeout(): + return int( + os.environ.get( + ray_constants.RAY_WORKER_PROCESS_SETUP_HOOK_LOAD_TIMEOUT_ENV_VAR, "60" + ) + ) + + +def _decode_function_key(key: bytes) -> str: + # b64encode only includes A-Z, a-z, 0-9, + and / characters + return RUNTIME_ENV_FUNC_IDENTIFIER + base64.b64encode(key).decode() + + +def _encode_function_key(key: str) -> bytes: + assert key.startswith(RUNTIME_ENV_FUNC_IDENTIFIER) + return base64.b64decode(key[len(RUNTIME_ENV_FUNC_IDENTIFIER) :]) + + +def export_setup_func_callable( + runtime_env: Union[Dict[str, Any], RuntimeEnv], + setup_func: Callable, + worker: "ray.Worker", +) -> Union[Dict[str, Any], RuntimeEnv]: + assert isinstance(setup_func, Callable) + try: + key = worker.function_actor_manager.export_setup_func( + setup_func, timeout=get_import_export_timeout() + ) + except Exception as e: + raise ray.exceptions.RuntimeEnvSetupError( + "Failed to export the setup function." + ) from e + env_vars = runtime_env.get("env_vars", {}) + assert ray_constants.WORKER_PROCESS_SETUP_HOOK_ENV_VAR not in env_vars, ( + f"The env var, {ray_constants.WORKER_PROCESS_SETUP_HOOK_ENV_VAR}, " + "is not permitted because it is reserved for the internal use." + ) + env_vars[ray_constants.WORKER_PROCESS_SETUP_HOOK_ENV_VAR] = _decode_function_key( + key + ) + runtime_env["env_vars"] = env_vars + # Note: This field is no-op. We don't have a plugin for the setup hook + # because we can implement it simply using an env var. + # This field is just for the observability purpose, so we store + # the name of the method. + runtime_env["worker_process_setup_hook"] = setup_func.__name__ + return runtime_env + + +def export_setup_func_module( + runtime_env: Union[Dict[str, Any], RuntimeEnv], + setup_func_module: str, +) -> Union[Dict[str, Any], RuntimeEnv]: + assert isinstance(setup_func_module, str) + env_vars = runtime_env.get("env_vars", {}) + assert ray_constants.WORKER_PROCESS_SETUP_HOOK_ENV_VAR not in env_vars, ( + f"The env var, {ray_constants.WORKER_PROCESS_SETUP_HOOK_ENV_VAR}, " + "is not permitted because it is reserved for the internal use." + ) + env_vars[ray_constants.WORKER_PROCESS_SETUP_HOOK_ENV_VAR] = setup_func_module + runtime_env["env_vars"] = env_vars + return runtime_env + + +def upload_worker_process_setup_hook_if_needed( + runtime_env: Union[Dict[str, Any], RuntimeEnv], + worker: "ray.Worker", +) -> Union[Dict[str, Any], RuntimeEnv]: + """Uploads the worker_process_setup_hook to GCS with a key. + + runtime_env["worker_process_setup_hook"] is converted to a decoded key + that can load the worker setup hook function from GCS. + i.e., you can use internalKV.Get(runtime_env["worker_process_setup_hook]) + to access the worker setup hook from GCS. + + Args: + runtime_env: The runtime_env. The value will be modified + when returned. + worker: ray.worker instance. + decoder: GCS requires the function key to be bytes. However, + we cannot json serialize (which is required to serialize + runtime env) the bytes. So the key should be decoded to + a string. The given decoder is used to decode the function + key. + """ + setup_func = runtime_env.get("worker_process_setup_hook") + + if setup_func is None: + return runtime_env + + if isinstance(setup_func, Callable): + return export_setup_func_callable(runtime_env, setup_func, worker) + elif isinstance(setup_func, str): + return export_setup_func_module(runtime_env, setup_func) + else: + raise TypeError( + "worker_process_setup_hook must be a function, " f"got {type(setup_func)}." + ) + + +def load_and_execute_setup_hook( + worker_process_setup_hook_key: str, +) -> Optional[str]: + """Load the setup hook from a given key and execute. + + Args: + worker_process_setup_hook_key: The key to import the setup hook + from GCS. + Returns: + An error message if it fails. None if it succeeds. + """ + assert worker_process_setup_hook_key is not None + if not worker_process_setup_hook_key.startswith(RUNTIME_ENV_FUNC_IDENTIFIER): + return load_and_execute_setup_hook_module(worker_process_setup_hook_key) + else: + return load_and_execute_setup_hook_func(worker_process_setup_hook_key) + + +def load_and_execute_setup_hook_module( + worker_process_setup_hook_key: str, +) -> Optional[str]: + try: + setup_func = load_class(worker_process_setup_hook_key) + setup_func() + return None + except Exception: + error_message = ( + "Failed to execute the setup hook method, " + f"{worker_process_setup_hook_key} " + "from ``ray.init(runtime_env=" + f"{{'worker_process_setup_hook': {worker_process_setup_hook_key}}})``. " + "Please make sure the given module exists and is available " + "from ray workers. For more details, see the error trace below.\n" + f"{traceback.format_exc()}" + ) + return error_message + + +def load_and_execute_setup_hook_func( + worker_process_setup_hook_key: str, +) -> Optional[str]: + worker = ray._private.worker.global_worker + assert worker.connected + func_manager = worker.function_actor_manager + try: + worker_setup_func_info = func_manager.fetch_registered_method( + _encode_function_key(worker_process_setup_hook_key), + timeout=get_import_export_timeout(), + ) + except Exception: + error_message = ( + "Failed to import setup hook within " + f"{get_import_export_timeout()} seconds.\n" + f"{traceback.format_exc()}" + ) + return error_message + + try: + setup_func = pickle.loads(worker_setup_func_info.function) + except Exception: + error_message = ( + "Failed to deserialize the setup hook method.\n" f"{traceback.format_exc()}" + ) + return error_message + + try: + setup_func() + except Exception: + error_message = ( + f"Failed to execute the setup hook method. Function name:" + f"{worker_setup_func_info.function_name}\n" + f"{traceback.format_exc()}" + ) + return error_message + + return None diff --git a/lib/python3.12/site-packages/ray/_private/runtime_env/uri_cache.py b/lib/python3.12/site-packages/ray/_private/runtime_env/uri_cache.py new file mode 100644 index 0000000000000000000000000000000000000000..430bcf6d7e4e1ffb62e55c7a6d49cf6d0cae0492 --- /dev/null +++ b/lib/python3.12/site-packages/ray/_private/runtime_env/uri_cache.py @@ -0,0 +1,115 @@ +import logging +from typing import Callable, Optional, Set + +default_logger = logging.getLogger(__name__) + +DEFAULT_MAX_URI_CACHE_SIZE_BYTES = (1024**3) * 10 # 10 GB + + +class URICache: + """Caches URIs up to a specified total size limit. + + URIs are represented by strings. Each URI has an associated size on disk. + When a URI is added to the URICache, it is marked as "in use". + When a URI is no longer in use, the user of this class should call + `mark_unused` to signal that the URI is safe for deletion. + + URIs in the cache can be marked as "in use" by calling `mark_used`. + + Deletion of URIs on disk does not occur until the size limit is exceeded. + When this happens, URIs that are not in use are deleted randomly until the + size limit is satisfied, or there are no more URIs that are not in use. + + It is possible for the total size on disk to exceed the size limit if all + the URIs are in use. + + """ + + def __init__( + self, + delete_fn: Optional[Callable[[str, logging.Logger], int]] = None, + max_total_size_bytes: int = DEFAULT_MAX_URI_CACHE_SIZE_BYTES, + debug_mode: bool = False, + ): + # Maps URIs to the size in bytes of their corresponding disk contents. + self._used_uris: Set[str] = set() + self._unused_uris: Set[str] = set() + + if delete_fn is None: + self._delete_fn = lambda uri, logger: 0 + else: + self._delete_fn = delete_fn + + # Total size of both used and unused URIs in the cache. + self._total_size_bytes = 0 + self.max_total_size_bytes = max_total_size_bytes + + # Used in `self._check_valid()` for testing. + self._debug_mode = debug_mode + + def mark_unused(self, uri: str, logger: logging.Logger = default_logger): + """Mark a URI as unused and okay to be deleted.""" + if uri not in self._used_uris: + logger.info(f"URI {uri} is already unused.") + else: + self._unused_uris.add(uri) + self._used_uris.remove(uri) + logger.info(f"Marked URI {uri} unused.") + self._evict_if_needed(logger) + self._check_valid() + + def mark_used(self, uri: str, logger: logging.Logger = default_logger): + """Mark a URI as in use. URIs in use will not be deleted.""" + if uri in self._used_uris: + return + elif uri in self._unused_uris: + self._used_uris.add(uri) + self._unused_uris.remove(uri) + else: + raise ValueError( + f"Got request to mark URI {uri} used, but this " + "URI is not present in the cache." + ) + logger.info(f"Marked URI {uri} used.") + self._check_valid() + + def add(self, uri: str, size_bytes: int, logger: logging.Logger = default_logger): + """Add a URI to the cache and mark it as in use.""" + if uri in self._unused_uris: + self._unused_uris.remove(uri) + + self._used_uris.add(uri) + self._total_size_bytes += size_bytes + + self._evict_if_needed(logger) + self._check_valid() + logger.info(f"Added URI {uri} with size {size_bytes}") + + def get_total_size_bytes(self) -> int: + return self._total_size_bytes + + def _evict_if_needed(self, logger: logging.Logger = default_logger): + """Evict unused URIs (if they exist) until total size <= max size.""" + while ( + self._unused_uris + and self.get_total_size_bytes() > self.max_total_size_bytes + ): + # TODO(architkulkarni): Evict least recently used URI instead + arbitrary_unused_uri = next(iter(self._unused_uris)) + self._unused_uris.remove(arbitrary_unused_uri) + num_bytes_deleted = self._delete_fn(arbitrary_unused_uri, logger) + self._total_size_bytes -= num_bytes_deleted + logger.info( + f"Deleted URI {arbitrary_unused_uri} with size " f"{num_bytes_deleted}." + ) + + def _check_valid(self): + """(Debug mode only) Check "used" and "unused" sets are disjoint.""" + if self._debug_mode: + assert self._used_uris & self._unused_uris == set() + + def __contains__(self, uri): + return uri in self._used_uris or uri in self._unused_uris + + def __repr__(self): + return str(self.__dict__) diff --git a/lib/python3.12/site-packages/ray/_private/runtime_env/utils.py b/lib/python3.12/site-packages/ray/_private/runtime_env/utils.py new file mode 100644 index 0000000000000000000000000000000000000000..a4387ff27b697d4918fdd6c4228a7dfcd7a59cc6 --- /dev/null +++ b/lib/python3.12/site-packages/ray/_private/runtime_env/utils.py @@ -0,0 +1,117 @@ +import asyncio +import itertools +import logging +import subprocess +import textwrap +import types +from typing import List + + +class SubprocessCalledProcessError(subprocess.CalledProcessError): + """The subprocess.CalledProcessError with stripped stdout.""" + + LAST_N_LINES = 50 + + def __init__(self, *args, cmd_index=None, **kwargs): + self.cmd_index = cmd_index + super().__init__(*args, **kwargs) + + @staticmethod + def _get_last_n_line(str_data: str, last_n_lines: int) -> str: + if last_n_lines < 0: + return str_data + lines = str_data.strip().split("\n") + return "\n".join(lines[-last_n_lines:]) + + def __str__(self): + str_list = ( + [] + if self.cmd_index is None + else [f"Run cmd[{self.cmd_index}] failed with the following details."] + ) + str_list.append(super().__str__()) + out = { + "stdout": self.stdout, + "stderr": self.stderr, + } + for name, s in out.items(): + if s: + subtitle = f"Last {self.LAST_N_LINES} lines of {name}:" + last_n_line_str = self._get_last_n_line(s, self.LAST_N_LINES).strip() + str_list.append( + f"{subtitle}\n{textwrap.indent(last_n_line_str, ' ' * 4)}" + ) + return "\n".join(str_list) + + +async def check_output_cmd( + cmd: List[str], + *, + logger: logging.Logger, + cmd_index_gen: types.GeneratorType = itertools.count(1), + **kwargs, +) -> str: + """Run command with arguments and return its output. + + If the return code was non-zero it raises a CalledProcessError. The + CalledProcessError object will have the return code in the returncode + attribute and any output in the output attribute. + + Args: + cmd: The cmdline should be a sequence of program arguments or else + a single string or path-like object. The program to execute is + the first item in cmd. + logger: The logger instance. + cmd_index_gen: The cmd index generator, default is itertools.count(1). + kwargs: All arguments are passed to the create_subprocess_exec. + + Returns: + The stdout of cmd. + + Raises: + CalledProcessError: If the return code of cmd is not 0. + """ + + cmd_index = next(cmd_index_gen) + logger.info("Run cmd[%s] %s", cmd_index, repr(cmd)) + + proc = None + try: + proc = await asyncio.create_subprocess_exec( + *cmd, + stdout=asyncio.subprocess.PIPE, + stderr=asyncio.subprocess.STDOUT, + **kwargs, + ) + # Use communicate instead of polling stdout: + # * Avoid deadlocks due to streams pausing reading or writing and blocking the + # child process. Please refer to: + # https://docs.python.org/3/library/asyncio-subprocess.html#asyncio.asyncio.subprocess.Process.stderr + # * Avoid mixing multiple outputs of concurrent cmds. + stdout, _ = await proc.communicate() + except asyncio.exceptions.CancelledError as e: + # since Python 3.9, when cancelled, the inner process needs to throw as it is + # for asyncio to timeout properly https://bugs.python.org/issue40607 + raise e + except BaseException as e: + raise RuntimeError(f"Run cmd[{cmd_index}] got exception.") from e + else: + stdout = stdout.decode("utf-8") + if stdout: + logger.info("Output of cmd[%s]: %s", cmd_index, stdout) + else: + logger.info("No output for cmd[%s]", cmd_index) + if proc.returncode != 0: + raise SubprocessCalledProcessError( + proc.returncode, cmd, output=stdout, cmd_index=cmd_index + ) + return stdout + finally: + if proc is not None: + # Kill process. + try: + proc.kill() + except ProcessLookupError: + pass + # Wait process exit. + await proc.wait() diff --git a/lib/python3.12/site-packages/ray/_private/runtime_env/uv.py b/lib/python3.12/site-packages/ray/_private/runtime_env/uv.py new file mode 100644 index 0000000000000000000000000000000000000000..e54954a996f0b69688a79da1a961d816d7000776 --- /dev/null +++ b/lib/python3.12/site-packages/ray/_private/runtime_env/uv.py @@ -0,0 +1,345 @@ +"""Util class to install packages via uv. +""" + +import asyncio +import hashlib +import json +import logging +import os +import shutil +import sys +from asyncio import create_task, get_running_loop +from typing import Dict, List, Optional + +from ray._common.utils import try_to_create_directory +from ray._private.runtime_env import dependency_utils, virtualenv_utils +from ray._private.runtime_env.packaging import Protocol, parse_uri +from ray._private.runtime_env.plugin import RuntimeEnvPlugin +from ray._private.runtime_env.utils import check_output_cmd +from ray._private.utils import get_directory_size_bytes + +default_logger = logging.getLogger(__name__) + + +def _get_uv_hash(uv_dict: Dict) -> str: + """Get a deterministic hash value for `uv` related runtime envs.""" + serialized_uv_spec = json.dumps(uv_dict, sort_keys=True) + hash_val = hashlib.sha1(serialized_uv_spec.encode("utf-8")).hexdigest() + return hash_val + + +def get_uri(runtime_env: Dict) -> Optional[str]: + """Return `"uv://"`, or None if no GC required.""" + uv = runtime_env.get("uv") + if uv is not None: + if isinstance(uv, dict): + uri = "uv://" + _get_uv_hash(uv_dict=uv) + elif isinstance(uv, list): + uri = "uv://" + _get_uv_hash(uv_dict=dict(packages=uv)) + else: + raise TypeError( + "uv field received by RuntimeEnvAgent must be " + f"list or dict, not {type(uv).__name__}." + ) + else: + uri = None + return uri + + +class UvProcessor: + def __init__( + self, + target_dir: str, + runtime_env: "RuntimeEnv", # noqa: F821 + logger: Optional[logging.Logger] = default_logger, + ): + try: + import virtualenv # noqa: F401 ensure virtualenv exists. + except ImportError: + raise RuntimeError( + f"Please install virtualenv " + f"`{sys.executable} -m pip install virtualenv`" + f"to enable uv runtime env." + ) + + logger.debug("Setting up uv for runtime_env: %s", runtime_env) + self._target_dir = target_dir + # An empty directory is created to execute cmd. + self._exec_cwd = os.path.join(self._target_dir, "exec_cwd") + self._runtime_env = runtime_env + self._logger = logger + + self._uv_config = self._runtime_env.uv_config() + self._uv_env = os.environ.copy() + self._uv_env.update(self._runtime_env.env_vars()) + + async def _install_uv( + self, path: str, cwd: str, pip_env: dict, logger: logging.Logger + ): + """Before package install, make sure the required version `uv` (if specifieds) + is installed. + """ + virtualenv_path = virtualenv_utils.get_virtualenv_path(path) + python = virtualenv_utils.get_virtualenv_python(path) + + def _get_uv_exec_to_install() -> str: + """Get `uv` executable with version to install.""" + uv_version = self._uv_config.get("uv_version", None) + if uv_version: + return f"uv{uv_version}" + # Use default version. + return "uv" + + uv_install_cmd = [ + python, + "-m", + "pip", + "install", + "--disable-pip-version-check", + "--no-cache-dir", + _get_uv_exec_to_install(), + ] + logger.info("Installing package uv to %s", virtualenv_path) + await check_output_cmd(uv_install_cmd, logger=logger, cwd=cwd, env=pip_env) + + async def _check_uv_existence( + self, path: str, cwd: str, env: dict, logger: logging.Logger + ) -> bool: + """Check and return the existence of `uv` in virtual env.""" + python = virtualenv_utils.get_virtualenv_python(path) + + check_existence_cmd = [ + python, + "-m", + "uv", + "version", + ] + + try: + # If `uv` doesn't exist, exception will be thrown. + await check_output_cmd(check_existence_cmd, logger=logger, cwd=cwd, env=env) + return True + except Exception: + return False + + async def _uv_check(sef, python: str, cwd: str, logger: logging.Logger) -> None: + """Check virtual env dependency compatibility. + If any incompatibility detected, exception will be thrown. + + param: + python: the path for python executable within virtual environment. + """ + cmd = [python, "-m", "uv", "pip", "check"] + await check_output_cmd( + cmd, + logger=logger, + cwd=cwd, + ) + + async def _install_uv_packages( + self, + path: str, + uv_packages: List[str], + cwd: str, + pip_env: Dict, + logger: logging.Logger, + ): + """Install required python packages via `uv`.""" + virtualenv_path = virtualenv_utils.get_virtualenv_path(path) + python = virtualenv_utils.get_virtualenv_python(path) + # TODO(fyrestone): Support -i, --no-deps, --no-cache-dir, ... + requirements_file = dependency_utils.get_requirements_file(path, uv_packages) + + # Check existence for `uv` and see if we could skip `uv` installation. + uv_exists = await self._check_uv_existence(path, cwd, pip_env, logger) + + # Install uv, which acts as the default package manager. + if (not uv_exists) or (self._uv_config.get("uv_version", None) is not None): + await self._install_uv(path, cwd, pip_env, logger) + + # Avoid blocking the event loop. + loop = get_running_loop() + await loop.run_in_executor( + None, dependency_utils.gen_requirements_txt, requirements_file, uv_packages + ) + + # Install all dependencies. + # + # Difference with pip: + # 1. `--disable-pip-version-check` has no effect for uv. + uv_install_cmd = [ + python, + "-m", + "uv", + "pip", + "install", + "-r", + requirements_file, + ] + + uv_opt_list = self._uv_config.get("uv_pip_install_options", ["--no-cache"]) + if uv_opt_list: + uv_install_cmd += uv_opt_list + + logger.info("Installing python requirements to %s", virtualenv_path) + await check_output_cmd(uv_install_cmd, logger=logger, cwd=cwd, env=pip_env) + + # Check python environment for conflicts. + if self._uv_config.get("uv_check", False): + await self._uv_check(python, cwd, logger) + + async def _run(self): + path = self._target_dir + logger = self._logger + uv_packages = self._uv_config["packages"] + # We create an empty directory for exec cmd so that the cmd will + # run more stable. e.g. if cwd has ray, then checking ray will + # look up ray in cwd instead of site packages. + os.makedirs(self._exec_cwd, exist_ok=True) + try: + await virtualenv_utils.create_or_get_virtualenv( + path, self._exec_cwd, logger + ) + python = virtualenv_utils.get_virtualenv_python(path) + async with dependency_utils.check_ray(python, self._exec_cwd, logger): + # Install packages with uv. + await self._install_uv_packages( + path, + uv_packages, + self._exec_cwd, + self._uv_env, + logger, + ) + except Exception: + logger.info("Delete incomplete virtualenv: %s", path) + shutil.rmtree(path, ignore_errors=True) + logger.exception("Failed to install uv packages.") + raise + + def __await__(self): + return self._run().__await__() + + +class UvPlugin(RuntimeEnvPlugin): + name = "uv" + + def __init__(self, resources_dir: str): + self._uv_resource_dir = os.path.join(resources_dir, "uv") + self._creating_task = {} + # Maps a URI to a lock that is used to prevent multiple concurrent + # installs of the same virtualenv, see #24513 + self._create_locks: Dict[str, asyncio.Lock] = {} + # Key: created hashes. Value: size of the uv dir. + self._created_hash_bytes: Dict[str, int] = {} + try_to_create_directory(self._uv_resource_dir) + + def _get_path_from_hash(self, hash_val: str) -> str: + """Generate a path from the hash of a uv spec. + + Example output: + /tmp/ray/session_2021-11-03_16-33-59_356303_41018/runtime_resources + /uv/ray-9a7972c3a75f55e976e620484f58410c920db091 + """ + return os.path.join(self._uv_resource_dir, hash_val) + + def get_uris(self, runtime_env: "RuntimeEnv") -> List[str]: # noqa: F821 + """Return the uv URI from the RuntimeEnv if it exists, else return [].""" + uv_uri = runtime_env.uv_uri() + if uv_uri: + return [uv_uri] + return [] + + def delete_uri( + self, uri: str, logger: Optional[logging.Logger] = default_logger + ) -> int: + """Delete URI and return the number of bytes deleted.""" + logger.info("Got request to delete uv URI %s", uri) + protocol, hash_val = parse_uri(uri) + if protocol != Protocol.UV: + raise ValueError( + "UvPlugin can only delete URIs with protocol " + f"uv. Received protocol {protocol}, URI {uri}" + ) + + # Cancel running create task. + task = self._creating_task.pop(hash_val, None) + if task is not None: + task.cancel() + + del self._created_hash_bytes[hash_val] + + uv_env_path = self._get_path_from_hash(hash_val) + local_dir_size = get_directory_size_bytes(uv_env_path) + del self._create_locks[uri] + try: + shutil.rmtree(uv_env_path) + except OSError as e: + logger.warning(f"Error when deleting uv env {uv_env_path}: {str(e)}") + return 0 + + return local_dir_size + + async def create( + self, + uri: str, + runtime_env: "RuntimeEnv", # noqa: F821 + context: "RuntimeEnvContext", # noqa: F821 + logger: Optional[logging.Logger] = default_logger, + ) -> int: + if not runtime_env.has_uv(): + return 0 + + protocol, hash_val = parse_uri(uri) + target_dir = self._get_path_from_hash(hash_val) + + async def _create_for_hash(): + await UvProcessor( + target_dir, + runtime_env, + logger, + ) + + loop = get_running_loop() + return await loop.run_in_executor( + None, get_directory_size_bytes, target_dir + ) + + if uri not in self._create_locks: + # async lock to prevent the same virtualenv being concurrently installed + self._create_locks[uri] = asyncio.Lock() + + async with self._create_locks[uri]: + if hash_val in self._created_hash_bytes: + return self._created_hash_bytes[hash_val] + self._creating_task[hash_val] = task = create_task(_create_for_hash()) + task.add_done_callback(lambda _: self._creating_task.pop(hash_val, None)) + uv_dir_bytes = await task + self._created_hash_bytes[hash_val] = uv_dir_bytes + return uv_dir_bytes + + def modify_context( + self, + uris: List[str], + runtime_env: "RuntimeEnv", # noqa: F821 + context: "RuntimeEnvContext", # noqa: F821 + logger: logging.Logger = default_logger, + ): + if not runtime_env.has_uv(): + return + # UvPlugin only uses a single URI. + uri = uris[0] + # Update py_executable. + protocol, hash_val = parse_uri(uri) + target_dir = self._get_path_from_hash(hash_val) + virtualenv_python = virtualenv_utils.get_virtualenv_python(target_dir) + + if not os.path.exists(virtualenv_python): + raise ValueError( + f"Local directory {target_dir} for URI {uri} does " + "not exist on the cluster. Something may have gone wrong while " + "installing the runtime_env `uv` packages." + ) + context.py_executable = virtualenv_python + context.command_prefix += virtualenv_utils.get_virtualenv_activate_command( + target_dir + ) diff --git a/lib/python3.12/site-packages/ray/_private/runtime_env/uv_runtime_env_hook.py b/lib/python3.12/site-packages/ray/_private/runtime_env/uv_runtime_env_hook.py new file mode 100644 index 0000000000000000000000000000000000000000..7a72107872f07ab5a1901bcbe07a569612a7c07f --- /dev/null +++ b/lib/python3.12/site-packages/ray/_private/runtime_env/uv_runtime_env_hook.py @@ -0,0 +1,383 @@ +import argparse +import copy +import optparse +import os +import sys +from pathlib import Path +from typing import Any, Dict, List, Optional, Tuple + +import psutil + + +def _create_uv_run_parser(): + """Create and return the argument parser for 'uv run' command.""" + + parser = optparse.OptionParser(prog="uv run", add_help_option=False) + + # Disable interspersed args to stop parsing when we hit the first + # argument that is not recognized by the parser. + parser.disable_interspersed_args() + + # Main options group + main_group = optparse.OptionGroup(parser, "Main options") + main_group.add_option("--extra", action="append", dest="extras") + main_group.add_option("--all-extras", action="store_true") + main_group.add_option("--no-extra", action="append", dest="no_extras") + main_group.add_option("--no-dev", action="store_true") + main_group.add_option("--group", action="append", dest="groups") + main_group.add_option("--no-group", action="append", dest="no_groups") + main_group.add_option("--no-default-groups", action="store_true") + main_group.add_option("--only-group", action="append", dest="only_groups") + main_group.add_option("--all-groups", action="store_true") + main_group.add_option("-m", "--module") + main_group.add_option("--only-dev", action="store_true") + main_group.add_option("--no-editable", action="store_true") + main_group.add_option("--exact", action="store_true") + main_group.add_option("--env-file", action="append", dest="env_files") + main_group.add_option("--no-env-file", action="store_true") + parser.add_option_group(main_group) + + # With options + with_group = optparse.OptionGroup(parser, "With options") + with_group.add_option("--with", action="append", dest="with_packages") + with_group.add_option("--with-editable", action="append", dest="with_editable") + with_group.add_option( + "--with-requirements", action="append", dest="with_requirements" + ) + parser.add_option_group(with_group) + + # Environment options + env_group = optparse.OptionGroup(parser, "Environment options") + env_group.add_option("--isolated", action="store_true") + env_group.add_option("--active", action="store_true") + env_group.add_option("--no-sync", action="store_true") + env_group.add_option("--locked", action="store_true") + env_group.add_option("--frozen", action="store_true") + parser.add_option_group(env_group) + + # Script options + script_group = optparse.OptionGroup(parser, "Script options") + script_group.add_option("-s", "--script", action="store_true") + script_group.add_option("--gui-script", action="store_true") + parser.add_option_group(script_group) + + # Workspace options + workspace_group = optparse.OptionGroup(parser, "Workspace options") + workspace_group.add_option("--all-packages", action="store_true") + workspace_group.add_option("--package") + workspace_group.add_option("--no-project", action="store_true") + parser.add_option_group(workspace_group) + + # Index options + index_group = optparse.OptionGroup(parser, "Index options") + index_group.add_option("--index", action="append", dest="indexes") + index_group.add_option("--default-index") + index_group.add_option("-i", "--index-url") + index_group.add_option( + "--extra-index-url", action="append", dest="extra_index_urls" + ) + index_group.add_option("-f", "--find-links", action="append", dest="find_links") + index_group.add_option("--no-index", action="store_true") + index_group.add_option( + "--index-strategy", + type="choice", + choices=["first-index", "unsafe-first-match", "unsafe-best-match"], + ) + index_group.add_option( + "--keyring-provider", type="choice", choices=["disabled", "subprocess"] + ) + parser.add_option_group(index_group) + + # Resolver options + resolver_group = optparse.OptionGroup(parser, "Resolver options") + resolver_group.add_option("-U", "--upgrade", action="store_true") + resolver_group.add_option( + "-P", "--upgrade-package", action="append", dest="upgrade_packages" + ) + resolver_group.add_option( + "--resolution", type="choice", choices=["highest", "lowest", "lowest-direct"] + ) + resolver_group.add_option( + "--prerelease", + type="choice", + choices=[ + "disallow", + "allow", + "if-necessary", + "explicit", + "if-necessary-or-explicit", + ], + ) + resolver_group.add_option( + "--fork-strategy", type="choice", choices=["fewest", "requires-python"] + ) + resolver_group.add_option("--exclude-newer") + resolver_group.add_option("--no-sources", action="store_true") + parser.add_option_group(resolver_group) + + # Installer options + installer_group = optparse.OptionGroup(parser, "Installer options") + installer_group.add_option("--reinstall", action="store_true") + installer_group.add_option( + "--reinstall-package", action="append", dest="reinstall_packages" + ) + installer_group.add_option( + "--link-mode", type="choice", choices=["clone", "copy", "hardlink", "symlink"] + ) + installer_group.add_option("--compile-bytecode", action="store_true") + parser.add_option_group(installer_group) + + # Build options + build_group = optparse.OptionGroup(parser, "Build options") + build_group.add_option( + "-C", "--config-setting", action="append", dest="config_settings" + ) + build_group.add_option("--no-build-isolation", action="store_true") + build_group.add_option( + "--no-build-isolation-package", + action="append", + dest="no_build_isolation_packages", + ) + build_group.add_option("--no-build", action="store_true") + build_group.add_option( + "--no-build-package", action="append", dest="no_build_packages" + ) + build_group.add_option("--no-binary", action="store_true") + build_group.add_option( + "--no-binary-package", action="append", dest="no_binary_packages" + ) + parser.add_option_group(build_group) + + # Cache options + cache_group = optparse.OptionGroup(parser, "Cache options") + cache_group.add_option("-n", "--no-cache", action="store_true") + cache_group.add_option("--cache-dir") + cache_group.add_option("--refresh", action="store_true") + cache_group.add_option( + "--refresh-package", action="append", dest="refresh_packages" + ) + parser.add_option_group(cache_group) + + # Python options + python_group = optparse.OptionGroup(parser, "Python options") + python_group.add_option("-p", "--python") + python_group.add_option("--managed-python", action="store_true") + python_group.add_option("--no-managed-python", action="store_true") + python_group.add_option("--no-python-downloads", action="store_true") + # note: the following is a legacy option and will be removed at some point + # https://github.com/astral-sh/uv/pull/12246 + python_group.add_option( + "--python-preference", + type="choice", + choices=["only-managed", "managed", "system", "only-system"], + ) + parser.add_option_group(python_group) + + # Global options + global_group = optparse.OptionGroup(parser, "Global options") + global_group.add_option("-q", "--quiet", action="count", default=0) + global_group.add_option("-v", "--verbose", action="count", default=0) + global_group.add_option( + "--color", type="choice", choices=["auto", "always", "never"] + ) + global_group.add_option("--native-tls", action="store_true") + global_group.add_option("--offline", action="store_true") + global_group.add_option( + "--allow-insecure-host", action="append", dest="insecure_hosts" + ) + global_group.add_option("--no-progress", action="store_true") + global_group.add_option("--directory") + global_group.add_option("--project") + global_group.add_option("--config-file") + global_group.add_option("--no-config", action="store_true") + parser.add_option_group(global_group) + + return parser + + +def _parse_args( + parser: optparse.OptionParser, args: List[str] +) -> Tuple[optparse.Values, List[str]]: + """ + Parse the command-line options found in 'args'. + + Replacement for parser.parse_args that handles unknown arguments + by keeping them in the command list instead of erroring and + discarding them. + """ + parser.rargs = args + parser.largs = [] + options = parser.get_default_values() + try: + parser._process_args(parser.largs, parser.rargs, options) + except optparse.BadOptionError as err: + # If we hit an argument that is not recognized, we put it + # back into the unconsumed arguments + parser.rargs = [err.opt_str] + parser.rargs + return options, parser.rargs + + +def _check_working_dir_files( + uv_run_args: optparse.Values, runtime_env: Dict[str, Any] +) -> None: + """ + Check that the files required by uv are local to the working_dir. This catches + the most common cases of how things are different in Ray, i.e. not the whole file + system will be available on the workers, only the working_dir. + + The function won't return anything, it just raises a RuntimeError if there is an error. + """ + working_dir = Path(runtime_env["working_dir"]).resolve() + + # Check if the requirements.txt file is in the working_dir + if uv_run_args.with_requirements: + for requirements_file in uv_run_args.with_requirements: + if not Path(requirements_file).resolve().is_relative_to(working_dir): + raise RuntimeError( + f"You specified --with-requirements={uv_run_args.with_requirements} but " + f"the requirements file is not in the working_dir {runtime_env['working_dir']}, " + "so the workers will not have access to the file. Make sure " + "the requirements file is in the working directory. " + "You can do so by specifying --directory in 'uv run', by changing the current " + "working directory before running 'uv run', or by using the 'working_dir' " + "parameter of the runtime_environment." + ) + + # Check if the pyproject.toml file is in the working_dir + pyproject = None + if uv_run_args.no_project: + pyproject = None + elif uv_run_args.project: + pyproject = Path(uv_run_args.project) + else: + # Walk up the directory tree until pyproject.toml is found + current_path = Path.cwd().resolve() + while current_path != current_path.parent: + if (current_path / "pyproject.toml").exists(): + pyproject = Path(current_path / "pyproject.toml") + break + current_path = current_path.parent + + if pyproject and not pyproject.resolve().is_relative_to(working_dir): + raise RuntimeError( + f"Your {pyproject.resolve()} is not in the working_dir {runtime_env['working_dir']}, " + "so the workers will not have access to the file. Make sure " + "the pyproject.toml file is in the working directory. " + "You can do so by specifying --directory in 'uv run', by changing the current " + "working directory before running 'uv run', or by using the 'working_dir' " + "parameter of the runtime_environment." + ) + + +def _get_uv_run_cmdline() -> Optional[List[str]]: + """ + Return the command line of the first ancestor process that was run with + "uv run" and None if there is no such ancestor. + + uv spawns the python process as a child process, so we first check the + parent process command line. We also check our parent's parents since + the Ray driver might be run as a subprocess of the 'uv run' process. + """ + parents = psutil.Process().parents() + for parent in parents: + try: + cmdline = parent.cmdline() + if ( + len(cmdline) > 1 + and os.path.basename(cmdline[0]) == "uv" + and cmdline[1] == "run" + ): + return cmdline + except psutil.NoSuchProcess: + continue + except psutil.AccessDenied: + continue + return None + + +def hook(runtime_env: Optional[Dict[str, Any]]) -> Dict[str, Any]: + """Hook that detects if the driver is run in 'uv run' and sets the runtime environment accordingly.""" + + runtime_env = copy.deepcopy(runtime_env) or {} + + cmdline = _get_uv_run_cmdline() + if not cmdline: + # This means the driver was not run in a 'uv run' environment -- in this case + # we leave the runtime environment unchanged + return runtime_env + + # First check that the "uv" and "pip" runtime environments are not used. + if "uv" in runtime_env or "pip" in runtime_env: + raise RuntimeError( + "You are using the 'pip' or 'uv' runtime environments together with " + "'uv run'. These are not compatible since 'uv run' will run the workers " + "in an isolated environment -- please add the 'pip' or 'uv' dependencies to your " + "'uv run' environment e.g. by including them in your pyproject.toml." + ) + + # Extract the arguments uv_run_args of 'uv run' that are not part of the command. + parser = _create_uv_run_parser() + (options, command) = _parse_args(parser, cmdline[2:]) + + if cmdline[-len(command) :] != command: + raise AssertionError( + f"uv run command {command} is not a suffix of command line {cmdline}" + ) + uv_run_args = cmdline[: -len(command)] + + # Remove the "--directory" argument since it has already been taken into + # account when setting the current working directory of the current process. + # Also remove the "--module" argument, since the default_worker.py is + # invoked as a script and not as a module. + parser = argparse.ArgumentParser() + parser.add_argument("--directory") + parser.add_argument("-m", "--module") + _, remaining_uv_run_args = parser.parse_known_args(uv_run_args) + + runtime_env["py_executable"] = " ".join(remaining_uv_run_args) + + # If the user specified a working_dir, we always honor it, otherwise + # use the same working_dir that uv run would use + if "working_dir" not in runtime_env: + runtime_env["working_dir"] = os.getcwd() + _check_working_dir_files(options, runtime_env) + + return runtime_env + + +# This __main__ is used for unit testing if the runtime_env_hook picks up the +# right settings. +if __name__ == "__main__": + import json + + test_parser = argparse.ArgumentParser() + test_parser.add_argument("--extra-args", action="store_true") + test_parser.add_argument("runtime_env") + args = test_parser.parse_args() + + # If the env variable is set, add one more level of subprocess indirection + if os.environ.get("RAY_TEST_UV_ADD_SUBPROCESS_INDIRECTION") == "1": + import subprocess + + env = os.environ.copy() + env.pop("RAY_TEST_UV_ADD_SUBPROCESS_INDIRECTION") + subprocess.check_call([sys.executable] + sys.argv, env=env) + sys.exit(0) + + # If the following env variable is set, we use multiprocessing + # spawn to start the subprocess, since it uses a different way to + # modify the command line than subprocess.check_call + if os.environ.get("RAY_TEST_UV_MULTIPROCESSING_SPAWN") == "1": + import multiprocessing + + multiprocessing.set_start_method("spawn") + pool = multiprocessing.Pool(processes=1) + runtime_env = json.loads(args.runtime_env) + print(json.dumps(pool.apply(hook, (runtime_env,)))) + sys.exit(0) + + # We purposefully modify sys.argv here to make sure the hook is robust + # against such modification. + sys.argv.pop(1) + runtime_env = json.loads(args.runtime_env) + print(json.dumps(hook(runtime_env))) diff --git a/lib/python3.12/site-packages/ray/_private/runtime_env/validation.py b/lib/python3.12/site-packages/ray/_private/runtime_env/validation.py new file mode 100644 index 0000000000000000000000000000000000000000..0d5670409ef4090aecd9c278a348bc7cb6178ba8 --- /dev/null +++ b/lib/python3.12/site-packages/ray/_private/runtime_env/validation.py @@ -0,0 +1,465 @@ +import logging +import sys +from collections import OrderedDict +from pathlib import Path +from typing import Dict, List, Optional, Union + +import yaml + +from ray._private.path_utils import is_path +from ray._private.runtime_env.packaging import parse_path + +logger = logging.getLogger(__name__) + + +def validate_path(path: str) -> None: + """Parse the path to ensure it is well-formed and exists.""" + parse_path(path) + + +def validate_uri(uri: str): + try: + from ray._private.runtime_env.packaging import Protocol, parse_uri + + protocol, path = parse_uri(uri) + except ValueError: + raise ValueError( + f"{uri} is not a valid URI. Passing directories or modules to " + "be dynamically uploaded is only supported at the job level " + "(i.e., passed to `ray.init`)." + ) + + if ( + protocol in Protocol.remote_protocols() + and not path.endswith(".zip") + and not path.endswith(".whl") + ): + raise ValueError("Only .zip or .whl files supported for remote URIs.") + + +def _handle_local_deps_requirement_file(requirements_file: str): + """Read the given [requirements_file], and return all required dependencies.""" + requirements_path = Path(requirements_file) + if not requirements_path.is_file(): + raise ValueError(f"{requirements_path} is not a valid file") + return requirements_path.read_text().strip().split("\n") + + +def validate_py_modules_uris(py_modules_uris: List[str]) -> List[str]: + """Parses and validates a 'py_modules' option. + + Expects py_modules to be a list of URIs. + """ + if not isinstance(py_modules_uris, list): + raise TypeError( + "`py_modules` must be a list of strings, got " f"{type(py_modules_uris)}." + ) + + for module in py_modules_uris: + + if not isinstance(module, str): + raise TypeError("`py_module` must be a string, got " f"{type(module)}.") + + validate_uri(module) + + +def parse_and_validate_py_modules(py_modules: List[str]) -> List[str]: + """Parses and validates a 'py_modules' option. + + Expects py_modules to be a list of local paths or URIs. + """ + if not isinstance(py_modules, list): + raise TypeError( + "`py_modules` must be a list of strings, got " f"{type(py_modules)}." + ) + + for module in py_modules: + + if not isinstance(module, str): + raise TypeError("`py_module` must be a string, got " f"{type(module)}.") + + if is_path(module): + validate_path(module) + else: + validate_uri(module) + + return py_modules + + +def validate_working_dir_uri(working_dir_uri: str) -> str: + """Parses and validates a 'working_dir' option.""" + if not isinstance(working_dir_uri, str): + raise TypeError( + "`working_dir` must be a string, got " f"{type(working_dir_uri)}." + ) + + validate_uri(working_dir_uri) + + +def parse_and_validate_working_dir(working_dir: str) -> str: + """Parses and validates a 'working_dir' option. + + This can be a URI or a path. + """ + assert working_dir is not None + + if not isinstance(working_dir, str): + raise TypeError("`working_dir` must be a string, got " f"{type(working_dir)}.") + + if is_path(working_dir): + validate_path(working_dir) + else: + validate_uri(working_dir) + + return working_dir + + +def parse_and_validate_conda(conda: Union[str, dict]) -> Union[str, dict]: + """Parses and validates a user-provided 'conda' option. + + Conda can be one of three cases: + 1) A dictionary describing the env. This is passed through directly. + 2) A string referring to the name of a preinstalled conda env. + 3) A string pointing to a local conda YAML file. This is detected + by looking for a '.yaml' or '.yml' suffix. In this case, the file + will be read as YAML and passed through as a dictionary. + """ + assert conda is not None + + if sys.platform == "win32": + logger.warning( + "runtime environment support is experimental on Windows. " + "If you run into issues please file a report at " + "https://github.com/ray-project/ray/issues." + ) + + result = conda + if isinstance(conda, str): + file_path = Path(conda) + if file_path.suffix in (".yaml", ".yml"): + if not file_path.is_file(): + raise ValueError(f"Can't find conda YAML file {file_path}.") + try: + result = yaml.safe_load(file_path.read_text()) + except Exception as e: + raise ValueError(f"Failed to read conda file {file_path}: {e}.") + elif file_path.is_absolute(): + if not file_path.is_dir(): + raise ValueError(f"Can't find conda env directory {file_path}.") + result = str(file_path) + elif isinstance(conda, dict): + result = conda + else: + raise TypeError( + "runtime_env['conda'] must be of type str or " f"dict, got {type(conda)}." + ) + + return result + + +def parse_and_validate_uv(uv: Union[str, List[str], Dict]) -> Optional[Dict]: + """Parses and validates a user-provided 'uv' option. + + The value of the input 'uv' field can be one of two cases: + 1) A List[str] describing the requirements. This is passed through. + Example usage: ["tensorflow", "requests"] + 2) a string containing the path to a local pip “requirements.txt” file. + 3) A python dictionary that has one field: + a) packages (required, List[str]): a list of uv packages, it same as 1). + b) uv_check (optional, bool): whether to enable pip check at the end of uv + install, default to False. + c) uv_version (optional, str): user provides a specific uv to use; if + unspecified, default version of uv will be used. + d) uv_pip_install_options (optional, List[str]): user-provided options for + `uv pip install` command, default to ["--no-cache"]. + + The returned parsed value will be a list of packages. If a Ray library + (e.g. "ray[serve]") is specified, it will be deleted and replaced by its + dependencies (e.g. "uvicorn", "requests"). + """ + assert uv is not None + if sys.platform == "win32": + logger.warning( + "runtime environment support is experimental on Windows. " + "If you run into issues please file a report at " + "https://github.com/ray-project/ray/issues." + ) + + result: str = "" + if isinstance(uv, str): + uv_list = _handle_local_deps_requirement_file(uv) + result = dict(packages=uv_list, uv_check=False) + elif isinstance(uv, list) and all(isinstance(dep, str) for dep in uv): + result = dict(packages=uv, uv_check=False) + elif isinstance(uv, dict): + if set(uv.keys()) - { + "packages", + "uv_check", + "uv_version", + "uv_pip_install_options", + }: + raise ValueError( + "runtime_env['uv'] can only have these fields: " + "packages, uv_check, uv_version and uv_pip_install_options, but got: " + f"{list(uv.keys())}" + ) + if "packages" not in uv: + raise ValueError( + f"runtime_env['uv'] must include field 'packages', but got {uv}" + ) + if "uv_check" in uv and not isinstance(uv["uv_check"], bool): + raise TypeError( + "runtime_env['uv']['uv_check'] must be of type bool, " + f"got {type(uv['uv_check'])}" + ) + if "uv_version" in uv and not isinstance(uv["uv_version"], str): + raise TypeError( + "runtime_env['uv']['uv_version'] must be of type str, " + f"got {type(uv['uv_version'])}" + ) + if "uv_pip_install_options" in uv: + if not isinstance(uv["uv_pip_install_options"], list): + raise TypeError( + "runtime_env['uv']['uv_pip_install_options'] must be of type " + f"list[str] got {type(uv['uv_pip_install_options'])}" + ) + # Check each item in installation option. + for idx, cur_opt in enumerate(uv["uv_pip_install_options"]): + if not isinstance(cur_opt, str): + raise TypeError( + "runtime_env['uv']['uv_pip_install_options'] must be of type " + f"list[str] got {type(cur_opt)} for {idx}-th item." + ) + + result = uv.copy() + result["uv_check"] = uv.get("uv_check", False) + result["uv_pip_install_options"] = uv.get( + "uv_pip_install_options", ["--no-cache"] + ) + if not isinstance(uv["packages"], list): + raise ValueError( + "runtime_env['uv']['packages'] must be of type list, " + f"got: {type(uv['packages'])}" + ) + else: + raise TypeError( + "runtime_env['uv'] must be of type " f"List[str], or dict, got {type(uv)}" + ) + + # Deduplicate packages for package lists. + result["packages"] = list(OrderedDict.fromkeys(result["packages"])) + + if len(result["packages"]) == 0: + result = None + logger.debug(f"Rewrote runtime_env `uv` field from {uv} to {result}.") + return result + + +def parse_and_validate_pip(pip: Union[str, List[str], Dict]) -> Optional[Dict]: + """Parses and validates a user-provided 'pip' option. + + The value of the input 'pip' field can be one of two cases: + 1) A List[str] describing the requirements. This is passed through. + 2) A string pointing to a local requirements file. In this case, the + file contents will be read split into a list. + 3) A python dictionary that has three fields: + a) packages (required, List[str]): a list of pip packages, it same as 1). + b) pip_check (optional, bool): whether to enable pip check at the end of pip + install, default to False. + c) pip_version (optional, str): the version of pip, ray will spell + the package name 'pip' in front of the `pip_version` to form the final + requirement string, the syntax of a requirement specifier is defined in + full in PEP 508. + d) pip_install_options (optional, List[str]): user-provided options for + `pip install` command, defaults to ["--disable-pip-version-check", "--no-cache-dir"]. + + The returned parsed value will be a list of pip packages. If a Ray library + (e.g. "ray[serve]") is specified, it will be deleted and replaced by its + dependencies (e.g. "uvicorn", "requests"). + """ + assert pip is not None + result = None + + if sys.platform == "win32": + logger.warning( + "runtime environment support is experimental on Windows. " + "If you run into issues please file a report at " + "https://github.com/ray-project/ray/issues." + ) + if isinstance(pip, str): + # We have been given a path to a requirements.txt file. + pip_list = _handle_local_deps_requirement_file(pip) + result = dict( + packages=pip_list, + pip_check=False, + ) + elif isinstance(pip, list) and all(isinstance(dep, str) for dep in pip): + result = dict(packages=pip, pip_check=False) + elif isinstance(pip, dict): + if set(pip.keys()) - { + "packages", + "pip_check", + "pip_install_options", + "pip_version", + }: + raise ValueError( + "runtime_env['pip'] can only have these fields: " + "packages, pip_check, pip_install_options and pip_version, but got: " + f"{list(pip.keys())}" + ) + + if "pip_check" in pip and not isinstance(pip["pip_check"], bool): + raise TypeError( + "runtime_env['pip']['pip_check'] must be of type bool, " + f"got {type(pip['pip_check'])}" + ) + if "pip_version" in pip: + if not isinstance(pip["pip_version"], str): + raise TypeError( + "runtime_env['pip']['pip_version'] must be of type str, " + f"got {type(pip['pip_version'])}" + ) + if "pip_install_options" in pip: + if not isinstance(pip["pip_install_options"], list): + raise TypeError( + "runtime_env['pip']['pip_install_options'] must be of type " + f"list[str] got {type(pip['pip_install_options'])}" + ) + # Check each item in installation option. + for idx, cur_opt in enumerate(pip["pip_install_options"]): + if not isinstance(cur_opt, str): + raise TypeError( + "runtime_env['pip']['pip_install_options'] must be of type " + f"list[str] got {type(cur_opt)} for {idx}-th item." + ) + + result = pip.copy() + # Contrary to pip_check, we do not insert the default value of pip_install_options. + # This is to maintain backwards compatibility with ray==2.0.1 + result["pip_check"] = pip.get("pip_check", False) + + if "packages" not in pip: + raise ValueError( + f"runtime_env['pip'] must include field 'packages', but got {pip}" + ) + elif isinstance(pip["packages"], str): + result["packages"] = _handle_local_deps_requirement_file(pip["packages"]) + elif not isinstance(pip["packages"], list): + raise ValueError( + "runtime_env['pip']['packages'] must be of type str of list, " + f"got: {type(pip['packages'])}" + ) + else: + raise TypeError( + "runtime_env['pip'] must be of type str or " f"List[str], got {type(pip)}" + ) + + # Eliminate duplicates to prevent `pip install` from erroring. Use + # OrderedDict to preserve the order of the list. This makes the output + # deterministic and easier to debug, because pip install can have + # different behavior depending on the order of the input. + result["packages"] = list(OrderedDict.fromkeys(result["packages"])) + + if len(result["packages"]) == 0: + result = None + + logger.debug(f"Rewrote runtime_env `pip` field from {pip} to {result}.") + + return result + + +def parse_and_validate_container(container: List[str]) -> List[str]: + """Parses and validates a user-provided 'container' option. + + This is passed through without validation (for now). + """ + assert container is not None + return container + + +def parse_and_validate_excludes(excludes: List[str]) -> List[str]: + """Parses and validates a user-provided 'excludes' option. + + This is validated to verify that it is of type List[str]. + + If an empty list is passed, we return `None` for consistency. + """ + assert excludes is not None + + if isinstance(excludes, list) and len(excludes) == 0: + return None + + if isinstance(excludes, list) and all(isinstance(path, str) for path in excludes): + return excludes + else: + raise TypeError( + "runtime_env['excludes'] must be of type " + f"List[str], got {type(excludes)}" + ) + + +def parse_and_validate_env_vars(env_vars: Dict[str, str]) -> Optional[Dict[str, str]]: + """Parses and validates a user-provided 'env_vars' option. + + This is validated to verify that all keys and vals are strings. + + If an empty dictionary is passed, we return `None` for consistency. + + Args: + env_vars: A dictionary of environment variables to set in the + runtime environment. + + Returns: + The validated env_vars dictionary, or None if it was empty. + + Raises: + TypeError: If the env_vars is not a dictionary of strings. The error message + will include the type of the invalid value. + """ + assert env_vars is not None + if len(env_vars) == 0: + return None + + if not isinstance(env_vars, dict): + raise TypeError( + "runtime_env['env_vars'] must be of type " + f"Dict[str, str], got {type(env_vars)}" + ) + + for key, val in env_vars.items(): + if not isinstance(key, str): + raise TypeError( + "runtime_env['env_vars'] must be of type " + f"Dict[str, str], but the key {key} is of type {type(key)}" + ) + if not isinstance(val, str): + raise TypeError( + "runtime_env['env_vars'] must be of type " + f"Dict[str, str], but the value {val} is of type {type(val)}" + ) + + return env_vars + + +# Dictionary mapping runtime_env options with the function to parse and +# validate them. +OPTION_TO_VALIDATION_FN = { + "py_modules": parse_and_validate_py_modules, + "working_dir": parse_and_validate_working_dir, + "excludes": parse_and_validate_excludes, + "conda": parse_and_validate_conda, + "pip": parse_and_validate_pip, + "uv": parse_and_validate_uv, + "env_vars": parse_and_validate_env_vars, + "container": parse_and_validate_container, +} + +# RuntimeEnv can be created with local paths +# for these options. However, after the packages +# for these options have been uploaded to GCS, +# they must be URIs. These functions provide the ability +# to validate that these options only contain well-formed URIs. +OPTION_TO_NO_PATH_VALIDATION_FN = { + "working_dir": validate_working_dir_uri, + "py_modules": validate_py_modules_uris, +} diff --git a/lib/python3.12/site-packages/ray/_private/runtime_env/virtualenv_utils.py b/lib/python3.12/site-packages/ray/_private/runtime_env/virtualenv_utils.py new file mode 100644 index 0000000000000000000000000000000000000000..afd3daa352ba07b894926c4cc8df8031fa382411 --- /dev/null +++ b/lib/python3.12/site-packages/ray/_private/runtime_env/virtualenv_utils.py @@ -0,0 +1,110 @@ +"""Utils to detect runtime environment.""" + +import logging +import os +import sys +from typing import List + +from ray._private.runtime_env.utils import check_output_cmd + +_WIN32 = os.name == "nt" + + +def is_in_virtualenv() -> bool: + # virtualenv <= 16.7.9 sets the real_prefix, + # virtualenv > 16.7.9 & venv set the base_prefix. + # So, we check both of them here. + # https://github.com/pypa/virtualenv/issues/1622#issuecomment-586186094 + return hasattr(sys, "real_prefix") or ( + hasattr(sys, "base_prefix") and sys.base_prefix != sys.prefix + ) + + +def get_virtualenv_path(target_dir: str) -> str: + """Get virtual environment path.""" + return os.path.join(target_dir, "virtualenv") + + +def get_virtualenv_python(target_dir: str) -> str: + virtualenv_path = get_virtualenv_path(target_dir) + if _WIN32: + return os.path.join(virtualenv_path, "Scripts", "python.exe") + else: + return os.path.join(virtualenv_path, "bin", "python") + + +def get_virtualenv_activate_command(target_dir: str) -> List[str]: + """Get the command to activate virtual environment.""" + virtualenv_path = get_virtualenv_path(target_dir) + if _WIN32: + cmd = [os.path.join(virtualenv_path, "Scripts", "activate.bat")] + else: + cmd = ["source", os.path.join(virtualenv_path, "bin/activate")] + return cmd + ["1>&2", "&&"] + + +async def create_or_get_virtualenv(path: str, cwd: str, logger: logging.Logger): + """Create or get a virtualenv from path.""" + python = sys.executable + virtualenv_path = os.path.join(path, "virtualenv") + virtualenv_app_data_path = os.path.join(path, "virtualenv_app_data") + + if _WIN32: + current_python_dir = sys.prefix + env = os.environ.copy() + else: + current_python_dir = os.path.abspath( + os.path.join(os.path.dirname(python), "..") + ) + env = {} + + if is_in_virtualenv(): + # virtualenv-clone homepage: + # https://github.com/edwardgeorge/virtualenv-clone + # virtualenv-clone Usage: + # virtualenv-clone /path/to/existing/venv /path/to/cloned/ven + # or + # python -m clonevirtualenv /path/to/existing/venv /path/to/cloned/ven + clonevirtualenv = os.path.join(os.path.dirname(__file__), "_clonevirtualenv.py") + create_venv_cmd = [ + python, + clonevirtualenv, + current_python_dir, + virtualenv_path, + ] + logger.info("Cloning virtualenv %s to %s", current_python_dir, virtualenv_path) + else: + # virtualenv options: + # https://virtualenv.pypa.io/en/latest/cli_interface.html + # + # --app-data + # --reset-app-data + # Set an empty separated app data folder for current virtualenv. + # + # --no-periodic-update + # Disable the periodic (once every 14 days) update of the embedded + # wheels. + # + # --system-site-packages + # Inherit site packages. + # + # --no-download + # Never download the latest pip/setuptools/wheel from PyPI. + create_venv_cmd = [ + python, + "-m", + "virtualenv", + "--app-data", + virtualenv_app_data_path, + "--reset-app-data", + "--no-periodic-update", + "--system-site-packages", + "--no-download", + virtualenv_path, + ] + logger.info( + "Creating virtualenv at %s, current python dir %s", + virtualenv_path, + virtualenv_path, + ) + await check_output_cmd(create_venv_cmd, logger=logger, cwd=cwd, env=env) diff --git a/lib/python3.12/site-packages/ray/_private/runtime_env/working_dir.py b/lib/python3.12/site-packages/ray/_private/runtime_env/working_dir.py new file mode 100644 index 0000000000000000000000000000000000000000..59fa66fac83b54a2a7b6e9bb2fc37fce9b3fe1a8 --- /dev/null +++ b/lib/python3.12/site-packages/ray/_private/runtime_env/working_dir.py @@ -0,0 +1,232 @@ +import logging +import os +from contextlib import contextmanager +from pathlib import Path +from typing import Any, Callable, Dict, List, Optional + +import ray._private.ray_constants as ray_constants +from ray._common.utils import try_to_create_directory +from ray._private.runtime_env.context import RuntimeEnvContext +from ray._private.runtime_env.packaging import ( + Protocol, + delete_package, + download_and_unpack_package, + get_local_dir_from_uri, + get_uri_for_directory, + get_uri_for_package, + parse_uri, + upload_package_if_needed, + upload_package_to_gcs, +) +from ray._private.runtime_env.plugin import RuntimeEnvPlugin +from ray._private.utils import get_directory_size_bytes +from ray._raylet import GcsClient +from ray.exceptions import RuntimeEnvSetupError + +default_logger = logging.getLogger(__name__) + +_WIN32 = os.name == "nt" + + +def upload_working_dir_if_needed( + runtime_env: Dict[str, Any], + scratch_dir: Optional[str] = os.getcwd(), + logger: Optional[logging.Logger] = default_logger, + upload_fn: Optional[Callable[[str, Optional[List[str]]], None]] = None, +) -> Dict[str, Any]: + """Uploads the working_dir and replaces it with a URI. + + If the working_dir is already a URI, this is a no-op. + """ + working_dir = runtime_env.get("working_dir") + if working_dir is None: + return runtime_env + + if not isinstance(working_dir, str) and not isinstance(working_dir, Path): + raise TypeError( + "working_dir must be a string or Path (either a local path " + f"or remote URI), got {type(working_dir)}." + ) + + if isinstance(working_dir, Path): + working_dir = str(working_dir) + + # working_dir is already a URI -- just pass it through. + try: + protocol, path = parse_uri(working_dir) + except ValueError: + protocol, path = None, None + + if protocol is not None: + if protocol in Protocol.remote_protocols() and not path.endswith(".zip"): + raise ValueError("Only .zip files supported for remote URIs.") + return runtime_env + + excludes = runtime_env.get("excludes", None) + try: + working_dir_uri = get_uri_for_directory(working_dir, excludes=excludes) + except ValueError: # working_dir is not a directory + package_path = Path(working_dir) + if not package_path.exists() or package_path.suffix != ".zip": + raise ValueError( + f"directory {package_path} must be an existing " + "directory or a zip package" + ) + + pkg_uri = get_uri_for_package(package_path) + try: + upload_package_to_gcs(pkg_uri, package_path.read_bytes()) + except Exception as e: + raise RuntimeEnvSetupError( + f"Failed to upload package {package_path} to the Ray cluster: {e}" + ) from e + runtime_env["working_dir"] = pkg_uri + return runtime_env + if upload_fn is None: + try: + upload_package_if_needed( + working_dir_uri, + scratch_dir, + working_dir, + include_parent_dir=False, + excludes=excludes, + logger=logger, + ) + except Exception as e: + raise RuntimeEnvSetupError( + f"Failed to upload working_dir {working_dir} to the Ray cluster: {e}" + ) from e + else: + upload_fn(working_dir, excludes=excludes) + + runtime_env["working_dir"] = working_dir_uri + return runtime_env + + +def set_pythonpath_in_context(python_path: str, context: RuntimeEnvContext): + """Insert the path as the first entry in PYTHONPATH in the runtime env. + + This is compatible with users providing their own PYTHONPATH in env_vars, + and is also compatible with the existing PYTHONPATH in the cluster. + + The import priority is as follows: + this python_path arg > env_vars PYTHONPATH > existing cluster env PYTHONPATH. + """ + if "PYTHONPATH" in context.env_vars: + python_path += os.pathsep + context.env_vars["PYTHONPATH"] + if "PYTHONPATH" in os.environ: + python_path += os.pathsep + os.environ["PYTHONPATH"] + context.env_vars["PYTHONPATH"] = python_path + + +class WorkingDirPlugin(RuntimeEnvPlugin): + + name = "working_dir" + + # Note working_dir is not following the priority order of other plugins. Instead + # it's specially treated to happen before all other plugins. + priority = 5 + + def __init__(self, resources_dir: str, gcs_client: GcsClient): + self._resources_dir = os.path.join(resources_dir, "working_dir_files") + self._gcs_client = gcs_client + try_to_create_directory(self._resources_dir) + + def delete_uri( + self, uri: str, logger: Optional[logging.Logger] = default_logger + ) -> int: + """Delete URI and return the number of bytes deleted.""" + logger.info("Got request to delete working dir URI %s", uri) + local_dir = get_local_dir_from_uri(uri, self._resources_dir) + local_dir_size = get_directory_size_bytes(local_dir) + + deleted = delete_package(uri, self._resources_dir) + if not deleted: + logger.warning(f"Tried to delete nonexistent URI: {uri}.") + return 0 + + return local_dir_size + + def get_uris(self, runtime_env: "RuntimeEnv") -> List[str]: # noqa: F821 + working_dir_uri = runtime_env.working_dir() + if working_dir_uri != "": + return [working_dir_uri] + return [] + + async def create( + self, + uri: Optional[str], + runtime_env: dict, + context: RuntimeEnvContext, + logger: logging.Logger = default_logger, + ) -> int: + local_dir = await download_and_unpack_package( + uri, + self._resources_dir, + self._gcs_client, + logger=logger, + overwrite=True, + ) + return get_directory_size_bytes(local_dir) + + def modify_context( + self, + uris: List[str], + runtime_env_dict: Dict, + context: RuntimeEnvContext, + logger: Optional[logging.Logger] = default_logger, + ): + if not uris: + return + + # WorkingDirPlugin uses a single URI. + uri = uris[0] + local_dir = get_local_dir_from_uri(uri, self._resources_dir) + if not local_dir.exists(): + raise ValueError( + f"Local directory {local_dir} for URI {uri} does " + "not exist on the cluster. Something may have gone wrong while " + "downloading or unpacking the working_dir." + ) + + if not _WIN32: + context.command_prefix += ["cd", str(local_dir), "&&"] + else: + # Include '/d' incase temp folder is on different drive than Ray install. + context.command_prefix += ["cd", "/d", f"{local_dir}", "&&"] + set_pythonpath_in_context(python_path=str(local_dir), context=context) + + @contextmanager + def with_working_dir_env(self, uri): + """ + If uri is not None, add the local working directory to the environment variable + as "RAY_RUNTIME_ENV_CREATE_WORKING_DIR". This is useful for other plugins to + create their environment with reference to the working directory. For example + `pip -r ${RAY_RUNTIME_ENV_CREATE_WORKING_DIR}/requirements.txt` + + The environment variable is removed after the context manager exits. + """ + if uri is None: + yield + else: + local_dir = get_local_dir_from_uri(uri, self._resources_dir) + if not local_dir.exists(): + raise ValueError( + f"Local directory {local_dir} for URI {uri} does " + "not exist on the cluster. Something may have gone wrong while " + "downloading or unpacking the working_dir." + ) + key = ray_constants.RAY_RUNTIME_ENV_CREATE_WORKING_DIR_ENV_VAR + prev = os.environ.get(key) + # Windows backslash paths are weird. When it's passed to the env var, and + # when Pip expands it, the backslashes are interpreted as escape characters + # and messes up the whole path. So we convert it to forward slashes. + # This works at least for all Python applications, including pip. + os.environ[key] = local_dir.as_posix() + try: + yield + finally: + if prev is None: + del os.environ[key] + else: + os.environ[key] = prev diff --git a/lib/python3.12/site-packages/ray/_private/services.py b/lib/python3.12/site-packages/ray/_private/services.py new file mode 100644 index 0000000000000000000000000000000000000000..d00a7bd30cad3bc8fbe6e120b4746e43431c9859 --- /dev/null +++ b/lib/python3.12/site-packages/ray/_private/services.py @@ -0,0 +1,2398 @@ +import base64 +import collections +import errno +import io +import json +import logging +import mmap +import multiprocessing +import os +import shutil +import signal +import socket +import subprocess +import sys +import time +from pathlib import Path +from typing import IO, AnyStr, List, Optional + +from filelock import FileLock + +# Ray modules +import ray +import ray._private.ray_constants as ray_constants +from ray._common.network_utils import ( + build_address, + get_localhost_ip, + is_ipv6, + node_ip_address_from_perspective, + parse_address, +) +from ray._private.ray_constants import RAY_NODE_IP_FILENAME +from ray._private.resource_isolation_config import ResourceIsolationConfig +from ray._raylet import GcsClient, GcsClientOptions +from ray.core.generated.common_pb2 import Language + +# Import psutil after ray so the packaged version is used. +import psutil + +resource = None +if sys.platform != "win32": + _timeout = 30 +else: + _timeout = 60 + +EXE_SUFFIX = ".exe" if sys.platform == "win32" else "" + +# True if processes are run in the valgrind profiler. +RUN_RAYLET_PROFILER = False + +# Location of the redis server. +RAY_HOME = os.path.join(os.path.dirname(os.path.dirname(__file__)), "..", "..") +RAY_PATH = os.path.abspath(os.path.dirname(os.path.dirname(__file__))) +RAY_PRIVATE_DIR = "_private" +AUTOSCALER_PRIVATE_DIR = os.path.join("autoscaler", "_private") +AUTOSCALER_V2_DIR = os.path.join("autoscaler", "v2") + +# Location of the raylet executables. +RAYLET_EXECUTABLE = os.path.join( + RAY_PATH, "core", "src", "ray", "raylet", "raylet" + EXE_SUFFIX +) +GCS_SERVER_EXECUTABLE = os.path.join( + RAY_PATH, "core", "src", "ray", "gcs", "gcs_server" + EXE_SUFFIX +) + +JEMALLOC_SO = os.path.join(RAY_PATH, "core", "libjemalloc.so") + +JEMALLOC_SO = JEMALLOC_SO if os.path.exists(JEMALLOC_SO) else None + +# Location of the cpp default worker executables. +DEFAULT_WORKER_EXECUTABLE = os.path.join(RAY_PATH, "cpp", "default_worker" + EXE_SUFFIX) + +# Location of the native libraries. +DEFAULT_NATIVE_LIBRARY_PATH = os.path.join(RAY_PATH, "cpp", "lib") + +DASHBOARD_DEPENDENCY_ERROR_MESSAGE = ( + "Not all Ray Dashboard dependencies were " + "found. To use the dashboard please " + "install Ray using `pip install " + "ray[default]`." +) + +RAY_JEMALLOC_LIB_PATH = "RAY_JEMALLOC_LIB_PATH" +RAY_JEMALLOC_CONF = "RAY_JEMALLOC_CONF" +RAY_JEMALLOC_PROFILE = "RAY_JEMALLOC_PROFILE" + +# Comma separated name of components that will run memory profiler. +# Ray uses `memray` to memory profile internal components. +# The name of the component must be one of ray_constants.PROCESS_TYPE*. +RAY_MEMRAY_PROFILE_COMPONENT_ENV = "RAY_INTERNAL_MEM_PROFILE_COMPONENTS" +# Options to specify for `memray run` command. See +# `memray run --help` for more details. +# Example: +# RAY_INTERNAL_MEM_PROFILE_OPTIONS="--live,--live-port,3456,-q," +# -> `memray run --live --live-port 3456 -q` +RAY_MEMRAY_PROFILE_OPTIONS_ENV = "RAY_INTERNAL_MEM_PROFILE_OPTIONS" + +# Logger for this module. It should be configured at the entry point +# into the program using Ray. Ray provides a default configuration at +# entry/init points. +logger = logging.getLogger(__name__) + +ProcessInfo = collections.namedtuple( + "ProcessInfo", + [ + "process", + "stdout_file", + "stderr_file", + "use_valgrind", + "use_gdb", + "use_valgrind_profiler", + "use_perftools_profiler", + "use_tmux", + ], +) + + +def _site_flags() -> List[str]: + """Detect whether flags related to site packages are enabled for the current + interpreter. To run Ray in hermetic build environments, it helps to pass these flags + down to Python workers. + """ + flags = [] + # sys.flags hidden behind helper methods for unit testing. + if _no_site(): + flags.append("-S") + if _no_user_site(): + flags.append("-s") + return flags + + +# sys.flags hidden behind helper methods for unit testing. +def _no_site(): + return sys.flags.no_site + + +# sys.flags hidden behind helper methods for unit testing. +def _no_user_site(): + return sys.flags.no_user_site + + +def _build_python_executable_command_memory_profileable( + component: str, session_dir: str, unbuffered: bool = True +): + """Build the Python executable command. + + It runs a memory profiler if env var is configured. + + Args: + component: Name of the component. It must be one of + ray_constants.PROCESS_TYPE*. + session_dir: The directory name of the Ray session. + unbuffered: If true, Python executable is started with unbuffered option. + e.g., `-u`. + It means the logs are flushed immediately (good when there's a failure), + but writing to a log file can be slower. + """ + command = [ + sys.executable, + ] + if unbuffered: + command.append("-u") + components_to_memory_profile = os.getenv(RAY_MEMRAY_PROFILE_COMPONENT_ENV, "") + if not components_to_memory_profile: + return command + + components_to_memory_profile = set(components_to_memory_profile.split(",")) + try: + import memray # noqa: F401 + except ImportError: + raise ImportError( + "Memray is required to memory profiler on components " + f"{components_to_memory_profile}. Run `pip install memray`." + ) + if component in components_to_memory_profile: + session_dir = Path(session_dir) + session_name = session_dir.name + profile_dir = session_dir / "profile" + profile_dir.mkdir(exist_ok=True) + output_file_path = profile_dir / f"{session_name}_memory_{component}.bin" + options = os.getenv(RAY_MEMRAY_PROFILE_OPTIONS_ENV, None) + options = options.split(",") if options else [] + # If neither --live nor any output option (-o/--output) is specified, add the default output path + if not any(opt in options for opt in ("--live", "-o", "--output")): + options[0:0] = ["-o", str(output_file_path)] + command.extend(["-m", "memray", "run", *options]) + + return command + + +def _get_gcs_client_options(gcs_server_address): + return GcsClientOptions.create( + gcs_server_address, + None, + allow_cluster_id_nil=True, + fetch_cluster_id_if_nil=False, + ) + + +def serialize_config(config): + return base64.b64encode(json.dumps(config).encode("utf-8")).decode("utf-8") + + +def propagate_jemalloc_env_var( + *, + jemalloc_path: str, + jemalloc_conf: str, + jemalloc_comps: List[str], + process_type: str, +): + """Read the jemalloc memory profiling related + env var and return the dictionary that translates + them to proper jemalloc related env vars. + + For example, if users specify `RAY_JEMALLOC_LIB_PATH`, + it is translated into `LD_PRELOAD` which is needed to + run Jemalloc as a shared library. + + Params: + jemalloc_path: The path to the jemalloc shared library. + jemalloc_conf: `,` separated string of jemalloc config. + jemalloc_comps: The list of Ray components + that we will profile. + process_type: The process type that needs jemalloc + env var for memory profiling. If it doesn't match one of + jemalloc_comps, the function will return an empty dict. + + Returns: + dictionary of {env_var: value} + that are needed to jemalloc profiling. The caller can + call `dict.update(return_value_of_this_func)` to + update the dict of env vars. If the process_type doesn't + match jemalloc_comps, it will return an empty dict. + """ + assert isinstance(jemalloc_comps, list) + assert process_type is not None + process_type = process_type.lower() + if not jemalloc_path: + return {} + + env_vars = { + "LD_PRELOAD": jemalloc_path, + "RAY_LD_PRELOAD_ON_WORKERS": os.environ.get("RAY_LD_PRELOAD_ON_WORKERS", "0"), + } + if process_type in jemalloc_comps and jemalloc_conf: + env_vars.update({"MALLOC_CONF": jemalloc_conf}) + return env_vars + + +class ConsolePopen(subprocess.Popen): + if sys.platform == "win32": + + def terminate(self): + if isinstance(self.stdin, io.IOBase): + self.stdin.close() + if self._use_signals: + self.send_signal(signal.CTRL_BREAK_EVENT) + else: + super(ConsolePopen, self).terminate() + + def __init__(self, *args, **kwargs): + # CREATE_NEW_PROCESS_GROUP is used to send Ctrl+C on Windows: + # https://docs.python.org/3/library/subprocess.html#subprocess.Popen.send_signal + new_pgroup = subprocess.CREATE_NEW_PROCESS_GROUP + flags_to_add = 0 + if ray._private.utils.detect_fate_sharing_support(): + # If we don't have kernel-mode fate-sharing, then don't do this + # because our children need to be in out process group for + # the process reaper to properly terminate them. + flags_to_add = new_pgroup + flags_key = "creationflags" + if flags_to_add: + kwargs[flags_key] = (kwargs.get(flags_key) or 0) | flags_to_add + self._use_signals = kwargs[flags_key] & new_pgroup + super(ConsolePopen, self).__init__(*args, **kwargs) + + +def _find_address_from_flag(flag: str): + """ + Attempts to find all valid Ray addresses on this node, specified by the + flag. + + Params: + flag: `--redis-address` or `--gcs-address` + Returns: + Set of detected addresses. + """ + # Using Redis address `--redis-address` as an example: + # Currently, this extracts the deprecated --redis-address from the command + # that launched the raylet running on this node, if any. Anyone looking to + # edit this function should be warned that these commands look like, for + # example: + # /usr/local/lib/python3.8/dist-packages/ray/core/src/ray/raylet/raylet + # --redis_address=123.456.78.910 --node_ip_address=123.456.78.910 + # --raylet_socket_name=... --store_socket_name=... --object_manager_port=0 + # --min_worker_port=10000 --max_worker_port=19999 + # --node_manager_port=58578 --redis_port=6379 + # --maximum_startup_concurrency=8 + # --static_resource_list=node:123.456.78.910,1.0,object_store_memory,66 + # --config_list=plasma_store_as_thread,True + # --python_worker_command=/usr/bin/python + # /usr/local/lib/python3.8/dist-packages/ray/workers/default_worker.py + # --redis-address=123.456.78.910:6379 + # --node-ip-address=123.456.78.910 --node-manager-port=58578 + # --object-store-name=... --raylet-name=... + # --temp-dir=/tmp/ray + # --metrics-agent-port=41856 --redis-password=[MASKED] + # --java_worker_command= --cpp_worker_command= + # --redis_password=[MASKED] --temp_dir=/tmp/ray --session_dir=... + # --metrics-agent-port=41856 --metrics_export_port=64229 + # --dashboard_agent_command=/usr/bin/python + # -u /usr/local/lib/python3.8/dist-packages/ray/dashboard/agent.py + # --redis-address=123.456.78.910:6379 --metrics-export-port=64229 + # --dashboard-agent-port=41856 --node-manager-port=58578 + # --object-store-name=... --raylet-name=... --temp-dir=/tmp/ray + # --log-dir=/tmp/ray/session_2020-11-08_14-29-07_199128_278000/logs + # --redis-password=[MASKED] --object_store_memory=5037192806 + # --plasma_directory=/tmp + # Longer arguments are elided with ... but all arguments from this instance + # are included, to provide a sense of what is in these. + # Indeed, we had to pull --redis-address to the front of each call to make + # this readable. + # As you can see, this is very long and complex, which is why we can't + # simply extract all the arguments using regular expressions and + # present a dict as if we never lost track of these arguments, for + # example. Picking out --redis-address below looks like it might grab the + # wrong thing, but double-checking that we're finding the correct process + # by checking that the contents look like we expect would probably be prone + # to choking in unexpected ways. + # Notice that --redis-address appears twice. This is not a copy-paste + # error; this is the reason why the for loop below attempts to pick out + # every appearance of --redis-address. + + # The --redis-address here is what is now called the --address, but it + # appears in the default_worker.py and agent.py calls as --redis-address. + addresses = set() + for proc in psutil.process_iter(["cmdline"]): + try: + # HACK: Workaround for UNIX idiosyncrasy + # Normally, cmdline() is supposed to return the argument list. + # But it in some cases (such as when setproctitle is called), + # an arbitrary string resembling a command-line is stored in + # the first argument. + # Explanation: https://unix.stackexchange.com/a/432681 + # More info: https://github.com/giampaolo/psutil/issues/1179 + cmdline = proc.info["cmdline"] + # NOTE(kfstorm): To support Windows, we can't use + # `os.path.basename(cmdline[0]) == "raylet"` here. + + if _is_raylet_process(cmdline): + for arglist in cmdline: + # Given we're merely seeking --redis-address, we just split + # every argument on spaces for now. + for arg in arglist.split(" "): + # TODO(ekl): Find a robust solution for locating Redis. + if arg.startswith(flag): + proc_addr = arg.split("=")[1] + # TODO(mwtian): remove this workaround after Ray + # no longer sets --redis-address to None. + if proc_addr != "" and proc_addr != "None": + addresses.add(proc_addr) + except psutil.AccessDenied: + pass + except psutil.NoSuchProcess: + pass + return addresses + + +def find_gcs_addresses(): + """Finds any local GCS processes based on grepping ps.""" + return _find_address_from_flag("--gcs-address") + + +def find_bootstrap_address(temp_dir: Optional[str]): + """Finds the latest Ray cluster address to connect to, if any. This is the + GCS address connected to by the last successful `ray start`.""" + return ray._private.utils.read_ray_address(temp_dir) + + +def get_ray_address_from_environment(addr: str, temp_dir: Optional[str]): + """Attempts to find the address of Ray cluster to use, in this order: + + 1. Use RAY_ADDRESS if defined and nonempty. + 2. If no address is provided or the provided address is "auto", use the + address in /tmp/ray/ray_current_cluster if available. This will error if + the specified address is None and there is no address found. For "auto", + we will fallback to connecting to any detected Ray cluster (legacy). + 3. Otherwise, use the provided address. + + Returns: + A string to pass into `ray.init(address=...)`, e.g. ip:port, `auto`. + """ + env_addr = os.environ.get(ray_constants.RAY_ADDRESS_ENVIRONMENT_VARIABLE) + if env_addr is not None and env_addr != "": + addr = env_addr + + if addr is not None and addr != "auto": + return addr + # We should try to automatically find an active local instance. + gcs_addrs = find_gcs_addresses() + bootstrap_addr = find_bootstrap_address(temp_dir) + + if len(gcs_addrs) > 1 and bootstrap_addr is not None: + logger.warning( + f"Found multiple active Ray instances: {gcs_addrs}. " + f"Connecting to latest cluster at {bootstrap_addr}. " + "You can override this by setting the `--address` flag " + "or `RAY_ADDRESS` environment variable." + ) + elif len(gcs_addrs) > 0 and addr == "auto": + # Preserve legacy "auto" behavior of connecting to any cluster, even if not + # started with ray start. However if addr is None, we will raise an error. + bootstrap_addr = list(gcs_addrs).pop() + + if bootstrap_addr is None: + if addr is None: + # Caller should start a new instance. + return None + else: + raise ConnectionError( + "Could not find any running Ray instance. " + "Please specify the one to connect to by setting `--address` flag " + "or `RAY_ADDRESS` environment variable." + ) + + return bootstrap_addr + + +def wait_for_node( + gcs_address: str, + node_plasma_store_socket_name: str, + timeout: int = _timeout, +): + """Wait until this node has appeared in the client table. + NOTE: Makes an RPC to the GCS up to every 0.1 seconds to + get all node info. Use only for testing. + + Args: + gcs_address: The gcs address + node_plasma_store_socket_name: The + plasma_store_socket_name for the given node which we wait for. + timeout: The amount of time in seconds to wait before raising an + exception. + + Raises: + TimeoutError: An exception is raised if the timeout expires before + the node appears in the client table. + """ + gcs_options = GcsClientOptions.create( + gcs_address, None, allow_cluster_id_nil=True, fetch_cluster_id_if_nil=False + ) + global_state = ray._private.state.GlobalState() + global_state._initialize_global_state(gcs_options) + start_time = time.time() + while time.time() - start_time < timeout: + clients = global_state.node_table() + object_store_socket_names = [ + client["ObjectStoreSocketName"] for client in clients + ] + if node_plasma_store_socket_name in object_store_socket_names: + return + else: + time.sleep(0.1) + raise TimeoutError( + f"Timed out after {timeout} seconds while waiting for node to startup. " + f"Did not find socket name {node_plasma_store_socket_name} in the list " + "of object store socket names." + ) + + +def get_node_to_connect_for_driver(gcs_address, node_ip_address): + # Get node table from global state accessor. + global_state = ray._private.state.GlobalState() + gcs_options = _get_gcs_client_options(gcs_address) + global_state._initialize_global_state(gcs_options) + return global_state.get_node_to_connect_for_driver(node_ip_address) + + +def get_node(gcs_address, node_id): + """ + Get the node information from the global state accessor. + """ + global_state = ray._private.state.GlobalState() + gcs_options = _get_gcs_client_options(gcs_address) + global_state._initialize_global_state(gcs_options) + return global_state.get_node(node_id) + + +def get_webui_url_from_internal_kv(): + assert ray.experimental.internal_kv._internal_kv_initialized() + webui_url = ray.experimental.internal_kv._internal_kv_get( + "webui:url", namespace=ray_constants.KV_NAMESPACE_DASHBOARD + ) + return ray._common.utils.decode(webui_url) if webui_url is not None else None + + +def remaining_processes_alive(): + """See if the remaining processes are alive or not. + + Note that this ignores processes that have been explicitly killed, + e.g., via a command like node.kill_raylet(). + + Returns: + True if the remaining processes started by ray.init() are alive and + False otherwise. + + Raises: + Exception: An exception is raised if the processes were not started by + ray.init(). + """ + if ray._private.worker._global_node is None: + raise RuntimeError( + "This process is not in a position to determine " + "whether all processes are alive or not." + ) + return ray._private.worker._global_node.remaining_processes_alive() + + +def canonicalize_bootstrap_address( + addr: str, temp_dir: Optional[str] = None +) -> Optional[str]: + """Canonicalizes Ray cluster bootstrap address to host:port. + Reads address from the environment if needed. + + This function should be used to process user supplied Ray cluster address, + via ray.init() or `--address` flags, before using the address to connect. + + Returns: + Ray cluster address string in format or None if the caller + should start a local Ray instance. + """ + if addr is None or addr == "auto": + addr = get_ray_address_from_environment(addr, temp_dir) + if addr is None or addr == "local": + return None + + parsed = parse_address(addr) + if parsed is None: + raise ValueError(f"Invalid address format: {addr}") + host, port = parsed + + try: + bootstrap_host = resolve_ip_for_localhost(host) + except Exception: + logger.exception(f"Failed to convert {addr} to host:port") + raise + return build_address(bootstrap_host, port) + + +def canonicalize_bootstrap_address_or_die( + addr: str, temp_dir: Optional[str] = None +) -> str: + """Canonicalizes Ray cluster bootstrap address to host:port. + + This function should be used when the caller expects there to be an active + and local Ray instance. If no address is provided or address="auto", this + will autodetect the latest Ray instance created with `ray start`. + + For convenience, if no address can be autodetected, this function will also + look for any running local GCS processes, based on pgrep output. This is to + allow easier use of Ray CLIs when debugging a local Ray instance (whose GCS + addresses are not recorded). + + Returns: + Ray cluster address string in format. Throws a + ConnectionError if zero or multiple active Ray instances are + autodetected. + """ + bootstrap_addr = canonicalize_bootstrap_address(addr, temp_dir=temp_dir) + if bootstrap_addr is not None: + return bootstrap_addr + + running_gcs_addresses = find_gcs_addresses() + if len(running_gcs_addresses) == 0: + raise ConnectionError( + "Could not find any running Ray instance. " + "Please specify the one to connect to by setting the `--address` " + "flag or `RAY_ADDRESS` environment variable." + ) + if len(running_gcs_addresses) > 1: + raise ConnectionError( + f"Found multiple active Ray instances: {running_gcs_addresses}. " + "Please specify the one to connect to by setting the `--address` " + "flag or `RAY_ADDRESS` environment variable." + ) + return running_gcs_addresses.pop() + + +def extract_ip_port(bootstrap_address: str): + ip_port = parse_address(bootstrap_address) + if ip_port is None: + raise ValueError( + f"Malformed address {bootstrap_address}. " f"Expected ':'." + ) + ip, port = ip_port + try: + port = int(port) + except ValueError: + raise ValueError(f"Malformed address port {port}. Must be an integer.") + if port < 1024 or port > 65535: + raise ValueError( + f"Invalid address port {port}. Must be between 1024 " + "and 65535 (inclusive)." + ) + return ip, port + + +def resolve_ip_for_localhost(host: str): + """Convert to a remotely reachable IP if the host is "localhost", + "127.0.0.1", or "::1". Otherwise do nothing. + + Args: + host: The hostname or IP address. + + Returns: + The same host but with the local host replaced by remotely + reachable IP. + """ + if not host: + raise ValueError(f"Malformed host: {host}") + if host == "127.0.0.1" or host == "::1" or host == "localhost": + # Make sure localhost isn't resolved to the loopback ip + return get_node_ip_address() + else: + return host + + +# NOTE: This API should not be used when you obtain the +# IP address when ray.init is not called because +# it cannot find the IP address if it is specified by +# ray start --node-ip-address. You should instead use +# get_cached_node_ip_address. +def get_node_ip_address(address=None): + if ray._private.worker._global_node is not None: + return ray._private.worker._global_node.node_ip_address + + if not ray_constants.ENABLE_RAY_CLUSTER: + # Use loopback IP as the local IP address to prevent bothersome + # firewall popups on OSX and Windows. + # https://github.com/ray-project/ray/issues/18730. + return get_localhost_ip() + + return node_ip_address_from_perspective(address) + + +def get_cached_node_ip_address(session_dir: str) -> str: + """Get a node address cached on this session. + + If a ray instance is started by `ray start --node-ip-address`, + the node ip address is cached to a file RAY_NODE_IP_FILENAME. + Otherwise, the file exists, but it is emptyl. + + This API is process-safe, meaning the file access is protected by + a file lock. + + Args: + session_dir: Path to the Ray session directory. + + Returns: + node_ip_address cached on the current node. None if the node + the file doesn't exist, meaning ray instance hasn't been + started on a current node. If node_ip_address is not written + to a file, it means --node-ip-address is not given, and in this + case, we find the IP address ourselves. + """ + file_path = Path(os.path.join(session_dir, RAY_NODE_IP_FILENAME)) + cached_node_ip_address = {} + + with FileLock(str(file_path.absolute()) + ".lock"): + if not file_path.exists(): + return None + + with file_path.open() as f: + cached_node_ip_address.update(json.load(f)) + + if "node_ip_address" in cached_node_ip_address: + return cached_node_ip_address["node_ip_address"] + else: + return ray.util.get_node_ip_address() + + +def write_node_ip_address(session_dir: str, node_ip_address: Optional[str]) -> None: + """Write a node ip address of the current session to + RAY_NODE_IP_FILENAME. + + If a ray instance is started by `ray start --node-ip-address`, + the node ip address is cached to a file RAY_NODE_IP_FILENAME. + + This API is process-safe, meaning the file access is protected by + a file lock. + + The file contains a single string node_ip_address. If nothing + is written, it means --node-ip-address was not given, and Ray + resolves the IP address on its own. It assumes in a single node, + you can have only 1 IP address (which is the assumption ray + has in general). + + node_ip_address is the ip address of the current node. + + Args: + session_dir: The path to Ray session directory. + node_ip_address: The node IP address of the current node. + If None, it means the node ip address is not given + by --node-ip-address. In this case, we don't write + anything to a file. + """ + file_path = Path(os.path.join(session_dir, RAY_NODE_IP_FILENAME)) + cached_node_ip_address = {} + + with FileLock(str(file_path.absolute()) + ".lock"): + if not file_path.exists(): + with file_path.open(mode="w") as f: + json.dump({}, f) + + with file_path.open() as f: + cached_node_ip_address.update(json.load(f)) + + cached_node_ip = cached_node_ip_address.get("node_ip_address") + + if node_ip_address is not None: + if cached_node_ip: + if cached_node_ip == node_ip_address: + # Nothing to do. + return + else: + logger.warning( + "The node IP address of the current host recorded " + f"in {RAY_NODE_IP_FILENAME} ({cached_node_ip}) " + "is different from the current IP address: " + f"{node_ip_address}. Ray will use {node_ip_address} " + "as the current node's IP address. " + "Creating 2 instances in the same host with different " + "IP address is not supported. " + "Please create an enhnacement request to" + "https://github.com/ray-project/ray/issues." + ) + + cached_node_ip_address["node_ip_address"] = node_ip_address + with file_path.open(mode="w") as f: + json.dump(cached_node_ip_address, f) + + +def get_node_instance_id(): + """Get the specified node instance id of the current node. + + Returns: + The node instance id of the current node. + """ + return os.getenv("RAY_CLOUD_INSTANCE_ID", "") + + +def create_redis_client(redis_address, password=None, username=None): + """Create a Redis client. + + Args: + redis_address: The IP address and port of the Redis server. + password: The password for Redis authentication. + username: The username for Redis authentication. + + Returns: + A Redis client. + """ + import redis + + if not hasattr(create_redis_client, "instances"): + create_redis_client.instances = {} + + num_retries = ray_constants.START_REDIS_WAIT_RETRIES + delay = 0.001 + for i in range(num_retries): + cli = create_redis_client.instances.get(redis_address) + if cli is None: + redis_ip_address, redis_port = extract_ip_port( + canonicalize_bootstrap_address_or_die(redis_address) + ) + cli = redis.StrictRedis( + host=redis_ip_address, + port=int(redis_port), + username=username, + password=password, + ) + create_redis_client.instances[redis_address] = cli + try: + cli.ping() + return cli + except Exception as e: + create_redis_client.instances.pop(redis_address) + if i >= num_retries - 1: + raise RuntimeError( + f"Unable to connect to Redis at {redis_address}: {e}" + ) + # Wait a little bit. + time.sleep(delay) + # Make sure the retry interval doesn't increase too large. + delay = min(1, delay * 2) + + +def start_ray_process( + command: List[str], + process_type: str, + fate_share: bool, + env_updates: Optional[dict] = None, + cwd: Optional[str] = None, + use_valgrind: bool = False, + use_gdb: bool = False, + use_valgrind_profiler: bool = False, + use_perftools_profiler: bool = False, + use_tmux: bool = False, + stdout_file: Optional[IO[AnyStr]] = None, + stderr_file: Optional[IO[AnyStr]] = None, + pipe_stdin: bool = False, +): + """Start one of the Ray processes. + + TODO(rkn): We need to figure out how these commands interact. For example, + it may only make sense to start a process in gdb if we also start it in + tmux. Similarly, certain combinations probably don't make sense, like + simultaneously running the process in valgrind and the profiler. + + Args: + command: The command to use to start the Ray process. + process_type: The type of the process that is being started + (e.g., "raylet"). + fate_share: If true, the child will be killed if its parent (us) dies. + True must only be passed after detection of this functionality. + env_updates: A dictionary of additional environment variables to + run the command with (in addition to the caller's environment + variables). + cwd: The directory to run the process in. + use_valgrind: True if we should start the process in valgrind. + use_gdb: True if we should start the process in gdb. + use_valgrind_profiler: True if we should start the process in + the valgrind profiler. + use_perftools_profiler: True if we should profile the process + using perftools. + use_tmux: True if we should start the process in tmux. + stdout_file: A file handle opened for writing to redirect stdout to. If + no redirection should happen, then this should be None. + stderr_file: A file handle opened for writing to redirect stderr to. If + no redirection should happen, then this should be None. + pipe_stdin: If true, subprocess.PIPE will be passed to the process as + stdin. + + Returns: + Information about the process that was started including a handle to + the process that was started. + """ + # Detect which flags are set through environment variables. + valgrind_env_var = f"RAY_{process_type.upper()}_VALGRIND" + if os.environ.get(valgrind_env_var) == "1": + logger.info("Detected environment variable '%s'.", valgrind_env_var) + use_valgrind = True + valgrind_profiler_env_var = f"RAY_{process_type.upper()}_VALGRIND_PROFILER" + if os.environ.get(valgrind_profiler_env_var) == "1": + logger.info("Detected environment variable '%s'.", valgrind_profiler_env_var) + use_valgrind_profiler = True + perftools_profiler_env_var = f"RAY_{process_type.upper()}_PERFTOOLS_PROFILER" + if os.environ.get(perftools_profiler_env_var) == "1": + logger.info("Detected environment variable '%s'.", perftools_profiler_env_var) + use_perftools_profiler = True + tmux_env_var = f"RAY_{process_type.upper()}_TMUX" + if os.environ.get(tmux_env_var) == "1": + logger.info("Detected environment variable '%s'.", tmux_env_var) + use_tmux = True + gdb_env_var = f"RAY_{process_type.upper()}_GDB" + if os.environ.get(gdb_env_var) == "1": + logger.info("Detected environment variable '%s'.", gdb_env_var) + use_gdb = True + # Jemalloc memory profiling. + if os.environ.get("LD_PRELOAD") is None: + jemalloc_lib_path = os.environ.get(RAY_JEMALLOC_LIB_PATH, JEMALLOC_SO) + jemalloc_conf = os.environ.get(RAY_JEMALLOC_CONF) + jemalloc_comps = os.environ.get(RAY_JEMALLOC_PROFILE) + jemalloc_comps = [] if not jemalloc_comps else jemalloc_comps.split(",") + jemalloc_env_vars = propagate_jemalloc_env_var( + jemalloc_path=jemalloc_lib_path, + jemalloc_conf=jemalloc_conf, + jemalloc_comps=jemalloc_comps, + process_type=process_type, + ) + else: + jemalloc_env_vars = {} + + use_jemalloc_mem_profiler = "MALLOC_CONF" in jemalloc_env_vars + + if ( + sum( + [ + use_gdb, + use_valgrind, + use_valgrind_profiler, + use_perftools_profiler, + use_jemalloc_mem_profiler, + ] + ) + > 1 + ): + raise ValueError( + "At most one of the 'use_gdb', 'use_valgrind', " + "'use_valgrind_profiler', 'use_perftools_profiler', " + "and 'use_jemalloc_mem_profiler' flags can " + "be used at a time." + ) + if env_updates is None: + env_updates = {} + if not isinstance(env_updates, dict): + raise ValueError("The 'env_updates' argument must be a dictionary.") + + modified_env = os.environ.copy() + modified_env.update(env_updates) + + if use_gdb: + if not use_tmux: + raise ValueError( + "If 'use_gdb' is true, then 'use_tmux' must be true as well." + ) + + # TODO(suquark): Any better temp file creation here? + gdb_init_path = os.path.join( + ray._common.utils.get_ray_temp_dir(), + f"gdb_init_{process_type}_{time.time()}", + ) + ray_process_path = command[0] + ray_process_args = command[1:] + run_args = " ".join(["'{}'".format(arg) for arg in ray_process_args]) + with open(gdb_init_path, "w") as gdb_init_file: + gdb_init_file.write(f"run {run_args}") + command = ["gdb", ray_process_path, "-x", gdb_init_path] + + if use_valgrind: + command = [ + "valgrind", + "--track-origins=yes", + "--leak-check=full", + "--show-leak-kinds=all", + "--leak-check-heuristics=stdstring", + "--error-exitcode=1", + ] + command + + if use_valgrind_profiler: + command = ["valgrind", "--tool=callgrind"] + command + + if use_perftools_profiler: + modified_env["LD_PRELOAD"] = os.environ["PERFTOOLS_PATH"] + modified_env["CPUPROFILE"] = os.environ["PERFTOOLS_LOGFILE"] + + modified_env.update(jemalloc_env_vars) + + if use_tmux: + # The command has to be created exactly as below to ensure that it + # works on all versions of tmux. (Tested with tmux 1.8-5, travis' + # version, and tmux 2.1) + command = ["tmux", "new-session", "-d", f"{' '.join(command)}"] + + if fate_share: + assert ray._private.utils.detect_fate_sharing_support(), ( + "kernel-level fate-sharing must only be specified if " + "detect_fate_sharing_support() has returned True" + ) + + def preexec_fn(): + import signal + + signal.pthread_sigmask(signal.SIG_BLOCK, {signal.SIGINT}) + if fate_share and sys.platform.startswith("linux"): + ray._private.utils.set_kill_on_parent_death_linux() + + win32_fate_sharing = fate_share and sys.platform == "win32" + # With Windows fate-sharing, we need special care: + # The process must be added to the job before it is allowed to execute. + # Otherwise, there's a race condition: the process might spawn children + # before the process itself is assigned to the job. + # After that point, its children will not be added to the job anymore. + CREATE_SUSPENDED = 0x00000004 # from Windows headers + if sys.platform == "win32": + # CreateProcess, which underlies Popen, is limited to + # 32,767 characters, including the Unicode terminating null + # character + total_chrs = sum([len(x) for x in command]) + if total_chrs > 31766: + raise ValueError( + f"command is limited to a total of 31767 characters, " + f"got {total_chrs}" + ) + + process = ConsolePopen( + command, + env=modified_env, + cwd=cwd, + stdout=stdout_file, + stderr=stderr_file, + stdin=subprocess.PIPE if pipe_stdin else None, + preexec_fn=preexec_fn if sys.platform != "win32" else None, + creationflags=CREATE_SUSPENDED if win32_fate_sharing else 0, + ) + + if win32_fate_sharing: + try: + ray._private.utils.set_kill_child_on_death_win32(process) + psutil.Process(process.pid).resume() + except (psutil.Error, OSError): + process.kill() + raise + + def _get_stream_name(stream): + if stream is not None: + try: + return stream.name + except AttributeError: + return str(stream) + return None + + return ProcessInfo( + process=process, + stdout_file=_get_stream_name(stdout_file), + stderr_file=_get_stream_name(stderr_file), + use_valgrind=use_valgrind, + use_gdb=use_gdb, + use_valgrind_profiler=use_valgrind_profiler, + use_perftools_profiler=use_perftools_profiler, + use_tmux=use_tmux, + ) + + +def start_reaper(fate_share=None): + """Start the reaper process. + + This is a lightweight process that simply + waits for its parent process to die and then terminates its own + process group. This allows us to ensure that ray processes are always + terminated properly so long as that process itself isn't SIGKILLed. + + Returns: + ProcessInfo for the process that was started. + """ + # Make ourselves a process group leader so that the reaper can clean + # up other ray processes without killing the process group of the + # process that started us. + try: + if sys.platform != "win32": + os.setpgrp() + except OSError as e: + errcode = e.errno + if errcode == errno.EPERM and os.getpgrp() == os.getpid(): + # Nothing to do; we're already a session leader. + pass + else: + logger.warning( + f"setpgrp failed, processes may not be cleaned up properly: {e}." + ) + # Don't start the reaper in this case as it could result in killing + # other user processes. + return None + + reaper_filepath = os.path.join(RAY_PATH, RAY_PRIVATE_DIR, "ray_process_reaper.py") + command = [sys.executable, "-u", reaper_filepath] + process_info = start_ray_process( + command, + ray_constants.PROCESS_TYPE_REAPER, + pipe_stdin=True, + fate_share=fate_share, + ) + return process_info + + +def start_log_monitor( + session_dir: str, + logs_dir: str, + gcs_address: str, + fate_share: Optional[bool] = None, + max_bytes: int = 0, + backup_count: int = 0, + stdout_filepath: Optional[str] = None, + stderr_filepath: Optional[str] = None, +): + """Start a log monitor process. + + Args: + session_dir: The session directory. + logs_dir: The directory of logging files. + gcs_address: GCS address for pubsub. + fate_share: Whether to share fate between log_monitor + and this process. + max_bytes: Log rotation parameter. Corresponding to + RotatingFileHandler's maxBytes. + backup_count: Log rotation parameter. Corresponding to + RotatingFileHandler's backupCount. + redirect_logging: Whether we should redirect logging to + the provided log directory. + stdout_filepath: The file path to dump log monitor stdout. + If None, stdout is not redirected. + stderr_filepath: The file path to dump log monitor stderr. + If None, stderr is not redirected. + + Returns: + ProcessInfo for the process that was started. + """ + log_monitor_filepath = os.path.join(RAY_PATH, RAY_PRIVATE_DIR, "log_monitor.py") + + command = [ + sys.executable, + "-u", + log_monitor_filepath, + f"--session-dir={session_dir}", + f"--logs-dir={logs_dir}", + f"--gcs-address={gcs_address}", + f"--logging-rotate-bytes={max_bytes}", + f"--logging-rotate-backup-count={backup_count}", + ] + + if stdout_filepath: + command.append(f"--stdout-filepath={stdout_filepath}") + if stderr_filepath: + command.append(f"--stderr-filepath={stderr_filepath}") + + if stdout_filepath is None and stderr_filepath is None: + # If not redirecting logging to files, unset log filename. + # This will cause log records to go to stderr. + command.append("--logging-filename=") + # Use stderr log format with the component name as a message prefix. + logging_format = ray_constants.LOGGER_FORMAT_STDERR.format( + component=ray_constants.PROCESS_TYPE_LOG_MONITOR + ) + command.append(f"--logging-format={logging_format}") + + stdout_file = None + if stdout_filepath: + stdout_file = open(os.devnull, "w") + + stderr_file = None + if stderr_filepath: + stderr_file = open(os.devnull, "w") + + process_info = start_ray_process( + command, + ray_constants.PROCESS_TYPE_LOG_MONITOR, + stdout_file=stdout_file, + stderr_file=stderr_file, + fate_share=fate_share, + ) + return process_info + + +def start_api_server( + include_dashboard: Optional[bool], + raise_on_failure: bool, + host: str, + gcs_address: str, + cluster_id_hex: str, + node_ip_address: str, + temp_dir: str, + logdir: str, + session_dir: str, + port: Optional[int] = None, + fate_share: Optional[bool] = None, + max_bytes: int = 0, + backup_count: int = 0, + stdout_filepath: Optional[str] = None, + stderr_filepath: Optional[str] = None, +): + """Start a API server process. + + Args: + include_dashboard: If true, this will load all dashboard-related modules + when starting the API server, or fail. If None, it will load all + dashboard-related modules conditioned on dependencies being present. + Otherwise, it will only start the modules that are not relevant to + the dashboard. + raise_on_failure: If true, this will raise an exception + if we fail to start the API server. Otherwise it will print + a warning if we fail to start the API server. + host: The host to bind the dashboard web server to. + gcs_address: The gcs address the dashboard should connect to + cluster_id_hex: Cluster ID in hex. + node_ip_address: The IP address where this is running. + temp_dir: The temporary directory used for log files and + information for this Ray session. + session_dir: The session directory under temp_dir. + It is used as a identifier of individual cluster. + logdir: The log directory used to generate dashboard log. + port: The port to bind the dashboard web server to. + Defaults to 8265. + max_bytes: Log rotation parameter. Corresponding to + RotatingFileHandler's maxBytes. + backup_count: Log rotation parameter. Corresponding to + RotatingFileHandler's backupCount. + stdout_filepath: The file path to dump dashboard stdout. + If None, stdout is not redirected. + stderr_filepath: The file path to dump dashboard stderr. + If None, stderr is not redirected. + + Returns: + A tuple of : + - Dashboard URL if dashboard enabled and started. + - ProcessInfo for the process that was started. + """ + try: + # Make sure port is available. + if port is None: + port_retries = 50 + port = ray_constants.DEFAULT_DASHBOARD_PORT + else: + port_retries = 0 + port_test_socket = socket.socket( + socket.AF_INET6 if is_ipv6(host) else socket.AF_INET, + socket.SOCK_STREAM, + ) + port_test_socket.setsockopt( + socket.SOL_SOCKET, + socket.SO_REUSEADDR, + 1, + ) + try: + port_test_socket.bind((host, port)) + port_test_socket.close() + except socket.error as e: + # 10013 on windows is a bit more broad than just + # "address in use": it can also indicate "permission denied". + # TODO: improve the error message? + if e.errno in {48, 98, 10013}: # address already in use. + raise ValueError( + f"Failed to bind to {host}:{port} because it's " + "already occupied. You can use `ray start " + "--dashboard-port ...` or `ray.init(dashboard_port=..." + ")` to select a different port." + ) + else: + raise e + # Make sure the process can start. + dashboard_dependency_error = ray._private.utils.get_dashboard_dependency_error() + + # Explicitly check here that when the user explicitly specifies + # dashboard inclusion, the install is not minimal. + if include_dashboard and dashboard_dependency_error: + logger.error( + f"Ray dashboard dependencies failed to install properly: {dashboard_dependency_error}.\n" + "Potential causes include:\n" + "1. --include-dashboard is not supported when minimal ray is used. " + "Download ray[default] to use the dashboard.\n" + "2. Dashboard dependencies are conflicting with your python environment. " + "Investigate your python environment and try reinstalling ray[default].\n" + ) + raise Exception("Cannot include dashboard with missing packages.") + + include_dash: bool = True if include_dashboard is None else include_dashboard + + # Start the dashboard process. + dashboard_dir = "dashboard" + dashboard_filepath = os.path.join(RAY_PATH, dashboard_dir, "dashboard.py") + + command = [ + *_build_python_executable_command_memory_profileable( + ray_constants.PROCESS_TYPE_DASHBOARD, + session_dir, + unbuffered=False, + ), + dashboard_filepath, + f"--host={host}", + f"--port={port}", + f"--port-retries={port_retries}", + f"--temp-dir={temp_dir}", + f"--log-dir={logdir}", + f"--session-dir={session_dir}", + f"--logging-rotate-bytes={max_bytes}", + f"--logging-rotate-backup-count={backup_count}", + f"--gcs-address={gcs_address}", + f"--cluster-id-hex={cluster_id_hex}", + f"--node-ip-address={node_ip_address}", + ] + + if stdout_filepath: + command.append(f"--stdout-filepath={stdout_filepath}") + if stderr_filepath: + command.append(f"--stderr-filepath={stderr_filepath}") + + if stdout_filepath is None and stderr_filepath is None: + # If not redirecting logging to files, unset log filename. + # This will cause log records to go to stderr. + command.append("--logging-filename=") + # Use stderr log format with the component name as a message prefix. + logging_format = ray_constants.LOGGER_FORMAT_STDERR.format( + component=ray_constants.PROCESS_TYPE_DASHBOARD + ) + command.append(f"--logging-format={logging_format}") + if dashboard_dependency_error is not None: + command.append("--minimal") + + if not include_dash: + # If dashboard is not included, load modules + # that are irrelevant to the dashboard. + # TODO(sang): Modules like job or state APIs should be + # loaded although dashboard is disabled. Fix it. + command.append("--modules-to-load=UsageStatsHead") + command.append("--disable-frontend") + + stdout_file = None + if stdout_filepath: + stdout_file = open(os.devnull, "w") + + stderr_file = None + if stderr_filepath: + stderr_file = open(os.devnull, "w") + + process_info = start_ray_process( + command, + ray_constants.PROCESS_TYPE_DASHBOARD, + stdout_file=stdout_file, + stderr_file=stderr_file, + fate_share=fate_share, + ) + + # Retrieve the dashboard url + gcs_client = GcsClient(address=gcs_address, cluster_id=cluster_id_hex) + ray.experimental.internal_kv._initialize_internal_kv(gcs_client) + dashboard_url = None + dashboard_returncode = None + start_time_s = time.time() + while ( + time.time() - start_time_s < ray_constants.RAY_DASHBOARD_STARTUP_TIMEOUT_S + ): + dashboard_url = ray.experimental.internal_kv._internal_kv_get( + ray_constants.DASHBOARD_ADDRESS, + namespace=ray_constants.KV_NAMESPACE_DASHBOARD, + ) + if dashboard_url is not None: + dashboard_url = dashboard_url.decode("utf-8") + break + dashboard_returncode = process_info.process.poll() + if dashboard_returncode is not None: + break + + # This is often on the critical path of ray.init() and ray start, + # so we need to poll often. + time.sleep(0.1) + + # Dashboard couldn't be started. + if dashboard_url is None: + returncode_str = ( + f", return code {dashboard_returncode}" + if dashboard_returncode is not None + else "" + ) + logger.error(f"Failed to start the dashboard {returncode_str}") + + def read_log(filename, lines_to_read): + """Read a log file and return the last 20 lines.""" + dashboard_log = os.path.join(logdir, filename) + # Read last n lines of dashboard log. The log file may be large. + lines_to_read = 20 + lines = [] + with open(dashboard_log, "rb") as f: + with mmap.mmap(f.fileno(), 0, access=mmap.ACCESS_READ) as mm: + end = mm.size() + for _ in range(lines_to_read): + sep = mm.rfind(b"\n", 0, end - 1) + if sep == -1: + break + lines.append(mm[sep + 1 : end].decode("utf-8")) + end = sep + lines.append( + f"The last {lines_to_read} lines of {dashboard_log} " + "(it contains the error message from the dashboard): " + ) + return lines + + if logdir: + lines_to_read = 20 + logger.error( + "Error should be written to 'dashboard.log' or " + "'dashboard.err'. We are printing the last " + f"{lines_to_read} lines for you. See " + "'https://docs.ray.io/en/master/ray-observability/user-guides/configure-logging.html#logging-directory-structure' " # noqa + "to find where the log file is." + ) + try: + lines = read_log("dashboard.log", lines_to_read=lines_to_read) + except Exception as e: + logger.error( + f"Couldn't read dashboard.log file. Error: {e}. " + "It means the dashboard is broken even before it " + "initializes the logger (mostly dependency issues). " + "Reading the dashboard.err file which contains stdout/stderr." + ) + # If we cannot read the .log file, we fallback to .err file. + # This is the case where dashboard couldn't be started at all + # and couldn't even initialize the logger to write logs to .log + # file. + try: + lines = read_log("dashboard.err", lines_to_read=lines_to_read) + except Exception as e: + raise Exception( + f"Failed to read dashboard.err file: {e}. " + "It is unexpected. Please report an issue to " + "Ray github. " + "https://github.com/ray-project/ray/issues" + ) + last_log_str = "\n" + "\n".join(reversed(lines[-lines_to_read:])) + raise Exception(last_log_str) + else: + # Is it reachable? + raise Exception("Failed to start a dashboard.") + + if dashboard_dependency_error is not None or not include_dash: + # If it is the minimal installation, the web url (dashboard url) + # shouldn't be configured because it doesn't start a server. + dashboard_url = "" + return dashboard_url, process_info + except Exception as e: + if raise_on_failure: + raise e from e + else: + logger.error(e) + return None, None + + +def get_address(redis_address): + parts = redis_address.split("://", 1) + enable_redis_ssl = False + if len(parts) == 1: + redis_ip_address, redis_port = parse_address(parts[0]) + else: + # rediss for SSL + if len(parts) != 2 or parts[0] not in ("redis", "rediss"): + raise ValueError( + f"Invalid redis address {redis_address}." + "Expected format is ip:port or redis://ip:port, " + "or rediss://ip:port for SSL." + ) + redis_ip_address, redis_port = parse_address(parts[1]) + if parts[0] == "rediss": + enable_redis_ssl = True + return redis_ip_address, redis_port, enable_redis_ssl + + +def start_gcs_server( + redis_address: str, + log_dir: str, + stdout_filepath: Optional[str], + stderr_filepath: Optional[str], + session_name: str, + redis_username: Optional[str] = None, + redis_password: Optional[str] = None, + config: Optional[dict] = None, + fate_share: Optional[bool] = None, + gcs_server_port: Optional[int] = None, + metrics_agent_port: Optional[int] = None, + node_ip_address: Optional[str] = None, +): + """Start a gcs server. + + Args: + redis_address: The address that the Redis server is listening on. + log_dir: The path of the dir where gcs log files are created. + stdout_filepath: The file path to dump gcs server stdout. + If None, stdout is not redirected. + stderr_filepath: The file path to dump gcs server stderr. + If None, stderr is not redirected. + session_name: The current Ray session name. + redis_username: The username of the Redis server. + redis_password: The password of the Redis server. + config: Optional configuration that will + override defaults in RayConfig. + gcs_server_port: Port number of the gcs server. + metrics_agent_port: The port where metrics agent is bound to. + node_ip_address: IP Address of a node where gcs server starts. + + Returns: + ProcessInfo for the process that was started. + """ + assert gcs_server_port > 0 + + command = [ + GCS_SERVER_EXECUTABLE, + f"--log_dir={log_dir}", + f"--config_list={serialize_config(config)}", + f"--gcs_server_port={gcs_server_port}", + f"--metrics-agent-port={metrics_agent_port}", + f"--node-ip-address={node_ip_address}", + f"--session-name={session_name}", + f"--ray-commit={ray.__commit__}", + ] + + if stdout_filepath: + command += [f"--stdout_filepath={stdout_filepath}"] + if stderr_filepath: + command += [f"--stderr_filepath={stderr_filepath}"] + + if redis_address: + redis_ip_address, redis_port, enable_redis_ssl = get_address(redis_address) + + command += [ + f"--redis_address={redis_ip_address}", + f"--redis_port={redis_port}", + f"--redis_enable_ssl={'true' if enable_redis_ssl else 'false'}", + ] + if redis_username: + command += [f"--redis_username={redis_username}"] + if redis_password: + command += [f"--redis_password={redis_password}"] + + stdout_file = None + if stdout_filepath: + stdout_file = open(os.devnull, "w") + + stderr_file = None + if stderr_filepath: + stderr_file = open(os.devnull, "w") + + process_info = start_ray_process( + command, + ray_constants.PROCESS_TYPE_GCS_SERVER, + stdout_file=stdout_file, + stderr_file=stderr_file, + fate_share=fate_share, + ) + return process_info + + +def start_raylet( + redis_address: str, + gcs_address: str, + node_id: str, + node_ip_address: str, + node_manager_port: int, + raylet_name: str, + plasma_store_name: str, + cluster_id: str, + worker_path: str, + setup_worker_path: str, + temp_dir: str, + session_dir: str, + resource_dir: str, + log_dir: str, + resource_and_label_spec, + plasma_directory: str, + fallback_directory: str, + object_store_memory: int, + session_name: str, + is_head_node: bool, + resource_isolation_config: ResourceIsolationConfig, + min_worker_port: Optional[int] = None, + max_worker_port: Optional[int] = None, + worker_port_list: Optional[List[int]] = None, + object_manager_port: Optional[int] = None, + redis_username: Optional[str] = None, + redis_password: Optional[str] = None, + metrics_agent_port: Optional[int] = None, + metrics_export_port: Optional[int] = None, + dashboard_agent_listen_port: Optional[int] = None, + runtime_env_agent_port: Optional[int] = None, + use_valgrind: bool = False, + use_profiler: bool = False, + raylet_stdout_filepath: Optional[str] = None, + raylet_stderr_filepath: Optional[str] = None, + dashboard_agent_stdout_filepath: Optional[str] = None, + dashboard_agent_stderr_filepath: Optional[str] = None, + runtime_env_agent_stdout_filepath: Optional[str] = None, + runtime_env_agent_stderr_filepath: Optional[str] = None, + huge_pages: bool = False, + fate_share: Optional[bool] = None, + socket_to_use: Optional[int] = None, + max_bytes: int = 0, + backup_count: int = 0, + ray_debugger_external: bool = False, + env_updates: Optional[dict] = None, + node_name: Optional[str] = None, + webui: Optional[str] = None, +): + """Start a raylet, which is a combined local scheduler and object manager. + + Args: + redis_address: The address of the primary Redis server. + gcs_address: The address of GCS server. + node_id: The hex ID of this node. + node_ip_address: The IP address of this node. + node_manager_port: The port to use for the node manager. If it's + 0, a random port will be used. + raylet_name: The name of the raylet socket to create. + plasma_store_name: The name of the plasma store socket to connect + to. + worker_path: The path of the Python file that new worker + processes will execute. + setup_worker_path: The path of the Python file that will set up + the environment for the worker process. + temp_dir: The path of the temporary directory Ray will use. + session_dir: The path of this session. + resource_dir: The path of resource of this session . + log_dir: The path of the dir where log files are created. + resource_and_label_spec: Resources and key-value labels for this raylet. + plasma_directory: A directory where the Plasma memory mapped files will + be created. + fallback_directory: A directory where the Object store fallback files will be created. + object_store_memory: The amount of memory (in bytes) to start the + object store with. + session_name: The current Ray session name. + resource_isolation_config: Resource isolation configuration for reserving + memory and cpu resources for ray system processes through cgroupv2 + is_head_node: whether this node is the head node. + min_worker_port: The lowest port number that workers will bind + on. If not set, random ports will be chosen. + max_worker_port: The highest port number that workers will bind + on. If set, min_worker_port must also be set. + worker_port_list: An explicit list of ports to be used for + workers (comma-separated). Overrides min_worker_port and + max_worker_port. + object_manager_port: The port to use for the object manager. If this is + None, then the object manager will choose its own port. + redis_username: The username to use when connecting to Redis. + redis_password: The password to use when connecting to Redis. + metrics_agent_port: The port where metrics agent is bound to. + metrics_export_port: The port at which metrics are exposed to. + dashboard_agent_listen_port: The port at which the dashboard agent + listens to for HTTP. + runtime_env_agent_port: The port at which the runtime env agent + listens to for HTTP. + use_valgrind: True if the raylet should be started inside + of valgrind. If this is True, use_profiler must be False. + use_profiler: True if the raylet should be started inside + a profiler. If this is True, use_valgrind must be False. + raylet_stdout_filepath: The file path to dump raylet stdout. + If None, stdout is not redirected. + raylet_stderr_filepath: The file path to dump raylet stderr. + If None, stderr is not redirected. + dashboard_agent_stdout_filepath: The file path to dump + dashboard agent stdout. If None, stdout is not redirected. + dashboard_agent_stderr_filepath: The file path to dump + dashboard agent stderr. If None, stderr is not redirected. + runtime_env_agent_stdout_filepath: The file path to dump + runtime env agent stdout. If None, stdout is not redirected. + runtime_env_agent_stderr_filepath: The file path to dump + runtime env agent stderr. If None, stderr is not redirected. + huge_pages: Boolean flag indicating whether to start the Object + Store with hugetlbfs support. Requires plasma_directory. + fate_share: Whether to share fate between raylet and this process. + max_bytes: Log rotation parameter. Corresponding to + RotatingFileHandler's maxBytes. + backup_count: Log rotation parameter. Corresponding to + RotatingFileHandler's backupCount. + ray_debugger_external: True if the Ray debugger should be made + available externally to this node. + env_updates: Environment variable overrides. + node_name: The name of the node. + webui: The url of the UI. + Returns: + ProcessInfo for the process that was started. + """ + assert node_manager_port is not None and type(node_manager_port) is int + + if use_valgrind and use_profiler: + raise ValueError("Cannot use valgrind and profiler at the same time.") + + # Get the static resources and labels from the resolved ResourceAndLabelSpec + static_resources = resource_and_label_spec.to_resource_dict() + labels = resource_and_label_spec.labels + + # Limit the number of workers that can be started in parallel by the + # raylet. However, make sure it is at least 1. + num_cpus_static = static_resources.get("CPU", 0) + maximum_startup_concurrency = max( + 1, min(multiprocessing.cpu_count(), num_cpus_static) + ) + + # Format the resource argument in a form like 'CPU,1.0,GPU,0,Custom,3'. + resource_argument = ",".join( + ["{},{}".format(*kv) for kv in static_resources.items()] + ) + + has_java_command = False + if shutil.which("java") is not None: + has_java_command = True + + ray_java_installed = False + try: + jars_dir = get_ray_jars_dir() + if os.path.exists(jars_dir): + ray_java_installed = True + except Exception: + pass + + include_java = has_java_command and ray_java_installed + if include_java is True: + java_worker_command = build_java_worker_command( + gcs_address, + plasma_store_name, + raylet_name, + redis_username, + redis_password, + session_dir, + node_ip_address, + setup_worker_path, + ) + else: + java_worker_command = [] + + if os.path.exists(DEFAULT_WORKER_EXECUTABLE): + cpp_worker_command = build_cpp_worker_command( + gcs_address, + plasma_store_name, + raylet_name, + redis_username, + redis_password, + session_dir, + log_dir, + node_ip_address, + setup_worker_path, + ) + else: + cpp_worker_command = [] + + # Create the command that the Raylet will use to start workers. + # TODO(architkulkarni): Pipe in setup worker args separately instead of + # inserting them into start_worker_command and later erasing them if + # needed. + start_worker_command = ( + [ + sys.executable, + setup_worker_path, + ] + + _site_flags() # Inherit "-S" and "-s" flags from current Python interpreter. + + [ + worker_path, + f"--node-ip-address={node_ip_address}", + "--node-manager-port=RAY_NODE_MANAGER_PORT_PLACEHOLDER", + f"--object-store-name={plasma_store_name}", + f"--raylet-name={raylet_name}", + f"--redis-address={redis_address}", + f"--metrics-agent-port={metrics_agent_port}", + f"--logging-rotate-bytes={max_bytes}", + f"--logging-rotate-backup-count={backup_count}", + f"--runtime-env-agent-port={runtime_env_agent_port}", + f"--gcs-address={gcs_address}", + f"--session-name={session_name}", + f"--temp-dir={temp_dir}", + f"--webui={webui}", + f"--cluster-id={cluster_id}", + ] + ) + + start_worker_command.append("RAY_WORKER_DYNAMIC_OPTION_PLACEHOLDER") + + if redis_username: + start_worker_command += [f"--redis-username={redis_username}"] + + if redis_password: + start_worker_command += [f"--redis-password={redis_password}"] + + # If the object manager port is None, then use 0 to cause the object + # manager to choose its own port. + if object_manager_port is None: + object_manager_port = 0 + + if min_worker_port is None: + min_worker_port = 0 + + if max_worker_port is None: + max_worker_port = 0 + + labels_json_str = "" + if labels: + labels_json_str = json.dumps(labels) + + dashboard_agent_command = [ + *_build_python_executable_command_memory_profileable( + ray_constants.PROCESS_TYPE_DASHBOARD_AGENT, session_dir + ), + os.path.join(RAY_PATH, "dashboard", "agent.py"), + f"--node-ip-address={node_ip_address}", + f"--metrics-export-port={metrics_export_port}", + f"--grpc-port={metrics_agent_port}", + f"--listen-port={dashboard_agent_listen_port}", + "--node-manager-port=RAY_NODE_MANAGER_PORT_PLACEHOLDER", + f"--object-store-name={plasma_store_name}", + f"--raylet-name={raylet_name}", + f"--temp-dir={temp_dir}", + f"--session-dir={session_dir}", + f"--log-dir={log_dir}", + f"--logging-rotate-bytes={max_bytes}", + f"--logging-rotate-backup-count={backup_count}", + f"--session-name={session_name}", + f"--gcs-address={gcs_address}", + f"--cluster-id-hex={cluster_id}", + ] + if dashboard_agent_stdout_filepath: + dashboard_agent_command.append( + f"--stdout-filepath={dashboard_agent_stdout_filepath}" + ) + if dashboard_agent_stderr_filepath: + dashboard_agent_command.append( + f"--stderr-filepath={dashboard_agent_stderr_filepath}" + ) + if ( + dashboard_agent_stdout_filepath is None + and dashboard_agent_stderr_filepath is None + ): + # If not redirecting logging to files, unset log filename. + # This will cause log records to go to stderr. + dashboard_agent_command.append("--logging-filename=") + # Use stderr log format with the component name as a message prefix. + logging_format = ray_constants.LOGGER_FORMAT_STDERR.format( + component=ray_constants.PROCESS_TYPE_DASHBOARD_AGENT + ) + dashboard_agent_command.append(f"--logging-format={logging_format}") + + if ray._private.utils.get_dashboard_dependency_error() is not None: + # If dependencies are not installed, it is the minimally packaged + # ray. We should restrict the features within dashboard agent + # that requires additional dependencies to be downloaded. + dashboard_agent_command.append("--minimal") + + runtime_env_agent_command = [ + *_build_python_executable_command_memory_profileable( + ray_constants.PROCESS_TYPE_RUNTIME_ENV_AGENT, session_dir + ), + os.path.join(RAY_PATH, "_private", "runtime_env", "agent", "main.py"), + f"--node-ip-address={node_ip_address}", + f"--runtime-env-agent-port={runtime_env_agent_port}", + f"--gcs-address={gcs_address}", + f"--cluster-id-hex={cluster_id}", + f"--runtime-env-dir={resource_dir}", + f"--logging-rotate-bytes={max_bytes}", + f"--logging-rotate-backup-count={backup_count}", + f"--log-dir={log_dir}", + f"--temp-dir={temp_dir}", + ] + if runtime_env_agent_stdout_filepath: + runtime_env_agent_command.append( + f"--stdout-filepath={runtime_env_agent_stdout_filepath}" + ) + if runtime_env_agent_stderr_filepath: + runtime_env_agent_command.append( + f"--stderr-filepath={runtime_env_agent_stderr_filepath}" + ) + if ( + runtime_env_agent_stdout_filepath is None + and runtime_env_agent_stderr_filepath is None + ): + # If not redirecting logging to files, unset log filename. + # This will cause log records to go to stderr. + runtime_env_agent_command.append("--logging-filename=") + # Use stderr log format with the component name as a message prefix. + logging_format = ray_constants.LOGGER_FORMAT_STDERR.format( + component=ray_constants.PROCESS_TYPE_RUNTIME_ENV_AGENT + ) + runtime_env_agent_command.append(f"--logging-format={logging_format}") + + command = [ + RAYLET_EXECUTABLE, + f"--raylet_socket_name={raylet_name}", + f"--store_socket_name={plasma_store_name}", + f"--object_manager_port={object_manager_port}", + f"--min_worker_port={min_worker_port}", + f"--max_worker_port={max_worker_port}", + f"--node_manager_port={node_manager_port}", + f"--node_id={node_id}", + f"--node_ip_address={node_ip_address}", + f"--maximum_startup_concurrency={maximum_startup_concurrency}", + f"--static_resource_list={resource_argument}", + f"--python_worker_command={subprocess.list2cmdline(start_worker_command)}", # noqa + f"--java_worker_command={subprocess.list2cmdline(java_worker_command)}", # noqa + f"--cpp_worker_command={subprocess.list2cmdline(cpp_worker_command)}", # noqa + f"--native_library_path={DEFAULT_NATIVE_LIBRARY_PATH}", + f"--temp_dir={temp_dir}", + f"--session_dir={session_dir}", + f"--log_dir={log_dir}", + f"--resource_dir={resource_dir}", + f"--metrics-agent-port={metrics_agent_port}", + f"--metrics_export_port={metrics_export_port}", + f"--runtime_env_agent_port={runtime_env_agent_port}", + f"--object_store_memory={object_store_memory}", + f"--plasma_directory={plasma_directory}", + f"--fallback_directory={fallback_directory}", + f"--ray-debugger-external={1 if ray_debugger_external else 0}", + f"--gcs-address={gcs_address}", + f"--session-name={session_name}", + f"--labels={labels_json_str}", + f"--cluster-id={cluster_id}", + ] + + if resource_isolation_config.is_enabled(): + logging.info( + f"Resource isolation enabled with cgroup_path={resource_isolation_config.cgroup_path}, " + f"system_reserved_cpu={resource_isolation_config.system_reserved_cpu_weight} " + f"system_reserved_memory={resource_isolation_config.system_reserved_memory}." + ) + command.append("--enable-resource-isolation") + command.append(f"--cgroup-path={resource_isolation_config.cgroup_path}") + command.append( + f"--system-reserved-cpu-weight={resource_isolation_config.system_reserved_cpu_weight}" + ) + command.append( + f"--system-reserved-memory-bytes={resource_isolation_config.system_reserved_memory}" + ) + command.append(f"--system-pids={resource_isolation_config.system_pids}") + + if raylet_stdout_filepath: + command.append(f"--stdout_filepath={raylet_stdout_filepath}") + if raylet_stderr_filepath: + command.append(f"--stderr_filepath={raylet_stderr_filepath}") + + if is_head_node: + command.append("--head") + + if worker_port_list is not None: + command.append(f"--worker_port_list={worker_port_list}") + command.append( + "--num_prestart_python_workers={}".format(int(resource_and_label_spec.num_cpus)) + ) + command.append( + "--dashboard_agent_command={}".format( + subprocess.list2cmdline(dashboard_agent_command) + ) + ) + command.append( + "--runtime_env_agent_command={}".format( + subprocess.list2cmdline(runtime_env_agent_command) + ) + ) + if huge_pages: + command.append("--huge_pages") + if socket_to_use: + socket_to_use.close() + if node_name is not None: + command.append( + f"--node-name={node_name}", + ) + + stdout_file = None + if raylet_stdout_filepath: + stdout_file = open(os.devnull, "w") + + stderr_file = None + if raylet_stderr_filepath: + stderr_file = open(os.devnull, "w") + + process_info = start_ray_process( + command, + ray_constants.PROCESS_TYPE_RAYLET, + use_valgrind=use_valgrind, + use_gdb=False, + use_valgrind_profiler=use_profiler, + use_perftools_profiler=("RAYLET_PERFTOOLS_PATH" in os.environ), + stdout_file=stdout_file, + stderr_file=stderr_file, + fate_share=fate_share, + env_updates=env_updates, + ) + return process_info + + +def get_ray_jars_dir(): + """Return a directory where all ray-related jars and + their dependencies locate.""" + current_dir = RAY_PATH + jars_dir = os.path.abspath(os.path.join(current_dir, "jars")) + if not os.path.exists(jars_dir): + raise RuntimeError( + "Ray jars is not packaged into ray. " + "Please build ray with java enabled " + "(set env var RAY_INSTALL_JAVA=1)" + ) + return os.path.abspath(os.path.join(current_dir, "jars")) + + +def build_java_worker_command( + bootstrap_address: str, + plasma_store_name: str, + raylet_name: str, + redis_username: str, + redis_password: str, + session_dir: str, + node_ip_address: str, + setup_worker_path: str, +): + """This method assembles the command used to start a Java worker. + + Args: + bootstrap_address: Bootstrap address of ray cluster. + plasma_store_name: The name of the plasma store socket to connect + to. + raylet_name: The name of the raylet socket to create. + redis_username: The username to connect to Redis. + redis_password: The password to connect to Redis. + session_dir: The path of this session. + node_ip_address: The IP address for this node. + setup_worker_path: The path of the Python file that will set up + the environment for the worker process. + Returns: + The command string for starting Java worker. + """ + pairs = [] + if bootstrap_address is not None: + pairs.append(("ray.address", bootstrap_address)) + pairs.append(("ray.raylet.node-manager-port", "RAY_NODE_MANAGER_PORT_PLACEHOLDER")) + + if plasma_store_name is not None: + pairs.append(("ray.object-store.socket-name", plasma_store_name)) + + if raylet_name is not None: + pairs.append(("ray.raylet.socket-name", raylet_name)) + + if redis_username is not None: + pairs.append(("ray.redis.username", redis_username)) + + if redis_password is not None: + pairs.append(("ray.redis.password", redis_password)) + + if node_ip_address is not None: + pairs.append(("ray.node-ip", node_ip_address)) + + pairs.append(("ray.home", RAY_HOME)) + pairs.append(("ray.logging.dir", os.path.join(session_dir, "logs"))) + pairs.append(("ray.session-dir", session_dir)) + command = ( + [sys.executable] + + [setup_worker_path] + + ["-D{}={}".format(*pair) for pair in pairs] + ) + + command += ["RAY_WORKER_DYNAMIC_OPTION_PLACEHOLDER"] + command += ["io.ray.runtime.runner.worker.DefaultWorker"] + + return command + + +def build_cpp_worker_command( + bootstrap_address: str, + plasma_store_name: str, + raylet_name: str, + redis_username: str, + redis_password: str, + session_dir: str, + log_dir: str, + node_ip_address: str, + setup_worker_path: str, +): + """This method assembles the command used to start a CPP worker. + + Args: + bootstrap_address: The bootstrap address of the cluster. + plasma_store_name: The name of the plasma store socket to connect + to. + raylet_name: The name of the raylet socket to create. + redis_username: The username to connect to Redis. + redis_password: The password to connect to Redis. + session_dir: The path of this session. + log_dir: The path of logs. + node_ip_address: The ip address for this node. + setup_worker_path: The path of the Python file that will set up + the environment for the worker process. + Returns: + The command string for starting CPP worker. + """ + + command = [ + sys.executable, + setup_worker_path, + DEFAULT_WORKER_EXECUTABLE, + f"--ray_plasma_store_socket_name={plasma_store_name}", + f"--ray_raylet_socket_name={raylet_name}", + "--ray_node_manager_port=RAY_NODE_MANAGER_PORT_PLACEHOLDER", + f"--ray_address={bootstrap_address}", + f"--ray_redis_username={redis_username}", + f"--ray_redis_password={redis_password}", + f"--ray_session_dir={session_dir}", + f"--ray_logs_dir={log_dir}", + f"--ray_node_ip_address={node_ip_address}", + "RAY_WORKER_DYNAMIC_OPTION_PLACEHOLDER", + ] + + return command + + +def determine_plasma_store_config( + object_store_memory: int, + temp_dir: str, + plasma_directory: Optional[str] = None, + fallback_directory: Optional[str] = None, + huge_pages: bool = False, +): + """Figure out how to configure the plasma object store. + + This will determine: + 1. which directory to use for the plasma store. On Linux, + we will try to use /dev/shm unless the shared memory file system is too + small, in which case we will fall back to /tmp. If any of the object store + memory or plasma directory parameters are specified by the user, then those + values will be preserved. + 2. which directory to use for the fallback files. It will default to the temp_dir + if it is not extracted from the object_spilling_config. + + Args: + object_store_memory: The object store memory to use. + plasma_directory: The user-specified plasma directory parameter. + fallback_directory: The path extracted from the object_spilling_config when the + object spilling config is set and the spilling type is to + filesystem. + huge_pages: The user-specified huge pages parameter. + + Returns: + A tuple of plasma directory to use, the fallback directory to use, and the + object store memory to use. If it is specified by the user, then that value will + be preserved. + """ + if not isinstance(object_store_memory, int): + object_store_memory = int(object_store_memory) + + if huge_pages and not (sys.platform == "linux" or sys.platform == "linux2"): + raise ValueError("The huge_pages argument is only supported on Linux.") + + system_memory = ray._common.utils.get_system_memory() + + # Determine which directory to use. By default, use /tmp on MacOS and + # /dev/shm on Linux, unless the shared-memory file system is too small, + # in which case we default to /tmp on Linux. + if plasma_directory is None: + if sys.platform == "linux" or sys.platform == "linux2": + shm_avail = ray._private.utils.get_shared_memory_bytes() + # Compare the requested memory size to the memory available in + # /dev/shm. + if shm_avail >= object_store_memory: + plasma_directory = "/dev/shm" + elif ( + not os.environ.get("RAY_OBJECT_STORE_ALLOW_SLOW_STORAGE") + and object_store_memory > ray_constants.REQUIRE_SHM_SIZE_THRESHOLD + ): + raise ValueError( + "The configured object store size ({} GB) exceeds " + "/dev/shm size ({} GB). This will harm performance. " + "Consider deleting files in /dev/shm or increasing its " + "size with " + "--shm-size in Docker. To ignore this warning, " + "set RAY_OBJECT_STORE_ALLOW_SLOW_STORAGE=1.".format( + object_store_memory / 1e9, shm_avail / 1e9 + ) + ) + else: + plasma_directory = ray._common.utils.get_user_temp_dir() + logger.warning( + "WARNING: The object store is using {} instead of " + "/dev/shm because /dev/shm has only {} bytes available. " + "This will harm performance! You may be able to free up " + "space by deleting files in /dev/shm. If you are inside a " + "Docker container, you can increase /dev/shm size by " + "passing '--shm-size={:.2f}gb' to 'docker run' (or add it " + "to the run_options list in a Ray cluster config). Make " + "sure to set this to more than 30% of available RAM.".format( + ray._common.utils.get_user_temp_dir(), + shm_avail, + object_store_memory * (1.1) / (2**30), + ) + ) + else: + plasma_directory = ray._common.utils.get_user_temp_dir() + + # Do some sanity checks. + if object_store_memory > system_memory: + raise ValueError( + "The requested object store memory size is greater " + "than the total available memory." + ) + else: + plasma_directory = os.path.abspath(plasma_directory) + logger.info("object_store_memory is not verified when plasma_directory is set.") + + if not os.path.isdir(plasma_directory): + raise ValueError( + f"The plasma directory file {plasma_directory} does not exist or is not a directory." + ) + + if huge_pages and plasma_directory is None: + raise ValueError( + "If huge_pages is True, then the " + "plasma_directory argument must be provided." + ) + + if object_store_memory < ray_constants.OBJECT_STORE_MINIMUM_MEMORY_BYTES: + raise ValueError( + "Attempting to cap object store memory usage at {} " + "bytes, but the minimum allowed is {} bytes.".format( + object_store_memory, ray_constants.OBJECT_STORE_MINIMUM_MEMORY_BYTES + ) + ) + + if ( + sys.platform == "darwin" + and object_store_memory > ray_constants.MAC_DEGRADED_PERF_MMAP_SIZE_LIMIT + and os.environ.get("RAY_ENABLE_MAC_LARGE_OBJECT_STORE") != "1" + ): + raise ValueError( + "The configured object store size ({:.4}GiB) exceeds " + "the optimal size on Mac ({:.4}GiB). " + "This will harm performance! There is a known issue where " + "Ray's performance degrades with object store size greater" + " than {:.4}GB on a Mac." + "To reduce the object store capacity, specify" + "`object_store_memory` when calling ray.init() or ray start." + "To ignore this warning, " + "set RAY_ENABLE_MAC_LARGE_OBJECT_STORE=1.".format( + object_store_memory / 2**30, + ray_constants.MAC_DEGRADED_PERF_MMAP_SIZE_LIMIT / 2**30, + ray_constants.MAC_DEGRADED_PERF_MMAP_SIZE_LIMIT / 2**30, + ) + ) + + if fallback_directory is None: + fallback_directory = temp_dir + else: + fallback_directory = os.path.abspath(fallback_directory) + + if not os.path.isdir(fallback_directory): + raise ValueError( + f"The fallback directory file {fallback_directory} does not exist or is not a directory." + ) + + # Print the object store memory using two decimal places. + logger.debug( + "Determine to start the Plasma object store with {} GB memory " + "using {} and fallback to {}".format( + round(object_store_memory / 10**9, 2), + plasma_directory, + fallback_directory, + ) + ) + + return plasma_directory, fallback_directory, object_store_memory + + +def start_monitor( + gcs_address: str, + logs_dir: str, + stdout_filepath: Optional[str] = None, + stderr_filepath: Optional[str] = None, + autoscaling_config: Optional[str] = None, + fate_share: Optional[bool] = None, + max_bytes: int = 0, + backup_count: int = 0, + monitor_ip: Optional[str] = None, + autoscaler_v2: bool = False, +): + """Run a process to monitor the other processes. + + Args: + gcs_address: The address of GCS server. + logs_dir: The path to the log directory. + stdout_filepath: The file path to dump monitor stdout. + If None, stdout is not redirected. + stderr_filepath: The file path to dump monitor stderr. + If None, stderr is not redirected. + autoscaling_config: path to autoscaling config file. + max_bytes: Log rotation parameter. Corresponding to + RotatingFileHandler's maxBytes. + backup_count: Log rotation parameter. Corresponding to + RotatingFileHandler's backupCount. + monitor_ip: IP address of the machine that the monitor will be + run on. Can be excluded, but required for autoscaler metrics. + Returns: + ProcessInfo for the process that was started. + """ + if autoscaler_v2: + entrypoint = os.path.join(RAY_PATH, AUTOSCALER_V2_DIR, "monitor.py") + else: + entrypoint = os.path.join(RAY_PATH, AUTOSCALER_PRIVATE_DIR, "monitor.py") + + command = [ + sys.executable, + "-u", + entrypoint, + f"--logs-dir={logs_dir}", + f"--logging-rotate-bytes={max_bytes}", + f"--logging-rotate-backup-count={backup_count}", + ] + assert gcs_address is not None + command.append(f"--gcs-address={gcs_address}") + + if stdout_filepath: + command.append(f"--stdout-filepath={stdout_filepath}") + if stderr_filepath: + command.append(f"--stderr-filepath={stderr_filepath}") + + if stdout_filepath is None and stderr_filepath is None: + # If not redirecting logging to files, unset log filename. + # This will cause log records to go to stderr. + command.append("--logging-filename=") + # Use stderr log format with the component name as a message prefix. + logging_format = ray_constants.LOGGER_FORMAT_STDERR.format( + component=ray_constants.PROCESS_TYPE_MONITOR + ) + command.append(f"--logging-format={logging_format}") + if autoscaling_config: + command.append("--autoscaling-config=" + str(autoscaling_config)) + if monitor_ip: + command.append("--monitor-ip=" + monitor_ip) + + stdout_file = None + if stdout_filepath: + stdout_file = open(os.devnull, "w") + + stderr_file = None + if stderr_filepath: + stderr_file = open(os.devnull, "w") + + process_info = start_ray_process( + command, + ray_constants.PROCESS_TYPE_MONITOR, + stdout_file=stdout_file, + stderr_file=stderr_file, + fate_share=fate_share, + ) + return process_info + + +def start_ray_client_server( + address: str, + ray_client_server_ip: str, + ray_client_server_port: int, + stdout_file: Optional[int] = None, + stderr_file: Optional[int] = None, + redis_username: Optional[int] = None, + redis_password: Optional[int] = None, + fate_share: Optional[bool] = None, + runtime_env_agent_address: Optional[str] = None, + server_type: str = "proxy", + serialized_runtime_env_context: Optional[str] = None, +): + """Run the server process of the Ray client. + + Args: + address: The address of the cluster. + ray_client_server_ip: Host IP the Ray client server listens on. + ray_client_server_port: Port the Ray client server listens on. + stdout_file: A file handle opened for writing to redirect stdout to. If + no redirection should happen, then this should be None. + stderr_file: A file handle opened for writing to redirect stderr to. If + no redirection should happen, then this should be None. + redis_username: The username of the Redis server. + redis_password: The password of the Redis server. + runtime_env_agent_address: Address to the Runtime Env Agent listens on via HTTP. + Only needed when server_type == "proxy". + server_type: Whether to start the proxy version of Ray Client. + serialized_runtime_env_context (str|None): If specified, the serialized + runtime_env_context to start the client server in. + + Returns: + ProcessInfo for the process that was started. + """ + root_ray_dir = Path(__file__).resolve().parents[1] + setup_worker_path = os.path.join( + root_ray_dir, "_private", "workers", ray_constants.SETUP_WORKER_FILENAME + ) + + ray_client_server_host = ray_client_server_ip + command = [ + sys.executable, + setup_worker_path, + "-m", + "ray.util.client.server", + f"--address={address}", + f"--host={ray_client_server_host}", + f"--port={ray_client_server_port}", + f"--mode={server_type}", + f"--language={Language.Name(Language.PYTHON)}", + ] + if redis_username: + command.append(f"--redis-username={redis_username}") + if redis_password: + command.append(f"--redis-password={redis_password}") + if serialized_runtime_env_context: + command.append( + f"--serialized-runtime-env-context={serialized_runtime_env_context}" # noqa: E501 + ) + if server_type == "proxy": + assert len(runtime_env_agent_address) > 0 + if runtime_env_agent_address: + command.append(f"--runtime-env-agent-address={runtime_env_agent_address}") + + process_info = start_ray_process( + command, + ray_constants.PROCESS_TYPE_RAY_CLIENT_SERVER, + stdout_file=stdout_file, + stderr_file=stderr_file, + fate_share=fate_share, + ) + return process_info + + +def _is_raylet_process(cmdline: Optional[List[str]]) -> bool: + """Check if the command line belongs to a raylet process. + + Args: + cmdline: List of command line arguments or None + + Returns: + bool: True if this is a raylet process, False otherwise + """ + if cmdline is None or len(cmdline) == 0: + return False + + executable = os.path.basename(cmdline[0]) + return "raylet" in executable diff --git a/lib/python3.12/site-packages/ray/_private/state.py b/lib/python3.12/site-packages/ray/_private/state.py new file mode 100644 index 0000000000000000000000000000000000000000..65e1da21c3e3f89e8f689d9e8c2c9398ce26b463 --- /dev/null +++ b/lib/python3.12/site-packages/ray/_private/state.py @@ -0,0 +1,1150 @@ +import json +import logging +import sys +from collections import defaultdict +from threading import Lock +from typing import Dict, Optional + +import ray +from ray._common.constants import HEAD_NODE_RESOURCE_NAME, NODE_ID_PREFIX +from ray._common.utils import binary_to_hex, decode, hex_to_binary +from ray._private.client_mode_hook import client_mode_hook +from ray._private.protobuf_compat import message_to_dict +from ray._private.utils import ( + validate_actor_state_name, +) +from ray._raylet import GlobalStateAccessor +from ray.core.generated import autoscaler_pb2, common_pb2, gcs_pb2 +from ray.util.annotations import DeveloperAPI + +logger = logging.getLogger(__name__) + + +class GlobalState: + """A class used to interface with the Ray control state. + + Attributes: + global_state_accessor: The client used to query gcs table from gcs + server. + """ + + def __init__(self): + """Create a GlobalState object.""" + # Args used for lazy init of this object. + self.gcs_options = None + self._global_state_accessor = None + self._init_lock = Lock() + + def _connect_and_get_accessor(self) -> GlobalStateAccessor: + """ + This lazily initializes clients needed for state accessors and returns a connected global state accessor. + + Returns: + GlobalStateAccessor: A connected global state accessor. + Raises: + RuntimeError: An exception is raised if ray.init() has not been called yet. + """ + with self._init_lock: + if self._global_state_accessor is not None: + return self._global_state_accessor + + if self.gcs_options is None: + raise ray.exceptions.RaySystemError( + "Ray has not been started yet. Trying to use state API before ray.init() has been called." + ) + + self._global_state_accessor = GlobalStateAccessor(self.gcs_options) + connected = self._global_state_accessor.connect() + if not connected: + self._global_state_accessor = None + raise ray.exceptions.RaySystemError( + "Failed to connect to GCS. Please check if the GCS server is running " + "and if this node can connect to the head node." + ) + return self._global_state_accessor + + def disconnect(self): + """Disconnect global state from GCS.""" + with self._init_lock: + self.gcs_options = None + if self._global_state_accessor is not None: + self._global_state_accessor = None + + def _initialize_global_state(self, gcs_options): + """Set args for lazily initialization of the GlobalState object. + + It's possible that certain keys in gcs kv may not have been fully + populated yet. In this case, we will retry this method until they have + been populated or we exceed a timeout. + + Args: + gcs_options: The client options for gcs + """ + + # Save args for lazy init of global state. This avoids opening extra + # gcs connections from each worker until needed. + with self._init_lock: + self.gcs_options = gcs_options + + def actor_table( + self, + actor_id: Optional[str], + job_id: Optional[ray.JobID] = None, + actor_state_name: Optional[str] = None, + ): + """Fetch and parse the actor table information for a single actor ID. + + Args: + actor_id: A hex string of the actor ID to fetch information about. + If this is None, then the actor table is fetched. + If this is not None, `job_id` and `actor_state_name` + will not take effect. + job_id: To filter actors by job_id, which is of type `ray.JobID`. + You can use the `ray.get_runtime_context().job_id` function + to get the current job ID + actor_state_name: To filter actors based on actor state, + which can be one of the following: "DEPENDENCIES_UNREADY", + "PENDING_CREATION", "ALIVE", "RESTARTING", or "DEAD". + Returns: + Information from the actor table. + """ + accessor = self._connect_and_get_accessor() + + if actor_id is not None: + actor_id = ray.ActorID(hex_to_binary(actor_id)) + actor_info = accessor.get_actor_info(actor_id) + if actor_info is None: + return {} + else: + actor_table_data = gcs_pb2.ActorTableData.FromString(actor_info) + return self._gen_actor_info(actor_table_data) + else: + validate_actor_state_name(actor_state_name) + actor_table = accessor.get_actor_table(job_id, actor_state_name) + results = {} + for i in range(len(actor_table)): + actor_table_data = gcs_pb2.ActorTableData.FromString(actor_table[i]) + results[ + binary_to_hex(actor_table_data.actor_id) + ] = self._gen_actor_info(actor_table_data) + + return results + + def _gen_actor_info(self, actor_table_data): + """Parse actor table data. + + Returns: + Information from actor table. + """ + actor_info = { + "ActorID": binary_to_hex(actor_table_data.actor_id), + "ActorClassName": actor_table_data.class_name, + "IsDetached": actor_table_data.is_detached, + "Name": actor_table_data.name, + "JobID": binary_to_hex(actor_table_data.job_id), + "Address": { + "IPAddress": actor_table_data.address.ip_address, + "Port": actor_table_data.address.port, + "NodeID": binary_to_hex(actor_table_data.address.node_id), + }, + "OwnerAddress": { + "IPAddress": actor_table_data.owner_address.ip_address, + "Port": actor_table_data.owner_address.port, + "NodeID": binary_to_hex(actor_table_data.owner_address.node_id), + }, + "State": gcs_pb2.ActorTableData.ActorState.DESCRIPTOR.values_by_number[ + actor_table_data.state + ].name, + "NumRestarts": actor_table_data.num_restarts, + "Timestamp": actor_table_data.timestamp, + "StartTime": actor_table_data.start_time, + "EndTime": actor_table_data.end_time, + "DeathCause": actor_table_data.death_cause, + "Pid": actor_table_data.pid, + } + return actor_info + + def node_table(self): + """Fetch and parse the Gcs node info table. + + Returns: + Information about the node in the cluster. + """ + accessor = self._connect_and_get_accessor() + + return accessor.get_node_table() + + def job_table(self): + """Fetch and parse the gcs job table. + + Returns: + Information about the Ray jobs in the cluster, + namely a list of dicts with keys: + - "JobID" (identifier for the job), + - "DriverIPAddress" (IP address of the driver for this job), + - "DriverPid" (process ID of the driver for this job), + - "StartTime" (UNIX timestamp of the start time of this job), + - "StopTime" (UNIX timestamp of the stop time of this job, if any) + """ + accessor = self._connect_and_get_accessor() + + job_table = accessor.get_job_table( + skip_submission_job_info_field=True, skip_is_running_tasks_field=True + ) + + results = [] + for i in range(len(job_table)): + entry = gcs_pb2.JobTableData.FromString(job_table[i]) + job_info = {} + job_info["JobID"] = entry.job_id.hex() + job_info["DriverIPAddress"] = entry.driver_address.ip_address + job_info["DriverPid"] = entry.driver_pid + job_info["Timestamp"] = entry.timestamp + job_info["StartTime"] = entry.start_time + job_info["EndTime"] = entry.end_time + job_info["IsDead"] = entry.is_dead + job_info["Entrypoint"] = entry.entrypoint + results.append(job_info) + + return results + + def next_job_id(self): + """Get next job id from GCS. + + Returns: + Next job id in the cluster. + """ + accessor = self._connect_and_get_accessor() + + return ray.JobID.from_int(accessor.get_next_job_id()) + + def profile_events(self): + """Retrieve and return task profiling events from GCS. + + Return: + Profiling events by component id (e.g. worker id). + { + : [ + { + event_type: , + component_id: , + node_ip_address: , + component_type: , + start_time: , + end_time: , + extra_data: , + } + ] + } + """ + accessor = self._connect_and_get_accessor() + + result = defaultdict(list) + task_events = accessor.get_task_events() + for i in range(len(task_events)): + event = gcs_pb2.TaskEvents.FromString(task_events[i]) + profile = event.profile_events + if not profile: + continue + + component_type = profile.component_type + component_id = binary_to_hex(profile.component_id) + node_ip_address = profile.node_ip_address + + for event in profile.events: + try: + extra_data = json.loads(event.extra_data) + except ValueError: + extra_data = {} + profile_event = { + "event_type": event.event_name, + "component_id": component_id, + "node_ip_address": node_ip_address, + "component_type": component_type, + "start_time": event.start_time, + "end_time": event.end_time, + "extra_data": extra_data, + } + + result[component_id].append(profile_event) + + return dict(result) + + def get_placement_group_by_name(self, placement_group_name, ray_namespace): + accessor = self._connect_and_get_accessor() + + placement_group_info = accessor.get_placement_group_by_name( + placement_group_name, ray_namespace + ) + if placement_group_info is None: + return None + else: + placement_group_table_data = gcs_pb2.PlacementGroupTableData.FromString( + placement_group_info + ) + return self._gen_placement_group_info(placement_group_table_data) + + def placement_group_table(self, placement_group_id=None): + accessor = self._connect_and_get_accessor() + + if placement_group_id is not None: + placement_group_id = ray.PlacementGroupID( + hex_to_binary(placement_group_id.hex()) + ) + placement_group_info = accessor.get_placement_group_info(placement_group_id) + if placement_group_info is None: + return {} + else: + placement_group_info = gcs_pb2.PlacementGroupTableData.FromString( + placement_group_info + ) + return self._gen_placement_group_info(placement_group_info) + else: + placement_group_table = accessor.get_placement_group_table() + results = {} + for placement_group_info in placement_group_table: + placement_group_table_data = gcs_pb2.PlacementGroupTableData.FromString( + placement_group_info + ) + placement_group_id = binary_to_hex( + placement_group_table_data.placement_group_id + ) + results[placement_group_id] = self._gen_placement_group_info( + placement_group_table_data + ) + + return results + + def _gen_placement_group_info(self, placement_group_info): + # This should be imported here, otherwise, it will error doc build. + from ray.core.generated.common_pb2 import PlacementStrategy + + def get_state(state): + if state == gcs_pb2.PlacementGroupTableData.PENDING: + return "PENDING" + elif state == gcs_pb2.PlacementGroupTableData.PREPARED: + return "PREPARED" + elif state == gcs_pb2.PlacementGroupTableData.CREATED: + return "CREATED" + elif state == gcs_pb2.PlacementGroupTableData.RESCHEDULING: + return "RESCHEDULING" + else: + return "REMOVED" + + def get_strategy(strategy): + if strategy == PlacementStrategy.PACK: + return "PACK" + elif strategy == PlacementStrategy.STRICT_PACK: + return "STRICT_PACK" + elif strategy == PlacementStrategy.STRICT_SPREAD: + return "STRICT_SPREAD" + elif strategy == PlacementStrategy.SPREAD: + return "SPREAD" + else: + raise ValueError(f"Invalid strategy returned: {PlacementStrategy}") + + stats = placement_group_info.stats + assert placement_group_info is not None + return { + "placement_group_id": binary_to_hex( + placement_group_info.placement_group_id + ), + "name": placement_group_info.name, + "bundles": { + # The value here is needs to be dictionarified + # otherwise, the payload becomes unserializable. + bundle.bundle_id.bundle_index: message_to_dict(bundle)["unitResources"] + for bundle in placement_group_info.bundles + }, + "bundles_to_node_id": { + bundle.bundle_id.bundle_index: binary_to_hex(bundle.node_id) + for bundle in placement_group_info.bundles + }, + "strategy": get_strategy(placement_group_info.strategy), + "state": get_state(placement_group_info.state), + "stats": { + "end_to_end_creation_latency_ms": ( + stats.end_to_end_creation_latency_us / 1000.0 + ), + "scheduling_latency_ms": (stats.scheduling_latency_us / 1000.0), + "scheduling_attempt": stats.scheduling_attempt, + "highest_retry_delay_ms": stats.highest_retry_delay_ms, + "scheduling_state": gcs_pb2.PlacementGroupStats.SchedulingState.DESCRIPTOR.values_by_number[ # noqa: E501 + stats.scheduling_state + ].name, + }, + } + + def _nanoseconds_to_microseconds(self, time_in_nanoseconds): + """A helper function for converting nanoseconds to microseconds.""" + time_in_microseconds = time_in_nanoseconds / 1000 + return time_in_microseconds + + # Colors are specified at + # https://github.com/catapult-project/catapult/blob/master/tracing/tracing/base/color_scheme.html. # noqa: E501 + _default_color_mapping = defaultdict( + lambda: "generic_work", + { + "worker_idle": "cq_build_abandoned", + "task": "rail_response", + "task:deserialize_arguments": "rail_load", + "task:execute": "rail_animation", + "task:store_outputs": "rail_idle", + "wait_for_function": "detailed_memory_dump", + "ray.get": "good", + "ray.put": "terrible", + "ray.wait": "vsync_highlight_color", + "submit_task": "background_memory_dump", + "fetch_and_run_function": "detailed_memory_dump", + "register_remote_function": "detailed_memory_dump", + }, + ) + + # These colors are for use in Chrome tracing. + _chrome_tracing_colors = [ + "thread_state_uninterruptible", + "thread_state_iowait", + "thread_state_running", + "thread_state_runnable", + "thread_state_sleeping", + "thread_state_unknown", + "background_memory_dump", + "light_memory_dump", + "detailed_memory_dump", + "vsync_highlight_color", + "generic_work", + "good", + "bad", + "terrible", + # "black", + # "grey", + # "white", + "yellow", + "olive", + "rail_response", + "rail_animation", + "rail_idle", + "rail_load", + "startup", + "heap_dump_stack_frame", + "heap_dump_object_type", + "heap_dump_child_node_arrow", + "cq_build_running", + "cq_build_passed", + "cq_build_failed", + "cq_build_abandoned", + "cq_build_attempt_runnig", + "cq_build_attempt_passed", + "cq_build_attempt_failed", + ] + + def chrome_tracing_dump(self, filename=None): + """Return a list of profiling events that can viewed as a timeline. + + To view this information as a timeline, simply dump it as a json file + by passing in "filename" or using using json.dump, and then load go to + chrome://tracing in the Chrome web browser and load the dumped file. + Make sure to enable "Flow events" in the "View Options" menu. + + Args: + filename: If a filename is provided, the timeline is dumped to that + file. + + Returns: + If filename is not provided, this returns a list of profiling + events. Each profile event is a dictionary. + """ + # TODO(rkn): Support including the task specification data in the + # timeline. + # TODO(rkn): This should support viewing just a window of time or a + # limited number of events. + + self._connect_and_get_accessor() + + # Add a small delay to account for propagation delay of events to the GCS. + # This should be harmless enough but prevents calls to timeline() from + # missing recent timeline data. + import time + + time.sleep(1) + + profile_events = self.profile_events() + all_events = [] + + for component_id_hex, component_events in profile_events.items(): + # Only consider workers and drivers. + component_type = component_events[0]["component_type"] + if component_type not in ["worker", "driver"]: + continue + + for event in component_events: + new_event = { + # The category of the event. + "cat": event["event_type"], + # The string displayed on the event. + "name": event["event_type"], + # The identifier for the group of rows that the event + # appears in. + "pid": event["node_ip_address"], + # The identifier for the row that the event appears in. + "tid": event["component_type"] + ":" + event["component_id"], + # The start time in microseconds. + "ts": self._nanoseconds_to_microseconds(event["start_time"]), + # The duration in microseconds. + "dur": self._nanoseconds_to_microseconds( + event["end_time"] - event["start_time"] + ), + # What is this? + "ph": "X", + # This is the name of the color to display the box in. + "cname": self._default_color_mapping[event["event_type"]], + # The extra user-defined data. + "args": event["extra_data"], + } + + # Modify the json with the additional user-defined extra data. + # This can be used to add fields or override existing fields. + if "cname" in event["extra_data"]: + new_event["cname"] = event["extra_data"]["cname"] + if "name" in event["extra_data"]: + new_event["name"] = event["extra_data"]["name"] + + all_events.append(new_event) + + if not all_events: + logger.warning( + "No profiling events found. Ray profiling must be enabled " + "by setting RAY_PROFILING=1, and make sure " + "RAY_task_events_report_interval_ms=0." + ) + + if filename is not None: + with open(filename, "w") as outfile: + json.dump(all_events, outfile) + else: + return all_events + + def chrome_tracing_object_transfer_dump(self, filename=None): + """Return a list of transfer events that can viewed as a timeline. + + To view this information as a timeline, simply dump it as a json file + by passing in "filename" or using json.dump, and then load go to + chrome://tracing in the Chrome web browser and load the dumped file. + Make sure to enable "Flow events" in the "View Options" menu. + + Args: + filename: If a filename is provided, the timeline is dumped to that + file. + + Returns: + If filename is not provided, this returns a list of profiling + events. Each profile event is a dictionary. + """ + self._connect_and_get_accessor() + + node_id_to_address = {} + for node_info in self.node_table(): + node_id_to_address[node_info["NodeID"]] = "{}:{}".format( + node_info["NodeManagerAddress"], node_info["ObjectManagerPort"] + ) + + all_events = [] + + for key, items in self.profile_events().items(): + # Only consider object manager events. + if items[0]["component_type"] != "object_manager": + continue + + for event in items: + if event["event_type"] == "transfer_send": + object_ref, remote_node_id, _, _ = event["extra_data"] + + elif event["event_type"] == "transfer_receive": + object_ref, remote_node_id, _ = event["extra_data"] + + elif event["event_type"] == "receive_pull_request": + object_ref, remote_node_id = event["extra_data"] + + else: + assert False, "This should be unreachable." + + # Choose a color by reading the first couple of hex digits of + # the object ref as an integer and turning that into a color. + object_ref_int = int(object_ref[:2], 16) + color = self._chrome_tracing_colors[ + object_ref_int % len(self._chrome_tracing_colors) + ] + + new_event = { + # The category of the event. + "cat": event["event_type"], + # The string displayed on the event. + "name": event["event_type"], + # The identifier for the group of rows that the event + # appears in. + "pid": node_id_to_address[key], + # The identifier for the row that the event appears in. + "tid": node_id_to_address[remote_node_id], + # The start time in microseconds. + "ts": self._nanoseconds_to_microseconds(event["start_time"]), + # The duration in microseconds. + "dur": self._nanoseconds_to_microseconds( + event["end_time"] - event["start_time"] + ), + # What is this? + "ph": "X", + # This is the name of the color to display the box in. + "cname": color, + # The extra user-defined data. + "args": event["extra_data"], + } + all_events.append(new_event) + + # Add another box with a color indicating whether it was a send + # or a receive event. + if event["event_type"] == "transfer_send": + additional_event = new_event.copy() + additional_event["cname"] = "black" + all_events.append(additional_event) + elif event["event_type"] == "transfer_receive": + additional_event = new_event.copy() + additional_event["cname"] = "grey" + all_events.append(additional_event) + else: + pass + + if filename is not None: + with open(filename, "w") as outfile: + json.dump(all_events, outfile) + else: + return all_events + + def workers(self): + """Get a dictionary mapping worker ID to worker information.""" + accessor = self._connect_and_get_accessor() + + # Get all data in worker table + worker_table = accessor.get_worker_table() + workers_data = {} + for i in range(len(worker_table)): + worker_table_data = gcs_pb2.WorkerTableData.FromString(worker_table[i]) + if ( + worker_table_data.is_alive + and worker_table_data.worker_type == common_pb2.WORKER + ): + worker_id = binary_to_hex(worker_table_data.worker_address.worker_id) + worker_info = worker_table_data.worker_info + + workers_data[worker_id] = { + "node_ip_address": decode(worker_info[b"node_ip_address"]), + "plasma_store_socket": decode(worker_info[b"plasma_store_socket"]), + } + if b"stderr_file" in worker_info: + workers_data[worker_id]["stderr_file"] = decode( + worker_info[b"stderr_file"] + ) + if b"stdout_file" in worker_info: + workers_data[worker_id]["stdout_file"] = decode( + worker_info[b"stdout_file"] + ) + return workers_data + + def add_worker(self, worker_id, worker_type, worker_info): + """Add a worker to the cluster. + + Args: + worker_id: ID of this worker. Type is bytes. + worker_type: Type of this worker. Value is common_pb2.DRIVER or + common_pb2.WORKER. + worker_info: Info of this worker. Type is dict{str: str}. + + Returns: + Is operation success + """ + accessor = self._connect_and_get_accessor() + + worker_data = gcs_pb2.WorkerTableData() + worker_data.is_alive = True + worker_data.worker_address.worker_id = worker_id + worker_data.worker_type = worker_type + for k, v in worker_info.items(): + worker_data.worker_info[k] = bytes(v, encoding="utf-8") + return accessor.add_worker_info(worker_data.SerializeToString()) + + def update_worker_debugger_port(self, worker_id, debugger_port): + """Update the debugger port of a worker. + + Args: + worker_id: ID of this worker. Type is bytes. + debugger_port: Port of the debugger. Type is int. + + Returns: + Is operation success + """ + accessor = self._connect_and_get_accessor() + + assert worker_id is not None, "worker_id is not valid" + assert ( + debugger_port is not None and debugger_port > 0 + ), "debugger_port is not valid" + + return accessor.update_worker_debugger_port(worker_id, debugger_port) + + def get_worker_debugger_port(self, worker_id): + """Get the debugger port of a worker. + + Args: + worker_id: ID of this worker. Type is bytes. + + Returns: + Debugger port of the worker. + """ + accessor = self._connect_and_get_accessor() + + assert worker_id is not None, "worker_id is not valid" + + return accessor.get_worker_debugger_port(worker_id) + + def update_worker_num_paused_threads(self, worker_id, num_paused_threads_delta): + """Updates the number of paused threads of a worker. + + Args: + worker_id: ID of this worker. Type is bytes. + num_paused_threads_delta: The delta of the number of paused threads. + + Returns: + Is operation success + """ + accessor = self._connect_and_get_accessor() + + assert worker_id is not None, "worker_id is not valid" + assert num_paused_threads_delta is not None, "worker_id is not valid" + + return accessor.update_worker_num_paused_threads( + worker_id, num_paused_threads_delta + ) + + def cluster_resources(self): + """Get the current total cluster resources. + + Note that this information can grow stale as nodes are added to or + removed from the cluster. + + Returns: + A dictionary mapping resource name to the total quantity of that + resource in the cluster. + """ + self._connect_and_get_accessor() + + # Calculate total resources. + total_resources = defaultdict(int) + for node_total_resources in self.total_resources_per_node().values(): + for resource_id, value in node_total_resources.items(): + total_resources[resource_id] += value + + return dict(total_resources) + + def _live_node_ids(self): + """Returns a set of node IDs corresponding to nodes still alive.""" + return set(self.total_resources_per_node().keys()) + + def available_resources_per_node(self): + """Returns a dictionary mapping node id to available resources.""" + accessor = self._connect_and_get_accessor() + available_resources_by_id = {} + + all_available_resources = accessor.get_all_available_resources() + for available_resource in all_available_resources: + message = gcs_pb2.AvailableResources.FromString(available_resource) + # Calculate available resources for this node. + dynamic_resources = {} + for resource_id, capacity in message.resources_available.items(): + dynamic_resources[resource_id] = capacity + # Update available resources for this node. + node_id = ray._common.utils.binary_to_hex(message.node_id) + available_resources_by_id[node_id] = dynamic_resources + + return available_resources_by_id + + # returns a dict that maps node_id(hex string) to a dict of {resource_id: capacity} + def total_resources_per_node(self) -> Dict[str, Dict[str, int]]: + accessor = self._connect_and_get_accessor() + total_resources_by_node = {} + + all_total_resources = accessor.get_all_total_resources() + for node_total_resources in all_total_resources: + message = gcs_pb2.TotalResources.FromString(node_total_resources) + # Calculate total resources for this node. + node_resources = {} + for resource_id, capacity in message.resources_total.items(): + node_resources[resource_id] = capacity + # Update total resources for this node. + node_id = ray._common.utils.binary_to_hex(message.node_id) + total_resources_by_node[node_id] = node_resources + + return total_resources_by_node + + def available_resources(self): + """Get the current available cluster resources. + + This is different from `cluster_resources` in that this will return + idle (available) resources rather than total resources. + + Note that this information can grow stale as tasks start and finish. + + Returns: + A dictionary mapping resource name to the total quantity of that + resource in the cluster. Note that if a resource (e.g., "CPU") + is currently not available (i.e., quantity is 0), it will not + be included in this dictionary. + """ + self._connect_and_get_accessor() + + available_resources_by_id = self.available_resources_per_node() + + # Calculate total available resources. + total_available_resources = defaultdict(int) + for available_resources in available_resources_by_id.values(): + for resource_id, num_available in available_resources.items(): + total_available_resources[resource_id] += num_available + + return dict(total_available_resources) + + def get_system_config(self): + """Get the system config of the cluster.""" + accessor = self._connect_and_get_accessor() + return json.loads(accessor.get_system_config()) + + def get_node_to_connect_for_driver(self, node_ip_address): + """Get the node to connect for a Ray driver.""" + accessor = self._connect_and_get_accessor() + return accessor.get_node_to_connect_for_driver(node_ip_address) + + def get_node(self, node_id: str): + """Get the node information for a node id.""" + accessor = self._connect_and_get_accessor() + return accessor.get_node(node_id) + + def get_draining_nodes(self) -> Dict[str, int]: + """Get all the hex ids of nodes that are being drained + and the corresponding draining deadline timestamps in ms. + + There is no deadline if the timestamp is 0. + """ + accessor = self._connect_and_get_accessor() + return accessor.get_draining_nodes() + + def get_cluster_config(self) -> autoscaler_pb2.ClusterConfig: + """Get the cluster config of the current cluster.""" + accessor = self._connect_and_get_accessor() + serialized_cluster_config = accessor.get_internal_kv( + ray._raylet.GCS_AUTOSCALER_STATE_NAMESPACE.encode(), + ray._raylet.GCS_AUTOSCALER_CLUSTER_CONFIG_KEY.encode(), + ) + if serialized_cluster_config: + return autoscaler_pb2.ClusterConfig.FromString(serialized_cluster_config) + return None + + @staticmethod + def _calculate_max_resource_from_cluster_config( + cluster_config: Optional[autoscaler_pb2.ClusterConfig], key: str + ) -> Optional[int]: + """Calculate the maximum available resources for a given resource type from cluster config. + If the resource type is not available, return None. + """ + if cluster_config is None: + return None + + max_value = 0 + for node_group_config in cluster_config.node_group_configs: + num_resources = node_group_config.resources.get(key, default=0) + num_nodes = node_group_config.max_count + if num_nodes == 0 or num_resources == 0: + continue + if num_nodes == -1 or num_resources == -1: + return sys.maxsize + max_value += num_nodes * num_resources + if max_value == 0: + return None + max_value_limit = cluster_config.max_resources.get(key, default=sys.maxsize) + return min(max_value, max_value_limit) + + def get_max_resources_from_cluster_config(self) -> Optional[Dict[str, int]]: + """Get the maximum available resources for all resource types from cluster config. + + Returns: + A dictionary mapping resource name to the maximum quantity of that + resource that could be available in the cluster based on the cluster config. + Returns None if the config is not available. + Values in the dictionary default to 0 if there is no such resource. + """ + all_resource_keys = set() + + config = self.get_cluster_config() + if config is None: + return None + + if config.node_group_configs: + for node_group_config in config.node_group_configs: + all_resource_keys.update(node_group_config.resources.keys()) + if len(all_resource_keys) == 0: + return None + + result = {} + for key in all_resource_keys: + max_value = self._calculate_max_resource_from_cluster_config(config, key) + result[key] = max_value if max_value is not None else 0 + + return result + + def get_actor_info(self, actor_id: ray.ActorID) -> Optional[str]: + """Get the actor info for a actor id.""" + accessor = self._connect_and_get_accessor() + return accessor.get_actor_info(actor_id) + + +state = GlobalState() +"""A global object used to access the cluster's global state.""" + + +def jobs(): + """Get a list of the jobs in the cluster (for debugging only). + + Returns: + Information from the job table, namely a list of dicts with keys: + - "JobID" (identifier for the job), + - "DriverIPAddress" (IP address of the driver for this job), + - "DriverPid" (process ID of the driver for this job), + - "StartTime" (UNIX timestamp of the start time of this job), + - "StopTime" (UNIX timestamp of the stop time of this job, if any) + """ + return state.job_table() + + +def next_job_id(): + """Get next job id from GCS. + + Returns: + Next job id in integer representation in the cluster. + """ + return state.next_job_id() + + +@DeveloperAPI +@client_mode_hook +def nodes(): + """Get a list of the nodes in the cluster (for debugging only). + + Returns: + Information about the Ray clients in the cluster. + """ + return state.node_table() + + +def workers(): + """Get a list of the workers in the cluster. + + Returns: + Information about the Ray workers in the cluster. + """ + return state.workers() + + +def current_node_id(): + """Return the node id of the current node. + + For example, "node:172.10.5.34". This can be used as a custom resource, + e.g., {node_id: 1} to reserve the whole node, or {node_id: 0.001} to + just force placement on the node. + + Returns: + Id of the current node. + """ + return NODE_ID_PREFIX + ray.util.get_node_ip_address() + + +def node_ids(): + """Get a list of the node ids in the cluster. + + For example, ["node:172.10.5.34", "node:172.42.3.77"]. These can be used + as custom resources, e.g., {node_id: 1} to reserve the whole node, or + {node_id: 0.001} to just force placement on the node. + + Returns: + List of the node resource ids. + """ + node_ids = [] + for node_total_resources in state.total_resources_per_node().values(): + for resource_id in node_total_resources.keys(): + if ( + resource_id.startswith(NODE_ID_PREFIX) + and resource_id != HEAD_NODE_RESOURCE_NAME + ): + node_ids.append(resource_id) + return node_ids + + +def actors( + actor_id: Optional[str] = None, + job_id: Optional[ray.JobID] = None, + actor_state_name: Optional[str] = None, +): + """Fetch actor info for one or more actor IDs (for debugging only). + + Args: + actor_id: A hex string of the actor ID to fetch information about. If + this is None, then all actor information is fetched. + If this is not None, `job_id` and `actor_state_name` + will not take effect. + job_id: To filter actors by job_id, which is of type `ray.JobID`. + You can use the `ray.get_runtime_context().job_id` function + to get the current job ID + actor_state_name: To filter actors based on actor state, + which can be one of the following: "DEPENDENCIES_UNREADY", + "PENDING_CREATION", "ALIVE", "RESTARTING", or "DEAD". + Returns: + Information about the actors. + """ + return state.actor_table( + actor_id=actor_id, job_id=job_id, actor_state_name=actor_state_name + ) + + +@DeveloperAPI +@client_mode_hook +def timeline(filename=None): + """Return a list of profiling events that can viewed as a timeline. + + Ray profiling must be enabled by setting the RAY_PROFILING=1 environment + variable prior to starting Ray, and set RAY_task_events_report_interval_ms=0 + + To view this information as a timeline, simply dump it as a json file by + passing in "filename" or using json.dump, and then load go to + chrome://tracing in the Chrome web browser and load the dumped file. + + Args: + filename: If a filename is provided, the timeline is dumped to that + file. + + Returns: + If filename is not provided, this returns a list of profiling events. + Each profile event is a dictionary. + """ + return state.chrome_tracing_dump(filename=filename) + + +def object_transfer_timeline(filename=None): + """Return a list of transfer events that can viewed as a timeline. + + To view this information as a timeline, simply dump it as a json file by + passing in "filename" or using json.dump, and then load go to + chrome://tracing in the Chrome web browser and load the dumped file. Make + sure to enable "Flow events" in the "View Options" menu. + + Args: + filename: If a filename is provided, the timeline is dumped to that + file. + + Returns: + If filename is not provided, this returns a list of profiling events. + Each profile event is a dictionary. + """ + return state.chrome_tracing_object_transfer_dump(filename=filename) + + +@DeveloperAPI +@client_mode_hook +def cluster_resources(): + """Get the current total cluster resources. + + Note that this information can grow stale as nodes are added to or removed + from the cluster. + + Returns: + A dictionary mapping resource name to the total quantity of that + resource in the cluster. + """ + return state.cluster_resources() + + +@DeveloperAPI +@client_mode_hook +def available_resources(): + """Get the current available cluster resources. + + This is different from `cluster_resources` in that this will return idle + (available) resources rather than total resources. + + Note that this information can grow stale as tasks start and finish. + + Returns: + A dictionary mapping resource name to the total quantity of that + resource in the cluster. Note that if a resource (e.g., "CPU") + is currently not available (i.e., quantity is 0), it will not + be included in this dictionary. + """ + return state.available_resources() + + +@DeveloperAPI +def available_resources_per_node(): + """Get the current available resources of each live node. + + Note that this information can grow stale as tasks start and finish. + + Returns: + A dictionary mapping node hex id to available resources dictionary. + """ + + return state.available_resources_per_node() + + +@DeveloperAPI +def total_resources_per_node(): + """Get the current total resources of each live node. + + Note that this information can grow stale as tasks start and finish. + + Returns: + A dictionary mapping node hex id to total resources dictionary. + """ + + return state.total_resources_per_node() + + +def update_worker_debugger_port(worker_id, debugger_port): + """Update the debugger port of a worker. + + Args: + worker_id: ID of this worker. Type is bytes. + debugger_port: Port of the debugger. Type is int. + + Returns: + Is operation success + """ + return state.update_worker_debugger_port(worker_id, debugger_port) + + +def update_worker_num_paused_threads(worker_id, num_paused_threads_delta): + """Update the number of paused threads of a worker. + + Args: + worker_id: ID of this worker. Type is bytes. + num_paused_threads_delta: The delta of the number of paused threads. + + Returns: + Is operation success + """ + return state.update_worker_num_paused_threads(worker_id, num_paused_threads_delta) + + +def get_worker_debugger_port(worker_id): + """Get the debugger port of a worker. + + Args: + worker_id: ID of this worker. Type is bytes. + + Returns: + Debugger port of the worker. + """ + return state.get_worker_debugger_port(worker_id) diff --git a/lib/python3.12/site-packages/ray/_private/state_api_test_utils.py b/lib/python3.12/site-packages/ray/_private/state_api_test_utils.py new file mode 100644 index 0000000000000000000000000000000000000000..a013842fa4586d4896651b4a711316dc5584a030 --- /dev/null +++ b/lib/python3.12/site-packages/ray/_private/state_api_test_utils.py @@ -0,0 +1,582 @@ +import asyncio +import concurrent.futures +import logging +import pprint +import sys +import time +import traceback +from collections import defaultdict +from concurrent.futures import ThreadPoolExecutor +from copy import deepcopy +from dataclasses import dataclass, field +from typing import Callable, Dict, List, Optional, Tuple, Union + +import numpy as np + +import ray +import ray._common.test_utils as test_utils +from ray._private.gcs_utils import GcsChannel +from ray._raylet import GcsClient +from ray.actor import ActorHandle +from ray.dashboard.state_aggregator import ( + StateAPIManager, +) +from ray.util.state import list_tasks, list_workers +from ray.util.state.common import ( + DEFAULT_LIMIT, + DEFAULT_RPC_TIMEOUT, + ListApiOptions, + PredicateType, + SupportedFilterType, +) +from ray.util.state.state_manager import StateDataSourceClient + +import psutil + + +@dataclass +class StateAPIMetric: + latency_sec: float + result_size: int + + +@dataclass +class StateAPICallSpec: + api: Callable + verify_cb: Callable + kwargs: Dict = field(default_factory=dict) + + +@dataclass +class StateAPIStats: + pending_calls: int = 0 + total_calls: int = 0 + calls: Dict = field(default_factory=lambda: defaultdict(list)) + + +GLOBAL_STATE_STATS = StateAPIStats() + +STATE_LIST_LIMIT = int(1e6) # 1m +STATE_LIST_TIMEOUT = 600 # 10min + + +def invoke_state_api( + verify_cb: Callable, + state_api_fn: Callable, + state_stats: StateAPIStats = GLOBAL_STATE_STATS, + key_suffix: Optional[str] = None, + print_result: Optional[bool] = False, + err_msg: Optional[str] = None, + **kwargs, +): + """Invoke a State API + + Args: + - verify_cb: Callback that takes in the response from `state_api_fn` and + returns a boolean, indicating the correctness of the results. + - state_api_fn: Function of the state API + - state_stats: Stats + - kwargs: Keyword arguments to be forwarded to the `state_api_fn` + """ + if "timeout" not in kwargs: + kwargs["timeout"] = STATE_LIST_TIMEOUT + + # Suppress missing output warning + kwargs["raise_on_missing_output"] = False + + res = None + try: + state_stats.total_calls += 1 + state_stats.pending_calls += 1 + + t_start = time.perf_counter() + res = state_api_fn(**kwargs) + t_end = time.perf_counter() + + if print_result: + pprint.pprint(res) + + metric = StateAPIMetric(t_end - t_start, len(res)) + if key_suffix: + key = f"{state_api_fn.__name__}_{key_suffix}" + else: + key = state_api_fn.__name__ + state_stats.calls[key].append(metric) + assert verify_cb( + res + ), f"Calling State API failed. len(res)=({len(res)}): {err_msg}" + except Exception as e: + traceback.print_exc() + assert ( + False + ), f"Calling {state_api_fn.__name__}({kwargs}) failed with {repr(e)}." + finally: + state_stats.pending_calls -= 1 + + return res + + +def invoke_state_api_n(*args, **kwargs): + def verify(): + NUM_API_CALL_SAMPLES = 10 + for _ in range(NUM_API_CALL_SAMPLES): + invoke_state_api(*args, **kwargs) + return True + + test_utils.wait_for_condition(verify, retry_interval_ms=2000, timeout=30) + + +def aggregate_perf_results(state_stats: StateAPIStats = GLOBAL_STATE_STATS): + """Aggregate stats of state API calls + + Return: + This returns a dict of below fields: + - max_{api_key_name}_latency_sec: + Max latency of call to {api_key_name} + - {api_key_name}_result_size_with_max_latency: + The size of the result (or the number of bytes for get_log API) + for the max latency invocation + - avg/p99/p95/p50_{api_key_name}_latency_sec: + The percentile latency stats + - avg_state_api_latency_sec: + The average latency of all the state apis tracked + """ + # Prevent iteration when modifying error + state_stats = deepcopy(state_stats) + perf_result = {} + for api_key_name, metrics in state_stats.calls.items(): + # Per api aggregation + # Max latency + latency_key = f"max_{api_key_name}_latency_sec" + size_key = f"{api_key_name}_result_size_with_max_latency" + metric = max(metrics, key=lambda metric: metric.latency_sec) + + perf_result[latency_key] = metric.latency_sec + perf_result[size_key] = metric.result_size + + latency_list = np.array([metric.latency_sec for metric in metrics]) + # avg latency + key = f"avg_{api_key_name}_latency_sec" + perf_result[key] = np.average(latency_list) + + # p99 latency + key = f"p99_{api_key_name}_latency_sec" + perf_result[key] = np.percentile(latency_list, 99) + + # p95 latency + key = f"p95_{api_key_name}_latency_sec" + perf_result[key] = np.percentile(latency_list, 95) + + # p50 latency + key = f"p50_{api_key_name}_latency_sec" + perf_result[key] = np.percentile(latency_list, 50) + + all_state_api_latency = sum( + metric.latency_sec + for metric_samples in state_stats.calls.values() + for metric in metric_samples + ) + + perf_result["avg_state_api_latency_sec"] = ( + (all_state_api_latency / state_stats.total_calls) + if state_stats.total_calls != 0 + else -1 + ) + + return perf_result + + +@ray.remote(num_cpus=0) +class StateAPIGeneratorActor: + def __init__( + self, + apis: List[StateAPICallSpec], + call_interval_s: float = 5.0, + print_interval_s: float = 20.0, + wait_after_stop: bool = True, + print_result: bool = False, + ) -> None: + """An actor that periodically issues state API + + Args: + - apis: List of StateAPICallSpec + - call_interval_s: State apis in the `apis` will be issued + every `call_interval_s` seconds. + - print_interval_s: How frequent state api stats will be dumped. + - wait_after_stop: When true, call to `ray.get(actor.stop.remote())` + will wait for all pending state APIs to return. + Setting it to `False` might miss some long-running state apis calls. + - print_result: True if result of each API call is printed. Default False. + """ + # Configs + self._apis = apis + self._call_interval_s = call_interval_s + self._print_interval_s = print_interval_s + self._wait_after_cancel = wait_after_stop + self._logger = logging.getLogger(self.__class__.__name__) + self._print_result = print_result + + # States + self._tasks = None + self._fut_queue = None + self._executor = None + self._loop = None + self._stopping = False + self._stopped = False + self._stats = StateAPIStats() + + async def start(self): + # Run the periodic api generator + self._fut_queue = asyncio.Queue() + self._executor = concurrent.futures.ThreadPoolExecutor() + + self._tasks = [ + asyncio.ensure_future(awt) + for awt in [ + self._run_generator(), + self._run_result_waiter(), + self._run_stats_reporter(), + ] + ] + await asyncio.gather(*self._tasks) + + def call(self, fn, verify_cb, **kwargs): + def run_fn(): + try: + self._logger.debug(f"calling {fn.__name__}({kwargs})") + return invoke_state_api( + verify_cb, + fn, + state_stats=self._stats, + print_result=self._print_result, + **kwargs, + ) + except Exception as e: + self._logger.warning(f"{fn.__name__}({kwargs}) failed with: {repr(e)}") + return None + + fut = asyncio.get_running_loop().run_in_executor(self._executor, run_fn) + return fut + + async def _run_stats_reporter(self): + while not self._stopped: + # Keep the reporter running until all pending apis finish and the bool + # `self._stopped` is then True + self._logger.info(pprint.pprint(aggregate_perf_results(self._stats))) + try: + await asyncio.sleep(self._print_interval_s) + except asyncio.CancelledError: + self._logger.info( + "_run_stats_reporter cancelled, " + f"waiting for all api {self._stats.pending_calls}calls to return..." + ) + + async def _run_generator(self): + try: + while not self._stopping: + # Run the state API in another thread + for api_spec in self._apis: + fut = self.call(api_spec.api, api_spec.verify_cb, **api_spec.kwargs) + self._fut_queue.put_nowait(fut) + + await asyncio.sleep(self._call_interval_s) + except asyncio.CancelledError: + # Stop running + self._logger.info("_run_generator cancelled, now stopping...") + return + + async def _run_result_waiter(self): + try: + while not self._stopping: + fut = await self._fut_queue.get() + await fut + except asyncio.CancelledError: + self._logger.info( + f"_run_result_waiter cancelled, cancelling {self._fut_queue.qsize()} " + "pending futures..." + ) + while not self._fut_queue.empty(): + fut = self._fut_queue.get_nowait() + if self._wait_after_cancel: + await fut + else: + # Ignore the queue futures if we are not + # waiting on them after stop() called + fut.cancel() + return + + def get_stats(self): + # deep copy to prevent race between reporting and modifying stats + return aggregate_perf_results(self._stats) + + def ready(self): + pass + + def stop(self): + self._stopping = True + self._logger.debug(f"calling stop, canceling {len(self._tasks)} tasks") + for task in self._tasks: + task.cancel() + + # This will block the stop() function until all futures are cancelled + # if _wait_after_cancel=True. When _wait_after_cancel=False, it will still + # wait for any in-progress futures. + # See: https://docs.python.org/3.8/library/concurrent.futures.html + self._executor.shutdown(wait=self._wait_after_cancel) + self._stopped = True + + +def periodic_invoke_state_apis_with_actor(*args, **kwargs) -> ActorHandle: + current_node_ip = ray._private.worker.global_worker.node_ip_address + # Schedule the actor on the current node. + actor = StateAPIGeneratorActor.options( + resources={f"node:{current_node_ip}": 0.001} + ).remote(*args, **kwargs) + print("Waiting for state api actor to be ready...") + ray.get(actor.ready.remote()) + print("State api actor is ready now.") + actor.start.remote() + return actor + + +def get_state_api_manager(gcs_address: str) -> StateAPIManager: + gcs_client = GcsClient(address=gcs_address) + gcs_channel = GcsChannel(gcs_address=gcs_address, aio=True) + gcs_channel.connect() + state_api_data_source_client = StateDataSourceClient( + gcs_channel.channel(), gcs_client + ) + return StateAPIManager( + state_api_data_source_client, + thread_pool_executor=ThreadPoolExecutor( + thread_name_prefix="state_api_test_utils" + ), + ) + + +def summarize_worker_startup_time(): + workers = list_workers( + detail=True, + filters=[("worker_type", "=", "WORKER")], + limit=10000, + raise_on_missing_output=False, + ) + time_to_launch = [] + time_to_initialize = [] + for worker in workers: + launch_time = worker.get("worker_launch_time_ms") + launched_time = worker.get("worker_launched_time_ms") + start_time = worker.get("start_time_ms") + + if launched_time > 0: + time_to_launch.append(launched_time - launch_time) + if start_time: + time_to_initialize.append(start_time - launched_time) + time_to_launch.sort() + time_to_initialize.sort() + + def print_latencies(latencies): + print(f"Avg: {round(sum(latencies) / len(latencies), 2)} ms") + print(f"P25: {round(latencies[int(len(latencies) * 0.25)], 2)} ms") + print(f"P50: {round(latencies[int(len(latencies) * 0.5)], 2)} ms") + print(f"P95: {round(latencies[int(len(latencies) * 0.95)], 2)} ms") + print(f"P99: {round(latencies[int(len(latencies) * 0.99)], 2)} ms") + + print("Time to launch workers") + print_latencies(time_to_launch) + print("=======================") + print("Time to initialize workers") + print_latencies(time_to_initialize) + + +def verify_failed_task( + name: str, error_type: str, error_message: Union[str, List[str], None] = None +) -> bool: + """ + Check if a task with 'name' has failed with the exact error type 'error_type' + and 'error_message' in the error message. + """ + tasks = list_tasks(filters=[("name", "=", name)], detail=True) + assert len(tasks) == 1, tasks + t = tasks[0] + assert t["state"] == "FAILED", t + assert t["error_type"] == error_type, t + if error_message is not None: + if isinstance(error_message, str): + error_message = [error_message] + for msg in error_message: + assert msg in t.get("error_message", None), t + return True + + +@ray.remote +class PidActor: + def __init__(self): + self.name_to_pid = {} + + def get_pids(self): + return self.name_to_pid + + def report_pid(self, name, pid, state=None): + self.name_to_pid[name] = (pid, state) + + +def _is_actor_task_running(actor_pid: int, task_name: str): + """ + Check whether the actor task `task_name` is running on the actor process + with pid `actor_pid`. + + Args: + actor_pid: The pid of the actor process. + task_name: The name of the actor task. + + Returns: + True if the actor task is running, False otherwise. + + Limitation: + If the actor task name is set using options.name and is a substring of + the actor name, this function may return true even if the task is not + running on the actor process. To resolve this issue, we can possibly + pass in the actor name. + """ + if not psutil.pid_exists(actor_pid): + return False + + """ + Why use both `psutil.Process.name()` and `psutil.Process.cmdline()`? + + 1. Core worker processes call `setproctitle` to set the process title before + and after executing tasks. However, the definition of "title" is a bit + complex. + + [ref]: https://github.com/dvarrazzo/py-setproctitle + + > The process title is usually visible in files such as /proc/PID/cmdline, + /proc/PID/status, /proc/PID/comm, depending on the operating system and + kernel version. This information is used by user-space tools such as ps + and top. + + Ideally, we would only need to check `psutil.Process.cmdline()`, but I decided + to check both `psutil.Process.name()` and `psutil.Process.cmdline()` based on + the definition of "title" stated above. + + 2. Additionally, the definition of `psutil.Process.name()` is not consistent + with the definition of "title" in `setproctitle`. The length of `/proc/PID/comm` and + the prefix of `/proc/PID/cmdline` affect the return value of + `psutil.Process.name()`. + + In addition, executing `setproctitle` in different threads within the same + process may result in different outcomes. + + To learn more details, please refer to the source code of `psutil`: + + [ref]: + https://github.com/giampaolo/psutil/blob/a17550784b0d3175da01cdb02cee1bc6b61637dc/psutil/__init__.py#L664-L693 + + 3. `/proc/PID/comm` will be truncated to TASK_COMM_LEN (16) characters + (including the terminating null byte). + + [ref]: + https://man7.org/linux/man-pages/man5/proc_pid_comm.5.html + """ + name = psutil.Process(actor_pid).name() + if task_name in name and name.startswith("ray::"): + return True + + cmdline = psutil.Process(actor_pid).cmdline() + # If `options.name` is set, the format is `ray::`. If not, + # the format is `ray::.`. + if cmdline and task_name in cmdline[0] and cmdline[0].startswith("ray::"): + return True + return False + + +def verify_tasks_running_or_terminated( + task_pids: Dict[str, Tuple[int, Optional[str]]], expect_num_tasks: int +): + """ + Check if the tasks in task_pids are in RUNNING state if pid exists + and running the task. + If the pid is missing or the task is not running the task, check if the task + is marked FAILED or FINISHED. + + Args: + task_pids: A dict of task name to (pid, expected terminal state). + + """ + assert len(task_pids) == expect_num_tasks, task_pids + for task_name, pid_and_state in task_pids.items(): + tasks = list_tasks(detail=True, filters=[("name", "=", task_name)]) + assert len(tasks) == 1, ( + f"One unique task with {task_name} should be found. " + "Use `options(name=)` when creating the task." + ) + task = tasks[0] + pid, expected_state = pid_and_state + + # If it's windows/macos, we don't have a way to check if the process + # is actually running the task since the process name is just python, + # rather than the actual task name. + if sys.platform in ["win32", "darwin"]: + if expected_state is not None: + assert task["state"] == expected_state, task + continue + if _is_actor_task_running(pid, task_name): + assert ( + "ray::IDLE" not in task["name"] + ), "One should not name it 'IDLE' since it's reserved in Ray" + assert task["state"] == "RUNNING", task + if expected_state is not None: + assert task["state"] == expected_state, task + else: + # Tasks no longer running. + if expected_state is None: + assert task["state"] in [ + "FAILED", + "FINISHED", + ], f"{task_name}: {task['task_id']} = {task['state']}" + else: + assert ( + task["state"] == expected_state + ), f"expect {expected_state} but {task['state']} for {task}" + + return True + + +def verify_schema(state, result_dict: dict, detail: bool = False): + """ + Verify the schema of the result_dict is the same as the state. + """ + state_fields_columns = set() + if detail: + state_fields_columns = state.columns() + else: + state_fields_columns = state.base_columns() + + for k in state_fields_columns: + assert k in result_dict + + for k in result_dict: + assert k in state_fields_columns + + # Make the field values can be converted without error as well + state(**result_dict) + + +def create_api_options( + timeout: int = DEFAULT_RPC_TIMEOUT, + limit: int = DEFAULT_LIMIT, + filters: List[Tuple[str, PredicateType, SupportedFilterType]] = None, + detail: bool = False, + exclude_driver: bool = True, +): + if not filters: + filters = [] + return ListApiOptions( + limit=limit, + timeout=timeout, + filters=filters, + server_timeout_multiplier=1.0, + detail=detail, + exclude_driver=exclude_driver, + ) diff --git a/lib/python3.12/site-packages/ray/_private/telemetry/__init__.py b/lib/python3.12/site-packages/ray/_private/telemetry/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/lib/python3.12/site-packages/ray/_private/telemetry/__pycache__/__init__.cpython-312.pyc b/lib/python3.12/site-packages/ray/_private/telemetry/__pycache__/__init__.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..6c69e97bb1aad29a8624f8695e943d2337a60f1f Binary files /dev/null and b/lib/python3.12/site-packages/ray/_private/telemetry/__pycache__/__init__.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/_private/telemetry/__pycache__/metric_cardinality.cpython-312.pyc b/lib/python3.12/site-packages/ray/_private/telemetry/__pycache__/metric_cardinality.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..6e7c6002e2d2e4362309b84322d1d8f0fe977447 Binary files /dev/null and b/lib/python3.12/site-packages/ray/_private/telemetry/__pycache__/metric_cardinality.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/_private/telemetry/__pycache__/open_telemetry_metric_recorder.cpython-312.pyc b/lib/python3.12/site-packages/ray/_private/telemetry/__pycache__/open_telemetry_metric_recorder.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..d96ae648403229e491c9be598953b1b74d613546 Binary files /dev/null and b/lib/python3.12/site-packages/ray/_private/telemetry/__pycache__/open_telemetry_metric_recorder.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/_private/telemetry/metric_cardinality.py b/lib/python3.12/site-packages/ray/_private/telemetry/metric_cardinality.py new file mode 100644 index 0000000000000000000000000000000000000000..17a8c5ce6bbbc464683a76b27b1b37c0b0964f7d --- /dev/null +++ b/lib/python3.12/site-packages/ray/_private/telemetry/metric_cardinality.py @@ -0,0 +1,62 @@ +from enum import Enum +from typing import Dict, List + +from ray._private.ray_constants import RAY_METRIC_CARDINALITY_LEVEL + +# Keep in sync with the WorkerIdKey in src/ray/stats/tag_defs.cc +WORKER_ID_TAG_KEY = "WorkerId" +# Keep in sync with the NameKey in src/ray/stats/metric_defs.cc +TASK_OR_ACTOR_NAME_TAG_KEY = "Name" + +_CARDINALITY_LEVEL = None +_HIGH_CARDINALITY_LABELS: Dict[str, List[str]] = {} + + +class MetricCardinality(str, Enum): + """Cardinality level configuration for all Ray metrics (ray_tasks, ray_actors, + etc.). This configurtion is used to determine whether to globally drop high + cardinality labels. This is important for high scale clusters that might consist + thousands of workers, millions of tasks. + + - LEGACY: Keep all labels. This is the default behavior. + - RECOMMENDED: Drop high cardinality labels. The set of high cardinality labels + are determined internally by Ray and not exposed to users. Currently, this includes + the following labels: WorkerId + - LOW: Same as RECOMMENDED, but also drop the Name label for tasks and actors. + """ + + LEGACY = "legacy" + RECOMMENDED = "recommended" + LOW = "low" + + @staticmethod + def get_cardinality_level() -> "MetricCardinality": + global _CARDINALITY_LEVEL + if _CARDINALITY_LEVEL is not None: + return _CARDINALITY_LEVEL + try: + _CARDINALITY_LEVEL = MetricCardinality(RAY_METRIC_CARDINALITY_LEVEL.lower()) + except ValueError: + _CARDINALITY_LEVEL = MetricCardinality.LEGACY + return _CARDINALITY_LEVEL + + @staticmethod + def get_high_cardinality_labels_to_drop(metric_name: str) -> List[str]: + """ + Get the high cardinality labels of the metric. + """ + if metric_name in _HIGH_CARDINALITY_LABELS: + return _HIGH_CARDINALITY_LABELS[metric_name] + + cardinality_level = MetricCardinality.get_cardinality_level() + if cardinality_level == MetricCardinality.LEGACY: + _HIGH_CARDINALITY_LABELS[metric_name] = [] + return [] + + _HIGH_CARDINALITY_LABELS[metric_name] = [WORKER_ID_TAG_KEY] + if cardinality_level == MetricCardinality.LOW and metric_name in [ + "tasks", + "actors", + ]: + _HIGH_CARDINALITY_LABELS[metric_name].append(TASK_OR_ACTOR_NAME_TAG_KEY) + return _HIGH_CARDINALITY_LABELS[metric_name] diff --git a/lib/python3.12/site-packages/ray/_private/telemetry/open_telemetry_metric_recorder.py b/lib/python3.12/site-packages/ray/_private/telemetry/open_telemetry_metric_recorder.py new file mode 100644 index 0000000000000000000000000000000000000000..d3b857e00eee00292bef1343a88781682c7be39e --- /dev/null +++ b/lib/python3.12/site-packages/ray/_private/telemetry/open_telemetry_metric_recorder.py @@ -0,0 +1,232 @@ +import logging +import threading +from collections import defaultdict +from typing import List + +from opentelemetry import metrics +from opentelemetry.exporter.prometheus import PrometheusMetricReader +from opentelemetry.metrics import Observation +from opentelemetry.sdk.metrics import MeterProvider + +from ray._private.metrics_agent import Record +from ray._private.telemetry.metric_cardinality import MetricCardinality + +logger = logging.getLogger(__name__) + +NAMESPACE = "ray" + + +class OpenTelemetryMetricRecorder: + """ + A class to record OpenTelemetry metrics. This is the main entry point for exporting + all ray telemetries to Prometheus server. + It uses OpenTelemetry's Prometheus exporter to export metrics. + """ + + _metrics_initialized = False + _metrics_initialized_lock = threading.Lock() + + def __init__(self): + self._lock = threading.Lock() + self._registered_instruments = {} + self._observations_by_name = defaultdict(dict) + self._histogram_bucket_midpoints = defaultdict(list) + self._init_metrics() + self.meter = metrics.get_meter(__name__) + + def _init_metrics(self): + # Initialize the global metrics provider and meter. We only do this once on + # the first initialization of the class, because re-setting the meter provider + # can result in loss of metrics. + with self._metrics_initialized_lock: + if self._metrics_initialized: + return + prometheus_reader = PrometheusMetricReader() + provider = MeterProvider(metric_readers=[prometheus_reader]) + metrics.set_meter_provider(provider) + self._metrics_initialized = True + + def register_gauge_metric(self, name: str, description: str) -> None: + with self._lock: + if name in self._registered_instruments: + # Gauge with the same name is already registered. + return + + # Register ObservableGauge with a dynamic callback. Callbacks are special + # features in OpenTelemetry that allow you to provide a function that will + # compute the telemetry at collection time. + def callback(options): + # Take snapshot of current observations. + with self._lock: + observations = self._observations_by_name[name] + # Clear the observations to avoid emitting dead observations. + self._observations_by_name[name] = {} + # Drop high cardinality from tag_set and sum up the value for + # same tag set after dropping + aggregated_observations = defaultdict(float) + high_cardinality_labels = ( + MetricCardinality.get_high_cardinality_labels_to_drop(name) + ) + for tag_set, val in observations.items(): + # Convert frozenset back to dict + tags_dict = dict(tag_set) + # Filter out high cardinality labels + filtered_tags = { + k: v + for k, v in tags_dict.items() + if k not in high_cardinality_labels + } + # Create a key for aggregation + filtered_key = frozenset(filtered_tags.items()) + # Sum up values for the same filtered tag set + aggregated_observations[filtered_key] += val + + return [ + Observation(val, attributes=dict(tag_set)) + for tag_set, val in aggregated_observations.items() + ] + + instrument = self.meter.create_observable_gauge( + name=f"{NAMESPACE}_{name}", + description=description, + unit="1", + callbacks=[callback], + ) + self._registered_instruments[name] = instrument + self._observations_by_name[name] = {} + + def register_counter_metric(self, name: str, description: str) -> None: + """ + Register a counter metric with the given name and description. + """ + with self._lock: + if name in self._registered_instruments: + # Counter with the same name is already registered. This is a common + # case when metrics are exported from multiple Ray components (e.g., + # raylet, worker, etc.) running in the same node. Since each component + # may export metrics with the same name, the same metric might be + # registered multiple times. + return + + instrument = self.meter.create_counter( + name=f"{NAMESPACE}_{name}", + description=description, + unit="1", + ) + self._registered_instruments[name] = instrument + + def register_sum_metric(self, name: str, description: str) -> None: + """ + Register a sum metric with the given name and description. + """ + with self._lock: + if name in self._registered_instruments: + # Sum with the same name is already registered. This is a common + # case when metrics are exported from multiple Ray components (e.g., + # raylet, worker, etc.) running in the same node. Since each component + # may export metrics with the same name, the same metric might be + # registered multiple times. + return + + instrument = self.meter.create_up_down_counter( + name=f"{NAMESPACE}_{name}", + description=description, + unit="1", + ) + self._registered_instruments[name] = instrument + + def register_histogram_metric( + self, name: str, description: str, buckets: List[float] + ) -> None: + """ + Register a histogram metric with the given name and description. + """ + with self._lock: + if name in self._registered_instruments: + # Histogram with the same name is already registered. This is a common + # case when metrics are exported from multiple Ray components (e.g., + # raylet, worker, etc.) running in the same node. Since each component + # may export metrics with the same name, the same metric might be + # registered multiple times. + return + + instrument = self.meter.create_histogram( + name=f"{NAMESPACE}_{name}", + description=description, + unit="1", + explicit_bucket_boundaries_advisory=buckets, + ) + self._registered_instruments[name] = instrument + + # calculate the bucket midpoints; this is used for converting histogram + # internal representation to approximated histogram data points. + for i in range(len(buckets)): + if i == 0: + lower_bound = 0.0 if buckets[0] > 0 else buckets[0] * 2.0 + self._histogram_bucket_midpoints[name].append( + (lower_bound + buckets[0]) / 2.0 + ) + else: + self._histogram_bucket_midpoints[name].append( + (buckets[i] + buckets[i - 1]) / 2.0 + ) + # Approximated mid point for Inf+ bucket. Inf+ bucket is an implicit bucket + # that is not part of buckets. + self._histogram_bucket_midpoints[name].append( + 1.0 if buckets[-1] <= 0 else buckets[-1] * 2.0 + ) + + def get_histogram_bucket_midpoints(self, name: str) -> List[float]: + """ + Get the bucket midpoints for a histogram metric with the given name. + """ + return self._histogram_bucket_midpoints[name] + + def set_metric_value(self, name: str, tags: dict, value: float): + """ + Set the value of a metric with the given name and tags. If the metric is not + registered, it lazily records the value for observable metrics or is a no-op for + synchronous metrics. + """ + with self._lock: + if self._observations_by_name.get(name) is not None: + # Set the value of an observable metric with the given name and tags. It + # lazily records the metric value by storing it in a dictionary until + # the value actually gets exported by OpenTelemetry. + self._observations_by_name[name][frozenset(tags.items())] = value + else: + instrument = self._registered_instruments.get(name) + tags = { + k: v + for k, v in tags.items() + if k + not in MetricCardinality.get_high_cardinality_labels_to_drop(name) + } + if isinstance(instrument, metrics.Counter): + instrument.add(value, attributes=tags) + elif isinstance(instrument, metrics.UpDownCounter): + instrument.add(value, attributes=tags) + elif isinstance(instrument, metrics.Histogram): + instrument.record(value, attributes=tags) + else: + logger.warning( + f"Unsupported synchronous instrument type for metric: {name}." + ) + + def record_and_export(self, records: List[Record], global_tags=None): + """ + Record a list of telemetry records and export them to Prometheus. + """ + global_tags = global_tags or {} + + for record in records: + gauge = record.gauge + value = record.value + tags = {**record.tags, **global_tags} + try: + self.register_gauge_metric(gauge.name, gauge.description or "") + self.set_metric_value(gauge.name, tags, value) + except Exception as e: + logger.error( + f"Failed to record metric {gauge.name} with value {value} with tags {tags!r} and global tags {global_tags!r} due to: {e!r}" + ) diff --git a/lib/python3.12/site-packages/ray/_private/test_utils.py b/lib/python3.12/site-packages/ray/_private/test_utils.py new file mode 100644 index 0000000000000000000000000000000000000000..bb3c1ab7e07494f71f5288f5115a0fdb1bb84826 --- /dev/null +++ b/lib/python3.12/site-packages/ray/_private/test_utils.py @@ -0,0 +1,2107 @@ +import asyncio +import fnmatch +import io +import json +import logging +import os +import pathlib +import random +import socket +import subprocess +import sys +import tempfile +import threading +import time +import timeit +import traceback +import uuid +from collections import defaultdict +from contextlib import contextmanager, redirect_stderr, redirect_stdout +from dataclasses import dataclass, field +from datetime import datetime +from typing import Any, Callable, Dict, List, Optional, Set, Tuple +from urllib.parse import quote + +import requests +import yaml + +import ray +import ray._private.memory_monitor as memory_monitor +import ray._private.services +import ray._private.services as services +import ray._private.utils +import ray.dashboard.consts as dashboard_consts +from ray._common.network_utils import build_address, parse_address +from ray._common.test_utils import wait_for_condition +from ray._common.utils import get_or_create_event_loop +from ray._private import ( + ray_constants, +) +from ray._private.internal_api import memory_summary +from ray._private.tls_utils import generate_self_signed_tls_certs +from ray._private.worker import RayContext +from ray._raylet import Config, GcsClient, GcsClientOptions, GlobalStateAccessor +from ray.core.generated import ( + gcs_pb2, + gcs_service_pb2, + node_manager_pb2, +) +from ray.util.queue import Empty, Queue, _QueueActor +from ray.util.scheduling_strategies import NodeAffinitySchedulingStrategy +from ray.util.state import get_actor, list_actors + +import psutil # We must import psutil after ray because we bundle it with ray. + +logger = logging.getLogger(__name__) + +EXE_SUFFIX = ".exe" if sys.platform == "win32" else "" +RAY_PATH = os.path.abspath(os.path.dirname(os.path.dirname(__file__))) +REDIS_EXECUTABLE = os.path.join( + RAY_PATH, "core/src/ray/thirdparty/redis/src/redis-server" + EXE_SUFFIX +) + +try: + from prometheus_client.core import Metric + from prometheus_client.parser import Sample, text_string_to_metric_families +except (ImportError, ModuleNotFoundError): + Metric = None + Sample = None + + def text_string_to_metric_families(*args, **kwargs): + raise ModuleNotFoundError("`prometheus_client` not found") + + +def make_global_state_accessor(ray_context): + gcs_options = GcsClientOptions.create( + ray_context.address_info["gcs_address"], + None, + allow_cluster_id_nil=True, + fetch_cluster_id_if_nil=False, + ) + global_state_accessor = GlobalStateAccessor(gcs_options) + global_state_accessor.connect() + return global_state_accessor + + +def external_redis_test_enabled(): + return os.environ.get("TEST_EXTERNAL_REDIS") == "1" + + +def redis_replicas(): + return int(os.environ.get("TEST_EXTERNAL_REDIS_REPLICAS", "1")) + + +def redis_sentinel_replicas(): + return int(os.environ.get("TEST_EXTERNAL_REDIS_SENTINEL_REPLICAS", "2")) + + +def get_redis_cli(port, enable_tls): + try: + # If there is no redis libs installed, skip the check. + # This could happen In minimal test, where we don't have + # redis. + import redis + except Exception: + return True + + params = {} + if enable_tls: + from ray._raylet import Config + + params = {"ssl": True, "ssl_cert_reqs": "required"} + if Config.REDIS_CA_CERT(): + params["ssl_ca_certs"] = Config.REDIS_CA_CERT() + if Config.REDIS_CLIENT_CERT(): + params["ssl_certfile"] = Config.REDIS_CLIENT_CERT() + if Config.REDIS_CLIENT_KEY(): + params["ssl_keyfile"] = Config.REDIS_CLIENT_KEY() + + return redis.Redis("localhost", str(port), **params) + + +def start_redis_sentinel_instance( + session_dir_path: str, + port: int, + redis_master_port: int, + password: Optional[str] = None, + enable_tls: bool = False, + db_dir=None, + free_port=0, +): + config_file = os.path.join( + session_dir_path, "redis-sentinel-" + uuid.uuid4().hex + ".conf" + ) + config_lines = [] + # Port for this Sentinel instance + if enable_tls: + config_lines.append(f"port {free_port}") + else: + config_lines.append(f"port {port}") + + # Monitor the Redis master + config_lines.append(f"sentinel monitor redis-test 127.0.0.1 {redis_master_port} 1") + config_lines.append( + "sentinel down-after-milliseconds redis-test 1000" + ) # failover after 1 second + config_lines.append("sentinel failover-timeout redis-test 5000") # + config_lines.append("sentinel parallel-syncs redis-test 1") + + if password: + config_lines.append(f"sentinel auth-pass redis-test {password}") + + if enable_tls: + config_lines.append(f"tls-port {port}") + if Config.REDIS_CA_CERT(): + config_lines.append(f"tls-ca-cert-file {Config.REDIS_CA_CERT()}") + # Check and add TLS client certificate file + if Config.REDIS_CLIENT_CERT(): + config_lines.append(f"tls-cert-file {Config.REDIS_CLIENT_CERT()}") + # Check and add TLS client key file + if Config.REDIS_CLIENT_KEY(): + config_lines.append(f"tls-key-file {Config.REDIS_CLIENT_KEY()}") + config_lines.append("tls-auth-clients no") + config_lines.append("sentinel tls-auth-clients redis-test no") + if db_dir: + config_lines.append(f"dir {db_dir}") + + with open(config_file, "w") as f: + f.write("\n".join(config_lines)) + + command = [REDIS_EXECUTABLE, config_file, "--sentinel"] + process_info = ray._private.services.start_ray_process( + command, + ray_constants.PROCESS_TYPE_REDIS_SERVER, + fate_share=False, + ) + return process_info + + +def start_redis_instance( + session_dir_path: str, + port: int, + redis_max_clients: Optional[int] = None, + num_retries: int = 20, + stdout_file: Optional[str] = None, + stderr_file: Optional[str] = None, + password: Optional[str] = None, + fate_share: Optional[bool] = None, + port_denylist: Optional[List[int]] = None, + listen_to_localhost_only: bool = False, + enable_tls: bool = False, + replica_of=None, + leader_id=None, + db_dir=None, + free_port=0, +): + """Start a single Redis server. + + Notes: + We will initially try to start the Redis instance at the given port, + and then try at most `num_retries - 1` times to start the Redis + instance at successive random ports. + + Args: + session_dir_path: Path to the session directory of + this Ray cluster. + port: Try to start a Redis server at this port. + redis_max_clients: If this is provided, Ray will attempt to configure + Redis with this maxclients number. + num_retries: The number of times to attempt to start Redis at + successive ports. + stdout_file: A file handle opened for writing to redirect stdout to. If + no redirection should happen, then this should be None. + stderr_file: A file handle opened for writing to redirect stderr to. If + no redirection should happen, then this should be None. + password: Prevents external clients without the password + from connecting to Redis if provided. + port_denylist: A set of denylist ports that shouldn't + be used when allocating a new port. + listen_to_localhost_only: Redis server only listens to + localhost (127.0.0.1) if it's true, + otherwise it listens to all network interfaces. + enable_tls: Enable the TLS/SSL in Redis or not + + Returns: + A tuple of the port used by Redis and ProcessInfo for the process that + was started. If a port is passed in, then the returned port value + is the same. + + Raises: + Exception: An exception is raised if Redis could not be started. + """ + + assert os.path.isfile(REDIS_EXECUTABLE) + + # Construct the command to start the Redis server. + command = [REDIS_EXECUTABLE] + if password: + if " " in password: + raise ValueError("Spaces not permitted in redis password.") + command += ["--requirepass", password] + if redis_replicas() > 1: + command += ["--cluster-enabled", "yes", "--cluster-config-file", f"node-{port}"] + if enable_tls: + command += [ + "--tls-port", + str(port), + "--loglevel", + "warning", + "--port", + str(free_port), + ] + else: + command += ["--port", str(port), "--loglevel", "warning"] + + if listen_to_localhost_only: + command += ["--bind", "127.0.0.1"] + pidfile = os.path.join(session_dir_path, "redis-" + uuid.uuid4().hex + ".pid") + command += ["--pidfile", pidfile] + if enable_tls: + if Config.REDIS_CA_CERT(): + command += ["--tls-ca-cert-file", Config.REDIS_CA_CERT()] + if Config.REDIS_CLIENT_CERT(): + command += ["--tls-cert-file", Config.REDIS_CLIENT_CERT()] + if Config.REDIS_CLIENT_KEY(): + command += ["--tls-key-file", Config.REDIS_CLIENT_KEY()] + if replica_of is not None: + command += ["--tls-replication", "yes"] + command += ["--tls-auth-clients", "no", "--tls-cluster", "yes"] + if sys.platform != "win32": + command += ["--save", "", "--appendonly", "no"] + if db_dir is not None: + command += ["--dir", str(db_dir)] + + process_info = ray._private.services.start_ray_process( + command, + ray_constants.PROCESS_TYPE_REDIS_SERVER, + stdout_file=stdout_file, + stderr_file=stderr_file, + fate_share=fate_share, + ) + node_id = None + if redis_replicas() > 1: + # Setup redis cluster + import redis + + while True: + try: + redis_cli = get_redis_cli(port, enable_tls) + if replica_of is None: + slots = [str(i) for i in range(16384)] + redis_cli.cluster("addslots", *slots) + else: + logger.info(redis_cli.cluster("meet", "127.0.0.1", str(replica_of))) + logger.info(redis_cli.cluster("replicate", leader_id)) + node_id = redis_cli.cluster("myid") + break + except ( + redis.exceptions.ConnectionError, + redis.exceptions.ResponseError, + ) as e: + from time import sleep + + logger.info( + f"Waiting for redis to be up. Check failed with error: {e}. " + "Will retry in 0.1s" + ) + + if process_info.process.poll() is not None: + raise Exception( + f"Redis process exited unexpectedly: {process_info}. " + f"Exit code: {process_info.process.returncode}" + ) + + sleep(0.1) + + logger.info( + f"Redis started with node_id {node_id} and pid {process_info.process.pid}" + ) + + return node_id, process_info + + +def _pid_alive(pid): + """Check if the process with this PID is alive or not. + + Args: + pid: The pid to check. + + Returns: + This returns false if the process is dead. Otherwise, it returns true. + """ + alive = True + try: + proc = psutil.Process(pid) + if proc.status() == psutil.STATUS_ZOMBIE: + alive = False + except psutil.NoSuchProcess: + alive = False + return alive + + +def _check_call_windows(main, argv, capture_stdout=False, capture_stderr=False): + # We use this function instead of calling the "ray" command to work around + # some deadlocks that occur when piping ray's output on Windows + stream = io.TextIOWrapper(io.BytesIO(), encoding=sys.stdout.encoding) + old_argv = sys.argv[:] + try: + sys.argv = argv[:] + try: + with redirect_stderr(stream if capture_stderr else sys.stderr): + with redirect_stdout(stream if capture_stdout else sys.stdout): + main() + finally: + stream.flush() + except SystemExit as ex: + if ex.code: + output = stream.buffer.getvalue() + raise subprocess.CalledProcessError(ex.code, argv, output) + except Exception as ex: + output = stream.buffer.getvalue() + raise subprocess.CalledProcessError(1, argv, output, ex.args[0]) + finally: + sys.argv = old_argv + if capture_stdout: + sys.stdout.buffer.write(stream.buffer.getvalue()) + elif capture_stderr: + sys.stderr.buffer.write(stream.buffer.getvalue()) + return stream.buffer.getvalue() + + +def check_call_subprocess(argv, capture_stdout=False, capture_stderr=False): + # We use this function instead of calling the "ray" command to work around + # some deadlocks that occur when piping ray's output on Windows + from ray.scripts.scripts import main as ray_main + + if sys.platform == "win32": + result = _check_call_windows( + ray_main, argv, capture_stdout=capture_stdout, capture_stderr=capture_stderr + ) + else: + stdout_redir = None + stderr_redir = None + if capture_stdout: + stdout_redir = subprocess.PIPE + if capture_stderr and capture_stdout: + stderr_redir = subprocess.STDOUT + elif capture_stderr: + stderr_redir = subprocess.PIPE + proc = subprocess.Popen(argv, stdout=stdout_redir, stderr=stderr_redir) + (stdout, stderr) = proc.communicate() + if proc.returncode: + raise subprocess.CalledProcessError(proc.returncode, argv, stdout, stderr) + result = b"".join([s for s in [stdout, stderr] if s is not None]) + return result + + +def check_call_ray(args, capture_stdout=False, capture_stderr=False): + check_call_subprocess(["ray"] + args, capture_stdout, capture_stderr) + + +def wait_for_dashboard_agent_available(cluster): + gcs_client = GcsClient(address=cluster.address) + + def get_dashboard_agent_address(): + return gcs_client.internal_kv_get( + f"{dashboard_consts.DASHBOARD_AGENT_ADDR_NODE_ID_PREFIX}{cluster.head_node.node_id}".encode(), + namespace=ray_constants.KV_NAMESPACE_DASHBOARD, + timeout=dashboard_consts.GCS_RPC_TIMEOUT_SECONDS, + ) + + wait_for_condition(lambda: get_dashboard_agent_address() is not None) + + +def wait_for_pid_to_exit(pid: int, timeout: float = 20): + start_time = time.time() + while time.time() - start_time < timeout: + if not _pid_alive(pid): + return + time.sleep(0.1) + raise TimeoutError(f"Timed out while waiting for process {pid} to exit.") + + +def wait_for_children_of_pid(pid, num_children=1, timeout=20): + p = psutil.Process(pid) + start_time = time.time() + alive = [] + while time.time() - start_time < timeout: + alive = p.children(recursive=False) + num_alive = len(alive) + if num_alive >= num_children: + return + time.sleep(0.1) + raise TimeoutError( + f"Timed out while waiting for process {pid} children to start " + f"({num_alive}/{num_children} started: {alive})." + ) + + +def wait_for_children_of_pid_to_exit(pid, timeout=20): + children = psutil.Process(pid).children() + if len(children) == 0: + return + + _, alive = psutil.wait_procs(children, timeout=timeout) + if len(alive) > 0: + raise TimeoutError( + "Timed out while waiting for process children to exit." + " Children still alive: {}.".format([p.name() for p in alive]) + ) + + +def kill_process_by_name(name, SIGKILL=False): + for p in psutil.process_iter(attrs=["name"]): + if p.info["name"] == name + ray._private.services.EXE_SUFFIX: + if SIGKILL: + p.kill() + else: + p.terminate() + + +def run_string_as_driver(driver_script: str, env: Dict = None, encode: str = "utf-8"): + """Run a driver as a separate process. + + Args: + driver_script: A string to run as a Python script. + env: The environment variables for the driver. + + Returns: + The script's output. + """ + proc = subprocess.Popen( + [sys.executable, "-"], + stdin=subprocess.PIPE, + stdout=subprocess.PIPE, + stderr=subprocess.STDOUT, + env=env, + ) + with proc: + output = proc.communicate(driver_script.encode(encoding=encode))[0] + if proc.returncode: + print(ray._common.utils.decode(output, encode_type=encode)) + logger.error(proc.stderr) + raise subprocess.CalledProcessError( + proc.returncode, proc.args, output, proc.stderr + ) + out = ray._common.utils.decode(output, encode_type=encode) + return out + + +def run_string_as_driver_stdout_stderr( + driver_script: str, env: Dict = None, encode: str = "utf-8" +) -> Tuple[str, str]: + """Run a driver as a separate process. + + Args: + driver_script: A string to run as a Python script. + env: The environment variables for the driver. + + Returns: + The script's stdout and stderr. + """ + proc = subprocess.Popen( + [sys.executable, "-"], + stdin=subprocess.PIPE, + stdout=subprocess.PIPE, + stderr=subprocess.PIPE, + env=env, + ) + with proc: + outputs_bytes = proc.communicate(driver_script.encode(encoding=encode)) + out_str, err_str = [ + ray._common.utils.decode(output, encode_type=encode) + for output in outputs_bytes + ] + if proc.returncode: + print(out_str) + print(err_str) + raise subprocess.CalledProcessError( + proc.returncode, proc.args, out_str, err_str + ) + return out_str, err_str + + +def run_string_as_driver_nonblocking(driver_script, env: Dict = None): + """Start a driver as a separate process and return immediately. + + Args: + driver_script: A string to run as a Python script. + + Returns: + A handle to the driver process. + """ + script = "; ".join( + [ + "import sys", + "script = sys.stdin.read()", + "sys.stdin.close()", + "del sys", + 'exec("del script\\n" + script)', + ] + ) + proc = subprocess.Popen( + [sys.executable, "-c", script], + stdin=subprocess.PIPE, + stdout=subprocess.PIPE, + stderr=subprocess.PIPE, + env=env, + ) + proc.stdin.write(driver_script.encode("ascii")) + proc.stdin.close() + return proc + + +def convert_actor_state(state): + if not state: + return None + return gcs_pb2.ActorTableData.ActorState.DESCRIPTOR.values_by_number[state].name + + +def wait_for_num_actors(num_actors, state=None, timeout=10): + state = convert_actor_state(state) + start_time = time.time() + while time.time() - start_time < timeout: + if ( + len( + list_actors( + filters=[("state", "=", state)] if state else None, + limit=num_actors, + ) + ) + >= num_actors + ): + return + time.sleep(0.1) + raise TimeoutError("Timed out while waiting for global state.") + + +def kill_actor_and_wait_for_failure(actor, timeout=10, retry_interval_ms=100): + actor_id = actor._actor_id.hex() + current_num_restarts = get_actor(id=actor_id).num_restarts + ray.kill(actor) + start = time.time() + while time.time() - start <= timeout: + actor_state = get_actor(id=actor_id) + if ( + actor_state.state == "DEAD" + or actor_state.num_restarts > current_num_restarts + ): + return + time.sleep(retry_interval_ms / 1000.0) + raise RuntimeError("It took too much time to kill an actor: {}".format(actor_id)) + + +def wait_for_assertion( + assertion_predictor: Callable, + timeout: int = 10, + retry_interval_ms: int = 100, + raise_exceptions: bool = False, + **kwargs: Any, +): + """Wait until an assertion is met or time out with an exception. + + Args: + assertion_predictor: A function that predicts the assertion. + timeout: Maximum timeout in seconds. + retry_interval_ms: Retry interval in milliseconds. + raise_exceptions: If true, exceptions that occur while executing + assertion_predictor won't be caught and instead will be raised. + **kwargs: Arguments to pass to the condition_predictor. + + Raises: + RuntimeError: If the assertion is not met before the timeout expires. + """ + + def _assertion_to_condition(): + try: + assertion_predictor(**kwargs) + return True + except AssertionError: + return False + + try: + wait_for_condition( + _assertion_to_condition, + timeout=timeout, + retry_interval_ms=retry_interval_ms, + raise_exceptions=raise_exceptions, + **kwargs, + ) + except RuntimeError: + assertion_predictor(**kwargs) # Should fail assert + + +@dataclass +class MetricSamplePattern: + name: Optional[str] = None + value: Optional[str] = None + partial_label_match: Optional[Dict[str, str]] = None + + def matches(self, sample: Sample): + if self.name is not None: + if self.name != sample.name: + return False + + if self.value is not None: + if self.value != sample.value: + return False + + if self.partial_label_match is not None: + for label, value in self.partial_label_match.items(): + if sample.labels.get(label) != value: + return False + + return True + + +@dataclass +class PrometheusTimeseries: + """A collection of timeseries from multiple addresses. Each timeseries is a + collection of samples with the same metric name and labels. Concretely: + - components_dict: a dictionary of addresses to the Component labels + - metric_descriptors: a dictionary of metric names to the Metric object + - metric_samples: the latest value of each label + """ + + components_dict: Dict[str, Set[str]] = field(default_factory=dict) + metric_descriptors: Dict[str, Metric] = field(default_factory=dict) + metric_samples: Dict[frozenset, Sample] = field(default_factory=dict) + + def flush(self): + self.components_dict.clear() + self.metric_descriptors.clear() + self.metric_samples.clear() + + +def get_metric_check_condition( + metrics_to_check: List[MetricSamplePattern], + timeseries: PrometheusTimeseries, + export_addr: Optional[str] = None, +) -> Callable[[], bool]: + """A condition to check if a prometheus metrics reach a certain value. + + This is a blocking check that can be passed into a `wait_for_condition` + style function. + + Args: + metrics_to_check: A list of MetricSamplePattern. The fields that + aren't `None` will be matched. + timeseries: A PrometheusTimeseries object to store the metrics. + export_addr: Optional address to export metrics to. + + Returns: + A function that returns True if all the metrics are emitted. + """ + node_info = ray.nodes()[0] + metrics_export_port = node_info["MetricsExportPort"] + addr = node_info["NodeManagerAddress"] + prom_addr = export_addr or build_address(addr, metrics_export_port) + + def f(): + for metric_pattern in metrics_to_check: + metric_samples = fetch_prometheus_timeseries( + [prom_addr], timeseries + ).metric_samples.values() + for metric_sample in metric_samples: + if metric_pattern.matches(metric_sample): + break + else: + print( + f"Didn't find {metric_pattern}", + "all samples", + metric_samples, + ) + return False + return True + + return f + + +def wait_until_succeeded_without_exception( + func, exceptions, *args, timeout_ms=1000, retry_interval_ms=100, raise_last_ex=False +): + """A helper function that waits until a given function + completes without exceptions. + + Args: + func: A function to run. + exceptions: Exceptions that are supposed to occur. + args: arguments to pass for a given func + timeout_ms: Maximum timeout in milliseconds. + retry_interval_ms: Retry interval in milliseconds. + raise_last_ex: Raise the last exception when timeout. + + Return: + Whether exception occurs within a timeout. + """ + if isinstance(type(exceptions), tuple): + raise Exception("exceptions arguments should be given as a tuple") + + time_elapsed = 0 + start = time.time() + last_ex = None + while time_elapsed <= timeout_ms: + try: + func(*args) + return True + except exceptions as ex: + last_ex = ex + time_elapsed = (time.time() - start) * 1000 + time.sleep(retry_interval_ms / 1000.0) + if raise_last_ex: + ex_stack = ( + traceback.format_exception(type(last_ex), last_ex, last_ex.__traceback__) + if last_ex + else [] + ) + ex_stack = "".join(ex_stack) + raise Exception(f"Timed out while testing, {ex_stack}") + return False + + +def recursive_fnmatch(dirpath, pattern): + """Looks at a file directory subtree for a filename pattern. + + Similar to glob.glob(..., recursive=True) but also supports 2.7 + """ + matches = [] + for root, dirnames, filenames in os.walk(dirpath): + for filename in fnmatch.filter(filenames, pattern): + matches.append(os.path.join(root, filename)) + return matches + + +def generate_system_config_map(**kwargs): + ray_kwargs = { + "_system_config": kwargs, + } + return ray_kwargs + + +def same_elements(elems_a, elems_b): + """Checks if two iterables (such as lists) contain the same elements. Elements + do not have to be hashable (this allows us to compare sets of dicts for + example). This comparison is not necessarily efficient. + """ + a = list(elems_a) + b = list(elems_b) + + for x in a: + if x not in b: + return False + + for x in b: + if x not in a: + return False + + return True + + +@ray.remote +def _put(obj): + return obj + + +def put_object(obj, use_ray_put): + if use_ray_put: + return ray.put(obj) + else: + return _put.remote(obj) + + +def wait_until_server_available(address, timeout_ms=5000, retry_interval_ms=100): + ip, port_str = parse_address(address) + port = int(port_str) + time_elapsed = 0 + start = time.time() + while time_elapsed <= timeout_ms: + s = socket.socket(socket.AF_INET, socket.SOCK_STREAM) + s.settimeout(1) + try: + s.connect((ip, port)) + except Exception: + time_elapsed = (time.time() - start) * 1000 + time.sleep(retry_interval_ms / 1000.0) + s.close() + continue + s.close() + return True + return False + + +def get_other_nodes(cluster, exclude_head=False): + """Get all nodes except the one that we're connected to.""" + return [ + node + for node in cluster.list_all_nodes() + if node._raylet_socket_name + != ray._private.worker._global_node._raylet_socket_name + and (exclude_head is False or node.head is False) + ] + + +def get_non_head_nodes(cluster): + """Get all non-head nodes.""" + return list(filter(lambda x: x.head is False, cluster.list_all_nodes())) + + +def init_error_pubsub(): + """Initialize error info pub/sub""" + s = ray._raylet.GcsErrorSubscriber( + address=ray._private.worker.global_worker.gcs_client.address + ) + s.subscribe() + return s + + +def get_error_message(subscriber, num=1e6, error_type=None, timeout=20): + """Gets errors from GCS subscriber. + + Returns maximum `num` error strings within `timeout`. + Only returns errors of `error_type` if specified. + """ + deadline = time.time() + timeout + msgs = [] + while time.time() < deadline and len(msgs) < num: + _, error_data = subscriber.poll(timeout=deadline - time.time()) + if not error_data: + # Timed out before any data is received. + break + if error_type is None or error_type == error_data["type"]: + msgs.append(error_data) + else: + time.sleep(0.01) + + return msgs + + +def init_log_pubsub(): + """Initialize log pub/sub""" + s = ray._raylet.GcsLogSubscriber( + address=ray._private.worker.global_worker.gcs_client.address + ) + s.subscribe() + return s + + +def get_log_data( + subscriber, + num: int = 1e6, + timeout: float = 20, + job_id: Optional[str] = None, + matcher=None, +) -> List[dict]: + deadline = time.time() + timeout + msgs = [] + while time.time() < deadline and len(msgs) < num: + logs_data = subscriber.poll(timeout=deadline - time.time()) + if not logs_data: + # Timed out before any data is received. + break + if job_id and job_id != logs_data["job"]: + continue + if matcher and all(not matcher(line) for line in logs_data["lines"]): + continue + msgs.append(logs_data) + return msgs + + +def get_log_message( + subscriber, + num: int = 1e6, + timeout: float = 20, + job_id: Optional[str] = None, + matcher=None, +) -> List[List[str]]: + """Gets log lines through GCS subscriber. + + Returns maximum `num` of log messages, within `timeout`. + + If `job_id` or `match` is specified, only returns log lines from `job_id` + or when `matcher` is true. + """ + msgs = get_log_data(subscriber, num, timeout, job_id, matcher) + return [msg["lines"] for msg in msgs] + + +def get_log_sources( + subscriber, + num: int = 1e6, + timeout: float = 20, + job_id: Optional[str] = None, + matcher=None, +): + """Get the source of all log messages""" + msgs = get_log_data(subscriber, num, timeout, job_id, matcher) + return {msg["pid"] for msg in msgs} + + +def get_log_batch( + subscriber, + num: int, + timeout: float = 20, + job_id: Optional[str] = None, + matcher=None, +) -> List[str]: + """Gets log batches through GCS subscriber. + + Returns maximum `num` batches of logs. Each batch is a dict that includes + metadata such as `pid`, `job_id`, and `lines` of log messages. + + If `job_id` or `match` is specified, only returns log batches from `job_id` + or when `matcher` is true. + """ + deadline = time.time() + timeout + batches = [] + while time.time() < deadline and len(batches) < num: + logs_data = subscriber.poll(timeout=deadline - time.time()) + if not logs_data: + # Timed out before any data is received. + break + if job_id and job_id != logs_data["job"]: + continue + if matcher and not matcher(logs_data): + continue + batches.append(logs_data) + + return batches + + +def format_web_url(url): + """Format web url.""" + url = url.replace("localhost", "http://127.0.0.1") + if not url.startswith("http://"): + return "http://" + url + return url + + +def client_test_enabled() -> bool: + return ray._private.client_mode_hook.is_client_mode_enabled + + +def object_memory_usage() -> bool: + """Returns the number of bytes used in the object store.""" + total = ray.cluster_resources().get("object_store_memory", 0) + avail = ray.available_resources().get("object_store_memory", 0) + return total - avail + + +def fetch_raw_prometheus(prom_addresses): + # Local import so minimal dependency tests can run without requests + import requests + + for address in prom_addresses: + try: + response = requests.get(f"http://{address}/metrics") + yield address, response.text + except requests.exceptions.ConnectionError: + continue + + +def fetch_prometheus(prom_addresses): + components_dict = {} + metric_descriptors = {} + metric_samples = [] + + for address in prom_addresses: + if address not in components_dict: + components_dict[address] = set() + + for address, response in fetch_raw_prometheus(prom_addresses): + for metric in text_string_to_metric_families(response): + for sample in metric.samples: + metric_descriptors[sample.name] = metric + metric_samples.append(sample) + if "Component" in sample.labels: + components_dict[address].add(sample.labels["Component"]) + return components_dict, metric_descriptors, metric_samples + + +def fetch_prometheus_timeseries( + prom_addreses: List[str], + result: PrometheusTimeseries, +) -> PrometheusTimeseries: + components_dict, metric_descriptors, metric_samples = fetch_prometheus( + prom_addreses + ) + for address, components in components_dict.items(): + if address not in result.components_dict: + result.components_dict[address] = set() + result.components_dict[address].update(components) + result.metric_descriptors.update(metric_descriptors) + for sample in metric_samples: + # udpate sample to the latest value + result.metric_samples[ + frozenset(list(sample.labels.items()) + [("_metric_name_", sample.name)]) + ] = sample + return result + + +def fetch_prometheus_metrics(prom_addresses: List[str]) -> Dict[str, List[Any]]: + """Return prometheus metrics from the given addresses. + + Args: + prom_addresses: List of metrics_agent addresses to collect metrics from. + + Returns: + Dict mapping from metric name to list of samples for the metric. + """ + _, _, samples = fetch_prometheus(prom_addresses) + samples_by_name = defaultdict(list) + for sample in samples: + samples_by_name[sample.name].append(sample) + return samples_by_name + + +def fetch_prometheus_metric_timeseries( + prom_addresses: List[str], result: PrometheusTimeseries +) -> Dict[str, List[Any]]: + samples = fetch_prometheus_timeseries( + prom_addresses, result + ).metric_samples.values() + samples_by_name = defaultdict(list) + for sample in samples: + samples_by_name[sample.name].append(sample) + return samples_by_name + + +def raw_metric_timeseries( + info: RayContext, result: PrometheusTimeseries +) -> Dict[str, List[Any]]: + """Return prometheus timeseries from a RayContext""" + metrics_page = "localhost:{}".format(info.address_info["metrics_export_port"]) + print("Fetch metrics from", metrics_page) + return fetch_prometheus_metric_timeseries([metrics_page], result) + + +def get_system_metric_for_component( + system_metric: str, component: str, prometheus_server_address: str +) -> List[float]: + """Get the system metric for a given component from a Prometheus server address. + Please note: + - This function requires the availability of the Prometheus server. Therefore, it + requires the server address. + - It assumes the system metric has a `Component` label and `pid` label. `pid` is the + process id, so it can be used to uniquely identify the process. + """ + session_name = os.path.basename( + ray._private.worker._global_node.get_session_dir_path() + ) + query = f"sum({system_metric}{{Component='{component}',SessionName='{session_name}'}}) by (pid)" + resp = requests.get( + f"{prometheus_server_address}/api/v1/query?query={quote(query)}" + ) + if resp.status_code != 200: + raise Exception(f"Failed to query Prometheus: {resp.status_code}") + result = resp.json() + return [float(item["value"][1]) for item in result["data"]["result"]] + + +def get_test_config_path(config_file_name): + """Resolve the test config path from the config file dir""" + here = os.path.realpath(__file__) + path = pathlib.Path(here) + grandparent = path.parent.parent + return os.path.join(grandparent, "tests/test_cli_patterns", config_file_name) + + +def load_test_config(config_file_name): + """Loads a config yaml from tests/test_cli_patterns.""" + config_path = get_test_config_path(config_file_name) + config = yaml.safe_load(open(config_path).read()) + return config + + +def set_setup_func(): + import ray._private.runtime_env as runtime_env + + runtime_env.VAR = "hello world" + + +class BatchQueue(Queue): + def __init__(self, maxsize: int = 0, actor_options: Optional[Dict] = None) -> None: + actor_options = actor_options or {} + self.maxsize = maxsize + self.actor = ( + ray.remote(_BatchQueueActor).options(**actor_options).remote(self.maxsize) + ) + + def get_batch( + self, + batch_size: int = None, + total_timeout: Optional[float] = None, + first_timeout: Optional[float] = None, + ) -> List[Any]: + """Gets batch of items from the queue and returns them in a + list in order. + + Raises: + Empty: if the queue does not contain the desired number of items + """ + return ray.get( + self.actor.get_batch.remote(batch_size, total_timeout, first_timeout) + ) + + +class _BatchQueueActor(_QueueActor): + async def get_batch(self, batch_size=None, total_timeout=None, first_timeout=None): + start = timeit.default_timer() + try: + first = await asyncio.wait_for(self.queue.get(), first_timeout) + batch = [first] + if total_timeout: + end = timeit.default_timer() + total_timeout = max(total_timeout - (end - start), 0) + except asyncio.TimeoutError: + raise Empty + if batch_size is None: + if total_timeout is None: + total_timeout = 0 + while True: + try: + start = timeit.default_timer() + batch.append( + await asyncio.wait_for(self.queue.get(), total_timeout) + ) + if total_timeout: + end = timeit.default_timer() + total_timeout = max(total_timeout - (end - start), 0) + except asyncio.TimeoutError: + break + else: + for _ in range(batch_size - 1): + try: + start = timeit.default_timer() + batch.append( + await asyncio.wait_for(self.queue.get(), total_timeout) + ) + if total_timeout: + end = timeit.default_timer() + total_timeout = max(total_timeout - (end - start), 0) + except asyncio.TimeoutError: + break + return batch + + +def is_placement_group_removed(pg): + table = ray.util.placement_group_table(pg) + if "state" not in table: + return False + return table["state"] == "REMOVED" + + +def placement_group_assert_no_leak(pgs_created): + for pg in pgs_created: + ray.util.remove_placement_group(pg) + + def wait_for_pg_removed(): + for pg_entry in ray.util.placement_group_table().values(): + if pg_entry["state"] != "REMOVED": + return False + return True + + wait_for_condition(wait_for_pg_removed) + + cluster_resources = ray.cluster_resources() + cluster_resources.pop("memory") + cluster_resources.pop("object_store_memory") + + def wait_for_resource_recovered(): + for resource, val in ray.available_resources().items(): + if resource in cluster_resources and cluster_resources[resource] != val: + return False + if "_group_" in resource: + return False + return True + + wait_for_condition(wait_for_resource_recovered) + + +def monitor_memory_usage( + print_interval_s: int = 30, + record_interval_s: int = 5, + warning_threshold: float = 0.9, +): + """Run the memory monitor actor that prints the memory usage. + + The monitor will run on the same node as this function is called. + + Params: + interval_s: The interval memory usage information is printed + warning_threshold: The threshold where the + memory usage warning is printed. + + Returns: + The memory monitor actor. + """ + assert ray.is_initialized(), "The API is only available when Ray is initialized." + + @ray.remote(num_cpus=0) + class MemoryMonitorActor: + def __init__( + self, + print_interval_s: float = 20, + record_interval_s: float = 5, + warning_threshold: float = 0.9, + n: int = 10, + ): + """The actor that monitor the memory usage of the cluster. + + Params: + print_interval_s: The interval where + memory usage is printed. + record_interval_s: The interval where + memory usage is recorded. + warning_threshold: The threshold where + memory warning is printed + n: When memory usage is printed, + top n entries are printed. + """ + # -- Interval the monitor prints the memory usage information. -- + self.print_interval_s = print_interval_s + # -- Interval the monitor records the memory usage information. -- + self.record_interval_s = record_interval_s + # -- Whether or not the monitor is running. -- + self.is_running = False + # -- The used_gb/total_gb threshold where warning message omits. -- + self.warning_threshold = warning_threshold + # -- The monitor that calculates the memory usage of the node. -- + self.monitor = memory_monitor.MemoryMonitor() + # -- The top n memory usage of processes are printed. -- + self.n = n + # -- The peak memory usage in GB during lifetime of monitor. -- + self.peak_memory_usage = 0 + # -- The top n memory usage of processes + # during peak memory usage. -- + self.peak_top_n_memory_usage = "" + # -- The last time memory usage was printed -- + self._last_print_time = 0 + # -- logger. -- + logging.basicConfig(level=logging.INFO) + + def ready(self): + pass + + async def run(self): + """Run the monitor.""" + self.is_running = True + while self.is_running: + now = time.time() + used_gb, total_gb = self.monitor.get_memory_usage() + top_n_memory_usage = memory_monitor.get_top_n_memory_usage(n=self.n) + if used_gb > self.peak_memory_usage: + self.peak_memory_usage = used_gb + self.peak_top_n_memory_usage = top_n_memory_usage + + if used_gb > total_gb * self.warning_threshold: + logging.warning( + "The memory usage is high: " f"{used_gb / total_gb * 100}%" + ) + if now - self._last_print_time > self.print_interval_s: + logging.info(f"Memory usage: {used_gb} / {total_gb}") + logging.info(f"Top {self.n} process memory usage:") + logging.info(top_n_memory_usage) + self._last_print_time = now + await asyncio.sleep(self.record_interval_s) + + async def stop_run(self): + """Stop running the monitor. + + Returns: + True if the monitor is stopped. False otherwise. + """ + was_running = self.is_running + self.is_running = False + return was_running + + async def get_peak_memory_info(self): + """Return the tuple of the peak memory usage and the + top n process information during the peak memory usage. + """ + return self.peak_memory_usage, self.peak_top_n_memory_usage + + current_node_ip = ray._private.worker.global_worker.node_ip_address + # Schedule the actor on the current node. + memory_monitor_actor = MemoryMonitorActor.options( + resources={f"node:{current_node_ip}": 0.001} + ).remote( + print_interval_s=print_interval_s, + record_interval_s=record_interval_s, + warning_threshold=warning_threshold, + ) + print("Waiting for memory monitor actor to be ready...") + ray.get(memory_monitor_actor.ready.remote()) + print("Memory monitor actor is ready now.") + memory_monitor_actor.run.remote() + return memory_monitor_actor + + +def setup_tls(): + """Sets up required environment variables for tls""" + import pytest + + if sys.platform == "darwin": + pytest.skip("Cryptography doesn't install in Mac build pipeline") + cert, key = generate_self_signed_tls_certs() + temp_dir = tempfile.mkdtemp("ray-test-certs") + cert_filepath = os.path.join(temp_dir, "server.crt") + key_filepath = os.path.join(temp_dir, "server.key") + with open(cert_filepath, "w") as fh: + fh.write(cert) + with open(key_filepath, "w") as fh: + fh.write(key) + + os.environ["RAY_USE_TLS"] = "1" + os.environ["RAY_TLS_SERVER_CERT"] = cert_filepath + os.environ["RAY_TLS_SERVER_KEY"] = key_filepath + os.environ["RAY_TLS_CA_CERT"] = cert_filepath + + return key_filepath, cert_filepath, temp_dir + + +def teardown_tls(key_filepath, cert_filepath, temp_dir): + os.remove(key_filepath) + os.remove(cert_filepath) + os.removedirs(temp_dir) + del os.environ["RAY_USE_TLS"] + del os.environ["RAY_TLS_SERVER_CERT"] + del os.environ["RAY_TLS_SERVER_KEY"] + del os.environ["RAY_TLS_CA_CERT"] + + +class ResourceKillerActor: + """Abstract base class used to implement resource killers for chaos testing. + + Subclasses should implement _find_resource_to_kill, which should find a resource + to kill. This method should return the args to _kill_resource, which is another + abstract method that should kill the resource and add it to the `killed` set. + """ + + def __init__( + self, + head_node_id, + kill_interval_s: float = 60, + kill_delay_s: float = 0, + max_to_kill: Optional[int] = 2, + batch_size_to_kill: int = 1, + kill_filter_fn: Optional[Callable] = None, + ): + self.kill_interval_s = kill_interval_s + self.kill_delay_s = kill_delay_s + self.is_running = False + self.head_node_id = head_node_id + self.killed = set() + self.done = get_or_create_event_loop().create_future() + self.max_to_kill = max_to_kill + self.batch_size_to_kill = batch_size_to_kill + self.kill_filter_fn = kill_filter_fn + self.kill_immediately_after_found = False + # -- logger. -- + logging.basicConfig(level=logging.INFO) + + def ready(self): + pass + + async def run(self): + self.is_running = True + + time.sleep(self.kill_delay_s) + + while self.is_running: + to_kills = await self._find_resources_to_kill() + + if not self.is_running: + break + + if self.kill_immediately_after_found: + sleep_interval = 0 + else: + sleep_interval = random.random() * self.kill_interval_s + time.sleep(sleep_interval) + + for to_kill in to_kills: + self._kill_resource(*to_kill) + if self.max_to_kill is not None and len(self.killed) >= self.max_to_kill: + break + await asyncio.sleep(self.kill_interval_s - sleep_interval) + + self.done.set_result(True) + await self.stop_run() + + async def _find_resources_to_kill(self): + raise NotImplementedError + + def _kill_resource(self, *args): + raise NotImplementedError + + async def stop_run(self): + was_running = self.is_running + if was_running: + self._cleanup() + + self.is_running = False + return was_running + + async def get_total_killed(self): + """Get the total number of killed resources""" + await self.done + return self.killed + + def _cleanup(self): + """Cleanup any resources created by the killer. + + Overriding this method is optional. + """ + pass + + +class NodeKillerBase(ResourceKillerActor): + async def _find_resources_to_kill(self): + nodes_to_kill = [] + while not nodes_to_kill and self.is_running: + worker_nodes = [ + node + for node in ray.nodes() + if node["Alive"] + and (node["NodeID"] != self.head_node_id) + and (node["NodeID"] not in self.killed) + ] + if self.kill_filter_fn: + candidates = list(filter(self.kill_filter_fn(), worker_nodes)) + else: + candidates = worker_nodes + + # Ensure at least one worker node remains alive. + if len(worker_nodes) < self.batch_size_to_kill + 1: + # Give the cluster some time to start. + await asyncio.sleep(1) + continue + + # Collect nodes to kill, limited by batch size. + for candidate in candidates[: self.batch_size_to_kill]: + nodes_to_kill.append( + ( + candidate["NodeID"], + candidate["NodeManagerAddress"], + candidate["NodeManagerPort"], + ) + ) + + return nodes_to_kill + + +@ray.remote(num_cpus=0) +class RayletKiller(NodeKillerBase): + def _kill_resource(self, node_id, node_to_kill_ip, node_to_kill_port): + if node_to_kill_port is not None: + try: + self._kill_raylet(node_to_kill_ip, node_to_kill_port, graceful=False) + except Exception: + pass + logging.info( + f"Killed node {node_id} at address: " + f"{node_to_kill_ip}, port: {node_to_kill_port}" + ) + self.killed.add(node_id) + + def _kill_raylet(self, ip, port, graceful=False): + import grpc + from grpc._channel import _InactiveRpcError + + from ray.core.generated import node_manager_pb2_grpc + + raylet_address = build_address(ip, port) + channel = grpc.insecure_channel(raylet_address) + stub = node_manager_pb2_grpc.NodeManagerServiceStub(channel) + try: + stub.ShutdownRaylet( + node_manager_pb2.ShutdownRayletRequest(graceful=graceful) + ) + except _InactiveRpcError: + assert not graceful + + +@ray.remote(num_cpus=0) +class EC2InstanceTerminator(NodeKillerBase): + def _kill_resource(self, node_id, node_to_kill_ip, _): + if node_to_kill_ip is not None: + try: + _terminate_ec2_instance(node_to_kill_ip) + except Exception: + pass + logging.info(f"Terminated instance, {node_id=}, address={node_to_kill_ip}") + self.killed.add(node_id) + + +@ray.remote(num_cpus=0) +class EC2InstanceTerminatorWithGracePeriod(NodeKillerBase): + def __init__(self, *args, grace_period_s: int = 30, **kwargs): + super().__init__(*args, **kwargs) + + self._grace_period_s = grace_period_s + self._kill_threads: Set[threading.Thread] = set() + + def _kill_resource(self, node_id, node_to_kill_ip, _): + assert node_id not in self.killed + + # Clean up any completed threads. + for thread in self._kill_threads.copy(): + if not thread.is_alive(): + thread.join() + self._kill_threads.remove(thread) + + def _kill_node_with_grace_period(node_id, node_to_kill_ip): + self._drain_node(node_id) + time.sleep(self._grace_period_s) + # Anyscale extends the drain deadline if you shut down the instance + # directly. To work around this, we force-stop Ray on the node. Anyscale + # should then terminate it shortly after without updating the drain + # deadline. + _execute_command_on_node("ray stop --force", node_to_kill_ip) + + logger.info(f"Starting killing thread {node_id=}, {node_to_kill_ip=}") + thread = threading.Thread( + target=_kill_node_with_grace_period, + args=(node_id, node_to_kill_ip), + daemon=True, + ) + thread.start() + self._kill_threads.add(thread) + self.killed.add(node_id) + + def _drain_node(self, node_id: str) -> None: + # We need to lazily import this object. Otherwise, Ray can't serialize the + # class. + from ray.core.generated import autoscaler_pb2 + + assert ray.NodeID.from_hex(node_id) != ray.NodeID.nil() + + logging.info(f"Draining node {node_id=}") + address = services.canonicalize_bootstrap_address_or_die(addr="auto") + gcs_client = ray._raylet.GcsClient(address=address) + deadline_timestamp_ms = (time.time_ns() // 1e6) + (self._grace_period_s * 1e3) + + try: + is_accepted, _ = gcs_client.drain_node( + node_id, + autoscaler_pb2.DrainNodeReason.Value("DRAIN_NODE_REASON_PREEMPTION"), + "", + deadline_timestamp_ms, + ) + except ray.exceptions.RayError as e: + logger.error(f"Failed to drain node {node_id=}") + raise e + + assert is_accepted, "Drain node request was rejected" + + def _cleanup(self): + for thread in self._kill_threads.copy(): + thread.join() + self._kill_threads.remove(thread) + + assert not self._kill_threads + + +@ray.remote(num_cpus=0) +class WorkerKillerActor(ResourceKillerActor): + def __init__(self, *args, **kwargs): + super().__init__(*args, **kwargs) + + # Kill worker immediately so that the task does + # not finish successfully on its own. + self.kill_immediately_after_found = True + + from ray.util.state.api import StateApiClient + from ray.util.state.common import ListApiOptions + + self.client = StateApiClient() + self.task_options = ListApiOptions( + filters=[ + ("state", "=", "RUNNING"), + ("name", "!=", "WorkerKillActor.run"), + ] + ) + + async def _find_resources_to_kill(self): + from ray.util.state.common import StateResource + + process_to_kill_task_id = None + process_to_kill_pid = None + process_to_kill_node_id = None + while process_to_kill_pid is None and self.is_running: + tasks = self.client.list( + StateResource.TASKS, + options=self.task_options, + raise_on_missing_output=False, + ) + if self.kill_filter_fn is not None: + tasks = list(filter(self.kill_filter_fn(), tasks)) + + for task in tasks: + if task.worker_id is not None and task.node_id is not None: + process_to_kill_task_id = task.task_id + process_to_kill_pid = task.worker_pid + process_to_kill_node_id = task.node_id + break + + # Give the cluster some time to start. + await asyncio.sleep(0.1) + + return [(process_to_kill_task_id, process_to_kill_pid, process_to_kill_node_id)] + + def _kill_resource( + self, process_to_kill_task_id, process_to_kill_pid, process_to_kill_node_id + ): + if process_to_kill_pid is not None: + + @ray.remote + def kill_process(pid): + import psutil + + proc = psutil.Process(pid) + proc.kill() + + scheduling_strategy = ( + ray.util.scheduling_strategies.NodeAffinitySchedulingStrategy( + node_id=process_to_kill_node_id, + soft=False, + ) + ) + kill_process.options(scheduling_strategy=scheduling_strategy).remote( + process_to_kill_pid + ) + logging.info( + f"Killing pid {process_to_kill_pid} on node {process_to_kill_node_id}" + ) + # Store both task_id and pid because retried tasks have same task_id. + self.killed.add((process_to_kill_task_id, process_to_kill_pid)) + + +def get_and_run_resource_killer( + resource_killer_cls, + kill_interval_s, + namespace=None, + lifetime=None, + no_start=False, + max_to_kill=2, + batch_size_to_kill=1, + kill_delay_s=0, + kill_filter_fn=None, +): + assert ray.is_initialized(), "The API is only available when Ray is initialized." + + head_node_id = ray.get_runtime_context().get_node_id() + # Schedule the actor on the current node. + resource_killer = resource_killer_cls.options( + scheduling_strategy=NodeAffinitySchedulingStrategy( + node_id=head_node_id, soft=False + ), + namespace=namespace, + name="ResourceKiller", + lifetime=lifetime, + ).remote( + head_node_id, + kill_interval_s=kill_interval_s, + kill_delay_s=kill_delay_s, + max_to_kill=max_to_kill, + batch_size_to_kill=batch_size_to_kill, + kill_filter_fn=kill_filter_fn, + ) + print("Waiting for ResourceKiller to be ready...") + ray.get(resource_killer.ready.remote()) + print("ResourceKiller is ready now.") + if not no_start: + resource_killer.run.remote() + return resource_killer + + +def get_actor_node_id(actor_handle: "ray.actor.ActorHandle") -> str: + return ray.get( + actor_handle.__ray_call__.remote( + lambda self: ray.get_runtime_context().get_node_id() + ) + ) + + +@contextmanager +def chdir(d: str): + old_dir = os.getcwd() + os.chdir(d) + try: + yield + finally: + os.chdir(old_dir) + + +def test_get_directory_size_bytes(): + with tempfile.TemporaryDirectory() as tmp_dir, chdir(tmp_dir): + assert ray._private.utils.get_directory_size_bytes(tmp_dir) == 0 + with open("test_file", "wb") as f: + f.write(os.urandom(100)) + assert ray._private.utils.get_directory_size_bytes(tmp_dir) == 100 + with open("test_file_2", "wb") as f: + f.write(os.urandom(50)) + assert ray._private.utils.get_directory_size_bytes(tmp_dir) == 150 + os.mkdir("subdir") + with open("subdir/subdir_file", "wb") as f: + f.write(os.urandom(2)) + assert ray._private.utils.get_directory_size_bytes(tmp_dir) == 152 + + +def check_local_files_gced(cluster): + for node in cluster.list_all_nodes(): + for subdir in ["conda", "pip", "working_dir_files", "py_modules_files"]: + all_files = os.listdir( + os.path.join(node.get_runtime_env_dir_path(), subdir) + ) + # Check that there are no files remaining except for .lock files + # and generated requirements.txt files. + # Note: On Windows the top folder is not deleted as it is in use. + # TODO(architkulkarni): these files should get cleaned up too! + items = list(filter(lambda f: not f.endswith((".lock", ".txt")), all_files)) + if len(items) > 0: + print(f"runtime_env files not GC'd from subdir '{subdir}': {items}") + return False + return True + + +def generate_runtime_env_dict(field, spec_format, tmp_path, pip_list=None): + if pip_list is None: + pip_list = ["pip-install-test==0.5"] + if field == "conda": + conda_dict = {"dependencies": ["pip", {"pip": pip_list}]} + if spec_format == "file": + conda_file = tmp_path / f"environment-{hash(str(pip_list))}.yml" + conda_file.write_text(yaml.dump(conda_dict)) + conda = str(conda_file) + elif spec_format == "python_object": + conda = conda_dict + runtime_env = {"conda": conda} + elif field == "pip": + if spec_format == "file": + pip_file = tmp_path / f"requirements-{hash(str(pip_list))}.txt" + pip_file.write_text("\n".join(pip_list)) + pip = str(pip_file) + elif spec_format == "python_object": + pip = pip_list + runtime_env = {"pip": pip} + return runtime_env + + +def check_spilled_mb(address, spilled=None, restored=None, fallback=None): + def ok(): + s = memory_summary(address=address["address"], stats_only=True) + print(s) + if restored: + if "Restored {} MiB".format(restored) not in s: + return False + else: + if "Restored" in s: + return False + if spilled: + if not isinstance(spilled, list): + spilled_lst = [spilled] + else: + spilled_lst = spilled + found = False + for n in spilled_lst: + if "Spilled {} MiB".format(n) in s: + found = True + if not found: + return False + else: + if "Spilled" in s: + return False + if fallback: + if "Plasma filesystem mmap usage: {} MiB".format(fallback) not in s: + return False + else: + if "Plasma filesystem mmap usage:" in s: + return False + return True + + wait_for_condition(ok, timeout=3, retry_interval_ms=1000) + + +def no_resource_leaks_excluding_node_resources(): + cluster_resources = ray.cluster_resources() + available_resources = ray.available_resources() + for r in ray.cluster_resources(): + if "node" in r: + del cluster_resources[r] + del available_resources[r] + + return cluster_resources == available_resources + + +def job_hook(**kwargs): + """Function called by reflection by test_cli_integration.""" + cmd = " ".join(kwargs["entrypoint"]) + print(f"hook intercepted: {cmd}") + sys.exit(0) + + +def wandb_setup_api_key_hook(): + """ + Example external hook to set up W&B API key in + WandbIntegrationTest.testWandbLoggerConfig + """ + return "abcd" + + +# Get node stats from node manager. +def get_node_stats(raylet, num_retry=5, timeout=2): + import grpc + + from ray._private.grpc_utils import init_grpc_channel + from ray.core.generated import node_manager_pb2_grpc + + raylet_address = build_address( + raylet["NodeManagerAddress"], raylet["NodeManagerPort"] + ) + channel = init_grpc_channel(raylet_address) + stub = node_manager_pb2_grpc.NodeManagerServiceStub(channel) + for _ in range(num_retry): + try: + reply = stub.GetNodeStats( + node_manager_pb2.GetNodeStatsRequest(), timeout=timeout + ) + break + except grpc.RpcError: + continue + assert reply is not None + return reply + + +# Gets resource usage assuming gcs is local. +def get_resource_usage(gcs_address, timeout=10): + from ray._private.grpc_utils import init_grpc_channel + from ray.core.generated import gcs_service_pb2_grpc + + if not gcs_address: + gcs_address = ray.worker._global_node.gcs_address + + gcs_channel = init_grpc_channel( + gcs_address, ray_constants.GLOBAL_GRPC_OPTIONS, asynchronous=False + ) + + gcs_node_resources_stub = gcs_service_pb2_grpc.NodeResourceInfoGcsServiceStub( + gcs_channel + ) + + request = gcs_service_pb2.GetAllResourceUsageRequest() + response = gcs_node_resources_stub.GetAllResourceUsage(request, timeout=timeout) + resources_batch_data = response.resource_usage_data + + return resources_batch_data + + +# Gets the load metrics report assuming gcs is local. +def get_load_metrics_report(webui_url): + webui_url = format_web_url(webui_url) + response = requests.get(f"{webui_url}/api/cluster_status") + response.raise_for_status() + return response.json()["data"]["clusterStatus"]["loadMetricsReport"] + + +# Send a RPC to the raylet to have it self-destruct its process. +def kill_raylet(raylet, graceful=False): + import grpc + from grpc._channel import _InactiveRpcError + + from ray.core.generated import node_manager_pb2_grpc + + raylet_address = build_address( + raylet["NodeManagerAddress"], raylet["NodeManagerPort"] + ) + channel = grpc.insecure_channel(raylet_address) + stub = node_manager_pb2_grpc.NodeManagerServiceStub(channel) + try: + stub.ShutdownRaylet(node_manager_pb2.ShutdownRayletRequest(graceful=graceful)) + except _InactiveRpcError: + assert not graceful + + +def get_gcs_memory_used(): + import psutil + + m = { + proc.info["name"]: proc.info["memory_info"].rss + for proc in psutil.process_iter(["status", "name", "memory_info"]) + if ( + proc.info["status"] not in (psutil.STATUS_ZOMBIE, psutil.STATUS_DEAD) + and proc.info["name"] in ("gcs_server", "redis-server") + ) + } + assert "gcs_server" in m + return sum(m.values()) + + +def safe_write_to_results_json( + result: dict, + default_file_name: str = "/tmp/release_test_output.json", + env_var: Optional[str] = "TEST_OUTPUT_JSON", +): + """ + Safe (atomic) write to file to guard against malforming the json + if the job gets interrupted in the middle of writing. + """ + test_output_json = os.environ.get(env_var, default_file_name) + test_output_json_tmp = f"{test_output_json}.tmp.{str(uuid.uuid4())}" + with open(test_output_json_tmp, "wt") as f: + json.dump(result, f) + f.flush() + os.replace(test_output_json_tmp, test_output_json) + logger.info(f"Wrote results to {test_output_json}") + logger.info(json.dumps(result)) + + +def get_current_unused_port(): + """ + Returns a port number that is not currently in use. + + This is useful for testing when we need to bind to a port but don't + care which one. + + Returns: + A port number that is not currently in use. (Note that this port + might become used by the time you try to bind to it.) + """ + sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM) + + # Bind the socket to a local address with a random port number + sock.bind(("localhost", 0)) + + port = sock.getsockname()[1] + sock.close() + return port + + +# Global counter to test different return values +# for external_ray_cluster_activity_hook1. +ray_cluster_activity_hook_counter = 0 +ray_cluster_activity_hook_5_counter = 0 + + +def external_ray_cluster_activity_hook1(): + """ + Example external hook for test_component_activities_hook. + Returns valid response and increments counter in `reason` + field on each call. + """ + global ray_cluster_activity_hook_counter + ray_cluster_activity_hook_counter += 1 + + from pydantic import BaseModel, Extra + + class TestRayActivityResponse(BaseModel, extra=Extra.allow): + """ + Redefinition of dashboard.modules.api.api_head.RayActivityResponse + used in test_component_activities_hook to mimic typical + usage of redefining or extending response type. + """ + + is_active: str + reason: Optional[str] = None + timestamp: float + + return { + "test_component1": TestRayActivityResponse( + is_active="ACTIVE", + reason=f"Counter: {ray_cluster_activity_hook_counter}", + timestamp=datetime.now().timestamp(), + ) + } + + +def external_ray_cluster_activity_hook2(): + """ + Example external hook for test_component_activities_hook. + Returns invalid output because the value of `test_component2` + should be of type RayActivityResponse. + """ + return {"test_component2": "bad_output"} + + +def external_ray_cluster_activity_hook3(): + """ + Example external hook for test_component_activities_hook. + Returns invalid output because return type is not + Dict[str, RayActivityResponse] + """ + return "bad_output" + + +def external_ray_cluster_activity_hook4(): + """ + Example external hook for test_component_activities_hook. + Errors during execution. + """ + raise Exception("Error in external cluster activity hook") + + +def external_ray_cluster_activity_hook5(): + """ + Example external hook for test_component_activities_hook. + Returns valid response and increments counter in `reason` + field on each call. + """ + global ray_cluster_activity_hook_5_counter + ray_cluster_activity_hook_5_counter += 1 + return { + "test_component5": { + "is_active": "ACTIVE", + "reason": f"Counter: {ray_cluster_activity_hook_5_counter}", + "timestamp": datetime.now().timestamp(), + } + } + + +# TODO(rickyx): We could remove this once we unify the autoscaler v1 and v2 +# code path for ray status +def reset_autoscaler_v2_enabled_cache(): + import ray.autoscaler.v2.utils as u + + u.cached_is_autoscaler_v2 = None + + +def _terminate_ec2_instance(node_ip: str) -> None: + logging.info(f"Terminating instance {node_ip}") + # This command uses IMDSv2 to get the host instance id and region. + # After that it terminates itself using aws cli. + command = ( + 'instanceId=$(curl -H "X-aws-ec2-metadata-token: $TOKEN" http://169.254.169.254/latest/meta-data/instance-id/);' # noqa: E501 + 'region=$(curl -H "X-aws-ec2-metadata-token: $TOKEN" http://169.254.169.254/latest/meta-data/placement/region);' # noqa: E501 + "aws ec2 terminate-instances --region $region --instance-ids $instanceId" # noqa: E501 + ) + _execute_command_on_node(command, node_ip) + + +def _execute_command_on_node(command: str, node_ip: str): + logging.debug(f"Executing command on node {node_ip}: {command}") + + multi_line_command = ( + 'TOKEN=$(curl -X PUT "http://169.254.169.254/latest/api/token" -H "X-aws-ec2-metadata-token-ttl-seconds: 21600");' # noqa: E501 + f"{command}" + ) + + # This is a feature on Anyscale platform that enables + # easy ssh access to worker nodes. + ssh_command = f"ssh -o StrictHostKeyChecking=no -o UserKnownHostsFile=/dev/null -p 2222 ray@{node_ip} '{multi_line_command}'" # noqa: E501 + + try: + subprocess.run( + ssh_command, shell=True, capture_output=True, text=True, check=True + ) + except subprocess.CalledProcessError as e: + print("Exit code:", e.returncode) + print("Stderr:", e.stderr) + + +RPC_FAILURE_MAP = { + "request": "100:0:0", + "response": "0:100:0", + "in_flight": "0:0:100", +} + +RPC_FAILURE_TYPES = list(RPC_FAILURE_MAP.keys()) diff --git a/lib/python3.12/site-packages/ray/_private/thirdparty/__init__.py b/lib/python3.12/site-packages/ray/_private/thirdparty/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/lib/python3.12/site-packages/ray/_private/thirdparty/__pycache__/__init__.cpython-312.pyc b/lib/python3.12/site-packages/ray/_private/thirdparty/__pycache__/__init__.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..cb4433c19961afbced24bde75e2e6b63742cb8a0 Binary files /dev/null and b/lib/python3.12/site-packages/ray/_private/thirdparty/__pycache__/__init__.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/_private/thirdparty/dacite/__init__.py b/lib/python3.12/site-packages/ray/_private/thirdparty/dacite/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..318cd9b60a2b4cd78072ac24c6fc669178fe8535 --- /dev/null +++ b/lib/python3.12/site-packages/ray/_private/thirdparty/dacite/__init__.py @@ -0,0 +1,3 @@ +from .config import Config +from .core import from_dict +from .exceptions import * diff --git a/lib/python3.12/site-packages/ray/_private/thirdparty/dacite/__pycache__/__init__.cpython-312.pyc b/lib/python3.12/site-packages/ray/_private/thirdparty/dacite/__pycache__/__init__.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..2e92e29e55814c52ddfdf790ea7ec1154cc5c1c9 Binary files /dev/null and b/lib/python3.12/site-packages/ray/_private/thirdparty/dacite/__pycache__/__init__.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/_private/thirdparty/dacite/__pycache__/config.cpython-312.pyc b/lib/python3.12/site-packages/ray/_private/thirdparty/dacite/__pycache__/config.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..d58ea66e80ab2d9a7152df3f4c3fa6be35e1a4e0 Binary files /dev/null and b/lib/python3.12/site-packages/ray/_private/thirdparty/dacite/__pycache__/config.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/_private/thirdparty/dacite/__pycache__/core.cpython-312.pyc b/lib/python3.12/site-packages/ray/_private/thirdparty/dacite/__pycache__/core.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..41774a3934b32f210b9711af343d30c23d98f3cd Binary files /dev/null and b/lib/python3.12/site-packages/ray/_private/thirdparty/dacite/__pycache__/core.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/_private/thirdparty/dacite/__pycache__/data.cpython-312.pyc b/lib/python3.12/site-packages/ray/_private/thirdparty/dacite/__pycache__/data.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..75573230912b849ef46e6c41f1192aea944c6db6 Binary files /dev/null and b/lib/python3.12/site-packages/ray/_private/thirdparty/dacite/__pycache__/data.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/_private/thirdparty/dacite/__pycache__/dataclasses.cpython-312.pyc b/lib/python3.12/site-packages/ray/_private/thirdparty/dacite/__pycache__/dataclasses.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..9bf0d1deb09f8860e2e27767a18f227120c995ad Binary files /dev/null and b/lib/python3.12/site-packages/ray/_private/thirdparty/dacite/__pycache__/dataclasses.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/_private/thirdparty/dacite/__pycache__/exceptions.cpython-312.pyc b/lib/python3.12/site-packages/ray/_private/thirdparty/dacite/__pycache__/exceptions.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..ea8742b2c2df7edba4bdd9070b81a09a613267c7 Binary files /dev/null and b/lib/python3.12/site-packages/ray/_private/thirdparty/dacite/__pycache__/exceptions.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/_private/thirdparty/dacite/__pycache__/types.cpython-312.pyc b/lib/python3.12/site-packages/ray/_private/thirdparty/dacite/__pycache__/types.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..508d269b9630d4a4e41478c9e8b20a72e419d9bc Binary files /dev/null and b/lib/python3.12/site-packages/ray/_private/thirdparty/dacite/__pycache__/types.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/_private/thirdparty/dacite/config.py b/lib/python3.12/site-packages/ray/_private/thirdparty/dacite/config.py new file mode 100644 index 0000000000000000000000000000000000000000..1766a68e2235cc9cde28116ba7e8310f048159f2 --- /dev/null +++ b/lib/python3.12/site-packages/ray/_private/thirdparty/dacite/config.py @@ -0,0 +1,12 @@ +from dataclasses import dataclass, field +from typing import Dict, Any, Callable, Optional, Type, List + + +@dataclass +class Config: + type_hooks: Dict[Type, Callable[[Any], Any]] = field(default_factory=dict) + cast: List[Type] = field(default_factory=list) + forward_references: Optional[Dict[str, Any]] = None + check_types: bool = True + strict: bool = False + strict_unions_match: bool = False diff --git a/lib/python3.12/site-packages/ray/_private/thirdparty/dacite/core.py b/lib/python3.12/site-packages/ray/_private/thirdparty/dacite/core.py new file mode 100644 index 0000000000000000000000000000000000000000..eccfe5e9640232039b8089c6f9756f2beed582ef --- /dev/null +++ b/lib/python3.12/site-packages/ray/_private/thirdparty/dacite/core.py @@ -0,0 +1,140 @@ +import copy +from dataclasses import is_dataclass +from itertools import zip_longest +from typing import TypeVar, Type, Optional, get_type_hints, Mapping, Any + +from .config import Config +from .data import Data +from .dataclasses import get_default_value_for_field, create_instance, DefaultValueNotFoundError, get_fields +from .exceptions import ( + ForwardReferenceError, + WrongTypeError, + DaciteError, + UnionMatchError, + MissingValueError, + DaciteFieldError, + UnexpectedDataError, + StrictUnionMatchError, +) +from .types import ( + is_instance, + is_generic_collection, + is_union, + extract_generic, + is_optional, + transform_value, + extract_origin_collection, + is_init_var, + extract_init_var, +) + +T = TypeVar("T") + + +def from_dict(data_class: Type[T], data: Data, config: Optional[Config] = None) -> T: + """Create a data class instance from a dictionary. + + :param data_class: a data class type + :param data: a dictionary of a input data + :param config: a configuration of the creation process + :return: an instance of a data class + """ + init_values: Data = {} + post_init_values: Data = {} + config = config or Config() + try: + data_class_hints = get_type_hints(data_class, globalns=config.forward_references) + except NameError as error: + raise ForwardReferenceError(str(error)) + data_class_fields = get_fields(data_class) + if config.strict: + extra_fields = set(data.keys()) - {f.name for f in data_class_fields} + if extra_fields: + raise UnexpectedDataError(keys=extra_fields) + for field in data_class_fields: + field = copy.copy(field) + field.type = data_class_hints[field.name] + try: + try: + field_data = data[field.name] + transformed_value = transform_value( + type_hooks=config.type_hooks, cast=config.cast, target_type=field.type, value=field_data + ) + value = _build_value(type_=field.type, data=transformed_value, config=config) + except DaciteFieldError as error: + error.update_path(field.name) + raise + if config.check_types and not is_instance(value, field.type): + raise WrongTypeError(field_path=field.name, field_type=field.type, value=value) + except KeyError: + try: + value = get_default_value_for_field(field) + except DefaultValueNotFoundError: + if not field.init: + continue + raise MissingValueError(field.name) + if field.init: + init_values[field.name] = value + else: + post_init_values[field.name] = value + + return create_instance(data_class=data_class, init_values=init_values, post_init_values=post_init_values) + + +def _build_value(type_: Type, data: Any, config: Config) -> Any: + if is_init_var(type_): + type_ = extract_init_var(type_) + if is_union(type_): + return _build_value_for_union(union=type_, data=data, config=config) + elif is_generic_collection(type_) and is_instance(data, extract_origin_collection(type_)): + return _build_value_for_collection(collection=type_, data=data, config=config) + elif is_dataclass(type_) and is_instance(data, Data): + return from_dict(data_class=type_, data=data, config=config) + return data + + +def _build_value_for_union(union: Type, data: Any, config: Config) -> Any: + types = extract_generic(union) + if is_optional(union) and len(types) == 2: + return _build_value(type_=types[0], data=data, config=config) + union_matches = {} + for inner_type in types: + try: + # noinspection PyBroadException + try: + data = transform_value( + type_hooks=config.type_hooks, cast=config.cast, target_type=inner_type, value=data + ) + except Exception: # pylint: disable=broad-except + continue + value = _build_value(type_=inner_type, data=data, config=config) + if is_instance(value, inner_type): + if config.strict_unions_match: + union_matches[inner_type] = value + else: + return value + except DaciteError: + pass + if config.strict_unions_match: + if len(union_matches) > 1: + raise StrictUnionMatchError(union_matches) + return union_matches.popitem()[1] + if not config.check_types: + return data + raise UnionMatchError(field_type=union, value=data) + + +def _build_value_for_collection(collection: Type, data: Any, config: Config) -> Any: + data_type = data.__class__ + if is_instance(data, Mapping): + item_type = extract_generic(collection, defaults=(Any, Any))[1] + return data_type((key, _build_value(type_=item_type, data=value, config=config)) for key, value in data.items()) + elif is_instance(data, tuple): + types = extract_generic(collection) + if len(types) == 2 and types[1] == Ellipsis: + return data_type(_build_value(type_=types[0], data=item, config=config) for item in data) + return data_type( + _build_value(type_=type_, data=item, config=config) for item, type_ in zip_longest(data, types) + ) + item_type = extract_generic(collection, defaults=(Any,))[0] + return data_type(_build_value(type_=item_type, data=item, config=config) for item in data) diff --git a/lib/python3.12/site-packages/ray/_private/thirdparty/dacite/data.py b/lib/python3.12/site-packages/ray/_private/thirdparty/dacite/data.py new file mode 100644 index 0000000000000000000000000000000000000000..560cdf8beddb082122f26bb548329ec95f23581d --- /dev/null +++ b/lib/python3.12/site-packages/ray/_private/thirdparty/dacite/data.py @@ -0,0 +1,3 @@ +from typing import Dict, Any + +Data = Dict[str, Any] diff --git a/lib/python3.12/site-packages/ray/_private/thirdparty/dacite/dataclasses.py b/lib/python3.12/site-packages/ray/_private/thirdparty/dacite/dataclasses.py new file mode 100644 index 0000000000000000000000000000000000000000..b5db7ac013444f4d5d72439dd396235278302ba7 --- /dev/null +++ b/lib/python3.12/site-packages/ray/_private/thirdparty/dacite/dataclasses.py @@ -0,0 +1,33 @@ +from dataclasses import Field, MISSING, _FIELDS, _FIELD, _FIELD_INITVAR # type: ignore +from typing import Type, Any, TypeVar, List + +from .data import Data +from .types import is_optional + +T = TypeVar("T", bound=Any) + + +class DefaultValueNotFoundError(Exception): + pass + + +def get_default_value_for_field(field: Field) -> Any: + if field.default != MISSING: + return field.default + elif field.default_factory != MISSING: # type: ignore + return field.default_factory() # type: ignore + elif is_optional(field.type): + return None + raise DefaultValueNotFoundError() + + +def create_instance(data_class: Type[T], init_values: Data, post_init_values: Data) -> T: + instance = data_class(**init_values) + for key, value in post_init_values.items(): + setattr(instance, key, value) + return instance + + +def get_fields(data_class: Type[T]) -> List[Field]: + fields = getattr(data_class, _FIELDS) + return [f for f in fields.values() if f._field_type is _FIELD or f._field_type is _FIELD_INITVAR] diff --git a/lib/python3.12/site-packages/ray/_private/thirdparty/dacite/exceptions.py b/lib/python3.12/site-packages/ray/_private/thirdparty/dacite/exceptions.py new file mode 100644 index 0000000000000000000000000000000000000000..871e73f1c29c8b301f04dd5785bd0d9bcca6e4d3 --- /dev/null +++ b/lib/python3.12/site-packages/ray/_private/thirdparty/dacite/exceptions.py @@ -0,0 +1,79 @@ +from typing import Any, Type, Optional, Set, Dict + + +def _name(type_: Type) -> str: + return type_.__name__ if hasattr(type_, "__name__") else str(type_) + + +class DaciteError(Exception): + pass + + +class DaciteFieldError(DaciteError): + def __init__(self, field_path: Optional[str] = None): + super().__init__() + self.field_path = field_path + + def update_path(self, parent_field_path: str) -> None: + if self.field_path: + self.field_path = f"{parent_field_path}.{self.field_path}" + else: + self.field_path = parent_field_path + + +class WrongTypeError(DaciteFieldError): + def __init__(self, field_type: Type, value: Any, field_path: Optional[str] = None) -> None: + super().__init__(field_path=field_path) + self.field_type = field_type + self.value = value + + def __str__(self) -> str: + return ( + f'wrong value type for field "{self.field_path}" - should be "{_name(self.field_type)}" ' + f'instead of value "{self.value}" of type "{_name(type(self.value))}"' + ) + + +class MissingValueError(DaciteFieldError): + def __init__(self, field_path: Optional[str] = None): + super().__init__(field_path=field_path) + + def __str__(self) -> str: + return f'missing value for field "{self.field_path}"' + + +class UnionMatchError(WrongTypeError): + def __str__(self) -> str: + return ( + f'can not match type "{_name(type(self.value))}" to any type ' + f'of "{self.field_path}" union: {_name(self.field_type)}' + ) + + +class StrictUnionMatchError(DaciteFieldError): + def __init__(self, union_matches: Dict[Type, Any], field_path: Optional[str] = None) -> None: + super().__init__(field_path=field_path) + self.union_matches = union_matches + + def __str__(self) -> str: + conflicting_types = ", ".join(_name(type_) for type_ in self.union_matches) + return f'can not choose between possible Union matches for field "{self.field_path}": {conflicting_types}' + + +class ForwardReferenceError(DaciteError): + def __init__(self, message: str) -> None: + super().__init__() + self.message = message + + def __str__(self) -> str: + return f"can not resolve forward reference: {self.message}" + + +class UnexpectedDataError(DaciteError): + def __init__(self, keys: Set[str]) -> None: + super().__init__() + self.keys = keys + + def __str__(self) -> str: + formatted_keys = ", ".join(f'"{key}"' for key in self.keys) + return f"can not match {formatted_keys} to any data class field" diff --git a/lib/python3.12/site-packages/ray/_private/thirdparty/dacite/py.typed b/lib/python3.12/site-packages/ray/_private/thirdparty/dacite/py.typed new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/lib/python3.12/site-packages/ray/_private/thirdparty/dacite/types.py b/lib/python3.12/site-packages/ray/_private/thirdparty/dacite/types.py new file mode 100644 index 0000000000000000000000000000000000000000..1d4dfea4f2a5b124fa00ef5237fc96567648b598 --- /dev/null +++ b/lib/python3.12/site-packages/ray/_private/thirdparty/dacite/types.py @@ -0,0 +1,172 @@ +from dataclasses import InitVar +from typing import Type, Any, Optional, Union, Collection, TypeVar, Dict, Callable, Mapping, List, Tuple + +T = TypeVar("T", bound=Any) + + +def transform_value( + type_hooks: Dict[Type, Callable[[Any], Any]], cast: List[Type], target_type: Type, value: Any +) -> Any: + if target_type in type_hooks: + value = type_hooks[target_type](value) + else: + for cast_type in cast: + if is_subclass(target_type, cast_type): + if is_generic_collection(target_type): + value = extract_origin_collection(target_type)(value) + else: + value = target_type(value) + break + if is_optional(target_type): + if value is None: + return None + target_type = extract_optional(target_type) + return transform_value(type_hooks, cast, target_type, value) + if is_generic_collection(target_type) and isinstance(value, extract_origin_collection(target_type)): + collection_cls = value.__class__ + if issubclass(collection_cls, dict): + key_cls, item_cls = extract_generic(target_type, defaults=(Any, Any)) + return collection_cls( + { + transform_value(type_hooks, cast, key_cls, key): transform_value(type_hooks, cast, item_cls, item) + for key, item in value.items() + } + ) + item_cls = extract_generic(target_type, defaults=(Any,))[0] + return collection_cls(transform_value(type_hooks, cast, item_cls, item) for item in value) + return value + + +def extract_origin_collection(collection: Type) -> Type: + try: + return collection.__extra__ + except AttributeError: + return collection.__origin__ + + +def is_optional(type_: Type) -> bool: + return is_union(type_) and type(None) in extract_generic(type_) + + +def extract_optional(optional: Type[Optional[T]]) -> T: + for type_ in extract_generic(optional): + if type_ is not type(None): + return type_ + raise ValueError("can not find not-none value") + + +def is_generic(type_: Type) -> bool: + return hasattr(type_, "__origin__") + + +def is_union(type_: Type) -> bool: + return is_generic(type_) and type_.__origin__ == Union + + +def is_literal(type_: Type) -> bool: + try: + from typing import Literal # type: ignore + + return is_generic(type_) and type_.__origin__ == Literal + except ImportError: + return False + + +def is_new_type(type_: Type) -> bool: + return hasattr(type_, "__supertype__") + + +def extract_new_type(type_: Type) -> Type: + return type_.__supertype__ + + +def is_init_var(type_: Type) -> bool: + return isinstance(type_, InitVar) or type_ is InitVar + + +def extract_init_var(type_: Type) -> Union[Type, Any]: + try: + return type_.type + except AttributeError: + return Any + + +def is_instance(value: Any, type_: Type) -> bool: + if type_ == Any: + return True + elif is_union(type_): + return any(is_instance(value, t) for t in extract_generic(type_)) + elif is_generic_collection(type_): + origin = extract_origin_collection(type_) + if not isinstance(value, origin): + return False + if not extract_generic(type_): + return True + if isinstance(value, tuple): + tuple_types = extract_generic(type_) + if len(tuple_types) == 1 and tuple_types[0] == (): + return len(value) == 0 + elif len(tuple_types) == 2 and tuple_types[1] is ...: + return all(is_instance(item, tuple_types[0]) for item in value) + else: + if len(tuple_types) != len(value): + return False + return all(is_instance(item, item_type) for item, item_type in zip(value, tuple_types)) + if isinstance(value, Mapping): + key_type, val_type = extract_generic(type_, defaults=(Any, Any)) + for key, val in value.items(): + if not is_instance(key, key_type) or not is_instance(val, val_type): + return False + return True + return all(is_instance(item, extract_generic(type_, defaults=(Any,))[0]) for item in value) + elif is_new_type(type_): + return is_instance(value, extract_new_type(type_)) + elif is_literal(type_): + return value in extract_generic(type_) + elif is_init_var(type_): + return is_instance(value, extract_init_var(type_)) + elif is_type_generic(type_): + return is_subclass(value, extract_generic(type_)[0]) + else: + try: + # As described in PEP 484 - section: "The numeric tower" + if isinstance(value, (int, float)) and type_ in [float, complex]: + return True + return isinstance(value, type_) + except TypeError: + return False + + +def is_generic_collection(type_: Type) -> bool: + if not is_generic(type_): + return False + origin = extract_origin_collection(type_) + try: + return bool(origin and issubclass(origin, Collection)) + except (TypeError, AttributeError): + return False + + +def extract_generic(type_: Type, defaults: Tuple = ()) -> tuple: + try: + if hasattr(type_, "_special") and type_._special: + return defaults + return type_.__args__ or defaults # type: ignore + except AttributeError: + return defaults + + +def is_subclass(sub_type: Type, base_type: Type) -> bool: + if is_generic_collection(sub_type): + sub_type = extract_origin_collection(sub_type) + try: + return issubclass(sub_type, base_type) + except TypeError: + return False + + +def is_type_generic(type_: Type) -> bool: + try: + return type_.__origin__ in (type, Type) + except AttributeError: + return False diff --git a/lib/python3.12/site-packages/ray/_private/thirdparty/pathspec/__init__.py b/lib/python3.12/site-packages/ray/_private/thirdparty/pathspec/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..423759fd5bcc5f50c44b86d04f6b98dbd32f354c --- /dev/null +++ b/lib/python3.12/site-packages/ray/_private/thirdparty/pathspec/__init__.py @@ -0,0 +1,69 @@ +# encoding: utf-8 +""" +The *pathspec* package provides pattern matching for file paths. So far +this only includes Git's wildmatch pattern matching (the style used for +".gitignore" files). + +The following classes are imported and made available from the root of +the `pathspec` package: + +- :class:`pathspec.pathspec.PathSpec` + +- :class:`pathspec.pattern.Pattern` + +- :class:`pathspec.pattern.RegexPattern` + +- :class:`pathspec.util.RecursionError` + +The following functions are also imported: + +- :func:`pathspec.util.iter_tree` +- :func:`pathspec.util.lookup_pattern` +- :func:`pathspec.util.match_files` +""" +from __future__ import unicode_literals + +__author__ = "Caleb P. Burns" +__copyright__ = "Copyright © 2013-2020 Caleb P. Burns" +__created__ = "2013-10-12" +__credits__ = [ + "dahlia ", + "highb ", + "029xue ", + "mikexstudios ", + "nhumrich ", + "davidfraser ", + "demurgos ", + "ghickman ", + "nvie ", + "adrienverge ", + "AndersBlomdell ", + "highb ", + "thmxv ", + "wimglenn ", + "hugovk ", + "dcecile ", + "mroutis ", + "jdufresne ", + "groodt ", + "ftrofin ", + "pykong ", + "nhhollander ", +] +__email__ = "cpburnz@gmail.com" +__license__ = "MPL 2.0" +__project__ = "pathspec" +__status__ = "Development" +__updated__ = "2020-11-07" +__version__ = "0.8.1" + +from .pathspec import PathSpec +from .pattern import Pattern, RegexPattern +from .util import iter_tree, lookup_pattern, match_files, RecursionError + +# Load pattern implementations. +from . import patterns + +# Expose `GitIgnorePattern` class in the root module for backward +# compatibility with v0.4. +from .patterns.gitwildmatch import GitIgnorePattern diff --git a/lib/python3.12/site-packages/ray/_private/thirdparty/pathspec/__pycache__/__init__.cpython-312.pyc b/lib/python3.12/site-packages/ray/_private/thirdparty/pathspec/__pycache__/__init__.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..ea13fed3fe2e226b69bdffe564274107bf56d3b6 Binary files /dev/null and b/lib/python3.12/site-packages/ray/_private/thirdparty/pathspec/__pycache__/__init__.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/_private/thirdparty/pathspec/__pycache__/compat.cpython-312.pyc b/lib/python3.12/site-packages/ray/_private/thirdparty/pathspec/__pycache__/compat.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..b431580724ae14d200851f953f1e5e67c63ce75a Binary files /dev/null and b/lib/python3.12/site-packages/ray/_private/thirdparty/pathspec/__pycache__/compat.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/_private/thirdparty/pathspec/__pycache__/pathspec.cpython-312.pyc b/lib/python3.12/site-packages/ray/_private/thirdparty/pathspec/__pycache__/pathspec.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..7a0dc0b4952fe687078fb06f8ed52f630a293c18 Binary files /dev/null and b/lib/python3.12/site-packages/ray/_private/thirdparty/pathspec/__pycache__/pathspec.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/_private/thirdparty/pathspec/__pycache__/pattern.cpython-312.pyc b/lib/python3.12/site-packages/ray/_private/thirdparty/pathspec/__pycache__/pattern.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..5c8ed2181b01b65d3c6ae5c82bbb2da3eb25fbf9 Binary files /dev/null and b/lib/python3.12/site-packages/ray/_private/thirdparty/pathspec/__pycache__/pattern.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/_private/thirdparty/pathspec/__pycache__/util.cpython-312.pyc b/lib/python3.12/site-packages/ray/_private/thirdparty/pathspec/__pycache__/util.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..c4aedc56946b7d7a4ada1e7705de54b2e3ac3bdb Binary files /dev/null and b/lib/python3.12/site-packages/ray/_private/thirdparty/pathspec/__pycache__/util.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/_private/thirdparty/pathspec/compat.py b/lib/python3.12/site-packages/ray/_private/thirdparty/pathspec/compat.py new file mode 100644 index 0000000000000000000000000000000000000000..37c6480510f4918fec453ad960be4e7e168dde7d --- /dev/null +++ b/lib/python3.12/site-packages/ray/_private/thirdparty/pathspec/compat.py @@ -0,0 +1,38 @@ +# encoding: utf-8 +""" +This module provides compatibility between Python 2 and 3. Hardly +anything is used by this project to constitute including `six`_. + +.. _`six`: http://pythonhosted.org/six +""" + +import sys + +if sys.version_info[0] < 3: + # Python 2. + unicode = unicode + string_types = (basestring,) + + from collections import Iterable + from itertools import izip_longest + + def iterkeys(mapping): + return mapping.iterkeys() + +else: + # Python 3. + unicode = str + string_types = (unicode,) + + from collections.abc import Iterable + from itertools import zip_longest as izip_longest + + def iterkeys(mapping): + return mapping.keys() + +try: + # Python 3.6+. + from collections.abc import Collection +except ImportError: + # Python 2.7 - 3.5. + from collections import Container as Collection diff --git a/lib/python3.12/site-packages/ray/_private/thirdparty/pathspec/pathspec.py b/lib/python3.12/site-packages/ray/_private/thirdparty/pathspec/pathspec.py new file mode 100644 index 0000000000000000000000000000000000000000..73250efeb4da23db302317e750e8e8fb3b691ebf --- /dev/null +++ b/lib/python3.12/site-packages/ray/_private/thirdparty/pathspec/pathspec.py @@ -0,0 +1,206 @@ +# encoding: utf-8 +""" +This module provides an object oriented interface for pattern matching +of files. +""" + +from . import util +from .compat import Collection, iterkeys, izip_longest, string_types, unicode + + +class PathSpec(object): + """ + The :class:`PathSpec` class is a wrapper around a list of compiled + :class:`.Pattern` instances. + """ + + def __init__(self, patterns): + """ + Initializes the :class:`PathSpec` instance. + + *patterns* (:class:`~collections.abc.Collection` or :class:`~collections.abc.Iterable`) + yields each compiled pattern (:class:`.Pattern`). + """ + + self.patterns = patterns if isinstance(patterns, Collection) else list(patterns) + """ + *patterns* (:class:`~collections.abc.Collection` of :class:`.Pattern`) + contains the compiled patterns. + """ + + def __eq__(self, other): + """ + Tests the equality of this path-spec with *other* (:class:`PathSpec`) + by comparing their :attr:`~PathSpec.patterns` attributes. + """ + if isinstance(other, PathSpec): + paired_patterns = izip_longest(self.patterns, other.patterns) + return all(a == b for a, b in paired_patterns) + else: + return NotImplemented + + def __len__(self): + """ + Returns the number of compiled patterns this path-spec contains + (:class:`int`). + """ + return len(self.patterns) + + def __add__(self, other): + """ + Combines the :attr:`Pathspec.patterns` patterns from two + :class:`PathSpec` instances. + """ + if isinstance(other, PathSpec): + return PathSpec(self.patterns + other.patterns) + else: + return NotImplemented + + def __iadd__(self, other): + """ + Adds the :attr:`Pathspec.patterns` patterns from one :class:`PathSpec` + instance to this instance. + """ + if isinstance(other, PathSpec): + self.patterns += other.patterns + return self + else: + return NotImplemented + + @classmethod + def from_lines(cls, pattern_factory, lines): + """ + Compiles the pattern lines. + + *pattern_factory* can be either the name of a registered pattern + factory (:class:`str`), or a :class:`~collections.abc.Callable` used + to compile patterns. It must accept an uncompiled pattern (:class:`str`) + and return the compiled pattern (:class:`.Pattern`). + + *lines* (:class:`~collections.abc.Iterable`) yields each uncompiled + pattern (:class:`str`). This simply has to yield each line so it can + be a :class:`file` (e.g., from :func:`open` or :class:`io.StringIO`) + or the result from :meth:`str.splitlines`. + + Returns the :class:`PathSpec` instance. + """ + if isinstance(pattern_factory, string_types): + pattern_factory = util.lookup_pattern(pattern_factory) + if not callable(pattern_factory): + raise TypeError("pattern_factory:{!r} is not callable.".format(pattern_factory)) + + if not util._is_iterable(lines): + raise TypeError("lines:{!r} is not an iterable.".format(lines)) + + lines = [pattern_factory(line) for line in lines if line] + return cls(lines) + + def match_file(self, file, separators=None): + """ + Matches the file to this path-spec. + + *file* (:class:`str` or :class:`~pathlib.PurePath`) is the file path + to be matched against :attr:`self.patterns `. + + *separators* (:class:`~collections.abc.Collection` of :class:`str`) + optionally contains the path separators to normalize. See + :func:`~pathspec.util.normalize_file` for more information. + + Returns :data:`True` if *file* matched; otherwise, :data:`False`. + """ + norm_file = util.normalize_file(file, separators=separators) + return util.match_file(self.patterns, norm_file) + + def match_entries(self, entries, separators=None): + """ + Matches the entries to this path-spec. + + *entries* (:class:`~collections.abc.Iterable` of :class:`~util.TreeEntry`) + contains the entries to be matched against :attr:`self.patterns `. + + *separators* (:class:`~collections.abc.Collection` of :class:`str`; + or :data:`None`) optionally contains the path separators to + normalize. See :func:`~pathspec.util.normalize_file` for more + information. + + Returns the matched entries (:class:`~collections.abc.Iterable` of + :class:`~util.TreeEntry`). + """ + if not util._is_iterable(entries): + raise TypeError("entries:{!r} is not an iterable.".format(entries)) + + entry_map = util._normalize_entries(entries, separators=separators) + match_paths = util.match_files(self.patterns, iterkeys(entry_map)) + for path in match_paths: + yield entry_map[path] + + def match_files(self, files, separators=None): + """ + Matches the files to this path-spec. + + *files* (:class:`~collections.abc.Iterable` of :class:`str; or + :class:`pathlib.PurePath`) contains the file paths to be matched + against :attr:`self.patterns `. + + *separators* (:class:`~collections.abc.Collection` of :class:`str`; + or :data:`None`) optionally contains the path separators to + normalize. See :func:`~pathspec.util.normalize_file` for more + information. + + Returns the matched files (:class:`~collections.abc.Iterable` of + :class:`str`). + """ + if not util._is_iterable(files): + raise TypeError("files:{!r} is not an iterable.".format(files)) + + file_map = util.normalize_files(files, separators=separators) + matched_files = util.match_files(self.patterns, iterkeys(file_map)) + for path in matched_files: + yield file_map[path] + + def match_tree_entries(self, root, on_error=None, follow_links=None): + """ + Walks the specified root path for all files and matches them to this + path-spec. + + *root* (:class:`str`; or :class:`pathlib.PurePath`) is the root + directory to search. + + *on_error* (:class:`~collections.abc.Callable` or :data:`None`) + optionally is the error handler for file-system exceptions. See + :func:`~pathspec.util.iter_tree_entries` for more information. + + *follow_links* (:class:`bool` or :data:`None`) optionally is whether + to walk symbolic links that resolve to directories. See + :func:`~pathspec.util.iter_tree_files` for more information. + + Returns the matched files (:class:`~collections.abc.Iterable` of + :class:`str`). + """ + entries = util.iter_tree_entries(root, on_error=on_error, follow_links=follow_links) + return self.match_entries(entries) + + def match_tree_files(self, root, on_error=None, follow_links=None): + """ + Walks the specified root path for all files and matches them to this + path-spec. + + *root* (:class:`str`; or :class:`pathlib.PurePath`) is the root + directory to search for files. + + *on_error* (:class:`~collections.abc.Callable` or :data:`None`) + optionally is the error handler for file-system exceptions. See + :func:`~pathspec.util.iter_tree_files` for more information. + + *follow_links* (:class:`bool` or :data:`None`) optionally is whether + to walk symbolic links that resolve to directories. See + :func:`~pathspec.util.iter_tree_files` for more information. + + Returns the matched files (:class:`~collections.abc.Iterable` of + :class:`str`). + """ + files = util.iter_tree_files(root, on_error=on_error, follow_links=follow_links) + return self.match_files(files) + + # Alias `match_tree_files()` as `match_tree()`. + match_tree = match_tree_files diff --git a/lib/python3.12/site-packages/ray/_private/thirdparty/pathspec/pattern.py b/lib/python3.12/site-packages/ray/_private/thirdparty/pathspec/pattern.py new file mode 100644 index 0000000000000000000000000000000000000000..4ba4edf790c84c51ad4ff79b76022a5f64bb8e65 --- /dev/null +++ b/lib/python3.12/site-packages/ray/_private/thirdparty/pathspec/pattern.py @@ -0,0 +1,146 @@ +# encoding: utf-8 +""" +This module provides the base definition for patterns. +""" + +import re + +from .compat import unicode + + +class Pattern(object): + """ + The :class:`Pattern` class is the abstract definition of a pattern. + """ + + # Make the class dict-less. + __slots__ = ('include',) + + def __init__(self, include): + """ + Initializes the :class:`Pattern` instance. + + *include* (:class:`bool` or :data:`None`) is whether the matched + files should be included (:data:`True`), excluded (:data:`False`), + or is a null-operation (:data:`None`). + """ + + self.include = include + """ + *include* (:class:`bool` or :data:`None`) is whether the matched + files should be included (:data:`True`), excluded (:data:`False`), + or is a null-operation (:data:`None`). + """ + + def match(self, files): + """ + Matches this pattern against the specified files. + + *files* (:class:`~collections.abc.Iterable` of :class:`str`) contains + each file relative to the root directory (e.g., ``"relative/path/to/file"``). + + Returns an :class:`~collections.abc.Iterable` yielding each matched + file path (:class:`str`). + """ + raise NotImplementedError("{}.{} must override match().".format(self.__class__.__module__, self.__class__.__name__)) + + +class RegexPattern(Pattern): + """ + The :class:`RegexPattern` class is an implementation of a pattern + using regular expressions. + """ + + # Make the class dict-less. + __slots__ = ('regex',) + + def __init__(self, pattern, include=None): + """ + Initializes the :class:`RegexPattern` instance. + + *pattern* (:class:`unicode`, :class:`bytes`, :class:`re.RegexObject`, + or :data:`None`) is the pattern to compile into a regular + expression. + + *include* (:class:`bool` or :data:`None`) must be :data:`None` + unless *pattern* is a precompiled regular expression (:class:`re.RegexObject`) + in which case it is whether matched files should be included + (:data:`True`), excluded (:data:`False`), or is a null operation + (:data:`None`). + + .. NOTE:: Subclasses do not need to support the *include* + parameter. + """ + + self.regex = None + """ + *regex* (:class:`re.RegexObject`) is the regular expression for the + pattern. + """ + + if isinstance(pattern, (unicode, bytes)): + assert include is None, "include:{!r} must be null when pattern:{!r} is a string.".format(include, pattern) + regex, include = self.pattern_to_regex(pattern) + # NOTE: Make sure to allow a null regular expression to be + # returned for a null-operation. + if include is not None: + regex = re.compile(regex) + + elif pattern is not None and hasattr(pattern, 'match'): + # Assume pattern is a precompiled regular expression. + # - NOTE: Used specified *include*. + regex = pattern + + elif pattern is None: + # NOTE: Make sure to allow a null pattern to be passed for a + # null-operation. + assert include is None, "include:{!r} must be null when pattern:{!r} is null.".format(include, pattern) + + else: + raise TypeError("pattern:{!r} is not a string, RegexObject, or None.".format(pattern)) + + super(RegexPattern, self).__init__(include) + self.regex = regex + + def __eq__(self, other): + """ + Tests the equality of this regex pattern with *other* (:class:`RegexPattern`) + by comparing their :attr:`~Pattern.include` and :attr:`~RegexPattern.regex` + attributes. + """ + if isinstance(other, RegexPattern): + return self.include == other.include and self.regex == other.regex + else: + return NotImplemented + + def match(self, files): + """ + Matches this pattern against the specified files. + + *files* (:class:`~collections.abc.Iterable` of :class:`str`) + contains each file relative to the root directory (e.g., "relative/path/to/file"). + + Returns an :class:`~collections.abc.Iterable` yielding each matched + file path (:class:`str`). + """ + if self.include is not None: + for path in files: + if self.regex.match(path) is not None: + yield path + + @classmethod + def pattern_to_regex(cls, pattern): + """ + Convert the pattern into an uncompiled regular expression. + + *pattern* (:class:`str`) is the pattern to convert into a regular + expression. + + Returns the uncompiled regular expression (:class:`str` or :data:`None`), + and whether matched files should be included (:data:`True`), + excluded (:data:`False`), or is a null-operation (:data:`None`). + + .. NOTE:: The default implementation simply returns *pattern* and + :data:`True`. + """ + return pattern, True diff --git a/lib/python3.12/site-packages/ray/_private/thirdparty/pathspec/patterns/__init__.py b/lib/python3.12/site-packages/ray/_private/thirdparty/pathspec/patterns/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..1a0d55ec74d901a6d08460e729a3ea8c321a5cd7 --- /dev/null +++ b/lib/python3.12/site-packages/ray/_private/thirdparty/pathspec/patterns/__init__.py @@ -0,0 +1,8 @@ +# encoding: utf-8 +""" +The *pathspec.patterns* package contains the pattern matching +implementations. +""" + +# Load pattern implementations. +from .gitwildmatch import GitWildMatchPattern diff --git a/lib/python3.12/site-packages/ray/_private/thirdparty/pathspec/patterns/__pycache__/__init__.cpython-312.pyc b/lib/python3.12/site-packages/ray/_private/thirdparty/pathspec/patterns/__pycache__/__init__.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..81e16b845c23ab1f9b84d797a62b3d1cc40741bd Binary files /dev/null and b/lib/python3.12/site-packages/ray/_private/thirdparty/pathspec/patterns/__pycache__/__init__.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/_private/thirdparty/pathspec/patterns/__pycache__/gitwildmatch.cpython-312.pyc b/lib/python3.12/site-packages/ray/_private/thirdparty/pathspec/patterns/__pycache__/gitwildmatch.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..016dd481f60030ea90f8ff4fb50bc49aee12b39c Binary files /dev/null and b/lib/python3.12/site-packages/ray/_private/thirdparty/pathspec/patterns/__pycache__/gitwildmatch.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/_private/thirdparty/pathspec/patterns/gitwildmatch.py b/lib/python3.12/site-packages/ray/_private/thirdparty/pathspec/patterns/gitwildmatch.py new file mode 100644 index 0000000000000000000000000000000000000000..07fd03880a9ad1f18a102f22421ffd73cf571d02 --- /dev/null +++ b/lib/python3.12/site-packages/ray/_private/thirdparty/pathspec/patterns/gitwildmatch.py @@ -0,0 +1,330 @@ +# encoding: utf-8 +""" +This module implements Git's wildmatch pattern matching which itself is +derived from Rsync's wildmatch. Git uses wildmatch for its ".gitignore" +files. +""" +from __future__ import unicode_literals + +import re +import warnings + +from .. import util +from ..compat import unicode +from ..pattern import RegexPattern + +#: The encoding to use when parsing a byte string pattern. +_BYTES_ENCODING = 'latin1' + + +class GitWildMatchPattern(RegexPattern): + """ + The :class:`GitWildMatchPattern` class represents a compiled Git + wildmatch pattern. + """ + + # Keep the dict-less class hierarchy. + __slots__ = () + + @classmethod + def pattern_to_regex(cls, pattern): + """ + Convert the pattern into a regular expression. + + *pattern* (:class:`unicode` or :class:`bytes`) is the pattern to + convert into a regular expression. + + Returns the uncompiled regular expression (:class:`unicode`, :class:`bytes`, + or :data:`None`), and whether matched files should be included + (:data:`True`), excluded (:data:`False`), or if it is a + null-operation (:data:`None`). + """ + if isinstance(pattern, unicode): + return_type = unicode + elif isinstance(pattern, bytes): + return_type = bytes + pattern = pattern.decode(_BYTES_ENCODING) + else: + raise TypeError("pattern:{!r} is not a unicode or byte string.".format(pattern)) + + pattern = pattern.strip() + + if pattern.startswith('#'): + # A pattern starting with a hash ('#') serves as a comment + # (neither includes nor excludes files). Escape the hash with a + # back-slash to match a literal hash (i.e., '\#'). + regex = None + include = None + + elif pattern == '/': + # EDGE CASE: According to `git check-ignore` (v2.4.1), a single + # '/' does not match any file. + regex = None + include = None + + elif pattern: + + if pattern.startswith('!'): + # A pattern starting with an exclamation mark ('!') negates the + # pattern (exclude instead of include). Escape the exclamation + # mark with a back-slash to match a literal exclamation mark + # (i.e., '\!'). + include = False + # Remove leading exclamation mark. + pattern = pattern[1:] + else: + include = True + + if pattern.startswith('\\'): + # Remove leading back-slash escape for escaped hash ('#') or + # exclamation mark ('!'). + pattern = pattern[1:] + + # Split pattern into segments. + pattern_segs = pattern.split('/') + + # Normalize pattern to make processing easier. + + if not pattern_segs[0]: + # A pattern beginning with a slash ('/') will only match paths + # directly on the root directory instead of any descendant + # paths. So, remove empty first segment to make pattern relative + # to root. + del pattern_segs[0] + + elif len(pattern_segs) == 1 or (len(pattern_segs) == 2 and not pattern_segs[1]): + # A single pattern without a beginning slash ('/') will match + # any descendant path. This is equivalent to "**/{pattern}". So, + # prepend with double-asterisks to make pattern relative to + # root. + # EDGE CASE: This also holds for a single pattern with a + # trailing slash (e.g. dir/). + if pattern_segs[0] != '**': + pattern_segs.insert(0, '**') + + else: + # EDGE CASE: A pattern without a beginning slash ('/') but + # contains at least one prepended directory (e.g. + # "dir/{pattern}") should not match "**/dir/{pattern}", + # according to `git check-ignore` (v2.4.1). + pass + + if not pattern_segs[-1] and len(pattern_segs) > 1: + # A pattern ending with a slash ('/') will match all descendant + # paths if it is a directory but not if it is a regular file. + # This is equivilent to "{pattern}/**". So, set last segment to + # double asterisks to include all descendants. + pattern_segs[-1] = '**' + + # Build regular expression from pattern. + output = ['^'] + need_slash = False + end = len(pattern_segs) - 1 + for i, seg in enumerate(pattern_segs): + if seg == '**': + if i == 0 and i == end: + # A pattern consisting solely of double-asterisks ('**') + # will match every path. + output.append('.+') + elif i == 0: + # A normalized pattern beginning with double-asterisks + # ('**') will match any leading path segments. + output.append('(?:.+/)?') + need_slash = False + elif i == end: + # A normalized pattern ending with double-asterisks ('**') + # will match any trailing path segments. + output.append('/.*') + else: + # A pattern with inner double-asterisks ('**') will match + # multiple (or zero) inner path segments. + output.append('(?:/.+)?') + need_slash = True + elif seg == '*': + # Match single path segment. + if need_slash: + output.append('/') + output.append('[^/]+') + need_slash = True + else: + # Match segment glob pattern. + if need_slash: + output.append('/') + output.append(cls._translate_segment_glob(seg)) + if i == end and include is True: + # A pattern ending without a slash ('/') will match a file + # or a directory (with paths underneath it). E.g., "foo" + # matches "foo", "foo/bar", "foo/bar/baz", etc. + # EDGE CASE: However, this does not hold for exclusion cases + # according to `git check-ignore` (v2.4.1). + output.append('(?:/.*)?') + need_slash = True + output.append('$') + regex = ''.join(output) + + else: + # A blank pattern is a null-operation (neither includes nor + # excludes files). + regex = None + include = None + + if regex is not None and return_type is bytes: + regex = regex.encode(_BYTES_ENCODING) + + return regex, include + + @staticmethod + def _translate_segment_glob(pattern): + """ + Translates the glob pattern to a regular expression. This is used in + the constructor to translate a path segment glob pattern to its + corresponding regular expression. + + *pattern* (:class:`str`) is the glob pattern. + + Returns the regular expression (:class:`str`). + """ + # NOTE: This is derived from `fnmatch.translate()` and is similar to + # the POSIX function `fnmatch()` with the `FNM_PATHNAME` flag set. + + escape = False + regex = '' + i, end = 0, len(pattern) + while i < end: + # Get next character. + char = pattern[i] + i += 1 + + if escape: + # Escape the character. + escape = False + regex += re.escape(char) + + elif char == '\\': + # Escape character, escape next character. + escape = True + + elif char == '*': + # Multi-character wildcard. Match any string (except slashes), + # including an empty string. + regex += '[^/]*' + + elif char == '?': + # Single-character wildcard. Match any single character (except + # a slash). + regex += '[^/]' + + elif char == '[': + # Braket expression wildcard. Except for the beginning + # exclamation mark, the whole braket expression can be used + # directly as regex but we have to find where the expression + # ends. + # - "[][!]" matchs ']', '[' and '!'. + # - "[]-]" matchs ']' and '-'. + # - "[!]a-]" matchs any character except ']', 'a' and '-'. + j = i + # Pass brack expression negation. + if j < end and pattern[j] == '!': + j += 1 + # Pass first closing braket if it is at the beginning of the + # expression. + if j < end and pattern[j] == ']': + j += 1 + # Find closing braket. Stop once we reach the end or find it. + while j < end and pattern[j] != ']': + j += 1 + + if j < end: + # Found end of braket expression. Increment j to be one past + # the closing braket: + # + # [...] + # ^ ^ + # i j + # + j += 1 + expr = '[' + + if pattern[i] == '!': + # Braket expression needs to be negated. + expr += '^' + i += 1 + elif pattern[i] == '^': + # POSIX declares that the regex braket expression negation + # "[^...]" is undefined in a glob pattern. Python's + # `fnmatch.translate()` escapes the caret ('^') as a + # literal. To maintain consistency with undefined behavior, + # I am escaping the '^' as well. + expr += '\\^' + i += 1 + + # Build regex braket expression. Escape slashes so they are + # treated as literal slashes by regex as defined by POSIX. + expr += pattern[i:j].replace('\\', '\\\\') + + # Add regex braket expression to regex result. + regex += expr + + # Set i to one past the closing braket. + i = j + + else: + # Failed to find closing braket, treat opening braket as a + # braket literal instead of as an expression. + regex += '\\[' + + else: + # Regular character, escape it for regex. + regex += re.escape(char) + + return regex + + @staticmethod + def escape(s): + """ + Escape special characters in the given string. + + *s* (:class:`unicode` or :class:`bytes`) a filename or a string + that you want to escape, usually before adding it to a `.gitignore` + + Returns the escaped string (:class:`unicode`, :class:`bytes`) + """ + # Reference: https://git-scm.com/docs/gitignore#_pattern_format + meta_characters = r"[]!*#?" + + return "".join("\\" + x if x in meta_characters else x for x in s) + +util.register_pattern('gitwildmatch', GitWildMatchPattern) + + +class GitIgnorePattern(GitWildMatchPattern): + """ + The :class:`GitIgnorePattern` class is deprecated by :class:`GitWildMatchPattern`. + This class only exists to maintain compatibility with v0.4. + """ + + def __init__(self, *args, **kw): + """ + Warn about deprecation. + """ + self._deprecated() + return super(GitIgnorePattern, self).__init__(*args, **kw) + + @staticmethod + def _deprecated(): + """ + Warn about deprecation. + """ + warnings.warn("GitIgnorePattern ('gitignore') is deprecated. Use GitWildMatchPattern ('gitwildmatch') instead.", DeprecationWarning, stacklevel=3) + + @classmethod + def pattern_to_regex(cls, *args, **kw): + """ + Warn about deprecation. + """ + cls._deprecated() + return super(GitIgnorePattern, cls).pattern_to_regex(*args, **kw) + +# Register `GitIgnorePattern` as "gitignore" for backward compatibility +# with v0.4. +util.register_pattern('gitignore', GitIgnorePattern) diff --git a/lib/python3.12/site-packages/ray/_private/thirdparty/pathspec/util.py b/lib/python3.12/site-packages/ray/_private/thirdparty/pathspec/util.py new file mode 100644 index 0000000000000000000000000000000000000000..bcba8783b61c9613aa4b0dec76a36a488a1b3b60 --- /dev/null +++ b/lib/python3.12/site-packages/ray/_private/thirdparty/pathspec/util.py @@ -0,0 +1,600 @@ +# encoding: utf-8 +""" +This module provides utility methods for dealing with path-specs. +""" + +import os +import os.path +import posixpath +import stat + +from .compat import Collection, Iterable, string_types, unicode + +NORMALIZE_PATH_SEPS = [sep for sep in [os.sep, os.altsep] if sep and sep != posixpath.sep] +""" +*NORMALIZE_PATH_SEPS* (:class:`list` of :class:`str`) contains the path +separators that need to be normalized to the POSIX separator for the +current operating system. The separators are determined by examining +:data:`os.sep` and :data:`os.altsep`. +""" + +_registered_patterns = {} +""" +*_registered_patterns* (:class:`dict`) maps a name (:class:`str`) to the +registered pattern factory (:class:`~collections.abc.Callable`). +""" + + +def detailed_match_files(patterns, files, all_matches=None): + """ + Matches the files to the patterns, and returns which patterns matched + the files. + + *patterns* (:class:`~collections.abc.Iterable` of :class:`~pathspec.pattern.Pattern`) + contains the patterns to use. + + *files* (:class:`~collections.abc.Iterable` of :class:`str`) contains + the normalized file paths to be matched against *patterns*. + + *all_matches* (:class:`boot` or :data:`None`) is whether to return all + matches patterns (:data:`True`), or only the last matched pattern + (:data:`False`). Default is :data:`None` for :data:`False`. + + Returns the matched files (:class:`dict`) which maps each matched file + (:class:`str`) to the patterns that matched in order (:class:`.MatchDetail`). + """ + all_files = files if isinstance(files, Collection) else list(files) + return_files = {} + for pattern in patterns: + if pattern.include is not None: + result_files = pattern.match(all_files) + if pattern.include: + # Add files and record pattern. + for result_file in result_files: + if result_file in return_files: + if all_matches: + return_files[result_file].patterns.append(pattern) + else: + return_files[result_file].patterns[0] = pattern + else: + return_files[result_file] = MatchDetail([pattern]) + + else: + # Remove files. + for file in result_files: + del return_files[file] + + return return_files + + +def _is_iterable(value): + """ + Check whether the value is an iterable (excludes strings). + + *value* is the value to check, + + Returns whether *value* is a iterable (:class:`bool`). + """ + return isinstance(value, Iterable) and not isinstance(value, (unicode, bytes)) + + +def iter_tree_entries(root, on_error=None, follow_links=None): + """ + Walks the specified directory for all files and directories. + + *root* (:class:`str`) is the root directory to search. + + *on_error* (:class:`~collections.abc.Callable` or :data:`None`) + optionally is the error handler for file-system exceptions. It will be + called with the exception (:exc:`OSError`). Reraise the exception to + abort the walk. Default is :data:`None` to ignore file-system + exceptions. + + *follow_links* (:class:`bool` or :data:`None`) optionally is whether + to walk symbolic links that resolve to directories. Default is + :data:`None` for :data:`True`. + + Raises :exc:`RecursionError` if recursion is detected. + + Returns an :class:`~collections.abc.Iterable` yielding each file or + directory entry (:class:`.TreeEntry`) relative to *root*. + """ + if on_error is not None and not callable(on_error): + raise TypeError("on_error:{!r} is not callable.".format(on_error)) + + if follow_links is None: + follow_links = True + + for entry in _iter_tree_entries_next(os.path.abspath(root), '', {}, on_error, follow_links): + yield entry + + +def iter_tree_files(root, on_error=None, follow_links=None): + """ + Walks the specified directory for all files. + + *root* (:class:`str`) is the root directory to search for files. + + *on_error* (:class:`~collections.abc.Callable` or :data:`None`) + optionally is the error handler for file-system exceptions. It will be + called with the exception (:exc:`OSError`). Reraise the exception to + abort the walk. Default is :data:`None` to ignore file-system + exceptions. + + *follow_links* (:class:`bool` or :data:`None`) optionally is whether + to walk symbolic links that resolve to directories. Default is + :data:`None` for :data:`True`. + + Raises :exc:`RecursionError` if recursion is detected. + + Returns an :class:`~collections.abc.Iterable` yielding the path to + each file (:class:`str`) relative to *root*. + """ + if on_error is not None and not callable(on_error): + raise TypeError("on_error:{!r} is not callable.".format(on_error)) + + if follow_links is None: + follow_links = True + + for entry in _iter_tree_entries_next(os.path.abspath(root), '', {}, on_error, follow_links): + if not entry.is_dir(follow_links): + yield entry.path + + +# Alias `iter_tree_files()` as `iter_tree()`. +iter_tree = iter_tree_files + + +def _iter_tree_entries_next(root_full, dir_rel, memo, on_error, follow_links): + """ + Scan the directory for all descendant files. + + *root_full* (:class:`str`) the absolute path to the root directory. + + *dir_rel* (:class:`str`) the path to the directory to scan relative to + *root_full*. + + *memo* (:class:`dict`) keeps track of ancestor directories + encountered. Maps each ancestor real path (:class:`str`) to relative + path (:class:`str`). + + *on_error* (:class:`~collections.abc.Callable` or :data:`None`) + optionally is the error handler for file-system exceptions. + + *follow_links* (:class:`bool`) is whether to walk symbolic links that + resolve to directories. + + Yields each entry (:class:`.TreeEntry`). + """ + dir_full = os.path.join(root_full, dir_rel) + dir_real = os.path.realpath(dir_full) + + # Remember each encountered ancestor directory and its canonical + # (real) path. If a canonical path is encountered more than once, + # recursion has occurred. + if dir_real not in memo: + memo[dir_real] = dir_rel + else: + raise RecursionError(real_path=dir_real, first_path=memo[dir_real], second_path=dir_rel) + + for node_name in os.listdir(dir_full): + node_rel = os.path.join(dir_rel, node_name) + node_full = os.path.join(root_full, node_rel) + + # Inspect child node. + try: + node_lstat = os.lstat(node_full) + except OSError as e: + if on_error is not None: + on_error(e) + continue + + if stat.S_ISLNK(node_lstat.st_mode): + # Child node is a link, inspect the target node. + is_link = True + try: + node_stat = os.stat(node_full) + except OSError as e: + if on_error is not None: + on_error(e) + continue + else: + is_link = False + node_stat = node_lstat + + if stat.S_ISDIR(node_stat.st_mode) and (follow_links or not is_link): + # Child node is a directory, recurse into it and yield its + # descendant files. + yield TreeEntry(node_name, node_rel, node_lstat, node_stat) + + for entry in _iter_tree_entries_next(root_full, node_rel, memo, on_error, follow_links): + yield entry + + elif stat.S_ISREG(node_stat.st_mode) or is_link: + # Child node is either a file or an unfollowed link, yield it. + yield TreeEntry(node_name, node_rel, node_lstat, node_stat) + + # NOTE: Make sure to remove the canonical (real) path of the directory + # from the ancestors memo once we are done with it. This allows the + # same directory to appear multiple times. If this is not done, the + # second occurrence of the directory will be incorrectly interpreted + # as a recursion. See . + del memo[dir_real] + + +def lookup_pattern(name): + """ + Lookups a registered pattern factory by name. + + *name* (:class:`str`) is the name of the pattern factory. + + Returns the registered pattern factory (:class:`~collections.abc.Callable`). + If no pattern factory is registered, raises :exc:`KeyError`. + """ + return _registered_patterns[name] + + +def match_file(patterns, file): + """ + Matches the file to the patterns. + + *patterns* (:class:`~collections.abc.Iterable` of :class:`~pathspec.pattern.Pattern`) + contains the patterns to use. + + *file* (:class:`str`) is the normalized file path to be matched + against *patterns*. + + Returns :data:`True` if *file* matched; otherwise, :data:`False`. + """ + matched = False + for pattern in patterns: + if pattern.include is not None: + if file in pattern.match((file,)): + matched = pattern.include + return matched + + +def match_files(patterns, files): + """ + Matches the files to the patterns. + + *patterns* (:class:`~collections.abc.Iterable` of :class:`~pathspec.pattern.Pattern`) + contains the patterns to use. + + *files* (:class:`~collections.abc.Iterable` of :class:`str`) contains + the normalized file paths to be matched against *patterns*. + + Returns the matched files (:class:`set` of :class:`str`). + """ + all_files = files if isinstance(files, Collection) else list(files) + return_files = set() + for pattern in patterns: + if pattern.include is not None: + result_files = pattern.match(all_files) + if pattern.include: + return_files.update(result_files) + else: + return_files.difference_update(result_files) + return return_files + + +def _normalize_entries(entries, separators=None): + """ + Normalizes the entry paths to use the POSIX path separator. + + *entries* (:class:`~collections.abc.Iterable` of :class:`.TreeEntry`) + contains the entries to be normalized. + + *separators* (:class:`~collections.abc.Collection` of :class:`str`; or + :data:`None`) optionally contains the path separators to normalize. + See :func:`normalize_file` for more information. + + Returns a :class:`dict` mapping the each normalized file path (:class:`str`) + to the entry (:class:`.TreeEntry`) + """ + norm_files = {} + for entry in entries: + norm_files[normalize_file(entry.path, separators=separators)] = entry + return norm_files + + +def normalize_file(file, separators=None): + """ + Normalizes the file path to use the POSIX path separator (i.e., ``'/'``). + + *file* (:class:`str` or :class:`pathlib.PurePath`) is the file path. + + *separators* (:class:`~collections.abc.Collection` of :class:`str`; or + :data:`None`) optionally contains the path separators to normalize. + This does not need to include the POSIX path separator (``'/'``), but + including it will not affect the results. Default is :data:`None` for + :data:`NORMALIZE_PATH_SEPS`. To prevent normalization, pass an empty + container (e.g., an empty tuple ``()``). + + Returns the normalized file path (:class:`str`). + """ + # Normalize path separators. + if separators is None: + separators = NORMALIZE_PATH_SEPS + + # Convert path object to string. + norm_file = str(file) + + for sep in separators: + norm_file = norm_file.replace(sep, posixpath.sep) + + # Remove current directory prefix. + if norm_file.startswith('./'): + norm_file = norm_file[2:] + + return norm_file + + +def normalize_files(files, separators=None): + """ + Normalizes the file paths to use the POSIX path separator. + + *files* (:class:`~collections.abc.Iterable` of :class:`str` or + :class:`pathlib.PurePath`) contains the file paths to be normalized. + + *separators* (:class:`~collections.abc.Collection` of :class:`str`; or + :data:`None`) optionally contains the path separators to normalize. + See :func:`normalize_file` for more information. + + Returns a :class:`dict` mapping the each normalized file path (:class:`str`) + to the original file path (:class:`str`) + """ + norm_files = {} + for path in files: + norm_files[normalize_file(path, separators=separators)] = path + return norm_files + + +def register_pattern(name, pattern_factory, override=None): + """ + Registers the specified pattern factory. + + *name* (:class:`str`) is the name to register the pattern factory + under. + + *pattern_factory* (:class:`~collections.abc.Callable`) is used to + compile patterns. It must accept an uncompiled pattern (:class:`str`) + and return the compiled pattern (:class:`.Pattern`). + + *override* (:class:`bool` or :data:`None`) optionally is whether to + allow overriding an already registered pattern under the same name + (:data:`True`), instead of raising an :exc:`AlreadyRegisteredError` + (:data:`False`). Default is :data:`None` for :data:`False`. + """ + if not isinstance(name, string_types): + raise TypeError("name:{!r} is not a string.".format(name)) + if not callable(pattern_factory): + raise TypeError("pattern_factory:{!r} is not callable.".format(pattern_factory)) + if name in _registered_patterns and not override: + raise AlreadyRegisteredError(name, _registered_patterns[name]) + _registered_patterns[name] = pattern_factory + + +class AlreadyRegisteredError(Exception): + """ + The :exc:`AlreadyRegisteredError` exception is raised when a pattern + factory is registered under a name already in use. + """ + + def __init__(self, name, pattern_factory): + """ + Initializes the :exc:`AlreadyRegisteredError` instance. + + *name* (:class:`str`) is the name of the registered pattern. + + *pattern_factory* (:class:`~collections.abc.Callable`) is the + registered pattern factory. + """ + super(AlreadyRegisteredError, self).__init__(name, pattern_factory) + + @property + def message(self): + """ + *message* (:class:`str`) is the error message. + """ + return "{name!r} is already registered for pattern factory:{pattern_factory!r}.".format( + name=self.name, + pattern_factory=self.pattern_factory, + ) + + @property + def name(self): + """ + *name* (:class:`str`) is the name of the registered pattern. + """ + return self.args[0] + + @property + def pattern_factory(self): + """ + *pattern_factory* (:class:`~collections.abc.Callable`) is the + registered pattern factory. + """ + return self.args[1] + + +class RecursionError(Exception): + """ + The :exc:`RecursionError` exception is raised when recursion is + detected. + """ + + def __init__(self, real_path, first_path, second_path): + """ + Initializes the :exc:`RecursionError` instance. + + *real_path* (:class:`str`) is the real path that recursion was + encountered on. + + *first_path* (:class:`str`) is the first path encountered for + *real_path*. + + *second_path* (:class:`str`) is the second path encountered for + *real_path*. + """ + super(RecursionError, self).__init__(real_path, first_path, second_path) + + @property + def first_path(self): + """ + *first_path* (:class:`str`) is the first path encountered for + :attr:`self.real_path `. + """ + return self.args[1] + + @property + def message(self): + """ + *message* (:class:`str`) is the error message. + """ + return "Real path {real!r} was encountered at {first!r} and then {second!r}.".format( + real=self.real_path, + first=self.first_path, + second=self.second_path, + ) + + @property + def real_path(self): + """ + *real_path* (:class:`str`) is the real path that recursion was + encountered on. + """ + return self.args[0] + + @property + def second_path(self): + """ + *second_path* (:class:`str`) is the second path encountered for + :attr:`self.real_path `. + """ + return self.args[2] + + +class MatchDetail(object): + """ + The :class:`.MatchDetail` class contains information about + """ + + #: Make the class dict-less. + __slots__ = ('patterns',) + + def __init__(self, patterns): + """ + Initialize the :class:`.MatchDetail` instance. + + *patterns* (:class:`~collections.abc.Sequence` of :class:`~pathspec.pattern.Pattern`) + contains the patterns that matched the file in the order they were + encountered. + """ + + self.patterns = patterns + """ + *patterns* (:class:`~collections.abc.Sequence` of :class:`~pathspec.pattern.Pattern`) + contains the patterns that matched the file in the order they were + encountered. + """ + + +class TreeEntry(object): + """ + The :class:`.TreeEntry` class contains information about a file-system + entry. + """ + + #: Make the class dict-less. + __slots__ = ('_lstat', 'name', 'path', '_stat') + + def __init__(self, name, path, lstat, stat): + """ + Initialize the :class:`.TreeEntry` instance. + + *name* (:class:`str`) is the base name of the entry. + + *path* (:class:`str`) is the relative path of the entry. + + *lstat* (:class:`~os.stat_result`) is the stat result of the direct + entry. + + *stat* (:class:`~os.stat_result`) is the stat result of the entry, + potentially linked. + """ + + self._lstat = lstat + """ + *_lstat* (:class:`~os.stat_result`) is the stat result of the direct + entry. + """ + + self.name = name + """ + *name* (:class:`str`) is the base name of the entry. + """ + + self.path = path + """ + *path* (:class:`str`) is the path of the entry. + """ + + self._stat = stat + """ + *_stat* (:class:`~os.stat_result`) is the stat result of the linked + entry. + """ + + def is_dir(self, follow_links=None): + """ + Get whether the entry is a directory. + + *follow_links* (:class:`bool` or :data:`None`) is whether to follow + symbolic links. If this is :data:`True`, a symlink to a directory + will result in :data:`True`. Default is :data:`None` for :data:`True`. + + Returns whether the entry is a directory (:class:`bool`). + """ + if follow_links is None: + follow_links = True + + node_stat = self._stat if follow_links else self._lstat + return stat.S_ISDIR(node_stat.st_mode) + + def is_file(self, follow_links=None): + """ + Get whether the entry is a regular file. + + *follow_links* (:class:`bool` or :data:`None`) is whether to follow + symbolic links. If this is :data:`True`, a symlink to a regular file + will result in :data:`True`. Default is :data:`None` for :data:`True`. + + Returns whether the entry is a regular file (:class:`bool`). + """ + if follow_links is None: + follow_links = True + + node_stat = self._stat if follow_links else self._lstat + return stat.S_ISREG(node_stat.st_mode) + + def is_symlink(self): + """ + Returns whether the entry is a symbolic link (:class:`bool`). + """ + return stat.S_ISLNK(self._lstat.st_mode) + + def stat(self, follow_links=None): + """ + Get the cached stat result for the entry. + + *follow_links* (:class:`bool` or :data:`None`) is whether to follow + symbolic links. If this is :data:`True`, the stat result of the + linked file will be returned. Default is :data:`None` for :data:`True`. + + Returns that stat result (:class:`~os.stat_result`). + """ + if follow_links is None: + follow_links = True + + return self._stat if follow_links else self._lstat diff --git a/lib/python3.12/site-packages/ray/_private/thirdparty/pyamdsmi/__init__.py b/lib/python3.12/site-packages/ray/_private/thirdparty/pyamdsmi/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..8af4746136833fd8fc591e122ed9efb8117018f6 --- /dev/null +++ b/lib/python3.12/site-packages/ray/_private/thirdparty/pyamdsmi/__init__.py @@ -0,0 +1 @@ +from ray._private.thirdparty.pyamdsmi.pyamdsmi import * diff --git a/lib/python3.12/site-packages/ray/_private/thirdparty/pyamdsmi/__pycache__/__init__.cpython-312.pyc b/lib/python3.12/site-packages/ray/_private/thirdparty/pyamdsmi/__pycache__/__init__.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..73082965b43340b419a20faf0f15ac2dbb6f129c Binary files /dev/null and b/lib/python3.12/site-packages/ray/_private/thirdparty/pyamdsmi/__pycache__/__init__.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/_private/thirdparty/pyamdsmi/__pycache__/pyamdsmi.cpython-312.pyc b/lib/python3.12/site-packages/ray/_private/thirdparty/pyamdsmi/__pycache__/pyamdsmi.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..1dd77d00fc5c6dcc18507cd0196ba01e4a0a1310 Binary files /dev/null and b/lib/python3.12/site-packages/ray/_private/thirdparty/pyamdsmi/__pycache__/pyamdsmi.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/_private/thirdparty/pyamdsmi/pyamdsmi.py b/lib/python3.12/site-packages/ray/_private/thirdparty/pyamdsmi/pyamdsmi.py new file mode 100644 index 0000000000000000000000000000000000000000..e28b8fe81eb00350a980613a8484e6684f1dab17 --- /dev/null +++ b/lib/python3.12/site-packages/ray/_private/thirdparty/pyamdsmi/pyamdsmi.py @@ -0,0 +1,568 @@ +# MIT License +# +# Copyright (c) 2023 Advanced Micro Devices, Inc. +# +# Permission is hereby granted, free of charge, to any person obtaining a copy +# of this software and associated documentation files (the "Software"), to deal +# in the Software without restriction, including without limitation the rights +# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +# copies of the Software, and to permit persons to whom the Software is +# furnished to do so, subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in all +# copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE +# SOFTWARE. + +# Python bindings for ROCm-SMI library +from ctypes import * +from os.path import join, realpath, isfile +import os +import logging +import subprocess +import sys +import threading +from enum import IntEnum, auto + + + +## Error checking +class ROCMLError_NotSupported(Exception): + pass + + +class ROCMLError_FunctionNotFound(Exception): + pass + + +class ROCMLError_LibraryNotFound(Exception): + pass + + +class ROCMLError_DriverNotLoaded(Exception): + pass + + +class ROCMLError_Unknown(Exception): + pass + + +class ROCMLError_Uninitialized(Exception): + pass + + +class ROCMLState(IntEnum): + UNINITIALIZED = auto() + """No attempt yet made to initialize PyROCML""" + INITIALIZED = auto() + """PyROCML was successfully initialized""" + DISABLED_PYROCML_NOT_AVAILABLE = auto() + """PyROCML not installed""" + DISABLED_CONFIG = auto() + """PyROCML diagnostics disabled by ``distributed.diagnostics.rocml`` config setting""" + DISABLED_LIBRARY_NOT_FOUND = auto() + """PyROCML available, but ROCML not installed""" + + +LIBROCM_NAME = 'librocm_smi64.so' +RSMI_MAX_BUFFER_LENGTH = 256 + +# Policy enums +RSMI_MAX_NUM_FREQUENCIES = 32 + + +class rsmi_status_t(c_int): + RSMI_STATUS_SUCCESS = 0x0 + RSMI_STATUS_INVALID_ARGS = 0x1 + RSMI_STATUS_NOT_SUPPORTED = 0x2 + RSMI_STATUS_FILE_ERROR = 0x3 + RSMI_STATUS_PERMISSION = 0x4 + RSMI_STATUS_OUT_OF_RESOURCES = 0x5 + RSMI_STATUS_INTERNAL_EXCEPTION = 0x6 + RSMI_STATUS_INPUT_OUT_OF_BOUNDS = 0x7 + RSMI_STATUS_INIT_ERROR = 0x8 + RSMI_INITIALIZATION_ERROR = RSMI_STATUS_INIT_ERROR + RSMI_STATUS_NOT_YET_IMPLEMENTED = 0x9 + RSMI_STATUS_NOT_FOUND = 0xA + RSMI_STATUS_INSUFFICIENT_SIZE = 0xB + RSMI_STATUS_INTERRUPT = 0xC + RSMI_STATUS_UNEXPECTED_SIZE = 0xD + RSMI_STATUS_NO_DATA = 0xE + RSMI_STATUS_UNKNOWN_ERROR = 0xFFFFFFFF + + +#Dictionary of rsmi ret codes and it's verbose output +rsmi_status_verbose_err_out = { + rsmi_status_t.RSMI_STATUS_SUCCESS: 'Operation was successful', + rsmi_status_t.RSMI_STATUS_INVALID_ARGS: 'Invalid arguments provided', + rsmi_status_t.RSMI_STATUS_NOT_SUPPORTED: 'Not supported on the given system', + rsmi_status_t.RSMI_STATUS_FILE_ERROR: 'Problem accessing a file', + rsmi_status_t.RSMI_STATUS_PERMISSION: 'Permission denied', + rsmi_status_t.RSMI_STATUS_OUT_OF_RESOURCES: 'Unable to acquire memory or other resource', + rsmi_status_t.RSMI_STATUS_INTERNAL_EXCEPTION: 'An internal exception was caught', + rsmi_status_t.RSMI_STATUS_INPUT_OUT_OF_BOUNDS: 'Provided input is out of allowable or safe range', + rsmi_status_t.RSMI_INITIALIZATION_ERROR: 'Error occurred during rsmi initialization', + rsmi_status_t.RSMI_STATUS_NOT_YET_IMPLEMENTED: 'Requested function is not implemented on this setup', + rsmi_status_t.RSMI_STATUS_NOT_FOUND: 'Item searched for but not found', + rsmi_status_t.RSMI_STATUS_INSUFFICIENT_SIZE: 'Insufficient resources available', + rsmi_status_t.RSMI_STATUS_INTERRUPT: 'Interrupt occurred during execution', + rsmi_status_t.RSMI_STATUS_UNEXPECTED_SIZE: 'Unexpected amount of data read', + rsmi_status_t.RSMI_STATUS_NO_DATA: 'No data found for the given input', + rsmi_status_t.RSMI_STATUS_UNKNOWN_ERROR: 'Unknown error occurred' +} + + +class rsmi_init_flags_t(c_int): + RSMI_INIT_FLAG_ALL_GPUS = 0x1 + + +class rsmi_memory_type_t(c_int): + RSMI_MEM_TYPE_FIRST = 0 + RSMI_MEM_TYPE_VRAM = RSMI_MEM_TYPE_FIRST + RSMI_MEM_TYPE_VIS_VRAM = 1 + RSMI_MEM_TYPE_GTT = 2 + RSMI_MEM_TYPE_LAST = RSMI_MEM_TYPE_GTT + + +# memory_type_l includes names for with rsmi_memory_type_t +# Usage example to get corresponding names: +# memory_type_l[rsmi_memory_type_t.RSMI_MEM_TYPE_VRAM] will return string 'vram' +memory_type_l = ['VRAM', 'VIS_VRAM', 'GTT'] + + +class rsmi_retired_page_record_t(Structure): + _fields_ = [('page_address', c_uint64), + ('page_size', c_uint64), + ('status', c_int)] + + +class rsmi_sw_component_t(c_int): + RSMI_SW_COMP_FIRST = 0x0 + RSMI_SW_COMP_DRIVER = RSMI_SW_COMP_FIRST + RSMI_SW_COMP_LAST = RSMI_SW_COMP_DRIVER + + +class rsmi_frequencies_t(Structure): + _fields_ = [('num_supported', c_int32), + ('current', c_uint32), + ('frequency', c_uint64 * RSMI_MAX_NUM_FREQUENCIES)] + + +class rsmi_pcie_bandwidth_t(Structure): + _fields_ = [('transfer_rate', rsmi_frequencies_t), + ('lanes', c_uint32 * RSMI_MAX_NUM_FREQUENCIES)] + + +class rsmi_process_info_t(Structure): + _fields_ = [('process_id', c_uint32), + ('pasid', c_uint32), # PSA: Power Spectrum Analysis ? + ('vram_usage', c_uint64), + ('sdma_usage', c_uint64), # SDMA: System Direct Memory Access + ('cu_occupancy', c_uint32)] + + +class rsmi_xgmi_status_t(c_int): + RSMI_XGMI_STATUS_NO_ERRORS = 0 + RSMI_XGMI_STATUS_ERROR = 1 + RSMI_XGMI_STATUS_MULTIPLE_ERRORS = 2 + + +class rsmi_io_link_type(c_int): + RSMI_IOLINK_TYPE_UNDEFINED = 0 + RSMI_IOLINK_TYPE_HYPERTRANSPORT = 1 + RSMI_IOLINK_TYPE_PCIEXPRESS = 2 + RSMI_IOLINK_TYPE_AMBA = 3 + RSMI_IOLINK_TYPE_MIPI = 4 + RSMI_IOLINK_TYPE_QPI_1_1 = 5 + RSMI_IOLINK_TYPE_RESERVED1 = 6 + RSMI_IOLINK_TYPE_RESERVED2 = 7 + RSMI_IOLINK_TYPE_RAPID_IO = 8 + RSMI_IOLINK_TYPE_INFINIBAND = 9 + RSMI_IOLINK_TYPE_RESERVED3 = 10 + RSMI_IOLINK_TYPE_XGMI = 11 + RSMI_IOLINK_TYPE_XGOP = 12 + RSMI_IOLINK_TYPE_GZ = 13 + RSMI_IOLINK_TYPE_ETHERNET_RDMA = 14 + RSMI_IOLINK_TYPE_RDMA_OTHER = 15 + RSMI_IOLINK_TYPE_OTHER = 16 + RSMI_IOLINK_TYPE_NUMIOLINKTYPES = 17 + RSMI_IOLINK_TYPE_SIZE = 0xFFFFFFFF + + +## Library loading +rocm_lib = None +lib_load_lock = threading.Lock() +_rocm_lib_refcount = 0 + + +## Function access, to prevent lib_load_lock deadlock +_rocml_get_function_ptr_cache = dict() + +def _rocml_get_function_ptr(name): + global rocm_lib + + if name in _rocml_get_function_ptr_cache: + return _rocml_get_function_ptr_cache[name] + + lib_load_lock.acquire() + try: + # ensure library was loaded + if rocm_lib == None: + raise ROCMLError_Uninitialized + try: + _rocml_get_function_ptr_cache[name] = getattr(rocm_lib, name) + return _rocml_get_function_ptr_cache[name] + except AttributeError: + raise ROCMLError_FunctionNotFound + finally: + # lock is always freed + lib_load_lock.release() + + +def _load_rocm_library(): + """Load ROCm library if not already loaded""" + global rocm_lib + + if rocm_lib == None: + + lib_load_lock.acquire() + + try: + if rocm_lib == None: + try: + if sys.platform[:3] == 'win': + raise ROCMLError_NotSupported('Windows platform is not supported yet') + else: + # assume linux + path_librocm = _find_lib_rocm() + cdll.LoadLibrary(path_librocm) + rocm_lib = CDLL(path_librocm) + except OSError: + raise ROCMLError_LibraryNotFound('ROCm library not found') + if rocm_lib == None: + raise ROCMLError_LibraryNotFound('ROCm library not found') + finally: + lib_load_lock.release() + + +def _find_lib_rocm(): + """search for librocm and returns path + if search fails, returns empty string + """ + rocm_path = os.environ.get('ROCM_PATH', '/opt/rocm') + rocm_lib_path = join(rocm_path, f'lib/{LIBROCM_NAME}') + return rocm_lib_path if isfile(rocm_lib_path) else '' + + +def _driver_initialized(): + """ Returns true if amdgpu is found in the list of initialized modules + """ + initialized = '' + try: + initialized = str(subprocess.check_output("cat /sys/module/amdgpu/initstate |grep live", shell=True, stderr=subprocess.DEVNULL)) + except subprocess.CalledProcessError: + pass + return len(initialized) > 0 + + +def smi_initialize(): + """Initialize ROCm binding of SMI""" + _load_rocm_library() + + if _driver_initialized(): + ret_init = rocm_lib.rsmi_init(0) + if ret_init != 0: + logging.debug("ROCm SMI init returned value: %s", ret_init) + raise RuntimeError('ROCm SMI initialization failed') + else: + raise RuntimeError('ROCm driver initilization failed') + + # update reference count + global _rocm_lib_refcount + lib_load_lock.acquire() + _rocm_lib_refcount += 1 + lib_load_lock.release() + + +def rsmi_ret_ok(my_ret): + """ Returns true if RSMI call status is 0 (success) + + @param device: DRM device identifier + @param my_ret: Return of RSMI call (rocm_smi_lib API) + @param metric: Parameter of GPU currently being analyzed + """ + if my_ret != rsmi_status_t.RSMI_STATUS_SUCCESS: + err_str = c_char_p() + rocm_lib.rsmi_status_string(my_ret, byref(err_str)) + logging.debug("ROCm RSMI error: %s", err_str.value.decode()) + return False + return True + + +def smi_shutdown(): + """leave the library loaded, but shutdown the interface""" + rsmi_ret_ok(rocm_lib.rsmi_shut_down()) + + # update reference count + global _rocm_lib_refcount + lib_load_lock.acquire() + _rocm_lib_refcount -= 1 + lib_load_lock.release() + + +def smi_get_kernel_version(): + """returns ROCm kernerl driver version""" + ver_str = create_string_buffer(256) + ret = rocm_lib.rsmi_version_str_get(rsmi_sw_component_t.RSMI_SW_COMP_DRIVER, ver_str, 256) + return ver_str.value.decode() if rsmi_ret_ok(ret) else '' + +def smi_get_device_id(dev): + """returns device id of the device as 64bit integer""" + uid = c_uint64() + ret = rocm_lib.rsmi_dev_id_get(dev, byref(uid)) + return uid.value if rsmi_ret_ok(ret) else -1 + +def smi_get_device_count(): + """returns a list of GPU devices """ + num_device = c_uint32(0) + ret = rocm_lib.rsmi_num_monitor_devices(byref(num_device)) + return num_device.value if rsmi_ret_ok(ret) else -1 + + +def smi_get_device_name(dev): + """returns the name of a GPU device""" + series = create_string_buffer(RSMI_MAX_BUFFER_LENGTH) + ret = rocm_lib.rsmi_dev_name_get(dev, series, RSMI_MAX_BUFFER_LENGTH) + return series.value.decode() if rsmi_ret_ok(ret) else '' + + +def smi_get_device_unique_id(dev): + """returns unique id of the device as 64bit integer""" + uid = c_uint64() + ret = rocm_lib.rsmi_dev_unique_id_get(dev, byref(uid)) + return uid.value if rsmi_ret_ok(ret) else -1 + +def smi_get_device_utilization(dev): + """returns GPU device busy percent of device_id dev""" + busy_percent = c_uint32() + ret = rocm_lib.rsmi_dev_busy_percent_get(dev, byref(busy_percent)) + return busy_percent.value if rsmi_ret_ok(ret) else -1 + + +def smi_get_device_memory_used(dev, type='VRAM'): + """returns used memory of device_id dev in bytes""" + type_idx = memory_type_l.index(type) + used = c_uint64() + ret = rocm_lib.rsmi_dev_memory_usage_get(dev, type_idx, byref(used)) + return used.value if rsmi_ret_ok(ret) else -1 + + +def smi_get_device_memory_total(dev, type='VRAM'): + """returns total memory of device_id dev in bytes""" + type_idx = memory_type_l.index(type) + total = c_uint64() + ret = rocm_lib.rsmi_dev_memory_total_get(dev, type_idx, byref(total)) + return total.value if rsmi_ret_ok(ret) else -1 + + +def smi_get_device_memory_busy(dev): + """returns percentage of time any device memory is being used""" + busy_percent = c_uint32() + ret = rocm_lib.rsmi_dev_memory_busy_percent_get(dev, byref(busy_percent)) + return busy_percent.value if rsmi_ret_ok(ret) else -1 + + +def smi_get_device_memory_reserved_pages(dev): + """returns info about reserved memory pages""" + num_pages = c_uint32() + records = rsmi_retired_page_record_t() + ret = rocm_lib.rsmi_dev_memory_reserved_pages_get(dev, byref(num_pages), byref(records)) + return (num_pages.value, records) if rsmi_ret_ok(ret) else -1 + + +# PCIE functions +def smi_get_device_pcie_bandwidth(dev): + """returns list of possible pcie bandwidths for the device in bytes/sec""" + bandwidth = rsmi_pcie_bandwidth_t() + ret = rocm_lib.rsmi_dev_pci_bandwidth_get(dev, byref(bandwidth)) + return bandwidth if rsmi_ret_ok(ret) else -1 + + +def smi_get_device_pci_id(dev): + """returns unique PCI ID of the device in 64bit Hex with format: + BDFID = ((DOMAIN & 0xffffffff) << 32) | ((BUS & 0xff) << 8) | + ((DEVICE & 0x1f) <<3 ) | (FUNCTION & 0x7) + """ + bdfid = c_uint64() + ret = rocm_lib.rsmi_dev_pci_id_get(dev, byref(bdfid)) + return bdfid.value if rsmi_ret_ok(ret) else -1 + + +def smi_get_device_topo_numa_affinity(dev): + """returns the NUMA node associated with the device""" + numa_node = c_uint32() + ret = reocm_lib.rsmi_topo_numa_affinity_get(dev, byref(numa_node)) + return numa_node.value if rsmi_ret_ok(ret) else -1 + + +def smi_get_device_pcie_throughput(dev): + """returns measured pcie throughput for the device in bytes/sec""" + sent = c_uint64() + recv = c_uint64() + max_pkt_sz = c_uint64() + ret = rocm_lib.rsmi_dev_pci_throughput_get(dev, byref(sent), byref(recv), byref(max_pkt_sz)) + return (recv.value + sent.value) * max_pkt_sz.value if rsmi_ret_ok(ret) else -1 + + +def smi_get_device_pci_replay_counter(dev): + """return PCIe replay counter of the device""" + counter = c_uint64() + ret = rocm_lib.rsmi_dev_pci_replay_counter_get(dev, byref(counter)) + return counter.value if rsmi_ret_ok(ret) else -1 + + +# Compute partition functions +def smi_get_device_compute_partition(dev): + """returns the compute partition of the device""" + partition = create_string_buffer(RSMI_MAX_BUFFER_LENGTH) + ret = rocm_lib.rsmi_dev_compute_partition_get(dev, byref(partition), RSMI_MAX_BUFFER_LENGTH) + return partition.value.decode() if rsmi_ret_ok(ret) else '' + + +def smi_set_device_compute_partition(dev, partition): + """modifies the compute partition of the selected device""" + ret = rocm_lib.rsmi_dev_compute_partition_set(dev, partition) + return rsmi_ret_ok(ret) + + +def smi_reset_device_compute_partition(dev): + """reverts the compute partition of the selected device to its boot state""" + ret = rocm_lib.rsmi_dev_compute_partition_reset(dev) + return rsmi_ret_ok(ret) + + +# Memory partition functions +def smi_get_device_memory_partition(dev): + """returns the memory partition of the device""" + partition = create_string_buffer(RSMI_MAX_BUFFER_LENGTH) + ret = rocm_lib.rsmi_dev_memory_partition_get(dev, byref(partition), RSMI_MAX_BUFFER_LENGTH) + return partition.value.decode() if rsmi_ret_ok(ret) else '' + + +def smi_set_device_memory_partition(dev, partition): + """modifies the memory partition of the selected device""" + ret = rocm_lib.rsmi_dev_memory_partition_set(dev, partition) + return rsmi_ret_ok(ret) + + +def smi_reset_device_memory_partition(dev): + """reverts the memory partition of the selected device to its boot state""" + ret = rocm_lib.rsmi_dev_memory_partition_reset(dev) + return rsmi_ret_ok(ret) + + +# Hardware Topology functions +def smi_get_device_topo_numa_node_number(dev): + """returns the NUMA node associated with the device""" + numa_node = c_uint32() + ret = rocm_lib.rsmi_topo_get_numa_node_number(dev, byref(numa_node)) + return numa_node.value if rsmi_ret_ok(ret) else -1 + + +def smi_get_device_topo_link_weight(dev_src, dev_dst): + """returns the weight of the link between two devices""" + weight = c_uint64() + ret = rocm_lib.rsmi_topo_get_link_weight(dev_src, dev_dst, byref(weight)) + return weight.value if rsmi_ret_ok(ret) else -1 + + +def smi_get_device_minmax_bandwidth(dev_src, dev_dst): + """returns the minimum and maximum io link bandwidth between two devices + API works if src and dst are connected via XGMI and are 1 hop away. + """ + assert smi_get_device_link_type(dev_src, dev_dst)[0] == 1, 'Devices must be 1 hop away' + min_bandwidth = c_uint64() + max_bandwidth = c_uint64() + ret = rocm_lib.rsmi_minmax_bandwidth_get(dev_src, dev_dst, byref(min_bandwidth), byref(max_bandwidth)) + return (min_bandwidth.value, max_bandwidth.value) if rsmi_ret_ok(ret) else -1 + + +def smi_get_device_link_type(dev_src, dev_dst): + """returns the hops and the type of link between two devices""" + hops = c_uint64() + link_type = rsmi_io_link_type() + ret = rocm_lib.rsmi_topo_get_link_type(dev_src, dev_dst, byref(hops), byref(link_type)) + return (hops.value, link_type.value) if rsmi_ret_ok(ret) else -1 + + +def smi_is_device_p2p_accessible(dev_src, dev_dst): + """returns true if two devices are p2p accessible""" + accessible = c_bool() + ret = rocm_lib.rsmi_is_P2P_accessible(dev_src, dev_dst, byref(accessible)) + return accessible.value if rsmi_ret_ok(ret) else -1 + + +def smi_get_device_compute_process(): + """returns list of process ids running compute on the system""" + num_procs = c_uint32() + ret = rocm_lib.rsmi_compute_process_info_get(None, byref(num_procs)) + if rsmi_ret_ok(ret): + buff_sz = num_procs.value + 10 + proc_info = (rsmi_process_info_t * buff_sz)() + ret2 = rocm_lib.rsmi_compute_process_info_get(byref(proc_info), byref(num_procs)) + + return [proc_info[i].process_id for i in range(num_procs.value)] if rsmi_ret_ok(ret2) else [] + else: + return [] + +def smi_get_compute_process_info_by_device(device_id: int, proc_ids: list) -> list: + """Returns list of process info running compute on the specified device by process IDs. + + Args: + device_id: The device index to query + proc_ids: List of process IDs to get info for + + Returns: + List of process info structures for the specified device and process IDs + """ + proc_infos = [] + for proc_id in proc_ids: + proc_info = rsmi_process_info_t() + ret = rocm_lib.rsmi_compute_process_info_by_device_get(proc_id, device_id, byref(proc_info)) + if rsmi_ret_ok(ret): + proc_infos.append(proc_info) + return proc_infos + +def smi_get_device_average_power(dev): + """returns average power of device_id dev""" + power = c_uint32() + ret = rocm_lib.rsmi_dev_power_ave_get(dev, 0, byref(power)) + + return power.value * 1e-6 if rsmi_ret_ok(ret) else -1 + + +# XGMI fuctions +def smi_get_device_xgmi_error_status(dev): + """returns XGMI error status for a device""" + status = rsmi_xgmi_status_t() + ret = rocm_lib.rsmi_dev_xgmi_error_status(dev, byref(status)) + return status.value if rsmi_ret_ok(ret) else -1 + + +def smi_reset_device_xgmi_error(dev): + """resets XGMI error status for a device""" + ret = rocm_lib.rsmi_dev_xgmi_error_reset(dev) + return rsmi_ret_ok(ret) + + +def smi_get_device_xgmi_hive_id(dev): + """returns XGMI hive ID for a device""" + hive_id = c_uint64() + ret = rocm_lib.rsmi_dev_xgmi_hive_id_get(dev, byref(hive_id)) + return hive_id.value if rsmi_ret_ok(ret) else -1 diff --git a/lib/python3.12/site-packages/ray/_private/thirdparty/pynvml/__init__.py b/lib/python3.12/site-packages/ray/_private/thirdparty/pynvml/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..fe773c4fca257093523381f489618e86b4cc436b --- /dev/null +++ b/lib/python3.12/site-packages/ray/_private/thirdparty/pynvml/__init__.py @@ -0,0 +1,3 @@ +from ray._private.thirdparty.pynvml.pynvml import * +# nvdia-ml-py version +__version__ = "13.580.65" diff --git a/lib/python3.12/site-packages/ray/_private/thirdparty/pynvml/__pycache__/__init__.cpython-312.pyc b/lib/python3.12/site-packages/ray/_private/thirdparty/pynvml/__pycache__/__init__.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..e7ca274dd460d8334cfdd35a1951223b431cfe6b Binary files /dev/null and b/lib/python3.12/site-packages/ray/_private/thirdparty/pynvml/__pycache__/__init__.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/_private/thirdparty/pynvml/pynvml.py b/lib/python3.12/site-packages/ray/_private/thirdparty/pynvml/pynvml.py new file mode 100644 index 0000000000000000000000000000000000000000..4db9754a1bb1a574c61e1ce054bf35133dc1aeb3 --- /dev/null +++ b/lib/python3.12/site-packages/ray/_private/thirdparty/pynvml/pynvml.py @@ -0,0 +1,6920 @@ +##### +# Copyright (c) 2011-2025, NVIDIA Corporation. All rights reserved. +# +# Redistribution and use in source and binary forms, with or without +# modification, are permitted provided that the following conditions are met: +# +# * Redistributions of source code must retain the above copyright notice, +# this list of conditions and the following disclaimer. +# * Redistributions in binary form must reproduce the above copyright +# notice, this list of conditions and the following disclaimer in the +# documentation and/or other materials provided with the distribution. +# * Neither the name of the NVIDIA Corporation nor the names of its +# contributors may be used to endorse or promote products derived from +# this software without specific prior written permission. +# +# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" +# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE +# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE +# ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE +# LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR +# CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF +# SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS +# INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN +# CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) +# ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF +# THE POSSIBILITY OF SUCH DAMAGE. +##### + +## +# Python bindings for the NVML library +## +from ctypes import * +from ctypes.util import find_library +from functools import wraps +import sys +import os +import threading +import string + +## C Type mappings ## +## Enums +_nvmlEnableState_t = c_uint +NVML_FEATURE_DISABLED = 0 +NVML_FEATURE_ENABLED = 1 + +_nvmlBrandType_t = c_uint +NVML_BRAND_UNKNOWN = 0 +NVML_BRAND_QUADRO = 1 +NVML_BRAND_TESLA = 2 +NVML_BRAND_NVS = 3 +NVML_BRAND_GRID = 4 # Deprecated from API reporting. Keeping definition for backward compatibility. +NVML_BRAND_GEFORCE = 5 +NVML_BRAND_TITAN = 6 +NVML_BRAND_NVIDIA_VAPPS = 7 # NVIDIA Virtual Applications +NVML_BRAND_NVIDIA_VPC = 8 # NVIDIA Virtual PC +NVML_BRAND_NVIDIA_VCS = 9 # NVIDIA Virtual Compute Server +NVML_BRAND_NVIDIA_VWS = 10 # NVIDIA RTX Virtual Workstation +NVML_BRAND_NVIDIA_CLOUD_GAMING = 11 # NVIDIA Cloud Gaming +NVML_BRAND_NVIDIA_VGAMING = NVML_BRAND_NVIDIA_CLOUD_GAMING # Deprecated from API reporting. Keeping definition for backward compatibility. +NVML_BRAND_QUADRO_RTX = 12 +NVML_BRAND_NVIDIA_RTX = 13 +NVML_BRAND_NVIDIA = 14 +NVML_BRAND_GEFORCE_RTX = 15 # Unused +NVML_BRAND_TITAN_RTX = 16 # Unused +NVML_BRAND_COUNT = 18 + +_nvmlTemperatureThresholds_t = c_uint +NVML_TEMPERATURE_THRESHOLD_SHUTDOWN = 0 +NVML_TEMPERATURE_THRESHOLD_SLOWDOWN = 1 +NVML_TEMPERATURE_THRESHOLD_MEM_MAX = 2 +NVML_TEMPERATURE_THRESHOLD_GPU_MAX = 3 +NVML_TEMPERATURE_THRESHOLD_ACOUSTIC_MIN = 4 +NVML_TEMPERATURE_THRESHOLD_ACOUSTIC_CURR = 5 +NVML_TEMPERATURE_THRESHOLD_ACOUSTIC_MAX = 6 +NVML_TEMPERATURE_THRESHOLD_GPS_CURR = 7 +NVML_TEMPERATURE_THRESHOLD_COUNT = 8 + +_nvmlTemperatureSensors_t = c_uint +NVML_TEMPERATURE_GPU = 0 +NVML_TEMPERATURE_COUNT = 1 + + +_nvmlComputeMode_t = c_uint +NVML_COMPUTEMODE_DEFAULT = 0 +NVML_COMPUTEMODE_EXCLUSIVE_THREAD = 1 ## Support Removed +NVML_COMPUTEMODE_PROHIBITED = 2 +NVML_COMPUTEMODE_EXCLUSIVE_PROCESS = 3 +NVML_COMPUTEMODE_COUNT = 4 + +_nvmlMemoryLocation_t = c_uint +NVML_MEMORY_LOCATION_L1_CACHE = 0 +NVML_MEMORY_LOCATION_L2_CACHE = 1 +NVML_MEMORY_LOCATION_DEVICE_MEMORY = 2 +NVML_MEMORY_LOCATION_DRAM = 2 +NVML_MEMORY_LOCATION_REGISTER_FILE = 3 +NVML_MEMORY_LOCATION_TEXTURE_MEMORY = 4 +NVML_MEMORY_LOCATION_TEXTURE_SHM = 5 +NVML_MEMORY_LOCATION_CBU = 6 +NVML_MEMORY_LOCATION_SRAM = 7 +NVML_MEMORY_LOCATION_COUNT = 8 + +NVML_NVLINK_MAX_LINKS = 18 + +# For backwards compatibility, maintain the incorrectly-named "LANES" define +NVML_NVLINK_MAX_LANES = NVML_NVLINK_MAX_LINKS + +_nvmlNvLinkErrorCounter_t = c_uint +NVML_NVLINK_ERROR_DL_REPLAY = 0 +NVML_NVLINK_ERROR_DL_RECOVERY = 1 +NVML_NVLINK_ERROR_DL_CRC_FLIT = 2 +NVML_NVLINK_ERROR_DL_CRC_DATA = 3 +NVML_NVLINK_ERROR_DL_ECC_DATA = 4 +NVML_NVLINK_ERROR_COUNT = 5 + +_nvmlNvLinkEccLaneErrorCounter_t = c_uint +NVML_NVLINK_ERROR_DL_ECC_LANE0 = 0 +NVML_NVLINK_ERROR_DL_ECC_LANE1 = 1 +NVML_NVLINK_ERROR_DL_ECC_LANE2 = 2 +NVML_NVLINK_ERROR_DL_ECC_LANE3 = 3 +NVML_NVLINK_ERROR_DL_ECC_COUNT = 5 + +_nvmlNvLinkCapability_t = c_uint +NVML_NVLINK_CAP_P2P_SUPPORTED = 0 +NVML_NVLINK_CAP_SYSMEM_ACCESS = 1 +NVML_NVLINK_CAP_P2P_ATOMICS = 2 +NVML_NVLINK_CAP_SYSMEM_ATOMICS= 3 +NVML_NVLINK_CAP_SLI_BRIDGE = 4 +NVML_NVLINK_CAP_VALID = 5 +NVML_NVLINK_CAP_COUNT = 6 + +_nvmlNvLinkUtilizationCountPktTypes_t = c_uint +NVML_NVLINK_COUNTER_PKTFILTER_NOP = 0x1 +NVML_NVLINK_COUNTER_PKTFILTER_READ = 0x2 +NVML_NVLINK_COUNTER_PKTFILTER_WRITE = 0x4 +NVML_NVLINK_COUNTER_PKTFILTER_RATOM = 0x8 +NVML_NVLINK_COUNTER_PKTFILTER_NRATOM = 0x10 +NVML_NVLINK_COUNTER_PKTFILTER_FLUSH = 0x20 +NVML_NVLINK_COUNTER_PKTFILTER_RESPDATA = 0x40 +NVML_NVLINK_COUNTER_PKTFILTER_RESPNODATA = 0x80 +NVML_NVLINK_COUNTER_PKTFILTER_ALL = 0xFF + +_nvmlNvLinkUtilizationCountUnits_t = c_uint +NVML_NVLINK_COUNTER_UNIT_CYCLES = 0 +NVML_NVLINK_COUNTER_UNIT_PACKETS = 1 +NVML_NVLINK_COUNTER_UNIT_BYTES = 2 +NVML_NVLINK_COUNTER_UNIT_RESERVED = 3 +NVML_NVLINK_COUNTER_UNIT_COUNT = 4 + +_nvmlNvLinkDeviceType_t = c_uint +NVML_NVLINK_DEVICE_TYPE_GPU = 0x00 +NVML_NVLINK_DEVICE_TYPE_IBMNPU = 0x01 +NVML_NVLINK_DEVICE_TYPE_SWITCH = 0x02 +NVML_NVLINK_DEVICE_TYPE_UNKNOWN = 0xFF + +# These are deprecated, instead use _nvmlMemoryErrorType_t +_nvmlEccBitType_t = c_uint +NVML_SINGLE_BIT_ECC = 0 +NVML_DOUBLE_BIT_ECC = 1 +NVML_ECC_ERROR_TYPE_COUNT = 2 + +_nvmlEccCounterType_t = c_uint +NVML_VOLATILE_ECC = 0 +NVML_AGGREGATE_ECC = 1 +NVML_ECC_COUNTER_TYPE_COUNT = 2 + +_nvmlMemoryErrorType_t = c_uint +NVML_MEMORY_ERROR_TYPE_CORRECTED = 0 +NVML_MEMORY_ERROR_TYPE_UNCORRECTED = 1 +NVML_MEMORY_ERROR_TYPE_COUNT = 2 + +_nvmlClockType_t = c_uint +NVML_CLOCK_GRAPHICS = 0 +NVML_CLOCK_SM = 1 +NVML_CLOCK_MEM = 2 +NVML_CLOCK_VIDEO = 3 +NVML_CLOCK_COUNT = 4 + +_nvmlClockId_t = c_uint +NVML_CLOCK_ID_CURRENT = 0 +NVML_CLOCK_ID_APP_CLOCK_TARGET = 1 +NVML_CLOCK_ID_APP_CLOCK_DEFAULT = 2 +NVML_CLOCK_ID_CUSTOMER_BOOST_MAX = 3 +NVML_CLOCK_ID_COUNT = 4 + +_nvmlDriverModel_t = c_uint +NVML_DRIVER_WDDM = 0 +NVML_DRIVER_WDM = 1 +NVML_DRIVER_MCDM = 2 + +NVML_MAX_GPU_PERF_PSTATES = 16 + +_nvmlPstates_t = c_uint +NVML_PSTATE_0 = 0 +NVML_PSTATE_1 = 1 +NVML_PSTATE_2 = 2 +NVML_PSTATE_3 = 3 +NVML_PSTATE_4 = 4 +NVML_PSTATE_5 = 5 +NVML_PSTATE_6 = 6 +NVML_PSTATE_7 = 7 +NVML_PSTATE_8 = 8 +NVML_PSTATE_9 = 9 +NVML_PSTATE_10 = 10 +NVML_PSTATE_11 = 11 +NVML_PSTATE_12 = 12 +NVML_PSTATE_13 = 13 +NVML_PSTATE_14 = 14 +NVML_PSTATE_15 = 15 +NVML_PSTATE_UNKNOWN = 32 + +_nvmlInforomObject_t = c_uint +NVML_INFOROM_OEM = 0 +NVML_INFOROM_ECC = 1 +NVML_INFOROM_POWER = 2 +NVML_INFOROM_DEN = 3 +NVML_INFOROM_COUNT = 4 + +_nvmlReturn_t = c_uint +NVML_SUCCESS = 0 +NVML_ERROR_UNINITIALIZED = 1 +NVML_ERROR_INVALID_ARGUMENT = 2 +NVML_ERROR_NOT_SUPPORTED = 3 +NVML_ERROR_NO_PERMISSION = 4 +NVML_ERROR_ALREADY_INITIALIZED = 5 +NVML_ERROR_NOT_FOUND = 6 +NVML_ERROR_INSUFFICIENT_SIZE = 7 +NVML_ERROR_INSUFFICIENT_POWER = 8 +NVML_ERROR_DRIVER_NOT_LOADED = 9 +NVML_ERROR_TIMEOUT = 10 +NVML_ERROR_IRQ_ISSUE = 11 +NVML_ERROR_LIBRARY_NOT_FOUND = 12 +NVML_ERROR_FUNCTION_NOT_FOUND = 13 +NVML_ERROR_CORRUPTED_INFOROM = 14 +NVML_ERROR_GPU_IS_LOST = 15 +NVML_ERROR_RESET_REQUIRED = 16 +NVML_ERROR_OPERATING_SYSTEM = 17 +NVML_ERROR_LIB_RM_VERSION_MISMATCH = 18 +NVML_ERROR_IN_USE = 19 +NVML_ERROR_MEMORY = 20 +NVML_ERROR_NO_DATA = 21 +NVML_ERROR_VGPU_ECC_NOT_SUPPORTED = 22 +NVML_ERROR_INSUFFICIENT_RESOURCES = 23 +NVML_ERROR_FREQ_NOT_SUPPORTED = 24 +NVML_ERROR_ARGUMENT_VERSION_MISMATCH = 25 +NVML_ERROR_DEPRECATED = 26 +NVML_ERROR_NOT_READY = 27 +NVML_ERROR_GPU_NOT_FOUND = 28 +NVML_ERROR_INVALID_STATE = 29 +NVML_ERROR_RESET_TYPE_NOT_SUPPORTED = 30 +NVML_ERROR_UNKNOWN = 999 + +_nvmlFanState_t = c_uint +NVML_FAN_NORMAL = 0 +NVML_FAN_FAILED = 1 + +_nvmlFanControlPolicy_t = c_uint +NVML_FAN_POLICY_TEMPERATURE_CONTINOUS_SW = 0 +NVML_FAN_POLICY_MANUAL = 1 + +_nvmlLedColor_t = c_uint +NVML_LED_COLOR_GREEN = 0 +NVML_LED_COLOR_AMBER = 1 + +_nvmlGpuOperationMode_t = c_uint +NVML_GOM_ALL_ON = 0 +NVML_GOM_COMPUTE = 1 +NVML_GOM_LOW_DP = 2 + +_nvmlPageRetirementCause_t = c_uint +NVML_PAGE_RETIREMENT_CAUSE_MULTIPLE_SINGLE_BIT_ECC_ERRORS = 0 +NVML_PAGE_RETIREMENT_CAUSE_DOUBLE_BIT_ECC_ERROR = 1 +NVML_PAGE_RETIREMENT_CAUSE_COUNT = 2 + +_nvmlRestrictedAPI_t = c_uint +NVML_RESTRICTED_API_SET_APPLICATION_CLOCKS = 0 +NVML_RESTRICTED_API_SET_AUTO_BOOSTED_CLOCKS = 1 +NVML_RESTRICTED_API_COUNT = 2 + +_nvmlBridgeChipType_t = c_uint +NVML_BRIDGE_CHIP_PLX = 0 +NVML_BRIDGE_CHIP_BRO4 = 1 +NVML_MAX_PHYSICAL_BRIDGE = 128 + +_nvmlValueType_t = c_uint +NVML_VALUE_TYPE_DOUBLE = 0 +NVML_VALUE_TYPE_UNSIGNED_INT = 1 +NVML_VALUE_TYPE_UNSIGNED_LONG = 2 +NVML_VALUE_TYPE_UNSIGNED_LONG_LONG = 3 +NVML_VALUE_TYPE_SIGNED_LONG_LONG = 4 +NVML_VALUE_TYPE_SIGNED_INT = 5 +NVML_VALUE_TYPE_UNSIGNED_SHORT = 6 +NVML_VALUE_TYPE_COUNT = 7 + +_nvmlNvlinkVersion_t = c_uint +NVML_NVLINK_VERSION_INVALID = 0 +NVML_NVLINK_VERSION_1_0 = 1 +NVML_NVLINK_VERSION_2_0 = 2 +NVML_NVLINK_VERSION_2_2 = 3 +NVML_NVLINK_VERSION_3_0 = 4 +NVML_NVLINK_VERSION_3_1 = 5 +NVML_NVLINK_VERSION_4_0 = 6 +NVML_NVLINK_VERSION_5_0 = 7 + +_nvmlPerfPolicyType_t = c_uint +NVML_PERF_POLICY_POWER = 0 +NVML_PERF_POLICY_THERMAL = 1 +NVML_PERF_POLICY_SYNC_BOOST = 2 +NVML_PERF_POLICY_BOARD_LIMIT = 3 +NVML_PERF_POLICY_LOW_UTILIZATION = 4 +NVML_PERF_POLICY_RELIABILITY = 5 +NVML_PERF_POLICY_TOTAL_APP_CLOCKS = 10 +NVML_PERF_POLICY_TOTAL_BASE_CLOCKS = 11 +NVML_PERF_POLICY_COUNT = 12 + +_nvmlEncoderQueryType_t = c_uint +NVML_ENCODER_QUERY_H264 = 0 +NVML_ENCODER_QUERY_HEVC = 1 +NVML_ENCODER_QUERY_AV1 = 2 +NVML_ENCODER_QUERY_UNKNOWN = 255 + +_nvmlFBCSessionType_t = c_uint +NVML_FBC_SESSION_TYPE_UNKNOWN = 0 +NVML_FBC_SESSION_TYPE_TOSYS = 1 +NVML_FBC_SESSION_TYPE_CUDA = 2 +NVML_FBC_SESSION_TYPE_VID = 3 +NVML_FBC_SESSION_TYPE_HWENC = 4 + +_nvmlDetachGpuState_t = c_uint +NVML_DETACH_GPU_KEEP = 0 +NVML_DETACH_GPU_REMOVE = 1 + +_nvmlPcieLinkState_t = c_uint +NVML_PCIE_LINK_KEEP = 0 +NVML_PCIE_LINK_SHUT_DOWN = 1 + +_nvmlSamplingType_t = c_uint +NVML_TOTAL_POWER_SAMPLES = 0 +NVML_GPU_UTILIZATION_SAMPLES = 1 +NVML_MEMORY_UTILIZATION_SAMPLES = 2 +NVML_ENC_UTILIZATION_SAMPLES = 3 +NVML_DEC_UTILIZATION_SAMPLES = 4 +NVML_PROCESSOR_CLK_SAMPLES = 5 +NVML_MEMORY_CLK_SAMPLES = 6 +NVML_MODULE_POWER_SAMPLES = 7 +NVML_JPG_UTILIZATION_SAMPLES = 8 +NVML_OFA_UTILIZATION_SAMPLES = 9 +NVML_SAMPLINGTYPE_COUNT = 10 + +_nvmlPcieUtilCounter_t = c_uint +NVML_PCIE_UTIL_TX_BYTES = 0 +NVML_PCIE_UTIL_RX_BYTES = 1 +NVML_PCIE_UTIL_COUNT = 2 + +_nvmlGpuTopologyLevel_t = c_uint +NVML_TOPOLOGY_INTERNAL = 0 +NVML_TOPOLOGY_SINGLE = 10 +NVML_TOPOLOGY_MULTIPLE = 20 +NVML_TOPOLOGY_HOSTBRIDGE = 30 +NVML_TOPOLOGY_NODE = 40 +NVML_TOPOLOGY_CPU = NVML_TOPOLOGY_NODE +NVML_TOPOLOGY_SYSTEM = 50 + +_nvmlGpuP2PCapsIndex_t = c_uint +NVML_P2P_CAPS_INDEX_READ = 0, +NVML_P2P_CAPS_INDEX_WRITE = 1 +NVML_P2P_CAPS_INDEX_NVLINK =2 +NVML_P2P_CAPS_INDEX_ATOMICS = 3 +# +# NVML_P2P_CAPS_INDEX_PROP is deprecated. +# Use NVML_P2P_CAPS_INDEX_PCI instead. +# +NVML_P2P_CAPS_INDEX_PROP = 4 +NVML_P2P_CAPS_INDEX_PCI = 4 +NVML_P2P_CAPS_INDEX_UNKNOWN = 5 + +_nvmlGpuP2PStatus_t = c_uint +NVML_P2P_STATUS_OK = 0 +NVML_P2P_STATUS_CHIPSET_NOT_SUPPORED = 1 +NVML_P2P_STATUS_CHIPSET_NOT_SUPPORTED = NVML_P2P_STATUS_CHIPSET_NOT_SUPPORED +NVML_P2P_STATUS_GPU_NOT_SUPPORTED = 2 +NVML_P2P_STATUS_IOH_TOPOLOGY_NOT_SUPPORTED =3 +NVML_P2P_STATUS_DISABLED_BY_REGKEY =4 +NVML_P2P_STATUS_NOT_SUPPORTED =5 +NVML_P2P_STATUS_UNKNOWN =6 + +_nvmlDeviceArchitecture_t = c_uint +NVML_DEVICE_ARCH_KEPLER = 2 +NVML_DEVICE_ARCH_MAXWELL = 3 +NVML_DEVICE_ARCH_PASCAL = 4 +NVML_DEVICE_ARCH_VOLTA = 5 +NVML_DEVICE_ARCH_TURING = 6 +NVML_DEVICE_ARCH_AMPERE = 7 +NVML_DEVICE_ARCH_ADA = 8 +NVML_DEVICE_ARCH_HOPPER = 9 +NVML_DEVICE_ARCH_BLACKWELL = 10 +NVML_DEVICE_ARCH_UNKNOWN = 0xffffffff + +# PCI bus Types +_nvmlBusType_t = c_uint +NVML_BUS_TYPE_UNKNOWN = 0 +NVML_BUS_TYPE_PCI = 1 +NVML_BUS_TYPE_PCIE = 2 +NVML_BUS_TYPE_FPCI = 3 +NVML_BUS_TYPE_AGP = 4 + +_nvmlPowerSource_t = c_uint +NVML_POWER_SOURCE_AC = 0x00000000 +NVML_POWER_SOURCE_BATTERY = 0x00000001 +NVML_POWER_SOURCE_UNDERSIZED = 0x00000002 + +_nvmlAdaptiveClockInfoStatus_t = c_uint +NVML_ADAPTIVE_CLOCKING_INFO_STATUS_DISABLED = 0x00000000 +NVML_ADAPTIVE_CLOCKING_INFO_STATUS_ENABLED = 0x00000001 + +_nvmlClockLimitId_t = c_uint +NVML_CLOCK_LIMIT_ID_RANGE_START = 0xffffff00 +NVML_CLOCK_LIMIT_ID_TDP = 0xffffff01 +NVML_CLOCK_LIMIT_ID_UNLIMITED = 0xffffff02 + +_nvmlPcieLinkMaxSpeed_t = c_uint +NVML_PCIE_LINK_MAX_SPEED_INVALID = 0x00000000 +NVML_PCIE_LINK_MAX_SPEED_2500MBPS = 0x00000001 +NVML_PCIE_LINK_MAX_SPEED_5000MBPS = 0x00000002 +NVML_PCIE_LINK_MAX_SPEED_8000MBPS = 0x00000003 +NVML_PCIE_LINK_MAX_SPEED_16000MBPS = 0x00000004 +NVML_PCIE_LINK_MAX_SPEED_32000MBPS = 0x00000005 +NVML_PCIE_LINK_MAX_SPEED_64000MBPS = 0x00000006 + +_nvmlPcieAtomicsCapability_t = c_uint +NVML_PCIE_ATOMICS_CAP_FETCHADD32 = 0x01 +NVML_PCIE_ATOMICS_CAP_FETCHADD64 = 0x02 +NVML_PCIE_ATOMICS_CAP_SWAP32 = 0x04 +NVML_PCIE_ATOMICS_CAP_SWAP64 = 0x08 +NVML_PCIE_ATOMICS_CAP_CAS32 = 0x10 +NVML_PCIE_ATOMICS_CAP_CAS64 = 0x20 +NVML_PCIE_ATOMICS_CAP_CAS128 = 0x40 +NVML_PCIE_ATOMICS_OPS_MAX = 7 + +_nvmlAffinityScope_t = c_uint +NVML_AFFINITY_SCOPE_NODE = 0 +NVML_AFFINITY_SCOPE_SOCKET = 1 + +_nvmlDeviceGpuRecoveryAction_t = c_uint +NVML_GPU_RECOVERY_ACTION_NONE = 0 +NVML_GPU_RECOVERY_ACTION_GPU_RESET = 1 +NVML_GPU_RECOVERY_ACTION_NODE_REBOOT = 2 +NVML_GPU_RECOVERY_ACTION_DRAIN_P2P = 3 +NVML_GPU_RECOVERY_ACTION_DRAIN_AND_RESET = 4 +NVML_GPU_RECOVERY_ACTION_GPU_RESET_BUS = 5 + +# C preprocessor defined values +nvmlFlagDefault = 0 +nvmlFlagForce = 1 +NVML_INIT_FLAG_NO_GPUS = 1 +NVML_INIT_FLAG_NO_ATTACH = 2 + +NVML_MAX_GPC_COUNT = 32 + +# buffer size +NVML_DEVICE_INFOROM_VERSION_BUFFER_SIZE = 16 +NVML_DEVICE_UUID_BUFFER_SIZE = 80 +NVML_DEVICE_UUID_V2_BUFFER_SIZE = 96 +NVML_SYSTEM_DRIVER_VERSION_BUFFER_SIZE = 80 +NVML_SYSTEM_NVML_VERSION_BUFFER_SIZE = 80 +NVML_DEVICE_NAME_BUFFER_SIZE = 64 +NVML_DEVICE_NAME_V2_BUFFER_SIZE = 96 +NVML_DEVICE_SERIAL_BUFFER_SIZE = 30 +NVML_DEVICE_PART_NUMBER_BUFFER_SIZE = 80 +NVML_DEVICE_GPU_PART_NUMBER_BUFFER_SIZE = 80 +NVML_DEVICE_VBIOS_VERSION_BUFFER_SIZE = 32 +NVML_DEVICE_PCI_BUS_ID_BUFFER_SIZE = 32 +NVML_DEVICE_PCI_BUS_ID_BUFFER_V2_SIZE = 16 +NVML_GRID_LICENSE_BUFFER_SIZE = 128 +NVML_VGPU_NAME_BUFFER_SIZE = 64 +NVML_GRID_LICENSE_FEATURE_MAX_COUNT = 3 +NVML_VGPU_METADATA_OPAQUE_DATA_SIZE = sizeof(c_uint) + 256 +NVML_VGPU_PGPU_METADATA_OPAQUE_DATA_SIZE = 256 +NVML_DEVICE_GPU_FRU_PART_NUMBER_BUFFER_SIZE = 0x14 # NV2080_GPU_MAX_PRODUCT_PART_NUMBER_LENGTH +NVML_PERF_MODES_BUFFER_SIZE = 2048 + +# Format strings +NVML_DEVICE_PCI_BUS_ID_LEGACY_FMT = "%04X:%02X:%02X.0" +NVML_DEVICE_PCI_BUS_ID_FMT = "%08X:%02X:%02X.0" + +NVML_VALUE_NOT_AVAILABLE_ulonglong = c_ulonglong(-1) +NVML_VALUE_NOT_AVAILABLE_uint = c_uint(-1) + +''' + Field Identifiers. + + All Identifiers pertain to a device. Each ID is only used once and is guaranteed never to change. +''' +NVML_FI_DEV_ECC_CURRENT = 1 # Current ECC mode. 1=Active. 0=Inactive +NVML_FI_DEV_ECC_PENDING = 2 # Pending ECC mode. 1=Active. 0=Inactive + +#ECC Count Totals +NVML_FI_DEV_ECC_SBE_VOL_TOTAL = 3 # Total single bit volatile ECC errors +NVML_FI_DEV_ECC_DBE_VOL_TOTAL = 4 # Total double bit volatile ECC errors +NVML_FI_DEV_ECC_SBE_AGG_TOTAL = 5 # Total single bit aggregate (persistent) ECC errors +NVML_FI_DEV_ECC_DBE_AGG_TOTAL = 6 # Total double bit aggregate (persistent) ECC errors +#Individual ECC locations +NVML_FI_DEV_ECC_SBE_VOL_L1 = 7 # L1 cache single bit volatile ECC errors +NVML_FI_DEV_ECC_DBE_VOL_L1 = 8 # L1 cache double bit volatile ECC errors +NVML_FI_DEV_ECC_SBE_VOL_L2 = 9 # L2 cache single bit volatile ECC errors +NVML_FI_DEV_ECC_DBE_VOL_L2 = 10 # L2 cache double bit volatile ECC errors +NVML_FI_DEV_ECC_SBE_VOL_DEV = 11 # Device memory single bit volatile ECC errors +NVML_FI_DEV_ECC_DBE_VOL_DEV = 12 # Device memory double bit volatile ECC errors +NVML_FI_DEV_ECC_SBE_VOL_REG = 13 # Register file single bit volatile ECC errors +NVML_FI_DEV_ECC_DBE_VOL_REG = 14 # Register file double bit volatile ECC errors +NVML_FI_DEV_ECC_SBE_VOL_TEX = 15 # Texture memory single bit volatile ECC errors +NVML_FI_DEV_ECC_DBE_VOL_TEX = 16 # Texture memory double bit volatile ECC errors +NVML_FI_DEV_ECC_DBE_VOL_CBU = 17 # CBU double bit volatile ECC errors +NVML_FI_DEV_ECC_SBE_AGG_L1 = 18 # L1 cache single bit aggregate (persistent) ECC errors +NVML_FI_DEV_ECC_DBE_AGG_L1 = 19 # L1 cache double bit aggregate (persistent) ECC errors +NVML_FI_DEV_ECC_SBE_AGG_L2 = 20 # L2 cache single bit aggregate (persistent) ECC errors +NVML_FI_DEV_ECC_DBE_AGG_L2 = 21 # L2 cache double bit aggregate (persistent) ECC errors +NVML_FI_DEV_ECC_SBE_AGG_DEV = 22 # Device memory single bit aggregate (persistent) ECC errors +NVML_FI_DEV_ECC_DBE_AGG_DEV = 23 # Device memory double bit aggregate (persistent) ECC errors +NVML_FI_DEV_ECC_SBE_AGG_REG = 24 # Register File single bit aggregate (persistent) ECC errors +NVML_FI_DEV_ECC_DBE_AGG_REG = 25 # Register File double bit aggregate (persistent) ECC errors +NVML_FI_DEV_ECC_SBE_AGG_TEX = 26 # Texture memory single bit aggregate (persistent) ECC errors +NVML_FI_DEV_ECC_DBE_AGG_TEX = 27 # Texture memory double bit aggregate (persistent) ECC errors +NVML_FI_DEV_ECC_DBE_AGG_CBU = 28 # CBU double bit aggregate ECC errors + +# Page Retirement +NVML_FI_DEV_RETIRED_SBE = 29 # Number of retired pages because of single bit errors +NVML_FI_DEV_RETIRED_DBE = 30 # Number of retired pages because of double bit errors +NVML_FI_DEV_RETIRED_PENDING = 31 # If any pages are pending retirement. 1=yes. 0=no. + +# NvLink Flit Error Counters +NVML_FI_DEV_NVLINK_CRC_FLIT_ERROR_COUNT_L0 = 32 # NVLink flow control CRC Error Counter for Lane 0 +NVML_FI_DEV_NVLINK_CRC_FLIT_ERROR_COUNT_L1 = 33 # NVLink flow control CRC Error Counter for Lane 1 +NVML_FI_DEV_NVLINK_CRC_FLIT_ERROR_COUNT_L2 = 34 # NVLink flow control CRC Error Counter for Lane 2 +NVML_FI_DEV_NVLINK_CRC_FLIT_ERROR_COUNT_L3 = 35 # NVLink flow control CRC Error Counter for Lane 3 +NVML_FI_DEV_NVLINK_CRC_FLIT_ERROR_COUNT_L4 = 36 # NVLink flow control CRC Error Counter for Lane 4 +NVML_FI_DEV_NVLINK_CRC_FLIT_ERROR_COUNT_L5 = 37 # NVLink flow control CRC Error Counter for Lane 5 +NVML_FI_DEV_NVLINK_CRC_FLIT_ERROR_COUNT_TOTAL = 38 # NVLink flow control CRC Error Counter total for all Lanes + +# NvLink CRC Data Error Counters +NVML_FI_DEV_NVLINK_CRC_DATA_ERROR_COUNT_L0 = 39 # NVLink data CRC Error Counter for Lane 0 +NVML_FI_DEV_NVLINK_CRC_DATA_ERROR_COUNT_L1 = 40 # NVLink data CRC Error Counter for Lane 1 +NVML_FI_DEV_NVLINK_CRC_DATA_ERROR_COUNT_L2 = 41 # NVLink data CRC Error Counter for Lane 2 +NVML_FI_DEV_NVLINK_CRC_DATA_ERROR_COUNT_L3 = 42 # NVLink data CRC Error Counter for Lane 3 +NVML_FI_DEV_NVLINK_CRC_DATA_ERROR_COUNT_L4 = 43 # NVLink data CRC Error Counter for Lane 4 +NVML_FI_DEV_NVLINK_CRC_DATA_ERROR_COUNT_L5 = 44 # NVLink data CRC Error Counter for Lane 5 +NVML_FI_DEV_NVLINK_CRC_DATA_ERROR_COUNT_TOTAL = 45 # NvLink data CRC Error Counter total for all Lanes + +# NvLink Replay Error Counters +NVML_FI_DEV_NVLINK_REPLAY_ERROR_COUNT_L0 = 46 # NVLink Replay Error Counter for Lane 0 +NVML_FI_DEV_NVLINK_REPLAY_ERROR_COUNT_L1 = 47 # NVLink Replay Error Counter for Lane 1 +NVML_FI_DEV_NVLINK_REPLAY_ERROR_COUNT_L2 = 48 # NVLink Replay Error Counter for Lane 2 +NVML_FI_DEV_NVLINK_REPLAY_ERROR_COUNT_L3 = 49 # NVLink Replay Error Counter for Lane 3 +NVML_FI_DEV_NVLINK_REPLAY_ERROR_COUNT_L4 = 50 # NVLink Replay Error Counter for Lane 4 +NVML_FI_DEV_NVLINK_REPLAY_ERROR_COUNT_L5 = 51 # NVLink Replay Error Counter for Lane 5 +NVML_FI_DEV_NVLINK_REPLAY_ERROR_COUNT_TOTAL = 52 # NVLink Replay Error Counter total for all Lanes + +# NvLink Recovery Error Counters +NVML_FI_DEV_NVLINK_RECOVERY_ERROR_COUNT_L0 = 53 # NVLink Recovery Error Counter for Lane 0 +NVML_FI_DEV_NVLINK_RECOVERY_ERROR_COUNT_L1 = 54 # NVLink Recovery Error Counter for Lane 1 +NVML_FI_DEV_NVLINK_RECOVERY_ERROR_COUNT_L2 = 55 # NVLink Recovery Error Counter for Lane 2 +NVML_FI_DEV_NVLINK_RECOVERY_ERROR_COUNT_L3 = 56 # NVLink Recovery Error Counter for Lane 3 +NVML_FI_DEV_NVLINK_RECOVERY_ERROR_COUNT_L4 = 57 # NVLink Recovery Error Counter for Lane 4 +NVML_FI_DEV_NVLINK_RECOVERY_ERROR_COUNT_L5 = 58 # NVLink Recovery Error Counter for Lane 5 +NVML_FI_DEV_NVLINK_RECOVERY_ERROR_COUNT_TOTAL = 59 # NVLink Recovery Error Counter total for all Lanes + +# NvLink Bandwidth Counters +NVML_FI_DEV_NVLINK_BANDWIDTH_C0_L0 = 60 # NVLink Bandwidth Counter for Counter Set 0, Lane 0 +NVML_FI_DEV_NVLINK_BANDWIDTH_C0_L1 = 61 # NVLink Bandwidth Counter for Counter Set 0, Lane 1 +NVML_FI_DEV_NVLINK_BANDWIDTH_C0_L2 = 62 # NVLink Bandwidth Counter for Counter Set 0, Lane 2 +NVML_FI_DEV_NVLINK_BANDWIDTH_C0_L3 = 63 # NVLink Bandwidth Counter for Counter Set 0, Lane 3 +NVML_FI_DEV_NVLINK_BANDWIDTH_C0_L4 = 64 # NVLink Bandwidth Counter for Counter Set 0, Lane 4 +NVML_FI_DEV_NVLINK_BANDWIDTH_C0_L5 = 65 # NVLink Bandwidth Counter for Counter Set 0, Lane 5 +NVML_FI_DEV_NVLINK_BANDWIDTH_C0_TOTAL = 66 # NVLink Bandwidth Counter Total for Counter Set 0, All Lanes + +# NvLink Bandwidth Counters +NVML_FI_DEV_NVLINK_BANDWIDTH_C1_L0 = 67 # NVLink Bandwidth Counter for Counter Set 1, Lane 0 +NVML_FI_DEV_NVLINK_BANDWIDTH_C1_L1 = 68 # NVLink Bandwidth Counter for Counter Set 1, Lane 1 +NVML_FI_DEV_NVLINK_BANDWIDTH_C1_L2 = 69 # NVLink Bandwidth Counter for Counter Set 1, Lane 2 +NVML_FI_DEV_NVLINK_BANDWIDTH_C1_L3 = 70 # NVLink Bandwidth Counter for Counter Set 1, Lane 3 +NVML_FI_DEV_NVLINK_BANDWIDTH_C1_L4 = 71 # NVLink Bandwidth Counter for Counter Set 1, Lane 4 +NVML_FI_DEV_NVLINK_BANDWIDTH_C1_L5 = 72 # NVLink Bandwidth Counter for Counter Set 1, Lane 5 +NVML_FI_DEV_NVLINK_BANDWIDTH_C1_TOTAL = 73 # NVLink Bandwidth Counter Total for Counter Set 1, All Lanes + +# Perf Policy Counters +NVML_FI_DEV_PERF_POLICY_POWER = 74 # Perf Policy Counter for Power Policy +NVML_FI_DEV_PERF_POLICY_THERMAL = 75 # Perf Policy Counter for Thermal Policy +NVML_FI_DEV_PERF_POLICY_SYNC_BOOST = 76 # Perf Policy Counter for Sync boost Policy +NVML_FI_DEV_PERF_POLICY_BOARD_LIMIT = 77 # Perf Policy Counter for Board Limit +NVML_FI_DEV_PERF_POLICY_LOW_UTILIZATION = 78 # Perf Policy Counter for Low GPU Utilization Policy +NVML_FI_DEV_PERF_POLICY_RELIABILITY = 79 # Perf Policy Counter for Reliability Policy +NVML_FI_DEV_PERF_POLICY_TOTAL_APP_CLOCKS = 80 # Perf Policy Counter for Total App Clock Policy +NVML_FI_DEV_PERF_POLICY_TOTAL_BASE_CLOCKS = 81 # Perf Policy Counter for Total Base Clocks Policy + +# Memory temperatures +NVML_FI_DEV_MEMORY_TEMP = 82 # Memory temperature for the device + +# Energy Counter +NVML_FI_DEV_TOTAL_ENERGY_CONSUMPTION = 83 # Total energy consumption for the GPU in mJ since the driver was last reloaded + +# NVLink Speed +NVML_FI_DEV_NVLINK_SPEED_MBPS_L0 = 84 +NVML_FI_DEV_NVLINK_SPEED_MBPS_L1 = 85 +NVML_FI_DEV_NVLINK_SPEED_MBPS_L2 = 86 +NVML_FI_DEV_NVLINK_SPEED_MBPS_L3 = 87 +NVML_FI_DEV_NVLINK_SPEED_MBPS_L4 = 88 +NVML_FI_DEV_NVLINK_SPEED_MBPS_L5 = 89 +NVML_FI_DEV_NVLINK_SPEED_MBPS_COMMON = 90 + +# NVLink Link Count +NVML_FI_DEV_NVLINK_LINK_COUNT = 91 + +# Page Retirement pending fields +NVML_FI_DEV_RETIRED_PENDING_SBE = 92 +NVML_FI_DEV_RETIRED_PENDING_DBE = 93 + +# PCIe replay and replay rollover counters +NVML_FI_DEV_PCIE_REPLAY_COUNTER = 94 +NVML_FI_DEV_PCIE_REPLAY_ROLLOVER_COUNTER = 95 + +# NvLink Flit Error Counters +NVML_FI_DEV_NVLINK_CRC_FLIT_ERROR_COUNT_L6 = 96 # NVLink flow control CRC Error Counter for Lane 6 +NVML_FI_DEV_NVLINK_CRC_FLIT_ERROR_COUNT_L7 = 97 # NVLink flow control CRC Error Counter for Lane 7 +NVML_FI_DEV_NVLINK_CRC_FLIT_ERROR_COUNT_L8 = 98 # NVLink flow control CRC Error Counter for Lane 8 +NVML_FI_DEV_NVLINK_CRC_FLIT_ERROR_COUNT_L9 = 99 # NVLink flow control CRC Error Counter for Lane 9 +NVML_FI_DEV_NVLINK_CRC_FLIT_ERROR_COUNT_L10 = 100 # NVLink flow control CRC Error Counter for Lane 10 +NVML_FI_DEV_NVLINK_CRC_FLIT_ERROR_COUNT_L11 = 101 # NVLink flow control CRC Error Counter for Lane 11 + +# NvLink CRC Data Error Counters +NVML_FI_DEV_NVLINK_CRC_DATA_ERROR_COUNT_L6 = 102 # NVLink data CRC Error Counter for Lane 6 +NVML_FI_DEV_NVLINK_CRC_DATA_ERROR_COUNT_L7 = 103 # NVLink data CRC Error Counter for Lane 7 +NVML_FI_DEV_NVLINK_CRC_DATA_ERROR_COUNT_L8 = 104 # NVLink data CRC Error Counter for Lane 8 +NVML_FI_DEV_NVLINK_CRC_DATA_ERROR_COUNT_L9 = 105 # NVLink data CRC Error Counter for Lane 9 +NVML_FI_DEV_NVLINK_CRC_DATA_ERROR_COUNT_L10 = 106 # NVLink data CRC Error Counter for Lane 10 +NVML_FI_DEV_NVLINK_CRC_DATA_ERROR_COUNT_L11 = 107 # NVLink data CRC Error Counter for Lane 11 + +# NvLink Replay Error Counters +NVML_FI_DEV_NVLINK_REPLAY_ERROR_COUNT_L6 = 108 # NVLink Replay Error Counter for Lane 6 +NVML_FI_DEV_NVLINK_REPLAY_ERROR_COUNT_L7 = 109 # NVLink Replay Error Counter for Lane 7 +NVML_FI_DEV_NVLINK_REPLAY_ERROR_COUNT_L8 = 110 # NVLink Replay Error Counter for Lane 8 +NVML_FI_DEV_NVLINK_REPLAY_ERROR_COUNT_L9 = 111 # NVLink Replay Error Counter for Lane 9 +NVML_FI_DEV_NVLINK_REPLAY_ERROR_COUNT_L10 = 112 # NVLink Replay Error Counter for Lane 10 +NVML_FI_DEV_NVLINK_REPLAY_ERROR_COUNT_L11 = 113 # NVLink Replay Error Counter for Lane 11 + +# NvLink Recovery Error Counters +NVML_FI_DEV_NVLINK_RECOVERY_ERROR_COUNT_L6 = 114 # NVLink Recovery Error Counter for Lane 6 +NVML_FI_DEV_NVLINK_RECOVERY_ERROR_COUNT_L7 = 115 # NVLink Recovery Error Counter for Lane 7 +NVML_FI_DEV_NVLINK_RECOVERY_ERROR_COUNT_L8 = 116 # NVLink Recovery Error Counter for Lane 8 +NVML_FI_DEV_NVLINK_RECOVERY_ERROR_COUNT_L9 = 117 # NVLink Recovery Error Counter for Lane 9 +NVML_FI_DEV_NVLINK_RECOVERY_ERROR_COUNT_L10 = 118 # NVLink Recovery Error Counter for Lane 10 +NVML_FI_DEV_NVLINK_RECOVERY_ERROR_COUNT_L11 = 119 # NVLink Recovery Error Counter for Lane 11 + +# NvLink Bandwidth Counters +NVML_FI_DEV_NVLINK_BANDWIDTH_C0_L6 = 120 # NVLink Bandwidth Counter for Counter Set 0, Lane 6 +NVML_FI_DEV_NVLINK_BANDWIDTH_C0_L7 = 121 # NVLink Bandwidth Counter for Counter Set 0, Lane 7 +NVML_FI_DEV_NVLINK_BANDWIDTH_C0_L8 = 122 # NVLink Bandwidth Counter for Counter Set 0, Lane 8 +NVML_FI_DEV_NVLINK_BANDWIDTH_C0_L9 = 123 # NVLink Bandwidth Counter for Counter Set 0, Lane 9 +NVML_FI_DEV_NVLINK_BANDWIDTH_C0_L10 = 124 # NVLink Bandwidth Counter for Counter Set 0, Lane 10 +NVML_FI_DEV_NVLINK_BANDWIDTH_C0_L11 = 125 # NVLink Bandwidth Counter for Counter Set 0, Lane 11 + +# NvLink Bandwidth Counters +NVML_FI_DEV_NVLINK_BANDWIDTH_C1_L6 = 126 # NVLink Bandwidth Counter for Counter Set 1, Lane 6 +NVML_FI_DEV_NVLINK_BANDWIDTH_C1_L7 = 127 # NVLink Bandwidth Counter for Counter Set 1, Lane 7 +NVML_FI_DEV_NVLINK_BANDWIDTH_C1_L8 = 128 # NVLink Bandwidth Counter for Counter Set 1, Lane 8 +NVML_FI_DEV_NVLINK_BANDWIDTH_C1_L9 = 129 # NVLink Bandwidth Counter for Counter Set 1, Lane 9 +NVML_FI_DEV_NVLINK_BANDWIDTH_C1_L10 = 130 # NVLink Bandwidth Counter for Counter Set 1, Lane 10 +NVML_FI_DEV_NVLINK_BANDWIDTH_C1_L11 = 131 # NVLink Bandwidth Counter for Counter Set 1, Lane 11 + +# NVLink Speed +NVML_FI_DEV_NVLINK_SPEED_MBPS_L6 = 132 +NVML_FI_DEV_NVLINK_SPEED_MBPS_L7 = 133 +NVML_FI_DEV_NVLINK_SPEED_MBPS_L8 = 134 +NVML_FI_DEV_NVLINK_SPEED_MBPS_L9 = 135 +NVML_FI_DEV_NVLINK_SPEED_MBPS_L10 = 136 +NVML_FI_DEV_NVLINK_SPEED_MBPS_L11 = 137 + +# NVLink Throughput Counters +NVML_FI_DEV_NVLINK_THROUGHPUT_DATA_TX = 138 # NVLink TX Data throughput in KiB +NVML_FI_DEV_NVLINK_THROUGHPUT_DATA_RX = 139 # NVLink RX Data throughput in KiB +NVML_FI_DEV_NVLINK_THROUGHPUT_RAW_TX = 140 # NVLink TX Data + protocol overhead in KiB +NVML_FI_DEV_NVLINK_THROUGHPUT_RAW_RX = 141 # NVLink RX Data + protocol overhead in KiB + +# Row Remapper +NVML_FI_DEV_REMAPPED_COR = 142 +NVML_FI_DEV_REMAPPED_UNC = 143 +NVML_FI_DEV_REMAPPED_PENDING = 144 +NVML_FI_DEV_REMAPPED_FAILURE = 145 + +#Remote device NVLink ID +NVML_FI_DEV_NVLINK_REMOTE_NVLINK_ID = 146 + +# Number of NVLinks connected to NVSwitch +NVML_FI_DEV_NVSWITCH_CONNECTED_LINK_COUNT = 147 + +# NvLink ECC Data Error Counters +NVML_FI_DEV_NVLINK_ECC_DATA_ERROR_COUNT_L0 = 148 #< NVLink data ECC Error Counter for Link 0 +NVML_FI_DEV_NVLINK_ECC_DATA_ERROR_COUNT_L1 = 149 #< NVLink data ECC Error Counter for Link 1 +NVML_FI_DEV_NVLINK_ECC_DATA_ERROR_COUNT_L2 = 150 #< NVLink data ECC Error Counter for Link 2 +NVML_FI_DEV_NVLINK_ECC_DATA_ERROR_COUNT_L3 = 151 #< NVLink data ECC Error Counter for Link 3 +NVML_FI_DEV_NVLINK_ECC_DATA_ERROR_COUNT_L4 = 152 #< NVLink data ECC Error Counter for Link 4 +NVML_FI_DEV_NVLINK_ECC_DATA_ERROR_COUNT_L5 = 153 #< NVLink data ECC Error Counter for Link 5 +NVML_FI_DEV_NVLINK_ECC_DATA_ERROR_COUNT_L6 = 154 #< NVLink data ECC Error Counter for Link 6 +NVML_FI_DEV_NVLINK_ECC_DATA_ERROR_COUNT_L7 = 155 #< NVLink data ECC Error Counter for Link 7 +NVML_FI_DEV_NVLINK_ECC_DATA_ERROR_COUNT_L8 = 156 #< NVLink data ECC Error Counter for Link 8 +NVML_FI_DEV_NVLINK_ECC_DATA_ERROR_COUNT_L9 = 157 #< NVLink data ECC Error Counter for Link 9 +NVML_FI_DEV_NVLINK_ECC_DATA_ERROR_COUNT_L10 = 158 #< NVLink data ECC Error Counter for Link 10 +NVML_FI_DEV_NVLINK_ECC_DATA_ERROR_COUNT_L11 = 159 #< NVLink data ECC Error Counter for Link 11 +NVML_FI_DEV_NVLINK_ECC_DATA_ERROR_COUNT_TOTAL = 160 #< NvLink data ECC Error Counter total for all Links + +NVML_FI_DEV_NVLINK_ERROR_DL_REPLAY = 161 +NVML_FI_DEV_NVLINK_ERROR_DL_RECOVERY = 162 +NVML_FI_DEV_NVLINK_ERROR_DL_CRC = 163 +NVML_FI_DEV_NVLINK_GET_SPEED = 164 +NVML_FI_DEV_NVLINK_GET_STATE = 165 +NVML_FI_DEV_NVLINK_GET_VERSION = 166 + +NVML_FI_DEV_NVLINK_GET_POWER_STATE = 167 +NVML_FI_DEV_NVLINK_GET_POWER_THRESHOLD = 168 + +NVML_FI_DEV_PCIE_L0_TO_RECOVERY_COUNTER = 169 + +NVML_FI_DEV_C2C_LINK_COUNT = 170 +NVML_FI_DEV_C2C_LINK_GET_STATUS = 171 +NVML_FI_DEV_C2C_LINK_GET_MAX_BW = 172 + +NVML_FI_DEV_PCIE_COUNT_CORRECTABLE_ERRORS = 173 +NVML_FI_DEV_PCIE_COUNT_NAKS_RECEIVED = 174 +NVML_FI_DEV_PCIE_COUNT_RECEIVER_ERROR = 175 +NVML_FI_DEV_PCIE_COUNT_BAD_TLP = 176 +NVML_FI_DEV_PCIE_COUNT_NAKS_SENT = 177 +NVML_FI_DEV_PCIE_COUNT_BAD_DLLP = 178 +NVML_FI_DEV_PCIE_COUNT_NON_FATAL_ERROR = 179 +NVML_FI_DEV_PCIE_COUNT_FATAL_ERROR = 180 +NVML_FI_DEV_PCIE_COUNT_UNSUPPORTED_REQ = 181 +NVML_FI_DEV_PCIE_COUNT_LCRC_ERROR = 182 +NVML_FI_DEV_PCIE_COUNT_LANE_ERROR = 183 + +NVML_FI_DEV_IS_RESETLESS_MIG_SUPPORTED = 184 + +NVML_FI_DEV_POWER_AVERAGE = 185 +NVML_FI_DEV_POWER_INSTANT = 186 +NVML_FI_DEV_POWER_MIN_LIMIT = 187 +NVML_FI_DEV_POWER_MAX_LIMIT = 188 +NVML_FI_DEV_POWER_DEFAULT_LIMIT = 189 +NVML_FI_DEV_POWER_CURRENT_LIMIT = 190 +NVML_FI_DEV_ENERGY = 191 +NVML_FI_DEV_POWER_REQUESTED_LIMIT = 192 + +NVML_FI_DEV_TEMPERATURE_SHUTDOWN_TLIMIT = 193 +NVML_FI_DEV_TEMPERATURE_SLOWDOWN_TLIMIT = 194 +NVML_FI_DEV_TEMPERATURE_MEM_MAX_TLIMIT = 195 +NVML_FI_DEV_TEMPERATURE_GPU_MAX_TLIMIT = 196 + +NVML_FI_DEV_PCIE_COUNT_TX_BYTES = 197 +NVML_FI_DEV_PCIE_COUNT_RX_BYTES = 198 + +NVML_FI_DEV_IS_MIG_MODE_INDEPENDENT_MIG_QUERY_CAPABLE = 199 + +NVML_FI_DEV_NVLINK_GET_POWER_THRESHOLD_MAX = 200 + +NVML_FI_DEV_NVLINK_COUNT_XMIT_PACKETS = 201 +NVML_FI_DEV_NVLINK_COUNT_XMIT_BYTES = 202 +NVML_FI_DEV_NVLINK_COUNT_RCV_PACKETS = 203 +NVML_FI_DEV_NVLINK_COUNT_RCV_BYTES = 204 +NVML_FI_DEV_NVLINK_COUNT_VL15_DROPPED = 205 # Deprecated, do not use +NVML_FI_DEV_NVLINK_COUNT_MALFORMED_PACKET_ERRORS = 206 +NVML_FI_DEV_NVLINK_COUNT_BUFFER_OVERRUN_ERRORS = 207 +NVML_FI_DEV_NVLINK_COUNT_RCV_ERRORS = 208 +NVML_FI_DEV_NVLINK_COUNT_RCV_REMOTE_ERRORS = 209 +NVML_FI_DEV_NVLINK_COUNT_RCV_GENERAL_ERRORS = 210 +NVML_FI_DEV_NVLINK_COUNT_LOCAL_LINK_INTEGRITY_ERRORS = 211 +NVML_FI_DEV_NVLINK_COUNT_XMIT_DISCARDS = 212 + +NVML_FI_DEV_NVLINK_COUNT_LINK_RECOVERY_SUCCESSFUL_EVENTS = 213 +NVML_FI_DEV_NVLINK_COUNT_LINK_RECOVERY_FAILED_EVENTS = 214 +NVML_FI_DEV_NVLINK_COUNT_LINK_RECOVERY_EVENTS = 215 + +NVML_FI_DEV_NVLINK_COUNT_RAW_BER_LANE0 = 216 # Deprecated, do not use +NVML_FI_DEV_NVLINK_COUNT_RAW_BER_LANE1 = 217 # Deprecated, do not use +NVML_FI_DEV_NVLINK_COUNT_RAW_BER = 218 # Deprecated, do not use +NVML_FI_DEV_NVLINK_COUNT_EFFECTIVE_ERRORS = 219 +NVML_FI_DEV_NVLINK_COUNT_EFFECTIVE_BER = 220 +NVML_FI_DEV_NVLINK_COUNT_SYMBOL_ERRORS = 221 +NVML_FI_DEV_NVLINK_COUNT_SYMBOL_BER = 222 + +NVML_FI_DEV_NVLINK_GET_POWER_THRESHOLD_MIN = 223 +NVML_FI_DEV_NVLINK_GET_POWER_THRESHOLD_UNITS = 224 # Values are in the form NVML_NVLINK_LOW_POWER_THRESHOLD_UNIT_* +NVML_FI_DEV_NVLINK_GET_POWER_THRESHOLD_SUPPORTED = 225 + +NVML_FI_DEV_RESET_STATUS = 226 # Deprecated use NVML_FI_DEV_GET_GPU_RECOVERY_ACTION instead +NVML_FI_DEV_DRAIN_AND_RESET_STATUS = 227 # Deprecated use NVML_FI_DEV_GET_GPU_RECOVERY_ACTION instead +NVML_FI_DEV_PCIE_OUTBOUND_ATOMICS_MASK = 228 +NVML_FI_DEV_PCIE_INBOUND_ATOMICS_MASK = 229 +NVML_FI_DEV_GET_GPU_RECOVERY_ACTION = 230 + +NVML_FI_DEV_C2C_LINK_ERROR_INTR = 231 +NVML_FI_DEV_C2C_LINK_ERROR_REPLAY = 232 +NVML_FI_DEV_C2C_LINK_ERROR_REPLAY_B2B = 233 +NVML_FI_DEV_C2C_LINK_POWER_STATE = 234 + +NVML_FI_DEV_NVLINK_COUNT_FEC_HISTORY_0 = 235 +NVML_FI_DEV_NVLINK_COUNT_FEC_HISTORY_1 = 236 +NVML_FI_DEV_NVLINK_COUNT_FEC_HISTORY_2 = 237 +NVML_FI_DEV_NVLINK_COUNT_FEC_HISTORY_3 = 238 +NVML_FI_DEV_NVLINK_COUNT_FEC_HISTORY_4 = 239 +NVML_FI_DEV_NVLINK_COUNT_FEC_HISTORY_5 = 240 +NVML_FI_DEV_NVLINK_COUNT_FEC_HISTORY_6 = 241 +NVML_FI_DEV_NVLINK_COUNT_FEC_HISTORY_7 = 242 +NVML_FI_DEV_NVLINK_COUNT_FEC_HISTORY_8 = 243 +NVML_FI_DEV_NVLINK_COUNT_FEC_HISTORY_9 = 244 +NVML_FI_DEV_NVLINK_COUNT_FEC_HISTORY_10 = 245 +NVML_FI_DEV_NVLINK_COUNT_FEC_HISTORY_11 = 246 +NVML_FI_DEV_NVLINK_COUNT_FEC_HISTORY_12 = 247 +NVML_FI_DEV_NVLINK_COUNT_FEC_HISTORY_13 = 248 +NVML_FI_DEV_NVLINK_COUNT_FEC_HISTORY_14 = 249 +NVML_FI_DEV_NVLINK_COUNT_FEC_HISTORY_15 = 250 +NVML_FI_DEV_CLOCKS_EVENT_REASON_SW_POWER_CAP = NVML_FI_DEV_PERF_POLICY_POWER +NVML_FI_DEV_CLOCKS_EVENT_REASON_SYNC_BOOST = NVML_FI_DEV_PERF_POLICY_SYNC_BOOST +NVML_FI_DEV_CLOCKS_EVENT_REASON_SW_THERM_SLOWDOWN = 251 +NVML_FI_DEV_CLOCKS_EVENT_REASON_HW_THERM_SLOWDOWN = 252 +NVML_FI_DEV_CLOCKS_EVENT_REASON_HW_POWER_BRAKE_SLOWDOWN = 253 +NVML_FI_DEV_POWER_SYNC_BALANCING_FREQ = 254 +NVML_FI_DEV_POWER_SYNC_BALANCING_AF = 255 +NVML_FI_PWR_SMOOTHING_ENABLED = 256 # Enablement (0/DISABLED or 1/ENABLED) +NVML_FI_PWR_SMOOTHING_PRIV_LVL = 257 # Current privilege level +NVML_FI_PWR_SMOOTHING_IMM_RAMP_DOWN_ENABLED = 258 # Immediate ramp down enablement (0/DISABLED or 1/ENABLED) +NVML_FI_PWR_SMOOTHING_APPLIED_TMP_CEIL = 259 # Applied TMP ceiling value +NVML_FI_PWR_SMOOTHING_APPLIED_TMP_FLOOR = 260 # Applied TMP floor value +NVML_FI_PWR_SMOOTHING_MAX_PERCENT_TMP_FLOOR_SETTING = 261 # Max % TMP Floor value +NVML_FI_PWR_SMOOTHING_MIN_PERCENT_TMP_FLOOR_SETTING = 262 # Min % TMP Floor value +NVML_FI_PWR_SMOOTHING_HW_CIRCUITRY_PERCENT_LIFETIME_REMAINING = 263 # HW Circuitry % lifetime remaining +NVML_FI_PWR_SMOOTHING_MAX_NUM_PRESET_PROFILES = 264 # Max number of preset profiles +NVML_FI_PWR_SMOOTHING_PROFILE_PERCENT_TMP_FLOOR = 265 # % TMP floor for a given profile +NVML_FI_PWR_SMOOTHING_PROFILE_RAMP_UP_RATE = 266 # Ramp up rate in mW/s for a given profile +NVML_FI_PWR_SMOOTHING_PROFILE_RAMP_DOWN_RATE = 267 # Ramp down rate in mW/s for a given profile +NVML_FI_PWR_SMOOTHING_PROFILE_RAMP_DOWN_HYST_VAL = 268 # Ramp down hysteresis value in ms for a given profile +NVML_FI_PWR_SMOOTHING_ACTIVE_PRESET_PROFILE = 269 # Active preset profile number +NVML_FI_PWR_SMOOTHING_ADMIN_OVERRIDE_PERCENT_TMP_FLOOR = 270 # % TMP floor for a given profile +NVML_FI_PWR_SMOOTHING_ADMIN_OVERRIDE_RAMP_UP_RATE = 271 # Ramp up rate in mW/s for a given profile +NVML_FI_PWR_SMOOTHING_ADMIN_OVERRIDE_RAMP_DOWN_RATE = 272 # Ramp down rate in mW/s for a given profile +NVML_FI_PWR_SMOOTHING_ADMIN_OVERRIDE_RAMP_DOWN_HYST_VAL = 273 # Ramp down hysteresis value in ms for a given profile +NVML_FI_MAX = 274 # One greater than the largest field ID defined above + +# NVML_FI_DEV_NVLINK_GET_STATE state enums +NVML_NVLINK_STATE_INACTIVE = 0x0 +NVML_NVLINK_STATE_ACTIVE = 0x1 +NVML_NVLINK_STATE_SLEEP = 0x2 + +NVML_NVLINK_LOW_POWER_THRESHOLD_UNIT_100US = 0 # NVML_FI_DEV_NVLINK_GET_POWER_THRESHOLD_UNITS +NVML_NVLINK_LOW_POWER_THRESHOLD_UNIT_50US = 1 # NVML_FI_DEV_NVLINK_GET_POWER_THRESHOLD_UNITS + +# NVML_FI_DEV_C2C_LINK_POWER_STATE state enums +NVML_C2C_POWER_STATE_FULL_POWER = 0 +NVML_C2C_POWER_STATE_LOW_POWER = 1 + +## Enums needed for the method nvmlDeviceGetVirtualizationMode and nvmlDeviceSetVirtualizationMode +NVML_GPU_VIRTUALIZATION_MODE_NONE = 0 # Represents Bare Metal GPU +NVML_GPU_VIRTUALIZATION_MODE_PASSTHROUGH = 1 # Device is associated with GPU-Passthorugh +NVML_GPU_VIRTUALIZATION_MODE_VGPU = 2 # Device is associated with vGPU inside virtual machine. +NVML_GPU_VIRTUALIZATION_MODE_HOST_VGPU = 3 # Device is associated with VGX hypervisor in vGPU mode +NVML_GPU_VIRTUALIZATION_MODE_HOST_VSGA = 4 # Device is associated with VGX hypervisor in vSGA mode + +## Lib loading ## +nvmlLib = None +libLoadLock = threading.Lock() +_nvmlLib_refcount = 0 # Incremented on each nvmlInit and decremented on nvmlShutdown + +## vGPU Management +_nvmlVgpuTypeId_t = c_uint +_nvmlVgpuInstance_t = c_uint + +_nvmlVgpuVmIdType_t = c_uint +NVML_VGPU_VM_ID_DOMAIN_ID = 0 +NVML_VGPU_VM_ID_UUID = 1 + +_nvmlGridLicenseFeatureCode_t = c_uint +NVML_GRID_LICENSE_FEATURE_CODE_UNKNOWN = 0 +NVML_GRID_LICENSE_FEATURE_CODE_VGPU = 1 +NVML_GRID_LICENSE_FEATURE_CODE_NVIDIA_RTX = 2 +NVML_GRID_LICENSE_FEATURE_CODE_VWORKSTATION = 2 # deprecated, use NVML_GRID_LICENSE_FEATURE_CODE_NVIDIA_RTX. +NVML_GRID_LICENSE_FEATURE_CODE_GAMING = 3 +NVML_GRID_LICENSE_FEATURE_CODE_COMPUTE = 4 + +_nvmlGridLicenseExpiryStatus_t = c_uint8 +NVML_GRID_LICENSE_EXPIRY_NOT_AVAILABLE = 0, # Expiry information not available +NVML_GRID_LICENSE_EXPIRY_INVALID = 1, # Invalid expiry or error fetching expiry +NVML_GRID_LICENSE_EXPIRY_VALID = 2, # Valid expiry +NVML_GRID_LICENSE_EXPIRY_NOT_APPLICABLE = 3, # Expiry not applicable +NVML_GRID_LICENSE_EXPIRY_PERMANENT = 4, # Permanent expiry + +_nvmlVgpuCapability_t = c_uint +NVML_VGPU_CAP_NVLINK_P2P = 0 # vGPU P2P over NVLink is supported +NVML_VGPU_CAP_GPUDIRECT = 1 # GPUDirect capability is supported +NVML_VGPU_CAP_MULTI_VGPU_EXCLUSIVE = 2 # vGPU profile cannot be mixed with other vGPU profiles in same VM +NVML_VGPU_CAP_EXCLUSIVE_TYPE = 3 # vGPU profile cannot run on a GPU alongside other profiles of different type +NVML_VGPU_CAP_EXCLUSIVE_SIZE = 4 # vGPU profile cannot run on a GPU alongside other profiles of different size +NVML_VGPU_CAP_COUNT = 5 + +_nvmlVgpuDriverCapability_t = c_uint +NVML_VGPU_DRIVER_CAP_HETEROGENEOUS_MULTI_VGPU = 0 # Supports mixing of different vGPU profiles within one guest VM +NVML_VGPU_DRIVER_CAP_WARM_UPDATE = 1 # Supports FSR and warm update of vGPU host driver without terminating the running guest VM +NVML_VGPU_DRIVER_CAP_COUNT = 2 + +_nvmlDeviceVgpuCapability_t = c_uint +NVML_DEVICE_VGPU_CAP_FRACTIONAL_MULTI_VGPU = 0 # Query whether the fractional vGPU profiles on this GPU can be used in multi-vGPU configurations +NVML_DEVICE_VGPU_CAP_HETEROGENEOUS_TIMESLICE_PROFILES = 1 # Query whether the GPU supports concurrent execution of timesliced vGPU profiles of differing types +NVML_DEVICE_VGPU_CAP_HETEROGENEOUS_TIMESLICE_SIZES = 2 # Query whether the GPU supports concurrent execution of timesliced vGPU profiles of differing framebuffer sizes +NVML_DEVICE_VGPU_CAP_READ_DEVICE_BUFFER_BW = 3 # Query the GPU's read_device_buffer expected bandwidth capacity in megabytes per second +NVML_DEVICE_VGPU_CAP_WRITE_DEVICE_BUFFER_BW = 4 # Query the GPU's write_device_buffer expected bandwidth capacity in megabytes per second +NVML_DEVICE_VGPU_CAP_DEVICE_STREAMING = 5 # Query whether the vGPU profiles on the GPU supports migration data streaming +NVML_DEVICE_VGPU_CAP_MINI_QUARTER_GPU = 6 # Set/Get support of mini-quarter vGPU profiles +NVML_DEVICE_VGPU_CAP_COMPUTE_MEDIA_ENGINE_GPU = 7 # Set/Get support for compute media engine vGPU profiles +NVML_DEVICE_VGPU_CAP_WARM_UPDATE = 8 # Query whether the GPU supports FSR and warm update +NVML_DEVICE_VGPU_CAP_HOMOGENEOUS_PLACEMENTS = 9 # Query whether the GPU supports reporting of placements of timesliced vGPU profiles with identical framebuffer sizes +NVML_DEVICE_VGPU_CAP_MIG_TIMESLICING_SUPPORTED = 10 # Query whether the GPU supports timesliced vGPU on MIG +NVML_DEVICE_VGPU_CAP_MIG_TIMESLICING_ENABLED = 11 # Set/Get MIG timesliced mode reporting, without impacting the underlying functionality +NVML_DEVICE_VGPU_CAP_COUNT = 12 + +_nvmlVgpuGuestInfoState_t = c_uint +NVML_VGPU_INSTANCE_GUEST_INFO_STATE_UNINITIALIZED = 0 +NVML_VGPU_INSTANCE_GUEST_INFO_STATE_INITIALIZED = 1 + +_nvmlVgpuVmCompatibility_t = c_uint +NVML_VGPU_VM_COMPATIBILITY_NONE = 0x0 +NVML_VGPU_VM_COMPATIBILITY_COLD = 0x1 +NVML_VGPU_VM_COMPATIBILITY_HIBERNATE = 0x2 +NVML_VGPU_VM_COMPATIBILITY_SLEEP = 0x4 +NVML_VGPU_VM_COMPATIBILITY_LIVE = 0x8 + +_nvmlVgpuPgpuCompatibilityLimitCode_t = c_uint +NVML_VGPU_COMPATIBILITY_LIMIT_NONE = 0x0 +NVML_VGPU_COMPATIBILITY_LIMIT_HOST_DRIVER = 0x1 +NVML_VGPU_COMPATIBILITY_LIMIT_GUEST_DRIVER = 0x2 +NVML_VGPU_COMPATIBILITY_LIMIT_GPU = 0x4 +NVML_VGPU_COMPATIBILITY_LIMIT_OTHER = 0x80000000 + +_nvmlHostVgpuMode_t = c_uint +NVML_HOST_VGPU_MODE_NON_SRIOV = 0 +NVML_HOST_VGPU_MODE_SRIOV = 1 + +_nvmlConfComputeGpusReadyState_t = c_uint +NVML_CC_ACCEPTING_CLIENT_REQUESTS_FALSE = 0 +NVML_CC_ACCEPTING_CLIENT_REQUESTS_TRUE = 1 + +_nvmlConfComputeGpuCaps_t = c_uint +NVML_CC_SYSTEM_GPUS_CC_NOT_CAPABLE = 0 +NVML_CC_SYSTEM_GPUS_CC_CAPABLE = 1 + +_nvmlConfComputeCpuCaps_t = c_uint +NVML_CC_SYSTEM_CPU_CAPS_NONE = 0 +NVML_CC_SYSTEM_CPU_CAPS_AMD_SEV = 1 +NVML_CC_SYSTEM_CPU_CAPS_INTEL_TDX = 2 +NVML_CC_SYSTEM_CPU_CAPS_AMD_SEV_SNP = 3 +NVML_CC_SYSTEM_CPU_CAPS_AMD_SNP_VTOM = 4 + +_nvmlConfComputeDevToolsMode_t = c_uint +NVML_CC_SYSTEM_DEVTOOLS_MODE_OFF = 0 +NVML_CC_SYSTEM_DEVTOOLS_MODE_ON = 1 + +NVML_CC_SYSTEM_MULTIGPU_NONE = 0 +NVML_CC_SYSTEM_MULTIGPU_PROTECTED_PCIE = 1 +NVML_CC_SYSTEM_MULTIGPU_NVLE = 2 + +NVML_CC_SYSTEM_ENVIRONMENT_UNAVAILABLE = 0 +NVML_CC_SYSTEM_ENVIRONMENT_SIM = 1 +NVML_CC_SYSTEM_ENVIRONMENT_PROD = 2 + +_nvmlConfComputeCcFeature_t = c_uint +NVML_CC_SYSTEM_FEATURE_DISABLED = 0 +NVML_CC_SYSTEM_FEATURE_ENABLED = 1 + +_nvmlConfComputeCcKeyRotationThreshAttackerAdv_t = c_uint +NVML_CC_KEY_ROTATION_THRESH_ATTACKER_ADVANTAGE_MIN = 50 +NVML_CC_KEY_ROTATION_THRESH_ATTACKER_ADVANTAGE_MAX = 65 + +# GSP firmware +NVML_GSP_FIRMWARE_VERSION_BUF_SIZE = 0x40 + +class NVMLLibraryMismatchError(Exception): + pass + +## Error Checking ## +class NVMLError(Exception): + _valClassMapping = dict() + # List of currently known error codes + _errcode_to_string = { + NVML_ERROR_UNINITIALIZED: "Uninitialized", + NVML_ERROR_INVALID_ARGUMENT: "Invalid Argument", + NVML_ERROR_NOT_SUPPORTED: "Not Supported", + NVML_ERROR_NO_PERMISSION: "Insufficient Permissions", + NVML_ERROR_ALREADY_INITIALIZED: "Already Initialized", + NVML_ERROR_NOT_FOUND: "Not Found", + NVML_ERROR_INSUFFICIENT_SIZE: "Insufficient Size", + NVML_ERROR_INSUFFICIENT_POWER: "Insufficient External Power", + NVML_ERROR_DRIVER_NOT_LOADED: "Driver Not Loaded", + NVML_ERROR_TIMEOUT: "Timeout", + NVML_ERROR_IRQ_ISSUE: "Interrupt Request Issue", + NVML_ERROR_LIBRARY_NOT_FOUND: "NVML Shared Library Not Found", + NVML_ERROR_FUNCTION_NOT_FOUND: "Function Not Found", + NVML_ERROR_CORRUPTED_INFOROM: "Corrupted infoROM", + NVML_ERROR_GPU_IS_LOST: "GPU is lost", + NVML_ERROR_RESET_REQUIRED: "GPU requires restart", + NVML_ERROR_OPERATING_SYSTEM: "The operating system has blocked the request.", + NVML_ERROR_LIB_RM_VERSION_MISMATCH: "RM has detected an NVML/RM version mismatch.", + NVML_ERROR_MEMORY: "Insufficient Memory", + NVML_ERROR_UNKNOWN: "Unknown Error", + } + def __new__(typ, value): + ''' + Maps value to a proper subclass of NVMLError. + See _extractNVMLErrorsAsClasses function for more details + ''' + if typ == NVMLError: + typ = NVMLError._valClassMapping.get(value, typ) + obj = Exception.__new__(typ) + obj.value = value + return obj + def __str__(self): + try: + if self.value not in NVMLError._errcode_to_string: + NVMLError._errcode_to_string[self.value] = str(nvmlErrorString(self.value)) + return NVMLError._errcode_to_string[self.value] + except NVMLError: + return "NVML Error with code %d" % self.value + def __eq__(self, other): + return self.value == other.value + +def nvmlExceptionClass(nvmlErrorCode): + if nvmlErrorCode not in NVMLError._valClassMapping: + raise ValueError('nvmlErrorCode %s is not valid' % nvmlErrorCode) + return NVMLError._valClassMapping[nvmlErrorCode] + +def _extractNVMLErrorsAsClasses(): + ''' + Generates a hierarchy of classes on top of NVMLError class. + + Each NVML Error gets a new NVMLError subclass. This way try,except blocks can filter appropriate + exceptions more easily. + + NVMLError is a parent class. Each NVML_ERROR_* gets it's own subclass. + e.g. NVML_ERROR_ALREADY_INITIALIZED will be turned into NVMLError_AlreadyInitialized + ''' + this_module = sys.modules[__name__] + nvmlErrorsNames = [x for x in dir(this_module) if x.startswith("NVML_ERROR_")] + for err_name in nvmlErrorsNames: + # e.g. Turn NVML_ERROR_ALREADY_INITIALIZED into NVMLError_AlreadyInitialized + class_name = "NVMLError_" + string.capwords(err_name.replace("NVML_ERROR_", ""), "_").replace("_", "") + err_val = getattr(this_module, err_name) + def gen_new(val): + def new(typ, *args): + obj = NVMLError.__new__(typ, val) + return obj + return new + new_error_class = type(class_name, (NVMLError,), {'__new__': gen_new(err_val)}) + new_error_class.__module__ = __name__ + setattr(this_module, class_name, new_error_class) + NVMLError._valClassMapping[err_val] = new_error_class +_extractNVMLErrorsAsClasses() + +def _nvmlCheckReturn(ret): + if (ret != NVML_SUCCESS): + raise NVMLError(ret) + return ret + +## Function access ## +_nvmlGetFunctionPointer_cache = dict() # function pointers are cached to prevent unnecessary libLoadLock locking +def _nvmlGetFunctionPointer(name): + global nvmlLib + + if name in _nvmlGetFunctionPointer_cache: + return _nvmlGetFunctionPointer_cache[name] + + libLoadLock.acquire() + try: + # ensure library was loaded + if (nvmlLib == None): + raise NVMLError(NVML_ERROR_UNINITIALIZED) + try: + _nvmlGetFunctionPointer_cache[name] = getattr(nvmlLib, name) + return _nvmlGetFunctionPointer_cache[name] + except AttributeError: + raise NVMLError(NVML_ERROR_FUNCTION_NOT_FOUND) + finally: + # lock is always freed + libLoadLock.release() + +## Alternative object +# Allows the object to be printed +# Allows mismatched types to be assigned +# - like None when the Structure variant requires c_uint +class nvmlFriendlyObject(object): + def __init__(self, dictionary): + for x in dictionary: + setattr(self, x, dictionary[x]) + def __str__(self): + return self.__dict__.__str__() + +def nvmlStructToFriendlyObject(struct): + d = {} + for x in struct._fields_: + key = x[0] + value = getattr(struct, key) + # only need to convert from bytes if bytes, no need to check python version. + d[key] = value.decode() if isinstance(value, bytes) else value + obj = nvmlFriendlyObject(d) + return obj + +# pack the object so it can be passed to the NVML library +def nvmlFriendlyObjectToStruct(obj, model): + for x in model._fields_: + key = x[0] + value = obj.__dict__[key] + # any c_char_p in python3 needs to be bytes, default encoding works fine. + if sys.version_info >= (3,): + setattr(model, key, value.encode()) + else: + setattr(model, key, value) + return model + +## Unit structures +class struct_c_nvmlUnit_t(Structure): + pass # opaque handle +c_nvmlUnit_t = POINTER(struct_c_nvmlUnit_t) + +class _PrintableStructure(Structure): + """ + Abstract class that produces nicer __str__ output than ctypes.Structure. + e.g. instead of: + + this class will print + class_name(field_name: formatted_value, field_name: formatted_value) + + _fmt_ dictionary of -> + e.g. class that has _field_ 'hex_value', c_uint could be formatted with + _fmt_ = {"hex_value" : "%08X"} + to produce nicer output. + Default fomratting string for all fields can be set with key "" like: + _fmt_ = {"" : "%d MHz"} # e.g all values are numbers in MHz. + If not set it's assumed to be just "%s" + + Exact format of returned str from this class is subject to change in the future. + """ + _fmt_ = {} + def __str__(self): + result = [] + for x in self._fields_: + key = x[0] + value = getattr(self, key) + fmt = "%s" + if key in self._fmt_: + fmt = self._fmt_[key] + elif "" in self._fmt_: + fmt = self._fmt_[""] + result.append(("%s: " + fmt) % (key, value)) + return self.__class__.__name__ + "(" + ", ".join(result) + ")" + + def __getattribute__(self, name): + res = super(_PrintableStructure, self).__getattribute__(name) + # need to convert bytes to unicode for python3 don't need to for python2 + # Python 2 strings are of both str and bytes + # Python 3 strings are not of type bytes + # ctypes should convert everything to the correct values otherwise + if isinstance(res, bytes): + if isinstance(res, str): + return res + return res.decode() + return res + + def __setattr__(self, name, value): + if isinstance(value, str): + # encoding a python2 string returns the same value, since python2 strings are bytes already + # bytes passed in python3 will be ignored. + value = value.encode() + super(_PrintableStructure, self).__setattr__(name, value) + +class c_nvmlUnitInfo_t(_PrintableStructure): + _fields_ = [ + ('name', c_char * 96), + ('id', c_char * 96), + ('serial', c_char * 96), + ('firmwareVersion', c_char * 96), + ] + +class c_nvmlC2cModeInfo_v1_t(_PrintableStructure): + _fields_ = [ + ('isC2cEnabled', c_uint) + ] + +nvmlC2cModeInfo_v1 = 0x1000008; + +class c_nvmlLedState_t(_PrintableStructure): + _fields_ = [ + ('cause', c_char * 256), + ('color', _nvmlLedColor_t), + ] + +class c_nvmlPSUInfo_t(_PrintableStructure): + _fields_ = [ + ('state', c_char * 256), + ('current', c_uint), + ('voltage', c_uint), + ('power', c_uint), + ] + +class c_nvmlUnitFanInfo_t(_PrintableStructure): + _fields_ = [ + ('speed', c_uint), + ('state', _nvmlFanState_t), + ] + +class c_nvmlUnitFanSpeeds_t(_PrintableStructure): + _fields_ = [ + ('fans', c_nvmlUnitFanInfo_t * 24), + ('count', c_uint) + ] + +## Device structures +class struct_c_nvmlDevice_t(Structure): + pass # opaque handle +c_nvmlDevice_t = POINTER(struct_c_nvmlDevice_t) + +class nvmlPciInfoExt_v1_t(_PrintableStructure): + _fields_ = [ + ('version', c_uint), + ('domain', c_uint), + ('bus', c_uint), + ('device', c_uint), + ('pciDeviceId', c_uint), + ('pciSubSystemId', c_uint), + ('baseClass', c_uint), + ('subClass', c_uint), + ('busId', c_char * NVML_DEVICE_PCI_BUS_ID_BUFFER_SIZE), + ] + _fmt_ = { + 'version' : "0x%04X", + 'domain' : "0x%04X", + 'bus' : "0x%02X", + 'device' : "0x%02X", + 'pciDeviceId' : "0x%08X", + 'pciSubSystemId' : "0x%08X", + 'baseClass' : "0x%01X", + 'subClass' : "0x%01X", + } + +nvmlPciInfoExt_v1 = 0x1000040 + +# Legacy pciInfo used for _v1 and _v2 +class nvmlPciInfo_v2_t(_PrintableStructure): + _fields_ = [ + ('busId', c_char * NVML_DEVICE_PCI_BUS_ID_BUFFER_V2_SIZE), + ('domain', c_uint), + ('bus', c_uint), + ('device', c_uint), + ('pciDeviceId', c_uint), + + # Added in 2.285 + ('pciSubSystemId', c_uint), + ('reserved0', c_uint), + ('reserved1', c_uint), + ('reserved2', c_uint), + ('reserved3', c_uint), + ] + _fmt_ = { + 'domain' : "0x%04X", + 'bus' : "0x%02X", + 'device' : "0x%02X", + 'pciDeviceId' : "0x%08X", + 'pciSubSystemId' : "0x%08X", + } + +class nvmlPciInfo_t(_PrintableStructure): + _fields_ = [ + # Moved to the new busId location below + ('busIdLegacy', c_char * NVML_DEVICE_PCI_BUS_ID_BUFFER_V2_SIZE), + ('domain', c_uint), + ('bus', c_uint), + ('device', c_uint), + ('pciDeviceId', c_uint), + + # Added in 2.285 + ('pciSubSystemId', c_uint), + # New busId replaced the long deprecated and reserved fields with a + # field of the same size in 9.0 + ('busId', c_char * NVML_DEVICE_PCI_BUS_ID_BUFFER_SIZE), + ] + _fmt_ = { + 'domain' : "0x%08X", + 'bus' : "0x%02X", + 'device' : "0x%02X", + 'pciDeviceId' : "0x%08X", + 'pciSubSystemId' : "0x%08X", + } + +class c_nvmlSystemDriverBranchInfo_v1_t(_PrintableStructure): + _fields_ = [ + ('version', c_uint), + ("branch", c_char * NVML_SYSTEM_DRIVER_VERSION_BUFFER_SIZE), + ] + +SystemDriverBranchInfo_v1 = 0x1000054 + +class c_nvmlExcludedDeviceInfo_t(_PrintableStructure): + _fields_ = [ + ('pci', nvmlPciInfo_t), + ('uuid', c_char * NVML_DEVICE_UUID_BUFFER_SIZE) + ] + +class nvmlNvLinkUtilizationControl_t(_PrintableStructure): + _fields_ = [ + ('units', _nvmlNvLinkUtilizationCountUnits_t), + ('pktfilter', _nvmlNvLinkUtilizationCountPktTypes_t), + ] + +class c_nvmlMemory_t(_PrintableStructure): + _fields_ = [ + ('total', c_ulonglong), + ('free', c_ulonglong), + ('used', c_ulonglong), + ] + _fmt_ = {'': "%d B"} + +class c_nvmlMemory_v2_t(_PrintableStructure): + _fields_ = [ + ('version', c_uint), + ('total', c_ulonglong), + ('reserved', c_ulonglong), + ('free', c_ulonglong), + ('used', c_ulonglong), + ] + _fmt_ = {'': "%d B"} + +nvmlMemory_v2 = 0x02000028 + +class c_nvmlBAR1Memory_t(_PrintableStructure): + _fields_ = [ + ('bar1Total', c_ulonglong), + ('bar1Free', c_ulonglong), + ('bar1Used', c_ulonglong), + ] + _fmt_ = {'': "%d B"} + +class nvmlClkMonFaultInfo_t(Structure): + _fields_ = [("clkApiDomain", c_uint), + ("clkDomainFaultMask", c_uint) + ] + +MAX_CLK_DOMAINS = 32 + +class nvmlClkMonStatus_t(Structure): + _fields_ = [("bGlobalStatus", c_uint), + ("clkMonListSize", c_uint), + ("clkMonList", nvmlClkMonFaultInfo_t * MAX_CLK_DOMAINS) + ] + +# On Windows with the WDDM driver, usedGpuMemory is reported as None +# Code that processes this structure should check for None, I.E. +# +# if (info.usedGpuMemory == None): +# # TODO handle the error +# pass +# else: +# print("Using %d MiB of memory" % (info.usedGpuMemory / 1024 / 1024)) +# endif +# +# See NVML documentation for more information +class c_nvmlProcessInfo_v2_t(_PrintableStructure): + _fields_ = [ + ('pid', c_uint), + ('usedGpuMemory', c_ulonglong), + ('gpuInstanceId', c_uint), + ('computeInstanceId', c_uint), + ] + _fmt_ = {'usedGpuMemory': "%d B"} + +c_nvmlProcessInfo_v3_t = c_nvmlProcessInfo_v2_t + +c_nvmlProcessInfo_t = c_nvmlProcessInfo_v3_t + +_nvmlProcessMode_t = c_uint +NVML_PROCESS_MODE_COMPUTE = 0 +NVML_PROCESS_MODE_GRAPHICS = 1 +NVML_PROCESS_MODE_MPS = 2 + +class c_nvmlProcessDetail_v1_t(Structure): + _fields_ = [ + ('pid', c_uint), + ('usedGpuMemory', c_ulonglong), + ('gpuInstanceId', c_uint), + ('computeInstanceId', c_uint), + ('usedGpuCcProtectedMemory', c_ulonglong), + ] + +class c_nvmlProcessDetailList_v1_t(_PrintableStructure): + _fields_ = [ + ('version', c_uint), + ('mode', _nvmlProcessMode_t), + ('numProcArrayEntries', c_uint), + ('procArray', POINTER(c_nvmlProcessDetail_v1_t)), + ] + _fmt_ = {'numProcArrayEntries': "%d B"} + +c_nvmlProcessDetailList_t = c_nvmlProcessDetailList_v1_t + +nvmlProcessDetailList_v1 = 0x1000018 + +class c_nvmlBridgeChipInfo_t(_PrintableStructure): + _fields_ = [ + ('type', _nvmlBridgeChipType_t), + ('fwVersion', c_uint), + ] + +class c_nvmlBridgeChipHierarchy_t(_PrintableStructure): + _fields_ = [ + ('bridgeCount', c_uint), + ('bridgeChipInfo', c_nvmlBridgeChipInfo_t * 128), + ] + +class c_nvmlEccErrorCounts_t(_PrintableStructure): + _fields_ = [ + ('l1Cache', c_ulonglong), + ('l2Cache', c_ulonglong), + ('deviceMemory', c_ulonglong), + ('registerFile', c_ulonglong), + ] + +class c_nvmlUtilization_t(_PrintableStructure): + _fields_ = [ + ('gpu', c_uint), + ('memory', c_uint), + ] + _fmt_ = {'': "%d %%"} + +# Added in 2.285 +class c_nvmlHwbcEntry_t(_PrintableStructure): + _fields_ = [ + ('hwbcId', c_uint), + ('firmwareVersion', c_char * 32), + ] + +class c_nvmlValue_t(Union): + _fields_ = [ + ('dVal', c_double), + ('uiVal', c_uint), + ('ulVal', c_ulong), + ('ullVal', c_ulonglong), + ('sllVal', c_longlong), + ('siVal', c_int), + ('usVal', c_ushort), + ] + +class c_nvmlSample_t(_PrintableStructure): + _fields_ = [ + ('timeStamp', c_ulonglong), + ('sampleValue', c_nvmlValue_t), + ] + +class c_nvmlViolationTime_t(_PrintableStructure): + _fields_ = [ + ('referenceTime', c_ulonglong), + ('violationTime', c_ulonglong), + ] + +class c_nvmlFieldValue_t(_PrintableStructure): + _fields_ = [ + ('fieldId', c_uint32), + ('scopeId', c_uint32), + ('timestamp', c_int64), + ('latencyUsec', c_int64), + ('valueType', _nvmlValueType_t), + ('nvmlReturn', _nvmlReturn_t), + ('value', c_nvmlValue_t) + ] + +NVML_NVLINK_TOTAL_SUPPORTED_BW_MODES = 23 + +nvmlNvlinkSupportedBwModes_v1 = 0x100001c +class c_nvmlNvlinkSupportedBwModes_v1_t(_PrintableStructure): + _fields_ = [ + ('version', c_uint), + ('bwModes', c_uint8 * NVML_NVLINK_TOTAL_SUPPORTED_BW_MODES), + ('totalBwModes', c_uint8) + ] + + def __init__(self): + super(c_nvmlNvlinkSupportedBwModes_v1_t, self).__init__(version=nvmlNvlinkSupportedBwModes_v1) + +nvmlNvlinkGetBwMode_v1 = 0x100000c +class c_nvmlNvlinkGetBwMode_v1_t(_PrintableStructure): + _fields_ = [ + ('version', c_uint), + ('bIsBest', c_uint), + ('bwMode', c_uint8) + ] + + def __init__(self): + super(c_nvmlNvlinkGetBwMode_v1_t, self).__init__(version=nvmlNvlinkGetBwMode_v1) + +nvmlNvlinkSetBwMode_v1 = 0x100000c +class c_nvmlNvlinkSetBwMode_v1_t(_PrintableStructure): + _fields_ = [ + ('version', c_uint), + ('bSetBest', c_uint), + ('bwMode', c_uint8) + ] + + def __init__(self): + super(c_nvmlNvlinkSetBwMode_v1_t, self).__init__(version=nvmlNvlinkSetBwMode_v1) + +class c_nvmlVgpuHeterogeneousMode_v1_t(_PrintableStructure): + _fields_ = [ + ('version', c_uint), + ('mode', c_uint), + ] + +VgpuHeterogeneousMode_v1 = 0x1000008 + +class c_nvmlVgpuPlacementId_v1_t(_PrintableStructure): + _fields_ = [ + ('version', c_uint), + ('placementId', c_uint), + ] + +VgpuPlacementId_v1 = 0x1000008 + +class c_nvmlVgpuPlacementList_v1_t(_PrintableStructure): + _fields_ = [ + ('version', c_uint), + ('count', c_uint), + ('placementSize', c_uint), + ('placementIds', POINTER(c_uint)), + ] + +VgpuPlacementList_v1 = 0x1000018 + +NVML_VGPU_PGPU_HETEROGENEOUS_MODE = 0 +NVML_VGPU_PGPU_HOMOGENEOUS_MODE = 1 + +class c_nvmlVgpuPlacementList_v2_t(_PrintableStructure): + _fields_ = [ + ('version', c_uint), + ('placementSize', c_uint), + ('count', c_uint), + ('placementIds', POINTER(c_uint)), + ('mode', c_uint), + ] + +VgpuPlacementList_v2 = 0x2000020 + +class c_nvmlVgpuTypeBar1Info_v1_t(_PrintableStructure): + _fields_ = [ + ('version', c_uint), + ('bar1Size', c_ulonglong), + ] + +VgpuTypeBar1Info_v1 = 0x1000010 + +class c_nvmlVgpuInstanceUtilizationSample_t(_PrintableStructure): + _fields_ = [ + ('vgpuInstance', _nvmlVgpuInstance_t), + ('timeStamp', c_ulonglong), + ('smUtil', c_nvmlValue_t), + ('memUtil', c_nvmlValue_t), + ('encUtil', c_nvmlValue_t), + ('decUtil', c_nvmlValue_t), + ] + +class c_nvmlVgpuInstanceUtilizationInfo_v1_t(_PrintableStructure): + _fields_ = [ + ('timeStamp', c_ulonglong), + ('vgpuInstance', _nvmlVgpuInstance_t), + ('smUtil', c_nvmlValue_t), + ('memUtil', c_nvmlValue_t), + ('encUtil', c_nvmlValue_t), + ('decUtil', c_nvmlValue_t), + ('jpgUtil', c_nvmlValue_t), + ('ofaUtil', c_nvmlValue_t), + ] + +class c_nvmlVgpuInstancesUtilizationInfo_v1_t(_PrintableStructure): + _fields_ = [ + ('version', c_uint), + ('sampleValType', _nvmlValueType_t), + ('vgpuInstanceCount', c_uint), + ('lastSeenTimeStamp', c_ulonglong), + ('vgpuUtilArray', POINTER(c_nvmlVgpuInstanceUtilizationInfo_v1_t)), + ] + +VgpuInstancesUtilizationInfo_v1 = 0x01000020 + +class c_nvmlVgpuProcessUtilizationSample_t(_PrintableStructure): + _fields_ = [ + ('vgpuInstance', _nvmlVgpuInstance_t), + ('pid', c_uint), + ('processName', c_char * NVML_VGPU_NAME_BUFFER_SIZE), + ('timeStamp', c_ulonglong), + ('smUtil', c_uint), + ('memUtil', c_uint), + ('encUtil', c_uint), + ('decUtil', c_uint), + ] + +class c_nvmlVgpuProcessUtilizationInfo_v1_t(_PrintableStructure): + _fields_ = [ + ('processName', c_char * NVML_VGPU_NAME_BUFFER_SIZE), + ('timeStamp', c_ulonglong), + ('vgpuInstance', _nvmlVgpuInstance_t), + ('pid', c_uint), + ('smUtil', c_uint), + ('memUtil', c_uint), + ('encUtil', c_uint), + ('decUtil', c_uint), + ('jpgUtil', c_uint), + ('ofaUtil', c_uint), + ] + +class c_nvmlVgpuProcessesUtilizationInfo_v1_t(_PrintableStructure): + _fields_ = [ + ('version', c_uint), + ('vgpuProcessCount', c_uint), + ('lastSeenTimeStamp', c_ulonglong), + ('vgpuProcUtilArray', POINTER(c_nvmlVgpuProcessUtilizationInfo_v1_t)), + ] + +VgpuProcessesUtilizationInfo_v1 = 0x01000018 + +class nvmlVgpuRuntimeState_v1_t(_PrintableStructure): + _fields_ = [ + ('version', c_uint), + ('size', c_ulonglong), + ] + +VgpuRuntimeState_v1 = 0x1000010 + +class c_nvmlVgpuLicenseExpiry_t(_PrintableStructure): + _fields_ = [ + ('year', c_uint32), + ('month', c_uint16), + ('day', c_uint16), + ('hour', c_uint16), + ('min', c_uint16), + ('sec', c_uint16), + ('status', c_uint8), + ] + +NVML_GRID_LICENSE_STATE_UNKNOWN = 0 +NVML_GRID_LICENSE_STATE_UNINITIALIZED = 1 +NVML_GRID_LICENSE_STATE_UNLICENSED_UNRESTRICTED = 2 +NVML_GRID_LICENSE_STATE_UNLICENSED_RESTRICTED = 3 +NVML_GRID_LICENSE_STATE_UNLICENSED = 4 +NVML_GRID_LICENSE_STATE_LICENSED = 5 + +class c_nvmlVgpuLicenseInfo_t(_PrintableStructure): + _fields_ = [ + ('isLicensed', c_uint8), + ('licenseExpiry', c_nvmlVgpuLicenseExpiry_t), + ('currentState', c_uint), + ] + +class c_nvmlEncoderSession_t(_PrintableStructure): + _fields_ = [ + ('sessionId', c_uint), + ('pid', c_uint), + ('vgpuInstance', _nvmlVgpuInstance_t), + ('codecType', c_uint), + ('hResolution', c_uint), + ('vResolution', c_uint), + ('averageFps', c_uint), + ('encodeLatency', c_uint), + ] + +class c_nvmlProcessUtilizationSample_t(_PrintableStructure): + _fields_ = [ + ('pid', c_uint), + ('timeStamp', c_ulonglong), + ('smUtil', c_uint), + ('memUtil', c_uint), + ('encUtil', c_uint), + ('decUtil', c_uint), + ] + +class c_nvmlProcessUtilizationInfo_v1_t(_PrintableStructure): + _fields_ = [ + ('timeStamp', c_ulonglong), + ('pid', c_uint), + ('smUtil', c_uint), + ('memUtil', c_uint), + ('encUtil', c_uint), + ('decUtil', c_uint), + ('jpgUtil', c_uint), + ('ofaUtil', c_uint), + ] + +class c_nvmlProcessesUtilizationInfo_v1_t(_PrintableStructure): + _fields_ = [ + ('version', c_uint), + ('processSamplesCount', c_uint), + ('lastSeenTimeStamp', c_ulonglong), + ('procUtilArray', POINTER(c_nvmlProcessUtilizationInfo_v1_t)), + ] + +ProcessesUtilizationInfo_v1 = 0x01000018 + +class c_nvmlGridLicenseExpiry_t(_PrintableStructure): + _fields_ = [ + ('year', c_uint32), + ('month', c_uint16), + ('day', c_uint16), + ('hour', c_uint16), + ('min', c_uint16), + ('sec', c_uint16), + ('status', c_uint8), + ] + +class c_nvmlGridLicensableFeature_v4_t(_PrintableStructure): + _fields_ = [ + ('featureCode', _nvmlGridLicenseFeatureCode_t), + ('featureState', c_uint), + ('licenseInfo', c_char * NVML_GRID_LICENSE_BUFFER_SIZE), + ('productName', c_char * NVML_GRID_LICENSE_BUFFER_SIZE), + ('featureEnabled', c_uint), + ('licenseExpiry', c_nvmlGridLicenseExpiry_t), + ] + +class c_nvmlGridLicensableFeatures_v4_t(_PrintableStructure): + _fields_ = [ + ('isGridLicenseSupported', c_int), + ('licensableFeaturesCount', c_uint), + ('gridLicensableFeatures', c_nvmlGridLicensableFeature_v4_t * NVML_GRID_LICENSE_FEATURE_MAX_COUNT), + ] + +class c_nvmlGridLicensableFeature_v3_t(_PrintableStructure): + _fields_ = [ + ('featureCode', _nvmlGridLicenseFeatureCode_t), + ('featureState', c_uint), + ('licenseInfo', c_char * NVML_GRID_LICENSE_BUFFER_SIZE), + ('productName', c_char * NVML_GRID_LICENSE_BUFFER_SIZE), + ('featureEnabled', c_uint), + ] + +class c_nvmlGridLicensableFeatures_v3_t(_PrintableStructure): + _fields_ = [ + ('isGridLicenseSupported', c_int), + ('licensableFeaturesCount', c_uint), + ('gridLicensableFeatures', c_nvmlGridLicensableFeature_v3_t * NVML_GRID_LICENSE_FEATURE_MAX_COUNT), + ] + +class c_nvmlGridLicensableFeature_v2_t(_PrintableStructure): + _fields_ = [ + ('featureCode', _nvmlGridLicenseFeatureCode_t), + ('featureState', c_uint), + ('licenseInfo', c_char * NVML_GRID_LICENSE_BUFFER_SIZE), + ('productName', c_char * NVML_GRID_LICENSE_BUFFER_SIZE), + ] + +class c_nvmlGridLicensableFeatures_v2_t(_PrintableStructure): + _fields_ = [ + ('isGridLicenseSupported', c_int), + ('licensableFeaturesCount', c_uint), + ('gridLicensableFeatures', c_nvmlGridLicensableFeature_v2_t * NVML_GRID_LICENSE_FEATURE_MAX_COUNT), + ] + +class c_nvmlGridLicensableFeature_t(_PrintableStructure): + _fields_ = [ + ('featureCode', _nvmlGridLicenseFeatureCode_t), + ('featureState', c_uint), + ('licenseInfo', c_char * NVML_GRID_LICENSE_BUFFER_SIZE), + ] + +class c_nvmlGridLicensableFeatures_t(_PrintableStructure): + _fields_ = [ + ('isGridLicenseSupported', c_int), + ('licensableFeaturesCount', c_uint), + ('gridLicensableFeatures', c_nvmlGridLicensableFeature_t * NVML_GRID_LICENSE_FEATURE_MAX_COUNT), + ] + +class c_nvmlMarginTemperature_v1_t(_PrintableStructure): + _fields_ = [ + ('version', c_uint), + ('marginTemperature', c_int), + ] + +nvmlMarginTemperature_v1 = 0x1000008 + +NVML_DEVICE_UUID_ASCII_LEN = 41 +NVML_DEVICE_UUID_BINARY_LEN = 16 + +NVML_UUID_TYPE_NONE = 0 +NVML_UUID_TYPE_ASCII = 1 +NVML_UUID_TYPE_BINARY = 2 + +class c_nvmlUUIDValue_t(Union): + _fields_ = [ + ('str', c_char * NVML_DEVICE_UUID_ASCII_LEN), + ('bytes', c_ubyte * NVML_DEVICE_UUID_BINARY_LEN), + ] + +nvmlUUID_v1 = 0x1000034 +class c_nvmlUUID_t(_PrintableStructure): + _fields_ = [ + ('version', c_uint), + ('type', c_uint), + ('value', c_nvmlUUIDValue_t), + ] + + def __init__(self): + super(c_nvmlUUID_t, self).__init__(version=nvmlUUID_v1) + +nvmlPdi_v1 = 0x1000010 +class c_nvmlPdi_t(_PrintableStructure): + _fields_ = [ + ('version', c_uint), + ('value', c_ulonglong), + ] + + def __init__(self): + super(c_nvmlPdi_t, self).__init__(version=nvmlPdi_v1) + +nvmlRepairStatus_v1 = 0x100000C +class c_nvmlRepairStatus_t(_PrintableStructure): + _fields_ = [ + ("version", c_uint), + ("bChannelRepairPending", c_uint), + ("bTpcRepairPending", c_uint), + ] + + def __init__(self): + super(c_nvmlRepairStatus_t, self).__init__(version=nvmlRepairStatus_v1) + +nvmlNvLinkInfo_v1 =0x1000008 +class c_nvmlNvLinkInfo_v1_t(_PrintableStructure): + _fields_ = [ + ("version", c_uint), + ("isNvleEnabled", c_uint), + ] + + def __init__(self): + super(c_nvmlNvLinkInfo_v1_t, self).__init__(version=nvmlNvLinkInfo_v1) + +NVML_NVLINK_FIRMWARE_UCODE_TYPE_MSE = 0x1 +NVML_NVLINK_FIRMWARE_UCODE_TYPE_NETIR = 0x2 +NVML_NVLINK_FIRMWARE_UCODE_TYPE_NETIR_UPHY = 0x3 +NVML_NVLINK_FIRMWARE_UCODE_TYPE_NETIR_CLN = 0x4 +NVML_NVLINK_FIRMWARE_UCODE_TYPE_NETIR_DLN = 0x5 +NVML_NVLINK_FIRMWARE_VERSION_LENGTH = 100 + +class c_nvmlNvlinkFirmwareVersion_t(_PrintableStructure): + _fields_ = [ + ("ucodeType", c_uint8), + ("major", c_uint), + ("minor", c_uint), + ("subMinor", c_uint) + ] + +class c_nvmlNvlinkFirmwareInfo_t(_PrintableStructure): + _fields_ = [ + ("firmwareVersion", c_nvmlNvlinkFirmwareVersion_t * NVML_NVLINK_FIRMWARE_VERSION_LENGTH), + ("numValidEntries", c_uint) + ] + +nvmlNvLinkInfo_v2 = 0x200064c +class c_nvmlNvLinkInfo_v2_t(_PrintableStructure): + _fields_ = [ + ("version", c_uint), + ("isNvleEnabled", c_uint), + ("firmwareInfo", c_nvmlNvlinkFirmwareInfo_t) + ] + + def __init__(self): + super(c_nvmlNvLinkInfo_v2_t, self).__init__(version=nvmlNvLinkInfo_v2) + + +NVML_PRM_DATA_MAX_SIZE = 496 +class c_nvmlPRMTLV_v1_t(_PrintableStructure): + _fields_ = [ + ('dataSize', c_uint32), + ('status', c_uint32), + ('data', c_ubyte * NVML_PRM_DATA_MAX_SIZE), + ] + def __init__(self, size=0): + super(c_nvmlPRMTLV_v1_t, self).__init__(dataSize=size, status=0) + +def nvmlDeviceReadWritePRM_v1(handle, c_info): + fn = _nvmlGetFunctionPointer("nvmlDeviceReadWritePRM_v1") + ret = fn(handle, byref(c_info)) + _nvmlCheckReturn(ret) + + +# Addressing modes +NVML_DEVICE_ADDRESSING_MODE_NONE = 0 +NVML_DEVICE_ADDRESSING_MODE_HMM = 1 +NVML_DEVICE_ADDRESSING_MODE_ATS = 2 + +nvmlDeviceAddressingMode_v1 = 0x1000008 +class c_nvmlDeviceAddressingMode_t(_PrintableStructure): + _fields_ = [ + ('version', c_uint), + ('value', c_uint), + ] + + def __init__(self): + super(c_nvmlDeviceAddressingMode_t, self).__init__(version=nvmlDeviceAddressingMode_v1) + +## Event structures +class struct_c_nvmlEventSet_t(Structure): + pass # opaque handle +c_nvmlEventSet_t = POINTER(struct_c_nvmlEventSet_t) + +nvmlEventTypeSingleBitEccError = 0x0000000000000001 +nvmlEventTypeDoubleBitEccError = 0x0000000000000002 +nvmlEventTypePState = 0x0000000000000004 +nvmlEventTypeXidCriticalError = 0x0000000000000008 +nvmlEventTypeClock = 0x0000000000000010 +nvmlEventTypePowerSourceChange = 0x0000000000000080 +nvmlEventMigConfigChange = 0x0000000000000100 +nvmlEventTypeSingleBitEccErrorStorm = 0x0000000000000200 +nvmlEventTypeDramRetirementEvent = 0x0000000000000400 +nvmlEventTypeDramRetirementFailure = 0x0000000000000800 +nvmlEventTypeNonFatalPoisonError = 0x0000000000001000 +nvmlEventTypeFatalPoisonError = 0x0000000000002000 +nvmlEventTypeGpuUnavailableError = 0x0000000000004000 +nvmlEventTypeGpuRecoveryAction = 0x0000000000008000 +nvmlEventTypeNone = 0x0000000000000000 +nvmlEventTypeAll = ( + nvmlEventTypeNone + | nvmlEventTypeSingleBitEccError + | nvmlEventTypeDoubleBitEccError + | nvmlEventTypePState + | nvmlEventTypeClock + | nvmlEventTypePowerSourceChange + | nvmlEventTypeXidCriticalError + | nvmlEventMigConfigChange + | nvmlEventTypeSingleBitEccErrorStorm + | nvmlEventTypeDramRetirementEvent + | nvmlEventTypeDramRetirementFailure + | nvmlEventTypeNonFatalPoisonError + | nvmlEventTypeFatalPoisonError + | nvmlEventTypeGpuUnavailableError + | nvmlEventTypeGpuRecoveryAction + ) + +## Clock Event Reasons defines +nvmlClocksEventReasonGpuIdle = 0x0000000000000001 +nvmlClocksEventReasonApplicationsClocksSetting = 0x0000000000000002 +nvmlClocksEventReasonUserDefinedClocks = nvmlClocksEventReasonApplicationsClocksSetting # deprecated, use nvmlClocksEventReasonApplicationsClocksSetting +nvmlClocksEventReasonSwPowerCap = 0x0000000000000004 +nvmlClocksEventReasonHwSlowdown = 0x0000000000000008 +nvmlClocksEventReasonSyncBoost = 0x0000000000000010 +nvmlClocksEventReasonSwThermalSlowdown = 0x0000000000000020 +nvmlClocksEventReasonHwThermalSlowdown = 0x0000000000000040 +nvmlClocksEventReasonHwPowerBrakeSlowdown = 0x0000000000000080 +nvmlClocksEventReasonDisplayClockSetting = 0x0000000000000100 +nvmlClocksEventReasonNone = 0x0000000000000000 +nvmlClocksEventReasonAll = ( + nvmlClocksEventReasonNone | + nvmlClocksEventReasonGpuIdle | + nvmlClocksEventReasonApplicationsClocksSetting | + nvmlClocksEventReasonSwPowerCap | + nvmlClocksEventReasonHwSlowdown | + nvmlClocksEventReasonSyncBoost | + nvmlClocksEventReasonSwThermalSlowdown | + nvmlClocksEventReasonHwThermalSlowdown | + nvmlClocksEventReasonHwPowerBrakeSlowdown | + nvmlClocksEventReasonDisplayClockSetting + ) + +## Following have been deprecated +nvmlClocksThrottleReasonGpuIdle = 0x0000000000000001 +nvmlClocksThrottleReasonApplicationsClocksSetting = 0x0000000000000002 +nvmlClocksThrottleReasonUserDefinedClocks = nvmlClocksThrottleReasonApplicationsClocksSetting # deprecated, use nvmlClocksThrottleReasonApplicationsClocksSetting +nvmlClocksThrottleReasonSwPowerCap = 0x0000000000000004 +nvmlClocksThrottleReasonHwSlowdown = 0x0000000000000008 +nvmlClocksThrottleReasonSyncBoost = 0x0000000000000010 +nvmlClocksThrottleReasonSwThermalSlowdown = 0x0000000000000020 +nvmlClocksThrottleReasonHwThermalSlowdown = 0x0000000000000040 +nvmlClocksThrottleReasonHwPowerBrakeSlowdown = 0x0000000000000080 +nvmlClocksThrottleReasonDisplayClockSetting = 0x0000000000000100 +nvmlClocksThrottleReasonNone = 0x0000000000000000 +nvmlClocksThrottleReasonAll = ( + nvmlClocksThrottleReasonNone | + nvmlClocksThrottleReasonGpuIdle | + nvmlClocksThrottleReasonApplicationsClocksSetting | + nvmlClocksThrottleReasonSwPowerCap | + nvmlClocksThrottleReasonHwSlowdown | + nvmlClocksThrottleReasonSyncBoost | + nvmlClocksThrottleReasonSwThermalSlowdown | + nvmlClocksThrottleReasonHwThermalSlowdown | + nvmlClocksThrottleReasonHwPowerBrakeSlowdown | + nvmlClocksThrottleReasonDisplayClockSetting + ) + +class c_nvmlEventData_t(_PrintableStructure): + _fields_ = [ + ('device', c_nvmlDevice_t), + ('eventType', c_ulonglong), + ('eventData', c_ulonglong), + ('gpuInstanceId', c_uint), + ('computeInstanceId', c_uint) + ] + _fmt_ = {'eventType': "0x%08X"} + +class struct_c_nvmlSystemEventSet_t(Structure): + pass # opaque handle +c_nvmlSystemEventSet_t = POINTER(struct_c_nvmlSystemEventSet_t) + +nvmlSystemEventTypeGpuDriverUnbind = 0x0000000000000001 +nvmlSystemEventTypeGpuDriverBind = 0x0000000000000002 + +nvmlSystemEventTypeCount = 2 + +nvmlSystemEventSetCreateRequest_v1 = 0x1000010 +class c_nvmlSystemEventSetCreateRequest_v1_t(_PrintableStructure): + _fields_ = [ + ('version', c_uint), + ('set', c_nvmlSystemEventSet_t) + ] + + def __init__(self): + super(c_nvmlSystemEventSetCreateRequest_v1_t, self).__init__(version=nvmlSystemEventSetCreateRequest_v1) + +nvmlSystemEventSetFreeRequest_v1 = 0x1000010 +class c_nvmlSystemEventSetFreeRequest_v1_t(_PrintableStructure): + _fields_ = [ + ('version', c_uint), + ('set', c_nvmlSystemEventSet_t) + ] + + def __init__(self): + super(c_nvmlSystemEventSetFreeRequest_v1_t, self).__init__(version=nvmlSystemEventSetFreeRequest_v1) + +nvmlSystemRegisterEventRequest_v1 = 0x1000018 +class c_nvmlSystemRegisterEventRequest_v1_t(_PrintableStructure): + _fields_ = [ + ('version', c_uint), + ('eventTypes', c_ulonglong), + ('set', c_nvmlSystemEventSet_t) + ] + + def __init__(self): + super(c_nvmlSystemRegisterEventRequest_v1_t, self).__init__(version=nvmlSystemRegisterEventRequest_v1) + +class c_nvmlSystemEventData_v1_t(_PrintableStructure): + _fields_ = [ + ('eventType', c_ulonglong), + ('gpuId', c_uint) + ] + _fmt_ = {'eventType': "0x%08X"} + +nvmlSystemEventSetWaitRequest_v1 = 0x1000020 +class c_nvmlSystemEventSetWaitRequest_v1_t(_PrintableStructure): + _fields_ = [ + ('version', c_uint), # input/output + ('timeoutms', c_uint), # input + ('set', c_nvmlSystemEventSet_t), # input + ('data', POINTER(c_nvmlSystemEventData_v1_t)), # input/output. Buffer owned by caller + ('dataSize', c_uint), # input + ('numEvent', c_uint) # output. Number of event recorded + ] + + def __init__(self): + super(c_nvmlSystemEventSetWaitRequest_v1_t, self).__init__(version=nvmlSystemEventSetWaitRequest_v1) + +# C APIs for system events +def c_nvmlSystemEventSetCreate(req): + """ + C API binding for c_nvmlSystemEventSetFree + :param req: reference of c_nvmlSystemEventSetCreateRequest_v1_t + """ + fn = _nvmlGetFunctionPointer("nvmlSystemEventSetCreate") + ret = fn(req) + _nvmlCheckReturn(ret) + +def c_nvmlSystemEventSetFree(req): + """ + C API binding for c_nvmlSystemEventSetFree + :param req: reference of c_nvmlSystemEventSetFreeRequest_v1_t + """ + fn = _nvmlGetFunctionPointer("nvmlSystemEventSetFree") + ret = fn(req) + _nvmlCheckReturn(ret) + +def c_nvmlSystemRegisterEvents(req): + """ + C API binding for nvmlSystemRegisterEvents + :param req: reference of c_nvmlSystemRegisterEventRequest_v1_t + """ + fn = _nvmlGetFunctionPointer("nvmlSystemRegisterEvents") + ret = fn(req) + _nvmlCheckReturn(ret) + +def c_nvmlSystemEventSetWait(req): + """ + C API binding for nvmlSystemEventSetWait + :param req: reference of c_nvmlSystemEventSetWaitRequest_v1_t + """ + fn = _nvmlGetFunctionPointer("nvmlSystemEventSetWait") + ret = fn(req) + # Accept TIMEOUT error + if ret == NVML_ERROR_TIMEOUT: + return ret + _nvmlCheckReturn(ret) + +# Pythonic APIs for system events +# No Pythonic API compability support +# When user upgrades to a new NVML Python binding, +# user should upgrade their code against API change. +def nvmlSystemEventSetCreate(version): + """ + Create an NVML system event set. + :param version: The version of the request. + :return: A ctypes pointer to the created c_nvmlSystemEventSet_t structure. + """ + if version != nvmlSystemEventSetCreateRequest_v1: + raise NVMLError(NVML_ERROR_ARGUMENT_VERSION_MISMATCH) + + # Prepare the request structure + req = c_nvmlSystemEventSetCreateRequest_v1_t() + req.version = version + req.set = None # Initialize as NULL + + c_nvmlSystemEventSetCreate(byref(req)) + + # Return the created set handle + if not req.set: + raise NVMLError(NVML_ERROR_MEMORY) + return req.set + +def nvmlSystemEventSetFree(version, setHandle): + """ + Free an NVML system event set. + :param version: The version of the request. + :param setHandle: A ctypes pointer to a c_nvmlSystemEventSet_t structure. + """ + if version != nvmlSystemEventSetFreeRequest_v1: + raise NVMLError(NVML_ERROR_ARGUMENT_VERSION_MISMATCH) + + req = c_nvmlSystemEventSetFreeRequest_v1_t() + req.version = version + req.set = setHandle + + c_nvmlSystemEventSetFree(byref(req)) + + +def nvmlSystemRegisterEvents(version, eventTypes, setHandle): + """ + Register events for an NVML system event set. + :param version: The version of the request. + :param eventTypes: A bitmask of event types to register. + :param setHandle: A ctypes pointer to a c_nvmlSystemEventSet_t structure. + """ + if version != nvmlSystemRegisterEventRequest_v1: + raise NVMLError(NVML_ERROR_ARGUMENT_VERSION_MISMATCH) + + req = c_nvmlSystemRegisterEventRequest_v1_t() + req.version = version + req.eventTypes = eventTypes + req.set = setHandle + + c_nvmlSystemRegisterEvents(byref(req)) + + +def nvmlSystemEventSetWait(version, setHandle, timeoutMs, eventData, eventDataSize): + """ + Wait for events in an NVML system event set. + :param version: The version of the request. + :param setHandle: A ctypes pointer to a c_nvmlSystemEventSet_t structure. + :param timeoutMs: Timeout in milliseconds. + :param eventData: A ctypes array of c_nvmlSystemEventData_v1_t for event data. + :param eventDataSize: Number of c_nvmlSystemEventData_v1_t in evenData array.. + :return: A list of dictionaries containing event data. + """ + if version != nvmlSystemEventSetWaitRequest_v1: + raise NVMLError(NVML_ERROR_ARGUMENT_VERSION_MISMATCH) + + req = c_nvmlSystemEventSetWaitRequest_v1_t() + req.version = version + req.timeoutms = timeoutMs + req.set = setHandle + req.data = eventData + req.dataSize = eventDataSize + + ret = c_nvmlSystemEventSetWait(byref(req)) + if ret == NVML_ERROR_TIMEOUT: + return [] # Timeout, no events occurred. + + # Extract event data from the buffer + events = [] + for i in range(req.numEvent): + events.append({ + 'eventType': eventData[i].eventType, + 'gpuId': eventData[i].gpuId + }) + return events + +class c_nvmlAccountingStats_t(_PrintableStructure): + _fields_ = [ + ('gpuUtilization', c_uint), + ('memoryUtilization', c_uint), + ('maxMemoryUsage', c_ulonglong), + ('time', c_ulonglong), + ('startTime', c_ulonglong), + ('isRunning', c_uint), + ('reserved', c_uint * 5) + ] + +class c_nvmlVgpuVersion_t(Structure): + _fields_ = [("minVersion", c_uint), + ("maxVersion", c_uint) + ] + +class c_nvmlVgpuMetadata_t(_PrintableStructure): + _fields_ = [("version", c_uint), + ("revision", c_uint), + ("guestInfoState", _nvmlVgpuGuestInfoState_t), + ("guestDriverVersion", c_char * NVML_SYSTEM_DRIVER_VERSION_BUFFER_SIZE), + ("hostDriverVersion", c_char * NVML_SYSTEM_DRIVER_VERSION_BUFFER_SIZE), + ("reserved", c_uint * 6), + ("vgpuVirtualizationCaps", c_uint), + ("guestVgpuVersion", c_uint), + ("opaqueDataSize", c_uint), + ("opaqueData", c_char * NVML_VGPU_METADATA_OPAQUE_DATA_SIZE) + ] + +class c_nvmlVgpuPgpuMetadata_t(_PrintableStructure): + _fields_ = [("version", c_uint), + ("revision", c_uint), + ("hostDriverVersion", c_char * NVML_SYSTEM_DRIVER_VERSION_BUFFER_SIZE), + ("pgpuVirtualizationCaps", c_uint), + ("reserved", c_uint * 5), + ("hostSupportedVgpuRange", c_nvmlVgpuVersion_t), + ("opaqueDataSize", c_uint), + ("opaqueData", c_char * NVML_VGPU_PGPU_METADATA_OPAQUE_DATA_SIZE) + ] + +class c_nvmlVgpuPgpuCompatibility_t(Structure): + _fields_ = [("vgpuVmCompatibility", _nvmlVgpuVmCompatibility_t), + ("compatibilityLimitCode", _nvmlVgpuPgpuCompatibilityLimitCode_t) + ] + +## vGPU scheduler policy defines +NVML_VGPU_SCHEDULER_POLICY_UNKNOWN = 0 +NVML_VGPU_SCHEDULER_POLICY_BEST_EFFORT = 1 +NVML_VGPU_SCHEDULER_POLICY_EQUAL_SHARE = 2 +NVML_VGPU_SCHEDULER_POLICY_FIXED_SHARE = 3 + +## Supported vGPU scheduler policy count +NVML_SUPPORTED_VGPU_SCHEDULER_POLICY_COUNT = 3 + +NVML_SCHEDULER_SW_MAX_LOG_ENTRIES = 200 + +NVML_VGPU_SCHEDULER_ARR_DEFAULT = 0 +NVML_VGPU_SCHEDULER_ARR_DISABLE = 1 +NVML_VGPU_SCHEDULER_ARR_ENABLE = 2 + +NVML_VGPU_SCHEDULER_ENGINE_TYPE_GRAPHICS = 1 + +class c_nvmlVgpuSchedDataWithARR_t(_PrintableStructure): + _fields_ = [ + ('avgFactor', c_uint), + ('timeslice', c_uint), + ] + +class c_nvmlVgpuSchedData_t(_PrintableStructure): + _fields_ = [ + ('timeslice', c_uint), + ] + +class c_nvmlVgpuSchedulerParams_t(Union): + _fields_ = [ + ('vgpuSchedDataWithARR', c_nvmlVgpuSchedDataWithARR_t), + ('vgpuSchedData', c_nvmlVgpuSchedData_t), + ] + +class c_nvmlVgpuSchedulerLogEntry_t(_PrintableStructure): + _fields_ = [ + ('timestamp', c_ulonglong), + ('timeRunTotal', c_ulonglong), + ('timeRun', c_ulonglong), + ('swRunlistId', c_uint), + ('targetTimeSlice', c_ulonglong), + ('cumulativePreemptionTime', c_ulonglong), + ] + +class c_nvmlVgpuSchedulerLog_t(_PrintableStructure): + _fields_ = [ + ('engineId', c_uint), + ('schedulerPolicy', c_uint), + ('arrMode', c_uint), + ('schedulerParams', c_nvmlVgpuSchedulerParams_t), + ('entriesCount', c_uint), + ('logEntries', c_nvmlVgpuSchedulerLogEntry_t * NVML_SCHEDULER_SW_MAX_LOG_ENTRIES), + ] + +class c_nvmlVgpuSchedulerGetState_t(_PrintableStructure): + _fields_ = [ + ('schedulerPolicy', c_uint), + ('arrMode', c_uint), + ('schedulerParams', c_nvmlVgpuSchedulerParams_t), + ] + +class c_nvmlVgpuSchedSetDataWithARR_t(_PrintableStructure): + _fields_ = [ + ('avgFactor', c_uint), + ('frequency', c_uint), + ] + +class c_nvmlVgpuSchedSetData_t(_PrintableStructure): + _fields_ = [ + ('timeslice', c_uint), + ] + +class c_nvmlVgpuSchedulerSetParams_t(Union): + _fields_ = [ + ('vgpuSchedDataWithARR', c_nvmlVgpuSchedSetDataWithARR_t), + ('vgpuSchedData', c_nvmlVgpuSchedSetData_t), + ] + +class c_nvmlVgpuSchedulerSetState_t(_PrintableStructure): + _fields_ = [ + ('schedulerPolicy', c_uint), + ('enableARRMode', c_uint), + ('schedulerParams', c_nvmlVgpuSchedulerSetParams_t), + ] + +class c_nvmlVgpuSchedulerCapabilities_t(_PrintableStructure): + _fields_ = [ + ('supportedSchedulers', c_uint * NVML_SUPPORTED_VGPU_SCHEDULER_POLICY_COUNT), + ('maxTimeslice', c_uint), + ('minTimeslice', c_uint), + ('isArrModeSupported', c_uint), + ('maxFrequencyForARR', c_uint), + ('minFrequencyForARR', c_uint), + ('maxAvgFactorForARR', c_uint), + ('minAvgFactorForARR', c_uint), + ] + +class c_nvmlVgpuTypeIdInfo_v1_t(_PrintableStructure): + _fields_ = [ + ('version', c_uint), + ('vgpuCount', c_uint), + ('vgpuTypeIds', POINTER(c_uint)), + ] + +nvmlVgpuTypeIdInfo_v1 = 0x1000010 + +class c_nvmlVgpuTypeMaxInstance_v1_t(_PrintableStructure): + _fields_ = [ + ('version', c_uint), + ('vgpuTypeId', c_uint), + ('maxInstancePerGI', c_uint), + ] + +nvmlVgpuTypeMaxInstance_v1 = 0x100000C + +class c_nvmlActiveVgpuInstanceInfo_v1_t(_PrintableStructure): + _fields_ = [ + ('version', c_uint), + ('vgpuCount', c_uint), + ('vgpuInstances', POINTER(c_uint)), + ] + +nvmlActiveVgpuInstanceInfo_v1 = 0x1000010 + +class c_nvmlVgpuSchedulerState_v1_t(_PrintableStructure): + _fields_ = [ + ('version', c_uint), + ('engineId', c_uint), + ('schedulerPolicy', c_uint), + ('enableARRMode', c_uint), + ('schedulerParams', c_nvmlVgpuSchedulerSetParams_t), + ] + +nvmlVgpuSchedulerState_v1 = 0x1000018 + +class c_nvmlVgpuSchedulerStateInfo_v1_t(_PrintableStructure): + _fields_ = [ + ('version', c_uint), # input + ('engineId', c_uint), # input. One of NVML_ENGINE_TYPE* + ('schedulerPolicy', c_uint), # output + ('arrMode', c_uint), # output + ('schedulerParams', c_nvmlVgpuSchedulerParams_t), # output + ] + +nvmlVgpuSchedulerStateInfo_v1 = 0x1000018 + +class c_nvmlVgpuSchedulerLogInfo_v1_t(_PrintableStructure): + _fields_ = [ + ('version', c_uint), # input + ('engineId', c_uint), # input. One of NVML_ENGINE_TYPE* + ('schedulerPolicy', c_uint), # output + ('arrMode', c_uint), # output + ('schedulerParams', c_nvmlVgpuSchedulerParams_t), # output + ('entriesCount', c_uint), # output + ('logEntries', c_nvmlVgpuSchedulerLogEntry_t * NVML_SCHEDULER_SW_MAX_LOG_ENTRIES), # output + ] + +nvmlVgpuSchedulerLogInfo_v1 = 0x10025A0 + +class c_nvmlVgpuCreatablePlacementInfo_v1_t(_PrintableStructure): + _fields_ = [ + ('version', c_uint), + ('vgpuTypeId', c_uint), + ('count', c_uint), + ('placementIds', POINTER(c_uint)), + ('placementSize', c_uint), + ] + +nvmlVgpuCreatablePlacementInfo_v1 = 0x1000020 + +class c_nvmlFBCStats_t(Structure): + _fields_ = [("sessionsCount", c_uint), + ("averageFPS", c_uint), + ("averageLatency", c_uint) + ] + +class c_nvmlFBCSession_t(_PrintableStructure): + _fields_ = [ + ('sessionId', c_uint), + ('pid', c_uint), + ('vgpuInstance', _nvmlVgpuInstance_t), + ('displayOrdinal', c_uint), + ('sessionType', c_uint), + ('sessionFlags', c_uint), + ('hMaxResolution', c_uint), + ('vMaxResolution', c_uint), + ('hResolution', c_uint), + ('vResolution', c_uint), + ('averageFPS', c_uint), + ('averageLatency', c_uint), + ] + +NVML_DEVICE_MIG_DISABLE = 0x0 +NVML_DEVICE_MIG_ENABLE = 0x1 + +NVML_GPU_INSTANCE_PROFILE_1_SLICE = 0x0 +NVML_GPU_INSTANCE_PROFILE_2_SLICE = 0x1 +NVML_GPU_INSTANCE_PROFILE_3_SLICE = 0x2 +NVML_GPU_INSTANCE_PROFILE_4_SLICE = 0x3 +NVML_GPU_INSTANCE_PROFILE_7_SLICE = 0x4 +NVML_GPU_INSTANCE_PROFILE_8_SLICE = 0x5 +NVML_GPU_INSTANCE_PROFILE_6_SLICE = 0x6 +NVML_GPU_INSTANCE_PROFILE_1_SLICE_REV1 = 0x7 +NVML_GPU_INSTANCE_PROFILE_2_SLICE_REV1 = 0x8 +NVML_GPU_INSTANCE_PROFILE_1_SLICE_REV2 = 0x9 +NVML_GPU_INSTANCE_PROFILE_1_SLICE_GFX = 0xA +NVML_GPU_INSTANCE_PROFILE_2_SLICE_GFX = 0xB +NVML_GPU_INSTANCE_PROFILE_4_SLICE_GFX = 0xC +NVML_GPU_INSTANCE_PROFILE_1_SLICE_NO_ME = 0xD +NVML_GPU_INSTANCE_PROFILE_2_SLICE_NO_ME = 0xE +NVML_GPU_INSTANCE_PROFILE_1_SLICE_ALL_ME = 0xF +NVML_GPU_INSTANCE_PROFILE_2_SLICE_ALL_ME = 0x10 +NVML_GPU_INSTANCE_PROFILE_COUNT = 0x11 + +class c_nvmlGpuInstancePlacement_t(Structure): + _fields_ = [("start", c_uint), + ("size", c_uint) + ] + +class c_nvmlGpuInstanceProfileInfo_t(Structure): + _fields_ = [("id", c_uint), + ("isP2pSupported", c_uint), + ("sliceCount", c_uint), + ("instanceCount", c_uint), + ("multiprocessorCount", c_uint), + ("copyEngineCount", c_uint), + ("decoderCount", c_uint), + ("encoderCount", c_uint), + ("jpegCount", c_uint), + ("ofaCount", c_uint), + ("memorySizeMB", c_ulonglong), + ] + +nvmlGpuInstanceProfileInfo_v2 = 0x02000098 + +class c_nvmlGpuInstanceProfileInfo_v2_t(_PrintableStructure): + _fields_ = [("version", c_uint), + ("id", c_uint), + ("isP2pSupported", c_uint), + ("sliceCount", c_uint), + ("instanceCount", c_uint), + ("multiprocessorCount", c_uint), + ("copyEngineCount", c_uint), + ("decoderCount", c_uint), + ("encoderCount", c_uint), + ("jpegCount", c_uint), + ("ofaCount", c_uint), + ("memorySizeMB", c_ulonglong), + ("name", c_char * NVML_DEVICE_NAME_V2_BUFFER_SIZE) + ] + + def __init__(self): + super(c_nvmlGpuInstanceProfileInfo_v2_t, self).__init__(version=nvmlGpuInstanceProfileInfo_v2) + +class c_nvmlGpuInstanceInfo_t(Structure): + _fields_ = [("device", c_nvmlDevice_t), + ("id", c_uint), + ("profileId", c_uint), + ("placement", c_nvmlGpuInstancePlacement_t) + ] + +class struct_c_nvmlGpuInstance_t(Structure): + pass # opaque handle +c_nvmlGpuInstance_t = POINTER(struct_c_nvmlGpuInstance_t) + +NVML_COMPUTE_INSTANCE_PROFILE_1_SLICE = 0x0 +NVML_COMPUTE_INSTANCE_PROFILE_2_SLICE = 0x1 +NVML_COMPUTE_INSTANCE_PROFILE_3_SLICE = 0x2 +NVML_COMPUTE_INSTANCE_PROFILE_4_SLICE = 0x3 +NVML_COMPUTE_INSTANCE_PROFILE_7_SLICE = 0x4 +NVML_COMPUTE_INSTANCE_PROFILE_8_SLICE = 0x5 +NVML_COMPUTE_INSTANCE_PROFILE_6_SLICE = 0x6 +NVML_COMPUTE_INSTANCE_PROFILE_1_SLICE_REV1 = 0x7 +NVML_COMPUTE_INSTANCE_PROFILE_COUNT = 0x8 + +NVML_COMPUTE_INSTANCE_ENGINE_PROFILE_SHARED = 0x0 +NVML_COMPUTE_INSTANCE_ENGINE_PROFILE_COUNT = 0x1 + +class c_nvmlComputeInstancePlacement_t(Structure): + _fields_ = [("start", c_uint), + ("size", c_uint) + ] + +class c_nvmlComputeInstanceProfileInfo_t(Structure): + _fields_ = [("id", c_uint), + ("sliceCount", c_uint), + ("instanceCount", c_uint), + ("multiprocessorCount", c_uint), + ("sharedCopyEngineCount", c_uint), + ("sharedDecoderCount", c_uint), + ("sharedEncoderCount", c_uint), + ("sharedJpegCount", c_uint), + ("sharedOfaCount", c_uint) + ] + +nvmlComputeInstanceProfileInfo_v2 = 0x02000088 + +class c_nvmlComputeInstanceProfileInfo_v2_t(_PrintableStructure): + _fields_ = [("version", c_uint), + ("id", c_uint), + ("sliceCount", c_uint), + ("instanceCount", c_uint), + ("multiprocessorCount", c_uint), + ("sharedCopyEngineCount", c_uint), + ("sharedDecoderCount", c_uint), + ("sharedEncoderCount", c_uint), + ("sharedJpegCount", c_uint), + ("sharedOfaCount", c_uint), + ("name", c_char * NVML_DEVICE_NAME_V2_BUFFER_SIZE) + ] + + def __init__(self): + super(c_nvmlComputeInstanceProfileInfo_v2_t, self).__init__(version=nvmlComputeInstanceProfileInfo_v2) + +class c_nvmlComputeInstanceInfo_t(Structure): + _fields_ = [("device", c_nvmlDevice_t), + ("gpuInstance", c_nvmlGpuInstance_t), + ("id", c_uint), + ("profileId", c_uint), + ("placement", c_nvmlComputeInstancePlacement_t) + ] + +NVML_MAX_GPU_UTILIZATIONS = 8 +NVML_GPU_UTILIZATION_DOMAIN_GPU = 0 +NVML_GPU_UTILIZATION_DOMAIN_FB = 1 +NVML_GPU_UTILIZATION_DOMAIN_VID = 2 +NVML_GPU_UTILIZATION_DOMAIN_BUS = 3 +class c_nvmlGpuDynamicPstatesUtilization_t(Structure): + _fields_ = [("bIsPresent", c_uint, 1), + ("percentage", c_uint), + ("incThreshold", c_uint), + ("decThreshold", c_uint)] +class c_nvmlGpuDynamicPstatesInfo_t(Structure): + _fields_ = [("flags", c_uint), + ("utilization", c_nvmlGpuDynamicPstatesUtilization_t * NVML_MAX_GPU_UTILIZATIONS)] + +NVML_MAX_THERMAL_SENSORS_PER_GPU = 3 + +NVML_THERMAL_TARGET_NONE = 0 +NVML_THERMAL_TARGET_GPU = 1 +NVML_THERMAL_TARGET_MEMORY = 2 +NVML_THERMAL_TARGET_POWER_SUPPLY = 4 +NVML_THERMAL_TARGET_BOARD = 8 +NVML_THERMAL_TARGET_VCD_BOARD = 9 +NVML_THERMAL_TARGET_VCD_INLET = 10 +NVML_THERMAL_TARGET_VCD_OUTLET = 11 +NVML_THERMAL_TARGET_ALL = 15 +NVML_THERMAL_TARGET_UNKNOWN = -1 + +NVML_THERMAL_CONTROLLER_NONE = 0 +NVML_THERMAL_CONTROLLER_GPU_INTERNAL = 1 +NVML_THERMAL_CONTROLLER_ADM1032 = 2 +NVML_THERMAL_CONTROLLER_ADT7461 = 3 +NVML_THERMAL_CONTROLLER_MAX6649 = 4 +NVML_THERMAL_CONTROLLER_MAX1617 = 5 +NVML_THERMAL_CONTROLLER_LM99 = 6 +NVML_THERMAL_CONTROLLER_LM89 = 7 +NVML_THERMAL_CONTROLLER_LM64 = 8 +NVML_THERMAL_CONTROLLER_G781 = 9 +NVML_THERMAL_CONTROLLER_ADT7473 = 10 +NVML_THERMAL_CONTROLLER_SBMAX6649 = 11 +NVML_THERMAL_CONTROLLER_VBIOSEVT = 12 +NVML_THERMAL_CONTROLLER_OS = 13 +NVML_THERMAL_CONTROLLER_NVSYSCON_CANOAS = 14 +NVML_THERMAL_CONTROLLER_NVSYSCON_E551 = 15 +NVML_THERMAL_CONTROLLER_MAX6649R = 16 +NVML_THERMAL_CONTROLLER_ADT7473S = 17 +NVML_THERMAL_CONTROLLER_UNKNOWN = -1 + +class c_nvmlGpuThermalSensor_t(Structure): + _fields_ = [("controller", c_int), + ("defaultMinTemp", c_int), + ("defaultMaxTemp", c_int), + ("currentTemp", c_int), + ("target", c_int)] +class c_nvmlGpuThermalSettings_t(Structure): + _fields_ = [("count", c_uint), + ("sensor", c_nvmlGpuThermalSensor_t * NVML_MAX_THERMAL_SENSORS_PER_GPU)] + +_nvmlCoolerControl_t = c_uint +NVML_THERMAL_COOLER_SIGNAL_NONE = 0 +NVML_THERMAL_COOLER_SIGNAL_TOGGLE = 1 +NVML_THERMAL_COOLER_SIGNAL_VARIABLE = 2 +NVML_THERMAL_COOLER_SIGNAL_COUNT = 3 + +_nvmlCoolerTarget_t = c_uint +NVML_THERMAL_COOLER_TARGET_NONE = (1 << 0) +NVML_THERMAL_COOLER_TARGET_GPU = (1 << 1) +NVML_THERMAL_COOLER_TARGET_MEMORY = (1 << 2) +NVML_THERMAL_COOLER_TARGET_POWER_SUPPLY = (1 << 3) +NVML_THERMAL_COOLER_TARGET_GPU_RELATED = (NVML_THERMAL_COOLER_TARGET_GPU | NVML_THERMAL_COOLER_TARGET_MEMORY | NVML_THERMAL_COOLER_TARGET_POWER_SUPPLY) + +class c_nvmlCoolerInfo_t(_PrintableStructure): + _fields_ = [("version", c_uint), + ("index", c_uint), + ("coolerControlType", _nvmlCoolerControl_t), + ("coolerTarget", _nvmlCoolerTarget_t) + ] + +nvmlCoolerInfo_v1 = 0x1000010 + +def nvmlDeviceGetCoolerInfo(handle): + c_coolerInfo = c_nvmlCoolerInfo_t() + c_coolerInfo.version = nvmlCoolerInfo_v1 + c_coolerInfo.index = 0 + fn = _nvmlGetFunctionPointer("nvmlDeviceGetCoolerInfo") + ret = fn(handle, byref(c_coolerInfo)) + _nvmlCheckReturn(ret) + return [c_coolerInfo.coolerControlType, c_coolerInfo.coolerTarget] + +class struct_c_nvmlComputeInstance_t(Structure): + pass # opaque handle +c_nvmlComputeInstance_t = POINTER(struct_c_nvmlComputeInstance_t) + +class c_nvmlDeviceAttributes(Structure): + _fields_ = [("multiprocessorCount", c_uint), + ("sharedCopyEngineCount", c_uint), + ("sharedDecoderCount", c_uint), + ("sharedEncoderCount", c_uint), + ("sharedJpegCount", c_uint), + ("sharedOfaCount", c_uint), + ("gpuInstanceSliceCount", c_uint), + ("computeInstanceSliceCount", c_uint), + ("memorySizeMB", c_ulonglong), + ] + +class c_nvmlRowRemapperHistogramValues(Structure): + _fields_ = [("max", c_uint), + ("high", c_uint), + ("partial", c_uint), + ("low", c_uint), + ("none", c_uint) + ] + +NVML_GPU_CERT_CHAIN_SIZE = 0x1000 +NVML_GPU_ATTESTATION_CERT_CHAIN_SIZE = 0x1400 +NVML_CC_GPU_CEC_NONCE_SIZE = 0x20 +NVML_CC_GPU_ATTESTATION_REPORT_SIZE = 0x2000 +NVML_CC_GPU_CEC_ATTESTATION_REPORT_SIZE = 0x1000 +NVML_CC_CEC_ATTESTATION_REPORT_NOT_PRESENT = 0 +NVML_CC_CEC_ATTESTATION_REPORT_PRESENT = 1 + +class c_nvmlConfComputeSystemState_t(Structure): + _fields_ = [('environment', c_uint), + ('ccFeature', c_uint), + ('devToolsMode', c_uint), + ] + +nvmlSystemConfComputeSettings_v1 = 0x1000014 + +class c_nvmlSystemConfComputeSettings_v1_t(Structure): + _fields_ = [('version', c_uint), + ('environment', c_uint), + ('ccFeature', c_uint), + ('devToolsMode', c_uint), + ('multiGpuMode', c_uint), + ] + def __init__(self): + super(c_nvmlSystemConfComputeSettings_v1_t, self).__init__(version=nvmlSystemConfComputeSettings_v1) + +class c_nvmlConfComputeSystemCaps_t(Structure): + _fields_ = [('cpuCaps', c_uint), + ('gpusCaps', c_uint), + ] + +class c_nvmlConfComputeMemSizeInfo_t(Structure): + _fields_ = [('protectedMemSizeKib', c_ulonglong), + ('unprotectedMemSizeKib', c_ulonglong), + ] + +class c_nvmlConfComputeGpuCertificate_t(Structure): + _fields_ = [('certChainSize', c_uint), + ('attestationCertChainSize', c_uint), + ('certChain', c_uint8 * NVML_GPU_CERT_CHAIN_SIZE), + ('attestationCertChain', c_uint8 * NVML_GPU_ATTESTATION_CERT_CHAIN_SIZE), + ] + +class c_nvmlConfComputeGpuAttestationReport_t(Structure): + _fields_ = [('isCecAttestationReportPresent', c_uint), # output + ('attestationReportSize', c_uint), # output + ('cecAttestationReportSize', c_uint), # output + ('nonce', c_uint8 * NVML_CC_GPU_CEC_NONCE_SIZE), # input: spdm supports 32 bytes on nonce + ('attestationReport', c_uint8 * NVML_CC_GPU_ATTESTATION_REPORT_SIZE), # output + ('cecAttestationReport', c_uint8 * NVML_CC_GPU_CEC_ATTESTATION_REPORT_SIZE), # output + ] + +class c_nvmlConfComputeSetKeyRotationThresholdInfo_t(Structure): + _fields_ = [('version', c_uint), + ('maxAttackerAdvantage', c_ulong), + ] +ConfComputeSetKeyRotationThresholdInfo_v1 = 0x1000010 + +class c_nvmlConfComputeGetKeyRotationThresholdInfo_t(Structure): + _fields_ = [('version', c_uint), + ('attackerAdvantage', c_ulong), + ] +ConfComputeGetKeyRotationThresholdInfo_v1 = 0x1000010 + +## string/bytes conversion for ease of use +def convertStrBytes(func): + ''' + In python 3, strings are unicode instead of bytes, and need to be converted for ctypes + Args from caller: (1, 'string', <__main__.c_nvmlDevice_t at 0xFFFFFFFF>) + Args passed to function: (1, b'string', <__main__.c_nvmlDevice_t at 0xFFFFFFFF)> + ---- + Returned from function: b'returned string' + Returned to caller: 'returned string' + ''' + @wraps(func) + def wrapper(*args, **kwargs): + # encoding a str returns bytes in python 2 and 3 + args = [arg.encode() if isinstance(arg, str) else arg for arg in args] + res = func(*args, **kwargs) + # In python 2, str and bytes are the same + # In python 3, str is unicode and should be decoded. + # Ctypes handles most conversions, this only effects c_char and char arrays. + if isinstance(res, bytes): + if isinstance(res, str): + return res + return res.decode() + return res + + if sys.version_info >= (3,): + return wrapper + return func + +def throwOnVersionMismatch(func): + @wraps(func) + def wrapper(*args, **kwargs): + try: + return func(*args, **kwargs) + except NVMLError_FunctionNotFound: + raise NVMLLibraryMismatchError("Unversioned function called and the " + "pyNVML version does not match the NVML lib version. " + "Either use matching pyNVML and NVML lib versions or " + "use a versioned function such as " + func.__name__ + "_v2") + return wrapper + +## C function wrappers ## +def nvmlInitWithFlags(flags): + _LoadNvmlLibrary() + + # + # Initialize the library + # + fn = _nvmlGetFunctionPointer("nvmlInitWithFlags") + ret = fn(flags) + _nvmlCheckReturn(ret) + + # Atomically update refcount + global _nvmlLib_refcount + libLoadLock.acquire() + _nvmlLib_refcount += 1 + libLoadLock.release() + return None + +def nvmlInit(): + nvmlInitWithFlags(0) + return None + +def _LoadNvmlLibrary(): + ''' + Load the library if it isn't loaded already + ''' + global nvmlLib + + if (nvmlLib == None): + # lock to ensure only one caller loads the library + libLoadLock.acquire() + + try: + # ensure the library still isn't loaded + if (nvmlLib == None): + try: + if (sys.platform[:3] == "win"): + # cdecl calling convention + try: + # Check for nvml.dll in System32 first for DCH drivers + nvmlLib = CDLL(os.path.join(os.getenv("WINDIR", "C:/Windows"), "System32/nvml.dll")) + except OSError as ose: + # If nvml.dll is not found in System32, it should be in ProgramFiles + # load nvml.dll from %ProgramFiles%/NVIDIA Corporation/NVSMI/nvml.dll + nvmlLib = CDLL(os.path.join(os.getenv("ProgramFiles", "C:/Program Files"), "NVIDIA Corporation/NVSMI/nvml.dll")) + else: + # assume linux + nvmlLib = CDLL("libnvidia-ml.so.1") + except OSError as ose: + _nvmlCheckReturn(NVML_ERROR_LIBRARY_NOT_FOUND) + if (nvmlLib == None): + _nvmlCheckReturn(NVML_ERROR_LIBRARY_NOT_FOUND) + finally: + # lock is always freed + libLoadLock.release() + +def nvmlShutdown(): + # + # Leave the library loaded, but shutdown the interface + # + fn = _nvmlGetFunctionPointer("nvmlShutdown") + ret = fn() + _nvmlCheckReturn(ret) + + # Atomically update refcount + global _nvmlLib_refcount + libLoadLock.acquire() + if (0 < _nvmlLib_refcount): + _nvmlLib_refcount -= 1 + libLoadLock.release() + return None + +# Added in 2.285 +@convertStrBytes +def nvmlErrorString(result): + fn = _nvmlGetFunctionPointer("nvmlErrorString") + fn.restype = c_char_p # otherwise return is an int + ret = fn(result) + return ret + +# Added in 2.285 +@convertStrBytes +def nvmlSystemGetNVMLVersion(): + c_version = create_string_buffer(NVML_SYSTEM_NVML_VERSION_BUFFER_SIZE) + fn = _nvmlGetFunctionPointer("nvmlSystemGetNVMLVersion") + ret = fn(c_version, c_uint(NVML_SYSTEM_NVML_VERSION_BUFFER_SIZE)) + _nvmlCheckReturn(ret) + return c_version.value + +def nvmlSystemGetCudaDriverVersion(): + c_cuda_version = c_int() + fn = _nvmlGetFunctionPointer("nvmlSystemGetCudaDriverVersion") + ret = fn(byref(c_cuda_version)) + _nvmlCheckReturn(ret) + return c_cuda_version.value + +def nvmlSystemGetCudaDriverVersion_v2(): + c_cuda_version = c_int() + fn = _nvmlGetFunctionPointer("nvmlSystemGetCudaDriverVersion_v2") + ret = fn(byref(c_cuda_version)) + _nvmlCheckReturn(ret) + return c_cuda_version.value + +# Added in 2.285 +@convertStrBytes +def nvmlSystemGetProcessName(pid): + c_name = create_string_buffer(1024) + fn = _nvmlGetFunctionPointer("nvmlSystemGetProcessName") + ret = fn(c_uint(pid), c_name, c_uint(1024)) + _nvmlCheckReturn(ret) + return c_name.value + +@convertStrBytes +def nvmlSystemGetDriverVersion(): + c_version = create_string_buffer(NVML_SYSTEM_DRIVER_VERSION_BUFFER_SIZE) + fn = _nvmlGetFunctionPointer("nvmlSystemGetDriverVersion") + ret = fn(c_version, c_uint(NVML_SYSTEM_DRIVER_VERSION_BUFFER_SIZE)) + _nvmlCheckReturn(ret) + return c_version.value + +# Added in 2.285 +def nvmlSystemGetHicVersion(): + c_count = c_uint(0) + hics = None + fn = _nvmlGetFunctionPointer("nvmlSystemGetHicVersion") + + # get the count + ret = fn(byref(c_count), None) + + # this should only fail with insufficient size + if ((ret != NVML_SUCCESS) and + (ret != NVML_ERROR_INSUFFICIENT_SIZE)): + raise NVMLError(ret) + + # If there are no hics + if (c_count.value == 0): + return [] + + hic_array = c_nvmlHwbcEntry_t * c_count.value + hics = hic_array() + ret = fn(byref(c_count), hics) + _nvmlCheckReturn(ret) + return hics + +def nvmlSystemGetDriverBranch(): + c_branchInfo = c_nvmlSystemDriverBranchInfo_v1_t(0) + c_branchInfo.version = SystemDriverBranchInfo_v1 + fn = _nvmlGetFunctionPointer("nvmlSystemGetDriverBranch") + ret = fn(byref(c_branchInfo), c_uint(NVML_SYSTEM_DRIVER_VERSION_BUFFER_SIZE)) + _nvmlCheckReturn(ret) + return c_branchInfo + +## Unit get functions +def nvmlUnitGetCount(): + c_count = c_uint() + fn = _nvmlGetFunctionPointer("nvmlUnitGetCount") + ret = fn(byref(c_count)) + _nvmlCheckReturn(ret) + return c_count.value + +def nvmlUnitGetHandleByIndex(index): + c_index = c_uint(index) + unit = c_nvmlUnit_t() + fn = _nvmlGetFunctionPointer("nvmlUnitGetHandleByIndex") + ret = fn(c_index, byref(unit)) + _nvmlCheckReturn(ret) + return unit + +def nvmlUnitGetUnitInfo(unit): + c_info = c_nvmlUnitInfo_t() + fn = _nvmlGetFunctionPointer("nvmlUnitGetUnitInfo") + ret = fn(unit, byref(c_info)) + _nvmlCheckReturn(ret) + return c_info + +def nvmlUnitGetLedState(unit): + c_state = c_nvmlLedState_t() + fn = _nvmlGetFunctionPointer("nvmlUnitGetLedState") + ret = fn(unit, byref(c_state)) + _nvmlCheckReturn(ret) + return c_state + +def nvmlUnitGetPsuInfo(unit): + c_info = c_nvmlPSUInfo_t() + fn = _nvmlGetFunctionPointer("nvmlUnitGetPsuInfo") + ret = fn(unit, byref(c_info)) + _nvmlCheckReturn(ret) + return c_info + +def nvmlUnitGetTemperature(unit, type): + c_temp = c_uint() + fn = _nvmlGetFunctionPointer("nvmlUnitGetTemperature") + ret = fn(unit, c_uint(type), byref(c_temp)) + _nvmlCheckReturn(ret) + return c_temp.value + +def nvmlUnitGetFanSpeedInfo(unit): + c_speeds = c_nvmlUnitFanSpeeds_t() + fn = _nvmlGetFunctionPointer("nvmlUnitGetFanSpeedInfo") + ret = fn(unit, byref(c_speeds)) + _nvmlCheckReturn(ret) + return c_speeds + +# added to API +def nvmlUnitGetDeviceCount(unit): + c_count = c_uint(0) + # query the unit to determine device count + fn = _nvmlGetFunctionPointer("nvmlUnitGetDevices") + ret = fn(unit, byref(c_count), None) + if (ret == NVML_ERROR_INSUFFICIENT_SIZE): + ret = NVML_SUCCESS + _nvmlCheckReturn(ret) + return c_count.value + +def nvmlUnitGetDevices(unit): + c_count = c_uint(nvmlUnitGetDeviceCount(unit)) + device_array = c_nvmlDevice_t * c_count.value + c_devices = device_array() + fn = _nvmlGetFunctionPointer("nvmlUnitGetDevices") + ret = fn(unit, byref(c_count), c_devices) + _nvmlCheckReturn(ret) + return c_devices + +## Device get functions +def nvmlDeviceGetCount(): + c_count = c_uint() + fn = _nvmlGetFunctionPointer("nvmlDeviceGetCount_v2") + ret = fn(byref(c_count)) + _nvmlCheckReturn(ret) + return c_count.value + +def nvmlDeviceGetHandleByIndex(index): + c_index = c_uint(index) + device = c_nvmlDevice_t() + fn = _nvmlGetFunctionPointer("nvmlDeviceGetHandleByIndex_v2") + ret = fn(c_index, byref(device)) + _nvmlCheckReturn(ret) + return device + +# Deprecated +@convertStrBytes +def nvmlDeviceGetHandleBySerial(serial): + c_serial = c_char_p(serial) + device = c_nvmlDevice_t() + fn = _nvmlGetFunctionPointer("nvmlDeviceGetHandleBySerial") + ret = fn(c_serial, byref(device)) + _nvmlCheckReturn(ret) + return device + +@convertStrBytes +def nvmlDeviceGetHandleByUUID(uuid): + c_uuid = c_char_p(uuid) + device = c_nvmlDevice_t() + fn = _nvmlGetFunctionPointer("nvmlDeviceGetHandleByUUID") + ret = fn(c_uuid, byref(device)) + _nvmlCheckReturn(ret) + return device + +@convertStrBytes +def nvmlDeviceGetHandleByUUIDV(uuid, type): + c_uuid = c_nvmlUUID_t() + c_uuid.type = type + if type == NVML_UUID_TYPE_ASCII: + c_uuid.value.str = uuid + elif type == NVML_UUID_TYPE_BINARY: + memmove(c_uuid.value.bytes, uuid, NVML_DEVICE_UUID_BINARY_LEN) + device = c_nvmlDevice_t() + fn = _nvmlGetFunctionPointer("nvmlDeviceGetHandleByUUIDV") + ret = fn(byref(c_uuid), byref(device)) + _nvmlCheckReturn(ret) + return device + +@convertStrBytes +def nvmlDeviceGetHandleByPciBusId(pciBusId): + c_busId = c_char_p(pciBusId) + device = c_nvmlDevice_t() + fn = _nvmlGetFunctionPointer("nvmlDeviceGetHandleByPciBusId_v2") + ret = fn(c_busId, byref(device)) + _nvmlCheckReturn(ret) + return device + +@convertStrBytes +def nvmlDeviceGetName(handle): + c_name = create_string_buffer(NVML_DEVICE_NAME_V2_BUFFER_SIZE) + fn = _nvmlGetFunctionPointer("nvmlDeviceGetName") + ret = fn(handle, c_name, c_uint(NVML_DEVICE_NAME_V2_BUFFER_SIZE)) + _nvmlCheckReturn(ret) + return c_name.value + +class c_nvmlDevicePerfModes_v1_t(_PrintableStructure): + _fields_ = [ + ('version', c_uint), + ('str', c_char * NVML_PERF_MODES_BUFFER_SIZE), + ] + +nvmlDevicePerfModes_v1 = 0x1000804 + +@convertStrBytes +def nvmlDeviceGetPerformanceModes(handle): + perfModes = c_nvmlDevicePerfModes_v1_t() + perfModes.version = nvmlDevicePerfModes_v1 + fn = _nvmlGetFunctionPointer("nvmlDeviceGetPerformanceModes") + ret = fn(handle, byref(perfModes)) + _nvmlCheckReturn(ret) + return perfModes.str + +class c_nvmlDeviceCurrentClockFreqs_v1_t(_PrintableStructure): + _fields_ = [ + ('version', c_uint), + ('str', c_char * NVML_PERF_MODES_BUFFER_SIZE), + ] + +nvmlDeviceCurrentClockFreqs_v1 = 0x1000804 + +@convertStrBytes +def nvmlDeviceGetCurrentClockFreqs(handle): + currentClockFreqs = c_nvmlDeviceCurrentClockFreqs_v1_t() + currentClockFreqs.version = nvmlDeviceCurrentClockFreqs_v1 + fn = _nvmlGetFunctionPointer("nvmlDeviceGetCurrentClockFreqs") + ret = fn(handle, byref(currentClockFreqs)) + _nvmlCheckReturn(ret) + return currentClockFreqs.str + +def nvmlDeviceGetBoardId(handle): + c_id = c_uint() + fn = _nvmlGetFunctionPointer("nvmlDeviceGetBoardId") + ret = fn(handle, byref(c_id)) + _nvmlCheckReturn(ret) + return c_id.value + +def nvmlDeviceGetMultiGpuBoard(handle): + c_multiGpu = c_uint() + fn = _nvmlGetFunctionPointer("nvmlDeviceGetMultiGpuBoard") + ret = fn(handle, byref(c_multiGpu)) + _nvmlCheckReturn(ret) + return c_multiGpu.value + +def nvmlDeviceGetBrand(handle): + c_type = _nvmlBrandType_t() + fn = _nvmlGetFunctionPointer("nvmlDeviceGetBrand") + ret = fn(handle, byref(c_type)) + _nvmlCheckReturn(ret) + return c_type.value + +def nvmlDeviceGetC2cModeInfoV1(handle): + c_info = c_nvmlC2cModeInfo_v1_t() + fn = _nvmlGetFunctionPointer("nvmlDeviceGetC2cModeInfoV") + ret = fn(handle, byref(c_info)) + _nvmlCheckReturn(ret) + return c_info + +def nvmlDeviceGetC2cModeInfoV(handle): + return nvmlDeviceGetC2cModeInfoV1(handle) + +@convertStrBytes +def nvmlDeviceGetBoardPartNumber(handle): + c_part_number = create_string_buffer(NVML_DEVICE_PART_NUMBER_BUFFER_SIZE) + fn = _nvmlGetFunctionPointer("nvmlDeviceGetBoardPartNumber") + ret = fn(handle, c_part_number, c_uint(NVML_DEVICE_PART_NUMBER_BUFFER_SIZE)) + _nvmlCheckReturn(ret) + return c_part_number.value + +@convertStrBytes +def nvmlDeviceGetSerial(handle): + c_serial = create_string_buffer(NVML_DEVICE_SERIAL_BUFFER_SIZE) + fn = _nvmlGetFunctionPointer("nvmlDeviceGetSerial") + ret = fn(handle, c_serial, c_uint(NVML_DEVICE_SERIAL_BUFFER_SIZE)) + _nvmlCheckReturn(ret) + return c_serial.value + +def nvmlDeviceGetModuleId(handle, moduleId=c_uint()): + isReference = type(moduleId) is not c_uint + moduleIdRef = moduleId if isReference else byref(moduleId) + fn = _nvmlGetFunctionPointer("nvmlDeviceGetModuleId") + ret = fn(handle, moduleIdRef) + if isReference: + return ret + else: + _nvmlCheckReturn(ret) + return moduleId.value + +def nvmlDeviceGetMemoryAffinity(handle, nodeSetSize, scope): + affinity_array = c_ulonglong * nodeSetSize + c_affinity = affinity_array() + fn = _nvmlGetFunctionPointer("nvmlDeviceGetMemoryAffinity") + ret = fn(handle, nodeSetSize, byref(c_affinity), _nvmlAffinityScope_t(scope)) + _nvmlCheckReturn(ret) + return c_affinity + +def nvmlDeviceGetCpuAffinityWithinScope(handle, cpuSetSize, scope): + affinity_array = c_ulonglong * cpuSetSize + c_affinity = affinity_array() + fn = _nvmlGetFunctionPointer("nvmlDeviceGetCpuAffinityWithinScope") + ret = fn(handle, cpuSetSize, byref(c_affinity), _nvmlAffinityScope_t(scope)) + _nvmlCheckReturn(ret) + return c_affinity + +def nvmlDeviceGetCpuAffinity(handle, cpuSetSize): + affinity_array = c_ulonglong * cpuSetSize + c_affinity = affinity_array() + fn = _nvmlGetFunctionPointer("nvmlDeviceGetCpuAffinity") + ret = fn(handle, cpuSetSize, byref(c_affinity)) + _nvmlCheckReturn(ret) + return c_affinity + +def nvmlDeviceSetCpuAffinity(handle): + fn = _nvmlGetFunctionPointer("nvmlDeviceSetCpuAffinity") + ret = fn(handle) + _nvmlCheckReturn(ret) + return None + +def nvmlDeviceClearCpuAffinity(handle): + fn = _nvmlGetFunctionPointer("nvmlDeviceClearCpuAffinity") + ret = fn(handle) + _nvmlCheckReturn(ret) + return None + +def nvmlDeviceGetNumaNodeId(handle): + fn = _nvmlGetFunctionPointer("nvmlDeviceGetNumaNodeId") + node = c_int() + ret = fn(handle, byref(node)) + _nvmlCheckReturn(ret) + return node.value + +def nvmlDeviceGetAddressingMode(device): + c_mode = c_nvmlDeviceAddressingMode_t() + fn = _nvmlGetFunctionPointer("nvmlDeviceGetAddressingMode") + ret = fn(device, byref(c_mode)) + _nvmlCheckReturn(ret) + return c_mode.value + +def nvmlDeviceGetMinorNumber(handle): + c_minor_number = c_uint() + fn = _nvmlGetFunctionPointer("nvmlDeviceGetMinorNumber") + ret = fn(handle, byref(c_minor_number)) + _nvmlCheckReturn(ret) + return c_minor_number.value + +@convertStrBytes +def nvmlDeviceGetUUID(handle): + c_uuid = create_string_buffer(NVML_DEVICE_UUID_V2_BUFFER_SIZE) + fn = _nvmlGetFunctionPointer("nvmlDeviceGetUUID") + ret = fn(handle, c_uuid, c_uint(NVML_DEVICE_UUID_V2_BUFFER_SIZE)) + _nvmlCheckReturn(ret) + return c_uuid.value + +@convertStrBytes +def nvmlDeviceGetInforomVersion(handle, infoRomObject): + c_version = create_string_buffer(NVML_DEVICE_INFOROM_VERSION_BUFFER_SIZE) + fn = _nvmlGetFunctionPointer("nvmlDeviceGetInforomVersion") + ret = fn(handle, _nvmlInforomObject_t(infoRomObject), + c_version, c_uint(NVML_DEVICE_INFOROM_VERSION_BUFFER_SIZE)) + _nvmlCheckReturn(ret) + return c_version.value + +# Added in 4.304 +@convertStrBytes +def nvmlDeviceGetInforomImageVersion(handle): + c_version = create_string_buffer(NVML_DEVICE_INFOROM_VERSION_BUFFER_SIZE) + fn = _nvmlGetFunctionPointer("nvmlDeviceGetInforomImageVersion") + ret = fn(handle, c_version, c_uint(NVML_DEVICE_INFOROM_VERSION_BUFFER_SIZE)) + _nvmlCheckReturn(ret) + return c_version.value + +# Added in 4.304 +def nvmlDeviceGetInforomConfigurationChecksum(handle): + c_checksum = c_uint() + fn = _nvmlGetFunctionPointer("nvmlDeviceGetInforomConfigurationChecksum") + ret = fn(handle, byref(c_checksum)) + _nvmlCheckReturn(ret) + return c_checksum.value + +# Added in 4.304 +def nvmlDeviceValidateInforom(handle): + fn = _nvmlGetFunctionPointer("nvmlDeviceValidateInforom") + ret = fn(handle) + _nvmlCheckReturn(ret) + return None + +def nvmlDeviceGetLastBBXFlushTime(handle): + c_timestamp = c_ulonglong() + c_durationUs = c_ulong() + fn = _nvmlGetFunctionPointer("nvmlDeviceGetLastBBXFlushTime") + ret = fn(handle, byref(c_timestamp), byref(c_durationUs)) + _nvmlCheckReturn(ret) + return [c_timestamp.value, c_durationUs.value] + +def nvmlDeviceGetDisplayMode(handle): + c_mode = _nvmlEnableState_t() + fn = _nvmlGetFunctionPointer("nvmlDeviceGetDisplayMode") + ret = fn(handle, byref(c_mode)) + _nvmlCheckReturn(ret) + return c_mode.value + +def nvmlDeviceGetDisplayActive(handle): + c_mode = _nvmlEnableState_t() + fn = _nvmlGetFunctionPointer("nvmlDeviceGetDisplayActive") + ret = fn(handle, byref(c_mode)) + _nvmlCheckReturn(ret) + return c_mode.value + + +def nvmlDeviceGetPersistenceMode(handle): + c_state = _nvmlEnableState_t() + fn = _nvmlGetFunctionPointer("nvmlDeviceGetPersistenceMode") + ret = fn(handle, byref(c_state)) + _nvmlCheckReturn(ret) + return c_state.value + +def nvmlDeviceGetPciInfoExt(handle, c_info): + fn = _nvmlGetFunctionPointer("nvmlDeviceGetPciInfoExt") + ret = fn(handle, c_info) + _nvmlCheckReturn(ret) + return None + +def nvmlDeviceGetPciInfo_v3(handle): + c_info = nvmlPciInfo_t() + fn = _nvmlGetFunctionPointer("nvmlDeviceGetPciInfo_v3") + ret = fn(handle, byref(c_info)) + _nvmlCheckReturn(ret) + return c_info + +def nvmlDeviceGetPciInfo(handle): + return nvmlDeviceGetPciInfo_v3(handle) + +def nvmlDeviceGetClockInfo(handle, type): + c_clock = c_uint() + fn = _nvmlGetFunctionPointer("nvmlDeviceGetClockInfo") + ret = fn(handle, _nvmlClockType_t(type), byref(c_clock)) + _nvmlCheckReturn(ret) + return c_clock.value + +# Added in 2.285 +def nvmlDeviceGetMaxClockInfo(handle, type): + c_clock = c_uint() + fn = _nvmlGetFunctionPointer("nvmlDeviceGetMaxClockInfo") + ret = fn(handle, _nvmlClockType_t(type), byref(c_clock)) + _nvmlCheckReturn(ret) + return c_clock.value + +# Added in 4.304 +# Deprecated +def nvmlDeviceGetApplicationsClock(handle, type): + c_clock = c_uint() + fn = _nvmlGetFunctionPointer("nvmlDeviceGetApplicationsClock") + ret = fn(handle, _nvmlClockType_t(type), byref(c_clock)) + _nvmlCheckReturn(ret) + return c_clock.value + +def nvmlDeviceGetMaxCustomerBoostClock(handle, type): + c_clock = c_uint() + fn = _nvmlGetFunctionPointer("nvmlDeviceGetMaxCustomerBoostClock") + ret = fn(handle, _nvmlClockType_t(type), byref(c_clock)) + _nvmlCheckReturn(ret) + return c_clock.value + +def nvmlDeviceGetClock(handle, type, id): + c_clock = c_uint() + fn = _nvmlGetFunctionPointer("nvmlDeviceGetClock") + ret = fn(handle, _nvmlClockType_t(type), _nvmlClockId_t(id), byref(c_clock)) + _nvmlCheckReturn(ret) + return c_clock.value + +# Added in 5.319 +# Deprecated +def nvmlDeviceGetDefaultApplicationsClock(handle, type): + c_clock = c_uint() + fn = _nvmlGetFunctionPointer("nvmlDeviceGetDefaultApplicationsClock") + ret = fn(handle, _nvmlClockType_t(type), byref(c_clock)) + _nvmlCheckReturn(ret) + return c_clock.value + +# Added in 4.304 +def nvmlDeviceGetSupportedMemoryClocks(handle): + # first call to get the size + c_count = c_uint(0) + fn = _nvmlGetFunctionPointer("nvmlDeviceGetSupportedMemoryClocks") + ret = fn(handle, byref(c_count), None) + + if (ret == NVML_SUCCESS): + # special case, no clocks + return [] + elif (ret == NVML_ERROR_INSUFFICIENT_SIZE): + # typical case + clocks_array = c_uint * c_count.value + c_clocks = clocks_array() + + # make the call again + ret = fn(handle, byref(c_count), c_clocks) + _nvmlCheckReturn(ret) + + procs = [] + for i in range(c_count.value): + procs.append(c_clocks[i]) + + return procs + else: + # error case + raise NVMLError(ret) + +# Added in 4.304 +def nvmlDeviceGetSupportedGraphicsClocks(handle, memoryClockMHz): + # first call to get the size + c_count = c_uint(0) + fn = _nvmlGetFunctionPointer("nvmlDeviceGetSupportedGraphicsClocks") + ret = fn(handle, c_uint(memoryClockMHz), byref(c_count), None) + + if (ret == NVML_SUCCESS): + # special case, no clocks + return [] + elif (ret == NVML_ERROR_INSUFFICIENT_SIZE): + # typical case + clocks_array = c_uint * c_count.value + c_clocks = clocks_array() + + # make the call again + ret = fn(handle, c_uint(memoryClockMHz), byref(c_count), c_clocks) + _nvmlCheckReturn(ret) + + procs = [] + for i in range(c_count.value): + procs.append(c_clocks[i]) + + return procs + else: + # error case + raise NVMLError(ret) + +def nvmlDeviceGetFanSpeed(handle): + c_speed = c_uint() + fn = _nvmlGetFunctionPointer("nvmlDeviceGetFanSpeed") + ret = fn(handle, byref(c_speed)) + _nvmlCheckReturn(ret) + return c_speed.value + +def nvmlDeviceGetFanSpeed_v2(handle, fan): + c_speed = c_uint() + fn = _nvmlGetFunctionPointer("nvmlDeviceGetFanSpeed_v2") + ret = fn(handle, fan, byref(c_speed)) + _nvmlCheckReturn(ret) + return c_speed.value + +class c_nvmlFanSpeedInfo_t(_PrintableStructure): + _fields_ = [ + ('version', c_uint), + ('fan', c_uint), + ('speed', c_uint), + ] + +nvmlFanSpeedInfo_v1 = 0x100000C + +def nvmlDeviceGetFanSpeedRPM(handle): + c_fanSpeed = c_nvmlFanSpeedInfo_t() + c_fanSpeed.fan = 0 + c_fanSpeed.version = nvmlFanSpeedInfo_v1 + fn = _nvmlGetFunctionPointer("nvmlDeviceGetFanSpeedRPM") + ret = fn(handle, byref(c_fanSpeed)) + _nvmlCheckReturn(ret) + return c_fanSpeed.speed + +def nvmlDeviceGetTargetFanSpeed(handle, fan): + c_speed = c_uint() + fn = _nvmlGetFunctionPointer("nvmlDeviceGetTargetFanSpeed") + ret = fn(handle, fan, byref(c_speed)) + _nvmlCheckReturn(ret) + return c_speed.value + +def nvmlDeviceGetNumFans(device): + c_numFans = c_uint() + fn = _nvmlGetFunctionPointer("nvmlDeviceGetNumFans") + ret = fn(device, byref(c_numFans)) + _nvmlCheckReturn(ret) + return c_numFans.value + +def nvmlDeviceSetDefaultFanSpeed_v2(handle, index): + fn = _nvmlGetFunctionPointer("nvmlDeviceSetDefaultFanSpeed_v2"); + ret = fn(handle, index) + _nvmlCheckReturn(ret) + return NVML_SUCCESS + +def nvmlDeviceGetMinMaxFanSpeed(handle, minSpeed=c_uint(), maxSpeed=c_uint()): + isReference = (type(minSpeed) is not c_uint) or (type(maxSpeed) is not c_uint) + minSpeedRef = minSpeed if isReference else byref(minSpeed) + maxSpeedRef = maxSpeed if isReference else byref(maxSpeed) + fn = _nvmlGetFunctionPointer("nvmlDeviceGetMinMaxFanSpeed") + ret = fn(handle, minSpeedRef, maxSpeedRef) + _nvmlCheckReturn(ret) + return NVML_SUCCESS if isReference else [minSpeed.value, maxSpeed.value] + +def nvmlDeviceGetFanControlPolicy_v2(handle, fan, fanControlPolicy=c_uint()): + isReference = type(fanControlPolicy) is not c_uint + fanControlPolicyRef = fanControlPolicy if isReference else byref(fanControlPolicy) + fn = _nvmlGetFunctionPointer("nvmlDeviceGetFanControlPolicy_v2") + ret = fn(handle, fan, fanControlPolicyRef) + _nvmlCheckReturn(ret) + return NVML_SUCCESS if isReference else fanControlPolicy.value + +def nvmlDeviceSetFanControlPolicy(handle, fan, fanControlPolicy): + fn = _nvmlGetFunctionPointer("nvmlDeviceSetFanControlPolicy") + ret = fn(handle, fan, _nvmlFanControlPolicy_t(fanControlPolicy)) + _nvmlCheckReturn(ret) + return NVML_SUCCESS + +class c_nvmlTemperature_v1_t(_PrintableStructure): + _fields_ = [ + ('version', c_uint), + ('sensorType', _nvmlTemperatureSensors_t), + ('temperature', c_int), + ] +nvmlTemperature_v1 = 0x100000C + +def nvmlDeviceGetTemperatureV1(handle, sensor): + c_temp = c_nvmlTemperature_v1_t() + c_temp.version = nvmlTemperature_v1 + c_temp.sensorType = _nvmlTemperatureSensors_t(sensor) + fn = _nvmlGetFunctionPointer("nvmlDeviceGetTemperatureV") + ret = fn(handle, byref(c_temp)) + _nvmlCheckReturn(ret) + return c_temp.temperature + +def nvmlDeviceGetTemperatureV(handle, sensor, version=nvmlTemperature_v1): + if version == nvmlTemperature_v1: + return nvmlDeviceGetTemperatureV1(handle, sensor) + else: + raise NVMLError(NVML_ERROR_ARGUMENT_VERSION_MISMATCH) + +# DEPRECATED use nvmlDeviceGetTemperatureV instead +def nvmlDeviceGetTemperature(handle, sensor): + c_temp = c_uint() + fn = _nvmlGetFunctionPointer("nvmlDeviceGetTemperature") + ret = fn(handle, _nvmlTemperatureSensors_t(sensor), byref(c_temp)) + _nvmlCheckReturn(ret) + return c_temp.value + +def nvmlDeviceGetTemperatureThreshold(handle, threshold): + c_temp = c_uint() + fn = _nvmlGetFunctionPointer("nvmlDeviceGetTemperatureThreshold") + ret = fn(handle, _nvmlTemperatureThresholds_t(threshold), byref(c_temp)) + _nvmlCheckReturn(ret) + return c_temp.value + +def nvmlDeviceSetTemperatureThreshold(handle, threshold, temp): + c_temp = c_uint() + c_temp.value = temp + fn = _nvmlGetFunctionPointer("nvmlDeviceSetTemperatureThreshold") + ret = fn(handle, _nvmlTemperatureThresholds_t(threshold), byref(c_temp)) + _nvmlCheckReturn(ret) + return None + +def nvmlDeviceGetMarginTemperature(handle): + c_marginTempInfo = c_nvmlMarginTemperature_v1_t() + c_marginTempInfo.version = nvmlMarginTemperature_v1 + fn = _nvmlGetFunctionPointer("nvmlDeviceGetMarginTemperature") + ret = fn(handle, byref(c_marginTempInfo)) + _nvmlCheckReturn(ret) + return c_marginTempInfo.marginTemperature + +# DEPRECATED use nvmlDeviceGetPerformanceState +def nvmlDeviceGetPowerState(handle): + c_pstate = _nvmlPstates_t() + fn = _nvmlGetFunctionPointer("nvmlDeviceGetPowerState") + ret = fn(handle, byref(c_pstate)) + _nvmlCheckReturn(ret) + return c_pstate.value + +def nvmlDeviceGetPerformanceState(handle): + c_pstate = _nvmlPstates_t() + fn = _nvmlGetFunctionPointer("nvmlDeviceGetPerformanceState") + ret = fn(handle, byref(c_pstate)) + _nvmlCheckReturn(ret) + return c_pstate.value + +# Deprecated +def nvmlDeviceGetPowerManagementMode(handle): + c_pcapMode = _nvmlEnableState_t() + fn = _nvmlGetFunctionPointer("nvmlDeviceGetPowerManagementMode") + ret = fn(handle, byref(c_pcapMode)) + _nvmlCheckReturn(ret) + return c_pcapMode.value + +def nvmlDeviceGetPowerManagementLimit(handle): + c_limit = c_uint() + fn = _nvmlGetFunctionPointer("nvmlDeviceGetPowerManagementLimit") + ret = fn(handle, byref(c_limit)) + _nvmlCheckReturn(ret) + return c_limit.value + +# Added in 4.304 +def nvmlDeviceGetPowerManagementLimitConstraints(handle): + c_minLimit = c_uint() + c_maxLimit = c_uint() + fn = _nvmlGetFunctionPointer("nvmlDeviceGetPowerManagementLimitConstraints") + ret = fn(handle, byref(c_minLimit), byref(c_maxLimit)) + _nvmlCheckReturn(ret) + return [c_minLimit.value, c_maxLimit.value] + +# Added in 4.304 +def nvmlDeviceGetPowerManagementDefaultLimit(handle): + c_limit = c_uint() + fn = _nvmlGetFunctionPointer("nvmlDeviceGetPowerManagementDefaultLimit") + ret = fn(handle, byref(c_limit)) + _nvmlCheckReturn(ret) + return c_limit.value + + +# Added in 331 +def nvmlDeviceGetEnforcedPowerLimit(handle): + c_limit = c_uint() + fn = _nvmlGetFunctionPointer("nvmlDeviceGetEnforcedPowerLimit") + ret = fn(handle, byref(c_limit)) + _nvmlCheckReturn(ret) + return c_limit.value + +def nvmlDeviceGetPowerUsage(handle): + c_watts = c_uint() + fn = _nvmlGetFunctionPointer("nvmlDeviceGetPowerUsage") + ret = fn(handle, byref(c_watts)) + _nvmlCheckReturn(ret) + return c_watts.value + +def nvmlDeviceGetTotalEnergyConsumption(handle): + c_millijoules = c_uint64() + fn = _nvmlGetFunctionPointer("nvmlDeviceGetTotalEnergyConsumption") + ret = fn(handle, byref(c_millijoules)) + _nvmlCheckReturn(ret) + return c_millijoules.value + +# Added in 4.304 +def nvmlDeviceGetGpuOperationMode(handle): + c_currState = _nvmlGpuOperationMode_t() + c_pendingState = _nvmlGpuOperationMode_t() + fn = _nvmlGetFunctionPointer("nvmlDeviceGetGpuOperationMode") + ret = fn(handle, byref(c_currState), byref(c_pendingState)) + _nvmlCheckReturn(ret) + return [c_currState.value, c_pendingState.value] + +# Added in 4.304 +def nvmlDeviceGetCurrentGpuOperationMode(handle): + return nvmlDeviceGetGpuOperationMode(handle)[0] + +# Added in 4.304 +def nvmlDeviceGetPendingGpuOperationMode(handle): + return nvmlDeviceGetGpuOperationMode(handle)[1] + +def nvmlDeviceGetMemoryInfo(handle, version=None): + if not version: + c_memory = c_nvmlMemory_t() + fn = _nvmlGetFunctionPointer("nvmlDeviceGetMemoryInfo") + else: + c_memory = c_nvmlMemory_v2_t() + c_memory.version = version + fn = _nvmlGetFunctionPointer("nvmlDeviceGetMemoryInfo_v2") + ret = fn(handle, byref(c_memory)) + _nvmlCheckReturn(ret) + return c_memory + +def nvmlDeviceGetBAR1MemoryInfo(handle): + c_bar1_memory = c_nvmlBAR1Memory_t() + fn = _nvmlGetFunctionPointer("nvmlDeviceGetBAR1MemoryInfo") + ret = fn(handle, byref(c_bar1_memory)) + _nvmlCheckReturn(ret) + return c_bar1_memory + +def nvmlDeviceGetComputeMode(handle): + c_mode = _nvmlComputeMode_t() + fn = _nvmlGetFunctionPointer("nvmlDeviceGetComputeMode") + ret = fn(handle, byref(c_mode)) + _nvmlCheckReturn(ret) + return c_mode.value + +def nvmlDeviceGetCudaComputeCapability(handle): + c_major = c_int() + c_minor = c_int() + fn = _nvmlGetFunctionPointer("nvmlDeviceGetCudaComputeCapability") + ret = fn(handle, byref(c_major), byref(c_minor)) + _nvmlCheckReturn(ret) + return (c_major.value, c_minor.value) + +def nvmlDeviceGetEccMode(handle): + c_currState = _nvmlEnableState_t() + c_pendingState = _nvmlEnableState_t() + fn = _nvmlGetFunctionPointer("nvmlDeviceGetEccMode") + ret = fn(handle, byref(c_currState), byref(c_pendingState)) + _nvmlCheckReturn(ret) + return [c_currState.value, c_pendingState.value] + +# added to API +def nvmlDeviceGetCurrentEccMode(handle): + return nvmlDeviceGetEccMode(handle)[0] + +# added to API +def nvmlDeviceGetPendingEccMode(handle): + return nvmlDeviceGetEccMode(handle)[1] + +def nvmlDeviceGetDefaultEccMode(handle): + c_defaultState = _nvmlEnableState_t() + fn = _nvmlGetFunctionPointer("nvmlDeviceGetDefaultEccMode") + ret = fn(handle, byref(c_defaultState)) + _nvmlCheckReturn(ret) + return [c_defaultState.value] + +def nvmlDeviceGetTotalEccErrors(handle, errorType, counterType): + c_count = c_ulonglong() + fn = _nvmlGetFunctionPointer("nvmlDeviceGetTotalEccErrors") + ret = fn(handle, _nvmlMemoryErrorType_t(errorType), + _nvmlEccCounterType_t(counterType), byref(c_count)) + _nvmlCheckReturn(ret) + return c_count.value + +# This is deprecated, instead use nvmlDeviceGetMemoryErrorCounter +def nvmlDeviceGetDetailedEccErrors(handle, errorType, counterType): + c_counts = c_nvmlEccErrorCounts_t() + fn = _nvmlGetFunctionPointer("nvmlDeviceGetDetailedEccErrors") + ret = fn(handle, _nvmlMemoryErrorType_t(errorType), + _nvmlEccCounterType_t(counterType), byref(c_counts)) + _nvmlCheckReturn(ret) + return c_counts + +# Added in 4.304 +def nvmlDeviceGetMemoryErrorCounter(handle, errorType, counterType, locationType): + c_count = c_ulonglong() + fn = _nvmlGetFunctionPointer("nvmlDeviceGetMemoryErrorCounter") + ret = fn(handle, + _nvmlMemoryErrorType_t(errorType), + _nvmlEccCounterType_t(counterType), + _nvmlMemoryLocation_t(locationType), + byref(c_count)) + _nvmlCheckReturn(ret) + return c_count.value + +def nvmlDeviceGetUtilizationRates(handle): + c_util = c_nvmlUtilization_t() + fn = _nvmlGetFunctionPointer("nvmlDeviceGetUtilizationRates") + ret = fn(handle, byref(c_util)) + _nvmlCheckReturn(ret) + return c_util + +def nvmlDeviceGetEncoderUtilization(handle): + c_util = c_uint() + c_samplingPeriod = c_uint() + fn = _nvmlGetFunctionPointer("nvmlDeviceGetEncoderUtilization") + ret = fn(handle, byref(c_util), byref(c_samplingPeriod)) + _nvmlCheckReturn(ret) + return [c_util.value, c_samplingPeriod.value] + +def nvmlDeviceGetDecoderUtilization(handle): + c_util = c_uint() + c_samplingPeriod = c_uint() + fn = _nvmlGetFunctionPointer("nvmlDeviceGetDecoderUtilization") + ret = fn(handle, byref(c_util), byref(c_samplingPeriod)) + _nvmlCheckReturn(ret) + return [c_util.value, c_samplingPeriod.value] + +def nvmlDeviceGetJpgUtilization(handle): + c_util = c_uint() + c_samplingPeriod = c_uint() + fn = _nvmlGetFunctionPointer("nvmlDeviceGetJpgUtilization") + ret = fn(handle, byref(c_util), byref(c_samplingPeriod)) + _nvmlCheckReturn(ret) + return [c_util.value, c_samplingPeriod.value] + +def nvmlDeviceGetOfaUtilization(handle): + c_util = c_uint() + c_samplingPeriod = c_uint() + fn = _nvmlGetFunctionPointer("nvmlDeviceGetOfaUtilization") + ret = fn(handle, byref(c_util), byref(c_samplingPeriod)) + _nvmlCheckReturn(ret) + return [c_util.value, c_samplingPeriod.value] + +def nvmlDeviceGetPcieReplayCounter(handle): + c_replay = c_uint() + fn = _nvmlGetFunctionPointer("nvmlDeviceGetPcieReplayCounter") + ret = fn(handle, byref(c_replay)) + _nvmlCheckReturn(ret) + return c_replay.value + +def nvmlDeviceGetDriverModel(handle): + c_currModel = _nvmlDriverModel_t() + c_pendingModel = _nvmlDriverModel_t() + fn = _nvmlGetFunctionPointer("nvmlDeviceGetDriverModel") + ret = fn(handle, byref(c_currModel), byref(c_pendingModel)) + _nvmlCheckReturn(ret) + return [c_currModel.value, c_pendingModel.value] + +# added to API +def nvmlDeviceGetCurrentDriverModel(handle): + return nvmlDeviceGetDriverModel(handle)[0] + +# added to API +def nvmlDeviceGetPendingDriverModel(handle): + return nvmlDeviceGetDriverModel(handle)[1] + +# Added in 2.285 +@convertStrBytes +def nvmlDeviceGetVbiosVersion(handle): + c_version = create_string_buffer(NVML_DEVICE_VBIOS_VERSION_BUFFER_SIZE) + fn = _nvmlGetFunctionPointer("nvmlDeviceGetVbiosVersion") + ret = fn(handle, c_version, c_uint(NVML_DEVICE_VBIOS_VERSION_BUFFER_SIZE)) + _nvmlCheckReturn(ret) + return c_version.value + +# Added in 2.285 +def nvmlDeviceGetComputeRunningProcesses_v2(handle): + # first call to get the size + c_count = c_uint(0) + fn = _nvmlGetFunctionPointer("nvmlDeviceGetComputeRunningProcesses_v2") + ret = fn(handle, byref(c_count), None) + if (ret == NVML_SUCCESS): + # special case, no running processes + return [] + elif (ret == NVML_ERROR_INSUFFICIENT_SIZE): + # typical case + # oversize the array incase more processes are created + c_count.value = c_count.value * 2 + 5 + proc_array = c_nvmlProcessInfo_v2_t * c_count.value + c_procs = proc_array() + # make the call again + ret = fn(handle, byref(c_count), c_procs) + _nvmlCheckReturn(ret) + procs = [] + for i in range(c_count.value): + # use an alternative struct for this object + obj = nvmlStructToFriendlyObject(c_procs[i]) + if (obj.usedGpuMemory == NVML_VALUE_NOT_AVAILABLE_ulonglong.value): + # special case for WDDM on Windows, see comment above + obj.usedGpuMemory = None + procs.append(obj) + return procs + else: + # error case + raise NVMLError(ret) + +# Added in 2.285 +def nvmlDeviceGetComputeRunningProcesses_v3(handle): + # first call to get the size + c_count = c_uint(0) + fn = _nvmlGetFunctionPointer("nvmlDeviceGetComputeRunningProcesses_v3") + ret = fn(handle, byref(c_count), None) + + if (ret == NVML_SUCCESS): + # special case, no running processes + return [] + elif (ret == NVML_ERROR_INSUFFICIENT_SIZE): + # typical case + # oversize the array incase more processes are created + c_count.value = c_count.value * 2 + 5 + proc_array = c_nvmlProcessInfo_v3_t * c_count.value + c_procs = proc_array() + + # make the call again + ret = fn(handle, byref(c_count), c_procs) + _nvmlCheckReturn(ret) + + procs = [] + for i in range(c_count.value): + # use an alternative struct for this object + obj = nvmlStructToFriendlyObject(c_procs[i]) + if (obj.usedGpuMemory == NVML_VALUE_NOT_AVAILABLE_ulonglong.value): + # special case for WDDM on Windows, see comment above + obj.usedGpuMemory = None + procs.append(obj) + + return procs + else: + # error case + raise NVMLError(ret) + +@throwOnVersionMismatch +def nvmlDeviceGetComputeRunningProcesses(handle): + return nvmlDeviceGetComputeRunningProcesses_v3(handle) + +def nvmlDeviceGetGraphicsRunningProcesses_v2(handle): + # first call to get the size + c_count = c_uint(0) + fn = _nvmlGetFunctionPointer("nvmlDeviceGetGraphicsRunningProcesses_v2") + ret = fn(handle, byref(c_count), None) + if (ret == NVML_SUCCESS): + # special case, no running processes + return [] + elif (ret == NVML_ERROR_INSUFFICIENT_SIZE): + # typical case + # oversize the array incase more processes are created + c_count.value = c_count.value * 2 + 5 + proc_array = c_nvmlProcessInfo_v2_t * c_count.value + c_procs = proc_array() + # make the call again + ret = fn(handle, byref(c_count), c_procs) + _nvmlCheckReturn(ret) + procs = [] + for i in range(c_count.value): + # use an alternative struct for this object + obj = nvmlStructToFriendlyObject(c_procs[i]) + if (obj.usedGpuMemory == NVML_VALUE_NOT_AVAILABLE_ulonglong.value): + # special case for WDDM on Windows, see comment above + obj.usedGpuMemory = None + procs.append(obj) + return procs + else: + # error case + raise NVMLError(ret) + +def nvmlDeviceGetGraphicsRunningProcesses_v3(handle): + # first call to get the size + c_count = c_uint(0) + fn = _nvmlGetFunctionPointer("nvmlDeviceGetGraphicsRunningProcesses_v3") + ret = fn(handle, byref(c_count), None) + + if (ret == NVML_SUCCESS): + # special case, no running processes + return [] + elif (ret == NVML_ERROR_INSUFFICIENT_SIZE): + # typical case + # oversize the array incase more processes are created + c_count.value = c_count.value * 2 + 5 + proc_array = c_nvmlProcessInfo_v3_t * c_count.value + c_procs = proc_array() + + # make the call again + ret = fn(handle, byref(c_count), c_procs) + _nvmlCheckReturn(ret) + + procs = [] + for i in range(c_count.value): + # use an alternative struct for this object + obj = nvmlStructToFriendlyObject(c_procs[i]) + if (obj.usedGpuMemory == NVML_VALUE_NOT_AVAILABLE_ulonglong.value): + # special case for WDDM on Windows, see comment above + obj.usedGpuMemory = None + procs.append(obj) + + return procs + else: + # error case + raise NVMLError(ret) + +@throwOnVersionMismatch +def nvmlDeviceGetGraphicsRunningProcesses(handle): + return nvmlDeviceGetGraphicsRunningProcesses_v3(handle) + +@throwOnVersionMismatch +def nvmlDeviceGetMPSComputeRunningProcesses(handle): + return nvmlDeviceGetMPSComputeRunningProcesses_v3(handle) + +def nvmlDeviceGetMPSComputeRunningProcesses_v2(handle): + # first call to get the size + c_count = c_uint(0) + fn = _nvmlGetFunctionPointer("nvmlDeviceGetMPSComputeRunningProcesses_v2") + ret = fn(handle, byref(c_count), None) + + if (ret == NVML_SUCCESS): + # special case, no running processes + return [] + elif (ret == NVML_ERROR_INSUFFICIENT_SIZE): + # typical case + # oversize the array incase more processes are created + c_count.value = c_count.value * 2 + 5 + proc_array = c_nvmlProcessInfo_v2_t * c_count.value + c_procs = proc_array() + + # make the call again + ret = fn(handle, byref(c_count), c_procs) + _nvmlCheckReturn(ret) + + procs = [] + for i in range(c_count.value): + # use an alternative struct for this object + obj = nvmlStructToFriendlyObject(c_procs[i]) + if (obj.usedGpuMemory == NVML_VALUE_NOT_AVAILABLE_ulonglong.value): + # special case for WDDM on Windows, see comment above + obj.usedGpuMemory = None + procs.append(obj) + + return procs + else: + # error case + raise NVMLError(ret) + +def nvmlDeviceGetMPSComputeRunningProcesses_v3(handle): + # first call to get the size + c_count = c_uint(0) + fn = _nvmlGetFunctionPointer("nvmlDeviceGetMPSComputeRunningProcesses_v3") + ret = fn(handle, byref(c_count), None) + + if (ret == NVML_SUCCESS): + # special case, no running processes + return [] + elif (ret == NVML_ERROR_INSUFFICIENT_SIZE): + # typical case + # oversize the array incase more processes are created + c_count.value = c_count.value * 2 + 5 + proc_array = c_nvmlProcessInfo_v3_t * c_count.value + c_procs = proc_array() + + # make the call again + ret = fn(handle, byref(c_count), c_procs) + _nvmlCheckReturn(ret) + + procs = [] + for i in range(c_count.value): + # use an alternative struct for this object + obj = nvmlStructToFriendlyObject(c_procs[i]) + if (obj.usedGpuMemory == NVML_VALUE_NOT_AVAILABLE_ulonglong.value): + # special case for WDDM on Windows, see comment above + obj.usedGpuMemory = None + procs.append(obj) + + return procs + else: + # error case + raise NVMLError(ret) + +def nvmlDeviceGetRunningProcessDetailList(handle, version, mode): + c_processDetailList = c_nvmlProcessDetailList_t() + c_processDetailList.version = version + c_processDetailList.mode = mode + + fn = _nvmlGetFunctionPointer("nvmlDeviceGetRunningProcessDetailList") + + # first call to get the size + ret = fn(handle, byref(c_processDetailList)) + if (ret == NVML_SUCCESS): + # special case, no running processes + return [] + elif (ret == NVML_ERROR_INSUFFICIENT_SIZE): + c_procs = c_nvmlProcessDetail_v1_t * c_processDetailList.numProcArrayEntries + c_processDetailList.procArray = cast((c_procs)(), POINTER(c_nvmlProcessDetail_v1_t)) + + # make the call again + ret = fn(handle, byref(c_processDetailList)) + _nvmlCheckReturn(ret) + + procs = [] + for i in range(c_processDetailList.numProcArrayEntries): + # use an alternative struct for this object + obj = c_processDetailList.procArray[i] + if (obj.usedGpuMemory == NVML_VALUE_NOT_AVAILABLE_ulonglong.value): + obj.usedGpuMemory = None + if (obj.usedGpuCcProtectedMemory == NVML_VALUE_NOT_AVAILABLE_ulonglong.value): + obj.usedGpuCcProtectedMemory = None + procs.append(obj) + + return procs + else: + # error case + raise NVMLError(ret) + +def nvmlDeviceGetAutoBoostedClocksEnabled(handle): + c_isEnabled = _nvmlEnableState_t() + c_defaultIsEnabled = _nvmlEnableState_t() + fn = _nvmlGetFunctionPointer("nvmlDeviceGetAutoBoostedClocksEnabled") + ret = fn(handle, byref(c_isEnabled), byref(c_defaultIsEnabled)) + _nvmlCheckReturn(ret) + return [c_isEnabled.value, c_defaultIsEnabled.value] + #Throws NVML_ERROR_NOT_SUPPORTED if hardware doesn't support setting auto boosted clocks + +## Set functions +def nvmlUnitSetLedState(unit, color): + fn = _nvmlGetFunctionPointer("nvmlUnitSetLedState") + ret = fn(unit, _nvmlLedColor_t(color)) + _nvmlCheckReturn(ret) + return None + +def nvmlDeviceSetPersistenceMode(handle, mode): + fn = _nvmlGetFunctionPointer("nvmlDeviceSetPersistenceMode") + ret = fn(handle, _nvmlEnableState_t(mode)) + _nvmlCheckReturn(ret) + return None + +def nvmlDeviceSetComputeMode(handle, mode): + fn = _nvmlGetFunctionPointer("nvmlDeviceSetComputeMode") + ret = fn(handle, _nvmlComputeMode_t(mode)) + _nvmlCheckReturn(ret) + return None + +def nvmlDeviceSetEccMode(handle, mode): + fn = _nvmlGetFunctionPointer("nvmlDeviceSetEccMode") + ret = fn(handle, _nvmlEnableState_t(mode)) + _nvmlCheckReturn(ret) + return None + +def nvmlDeviceClearEccErrorCounts(handle, counterType): + fn = _nvmlGetFunctionPointer("nvmlDeviceClearEccErrorCounts") + ret = fn(handle, _nvmlEccCounterType_t(counterType)) + _nvmlCheckReturn(ret) + return None + +def nvmlDeviceSetDriverModel(handle, model): + fn = _nvmlGetFunctionPointer("nvmlDeviceSetDriverModel") + ret = fn(handle, _nvmlDriverModel_t(model)) + _nvmlCheckReturn(ret) + return None + +def nvmlDeviceSetAutoBoostedClocksEnabled(handle, enabled): + fn = _nvmlGetFunctionPointer("nvmlDeviceSetAutoBoostedClocksEnabled") + ret = fn(handle, _nvmlEnableState_t(enabled)) + _nvmlCheckReturn(ret) + return None + #Throws NVML_ERROR_NOT_SUPPORTED if hardware doesn't support setting auto boosted clocks + +def nvmlDeviceSetDefaultAutoBoostedClocksEnabled(handle, enabled, flags): + fn = _nvmlGetFunctionPointer("nvmlDeviceSetDefaultAutoBoostedClocksEnabled") + ret = fn(handle, _nvmlEnableState_t(enabled), c_uint(flags)) + _nvmlCheckReturn(ret) + return None + #Throws NVML_ERROR_NOT_SUPPORTED if hardware doesn't support setting auto boosted clocks + +def nvmlDeviceSetGpuLockedClocks(handle, minGpuClockMHz, maxGpuClockMHz): + fn = _nvmlGetFunctionPointer("nvmlDeviceSetGpuLockedClocks") + ret = fn(handle, c_uint(minGpuClockMHz), c_uint(maxGpuClockMHz)) + _nvmlCheckReturn(ret) + return None + +def nvmlDeviceResetGpuLockedClocks(handle): + fn = _nvmlGetFunctionPointer("nvmlDeviceResetGpuLockedClocks") + ret = fn(handle) + _nvmlCheckReturn(ret) + return None + +def nvmlDeviceSetMemoryLockedClocks(handle, minMemClockMHz, maxMemClockMHz): + fn = _nvmlGetFunctionPointer("nvmlDeviceSetMemoryLockedClocks") + ret = fn(handle, c_uint(minMemClockMHz), c_uint(maxMemClockMHz)) + _nvmlCheckReturn(ret) + return None + +def nvmlDeviceResetMemoryLockedClocks(handle): + fn = _nvmlGetFunctionPointer("nvmlDeviceResetMemoryLockedClocks") + ret = fn(handle) + _nvmlCheckReturn(ret) + return None + +def nvmlDeviceGetClkMonStatus(handle, c_clkMonInfo=nvmlClkMonStatus_t()): + isReference = type(c_clkMonInfo) is not nvmlClkMonStatus_t + c_clkMonInfoRef = c_clkMonInfo if isReference else byref(c_clkMonInfo) + fn = _nvmlGetFunctionPointer("nvmlDeviceGetClkMonStatus") + ret = fn(handle, c_clkMonInfoRef) + _nvmlCheckReturn(ret) + return NVML_SUCCESS if isReference else c_clkMonInfo + +# Added in 4.304 +# Deprecated +def nvmlDeviceSetApplicationsClocks(handle, maxMemClockMHz, maxGraphicsClockMHz): + fn = _nvmlGetFunctionPointer("nvmlDeviceSetApplicationsClocks") + ret = fn(handle, c_uint(maxMemClockMHz), c_uint(maxGraphicsClockMHz)) + _nvmlCheckReturn(ret) + return None + +# Added in 4.304 +# Deprecated +def nvmlDeviceResetApplicationsClocks(handle): + fn = _nvmlGetFunctionPointer("nvmlDeviceResetApplicationsClocks") + ret = fn(handle) + _nvmlCheckReturn(ret) + return None + +# Added in 4.304 +def nvmlDeviceSetPowerManagementLimit(handle, limit): + fn = _nvmlGetFunctionPointer("nvmlDeviceSetPowerManagementLimit") + ret = fn(handle, c_uint(limit)) + _nvmlCheckReturn(ret) + return None + +# Added in 4.304 +def nvmlDeviceSetGpuOperationMode(handle, mode): + fn = _nvmlGetFunctionPointer("nvmlDeviceSetGpuOperationMode") + ret = fn(handle, _nvmlGpuOperationMode_t(mode)) + _nvmlCheckReturn(ret) + return None + +# Added in 2.285 +def nvmlEventSetCreate(): + fn = _nvmlGetFunctionPointer("nvmlEventSetCreate") + eventSet = c_nvmlEventSet_t() + ret = fn(byref(eventSet)) + _nvmlCheckReturn(ret) + return eventSet + +# Added in 2.285 +def nvmlDeviceRegisterEvents(handle, eventTypes, eventSet): + fn = _nvmlGetFunctionPointer("nvmlDeviceRegisterEvents") + ret = fn(handle, c_ulonglong(eventTypes), eventSet) + _nvmlCheckReturn(ret) + return None + +# Added in 2.285 +def nvmlDeviceGetSupportedEventTypes(handle): + c_eventTypes = c_ulonglong() + fn = _nvmlGetFunctionPointer("nvmlDeviceGetSupportedEventTypes") + ret = fn(handle, byref(c_eventTypes)) + _nvmlCheckReturn(ret) + return c_eventTypes.value + +# raises NVML_ERROR_TIMEOUT exception on timeout +def nvmlEventSetWait_v2(eventSet, timeoutms): + fn = _nvmlGetFunctionPointer("nvmlEventSetWait_v2") + data = c_nvmlEventData_t() + ret = fn(eventSet, byref(data), c_uint(timeoutms)) + _nvmlCheckReturn(ret) + return data + +def nvmlEventSetWait(eventSet, timeoutms): + return nvmlEventSetWait_v2(eventSet, timeoutms) + +# Added in 2.285 +def nvmlEventSetFree(eventSet): + fn = _nvmlGetFunctionPointer("nvmlEventSetFree") + ret = fn(eventSet) + _nvmlCheckReturn(ret) + return None + +# Added in 3.295 +def nvmlDeviceOnSameBoard(handle1, handle2): + fn = _nvmlGetFunctionPointer("nvmlDeviceOnSameBoard") + onSameBoard = c_int() + ret = fn(handle1, handle2, byref(onSameBoard)) + _nvmlCheckReturn(ret) + return (onSameBoard.value != 0) + +# Added in 3.295 +def nvmlDeviceGetCurrPcieLinkGeneration(handle): + fn = _nvmlGetFunctionPointer("nvmlDeviceGetCurrPcieLinkGeneration") + gen = c_uint() + ret = fn(handle, byref(gen)) + _nvmlCheckReturn(ret) + return gen.value + +# Added in 3.295 +def nvmlDeviceGetMaxPcieLinkGeneration(handle): + fn = _nvmlGetFunctionPointer("nvmlDeviceGetMaxPcieLinkGeneration") + gen = c_uint() + ret = fn(handle, byref(gen)) + _nvmlCheckReturn(ret) + return gen.value + +# Added in 3.295 +def nvmlDeviceGetCurrPcieLinkWidth(handle): + fn = _nvmlGetFunctionPointer("nvmlDeviceGetCurrPcieLinkWidth") + width = c_uint() + ret = fn(handle, byref(width)) + _nvmlCheckReturn(ret) + return width.value + +# Added in 3.295 +def nvmlDeviceGetMaxPcieLinkWidth(handle): + fn = _nvmlGetFunctionPointer("nvmlDeviceGetMaxPcieLinkWidth") + width = c_uint() + ret = fn(handle, byref(width)) + _nvmlCheckReturn(ret) + return width.value + +def nvmlDeviceGetGpuMaxPcieLinkGeneration(handle): + fn = _nvmlGetFunctionPointer("nvmlDeviceGetGpuMaxPcieLinkGeneration") + gen = c_uint() + ret = fn(handle, byref(gen)) + _nvmlCheckReturn(ret) + return gen.value + +# Added in 4.304 +# Deprecated +def nvmlDeviceGetSupportedClocksThrottleReasons(handle): + c_reasons= c_ulonglong() + fn = _nvmlGetFunctionPointer("nvmlDeviceGetSupportedClocksThrottleReasons") + ret = fn(handle, byref(c_reasons)) + _nvmlCheckReturn(ret) + return c_reasons.value + +def nvmlDeviceGetSupportedClocksEventReasons(handle): + c_reasons= c_ulonglong() + fn = _nvmlGetFunctionPointer("nvmlDeviceGetSupportedClocksEventReasons") + ret = fn(handle, byref(c_reasons)) + _nvmlCheckReturn(ret) + return c_reasons.value + +# Added in 4.304 +# Deprecated +def nvmlDeviceGetCurrentClocksThrottleReasons(handle): + c_reasons= c_ulonglong() + fn = _nvmlGetFunctionPointer("nvmlDeviceGetCurrentClocksThrottleReasons") + ret = fn(handle, byref(c_reasons)) + _nvmlCheckReturn(ret) + return c_reasons.value + +def nvmlDeviceGetCurrentClocksEventReasons(handle): + c_reasons= c_ulonglong() + fn = _nvmlGetFunctionPointer("nvmlDeviceGetCurrentClocksEventReasons") + ret = fn(handle, byref(c_reasons)) + _nvmlCheckReturn(ret) + return c_reasons.value + +# Added in 5.319 +def nvmlDeviceGetIndex(handle): + fn = _nvmlGetFunctionPointer("nvmlDeviceGetIndex") + c_index = c_uint() + ret = fn(handle, byref(c_index)) + _nvmlCheckReturn(ret) + return c_index.value + +# Added in 5.319 +def nvmlDeviceGetAccountingMode(handle): + c_mode = _nvmlEnableState_t() + fn = _nvmlGetFunctionPointer("nvmlDeviceGetAccountingMode") + ret = fn(handle, byref(c_mode)) + _nvmlCheckReturn(ret) + return c_mode.value + +def nvmlDeviceSetAccountingMode(handle, mode): + fn = _nvmlGetFunctionPointer("nvmlDeviceSetAccountingMode") + ret = fn(handle, _nvmlEnableState_t(mode)) + _nvmlCheckReturn(ret) + return None + +def nvmlDeviceClearAccountingPids(handle): + fn = _nvmlGetFunctionPointer("nvmlDeviceClearAccountingPids") + ret = fn(handle) + _nvmlCheckReturn(ret) + return None + +def nvmlDeviceGetAccountingStats(handle, pid): + stats = c_nvmlAccountingStats_t() + fn = _nvmlGetFunctionPointer("nvmlDeviceGetAccountingStats") + ret = fn(handle, c_uint(pid), byref(stats)) + _nvmlCheckReturn(ret) + if (stats.maxMemoryUsage == NVML_VALUE_NOT_AVAILABLE_ulonglong.value): + # special case for WDDM on Windows, see comment above + stats.maxMemoryUsage = None + return stats + +def nvmlDeviceGetAccountingPids(handle): + count = c_uint(nvmlDeviceGetAccountingBufferSize(handle)) + pids = (c_uint * count.value)() + fn = _nvmlGetFunctionPointer("nvmlDeviceGetAccountingPids") + ret = fn(handle, byref(count), pids) + _nvmlCheckReturn(ret) + return list(map(int, pids[0:count.value])) + +def nvmlDeviceGetAccountingBufferSize(handle): + bufferSize = c_uint() + fn = _nvmlGetFunctionPointer("nvmlDeviceGetAccountingBufferSize") + ret = fn(handle, byref(bufferSize)) + _nvmlCheckReturn(ret) + return int(bufferSize.value) + +def nvmlDeviceGetRetiredPages(device, sourceFilter): + c_source = _nvmlPageRetirementCause_t(sourceFilter) + c_count = c_uint(0) + fn = _nvmlGetFunctionPointer("nvmlDeviceGetRetiredPages") + + # First call will get the size + ret = fn(device, c_source, byref(c_count), None) + + # this should only fail with insufficient size + if ((ret != NVML_SUCCESS) and + (ret != NVML_ERROR_INSUFFICIENT_SIZE)): + raise NVMLError(ret) + + # call again with a buffer + # oversize the array for the rare cases where additional pages + # are retired between NVML calls + c_count.value = c_count.value * 2 + 5 + page_array = c_ulonglong * c_count.value + c_pages = page_array() + ret = fn(device, c_source, byref(c_count), c_pages) + _nvmlCheckReturn(ret) + return list(map(int, c_pages[0:c_count.value])) + +def nvmlDeviceGetRetiredPages_v2(device, sourceFilter): + c_source = _nvmlPageRetirementCause_t(sourceFilter) + c_count = c_uint(0) + fn = _nvmlGetFunctionPointer("nvmlDeviceGetRetiredPages_v2") + + # First call will get the size + ret = fn(device, c_source, byref(c_count), None) + + # this should only fail with insufficient size + if ((ret != NVML_SUCCESS) and + (ret != NVML_ERROR_INSUFFICIENT_SIZE)): + raise NVMLError(ret) + + # call again with a buffer + # oversize the array for the rare cases where additional pages + # are retired between NVML calls + c_count.value = c_count.value * 2 + 5 + page_array = c_ulonglong * c_count.value + c_pages = page_array() + times_array = c_ulonglong * c_count.value + c_times = times_array() + ret = fn(device, c_source, byref(c_count), c_pages, c_times) + _nvmlCheckReturn(ret) + return [ { 'address': int(c_pages[i]), 'timestamp': int(c_times[i]) } for i in range(c_count.value) ]; + +def nvmlDeviceGetRetiredPagesPendingStatus(device): + c_pending = _nvmlEnableState_t() + fn = _nvmlGetFunctionPointer("nvmlDeviceGetRetiredPagesPendingStatus") + ret = fn(device, byref(c_pending)) + _nvmlCheckReturn(ret) + return int(c_pending.value) + +def nvmlDeviceGetAPIRestriction(device, apiType): + c_permission = _nvmlEnableState_t() + fn = _nvmlGetFunctionPointer("nvmlDeviceGetAPIRestriction") + ret = fn(device, _nvmlRestrictedAPI_t(apiType), byref(c_permission)) + _nvmlCheckReturn(ret) + return int(c_permission.value) + +def nvmlDeviceSetAPIRestriction(handle, apiType, isRestricted): + fn = _nvmlGetFunctionPointer("nvmlDeviceSetAPIRestriction") + ret = fn(handle, _nvmlRestrictedAPI_t(apiType), _nvmlEnableState_t(isRestricted)) + _nvmlCheckReturn(ret) + return None + +def nvmlDeviceGetBridgeChipInfo(handle): + bridgeHierarchy = c_nvmlBridgeChipHierarchy_t() + fn = _nvmlGetFunctionPointer("nvmlDeviceGetBridgeChipInfo") + ret = fn(handle, byref(bridgeHierarchy)) + _nvmlCheckReturn(ret) + return bridgeHierarchy + +def nvmlDeviceGetSamples(device, sampling_type, timeStamp): + c_sampling_type = _nvmlSamplingType_t(sampling_type) + c_time_stamp = c_ulonglong(timeStamp) + c_sample_count = c_uint(0) + c_sample_value_type = _nvmlValueType_t() + fn = _nvmlGetFunctionPointer("nvmlDeviceGetSamples") + + ## First Call gets the size + ret = fn(device, c_sampling_type, c_time_stamp, byref(c_sample_value_type), byref(c_sample_count), None) + + # Stop if this fails + if (ret != NVML_SUCCESS): + raise NVMLError(ret) + + sampleArray = c_sample_count.value * c_nvmlSample_t + c_samples = sampleArray() + ret = fn(device, c_sampling_type, c_time_stamp, byref(c_sample_value_type), byref(c_sample_count), c_samples) + _nvmlCheckReturn(ret) + return (c_sample_value_type.value, c_samples[0:c_sample_count.value]) + +# Deprecated +def nvmlDeviceGetViolationStatus(device, perfPolicyType): + c_perfPolicy_type = _nvmlPerfPolicyType_t(perfPolicyType) + c_violTime = c_nvmlViolationTime_t() + fn = _nvmlGetFunctionPointer("nvmlDeviceGetViolationStatus") + + ## Invoke the method to get violation time + ret = fn(device, c_perfPolicy_type, byref(c_violTime)) + _nvmlCheckReturn(ret) + return c_violTime + +def nvmlDeviceGetPcieThroughput(device, counter): + c_util = c_uint() + fn = _nvmlGetFunctionPointer("nvmlDeviceGetPcieThroughput") + ret = fn(device, _nvmlPcieUtilCounter_t(counter), byref(c_util)) + _nvmlCheckReturn(ret) + return c_util.value + +def nvmlSystemGetTopologyGpuSet(cpuNumber): + c_count = c_uint(0) + fn = _nvmlGetFunctionPointer("nvmlSystemGetTopologyGpuSet") + + # First call will get the size + ret = fn(cpuNumber, byref(c_count), None) + + if ret != NVML_SUCCESS: + raise NVMLError(ret) + # call again with a buffer + device_array = c_nvmlDevice_t * c_count.value + c_devices = device_array() + ret = fn(cpuNumber, byref(c_count), c_devices) + _nvmlCheckReturn(ret) + return list(c_devices[0:c_count.value]) + +def nvmlDeviceGetTopologyNearestGpus(device, level): + c_count = c_uint(0) + fn = _nvmlGetFunctionPointer("nvmlDeviceGetTopologyNearestGpus") + + # First call will get the size + ret = fn(device, level, byref(c_count), None) + + if ret != NVML_SUCCESS: + raise NVMLError(ret) + + # call again with a buffer + device_array = c_nvmlDevice_t * c_count.value + c_devices = device_array() + ret = fn(device, level, byref(c_count), c_devices) + _nvmlCheckReturn(ret) + return list(c_devices[0:c_count.value]) + +def nvmlDeviceGetTopologyCommonAncestor(device1, device2): + c_level = _nvmlGpuTopologyLevel_t() + fn = _nvmlGetFunctionPointer("nvmlDeviceGetTopologyCommonAncestor") + ret = fn(device1, device2, byref(c_level)) + _nvmlCheckReturn(ret) + return c_level.value + +# Deprecated +def nvmlDeviceGetNvLinkUtilizationCounter(device, link, counter): + c_rxcounter = c_ulonglong() + c_txcounter = c_ulonglong() + fn = _nvmlGetFunctionPointer("nvmlDeviceGetNvLinkUtilizationCounter") + ret = fn(device, link, counter, byref(c_rxcounter), byref(c_txcounter)) + _nvmlCheckReturn(ret) + return (c_rxcounter.value, c_txcounter.value) + +# Deprecated +def nvmlDeviceFreezeNvLinkUtilizationCounter(device, link, counter, freeze): + fn = _nvmlGetFunctionPointer("nvmlDeviceFreezeNvLinkUtilizationCounter") + ret = fn(device, link, counter, freeze) + _nvmlCheckReturn(ret) + return None + +# Deprecated +def nvmlDeviceResetNvLinkUtilizationCounter(device, link, counter): + fn = _nvmlGetFunctionPointer("nvmlDeviceResetNvLinkUtilizationCounter") + ret = fn(device, link, counter) + _nvmlCheckReturn(ret) + return None + +# Deprecated +def nvmlDeviceSetNvLinkUtilizationControl(device, link, counter, control, reset): + fn = _nvmlGetFunctionPointer("nvmlDeviceSetNvLinkUtilizationControl") + ret = fn(device, link, counter, byref(control), reset) + _nvmlCheckReturn(ret) + return None + +# Deprecated +def nvmlDeviceGetNvLinkUtilizationControl(device, link, counter): + c_control = nvmlNvLinkUtilizationControl_t() + fn = _nvmlGetFunctionPointer("nvmlDeviceGetNvLinkUtilizationControl") + ret = fn(device, link, counter, byref(c_control)) + _nvmlCheckReturn(ret) + return c_control + +def nvmlDeviceGetNvLinkCapability(device, link, capability): + c_capResult = c_uint() + fn = _nvmlGetFunctionPointer("nvmlDeviceGetNvLinkCapability") + ret = fn(device, link, capability, byref(c_capResult)) + _nvmlCheckReturn(ret) + return c_capResult.value + +def nvmlDeviceGetNvLinkErrorCounter(device, link, counter): + c_result = c_ulonglong() + fn = _nvmlGetFunctionPointer("nvmlDeviceGetNvLinkErrorCounter") + ret = fn(device, link, counter, byref(c_result)) + _nvmlCheckReturn(ret) + return c_result.value + +def nvmlDeviceResetNvLinkErrorCounters(device, link): + fn = _nvmlGetFunctionPointer("nvmlDeviceResetNvLinkErrorCounters") + ret = fn(device, link) + _nvmlCheckReturn(ret) + return None + +def nvmlDeviceGetNvLinkRemotePciInfo(device, link): + c_pci = nvmlPciInfo_t() + fn = _nvmlGetFunctionPointer("nvmlDeviceGetNvLinkRemotePciInfo_v2") + ret = fn(device, link, byref(c_pci)) + _nvmlCheckReturn(ret) + return c_pci + +def nvmlDeviceGetNvLinkRemoteDeviceType(handle, link): + c_type = _nvmlNvLinkDeviceType_t() + fn = _nvmlGetFunctionPointer("nvmlDeviceGetNvLinkRemoteDeviceType") + ret = fn(handle, link, byref(c_type)) + _nvmlCheckReturn(ret) + return c_type.value + +def nvmlDeviceGetNvLinkState(device, link): + c_isActive = c_uint() + fn = _nvmlGetFunctionPointer("nvmlDeviceGetNvLinkState") + ret = fn(device, link, byref(c_isActive)) + _nvmlCheckReturn(ret) + return c_isActive.value + +def nvmlDeviceGetNvLinkVersion(device, link): + c_version = c_uint() + fn = _nvmlGetFunctionPointer("nvmlDeviceGetNvLinkVersion") + ret = fn(device, link, byref(c_version)) + _nvmlCheckReturn(ret) + return c_version.value + +def nvmlDeviceModifyDrainState(pciInfo, newState): + fn = _nvmlGetFunctionPointer("nvmlDeviceModifyDrainState") + ret = fn(pointer(pciInfo), newState) + _nvmlCheckReturn(ret) + return None + +def nvmlDeviceQueryDrainState(pciInfo): + c_newState = c_uint() + fn = _nvmlGetFunctionPointer("nvmlDeviceQueryDrainState") + ret = fn(pointer(pciInfo), byref(c_newState)) + _nvmlCheckReturn(ret) + return c_newState.value + +def nvmlDeviceRemoveGpu(pciInfo): + fn = _nvmlGetFunctionPointer("nvmlDeviceRemoveGpu") + ret = fn(pointer(pciInfo)) + _nvmlCheckReturn(ret) + return None + +def nvmlDeviceDiscoverGpus(pciInfo): + fn = _nvmlGetFunctionPointer("nvmlDeviceDiscoverGpus") + ret = fn(pointer(pciInfo)) + _nvmlCheckReturn(ret) + return None + +def nvmlDeviceGetFieldValues(handle, fieldIds): + values_arr = c_nvmlFieldValue_t * len(fieldIds) + values = values_arr() + fn = _nvmlGetFunctionPointer("nvmlDeviceGetFieldValues") + + for i, fieldId in enumerate(fieldIds): + try: + (values[i].fieldId, values[i].scopeId) = fieldId + except TypeError: + values[i].fieldId = fieldId + + ret = fn(handle, c_int32(len(fieldIds)), byref(values)) + _nvmlCheckReturn(ret) + return values + +def nvmlDeviceClearFieldValues(handle, fieldIds): + values_arr = c_nvmlFieldValue_t * len(fieldIds) + values = values_arr() + fn = _nvmlGetFunctionPointer("nvmlDeviceClearFieldValues") + + for i, fieldId in enumerate(fieldIds): + try: + (values[i].fieldId, values[i].scopeId) = fieldId + except TypeError: + values[i].fieldId = fieldId + + ret = fn(handle, c_int32(len(fieldIds)), byref(values)) + _nvmlCheckReturn(ret) + return values + +def nvmlDeviceGetVirtualizationMode(handle): + c_virtualization_mode = c_ulonglong() + fn = _nvmlGetFunctionPointer("nvmlDeviceGetVirtualizationMode") + ret = fn(handle, byref(c_virtualization_mode)) + _nvmlCheckReturn(ret) + return c_virtualization_mode.value + +def nvmlDeviceSetVirtualizationMode(handle, virtualization_mode): + fn = _nvmlGetFunctionPointer("nvmlDeviceSetVirtualizationMode") + return fn(handle, virtualization_mode) + +def nvmlDeviceGetVgpuHeterogeneousMode(handle): + c_vgpuHeterogeneousMode = c_nvmlVgpuHeterogeneousMode_v1_t(0) + c_vgpuHeterogeneousMode.version = VgpuHeterogeneousMode_v1 + fn = _nvmlGetFunctionPointer("nvmlDeviceGetVgpuHeterogeneousMode") + ret = fn(handle, byref(c_vgpuHeterogeneousMode)) + _nvmlCheckReturn(ret) + return c_vgpuHeterogeneousMode.mode + +def nvmlDeviceSetVgpuHeterogeneousMode(handle, heterogeneous_mode): + c_vgpuHeterogeneousMode = c_nvmlVgpuHeterogeneousMode_v1_t(0) + c_vgpuHeterogeneousMode.version = VgpuHeterogeneousMode_v1 + c_vgpuHeterogeneousMode.mode = heterogeneous_mode + fn = _nvmlGetFunctionPointer("nvmlDeviceSetVgpuHeterogeneousMode") + ret = fn(handle, byref(c_vgpuHeterogeneousMode)) + _nvmlCheckReturn(ret) + return NVML_SUCCESS + +def nvmlVgpuInstanceGetPlacementId(vgpuInstance): + c_placement = c_nvmlVgpuPlacementId_v1_t(0) + c_placement.version = VgpuPlacementId_v1 + fn = _nvmlGetFunctionPointer("nvmlVgpuInstanceGetPlacementId") + ret = fn(vgpuInstance, byref(c_placement)) + _nvmlCheckReturn(ret) + return c_placement.placementId + +def nvmlDeviceGetVgpuTypeSupportedPlacements(handle, vgpuTypeId, mode=0, version=1): + c_max_instances = c_uint(0) + fn = _nvmlGetFunctionPointer("nvmlVgpuTypeGetMaxInstances") + ret = fn(handle, vgpuTypeId, byref(c_max_instances)) + _nvmlCheckReturn(ret) + + if version == 2: + c_vgpu_placements = c_nvmlVgpuPlacementList_v2_t() + c_vgpu_placements.version = VgpuPlacementList_v2 + c_vgpu_placements.count = c_max_instances.value + c_vgpu_placements.mode = mode + elif version == 1: + c_vgpu_placements = c_nvmlVgpuPlacementList_v1_t() + c_vgpu_placements.version = VgpuPlacementList_v1 + else: + raise NVMLError(NVML_ERROR_ARGUMENT_VERSION_MISMATCH) + + c_placements = c_uint * c_max_instances.value + c_vgpu_placements.placementIds = c_placements() + fn = _nvmlGetFunctionPointer("nvmlDeviceGetVgpuTypeSupportedPlacements") + ret = fn(handle, vgpuTypeId, byref(c_vgpu_placements)) + _nvmlCheckReturn(ret) + return c_vgpu_placements + +def nvmlDeviceGetVgpuTypeCreatablePlacements(handle, vgpuTypeId, version=1): + c_max_instances = c_uint(0) + fn = _nvmlGetFunctionPointer("nvmlVgpuTypeGetMaxInstances") + ret = fn(handle, vgpuTypeId, byref(c_max_instances)) + _nvmlCheckReturn(ret) + + if version == 2: + c_vgpu_placements = c_nvmlVgpuPlacementList_v2_t() + c_vgpu_placements.version = VgpuPlacementList_v2 + c_vgpu_placements.count = c_max_instances.value + elif version == 1: + c_vgpu_placements = c_nvmlVgpuPlacementList_v1_t() + c_vgpu_placements.version = VgpuPlacementList_v1 + + c_placements = c_uint * c_max_instances.value + c_vgpu_placements.placementIds = c_placements() + fn = _nvmlGetFunctionPointer("nvmlDeviceGetVgpuTypeCreatablePlacements") + ret = fn(handle, vgpuTypeId, byref(c_vgpu_placements)) + _nvmlCheckReturn(ret) + return c_vgpu_placements + +def nvmlGetVgpuDriverCapabilities(capability): + c_capResult = c_uint() + fn = _nvmlGetFunctionPointer("nvmlGetVgpuDriverCapabilities") + ret = fn(_nvmlVgpuDriverCapability_t(capability), byref(c_capResult)) + _nvmlCheckReturn(ret) + return c_capResult.value + +def nvmlDeviceGetVgpuCapabilities(handle, capability): + c_capResult = c_uint() + fn = _nvmlGetFunctionPointer("nvmlDeviceGetVgpuCapabilities") + ret = fn(handle, _nvmlDeviceVgpuCapability_t(capability), byref(c_capResult)) + _nvmlCheckReturn(ret) + return c_capResult.value + +def nvmlDeviceSetVgpuCapabilities(handle, capability, state): + fn = _nvmlGetFunctionPointer("nvmlDeviceSetVgpuCapabilities") + ret = fn(handle, _nvmlDeviceVgpuCapability_t(capability), state) + _nvmlCheckReturn(ret) + return NVML_SUCCESS + +def nvmlDeviceGetSupportedVgpus(handle): + # first call to get the size + c_vgpu_count = c_uint(0) + + fn = _nvmlGetFunctionPointer("nvmlDeviceGetSupportedVgpus") + ret = fn(handle, byref(c_vgpu_count), None) + + if (ret == NVML_SUCCESS): + # special case, no supported vGPUs + return [] + elif (ret == NVML_ERROR_INSUFFICIENT_SIZE): + # typical case + vgpu_type_ids_array = _nvmlVgpuTypeId_t * c_vgpu_count.value + c_vgpu_type_ids = vgpu_type_ids_array() + + # make the call again + ret = fn(handle, byref(c_vgpu_count), c_vgpu_type_ids) + _nvmlCheckReturn(ret) + vgpus = [] + for i in range(c_vgpu_count.value): + vgpus.append(c_vgpu_type_ids[i]) + return vgpus + else: + # error case + raise NVMLError(ret) + +def nvmlDeviceGetCreatableVgpus(handle): + # first call to get the size + c_vgpu_count = c_uint(0) + + fn = _nvmlGetFunctionPointer("nvmlDeviceGetCreatableVgpus") + ret = fn(handle, byref(c_vgpu_count), None) + + if (ret == NVML_SUCCESS): + # special case, no supported vGPUs + return [] + elif (ret == NVML_ERROR_INSUFFICIENT_SIZE): + # typical case + vgpu_type_ids_array = _nvmlVgpuTypeId_t * c_vgpu_count.value + c_vgpu_type_ids = vgpu_type_ids_array() + + # make the call again + ret = fn(handle, byref(c_vgpu_count), c_vgpu_type_ids) + _nvmlCheckReturn(ret) + vgpus = [] + for i in range(c_vgpu_count.value): + vgpus.append(c_vgpu_type_ids[i]) + return vgpus + else: + # error case + raise NVMLError(ret) + +def nvmlVgpuTypeGetGpuInstanceProfileId(vgpuTypeId): + c_profile_id = c_uint(0) + fn = _nvmlGetFunctionPointer("nvmlVgpuTypeGetGpuInstanceProfileId") + ret = fn(vgpuTypeId, byref(c_profile_id)) + _nvmlCheckReturn(ret) + return (c_profile_id.value) + +@convertStrBytes +def nvmlVgpuTypeGetClass(vgpuTypeId): + c_class = create_string_buffer(NVML_DEVICE_NAME_BUFFER_SIZE) + c_buffer_size = c_uint(NVML_DEVICE_NAME_BUFFER_SIZE) + fn = _nvmlGetFunctionPointer("nvmlVgpuTypeGetClass") + ret = fn(vgpuTypeId, c_class, byref(c_buffer_size)) + _nvmlCheckReturn(ret) + return c_class.value + +@convertStrBytes +def nvmlVgpuTypeGetName(vgpuTypeId): + c_name = create_string_buffer(NVML_DEVICE_NAME_BUFFER_SIZE) + c_buffer_size = c_uint(NVML_DEVICE_NAME_BUFFER_SIZE) + fn = _nvmlGetFunctionPointer("nvmlVgpuTypeGetName") + ret = fn(vgpuTypeId, c_name, byref(c_buffer_size)) + _nvmlCheckReturn(ret) + return c_name.value + +def nvmlVgpuTypeGetDeviceID(vgpuTypeId): + c_device_id = c_ulonglong(0) + c_subsystem_id = c_ulonglong(0) + fn = _nvmlGetFunctionPointer("nvmlVgpuTypeGetDeviceID") + ret = fn(vgpuTypeId, byref(c_device_id), byref(c_subsystem_id)) + _nvmlCheckReturn(ret) + return (c_device_id.value, c_subsystem_id.value) + +def nvmlVgpuTypeGetFramebufferSize(vgpuTypeId): + c_fb_size = c_ulonglong(0) + fn = _nvmlGetFunctionPointer("nvmlVgpuTypeGetFramebufferSize") + ret = fn(vgpuTypeId, byref(c_fb_size)) + _nvmlCheckReturn(ret) + return c_fb_size.value + +def nvmlVgpuTypeGetNumDisplayHeads(vgpuTypeId): + c_num_heads = c_uint(0) + fn = _nvmlGetFunctionPointer("nvmlVgpuTypeGetNumDisplayHeads") + ret = fn(vgpuTypeId, byref(c_num_heads)) + _nvmlCheckReturn(ret) + return c_num_heads.value + +def nvmlVgpuTypeGetResolution(vgpuTypeId): + c_xdim = c_uint(0) + c_ydim = c_uint(0) + fn = _nvmlGetFunctionPointer("nvmlVgpuTypeGetResolution") + ret = fn(vgpuTypeId, 0, byref(c_xdim), byref(c_ydim)) + _nvmlCheckReturn(ret) + return (c_xdim.value, c_ydim.value) + +@convertStrBytes +def nvmlVgpuTypeGetLicense(vgpuTypeId): + c_license = create_string_buffer(NVML_GRID_LICENSE_BUFFER_SIZE) + c_buffer_size = c_uint(NVML_GRID_LICENSE_BUFFER_SIZE) + fn = _nvmlGetFunctionPointer("nvmlVgpuTypeGetLicense") + ret = fn(vgpuTypeId, c_license, c_buffer_size) + _nvmlCheckReturn(ret) + return c_license.value + +def nvmlVgpuTypeGetFrameRateLimit(vgpuTypeId): + c_frl_config = c_uint(0) + fn = _nvmlGetFunctionPointer("nvmlVgpuTypeGetFrameRateLimit") + ret = fn(vgpuTypeId, byref(c_frl_config)) + _nvmlCheckReturn(ret) + return c_frl_config.value + +def nvmlVgpuTypeGetGspHeapSize(vgpuTypeId): + c_gsp_heap = c_uint(0) + fn = _nvmlGetFunctionPointer("nvmlVgpuTypeGetGspHeapSize") + ret = fn(vgpuTypeId, byref(c_gsp_heap)) + _nvmlCheckReturn(ret) + return c_gsp_heap.value + +def nvmlVgpuTypeGetFbReservation(vgpuTypeId): + c_fb_reservation = c_uint(0) + fn = _nvmlGetFunctionPointer("nvmlVgpuTypeGetFbReservation") + ret = fn(vgpuTypeId, byref(c_fb_reservation)) + _nvmlCheckReturn(ret) + return c_fb_reservation.value + +def nvmlVgpuInstanceGetRuntimeStateSize(vgpuInstance): + c_runtime_state = nvmlVgpuRuntimeState_v1_t() + c_runtime_state.version = VgpuRuntimeState_v1 + fn = _nvmlGetFunctionPointer("nvmlVgpuInstanceGetRuntimeStateSize") + ret = fn(vgpuInstance, byref(c_runtime_state)) + _nvmlCheckReturn(ret) + return c_runtime_state + +def nvmlVgpuTypeGetMaxInstances(handle, vgpuTypeId): + c_max_instances = c_uint(0) + fn = _nvmlGetFunctionPointer("nvmlVgpuTypeGetMaxInstances") + ret = fn(handle, vgpuTypeId, byref(c_max_instances)) + _nvmlCheckReturn(ret) + return c_max_instances.value + +def nvmlVgpuTypeGetMaxInstancesPerVm(vgpuTypeId): + c_max_instances_per_vm = c_uint(0) + fn = _nvmlGetFunctionPointer("nvmlVgpuTypeGetMaxInstancesPerVm") + ret = fn(vgpuTypeId, byref(c_max_instances_per_vm)) + _nvmlCheckReturn(ret) + return c_max_instances_per_vm.value + +def nvmlVgpuTypeGetBAR1Info(vgpuTypeId): + c_bar1Info = c_nvmlVgpuTypeBar1Info_v1_t(0) + c_bar1Info.version = VgpuTypeBar1Info_v1 + fn = _nvmlGetFunctionPointer("nvmlVgpuTypeGetBAR1Info") + ret = fn(vgpuTypeId, byref(c_bar1Info)) + _nvmlCheckReturn(ret) + return c_bar1Info + +def nvmlDeviceGetActiveVgpus(handle): + # first call to get the size + c_vgpu_count = c_uint(0) + + fn = _nvmlGetFunctionPointer("nvmlDeviceGetActiveVgpus") + ret = fn(handle, byref(c_vgpu_count), None) + + if (ret == NVML_SUCCESS): + # special case, no active vGPUs + return [] + elif (ret == NVML_ERROR_INSUFFICIENT_SIZE): + # typical case + vgpu_instance_array = _nvmlVgpuInstance_t * c_vgpu_count.value + c_vgpu_instances = vgpu_instance_array() + + # make the call again + ret = fn(handle, byref(c_vgpu_count), c_vgpu_instances) + _nvmlCheckReturn(ret) + vgpus = [] + for i in range(c_vgpu_count.value): + vgpus.append(c_vgpu_instances[i]) + return vgpus + else: + # error case + raise NVMLError(ret) + +@convertStrBytes +def nvmlVgpuInstanceGetVmID(vgpuInstance): + c_vm_id = create_string_buffer(NVML_DEVICE_UUID_BUFFER_SIZE) + c_buffer_size = c_uint(NVML_GRID_LICENSE_BUFFER_SIZE) + c_vm_id_type = c_uint(0) + fn = _nvmlGetFunctionPointer("nvmlVgpuInstanceGetVmID") + ret = fn(vgpuInstance, byref(c_vm_id), c_buffer_size, byref(c_vm_id_type)) + _nvmlCheckReturn(ret) + return (c_vm_id.value, c_vm_id_type.value) + +@convertStrBytes +def nvmlVgpuInstanceGetUUID(vgpuInstance): + c_uuid = create_string_buffer(NVML_DEVICE_UUID_BUFFER_SIZE) + c_buffer_size = c_uint(NVML_DEVICE_UUID_BUFFER_SIZE) + fn = _nvmlGetFunctionPointer("nvmlVgpuInstanceGetUUID") + ret = fn(vgpuInstance, byref(c_uuid), c_buffer_size) + _nvmlCheckReturn(ret) + return c_uuid.value + +@convertStrBytes +def nvmlVgpuInstanceGetMdevUUID(vgpuInstance): + c_uuid = create_string_buffer(NVML_DEVICE_UUID_BUFFER_SIZE) + c_buffer_size = c_uint(NVML_DEVICE_UUID_BUFFER_SIZE) + fn = _nvmlGetFunctionPointer("nvmlVgpuInstanceGetMdevUUID") + ret = fn(vgpuInstance, byref(c_uuid), c_buffer_size) + _nvmlCheckReturn(ret) + return c_uuid.value + +@convertStrBytes +def nvmlVgpuInstanceGetVmDriverVersion(vgpuInstance): + c_driver_version = create_string_buffer(NVML_SYSTEM_DRIVER_VERSION_BUFFER_SIZE) + c_buffer_size = c_uint(NVML_SYSTEM_DRIVER_VERSION_BUFFER_SIZE) + fn = _nvmlGetFunctionPointer("nvmlVgpuInstanceGetVmDriverVersion") + ret = fn(vgpuInstance, byref(c_driver_version), c_buffer_size) + _nvmlCheckReturn(ret) + return c_driver_version.value + +# Deprecated +def nvmlVgpuInstanceGetLicenseStatus(vgpuInstance): + c_license_status = c_uint(0) + fn = _nvmlGetFunctionPointer("nvmlVgpuInstanceGetLicenseStatus") + ret = fn(vgpuInstance, byref(c_license_status)) + _nvmlCheckReturn(ret) + return c_license_status.value + +def nvmlVgpuInstanceGetLicenseInfo_v2(vgpuInstance): + fn = _nvmlGetFunctionPointer("nvmlVgpuInstanceGetLicenseInfo_v2") + c_license_info = c_nvmlVgpuLicenseInfo_t() + ret = fn(vgpuInstance, byref(c_license_info)) + _nvmlCheckReturn(ret) + return c_license_info + +def nvmlVgpuInstanceGetLicenseInfo(vgpuInstance): + return nvmlVgpuInstanceGetLicenseInfo_v2(vgpuInstance) + +def nvmlVgpuInstanceGetFrameRateLimit(vgpuInstance): + c_frl = c_uint(0) + fn = _nvmlGetFunctionPointer("nvmlVgpuInstanceGetFrameRateLimit") + ret = fn(vgpuInstance, byref(c_frl)) + _nvmlCheckReturn(ret) + return c_frl.value + +def nvmlVgpuInstanceGetEccMode(vgpuInstance): + c_mode = _nvmlEnableState_t() + fn = _nvmlGetFunctionPointer("nvmlVgpuInstanceGetEccMode") + ret = fn(vgpuInstance, byref(c_mode)) + _nvmlCheckReturn(ret) + return c_mode.value + +def nvmlVgpuInstanceGetType(vgpuInstance): + c_vgpu_type = c_uint(0) + fn = _nvmlGetFunctionPointer("nvmlVgpuInstanceGetType") + ret = fn(vgpuInstance, byref(c_vgpu_type)) + _nvmlCheckReturn(ret) + return c_vgpu_type.value + +def nvmlVgpuInstanceGetEncoderCapacity(vgpuInstance): + c_encoder_capacity = c_ulonglong(0) + fn = _nvmlGetFunctionPointer("nvmlVgpuInstanceGetEncoderCapacity") + ret = fn(vgpuInstance, byref(c_encoder_capacity)) + _nvmlCheckReturn(ret) + return c_encoder_capacity.value + +def nvmlVgpuInstanceSetEncoderCapacity(vgpuInstance, encoder_capacity): + fn = _nvmlGetFunctionPointer("nvmlVgpuInstanceSetEncoderCapacity") + return fn(vgpuInstance, encoder_capacity) + +def nvmlVgpuInstanceGetFbUsage(vgpuInstance): + c_fb_usage = c_uint(0) + fn = _nvmlGetFunctionPointer("nvmlVgpuInstanceGetFbUsage") + ret = fn(vgpuInstance, byref(c_fb_usage)) + _nvmlCheckReturn(ret) + return c_fb_usage.value + +def nvmlVgpuTypeGetCapabilities(vgpuTypeId, capability): + c_cap_result = c_uint(0) + fn = _nvmlGetFunctionPointer("nvmlVgpuTypeGetCapabilities") + ret = fn(vgpuTypeId, _nvmlVgpuCapability_t(capability), byref(c_cap_result)) + _nvmlCheckReturn(ret) + return (c_cap_result.value) + +def nvmlVgpuInstanceGetGpuInstanceId(vgpuInstance): + c_id = c_uint(0) + fn = _nvmlGetFunctionPointer("nvmlVgpuInstanceGetGpuInstanceId") + ret = fn(vgpuInstance, byref(c_id)) + _nvmlCheckReturn(ret) + return (c_id.value) + +@convertStrBytes +def nvmlVgpuInstanceGetGpuPciId(vgpuInstance): + c_vgpuPciId = create_string_buffer(NVML_DEVICE_PCI_BUS_ID_BUFFER_SIZE) + fn = _nvmlGetFunctionPointer("nvmlVgpuInstanceGetGpuPciId") + ret = fn(vgpuInstance, c_vgpuPciId, byref(c_uint(NVML_DEVICE_PCI_BUS_ID_BUFFER_SIZE))) + _nvmlCheckReturn(ret) + return c_vgpuPciId.value + +def nvmlDeviceGetVgpuUtilization(handle, timeStamp): + # first call to get the size + c_vgpu_count = c_uint(0) + c_time_stamp = c_ulonglong(timeStamp) + c_sample_value_type = _nvmlValueType_t() + + fn = _nvmlGetFunctionPointer("nvmlDeviceGetVgpuUtilization") + ret = fn(handle, c_time_stamp, byref(c_sample_value_type), byref(c_vgpu_count), None) + + if (ret == NVML_SUCCESS): + # special case, no active vGPUs + return [] + elif (ret == NVML_ERROR_INSUFFICIENT_SIZE): + # typical case + sampleArray = c_vgpu_count.value * c_nvmlVgpuInstanceUtilizationSample_t + c_samples = sampleArray() + + # make the call again + ret = fn(handle, c_time_stamp, byref(c_sample_value_type), byref(c_vgpu_count), c_samples) + _nvmlCheckReturn(ret) + + return c_samples[0:c_vgpu_count.value] + else: + # error case + raise NVMLError(ret) + +def nvmlDeviceGetVgpuInstancesUtilizationInfo(handle, timeStamp): + # first call to get the size + c_time_stamp = c_ulonglong(timeStamp) + c_vgpuUtilInfo = c_nvmlVgpuInstancesUtilizationInfo_v1_t(0) + c_vgpuUtilInfo.version = VgpuInstancesUtilizationInfo_v1 + c_vgpuUtilInfo.sampleValType = _nvmlValueType_t() + c_vgpuUtilInfo.vgpuInstanceCount = c_uint(0) + c_vgpuUtilInfo.lastSeenTimeStamp = c_time_stamp + + fn = _nvmlGetFunctionPointer("nvmlDeviceGetVgpuInstancesUtilizationInfo") + ret = fn(handle, byref(c_vgpuUtilInfo)) + + if (ret == NVML_SUCCESS): + # special case, no active vGPUs + return [] + elif (ret == NVML_ERROR_INSUFFICIENT_SIZE): + # typical case + sampleArray = c_vgpuUtilInfo.vgpuInstanceCount * c_nvmlVgpuInstanceUtilizationInfo_v1_t + c_samples = sampleArray() + c_vgpuUtilInfo.vgpuUtilArray = c_samples + + # make the call again + ret = fn(handle, byref(c_vgpuUtilInfo)) + _nvmlCheckReturn(ret) + + return c_samples[0:c_vgpuUtilInfo.vgpuInstanceCount] + else: + # error case + raise NVMLError(ret) + +def nvmlDeviceGetP2PStatus(device1, device2, p2pIndex): + c_p2pstatus = _nvmlGpuP2PStatus_t() + fn = _nvmlGetFunctionPointer("nvmlDeviceGetP2PStatus") + ret = fn(device1, device2,p2pIndex, byref(c_p2pstatus)) + _nvmlCheckReturn(ret) + return c_p2pstatus.value + +def nvmlDeviceGetGridLicensableFeatures_v4(handle): + c_get_grid_licensable_features = c_nvmlGridLicensableFeatures_v4_t() + fn = _nvmlGetFunctionPointer("nvmlDeviceGetGridLicensableFeatures_v4") + ret = fn(handle, byref(c_get_grid_licensable_features)) + _nvmlCheckReturn(ret) + + return (c_get_grid_licensable_features) + +def nvmlDeviceGetGridLicensableFeatures(handle): + return nvmlDeviceGetGridLicensableFeatures_v4(handle) + +def nvmlDeviceGetGspFirmwareVersion(handle, version=None): + isUserDefined = version is not None + if not isUserDefined: + version = (c_char * NVML_GSP_FIRMWARE_VERSION_BUF_SIZE)() + fn = _nvmlGetFunctionPointer("nvmlDeviceGetGspFirmwareVersion") + ret = fn(handle, version) + _nvmlCheckReturn(ret) + return NVML_SUCCESS if isUserDefined else version.value + +def nvmlDeviceGetGspFirmwareMode(handle, isEnabled=c_uint(), defaultMode=c_uint()): + isReference = type(isEnabled) is not c_uint + isEnabledRef = isEnabled if isReference else byref(isEnabled) + defaultModeRef = defaultMode if isReference else byref(defaultMode) + fn = _nvmlGetFunctionPointer("nvmlDeviceGetGspFirmwareMode") + ret = fn(handle, isEnabledRef, defaultModeRef) + _nvmlCheckReturn(ret) + return NVML_SUCCESS if isReference else [isEnabled.value, defaultMode.value] + +def nvmlDeviceGetEncoderCapacity(handle, encoderQueryType): + c_encoder_capacity = c_ulonglong(0) + c_encoderQuery_type = _nvmlEncoderQueryType_t(encoderQueryType) + + fn = _nvmlGetFunctionPointer("nvmlDeviceGetEncoderCapacity") + ret = fn(handle, c_encoderQuery_type, byref(c_encoder_capacity)) + _nvmlCheckReturn(ret) + return c_encoder_capacity.value + +def nvmlDeviceGetVgpuProcessUtilization(handle, timeStamp): + # first call to get the size + c_vgpu_count = c_uint(0) + c_time_stamp = c_ulonglong(timeStamp) + + fn = _nvmlGetFunctionPointer("nvmlDeviceGetVgpuProcessUtilization") + ret = fn(handle, c_time_stamp, byref(c_vgpu_count), None) + + if (ret == NVML_SUCCESS): + # special case, no active vGPUs + return [] + elif (ret == NVML_ERROR_INSUFFICIENT_SIZE): + # typical case + sampleArray = c_vgpu_count.value * c_nvmlVgpuProcessUtilizationSample_t + c_samples = sampleArray() + + # make the call again + ret = fn(handle, c_time_stamp, byref(c_vgpu_count), c_samples) + _nvmlCheckReturn(ret) + + return c_samples[0:c_vgpu_count.value] + else: + # error case + raise NVMLError(ret) + +def nvmlDeviceGetVgpuProcessesUtilizationInfo(handle, timeStamp): + # first call to get the size + c_time_stamp = c_ulonglong(timeStamp) + c_vgpuProcUtilInfo = c_nvmlVgpuProcessesUtilizationInfo_v1_t(0) + c_vgpuProcUtilInfo.version = VgpuProcessesUtilizationInfo_v1 + c_vgpuProcUtilInfo.vgpuProcessCount = c_uint(0) + c_vgpuProcUtilInfo.lastSeenTimeStamp = c_time_stamp + + fn = _nvmlGetFunctionPointer("nvmlDeviceGetVgpuProcessesUtilizationInfo") + ret = fn(handle, byref(c_vgpuProcUtilInfo)) + + if (ret == NVML_SUCCESS): + # special case, no active vGPUs + return [] + elif (ret == NVML_ERROR_INSUFFICIENT_SIZE): + # typical case + sampleArray = c_vgpuProcUtilInfo.vgpuProcessCount * c_nvmlVgpuProcessUtilizationInfo_v1_t + c_samples = sampleArray() + c_vgpuProcUtilInfo.vgpuProcUtilArray = c_samples + + # make the call again + ret = fn(handle, byref(c_vgpuProcUtilInfo)) + _nvmlCheckReturn(ret) + + return c_samples[0:c_vgpuProcUtilInfo.vgpuProcessCount] + else: + # error case + raise NVMLError(ret) + +def nvmlDeviceGetEncoderStats(handle): + c_encoderCount = c_ulonglong(0) + c_encodeFps = c_ulonglong(0) + c_encoderLatency = c_ulonglong(0) + fn = _nvmlGetFunctionPointer("nvmlDeviceGetEncoderStats") + ret = fn(handle, byref(c_encoderCount), byref(c_encodeFps), byref(c_encoderLatency)) + _nvmlCheckReturn(ret) + return (c_encoderCount.value, c_encodeFps.value, c_encoderLatency.value) + +def nvmlDeviceGetEncoderSessions(handle): + # first call to get the size + c_session_count = c_uint(0) + + fn = _nvmlGetFunctionPointer("nvmlDeviceGetEncoderSessions") + ret = fn(handle, byref(c_session_count), None) + + if (ret == NVML_SUCCESS): + if (c_session_count.value != 0): + # typical case + session_array = c_nvmlEncoderSession_t * c_session_count.value + c_sessions = session_array() + + # make the call again + ret = fn(handle, byref(c_session_count), c_sessions) + _nvmlCheckReturn(ret) + sessions = [] + for i in range(c_session_count.value): + sessions.append(c_sessions[i]) + return sessions + else: + return [] # no active sessions + else: + # error case + raise NVMLError(ret) + +def nvmlDeviceGetFBCStats(handle): + c_fbcStats = c_nvmlFBCStats_t() + fn = _nvmlGetFunctionPointer("nvmlDeviceGetFBCStats") + ret = fn(handle, byref(c_fbcStats)) + _nvmlCheckReturn(ret) + return c_fbcStats + +def nvmlDeviceGetFBCSessions(handle): + # first call to get the size + c_session_count = c_uint(0) + + fn = _nvmlGetFunctionPointer("nvmlDeviceGetFBCSessions") + ret = fn(handle, byref(c_session_count), None) + + if (ret == NVML_SUCCESS): + if (c_session_count.value != 0): + # typical case + session_array = c_nvmlFBCSession_t * c_session_count.value + c_sessions = session_array() + + # make the call again + ret = fn(handle, byref(c_session_count), c_sessions) + _nvmlCheckReturn(ret) + sessions = [] + for i in range(c_session_count.value): + sessions.append(c_sessions[i]) + return sessions + else: + return [] # no active sessions + else: + # error case + raise NVMLError(ret) + +def nvmlVgpuInstanceGetEncoderStats(vgpuInstance): + c_encoderCount = c_ulonglong(0) + c_encodeFps = c_ulonglong(0) + c_encoderLatency = c_ulonglong(0) + fn = _nvmlGetFunctionPointer("nvmlVgpuInstanceGetEncoderStats") + ret = fn(vgpuInstance, byref(c_encoderCount), byref(c_encodeFps), byref(c_encoderLatency)) + _nvmlCheckReturn(ret) + return (c_encoderCount.value, c_encodeFps.value, c_encoderLatency.value) + +def nvmlVgpuInstanceGetEncoderSessions(vgpuInstance): + # first call to get the size + c_session_count = c_uint(0) + + fn = _nvmlGetFunctionPointer("nvmlVgpuInstanceGetEncoderSessions") + ret = fn(vgpuInstance, byref(c_session_count), None) + + if (ret == NVML_SUCCESS): + if (c_session_count.value != 0): + # typical case + session_array = c_nvmlEncoderSession_t * c_session_count.value + c_sessions = session_array() + + # make the call again + ret = fn(vgpuInstance, byref(c_session_count), c_sessions) + _nvmlCheckReturn(ret) + sessions = [] + for i in range(c_session_count.value): + sessions.append(c_sessions[i]) + return sessions + else: + return [] # no active sessions + else: + # error case + raise NVMLError(ret) + +def nvmlVgpuInstanceGetFBCStats(vgpuInstance): + c_fbcStats = c_nvmlFBCStats_t() + fn = _nvmlGetFunctionPointer("nvmlVgpuInstanceGetFBCStats") + ret = fn(vgpuInstance, byref(c_fbcStats)) + _nvmlCheckReturn(ret) + return c_fbcStats + +def nvmlVgpuInstanceGetFBCSessions(vgpuInstance): + # first call to get the size + c_session_count = c_uint(0) + + fn = _nvmlGetFunctionPointer("nvmlVgpuInstanceGetFBCSessions") + ret = fn(vgpuInstance, byref(c_session_count), None) + + if (ret == NVML_SUCCESS): + if (c_session_count.value != 0): + # typical case + session_array = c_nvmlFBCSession_t * c_session_count.value + c_sessions = session_array() + + # make the call again + ret = fn(vgpuInstance, byref(c_session_count), c_sessions) + _nvmlCheckReturn(ret) + sessions = [] + for i in range(c_session_count.value): + sessions.append(c_sessions[i]) + return sessions + else: + return [] # no active sessions + else: + # error case + raise NVMLError(ret) + +def nvmlDeviceGetProcessUtilization(handle, timeStamp): + # first call to get the size + c_count = c_uint(0) + c_time_stamp = c_ulonglong(timeStamp) + + fn = _nvmlGetFunctionPointer("nvmlDeviceGetProcessUtilization") + ret = fn(handle, None, byref(c_count), c_time_stamp) + + if (ret == NVML_ERROR_INSUFFICIENT_SIZE): + # typical case + sampleArray = c_count.value * c_nvmlProcessUtilizationSample_t + c_samples = sampleArray() + + # make the call again + ret = fn(handle, c_samples, byref(c_count), c_time_stamp) + _nvmlCheckReturn(ret) + + return c_samples[0:c_count.value] + else: + # error case + raise NVMLError(ret) + +def nvmlDeviceGetProcessesUtilizationInfo(handle, timeStamp): + # first call to get the size + c_time_stamp = c_ulonglong(timeStamp) + c_processesUtilInfo = c_nvmlProcessesUtilizationInfo_v1_t(0) + c_processesUtilInfo.version = ProcessesUtilizationInfo_v1 + c_processesUtilInfo.processSamplesCount = c_uint(0) + c_processesUtilInfo.lastSeenTimeStamp = c_time_stamp + + fn = _nvmlGetFunctionPointer("nvmlDeviceGetProcessesUtilizationInfo") + ret = fn(handle, byref(c_processesUtilInfo)) + + if (ret == NVML_ERROR_INSUFFICIENT_SIZE): + # typical case + sampleArray = c_processesUtilInfo.processSamplesCount * c_nvmlProcessUtilizationInfo_v1_t + c_samples = sampleArray() + c_processesUtilInfo.procUtilArray = c_samples + + # make the call again + ret = fn(handle, byref(c_processesUtilInfo)) + _nvmlCheckReturn(ret) + + return c_samples[0:c_processesUtilInfo.processSamplesCount] + else: + # error case + raise NVMLError(ret) + +def nvmlVgpuInstanceGetMetadata(vgpuInstance): + fn = _nvmlGetFunctionPointer("nvmlVgpuInstanceGetMetadata") + c_vgpuMetadata = c_nvmlVgpuMetadata_t() + c_bufferSize = c_uint(0) + # Make the first NVML API call to get the c_bufferSize value. + # We have already allocated required buffer above. + ret = fn(vgpuInstance, byref(c_vgpuMetadata), byref(c_bufferSize)) + if (ret == NVML_ERROR_INSUFFICIENT_SIZE): + ret = fn(vgpuInstance, byref(c_vgpuMetadata), byref(c_bufferSize)) + _nvmlCheckReturn(ret) + else: + raise NVMLError(ret) + return c_vgpuMetadata + +def nvmlDeviceGetVgpuMetadata(handle): + fn = _nvmlGetFunctionPointer("nvmlDeviceGetVgpuMetadata") + c_vgpuPgpuMetadata = c_nvmlVgpuPgpuMetadata_t() + c_bufferSize = c_uint(0) + # Make the first NVML API call to get the c_bufferSize value. + # We have already allocated required buffer above. + ret = fn(handle, byref(c_vgpuPgpuMetadata), byref(c_bufferSize)) + if (ret == NVML_ERROR_INSUFFICIENT_SIZE): + ret = fn(handle, byref(c_vgpuPgpuMetadata), byref(c_bufferSize)) + _nvmlCheckReturn(ret) + else: + raise NVMLError(ret) + return c_vgpuPgpuMetadata + +def nvmlGetVgpuCompatibility(vgpuMetadata, pgpuMetadata): + fn = _nvmlGetFunctionPointer("nvmlGetVgpuCompatibility") + c_vgpuPgpuCompatibility = c_nvmlVgpuPgpuCompatibility_t() + ret = fn(byref(vgpuMetadata), byref(pgpuMetadata), byref(c_vgpuPgpuCompatibility)) + _nvmlCheckReturn(ret) + return c_vgpuPgpuCompatibility + +@convertStrBytes +def nvmlDeviceGetPgpuMetadataString(handle): + fn = _nvmlGetFunctionPointer("nvmlDeviceGetPgpuMetadataString") + c_pgpuMetadata = create_string_buffer(NVML_VGPU_PGPU_METADATA_OPAQUE_DATA_SIZE) + c_bufferSize = c_uint(0) + # Make the first NVML API call to get the c_bufferSize value. + # We have already allocated required buffer above. + ret = fn(handle, byref(c_pgpuMetadata), byref(c_bufferSize)) + if (ret == NVML_ERROR_INSUFFICIENT_SIZE): + ret = fn(handle, byref(c_pgpuMetadata), byref(c_bufferSize)) + _nvmlCheckReturn(ret) + else: + raise NVMLError(ret) + return (c_pgpuMetadata.value, c_bufferSize.value) + +def nvmlDeviceGetVgpuSchedulerLog(handle): + c_vgpu_sched_log = c_nvmlVgpuSchedulerLog_t() + fn = _nvmlGetFunctionPointer("nvmlDeviceGetVgpuSchedulerLog") + ret = fn(handle, byref(c_vgpu_sched_log)) + _nvmlCheckReturn(ret) + return c_vgpu_sched_log + +def nvmlDeviceGetVgpuSchedulerState(handle): + c_vgpu_sched_state = c_nvmlVgpuSchedulerGetState_t() + fn = _nvmlGetFunctionPointer("nvmlDeviceGetVgpuSchedulerState") + ret = fn(handle, byref(c_vgpu_sched_state)) + _nvmlCheckReturn(ret) + return c_vgpu_sched_state + +def nvmlDeviceGetVgpuSchedulerCapabilities(handle): + c_vgpu_sched_caps = c_nvmlVgpuSchedulerCapabilities_t() + fn = _nvmlGetFunctionPointer("nvmlDeviceGetVgpuSchedulerCapabilities") + ret = fn(handle, byref(c_vgpu_sched_caps)) + _nvmlCheckReturn(ret) + return c_vgpu_sched_caps + +def nvmlDeviceSetVgpuSchedulerState(handle, sched_state): + fn = _nvmlGetFunctionPointer("nvmlDeviceSetVgpuSchedulerState") + ret = fn(handle, byref(sched_state)) + _nvmlCheckReturn(ret) + return NVML_SUCCESS + +def nvmlSetVgpuVersion(vgpuVersion): + fn = _nvmlGetFunctionPointer("nvmlSetVgpuVersion") + ret = fn(byref(vgpuVersion)) + _nvmlCheckReturn(ret) + return NVML_SUCCESS + +def nvmlGetVgpuVersion(supported=None, current=None): + isUserDefined = (supported is not None) or (current is not None) + if not isUserDefined: + supported = c_nvmlVgpuVersion_t() + current = c_nvmlVgpuVersion_t() + fn = _nvmlGetFunctionPointer("nvmlGetVgpuVersion") + ret = fn(byref(supported), byref(current)) + _nvmlCheckReturn(ret) + return NVML_SUCCESS if isUserDefined else [(supported.minVersion, + supported.maxVersion), + (current.minVersion, + current.maxVersion)] + +def nvmlVgpuInstanceGetAccountingMode(vgpuInstance): + c_mode = _nvmlEnableState_t() + fn = _nvmlGetFunctionPointer("nvmlVgpuInstanceGetAccountingMode") + ret = fn(vgpuInstance, byref(c_mode)) + _nvmlCheckReturn(ret) + return c_mode.value + +def nvmlVgpuInstanceGetAccountingPids(vgpuInstance): + c_pidCount = c_uint() + fn = _nvmlGetFunctionPointer("nvmlVgpuInstanceGetAccountingPids") + ret = fn(vgpuInstance, byref(c_pidCount), None) + if (ret == NVML_ERROR_INSUFFICIENT_SIZE): + sampleArray = c_pidCount.value * c_uint + c_pidArray = sampleArray() + ret = fn(vgpuInstance, byref(c_pidCount), byref(c_pidArray)) + _nvmlCheckReturn(ret) + else: + raise NVMLError(ret) + return (c_pidCount, c_pidArray) + +def nvmlVgpuInstanceGetAccountingStats(vgpuInstance, pid): + c_accountingStats = c_nvmlAccountingStats_t() + fn = _nvmlGetFunctionPointer("nvmlVgpuInstanceGetAccountingStats") + ret = fn(vgpuInstance, pid, byref(c_accountingStats)) + _nvmlCheckReturn(ret) + return c_accountingStats + +def nvmlVgpuInstanceClearAccountingPids(vgpuInstance): + fn = _nvmlGetFunctionPointer("nvmlVgpuInstanceClearAccountingPids") + ret = fn(vgpuInstance) + _nvmlCheckReturn(ret) + return NVML_SUCCESS + +def nvmlGpuInstanceGetCreatableVgpus(gpuInstance, c_vgpus): + c_vgpus.vgpuCount = 0; + fn = _nvmlGetFunctionPointer("nvmlGpuInstanceGetCreatableVgpus") + ret = fn(gpuInstance, byref(c_vgpus)) + + if (ret == NVML_SUCCESS): + # special case, no supported vGPUs + return c_vgpus + elif (ret == NVML_ERROR_INSUFFICIENT_SIZE): + vgpu_type_ids_array = _nvmlVgpuTypeId_t * c_vgpus.vgpuCount + c_vgpus.vgpuTypeIds = vgpu_type_ids_array() + ret = fn(gpuInstance, byref(c_vgpus)) + _nvmlCheckReturn(ret) + + return c_vgpus + else: + raise NVMLError(ret) + +def nvmlVgpuTypeGetMaxInstancesPerGpuInstance(vgpuTypeId): + c_max_instances = c_nvmlVgpuTypeMaxInstance_v1_t() + c_max_instances.version = nvmlVgpuTypeMaxInstance_v1 + c_max_instances.vgpuTypeId = vgpuTypeId + fn = _nvmlGetFunctionPointer("nvmlVgpuTypeGetMaxInstancesPerGpuInstance") + ret = fn(byref(c_max_instances)) + _nvmlCheckReturn(ret) + return c_max_instances.maxInstancePerGI + +def nvmlGpuInstanceGetActiveVgpus(gpuInstance, c_vgpu_instance_info): + c_vgpu_instance_info.vgpuCount = c_uint(0) + + fn = _nvmlGetFunctionPointer("nvmlGpuInstanceGetActiveVgpus") + ret = fn(gpuInstance, byref(c_vgpu_instance_info)) + + if (ret == NVML_SUCCESS): + # special case, no active vGPUs + return c_vgpu_instance_info + elif (ret == NVML_ERROR_INSUFFICIENT_SIZE): + # typical case + vgpu_instance_array = _nvmlVgpuInstance_t * c_vgpu_instance_info.vgpuCount + c_vgpu_instance_info.vgpuInstances = vgpu_instance_array() + + # make the call again + ret = fn(gpuInstance, byref(c_vgpu_instance_info)) + _nvmlCheckReturn(ret) + return c_vgpu_instance_info + else: + raise NVMLError(ret) + +def nvmlGpuInstanceSetVgpuSchedulerState(gpuInstance, sched_state): + fn = _nvmlGetFunctionPointer("nvmlGpuInstanceSetVgpuSchedulerState") + ret = fn(gpuInstance, byref(sched_state)) + _nvmlCheckReturn(ret) + return ret + +def nvmlGpuInstanceGetVgpuSchedulerState(gpuInstance, c_vgpu_sched_state_info): + fn = _nvmlGetFunctionPointer("nvmlGpuInstanceGetVgpuSchedulerState") + ret = fn(gpuInstance, byref(c_vgpu_sched_state_info)) + _nvmlCheckReturn(ret) + return c_vgpu_sched_state_info + +def nvmlGpuInstanceGetVgpuSchedulerLog(gpuInstance, c_vgpu_sched_log_info): + fn = _nvmlGetFunctionPointer("nvmlGpuInstanceGetVgpuSchedulerLog") + ret = fn(gpuInstance, byref(c_vgpu_sched_log_info)) + _nvmlCheckReturn(ret) + return c_vgpu_sched_log_info + +def nvmlGpuInstanceGetVgpuTypeCreatablePlacements(gpuInstance, c_vgpu_placements): + c_max_instances = c_nvmlVgpuTypeMaxInstance_v1_t() + c_max_instances.version = nvmlVgpuTypeMaxInstance_v1 + c_max_instances.vgpuTypeId = c_vgpu_placements.vgpuTypeId + fn = _nvmlGetFunctionPointer("nvmlVgpuTypeGetMaxInstancesPerGpuInstance") + ret = fn(byref(c_max_instances)) + _nvmlCheckReturn(ret) + + c_vgpu_placements.count = c_max_instances.maxInstancePerGI + + c_placements = c_uint * c_max_instances.maxInstancePerGI + c_vgpu_placements.placementIds = c_placements() + fn = _nvmlGetFunctionPointer("nvmlGpuInstanceGetVgpuTypeCreatablePlacements") + ret = fn(gpuInstance, byref(c_vgpu_placements)) + _nvmlCheckReturn(ret) + return c_vgpu_placements + +def nvmlGpuInstanceGetVgpuHeterogeneousMode(gpuInstance): + c_vgpuHeterogeneousMode = c_nvmlVgpuHeterogeneousMode_v1_t(0) + c_vgpuHeterogeneousMode.version = VgpuHeterogeneousMode_v1 + fn = _nvmlGetFunctionPointer("nvmlGpuInstanceGetVgpuHeterogeneousMode") + ret = fn(gpuInstance, byref(c_vgpuHeterogeneousMode)) + _nvmlCheckReturn(ret) + return c_vgpuHeterogeneousMode.mode + +def nvmlGpuInstanceSetVgpuHeterogeneousMode(gpuInstance, heterogeneous_mode): + c_vgpuHeterogeneousMode = c_nvmlVgpuHeterogeneousMode_v1_t(0) + c_vgpuHeterogeneousMode.version = VgpuHeterogeneousMode_v1 + c_vgpuHeterogeneousMode.mode = heterogeneous_mode + fn = _nvmlGetFunctionPointer("nvmlGpuInstanceSetVgpuHeterogeneousMode") + ret = fn(gpuInstance, byref(c_vgpuHeterogeneousMode)) + _nvmlCheckReturn(ret) + return NVML_SUCCESS + +def nvmlGetExcludedDeviceCount(): + c_count = c_uint() + fn = _nvmlGetFunctionPointer("nvmlGetExcludedDeviceCount") + ret = fn(byref(c_count)) + _nvmlCheckReturn(ret) + return c_count.value + +def nvmlGetExcludedDeviceInfoByIndex(index): + c_index = c_uint(index) + info = c_nvmlExcludedDeviceInfo_t() + fn = _nvmlGetFunctionPointer("nvmlGetExcludedDeviceInfoByIndex") + ret = fn(c_index, byref(info)) + _nvmlCheckReturn(ret) + return info + +def nvmlDeviceGetHostVgpuMode(handle): + c_host_vgpu_mode = _nvmlHostVgpuMode_t() + fn = _nvmlGetFunctionPointer("nvmlDeviceGetHostVgpuMode") + ret = fn(handle, byref(c_host_vgpu_mode)) + _nvmlCheckReturn(ret) + return c_host_vgpu_mode.value + +def nvmlDeviceSetMigMode(device, mode): + c_activationStatus = c_uint() + fn = _nvmlGetFunctionPointer("nvmlDeviceSetMigMode") + ret = fn(device, mode, byref(c_activationStatus)) + _nvmlCheckReturn(ret) + return c_activationStatus.value + +def nvmlDeviceGetMigMode(device): + c_currentMode = c_uint() + c_pendingMode = c_uint() + fn = _nvmlGetFunctionPointer("nvmlDeviceGetMigMode") + ret = fn(device, byref(c_currentMode), byref(c_pendingMode)) + _nvmlCheckReturn(ret) + return [c_currentMode.value, c_pendingMode.value] + +def nvmlDeviceGetGpuInstanceProfileInfo(device, profile, version=2): + if version == 2: + c_info = c_nvmlGpuInstanceProfileInfo_v2_t() + fn = _nvmlGetFunctionPointer("nvmlDeviceGetGpuInstanceProfileInfoV") + elif version == 1: + c_info = c_nvmlGpuInstanceProfileInfo_t() + fn = _nvmlGetFunctionPointer("nvmlDeviceGetGpuInstanceProfileInfo") + else: + raise NVMLError(NVML_ERROR_FUNCTION_NOT_FOUND) + ret = fn(device, profile, byref(c_info)) + _nvmlCheckReturn(ret) + return c_info + +def nvmlDeviceGetGpuInstanceProfileInfoById(device, profileId): + c_info = c_nvmlGpuInstanceProfileInfo_v2_t() + fn = _nvmlGetFunctionPointer("nvmlDeviceGetGpuInstanceProfileInfoByIdV") + + ret = fn(device, profileId, byref(c_info)) + _nvmlCheckReturn(ret) + return c_info + +# Define function alias for the API exposed by NVML +nvmlDeviceGetGpuInstanceProfileInfoV = nvmlDeviceGetGpuInstanceProfileInfo +nvmlDeviceGetGpuInstanceProfileInfoByIdV = nvmlDeviceGetGpuInstanceProfileInfoById + +def nvmlDeviceGetGpuInstanceRemainingCapacity(device, profileId): + c_count = c_uint() + fn = _nvmlGetFunctionPointer("nvmlDeviceGetGpuInstanceRemainingCapacity") + ret = fn(device, profileId, byref(c_count)) + _nvmlCheckReturn(ret) + return c_count.value + +def nvmlDeviceGetGpuInstancePossiblePlacements(device, profileId, placementsRef, countRef): + fn = _nvmlGetFunctionPointer("nvmlDeviceGetGpuInstancePossiblePlacements_v2") + ret = fn(device, profileId, placementsRef, countRef) + _nvmlCheckReturn(ret) + return NVML_SUCCESS + +def nvmlDeviceCreateGpuInstance(device, profileId): + c_instance = c_nvmlGpuInstance_t() + fn = _nvmlGetFunctionPointer("nvmlDeviceCreateGpuInstance") + ret = fn(device, profileId, byref(c_instance)) + _nvmlCheckReturn(ret) + return c_instance + +def nvmlDeviceCreateGpuInstanceWithPlacement(device, profileId, placement): + c_instance = c_nvmlGpuInstance_t() + fn = _nvmlGetFunctionPointer("nvmlDeviceCreateGpuInstanceWithPlacement") + ret = fn(device, profileId, placement, byref(c_instance)) + _nvmlCheckReturn(ret) + return c_instance + +def nvmlGpuInstanceDestroy(gpuInstance): + fn = _nvmlGetFunctionPointer("nvmlGpuInstanceDestroy") + ret = fn(gpuInstance) + _nvmlCheckReturn(ret) + return NVML_SUCCESS + +def nvmlDeviceGetGpuInstances(device, profileId, gpuInstancesRef, countRef): + fn = _nvmlGetFunctionPointer("nvmlDeviceGetGpuInstances") + ret = fn(device, profileId, gpuInstancesRef, countRef) + _nvmlCheckReturn(ret) + return NVML_SUCCESS + +def nvmlDeviceGetGpuInstanceById(device, gpuInstanceId): + c_instance = c_nvmlGpuInstance_t() + fn = _nvmlGetFunctionPointer("nvmlDeviceGetGpuInstanceById") + ret = fn(device, gpuInstanceId, byref(c_instance)) + _nvmlCheckReturn(ret) + return c_instance + +def nvmlGpuInstanceGetInfo(gpuInstance): + c_info = c_nvmlGpuInstanceInfo_t() + fn = _nvmlGetFunctionPointer("nvmlGpuInstanceGetInfo") + ret = fn(gpuInstance, byref(c_info)) + _nvmlCheckReturn(ret) + return c_info + +def nvmlGpuInstanceGetComputeInstanceProfileInfo(device, profile, engProfile, version=2): + if version == 2: + c_info = c_nvmlComputeInstanceProfileInfo_v2_t() + fn = _nvmlGetFunctionPointer("nvmlGpuInstanceGetComputeInstanceProfileInfoV") + elif version == 1: + c_info = c_nvmlComputeInstanceProfileInfo_t() + fn = _nvmlGetFunctionPointer("nvmlGpuInstanceGetComputeInstanceProfileInfo") + else: + raise NVMLError(NVML_ERROR_FUNCTION_NOT_FOUND) + ret = fn(device, profile, engProfile, byref(c_info)) + _nvmlCheckReturn(ret) + return c_info + +# Define function alias for the API exposed by NVML +nvmlGpuInstanceGetComputeInstanceProfileInfoV = nvmlGpuInstanceGetComputeInstanceProfileInfo + +def nvmlGpuInstanceGetComputeInstanceRemainingCapacity(gpuInstance, profileId): + c_count = c_uint() + fn = _nvmlGetFunctionPointer("nvmlGpuInstanceGetComputeInstanceRemainingCapacity") + ret = fn(gpuInstance, profileId, byref(c_count)) + _nvmlCheckReturn(ret) + return c_count.value + +def nvmlGpuInstanceGetComputeInstancePossiblePlacements(gpuInstance, profileId, placementsRef, countRef): + fn = _nvmlGetFunctionPointer("nvmlGpuInstanceGetComputeInstancePossiblePlacements") + ret = fn(gpuInstance, profileId, placementsRef, countRef) + _nvmlCheckReturn(ret) + return NVML_SUCCESS + +def nvmlGpuInstanceCreateComputeInstance(gpuInstance, profileId): + c_instance = c_nvmlComputeInstance_t() + fn = _nvmlGetFunctionPointer("nvmlGpuInstanceCreateComputeInstance") + ret = fn(gpuInstance, profileId, byref(c_instance)) + _nvmlCheckReturn(ret) + return c_instance + +def nvmlGpuInstanceCreateComputeInstanceWithPlacement(gpuInstance, profileId, placement): + c_instance = c_nvmlComputeInstance_t() + fn = _nvmlGetFunctionPointer("nvmlGpuInstanceCreateComputeInstanceWithPlacement") + ret = fn(gpuInstance, profileId, placement, byref(c_instance)) + _nvmlCheckReturn(ret) + return c_instance + +def nvmlComputeInstanceDestroy(computeInstance): + fn = _nvmlGetFunctionPointer("nvmlComputeInstanceDestroy") + ret = fn(computeInstance) + _nvmlCheckReturn(ret) + return NVML_SUCCESS + +def nvmlGpuInstanceGetComputeInstances(gpuInstance, profileId, computeInstancesRef, countRef): + fn = _nvmlGetFunctionPointer("nvmlGpuInstanceGetComputeInstances") + ret = fn(gpuInstance, profileId, computeInstancesRef, countRef) + _nvmlCheckReturn(ret) + return NVML_SUCCESS + +def nvmlGpuInstanceGetComputeInstanceById(gpuInstance, computeInstanceId): + c_instance = c_nvmlComputeInstance_t() + fn = _nvmlGetFunctionPointer("nvmlGpuInstanceGetComputeInstanceById") + ret = fn(gpuInstance, computeInstanceId, byref(c_instance)) + _nvmlCheckReturn(ret) + return c_instance + +def nvmlComputeInstanceGetInfo_v2(computeInstance): + c_info = c_nvmlComputeInstanceInfo_t() + fn = _nvmlGetFunctionPointer("nvmlComputeInstanceGetInfo_v2") + ret = fn(computeInstance, byref(c_info)) + _nvmlCheckReturn(ret) + return c_info + +def nvmlComputeInstanceGetInfo(computeInstance): + return nvmlComputeInstanceGetInfo_v2(computeInstance) + +def nvmlDeviceIsMigDeviceHandle(device): + c_isMigDevice = c_uint() + fn = _nvmlGetFunctionPointer("nvmlDeviceIsMigDeviceHandle") + ret = fn(device, byref(c_isMigDevice)) + _nvmlCheckReturn(ret) + return c_isMigDevice + +def nvmlDeviceGetGpuInstanceId(device): + c_gpuInstanceId = c_uint() + fn = _nvmlGetFunctionPointer("nvmlDeviceGetGpuInstanceId") + ret = fn(device, byref(c_gpuInstanceId)) + _nvmlCheckReturn(ret) + return c_gpuInstanceId.value + +def nvmlDeviceGetComputeInstanceId(device): + c_computeInstanceId = c_uint() + fn = _nvmlGetFunctionPointer("nvmlDeviceGetComputeInstanceId") + ret = fn(device, byref(c_computeInstanceId)) + _nvmlCheckReturn(ret) + return c_computeInstanceId.value + +def nvmlDeviceGetMaxMigDeviceCount(device): + c_count = c_uint() + fn = _nvmlGetFunctionPointer("nvmlDeviceGetMaxMigDeviceCount") + ret = fn(device, byref(c_count)) + _nvmlCheckReturn(ret) + return c_count.value + +def nvmlDeviceGetMigDeviceHandleByIndex(device, index): + c_index = c_uint(index) + migDevice = c_nvmlDevice_t() + fn = _nvmlGetFunctionPointer("nvmlDeviceGetMigDeviceHandleByIndex") + ret = fn(device, c_index, byref(migDevice)) + _nvmlCheckReturn(ret) + return migDevice + +def nvmlDeviceGetDeviceHandleFromMigDeviceHandle(migDevice): + device = c_nvmlDevice_t() + fn = _nvmlGetFunctionPointer("nvmlDeviceGetDeviceHandleFromMigDeviceHandle") + ret = fn(migDevice, byref(device)) + _nvmlCheckReturn(ret) + return device + +def nvmlDeviceGetAttributes_v2(device): + c_attrs = c_nvmlDeviceAttributes() + fn = _nvmlGetFunctionPointer("nvmlDeviceGetAttributes_v2") + ret = fn(device, byref(c_attrs)) + _nvmlCheckReturn(ret) + return c_attrs + +def nvmlDeviceGetAttributes(device): + return nvmlDeviceGetAttributes_v2(device) + +def nvmlDeviceGetRemappedRows(device): + fn = _nvmlGetFunctionPointer("nvmlDeviceGetRemappedRows") + c_corr = c_uint() + c_unc = c_uint() + c_bpending = c_uint() + c_bfailure = c_uint() + ret = fn(device, byref(c_corr), byref(c_unc), byref(c_bpending), byref(c_bfailure)) + _nvmlCheckReturn(ret) + return (c_corr.value, c_unc.value, c_bpending.value, c_bfailure.value) + +def nvmlDeviceGetRowRemapperHistogram(device): + c_vals = c_nvmlRowRemapperHistogramValues() + fn = _nvmlGetFunctionPointer("nvmlDeviceGetRowRemapperHistogram") + ret = fn(device, byref(c_vals)) + _nvmlCheckReturn(ret) + return c_vals + +def nvmlDeviceGetArchitecture(device): + arch = _nvmlDeviceArchitecture_t() + fn = _nvmlGetFunctionPointer("nvmlDeviceGetArchitecture") + ret = fn(device, byref(arch)) + _nvmlCheckReturn(ret) + return arch.value + +def nvmlDeviceGetBusType(device): + c_busType = _nvmlBusType_t() + fn = _nvmlGetFunctionPointer("nvmlDeviceGetBusType") + ret = fn(device, byref(c_busType)) + _nvmlCheckReturn(ret) + return c_busType.value + +def nvmlDeviceGetIrqNum(device): + c_irqNum = c_uint() + fn = _nvmlGetFunctionPointer("nvmlDeviceGetIrqNum") + ret = fn(device, byref(c_irqNum)) + _nvmlCheckReturn(ret) + return c_irqNum.value + +def nvmlDeviceGetNumGpuCores(device): + c_numCores = c_uint() + fn = _nvmlGetFunctionPointer("nvmlDeviceGetNumGpuCores") + ret = fn(device, byref(c_numCores)) + _nvmlCheckReturn(ret) + return c_numCores.value + +def nvmlDeviceGetPowerSource(device): + c_powerSource = _nvmlPowerSource_t() + fn = _nvmlGetFunctionPointer("nvmlDeviceGetPowerSource") + ret = fn(device, byref(c_powerSource)) + _nvmlCheckReturn(ret) + return c_powerSource.value + +def nvmlDeviceGetMemoryBusWidth(device): + c_memBusWidth = c_uint() + fn = _nvmlGetFunctionPointer("nvmlDeviceGetMemoryBusWidth") + ret = fn(device, byref(c_memBusWidth)) + _nvmlCheckReturn(ret) + return c_memBusWidth.value + +def nvmlDeviceGetPcieLinkMaxSpeed(device): + c_speed = _nvmlPcieLinkMaxSpeed_t() + fn = _nvmlGetFunctionPointer("nvmlDeviceGetPcieLinkMaxSpeed") + ret = fn(device, byref(c_speed)) + _nvmlCheckReturn(ret) + return c_speed.value + +def nvmlDeviceGetAdaptiveClockInfoStatus(device): + c_adaptiveClockInfoStatus = _nvmlAdaptiveClockInfoStatus_t() + fn = _nvmlGetFunctionPointer("nvmlDeviceGetAdaptiveClockInfoStatus") + ret = fn(device, byref(c_adaptiveClockInfoStatus)) + _nvmlCheckReturn(ret) + return c_adaptiveClockInfoStatus.value + +def nvmlDeviceGetPcieSpeed(device): + c_speed = c_uint() + fn = _nvmlGetFunctionPointer("nvmlDeviceGetPcieSpeed") + ret = fn(device, byref(c_speed)) + _nvmlCheckReturn(ret) + return c_speed.value + +def nvmlDeviceGetDynamicPstatesInfo(device, c_dynamicpstatesinfo=c_nvmlGpuDynamicPstatesInfo_t()): + isReference = type(c_dynamicpstatesinfo) is not c_nvmlGpuDynamicPstatesInfo_t + dynamicpstatesinfoRef = c_dynamicpstatesinfo if isReference else byref(c_dynamicpstatesinfo) + + fn = _nvmlGetFunctionPointer("nvmlDeviceGetDynamicPstatesInfo"); + ret = fn(device, dynamicpstatesinfoRef) + _nvmlCheckReturn(ret) + return NVML_SUCCESS if isReference else c_dynamicpstatesinfo + +def nvmlDeviceSetFanSpeed_v2(handle, index, speed): + fn = _nvmlGetFunctionPointer("nvmlDeviceSetFanSpeed_v2"); + ret = fn(handle, index, speed) + _nvmlCheckReturn(ret) + return NVML_SUCCESS + +def nvmlDeviceGetThermalSettings(device, sensorindex, c_thermalsettings=c_nvmlGpuThermalSettings_t()): + isReference = type(c_thermalsettings) is not c_nvmlGpuThermalSettings_t + thermalsettingsRef = c_thermalsettings if isReference else byref(c_thermalsettings) + fn = _nvmlGetFunctionPointer("nvmlDeviceGetThermalSettings"); + ret = fn(device, sensorindex, thermalsettingsRef) + _nvmlCheckReturn(ret) + return NVML_SUCCESS if isReference else c_thermalsettings.sensor[:] + +def nvmlDeviceGetMinMaxClockOfPState(device, clockType, pstate, minClockMHz=c_uint(), maxClockMHz=c_uint()): + isReference = (type(minClockMHz) is not c_uint) or (type(maxClockMHz) is not c_uint) + minClockMHzRef = minClockMHz if isReference else byref(minClockMHz) + maxClockMHzRef = maxClockMHz if isReference else byref(maxClockMHz) + fn = _nvmlGetFunctionPointer("nvmlDeviceGetMinMaxClockOfPState"); + ret = fn(device, _nvmlClockType_t(clockType), _nvmlClockType_t(pstate), minClockMHzRef, maxClockMHzRef) + _nvmlCheckReturn(ret) + return NVML_SUCCESS if isReference else (minClockMHz.value, maxClockMHz.value) + +_nvmlPowerMizerMode_t = c_uint +NVML_POWER_MIZER_MODE_ADAPTIVE = 0 +NVML_POWER_MIZER_MODE_PREFER_MAXIMUM_PERFORMANCE = 1 +NVML_POWER_MIZER_MODE_AUTO = 2 +NVML_POWER_MIZER_MODE_PREFER_CONSISTENT_PERFORMANCE = 3 + +class c_nvmlDevicePowerMizerModes_v1_t(_PrintableStructure): + _fields_ = [ + ('currentMode', _nvmlPowerMizerMode_t), + ('mode', _nvmlPowerMizerMode_t), + ('supportedPowerMizerModes', _nvmlPowerMizerMode_t), + ] + +def nvmlDeviceGetPowerMizerMode_v1(device, info): + fn = _nvmlGetFunctionPointer("nvmlDeviceGetPowerMizerMode_v1"); + ret = fn(device, info) + return ret + +def nvmlDeviceSetPowerMizerMode_v1(device, info): + fn = _nvmlGetFunctionPointer("nvmlDeviceSetPowerMizerMode_v1"); + ret = fn(device, info) + return ret + +class c_nvmlClockOffset_t(_PrintableStructure): + _fields_ = [ + ('version', c_uint), + ('type', _nvmlClockType_t), + ('pstate', _nvmlPstates_t), + ('clockOffsetMHz', c_int), + ('minClockOffsetMHz', c_int), + ('maxClockOffsetMHz', c_int), + ] + +nvmlClockOffset_v1 = 0x1000018 + +def nvmlDeviceGetClockOffsets(device, info): + fn = _nvmlGetFunctionPointer("nvmlDeviceGetClockOffsets"); + ret = fn(device, info) + return NVML_SUCCESS + +def nvmlDeviceSetClockOffsets(device, info): + fn = _nvmlGetFunctionPointer("nvmlDeviceSetClockOffsets"); + ret = fn(device, info) + return NVML_SUCCESS + +def nvmlDeviceGetSupportedPerformanceStates(device): + pstates = [] + c_count = c_uint(NVML_MAX_GPU_PERF_PSTATES) + c_size = sizeof(c_uint)*c_count.value + + # NOTE: use 'c_uint' to represent the size of the nvmlPstate_t enumeration. + pstates_array = _nvmlPstates_t * c_count.value + c_pstates = pstates_array() + + fn = _nvmlGetFunctionPointer("nvmlDeviceGetSupportedPerformanceStates") + ret = fn(device, c_pstates, c_size) + _nvmlCheckReturn(ret) + + for value in c_pstates: + if value != NVML_PSTATE_UNKNOWN: + pstates.append(value) + + return pstates + +def nvmlDeviceGetGpcClkVfOffset(device): + offset = c_int32() + fn = _nvmlGetFunctionPointer("nvmlDeviceGetGpcClkVfOffset") + ret = fn(device, byref(offset)) + _nvmlCheckReturn(ret) + return offset.value + +# Deprecated +def nvmlDeviceSetGpcClkVfOffset(device, offset): + c_offset = c_int32(offset) + fn = _nvmlGetFunctionPointer("nvmlDeviceSetGpcClkVfOffset") + ret = fn(device, c_offset) + _nvmlCheckReturn(ret) + return NVML_SUCCESS + +def nvmlDeviceGetGpcClkMinMaxVfOffset(device, minOffset=c_int(), maxOffset=c_int()): + isReference = (type(minOffset) is not c_int) or (type(maxOffset) is not c_int) + minOffsetRef = minOffset if isReference else byref(minOffset) + maxOffsetRef = maxOffset if isReference else byref(maxOffset) + fn = _nvmlGetFunctionPointer("nvmlDeviceGetGpcClkMinMaxVfOffset") + ret = fn(device, minOffsetRef, maxOffsetRef) + _nvmlCheckReturn(ret) + return NVML_SUCCESS if isReference else (minOffset.value, maxOffset.value) + +def nvmlDeviceGetMemClkVfOffset(device): + offset = c_int32() + fn = _nvmlGetFunctionPointer("nvmlDeviceGetMemClkVfOffset") + ret = fn(device, byref(offset)) + _nvmlCheckReturn(ret) + return offset.value + +# Deprecated +def nvmlDeviceSetMemClkVfOffset(device, offset): + c_offset = c_int32(offset) + fn = _nvmlGetFunctionPointer("nvmlDeviceSetMemClkVfOffset") + ret = fn(device, c_offset) + _nvmlCheckReturn(ret) + return NVML_SUCCESS + +def nvmlDeviceGetMemClkMinMaxVfOffset(device, minOffset=c_int(), maxOffset=c_int()): + isReference = (type(minOffset) is not c_int) or (type(maxOffset) is not c_int) + minOffsetRef = minOffset if isReference else byref(minOffset) + maxOffsetRef = maxOffset if isReference else byref(maxOffset) + + fn = _nvmlGetFunctionPointer("nvmlDeviceGetMemClkMinMaxVfOffset") + ret = fn(device, minOffsetRef, maxOffsetRef) + _nvmlCheckReturn(ret) + return NVML_SUCCESS if isReference else (minOffset.value, maxOffset.value) + +def nvmlSystemSetConfComputeGpusReadyState(state): + c_state = c_uint(state) + fn = _nvmlGetFunctionPointer("nvmlSystemSetConfComputeGpusReadyState") + ret = fn(c_state) + _nvmlCheckReturn(ret) + return NVML_SUCCESS + +def nvmlSystemGetConfComputeGpusReadyState(): + c_state = c_uint() + fn = _nvmlGetFunctionPointer("nvmlSystemGetConfComputeGpusReadyState") + ret = fn(byref(c_state)) + _nvmlCheckReturn(ret) + return c_state.value + +def nvmlSystemGetConfComputeCapabilities(): + c_ccSysCaps = c_nvmlConfComputeSystemCaps_t() + fn = _nvmlGetFunctionPointer("nvmlSystemGetConfComputeCapabilities") + ret = fn(byref(c_ccSysCaps)) + _nvmlCheckReturn(ret) + return c_ccSysCaps + +def nvmlSystemGetConfComputeState(): + c_state = c_nvmlConfComputeSystemState_t() + fn = _nvmlGetFunctionPointer("nvmlSystemGetConfComputeState") + ret = fn(byref(c_state)) + _nvmlCheckReturn(ret) + return c_state + +def nvmlSystemGetConfComputeSettings(settings): + fn = _nvmlGetFunctionPointer("nvmlSystemGetConfComputeSettings") + return fn(settings) + +def nvmlDeviceSetConfComputeUnprotectedMemSize(device, c_ccMemSize): + fn = _nvmlGetFunctionPointer("nvmlDeviceSetConfComputeUnprotectedMemSize") + ret = fn(device, c_ccMemSize) + _nvmlCheckReturn(ret) + return NVML_SUCCESS + +def nvmlDeviceGetConfComputeMemSizeInfo(device): + c_ccMemSize = c_nvmlConfComputeMemSizeInfo_t() + fn = _nvmlGetFunctionPointer("nvmlDeviceGetConfComputeMemSizeInfo") + ret = fn(device, byref(c_ccMemSize)) + _nvmlCheckReturn(ret) + return c_ccMemSize + +def nvmlDeviceGetConfComputeProtectedMemoryUsage(device): + c_memory = c_nvmlMemory_t() + fn = _nvmlGetFunctionPointer("nvmlDeviceGetConfComputeProtectedMemoryUsage") + ret = fn(device, byref(c_memory)) + _nvmlCheckReturn(ret) + return c_memory + +def nvmlDeviceGetConfComputeGpuCertificate(device): + c_cert = c_nvmlConfComputeGpuCertificate_t() + fn = _nvmlGetFunctionPointer("nvmlDeviceGetConfComputeGpuCertificate") + ret = fn(device, byref(c_cert)) + _nvmlCheckReturn(ret) + return c_cert + +def nvmlDeviceGetConfComputeGpuAttestationReport(device, c_nonce): + c_attestReport = c_nvmlConfComputeGpuAttestationReport_t() + c_nonce_arr = (c_uint8 * len(c_nonce))(*(c_nonce)) + setattr(c_attestReport, 'nonce', c_nonce_arr) + fn = _nvmlGetFunctionPointer("nvmlDeviceGetConfComputeGpuAttestationReport") + ret = fn(device, byref(c_attestReport)) + _nvmlCheckReturn(ret) + return c_attestReport + +def nvmlSystemSetConfComputeKeyRotationThresholdInfo(max_atk_adv): + c_keyRotationThrInfo = c_nvmlConfComputeSetKeyRotationThresholdInfo_t(0) + c_keyRotationThrInfo.version = ConfComputeSetKeyRotationThresholdInfo_v1 + c_keyRotationThrInfo.maxAttackerAdvantage = max_atk_adv + fn = _nvmlGetFunctionPointer("nvmlSystemSetConfComputeKeyRotationThresholdInfo") + ret = fn(byref(c_keyRotationThrInfo)) + _nvmlCheckReturn(ret) + return NVML_SUCCESS + +def nvmlSystemGetConfComputeKeyRotationThresholdInfo(): + c_keyRotationThrInfo = c_nvmlConfComputeGetKeyRotationThresholdInfo_t(0) + c_keyRotationThrInfo.version = ConfComputeGetKeyRotationThresholdInfo_v1 + fn = _nvmlGetFunctionPointer("nvmlSystemGetConfComputeKeyRotationThresholdInfo") + ret = fn(byref(c_keyRotationThrInfo)) + _nvmlCheckReturn(ret) + return c_keyRotationThrInfo + +## GPM ## +######### + +## Enums/defines + +#### GPM Metric Identifiers +NVML_GPM_METRIC_GRAPHICS_UTIL = 1 # Percentage of time any compute/graphics app was active on the GPU. 0.0 - 100.0 +NVML_GPM_METRIC_SM_UTIL = 2 # Percentage of SMs that were busy. 0.0 - 100.0 +NVML_GPM_METRIC_SM_OCCUPANCY = 3 # Percentage of warps that were active vs theoretical maximum. 0.0 - 100.0 +NVML_GPM_METRIC_INTEGER_UTIL = 4 # Percentage of time the GPU's SMs were doing integer operations. 0.0 - 100.0 +NVML_GPM_METRIC_ANY_TENSOR_UTIL = 5 # Percentage of time the GPU's SMs were doing ANY tensor operations. 0.0 - 100.0 +NVML_GPM_METRIC_DFMA_TENSOR_UTIL = 6 # Percentage of time the GPU's SMs were doing DFMA tensor operations. 0.0 - 100.0 +NVML_GPM_METRIC_HMMA_TENSOR_UTIL = 7 # Percentage of time the GPU's SMs were doing HMMA tensor operations. 0.0 - 100.0 +NVML_GPM_METRIC_IMMA_TENSOR_UTIL = 9 # Percentage of time the GPU's SMs were doing IMMA tensor operations. 0.0 - 100.0 +NVML_GPM_METRIC_DRAM_BW_UTIL = 10 # Percentage of DRAM bw used vs theoretical maximum. 0.0 - 100.0 +NVML_GPM_METRIC_FP64_UTIL = 11 # Percentage of time the GPU's SMs were doing non-tensor FP64 math. 0.0 - 100.0 +NVML_GPM_METRIC_FP32_UTIL = 12 # Percentage of time the GPU's SMs were doing non-tensor FP32 math. 0.0 - 100.0 +NVML_GPM_METRIC_FP16_UTIL = 13 # Percentage of time the GPU's SMs were doing non-tensor FP16 math. 0.0 - 100.0 +NVML_GPM_METRIC_PCIE_TX_PER_SEC = 20 # PCIe traffic from this GPU in MiB/sec +NVML_GPM_METRIC_PCIE_RX_PER_SEC = 21 # PCIe traffic to this GPU in MiB/sec +NVML_GPM_METRIC_NVDEC_0_UTIL = 30 # Percent utilization of NVDEC 0. 0.0 - 100.0 +NVML_GPM_METRIC_NVDEC_1_UTIL = 31 # Percent utilization of NVDEC 1. 0.0 - 100.0 +NVML_GPM_METRIC_NVDEC_2_UTIL = 32 # Percent utilization of NVDEC 2. 0.0 - 100.0 +NVML_GPM_METRIC_NVDEC_3_UTIL = 33 # Percent utilization of NVDEC 3. 0.0 - 100.0 +NVML_GPM_METRIC_NVDEC_4_UTIL = 34 # Percent utilization of NVDEC 4. 0.0 - 100.0 +NVML_GPM_METRIC_NVDEC_5_UTIL = 35 # Percent utilization of NVDEC 5. 0.0 - 100.0 +NVML_GPM_METRIC_NVDEC_6_UTIL = 36 # Percent utilization of NVDEC 6. 0.0 - 100.0 +NVML_GPM_METRIC_NVDEC_7_UTIL = 37 # Percent utilization of NVDEC 7. 0.0 - 100.0 +NVML_GPM_METRIC_NVJPG_0_UTIL = 40 # Percent utilization of NVJPG 0. 0.0 - 100.0 +NVML_GPM_METRIC_NVJPG_1_UTIL = 41 # Percent utilization of NVJPG 1. 0.0 - 100.0 +NVML_GPM_METRIC_NVJPG_2_UTIL = 42 # Percent utilization of NVJPG 2. 0.0 - 100.0 +NVML_GPM_METRIC_NVJPG_3_UTIL = 43 # Percent utilization of NVJPG 3. 0.0 - 100.0 +NVML_GPM_METRIC_NVJPG_4_UTIL = 44 # Percent utilization of NVJPG 4. 0.0 - 100.0 +NVML_GPM_METRIC_NVJPG_5_UTIL = 45 # Percent utilization of NVJPG 5. 0.0 - 100.0 +NVML_GPM_METRIC_NVJPG_6_UTIL = 46 # Percent utilization of NVJPG 6. 0.0 - 100.0 +NVML_GPM_METRIC_NVJPG_7_UTIL = 47 # Percent utilization of NVJPG 7. 0.0 - 100.0 +NVML_GPM_METRIC_NVOFA_0_UTIL = 50 # Percent utilization of NVOFA 0. 0.0 - 100.0 +NVML_GPM_METRIC_NVOFA_1_UTIL = 51 # Percent utilization of NVOFA 1. 0.0 - 100.0 +NVML_GPM_METRIC_NVLINK_TOTAL_RX_PER_SEC = 60 # NvLink read bandwidth for all links in MiB/sec +NVML_GPM_METRIC_NVLINK_TOTAL_TX_PER_SEC = 61 # NvLink write bandwidth for all links in MiB/sec +NVML_GPM_METRIC_NVLINK_L0_RX_PER_SEC = 62 # NvLink read bandwidth for link 0 in MiB/sec +NVML_GPM_METRIC_NVLINK_L0_TX_PER_SEC = 63 # NvLink write bandwidth for link 0 in MiB/sec +NVML_GPM_METRIC_NVLINK_L1_RX_PER_SEC = 64 # NvLink read bandwidth for link 1 in MiB/sec +NVML_GPM_METRIC_NVLINK_L1_TX_PER_SEC = 65 # NvLink write bandwidth for link 1 in MiB/sec +NVML_GPM_METRIC_NVLINK_L2_RX_PER_SEC = 66 # NvLink read bandwidth for link 2 in MiB/sec +NVML_GPM_METRIC_NVLINK_L2_TX_PER_SEC = 67 # NvLink write bandwidth for link 2 in MiB/sec +NVML_GPM_METRIC_NVLINK_L3_RX_PER_SEC = 68 # NvLink read bandwidth for link 3 in MiB/sec +NVML_GPM_METRIC_NVLINK_L3_TX_PER_SEC = 69 # NvLink write bandwidth for link 3 in MiB/sec +NVML_GPM_METRIC_NVLINK_L4_RX_PER_SEC = 70 # NvLink read bandwidth for link 4 in MiB/sec +NVML_GPM_METRIC_NVLINK_L4_TX_PER_SEC = 71 # NvLink write bandwidth for link 4 in MiB/sec +NVML_GPM_METRIC_NVLINK_L5_RX_PER_SEC = 72 # NvLink read bandwidth for link 5 in MiB/sec +NVML_GPM_METRIC_NVLINK_L5_TX_PER_SEC = 73 # NvLink write bandwidth for link 5 in MiB/sec +NVML_GPM_METRIC_NVLINK_L6_RX_PER_SEC = 74 # NvLink read bandwidth for link 6 in MiB/sec +NVML_GPM_METRIC_NVLINK_L6_TX_PER_SEC = 75 # NvLink write bandwidth for link 6 in MiB/sec +NVML_GPM_METRIC_NVLINK_L7_RX_PER_SEC = 76 # NvLink read bandwidth for link 7 in MiB/sec +NVML_GPM_METRIC_NVLINK_L7_TX_PER_SEC = 77 # NvLink write bandwidth for link 7 in MiB/sec +NVML_GPM_METRIC_NVLINK_L8_RX_PER_SEC = 78 # NvLink read bandwidth for link 8 in MiB/sec +NVML_GPM_METRIC_NVLINK_L8_TX_PER_SEC = 79 # NvLink write bandwidth for link 8 in MiB/sec +NVML_GPM_METRIC_NVLINK_L9_RX_PER_SEC = 80 # NvLink read bandwidth for link 9 in MiB/sec +NVML_GPM_METRIC_NVLINK_L9_TX_PER_SEC = 81 # NvLink write bandwidth for link 9 in MiB/sec +NVML_GPM_METRIC_NVLINK_L10_RX_PER_SEC = 82 # NvLink read bandwidth for link 10 in MiB/sec +NVML_GPM_METRIC_NVLINK_L10_TX_PER_SEC = 83 # NvLink write bandwidth for link 10 in MiB/sec +NVML_GPM_METRIC_NVLINK_L11_RX_PER_SEC = 84 # NvLink read bandwidth for link 11 in MiB/sec +NVML_GPM_METRIC_NVLINK_L11_TX_PER_SEC = 85 # NvLink write bandwidth for link 11 in MiB/sec +NVML_GPM_METRIC_NVLINK_L12_RX_PER_SEC = 86 # NvLink read bandwidth for link 12 in MiB/sec +NVML_GPM_METRIC_NVLINK_L12_TX_PER_SEC = 87 # NvLink write bandwidth for link 12 in MiB/sec +NVML_GPM_METRIC_NVLINK_L13_RX_PER_SEC = 88 # NvLink read bandwidth for link 13 in MiB/sec +NVML_GPM_METRIC_NVLINK_L13_TX_PER_SEC = 89 # NvLink write bandwidth for link 13 in MiB/sec +NVML_GPM_METRIC_NVLINK_L14_RX_PER_SEC = 90 # NvLink read bandwidth for link 14 in MiB/sec +NVML_GPM_METRIC_NVLINK_L14_TX_PER_SEC = 91 # NvLink write bandwidth for link 14 in MiB/sec +NVML_GPM_METRIC_NVLINK_L15_RX_PER_SEC = 92 # NvLink read bandwidth for link 15 in MiB/sec +NVML_GPM_METRIC_NVLINK_L15_TX_PER_SEC = 93 # NvLink write bandwidth for link 15 in MiB/sec +NVML_GPM_METRIC_NVLINK_L16_RX_PER_SEC = 94 # NvLink read bandwidth for link 16 in MiB/sec +NVML_GPM_METRIC_NVLINK_L16_TX_PER_SEC = 95 # NvLink write bandwidth for link 16 in MiB/sec +NVML_GPM_METRIC_NVLINK_L17_RX_PER_SEC = 96 # NvLink read bandwidth for link 17 in MiB/sec +NVML_GPM_METRIC_NVLINK_L17_TX_PER_SEC = 97 # NvLink write bandwidth for link 17 in MiB/sec +NVML_GPM_METRIC_C2C_TOTAL_TX_PER_SEC = 100 +NVML_GPM_METRIC_C2C_TOTAL_RX_PER_SEC = 101 +NVML_GPM_METRIC_C2C_DATA_TX_PER_SEC = 102 +NVML_GPM_METRIC_C2C_DATA_RX_PER_SEC = 103 +NVML_GPM_METRIC_C2C_LINK0_TOTAL_TX_PER_SEC = 104 +NVML_GPM_METRIC_C2C_LINK0_TOTAL_RX_PER_SEC = 105 +NVML_GPM_METRIC_C2C_LINK0_DATA_TX_PER_SEC = 106 +NVML_GPM_METRIC_C2C_LINK0_DATA_RX_PER_SEC = 107 +NVML_GPM_METRIC_C2C_LINK1_TOTAL_TX_PER_SEC = 108 +NVML_GPM_METRIC_C2C_LINK1_TOTAL_RX_PER_SEC = 109 +NVML_GPM_METRIC_C2C_LINK1_DATA_TX_PER_SEC = 110 +NVML_GPM_METRIC_C2C_LINK1_DATA_RX_PER_SEC = 111 +NVML_GPM_METRIC_C2C_LINK2_TOTAL_TX_PER_SEC = 112 +NVML_GPM_METRIC_C2C_LINK2_TOTAL_RX_PER_SEC = 113 +NVML_GPM_METRIC_C2C_LINK2_DATA_TX_PER_SEC = 114 +NVML_GPM_METRIC_C2C_LINK2_DATA_RX_PER_SEC = 115 +NVML_GPM_METRIC_C2C_LINK3_TOTAL_TX_PER_SEC = 116 +NVML_GPM_METRIC_C2C_LINK3_TOTAL_RX_PER_SEC = 117 +NVML_GPM_METRIC_C2C_LINK3_DATA_TX_PER_SEC = 118 +NVML_GPM_METRIC_C2C_LINK3_DATA_RX_PER_SEC = 119 +NVML_GPM_METRIC_C2C_LINK4_TOTAL_TX_PER_SEC = 120 +NVML_GPM_METRIC_C2C_LINK4_TOTAL_RX_PER_SEC = 121 +NVML_GPM_METRIC_C2C_LINK4_DATA_TX_PER_SEC = 122 +NVML_GPM_METRIC_C2C_LINK4_DATA_RX_PER_SEC = 123 +NVML_GPM_METRIC_C2C_LINK5_TOTAL_TX_PER_SEC = 124 +NVML_GPM_METRIC_C2C_LINK5_TOTAL_RX_PER_SEC = 125 +NVML_GPM_METRIC_C2C_LINK5_DATA_TX_PER_SEC = 126 +NVML_GPM_METRIC_C2C_LINK5_DATA_RX_PER_SEC = 127 +NVML_GPM_METRIC_C2C_LINK6_TOTAL_TX_PER_SEC = 128 +NVML_GPM_METRIC_C2C_LINK6_TOTAL_RX_PER_SEC = 129 +NVML_GPM_METRIC_C2C_LINK6_DATA_TX_PER_SEC = 130 +NVML_GPM_METRIC_C2C_LINK6_DATA_RX_PER_SEC = 131 +NVML_GPM_METRIC_C2C_LINK7_TOTAL_TX_PER_SEC = 132 +NVML_GPM_METRIC_C2C_LINK7_TOTAL_RX_PER_SEC = 133 +NVML_GPM_METRIC_C2C_LINK7_DATA_TX_PER_SEC = 134 +NVML_GPM_METRIC_C2C_LINK7_DATA_RX_PER_SEC = 135 +NVML_GPM_METRIC_C2C_LINK8_TOTAL_TX_PER_SEC = 136 +NVML_GPM_METRIC_C2C_LINK8_TOTAL_RX_PER_SEC = 137 +NVML_GPM_METRIC_C2C_LINK8_DATA_TX_PER_SEC = 138 +NVML_GPM_METRIC_C2C_LINK8_DATA_RX_PER_SEC = 139 +NVML_GPM_METRIC_C2C_LINK9_TOTAL_TX_PER_SEC = 140 +NVML_GPM_METRIC_C2C_LINK9_TOTAL_RX_PER_SEC = 141 +NVML_GPM_METRIC_C2C_LINK9_DATA_TX_PER_SEC = 142 +NVML_GPM_METRIC_C2C_LINK9_DATA_RX_PER_SEC = 143 +NVML_GPM_METRIC_C2C_LINK10_TOTAL_TX_PER_SEC = 144 +NVML_GPM_METRIC_C2C_LINK10_TOTAL_RX_PER_SEC = 145 +NVML_GPM_METRIC_C2C_LINK10_DATA_TX_PER_SEC = 146 +NVML_GPM_METRIC_C2C_LINK10_DATA_RX_PER_SEC = 147 +NVML_GPM_METRIC_C2C_LINK11_TOTAL_TX_PER_SEC = 148 +NVML_GPM_METRIC_C2C_LINK11_TOTAL_RX_PER_SEC = 149 +NVML_GPM_METRIC_C2C_LINK11_DATA_TX_PER_SEC = 150 +NVML_GPM_METRIC_C2C_LINK11_DATA_RX_PER_SEC = 151 +NVML_GPM_METRIC_C2C_LINK12_TOTAL_TX_PER_SEC = 152 +NVML_GPM_METRIC_C2C_LINK12_TOTAL_RX_PER_SEC = 153 +NVML_GPM_METRIC_C2C_LINK12_DATA_TX_PER_SEC = 154 +NVML_GPM_METRIC_C2C_LINK12_DATA_RX_PER_SEC = 155 +NVML_GPM_METRIC_C2C_LINK13_TOTAL_TX_PER_SEC = 156 +NVML_GPM_METRIC_C2C_LINK13_TOTAL_RX_PER_SEC = 157 +NVML_GPM_METRIC_C2C_LINK13_DATA_TX_PER_SEC = 158 +NVML_GPM_METRIC_C2C_LINK13_DATA_RX_PER_SEC = 159 +NVML_GPM_METRIC_HOSTMEM_CACHE_HIT = 160 +NVML_GPM_METRIC_HOSTMEM_CACHE_MISS = 161 +NVML_GPM_METRIC_PEERMEM_CACHE_HIT = 162 +NVML_GPM_METRIC_PEERMEM_CACHE_MISS = 163 +NVML_GPM_METRIC_DRAM_CACHE_HIT = 164 +NVML_GPM_METRIC_DRAM_CACHE_MISS = 165 +NVML_GPM_METRIC_NVENC_0_UTIL = 166, +NVML_GPM_METRIC_NVENC_1_UTIL = 167, +NVML_GPM_METRIC_NVENC_2_UTIL = 168, +NVML_GPM_METRIC_NVENC_3_UTIL = 169, +NVML_GPM_METRIC_GR0_CTXSW_CYCLES_ELAPSED = 170, +NVML_GPM_METRIC_GR0_CTXSW_CYCLES_ACTIVE = 171, +NVML_GPM_METRIC_GR0_CTXSW_REQUESTS = 172, +NVML_GPM_METRIC_GR0_CTXSW_CYCLES_PER_REQ = 173, +NVML_GPM_METRIC_GR0_CTXSW_ACTIVE_PCT = 174, +NVML_GPM_METRIC_GR1_CTXSW_CYCLES_ELAPSED = 175, +NVML_GPM_METRIC_GR1_CTXSW_CYCLES_ACTIVE = 176, +NVML_GPM_METRIC_GR1_CTXSW_REQUESTS = 177, +NVML_GPM_METRIC_GR1_CTXSW_CYCLES_PER_REQ = 178, +NVML_GPM_METRIC_GR1_CTXSW_ACTIVE_PCT = 179, +NVML_GPM_METRIC_GR2_CTXSW_CYCLES_ELAPSED = 180, +NVML_GPM_METRIC_GR2_CTXSW_CYCLES_ACTIVE = 181, +NVML_GPM_METRIC_GR2_CTXSW_REQUESTS = 182, +NVML_GPM_METRIC_GR2_CTXSW_CYCLES_PER_REQ = 183, +NVML_GPM_METRIC_GR2_CTXSW_ACTIVE_PCT = 184, +NVML_GPM_METRIC_GR3_CTXSW_CYCLES_ELAPSED = 185, +NVML_GPM_METRIC_GR3_CTXSW_CYCLES_ACTIVE = 186, +NVML_GPM_METRIC_GR3_CTXSW_REQUESTS = 187, +NVML_GPM_METRIC_GR3_CTXSW_CYCLES_PER_REQ = 188, +NVML_GPM_METRIC_GR3_CTXSW_ACTIVE_PCT = 189, +NVML_GPM_METRIC_GR4_CTXSW_CYCLES_ELAPSED = 190, +NVML_GPM_METRIC_GR4_CTXSW_CYCLES_ACTIVE = 191, +NVML_GPM_METRIC_GR4_CTXSW_REQUESTS = 192, +NVML_GPM_METRIC_GR4_CTXSW_CYCLES_PER_REQ = 193, +NVML_GPM_METRIC_GR4_CTXSW_ACTIVE_PCT = 194, +NVML_GPM_METRIC_GR5_CTXSW_CYCLES_ELAPSED = 195, +NVML_GPM_METRIC_GR5_CTXSW_CYCLES_ACTIVE = 196, +NVML_GPM_METRIC_GR5_CTXSW_REQUESTS = 197, +NVML_GPM_METRIC_GR5_CTXSW_CYCLES_PER_REQ = 198, +NVML_GPM_METRIC_GR5_CTXSW_ACTIVE_PCT = 199, +NVML_GPM_METRIC_GR6_CTXSW_CYCLES_ELAPSED = 200, +NVML_GPM_METRIC_GR6_CTXSW_CYCLES_ACTIVE = 201, +NVML_GPM_METRIC_GR6_CTXSW_REQUESTS = 202, +NVML_GPM_METRIC_GR6_CTXSW_CYCLES_PER_REQ = 203, +NVML_GPM_METRIC_GR6_CTXSW_ACTIVE_PCT = 204, +NVML_GPM_METRIC_GR7_CTXSW_CYCLES_ELAPSED = 205, +NVML_GPM_METRIC_GR7_CTXSW_CYCLES_ACTIVE = 206, +NVML_GPM_METRIC_GR7_CTXSW_REQUESTS = 207, +NVML_GPM_METRIC_GR7_CTXSW_CYCLES_PER_REQ = 208, +NVML_GPM_METRIC_GR7_CTXSW_ACTIVE_PCT = 209, +NVML_GPM_METRIC_MAX = 210 + +## Structs + +class c_nvmlUnitInfo_t(_PrintableStructure): + _fields_ = [ + ('name', c_char * 96), + ('id', c_char * 96), + ('serial', c_char * 96), + ('firmwareVersion', c_char * 96), + ] + +class struct_c_nvmlGpmSample_t(Structure): + pass # opaque handle +c_nvmlGpmSample_t = POINTER(struct_c_nvmlGpmSample_t) + +class c_metricInfo_t(Structure): + _fields_ = [ + ("shortName", c_char_p), + ("longName", c_char_p), + ("unit", c_char_p), + ] + +class c_nvmlGpmMetric_t(_PrintableStructure): + _fields_ = [ + ('metricId', c_uint), + ('nvmlReturn', _nvmlReturn_t), + ('value', c_double), + ('metricInfo', c_metricInfo_t) + ] + +class c_nvmlGpmMetricsGet_t(_PrintableStructure): + _fields_ = [ + ('version', c_uint), + ('numMetrics', c_uint), + ('sample1', c_nvmlGpmSample_t), + ('sample2', c_nvmlGpmSample_t), + ('metrics', c_nvmlGpmMetric_t * NVML_GPM_METRIC_MAX) + ] + +NVML_GPM_METRICS_GET_VERSION = 1 + +class c_nvmlGpmSupport_t(_PrintableStructure): + _fields_ = [ + ('version', c_uint), + ('isSupportedDevice', c_uint), + ] + +NVML_GPM_SUPPORT_VERSION = 1 + +## Functions + +def nvmlGpmMetricsGet(metricsGet): + fn = _nvmlGetFunctionPointer("nvmlGpmMetricsGet") + ret = fn(byref(metricsGet)) + _nvmlCheckReturn(ret) + return metricsGet + +def nvmlGpmSampleFree(gpmSample): + fn = _nvmlGetFunctionPointer("nvmlGpmSampleFree") + ret = fn(gpmSample) + _nvmlCheckReturn(ret) + return + +def nvmlGpmSampleAlloc(): + gpmSample = c_nvmlGpmSample_t() + fn = _nvmlGetFunctionPointer("nvmlGpmSampleAlloc") + ret = fn(byref(gpmSample)) + _nvmlCheckReturn(ret) + return gpmSample + +def nvmlGpmSampleGet(device, gpmSample): + fn = _nvmlGetFunctionPointer("nvmlGpmSampleGet") + ret = fn(device, gpmSample) + _nvmlCheckReturn(ret) + return gpmSample + +def nvmlGpmMigSampleGet(device, gpuInstanceId, gpmSample): + fn = _nvmlGetFunctionPointer("nvmlGpmMigSampleGet") + ret = fn(device, gpuInstanceId, gpmSample) + _nvmlCheckReturn(ret) + return gpmSample + +def nvmlGpmQueryDeviceSupport(device): + gpmSupport = c_nvmlGpmSupport_t() + gpmSupport.version = NVML_GPM_SUPPORT_VERSION + fn = _nvmlGetFunctionPointer("nvmlGpmQueryDeviceSupport") + ret = fn(device, byref(gpmSupport)) + _nvmlCheckReturn(ret) + return gpmSupport + +def nvmlGpmSetStreamingEnabled(device, state): + c_state = c_uint(state) + fn = _nvmlGetFunctionPointer("nvmlGpmSetStreamingEnabled") + ret = fn(device, c_state) + _nvmlCheckReturn(ret) + return NVML_SUCCESS + +def nvmlGpmQueryIfStreamingEnabled(device): + c_state = c_uint() + fn = _nvmlGetFunctionPointer("nvmlGpmQueryIfStreamingEnabled") + ret = fn(device, byref(c_state)) + _nvmlCheckReturn(ret) + return c_state.value + +# Low Power Structure and Function + +NVML_NVLINK_POWER_STATE_HIGH_SPEED = 0x0 +NVML_NVLINK_POWER_STATE_LOW = 0x1 + +NVML_NVLINK_LOW_POWER_THRESHOLD_MIN = 0x1 +NVML_NVLINK_LOW_POWER_THRESHOLD_MAX = 0x1FFF +NVML_NVLINK_LOW_POWER_THRESHOLD_RESET = 0xFFFFFFFF +NVML_NVLINK_LOW_POWER_THRESHOLD_DEFAULT = NVML_NVLINK_LOW_POWER_THRESHOLD_RESET + +class c_nvmlNvLinkPowerThres_t(Structure): + _fields_ = [ + ("lowPwrThreshold", c_uint), + ] + +def nvmlDeviceSetNvLinkDeviceLowPowerThreshold(device, l1threshold): + c_info = c_nvmlNvLinkPowerThres_t() + c_info.lowPwrThreshold = l1threshold + fn = _nvmlGetFunctionPointer("nvmlDeviceSetNvLinkDeviceLowPowerThreshold") + ret = fn(device, byref(c_info)) + _nvmlCheckReturn(ret) + return NVML_SUCCESS + +NVML_GPU_FABRIC_UUID_LEN = 16 + +_nvmlGpuFabricState_t = c_uint +NVML_GPU_FABRIC_STATE_NOT_SUPPORTED = 0 +NVML_GPU_FABRIC_STATE_NOT_STARTED = 1 +NVML_GPU_FABRIC_STATE_IN_PROGRESS = 2 +NVML_GPU_FABRIC_STATE_COMPLETED = 3 + +class c_nvmlGpuFabricInfo_t(_PrintableStructure): + _fields_ = [ + ("clusterUuid", c_uint8 * NVML_DEVICE_UUID_BUFFER_SIZE), + ("status", _nvmlReturn_t), + ("cliqueId", c_uint32), + ("state", _nvmlGpuFabricState_t) + ] + +NVML_GPU_FABRIC_HEALTH_MASK_DEGRADED_BW_NOT_SUPPORTED = 0 +NVML_GPU_FABRIC_HEALTH_MASK_DEGRADED_BW_TRUE = 1 +NVML_GPU_FABRIC_HEALTH_MASK_DEGRADED_BW_FALSE = 2 +NVML_GPU_FABRIC_HEALTH_MASK_SHIFT_DEGRADED_BW = 0 +NVML_GPU_FABRIC_HEALTH_MASK_WIDTH_DEGRADED_BW = 0x11 + +NVML_GPU_FABRIC_HEALTH_MASK_ROUTE_RECOVERY_NOT_SUPPORTED = 0 +NVML_GPU_FABRIC_HEALTH_MASK_ROUTE_RECOVERY_TRUE = 1 +NVML_GPU_FABRIC_HEALTH_MASK_ROUTE_RECOVERY_FALSE = 2 +NVML_GPU_FABRIC_HEALTH_MASK_SHIFT_ROUTE_RECOVERY = 2 +NVML_GPU_FABRIC_HEALTH_MASK_WIDTH_ROUTE_RECOVERY = 0x11 + +NVML_GPU_FABRIC_HEALTH_MASK_ROUTE_UNHEALTHY_NOT_SUPPORTED = 0 +NVML_GPU_FABRIC_HEALTH_MASK_ROUTE_UNHEALTHY_TRUE = 1 +NVML_GPU_FABRIC_HEALTH_MASK_ROUTE_UNHEALTHY_FALSE = 2 +NVML_GPU_FABRIC_HEALTH_MASK_SHIFT_ROUTE_UNHEALTHY = 4 +NVML_GPU_FABRIC_HEALTH_MASK_WIDTH_ROUTE_UNHEALTHY = 0x11 + +NVML_GPU_FABRIC_HEALTH_MASK_ACCESS_TIMEOUT_RECOVERY_NOT_SUPPORTED = 0 +NVML_GPU_FABRIC_HEALTH_MASK_ACCESS_TIMEOUT_RECOVERY_TRUE = 1 +NVML_GPU_FABRIC_HEALTH_MASK_ACCESS_TIMEOUT_RECOVERY_FALSE = 2 +NVML_GPU_FABRIC_HEALTH_MASK_SHIFT_ACCESS_TIMEOUT_RECOVERY = 6 +NVML_GPU_FABRIC_HEALTH_MASK_WIDTH_ACCESS_TIMEOUT_RECOVERY = 0x11 + +NVML_GPU_FABRIC_HEALTH_MASK_INCORRECT_CONFIGURATION_NOT_SUPPORTED = 0 +NVML_GPU_FABRIC_HEALTH_MASK_INCORRECT_CONFIGURATION_NONE = 1 +NVML_GPU_FABRIC_HEALTH_MASK_INCORRECT_CONFIGURATION_INCORRECT_SYSGUID = 2 +NVML_GPU_FABRIC_HEALTH_MASK_INCORRECT_CONFIGURATION_INCORRECT_CHASSIS_SN = 3 +NVML_GPU_FABRIC_HEALTH_MASK_INCORRECT_CONFIGURATION_NO_PARTITION = 4 +NVML_GPU_FABRIC_HEALTH_MASK_INCORRECT_CONFIGURATION_INSUFFICIENT_NVLINKS = 5 +NVML_GPU_FABRIC_HEALTH_MASK_SHIFT_INCORRECT_CONFIGURATION = 8 +NVML_GPU_FABRIC_HEALTH_MASK_WIDTH_INCORRECT_CONFIGURATION = 0xf + +NVML_GPU_FABRIC_HEALTH_SUMMARY_NOT_SUPPORTED = 0 +NVML_GPU_FABRIC_HEALTH_SUMMARY_HEALTHY = 1 +NVML_GPU_FABRIC_HEALTH_SUMMARY_UNHEALTHY = 2 +NVML_GPU_FABRIC_HEALTH_SUMMARY_LIMITED_CAPACITY = 3 + +nvmlGpuFabricInfo_v2 = 0x02000024 + +class c_nvmlGpuFabricInfo_v2_t(_PrintableStructure): + _fields_ = [ + ("version", c_uint), + ("clusterUuid", c_uint8 * NVML_GPU_FABRIC_UUID_LEN), + ("status", _nvmlReturn_t), + ("cliqueId", c_uint32), + ("state", _nvmlGpuFabricState_t), + ("healthMask", c_uint32) + ] + + def __init__(self): + super(c_nvmlGpuFabricInfo_v2_t, self).__init__(version=nvmlGpuFabricInfo_v2) + +c_nvmlGpuFabricInfoV_t = c_nvmlGpuFabricInfo_v2_t + +class c_nvmlGpuFabricInfo_v3_t(_PrintableStructure): + _fields_ = [ + ("version", c_uint), + ("clusterUuid", c_uint8 * NVML_GPU_FABRIC_UUID_LEN), + ("status", _nvmlReturn_t), + ("cliqueId", c_uint32), + ("state", _nvmlGpuFabricState_t), + ("healthMask", c_uint32), + ("healthSummary", c_uint8) + ] + + def __init__(self): + super(c_nvmlGpuFabricInfo_v3_t, self).__init__(version=nvmlGpuFabricInfo_v3) + +nvmlGpuFabricInfo_v3 = 0x3000028 + +# Deprecated +def nvmlDeviceGetGpuFabricInfo(device, gpuFabricInfo): + fn = _nvmlGetFunctionPointer("nvmlDeviceGetGpuFabricInfo"); + ret = fn(device, gpuFabricInfo) + _nvmlCheckReturn(ret) + return NVML_SUCCESS + +def nvmlDeviceGetGpuFabricInfoV(device, gpuFabricInfo): + fn = _nvmlGetFunctionPointer("nvmlDeviceGetGpuFabricInfoV"); + ret = fn(device, gpuFabricInfo) + _nvmlCheckReturn(ret) + return NVML_SUCCESS + +###################### +## Enums/defines +#### NVML GPU NVLINK BW MODE +NVML_GPU_NVLINK_BW_MODE_FULL = 0x0 +NVML_GPU_NVLINK_BW_MODE_OFF = 0x1 +NVML_GPU_NVLINK_BW_MODE_MIN = 0x2 +NVML_GPU_NVLINK_BW_MODE_HALF = 0x3 +NVML_GPU_NVLINK_BW_MODE_3QUARTER = 0x4 +NVML_GPU_NVLINK_BW_MODE_COUNT = 0x5 + +def nvmlSystemSetNvlinkBwMode(mode): + fn = _nvmlGetFunctionPointer("nvmlSystemSetNvlinkBwMode") + ret = fn(mode) + _nvmlCheckReturn(ret) + return NVML_SUCCESS + +def nvmlSystemGetNvlinkBwMode(): + mode = c_uint() + fn = _nvmlGetFunctionPointer("nvmlSystemGetNvlinkBwMode") + ret = fn(byref(mode)) + _nvmlCheckReturn(ret) + return mode.value + +_nvmlPowerScopeType_t = c_uint +NVML_POWER_SCOPE_GPU = 0 +NVML_POWER_SCOPE_MODULE = 1 +NVML_POWER_SCOPE_MEMORY = 2 + +class c_nvmlPowerValue_v2_t(_PrintableStructure): + _fields_ = [ + ('version', c_uint), + ('powerScope', _nvmlPowerScopeType_t), + ('powerValueMw', c_uint), + ] + _fmt_ = {'': "%d B"} + +nvmlPowerValue_v2 = 0x0200000C + +def nvmlDeviceSetPowerManagementLimit_v2(device, powerScope, powerLimit, version=nvmlPowerValue_v2): + c_powerScope = _nvmlPowerScopeType_t(powerScope) + c_powerValue = c_nvmlPowerValue_v2_t() + c_powerValue.version = c_uint(version) + c_powerValue.powerScope = c_powerScope + c_powerValue.powerValueMw = c_uint(powerLimit) + fn = _nvmlGetFunctionPointer("nvmlDeviceSetPowerManagementLimit_v2") + ret = fn(device, byref(c_powerValue)) + return NVML_SUCCESS + +class c_nvmlEccSramErrorStatus_v1_t(_PrintableStructure): + _fields_ = [ + ('version', c_uint), + ('aggregateUncParity', c_ulonglong), + ('aggregateUncSecDed', c_ulonglong), + ('aggregateCor', c_ulonglong), + ('volatileUncParity', c_ulonglong), + ('volatileUncSecDed', c_ulonglong), + ('volatileCor', c_ulonglong), + ('aggregateUncBucketL2', c_ulonglong), + ('aggregateUncBucketSm', c_ulonglong), + ('aggregateUncBucketPcie', c_ulonglong), + ('aggregateUncBucketMcu', c_ulonglong), + ('aggregateUncBucketOther', c_ulonglong), + ('bThresholdExceeded', c_uint) + ] + + def __init__(self): + super(c_nvmlEccSramErrorStatus_v1_t, self).__init__(version=nvmlEccSramErrorStatus_v1) + +nvmlEccSramErrorStatus_v1 = 0x1000068 +def nvmlDeviceGetSramEccErrorStatus(device, status): + fn = _nvmlGetFunctionPointer("nvmlDeviceGetSramEccErrorStatus") + ret = fn(device, status) + _nvmlCheckReturn(ret) + return NVML_SUCCESS + +NVML_DEV_CAP_EGM = (1 << 0) +nvmlDeviceCapabilities_v1 = 0x1000008 + +class c_nvmlDeviceCapabilities_v1_t(_PrintableStructure): + _fields_ = [ + ('version', c_uint), + ('capMask', c_uint), + ] + + def __init__(self): + super(c_nvmlDeviceCapabilities_v1_t, self).__init__(version=nvmlDeviceCapabilities_v1) + + +def nvmlDeviceGetCapabilities(device, caps): + fn = _nvmlGetFunctionPointer("nvmlDeviceGetCapabilities") + return fn(device, caps) + +class c_nvmlPlatformInfo_v1_t(_PrintableStructure): + _fields_ = [ + ('version', c_uint), + ('ibGuid', c_char * 16), + ('rackGuid', c_char * 16), + ('chassisPhysicalSlotNumber', c_char), + ('computeSlotIndex', c_char), + ('nodeIndex', c_char), + ('peerType', c_char), + ('moduleId', c_char) + ] + + def __init__(self): + super(c_nvmlPlatformInfo_v1_t, self).__init__(version=nvmlPlatformInfo_v1) + +class c_nvmlPlatformInfo_v2_t(_PrintableStructure): + _fields_ = [ + ('version', c_uint), + ('ibGuid', c_char * 16), + ('chassisSerialNumber', c_char * 16), + ('slotNumber', c_char), + ('trayIndex', c_char), + ('hostId', c_char), + ('peerType', c_char), + ('moduleId', c_char) + ] + + def __init__(self): + super(c_nvmlPlatformInfo_v2_t, self).__init__(version=nvmlPlatformInfo_v2) + +nvmlPlatformInfo_v1 = 0x100002c +nvmlPlatformInfo_v2 = 0x200002c + +def nvmlDeviceGetPlatformInfo(device, platformInfo): + fn = _nvmlGetFunctionPointer("nvmlDeviceGetPlatformInfo") + ret = fn(device, platformInfo) + _nvmlCheckReturn(ret) + return NVML_SUCCESS + +class c_nvmlMask255_t(_PrintableStructure): + _fields_ = [ + ('mask', c_uint * 8), + ] + +NVML_WORKLOAD_POWER_MAX_PROFILES = 255 +NVML_POWER_PROFILE_MAX_P = 0 +NVML_POWER_PROFILE_MAX_Q = 1 +NVML_POWER_PROFILE_COMPUTE = 2 +NVML_POWER_PROFILE_MEMORY_BOUND = 3 +NVML_POWER_PROFILE_NETWORK = 4 +NVML_POWER_PROFILE_BALANCED = 5 +NVML_POWER_PROFILE_LLM_INFERENCE = 6 +NVML_POWER_PROFILE_LLM_TRAINING = 7 +NVML_POWER_PROFILE_RBM = 8 +NVML_POWER_PROFILE_DCPCIE = 9 +NVML_POWER_PROFILE_HMMA_SPARSE = 10 +NVML_POWER_PROFILE_HMMA_DENSE = 11 +NVML_POWER_PROFILE_SYNC_BALANCED = 12 +NVML_POWER_PROFILE_HPC = 13 +NVML_POWER_PROFILE_MIG = 14 +NVML_POWER_PROFILE_MAX = 15 + +nvmlWorkloadPowerProfileInfo_v1 = 0x100002c +class c_nvmlWorkloadPowerProfileInfo_v1_t(_PrintableStructure): + _fields_ = [ + ('version', c_uint), + ('profileId', c_uint), + ('priority', c_uint), + ('conflictingmask', c_nvmlMask255_t) + ] + + def __init__(self): + super(c_nvmlWorkloadPowerProfileInfo_v1_t, self).__init__(version=nvmlWorkloadPowerProfileInfo_v1) + +nvmlWorkloadPowerProfileProfilesInfo_v1 = 0x1002bf8 +class c_nvmlWorkloadPowerProfileProfilesInfo_v1_t(_PrintableStructure): + _fields_ = [ + ('version', c_uint), + ('perfProfilesMask', c_nvmlMask255_t), + ('perfProfile', c_nvmlWorkloadPowerProfileInfo_v1_t * NVML_WORKLOAD_POWER_MAX_PROFILES) + ] + + def __init__(self): + super(c_nvmlWorkloadPowerProfileProfilesInfo_v1_t, self).__init__(version=nvmlWorkloadPowerProfileProfilesInfo_v1) + +nvmlWorkloadPowerProfileCurrentProfiles_v1 = 0x1000064 +class c_nvmlWorkloadPowerProfileCurrentProfiles_v1_t(_PrintableStructure): + _fields_ = [ + ('version', c_uint), + ('perfProfilesMask', c_nvmlMask255_t), + ('requestedProfilesMask', c_nvmlMask255_t), + ('enforcedProfilesMask', c_nvmlMask255_t) + ] + + def __init__(self): + super(c_nvmlWorkloadPowerProfileCurrentProfiles_v1_t, self).__init__(version=nvmlWorkloadPowerProfileCurrentProfiles_v1) + +nvmlWorkloadPowerProfileRequestedProfiles_v1 = 0x1000024 +class c_nvmlWorkloadPowerProfileRequestedProfiles_v1_t(_PrintableStructure): + _fields_ = [ + ('version', c_uint), + ('requestedProfilesMask', c_nvmlMask255_t), + ] + + def __init__(self): + super(c_nvmlWorkloadPowerProfileRequestedProfiles_v1_t, self).__init__(version=nvmlWorkloadPowerProfileRequestedProfiles_v1) + +def nvmlDeviceWorkloadPowerProfileGetProfilesInfo(device, profilesInfo): + fn = _nvmlGetFunctionPointer("nvmlDeviceWorkloadPowerProfileGetProfilesInfo") + ret = fn(device, profilesInfo) + _nvmlCheckReturn(ret) + return NVML_SUCCESS + +def nvmlDeviceWorkloadPowerProfileGetCurrentProfiles(device, currentProfiles): + fn = _nvmlGetFunctionPointer("nvmlDeviceWorkloadPowerProfileGetCurrentProfiles") + ret = fn(device, currentProfiles) + _nvmlCheckReturn(ret) + return NVML_SUCCESS + +def nvmlDeviceWorkloadPowerProfileSetRequestedProfiles(device, requestedProfiles): + fn = _nvmlGetFunctionPointer("nvmlDeviceWorkloadPowerProfileSetRequestedProfiles") + ret = fn(device, requestedProfiles) + _nvmlCheckReturn(ret) + return NVML_SUCCESS + +def nvmlDeviceWorkloadPowerProfileClearRequestedProfiles(device, requestedProfiles): + fn = _nvmlGetFunctionPointer("nvmlDeviceWorkloadPowerProfileClearRequestedProfiles") + ret = fn(device, requestedProfiles) + _nvmlCheckReturn(ret) + return NVML_SUCCESS + +def nvmlDeviceGetNvlinkSupportedBwModes(device, supportedBwModes): + fn = _nvmlGetFunctionPointer("nvmlDeviceGetNvlinkSupportedBwModes") + ret = fn(device, supportedBwModes) + _nvmlCheckReturn(ret) + return NVML_SUCCESS + +def nvmlDeviceGetNvlinkBwMode(device, getBwMode): + fn = _nvmlGetFunctionPointer("nvmlDeviceGetNvlinkBwMode") + ret = fn(device, getBwMode) + _nvmlCheckReturn(ret) + return NVML_SUCCESS + +def nvmlDeviceSetNvlinkBwMode(device, setBwMode): + fn = _nvmlGetFunctionPointer("nvmlDeviceSetNvlinkBwMode") + ret = fn(device, setBwMode) + _nvmlCheckReturn(ret) + return NVML_SUCCESS + +nvmlDramEncryptionInfo_v1 = 0x01000008 + +class c_nvmlDramEncryptionInfo_t(_PrintableStructure): + _fields_ = [ + ('version', c_uint), + ('encryptionState', _nvmlEnableState_t), + ] + + def __init__(self): + super(c_nvmlDramEncryptionInfo_t, self).__init__(version=nvmlDramEncryptionInfo_v1) + +def nvmlDeviceGetDramEncryptionMode(handle): + c_currState = c_nvmlDramEncryptionInfo_t() + c_pendingState = c_nvmlDramEncryptionInfo_t() + fn = _nvmlGetFunctionPointer("nvmlDeviceGetDramEncryptionMode") + ret = fn(handle, byref(c_currState), byref(c_pendingState)) + _nvmlCheckReturn(ret) + return [c_currState.encryptionState, c_pendingState.encryptionState] + +# added to API +def nvmlDeviceGetCurrentDramEncryptionMode(handle): + return nvmlDeviceGetDramEncryptionMode(handle)[0] + +# added to API +def nvmlDeviceGetPendingDramEncryptionMode(handle): + return nvmlDeviceGetDramEncryptionMode(handle)[1] + +def nvmlDeviceSetDramEncryptionMode(handle, mode): + fn = _nvmlGetFunctionPointer("nvmlDeviceSetDramEncryptionMode") + c_dramEncryptionMode = c_nvmlDramEncryptionInfo_t() + c_dramEncryptionMode.encryptionState = mode; + ret = fn(handle, byref(c_dramEncryptionMode)) + _nvmlCheckReturn(ret) + return None + +# Power Smoothing defines +NVML_POWER_SMOOTHING_MAX_NUM_PROFILES = 5 +NVML_POWER_SMOOTHING_ADMIN_OVERRIDE_NOT_SET = 0xFFFFFFFF +NVML_POWER_SMOOTHING_PROFILE_PARAM_PERCENT_TMP_FLOOR = 0 +NVML_POWER_SMOOTHING_PROFILE_PARAM_RAMP_UP_RATE = 1 +NVML_POWER_SMOOTHING_PROFILE_PARAM_RAMP_DOWN_RATE = 2 +NVML_POWER_SMOOTHING_PROFILE_PARAM_RAMP_DOWN_HYSTERESIS = 3 + +nvmlPowerSmoothingState_v1=0x1000008 +class c_nvmlPowerSmoothingState_v1_t(_PrintableStructure): + _fields_ = [ + ('version', c_uint), + ('state', c_uint), + ] + + def __init__(self): + super(c_nvmlPowerSmoothingState_v1_t, self).__init__(version=nvmlPowerSmoothingState_v1) + +nvmlPowerSmoothingProfile_v1=0x1000018 +class c_nvmlPowerSmoothingProfile_v1_t(_PrintableStructure): + _fields_ = [ + ('version', c_uint), + ('profileId', c_uint), + ('paramId', c_uint), + ('value', c_double), + ] + + def __init__(self): + super(c_nvmlPowerSmoothingProfile_v1_t, self).__init__(version=nvmlPowerSmoothingProfile_v1) + +def nvmlDevicePowerSmoothingActivatePresetProfile(device, profile): + fn = _nvmlGetFunctionPointer("nvmlDevicePowerSmoothingActivatePresetProfile") + ret = fn(device, profile) + _nvmlCheckReturn(ret) + +def nvmlDevicePowerSmoothingUpdatePresetProfileParam(device, profile): + fn = _nvmlGetFunctionPointer("nvmlDevicePowerSmoothingUpdatePresetProfileParam") + ret = fn(device, profile) + _nvmlCheckReturn(ret) + +def nvmlDevicePowerSmoothingSetState(device, state): + fn = _nvmlGetFunctionPointer("nvmlDevicePowerSmoothingSetState") + ret = fn(device, state) + _nvmlCheckReturn(ret) + +class c_nvmlEccSramUniqueUncorrectedErrorEntry_v1_t(_PrintableStructure): + _fields_ = [ + ('unit', c_uint), + ('location', c_uint), + ('sublocation', c_uint), + ('extlocation', c_uint), + ('address', c_uint), + ('isParity', c_uint), + ('count', c_uint) + ] + +class c_nvmlEccSramUniqueUncorrectedErrorCounts_v1_t(_PrintableStructure): + _fields_ = [ + ('version', c_uint), + ('entryCount', c_uint), + ('entries', POINTER(c_nvmlEccSramUniqueUncorrectedErrorEntry_v1_t)) + ] + + def __init__(self): + super(c_nvmlEccSramUniqueUncorrectedErrorCounts_v1_t, self).__init__(version=nvmlEccSramUniqueUncorrectedErrorCounts_v1) + +nvmlEccSramUniqueUncorrectedErrorCounts_v1 = 0x1000010 +def nvmlDeviceGetSramUniqueUncorrectedEccErrorCounts(device, counts): + fn = _nvmlGetFunctionPointer("nvmlDeviceGetSramUniqueUncorrectedEccErrorCounts") + ret = fn(device, counts) + _nvmlCheckReturn(ret) + return NVML_SUCCESS + +def nvmlDeviceGetPdi(device): + c_pdi = c_nvmlPdi_t() + fn = _nvmlGetFunctionPointer("nvmlDeviceGetPdi") + ret = fn(device, byref(c_pdi)) + _nvmlCheckReturn(ret) + return c_pdi.value + +def nvmlDeviceGetNvLinkInfo(device, info): + fn = _nvmlGetFunctionPointer("nvmlDeviceGetNvLinkInfo"); + ret = fn(device, info) + _nvmlCheckReturn(ret) + return NVML_SUCCESS + +def nvmlDeviceGetRepairStatus(device): + c_status = c_nvmlRepairStatus_t() + fn = _nvmlGetFunctionPointer("nvmlDeviceGetRepairStatus") + ret = fn(device, byref(c_status)) + _nvmlCheckReturn(ret) + return [c_status.bChannelRepairPending, c_status.bTpcRepairPending] diff --git a/lib/python3.12/site-packages/ray/_private/thirdparty/tabulate/__init__.py b/lib/python3.12/site-packages/ray/_private/thirdparty/tabulate/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/lib/python3.12/site-packages/ray/_private/thirdparty/tabulate/__pycache__/__init__.cpython-312.pyc b/lib/python3.12/site-packages/ray/_private/thirdparty/tabulate/__pycache__/__init__.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..debd129eed9e2294f36388321f416f1f934873a5 Binary files /dev/null and b/lib/python3.12/site-packages/ray/_private/thirdparty/tabulate/__pycache__/__init__.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/_private/thirdparty/tabulate/__pycache__/tabulate.cpython-312.pyc b/lib/python3.12/site-packages/ray/_private/thirdparty/tabulate/__pycache__/tabulate.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..ac2379a424a73020fc38558ae1ddbf04540f8182 Binary files /dev/null and b/lib/python3.12/site-packages/ray/_private/thirdparty/tabulate/__pycache__/tabulate.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/_private/thirdparty/tabulate/tabulate.py b/lib/python3.12/site-packages/ray/_private/thirdparty/tabulate/tabulate.py new file mode 100644 index 0000000000000000000000000000000000000000..83b1090ffaf9aabee2f874360c52415ebfeaf1fd --- /dev/null +++ b/lib/python3.12/site-packages/ray/_private/thirdparty/tabulate/tabulate.py @@ -0,0 +1,2720 @@ +# -*- coding: utf-8 -*- + +# Version 0.9.0, commit bf58e37e6b35e3cc9a0bd740f752abfd32b6e6f8 + +"""Pretty-print tabular data.""" + +from collections import namedtuple +from collections.abc import Iterable, Sized +from html import escape as htmlescape +from itertools import chain, zip_longest as izip_longest +from functools import reduce, partial +import io +import re +import math +import textwrap +import dataclasses + +try: + import wcwidth # optional wide-character (CJK) support +except ImportError: + wcwidth = None + + +def _is_file(f): + return isinstance(f, io.IOBase) + + +__all__ = ["tabulate", "tabulate_formats", "simple_separated_format"] +try: + from .version import version as __version__ # noqa: F401 +except ImportError: + pass # running __init__.py as a script, AppVeyor pytests + + +# minimum extra space in headers +MIN_PADDING = 2 + +# Whether or not to preserve leading/trailing whitespace in data. +PRESERVE_WHITESPACE = False + +_DEFAULT_FLOATFMT = "g" +_DEFAULT_INTFMT = "" +_DEFAULT_MISSINGVAL = "" +# default align will be overwritten by "left", "center" or "decimal" +# depending on the formatter +_DEFAULT_ALIGN = "default" + + +# if True, enable wide-character (CJK) support +WIDE_CHARS_MODE = wcwidth is not None + +# Constant that can be used as part of passed rows to generate a separating line +# It is purposely an unprintable character, very unlikely to be used in a table +SEPARATING_LINE = "\001" + +Line = namedtuple("Line", ["begin", "hline", "sep", "end"]) + + +DataRow = namedtuple("DataRow", ["begin", "sep", "end"]) + + +# A table structure is supposed to be: +# +# --- lineabove --------- +# headerrow +# --- linebelowheader --- +# datarow +# --- linebetweenrows --- +# ... (more datarows) ... +# --- linebetweenrows --- +# last datarow +# --- linebelow --------- +# +# TableFormat's line* elements can be +# +# - either None, if the element is not used, +# - or a Line tuple, +# - or a function: [col_widths], [col_alignments] -> string. +# +# TableFormat's *row elements can be +# +# - either None, if the element is not used, +# - or a DataRow tuple, +# - or a function: [cell_values], [col_widths], [col_alignments] -> string. +# +# padding (an integer) is the amount of white space around data values. +# +# with_header_hide: +# +# - either None, to display all table elements unconditionally, +# - or a list of elements not to be displayed if the table has column headers. +# +TableFormat = namedtuple( + "TableFormat", + [ + "lineabove", + "linebelowheader", + "linebetweenrows", + "linebelow", + "headerrow", + "datarow", + "padding", + "with_header_hide", + ], +) + + +def _is_separating_line(row): + row_type = type(row) + is_sl = (row_type == list or row_type == str) and ( + (len(row) >= 1 and row[0] == SEPARATING_LINE) + or (len(row) >= 2 and row[1] == SEPARATING_LINE) + ) + return is_sl + + +def _pipe_segment_with_colons(align, colwidth): + """Return a segment of a horizontal line with optional colons which + indicate column's alignment (as in `pipe` output format).""" + w = colwidth + if align in ["right", "decimal"]: + return ("-" * (w - 1)) + ":" + elif align == "center": + return ":" + ("-" * (w - 2)) + ":" + elif align == "left": + return ":" + ("-" * (w - 1)) + else: + return "-" * w + + +def _pipe_line_with_colons(colwidths, colaligns): + """Return a horizontal line with optional colons to indicate column's + alignment (as in `pipe` output format).""" + if not colaligns: # e.g. printing an empty data frame (github issue #15) + colaligns = [""] * len(colwidths) + segments = [_pipe_segment_with_colons(a, w) for a, w in zip(colaligns, colwidths)] + return "|" + "|".join(segments) + "|" + + +def _mediawiki_row_with_attrs(separator, cell_values, colwidths, colaligns): + alignment = { + "left": "", + "right": 'align="right"| ', + "center": 'align="center"| ', + "decimal": 'align="right"| ', + } + # hard-coded padding _around_ align attribute and value together + # rather than padding parameter which affects only the value + values_with_attrs = [ + " " + alignment.get(a, "") + c + " " for c, a in zip(cell_values, colaligns) + ] + colsep = separator * 2 + return (separator + colsep.join(values_with_attrs)).rstrip() + + +def _textile_row_with_attrs(cell_values, colwidths, colaligns): + cell_values[0] += " " + alignment = {"left": "<.", "right": ">.", "center": "=.", "decimal": ">."} + values = (alignment.get(a, "") + v for a, v in zip(colaligns, cell_values)) + return "|" + "|".join(values) + "|" + + +def _html_begin_table_without_header(colwidths_ignore, colaligns_ignore): + # this table header will be suppressed if there is a header row + return "\n" + + +def _html_row_with_attrs(celltag, unsafe, cell_values, colwidths, colaligns): + alignment = { + "left": "", + "right": ' style="text-align: right;"', + "center": ' style="text-align: center;"', + "decimal": ' style="text-align: right;"', + } + if unsafe: + values_with_attrs = [ + "<{0}{1}>{2}".format(celltag, alignment.get(a, ""), c) + for c, a in zip(cell_values, colaligns) + ] + else: + values_with_attrs = [ + "<{0}{1}>{2}".format(celltag, alignment.get(a, ""), htmlescape(c)) + for c, a in zip(cell_values, colaligns) + ] + rowhtml = "{}".format("".join(values_with_attrs).rstrip()) + if celltag == "th": # it's a header row, create a new table header + rowhtml = f"
\n\n{rowhtml}\n\n" + return rowhtml + + +def _moin_row_with_attrs(celltag, cell_values, colwidths, colaligns, header=""): + alignment = { + "left": "", + "right": '', + "center": '', + "decimal": '', + } + values_with_attrs = [ + "{}{} {} ".format(celltag, alignment.get(a, ""), header + c + header) + for c, a in zip(cell_values, colaligns) + ] + return "".join(values_with_attrs) + "||" + + +def _latex_line_begin_tabular(colwidths, colaligns, booktabs=False, longtable=False): + alignment = {"left": "l", "right": "r", "center": "c", "decimal": "r"} + tabular_columns_fmt = "".join([alignment.get(a, "l") for a in colaligns]) + return "\n".join( + [ + ("\\begin{tabular}{" if not longtable else "\\begin{longtable}{") + + tabular_columns_fmt + + "}", + "\\toprule" if booktabs else "\\hline", + ] + ) + + +def _asciidoc_row(is_header, *args): + """handle header and data rows for asciidoc format""" + + def make_header_line(is_header, colwidths, colaligns): + # generate the column specifiers + + alignment = {"left": "<", "right": ">", "center": "^", "decimal": ">"} + # use the column widths generated by tabulate for the asciidoc column width specifiers + asciidoc_alignments = zip( + colwidths, [alignment[colalign] for colalign in colaligns] + ) + asciidoc_column_specifiers = [ + "{:d}{}".format(width, align) for width, align in asciidoc_alignments + ] + header_list = ['cols="' + (",".join(asciidoc_column_specifiers)) + '"'] + + # generate the list of options (currently only "header") + options_list = [] + + if is_header: + options_list.append("header") + + if options_list: + header_list += ['options="' + ",".join(options_list) + '"'] + + # generate the list of entries in the table header field + + return "[{}]\n|====".format(",".join(header_list)) + + if len(args) == 2: + # two arguments are passed if called in the context of aboveline + # print the table header with column widths and optional header tag + return make_header_line(False, *args) + + elif len(args) == 3: + # three arguments are passed if called in the context of dataline or headerline + # print the table line and make the aboveline if it is a header + + cell_values, colwidths, colaligns = args + data_line = "|" + "|".join(cell_values) + + if is_header: + return make_header_line(True, colwidths, colaligns) + "\n" + data_line + else: + return data_line + + else: + raise ValueError( + " _asciidoc_row() requires two (colwidths, colaligns) " + + "or three (cell_values, colwidths, colaligns) arguments) " + ) + + +LATEX_ESCAPE_RULES = { + r"&": r"\&", + r"%": r"\%", + r"$": r"\$", + r"#": r"\#", + r"_": r"\_", + r"^": r"\^{}", + r"{": r"\{", + r"}": r"\}", + r"~": r"\textasciitilde{}", + "\\": r"\textbackslash{}", + r"<": r"\ensuremath{<}", + r">": r"\ensuremath{>}", +} + + +def _latex_row(cell_values, colwidths, colaligns, escrules=LATEX_ESCAPE_RULES): + def escape_char(c): + return escrules.get(c, c) + + escaped_values = ["".join(map(escape_char, cell)) for cell in cell_values] + rowfmt = DataRow("", "&", "\\\\") + return _build_simple_row(escaped_values, rowfmt) + + +def _rst_escape_first_column(rows, headers): + def escape_empty(val): + if isinstance(val, (str, bytes)) and not val.strip(): + return ".." + else: + return val + + new_headers = list(headers) + new_rows = [] + if headers: + new_headers[0] = escape_empty(headers[0]) + for row in rows: + new_row = list(row) + if new_row: + new_row[0] = escape_empty(row[0]) + new_rows.append(new_row) + return new_rows, new_headers + + +_table_formats = { + "simple": TableFormat( + lineabove=Line("", "-", " ", ""), + linebelowheader=Line("", "-", " ", ""), + linebetweenrows=None, + linebelow=Line("", "-", " ", ""), + headerrow=DataRow("", " ", ""), + datarow=DataRow("", " ", ""), + padding=0, + with_header_hide=["lineabove", "linebelow"], + ), + "plain": TableFormat( + lineabove=None, + linebelowheader=None, + linebetweenrows=None, + linebelow=None, + headerrow=DataRow("", " ", ""), + datarow=DataRow("", " ", ""), + padding=0, + with_header_hide=None, + ), + "grid": TableFormat( + lineabove=Line("+", "-", "+", "+"), + linebelowheader=Line("+", "=", "+", "+"), + linebetweenrows=Line("+", "-", "+", "+"), + linebelow=Line("+", "-", "+", "+"), + headerrow=DataRow("|", "|", "|"), + datarow=DataRow("|", "|", "|"), + padding=1, + with_header_hide=None, + ), + "simple_grid": TableFormat( + lineabove=Line("┌", "─", "┬", "┐"), + linebelowheader=Line("├", "─", "┼", "┤"), + linebetweenrows=Line("├", "─", "┼", "┤"), + linebelow=Line("└", "─", "┴", "┘"), + headerrow=DataRow("│", "│", "│"), + datarow=DataRow("│", "│", "│"), + padding=1, + with_header_hide=None, + ), + "rounded_grid": TableFormat( + lineabove=Line("╭", "─", "┬", "╮"), + linebelowheader=Line("├", "─", "┼", "┤"), + linebetweenrows=Line("├", "─", "┼", "┤"), + linebelow=Line("╰", "─", "┴", "╯"), + headerrow=DataRow("│", "│", "│"), + datarow=DataRow("│", "│", "│"), + padding=1, + with_header_hide=None, + ), + "heavy_grid": TableFormat( + lineabove=Line("┏", "━", "┳", "┓"), + linebelowheader=Line("┣", "━", "╋", "┫"), + linebetweenrows=Line("┣", "━", "╋", "┫"), + linebelow=Line("┗", "━", "┻", "┛"), + headerrow=DataRow("┃", "┃", "┃"), + datarow=DataRow("┃", "┃", "┃"), + padding=1, + with_header_hide=None, + ), + "mixed_grid": TableFormat( + lineabove=Line("┍", "━", "┯", "┑"), + linebelowheader=Line("┝", "━", "┿", "┥"), + linebetweenrows=Line("├", "─", "┼", "┤"), + linebelow=Line("┕", "━", "┷", "┙"), + headerrow=DataRow("│", "│", "│"), + datarow=DataRow("│", "│", "│"), + padding=1, + with_header_hide=None, + ), + "double_grid": TableFormat( + lineabove=Line("╔", "═", "╦", "╗"), + linebelowheader=Line("╠", "═", "╬", "╣"), + linebetweenrows=Line("╠", "═", "╬", "╣"), + linebelow=Line("╚", "═", "╩", "╝"), + headerrow=DataRow("║", "║", "║"), + datarow=DataRow("║", "║", "║"), + padding=1, + with_header_hide=None, + ), + "fancy_grid": TableFormat( + lineabove=Line("╒", "═", "╤", "╕"), + linebelowheader=Line("╞", "═", "╪", "╡"), + linebetweenrows=Line("├", "─", "┼", "┤"), + linebelow=Line("╘", "═", "╧", "╛"), + headerrow=DataRow("│", "│", "│"), + datarow=DataRow("│", "│", "│"), + padding=1, + with_header_hide=None, + ), + "outline": TableFormat( + lineabove=Line("+", "-", "+", "+"), + linebelowheader=Line("+", "=", "+", "+"), + linebetweenrows=None, + linebelow=Line("+", "-", "+", "+"), + headerrow=DataRow("|", "|", "|"), + datarow=DataRow("|", "|", "|"), + padding=1, + with_header_hide=None, + ), + "simple_outline": TableFormat( + lineabove=Line("┌", "─", "┬", "┐"), + linebelowheader=Line("├", "─", "┼", "┤"), + linebetweenrows=None, + linebelow=Line("└", "─", "┴", "┘"), + headerrow=DataRow("│", "│", "│"), + datarow=DataRow("│", "│", "│"), + padding=1, + with_header_hide=None, + ), + "rounded_outline": TableFormat( + lineabove=Line("╭", "─", "┬", "╮"), + linebelowheader=Line("├", "─", "┼", "┤"), + linebetweenrows=None, + linebelow=Line("╰", "─", "┴", "╯"), + headerrow=DataRow("│", "│", "│"), + datarow=DataRow("│", "│", "│"), + padding=1, + with_header_hide=None, + ), + "heavy_outline": TableFormat( + lineabove=Line("┏", "━", "┳", "┓"), + linebelowheader=Line("┣", "━", "╋", "┫"), + linebetweenrows=None, + linebelow=Line("┗", "━", "┻", "┛"), + headerrow=DataRow("┃", "┃", "┃"), + datarow=DataRow("┃", "┃", "┃"), + padding=1, + with_header_hide=None, + ), + "mixed_outline": TableFormat( + lineabove=Line("┍", "━", "┯", "┑"), + linebelowheader=Line("┝", "━", "┿", "┥"), + linebetweenrows=None, + linebelow=Line("┕", "━", "┷", "┙"), + headerrow=DataRow("│", "│", "│"), + datarow=DataRow("│", "│", "│"), + padding=1, + with_header_hide=None, + ), + "double_outline": TableFormat( + lineabove=Line("╔", "═", "╦", "╗"), + linebelowheader=Line("╠", "═", "╬", "╣"), + linebetweenrows=None, + linebelow=Line("╚", "═", "╩", "╝"), + headerrow=DataRow("║", "║", "║"), + datarow=DataRow("║", "║", "║"), + padding=1, + with_header_hide=None, + ), + "fancy_outline": TableFormat( + lineabove=Line("╒", "═", "╤", "╕"), + linebelowheader=Line("╞", "═", "╪", "╡"), + linebetweenrows=None, + linebelow=Line("╘", "═", "╧", "╛"), + headerrow=DataRow("│", "│", "│"), + datarow=DataRow("│", "│", "│"), + padding=1, + with_header_hide=None, + ), + "github": TableFormat( + lineabove=Line("|", "-", "|", "|"), + linebelowheader=Line("|", "-", "|", "|"), + linebetweenrows=None, + linebelow=None, + headerrow=DataRow("|", "|", "|"), + datarow=DataRow("|", "|", "|"), + padding=1, + with_header_hide=["lineabove"], + ), + "pipe": TableFormat( + lineabove=_pipe_line_with_colons, + linebelowheader=_pipe_line_with_colons, + linebetweenrows=None, + linebelow=None, + headerrow=DataRow("|", "|", "|"), + datarow=DataRow("|", "|", "|"), + padding=1, + with_header_hide=["lineabove"], + ), + "orgtbl": TableFormat( + lineabove=None, + linebelowheader=Line("|", "-", "+", "|"), + linebetweenrows=None, + linebelow=None, + headerrow=DataRow("|", "|", "|"), + datarow=DataRow("|", "|", "|"), + padding=1, + with_header_hide=None, + ), + "jira": TableFormat( + lineabove=None, + linebelowheader=None, + linebetweenrows=None, + linebelow=None, + headerrow=DataRow("||", "||", "||"), + datarow=DataRow("|", "|", "|"), + padding=1, + with_header_hide=None, + ), + "presto": TableFormat( + lineabove=None, + linebelowheader=Line("", "-", "+", ""), + linebetweenrows=None, + linebelow=None, + headerrow=DataRow("", "|", ""), + datarow=DataRow("", "|", ""), + padding=1, + with_header_hide=None, + ), + "pretty": TableFormat( + lineabove=Line("+", "-", "+", "+"), + linebelowheader=Line("+", "-", "+", "+"), + linebetweenrows=None, + linebelow=Line("+", "-", "+", "+"), + headerrow=DataRow("|", "|", "|"), + datarow=DataRow("|", "|", "|"), + padding=1, + with_header_hide=None, + ), + "psql": TableFormat( + lineabove=Line("+", "-", "+", "+"), + linebelowheader=Line("|", "-", "+", "|"), + linebetweenrows=None, + linebelow=Line("+", "-", "+", "+"), + headerrow=DataRow("|", "|", "|"), + datarow=DataRow("|", "|", "|"), + padding=1, + with_header_hide=None, + ), + "rst": TableFormat( + lineabove=Line("", "=", " ", ""), + linebelowheader=Line("", "=", " ", ""), + linebetweenrows=None, + linebelow=Line("", "=", " ", ""), + headerrow=DataRow("", " ", ""), + datarow=DataRow("", " ", ""), + padding=0, + with_header_hide=None, + ), + "mediawiki": TableFormat( + lineabove=Line( + '{| class="wikitable" style="text-align: left;"', + "", + "", + "\n|+ \n|-", + ), + linebelowheader=Line("|-", "", "", ""), + linebetweenrows=Line("|-", "", "", ""), + linebelow=Line("|}", "", "", ""), + headerrow=partial(_mediawiki_row_with_attrs, "!"), + datarow=partial(_mediawiki_row_with_attrs, "|"), + padding=0, + with_header_hide=None, + ), + "moinmoin": TableFormat( + lineabove=None, + linebelowheader=None, + linebetweenrows=None, + linebelow=None, + headerrow=partial(_moin_row_with_attrs, "||", header="'''"), + datarow=partial(_moin_row_with_attrs, "||"), + padding=1, + with_header_hide=None, + ), + "youtrack": TableFormat( + lineabove=None, + linebelowheader=None, + linebetweenrows=None, + linebelow=None, + headerrow=DataRow("|| ", " || ", " || "), + datarow=DataRow("| ", " | ", " |"), + padding=1, + with_header_hide=None, + ), + "html": TableFormat( + lineabove=_html_begin_table_without_header, + linebelowheader="", + linebetweenrows=None, + linebelow=Line("\n
", "", "", ""), + headerrow=partial(_html_row_with_attrs, "th", False), + datarow=partial(_html_row_with_attrs, "td", False), + padding=0, + with_header_hide=["lineabove"], + ), + "unsafehtml": TableFormat( + lineabove=_html_begin_table_without_header, + linebelowheader="", + linebetweenrows=None, + linebelow=Line("\n", "", "", ""), + headerrow=partial(_html_row_with_attrs, "th", True), + datarow=partial(_html_row_with_attrs, "td", True), + padding=0, + with_header_hide=["lineabove"], + ), + "latex": TableFormat( + lineabove=_latex_line_begin_tabular, + linebelowheader=Line("\\hline", "", "", ""), + linebetweenrows=None, + linebelow=Line("\\hline\n\\end{tabular}", "", "", ""), + headerrow=_latex_row, + datarow=_latex_row, + padding=1, + with_header_hide=None, + ), + "latex_raw": TableFormat( + lineabove=_latex_line_begin_tabular, + linebelowheader=Line("\\hline", "", "", ""), + linebetweenrows=None, + linebelow=Line("\\hline\n\\end{tabular}", "", "", ""), + headerrow=partial(_latex_row, escrules={}), + datarow=partial(_latex_row, escrules={}), + padding=1, + with_header_hide=None, + ), + "latex_booktabs": TableFormat( + lineabove=partial(_latex_line_begin_tabular, booktabs=True), + linebelowheader=Line("\\midrule", "", "", ""), + linebetweenrows=None, + linebelow=Line("\\bottomrule\n\\end{tabular}", "", "", ""), + headerrow=_latex_row, + datarow=_latex_row, + padding=1, + with_header_hide=None, + ), + "latex_longtable": TableFormat( + lineabove=partial(_latex_line_begin_tabular, longtable=True), + linebelowheader=Line("\\hline\n\\endhead", "", "", ""), + linebetweenrows=None, + linebelow=Line("\\hline\n\\end{longtable}", "", "", ""), + headerrow=_latex_row, + datarow=_latex_row, + padding=1, + with_header_hide=None, + ), + "tsv": TableFormat( + lineabove=None, + linebelowheader=None, + linebetweenrows=None, + linebelow=None, + headerrow=DataRow("", "\t", ""), + datarow=DataRow("", "\t", ""), + padding=0, + with_header_hide=None, + ), + "textile": TableFormat( + lineabove=None, + linebelowheader=None, + linebetweenrows=None, + linebelow=None, + headerrow=DataRow("|_. ", "|_.", "|"), + datarow=_textile_row_with_attrs, + padding=1, + with_header_hide=None, + ), + "asciidoc": TableFormat( + lineabove=partial(_asciidoc_row, False), + linebelowheader=None, + linebetweenrows=None, + linebelow=Line("|====", "", "", ""), + headerrow=partial(_asciidoc_row, True), + datarow=partial(_asciidoc_row, False), + padding=1, + with_header_hide=["lineabove"], + ), +} + + +tabulate_formats = list(sorted(_table_formats.keys())) + +# The table formats for which multiline cells will be folded into subsequent +# table rows. The key is the original format specified at the API. The value is +# the format that will be used to represent the original format. +multiline_formats = { + "plain": "plain", + "simple": "simple", + "grid": "grid", + "simple_grid": "simple_grid", + "rounded_grid": "rounded_grid", + "heavy_grid": "heavy_grid", + "mixed_grid": "mixed_grid", + "double_grid": "double_grid", + "fancy_grid": "fancy_grid", + "pipe": "pipe", + "orgtbl": "orgtbl", + "jira": "jira", + "presto": "presto", + "pretty": "pretty", + "psql": "psql", + "rst": "rst", +} + +# TODO: Add multiline support for the remaining table formats: +# - mediawiki: Replace \n with
+# - moinmoin: TBD +# - youtrack: TBD +# - html: Replace \n with
+# - latex*: Use "makecell" package: In header, replace X\nY with +# \thead{X\\Y} and in data row, replace X\nY with \makecell{X\\Y} +# - tsv: TBD +# - textile: Replace \n with
(must be well-formed XML) + +_multiline_codes = re.compile(r"\r|\n|\r\n") +_multiline_codes_bytes = re.compile(b"\r|\n|\r\n") + +# Handle ANSI escape sequences for both control sequence introducer (CSI) and +# operating system command (OSC). Both of these begin with 0x1b (or octal 033), +# which will be shown below as ESC. +# +# CSI ANSI escape codes have the following format, defined in section 5.4 of ECMA-48: +# +# CSI: ESC followed by the '[' character (0x5b) +# Parameter Bytes: 0..n bytes in the range 0x30-0x3f +# Intermediate Bytes: 0..n bytes in the range 0x20-0x2f +# Final Byte: a single byte in the range 0x40-0x7e +# +# Also include the terminal hyperlink sequences as described here: +# https://gist.github.com/egmontkob/eb114294efbcd5adb1944c9f3cb5feda +# +# OSC 8 ; params ; uri ST display_text OSC 8 ;; ST +# +# Example: \x1b]8;;https://example.com\x5ctext to show\x1b]8;;\x5c +# +# Where: +# OSC: ESC followed by the ']' character (0x5d) +# params: 0..n optional key value pairs separated by ':' (e.g. foo=bar:baz=qux:abc=123) +# URI: the actual URI with protocol scheme (e.g. https://, file://, ftp://) +# ST: ESC followed by the '\' character (0x5c) +_esc = r"\x1b" +_csi = rf"{_esc}\[" +_osc = rf"{_esc}\]" +_st = rf"{_esc}\\" + +_ansi_escape_pat = rf""" + ( + # terminal colors, etc + {_csi} # CSI + [\x30-\x3f]* # parameter bytes + [\x20-\x2f]* # intermediate bytes + [\x40-\x7e] # final byte + | + # terminal hyperlinks + {_osc}8; # OSC opening + (\w+=\w+:?)* # key=value params list (submatch 2) + ; # delimiter + ([^{_esc}]+) # URI - anything but ESC (submatch 3) + {_st} # ST + ([^{_esc}]+) # link text - anything but ESC (submatch 4) + {_osc}8;;{_st} # "closing" OSC sequence + ) +""" +_ansi_codes = re.compile(_ansi_escape_pat, re.VERBOSE) +_ansi_codes_bytes = re.compile(_ansi_escape_pat.encode("utf8"), re.VERBOSE) +_ansi_color_reset_code = "\033[0m" + +_float_with_thousands_separators = re.compile( + r"^(([+-]?[0-9]{1,3})(?:,([0-9]{3}))*)?(?(1)\.[0-9]*|\.[0-9]+)?$" +) + + +def simple_separated_format(separator): + """Construct a simple TableFormat with columns separated by a separator. + + >>> tsv = simple_separated_format("\\t") ; \ + tabulate([["foo", 1], ["spam", 23]], tablefmt=tsv) == 'foo \\t 1\\nspam\\t23' + True + + """ + return TableFormat( + None, + None, + None, + None, + headerrow=DataRow("", separator, ""), + datarow=DataRow("", separator, ""), + padding=0, + with_header_hide=None, + ) + + +def _isnumber_with_thousands_separator(string): + """ + >>> _isnumber_with_thousands_separator(".") + False + >>> _isnumber_with_thousands_separator("1") + True + >>> _isnumber_with_thousands_separator("1.") + True + >>> _isnumber_with_thousands_separator(".1") + True + >>> _isnumber_with_thousands_separator("1000") + False + >>> _isnumber_with_thousands_separator("1,000") + True + >>> _isnumber_with_thousands_separator("1,0000") + False + >>> _isnumber_with_thousands_separator("1,000.1234") + True + >>> _isnumber_with_thousands_separator(b"1,000.1234") + True + >>> _isnumber_with_thousands_separator("+1,000.1234") + True + >>> _isnumber_with_thousands_separator("-1,000.1234") + True + """ + try: + string = string.decode() + except (UnicodeDecodeError, AttributeError): + pass + + return bool(re.match(_float_with_thousands_separators, string)) + + +def _isconvertible(conv, string): + try: + conv(string) + return True + except (ValueError, TypeError): + return False + + +def _isnumber(string): + """ + >>> _isnumber("123.45") + True + >>> _isnumber("123") + True + >>> _isnumber("spam") + False + >>> _isnumber("123e45678") + False + >>> _isnumber("inf") + True + """ + if not _isconvertible(float, string): + return False + elif isinstance(string, (str, bytes)) and ( + math.isinf(float(string)) or math.isnan(float(string)) + ): + return string.lower() in ["inf", "-inf", "nan"] + return True + + +def _isint(string, inttype=int): + """ + >>> _isint("123") + True + >>> _isint("123.45") + False + """ + return ( + type(string) is inttype + or isinstance(string, (bytes, str)) + and _isconvertible(inttype, string) + ) + + +def _isbool(string): + """ + >>> _isbool(True) + True + >>> _isbool("False") + True + >>> _isbool(1) + False + """ + return type(string) is bool or ( + isinstance(string, (bytes, str)) and string in ("True", "False") + ) + + +def _type(string, has_invisible=True, numparse=True): + """The least generic type (type(None), int, float, str, unicode). + + >>> _type(None) is type(None) + True + >>> _type("foo") is type("") + True + >>> _type("1") is type(1) + True + >>> _type('\x1b[31m42\x1b[0m') is type(42) + True + >>> _type('\x1b[31m42\x1b[0m') is type(42) + True + + """ + + if has_invisible and isinstance(string, (str, bytes)): + string = _strip_ansi(string) + + if string is None: + return type(None) + elif hasattr(string, "isoformat"): # datetime.datetime, date, and time + return str + elif _isbool(string): + return bool + elif _isint(string) and numparse: + return int + elif _isnumber(string) and numparse: + return float + elif isinstance(string, bytes): + return bytes + else: + return str + + +def _afterpoint(string): + """Symbols after a decimal point, -1 if the string lacks the decimal point. + + >>> _afterpoint("123.45") + 2 + >>> _afterpoint("1001") + -1 + >>> _afterpoint("eggs") + -1 + >>> _afterpoint("123e45") + 2 + >>> _afterpoint("123,456.78") + 2 + + """ + if _isnumber(string) or _isnumber_with_thousands_separator(string): + if _isint(string): + return -1 + else: + pos = string.rfind(".") + pos = string.lower().rfind("e") if pos < 0 else pos + if pos >= 0: + return len(string) - pos - 1 + else: + return -1 # no point + else: + return -1 # not a number + + +def _padleft(width, s): + """Flush right. + + >>> _padleft(6, '\u044f\u0439\u0446\u0430') == ' \u044f\u0439\u0446\u0430' + True + + """ + fmt = "{0:>%ds}" % width + return fmt.format(s) + + +def _padright(width, s): + """Flush left. + + >>> _padright(6, '\u044f\u0439\u0446\u0430') == '\u044f\u0439\u0446\u0430 ' + True + + """ + fmt = "{0:<%ds}" % width + return fmt.format(s) + + +def _padboth(width, s): + """Center string. + + >>> _padboth(6, '\u044f\u0439\u0446\u0430') == ' \u044f\u0439\u0446\u0430 ' + True + + """ + fmt = "{0:^%ds}" % width + return fmt.format(s) + + +def _padnone(ignore_width, s): + return s + + +def _strip_ansi(s): + r"""Remove ANSI escape sequences, both CSI (color codes, etc) and OSC hyperlinks. + + CSI sequences are simply removed from the output, while OSC hyperlinks are replaced + with the link text. Note: it may be desirable to show the URI instead but this is not + supported. + + >>> repr(_strip_ansi('\x1B]8;;https://example.com\x1B\\This is a link\x1B]8;;\x1B\\')) + "'This is a link'" + + >>> repr(_strip_ansi('\x1b[31mred\x1b[0m text')) + "'red text'" + + """ + if isinstance(s, str): + return _ansi_codes.sub(r"\4", s) + else: # a bytestring + return _ansi_codes_bytes.sub(r"\4", s) + + +def _visible_width(s): + """Visible width of a printed string. ANSI color codes are removed. + + >>> _visible_width('\x1b[31mhello\x1b[0m'), _visible_width("world") + (5, 5) + + """ + # optional wide-character support + if wcwidth is not None and WIDE_CHARS_MODE: + len_fn = wcwidth.wcswidth + else: + len_fn = len + if isinstance(s, (str, bytes)): + return len_fn(_strip_ansi(s)) + else: + return len_fn(str(s)) + + +def _is_multiline(s): + if isinstance(s, str): + return bool(re.search(_multiline_codes, s)) + else: # a bytestring + return bool(re.search(_multiline_codes_bytes, s)) + + +def _multiline_width(multiline_s, line_width_fn=len): + """Visible width of a potentially multiline content.""" + return max(map(line_width_fn, re.split("[\r\n]", multiline_s))) + + +def _choose_width_fn(has_invisible, enable_widechars, is_multiline): + """Return a function to calculate visible cell width.""" + if has_invisible: + line_width_fn = _visible_width + elif enable_widechars: # optional wide-character support if available + line_width_fn = wcwidth.wcswidth + else: + line_width_fn = len + if is_multiline: + width_fn = lambda s: _multiline_width(s, line_width_fn) # noqa + else: + width_fn = line_width_fn + return width_fn + + +def _align_column_choose_padfn(strings, alignment, has_invisible): + if alignment == "right": + if not PRESERVE_WHITESPACE: + strings = [s.strip() for s in strings] + padfn = _padleft + elif alignment == "center": + if not PRESERVE_WHITESPACE: + strings = [s.strip() for s in strings] + padfn = _padboth + elif alignment == "decimal": + if has_invisible: + decimals = [_afterpoint(_strip_ansi(s)) for s in strings] + else: + decimals = [_afterpoint(s) for s in strings] + maxdecimals = max(decimals) + strings = [s + (maxdecimals - decs) * " " for s, decs in zip(strings, decimals)] + padfn = _padleft + elif not alignment: + padfn = _padnone + else: + if not PRESERVE_WHITESPACE: + strings = [s.strip() for s in strings] + padfn = _padright + return strings, padfn + + +def _align_column_choose_width_fn(has_invisible, enable_widechars, is_multiline): + if has_invisible: + line_width_fn = _visible_width + elif enable_widechars: # optional wide-character support if available + line_width_fn = wcwidth.wcswidth + else: + line_width_fn = len + if is_multiline: + width_fn = lambda s: _align_column_multiline_width(s, line_width_fn) # noqa + else: + width_fn = line_width_fn + return width_fn + + +def _align_column_multiline_width(multiline_s, line_width_fn=len): + """Visible width of a potentially multiline content.""" + return list(map(line_width_fn, re.split("[\r\n]", multiline_s))) + + +def _flat_list(nested_list): + ret = [] + for item in nested_list: + if isinstance(item, list): + for subitem in item: + ret.append(subitem) + else: + ret.append(item) + return ret + + +def _align_column( + strings, + alignment, + minwidth=0, + has_invisible=True, + enable_widechars=False, + is_multiline=False, +): + """[string] -> [padded_string]""" + strings, padfn = _align_column_choose_padfn(strings, alignment, has_invisible) + width_fn = _align_column_choose_width_fn( + has_invisible, enable_widechars, is_multiline + ) + + s_widths = list(map(width_fn, strings)) + maxwidth = max(max(_flat_list(s_widths)), minwidth) + # TODO: refactor column alignment in single-line and multiline modes + if is_multiline: + if not enable_widechars and not has_invisible: + padded_strings = [ + "\n".join([padfn(maxwidth, s) for s in ms.splitlines()]) + for ms in strings + ] + else: + # enable wide-character width corrections + s_lens = [[len(s) for s in re.split("[\r\n]", ms)] for ms in strings] + visible_widths = [ + [maxwidth - (w - l) for w, l in zip(mw, ml)] + for mw, ml in zip(s_widths, s_lens) + ] + # wcswidth and _visible_width don't count invisible characters; + # padfn doesn't need to apply another correction + padded_strings = [ + "\n".join([padfn(w, s) for s, w in zip((ms.splitlines() or ms), mw)]) + for ms, mw in zip(strings, visible_widths) + ] + else: # single-line cell values + if not enable_widechars and not has_invisible: + padded_strings = [padfn(maxwidth, s) for s in strings] + else: + # enable wide-character width corrections + s_lens = list(map(len, strings)) + visible_widths = [maxwidth - (w - l) for w, l in zip(s_widths, s_lens)] + # wcswidth and _visible_width don't count invisible characters; + # padfn doesn't need to apply another correction + padded_strings = [padfn(w, s) for s, w in zip(strings, visible_widths)] + return padded_strings + + +def _more_generic(type1, type2): + types = { + type(None): 0, # noqa + bool: 1, + int: 2, + float: 3, + bytes: 4, + str: 5, + } + invtypes = { + 5: str, + 4: bytes, + 3: float, + 2: int, + 1: bool, + 0: type(None), + } + moregeneric = max(types.get(type1, 5), types.get(type2, 5)) + return invtypes[moregeneric] + + +def _column_type(strings, has_invisible=True, numparse=True): + """The least generic type all column values are convertible to. + + >>> _column_type([True, False]) is bool + True + >>> _column_type(["1", "2"]) is int + True + >>> _column_type(["1", "2.3"]) is float + True + >>> _column_type(["1", "2.3", "four"]) is str + True + >>> _column_type(["four", '\u043f\u044f\u0442\u044c']) is str + True + >>> _column_type([None, "brux"]) is str + True + >>> _column_type([1, 2, None]) is int + True + >>> import datetime as dt + >>> _column_type([dt.datetime(1991,2,19), dt.time(17,35)]) is str + True + + """ + types = [_type(s, has_invisible, numparse) for s in strings] + return reduce(_more_generic, types, bool) + + +def _format(val, valtype, floatfmt, intfmt, missingval="", has_invisible=True): + """Format a value according to its type. + + Unicode is supported: + + >>> hrow = ['\u0431\u0443\u043a\u0432\u0430', '\u0446\u0438\u0444\u0440\u0430'] ; \ + tbl = [['\u0430\u0437', 2], ['\u0431\u0443\u043a\u0438', 4]] ; \ + good_result = '\\u0431\\u0443\\u043a\\u0432\\u0430 \\u0446\\u0438\\u0444\\u0440\\u0430\\n------- -------\\n\\u0430\\u0437 2\\n\\u0431\\u0443\\u043a\\u0438 4' ; \ + tabulate(tbl, headers=hrow) == good_result + True + + """ # noqa + if val is None: + return missingval + + if valtype is str: + return f"{val}" + elif valtype is int: + return format(val, intfmt) + elif valtype is bytes: + try: + return str(val, "ascii") + except (TypeError, UnicodeDecodeError): + return str(val) + elif valtype is float: + is_a_colored_number = has_invisible and isinstance(val, (str, bytes)) + if is_a_colored_number: + raw_val = _strip_ansi(val) + formatted_val = format(float(raw_val), floatfmt) + return val.replace(raw_val, formatted_val) + else: + return format(float(val), floatfmt) + else: + return f"{val}" + + +def _align_header( + header, alignment, width, visible_width, is_multiline=False, width_fn=None +): + "Pad string header to width chars given known visible_width of the header." + if is_multiline: + header_lines = re.split(_multiline_codes, header) + padded_lines = [ + _align_header(h, alignment, width, width_fn(h)) for h in header_lines + ] + return "\n".join(padded_lines) + # else: not multiline + ninvisible = len(header) - visible_width + width += ninvisible + if alignment == "left": + return _padright(width, header) + elif alignment == "center": + return _padboth(width, header) + elif not alignment: + return f"{header}" + else: + return _padleft(width, header) + + +def _remove_separating_lines(rows): + if type(rows) == list: + separating_lines = [] + sans_rows = [] + for index, row in enumerate(rows): + if _is_separating_line(row): + separating_lines.append(index) + else: + sans_rows.append(row) + return sans_rows, separating_lines + else: + return rows, None + + +def _reinsert_separating_lines(rows, separating_lines): + if separating_lines: + for index in separating_lines: + rows.insert(index, SEPARATING_LINE) + + +def _prepend_row_index(rows, index): + """Add a left-most index column.""" + if index is None or index is False: + return rows + if isinstance(index, Sized) and len(index) != len(rows): + raise ValueError( + "index must be as long as the number of data rows: " + + "len(index)={} len(rows)={}".format(len(index), len(rows)) + ) + sans_rows, separating_lines = _remove_separating_lines(rows) + new_rows = [] + index_iter = iter(index) + for row in sans_rows: + index_v = next(index_iter) + new_rows.append([index_v] + list(row)) + rows = new_rows + _reinsert_separating_lines(rows, separating_lines) + return rows + + +def _bool(val): + "A wrapper around standard bool() which doesn't throw on NumPy arrays" + try: + return bool(val) + except ValueError: # val is likely to be a numpy array with many elements + return False + + +def _normalize_tabular_data(tabular_data, headers, showindex="default"): + """Transform a supported data type to a list of lists, and a list of headers. + + Supported tabular data types: + + * list-of-lists or another iterable of iterables + + * list of named tuples (usually used with headers="keys") + + * list of dicts (usually used with headers="keys") + + * list of OrderedDicts (usually used with headers="keys") + + * list of dataclasses (Python 3.7+ only, usually used with headers="keys") + + * 2D NumPy arrays + + * NumPy record arrays (usually used with headers="keys") + + * dict of iterables (usually used with headers="keys") + + * pandas.DataFrame (usually used with headers="keys") + + The first row can be used as headers if headers="firstrow", + column indices can be used as headers if headers="keys". + + If showindex="default", show row indices of the pandas.DataFrame. + If showindex="always", show row indices for all types of data. + If showindex="never", don't show row indices for all types of data. + If showindex is an iterable, show its values as row indices. + + """ + + try: + bool(headers) + is_headers2bool_broken = False # noqa + except ValueError: # numpy.ndarray, pandas.core.index.Index, ... + is_headers2bool_broken = True # noqa + headers = list(headers) + + index = None + if hasattr(tabular_data, "keys") and hasattr(tabular_data, "values"): + # dict-like and pandas.DataFrame? + if hasattr(tabular_data.values, "__call__"): + # likely a conventional dict + keys = tabular_data.keys() + rows = list( + izip_longest(*tabular_data.values()) + ) # columns have to be transposed + elif hasattr(tabular_data, "index"): + # values is a property, has .index => it's likely a pandas.DataFrame (pandas 0.11.0) + keys = list(tabular_data) + if ( + showindex in ["default", "always", True] + and tabular_data.index.name is not None + ): + if isinstance(tabular_data.index.name, list): + keys[:0] = tabular_data.index.name + else: + keys[:0] = [tabular_data.index.name] + vals = tabular_data.values # values matrix doesn't need to be transposed + # for DataFrames add an index per default + index = list(tabular_data.index) + rows = [list(row) for row in vals] + else: + raise ValueError("tabular data doesn't appear to be a dict or a DataFrame") + + if headers == "keys": + headers = list(map(str, keys)) # headers should be strings + + else: # it's a usual iterable of iterables, or a NumPy array, or an iterable of dataclasses + rows = list(tabular_data) + + if headers == "keys" and not rows: + # an empty table (issue #81) + headers = [] + elif ( + headers == "keys" + and hasattr(tabular_data, "dtype") + and getattr(tabular_data.dtype, "names") + ): + # numpy record array + headers = tabular_data.dtype.names + elif ( + headers == "keys" + and len(rows) > 0 + and isinstance(rows[0], tuple) + and hasattr(rows[0], "_fields") + ): + # namedtuple + headers = list(map(str, rows[0]._fields)) + elif len(rows) > 0 and hasattr(rows[0], "keys") and hasattr(rows[0], "values"): + # dict-like object + uniq_keys = set() # implements hashed lookup + keys = [] # storage for set + if headers == "firstrow": + firstdict = rows[0] if len(rows) > 0 else {} + keys.extend(firstdict.keys()) + uniq_keys.update(keys) + rows = rows[1:] + for row in rows: + for k in row.keys(): + # Save unique items in input order + if k not in uniq_keys: + keys.append(k) + uniq_keys.add(k) + if headers == "keys": + headers = keys + elif isinstance(headers, dict): + # a dict of headers for a list of dicts + headers = [headers.get(k, k) for k in keys] + headers = list(map(str, headers)) + elif headers == "firstrow": + if len(rows) > 0: + headers = [firstdict.get(k, k) for k in keys] + headers = list(map(str, headers)) + else: + headers = [] + elif headers: + raise ValueError( + "headers for a list of dicts is not a dict or a keyword" + ) + rows = [[row.get(k) for k in keys] for row in rows] + + elif ( + headers == "keys" + and hasattr(tabular_data, "description") + and hasattr(tabular_data, "fetchone") + and hasattr(tabular_data, "rowcount") + ): + # Python Database API cursor object (PEP 0249) + # print tabulate(cursor, headers='keys') + headers = [column[0] for column in tabular_data.description] + + elif ( + dataclasses is not None + and len(rows) > 0 + and dataclasses.is_dataclass(rows[0]) + ): + # Python 3.7+'s dataclass + field_names = [field.name for field in dataclasses.fields(rows[0])] + if headers == "keys": + headers = field_names + rows = [[getattr(row, f) for f in field_names] for row in rows] + + elif headers == "keys" and len(rows) > 0: + # keys are column indices + headers = list(map(str, range(len(rows[0])))) + + # take headers from the first row if necessary + if headers == "firstrow" and len(rows) > 0: + if index is not None: + headers = [index[0]] + list(rows[0]) + index = index[1:] + else: + headers = rows[0] + headers = list(map(str, headers)) # headers should be strings + rows = rows[1:] + elif headers == "firstrow": + headers = [] + + headers = list(map(str, headers)) + # rows = list(map(list, rows)) + rows = list(map(lambda r: r if _is_separating_line(r) else list(r), rows)) + + # add or remove an index column + showindex_is_a_str = type(showindex) in [str, bytes] + if showindex == "default" and index is not None: + rows = _prepend_row_index(rows, index) + elif isinstance(showindex, Sized) and not showindex_is_a_str: + rows = _prepend_row_index(rows, list(showindex)) + elif isinstance(showindex, Iterable) and not showindex_is_a_str: + rows = _prepend_row_index(rows, showindex) + elif showindex == "always" or (_bool(showindex) and not showindex_is_a_str): + if index is None: + index = list(range(len(rows))) + rows = _prepend_row_index(rows, index) + elif showindex == "never" or (not _bool(showindex) and not showindex_is_a_str): + pass + + # pad with empty headers for initial columns if necessary + if headers and len(rows) > 0: + nhs = len(headers) + ncols = len(rows[0]) + if nhs < ncols: + headers = [""] * (ncols - nhs) + headers + + return rows, headers + + +def _wrap_text_to_colwidths(list_of_lists, colwidths, numparses=True): + numparses = _expand_iterable(numparses, len(list_of_lists[0]), True) + + result = [] + + for row in list_of_lists: + new_row = [] + for cell, width, numparse in zip(row, colwidths, numparses): + if _isnumber(cell) and numparse: + new_row.append(cell) + continue + + if width is not None: + wrapper = _CustomTextWrap(width=width) + # Cast based on our internal type handling + # Any future custom formatting of types (such as datetimes) + # may need to be more explicit than just `str` of the object + casted_cell = ( + str(cell) if _isnumber(cell) else _type(cell, numparse)(cell) + ) + wrapped = wrapper.wrap(casted_cell) + new_row.append("\n".join(wrapped)) + else: + new_row.append(cell) + result.append(new_row) + + return result + + +def _to_str(s, encoding="utf8", errors="ignore"): + """ + A type safe wrapper for converting a bytestring to str. This is essentially just + a wrapper around .decode() intended for use with things like map(), but with some + specific behavior: + + 1. if the given parameter is not a bytestring, it is returned unmodified + 2. decode() is called for the given parameter and assumes utf8 encoding, but the + default error behavior is changed from 'strict' to 'ignore' + + >>> repr(_to_str(b'foo')) + "'foo'" + + >>> repr(_to_str('foo')) + "'foo'" + + >>> repr(_to_str(42)) + "'42'" + + """ + if isinstance(s, bytes): + return s.decode(encoding=encoding, errors=errors) + return str(s) + + +def tabulate( + tabular_data, + headers=(), + tablefmt="simple", + floatfmt=_DEFAULT_FLOATFMT, + intfmt=_DEFAULT_INTFMT, + numalign=_DEFAULT_ALIGN, + stralign=_DEFAULT_ALIGN, + missingval=_DEFAULT_MISSINGVAL, + showindex="default", + disable_numparse=False, + colalign=None, + maxcolwidths=None, + rowalign=None, + maxheadercolwidths=None, +): + """Format a fixed width table for pretty printing. + + >>> print(tabulate([[1, 2.34], [-56, "8.999"], ["2", "10001"]])) + --- --------- + 1 2.34 + -56 8.999 + 2 10001 + --- --------- + + The first required argument (`tabular_data`) can be a + list-of-lists (or another iterable of iterables), a list of named + tuples, a dictionary of iterables, an iterable of dictionaries, + an iterable of dataclasses (Python 3.7+), a two-dimensional NumPy array, + NumPy record array, or a Pandas' dataframe. + + + Table headers + ------------- + + To print nice column headers, supply the second argument (`headers`): + + - `headers` can be an explicit list of column headers + - if `headers="firstrow"`, then the first row of data is used + - if `headers="keys"`, then dictionary keys or column indices are used + + Otherwise a headerless table is produced. + + If the number of headers is less than the number of columns, they + are supposed to be names of the last columns. This is consistent + with the plain-text format of R and Pandas' dataframes. + + >>> print(tabulate([["sex","age"],["Alice","F",24],["Bob","M",19]], + ... headers="firstrow")) + sex age + ----- ----- ----- + Alice F 24 + Bob M 19 + + By default, pandas.DataFrame data have an additional column called + row index. To add a similar column to all other types of data, + use `showindex="always"` or `showindex=True`. To suppress row indices + for all types of data, pass `showindex="never" or `showindex=False`. + To add a custom row index column, pass `showindex=some_iterable`. + + >>> print(tabulate([["F",24],["M",19]], showindex="always")) + - - -- + 0 F 24 + 1 M 19 + - - -- + + + Column alignment + ---------------- + + `tabulate` tries to detect column types automatically, and aligns + the values properly. By default it aligns decimal points of the + numbers (or flushes integer numbers to the right), and flushes + everything else to the left. Possible column alignments + (`numalign`, `stralign`) are: "right", "center", "left", "decimal" + (only for `numalign`), and None (to disable alignment). + + + Table formats + ------------- + + `intfmt` is a format specification used for columns which + contain numeric data without a decimal point. This can also be + a list or tuple of format strings, one per column. + + `floatfmt` is a format specification used for columns which + contain numeric data with a decimal point. This can also be + a list or tuple of format strings, one per column. + + `None` values are replaced with a `missingval` string (like + `floatfmt`, this can also be a list of values for different + columns): + + >>> print(tabulate([["spam", 1, None], + ... ["eggs", 42, 3.14], + ... ["other", None, 2.7]], missingval="?")) + ----- -- ---- + spam 1 ? + eggs 42 3.14 + other ? 2.7 + ----- -- ---- + + Various plain-text table formats (`tablefmt`) are supported: + 'plain', 'simple', 'grid', 'pipe', 'orgtbl', 'rst', 'mediawiki', + 'latex', 'latex_raw', 'latex_booktabs', 'latex_longtable' and tsv. + Variable `tabulate_formats`contains the list of currently supported formats. + + "plain" format doesn't use any pseudographics to draw tables, + it separates columns with a double space: + + >>> print(tabulate([["spam", 41.9999], ["eggs", "451.0"]], + ... ["strings", "numbers"], "plain")) + strings numbers + spam 41.9999 + eggs 451 + + >>> print(tabulate([["spam", 41.9999], ["eggs", "451.0"]], tablefmt="plain")) + spam 41.9999 + eggs 451 + + "simple" format is like Pandoc simple_tables: + + >>> print(tabulate([["spam", 41.9999], ["eggs", "451.0"]], + ... ["strings", "numbers"], "simple")) + strings numbers + --------- --------- + spam 41.9999 + eggs 451 + + >>> print(tabulate([["spam", 41.9999], ["eggs", "451.0"]], tablefmt="simple")) + ---- -------- + spam 41.9999 + eggs 451 + ---- -------- + + "grid" is similar to tables produced by Emacs table.el package or + Pandoc grid_tables: + + >>> print(tabulate([["spam", 41.9999], ["eggs", "451.0"]], + ... ["strings", "numbers"], "grid")) + +-----------+-----------+ + | strings | numbers | + +===========+===========+ + | spam | 41.9999 | + +-----------+-----------+ + | eggs | 451 | + +-----------+-----------+ + + >>> print(tabulate([["spam", 41.9999], ["eggs", "451.0"]], tablefmt="grid")) + +------+----------+ + | spam | 41.9999 | + +------+----------+ + | eggs | 451 | + +------+----------+ + + "simple_grid" draws a grid using single-line box-drawing + characters: + + >>> print(tabulate([["spam", 41.9999], ["eggs", "451.0"]], + ... ["strings", "numbers"], "simple_grid")) + ┌───────────┬───────────┐ + │ strings │ numbers │ + ├───────────┼───────────┤ + │ spam │ 41.9999 │ + ├───────────┼───────────┤ + │ eggs │ 451 │ + └───────────┴───────────┘ + + "rounded_grid" draws a grid using single-line box-drawing + characters with rounded corners: + + >>> print(tabulate([["spam", 41.9999], ["eggs", "451.0"]], + ... ["strings", "numbers"], "rounded_grid")) + ╭───────────┬───────────╮ + │ strings │ numbers │ + ├───────────┼───────────┤ + │ spam │ 41.9999 │ + ├───────────┼───────────┤ + │ eggs │ 451 │ + ╰───────────┴───────────╯ + + "heavy_grid" draws a grid using bold (thick) single-line box-drawing + characters: + + >>> print(tabulate([["spam", 41.9999], ["eggs", "451.0"]], + ... ["strings", "numbers"], "heavy_grid")) + ┏━━━━━━━━━━━┳━━━━━━━━━━━┓ + ┃ strings ┃ numbers ┃ + ┣━━━━━━━━━━━╋━━━━━━━━━━━┫ + ┃ spam ┃ 41.9999 ┃ + ┣━━━━━━━━━━━╋━━━━━━━━━━━┫ + ┃ eggs ┃ 451 ┃ + ┗━━━━━━━━━━━┻━━━━━━━━━━━┛ + + "mixed_grid" draws a grid using a mix of light (thin) and heavy (thick) lines + box-drawing characters: + + >>> print(tabulate([["spam", 41.9999], ["eggs", "451.0"]], + ... ["strings", "numbers"], "mixed_grid")) + ┍━━━━━━━━━━━┯━━━━━━━━━━━┑ + │ strings │ numbers │ + ┝━━━━━━━━━━━┿━━━━━━━━━━━┥ + │ spam │ 41.9999 │ + ├───────────┼───────────┤ + │ eggs │ 451 │ + ┕━━━━━━━━━━━┷━━━━━━━━━━━┙ + + "double_grid" draws a grid using double-line box-drawing + characters: + + >>> print(tabulate([["spam", 41.9999], ["eggs", "451.0"]], + ... ["strings", "numbers"], "double_grid")) + ╔═══════════╦═══════════╗ + ║ strings ║ numbers ║ + ╠═══════════╬═══════════╣ + ║ spam ║ 41.9999 ║ + ╠═══════════╬═══════════╣ + ║ eggs ║ 451 ║ + ╚═══════════╩═══════════╝ + + "fancy_grid" draws a grid using a mix of single and + double-line box-drawing characters: + + >>> print(tabulate([["spam", 41.9999], ["eggs", "451.0"]], + ... ["strings", "numbers"], "fancy_grid")) + ╒═══════════╤═══════════╕ + │ strings │ numbers │ + ╞═══════════╪═══════════╡ + │ spam │ 41.9999 │ + ├───────────┼───────────┤ + │ eggs │ 451 │ + ╘═══════════╧═══════════╛ + + "outline" is the same as the "grid" format but doesn't draw lines between rows: + + >>> print(tabulate([["spam", 41.9999], ["eggs", "451.0"]], + ... ["strings", "numbers"], "outline")) + +-----------+-----------+ + | strings | numbers | + +===========+===========+ + | spam | 41.9999 | + | eggs | 451 | + +-----------+-----------+ + + >>> print(tabulate([["spam", 41.9999], ["eggs", "451.0"]], tablefmt="outline")) + +------+----------+ + | spam | 41.9999 | + | eggs | 451 | + +------+----------+ + + "simple_outline" is the same as the "simple_grid" format but doesn't draw lines between rows: + + >>> print(tabulate([["spam", 41.9999], ["eggs", "451.0"]], + ... ["strings", "numbers"], "simple_outline")) + ┌───────────┬───────────┐ + │ strings │ numbers │ + ├───────────┼───────────┤ + │ spam │ 41.9999 │ + │ eggs │ 451 │ + └───────────┴───────────┘ + + "rounded_outline" is the same as the "rounded_grid" format but doesn't draw lines between rows: + + >>> print(tabulate([["spam", 41.9999], ["eggs", "451.0"]], + ... ["strings", "numbers"], "rounded_outline")) + ╭───────────┬───────────╮ + │ strings │ numbers │ + ├───────────┼───────────┤ + │ spam │ 41.9999 │ + │ eggs │ 451 │ + ╰───────────┴───────────╯ + + "heavy_outline" is the same as the "heavy_grid" format but doesn't draw lines between rows: + + >>> print(tabulate([["spam", 41.9999], ["eggs", "451.0"]], + ... ["strings", "numbers"], "heavy_outline")) + ┏━━━━━━━━━━━┳━━━━━━━━━━━┓ + ┃ strings ┃ numbers ┃ + ┣━━━━━━━━━━━╋━━━━━━━━━━━┫ + ┃ spam ┃ 41.9999 ┃ + ┃ eggs ┃ 451 ┃ + ┗━━━━━━━━━━━┻━━━━━━━━━━━┛ + + "mixed_outline" is the same as the "mixed_grid" format but doesn't draw lines between rows: + + >>> print(tabulate([["spam", 41.9999], ["eggs", "451.0"]], + ... ["strings", "numbers"], "mixed_outline")) + ┍━━━━━━━━━━━┯━━━━━━━━━━━┑ + │ strings │ numbers │ + ┝━━━━━━━━━━━┿━━━━━━━━━━━┥ + │ spam │ 41.9999 │ + │ eggs │ 451 │ + ┕━━━━━━━━━━━┷━━━━━━━━━━━┙ + + "double_outline" is the same as the "double_grid" format but doesn't draw lines between rows: + + >>> print(tabulate([["spam", 41.9999], ["eggs", "451.0"]], + ... ["strings", "numbers"], "double_outline")) + ╔═══════════╦═══════════╗ + ║ strings ║ numbers ║ + ╠═══════════╬═══════════╣ + ║ spam ║ 41.9999 ║ + ║ eggs ║ 451 ║ + ╚═══════════╩═══════════╝ + + "fancy_outline" is the same as the "fancy_grid" format but doesn't draw lines between rows: + + >>> print(tabulate([["spam", 41.9999], ["eggs", "451.0"]], + ... ["strings", "numbers"], "fancy_outline")) + ╒═══════════╤═══════════╕ + │ strings │ numbers │ + ╞═══════════╪═══════════╡ + │ spam │ 41.9999 │ + │ eggs │ 451 │ + ╘═══════════╧═══════════╛ + + "pipe" is like tables in PHP Markdown Extra extension or Pandoc + pipe_tables: + + >>> print(tabulate([["spam", 41.9999], ["eggs", "451.0"]], + ... ["strings", "numbers"], "pipe")) + | strings | numbers | + |:----------|----------:| + | spam | 41.9999 | + | eggs | 451 | + + "presto" is like tables produce by the Presto CLI: + + >>> print(tabulate([["spam", 41.9999], ["eggs", "451.0"]], + ... ["strings", "numbers"], "presto")) + strings | numbers + -----------+----------- + spam | 41.9999 + eggs | 451 + + >>> print(tabulate([["spam", 41.9999], ["eggs", "451.0"]], tablefmt="pipe")) + |:-----|---------:| + | spam | 41.9999 | + | eggs | 451 | + + "orgtbl" is like tables in Emacs org-mode and orgtbl-mode. They + are slightly different from "pipe" format by not using colons to + define column alignment, and using a "+" sign to indicate line + intersections: + + >>> print(tabulate([["spam", 41.9999], ["eggs", "451.0"]], + ... ["strings", "numbers"], "orgtbl")) + | strings | numbers | + |-----------+-----------| + | spam | 41.9999 | + | eggs | 451 | + + + >>> print(tabulate([["spam", 41.9999], ["eggs", "451.0"]], tablefmt="orgtbl")) + | spam | 41.9999 | + | eggs | 451 | + + "rst" is like a simple table format from reStructuredText; please + note that reStructuredText accepts also "grid" tables: + + >>> print(tabulate([["spam", 41.9999], ["eggs", "451.0"]], + ... ["strings", "numbers"], "rst")) + ========= ========= + strings numbers + ========= ========= + spam 41.9999 + eggs 451 + ========= ========= + + >>> print(tabulate([["spam", 41.9999], ["eggs", "451.0"]], tablefmt="rst")) + ==== ======== + spam 41.9999 + eggs 451 + ==== ======== + + "mediawiki" produces a table markup used in Wikipedia and on other + MediaWiki-based sites: + + >>> print(tabulate([["strings", "numbers"], ["spam", 41.9999], ["eggs", "451.0"]], + ... headers="firstrow", tablefmt="mediawiki")) + {| class="wikitable" style="text-align: left;" + |+ + |- + ! strings !! align="right"| numbers + |- + | spam || align="right"| 41.9999 + |- + | eggs || align="right"| 451 + |} + + "html" produces HTML markup as an html.escape'd str + with a ._repr_html_ method so that Jupyter Lab and Notebook display the HTML + and a .str property so that the raw HTML remains accessible + the unsafehtml table format can be used if an unescaped HTML format is required: + + >>> print(tabulate([["strings", "numbers"], ["spam", 41.9999], ["eggs", "451.0"]], + ... headers="firstrow", tablefmt="html")) + + + + + + + + +
strings numbers
spam 41.9999
eggs 451
+ + "latex" produces a tabular environment of LaTeX document markup: + + >>> print(tabulate([["spam", 41.9999], ["eggs", "451.0"]], tablefmt="latex")) + \\begin{tabular}{lr} + \\hline + spam & 41.9999 \\\\ + eggs & 451 \\\\ + \\hline + \\end{tabular} + + "latex_raw" is similar to "latex", but doesn't escape special characters, + such as backslash and underscore, so LaTeX commands may embedded into + cells' values: + + >>> print(tabulate([["spam$_9$", 41.9999], ["\\\\emph{eggs}", "451.0"]], tablefmt="latex_raw")) + \\begin{tabular}{lr} + \\hline + spam$_9$ & 41.9999 \\\\ + \\emph{eggs} & 451 \\\\ + \\hline + \\end{tabular} + + "latex_booktabs" produces a tabular environment of LaTeX document markup + using the booktabs.sty package: + + >>> print(tabulate([["spam", 41.9999], ["eggs", "451.0"]], tablefmt="latex_booktabs")) + \\begin{tabular}{lr} + \\toprule + spam & 41.9999 \\\\ + eggs & 451 \\\\ + \\bottomrule + \\end{tabular} + + "latex_longtable" produces a tabular environment that can stretch along + multiple pages, using the longtable package for LaTeX. + + >>> print(tabulate([["spam", 41.9999], ["eggs", "451.0"]], tablefmt="latex_longtable")) + \\begin{longtable}{lr} + \\hline + spam & 41.9999 \\\\ + eggs & 451 \\\\ + \\hline + \\end{longtable} + + + Number parsing + -------------- + By default, anything which can be parsed as a number is a number. + This ensures numbers represented as strings are aligned properly. + This can lead to weird results for particular strings such as + specific git SHAs e.g. "42992e1" will be parsed into the number + 429920 and aligned as such. + + To completely disable number parsing (and alignment), use + `disable_numparse=True`. For more fine grained control, a list column + indices is used to disable number parsing only on those columns + e.g. `disable_numparse=[0, 2]` would disable number parsing only on the + first and third columns. + + Column Widths and Auto Line Wrapping + ------------------------------------ + Tabulate will, by default, set the width of each column to the length of the + longest element in that column. However, in situations where fields are expected + to reasonably be too long to look good as a single line, tabulate can help automate + word wrapping long fields for you. Use the parameter `maxcolwidth` to provide a + list of maximal column widths + + >>> print(tabulate( \ + [('1', 'John Smith', \ + 'This is a rather long description that might look better if it is wrapped a bit')], \ + headers=("Issue Id", "Author", "Description"), \ + maxcolwidths=[None, None, 30], \ + tablefmt="grid" \ + )) + +------------+------------+-------------------------------+ + | Issue Id | Author | Description | + +============+============+===============================+ + | 1 | John Smith | This is a rather long | + | | | description that might look | + | | | better if it is wrapped a bit | + +------------+------------+-------------------------------+ + + Header column width can be specified in a similar way using `maxheadercolwidth` + + """ + + if tabular_data is None: + tabular_data = [] + + list_of_lists, headers = _normalize_tabular_data( + tabular_data, headers, showindex=showindex + ) + list_of_lists, separating_lines = _remove_separating_lines(list_of_lists) + + if maxcolwidths is not None: + num_cols = len(list_of_lists[0]) + if isinstance(maxcolwidths, int): # Expand scalar for all columns + maxcolwidths = _expand_iterable(maxcolwidths, num_cols, maxcolwidths) + else: # Ignore col width for any 'trailing' columns + maxcolwidths = _expand_iterable(maxcolwidths, num_cols, None) + + numparses = _expand_numparse(disable_numparse, num_cols) + list_of_lists = _wrap_text_to_colwidths( + list_of_lists, maxcolwidths, numparses=numparses + ) + + if maxheadercolwidths is not None: + num_cols = len(list_of_lists[0]) + if isinstance(maxheadercolwidths, int): # Expand scalar for all columns + maxheadercolwidths = _expand_iterable( + maxheadercolwidths, num_cols, maxheadercolwidths + ) + else: # Ignore col width for any 'trailing' columns + maxheadercolwidths = _expand_iterable(maxheadercolwidths, num_cols, None) + + numparses = _expand_numparse(disable_numparse, num_cols) + headers = _wrap_text_to_colwidths( + [headers], maxheadercolwidths, numparses=numparses + )[0] + + # empty values in the first column of RST tables should be escaped (issue #82) + # "" should be escaped as "\\ " or ".." + if tablefmt == "rst": + list_of_lists, headers = _rst_escape_first_column(list_of_lists, headers) + + # PrettyTable formatting does not use any extra padding. + # Numbers are not parsed and are treated the same as strings for alignment. + # Check if pretty is the format being used and override the defaults so it + # does not impact other formats. + min_padding = MIN_PADDING + if tablefmt == "pretty": + min_padding = 0 + disable_numparse = True + numalign = "center" if numalign == _DEFAULT_ALIGN else numalign + stralign = "center" if stralign == _DEFAULT_ALIGN else stralign + else: + numalign = "decimal" if numalign == _DEFAULT_ALIGN else numalign + stralign = "left" if stralign == _DEFAULT_ALIGN else stralign + + # optimization: look for ANSI control codes once, + # enable smart width functions only if a control code is found + # + # convert the headers and rows into a single, tab-delimited string ensuring + # that any bytestrings are decoded safely (i.e. errors ignored) + plain_text = "\t".join( + chain( + # headers + map(_to_str, headers), + # rows: chain the rows together into a single iterable after mapping + # the bytestring conversino to each cell value + chain.from_iterable(map(_to_str, row) for row in list_of_lists), + ) + ) + + has_invisible = _ansi_codes.search(plain_text) is not None + + enable_widechars = wcwidth is not None and WIDE_CHARS_MODE + if ( + not isinstance(tablefmt, TableFormat) + and tablefmt in multiline_formats + and _is_multiline(plain_text) + ): + tablefmt = multiline_formats.get(tablefmt, tablefmt) + is_multiline = True + else: + is_multiline = False + width_fn = _choose_width_fn(has_invisible, enable_widechars, is_multiline) + + # format rows and columns, convert numeric values to strings + cols = list(izip_longest(*list_of_lists)) + numparses = _expand_numparse(disable_numparse, len(cols)) + coltypes = [_column_type(col, numparse=np) for col, np in zip(cols, numparses)] + if isinstance(floatfmt, str): # old version + float_formats = len(cols) * [ + floatfmt + ] # just duplicate the string to use in each column + else: # if floatfmt is list, tuple etc we have one per column + float_formats = list(floatfmt) + if len(float_formats) < len(cols): + float_formats.extend((len(cols) - len(float_formats)) * [_DEFAULT_FLOATFMT]) + if isinstance(intfmt, str): # old version + int_formats = len(cols) * [ + intfmt + ] # just duplicate the string to use in each column + else: # if intfmt is list, tuple etc we have one per column + int_formats = list(intfmt) + if len(int_formats) < len(cols): + int_formats.extend((len(cols) - len(int_formats)) * [_DEFAULT_INTFMT]) + if isinstance(missingval, str): + missing_vals = len(cols) * [missingval] + else: + missing_vals = list(missingval) + if len(missing_vals) < len(cols): + missing_vals.extend((len(cols) - len(missing_vals)) * [_DEFAULT_MISSINGVAL]) + cols = [ + [_format(v, ct, fl_fmt, int_fmt, miss_v, has_invisible) for v in c] + for c, ct, fl_fmt, int_fmt, miss_v in zip( + cols, coltypes, float_formats, int_formats, missing_vals + ) + ] + + # align columns + aligns = [numalign if ct in [int, float] else stralign for ct in coltypes] + if colalign is not None: + assert isinstance(colalign, Iterable) + for idx, align in enumerate(colalign): + aligns[idx] = align + minwidths = ( + [width_fn(h) + min_padding for h in headers] if headers else [0] * len(cols) + ) + cols = [ + _align_column(c, a, minw, has_invisible, enable_widechars, is_multiline) + for c, a, minw in zip(cols, aligns, minwidths) + ] + + if headers: + # align headers and add headers + t_cols = cols or [[""]] * len(headers) + t_aligns = aligns or [stralign] * len(headers) + minwidths = [ + max(minw, max(width_fn(cl) for cl in c)) + for minw, c in zip(minwidths, t_cols) + ] + headers = [ + _align_header(h, a, minw, width_fn(h), is_multiline, width_fn) + for h, a, minw in zip(headers, t_aligns, minwidths) + ] + rows = list(zip(*cols)) + else: + minwidths = [max(width_fn(cl) for cl in c) for c in cols] + rows = list(zip(*cols)) + + if not isinstance(tablefmt, TableFormat): + tablefmt = _table_formats.get(tablefmt, _table_formats["simple"]) + + ra_default = rowalign if isinstance(rowalign, str) else None + rowaligns = _expand_iterable(rowalign, len(rows), ra_default) + _reinsert_separating_lines(rows, separating_lines) + + return _format_table( + tablefmt, headers, rows, minwidths, aligns, is_multiline, rowaligns=rowaligns + ) + + +def _expand_numparse(disable_numparse, column_count): + """ + Return a list of bools of length `column_count` which indicates whether + number parsing should be used on each column. + If `disable_numparse` is a list of indices, each of those indices are False, + and everything else is True. + If `disable_numparse` is a bool, then the returned list is all the same. + """ + if isinstance(disable_numparse, Iterable): + numparses = [True] * column_count + for index in disable_numparse: + numparses[index] = False + return numparses + else: + return [not disable_numparse] * column_count + + +def _expand_iterable(original, num_desired, default): + """ + Expands the `original` argument to return a return a list of + length `num_desired`. If `original` is shorter than `num_desired`, it will + be padded with the value in `default`. + If `original` is not a list to begin with (i.e. scalar value) a list of + length `num_desired` completely populated with `default will be returned + """ + if isinstance(original, Iterable) and not isinstance(original, str): + return original + [default] * (num_desired - len(original)) + else: + return [default] * num_desired + + +def _pad_row(cells, padding): + if cells: + pad = " " * padding + padded_cells = [pad + cell + pad for cell in cells] + return padded_cells + else: + return cells + + +def _build_simple_row(padded_cells, rowfmt): + "Format row according to DataRow format without padding." + begin, sep, end = rowfmt + return (begin + sep.join(padded_cells) + end).rstrip() + + +def _build_row(padded_cells, colwidths, colaligns, rowfmt): + "Return a string which represents a row of data cells." + if not rowfmt: + return None + if hasattr(rowfmt, "__call__"): + return rowfmt(padded_cells, colwidths, colaligns) + else: + return _build_simple_row(padded_cells, rowfmt) + + +def _append_basic_row(lines, padded_cells, colwidths, colaligns, rowfmt, rowalign=None): + # NOTE: rowalign is ignored and exists for api compatibility with _append_multiline_row + lines.append(_build_row(padded_cells, colwidths, colaligns, rowfmt)) + return lines + + +def _align_cell_veritically(text_lines, num_lines, column_width, row_alignment): + delta_lines = num_lines - len(text_lines) + blank = [" " * column_width] + if row_alignment == "bottom": + return blank * delta_lines + text_lines + elif row_alignment == "center": + top_delta = delta_lines // 2 + bottom_delta = delta_lines - top_delta + return top_delta * blank + text_lines + bottom_delta * blank + else: + return text_lines + blank * delta_lines + + +def _append_multiline_row( + lines, padded_multiline_cells, padded_widths, colaligns, rowfmt, pad, rowalign=None +): + colwidths = [w - 2 * pad for w in padded_widths] + cells_lines = [c.splitlines() for c in padded_multiline_cells] + nlines = max(map(len, cells_lines)) # number of lines in the row + # vertically pad cells where some lines are missing + # cells_lines = [ + # (cl + [" " * w] * (nlines - len(cl))) for cl, w in zip(cells_lines, colwidths) + # ] + + cells_lines = [ + _align_cell_veritically(cl, nlines, w, rowalign) + for cl, w in zip(cells_lines, colwidths) + ] + lines_cells = [[cl[i] for cl in cells_lines] for i in range(nlines)] + for ln in lines_cells: + padded_ln = _pad_row(ln, pad) + _append_basic_row(lines, padded_ln, colwidths, colaligns, rowfmt) + return lines + + +def _build_line(colwidths, colaligns, linefmt): + "Return a string which represents a horizontal line." + if not linefmt: + return None + if hasattr(linefmt, "__call__"): + return linefmt(colwidths, colaligns) + else: + begin, fill, sep, end = linefmt + cells = [fill * w for w in colwidths] + return _build_simple_row(cells, (begin, sep, end)) + + +def _append_line(lines, colwidths, colaligns, linefmt): + lines.append(_build_line(colwidths, colaligns, linefmt)) + return lines + + +class JupyterHTMLStr(str): + """Wrap the string with a _repr_html_ method so that Jupyter + displays the HTML table""" + + def _repr_html_(self): + return self + + @property + def str(self): + """add a .str property so that the raw string is still accessible""" + return self + + +def _format_table(fmt, headers, rows, colwidths, colaligns, is_multiline, rowaligns): + """Produce a plain-text representation of the table.""" + lines = [] + hidden = fmt.with_header_hide if (headers and fmt.with_header_hide) else [] + pad = fmt.padding + headerrow = fmt.headerrow + + padded_widths = [(w + 2 * pad) for w in colwidths] + if is_multiline: + pad_row = lambda row, _: row # noqa do it later, in _append_multiline_row + append_row = partial(_append_multiline_row, pad=pad) + else: + pad_row = _pad_row + append_row = _append_basic_row + + padded_headers = pad_row(headers, pad) + padded_rows = [pad_row(row, pad) for row in rows] + + if fmt.lineabove and "lineabove" not in hidden: + _append_line(lines, padded_widths, colaligns, fmt.lineabove) + + if padded_headers: + append_row(lines, padded_headers, padded_widths, colaligns, headerrow) + if fmt.linebelowheader and "linebelowheader" not in hidden: + _append_line(lines, padded_widths, colaligns, fmt.linebelowheader) + + if padded_rows and fmt.linebetweenrows and "linebetweenrows" not in hidden: + # initial rows with a line below + for row, ralign in zip(padded_rows[:-1], rowaligns): + append_row( + lines, row, padded_widths, colaligns, fmt.datarow, rowalign=ralign + ) + _append_line(lines, padded_widths, colaligns, fmt.linebetweenrows) + # the last row without a line below + append_row( + lines, + padded_rows[-1], + padded_widths, + colaligns, + fmt.datarow, + rowalign=rowaligns[-1], + ) + else: + separating_line = ( + fmt.linebetweenrows + or fmt.linebelowheader + or fmt.linebelow + or fmt.lineabove + or Line("", "", "", "") + ) + for row in padded_rows: + # test to see if either the 1st column or the 2nd column (account for showindex) has + # the SEPARATING_LINE flag + if _is_separating_line(row): + _append_line(lines, padded_widths, colaligns, separating_line) + else: + append_row(lines, row, padded_widths, colaligns, fmt.datarow) + + if fmt.linebelow and "linebelow" not in hidden: + _append_line(lines, padded_widths, colaligns, fmt.linebelow) + + if headers or rows: + output = "\n".join(lines) + if fmt.lineabove == _html_begin_table_without_header: + return JupyterHTMLStr(output) + else: + return output + else: # a completely empty table + return "" + + +class _CustomTextWrap(textwrap.TextWrapper): + """A custom implementation of CPython's textwrap.TextWrapper. This supports + both wide characters (Korea, Japanese, Chinese) - including mixed string. + For the most part, the `_handle_long_word` and `_wrap_chunks` functions were + copy pasted out of the CPython baseline, and updated with our custom length + and line appending logic. + """ + + def __init__(self, *args, **kwargs): + self._active_codes = [] + self.max_lines = None # For python2 compatibility + textwrap.TextWrapper.__init__(self, *args, **kwargs) + + @staticmethod + def _len(item): + """Custom len that gets console column width for wide + and non-wide characters as well as ignores color codes""" + stripped = _strip_ansi(item) + if wcwidth: + return wcwidth.wcswidth(stripped) + else: + return len(stripped) + + def _update_lines(self, lines, new_line): + """Adds a new line to the list of lines the text is being wrapped into + This function will also track any ANSI color codes in this string as well + as add any colors from previous lines order to preserve the same formatting + as a single unwrapped string. + """ + code_matches = [x for x in _ansi_codes.finditer(new_line)] + color_codes = [ + code.string[code.span()[0] : code.span()[1]] for code in code_matches + ] + + # Add color codes from earlier in the unwrapped line, and then track any new ones we add. + new_line = "".join(self._active_codes) + new_line + + for code in color_codes: + if code != _ansi_color_reset_code: + self._active_codes.append(code) + else: # A single reset code resets everything + self._active_codes = [] + + # Always ensure each line is color terminted if any colors are + # still active, otherwise colors will bleed into other cells on the console + if len(self._active_codes) > 0: + new_line = new_line + _ansi_color_reset_code + + lines.append(new_line) + + def _handle_long_word(self, reversed_chunks, cur_line, cur_len, width): + """_handle_long_word(chunks : [string], + cur_line : [string], + cur_len : int, width : int) + Handle a chunk of text (most likely a word, not whitespace) that + is too long to fit in any line. + """ + # Figure out when indent is larger than the specified width, and make + # sure at least one character is stripped off on every pass + if width < 1: + space_left = 1 + else: + space_left = width - cur_len + + # If we're allowed to break long words, then do so: put as much + # of the next chunk onto the current line as will fit. + if self.break_long_words: + # Tabulate Custom: Build the string up piece-by-piece in order to + # take each charcter's width into account + chunk = reversed_chunks[-1] + i = 1 + while self._len(chunk[:i]) <= space_left: + i = i + 1 + cur_line.append(chunk[: i - 1]) + reversed_chunks[-1] = chunk[i - 1 :] + + # Otherwise, we have to preserve the long word intact. Only add + # it to the current line if there's nothing already there -- + # that minimizes how much we violate the width constraint. + elif not cur_line: + cur_line.append(reversed_chunks.pop()) + + # If we're not allowed to break long words, and there's already + # text on the current line, do nothing. Next time through the + # main loop of _wrap_chunks(), we'll wind up here again, but + # cur_len will be zero, so the next line will be entirely + # devoted to the long word that we can't handle right now. + + def _wrap_chunks(self, chunks): + """_wrap_chunks(chunks : [string]) -> [string] + Wrap a sequence of text chunks and return a list of lines of + length 'self.width' or less. (If 'break_long_words' is false, + some lines may be longer than this.) Chunks correspond roughly + to words and the whitespace between them: each chunk is + indivisible (modulo 'break_long_words'), but a line break can + come between any two chunks. Chunks should not have internal + whitespace; ie. a chunk is either all whitespace or a "word". + Whitespace chunks will be removed from the beginning and end of + lines, but apart from that whitespace is preserved. + """ + lines = [] + if self.width <= 0: + raise ValueError("invalid width %r (must be > 0)" % self.width) + if self.max_lines is not None: + if self.max_lines > 1: + indent = self.subsequent_indent + else: + indent = self.initial_indent + if self._len(indent) + self._len(self.placeholder.lstrip()) > self.width: + raise ValueError("placeholder too large for max width") + + # Arrange in reverse order so items can be efficiently popped + # from a stack of chucks. + chunks.reverse() + + while chunks: + + # Start the list of chunks that will make up the current line. + # cur_len is just the length of all the chunks in cur_line. + cur_line = [] + cur_len = 0 + + # Figure out which static string will prefix this line. + if lines: + indent = self.subsequent_indent + else: + indent = self.initial_indent + + # Maximum width for this line. + width = self.width - self._len(indent) + + # First chunk on line is whitespace -- drop it, unless this + # is the very beginning of the text (ie. no lines started yet). + if self.drop_whitespace and chunks[-1].strip() == "" and lines: + del chunks[-1] + + while chunks: + chunk_len = self._len(chunks[-1]) + + # Can at least squeeze this chunk onto the current line. + if cur_len + chunk_len <= width: + cur_line.append(chunks.pop()) + cur_len += chunk_len + + # Nope, this line is full. + else: + break + + # The current line is full, and the next chunk is too big to + # fit on *any* line (not just this one). + if chunks and self._len(chunks[-1]) > width: + self._handle_long_word(chunks, cur_line, cur_len, width) + cur_len = sum(map(self._len, cur_line)) + + # If the last chunk on this line is all whitespace, drop it. + if self.drop_whitespace and cur_line and cur_line[-1].strip() == "": + cur_len -= self._len(cur_line[-1]) + del cur_line[-1] + + if cur_line: + if ( + self.max_lines is None + or len(lines) + 1 < self.max_lines + or ( + not chunks + or self.drop_whitespace + and len(chunks) == 1 + and not chunks[0].strip() + ) + and cur_len <= width + ): + # Convert current line back to a string and store it in + # list of all lines (return value). + self._update_lines(lines, indent + "".join(cur_line)) + else: + while cur_line: + if ( + cur_line[-1].strip() + and cur_len + self._len(self.placeholder) <= width + ): + cur_line.append(self.placeholder) + self._update_lines(lines, indent + "".join(cur_line)) + break + cur_len -= self._len(cur_line[-1]) + del cur_line[-1] + else: + if lines: + prev_line = lines[-1].rstrip() + if ( + self._len(prev_line) + self._len(self.placeholder) + <= self.width + ): + lines[-1] = prev_line + self.placeholder + break + self._update_lines(lines, indent + self.placeholder.lstrip()) + break + + return lines + + +def _main(): + """\ + Usage: tabulate [options] [FILE ...] + + Pretty-print tabular data. + See also https://github.com/astanin/python-tabulate + + FILE a filename of the file with tabular data; + if "-" or missing, read data from stdin. + + Options: + + -h, --help show this message + -1, --header use the first row of data as a table header + -o FILE, --output FILE print table to FILE (default: stdout) + -s REGEXP, --sep REGEXP use a custom column separator (default: whitespace) + -F FPFMT, --float FPFMT floating point number format (default: g) + -I INTFMT, --int INTFMT integer point number format (default: "") + -f FMT, --format FMT set output table format; supported formats: + plain, simple, grid, fancy_grid, pipe, orgtbl, + rst, mediawiki, html, latex, latex_raw, + latex_booktabs, latex_longtable, tsv + (default: simple) + """ + import getopt + import sys + import textwrap + + usage = textwrap.dedent(_main.__doc__) + try: + opts, args = getopt.getopt( + sys.argv[1:], + "h1o:s:F:A:f:", + ["help", "header", "output", "sep=", "float=", "int=", "align=", "format="], + ) + except getopt.GetoptError as e: + print(e) + print(usage) + sys.exit(2) + headers = [] + floatfmt = _DEFAULT_FLOATFMT + intfmt = _DEFAULT_INTFMT + colalign = None + tablefmt = "simple" + sep = r"\s+" + outfile = "-" + for opt, value in opts: + if opt in ["-1", "--header"]: + headers = "firstrow" + elif opt in ["-o", "--output"]: + outfile = value + elif opt in ["-F", "--float"]: + floatfmt = value + elif opt in ["-I", "--int"]: + intfmt = value + elif opt in ["-C", "--colalign"]: + colalign = value.split() + elif opt in ["-f", "--format"]: + if value not in tabulate_formats: + print("%s is not a supported table format" % value) + print(usage) + sys.exit(3) + tablefmt = value + elif opt in ["-s", "--sep"]: + sep = value + elif opt in ["-h", "--help"]: + print(usage) + sys.exit(0) + files = [sys.stdin] if not args else args + with (sys.stdout if outfile == "-" else open(outfile, "w")) as out: + for f in files: + if f == "-": + f = sys.stdin + if _is_file(f): + _pprint_file( + f, + headers=headers, + tablefmt=tablefmt, + sep=sep, + floatfmt=floatfmt, + intfmt=intfmt, + file=out, + colalign=colalign, + ) + else: + with open(f) as fobj: + _pprint_file( + fobj, + headers=headers, + tablefmt=tablefmt, + sep=sep, + floatfmt=floatfmt, + intfmt=intfmt, + file=out, + colalign=colalign, + ) + + +def _pprint_file(fobject, headers, tablefmt, sep, floatfmt, intfmt, file, colalign): + rows = fobject.readlines() + table = [re.split(sep, r.rstrip()) for r in rows if r.strip()] + print( + tabulate( + table, + headers, + tablefmt, + floatfmt=floatfmt, + intfmt=intfmt, + colalign=colalign, + ), + file=file, + ) + + +if __name__ == "__main__": + _main() diff --git a/lib/python3.12/site-packages/ray/_private/tls_utils.py b/lib/python3.12/site-packages/ray/_private/tls_utils.py new file mode 100644 index 0000000000000000000000000000000000000000..12c65b50905cfec0a523ccabf74c33764be2ecfd --- /dev/null +++ b/lib/python3.12/site-packages/ray/_private/tls_utils.py @@ -0,0 +1,96 @@ +import datetime +import os +import socket + +from ray._common.network_utils import ( + get_localhost_ip, + node_ip_address_from_perspective, +) + + +def generate_self_signed_tls_certs(): + """Create self-signed key/cert pair for testing. + + This method requires the library ``cryptography`` be installed. + """ + try: + from cryptography import x509 + from cryptography.hazmat.backends import default_backend + from cryptography.hazmat.primitives import hashes, serialization + from cryptography.hazmat.primitives.asymmetric import rsa + from cryptography.x509.oid import NameOID + except ImportError: + raise ImportError( + "Using `Security.temporary` requires `cryptography`, please " + "install it using either pip or conda" + ) + key = rsa.generate_private_key( + public_exponent=65537, key_size=2048, backend=default_backend() + ) + key_contents = key.private_bytes( + encoding=serialization.Encoding.PEM, + format=serialization.PrivateFormat.PKCS8, + encryption_algorithm=serialization.NoEncryption(), + ).decode() + + ray_interal = x509.Name([x509.NameAttribute(NameOID.COMMON_NAME, "ray-internal")]) + altnames = x509.SubjectAlternativeName( + [ + x509.DNSName( + socket.gethostbyname(socket.gethostname()) + ), # Probably 127.0.0.1 or ::1 + x509.DNSName(get_localhost_ip()), + x509.DNSName(node_ip_address_from_perspective()), + x509.DNSName("localhost"), + ] + ) + now = datetime.datetime.utcnow() + cert = ( + x509.CertificateBuilder() + .subject_name(ray_interal) + .issuer_name(ray_interal) + .add_extension(altnames, critical=False) + .public_key(key.public_key()) + .serial_number(x509.random_serial_number()) + .not_valid_before(now) + .not_valid_after(now + datetime.timedelta(days=365)) + .sign(key, hashes.SHA256(), default_backend()) + ) + + cert_contents = cert.public_bytes(serialization.Encoding.PEM).decode() + + return cert_contents, key_contents + + +def add_port_to_grpc_server(server, address): + import grpc + + if os.environ.get("RAY_USE_TLS", "0").lower() in ("1", "true"): + server_cert_chain, private_key, ca_cert = load_certs_from_env() + credentials = grpc.ssl_server_credentials( + [(private_key, server_cert_chain)], + root_certificates=ca_cert, + require_client_auth=ca_cert is not None, + ) + return server.add_secure_port(address, credentials) + else: + return server.add_insecure_port(address) + + +def load_certs_from_env(): + tls_env_vars = ["RAY_TLS_SERVER_CERT", "RAY_TLS_SERVER_KEY", "RAY_TLS_CA_CERT"] + if any(v not in os.environ for v in tls_env_vars): + raise RuntimeError( + "If the environment variable RAY_USE_TLS is set to true " + "then RAY_TLS_SERVER_CERT, RAY_TLS_SERVER_KEY and " + "RAY_TLS_CA_CERT must also be set." + ) + + with open(os.environ["RAY_TLS_SERVER_CERT"], "rb") as f: + server_cert_chain = f.read() + with open(os.environ["RAY_TLS_SERVER_KEY"], "rb") as f: + private_key = f.read() + with open(os.environ["RAY_TLS_CA_CERT"], "rb") as f: + ca_cert = f.read() + + return server_cert_chain, private_key, ca_cert diff --git a/lib/python3.12/site-packages/ray/_private/utils.py b/lib/python3.12/site-packages/ray/_private/utils.py new file mode 100644 index 0000000000000000000000000000000000000000..f7ba9af1465d1e6ab0e9643aebf2dff944e6d725 --- /dev/null +++ b/lib/python3.12/site-packages/ray/_private/utils.py @@ -0,0 +1,1626 @@ +import contextlib +import importlib +import json +import logging +import multiprocessing +import os +import platform +import re +import signal +import subprocess +import sys +import threading +import time +from collections import defaultdict +from pathlib import Path +from subprocess import list2cmdline +from typing import ( + TYPE_CHECKING, + Dict, + List, + Mapping, + Optional, + Tuple, + Union, +) + +from google.protobuf import json_format + +import ray +import ray._private.ray_constants as ray_constants +from ray._common.utils import ( + PLACEMENT_GROUP_BUNDLE_RESOURCE_NAME, + get_ray_address_file, + get_system_memory, +) +from ray.core.generated.runtime_environment_pb2 import ( + RuntimeEnvInfo as ProtoRuntimeEnvInfo, +) + +# Import psutil after ray so the packaged version is used. +import psutil + +if TYPE_CHECKING: + from ray.runtime_env import RuntimeEnv + + +INT32_MAX = (2**31) - 1 + + +pwd = None +if sys.platform != "win32": + import pwd + +logger = logging.getLogger(__name__) + +# Linux can bind child processes' lifetimes to that of their parents via prctl. +# prctl support is detected dynamically once, and assumed thereafter. +linux_prctl = None + +# Windows can bind processes' lifetimes to that of kernel-level "job objects". +# We keep a global job object to tie its lifetime to that of our own process. +win32_job = None +win32_AssignProcessToJobObject = None + +ENV_DISABLE_DOCKER_CPU_WARNING = "RAY_DISABLE_DOCKER_CPU_WARNING" in os.environ + +# This global variable is used for testing only +_CALLED_FREQ = defaultdict(lambda: 0) +_CALLED_FREQ_LOCK = threading.Lock() + +PLACEMENT_GROUP_INDEXED_BUNDLED_RESOURCE_PATTERN = re.compile( + r"(.+)_group_(\d+)_([0-9a-zA-Z]+)" +) +PLACEMENT_GROUP_WILDCARD_RESOURCE_PATTERN = re.compile(r"(.+)_group_([0-9a-zA-Z]+)") + + +def write_ray_address(ray_address: str, temp_dir: Optional[str] = None): + address_file = get_ray_address_file(temp_dir) + if os.path.exists(address_file): + with open(address_file, "r") as f: + prev_address = f.read() + if prev_address == ray_address: + return + + logger.info( + f"Overwriting previous Ray address ({prev_address}). " + "Running ray.init() on this node will now connect to the new " + f"instance at {ray_address}. To override this behavior, pass " + f"address={prev_address} to ray.init()." + ) + + with open(address_file, "w+") as f: + f.write(ray_address) + + +def read_ray_address(temp_dir: Optional[str] = None) -> str: + address_file = get_ray_address_file(temp_dir) + if not os.path.exists(address_file): + return None + with open(address_file, "r") as f: + return f.read().strip() + + +def format_error_message(exception_message: str, task_exception: bool = False): + """Improve the formatting of an exception thrown by a remote function. + + This method takes a traceback from an exception and makes it nicer by + removing a few uninformative lines and adding some space to indent the + remaining lines nicely. + + Args: + exception_message: A message generated by traceback.format_exc(). + + Returns: + A string of the formatted exception message. + """ + lines = exception_message.split("\n") + if task_exception: + # For errors that occur inside of tasks, remove lines 1 and 2 which are + # always the same, they just contain information about the worker code. + lines = lines[0:1] + lines[3:] + pass + return "\n".join(lines) + + +def push_error_to_driver( + worker, error_type: str, message: str, job_id: Optional[str] = None +): + """Push an error message to the driver to be printed in the background. + + Args: + worker: The worker to use. + error_type: The type of the error. + message: The message that will be printed in the background + on the driver. + job_id: The ID of the driver to push the error message to. If this + is None, then the message will be pushed to all drivers. + """ + if job_id is None: + job_id = ray.JobID.nil() + assert isinstance(job_id, ray.JobID) + worker.core_worker.push_error(job_id, error_type, message, time.time()) + + +def publish_error_to_driver( + error_type: str, + message: str, + gcs_client, + job_id=None, +): + """Push an error message to the driver to be printed in the background. + + Normally the push_error_to_driver function should be used. However, in some + instances, the raylet client is not available, e.g., because the + error happens in Python before the driver or worker has connected to the + backend processes. + + Args: + error_type: The type of the error. + message: The message that will be printed in the background + on the driver. + gcs_client: The GCS client to use. + job_id: The ID of the driver to push the error message to. If this + is None, then the message will be pushed to all drivers. + """ + if job_id is None: + job_id = ray.JobID.nil() + assert isinstance(job_id, ray.JobID) + try: + gcs_client.publish_error(job_id.hex().encode(), error_type, message, job_id, 60) + except Exception: + logger.exception(f"Failed to publish error: {message} [type {error_type}]") + + +def ensure_str(s, encoding="utf-8", errors="strict"): + """Coerce *s* to `str`. + + - `str` -> `str` + - `bytes` -> decoded to `str` + """ + if isinstance(s, str): + return s + else: + assert isinstance(s, bytes), f"Expected str or bytes, got {type(s)}" + return s.decode(encoding, errors) + + +def binary_to_object_ref(binary_object_ref): + return ray.ObjectRef(binary_object_ref) + + +def binary_to_task_id(binary_task_id): + return ray.TaskID(binary_task_id) + + +# TODO(qwang): Remove these hepler functions +# once we separate `WorkerID` from `UniqueID`. +def compute_job_id_from_driver(driver_id): + assert isinstance(driver_id, ray.WorkerID) + return ray.JobID(driver_id.binary()[0 : ray.JobID.size()]) + + +def compute_driver_id_from_job(job_id): + assert isinstance(job_id, ray.JobID) + rest_length = ray_constants.ID_SIZE - job_id.size() + driver_id_str = job_id.binary() + (rest_length * b"\xff") + return ray.WorkerID(driver_id_str) + + +def get_visible_accelerator_ids() -> Mapping[str, Optional[List[str]]]: + """Get the mapping from accelerator resource name + to the visible ids.""" + + from ray._private.accelerators import ( + get_accelerator_manager_for_resource, + get_all_accelerator_resource_names, + ) + + return { + accelerator_resource_name: get_accelerator_manager_for_resource( + accelerator_resource_name + ).get_current_process_visible_accelerator_ids() + for accelerator_resource_name in get_all_accelerator_resource_names() + } + + +def set_omp_num_threads_if_unset() -> bool: + """Set the OMP_NUM_THREADS to default to num cpus assigned to the worker + + This function sets the environment variable OMP_NUM_THREADS for the worker, + if the env is not previously set and it's running in worker (WORKER_MODE). + + Returns True if OMP_NUM_THREADS is set in this function. + + """ + num_threads_from_env = os.environ.get("OMP_NUM_THREADS") + if num_threads_from_env is not None: + # No ops if it's set + return False + + # If unset, try setting the correct CPU count assigned. + runtime_ctx = ray.get_runtime_context() + if runtime_ctx.worker.mode != ray._private.worker.WORKER_MODE: + # Non worker mode, no ops. + return False + + num_assigned_cpus = runtime_ctx.get_assigned_resources().get("CPU") + + if num_assigned_cpus is None: + # This is an actor task w/o any num_cpus specified, set it to 1 + logger.debug( + "[ray] Forcing OMP_NUM_THREADS=1 to avoid performance " + "degradation with many workers (issue #6998). You can override this " + "by explicitly setting OMP_NUM_THREADS, or changing num_cpus." + ) + num_assigned_cpus = 1 + + import math + + # For num_cpu < 1: Set to 1. + # For num_cpus >= 1: Set to the floor of the actual assigned cpus. + omp_num_threads = max(math.floor(num_assigned_cpus), 1) + os.environ["OMP_NUM_THREADS"] = str(omp_num_threads) + return True + + +def set_visible_accelerator_ids() -> Mapping[str, Optional[str]]: + """Set (CUDA_VISIBLE_DEVICES, ONEAPI_DEVICE_SELECTOR, HIP_VISIBLE_DEVICES, + NEURON_RT_VISIBLE_CORES, TPU_VISIBLE_CHIPS , HABANA_VISIBLE_MODULES ,...) + environment variables based on the accelerator runtime. Return the original + environment variables. + """ + from ray._private.ray_constants import env_bool + + original_visible_accelerator_env_vars = {} + override_on_zero = env_bool( + ray._private.accelerators.RAY_ACCEL_ENV_VAR_OVERRIDE_ON_ZERO_ENV_VAR, + True, + ) + for resource_name, accelerator_ids in ( + ray.get_runtime_context().get_accelerator_ids().items() + ): + # If no accelerator ids are set, skip overriding the environment variable. + if not override_on_zero and len(accelerator_ids) == 0: + continue + env_var = ray._private.accelerators.get_accelerator_manager_for_resource( + resource_name + ).get_visible_accelerator_ids_env_var() + original_visible_accelerator_env_vars[env_var] = os.environ.get(env_var, None) + ray._private.accelerators.get_accelerator_manager_for_resource( + resource_name + ).set_current_process_visible_accelerator_ids(accelerator_ids) + + return original_visible_accelerator_env_vars + + +def reset_visible_accelerator_env_vars( + original_visible_accelerator_env_vars: Mapping[str, Optional[str]] +) -> None: + """Reset the visible accelerator env vars to the original values.""" + for env_var, env_value in original_visible_accelerator_env_vars.items(): + if env_value is None: + os.environ.pop(env_var, None) + else: + os.environ[env_var] = env_value + + +class Unbuffered(object): + """There's no "built-in" solution to programatically disabling buffering of + text files. Ray expects stdout/err to be text files, so creating an + unbuffered binary file is unacceptable. + + See + https://mail.python.org/pipermail/tutor/2003-November/026645.html. + https://docs.python.org/3/library/functions.html#open + + """ + + def __init__(self, stream): + self.stream = stream + + def write(self, data): + self.stream.write(data) + self.stream.flush() + + def writelines(self, datas): + self.stream.writelines(datas) + self.stream.flush() + + def __getattr__(self, attr): + # Avoid endless loop when get `stream` attribute + if attr == "stream": + return super().__getattribute__("stream") + return getattr(self.stream, attr) + + +def open_log(path, unbuffered=False, **kwargs): + """ + Opens the log file at `path`, with the provided kwargs being given to + `open`. + """ + # Disable buffering, see test_advanced_3.py::test_logging_to_driver + kwargs.setdefault("buffering", 1) + kwargs.setdefault("mode", "a") + kwargs.setdefault("encoding", "utf-8") + stream = open(path, **kwargs) + if unbuffered: + return Unbuffered(stream) + else: + return stream + + +def _get_docker_cpus( + cpu_quota_file_name="/sys/fs/cgroup/cpu/cpu.cfs_quota_us", + cpu_period_file_name="/sys/fs/cgroup/cpu/cpu.cfs_period_us", + cpuset_file_name="/sys/fs/cgroup/cpuset/cpuset.cpus", + cpu_max_file_name="/sys/fs/cgroup/cpu.max", +) -> Optional[float]: + # TODO (Alex): Don't implement this logic oursleves. + # Docker has 2 underyling ways of implementing CPU limits: + # https://docs.docker.com/config/containers/resource_constraints/#configure-the-default-cfs-scheduler + # 1. --cpuset-cpus 2. --cpus or --cpu-quota/--cpu-period (--cpu-shares is a + # soft limit so we don't worry about it). For Ray's purposes, if we use + # docker, the number of vCPUs on a machine is whichever is set (ties broken + # by smaller value). + + cpu_quota = None + # See: https://bugs.openjdk.java.net/browse/JDK-8146115 + if os.path.exists(cpu_quota_file_name) and os.path.exists(cpu_period_file_name): + try: + with ( + open(cpu_quota_file_name, "r") as quota_file, + open(cpu_period_file_name, "r") as period_file, + ): + cpu_quota = float(quota_file.read()) / float(period_file.read()) + except Exception: + logger.exception("Unexpected error calculating docker cpu quota.") + # Look at cpu.max for cgroups v2 + elif os.path.exists(cpu_max_file_name): + try: + max_file = open(cpu_max_file_name).read() + quota_str, period_str = max_file.split() + if quota_str.isnumeric() and period_str.isnumeric(): + cpu_quota = float(quota_str) / float(period_str) + else: + # quota_str is "max" meaning the cpu quota is unset + cpu_quota = None + except Exception: + logger.exception("Unexpected error calculating docker cpu quota.") + if (cpu_quota is not None) and (cpu_quota < 0): + cpu_quota = None + elif cpu_quota == 0: + # Round up in case the cpu limit is less than 1. + cpu_quota = 1 + + cpuset_num = None + if os.path.exists(cpuset_file_name): + try: + with open(cpuset_file_name) as cpuset_file: + ranges_as_string = cpuset_file.read() + ranges = ranges_as_string.split(",") + cpu_ids = [] + for num_or_range in ranges: + if "-" in num_or_range: + start, end = num_or_range.split("-") + cpu_ids.extend(list(range(int(start), int(end) + 1))) + else: + cpu_ids.append(int(num_or_range)) + cpuset_num = len(cpu_ids) + except Exception: + logger.exception("Unexpected error calculating docker cpuset ids.") + # Possible to-do: Parse cgroups v2's cpuset.cpus.effective for the number + # of accessible CPUs. + + if cpu_quota and cpuset_num: + return min(cpu_quota, cpuset_num) + return cpu_quota or cpuset_num + + +def get_num_cpus( + override_docker_cpu_warning: bool = ENV_DISABLE_DOCKER_CPU_WARNING, + truncate: bool = True, +) -> float: + """ + Get the number of CPUs available on this node. + Depending on the situation, use multiprocessing.cpu_count() or cgroups. + + Args: + override_docker_cpu_warning: An extra flag to explicitly turn off the Docker + warning. Setting this flag True has the same effect as setting the env + RAY_DISABLE_DOCKER_CPU_WARNING. By default, whether or not to log + the warning is determined by the env variable + RAY_DISABLE_DOCKER_CPU_WARNING. + truncate: truncates the return value and drops the decimal part. + """ + cpu_count = multiprocessing.cpu_count() + if os.environ.get("RAY_USE_MULTIPROCESSING_CPU_COUNT"): + logger.info( + "Detected RAY_USE_MULTIPROCESSING_CPU_COUNT=1: Using " + "multiprocessing.cpu_count() to detect the number of CPUs. " + "This may be inconsistent when used inside docker. " + "To correctly detect CPUs, unset the env var: " + "`RAY_USE_MULTIPROCESSING_CPU_COUNT`." + ) + return cpu_count + try: + # Not easy to get cpu count in docker, see: + # https://bugs.python.org/issue36054 + docker_count = _get_docker_cpus() + if docker_count is not None and docker_count != cpu_count: + # Don't log this warning if we're on K8s or if the warning is + # explicitly disabled. + if ( + "KUBERNETES_SERVICE_HOST" not in os.environ + and not ENV_DISABLE_DOCKER_CPU_WARNING + and not override_docker_cpu_warning + ): + logger.warning( + "Detecting docker specified CPUs. In " + "previous versions of Ray, CPU detection in containers " + "was incorrect. Please ensure that Ray has enough CPUs " + "allocated. As a temporary workaround to revert to the " + "prior behavior, set " + "`RAY_USE_MULTIPROCESSING_CPU_COUNT=1` as an env var " + "before starting Ray. Set the env var: " + "`RAY_DISABLE_DOCKER_CPU_WARNING=1` to mute this warning." + ) + # TODO (Alex): We should probably add support for fractional cpus. + if int(docker_count) != float(docker_count): + logger.warning( + f"Ray currently does not support initializing Ray " + f"with fractional cpus. Your num_cpus will be " + f"truncated from {docker_count} to " + f"{int(docker_count)}." + ) + if truncate: + docker_count = int(docker_count) + cpu_count = docker_count + + except Exception: + # `nproc` and cgroup are linux-only. If docker only works on linux + # (will run in a linux VM on other platforms), so this is fine. + pass + + return cpu_count + + +# TODO(clarng): merge code with c++ +def get_cgroup_used_memory( + memory_stat_filename: str, + memory_usage_filename: str, + inactive_file_key: str, + active_file_key: str, +): + """ + The calculation logic is the same with `GetCGroupMemoryUsedBytes` + in `memory_monitor.cc` file. + """ + inactive_file_bytes = -1 + active_file_bytes = -1 + with open(memory_stat_filename, "r") as f: + lines = f.readlines() + for line in lines: + if f"{inactive_file_key} " in line: + inactive_file_bytes = int(line.split()[1]) + elif f"{active_file_key} " in line: + active_file_bytes = int(line.split()[1]) + + with open(memory_usage_filename, "r") as f: + lines = f.readlines() + cgroup_usage_in_bytes = int(lines[0].strip()) + + if ( + inactive_file_bytes == -1 + or cgroup_usage_in_bytes == -1 + or active_file_bytes == -1 + ): + return None + + return cgroup_usage_in_bytes - inactive_file_bytes - active_file_bytes + + +def get_used_memory(): + """Return the currently used system memory in bytes + + Returns: + The total amount of used memory + """ + # Try to accurately figure out the memory usage if we are in a docker + # container. + docker_usage = None + # For cgroups v1: + memory_usage_filename_v1 = "/sys/fs/cgroup/memory/memory.usage_in_bytes" + memory_stat_filename_v1 = "/sys/fs/cgroup/memory/memory.stat" + # For cgroups v2: + memory_usage_filename_v2 = "/sys/fs/cgroup/memory.current" + memory_stat_filename_v2 = "/sys/fs/cgroup/memory.stat" + if os.path.exists(memory_usage_filename_v1) and os.path.exists( + memory_stat_filename_v1 + ): + docker_usage = get_cgroup_used_memory( + memory_stat_filename_v1, + memory_usage_filename_v1, + "total_inactive_file", + "total_active_file", + ) + elif os.path.exists(memory_usage_filename_v2) and os.path.exists( + memory_stat_filename_v2 + ): + docker_usage = get_cgroup_used_memory( + memory_stat_filename_v2, + memory_usage_filename_v2, + "inactive_file", + "active_file", + ) + + if docker_usage is not None: + return docker_usage + return psutil.virtual_memory().used + + +def estimate_available_memory(): + """Return the currently available amount of system memory in bytes. + + Returns: + The total amount of available memory in bytes. Based on the used + and total memory. + + """ + return get_system_memory() - get_used_memory() + + +def get_shared_memory_bytes(): + """Get the size of the shared memory file system. + + Returns: + The size of the shared memory file system in bytes. + """ + # Make sure this is only called on Linux. + assert sys.platform == "linux" or sys.platform == "linux2" + + shm_fd = os.open("/dev/shm", os.O_RDONLY) + try: + shm_fs_stats = os.fstatvfs(shm_fd) + # The value shm_fs_stats.f_bsize is the block size and the + # value shm_fs_stats.f_bavail is the number of available + # blocks. + shm_avail = shm_fs_stats.f_bsize * shm_fs_stats.f_bavail + finally: + os.close(shm_fd) + + return shm_avail + + +def check_oversized_function( + pickled: bytes, name: str, obj_type: str, worker: "ray.Worker" +) -> None: + """Send a warning message if the pickled function is too large. + + Args: + pickled: the pickled function. + name: name of the pickled object. + obj_type: type of the pickled object, can be 'function', + 'remote function', or 'actor'. + worker: the worker used to send warning message. message will be logged + locally if None. + """ + length = len(pickled) + if length <= ray_constants.FUNCTION_SIZE_WARN_THRESHOLD: + return + elif length < ray_constants.FUNCTION_SIZE_ERROR_THRESHOLD: + warning_message = ( + "The {} {} is very large ({} MiB). " + "Check that its definition is not implicitly capturing a large " + "array or other object in scope. Tip: use ray.put() to put large " + "objects in the Ray object store." + ).format(obj_type, name, length // (1024 * 1024)) + if worker: + push_error_to_driver( + worker, + ray_constants.PICKLING_LARGE_OBJECT_PUSH_ERROR, + "Warning: " + warning_message, + job_id=worker.current_job_id, + ) + else: + error = ( + "The {} {} is too large ({} MiB > FUNCTION_SIZE_ERROR_THRESHOLD={}" + " MiB). Check that its definition is not implicitly capturing a " + "large array or other object in scope. Tip: use ray.put() to " + "put large objects in the Ray object store." + ).format( + obj_type, + name, + length // (1024 * 1024), + ray_constants.FUNCTION_SIZE_ERROR_THRESHOLD // (1024 * 1024), + ) + raise ValueError(error) + + +def is_main_thread(): + return threading.current_thread().getName() == "MainThread" + + +def detect_fate_sharing_support_win32(): + global win32_job, win32_AssignProcessToJobObject + if win32_job is None and sys.platform == "win32": + import ctypes + + try: + from ctypes.wintypes import BOOL, DWORD, HANDLE, LPCWSTR, LPVOID + + kernel32 = ctypes.WinDLL("kernel32") + kernel32.CreateJobObjectW.argtypes = (LPVOID, LPCWSTR) + kernel32.CreateJobObjectW.restype = HANDLE + sijo_argtypes = (HANDLE, ctypes.c_int, LPVOID, DWORD) + kernel32.SetInformationJobObject.argtypes = sijo_argtypes + kernel32.SetInformationJobObject.restype = BOOL + kernel32.AssignProcessToJobObject.argtypes = (HANDLE, HANDLE) + kernel32.AssignProcessToJobObject.restype = BOOL + kernel32.IsDebuggerPresent.argtypes = () + kernel32.IsDebuggerPresent.restype = BOOL + except (AttributeError, TypeError, ImportError): + kernel32 = None + job = kernel32.CreateJobObjectW(None, None) if kernel32 else None + job = subprocess.Handle(job) if job else job + if job: + from ctypes.wintypes import DWORD, LARGE_INTEGER, ULARGE_INTEGER + + class JOBOBJECT_BASIC_LIMIT_INFORMATION(ctypes.Structure): + _fields_ = [ + ("PerProcessUserTimeLimit", LARGE_INTEGER), + ("PerJobUserTimeLimit", LARGE_INTEGER), + ("LimitFlags", DWORD), + ("MinimumWorkingSetSize", ctypes.c_size_t), + ("MaximumWorkingSetSize", ctypes.c_size_t), + ("ActiveProcessLimit", DWORD), + ("Affinity", ctypes.c_size_t), + ("PriorityClass", DWORD), + ("SchedulingClass", DWORD), + ] + + class IO_COUNTERS(ctypes.Structure): + _fields_ = [ + ("ReadOperationCount", ULARGE_INTEGER), + ("WriteOperationCount", ULARGE_INTEGER), + ("OtherOperationCount", ULARGE_INTEGER), + ("ReadTransferCount", ULARGE_INTEGER), + ("WriteTransferCount", ULARGE_INTEGER), + ("OtherTransferCount", ULARGE_INTEGER), + ] + + class JOBOBJECT_EXTENDED_LIMIT_INFORMATION(ctypes.Structure): + _fields_ = [ + ("BasicLimitInformation", JOBOBJECT_BASIC_LIMIT_INFORMATION), + ("IoInfo", IO_COUNTERS), + ("ProcessMemoryLimit", ctypes.c_size_t), + ("JobMemoryLimit", ctypes.c_size_t), + ("PeakProcessMemoryUsed", ctypes.c_size_t), + ("PeakJobMemoryUsed", ctypes.c_size_t), + ] + + debug = kernel32.IsDebuggerPresent() + + # Defined in ; also available here: + # https://docs.microsoft.com/en-us/windows/win32/api/jobapi2/nf-jobapi2-setinformationjobobject + JobObjectExtendedLimitInformation = 9 + JOB_OBJECT_LIMIT_BREAKAWAY_OK = 0x00000800 + JOB_OBJECT_LIMIT_DIE_ON_UNHANDLED_EXCEPTION = 0x00000400 + JOB_OBJECT_LIMIT_KILL_ON_JOB_CLOSE = 0x00002000 + buf = JOBOBJECT_EXTENDED_LIMIT_INFORMATION() + buf.BasicLimitInformation.LimitFlags = ( + (0 if debug else JOB_OBJECT_LIMIT_KILL_ON_JOB_CLOSE) + | JOB_OBJECT_LIMIT_DIE_ON_UNHANDLED_EXCEPTION + | JOB_OBJECT_LIMIT_BREAKAWAY_OK + ) + infoclass = JobObjectExtendedLimitInformation + if not kernel32.SetInformationJobObject( + job, infoclass, ctypes.byref(buf), ctypes.sizeof(buf) + ): + job = None + win32_AssignProcessToJobObject = ( + kernel32.AssignProcessToJobObject if kernel32 is not None else False + ) + win32_job = job if job else False + return bool(win32_job) + + +def detect_fate_sharing_support_linux(): + global linux_prctl + if linux_prctl is None and sys.platform.startswith("linux"): + try: + from ctypes import CDLL, c_int, c_ulong + + prctl = CDLL(None).prctl + prctl.restype = c_int + prctl.argtypes = [c_int, c_ulong, c_ulong, c_ulong, c_ulong] + except (AttributeError, TypeError): + prctl = None + linux_prctl = prctl if prctl else False + return bool(linux_prctl) + + +def detect_fate_sharing_support(): + result = None + if sys.platform == "win32": + result = detect_fate_sharing_support_win32() + elif sys.platform.startswith("linux"): + result = detect_fate_sharing_support_linux() + return result + + +def set_kill_on_parent_death_linux(): + """Ensures this process dies if its parent dies (fate-sharing). + + Linux-only. Must be called in preexec_fn (i.e. by the child). + """ + if detect_fate_sharing_support_linux(): + import signal + + PR_SET_PDEATHSIG = 1 + if linux_prctl(PR_SET_PDEATHSIG, signal.SIGKILL, 0, 0, 0) != 0: + import ctypes + + raise OSError(ctypes.get_errno(), "prctl(PR_SET_PDEATHSIG) failed") + else: + assert False, "PR_SET_PDEATHSIG used despite being unavailable" + + +def set_kill_child_on_death_win32(child_proc): + """Ensures the child process dies if this process dies (fate-sharing). + + Windows-only. Must be called by the parent, after spawning the child. + + Args: + child_proc: The subprocess.Popen or subprocess.Handle object. + """ + + if isinstance(child_proc, subprocess.Popen): + child_proc = child_proc._handle + assert isinstance(child_proc, subprocess.Handle) + + if detect_fate_sharing_support_win32(): + if not win32_AssignProcessToJobObject(win32_job, int(child_proc)): + import ctypes + + raise OSError(ctypes.get_last_error(), "AssignProcessToJobObject() failed") + else: + assert False, "AssignProcessToJobObject used despite being unavailable" + + +def set_sigterm_handler(sigterm_handler): + """Registers a handler for SIGTERM in a platform-compatible manner.""" + if sys.platform == "win32": + # Note that these signal handlers only work for console applications. + # TODO(mehrdadn): implement graceful process termination mechanism + # SIGINT is Ctrl+C, SIGBREAK is Ctrl+Break. + signal.signal(signal.SIGBREAK, sigterm_handler) + else: + signal.signal(signal.SIGTERM, sigterm_handler) + + +def try_to_symlink(symlink_path, target_path): + """Attempt to create a symlink. + + If the symlink path exists and isn't a symlink, the symlink will not be + created. If a symlink exists in the path, it will be attempted to be + removed and replaced. + + Args: + symlink_path: The path at which to create the symlink. + target_path: The path the symlink should point to. + """ + symlink_path = os.path.expanduser(symlink_path) + target_path = os.path.expanduser(target_path) + + if os.path.exists(symlink_path): + if os.path.islink(symlink_path): + # Try to remove existing symlink. + try: + os.remove(symlink_path) + except OSError: + return + else: + # There's an existing non-symlink file, don't overwrite it. + return + + try: + os.symlink(target_path, symlink_path) + except OSError: + return + + +def get_user(): + if pwd is None: + return "" + try: + return pwd.getpwuid(os.getuid()).pw_name + except Exception: + return "" + + +def get_conda_bin_executable(executable_name): + """ + Return path to the specified executable, assumed to be discoverable within + the 'bin' subdirectory of a conda installation. Adapted from + https://github.com/mlflow/mlflow. + """ + + # Use CONDA_EXE as per https://github.com/conda/conda/issues/7126 + if "CONDA_EXE" in os.environ: + conda_bin_dir = os.path.dirname(os.environ["CONDA_EXE"]) + return os.path.join(conda_bin_dir, executable_name) + return executable_name + + +def get_conda_env_dir(env_name): + """Find and validate the conda directory for a given conda environment. + + For example, given the environment name `tf1`, this function checks + the existence of the corresponding conda directory, e.g. + `/Users/scaly/anaconda3/envs/tf1`, and returns it. + """ + conda_prefix = os.environ.get("CONDA_PREFIX") + if conda_prefix is None: + # The caller is neither in a conda env or in (base) env. This is rare + # because by default, new terminals start in (base), but we can still + # support this case. + conda_exe = os.environ.get("CONDA_EXE") + if conda_exe is None: + raise ValueError( + "Cannot find environment variables set by conda. " + "Please verify conda is installed." + ) + # Example: CONDA_EXE=$HOME/anaconda3/bin/python + # Strip out /bin/python by going up two parent directories. + conda_prefix = str(Path(conda_exe).parent.parent) + + # There are two cases: + # 1. We are in a conda (base) env: CONDA_DEFAULT_ENV=base and + # CONDA_PREFIX=$HOME/anaconda3 + # 2. We are in a user-created conda env: CONDA_DEFAULT_ENV=$env_name and + # CONDA_PREFIX=$HOME/anaconda3/envs/$current_env_name + if os.environ.get("CONDA_DEFAULT_ENV") == "base": + # Caller's curent environment is (base). + # Not recommended by conda, but we can still support it. + if env_name == "base": + # Desired environment is (base), located at e.g. $HOME/anaconda3 + env_dir = conda_prefix + else: + # Desired environment is user-created, e.g. + # $HOME/anaconda3/envs/$env_name + env_dir = os.path.join(conda_prefix, "envs", env_name) + else: + # Now `conda_prefix` should be something like + # $HOME/anaconda3/envs/$current_env_name + # We want to replace the last component with the desired env name. + conda_envs_dir = os.path.split(conda_prefix)[0] + env_dir = os.path.join(conda_envs_dir, env_name) + if not os.path.isdir(env_dir): + raise ValueError( + "conda env " + + env_name + + " not found in conda envs directory. Run `conda env list` to " + + "verify the name is correct." + ) + return env_dir + + +def get_ray_doc_version(): + """Get the docs.ray.io version corresponding to the ray.__version__.""" + # The ray.__version__ can be official Ray release (such as 1.12.0), or + # dev (3.0.0dev0) or release candidate (2.0.0rc0). For the later we map + # to the master doc version at docs.ray.io. + if re.match(r"^\d+\.\d+\.\d+$", ray.__version__) is None: + return "master" + # For the former (official Ray release), we have corresponding doc version + # released as well. + return f"releases-{ray.__version__}" + + +# Used to only print a deprecation warning once for a given function if we +# don't wish to spam the caller. +def get_wheel_filename( + sys_platform: str = sys.platform, + ray_version: str = ray.__version__, + py_version: Tuple[int, int] = (sys.version_info.major, sys.version_info.minor), + architecture: Optional[str] = None, +) -> str: + """Returns the filename used for the nightly Ray wheel. + + Args: + sys_platform: The platform as returned by sys.platform. Examples: + "darwin", "linux", "win32" + ray_version: The Ray version as returned by ray.__version__ or + `ray --version`. Examples: "3.0.0.dev0" + py_version: The Python version as returned by sys.version_info. A + tuple of (major, minor). Examples: (3, 8) + architecture: Architecture, e.g. ``x86_64`` or ``aarch64``. If None, will + be determined by calling ``platform.processor()``. + + Returns: + The wheel file name. Examples: + ray-3.0.0.dev0-cp38-cp38-manylinux2014_x86_64.whl + """ + assert py_version in ray_constants.RUNTIME_ENV_CONDA_PY_VERSIONS, py_version + + py_version_str = "".join(map(str, py_version)) + + architecture = architecture or platform.processor() + + assert sys_platform in ["darwin", "linux", "win32"], sys_platform + + if sys_platform == "darwin": + if architecture == "x86_64": + os_string = "macosx_12_0_x86_64" + else: + os_string = "macosx_12_0_arm64" + elif sys_platform == "linux": + if architecture == "aarch64" or architecture == "arm64": + os_string = "manylinux2014_aarch64" + else: + os_string = "manylinux2014_x86_64" + elif sys_platform == "win32": + os_string = "win_amd64" + + wheel_filename = ( + f"ray-{ray_version}-cp{py_version_str}-" + f"cp{py_version_str}{'m' if py_version_str in ['37'] else ''}" + f"-{os_string}.whl" + ) + + return wheel_filename + + +def get_master_wheel_url( + ray_commit: str = ray.__commit__, + sys_platform: str = sys.platform, + ray_version: str = ray.__version__, + py_version: Tuple[int, int] = sys.version_info[:2], +) -> str: + """Return the URL for the wheel from a specific commit.""" + filename = get_wheel_filename( + sys_platform=sys_platform, ray_version=ray_version, py_version=py_version + ) + return ( + f"https://s3-us-west-2.amazonaws.com/ray-wheels/master/" + f"{ray_commit}/{filename}" + ) + + +def get_release_wheel_url( + ray_commit: str = ray.__commit__, + sys_platform: str = sys.platform, + ray_version: str = ray.__version__, + py_version: Tuple[int, int] = sys.version_info[:2], +) -> str: + """Return the URL for the wheel for a specific release.""" + filename = get_wheel_filename( + sys_platform=sys_platform, ray_version=ray_version, py_version=py_version + ) + return ( + f"https://ray-wheels.s3-us-west-2.amazonaws.com/releases/" + f"{ray_version}/{ray_commit}/{filename}" + ) + # e.g. https://ray-wheels.s3-us-west-2.amazonaws.com/releases/1.4.0rc1/e7c7 + # f6371a69eb727fa469e4cd6f4fbefd143b4c/ray-1.4.0rc1-cp36-cp36m-manylinux201 + # 4_x86_64.whl + + +def validate_namespace(namespace: str): + if not isinstance(namespace, str): + raise TypeError("namespace must be None or a string.") + elif namespace == "": + raise ValueError( + '"" is not a valid namespace. ' "Pass None to not specify a namespace." + ) + + +def get_dashboard_dependency_error() -> Optional[ImportError]: + """Returns the exception error if Ray Dashboard dependencies are not installed. + None if they are installed. + + Checks to see if we should start the dashboard agent or not based on the + Ray installation version the user has installed (ray vs. ray[default]). + Unfortunately there doesn't seem to be a cleaner way to detect this other + than just blindly importing the relevant packages. + """ + try: + import ray.dashboard.optional_deps # noqa: F401 + + return None + except ImportError as e: + return e + + +def get_ray_client_dependency_error() -> Optional[ImportError]: + """Returns the exception error if Ray Client dependencies are not installed. + None if they are installed. + """ + try: + import grpc # noqa: F401 + + return None + except ImportError as e: + return e + + +connect_error = ( + "Unable to connect to GCS (ray head) at {}. " + "Check that (1) Ray with matching version started " + "successfully at the specified address, (2) this " + "node can reach the specified address, and (3) there is " + "no firewall setting preventing access." +) + + +def internal_kv_list_with_retry(gcs_client, prefix, namespace, num_retries=20): + result = None + if isinstance(prefix, str): + prefix = prefix.encode() + if isinstance(namespace, str): + namespace = namespace.encode() + for _ in range(num_retries): + try: + result = gcs_client.internal_kv_keys(prefix, namespace) + except Exception as e: + if isinstance(e, ray.exceptions.RpcError) and e.rpc_code in ( + ray._raylet.GRPC_STATUS_CODE_UNAVAILABLE, + ray._raylet.GRPC_STATUS_CODE_UNKNOWN, + ): + logger.warning(connect_error.format(gcs_client.address)) + else: + logger.exception("Internal KV List failed") + result = None + + if result is not None: + break + else: + logger.debug(f"Fetched {prefix}=None from KV. Retrying.") + time.sleep(2) + if result is None: + raise ConnectionError( + f"Could not list '{prefix}' from GCS. Did GCS start successfully?" + ) + return result + + +def internal_kv_get_with_retry(gcs_client, key, namespace, num_retries=20): + result = None + if isinstance(key, str): + key = key.encode() + for _ in range(num_retries): + try: + result = gcs_client.internal_kv_get(key, namespace) + except Exception as e: + if isinstance(e, ray.exceptions.RpcError) and e.rpc_code in ( + ray._raylet.GRPC_STATUS_CODE_UNAVAILABLE, + ray._raylet.GRPC_STATUS_CODE_UNKNOWN, + ): + logger.warning(connect_error.format(gcs_client.address)) + else: + logger.exception("Internal KV Get failed") + result = None + + if result is not None: + break + else: + logger.debug(f"Fetched {key}=None from KV. Retrying.") + time.sleep(2) + if not result: + raise ConnectionError( + f"Could not read '{key.decode()}' from GCS. Did GCS start successfully?" + ) + return result + + +def parse_resources_json( + resources: str, cli_logger, cf, command_arg="--resources" +) -> Dict[str, float]: + try: + resources = json.loads(resources) + if not isinstance(resources, dict): + raise ValueError("The format after deserialization is not a dict") + except Exception as e: + cli_logger.error( + "`{}` is not a valid JSON string, detail error:{}", + cf.bold(f"{command_arg}={resources}"), + str(e), + ) + cli_logger.abort( + "Valid values look like this: `{}`", + cf.bold( + f'{command_arg}=\'{{"CustomResource3": 1, "CustomResource2": 2}}\'' + ), + ) + return resources + + +def parse_metadata_json( + metadata: str, cli_logger, cf, command_arg="--metadata-json" +) -> Dict[str, str]: + try: + metadata = json.loads(metadata) + if not isinstance(metadata, dict): + raise ValueError("The format after deserialization is not a dict") + except Exception as e: + cli_logger.error( + "`{}` is not a valid JSON string, detail error:{}", + cf.bold(f"{command_arg}={metadata}"), + str(e), + ) + cli_logger.abort( + "Valid values look like this: `{}`", + cf.bold(f'{command_arg}=\'{{"key1": "value1", "key2": "value2"}}\''), + ) + return metadata + + +def internal_kv_put_with_retry(gcs_client, key, value, namespace, num_retries=20): + if isinstance(key, str): + key = key.encode() + if isinstance(value, str): + value = value.encode() + if isinstance(namespace, str): + namespace = namespace.encode() + error = None + for _ in range(num_retries): + try: + return gcs_client.internal_kv_put( + key, value, overwrite=True, namespace=namespace + ) + except ray.exceptions.RpcError as e: + if e.rpc_code in ( + ray._raylet.GRPC_STATUS_CODE_UNAVAILABLE, + ray._raylet.GRPC_STATUS_CODE_UNKNOWN, + ): + logger.warning(connect_error.format(gcs_client.address)) + else: + logger.exception("Internal KV Put failed") + time.sleep(2) + error = e + # Reraise the last error. + raise error + + +def compute_version_info(): + """Compute the versions of Python, and Ray. + + Returns: + A tuple containing the version information. + """ + ray_version = ray.__version__ + python_version = ".".join(map(str, sys.version_info[:3])) + return ray_version, python_version + + +def get_directory_size_bytes(path: Union[str, Path] = ".") -> int: + """Get the total size of a directory in bytes, including subdirectories.""" + total_size_bytes = 0 + for dirpath, dirnames, filenames in os.walk(path): + for f in filenames: + fp = os.path.join(dirpath, f) + # skip if it is a symbolic link or a .pyc file + if not os.path.islink(fp) and not f.endswith(".pyc"): + total_size_bytes += os.path.getsize(fp) + + return total_size_bytes + + +def check_version_info( + cluster_metadata, + this_process_address, + raise_on_mismatch=True, + python_version_match_level=None, +): + """Check if the Python and Ray versions stored in GCS matches this process. + Args: + cluster_metadata: Ray cluster metadata from GCS. + this_process_address: Informational only. The address of this process. + e.g. "node address:port" or "Ray Client". + raise_on_mismatch: Raise an exception on True, log a warning otherwise. + python_version_match_level: "minor" or "patch". To which python version level we + try to match. Note if "minor" and the patch is different, we will still log + a warning. Default value is `RAY_DEFAULT_PYTHON_VERSION_MATCH_LEVEL` if it + exists, otherwise "patch" + + Behavior: + - We raise or log a warning, based on raise_on_mismatch, if: + - Ray versions do not match; OR + - Python (major, minor) versions do not match, + if python_version_match_level == 'minor'; OR + - Python (major, minor, patch) versions do not match, + if python_version_match_level == 'patch'. + - We also log a warning if: + - Python (major, minor) versions match, AND + - Python patch versions do not match, AND + - python_version_match_level == 'minor' AND + - raise_on_mismatch == False. + + Raises: + Exception: An exception is raised if there is a version mismatch. + """ + if python_version_match_level is None: + python_version_match_level = os.environ.get( + "RAY_DEFAULT_PYTHON_VERSION_MATCH_LEVEL", "patch" + ) + + cluster_version_info = ( + cluster_metadata["ray_version"], + cluster_metadata["python_version"], + ) + my_version_info = compute_version_info() + + # Calculate: ray_matches, python_matches, python_full_matches + ray_matches = cluster_version_info[0] == my_version_info[0] + python_full_matches = cluster_version_info[1] == my_version_info[1] + if python_version_match_level == "patch": + python_matches = cluster_version_info[1] == my_version_info[1] + elif python_version_match_level == "minor": + my_python_versions = my_version_info[1].split(".") + cluster_python_versions = cluster_version_info[1].split(".") + python_matches = my_python_versions[:2] == cluster_python_versions[:2] + else: + raise ValueError( + f"Invalid python_version_match_level: {python_version_match_level}, " + "want: 'minor' or 'patch'" + ) + + mismatch_msg = ( + "The cluster was started with:\n" + f" Ray: {cluster_version_info[0]}\n" + f" Python: {cluster_version_info[1]}\n" + f"This process on {this_process_address} was started with:\n" + f" Ray: {my_version_info[0]}\n" + f" Python: {my_version_info[1]}\n" + ) + + if ray_matches and python_matches: + if not python_full_matches: + logger.warning(f"Python patch version mismatch: {mismatch_msg}") + else: + error_message = f"Version mismatch: {mismatch_msg}" + if raise_on_mismatch: + raise RuntimeError(error_message) + else: + logger.warning(error_message) + + +def get_runtime_env_info( + runtime_env: "RuntimeEnv", + *, + is_job_runtime_env: bool = False, + serialize: bool = False, +): + """Create runtime env info from runtime env. + + In the user interface, the argument `runtime_env` contains some fields + which not contained in `ProtoRuntimeEnv` but in `ProtoRuntimeEnvInfo`, + such as `eager_install`. This function will extract those fields from + `RuntimeEnv` and create a new `ProtoRuntimeEnvInfo`, and serialize it + into json format. + """ + from ray.runtime_env import RuntimeEnvConfig + + proto_runtime_env_info = ProtoRuntimeEnvInfo() + + if runtime_env.working_dir_uri(): + proto_runtime_env_info.uris.working_dir_uri = runtime_env.working_dir_uri() + if len(runtime_env.py_modules_uris()) > 0: + proto_runtime_env_info.uris.py_modules_uris[:] = runtime_env.py_modules_uris() + + # TODO(Catch-Bull): overload `__setitem__` for `RuntimeEnv`, change the + # runtime_env of all internal code from dict to RuntimeEnv. + + runtime_env_config = runtime_env.get("config") + if runtime_env_config is None: + runtime_env_config = RuntimeEnvConfig.default_config() + else: + runtime_env_config = RuntimeEnvConfig.parse_and_validate_runtime_env_config( + runtime_env_config + ) + + proto_runtime_env_info.runtime_env_config.CopyFrom( + runtime_env_config.build_proto_runtime_env_config() + ) + + # Normally, `RuntimeEnv` should guarantee the accuracy of field eager_install, + # but so far, the internal code has not completely prohibited direct + # modification of fields in RuntimeEnv, so we should check it for insurance. + eager_install = ( + runtime_env_config.get("eager_install") + if runtime_env_config is not None + else None + ) + if is_job_runtime_env or eager_install is not None: + if eager_install is None: + eager_install = True + elif not isinstance(eager_install, bool): + raise TypeError( + f"eager_install must be a boolean. got {type(eager_install)}" + ) + proto_runtime_env_info.runtime_env_config.eager_install = eager_install + + proto_runtime_env_info.serialized_runtime_env = runtime_env.serialize() + + if not serialize: + return proto_runtime_env_info + + return json_format.MessageToJson(proto_runtime_env_info) + + +def parse_runtime_env_for_task_or_actor( + runtime_env: Optional[Union[Dict, "RuntimeEnv"]] +): + from ray.runtime_env import RuntimeEnv + from ray.runtime_env.runtime_env import _validate_no_local_paths + + # Parse local pip/conda config files here. If we instead did it in + # .remote(), it would get run in the Ray Client server, which runs on + # a remote node where the files aren't available. + if runtime_env: + if isinstance(runtime_env, dict): + runtime_env = RuntimeEnv(**(runtime_env or {})) + _validate_no_local_paths(runtime_env) + return runtime_env + raise TypeError( + "runtime_env must be dict or RuntimeEnv, ", + f"but got: {type(runtime_env)}", + ) + else: + # Keep the new_runtime_env as None. In .remote(), we need to know + # if runtime_env is None to know whether or not to fall back to the + # runtime_env specified in the @ray.remote decorator. + return None + + +def split_address(address: str) -> Tuple[str, str]: + """Splits address into a module string (scheme) and an inner_address. + + We use a custom splitting function instead of urllib because + PEP allows "underscores" in a module names, while URL schemes do not + allow them. + + Args: + address: The address to split. + + Returns: + A tuple of (scheme, inner_address). + + Raises: + ValueError: If the address does not contain '://'. + + Examples: + >>> split_address("ray://my_cluster") + ('ray', 'my_cluster') + """ + if "://" not in address: + raise ValueError("Address must contain '://'") + + module_string, inner_address = address.split("://", maxsplit=1) + return (module_string, inner_address) + + +def get_entrypoint_name(): + """Get the entrypoint of the current script.""" + prefix = "" + try: + curr = psutil.Process() + # Prepend `interactive_shell` for interactive shell scripts. + # https://stackoverflow.com/questions/2356399/tell-if-python-is-in-interactive-mode # noqa + if hasattr(sys, "ps1"): + prefix = "(interactive_shell) " + + return prefix + list2cmdline(curr.cmdline()) + except Exception: + return "unknown" + + +class DeferSigint(contextlib.AbstractContextManager): + """Context manager that defers SIGINT signals until the context is left.""" + + # This is used by Ray's task cancellation to defer cancellation interrupts during + # problematic areas, e.g. task argument deserialization. + def __init__(self): + # Whether a SIGINT signal was received during the context. + self.signal_received = False + # The overridden SIGINT handler + self.overridden_sigint_handler = None + # The original signal method. + self.orig_signal = None + + @classmethod + def create_if_main_thread(cls) -> contextlib.AbstractContextManager: + """Creates a DeferSigint context manager if running on the main thread, + returns a no-op context manager otherwise. + """ + if threading.current_thread() == threading.main_thread(): + return cls() + else: + return contextlib.nullcontext() + + def _set_signal_received(self, signum, frame): + """SIGINT handler that defers the signal.""" + self.signal_received = True + + def _signal_monkey_patch(self, signum, handler): + """Monkey patch for signal.signal that defers the setting of new signal + handler after the DeferSigint context exits.""" + # Only handle it in the main thread because if setting a handler in a non-main + # thread, signal.signal will raise an error because Python doesn't allow it. + if ( + threading.current_thread() == threading.main_thread() + and signum == signal.SIGINT + ): + orig_sigint_handler = self.overridden_sigint_handler + self.overridden_sigint_handler = handler + return orig_sigint_handler + return self.orig_signal(signum, handler) + + def __enter__(self): + # Save original SIGINT handler for later restoration. + self.overridden_sigint_handler = signal.getsignal(signal.SIGINT) + # Set SIGINT signal handler that defers the signal. + signal.signal(signal.SIGINT, self._set_signal_received) + # Monkey patch signal.signal to raise an error if a SIGINT handler is registered + # within the context. + self.orig_signal = signal.signal + signal.signal = self._signal_monkey_patch + return self + + def __exit__(self, exc_type, exc, exc_tb): + assert self.overridden_sigint_handler is not None + assert self.orig_signal is not None + # Restore original signal.signal function. + signal.signal = self.orig_signal + # Restore overridden SIGINT handler. + signal.signal(signal.SIGINT, self.overridden_sigint_handler) + if exc_type is None and self.signal_received: + # No exception raised in context, call the original SIGINT handler. + # By default, this means raising KeyboardInterrupt. + self.overridden_sigint_handler(signal.SIGINT, None) + else: + # If exception was raised in context, returning False will cause it to be + # reraised. + return False + + +def try_import_each_module(module_names_to_import: List[str]) -> None: + """ + Make a best-effort attempt to import each named Python module. + This is used by the Python default_worker.py to preload modules. + """ + for module_to_preload in module_names_to_import: + try: + importlib.import_module(module_to_preload) + except ImportError: + logger.exception(f'Failed to preload the module "{module_to_preload}"') + + +def remove_ray_internal_flags_from_env(env: dict): + """ + Remove Ray internal flags from `env`. + Defined in ray/common/ray_internal_flag_def.h + """ + for flag in ray_constants.RAY_INTERNAL_FLAGS: + env.pop(flag, None) + + +def update_envs(env_vars: Dict[str, str]): + """ + When updating the environment variable, if there is ${X}, + it will be replaced with the current environment variable. + """ + if not env_vars: + return + + for key, value in env_vars.items(): + expanded = os.path.expandvars(value) + # Replace non-existant env vars with an empty string. + result = re.sub(r"\$\{[A-Z0-9_]+\}", "", expanded) + os.environ[key] = result + + +def parse_pg_formatted_resources_to_original( + pg_formatted_resources: Dict[str, float] +) -> Dict[str, float]: + original_resources = {} + + for key, value in pg_formatted_resources.items(): + result = PLACEMENT_GROUP_INDEXED_BUNDLED_RESOURCE_PATTERN.match(key) + if result and len(result.groups()) == 3: + # Filter out resources that have bundle_group_[pg_id] since + # it is an implementation detail. + # This resource is automatically added to the resource + # request for all tasks that require placement groups. + if result.group(1) == PLACEMENT_GROUP_BUNDLE_RESOURCE_NAME: + continue + + original_resources[result.group(1)] = value + continue + + result = PLACEMENT_GROUP_WILDCARD_RESOURCE_PATTERN.match(key) + if result and len(result.groups()) == 2: + if result.group(1) == "bundle": + continue + + original_resources[result.group(1)] = value + continue + original_resources[key] = value + + return original_resources + + +def validate_actor_state_name(actor_state_name): + if actor_state_name is None: + return + actor_state_names = [ + "DEPENDENCIES_UNREADY", + "PENDING_CREATION", + "ALIVE", + "RESTARTING", + "DEAD", + ] + if actor_state_name not in actor_state_names: + raise ValueError( + f'"{actor_state_name}" is not a valid actor state name, ' + 'it must be one of the following: "DEPENDENCIES_UNREADY", ' + '"PENDING_CREATION", "ALIVE", "RESTARTING", or "DEAD"' + ) + + +def get_current_node_cpu_model_name() -> Optional[str]: + if not sys.platform.startswith("linux"): + return None + + try: + """ + /proc/cpuinfo content example: + + processor : 0 + vendor_id : GenuineIntel + cpu family : 6 + model : 85 + model name : Intel(R) Xeon(R) Platinum 8259CL CPU @ 2.50GHz + stepping : 7 + """ + with open("/proc/cpuinfo", "r") as f: + for line in f: + if line.startswith("model name"): + return line.split(":")[1].strip() + return None + except Exception: + logger.debug("Failed to get CPU model name", exc_info=True) + return None + + +def validate_socket_filepath(filepath: str): + """ + Validate the provided filename is a valid Unix socket filename. + """ + # Don't check for Windows as it doesn't support Unix sockets. + if sys.platform == "win32": + return + is_mac = sys.platform.startswith("darwin") + maxlen = (104 if is_mac else 108) - 1 + if len(filepath.encode("utf-8")) > maxlen: + raise OSError( + f"validate_socket_filename failed: AF_UNIX path length cannot exceed {maxlen} bytes: {filepath}" + ) + + +# Whether we're currently running in a test, either local or CI. +in_test = None + + +def is_in_test(): + global in_test + + if in_test is None: + in_test = any( + env_var in os.environ + # These environment variables are always set by pytest and Buildkite, + # respectively. + for env_var in ("PYTEST_CURRENT_TEST", "BUILDKITE") + ) + return in_test diff --git a/lib/python3.12/site-packages/ray/_private/worker.py b/lib/python3.12/site-packages/ray/_private/worker.py new file mode 100644 index 0000000000000000000000000000000000000000..379fcc550af724e78c29aab14be0b41f2c0e6426 --- /dev/null +++ b/lib/python3.12/site-packages/ray/_private/worker.py @@ -0,0 +1,3795 @@ +import atexit +import faulthandler +import functools +import inspect +import io +import json +import logging +import os +import sys +import threading +import time +import traceback +import urllib +import warnings +from abc import ABCMeta, abstractmethod +from collections.abc import Mapping +from contextlib import contextmanager +from dataclasses import dataclass +from functools import wraps +from typing import ( + TYPE_CHECKING, + Any, + AnyStr, + Callable, + Dict, + Generic, + Iterator, + List, + Literal, + Optional, + Protocol, + Sequence, + Tuple, + Type, + TypeVar, + Union, + overload, +) +from urllib.parse import urlparse + +if TYPE_CHECKING: + import torch + +import colorama + +import ray +import ray._private.node +import ray._private.parameter +import ray._private.profiling as profiling +import ray._private.ray_constants as ray_constants +import ray._private.serialization as serialization +import ray._private.services as services +import ray._private.state +import ray._private.worker + +# Ray modules +import ray.actor +import ray.cloudpickle as pickle # noqa +import ray.job_config +import ray.remote_function +from ray import ActorID, JobID, Language, ObjectRef +from ray._common import ray_option_utils +from ray._common.constants import RAY_WARN_BLOCKING_GET_INSIDE_ASYNC_ENV_VAR +from ray._common.utils import load_class +from ray._private.authentication.authentication_token_setup import ( + ensure_token_if_auth_enabled, +) +from ray._private.client_mode_hook import client_mode_hook +from ray._private.custom_types import TensorTransportEnum +from ray._private.function_manager import FunctionActorManager +from ray._private.inspect_util import is_cython +from ray._private.ray_logging import ( + global_worker_stdstream_dispatcher, + setup_logger, + stderr_deduplicator, + stdout_deduplicator, +) +from ray._private.ray_logging.logging_config import LoggingConfig +from ray._private.resource_isolation_config import ResourceIsolationConfig +from ray._private.runtime_env.constants import RAY_JOB_CONFIG_JSON_ENV_VAR +from ray._private.runtime_env.py_modules import upload_py_modules_if_needed +from ray._private.runtime_env.setup_hook import ( + upload_worker_process_setup_hook_if_needed, +) +from ray._private.runtime_env.working_dir import upload_working_dir_if_needed +from ray._private.utils import get_ray_doc_version +from ray._raylet import ( + ObjectRefGenerator, + TaskID, + raise_sys_exit_with_custom_error_message, +) +from ray.actor import ActorClass +from ray.exceptions import ObjectStoreFullError, RayError, RaySystemError, RayTaskError +from ray.experimental import tqdm_ray +from ray.experimental.compiled_dag_ref import CompiledDAGRef +from ray.experimental.internal_kv import ( + _initialize_internal_kv, + _internal_kv_get, + _internal_kv_initialized, + _internal_kv_reset, +) +from ray.experimental.tqdm_ray import RAY_TQDM_MAGIC +from ray.runtime_env.runtime_env import _merge_runtime_env +from ray.util.annotations import Deprecated, DeveloperAPI, PublicAPI +from ray.util.debug import log_once +from ray.util.scheduling_strategies import PlacementGroupSchedulingStrategy +from ray.util.tracing.tracing_helper import _import_from_string +from ray.widgets import Template +from ray.widgets.util import repr_with_fallback + +SCRIPT_MODE = 0 +WORKER_MODE = 1 +LOCAL_MODE = 2 +SPILL_WORKER_MODE = 3 +RESTORE_WORKER_MODE = 4 + +# Logger for this module. It should be configured at the entry point +# into the program using Ray. Ray provides a default configuration at +# entry/init points. +logger = logging.getLogger(__name__) + + +T = TypeVar("T") +T0 = TypeVar("T0") +T1 = TypeVar("T1") +T2 = TypeVar("T2") +T3 = TypeVar("T3") +T4 = TypeVar("T4") +T5 = TypeVar("T5") +T6 = TypeVar("T6") +T7 = TypeVar("T7") +T8 = TypeVar("T8") +T9 = TypeVar("T9") +R = TypeVar("R") + +DAGNode = TypeVar("DAGNode") + + +# Only used for type annotations as a placeholder +Undefined: Any = object() + + +# TypeVar for self-referential generics in `RemoteFunction[N]`. +RF = TypeVar("RF", bound="HasOptions") + + +class HasOptions(Protocol): + def options(self: RF, **task_options) -> RF: + ... + + +class RemoteFunctionNoArgs(HasOptions, Generic[R]): + def __init__(self, function: Callable[[], R]) -> None: + pass + + def remote( + self, + ) -> "ObjectRef[R]": + ... + + def bind( + self, + ) -> "DAGNode[R]": + ... + + +class RemoteFunction0(HasOptions, Generic[R, T0]): + def __init__(self, function: Callable[[T0], R]) -> None: + pass + + def remote( + self, + __arg0: "Union[T0, ObjectRef[T0]]", + ) -> "ObjectRef[R]": + ... + + def bind( + self, + __arg0: "Union[T0, DAGNode[T0]]", + ) -> "DAGNode[R]": + ... + + +class RemoteFunction1(HasOptions, Generic[R, T0, T1]): + def __init__(self, function: Callable[[T0, T1], R]) -> None: + pass + + def remote( + self, + __arg0: "Union[T0, ObjectRef[T0]]", + __arg1: "Union[T1, ObjectRef[T1]]", + ) -> "ObjectRef[R]": + ... + + def bind( + self, + __arg0: "Union[T0, DAGNode[T0]]", + __arg1: "Union[T1, DAGNode[T1]]", + ) -> "DAGNode[R]": + ... + + +class RemoteFunction2(HasOptions, Generic[R, T0, T1, T2]): + def __init__(self, function: Callable[[T0, T1, T2], R]) -> None: + pass + + def remote( + self, + __arg0: "Union[T0, ObjectRef[T0]]", + __arg1: "Union[T1, ObjectRef[T1]]", + __arg2: "Union[T2, ObjectRef[T2]]", + ) -> "ObjectRef[R]": + ... + + def bind( + self, + __arg0: "Union[T0, DAGNode[T0]]", + __arg1: "Union[T1, DAGNode[T1]]", + __arg2: "Union[T2, DAGNode[T2]]", + ) -> "DAGNode[R]": + ... + + +class RemoteFunction3(HasOptions, Generic[R, T0, T1, T2, T3]): + def __init__(self, function: Callable[[T0, T1, T2, T3], R]) -> None: + pass + + def remote( + self, + __arg0: "Union[T0, ObjectRef[T0]]", + __arg1: "Union[T1, ObjectRef[T1]]", + __arg2: "Union[T2, ObjectRef[T2]]", + __arg3: "Union[T3, ObjectRef[T3]]", + ) -> "ObjectRef[R]": + ... + + def bind( + self, + __arg0: "Union[T0, DAGNode[T0]]", + __arg1: "Union[T1, DAGNode[T1]]", + __arg2: "Union[T2, DAGNode[T2]]", + __arg3: "Union[T3, DAGNode[T3]]", + ) -> "DAGNode[R]": + ... + + +class RemoteFunction4(HasOptions, Generic[R, T0, T1, T2, T3, T4]): + def __init__(self, function: Callable[[T0, T1, T2, T3, T4], R]) -> None: + pass + + def remote( + self, + __arg0: "Union[T0, ObjectRef[T0]]", + __arg1: "Union[T1, ObjectRef[T1]]", + __arg2: "Union[T2, ObjectRef[T2]]", + __arg3: "Union[T3, ObjectRef[T3]]", + __arg4: "Union[T4, ObjectRef[T4]]", + ) -> "ObjectRef[R]": + ... + + def bind( + self, + __arg0: "Union[T0, DAGNode[T0]]", + __arg1: "Union[T1, DAGNode[T1]]", + __arg2: "Union[T2, DAGNode[T2]]", + __arg3: "Union[T3, DAGNode[T3]]", + __arg4: "Union[T4, DAGNode[T4]]", + ) -> "DAGNode[R]": + ... + + +class RemoteFunction5(HasOptions, Generic[R, T0, T1, T2, T3, T4, T5]): + def __init__(self, function: Callable[[T0, T1, T2, T3, T4, T5], R]) -> None: + pass + + def remote( + self, + __arg0: "Union[T0, ObjectRef[T0]]", + __arg1: "Union[T1, ObjectRef[T1]]", + __arg2: "Union[T2, ObjectRef[T2]]", + __arg3: "Union[T3, ObjectRef[T3]]", + __arg4: "Union[T4, ObjectRef[T4]]", + __arg5: "Union[T5, ObjectRef[T5]]", + ) -> "ObjectRef[R]": + ... + + def bind( + self, + __arg0: "Union[T0, DAGNode[T0]]", + __arg1: "Union[T1, DAGNode[T1]]", + __arg2: "Union[T2, DAGNode[T2]]", + __arg3: "Union[T3, DAGNode[T3]]", + __arg4: "Union[T4, DAGNode[T4]]", + __arg5: "Union[T5, DAGNode[T5]]", + ) -> "DAGNode[R]": + ... + + +class RemoteFunction6(HasOptions, Generic[R, T0, T1, T2, T3, T4, T5, T6]): + def __init__(self, function: Callable[[T0, T1, T2, T3, T4, T5, T6], R]) -> None: + pass + + def remote( + self, + __arg0: "Union[T0, ObjectRef[T0]]", + __arg1: "Union[T1, ObjectRef[T1]]", + __arg2: "Union[T2, ObjectRef[T2]]", + __arg3: "Union[T3, ObjectRef[T3]]", + __arg4: "Union[T4, ObjectRef[T4]]", + __arg5: "Union[T5, ObjectRef[T5]]", + __arg6: "Union[T6, ObjectRef[T6]]", + ) -> "ObjectRef[R]": + ... + + def bind( + self, + __arg0: "Union[T0, DAGNode[T0]]", + __arg1: "Union[T1, DAGNode[T1]]", + __arg2: "Union[T2, DAGNode[T2]]", + __arg3: "Union[T3, DAGNode[T3]]", + __arg4: "Union[T4, DAGNode[T4]]", + __arg5: "Union[T5, DAGNode[T5]]", + __arg6: "Union[T6, DAGNode[T6]]", + ) -> "DAGNode[R]": + ... + + +class RemoteFunction7(HasOptions, Generic[R, T0, T1, T2, T3, T4, T5, T6, T7]): + def __init__(self, function: Callable[[T0, T1, T2, T3, T4, T5, T6, T7], R]) -> None: + pass + + def remote( + self, + __arg0: "Union[T0, ObjectRef[T0]]", + __arg1: "Union[T1, ObjectRef[T1]]", + __arg2: "Union[T2, ObjectRef[T2]]", + __arg3: "Union[T3, ObjectRef[T3]]", + __arg4: "Union[T4, ObjectRef[T4]]", + __arg5: "Union[T5, ObjectRef[T5]]", + __arg6: "Union[T6, ObjectRef[T6]]", + __arg7: "Union[T7, ObjectRef[T7]]", + ) -> "ObjectRef[R]": + ... + + def bind( + self, + __arg0: "Union[T0, DAGNode[T0]]", + __arg1: "Union[T1, DAGNode[T1]]", + __arg2: "Union[T2, DAGNode[T2]]", + __arg3: "Union[T3, DAGNode[T3]]", + __arg4: "Union[T4, DAGNode[T4]]", + __arg5: "Union[T5, DAGNode[T5]]", + __arg6: "Union[T6, DAGNode[T6]]", + __arg7: "Union[T7, DAGNode[T7]]", + ) -> "DAGNode[R]": + ... + + +class RemoteFunction8(HasOptions, Generic[R, T0, T1, T2, T3, T4, T5, T6, T7, T8]): + def __init__( + self, function: Callable[[T0, T1, T2, T3, T4, T5, T6, T7, T8], R] + ) -> None: + pass + + def remote( + self, + __arg0: "Union[T0, ObjectRef[T0]]", + __arg1: "Union[T1, ObjectRef[T1]]", + __arg2: "Union[T2, ObjectRef[T2]]", + __arg3: "Union[T3, ObjectRef[T3]]", + __arg4: "Union[T4, ObjectRef[T4]]", + __arg5: "Union[T5, ObjectRef[T5]]", + __arg6: "Union[T6, ObjectRef[T6]]", + __arg7: "Union[T7, ObjectRef[T7]]", + __arg8: "Union[T8, ObjectRef[T8]]", + ) -> "ObjectRef[R]": + ... + + def bind( + self, + __arg0: "Union[T0, DAGNode[T0]]", + __arg1: "Union[T1, DAGNode[T1]]", + __arg2: "Union[T2, DAGNode[T2]]", + __arg3: "Union[T3, DAGNode[T3]]", + __arg4: "Union[T4, DAGNode[T4]]", + __arg5: "Union[T5, DAGNode[T5]]", + __arg6: "Union[T6, DAGNode[T6]]", + __arg7: "Union[T7, DAGNode[T7]]", + __arg8: "Union[T8, DAGNode[T8]]", + ) -> "DAGNode[R]": + ... + + +class RemoteFunction9(HasOptions, Generic[R, T0, T1, T2, T3, T4, T5, T6, T7, T8, T9]): + def __init__( + self, function: Callable[[T0, T1, T2, T3, T4, T5, T6, T7, T8, T9], R] + ) -> None: + pass + + def remote( + self, + __arg0: "Union[T0, ObjectRef[T0]]", + __arg1: "Union[T1, ObjectRef[T1]]", + __arg2: "Union[T2, ObjectRef[T2]]", + __arg3: "Union[T3, ObjectRef[T3]]", + __arg4: "Union[T4, ObjectRef[T4]]", + __arg5: "Union[T5, ObjectRef[T5]]", + __arg6: "Union[T6, ObjectRef[T6]]", + __arg7: "Union[T7, ObjectRef[T7]]", + __arg8: "Union[T8, ObjectRef[T8]]", + __arg9: "Union[T9, ObjectRef[T9]]", + ) -> "ObjectRef[R]": + ... + + def bind( + self, + __arg0: "Union[T0, DAGNode[T0]]", + __arg1: "Union[T1, DAGNode[T1]]", + __arg2: "Union[T2, DAGNode[T2]]", + __arg3: "Union[T3, DAGNode[T3]]", + __arg4: "Union[T4, DAGNode[T4]]", + __arg5: "Union[T5, DAGNode[T5]]", + __arg6: "Union[T6, DAGNode[T6]]", + __arg7: "Union[T7, DAGNode[T7]]", + __arg8: "Union[T8, DAGNode[T8]]", + __arg9: "Union[T9, DAGNode[T9]]", + ) -> "DAGNode[R]": + ... + + +# Visible for testing. +def _unhandled_error_handler(e: Exception): + logger.error( + f"Unhandled error (suppress with 'RAY_IGNORE_UNHANDLED_ERRORS=1'): {e}" + ) + + +class Worker: + """A class used to define the control flow of a worker process. + + Note: + The methods in this class are considered unexposed to the user. The + functions outside of this class are considered exposed. + + Attributes: + node (ray._private.node.Node): The node this worker is attached to. + mode: The mode of the worker. One of SCRIPT_MODE, LOCAL_MODE, and + WORKER_MODE. + """ + + def __init__(self): + """Initialize a Worker object.""" + self.node = None + self.mode = None + self.actors = {} + # GPU object manager to manage GPU object lifecycles, including coordinating out-of-band + # tensor transfers between actors, storing and retrieving GPU objects, and garbage collection. + # We create the GPU object manager lazily, if a user specifies a + # non-default tensor_transport, to avoid circular import and because it + # imports third-party dependencies like PyTorch. + self._gpu_object_manager = None + # When the worker is constructed. Record the original value of the + # (CUDA_VISIBLE_DEVICES, ONEAPI_DEVICE_SELECTOR, HIP_VISIBLE_DEVICES, + # NEURON_RT_VISIBLE_CORES, TPU_VISIBLE_CHIPS, ..) environment variables. + self.original_visible_accelerator_ids = ( + ray._private.utils.get_visible_accelerator_ids() + ) + # A dictionary that maps from driver id to SerializationContext + # TODO: clean up the SerializationContext once the job finished. + self.serialization_context_map = {} + self.function_actor_manager = FunctionActorManager(self) + # This event is checked regularly by all of the threads so that they + # know when to exit. + self.threads_stopped = threading.Event() + # If this is set, the next .remote call should drop into the + # debugger, at the specified breakpoint ID. + self.debugger_breakpoint = b"" + # If this is set, ray.get calls invoked on the object ID returned + # by the worker should drop into the debugger at the specified + # breakpoint ID. + self.debugger_get_breakpoint = b"" + # If True, make the debugger external to the node this worker is + # running on. + self.ray_debugger_external = False + self._load_code_from_local = False + # Opened file descriptor to stdout/stderr for this python worker. + self._enable_record_actor_task_log = ( + ray_constants.RAY_ENABLE_RECORD_ACTOR_TASK_LOGGING + ) + # Whether rotation is enabled for out file and err file, task log position report will be skipped if rotation enabled, since the position cannot be accurate. + self._file_rotation_enabled = False + self._out_filepath = None + self._err_filepath = None + # Create the lock here because the serializer will use it before + # initializing Ray. + self.lock = threading.RLock() + # By default, don't show logs from other drivers. This is set to true by Serve + # in order to stream logs from the controller and replica actors across + # different drivers that connect to the same Serve instance. + # See https://github.com/ray-project/ray/pull/35070. + self._filter_logs_by_job = True + # the debugger port for this worker + self._debugger_port = None + # Cache the job id from initialize_job_config() to optimize lookups. + # This is on the critical path of ray.get()/put() calls. + self._cached_job_id = None + # Indicates whether the worker is connected to the Ray cluster. + # It should be set to True in `connect` and False in `disconnect`. + self._is_connected: bool = False + + @property + def gpu_object_manager(self) -> "ray.experimental.GPUObjectManager": + if self._gpu_object_manager is None: + # We create the GPU object manager lazily, if a user specifies a + # non-default tensor_transport, to avoid circular import and because it + # imports third-party dependencies like PyTorch. + from ray.experimental import GPUObjectManager + + self._gpu_object_manager = GPUObjectManager() + return self._gpu_object_manager + + @property + def connected(self): + """bool: True if Ray has been started and False otherwise.""" + return self._is_connected + + def set_is_connected(self, is_connected: bool): + self._is_connected = is_connected + + @property + def node_ip_address(self): + self.check_connected() + return self.node.node_ip_address + + @property + def load_code_from_local(self): + self.check_connected() + return self._load_code_from_local + + @property + def current_job_id(self): + if self._cached_job_id is not None: + return self._cached_job_id + elif hasattr(self, "core_worker"): + return self.core_worker.get_current_job_id() + return JobID.nil() + + @property + def actor_id(self): + if hasattr(self, "core_worker"): + return self.core_worker.get_actor_id() + return ActorID.nil() + + @property + def actor_name(self): + if hasattr(self, "core_worker"): + return self.core_worker.get_actor_name().decode("utf-8") + return None + + @property + def current_task_id(self): + return self.core_worker.get_current_task_id() + + @property + def current_task_name(self): + return self.core_worker.get_current_task_name() + + @property + def current_task_function_name(self): + return self.core_worker.get_current_task_function_name() + + @property + def current_node_id(self): + return self.core_worker.get_current_node_id() + + @property + def task_depth(self): + return self.core_worker.get_task_depth() + + @property + def namespace(self): + return self.core_worker.get_job_config().ray_namespace + + @property + def placement_group_id(self): + return self.core_worker.get_placement_group_id() + + @property + def worker_id(self): + return self.core_worker.get_worker_id().binary() + + @property + def should_capture_child_tasks_in_placement_group(self): + return self.core_worker.should_capture_child_tasks_in_placement_group() + + @property + def current_cluster_and_job(self): + """Get the current session index and job id as pair.""" + assert isinstance(self.node.cluster_id, ray.ClusterID) + assert isinstance(self.current_job_id, ray.JobID) + return self.node.cluster_id, self.current_job_id + + @property + def runtime_env(self): + """Get the runtime env in json format""" + return self.core_worker.get_current_runtime_env() + + @property + def debugger_port(self): + """Get the debugger port for this worker""" + worker_id = self.core_worker.get_worker_id() + return ray._private.state.get_worker_debugger_port(worker_id) + + @property + def job_logging_config(self): + """Get the job's logging config for this worker""" + if not hasattr(self, "core_worker"): + return None + job_config = self.core_worker.get_job_config() + if not job_config.serialized_py_logging_config: + return None + logging_config = pickle.loads(job_config.serialized_py_logging_config) + return logging_config + + @property + def current_node_labels(self): + # Return the node labels of this worker's current node. + return self.node.node_labels + + def set_debugger_port(self, port): + worker_id = self.core_worker.get_worker_id() + ray._private.state.update_worker_debugger_port(worker_id, port) + + def set_cached_job_id(self, job_id): + """Set the cached job id to speed `current_job_id()`.""" + self._cached_job_id = job_id + + @contextmanager + def task_paused_by_debugger(self): + """Use while the task is paused by debugger""" + try: + self.core_worker.update_task_is_debugger_paused( + ray.get_runtime_context()._get_current_task_id(), True + ) + yield + finally: + self.core_worker.update_task_is_debugger_paused( + ray.get_runtime_context()._get_current_task_id(), False + ) + + @contextmanager + def worker_paused_by_debugger(self): + """ + Updates the worker num paused threads when the worker is paused by debugger + """ + try: + worker_id = self.core_worker.get_worker_id() + ray._private.state.update_worker_num_paused_threads(worker_id, 1) + yield + finally: + ray._private.state.update_worker_num_paused_threads(worker_id, -1) + + def set_file_rotation_enabled(self, rotation_enabled: bool) -> None: + """Set whether rotation is enabled for outfile and errfile.""" + self._file_rotation_enabled = rotation_enabled + + def set_err_file(self, err_filepath=Optional[AnyStr]) -> None: + """Set the worker's err file where stderr is redirected to""" + self._err_filepath = err_filepath + + def set_out_file(self, out_filepath=Optional[AnyStr]) -> None: + """Set the worker's out file where stdout is redirected to""" + self._out_filepath = out_filepath + + def record_task_log_start(self, task_id: TaskID, attempt_number: int): + """Record the task log info when task starts executing for + non concurrent actor tasks.""" + if not self._enable_record_actor_task_log and not self.actor_id.is_nil(): + # We are not recording actor task log if not enabled explicitly. + # Recording actor task log is expensive and should be enabled only + # when needed. + # https://github.com/ray-project/ray/issues/35598 + return + + if not hasattr(self, "core_worker"): + return + if self._file_rotation_enabled: + return + + self.core_worker.record_task_log_start( + task_id, + attempt_number, + self.get_out_file_path(), + self.get_err_file_path(), + self.get_current_out_offset(), + self.get_current_err_offset(), + ) + + def record_task_log_end(self, task_id: TaskID, attempt_number: int): + """Record the task log info when task finishes executing for + non concurrent actor tasks.""" + if not self._enable_record_actor_task_log and not self.actor_id.is_nil(): + # We are not recording actor task log if not enabled explicitly. + # Recording actor task log is expensive and should be enabled only + # when needed. + # https://github.com/ray-project/ray/issues/35598 + return + + if not hasattr(self, "core_worker"): + return + + # Disable file offset fetch if rotation enabled (since file offset doesn't make sense for rotated files). + if self._file_rotation_enabled: + return + + self.core_worker.record_task_log_end( + task_id, + attempt_number, + self.get_current_out_offset(), + self.get_current_err_offset(), + ) + + def get_out_file_path(self) -> str: + """Get the out log file path""" + return self._out_filepath if self._out_filepath is not None else "" + + def get_err_file_path(self) -> str: + """Get the err log file path""" + return self._err_filepath if self._err_filepath is not None else "" + + def get_current_out_offset(self) -> int: + """Get the current offset of the out file if seekable, else 0""" + if self._out_filepath is not None: + return os.path.getsize(self._out_filepath) + return 0 + + def get_current_err_offset(self) -> int: + """Get the current offset of the err file if seekable, else 0""" + if self._err_filepath is not None: + return os.path.getsize(self._err_filepath) + return 0 + + def get_serialization_context(self): + """Get the SerializationContext of the job that this worker is processing. + + Returns: + The serialization context of the given job. + """ + # This function needs to be protected by a lock, because it will be + # called by`register_class_for_serialization`, as well as the import + # thread, from different threads. Also, this function will recursively + # call itself, so we use RLock here. + job_id = self.current_job_id + context_map = self.serialization_context_map + with self.lock: + if job_id not in context_map: + # The job ID is nil before initializing Ray. + if JobID.nil() in context_map: + # Transfer the serializer context used before initializing Ray. + context_map[job_id] = context_map.pop(JobID.nil()) + else: + context_map[job_id] = serialization.SerializationContext(self) + return context_map[job_id] + + def check_connected(self): + """Check if the worker is connected. + + Raises: + Exception: An exception is raised if the worker is not connected. + """ + if not self.connected: + raise RaySystemError( + "Ray has not been started yet. You can start Ray with 'ray.init()'." + ) + + def set_mode(self, mode): + """Set the mode of the worker. + + The mode SCRIPT_MODE should be used if this Worker is a driver that is + being run as a Python script or interactively in a shell. It will print + information about task failures. + + The mode WORKER_MODE should be used if this Worker is not a driver. It + will not print information about tasks. + + The mode LOCAL_MODE should be used if this Worker is a driver and if + you want to run the driver in a manner equivalent to serial Python for + debugging purposes. It will not send remote function calls to the + scheduler and will instead execute them in a blocking fashion. + + Args: + mode: One of SCRIPT_MODE, WORKER_MODE, and LOCAL_MODE. + """ + self.mode = mode + + def set_load_code_from_local(self, load_code_from_local): + self._load_code_from_local = load_code_from_local + + def put_object( + self, + value: Any, + owner_address: Optional[str] = None, + _is_experimental_channel: bool = False, + _tensor_transport: str = "object_store", + ): + """Put value in the local object store. + + If the plasma store is full, the worker will automatically + retry up to DEFAULT_PUT_OBJECT_RETRIES times. Each retry + will delay for an exponentially doubling amount of time, + starting with DEFAULT_PUT_OBJECT_DELAY. After this, exception + will be raised. + + Args: + value: The value to put in the object store. + owner_address: The serialized address of object's owner. + _is_experimental_channel: An experimental flag for mutable + objects. If True, then the returned object will not have a + valid value. The object must be written to using the + ray.experimental.channel API before readers can read. + _tensor_transport: [Alpha] The tensor transport backend to use. Currently, this supports "object_store" and "nixl". + Returns: + ObjectRef: The object ref the object was put under. + + Raises: + ray.exceptions.ObjectStoreFullError: This is raised if the attempt + to store the object fails because the object store is full even + after multiple retries. + """ + # Make sure that the value is not an object ref. + if isinstance(value, ObjectRef): + raise TypeError( + "Calling 'put' on an ray.ObjectRef is not allowed. " + "If you really want to do this, you can wrap the " + "ray.ObjectRef in a list and call 'put' on it." + ) + tensors = None + tensor_transport: TensorTransportEnum = TensorTransportEnum.from_str( + _tensor_transport + ) + if tensor_transport not in [ + TensorTransportEnum.OBJECT_STORE, + TensorTransportEnum.NIXL, + ]: + raise ValueError( + "Currently, Ray Direct Transport only supports 'object_store' and 'nixl' for tensor transport in ray.put()." + ) + try: + if tensor_transport != TensorTransportEnum.OBJECT_STORE: + ( + serialized_value, + tensors, + ) = self.get_serialization_context().serialize_gpu_objects(value) + else: + serialized_value = self.get_serialization_context().serialize(value) + except TypeError as e: + sio = io.StringIO() + ray.util.inspect_serializability(value, print_file=sio) + msg = ( + "Could not serialize the put value " + f"{repr(value)}:\n" + f"{sio.getvalue()}" + ) + raise TypeError(msg) from e + + # If the object is mutable, then the raylet should never read the + # object. Instead, clients will keep the object pinned. + pin_object = not _is_experimental_channel + + # This *must* be the first place that we construct this python + # ObjectRef because an entry with 0 local references is created when + # the object is Put() in the core worker, expecting that this python + # reference will be created. If another reference is created and + # removed before this one, it will corrupt the state in the + # reference counter. + ret = self.core_worker.put_object( + serialized_value, + pin_object=pin_object, + owner_address=owner_address, + inline_small_object=True, + _is_experimental_channel=_is_experimental_channel, + tensor_transport_val=tensor_transport.value, + ) + if tensors: + self.gpu_object_manager.put_object(ret, tensor_transport, tensors) + return ret + + def raise_errors(self, serialized_objects, object_refs): + out = self.deserialize_objects(serialized_objects, object_refs) + if "RAY_IGNORE_UNHANDLED_ERRORS" in os.environ: + return + for e in out: + _unhandled_error_handler(e) + + def deserialize_objects( + self, + serialized_objects, + object_refs, + tensor_transport_hint: Optional[TensorTransportEnum] = None, + ): + gpu_objects: Dict[str, List["torch.Tensor"]] = {} + for obj_ref, (_, _, tensor_transport) in zip(object_refs, serialized_objects): + # TODO: Here tensor_transport_hint is set by the user in ray.get(), tensor_transport is set + # in serialize_objects by ray.method(tensor_transport="xxx"), and obj_ref.tensor_transport() + # is set by ray.put(). We may clean up this logic in the future. + if ( + tensor_transport is None + or tensor_transport == TensorTransportEnum.OBJECT_STORE + ) and ( + obj_ref is None + or obj_ref.tensor_transport() == TensorTransportEnum.OBJECT_STORE.value + ): + # The object is not a gpu object, so we cannot use other external transport to + # fetch it. + continue + + # If the object is a gpu object, we can choose to use the object store or other external + # transport to fetch it. The `tensor_transport_hint` has the highest priority, then the + # tensor_transport in obj_ref.tensor_transport(), then the tensor_transport in serialize_objects, + # then the default value `OBJECT_STORE`. + chosen_tensor_transport = ( + tensor_transport_hint + or ( + TensorTransportEnum(obj_ref.tensor_transport()) if obj_ref else None + ) + or tensor_transport + or TensorTransportEnum.OBJECT_STORE + ) + + object_id = obj_ref.hex() + if object_id not in gpu_objects: + # If using a non-object store transport, then tensors will be sent + # out-of-band. Get them before deserializing the object store data. + gpu_objects[object_id] = self.gpu_object_manager.get_gpu_object( + object_id, tensor_transport=chosen_tensor_transport + ) + + # Function actor manager or the import thread may call pickle.loads + # at the same time which can lead to failed imports + # TODO: We may be better off locking on all imports or injecting a lock + # into pickle.loads (https://github.com/ray-project/ray/issues/16304) + with self.function_actor_manager.lock: + context = self.get_serialization_context() + return context.deserialize_objects( + serialized_objects, object_refs, gpu_objects + ) + + def get_objects( + self, + object_refs: list, + timeout: Optional[float] = None, + return_exceptions: bool = False, + skip_deserialization: bool = False, + _tensor_transport: Optional[str] = None, + ) -> Tuple[List[serialization.SerializedRayObject], bytes]: + """Get the values in the object store associated with the IDs. + + Return the values from the local object store for object_refs. This + will block until all the values for object_refs have been written to + the local object store. + + Args: + object_refs: A list of the object refs + whose values should be retrieved. + timeout: The maximum amount of time in + seconds to wait before returning. + return_exceptions: If any of the objects deserialize to an + Exception object, whether to return them as values in the + returned list. If False, then the first found exception will be + raised. + skip_deserialization: If true, only the buffer will be released and + the object associated with the buffer will not be deserialized. + _tensor_transport: [Alpha] The tensor transport to use to fetch `torch.Tensors` found in the Ray Direct Transport object. Currently, this supports "object_store" and "nixl". + Returns: + list: List of deserialized objects or None if skip_deserialization is True. + bytes: UUID of the debugger breakpoint we should drop + into or b"" if there is no breakpoint. + """ + # Make sure that the values are object refs. + for object_ref in object_refs: + if not isinstance(object_ref, ObjectRef): + raise TypeError( + f"Attempting to call `get` on the value {object_ref}, " + "which is not an ray.ObjectRef." + ) + tensor_transport: TensorTransportEnum = ( + TensorTransportEnum.from_str(_tensor_transport) + if _tensor_transport is not None + else None + ) + assert tensor_transport in [ + TensorTransportEnum.OBJECT_STORE, + TensorTransportEnum.NIXL, + None, + ], "Currently, RDT only supports 'object_store' and 'nixl' for tensor transport in ray.get()." + timeout_ms = ( + int(timeout * 1000) if timeout is not None and timeout != -1 else -1 + ) + serialized_objects: List[ + serialization.SerializedRayObject + ] = self.core_worker.get_objects( + object_refs, + timeout_ms, + ) + + debugger_breakpoint = b"" + for data, metadata, _ in serialized_objects: + if metadata: + metadata_fields = metadata.split(b",") + if len(metadata_fields) >= 2 and metadata_fields[1].startswith( + ray_constants.OBJECT_METADATA_DEBUG_PREFIX + ): + debugger_breakpoint = metadata_fields[1][ + len(ray_constants.OBJECT_METADATA_DEBUG_PREFIX) : + ] + if skip_deserialization: + return None, debugger_breakpoint + + values = self.deserialize_objects( + serialized_objects, object_refs, tensor_transport_hint=tensor_transport + ) + if not return_exceptions: + # Raise exceptions instead of returning them to the user. + for i, value in enumerate(values): + if isinstance(value, RayError): + if isinstance( + value, ray.exceptions.ObjectLostError + ) and not isinstance(value, ray.exceptions.OwnerDiedError): + global_worker.core_worker.log_plasma_usage() + if isinstance(value, RayTaskError): + raise value.as_instanceof_cause() + else: + raise value + + return values, debugger_breakpoint + + def main_loop(self): + """The main loop a worker runs to receive and execute tasks.""" + + def sigterm_handler(signum, frame): + raise_sys_exit_with_custom_error_message( + "The process receives a SIGTERM.", exit_code=1 + ) + # Note: shutdown() function is called from atexit handler. + + ray._private.utils.set_sigterm_handler(sigterm_handler) + self.core_worker.run_task_loop() + sys.exit(0) + + def print_logs(self): + """Prints log messages from workers on all nodes in the same job.""" + subscriber = self.gcs_log_subscriber + subscriber.subscribe() + exception_type = ray.exceptions.RpcError + localhost = services.get_node_ip_address() + try: + # Number of messages received from the last polling. When the batch + # size exceeds 100 and keeps increasing, the worker and the user + # probably will not be able to consume the log messages as rapidly + # as they are coming in. + # This is meaningful only for GCS subscriber. + last_polling_batch_size = 0 + job_id_hex = self.current_job_id.hex() + while True: + # Exit if we received a signal that we should stop. + if self.threads_stopped.is_set(): + return + + data = subscriber.poll() + # GCS subscriber only returns None on unavailability. + if data is None: + last_polling_batch_size = 0 + continue + + if ( + self._filter_logs_by_job + and data["job"] + and data["job"] != job_id_hex + ): + last_polling_batch_size = 0 + continue + + data["localhost"] = localhost + global_worker_stdstream_dispatcher.emit(data) + + lagging = 100 <= last_polling_batch_size < subscriber.last_batch_size + if lagging: + logger.warning( + "The driver may not be able to keep up with the " + "stdout/stderr of the workers. To avoid forwarding " + "logs to the driver, use " + "'ray.init(log_to_driver=False)'." + ) + + last_polling_batch_size = subscriber.last_batch_size + + except (OSError, exception_type) as e: + logger.error(f"print_logs: {e}") + finally: + # Close the pubsub client to avoid leaking file descriptors. + subscriber.close() + + def get_accelerator_ids_for_accelerator_resource( + self, resource_name: str, resource_regex: str + ) -> Union[List[str], List[int]]: + """Get the accelerator IDs that are assigned to the given accelerator resource. + + Args: + resource_name: The name of the resource. + resource_regex: The regex of the resource. + + Returns: + (List[str]) The IDs that are assigned to the given resource pre-configured. + (List[int]) The IDs that are assigned to the given resource. + """ + resource_ids = self.core_worker.resource_ids() + assigned_ids = set() + # Handle both normal and placement group accelerator resources. + # Note: We should only get the accelerator ids from the placement + # group resource that does not contain the bundle index! + import re + + for resource, assignment in resource_ids.items(): + if resource == resource_name or re.match(resource_regex, resource): + for resource_id, _ in assignment: + assigned_ids.add(resource_id) + + # If the user had already set the environment variables + # (CUDA_VISIBLE_DEVICES, ONEAPI_DEVICE_SELECTOR, NEURON_RT_VISIBLE_CORES, + # TPU_VISIBLE_CHIPS, ..) then respect that in the sense that only IDs + # that appear in (CUDA_VISIBLE_DEVICES, ONEAPI_DEVICE_SELECTOR, + # HIP_VISIBLE_DEVICES, NEURON_RT_VISIBLE_CORES, TPU_VISIBLE_CHIPS, ..) + # should be returned. + if self.original_visible_accelerator_ids.get(resource_name, None) is not None: + original_ids = self.original_visible_accelerator_ids[resource_name] + assigned_ids = {str(original_ids[i]) for i in assigned_ids} + # Give all accelerator ids in local_mode. + if self.mode == LOCAL_MODE: + if resource_name == ray_constants.GPU: + max_accelerators = self.node.get_resource_and_label_spec().num_gpus + else: + max_accelerators = ( + self.node.get_resource_and_label_spec().resources.get( + resource_name, None + ) + ) + if max_accelerators: + assigned_ids = original_ids[:max_accelerators] + return list(assigned_ids) + + def shutdown_gpu_object_manager(self): + if self._gpu_object_manager: + self._gpu_object_manager.shutdown() + + +_connect_or_shutdown_lock = threading.RLock() + + +def with_connect_or_shutdown_lock(func: Callable) -> Callable: + @wraps(func) + def wrapper(*args, **kwargs): + with _connect_or_shutdown_lock: + return func(*args, **kwargs) + + return wrapper + + +@PublicAPI +@client_mode_hook +def get_gpu_ids() -> Union[List[int], List[str]]: + """Get the IDs of the GPUs that are available to the worker. + + This method should only be called inside of a task or actor, and not a driver. + + If the CUDA_VISIBLE_DEVICES environment variable was set when the worker + started up, then the IDs returned by this method will be a subset of the + IDs in CUDA_VISIBLE_DEVICES. If not, the IDs will fall in the range + [0, NUM_GPUS - 1], where NUM_GPUS is the number of GPUs that the node has. + + Returns: + A list of GPU IDs. + """ + worker = global_worker + worker.check_connected() + return worker.get_accelerator_ids_for_accelerator_resource( + ray_constants.GPU, f"^{ray_constants.GPU}_group_[0-9A-Za-z]+$" + ) + + +@Deprecated( + message="Use ray.get_runtime_context().get_assigned_resources() instead.", + warning=True, +) +def get_resource_ids(): + """Get the IDs of the resources that are available to the worker. + + Returns: + A dictionary mapping the name of a resource to a list of pairs, where + each pair consists of the ID of a resource and the fraction of that + resource reserved for this worker. + """ + worker = global_worker + worker.check_connected() + + if _mode() == LOCAL_MODE: + raise RuntimeError( + "ray._private.worker.get_resource_ids() does not work in local_mode." + ) + + return global_worker.core_worker.resource_ids() + + +@Deprecated(message="Use ray.init().address_info['webui_url'] instead.") +def get_dashboard_url(): + """Get the URL to access the Ray dashboard. + + Note that the URL does not specify which node the dashboard is on. + + Returns: + The URL of the dashboard as a string. + """ + if ray_constants.RAY_OVERRIDE_DASHBOARD_URL in os.environ: + return _remove_protocol_from_url( + os.environ.get(ray_constants.RAY_OVERRIDE_DASHBOARD_URL) + ) + else: + worker = global_worker + worker.check_connected() + return _global_node.webui_url + + +def _remove_protocol_from_url(url: Optional[str]) -> str: + """ + Helper function to remove protocol from URL if it exists. + """ + if not url: + return url + parsed_url = urllib.parse.urlparse(url) + if parsed_url.scheme: + # Construct URL without protocol + scheme = f"{parsed_url.scheme}://" + return parsed_url.geturl().replace(scheme, "", 1) + return url + + +class BaseContext(metaclass=ABCMeta): + """ + Base class for RayContext and ClientContext + """ + + dashboard_url: Optional[str] + python_version: str + ray_version: str + + @abstractmethod + def disconnect(self): + """ + If this context is for directly attaching to a cluster, disconnect + will call ray.shutdown(). Otherwise, if the context is for a ray + client connection, the client will be disconnected. + """ + pass + + @abstractmethod + def __enter__(self): + pass + + @abstractmethod + def __exit__(self): + pass + + def _context_table_template(self): + if self.dashboard_url: + dashboard_row = Template("context_dashrow.html.j2").render( + dashboard_url="http://" + self.dashboard_url + ) + else: + dashboard_row = None + + return Template("context_table.html.j2").render( + python_version=self.python_version, + ray_version=self.ray_version, + dashboard_row=dashboard_row, + ) + + def _repr_html_(self): + return Template("context.html.j2").render( + context_logo=Template("context_logo.html.j2").render(), + context_table=self._context_table_template(), + ) + + @repr_with_fallback(["ipywidgets", "8"]) + def _get_widget_bundle(self, **kwargs) -> Dict[str, Any]: + """Get the mimebundle for the widget representation of the context. + + Args: + **kwargs: Passed to the _repr_mimebundle_() function for the widget + + Returns: + Dictionary ("mimebundle") of the widget representation of the context. + """ + import ipywidgets + + disconnect_button = ipywidgets.Button( + description="Disconnect", + disabled=False, + button_style="", + tooltip="Disconnect from the Ray cluster", + layout=ipywidgets.Layout(margin="auto 0px 0px 0px"), + ) + + def disconnect_callback(button): + button.disabled = True + button.description = "Disconnecting..." + self.disconnect() + button.description = "Disconnected" + + disconnect_button.on_click(disconnect_callback) + left_content = ipywidgets.VBox( + [ + ipywidgets.HTML(Template("context_logo.html.j2").render()), + disconnect_button, + ], + layout=ipywidgets.Layout(), + ) + right_content = ipywidgets.HTML(self._context_table_template()) + widget = ipywidgets.HBox( + [left_content, right_content], layout=ipywidgets.Layout(width="100%") + ) + return widget._repr_mimebundle_(**kwargs) + + def _repr_mimebundle_(self, **kwargs): + bundle = self._get_widget_bundle(**kwargs) + + # Overwrite the widget html repr and default repr with those of the BaseContext + bundle.update({"text/html": self._repr_html_(), "text/plain": repr(self)}) + return bundle + + +@dataclass +class RayContext(BaseContext, Mapping): + """ + Context manager for attached drivers. + """ + + dashboard_url: Optional[str] + python_version: str + ray_version: str + ray_commit: str + + def __init__(self, address_info: Dict[str, Optional[str]]): + super().__init__() + self.dashboard_url = get_dashboard_url() + self.python_version = "{}.{}.{}".format(*sys.version_info[:3]) + self.ray_version = ray.__version__ + self.ray_commit = ray.__commit__ + self.address_info = address_info + + def __getitem__(self, key): + if log_once("ray_context_getitem"): + warnings.warn( + f'Accessing values through ctx["{key}"] is deprecated. ' + f'Use ctx.address_info["{key}"] instead.', + DeprecationWarning, + stacklevel=2, + ) + return self.address_info[key] + + def __len__(self): + if log_once("ray_context_len"): + warnings.warn("len(ctx) is deprecated. Use len(ctx.address_info) instead.") + return len(self.address_info) + + def __iter__(self): + if log_once("ray_context_len"): + warnings.warn( + "iter(ctx) is deprecated. Use iter(ctx.address_info) instead." + ) + return iter(self.address_info) + + def __enter__(self) -> "RayContext": + return self + + def __exit__(self, *exc): + ray.shutdown() + + def disconnect(self): + # Include disconnect() to stay consistent with ClientContext + ray.shutdown() + + +global_worker = Worker() +"""Worker: The global Worker object for this worker process. + +We use a global Worker object to ensure that there is a single worker object +per worker process. +""" + +_global_node = None +"""ray._private.node.Node: The global node object that is created by ray.init().""" + + +def _maybe_modify_runtime_env( + runtime_env: Optional[Dict[str, Any]], _skip_env_hook: bool +) -> Dict[str, Any]: + """ + If you set RAY_ENABLE_UV_RUN_RUNTIME_ENV, which is the default, and run the driver with `uv run`, + this function sets up a runtime environment that replicates the driver's environment to the + workers. Otherwise, if a runtime environment hook is present it will modify the runtime environment. + """ + + if ray_constants.RAY_ENABLE_UV_RUN_RUNTIME_ENV: + from ray._private.runtime_env.uv_runtime_env_hook import ( + _get_uv_run_cmdline, + hook, + ) + + cmdline = _get_uv_run_cmdline() + if cmdline: + # This means the current driver is running in `uv run`, in which case we want + # to propagate the uv environment to the workers. + return hook(runtime_env) + + if ray_constants.RAY_RUNTIME_ENV_HOOK in os.environ and not _skip_env_hook: + return load_class(os.environ[ray_constants.RAY_RUNTIME_ENV_HOOK])(runtime_env) + + return runtime_env + + +@PublicAPI +@client_mode_hook +def init( + address: Optional[str] = None, + *, + num_cpus: Optional[int] = None, + num_gpus: Optional[int] = None, + resources: Optional[Dict[str, float]] = None, + labels: Optional[Dict[str, str]] = None, + object_store_memory: Optional[int] = None, + local_mode: bool = False, + ignore_reinit_error: bool = False, + include_dashboard: Optional[bool] = None, + dashboard_host: str = ray_constants.DEFAULT_DASHBOARD_IP, + dashboard_port: Optional[int] = None, + job_config: "ray.job_config.JobConfig" = None, + configure_logging: bool = True, + logging_level: int = ray_constants.LOGGER_LEVEL, + logging_format: Optional[str] = None, + logging_config: Optional[LoggingConfig] = None, + log_to_driver: Optional[bool] = None, + namespace: Optional[str] = None, + runtime_env: Optional[Union[Dict[str, Any], "RuntimeEnv"]] = None, # noqa: F821 + enable_resource_isolation: bool = False, + system_reserved_cpu: Optional[float] = None, + system_reserved_memory: Optional[int] = None, + **kwargs, +) -> BaseContext: + """ + Connect to an existing Ray cluster or start one and connect to it. + + This method handles two cases; either a Ray cluster already exists and we + just attach this driver to it or we start all of the processes associated + with a Ray cluster and attach to the newly started cluster. + Note: This method overwrite sigterm handler of the driver process. + + In most cases, it is enough to just call this method with no arguments. + This will autodetect an existing Ray cluster or start a new Ray instance if + no existing cluster is found: + + .. testcode:: + + ray.init() + + To explicitly connect to an existing local cluster, use this as follows. A + ConnectionError will be thrown if no existing local cluster is found. + + .. testcode:: + :skipif: True + + ray.init(address="auto") + + To connect to an existing remote cluster, use this as follows (substituting + in the appropriate address). Note the addition of "ray://" at the beginning + of the address. This requires `ray[client]`. + + .. testcode:: + :skipif: True + + ray.init(address="ray://123.45.67.89:10001") + + More details for starting and connecting to a remote cluster can be found + here: https://docs.ray.io/en/master/cluster/getting-started.html + + You can also define an environment variable called `RAY_ADDRESS` in + the same format as the `address` parameter to connect to an existing + cluster with ray.init() or ray.init(address="auto"). + + Args: + address: The address of the Ray cluster to connect to. The provided + address is resolved as follows: + 1. If a concrete address (e.g., localhost:) is provided, try to + connect to it. Concrete addresses can be prefixed with "ray://" to + connect to a remote cluster. For example, passing in the address + "ray://123.45.67.89:50005" will connect to the cluster at the given + address. + 2. If no address is provided, try to find an existing Ray instance + to connect to. This is done by first checking the environment + variable `RAY_ADDRESS`. If this is not defined, check the address + of the latest cluster started (found in + /tmp/ray/ray_current_cluster) if available. If this is also empty, + then start a new local Ray instance. + 3. If the provided address is "auto", then follow the same process + as above. However, if there is no existing cluster found, this will + throw a ConnectionError instead of starting a new local Ray + instance. + 4. If the provided address is "local", start a new local Ray + instance, even if there is already an existing local Ray instance. + num_cpus: Number of CPUs the user wishes to assign to each + raylet. By default, this is set based on virtual cores. + num_gpus: Number of GPUs the user wishes to assign to each + raylet. By default, this is set based on detected GPUs. + resources: A dictionary mapping the names of custom resources to the + quantities for them available. + labels: [Experimental] The key-value labels of the node. + object_store_memory: The amount of memory (in bytes) to start the + object store with. + By default, this is 30% of available system memory capped by + the shm size and 200G but can be set higher. + local_mode: Deprecated: consider using the Ray Distributed Debugger instead. + ignore_reinit_error: If true, Ray suppresses errors from calling + ray.init() a second time. Ray won't be restarted. + include_dashboard: Boolean flag indicating whether or not to start the + Ray dashboard, which displays the status of the Ray + cluster. If this argument is None, then the UI will be started if + the relevant dependencies are present. + dashboard_host: The host to bind the dashboard server to. Can either be + localhost (127.0.0.1) or 0.0.0.0 (available from all interfaces). + By default, this is set to localhost to prevent access from + external machines. + dashboard_port(int, None): The port to bind the dashboard server to. + Defaults to 8265 and Ray will automatically find a free port if + 8265 is not available. + job_config (ray.job_config.JobConfig): The job configuration. + configure_logging: True (default) if configuration of logging is + allowed here. Otherwise, the user may want to configure it + separately. + logging_level: Logging level for the "ray" logger of the driver process, + defaults to logging.INFO. Ignored unless "configure_logging" is true. + logging_format: Logging format for the "ray" logger of the driver process, + defaults to a string containing a timestamp, filename, line number, and + message. See the source file ray_constants.py for details. Ignored unless + "configure_logging" is true. + logging_config: [Experimental] Logging configuration will be applied to the + root loggers for both the driver process and all worker processes belonging + to the current job. See :class:`~ray.LoggingConfig` for details. + log_to_driver: If true, the output from all of the worker + processes on all nodes will be directed to the driver. + namespace: A namespace is a logical grouping of jobs and named actors. + runtime_env: The runtime environment to use + for this job (see :ref:`runtime-environments` for details). + object_spilling_directory: The path to spill objects to. The same path will + be used as the object store fallback directory as well. + enable_resource_isolation: Enable resource isolation through cgroupv2 by reserving + memory and cpu resources for ray system processes. To use, only cgroupv2 (not cgroupv1) + must be enabled with read and write permissions for the raylet. Cgroup memory and + cpu controllers must also be enabled. + system_reserved_cpu: The number of cpu cores to reserve for ray system processes. + Cores can be fractional i.e. 1.5 means one and a half a cpu core. + By default, the value will be atleast 1 core, and at maximum 3 cores. The default value + is calculated using the formula min(3.0, max(1.0, 0.05 * num_cores_on_the_system)) + This option only works if enable_resource_isolation is True. + system_reserved_memory: The amount of memory (in bytes) to reserve for ray system processes. + By default, the value will be atleast 500MB, and at most 10GB. The default value is + calculated using the formula min(10GB, max(500MB, 0.10 * memory_available_on_the_system)) + This option only works if enable_resource_isolation is True. + _cgroup_path: The path for the cgroup the raylet should use to enforce resource isolation. + By default, the cgroup used for resource isolation will be /sys/fs/cgroup. + The process starting ray must have read/write permissions to this path. + Cgroup memory and cpu controllers be enabled for this cgroup. + This option only works if enable_resource_isolation is True. + _enable_object_reconstruction: If True, when an object stored in + the distributed plasma store is lost due to node failure, Ray will + attempt to reconstruct the object by re-executing the task that + created the object. Arguments to the task will be recursively + reconstructed. If False, then ray.ObjectLostError will be + thrown. + _plasma_directory: Override the plasma mmap file directory. + _node_ip_address: The IP address of the node that we are on. + _driver_object_store_memory: Deprecated. + _memory: Amount of reservable memory resource in bytes rounded + down to the nearest integer. + _redis_username: Prevents external clients without the username + from connecting to Redis if provided. + _redis_password: Prevents external clients without the password + from connecting to Redis if provided. + _temp_dir: If provided, specifies the root temporary + directory for the Ray process. Must be an absolute path. Defaults to an + OS-specific conventional location, e.g., "/tmp/ray". + _metrics_export_port: Port number Ray exposes system metrics + through a Prometheus endpoint. It is currently under active + development, and the API is subject to change. + _system_config: Configuration for overriding + RayConfig defaults. For testing purposes ONLY. + _tracing_startup_hook: If provided, turns on and sets up tracing + for Ray. Must be the name of a function that takes no arguments and + sets up a Tracer Provider, Remote Span Processors, and + (optional) additional instruments. See more at + docs.ray.io/tracing.html. It is currently under active development, + and the API is subject to change. + _node_name: User-provided node name or identifier. Defaults to + the node IP address. + + Returns: + If the provided address includes a protocol, for example by prepending + "ray://" to the address to get "ray://1.2.3.4:10001", then a + ClientContext is returned with information such as settings, server + versions for ray and python, and the dashboard_url. Otherwise, + a RayContext is returned with ray and python versions, and address + information about the started processes. + + Raises: + Exception: An exception is raised if an inappropriate combination of + arguments is passed in. + """ + if log_to_driver is None: + log_to_driver = ray_constants.RAY_LOG_TO_DRIVER + + # Configure the "ray" logger for the driver process. + if configure_logging: + setup_logger(logging_level, logging_format or ray_constants.LOGGER_FORMAT) + else: + logging.getLogger("ray").handlers.clear() + + # Configure the logging settings for the driver process. + if logging_config or ray_constants.RAY_LOGGING_CONFIG_ENCODING: + logging_config = logging_config or LoggingConfig( + encoding=ray_constants.RAY_LOGGING_CONFIG_ENCODING + ) + logging_config._apply() + + # Parse the hidden options + _cgroup_path: str = kwargs.pop("_cgroup_path", None) + + _enable_object_reconstruction: bool = kwargs.pop( + "_enable_object_reconstruction", False + ) + _plasma_directory: Optional[str] = kwargs.pop("_plasma_directory", None) + _object_spilling_directory: Optional[str] = kwargs.pop( + "object_spilling_directory", None + ) + _node_ip_address: str = kwargs.pop("_node_ip_address", None) + _driver_object_store_memory: Optional[int] = kwargs.pop( + "_driver_object_store_memory", None + ) + _memory: Optional[int] = kwargs.pop("_memory", None) + _redis_username: str = kwargs.pop( + "_redis_username", ray_constants.REDIS_DEFAULT_USERNAME + ) + _redis_password: str = kwargs.pop( + "_redis_password", ray_constants.REDIS_DEFAULT_PASSWORD + ) + _temp_dir: Optional[str] = kwargs.pop("_temp_dir", None) + _metrics_export_port: Optional[int] = kwargs.pop("_metrics_export_port", None) + _system_config: Optional[Dict[str, str]] = kwargs.pop("_system_config", None) + _tracing_startup_hook: Optional[Callable] = kwargs.pop( + "_tracing_startup_hook", None + ) + _node_name: str = kwargs.pop("_node_name", None) + # Fix for https://github.com/ray-project/ray/issues/26729 + _skip_env_hook: bool = kwargs.pop("_skip_env_hook", False) + + resource_isolation_config = ResourceIsolationConfig( + enable_resource_isolation=enable_resource_isolation, + cgroup_path=_cgroup_path, + system_reserved_cpu=system_reserved_cpu, + system_reserved_memory=system_reserved_memory, + ) + + # terminate any signal before connecting driver + def sigterm_handler(signum, frame): + sys.exit(signum) + + if threading.current_thread() is threading.main_thread(): + ray._private.utils.set_sigterm_handler(sigterm_handler) + else: + logger.warning( + "SIGTERM handler is not set because current thread " + "is not the main thread." + ) + + # If available, use RAY_ADDRESS to override if the address was left + # unspecified, or set to "auto" in the call to init + address_env_var = os.environ.get(ray_constants.RAY_ADDRESS_ENVIRONMENT_VARIABLE) + if address_env_var and (address is None or address == "auto"): + address = address_env_var + logger.info( + f"Using address {address_env_var} set in the environment " + f"variable {ray_constants.RAY_ADDRESS_ENVIRONMENT_VARIABLE}" + ) + + if address is not None and "://" in address: + # Address specified a protocol, use ray client + builder = ray.client(address, _deprecation_warn_enabled=False) + + # Forward any keyword arguments that were changed from their default + # values to the builder + init_sig = inspect.signature(init) + passed_kwargs = {} + for argument_name, param_obj in init_sig.parameters.items(): + if argument_name in {"kwargs", "address"}: + # kwargs and address are handled separately + continue + default_value = param_obj.default + passed_value = locals()[argument_name] + if passed_value != default_value: + # passed value is different than default, pass to the client + # builder + passed_kwargs[argument_name] = passed_value + passed_kwargs.update(kwargs) + builder._init_args(**passed_kwargs) + ctx = builder.connect() + from ray._common.usage import usage_lib + + if passed_kwargs.get("allow_multiple") is True: + with ctx: + usage_lib.put_pre_init_usage_stats() + else: + usage_lib.put_pre_init_usage_stats() + + usage_lib.record_library_usage("client") + return ctx + + if kwargs.get("allow_multiple"): + raise RuntimeError( + "`allow_multiple` argument is passed to `ray.init` when the " + "ray client is not used (" + f"https://docs.ray.io/en/{get_ray_doc_version()}/cluster" + "/running-applications/job-submission" + "/ray-client.html#connect-to-multiple-ray-clusters-experimental). " + "Do not pass the `allow_multiple` to `ray.init` to fix the issue." + ) + + if kwargs.get("storage"): + raise RuntimeError( + "Cluster-wide storage configuration has been removed. " + "The last Ray version supporting the `storage` argument is `ray==2.47`." + ) + + if kwargs: + # User passed in extra keyword arguments but isn't connecting through + # ray client. Raise an error, since most likely a typo in keyword + unknown = ", ".join(kwargs) + raise RuntimeError(f"Unknown keyword argument(s): {unknown}") + + # Try to increase the file descriptor limit, which is too low by + # default for Ray: https://github.com/ray-project/ray/issues/11239 + try: + import resource + + soft, hard = resource.getrlimit(resource.RLIMIT_NOFILE) + if soft < hard: + # https://github.com/ray-project/ray/issues/12059 + soft = max(soft, min(hard, 65536)) + logger.debug( + f"Automatically increasing RLIMIT_NOFILE to max value of {hard}" + ) + try: + resource.setrlimit(resource.RLIMIT_NOFILE, (soft, hard)) + except ValueError: + logger.debug("Failed to raise limit.") + soft, _ = resource.getrlimit(resource.RLIMIT_NOFILE) + if soft < 4096: + logger.warning( + "File descriptor limit {} is too low for production " + "servers and may result in connection errors. " + "At least 8192 is recommended. --- " + "Fix with 'ulimit -n 8192'".format(soft) + ) + except ImportError: + logger.debug("Could not import resource module (on Windows)") + pass + + if job_config is None: + job_config = ray.job_config.JobConfig() + + if RAY_JOB_CONFIG_JSON_ENV_VAR in os.environ: + injected_job_config_json = json.loads( + os.environ.get(RAY_JOB_CONFIG_JSON_ENV_VAR) + ) + injected_job_config: ray.job_config.JobConfig = ( + ray.job_config.JobConfig.from_json(injected_job_config_json) + ) + driver_runtime_env = runtime_env + runtime_env = _merge_runtime_env( + injected_job_config.runtime_env, + driver_runtime_env, + override=os.getenv("RAY_OVERRIDE_JOB_RUNTIME_ENV") == "1", + ) + if runtime_env is None: + # None means there was a conflict. + raise ValueError( + "Failed to merge the Job's runtime env " + f"{injected_job_config.runtime_env} with " + f"a ray.init's runtime env {driver_runtime_env} because " + "of a conflict. Specifying the same runtime_env fields " + "or the same environment variable keys is not allowed. " + "Use RAY_OVERRIDE_JOB_RUNTIME_ENV=1 to instruct Ray to " + "combine Job and Driver's runtime environment in the event of " + "a conflict." + ) + + runtime_env = _maybe_modify_runtime_env(runtime_env, _skip_env_hook) + + job_config.set_runtime_env(runtime_env) + # Similarly, we prefer metadata provided via job submission API + for key, value in injected_job_config.metadata.items(): + job_config.set_metadata(key, value) + + # RAY_JOB_CONFIG_JSON_ENV_VAR is only set at ray job manager level and has + # higher priority in case user also provided runtime_env for ray.init() + else: + runtime_env = _maybe_modify_runtime_env(runtime_env, _skip_env_hook) + + if runtime_env: + # Set runtime_env in job_config if passed in as part of ray.init() + job_config.set_runtime_env(runtime_env) + + # Pass the logging_config to job_config to configure loggers of all worker + # processes belonging to the job. + if logging_config is not None: + job_config.set_py_logging_config(logging_config) + + redis_address, gcs_address = None, None + bootstrap_address = services.canonicalize_bootstrap_address(address, _temp_dir) + if bootstrap_address is not None: + gcs_address = bootstrap_address + logger.info("Connecting to existing Ray cluster at address: %s...", gcs_address) + + if local_mode: + driver_mode = LOCAL_MODE + warnings.warn( + "`local_mode` is an experimental feature that is no " + "longer maintained and will be removed in the near future. " + "For debugging consider using the Ray distributed debugger.", + FutureWarning, + stacklevel=2, + ) + else: + driver_mode = SCRIPT_MODE + + global _global_node + + if global_worker.connected: + if ignore_reinit_error: + logger.info("Calling ray.init() again after it has already been called.") + node_id = global_worker.core_worker.get_current_node_id() + return RayContext(dict(_global_node.address_info, node_id=node_id.hex())) + else: + raise RuntimeError( + "Maybe you called ray.init twice by accident? " + "This error can be suppressed by passing in " + "'ignore_reinit_error=True' or by calling " + "'ray.shutdown()' prior to 'ray.init()'." + ) + + _system_config = _system_config or {} + if not isinstance(_system_config, dict): + raise TypeError("The _system_config must be a dict.") + + if bootstrap_address is None: + # In this case, we need to start a new cluster. + + # Setup and verify authentication for new cluster + ensure_token_if_auth_enabled(_system_config, create_token_if_missing=True) + + # Don't collect usage stats in ray.init() unless it's a nightly wheel. + from ray._common.usage import usage_lib + + if usage_lib.is_nightly_wheel(): + usage_lib.show_usage_stats_prompt(cli=False) + else: + usage_lib.set_usage_stats_enabled_via_env_var(False) + + # Use a random port by not specifying Redis port / GCS server port. + ray_params = ray._private.parameter.RayParams( + node_ip_address=_node_ip_address, + driver_mode=driver_mode, + redirect_output=None, + num_cpus=num_cpus, + num_gpus=num_gpus, + resources=resources, + labels=labels, + num_redis_shards=None, + redis_max_clients=None, + redis_username=_redis_username, + redis_password=_redis_password, + plasma_directory=_plasma_directory, + object_spilling_directory=_object_spilling_directory, + huge_pages=None, + include_dashboard=include_dashboard, + dashboard_host=dashboard_host, + dashboard_port=dashboard_port, + memory=_memory, + object_store_memory=object_store_memory, + plasma_store_socket_name=None, + temp_dir=_temp_dir, + _system_config=_system_config, + enable_object_reconstruction=_enable_object_reconstruction, + metrics_export_port=_metrics_export_port, + tracing_startup_hook=_tracing_startup_hook, + node_name=_node_name, + resource_isolation_config=resource_isolation_config, + ) + # Start the Ray processes. We set shutdown_at_exit=False because we + # shutdown the node in the ray.shutdown call that happens in the atexit + # handler. We still spawn a reaper process in case the atexit handler + # isn't called. + _global_node = ray._private.node.Node( + ray_params=ray_params, + head=True, + shutdown_at_exit=False, + spawn_reaper=True, + ray_init_cluster=True, + ) + else: + # In this case, we are connecting to an existing cluster. + if num_cpus is not None or num_gpus is not None: + raise ValueError( + "When connecting to an existing cluster, num_cpus " + "and num_gpus must not be provided." + ) + if resources is not None: + raise ValueError( + "When connecting to an existing cluster, " + "resources must not be provided." + ) + if labels is not None: + raise ValueError( + "When connecting to an existing cluster, " + "labels must not be provided." + ) + if object_store_memory is not None: + raise ValueError( + "When connecting to an existing cluster, " + "object_store_memory must not be provided." + ) + if _system_config is not None and len(_system_config) != 0: + raise ValueError( + "When connecting to an existing cluster, " + "_system_config must not be provided." + ) + if _enable_object_reconstruction: + raise ValueError( + "When connecting to an existing cluster, " + "_enable_object_reconstruction must not be provided." + ) + if _node_name is not None: + raise ValueError( + "_node_name cannot be configured when connecting to " + "an existing cluster." + ) + + # Setup and verify authentication for connecting to existing cluster + ensure_token_if_auth_enabled(_system_config, create_token_if_missing=False) + + # In this case, we only need to connect the node. + ray_params = ray._private.parameter.RayParams( + node_ip_address=_node_ip_address, + gcs_address=gcs_address, + redis_address=redis_address, + redis_username=_redis_username, + redis_password=_redis_password, + temp_dir=_temp_dir, + _system_config=_system_config, + enable_object_reconstruction=_enable_object_reconstruction, + metrics_export_port=_metrics_export_port, + ) + try: + _global_node = ray._private.node.Node( + ray_params, + head=False, + shutdown_at_exit=False, + spawn_reaper=False, + connect_only=True, + ) + except (ConnectionError, RuntimeError): + if gcs_address == ray._private.utils.read_ray_address(_temp_dir): + logger.info( + "Failed to connect to the default Ray cluster address at " + f"{gcs_address}. This is most likely due to a previous Ray " + "instance that has since crashed. To reset the default " + "address to connect to, run `ray stop` or restart Ray with " + "`ray start`." + ) + raise ConnectionError + + # Log a message to find the Ray address that we connected to and the + # dashboard URL. + if ray_constants.RAY_OVERRIDE_DASHBOARD_URL in os.environ: + dashboard_url = os.environ.get(ray_constants.RAY_OVERRIDE_DASHBOARD_URL) + else: + dashboard_url = _global_node.webui_url + # Add http protocol to dashboard URL if it doesn't + # already contain a protocol. + if dashboard_url and not urlparse(dashboard_url).scheme: + dashboard_url = "http://" + dashboard_url + + # We logged the address before attempting the connection, so we don't need + # to log it again. + info_str = "Connected to Ray cluster." + if gcs_address is None: + info_str = "Started a local Ray instance." + if dashboard_url: + logger.info( + info_str + " View the dashboard at %s%s%s %s%s", + colorama.Style.BRIGHT, + colorama.Fore.GREEN, + dashboard_url, + colorama.Fore.RESET, + colorama.Style.NORMAL, + ) + else: + logger.info(info_str) + + connect( + _global_node, + _global_node.session_name, + mode=driver_mode, + log_to_driver=log_to_driver, + worker=global_worker, + driver_object_store_memory=_driver_object_store_memory, + job_id=None, + namespace=namespace, + job_config=job_config, + entrypoint=ray._private.utils.get_entrypoint_name(), + ) + if job_config and job_config.code_search_path: + global_worker.set_load_code_from_local(True) + else: + # Because `ray.shutdown()` doesn't reset this flag, for multiple + # sessions in one process, the 2nd `ray.init()` will reuse the + # flag of last session. For example: + # ray.init(load_code_from_local=True) + # ray.shutdown() + # ray.init() + # # Here the flag `load_code_from_local` is still True if we + # # doesn't have this `else` branch. + # ray.shutdown() + global_worker.set_load_code_from_local(False) + + for hook in _post_init_hooks: + hook() + + # Check and show accelerator override warning during driver initialization + from ray._private.ray_constants import env_bool + + override_on_zero = env_bool( + ray._private.accelerators.RAY_ACCEL_ENV_VAR_OVERRIDE_ON_ZERO_ENV_VAR, + True, + ) + if override_on_zero and log_once("ray_accel_env_var_override_on_zero"): + warnings.warn( + "Tip: In future versions of Ray, Ray will no longer override accelerator " + "visible devices env var if num_gpus=0 or num_gpus=None (default). To enable " + "this behavior and turn off this error message, set RAY_ACCEL_ENV_VAR_OVERRIDE_ON_ZERO=0", + FutureWarning, + ) + + node_id = global_worker.core_worker.get_current_node_id() + global_node_address_info = _global_node.address_info.copy() + global_node_address_info["webui_url"] = _remove_protocol_from_url(dashboard_url) + return RayContext(dict(global_node_address_info, node_id=node_id.hex())) + + +# Functions to run as callback after a successful ray init. +_post_init_hooks = [] + + +@PublicAPI +@client_mode_hook +@with_connect_or_shutdown_lock +def shutdown(_exiting_interpreter: bool = False): + """Disconnect the worker, and terminate processes started by ray.init(). + + This will automatically run at the end when a Python process that uses Ray + exits. It is ok to run this twice in a row. The primary use case for this + function is to cleanup state between tests. + + Note that this will clear any remote function definitions, actor + definitions, and existing actors, so if you wish to use any previously + defined remote functions or actors after calling ray.shutdown(), then you + need to redefine them. If they were defined in an imported module, then you + will need to reload the module. + + Args: + _exiting_interpreter: True if this is called by the atexit hook + and false otherwise. If we are exiting the interpreter, we will + wait a little while to print any extra error messages. + """ + # Make sure to clean up compiled dag node if exists. + from ray.dag.compiled_dag_node import _shutdown_all_compiled_dags + + _shutdown_all_compiled_dags() + global_worker.shutdown_gpu_object_manager() + + if _exiting_interpreter and global_worker.mode == SCRIPT_MODE: + # This is a duration to sleep before shutting down everything in order + # to make sure that log messages finish printing. + time.sleep(0.5) + disconnect(_exiting_interpreter) + + # disconnect internal kv + if hasattr(global_worker, "gcs_client"): + del global_worker.gcs_client + _internal_kv_reset() + + # We need to destruct the core worker here because after this function, + # we will tear down any processes spawned by ray.init() and the background + # IO thread in the core worker doesn't currently handle that gracefully. + if hasattr(global_worker, "core_worker"): + if global_worker.mode == SCRIPT_MODE or global_worker.mode == LOCAL_MODE: + global_worker.core_worker.shutdown_driver() + del global_worker.core_worker + # We need to reset function actor manager to clear the context + global_worker.function_actor_manager = FunctionActorManager(global_worker) + # Disconnect global state from GCS. + ray._private.state.state.disconnect() + + # Shut down the Ray processes. + global _global_node + if _global_node is not None: + if _global_node.is_head(): + _global_node.destroy_external_storage() + _global_node.kill_all_processes(check_alive=False, allow_graceful=True) + _global_node = None + + # TODO(rkn): Instead of manually resetting some of the worker fields, we + # should simply set "global_worker" to equal "None" or something like that. + global_worker.set_mode(None) + global_worker.set_cached_job_id(None) + + +atexit.register(shutdown, True) + +# Define a custom excepthook so that if the driver exits with an exception, we +# can push that exception to Redis. +normal_excepthook = sys.excepthook + + +def custom_excepthook(type, value, tb): + import ray.core.generated.common_pb2 as common_pb2 + + # If this is a driver, push the exception to GCS worker table. + if global_worker.mode == SCRIPT_MODE and hasattr(global_worker, "worker_id"): + error_message = "".join(traceback.format_tb(tb)) + worker_id = global_worker.worker_id + worker_type = common_pb2.DRIVER + worker_info = {"exception": error_message} + + ray._private.state.state._connect_and_get_accessor() + ray._private.state.state.add_worker(worker_id, worker_type, worker_info) + # Call the normal excepthook. + normal_excepthook(type, value, tb) + + +sys.excepthook = custom_excepthook + + +def print_to_stdstream(data, ignore_prefix: bool): + should_dedup = data.get("pid") not in ["autoscaler"] + + if data["is_err"]: + if should_dedup: + batches = stderr_deduplicator.deduplicate(data) + else: + batches = [data] + sink = sys.stderr + else: + if should_dedup: + batches = stdout_deduplicator.deduplicate(data) + else: + batches = [data] + sink = sys.stdout + + for batch in batches: + print_worker_logs(batch, sink, ignore_prefix) + + +# Start time of this process, used for relative time logs. +t0 = time.time() +autoscaler_log_fyi_printed = False + + +def filter_autoscaler_events(lines: List[str]) -> Iterator[str]: + """Given raw log lines from the monitor, return only autoscaler events. + + For Autoscaler V1: + Autoscaler events are denoted by the ":event_summary:" magic token. + For Autoscaler V2: + Autoscaler events are published from log_monitor.py which read + them from the `event_AUTOSCALER.log`. + """ + + if not ray_constants.AUTOSCALER_EVENTS: + return + + AUTOSCALER_LOG_FYI = ( + "Tip: use `ray status` to view detailed " + "cluster status. To disable these " + "messages, set RAY_SCHEDULER_EVENTS=0." + ) + + def autoscaler_log_fyi_needed() -> bool: + global autoscaler_log_fyi_printed + if not autoscaler_log_fyi_printed: + autoscaler_log_fyi_printed = True + return True + return False + + from ray.autoscaler.v2.utils import is_autoscaler_v2 + + if is_autoscaler_v2(): + from ray._private.event.event_logger import filter_event_by_level, parse_event + + for event_line in lines: + if autoscaler_log_fyi_needed(): + yield AUTOSCALER_LOG_FYI + + event = parse_event(event_line) + if not event or not event.message: + continue + + if filter_event_by_level( + event, ray_constants.RAY_LOG_TO_DRIVER_EVENT_LEVEL + ): + continue + + yield event.message + else: + # Print out autoscaler events only, ignoring other messages. + for line in lines: + if ray_constants.LOG_PREFIX_EVENT_SUMMARY in line: + if autoscaler_log_fyi_needed(): + yield AUTOSCALER_LOG_FYI + # The event text immediately follows the ":event_summary:" + # magic token. + yield line.split(ray_constants.LOG_PREFIX_EVENT_SUMMARY)[1] + + +def time_string() -> str: + """Return the relative time from the start of this job. + + For example, 15m30s. + """ + delta = time.time() - t0 + hours = 0 + minutes = 0 + while delta > 3600: + hours += 1 + delta -= 3600 + while delta > 60: + minutes += 1 + delta -= 60 + output = "" + if hours: + output += f"{hours}h" + if minutes: + output += f"{minutes}m" + output += f"{int(delta)}s" + return output + + +# When we enter a breakpoint, worker logs are automatically disabled via this. +_worker_logs_enabled = True + + +def print_worker_logs( + data: Dict[str, str], print_file: Any, ignore_prefix: bool = False +): + if not _worker_logs_enabled: + return + + def prefix_for(data: Dict[str, str]) -> str: + """The PID prefix for this log line.""" + if data.get("pid") in ["autoscaler", "raylet"]: + return "" + else: + res = "pid=" + if data.get("actor_name"): + res = f"{data['actor_name']} {res}" + elif data.get("task_name"): + res = f"{data['task_name']} {res}" + return res + + def message_for(data: Dict[str, str], line: str) -> str: + """The printed message of this log line.""" + if ray_constants.LOG_PREFIX_INFO_MESSAGE in line: + return line.split(ray_constants.LOG_PREFIX_INFO_MESSAGE)[1] + return line + + def color_for(data: Dict[str, str], line: str) -> str: + """The color for this log line.""" + if ( + data.get("pid") == "raylet" + and ray_constants.LOG_PREFIX_INFO_MESSAGE not in line + ): + return colorama.Fore.YELLOW + elif data.get("pid") == "autoscaler": + if "Error:" in line or "Warning:" in line: + return colorama.Fore.YELLOW + else: + return colorama.Fore.CYAN + elif os.getenv("RAY_COLOR_PREFIX") == "1": + colors = [ + # colorama.Fore.BLUE, # Too dark + colorama.Fore.MAGENTA, + colorama.Fore.CYAN, + colorama.Fore.GREEN, + # colorama.Fore.WHITE, # Too light + # colorama.Fore.RED, + colorama.Fore.LIGHTBLACK_EX, + colorama.Fore.LIGHTBLUE_EX, + # colorama.Fore.LIGHTCYAN_EX, # Too light + # colorama.Fore.LIGHTGREEN_EX, # Too light + colorama.Fore.LIGHTMAGENTA_EX, + # colorama.Fore.LIGHTWHITE_EX, # Too light + # colorama.Fore.LIGHTYELLOW_EX, # Too light + ] + pid = data.get("pid", 0) + try: + i = int(pid) + except ValueError: + i = 0 + return colors[i % len(colors)] + else: + return colorama.Fore.CYAN + + if data.get("pid") == "autoscaler": + pid = "autoscaler +{}".format(time_string()) + lines = filter_autoscaler_events(data.get("lines", [])) + else: + pid = data.get("pid") + lines = data.get("lines", []) + + ip = data.get("ip") + ip_prefix = "" if ip == data.get("localhost") else f", ip={ip}" + for line in lines: + if RAY_TQDM_MAGIC in line: + process_tqdm(line) + else: + hide_tqdm() + # If RAY_COLOR_PREFIX=0, do not wrap with any color codes + if os.getenv("RAY_COLOR_PREFIX") == "0": + color_pre = "" + color_post = "" + else: + color_pre = color_for(data, line) + color_post = colorama.Style.RESET_ALL + + if ignore_prefix: + print( + f"{message_for(data, line)}", + file=print_file, + ) + else: + print( + f"{color_pre}({prefix_for(data)}{pid}{ip_prefix}){color_post} " + f"{message_for(data, line)}", + file=print_file, + ) + + # Restore once at end of batch to avoid excess hiding/unhiding of tqdm. + restore_tqdm() + + +def process_tqdm(line): + """Experimental distributed tqdm: see ray.experimental.tqdm_ray.""" + try: + data = json.loads(line) + tqdm_ray.instance().process_state_update(data) + except Exception: + if log_once("tqdm_corruption"): + logger.warning( + f"[tqdm_ray] Failed to decode {line}, this may be due to " + "logging too fast. This warning will not be printed again." + ) + + +def hide_tqdm(): + """Hide distributed tqdm bars temporarily to avoid conflicts with other logs.""" + tqdm_ray.instance().hide_bars() + + +def restore_tqdm(): + """Undo hide_tqdm().""" + tqdm_ray.instance().unhide_bars() + + +def listen_error_messages(worker, threads_stopped): + """Listen to error messages in the background on the driver. + + This runs in a separate thread on the driver and pushes (error, time) + tuples to be published. + + Args: + worker: The worker class that this thread belongs to. + threads_stopped (threading.Event): A threading event used to signal to + the thread that it should exit. + """ + + # TODO: we should just subscribe to the errors for this specific job. + worker.gcs_error_subscriber.subscribe() + + try: + if _internal_kv_initialized(): + # Get any autoscaler errors that occurred before the call to + # subscribe. + error_message = _internal_kv_get(ray_constants.DEBUG_AUTOSCALING_ERROR) + if error_message is not None: + logger.warning(error_message.decode()) + while True: + # Exit if received a signal that the thread should stop. + if threads_stopped.is_set(): + return + + _, error_data = worker.gcs_error_subscriber.poll() + if error_data is None: + continue + if error_data["job_id"] is not None and error_data["job_id"] not in [ + worker.current_job_id.binary(), + JobID.nil().binary(), + ]: + continue + + error_message = error_data["error_message"] + print_to_stdstream( + { + "lines": [error_message], + "pid": "raylet", + "is_err": False, + }, + ignore_prefix=False, + ) + except (OSError, ConnectionError) as e: + logger.error(f"listen_error_messages: {e}") + + +@PublicAPI +@client_mode_hook +def is_initialized() -> bool: + """Check if ray.init has been called yet. + + Returns: + True if ray.init has already been called and false otherwise. + """ + return ray._private.worker.global_worker.connected + + +@with_connect_or_shutdown_lock +def connect( + node, + session_name: str, + mode=WORKER_MODE, + log_to_driver: bool = False, + worker=global_worker, + driver_object_store_memory: Optional[int] = None, + job_id=None, + namespace: Optional[str] = None, + job_config=None, + runtime_env_hash: int = 0, + startup_token: int = 0, + ray_debugger_external: bool = False, + entrypoint: str = "", + worker_launch_time_ms: int = -1, + worker_launched_time_ms: int = -1, + debug_source: str = "", +): + """Connect this worker to the raylet, to Plasma, and to GCS. + + Args: + node (ray._private.node.Node): The node to connect. + session_name: The current Ray session name. + mode: The mode of the worker. One of SCRIPT_MODE, WORKER_MODE, and LOCAL_MODE. + log_to_driver: If true, then output from all of the worker + processes on all nodes will be directed to the driver. + worker: The ray.Worker instance. + driver_object_store_memory: Deprecated. + job_id: The ID of job. If it's None, then we will generate one. + namespace: Namespace to use. + job_config (ray.job_config.JobConfig): The job configuration. + runtime_env_hash: The hash of the runtime env for this worker. + startup_token: The startup token of the process assigned to + it during startup as a command line argument. + ray_debugger_external: If True, make the debugger external to the + node this worker is running on. + entrypoint: The name of the entrypoint script. Ignored if the + mode != SCRIPT_MODE + worker_launch_time_ms: The time when the worker process for this worker + is launched. If the worker is not launched by raylet (e.g., + driver), this must be -1 (default value). + worker_launched_time_ms: The time when the worker process for this worker + finshes launching. If the worker is not launched by raylet (e.g., + driver), this must be -1 (default value). + debug_source: Source information for `CoreWorker`, used for debugging and informational purpose, rather than functional purpose. + """ + # Do some basic checking to make sure we didn't call ray.init twice. + error_message = "Perhaps you called ray.init twice by accident?" + assert not worker.connected, error_message + + # Enable nice stack traces on SIGSEGV etc. + try: + if not faulthandler.is_enabled(): + faulthandler.enable(all_threads=False) + except io.UnsupportedOperation: + pass # ignore + + worker.gcs_client = node.get_gcs_client() + assert worker.gcs_client is not None + _initialize_internal_kv(worker.gcs_client) + ray._private.state.state._initialize_global_state( + ray._raylet.GcsClientOptions.create( + node.gcs_address, + node.cluster_id.hex(), + allow_cluster_id_nil=False, + fetch_cluster_id_if_nil=False, + ) + ) + # Initialize some fields. + if mode in (WORKER_MODE, RESTORE_WORKER_MODE, SPILL_WORKER_MODE): + # We should not specify the job_id if it's `WORKER_MODE`. + assert job_id is None + job_id = JobID.nil() + else: + # This is the code path of driver mode. + if job_id is None: + job_id = ray._private.state.next_job_id() + + if mode is not SCRIPT_MODE and mode is not LOCAL_MODE: + process_name = ray_constants.WORKER_PROCESS_TYPE_IDLE_WORKER + if mode is SPILL_WORKER_MODE: + process_name = ray_constants.WORKER_PROCESS_TYPE_SPILL_WORKER_IDLE + elif mode is RESTORE_WORKER_MODE: + process_name = ray_constants.WORKER_PROCESS_TYPE_RESTORE_WORKER_IDLE + ray._raylet.setproctitle(process_name) + + if not isinstance(job_id, JobID): + raise TypeError("The type of given job id must be JobID.") + + # All workers start out as non-actors. A worker can be turned into an actor + # after it is created. + worker.node = node + worker.set_mode(mode) + + # For driver's check that the version information matches the version + # information that the Ray cluster was started with. + try: + node.check_version_info() + except Exception as e: + if mode == SCRIPT_MODE: + raise e + elif mode == WORKER_MODE: + traceback_str = traceback.format_exc() + ray._private.utils.publish_error_to_driver( + ray_constants.VERSION_MISMATCH_PUSH_ERROR, + traceback_str, + gcs_client=worker.gcs_client, + ) + + driver_name = "" + interactive_mode = False + if mode == SCRIPT_MODE: + import __main__ as main + + if hasattr(main, "__file__"): + driver_name = main.__file__ + else: + interactive_mode = True + driver_name = "INTERACTIVE MODE" + elif not LOCAL_MODE: + raise ValueError("Invalid worker mode. Expected DRIVER, WORKER or LOCAL.") + + gcs_options = ray._raylet.GcsClientOptions.create( + node.gcs_address, + node.cluster_id.hex(), + allow_cluster_id_nil=False, + fetch_cluster_id_if_nil=False, + ) + if job_config is None: + job_config = ray.job_config.JobConfig() + + if namespace is not None: + ray._private.utils.validate_namespace(namespace) + + # The namespace field of job config may have already been set in code + # paths such as the client. + job_config.set_ray_namespace(namespace) + + # Make sure breakpoint() in the user's code will + # invoke the Ray debugger if we are in a worker or actor process + # (but not on the driver). + if mode == WORKER_MODE: + os.environ["PYTHONBREAKPOINT"] = "ray.util.rpdb.set_trace" + else: + # Add hook to suppress worker logs during breakpoint. + os.environ["PYTHONBREAKPOINT"] = "ray.util.rpdb._driver_set_trace" + + worker.ray_debugger_external = ray_debugger_external + + # If it's a driver and it's not coming from ray client, we'll prepare the + # environment here. If it's ray client, the environment will be prepared + # at the server side. + if mode == SCRIPT_MODE and not job_config._client_job and job_config.runtime_env: + scratch_dir: str = worker.node.get_runtime_env_dir_path() + runtime_env = job_config.runtime_env or {} + runtime_env = upload_py_modules_if_needed( + runtime_env, scratch_dir, logger=logger + ) + runtime_env = upload_working_dir_if_needed( + runtime_env, scratch_dir, logger=logger + ) + runtime_env = upload_worker_process_setup_hook_if_needed( + runtime_env, + worker, + ) + # Remove excludes, it isn't relevant after the upload step. + runtime_env.pop("excludes", None) + job_config.set_runtime_env(runtime_env, validate=True) + + if mode == SCRIPT_MODE: + # Add the directory containing the script that is running to the Python + # paths of the workers. Also add the current directory. Note that this + # assumes that the directory structures on the machines in the clusters + # are the same. + # When using an interactive shell, there is no script directory. + # We also want to skip adding script directory when running from dashboard. + code_paths = [] + if not interactive_mode and not ( + namespace and namespace == ray._raylet.RAY_INTERNAL_DASHBOARD_NAMESPACE + ): + script_directory = os.path.dirname(os.path.realpath(sys.argv[0])) + # If driver's sys.path doesn't include the script directory + # (e.g driver is started via `python -m`, + # see https://peps.python.org/pep-0338/), + # then we shouldn't add it to the workers. + if script_directory in sys.path: + code_paths.append(script_directory) + # In client mode, if we use runtime envs with "working_dir", then + # it'll be handled automatically. Otherwise, add the current dir. + if not job_config._client_job and not job_config._runtime_env_has_working_dir(): + current_directory = os.path.abspath(os.path.curdir) + code_paths.append(current_directory) + if len(code_paths) != 0: + job_config._py_driver_sys_path.extend(code_paths) + + serialized_job_config = job_config._serialize() + if not node.should_redirect_logs(): + # Logging to stderr, so give core worker empty logs directory. + logs_dir = "" + else: + logs_dir = node.get_logs_dir_path() + + worker.core_worker = ray._raylet.CoreWorker( + mode, + node.plasma_store_socket_name, + node.raylet_socket_name, + job_id, + gcs_options, + logs_dir, + node.node_ip_address, + node.node_manager_port, + (mode == LOCAL_MODE), + driver_name, + serialized_job_config, + node.metrics_agent_port, + runtime_env_hash, + startup_token, + session_name, + node.cluster_id.hex(), + "" if mode != SCRIPT_MODE else entrypoint, + worker_launch_time_ms, + worker_launched_time_ms, + debug_source, + ) + + if mode == SCRIPT_MODE: + worker_id = worker.worker_id + worker.gcs_error_subscriber = ray._raylet.GcsErrorSubscriber( + worker_id=worker_id, address=worker.gcs_client.address + ) + worker.gcs_log_subscriber = ray._raylet.GcsLogSubscriber( + worker_id=worker_id, address=worker.gcs_client.address + ) + + if driver_object_store_memory is not None: + logger.warning( + "`driver_object_store_memory` is deprecated" + " and will be removed in the future." + ) + + # If this is a driver running in SCRIPT_MODE, start a thread to print error + # messages asynchronously in the background. Ideally the scheduler would + # push messages to the driver's worker service, but we ran into bugs when + # trying to properly shutdown the driver's worker service, so we are + # temporarily using this implementation which constantly queries the + # scheduler for new error messages. + if mode == SCRIPT_MODE: + worker.listener_thread = threading.Thread( + target=listen_error_messages, + name="ray_listen_error_messages", + args=(worker, worker.threads_stopped), + ) + worker.listener_thread.daemon = True + worker.listener_thread.start() + # If the job's logging config is set, don't add the prefix + # (task/actor's name and its PID) to the logs. + ignore_prefix = global_worker.job_logging_config is not None + + if log_to_driver: + global_worker_stdstream_dispatcher.add_handler( + "ray_print_logs", + functools.partial(print_to_stdstream, ignore_prefix=ignore_prefix), + ) + worker.logger_thread = threading.Thread( + target=worker.print_logs, name="ray_print_logs" + ) + worker.logger_thread.daemon = True + worker.logger_thread.start() + + # Setup tracing here + tracing_hook_val = worker.gcs_client.internal_kv_get( + b"tracing_startup_hook", ray_constants.KV_NAMESPACE_TRACING + ) + if tracing_hook_val is not None: + ray.util.tracing.tracing_helper._enable_tracing() + if not getattr(ray, "__traced__", False): + _setup_tracing = _import_from_string(tracing_hook_val.decode("utf-8")) + _setup_tracing() + ray.__traced__ = True + + # Mark the worker as connected. + worker.set_is_connected(True) + + +def disconnect(exiting_interpreter=False): + """Disconnect this worker from the raylet and object store.""" + # Reset the list of cached remote functions and actors so that if more + # remote functions or actors are defined and then connect is called again, + # the remote functions will be exported. This is mostly relevant for the + # tests. + worker = global_worker + if worker.connected: + # Shutdown all of the threads that we've started. TODO(rkn): This + # should be handled cleanly in the worker object's destructor and not + # in this disconnect method. + worker.threads_stopped.set() + if hasattr(worker, "gcs_error_subscriber"): + worker.gcs_error_subscriber.close() + if hasattr(worker, "gcs_log_subscriber"): + worker.gcs_log_subscriber.close() + if hasattr(worker, "listener_thread"): + worker.listener_thread.join() + if hasattr(worker, "logger_thread"): + worker.logger_thread.join() + worker.threads_stopped.clear() + + # Ignore the prefix if the logging config is set. + ignore_prefix = worker.job_logging_config is not None + for leftover in stdout_deduplicator.flush(): + print_worker_logs(leftover, sys.stdout, ignore_prefix) + for leftover in stderr_deduplicator.flush(): + print_worker_logs(leftover, sys.stderr, ignore_prefix) + global_worker_stdstream_dispatcher.remove_handler("ray_print_logs") + + worker.node = None # Disconnect the worker from the node. + worker.serialization_context_map.clear() + try: + ray_actor = ray.actor + except AttributeError: + ray_actor = None # This can occur during program termination + if ray_actor is not None: + ray_actor._ActorClassMethodMetadata.reset_cache() + + # Mark the worker as disconnected. + worker.set_is_connected(False) + + +@contextmanager +def _changeproctitle(title, next_title): + if _mode() is not LOCAL_MODE: + ray._raylet.setproctitle(title) + try: + yield + finally: + if _mode() is not LOCAL_MODE: + ray._raylet.setproctitle(next_title) + + +@DeveloperAPI +def show_in_dashboard(message: str, key: str = "", dtype: str = "text"): + """Display message in dashboard. + + Display message for the current task or actor in the dashboard. + For example, this can be used to display the status of a long-running + computation. + + Args: + message: Message to be displayed. + key: The key name for the message. Multiple message under + different keys will be displayed at the same time. Messages + under the same key will be overridden. + dtype: The type of message for rendering. One of the + following: text, html. + """ + worker = global_worker + worker.check_connected() + + acceptable_dtypes = {"text", "html"} + assert dtype in acceptable_dtypes, f"dtype accepts only: {acceptable_dtypes}" + + message_wrapped = {"message": message, "dtype": dtype} + message_encoded = json.dumps(message_wrapped).encode() + + worker.core_worker.set_webui_display(key.encode(), message_encoded) + + +# Global variable to make sure we only send out the warning once. +blocking_get_inside_async_warned = False + + +@overload +def get( + object_refs: "Sequence[ObjectRef[Any]]", *, timeout: Optional[float] = None +) -> List[Any]: + ... + + +@overload +def get( + object_refs: "Sequence[ObjectRef[R]]", *, timeout: Optional[float] = None +) -> List[R]: + ... + + +@overload +def get(object_refs: "ObjectRef[R]", *, timeout: Optional[float] = None) -> R: + ... + + +@overload +def get( + object_refs: Sequence[CompiledDAGRef], *, timeout: Optional[float] = None +) -> List[Any]: + ... + + +@overload +def get(object_refs: CompiledDAGRef, *, timeout: Optional[float] = None) -> Any: + ... + + +@PublicAPI +@client_mode_hook +def get( + object_refs: Union[ + "ObjectRef[Any]", + Sequence["ObjectRef[Any]"], + CompiledDAGRef, + Sequence[CompiledDAGRef], + ], + *, + timeout: Optional[float] = None, + _tensor_transport: Optional[str] = None, +) -> Union[Any, List[Any]]: + """Get a remote object or a list of remote objects from the object store. + + This method blocks until the object corresponding to the object ref is + available in the local object store. If this object is not in the local + object store, it will be shipped from an object store that has it (once the + object has been created). If object_refs is a list, then the objects + corresponding to each object in the list will be returned. + + Ordering for an input list of object refs is preserved for each object + returned. That is, if an object ref to A precedes an object ref to B in the + input list, then A will precede B in the returned list. + + This method will issue a warning if it's running inside async context, + you can use ``await object_ref`` instead of ``ray.get(object_ref)``. For + a list of object refs, you can use ``await asyncio.gather(*object_refs)``. + + Passing :class:`~ObjectRefGenerator` is not allowed. + + Related patterns and anti-patterns: + + - :doc:`/ray-core/patterns/ray-get-loop` + - :doc:`/ray-core/patterns/unnecessary-ray-get` + - :doc:`/ray-core/patterns/ray-get-submission-order` + - :doc:`/ray-core/patterns/ray-get-too-many-objects` + + + Args: + object_refs: Object ref of the object to get or a list of object refs + to get. + timeout (Optional[float]): The maximum amount of time in seconds to + wait before returning. Set this to None will block until the + corresponding object becomes available. Setting ``timeout=0`` will + return the object immediately if it's available, else raise + GetTimeoutError in accordance with the above docstring. + _tensor_transport: [Alpha] The tensor transport to use to fetch `torch.Tensors` found in the Ray Direct Transport object. Currently, this supports "object_store" and "nixl". + + Returns: + A Python object or a list of Python objects. + + Raises: + GetTimeoutError: A GetTimeoutError is raised if a timeout is set and + the get takes longer than timeout to return. + Exception: An exception is raised immediately if any task that created + the object or that created one of the objects raised an exception, + without waiting for the remaining ones to finish. + """ + worker = global_worker + worker.check_connected() + + if hasattr(worker, "core_worker") and worker.core_worker.current_actor_is_asyncio(): + global blocking_get_inside_async_warned + if not blocking_get_inside_async_warned: + if ray_constants.env_bool( + RAY_WARN_BLOCKING_GET_INSIDE_ASYNC_ENV_VAR, + True, + ): + logger.warning( + "Using blocking ray.get inside async actor. " + "This blocks the event loop. Please use `await` " + "on object ref with asyncio.gather if you want to " + "yield execution to the event loop instead." + ) + blocking_get_inside_async_warned = True + + with profiling.profile("ray.get"): + # TODO(sang): Should make ObjectRefGenerator + # compatible to ray.get for dataset. + if isinstance(object_refs, ObjectRefGenerator): + return object_refs + + if isinstance(object_refs, CompiledDAGRef): + return object_refs.get(timeout=timeout) + + if isinstance(object_refs, list): + all_compiled_dag_refs = True + any_compiled_dag_refs = False + for object_ref in object_refs: + is_dag_ref = isinstance(object_ref, CompiledDAGRef) + all_compiled_dag_refs = all_compiled_dag_refs and is_dag_ref + any_compiled_dag_refs = any_compiled_dag_refs or is_dag_ref + if all_compiled_dag_refs: + return [object_ref.get(timeout=timeout) for object_ref in object_refs] + elif any_compiled_dag_refs: + raise ValueError( + "Invalid type of object refs. 'object_refs' must be a list of " + "CompiledDAGRefs if there is any CompiledDAGRef within it. " + ) + + is_individual_id = isinstance(object_refs, ray.ObjectRef) + if is_individual_id: + object_refs = [object_refs] + + if not isinstance(object_refs, list): + raise ValueError( + f"Invalid type of object refs, {type(object_refs)}, is given. " + "'object_refs' must either be an ObjectRef or a list of ObjectRefs. " + ) + + values, debugger_breakpoint = worker.get_objects( + object_refs, timeout=timeout, _tensor_transport=_tensor_transport + ) + for i, value in enumerate(values): + if isinstance(value, RayError): + # If the object was lost and it wasn't due to owner death, it may be + # because the object store is full and objects needed to be evicted. + if isinstance(value, ray.exceptions.ObjectLostError) and not isinstance( + value, ray.exceptions.OwnerDiedError + ): + worker.core_worker.log_plasma_usage() + if isinstance(value, RayTaskError): + raise value.as_instanceof_cause() + else: + raise value + + if is_individual_id: + values = values[0] + + if debugger_breakpoint != b"": + frame = sys._getframe().f_back + rdb = ray.util.pdb._connect_ray_pdb( + host=None, + port=None, + patch_stdstreams=False, + quiet=None, + breakpoint_uuid=( + debugger_breakpoint.decode() if debugger_breakpoint else None + ), + debugger_external=worker.ray_debugger_external, + ) + rdb.set_trace(frame=frame) + + return values + + +@PublicAPI +@client_mode_hook +def put( + value: Any, + *, + _owner: Optional["ray.actor.ActorHandle"] = None, + _tensor_transport: str = "object_store", +) -> "ray.ObjectRef": + """Store an object in the object store. + + The object may not be evicted while a reference to the returned ID exists. + + Related patterns and anti-patterns: + + - :doc:`/ray-core/patterns/return-ray-put` + - :doc:`/ray-core/patterns/pass-large-arg-by-value` + - :doc:`/ray-core/patterns/closure-capture-large-objects` + + Args: + value: The Python object to be stored. + _owner [Experimental]: The actor that should own this object. This + allows creating objects with lifetimes decoupled from that of the + creating process. The owner actor must be passed a reference to the + object prior to the object creator exiting, otherwise the reference + will still be lost. *Note that this argument is an experimental API + and should be avoided if possible.* + _tensor_transport: [Alpha] The tensor transport to use for the GPU object. Currently, this supports "object_store" and "nixl" for tensor transport in ray.put(). + + Returns: + The object ref assigned to this value. + """ + worker = global_worker + worker.check_connected() + + if _owner is None: + serialize_owner_address = None + elif isinstance(_owner, ray.actor.ActorHandle): + # Ensure GlobalState is connected + ray._private.state.state._connect_and_get_accessor() + serialize_owner_address = ( + ray._raylet._get_actor_serialized_owner_address_or_none( + ray._private.state.state.get_actor_info(_owner._actor_id) + ) + ) + if not serialize_owner_address: + raise RuntimeError(f"{_owner} is not alive, it's worker_id is empty!") + else: + raise TypeError(f"Expect an `ray.actor.ActorHandle`, but got: {type(_owner)}") + + with profiling.profile("ray.put"): + try: + object_ref = worker.put_object( + value, + owner_address=serialize_owner_address, + _tensor_transport=_tensor_transport, + ) + except ObjectStoreFullError: + logger.info( + "Put failed since the value was either too large or the " + "store was full of pinned objects." + ) + raise + return object_ref + + +# Global variable to make sure we only send out the warning once. +blocking_wait_inside_async_warned = False + + +@PublicAPI +@client_mode_hook +def wait( + ray_waitables: List[Union[ObjectRef, ObjectRefGenerator]], + *, + num_returns: int = 1, + timeout: Optional[float] = None, + fetch_local: bool = True, +) -> Tuple[ + List[Union[ObjectRef, ObjectRefGenerator]], + List[Union[ObjectRef, ObjectRefGenerator]], +]: + """Return a list of IDs that are ready and a list of IDs that are not. + + If timeout is set, the function returns either when the requested number of + IDs are ready or when the timeout is reached, whichever occurs first. If it + is not set, the function simply waits until that number of objects is ready + and returns that exact number of object refs. + + `ray_waitables` is a list of :class:`~ray.ObjectRef` and + :class:`~ray.ObjectRefGenerator`. + + The method returns two lists, ready and unready `ray_waitables`. + + ObjectRef: + object refs that correspond to objects that are available + in the object store are in the first list. + The rest of the object refs are in the second list. + + ObjectRefGenerator: + Generators whose next reference (that will be obtained + via `next(generator)`) has a corresponding object available + in the object store are in the first list. + All other generators are placed in the second list. + + Ordering of the input list of ray_waitables is preserved. That is, if A + precedes B in the input list, and both are in the ready list, then A will + precede B in the ready list. This also holds true if A and B are both in + the remaining list. + + This method will issue a warning if it's running inside an async context. + Instead of ``ray.wait(ray_waitables)``, you can use + ``await asyncio.wait(ray_waitables)``. + + Related patterns and anti-patterns: + + - :doc:`/ray-core/patterns/limit-pending-tasks` + - :doc:`/ray-core/patterns/ray-get-submission-order` + + Args: + ray_waitables: List of :class:`~ObjectRef` or + :class:`~ObjectRefGenerator` for objects that may or may + not be ready. Note that these must be unique. + num_returns: The number of ray_waitables that should be returned. + timeout: The maximum amount of time in seconds to wait before + returning. + fetch_local: If True, wait for the object to be downloaded onto + the local node before returning it as ready. If the `ray_waitable` + is a generator, it will wait until the next object in the generator + is downloaed. If False, ray.wait() will not trigger fetching of + objects to the local node and will return immediately once the + object is available anywhere in the cluster. + + Returns: + A list of object refs that are ready and a list of the remaining object + IDs. + """ + worker = global_worker + worker.check_connected() + + if ( + hasattr(worker, "core_worker") + and worker.core_worker.current_actor_is_asyncio() + and timeout != 0 + ): + global blocking_wait_inside_async_warned + if not blocking_wait_inside_async_warned: + logger.debug( + "Using blocking ray.wait inside async method. " + "This blocks the event loop. Please use `await` " + "on object ref with asyncio.wait. " + ) + blocking_wait_inside_async_warned = True + + if isinstance(ray_waitables, ObjectRef) or isinstance( + ray_waitables, ObjectRefGenerator + ): + raise TypeError( + "wait() expected a list of ray.ObjectRef or ray.ObjectRefGenerator" + ", got a single ray.ObjectRef or ray.ObjectRefGenerator " + f"{ray_waitables}" + ) + + if not isinstance(ray_waitables, list): + raise TypeError( + "wait() expected a list of ray.ObjectRef or " + "ray.ObjectRefGenerator, " + f"got {type(ray_waitables)}" + ) + + if timeout is not None and timeout < 0: + raise ValueError( + "The 'timeout' argument must be nonnegative. " f"Received {timeout}" + ) + + for ray_waitable in ray_waitables: + if not isinstance(ray_waitable, ObjectRef) and not isinstance( + ray_waitable, ObjectRefGenerator + ): + raise TypeError( + "wait() expected a list of ray.ObjectRef or " + "ray.ObjectRefGenerator, " + f"got list containing {type(ray_waitable)}" + ) + worker.check_connected() + + # TODO(swang): Check main thread. + with profiling.profile("ray.wait"): + # TODO(rkn): This is a temporary workaround for + # https://github.com/ray-project/ray/issues/997. However, it should be + # fixed in Arrow instead of here. + if len(ray_waitables) == 0: + return [], [] + + if len(ray_waitables) != len(set(ray_waitables)): + raise ValueError("Wait requires a list of unique ray_waitables.") + if num_returns <= 0: + raise ValueError("Invalid number of objects to return %d." % num_returns) + if num_returns > len(ray_waitables): + raise ValueError( + "num_returns cannot be greater than the number " + "of ray_waitables provided to ray.wait." + ) + + timeout = timeout if timeout is not None else 10**6 + timeout_milliseconds = int(timeout * 1000) + ready_ids, remaining_ids = worker.core_worker.wait( + ray_waitables, + num_returns, + timeout_milliseconds, + fetch_local, + ) + return ready_ids, remaining_ids + + +@PublicAPI +@client_mode_hook +def get_actor(name: str, namespace: Optional[str] = None) -> "ray.actor.ActorHandle": + """Get a handle to a named actor. + + Gets a handle to an actor with the given name. The actor must + have been created with Actor.options(name="name").remote(). This + works for both detached & non-detached actors. + + This method is a sync call and it'll timeout after 60s. This can be modified + by setting OS env RAY_gcs_server_request_timeout_seconds before starting + the cluster. + + Args: + name: The name of the actor. + namespace: The namespace of the actor, or None to specify the current + namespace. + + Returns: + ActorHandle to the actor. + + Raises: + ValueError: if the named actor does not exist. + """ + if not name: + raise ValueError("Please supply a non-empty value to get_actor") + + if namespace is not None: + ray._private.utils.validate_namespace(namespace) + + worker = global_worker + worker.check_connected() + return worker.core_worker.get_named_actor_handle(name, namespace or "") + + +@PublicAPI +@client_mode_hook +def kill(actor: "ray.actor.ActorHandle", *, no_restart: bool = True): + """Kill an actor forcefully. + + This will interrupt any running tasks on the actor, causing them to fail + immediately. ``atexit`` handlers installed in the actor will not be run. + + If you want to kill the actor but let pending tasks finish, + you can call ``actor.__ray_terminate__.remote()`` instead to queue a + termination task. Any ``atexit`` handlers installed in the actor *will* + be run in this case. + + If the actor is a detached actor, subsequent calls to get its handle via + ray.get_actor will fail. + + Args: + actor: Handle to the actor to kill. + no_restart: Whether or not this actor should be restarted if + it's a restartable actor. + """ + worker = global_worker + worker.check_connected() + if not isinstance(actor, ray.actor.ActorHandle): + raise ValueError( + "ray.kill() only supported for actors. For tasks, try ray.cancel(). " + "Got: {}.".format(type(actor)) + ) + worker.core_worker.kill_actor(actor._ray_actor_id, no_restart) + + +@PublicAPI +@client_mode_hook +def cancel( + ray_waitable: Union["ObjectRef[R]", "ObjectRefGenerator[R]"], + *, + force: bool = False, + recursive: bool = True, +) -> None: + """Cancels a task. + + Cancel API has a different behavior depending on if it is a remote function + (Task) or a remote Actor method (Actor Task). + + Task: + If the specified Task is pending execution, it is cancelled and not + executed. If the Task is currently executing, the behavior depends + on the `force` flag. When `force=False`, a KeyboardInterrupt is + raised in Python and when `force=True`, the executing Task + immediately exits. If the Task is already finished, nothing happens. + + Cancelled Tasks aren't retried. `max_task_retries` aren't respected. + + Calling ray.get on a cancelled Task raises a TaskCancelledError + if the Task has been scheduled or interrupted. + It raises a WorkerCrashedError if `force=True`. + + If `recursive=True`, all the child Tasks and Actor Tasks + are cancelled. If `force=True` and `recursive=True`, `force=True` + is ignored for child Actor Tasks. + + Actor Task: + If the specified Task is pending execution, it is cancelled and not + executed. If the Task is currently executing, the behavior depends + on the execution model of an Actor. If it is a regular Actor + or a threaded Actor, the execution isn't cancelled. + Actor Tasks cannot be interrupted because Actors have + states. If it is an async Actor, Ray cancels a `asyncio.Task`. + The semantic of cancellation is equivalent to asyncio's cancellation. + https://docs.python.org/3/library/asyncio-task.html#task-cancellation + If the Task has finished, nothing happens. + + Only `force=False` is allowed for an Actor Task. Otherwise, it raises + `ValueError`. Use `ray.kill(actor)` instead to kill an Actor. + + Cancelled Tasks aren't retried. `max_task_retries` aren't respected. + + Calling ray.get on a cancelled Task raises a TaskCancelledError + if the Task has been scheduled or interrupted. Also note that + only async actor tasks can be interrupted. + + If `recursive=True`, all the child Tasks and actor Tasks + are cancelled. + + Args: + ray_waitable: :class:`~ObjectRef` and + :class:`~ObjectRefGenerator` + returned by the task that should be canceled. + force: Whether to force-kill a running task by killing + the worker that is running the task. + recursive: Whether to try to cancel tasks submitted by the + task specified. + """ + worker = ray._private.worker.global_worker + worker.check_connected() + + if isinstance(ray_waitable, ray._raylet.ObjectRefGenerator): + assert hasattr(ray_waitable, "_generator_ref") + ray_waitable = ray_waitable._generator_ref + + if not isinstance(ray_waitable, ray.ObjectRef): + raise TypeError( + "ray.cancel() only supported for object refs. " + f"For actors, try ray.kill(). Got: {type(ray_waitable)}." + ) + return worker.core_worker.cancel_task(ray_waitable, force, recursive) + + +def _mode(worker=global_worker): + """This is a wrapper around worker.mode. + + We use this wrapper so that in the remote decorator, we can call _mode() + instead of worker.mode. The difference is that when we attempt to + serialize remote functions, we don't attempt to serialize the worker + object, which cannot be serialized. + """ + return worker.mode + + +def _make_remote(function_or_class, options): + if not function_or_class.__module__: + function_or_class.__module__ = "global" + + if inspect.isfunction(function_or_class) or is_cython(function_or_class): + ray_option_utils.validate_task_options(options, in_options=False) + return ray.remote_function.RemoteFunction( + Language.PYTHON, + function_or_class, + None, + options, + ) + + if inspect.isclass(function_or_class): + ray_option_utils.validate_actor_options(options, in_options=False) + return ray.actor._make_actor(function_or_class, options) + + raise TypeError( + "The @ray.remote decorator must be applied to either a function or a class." + ) + + +class RemoteDecorator(Protocol): + @overload + def __call__(self, __t: Type[T]) -> ActorClass[T]: + ... + + @overload + def __call__(self, __function: Callable[[], R]) -> RemoteFunctionNoArgs[R]: + ... + + @overload + def __call__(self, __function: Callable[[T0], R]) -> RemoteFunction0[R, T0]: + ... + + @overload + def __call__(self, __function: Callable[[T0, T1], R]) -> RemoteFunction1[R, T0, T1]: + ... + + @overload + def __call__( + self, __function: Callable[[T0, T1, T2], R] + ) -> RemoteFunction2[R, T0, T1, T2]: + ... + + @overload + def __call__( + self, __function: Callable[[T0, T1, T2, T3], R] + ) -> RemoteFunction3[R, T0, T1, T2, T3]: + ... + + @overload + def __call__( + self, __function: Callable[[T0, T1, T2, T3, T4], R] + ) -> RemoteFunction4[R, T0, T1, T2, T3, T4]: + ... + + @overload + def __call__( + self, __function: Callable[[T0, T1, T2, T3, T4, T5], R] + ) -> RemoteFunction5[R, T0, T1, T2, T3, T4, T5]: + ... + + @overload + def __call__( + self, __function: Callable[[T0, T1, T2, T3, T4, T5, T6], R] + ) -> RemoteFunction6[R, T0, T1, T2, T3, T4, T5, T6]: + ... + + @overload + def __call__( + self, __function: Callable[[T0, T1, T2, T3, T4, T5, T6, T7], R] + ) -> RemoteFunction7[R, T0, T1, T2, T3, T4, T5, T6, T7]: + ... + + @overload + def __call__( + self, __function: Callable[[T0, T1, T2, T3, T4, T5, T6, T7, T8], R] + ) -> RemoteFunction8[R, T0, T1, T2, T3, T4, T5, T6, T7, T8]: + ... + + @overload + def __call__( + self, __function: Callable[[T0, T1, T2, T3, T4, T5, T6, T7, T8, T9], R] + ) -> RemoteFunction9[R, T0, T1, T2, T3, T4, T5, T6, T7, T8, T9]: + ... + + +@overload +def remote(__t: Type[T]) -> ActorClass[T]: + ... + + +@overload +def remote(__function: Callable[[], R]) -> RemoteFunctionNoArgs[R]: + ... + + +@overload +def remote(__function: Callable[[T0], R]) -> RemoteFunction0[R, T0]: + ... + + +@overload +def remote(__function: Callable[[T0, T1], R]) -> RemoteFunction1[R, T0, T1]: + ... + + +@overload +def remote(__function: Callable[[T0, T1, T2], R]) -> RemoteFunction2[R, T0, T1, T2]: + ... + + +@overload +def remote( + __function: Callable[[T0, T1, T2, T3], R] +) -> RemoteFunction3[R, T0, T1, T2, T3]: + ... + + +@overload +def remote( + __function: Callable[[T0, T1, T2, T3, T4], R] +) -> RemoteFunction4[R, T0, T1, T2, T3, T4]: + ... + + +@overload +def remote( + __function: Callable[[T0, T1, T2, T3, T4, T5], R] +) -> RemoteFunction5[R, T0, T1, T2, T3, T4, T5]: + ... + + +@overload +def remote( + __function: Callable[[T0, T1, T2, T3, T4, T5, T6], R] +) -> RemoteFunction6[R, T0, T1, T2, T3, T4, T5, T6]: + ... + + +@overload +def remote( + __function: Callable[[T0, T1, T2, T3, T4, T5, T6, T7], R] +) -> RemoteFunction7[R, T0, T1, T2, T3, T4, T5, T6, T7]: + ... + + +@overload +def remote( + __function: Callable[[T0, T1, T2, T3, T4, T5, T6, T7, T8], R] +) -> RemoteFunction8[R, T0, T1, T2, T3, T4, T5, T6, T7, T8]: + ... + + +@overload +def remote( + __function: Callable[[T0, T1, T2, T3, T4, T5, T6, T7, T8, T9], R] +) -> RemoteFunction9[R, T0, T1, T2, T3, T4, T5, T6, T7, T8, T9]: + ... + + +# Passing options +@overload +def remote( + *, + num_returns: Union[int, Literal["streaming"]] = Undefined, + num_cpus: Union[int, float] = Undefined, + num_gpus: Union[int, float] = Undefined, + resources: Dict[str, float] = Undefined, + accelerator_type: str = Undefined, + memory: Union[int, float] = Undefined, + max_calls: int = Undefined, + max_restarts: int = Undefined, + max_task_retries: int = Undefined, + max_retries: int = Undefined, + runtime_env: Dict[str, Any] = Undefined, + retry_exceptions: bool = Undefined, + scheduling_strategy: Union[ + None, Literal["DEFAULT"], Literal["SPREAD"], PlacementGroupSchedulingStrategy + ] = Undefined, + label_selector: Dict[str, str] = Undefined, + fallback_strategy: List[Dict[str, Any]] = Undefined, +) -> RemoteDecorator: + ... + + +@PublicAPI +def remote( + *args, **kwargs +) -> Union[ray.remote_function.RemoteFunction, ray.actor.ActorClass]: + """Defines a remote function or an actor class. + + This function can be used as a decorator with no arguments + to define a remote function or actor as follows: + + .. testcode:: + + import ray + + @ray.remote + def f(a, b, c): + return a + b + c + + object_ref = f.remote(1, 2, 3) + result = ray.get(object_ref) + assert result == (1 + 2 + 3) + + @ray.remote + class Foo: + def __init__(self, arg): + self.x = arg + + def method(self, a): + return self.x + a + + actor_handle = Foo.remote(123) + object_ref = actor_handle.method.remote(321) + result = ray.get(object_ref) + assert result == (123 + 321) + + Equivalently, use a function call to create a remote function or actor. + + .. testcode:: + + def g(a, b, c): + return a + b + c + + remote_g = ray.remote(g) + object_ref = remote_g.remote(1, 2, 3) + assert ray.get(object_ref) == (1 + 2 + 3) + + class Bar: + def __init__(self, arg): + self.x = arg + + def method(self, a): + return self.x + a + + RemoteBar = ray.remote(Bar) + actor_handle = RemoteBar.remote(123) + object_ref = actor_handle.method.remote(321) + result = ray.get(object_ref) + assert result == (123 + 321) + + + It can also be used with specific keyword arguments as follows: + + .. testcode:: + + @ray.remote(num_gpus=1, max_calls=1, num_returns=2) + def f(): + return 1, 2 + + @ray.remote(num_cpus=2, resources={"CustomResource": 1}) + class Foo: + def method(self): + return 1 + + Remote task and actor objects returned by @ray.remote can also be + dynamically modified with the same arguments as above using + ``.options()`` as follows: + + .. testcode:: + :hide: + + ray.shutdown() + + ray.init(num_cpus=5, num_gpus=5) + + .. testcode:: + + @ray.remote(num_gpus=1, max_calls=1, num_returns=2) + def f(): + return 1, 2 + + f_with_2_gpus = f.options(num_gpus=2) + object_refs = f_with_2_gpus.remote() + assert ray.get(object_refs) == [1, 2] + + @ray.remote(num_cpus=2, resources={"CustomResource": 1}) + class Foo: + def method(self): + return 1 + + Foo_with_no_resources = Foo.options(num_cpus=1, resources=None) + foo_actor = Foo_with_no_resources.remote() + assert ray.get(foo_actor.method.remote()) == 1 + + + A remote actor will be terminated when all actor handle to it + in Python is deleted, which will cause them to complete any outstanding + work and then shut down. If you only have 1 reference to an actor handle, + calling ``del actor`` *could* trigger actor deletion. Note that your program + may have multiple references to the same ActorHandle, and actor termination + will not occur until the reference count goes to 0. See the Python + documentation for more context about object deletion. + https://docs.python.org/3.9/reference/datamodel.html#object.__del__ + + If you want to kill actors immediately, you can also call ``ray.kill(actor)``. + + .. tip:: + Avoid repeatedly passing in large arguments to remote task or method calls. + + Instead, use ray.put to create a copy of the object in the object store. + + See :ref:`more info here `. + + Args: + num_returns: This is only for *remote functions*. It specifies + the number of object refs returned by the remote function + invocation. The default value is 1. + Pass "dynamic" to allow the task to decide how many + return values to return during execution, and the caller will + receive an ObjectRef[DynamicObjectRefGenerator]. + See :ref:`dynamic generators ` for more details. + num_cpus: The quantity of CPU resources to reserve + for this task or for the lifetime of the actor. + By default, tasks use 1 CPU resource and actors use 1 CPU + for scheduling and 0 CPU for running + (This means, by default, actors cannot get scheduled on a zero-cpu node, + but an infinite number of them can run on any non-zero cpu node. + The default value for actors was chosen for historical reasons. + It's recommended to always explicitly set num_cpus for actors + to avoid any surprises. + If resources are specified explicitly, + they are required for both scheduling and running.) + See :ref:`specifying resource requirements ` + for more details. + num_gpus: The quantity of GPU resources to reserve + for this task or for the lifetime of the actor. + The default value is 0. + See :ref:`Ray GPU support ` for more details. + resources (Dict[str, float]): The quantity of various + :ref:`custom resources ` + to reserve for this task or for the lifetime of the actor. + This is a dictionary mapping strings (resource names) to floats. + By default it is empty. + label_selector: [Experimental] If specified, the labels required for the node on + which this actor can be scheduled on. The label selector consist of key-value pairs, + where the keys are label names and the value are expressions consisting of an operator + with label values or just a value to indicate equality. + fallback_strategy: [Experimental] If specified, expresses soft constraints for scheduling + through a list of dicts of decorator options to fall back on when scheduling on a node. + Decorator options are evaluated together during scheduling. The first satisfied + dict of options is used. Currently only `label_selector` is a supported option. + accelerator_type: If specified, requires that the task or actor run + on a node with the specified type of accelerator. + See :ref:`accelerator types `. + memory: The heap memory request in bytes for this task/actor, + rounded down to the nearest integer. + max_calls: Only for *remote functions*. This specifies the + maximum number of times that a given worker can execute + the given remote function before it must exit + (this can be used to address :ref:`memory leaks ` in third-party + libraries or to reclaim resources that cannot easily be + released, e.g., GPU memory that was acquired by TensorFlow). + By default this is infinite for CPU tasks and 1 for GPU tasks + (to force GPU tasks to release resources after finishing). + max_restarts: Only for *actors*. This specifies the maximum + number of times that the actor should be restarted when it dies + unexpectedly. The minimum valid value is 0 (default), + which indicates that the actor doesn't need to be restarted. + A value of -1 indicates that an actor should be restarted + indefinitely. + See :ref:`actor fault tolerance ` for more details. + max_task_retries: Only for *actors*. How many times to + retry an actor task if the task fails due to a system error, + e.g., the actor has died. If set to -1, the system will + retry the failed task until the task succeeds, or the actor + has reached its max_restarts limit. If set to `n > 0`, the + system will retry the failed task up to n times, after which the + task will throw a `RayActorError` exception upon :obj:`ray.get`. + Note that Python exceptions are not considered system errors + and will not trigger retries. + The default value is 0. + See :ref:`actor fault tolerance ` for more details. + max_retries: Only for *remote functions*. This specifies + the maximum number of times that the remote function + should be rerun when the worker process executing it + crashes unexpectedly. The minimum valid value is 0, + the default value is 3, and a value of -1 indicates + infinite retries. + See :ref:`task fault tolerance ` for more details. + allow_out_of_order_execution: Only for *actors*. Whether Ray executes actor + tasks out of order. If you're using multi-threaded (``max_concurrency > 1``) + or async actors, you can't set this to False. Defaults to True if you're + using multi-threaded or async actors, and False otherwise. Actor task + retries are always executed out of order. + runtime_env (Dict[str, Any]): Specifies the runtime environment for + this actor or task and its children. See + :ref:`runtime-environments` for detailed documentation. + retry_exceptions: Only for *remote functions*. This specifies whether + application-level errors should be retried up to max_retries times. + This can be a boolean or a list of exceptions that should be retried. + See :ref:`task fault tolerance ` for more details. + scheduling_strategy: Strategy about how to + schedule a remote function or actor. Possible values are + None: ray will figure out the scheduling strategy to use, it + will either be the PlacementGroupSchedulingStrategy using parent's + placement group if parent has one and has + placement_group_capture_child_tasks set to true, + or "DEFAULT"; + "DEFAULT": default hybrid scheduling; + "SPREAD": best effort spread scheduling; + `PlacementGroupSchedulingStrategy`: + placement group based scheduling; + `NodeAffinitySchedulingStrategy`: + node id based affinity scheduling. + See :ref:`Ray scheduling strategies ` + for more details. + _labels: The key-value labels of a task or actor. + """ + # "callable" returns true for both function and class. + if len(args) == 1 and len(kwargs) == 0 and callable(args[0]): + # This is the case where the decorator is just @ray.remote. + # "args[0]" is the class or function under the decorator. + return _make_remote(args[0], {}) + assert len(args) == 0 and len(kwargs) > 0, ray_option_utils.remote_args_error_string + return functools.partial(_make_remote, options=kwargs) diff --git a/lib/python3.12/site-packages/ray/_private/workers/__init__.py b/lib/python3.12/site-packages/ray/_private/workers/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/lib/python3.12/site-packages/ray/_private/workers/__pycache__/__init__.cpython-312.pyc b/lib/python3.12/site-packages/ray/_private/workers/__pycache__/__init__.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..1c8a7e4595cfb8878328c20f30964507a6acd5ef Binary files /dev/null and b/lib/python3.12/site-packages/ray/_private/workers/__pycache__/__init__.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/_private/workers/__pycache__/default_worker.cpython-312.pyc b/lib/python3.12/site-packages/ray/_private/workers/__pycache__/default_worker.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..f38236e79ec8571fcb569270be53f7307616db7b Binary files /dev/null and b/lib/python3.12/site-packages/ray/_private/workers/__pycache__/default_worker.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/_private/workers/__pycache__/setup_worker.cpython-312.pyc b/lib/python3.12/site-packages/ray/_private/workers/__pycache__/setup_worker.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..0b96e8680e721701623a87f32b18ed6348452678 Binary files /dev/null and b/lib/python3.12/site-packages/ray/_private/workers/__pycache__/setup_worker.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/_private/workers/default_worker.py b/lib/python3.12/site-packages/ray/_private/workers/default_worker.py new file mode 100644 index 0000000000000000000000000000000000000000..12cf83040574e8f07e6c3eda1a0d7beb746701d9 --- /dev/null +++ b/lib/python3.12/site-packages/ray/_private/workers/default_worker.py @@ -0,0 +1,329 @@ +import argparse +import base64 +import json +import os +import sys +import time + +import ray +import ray._private.node +import ray._private.ray_constants as ray_constants +import ray._private.utils +import ray.actor +from ray._common.ray_constants import ( + LOGGING_ROTATE_BACKUP_COUNT, + LOGGING_ROTATE_BYTES, +) +from ray._private.async_compat import try_install_uvloop +from ray._private.parameter import RayParams +from ray._private.ray_logging import get_worker_log_file_name +from ray._private.runtime_env.setup_hook import load_and_execute_setup_hook + +parser = argparse.ArgumentParser( + description=("Parse addresses for the worker to connect to.") +) +parser.add_argument( + "--cluster-id", + required=True, + type=str, + help="the auto-generated ID of the cluster", +) +parser.add_argument( + "--node-id", + required=True, + type=str, + help="the auto-generated ID of the node", +) +parser.add_argument( + "--node-ip-address", + required=True, + type=str, + help="the ip address of the worker's node", +) +parser.add_argument( + "--node-manager-port", required=True, type=int, help="the port of the worker's node" +) +parser.add_argument( + "--raylet-ip-address", + required=False, + type=str, + default=None, + help="the ip address of the worker's raylet", +) +parser.add_argument( + "--redis-address", required=True, type=str, help="the address to use for Redis" +) +parser.add_argument( + "--gcs-address", required=True, type=str, help="the address to use for GCS" +) +parser.add_argument( + "--redis-username", + required=False, + type=str, + default=None, + help="the username to use for Redis", +) +parser.add_argument( + "--redis-password", + required=False, + type=str, + default=None, + help="the password to use for Redis", +) +parser.add_argument( + "--object-store-name", required=True, type=str, help="the object store's name" +) +parser.add_argument("--raylet-name", required=False, type=str, help="the raylet's name") +parser.add_argument( + "--logging-level", + required=False, + type=str, + default=ray_constants.LOGGER_LEVEL, + choices=ray_constants.LOGGER_LEVEL_CHOICES, + help=ray_constants.LOGGER_LEVEL_HELP, +) +parser.add_argument( + "--logging-format", + required=False, + type=str, + default=ray_constants.LOGGER_FORMAT, + help=ray_constants.LOGGER_FORMAT_HELP, +) +parser.add_argument( + "--temp-dir", + required=False, + type=str, + default=None, + help="Specify the path of the temporary directory use by Ray process.", +) +parser.add_argument( + "--load-code-from-local", + default=False, + action="store_true", + help="True if code is loaded from local files, as opposed to the GCS.", +) +parser.add_argument( + "--worker-type", + required=False, + type=str, + default="WORKER", + help="Specify the type of the worker process", +) +parser.add_argument( + "--metrics-agent-port", + required=True, + type=int, + help="the port of the node's metric agent.", +) +parser.add_argument( + "--runtime-env-agent-port", + required=True, + type=int, + default=None, + help="The port on which the runtime env agent receives HTTP requests.", +) +parser.add_argument( + "--object-spilling-config", + required=False, + type=str, + default="", + help="The configuration of object spilling. Only used by I/O workers.", +) +parser.add_argument( + "--logging-rotate-bytes", + required=False, + type=int, + default=LOGGING_ROTATE_BYTES, + help="Specify the max bytes for rotating " + "log file, default is " + f"{LOGGING_ROTATE_BYTES} bytes.", +) +parser.add_argument( + "--logging-rotate-backup-count", + required=False, + type=int, + default=LOGGING_ROTATE_BACKUP_COUNT, + help="Specify the backup count of rotated log file, default is " + f"{LOGGING_ROTATE_BACKUP_COUNT}.", +) +parser.add_argument( + "--runtime-env-hash", + required=False, + type=int, + default=0, + help="The computed hash of the runtime env for this worker.", +) +parser.add_argument( + "--startup-token", + required=True, + type=int, + help="The startup token assigned to this worker process by the raylet.", +) +parser.add_argument( + "--ray-debugger-external", + default=False, + action="store_true", + help="True if Ray debugger is made available externally.", +) +parser.add_argument( + "--session-name", required=False, help="The current Ray session name" +) +parser.add_argument( + "--webui", + required=False, + help="The address of web ui", +) +parser.add_argument( + "--worker-launch-time-ms", + required=True, + type=int, + help="The time when raylet starts to launch the worker process.", +) + +parser.add_argument( + "--worker-preload-modules", + type=str, + required=False, + help=( + "A comma-separated list of Python module names " + "to import before accepting work." + ), +) +parser.add_argument( + "--enable-resource-isolation", + type=bool, + required=False, + default=False, + help=( + "If true, core worker enables resource isolation by adding itself into appropriate cgroup." + ), +) + +if __name__ == "__main__": + # NOTE(sang): For some reason, if we move the code below + # to a separate function, tensorflow will capture that method + # as a step function. For more details, check out + # https://github.com/ray-project/ray/pull/12225#issue-525059663. + args = parser.parse_args() + ray._private.ray_logging.setup_logger(args.logging_level, args.logging_format) + worker_launched_time_ms = time.time_ns() // 1e6 + if args.worker_type == "WORKER": + mode = ray.WORKER_MODE + elif args.worker_type == "SPILL_WORKER": + mode = ray.SPILL_WORKER_MODE + elif args.worker_type == "RESTORE_WORKER": + mode = ray.RESTORE_WORKER_MODE + else: + raise ValueError("Unknown worker type: " + args.worker_type) + + # Try installing uvloop as default event-loop implementation + # for asyncio + try_install_uvloop() + + ray_params = RayParams( + node_ip_address=args.node_ip_address, + node_manager_port=args.node_manager_port, + redis_address=args.redis_address, + redis_username=args.redis_username, + redis_password=args.redis_password, + plasma_store_socket_name=args.object_store_name, + raylet_socket_name=args.raylet_name, + temp_dir=args.temp_dir, + metrics_agent_port=args.metrics_agent_port, + runtime_env_agent_port=args.runtime_env_agent_port, + gcs_address=args.gcs_address, + session_name=args.session_name, + webui=args.webui, + cluster_id=args.cluster_id, + node_id=args.node_id, + ) + node = ray._private.node.Node( + ray_params, + head=False, + shutdown_at_exit=False, + spawn_reaper=False, + connect_only=True, + default_worker=True, + ) + + # NOTE(suquark): We must initialize the external storage before we + # connect to raylet. Otherwise we may receive requests before the + # external storage is initialized. + if mode == ray.RESTORE_WORKER_MODE or mode == ray.SPILL_WORKER_MODE: + from ray._private import external_storage + + if args.object_spilling_config: + object_spilling_config = base64.b64decode(args.object_spilling_config) + object_spilling_config = json.loads(object_spilling_config) + else: + object_spilling_config = {} + external_storage.setup_external_storage( + object_spilling_config, node.node_id, node.session_name + ) + + ray._private.worker._global_node = node + ray._private.worker.connect( + node, + node.session_name, + mode=mode, + runtime_env_hash=args.runtime_env_hash, + startup_token=args.startup_token, + ray_debugger_external=args.ray_debugger_external, + worker_launch_time_ms=args.worker_launch_time_ms, + worker_launched_time_ms=worker_launched_time_ms, + ) + + worker = ray._private.worker.global_worker + + stdout_fileno = sys.stdout.fileno() + stderr_fileno = sys.stderr.fileno() + # We also manually set sys.stdout and sys.stderr because that seems to + # have an effect on the output buffering. Without doing this, stdout + # and stderr are heavily buffered resulting in seemingly lost logging + # statements. We never want to close the stdout file descriptor, dup2 will + # close it when necessary and we don't want python's GC to close it. + sys.stdout = ray._private.utils.open_log( + stdout_fileno, unbuffered=True, closefd=False + ) + sys.stderr = ray._private.utils.open_log( + stderr_fileno, unbuffered=True, closefd=False + ) + + # Setup log file. + out_filepath, err_filepath = node.get_log_file_names( + get_worker_log_file_name(args.worker_type), + unique=False, # C++ core worker process already creates the file, should use a deterministic function to get the same file path. + create_out=True, + create_err=True, + ) + worker.set_out_file(out_filepath) + worker.set_err_file(err_filepath) + + rotation_max_bytes = os.getenv("RAY_ROTATION_MAX_BYTES", None) + + # Log rotation is disabled on windows platform. + if sys.platform != "win32" and rotation_max_bytes and int(rotation_max_bytes) > 0: + worker.set_file_rotation_enabled(True) + + if mode == ray.WORKER_MODE and args.worker_preload_modules: + module_names_to_import = args.worker_preload_modules.split(",") + ray._private.utils.try_import_each_module(module_names_to_import) + + # If the worker setup function is configured, run it. + worker_process_setup_hook_key = os.getenv( + ray_constants.WORKER_PROCESS_SETUP_HOOK_ENV_VAR + ) + if worker_process_setup_hook_key: + error = load_and_execute_setup_hook(worker_process_setup_hook_key) + if error is not None: + worker.core_worker.drain_and_exit_worker("system", error) + + if mode == ray.WORKER_MODE: + worker.main_loop() + elif mode in [ray.RESTORE_WORKER_MODE, ray.SPILL_WORKER_MODE]: + # It is handled by another thread in the C++ core worker. + # We just need to keep the worker alive. + while True: + time.sleep(100000) + else: + raise ValueError(f"Unexcepted worker mode: {mode}") diff --git a/lib/python3.12/site-packages/ray/_private/workers/setup_worker.py b/lib/python3.12/site-packages/ray/_private/workers/setup_worker.py new file mode 100644 index 0000000000000000000000000000000000000000..23ba980a5bb22b2ce648c6df45f8ba0b7e73f3dd --- /dev/null +++ b/lib/python3.12/site-packages/ray/_private/workers/setup_worker.py @@ -0,0 +1,33 @@ +import argparse +import logging + +from ray._private.ray_constants import LOGGER_FORMAT, LOGGER_LEVEL +from ray._private.ray_logging import setup_logger +from ray._private.runtime_env.context import RuntimeEnvContext +from ray.core.generated.common_pb2 import Language + +logger = logging.getLogger(__name__) + +parser = argparse.ArgumentParser( + description=("Set up the environment for a Ray worker and launch the worker.") +) + +parser.add_argument( + "--serialized-runtime-env-context", + type=str, + help="the serialized runtime env context", +) + +parser.add_argument("--language", type=str, help="the language type of the worker") + + +if __name__ == "__main__": + setup_logger(LOGGER_LEVEL, LOGGER_FORMAT) + args, remaining_args = parser.parse_known_args() + # NOTE(edoakes): args.serialized_runtime_env_context is only None when + # we're starting the main Ray client proxy server. That case should + # probably not even go through this codepath. + runtime_env_context = RuntimeEnvContext.deserialize( + args.serialized_runtime_env_context or "{}" + ) + runtime_env_context.exec_worker(remaining_args, Language.Value(args.language))