o `۷i@sUddlZddlmZmZmZmZddlmZddlm Z ddl m Z m Z ddl mZmZddlmZmZmZer?ddlmZdad ed <eZd e Zd e Zed dGdddZed ddefddZdS)N) TYPE_CHECKINGAnyDictOptional)session)StorageContext)V2_MIGRATION_GUIDE_MESSAGE_v2_migration_warnings_enabled) _copy_doc_log_deprecation_warning) Deprecated DeveloperAPI PublicAPI)PlacementGroupFactoryzOptional[TrainContext]_default_contextz`get_metadata` was an experimental API that accessed the metadata passed to `Trainer(metadata=...)`. This API can be replaced by passing the metadata directly to the training function (e.g., via `train_loop_config`). z`{}` is deprecated because the concept of a `Trial` will soon be removed in Ray Train.Ray Train will no longer assume that it's running within a Ray Tune `Trial` in the future. stable) stabilityc@seZdZdZeejdefddZeejde fddZeej de fddZ eej de fd d Z eej de fd d Z eej de fd dZ eeejdefddZeeedeejdeeeffddZeededeejdefddZeededeejdefddZeededeejd"ddZeededeejdefdd Zd!S)# TrainContextzJContext containing metadata that can be accessed within Ray Train workers.returncCtSN)rget_experiment_nameselfrG/home/ubuntu/vllm_env/lib/python3.10/site-packages/ray/train/context.pyr*z TrainContext.get_experiment_namecCrr)rget_world_sizerrrrr.rzTrainContext.get_world_sizecCrr)rget_world_rankrrrrr2rzTrainContext.get_world_rankcCrr)rget_local_rankrrrrr6rzTrainContext.get_local_rankcCrr)rget_local_world_sizerrrrr :rz!TrainContext.get_local_world_sizecCrr)r get_node_rankrrrrr!>rzTrainContext.get_node_rankcCrr)r get_storagerrrrr"BszTrainContext.get_storage)messagewarningcCrr)r get_metadatarrrrr%IzTrainContext.get_metadataget_trial_namecCrr)rr'rrrrr'Qr&zTrainContext.get_trial_name get_trial_idcCrr)rr(rrrrr(Yr&zTrainContext.get_trial_idget_trial_resourcesrcCrr)rr)rrrrr)asz TrainContext.get_trial_resources get_trial_dircCrr)rr*rrrrr*kr&zTrainContext.get_trial_dirN)rr)__name__ __module__ __qualname____doc__r rrstrrintrrr r!r r"rr !_GET_METADATA_DEPRECATION_MESSAGEr r%rr*_TUNE_SPECIFIC_CONTEXT_DEPRECATION_MESSAGEformatr'r(r)r*rrrrr&sbrrcCsrddlm}|rddlm}trtdt|Sttdur&t atWdS1s2wYdS)zGet or create a singleton training context. The context is only available within a function passed to Ray Train. See the :class:`~ray.train.TrainContext` API reference to see available methods. r)_in_tune_session) get_contextz`ray.train.get_context()` should be switched to `ray.tune.get_context()` when running in a function passed to Ray Tune. This will be an error in the future. N) %ray.tune.trainable.trainable_fn_utilsr4ray.tuner5r r r _context_lockrr)r4get_tune_contextrrrr5ts  $r5) threadingtypingrrrrray.train._internalrray.train._internal.storagerray.train.constantsrr ray.train.utilsr r ray.util.annotationsr r r#ray.tune.execution.placement_groupsrr__annotations__Lockr8r1r2rr5rrrrs.      M